Proper generation and shrinking of built-in types

plutus-core/plutus-core/src/PlutusCore/Builtin/KnownKind.hs

@@ -10,9 +10,12 @@ module PlutusCore.Builtin.KnownKind where
 import PlutusCore.Core
 
 import Data.Kind as GHC
+import Data.Proxy
 import GHC.Types
 import Universe
 
+infixr 5 `SingKindArrow`
+
 -- | The type of singletonized Haskell kinds representing Plutus kinds.
 -- Indexing by a Haskell kind allows us to avoid an 'error' call in the 'ToKind' instance of
 -- 'DefaultUni' and not worry about proxies in the type of 'knownKind'.
@@ -31,6 +34,13 @@ instance rep ~ LiftedRep => KnownKind (TYPE rep) where
 instance (KnownKind dom, KnownKind cod) => KnownKind (dom -> cod) where
     knownKind = SingKindArrow (knownKind @dom) (knownKind @cod)
 
+withKnownKind :: Kind ann -> (forall k. KnownKind k => Proxy k -> r) -> r
+withKnownKind (Type _) k = k $ Proxy @GHC.Type
+withKnownKind (KindArrow _ dom cod) k =
+    withKnownKind dom $ \(_ :: Proxy dom) ->
+        withKnownKind cod $ \(_ :: Proxy cod) ->
+            k $ Proxy @(dom -> cod)
+
 -- We need this for type checking Plutus, however we get Plutus types/terms/programs by either
 -- producing them directly or by parsing/decoding them and in both the cases we have access to the
 -- @Typeable@ information for the Haskell kind of a type and so we could just keep it around

plutus-core/testlib/PlutusCore/Generators/PIR/Builtin.hs

@@ -1,4 +1,6 @@
 {-# LANGUAGE DefaultSignatures #-}
+{-# LANGUAGE GADTs             #-}
+{-# LANGUAGE PolyKinds         #-}
 {-# LANGUAGE TypeApplications  #-}
 
 {-# OPTIONS_GHC -fno-warn-orphans #-}
@@ -7,11 +9,15 @@ module PlutusCore.Generators.PIR.Builtin where
 
 import Data.ByteString (ByteString)
 import Data.Coerce
+import Data.Kind qualified as GHC
+import Data.Proxy
 import Data.Text (Text)
 import Data.Text qualified as Text
 import Data.Text.Encoding qualified as Text
 import Test.QuickCheck
 
+import PlutusCore
+import PlutusCore.Builtin
 import PlutusCore.Data
 
 instance Arbitrary Data where
@@ -62,3 +68,77 @@ instance ArbitraryBuiltin a => ArbitraryBuiltin [a] where
 instance (ArbitraryBuiltin a, ArbitraryBuiltin b) => ArbitraryBuiltin (a, b) where
     arbitraryBuiltin = coerce $ arbitrary @(AsArbitraryBuiltin a, AsArbitraryBuiltin b)
     shrinkBuiltin = coerce $ shrink @(AsArbitraryBuiltin a, AsArbitraryBuiltin b)
+
+-- | Either a fail to generate anything or a built-in type of a given kind.
+data MaybeSomeTypeOf k
+    = NothingSomeType
+    | forall (a :: k). JustSomeType (DefaultUni (Esc a))
+
+-- | Generate a 'DefaultUniApply' if possible.
+genDefaultUniApply :: KnownKind k => Gen (MaybeSomeTypeOf k)
+genDefaultUniApply = do
+    mayFun <- scale (`div` 2) arbitrary
+    mayArg <- scale (`div` 2) arbitrary :: Gen (MaybeSomeTypeOf GHC.Type)
+    pure $ case (mayFun, mayArg) of
+        (JustSomeType fun, JustSomeType arg) -> JustSomeType $ fun `DefaultUniApply` arg
+        _                                    -> NothingSomeType
+
+-- | Shrink a 'DefaultUniApply' to one of the elements of the spine and throw away the head
+-- (because the head of an application can't be of the same kind as the whole application).
+-- We don't have higher-kinded built-in types, so we don't do this kind of shrinking for any kinds
+-- other than *.
+shrinkDefaultUniApplyStar :: DefaultUni (Esc b) -> [MaybeSomeTypeOf GHC.Type]
+shrinkDefaultUniApplyStar (fun `DefaultUniApply` arg) =
+    [JustSomeType arg | SingType <- [toSingKind arg]] ++ shrinkDefaultUniApplyStar fun
+shrinkDefaultUniApplyStar _ = []
+
+instance KnownKind k => Arbitrary (MaybeSomeTypeOf k) where
+   arbitrary = do
+       size <- getSize
+       oneof $ case knownKind @k of
+           SingType ->
+               [genDefaultUniApply | size > 10] ++
+               [ pure $ JustSomeType DefaultUniInteger
+               , pure $ JustSomeType DefaultUniByteString
+               , pure $ JustSomeType DefaultUniString
+               , pure $ JustSomeType DefaultUniUnit
+               , pure $ JustSomeType DefaultUniBool
+               -- Commented out, because the 'Arbitrary' instance for 'Data' isn't implemented
+               -- (errors out).
+               -- , pure $ JustSomeType DefaultUniData
+               ]
+           SingType `SingKindArrow` SingType ->
+               [genDefaultUniApply | size > 10] ++
+               [pure $ JustSomeType DefaultUniProtoList]
+           SingType `SingKindArrow` SingType `SingKindArrow` SingType ->
+               -- No 'genDefaultUniApply', because we don't have any built-in type constructors
+               -- taking three or more arguments.
+               [pure $ JustSomeType DefaultUniProtoPair]
+           _ -> [pure NothingSomeType]
+
+   -- No shrinks if you don't have anything to shrink.
+   shrink NothingSomeType    = []
+   shrink (JustSomeType uni) =
+       case knownKind @k of
+           SingType -> case uni of
+               -- 'DefaultUniUnit' is the "minimal" built-in type, can't shrink it any further.
+               DefaultUniUnit -> []
+               -- Any other built-in type of kind * shrinks to 'DefaultUniUnit' and if it happens to
+               -- be a built-in type application, then also all suitable arguments of the
+               -- application.
+               _              -> JustSomeType DefaultUniUnit : shrinkDefaultUniApplyStar uni
+           -- No suitable builtins to shrink to for non-* kinds.
+           _ -> []
+
+-- | Convert a @MaybeSomeTypeOf kind@ to a @Maybe (SomeTypeIn DefaultUni)@ and forget the @kind@
+-- (which was only used for generating well-kinded built-in type applications, which we have to do
+-- due to having intrinsic kinding of builtins).
+fromMaybeSomeTypeOf :: MaybeSomeTypeOf kind -> Maybe (SomeTypeIn DefaultUni)
+fromMaybeSomeTypeOf NothingSomeType    = Nothing
+fromMaybeSomeTypeOf (JustSomeType uni) = Just $ SomeTypeIn uni
+
+-- | Generate a built-in type of a given kind.
+genBuiltinTypeOf :: Kind () -> Gen (Maybe (SomeTypeIn DefaultUni))
+genBuiltinTypeOf kind =
+    withKnownKind kind $ \(_ :: Proxy kind) ->
+        fromMaybeSomeTypeOf <$> arbitrary @(MaybeSomeTypeOf kind)

plutus-core/testlib/PlutusCore/Generators/PIR/GenTm.hs

@@ -19,6 +19,7 @@ import Control.Monad.Reader
 import Data.List.NonEmpty (NonEmpty (..))
 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)
@@ -112,6 +113,12 @@ runGenTmCustom f cg g = do
           }
     flip runReader env <$> runGenT g
 
+suchThatMap :: Monad m => GenT m a -> (a -> Maybe b) -> GenT m b
+gen `suchThatMap` f = fmap fromJust $ fmap f gen `suchThat` isJust
+
+deliver :: Monad m => GenT m (Maybe a) -> GenT m a
+deliver gen = gen `suchThatMap` id
+
 -- * Utility functions
 
 -- | Don't allow types to escape from a generator.

plutus-core/testlib/PlutusCore/Generators/PIR/GenerateTerms.hs

@@ -313,8 +313,8 @@ genTerm mty = checkInvariants $ do
     canConst (Just _)           = False
 
     genConst Nothing = do
-      b <- liftGen $ elements $ builtinTysAt $ Type ()
-      (TyBuiltin () b,) <$> genConstant b
+      b <- deliver . liftGen . genBuiltinTypeOf $ Type ()
+      (TyBuiltin () b, ) <$> genConstant b
     genConst (Just ty@(TyBuiltin _ b)) = (ty,) <$> genConstant b
     genConst _ = error "genConst: impossible"
 

plutus-core/testlib/PlutusCore/Generators/PIR/GenerateTypes.hs

@@ -6,6 +6,7 @@ module PlutusCore.Generators.PIR.GenerateTypes where
 import Control.Monad.Reader
 
 import Data.Map.Strict qualified as Map
+import Data.Maybe
 import Data.String
 import Test.QuickCheck (shuffle)
 import Test.QuickCheck.GenT
@@ -20,6 +21,7 @@ import PlutusCore.Quote (runQuote)
 import PlutusIR
 import PlutusIR.Core.Instance.Pretty.Readable
 
+import PlutusCore.Generators.PIR.Builtin
 import PlutusCore.Generators.PIR.GenTm
 import PlutusCore.Generators.PIR.GenerateKinds ()
 
@@ -44,27 +46,6 @@ import PlutusCore.Generators.PIR.GenerateKinds ()
 
 -- * Generators for well-kinded types
 
--- TODO: make this a proper generator running in 'GenTm'.
--- | Get the types of builtins at a given kind
-builtinTysAt :: Kind () -> [SomeTypeIn DefaultUni]
-builtinTysAt (Type _) =
-  -- TODO: add 'DefaultUniData' once it has a non-throwing 'ArbitraryBuiltin' instance.
-  [ SomeTypeIn DefaultUniInteger
-  , SomeTypeIn DefaultUniUnit
-  , SomeTypeIn DefaultUniBool
-  , SomeTypeIn DefaultUniByteString
-  , SomeTypeIn DefaultUniString
-  -- TODO: should we have more examples of lists and pairs here?
-  , SomeTypeIn $ DefaultUniList DefaultUniBool
-  , SomeTypeIn $ DefaultUniPair DefaultUniInteger DefaultUniUnit]
-builtinTysAt (KindArrow _ (Type _) (Type _)) =
-  [ SomeTypeIn DefaultUniProtoList
-  , SomeTypeIn $ DefaultUniProtoPair `DefaultUniApply` DefaultUniString
-  ]
-builtinTysAt (KindArrow _ (Type _) (KindArrow () (Type _) (Type _))) =
-  [SomeTypeIn DefaultUniProtoPair]
-builtinTysAt _ = []
-
 -- | Generate "small" types at a given kind such as builtins, bound variables, bound datatypes,
 -- and lambda abstractions \ t0 ... tn. T
 genAtomicType :: Kind () -> GenTm (Type TyName DefaultUni ())
@@ -73,15 +54,15 @@ genAtomicType k = do
   dts <- asks geDatas
   let atoms = [ TyVar () x | (x, k') <- Map.toList tys, k == k' ] ++
               [ TyVar () x | (x, Datatype _ (TyVarDecl _ _ k') _ _ _) <- Map.toList dts, k == k' ]
-      builtins = map (TyBuiltin ()) $ builtinTysAt k
+      genBuiltin = fmap (TyBuiltin ()) <$> genBuiltinTypeOf k
       lam k1 k2 = do
         x <- genMaybeFreshTyName "a"
         TyLam () x k1 <$> bindTyName x k1 (genAtomicType k2)
-  -- TODO: probably should be 'frequencyTm'?
-  oneof $ concat
-    [ map pure atoms
-    , [elements builtins | not $ null builtins]
-    , [lam k1 k2 | KindArrow _ k1 k2 <- [k]]
+  -- There's always an atomic type of a given type, hence the usage of 'deliverOneof'.
+  deliver $ frequency
+    [ (7, if null atoms then pure Nothing else Just <$> elements atoms)
+    , (1, liftGen genBuiltin)
+    , (3, sequence $ listToMaybe [lam k1 k2 | KindArrow _ k1 k2 <- [k]])
     ]
 
 -- | Generate a type at a given kind