{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

module HaskellWorks.Data.Vector.Storable
  ( padded
  , foldMap
  , mapAccumL
  , mmap
  , constructSI
  , construct2N
  , construct64UnzipN
  , unzipFromListN2
  ) where

import Control.Monad.ST                   (ST, runST)
import Data.Vector.Storable               (Storable)
import Data.Word
import Foreign.ForeignPtr
import HaskellWorks.Data.Vector.AsVector8
import Prelude                            hiding (abs, foldMap)

import qualified Data.ByteString              as BS
import qualified Data.Vector.Generic          as DVG
import qualified Data.Vector.Storable         as DVS
import qualified Data.Vector.Storable.Mutable as DVSM
import qualified System.IO.MMap               as IO

{- HLINT ignore "Redundant do" -}

padded :: Int -> DVS.Vector Word8 -> DVS.Vector Word8
padded :: Int -> Vector Word8 -> Vector Word8
padded Int
n Vector Word8
v = Vector Word8
v Vector Word8 -> Vector Word8 -> Vector Word8
forall a. Semigroup a => a -> a -> a
<> Int -> Word8 -> Vector Word8
forall a. Storable a => Int -> a -> Vector a
DVS.replicate ((Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector Word8 -> Int
forall a. Storable a => Vector a -> Int
DVS.length Vector Word8
v) Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` Int
0) Word8
0
{-# INLINE padded #-}

foldMap :: (DVS.Storable a, Monoid m) => (a -> m) -> DVS.Vector a -> m
foldMap :: (a -> m) -> Vector a -> m
foldMap a -> m
f = (m -> a -> m) -> m -> Vector a -> m
forall b a. Storable b => (a -> b -> a) -> a -> Vector b -> a
DVS.foldl' (\m
a a
b -> m
a m -> m -> m
forall a. Semigroup a => a -> a -> a
<> a -> m
f a
b) m
forall a. Monoid a => a
mempty
{-# INLINE foldMap #-}

mapAccumL :: forall a b c. (Storable b, Storable c)
  => (a -> b -> (a, c))
  -> a
  -> DVS.Vector b
  -> (a, DVS.Vector c)
mapAccumL :: (a -> b -> (a, c)) -> a -> Vector b -> (a, Vector c)
mapAccumL a -> b -> (a, c)
f a
a Vector b
vb = (forall s. ST s (a, MVector s c)) -> (a, Vector c)
forall (f :: * -> *) a.
(Traversable f, Storable a) =>
(forall s. ST s (f (MVector s a))) -> f (Vector a)
DVS.createT ((forall s. ST s (a, MVector s c)) -> (a, Vector c))
-> (forall s. ST s (a, MVector s c)) -> (a, Vector c)
forall a b. (a -> b) -> a -> b
$ do
  MVector s c
vc <- Int -> ST s (MVector (PrimState (ST s)) c)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
DVSM.unsafeNew (Vector b -> Int
forall a. Storable a => Vector a -> Int
DVS.length Vector b
vb)
  a
a' <- Int -> a -> MVector s c -> ST s a
forall s. Int -> a -> MVector s c -> ST s a
go Int
0 a
a MVector s c
vc
  (a, MVector s c) -> ST s (a, MVector s c)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a', MVector s c
vc)
  where go :: Int -> a -> DVS.MVector s c -> ST s a
        go :: Int -> a -> MVector s c -> ST s a
go Int
i a
a0 MVector s c
vc = if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Vector b -> Int
forall a. Storable a => Vector a -> Int
DVS.length Vector b
vb
          then do
            let (a
a1, c
c1) = a -> b -> (a, c)
f a
a0 (Vector b -> Int -> b
forall a. Storable a => Vector a -> Int -> a
DVS.unsafeIndex Vector b
vb Int
i)
            MVector (PrimState (ST s)) c -> Int -> c -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
DVSM.unsafeWrite MVector s c
MVector (PrimState (ST s)) c
vc Int
i c
c1
            Int -> a -> MVector s c -> ST s a
forall s. Int -> a -> MVector s c -> ST s a
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
a1 MVector s c
vc
          else a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a0
{-# INLINE mapAccumL #-}

-- | MMap the file as a storable vector.  If the size of the file is not a multiple of the element size
-- in bytes, then the last few bytes of the file will not be included in the vector.
mmap :: Storable a => FilePath -> IO (DVS.Vector a)
mmap :: FilePath -> IO (Vector a)
mmap FilePath
filepath = do
  (ForeignPtr Word8
fptr :: ForeignPtr Word8, Int
offset, Int
size) <- FilePath
-> Mode -> Maybe (Int64, Int) -> IO (ForeignPtr Word8, Int, Int)
forall a.
FilePath
-> Mode -> Maybe (Int64, Int) -> IO (ForeignPtr a, Int, Int)
IO.mmapFileForeignPtr FilePath
filepath Mode
IO.ReadOnly Maybe (Int64, Int)
forall a. Maybe a
Nothing
  let !v :: Vector Word8
v = ForeignPtr Word8 -> Int -> Int -> Vector Word8
forall a. Storable a => ForeignPtr a -> Int -> Int -> Vector a
DVS.unsafeFromForeignPtr ForeignPtr Word8
fptr Int
offset Int
size
  Vector a -> IO (Vector a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector Word8 -> Vector a
forall a b. (Storable a, Storable b) => Vector a -> Vector b
DVS.unsafeCast Vector Word8
v)

-- | Construct a vector statefully with index
constructSI :: forall a s. Storable a => Int -> (Int -> s -> (s, a)) -> s -> (s, DVS.Vector a)
constructSI :: Int -> (Int -> s -> (s, a)) -> s -> (s, Vector a)
constructSI Int
n Int -> s -> (s, a)
f s
state = (forall s. ST s (s, MVector s a)) -> (s, Vector a)
forall (f :: * -> *) a.
(Traversable f, Storable a) =>
(forall s. ST s (f (MVector s a))) -> f (Vector a)
DVS.createT ((forall s. ST s (s, MVector s a)) -> (s, Vector a))
-> (forall s. ST s (s, MVector s a)) -> (s, Vector a)
forall a b. (a -> b) -> a -> b
$ do
  MVector s a
mv <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
DVSM.unsafeNew Int
n
  s
state' <- Int -> s -> MVector s a -> ST s s
forall t. Int -> s -> MVector t a -> ST t s
go Int
0 s
state MVector s a
mv
  (s, MVector s a) -> ST s (s, MVector s a)
forall (m :: * -> *) a. Monad m => a -> m a
return (s
state', MVector s a
mv)
  where go :: Int -> s -> DVSM.MVector t a -> ST t s
        go :: Int -> s -> MVector t a -> ST t s
go Int
i s
s MVector t a
mv = if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< MVector t a -> Int
forall a s. Storable a => MVector s a -> Int
DVSM.length MVector t a
mv
          then do
            let (s
s', a
a) = Int -> s -> (s, a)
f Int
i s
s
            MVector (PrimState (ST t)) a -> Int -> a -> ST t ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
DVSM.unsafeWrite MVector t a
MVector (PrimState (ST t)) a
mv Int
i a
a
            Int -> s -> MVector t a -> ST t s
forall t. Int -> s -> MVector t a -> ST t s
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) s
s' MVector t a
mv
          else s -> ST t s
forall (m :: * -> *) a. Monad m => a -> m a
return s
s

construct2N :: (Storable b, Storable c)
  => Int
  -> (forall s. a -> DVSM.MVector s b -> ST s Int)
  -> Int
  -> (forall s. a -> DVSM.MVector s c -> ST s Int)
  -> [a]
  -> (DVS.Vector b, DVS.Vector c)
construct2N :: Int
-> (forall s. a -> MVector s b -> ST s Int)
-> Int
-> (forall s. a -> MVector s c -> ST s Int)
-> [a]
-> (Vector b, Vector c)
construct2N Int
nb forall s. a -> MVector s b -> ST s Int
fb Int
nc forall s. a -> MVector s c -> ST s Int
fc [a]
as = (forall s. ST s (Vector b, Vector c)) -> (Vector b, Vector c)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector b, Vector c)) -> (Vector b, Vector c))
-> (forall s. ST s (Vector b, Vector c)) -> (Vector b, Vector c)
forall a b. (a -> b) -> a -> b
$ do
  MVector s b
mbs <- Int -> ST s (MVector (PrimState (ST s)) b)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
DVSM.unsafeNew Int
nb
  MVector s c
mcs <- Int -> ST s (MVector (PrimState (ST s)) c)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
DVSM.unsafeNew Int
nc
  (MVector s b
mbs2, MVector s c
mcs2) <- (forall s. a -> MVector s b -> ST s Int)
-> Int
-> MVector s b
-> (forall s. a -> MVector s c -> ST s Int)
-> Int
-> MVector s c
-> [a]
-> ST s (MVector s b, MVector s c)
forall b c a s.
(Storable b, Storable c) =>
(forall t. a -> MVector t b -> ST t Int)
-> Int
-> MVector s b
-> (forall t. a -> MVector t c -> ST t Int)
-> Int
-> MVector s c
-> [a]
-> ST s (MVector s b, MVector s c)
go forall s. a -> MVector s b -> ST s Int
fb Int
0 MVector s b
mbs forall s. a -> MVector s c -> ST s Int
fc Int
0 MVector s c
mcs [a]
as
  Vector b
bs <- Mutable Vector (PrimState (ST s)) b -> ST s (Vector b)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
DVG.unsafeFreeze MVector s b
Mutable Vector (PrimState (ST s)) b
mbs2
  Vector c
cs <- Mutable Vector (PrimState (ST s)) c -> ST s (Vector c)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
DVG.unsafeFreeze MVector s c
Mutable Vector (PrimState (ST s)) c
mcs2
  (Vector b, Vector c) -> ST s (Vector b, Vector c)
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector b
bs, Vector c
cs)
  where go :: (Storable b, Storable c)
          => (forall t. a -> DVSM.MVector t b -> ST t Int)
          -> Int
          -> DVSM.MVector s b
          -> (forall t. a -> DVSM.MVector t c -> ST t Int)
          -> Int
          -> DVSM.MVector s c
          -> [a]
          -> ST s (DVSM.MVector s b, DVSM.MVector s c)
        go :: (forall t. a -> MVector t b -> ST t Int)
-> Int
-> MVector s b
-> (forall t. a -> MVector t c -> ST t Int)
-> Int
-> MVector s c
-> [a]
-> ST s (MVector s b, MVector s c)
go   forall t. a -> MVector t b -> ST t Int
_ Int
bn MVector s b
mbs   forall t. a -> MVector t c -> ST t Int
_ Int
cn MVector s c
mcs []     = (MVector s b, MVector s c) -> ST s (MVector s b, MVector s c)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> MVector s b -> MVector s b
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.take Int
bn MVector s b
mbs, Int -> MVector s c -> MVector s c
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.take Int
cn MVector s c
mcs)
        go forall t. a -> MVector t b -> ST t Int
fb' Int
bn MVector s b
mbs forall t. a -> MVector t c -> ST t Int
fc' Int
cn MVector s c
mcs (a
d:[a]
ds) = do
          Int
bi <- a -> MVector s b -> ST s Int
forall t. a -> MVector t b -> ST t Int
fb' a
d (Int -> MVector s b -> MVector s b
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.drop Int
bn MVector s b
mbs)
          Int
ci <- a -> MVector s c -> ST s Int
forall t. a -> MVector t c -> ST t Int
fc' a
d (Int -> MVector s c -> MVector s c
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.drop Int
cn MVector s c
mcs)
          (forall t. a -> MVector t b -> ST t Int)
-> Int
-> MVector s b
-> (forall t. a -> MVector t c -> ST t Int)
-> Int
-> MVector s c
-> [a]
-> ST s (MVector s b, MVector s c)
forall b c a s.
(Storable b, Storable c) =>
(forall t. a -> MVector t b -> ST t Int)
-> Int
-> MVector s b
-> (forall t. a -> MVector t c -> ST t Int)
-> Int
-> MVector s c
-> [a]
-> ST s (MVector s b, MVector s c)
go forall t. a -> MVector t b -> ST t Int
fb' (Int
bn Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bi) MVector s b
mbs forall t. a -> MVector t c -> ST t Int
fc' (Int
cn Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
ci) MVector s c
mcs [a]
ds

construct64UnzipN :: Int -> [(BS.ByteString, BS.ByteString)] -> (DVS.Vector Word64, DVS.Vector Word64)
construct64UnzipN :: Int -> [(ByteString, ByteString)] -> (Vector Word64, Vector Word64)
construct64UnzipN Int
nBytes [(ByteString, ByteString)]
xs = (Vector Word8 -> Vector Word64
forall a b. (Storable a, Storable b) => Vector a -> Vector b
DVS.unsafeCast Vector Word8
ibv, Vector Word8 -> Vector Word64
forall a b. (Storable a, Storable b) => Vector a -> Vector b
DVS.unsafeCast Vector Word8
bpv)
  where [Vector Word8
ibv, Vector Word8
bpv] = (forall s. ST s [MVector s Word8]) -> [Vector Word8]
forall (f :: * -> *) a.
(Traversable f, Storable a) =>
(forall s. ST s (f (MVector s a))) -> f (Vector a)
DVS.createT ((forall s. ST s [MVector s Word8]) -> [Vector Word8])
-> (forall s. ST s [MVector s Word8]) -> [Vector Word8]
forall a b. (a -> b) -> a -> b
$ do
          let nW8s :: Int
nW8s     = (Int
nBytes Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
8
          MVector s Word8
ibmv <- Int -> ST s (MVector (PrimState (ST s)) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
DVSM.new Int
nW8s
          MVector s Word8
bpmv <- Int -> ST s (MVector (PrimState (ST s)) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
DVSM.new Int
nW8s
          (Int
ibmvRemaining, Int
bpmvRemaining) <- MVector s Word8
-> MVector s Word8 -> [(ByteString, ByteString)] -> ST s (Int, Int)
forall s.
MVector s Word8
-> MVector s Word8 -> [(ByteString, ByteString)] -> ST s (Int, Int)
go MVector s Word8
ibmv MVector s Word8
bpmv [(ByteString, ByteString)]
xs
          let ibl :: Int
ibl = ((MVector s Word8 -> Int
forall a s. Storable a => MVector s a -> Int
DVSM.length MVector s Word8
ibmv Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ibmvRemaining Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
8) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
8
          let bpl :: Int
bpl = ((MVector s Word8 -> Int
forall a s. Storable a => MVector s a -> Int
DVSM.length MVector s Word8
bpmv Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bpmvRemaining Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
8) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
8
          [MVector s Word8] -> ST s [MVector s Word8]
forall (m :: * -> *) a. Monad m => a -> m a
return [Int -> MVector s Word8 -> MVector s Word8
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.take Int
ibl MVector s Word8
ibmv, Int -> MVector s Word8 -> MVector s Word8
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.take Int
bpl MVector s Word8
bpmv]
        go :: DVSM.MVector s Word8 -> DVSM.MVector s Word8 -> [(BS.ByteString, BS.ByteString)] -> ST s (Int, Int)
        go :: MVector s Word8
-> MVector s Word8 -> [(ByteString, ByteString)] -> ST s (Int, Int)
go MVector s Word8
ibmv MVector s Word8
bpmv ((ByteString
ib, ByteString
bp):[(ByteString, ByteString)]
ys) = do
          MVector (PrimState (ST s)) Word8 -> Vector Word8 -> ST s ()
forall a (m :: * -> *).
(Storable a, PrimMonad m) =>
MVector (PrimState m) a -> Vector a -> m ()
DVS.copy (Int -> MVector s Word8 -> MVector s Word8
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.take (ByteString -> Int
BS.length ByteString
ib) MVector s Word8
ibmv) (ByteString -> Vector Word8
forall a. AsVector8 a => a -> Vector Word8
asVector8 ByteString
ib)
          MVector (PrimState (ST s)) Word8 -> Vector Word8 -> ST s ()
forall a (m :: * -> *).
(Storable a, PrimMonad m) =>
MVector (PrimState m) a -> Vector a -> m ()
DVS.copy (Int -> MVector s Word8 -> MVector s Word8
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.take (ByteString -> Int
BS.length ByteString
bp) MVector s Word8
bpmv) (ByteString -> Vector Word8
forall a. AsVector8 a => a -> Vector Word8
asVector8 ByteString
bp)
          MVector s Word8
-> MVector s Word8 -> [(ByteString, ByteString)] -> ST s (Int, Int)
forall s.
MVector s Word8
-> MVector s Word8 -> [(ByteString, ByteString)] -> ST s (Int, Int)
go (Int -> MVector s Word8 -> MVector s Word8
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.drop (ByteString -> Int
BS.length ByteString
ib) MVector s Word8
ibmv) (Int -> MVector s Word8 -> MVector s Word8
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.drop (ByteString -> Int
BS.length ByteString
bp) MVector s Word8
bpmv) [(ByteString, ByteString)]
ys
        go MVector s Word8
ibmv MVector s Word8
bpmv [] = do
          let ibl :: Int
ibl = MVector s Word8 -> Int
forall a s. Storable a => MVector s a -> Int
DVSM.length MVector s Word8
ibmv Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
8
          let bpl :: Int
bpl = MVector s Word8 -> Int
forall a s. Storable a => MVector s a -> Int
DVSM.length MVector s Word8
bpmv Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
8
          MVector (PrimState (ST s)) Word8 -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> a -> m ()
DVSM.set (Int -> MVector s Word8 -> MVector s Word8
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.take Int
ibl MVector s Word8
ibmv) Word8
0
          MVector (PrimState (ST s)) Word8 -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> a -> m ()
DVSM.set (Int -> MVector s Word8 -> MVector s Word8
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.take Int
bpl MVector s Word8
bpmv) Word8
0
          (Int, Int) -> ST s (Int, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (MVector s Word8 -> Int
forall a s. Storable a => MVector s a -> Int
DVSM.length (Int -> MVector s Word8 -> MVector s Word8
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.drop Int
ibl MVector s Word8
ibmv), MVector s Word8 -> Int
forall a s. Storable a => MVector s a -> Int
DVSM.length (Int -> MVector s Word8 -> MVector s Word8
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.drop Int
bpl MVector s Word8
bpmv))

unzipFromListN2 :: (Storable a, Storable b) => Int -> [(a, b)] -> (DVS.Vector a, DVS.Vector b)
unzipFromListN2 :: Int -> [(a, b)] -> (Vector a, Vector b)
unzipFromListN2 Int
n [(a, b)]
abs = (forall s. ST s (Vector a, Vector b)) -> (Vector a, Vector b)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector a, Vector b)) -> (Vector a, Vector b))
-> (forall s. ST s (Vector a, Vector b)) -> (Vector a, Vector b)
forall a b. (a -> b) -> a -> b
$ do
  MVector s a
mas <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
DVSM.unsafeNew Int
n
  MVector s b
mbs <- Int -> ST s (MVector (PrimState (ST s)) b)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
DVSM.unsafeNew Int
n
  Int
len <- Int -> MVector s a -> MVector s b -> [(a, b)] -> ST s Int
forall c d s.
(Storable c, Storable d) =>
Int -> MVector s c -> MVector s d -> [(c, d)] -> ST s Int
go Int
0 MVector s a
mas MVector s b
mbs [(a, b)]
abs
  Vector a
as <- Mutable Vector (PrimState (ST s)) a -> ST s (Vector a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
DVG.unsafeFreeze (Int -> MVector s a -> MVector s a
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.take Int
len MVector s a
mas)
  Vector b
bs <- Mutable Vector (PrimState (ST s)) b -> ST s (Vector b)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
DVG.unsafeFreeze (Int -> MVector s b -> MVector s b
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.take Int
len MVector s b
mbs)
  (Vector a, Vector b) -> ST s (Vector a, Vector b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector a
as, Vector b
bs)
  where go :: (Storable c, Storable d)
          => Int
          -> DVSM.MVector s c
          -> DVSM.MVector s d
          -> [(c, d)]
          -> ST s Int
        go :: Int -> MVector s c -> MVector s d -> [(c, d)] -> ST s Int
go Int
i MVector s c
_   MVector s d
_   []           = Int -> ST s Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
i
        go Int
i MVector s c
mvc MVector s d
mvd ((c
c, d
d):[(c, d)]
cds) = if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
          then do
            MVector (PrimState (ST s)) c -> Int -> c -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
DVSM.write MVector s c
MVector (PrimState (ST s)) c
mvc Int
i c
c
            MVector (PrimState (ST s)) d -> Int -> d -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
DVSM.write MVector s d
MVector (PrimState (ST s)) d
mvd Int
i d
d
            Int -> MVector s c -> MVector s d -> [(c, d)] -> ST s Int
forall c d s.
(Storable c, Storable d) =>
Int -> MVector s c -> MVector s d -> [(c, d)] -> ST s Int
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) MVector s c
mvc MVector s d
mvd [(c, d)]
cds
          else Int -> ST s Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
i