diff options
| -rw-r--r-- | src/Data/Array/Mixed/Shape.hs | 21 | ||||
| -rw-r--r-- | src/Data/Array/Nested/Internal/Shape.hs | 21 |
2 files changed, 24 insertions, 18 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 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)) |