Fix slice typeHEADmaster

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