-- NOTE: Usage of 'trace' in 'spy' is accepted here.
{-# OPTIONS_GHC -Wno-deprecations #-}
-- NOTE: For the 'NominalDiffTime' CBOR instances below.
{-# OPTIONS_GHC -Wno-orphans #-}

module Hydra.Prelude (
  module Relude,
  MonadLabelledSTM,
  MonadSTM,
  STM,
  atomically,
  module Control.Monad.Class.MonadTime.SI,
  module Control.Monad.Class.MonadST,
  MonadAsync,
  Async,
  module Control.Monad.Class.MonadEventlog,
  module Control.Monad.Class.MonadTimer.SI,
  Control.Monad.Class.MonadFork.MonadFork,
  Control.Monad.Class.MonadFork.MonadThread,
  Control.Monad.Class.MonadFork.myThreadId,
  labelThisThread,
  module Control.Monad.Class.MonadThrow,
  module Control.Concurrent.Class.MonadSTM.TBQueue,
  module Control.Concurrent.Class.MonadSTM.TMVar,
  module Control.Concurrent.Class.MonadSTM.TQueue,
  module Control.Concurrent.Class.MonadSTM.TVar,
  StaticMap (..),
  DynamicMap (..),
  keys,
  elems,
  FromCBOR (..),
  ToCBOR (..),
  genericToCBOR,
  genericFromCBOR,
  FromJSON (..),
  ToJSON (..),
  encodePretty,
  padRight,
  Except,
  encodeBase16,
  decodeBase16,
  (?>),
  withFile,
  spy,
  spy',
  newLabelledTVar,
  newLabelledTVarIO,
  newLabelledEmptyTMVar,
  newLabelledTQueueIO,
  newLabelledEmptyTMVarIO,
  concurrentlyLabelled,
  concurrentlyLabelled_,
  asyncLabelled,
  raceLabelled,
  raceLabelled_,
  withAsyncLabelled,
  newLabelledTQueue,
  newLabelledTBQueue,
  newLabelledTBQueueIO,
) where

import Cardano.Binary (
  FromCBOR (..),
  ToCBOR (..),
 )
import Control.Concurrent.Class.MonadSTM (MonadLabelledSTM (..), MonadSTM (..))
import Control.Concurrent.Class.MonadSTM.TBQueue (TBQueue)
import Control.Concurrent.Class.MonadSTM.TMVar (TMVar)
import Control.Concurrent.Class.MonadSTM.TQueue (TQueue)
import Control.Concurrent.Class.MonadSTM.TVar (TVar, readTVar)
import Control.Exception (IOException)
import Control.Monad.Class.MonadAsync (
  Async,
  MonadAsync (async, concurrently, race, withAsync),
 )
import Control.Monad.Class.MonadEventlog (
  MonadEventlog,
 )
import Control.Monad.Class.MonadFork (MonadFork, MonadThread, labelThisThread, myThreadId)
import Control.Monad.Class.MonadST (
  MonadST,
 )
import Control.Monad.Class.MonadSTM ()
import Control.Monad.Class.MonadThrow (
  MonadCatch (..),
  MonadEvaluate (..),
  MonadMask (..),
  MonadThrow (..),
 )
import Control.Monad.Class.MonadTime.SI (
  DiffTime,
  MonadMonotonicTime (..),
  MonadTime (..),
  NominalDiffTime,
  Time (..),
  UTCTime,
  addTime,
  addUTCTime,
  diffTime,
  diffUTCTime,
 )
import Control.Monad.Class.MonadTimer.SI (
  MonadDelay (..),
  MonadTimer (..),
 )
import Control.Monad.Trans.Except (Except)
import Data.Aeson (
  FromJSON (..),
  ToJSON (..),
 )
import Data.Aeson.Encode.Pretty (
  encodePretty,
 )
import Data.ByteString.Base16 qualified as Base16
import Data.Fixed (Fixed (..))
import Data.Text qualified as T
import Data.Time.Clock (nominalDiffTimeToSeconds, secondsToNominalDiffTime)
import Hydra.CBOR.Generic (genericFromCBOR, genericToCBOR)
import Relude hiding (
  MVar,
  Nat,
  STM,
  TMVar,
  TVar,
  atomically,
  catchSTM,
  isEmptyTMVar,
  mkWeakTMVar,
  modifyTVar',
  newEmptyMVar,
  newEmptyTMVar,
  newEmptyTMVarIO,
  newMVar,
  newTMVar,
  newTMVarIO,
  newTVar,
  newTVarIO,
  putMVar,
  putTMVar,
  readMVar,
  readTMVar,
  readTVar,
  readTVarIO,
  swapMVar,
  swapTMVar,
  takeMVar,
  takeTMVar,
  throwSTM,
  traceM,
  tryPutMVar,
  tryPutTMVar,
  tryReadMVar,
  tryReadTMVar,
  tryTakeMVar,
  tryTakeTMVar,
  withFile,
  writeTVar,
 )
import Relude.Extra.Map (
  DynamicMap (..),
  StaticMap (..),
  elems,
  keys,
 )
import System.IO qualified
import Text.Pretty.Simple (pShow)

-- Orphan CBOR instances: cardano-binary provides 'UTCTime' instances, but
-- lacks 'NominalDiffTime'. Encoded as integer picoseconds, which is lossless
-- since 'NominalDiffTime' has fixed picosecond resolution.
instance ToCBOR NominalDiffTime where
  toCBOR :: NominalDiffTime -> Encoding
toCBOR NominalDiffTime
ndt = case NominalDiffTime -> Pico
nominalDiffTimeToSeconds NominalDiffTime
ndt of MkFixed Integer
i -> Integer -> Encoding
forall a. ToCBOR a => a -> Encoding
toCBOR Integer
i

instance FromCBOR NominalDiffTime where
  fromCBOR :: forall s. Decoder s NominalDiffTime
fromCBOR = Pico -> NominalDiffTime
secondsToNominalDiffTime (Pico -> NominalDiffTime)
-> (Integer -> Pico) -> Integer -> NominalDiffTime
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Pico
forall k (a :: k). Integer -> Fixed a
MkFixed (Integer -> NominalDiffTime)
-> Decoder s Integer -> Decoder s NominalDiffTime
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Decoder s Integer
forall s. Decoder s Integer
forall a s. FromCBOR a => Decoder s a
fromCBOR

-- | Pad a text-string to right with the given character until it reaches the given
-- length.
--
-- NOTE: Truncate the string if longer than the given length.
padRight :: Char -> Int -> Text -> Text
padRight :: Char -> Int -> Text -> Text
padRight Char
c Int
n Text
str = Int -> Text -> Text
T.take Int
n (Text
str Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text -> Text
T.replicate Int
n (Char -> Text
T.singleton Char
c))

-- | Encode some bytes to hex-encoded text.
--
-- >>> encodeBase16 "ab"
-- "4142"
encodeBase16 :: ByteString -> Text
encodeBase16 :: ByteString -> Text
encodeBase16 =
  ByteString -> Text
forall a b. ConvertUtf8 a b => b -> a
decodeUtf8 (ByteString -> Text)
-> (ByteString -> ByteString) -> ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
Base16.encode

-- | Decode some hex-encoded text string to raw bytes.
--
-- >>> decodeBase16 "dflkgjdjgdh"
-- Left "Not base 16"
decodeBase16 :: MonadFail f => Text -> f ByteString
decodeBase16 :: forall (f :: * -> *). MonadFail f => Text -> f ByteString
decodeBase16 =
  (String -> f ByteString)
-> (ByteString -> f ByteString)
-> Either String ByteString
-> f ByteString
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either String -> f ByteString
forall a. String -> f a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail ByteString -> f ByteString
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either String ByteString -> f ByteString)
-> (Text -> Either String ByteString) -> Text -> f ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Either String ByteString
Base16.decode (ByteString -> Either String ByteString)
-> (Text -> ByteString) -> Text -> Either String ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
forall a b. ConvertUtf8 a b => a -> b
encodeUtf8

infixl 4 ?>

-- | If 'Nothing' use given 'e' as 'Left'. Infix version of `maybeToEither`.
(?>) :: Maybe a -> e -> Either e a
?> :: forall a e. Maybe a -> e -> Either e a
(?>) Maybe a
m e
e =
  case Maybe a
m of
    Maybe a
Nothing -> e -> Either e a
forall a b. a -> Either a b
Left e
e
    Just a
a -> a -> Either e a
forall a b. b -> Either a b
Right a
a

-- | Like 'withFile' from 'base', but without annotating errors originating from
-- enclosed action.
--
-- XXX: This should be fixed upstream in 'base'.
withFile :: FilePath -> IOMode -> (Handle -> IO a) -> IO a
withFile :: forall a. String -> IOMode -> (Handle -> IO a) -> IO a
withFile String
fp IOMode
mode Handle -> IO a
action =
  String
-> IOMode
-> (Handle -> IO (Either IOException a))
-> IO (Either IOException a)
forall a. String -> IOMode -> (Handle -> IO a) -> IO a
System.IO.withFile String
fp IOMode
mode (IO a -> IO (Either IOException a)
forall e a. Exception e => IO a -> IO (Either e a)
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
try (IO a -> IO (Either IOException a))
-> (Handle -> IO a) -> Handle -> IO (Either IOException a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Handle -> IO a
action) IO (Either IOException a) -> (Either IOException a -> IO a) -> IO a
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Left (IOException
e :: IOException) -> IOException -> IO a
forall e a. Exception e => e -> IO a
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwIO IOException
e
    Right a
x -> a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x

-- | Like 'traceShow', but with pretty printing of the value.
{-# WARNING spy "Use for debugging purposes only" #-}
spy :: Show a => a -> a
spy :: forall a. Show a => a -> a
spy a
a = String -> a -> a
forall a. String -> a -> a
trace (Text -> String
forall a. ToString a => a -> String
toString (Text -> String) -> Text -> String
forall a b. (a -> b) -> a -> b
$ a -> Text
forall a. Show a => a -> Text
pShow a
a) a
a

-- | Like 'spy' but prefixed with a label.
{-# WARNING spy' "Use for debugging purposes only" #-}
spy' :: Show a => String -> a -> a
spy' :: forall a. Show a => String -> a -> a
spy' String
msg a
a = String -> a -> a
forall a. String -> a -> a
trace (String
msg String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
forall a. ToString a => a -> String
toString (a -> Text
forall a. Show a => a -> Text
pShow a
a)) a
a

-- * Helpers for labeling TVar

newLabelledTVar :: MonadLabelledSTM m => String -> a -> STM m (TVar m a)
newLabelledTVar :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> a -> STM m (TVar m a)
newLabelledTVar String
lbl a
val = do
  TVar m a
tv <- a -> STM m (TVar m a)
forall a. a -> STM m (TVar m a)
forall (m :: * -> *) a. MonadSTM m => a -> STM m (TVar m a)
newTVar a
val
  TVar m a -> String -> STM m ()
forall a. TVar m a -> String -> STM m ()
forall (m :: * -> *) a.
MonadLabelledSTM m =>
TVar m a -> String -> STM m ()
labelTVar TVar m a
tv String
lbl
  TVar m a -> STM m (TVar m a)
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TVar m a
tv

newLabelledTVarIO :: MonadLabelledSTM m => String -> a -> m (TVar m a)
newLabelledTVarIO :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> a -> m (TVar m a)
newLabelledTVarIO = (STM m (TVar m a) -> m (TVar m a)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically .) ((a -> STM m (TVar m a)) -> a -> m (TVar m a))
-> (String -> a -> STM m (TVar m a)) -> String -> a -> m (TVar m a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> a -> STM m (TVar m a)
forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> a -> STM m (TVar m a)
newLabelledTVar

-- * Helpers for labeling TMVar

newLabelledEmptyTMVar :: MonadLabelledSTM m => String -> STM m (TMVar m a)
newLabelledEmptyTMVar :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> STM m (TMVar m a)
newLabelledEmptyTMVar String
lbl = do
  TMVar m a
tmv <- STM m (TMVar m a)
forall a. STM m (TMVar m a)
forall (m :: * -> *) a. MonadSTM m => STM m (TMVar m a)
newEmptyTMVar
  TMVar m a -> String -> STM m ()
forall a. TMVar m a -> String -> STM m ()
forall (m :: * -> *) a.
MonadLabelledSTM m =>
TMVar m a -> String -> STM m ()
labelTMVar TMVar m a
tmv String
lbl
  TMVar m a -> STM m (TMVar m a)
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TMVar m a
tmv

newLabelledEmptyTMVarIO :: MonadLabelledSTM m => String -> m (TMVar m a)
newLabelledEmptyTMVarIO :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> m (TMVar m a)
newLabelledEmptyTMVarIO = STM m (TMVar m a) -> m (TMVar m a)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (TMVar m a) -> m (TMVar m a))
-> (String -> STM m (TMVar m a)) -> String -> m (TMVar m a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> STM m (TMVar m a)
forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> STM m (TMVar m a)
newLabelledEmptyTMVar

-- * Helpers for labeling TQueue

newLabelledTQueue :: MonadLabelledSTM m => String -> STM m (TQueue m a)
newLabelledTQueue :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> STM m (TQueue m a)
newLabelledTQueue String
lbl = do
  TQueue m a
q <- STM m (TQueue m a)
forall a. STM m (TQueue m a)
forall (m :: * -> *) a. MonadSTM m => STM m (TQueue m a)
newTQueue
  TQueue m a -> String -> STM m ()
forall a. TQueue m a -> String -> STM m ()
forall (m :: * -> *) a.
MonadLabelledSTM m =>
TQueue m a -> String -> STM m ()
labelTQueue TQueue m a
q String
lbl
  TQueue m a -> STM m (TQueue m a)
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TQueue m a
q

newLabelledTQueueIO :: MonadLabelledSTM m => String -> m (TQueue m a)
newLabelledTQueueIO :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> m (TQueue m a)
newLabelledTQueueIO = STM m (TQueue m a) -> m (TQueue m a)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (TQueue m a) -> m (TQueue m a))
-> (String -> STM m (TQueue m a)) -> String -> m (TQueue m a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> STM m (TQueue m a)
forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> STM m (TQueue m a)
newLabelledTQueue

-- * Helpers for labeling TBQueue

newLabelledTBQueue :: MonadLabelledSTM m => String -> Natural -> STM m (TBQueue m a)
newLabelledTBQueue :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> Natural -> STM m (TBQueue m a)
newLabelledTBQueue String
lbl Natural
capacity = do
  TBQueue m a
bq <- Natural -> STM m (TBQueue m a)
forall a. Natural -> STM m (TBQueue m a)
forall (m :: * -> *) a.
MonadSTM m =>
Natural -> STM m (TBQueue m a)
newTBQueue Natural
capacity
  TBQueue m a -> String -> STM m ()
forall a. TBQueue m a -> String -> STM m ()
forall (m :: * -> *) a.
MonadLabelledSTM m =>
TBQueue m a -> String -> STM m ()
labelTBQueue TBQueue m a
bq String
lbl
  TBQueue m a -> STM m (TBQueue m a)
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TBQueue m a
bq

newLabelledTBQueueIO :: MonadLabelledSTM m => String -> Natural -> m (TBQueue m a)
newLabelledTBQueueIO :: forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> Natural -> m (TBQueue m a)
newLabelledTBQueueIO = (STM m (TBQueue m a) -> m (TBQueue m a)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically .) ((Natural -> STM m (TBQueue m a)) -> Natural -> m (TBQueue m a))
-> (String -> Natural -> STM m (TBQueue m a))
-> String
-> Natural
-> m (TBQueue m a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Natural -> STM m (TBQueue m a)
forall (m :: * -> *) a.
MonadLabelledSTM m =>
String -> Natural -> STM m (TBQueue m a)
newLabelledTBQueue

-- * Helpers for labeling Threads

raceLabelled :: MonadAsync m => (String, m a) -> (String, m b) -> m (Either a b)
raceLabelled :: forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (String, m b) -> m (Either a b)
raceLabelled (String
lblA, m a
mA) (String
lblB, m b
mB) =
  m a -> m b -> m (Either a b)
forall a b. m a -> m b -> m (Either a b)
forall (m :: * -> *) a b.
MonadAsync m =>
m a -> m b -> m (Either a b)
race
    (String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
lblA m () -> m a -> m a
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m a
mA)
    (String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
lblB m () -> m b -> m b
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m b
mB)

raceLabelled_ :: MonadAsync m => (String, m a) -> (String, m b) -> m ()
raceLabelled_ :: forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (String, m b) -> m ()
raceLabelled_ = (m (Either a b) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void .) (((String, m b) -> m (Either a b)) -> (String, m b) -> m ())
-> ((String, m a) -> (String, m b) -> m (Either a b))
-> (String, m a)
-> (String, m b)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String, m a) -> (String, m b) -> m (Either a b)
forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (String, m b) -> m (Either a b)
raceLabelled

withAsyncLabelled :: MonadAsync m => (String, m a) -> (Async m a -> m b) -> m b
withAsyncLabelled :: forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (Async m a -> m b) -> m b
withAsyncLabelled (String
lbl, m a
ma) = m a -> (Async m a -> m b) -> m b
forall a b. m a -> (Async m a -> m b) -> m b
forall (m :: * -> *) a b.
MonadAsync m =>
m a -> (Async m a -> m b) -> m b
withAsync (String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
lbl m () -> m a -> m a
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m a
ma)

concurrentlyLabelled :: MonadAsync m => (String, m a) -> (String, m b) -> m (a, b)
concurrentlyLabelled :: forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (String, m b) -> m (a, b)
concurrentlyLabelled (String
lblA, m a
mA) (String
lblB, m b
mB) =
  m a -> m b -> m (a, b)
forall a b. m a -> m b -> m (a, b)
forall (m :: * -> *) a b. MonadAsync m => m a -> m b -> m (a, b)
concurrently
    (String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
lblA m () -> m a -> m a
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m a
mA)
    (String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
lblB m () -> m b -> m b
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m b
mB)

concurrentlyLabelled_ :: MonadAsync m => (String, m a) -> (String, m b) -> m ()
concurrentlyLabelled_ :: forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (String, m b) -> m ()
concurrentlyLabelled_ = (m (a, b) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void .) (((String, m b) -> m (a, b)) -> (String, m b) -> m ())
-> ((String, m a) -> (String, m b) -> m (a, b))
-> (String, m a)
-> (String, m b)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String, m a) -> (String, m b) -> m (a, b)
forall (m :: * -> *) a b.
MonadAsync m =>
(String, m a) -> (String, m b) -> m (a, b)
concurrentlyLabelled

asyncLabelled :: MonadAsync m => String -> m a -> m (Async m a)
asyncLabelled :: forall (m :: * -> *) a.
MonadAsync m =>
String -> m a -> m (Async m a)
asyncLabelled String
lbl m a
mA = m a -> m (Async m a)
forall a. m a -> m (Async m a)
forall (m :: * -> *) a. MonadAsync m => m a -> m (Async m a)
async (m a -> m (Async m a)) -> m a -> m (Async m a)
forall a b. (a -> b) -> a -> b
$ String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
lbl m () -> m a -> m a
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m a
mA