1{-# LANGUAGE BangPatterns #-} 2{-# LANGUAGE CPP #-} 3 4module Data.HashTable.Internal.Linear.Bucket 5( Bucket, 6 newBucketArray, 7 newBucketSize, 8 emptyWithSize, 9 growBucketTo, 10 snoc, 11 size, 12 lookup, 13 lookupIndex, 14 elemAt, 15 delete, 16 mutate, 17 mutateST, 18 toList, 19 fromList, 20 mapM_, 21 foldM, 22 expandBucketArray, 23 expandArray, 24 nelemsAndOverheadInWords, 25 bucketSplitSize 26) where 27 28 29------------------------------------------------------------------------------ 30#if !MIN_VERSION_base(4,8,0) 31import Control.Applicative 32#endif 33import Control.Monad hiding (foldM, mapM_) 34import qualified Control.Monad 35import Control.Monad.ST (ST) 36#ifdef DEBUG 37import Data.HashTable.Internal.Utils (unsafeIOToST) 38#endif 39import Data.HashTable.Internal.Array 40import Data.Maybe (fromMaybe) 41import Data.STRef 42import Prelude hiding (lookup, mapM_) 43------------------------------------------------------------------------------ 44import Data.HashTable.Internal.UnsafeTricks 45 46 47#ifdef DEBUG 48import System.IO 49#endif 50 51 52type Bucket s k v = Key (Bucket_ s k v) 53 54------------------------------------------------------------------------------ 55data Bucket_ s k v = Bucket { _bucketSize :: {-# UNPACK #-} !Int 56 , _highwater :: {-# UNPACK #-} !(STRef s Int) 57 , _keys :: {-# UNPACK #-} !(MutableArray s k) 58 , _values :: {-# UNPACK #-} !(MutableArray s v) 59 } 60 61 62------------------------------------------------------------------------------ 63bucketSplitSize :: Int 64bucketSplitSize = 16 65 66 67------------------------------------------------------------------------------ 68newBucketArray :: Int -> ST s (MutableArray s (Bucket s k v)) 69newBucketArray k = newArray k emptyRecord 70 71------------------------------------------------------------------------------ 72nelemsAndOverheadInWords :: Bucket s k v -> ST s (Int,Int) 73nelemsAndOverheadInWords bKey = do 74 if (not $ keyIsEmpty bKey) 75 then do 76 !hw <- readSTRef hwRef 77 let !w = sz - hw 78 return (hw, constOverhead + 2*w) 79 else 80 return (0, 0) 81 82 where 83 constOverhead = 8 84 b = fromKey bKey 85 sz = _bucketSize b 86 hwRef = _highwater b 87 88 89------------------------------------------------------------------------------ 90emptyWithSize :: Int -> ST s (Bucket s k v) 91emptyWithSize !sz = do 92 !keys <- newArray sz undefined 93 !values <- newArray sz undefined 94 !ref <- newSTRef 0 95 96 return $ toKey $ Bucket sz ref keys values 97 98 99------------------------------------------------------------------------------ 100newBucketSize :: Int 101newBucketSize = 4 102 103 104------------------------------------------------------------------------------ 105expandArray :: a -- ^ default value 106 -> Int -- ^ new size 107 -> Int -- ^ number of elements to copy 108 -> MutableArray s a -- ^ old array 109 -> ST s (MutableArray s a) 110expandArray def !sz !hw !arr = do 111 newArr <- newArray sz def 112 cp newArr 113 114 where 115 cp !newArr = go 0 116 where 117 go !i 118 | i >= hw = return newArr 119 | otherwise = do 120 readArray arr i >>= writeArray newArr i 121 go (i+1) 122 123 124------------------------------------------------------------------------------ 125expandBucketArray :: Int 126 -> Int 127 -> MutableArray s (Bucket s k v) 128 -> ST s (MutableArray s (Bucket s k v)) 129expandBucketArray = expandArray emptyRecord 130 131 132------------------------------------------------------------------------------ 133growBucketTo :: Int -> Bucket s k v -> ST s (Bucket s k v) 134growBucketTo !sz bk | keyIsEmpty bk = emptyWithSize sz 135 | otherwise = do 136 if osz >= sz 137 then return bk 138 else do 139 hw <- readSTRef hwRef 140 k' <- expandArray undefined sz hw keys 141 v' <- expandArray undefined sz hw values 142 return $ toKey $ Bucket sz hwRef k' v' 143 144 where 145 bucket = fromKey bk 146 osz = _bucketSize bucket 147 hwRef = _highwater bucket 148 keys = _keys bucket 149 values = _values bucket 150 151 152------------------------------------------------------------------------------ 153{-# INLINE snoc #-} 154-- Just return == new bucket object 155snoc :: Bucket s k v -> k -> v -> ST s (Int, Maybe (Bucket s k v)) 156snoc bucket | keyIsEmpty bucket = mkNew 157 | otherwise = snoc' (fromKey bucket) 158 where 159 mkNew !k !v = do 160 debug "Bucket.snoc: mkNew" 161 keys <- newArray newBucketSize undefined 162 values <- newArray newBucketSize undefined 163 164 writeArray keys 0 k 165 writeArray values 0 v 166 ref <- newSTRef 1 167 return (1, Just $ toKey $ Bucket newBucketSize ref keys values) 168 169 snoc' (Bucket bsz hwRef keys values) !k !v = 170 readSTRef hwRef >>= check 171 where 172 check !hw 173 | hw < bsz = bump hw 174 | otherwise = spill hw 175 176 bump hw = do 177 debug $ "Bucket.snoc: bumping hw, bsz=" ++ show bsz ++ ", hw=" 178 ++ show hw 179 180 writeArray keys hw k 181 writeArray values hw v 182 let !hw' = hw + 1 183 writeSTRef hwRef hw' 184 debug "Bucket.snoc: finished" 185 return (hw', Nothing) 186 187 doublingThreshold = bucketSplitSize `div` 2 188 growFactor = 1.5 :: Double 189 newSize z | z == 0 = newBucketSize 190 | z < doublingThreshold = z * 2 191 | otherwise = ceiling $ growFactor * fromIntegral z 192 193 spill !hw = do 194 let sz = newSize bsz 195 debug $ "Bucket.snoc: spilling, old size=" ++ show bsz ++ ", new size=" 196 ++ show sz 197 198 bk <- growBucketTo sz bucket 199 200 debug "Bucket.snoc: spill finished, snoccing element" 201 let (Bucket _ hwRef' keys' values') = fromKey bk 202 203 let !hw' = hw+1 204 writeArray keys' hw k 205 writeArray values' hw v 206 writeSTRef hwRef' hw' 207 208 return (hw', Just bk) 209 210 211 212------------------------------------------------------------------------------ 213{-# INLINE size #-} 214size :: Bucket s k v -> ST s Int 215size b | keyIsEmpty b = return 0 216 | otherwise = readSTRef $ _highwater $ fromKey b 217 218 219------------------------------------------------------------------------------ 220-- note: search in reverse order! We prefer recently snoc'd keys. 221lookup :: (Eq k) => Bucket s k v -> k -> ST s (Maybe v) 222lookup bucketKey !k | keyIsEmpty bucketKey = return Nothing 223 | otherwise = lookup' $ fromKey bucketKey 224 where 225 lookup' (Bucket _ hwRef keys values) = do 226 hw <- readSTRef hwRef 227 go (hw-1) 228 where 229 go !i 230 | i < 0 = return Nothing 231 | otherwise = do 232 k' <- readArray keys i 233 if k == k' 234 then do 235 !v <- readArray values i 236 return $! Just v 237 else go (i-1) 238 239------------------------------------------------------------------------------ 240-- note: search in reverse order! We prefer recently snoc'd keys. 241lookupIndex :: (Eq k) => Bucket s k v -> k -> ST s (Maybe Int) 242lookupIndex bucketKey !k 243 | keyIsEmpty bucketKey = return Nothing 244 | otherwise = lookup' $ fromKey bucketKey 245 where 246 lookup' (Bucket _ hwRef keys _values) = do 247 hw <- readSTRef hwRef 248 go (hw-1) 249 where 250 go !i 251 | i < 0 = return Nothing 252 | otherwise = do 253 k' <- readArray keys i 254 if k == k' 255 then return (Just i) 256 else go (i-1) 257 258elemAt :: Bucket s k v -> Int -> ST s (Maybe (k,v)) 259elemAt bucketKey ix 260 | keyIsEmpty bucketKey = return Nothing 261 | otherwise = lookup' $ fromKey bucketKey 262 where 263 lookup' (Bucket _ hwRef keys values) = do 264 hw <- readSTRef hwRef 265 if 0 <= ix && ix < hw 266 then do k <- readArray keys ix 267 v <- readArray values ix 268 return (Just (k,v)) 269 else return Nothing 270 271------------------------------------------------------------------------------ 272{-# INLINE toList #-} 273toList :: Bucket s k v -> ST s [(k,v)] 274toList bucketKey | keyIsEmpty bucketKey = return [] 275 | otherwise = toList' $ fromKey bucketKey 276 where 277 toList' (Bucket _ hwRef keys values) = do 278 hw <- readSTRef hwRef 279 go [] hw 0 280 where 281 go !l !hw !i | i >= hw = return l 282 | otherwise = do 283 k <- readArray keys i 284 v <- readArray values i 285 go ((k,v):l) hw $ i+1 286 287 288------------------------------------------------------------------------------ 289-- fromList needs to reverse the input in order to make fromList . toList == id 290{-# INLINE fromList #-} 291fromList :: [(k,v)] -> ST s (Bucket s k v) 292fromList l = Control.Monad.foldM f emptyRecord (reverse l) 293 where 294 f bucket (k,v) = do 295 (_,m) <- snoc bucket k v 296 return $ fromMaybe bucket m 297 298------------------------------------------------------------------------------ 299delete :: (Eq k) => Bucket s k v -> k -> ST s Bool 300delete bucketKey !k | keyIsEmpty bucketKey = do 301 debug $ "Bucket.delete: empty bucket" 302 return False 303 | otherwise = do 304 debug "Bucket.delete: start" 305 del $ fromKey bucketKey 306 where 307 del (Bucket sz hwRef keys values) = do 308 hw <- readSTRef hwRef 309 debug $ "Bucket.delete: hw=" ++ show hw ++ ", sz=" ++ show sz 310 go hw $ hw - 1 311 312 where 313 go !hw !i | i < 0 = return False 314 | otherwise = do 315 k' <- readArray keys i 316 if k == k' 317 then do 318 debug $ "found entry to delete at " ++ show i 319 move (hw-1) i keys 320 move (hw-1) i values 321 let !hw' = hw-1 322 writeSTRef hwRef hw' 323 return True 324 else go hw (i-1) 325 326 327------------------------------------------------------------------------------ 328mutate :: (Eq k) => 329 Bucket s k v 330 -> k 331 -> (Maybe v -> (Maybe v, a)) 332 -> ST s (Int, Maybe (Bucket s k v), a) 333mutate bucketKey !k !f = mutateST bucketKey k (pure . f) 334{-# INLINE mutate #-} 335 336 337------------------------------------------------------------------------------ 338mutateST :: (Eq k) => 339 Bucket s k v 340 -> k 341 -> (Maybe v -> ST s (Maybe v, a)) 342 -> ST s (Int, Maybe (Bucket s k v), a) 343mutateST bucketKey !k !f 344 | keyIsEmpty bucketKey = do 345 fRes <- f Nothing 346 case fRes of 347 (Nothing, a) -> return (0, Nothing, a) 348 (Just v', a) -> do 349 (!hw', mbk) <- snoc bucketKey k v' 350 return (hw', mbk, a) 351 | otherwise = mutate' $ fromKey bucketKey 352 where 353 mutate' (Bucket _sz hwRef keys values) = do 354 hw <- readSTRef hwRef 355 pos <- findPosition hw (hw-1) 356 mv <- do 357 if pos < 0 358 then return Nothing 359 else readArray values pos >>= return . Just 360 fRes <- f mv 361 case (mv, fRes) of 362 (Nothing, (Nothing, a)) -> return (hw, Nothing, a) 363 (Nothing, (Just v', a)) -> do 364 (!hw', mbk) <- snoc bucketKey k v' 365 return (hw', mbk, a) 366 (Just _v, (Just v', a)) -> do 367 writeArray values pos v' 368 return (hw, Nothing, a) 369 (Just _v, (Nothing, a)) -> do 370 move (hw-1) pos keys 371 move (hw-1) pos values 372 let !hw' = hw-1 373 writeSTRef hwRef hw' 374 return (hw', Nothing, a) 375 where 376 findPosition !hw !i 377 | i < 0 = return (-1) 378 | otherwise = do 379 k' <- readArray keys i 380 if k == k' 381 then return i 382 else findPosition hw (i-1) 383 384 385------------------------------------------------------------------------------ 386{-# INLINE mapM_ #-} 387mapM_ :: ((k,v) -> ST s a) -> Bucket s k v -> ST s () 388mapM_ f bucketKey 389 | keyIsEmpty bucketKey = do 390 debug $ "Bucket.mapM_: bucket was empty" 391 return () 392 | otherwise = doMap $ fromKey bucketKey 393 where 394 doMap (Bucket sz hwRef keys values) = do 395 hw <- readSTRef hwRef 396 debug $ "Bucket.mapM_: hw was " ++ show hw ++ ", sz was " ++ show sz 397 go hw 0 398 where 399 go !hw !i | i >= hw = return () 400 | otherwise = do 401 k <- readArray keys i 402 v <- readArray values i 403 _ <- f (k,v) 404 go hw $ i+1 405 406 407------------------------------------------------------------------------------ 408{-# INLINE foldM #-} 409foldM :: (a -> (k,v) -> ST s a) -> a -> Bucket s k v -> ST s a 410foldM f !seed0 bucketKey 411 | keyIsEmpty bucketKey = return seed0 412 | otherwise = doMap $ fromKey bucketKey 413 where 414 doMap (Bucket _ hwRef keys values) = do 415 hw <- readSTRef hwRef 416 go hw seed0 0 417 where 418 go !hw !seed !i | i >= hw = return seed 419 | otherwise = do 420 k <- readArray keys i 421 v <- readArray values i 422 seed' <- f seed (k,v) 423 go hw seed' (i+1) 424 425 426------------------------------------------------------------------------------ 427-- move i into j 428move :: Int -> Int -> MutableArray s a -> ST s () 429move i j arr | i == j = do 430 debug $ "move " ++ show i ++ " into " ++ show j 431 return () 432 | otherwise = do 433 debug $ "move " ++ show i ++ " into " ++ show j 434 readArray arr i >>= writeArray arr j 435 436 437 438{-# INLINE debug #-} 439debug :: String -> ST s () 440 441#ifdef DEBUG 442debug s = unsafeIOToST $ do 443 putStrLn s 444 hFlush stdout 445#else 446#ifdef TESTSUITE 447debug !s = do 448 let !_ = length s 449 return $! () 450#else 451debug _ = return () 452#endif 453#endif 454 455