Save some VS.concat by using toVectorListT

1 files changed, 9 insertions, 6 deletions

diff --git a/src/Data/Array/XArray.hs b/src/Data/Array/XArray.hs
index f779a53..38ccee6 100644
--- a/src/Data/Array/XArray.hs
+++ b/src/Data/Array/XArray.hs
@@ -17,7 +17,7 @@
 module Data.Array.XArray where
 
 import Control.DeepSeq (NFData)
-import Control.Monad (foldM)
+import Control.Monad (foldM_, foldM)
 import Control.Monad.ST
 import Data.Array.Internal qualified as OI
 import Data.Array.Internal.RankedG qualified as ORG
@@ -331,16 +331,19 @@ ravelOuterN k as@(a0 :| _) = runST $ do
       len = product sh0
       vecSize = k * len
   vec <- VSM.unsafeNew vecSize
-  let f !n a =
+  let f !n (ORS.A (ORG.A sht t)) =
         if | n >= k ->
              error $ "ravelOuterN: list too long " ++ show (n, k)
                -- if we do this check just once at the end, we may
                -- crash instead of producing an accurate error message
-           | S.shapeL a == sh0 -> do
-             VS.unsafeCopy (VSM.slice (n * len) len vec) (S.toVector a)
-             return $! n + 1
+           | sht == sh0 -> do
+               let g off el = do
+                     VS.unsafeCopy (VSM.slice off (VS.length el) vec) el
+                     return $! off + VS.length el
+               foldM_ g (n * len) (OI.toVectorListT sht t)
+               return $! n + 1
            | otherwise ->
-             error $ "ravelOuterN: unequal shapes " ++ show (S.shapeL a, sh0)
+             error $ "ravelOuterN: unequal shapes " ++ show (sht, sh0)
   nFinal <- foldM f 0 as
   if nFinal == k
   then S.fromVector (k : sh0) <$> VS.unsafeFreeze vec