From dd21ebafeeeb2abc222b3525e136d4ba6a769176 Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Fri, 22 Nov 2024 18:54:01 +0100
Subject: [mrs]{un,}zip

---
 src/Data/Array/Nested/Internal/Mixed.hs  | 6 ++++++
 src/Data/Array/Nested/Internal/Ranked.hs | 6 ++++++
 src/Data/Array/Nested/Internal/Shaped.hs | 6 ++++++
 3 files changed, 18 insertions(+)

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

diff --git a/src/Data/Array/Nested/Internal/Mixed.hs b/src/Data/Array/Nested/Internal/Mixed.hs
index 8d239cf..619d9bc 100644
--- a/src/Data/Array/Nested/Internal/Mixed.hs
+++ b/src/Data/Array/Nested/Internal/Mixed.hs
@@ -790,6 +790,12 @@ mnest ssh arr = M_Nest (fst (shxSplitApp (Proxy @sh') ssh (mshape arr))) arr
 munNest :: Mixed sh (Mixed sh' a) -> Mixed (sh ++ sh') a
 munNest (M_Nest _ arr) = arr
 
+mzip :: Mixed sh a -> Mixed sh b -> Mixed sh (a, b)
+mzip = M_Tup2
+
+munzip :: Mixed sh (a, b) -> (Mixed sh a, Mixed sh b)
+munzip (M_Tup2 a b) = (a, b)
+
 mrerankP :: forall sh1 sh2 sh a b. (Storable a, Storable b)
          => StaticShX sh -> IShX sh2
          -> (Mixed sh1 (Primitive a) -> Mixed sh2 (Primitive b))
diff --git a/src/Data/Array/Nested/Internal/Ranked.hs b/src/Data/Array/Nested/Internal/Ranked.hs
index 9e8a7b2..1966270 100644
--- a/src/Data/Array/Nested/Internal/Ranked.hs
+++ b/src/Data/Array/Nested/Internal/Ranked.hs
@@ -390,6 +390,12 @@ runNest rarr@(Ranked (M_Ranked (M_Nest _ arr)))
   | Refl <- lemReplicatePlusApp (rrank rarr) (Proxy @m) (Proxy @(Nothing @Nat))
   = Ranked arr
 
+rzip :: Ranked n a -> Ranked n b -> Ranked n (a, b)
+rzip = coerce mzip
+
+runzip :: Ranked n (a, b) -> (Ranked n a, Ranked n b)
+runzip = coerce munzip
+
 rrerankP :: forall n1 n2 n a b. (Storable a, Storable b)
          => SNat n -> IShR n2
          -> (Ranked n1 (Primitive a) -> Ranked n2 (Primitive b))
diff --git a/src/Data/Array/Nested/Internal/Shaped.hs b/src/Data/Array/Nested/Internal/Shaped.hs
index 228d800..5d3eba1 100644
--- a/src/Data/Array/Nested/Internal/Shaped.hs
+++ b/src/Data/Array/Nested/Internal/Shaped.hs
@@ -376,6 +376,12 @@ sunNest sarr@(Shaped (M_Shaped (M_Nest _ arr)))
   | Refl <- lemMapJustApp (sshape sarr) (Proxy @sh')
   = Shaped arr
 
+szip :: Shaped sh a -> Shaped sh b -> Shaped sh (a, b)
+szip = coerce mzip
+
+sunzip :: Shaped sh (a, b) -> (Shaped sh a, Shaped sh b)
+sunzip = coerce munzip
+
 srerankP :: forall sh1 sh2 sh a b. (Storable a, Storable b)
          => ShS sh -> ShS sh2
          -> (Shaped sh1 (Primitive a) -> Shaped sh2 (Primitive b))
-- 
cgit v1.2.3-70-g09d2