Benchmark and improve ixxFromLinear

2 files changed, 35 insertions, 20 deletions

diff --git a/bench/Main.hs b/bench/Main.hs
index 8df09a9..01d9a3f 100644
--- a/bench/Main.hs
+++ b/bench/Main.hs
@@ -17,19 +17,12 @@ import Text.Show (showListWith)
 
 import Data.Array.Nested
 import Data.Array.Nested.Mixed (Mixed(M_Primitive), mliftPrim, mliftPrim2, toPrimitive)
+import Data.Array.Nested.Mixed.Shape
 import Data.Array.Nested.Ranked (liftRanked1, liftRanked2)
 import Data.Array.Strided.Arith.Internal qualified as Arith
 import Data.Array.XArray (XArray(..))
 
 
-enableMisc :: Bool
-enableMisc = False
-
-bgroupIf :: Bool -> String -> [Benchmark] -> Benchmark
-bgroupIf True = bgroup
-bgroupIf False = \name _ -> bgroup name []
-
-
 main :: IO ()
 main = do
   let enable = False
@@ -104,7 +97,7 @@ main_tests = defaultMain
       in nf (\a -> RS.normalize a)
             (RS.rev [0] (RS.rev [0] (RS.iota @Double n)))
     ]
-  ,bgroupIf enableMisc "misc"
+  ,bgroup "misc"
     [let n = 1000
          k = 1000
      in bgroup ("fusion [" ++ show k ++ "]*" ++ show n)
@@ -148,6 +141,16 @@ main_tests = defaultMain
           | ki <- [1::Int ..5]]
         ]
       ]
+    ,bench "ixxFromLinear 10000x" $
+      let n = 10000
+          sh0 = SUnknown 10 :$% SUnknown 10 :$% SUnknown 10 :$% SUnknown 10 :$% SUnknown 10 :$% ZSX
+      in nf (\sh -> [ixxFromLinear sh i | i <- [1..n]]) sh0
+    ,bench "ixxFromLinear 1x" $
+      let sh0 = SUnknown 10 :$% SUnknown 10 :$% SUnknown 10 :$% SUnknown 10 :$% SUnknown 10 :$% ZSX
+      in nf (\sh -> ixxFromLinear sh 1234) sh0
+    ,bench "shxEnum" $
+      let sh0 = SUnknown 10 :$% SUnknown 10 :$% SUnknown 10 :$% SUnknown 10 :$% SUnknown 10 :$% ZSX
+      in nf (\sh -> shxEnum sh) sh0
     ]
   ]
 
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)