From a3149a6e039516e984f745cd6260ce95bf80dee4 Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Sat, 20 Apr 2024 22:52:23 +0200
Subject: slice

---
 src/Data/Array/Mixed.hs           | 3 +++
 src/Data/Array/Nested.hs          | 4 +++-
 src/Data/Array/Nested/Internal.hs | 9 +++++++++
 3 files changed, 15 insertions(+), 1 deletion(-)

(limited to 'src')

diff --git a/src/Data/Array/Mixed.hs b/src/Data/Array/Mixed.hs
index e7d2b69..164f832 100644
--- a/src/Data/Array/Mixed.hs
+++ b/src/Data/Array/Mixed.hs
@@ -363,3 +363,6 @@ fromList ssh l
 
 toList :: Storable a => XArray (n : sh) a -> [XArray sh a]
 toList (XArray arr) = coerce (ORB.toList (S.unravel arr))
+
+slice :: [(Int, Int)] -> XArray sh a -> XArray sh a
+slice ivs (XArray arr) = XArray (S.slice ivs arr)
diff --git a/src/Data/Array/Nested.hs b/src/Data/Array/Nested.hs
index ee007e9..ada8751 100644
--- a/src/Data/Array/Nested.hs
+++ b/src/Data/Array/Nested.hs
@@ -6,6 +6,7 @@ module Data.Array.Nested (
   rshape, rindex, rindexPartial, rgenerate, rsumOuter1,
   rtranspose, rappend, rscalar, rfromVector, runScalar,
   rconstant, rfromList, rfromList1, rtoList, rtoList1,
+  rslice,
   -- ** Lifting orthotope operations to 'Ranked' arrays
   rlift,
 
@@ -16,6 +17,7 @@ module Data.Array.Nested (
   sshape, sindex, sindexPartial, sgenerate, ssumOuter1,
   stranspose, sappend, sscalar, sfromVector, sunScalar,
   sconstant, sfromList, sfromList1, stoList, stoList1,
+  sslice,
   -- ** Lifting orthotope operations to 'Shaped' arrays
   slift,
 
@@ -24,7 +26,7 @@ module Data.Array.Nested (
   IxX(..),
   KnownShapeX(..), StaticShapeX(..),
   mgenerate, mtranspose, mappend, mfromVector, munScalar,
-  mconstant, mfromList1, mtoList1,
+  mconstant, mfromList1, mtoList1, mslice,
 
   -- * Array elements
   Elt(mshape, mindex, mindexPartial, mscalar, mfromList, mtoList, mlift, mlift2),
diff --git a/src/Data/Array/Nested/Internal.hs b/src/Data/Array/Nested/Internal.hs
index c8a0670..4e78647 100644
--- a/src/Data/Array/Nested/Internal.hs
+++ b/src/Data/Array/Nested/Internal.hs
@@ -517,6 +517,9 @@ mconstant :: forall sh a. (KnownShapeX sh, Storable a, Coercible (Mixed sh (Prim
           => IxX sh -> a -> Mixed sh a
 mconstant sh x = coerce (mconstantP sh x)
 
+mslice :: (KnownShapeX sh, Elt a) => [(Int, Int)] -> Mixed sh a -> Mixed sh a
+mslice ivs = mlift $ \_ -> X.slice ivs
+
 mliftPrim :: (KnownShapeX sh, Storable a)
           => (a -> a)
           -> Mixed sh (Primitive a) -> Mixed sh (Primitive a)
@@ -959,6 +962,9 @@ rconstant :: forall n a. (KnownINat n, Storable a, Coercible (Mixed (Replicate n
           => IxR n -> a -> Ranked n a
 rconstant sh x = coerce (rconstantP sh x)
 
+rslice :: (KnownINat n, Elt a) => [(Int, Int)] -> Ranked n a -> Ranked n a
+rslice ivs = rlift $ \_ -> X.slice ivs
+
 rfromList :: forall n a. (KnownINat n, Elt a) => NonEmpty (Ranked n a) -> Ranked (S n) a
 rfromList l
   | Dict <- lemKnownReplicate (Proxy @n)
@@ -1101,6 +1107,9 @@ sconstant :: forall sh a. (KnownShape sh, Storable a, Coercible (Mixed (MapJust
           => a -> Shaped sh a
 sconstant x = coerce (sconstantP @sh x)
 
+sslice :: (KnownShape sh, Elt a) => [(Int, Int)] -> Shaped sh a -> Shaped sh a
+sslice ivs = slift $ \_ -> X.slice ivs
+
 sfromList :: forall n sh a. (KnownNat n, KnownShape sh, Elt a)
           => NonEmpty (Shaped sh a) -> Shaped (n : sh) a
 sfromList l
-- 
cgit v1.2.3-70-g09d2