[mrs]{un,}zipsingletons

4 files changed, 21 insertions, 3 deletions

diff --git a/src/Data/Array/Nested.hs b/src/Data/Array/Nested.hs
index b825691..499fdda 100644
--- a/src/Data/Array/Nested.hs
+++ b/src/Data/Array/Nested.hs
@@ -14,7 +14,7 @@ module Data.Array.Nested (
   rfromListLinear, rfromListPrimLinear, rtoListLinear,
   rslice, rrev1, rreshape, rflatten, riota,
   rminIndexPrim, rmaxIndexPrim, rdot1Inner, rdot,
-  rnest, runNest,
+  rnest, runNest, rzip, runzip,
   -- ** Lifting orthotope operations to 'Ranked' arrays
   rlift, rlift2,
   -- ** Conversions
@@ -36,7 +36,7 @@ module Data.Array.Nested (
   sfromListLinear, sfromListPrimLinear, stoListLinear,
   sslice, srev1, sreshape, sflatten, siota,
   sminIndexPrim, smaxIndexPrim, sdot1Inner, sdot,
-  snest, sunNest,
+  snest, sunNest, szip, sunzip,
   -- ** Lifting orthotope operations to 'Shaped' arrays
   slift, slift2,
   -- ** Conversions
@@ -58,7 +58,7 @@ module Data.Array.Nested (
   mfromListLinear, mfromListPrimLinear, mtoListLinear,
   mslice, mrev1, mreshape, mflatten, miota,
   mminIndexPrim, mmaxIndexPrim, mdot1Inner, mdot,
-  mnest, munNest,
+  mnest, munNest, mzip, munzip,
   -- ** Lifting orthotope operations to 'Mixed' arrays
   mlift, mlift2,
   -- ** Conversions
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))