Skip to content

Commit 98b609e

Browse files
ngernestngernest
authored
Merge pull request leanprover-community#6 from ngernest/instance_deriving
Deriving Handler Frontend for `Arbitrary` Typeclass
2 parents 5267897 + 97050fe commit 98b609e

8 files changed

Lines changed: 303 additions & 124 deletions

‎Plausible/New/DeriveArbitrary.lean‎

Lines changed: 134 additions & 108 deletions
Original file line numberDiff line numberDiff line change
@@ -55,10 +55,123 @@ def getCtorArgsNamesAndTypes (ctorName : Name) : MetaM (Array (Name × Expr)) :=
5555

5656
return argNamesAndTypes
5757

58+
/-- Creates an instance of the `ArbitrarySized` typeclass for an inductive type
59+
whose name is given by `targetTypeName`.
60+
61+
(Note: the main logic for determining the structure of the derived generator
62+
is contained in this function.) -/
63+
def mkArbitrarySizedInstance (targetTypeName : Name) : CommandElabM (TSyntax `command) := do
64+
-- Obtain Lean's `InductiveVal` data structure, which contains metadata about
65+
-- the type corresponding to `targetTypeName`
66+
let inductiveVal ← getConstInfoInduct targetTypeName
67+
68+
-- Fetch the ambient local context, which we need to produce user-accessible fresh names
69+
let localCtx ← liftTermElabM $ getLCtx
70+
71+
-- Produce a fresh name for the `size` argument for the lambda
72+
-- at the end of the generator function, as well as the `aux_arb` inner helper function
73+
let freshSizeIdent := mkFreshAccessibleIdent localCtx `size
74+
let freshSize' := mkFreshAccessibleIdent localCtx `size'
75+
let auxArbIdent := mkFreshAccessibleIdent localCtx `aux_arb
76+
77+
let mut nonRecursiveGenerators := #[]
78+
let mut recursiveGenerators := #[]
79+
for ctorName in inductiveVal.ctors do
80+
let ctorIdent := mkIdent ctorName
81+
82+
let ctorArgNamesTypes ← liftTermElabM $ getCtorArgsNamesAndTypes ctorName
83+
84+
if ctorArgNamesTypes.isEmpty then
85+
-- Constructor is nullary, we can just use a generator of the form `pure ...`
86+
let pureGen ← `($pureFn $ctorIdent)
87+
nonRecursiveGenerators := nonRecursiveGenerators.push pureGen
88+
else
89+
-- Produce a fresh name for each of the args to the constructor
90+
let ctorArgNames := Prod.fst <$> ctorArgNamesTypes
91+
let freshArgIdents := Lean.mkIdent <$> genFreshNames (existingNames := ctorArgNames) (namePrefixes := ctorArgNames)
92+
93+
let mut doElems := #[]
94+
95+
-- Determine whether the constructor has any recursive arguments
96+
let ctorIsRecursive ← liftTermElabM $ isConstructorRecursive targetTypeName ctorName
97+
if !ctorIsRecursive then
98+
-- Call `arbitrary` to generate a random value for each of the arguments
99+
for freshIdent in freshArgIdents do
100+
let bindExpr ← liftTermElabM $ mkLetBind freshIdent #[arbitraryFn]
101+
doElems := doElems.push bindExpr
102+
else
103+
-- For recursive constructors, we need to examine each argument to see which of them require
104+
-- recursive calls to the generator
105+
let freshArgIdentsTypes := Array.zip freshArgIdents (Prod.snd <$> ctorArgNamesTypes)
106+
for (freshIdent, argType) in freshArgIdentsTypes do
107+
-- If the argument's type is the same as the target type,
108+
-- produce a recursive call to the generator using `aux_arb`,
109+
-- otherwise generate a value using `arbitrary`
110+
let bindExpr ←
111+
liftTermElabM $
112+
if argType.getAppFn.constName == targetTypeName then
113+
mkLetBind freshIdent #[auxArbFn, freshSize']
114+
else
115+
mkLetBind freshIdent #[arbitraryFn]
116+
doElems := doElems.push bindExpr
117+
118+
-- Create an expression `return C x1 ... xn` at the end of the generator, where
119+
-- `C` is the constructor name and the `xi` are the generated values for the args
120+
let pureExpr ← `(doElem| return $ctorIdent $freshArgIdents*)
121+
doElems := doElems.push pureExpr
122+
123+
-- Put the body of the generator together
124+
let generatorBody ← liftTermElabM $ mkDoBlock doElems
125+
if !ctorIsRecursive then
126+
nonRecursiveGenerators := nonRecursiveGenerators.push generatorBody
127+
else
128+
recursiveGenerators := recursiveGenerators.push generatorBody
129+
130+
-- Just use the first non-recursive generator as the default generator
131+
let defaultGenerator := nonRecursiveGenerators[0]!
132+
133+
-- Turn each generator into a thunked function and associate each generator with its weight
134+
-- (1 for non-recursive generators, `.succ size'` for recursive generators)
135+
let thunkedNonRecursiveGenerators ←
136+
Array.mapM (fun generatorBody => `($generatorCombinatorsThunkGenFn (fun _ => $generatorBody))) nonRecursiveGenerators
137+
138+
let mut weightedThunkedNonRecursiveGens := #[]
139+
for thunkedGen in thunkedNonRecursiveGenerators do
140+
let thunkedGen ← `((1, $thunkedGen))
141+
weightedThunkedNonRecursiveGens := weightedThunkedNonRecursiveGens.push thunkedGen
142+
143+
let mut weightedThunkedRecursiveGens := #[]
144+
for recursiveGen in recursiveGenerators do
145+
let thunkedWeightedGen ← `(($succIdent $freshSize', $generatorCombinatorsThunkGenFn (fun _ => $recursiveGen)))
146+
weightedThunkedRecursiveGens := weightedThunkedRecursiveGens.push thunkedWeightedGen
147+
148+
-- Create the cases for the pattern-match on the size argument
149+
-- If `size = 0`, pick one of the thunked non-recursive generators
150+
let mut caseExprs := #[]
151+
let zeroCase ← `(Term.matchAltExpr| | $zeroIdent => $oneOfWithDefaultFn $defaultGenerator [$thunkedNonRecursiveGenerators,*])
152+
caseExprs := caseExprs.push zeroCase
153+
154+
-- If `size = .succ size'`, pick a generator (it can be non-recursive or recursive)
155+
let mut allThunkedWeightedGenerators ← `([$weightedThunkedNonRecursiveGens,*, $weightedThunkedRecursiveGens,*])
156+
let succCase ← `(Term.matchAltExpr| | $succIdent $freshSize' => $frequencyFn $defaultGenerator $allThunkedWeightedGenerators)
157+
caseExprs := caseExprs.push succCase
158+
159+
-- Create function argument for the generator size
160+
let sizeParam ← `(Term.letIdBinder| ($sizeIdent : $natIdent))
161+
let matchExpr ← liftTermElabM $ mkMatchExpr sizeIdent caseExprs
162+
163+
-- Create an instance of the `ArbitrarySized` typeclass
164+
let targetTypeIdent := mkIdent targetTypeName
165+
let generatorType ← `($genIdent $targetTypeIdent)
166+
`(instance : $arbitrarySizedTypeclass $targetTypeIdent where
167+
$unqualifiedArbitrarySizedFn:ident :=
168+
let rec $auxArbIdent:ident $sizeParam : $generatorType :=
169+
$matchExpr
170+
fun $freshSizeIdent => $auxArbIdent $freshSizeIdent)
58171

59172
syntax (name := derive_arbitrary) "#derive_arbitrary" term : command
60173

61-
/-- Derives an instance of the `ArbitrarySized` typeclass -/
174+
/-- Command elaborator which derives an instance of the `ArbitrarySized` typeclass -/
62175
@[command_elab derive_arbitrary]
63176
def elabDeriveArbitrary : CommandElab := fun stx => do
64177
match stx with
@@ -73,121 +186,34 @@ def elabDeriveArbitrary : CommandElab := fun stx => do
73186

74187
let isInductiveType ← isInductive targetTypeName
75188
if isInductiveType then
76-
let inductiveVal ← getConstInfoInduct targetTypeName
77-
78-
-- Fetch the ambient local context, which we need to produce user-accessible fresh names
79-
let localCtx ← liftTermElabM $ getLCtx
80-
81-
-- Produce a fresh name for the `size` argument for the lambda
82-
-- at the end of the generator function, as well as the `aux_arb` inner helper function
83-
let freshSizeIdent := mkFreshAccessibleIdent localCtx `size
84-
let freshSize' := mkFreshAccessibleIdent localCtx `size'
85-
let auxArbIdent := mkFreshAccessibleIdent localCtx `aux_arb
86-
87-
let mut nonRecursiveGenerators := #[]
88-
let mut recursiveGenerators := #[]
89-
for ctorName in inductiveVal.ctors do
90-
let ctorIdent := mkIdent ctorName
91-
92-
let ctorArgNamesTypes ← liftTermElabM $ getCtorArgsNamesAndTypes ctorName
93-
94-
if ctorArgNamesTypes.isEmpty then
95-
-- Constructor is nullary, we can just use a generator of the form `pure ...`
96-
let pureGen ← `($pureFn $ctorIdent)
97-
nonRecursiveGenerators := nonRecursiveGenerators.push pureGen
98-
else
99-
-- Produce a fresh name for each of the args to the constructor
100-
let ctorArgNames := Prod.fst <$> ctorArgNamesTypes
101-
let freshArgIdents := Lean.mkIdent <$> genFreshNames (existingNames := ctorArgNames) (namePrefixes := ctorArgNames)
102-
103-
let mut doElems := #[]
104-
105-
-- Determine whether the constructor has any recursive arguments
106-
let ctorIsRecursive ← liftTermElabM $ isConstructorRecursive targetTypeName ctorName
107-
if !ctorIsRecursive then
108-
-- Call `arbitrary` to generate a random value for each of the arguments
109-
for freshIdent in freshArgIdents do
110-
let bindExpr ← liftTermElabM $ mkLetBind freshIdent #[arbitraryFn]
111-
doElems := doElems.push bindExpr
112-
else
113-
-- For recursive constructors, we need to examine each argument to see which of them require
114-
-- recursive calls to the generator
115-
let freshArgIdentsTypes := Array.zip freshArgIdents (Prod.snd <$> ctorArgNamesTypes)
116-
for (freshIdent, argType) in freshArgIdentsTypes do
117-
-- If the argument's type is the same as the target type,
118-
-- produce a recursive call to the generator using `aux_arb`,
119-
-- otherwise generate a value using `arbitrary`
120-
let bindExpr ←
121-
liftTermElabM $
122-
if argType.getAppFn.constName == targetTypeName then
123-
mkLetBind freshIdent #[auxArbFn, freshSize']
124-
else
125-
mkLetBind freshIdent #[arbitraryFn]
126-
doElems := doElems.push bindExpr
127-
128-
-- Create an expression `return C x1 ... xn` at the end of the generator, where
129-
-- `C` is the constructor name and the `xi` are the generated values for the args
130-
let pureExpr ← `(doElem| return $ctorIdent $freshArgIdents*)
131-
doElems := doElems.push pureExpr
132-
133-
-- Put the body of the generator together
134-
let generatorBody ← liftTermElabM $ mkDoBlock doElems
135-
if !ctorIsRecursive then
136-
nonRecursiveGenerators := nonRecursiveGenerators.push generatorBody
137-
else
138-
recursiveGenerators := recursiveGenerators.push generatorBody
139-
140-
-- Just use the first non-recursive generator as the default generator
141-
let defaultGenerator := nonRecursiveGenerators[0]!
142-
143-
-- Turn each generator into a thunked function and associate each generator with its weight
144-
-- (1 for non-recursive generators, `.succ size'` for recursive generators)
145-
let thunkedNonRecursiveGenerators ←
146-
Array.mapM (fun generatorBody => `($generatorCombinatorsThunkGenFn (fun _ => $generatorBody))) nonRecursiveGenerators
147-
148-
let mut weightedThunkedNonRecursiveGens := #[]
149-
for thunkedGen in thunkedNonRecursiveGenerators do
150-
let thunkedGen ← `((1, $thunkedGen))
151-
weightedThunkedNonRecursiveGens := weightedThunkedNonRecursiveGens.push thunkedGen
152-
153-
let mut weightedThunkedRecursiveGens := #[]
154-
for recursiveGen in recursiveGenerators do
155-
let thunkedWeightedGen ← `(($succIdent $freshSize', $generatorCombinatorsThunkGenFn (fun _ => $recursiveGen)))
156-
weightedThunkedRecursiveGens := weightedThunkedRecursiveGens.push thunkedWeightedGen
157-
158-
-- Create the cases for the pattern-match on the size argument
159-
-- If `size = 0`, pick one of the thunked non-recursive generators
160-
let mut caseExprs := #[]
161-
let zeroCase ← `(Term.matchAltExpr| | $zeroIdent => $oneOfWithDefaultFn $defaultGenerator [$thunkedNonRecursiveGenerators,*])
162-
caseExprs := caseExprs.push zeroCase
163-
164-
-- If `size = .succ size'`, pick a generator (it can be non-recursive or recursive)
165-
let mut allThunkedWeightedGenerators ← `([$weightedThunkedNonRecursiveGens,*, $weightedThunkedRecursiveGens,*])
166-
let succCase ← `(Term.matchAltExpr| | $succIdent $freshSize' => $frequencyFn $defaultGenerator $allThunkedWeightedGenerators)
167-
caseExprs := caseExprs.push succCase
168-
169-
-- Create function argument for the generator size
170-
let sizeParam ← `(Term.letIdBinder| ($sizeIdent : $natIdent))
171-
let matchExpr ← liftTermElabM $ mkMatchExpr sizeIdent caseExprs
172-
173-
-- Create an instance of the `ArbitrarySized` typeclass
174-
let generatorType ← `($genIdent $targetTypeIdent)
175-
let typeclassInstance ←
176-
`(instance : $arbitrarySizedTypeclass $targetTypeIdent where
177-
$unqualifiedArbitrarySizedFn:ident :=
178-
let rec $auxArbIdent:ident $sizeParam : $generatorType :=
179-
$matchExpr
180-
fun $freshSizeIdent => $auxArbIdent $freshSizeIdent)
189+
let typeClassInstance ← mkArbitrarySizedInstance targetTypeName
181190

182191
-- Pretty-print the derived generator
183-
let genFormat ← liftCoreM (PrettyPrinter.ppCommand typeclassInstance)
192+
let genFormat ← liftCoreM (PrettyPrinter.ppCommand typeClassInstance)
184193

185194
-- Display the code for the derived typeclass instance to the user
186195
-- & prompt the user to accept it in the VS Code side panel
187196
liftTermElabM $ Tactic.TryThis.addSuggestion stx
188197
(Format.pretty genFormat) (header := "Try this generator: ")
189198

190199
-- Elaborate the typeclass instance and add it to the local context
191-
elabCommand typeclassInstance
200+
elabCommand typeClassInstance
201+
else
202+
throwError "Cannot derive Arbitrary instance for non-inductive types"
192203

193204
| _ => throwUnsupportedSyntax
205+
206+
/-- Deriving handler which produces an instance of the `ArbitrarySized` typeclass for
207+
each type specified in `declNames` -/
208+
def deriveArbitraryInstanceHandler (declNames : Array Name) : CommandElabM Bool := do
209+
if (← declNames.allM isInductive) then
210+
for targetTypeName in declNames do
211+
let typeClassInstance ← mkArbitrarySizedInstance targetTypeName
212+
elabCommand typeClassInstance
213+
return true
214+
else
215+
throwError "Cannot derive instance of Arbitrary typeclass for non-inductive types"
216+
return false
217+
218+
initialize
219+
registerDerivingHandler ``Arbitrary deriveArbitraryInstanceHandler

‎Plausible/New/Tests.lean‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,15 +12,22 @@ open Plausible ArbitrarySizedSuchThat OptionTGen
1212

1313
/-- Dummy inductive relation for testing purposes -/
1414
inductive RGB where
15-
| Red
16-
| Green
17-
| Blue
15+
| Red
16+
| Green
17+
| Blue
18+
deriving Arbitrary
1819

1920
inductive Value where
2021
| none
2122
| bool (b : Bool)
2223
| int (i : Int)
2324
| tensor (shape : List Nat) (dtype : String)
25+
deriving Arbitrary
26+
27+
inductive Foo where
28+
| FromBitVec : ∀ (n : Nat), BitVec n → String → Foo
29+
deriving Arbitrary
30+
2431

2532
inductive MyList where
2633
| Nil
@@ -30,6 +37,12 @@ inductive MyListAnon where
3037
| Nil : MyListAnon
3138
| Cons : Nat -> MyListAnon -> MyListAnon
3239

40+
-- deriving instance Arbitrary for MyList, MyListAnon
41+
42+
-- #synth Arbitrary MyList
43+
-- #synth Arbitrary MyListAnon
44+
45+
3346
-- #derive_arbitrary MyListAnon
3447

3548
-- #derive_arbitrary Tree

‎README.md‎

Lines changed: 54 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,24 +3,74 @@ A property testing framework for Lean 4 that integrates into the tactic framewor
33

44
## New Metaprogramming Code
55
See the [`New`](./Plausible/New/) subdirectory for code that uses Lean's metaprogramming facilities (`TSyntax`)
6-
to automatically derive generators/checkers for inductive relations, à la [Paraskevopoulou et al. 2022](https://lemonidas.github.io/pdf/ComputingCorrectly.pdf).
6+
to automatically derive generators/checkers for inductive relations.
7+
8+
Our design is heavily inspired by [Coq/Rocq's QuickChick](https://github.com/QuickChick/QuickChick) library and the following papers:
9+
- [Computing Correctly with Inductive Relation (PLDI 2022)](https://lemonidas.github.io/pdf/ComputingCorrectly.pdf)
10+
- [Generating Good Generators for Inductive Relations (POPL 2018)](https://lemonidas.github.io/pdf/GeneratingGoodGenerators.pdf)
11+
12+
Like QuickChick & [Haskell QuickChick](https://hackage.haskell.org/package/QuickCheck), we provide the following typeclasses for random generation:
13+
- `Arbitrary`: random generators for inhabitants of algebraic data types
14+
- `ArbitrarySuchThat`: generators which only produce random values that satisfy a user-supplied inductive relation
15+
- `ArbitrarySized`, `ArbitrarySizedSuchThat`: versions of the two typeclasses above where the generator's size parameter is made explicit
16+
17+
We provide two top-level commands which automatically derive generators for Lean `inductive`s:
18+
19+
**1. Deriving unconstrained generators**
20+
An *unconstrained* generator produces random inhabitants of an algebraic data type.
21+
We provide two frontends which derive instances of `Arbitrary` & `ArbitrarySuchThat` respectively:
22+
23+
**1a. Deriving Instance** (for algebraic data types)
24+
Users can write `deriving Arbitrary` after an inductive type definition:
25+
26+
```lean
27+
inductive Tree where
28+
...
29+
deriving Arbitrary
30+
```
31+
32+
Alternatively, users can also write `deriving instance Arbitrary for T1, ..., Tn` as a top-level command
33+
to derive `Arbitrary` instances for types `T1, ..., Tn` simultaneously.
34+
35+
**1b. Command Elaborator**
36+
We provide a command elaborator which elaborates the `#derive_arbitrary` command:
737

8-
We provide two commands which automatically derive generators for Lean inductives:
938
```lean
10-
-- `#derive_arbitrary` derives an unconstrained generator for a `Tree` algebraic data type
39+
-- `#derive_arbitrary` derives an instance of `Arbitrary` for the `Tree` datatype
1140
#derive_arbitrary Tree
41+
```
1242

43+
Regardless of which frontend is used, to sample from the derived generator, users can simply call `runArbitrary` and specify some
44+
`Nat` to act as the generator's size parameter (`10` in the example below):
45+
46+
```lean
47+
#eval runArbitrary (α := Tree) 10
48+
```
49+
50+
**2. Deriving constrained generators** (for inductive relations)
51+
A *constrained* generator only produces random values that satisfy a user-specified inductive relation.
52+
We provide a command elaborator which elaborates the `#derive_generator` command:
53+
54+
```lean
1355
-- `#derive_generator` derives a constrained generator for `Tree`s that are balanced at some height `n`,
1456
-- where `balanced n t` is a user-defined inductive relation
1557
#derive_generator (fun (t : Tree) => balanced n t)
58+
``
59+
60+
To sample from the derived generator, users invoke `runSizedGen` & specify the right
61+
instance of the `ArbitrarySizedSuchThat` typeclass (along with some `Nat` to act as the generator size):
62+
63+
```lean
64+
#eval runSizedGen (ArbitrarySizedSuchThat.arbitrarySizedST (fun t => balanced 5 t)) 10
1665
```
1766

67+
1868
**Repo overview**:
1969

2070
- [`OptionTGen.lean`](./Plausible/New/OptionTGen.lean): Generator combinators that work over the `OptionT Gen` monad transformer (representing generators that may fail)
2171
- [`DecOpt.lean`](./Plausible/New/DecOpt.lean): The `DecOpt` typeclass for partially decidable propositions, adapted from QuickChick
2272
- [`Arbitrary.lean`](./Plausible/New/Arbitrary.lean): The `Arbitrary` & `ArbitrarySized` typeclasses for unconstrained generators, adapted from QuickChick
23-
- [`ArbitrarySizedSuchThat.lean`](./Plausible/New/ArbitrarySizedSuchThat.lean): The `ArbitrarySuchThat` & `ArbitrarySizedSuchThat` typeclasses for constrained generators (generators of values satisfying a proposition), adapted from QuickChick
73+
- [`ArbitrarySizedSuchThat.lean`](./Plausible/New/ArbitrarySizedSuchThat.lean): The `ArbitrarySuchThat` & `ArbitrarySizedSuchThat` typeclasses for constrained generators, adapted from QuickChick
2474
- [`GeneratorCombinators.lean`](./Plausible/New/GeneratorCombinators.lean): Extra combinators for Plausible generators (e.g. analogs of the `sized` and `frequency` combinators from Haskell QuickCheck)
2575
- [`DeriveArbitrary.lean`](./Plausible/New/DeriveArbitrary.lean): Metaprogramming infrastructure for deriving *unconstrained* generators (instances of the `ArbitrarySized` typeclass)
2676
- [`DeriveGenerator.lean`](./Plausible/New/DeriveGenerator.lean): Metaprogramming infrastructure for deriving *constrained* generators (instances of the `ArbitrarySizedSuchThat` typeclass)

0 commit comments

Comments
 (0)