From 429416f327a94947c0d42ccea8906cd22bae64b4 Mon Sep 17 00:00:00 2001
From: Mikolaj Konarski <mikolaj.konarski@funktory.com>
Date: Wed, 17 Dec 2025 12:51:46 +0100
Subject: Optimize slightly sh?FromList and sh?ToList

---
 src/Data/Array/Nested/Ranked/Shape.hs | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

(limited to 'src/Data/Array/Nested/Ranked')

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)
-- 
cgit v1.2.3-70-g09d2