From d4105e1e16d3398793b893a7776dd279abe2228f Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Wed, 5 Sep 2018 10:35:57 -0700
Subject: [PATCH] chore(library/init/lean/parser/basic): avoid nesting `rbnode`
 in `trie`

The new two-layer rbnode/list structure is a little faster
---
 library/init/data/list/basic.lean   |  4 +++
 library/init/lean/parser/basic.lean | 44 ++++++++++++++++++-----------
 2 files changed, 32 insertions(+), 16 deletions(-)

diff --git a/library/init/data/list/basic.lean b/library/init/data/list/basic.lean
index b7e6becca1..09d5e7e9c2 100644
--- a/library/init/data/list/basic.lean
+++ b/library/init/data/list/basic.lean
@@ -155,6 +155,10 @@ def find_index (p : α → Prop) [decidable_pred p] : list α → nat
 
 def index_of [decidable_eq α] (a : α) : list α → nat := find_index (eq a)
 
+def assoc [decidable_eq α] : list (α × β) → α → option β
+| []         _  := none
+| ((a,b)::l) a' := if a = a' then some b else assoc l a'
+
 def remove_all [decidable_eq α] (xs ys : list α) : list α :=
 filter (∉ ys) xs
 
diff --git a/library/init/lean/parser/basic.lean b/library/init/lean/parser/basic.lean
index 1bfbb4e170..ef3b41f8c9 100644
--- a/library/init/lean/parser/basic.lean
+++ b/library/init/lean/parser/basic.lean
@@ -59,33 +59,45 @@ abbreviation message := string
 
 namespace parser
 
-inductive trie (α : Type)
-| node : option α → rbmap char trie (<) → trie
+inductive trie.node (α : Type)
+| mk : option α → list (char × trie.node) → trie.node
+
+def trie (α : Type) := rbmap char (trie.node α) (<)
 
 namespace trie
 variables {α : Type}
 
-protected def mk : trie α := ⟨none, mk_rbmap _ _ _⟩
+protected def node.empty : trie.node α := ⟨none, []⟩
 
-private def insert_aux (val : α) : nat → trie α → string.iterator → trie α
-| 0 (trie.node _ map) _ := trie.node (some val) map   -- NOTE: overrides old value
-| (n+1) (trie.node val map) it :=
-  let t' := (map.find it.curr).get_or_else trie.mk in
-  trie.node val (map.insert it.curr (insert_aux n t' it.next))
+protected def mk : trie α := mk_rbmap _ _ _
+
+private def update_child (c : char) (f : trie.node α → trie.node α) : list (char × trie.node α) → list (char × trie.node α)
+| []           := [(c, f trie.node.empty)]
+| ((c',t)::ts) := if c = c' then (c', f t)::ts else (c',t)::update_child ts
+
+private def insert_aux (val : α) : nat → trie.node α → string.iterator → trie.node α
+| 0     ⟨_,   ts⟩ _  := ⟨some val, ts⟩   -- NOTE: overrides old value
+| (n+1) ⟨val, ts⟩ it :=
+  ⟨val, update_child it.curr (λ t, insert_aux n t it.next) ts⟩
 
 def insert (t : trie α) (s : string) (val : α) : trie α :=
-insert_aux val s.length t s.mk_iterator
+let it := s.mk_iterator in
+let t' := (t.find it.curr).get_or_else trie.node.empty in
+let it' := it.next in
+t.insert it.curr (insert_aux val it'.remaining t' it')
 
-private def match_prefix_aux : nat → trie α → string.iterator → option (string.iterator × α) → option (string.iterator × α)
-| 0 (trie.node val map) it acc := prod.mk it <$> val <|> acc
-| (n+1) (trie.node val map) it acc :=
+private def match_prefix_aux : nat → trie.node α → string.iterator → option (string.iterator × α) → option (string.iterator × α)
+| 0 ⟨val, ts⟩ it acc := prod.mk it <$> val <|> acc
+| (n+1) ⟨val, ts⟩ it acc :=
   let acc' := prod.mk it <$> val <|> acc in
-  match map.find it.curr with
-    | some t := match_prefix_aux n t it.next acc'
-    | none   := acc'
+  match ts.assoc it.curr with
+  | some t := match_prefix_aux n t it.next acc'
+  | none   := acc'
 
 def match_prefix {α : Type} (t : trie α) (it : string.iterator) : option (string.iterator × α) :=
-match_prefix_aux it.remaining t it none
+do t' ← t.find it.curr,
+   let it' := it.next,
+   match_prefix_aux it'.remaining t' it' none
 end trie