Add a couple of new zip operations that obviously make sense

1 files changed, 7 insertions, 2 deletions

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