{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ExistentialQuantification #-}

module Test.MockCat.Internal.Registry.Core
  ( attachVerifierToFn
  , lookupVerifierForFn
  , attachDynamicVerifierToFn
  , createOverlay
  , installOverlay
  , clearOverlay
  , registerUnitMeta
  , lookupUnitMeta
  , UnitMeta
  , withUnitGuard
  , withAllUnitGuards
  , markUnitUsed
  , isGuardActive
  ) where

import Control.Concurrent.STM
  ( TVar
  , atomically
  , modifyTVar'
  , newTVarIO
  , readTVar
  , readTVarIO
  , writeTVar
  )
import Control.Exception (bracket_)
import Control.Monad (forM_)
import Data.Dynamic (Dynamic, toDyn)
import Data.Typeable (Typeable)
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import System.IO.Unsafe (unsafePerformIO)
import Test.MockCat.Internal.Types (MockName, InvocationRecorder(..))
import System.Mem.StableName (StableName, eqStableName, hashStableName, makeStableName)

data SomeStableName = forall a. SomeStableName (StableName a)

instance Eq SomeStableName where
  (SomeStableName StableName a
sn1) == :: UnitStableName -> UnitStableName -> Bool
== (SomeStableName StableName a
sn2) = StableName a
sn1 StableName a -> StableName a -> Bool
forall a b. StableName a -> StableName b -> Bool
`eqStableName` StableName a
sn2

type FnStableName = SomeStableName

data Entry
  = Entry !FnStableName !Dynamic
  | NamedEntry !FnStableName !MockName !Dynamic

stableFnName :: Entry -> FnStableName
stableFnName :: Entry -> UnitStableName
stableFnName (Entry UnitStableName
fn Dynamic
_) = UnitStableName
fn
stableFnName (NamedEntry UnitStableName
fn MockName
_ Dynamic
_) = UnitStableName
fn

mockName :: Entry -> Maybe MockName
mockName :: Entry -> Maybe MockName
mockName (NamedEntry UnitStableName
_ MockName
name Dynamic
_) = MockName -> Maybe MockName
forall a. a -> Maybe a
Just MockName
name
mockName Entry
_ = Maybe MockName
forall a. Maybe a
Nothing

entryPayload :: Entry -> Dynamic
entryPayload :: Entry -> Dynamic
entryPayload (Entry UnitStableName
_ Dynamic
payload) = Dynamic
payload
entryPayload (NamedEntry UnitStableName
_ MockName
_ Dynamic
payload) = Dynamic
payload

toFnStable :: forall a. StableName a -> FnStableName
toFnStable :: forall a. StableName a -> UnitStableName
toFnStable = StableName a -> UnitStableName
forall a. StableName a -> UnitStableName
SomeStableName

sameFnStable :: FnStableName -> FnStableName -> Bool
sameFnStable :: UnitStableName -> UnitStableName -> Bool
sameFnStable UnitStableName
a UnitStableName
b = UnitStableName
a UnitStableName -> UnitStableName -> Bool
forall a. Eq a => a -> a -> Bool
== UnitStableName
b

type Registry = IntMap [Entry]

registry :: TVar Registry
registry :: TVar Registry
registry = (IO (TVar Registry) -> TVar Registry
forall a. IO a -> a
unsafePerformIO (IO (TVar Registry) -> TVar Registry)
-> IO (TVar Registry) -> TVar Registry
forall a b. (a -> b) -> a -> b
$ Registry -> IO (TVar Registry)
forall a. a -> IO (TVar a)
newTVarIO Registry
forall a. IntMap a
IntMap.empty) :: TVar Registry


attachVerifierToFn ::
  forall fn params.
  (Typeable (InvocationRecorder params)) =>
  fn ->
  (Maybe MockName, InvocationRecorder params) ->
  IO ()
attachVerifierToFn :: forall fn params.
Typeable (InvocationRecorder params) =>
fn -> (Maybe MockName, InvocationRecorder params) -> IO ()
attachVerifierToFn fn
fn (Maybe MockName
name, InvocationRecorder params
payload) = fn -> (Maybe MockName, Dynamic) -> IO ()
forall fn. fn -> (Maybe MockName, Dynamic) -> IO ()
attachDynamicVerifierToFn fn
fn (Maybe MockName
name, InvocationRecorder params -> Dynamic
forall a. Typeable a => a -> Dynamic
toDyn InvocationRecorder params
payload)

lookupVerifierForFn ::
  forall fn.
  fn ->
  IO (Maybe (Maybe MockName, Dynamic))
lookupVerifierForFn :: forall fn. fn -> IO (Maybe (Maybe MockName, Dynamic))
lookupVerifierForFn fn
fn = do
  StableName fn
stable <- fn -> IO (StableName fn)
forall a. a -> IO (StableName a)
makeStableName fn
fn
  let
    key :: Int
key = StableName fn -> Int
forall a. StableName a -> Int
hashStableName StableName fn
stable
    stableFn :: UnitStableName
stableFn = StableName fn -> UnitStableName
forall a. StableName a -> UnitStableName
toFnStable StableName fn
stable
  Registry
store <- TVar Registry -> IO Registry
forall a. TVar a -> IO a
readTVarIO TVar Registry
registry
  -- ONLY use direct StableName matching in the registry.
  -- Name-based resolution via lookupNameByHash here is unsafe because it can
  -- return a verifier from a previous session if the hash collided or was reused.
  case Int -> Registry -> Maybe [Entry]
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
key Registry
store Maybe [Entry]
-> ([Entry] -> Maybe (Maybe MockName, Dynamic))
-> Maybe (Maybe MockName, Dynamic)
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= UnitStableName -> [Entry] -> Maybe (Maybe MockName, Dynamic)
findMatch UnitStableName
stableFn of
    Just (Maybe MockName, Dynamic)
res -> Maybe (Maybe MockName, Dynamic)
-> IO (Maybe (Maybe MockName, Dynamic))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Maybe MockName, Dynamic) -> Maybe (Maybe MockName, Dynamic)
forall a. a -> Maybe a
Just (Maybe MockName, Dynamic)
res)
    Maybe (Maybe MockName, Dynamic)
Nothing -> Maybe (Maybe MockName, Dynamic)
-> IO (Maybe (Maybe MockName, Dynamic))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Maybe MockName, Dynamic)
forall a. Maybe a
Nothing

attachDynamicVerifierToFn :: forall fn. fn -> (Maybe MockName, Dynamic) -> IO ()
attachDynamicVerifierToFn :: forall fn. fn -> (Maybe MockName, Dynamic) -> IO ()
attachDynamicVerifierToFn fn
fn (Maybe MockName
name, Dynamic
payload) = do
  -- Record stable-name of the passed function
  StableName fn
stable <- fn -> IO (StableName fn)
forall a. a -> IO (StableName a)
makeStableName fn
fn
  let stableFn :: UnitStableName
stableFn = StableName fn -> UnitStableName
forall a. StableName a -> UnitStableName
toFnStable StableName fn
stable
  let passedKey :: Int
passedKey = StableName fn -> Int
forall a. StableName a -> Int
hashStableName StableName fn
stable
  -- Always attach to the passed function stable-name directly.
  -- Avoiding lookupFnByName here prevents cross-session identity pollution.
  let key :: Int
key = Int
passedKey
  let entry :: Entry
entry = Maybe MockName -> UnitStableName -> Dynamic -> Entry
toEntry Maybe MockName
name UnitStableName
stableFn Dynamic
payload
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$
    TVar Registry -> (Registry -> Registry) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Registry
registry ((Registry -> Registry) -> STM ())
-> (Registry -> Registry) -> STM ()
forall a b. (a -> b) -> a -> b
$ \Registry
m -> (Maybe [Entry] -> Maybe [Entry]) -> Int -> Registry -> Registry
forall a. (Maybe a -> Maybe a) -> Int -> IntMap a -> IntMap a
IntMap.alter (Entry -> UnitStableName -> Maybe [Entry] -> Maybe [Entry]
updateEntries Entry
entry UnitStableName
stableFn) Int
key Registry
m

toEntry :: Maybe MockName -> FnStableName -> Dynamic -> Entry
toEntry :: Maybe MockName -> UnitStableName -> Dynamic -> Entry
toEntry (Just MockName
n) UnitStableName
stableFn Dynamic
p = UnitStableName -> MockName -> Dynamic -> Entry
NamedEntry UnitStableName
stableFn MockName
n Dynamic
p
toEntry Maybe MockName
Nothing UnitStableName
stableFn Dynamic
p = UnitStableName -> Dynamic -> Entry
Entry UnitStableName
stableFn Dynamic
p

updateEntries :: Entry -> FnStableName -> Maybe [Entry] -> Maybe [Entry]
updateEntries :: Entry -> UnitStableName -> Maybe [Entry] -> Maybe [Entry]
updateEntries Entry
entry UnitStableName
stableFn (Just [Entry]
entries) = [Entry] -> Maybe [Entry]
forall a. a -> Maybe a
Just ([Entry] -> Maybe [Entry]) -> [Entry] -> Maybe [Entry]
forall a b. (a -> b) -> a -> b
$ Entry
entry Entry -> [Entry] -> [Entry]
forall a. a -> [a] -> [a]
: UnitStableName -> [Entry] -> [Entry]
filterSameFnStable UnitStableName
stableFn [Entry]
entries
updateEntries Entry
entry UnitStableName
_        Maybe [Entry]
Nothing        = [Entry] -> Maybe [Entry]
forall a. a -> Maybe a
Just [Entry
entry]

filterSameFnStable :: FnStableName -> [Entry] -> [Entry]
filterSameFnStable :: UnitStableName -> [Entry] -> [Entry]
filterSameFnStable UnitStableName
stableFn = (Entry -> Bool) -> [Entry] -> [Entry]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Entry -> Bool) -> Entry -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnitStableName -> UnitStableName -> Bool
sameFnStable UnitStableName
stableFn (UnitStableName -> Bool)
-> (Entry -> UnitStableName) -> Entry -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry -> UnitStableName
stableFnName)

findMatch :: FnStableName -> [Entry] -> Maybe (Maybe MockName, Dynamic)
findMatch :: UnitStableName -> [Entry] -> Maybe (Maybe MockName, Dynamic)
findMatch UnitStableName
_ [] = Maybe (Maybe MockName, Dynamic)
forall a. Maybe a
Nothing
findMatch UnitStableName
target  (Entry
entry : [Entry]
rest)
  | UnitStableName -> UnitStableName -> Bool
sameFnStable UnitStableName
target (Entry -> UnitStableName
stableFnName Entry
entry) = (Maybe MockName, Dynamic) -> Maybe (Maybe MockName, Dynamic)
forall a. a -> Maybe a
Just (Entry -> Maybe MockName
mockName Entry
entry, Entry -> Dynamic
entryPayload Entry
entry)
  | Bool
otherwise = UnitStableName -> [Entry] -> Maybe (Maybe MockName, Dynamic)
findMatch UnitStableName
target [Entry]
rest







type UnitStableName = SomeStableName

data UnitMeta = UnitMeta
  { UnitMeta -> TVar Bool
unitGuardRef :: TVar Bool
  , UnitMeta -> TVar Bool
unitUsedRef :: TVar Bool
  }

data UnitEntry = UnitEntry !UnitStableName !UnitMeta

unitEntryStable :: UnitEntry -> UnitStableName
unitEntryStable :: UnitEntry -> UnitStableName
unitEntryStable (UnitEntry UnitStableName
stable UnitMeta
_) = UnitStableName
stable

unitEntryMeta :: UnitEntry -> UnitMeta
unitEntryMeta :: UnitEntry -> UnitMeta
unitEntryMeta (UnitEntry UnitStableName
_ UnitMeta
meta) = UnitMeta
meta

toUnitStable :: forall a. StableName a -> UnitStableName
toUnitStable :: forall a. StableName a -> UnitStableName
toUnitStable = StableName a -> UnitStableName
forall a. StableName a -> UnitStableName
SomeStableName

sameUnitStable :: UnitStableName -> UnitStableName -> Bool
sameUnitStable :: UnitStableName -> UnitStableName -> Bool
sameUnitStable UnitStableName
a UnitStableName
b = UnitStableName
a UnitStableName -> UnitStableName -> Bool
forall a. Eq a => a -> a -> Bool
== UnitStableName
b

type UnitRegistry = IntMap [UnitEntry]

unitRegistry :: TVar UnitRegistry
unitRegistry :: TVar UnitRegistry
unitRegistry = (IO (TVar UnitRegistry) -> TVar UnitRegistry
forall a. IO a -> a
unsafePerformIO (IO (TVar UnitRegistry) -> TVar UnitRegistry)
-> IO (TVar UnitRegistry) -> TVar UnitRegistry
forall a b. (a -> b) -> a -> b
$ UnitRegistry -> IO (TVar UnitRegistry)
forall a. a -> IO (TVar a)
newTVarIO UnitRegistry
forall a. IntMap a
IntMap.empty) :: TVar UnitRegistry

-- | Per-run overlay registry (optional).
data Overlay = Overlay

-- | Run the given IO action with a per-run overlay registry active.
-- The overlay is cleaned up after the action completes.
createOverlay :: IO Overlay
createOverlay :: IO Overlay
createOverlay = Overlay -> IO Overlay
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Overlay
Overlay

installOverlay :: Overlay -> IO ()
installOverlay :: Overlay -> IO ()
installOverlay Overlay
_ = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

clearOverlay :: IO ()
clearOverlay :: IO ()
clearOverlay = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()



registerUnitMeta :: TVar ref -> IO UnitMeta
registerUnitMeta :: forall ref. TVar ref -> IO UnitMeta
registerUnitMeta TVar ref
ref = do
  StableName (TVar ref)
stable <- TVar ref -> IO (StableName (TVar ref))
forall a. a -> IO (StableName a)
makeStableName TVar ref
ref
  let key :: Int
key = StableName (TVar ref) -> Int
forall a. StableName a -> Int
hashStableName StableName (TVar ref)
stable
      unitStable :: UnitStableName
unitStable = StableName (TVar ref) -> UnitStableName
forall a. StableName a -> UnitStableName
toUnitStable StableName (TVar ref)
stable
  UnitMeta
fresh <- IO UnitMeta
createUnitMeta
  STM UnitMeta -> IO UnitMeta
forall a. STM a -> IO a
atomically (STM UnitMeta -> IO UnitMeta) -> STM UnitMeta -> IO UnitMeta
forall a b. (a -> b) -> a -> b
$ do
    UnitRegistry
store <- TVar UnitRegistry -> STM UnitRegistry
forall a. TVar a -> STM a
readTVar TVar UnitRegistry
unitRegistry
    case Int -> UnitRegistry -> Maybe [UnitEntry]
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
key UnitRegistry
store of
      Just [UnitEntry]
entries ->
        case UnitStableName -> [UnitEntry] -> Maybe UnitMeta
findUnit UnitStableName
unitStable [UnitEntry]
entries of
          Just UnitMeta
existing -> UnitMeta -> STM UnitMeta
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure UnitMeta
existing
          Maybe UnitMeta
Nothing -> do
            let newEntries :: [UnitEntry]
newEntries = UnitStableName -> UnitMeta -> UnitEntry
UnitEntry UnitStableName
unitStable UnitMeta
fresh UnitEntry -> [UnitEntry] -> [UnitEntry]
forall a. a -> [a] -> [a]
: [UnitEntry]
entries
            TVar UnitRegistry -> UnitRegistry -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar UnitRegistry
unitRegistry (Int -> [UnitEntry] -> UnitRegistry -> UnitRegistry
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
key [UnitEntry]
newEntries UnitRegistry
store)
            UnitMeta -> STM UnitMeta
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure UnitMeta
fresh
      Maybe [UnitEntry]
Nothing -> do
        TVar UnitRegistry -> UnitRegistry -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar UnitRegistry
unitRegistry (Int -> [UnitEntry] -> UnitRegistry -> UnitRegistry
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
key [UnitStableName -> UnitMeta -> UnitEntry
UnitEntry UnitStableName
unitStable UnitMeta
fresh] UnitRegistry
store)
        UnitMeta -> STM UnitMeta
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure UnitMeta
fresh

lookupUnitMeta :: TVar ref -> IO (Maybe UnitMeta)
lookupUnitMeta :: forall ref. TVar ref -> IO (Maybe UnitMeta)
lookupUnitMeta TVar ref
ref = do
  StableName (TVar ref)
stable <- TVar ref -> IO (StableName (TVar ref))
forall a. a -> IO (StableName a)
makeStableName TVar ref
ref
  let key :: Int
key = StableName (TVar ref) -> Int
forall a. StableName a -> Int
hashStableName StableName (TVar ref)
stable
      unitStable :: UnitStableName
unitStable = StableName (TVar ref) -> UnitStableName
forall a. StableName a -> UnitStableName
toUnitStable StableName (TVar ref)
stable
  UnitRegistry
store <- TVar UnitRegistry -> IO UnitRegistry
forall a. TVar a -> IO a
readTVarIO TVar UnitRegistry
unitRegistry
  Maybe UnitMeta -> IO (Maybe UnitMeta)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe UnitMeta -> IO (Maybe UnitMeta))
-> Maybe UnitMeta -> IO (Maybe UnitMeta)
forall a b. (a -> b) -> a -> b
$ Int -> UnitRegistry -> Maybe [UnitEntry]
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
key UnitRegistry
store Maybe [UnitEntry]
-> ([UnitEntry] -> Maybe UnitMeta) -> Maybe UnitMeta
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= UnitStableName -> [UnitEntry] -> Maybe UnitMeta
findUnit UnitStableName
unitStable

withUnitGuard :: UnitMeta -> IO a -> IO a
withUnitGuard :: forall a. UnitMeta -> IO a -> IO a
withUnitGuard UnitMeta
meta =
  IO () -> IO () -> IO a -> IO a
forall a b c. IO a -> IO b -> IO c -> IO c
bracket_
    (STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (UnitMeta -> TVar Bool
unitGuardRef UnitMeta
meta) Bool
True)
    (STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (UnitMeta -> TVar Bool
unitGuardRef UnitMeta
meta) Bool
False)

withAllUnitGuards :: IO a -> IO a
withAllUnitGuards :: forall a. IO a -> IO a
withAllUnitGuards = IO () -> IO () -> IO a -> IO a
forall a b c. IO a -> IO b -> IO c -> IO c
bracket_ (Bool -> IO ()
setAllUnitGuards Bool
True) (Bool -> IO ()
setAllUnitGuards Bool
False)

markUnitUsed :: UnitMeta -> IO ()
markUnitUsed :: UnitMeta -> IO ()
markUnitUsed UnitMeta
meta = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (UnitMeta -> TVar Bool
unitUsedRef UnitMeta
meta) Bool
True

isGuardActive :: UnitMeta -> IO Bool
isGuardActive :: UnitMeta -> IO Bool
isGuardActive UnitMeta
meta = TVar Bool -> IO Bool
forall a. TVar a -> IO a
readTVarIO (UnitMeta -> TVar Bool
unitGuardRef UnitMeta
meta)

createUnitMeta :: IO UnitMeta
createUnitMeta :: IO UnitMeta
createUnitMeta = do
  TVar Bool
guardRef <- Bool -> IO (TVar Bool)
forall a. a -> IO (TVar a)
newTVarIO Bool
False
  TVar Bool
usedRef <- Bool -> IO (TVar Bool)
forall a. a -> IO (TVar a)
newTVarIO Bool
False
  UnitMeta -> IO UnitMeta
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure UnitMeta {unitGuardRef :: TVar Bool
unitGuardRef = TVar Bool
guardRef, unitUsedRef :: TVar Bool
unitUsedRef = TVar Bool
usedRef}

findUnit :: UnitStableName -> [UnitEntry] -> Maybe UnitMeta
findUnit :: UnitStableName -> [UnitEntry] -> Maybe UnitMeta
findUnit UnitStableName
_ [] = Maybe UnitMeta
forall a. Maybe a
Nothing
findUnit UnitStableName
target (UnitEntry
entry : [UnitEntry]
rest)
  | UnitStableName -> UnitStableName -> Bool
sameUnitStable UnitStableName
target (UnitEntry -> UnitStableName
unitEntryStable UnitEntry
entry) = UnitMeta -> Maybe UnitMeta
forall a. a -> Maybe a
Just (UnitEntry -> UnitMeta
unitEntryMeta UnitEntry
entry)
  | Bool
otherwise = UnitStableName -> [UnitEntry] -> Maybe UnitMeta
findUnit UnitStableName
target [UnitEntry]
rest

setAllUnitGuards :: Bool -> IO ()
setAllUnitGuards :: Bool -> IO ()
setAllUnitGuards Bool
flag =
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    UnitRegistry
store <- TVar UnitRegistry -> STM UnitRegistry
forall a. TVar a -> STM a
readTVar TVar UnitRegistry
unitRegistry
    [UnitEntry] -> (UnitEntry -> STM ()) -> STM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([[UnitEntry]] -> [UnitEntry]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (UnitRegistry -> [[UnitEntry]]
forall a. IntMap a -> [a]
IntMap.elems UnitRegistry
store)) ((UnitEntry -> STM ()) -> STM ())
-> (UnitEntry -> STM ()) -> STM ()
forall a b. (a -> b) -> a -> b
$ \UnitEntry
entry ->
      TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (UnitMeta -> TVar Bool
unitGuardRef (UnitEntry -> UnitMeta
unitEntryMeta UnitEntry
entry)) Bool
flag