{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

-- | This module corresponds to "Control.Concurrent.MVar" in the @base@ package.
--
-- This module can be used as a drop-in replacement for
-- "Control.Concurrent.Class.MonadMVar.Strict", but not the other way around.
module Control.Concurrent.Class.MonadMVar.Strict.Checked
  ( -- * StrictMVar
    LazyMVar
  , StrictMVar
  , castStrictMVar
  , fromLazyMVar
  , isEmptyMVar
  , modifyMVar
  , modifyMVarMasked
  , modifyMVarMasked_
  , modifyMVar_
  , newEmptyMVar
  , newEmptyMVarWithInvariant
  , newMVar
  , newMVarWithInvariant
  , putMVar
  , readMVar
  , swapMVar
  , takeMVar
  , toLazyMVar
  , tryPutMVar
  , tryReadMVar
  , tryTakeMVar
  , unsafeToUncheckedStrictMVar
  , withMVar
  , withMVarMasked

    -- * Invariant
  , checkInvariant

    -- * Re-exports
  , MonadMVar
  ) where

import Control.Concurrent.Class.MonadMVar.Strict (LazyMVar, MonadMVar)
import Control.Concurrent.Class.MonadMVar.Strict qualified as Strict
import Data.Kind (Type)
import GHC.Stack (HasCallStack)

{-------------------------------------------------------------------------------
  StrictMVar
-------------------------------------------------------------------------------}

-- | A strict MVar with invariant checking.
--
-- There is a weaker invariant for a 'StrictMVar' than for a 'StrictTVar':
-- although all functions that modify the 'StrictMVar' check the invariant, we
-- do /not/ guarantee that the value inside the 'StrictMVar' always satisfies
-- the invariant. Instead, we /do/ guarantee that if the 'StrictMVar' is updated
-- with a value that does not satisfy the invariant, an exception is thrown. The
-- reason for this weaker guarantee is that leaving an 'MVar' empty can lead to
-- very hard to debug "blocked indefinitely" problems.
type StrictMVar :: (Type -> Type) -> Type -> Type
#if CHECK_MVAR_INVARIANTS
data StrictMVar m a = StrictMVar {
    -- | The invariant that is checked whenever the 'StrictMVar' is updated.
    forall (m :: * -> *) a. StrictMVar m a -> a -> Maybe String
invariant :: !(a -> Maybe String)
  , forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar      :: !(Strict.StrictMVar m a)
  }
#else
newtype StrictMVar m a = StrictMVar {
    mvar :: Strict.StrictMVar m a
  }
#endif

castStrictMVar ::
  LazyMVar m ~ LazyMVar n =>
  StrictMVar m a -> StrictMVar n a
castStrictMVar :: forall (m :: * -> *) (n :: * -> *) a.
(LazyMVar m ~ LazyMVar n) =>
StrictMVar m a -> StrictMVar n a
castStrictMVar StrictMVar m a
v = (a -> Maybe String) -> StrictMVar n a -> StrictMVar n a
forall a (m :: * -> *).
(a -> Maybe String) -> StrictMVar m a -> StrictMVar m a
mkStrictMVar (StrictMVar m a -> a -> Maybe String
forall (m :: * -> *) a. StrictMVar m a -> a -> Maybe String
getInvariant StrictMVar m a
v) (StrictMVar m a -> StrictMVar n a
forall (m :: * -> *) (n :: * -> *) a.
(LazyMVar m ~ LazyMVar n) =>
StrictMVar m a -> StrictMVar n a
Strict.castStrictMVar (StrictMVar m a -> StrictMVar n a)
-> StrictMVar m a -> StrictMVar n a
forall a b. (a -> b) -> a -> b
$ StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar StrictMVar m a
v)

-- | Get the underlying @MVar@
--
-- Since we obviously can not guarantee that updates to this 'LazyMVar' will be
-- strict, this should be used with caution.
--
-- Similarly, we can not guarantee that updates to this 'LazyMVar' do not break
-- the original invariant that the 'StrictMVar' held.
toLazyMVar :: StrictMVar m a -> LazyMVar m a
toLazyMVar :: forall (m :: * -> *) a. StrictMVar m a -> LazyMVar m a
toLazyMVar = StrictMVar m a -> MVar m a
forall (m :: * -> *) a. StrictMVar m a -> LazyMVar m a
Strict.toLazyMVar (StrictMVar m a -> MVar m a)
-> (StrictMVar m a -> StrictMVar m a) -> StrictMVar m a -> MVar m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar

-- | Create a 'StrictMVar' from a 'LazyMVar'
--
-- It is not guaranteed that the 'LazyMVar' contains a value that is in WHNF, so
-- there is no guarantee that the resulting 'StrictMVar' contains a value that
-- is in WHNF. This should be used with caution.
--
-- The resulting 'StrictMVar' has a trivial invariant.
fromLazyMVar :: LazyMVar m a -> StrictMVar m a
fromLazyMVar :: forall (m :: * -> *) a. LazyMVar m a -> StrictMVar m a
fromLazyMVar = (a -> Maybe String) -> StrictMVar m a -> StrictMVar m a
forall a (m :: * -> *).
(a -> Maybe String) -> StrictMVar m a -> StrictMVar m a
mkStrictMVar (Maybe String -> a -> Maybe String
forall a b. a -> b -> a
const Maybe String
forall a. Maybe a
Nothing) (StrictMVar m a -> StrictMVar m a)
-> (MVar m a -> StrictMVar m a) -> MVar m a -> StrictMVar m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MVar m a -> StrictMVar m a
forall (m :: * -> *) a. MVar m a -> StrictMVar m a
Strict.fromLazyMVar

-- | Create an unchecked reference to the given checked 'StrictMVar'.
--
-- Note that the invariant is only guaranteed when modifying the checked MVar.
-- Any modification to the unchecked reference might break the invariants.
unsafeToUncheckedStrictMVar :: StrictMVar m a -> Strict.StrictMVar m a
unsafeToUncheckedStrictMVar :: forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
unsafeToUncheckedStrictMVar = StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar

newEmptyMVar :: MonadMVar m => m (StrictMVar m a)
newEmptyMVar :: forall (m :: * -> *) a. MonadMVar m => m (StrictMVar m a)
newEmptyMVar = (a -> Maybe String) -> StrictMVar m a -> StrictMVar m a
forall a (m :: * -> *).
(a -> Maybe String) -> StrictMVar m a -> StrictMVar m a
mkStrictMVar (Maybe String -> a -> Maybe String
forall a b. a -> b -> a
const Maybe String
forall a. Maybe a
Nothing) (StrictMVar m a -> StrictMVar m a)
-> m (StrictMVar m a) -> m (StrictMVar m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (StrictMVar m a)
forall (m :: * -> *) a. MonadMVar m => m (StrictMVar m a)
Strict.newEmptyMVar

newEmptyMVarWithInvariant ::
  MonadMVar m =>
  (a -> Maybe String) ->
  m (StrictMVar m a)
newEmptyMVarWithInvariant :: forall (m :: * -> *) a.
MonadMVar m =>
(a -> Maybe String) -> m (StrictMVar m a)
newEmptyMVarWithInvariant a -> Maybe String
inv = (a -> Maybe String) -> StrictMVar m a -> StrictMVar m a
forall a (m :: * -> *).
(a -> Maybe String) -> StrictMVar m a -> StrictMVar m a
mkStrictMVar a -> Maybe String
inv (StrictMVar m a -> StrictMVar m a)
-> m (StrictMVar m a) -> m (StrictMVar m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (StrictMVar m a)
forall (m :: * -> *) a. MonadMVar m => m (StrictMVar m a)
Strict.newEmptyMVar

newMVar :: MonadMVar m => a -> m (StrictMVar m a)
newMVar :: forall (m :: * -> *) a. MonadMVar m => a -> m (StrictMVar m a)
newMVar a
a = (a -> Maybe String) -> StrictMVar m a -> StrictMVar m a
forall a (m :: * -> *).
(a -> Maybe String) -> StrictMVar m a -> StrictMVar m a
mkStrictMVar (Maybe String -> a -> Maybe String
forall a b. a -> b -> a
const Maybe String
forall a. Maybe a
Nothing) (StrictMVar m a -> StrictMVar m a)
-> m (StrictMVar m a) -> m (StrictMVar m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m (StrictMVar m a)
forall (m :: * -> *) a. MonadMVar m => a -> m (StrictMVar m a)
Strict.newMVar a
a

-- | Create a 'StrictMVar' with an invariant.
--
-- Contrary to functions that modify a 'StrictMVar', this function checks the
-- invariant /before/ putting the value in a new 'StrictMVar'.
newMVarWithInvariant ::
  (HasCallStack, MonadMVar m) =>
  (a -> Maybe String) ->
  a ->
  m (StrictMVar m a)
newMVarWithInvariant :: forall (m :: * -> *) a.
(HasCallStack, MonadMVar m) =>
(a -> Maybe String) -> a -> m (StrictMVar m a)
newMVarWithInvariant a -> Maybe String
inv !a
a =
  Maybe String -> m (StrictMVar m a) -> m (StrictMVar m a)
forall a. HasCallStack => Maybe String -> a -> a
checkInvariant (a -> Maybe String
inv a
a) (m (StrictMVar m a) -> m (StrictMVar m a))
-> m (StrictMVar m a) -> m (StrictMVar m a)
forall a b. (a -> b) -> a -> b
$
    (a -> Maybe String) -> StrictMVar m a -> StrictMVar m a
forall a (m :: * -> *).
(a -> Maybe String) -> StrictMVar m a -> StrictMVar m a
mkStrictMVar a -> Maybe String
inv (StrictMVar m a -> StrictMVar m a)
-> m (StrictMVar m a) -> m (StrictMVar m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m (StrictMVar m a)
forall (m :: * -> *) a. MonadMVar m => a -> m (StrictMVar m a)
Strict.newMVar a
a

takeMVar :: MonadMVar m => StrictMVar m a -> m a
takeMVar :: forall (m :: * -> *) a. MonadMVar m => StrictMVar m a -> m a
takeMVar = StrictMVar m a -> m a
forall (m :: * -> *) a. MonadMVar m => StrictMVar m a -> m a
Strict.takeMVar (StrictMVar m a -> m a)
-> (StrictMVar m a -> StrictMVar m a) -> StrictMVar m a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar

putMVar :: (HasCallStack, MonadMVar m) => StrictMVar m a -> a -> m ()
putMVar :: forall (m :: * -> *) a.
(HasCallStack, MonadMVar m) =>
StrictMVar m a -> a -> m ()
putMVar StrictMVar m a
v a
a = do
  StrictMVar m a -> a -> m ()
forall (m :: * -> *) a. MonadMVar m => StrictMVar m a -> a -> m ()
Strict.putMVar (StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar StrictMVar m a
v) a
a
  Maybe String -> m () -> m ()
forall a. HasCallStack => Maybe String -> a -> a
checkInvariant (StrictMVar m a -> a -> Maybe String
forall (m :: * -> *) a. StrictMVar m a -> a -> Maybe String
getInvariant StrictMVar m a
v a
a) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

readMVar :: MonadMVar m => StrictMVar m a -> m a
readMVar :: forall (m :: * -> *) a. MonadMVar m => StrictMVar m a -> m a
readMVar StrictMVar m a
v = StrictMVar m a -> m a
forall (m :: * -> *) a. MonadMVar m => StrictMVar m a -> m a
Strict.readMVar (StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar StrictMVar m a
v)

swapMVar :: (HasCallStack, MonadMVar m) => StrictMVar m a -> a -> m a
swapMVar :: forall (m :: * -> *) a.
(HasCallStack, MonadMVar m) =>
StrictMVar m a -> a -> m a
swapMVar StrictMVar m a
v a
a = do
  oldValue <- StrictMVar m a -> a -> m a
forall (m :: * -> *) a. MonadMVar m => StrictMVar m a -> a -> m a
Strict.swapMVar (StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar StrictMVar m a
v) a
a
  checkInvariant (getInvariant v a) $ pure oldValue

tryTakeMVar :: MonadMVar m => StrictMVar m a -> m (Maybe a)
tryTakeMVar :: forall (m :: * -> *) a.
MonadMVar m =>
StrictMVar m a -> m (Maybe a)
tryTakeMVar StrictMVar m a
v = StrictMVar m a -> m (Maybe a)
forall (m :: * -> *) a.
MonadMVar m =>
StrictMVar m a -> m (Maybe a)
Strict.tryTakeMVar (StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar StrictMVar m a
v)

tryPutMVar :: (HasCallStack, MonadMVar m) => StrictMVar m a -> a -> m Bool
tryPutMVar :: forall (m :: * -> *) a.
(HasCallStack, MonadMVar m) =>
StrictMVar m a -> a -> m Bool
tryPutMVar StrictMVar m a
v a
a = do
  didPut <- StrictMVar m a -> a -> m Bool
forall (m :: * -> *) a.
MonadMVar m =>
StrictMVar m a -> a -> m Bool
Strict.tryPutMVar (StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar StrictMVar m a
v) a
a
  checkInvariant (getInvariant v a) $ pure didPut

isEmptyMVar :: MonadMVar m => StrictMVar m a -> m Bool
isEmptyMVar :: forall (m :: * -> *) a. MonadMVar m => StrictMVar m a -> m Bool
isEmptyMVar StrictMVar m a
v = StrictMVar m a -> m Bool
forall (m :: * -> *) a. MonadMVar m => StrictMVar m a -> m Bool
Strict.isEmptyMVar (StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar StrictMVar m a
v)

withMVar :: MonadMVar m => StrictMVar m a -> (a -> m b) -> m b
withMVar :: forall (m :: * -> *) a b.
MonadMVar m =>
StrictMVar m a -> (a -> m b) -> m b
withMVar StrictMVar m a
v = StrictMVar m a -> (a -> m b) -> m b
forall (m :: * -> *) a b.
MonadMVar m =>
StrictMVar m a -> (a -> m b) -> m b
Strict.withMVar (StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar StrictMVar m a
v)

withMVarMasked :: MonadMVar m => StrictMVar m a -> (a -> m b) -> m b
withMVarMasked :: forall (m :: * -> *) a b.
MonadMVar m =>
StrictMVar m a -> (a -> m b) -> m b
withMVarMasked StrictMVar m a
v = StrictMVar m a -> (a -> m b) -> m b
forall (m :: * -> *) a b.
MonadMVar m =>
StrictMVar m a -> (a -> m b) -> m b
Strict.withMVarMasked (StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar StrictMVar m a
v)

-- | 'modifyMVar_' is defined in terms of 'modifyMVar'.
modifyMVar_ ::
  (HasCallStack, MonadMVar m) =>
  StrictMVar m a ->
  (a -> m a) ->
  m ()
modifyMVar_ :: forall (m :: * -> *) a.
(HasCallStack, MonadMVar m) =>
StrictMVar m a -> (a -> m a) -> m ()
modifyMVar_ StrictMVar m a
v a -> m a
io = StrictMVar m a -> (a -> m (a, ())) -> m ()
forall (m :: * -> *) a b.
(HasCallStack, MonadMVar m) =>
StrictMVar m a -> (a -> m (a, b)) -> m b
modifyMVar StrictMVar m a
v a -> m (a, ())
io'
 where
  io' :: a -> m (a, ())
io' a
a = (,()) (a -> (a, ())) -> m a -> m (a, ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m a
io a
a

modifyMVar ::
  (HasCallStack, MonadMVar m) =>
  StrictMVar m a ->
  (a -> m (a, b)) ->
  m b
modifyMVar :: forall (m :: * -> *) a b.
(HasCallStack, MonadMVar m) =>
StrictMVar m a -> (a -> m (a, b)) -> m b
modifyMVar StrictMVar m a
v a -> m (a, b)
io = do
  (a', b) <- StrictMVar m a -> (a -> m (a, (a, b))) -> m (a, b)
forall (m :: * -> *) a b.
MonadMVar m =>
StrictMVar m a -> (a -> m (a, b)) -> m b
Strict.modifyMVar (StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar StrictMVar m a
v) a -> m (a, (a, b))
io'
  checkInvariant (getInvariant v a') $ pure b
 where
  io' :: a -> m (a, (a, b))
io' a
a = do
    (a', b) <- a -> m (a, b)
io a
a
    -- Returning @a'@ along with @b@ allows us to check the invariant /after/
    -- filling in the MVar.
    pure (a', (a', b))

-- | 'modifyMVarMasked_' is defined in terms of 'modifyMVarMasked'.
modifyMVarMasked_ ::
  (HasCallStack, MonadMVar m) =>
  StrictMVar m a ->
  (a -> m a) ->
  m ()
modifyMVarMasked_ :: forall (m :: * -> *) a.
(HasCallStack, MonadMVar m) =>
StrictMVar m a -> (a -> m a) -> m ()
modifyMVarMasked_ StrictMVar m a
v a -> m a
io = StrictMVar m a -> (a -> m (a, ())) -> m ()
forall (m :: * -> *) a b.
(HasCallStack, MonadMVar m) =>
StrictMVar m a -> (a -> m (a, b)) -> m b
modifyMVarMasked StrictMVar m a
v a -> m (a, ())
io'
 where
  io' :: a -> m (a, ())
io' a
a = (,()) (a -> (a, ())) -> m a -> m (a, ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m a
io a
a

modifyMVarMasked ::
  (HasCallStack, MonadMVar m) =>
  StrictMVar m a ->
  (a -> m (a, b)) ->
  m b
modifyMVarMasked :: forall (m :: * -> *) a b.
(HasCallStack, MonadMVar m) =>
StrictMVar m a -> (a -> m (a, b)) -> m b
modifyMVarMasked StrictMVar m a
v a -> m (a, b)
io = do
  (a', b) <- StrictMVar m a -> (a -> m (a, (a, b))) -> m (a, b)
forall (m :: * -> *) a b.
MonadMVar m =>
StrictMVar m a -> (a -> m (a, b)) -> m b
Strict.modifyMVarMasked (StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar StrictMVar m a
v) a -> m (a, (a, b))
io'
  checkInvariant (getInvariant v a') $ pure b
 where
  io' :: a -> m (a, (a, b))
io' a
a = do
    (a', b) <- a -> m (a, b)
io a
a
    -- Returning @a'@ along with @b@ allows us to check the invariant /after/
    -- filling in the MVar.
    pure (a', (a', b))

tryReadMVar :: MonadMVar m => StrictMVar m a -> m (Maybe a)
tryReadMVar :: forall (m :: * -> *) a.
MonadMVar m =>
StrictMVar m a -> m (Maybe a)
tryReadMVar StrictMVar m a
v = StrictMVar m a -> m (Maybe a)
forall (m :: * -> *) a.
MonadMVar m =>
StrictMVar m a -> m (Maybe a)
Strict.tryReadMVar (StrictMVar m a -> StrictMVar m a
forall (m :: * -> *) a. StrictMVar m a -> StrictMVar m a
mvar StrictMVar m a
v)

--
-- Dealing with invariants
--

-- | Check invariant (if enabled) before continuing
--
-- @checkInvariant mErr x@ is equal to @x@ if @mErr == Nothing@, and throws
-- an error @err@ if @mErr == Just err@.
--
-- This is exported so that other code that wants to conditionally check
-- invariants can reuse the same logic, rather than having to introduce new
-- per-package flags.
checkInvariant :: HasCallStack => Maybe String -> a -> a
getInvariant :: StrictMVar m a -> a -> Maybe String
mkStrictMVar :: (a -> Maybe String) -> Strict.StrictMVar m a -> StrictMVar m a

#if CHECK_MVAR_INVARIANTS
checkInvariant :: forall a. HasCallStack => Maybe String -> a -> a
checkInvariant Maybe String
Nothing    a
k = a
k
checkInvariant (Just String
err) a
_ = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ String
"StrictMVar invariant violation: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
err
getInvariant :: forall (m :: * -> *) a. StrictMVar m a -> a -> Maybe String
getInvariant StrictMVar {a -> Maybe String
invariant :: forall (m :: * -> *) a. StrictMVar m a -> a -> Maybe String
invariant :: a -> Maybe String
invariant} = a -> Maybe String
invariant
mkStrictMVar :: forall a (m :: * -> *).
(a -> Maybe String) -> StrictMVar m a -> StrictMVar m a
mkStrictMVar a -> Maybe String
invariant  StrictMVar m a
mvar        = StrictMVar {a -> Maybe String
invariant :: a -> Maybe String
invariant :: a -> Maybe String
invariant, StrictMVar m a
mvar :: StrictMVar m a
mvar :: StrictMVar m a
mvar}
#else
checkInvariant _err       k  = k
getInvariant _               = const Nothing
mkStrictMVar _invariant mvar = StrictMVar {mvar}
#endif