module SDR.VectorUtils (
mapAccumMV,
stride,
fill,
) where
import Control.Monad
import Control.Monad.Primitive
import Data.Vector.Generic as VG hiding ((++))
import qualified Data.Vector.Generic.Mutable as VGM
import Data.Vector.Fusion.Stream.Monadic hiding ((++))
import qualified Data.Vector.Fusion.Stream as VFS hiding ((++))
import qualified Data.Vector.Fusion.Stream.Monadic as VFSM hiding ((++))
mapAccumMV :: (Monad m)
=> (acc -> x -> m (acc, y))
-> acc
-> Stream m x
-> Stream m y
mapAccumMV func z (Stream step s sz) = Stream step' (s, z) sz
where
step' (s, acc) = do
r <- step s
case r of
Yield y s' -> do
(!acc', !res) <- func acc y
return $ Yield res (s', acc')
Skip s' -> return $ Skip (s', acc)
Done -> return Done
stride :: VG.Vector v a
=> Int
-> v a
-> v a
stride str inv = VG.unstream $ VFS.unfoldr func 0
where
len = VG.length inv
func i | i >= len = Nothing
| otherwise = Just (VG.unsafeIndex inv i, i + str)
fill :: (PrimMonad m, Functor m, VGM.MVector vm a)
=> VFS.MStream m a
-> vm (PrimState m) a
-> m ()
fill str outBuf = void $ VFSM.foldM' put 0 str
where
put i x = do
VGM.unsafeWrite outBuf i x
return $ i + 1