16 withIStore2, |
16 withIStore2, |
17 (!), |
17 (!), |
18 indices |
18 indices |
19 ) where |
19 ) where |
20 |
20 |
21 import qualified Data.Array.IArray as IA |
|
22 import qualified Data.Array.IO as IOA |
|
23 import qualified Data.IntSet as IntSet |
21 import qualified Data.IntSet as IntSet |
|
22 import qualified Data.Vector as V |
|
23 import qualified Data.Vector.Mutable as MV |
24 import Data.IORef |
24 import Data.IORef |
25 import Control.Monad |
25 import Control.Monad |
26 import Control.DeepSeq |
26 import Control.DeepSeq |
27 |
27 |
28 |
28 |
29 newtype ElemIndex = ElemIndex Int |
29 newtype ElemIndex = ElemIndex Int |
30 deriving (Eq, Show, Read, Ord, NFData) |
30 deriving (Eq, Show, Read, Ord, NFData) |
31 newtype MStore e = MStore (IORef (IntSet.IntSet, IntSet.IntSet, IOA.IOArray Int e)) |
31 newtype MStore e = MStore (IORef (IntSet.IntSet, IntSet.IntSet, MV.IOVector e)) |
32 newtype IStore e = IStore (IntSet.IntSet, IA.Array Int e) |
32 newtype IStore e = IStore (IntSet.IntSet, V.Vector e) |
33 |
33 |
34 |
34 |
35 firstIndex :: ElemIndex |
35 firstIndex :: ElemIndex |
36 firstIndex = ElemIndex 0 |
36 firstIndex = ElemIndex 0 |
37 |
37 |
48 | otherwise = a |
48 | otherwise = a |
49 |
49 |
50 |
50 |
51 newStore :: IO (MStore e) |
51 newStore :: IO (MStore e) |
52 newStore = do |
52 newStore = do |
53 newar <- IOA.newArray_ (0, initialSize - 1) |
53 newar <- MV.new initialSize |
54 new <- newIORef (IntSet.empty, IntSet.fromAscList [0..initialSize - 1], newar) |
54 new <- newIORef (IntSet.empty, IntSet.fromAscList [0..initialSize - 1], newar) |
55 return (MStore new) |
55 return (MStore new) |
56 |
56 |
57 |
57 |
58 growStore :: MStore e -> IO () |
58 growStore :: MStore e -> IO () |
59 growStore (MStore ref) = do |
59 growStore (MStore ref) = do |
60 (busyElems, freeElems, arr) <- readIORef ref |
60 (busyElems, freeElems, arr) <- readIORef ref |
61 (_, m') <- IOA.getBounds arr |
61 let oldSize = MV.length arr |
62 let newM' = growFunc (m' + 1) - 1 |
62 let newSize = growFunc oldSize |
63 newArr <- IOA.newArray_ (0, newM') |
63 newArr <- MV.grow arr (newSize - oldSize) |
64 sequence_ [IOA.readArray arr i >>= IOA.writeArray newArr i | i <- [0..m']] |
64 writeIORef ref (busyElems, freeElems `IntSet.union` IntSet.fromAscList [oldSize .. newSize-1], newArr) |
65 writeIORef ref (busyElems, freeElems `IntSet.union` IntSet.fromAscList [m'+1..newM'], newArr) |
|
66 |
65 |
67 |
66 |
68 growIfNeeded :: MStore e -> IO () |
67 growIfNeeded :: MStore e -> IO () |
69 growIfNeeded m@(MStore ref) = do |
68 growIfNeeded m@(MStore ref) = do |
70 (_, freeElems, _) <- readIORef ref |
69 (_, freeElems, _) <- readIORef ref |
72 |
71 |
73 |
72 |
74 truncateIfNeeded :: MStore e -> IO () |
73 truncateIfNeeded :: MStore e -> IO () |
75 truncateIfNeeded (MStore ref) = do |
74 truncateIfNeeded (MStore ref) = do |
76 (busyElems, _, arr) <- readIORef ref |
75 (busyElems, _, arr) <- readIORef ref |
77 (_, m') <- IOA.getBounds arr |
76 let oldSize = MV.length arr |
78 let newM' = truncFunc (m' + 1) - 1 |
77 let newSize = truncFunc oldSize |
79 when (newM' < m' && (not $ IntSet.null busyElems) && IntSet.findMax busyElems <= newM') $ do |
78 when (newSize < oldSize && (not $ IntSet.null busyElems) && IntSet.findMax busyElems < newSize) $ do |
80 newArr <- IOA.newArray_ (0, newM') |
79 writeIORef ref (busyElems, IntSet.fromAscList [0..newSize - 1] `IntSet.difference` busyElems, MV.take newSize arr) |
81 sequence_ [IOA.readArray arr i >>= IOA.writeArray newArr i | i <- IntSet.toList busyElems] |
|
82 writeIORef ref (busyElems, IntSet.fromAscList [0..newM'] `IntSet.difference` busyElems, newArr) |
|
83 |
80 |
84 |
81 |
85 addElem :: MStore e -> e -> IO ElemIndex |
82 addElem :: MStore e -> e -> IO ElemIndex |
86 addElem m@(MStore ref) element = do |
83 addElem m@(MStore ref) element = do |
87 growIfNeeded m |
84 growIfNeeded m |
88 (busyElems, freeElems, arr) <- readIORef ref |
85 (busyElems, freeElems, arr) <- readIORef ref |
89 let (!n, freeElems') = IntSet.deleteFindMin freeElems |
86 let (!n, freeElems') = IntSet.deleteFindMin freeElems |
90 IOA.writeArray arr n element |
87 MV.write arr n element |
91 writeIORef ref (IntSet.insert n busyElems, freeElems', arr) |
88 writeIORef ref (IntSet.insert n busyElems, freeElems', arr) |
92 return $ ElemIndex n |
89 return $ ElemIndex n |
93 |
90 |
94 |
91 |
95 removeElem :: MStore e -> ElemIndex -> IO () |
92 removeElem :: MStore e -> ElemIndex -> IO () |
96 removeElem m@(MStore ref) (ElemIndex n) = do |
93 removeElem m@(MStore ref) (ElemIndex n) = do |
97 (busyElems, freeElems, arr) <- readIORef ref |
94 (busyElems, freeElems, arr) <- readIORef ref |
98 IOA.writeArray arr n (error $ "Store: no element " ++ show n) |
95 MV.write arr n (error $ "Store: no element " ++ show n) |
99 writeIORef ref (IntSet.delete n busyElems, IntSet.insert n freeElems, arr) |
96 writeIORef ref (IntSet.delete n busyElems, IntSet.insert n freeElems, arr) |
100 truncateIfNeeded m |
97 truncateIfNeeded m |
101 |
98 |
102 |
99 |
103 readElem :: MStore e -> ElemIndex -> IO e |
100 readElem :: MStore e -> ElemIndex -> IO e |
104 readElem (MStore ref) (ElemIndex n) = readIORef ref >>= \(_, _, arr) -> IOA.readArray arr n |
101 readElem (MStore ref) (ElemIndex n) = readIORef ref >>= \(_, _, arr) -> MV.read arr n |
105 |
102 |
106 |
103 |
107 writeElem :: MStore e -> ElemIndex -> e -> IO () |
104 writeElem :: MStore e -> ElemIndex -> e -> IO () |
108 writeElem (MStore ref) (ElemIndex n) el = readIORef ref >>= \(_, _, arr) -> IOA.writeArray arr n el |
105 writeElem (MStore ref) (ElemIndex n) el = readIORef ref >>= \(_, _, arr) -> MV.write arr n el |
109 |
106 |
110 |
107 |
111 modifyElem :: MStore e -> (e -> e) -> ElemIndex -> IO () |
108 modifyElem :: MStore e -> (e -> e) -> ElemIndex -> IO () |
112 modifyElem (MStore ref) f (ElemIndex n) = do |
109 modifyElem (MStore ref) f (ElemIndex n) = do |
113 (_, _, arr) <- readIORef ref |
110 (_, _, arr) <- readIORef ref |
114 IOA.readArray arr n >>= IOA.writeArray arr n . f |
111 MV.read arr n >>= MV.write arr n . f |
115 |
112 |
116 elemExists :: MStore e -> ElemIndex -> IO Bool |
113 elemExists :: MStore e -> ElemIndex -> IO Bool |
117 elemExists (MStore ref) (ElemIndex n) = do |
114 elemExists (MStore ref) (ElemIndex n) = do |
118 (_, !free, _) <- readIORef ref |
115 (_, !free, _) <- readIORef ref |
119 return $ n `IntSet.notMember` free |
116 return $ n `IntSet.notMember` free |
126 |
123 |
127 -- A way to see MStore elements in pure code via IStore |
124 -- A way to see MStore elements in pure code via IStore |
128 m2i :: MStore e -> IO (IStore e) |
125 m2i :: MStore e -> IO (IStore e) |
129 m2i (MStore ref) = do |
126 m2i (MStore ref) = do |
130 (a, _, c') <- readIORef ref |
127 (a, _, c') <- readIORef ref |
131 c <- IOA.unsafeFreeze c' |
128 c <- V.unsafeFreeze c' |
132 return $ IStore (a, c) |
129 return $ IStore (a, c) |
133 |
130 |
134 i2m :: MStore e -> IStore e -> IO () |
131 i2m :: MStore e -> IStore e -> IO () |
135 i2m (MStore ref) (IStore (_, arr)) = do |
132 i2m (MStore ref) (IStore (_, arr)) = do |
136 (b, e, _) <- readIORef ref |
133 (b, e, _) <- readIORef ref |
137 a <- IOA.unsafeThaw arr |
134 a <- V.unsafeThaw arr |
138 writeIORef ref (b, e, a) |
135 writeIORef ref (b, e, a) |
139 |
136 |
140 withIStore :: MStore e -> (IStore e -> a) -> IO a |
137 withIStore :: MStore e -> (IStore e -> a) -> IO a |
141 withIStore m f = do |
138 withIStore m f = do |
142 i <- m2i m |
139 i <- m2i m |