{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE UndecidableInstances  #-}

{-# OPTIONS_GHC -Wno-orphans #-}

module PlutusCore.Generators.QuickCheck.Common where

import PlutusCore.Default
import PlutusCore.Name.Unique
import PlutusCore.TypeCheck (defKindCheckConfig)
import PlutusCore.TypeCheck.Internal (inferKindM, runTypeCheckM, withTyVar)
import PlutusIR
import PlutusIR.Error

import Data.Bifunctor
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import GHC.Stack
import Test.QuickCheck.Gen (Gen)
import Test.QuickCheck.Gen qualified as Gen
import Test.QuickCheck.Modifiers (NonNegative (..))
import Test.QuickCheck.Property
import Text.Pretty
import Text.PrettyBy
import Text.PrettyBy.Internal

instance Testable (Either String ()) where
    property :: Either String () -> Property
property = Result -> Property
forall prop. Testable prop => prop -> Property
property (Result -> Property)
-> (Either String () -> Result) -> Either String () -> Property
forall b c a. (b -> c) -> (a -> b) -> a -> c
. \case
        Left String
err -> Result
failed { reason = err }
        Right () -> Result
succeeded

deriving newtype instance Pretty i => Pretty (NonNegative i)
instance PrettyBy config i => DefaultPrettyBy config (NonNegative i)
deriving via PrettyCommon (NonNegative i)
    instance PrettyDefaultBy config (NonNegative i) => PrettyBy config (NonNegative i)

type TypeCtx = Map TyName (Kind ())

-- | Infer the kind of a type in a given kind context
inferKind :: TypeCtx -> Type TyName DefaultUni () -> Either String (Kind ())
inferKind :: TypeCtx -> Type TyName DefaultUni () -> Either String (Kind ())
inferKind TypeCtx
ctx Type TyName DefaultUni ()
ty =
    (Error DefaultUni DefaultFun () -> String)
-> Either (Error DefaultUni DefaultFun ()) (Kind ())
-> Either String (Kind ())
forall a b c. (a -> b) -> Either a c -> Either b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Error DefaultUni DefaultFun () -> String
forall str a. (Pretty a, Render str) => a -> str
display (Either (Error DefaultUni DefaultFun ()) (Kind ())
 -> Either String (Kind ()))
-> (TypeCheckT
      DefaultUni
      DefaultFun
      KindCheckConfig
      (Either (Error DefaultUni DefaultFun ()))
      (Kind ())
    -> Either (Error DefaultUni DefaultFun ()) (Kind ()))
-> TypeCheckT
     DefaultUni
     DefaultFun
     KindCheckConfig
     (Either (Error DefaultUni DefaultFun ()))
     (Kind ())
-> Either String (Kind ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KindCheckConfig
-> TypeCheckT
     DefaultUni
     DefaultFun
     KindCheckConfig
     (Either (Error DefaultUni DefaultFun ()))
     (Kind ())
-> Either (Error DefaultUni DefaultFun ()) (Kind ())
forall cfg (uni :: * -> *) fun (m :: * -> *) a.
cfg -> TypeCheckT uni fun cfg m a -> m a
runTypeCheckM KindCheckConfig
defKindCheckConfig (TypeCheckT
   DefaultUni
   DefaultFun
   KindCheckConfig
   (Either (Error DefaultUni DefaultFun ()))
   (Kind ())
 -> Either String (Kind ()))
-> TypeCheckT
     DefaultUni
     DefaultFun
     KindCheckConfig
     (Either (Error DefaultUni DefaultFun ()))
     (Kind ())
-> Either String (Kind ())
forall a b. (a -> b) -> a -> b
$
        ((TyName, Kind ())
 -> TypeCheckT
      DefaultUni
      DefaultFun
      KindCheckConfig
      (Either (Error DefaultUni DefaultFun ()))
      (Kind ())
 -> TypeCheckT
      DefaultUni
      DefaultFun
      KindCheckConfig
      (Either (Error DefaultUni DefaultFun ()))
      (Kind ()))
-> TypeCheckT
     DefaultUni
     DefaultFun
     KindCheckConfig
     (Either (Error DefaultUni DefaultFun ()))
     (Kind ())
-> [(TyName, Kind ())]
-> TypeCheckT
     DefaultUni
     DefaultFun
     KindCheckConfig
     (Either (Error DefaultUni DefaultFun ()))
     (Kind ())
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
            ((TyName
 -> Kind ()
 -> TypeCheckT
      DefaultUni
      DefaultFun
      KindCheckConfig
      (Either (Error DefaultUni DefaultFun ()))
      (Kind ())
 -> TypeCheckT
      DefaultUni
      DefaultFun
      KindCheckConfig
      (Either (Error DefaultUni DefaultFun ()))
      (Kind ()))
-> (TyName, Kind ())
-> TypeCheckT
     DefaultUni
     DefaultFun
     KindCheckConfig
     (Either (Error DefaultUni DefaultFun ()))
     (Kind ())
-> TypeCheckT
     DefaultUni
     DefaultFun
     KindCheckConfig
     (Either (Error DefaultUni DefaultFun ()))
     (Kind ())
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry TyName
-> Kind ()
-> TypeCheckT
     DefaultUni
     DefaultFun
     KindCheckConfig
     (Either (Error DefaultUni DefaultFun ()))
     (Kind ())
-> TypeCheckT
     DefaultUni
     DefaultFun
     KindCheckConfig
     (Either (Error DefaultUni DefaultFun ()))
     (Kind ())
forall (uni :: * -> *) fun cfg (m :: * -> *) a.
TyName
-> Kind ()
-> TypeCheckT uni fun cfg m a
-> TypeCheckT uni fun cfg m a
withTyVar)
            (forall err term (uni :: * -> *) fun ann (m :: * -> *) cfg.
(MonadKindCheck err term uni fun ann m, HasKindCheckConfig cfg) =>
Type TyName uni ann -> TypeCheckT uni fun cfg m (Kind ())
inferKindM @(Error DefaultUni DefaultFun ()) Type TyName DefaultUni ()
ty)
            (TypeCtx -> [(TyName, Kind ())]
forall k a. Map k a -> [(k, a)]
Map.toList TypeCtx
ctx)

-- | Partial unsafeInferKind, useful for context where invariants are set up to guarantee
-- that types are well-kinded.
unsafeInferKind :: HasCallStack => TypeCtx -> Type TyName DefaultUni () -> Kind ()
unsafeInferKind :: HasCallStack => TypeCtx -> Type TyName DefaultUni () -> Kind ()
unsafeInferKind TypeCtx
ctx Type TyName DefaultUni ()
ty =
  case TypeCtx -> Type TyName DefaultUni () -> Either String (Kind ())
inferKind TypeCtx
ctx Type TyName DefaultUni ()
ty of
    Left String
msg -> String -> Kind ()
forall a. HasCallStack => String -> a
error String
msg
    Right Kind ()
k  -> Kind ()
k

-- | Check well-kindedness of a type in a context
checkKind :: TypeCtx -> Type TyName DefaultUni () -> Kind () -> Either String ()
checkKind :: TypeCtx -> Type TyName DefaultUni () -> Kind () -> Either String ()
checkKind TypeCtx
ctx Type TyName DefaultUni ()
ty Kind ()
kExp =
    if Either String (Kind ())
kInf Either String (Kind ()) -> Either String (Kind ()) -> Bool
forall a. Eq a => a -> a -> Bool
== Kind () -> Either String (Kind ())
forall a b. b -> Either a b
Right Kind ()
kExp
      then () -> Either String ()
forall a b. b -> Either a b
Right ()
      else String -> Either String ()
forall a b. a -> Either a b
Left (String -> Either String ()) -> String -> Either String ()
forall a b. (a -> b) -> a -> b
$ [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
        [ String
"Inferred kind is "
        , Either String (Kind ()) -> String
forall str a. (Pretty a, Render str) => a -> str
display Either String (Kind ())
kInf
        , String
" while expected "
        , Kind () -> String
forall str a. (Pretty a, Render str) => a -> str
display Kind ()
kExp
        ]
  where
    kInf :: Either String (Kind ())
kInf = TypeCtx -> Type TyName DefaultUni () -> Either String (Kind ())
inferKind TypeCtx
ctx Type TyName DefaultUni ()
ty

-- | Generate a list with the given minimum and maximum lengths.
-- It is similar to @Hedgehog.Internal.Gen.list@.
--
-- Note that @genList 0 n gen@ behaves differently than @resize n (listOf gen)@, because
--
-- @
--    resize m (genList 0 n gen) = genList 0 n (resize m gen)
-- @
--
-- whereas
--
-- @
--    resize m (resize n (listOf gen)) = resize n (listOf gen)
-- @
genList :: Int -> Int -> Gen a -> Gen [a]
genList :: forall a. Int -> Int -> Gen a -> Gen [a]
genList Int
lb Int
ub Gen a
gen = do
    Int
len <- (Int, Int) -> Gen Int
Gen.chooseInt (Int
lb, Int
ub)
    Int -> Gen a -> Gen [a]
forall a. Int -> Gen a -> Gen [a]
Gen.vectorOf Int
len Gen a
gen