Commit Tom's new code for sh?Enum functions

1 files changed, 13 insertions, 5 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