Optimize slightly sh?FromList and sh?ToList

1 files changed, 17 insertions, 2 deletions

diff --git a/src/Data/Array/Nested/Ranked/Shape.hs b/src/Data/Array/Nested/Ranked/Shape.hs
index 74f4c5c..2f20e1a 100644
--- a/src/Data/Array/Nested/Ranked/Shape.hs
+++ b/src/Data/Array/Nested/Ranked/Shape.hs
@@ -390,12 +390,27 @@ shrRank (ShR sh) | Refl <- lemRankReplicate (Proxy @n) = shxRank sh
 shrSize :: IShR n -> Int
 shrSize (ShR sh) = shxSize sh
 
+-- This is equivalent to but faster than @coerce (shxFromList (ssxReplicate snat))@.
+-- We don't report the size of the list in case of errors in order not to retain the list.
+{-# INLINEABLE shrFromList #-}
 shrFromList :: SNat n -> [Int] -> IShR n
-shrFromList snat = coerce (shxFromList (ssxReplicate snat))
+shrFromList snat topl = ShR $ ShX $ go snat topl
+  where
+    go :: SNat n -> [Int] -> ListH (Replicate n Nothing) Int
+    go SZ [] = ZH
+    go SZ _ = error $ "shrFromList: List too long (type says " ++ show (fromSNat' snat) ++ ")"
+    go (SS sn :: SNat n1) (i : is) | Refl <- lemReplicateSucc2 (Proxy @n1) Refl = ConsUnknown i (go sn is)
+    go _ _ = error $ "shrFromList: List too short (type says " ++ show (fromSNat' snat) ++ ")"
 
+-- This is equivalent to but faster than @coerce shxToList@.
 {-# INLINEABLE shrToList #-}
 shrToList :: IShR n -> [Int]
-shrToList = coerce shxToList
+shrToList (ShR (ShX l)) = build (\(cons :: i -> is -> is) (nil :: is) ->
+  let go :: ListH sh Int -> is
+      go ZH = nil
+      go (ConsUnknown i rest) = i `cons` go rest
+      go ConsKnown{} = error "shrToList: impossible case"
+  in go l)
 
 shrHead :: forall n i. ShR (n + 1) i -> i
 shrHead (ShR sh)