Generalize ix?ToLinear and speed it up a bit

1 files changed, 7 insertions, 8 deletions

diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
index 145ea5f..11ef757 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -284,16 +284,15 @@ ixxZipWith :: (i -> j -> k) -> IxX sh i -> IxX sh j -> IxX sh k
 ixxZipWith _ ZIX ZIX = ZIX
 ixxZipWith f (i :.% is) (j :.% js) = f i j :.% ixxZipWith f is js
 
+-- | Given a multidimensional index, get the corresponding linear
+-- index into the buffer.
 {-# INLINEABLE ixxToLinear #-}
-ixxToLinear :: IShX sh -> IIxX sh -> Int
-ixxToLinear = \sh i -> fst (go sh i)
+ixxToLinear :: Num i => IShX sh -> IxX sh i -> i
+ixxToLinear = \sh i -> go sh i 0
   where
-    -- returns (index in subarray, size of subarray)
-    go :: IShX sh -> IIxX sh -> (Int, Int)
-    go ZSX ZIX = (0, 1)
-    go (n :$% sh) (i :.% ix) =
-      let (lidx, sz) = go sh ix
-      in (sz * i + lidx, fromSMayNat' n * sz)
+    go :: Num i => IShX sh -> IxX sh i -> i -> i
+    go ZSX ZIX a = a
+    go (n :$% sh) (i :.% ix) a = go sh ix (fromIntegral (fromSMayNat' n) * a + i)
 
 
 -- * Mixed shapes