-- | This module implements per-thread variables
module Util.ThreadDict(
   ThreadDict, -- contains all the thread variables
   newThreadDict, -- :: IO (ThreadDict a)
   writeThreadDict, -- :: ThreadDict a -> a -> IO ()
   readThreadDict, -- :: ThreadDict a -> IO (Maybe a)
   modifyThreadDict, -- :: ThreadDict a -> (Maybe a -> IO (Maybe a,b)) -> IO b
   ) where

import Control.Concurrent

import qualified Data.Map as Map
import Data.IORef

-- -------------------------------------------------------------------------
-- Data types
-- -------------------------------------------------------------------------

newtype ThreadDict a = ThreadDict (IORef (Map.Map ThreadId a))

-- -------------------------------------------------------------------------
-- Functions
-- -------------------------------------------------------------------------

newThreadDict :: IO (ThreadDict a)
newThreadDict :: IO (ThreadDict a)
newThreadDict = do
  IORef (Map ThreadId a)
m <- Map ThreadId a -> IO (IORef (Map ThreadId a))
forall a. a -> IO (IORef a)
newIORef Map ThreadId a
forall k a. Map k a
Map.empty
  ThreadDict a -> IO (ThreadDict a)
forall (m :: * -> *) a. Monad m => a -> m a
return (IORef (Map ThreadId a) -> ThreadDict a
forall a. IORef (Map ThreadId a) -> ThreadDict a
ThreadDict IORef (Map ThreadId a)
m)

writeThreadDict :: ThreadDict a -> a -> IO ()
writeThreadDict :: ThreadDict a -> a -> IO ()
writeThreadDict (ThreadDict IORef (Map ThreadId a)
table) a
a =
   do
      ThreadId
ti <- IO ThreadId
myThreadId
      IORef (Map ThreadId a)
-> (Map ThreadId a -> (Map ThreadId a, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef (Map ThreadId a)
table ((Map ThreadId a -> (Map ThreadId a, ())) -> IO ())
-> (Map ThreadId a -> (Map ThreadId a, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Map ThreadId a
m -> (ThreadId -> a -> Map ThreadId a -> Map ThreadId a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert ThreadId
ti a
a Map ThreadId a
m, ())

readThreadDict :: ThreadDict a -> IO (Maybe a)
readThreadDict :: ThreadDict a -> IO (Maybe a)
readThreadDict (ThreadDict IORef (Map ThreadId a)
table) =
   do
      ThreadId
ti <- IO ThreadId
myThreadId
      Map ThreadId a
m <- IORef (Map ThreadId a) -> IO (Map ThreadId a)
forall a. IORef a -> IO a
readIORef IORef (Map ThreadId a)
table
      Maybe a -> IO (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> IO (Maybe a)) -> Maybe a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ ThreadId -> Map ThreadId a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ThreadId
ti Map ThreadId a
m

modifyThreadDict :: ThreadDict a -> (Maybe a -> IO (Maybe a, b)) -> IO b
modifyThreadDict :: ThreadDict a -> (Maybe a -> IO (Maybe a, b)) -> IO b
modifyThreadDict (ThreadDict IORef (Map ThreadId a)
table) Maybe a -> IO (Maybe a, b)
updateFn =
   do
      ThreadId
ti <- IO ThreadId
myThreadId
      Map ThreadId a
m <- IORef (Map ThreadId a) -> IO (Map ThreadId a)
forall a. IORef a -> IO a
readIORef IORef (Map ThreadId a)
table
      (Maybe a
aOpt1, b
b) <- Maybe a -> IO (Maybe a, b)
updateFn (Maybe a -> IO (Maybe a, b)) -> Maybe a -> IO (Maybe a, b)
forall a b. (a -> b) -> a -> b
$ ThreadId -> Map ThreadId a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ThreadId
ti Map ThreadId a
m
      IORef (Map ThreadId a)
-> (Map ThreadId a -> (Map ThreadId a, b)) -> IO b
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef (Map ThreadId a)
table ((Map ThreadId a -> (Map ThreadId a, b)) -> IO b)
-> (Map ThreadId a -> (Map ThreadId a, b)) -> IO b
forall a b. (a -> b) -> a -> b
$ \ Map ThreadId a
im -> ((case Maybe a
aOpt1 of
            Maybe a
Nothing -> ThreadId -> Map ThreadId a -> Map ThreadId a
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete ThreadId
ti
            Just a
a -> ThreadId -> a -> Map ThreadId a -> Map ThreadId a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert ThreadId
ti a
a) Map ThreadId a
im, b
b)