From 280e7bb0067ba895b7eaf1cd239b68edea87e6a6 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 5 Jun 2019 21:54:28 -0700
Subject: [PATCH] feat(library/init/lean/attributes): add `TagAttribute`s

`TagAttribute`s are implemented on top of the low level Attribute API,
and `PersistentEnvExtension`.
This is just the first attribute on a series of attributes we are
going to implement using Lean itself.
---
 library/init/lean/attributes.lean | 41 +++++++++++++++++++++++++++++++
 library/init/lean/syntax.lean     |  4 +++
 2 files changed, 45 insertions(+)

diff --git a/library/init/lean/attributes.lean b/library/init/lean/attributes.lean
index 3564edc187..663270340b 100644
--- a/library/init/lean/attributes.lean
+++ b/library/init/lean/attributes.lean
@@ -225,4 +225,45 @@ do let env := env.popScopeCore,
    attrs.mfoldl (λ env attr, attr.popScope env) env
 
 end Environment
+
+/--
+   Tag attributes are simple and efficient. They are useful for marking declarations in the modules where
+   they were defined.
+
+   The startup cost for this kind of attribute is very small since `addImportedFn` is a constant function.
+
+   They provide the predicate `tagAttr.hasTag env decl` which returns true iff declaration `decl`
+   is tagged in the environment `env`. -/
+structure TagAttribute :=
+(attr : AttributeImpl)
+(ext  : PersistentEnvExtension Name NameSet)
+
+def registerTagAttribute (name : Name) (descr : String) : IO TagAttribute :=
+do
+ext : PersistentEnvExtension Name NameSet ← registerPersistentEnvExtension {
+  name            := name,
+  addImportedFn   := λ _, {},
+  addEntryFn      := λ (s : NameSet) n, s.insert n,
+  exportEntriesFn := λ es,
+    let r : Array Name := es.fold (λ a e, a.push e) Array.empty in
+    r.qsort Name.quickLt
+},
+let attrImpl : AttributeImpl := {
+  name  := name,
+  descr := descr,
+  add   := λ env decl args persistent, do
+    unless args.isMissing $ throw (IO.userError ("invalid attribute '" ++ toString name ++ "', unexpected argument")),
+    unless persistent $ throw (IO.userError ("invalid attribute '" ++ toString name ++ "', must be persistent")),
+    unless (env.getModuleIdxFor decl).isNone $
+      throw (IO.userError ("invalid attribute '" ++ toString name ++ "', declaration is in an imported module")),
+    pure $ ext.addEntry env decl
+},
+registerAttribute attrImpl,
+pure { attr := attrImpl, ext := ext }
+
+def TagAttribute.hasTag (attr : TagAttribute) (env : Environment) (decl : Name) : Bool :=
+match env.getModuleIdxFor decl with
+| some modIdx := (attr.ext.getModuleEntries env modIdx).binSearchContains decl Name.quickLt
+| none        := (attr.ext.getState env).contains decl
+
 end Lean
diff --git a/library/init/lean/syntax.lean b/library/init/lean/syntax.lean
index afa786b243..61d925aa95 100644
--- a/library/init/lean/syntax.lean
+++ b/library/init/lean/syntax.lean
@@ -78,6 +78,10 @@ inductive Syntax
 instance stxInh : Inhabited Syntax :=
 ⟨Syntax.missing⟩
 
+def Syntax.isMissing : Syntax → Bool
+| Syntax.missing := true
+| _ := false
+
 def SyntaxNodeKind.fix : SyntaxNodeKind → IO SyntaxNodeKind
 | {name := n, ..} := do
   m ← nameToKindTable.get,