From 526cbfbcd0ffdfb76e6d31e909904ec83ee49a7a Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Tue, 3 Aug 2021 14:29:36 -0700
Subject: [PATCH] refactor: add `ImportingFlag.lean`

---
 src/Lean/Attributes.lean    |  3 ++-
 src/Lean/Data/Options.lean  |  3 +++
 src/Lean/Environment.lean   | 15 +++------------
 src/Lean/ImportingFlag.lean | 22 ++++++++++++++++++++++
 4 files changed, 30 insertions(+), 13 deletions(-)
 create mode 100644 src/Lean/ImportingFlag.lean

diff --git a/src/Lean/Attributes.lean b/src/Lean/Attributes.lean
index 0230f7c35b..9b2bcd1c60 100644
--- a/src/Lean/Attributes.lean
+++ b/src/Lean/Attributes.lean
@@ -47,7 +47,8 @@ builtin_initialize attributeMapRef : IO.Ref (PersistentHashMap Name AttributeImp
 def registerBuiltinAttribute (attr : AttributeImpl) : IO Unit := do
   let m ← attributeMapRef.get
   if m.contains attr.name then throw (IO.userError ("invalid builtin attribute declaration, '" ++ toString attr.name ++ "' has already been used"))
-  unless (← IO.initializing) || (← importing) do throw (IO.userError "failed to register attribute, attributes can only be registered during initialization")
+  unless (← initializing) do
+    throw (IO.userError "failed to register attribute, attributes can only be registered during initialization")
   attributeMapRef.modify fun m => m.insert attr.name attr
 
 /-
diff --git a/src/Lean/Data/Options.lean b/src/Lean/Data/Options.lean
index bfc164804f..fb4d1cac0f 100644
--- a/src/Lean/Data/Options.lean
+++ b/src/Lean/Data/Options.lean
@@ -3,6 +3,7 @@ Copyright (c) 2018 Microsoft Corporation. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Sebastian Ullrich and Leonardo de Moura
 -/
+import Lean.ImportingFlag
 import Lean.Data.KVMap
 
 namespace Lean
@@ -33,6 +34,8 @@ private constant optionDeclsRef : IO.Ref OptionDecls
 
 @[export lean_register_option]
 def registerOption (name : Name) (decl : OptionDecl) : IO Unit := do
+  unless (← initializing) do
+    throw (IO.userError "failed to register option, options can only be registered during initialization")
   let decls ← optionDeclsRef.get
   if decls.contains name then
     throw $ IO.userError s!"invalid option declaration '{name}', option already exists"
diff --git a/src/Lean/Environment.lean b/src/Lean/Environment.lean
index e09055739c..be2dfc68fe 100644
--- a/src/Lean/Environment.lean
+++ b/src/Lean/Environment.lean
@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 import Std.Data.HashMap
+import Lean.ImportingFlag
 import Lean.Data.SMap
 import Lean.Declaration
 import Lean.LocalContext
@@ -12,13 +13,6 @@ import Lean.Util.FindExpr
 import Lean.Util.Profile
 
 namespace Lean
-builtin_initialize importingRef : IO.Ref Bool ← IO.mkRef false
-
-/- True while modules are being imported. We use this flag to test check whether environment extensions are registered only
-   during initialization (builtin ones), and importing (user defined ones). -/
-def importing : IO Bool :=
-  importingRef.get
-
 /- Opaque environment extension state. -/
 constant EnvExtensionStateSpec : PointedType.{0}
 def EnvExtensionState : Type := EnvExtensionStateSpec.type
@@ -224,7 +218,7 @@ unsafe def getState {σ} [Inhabited σ] (ext : Ext σ) (env : Environment) : σ
     panic! invalidExtMsg
 
 unsafe def registerExt {σ} (mkInitial : IO σ) : IO (Ext σ) := do
-  unless (← IO.initializing) || (← importing) do
+  unless (← initializing) do
     throw (IO.userError "failed to register environment, extensions can only be registered during initialization")
   let exts ← envExtensionsRef.get
   let idx := exts.size
@@ -613,8 +607,7 @@ structure ImportState where
 
 @[export lean_import_modules]
 partial def importModules (imports : List Import) (opts : Options) (trustLevel : UInt32 := 0) : IO Environment := profileitIO "import" opts do
-  try
-    importingRef.set true
+  withImporting do
     let (_, s) ← importMods imports |>.run {}
     let mut numConsts := 0
     for mod in s.moduleData do
@@ -648,8 +641,6 @@ partial def importModules (imports : List Import) (opts : Options) (trustLevel :
     let env ← setImportedEntries env s.moduleData
     let env ← finalizePersistentExtensions env s.moduleData opts
     pure env
-  finally
-    importingRef.set false
 where
   importMods : List Import → StateRefT ImportState IO Unit
   | []    => pure ()
diff --git a/src/Lean/ImportingFlag.lean b/src/Lean/ImportingFlag.lean
new file mode 100644
index 0000000000..250bb7c355
--- /dev/null
+++ b/src/Lean/ImportingFlag.lean
@@ -0,0 +1,22 @@
+/-
+Copyright (c) 2021 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+namespace Lean
+
+builtin_initialize importingRef : IO.Ref Bool ← IO.mkRef false
+
+/- We say Lean is "initializing" when it is executing `builtin_initialize` declarations and importing modules.
+   Recall that Lean excutes `initialize` declarations while importing modules. -/
+def initializing : IO Bool :=
+  IO.initializing <||> importingRef.get
+
+def withImporting (x : IO α) : IO α :=
+  try
+    importingRef.set true
+    x
+  finally
+    importingRef.set false
+
+end Lean