feat: reduce etas in structure instances

src/Lean/Elab/StructInst.lean

@@ -505,8 +505,16 @@ mutual
             return { ref, lhs := [FieldLHS.fieldName ref fieldName], val := val } :: fields
           match Lean.isSubobjectField? env s.structName fieldName with
           | some substructName =>
-            -- If one of the sources has the subobject field, use it
-            if let some val ← s.source.explicit.findSomeM? fun source => mkProjStx? source.stx source.structName fieldName then
+            -- If src is a term for a parent field and the field is that parent projection, use it
+            -- This is necessary as `findField? `A `toA` returns none hence so will `mkProjStx?`
+            -- The resulting extra eta expansion has a surprising effect on performance. Without
+            -- this change, the terms resulting from elaboration are larger. Worse, unification
+            -- must pull them apart when often, without the eta expansion, that would not be
+            -- necessary. This change occurred in #2478.
+            if let some stx := s.source.explicit.find? (·.structName == substructName) |>.map (·.stx) then
+              addField (FieldVal.term stx)
+            -- If one of the sources has the subobject field as a field, use it
+            else if let some val ← s.source.explicit.findSomeM? fun source => mkProjStx? source.stx source.structName fieldName then
               addField (FieldVal.term val)
             else
               let substruct := Struct.mk ref substructName #[] [] s.source

tests/lean/run/structInstFast.lean

@@ -0,0 +1,95 @@
+import Lean
+
+/-! This test uses a macro to generate a structure with `depth` nested structures and `width` fields
+   at each level. The structure is initialized with `val : α` at every field using the `with` construction pattern.
+   ```
+      structure A0 extends A1 where
+      ...
+      def a0 : A0 := { a1 with xm_1 := val, ..., xm_(width) := val}
+      def a'0: A0 := { a'1 with ...}
+      ...
+      structure A(depth) extends Base where
+      ...
+      structure Base where
+        base_1 : α
+        ...
+      ---
+      def base : Base := ⟨val, ..., val⟩
+    ```
+  The test is then
+  ```
+    set_option maxHeartbeats 200 in
+    example : a0 = a'0 := rfl
+  ```
+  As this type of unification occurs over and over and over when type checking declarations, it needs
+  to be and stay fast. Before #2478, this test took over 700 heartbeats.
+-/
+
+open Lean Parser Command
+
+def mkIdent' (s : String) (i : Nat) := mkIdent <| Name.mkSimple <| s++s!"{i}"
+
+def mkFieldsStx (type : TSyntax `ident) (s : String) (width : Nat) : MacroM (TSyntax ``structFields) :=
+  go type s width []
+where go (type : TSyntax `ident) (s : String) (width : Nat)
+    (binds : List <| TSyntax ``structExplicitBinder) : MacroM (TSyntax ``structFields) := do
+  match width with
+  | 0 => `(structFields| $binds.reverse.toArray*)
+  | m+1 =>
+    let new ← `(structExplicitBinder| ($(mkIdent' s m) : $type))
+    go type s m (new::binds)
+
+def baseIdent := mkIdent (Name.mkSimple "base")
+
+def baseTypeIdent := mkIdent (Name.mkSimple "Base")
+
+def init (type val : TSyntax `ident) (width : Nat) : MacroM (TSyntax `command) := do
+  let fieldsStx ← mkFieldsStx type "base_" width
+  let vals := Array.mkArray width val
+  `(structure $baseTypeIdent where
+    $fieldsStx:structFields
+    def $baseIdent : $baseTypeIdent := ⟨$vals,*⟩)
+
+def mkStructAndInstStx (type val : TSyntax `ident) (width depth: Nat) : MacroM <| TSyntax `command := do
+  let init ← init type val width
+  go val width depth #[init]
+where go (val : TSyntax `ident) (width depth : Nat) (cmds : Array <| TSyntax `command) :
+    MacroM <| TSyntax `command := do
+  match depth with
+  | 0 =>
+    let cmd : TSyntax `command := ⟨mkNullNode cmds⟩
+    `($cmd:command)
+  | m+1 =>
+    let len := cmds.data.length
+    let newTerm (s : String) := if len = 1 then baseTypeIdent else mkIdent' s (m+1)
+    let newTerm' (s : String) := if len = 1 then baseIdent else mkIdent' s (m+1)
+    let fieldsStx ← mkFieldsStx type (s!"x{m}_") width
+    let nextStruct ←
+      `(structure $(mkIdent' "A" m) extends $(newTerm "A") where
+        $fieldsStx:structFields)
+    let structVals ← (List.range width).mapM fun j =>
+      `(Term.structInstField| $(mkIdent' s!"x{m}_" j):ident := $val)
+    -- Note here we index backwards to `A0` is the deepest structure with terms `a0` and `a'0`
+    let nextStructInst ←
+      `(def $(mkIdent' "a" m) : $(mkIdent' "A" m) := { $(newTerm' "a"):ident with $structVals.toArray,* })
+    let nextStructInst' ←
+      `(def $(mkIdent' "a'" m) : $(mkIdent' "A" m) := { $(newTerm' "a'"):ident with $structVals.toArray,* })
+    let newCmd : TSyntax `command ←
+      `($nextStruct:command
+        $nextStructInst
+        $nextStructInst')
+    go val width m (cmds.push newCmd)
+
+/-- We build a structure with `depth` nested structures and `width` fields at each level.
+    The structure is initialized with `val` at every field using the `with` construction pattern. -/
+macro "deep_wide_struct_inst_with" type:ident val:ident depth:num width:num : command =>
+  mkStructAndInstStx type val width.getNat depth.getNat
+
+def foo : String := "foo"
+
+deep_wide_struct_inst_with String foo 50 20
+
+/- Structure instances using the `with` pattern should be fast. Without #2478, this takes over 700
+  heartbeats. -/
+set_option maxHeartbeats 200 in
+example : a0 = a'0 := rfl