From 2387f3c2a22da8a27a4c643b67a0330597572999 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 22 Jul 2019 08:14:35 -0700
Subject: [PATCH] feat(library/init/lean/elaborator): improve `namespace`,
 `scope`, and `end` commands

---
 library/init/lean/elaborator/basic.lean   |  7 ++++
 library/init/lean/elaborator/command.lean | 49 ++++++++++++++++-------
 library/init/lean/name.lean               |  9 +++++
 3 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/library/init/lean/elaborator/basic.lean b/library/init/lean/elaborator/basic.lean
index 05414a9172..820b524268 100644
--- a/library/init/lean/elaborator/basic.lean
+++ b/library/init/lean/elaborator/basic.lean
@@ -19,6 +19,7 @@ structure ElabScope :=
 (cmd     : String)
 (header  : Name)
 (options : Options := {})
+(ns      : Name := Name.anonymous) -- current namespace
 
 structure ElabState :=
 (env      : Environment)
@@ -315,6 +316,12 @@ namespace Elab
 instance {α} : Inhabited (Elab α) :=
 ⟨fun _ => default _⟩
 
+def getNamespace : Elab Name :=
+do s ← get;
+   match s.scopes with
+   | []      => pure Name.anonymous
+   | (sc::_) => pure sc.ns
+
 /- Remark: in an ideal world where performance doesn't matter, we would define `Elab` as
    ```
    ExceptT ElabException (StateT ElabException IO)
diff --git a/library/init/lean/elaborator/command.lean b/library/init/lean/elaborator/command.lean
index 1e3d33e86f..578340b761 100644
--- a/library/init/lean/elaborator/command.lean
+++ b/library/init/lean/elaborator/command.lean
@@ -9,31 +9,52 @@ import init.lean.elaborator.basic
 namespace Lean
 namespace Elab
 
+private def addScopes (cmd : String) : Name → Name → List ElabScope → List ElabScope
+| Name.anonymous ns scopes := scopes
+| (Name.mkString p h) ns scopes :=
+  let ns := Name.mkString ns h;
+  { cmd := cmd, header := h, ns := ns.mkString h } :: addScopes p ns scopes
+| _ _ _ := [] -- unreachable
+
 @[builtinCommandElab «namespace»] def elabNamespace : CommandElab :=
 fun n => do
+  ns     ← getNamespace;
   header ← (n.getArg 1).getIdentVal;
-  modify $ fun s =>
-    let scope : ElabScope := { cmd := "namespace", header := header };
-    { scopes := scope :: s.scopes, .. s }
+  modify $ fun s => { scopes := addScopes "namespace" header ns s.scopes, .. s }
 
 @[builtinCommandElab «section»] def elabSection : CommandElab :=
 fun n => do
+  ns     ← getNamespace;
   header ← (n.getArg 1).getOptionalIdent;
-  let header := header.getOrElse Name.anonymous;
   modify $ fun s =>
-    let scope : ElabScope := { cmd := "section", header := header };
-    { scopes := scope :: s.scopes, .. s }
+    match header with
+    | some header => { scopes := addScopes "section" header ns s.scopes, .. s }
+    | none        => { scopes := { cmd := "section", header := Name.anonymous, ns := ns } :: s.scopes, .. s }
+
+private def getNumEndScopes : Option Name → Nat
+| none     := 1
+| (some n) := n.getNumParts
+
+private def checkAnonymousScope : List ElabScope → Bool
+| ({ header := Name.anonymous, .. } :: _) := true
+| _ := false
+
+private def checkEndHeader : Name → List ElabScope → Bool
+| Name.anonymous _ := true
+| (Name.mkString p s) ({ header := h, .. } :: scopes) := h.eqStr s && checkEndHeader p scopes
+| _ _ := false
 
 @[builtinCommandElab «end»] def elabEnd : CommandElab :=
 fun n => do
-  s ← get;
-  match s.scopes with
-  | [] => throw "invalid 'end', there is no open scope"
-  | (scope::scopes) => do
-    modify $ fun s => { scopes := scopes, .. s };
-    header ← (n.getArg 1).getOptionalIdent;
-    let header := header.getOrElse Name.anonymous;
-    unless (scope.header == header) $ throw "invalid 'end', name mismatch"
+  s      ← get;
+  header ← (n.getArg 1).getOptionalIdent;
+  let num    := getNumEndScopes header;
+  let scopes := s.scopes;
+  modify $ fun s => { scopes := s.scopes.drop num, .. s };
+  when (num > scopes.length) $ throw "invalid 'end', insufficient scopes";
+  match header with
+  | none => unless (checkAnonymousScope scopes) $ throw "invalid 'end', name is missing"
+  | some header => unless (checkEndHeader header scopes) $ throw "invalid 'end', name mismatch"
 
 end Elab
 end Lean
diff --git a/library/init/lean/name.lean b/library/init/lean/name.lean
index c9396e1eea..e2d127e606 100644
--- a/library/init/lean/name.lean
+++ b/library/init/lean/name.lean
@@ -43,6 +43,11 @@ def getPrefix : Name → Name
 | (mkString p s)  := p
 | (mkNumeral p s) := p
 
+def getNumParts : Name → Nat
+| anonymous       := 0
+| (mkString p _)  := getNumParts p + 1
+| (mkNumeral p _) := getNumParts p + 1
+
 def updatePrefix : Name → Name → Name
 | anonymous       newP := anonymous
 | (mkString p s)  newP := mkString newP s
@@ -81,6 +86,10 @@ protected def decEq : ∀ (a b : @& Name), Decidable (a = b)
 instance : DecidableEq Name :=
 {decEq := Name.decEq}
 
+def eqStr : Name → String → Bool
+| (mkString Name.anonymous s) s' := s == s'
+| _ _ := false
+
 protected def append : Name → Name → Name
 | n anonymous := n
 | n (mkString p s) :=