3 files changed, 21 insertions, 0 deletions
diff --git a/src/Data/Array/Nested/Mixed.hs b/src/Data/Array/Nested/Mixed.hs
index cf62417..b2a400a 100644
--- a/src/Data/Array/Nested/Mixed.hs
+++ b/src/Data/Array/Nested/Mixed.hs
@@ -397,6 +397,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)
 
 
@@ -497,6 +501,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]
@@ -559,6 +564,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.
@@ -704,6 +710,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 defd72f..96a9360 100644
--- a/src/Data/Array/Nested/Ranked/Base.hs
+++ b/src/Data/Array/Nested/Ranked/Base.hs
@@ -196,6 +196,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 7ebc685..686f5fc 100644
--- a/src/Data/Array/Nested/Shaped/Base.hs
+++ b/src/Data/Array/Nested/Shaped/Base.hs
@@ -189,6 +189,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))