76 lines
2.3 KiB
Text
76 lines
2.3 KiB
Text
/-
|
|
Copyright (c) 2021 Microsoft Corporation. All rights reserved.
|
|
Released under Apache 2.0 license as described in the file LICENSE.
|
|
Authors: Leonardo de Moura
|
|
-/
|
|
prelude
|
|
import Lean.ScopedEnvExtension
|
|
import Lean.Util.Recognizers
|
|
import Lean.Util.ReplaceExpr
|
|
|
|
namespace Lean.Compiler
|
|
namespace CSimp
|
|
|
|
structure Entry where
|
|
fromDeclName : Name
|
|
toDeclName : Name
|
|
thmName : Name
|
|
deriving Inhabited
|
|
|
|
structure State where
|
|
map : SMap Name Name := {}
|
|
thmNames : SSet Name := {}
|
|
deriving Inhabited
|
|
|
|
def State.switch : State → State
|
|
| { map, thmNames } => { map := map.switch, thmNames := thmNames.switch }
|
|
|
|
builtin_initialize ext : SimpleScopedEnvExtension Entry State ←
|
|
registerSimpleScopedEnvExtension {
|
|
initial := {}
|
|
addEntry := fun { map, thmNames } { fromDeclName, toDeclName, thmName } => { map := map.insert fromDeclName toDeclName, thmNames := thmNames.insert thmName }
|
|
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, thmName := declName }
|
|
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 s := ext.getState env
|
|
e.replace fun e =>
|
|
if e.isConst then
|
|
match s.map.find? e.constName! with
|
|
| some declNameNew => some (mkConst declNameNew e.constLevels!)
|
|
| none => none
|
|
else
|
|
none
|
|
|
|
end CSimp
|
|
|
|
def hasCSimpAttribute (env : Environment) (declName : Name) : Bool :=
|
|
CSimp.ext.getState env |>.thmNames.contains declName
|
|
|
|
end Lean.Compiler
|