De-triplicate sh?Enum

1 files changed, 17 insertions, 16 deletions

diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
index 057418e..7422b0a 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -328,6 +328,23 @@ ixxFromLinear = \sh ->  -- give this function arity 1 so that suffixes is shared
     outrange sh i = error $ "ixxFromLinear: out of range (" ++ show i ++
                             " in array of shape " ++ show sh ++ ")"
 
+shxEnum :: IShX sh -> [IIxX sh]
+shxEnum = shxEnum'
+
+{-# INLINABLE shxEnum' #-}  -- ensure this can be specialised at use site
+shxEnum' :: Num i => IShX sh -> [IxX sh i]
+shxEnum' sh = [fromLin sh suffixes li# | I# li# <- [0 .. shxSize sh - 1]]
+  where
+    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"
+
+
 -- * Mixed shape-like lists to be used for ShX and StaticShX
 
 data SMayNat i n where
@@ -649,22 +666,6 @@ shxSplitApp :: proxy sh' -> StaticShX sh -> ShX (sh ++ sh') i -> (ShX sh i, ShX
 shxSplitApp _ ZKX idx = (ZSX, idx)
 shxSplitApp p (_ :!% ssh) (i :$% idx) = first (i :$%) (shxSplitApp p ssh idx)
 
-shxEnum :: IShX sh -> [IIxX sh]
-shxEnum = shxEnum'
-
-{-# INLINABLE shxEnum' #-}  -- ensure this can be specialised at use site
-shxEnum' :: Num i => IShX sh -> [IxX sh i]
-shxEnum' sh = [fromLin sh suffixes li# | I# li# <- [0 .. shxSize sh - 1]]
-  where
-    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
 shxCast (SKnown m    :!% ssh) (SKnown n   :$% sh) | Just Refl <- testEquality n m = (SKnown n :$%) <$> shxCast ssh sh