De-triplicate sh?Enum

1 files changed, 8 insertions, 18 deletions

diff --git a/src/Data/Array/Nested/Shaped/Shape.hs b/src/Data/Array/Nested/Shaped/Shape.hs
index f57e7dd..39be729 100644
--- a/src/Data/Array/Nested/Shaped/Shape.hs
+++ b/src/Data/Array/Nested/Shaped/Shape.hs
@@ -1,4 +1,3 @@
-{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE FlexibleInstances #-}
@@ -37,7 +36,7 @@ import Data.Kind (Constraint, Type)
 import Data.Monoid (Sum(..))
 import Data.Proxy
 import Data.Type.Equality
-import GHC.Exts (Int(..), Int#, build, quotRemInt#, withDict)
+import GHC.Exts (build, withDict)
 import GHC.IsList (IsList)
 import GHC.IsList qualified as IsList
 import GHC.TypeLits
@@ -325,6 +324,13 @@ ixsFromLinear (ShS sh) i = ixsFromIxX $ ixxFromLinear sh i
 ixsFromIxX :: IxX (MapJust sh) i -> IxS sh i
 ixsFromIxX = unsafeCoerce
 
+shsEnum :: ShS sh -> [IIxS sh]
+shsEnum = shsEnum'
+
+{-# INLINABLE shsEnum' #-}  -- ensure this can be specialised at use site
+shsEnum' :: Num i => ShS sh -> [IxS sh i]
+shsEnum' (ShS sh) = (unsafeCoerce :: [IxX (MapJust sh) i] -> [IxS sh i]) $ shxEnum' sh
+
 
 -- * Shaped shapes
 
@@ -506,22 +512,6 @@ shsOrthotopeShape :: ShS sh -> Dict O.Shape sh
 shsOrthotopeShape ZSS = Dict
 shsOrthotopeShape (SNat :$$ sh) | Dict <- shsOrthotopeShape sh = Dict
 
-shsEnum :: ShS sh -> [IIxS sh]
-shsEnum = shsEnum'
-
-{-# INLINABLE shsEnum' #-}  -- ensure this can be specialised at use site
-shsEnum' :: Num i => ShS sh -> [IxS sh i]
-shsEnum' sh = [fromLin sh suffixes li# | I# li# <- [0 .. shsSize sh - 1]]
-  where
-    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.
 instance KnownShS sh => IsList (ListS sh (Const i)) where