{-# LANGUAGE RecordWildCards #-}
module Data.SlotMap
( Key
, SlotMap
, empty
, clone
, Data.SlotMap.lookup
, delete
, unsafeDelete
, insert
, update
, unsafeUpdate
, Data.SlotMap.null
, capacity
, size
, Data.SlotMap.foldr
, Data.SlotMap.map
, elems
) where
import Data.Primitive.MutVar
import Control.Monad.Primitive (PrimMonad, PrimState)
import Data.Vector.Mutable (MVector)
import qualified Data.Vector.Mutable as V
data Key = Key
{ keyIndex :: Int
, keyGeneration :: Int
}
data Slot a = Slot
{ slotItem :: Maybe a
, slotGeneration :: Int
, slotFreeIndex :: Int
}
data SlotMap m a = SlotMap
{ slotMapFree :: MutVar (PrimState m) Int
, slotMapCount :: MutVar (PrimState m) Int
, slotMapItems :: MutVar (PrimState m) (MVector (PrimState m) (Slot a))
}
empty :: PrimMonad m => m (SlotMap m a)
empty = do
slotMapFree <- newMutVar 0
slotMapCount <- newMutVar 0
slotMapItems <- V.new 0 >>= newMutVar
return SlotMap{..}
clone :: PrimMonad m => (SlotMap m a) -> m (SlotMap m a)
clone map = do
slotMapFree <- readMutVar (slotMapFree map) >>= newMutVar
slotMapCount <- readMutVar (slotMapCount map) >>= newMutVar
slotMapItems <- readMutVar (slotMapItems map) >>= V.clone >>= newMutVar
return SlotMap{..}
lookup :: PrimMonad m => Key -> SlotMap m a -> m (Maybe a)
lookup Key{..} SlotMap{..} = do
items <- readMutVar slotMapItems
if keyIndex >= V.length items then pure Nothing else do
slot <- V.unsafeRead items keyIndex
if slotGeneration slot /= keyGeneration then pure Nothing else
pure $ slotItem slot
delete :: PrimMonad m => Key -> SlotMap m a -> m ()
delete key map = do
item <- Data.SlotMap.lookup key map
case item of
Nothing -> pure ()
Just _ -> unsafeDelete key map
unsafeDelete :: PrimMonad m => Key -> SlotMap m a -> m ()
unsafeDelete Key{..} SlotMap{..} = do
modifyMutVar slotMapCount (flip (-) 1)
writeMutVar slotMapFree keyIndex
let eraseItem x = x { slotItem = Nothing }
items <- readMutVar slotMapItems
V.unsafeModify items eraseItem keyIndex
insert :: PrimMonad m => a -> SlotMap m a -> m Key
insert item SlotMap{..} = do
free <- readMutVar slotMapFree
items <- readMutVar slotMapItems
potentialcurrent <- readMaybe items free
case potentialcurrent of
Nothing -> do
expanded <- V.unsafeGrow items 1
let slot = Slot { slotItem = Just item, slotGeneration = 1, slotFreeIndex = 0 }
V.unsafeWrite expanded (V.length expanded - 1) slot
writeMutVar slotMapFree (V.length expanded)
modifyMutVar slotMapCount (1 +)
writeMutVar slotMapItems expanded
return Key
{ keyIndex = free
, keyGeneration = 1
}
Just current -> do
let next = slotGeneration current + 1
let slot = current { slotItem = Just item, slotGeneration = next }
V.unsafeWrite items free current
writeMutVar slotMapFree (slotFreeIndex current)
modifyMutVar slotMapCount (1 +)
return Key
{ keyIndex = free
, keyGeneration = slotGeneration current + 1
}
update :: PrimMonad m => SlotMap m a -> (a -> a) -> Key -> m ()
update map operation key = do
potentialitem <- Data.SlotMap.lookup key map
case potentialitem of
Nothing -> pure ()
Just _ -> unsafeUpdate map operation key
unsafeUpdate :: PrimMonad m => SlotMap m a -> (a -> a) -> Key -> m ()
unsafeUpdate SlotMap{..} operation Key{..} = do
items <- readMutVar slotMapItems
V.unsafeModify items (fmap operation) keyIndex
null :: PrimMonad m => SlotMap m a -> m Bool
null map = Data.SlotMap.size map >>= pure . (==) 0
capacity :: PrimMonad m => SlotMap m a -> m Int
capacity SlotMap{..} = readMutVar slotMapItems >>= pure . V.length
size :: PrimMonad m => SlotMap m a -> m Int
size SlotMap{..} = readMutVar slotMapCount
foldr :: PrimMonad m => (a -> b -> b) -> b -> SlotMap m a -> m b
foldr op initial SlotMap{..} = readMutVar slotMapItems >>= foldrVector (foldOpMaybe op . slotItem) initial
map :: PrimMonad m => (a -> a) -> SlotMap m a -> m ()
map operation SlotMap{..} = readMutVar slotMapItems >>= mapVector (fmap operation)
elems :: PrimMonad m => SlotMap m a -> m [a]
elems SlotMap{..} = readMutVar slotMapItems >>= foldrVector (foldOpMaybe (:) . slotItem) []
foldOpMaybe :: (a -> b -> b) -> Maybe a -> b -> b
foldOpMaybe f Nothing x = x
foldOpMaybe f (Just x) y = f x y
readMaybe :: PrimMonad m => MVector (PrimState m) a -> Int -> m (Maybe a)
readMaybe vec index = if index >= V.length vec then pure Nothing else do
V.unsafeRead vec index >>= pure . Just
mapVector :: (PrimMonad m) => (a -> a) -> MVector (PrimState m) a -> m ()
mapVector operation vector = go 0 where
go iter = if iter == V.length vector then pure () else do
V.unsafeModify vector operation iter
go $ iter + 1
foldrVector :: PrimMonad m => (a -> b -> b) -> b -> MVector (PrimState m) a -> m b
foldrVector operation initial vector = go 0 where
go iter = if iter == V.length vector then pure initial else do
item <- V.unsafeRead vector iter
rest <- go $ iter + 1
pure $ operation item rest
instance Functor Slot where
fmap operation slot = slot { slotItem = fmap operation (slotItem slot) }