aboutsummaryrefslogtreecommitdiff
path: root/src/Data/Array/Nested
diff options
context:
space:
mode:
authorTom Smeding <tom@tomsmeding.com>2024-05-17 13:31:39 +0200
committerTom Smeding <tom@tomsmeding.com>2024-05-17 13:31:39 +0200
commit43de4dd2e730273eb04bdf7d0ac62ac1e1422880 (patch)
tree7a91b756734222842999d767378c7544aac1e3bf /src/Data/Array/Nested
parent03879dfc925510c13c46664ca3438cc1d4872048 (diff)
Fix slice typeHEADmaster
Diffstat (limited to 'src/Data/Array/Nested')
-rw-r--r--src/Data/Array/Nested/Internal.hs17
1 files changed, 11 insertions, 6 deletions
diff --git a/src/Data/Array/Nested/Internal.hs b/src/Data/Array/Nested/Internal.hs
index 27e1a30..7bd6565 100644
--- a/src/Data/Array/Nested/Internal.hs
+++ b/src/Data/Array/Nested/Internal.hs
@@ -594,8 +594,11 @@ mconstant :: forall sh a. (KnownShapeX sh, Storable a, PrimElt a)
=> IShX sh -> a -> Mixed sh a
mconstant sh x = fromPrimitive (mconstantP sh x)
-mslice :: (KnownShapeX sh, Elt a) => [(Int, Int)] -> Mixed sh a -> Mixed sh a
-mslice ivs = mlift $ \_ -> X.slice ivs
+mslice :: (KnownShapeX sh, Elt a) => SNat i -> SNat n -> Mixed (Just (i + n + k) : sh) a -> Mixed (Just n : sh) a
+mslice i n = withKnownNat n $ mlift $ \_ -> X.slice i n
+
+msliceU :: (KnownShapeX sh, Elt a) => Int -> Int -> Mixed (Nothing : sh) a -> Mixed (Nothing : sh) a
+msliceU i n = mlift $ \_ -> X.sliceU i n
mrev1 :: (KnownShapeX (n : sh), Elt a) => Mixed (n : sh) a -> Mixed (n : sh) a
mrev1 = mlift $ \_ -> X.rev1
@@ -1158,8 +1161,10 @@ rconstant :: forall n a. (KnownNat n, Storable a, PrimElt a)
=> IShR n -> a -> Ranked n a
rconstant sh x = coerce fromPrimitive (rconstantP sh x)
-rslice :: (KnownNat n, Elt a) => [(Int, Int)] -> Ranked n a -> Ranked n a
-rslice ivs = rlift $ \_ -> X.slice ivs
+rslice :: forall n a. (KnownNat n, Elt a) => Int -> Int -> Ranked (n + 1) a -> Ranked (n + 1) a
+rslice i n
+ | Refl <- X.lemReplicateSucc @(Nothing @Nat) @n
+ = rlift $ \_ -> X.sliceU i n
rrev1 :: forall n a. (KnownNat n, Elt a) => Ranked (n + 1) a -> Ranked (n + 1) a
rrev1 = rlift $ \(Proxy @sh') ->
@@ -1436,8 +1441,8 @@ sconstant :: forall sh a. (KnownShape sh, Storable a, PrimElt a)
=> a -> Shaped sh a
sconstant x = coerce fromPrimitive (sconstantP @sh x)
-sslice :: (KnownShape sh, Elt a) => [(Int, Int)] -> Shaped sh a -> Shaped sh a
-sslice ivs = slift $ \_ -> X.slice ivs
+sslice :: (KnownShape sh, Elt a) => SNat i -> SNat n -> Shaped (i + n + k : sh) a -> Shaped (n : sh) a
+sslice i n = withKnownNat n $ slift $ \_ -> X.slice i n
srev1 :: (KnownNat n, KnownShape sh, Elt a) => Shaped (n : sh) a -> Shaped (n : sh) a
srev1 = slift $ \_ -> X.rev1