From f8b43630a665401f2142549c5342b35491befe1d Mon Sep 17 00:00:00 2001
From: Gabriel Ebner <gebner@gebner.org>
Date: Mon, 7 Feb 2022 13:40:39 +0100
Subject: [PATCH] fix: refs to copied subobjects in diamond extension

---
 src/Lean/Elab/Structure.lean           | 25 ++++++++++++++++++-------
 tests/lean/diamond10.lean              | 12 ++++++++++++
 tests/lean/diamond10.lean.expected.out |  1 +
 3 files changed, 31 insertions(+), 7 deletions(-)
 create mode 100644 tests/lean/diamond10.lean
 create mode 100644 tests/lean/diamond10.lean.expected.out

diff --git a/src/Lean/Elab/Structure.lean b/src/Lean/Elab/Structure.lean
index 3227436d6e..79634cea56 100644
--- a/src/Lean/Elab/Structure.lean
+++ b/src/Lean/Elab/Structure.lean
@@ -61,7 +61,7 @@ structure StructView where
 
 inductive StructFieldKind where
   | newField | copiedField | fromParent | subobject
-  deriving Inhabited, BEq
+  deriving Inhabited, DecidableEq, Repr
 
 structure StructFieldInfo where
   name     : Name
@@ -70,7 +70,7 @@ structure StructFieldInfo where
   kind     : StructFieldKind
   inferMod : Bool := false
   value?   : Option Expr := none
-  deriving Inhabited
+  deriving Inhabited, Repr
 
 def StructFieldInfo.isFromParent (info : StructFieldInfo) : Bool :=
   match info.kind with
@@ -280,6 +280,14 @@ where
     else
       k infos
 
+/-- Given `obj.foo.bar.baz`, return `obj`. -/
+private partial def getNestedProjectionArg (e : Expr) : MetaM Expr := do
+  if let Expr.const subProjName .. := e.getAppFn then
+    if let some { numParams, .. } ← getProjectionFnInfo? subProjName then
+      if e.getAppNumArgs == numParams + 1 then
+        return ← getNestedProjectionArg e.appArg!
+  return e
+
 /-- Get field type of `fieldName` in `parentStructName`, but replace references
   to other fields of that structure by existing field fvars.
   Auxiliary method for `copyNewFieldsFrom`. -/
@@ -287,17 +295,20 @@ private def getFieldType (infos : Array StructFieldInfo) (parentStructName : Nam
   withLocalDeclD (← mkFreshId) parentType fun parent => do
     let proj ← mkProjection parent fieldName
     let projType ← inferType proj
-    /- Eliminate occurrences of `parent.field`. This happens when the structure contains dependent fields. -/
+    /- Eliminate occurrences of `parent.field`. This happens when the structure contains dependent fields.
+    If the copied parent extended another structure via a subobject,
+    then the occurrence can also look like `parent.toGrandparent.field`
+    (where `toGrandparent` is not a field of the current structure). -/
     let visit (e : Expr) : MetaM TransformStep := do
       if let Expr.const subProjName .. := e.getAppFn then
         if let some { ctorName, numParams, .. } ← getProjectionFnInfo? subProjName then
           let Name.str subStructName subFieldName .. := subProjName
             | throwError "invalid projection name {subProjName}"
           let args := e.getAppArgs
-          if args.get? numParams == parent then
-            let some existingFieldInfo := findFieldInfo? infos subFieldName
-              | throwError "unexpected field access to {fieldName} in{indentExpr e}"
-            return TransformStep.done <| mkAppN existingFieldInfo.fvar args[numParams+1:args.size]
+          if let some major := args.get? numParams then
+            if (← getNestedProjectionArg major) == parent then
+              if let some existingFieldInfo := findFieldInfo? infos subFieldName then
+                return TransformStep.done <| mkAppN existingFieldInfo.fvar args[numParams+1:args.size]
       return TransformStep.done e
     let projType ← Meta.transform projType (post := visit)
     if projType.containsFVar parent.fvarId! then
diff --git a/tests/lean/diamond10.lean b/tests/lean/diamond10.lean
new file mode 100644
index 0000000000..414e518856
--- /dev/null
+++ b/tests/lean/diamond10.lean
@@ -0,0 +1,12 @@
+class Zero (A : Type u) where zero : A
+instance {A} [Zero A] : OfNat A (nat_lit 0) := ⟨Zero.zero⟩
+
+class AddMonoid (A : Type u) extends Add A, Zero A
+class Semiring (R : Type u) extends AddMonoid R
+class SubNegMonoid (A : Type u) extends AddMonoid A, Neg A
+class AddGroup (A : Type u) extends SubNegMonoid A where
+  add_left_neg (a : A) : -a + a = 0
+
+class Ring (R : Type u) extends Semiring R, AddGroup R
+
+#print Ring.mk
diff --git a/tests/lean/diamond10.lean.expected.out b/tests/lean/diamond10.lean.expected.out
new file mode 100644
index 0000000000..ee0a56c1df
--- /dev/null
+++ b/tests/lean/diamond10.lean.expected.out
@@ -0,0 +1 @@
+constructor Ring.mk.{u} : {R : Type u} → [toSemiring : Semiring R] → [toNeg : Neg R] → (∀ (a : R), -a + a = 0) → Ring R