dbs_allocator/
interval_tree.rs

1// Copyright (C) 2019 Alibaba Cloud. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! An interval tree implementation specialized for VMM resource management.
5//!
6//! It's not designed as a generic interval tree, but specialized for VMM resource management.
7//! In addition to the normal get()/insert()/delete()/update() tree operations, it also implements
8//! allocate()/free() for resource allocation.
9//!
10//! # Examples
11//! ```rust
12//! extern crate dbs_allocator;
13//! use dbs_allocator::{Constraint, IntervalTree, NodeState, Range};
14//!
15//! // Create an interval tree and add available resources.
16//! let mut tree = IntervalTree::<u64>::new();
17//! tree.insert(Range::new(0x100u32, 0x100u32), None);
18//! tree.insert(Range::new(0x200u16, 0x2ffu16), None);
19//!
20//! // Allocate a range with constraints.
21//! let mut constraint = Constraint::new(8u64);
22//! constraint.min = 0x211;
23//! constraint.max = 0x21f;
24//! constraint.align = 0x8;
25//!
26//! let key = tree.allocate(&constraint);
27//! assert_eq!(key, Some(Range::new(0x218u64, 0x21fu64)));
28//! let val = tree.get(&Range::new(0x218u64, 0x21fu64));
29//! assert_eq!(val, Some(NodeState::Allocated));
30//!
31//! // Associate data with the allocated range and mark the range as occupied.
32//! // Note: caller needs to protect from concurrent access between allocate() and the first call
33//! // to update() to mark range as occupied.
34//! let old = tree.update(&Range::new(0x218u32, 0x21fu32), 2);
35//! assert_eq!(old, None);
36//! let old = tree.update(&Range::new(0x218u32, 0x21fu32), 3);
37//! assert_eq!(old, Some(2));
38//! let val = tree.get(&Range::new(0x218u32, 0x21fu32));
39//! assert_eq!(val, Some(NodeState::Valued(&3)));
40//!
41//! // Free allocated resource.
42//! let old = tree.free(key.as_ref().unwrap());
43//! assert_eq!(old, Some(3));
44//! ```
45
46use std::cmp::{max, min, Ordering};
47
48use crate::{AllocPolicy, Constraint};
49
50/// Represent a closed range `[min, max]`.
51#[allow(missing_docs)]
52#[derive(Copy, Clone, PartialEq, Eq, Hash)]
53pub struct Range {
54    pub min: u64,
55    pub max: u64,
56}
57
58impl std::fmt::Debug for Range {
59    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
60        write!(f, "[ {:016x}, {:016x} ]", self.min, self.max)
61    }
62}
63
64impl Range {
65    /// Create a instance of [`Range`] with given `min` and `max`.
66    ///
67    /// ## Panic
68    /// - if min is bigger than max
69    /// - if min == 0 && max == u64:MAX
70    pub fn new<T>(min: T, max: T) -> Self
71    where
72        u64: From<T>,
73    {
74        let umin = u64::from(min);
75        let umax = u64::from(max);
76        if umin > umax || (umin == 0 && umax == u64::MAX) {
77            panic!("interval_tree: Range({}, {}) is invalid", umin, umax);
78        }
79        Range {
80            min: umin,
81            max: umax,
82        }
83    }
84
85    /// Create a instance of [`Range`] with given base and size.
86    ///
87    /// ## Panic
88    /// - if base + size wraps around
89    /// - if base == 0 && size == u64::MAX
90    pub fn with_size<T>(base: T, size: T) -> Self
91    where
92        u64: From<T>,
93    {
94        let umin = u64::from(base);
95        let umax = u64::from(size).checked_add(umin).unwrap();
96        if umin > umax || (umin == 0 && umax == std::u64::MAX) {
97            panic!("interval_tree: Range({}, {}) is invalid", umin, umax);
98        }
99        Range {
100            min: umin,
101            max: umax,
102        }
103    }
104
105    /// Create a instance of [`Range`] containing only the point `value`.
106    pub fn new_point<T>(value: T) -> Self
107    where
108        u64: From<T>,
109    {
110        let val = u64::from(value);
111        Range { min: val, max: val }
112    }
113
114    /// Get size of the range.
115    pub fn len(&self) -> u64 {
116        self.max - self.min + 1
117    }
118
119    /// Check whether the range is empty.
120    pub fn is_empty(&self) -> bool {
121        false
122    }
123
124    /// Check whether two Range objects intersect with each other.
125    pub fn intersect(&self, other: &Range) -> bool {
126        max(self.min, other.min) <= min(self.max, other.max)
127    }
128
129    /// Check whether another [Range] object is fully covered by this range.
130    pub fn contain(&self, other: &Range) -> bool {
131        self.min <= other.min && self.max >= other.max
132    }
133
134    /// Create a new instance of [Range] with `min` aligned to `align`.
135    ///
136    /// # Examples
137    /// ```rust
138    /// extern crate dbs_allocator;
139    /// use dbs_allocator::Range;
140    ///
141    /// let a = Range::new(2u32, 6u32);
142    /// assert_eq!(a.align_to(0), Some(Range::new(2u32, 6u32)));
143    /// assert_eq!(a.align_to(1), Some(Range::new(2u16, 6u16)));
144    /// assert_eq!(a.align_to(2), Some(Range::new(2u64, 6u64)));
145    /// assert_eq!(a.align_to(4), Some(Range::new(4u8, 6u8)));
146    /// assert_eq!(a.align_to(8), None);
147    /// assert_eq!(a.align_to(3), None);
148    /// let b = Range::new(2u8, 2u8);
149    /// assert_eq!(b.align_to(2), Some(Range::new(2u8, 2u8)));
150    /// ```
151    pub fn align_to(&self, align: u64) -> Option<Range> {
152        match align {
153            0 | 1 => Some(*self),
154            _ => {
155                if align & (align - 1) != 0 {
156                    return None;
157                }
158                if let Some(min) = self.min.checked_add(align - 1).map(|v| v & !(align - 1)) {
159                    if min <= self.max {
160                        return Some(Range::new(min, self.max));
161                    }
162                }
163                None
164            }
165        }
166    }
167}
168
169impl PartialOrd for Range {
170    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
171        match self.min.cmp(&other.min) {
172            Ordering::Equal => Some(self.max.cmp(&other.max)),
173            res => Some(res),
174        }
175    }
176}
177
178impl Ord for Range {
179    fn cmp(&self, other: &Self) -> Ordering {
180        match self.min.cmp(&other.min) {
181            Ordering::Equal => self.max.cmp(&other.max),
182            res => res,
183        }
184    }
185}
186
187/// State of interval tree node.
188///
189/// Valid state transitions:
190/// - None -> Free: [IntervalTree::insert()]
191/// - None -> Valued: [IntervalTree::insert()]
192/// - Free -> Allocated: [IntervalTree::allocate()]
193/// - Allocated -> Valued(T): [IntervalTree::update()]
194/// - Valued -> Valued(T): [IntervalTree::update()]
195/// - Allocated -> Free: [IntervalTree::free()]
196/// - Valued(T) -> Free: [IntervalTree::free()]
197/// - * -> None: [IntervalTree::delete()]
198#[derive(Clone, Debug, PartialEq, PartialOrd, Eq, Ord)]
199pub enum NodeState<T> {
200    /// Node is free
201    Free,
202    /// Node is allocated but without associated data
203    Allocated,
204    /// Node is allocated with associated data.
205    Valued(T),
206}
207
208impl<T> NodeState<T> {
209    fn take(&mut self) -> Self {
210        std::mem::replace(self, NodeState::<T>::Free)
211    }
212
213    fn replace(&mut self, value: NodeState<T>) -> Self {
214        std::mem::replace(self, value)
215    }
216
217    fn as_ref(&self) -> NodeState<&T> {
218        match self {
219            NodeState::<T>::Valued(ref x) => NodeState::<&T>::Valued(x),
220            NodeState::<T>::Allocated => NodeState::<&T>::Allocated,
221            NodeState::<T>::Free => NodeState::<&T>::Free,
222        }
223    }
224
225    fn as_mut(&mut self) -> NodeState<&mut T> {
226        match self {
227            NodeState::<T>::Valued(ref mut x) => NodeState::<&mut T>::Valued(x),
228            NodeState::<T>::Allocated => NodeState::<&mut T>::Allocated,
229            NodeState::<T>::Free => NodeState::<&mut T>::Free,
230        }
231    }
232
233    fn is_free(&self) -> bool {
234        matches!(self, NodeState::<T>::Free)
235    }
236}
237
238impl<T> From<NodeState<T>> for Option<T> {
239    fn from(n: NodeState<T>) -> Option<T> {
240        match n {
241            NodeState::<T>::Free | NodeState::<T>::Allocated => None,
242            NodeState::<T>::Valued(data) => Some(data),
243        }
244    }
245}
246
247/// Internal tree node to implement interval tree.
248#[derive(Debug, PartialEq, Eq)]
249struct InnerNode<T> {
250    /// Interval handled by this node.
251    key: Range,
252    /// Optional contained data, None if the node is free.
253    data: NodeState<T>,
254    /// Optional left child of current node.
255    left: Option<Node<T>>,
256    /// Optional right child of current node.
257    right: Option<Node<T>>,
258    /// Cached height of the node.
259    height: u32,
260    /// Cached maximum valued covered by this node.
261    max_key: u64,
262}
263
264impl<T> InnerNode<T> {
265    fn new(key: Range, data: NodeState<T>) -> Self {
266        InnerNode {
267            key,
268            data,
269            left: None,
270            right: None,
271            height: 1,
272            max_key: key.max,
273        }
274    }
275}
276
277/// Newtype for interval tree nodes.
278#[derive(Debug, PartialEq, Eq)]
279struct Node<T>(Box<InnerNode<T>>);
280
281impl<T> Node<T> {
282    fn new(key: Range, data: Option<T>) -> Self {
283        let value = if let Some(t) = data {
284            NodeState::Valued(t)
285        } else {
286            NodeState::Free
287        };
288        Node(Box::new(InnerNode::new(key, value)))
289    }
290
291    /// Returns a readonly reference to the node associated with the `key` or None if not found.
292    fn search(&self, key: &Range) -> Option<&Self> {
293        match self.0.key.cmp(key) {
294            Ordering::Equal => Some(self),
295            Ordering::Less => self.0.right.as_ref().and_then(|node| node.search(key)),
296            Ordering::Greater => self.0.left.as_ref().and_then(|node| node.search(key)),
297        }
298    }
299
300    /// Returns a shared reference to the node covers full range of the `key`.
301    fn search_superset(&self, key: &Range) -> Option<&Self> {
302        if self.0.key.contain(key) {
303            Some(self)
304        } else if key.max < self.0.key.min && self.0.left.is_some() {
305            // Safe to unwrap() because we have just checked it.
306            self.0.left.as_ref().unwrap().search_superset(key)
307        } else if key.min > self.0.key.max && self.0.right.is_some() {
308            // Safe to unwrap() because we have just checked it.
309            self.0.right.as_ref().unwrap().search_superset(key)
310        } else {
311            None
312        }
313    }
314
315    /// Returns a mutable reference to the node covers full range of the `key`.
316    fn search_superset_mut(&mut self, key: &Range) -> Option<&mut Self> {
317        if self.0.key.contain(key) {
318            Some(self)
319        } else if key.max < self.0.key.min && self.0.left.is_some() {
320            // Safe to unwrap() because we have just checked it.
321            self.0.left.as_mut().unwrap().search_superset_mut(key)
322        } else if key.min > self.0.key.max && self.0.right.is_some() {
323            // Safe to unwrap() because we have just checked it.
324            self.0.right.as_mut().unwrap().search_superset_mut(key)
325        } else {
326            None
327        }
328    }
329
330    /// Insert a new (key, data) pair into the subtree.
331    ///
332    /// Note: it will panic if the new key intersects with existing nodes.
333    fn insert(mut self, key: Range, data: Option<T>) -> Self {
334        match self.0.key.cmp(&key) {
335            Ordering::Equal => {
336                panic!("interval_tree: key {:?} exists", key);
337            }
338            Ordering::Less => {
339                if self.0.key.intersect(&key) {
340                    panic!(
341                        "interval_tree: key {:?} intersects with existing {:?}",
342                        key, self.0.key
343                    );
344                }
345                match self.0.right {
346                    None => self.0.right = Some(Node::new(key, data)),
347                    Some(_) => self.0.right = self.0.right.take().map(|n| n.insert(key, data)),
348                }
349            }
350            Ordering::Greater => {
351                if self.0.key.intersect(&key) {
352                    panic!(
353                        "interval_tree: key {:?} intersects with existing {:?}",
354                        key, self.0.key
355                    );
356                }
357                match self.0.left {
358                    None => self.0.left = Some(Node::new(key, data)),
359                    Some(_) => self.0.left = self.0.left.take().map(|n| n.insert(key, data)),
360                }
361            }
362        }
363        self.updated_node()
364    }
365
366    /// Update an existing entry and return the old value.
367    fn update(&mut self, key: &Range, data: NodeState<T>) -> Option<T> {
368        match self.0.key.cmp(key) {
369            Ordering::Equal => {
370                match (self.0.data.as_ref(), data.as_ref()) {
371                    (NodeState::<&T>::Free, NodeState::<&T>::Free)
372                    | (NodeState::<&T>::Free, NodeState::<&T>::Valued(_))
373                    | (NodeState::<&T>::Allocated, NodeState::<&T>::Free)
374                    | (NodeState::<&T>::Allocated, NodeState::<&T>::Allocated)
375                    | (NodeState::<&T>::Valued(_), NodeState::<&T>::Free)
376                    | (NodeState::<&T>::Valued(_), NodeState::<&T>::Allocated) => {
377                        panic!("try to update unallocated interval tree node");
378                    }
379                    _ => {}
380                }
381                self.0.data.replace(data).into()
382            }
383            Ordering::Less => match self.0.right.as_mut() {
384                None => None,
385                Some(node) => node.update(key, data),
386            },
387            Ordering::Greater => match self.0.left.as_mut() {
388                None => None,
389                Some(node) => node.update(key, data),
390            },
391        }
392    }
393
394    /// Delete `key` from the subtree.
395    ///
396    /// Note: it doesn't return whether the key exists in the subtree, so caller need to ensure the
397    /// logic.
398    fn delete(mut self, key: &Range) -> (Option<T>, Option<Self>) {
399        match self.0.key.cmp(key) {
400            Ordering::Equal => {
401                let data = self.0.data.take();
402                return (data.into(), self.delete_root());
403            }
404            Ordering::Less => {
405                if let Some(node) = self.0.right.take() {
406                    let (data, right) = node.delete(key);
407                    self.0.right = right;
408                    return (data, Some(self.updated_node()));
409                }
410            }
411            Ordering::Greater => {
412                if let Some(node) = self.0.left.take() {
413                    let (data, left) = node.delete(key);
414                    self.0.left = left;
415                    return (data, Some(self.updated_node()));
416                }
417            }
418        }
419        (None, Some(self))
420    }
421
422    /// Rotate the node if necessary to keep balance.
423    fn rotate(self) -> Self {
424        let l = height(&self.0.left);
425        let r = height(&self.0.right);
426        match (l as i32) - (r as i32) {
427            1 | 0 | -1 => self,
428            2 => self.rotate_left_successor(),
429            -2 => self.rotate_right_successor(),
430            _ => unreachable!(),
431        }
432    }
433
434    /// Perform a single left rotation on this node.
435    fn rotate_left(mut self) -> Self {
436        let mut new_root = self.0.right.take().expect("Node is broken");
437        self.0.right = new_root.0.left.take();
438        self.update_cached_info();
439        new_root.0.left = Some(self);
440        new_root.update_cached_info();
441        new_root
442    }
443
444    /// Perform a single right rotation on this node.
445    fn rotate_right(mut self) -> Self {
446        let mut new_root = self.0.left.take().expect("Node is broken");
447        self.0.left = new_root.0.right.take();
448        self.update_cached_info();
449        new_root.0.right = Some(self);
450        new_root.update_cached_info();
451        new_root
452    }
453
454    /// Performs a rotation when the left successor is too high.
455    fn rotate_left_successor(mut self) -> Self {
456        let left = self.0.left.take().expect("Node is broken");
457        if height(&left.0.left) < height(&left.0.right) {
458            let rotated = left.rotate_left();
459            self.0.left = Some(rotated);
460            self.update_cached_info();
461        } else {
462            self.0.left = Some(left);
463        }
464        self.rotate_right()
465    }
466
467    /// Performs a rotation when the right successor is too high.
468    fn rotate_right_successor(mut self) -> Self {
469        let right = self.0.right.take().expect("Node is broken");
470        if height(&right.0.left) > height(&right.0.right) {
471            let rotated = right.rotate_right();
472            self.0.right = Some(rotated);
473            self.update_cached_info();
474        } else {
475            self.0.right = Some(right);
476        }
477        self.rotate_left()
478    }
479
480    fn delete_root(mut self) -> Option<Self> {
481        match (self.0.left.take(), self.0.right.take()) {
482            (None, None) => None,
483            (Some(l), None) => Some(l),
484            (None, Some(r)) => Some(r),
485            (Some(l), Some(r)) => Some(Self::combine_subtrees(l, r)),
486        }
487    }
488
489    /// Find the minimal key below the tree and returns a new optional tree where the minimal
490    /// value has been removed and the (optional) minimal node as tuple (min_node, remaining)
491    fn get_new_root(mut self) -> (Self, Option<Self>) {
492        match self.0.left.take() {
493            None => {
494                let remaining = self.0.right.take();
495                (self, remaining)
496            }
497            Some(left) => {
498                let (min_node, left) = left.get_new_root();
499                self.0.left = left;
500                (min_node, Some(self.updated_node()))
501            }
502        }
503    }
504
505    fn combine_subtrees(l: Self, r: Self) -> Self {
506        let (mut new_root, remaining) = r.get_new_root();
507        new_root.0.left = Some(l);
508        new_root.0.right = remaining;
509        new_root.updated_node()
510    }
511
512    fn find_candidate(&self, constraint: &Constraint) -> Option<&Self> {
513        match constraint.policy {
514            AllocPolicy::FirstMatch => self.first_match(constraint),
515            AllocPolicy::Default => self.first_match(constraint),
516        }
517    }
518
519    fn first_match(&self, constraint: &Constraint) -> Option<&Self> {
520        let mut candidate = if self.0.left.is_some() {
521            self.0.left.as_ref().unwrap().first_match(constraint)
522        } else {
523            None
524        };
525
526        if candidate.is_none() && self.check_constraint(constraint) {
527            candidate = Some(self);
528        }
529        if candidate.is_none() && self.0.right.is_some() {
530            candidate = self.0.right.as_ref().unwrap().first_match(constraint);
531        }
532        candidate
533    }
534
535    fn check_constraint(&self, constraint: &Constraint) -> bool {
536        if self.0.data.is_free() {
537            let min = std::cmp::max(self.0.key.min, constraint.min);
538            let max = std::cmp::min(self.0.key.max, constraint.max);
539            if min <= max {
540                let key = Range::new(min, max);
541                if constraint.align == 0 || constraint.align == 1 {
542                    return key.len() >= constraint.size;
543                }
544                return match key.align_to(constraint.align) {
545                    None => false,
546                    Some(aligned_key) => aligned_key.len() >= constraint.size,
547                };
548            }
549        }
550        false
551    }
552
553    /// Update cached information of the node.
554    /// Please make sure that the cached values of both children are up to date.
555    fn update_cached_info(&mut self) {
556        self.0.height = max(height(&self.0.left), height(&self.0.right)) + 1;
557        self.0.max_key = max(
558            max_key(&self.0.left),
559            max(max_key(&self.0.right), self.0.key.max),
560        );
561    }
562
563    /// Update the sub-tree to keep balance.
564    fn updated_node(mut self) -> Self {
565        self.update_cached_info();
566        self.rotate()
567    }
568}
569
570/// Compute height of the optional sub-tree.
571fn height<T>(node: &Option<Node<T>>) -> u32 {
572    node.as_ref().map_or(0, |n| n.0.height)
573}
574
575/// Compute maximum key value covered by the optional sub-tree.
576fn max_key<T>(node: &Option<Node<T>>) -> u64 {
577    node.as_ref().map_or(0, |n| n.0.max_key)
578}
579
580/// An interval tree implementation specialized for VMM resource management.
581#[derive(Debug, Default, PartialEq, Eq)]
582pub struct IntervalTree<T> {
583    root: Option<Node<T>>,
584}
585
586impl<T> IntervalTree<T> {
587    /// Construct a default empty [IntervalTree] object.
588    ///
589    /// # Examples
590    /// ```rust
591    /// extern crate dbs_allocator;
592    ///
593    /// let tree = dbs_allocator::IntervalTree::<u64>::new();
594    /// ```
595    pub fn new() -> Self {
596        IntervalTree { root: None }
597    }
598
599    /// Check whether the interval tree is empty.
600    pub fn is_empty(&self) -> bool {
601        self.root.is_none()
602    }
603
604    /// Get the data item associated with the key, or return None if no match found.
605    ///
606    /// # Examples
607    /// ```rust
608    /// extern crate dbs_allocator;
609    /// use dbs_allocator::{IntervalTree, NodeState, Range};
610    ///
611    /// let mut tree = dbs_allocator::IntervalTree::<u64>::new();
612    /// assert!(tree.is_empty());
613    /// assert_eq!(tree.get(&Range::new(0x101u64, 0x101u64)), None);
614    /// tree.insert(Range::new(0x100u64, 0x100u64), Some(1));
615    /// tree.insert(Range::new(0x200u64, 0x2ffu64), None);
616    /// assert!(!tree.is_empty());
617    /// assert_eq!(
618    ///     tree.get(&Range::new(0x100u64, 0x100u64)),
619    ///     Some(NodeState::Valued(&1))
620    /// );
621    /// assert_eq!(
622    ///     tree.get(&Range::new(0x200u64, 0x2ffu64)),
623    ///     Some(NodeState::Free)
624    /// );
625    /// assert_eq!(tree.get(&Range::new(0x101u64, 0x101u64)), None);
626    /// assert_eq!(tree.get(&Range::new(0x100u64, 0x101u64)), None);
627    /// ```
628    pub fn get(&self, key: &Range) -> Option<NodeState<&T>> {
629        match self.root {
630            None => None,
631            Some(ref node) => node.search(key).map(|n| n.0.data.as_ref()),
632        }
633    }
634
635    /// Get a shared reference to the node fully covering the entire key range.
636    ///
637    /// # Examples
638    /// ```rust
639    /// extern crate dbs_allocator;
640    /// use dbs_allocator::{IntervalTree, NodeState, Range};
641    ///
642    /// let mut tree = IntervalTree::<u64>::new();
643    /// tree.insert(Range::new(0x100u32, 0x100u32), Some(1));
644    /// tree.insert(Range::new(0x200u32, 0x2ffu32), None);
645    /// assert_eq!(
646    ///     tree.get_superset(&Range::new(0x100u32, 0x100u32)),
647    ///     Some((&Range::new(0x100u32, 0x100u32), NodeState::Valued(&1)))
648    /// );
649    /// assert_eq!(
650    ///     tree.get_superset(&Range::new(0x210u32, 0x210u32)),
651    ///     Some((&Range::new(0x200u32, 0x2ffu32), NodeState::Free))
652    /// );
653    /// assert_eq!(
654    ///     tree.get_superset(&Range::new(0x2ffu32, 0x2ffu32)),
655    ///     Some((&Range::new(0x200u32, 0x2ffu32), NodeState::Free))
656    /// );
657    /// ```
658    pub fn get_superset(&self, key: &Range) -> Option<(&Range, NodeState<&T>)> {
659        match self.root {
660            None => None,
661            Some(ref node) => node
662                .search_superset(key)
663                .map(|n| (&n.0.key, n.0.data.as_ref())),
664        }
665    }
666
667    /// Get a mutable reference to the node fully covering the entire key range.
668    ///
669    /// # Examples
670    /// ```rust
671    /// extern crate dbs_allocator;
672    /// use dbs_allocator::{IntervalTree, NodeState, Range};
673    ///
674    /// let mut tree = IntervalTree::<u64>::new();
675    /// tree.insert(Range::new(0x100u32, 0x100u32), Some(1));
676    /// tree.insert(Range::new(0x200u32, 0x2ffu32), None);
677    /// assert_eq!(
678    ///     tree.get_superset_mut(&Range::new(0x100u32, 0x100u32)),
679    ///     Some((&Range::new(0x100u32, 0x100u32), NodeState::Valued(&mut 1)))
680    /// );
681    /// assert_eq!(
682    ///     tree.get_superset_mut(&Range::new(0x210u32, 0x210u32)),
683    ///     Some((&Range::new(0x200u32, 0x2ffu32), NodeState::Free))
684    /// );
685    /// assert_eq!(
686    ///     tree.get_superset_mut(&Range::new(0x2ffu32, 0x2ffu32)),
687    ///     Some((&Range::new(0x200u32, 0x2ffu32), NodeState::Free))
688    /// );
689    /// ```
690    pub fn get_superset_mut(&mut self, key: &Range) -> Option<(&Range, NodeState<&mut T>)> {
691        match self.root {
692            None => None,
693            Some(ref mut node) => node
694                .search_superset_mut(key)
695                .map(|n| (&n.0.key, n.0.data.as_mut())),
696        }
697    }
698
699    /// Get a shared reference to the value associated with the id.
700    ///
701    /// # Examples
702    /// ```rust
703    /// extern crate dbs_allocator;
704    /// use dbs_allocator::{IntervalTree, NodeState, Range};
705    ///
706    /// let mut tree = IntervalTree::<u32>::new();
707    /// tree.insert(Range::new(0x100u16, 0x100u16), Some(1));
708    /// tree.insert(Range::new(0x200u16, 0x2ffu16), None);
709    /// assert_eq!(tree.get_by_id(0x100u16), Some(&1));
710    /// assert_eq!(tree.get_by_id(0x210u32), None);
711    /// assert_eq!(tree.get_by_id(0x2ffu64), None);
712    /// ```
713    pub fn get_by_id<U>(&self, id: U) -> Option<&T>
714    where
715        u64: From<U>,
716    {
717        match self.root {
718            None => None,
719            Some(ref node) => {
720                let key = Range::new_point(id);
721                match node.search_superset(&key) {
722                    Some(node) => node.0.data.as_ref().into(),
723                    None => None,
724                }
725            }
726        }
727    }
728
729    /// Get a mutable reference to the value associated with the id.
730    ///
731    /// # Examples
732    /// ```rust
733    /// extern crate dbs_allocator;
734    /// use dbs_allocator::{IntervalTree, NodeState, Range};
735    ///
736    /// let mut tree = IntervalTree::<u32>::new();
737    /// tree.insert(Range::new(0x100u16, 0x100u16), Some(1));
738    /// tree.insert(Range::new(0x200u16, 0x2ffu16), None);
739    /// assert_eq!(tree.get_by_id_mut(0x100u16), Some(&mut 1));
740    /// assert_eq!(tree.get_by_id_mut(0x210u32), None);
741    /// assert_eq!(tree.get_by_id_mut(0x2ffu64), None);
742    /// ```
743    pub fn get_by_id_mut<U>(&mut self, id: U) -> Option<&mut T>
744    where
745        u64: From<U>,
746    {
747        match self.root {
748            None => None,
749            Some(ref mut node) => {
750                let key = Range::new_point(id);
751                match node.search_superset_mut(&key) {
752                    Some(node) => node.0.data.as_mut().into(),
753                    None => None,
754                }
755            }
756        }
757    }
758
759    /// Insert the (key, data) pair into the interval tree, panic if intersects with existing nodes.
760    ///
761    /// # Examples
762    /// ```rust
763    /// extern crate dbs_allocator;
764    /// use dbs_allocator::{IntervalTree, NodeState, Range};
765    ///
766    /// let mut tree = IntervalTree::<u64>::new();
767    /// tree.insert(Range::new(0x100u32, 0x100u32), Some(1));
768    /// tree.insert(Range::new(0x200u32, 0x2ffu32), None);
769    /// assert_eq!(
770    ///     tree.get(&Range::new(0x100u64, 0x100u64)),
771    ///     Some(NodeState::Valued(&1))
772    /// );
773    /// assert_eq!(
774    ///     tree.get(&Range::new(0x200u64, 0x2ffu64)),
775    ///     Some(NodeState::Free)
776    /// );
777    /// ```
778    pub fn insert(&mut self, key: Range, data: Option<T>) {
779        match self.root.take() {
780            None => self.root = Some(Node::new(key, data)),
781            Some(node) => self.root = Some(node.insert(key, data)),
782        }
783    }
784
785    /// Update an existing entry and return the old value.
786    ///
787    /// # Examples
788    /// ```rust
789    /// extern crate dbs_allocator;
790    /// use dbs_allocator::{Constraint, IntervalTree, Range};
791    ///
792    /// let mut tree = IntervalTree::<u64>::new();
793    /// tree.insert(Range::new(0x100u64, 0x100u64), None);
794    /// tree.insert(Range::new(0x200u64, 0x2ffu64), None);
795    ///
796    /// let constraint = Constraint::new(2u32);
797    /// let key = tree.allocate(&constraint);
798    /// assert_eq!(key, Some(Range::new(0x200u64, 0x201u64)));
799    /// let old = tree.update(&Range::new(0x200u64, 0x201u64), 2);
800    /// assert_eq!(old, None);
801    /// let old = tree.update(&Range::new(0x200u64, 0x201u64), 3);
802    /// assert_eq!(old, Some(2));
803    /// ```
804    pub fn update(&mut self, key: &Range, data: T) -> Option<T> {
805        match self.root.as_mut() {
806            None => None,
807            Some(node) => node.update(key, NodeState::<T>::Valued(data)),
808        }
809    }
810
811    /// Remove the `key` from the tree and return the associated data.
812    ///
813    /// # Examples
814    /// ```rust
815    /// extern crate dbs_allocator;
816    /// use dbs_allocator::{IntervalTree, Range};
817    ///
818    /// let mut tree = IntervalTree::<u64>::new();
819    /// tree.insert(Range::new(0x100u64, 0x100u64), Some(1));
820    /// tree.insert(Range::new(0x200u64, 0x2ffu64), None);
821    /// let old = tree.delete(&Range::new(0x100u64, 0x100u64));
822    /// assert_eq!(old, Some(1));
823    /// let old = tree.delete(&Range::new(0x200u64, 0x2ffu64));
824    /// assert_eq!(old, None);
825    /// ```
826    pub fn delete(&mut self, key: &Range) -> Option<T> {
827        match self.root.take() {
828            Some(node) => {
829                let (data, root) = node.delete(key);
830                self.root = root;
831                data
832            }
833            None => None,
834        }
835    }
836
837    /// Allocate a resource range according the allocation constraints.
838    ///
839    /// # Examples
840    /// ```rust
841    /// extern crate dbs_allocator;
842    /// use dbs_allocator::{Constraint, IntervalTree, Range};
843    ///
844    /// let mut tree = IntervalTree::<u64>::new();
845    /// tree.insert(Range::new(0x100u64, 0x100u64), None);
846    /// tree.insert(Range::new(0x200u64, 0x2ffu64), None);
847    ///
848    /// let constraint = Constraint::new(2u8);
849    /// let key = tree.allocate(&constraint);
850    /// assert_eq!(key, Some(Range::new(0x200u64, 0x201u64)));
851    /// tree.update(&Range::new(0x200u64, 0x201u64), 2);
852    /// ```
853    pub fn allocate(&mut self, constraint: &Constraint) -> Option<Range> {
854        if constraint.size == 0 {
855            return None;
856        }
857        let candidate = match self.root.as_mut() {
858            None => None,
859            Some(node) => node.find_candidate(constraint),
860        };
861
862        match candidate {
863            None => None,
864            Some(node) => {
865                let node_key = node.0.key;
866                let range = Range::new(
867                    max(node_key.min, constraint.min),
868                    min(node_key.max, constraint.max),
869                );
870                // Safe to unwrap because candidate satisfy the constraints.
871                let aligned_key = range.align_to(constraint.align).unwrap();
872                let result = Range::new(aligned_key.min, aligned_key.min + constraint.size - 1);
873
874                // Allocate a resource from the node, no need to split the candidate node.
875                if node_key.min == aligned_key.min && node_key.len() == constraint.size {
876                    self.root
877                        .as_mut()
878                        .unwrap()
879                        .update(&node_key, NodeState::<T>::Allocated);
880                    return Some(node_key);
881                }
882
883                // Split the candidate node.
884                // TODO: following algorithm is not optimal in preference of simplicity.
885                self.delete(&node_key);
886                if aligned_key.min > node_key.min {
887                    self.insert(Range::new(node_key.min, aligned_key.min - 1), None);
888                }
889                self.insert(result, None);
890                if result.max < node_key.max {
891                    self.insert(Range::new(result.max + 1, node_key.max), None);
892                }
893
894                self.root
895                    .as_mut()
896                    .unwrap()
897                    .update(&result, NodeState::<T>::Allocated);
898                Some(result)
899            }
900        }
901    }
902
903    /// Free an allocated range and return the associated data.
904    pub fn free(&mut self, key: &Range) -> Option<T> {
905        let result = self.delete(key);
906        let mut range = *key;
907
908        // Try to merge with adjacent free nodes.
909        if range.min > 0 {
910            if let Some((r, v)) = self.get_superset(&Range::new(range.min - 1, range.min - 1)) {
911                if v.is_free() {
912                    range.min = r.min;
913                }
914            }
915        }
916        if range.max < std::u64::MAX {
917            if let Some((r, v)) = self.get_superset(&Range::new(range.max + 1, range.max + 1)) {
918                if v.is_free() {
919                    range.max = r.max;
920                }
921            }
922        }
923
924        if range.min < key.min {
925            self.delete(&Range::new(range.min, key.min - 1));
926        }
927        if range.max > key.max {
928            self.delete(&Range::new(key.max + 1, range.max));
929        }
930        self.insert(range, None);
931
932        result
933    }
934}
935
936#[cfg(test)]
937mod tests {
938    use super::*;
939
940    #[test]
941    #[should_panic]
942    fn test_new_range() {
943        let _ = Range::new(2u8, 1u8);
944    }
945
946    #[test]
947    #[should_panic]
948    fn test_new_range_overflow() {
949        let _ = Range::new(0u64, std::u64::MAX);
950    }
951
952    #[test]
953    fn test_range_intersect() {
954        let range_a = Range::new(1u8, 4u8);
955        let range_b = Range::new(4u16, 6u16);
956        let range_c = Range::new(2u32, 3u32);
957        let range_d = Range::new(4u64, 4u64);
958        let range_e = Range::new(5u32, 6u32);
959
960        assert!(range_a.intersect(&range_b));
961        assert!(range_b.intersect(&range_a));
962        assert!(range_a.intersect(&range_c));
963        assert!(range_c.intersect(&range_a));
964        assert!(range_a.intersect(&range_d));
965        assert!(range_d.intersect(&range_a));
966        assert!(!range_a.intersect(&range_e));
967        assert!(!range_e.intersect(&range_a));
968
969        assert_eq!(range_a.len(), 4);
970        assert_eq!(range_d.len(), 1);
971    }
972
973    #[test]
974    fn test_range_contain() {
975        let range_a = Range::new(2u8, 6u8);
976        assert!(range_a.contain(&Range::new(2u8, 3u8)));
977        assert!(range_a.contain(&Range::new(3u8, 4u8)));
978        assert!(range_a.contain(&Range::new(5u8, 5u8)));
979        assert!(range_a.contain(&Range::new(5u8, 6u8)));
980        assert!(range_a.contain(&Range::new(6u8, 6u8)));
981        assert!(!range_a.contain(&Range::new(1u8, 1u8)));
982        assert!(!range_a.contain(&Range::new(1u8, 2u8)));
983        assert!(!range_a.contain(&Range::new(1u8, 3u8)));
984        assert!(!range_a.contain(&Range::new(1u8, 7u8)));
985        assert!(!range_a.contain(&Range::new(7u8, 8u8)));
986        assert!(!range_a.contain(&Range::new(6u8, 7u8)));
987        assert!(!range_a.contain(&Range::new(7u8, 8u8)));
988    }
989
990    #[test]
991    fn test_range_align_to() {
992        let range_a = Range::new(2u32, 6);
993        assert_eq!(range_a.align_to(0), Some(Range::new(2u64, 6u64)));
994        assert_eq!(range_a.align_to(1), Some(Range::new(2u8, 6u8)));
995        assert_eq!(range_a.align_to(2), Some(Range::new(2u16, 6u16)));
996        assert_eq!(range_a.align_to(4), Some(Range::new(4u32, 6u32)));
997        assert_eq!(range_a.align_to(8), None);
998        assert_eq!(range_a.align_to(3), None);
999
1000        let range_b = Range::new(0xFFFF_FFFF_FFFF_FFFDu64, 0xFFFF_FFFF_FFFF_FFFFu64);
1001        assert_eq!(
1002            range_b.align_to(2),
1003            Some(Range::new(0xFFFF_FFFF_FFFF_FFFEu64, 0xFFFF_FFFF_FFFF_FFFF))
1004        );
1005        assert_eq!(range_b.align_to(4), None);
1006    }
1007
1008    #[test]
1009    fn test_range_ord() {
1010        let range_a = Range::new(1u32, 4u32);
1011        let range_b = Range::new(1u32, 4u32);
1012        let range_c = Range::new(1u32, 3u32);
1013        let range_d = Range::new(1u32, 5u32);
1014        let range_e = Range::new(2u32, 2u32);
1015
1016        assert_eq!(range_a, range_b);
1017        assert_eq!(range_b, range_a);
1018        assert!(range_a > range_c);
1019        assert!(range_c < range_a);
1020        assert!(range_a < range_d);
1021        assert!(range_d > range_a);
1022        assert!(range_a < range_e);
1023        assert!(range_e > range_a);
1024    }
1025
1026    #[should_panic]
1027    #[test]
1028    fn test_tree_insert_equal() {
1029        let mut tree = IntervalTree::<u64>::new();
1030        tree.insert(Range::new(0x100u16, 0x200), Some(1));
1031        tree.insert(Range::new(0x100u32, 0x200), None);
1032    }
1033
1034    #[should_panic]
1035    #[test]
1036    fn test_tree_insert_intersect_on_right() {
1037        let mut tree = IntervalTree::<u64>::new();
1038        tree.insert(Range::new(0x100, 0x200u32), Some(1));
1039        tree.insert(Range::new(0x200, 0x2ffu64), None);
1040    }
1041
1042    #[should_panic]
1043    #[test]
1044    fn test_tree_insert_intersect_on_left() {
1045        let mut tree = IntervalTree::<u64>::new();
1046        tree.insert(Range::new(0x100, 0x200u32), Some(1));
1047        tree.insert(Range::new(0x000, 0x100u64), None);
1048    }
1049
1050    #[test]
1051    fn test_tree_get_superset() {
1052        let mut tree = IntervalTree::<u64>::new();
1053        tree.insert(Range::new(0x100u32, 0x100u32), Some(1));
1054        tree.insert(Range::new(0x001u16, 0x008u16), None);
1055        tree.insert(Range::new(0x009u16, 0x00fu16), None);
1056        tree.insert(Range::new(0x200u16, 0x2ffu16), None);
1057        let mut constraint = Constraint::new(8u64);
1058        constraint.min = 0x211;
1059        constraint.max = 0x21f;
1060        constraint.align = 0x8;
1061        tree.allocate(&constraint);
1062
1063        // Valued case.
1064        assert_eq!(
1065            tree.get_superset(&Range::new(0x100u32, 0x100)),
1066            Some((&Range::new(0x100, 0x100u32), NodeState::Valued(&1)))
1067        );
1068
1069        // Free case.
1070        assert_eq!(
1071            tree.get_superset(&Range::new(0x200u16, 0x200)),
1072            Some((&Range::new(0x200, 0x217u64), NodeState::Free))
1073        );
1074        assert_eq!(
1075            tree.get_superset(&Range::new(0x2ffu32, 0x2ff)),
1076            Some((&Range::new(0x220, 0x2ffu32), NodeState::Free))
1077        );
1078
1079        // Allocated case.
1080        assert_eq!(
1081            tree.get_superset(&Range::new(0x218u16, 0x21f)),
1082            Some((&Range::new(0x218, 0x21fu16), NodeState::Allocated))
1083        );
1084
1085        // None case.
1086        assert_eq!(tree.get_superset(&Range::new(0x2ffu32, 0x300)), None);
1087        assert_eq!(tree.get_superset(&Range::new(0x300u32, 0x300)), None);
1088        assert_eq!(tree.get_superset(&Range::new(0x1ffu32, 0x300)), None);
1089    }
1090
1091    #[test]
1092    fn test_tree_get_superset_mut() {
1093        let mut tree = IntervalTree::<u64>::new();
1094        tree.insert(Range::new(0x100u32, 0x100u32), Some(1));
1095        tree.insert(Range::new(0x200u16, 0x2ffu16), None);
1096        let mut constraint = Constraint::new(8u64);
1097        constraint.min = 0x211;
1098        constraint.max = 0x21f;
1099        constraint.align = 0x8;
1100        tree.allocate(&constraint);
1101
1102        // Valued case.
1103        assert_eq!(
1104            tree.get_superset_mut(&Range::new(0x100u32, 0x100u32)),
1105            Some((&Range::new(0x100u32, 0x100u32), NodeState::Valued(&mut 1)))
1106        );
1107
1108        // Allocated case.
1109        assert_eq!(
1110            tree.get_superset_mut(&Range::new(0x218u64, 0x21fu64)),
1111            Some((&Range::new(0x218u64, 0x21fu64), NodeState::Allocated))
1112        );
1113
1114        // Free case.
1115        assert_eq!(
1116            tree.get_superset_mut(&Range::new(0x2ffu32, 0x2ffu32)),
1117            Some((&Range::new(0x220u32, 0x2ffu32), NodeState::Free))
1118        );
1119
1120        // None case.
1121        assert_eq!(tree.get_superset(&Range::new(0x2ffu32, 0x300)), None);
1122        assert_eq!(tree.get_superset(&Range::new(0x300u32, 0x300)), None);
1123        assert_eq!(tree.get_superset(&Range::new(0x1ffu32, 0x300)), None);
1124    }
1125
1126    #[test]
1127    fn test_tree_update() {
1128        let mut tree = IntervalTree::<u64>::new();
1129        tree.insert(Range::new(0x100u32, 0x100u32), None);
1130        tree.insert(Range::new(0x200u32, 0x2ffu32), None);
1131
1132        let constraint = Constraint::new(2u32);
1133        let key = tree.allocate(&constraint);
1134        assert_eq!(key, Some(Range::new(0x200u32, 0x201u32)));
1135        let old = tree.update(&Range::new(0x200u32, 0x201u32), 2);
1136        assert_eq!(old, None);
1137        let old = tree.update(&Range::new(0x200u32, 0x201u32), 3);
1138        assert_eq!(old, Some(2));
1139        let old = tree.update(&Range::new(0x200u32, 0x200u32), 4);
1140        assert_eq!(old, None);
1141        let old = tree.update(&Range::new(0x200u32, 0x203u32), 5);
1142        assert_eq!(old, None);
1143
1144        tree.delete(&Range::new(0x200u32, 0x201u32));
1145        let old = tree.update(&Range::new(0x200u32, 0x201u32), 2);
1146        assert_eq!(old, None);
1147    }
1148
1149    #[test]
1150    fn test_tree_delete() {
1151        let mut tree = IntervalTree::<u64>::new();
1152        assert_eq!(tree.get(&Range::new(0x101u32, 0x101u32)), None);
1153        assert!(tree.is_empty());
1154        tree.insert(Range::new(0x100u32, 0x100u32), Some(1));
1155        tree.insert(Range::new(0x001u16, 0x00fu16), None);
1156        tree.insert(Range::new(0x200u32, 0x2ffu32), None);
1157        assert!(!tree.is_empty());
1158        assert_eq!(
1159            tree.get(&Range::new(0x100u32, 0x100u32)),
1160            Some(NodeState::Valued(&1))
1161        );
1162        assert_eq!(
1163            tree.get(&Range::new(0x200u32, 0x2ffu32)),
1164            Some(NodeState::Free)
1165        );
1166        assert_eq!(tree.get(&Range::new(0x101u32, 0x101u32)), None);
1167
1168        let old = tree.delete(&Range::new(0x001u16, 0x00fu16));
1169        assert_eq!(old, None);
1170        let old = tree.delete(&Range::new(0x100u32, 0x100u32));
1171        assert_eq!(old, Some(1));
1172        let old = tree.delete(&Range::new(0x200u32, 0x2ffu32));
1173        assert_eq!(old, None);
1174
1175        assert!(tree.is_empty());
1176        assert_eq!(tree.get(&Range::new(0x100u32, 0x100u32)), None);
1177        assert_eq!(tree.get(&Range::new(0x200u32, 0x2ffu32)), None);
1178    }
1179
1180    #[test]
1181    fn test_allocate_free() {
1182        let mut tree = IntervalTree::<u64>::new();
1183        let mut constraint = Constraint::new(1u8);
1184
1185        assert_eq!(tree.allocate(&constraint), None);
1186        tree.insert(Range::new(0x100u16, 0x100u16), None);
1187        tree.insert(Range::new(0x200u16, 0x2ffu16), None);
1188
1189        let key = tree.allocate(&constraint);
1190        assert_eq!(key, Some(Range::new(0x100u16, 0x100u16)));
1191        let old = tree.update(&Range::new(0x100u16, 0x100u16), 2);
1192        assert_eq!(old, None);
1193        let val = tree.get(&Range::new(0x100u16, 0x100u16));
1194        assert_eq!(val, Some(NodeState::Valued(&2)));
1195
1196        constraint.min = 0x100;
1197        constraint.max = 0x100;
1198        assert_eq!(tree.allocate(&constraint), None);
1199
1200        constraint.min = 0x201;
1201        constraint.max = 0x300;
1202        constraint.align = 0x8;
1203        constraint.size = 0x10;
1204        assert_eq!(
1205            tree.allocate(&constraint),
1206            Some(Range::new(0x208u16, 0x217u16))
1207        );
1208
1209        // Free the node when it's still in 'Allocated' state.
1210        let old = tree.free(&Range::new(0x208u16, 0x217u16));
1211        assert_eq!(old, None);
1212
1213        // Reallocate the freed resource.
1214        assert_eq!(
1215            tree.allocate(&constraint),
1216            Some(Range::new(0x208u16, 0x217u16))
1217        );
1218
1219        constraint.size = 0x100;
1220        assert_eq!(tree.allocate(&constraint), None);
1221
1222        // Verify that allocating a bigger range with smaller allocated range fails.
1223        constraint.min = 0x200;
1224        constraint.max = 0x2ff;
1225        constraint.align = 0x8;
1226        constraint.size = 0x100;
1227        assert_eq!(tree.allocate(&constraint), None);
1228
1229        // Free the node when it's in 'Valued' state.
1230        tree.update(&Range::new(0x208u16, 0x217u16), 0x10);
1231        assert_eq!(tree.allocate(&constraint), None);
1232        let old = tree.free(&Range::new(0x208u16, 0x217u16));
1233        assert_eq!(old, Some(0x10));
1234
1235        // Reallocate the freed resource, verify that adjacent free nodes have been merged.
1236        assert_eq!(
1237            tree.allocate(&constraint),
1238            Some(Range::new(0x200u32, 0x2ffu32))
1239        );
1240    }
1241
1242    #[test]
1243    fn test_with_size() {
1244        let range_a = Range::with_size(1u8, 3u8);
1245        let range_b = Range::with_size(4u16, 2u16);
1246        let range_c = Range::with_size(2u32, 1u32);
1247        let range_d = Range::with_size(4u64, 0u64);
1248        let range_e = Range::with_size(5u32, 1u32);
1249
1250        assert_eq!(range_a, Range::new(1u8, 4u8));
1251        assert_eq!(range_b, Range::new(4u16, 6u16));
1252        assert_eq!(range_c, Range::new(2u32, 3u32));
1253        assert_eq!(range_d, Range::new(4u64, 4u64));
1254        assert_eq!(range_e, Range::new(5u32, 6u32));
1255    }
1256
1257    #[test]
1258    fn test_new_point() {
1259        let range_a = Range::new_point(1u8);
1260        let range_b = Range::new_point(2u16);
1261        let range_c = Range::new_point(3u32);
1262        let range_d = Range::new_point(4u64);
1263        let range_e = Range::new_point(5u32);
1264
1265        assert_eq!(range_a, Range::with_size(1u8, 0u8));
1266        assert_eq!(range_b, Range::with_size(2u16, 0u16));
1267        assert_eq!(range_c, Range::with_size(3u32, 0u32));
1268        assert_eq!(range_d, Range::with_size(4u64, 0u64));
1269        assert_eq!(range_e, Range::with_size(5u32, 0u32));
1270    }
1271
1272    #[test]
1273    fn test_get_by_id() {
1274        let mut tree = IntervalTree::<u32>::new();
1275        tree.insert(Range::new(0x100u16, 0x100u16), Some(1));
1276        tree.insert(Range::new(0x001u32, 0x005u32), Some(2));
1277        tree.insert(Range::new(0x200u16, 0x2ffu16), None);
1278
1279        assert_eq!(tree.get_by_id(0x100u16), Some(&1));
1280        assert_eq!(tree.get_by_id(0x002u32), Some(&2));
1281        assert_eq!(tree.get_by_id(0x210u32), None);
1282        assert_eq!(tree.get_by_id(0x2ffu64), None);
1283    }
1284
1285    #[test]
1286    fn test_get_by_id_mut() {
1287        let mut tree = IntervalTree::<u32>::new();
1288        tree.insert(Range::new(0x100u16, 0x100u16), Some(1));
1289        tree.insert(Range::new(0x001u32, 0x005u32), Some(2));
1290        tree.insert(Range::new(0x200u16, 0x2ffu16), None);
1291
1292        assert_eq!(tree.get_by_id_mut(0x100u16), Some(&mut 1));
1293        assert_eq!(tree.get_by_id_mut(0x002u32), Some(&mut 2));
1294        assert_eq!(tree.get_by_id_mut(0x210u32), None);
1295        assert_eq!(tree.get_by_id_mut(0x2ffu64), None);
1296    }
1297}
dbs_allocator/interval_tree.rs

dbs_allocator/
interval_tree.rs
