1use std::{
6 borrow::Borrow,
7 cmp::Ordering,
8 collections::{btree_map, BTreeMap},
9 fmt::{self, Debug, Formatter},
10 hash::{Hash, Hasher},
11 iter::FusedIterator,
12 mem,
13 ops::{Deref, DerefMut, Index},
14};
15
16use crate::{Commonality, DefaultCommonality, PhantomPtr};
17
18pub struct TotalBTreeMap<K, V, C = DefaultCommonality> {
30 inner: BTreeMap<K, V>,
31 common: V, _commonality: PhantomPtr<C>,
33}
34
35impl<K: Clone, V: Clone, C> Clone for TotalBTreeMap<K, V, C> {
36 fn clone(&self) -> Self {
37 Self {
38 inner: self.inner.clone(),
39 common: self.common.clone(),
40 _commonality: PhantomPtr::default(),
41 }
42 }
43}
44
45impl<K, V, C: Commonality<V>> Default for TotalBTreeMap<K, V, C> {
46 fn default() -> Self {
47 Self {
48 inner: BTreeMap::default(),
49 common: C::common(),
50 _commonality: PhantomPtr::default(),
51 }
52 }
53}
54impl<K, V, C: Commonality<V>> TotalBTreeMap<K, V, C> {
55 pub fn new() -> Self {
57 Self::default()
58 }
59}
60
61impl<K, V, C> TotalBTreeMap<K, V, C> {
62 pub fn len(&self) -> usize {
64 self.inner.len()
65 }
66 pub fn is_empty(&self) -> bool {
68 self.inner.is_empty()
69 }
70 pub fn clear(&mut self) {
72 self.inner.clear()
73 }
74}
75
76impl<K: Ord, V, C> TotalBTreeMap<K, V, C> {
80 pub fn get<Q>(&self, key: &Q) -> &V
82 where
83 K: Borrow<Q>,
84 Q: Ord + ?Sized,
85 {
86 self.inner.get(key).unwrap_or(&self.common)
87 }
88 pub fn contains_key<Q>(&self, key: &Q) -> bool
90 where
91 K: Borrow<Q>,
92 Q: Ord + ?Sized,
93 {
94 self.inner.contains_key(key)
95 }
96}
97
98impl<K: Borrow<Q> + Ord, Q: Ord + ?Sized, V, C> Index<&Q> for TotalBTreeMap<K, V, C> {
99 type Output = V;
100 fn index(&self, index: &Q) -> &Self::Output {
101 self.get(index)
102 }
103}
104
105impl<K: Ord, V, C: Commonality<V>> TotalBTreeMap<K, V, C> {
106 pub fn insert(&mut self, key: K, value: V) -> V {
109 if C::is_common(&value) { self.inner.remove(&key) } else { self.inner.insert(key, value) }
110 .unwrap_or_else(C::common)
111 }
112
113 pub fn remove<Q>(&mut self, key: &Q) -> V
116 where
117 K: Borrow<Q>,
118 Q: Ord + ?Sized,
119 {
120 self.inner.remove(key).unwrap_or_else(C::common)
121 }
122
123 pub fn entry(&mut self, key: K) -> Entry<'_, K, K, V, C> {
125 Entry {
126 inner: match self.inner.entry(key) {
127 btree_map::Entry::Occupied(inner) => EntryInner::Occupied { inner },
128 btree_map::Entry::Vacant(inner) => EntryInner::Vacant { inner, value: C::common() },
129 },
130 }
131 }
132
133 pub fn uncommon_entry<'a, Q>(&'a mut self, key: &'a Q) -> Option<Entry<'a, Q, K, V, C>>
138 where
139 K: Borrow<Q>,
140 Q: Ord + ?Sized,
141 {
142 let map = self as *mut _;
143 let value = self.inner.get_mut(key)?;
144 Some(Entry { inner: EntryInner::ByRef { map, key, value } })
145 }
146}
147
148pub struct Entry<'a, Q, K, V, C = DefaultCommonality>
152where
153 Q: Ord + ?Sized,
154 K: Ord + Borrow<Q>,
155 C: Commonality<V>,
156{
157 inner: EntryInner<'a, Q, K, V, C>,
158}
159
160impl<Q, K, V, C> Deref for Entry<'_, Q, K, V, C>
161where
162 Q: Ord + ?Sized,
163 K: Ord + Borrow<Q>,
164 C: Commonality<V>,
165{
166 type Target = V;
167 fn deref(&self) -> &Self::Target {
168 match &self.inner {
169 EntryInner::Occupied { inner } => inner.get(),
170 EntryInner::Vacant { value, .. } => value,
171 EntryInner::ByRef { value, .. } => value,
172 EntryInner::Dropping => unreachable!(),
173 }
174 }
175}
176impl<Q, K, V, C> DerefMut for Entry<'_, Q, K, V, C>
177where
178 Q: Ord + ?Sized,
179 K: Ord + Borrow<Q>,
180 C: Commonality<V>,
181{
182 fn deref_mut(&mut self) -> &mut Self::Target {
183 match &mut self.inner {
184 EntryInner::Occupied { inner } => inner.get_mut(),
185 EntryInner::Vacant { value, .. } => value,
186 EntryInner::ByRef { value, .. } => value,
187 EntryInner::Dropping => unreachable!(),
188 }
189 }
190}
191
192impl<Q, K, V, C> Drop for Entry<'_, Q, K, V, C>
193where
194 Q: Ord + ?Sized,
195 K: Ord + Borrow<Q>,
196 C: Commonality<V>,
197{
198 fn drop(&mut self) {
199 match mem::replace(&mut self.inner, EntryInner::Dropping) {
200 EntryInner::Occupied { inner } => {
201 if C::is_common(inner.get()) {
202 inner.remove();
203 }
204 }
205 EntryInner::Vacant { inner, value } => {
206 if !C::is_common(&value) {
207 inner.insert(value);
208 }
209 }
210 EntryInner::ByRef { map, key, value } => {
211 if C::is_common(value) {
212 unsafe { &mut *map }.remove(key);
213 }
214 }
215 EntryInner::Dropping => unreachable!(),
216 }
217 }
218}
219
220impl<'a, Q, K, V, C> Debug for Entry<'a, Q, K, V, C>
221where
222 Q: Debug + Ord + ?Sized,
223 K: Debug + Ord + Borrow<Q>,
224 V: Debug,
225 C: Commonality<V>,
226{
227 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
228 let mut f = f.debug_tuple("Entry");
229 match &self.inner {
230 EntryInner::Occupied { inner } => f.field(inner.key()).field(inner.get()),
231 EntryInner::Vacant { inner, value } => f.field(inner.key()).field(value),
232 EntryInner::ByRef { key, value, .. } => f.field(key).field(value),
233 EntryInner::Dropping => &mut f,
234 };
235 f.finish()
236 }
237}
238
239enum EntryInner<'a, Q: ?Sized, K, V, C> {
240 Occupied { inner: btree_map::OccupiedEntry<'a, K, V> },
241 Vacant { inner: btree_map::VacantEntry<'a, K, V>, value: V },
242 ByRef { map: *mut TotalBTreeMap<K, V, C>, key: &'a Q, value: &'a mut V },
243 Dropping,
244}
245
246impl<K, V, C> TotalBTreeMap<K, V, C> {
250 pub fn keys(&self) -> Keys<'_, K, V> {
252 Keys(self.inner.keys())
253 }
254 pub fn into_keys(self) -> IntoKeys<K, V> {
257 IntoKeys(self.inner.into_keys())
258 }
259 pub fn values(&self) -> Values<'_, K, V> {
261 Values(self.inner.values())
262 }
263 pub fn into_values(self) -> IntoValues<K, V> {
265 IntoValues(self.inner.into_values())
266 }
267 pub fn iter(&self) -> Iter<'_, K, V> {
269 Iter(self.inner.iter())
270 }
271}
272
273impl<K, V, C> IntoIterator for TotalBTreeMap<K, V, C> {
274 type Item = (K, V);
275 type IntoIter = IntoIter<K, V>;
276 fn into_iter(self) -> Self::IntoIter {
277 IntoIter(self.inner.into_iter())
278 }
279}
280impl<'a, K, V, C> IntoIterator for &'a TotalBTreeMap<K, V, C> {
281 type Item = (&'a K, &'a V);
282 type IntoIter = Iter<'a, K, V>;
283 fn into_iter(self) -> Self::IntoIter {
284 self.iter()
285 }
286}
287
288pub struct Keys<'a, K, V>(btree_map::Keys<'a, K, V>);
292impl<K, V> Clone for Keys<'_, K, V> {
293 fn clone(&self) -> Self {
294 Self(self.0.clone())
295 }
296}
297impl<'a, K, V> Default for Keys<'a, K, V> {
298 fn default() -> Self {
299 Self(Default::default())
300 }
301}
302impl<'a, K, V> Iterator for Keys<'a, K, V> {
303 type Item = &'a K;
304 fn next(&mut self) -> Option<Self::Item> {
305 self.0.next()
306 }
307 fn size_hint(&self) -> (usize, Option<usize>) {
308 self.0.size_hint()
309 }
310}
311impl<K, V> DoubleEndedIterator for Keys<'_, K, V> {
312 fn next_back(&mut self) -> Option<Self::Item> {
313 self.0.next_back()
314 }
315}
316impl<K, V> ExactSizeIterator for Keys<'_, K, V> {
317 fn len(&self) -> usize {
318 self.0.len()
319 }
320}
321impl<K, V> FusedIterator for Keys<'_, K, V> {}
322impl<K: Debug, V: Debug> Debug for Keys<'_, K, V> {
323 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
324 f.debug_list().entries(self.clone()).finish()
325 }
326}
327
328pub struct IntoKeys<K, V>(btree_map::IntoKeys<K, V>);
332impl<K, V> Default for IntoKeys<K, V> {
333 fn default() -> Self {
334 Self(Default::default())
335 }
336}
337impl<K, V> Iterator for IntoKeys<K, V> {
338 type Item = K;
339 fn next(&mut self) -> Option<Self::Item> {
340 self.0.next()
341 }
342 fn size_hint(&self) -> (usize, Option<usize>) {
343 self.0.size_hint()
344 }
345}
346impl<K, V> DoubleEndedIterator for IntoKeys<K, V> {
347 fn next_back(&mut self) -> Option<Self::Item> {
348 self.0.next_back()
349 }
350}
351impl<K, V> ExactSizeIterator for IntoKeys<K, V> {
352 fn len(&self) -> usize {
353 self.0.len()
354 }
355}
356impl<K, V> FusedIterator for IntoKeys<K, V> {}
357
358pub struct Values<'a, K, V>(btree_map::Values<'a, K, V>);
362impl<K, V> Clone for Values<'_, K, V> {
363 fn clone(&self) -> Self {
364 Self(self.0.clone())
365 }
366}
367impl<'a, K, V> Default for Values<'a, K, V> {
368 fn default() -> Self {
369 Self(Default::default())
370 }
371}
372impl<'a, K, V> Iterator for Values<'a, K, V> {
373 type Item = &'a V;
374 fn next(&mut self) -> Option<Self::Item> {
375 self.0.next()
376 }
377 fn size_hint(&self) -> (usize, Option<usize>) {
378 self.0.size_hint()
379 }
380}
381impl<'a, K, V> DoubleEndedIterator for Values<'a, K, V> {
382 fn next_back(&mut self) -> Option<Self::Item> {
383 self.0.next_back()
384 }
385}
386impl<K, V> ExactSizeIterator for Values<'_, K, V> {
387 fn len(&self) -> usize {
388 self.0.len()
389 }
390}
391impl<K, V> FusedIterator for Values<'_, K, V> {}
392impl<K: Debug, V: Debug> Debug for Values<'_, K, V> {
393 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
394 f.debug_list().entries(self.clone()).finish()
395 }
396}
397
398pub struct IntoValues<K, V>(btree_map::IntoValues<K, V>);
402impl<K, V> Default for IntoValues<K, V> {
403 fn default() -> Self {
404 Self(Default::default())
405 }
406}
407impl<K, V> Iterator for IntoValues<K, V> {
408 type Item = V;
409 fn next(&mut self) -> Option<Self::Item> {
410 self.0.next()
411 }
412 fn size_hint(&self) -> (usize, Option<usize>) {
413 self.0.size_hint()
414 }
415}
416impl<K, V> DoubleEndedIterator for IntoValues<K, V> {
417 fn next_back(&mut self) -> Option<Self::Item> {
418 self.0.next_back()
419 }
420}
421impl<K, V> ExactSizeIterator for IntoValues<K, V> {
422 fn len(&self) -> usize {
423 self.0.len()
424 }
425}
426impl<K, V> FusedIterator for IntoValues<K, V> {}
427
428pub struct Iter<'a, K, V>(btree_map::Iter<'a, K, V>);
432impl<K, V> Clone for Iter<'_, K, V> {
433 fn clone(&self) -> Self {
434 Self(self.0.clone())
435 }
436}
437impl<'a, K, V> Default for Iter<'a, K, V> {
438 fn default() -> Self {
439 Self(Default::default())
440 }
441}
442impl<'a, K, V> Iterator for Iter<'a, K, V> {
443 type Item = (&'a K, &'a V);
444 fn next(&mut self) -> Option<Self::Item> {
445 self.0.next()
446 }
447 fn size_hint(&self) -> (usize, Option<usize>) {
448 self.0.size_hint()
449 }
450}
451impl<'a, K, V> DoubleEndedIterator for Iter<'a, K, V> {
452 fn next_back(&mut self) -> Option<Self::Item> {
453 self.0.next_back()
454 }
455}
456impl<K, V> ExactSizeIterator for Iter<'_, K, V> {
457 fn len(&self) -> usize {
458 self.0.len()
459 }
460}
461impl<K, V> FusedIterator for Iter<'_, K, V> {}
462impl<K: Debug, V: Debug> Debug for Iter<'_, K, V> {
463 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
464 f.debug_list().entries(self.clone()).finish()
465 }
466}
467
468pub struct IntoIter<K, V>(btree_map::IntoIter<K, V>);
472impl<K, V> Default for IntoIter<K, V> {
473 fn default() -> Self {
474 Self(Default::default())
475 }
476}
477impl<K, V> Iterator for IntoIter<K, V> {
478 type Item = (K, V);
479 fn next(&mut self) -> Option<Self::Item> {
480 self.0.next()
481 }
482 fn size_hint(&self) -> (usize, Option<usize>) {
483 self.0.size_hint()
484 }
485}
486impl<K, V> DoubleEndedIterator for IntoIter<K, V> {
487 fn next_back(&mut self) -> Option<Self::Item> {
488 self.0.next_back()
489 }
490}
491impl<K, V> ExactSizeIterator for IntoIter<K, V> {
492 fn len(&self) -> usize {
493 self.0.len()
494 }
495}
496impl<K, V> FusedIterator for IntoIter<K, V> {}
497
498impl<K: Ord, V, C: Commonality<V>> Extend<(K, V)> for TotalBTreeMap<K, V, C> {
502 fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
503 for (key, value) in iter {
504 self.insert(key, value);
505 }
506 }
507}
508impl<K: Ord, V, C: Commonality<V>> FromIterator<(K, V)> for TotalBTreeMap<K, V, C> {
509 fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
510 let mut this = Self::default();
511 this.extend(iter);
512 this
513 }
514}
515
516impl<K, V, C> TotalBTreeMap<K, V, C> {
520 pub fn as_btree_map(&self) -> &BTreeMap<K, V> {
523 &self.inner
524 }
525}
526
527impl<K: Ord, V, C: Commonality<V>> TotalBTreeMap<K, V, C> {
528 pub fn as_btree_map_mut(&mut self) -> AsBTreeMapMut<'_, K, V, C> {
539 AsBTreeMapMut { map: &mut self.inner, _commonality: PhantomPtr::default() }
540 }
541}
542
543pub struct AsBTreeMapMut<'a, K: Ord, V, C: Commonality<V> = DefaultCommonality> {
547 map: &'a mut BTreeMap<K, V>,
548 _commonality: PhantomPtr<C>,
549}
550
551impl<K: Ord, V, C: Commonality<V>> Deref for AsBTreeMapMut<'_, K, V, C> {
552 type Target = BTreeMap<K, V>;
553 fn deref(&self) -> &Self::Target {
554 self.map
555 }
556}
557impl<K: Ord, V, C: Commonality<V>> DerefMut for AsBTreeMapMut<'_, K, V, C> {
558 fn deref_mut(&mut self) -> &mut Self::Target {
559 self.map
560 }
561}
562
563impl<K: Ord, V, C: Commonality<V>> Drop for AsBTreeMapMut<'_, K, V, C> {
564 fn drop(&mut self) {
565 self.map.retain(|_, value| !C::is_common(value));
566 }
567}
568
569impl<K: Ord, V: PartialEq, C: Commonality<V>> PartialEq for AsBTreeMapMut<'_, K, V, C> {
570 fn eq(&self, other: &Self) -> bool {
571 self.map == other.map
572 }
574}
575impl<K: Ord, V: Eq, C: Commonality<V>> Eq for AsBTreeMapMut<'_, K, V, C> {}
576impl<K: Ord, V: PartialOrd, C: Commonality<V>> PartialOrd for AsBTreeMapMut<'_, K, V, C> {
577 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
578 self.map.partial_cmp(&other.map)
579 }
581}
582impl<K: Ord, V: Ord, C: Commonality<V>> Ord for AsBTreeMapMut<'_, K, V, C> {
583 fn cmp(&self, other: &Self) -> Ordering {
584 self.map.cmp(&other.map)
585 }
587}
588impl<K: Ord + Hash, V: Hash, C: Commonality<V>> Hash for AsBTreeMapMut<'_, K, V, C> {
589 fn hash<H: Hasher>(&self, state: &mut H) {
590 self.map.hash(state);
591 }
593}
594impl<K: Ord + Debug, V: Debug, C: Commonality<V>> Debug for AsBTreeMapMut<'_, K, V, C> {
595 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
596 f.debug_tuple("AsBTreeMapMut").field(&self.map).finish()
597 }
598}
599
600impl<K: PartialEq, V: PartialEq, C> PartialEq for TotalBTreeMap<K, V, C> {
604 fn eq(&self, other: &Self) -> bool {
605 self.common == other.common && self.inner == other.inner
607 }
608}
609impl<K: Eq, V: Eq, C> Eq for TotalBTreeMap<K, V, C> {}
610
611impl<K: PartialOrd, V: PartialOrd, C> PartialOrd for TotalBTreeMap<K, V, C> {
612 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
613 match self.common.partial_cmp(&other.common) {
615 Some(Ordering::Equal) => {}
616 ord => return ord,
617 };
618 self.inner.partial_cmp(&other.inner)
619 }
620}
621impl<K: Ord, V: Ord, C> Ord for TotalBTreeMap<K, V, C> {
622 fn cmp(&self, other: &Self) -> Ordering {
623 self.common.cmp(&other.common).then_with(|| self.inner.cmp(&other.inner))
625 }
626}
627
628impl<K: Hash, V: Hash, C> Hash for TotalBTreeMap<K, V, C> {
629 fn hash<H: Hasher>(&self, state: &mut H) {
630 self.common.hash(state);
632 self.inner.hash(state);
633 }
634}
635
636impl<K: Debug, V: Debug, C> Debug for TotalBTreeMap<K, V, C> {
637 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
638 struct Rest;
639 impl Debug for Rest {
640 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
641 write!(f, "...")
642 }
643 }
644 f.debug_map().entries(self.iter()).entry(&Rest, &self.common).finish()
645 }
646}