diff options
author | Tom Smeding <tom@tomsmeding.com> | 2024-05-17 13:31:39 +0200 |
---|---|---|
committer | Tom Smeding <tom@tomsmeding.com> | 2024-05-17 13:31:39 +0200 |
commit | 43de4dd2e730273eb04bdf7d0ac62ac1e1422880 (patch) | |
tree | 7a91b756734222842999d767378c7544aac1e3bf | |
parent | 03879dfc925510c13c46664ca3438cc1d4872048 (diff) |
-rw-r--r-- | src/Data/Array/Mixed.hs | 7 | ||||
-rw-r--r-- | src/Data/Array/Nested/Internal.hs | 17 |
2 files changed, 16 insertions, 8 deletions
diff --git a/src/Data/Array/Mixed.hs b/src/Data/Array/Mixed.hs index 5fbc46f..ce18431 100644 --- a/src/Data/Array/Mixed.hs +++ b/src/Data/Array/Mixed.hs @@ -558,8 +558,11 @@ empty sh = XArray (S.constant (shapeLshape sh) (error "Data.Array.Mixed.empty: shape was not empty")) -slice :: [(Int, Int)] -> XArray sh a -> XArray sh a -slice ivs (XArray arr) = XArray (S.slice ivs arr) +slice :: SNat i -> SNat n -> XArray (Just (i + n + k) : sh) a -> XArray (Just n : sh) a +slice i n (XArray arr) = XArray (S.slice [(fromSNat' i, fromSNat' n)] arr) + +sliceU :: Int -> Int -> XArray (Nothing : sh) a -> XArray (Nothing : sh) a +sliceU i n (XArray arr) = XArray (S.slice [(i, n)] arr) rev1 :: XArray (n : sh) a -> XArray (n : sh) a rev1 (XArray arr) = XArray (S.rev [0] arr) 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 |