Skip to content

Commit

Permalink
feat: add cleanupAnnotations parameter to forallTelescope methods (
Browse files Browse the repository at this point in the history
  • Loading branch information
leodemoura authored May 15, 2024
1 parent c774160 commit 9a8e7a6
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 20 deletions.
51 changes: 32 additions & 19 deletions src/Lean/Meta/Basic.lean
Original file line number Diff line number Diff line change
Expand Up @@ -1087,16 +1087,20 @@ mutual
if `maxFVars?` is `some max`, then we interrupt the telescope construction
when `fvars.size == max`
If `cleanupAnnotations` is `true`, we apply `Expr.cleanupAnnotations` to each type in the telescope.
-/
private partial def forallTelescopeReducingAuxAux
(reducing : Bool) (maxFVars? : Option Nat)
(type : Expr)
(k : Array Expr → Expr → MetaM α) : MetaM α := do
(k : Array Expr → Expr → MetaM α) (cleanupAnnotations : Bool) : MetaM α := do
let rec process (lctx : LocalContext) (fvars : Array Expr) (j : Nat) (type : Expr) : MetaM α := do
match type with
| .forallE n d b bi =>
if fvarsSizeLtMaxFVars fvars maxFVars? then
let d := d.instantiateRevRange j fvars.size fvars
let d := if cleanupAnnotations then d.cleanupAnnotations else d
let fvarId ← mkFreshFVarId
let lctx := lctx.mkLocalDecl fvarId n d bi
let fvar := mkFVar fvarId
Expand All @@ -1121,13 +1125,13 @@ mutual
k fvars type
process (← getLCtx) #[] 0 type

private partial def forallTelescopeReducingAux (type : Expr) (maxFVars? : Option Nat) (k : Array Expr → Expr → MetaM α) : MetaM α := do
private partial def forallTelescopeReducingAux (type : Expr) (maxFVars? : Option Nat) (k : Array Expr → Expr → MetaM α) (cleanupAnnotations : Bool) : MetaM α := do
match maxFVars? with
| some 0 => k #[] type
| _ => do
let newType ← whnf type
if newType.isForall then
forallTelescopeReducingAuxAux true maxFVars? newType k
forallTelescopeReducingAuxAux true maxFVars? newType k cleanupAnnotations
else
k #[] type

Expand All @@ -1151,7 +1155,7 @@ mutual

private partial def isClassExpensive? (type : Expr) : MetaM (Option Name) :=
withReducible do -- when testing whether a type is a type class, we only unfold reducible constants.
forallTelescopeReducingAux type none fun _ type => isClassApp? type
forallTelescopeReducingAux type none (cleanupAnnotations := false) fun _ type => isClassApp? type

private partial def isClassImp? (type : Expr) : MetaM (Option Name) := do
match (← isClassQuick? type) with
Expand Down Expand Up @@ -1180,15 +1184,18 @@ private def withNewLocalInstancesImpAux (fvars : Array Expr) (j : Nat) : n α
partial def withNewLocalInstances (fvars : Array Expr) (j : Nat) : n α → n α :=
mapMetaM <| withNewLocalInstancesImpAux fvars j

@[inline] private def forallTelescopeImp (type : Expr) (k : Array Expr → Expr → MetaM α) : MetaM α := do
forallTelescopeReducingAuxAux (reducing := false) (maxFVars? := none) type k
@[inline] private def forallTelescopeImp (type : Expr) (k : Array Expr → Expr → MetaM α) (cleanupAnnotations : Bool) : MetaM α := do
forallTelescopeReducingAuxAux (reducing := false) (maxFVars? := none) type k cleanupAnnotations

/--
Given `type` of the form `forall xs, A`, execute `k xs A`.
This combinator will declare local declarations, create free variables for them,
execute `k` with updated local context, and make sure the cache is restored after executing `k`. -/
def forallTelescope (type : Expr) (k : Array Expr → Expr → n α) : n α :=
map2MetaM (fun k => forallTelescopeImp type k) k
execute `k` with updated local context, and make sure the cache is restored after executing `k`.
If `cleanupAnnotations` is `true`, we apply `Expr.cleanupAnnotations` to each type in the telescope.
-/
def forallTelescope (type : Expr) (k : Array Expr → Expr → n α) (cleanupAnnotations := false) : n α :=
map2MetaM (fun k => forallTelescopeImp type k cleanupAnnotations) k

/--
Given a monadic function `f` that takes a type and a term of that type and produces a new term,
Expand All @@ -1207,23 +1214,29 @@ and then builds the lambda telescope term for the new term.
def mapForallTelescope (f : Expr → MetaM Expr) (forallTerm : Expr) : MetaM Expr := do
mapForallTelescope' (fun _ e => f e) forallTerm

private def forallTelescopeReducingImp (type : Expr) (k : Array Expr → Expr → MetaM α) : MetaM α :=
forallTelescopeReducingAux type (maxFVars? := none) k
private def forallTelescopeReducingImp (type : Expr) (k : Array Expr → Expr → MetaM α) (cleanupAnnotations : Bool) : MetaM α :=
forallTelescopeReducingAux type (maxFVars? := none) k cleanupAnnotations

/--
Similar to `forallTelescope`, but given `type` of the form `forall xs, A`,
it reduces `A` and continues building the telescope if it is a `forall`. -/
def forallTelescopeReducing (type : Expr) (k : Array Expr → Expr → n α) : n α :=
map2MetaM (fun k => forallTelescopeReducingImp type k) k
it reduces `A` and continues building the telescope if it is a `forall`.
private def forallBoundedTelescopeImp (type : Expr) (maxFVars? : Option Nat) (k : Array Expr → Expr → MetaM α) : MetaM α :=
forallTelescopeReducingAux type maxFVars? k
If `cleanupAnnotations` is `true`, we apply `Expr.cleanupAnnotations` to each type in the telescope.
-/
def forallTelescopeReducing (type : Expr) (k : Array Expr → Expr → n α) (cleanupAnnotations := false) : n α :=
map2MetaM (fun k => forallTelescopeReducingImp type k cleanupAnnotations) k

private def forallBoundedTelescopeImp (type : Expr) (maxFVars? : Option Nat) (k : Array Expr → Expr → MetaM α) (cleanupAnnotations : Bool) : MetaM α :=
forallTelescopeReducingAux type maxFVars? k cleanupAnnotations

/--
Similar to `forallTelescopeReducing`, stops constructing the telescope when
it reaches size `maxFVars`. -/
def forallBoundedTelescope (type : Expr) (maxFVars? : Option Nat) (k : Array Expr → Expr → n α) : n α :=
map2MetaM (fun k => forallBoundedTelescopeImp type maxFVars? k) k
it reaches size `maxFVars`.
If `cleanupAnnotations` is `true`, we apply `Expr.cleanupAnnotations` to each type in the telescope.
-/
def forallBoundedTelescope (type : Expr) (maxFVars? : Option Nat) (k : Array Expr → Expr → n α) (cleanupAnnotations := false) : n α :=
map2MetaM (fun k => forallBoundedTelescopeImp type maxFVars? k cleanupAnnotations) k

private partial def lambdaTelescopeImp (e : Expr) (consumeLet : Bool) (k : Array Expr → Expr → MetaM α) (cleanupAnnotations := false) : MetaM α := do
process consumeLet (← getLCtx) #[] 0 e
Expand Down
2 changes: 1 addition & 1 deletion tests/lean/docStr.lean.expected.out
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ doc string for 'g' is not available
"let rec documentation at g "
"Gadget for optional parameter support.\n\nA binder like `(x : α := default)` in a declaration is syntax sugar for\n`x : optParam α default`, and triggers the elaborator to attempt to use\n`default` to supply the argument if it is not supplied.\n"
"Auxiliary declaration used to implement named patterns like `x@h:p`. "
"Similar to `forallTelescope`, but given `type` of the form `forall xs, A`,\nit reduces `A` and continues building the telescope if it is a `forall`. "
"Similar to `forallTelescope`, but given `type` of the form `forall xs, A`,\nit reduces `A` and continues building the telescope if it is a `forall`.\n\nIf `cleanupAnnotations` is `true`, we apply `Expr.cleanupAnnotations` to each type in the telescope.\n"
Foo :=
{ range := { pos := { line := 4, column := 0 },
charUtf16 := 0,
Expand Down
37 changes: 37 additions & 0 deletions tests/lean/run/cleanup_forallTelescope.lean
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
import Lean

open Lean Meta

def gen (declName : Name) : MetaM Unit := do
let info ← getConstInfo declName
IO.println (← ppExpr info.type)
IO.println "---"
forallTelescope info.type (cleanupAnnotations := true) fun xs type => do
let r ← mkForallFVars xs type
IO.println (← ppExpr r)

/--
info: {coll : Type u} →
{idx : Type v} →
{elem : outParam (Type w)} →
{valid : outParam (coll → idx → Prop)} →
[self : GetElem coll idx elem valid] → (xs : coll) → (i : idx) → valid xs i → elem
---
{coll : Type u} →
{idx : Type v} →
{elem : Type w} →
{valid : coll → idx → Prop} → [self : GetElem coll idx elem valid] → (xs : coll) → (i : idx) → valid xs i → elem
-/
#guard_msgs in
#eval gen ``GetElem.getElem

def f (x := 0) (_ : x = x := by rfl) : Nat := x+1


/--
info: (x : optParam Nat 0) → autoParam (x = x) _auto✝ → Nat
---
(x : Nat) → x = x → Nat
-/
#guard_msgs in
#eval gen ``f

0 comments on commit 9a8e7a6

Please sign in to comment.