lz_diet/
lib.rs

1//! This crate provides a Discrete Interval Encoding Tree (DIET) implementation.
2//!
3//! At a small performance cost on mutation of the tree, large memory savings
4//! are made and performance gains may be made on lookup of elements.
5//!
6//! An example of using the `Diet<T>`:
7//!
8//! ```
9//! use lz_diet::{Diet, Interval};
10//!
11//! let mut diet = Diet::new();
12//! diet.insert(5);
13//! diet.insert(7);
14//!
15//! diet.insert(6);
16//!
17//! let intervals: Vec<_> = diet.into_iter().collect();
18//! assert_eq!(&intervals[..], &[Interval::from(5..8)]);
19//! ```
20//!
21//! Interval endpoints are evaluated to determine whether the interval is
22//! discrete, this check implemented with the `AdjacentBound` trait and may be
23//! quickly implemented for types with a scalar interval using the
24//! `adjacent_bound_impl` macro provided.
25//!
26
27use log::debug;
28use matches::matches;
29
30mod interval;
31mod node_mut_ext;
32pub use crate::interval::Interval;
33
34mod adjacent_bound;
35pub use crate::adjacent_bound::AdjacentBound;
36
37mod walk_direction;
38pub(crate) use crate::walk_direction::WalkDirection;
39
40mod split_result;
41pub(crate) use crate::split_result::SplitResult;
42
43mod diet_node;
44pub use crate::diet_node::{DietNode, DietNodePtr};
45
46mod iterators;
47pub use crate::iterators::{IntoIter, Iter};
48
49use binary_tree::iter::IntoIter as GenIntoIter;
50use binary_tree::{BinaryTree, Node, NodeMut, WalkAction};
51use std::borrow::{Borrow, Cow};
52use std::hash::{Hash, Hasher};
53use std::iter::FromIterator;
54
55/// An AVL balanced Discrete Interval Encoding Tree where each node represents
56/// a discrete (non-touching) interval.
57#[derive(Debug, Clone, Eq)]
58pub struct Diet<T> {
59    root: Option<Box<DietNode<T>>>,
60}
61
62impl<T> Drop for Diet<T> {
63    fn drop(&mut self) {
64        self.clear();
65    }
66}
67
68impl<T> Diet<T> {
69    /// Creates a new `Diet<T>`.
70    ///
71    /// ```
72    /// use lz_diet::Diet;
73    ///
74    /// let diet : Diet<u32> = Diet::new();
75    /// ```
76    pub fn new() -> Self {
77        Self { root: None }
78    }
79
80    /// Iterate over a borrow of the `Interval<T>` values.
81    ///
82    /// ```
83    /// use lz_diet::Diet;
84    ///
85    /// let mut diet = Diet::new();
86    /// diet.insert(5u32);
87    ///
88    /// let intervals : Vec<_> = diet.iter().collect();
89    /// assert_eq!(intervals, vec![&(5..6).into()]);
90    /// ```
91    pub fn iter(&self) -> Iter<T> {
92        self.into_iter()
93    }
94
95    /// Clears the `Diet<T>`.
96    ///
97    /// ```
98    /// use lz_diet::Diet;
99    ///
100    /// let mut diet = Diet::new();
101    /// diet.insert(7u32);
102    /// assert_eq!(diet.len(), 1);
103    /// assert_eq!(diet.is_empty(), false);
104    ///
105    /// diet.clear();
106    /// assert_eq!(diet.len(), 0);
107    /// assert_eq!(diet.is_empty(), true);
108    /// ```
109    pub fn clear(&mut self) {
110        debug!("clearing Diet");
111
112        {
113            // The iterator ensures we don't get a stackoverflow for a large tree
114            // as its drop implementation iterates and drops each node individually
115            let _: GenIntoIter<DietNode<_>> = GenIntoIter::new(self.root.take());
116        }
117
118        debug!("cleared Diet");
119    }
120
121    /// Gets the number of `Interval<T>` values in the tree.
122    ///
123    /// ```
124    /// use lz_diet::Diet;
125    ///
126    /// let mut diet = Diet::new();
127    ///
128    /// assert_eq!(diet.len(), 0);
129    ///
130    /// diet.insert(8u32);
131    /// assert_eq!(diet.len(), 1);
132    ///
133    /// diet.insert(10);
134    /// assert_eq!(diet.len(), 2, "a new interval should have started");
135    ///
136    /// diet.insert(9);
137    /// assert_eq!(diet.len(), 1, "the previous intervals should have merged");
138    /// ```
139    pub fn len(&self) -> usize {
140        self.root().map_or(0, |node| node.len())
141    }
142
143    /// Gets whether the `Diet<T>` is empty or contains `Interval<T>` values.
144    ///
145    /// ```
146    /// use lz_diet::Diet;
147    ///
148    /// let mut diet = Diet::new();
149    /// assert!(diet.is_empty(), "a new tree should always be empty");
150    ///
151    /// diet.insert(7u32);
152    /// assert_eq!(diet.is_empty(), false);
153    ///
154    /// diet.clear();
155    /// assert_eq!(diet.is_empty(), true);
156    /// ```
157    pub fn is_empty(&self) -> bool {
158        self.root().is_none()
159    }
160
161    /// Gets whether the `Diet<T>` contains the specified value.
162    ///
163    /// ```
164    /// use lz_diet::Diet;
165    ///
166    /// let mut diet = Diet::new();
167    /// assert_eq!(diet.contains(&5), false);
168    ///
169    /// diet.insert(5u32);
170    /// assert!(diet.contains(&5));
171    /// ```
172    pub fn contains<Q>(&self, value: &Q) -> bool
173    where
174        T: Borrow<Q>,
175        Q: ?Sized + Ord,
176    {
177        if let Some(ref root) = self.root {
178            let mut contains = false;
179            root.walk(|node| {
180                let walk_action = node
181                    .calculate_walk_direction(value)
182                    .ok()
183                    .map(|direction| direction.into())
184                    .unwrap_or(WalkAction::Stop);
185
186                if matches!(walk_action, WalkAction::Stop) {
187                    contains = true;
188                }
189                walk_action
190            });
191
192            contains
193        } else {
194            false
195        }
196    }
197
198    /// Get the next value outside the `Diet<T>`
199    ///
200    /// Returns a value `r` that's greater or equal to `from`
201    /// and for which `diet.contains(r)` returns `false`.
202    pub fn find_next_gap<'a, 'b, Q>(&'b self, from: &'a Q) -> &'a Q
203    where
204        T: Borrow<Q>,
205        Q: ?Sized + Ord,
206        'b: 'a,
207    {
208        if let Some(ref root) = self.root {
209            let mut next = from;
210            root.walk(|node| {
211                let walk_action = node.calculate_walk_direction(from)
212                    .ok()
213                    .map(|direction| direction.into())
214                    .unwrap_or(WalkAction::Stop);
215
216                if matches!(walk_action, WalkAction::Stop) {
217                    next = node.value().exclusive_end().borrow();
218                }
219                walk_action
220            });
221
222            next
223        } else {
224            from
225        }
226    }
227}
228
229impl<A: AdjacentBound> FromIterator<A> for Diet<A> {
230    fn from_iter<T>(iter: T) -> Self
231    where
232        T: IntoIterator<Item = A>,
233    {
234        let mut diet = Diet::new();
235
236        for value in iter {
237            diet.insert(value);
238        }
239
240        diet
241    }
242}
243
244impl<T> BinaryTree for Diet<T> {
245    type Node = DietNode<T>;
246
247    fn root(&self) -> Option<&Self::Node> {
248        self.root.as_ref().map(|n| &**n)
249    }
250}
251
252impl<T: PartialEq> PartialEq for Diet<T> {
253    fn eq(&self, other: &Self) -> bool {
254        let self_intervals = self.into_iter();
255        let other_intervals = other.into_iter();
256
257        self_intervals.eq(other_intervals)
258    }
259}
260impl<T: Hash> Hash for Diet<T> {
261    fn hash<H: Hasher>(&self, state: &mut H) {
262        let intervals: Vec<_> = self.into_iter().collect();
263
264        intervals.hash(state);
265    }
266}
267
268impl<'a, T> IntoIterator for &'a Diet<T> {
269    type Item = &'a Interval<T>;
270    type IntoIter = Iter<'a, T>;
271
272    fn into_iter(self) -> Self::IntoIter {
273        Iter::new(self.root())
274    }
275}
276
277impl<T> IntoIterator for Diet<T> {
278    type Item = Interval<T>;
279    type IntoIter = IntoIter<T>;
280
281    fn into_iter(mut self) -> Self::IntoIter {
282        IntoIter::new(self.root.take())
283    }
284}
285
286impl<T: AdjacentBound> Diet<T> {
287    /// Inserts a new value into the `Diet<T>`.
288    ///
289    /// # Returns
290    /// true - if the value was inserted.
291    /// false - if the value already existed contained.
292    ///
293    /// ```
294    /// use lz_diet::Diet;
295    ///
296    /// let mut diet = Diet::new();
297    /// assert!(diet.insert(5u32));
298    /// assert_eq!(diet.insert(5), false);
299    /// assert_eq!(diet.iter().collect::<Vec<_>>(), vec![&(5..6).into()]);
300    /// assert!(diet.insert(15u32));
301    /// assert_eq!(diet.insert(15), false);
302    /// ```
303    pub fn insert(&mut self, value: T) -> bool {
304        if let Some(ref mut root) = self.root {
305            root.insert(value)
306        } else {
307            let exclusive_end = value.increment();
308
309            let new_node = Box::new(DietNode::new(value..exclusive_end));
310
311            self.root = Some(new_node);
312            true
313        }
314    }
315
316    /// Removes a value from the `Diet<T>`.
317    ///
318    /// This takes a `Cow<T>` as an owned value is only required if the value
319    /// is in the middle of an `Interval<T>`.
320    ///
321    /// # Returns
322    /// true - if the value was found and removed.
323    /// false - if the value was not found.
324    ///
325    /// ```
326    /// use lz_diet::Diet;
327    /// use std::borrow::Cow;
328    ///
329    /// let mut diet = Diet::new();
330    ///
331    /// let to_remove = 5u32;
332    /// // for a u32 we would probably just use Owned but this demonstrates how
333    /// // borrowed values may also be used.
334    /// assert_eq!(diet.remove(Cow::Borrowed(&to_remove)), false);
335    ///
336    /// diet.insert(5u32);
337    /// assert!(diet.remove(Cow::Borrowed(&to_remove)));
338    /// ```
339    pub fn remove<Q>(&mut self, value: Cow<Q>) -> bool
340    where
341        T: Borrow<Q>,
342        Q: ?Sized + Ord + ToOwned<Owned = T> + AdjacentBound,
343    {
344        let remove_result = self
345            .root
346            .as_mut()
347            .map(|root| root.remove(value))
348            .unwrap_or(Err(()));
349
350        match remove_result {
351            Ok(true) => {
352                if self
353                    .root
354                    .as_mut()
355                    .expect("there must be a root node to be removed")
356                    .try_remove(|node, _| node.rebalance())
357                    .is_none()
358                {
359                    self.root = None;
360                }
361
362                true
363            }
364            Ok(false) => true,
365            Err(()) => false,
366        }
367    }
368
369    /// Splits a `Diet<T>` on a value. This will drain the elements from self.
370    ///
371    /// # Returns
372    /// Two `Diet<T>` values where the first contains children less than the
373    /// value and the second is all children greater than the value.
374    ///
375    /// ```
376    /// use lz_diet::Diet;
377    /// use std::borrow::Cow;
378    ///
379    /// let mut diet = Diet::new();
380    /// diet.extend_from_slice(&[6u32, 7, 10, 11, 15, 16, 17]);
381    ///
382    /// let (left, right) = diet.split(Cow::Owned(16));
383    /// assert_eq!(left.into_iter().collect::<Vec<_>>(), vec![(6..8).into(), (10..12).into(), (15..16).into()]);
384    /// assert_eq!(right.into_iter().collect::<Vec<_>>(), vec![(17..18).into()]);
385    /// ```
386    pub fn split<Q>(&mut self, value: Cow<Q>) -> (Diet<T>, Diet<T>)
387    where
388        T: Borrow<Q>,
389        Q: ?Sized + Ord + ToOwned<Owned = T> + AdjacentBound,
390    {
391        let split_result = self
392            .root
393            .take()
394            .map(|node| node.split(value))
395            .unwrap_or(SplitResult::None);
396
397        match split_result {
398            SplitResult::Split(left, right) => (
399                Diet {
400                    root: Some(Box::new(left)),
401                },
402                Diet {
403                    root: Some(Box::new(right)),
404                },
405            ),
406            SplitResult::Single(node) => (
407                Diet {
408                    root: Some(Box::new(node)),
409                },
410                Diet::new(),
411            ),
412            SplitResult::None => (Diet::new(), Diet::new()),
413        }
414    }
415}
416
417impl<T: AdjacentBound + Clone> Diet<T> {
418    pub fn extend_from_slice(&mut self, other: &[T]) {
419        for val in other.into_iter().cloned() {
420            self.insert(val);
421        }
422    }
423}
424
425impl<T> Default for Diet<T> {
426    fn default() -> Self {
427        Self::new()
428    }
429}
430
431#[cfg(test)]
432mod tests {
433    use super::*;
434    use std::borrow::Cow;
435
436    #[test]
437    fn contains_returns_false_for_default() {
438        let diet = Diet::<u32>::default();
439
440        assert!(!diet.contains(&5));
441    }
442
443    #[test]
444    fn contains_returns_true_for_existing_value() {
445        let diet = Diet::from_iter([3, 1, 5].iter().cloned());
446
447        assert!(diet.contains(&5));
448    }
449
450    #[test]
451    fn find_next_gap_for_empty() {
452        let diet = Diet::<u32>::default();
453
454        assert!(diet.find_next_gap(&5) == &5);
455    }
456
457    #[test]
458    fn find_next_gap_from_outside() {
459        let diet = Diet::from_iter([3, 1, 5].iter().cloned());
460
461        assert!(diet.find_next_gap(&4) == &4);
462    }
463
464    #[test]
465    fn find_next_gap_from_inside() {
466        let diet = Diet::from_iter([3, 1, 5].iter().cloned());
467
468        assert!(diet.find_next_gap(&5) == &6);
469    }
470
471    #[test]
472    fn len_returns_zero_for_default() {
473        let diet = Diet::<u32>::default();
474
475        assert_eq!(diet.len(), 0);
476    }
477
478    #[test]
479    fn insert_returns_true_for_new_value() {
480        let mut diet = Diet::default();
481
482        assert!(diet.insert(1));
483    }
484
485    #[test]
486    fn insert_returns_false_for_existing_value() {
487        let mut diet = Diet::default();
488
489        diet.insert(1);
490        assert!(!diet.insert(1));
491    }
492
493    #[test]
494    fn insert_extends_existing_ranges() {
495        let mut diet = Diet::from_iter([1, 5].iter().cloned());
496
497        diet.insert(2);
498        diet.insert(4);
499
500        assert_eq!(diet.len(), 2);
501    }
502
503    #[test]
504    fn insert_combines_range() {
505        let mut diet = Diet::from_iter([1, 3].iter().cloned());
506
507        diet.insert(2);
508
509        assert_eq!(diet.len(), 1);
510    }
511
512    #[test]
513    fn insert_combines_ranges() {
514        let mut diet = Diet::from_iter([3, 1, 5, 8].iter().cloned());
515
516        diet.insert(2);
517        diet.insert(4);
518        diet.insert(6);
519        diet.insert(7);
520
521        assert_eq!(diet.len(), 1);
522    }
523
524    #[test]
525    fn remove_returns_false_for_default() {
526        let mut diet = Diet::<u32>::default();
527
528        assert!(!diet.remove(Cow::Owned(5)));
529    }
530
531    #[test]
532    fn remove_returns_false_for_non_existant_value() {
533        let mut diet = Diet::from_iter([1, 2, 3, 6].iter().cloned());
534
535        assert!(!diet.remove(Cow::Owned(10)));
536    }
537
538    #[test]
539    fn remove_of_adjacent_returns_false() {
540        let mut diet = Diet::from_iter([4, 5, 6].iter().cloned());
541
542        assert!(!diet.remove(Cow::Owned(3)));
543    }
544
545    #[test]
546    fn remove_of_lower_bounds() {
547        let mut diet = Diet::from_iter([50, 51, 52, 1, 2, 20, 21].iter().cloned());
548
549        assert!(diet.remove(Cow::Owned(50)));
550        assert!(!diet.contains(&50));
551
552        assert!(diet.remove(Cow::Owned(51)));
553        assert!(!diet.contains(&51));
554
555        assert!(diet.remove(Cow::Owned(1)));
556        assert!(!diet.contains(&1));
557
558        assert!(diet.remove(Cow::Owned(20)));
559        assert!(!diet.contains(&20));
560
561        assert_eq!(diet.len(), 3);
562    }
563
564    #[test]
565    fn remove_of_upper_bounds() {
566        let mut diet = Diet::from_iter([50, 51, 52, 1, 2, 20, 21].iter().cloned());
567
568        assert!(diet.remove(Cow::Owned(52)));
569        assert!(!diet.contains(&52));
570
571        assert!(diet.remove(Cow::Owned(51)));
572        assert!(!diet.contains(&51));
573
574        assert!(diet.remove(Cow::Owned(2)));
575        assert!(!diet.contains(&2));
576
577        assert!(diet.remove(Cow::Owned(21)));
578        assert!(!diet.contains(&21));
579
580        assert_eq!(diet.len(), 3);
581    }
582
583    #[test]
584    fn remove_root_node() {
585        let mut diet = Diet::from_iter([50, 51, 1, 2, 10, 20, 21].iter().cloned());
586
587        assert!(diet.remove(Cow::Owned(50)));
588        assert!(!diet.contains(&50));
589        assert!(diet.remove(Cow::Owned(51)));
590        assert!(!diet.contains(&51));
591
592        assert_eq!(diet.len(), 3);
593
594        assert!(diet.contains(&1));
595        assert!(diet.contains(&2));
596        assert!(diet.contains(&10));
597        assert!(diet.contains(&20));
598        assert!(diet.contains(&21));
599    }
600
601    #[test]
602    fn remove_within_interval() {
603        let mut diet = Diet::from_iter([1, 2, 3, 5].iter().cloned());
604
605        assert!(diet.remove(Cow::Owned(2)));
606
607        assert!(!diet.contains(&2));
608
609        assert!(diet.contains(&1));
610        assert!(diet.contains(&3));
611        assert!(diet.contains(&5));
612
613        assert_eq!(diet.len(), 3);
614    }
615
616    #[test]
617    fn remove_entire_node() {
618        let mut diet = Diet::from_iter([1, 3, 5].iter().cloned());
619
620        assert!(diet.remove(Cow::Owned(3)));
621        assert_eq!(diet.len(), 2);
622        assert!(diet.contains(&1));
623        assert!(diet.contains(&5));
624    }
625
626    #[test]
627    fn extend_from_slice_inserts_values() {
628        let mut diet = Diet::default();
629
630        diet.extend_from_slice(&[1, 5, 3]);
631
632        assert!(diet.contains(&1));
633        assert!(diet.contains(&5));
634        assert!(diet.contains(&3));
635    }
636
637    #[test]
638    fn equals_with_different_insertion_order() {
639        let first = Diet::from_iter([1, 2, 5].iter().cloned());
640        let second = Diet::from_iter([5, 1, 2].iter().cloned());
641
642        assert_eq!(first, second);
643    }
644
645    #[test]
646    fn clone() {
647        let diet = Diet::from_iter([1, 2, 5].iter().cloned());
648        let cloned = diet.clone();
649
650        assert_eq!(diet, cloned);
651    }
652}