Benchmark and improve ixxFromLinear

1 files changed, 23 insertions, 11 deletions

diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
index 1b008e5..f127e3a 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -262,18 +262,30 @@ ixxZipWith _ ZIX ZIX = ZIX
 ixxZipWith f (i :.% is) (j :.% js) = f i j :.% ixxZipWith f is js
 
 ixxFromLinear :: IShX sh -> Int -> IIxX sh
-ixxFromLinear = \sh i -> case go sh i of
-  (idx, 0) -> idx
-  _ -> error $ "ixxFromLinear: out of range (" ++ show i ++
-               " in array of shape " ++ show sh ++ ")"
+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 ->
+      if i < 0 then outrange sh i else
+      case (sh, suffixes) of  -- unfold first iteration of fromLin to do the range check
+        (ZSX, _) | i > 0 -> outrange sh i
+                 | otherwise -> ZIX
+        (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
+        _ -> error "impossible"
   where
-    -- returns (index in subarray, remaining index in enclosing array)
-    go :: IShX sh -> Int -> (IIxX sh, Int)
-    go ZSX i = (ZIX, i)
-    go (n :$% sh) i =
-      let (idx, i') = go sh i
-          (upi, locali) = i' `quotRem` fromSMayNat' n
-      in (locali :.% idx, upi)
+    fromLin :: IShX sh -> [Int] -> Int -> IxX sh Int
+    fromLin ZSX _ !_ = ZIX
+    fromLin (_ :$% sh') (suff : suffs) i =
+      let (q, r) = i `quotRem` suff  -- suff == shrSize sh'
+      in q :.% fromLin sh' suffs r
+    fromLin _ _ _ = error "impossible"
+
+    {-# NOINLINE outrange #-}
+    outrange :: IShX sh -> Int -> a
+    outrange sh i = error $ "ixxFromLinear: out of range (" ++ show i ++
+                            " in array of shape " ++ show sh ++ ")"
 
 ixxToLinear :: IShX sh -> IIxX sh -> Int
 ixxToLinear = \sh i -> fst (go sh i)