feat: do not instantiate metavariables in kabstract/rw for disallowed occurrences (#2539)

Fixes #2538.

src/Init/Tactics.lean

@@ -355,9 +355,9 @@ Using `rw (config := {occs := .pos L}) [e]`,
 where `L : List Nat`, you can control which "occurrences" are rewritten.
 (This option applies to each rule, so usually this will only be used with a single rule.)
 Occurrences count from `1`.
-At the first occurrence, whether allowed or not,
-arguments of the rewrite rule `e` may be instantiated,
+At each allowed occurrence, arguments of the rewrite rule `e` may be instantiated,
 restricting which later rewrites can be found.
+(Disallowed occurrences do not result in instantiation.)
 `{occs := .neg L}` allows skipping specified occurrences.
 -/
 syntax (name := rewriteSeq) "rewrite" (config)? rwRuleSeq (location)? : tactic

src/Lean/Meta/KAbstract.lean

@@ -16,9 +16,11 @@ By default, all occurrences are abstracted,
 but this behavior can be controlled using the `occs` parameter.
 
 All matches of `p` in `e` are considered for occurrences,
-but at the first match (whether included in the occurrences or not)
+but for each match that is included by the `occs` parameter,
 metavariables appearing in `p` (or `e`) may become instantiated,
 affecting the possibility of subsequent matches.
+For matches that are not included in the `occs` parameter, the metavariable context is rolled back
+to prevent blocking subsequent matches which require different instantiations.
 -/
 def kabstract (e : Expr) (p : Expr) (occs : Occurrences := .all) : MetaM Expr := do
   let e ← instantiateMVars e
@@ -41,15 +43,22 @@ def kabstract (e : Expr) (p : Expr) (occs : Occurrences := .all) : MetaM Expr :=
         visitChildren ()
       else if e.toHeadIndex != pHeadIdx || e.headNumArgs != pNumArgs then
         visitChildren ()
-      else if (← isDefEq e p) then
-        let i ← get
-        set (i+1)
-        if occs.contains i then
-          return mkBVar offset
+      else
+        -- We save the metavariable context here,
+        -- so that it can be rolled back unless `occs.contains i`.
+        let mctx ← getMCtx
+        if (← isDefEq e p) then
+          let i ← get
+          set (i+1)
+          if occs.contains i then
+            return mkBVar offset
+          else
+            -- Revert the metavariable context,
+            -- so that other matches are still possible.
+            setMCtx mctx
+            visitChildren ()
         else
           visitChildren ()
-      else
-        visitChildren ()
     visit e 0 |>.run' 1
 
 end Lean.Meta

tests/lean/interactive/1403.lean.expected.out

@@ -4,5 +4,5 @@
  {"start": {"line": 1, "character": 2}, "end": {"line": 1, "character": 12}},
  "contents":
  {"value":
-  "`rewrite [e]` applies identity `e` as a rewrite rule to the target of the main goal.\nIf `e` is preceded by left arrow (`←` or `<-`), the rewrite is applied in the reverse direction.\nIf `e` is a defined constant, then the equational theorems associated with `e` are used.\nThis provides a convenient way to unfold `e`.\n- `rewrite [e₁, ..., eₙ]` applies the given rules sequentially.\n- `rewrite [e] at l` rewrites `e` at location(s) `l`, where `l` is either `*` or a\n  list of hypotheses in the local context. In the latter case, a turnstile `⊢` or `|-`\n  can also be used, to signify the target of the goal.\n\nUsing `rw (config := {occs := .pos L}) [e]`,\nwhere `L : List Nat`, you can control which \"occurrences\" are rewritten.\n(This option applies to each rule, so usually this will only be used with a single rule.)\nOccurrences count from `1`.\nAt the first occurrence, whether allowed or not,\narguments of the rewrite rule `e` may be instantiated,\nrestricting which later rewrites can be found.\n`{occs := .neg L}` allows skipping specified occurrences.\n",
+  "`rewrite [e]` applies identity `e` as a rewrite rule to the target of the main goal.\nIf `e` is preceded by left arrow (`←` or `<-`), the rewrite is applied in the reverse direction.\nIf `e` is a defined constant, then the equational theorems associated with `e` are used.\nThis provides a convenient way to unfold `e`.\n- `rewrite [e₁, ..., eₙ]` applies the given rules sequentially.\n- `rewrite [e] at l` rewrites `e` at location(s) `l`, where `l` is either `*` or a\n  list of hypotheses in the local context. In the latter case, a turnstile `⊢` or `|-`\n  can also be used, to signify the target of the goal.\n\nUsing `rw (config := {occs := .pos L}) [e]`,\nwhere `L : List Nat`, you can control which \"occurrences\" are rewritten.\n(This option applies to each rule, so usually this will only be used with a single rule.)\nOccurrences count from `1`.\nAt each allowed occurrence, arguments of the rewrite rule `e` may be instantiated,\nrestricting which later rewrites can be found.\n(Disallowed occurrences do not result in instantiation.)\n`{occs := .neg L}` allows skipping specified occurrences.\n",
   "kind": "markdown"}}

tests/lean/rewrite.lean

@@ -63,11 +63,7 @@ variable (f : Nat → Nat) (w : ∀ n, f n = 0)
 
 example : [f 1, f 2, f 1, f 2] = [0, 0, 0, 0] := by
   rw (config := {occs := .pos [2]}) [w]
-  -- Because the metavariables are instantiated after finding the first occurrence,
-  -- even though this is rejected, with the current behaviour this rewrites the second `f 1`,
-  -- rather than the first `f 2`.
-  -- Arguably this behaviour should be changed.
-  trace_state
+  trace_state -- make sure we rewrote the first `f 2`, rather than the second `f 1`.
   rw [w, w]
 
 example : [f 1, f 2, f 1, f 2] = [0, 0, 0, 0] := by
@@ -84,23 +80,34 @@ example : [f 1, f 2, f 1, f 2] = [0, 0, 0, 0] := by
 
 example : [f 1, f 2, f 1, f 2] = [0, 0, 0, 0] := by
   rw (config := {occs := .neg [1]}) [w]
-  -- Again, the rejected first occurrence nevertheless instantiates the metavariables.
-  -- Arguably the state here should be `[f 1, 0, f 1, 0] = [0, 0, 0, 0]`,
-  -- but for now if `[f 1, f 2, 0, f 2] = [0, 0, 0, 0]`
-  trace_state
-  rw [w, w]
+  trace_state -- expecting [f 1, 0, f 1, 0] = [0, 0, 0, 0]
+  rw [w]
 
 example : [f 1, f 2, f 1, f 2] = [0, 0, 0, 0] := by
-  -- Arguably should succeed giving `[f 1, f 2, 0, f 2] = [0, 0, 0, 0]`
-  -- but currently fails.
-  fail_if_success rw (config := {occs := .neg [1, 2]}) [w]
-  trace_state
+  rw (config := {occs := .neg [1, 2]}) [w]
+  trace_state -- expecting [f 1, f 2, 0, f 2] = [0, 0, 0, 0]
   rw [w, w]
 
 example : [f 1, f 2, f 1, f 2] = [0, 0, 0, 0] := by
   rw (config := {occs := .neg [1, 3]}) [w]
-  -- Arguably should give `[f 1, 0, f 1, f 2] = [0, 0, 0, 0]`
-  -- but currently gives `[f 1, f 2, 0, f 2] = [0, 0, 0, 0]`
-  trace_state
+  trace_state -- expecting [f 1, 0, f 1, f 2] = [0, 0, 0, 0]
+  -- This one is slightly confusing:
+  -- We skipped the first `f 1`, because `1 ∈ [1,3]`.
+  -- We then rewrite the first `f 2 = 0`, because it is the second match.
+  -- Now the second `f 1` is no longer eligible, because we have instantiated the argument of `w` to `2`.
+  -- Finally we skip the second `f 2` because of `3 ∈ [1,3]`.
   rw [w, w]
 
+-- A slightly trickier example, where there are already metavariables in the goal,
+-- and we ensure that we don't unify metavariables in positions skipped by `occs`.
+example : { x : Nat × Nat // x.1 = 0 ∧ x.2 = 0 } := by
+  refine ⟨(f ?_, f ?_), ?_⟩
+  rotate_right
+  dsimp
+  rw (config := {occs := .pos [2]}) [w]
+  -- TODO: note that `rw` duplicates goals when there are metavariables.
+  -- This should be fixed.
+  rw [w]
+  exact ⟨rfl, rfl⟩
+  exact 37
+  exact 0

tests/lean/rewrite.lean.expected.out

@@ -33,7 +33,7 @@ h : p → a = b
 ⊢ p
 f : Nat → Nat
 w : ∀ (n : Nat), f n = 0
-⊢ [f 1, f 2, 0, f 2] = [0, 0, 0, 0]
+⊢ [f 1, 0, f 1, f 2] = [0, 0, 0, 0]
 f : Nat → Nat
 w : ∀ (n : Nat), f n = 0
 ⊢ [0, f 2, 0, f 2] = [0, 0, 0, 0]
@@ -42,10 +42,10 @@ w : ∀ (n : Nat), f n = 0
 ⊢ [0, f 2, 0, f 2] = [0, 0, 0, 0]
 f : Nat → Nat
 w : ∀ (n : Nat), f n = 0
-⊢ [f 1, f 2, 0, f 2] = [0, 0, 0, 0]
-f : Nat → Nat
-w : ∀ (n : Nat), f n = 0
-⊢ [f 1, f 2, f 1, f 2] = [0, 0, 0, 0]
+⊢ [f 1, 0, f 1, 0] = [0, 0, 0, 0]
 f : Nat → Nat
 w : ∀ (n : Nat), f n = 0
 ⊢ [f 1, f 2, 0, f 2] = [0, 0, 0, 0]
+f : Nat → Nat
+w : ∀ (n : Nat), f n = 0
+⊢ [f 1, 0, f 1, f 2] = [0, 0, 0, 0]