From c4e6e48690ce15cff5c37dffe233ed2ec5ad925b Mon Sep 17 00:00:00 2001 From: Leonardo de Moura Date: Mon, 29 Jan 2024 15:36:10 -0800 Subject: [PATCH] feat: builtin seval simproc attribute --- src/Init/Simproc.lean | 15 +++++- src/Lean/Meta/Tactic/Simp/Simproc.lean | 67 +++++++++++++++++++++----- 2 files changed, 68 insertions(+), 14 deletions(-) diff --git a/src/Init/Simproc.lean b/src/Init/Simproc.lean index 3625b67e5e..8b2acd7dac 100644 --- a/src/Init/Simproc.lean +++ b/src/Init/Simproc.lean @@ -40,7 +40,7 @@ syntax (docComment)? "simproc_decl " ident " (" term ")" " := " term : command /-- A builtin simplification procedure. -/ -syntax (docComment)? attrKind "builtin_simproc " (Tactic.simpPre <|> Tactic.simpPost)? ident " (" term ")" " := " term : command +syntax (docComment)? attrKind "builtin_simproc " (Tactic.simpPre <|> Tactic.simpPost)? ("[" ident,* "]")? ident " (" term ")" " := " term : command /-- A builtin simplification procedure declaration. @@ -72,6 +72,12 @@ syntax (name := sevalprocAttr) "sevalproc" (Tactic.simpPre <|> Tactic.simpPost)? Auxiliary attribute for builtin simplification procedures. -/ syntax (name := simprocBuiltinAttr) "builtin_simproc" (Tactic.simpPre <|> Tactic.simpPost)? : attr + +/-- +Auxiliary attribute for builtin symbolic evaluation procedures. +-/ +syntax (name := sevalprocBuiltinAttr) "builtin_sevalproc" (Tactic.simpPre <|> Tactic.simpPost)? : attr + end Attr macro_rules @@ -112,5 +118,12 @@ macro_rules | `($[$doc?:docComment]? $kind:attrKind builtin_simproc $[$pre?]? $n:ident ($pattern:term) := $body) => do `(builtin_simproc_decl $n ($pattern) := $body attribute [$kind builtin_simproc $[$pre?]?] $n) + | `($[$doc?:docComment]? $kind:attrKind builtin_simproc $[$pre?]? [seval] $n:ident ($pattern:term) := $body) => do + `(builtin_simproc_decl $n ($pattern) := $body + attribute [$kind builtin_sevalproc $[$pre?]?] $n) + | `($[$doc?:docComment]? $kind:attrKind builtin_simproc $[$pre?]? [simp, seval] $n:ident ($pattern:term) := $body) => do + `(builtin_simproc_decl $n ($pattern) := $body + attribute [$kind builtin_simproc $[$pre?]?] $n + attribute [$kind builtin_sevalproc $[$pre?]?] $n) end Lean.Parser diff --git a/src/Lean/Meta/Tactic/Simp/Simproc.lean b/src/Lean/Meta/Tactic/Simp/Simproc.lean index f6fd8b6f59..7032cea9d4 100644 --- a/src/Lean/Meta/Tactic/Simp/Simproc.lean +++ b/src/Lean/Meta/Tactic/Simp/Simproc.lean @@ -10,11 +10,24 @@ import Lean.Meta.Tactic.Simp.Types namespace Lean.Meta.Simp +/-- +Builtin simproc declaration extension state. + +It contains: +- The discrimination tree keys for each simproc (including symbolic evaluation procedures) name. +- The actual procedure associated with a name. +-/ structure BuiltinSimprocs where keys : HashMap Name (Array SimpTheoremKey) := {} procs : HashMap Name Simproc := {} deriving Inhabited +/-- +This global reference is populated using the command `builtin_simproc_pattern%`. + +This reference is then used to process the attributes `builtin_simproc` and `builtin_sevalproc`. +Both attributes need the keys and the actual procedure registered using the command `builtin_simproc_pattern%`. +-/ builtin_initialize builtinSimprocDeclsRef : IO.Ref BuiltinSimprocs ← IO.mkRef {} structure SimprocDecl where @@ -60,6 +73,11 @@ def isBuiltinSimproc (declName : Name) : CoreM Bool := do def isSimproc (declName : Name) : CoreM Bool := return (← getSimprocDeclKeys? declName).isSome +/-- +Given a declaration name `declName`, store the discrimination tree keys and the actual procedure. + +This method is invoked by the command `builtin_simproc_pattern%` elaborator. +-/ def registerBuiltinSimproc (declName : Name) (key : Array SimpTheoremKey) (proc : Simproc) : IO Unit := do unless (← initializing) do throw (IO.userError s!"invalid builtin simproc declaration, it can only be registered during initialization") @@ -85,9 +103,11 @@ instance : ToFormat SimprocEntry where def Simprocs.erase (s : Simprocs) (declName : Name) : Simprocs := { s with erased := s.erased.insert declName, simprocNames := s.simprocNames.erase declName } +/-- Builtin simprocs. -/ builtin_initialize builtinSimprocsRef : IO.Ref Simprocs ← IO.mkRef {} -builtin_initialize builtinSimprocsSEvalRef : IO.Ref Simprocs ← IO.mkRef {} +/-- Builtin symbolic evaluation procedures. -/ +builtin_initialize builtinSEvalprocsRef : IO.Ref Simprocs ← IO.mkRef {} abbrev SimprocExtension := ScopedEnvExtension SimprocOLeanEntry SimprocEntry Simprocs @@ -115,13 +135,22 @@ def addSimprocAttr (ext : SimprocExtension) (declName : Name) (kind : AttributeK throwError "invalid [simproc] attribute, '{declName}' is not a simproc" ext.add { declName, post, keys, proc } kind -def addSimprocBuiltinAttr (declName : Name) (post : Bool) (proc : Simproc) : IO Unit := do +/-- +Implements attributes `builtin_simproc` and `builtin_sevalproc`. +-/ +def addSimprocBuiltinAttrCore (ref : IO.Ref Simprocs) (declName : Name) (post : Bool) (proc : Simproc) : IO Unit := do let some keys := (← builtinSimprocDeclsRef.get).keys.find? declName | throw (IO.userError "invalid [builtin_simproc] attribute, '{declName}' is not a builtin simproc") if post then - builtinSimprocsRef.modify fun s => { s with post := s.post.insertCore keys { declName, keys, post, proc } } + ref.modify fun s => { s with post := s.post.insertCore keys { declName, keys, post, proc } } else - builtinSimprocsRef.modify fun s => { s with pre := s.pre.insertCore keys { declName, keys, post, proc } } + ref.modify fun s => { s with pre := s.pre.insertCore keys { declName, keys, post, proc } } + +def addSimprocBuiltinAttr (declName : Name) (post : Bool) (proc : Simproc) : IO Unit := + addSimprocBuiltinAttrCore builtinSimprocsRef declName post proc + +def addSEvalprocBuiltinAttr (declName : Name) (post : Bool) (proc : Simproc) : IO Unit := + addSimprocBuiltinAttrCore builtinSEvalprocsRef declName post proc def Simprocs.add (s : Simprocs) (declName : Name) (post : Bool) : CoreM Simprocs := do let proc ← @@ -276,22 +305,34 @@ def registerSimprocAttr (attrName : Name) (attrDescr : String) (ref? : Option (I builtin_initialize simprocExtension : SimprocExtension ← registerSimprocAttr `simprocAttr "Simplification procedure" (some builtinSimprocsRef) -builtin_initialize simprocSEvalExtension : SimprocExtension ← registerSimprocAttr `sevalprocAttr "Symbolic evaluator procedure" (some builtinSimprocsSEvalRef) +builtin_initialize simprocSEvalExtension : SimprocExtension ← registerSimprocAttr `sevalprocAttr "Symbolic evaluator procedure" (some builtinSEvalprocsRef) + +private def addBuiltin (declName : Name) (stx : Syntax) (addDeclName : Name) : AttrM Unit := do + let go : MetaM Unit := do + let post := if stx[1].isNone then true else stx[1][0].getKind == ``Lean.Parser.Tactic.simpPost + let val := mkAppN (mkConst addDeclName) #[toExpr declName, toExpr post, mkConst declName] + let initDeclName ← mkFreshUserName (declName ++ `declare) + declareBuiltin initDeclName val + go.run' {} builtin_initialize registerBuiltinAttribute { ref := by exact decl_name% name := `simprocBuiltinAttr descr := "Builtin simplification procedure" - erase := fun _ => pure () -- FIX: eraseSimprocAttr - add := fun declName stx _ => do - let go : MetaM Unit := do - let post := if stx[1].isNone then true else stx[1][0].getKind == ``Lean.Parser.Tactic.simpPost - let val := mkAppN (mkConst ``addSimprocBuiltinAttr) #[toExpr declName, toExpr post, mkConst declName] - let initDeclName ← mkFreshUserName (declName ++ `declare) - declareBuiltin initDeclName val - go.run' {} applicationTime := AttributeApplicationTime.afterCompilation + erase := fun _ => throwError "Not implemented yet, [-builtin_simproc]" + add := fun declName stx _ => addBuiltin declName stx ``addSimprocBuiltinAttr + } + +builtin_initialize + registerBuiltinAttribute { + ref := by exact decl_name% + name := `sevalprocBuiltinAttr + descr := "Builtin symbolic evaluation procedure" + applicationTime := AttributeApplicationTime.afterCompilation + erase := fun _ => throwError "Not implemented yet, [-builtin_sevalproc]" + add := fun declName stx _ => addBuiltin declName stx ``addSEvalprocBuiltinAttr } def getSimprocs : CoreM Simprocs :=