{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TypeApplications #-}

module Ki.Ctx
  ( Ctx (..),
    newCtxSTM,
    deriveCtx,
    cancelCtx,
    cancelCtxSTM,
    ctxCancelToken,
  )
where

import qualified Data.IntMap.Strict as IntMap
import Ki.CancelToken
import Ki.Prelude

data Ctx = Ctx
  { Ctx -> TVar (Maybe CancelToken)
cancelTokenVar :: TVar (Maybe CancelToken),
    Ctx -> TVar (IntMap Ctx)
childrenVar :: TVar (IntMap Ctx),
    -- | The next id to assign to a child context. The child needs a unique identifier so it can delete itself from its
    -- parent's children map if it's cancelled independently. Wrap-around seems ok; that's a *lot* of children for one
    -- parent to have.
    Ctx -> TVar Int
nextIdVar :: TVar Int,
    -- | When I'm cancelled, this action removes myself from my parent's context. This isn't simply a pointer to the
    -- parent 'Ctx' for three reasons:
    --
    --   * "Root" contexts don't have a parent, so it'd have to be a Maybe (one more pointer indirection)
    --   * We don't really need a reference to the parent, because we only want to be able to remove ourselves from its
    --     children map, so just storing the STM action that does exactly seems a bit safer, even if conceptually it's
    --     a bit indirect.
    --   * If we stored a reference to the parent, we'd also have to store our own id, rather than just currying it into
    --     this action.
    Ctx -> STM ()
onCancel :: STM ()
  }

newCtxSTM :: STM Ctx
newCtxSTM :: STM Ctx
newCtxSTM =
  STM () -> STM Ctx
newCtxSTM_ (() -> STM ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

newCtxSTM_ :: STM () -> STM Ctx
newCtxSTM_ :: STM () -> STM Ctx
newCtxSTM_ STM ()
onCancel = do
  TVar (Maybe CancelToken)
cancelTokenVar <- Maybe CancelToken -> STM (TVar (Maybe CancelToken))
forall a. a -> STM (TVar a)
newTVar Maybe CancelToken
forall a. Maybe a
Nothing
  TVar (IntMap Ctx)
childrenVar <- IntMap Ctx -> STM (TVar (IntMap Ctx))
forall a. a -> STM (TVar a)
newTVar IntMap Ctx
forall a. IntMap a
IntMap.empty
  TVar Int
nextIdVar <- Int -> STM (TVar Int)
forall a. a -> STM (TVar a)
newTVar Int
0
  Ctx -> STM Ctx
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ctx :: TVar (Maybe CancelToken)
-> TVar (IntMap Ctx) -> TVar Int -> STM () -> Ctx
Ctx {TVar (Maybe CancelToken)
cancelTokenVar :: TVar (Maybe CancelToken)
$sel:cancelTokenVar:Ctx :: TVar (Maybe CancelToken)
cancelTokenVar, TVar (IntMap Ctx)
childrenVar :: TVar (IntMap Ctx)
$sel:childrenVar:Ctx :: TVar (IntMap Ctx)
childrenVar, TVar Int
nextIdVar :: TVar Int
$sel:nextIdVar:Ctx :: TVar Int
nextIdVar, STM ()
onCancel :: STM ()
$sel:onCancel:Ctx :: STM ()
onCancel}

deriveCtx :: Ctx -> STM Ctx
deriveCtx :: Ctx -> STM Ctx
deriveCtx context :: Ctx
context@Ctx {TVar (Maybe CancelToken)
cancelTokenVar :: TVar (Maybe CancelToken)
$sel:cancelTokenVar:Ctx :: Ctx -> TVar (Maybe CancelToken)
cancelTokenVar, TVar (IntMap Ctx)
childrenVar :: TVar (IntMap Ctx)
$sel:childrenVar:Ctx :: Ctx -> TVar (IntMap Ctx)
childrenVar, TVar Int
nextIdVar :: TVar Int
$sel:nextIdVar:Ctx :: Ctx -> TVar Int
nextIdVar} =
  TVar (Maybe CancelToken) -> STM (Maybe CancelToken)
forall a. TVar a -> STM a
readTVar TVar (Maybe CancelToken)
cancelTokenVar STM (Maybe CancelToken)
-> (Maybe CancelToken -> STM Ctx) -> STM Ctx
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe CancelToken
Nothing -> do
      Int
childId <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
nextIdVar
      TVar Int -> Int -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Int
nextIdVar (Int -> STM ()) -> Int -> STM ()
forall a b. (a -> b) -> a -> b
$! Int
childId Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
      Ctx
child <- STM () -> STM Ctx
newCtxSTM_ (TVar (IntMap Ctx) -> (IntMap Ctx -> IntMap Ctx) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (IntMap Ctx)
childrenVar (Int -> IntMap Ctx -> IntMap Ctx
forall a. Int -> IntMap a -> IntMap a
IntMap.delete Int
childId))
      IntMap Ctx
children <- TVar (IntMap Ctx) -> STM (IntMap Ctx)
forall a. TVar a -> STM a
readTVar TVar (IntMap Ctx)
childrenVar
      TVar (IntMap Ctx) -> IntMap Ctx -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (IntMap Ctx)
childrenVar (IntMap Ctx -> STM ()) -> IntMap Ctx -> STM ()
forall a b. (a -> b) -> a -> b
$! Int -> Ctx -> IntMap Ctx -> IntMap Ctx
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
childId Ctx
child IntMap Ctx
children
      Ctx -> STM Ctx
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ctx
child
    Just (CancelToken Int
_) -> Ctx -> STM Ctx
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ctx
context

cancelCtx :: Ctx -> IO ()
cancelCtx :: Ctx -> IO ()
cancelCtx Ctx
context = do
  CancelToken
token <- IO CancelToken
newCancelToken
  STM () -> IO ()
forall a. STM a -> IO a
atomically (Ctx -> CancelToken -> STM ()
cancelCtxSTM Ctx
context CancelToken
token)

cancelCtxSTM :: Ctx -> CancelToken -> STM ()
cancelCtxSTM :: Ctx -> CancelToken -> STM ()
cancelCtxSTM ctx :: Ctx
ctx@Ctx {STM ()
onCancel :: STM ()
$sel:onCancel:Ctx :: Ctx -> STM ()
onCancel} CancelToken
token =
  Ctx -> CancelToken -> STM () -> STM ()
whenCanceling Ctx
ctx CancelToken
token do
    Ctx -> CancelToken -> STM ()
cancelChildren Ctx
ctx CancelToken
token
    STM ()
onCancel

ctxCancelSTM_ :: CancelToken -> Ctx -> STM ()
ctxCancelSTM_ :: CancelToken -> Ctx -> STM ()
ctxCancelSTM_ CancelToken
token Ctx
ctx =
  Ctx -> CancelToken -> STM () -> STM ()
whenCanceling Ctx
ctx CancelToken
token (Ctx -> CancelToken -> STM ()
cancelChildren Ctx
ctx CancelToken
token)

whenCanceling :: Ctx -> CancelToken -> STM () -> STM ()
whenCanceling :: Ctx -> CancelToken -> STM () -> STM ()
whenCanceling Ctx {TVar (Maybe CancelToken)
cancelTokenVar :: TVar (Maybe CancelToken)
$sel:cancelTokenVar:Ctx :: Ctx -> TVar (Maybe CancelToken)
cancelTokenVar} CancelToken
token STM ()
action =
  TVar (Maybe CancelToken) -> STM (Maybe CancelToken)
forall a. TVar a -> STM a
readTVar TVar (Maybe CancelToken)
cancelTokenVar STM (Maybe CancelToken) -> (Maybe CancelToken -> STM ()) -> STM ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe CancelToken
Nothing -> do
      TVar (Maybe CancelToken) -> Maybe CancelToken -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Maybe CancelToken)
cancelTokenVar (Maybe CancelToken -> STM ()) -> Maybe CancelToken -> STM ()
forall a b. (a -> b) -> a -> b
$! CancelToken -> Maybe CancelToken
forall a. a -> Maybe a
Just CancelToken
token
      STM ()
action
    Just (CancelToken Int
_) -> () -> STM ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

cancelChildren :: Ctx -> CancelToken -> STM ()
cancelChildren :: Ctx -> CancelToken -> STM ()
cancelChildren Ctx {TVar (IntMap Ctx)
childrenVar :: TVar (IntMap Ctx)
$sel:childrenVar:Ctx :: Ctx -> TVar (IntMap Ctx)
childrenVar} CancelToken
token = do
  IntMap Ctx
children <- TVar (IntMap Ctx) -> STM (IntMap Ctx)
forall a. TVar a -> STM a
readTVar TVar (IntMap Ctx)
childrenVar
  [Ctx] -> (Ctx -> STM ()) -> STM ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (IntMap Ctx -> [Ctx]
forall a. IntMap a -> [a]
IntMap.elems IntMap Ctx
children) (CancelToken -> Ctx -> STM ()
ctxCancelSTM_ CancelToken
token)

ctxCancelToken :: Ctx -> STM CancelToken
ctxCancelToken :: Ctx -> STM CancelToken
ctxCancelToken Ctx {TVar (Maybe CancelToken)
cancelTokenVar :: TVar (Maybe CancelToken)
$sel:cancelTokenVar:Ctx :: Ctx -> TVar (Maybe CancelToken)
cancelTokenVar} =
  TVar (Maybe CancelToken) -> STM (Maybe CancelToken)
forall a. TVar a -> STM a
readTVar TVar (Maybe CancelToken)
cancelTokenVar STM (Maybe CancelToken)
-> (Maybe CancelToken -> STM CancelToken) -> STM CancelToken
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe CancelToken
Nothing -> STM CancelToken
forall a. STM a
retry
    Just CancelToken
token -> CancelToken -> STM CancelToken
forall (f :: * -> *) a. Applicative f => a -> f a
pure CancelToken
token