{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies      #-}

module PlutusCore.Generators.QuickCheck.Utils
    ( module PlutusCore.Generators.QuickCheck.Utils
    , module Export
    ) where

import PlutusCore.Default
import PlutusCore.Generators.QuickCheck.Split as Export
import PlutusCore.MkPlc hiding (error)
import PlutusCore.Name.Unique
import PlutusCore.Pretty
import PlutusCore.Quote
import PlutusIR
import PlutusIR.Compiler.Datatype
import PlutusIR.Subst

import Data.Kind qualified as GHC
import Data.List.NonEmpty (NonEmpty (..))
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Set.Lens (setOf)
import Data.String
import Prettyprinter
import Test.QuickCheck

-- | Generate a list of the given length, all arguments of which are distinct. Takes O(n^2) time
-- or more if the generator is likely to generate equal values.
uniqueVectorOf :: Eq a => Int -> Gen a -> Gen [a]
uniqueVectorOf :: forall a. Eq a => Int -> Gen a -> Gen [a]
uniqueVectorOf Int
i0 Gen a
genX = [a] -> Int -> Gen [a]
go [] Int
i0 where
    go :: [a] -> Int -> Gen [a]
go [a]
acc Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0    = [a] -> Gen [a]
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
acc
        | Bool
otherwise = do
              a
x <- Gen a
genX Gen a -> (a -> Bool) -> Gen a
forall a. Gen a -> (a -> Bool) -> Gen a
`suchThat` (a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [a]
acc)
              [a] -> Int -> Gen [a]
go (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

-- | Show a `Doc` when a property fails.
ceDoc :: Testable t => Doc ann -> t -> Property
ceDoc :: forall t ann. Testable t => Doc ann -> t -> Property
ceDoc Doc ann
d = String -> t -> Property
forall prop. Testable prop => String -> prop -> Property
counterexample (Doc ann -> String
forall ann. Doc ann -> String
forall str ann. Render str => Doc ann -> str
render Doc ann
d)

-- | Bind a value to a name in a property so that
-- it is displayed as a `name = thing` binding if the
-- property fails.
letCE :: (PrettyReadable a, Testable p)
      => String
      -> a
      -> (a -> p)
      -> Property
letCE :: forall a p.
(PrettyReadable a, Testable p) =>
String -> a -> (a -> p) -> Property
letCE String
name a
x a -> p
k = Doc Any -> p -> Property
forall t ann. Testable t => Doc ann -> t -> Property
ceDoc (String -> Doc Any
forall a. IsString a => String -> a
fromString String
name Doc Any -> Doc Any -> Doc Any
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc Any
"=" Doc Any -> Doc Any -> Doc Any
forall ann. Doc ann -> Doc ann -> Doc ann
<+> a -> Doc Any
forall a ann. PrettyReadable a => a -> Doc ann
prettyReadable a
x) (a -> p
k a
x)

-- | Like `forAllShrink` but displays the bound value as
-- a named pretty-printed binding like `letCE`
forAllDoc :: (PrettyReadable a, Testable p)
          => String
          -> Gen a
          -> (a -> [a])
          -> (a -> p)
          -> Property
forAllDoc :: forall a p.
(PrettyReadable a, Testable p) =>
String -> Gen a -> (a -> [a]) -> (a -> p) -> Property
forAllDoc String
name Gen a
g a -> [a]
shr a -> p
k =
  Gen a -> (a -> [a]) -> (a -> Property) -> Property
forall prop a.
Testable prop =>
Gen a -> (a -> [a]) -> (a -> prop) -> Property
forAllShrinkBlind Gen a
g a -> [a]
shr ((a -> Property) -> Property) -> (a -> Property) -> Property
forall a b. (a -> b) -> a -> b
$ \ a
x ->
    Doc Any -> p -> Property
forall t ann. Testable t => Doc ann -> t -> Property
ceDoc (String -> Doc Any
forall a. IsString a => String -> a
fromString String
name Doc Any -> Doc Any -> Doc Any
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc Any
"=" Doc Any -> Doc Any -> Doc Any
forall ann. Doc ann -> Doc ann -> Doc ann
<+> a -> Doc Any
forall a ann. PrettyReadable a => a -> Doc ann
prettyReadable a
x)
          (a -> p
k a
x)

-- | Check that a list of potential counterexamples is empty and display the
-- list as a QuickCheck counterexample if its not.
assertNoCounterexamples :: (PrettyReadable a) => [a] -> Property
assertNoCounterexamples :: forall a. PrettyReadable a => [a] -> Property
assertNoCounterexamples []  = Bool -> Property
forall prop. Testable prop => prop -> Property
property Bool
True
assertNoCounterexamples [a]
bad = Doc Any -> Bool -> Property
forall t ann. Testable t => Doc ann -> t -> Property
ceDoc ([a] -> Doc Any
forall a ann. PrettyReadable a => a -> Doc ann
prettyReadable [a]
bad) Bool
False

-- * Containers (zipper-ish, very useful for shrinking).

-- There's https://hackage.haskell.org/package/functor-combo-0.3.6/docs/FunctorCombo-Holey.html
-- (explained in http://conal.net/blog/posts/differentiation-of-higher-order-types), but it's kinda
-- easier to roll out our own version than to dig through the generics in that library. Plus, the
-- library has multiple interfaces looking slightly different and it's not immediately clear how
-- they relate to each other.
-- | A type is a container for the purposes of shrinking if it has:
class Container f where
  data OneHoleContext f :: GHC.Type -> GHC.Type
  -- ^ One hole context where we can shrink a single "element" of the container
  oneHoleContexts :: f a -> [(OneHoleContext f a, a)]
  -- ^ A way of getting all the one hole contexts of an `f a`
  plugHole :: OneHoleContext f a -> a -> f a
  -- ^ A way to plug the hole with a new, shrunk, term

-- | Containers for lists is zipper like, a hole is a specific position in the list
instance Container [] where
  data OneHoleContext [] a = ListContext [a] [a]

  oneHoleContexts :: forall a. [a] -> [(OneHoleContext [] a, a)]
oneHoleContexts [] = []
  oneHoleContexts (a
x : [a]
xs)
    = ([a] -> [a] -> OneHoleContext [] a
forall a. [a] -> [a] -> OneHoleContext [] a
ListContext [] [a]
xs, a
x)
    (OneHoleContext [] a, a)
-> [(OneHoleContext [] a, a)] -> [(OneHoleContext [] a, a)]
forall a. a -> [a] -> [a]
: [ ([a] -> [a] -> OneHoleContext [] a
forall a. [a] -> [a] -> OneHoleContext [] a
ListContext (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ys) [a]
zs, a
y)
      | (ListContext [a]
ys [a]
zs, a
y) <- [a] -> [(OneHoleContext [] a, a)]
forall a. [a] -> [(OneHoleContext [] a, a)]
forall (f :: * -> *) a.
Container f =>
f a -> [(OneHoleContext f a, a)]
oneHoleContexts [a]
xs
      ]

  plugHole :: forall a. OneHoleContext [] a -> a -> [a]
plugHole (ListContext [a]
xs [a]
ys) a
z = [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
z] [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
ys

-- | An analogous implementation of `Container` as for lists
instance Container NonEmpty where
  data OneHoleContext NonEmpty a = NonEmptyContext [a] [a]

  oneHoleContexts :: forall a. NonEmpty a -> [(OneHoleContext NonEmpty a, a)]
oneHoleContexts (a
x :| [a]
xs)
    = ([a] -> [a] -> OneHoleContext NonEmpty a
forall a. [a] -> [a] -> OneHoleContext NonEmpty a
NonEmptyContext [] [a]
xs, a
x)
    (OneHoleContext NonEmpty a, a)
-> [(OneHoleContext NonEmpty a, a)]
-> [(OneHoleContext NonEmpty a, a)]
forall a. a -> [a] -> [a]
: [ ([a] -> [a] -> OneHoleContext NonEmpty a
forall a. [a] -> [a] -> OneHoleContext NonEmpty a
NonEmptyContext (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ys) [a]
zs, a
y)
      | (ListContext [a]
ys [a]
zs, a
y) <- [a] -> [(OneHoleContext [] a, a)]
forall a. [a] -> [(OneHoleContext [] a, a)]
forall (f :: * -> *) a.
Container f =>
f a -> [(OneHoleContext f a, a)]
oneHoleContexts [a]
xs
      ]

  plugHole :: forall a. OneHoleContext NonEmpty a -> a -> NonEmpty a
plugHole (NonEmptyContext []       [a]
ys) a
z = a
z a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
ys
  plugHole (NonEmptyContext (a
x : [a]
xs) [a]
ys) a
z = a
x a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
z] [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
ys

-- Note that this doesn't have any relation to 'freshenTyName'. Both functions change the unique,
-- but do it in incomparably different ways.
-- | Freshen a TyName so that it does not equal any of the names in the set.
freshenTyNameWith :: Set TyName -> TyName -> TyName
freshenTyNameWith :: Set TyName -> TyName -> TyName
freshenTyNameWith Set TyName
fvs (TyName (Name Text
x Unique
j)) = Name -> TyName
TyName (Text -> Unique -> Name
Name Text
x Unique
i) where
  i :: Unique
i  = Unique -> Unique
forall a. Enum a => a -> a
succ (Unique -> Unique) -> Unique -> Unique
forall a b. (a -> b) -> a -> b
$ Set Unique -> Unique
forall a. Set a -> a
Set.findMax Set Unique
is
  is :: Set Unique
is = Unique -> Set Unique -> Set Unique
forall a. Ord a => a -> Set a -> Set a
Set.insert Unique
j (Set Unique -> Set Unique) -> Set Unique -> Set 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
forall a. Enum a => Int -> a
toEnum Int
0) (Set Unique -> Set Unique) -> Set Unique -> 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) Set TyName
fvs

maxUsedUnique :: Set TyName -> Unique
maxUsedUnique :: Set TyName -> Unique
maxUsedUnique Set TyName
fvs = Unique
i
  where
    i :: Unique
i  = Set Unique -> Unique
forall a. Set a -> a
Set.findMax Set Unique
is
    is :: Set Unique
is = Unique -> Set Unique -> Set Unique
forall a. Ord a => a -> Set a -> Set a
Set.insert (Int -> Unique
forall a. Enum a => Int -> a
toEnum Int
0) (Set Unique -> Set Unique) -> Set Unique -> 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) Set TyName
fvs

-- | Get the names and types of the constructors of a datatype.
constrTypes :: Datatype TyName Name DefaultUni () -> [(Name, Type TyName DefaultUni ())]
constrTypes :: Datatype TyName Name DefaultUni ()
-> [(Name, Type TyName DefaultUni ())]
constrTypes (Datatype ()
_ TyVarDecl TyName ()
_ [TyVarDecl TyName ()]
xs Name
_ [VarDecl TyName Name DefaultUni ()]
cs) = [ (Name
c, [TyVarDecl TyName ()]
-> Type TyName DefaultUni () -> Type TyName DefaultUni ()
forall tyname ann (uni :: * -> *).
[TyVarDecl tyname ann]
-> Type tyname uni ann -> Type tyname uni ann
mkIterTyForall [TyVarDecl TyName ()]
xs Type TyName DefaultUni ()
ty) | VarDecl ()
_ Name
c Type TyName DefaultUni ()
ty <- [VarDecl TyName Name DefaultUni ()]
cs ]

-- | Get the name and type of the match function for a given datatype.
matchType :: Datatype TyName Name DefaultUni () -> (Name, Type TyName DefaultUni ())
matchType :: Datatype TyName Name DefaultUni ()
-> (Name, Type TyName DefaultUni ())
matchType d :: Datatype TyName Name DefaultUni ()
d@(Datatype ()
_ (TyVarDecl ()
_ TyName
a Kind ()
_) [TyVarDecl TyName ()]
xs Name
m [VarDecl TyName Name DefaultUni ()]
cs) = (Name
m, Type TyName DefaultUni ()
destrTy)
  where
    fvs :: Set TyName
fvs = [TyName] -> Set TyName
forall a. Ord a => [a] -> Set a
Set.fromList (TyName
a TyName -> [TyName] -> [TyName]
forall a. a -> [a] -> [a]
: [TyName
x | TyVarDecl ()
_ TyName
x Kind ()
_ <- [TyVarDecl TyName ()]
xs]) Set TyName -> Set TyName -> Set TyName
forall a. Semigroup a => a -> a -> a
<>
          [Set TyName] -> Set TyName
forall a. Monoid a => [a] -> a
mconcat [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 | VarDecl ()
_ Name
_ Type TyName DefaultUni ()
ty <- [VarDecl TyName Name DefaultUni ()]
cs]
    maxUsed :: Unique
maxUsed = Set TyName -> Unique
maxUsedUnique Set TyName
fvs
    destrTy :: Type TyName DefaultUni ()
destrTy = Quote (Type TyName DefaultUni ()) -> Type TyName DefaultUni ()
forall a. Quote a -> a
runQuote (Quote (Type TyName DefaultUni ()) -> Type TyName DefaultUni ())
-> Quote (Type TyName DefaultUni ()) -> Type TyName DefaultUni ()
forall a b. (a -> b) -> a -> b
$ Unique -> QuoteT Identity ()
forall (m :: * -> *). MonadQuote m => Unique -> m ()
markNonFresh Unique
maxUsed QuoteT Identity ()
-> Quote (Type TyName DefaultUni ())
-> Quote (Type TyName DefaultUni ())
forall a b.
QuoteT Identity a -> QuoteT Identity b -> QuoteT Identity b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Datatype TyName Name DefaultUni ()
-> Quote (Type TyName DefaultUni ())
forall (m :: * -> *) (uni :: * -> *) a.
MonadQuote m =>
Datatype TyName Name uni a -> m (Type TyName uni a)
mkDestructorTy Datatype TyName Name DefaultUni ()
d