-- NOTE: Usage of 'trace' in 'spy' is accepted here.
{-# OPTIONS_GHC -Wno-deprecations #-}

module Hydra.Prelude (
  module Relude,
  module Control.Monad.Class.MonadSTM,
  module Control.Monad.Class.MonadTime.SI,
  module Control.Monad.Class.MonadST,
  module Control.Monad.Class.MonadAsync,
  module Control.Monad.Class.MonadEventlog,
  module Control.Monad.Class.MonadTimer.SI,
  module Control.Monad.Class.MonadFork,
  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 (..),
  FromJSON (..),
  ToJSON (..),
  encodePretty,
  Gen,
  Arbitrary (..),
  genericArbitrary,
  genericShrink,
  generateWith,
  shrinkListAggressively,
  reasonablySized,
  ReasonablySized (..),
  padLeft,
  padRight,
  Except,
  decodeBase16,
  (?>),
  withFile,
  spy,
  spy',
) where

import Cardano.Binary (
  FromCBOR (..),
  ToCBOR (..),
 )
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 (
  MonadAsync (concurrently, concurrently_, race, race_, withAsync),
 )
import Control.Monad.Class.MonadEventlog (
  MonadEventlog,
 )
import Control.Monad.Class.MonadFork (
  MonadFork,
  MonadThread,
 )
import Control.Monad.Class.MonadST (
  MonadST,
 )
import Control.Monad.Class.MonadSTM (
  MonadSTM,
  STM,
  atomically,
 )
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.Text qualified as T
import GHC.Generics (Rep)
import Generic.Random qualified as Random
import Generic.Random.Internal.Generic qualified as Random
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 Test.QuickCheck (
  Arbitrary (..),
  Gen,
  genericShrink,
  scale,
 )
import Test.QuickCheck.Gen (Gen (..))
import Test.QuickCheck.Random (mkQCGen)
import Text.Pretty.Simple (pShow)

-- | Provides a sensible way of automatically deriving generic 'Arbitrary'
-- instances for data-types. In the case where more advanced or tailored
-- generators are needed, custom hand-written generators should be used with
-- functions such as `forAll` or `forAllShrink`.
genericArbitrary ::
  ( Generic a
  , Random.GA Random.UnsizedOpts (Rep a)
  , Random.UniformWeight (Random.Weights_ (Rep a))
  ) =>
  Gen a
genericArbitrary :: forall a.
(Generic a, GA UnsizedOpts (Rep a),
 UniformWeight (Weights_ (Rep a))) =>
Gen a
genericArbitrary =
  Weights a -> Gen a
forall a. GArbitrary UnsizedOpts a => Weights a -> Gen a
Random.genericArbitrary Weights a
forall a. UniformWeight_ (Rep a) => Weights a
Random.uniform

-- | A seeded, deterministic, generator
generateWith :: Gen a -> Int -> a
generateWith :: forall a. Gen a -> Int -> a
generateWith (MkGen QCGen -> Int -> a
runGen) Int
seed =
  QCGen -> Int -> a
runGen (Int -> QCGen
mkQCGen Int
seed) Int
30

-- | Like 'shrinkList', but more aggressive :)
--
-- Useful for shrinking large nested Map or Lists where the shrinker "don't have
-- time" to go through many cases.
shrinkListAggressively :: [a] -> [[a]]
shrinkListAggressively :: forall a. [a] -> [[a]]
shrinkListAggressively = \case
  [] -> []
  [a]
xs -> [[], Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [a]
xs, Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
1 [a]
xs]

-- | Resize a generator to grow with the size parameter, but remains reasonably
-- sized. That is handy when testing on data-structures that can be arbitrarily
-- large and, when large entities don't really bring any value to the test
-- itself.
--
-- It uses a square root function which makes the size parameter grows
-- quadratically slower than normal. That is,
--
--     +-------------+------------------+
--     | Normal Size | Reasonable Size  |
--     | ----------- + ---------------- +
--     | 0           | 0                |
--     | 1           | 1                |
--     | 10          | 3                |
--     | 100         | 10               |
--     | 1000        | 31               |
--     +-------------+------------------+
reasonablySized :: Gen a -> Gen a
reasonablySized :: forall a. Gen a -> Gen a
reasonablySized = (Int -> Int) -> Gen a -> Gen a
forall a. (Int -> Int) -> Gen a -> Gen a
scale (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Double -> Int) -> (Int -> Double) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
sqrt @Double (Double -> Double) -> (Int -> Double) -> Int -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral)

-- | A QuickCheck modifier to make use of `reasonablySized` on existing types.
newtype ReasonablySized a = ReasonablySized a
  deriving newtype (Int -> ReasonablySized a -> ShowS
[ReasonablySized a] -> ShowS
ReasonablySized a -> String
(Int -> ReasonablySized a -> ShowS)
-> (ReasonablySized a -> String)
-> ([ReasonablySized a] -> ShowS)
-> Show (ReasonablySized a)
forall a. Show a => Int -> ReasonablySized a -> ShowS
forall a. Show a => [ReasonablySized a] -> ShowS
forall a. Show a => ReasonablySized a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> ReasonablySized a -> ShowS
showsPrec :: Int -> ReasonablySized a -> ShowS
$cshow :: forall a. Show a => ReasonablySized a -> String
show :: ReasonablySized a -> String
$cshowList :: forall a. Show a => [ReasonablySized a] -> ShowS
showList :: [ReasonablySized a] -> ShowS
Show, [ReasonablySized a] -> Encoding
[ReasonablySized a] -> Value
ReasonablySized a -> Bool
ReasonablySized a -> Encoding
ReasonablySized a -> Value
(ReasonablySized a -> Value)
-> (ReasonablySized a -> Encoding)
-> ([ReasonablySized a] -> Value)
-> ([ReasonablySized a] -> Encoding)
-> (ReasonablySized a -> Bool)
-> ToJSON (ReasonablySized a)
forall a. ToJSON a => [ReasonablySized a] -> Encoding
forall a. ToJSON a => [ReasonablySized a] -> Value
forall a. ToJSON a => ReasonablySized a -> Bool
forall a. ToJSON a => ReasonablySized a -> Encoding
forall a. ToJSON a => ReasonablySized a -> Value
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> (a -> Bool)
-> ToJSON a
$ctoJSON :: forall a. ToJSON a => ReasonablySized a -> Value
toJSON :: ReasonablySized a -> Value
$ctoEncoding :: forall a. ToJSON a => ReasonablySized a -> Encoding
toEncoding :: ReasonablySized a -> Encoding
$ctoJSONList :: forall a. ToJSON a => [ReasonablySized a] -> Value
toJSONList :: [ReasonablySized a] -> Value
$ctoEncodingList :: forall a. ToJSON a => [ReasonablySized a] -> Encoding
toEncodingList :: [ReasonablySized a] -> Encoding
$comitField :: forall a. ToJSON a => ReasonablySized a -> Bool
omitField :: ReasonablySized a -> Bool
ToJSON, Maybe (ReasonablySized a)
Value -> Parser [ReasonablySized a]
Value -> Parser (ReasonablySized a)
(Value -> Parser (ReasonablySized a))
-> (Value -> Parser [ReasonablySized a])
-> Maybe (ReasonablySized a)
-> FromJSON (ReasonablySized a)
forall a. FromJSON a => Maybe (ReasonablySized a)
forall a. FromJSON a => Value -> Parser [ReasonablySized a]
forall a. FromJSON a => Value -> Parser (ReasonablySized a)
forall a.
(Value -> Parser a)
-> (Value -> Parser [a]) -> Maybe a -> FromJSON a
$cparseJSON :: forall a. FromJSON a => Value -> Parser (ReasonablySized a)
parseJSON :: Value -> Parser (ReasonablySized a)
$cparseJSONList :: forall a. FromJSON a => Value -> Parser [ReasonablySized a]
parseJSONList :: Value -> Parser [ReasonablySized a]
$comittedField :: forall a. FromJSON a => Maybe (ReasonablySized a)
omittedField :: Maybe (ReasonablySized a)
FromJSON)

instance Arbitrary a => Arbitrary (ReasonablySized a) where
  arbitrary :: Gen (ReasonablySized a)
arbitrary = a -> ReasonablySized a
forall a. a -> ReasonablySized a
ReasonablySized (a -> ReasonablySized a) -> Gen a -> Gen (ReasonablySized a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen a -> Gen a
forall a. Gen a -> Gen a
reasonablySized Gen a
forall a. Arbitrary a => Gen a
arbitrary

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

-- | 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))

-- | 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.
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.
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 -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
": " String -> ShowS
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