feat: use structure extension to implement Structure.lean

src/Lean/Elab/Structure.lean

@@ -209,8 +209,6 @@ private def expandFields (structStx : Syntax) (structModifiers : Modifiers) (str
     let idents := fieldBinder[2].getArgs
     idents.foldlM (init := views) fun (views : Array StructFieldView) ident => withRef ident do
       let name     := ident.getId
-      if isInternalSubobjectFieldName name then
-        throwError "invalid field name '{name}', identifiers starting with '_' are reserved to the system"
       let declName := structDeclName ++ name
       let declName ← applyVisibility fieldModifiers.visibility declName
       addDocString' declName fieldModifiers.docString?

src/Lean/Expr.lean

@@ -446,6 +446,10 @@ def isLit : Expr → Bool
   | lit .. => true
   | _      => false
 
+def getForallBody : Expr → Expr
+  | forallE _ _ b .. => getForallBody b
+  | e                => e
+
 def getAppFn : Expr → Expr
   | app f a _ => getAppFn f
   | e         => e

src/Lean/Structure.lean

@@ -8,35 +8,36 @@ Helper functions for retrieving structure information.
 import Lean.Environment
 import Lean.ProjFns
 
-/- TODO: We currently assume that the projection function for field `fieldName` at structure `structName` is
-   `structName ++ fieldName`. This is incorrect for private projections.
-   We will fix this by storing a mapping from `structure` + `fieldName` to projection function name in the environment.
-   This modification will impact functions such as `getStructureFields` and `getProjFnForField?` -/
-
 namespace Lean
 
 structure StructureFieldInfo where
   fieldName : Name
   projFn    : Name
-  subobject : Bool
+  subobject : Bool -- TODO: use `Option Name`
+  deriving Inhabited, Repr
+
+def StructureFieldInfo.lt (i₁ i₂ : StructureFieldInfo) : Bool :=
+  Name.quickLt i₁.fieldName i₂.fieldName
+
+structure StructureInfo where
+  structName : Name
+  fieldNames : Array Name := #[] -- sorted by field position in the structure
+  fieldInfo  : Array StructureFieldInfo := #[] -- sorted by `fieldName`
   deriving Inhabited
 
-private structure StructureEntry where
-  structName : Name
-  fieldNames : Array Name -- sorted by field position in the structure
-  fieldInfo  : Array StructureFieldInfo -- sorted by `fieldName`
-  deriving Inhabited
+def StructureInfo.lt (i₁ i₂ : StructureInfo) : Bool :=
+  Name.quickLt i₁.structName i₂.structName
 
 /-- Auxiliary state for structures defined in the current module. -/
 private structure StructureState where
-  map : Std.PersistentHashMap Name StructureEntry := {}
+  map : Std.PersistentHashMap Name StructureInfo := {}
   deriving Inhabited
 
-builtin_initialize structureExt : SimplePersistentEnvExtension StructureEntry StructureState ← registerSimplePersistentEnvExtension {
+builtin_initialize structureExt : SimplePersistentEnvExtension StructureInfo StructureState ← registerSimplePersistentEnvExtension {
   name          := `structExt
   addImportedFn := fun _ => {}
   addEntryFn    := fun s e => { s with map := s.map.insert e.structName e }
-  toArrayFn     := fun es => es.toArray.qsort fun e₁ e₂ => Name.quickLt e₁.structName e₂.structName
+  toArrayFn     := fun es => es.toArray.qsort StructureInfo.lt
 }
 
 structure StructureDescr where
@@ -48,9 +49,14 @@ def registerStructure (env : Environment) (e : StructureDescr) : Environment :=
   structureExt.addEntry env {
     structName := e.structName
     fieldNames := e.fields.map fun e => e.fieldName
-    fieldInfo  := e.fields.qsort fun e₁ e₂ => Name.quickLt e₁.fieldName e₂.fieldName
+    fieldInfo  := e.fields.qsort StructureFieldInfo.lt
   }
 
+def getStructureInfo? (env : Environment) (structName : Name) : Option StructureInfo :=
+  match env.getModuleIdxFor? structName with
+  | some modIdx => structureExt.getModuleEntries env modIdx |>.binSearch { structName } StructureInfo.lt
+  | none        => structureExt.getState env |>.map.find? structName
+
 /--
   Return true iff `constName` is the a non-recursive inductive
   datatype that has only one constructor. -/
@@ -63,14 +69,6 @@ def isStructureLike (env : Environment) (constName : Name) : Bool :=
 def mkInternalSubobjectFieldName (fieldName : Name) : Name :=
   fieldName.appendBefore "_"
 
-def isInternalSubobjectFieldName : Name → Bool
-  | Name.str _ s _ => s.length > 0 && s.get 0 == '_'
-  | _              => false
-
-def deinternalizeFieldName : Name → Name
-  | n@(Name.str p s _) => if s.length > 0 && s.get 0 == '_' then Name.mkStr p (s.drop 1) else n
-  | n                  => n
-
 def getStructureCtor (env : Environment) (constName : Name) : ConstructorVal :=
   match env.find? constName with
   | some (ConstantInfo.inductInfo { isRec := false, ctors := [ctorName], .. }) =>
@@ -79,38 +77,33 @@ def getStructureCtor (env : Environment) (constName : Name) : ConstructorVal :=
     | _ => panic! "ill-formed environment"
   | _ => panic! "structure expected"
 
-private def getStructureFieldsAux (numParams : Nat) : Nat → Expr → Array Name → Array Name
-  | i, Expr.forallE n d b _, fieldNames =>
-    if i < numParams then
-      getStructureFieldsAux numParams (i+1) b fieldNames
-    else
-      getStructureFieldsAux numParams (i+1) b <| fieldNames.push <| deinternalizeFieldName n
-  | _, _, fieldNames => fieldNames
-
--- TODO: fix. See comment in the beginning of the file
-
 /-- Get direct field names for the given structure. -/
 def getStructureFields (env : Environment) (structName : Name) : Array Name :=
-  let ctor := getStructureCtor env structName;
-  getStructureFieldsAux ctor.numParams 0 ctor.type #[]
+  if let some info := getStructureInfo? env structName then
+    info.fieldNames
+  else
+    panic! "structure expected"
 
-private def isSubobjectFieldAux (numParams : Nat) (target : Name) : Nat → Expr → Option Name
-  | i, Expr.forallE n d b _ =>
-    if i < numParams then
-      isSubobjectFieldAux numParams target (i+1) b
-    else if n == target then
-      match d.getAppFn with
-      | Expr.const parentStructName _ _ => some parentStructName
-      | _ => panic! "ill-formed structure"
-    else
-      isSubobjectFieldAux numParams target (i+1) b
-  | _, _ => none
+def getFieldInfo? (env : Environment) (structName : Name) (fieldName : Name) : Option StructureFieldInfo :=
+  if let some info := getStructureInfo? env structName then
+    info.fieldInfo.binSearch { fieldName := fieldName, projFn := arbitrary, subobject := arbitrary } StructureFieldInfo.lt
+  else
+    none
 
--- TODO: fix. See comment in the beginning of the file
 /-- If `fieldName` represents the relation to a parent structure `S`, return `S` -/
 def isSubobjectField? (env : Environment) (structName : Name) (fieldName : Name) : Option Name :=
-  let ctor := getStructureCtor env structName;
-  isSubobjectFieldAux ctor.numParams (mkInternalSubobjectFieldName fieldName) 0 ctor.type
+  if let some fieldInfo := getFieldInfo? env structName fieldName then
+    if fieldInfo.subobject then
+      match env.find? fieldInfo.projFn with
+      | some (ConstantInfo.defnInfo val) =>
+        match val.type.getForallBody.getAppFn with
+        | Expr.const n .. => some n
+        | _ => panic! "ill-formed structure"
+      | _ => panic! "ill-formed environment"
+    else
+      none
+  else
+    none
 
 /-- Return immediate parent structures -/
 def getParentStructures (env : Environment) (structName : Name) : Array Name :=
@@ -147,16 +140,6 @@ private partial def getStructureFieldsFlattenedAux (env : Environment) (structNa
 def getStructureFieldsFlattened (env : Environment) (structName : Name) (includeSubobjectFields := true) : Array Name :=
   getStructureFieldsFlattenedAux env structName #[] includeSubobjectFields
 
--- TODO: fix. See comment in the beginning of the file
-private def hasProjFn (env : Environment) (structName : Name) (numParams : Nat) : Nat → Expr → Bool
-  | i, Expr.forallE n d b _ =>
-    if i < numParams then
-      hasProjFn env structName numParams (i+1) b
-    else
-      let fullFieldName := structName ++ deinternalizeFieldName n;
-      env.isProjectionFn fullFieldName
-  | _, _ => false
-
 /--
   Return true if `constName` is the name of an inductive datatype
   created using the `structure` or `class` commands.
@@ -164,17 +147,11 @@ private def hasProjFn (env : Environment) (structName : Name) (numParams : Nat)
   We perform the check by testing whether auxiliary projection functions
   have been created. -/
 def isStructure (env : Environment) (constName : Name) : Bool :=
-  if isStructureLike env constName then
-    let ctor := getStructureCtor env constName;
-    hasProjFn env constName ctor.numParams 0 ctor.type
-  else
-    false
+  getStructureInfo? env constName |>.isSome
 
--- TODO: fix. See comment in the beginning of the file
 def getProjFnForField? (env : Environment) (structName : Name) (fieldName : Name) : Option Name :=
-  let fieldNames := getStructureFields env structName;
-  if fieldNames.any fun n => fieldName == n then
-    some (structName ++ fieldName)
+  if let some fieldInfo := getFieldInfo? env structName fieldName then
+    some fieldInfo.projFn
   else
     none
 

tests/lean/interactive/goTo.lean

@@ -16,10 +16,10 @@ def mkFoo₁ : Foo := {
   bar := ⟨⟩
 }
 
-inductive HandWrittenStruct where
-  | mk (n : Nat)
+structure HandWrittenStruct where
+  n : Nat
 
-def HandWrittenStruct.n := fun | mk n => n
+-- def HandWrittenStruct.n := fun | mk n => n
 
           --v textDocument/definition
 def hws : HandWrittenStruct := {

tests/lean/interactive/goTo.lean.expected.out

@@ -32,7 +32,7 @@
   {"start": {"line": 18, "character": 10},
    "end": {"line": 18, "character": 27}},
   "targetRange":
-  {"start": {"line": 18, "character": 0}, "end": {"line": 19, "character": 16}},
+  {"start": {"line": 18, "character": 0}, "end": {"line": 19, "character": 9}},
   "originSelectionRange":
   {"start": {"line": 24, "character": 10},
    "end": {"line": 24, "character": 27}}}]
@@ -40,9 +40,9 @@
  "position": {"line": 26, "character": 2}}
 [{"targetUri": "file://goTo.lean",
   "targetSelectionRange":
-  {"start": {"line": 21, "character": 4}, "end": {"line": 21, "character": 23}},
+  {"start": {"line": 19, "character": 2}, "end": {"line": 19, "character": 3}},
   "targetRange":
-  {"start": {"line": 21, "character": 0}, "end": {"line": 21, "character": 42}},
+  {"start": {"line": 19, "character": 2}, "end": {"line": 19, "character": 3}},
   "originSelectionRange":
   {"start": {"line": 26, "character": 2}, "end": {"line": 26, "character": 3}}}]
 {"textDocument": {"uri": "file://goTo.lean"},

tests/lean/struct1.lean

@@ -1,4 +1,4 @@
-
+--
 
 structure A (α : Type) :=
 (x : α)
@@ -18,10 +18,10 @@ structure S extends A Nat, A Bool := -- error field toA already declared
 structure S extends A Nat, B Bool := -- error field `x` from `B` has already been declared
 (x : Nat)
 
-structure S := -- error '_' is not allowed
+structure S1 :=
 (_x : Nat)
 
-structure S := -- error '_' is not allowed
+structure S2 :=
 (x _y : Nat)
 
 structure S :=

tests/lean/struct1.lean.expected.out

@@ -2,8 +2,6 @@ struct1.lean:9:14-9:17: error: expected Type
 struct1.lean:12:20-12:29: error: expected structure
 struct1.lean:15:27-15:33: error: field 'toA' has already been declared
 struct1.lean:18:27-18:33: error: field 'x' from 'B' has already been declared
-struct1.lean:22:1-22:3: error: invalid field name '_x', identifiers starting with '_' are reserved to the system
-struct1.lean:25:3-25:5: error: invalid field name '_y', identifiers starting with '_' are reserved to the system
 struct1.lean:29:1-29:2: error: field 'x' has already been declared
 struct1.lean:32:1-32:2: error: field 'x' has been declared in parent structure
 struct1.lean:35:6-35:10: error: type mismatch