feat: multiple scoped in notation3 (#6793)

See the [discussion on Zulip](https://leanprover.zulipchat.com/#narrow/stream/287929-mathlib4/topic/mathport.20fails/near/387335863). `notation3` did not support more than one occurrence of `scoped` due to an unnecessary limitation in the generated delaborator. It now handles this case correctly.

Author: digama0

Mathlib/Mathport/Notation.lean

@@ -97,8 +97,8 @@ structure MatchState where
   We store the contexts since we need to delaborate expressions after we leave
   scoping constructs. -/
   vars : HashMap Name (SubExpr × LocalContext × LocalInstances)
-  /-- The binders accumulated when matching a `scoped` expression. -/
-  scopeState : Array (TSyntax ``extBinderParenthesized)
+  /-- The binders accumulated while matching a `scoped` expression. -/
+  scopeState : Option (Array (TSyntax ``extBinderParenthesized))
   /-- The arrays of delaborated `Term`s accumulated while matching
   `foldl` and `foldr` expressions. For `foldl`, the arrays are stored in reverse order. -/
   foldState : HashMap Name (Array Term)
@@ -110,7 +110,7 @@ def Matcher := MatchState → DelabM MatchState
 /-- The initial state. -/
 def MatchState.empty : MatchState where
   vars := {}
-  scopeState := #[]
+  scopeState := none
   foldState := {}
 
 /-- Evaluate `f` with the given variable's value as the `SubExpr` and within that subexpression's
@@ -134,19 +134,16 @@ def MatchState.delabVar (s : MatchState) (name : Name) (checkNot? : Option Expr
 def MatchState.captureSubexpr (s : MatchState) (name : Name) : DelabM MatchState := do
   return {s with vars := s.vars.insert name (← readThe SubExpr, ← getLCtx, ← getLocalInstances)}
 
-/-- Push a binder onto the binder array. For `scoped`. -/
-def MatchState.pushBinder (s : MatchState) (b : TSyntax ``extBinderParenthesized) :
-    DelabM MatchState := do
-  let binders := s.scopeState
-  -- TODO merge binders as an inverse to `satisfies_binder_pred%`
-  let binders := binders.push b
-  return {s with scopeState := binders}
-
 /-- Get the accumulated array of delaborated terms for a given foldr/foldl.
 Returns `#[]` if nothing has been pushed yet. -/
 def MatchState.getFoldArray (s : MatchState) (name : Name) : Array Term :=
   (s.foldState.find? name).getD #[]
 
+/-- Get the accumulated array of delaborated terms for a given foldr/foldl.
+Returns `#[]` if nothing has been pushed yet. -/
+def MatchState.getBinders (s : MatchState) : Array (TSyntax ``extBinderParenthesized) :=
+  s.scopeState.getD #[]
+
 /-- Push a delaborated term onto a foldr/foldl array. -/
 def MatchState.pushFold (s : MatchState) (name : Name) (t : Term) : MatchState :=
   let ts := (s.getFoldArray name).push t
@@ -268,46 +265,44 @@ where
 against is in the `lit` variable.
 
 Runs `smatcher`, extracts the resulting `scopeId` variable, processes this value
-(which must be a lambda) to produce a binder, and loops.
-
-Succeeds even if it matches nothing, so it is up to the caller to decide if the
-empty scope state is ok. -/
-partial def matchScoped (lit scopeId : Name) (smatcher : Matcher) : Matcher := fun s => do
-  -- `lit` is bound to the SubExpr that the `scoped` syntax produced
-  s.withVar lit do
-  try
-    -- Run `smatcher` at `lit`, clearing the `scopeId` variable so that it can get a fresh value
-    let s ← smatcher {s with vars := s.vars.erase scopeId}
-    s.withVar scopeId do
-      guard (← getExpr).isLambda
-      let prop ← try Meta.isProp (← getExpr).bindingDomain! catch _ => pure false
-      let isDep := (← getExpr).bindingBody!.hasLooseBVar 0
-      let ppTypes ← getPPOption getPPPiBinderTypes -- the same option controlling ∀
-      let dom ← withBindingDomain delab
-      withBindingBodyUnusedName <| fun x => do
-        let x : Ident := ⟨x⟩
-        let binder ←
-          if prop && !isDep then
-            -- this underscore is used to support binder predicates, since it indicates
-            -- the variable is unused and this binder is safe to merge into another
-            `(extBinderParenthesized|(_ : $dom))
-          else if prop || ppTypes then
-            `(extBinderParenthesized|($x:ident : $dom))
-          else
-            `(extBinderParenthesized|($x:ident))
-        -- Now use the body of the lambda for `lit` for the next iteration
-        let s ← s.captureSubexpr lit
-        let s ← s.pushBinder binder
-        matchScoped lit scopeId smatcher s
-  catch _ =>
-    return s
-
-/-- Like `matchScoped` but ensures that it matches at least one binder. -/
-partial def matchScoped' (lit scopeId : Name) (smatcher : Matcher) : Matcher := fun s => do
-  guard <| s.scopeState.isEmpty
-  let s ← matchScoped lit scopeId smatcher s
-  guard <| !s.scopeState.isEmpty
-  return s
+(which must be a lambda) to produce a binder, and loops. -/
+partial def matchScoped (lit scopeId : Name) (smatcher : Matcher) : Matcher := go #[] where
+  /-- Variant of `matchScoped` after some number of `binders` have already been captured. -/
+  go (binders : Array (TSyntax ``extBinderParenthesized)) : Matcher := fun s => do
+    -- `lit` is bound to the SubExpr that the `scoped` syntax produced
+    s.withVar lit do
+    try
+      -- Run `smatcher` at `lit`, clearing the `scopeId` variable so that it can get a fresh value
+      let s ← smatcher {s with vars := s.vars.erase scopeId}
+      s.withVar scopeId do
+        guard (← getExpr).isLambda
+        let prop ← try Meta.isProp (← getExpr).bindingDomain! catch _ => pure false
+        let isDep := (← getExpr).bindingBody!.hasLooseBVar 0
+        let ppTypes ← getPPOption getPPPiBinderTypes -- the same option controlling ∀
+        let dom ← withBindingDomain delab
+        withBindingBodyUnusedName <| fun x => do
+          let x : Ident := ⟨x⟩
+          let binder ←
+            if prop && !isDep then
+              -- this underscore is used to support binder predicates, since it indicates
+              -- the variable is unused and this binder is safe to merge into another
+              `(extBinderParenthesized|(_ : $dom))
+            else if prop || ppTypes then
+              `(extBinderParenthesized|($x:ident : $dom))
+            else
+              `(extBinderParenthesized|($x:ident))
+          -- Now use the body of the lambda for `lit` for the next iteration
+          let s ← s.captureSubexpr lit
+          -- TODO merge binders as an inverse to `satisfies_binder_pred%`
+          let binders := binders.push binder
+          go binders s
+    catch _ =>
+      guard <| !binders.isEmpty
+      if let some binders₂ := s.scopeState then
+        guard <| binders == binders₂ -- TODO: this might be a bit too strict, but it seems to work
+        return s
+      else
+        return {s with scopeState := binders}
 
 /- Create a `Term` that represents a matcher for `scoped` notation.
 Fails in the `OptionT` sense if a matcher couldn't be constructed.
@@ -317,7 +312,7 @@ partial def mkScopedMatcher (lit scopeId : Name) (scopedTerm : Term) (boundNames
     OptionT TermElabM (List Name × Term) := do
   -- Build the matcher for `scopedTerm` with `scopeId` as an additional variable
   let (keys, smatcher) ← mkExprMatcher scopedTerm (boundNames.insert scopeId)
-  return (keys, ← ``(matchScoped' $(quote lit) $(quote scopeId) $smatcher))
+  return (keys, ← ``(matchScoped $(quote lit) $(quote scopeId) $smatcher))
 
 /-- Matcher for expressions produced by `foldl`. -/
 partial def matchFoldl (lit x y : Name) (smatcher : Matcher) (sinit : Matcher) :
@@ -482,8 +477,6 @@ elab doc:(docComment)? attrs?:(Parser.Term.attributes)? attrKind:Term.attrKind
             mkFoldrMatcher id.getId x.getId y.getId scopedTerm init (getBoundNames boundValues)
         | _ => throwUnsupportedSyntax
     | `(notation3Item| $lit:ident $(prec?)? : (scoped $scopedId:ident => $scopedTerm)) =>
-      if hasScoped then
-        throwErrorAt item "Cannot have more than one `scoped` item."
       hasScoped := true
       (syntaxArgs, pattArgs) ← pushMacro syntaxArgs pattArgs <|←
         `(macroArg| $lit:ident:term $(prec?)?)
@@ -547,7 +540,7 @@ elab doc:(docComment)? attrs?:(Parser.Term.attributes)? attrKind:Term.attrKind
         | .foldr => result ←
           `(let $id := MatchState.getFoldArray s $(quote name); $result)
       if hasBindersItem then
-        result ← `(`(extBinders| $$(MatchState.scopeState s)*) >>= fun binders => $result)
+        result ← `(`(extBinders| $$(MatchState.getBinders s)*) >>= fun binders => $result)
       elabCommand <| ← `(command|
         def $(Lean.mkIdent delabName) : Delab := whenPPOption getPPNotation <|
           getExpr >>= fun e => $matcher MatchState.empty >>= fun s => $result)

test/notation3.lean

@@ -1,6 +1,8 @@
+import Std.Tactic.GuardMsgs
 import Mathlib.Mathport.Notation
 import Mathlib.Init.Data.Nat.Lemmas
 
+set_option pp.unicode.fun true
 set_option autoImplicit true
 
 namespace Test
@@ -13,45 +15,71 @@ def Filter.eventually (p : α → Prop) (f : Filter α) := f p
 
 notation3 "∀ᶠ " (...) " in " f ", " r:(scoped p => Filter.eventually p f) => r
 
-#check ∀ᶠ (x : Nat) (y) in Filter.atTop, x < y
-#check ∀ᶠ x in Filter.atTop, x < 3
+/-- info: ∀ᶠ (x : ℕ) (y : ℕ) in Filter.atTop, x < y : Prop -/
+#guard_msgs in #check ∀ᶠ (x : Nat) (y) in Filter.atTop, x < y
+/-- info: ∀ᶠ (x : ℕ) in Filter.atTop, x < 3 : Prop -/
+#guard_msgs in #check ∀ᶠ x in Filter.atTop, x < 3
 
+def foobar (p : α → Prop) (f : Prop) := ∀ x, p x = f
+
+notation3 "∀ᶠᶠ " (...) " in " f ": "
+  r1:(scoped p => Filter.eventually p f) ", " r2:(scoped p => foobar p r1) => r2
+
+/-- info: ∀ᶠᶠ (x : ℕ) (y : ℕ) in Filter.atTop: x < y, x = y : Prop -/
+#guard_msgs in #check ∀ᶠᶠ (x : Nat) (y) in Filter.atTop: x < y, x = y
+/-- info: ∀ᶠᶠ (x : ℕ) in Filter.atTop: x < 3, x = 1 : Prop -/
+#guard_msgs in #check ∀ᶠᶠ x in Filter.atTop: x < 3, x = 1
+/-- info: ∀ᶠᶠ (x : ℕ) in Filter.atTop: x < 3, x = 1 : Prop -/
+#guard_msgs in #check foobar (fun x ↦ Eq x 1) (Filter.atTop.eventually fun x ↦ LT.lt x 3)
+/-- info: foobar (fun y ↦ y = 1) (∀ᶠ (x : ℕ) in Filter.atTop, x < 3) : Prop -/
+#guard_msgs in #check foobar (fun y ↦ Eq y 1) (Filter.atTop.eventually fun x ↦ LT.lt x 3)
 
 notation3 "∃' " (...) ", " r:(scoped p => Exists p) => r
-#check ∃' x < 3, x < 3
+/-- info: ∃' (x : ℕ) (_ : x < 3), x < 3 : Prop -/
+#guard_msgs in #check ∃' x < 3, x < 3
 
 def func (x : α) : α := x
 notation3 "func! " (...) ", " r:(scoped p => func p) => r
 -- Make sure it handles additional arguments. Should not consume `(· * 2)`.
 -- Note: right now this causes the notation to not pretty print at all.
-#check (func! (x : Nat → Nat), x) (· * 2)
+/-- info: func (fun x ↦ x) fun x ↦ x * 2 : ℕ → ℕ -/
+#guard_msgs in #check (func! (x : Nat → Nat), x) (· * 2)
 
-structure MyUnit
+structure MyUnit where
 notation3 "~{" (x"; "* => foldl (a b => Prod.mk a b) MyUnit) "}~" => x
-#check ~{1; true; ~{2}~}~
-#check ~{}~
+/-- info: ~{1; true; ~{2}~}~ : ((Type × ℕ) × Bool) × Type × ℕ -/
+#guard_msgs in #check ~{1; true; ~{2}~}~
+/-- info: ~{}~ : Type -/
+#guard_msgs in #check ~{}~
 
 notation3 "%[" (x", "* => foldr (a b => List.cons a b) List.nil) "]" => x
-#check %[1, 2, 3]
+/-- info: %[1, 2, 3] : List ℕ -/
+#guard_msgs in #check %[1, 2, 3]
 
 def foo (a : Nat) (f : Nat → Nat) := a + f a
 def bar (a b : Nat) := a * b
 notation3 "*[" x "] " (...) ", " v:(scoped c => bar x (foo x c)) => v
-#check *[1] (x) (y), x + y
-#check bar 1
+/-- info: *[1] (x : ℕ) (y : ℕ), x + y : ℕ -/
+#guard_msgs in #check *[1] (x) (y), x + y
+/-- info: bar 1 : ℕ → ℕ -/
+#guard_msgs in #check bar 1
 
 -- Checking that the `<|` macro is expanded when making matcher
 def foo' (a : Nat) (f : Nat → Nat) := a + f a
 def bar' (a b : Nat) := a * b
 notation3 "*'[" x "] " (...) ", " v:(scoped c => bar' x <| foo' x c) => v
-#check *'[1] (x) (y), x + y
-#check bar' 1
+/-- info: *'[1] (x : ℕ) (y : ℕ), x + y : ℕ -/
+#guard_msgs in #check *'[1] (x) (y), x + y
+/-- info: bar' 1 : ℕ → ℕ -/
+#guard_msgs in #check bar' 1
 
 -- Currently does not pretty print due to pi type
 notation3 (prettyPrint := false) "MyPi " (...) ", " r:(scoped p => (x : _) → p x) => r
-#check MyPi (x : Nat) (y : Nat), x < y
+/-- info: ∀ (x : ℕ), (fun x ↦ ∀ (x_1 : ℕ), (fun y ↦ x < y) x_1) x : Prop -/
+#guard_msgs in #check MyPi (x : Nat) (y : Nat), x < y
 
 -- The notation parses fine, but the delaborator never succeeds, which is expected
 def myId (x : α) := x
 notation3 "BAD " c "; " (x", "* => foldl (a b => b) c) " DAB" => myId x
-#check BAD 1; 2, 3 DAB
+/-- info: myId 3 : ℕ -/
+#guard_msgs in #check BAD 1; 2, 3 DAB