Generalize mgeneratePrim to potentially avoid @fmap fromIntegral@

1 files changed, 4 insertions, 4 deletions

diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
index 066ae8e..ed03310 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -275,7 +275,7 @@ ixxZipWith _ ZIX ZIX = ZIX
 ixxZipWith f (i :.% is) (j :.% js) = f i j :.% ixxZipWith f is js
 
 {-# INLINEABLE ixxFromLinear #-}
-ixxFromLinear :: IShX sh -> Int -> IIxX sh
+ixxFromLinear :: Num i => IShX sh -> Int -> IxX sh i
 ixxFromLinear = \sh ->  -- give this function arity 1 so that suffixes is shared when it's called many times
   let suffixes = drop 1 (scanr (*) 1 (shxToList sh))
   in \i ->
@@ -286,14 +286,14 @@ ixxFromLinear = \sh ->  -- give this function arity 1 so that suffixes is shared
         (n :$% sh', suff : suffs) ->
           let (q, r) = i `quotRem` suff
           in if q >= fromSMayNat' n then outrange sh i else
-             q :.% fromLin sh' suffs r
+             fromIntegral q :.% fromLin sh' suffs r
         _ -> error "impossible"
   where
-    fromLin :: IShX sh -> [Int] -> Int -> IxX sh Int
+    fromLin :: Num i => IShX sh -> [Int] -> Int -> IxX sh i
     fromLin ZSX _ !_ = ZIX
     fromLin (_ :$% sh') (suff : suffs) i =
       let (q, r) = i `quotRem` suff  -- suff == shrSize sh'
-      in q :.% fromLin sh' suffs r
+      in fromIntegral q :.% fromLin sh' suffs r
     fromLin _ _ _ = error "impossible"
 
     {-# NOINLINE outrange #-}