SPEC magicSPEC-magic

1 files changed, 9 insertions, 3 deletions

diff --git a/src/Data/Array/Nested/Mixed/Shape/Internal.hs b/src/Data/Array/Nested/Mixed/Shape/Internal.hs
index 2a86ac1..ee97257 100644
--- a/src/Data/Array/Nested/Mixed/Shape/Internal.hs
+++ b/src/Data/Array/Nested/Mixed/Shape/Internal.hs
@@ -3,6 +3,7 @@
 {-# LANGUAGE TemplateHaskell #-}
 module Data.Array.Nested.Mixed.Shape.Internal where
 
+import GHC.Exts (SPEC(SPEC))
 import Language.Haskell.TH
 
 
@@ -24,9 +25,14 @@ ixFromLinearStub fname' ishty ixty zshC consshC ixz ixcons shtolist = do
     -- function and realise that 'suffixes' is shared. But then later inline it
     -- anyway, to avoid the function call. Removing the pragma makes GHC
     -- somehow unable to recognise that 'suffixes' can be shared in a loop.
+    -- Making specialization more aggressive via SPEC helps inline the outer
+    -- function, lowering runtime by 20% in benchmark "ixxFromLinear 10000x"
+    -- with ghc-9.15.20251127. A more brutal way of obtaining the same result
+    -- is setting INLINE for the outer function(s), but the function code is
+    -- rather large, so this may be counterproductive in some contexts.
     {-# NOINLINE [0] fromLin0 #-}
-    fromLin0 :: Num i => $ishty sh -> [Int] -> Int -> $ixty sh i
-    fromLin0 sh suffixes i =
+    fromLin0 :: Num i => SPEC -> $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
@@ -53,7 +59,7 @@ ixFromLinearStub fname' ishty ixty zshC consshC ixz ixcons shtolist = do
     \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 fromLin0 sh suffixes |]
+      in fromLin0 SPEC sh suffixes |]
 
   return [SigD fname typesig
          ,FunD fname [Clause [] (NormalB body) locals]]