Commit Tom's new code for sh?Enum functions

3 files changed, 39 insertions, 14 deletions

diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
index 0464f1f..1b008e5 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -1,9 +1,11 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE NoStarIsType #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE PolyKinds #-}
@@ -17,6 +19,7 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE UndecidableInstances #-}
 {-# LANGUAGE ViewPatterns #-}
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
@@ -32,7 +35,7 @@ import Data.Functor.Product
 import Data.Kind (Constraint, Type)
 import Data.Monoid (Sum(..))
 import Data.Type.Equality
-import GHC.Exts (withDict)
+import GHC.Exts (Int(..), Int#, quotRemInt#, withDict)
 import GHC.Generics (Generic)
 import GHC.IsList (IsList)
 import GHC.IsList qualified as IsList
@@ -476,11 +479,16 @@ shxEnum = shxEnum'
 
 {-# INLINABLE shxEnum' #-}  -- ensure this can be specialised at use site
 shxEnum' :: Num i => IShX sh -> [IxX sh i]
-shxEnum' = \sh -> go sh id []
+shxEnum' sh = [fromLin sh suffixes li# | I# li# <- [0 .. shxSize sh - 1]]
   where
-    go :: Num i => IShX sh -> (IxX sh i -> a) -> [a] -> [a]
-    go ZSX f = (f ZIX :)
-    go (n :$% sh) f = foldr (.) id [go sh (f . (fromIntegral i :.%)) | i <- [0 .. fromSMayNat' n - 1]]
+    suffixes = drop 1 (scanr (*) 1 (shxToList sh))
+
+    fromLin :: Num i => IShX sh -> [Int] -> Int# -> IxX sh i
+    fromLin ZSX _ _ = ZIX
+    fromLin (_ :$% sh') (I# suff# : suffs) i# =
+      let !(# q#, r# #) = i# `quotRemInt#` suff#  -- suff == shrSize sh'
+      in fromIntegral (I# q#) :.% fromLin sh' suffs r#
+    fromLin _ _ _ = error "impossible"
 
 shxCast :: StaticShX sh' -> IShX sh -> Maybe (IShX sh')
 shxCast ZKX ZSX = Just ZSX
diff --git a/src/Data/Array/Nested/Ranked/Shape.hs b/src/Data/Array/Nested/Ranked/Shape.hs
index 200fac5..50338d2 100644
--- a/src/Data/Array/Nested/Ranked/Shape.hs
+++ b/src/Data/Array/Nested/Ranked/Shape.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DeriveFoldable #-}
@@ -8,6 +9,7 @@
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE NoStarIsType #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE PolyKinds #-}
@@ -21,6 +23,7 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE UndecidableInstances #-}
 {-# LANGUAGE ViewPatterns #-}
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
@@ -33,6 +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#)
 import GHC.Generics (Generic)
 import GHC.IsList (IsList)
 import GHC.IsList qualified as IsList
@@ -335,11 +339,16 @@ shrEnum = shrEnum'
 
 {-# INLINABLE shrEnum' #-}  -- ensure this can be specialised at use site
 shrEnum' :: Num i => IShR sh -> [IxR sh i]
-shrEnum' = \sh -> go sh id []
+shrEnum' sh = [fromLin sh suffixes li# | I# li# <- [0 .. shrSize sh - 1]]
   where
-    go :: Num i => IShR sh -> (IxR sh i -> a) -> [a] -> [a]
-    go ZSR f = (f ZIR :)
-    go (n :$: sh) f = foldr (.) id [go sh (f . (fromIntegral i :.:)) | i <- [0 .. n - 1]]
+    suffixes = drop 1 (scanr (*) 1 (Foldable.toList sh))
+
+    fromLin :: Num i => IShR sh -> [Int] -> Int# -> IxR sh i
+    fromLin ZSR _ _ = ZIR
+    fromLin (_ :$: sh') (I# suff# : suffs) i# =
+      let !(# q#, r# #) = i# `quotRemInt#` suff#  -- suff == shrSize sh'
+      in fromIntegral (I# q#) :.: fromLin sh' suffs r#
+    fromLin _ _ _ = error "impossible"
 
 
 -- | Untyped: length is checked at runtime.
diff --git a/src/Data/Array/Nested/Shaped/Shape.hs b/src/Data/Array/Nested/Shaped/Shape.hs
index 582cc50..218caaa 100644
--- a/src/Data/Array/Nested/Shaped/Shape.hs
+++ b/src/Data/Array/Nested/Shaped/Shape.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DeriveGeneric #-}
@@ -5,6 +6,7 @@
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE NoStarIsType #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE PolyKinds #-}
@@ -18,6 +20,7 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE UndecidableInstances #-}
 {-# LANGUAGE ViewPatterns #-}
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
@@ -34,7 +37,7 @@ import Data.Kind (Constraint, Type)
 import Data.Monoid (Sum(..))
 import Data.Proxy
 import Data.Type.Equality
-import GHC.Exts (withDict)
+import GHC.Exts (Int(..), Int#, quotRemInt#, withDict)
 import GHC.Generics (Generic)
 import GHC.IsList (IsList)
 import GHC.IsList qualified as IsList
@@ -394,11 +397,16 @@ shsEnum = shsEnum'
 
 {-# INLINABLE shsEnum' #-}  -- ensure this can be specialised at use site
 shsEnum' :: Num i => ShS sh -> [IxS sh i]
-shsEnum' = \sh -> go sh id []
+shsEnum' sh = [fromLin sh suffixes li# | I# li# <- [0 .. shsSize sh - 1]]
   where
-    go :: Num i => ShS sh -> (IxS sh i -> a) -> [a] -> [a]
-    go ZSS f = (f ZIS :)
-    go (n :$$ sh) f = foldr (.) id [go sh (f . (fromIntegral i :.$)) | i <- [0 .. fromSNat' n - 1]]
+    suffixes = drop 1 (scanr (*) 1 (shsToList sh))
+
+    fromLin :: Num i => ShS sh -> [Int] -> Int# -> IxS sh i
+    fromLin ZSS _ _ = ZIS
+    fromLin (_ :$$ sh') (I# suff# : suffs) i# =
+      let !(# q#, r# #) = i# `quotRemInt#` suff#  -- suff == shsSize sh'
+      in fromIntegral (I# q#) :.$ fromLin sh' suffs r#
+    fromLin _ _ _ = error "impossible"
 
 
 -- | Untyped: length is checked at runtime.