perf: instantiateExprMVars

src/Lean/MetavarContext.lean

@@ -541,20 +541,10 @@ This operation is performed at `instantiateExprMVars`, `elimMVarDeps`, and `leve
 @[extern "lean_instantiate_level_mvars"]
 opaque instantiateLevelMVarsImp (mctx : MetavarContext) (l : Level) : MetavarContext × Level
 
-partial def instantiateLevelMVars [Monad m] [MonadMCtx m] : Level → m Level
-  | lvl@(Level.succ lvl₁)      => return Level.updateSucc! lvl (← instantiateLevelMVars lvl₁)
-  | lvl@(Level.max lvl₁ lvl₂)  => return Level.updateMax! lvl (← instantiateLevelMVars lvl₁) (← instantiateLevelMVars lvl₂)
-  | lvl@(Level.imax lvl₁ lvl₂) => return Level.updateIMax! lvl (← instantiateLevelMVars lvl₁) (← instantiateLevelMVars lvl₂)
-  | lvl@(Level.mvar mvarId)    => do
-    match (← getLevelMVarAssignment? mvarId) with
-    | some newLvl =>
-      if !newLvl.hasMVar then pure newLvl
-      else do
-        let newLvl' ← instantiateLevelMVars newLvl
-        assignLevelMVar mvarId newLvl'
-        pure newLvl'
-    | none        => pure lvl
-  | lvl => pure lvl
+partial def instantiateLevelMVars [Monad m] [MonadMCtx m] (l : Level) : m Level := do
+  let (mctx, lNew) := instantiateLevelMVarsImp (← getMCtx) l
+  setMCtx mctx
+  return lNew
 
 @[extern "lean_instantiate_expr_mvars"]
 opaque instantiateExprMVarsImp (mctx : MetavarContext) (e : Expr) : MetavarContext × Expr

src/kernel/instantiate.cpp

@@ -165,22 +165,38 @@ bool is_head_beta(expr const & t) {
     return is_app(t) && is_lambda(get_app_fn(t));
 }
 
-expr apply_beta(expr f, unsigned num_args, expr const * args) {
-    if (num_args == 0) {
-        return f;
-    } else if (!is_lambda(f)) {
-        return mk_rev_app(f, num_args, args);
-    } else {
-        unsigned m = 1;
-        while (is_lambda(binding_body(f)) && m < num_args) {
-            f = binding_body(f);
-            m++;
+static expr apply_beta_rec(expr e, unsigned i, unsigned num_rev_args, expr const * rev_args, bool preserve_data, bool zeta) {
+    if (is_lambda(e)) {
+        if (i + 1 < num_rev_args) {
+            return apply_beta_rec(binding_body(e), i+1, num_rev_args, rev_args, preserve_data, zeta);
+        } else {
+            return instantiate(binding_body(e), num_rev_args, rev_args);
+        }
+    } else if (is_let(e)) {
+        if (zeta && i < num_rev_args) {
+            return apply_beta_rec(instantiate(let_body(e), let_value(e)), i, num_rev_args, rev_args, preserve_data, zeta);
+        } else {
+            unsigned n = num_rev_args - i;
+            return mk_rev_app(instantiate(e, i, rev_args + n), n, rev_args);
+        }
+    } else if (is_mdata(e)) {
+        if (preserve_data) {
+            unsigned n = num_rev_args - i;
+            return mk_rev_app(instantiate(e, i, rev_args + n), n, rev_args);
+        } else {
+            return apply_beta_rec(mdata_expr(e), i, num_rev_args, rev_args, preserve_data, zeta);
         }
-        lean_assert(m <= num_args);
-        return mk_rev_app(instantiate(binding_body(f), m, args + (num_args - m)), num_args - m, args);
+    } else {
+        unsigned n = num_rev_args - i;
+        return mk_rev_app(instantiate(e, i, rev_args + n), n, rev_args);
     }
 }
 
+expr apply_beta(expr f, unsigned num_rev_args, expr const * rev_args, bool preserve_data, bool zeta) {
+    if (num_rev_args == 0) return f;
+    return apply_beta_rec(f, 0, num_rev_args, rev_args, preserve_data, zeta);
+}
+
 expr head_beta_reduce(expr const & t) {
     if (!is_head_beta(t)) {
         return t;

src/kernel/instantiate.h

@@ -24,7 +24,7 @@ inline expr instantiate_rev(expr const & e, buffer<expr> const & s) {
     return instantiate_rev(e, s.size(), s.data());
 }
 
-expr apply_beta(expr f, unsigned num_rev_args, expr const * rev_args);
+expr apply_beta(expr f, unsigned num_rev_args, expr const * rev_args, bool preserve_data = true, bool zeta = false);
 bool is_head_beta(expr const & t);
 expr head_beta_reduce(expr const & t);
 /* If `e` is of the form `(fun x, t) a` return `head_beta_const_fn(t)` if `t` does not depend on `x`,