iota

1 files changed, 10 insertions, 0 deletions

diff --git a/src/Data/Array/Nested/Internal.hs b/src/Data/Array/Nested/Internal.hs
index e402f1e..308f8ce 100644
--- a/src/Data/Array/Nested/Internal.hs
+++ b/src/Data/Array/Nested/Internal.hs
@@ -55,6 +55,7 @@ import Foreign.Storable (Storable)
 import GHC.IsList (IsList)
 import qualified GHC.IsList as IsList
 import GHC.TypeLits
+import qualified GHC.TypeNats as TypeNats
 import Unsafe.Coerce
 
 import Data.Array.Mixed
@@ -951,6 +952,9 @@ mreshape sh' arr =
         (\sshIn -> X.reshapePartial (X.staticShapeFrom (mshape arr)) sshIn sh')
         arr
 
+miota :: (Enum a, PrimElt a) => SNat n -> Mixed '[Just n] a
+miota sn = fromPrimitive $ M_Primitive (SKnown sn :$% ZSX) (X.iota sn)
+
 masXArrayPrimP :: Mixed sh (Primitive a) -> (IShX sh, XArray sh a)
 masXArrayPrimP (M_Primitive sh arr) = (sh, arr)
 
@@ -1489,6 +1493,9 @@ rreshape sh' rarr@(Ranked arr)
   , Dict <- lemKnownReplicate (snatFromShR sh')
   = Ranked (mreshape (shCvtRX sh') arr)
 
+riota :: (Enum a, PrimElt a, Elt a) => Int -> Ranked 1 a
+riota n = TypeNats.withSomeSNat (fromIntegral n) $ mtoRanked . miota
+
 rasXArrayPrimP :: Ranked n (Primitive a) -> (IShR n, XArray (Replicate n Nothing) a)
 rasXArrayPrimP (Ranked arr) = first shCvtXR' (masXArrayPrimP arr)
 
@@ -1749,6 +1756,9 @@ srev1 arr = slift (sshape arr) (\_ -> X.rev1) arr
 sreshape :: Elt a => ShS sh' -> Shaped sh a -> Shaped sh' a
 sreshape sh' (Shaped arr) = Shaped (mreshape (shCvtSX sh') arr)
 
+siota :: (Enum a, PrimElt a) => SNat n -> Shaped '[n] a
+siota sn = Shaped (miota sn)
+
 sasXArrayPrimP :: Shaped sh (Primitive a) -> (ShS sh, XArray (MapJust sh) a)
 sasXArrayPrimP (Shaped arr) = first shCvtXS' (masXArrayPrimP arr)