Optimize slightly sh?FromList and sh?ToList

3 files changed, 52 insertions, 30 deletions

diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
index fd8c4ce..802c71e 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -584,27 +584,28 @@ shxSize :: IShX sh -> Int
 shxSize ZSX = 1
 shxSize (n :$% sh) = fromSMayNat' n * shxSize sh
 
+-- We don't report the size of the list in case of errors in order not to retain the list.
+{-# INLINEABLE shxFromList #-}
 shxFromList :: StaticShX sh -> [Int] -> IShX sh
-shxFromList topssh topl = go topssh topl
+shxFromList (StaticShX topssh) topl = ShX $ go topssh topl
   where
-    go :: StaticShX sh' -> [Int] -> IShX sh'
-    go ZKX [] = ZSX
-    go (SKnown sn :!% sh) (i : is)
-      | i == fromSNat' sn = SKnown sn :$% go sh is
-      | otherwise = error $ "shxFromList: Value does not match typing (type says "
-                              ++ show (fromSNat' sn) ++ ", list contains " ++ show i ++ ")"
-    go (SUnknown () :!% sh) (i : is) = SUnknown i :$% go sh is
-    go _ _ = error $ "shxFromList: Mismatched list length (type says "
-                       ++ show (ssxLength topssh) ++ ", list has length "
-                       ++ show (length topl) ++ ")"
+    go :: ListH sh' () -> [Int] -> ListH sh' Int
+    go ZH [] = ZH
+    go ZH _ = error $ "shxFromList: List too long (type says " ++ show (listhLength topssh) ++ ")"
+    go (ConsKnown sn sh) (i : is)
+      | i == fromSNat' sn = ConsKnown sn (go sh is)
+      | otherwise = error $ "shxFromList: Value does not match typing"
+    go (ConsUnknown () sh) (i : is) = ConsUnknown i (go sh is)
+    go _ _ = error $ "shxFromList: List too short (type says " ++ show (listhLength topssh) ++ ")"
 
 {-# INLINEABLE shxToList #-}
 shxToList :: IShX sh -> [Int]
-shxToList sh0 = build (\(cons :: i -> is -> is) (nil :: is) ->
-  let go :: IShX sh -> is
-      go ZSX = nil
-      go (smn :$% sh) = fromSMayNat' smn `cons` go sh
-  in go sh0)
+shxToList (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 sn rest) = fromSNat' sn `cons` go rest
+  in go l)
 
 -- If it ever matters for performance, this is unsafeCoercible.
 shxFromSSX :: StaticShX (MapJust sh) -> ShX (MapJust sh) i
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)
diff --git a/src/Data/Array/Nested/Shaped/Shape.hs b/src/Data/Array/Nested/Shaped/Shape.hs
index 634c854..afd2227 100644
--- a/src/Data/Array/Nested/Shaped/Shape.hs
+++ b/src/Data/Array/Nested/Shaped/Shape.hs
@@ -402,24 +402,30 @@ shsSize :: ShS sh -> Int
 shsSize (ShS sh) = shxSize sh
 
 -- | This is a partial @const@ that fails when the second argument
--- doesn't match the first. It also has a better error message comparing
--- to just coercing 'shxFromList'.
+-- doesn't match the first. We don't report the size of the list
+-- in case of errors in order not to retain the list.
+{-# INLINEABLE shsFromList #-}
 shsFromList :: ShS sh -> [Int] -> ShS sh
-shsFromList topsh topl = go topsh topl `seq` topsh
+shsFromList sh0@(ShS (ShX topsh)) topl = go topsh topl `seq` sh0
   where
-    go :: ShS sh' -> [Int] -> ()
-    go ZSS [] = ()
-    go (sn :$$ sh) (i : is)
+    go :: ListH sh' Int -> [Int] -> ()
+    go ZH [] = ()
+    go ZH _ = error $ "shsFromList: List too long (type says " ++ show (listhLength topsh) ++ ")"
+    go (ConsKnown sn sh) (i : is)
       | i == fromSNat' sn = go sh is
-      | otherwise = error $ "shsFromList: Value does not match typing (type says "
-                            ++ show (fromSNat' sn) ++ ", list contains " ++ show i ++ ")"
-    go _ _ = error $ "shsFromList: Mismatched list length (type says "
-                     ++ show (shsLength topsh) ++ ", list has length "
-                     ++ show (length topl) ++ ")"
+      | otherwise = error $ "shsFromList: Value does not match typing"
+    go ConsUnknown{} _ = error "shsFromList: impossible case"
+    go _ _ = error $ "shsFromList: List too short (type says " ++ show (listhLength topsh) ++ ")"
 
+-- This is equivalent to but faster than @coerce shxToList@.
 {-# INLINEABLE shsToList #-}
 shsToList :: ShS sh -> [Int]
-shsToList = coerce shxToList
+shsToList (ShS (ShX l)) = build (\(cons :: i -> is -> is) (nil :: is) ->
+  let go :: ListH sh Int -> is
+      go ZH = nil
+      go ConsUnknown{} = error "shsToList: impossible case"
+      go (ConsKnown snat rest) = fromSNat' snat `cons` go rest
+  in go l)
 
 shsHead :: ShS (n : sh) -> SNat n
 shsHead (ShS shx) = case shxHead shx of