{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

module PlutusCore.Generators.Hedgehog.Builtin (
    GenTypedTerm (..),
    GenArbitraryTerm (..),
    genConstant
) where

import PlutusCore hiding (Constr)
import PlutusCore.Builtin
import PlutusCore.Crypto.BLS12_381.G1 qualified as BLS12_381.G1
import PlutusCore.Crypto.BLS12_381.G2 qualified as BLS12_381.G2
import PlutusCore.Crypto.BLS12_381.Pairing qualified as BLS12_381.Pairing
import PlutusCore.Data (Data (..))
import PlutusCore.Evaluation.Machine.ExMemoryUsage (IntegerCostedLiterally, ListCostedByLength,
                                                    NumBytesCostedAsNumWords)
import PlutusCore.Generators.Hedgehog.AST hiding (genConstant)
import PlutusCore.Generators.QuickCheck.Builtin

import Data.ByteString qualified as BS
import Data.Kind qualified as GHC
import Data.Text (Text)
import Data.Type.Equality
import Data.Word (Word8)
import GHC.Natural
import Hedgehog hiding (Opaque, Var, eval)
import Hedgehog.Gen qualified as Gen
import Hedgehog.Gen.QuickCheck (arbitrary)
import Hedgehog.Range qualified as Range
import Type.Reflection

-- | This class exists so we can provide an ad-hoc typed term generator
-- for various universes. We usually rely-on a universe-specific generator
-- for well-typed constants within that universe.
--
-- TODO: Move this to "PlutusIR.Generators.AST", and merge `genConstant` with
-- `PlutusIR.Generators.AST.genConstant`.
class GenTypedTerm uni where
    -- | Generate a `Term` in @uni@ with the given type.
    genTypedTerm ::
        forall (a :: GHC.Type) tyname name fun.
        TypeRep a ->
        Gen (Term tyname name uni fun ())

instance GenTypedTerm DefaultUni where
    -- TODO: currently it only generates constant terms.
    genTypedTerm :: forall a tyname name fun.
TypeRep a -> Gen (Term tyname name DefaultUni fun ())
genTypedTerm TypeRep a
tr = case TypeRep a -> SomeGen DefaultUni
forall a. TypeRep a -> SomeGen DefaultUni
genConstant TypeRep a
tr of
        SomeGen Gen a
gen -> ()
-> Some (ValueOf DefaultUni) -> Term tyname name DefaultUni fun ()
forall tyname name (uni :: * -> *) fun ann.
ann -> Some (ValueOf uni) -> Term tyname name uni fun ann
Constant () (Some (ValueOf DefaultUni) -> Term tyname name DefaultUni fun ())
-> (a -> Some (ValueOf DefaultUni))
-> a
-> Term tyname name DefaultUni fun ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Some (ValueOf DefaultUni)
forall a (uni :: * -> *). Contains uni a => a -> Some (ValueOf uni)
someValue (a -> Term tyname name DefaultUni fun ())
-> Gen a -> Gen (Term tyname name DefaultUni fun ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen a
gen

-- | This class exists so we can provide an ad-hoc arbitrary term generator
-- for various universes.
class GenArbitraryTerm uni where
    -- | Generate an arbitrary `Term` in @uni@.
    genArbitraryTerm ::
        forall fun.
        (Bounded fun, Enum fun) =>
        Gen (Term TyName Name uni fun ())

instance GenArbitraryTerm DefaultUni where
    genArbitraryTerm :: forall fun.
(Bounded fun, Enum fun) =>
Gen (Term TyName Name DefaultUni fun ())
genArbitraryTerm = AstGen (Term TyName Name DefaultUni fun ())
-> GenT Identity (Term TyName Name DefaultUni fun ())
forall (m :: * -> *) a. MonadGen m => AstGen a -> m a
runAstGen AstGen (Term TyName Name DefaultUni fun ())
forall fun.
(Bounded fun, Enum fun) =>
AstGen (Term TyName Name DefaultUni fun ())
genTerm

data SomeGen uni = forall a. uni `HasTermLevel` a => SomeGen (Gen a)

genArbitraryBuiltin
    :: forall a. (ArbitraryBuiltin a, DefaultUni `HasTermLevel` a) => SomeGen DefaultUni
genArbitraryBuiltin :: forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin = Gen a -> SomeGen DefaultUni
forall (uni :: * -> *) a.
HasTermLevel uni a =>
Gen a -> SomeGen uni
SomeGen (Gen a -> SomeGen DefaultUni) -> Gen a -> SomeGen DefaultUni
forall a b. (a -> b) -> a -> b
$ AsArbitraryBuiltin a -> a
forall a. AsArbitraryBuiltin a -> a
unAsArbitraryBuiltin (AsArbitraryBuiltin a -> a)
-> GenT Identity (AsArbitraryBuiltin a) -> Gen a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a (m :: * -> *). (Arbitrary a, MonadGen m) => m a
arbitrary @(AsArbitraryBuiltin a)

genConstant :: forall (a :: GHC.Type). TypeRep a -> SomeGen DefaultUni
genConstant :: forall a. TypeRep a -> SomeGen DefaultUni
genConstant TypeRep a
tr
    | Just a :~~: ()
HRefl <- TypeRep a -> TypeRep () -> Maybe (a :~~: ())
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @()) = forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @()
    | Just a :~~: Integer
HRefl <- TypeRep a -> TypeRep Integer -> Maybe (a :~~: Integer)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Integer) = forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @Integer
    | Just a :~~: Int
HRefl <- TypeRep a -> TypeRep Int -> Maybe (a :~~: Int)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Int) = forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @Integer
    | Just a :~~: Word8
HRefl <- TypeRep a -> TypeRep Word8 -> Maybe (a :~~: Word8)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Word8) = forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @Integer
    | Just a :~~: Natural
HRefl <- TypeRep a -> TypeRep Natural -> Maybe (a :~~: Natural)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Natural) = forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @Integer
    | Just a :~~: NumBytesCostedAsNumWords
HRefl <- TypeRep a
-> TypeRep NumBytesCostedAsNumWords
-> Maybe (a :~~: NumBytesCostedAsNumWords)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @NumBytesCostedAsNumWords) = forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @Integer
    | Just a :~~: IntegerCostedLiterally
HRefl <- TypeRep a
-> TypeRep IntegerCostedLiterally
-> Maybe (a :~~: IntegerCostedLiterally)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @IntegerCostedLiterally) = forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @Integer
    | Just a :~~: Bool
HRefl <- TypeRep a -> TypeRep Bool -> Maybe (a :~~: Bool)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Bool) = forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @Bool
    | Just a :~~: ByteString
HRefl <- TypeRep a -> TypeRep ByteString -> Maybe (a :~~: ByteString)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @BS.ByteString) = forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @BS.ByteString
    | Just a :~~: Text
HRefl <- TypeRep a -> TypeRep Text -> Maybe (a :~~: Text)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Text) = forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @Text
    | Just a :~~: Data
HRefl <- TypeRep a -> TypeRep Data -> Maybe (a :~~: Data)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Data) = forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @Data
    | Just a :~~: Element
HRefl <- TypeRep a -> TypeRep Element -> Maybe (a :~~: Element)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @BLS12_381.G1.Element) =
          forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @BLS12_381.G1.Element
    | Just a :~~: Element
HRefl <- TypeRep a -> TypeRep Element -> Maybe (a :~~: Element)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @BLS12_381.G2.Element) =
          forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @BLS12_381.G2.Element
    | Just a :~~: MlResult
HRefl <- TypeRep a -> TypeRep MlResult -> Maybe (a :~~: MlResult)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
tr (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @BLS12_381.Pairing.MlResult) =
          forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @BLS12_381.Pairing.MlResult
    | TypeRep a
trPair `App` TypeRep b
tr1 `App` TypeRep b
tr2 <- TypeRep a
tr
    , Just a :~~: (,)
HRefl <- TypeRep a -> TypeRep (,) -> Maybe (a :~~: (,))
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
trPair (forall {k} (a :: k). Typeable a => TypeRep a
forall (a :: * -> * -> *). Typeable a => TypeRep a
typeRep @(,)) =
        -- We can perhaps use the @QuickCheck@ generator here too, but this seems rather hard.
        -- Maybe we should simply copy the logic. Should we halve the size explicitly here?
        case (TypeRep b -> SomeGen DefaultUni
forall a. TypeRep a -> SomeGen DefaultUni
genConstant TypeRep b
TypeRep b
tr1, TypeRep b -> SomeGen DefaultUni
forall a. TypeRep a -> SomeGen DefaultUni
genConstant TypeRep b
TypeRep b
tr2) of
            (SomeGen Gen a
g1, SomeGen Gen a
g2) -> Gen (a, a) -> SomeGen DefaultUni
forall (uni :: * -> *) a.
HasTermLevel uni a =>
Gen a -> SomeGen uni
SomeGen (Gen (a, a) -> SomeGen DefaultUni)
-> Gen (a, a) -> SomeGen DefaultUni
forall a b. (a -> b) -> a -> b
$ (,) (a -> a -> (a, a)) -> Gen a -> GenT Identity (a -> (a, a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen a
g1 GenT Identity (a -> (a, a)) -> Gen a -> Gen (a, a)
forall a b.
GenT Identity (a -> b) -> GenT Identity a -> GenT Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen a
g2
    | TypeRep a
trList `App` TypeRep b
trElem <- TypeRep a
tr
    , Just a :~~: []
HRefl <- TypeRep a -> TypeRep [] -> Maybe (a :~~: [])
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
trList (forall {k} (a :: k). Typeable a => TypeRep a
forall (a :: * -> *). Typeable a => TypeRep a
typeRep @[]) =
        case TypeRep b -> SomeGen DefaultUni
forall a. TypeRep a -> SomeGen DefaultUni
genConstant TypeRep b
TypeRep b
trElem of
            SomeGen Gen a
genElem -> Gen [a] -> SomeGen DefaultUni
forall (uni :: * -> *) a.
HasTermLevel uni a =>
Gen a -> SomeGen uni
SomeGen (Gen [a] -> SomeGen DefaultUni) -> Gen [a] -> SomeGen DefaultUni
forall a b. (a -> b) -> a -> b
$ Range Int -> Gen a -> Gen [a]
forall (m :: * -> *) a. MonadGen m => Range Int -> m a -> m [a]
Gen.list (Int -> Int -> Range Int
forall a. Integral a => a -> a -> Range a
Range.linear Int
0 Int
10) Gen a
genElem
    | TypeRep a
trList' `App` TypeRep b
trElem <- TypeRep a
tr
    , Just a :~~: ListCostedByLength
HRefl <- TypeRep a
-> TypeRep ListCostedByLength -> Maybe (a :~~: ListCostedByLength)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
trList' (forall {k} (a :: k). Typeable a => TypeRep a
forall (a :: * -> *). Typeable a => TypeRep a
typeRep @ListCostedByLength) =
        case TypeRep b -> SomeGen DefaultUni
forall a. TypeRep a -> SomeGen DefaultUni
genConstant TypeRep b
TypeRep b
trElem of
          SomeGen Gen a
genElem -> Gen [a] -> SomeGen DefaultUni
forall (uni :: * -> *) a.
HasTermLevel uni a =>
Gen a -> SomeGen uni
SomeGen (Gen [a] -> SomeGen DefaultUni) -> Gen [a] -> SomeGen DefaultUni
forall a b. (a -> b) -> a -> b
$ Range Int -> Gen a -> Gen [a]
forall (m :: * -> *) a. MonadGen m => Range Int -> m a -> m [a]
Gen.list (Int -> Int -> Range Int
forall a. Integral a => a -> a -> Range a
Range.linear Int
0 Int
10) Gen a
genElem
    | TypeRep a
trSomeConstant `App` TypeRep b
_ `App` TypeRep b
trEl <- TypeRep a
tr
    , Just a :~~: SomeConstant
HRefl <- TypeRep a -> TypeRep SomeConstant -> Maybe (a :~~: SomeConstant)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
trSomeConstant (forall {k} (a :: k). Typeable a => TypeRep a
forall (a :: (* -> *) -> * -> *). Typeable a => TypeRep a
typeRep @SomeConstant) =
        TypeRep b -> SomeGen DefaultUni
forall a. TypeRep a -> SomeGen DefaultUni
genConstant TypeRep b
TypeRep b
trEl
    | TypeRep a
trOpaque `App` TypeRep b
_ `App` TypeRep b
trEl <- TypeRep a
tr
    , Just a :~~: Opaque
HRefl <- TypeRep a -> TypeRep Opaque -> Maybe (a :~~: Opaque)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
trOpaque (forall {k} (a :: k). Typeable a => TypeRep a
forall (a :: * -> * -> *). Typeable a => TypeRep a
typeRep @Opaque) =
        TypeRep b -> SomeGen DefaultUni
forall a. TypeRep a -> SomeGen DefaultUni
genConstant TypeRep b
TypeRep b
trEl
    | TypeRep a
trTyVarRep `App` TypeRep b
_ <- TypeRep a
tr
    , Just a :~~: TyVarRep
HRefl <- TypeRep a -> TypeRep TyVarRep -> Maybe (a :~~: TyVarRep)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
trTyVarRep (forall {k} (a :: k). Typeable a => TypeRep a
forall (a :: TyNameRep (*) -> *). Typeable a => TypeRep a
typeRep @(TyVarRep @GHC.Type)) =
        -- In the current implementation, all type variables are instantiated
        -- to `Integer` (TODO: change this?).
        forall a.
(ArbitraryBuiltin a, HasTermLevel DefaultUni a) =>
SomeGen DefaultUni
genArbitraryBuiltin @Integer
    | Bool
otherwise =
        [Char] -> SomeGen DefaultUni
forall a. HasCallStack => [Char] -> a
error ([Char] -> SomeGen DefaultUni) -> [Char] -> SomeGen DefaultUni
forall a b. (a -> b) -> a -> b
$
            [Char]
"genConstant: I don't know how to generate constants of this type: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> TypeRep a -> [Char]
forall a. Show a => a -> [Char]
show TypeRep a
tr