Low-tech concurrent hashmap

concurrent-map.hs

module ConcurrentMap where

import Control.Concurrent.STM.TVar (TVar)
import Control.Concurrent.STM (STM)
import Data.Hashable (Hashable)
import Data.HashMap.Strict (HashMap)
import Data.Vector (Vector)

import qualified Control.Concurrent.STM.TVar as TVar
import qualified Data.Hashable               as Hashable
import qualified Data.HashMap.Strict         as HashMap
import qualified Data.Vector                 as Vector

newtype Map k v = Map { unMap :: Vector (TVar (HashMap k v)) }

new :: Int -> STM (Map k v)
new n = do
    vector <- Vector.replicateM (2^n) (TVar.newTVar HashMap.empty)
    
    return (Map vector)
    
insert :: (Eq k, Hashable k) => k -> v -> Map k v -> STM ()
insert key value (Map vector) = do
    let index = Hashable.hash key `rem` Vector.length vector
    
    TVar.modifyTVar' (Vector.unsafeIndex vector index) (HashMap.insert key value)
    
lookup :: (Eq k, Hashable k) => k -> Map k v -> STM (Maybe v)
lookup key (Map vector) = do
    let index = Hashable.hash key `rem` Vector.length vector
    
    m <- TVar.readTVar (Vector.unsafeIndex vector index)
    
    return (HashMap.lookup key m)