diff --git a/lib/src/hypertree.rs b/lib/src/hypertree.rs index 7811ef577..30c01c63e 100644 --- a/lib/src/hypertree.rs +++ b/lib/src/hypertree.rs @@ -2,7 +2,7 @@ use std::fmt::Display; use bytemuck::{Pod, Zeroable}; -use crate::DataIndex; +use crate::{DataIndex, RedBlackTreeRangeIterator}; // Set to less than DataIndex::MAX because formal verification required it. It // would be better to set it fully to DataIndex::MAX, but not a major concern @@ -62,6 +62,7 @@ impl Payload fo // work with it in the same trait. pub trait HyperTreeReadOperations<'a> { fn lookup_index(&'a self, value: &V) -> DataIndex; + fn range(&'a self, min: &V, max: &V) -> RedBlackTreeRangeIterator<'a, Self, V>; fn lookup_max_index(&'a self) -> DataIndex; fn get_max_index(&self) -> DataIndex; fn get_root_index(&self) -> DataIndex; diff --git a/lib/src/red_black_tree.rs b/lib/src/red_black_tree.rs index e87dbeba3..4f3c40eaf 100644 --- a/lib/src/red_black_tree.rs +++ b/lib/src/red_black_tree.rs @@ -133,6 +133,47 @@ impl<'a, V: Payload> RedBlackTreeReadOnly<'a, V> { } } +pub struct RedBlackTreeRangeIterator<'a, T: GetRedBlackTreeReadOnlyData<'a>, V: Payload> { + pub(crate) tree: &'a T, + pub(crate) min: &'a V, + pub(crate) max: &'a V, + pub(crate) current_index: DataIndex, + pub(crate) phantom: std::marker::PhantomData<&'a V>, +} + +impl<'a, T, V> Iterator for RedBlackTreeRangeIterator<'a, T, V> +where + T: GetRedBlackTreeReadOnlyData<'a> + HyperTreeReadOperations<'a>, + V: Payload, +{ + type Item = (DataIndex, &'a V); + + fn next(&mut self) -> Option { + while self.current_index != NIL { + let current_value = self.tree.get_value::(self.current_index); + + if current_value >= self.min && current_value <= self.max { + // Store the current index to return + let result_index = self.current_index; + let result_value = current_value; + + // Move to the next lower index + self.current_index = self.tree.get_next_lower_index::(self.current_index); + + return Some((result_index, result_value)); + } + + // If the current value is out of range, move to the next relevant node + if current_value < self.min { + self.current_index = self.tree.get_next_higher_index::(self.current_index); + } else { + self.current_index = self.tree.get_next_lower_index::(self.current_index); + } + } + None + } +} + // Specific to red black trees and not all data structures. Implementing this // gets a lot of other stuff for free. pub trait GetRedBlackTreeReadOnlyData<'a> { @@ -610,6 +651,21 @@ where current_index } + fn range(&'a self, min: &V, max: &V) -> RedBlackTreeRangeIterator<'a, T, V> { + let root = self.get_root_index(); + RedBlackTreeRangeIterator { + tree: self, + min, + max, + current_index: if root != NIL { + self.lookup_index(min) + } else { + NIL + }, + phantom: std::marker::PhantomData, + } + } + fn lookup_max_index(&'a self) -> DataIndex { let mut current_index = self.root_index(); if current_index == NIL { @@ -2734,3 +2790,26 @@ pub(crate) mod test { tree.lookup_index(&TestOrder2::new(1_000, 7890)); } } + +#[test] +fn test_hypertree_range_query() { + let mut data: [u8; 100000] = [0; 100000]; + let tree = RedBlackTree::new(&mut data, NIL, NIL); + let hypertree = HyperTree::new(tree); + + hypertree.insert(0, TestOrderBid::new(1000)); + hypertree.insert(1, TestOrderBid::new(2000)); + hypertree.insert(2, TestOrderBid::new(3000)); + hypertree.insert(3, TestOrderBid::new(4000)); + hypertree.insert(4, TestOrderBid::new(5000)); + + let range_min = TestOrderBid::new(2000); + let range_max = TestOrderBid::new(4000); + + let results: Vec<_> = hypertree.range(&range_min, &range_max).collect(); + + assert_eq!(results.len(), 3); + assert_eq!(results[0].1.order_id, 2000); + assert_eq!(results[1].1.order_id, 3000); + assert_eq!(results[2].1.order_id, 4000); +} \ No newline at end of file