Ticket #7185: Main.hs

File Main.hs, 1.5 KB (added by waldheinz, 3 years ago)

self-contained example

Line 
1
2import Data.Primitive.MutVar
3import qualified Data.Vector.Generic as GV
4import qualified Data.Vector.Generic.Mutable as GMV
5
6
7import Debug.Trace
8import Control.Monad
9import Control.Monad.Primitive
10import qualified Data.Vector.Mutable as MV
11import qualified Data.Vector as V
12import qualified System.Random.MWC as MWC
13
14data GrowVec v s a = GV ! (MutVar s (v s a)) ! (MutVar s Int)
15
16gvAdd :: (PrimMonad m, GMV.MVector v a) => GrowVec v (PrimState m) a -> a -> m ()
17{-# INLINE gvAdd #-}
18gvAdd (GV vr cntr) e = do
19   v <- readMutVar vr
20   cnt <- readMutVar cntr
21   
22   let l = GMV.length v
23   when (l < (cnt + 1)) $ (traceShow l) $ GMV.grow v l >>= writeMutVar vr
24   
25   v' <- readMutVar vr
26   GMV.write v' cnt e
27   modifyMutVar cntr (+1)
28
29gvNew :: (PrimMonad m, GMV.MVector v a) => m (GrowVec v (PrimState m) a)
30{-# INLINE gvNew #-}
31gvNew = do
32   v <- GMV.new 1
33   vr <- newMutVar v
34   cr <- newMutVar 0
35   return $! GV vr cr
36
37gvFreeze :: (GV.Vector v a, PrimMonad m) => GrowVec (GV.Mutable v) (PrimState m) a -> m (v a)
38{-# INLINE gvFreeze #-}
39gvFreeze (GV vr cr) = do
40   v <- readMutVar vr
41   c <- readMutVar cr
42   GV.freeze (GMV.take c v)
43
44
45main :: IO ()
46main = do
47   let len = 500
48   v <- MV.replicateM len (gvNew :: IO (GrowVec (MV.MVector) (PrimState IO) Float))
49   rng <- (MWC.create :: IO MWC.GenIO)
50   
51   forM_ [0..(len - 1)] $ \ i -> do
52      x <- MWC.uniform rng
53      MV.read v i >>= \gv -> (replicateM_ i (gvAdd gv x))
54   
55   v' <- trace "convert" $ V.generateM (MV.length v) $ \i -> (MV.read v i >>= \x -> (gvFreeze x :: IO (V.Vector Float)))
56   return ()
57