Write a custom map type that allows duplicates

6 files changed, 217 insertions, 39 deletions

diff --git a/dhall_syntax/Cargo.toml b/dhall_syntax/Cargo.toml
index 3e32930..b25a0fe 100644
--- a/dhall_syntax/Cargo.toml
+++ b/dhall_syntax/Cargo.toml
@@ -12,5 +12,7 @@ doctest = false
 itertools = "0.8.0"
 percent-encoding = "1.0.1"
 pest = "2.1"
+either = "1.5.2"
+take_mut = "0.2.2"
 dhall_generated_parser = { path = "../dhall_generated_parser" }
 improved_slice_patterns = { version = "2.0.0", path = "../improved_slice_patterns" }
diff --git a/dhall_syntax/src/core/expr.rs b/dhall_syntax/src/core/expr.rs
index c509bae..f8b0215 100644
--- a/dhall_syntax/src/core/expr.rs
+++ b/dhall_syntax/src/core/expr.rs
@@ -1,6 +1,7 @@
 #![allow(non_snake_case)]
 use std::rc::Rc;
 
+use crate::map::DupTreeMap;
 use crate::visitor;
 use crate::*;
 
@@ -201,13 +202,13 @@ pub enum ExprF<SubExpr, Embed> {
     ///  `Some e`
     SomeLit(SubExpr),
     ///  `{ k1 : t1, k2 : t1 }`
-    RecordType(Vec<(Label, SubExpr)>),
+    RecordType(DupTreeMap<Label, SubExpr>),
     ///  `{ k1 = v1, k2 = v2 }`
-    RecordLit(Vec<(Label, SubExpr)>),
+    RecordLit(DupTreeMap<Label, SubExpr>),
     ///  `< k1 : t1, k2 >`
-    UnionType(Vec<(Label, Option<SubExpr>)>),
+    UnionType(DupTreeMap<Label, Option<SubExpr>>),
     ///  `< k1 = t1, k2 : t2, k3 >`
-    UnionLit(Label, SubExpr, Vec<(Label, Option<SubExpr>)>),
+    UnionLit(Label, SubExpr, DupTreeMap<Label, Option<SubExpr>>),
     ///  `merge x y : t`
     Merge(SubExpr, SubExpr, Option<SubExpr>),
     ///  `e.x`
diff --git a/dhall_syntax/src/core/map.rs b/dhall_syntax/src/core/map.rs
new file mode 100644
index 0000000..c76ac34
--- /dev/null
+++ b/dhall_syntax/src/core/map.rs
@@ -0,0 +1,178 @@
+/// A sorted map that allows multiple values for each key.
+pub use dup_tree_map::DupTreeMap;
+
+mod one_or_more {
+    use either::Either;
+    use std::{iter, slice, vec};
+
+    #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+    pub enum OneOrMore<T> {
+        One(T),
+        More(Vec<T>),
+    }
+
+    pub type Iter<'a, T> = Either<slice::Iter<'a, T>, iter::Once<&'a T>>;
+    pub type IntoIter<T> = Either<vec::IntoIter<T>, iter::Once<T>>;
+
+    impl<T> OneOrMore<T> {
+        pub fn new(x: T) -> Self {
+            OneOrMore::One(x)
+        }
+
+        pub fn push(&mut self, x: T) {
+            take_mut::take(self, |sef| match sef {
+                OneOrMore::More(mut vec) => {
+                    vec.push(x);
+                    OneOrMore::More(vec)
+                }
+                OneOrMore::One(one) => OneOrMore::More(vec![one, x]),
+            })
+        }
+
+        pub fn iter(&self) -> Iter<'_, T> {
+            match self {
+                OneOrMore::More(vec) => Either::Left(vec.iter()),
+                OneOrMore::One(x) => Either::Right(iter::once(x)),
+            }
+        }
+    }
+
+    impl<T> IntoIterator for OneOrMore<T> {
+        type Item = T;
+        type IntoIter = IntoIter<T>;
+
+        fn into_iter(self) -> Self::IntoIter {
+            match self {
+                OneOrMore::More(vec) => Either::Left(vec.into_iter()),
+                OneOrMore::One(x) => Either::Right(iter::once(x)),
+            }
+        }
+    }
+}
+
+mod dup_tree_map {
+    use super::one_or_more;
+    use super::one_or_more::OneOrMore;
+    use std::collections::{btree_map, BTreeMap};
+    use std::iter;
+
+    #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+    pub struct DupTreeMap<K, V> {
+        map: BTreeMap<K, OneOrMore<V>>,
+        size: usize,
+    }
+
+    pub type IterInternal<'a, K, V> =
+        iter::Zip<iter::Repeat<&'a K>, one_or_more::Iter<'a, V>>;
+    pub type Iter<'a, K, V> = iter::FlatMap<
+        btree_map::Iter<'a, K, OneOrMore<V>>,
+        IterInternal<'a, K, V>,
+        for<'b> fn((&'b K, &'b OneOrMore<V>)) -> IterInternal<'b, K, V>,
+    >;
+    pub type IntoIterInternal<K, V> =
+        iter::Zip<iter::Repeat<K>, one_or_more::IntoIter<V>>;
+    pub type IntoIter<K, V> = iter::FlatMap<
+        btree_map::IntoIter<K, OneOrMore<V>>,
+        IntoIterInternal<K, V>,
+        fn((K, OneOrMore<V>)) -> IntoIterInternal<K, V>,
+    >;
+
+    impl<K, V> DupTreeMap<K, V> {
+        pub fn new() -> Self
+        where
+            K: Ord,
+        {
+            DupTreeMap {
+                map: BTreeMap::new(),
+                size: 0,
+            }
+        }
+
+        pub fn insert(&mut self, key: K, value: V)
+        where
+            K: Ord,
+        {
+            use std::collections::btree_map::Entry;
+            match self.map.entry(key) {
+                Entry::Vacant(e) => {
+                    e.insert(OneOrMore::new(value));
+                }
+                Entry::Occupied(mut e) => e.get_mut().push(value),
+            }
+        }
+
+        pub fn len(&self) -> usize {
+            self.size
+        }
+        pub fn is_empty(&self) -> bool {
+            self.size == 0
+        }
+
+        pub fn iter(&self) -> Iter<'_, K, V>
+        where
+            K: Ord,
+        {
+            fn foo<'a, K, V>(
+                (k, oom): (&'a K, &'a OneOrMore<V>),
+            ) -> IterInternal<'a, K, V> {
+                iter::repeat(k).zip(oom.iter())
+            }
+            self.map.iter().flat_map(foo)
+        }
+    }
+
+    impl<K, V> Default for DupTreeMap<K, V>
+    where
+        K: Ord,
+    {
+        fn default() -> Self {
+            Self::new()
+        }
+    }
+
+    impl<K, V> IntoIterator for DupTreeMap<K, V>
+    where
+        K: Ord + Clone,
+    {
+        type Item = (K, V);
+        type IntoIter = IntoIter<K, V>;
+
+        fn into_iter(self) -> Self::IntoIter {
+            fn foo<K, V>((k, oom): (K, OneOrMore<V>)) -> IntoIterInternal<K, V>
+            where
+                K: Clone,
+            {
+                iter::repeat(k).zip(oom.into_iter())
+            }
+            self.map.into_iter().flat_map(foo)
+        }
+    }
+
+    impl<'a, K, V> IntoIterator for &'a DupTreeMap<K, V>
+    where
+        K: Ord,
+    {
+        type Item = (&'a K, &'a V);
+        type IntoIter = Iter<'a, K, V>;
+
+        fn into_iter(self) -> Self::IntoIter {
+            self.iter()
+        }
+    }
+
+    impl<K, V> iter::FromIterator<(K, V)> for DupTreeMap<K, V>
+    where
+        K: Ord,
+    {
+        fn from_iter<T>(iter: T) -> Self
+        where
+            T: IntoIterator<Item = (K, V)>,
+        {
+            let mut map = DupTreeMap::new();
+            for (k, v) in iter {
+                map.insert(k, v);
+            }
+            map
+        }
+    }
+}
diff --git a/dhall_syntax/src/core/mod.rs b/dhall_syntax/src/core/mod.rs
index 7762479..fe2c0be 100644
--- a/dhall_syntax/src/core/mod.rs
+++ b/dhall_syntax/src/core/mod.rs
@@ -7,4 +7,5 @@ pub use label::*;
 mod text;
 pub use text::*;
 pub mod context;
+pub mod map;
 pub mod visitor;
diff --git a/dhall_syntax/src/core/visitor.rs b/dhall_syntax/src/core/visitor.rs
index dd539e3..99a9c11 100644
--- a/dhall_syntax/src/core/visitor.rs
+++ b/dhall_syntax/src/core/visitor.rs
@@ -1,4 +1,5 @@
 use crate::*;
+use std::iter::FromIterator;
 
 /// A way too generic Visitor trait.
 pub trait GenericVisitor<Input, Return>: Sized {
@@ -62,25 +63,29 @@ where
                 None => None,
             })
         }
-        fn vecmap<'a, V, Ret, SE, E>(
-            x: &'a Vec<(Label, SE)>,
+        fn dupmap<'a, V, Ret, SE, E, T>(
+            x: impl IntoIterator<Item = (&'a Label, &'a SE)>,
             mut v: V,
-        ) -> Result<Vec<(Label, V::SE2)>, V::Error>
+        ) -> Result<T, V::Error>
         where
+            SE: 'a,
+            T: FromIterator<(Label, V::SE2)>,
             V: ExprFVeryGenericVisitor<'a, Ret, SE, E>,
         {
-            x.iter()
+            x.into_iter()
                 .map(|(k, x)| Ok((k.clone(), v.visit_subexpr(x)?)))
                 .collect()
         }
-        fn vecoptmap<'a, V, Ret, SE, E>(
-            x: &'a Vec<(Label, Option<SE>)>,
+        fn optdupmap<'a, V, Ret, SE, E, T>(
+            x: impl IntoIterator<Item = (&'a Label, &'a Option<SE>)>,
             mut v: V,
-        ) -> Result<Vec<(Label, Option<V::SE2>)>, V::Error>
+        ) -> Result<T, V::Error>
         where
+            SE: 'a,
+            T: FromIterator<(Label, Option<V::SE2>)>,
             V: ExprFVeryGenericVisitor<'a, Ret, SE, E>,
         {
-            x.iter()
+            x.into_iter()
                 .map(|(k, x)| {
                     Ok((
                         k.clone(),
@@ -137,11 +142,11 @@ where
                 v.visit_subexpr(t)?,
             ),
             SomeLit(e) => SomeLit(v.visit_subexpr(e)?),
-            RecordType(kts) => RecordType(vecmap(kts, v)?),
-            RecordLit(kvs) => RecordLit(vecmap(kvs, v)?),
-            UnionType(kts) => UnionType(vecoptmap(kts, v)?),
+            RecordType(kts) => RecordType(dupmap(kts, v)?),
+            RecordLit(kvs) => RecordLit(dupmap(kvs, v)?),
+            UnionType(kts) => UnionType(optdupmap(kts, v)?),
             UnionLit(k, x, kts) => {
-                UnionLit(k.clone(), v.visit_subexpr(x)?, vecoptmap(kts, v)?)
+                UnionLit(k.clone(), v.visit_subexpr(x)?, optdupmap(kts, v)?)
             }
             Merge(x, y, t) => Merge(
                 v.visit_subexpr(x)?,
diff --git a/dhall_syntax/src/parser.rs b/dhall_syntax/src/parser.rs
index 276510e..c847b29 100644
--- a/dhall_syntax/src/parser.rs
+++ b/dhall_syntax/src/parser.rs
@@ -7,6 +7,7 @@ use std::rc::Rc;
 
 use dhall_generated_parser::{DhallParser, Rule};
 
+use crate::map::DupTreeMap;
 use crate::ExprF::*;
 use crate::*;
 
@@ -886,33 +887,29 @@ make_parser! {
     ));
 
     rule!(empty_record_literal<ParsedSubExpr> as expression; span;
-        captured_str!(_) => spanned(span, RecordLit(Vec::new()))
+        captured_str!(_) => spanned(span, RecordLit(Default::default()))
     );
 
     rule!(empty_record_type<ParsedSubExpr> as expression; span;
-        captured_str!(_) => spanned(span, RecordType(Vec::new()))
+        captured_str!(_) => spanned(span, RecordType(Default::default()))
     );
 
     rule!(non_empty_record_type_or_literal<ParsedSubExpr> as expression; span;
           children!(
         [label(first_label), non_empty_record_type(rest)] => {
             let (first_expr, mut map) = rest;
-            map.push((first_label, first_expr));
-            // Sort until we stop using binary decode for parser tests
-            map.sort_by(|(l1, _), (l2, _)| l1.cmp(l2));
+            map.insert(first_label, first_expr);
             spanned(span, RecordType(map))
         },
         [label(first_label), non_empty_record_literal(rest)] => {
             let (first_expr, mut map) = rest;
-            map.push((first_label, first_expr));
-            // Sort until we stop using binary decode for parser tests
-            map.sort_by(|(l1, _), (l2, _)| l1.cmp(l2));
+            map.insert(first_label, first_expr);
             spanned(span, RecordLit(map))
         },
     ));
 
     rule!(non_empty_record_type
-          <(ParsedSubExpr, Vec<(Label, ParsedSubExpr)>)>; children!(
+          <(ParsedSubExpr, DupTreeMap<Label, ParsedSubExpr>)>; children!(
         [expression(expr), record_type_entry(entries)..] => {
             (expr, entries.collect())
         }
@@ -923,7 +920,7 @@ make_parser! {
     ));
 
     rule!(non_empty_record_literal
-          <(ParsedSubExpr, Vec<(Label, ParsedSubExpr)>)>; children!(
+          <(ParsedSubExpr, DupTreeMap<Label, ParsedSubExpr>)>; children!(
         [expression(expr), record_literal_entry(entries)..] => {
             (expr, entries.collect())
         }
@@ -935,18 +932,12 @@ make_parser! {
 
     rule!(union_type_or_literal<ParsedSubExpr> as expression; span; children!(
         [empty_union_type(_)] => {
-            spanned(span, UnionType(Vec::new()))
+            spanned(span, UnionType(Default::default()))
         },
         [non_empty_union_type_or_literal((Some((l, e)), entries))] => {
-            let mut entries = entries;
-            // Sort until we stop using binary decode for parser tests
-            entries.sort_by(|(l1, _), (l2, _)| l1.cmp(l2));
             spanned(span, UnionLit(l, e, entries))
         },
         [non_empty_union_type_or_literal((None, entries))] => {
-            let mut entries = entries;
-            // Sort until we stop using binary decode for parser tests
-            entries.sort_by(|(l1, _), (l2, _)| l1.cmp(l2));
             spanned(span, UnionType(entries))
         },
     ));
@@ -955,20 +946,20 @@ make_parser! {
 
     rule!(non_empty_union_type_or_literal
           <(Option<(Label, ParsedSubExpr)>,
-            Vec<(Label, Option<ParsedSubExpr>)>)>;
+            DupTreeMap<Label, Option<ParsedSubExpr>>)>;
             children!(
         [label(l), union_literal_variant_value((e, entries))] => {
             (Some((l, e)), entries)
         },
         [label(l), union_type_or_literal_variant_type((e, rest))] => {
             let (x, mut entries) = rest;
-            entries.push((l, e));
+            entries.insert(l, e);
             (x, entries)
         },
     ));
 
     rule!(union_literal_variant_value
-          <(ParsedSubExpr, Vec<(Label, Option<ParsedSubExpr>)>)>;
+          <(ParsedSubExpr, DupTreeMap<Label, Option<ParsedSubExpr>>)>;
             children!(
         [expression(e), union_type_entry(entries)..] => {
             (e, entries.collect())
@@ -984,19 +975,19 @@ make_parser! {
     rule!(union_type_or_literal_variant_type
           <(Option<ParsedSubExpr>,
             (Option<(Label, ParsedSubExpr)>,
-             Vec<(Label, Option<ParsedSubExpr>)>))>;
+             DupTreeMap<Label, Option<ParsedSubExpr>>))>;
                 children!(
         [expression(e), non_empty_union_type_or_literal(rest)] => {
             (Some(e), rest)
         },
         [expression(e)] => {
-            (Some(e), (None, Vec::new()))
+            (Some(e), (None, Default::default()))
         },
         [non_empty_union_type_or_literal(rest)] => {
             (None, rest)
         },
         [] => {
-            (None, (None, Vec::new()))
+            (None, (None, Default::default()))
         },
     ));