Add mvecsReplicatemvecsReplicate

3 files changed, 21 insertions, 0 deletions

diff --git a/src/Data/Array/Nested/Mixed.hs b/src/Data/Array/Nested/Mixed.hs
index 182943d..654ce3c 100644
--- a/src/Data/Array/Nested/Mixed.hs
+++ b/src/Data/Array/Nested/Mixed.hs
@@ -393,6 +393,10 @@ class Elt a => KnownElt a where
   -- this vector and an example for the contents.
   mvecsUnsafeNew :: IShX sh -> a -> ST s (MixedVecs s sh a)
 
+  -- | Create initialised vectors for this array type, given the shape of
+  -- this vector and the chosen element.
+  mvecsReplicate :: IShX sh -> a -> ST s (MixedVecs s sh a)
+
   mvecsNewEmpty :: Proxy a -> ST s (MixedVecs s sh a)
 
 
@@ -492,6 +496,7 @@ deriving via Primitive () instance Elt ()
 instance Storable a => KnownElt (Primitive a) where
   memptyArrayUnsafe sh = M_Primitive sh (X.empty sh)
   mvecsUnsafeNew sh _ = MV_Primitive <$> VSM.unsafeNew (shxSize sh)
+  mvecsReplicate sh (Primitive a) = MV_Primitive <$> VSM.replicate (shxSize sh) a
   mvecsNewEmpty _ = MV_Primitive <$> VSM.unsafeNew 0
 
 -- [PRIMITIVE ELEMENT TYPES LIST]
@@ -553,6 +558,7 @@ instance (Elt a, Elt b) => Elt (a, b) where
 instance (KnownElt a, KnownElt b) => KnownElt (a, b) where
   memptyArrayUnsafe sh = M_Tup2 (memptyArrayUnsafe sh) (memptyArrayUnsafe sh)
   mvecsUnsafeNew sh (x, y) = MV_Tup2 <$> mvecsUnsafeNew sh x <*> mvecsUnsafeNew sh y
+  mvecsReplicate sh (x, y) = MV_Tup2 <$> mvecsReplicate sh x <*> mvecsReplicate sh y
   mvecsNewEmpty _ = MV_Tup2 <$> mvecsNewEmpty (Proxy @a) <*> mvecsNewEmpty (Proxy @b)
 
 -- Arrays of arrays are just arrays, but with more dimensions.
@@ -697,6 +703,13 @@ instance (KnownShX sh', KnownElt a) => KnownElt (Mixed sh' a) where
     where
       sh' = mshape example
 
+  mvecsReplicate sh example = do
+    vecs <- mvecsUnsafeNew sh example
+    forM_ (shxEnum sh) $ \idx -> mvecsWrite sh idx example vecs
+      -- this is a slow case, but the alternative, mvecsUnsafeNew with manual
+      -- writing in a loop, leads to every case being as slow
+    return vecs
+
   mvecsNewEmpty _ = MV_Nest (shxCompleteZeros (knownShX @sh')) <$> mvecsNewEmpty (Proxy @a)
 
 
diff --git a/src/Data/Array/Nested/Ranked/Base.hs b/src/Data/Array/Nested/Ranked/Base.hs
index 11a8ffb..ef3af31 100644
--- a/src/Data/Array/Nested/Ranked/Base.hs
+++ b/src/Data/Array/Nested/Ranked/Base.hs
@@ -188,6 +188,10 @@ instance (KnownNat n, KnownElt a) => KnownElt (Ranked n a) where
     | Dict <- lemKnownReplicate (SNat @n)
     = MV_Ranked <$> mvecsUnsafeNew idx arr
 
+  mvecsReplicate idx (Ranked arr)
+    | Dict <- lemKnownReplicate (SNat @n)
+    = MV_Ranked <$> mvecsReplicate idx arr
+
   mvecsNewEmpty _
     | Dict <- lemKnownReplicate (SNat @n)
     = MV_Ranked <$> mvecsNewEmpty (Proxy @(Mixed (Replicate n Nothing) a))
diff --git a/src/Data/Array/Nested/Shaped/Base.hs b/src/Data/Array/Nested/Shaped/Base.hs
index 98f1241..ded1175 100644
--- a/src/Data/Array/Nested/Shaped/Base.hs
+++ b/src/Data/Array/Nested/Shaped/Base.hs
@@ -181,6 +181,10 @@ instance (KnownShS sh, KnownElt a) => KnownElt (Shaped sh a) where
     | Dict <- lemKnownMapJust (Proxy @sh)
     = MV_Shaped <$> mvecsUnsafeNew idx arr
 
+  mvecsReplicate idx (Shaped arr)
+    | Dict <- lemKnownMapJust (Proxy @sh)
+    = MV_Shaped <$> mvecsReplicate idx arr
+
   mvecsNewEmpty _
     | Dict <- lemKnownMapJust (Proxy @sh)
     = MV_Shaped <$> mvecsNewEmpty (Proxy @(Mixed (MapJust sh) a))