diff options
Diffstat (limited to 'src/Data/Array/Mixed/Shape.hs')
| -rw-r--r-- | src/Data/Array/Mixed/Shape.hs | 20 |
1 files changed, 20 insertions, 0 deletions
diff --git a/src/Data/Array/Mixed/Shape.hs b/src/Data/Array/Mixed/Shape.hs index b49e005..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,6 +153,17 @@ listxLast :: forall f n sh. ListX (n : sh) f -> f (Last (n : sh)) listxLast (_ ::% sh@(_ ::% _)) = listxLast sh listxLast (x ::% ZX) = x +listxZip :: ListX sh f -> ListX sh g -> ListX sh (Product f g) +listxZip ZX ZX = ZX +listxZip (i ::% irest) (j ::% jrest) = + Pair i j ::% listxZip irest jrest + +listxZipWith :: (forall a. f a -> g a -> h a) -> ListX sh f -> ListX sh g + -> ListX sh h +listxZipWith _ ZX ZX = ZX +listxZipWith f (i ::% is) (j ::% js) = + f i j ::% listxZipWith f is js + -- * Mixed indices @@ -222,6 +234,14 @@ 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 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 _ 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 (idx, 0) -> idx |