aboutsummaryrefslogtreecommitdiff
path: root/src/Data/Array/Mixed/Shape.hs
diff options
context:
space:
mode:
Diffstat (limited to 'src/Data/Array/Mixed/Shape.hs')
-rw-r--r--src/Data/Array/Mixed/Shape.hs21
1 files changed, 12 insertions, 9 deletions
diff --git a/src/Data/Array/Mixed/Shape.hs b/src/Data/Array/Mixed/Shape.hs
index 77d6a8b..809ed9e 100644
--- a/src/Data/Array/Mixed/Shape.hs
+++ b/src/Data/Array/Mixed/Shape.hs
@@ -27,6 +27,7 @@ import Data.Bifunctor (first)
import Data.Coerce
import Data.Foldable qualified as Foldable
import Data.Functor.Const
+import Data.Functor.Product
import Data.Kind (Type, Constraint)
import Data.Monoid (Sum(..))
import Data.Proxy
@@ -152,16 +153,16 @@ listxLast :: forall f n sh. ListX (n : sh) f -> f (Last (n : sh))
listxLast (_ ::% sh@(_ ::% _)) = listxLast sh
listxLast (x ::% ZX) = x
-listxZip :: ListX sh (Const i) -> ListX sh (Const j) -> ListX sh (Const (i, j))
+listxZip :: ListX sh f -> ListX sh g -> ListX sh (Product f g)
listxZip ZX ZX = ZX
-listxZip (Const i ::% irest) (Const j ::% jrest) =
- Const (i, j) ::% listxZip irest jrest
+listxZip (i ::% irest) (j ::% jrest) =
+ Pair i j ::% listxZip irest jrest
-listxZipWith :: (i -> j -> k) -> ListX sh (Const i) -> ListX sh (Const j)
- -> ListX sh (Const k)
+listxZipWith :: (forall a. f a -> g a -> h a) -> ListX sh f -> ListX sh g
+ -> ListX sh h
listxZipWith _ ZX ZX = ZX
-listxZipWith f (Const i ::% irest) (Const j ::% jrest) =
- Const (f i j) ::% listxZipWith f irest jrest
+listxZipWith f (i ::% is) (j ::% js) =
+ f i j ::% listxZipWith f is js
-- * Mixed indices
@@ -234,10 +235,12 @@ 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 l1) (IxX l2) = IxX $ listxZip l1 l2
+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 f (IxX l1) (IxX l2) = IxX $ listxZipWith f l1 l2
+ixxZipWith _ ZIX ZIX = ZIX
+ixxZipWith f (i :.% is) (j :.% js) = f i j :.% ixxZipWith f is js
ixxFromLinear :: IShX sh -> Int -> IIxX sh
ixxFromLinear = \sh i -> case go sh i of