diff options
author | Nadrieril | 2019-05-10 19:49:38 +0200 |
---|---|---|
committer | Nadrieril | 2019-05-10 19:49:38 +0200 |
commit | a8e696f62f14296b94964adb1946d7e2b5ef5ebd (patch) | |
tree | 8ac9c17c4e0cf7470b047c46ac7042d1db1bb1e9 /dhall_syntax/src/core/map.rs | |
parent | 36bcec6c91d3192b5c84c96af96961ff6b79f0f0 (diff) |
Write a custom map type that allows duplicates
Diffstat (limited to 'dhall_syntax/src/core/map.rs')
-rw-r--r-- | dhall_syntax/src/core/map.rs | 178 |
1 files changed, 178 insertions, 0 deletions
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 + } + } +} |