From 9e246b365ea2e2241279577d2002ed2081b8aa7a Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 9 May 2019 07:35:47 -0700
Subject: [PATCH] feat(library/init/lean): add `SMap` for implementing
 `Environment`

---
 library/init/lean/smap.lean | 59 +++++++++++++++++++++++++++++++++++++
 tests/playground/smap.lean  | 16 ++++++++++
 2 files changed, 75 insertions(+)
 create mode 100644 library/init/lean/smap.lean
 create mode 100644 tests/playground/smap.lean

diff --git a/library/init/lean/smap.lean b/library/init/lean/smap.lean
new file mode 100644
index 0000000000..84c9c56db4
--- /dev/null
+++ b/library/init/lean/smap.lean
@@ -0,0 +1,59 @@
+/-
+Copyright (c) 2019 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import init.data.hashmap init.data.rbmap
+universes u v w w'
+
+namespace Lean
+
+/- Staged map for implementing the Environment. The idea is to store
+   imported entries into a hashtable and local entries into a red-black tree.
+
+   Hypotheses:
+   - The number of entries (i.e., declarations) coming from imported files is much bigger than
+     the number of entries in the current file.
+   - HashMap is faster than RBMap.
+   - When we are reading imported files, we have exclusive access to the map, and efficient
+     destructive updates are performed.
+
+   Remarks:
+   - We never remove declarations from the Environment. In principle, we could support
+     deletion by using `(RBMap α (Option β) lt)` where the value `none` would indicate
+     that an entry was "removed" from the hashtable.
+   - We do not need additional bookkeeping for extracting the local entries.
+-/
+structure SMap (α : Type u) (β : Type v) (lt : α → α → Bool) [HasBeq α] [Hashable α] :=
+(stage₁ : Bool         := true)
+(map₁   : HashMap α β  := {})
+(map₂   : RBMap α β lt := {})
+
+namespace SMap
+variables {α : Type u} {β : Type v} {lt : α → α → Bool} [HasBeq α] [Hashable α]
+
+def empty : SMap α β lt := {}
+instance : HasEmptyc (SMap α β lt) := ⟨SMap.empty⟩
+
+@[specialize] def insert : SMap α β lt → α → β → SMap α β lt
+| ⟨true, m₁, m₂⟩ k v  := ⟨true, m₁.insert k v, m₂⟩
+| ⟨false, m₁, m₂⟩ k v := ⟨false, m₁, m₂.insert k v⟩
+
+@[specialize] def find : SMap α β lt → α → Option β
+| ⟨true, m₁, _⟩ k   := m₁.find k
+| ⟨false, m₁, m₂⟩ k := (m₂.find k).orelse (m₁.find k)
+
+@[specialize] def contains : SMap α β lt → α → Bool
+| ⟨true, m₁, _⟩ k   := m₁.contains k
+| ⟨false, m₁, m₂⟩ k := m₂.contains k || m₁.contains k
+
+/- Move from stage 1 into stage 2. -/
+def switch (m : SMap α β lt) : SMap α β lt :=
+{ stage₁ := false, .. m }
+
+@[inline] def foldStage2 {σ : Type w} (f : σ → α → β → σ) (s : σ) (m : SMap α β lt) : σ :=
+m.map₂.fold f s
+
+end SMap
+end Lean
diff --git a/tests/playground/smap.lean b/tests/playground/smap.lean
new file mode 100644
index 0000000000..c711da5355
--- /dev/null
+++ b/tests/playground/smap.lean
@@ -0,0 +1,16 @@
+import init.lean.smap
+
+abbrev Map : Type := Lean.SMap Nat Bool (λ a b, a < b)
+
+def test1 (n₁ n₂ : Nat) : IO Unit :=
+let m : Map := {} in
+let m := n₁.for (λ i (m : Map), m.insert i (i % 2 == 0)) m in
+let m := m.switch in
+let m := n₂.for (λ i (m : Map), m.insert (i+n₁) (i % 3 == 0)) m in
+do
+n₁.mfor $ λ i, IO.println (i, (m.find i)),
+n₂.mfor $ λ i, IO.println (i+n₁, (m.find (i+n₁))),
+IO.println (m.foldStage2 (λ kvs k v, (k, v)::kvs) [])
+
+def main (xs : List String) : IO Unit :=
+test1 xs.head.toNat xs.tail.head.toNat