module UntypedPlutusCore.Generators.Hedgehog.AST
  ( regenConstantsUntil
  , PLC.AstGen
  , PLC.runAstGen
  , PLC.genVersion
  , genTerm
  , genProgram
  , mangleNames
  ) where

import PlutusPrelude

import PlutusCore.Generators.Hedgehog.AST qualified as PLC

import PlutusCore.Compiler.Erase
import UntypedPlutusCore as UPLC

import Data.Set.Lens (setOf)
import Hedgehog
import Universe

regenConstantsUntil
  :: MonadGen m
  => (Some (ValueOf DefaultUni) -> Bool)
  -> Program name DefaultUni fun ann
  -> m (Program name DefaultUni fun ann)
regenConstantsUntil :: forall (m :: * -> *) name fun ann.
MonadGen m =>
(Some (ValueOf DefaultUni) -> Bool)
-> Program name DefaultUni fun ann
-> m (Program name DefaultUni fun ann)
regenConstantsUntil Some (ValueOf DefaultUni) -> Bool
p =
  (Term name DefaultUni fun ann -> m (Term name DefaultUni fun ann))
-> Program name DefaultUni fun ann
-> m (Program name DefaultUni fun ann)
forall name1 (uni1 :: * -> *) fun1 ann name2 (uni2 :: * -> *) fun2
       (f :: * -> *).
Functor f =>
(Term name1 uni1 fun1 ann -> f (Term name2 uni2 fun2 ann))
-> Program name1 uni1 fun1 ann -> f (Program name2 uni2 fun2 ann)
progTerm ((Term name DefaultUni fun ann -> m (Term name DefaultUni fun ann))
 -> Program name DefaultUni fun ann
 -> m (Program name DefaultUni fun ann))
-> ((ann
     -> Some (ValueOf DefaultUni)
     -> m (Maybe (Term name DefaultUni fun ann)))
    -> Term name DefaultUni fun ann
    -> m (Term name DefaultUni fun ann))
-> (ann
    -> Some (ValueOf DefaultUni)
    -> m (Maybe (Term name DefaultUni fun ann)))
-> Program name DefaultUni fun ann
-> m (Program name DefaultUni fun ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ann
 -> Some (ValueOf DefaultUni)
 -> m (Maybe (Term name DefaultUni fun ann)))
-> Term name DefaultUni fun ann -> m (Term name DefaultUni fun ann)
forall (m :: * -> *) ann (uni :: * -> *) name fun.
Monad m =>
(ann -> Some (ValueOf uni) -> m (Maybe (Term name uni fun ann)))
-> Term name uni fun ann -> m (Term name uni fun ann)
termSubstConstantsM ((ann
  -> Some (ValueOf DefaultUni)
  -> m (Maybe (Term name DefaultUni fun ann)))
 -> Program name DefaultUni fun ann
 -> m (Program name DefaultUni fun ann))
-> (ann
    -> Some (ValueOf DefaultUni)
    -> m (Maybe (Term name DefaultUni fun ann)))
-> Program name DefaultUni fun ann
-> m (Program name DefaultUni fun ann)
forall a b. (a -> b) -> a -> b
$ \ann
ann -> (Maybe (Some (ValueOf DefaultUni))
 -> Maybe (Term name DefaultUni fun ann))
-> m (Maybe (Some (ValueOf DefaultUni)))
-> m (Maybe (Term name DefaultUni fun ann))
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Some (ValueOf DefaultUni) -> Term name DefaultUni fun ann)
-> Maybe (Some (ValueOf DefaultUni))
-> Maybe (Term name DefaultUni fun ann)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Some (ValueOf DefaultUni) -> Term name DefaultUni fun ann)
 -> Maybe (Some (ValueOf DefaultUni))
 -> Maybe (Term name DefaultUni fun ann))
-> (Some (ValueOf DefaultUni) -> Term name DefaultUni fun ann)
-> Maybe (Some (ValueOf DefaultUni))
-> Maybe (Term name DefaultUni fun ann)
forall a b. (a -> b) -> a -> b
$ ann -> Some (ValueOf DefaultUni) -> Term name DefaultUni fun ann
forall name (uni :: * -> *) fun ann.
ann -> Some (ValueOf uni) -> Term name uni fun ann
Constant ann
ann) (m (Maybe (Some (ValueOf DefaultUni)))
 -> m (Maybe (Term name DefaultUni fun ann)))
-> (Some (ValueOf DefaultUni)
    -> m (Maybe (Some (ValueOf DefaultUni))))
-> Some (ValueOf DefaultUni)
-> m (Maybe (Term name DefaultUni fun ann))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Some (ValueOf DefaultUni) -> Bool)
-> Some (ValueOf DefaultUni)
-> m (Maybe (Some (ValueOf DefaultUni)))
forall (m :: * -> *).
MonadGen m =>
(Some (ValueOf DefaultUni) -> Bool)
-> Some (ValueOf DefaultUni)
-> m (Maybe (Some (ValueOf DefaultUni)))
PLC.regenConstantUntil Some (ValueOf DefaultUni) -> Bool
p

genTerm
  :: forall fun
   . (Bounded fun, Enum fun)
  => PLC.AstGen (Term Name DefaultUni fun ())
genTerm :: forall fun.
(Bounded fun, Enum fun) =>
AstGen (Term Name DefaultUni fun ())
genTerm = (Term TyName Name DefaultUni fun () -> Term Name DefaultUni fun ())
-> GenT (Reader [Name]) (Term TyName Name DefaultUni fun ())
-> GenT (Reader [Name]) (Term Name DefaultUni fun ())
forall a b.
(a -> b) -> GenT (Reader [Name]) a -> GenT (Reader [Name]) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Term TyName Name DefaultUni fun () -> Term Name DefaultUni fun ()
forall name tyname (uni :: * -> *) fun ann.
HasUnique name TermUnique =>
Term tyname name uni fun ann -> Term name uni fun ann
eraseTerm GenT (Reader [Name]) (Term TyName Name DefaultUni fun ())
forall fun.
(Bounded fun, Enum fun) =>
AstGen (Term TyName Name DefaultUni fun ())
PLC.genTerm

genProgram
  :: forall fun
   . (Bounded fun, Enum fun) => PLC.AstGen (Program Name DefaultUni fun ())
genProgram :: forall fun.
(Bounded fun, Enum fun) =>
AstGen (Program Name DefaultUni fun ())
genProgram = (Program TyName Name DefaultUni fun ()
 -> Program Name DefaultUni fun ())
-> GenT (Reader [Name]) (Program TyName Name DefaultUni fun ())
-> GenT (Reader [Name]) (Program Name DefaultUni fun ())
forall a b.
(a -> b) -> GenT (Reader [Name]) a -> GenT (Reader [Name]) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Program TyName Name DefaultUni fun ()
-> Program Name DefaultUni fun ()
forall name tyname (uni :: * -> *) fun ann.
HasUnique name TermUnique =>
Program tyname name uni fun ann -> Program name uni fun ann
eraseProgram GenT (Reader [Name]) (Program TyName Name DefaultUni fun ())
forall fun.
(Bounded fun, Enum fun) =>
AstGen (Program TyName Name DefaultUni fun ())
PLC.genProgram

-- See Note [Name mangling]
mangleNames
  :: Term Name DefaultUni DefaultFun ()
  -> PLC.AstGen (Maybe (Term Name DefaultUni DefaultFun ()))
mangleNames :: Term Name DefaultUni DefaultFun ()
-> AstGen (Maybe (Term Name DefaultUni DefaultFun ()))
mangleNames Term Name DefaultUni DefaultFun ()
term = do
  Maybe (Name -> AstGen (Maybe Name))
mayMang <- Set Name -> AstGen (Maybe (Name -> AstGen (Maybe Name)))
PLC.genNameMangler (Set Name -> AstGen (Maybe (Name -> AstGen (Maybe Name))))
-> Set Name -> AstGen (Maybe (Name -> AstGen (Maybe Name)))
forall a b. (a -> b) -> a -> b
$ Getting (Set Name) (Term Name DefaultUni DefaultFun ()) Name
-> Term Name DefaultUni DefaultFun () -> Set Name
forall a s. Getting (Set a) s a -> s -> Set a
setOf Getting (Set Name) (Term Name DefaultUni DefaultFun ()) Name
forall name (uni :: * -> *) fun ann (f :: * -> *).
(Contravariant f, Applicative f) =>
(name -> f name)
-> Term name uni fun ann -> f (Term name uni fun ann)
vTerm Term Name DefaultUni DefaultFun ()
term
  Maybe (Name -> AstGen (Maybe Name))
-> ((Name -> AstGen (Maybe Name))
    -> GenT (Reader [Name]) (Term Name DefaultUni DefaultFun ()))
-> AstGen (Maybe (Term Name DefaultUni DefaultFun ()))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for Maybe (Name -> AstGen (Maybe Name))
mayMang (((Name -> AstGen (Maybe Name))
  -> GenT (Reader [Name]) (Term Name DefaultUni DefaultFun ()))
 -> AstGen (Maybe (Term Name DefaultUni DefaultFun ())))
-> ((Name -> AstGen (Maybe Name))
    -> GenT (Reader [Name]) (Term Name DefaultUni DefaultFun ()))
-> AstGen (Maybe (Term Name DefaultUni DefaultFun ()))
forall a b. (a -> b) -> a -> b
$ \Name -> AstGen (Maybe Name)
mang -> (Name -> () -> AstGen (Maybe (Term Name DefaultUni DefaultFun ())))
-> Term Name DefaultUni DefaultFun ()
-> GenT (Reader [Name]) (Term Name DefaultUni DefaultFun ())
forall (m :: * -> *) name ann (uni :: * -> *) fun.
Monad m =>
(name -> ann -> m (Maybe (Term name uni fun ann)))
-> Term name uni fun ann -> m (Term name uni fun ann)
termSubstNamesM (AstGen (Maybe (Term Name DefaultUni DefaultFun ()))
-> () -> AstGen (Maybe (Term Name DefaultUni DefaultFun ()))
forall a b. a -> b -> a
const (AstGen (Maybe (Term Name DefaultUni DefaultFun ()))
 -> () -> AstGen (Maybe (Term Name DefaultUni DefaultFun ())))
-> (Name -> AstGen (Maybe (Term Name DefaultUni DefaultFun ())))
-> Name
-> ()
-> AstGen (Maybe (Term Name DefaultUni DefaultFun ()))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe Name -> Maybe (Term Name DefaultUni DefaultFun ()))
-> AstGen (Maybe Name)
-> AstGen (Maybe (Term Name DefaultUni DefaultFun ()))
forall a b.
(a -> b) -> GenT (Reader [Name]) a -> GenT (Reader [Name]) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Name -> Term Name DefaultUni DefaultFun ())
-> Maybe Name -> Maybe (Term Name DefaultUni DefaultFun ())
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Name -> Term Name DefaultUni DefaultFun ())
 -> Maybe Name -> Maybe (Term Name DefaultUni DefaultFun ()))
-> (Name -> Term Name DefaultUni DefaultFun ())
-> Maybe Name
-> Maybe (Term Name DefaultUni DefaultFun ())
forall a b. (a -> b) -> a -> b
$ () -> Name -> Term Name DefaultUni DefaultFun ()
forall name (uni :: * -> *) fun ann.
ann -> name -> Term name uni fun ann
UPLC.Var ()) (AstGen (Maybe Name)
 -> AstGen (Maybe (Term Name DefaultUni DefaultFun ())))
-> (Name -> AstGen (Maybe Name))
-> Name
-> AstGen (Maybe (Term Name DefaultUni DefaultFun ()))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> AstGen (Maybe Name)
mang) Term Name DefaultUni DefaultFun ()
term