From 31eff15c7f8f062748cda6c7b86c5b38818d9035 Mon Sep 17 00:00:00 2001
From: Mikolaj Konarski <mikolaj.konarski@funktory.com>
Date: Tue, 13 May 2025 16:37:23 +0200
Subject: Add a couple of new zip operations that obviously make sense

---
 src/Data/Array/Mixed/Shape.hs           | 9 +++++++--
 src/Data/Array/Nested/Internal/Shape.hs | 6 ++++++
 2 files changed, 13 insertions(+), 2 deletions(-)

(limited to 'src')

diff --git a/src/Data/Array/Mixed/Shape.hs b/src/Data/Array/Mixed/Shape.hs
index 809ed9e..26f938d 100644
--- a/src/Data/Array/Mixed/Shape.hs
+++ b/src/Data/Array/Mixed/Shape.hs
@@ -234,11 +234,11 @@ ixxInit = coerce (listxInit @(Const i))
 ixxLast :: forall n sh i. IxX (n : sh) i -> i
 ixxLast = coerce (listxLast @(Const i))
 
-ixxZip :: IxX n i -> IxX n j -> IxX n (i, j)
+ixxZip :: IxX sh i -> IxX sh j -> IxX sh (i, j)
 ixxZip ZIX ZIX = ZIX
 ixxZip (i :.% is) (j :.% js) = (i, j) :.% ixxZip is js
 
-ixxZipWith :: (i -> j -> k) -> IxX n i -> IxX n j -> IxX n k
+ixxZipWith :: (i -> j -> k) -> IxX sh i -> IxX sh j -> IxX sh k
 ixxZipWith _ ZIX ZIX = ZIX
 ixxZipWith f (i :.% is) (j :.% js) = f i j :.% ixxZipWith f is js
 
@@ -426,6 +426,11 @@ shxTakeSSX _ = flip go
     go ZKX _ = ZSX
     go (_ :!% ssh1) (n :$% sh) = n :$% go ssh1 sh
 
+shxZipWith :: (forall n. SMayNat i SNat n -> SMayNat j SNat n -> SMayNat k SNat n)
+           -> ShX sh i -> ShX sh j -> ShX sh k
+shxZipWith _ ZSX ZSX = ZSX
+shxZipWith f (i :$% is) (j :$% js) = f i j :$% shxZipWith f is js
+
 -- This is a weird operation, so it has a long name
 shxCompleteZeros :: StaticShX sh -> IShX sh
 shxCompleteZeros ZKX = ZSX
diff --git a/src/Data/Array/Nested/Internal/Shape.hs b/src/Data/Array/Nested/Internal/Shape.hs
index 82b7966..490c9a8 100644
--- a/src/Data/Array/Nested/Internal/Shape.hs
+++ b/src/Data/Array/Nested/Internal/Shape.hs
@@ -333,6 +333,12 @@ shrLast (ShR list) = listrLast list
 shrAppend :: forall n m i. ShR n i -> ShR m i -> ShR (n + m) i
 shrAppend = coerce (listrAppend @_ @i)
 
+shrZip :: ShR n i -> ShR n j -> ShR n (i, j)
+shrZip (ShR l1) (ShR l2) = ShR $ listrZip l1 l2
+
+shrZipWith :: (i -> j -> k) -> ShR n i -> ShR n j -> ShR n k
+shrZipWith f (ShR l1) (ShR l2) = ShR $ listrZipWith f l1 l2
+
 shrPermutePrefix :: forall n i. [Int] -> ShR n i -> ShR n i
 shrPermutePrefix = coerce (listrPermutePrefix @i)
 
-- 
cgit v1.2.3-70-g09d2