From 43de4dd2e730273eb04bdf7d0ac62ac1e1422880 Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Fri, 17 May 2024 13:31:39 +0200
Subject: Fix slice type

---
 src/Data/Array/Nested/Internal.hs | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

(limited to 'src/Data/Array/Nested')

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
-- 
cgit v1.2.3-70-g09d2