module Data.BTree.HashTable.STM
( HashTableSTM
, newSized
, insert
, lookup
, delete
, size
, keys
, toList
) where
import Prelude hiding (lookup)
import Control.Concurrent.STM
import Control.Monad
import Data.Array
import Data.Hashable (Hashable(..))
import Data.Maybe
import qualified Data.List as L
import qualified Data.Map as M
type Bucket k v = M.Map k v
data HashTableSTM k v = HashTableSTM { buckets :: !Int
, count :: TVar Int
, getArray :: Array Int (TVar (Bucket k v)) }
bucket (HashTableSTM n _ arr) k = arr ! (hash k `mod` n)
newSized :: Int -> IO (HashTableSTM k v)
newSized n =
do co <- newTVarIO 0
vs <- mapM (\i -> (i,) `fmap` newTVarIO M.empty) [0..n1]
return $! HashTableSTM n co $! array (0, n1) vs
modifyCount h f = do
n <- readTVar $ count h
writeTVar (count h) $! f n
insert :: (Hashable k, Ord k) => HashTableSTM k v -> k -> v -> STM ()
insert h !k !v = do
m <- readTVar b
let m' = M.insert k v m
writeTVar b $! m'
modifyCount h (+1)
where
b = bucket h k
lookup :: (Hashable k, Ord k) => HashTableSTM k v -> k -> STM (Maybe v)
lookup h !k =
do lst <- readTVar $ bucket h k
return $! M.lookup k lst
delete :: (Hashable k, Ord k) => HashTableSTM k v -> k -> STM ()
delete h !k =
do m <- readTVar b
let (a, m') = M.updateLookupWithKey (\_ _ -> Nothing) k m
when (isJust a) $
modifyCount h (subtract 1)
M.lookup k m' `seq` writeTVar b m'
where
b = bucket h k
size :: HashTableSTM k v -> STM Int
size h = readTVar c
where
c = count h
toList :: HashTableSTM k v -> STM [(k, v)]
toList h =
do bs <- buckets h
return $ concat $ map M.toList bs
where
buckets (HashTableSTM n _ arr) =
mapM readTVar $ elems arr
keys :: HashTableSTM k v -> STM [k]
keys h =
do l <- toList h
return $ map fst l