15 | | import Control.Monad ((=<<), liftM, unless, when) |
16 | | import Data.Bits ((.&.), shiftL, shiftR) |
17 | | import Data.IORef (IORef, newIORef, readIORef, writeIORef) |
18 | | import Data.Maybe (Maybe(..), isJust) |
19 | | import Foreign.ForeignPtr (ForeignPtr, mallocForeignPtr, withForeignPtr) |
20 | | import Foreign.Storable (peek, poke) |
21 | | import GHC.Base (Monad(..), ($), const, otherwise) |
22 | | import GHC.Classes (Eq(..), Ord(..)) |
| 13 | import Prelude (($!)) |
| 14 | import Control.Monad |
| 15 | import Data.Bits |
| 16 | import Data.IORef |
| 17 | import Data.Maybe |
| 18 | import Foreign.ForeignPtr |
| 19 | import Foreign.Storable |
48 | | lookup k (IntTable ref) = do |
49 | | let go Bucket{..} |
50 | | | bucketKey == k = return (Just bucketValue) |
51 | | | otherwise = go bucketNext |
52 | | go _ = return Nothing |
53 | | it@IT{..} <- readIORef ref |
54 | | go =<< Arr.read tabArr (indexOf k it) |
| 46 | lookup k IntTable{..} = do |
| 47 | arr <- readIORef tblRef |
| 48 | bkt <- A.read arr (indexOf k arr) |
| 49 | return (go bkt) |
| 50 | where go Empty = Nothing |
| 51 | go Bucket{..} |
| 52 | | bucketKey == k = Just bucketVal |
| 53 | | otherwise = go bucketNext |
68 | | grow :: IT a -> IORef (IT a) -> Int -> IO () |
69 | | grow oldit ref size = do |
70 | | newit <- new_ (Arr.size (tabArr oldit) `shiftL` 1) |
71 | | let copySlot n !i |
72 | | | n == size = return () |
73 | | | otherwise = do |
74 | | let copyBucket !m Empty = copySlot m (i+1) |
75 | | copyBucket m bkt@Bucket{..} = do |
76 | | let idx = indexOf bucketKey newit |
77 | | next <- Arr.read (tabArr newit) idx |
78 | | Arr.write (tabArr newit) idx bkt { bucketNext = next } |
79 | | copyBucket (m+1) bucketNext |
80 | | copyBucket n =<< Arr.read (tabArr oldit) i |
81 | | copySlot 0 0 |
82 | | withForeignPtr (tabSize newit) $ \ptr -> poke ptr size |
83 | | writeIORef ref newit |
| 63 | copy :: Arr (Bucket a) -> Arr (Bucket a) -> Int -> IO () |
| 64 | copy src dst size = A.read src 0 >>= go 1 0 |
| 65 | where go i n Empty |
| 66 | | n == size = return () |
| 67 | | otherwise = A.read src i >>= go (i + 1) n |
| 68 | go i n bkt@Bucket{..} = do |
| 69 | let ix = bucketKey .&. (A.size dst - 1) |
| 70 | next <- A.read dst ix |
| 71 | A.write dst ix (bkt { bucketNext = next }) |
| 72 | go i (n + 1) bucketNext |
86 | | insertWith f k v inttable@(IntTable ref) = do |
87 | | it@IT{..} <- readIORef ref |
88 | | let idx = indexOf k it |
89 | | go seen bkt@Bucket{..} |
90 | | | bucketKey == k = do |
91 | | let !v' = f v bucketValue |
92 | | !next = seen <> bucketNext |
93 | | Empty <> bs = bs |
94 | | b@Bucket{..} <> bs = b { bucketNext = bucketNext <> bs } |
95 | | Arr.write tabArr idx (Bucket k v' next) |
96 | | return (Just bucketValue) |
97 | | | otherwise = go bkt { bucketNext = seen } bucketNext |
98 | | go seen _ = withForeignPtr tabSize $ \ptr -> do |
99 | | size <- peek ptr |
100 | | if size + 1 >= Arr.size tabArr - (Arr.size tabArr `shiftR` 2) |
101 | | then grow it ref size >> insertWith f k v inttable |
102 | | else do |
103 | | v `seq` Arr.write tabArr idx (Bucket k v seen) |
104 | | poke ptr (size + 1) |
105 | | return Nothing |
106 | | go Empty =<< Arr.read tabArr idx |
107 | | {-# INLINABLE insertWith #-} |
| 75 | insertWith f k v tbl@IntTable{..} = do |
| 76 | arr <- readIORef tblRef |
| 77 | let ix = indexOf k arr |
| 78 | go cont Empty = do |
| 79 | size <- withForeignPtr tblSize peek |
| 80 | if size + 1 >= A.size arr - (A.size arr `unsafeShiftR` 2) |
| 81 | then do arr' <- A.new Empty (A.size arr `unsafeShiftL` 1) |
| 82 | copy arr arr' size |
| 83 | writeIORef tblRef arr' |
| 84 | insertWith f k v tbl |
| 85 | else do A.write arr ix $! cont (Bucket k v Empty) |
| 86 | withForeignPtr tblSize (`poke` (size + 1)) |
| 87 | return Nothing |
| 88 | go cont bkt@Bucket{..} |
| 89 | | bucketKey == k = do |
| 90 | let !v' = f v bucketVal |
| 91 | A.write arr ix $! cont (bkt { bucketVal = v' }) |
| 92 | return (Just bucketVal) |
| 93 | | otherwise = go (\x -> cont (bkt { bucketNext = x })) bucketNext |
| 94 | A.read arr ix >>= go id |
121 | | updateWith f k (IntTable ref) = do |
122 | | it@IT{..} <- readIORef ref |
123 | | let idx = indexOf k it |
124 | | go changed bkt@Bucket{..} |
125 | | | bucketKey == k = |
126 | | let fbv = f bucketValue |
127 | | !nb = case fbv of |
128 | | Just val -> bkt { bucketValue = val } |
129 | | Nothing -> bucketNext |
130 | | in (fbv, Just bucketValue, nb) |
131 | | | otherwise = case go changed bucketNext of |
132 | | (fbv, ov, nb) -> (fbv, ov, bkt { bucketNext = nb }) |
133 | | go _ e = (Nothing, Nothing, e) |
134 | | (fbv, oldVal, newBucket) <- go False `liftM` Arr.read tabArr idx |
135 | | when (isJust oldVal) $ do |
136 | | Arr.write tabArr idx newBucket |
137 | | unless (isJust fbv) $ |
138 | | withForeignPtr tabSize $ \ptr -> do |
139 | | size <- peek ptr |
140 | | poke ptr (size - 1) |
141 | | return oldVal |
| 104 | updateWith f k IntTable{..} = do |
| 105 | arr <- readIORef tblRef |
| 106 | let ix = indexOf k arr |
| 107 | go _ Empty = (False, Nothing, Empty) |
| 108 | go cont bkt@Bucket{..} |
| 109 | | bucketKey == k = case f bucketVal of |
| 110 | Nothing -> (True, Just bucketVal, cont bucketNext) |
| 111 | Just v -> (False, Just bucketVal, cont (bkt { bucketVal = v })) |
| 112 | | otherwise = go (\x -> cont (bkt { bucketNext = x })) bucketNext |
| 113 | (del, old, bkt) <- A.read arr ix >>= return . go id |
| 114 | when (isJust old) $ do |
| 115 | A.write arr ix bkt |
| 116 | when del . withForeignPtr tblSize $ \ptr -> do |
| 117 | size <- peek ptr |
| 118 | poke ptr (size - 1) |
| 119 | return old |