{-# OPTIONS -fno-cse #-}

{- Author: Robert Dockins.
   This work is hereby released into the public domain.
-}

module ThreadLocal 
( TL
, TLRef
, createThreadLocal
, tryReadThreadLocal
, readThreadLocal
, writeThreadLocal
, clearThreadLocal
, forkIO'
, tryReadTL
, readTL
, newIORefTL
, newMVarTL
, newEmptyMVarTL
, initMainThread
, initTL
, newThreadLocal
, newThreadLocalMVar
, newThreadLocalEmptyMVar
, newThreadLocalIORef
)
where

import System.IO.Unsafe
import System.Mem.StableName
import Data.IORef
import qualified Data.Map as Map
import Control.Concurrent
import Control.Concurrent.MVar
import Control.Monad.Reader
import qualified Control.Exception as Ex

newtype TLRef a = TLRef Integer
newtype TL a = TL { unTL :: ReaderT (Maybe ThreadId) IO a }

instance Monad TL where
  m >>= f = TL (unTL m >>= unTL . f)
  return  = TL . return
  fail    = TL . fail

-- the typing is going to get a little funny here.....
-- we use unsafePerformIO to break typing in the usual way
-- and we rely on the way we access the map to ensure
-- things don't break.

type ThreadLocalBank a = Map.Map Integer a

{-# NOINLINE threadLocalIds #-}
threadLocalIds :: MVar Integer
threadLocalIds = unsafePerformIO (newMVar 0)

{-# NOINLINE threadLocalBaseMap #-}
threadLocalBaseMap :: MVar (Map.Map (Maybe ThreadId) (ThreadLocalBank a))
threadLocalBaseMap = unsafePerformIO (newMVar Map.empty)

{-# NOINLINE threadLocalDefs #-}
threadLocalDefs :: MVar (Map.Map Integer (TL a))
threadLocalDefs = unsafePerformIO (newMVar Map.empty)

createThreadLocal :: TL a -> IO (TLRef a)
createThreadLocal tlInit =
   modifyMVar threadLocalIds (\id -> do
      modifyMVar_ threadLocalDefs (\defs -> return (Map.insert id tlInit defs))
      return (id+1,TLRef id))

tryReadThreadLocal :: TLRef a -> IO (Maybe a)
tryReadThreadLocal (TLRef i) = do
   tid <- myThreadId
   withMVar threadLocalBaseMap (\baseMap -> do
      let vars = Map.findWithDefault Map.empty (Just tid) baseMap
      return (Map.lookup i vars))

readThreadLocal :: TLRef a -> IO a
readThreadLocal tlref = do
   x <- tryReadThreadLocal tlref
   case x of
     Nothing -> fail "attempt to read from empty thread local var"
     Just z  -> return z

writeThreadLocal :: a -> TLRef a -> IO ()
writeThreadLocal x (TLRef i) = do
   tid <- myThreadId
   modifyMVar_ threadLocalBaseMap (\baseMap -> do
      let vars  = Map.findWithDefault Map.empty (Just tid) baseMap
      let vars' = Map.insert i x vars
      let baseMap' = Map.insert (Just tid) vars' baseMap
      return baseMap')

clearThreadLocal :: TLRef a -> IO ()
clearThreadLocal (TLRef i) = do
   tid <- myThreadId
   modifyMVar_ threadLocalBaseMap (\baseMap -> do
      let vars  = Map.findWithDefault Map.empty (Just tid) baseMap
      let vars' = Map.delete i vars
      let baseMap' = Map.insert (Just tid) vars' baseMap
      return baseMap')

runTL :: Maybe ThreadId -> TL a -> IO a
runTL tid tl = runReaderT (unTL tl) tid

initThreadLocal :: Maybe ThreadId -> ThreadId -> IO ()
initThreadLocal oldTid newTid = do
  withMVar threadLocalDefs (\defs -> do
    newVarList <- sequence [ runTL oldTid tlinit >>= return . ((,) i) | (i,tlinit) <- Map.toList defs ]
    let newVars = Map.fromList newVarList
    modifyMVar_ threadLocalBaseMap (\baseMap -> do
       return (Map.insert (Just newTid) newVars baseMap)))

initMainThread :: IO ()
initMainThread = do
   tid <- myThreadId
   initThreadLocal Nothing tid


forkIO' :: IO () -> IO ThreadId
forkIO' m = do
  mvar <- newEmptyMVar
  oldTid <- myThreadId
  tid <- forkIO (do
     newTid <- myThreadId
     initResult <- Ex.try (initThreadLocal (Just oldTid) newTid)
     case initResult of
        Left  ex  -> putMVar mvar (Just ex)
        Right ()  -> putMVar mvar Nothing >> m)

  x <- takeMVar mvar
  case x of
     Nothing -> return tid
     Just ex -> Ex.throwIO ex


tryReadTL :: TLRef a -> TL (Maybe a)
tryReadTL (TLRef i) = TL (do
   tid <- ask
   liftIO (withMVar threadLocalBaseMap (\baseMap -> do
      let vars = Map.findWithDefault Map.empty tid baseMap
      return (Map.lookup i vars))))

readTL :: TLRef a -> TL a
readTL ref = do
   x <- tryReadTL ref
   case x of
     Nothing -> fail "attempt to read from empty thread local"
     Just z  -> return z

newIORefTL :: a -> TL (IORef a)
newIORefTL x = TL (liftIO (newIORef x))

newMVarTL :: a -> TL (MVar a)
newMVarTL x = TL (liftIO (newMVar x))

newEmptyMVarTL :: TL (MVar a)
newEmptyMVarTL = TL (liftIO (newEmptyMVar))

initTL :: TL a -> TLRef a -> TL a
initTL m ref = do
   x <- tryReadTL ref
   case x of
     Nothing -> m
     Just z  -> return z

{-# NOINLINE newThreadLocal #-}
newThreadLocal :: a -> TLRef a
newThreadLocal a = ref
  where {-# NOINLINE ref #-}
        ref  = unsafePerformIO (createThreadLocal init)
        init = initTL (return a) ref

{-# NOINLINE newThreadLocalMVar #-}
newThreadLocalMVar :: a -> TLRef (MVar a)
newThreadLocalMVar a = ref
  where {-# NOINLIKNE ref #-}
        ref  = unsafePerformIO (createThreadLocal init)
        init = initTL (newMVarTL a) ref

{-# NOINLINE newThreadLocalEmptyMVar #-}
newThreadLocalEmptyMVar :: TLRef (MVar a)
newThreadLocalEmptyMVar = ref
  where {-# NOINLINE ref #-}
        ref  = unsafePerformIO (createThreadLocal init)
        init = initTL newEmptyMVarTL ref


{-# NOINLINE newThreadLocalIORef #-}
newThreadLocalIORef :: a -> TLRef (IORef a)
newThreadLocalIORef a = ref
  where {-# NOINLINE ref #-}
        ref  = unsafePerformIO (createThreadLocal init)
        init = initTL (newIORefTL a) ref