trie_me/
map.rs

1use std::borrow::Borrow;
2use std::collections::HashMap;
3use std::hash::Hash;
4use std::iter::FusedIterator;
5use std::marker::PhantomData;
6use std::ops::Index;
7use std::ops::IndexMut;
8
9pub struct TrieMap<K, V, A = u8> {
10	_key: PhantomData<K>,
11	val: Option<V>,
12	subtrie: HashMap<A, TrieMap<K, V, A>>,
13	len: usize,
14}
15
16impl<K, V, A: Clone + Eq + Hash> TrieMap<K, V, A> {
17	pub fn new() -> Self {
18		Self::default()
19	}
20
21	pub fn clear(&mut self) {
22		self.val.take();
23		self.subtrie.clear();
24		self.len = 0;
25	}
26
27	pub fn contains<T: AsRef<[A]> + ?Sized>(&self, key: &T) -> bool
28	where K: Borrow<T> {
29		self.get(key).is_some()
30	}
31
32	pub fn get<T: AsRef<[A]> + ?Sized>(&self, key: &T) -> Option<&V>
33	where K: Borrow<T> {
34		self.get_helper(key.as_ref())
35	}
36
37	pub fn get_mut<T: AsRef<[A]> + ?Sized>(&mut self, key: &T) -> Option<&mut V>
38	where K: Borrow<T> {
39		self.get_mut_helper(key.as_ref())
40	}
41
42	pub fn insert<T: AsRef<[A]> + ?Sized>(&mut self, key: &T, val: V) -> Option<V>
43	where K: Borrow<T> {
44		self.insert_helper(key.as_ref(), val)
45	}
46
47	pub fn is_empty(&self) -> bool {
48		self.len == 0
49	}
50
51	pub fn iter(&self) -> Iter<&V> {
52		self.iter_helper(&[])
53	}
54
55	pub fn iter_mut(&mut self) -> Iter<&mut V> {
56		self.iter_mut_helper(&[])
57	}
58
59	pub fn iter_prefix<T: AsRef<[A]> + ?Sized>(&self, prefix: &T) -> Iter<&V>
60	where K: Borrow<T> {
61		self.iter_helper(prefix.as_ref())
62	}
63
64	pub fn iter_prefix_mut<T: AsRef<[A]> + ?Sized>(&mut self, prefix: &T) -> Iter<&mut V>
65	where K: Borrow<T> {
66		self.iter_mut_helper(prefix.as_ref())
67	}
68
69	pub fn len(&self) -> usize {
70		self.len
71	}
72
73	fn get_helper(&self, key: &[A]) -> Option<&V> {
74		if let Some(atom) = key.first() {
75			self.subtrie.get(atom)?.get_helper(&key[1..])
76		} else {
77			self.val.as_ref()
78		}
79	}
80
81	fn get_mut_helper(&mut self, key: &[A]) -> Option<&mut V> {
82		if let Some(atom) = key.first() {
83			self.subtrie.get_mut(atom)?.get_mut_helper(&key[1..])
84		} else {
85			self.val.as_mut()
86		}
87	}
88
89	fn insert_helper(&mut self, key: &[A], val: V) -> Option<V> {
90		let existing = if let Some(atom) = key.first().cloned() {
91			self.subtrie.entry(atom).or_default().insert_helper(&key[1..], val)
92		} else {
93			self.val.replace(val)
94		};
95
96		if existing.is_none() {
97			self.len += 1;
98		}
99
100		existing
101	}
102
103	fn iter_helper(&self, key: &[A]) -> Iter<&V> {
104		use std::iter;
105
106		if let Some(atom) = key.first() {
107			if let Some(subtrie) = self.subtrie.get(atom) {
108				subtrie.iter_helper(&key[1..])
109			} else {
110				Box::new(iter::empty())
111			}
112		} else {
113			Box::new(self.val.iter().chain(self.subtrie.iter().flat_map(|(_, subtrie)|
114				subtrie.iter_helper(&[]))
115			))
116		}
117	}
118
119	fn iter_mut_helper(&mut self, key: &[A]) -> Iter<&mut V> {
120		use std::iter;
121
122		if let Some(atom) = key.first() {
123			if let Some(subtrie) = self.subtrie.get_mut(atom) {
124				subtrie.iter_mut_helper(&key[1..])
125			} else {
126				Box::new(iter::empty())
127			}
128		} else {
129			Box::new(self.val.iter_mut().chain(self.subtrie.iter_mut().flat_map(|(_, subtrie)|
130				subtrie.iter_mut_helper(&[]))
131			))
132		}
133	}
134}
135
136impl<K, V, A: Clone + Eq + Hash> Default for TrieMap<K, V, A> {
137	fn default() -> Self {
138		Self {
139			_key: PhantomData::default(),
140			val: None,
141			subtrie: HashMap::default(),
142			len: 0,
143		}
144	}
145}
146
147impl<A: Clone + Eq + Hash, T: AsRef<[A]> + ?Sized, K: Borrow<T>, V> Index<&T> for TrieMap<K, V, A> {
148	type Output = V;
149
150	fn index(&self, key: &T) -> &Self::Output {
151		self.get(key).unwrap()
152	}
153}
154
155impl<A: Clone + Eq + Hash, T: AsRef<[A]> + ?Sized, K: Borrow<T>, V> IndexMut<&T> for TrieMap<K, V, A> {
156	fn index_mut(&mut self, key: &T) -> &mut Self::Output {
157		self.get_mut(key).unwrap()
158	}
159}
160
161pub type Iter<'a, T> = Box<dyn FusedIterator<Item = T> + 'a>;
trie_me/map.rs

trie_me/
map.rs
