feat: add basic support for csimp

src/Lean/Compiler.lean

@@ -11,3 +11,4 @@ import Lean.Compiler.ExternAttr
 import Lean.Compiler.ImplementedByAttr
 import Lean.Compiler.NeverExtractAttr
 import Lean.Compiler.IR
+import Lean.Compiler.CSimpAttr

src/Lean/Compiler/CSimpAttr.lean

@@ -0,0 +1,65 @@
+/-
+Copyright (c) 2021 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+import Lean.ScopedEnvExtension
+import Lean.Util.Recognizers
+import Lean.Util.ReplaceExpr
+
+namespace Lean.Compiler
+namespace CSimp
+
+structure Entry where
+  fromDeclName : Name
+  toDeclName   : Name
+  deriving Inhabited
+
+abbrev State := SMap Name Name
+
+builtin_initialize ext : SimpleScopedEnvExtension Entry State ←
+  registerSimpleScopedEnvExtension {
+    name           := `csimp
+    initial        := {}
+    addEntry       := fun s { fromDeclName, toDeclName } => s.insert fromDeclName toDeclName
+    finalizeImport := fun s => s.switch
+  }
+
+private def isConstantReplacement? (declName : Name) : CoreM (Option Entry) := do
+  let info ← getConstInfo declName
+  match info.type.eq? with
+  | some (_, Expr.const fromDeclName us .., Expr.const toDeclName vs ..) =>
+    if us == vs then
+      return some { fromDeclName, toDeclName }
+    else
+      return none
+  | _ => return none
+
+def add (declName : Name) (kind : AttributeKind) : CoreM Unit := do
+  if let some entry ← isConstantReplacement? declName then
+    ext.add entry kind
+  else
+    throwError "invalid 'csimp' theorem, only constant replacement theorems (e.g., `@f = @g`) are currently supported."
+
+builtin_initialize
+  registerBuiltinAttribute {
+    name  := `csimp
+    descr := "simplification theorem for the compiler"
+    add   := fun declName stx attrKind => do
+      Attribute.Builtin.ensureNoArgs stx
+      discard <| add declName attrKind
+  }
+
+@[export lean_csimp_replace_constants]
+def replaceConstants (env : Environment) (e : Expr) : Expr :=
+  let map := ext.getState env
+  e.replace fun e =>
+    if e.isConst then
+      match map.find? e.constName! with
+      | some declNameNew => some (mkConst declNameNew e.constLevels!)
+      | none => none
+    else
+      none
+
+end CSimp
+end Lean.Compiler