Commit Tom's new code for sh?Enum functions

1 files changed, 13 insertions, 5 deletions

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.