diff --git a/src/Lean/Meta/InferType.lean b/src/Lean/Meta/InferType.lean index 3f8c4a081c..d5d1c14d63 100644 --- a/src/Lean/Meta/InferType.lean +++ b/src/Lean/Meta/InferType.lean @@ -19,6 +19,9 @@ Auxiliary function for instantiating the loose bound variables in `e` with `args This function is similar to `instantiateRevRange`, but it applies beta-reduction when we instantiate a bound variable with a lambda expression. +If `args` contains no lambda expressions, it is equivalent to `instantiateRevRange`, and in fact +it will call `instantiateRevRange` for efficiency. + Example: Given the term `#0 a`, and `start := 0, stop := 1, args := #[fun x => x]` the result is `a` instead of `(fun x => x) a`. This reduction is useful when we are inferring the type of eliminator-like applications. @@ -32,10 +35,38 @@ We use this to implement `inferAppType`. partial def Expr.instantiateBetaRevRange (e : Expr) (start : Nat) (stop : Nat) (args : Array Expr) : Expr := if e.hasLooseBVars && stop > start then assert! stop ≤ args.size - visit e 0 |>.run + if args.any (·.consumeMData.isLambda) start stop then + visit e 0 |>.run + else + -- If there are no lambdas, then `instantiateRevRange` suffices. + instantiateRevRange e start stop args else e where + /-- + Visit a bvar `e := .bvar vidx`, assuming `offset < e.looseBVarRange`. + -/ + visitBVar (vidx : Nat) (offset : Nat) : Expr := + -- Recall that `looseBVarRange` for `Expr.bvar` is `vidx+1`. + -- So, we must have `offset ≤ vidx`, since `offset < e.looseBVarRange` + let n := stop - start + if vidx < offset + n then + args[stop - (vidx - offset) - 1]!.liftLooseBVars 0 offset + else + Expr.bvar (vidx - n) + visitWithoutBeta (e : Expr) (offset : Nat) : MonadStateCacheT (ExprStructEq × Nat) Expr Id Expr := do + if offset >= e.looseBVarRange then + -- `e` doesn't have free variables + return e + else + match e with + | .app f a => + -- Check the cache only here, since in the other alternative `visit` will check the cache. + checkCache ({ val := e : ExprStructEq }, offset) fun _ => visitApp e f a offset + | e => visit e offset + /-- Visit an application without beta reducing the head -/ + visitApp (e f a : Expr) (offset : Nat) : MonadStateCacheT (ExprStructEq × Nat) Expr Id Expr := + return e.updateApp! (← visitWithoutBeta f offset) (← visit a offset) visit (e : Expr) (offset : Nat) : MonadStateCacheT (ExprStructEq × Nat) Expr Id Expr := if offset >= e.looseBVarRange then -- `e` doesn't have free variables @@ -47,23 +78,17 @@ where | .letE _ t v b _ => return e.updateLetE! (← visit t offset) (← visit v offset) (← visit b (offset+1)) | .mdata _ b => return e.updateMData! (← visit b offset) | .proj _ _ b => return e.updateProj! (← visit b offset) - | .app .. => - e.withAppRev fun f revArgs => do - let fNew ← visit f offset - let revArgs ← revArgs.mapM (visit · offset) - if f.isBVar then - -- try to beta reduce if `f` was a bound variable - return fNew.betaRev revArgs + | .bvar vidx => return visitBVar vidx offset + | .app f a => + let head := e.getAppFn + -- try to beta reduce if the head is a bound variable + if head.isBVar then + -- using `visit` instead of `visitBVar` for the `offset >= vidx` check and for caching `liftLooseBVars` + let head ← visit head offset + let revArgs ← e.getAppRevArgs.mapM (visit · offset) + return head.betaRev revArgs else - return mkAppRev fNew revArgs - | Expr.bvar vidx => - -- Recall that looseBVarRange for `Expr.bvar` is `vidx+1`. - -- So, we must have offset ≤ vidx, since we are in the "else" branch of `if offset >= e.looseBVarRange` - let n := stop - start - if vidx < offset + n then - return args[stop - (vidx - offset) - 1]!.liftLooseBVars 0 offset - else - return mkBVar (vidx - n) + visitApp e f a offset -- The following cases are unreachable because they never contain loose bound variables | .const .. => unreachable! | .fvar .. => unreachable! @@ -79,8 +104,6 @@ def throwFunctionExpected {α} (f : Expr) : MetaM α := private def inferAppType (f : Expr) (args : Array Expr) : MetaM Expr := do let mut fType ← inferType f let mut j := 0 - /- TODO: check whether `instantiateBetaRevRange` is too expensive, and - use it only when `args` contains a lambda expression. -/ for i in *...args.size do match fType with | Expr.forallE _ _ b _ => fType := b