fenwick_bit_tree/
lib.rs

1//! # Slighly over-engineered Fenwick Tree implmentation.
2//! 
3//! Allows efficient prefix sum calculation.
4//! 
5//! Created for training purposes to test: 
6//!
7//! 1. rust typesystem, default trait implmentation, enums as a way for polymorphism
8//! 2. memory management and consumption of value
9//! 3. cargo tools, docs, tests, clippy and benchmarks, build and publish.
10//!
11//! Code is free to do whatever you feel like.
12//! 
13//! Provides abstraction for Fenwick tree data structure and 2 implmentations:
14//!  - [`prelude::FixedSizeFenwickTree`]
15//!  - [`prelude::GrowingFenwickTree`]
16//! 
17//! Key space for a tree lies within [`usize`] range. Tree support any value that 
18//! implements [`FenwickTreeValue`] trait. [`FenwickTreeValue`] is automatically 
19//! implmented for all primitive numeric types that support [`std::ops::AddAssign`], 
20//! [`std::ops::Sub`], [`core::cmp::PartialEq`] and [`Copy`] traits.
21//!
22//! ## Installation  
23//!
24//! ```bash
25//! cargo add fenwick-bit-tree
26//! ```
27//! 
28//! ## Test
29//! 
30//! ```bash
31//! cargo test
32//! ```
33//! 
34//! ## Benchmarks
35//! 
36//! ```bash
37//! cargo bench --features benchmarks
38//! ```
39//! 
40//! ## Basic usage:
41//! 
42//! ```rust
43//! use fenwick_bit_tree::prelude::*;
44//! 
45//! // Create the tree with capacity for 32 aggregated [`i32`] data points. 
46//! // One can use whole usize range to store datapoints for unicode timestamps
47//! let mut tree = FixedSizeFenwickTree::<i32>::new(32);
48//!
49//! // Add values
50//! 
51//! tree.update(0, 1); 
52//! tree.update(0, 4); // Will aggregate value at index 0 so it would be 5
53//! tree.update(10, 10);
54//! tree.update(20, 10);
55//! tree.update(30, 10);
56//! 
57//! // Now you can query data. 
58//! // NOTE: FixedSizeFenwickTree will raise error when query goes out of bounds.
59//! //       GrowingFenwickTree will automatically truncate the range to the rightmost index. 
60//! 
61//! assert_eq!(tree.query(4).unwrap(), 5); 
62//! assert_eq!(tree.query(15).unwrap(), 15);
63//! assert_eq!(tree.query(31).unwrap(), 35);
64//!
65//! // Also allows making range queries
66//! 
67//! let val = tree.range_query(2, 16).unwrap(); // Will return aggregated sum of all values between those keys.
68//! assert_eq!(val, 10);
69//! ```
70
71#![forbid(unsafe_code)]
72#![feature(test)]
73
74use std::ops::{Deref, DerefMut};
75
76mod fixed_size_tree;
77mod growing_tree;
78
79pub use fixed_size_tree::FixedSizeFenwickTree;
80pub use growing_tree::GrowingFenwickTree;
81
82/// Contains all public types
83pub mod prelude {
84    pub use crate::FenwickTreeValue;
85    pub use crate::fixed_size_tree::FixedSizeFenwickTree;
86    pub use crate::growing_tree::GrowingFenwickTree;
87    pub use crate::FenwickTree;
88    pub use crate::TreeError;
89}
90
91fn least_significant_bit(idx: usize) -> usize {
92    let int_idx = idx as i32;
93    (int_idx & -int_idx) as usize
94}
95
96/// Types that implement that trait can be stored and aggregated within Fenwick tree.
97pub trait FenwickTreeValue:
98    Default + Clone //
99    + core::cmp::PartialEq 
100{
101    fn store_value(&mut self, other: &Self);
102    fn substract(self, other: Self) -> Self;
103}
104
105impl<T> FenwickTreeValue for T 
106where T: Default + Copy //
107    + std::ops::AddAssign
108    + std::ops::Sub<Output = Self>
109    + core::cmp::PartialEq 
110{
111    fn store_value(&mut self, other: &Self) {
112        *self += *other
113    }
114
115    fn substract(self, other: Self) -> Self {
116        self - other
117    }
118}
119
120/// Fenwick tree trait, API of that data structure
121pub trait FenwickTree {
122    type Value: FenwickTreeValue;
123
124    /// Returns sum of values across all indexes lesser or equal than `idx`.
125    ///
126    /// # Errors
127    ///
128    /// This function will returns an error if idx is out of bounds.
129    /// GrowingFenwick tree implementation never returns error.
130    /// 
131    fn query(&self, idx: usize) -> Result<Self::Value, TreeError>;
132    
133    /// Add new value to the `idx` stored value, which is 0 by default. 
134    ///
135    /// # Errors
136    ///
137    /// This function will return an error if idx is out of bounds.
138    /// GrowingFenwick tree implementation never returns error.
139    /// 
140    fn update(&mut self, idx: usize, value: Self::Value) -> Result<(), TreeError>;
141
142    /// Returns sum of values across all indexes in between `from` and `to` indexes 
143    /// (including edges).
144    ///
145    /// # Errors
146    ///
147    /// This function will return an error if any index is out of bounds.
148    /// GrowingFenwick tree implementation never return error.
149    /// 
150    fn range_query(&self, from: usize, to: usize) -> Result<Self::Value, TreeError> {
151        let from_sum = self.query(from)?;
152        let to_sum = self.query(to)?;
153        Ok(to_sum.substract(from_sum))
154    }
155}
156
157/// For the sake of clarity Tree supports 2 types of indexing. [`TreeIndex::External`] is meant to be used 
158/// by library consumer. While [`TreeIndex::Internal`] is used for purposes to make tree reindexing code more
159/// understable and maintainable. [`usize`] can be automatically converted using `into()` into the [`TreeIndex::External`]
160#[derive(Debug, Clone, Copy)]
161enum TreeIndex {
162    Internal { val: usize },
163    External { val: usize },
164}
165
166#[derive(Debug, PartialEq)]
167pub enum TreeError {
168    IndexOutOfBounds( usize )
169}
170
171impl TreeIndex {
172
173    fn to_internal(self) -> Self {
174        match self {
175            TreeIndex::Internal { val: _ } => self,
176            TreeIndex::External { val } => TreeIndex::Internal { val: val + 1 },
177        }
178    }
179
180    fn to_external(self) -> Result<Self, String> {
181        match self {
182            TreeIndex::Internal { val } => {
183                if val == 0 {
184                    return Err("Index is out of bounds.".to_string());
185                }
186                Ok(TreeIndex::External { val: val - 1 })
187            }
188            TreeIndex::External { val: _ } => Ok(self),
189        }
190    }
191
192    /// Starts with the initial value and then moves down to zero returning result of
193    /// deduction of the least significant bit
194    fn lsb_descending(self) -> LeastSignificantBitDescentingChain {
195        LeastSignificantBitDescentingChain {
196            idx: self.to_internal(),
197        }
198    }
199
200    /// Starts with the initial value and then moves up until upper bound is reached 
201    /// returning the result of deduction of the least significant bit
202    fn lsb_ascending(self, upper_bound: usize) -> LeastSignificantBitAscendingChain {
203        LeastSignificantBitAscendingChain {
204            idx: self.to_internal(),
205            max: upper_bound,
206        }
207    }
208
209    fn is_power_of_2(self) -> bool {
210        let idx = *self;
211        idx.is_power_of_two()
212    }
213
214}
215
216impl From<usize> for TreeIndex {
217    fn from(value: usize) -> Self {
218        Self::External { val: value }
219    }
220}
221
222impl Deref for TreeIndex {
223    type Target = usize;
224
225    fn deref(&self) -> &Self::Target {
226        match self {
227            TreeIndex::External { val } => val,
228            TreeIndex::Internal { val } => val,
229        }
230    }
231}
232
233impl PartialEq for TreeIndex {
234    fn eq(&self, other: &Self) -> bool {
235        match (self, other) {
236            (Self::Internal { val: l_val }, Self::Internal { val: r_val }) => l_val == r_val,
237            (Self::External { val: l_val }, Self::External { val: r_val }) => l_val == r_val,
238            _ => false,
239        }
240    }
241}
242
243impl DerefMut for TreeIndex {
244    fn deref_mut(&mut self) -> &mut Self::Target {
245        match self {
246            TreeIndex::External { val } => val,
247            TreeIndex::Internal { val } => val,
248        }
249    }
250}
251
252/// Iterator that implements changing value by deduction of the least significant bit and 
253/// returning result
254struct LeastSignificantBitDescentingChain {
255    idx: TreeIndex,
256}
257
258impl Iterator for LeastSignificantBitDescentingChain {
259    type Item = TreeIndex;
260
261    fn next(&mut self) -> Option<Self::Item> {
262        if *self.idx == 0 {
263            return None;
264        }
265        // TODO: implement COpy?
266        let res = TreeIndex::Internal { val: *self.idx };
267        *self.idx -= least_significant_bit(*self.idx);
268        Some(res)
269    }
270}
271
272/// Iterator that implements changing value by addition of the least significant bit and 
273/// returning result
274struct LeastSignificantBitAscendingChain {
275    idx: TreeIndex,
276    max: usize,
277}
278
279impl Iterator for LeastSignificantBitAscendingChain {
280    type Item = TreeIndex;
281
282    fn next(&mut self) -> Option<Self::Item> {
283        if *self.idx > self.max {
284            return None;
285        }
286        // TODO: implement COpy?
287        let res = TreeIndex::Internal { val: *self.idx };
288        *self.idx += least_significant_bit(*self.idx);
289        Some(res)
290    }
291}
292
293#[cfg(test)]
294mod tests {
295
296    use pretty_assertions::assert_eq;
297
298    use crate::{least_significant_bit, TreeIndex};
299
300    fn to_internal_index_vec(indexes: &[usize]) -> Vec<TreeIndex> {
301        indexes
302            .into_iter()
303            .map(|i| TreeIndex::Internal { val: *i })
304            .collect::<Vec<TreeIndex>>()
305    }
306
307    #[test]
308    fn test_index_transform_from_internal_to_external_with_error() {
309        let idx = TreeIndex::Internal { val: 0 };
310        idx.to_external().expect_err("Index is out of bounds.");
311    }
312
313    #[test]
314    fn test_index_transform_from_internal_to_external() {
315        for val in 1..100 {
316            let idx = TreeIndex::Internal { val: val };
317            assert_eq!(
318                idx.to_external().unwrap(),
319                TreeIndex::External { val: val - 1 }
320            );
321        }
322    }
323
324    #[test]
325    fn test_index_transform_from_external_to_internal() {
326        for val in 0..100 {
327            let idx = TreeIndex::External { val: val };
328            assert_eq!(idx.to_internal(), TreeIndex::Internal { val: val + 1 });
329        }
330    }
331
332    #[test]
333    fn test_index_transform_to_itseld() {
334        for val in 0..100 {
335            let idx = TreeIndex::External { val: val };
336            assert_eq!(idx.to_external().unwrap(), TreeIndex::External { val });
337        }
338
339        for val in 0..100 {
340            let idx = TreeIndex::Internal { val: val };
341            assert_eq!(idx.to_internal(), TreeIndex::Internal { val: val });
342        }
343    }
344
345    #[test]
346    fn test_ascending_lsb_chain() {
347        let idx: TreeIndex = 0.into();
348        assert_eq!(
349            idx.lsb_ascending(64).collect::<Vec<TreeIndex>>(),
350            to_internal_index_vec(&[1, 2, 4, 8, 16, 32, 64])
351        );
352
353        let idx: TreeIndex = 1.into();
354        assert_eq!(
355            idx.lsb_ascending(64).collect::<Vec<TreeIndex>>(),
356            to_internal_index_vec(&[2, 4, 8, 16, 32, 64])
357        );
358
359        let idx: TreeIndex = 6.into();
360        assert_eq!(
361            idx.lsb_ascending(64).collect::<Vec<TreeIndex>>(),
362            to_internal_index_vec(&[7, 8, 16, 32, 64])
363        );
364
365        let idx: TreeIndex = 6.into();
366        assert_eq!(idx.lsb_ascending(0).collect::<Vec<TreeIndex>>(), vec![]);
367    }
368
369    #[test]
370    fn test_descending_lsb_chain() {
371        let idx: TreeIndex = 5.into();
372        assert_eq!(idx, TreeIndex::External { val: 5 });
373        assert_eq!(
374            idx.lsb_descending().collect::<Vec<TreeIndex>>(),
375            to_internal_index_vec(&[6, 4])
376        );
377
378        let idx: TreeIndex = 4.into();
379        assert_eq!(
380            idx.lsb_descending().collect::<Vec<TreeIndex>>(),
381            to_internal_index_vec(&[5, 4])
382        );
383
384        let idx = TreeIndex::Internal { val: 3 };
385        assert_eq!(
386            idx.lsb_descending().collect::<Vec<TreeIndex>>(),
387            to_internal_index_vec(&[3, 2])
388        );
389
390        let idx = TreeIndex::Internal { val: 12 };
391        assert_eq!(
392            idx.lsb_descending().collect::<Vec<TreeIndex>>(),
393            to_internal_index_vec(&[12, 8])
394        );
395    }
396
397    #[test]
398    fn test_lsb() {
399        assert_eq!(least_significant_bit(12), 4)
400    }
401
402    #[test]
403    fn test_bitwise_op() {
404        assert_eq!(12usize.next_power_of_two(), 16);
405        assert_eq!(12usize.next_power_of_two() >> 1, 8);
406    }
407}