Restore suffixes sharing in *FromLinearHEADmaster

Thanks Mikolaj for spotting it's broken

1 files changed, 27 insertions, 20 deletions

diff --git a/src/Data/Array/Nested/Mixed/Shape/Internal.hs b/src/Data/Array/Nested/Mixed/Shape/Internal.hs
index cf44522..9997b0f 100644
--- a/src/Data/Array/Nested/Mixed/Shape/Internal.hs
+++ b/src/Data/Array/Nested/Mixed/Shape/Internal.hs
@@ -19,32 +19,39 @@ ixFromLinearStub fname' ishty ixty zshC consshC ixz ixcons shtolist = do
   typesig <- [t| forall i sh. Num i => $ishty sh -> Int -> $ixty sh i |]
 
   locals <- [d|
-      fromLin :: Num i => $ishty sh -> [Int] -> Int -> $ixty sh i
-      fromLin $zshC _ !_ = $ixz
-      fromLin ($(consshC wildP (varP (mkName "sh'")))) (suff : suffs) i =
-        let (q, r) = i `quotRem` suff  -- suff == shrSize sh'
-        in $ixcons (fromIntegral q) (fromLin sh' suffs r)
-      fromLin _ _ _ = error "impossible"
+    -- Unfold first iteration of fromLin to do the range check.
+    -- Don't inline because if this is inlined, GHC seems to stop sharing
+    -- 'suffixes' over multiple calls, which breaks performance in sh*Enum.
+    {-# NOINLINE fromLin0 #-}
+    fromLin0 :: Num i => $ishty sh -> [Int] -> Int -> $ixty sh i
+    fromLin0 sh suffixes i =
+        if i < 0 then outrange sh i else
+        case (sh, suffixes) of
+          ($zshC, _) | i > 0 -> outrange sh i
+                     | otherwise -> $ixz
+          ($(consshC (varP (mkName "n")) (varP (mkName "sh'"))), suff : suffs) ->
+            let (q, r) = i `quotRem` suff
+            in if q >= n then outrange sh i else
+               $ixcons (fromIntegral q) (fromLin sh' suffs r)
+          _ -> error "impossible"
 
-      {-# NOINLINE outrange #-}
-      outrange :: $ishty sh -> Int -> a
-      outrange sh i = error $ fname' ++ ": out of range (" ++ show i ++
-                              " in array of shape " ++ show sh ++ ")" |]
+    fromLin :: Num i => $ishty sh -> [Int] -> Int -> $ixty sh i
+    fromLin $zshC _ !_ = $ixz
+    fromLin ($(consshC wildP (varP (mkName "sh'")))) (suff : suffs) i =
+      let (q, r) = i `quotRem` suff  -- suff == shrSize sh'
+      in $ixcons (fromIntegral q) (fromLin sh' suffs r)
+    fromLin _ _ _ = error "impossible"
+
+    {-# NOINLINE outrange #-}
+    outrange :: $ishty sh -> Int -> a
+    outrange sh i = error $ fname' ++ ": out of range (" ++ show i ++
+                            " in array of shape " ++ show sh ++ ")" |]
 
   body <- [|
     \sh ->  -- give this function arity 1 so that 'suffixes' is shared when
             -- it's called many times
       let suffixes = drop 1 (scanr (*) 1 ($shtolist 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
-            ($zshC, _) | i > 0 -> outrange sh i
-                       | otherwise -> $ixz
-            ($(consshC (varP (mkName "n")) (varP (mkName "sh'"))), suff : suffs) ->
-              let (q, r) = i `quotRem` suff
-              in if q >= n then outrange sh i else
-                 $ixcons (fromIntegral q) (fromLin sh' suffs r)
-            _ -> error "impossible" |]
+      in fromLin0 sh suffixes |]
 
   return [SigD fname typesig
          ,FunD fname [Clause [] (NormalB body) locals]]