[Test] Improve distribution of generated integers (#6315)

This improves distribution of generated integers, so that we more often hit important edge cases such as `2 ^ 16`, `2 ^ 32 - 1`, `2 ^ 32`, `2 ^ 64` etc.

plutus-core/plutus-core.cabal

@@ -105,6 +105,7 @@ library
  PlutusCore.Compiler.Opts
  PlutusCore.Compiler.Types
  PlutusCore.Core
+ PlutusCore.Core.Plated
  PlutusCore.Crypto.BLS12_381.Error
  PlutusCore.Crypto.BLS12_381.G1
  PlutusCore.Crypto.BLS12_381.G2
@@ -192,6 +193,7 @@ library
  UntypedPlutusCore.Check.Scope
  UntypedPlutusCore.Check.Uniques
  UntypedPlutusCore.Core
+ UntypedPlutusCore.Core.Plated
  UntypedPlutusCore.Core.Type
  UntypedPlutusCore.Core.Zip
  UntypedPlutusCore.DeBruijn
@@ -229,7 +231,6 @@ library
  PlutusCore.Core.Instance.Pretty.Plc
  PlutusCore.Core.Instance.Pretty.Readable
  PlutusCore.Core.Instance.Scoping
- PlutusCore.Core.Plated
  PlutusCore.Core.Type
  PlutusCore.Crypto.Utils
  PlutusCore.Default.Universe
@@ -256,7 +257,6 @@ library
  UntypedPlutusCore.Core.Instance.Pretty.Default
  UntypedPlutusCore.Core.Instance.Pretty.Plc
  UntypedPlutusCore.Core.Instance.Pretty.Readable
- UntypedPlutusCore.Core.Plated
  UntypedPlutusCore.Evaluation.Machine.Cek.EmitterMode
  UntypedPlutusCore.Evaluation.Machine.Cek.ExBudgetMode
  UntypedPlutusCore.Evaluation.Machine.CommonAPI
@@ -817,6 +817,7 @@ library plutus-core-testlib
  PlutusIR.Pass.Test
  PlutusIR.Test
  Test.Tasty.Extras
+ UntypedPlutusCore.Generators.Hedgehog
  UntypedPlutusCore.Test.DeBruijn.Bad
  UntypedPlutusCore.Test.DeBruijn.Good
 
@@ -831,6 +832,7 @@ library plutus-core-testlib
  , free
  , hashable
  , hedgehog >=1.0
+ , hedgehog-quickcheck
  , lazy-search
  , lens
  , mmorph

plutus-core/plutus-core/src/PlutusCore/Core/Plated.hs

@@ -12,15 +12,17 @@ module PlutusCore.Core.Plated
  , typeSubtypes
  , typeSubtypesDeep
  , varDeclSubtypes
+ , termConstants
  , termTyBinds
  , termBinds
  , termVars
  , termUniques
  , termSubkinds
  , termSubtypes
+ , termSubtermsDeep
  , termSubtypesDeep
+ , termConstantsDeep
  , termSubterms
- , termSubtermsDeep
  , typeUniquesDeep
  , termUniquesDeep
  ) where
@@ -31,6 +33,7 @@ import PlutusCore.Core.Type
 import PlutusCore.Name.Unique
 
 import Control.Lens
+import Universe
 
 kindSubkinds :: Traversal' (Kind ann) (Kind ann)
 kindSubkinds f kind0 = case kind0 of
@@ -116,6 +119,22 @@ typeSubtypesDeep = cosmosOf typeSubtypes
 varDeclSubtypes :: Traversal' (VarDecl tyname name uni a) (Type tyname uni a)
 varDeclSubtypes f (VarDecl a n ty) = VarDecl a n <$> f ty
 
+-- | Get all the direct constants of the given 'Term' from 'Constant's.
+termConstants :: Traversal' (Term tyname name uni fun ann) (Some (ValueOf uni))
+termConstants f term0 = case term0 of
+ Constant ann val -> Constant ann <$> f val
+ Var{} -> pure term0
+ TyAbs{} -> pure term0
+ LamAbs{} -> pure term0
+ TyInst{} -> pure term0
+ IWrap{} -> pure term0
+ Error{} -> pure term0
+ Apply{} -> pure term0
+ Unwrap{} -> pure term0
+ Builtin{} -> pure term0
+ Constr{} -> pure term0
+ Case{} -> pure term0
+
 -- | Get all the direct child 'tyname a's of the given 'Term' from 'TyAbs'es.
 termTyBinds :: Traversal' (Term tyname name uni fun ann) tyname
 termTyBinds f term0 = case term0 of
@@ -214,6 +233,10 @@ termSubtypes f term0 = case term0 of
  Constant{} -> pure term0
  Builtin{} -> pure term0
 
+-- | Get all the transitive child 'Constant's of the given 'Term'.
+termConstantsDeep :: Fold (Term tyname name uni fun ann) (Some (ValueOf uni))
+termConstantsDeep = termSubtermsDeep . termConstants
+
 -- | Get all the transitive child 'Type's of the given 'Term'.
 termSubtypesDeep :: Fold (Term tyname name uni fun ann) (Type tyname uni ann)
 termSubtypesDeep = termSubtermsDeep . termSubtypes . typeSubtypesDeep

plutus-core/plutus-core/test/Parser/Spec.hs

@@ -7,6 +7,7 @@ module Parser.Spec (tests) where
 import PlutusCore
 import PlutusCore.Error (ParserErrorBundle)
 import PlutusCore.Generators.Hedgehog.AST
+import PlutusCore.Test (isSerialisable)
 import PlutusPrelude
 
 import Data.Text qualified as T
@@ -19,7 +20,8 @@ import Test.Tasty.Hedgehog
 -- | The `SrcSpan` of a parsed `Term` should not including trailing whitespaces.
 propTermSrcSpan :: Property
 propTermSrcSpan = property $ do
- term <- forAllWith display (runAstGen genTerm)
+ term <- _progTerm <$>
+ forAllWith display (runAstGen $ discardIfAnyConstant (not . isSerialisable) genProgram)
  let code = display (term :: Term TyName Name DefaultUni DefaultFun ())
  let (endingLine, endingCol) = length &&& T.length . last $ T.lines code
  trailingSpaces <- forAll $ Gen.text (Range.linear 0 10) (Gen.element [' ', '\n'])

plutus-core/plutus-core/test/Spec.hs

@@ -140,7 +140,8 @@ instance (Eq a) => Eq (TextualProgram a) where
 
 propFlat :: Property
 propFlat = property $ do
- prog <- forAllPretty $ runAstGen (genProgram @DefaultFun)
+ prog <- forAllPretty . runAstGen $
+ discardIfAnyConstant (not . isSerialisable) $ genProgram @DefaultFun
  Hedgehog.tripping prog Flat.flat Flat.unflat
 
 {- The following tests check that (A) the parser can
@@ -222,7 +223,8 @@ text, hopefully returning the same thing.
 -}
 propParser :: Property
 propParser = property $ do
- prog <- TextualProgram <$> forAllPretty (runAstGen genProgram)
+ prog <- TextualProgram <$>
+ forAllPretty (runAstGen $ discardIfAnyConstant (not . isSerialisable) genProgram)
  Hedgehog.tripping
  prog
  (displayPlc . unTextualProgram)

plutus-core/plutus-ir/src/PlutusIR.hs

@@ -1,6 +1,9 @@
 module PlutusIR (
  -- * AST
  Term (..),
+ progAnn,
+ progVer,
+ progTerm,
  termSubterms,
  termSubtypes,
  termBindings,

plutus-core/plutus-ir/src/PlutusIR/Core/Plated.hs

@@ -10,6 +10,8 @@ module PlutusIR.Core.Plated
  , termSubkinds
  , termBindings
  , termVars
+ , termConstants
+ , termConstantsDeep
  , typeSubtypes
  , typeSubtypesDeep
  , typeSubkinds
@@ -43,6 +45,7 @@ import PlutusIR.Core.Type
 import Control.Lens hiding (Strict, (<.>))
 import Data.Functor.Apply
 import Data.Functor.Bind.Class
+import Universe
 
 infixr 6 <^>
 
@@ -115,6 +118,23 @@ bindingIds f = \case
  <.> PLC.theUnique f n
  <.*> traverse1Maybe ((PLC.varDeclName . PLC.theUnique) f) vdecls)
 
+-- | Get all the direct constants of the given 'Term' from 'Constant's.
+termConstants :: Traversal' (Term tyname name uni fun ann) (Some (ValueOf uni))
+termConstants f term0 = case term0 of
+ Constant ann val -> Constant ann <$> f val
+ Let{} -> pure term0
+ Var{} -> pure term0
+ TyAbs{} -> pure term0
+ LamAbs{} -> pure term0
+ TyInst{} -> pure term0
+ IWrap{} -> pure term0
+ Error{} -> pure term0
+ Apply{} -> pure term0
+ Unwrap{} -> pure term0
+ Builtin{} -> pure term0
+ Constr{} -> pure term0
+ Case{} -> pure term0
+
 {-# INLINE termSubkinds #-}
 -- | Get all the direct child 'Kind's of the given 'Term'.
 termSubkinds :: Traversal' (Term tyname name uni fun ann) (Kind ann)
@@ -209,6 +229,10 @@ termVars f term0 = case term0 of
  Var ann n -> Var ann <$> f n
  t -> pure t
 
+-- | Get all the transitive child 'Constant's of the given 'Term'.
+termConstantsDeep :: Fold (Term tyname name uni fun ann) (Some (ValueOf uni))
+termConstantsDeep = termSubtermsDeep . termConstants
+
 -- | Get all the transitive child 'Unique's of the given 'Term' (including the type-level ones).
 termUniquesDeep
  :: PLC.HasUniques (Term tyname name uni fun ann)

plutus-core/plutus-ir/src/PlutusIR/Core/Type.hs

@@ -27,7 +27,7 @@ module PlutusIR.Core.Type
  , termAnn
  , bindingAnn
  , progAnn
- , progVersion
+ , progVer
  , progTerm
  ) where
 
@@ -194,10 +194,10 @@ instance TermLike (Term tyname name uni fun) tyname name uni fun where
  typeLet x (Def vd bind) = Let x NonRec (pure $ TypeBind x vd bind)
 
 data Program tyname name uni fun ann = Program
- { _progAnn  :: ann
- , _progVersion :: Version
+ { _progAnn :: ann
+ , _progVer  :: Version
  -- ^ The version of the program. This corresponds to the underlying Plutus Core version.
- , _progTerm  :: Term tyname name uni fun ann
+ , _progTerm :: Term tyname name uni fun ann
  }
  deriving stock (Functor, Generic)
 makeLenses ''Program
@@ -237,22 +237,22 @@ applyProgram (Program _a1 v1 _t1) (Program _a2 v2 _t2) =
 
 termAnn :: Term tyname name uni fun a -> a
 termAnn = \case
- Let a _ _ _ -> a
- Var a _ -> a
- TyAbs a _ _ _ -> a
+ Let a _ _ _  -> a
+ Var a _  -> a
+ TyAbs a _ _ _  -> a
  LamAbs a _ _ _ -> a
- Apply a _ _ -> a
- Constant a _ -> a
- Builtin a _ -> a
- TyInst a _ _ -> a
- Error a _ -> a
- IWrap a _ _ _ -> a
- Unwrap a _ -> a
+ Apply a _ _  -> a
+ Constant a _  -> a
+ Builtin a _  -> a
+ TyInst a _ _  -> a
+ Error a _  -> a
+ IWrap a _ _ _  -> a
+ Unwrap a _  -> a
  Constr a _ _ _ -> a
- Case a _ _ _ -> a
+ Case a _ _ _  -> a
 
 bindingAnn :: Binding tyname name uni fun a -> a
 bindingAnn = \case
  TermBind a _ _ _ -> a
- TypeBind a _ _ -> a
+ TypeBind a _ _  -> a
  DatatypeBind a _ -> a

plutus-core/plutus-ir/test/PlutusIR/Parser/Tests.hs

@@ -1,29 +1,29 @@
 -- editorconfig-checker-disable-file
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE PolyKinds #-}
 
 -- | Tests for PIR parser.
 module PlutusIR.Parser.Tests where
 
 import PlutusPrelude
 
-import Data.Char
-import Data.Text qualified as T
-
-import PlutusCore (runQuoteT)
+import PlutusCore qualified as PLC
 import PlutusCore.Annotation
-import PlutusCore.Default qualified as PLC
+import PlutusCore.Default (noMoreTypeFunctions)
 import PlutusCore.Error (ParserErrorBundle)
-import PlutusCore.Test (mapTestLimitAtLeast)
+import PlutusCore.Test (isSerialisable, mapTestLimitAtLeast)
 import PlutusIR
 import PlutusIR.Generators.AST
 import PlutusIR.Parser
 
+import Data.Char
+import Data.Text qualified as T
 import Hedgehog hiding (Var)
 import Hedgehog.Gen qualified as Gen
 import Hedgehog.Range qualified as Range
-
 import Test.Tasty
 import Test.Tasty.Hedgehog
 
@@ -79,23 +79,46 @@ aroundSeparators = go False False
  pure $ a : s1 ++ b : s2 ++ rest
  | otherwise -> (a :) <$> go inQuotedName inUnique splice (b : l)
 
+-- | Check whether the given constant can be scrambled (in the sense of 'genScrambledWith').
+isScramblable :: PLC.Some (PLC.ValueOf PLC.DefaultUni) -> Bool
+isScramblable (PLC.Some (PLC.ValueOf uni0 x0)) = go uni0 x0 where
+ go :: PLC.DefaultUni (PLC.Esc a) -> a -> Bool
+ go PLC.DefaultUniInteger _ = True
+ go PLC.DefaultUniByteString _ = True
+ -- Keep in sync with 'aroundSeparators'.
+ go PLC.DefaultUniString text = T.all (\c -> not (separator c) && c /= '`') text
+ go PLC.DefaultUniUnit _ = True
+ go PLC.DefaultUniBool _ = True
+ go (PLC.DefaultUniProtoList `PLC.DefaultUniApply` uniA) xs =
+ all (go uniA) xs
+ go (PLC.DefaultUniProtoPair `PLC.DefaultUniApply` uniA `PLC.DefaultUniApply` uniB) (x, y) =
+ go uniA x && go uniB y
+ go (f `PLC.DefaultUniApply` _ `PLC.DefaultUniApply` _ `PLC.DefaultUniApply` _) _ =
+ noMoreTypeFunctions f
+ go PLC.DefaultUniData _ = True
+ go PLC.DefaultUniBLS12_381_G1_Element _ = False
+ go PLC.DefaultUniBLS12_381_G2_Element _ = False
+ go PLC.DefaultUniBLS12_381_MlResult _ = False
+
 genScrambledWith :: MonadGen m => m String -> m (String, String)
 genScrambledWith splice = do
- original <- display <$> runAstGen genProgram
+ original <- display <$> runAstGen (discardIfAnyConstant (not . isScramblable) genProgram)
  scrambled <- aroundSeparators splice original
  return (original, scrambled)
 
 propRoundTrip :: Property
 propRoundTrip = property $ do
- code <- display <$> forAllWith display (runAstGen genProgram)
+ code <- display <$>
+ forAllWith display (runAstGen $ discardIfAnyConstant (not . isSerialisable) genProgram)
  let backward = fmap (display . prog)
  forward = fmap PrettyProg . parseProg
  tripping code forward backward
 
 -- | The `SrcSpan` of a parsed `Term` should not including trailing whitespaces.
 propTermSrcSpan :: Property
 propTermSrcSpan = property $ do
- code <- display <$> forAllWith display (runAstGen genTerm)
+ code <- display . _progTerm <$>
+ forAllWith display (runAstGen $ discardIfAnyConstant (not . isSerialisable) genProgram)
  let (endingLine, endingCol) = length &&& T.length . last $ T.lines code
  trailingSpaces <- forAll $ Gen.text (Range.linear 0 10) (Gen.element [' ', '\n'])
  case parseTerm (code <> trailingSpaces) of
@@ -110,15 +133,15 @@ parseProg ::
  ParserErrorBundle
  (Program TyName Name PLC.DefaultUni PLC.DefaultFun SrcSpan)
 parseProg p =
- runQuoteT $ parse program "test" p
+ PLC.runQuoteT $ parse program "test" p
 
 parseTerm ::
  T.Text ->
  Either
  ParserErrorBundle
  (Term TyName Name PLC.DefaultUni PLC.DefaultFun SrcSpan)
 parseTerm p =
- runQuoteT $ parse pTerm "test" p
+ PLC.runQuoteT $ parse pTerm "test" p
 
 propIgnores :: Gen String -> Property
 propIgnores splice = property $ do