Use lazilyConcat for listxAppend

3 files changed, 10 insertions, 10 deletions

diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
index 611ec19..54a945b 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -71,9 +71,8 @@ listxHead (i ::% _) = i
 listxTail :: ListX (n : sh) i -> ListX sh i
 listxTail (_ ::% sh) = sh
 
-listxAppend :: ListX sh i -> ListX sh' i -> ListX (sh ++ sh') i
-listxAppend ZX idx' = idx'
-listxAppend (i ::% idx) idx' = i ::% listxAppend idx idx'
+listxAppend :: forall sh sh' i. ListX sh i -> ListX sh' i -> ListX (sh ++ sh') i
+listxAppend = lazilyConcat (++)
 
 listxDrop :: forall i j sh sh'. ListX sh j -> ListX (sh ++ sh') i -> ListX sh' i
 listxDrop ZX long = long
@@ -147,7 +146,7 @@ ixxTail :: IxX (n : sh) i -> IxX sh i
 ixxTail (IxX list) = IxX (listxTail list)
 
 ixxAppend :: forall sh sh' i. IxX sh i -> IxX sh' i -> IxX (sh ++ sh') i
-ixxAppend = coerce (listxAppend @_ @i)
+ixxAppend = coerce (listxAppend @_ @_ @i)
 
 ixxDrop :: forall sh sh' i. IxX sh i -> IxX (sh ++ sh') i -> IxX sh' i
 ixxDrop = coerce (listxDrop @i @i)
diff --git a/src/Data/Array/Nested/Ranked/Shape.hs b/src/Data/Array/Nested/Ranked/Shape.hs
index 4eef090..687df69 100644
--- a/src/Data/Array/Nested/Ranked/Shape.hs
+++ b/src/Data/Array/Nested/Ranked/Shape.hs
@@ -119,9 +119,10 @@ listrRank :: ListR n i -> SNat n
 listrRank ZR = SNat
 listrRank (_ ::: sh) = snatSucc (listrRank sh)
 
+-- lemReplicatePlusApp requires SNat that would cause overhead (not benchmarked)
 listrAppend :: forall n m i. ListR n i -> ListR m i -> ListR (n + m) i
-listrAppend ZR sh = sh
-listrAppend (x ::: xs) sh = x ::: listrAppend xs sh
+listrAppend = gcastWith (unsafeCoerceRefl :: Replicate (n + m) (Nothing @Nat) :~: Replicate n Nothing ++ Replicate m Nothing) $
+              coerce (listxAppend @_ @_ @i)
 
 {-# INLINE listrFromList #-}
 listrFromList :: SNat n -> [i] -> ListR n i
diff --git a/src/Data/Array/Nested/Shaped/Shape.hs b/src/Data/Array/Nested/Shaped/Shape.hs
index 13596a7..fe55d1a 100644
--- a/src/Data/Array/Nested/Shaped/Shape.hs
+++ b/src/Data/Array/Nested/Shaped/Shape.hs
@@ -119,9 +119,9 @@ listsLast :: ListS (n : sh) i -> i
 listsLast (_ ::$ sh@(_ ::$ _)) = listsLast sh
 listsLast (n ::$ ZS) = n
 
-listsAppend :: ListS sh i -> ListS sh' i -> ListS (sh ++ sh') i
-listsAppend ZS idx' = idx'
-listsAppend (i ::$ idx) idx' = i ::$ listsAppend idx idx'
+listsAppend :: forall sh sh' i. ListS sh i -> ListS sh' i -> ListS (sh ++ sh') i
+listsAppend = gcastWith (unsafeCoerceRefl :: MapJust (sh ++ sh') :~: MapJust sh ++ MapJust sh') $
+              coerce (listxAppend @_ @_ @i)
 
 listsZip :: ListS sh i -> ListS sh j -> ListS sh (i, j)
 listsZip ZS ZS = ZS
@@ -219,7 +219,7 @@ ixsCast ZIS = ZIS
 ixsCast (i :.$ idx) = i :.$ ixsCast idx
 
 ixsAppend :: forall sh sh' i. IxS sh i -> IxS sh' i -> IxS (sh ++ sh') i
-ixsAppend = coerce (listsAppend @_ @i)
+ixsAppend = coerce (listsAppend @_ @_ @i)
 
 ixsZip :: IxS sh i -> IxS sh j -> IxS sh (i, j)
 ixsZip ZIS ZIS = ZIS