Generalise list{x,s}Zip{,With} from Const to fHEADmaster

1 files changed, 12 insertions, 9 deletions

diff --git a/src/Data/Array/Nested/Internal/Shape.hs b/src/Data/Array/Nested/Internal/Shape.hs
index 438f049..82b7966 100644
--- a/src/Data/Array/Nested/Internal/Shape.hs
+++ b/src/Data/Array/Nested/Internal/Shape.hs
@@ -32,6 +32,7 @@ import Data.Array.Mixed.Types
 import Data.Coerce (coerce)
 import Data.Foldable qualified as Foldable
 import Data.Functor.Const
+import Data.Functor.Product qualified as Fun
 import Data.Kind (Type, Constraint)
 import Data.Monoid (Sum(..))
 import Data.Proxy
@@ -452,16 +453,16 @@ listsAppend :: ListS sh f -> ListS sh' f -> ListS (sh ++ sh') f
 listsAppend ZS idx' = idx'
 listsAppend (i ::$ idx) idx' = i ::$ listsAppend idx idx'
 
-listsZip :: ListS sh (Const i) -> ListS sh (Const j) -> ListS sh (Const (i, j))
+listsZip :: ListS sh f -> ListS sh g -> ListS sh (Fun.Product f g)
 listsZip ZS ZS = ZS
-listsZip (Const i ::$ irest) (Const j ::$ jrest) =
-  Const (i, j) ::$ listsZip irest jrest
+listsZip (i ::$ is) (j ::$ js) =
+  Fun.Pair i j ::$ listsZip is js
 
-listsZipWith :: (i -> j -> k) -> ListS sh (Const i) -> ListS sh (Const j)
-             -> ListS sh (Const k)
+listsZipWith :: (forall a. f a -> g a -> h a) -> ListS sh f -> ListS sh g
+             -> ListS sh h
 listsZipWith _ ZS ZS = ZS
-listsZipWith f (Const i ::$ irest) (Const j ::$ jrest) =
-  Const (f i j) ::$ listsZipWith f irest jrest
+listsZipWith f (i ::$ is) (j ::$ js) =
+  f i j ::$ listsZipWith f is js
 
 listsTakeLenPerm :: forall f is sh. Perm is -> ListS sh f -> ListS (TakeLen is sh) f
 listsTakeLenPerm PNil _ = ZS
@@ -560,10 +561,12 @@ ixsAppend :: forall sh sh' i. IxS sh i -> IxS sh' i -> IxS (sh ++ sh') i
 ixsAppend = coerce (listsAppend @_ @(Const i))
 
 ixsZip :: IxS n i -> IxS n j -> IxS n (i, j)
-ixsZip (IxS l1) (IxS l2) = IxS $ listsZip l1 l2
+ixsZip ZIS ZIS = ZIS
+ixsZip (i :.$ is) (j :.$ js) = (i, j) :.$ ixsZip is js
 
 ixsZipWith :: (i -> j -> k) -> IxS n i -> IxS n j -> IxS n k
-ixsZipWith f (IxS l1) (IxS l2) = IxS $ listsZipWith f l1 l2
+ixsZipWith _ ZIS ZIS = ZIS
+ixsZipWith f (i :.$ is) (j :.$ js) = f i j :.$ ixsZipWith f is js
 
 ixsPermutePrefix :: forall i is sh. Perm is -> IxS sh i -> IxS (PermutePrefix is sh) i
 ixsPermutePrefix = coerce (listsPermutePrefix @(Const i))