Save some VS.concat by using toVectorListT

1 files changed, 10 insertions, 4 deletions

diff --git a/src/Data/Array/Nested/Mixed.hs b/src/Data/Array/Nested/Mixed.hs
index f0dadb4..39f00fa 100644
--- a/src/Data/Array/Nested/Mixed.hs
+++ b/src/Data/Array/Nested/Mixed.hs
@@ -23,8 +23,11 @@ module Data.Array.Nested.Mixed where
 import Prelude hiding (mconcat)
 
 import Control.DeepSeq (NFData(..))
-import Control.Monad (forM_, when)
+import Control.Monad (foldM_, forM_, when)
 import Control.Monad.ST
+import Data.Array.Internal qualified as OI
+import Data.Array.Internal.RankedG qualified as ORG
+import Data.Array.Internal.RankedS qualified as ORS
 import Data.Array.RankedS qualified as S
 import Data.Bifunctor (bimap)
 import Data.Coerce
@@ -485,14 +488,17 @@ instance Storable a => Elt (Primitive a) where
   marrayStrides (M_Primitive _ arr) = BOne (X.arrayStrides arr)
   mvecsWriteLinear i (Primitive x) (MV_Primitive v) = VSM.write v i x
 
-  -- TODO: this use of toVector is suboptimal
+  -- TODO: this use of toVectorListT is suboptimal
   mvecsWritePartialLinear
     :: forall sh' sh s.
        Proxy sh -> Int -> Mixed sh' (Primitive a) -> MixedVecs s (sh ++ sh') (Primitive a) -> ST s ()
-  mvecsWritePartialLinear _ i (M_Primitive sh' arr) (MV_Primitive v) = do
+  mvecsWritePartialLinear _ i (M_Primitive sh' arr@(XArray (ORS.A (ORG.A sht t)))) (MV_Primitive v) = do
     let arrsh = X.shape (ssxFromShX sh') arr
         offset = i * shxSize arrsh
-    VS.copy (VSM.slice offset (shxSize arrsh) v) (X.toVector arr)
+        f off el = do
+          VS.copy (VSM.slice off (VS.length el) v) el
+          return $! off + VS.length el
+    foldM_ f offset (OI.toVectorListT sht t)
 
   mvecsFreeze sh (MV_Primitive v) = M_Primitive sh . X.fromVector sh <$> VS.freeze v
   mvecsUnsafeFreeze sh (MV_Primitive v) = M_Primitive sh . X.fromVector sh <$> VS.unsafeFreeze v