From 5064115ff1cf14d04eb41b1ffac86a33ea818871 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 3 Dec 2020 16:40:00 -0800
Subject: [PATCH] feat: trie for hierarchical names

---
 src/Lean/Data.lean            |   2 +
 src/Lean/Data/NameTrie.lean   |  66 ++++++++++++++++++++
 src/Lean/Data/PrefixTree.lean | 111 ++++++++++++++++++++++++++++++++++
 3 files changed, 179 insertions(+)
 create mode 100644 src/Lean/Data/NameTrie.lean
 create mode 100644 src/Lean/Data/PrefixTree.lean

diff --git a/src/Lean/Data.lean b/src/Lean/Data.lean
index ddb5732a95..61f35a40fe 100644
--- a/src/Lean/Data.lean
+++ b/src/Lean/Data.lean
@@ -17,3 +17,5 @@ import Lean.Data.Options
 import Lean.Data.Position
 import Lean.Data.SMap
 import Lean.Data.Trie
+import Lean.Data.PrefixTree
+import Lean.Data.NameTrie
diff --git a/src/Lean/Data/NameTrie.lean b/src/Lean/Data/NameTrie.lean
new file mode 100644
index 0000000000..e932c5ebdb
--- /dev/null
+++ b/src/Lean/Data/NameTrie.lean
@@ -0,0 +1,66 @@
+/-
+Copyright (c) 2020 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+import Lean.Data.PrefixTree
+
+namespace Lean
+
+inductive NamePart
+  | str (s : String)
+  | num (n : Nat)
+
+instance : ToString NamePart where
+  toString
+    | NamePart.str s => s
+    | NamePart.num n => toString n
+
+def NamePart.lt : NamePart → NamePart → Bool
+  | NamePart.str a, NamePart.str b => a < b
+  | NamePart.num a, NamePart.num b => a < b
+  | NamePart.num _, NamePart.str _ => true
+  | _, _ => false
+
+def NameTrie (β : Type u) := PrefixTree NamePart β NamePart.lt
+
+private def toKey (n : Name) : List NamePart :=
+  loop n []
+where
+  loop
+    | Name.str p s _, parts   => loop p (NamePart.str s :: parts)
+    | Name.num p n _, parts   => loop p (NamePart.num n :: parts)
+    | Name.anonymous, parts => parts
+
+def NameTrie.insert (t : NameTrie β) (n : Name) (b : β) : NameTrie β :=
+  PrefixTree.insert t (toKey n) b
+
+def NameTrie.empty : NameTrie β :=
+  PrefixTree.empty
+
+instance : Inhabited (NameTrie β) where
+  default := NameTrie.empty
+
+instance : EmptyCollection (NameTrie β) where
+  emptyCollection := NameTrie.empty
+
+def NameTrie.find? (t : NameTrie β) (k : Name) : Option β :=
+  PrefixTree.find? t (toKey k)
+
+@[inline]
+def NameTrie.foldMatchingM [Monad m] (t : NameTrie β) (k : Name) (init : σ) (f : β → σ → m σ) : m σ :=
+  PrefixTree.foldMatchingM t (toKey k) init f
+
+@[inline]
+def NameTrie.foldM [Monad m] (t : NameTrie β) (init : σ) (f : β → σ → m σ) : m σ :=
+  t.foldMatchingM Name.anonymous init f
+
+@[inline]
+def NameTrie.forMatchingM [Monad m] (t : NameTrie β) (k : Name) (f : β → m Unit) : m Unit :=
+  PrefixTree.forMatchingM t (toKey k) f
+
+@[inline]
+def NameTrie.forM [Monad m] (t : NameTrie β) (f : β → m Unit) : m Unit :=
+  t.forMatchingM Name.anonymous f
+
+end Lean
diff --git a/src/Lean/Data/PrefixTree.lean b/src/Lean/Data/PrefixTree.lean
new file mode 100644
index 0000000000..e02b3011f8
--- /dev/null
+++ b/src/Lean/Data/PrefixTree.lean
@@ -0,0 +1,111 @@
+/-
+Copyright (c) 2020 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+import Std.Data.RBMap
+
+namespace Lean
+open Std
+
+/- Similar to trie, but for arbitrary keys -/
+inductive PrefixTreeNode (α : Type u) (β : Type v) where
+  | Node : Option β → RBNode α (fun _ => PrefixTreeNode α β) → PrefixTreeNode α β
+
+instance : Inhabited (PrefixTreeNode α β) where
+  default := PrefixTreeNode.Node none RBNode.leaf
+
+namespace PrefixTreeNode
+
+def empty : PrefixTreeNode α β :=
+  PrefixTreeNode.Node none RBNode.leaf
+
+@[specialize]
+partial def insert (t : PrefixTreeNode α β) (lt : α → α → Bool) (k : List α) (val : β) : PrefixTreeNode α β :=
+  let rec insertEmpty (k : List α) : PrefixTreeNode α β :=
+    match k with
+    | [] => PrefixTreeNode.Node (some val) RBNode.leaf
+    | k :: ks =>
+      let t := insertEmpty ks
+      PrefixTreeNode.Node none (RBNode.singleton k t)
+  let rec loop
+    | PrefixTreeNode.Node v m, [] =>
+      PrefixTreeNode.Node (some val) m -- overrides old value
+    | PrefixTreeNode.Node v m, k :: ks =>
+      let t := match RBNode.find lt m k with
+        | none   => insertEmpty ks
+        | some t => loop t ks
+      PrefixTreeNode.Node v (RBNode.insert lt m k t)
+  loop t k
+
+@[specialize]
+partial def find? (t : PrefixTreeNode α β) (lt : α → α → Bool) (k : List α) : Option β :=
+  let rec loop
+    | PrefixTreeNode.Node val m, [] => val
+    | PrefixTreeNode.Node val m, k :: ks =>
+      match RBNode.find lt m k with
+      | none   => none
+      | some t => loop t ks
+  loop t k
+
+@[specialize]
+partial def foldMatchingM [Monad m] (t : PrefixTreeNode α β) (lt : α → α → Bool) (k : List α) (init : σ) (f : β → σ → m σ) : m σ :=
+  let rec fold : PrefixTreeNode α β → σ → m σ
+    | PrefixTreeNode.Node b? n, d => do
+      let d ← match b? with
+        | none   => pure d
+        | some b => f b d
+      n.foldM (init := d) fun d _ t => fold t d
+  let rec find : List α → PrefixTreeNode α β → σ → m σ
+    | [],    t, d => fold t d
+    | k::ks, PrefixTreeNode.Node _ m, d =>
+      match RBNode.find lt m k with
+      | none   => pure init
+      | some t => find ks t d
+  find k t init
+
+inductive WellFormed (lt : α → α → Bool) : PrefixTreeNode α β → Prop where
+  | emptyWff    : WellFormed lt empty
+  | insertWff  {t : PrefixTreeNode α β} {k : List α} {val : β} : WellFormed lt t → WellFormed lt (insert t lt k val)
+
+end PrefixTreeNode
+
+def PrefixTree (α : Type u) (β : Type v) (lt : α → α → Bool) : Type (max u v) :=
+  { t : PrefixTreeNode α β // t.WellFormed lt }
+
+open PrefixTreeNode
+
+def PrefixTree.empty : PrefixTree α β p :=
+  ⟨PrefixTreeNode.empty, WellFormed.emptyWff⟩
+
+instance : Inhabited (PrefixTree α β p) where
+  default := PrefixTree.empty
+
+instance : EmptyCollection (PrefixTree α β p) where
+  emptyCollection := PrefixTree.empty
+
+@[inline]
+def PrefixTree.insert (t : PrefixTree α β p) (k : List α) (v : β) : PrefixTree α β p :=
+  ⟨t.val.insert p k v, WellFormed.insertWff t.property⟩
+
+@[inline]
+def PrefixTree.find? (t : PrefixTree α β p) (k : List α) : Option β :=
+  t.val.find? p k
+
+@[inline]
+def PrefixTree.foldMatchingM [Monad m] (t : PrefixTree α β p) (k : List α) (init : σ) (f : β → σ → m σ) : m σ :=
+  t.val.foldMatchingM p k init f
+
+@[inline]
+def PrefixTree.foldM [Monad m] (t : PrefixTree α β p) (init : σ) (f : β → σ → m σ) : m σ :=
+  t.foldMatchingM [] init f
+
+@[inline]
+def PrefixTree.forMatchingM [Monad m] (t : PrefixTree α β p) (k : List α) (f : β → m Unit) : m Unit :=
+  t.val.foldMatchingM p k () (fun b _ => f b)
+
+@[inline]
+def PrefixTree.forM [Monad m] (t : PrefixTree α β p) (f : β → m Unit) : m Unit :=
+  t.forMatchingM [] f
+
+end Lean