augmented_rbtree/
augmentations.rs

1//! Ready-to-use [`Augment`] implementations for common use cases.
2//!
3//! These types can be used directly as the `G` type parameter of
4//! [`AugmentedRBTree`](crate::AugmentedRBTree) without implementing the trait yourself.
5//!
6//! # Examples
7//!
8//! ```
9//! use augmented_rbtree::{AugmentedRBTree, augmentations::SubtreeSize};
10//!
11//! let mut tree = AugmentedRBTree::<i32, f64, SubtreeSize>::new();
12//! tree.insert(1, 1.0);
13//! tree.insert(2, 2.0);
14//! tree.insert(3, 3.0);
15//!
16//! assert_eq!(tree.root_stats(), Some(&3));
17//! ```
18
19use core::{marker::PhantomData, ops::Add};
20
21use crate::Augment;
22
23// ============================================================================
24// UnitAugmentation
25// ============================================================================
26
27/// Augmentation that does not store any additional data.
28#[derive(Debug)]
29pub struct Unit;
30
31impl<K, V> Augment<K, V> for Unit {
32    type Stats = ();
33
34    fn compute(
35        _key: &K,
36        _value: &V,
37        _left: Option<(&K, &V, &Self::Stats)>,
38        _right: Option<(&K, &V, &Self::Stats)>,
39    ) -> Self::Stats {
40    }
41}
42
43// ============================================================================
44// SubtreeSize
45// ============================================================================
46
47/// Augmentation that tracks the number of nodes in each subtree.
48///
49/// This enables O(log n) rank queries (find the k-th smallest element) and
50/// select operations when combined with custom traversal.
51///
52/// # Examples
53///
54/// ```
55/// use augmented_rbtree::{AugmentedRBTree, augmentations::SubtreeSize};
56///
57/// let mut tree = AugmentedRBTree::<i32, &str, SubtreeSize>::new();
58/// tree.insert(3, "c");
59/// tree.insert(1, "a");
60/// tree.insert(2, "b");
61///
62/// assert_eq!(tree.root_stats(), Some(&3));
63/// assert_eq!(tree.len(), 3);
64/// ```
65#[derive(Debug, Default)]
66pub struct SubtreeSize<S = usize>(PhantomData<S>);
67
68impl<K, V, S> Augment<K, V> for SubtreeSize<S>
69where
70    S: Add<Output = S> + From<usize> + Copy,
71{
72    type Stats = S;
73
74    fn compute(_key: &K, _value: &V, left: Option<(&K, &V, &S)>, right: Option<(&K, &V, &S)>) -> S {
75        S::from(1usize)
76            + left.map_or(S::from(0usize), |(_, _, &c)| c)
77            + right.map_or(S::from(0), |(_, _, &c)| c)
78    }
79}
80
81// ============================================================================
82// SumAugmentation
83// ============================================================================
84
85/// Augmentation that tracks the sum of all values in each subtree.
86///
87/// Useful for range-sum queries: after locating the range boundaries, access
88/// the subtree statistics to get aggregate sums.
89///
90/// The value type `V` must implement [`core::ops::Add`], [`Copy`], and provide a
91/// zero element via [`Default`].
92///
93/// # Examples
94///
95/// ```
96/// use augmented_rbtree::{AugmentedRBTree, augmentations::SumAugmentation};
97///
98/// let mut tree = AugmentedRBTree::<i32, i64, SumAugmentation>::new();
99/// tree.insert(1, 10);
100/// tree.insert(2, 20);
101/// tree.insert(3, 30);
102///
103/// // Root stats holds the sum of all values
104/// assert_eq!(tree.root_stats(), Some(&60));
105/// ```
106
107#[derive(Debug)]
108pub struct SumAugmentation;
109
110impl<K, V> Augment<K, V> for SumAugmentation
111where
112    V: core::ops::Add<Output = V> + Copy + Default,
113{
114    type Stats = V;
115
116    fn compute(_key: &K, value: &V, left: Option<(&K, &V, &V)>, right: Option<(&K, &V, &V)>) -> V {
117        let left_sum = left.map(|(_, _, &s)| s).unwrap_or_default();
118        let right_sum = right.map(|(_, _, &s)| s).unwrap_or_default();
119        left_sum + *value + right_sum
120    }
121}
122
123// ============================================================================
124// MaxAugmentation
125// ============================================================================
126
127/// Augmentation that tracks the maximum value in each subtree.
128///
129/// Useful for segment-tree-style range-max queries.
130///
131/// # Examples
132///
133/// ```
134/// use augmented_rbtree::{AugmentedRBTree, augmentations::MaxAugmentation};
135///
136/// let mut tree = AugmentedRBTree::<i32, i32, MaxAugmentation>::new();
137/// tree.insert(1, 5);
138/// tree.insert(2, 12);
139/// tree.insert(3, 3);
140///
141/// assert_eq!(tree.root_stats(), Some(&Some(12)));
142/// ```
143#[derive(Debug)]
144pub struct MaxAugmentation;
145
146impl<K, V> Augment<K, V> for MaxAugmentation
147where
148    V: Ord + Copy,
149{
150    type Stats = Option<V>;
151
152    fn compute(
153        _key: &K,
154        value: &V,
155        left: Option<(&K, &V, &Option<V>)>,
156        right: Option<(&K, &V, &Option<V>)>,
157    ) -> Option<V> {
158        let mut max = *value;
159        if let Some((_, _, Some(ls))) = left {
160            if *ls > max {
161                max = *ls;
162            }
163        }
164        if let Some((_, _, Some(rs))) = right {
165            if *rs > max {
166                max = *rs;
167            }
168        }
169        Some(max)
170    }
171}
172
173// ============================================================================
174// MinAugmentation
175// ============================================================================
176
177/// Augmentation that tracks the minimum value in each subtree.
178///
179/// # Examples
180///
181/// ```
182/// use augmented_rbtree::{AugmentedRBTree, augmentations::MinAugmentation};
183///
184/// let mut tree = AugmentedRBTree::<i32, i32, MinAugmentation>::new();
185/// tree.insert(1, 5);
186/// tree.insert(2, 1);
187/// tree.insert(3, 8);
188///
189/// assert_eq!(tree.root_stats(), Some(&Some(1)));
190/// ```
191#[derive(Debug)]
192pub struct MinAugmentation;
193
194impl<K, V> Augment<K, V> for MinAugmentation
195where
196    V: Ord + Copy,
197{
198    type Stats = Option<V>;
199
200    fn compute(
201        _key: &K,
202        value: &V,
203        left: Option<(&K, &V, &Option<V>)>,
204        right: Option<(&K, &V, &Option<V>)>,
205    ) -> Option<V> {
206        let mut min = *value;
207        if let Some((_, _, Some(ls))) = left {
208            if *ls < min {
209                min = *ls;
210            }
211        }
212        if let Some((_, _, Some(rs))) = right {
213            if *rs < min {
214                min = *rs;
215            }
216        }
217        Some(min)
218    }
219}
220
221// ============================================================================
222// IntervalMaxEnd
223// ============================================================================
224
225/// Augmentation for interval trees: tracks the maximum interval endpoint in each subtree.
226///
227/// When keys are interval start points and values are interval end points, this augmentation
228/// allows efficient overlap queries: if `root_stats()` < `query_start`, no interval overlaps.
229///
230/// # Examples
231///
232/// ```
233/// use augmented_rbtree::{AugmentedRBTree, augmentations::IntervalMaxEnd};
234///
235/// // Key = interval start, Value = interval end
236/// let mut tree = AugmentedRBTree::<i32, i32, IntervalMaxEnd>::new();
237/// tree.insert(1, 5);   // interval [1, 5]
238/// tree.insert(3, 10);  // interval [3, 10]
239/// tree.insert(8, 12);  // interval [8, 12]
240///
241/// // Max endpoint in the whole tree
242/// assert_eq!(tree.root_stats(), Some(&Some(12)));
243/// ```
244#[derive(Debug)]
245pub struct IntervalMaxEnd;
246
247impl<K> Augment<K, K> for IntervalMaxEnd
248where
249    K: Ord + Copy,
250{
251    type Stats = Option<K>;
252
253    fn compute(
254        _key: &K,
255        value: &K,
256        left: Option<(&K, &K, &Option<K>)>,
257        right: Option<(&K, &K, &Option<K>)>,
258    ) -> Option<K> {
259        let mut max_end = *value;
260        if let Some((_, _, Some(ls))) = left {
261            if *ls > max_end {
262                max_end = *ls;
263            }
264        }
265        if let Some((_, _, Some(rs))) = right {
266            if *rs > max_end {
267                max_end = *rs;
268            }
269        }
270        Some(max_end)
271    }
272}
273
274/// Dynamically generates a custom constant [`Augment`] type.
275///
276/// This eliminates the need to manage external traits or deal with
277/// const-generic limitations when configuring a tree with a fixed baseline value.
278///
279/// # Examples
280///
281/// ```
282/// # use augmented_rbtree::constant_augment;
283/// constant_augment!(MyConstantAugment, i32, 42);
284/// ```
285#[macro_export]
286macro_rules! constant_augment {
287    ($name:ident, $type:ty, $val:expr) => {
288        #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
289        pub struct $name;
290
291        impl $crate::Augment<$type, $type> for $name {
292            type Stats = $type;
293
294            #[inline]
295            fn compute(
296                _key: &$type,
297                _value: &$type,
298                _left: Option<(&$type, &$type, &Self::Stats)>,
299                _right: Option<(&$type, &$type, &Self::Stats)>,
300            ) -> Self::Stats {
301                $val
302            }
303        }
304    };
305}
augmented_rbtree/augmentations.rs

augmented_rbtree/
augmentations.rs
