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.hs31
1 files changed, 31 insertions, 0 deletions
diff --git a/src/Data/Array/Mixed/Shape.hs b/src/Data/Array/Mixed/Shape.hs
index 99a137d..4dd0aa6 100644
--- a/src/Data/Array/Mixed/Shape.hs
+++ b/src/Data/Array/Mixed/Shape.hs
@@ -130,6 +130,9 @@ listxToList :: ListX sh' (Const i) -> [i]
listxToList ZX = []
listxToList (Const i ::% is) = i : listxToList is
+listxHead :: ListX (mn ': sh) f -> f mn
+listxHead (i ::% _) = i
+
listxTail :: ListX (n : sh) i -> ListX sh i
listxTail (_ ::% sh) = sh
@@ -149,6 +152,19 @@ 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 ZX ZX = ZX
+listxZip (Const i ::% irest) (Const j ::% jrest) =
+ Const (i, j) ::% listxZip irest jrest
+--listxZip _ _ = error "listxZip: impossible pattern needlessly required"
+
+listxZipWith :: (i -> j -> k) -> ListX sh (Const i) -> ListX sh (Const j)
+ -> ListX sh (Const k)
+listxZipWith _ ZX ZX = ZX
+listxZipWith f (Const i ::% irest) (Const j ::% jrest) =
+ Const (f i j) ::% listxZipWith f irest jrest
+--listxZipWith _ _ _ = error "listxZipWith: impossible pattern needlessly required"
+
-- * Mixed indices
@@ -201,6 +217,9 @@ ixxZero' (_ :$% sh) = 0 :.% ixxZero' sh
ixxFromList :: forall sh i. StaticShX sh -> [i] -> IxX sh i
ixxFromList = coerce (listxFromList @_ @i)
+ixxHead :: IxX (n : sh) i -> i
+ixxHead (IxX list) = getConst (listxHead list)
+
ixxTail :: IxX (n : sh) i -> IxX sh i
ixxTail (IxX list) = IxX (listxTail list)
@@ -216,6 +235,12 @@ 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 (IxX l1) (IxX l2) = IxX $ listxZip l1 l2
+
+ixxZipWith :: (i -> j -> k) -> IxX n i -> IxX n j -> IxX n k
+ixxZipWith f (IxX l1) (IxX l2) = IxX $ listxZipWith f l1 l2
+
ixxFromLinear :: IShX sh -> Int -> IIxX sh
ixxFromLinear = \sh i -> case go sh i of
(idx, 0) -> idx
@@ -372,6 +397,9 @@ shxToList (smn :$% sh) = fromSMayNat' smn : shxToList sh
shxAppend :: forall sh sh' i. ShX sh i -> ShX sh' i -> ShX (sh ++ sh') i
shxAppend = coerce (listxAppend @_ @(SMayNat i SNat))
+shxHead :: ShX (n : sh) i -> SMayNat i SNat n
+shxHead (ShX list) = listxHead list
+
shxTail :: ShX (n : sh) i -> ShX sh i
shxTail (ShX list) = ShX (listxTail list)
@@ -474,6 +502,9 @@ ssxAppend :: StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend ZKX sh' = sh'
ssxAppend (n :!% sh) sh' = n :!% ssxAppend sh sh'
+ssxHead :: StaticShX (n : sh) -> SMayNat () SNat n
+ssxHead (StaticShX list) = listxHead list
+
ssxTail :: StaticShX (n : sh) -> StaticShX sh
ssxTail (_ :!% ssh) = ssh