{-# LANGUAGE DeriveAnyClass             #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE DerivingStrategies         #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}

module Ouroboros.Consensus.Util.STM (
    -- * Misc
    blockUntilChanged
  , onEachChange
  , runWhenJust
  , blockUntilJust
  , blockUntilAllJust
  , Fingerprint (..)
  , WithFingerprint (..)
    -- * Simulate various monad stacks in STM
  , Sim(..)
  , simId
  , simStateT
  ) where

import           Control.Monad.State
import           Data.Void
import           Data.Word (Word64)
import           GHC.Generics (Generic)
import           GHC.Stack

import           Ouroboros.Consensus.Util.IOLike
import           Ouroboros.Consensus.Util.ResourceRegistry

{-------------------------------------------------------------------------------
  Misc
-------------------------------------------------------------------------------}

-- | Wait until the TVar changed
blockUntilChanged :: forall stm a b. (MonadSTMTx stm, Eq b)
                  => (a -> b) -> b -> stm a -> stm (a, b)
blockUntilChanged :: (a -> b) -> b -> stm a -> stm (a, b)
blockUntilChanged a -> b
f b
b stm a
getA = do
    a
a <- stm a
getA
    let b' :: b
b' = a -> b
f a
a
    if b
b' b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
b
      then stm (a, b)
forall (stm :: * -> *) a. MonadSTMTx stm => stm a
retry
      else (a, b) -> stm (a, b)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, b
b')

-- | Spawn a new thread that runs an action each time an STM value changes.
--
-- NOTE: STM does not guarantee that 'onEachChange' will /literally/ be called
-- on /every/ change: when the system is under heavy load, some updates may
-- be missed.
--
-- The thread will be linked to the registry.
onEachChange :: forall m a b. (IOLike m, Eq b, HasCallStack)
             => ResourceRegistry m
             -> String    -- ^ Label for the thread
             -> (a -> b)  -- ^ Obtain a fingerprint
             -> Maybe b   -- ^ Optional initial fingerprint
                          -- If 'Nothing', the action is executed once
                          -- immediately to obtain the initial fingerprint.
             -> STM m a
             -> (a -> m ())
             -> m (Thread m Void)
onEachChange :: ResourceRegistry m
-> String
-> (a -> b)
-> Maybe b
-> STM m a
-> (a -> m ())
-> m (Thread m Void)
onEachChange ResourceRegistry m
registry String
label a -> b
f Maybe b
mbInitB STM m a
getA a -> m ()
notify =
    ResourceRegistry m -> String -> m Void -> m (Thread m Void)
forall (m :: * -> *) a.
(IOLike m, HasCallStack) =>
ResourceRegistry m -> String -> m a -> m (Thread m a)
forkLinkedThread ResourceRegistry m
registry String
label m Void
body
  where
    body :: m Void
    body :: m Void
body = do
        b
initB <- case Maybe b
mbInitB of
          Just b
initB -> b -> m b
forall (m :: * -> *) a. Monad m => a -> m a
return b
initB
          Maybe b
Nothing    -> do
            a
a <- STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically STM m a
getA
            a -> m ()
notify a
a
            b -> m b
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> m b) -> b -> m b
forall a b. (a -> b) -> a -> b
$ a -> b
f a
a
        b -> m Void
loop b
initB

    loop :: b -> m Void
    loop :: b -> m Void
loop b
b = do
      (a
a, b
b') <- STM m (a, b) -> m (a, b)
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (a, b) -> m (a, b)) -> STM m (a, b) -> m (a, b)
forall a b. (a -> b) -> a -> b
$ (a -> b) -> b -> STM m a -> STM m (a, b)
forall (stm :: * -> *) a b.
(MonadSTMTx stm, Eq b) =>
(a -> b) -> b -> stm a -> stm (a, b)
blockUntilChanged a -> b
f b
b STM m a
getA
      a -> m ()
notify a
a
      b -> m Void
loop b
b'

-- | Spawn a new thread that waits for an STM value to become 'Just'
--
-- The thread will be linked to the registry.
runWhenJust :: IOLike m
            => ResourceRegistry m
            -> String  -- ^ Label for the thread
            -> STM m (Maybe a)
            -> (a -> m ())
            -> m ()
runWhenJust :: ResourceRegistry m
-> String -> STM m (Maybe a) -> (a -> m ()) -> m ()
runWhenJust ResourceRegistry m
registry String
label STM m (Maybe a)
getMaybeA a -> m ()
action =
    m (Thread m ()) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m (Thread m ()) -> m ()) -> m (Thread m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ ResourceRegistry m -> String -> m () -> m (Thread m ())
forall (m :: * -> *) a.
(IOLike m, HasCallStack) =>
ResourceRegistry m -> String -> m a -> m (Thread m a)
forkLinkedThread ResourceRegistry m
registry String
label (m () -> m (Thread m ())) -> m () -> m (Thread m ())
forall a b. (a -> b) -> a -> b
$
      a -> m ()
action (a -> m ()) -> m a -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (Maybe a) -> STM m a
forall (stm :: * -> *) a. MonadSTMTx stm => stm (Maybe a) -> stm a
blockUntilJust STM m (Maybe a)
getMaybeA)

blockUntilJust :: MonadSTMTx stm => stm (Maybe a) -> stm a
blockUntilJust :: stm (Maybe a) -> stm a
blockUntilJust stm (Maybe a)
getMaybeA = do
    Maybe a
ma <- stm (Maybe a)
getMaybeA
    case Maybe a
ma of
      Maybe a
Nothing -> stm a
forall (stm :: * -> *) a. MonadSTMTx stm => stm a
retry
      Just a
a  -> a -> stm a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a

blockUntilAllJust :: MonadSTMTx stm => [stm (Maybe a)] -> stm [a]
blockUntilAllJust :: [stm (Maybe a)] -> stm [a]
blockUntilAllJust = (stm (Maybe a) -> stm a) -> [stm (Maybe a)] -> stm [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM stm (Maybe a) -> stm a
forall (stm :: * -> *) a. MonadSTMTx stm => stm (Maybe a) -> stm a
blockUntilJust

-- | Simple type that can be used to indicate something in a @TVar@ is
-- changed.
newtype Fingerprint = Fingerprint Word64
  deriving stock    (Int -> Fingerprint -> ShowS
[Fingerprint] -> ShowS
Fingerprint -> String
(Int -> Fingerprint -> ShowS)
-> (Fingerprint -> String)
-> ([Fingerprint] -> ShowS)
-> Show Fingerprint
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Fingerprint] -> ShowS
$cshowList :: [Fingerprint] -> ShowS
show :: Fingerprint -> String
$cshow :: Fingerprint -> String
showsPrec :: Int -> Fingerprint -> ShowS
$cshowsPrec :: Int -> Fingerprint -> ShowS
Show, Fingerprint -> Fingerprint -> Bool
(Fingerprint -> Fingerprint -> Bool)
-> (Fingerprint -> Fingerprint -> Bool) -> Eq Fingerprint
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Fingerprint -> Fingerprint -> Bool
$c/= :: Fingerprint -> Fingerprint -> Bool
== :: Fingerprint -> Fingerprint -> Bool
$c== :: Fingerprint -> Fingerprint -> Bool
Eq, (forall x. Fingerprint -> Rep Fingerprint x)
-> (forall x. Rep Fingerprint x -> Fingerprint)
-> Generic Fingerprint
forall x. Rep Fingerprint x -> Fingerprint
forall x. Fingerprint -> Rep Fingerprint x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Fingerprint x -> Fingerprint
$cfrom :: forall x. Fingerprint -> Rep Fingerprint x
Generic)
  deriving newtype  (Int -> Fingerprint
Fingerprint -> Int
Fingerprint -> [Fingerprint]
Fingerprint -> Fingerprint
Fingerprint -> Fingerprint -> [Fingerprint]
Fingerprint -> Fingerprint -> Fingerprint -> [Fingerprint]
(Fingerprint -> Fingerprint)
-> (Fingerprint -> Fingerprint)
-> (Int -> Fingerprint)
-> (Fingerprint -> Int)
-> (Fingerprint -> [Fingerprint])
-> (Fingerprint -> Fingerprint -> [Fingerprint])
-> (Fingerprint -> Fingerprint -> [Fingerprint])
-> (Fingerprint -> Fingerprint -> Fingerprint -> [Fingerprint])
-> Enum Fingerprint
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
enumFromThenTo :: Fingerprint -> Fingerprint -> Fingerprint -> [Fingerprint]
$cenumFromThenTo :: Fingerprint -> Fingerprint -> Fingerprint -> [Fingerprint]
enumFromTo :: Fingerprint -> Fingerprint -> [Fingerprint]
$cenumFromTo :: Fingerprint -> Fingerprint -> [Fingerprint]
enumFromThen :: Fingerprint -> Fingerprint -> [Fingerprint]
$cenumFromThen :: Fingerprint -> Fingerprint -> [Fingerprint]
enumFrom :: Fingerprint -> [Fingerprint]
$cenumFrom :: Fingerprint -> [Fingerprint]
fromEnum :: Fingerprint -> Int
$cfromEnum :: Fingerprint -> Int
toEnum :: Int -> Fingerprint
$ctoEnum :: Int -> Fingerprint
pred :: Fingerprint -> Fingerprint
$cpred :: Fingerprint -> Fingerprint
succ :: Fingerprint -> Fingerprint
$csucc :: Fingerprint -> Fingerprint
Enum)
  deriving anyclass (Context -> Fingerprint -> IO (Maybe ThunkInfo)
Proxy Fingerprint -> String
(Context -> Fingerprint -> IO (Maybe ThunkInfo))
-> (Context -> Fingerprint -> IO (Maybe ThunkInfo))
-> (Proxy Fingerprint -> String)
-> NoThunks Fingerprint
forall a.
(Context -> a -> IO (Maybe ThunkInfo))
-> (Context -> a -> IO (Maybe ThunkInfo))
-> (Proxy a -> String)
-> NoThunks a
showTypeOf :: Proxy Fingerprint -> String
$cshowTypeOf :: Proxy Fingerprint -> String
wNoThunks :: Context -> Fingerprint -> IO (Maybe ThunkInfo)
$cwNoThunks :: Context -> Fingerprint -> IO (Maybe ThunkInfo)
noThunks :: Context -> Fingerprint -> IO (Maybe ThunkInfo)
$cnoThunks :: Context -> Fingerprint -> IO (Maybe ThunkInfo)
NoThunks)

-- | Store a value together with its fingerprint.
data WithFingerprint a = WithFingerprint
  { WithFingerprint a -> a
forgetFingerprint :: !a
  , WithFingerprint a -> Fingerprint
getFingerprint    :: !Fingerprint
  } deriving (Int -> WithFingerprint a -> ShowS
[WithFingerprint a] -> ShowS
WithFingerprint a -> String
(Int -> WithFingerprint a -> ShowS)
-> (WithFingerprint a -> String)
-> ([WithFingerprint a] -> ShowS)
-> Show (WithFingerprint a)
forall a. Show a => Int -> WithFingerprint a -> ShowS
forall a. Show a => [WithFingerprint a] -> ShowS
forall a. Show a => WithFingerprint a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [WithFingerprint a] -> ShowS
$cshowList :: forall a. Show a => [WithFingerprint a] -> ShowS
show :: WithFingerprint a -> String
$cshow :: forall a. Show a => WithFingerprint a -> String
showsPrec :: Int -> WithFingerprint a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> WithFingerprint a -> ShowS
Show, WithFingerprint a -> WithFingerprint a -> Bool
(WithFingerprint a -> WithFingerprint a -> Bool)
-> (WithFingerprint a -> WithFingerprint a -> Bool)
-> Eq (WithFingerprint a)
forall a. Eq a => WithFingerprint a -> WithFingerprint a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: WithFingerprint a -> WithFingerprint a -> Bool
$c/= :: forall a. Eq a => WithFingerprint a -> WithFingerprint a -> Bool
== :: WithFingerprint a -> WithFingerprint a -> Bool
$c== :: forall a. Eq a => WithFingerprint a -> WithFingerprint a -> Bool
Eq, a -> WithFingerprint b -> WithFingerprint a
(a -> b) -> WithFingerprint a -> WithFingerprint b
(forall a b. (a -> b) -> WithFingerprint a -> WithFingerprint b)
-> (forall a b. a -> WithFingerprint b -> WithFingerprint a)
-> Functor WithFingerprint
forall a b. a -> WithFingerprint b -> WithFingerprint a
forall a b. (a -> b) -> WithFingerprint a -> WithFingerprint b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> WithFingerprint b -> WithFingerprint a
$c<$ :: forall a b. a -> WithFingerprint b -> WithFingerprint a
fmap :: (a -> b) -> WithFingerprint a -> WithFingerprint b
$cfmap :: forall a b. (a -> b) -> WithFingerprint a -> WithFingerprint b
Functor, (forall x. WithFingerprint a -> Rep (WithFingerprint a) x)
-> (forall x. Rep (WithFingerprint a) x -> WithFingerprint a)
-> Generic (WithFingerprint a)
forall x. Rep (WithFingerprint a) x -> WithFingerprint a
forall x. WithFingerprint a -> Rep (WithFingerprint a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (WithFingerprint a) x -> WithFingerprint a
forall a x. WithFingerprint a -> Rep (WithFingerprint a) x
$cto :: forall a x. Rep (WithFingerprint a) x -> WithFingerprint a
$cfrom :: forall a x. WithFingerprint a -> Rep (WithFingerprint a) x
Generic, Context -> WithFingerprint a -> IO (Maybe ThunkInfo)
Proxy (WithFingerprint a) -> String
(Context -> WithFingerprint a -> IO (Maybe ThunkInfo))
-> (Context -> WithFingerprint a -> IO (Maybe ThunkInfo))
-> (Proxy (WithFingerprint a) -> String)
-> NoThunks (WithFingerprint a)
forall a.
NoThunks a =>
Context -> WithFingerprint a -> IO (Maybe ThunkInfo)
forall a. NoThunks a => Proxy (WithFingerprint a) -> String
forall a.
(Context -> a -> IO (Maybe ThunkInfo))
-> (Context -> a -> IO (Maybe ThunkInfo))
-> (Proxy a -> String)
-> NoThunks a
showTypeOf :: Proxy (WithFingerprint a) -> String
$cshowTypeOf :: forall a. NoThunks a => Proxy (WithFingerprint a) -> String
wNoThunks :: Context -> WithFingerprint a -> IO (Maybe ThunkInfo)
$cwNoThunks :: forall a.
NoThunks a =>
Context -> WithFingerprint a -> IO (Maybe ThunkInfo)
noThunks :: Context -> WithFingerprint a -> IO (Maybe ThunkInfo)
$cnoThunks :: forall a.
NoThunks a =>
Context -> WithFingerprint a -> IO (Maybe ThunkInfo)
NoThunks)

{-------------------------------------------------------------------------------
  Simulate monad stacks
-------------------------------------------------------------------------------}

newtype Sim n m = Sim { Sim n m -> forall a. n a -> STM m a
runSim :: forall a. n a -> STM m a }

simId :: Sim (STM m) m
simId :: Sim (STM m) m
simId = (forall a. STM m a -> STM m a) -> Sim (STM m) m
forall (n :: * -> *) (m :: * -> *).
(forall a. n a -> STM m a) -> Sim n m
Sim forall a. a -> a
forall a. STM m a -> STM m a
id

simStateT :: IOLike m => StrictTVar m st -> Sim n m -> Sim (StateT st n) m
simStateT :: StrictTVar m st -> Sim n m -> Sim (StateT st n) m
simStateT StrictTVar m st
stVar (Sim forall a. n a -> STM m a
k) = (forall a. StateT st n a -> STM m a) -> Sim (StateT st n) m
forall (n :: * -> *) (m :: * -> *).
(forall a. n a -> STM m a) -> Sim n m
Sim ((forall a. StateT st n a -> STM m a) -> Sim (StateT st n) m)
-> (forall a. StateT st n a -> STM m a) -> Sim (StateT st n) m
forall a b. (a -> b) -> a -> b
$ \(StateT f) -> do
    st
st       <- StrictTVar m st -> STM m st
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> STM m a
readTVar StrictTVar m st
stVar
    (a
a, st
st') <- n (a, st) -> STM m (a, st)
forall a. n a -> STM m a
k (st -> n (a, st)
f st
st)
    StrictTVar m st -> st -> STM m ()
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
StrictTVar m a -> a -> STM m ()
writeTVar StrictTVar m st
stVar st
st'
    a -> STM m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a