Optimize slightly sh?FromList and sh?ToList

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 7422b0a..4b04e0e 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -585,27 +585,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