De-triplicate sh?Enum

1 files changed, 10 insertions, 18 deletions

diff --git a/src/Data/Array/Nested/Ranked/Shape.hs b/src/Data/Array/Nested/Ranked/Shape.hs
index 59289fb..0ac980e 100644
--- a/src/Data/Array/Nested/Ranked/Shape.hs
+++ b/src/Data/Array/Nested/Ranked/Shape.hs
@@ -1,4 +1,3 @@
-{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE FlexibleInstances #-}
@@ -34,7 +33,7 @@ import Data.Foldable qualified as Foldable
 import Data.Kind (Type)
 import Data.Proxy
 import Data.Type.Equality
-import GHC.Exts (Int(..), Int#, build, quotRemInt#)
+import GHC.Exts (build)
 import GHC.IsList (IsList)
 import GHC.IsList qualified as IsList
 import GHC.TypeLits
@@ -309,6 +308,15 @@ ixrFromLinear (ShR sh) i
 ixrFromIxX :: IxX sh i -> IxR (Rank sh) i
 ixrFromIxX = unsafeCoerce
 
+shrEnum :: IShR n -> [IIxR n]
+shrEnum = shrEnum'
+
+{-# INLINABLE shrEnum' #-}  -- ensure this can be specialised at use site
+shrEnum' :: forall i n. Num i => IShR n -> [IxR n i]
+shrEnum' (ShR sh)
+  | Refl <- lemRankReplicate (Proxy @n)
+  = (unsafeCoerce :: [IxX (Replicate n Nothing) i] -> [IxR n i]) $ shxEnum' sh
+
 
 -- * Ranked shapes
 
@@ -472,22 +480,6 @@ shrPermutePrefix = \perm sh ->
           EQI -> shrIndex si l :$: applyPermRFull sm perm l
           GTI -> error "shrPermutePrefix: Index in permutation out of range"
 
-shrEnum :: IShR sh -> [IIxR sh]
-shrEnum = shrEnum'
-
-{-# INLINABLE shrEnum' #-}  -- ensure this can be specialised at use site
-shrEnum' :: Num i => IShR sh -> [IxR sh i]
-shrEnum' sh = [fromLin sh suffixes li# | I# li# <- [0 .. shrSize sh - 1]]
-  where
-    suffixes = drop 1 (scanr (*) 1 (shrToList 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.
 instance KnownNat n => IsList (ListR n i) where