{-# LANGUAGE CPP #-}
{-# LANGUAGE PatternGuards #-}
{-# OPTIONS_GHC -funbox-strict-fields #-}

-----------------------------------------------------------------------
-- | A contention-free STM hash map.
-- \"Contention-free\" means that the map will never cause spurious conflicts.
-- A transaction operating on the map will only ever have to retry if
-- another transaction is operating on the same key at the same time.
-----------------------------------------------------------------------

module Control.Concurrent.STM.Map
    ( Map

      -- * Construction
    , empty

      -- * Modification
    , insert
    , delete
    , unsafeDelete

      -- * Query
    , lookup
    , phantomLookup
    , member

      -- * Lists
    , fromList
    , unsafeToList
    ) where

#if __GLASGOW_HASKELL__ < 710
import Control.Applicative ((<$>))
#endif
import Control.Concurrent.STM
import Control.Monad
import Data.Atomics
import Data.IORef
import Data.Maybe
import GHC.Conc.Sync (unsafeIOToSTM)
import Prelude hiding (lookup)

import Data.SparseArray

-----------------------------------------------------------------------

-- | A map from keys @k@ to values @v@.
newtype Map k v = Map (INode k v)

type INode k v = IORef (Node k v)

data Node k v = Array !(SparseArray (Branch k v))
              | List  ![Leaf k v]
              | Tomb  !(Leaf k v)

data Branch k v = I !(INode k v)
                | L !(Leaf k v)

data Leaf k v = Leaf !k !(TVar (Maybe v))

-----------------------------------------------------------------------

-- | /O(1)/. Construct an empty map.
empty :: STM (Map k v)
empty :: STM (Map k v)
empty = IO (Map k v) -> STM (Map k v)
forall a. IO a -> STM a
unsafeIOToSTM (IO (Map k v) -> STM (Map k v)) -> IO (Map k v) -> STM (Map k v)
forall a b. (a -> b) -> a -> b
$ INode k v -> Map k v
forall k v. INode k v -> Map k v
Map (INode k v -> Map k v) -> IO (INode k v) -> IO (Map k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Node k v -> IO (INode k v)
forall a. a -> IO (IORef a)
newIORef (SparseArray (Branch k v) -> Node k v
forall k v. SparseArray (Branch k v) -> Node k v
Array SparseArray (Branch k v)
forall a. SparseArray a
emptyArray)
{-# INLINE empty #-}

-- | /O(log n)/. Associate the given value with the given key.
-- If the key is already present in the map, the old value is replaced.
insert :: (Eq k, Hashable k) => k -> v -> Map k v -> STM ()
insert :: k -> v -> Map k v -> STM ()
insert k
k v
v Map k v
m = do TVar (Maybe v)
var <- k -> Map k v -> STM (TVar (Maybe v))
forall k v.
(Eq k, Hashable k) =>
k -> Map k v -> STM (TVar (Maybe v))
getTVar k
k Map k v
m
                  TVar (Maybe v) -> Maybe v -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Maybe v)
var (v -> Maybe v
forall a. a -> Maybe a
Just v
v)
{-# INLINABLE insert #-}

-- | /O(log n)/. Return the value associated with the given key, or 'Nothing'.
--
-- __Note__: This might increase the map's memory consumption
-- by putting the key into the map.
-- If that is not acceptable, use 'phantomLookup'.
lookup :: (Eq k, Hashable k) => k -> Map k v -> STM (Maybe v)
lookup :: k -> Map k v -> STM (Maybe v)
lookup k
k Map k v
m = do TVar (Maybe v)
var <- k -> Map k v -> STM (TVar (Maybe v))
forall k v.
(Eq k, Hashable k) =>
k -> Map k v -> STM (TVar (Maybe v))
getTVar k
k Map k v
m
                TVar (Maybe v) -> STM (Maybe v)
forall a. TVar a -> STM a
readTVar TVar (Maybe v)
var
{-# INLINABLE lookup #-}

-- | /O(log n)/. Remove the value associated with a given key from the map,
-- if present.
--
-- __Note__: This does not actually remove the key from the map.
-- In fact, it might actually increase the map's memory consumption
-- by putting the key into the map.
-- To completely delete an entry, including its key, use 'unsafeDelete'.
delete :: (Eq k, Hashable k) => k -> Map k v -> STM ()
delete :: k -> Map k v -> STM ()
delete k
k Map k v
m = do TVar (Maybe v)
var <- k -> Map k v -> STM (TVar (Maybe v))
forall k v.
(Eq k, Hashable k) =>
k -> Map k v -> STM (TVar (Maybe v))
getTVar k
k Map k v
m
                TVar (Maybe v) -> Maybe v -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Maybe v)
var Maybe v
forall a. Maybe a
Nothing
{-# INLINABLE delete #-}

-----------------------------------------------------------------------

-- |/O(log n)/. Is the key a member of the map?
member :: (Eq k, Hashable k) => k -> Map k v -> STM Bool
member :: k -> Map k v -> STM Bool
member k
k Map k v
m = do
    Maybe v
v <- k -> Map k v -> STM (Maybe v)
forall k v. (Eq k, Hashable k) => k -> Map k v -> STM (Maybe v)
lookup k
k Map k v
m
    case Maybe v
v of
        Maybe v
Nothing -> Bool -> STM Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
        Just v
_  -> Bool -> STM Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True

-----------------------------------------------------------------------

getTVar :: (Eq k, Hashable k) => k -> Map k v -> STM (TVar (Maybe v))
getTVar :: k -> Map k v -> STM (TVar (Maybe v))
getTVar k
k (Map INode k v
root) = INode k v -> Level -> INode k v -> STM (TVar (Maybe v))
go INode k v
root Level
0 INode k v
forall a. HasCallStack => a
undefined
  where
    h :: Hash
h = k -> Hash
forall a. Hashable a => a -> Hash
hash k
k

    go :: INode k v -> Level -> INode k v -> STM (TVar (Maybe v))
go INode k v
inode Level
level INode k v
parent = do
        Ticket (Node k v)
ticket <- IO (Ticket (Node k v)) -> STM (Ticket (Node k v))
forall a. IO a -> STM a
unsafeIOToSTM (IO (Ticket (Node k v)) -> STM (Ticket (Node k v)))
-> IO (Ticket (Node k v)) -> STM (Ticket (Node k v))
forall a b. (a -> b) -> a -> b
$ INode k v -> IO (Ticket (Node k v))
forall a. IORef a -> IO (Ticket a)
readForCAS INode k v
inode
        case Ticket (Node k v) -> Node k v
forall a. Ticket a -> a
peekTicket Ticket (Node k v)
ticket of
            Array SparseArray (Branch k v)
a -> case Level -> Hash -> SparseArray (Branch k v) -> Maybe (Branch k v)
forall a. Level -> Hash -> SparseArray a -> Maybe a
arrayLookup Level
level Hash
h SparseArray (Branch k v)
a of
                Just (I INode k v
inode2) -> INode k v -> Level -> INode k v -> STM (TVar (Maybe v))
go INode k v
inode2 (Level -> Level
down Level
level) INode k v
inode
                Just (L leaf2 :: Leaf k v
leaf2@(Leaf k
k2 TVar (Maybe v)
var))
                    | k
k k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
k2   -> TVar (Maybe v) -> STM (TVar (Maybe v))
forall (m :: * -> *) a. Monad m => a -> m a
return TVar (Maybe v)
var
                    | Bool
otherwise -> INode k v
-> Ticket (Node k v)
-> (Leaf k v -> STM (Node k v))
-> STM (TVar (Maybe v))
cas INode k v
inode Ticket (Node k v)
ticket (Level
-> SparseArray (Branch k v)
-> Hash
-> Leaf k v
-> Leaf k v
-> STM (Node k v)
forall k v.
Level
-> SparseArray (Branch k v)
-> Hash
-> Leaf k v
-> Leaf k v
-> STM (Node k v)
growTrie Level
level SparseArray (Branch k v)
a (k -> Hash
forall a. Hashable a => a -> Hash
hash k
k2) Leaf k v
leaf2)
                Maybe (Branch k v)
Nothing         -> INode k v
-> Ticket (Node k v)
-> (Leaf k v -> STM (Node k v))
-> STM (TVar (Maybe v))
cas INode k v
inode Ticket (Node k v)
ticket (Level -> SparseArray (Branch k v) -> Leaf k v -> STM (Node k v)
forall (m :: * -> *) k v.
Monad m =>
Level -> SparseArray (Branch k v) -> Leaf k v -> m (Node k v)
insertLeaf Level
level SparseArray (Branch k v)
a)
            List [Leaf k v]
xs -> case k -> [Leaf k v] -> Maybe (TVar (Maybe v))
forall k v. Eq k => k -> [Leaf k v] -> Maybe (TVar (Maybe v))
listLookup k
k [Leaf k v]
xs of
                Just TVar (Maybe v)
var        -> TVar (Maybe v) -> STM (TVar (Maybe v))
forall (m :: * -> *) a. Monad m => a -> m a
return TVar (Maybe v)
var
                Maybe (TVar (Maybe v))
Nothing         -> INode k v
-> Ticket (Node k v)
-> (Leaf k v -> STM (Node k v))
-> STM (TVar (Maybe v))
cas INode k v
inode Ticket (Node k v)
ticket (Node k v -> STM (Node k v)
forall (m :: * -> *) a. Monad m => a -> m a
return (Node k v -> STM (Node k v))
-> (Leaf k v -> Node k v) -> Leaf k v -> STM (Node k v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Leaf k v] -> Node k v
forall k v. [Leaf k v] -> Node k v
List ([Leaf k v] -> Node k v)
-> (Leaf k v -> [Leaf k v]) -> Leaf k v -> Node k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Leaf k v -> [Leaf k v] -> [Leaf k v]
forall a. a -> [a] -> [a]
:[Leaf k v]
xs))
            Tomb Leaf k v
_  -> IO () -> STM ()
forall a. IO a -> STM a
unsafeIOToSTM (INode k v -> Level -> IO ()
forall k v. INode k v -> Level -> IO ()
clean INode k v
parent (Level -> Level
up Level
level)) STM () -> STM (TVar (Maybe v)) -> STM (TVar (Maybe v))
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> INode k v -> Level -> INode k v -> STM (TVar (Maybe v))
go INode k v
root Level
0 INode k v
forall a. HasCallStack => a
undefined

    cas :: INode k v
-> Ticket (Node k v)
-> (Leaf k v -> STM (Node k v))
-> STM (TVar (Maybe v))
cas INode k v
inode Ticket (Node k v)
ticket Leaf k v -> STM (Node k v)
f = do
        TVar (Maybe v)
var <- Maybe v -> STM (TVar (Maybe v))
forall a. a -> STM (TVar a)
newTVar Maybe v
forall a. Maybe a
Nothing
        Node k v
node <- Leaf k v -> STM (Node k v)
f (k -> TVar (Maybe v) -> Leaf k v
forall k v. k -> TVar (Maybe v) -> Leaf k v
Leaf k
k TVar (Maybe v)
var)
        (Bool
ok,Ticket (Node k v)
_) <- IO (Bool, Ticket (Node k v)) -> STM (Bool, Ticket (Node k v))
forall a. IO a -> STM a
unsafeIOToSTM (IO (Bool, Ticket (Node k v)) -> STM (Bool, Ticket (Node k v)))
-> IO (Bool, Ticket (Node k v)) -> STM (Bool, Ticket (Node k v))
forall a b. (a -> b) -> a -> b
$ INode k v
-> Ticket (Node k v) -> Node k v -> IO (Bool, Ticket (Node k v))
forall a. IORef a -> Ticket a -> a -> IO (Bool, Ticket a)
casIORef INode k v
inode Ticket (Node k v)
ticket Node k v
node
        if Bool
ok then TVar (Maybe v) -> STM (TVar (Maybe v))
forall (m :: * -> *) a. Monad m => a -> m a
return TVar (Maybe v)
var
              else INode k v -> Level -> INode k v -> STM (TVar (Maybe v))
go INode k v
root Level
0 INode k v
forall a. HasCallStack => a
undefined

    insertLeaf :: Level -> SparseArray (Branch k v) -> Leaf k v -> m (Node k v)
insertLeaf Level
level SparseArray (Branch k v)
a Leaf k v
leaf = do
        let a' :: SparseArray (Branch k v)
a' = Level
-> Hash
-> Branch k v
-> SparseArray (Branch k v)
-> SparseArray (Branch k v)
forall a. Level -> Hash -> a -> SparseArray a -> SparseArray a
arrayInsert Level
level Hash
h (Leaf k v -> Branch k v
forall k v. Leaf k v -> Branch k v
L Leaf k v
leaf) SparseArray (Branch k v)
a
        Node k v -> m (Node k v)
forall (m :: * -> *) a. Monad m => a -> m a
return (SparseArray (Branch k v) -> Node k v
forall k v. SparseArray (Branch k v) -> Node k v
Array SparseArray (Branch k v)
a')

    growTrie :: Level
-> SparseArray (Branch k v)
-> Hash
-> Leaf k v
-> Leaf k v
-> STM (Node k v)
growTrie Level
level SparseArray (Branch k v)
a Hash
h2 Leaf k v
leaf2 Leaf k v
leaf1 = do
        IORef (Node k v)
inode2 <- IO (IORef (Node k v)) -> STM (IORef (Node k v))
forall a. IO a -> STM a
unsafeIOToSTM (IO (IORef (Node k v)) -> STM (IORef (Node k v)))
-> IO (IORef (Node k v)) -> STM (IORef (Node k v))
forall a b. (a -> b) -> a -> b
$ Level
-> Hash -> Leaf k v -> Hash -> Leaf k v -> IO (IORef (Node k v))
forall t k v.
Level -> t -> Leaf k v -> Hash -> Leaf k v -> IO (IORef (Node k v))
combineLeaves (Level -> Level
down Level
level) Hash
h Leaf k v
leaf1 Hash
h2 Leaf k v
leaf2
        let a' :: SparseArray (Branch k v)
a' = Level
-> Hash
-> Branch k v
-> SparseArray (Branch k v)
-> SparseArray (Branch k v)
forall a. Level -> Hash -> a -> SparseArray a -> SparseArray a
arrayUpdate Level
level Hash
h (IORef (Node k v) -> Branch k v
forall k v. INode k v -> Branch k v
I IORef (Node k v)
inode2) SparseArray (Branch k v)
a
        Node k v -> STM (Node k v)
forall (m :: * -> *) a. Monad m => a -> m a
return (SparseArray (Branch k v) -> Node k v
forall k v. SparseArray (Branch k v) -> Node k v
Array SparseArray (Branch k v)
a')

    combineLeaves :: Level -> t -> Leaf k v -> Hash -> Leaf k v -> IO (IORef (Node k v))
combineLeaves Level
level t
h1 Leaf k v
leaf1 Hash
h2 Leaf k v
leaf2
        | Level
level Level -> Level -> Bool
forall a. Ord a => a -> a -> Bool
>= Level
lastLevel = Node k v -> IO (IORef (Node k v))
forall a. a -> IO (IORef a)
newIORef ([Leaf k v] -> Node k v
forall k v. [Leaf k v] -> Node k v
List [Leaf k v
leaf1, Leaf k v
leaf2])
        | Bool
otherwise =
            case Level
-> Hash
-> Branch k v
-> Hash
-> Branch k v
-> Maybe (SparseArray (Branch k v))
forall a. Level -> Hash -> a -> Hash -> a -> Maybe (SparseArray a)
mkPair Level
level Hash
h (Leaf k v -> Branch k v
forall k v. Leaf k v -> Branch k v
L Leaf k v
leaf1) Hash
h2 (Leaf k v -> Branch k v
forall k v. Leaf k v -> Branch k v
L Leaf k v
leaf2) of
                Just SparseArray (Branch k v)
pair -> Node k v -> IO (IORef (Node k v))
forall a. a -> IO (IORef a)
newIORef (SparseArray (Branch k v) -> Node k v
forall k v. SparseArray (Branch k v) -> Node k v
Array SparseArray (Branch k v)
pair)
                Maybe (SparseArray (Branch k v))
Nothing -> do
                    IORef (Node k v)
inode <- Level -> t -> Leaf k v -> Hash -> Leaf k v -> IO (IORef (Node k v))
combineLeaves (Level -> Level
down Level
level) t
h1 Leaf k v
leaf1 Hash
h2 Leaf k v
leaf2
                    let a :: SparseArray (Branch k v)
a = Level -> Hash -> Branch k v -> SparseArray (Branch k v)
forall a. Level -> Hash -> a -> SparseArray a
mkSingleton Level
level Hash
h (IORef (Node k v) -> Branch k v
forall k v. INode k v -> Branch k v
I IORef (Node k v)
inode)
                    Node k v -> IO (IORef (Node k v))
forall a. a -> IO (IORef a)
newIORef (SparseArray (Branch k v) -> Node k v
forall k v. SparseArray (Branch k v) -> Node k v
Array SparseArray (Branch k v)
a)

{-# INLINE getTVar #-}


-- | /O(log n)/. Return the value associated with the given key, or 'Nothing'.
--
-- In contrast to 'lookup', this will never increase the map's memory consumption.
-- However, it might allow /phantom reads/ to occur.
-- Consider the following situation:
--
-- > f = atomically $ do v1 <- phantomLookup k m
-- >                     v2 <- phantomLookup k m
-- >                     return (v1 == v2)
--
-- Under certain circumstances @f@ might actually return @False@, in particular
-- if the first @phantomLookup@ happens on an empty map
-- and some other transaction inserts a value for @k@ before the second call
-- to @phantomLookup@.
phantomLookup :: (Eq k, Hashable k) => k -> Map k v -> STM (Maybe v)
phantomLookup :: k -> Map k v -> STM (Maybe v)
phantomLookup k
k (Map INode k v
root) = INode k v -> Level -> INode k v -> STM (Maybe v)
go INode k v
root Level
0 INode k v
forall a. HasCallStack => a
undefined
  where
    h :: Hash
h = k -> Hash
forall a. Hashable a => a -> Hash
hash k
k
    go :: INode k v -> Level -> INode k v -> STM (Maybe v)
go INode k v
inode Level
level INode k v
parent = do
        Node k v
node <- IO (Node k v) -> STM (Node k v)
forall a. IO a -> STM a
unsafeIOToSTM (IO (Node k v) -> STM (Node k v))
-> IO (Node k v) -> STM (Node k v)
forall a b. (a -> b) -> a -> b
$ INode k v -> IO (Node k v)
forall a. IORef a -> IO a
readIORef INode k v
inode
        case Node k v
node of
            Array SparseArray (Branch k v)
a -> case Level -> Hash -> SparseArray (Branch k v) -> Maybe (Branch k v)
forall a. Level -> Hash -> SparseArray a -> Maybe a
arrayLookup Level
level Hash
h SparseArray (Branch k v)
a of
                Just (I INode k v
inode2) -> INode k v -> Level -> INode k v -> STM (Maybe v)
go INode k v
inode2 (Level -> Level
down Level
level) INode k v
inode
                Just (L (Leaf k
k2 TVar (Maybe v)
var))
                    | k
k k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
k2   -> TVar (Maybe v) -> STM (Maybe v)
forall a. TVar a -> STM a
readTVar TVar (Maybe v)
var
                    | Bool
otherwise -> Maybe v -> STM (Maybe v)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe v
forall a. Maybe a
Nothing
                Maybe (Branch k v)
Nothing         -> Maybe v -> STM (Maybe v)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe v
forall a. Maybe a
Nothing
            List [Leaf k v]
xs -> case k -> [Leaf k v] -> Maybe (TVar (Maybe v))
forall k v. Eq k => k -> [Leaf k v] -> Maybe (TVar (Maybe v))
listLookup k
k [Leaf k v]
xs of
                Just TVar (Maybe v)
var -> TVar (Maybe v) -> STM (Maybe v)
forall a. TVar a -> STM a
readTVar TVar (Maybe v)
var
                Maybe (TVar (Maybe v))
Nothing  -> Maybe v -> STM (Maybe v)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe v
forall a. Maybe a
Nothing
            Tomb Leaf k v
_  -> IO () -> STM ()
forall a. IO a -> STM a
unsafeIOToSTM (INode k v -> Level -> IO ()
forall k v. INode k v -> Level -> IO ()
clean INode k v
parent (Level -> Level
up Level
level)) STM () -> STM (Maybe v) -> STM (Maybe v)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> INode k v -> Level -> INode k v -> STM (Maybe v)
go INode k v
root Level
0 INode k v
forall a. HasCallStack => a
undefined

{-# INLINABLE phantomLookup #-}


-- | /O(log n)/. This will completely remove a given key
-- and its associated value from the map, if present.
-- This is not an atomic operation, however. __Use with caution!__
unsafeDelete :: (Eq k, Hashable k) => k -> Map k v -> IO ()
unsafeDelete :: k -> Map k v -> IO ()
unsafeDelete k
k m :: Map k v
m@(Map INode k v
root) = do
    Bool
ok <- INode k v -> Level -> INode k v -> IO Bool
forall v. INode k v -> Level -> INode k v -> IO Bool
go INode k v
root Level
0 INode k v
forall a. HasCallStack => a
undefined
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
ok (k -> Map k v -> IO ()
forall k v. (Eq k, Hashable k) => k -> Map k v -> IO ()
unsafeDelete k
k Map k v
m)
  where
    h :: Hash
h = k -> Hash
forall a. Hashable a => a -> Hash
hash k
k
    go :: INode k v -> Level -> INode k v -> IO Bool
go INode k v
inode Level
level INode k v
parent = do
        Ticket (Node k v)
ticket <- INode k v -> IO (Ticket (Node k v))
forall a. IORef a -> IO (Ticket a)
readForCAS INode k v
inode
        case Ticket (Node k v) -> Node k v
forall a. Ticket a -> a
peekTicket Ticket (Node k v)
ticket of
            Array SparseArray (Branch k v)
a -> do
                Bool
ok <- case Level -> Hash -> SparseArray (Branch k v) -> Maybe (Branch k v)
forall a. Level -> Hash -> SparseArray a -> Maybe a
arrayLookup Level
level Hash
h SparseArray (Branch k v)
a of
                    Just (I INode k v
inode2) -> INode k v -> Level -> INode k v -> IO Bool
go INode k v
inode2 (Level -> Level
down Level
level) INode k v
inode
                    Just (L (Leaf k
k2 TVar (Maybe v)
_))
                        | k
k k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
k2   -> INode k v
-> Ticket (Node k v)
-> Level
-> SparseArray (Branch k v)
-> IO Bool
forall k v.
IORef (Node k v)
-> Ticket (Node k v)
-> Level
-> SparseArray (Branch k v)
-> IO Bool
casArrayDelete INode k v
inode Ticket (Node k v)
ticket Level
level SparseArray (Branch k v)
a
                        | Bool
otherwise -> Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
                    Maybe (Branch k v)
Nothing         -> Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
                Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
ok (Level -> INode k v -> INode k v -> IO ()
forall k v. Level -> IORef (Node k v) -> IORef (Node k v) -> IO ()
compressIfPossible Level
level INode k v
inode INode k v
parent)
                Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
ok
            List [Leaf k v]
xs -> INode k v -> Ticket (Node k v) -> [Leaf k v] -> IO Bool
forall v.
IORef (Node k v) -> Ticket (Node k v) -> [Leaf k v] -> IO Bool
casListDelete INode k v
inode Ticket (Node k v)
ticket [Leaf k v]
xs
            Tomb Leaf k v
_  -> INode k v -> Level -> IO ()
forall k v. INode k v -> Level -> IO ()
clean INode k v
parent (Level -> Level
up Level
level) IO () -> IO Bool -> IO Bool
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

    compressIfPossible :: Level -> IORef (Node k v) -> IORef (Node k v) -> IO ()
compressIfPossible Level
level IORef (Node k v)
inode IORef (Node k v)
parent = do
        Node k v
n <- IORef (Node k v) -> IO (Node k v)
forall a. IORef a -> IO a
readIORef IORef (Node k v)
inode
        case Node k v
n of
            Tomb Leaf k v
_ -> IORef (Node k v) -> IORef (Node k v) -> Hash -> Level -> IO ()
forall k v. INode k v -> INode k v -> Hash -> Level -> IO ()
cleanParent IORef (Node k v)
parent IORef (Node k v)
inode Hash
h (Level -> Level
up Level
level)
            Node k v
_      -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

    casArrayDelete :: IORef (Node k v)
-> Ticket (Node k v)
-> Level
-> SparseArray (Branch k v)
-> IO Bool
casArrayDelete IORef (Node k v)
inode Ticket (Node k v)
ticket Level
level SparseArray (Branch k v)
a = do
        let a' :: SparseArray (Branch k v)
a' = Level
-> Hash -> SparseArray (Branch k v) -> SparseArray (Branch k v)
forall a. Level -> Hash -> SparseArray a -> SparseArray a
arrayDelete Level
level Hash
h SparseArray (Branch k v)
a
            n :: Node k v
n  = Level -> Node k v -> Node k v
forall k v. Level -> Node k v -> Node k v
contract Level
level (SparseArray (Branch k v) -> Node k v
forall k v. SparseArray (Branch k v) -> Node k v
Array SparseArray (Branch k v)
a')
        (Bool
ok,Ticket (Node k v)
_) <- IORef (Node k v)
-> Ticket (Node k v) -> Node k v -> IO (Bool, Ticket (Node k v))
forall a. IORef a -> Ticket a -> a -> IO (Bool, Ticket a)
casIORef IORef (Node k v)
inode Ticket (Node k v)
ticket Node k v
n
        Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
ok

    casListDelete :: IORef (Node k v) -> Ticket (Node k v) -> [Leaf k v] -> IO Bool
casListDelete IORef (Node k v)
inode Ticket (Node k v)
ticket [Leaf k v]
xs = do
        let xs' :: [Leaf k v]
xs' = k -> [Leaf k v] -> [Leaf k v]
forall k v. Eq k => k -> [Leaf k v] -> [Leaf k v]
listDelete k
k [Leaf k v]
xs
            n :: Node k v
n | [Leaf k v
l] <- [Leaf k v]
xs' = Leaf k v -> Node k v
forall k v. Leaf k v -> Node k v
Tomb Leaf k v
l
              | Bool
otherwise  = [Leaf k v] -> Node k v
forall k v. [Leaf k v] -> Node k v
List [Leaf k v]
xs'
        (Bool
ok,Ticket (Node k v)
_) <- IORef (Node k v)
-> Ticket (Node k v) -> Node k v -> IO (Bool, Ticket (Node k v))
forall a. IORef a -> Ticket a -> a -> IO (Bool, Ticket a)
casIORef IORef (Node k v)
inode Ticket (Node k v)
ticket Node k v
n
        Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
ok

{-# INLINABLE unsafeDelete #-}

-----------------------------------------------------------------------

clean :: INode k v -> Level -> IO ()
clean :: INode k v -> Level -> IO ()
clean INode k v
inode Level
level = do
    Ticket (Node k v)
ticket <- INode k v -> IO (Ticket (Node k v))
forall a. IORef a -> IO (Ticket a)
readForCAS INode k v
inode
    case Ticket (Node k v) -> Node k v
forall a. Ticket a -> a
peekTicket Ticket (Node k v)
ticket of
        n :: Node k v
n@(Array SparseArray (Branch k v)
_) -> do
            Node k v
n' <- Level -> Node k v -> IO (Node k v)
forall k v. Level -> Node k v -> IO (Node k v)
compress Level
level Node k v
n
            IO (Bool, Ticket (Node k v)) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Bool, Ticket (Node k v)) -> IO ())
-> IO (Bool, Ticket (Node k v)) -> IO ()
forall a b. (a -> b) -> a -> b
$ INode k v
-> Ticket (Node k v) -> Node k v -> IO (Bool, Ticket (Node k v))
forall a. IORef a -> Ticket a -> a -> IO (Bool, Ticket a)
casIORef INode k v
inode Ticket (Node k v)
ticket Node k v
n'
        Node k v
_ -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
{-# INLINE clean #-}

cleanParent :: INode k v -> INode k v -> Hash -> Level -> IO ()
cleanParent :: INode k v -> INode k v -> Hash -> Level -> IO ()
cleanParent INode k v
parent INode k v
inode Hash
h Level
level = do
    Ticket (Node k v)
ticket <- INode k v -> IO (Ticket (Node k v))
forall a. IORef a -> IO (Ticket a)
readForCAS INode k v
parent
    case Ticket (Node k v) -> Node k v
forall a. Ticket a -> a
peekTicket Ticket (Node k v)
ticket of
        n :: Node k v
n@(Array SparseArray (Branch k v)
a) -> case Level -> Hash -> SparseArray (Branch k v) -> Maybe (Branch k v)
forall a. Level -> Hash -> SparseArray a -> Maybe a
arrayLookup Level
level Hash
h SparseArray (Branch k v)
a of
            Just (I INode k v
inode2) | INode k v
inode2 INode k v -> INode k v -> Bool
forall a. Eq a => a -> a -> Bool
== INode k v
inode -> do
                Node k v
n2 <- INode k v -> IO (Node k v)
forall a. IORef a -> IO a
readIORef INode k v
inode
                case Node k v
n2 of
                    Tomb Leaf k v
_ -> do
                        Node k v
n' <- Level -> Node k v -> IO (Node k v)
forall k v. Level -> Node k v -> IO (Node k v)
compress Level
level Node k v
n
                        (Bool
ok,Ticket (Node k v)
_) <- INode k v
-> Ticket (Node k v) -> Node k v -> IO (Bool, Ticket (Node k v))
forall a. IORef a -> Ticket a -> a -> IO (Bool, Ticket a)
casIORef INode k v
parent Ticket (Node k v)
ticket Node k v
n'
                        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
ok (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ INode k v -> INode k v -> Hash -> Level -> IO ()
forall k v. INode k v -> INode k v -> Hash -> Level -> IO ()
cleanParent INode k v
parent INode k v
inode Hash
h Level
level
                    Node k v
_ -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            Maybe (Branch k v)
_ -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        Node k v
_ -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

compress :: Level -> Node k v -> IO (Node k v)
compress :: Level -> Node k v -> IO (Node k v)
compress Level
level (Array SparseArray (Branch k v)
a) = Level -> Node k v -> Node k v
forall k v. Level -> Node k v -> Node k v
contract Level
level (Node k v -> Node k v)
-> (SparseArray (Branch k v) -> Node k v)
-> SparseArray (Branch k v)
-> Node k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SparseArray (Branch k v) -> Node k v
forall k v. SparseArray (Branch k v) -> Node k v
Array (SparseArray (Branch k v) -> Node k v)
-> IO (SparseArray (Branch k v)) -> IO (Node k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Branch k v -> IO (Branch k v))
-> SparseArray (Branch k v) -> IO (SparseArray (Branch k v))
forall a. (a -> IO a) -> SparseArray a -> IO (SparseArray a)
arrayMapM Branch k v -> IO (Branch k v)
forall k v. Branch k v -> IO (Branch k v)
resurrect SparseArray (Branch k v)
a
compress Level
_     Node k v
n         = Node k v -> IO (Node k v)
forall (m :: * -> *) a. Monad m => a -> m a
return Node k v
n
{-# INLINE compress #-}

resurrect :: Branch k v -> IO (Branch k v)
resurrect :: Branch k v -> IO (Branch k v)
resurrect b :: Branch k v
b@(I INode k v
inode) = do Node k v
n <- INode k v -> IO (Node k v)
forall a. IORef a -> IO a
readIORef INode k v
inode
                           case Node k v
n of
                               Tomb Leaf k v
leaf -> Branch k v -> IO (Branch k v)
forall (m :: * -> *) a. Monad m => a -> m a
return (Leaf k v -> Branch k v
forall k v. Leaf k v -> Branch k v
L Leaf k v
leaf)
                               Node k v
_         -> Branch k v -> IO (Branch k v)
forall (m :: * -> *) a. Monad m => a -> m a
return Branch k v
b
resurrect Branch k v
b           = Branch k v -> IO (Branch k v)
forall (m :: * -> *) a. Monad m => a -> m a
return Branch k v
b
{-# INLINE resurrect #-}

contract :: Level -> Node k v -> Node k v
contract :: Level -> Node k v -> Node k v
contract Level
level (Array SparseArray (Branch k v)
a) | Level
level Level -> Level -> Bool
forall a. Ord a => a -> a -> Bool
> Level
0
                         , Just (L Leaf k v
leaf) <- SparseArray (Branch k v) -> Maybe (Branch k v)
forall a. SparseArray a -> Maybe a
arrayToMaybe SparseArray (Branch k v)
a
                         = Leaf k v -> Node k v
forall k v. Leaf k v -> Node k v
Tomb Leaf k v
leaf
contract Level
_     Node k v
n         = Node k v
n
{-# INLINE contract #-}

-----------------------------------------------------------------------

listLookup :: Eq k => k -> [Leaf k v] -> Maybe (TVar (Maybe v))
listLookup :: k -> [Leaf k v] -> Maybe (TVar (Maybe v))
listLookup k
k1 = [Leaf k v] -> Maybe (TVar (Maybe v))
forall v. [Leaf k v] -> Maybe (TVar (Maybe v))
go
  where
    go :: [Leaf k v] -> Maybe (TVar (Maybe v))
go []                             = Maybe (TVar (Maybe v))
forall a. Maybe a
Nothing
    go (Leaf k
k2 TVar (Maybe v)
var : [Leaf k v]
xs) | k
k1 k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
k2  = TVar (Maybe v) -> Maybe (TVar (Maybe v))
forall a. a -> Maybe a
Just TVar (Maybe v)
var
                          | Bool
otherwise = [Leaf k v] -> Maybe (TVar (Maybe v))
go [Leaf k v]
xs

listDelete :: Eq k => k -> [Leaf k v] -> [Leaf k v]
listDelete :: k -> [Leaf k v] -> [Leaf k v]
listDelete k
k1 = [Leaf k v] -> [Leaf k v]
forall v. [Leaf k v] -> [Leaf k v]
go
  where
    go :: [Leaf k v] -> [Leaf k v]
go [] = []
    go (x :: Leaf k v
x@(Leaf k
k2 TVar (Maybe v)
_):[Leaf k v]
xs) | k
k1 k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
k2  = [Leaf k v]
xs
                          | Bool
otherwise = Leaf k v
x Leaf k v -> [Leaf k v] -> [Leaf k v]
forall a. a -> [a] -> [a]
: [Leaf k v] -> [Leaf k v]
go [Leaf k v]
xs

-----------------------------------------------------------------------

-- | /O(n * log n)/. Construct a map from a list of key/value pairs.
fromList :: (Eq k, Hashable k) => [(k,v)] -> IO (Map k v)
fromList :: [(k, v)] -> IO (Map k v)
fromList [(k, v)]
xs = do
    Map k v
m <- STM (Map k v) -> IO (Map k v)
forall a. STM a -> IO a
atomically STM (Map k v)
forall k v. STM (Map k v)
empty
    [(k, v)] -> ((k, v) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(k, v)]
xs (((k, v) -> IO ()) -> IO ()) -> ((k, v) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(k
k,v
v) -> STM () -> IO ()
forall a. STM a -> IO a
atomically (k -> v -> Map k v -> STM ()
forall k v. (Eq k, Hashable k) => k -> v -> Map k v -> STM ()
insert k
k v
v Map k v
m)
    Map k v -> IO (Map k v)
forall (m :: * -> *) a. Monad m => a -> m a
return Map k v
m

-- | /O(n)/. Unsafely convert the map to a list of key/value pairs.
--
-- __Warning__: 'unsafeToList' makes no atomicity guarantees. Concurrent
-- changes to the map will lead to inconsistent results.
unsafeToList :: Map k v -> IO [(k,v)]
unsafeToList :: Map k v -> IO [(k, v)]
unsafeToList (Map INode k v
root) = INode k v -> IO [(k, v)]
forall a b. IORef (Node a b) -> IO [(a, b)]
go INode k v
root
  where
    go :: IORef (Node a b) -> IO [(a, b)]
go IORef (Node a b)
inode = do
        Node a b
node <- IORef (Node a b) -> IO (Node a b)
forall a. IORef a -> IO a
readIORef IORef (Node a b)
inode
        case Node a b
node of
            Array SparseArray (Branch a b)
a -> ([(a, b)] -> Branch a b -> IO [(a, b)])
-> [(a, b)] -> SparseArray (Branch a b) -> IO [(a, b)]
forall b a. (b -> a -> IO b) -> b -> SparseArray a -> IO b
arrayFoldM' [(a, b)] -> Branch a b -> IO [(a, b)]
go2 [] SparseArray (Branch a b)
a
            List [Leaf a b]
xs -> ([(a, b)] -> Leaf a b -> IO [(a, b)])
-> [(a, b)] -> [Leaf a b] -> IO [(a, b)]
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM [(a, b)] -> Leaf a b -> IO [(a, b)]
forall a b. [(a, b)] -> Leaf a b -> IO [(a, b)]
go3 [] [Leaf a b]
xs
            Tomb Leaf a b
leaf -> [(a, b)] -> Leaf a b -> IO [(a, b)]
forall a b. [(a, b)] -> Leaf a b -> IO [(a, b)]
go3 [] Leaf a b
leaf

    go2 :: [(a, b)] -> Branch a b -> IO [(a, b)]
go2 [(a, b)]
xs (I IORef (Node a b)
inode) = IORef (Node a b) -> IO [(a, b)]
go IORef (Node a b)
inode IO [(a, b)] -> ([(a, b)] -> IO [(a, b)]) -> IO [(a, b)]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \[(a, b)]
ys -> [(a, b)] -> IO [(a, b)]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(a, b)]
ys [(a, b)] -> [(a, b)] -> [(a, b)]
forall a. [a] -> [a] -> [a]
++ [(a, b)]
xs)
    go2 [(a, b)]
xs (L Leaf a b
leaf) = [(a, b)] -> Leaf a b -> IO [(a, b)]
forall a b. [(a, b)] -> Leaf a b -> IO [(a, b)]
go3 [(a, b)]
xs Leaf a b
leaf

    go3 :: [(a, b)] -> Leaf a b -> IO [(a, b)]
go3 [(a, b)]
xs (Leaf a
k TVar (Maybe b)
var) = do
        Maybe b
v <- TVar (Maybe b) -> IO (Maybe b)
forall a. TVar a -> IO a
readTVarIO TVar (Maybe b)
var
        case Maybe b
v of
            Maybe b
Nothing -> [(a, b)] -> IO [(a, b)]
forall (m :: * -> *) a. Monad m => a -> m a
return [(a, b)]
xs
            Just b
v' -> [(a, b)] -> IO [(a, b)]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(a, b)] -> IO [(a, b)]) -> [(a, b)] -> IO [(a, b)]
forall a b. (a -> b) -> a -> b
$ (a
k,b
v') (a, b) -> [(a, b)] -> [(a, b)]
forall a. a -> [a] -> [a]
: [(a, b)]
xs