fix: make cdot anonymous function notation handle ambiguous notation

Fixes an issue where each alternative in choice nodes would get their own arguments. Now cdot function expansion is aware of choice nodes.

Also modifies the variable naming so that multi-argument functions like `(· + ·)` expand as `fun x1 x2 => x1 + x2` rather than `fun x x_1 => x + x_1`.

Closes #4832

Author: kmill

src/Lean/Elab/BuiltinNotation.lean

@@ -261,27 +261,62 @@ partial def hasCDot : Syntax → Bool
   - `f · · b` => `fun _a_1 _a_2 => f _a_1 _a_2 b` -/
 partial def expandCDot? (stx : Term) : MacroM (Option Term) := do
   if hasCDot stx then
-    let (newStx, binders) ← (go stx).run #[]
+    let numArgs ← countArgs stx
+    let (newStx, binders) ← withFreshMacroScope <| (go numArgs stx).run #[]
     `(fun $binders* => $(⟨newStx⟩))
   else
     pure none
 where
   /--
-    Auxiliary function for expanding the `·` notation.
-    The extra state `Array Syntax` contains the new binder names.
-    If `stx` is a `·`, we create a fresh identifier, store in the
-    extra state, and return it. Otherwise, we just return `stx`. -/
-  go : Syntax → StateT (Array Ident) MacroM Syntax
-    | stx@`(($(_))) => pure stx
-    | stx@`(·) => withFreshMacroScope do
-      let id ← mkFreshIdent stx (canonical := true)
-      modify (·.push id)
-      pure id
-    | stx => match stx with
-      | .node _ k args => do
-        let args ← args.mapM go
-        return .node (.fromRef stx (canonical := true)) k args
-      | _ => pure stx
+  Count arguments, taking into consideration choice nodes
+  and validating that each alternative has the same number of arguments.
+  -/
+  countArgs : Syntax → MacroM Nat
+  | `(($_)) => return 0
+  | `(·) => return 1
+  | stx =>
+    match stx with
+    | .node _ k args => do
+      if k == choiceKind then
+        if args.isEmpty then
+          Macro.throwUnsupported
+        let counts ← args.mapM countArgs
+        unless counts.all (fun c => counts[0]! == c) do
+          Macro.throwErrorAt stx "Ambiguous notation in cdot function has different numbers of '·' arguments in each alternative."
+        return counts[0]!
+      else
+        args.foldlM (init := 0) (fun acc arg => return acc + (← countArgs arg))
+    | _ => return 0
+  /--
+  Auxiliary function for expanding the `·` notation.
+  The extra state `Array Syntax` contains the new binder names.
+  If `stx` is a `·`, we create a fresh identifier, store it in the
+  extra state, and return it. Otherwise, we just return `stx`.
+  -/
+  go (numArgs : Nat) : Syntax → StateT (Array Ident) MacroM Syntax
+  | stx@`(($(_))) => pure stx
+  | stx@`(·) => do
+    let i := (← get).size
+    let name ← MonadQuotation.addMacroScope <|
+      if numArgs == 1 then
+        `x
+      else
+        Name.mkSimple s!"x{i + 1}"
+    let id := mkIdentFrom stx name (canonical := true)
+    modify (fun s => s.push id)
+    pure id
+  | stx => match stx with
+    | .node _ k args => do
+      let args ←
+        if k == choiceKind then
+          let s ← get
+          let args' ← args.mapM (fun arg => go numArgs arg |>.run s)
+          set args'[0]!.2
+          pure <| args'.map Prod.fst
+        else
+          args.mapM (go numArgs)
+      return .node (.fromRef stx (canonical := true)) k args
+    | _ => pure stx
 
 /--
   Helper method for elaborating terms such as `(.+.)` where a constant name is expected.

tests/lean/diamond1.lean.expected.out

@@ -11,5 +11,5 @@ toBar : Bar (α → α)
 c : Bool → α
 d : Nat
 def f : Nat → Foo Nat :=
-fun x => { a := fun y => x + y, b := fun x x_1 => x + x_1, c := fun x_1 => x, d := x }
+fun x => { a := fun y => x + y, b := fun x1 x2 => x1 + x2, c := fun x_1 => x, d := x }
 diamond1.lean:27:47-27:52: warning: field 'a' from 'Baz' has already been declared

tests/lean/interactive/hover.lean.expected.out

@@ -407,7 +407,7 @@
   "end": {"line": 206, "character": 13}},
  "contents":
  {"value":
-  "```lean\n?m x✝¹ x✝\n```\n***\n`a + b` computes the sum of `a` and `b`.\nThe meaning of this notation is type-dependent. ",
+  "```lean\n?m x1✝ x2✝\n```\n***\n`a + b` computes the sum of `a` and `b`.\nThe meaning of this notation is type-dependent. ",
   "kind": "markdown"}}
 {"textDocument": {"uri": "file:///hover.lean"},
  "position": {"line": 215, "character": 28}}

tests/lean/run/cdotTests.lean

@@ -31,3 +31,15 @@ rfl
 
 theorem ex4 : sum [1, 2, 3, 4] = 10 :=
 rfl
+
+/-
+Check that ambiguous notation inside cdot notation still has only a single argument.
+(Need to process choice nodes specially.)
+-/
+
+def tag (_ : α) (y : α) := y
+notation "f" x => tag 1 x
+notation "f" x => tag "2" x
+/-- info: fun x => (f x).length : String → Nat -/
+#guard_msgs in
+#check (String.length <| f ·)