Optimize slightly sh?FromList and sh?ToList

1 files changed, 18 insertions, 12 deletions

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