-- editorconfig-checker-disable-file
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances  #-}

{-# OPTIONS_GHC -Wno-orphans #-}

module PlutusCore.Generators.QuickCheck.GenTm
  ( module PlutusCore.Generators.QuickCheck.GenTm
  , module Export
  , Arbitrary (..)
  , Gen
  ) where

import PlutusCore.Generators.QuickCheck.Common
import PlutusCore.Generators.QuickCheck.Utils

import PlutusCore.Default
import PlutusCore.Name.Unique
import PlutusIR
import PlutusIR.Compiler
import PlutusIR.Subst

import Control.Monad.Reader
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Set.Lens (setOf)
import Data.String
import Test.QuickCheck (Arbitrary (..), Gen)
import Test.QuickCheck qualified as QC
import Test.QuickCheck.GenT as Export hiding (var)

instance MonadReader r m => MonadReader r (GenT m) where
    ask :: GenT m r
ask = m r -> GenT m r
forall (m :: * -> *) a. Monad m => m a -> GenT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m r
forall r (m :: * -> *). MonadReader r m => m r
ask
    local :: forall a. (r -> r) -> GenT m a -> GenT m a
local r -> r
f (GenT QCGen -> Int -> m a
k) = (QCGen -> Int -> m a) -> GenT m a
forall (m :: * -> *) a. (QCGen -> Int -> m a) -> GenT m a
GenT ((QCGen -> Int -> m a) -> GenT m a)
-> (QCGen -> Int -> m a) -> GenT m a
forall a b. (a -> b) -> a -> b
$ \QCGen
qc Int
size -> (r -> r) -> m a -> m a
forall a. (r -> r) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local r -> r
f (m a -> m a) -> m a -> m a
forall a b. (a -> b) -> a -> b
$ QCGen -> Int -> m a
k QCGen
qc Int
size

-- | Term generators carry around a context to know
-- e.g. what types and terms are in scope.
type GenTm = GenT (Reader GenEnv)

data GenEnv = GenEnv
  { GenEnv -> Int
geAstSize            :: Int
  -- ^ Generator size bound. See Note [Recursion Control and geAstSize]
  , GenEnv -> Map TyName (Datatype TyName Name DefaultUni ())
geDatas              :: Map TyName (Datatype TyName Name DefaultUni ())
  -- ^ Datatype context.
  -- TODO: it's a little weird, 'cause it maps the datatype name to the datatype and the datatype
  -- introduces multiple names. It's probably fine to have something like that, though.
  , GenEnv -> TypeCtx
geTypes              :: TypeCtx
  -- ^ Type context
  , GenEnv -> Map Name (Type TyName DefaultUni ())
geTerms              :: Map Name (Type TyName DefaultUni ())
  -- ^ Term context
  , GenEnv -> Set TyName
geUnboundUsedTyNames :: Set TyName
  -- ^ Names that we have generated and don't want to shadow but haven't bound yet.
  , GenEnv -> AllowEscape
geEscaping           :: AllowEscape
  -- ^ Are we in a place where we are allowed to generate a datatype binding?
  , GenEnv
-> Maybe (Type TyName DefaultUni ())
-> GenTm
     (Type TyName DefaultUni (),
      Term TyName Name DefaultUni DefaultFun ())
geCustomGen          :: Maybe (Type TyName DefaultUni ())
                         -- TODO: this could return a `Maybe` perhaps. Or it
                         -- could be `Maybe (Maybe Type -> ...)`.
                         -> GenTm (Type TyName DefaultUni (), Term TyName Name DefaultUni DefaultFun ())
  -- ^ A custom user-controlled generator for terms - useful for situations where
  -- we want to e.g. generate custom strings for coverage or test some specific
  -- pattern that generates a special case for the compiler.
  , GenEnv -> Int
geCustomFreq         :: Int
  -- ^ How often do we use the custom generator -
  -- values in the range of 10-30 are usually reasonable.
  , GenEnv -> Bool
geDebug              :: Bool
  -- ^ Are we currently running in debug-mode (to debug our generators)
  }

{- Note [Recursion Control and geAstSize]
One would be forgiven for thinking that you don't need `geAstSize` in `GenTm` because size is
built-in to 'Gen'. However, if you use 'Gen's built-in size to control the size of both the terms
you generate *and* the size of the constants in the terms you will end up with skewed
terms. Constants near the top of the term will be big and constants near the bottom of the term will
be small. For this reason we follow QuickCheck best practice and keep track of the "recursion
control size" separately from 'Gen's size in the 'geAstSize' field of the 'GenEnv'
environment. I.e. we let the QuickCheck's size parameter to be responsible for the size of constants
at the leaves of the AST and use 'geAstSize' to control the size of the AST itself.
-}

-- | Run a generator in debug-mode.
withDebug :: GenTm a -> GenTm a
withDebug :: forall a. GenTm a -> GenTm a
withDebug GenTm a
gen = (GenEnv -> GenEnv) -> GenTm a -> GenTm a
forall a.
(GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (\GenEnv
env -> GenEnv
env { geDebug = True }) GenTm a
gen

-- | Run a `GenTm  generator in a top-level empty context where we are allowed to generate
-- datatypes.
runGenTm :: GenTm a -> Gen a
runGenTm :: forall a. GenTm a -> Gen a
runGenTm = Int
-> (Maybe (Type TyName DefaultUni ())
    -> GenTm
         (Type TyName DefaultUni (),
          Term TyName Name DefaultUni DefaultFun ()))
-> GenTm a
-> Gen a
forall a.
Int
-> (Maybe (Type TyName DefaultUni ())
    -> GenTm
         (Type TyName DefaultUni (),
          Term TyName Name DefaultUni DefaultFun ()))
-> GenTm a
-> Gen a
runGenTmCustom Int
0 ([Char]
-> Maybe (Type TyName DefaultUni ())
-> GenTm
     (Type TyName DefaultUni (),
      Term TyName Name DefaultUni DefaultFun ())
forall a. HasCallStack => [Char] -> a
error [Char]
"No custom generator.")

-- | Run a `GenTm` generator with a plug-in custom generator for terms that is included with
-- the other generators.
runGenTmCustom :: Int
               -> (Maybe (Type TyName DefaultUni ())
                  -> GenTm (Type TyName DefaultUni (), Term TyName Name DefaultUni DefaultFun ()))
               -> GenTm a
               -> Gen a
runGenTmCustom :: forall a.
Int
-> (Maybe (Type TyName DefaultUni ())
    -> GenTm
         (Type TyName DefaultUni (),
          Term TyName Name DefaultUni DefaultFun ()))
-> GenTm a
-> Gen a
runGenTmCustom Int
f Maybe (Type TyName DefaultUni ())
-> GenTm
     (Type TyName DefaultUni (),
      Term TyName Name DefaultUni DefaultFun ())
cg GenTm a
g = do
  (Int -> Gen a) -> Gen a
forall a. (Int -> Gen a) -> Gen a
forall (g :: * -> *) a. MonadGen g => (Int -> g a) -> g a
sized ((Int -> Gen a) -> Gen a) -> (Int -> Gen a) -> Gen a
forall a b. (a -> b) -> a -> b
$ \ Int
n -> do
    let env :: GenEnv
env = GenEnv
          { -- Duplicating the QC size parameter that we use to control the size of constants as
            -- the initial AST size parameter.
            geAstSize :: Int
geAstSize            = Int
n
          , geDatas :: Map TyName (Datatype TyName Name DefaultUni ())
geDatas              = Map TyName (Datatype TyName Name DefaultUni ())
forall k a. Map k a
Map.empty
          , geTypes :: TypeCtx
geTypes              = TypeCtx
forall k a. Map k a
Map.empty
          , geTerms :: Map Name (Type TyName DefaultUni ())
geTerms              = Map Name (Type TyName DefaultUni ())
forall k a. Map k a
Map.empty
          , geUnboundUsedTyNames :: Set TyName
geUnboundUsedTyNames = Set TyName
forall a. Set a
Set.empty
          , geEscaping :: AllowEscape
geEscaping           = AllowEscape
YesEscape
          , geCustomGen :: Maybe (Type TyName DefaultUni ())
-> GenTm
     (Type TyName DefaultUni (),
      Term TyName Name DefaultUni DefaultFun ())
geCustomGen          = Maybe (Type TyName DefaultUni ())
-> GenTm
     (Type TyName DefaultUni (),
      Term TyName Name DefaultUni DefaultFun ())
cg
          , geCustomFreq :: Int
geCustomFreq         = Int
f
          , geDebug :: Bool
geDebug              = Bool
False
          }
    (Reader GenEnv a -> GenEnv -> a) -> GenEnv -> Reader GenEnv a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Reader GenEnv a -> GenEnv -> a
forall r a. Reader r a -> r -> a
runReader GenEnv
env (Reader GenEnv a -> a) -> Gen (Reader GenEnv a) -> Gen a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GenTm a -> Gen (Reader GenEnv a)
forall (m :: * -> *) a. GenT m a -> Gen (m a)
runGenT GenTm a
g

-- | Create a generator that runs the given generator and applies the given function to produced
-- values until the result is a @Just y@, returning the @y@.
suchThatMap :: Monad m => GenT m a -> (a -> Maybe b) -> GenT m b
GenT m a
gen suchThatMap :: forall (m :: * -> *) a b.
Monad m =>
GenT m a -> (a -> Maybe b) -> GenT m b
`suchThatMap` a -> Maybe b
f = (Maybe b -> b) -> GenT m (Maybe b) -> GenT m b
forall a b. (a -> b) -> GenT m a -> GenT m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Maybe b -> b
forall a. HasCallStack => Maybe a -> a
fromJust (GenT m (Maybe b) -> GenT m b) -> GenT m (Maybe b) -> GenT m b
forall a b. (a -> b) -> a -> b
$ (a -> Maybe b) -> GenT m a -> GenT m (Maybe b)
forall a b. (a -> b) -> GenT m a -> GenT m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Maybe b
f GenT m a
gen GenT m (Maybe b) -> (Maybe b -> Bool) -> GenT m (Maybe b)
forall (m :: * -> *) a. MonadGen m => m a -> (a -> Bool) -> m a
`suchThat` Maybe b -> Bool
forall a. Maybe a -> Bool
isJust

-- | Create a generator that runs the given generator until the result is a @Just x@,
-- returning the @x@.
deliver :: Monad m => GenT m (Maybe a) -> GenT m a
deliver :: forall (m :: * -> *) a. Monad m => GenT m (Maybe a) -> GenT m a
deliver GenT m (Maybe a)
gen = GenT m (Maybe a)
gen GenT m (Maybe a) -> (Maybe a -> Maybe a) -> GenT m a
forall (m :: * -> *) a b.
Monad m =>
GenT m a -> (a -> Maybe b) -> GenT m b
`suchThatMap` Maybe a -> Maybe a
forall a. a -> a
id

-- * Utility functions

-- | Don't allow types to escape from a generator.
withNoEscape :: GenTm a -> GenTm a
withNoEscape :: forall a. GenTm a -> GenTm a
withNoEscape = (GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a
forall a.
(GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((GenEnv -> GenEnv)
 -> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a)
-> (GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a
-> GenT (Reader GenEnv) a
forall a b. (a -> b) -> a -> b
$ \GenEnv
env -> GenEnv
env { geEscaping = NoEscape }

-- * Dealing with size

-- | Map a function over the generator size
onAstSize :: (Int -> Int) -> GenTm a -> GenTm a
onAstSize :: forall a. (Int -> Int) -> GenTm a -> GenTm a
onAstSize Int -> Int
f = (GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a
forall a.
(GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((GenEnv -> GenEnv)
 -> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a)
-> (GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a
-> GenT (Reader GenEnv) a
forall a b. (a -> b) -> a -> b
$ \ GenEnv
env -> GenEnv
env { geAstSize = f (geAstSize env) }

-- | Default to the first generator if the size is zero (or negative),
-- use the second generator otherwise.
ifAstSizeZero :: GenTm a -> GenTm a -> GenTm a
ifAstSizeZero :: forall a. GenTm a -> GenTm a -> GenTm a
ifAstSizeZero GenTm a
ifZ GenTm a
nonZ = do
  Int
n <- (GenEnv -> Int) -> GenT (Reader GenEnv) Int
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks GenEnv -> Int
geAstSize
  if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 then GenTm a
ifZ else GenTm a
nonZ

-- | Locally set the size in a generator
withAstSize :: Int -> GenTm a -> GenTm a
withAstSize :: forall a. Int -> GenTm a -> GenTm a
withAstSize = (Int -> Int) -> GenTm a -> GenTm a
forall a. (Int -> Int) -> GenTm a -> GenTm a
onAstSize ((Int -> Int) -> GenTm a -> GenTm a)
-> (Int -> Int -> Int) -> Int -> GenTm a -> GenTm a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a b. a -> b -> a
const

-- | Split the size between two generators in the ratio specified by
-- the first two arguments.
astSizeSplit_ :: Int -> Int -> GenTm a -> GenTm b -> GenTm (a, b)
astSizeSplit_ :: forall a b. Int -> Int -> GenTm a -> GenTm b -> GenTm (a, b)
astSizeSplit_ Int
a Int
b GenTm a
ga GenTm b
gb = Int -> Int -> GenTm a -> (a -> GenTm b) -> GenTm (a, b)
forall a b. Int -> Int -> GenTm a -> (a -> GenTm b) -> GenTm (a, b)
astSizeSplit Int
a Int
b GenTm a
ga (GenTm b -> a -> GenTm b
forall a b. a -> b -> a
const GenTm b
gb)

-- | Split the size between two generators in the ratio specified by
-- the first two arguments and use the result of the first generator
-- in the second.
astSizeSplit :: Int -> Int -> GenTm a -> (a -> GenTm b) -> GenTm (a, b)
astSizeSplit :: forall a b. Int -> Int -> GenTm a -> (a -> GenTm b) -> GenTm (a, b)
astSizeSplit Int
a Int
b GenTm a
ga a -> GenTm b
gb = do
  Int
n <- (GenEnv -> Int) -> GenT (Reader GenEnv) Int
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks GenEnv -> Int
geAstSize
  let na :: Int
na = (Int
a Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` (Int
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
b)
      nb :: Int
nb = (Int
b Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` (Int
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
b)
  a
x <- Int -> GenTm a -> GenTm a
forall a. Int -> GenTm a -> GenTm a
withAstSize Int
na GenTm a
ga
  (,) a
x (b -> (a, b)) -> GenTm b -> GenTm (a, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> GenTm b -> GenTm b
forall a. Int -> GenTm a -> GenTm a
withAstSize Int
nb (a -> GenTm b
gb a
x)

-- * Names

{- Note [Chaotic Good fresh name generation]
Currently we don't use 'Quote' for generating fresh names in the QuickCheck generators. It's mostly
for historical reasons and it's very non-trivial to change the current alignment to use 'Quote' at
this point. It does have positive sides to it, though, which we'll discuss below.

So what we do right now is track all names that are in the scope (not all actually and not in the
scope, more on that later, for now let's pretend the statement is correct) in separate maps:

- one for type names
- one for term names
- one for datatype-associated names (the name of the data type, its matcher and constructors)

and all these maps are processed in a linear fashion each time we need to generate a bunch of fresh
names. The processing is done via collecting in a set all uniques of names found in the maps: the
key sets of all the three maps and the names in the elements of the data map. Why do we need the
latter given that the names there are present as keys in the term map? No idea.

What makes it hard to change the current alignment to use 'Quote' is that names are generated all
over the place and put into the current context, inside and outside of the 'GenTm' monad. Any
operation on the context automatically makes the name generation machinery aware of the changes,
because the machinery looks at the whole context each time. It's really hard to move all of that to
run in 'Quote' simply because it’s trivial to miss a call here and there and types don't help much.

Another challenge is that the current fresh name generation is pretty fancy and doesn’t just always
pick the next fresh name. It can pick the next 10th fresh name, then go back, then again forward etc
– all of that between different calls to the function generating fresh names. We want to keep this
logic as it allows us to produce more diverse terms. We probably can simulate this kind of logic
with 'Quote', but it's non-trivial for sure.

Finally, we do not have any function that generates a definitely fresh name or a function that
generates a definitely non-fresh name -- both for very technical reasons (see the docs inline).
This likely isn't intentional, but it's hard to change that at this point and we don't care much
anyway, name clashes are fine. Moreover, the term map rather than containing all term names that
are in the scope only contains those that can be correctly inserted into the program (i.e. don't
reference shadowed type variables in their type). However the data map isn't filtered the same way
as the term map and so there can be constructors whose type references shadowed type variables,
which we never notice because those are only used for fresh name generation... which I guess is how
looking at the elements of the data map isn't redundant in the end.

Overall, this chaotic goodness needs to be sorted out.
-}

-- | Get all uniques we have generated and are used in the current context.
getUniques :: GenTm (Set Unique)
getUniques :: GenTm (Set Unique)
getUniques = do
  GenEnv{geDatas :: GenEnv -> Map TyName (Datatype TyName Name DefaultUni ())
geDatas = Map TyName (Datatype TyName Name DefaultUni ())
dts, geTypes :: GenEnv -> TypeCtx
geTypes = TypeCtx
tys, geTerms :: GenEnv -> Map Name (Type TyName DefaultUni ())
geTerms = Map Name (Type TyName DefaultUni ())
tms, geUnboundUsedTyNames :: GenEnv -> Set TyName
geUnboundUsedTyNames = Set TyName
used} <- GenT (Reader GenEnv) GenEnv
forall r (m :: * -> *). MonadReader r m => m r
ask
  Set Unique -> GenTm (Set Unique)
forall a. a -> GenT (Reader GenEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return (Set Unique -> GenTm (Set Unique))
-> Set Unique -> GenTm (Set Unique)
forall a b. (a -> b) -> a -> b
$ (TyName -> Unique) -> Set TyName -> Set Unique
forall a b. (a -> b) -> Set a -> Set b
Set.mapMonotonic (Name -> Unique
_nameUnique (Name -> Unique) -> (TyName -> Name) -> TyName -> Unique
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyName -> Name
unTyName) (Map TyName (Datatype TyName Name DefaultUni ()) -> Set TyName
forall k a. Map k a -> Set k
Map.keysSet Map TyName (Datatype TyName Name DefaultUni ())
dts Set TyName -> Set TyName -> Set TyName
forall a. Semigroup a => a -> a -> a
<> TypeCtx -> Set TyName
forall k a. Map k a -> Set k
Map.keysSet TypeCtx
tys Set TyName -> Set TyName -> Set TyName
forall a. Semigroup a => a -> a -> a
<> Set TyName
used) Set Unique -> Set Unique -> Set Unique
forall a. Semigroup a => a -> a -> a
<>
           (Name -> Unique) -> Set Name -> Set Unique
forall a b. (a -> b) -> Set a -> Set b
Set.mapMonotonic Name -> Unique
_nameUnique (Map Name (Type TyName DefaultUni ()) -> Set Name
forall k a. Map k a -> Set k
Map.keysSet Map Name (Type TyName DefaultUni ())
tms) Set Unique -> Set Unique -> Set Unique
forall a. Semigroup a => a -> a -> a
<>
           [Set Unique] -> Set Unique
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions [ Datatype TyName Name DefaultUni () -> Set Unique
forall {tyname} {uni :: * -> *} {a}.
Datatype tyname Name uni a -> Set Unique
names Datatype TyName Name DefaultUni ()
d | Datatype TyName Name DefaultUni ()
d <- Map TyName (Datatype TyName Name DefaultUni ())
-> [Datatype TyName Name DefaultUni ()]
forall k a. Map k a -> [a]
Map.elems Map TyName (Datatype TyName Name DefaultUni ())
dts ]
  where
    names :: Datatype tyname Name uni a -> Set Unique
names (Datatype a
_ TyVarDecl tyname a
_ [TyVarDecl tyname a]
_ Name
m [VarDecl tyname Name uni a]
cs) = [Unique] -> Set Unique
forall a. Ord a => [a] -> Set a
Set.fromList ([Unique] -> Set Unique) -> [Unique] -> Set Unique
forall a b. (a -> b) -> a -> b
$ Name -> Unique
_nameUnique Name
m Unique -> [Unique] -> [Unique]
forall a. a -> [a] -> [a]
: [ Name -> Unique
_nameUnique Name
c | VarDecl a
_ Name
c Type tyname uni a
_ <- [VarDecl tyname Name uni a]
cs ]

{- Note [Warning about generating fresh names]
Since 'GenTm' is a *reader* monad names are not immediately put into any state when generated.
In order to bind generated names you need to use functions like 'bindTyName' and 'bindTmName', e.g.

    genLam k1 k2 = do
        x <- genMaybeFreshTyName "a"
        --                                              v--- LOOK HERE!
        fmap (TyLam () x k1) $ onAstSize (subtract 1) $ bindTyName x k1 (genType k2)
-}

-- | Generate a likely fresh name. See Note [Warning about generating fresh names].
genLikelyFreshName :: String -> GenTm Name
genLikelyFreshName :: [Char] -> GenTm Name
genLikelyFreshName [Char]
s = [Name] -> Name
forall a. HasCallStack => [a] -> a
head ([Name] -> Name) -> GenT (Reader GenEnv) [Name] -> GenTm Name
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[Char]] -> GenT (Reader GenEnv) [Name]
genLikelyFreshNames [[Char]
s]

-- See Note [Chaotic Good fresh name generation].
-- See Note [Warning about generating fresh names].
-- | Generate one likely fresh name per string in the input list.
-- Note that this may not give you a fresh name, if it happens to generate a name that was removed
-- from the terms map in 'bindTyName' (due to referencing a now-shadowed type variable) but is still
-- in the scope.
genLikelyFreshNames :: [String] -> GenTm [Name]
genLikelyFreshNames :: [[Char]] -> GenT (Reader GenEnv) [Name]
genLikelyFreshNames [[Char]]
ss = do
  Set Unique
used <- GenTm (Set Unique)
getUniques
  let i :: Int
i = Unique -> Int
forall a. Enum a => a -> Int
fromEnum (Unique -> Int) -> Unique -> Int
forall a b. (a -> b) -> a -> b
$ Set Unique -> Unique
forall a. Set a -> a
Set.findMax (Set Unique -> Unique) -> Set Unique -> Unique
forall a b. (a -> b) -> a -> b
$ Unique -> Set Unique -> Set Unique
forall a. Ord a => a -> Set a -> Set a
Set.insert (Int -> Unique
Unique Int
0) Set Unique
used
      js :: [Int]
js = [ Int
j | Int
j <- [Int
1..Int
i], Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Int -> Unique
Unique Int
j Unique -> Set Unique -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set Unique
used ]
      is :: [Int]
is = [Int]
js [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take ([[Char]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Char]]
ss Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
10) [Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1..]
  [Int]
is' <- Gen [Int] -> GenT (Reader GenEnv) [Int]
forall a. Gen a -> GenT (Reader GenEnv) a
forall (g :: * -> *) a. MonadGen g => Gen a -> g a
liftGen (Gen [Int] -> GenT (Reader GenEnv) [Int])
-> Gen [Int] -> GenT (Reader GenEnv) [Int]
forall a b. (a -> b) -> a -> b
$ [Int] -> Gen [Int]
forall a. [a] -> Gen [a]
QC.shuffle [Int]
is
  [Name] -> GenT (Reader GenEnv) [Name]
forall a. a -> GenT (Reader GenEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return [Text -> Unique -> Name
Name ([Char] -> Text
forall a. IsString a => [Char] -> a
fromString ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ [Char]
s [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
j) (Int -> Unique
forall a. Enum a => Int -> a
toEnum Int
j) | ([Char]
s, Int
j) <- [[Char]] -> [Int] -> [([Char], Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [[Char]]
ss [Int]
is']

-- | Same as 'genLikelyFreshName', except gives you a 'TyName'.
genLikelyFreshTyName :: String -> GenTm TyName
genLikelyFreshTyName :: [Char] -> GenTm TyName
genLikelyFreshTyName [Char]
s = Name -> TyName
TyName (Name -> TyName) -> GenTm Name -> GenTm TyName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> GenTm Name
genLikelyFreshName [Char]
s

-- | Same as 'genLikelyFreshNames', except gives you 'TyName's.
genLikelyFreshTyNames :: [String] -> GenTm [TyName]
genLikelyFreshTyNames :: [[Char]] -> GenTm [TyName]
genLikelyFreshTyNames [[Char]]
ss = (Name -> TyName) -> [Name] -> [TyName]
forall a b. (a -> b) -> [a] -> [b]
map Name -> TyName
TyName ([Name] -> [TyName])
-> GenT (Reader GenEnv) [Name] -> GenTm [TyName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[Char]] -> GenT (Reader GenEnv) [Name]
genLikelyFreshNames [[Char]]
ss

-- See Note [Chaotic Good fresh name generation].
-- | Generate a name that likely overlaps with existing names on purpose. If there are no existing
-- names, generate a fresh name. This function doesn't distinguish between the type- and term-level
-- scopes, hence it may generate a 'Name' \"clashing\" with a previously generated 'TyName' and not
-- clashing with any previously generated 'Name'. Which is a good thing, because we want to check
-- that Plutus is able to handle such spurious name clashes. Generating weird stuff is useful for a
-- testing machinery! We don't need any \"definitely non-fresh name\" anyway, since we get enough
-- non-fresh names already.
genLikelyNotFreshName :: String -> GenTm Name
genLikelyNotFreshName :: [Char] -> GenTm Name
genLikelyNotFreshName [Char]
s = do
  [Unique]
used <- Set Unique -> [Unique]
forall a. Set a -> [a]
Set.toList (Set Unique -> [Unique])
-> GenTm (Set Unique) -> GenT (Reader GenEnv) [Unique]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GenTm (Set Unique)
getUniques
  case [Unique]
used of
    [] -> [Char] -> GenTm Name
genLikelyFreshName [Char]
s
    [Unique]
_  -> Gen Name -> GenTm Name
forall a. Gen a -> GenT (Reader GenEnv) a
forall (g :: * -> *) a. MonadGen g => Gen a -> g a
liftGen (Gen Name -> GenTm Name) -> Gen Name -> GenTm Name
forall a b. (a -> b) -> a -> b
$ [Name] -> Gen Name
forall (m :: * -> *) a. MonadGen m => [a] -> m a
elements [ Text -> Unique -> Name
Name ([Char] -> Text
forall a. IsString a => [Char] -> a
fromString ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ [Char]
s [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show (Unique -> Int
unUnique Unique
i)) Unique
i | Unique
i <- [Unique]
used ]

-- | Generate a fresh name most (roughly 75%) of the time and otherwise
-- generate an already bound name. When there are no bound names generate a fresh name.
genMaybeFreshName :: String -> GenTm Name
genMaybeFreshName :: [Char] -> GenTm Name
genMaybeFreshName [Char]
s = [(Int, GenTm Name)] -> GenTm Name
forall (m :: * -> *) a. MonadGen m => [(Int, m a)] -> m a
frequency [(Int
3, [Char] -> GenTm Name
genLikelyFreshName [Char]
s), (Int
1, [Char] -> GenTm Name
genLikelyNotFreshName [Char]
s)]

-- | See `genMaybeFreshName`
genMaybeFreshTyName :: String -> GenTm TyName
genMaybeFreshTyName :: [Char] -> GenTm TyName
genMaybeFreshTyName [Char]
s = Name -> TyName
TyName (Name -> TyName) -> GenTm Name -> GenTm TyName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> GenTm Name
genMaybeFreshName [Char]
s

-- | Bind a type name to a kind and avoid capturing free type variables.
bindTyName :: TyName -> Kind () -> GenTm a -> GenTm a
bindTyName :: forall a. TyName -> Kind () -> GenTm a -> GenTm a
bindTyName TyName
x Kind ()
k = (GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a
forall a.
(GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((GenEnv -> GenEnv)
 -> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a)
-> (GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a
-> GenT (Reader GenEnv) a
forall a b. (a -> b) -> a -> b
$ \ GenEnv
e -> GenEnv
e
    { geTypes = Map.insert x k (geTypes e)
      -- See Note [Chaotic Good fresh name generation].
    , geTerms = Map.filter (\Type TyName DefaultUni ()
ty -> Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ TyName
x TyName -> Set TyName -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Getting (Set TyName) (Type TyName DefaultUni ()) TyName
-> Type TyName DefaultUni () -> Set TyName
forall a s. Getting (Set a) s a -> s -> Set a
setOf Getting (Set TyName) (Type TyName DefaultUni ()) TyName
forall tyname unique (uni :: * -> *) ann.
HasUnique tyname unique =>
Traversal' (Type tyname uni ann) tyname
Traversal' (Type TyName DefaultUni ()) TyName
ftvTy Type TyName DefaultUni ()
ty) (geTerms e)
    , geDatas = Map.delete x (geDatas e)
    }

-- | Bind type names
bindTyNames :: [(TyName, Kind ())] -> GenTm a -> GenTm a
bindTyNames :: forall a. [(TyName, Kind ())] -> GenTm a -> GenTm a
bindTyNames = (GenTm a -> [(TyName, Kind ())] -> GenTm a)
-> [(TyName, Kind ())] -> GenTm a -> GenTm a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((GenTm a -> [(TyName, Kind ())] -> GenTm a)
 -> [(TyName, Kind ())] -> GenTm a -> GenTm a)
-> (GenTm a -> [(TyName, Kind ())] -> GenTm a)
-> [(TyName, Kind ())]
-> GenTm a
-> GenTm a
forall a b. (a -> b) -> a -> b
$ ((TyName, Kind ()) -> GenTm a -> GenTm a)
-> GenTm a -> [(TyName, Kind ())] -> GenTm a
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 () -> GenTm a -> GenTm a)
-> (TyName, Kind ()) -> GenTm a -> GenTm a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry TyName -> Kind () -> GenTm a -> GenTm a
forall a. TyName -> Kind () -> GenTm a -> GenTm a
bindTyName)

-- | Remember that we have generated a type name locally but don't bind it.
-- Useful for non-recursive definitions where we want to control name overlap.
registerTyName :: TyName -> GenTm a -> GenTm a
registerTyName :: forall a. TyName -> GenTm a -> GenTm a
registerTyName TyName
n = (GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a
forall a.
(GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((GenEnv -> GenEnv)
 -> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a)
-> (GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a
-> GenT (Reader GenEnv) a
forall a b. (a -> b) -> a -> b
$ \ GenEnv
e -> GenEnv
e { geUnboundUsedTyNames = Set.insert n (geUnboundUsedTyNames e) }

-- | Bind a term to a type in a generator.
bindTmName :: Name -> Type TyName DefaultUni () -> GenTm a -> GenTm a
bindTmName :: forall a. Name -> Type TyName DefaultUni () -> GenTm a -> GenTm a
bindTmName Name
x Type TyName DefaultUni ()
ty = (GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a
forall a.
(GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((GenEnv -> GenEnv)
 -> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a)
-> (GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a
-> GenT (Reader GenEnv) a
forall a b. (a -> b) -> a -> b
$ \ GenEnv
e -> GenEnv
e { geTerms = Map.insert x ty (geTerms e) }

-- | Bind term names
bindTmNames :: [(Name, Type TyName DefaultUni ())] -> GenTm a -> GenTm a
bindTmNames :: forall a. [(Name, Type TyName DefaultUni ())] -> GenTm a -> GenTm a
bindTmNames = (GenTm a -> [(Name, Type TyName DefaultUni ())] -> GenTm a)
-> [(Name, Type TyName DefaultUni ())] -> GenTm a -> GenTm a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((GenTm a -> [(Name, Type TyName DefaultUni ())] -> GenTm a)
 -> [(Name, Type TyName DefaultUni ())] -> GenTm a -> GenTm a)
-> (GenTm a -> [(Name, Type TyName DefaultUni ())] -> GenTm a)
-> [(Name, Type TyName DefaultUni ())]
-> GenTm a
-> GenTm a
forall a b. (a -> b) -> a -> b
$ ((Name, Type TyName DefaultUni ()) -> GenTm a -> GenTm a)
-> GenTm a -> [(Name, Type TyName DefaultUni ())] -> GenTm a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Name -> Type TyName DefaultUni () -> GenTm a -> GenTm a)
-> (Name, Type TyName DefaultUni ()) -> GenTm a -> GenTm a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Name -> Type TyName DefaultUni () -> GenTm a -> GenTm a
forall a. Name -> Type TyName DefaultUni () -> GenTm a -> GenTm a
bindTmName)

-- | Create a fresh term name, bind it to a type, and use it in a generator.
bindLikelyFreshTmName :: String -> Type TyName DefaultUni () -> (Name -> GenTm a) -> GenTm a
bindLikelyFreshTmName :: forall a.
[Char] -> Type TyName DefaultUni () -> (Name -> GenTm a) -> GenTm a
bindLikelyFreshTmName [Char]
name Type TyName DefaultUni ()
ty Name -> GenTm a
k = do
  Name
x <- [Char] -> GenTm Name
genLikelyFreshName [Char]
name
  Name -> Type TyName DefaultUni () -> GenTm a -> GenTm a
forall a. Name -> Type TyName DefaultUni () -> GenTm a -> GenTm a
bindTmName Name
x Type TyName DefaultUni ()
ty (Name -> GenTm a
k Name
x)

-- | Bind a datatype declaration in a generator.
bindDat :: Datatype TyName Name DefaultUni ()
        -> GenTm a
        -> GenTm a
bindDat :: forall a. Datatype TyName Name DefaultUni () -> GenTm a -> GenTm a
bindDat dat :: Datatype TyName Name DefaultUni ()
dat@(Datatype ()
_ (TyVarDecl ()
_ TyName
a Kind ()
k) [TyVarDecl TyName ()]
_ Name
_ [VarDecl TyName Name DefaultUni ()]
_) GenTm a
cont =
  TyName -> Kind () -> GenTm a -> GenTm a
forall a. TyName -> Kind () -> GenTm a -> GenTm a
bindTyName TyName
a Kind ()
k (GenTm a -> GenTm a) -> GenTm a -> GenTm a
forall a b. (a -> b) -> a -> b
$
  (GenEnv -> GenEnv) -> GenTm a -> GenTm a
forall a.
(GenEnv -> GenEnv)
-> GenT (Reader GenEnv) a -> GenT (Reader GenEnv) a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (\ GenEnv
e -> GenEnv
e { geDatas = Map.insert a dat (geDatas e) }) (GenTm a -> GenTm a) -> GenTm a -> GenTm a
forall a b. (a -> b) -> a -> b
$
  ((Name, Type TyName DefaultUni ()) -> GenTm a -> GenTm a)
-> GenTm a -> [(Name, Type TyName DefaultUni ())] -> GenTm a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Name -> Type TyName DefaultUni () -> GenTm a -> GenTm a)
-> (Name, Type TyName DefaultUni ()) -> GenTm a -> GenTm a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Name -> Type TyName DefaultUni () -> GenTm a -> GenTm a
forall a. Name -> Type TyName DefaultUni () -> GenTm a -> GenTm a
bindTmName) GenTm a
cont (Datatype TyName Name DefaultUni ()
-> (Name, Type TyName DefaultUni ())
matchType Datatype TyName Name DefaultUni ()
dat (Name, Type TyName DefaultUni ())
-> [(Name, Type TyName DefaultUni ())]
-> [(Name, Type TyName DefaultUni ())]
forall a. a -> [a] -> [a]
: Datatype TyName Name DefaultUni ()
-> [(Name, Type TyName DefaultUni ())]
constrTypes Datatype TyName Name DefaultUni ()
dat)

-- | Bind a binding.
bindBind :: Binding TyName Name DefaultUni DefaultFun ()
         -> GenTm a
         -> GenTm a
bindBind :: forall a.
Binding TyName Name DefaultUni DefaultFun () -> GenTm a -> GenTm a
bindBind (DatatypeBind ()
_ Datatype TyName Name DefaultUni ()
dat)              = Datatype TyName Name DefaultUni () -> GenTm a -> GenTm a
forall a. Datatype TyName Name DefaultUni () -> GenTm a -> GenTm a
bindDat Datatype TyName Name DefaultUni ()
dat
bindBind (TermBind ()
_ Strictness
_ (VarDecl ()
_ Name
x Type TyName DefaultUni ()
ty) Term TyName Name DefaultUni DefaultFun ()
_) = Name -> Type TyName DefaultUni () -> GenTm a -> GenTm a
forall a. Name -> Type TyName DefaultUni () -> GenTm a -> GenTm a
bindTmName Name
x Type TyName DefaultUni ()
ty
-- TODO: We should generate type bindings
bindBind Binding TyName Name DefaultUni DefaultFun ()
_                                 = [Char] -> GenTm a -> GenTm a
forall a. HasCallStack => [Char] -> a
error [Char]
"unreachable"

-- | Bind multiple bindings
bindBinds :: Foldable f => f (Binding TyName Name DefaultUni DefaultFun ()) -> GenTm a -> GenTm a
bindBinds :: forall (f :: * -> *) a.
Foldable f =>
f (Binding TyName Name DefaultUni DefaultFun ())
-> GenTm a -> GenTm a
bindBinds = (GenTm a
 -> f (Binding TyName Name DefaultUni DefaultFun ()) -> GenTm a)
-> f (Binding TyName Name DefaultUni DefaultFun ())
-> GenTm a
-> GenTm a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Binding TyName Name DefaultUni DefaultFun ()
 -> GenTm a -> GenTm a)
-> GenTm a
-> f (Binding TyName Name DefaultUni DefaultFun ())
-> GenTm a
forall a b. (a -> b -> b) -> b -> f a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Binding TyName Name DefaultUni DefaultFun () -> GenTm a -> GenTm a
forall a.
Binding TyName Name DefaultUni DefaultFun () -> GenTm a -> GenTm a
bindBind)