SPEC magicSPEC-magic

4 files changed, 12 insertions, 6 deletions

diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
index c999853..8aadd7f 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -36,7 +36,7 @@ import Data.Functor.Product
 import Data.Kind (Constraint, Type)
 import Data.Monoid (Sum(..))
 import Data.Type.Equality
-import GHC.Exts (Int(..), Int#, quotRemInt#, withDict, build)
+import GHC.Exts (Int(..), Int#, build, quotRemInt#, withDict)
 import GHC.Generics (Generic)
 import GHC.IsList (IsList)
 import GHC.IsList qualified as IsList
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]]
diff --git a/src/Data/Array/Nested/Ranked/Shape.hs b/src/Data/Array/Nested/Ranked/Shape.hs
index 6d61bd5..6999b4d 100644
--- a/src/Data/Array/Nested/Ranked/Shape.hs
+++ b/src/Data/Array/Nested/Ranked/Shape.hs
@@ -36,7 +36,7 @@ import Data.Foldable qualified as Foldable
 import Data.Kind (Type)
 import Data.Proxy
 import Data.Type.Equality
-import GHC.Exts (Int(..), Int#, quotRemInt#, build)
+import GHC.Exts (Int(..), Int#, build, quotRemInt#)
 import GHC.Generics (Generic)
 import GHC.IsList (IsList)
 import GHC.IsList qualified as IsList
diff --git a/src/Data/Array/Nested/Shaped/Shape.hs b/src/Data/Array/Nested/Shaped/Shape.hs
index 0d90e91..0f43126 100644
--- a/src/Data/Array/Nested/Shaped/Shape.hs
+++ b/src/Data/Array/Nested/Shaped/Shape.hs
@@ -38,7 +38,7 @@ import Data.Kind (Constraint, Type)
 import Data.Monoid (Sum(..))
 import Data.Proxy
 import Data.Type.Equality
-import GHC.Exts (Int(..), Int#, quotRemInt#, withDict, build)
+import GHC.Exts (Int(..), Int#, build, quotRemInt#, withDict)
 import GHC.Generics (Generic)
 import GHC.IsList (IsList)
 import GHC.IsList qualified as IsList