From a7e8a82e89f7a7f95f13a1b61e29c99a3ce22747 Mon Sep 17 00:00:00 2001
From: Gabriel Ebner <gebner@gebner.org>
Date: Mon, 11 Jul 2022 13:26:30 +0200
Subject: [PATCH] chore: require @[computedField] attribute

---
 src/Lean/Elab/ComputedFields.lean        | 8 ++++++++
 src/Lean/Elab/Inductive.lean             | 6 +++++-
 tests/compiler/computedFieldsExtern.lean | 2 +-
 tests/lean/computedFieldsCode.lean       | 2 +-
 tests/lean/run/computedFields.lean       | 8 ++++----
 5 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/src/Lean/Elab/ComputedFields.lean b/src/Lean/Elab/ComputedFields.lean
index dcd7b9d462..f07f4f23b3 100644
--- a/src/Lean/Elab/ComputedFields.lean
+++ b/src/Lean/Elab/ComputedFields.lean
@@ -32,6 +32,11 @@ This file implements the computed fields feature by simulating it via
 namespace Lean.Elab.ComputedFields
 open Meta
 
+builtin_initialize computedFieldAttr : TagAttribute ←
+  registerTagAttribute `computedField "Marks a function as a computed field of an inductive" fun _ => do
+    unless (← getOptions).getBool `elaboratingComputedFields do
+      throwError "The @[computedField] attribute can only be used in the with-block of an inductive"
+
 def mkUnsafeCastTo (expectedType : Expr) (e : Expr) : MetaM Expr :=
   mkAppOptM ``unsafeCast #[none, expectedType, e]
 
@@ -190,6 +195,9 @@ and the names of the associated computed fields.
 -/
 def setComputedFields (computedFields : Array (Name × Array Name)) : MetaM Unit := do
   for (indName, computedFieldNames) in computedFields do
+    for computedFieldName in computedFieldNames do
+      unless computedFieldAttr.hasTag (← getEnv) computedFieldName do
+        logError m!"'{computedFieldName}' must be tagged with @[computedField]"
     mkComputedFieldOverrides indName computedFieldNames
 
   -- Once all the implementedBy infrastructure is set up, compile everything.
diff --git a/src/Lean/Elab/Inductive.lean b/src/Lean/Elab/Inductive.lean
index 93dbb42837..4d2792e823 100644
--- a/src/Lean/Elab/Inductive.lean
+++ b/src/Lean/Elab/Inductive.lean
@@ -825,7 +825,11 @@ private def applyComputedFields (indViews : Array InductiveView) : CommandElabM
           def%$ref $(mkIdent <| `_root_ ++ declName ++ fieldId):ident : $type $matchAlts:matchAlts)
     let computedFieldNames := indView.computedFields.map fun {fieldId, ..} => declName ++ fieldId
     computedFields := computedFields.push (declName, computedFieldNames)
-  elabCommand <| ← `(set_option bootstrap.genMatcherCode false in mutual $computedFieldDefs* end)
+  withScope (fun scope => { scope with
+      opts := scope.opts
+        |>.setBool `bootstrap.genMatcherCode false
+        |>.setBool `elaboratingComputedFields true}) <|
+    elabCommand <| ← `(mutual $computedFieldDefs* end)
 
   liftTermElabM indViews[0]!.declName do
     ComputedFields.setComputedFields computedFields
diff --git a/tests/compiler/computedFieldsExtern.lean b/tests/compiler/computedFieldsExtern.lean
index 4e9f8bf2aa..b8898dd683 100644
--- a/tests/compiler/computedFieldsExtern.lean
+++ b/tests/compiler/computedFieldsExtern.lean
@@ -2,7 +2,7 @@ inductive Exp
   | var (i : Nat)
   | app (a b : Exp)
 with
-  @[extern c inline "(lean_ctor_get_uint64(#1, lean_ctor_num_objs(#1)*sizeof(void*)) + 40)"]
+  @[computedField, extern c inline "(lean_ctor_get_uint64(#1, lean_ctor_num_objs(#1)*sizeof(void*)) + 40)"]
   hash : Exp → UInt64
     | .var i => Hashable.hash i
     | .app a b => a.hash + b.hash
diff --git a/tests/lean/computedFieldsCode.lean b/tests/lean/computedFieldsCode.lean
index 027af2d3fb..fd12fca632 100644
--- a/tests/lean/computedFieldsCode.lean
+++ b/tests/lean/computedFieldsCode.lean
@@ -11,7 +11,7 @@ inductive Exp
   | a4
   | a5
 with
-  hash : Exp → UInt64
+  @[computedField] hash : Exp → UInt64
     | .var i => Hashable.hash i + 1000
     | .app a b => mixHash (hash a) (hash b)
     | _ => 42
diff --git a/tests/lean/run/computedFields.lean b/tests/lean/run/computedFields.lean
index a8f4f1c8de..735d60cdbf 100644
--- a/tests/lean/run/computedFields.lean
+++ b/tests/lean/run/computedFields.lean
@@ -3,7 +3,7 @@ inductive Exp
   | var (i : UInt32)
   | app (a b : Exp)
 with
-  /-- Computes the hash -/ @[simp] protected hash : Exp → UInt64
+  /-- Computes the hash -/ @[simp, computedField] protected hash : Exp → UInt64
     | .var i => Hashable.hash i
     | .app a b => mixHash a.hash b.hash
     | .hole => 32
@@ -30,7 +30,7 @@ inductive B.C (α : Type u) : Nat → Type u
   | a : C α 0
   | b (c : C α n) {d : C α (n-1)} : C α (n+1)
 with
-  hash : ∀ α i, C α i → UInt64
+  @[computedField] hash : ∀ α i, C α i → UInt64
     | _, _, .a => 1
     | _, _, .b c => 42 + c.hash
 
@@ -50,7 +50,7 @@ mutual
     | a (b : B)
     | b (b : B)
   with
-    f : A → Nat
+    @[computedField] f : A → Nat
       | .a c => 32 + c.f
       | .b c => 42 + 2*c.f
 
@@ -58,7 +58,7 @@ mutual
     | c (a : A)
     | d
   with
-    f : B → Nat
+    @[computedField] f : B → Nat
       | .c a => a.f
       | .d => 0
 end