[rsm]emptyArray

1 files changed, 9 insertions, 6 deletions

diff --git a/src/Data/Array/Nested/Internal/Mixed.hs b/src/Data/Array/Nested/Internal/Mixed.hs
index 0e4f5e6..8d239cf 100644
--- a/src/Data/Array/Nested/Internal/Mixed.hs
+++ b/src/Data/Array/Nested/Internal/Mixed.hs
@@ -366,7 +366,7 @@ class Elt a where
 -- 'Data.Array.Nested.Shaped.sgenerate'.
 class Elt a => KnownElt a where
   -- | Create an empty array. The given shape must have size zero; this may or may not be checked.
-  memptyArray :: IShX sh -> Mixed sh a
+  memptyArrayUnsafe :: IShX sh -> Mixed sh a
 
   -- | Create uninitialised vectors for this array type, given the shape of
   -- this vector and an example for the contents.
@@ -461,7 +461,7 @@ deriving via Primitive Float instance Elt Float
 deriving via Primitive () instance Elt ()
 
 instance Storable a => KnownElt (Primitive a) where
-  memptyArray sh = M_Primitive sh (X.empty sh)
+  memptyArrayUnsafe sh = M_Primitive sh (X.empty sh)
   mvecsUnsafeNew sh _ = MV_Primitive <$> VSM.unsafeNew (shxSize sh)
   mvecsNewEmpty _ = MV_Primitive <$> VSM.unsafeNew 0
 
@@ -517,7 +517,7 @@ instance (Elt a, Elt b) => Elt (a, b) where
   mvecsFreeze sh (MV_Tup2 a b) = M_Tup2 <$> mvecsFreeze sh a <*> mvecsFreeze sh b
 
 instance (KnownElt a, KnownElt b) => KnownElt (a, b) where
-  memptyArray sh = M_Tup2 (memptyArray sh) (memptyArray sh)
+  memptyArrayUnsafe sh = M_Tup2 (memptyArrayUnsafe sh) (memptyArrayUnsafe sh)
   mvecsUnsafeNew sh (x, y) = MV_Tup2 <$> mvecsUnsafeNew sh x <*> mvecsUnsafeNew sh y
   mvecsNewEmpty _ = MV_Tup2 <$> mvecsNewEmpty (Proxy @a) <*> mvecsNewEmpty (Proxy @b)
 
@@ -650,7 +650,7 @@ instance Elt a => Elt (Mixed sh' a) where
   mvecsFreeze sh (MV_Nest sh' vecs) = M_Nest sh <$> mvecsFreeze (shxAppend sh sh') vecs
 
 instance (KnownShX sh', KnownElt a) => KnownElt (Mixed sh' a) where
-  memptyArray sh = M_Nest sh (memptyArray (shxAppend sh (shxCompleteZeros (knownShX @sh'))))
+  memptyArrayUnsafe sh = M_Nest sh (memptyArrayUnsafe (shxAppend sh (shxCompleteZeros (knownShX @sh'))))
 
   mvecsUnsafeNew sh example
     | shxSize sh' == 0 = mvecsNewEmpty (Proxy @(Mixed sh' a))
@@ -661,6 +661,9 @@ instance (KnownShX sh', KnownElt a) => KnownElt (Mixed sh' a) where
   mvecsNewEmpty _ = MV_Nest (shxCompleteZeros (knownShX @sh')) <$> mvecsNewEmpty (Proxy @a)
 
 
+memptyArray :: KnownElt a => IShX sh -> Mixed (Just 0 : sh) a
+memptyArray sh = memptyArrayUnsafe (SKnown SNat :$% sh)
+
 mrank :: Elt a => Mixed sh a -> SNat (Rank sh)
 mrank = shxRank . mshape
 
@@ -687,12 +690,12 @@ msize = shxSize . mshape
 -- easily, hence the runtime check.
 mgenerate :: forall sh a. KnownElt a => IShX sh -> (IIxX sh -> a) -> Mixed sh a
 mgenerate sh f = case shxEnum sh of
-  [] -> memptyArray sh
+  [] -> memptyArrayUnsafe sh
   firstidx : restidxs ->
     let firstelem = f (ixxZero' sh)
         shapetree = mshapeTree firstelem
     in if mshapeTreeEmpty (Proxy @a) shapetree
-         then memptyArray sh
+         then memptyArrayUnsafe sh
          else runST $ do
                 vecs <- mvecsUnsafeNew sh firstelem
                 mvecsWrite sh firstidx firstelem vecs