Ticket #7185: Main.hs

import Data.Primitive.MutVar
import qualified Data.Vector.Generic as GV
import qualified Data.Vector.Generic.Mutable as GMV


import Debug.Trace
import Control.Monad
import Control.Monad.Primitive
import qualified Data.Vector.Mutable as MV
import qualified Data.Vector as V
import qualified System.Random.MWC as MWC

data GrowVec v s a = GV ! (MutVar s (v s a)) ! (MutVar s Int)

gvAdd :: (PrimMonad m, GMV.MVector v a) => GrowVec v (PrimState m) a -> a -> m ()
{-# INLINE gvAdd #-}
gvAdd (GV vr cntr) e = do
   v <- readMutVar vr
   cnt <- readMutVar cntr
   
   let l = GMV.length v
   when (l < (cnt + 1)) $ (traceShow l) $ GMV.grow v l >>= writeMutVar vr
   
   v' <- readMutVar vr
   GMV.write v' cnt e
   modifyMutVar cntr (+1)

gvNew :: (PrimMonad m, GMV.MVector v a) => m (GrowVec v (PrimState m) a)
{-# INLINE gvNew #-}
gvNew = do
   v <- GMV.new 1
   vr <- newMutVar v
   cr <- newMutVar 0
   return $! GV vr cr

gvFreeze :: (GV.Vector v a, PrimMonad m) => GrowVec (GV.Mutable v) (PrimState m) a -> m (v a)
{-# INLINE gvFreeze #-}
gvFreeze (GV vr cr) = do
   v <- readMutVar vr
   c <- readMutVar cr
   GV.freeze (GMV.take c v)


main :: IO ()
main = do
   let len = 500
   v <- MV.replicateM len (gvNew :: IO (GrowVec (MV.MVector) (PrimState IO) Float))
   rng <- (MWC.create :: IO MWC.GenIO)
   
   forM_ [0..(len - 1)] $ \ i -> do
      x <- MWC.uniform rng
      MV.read v i >>= \gv -> (replicateM_ i (gvAdd gv x))
   
   v' <- trace "convert" $ V.generateM (MV.length v) $ \i -> (MV.read v i >>= \x -> (gvFreeze x :: IO (V.Vector Float)))
   return ()