sum_tree: Implement find functions iteratively (#47013)

Lukas Wirth created 1 week ago

Recursion here is unnecessary as we do not make use of the stack frames,
so iterating is even cheaper as we do not need to do any stack
bookkeeping either way.

Release Notes:

- N/A *or* Added/Fixed/Improved ...

Change summary

crates/sum_tree/src/sum_tree.rs | 180 +++++++++++++++++-----------------
1 file changed, 89 insertions(+), 91 deletions(-)

Detailed changes

crates/sum_tree/src/sum_tree.rs

@@ -401,7 +401,7 @@ impl<T: Item> SumTree<T> {
         }
 
         let mut pos = D::zero(cx);
-        return match Self::find_recurse::<_, _, true>(cx, target, bias, &mut pos, self) {
+        return match Self::find_iterate::<_, _, true>(cx, target, bias, &mut pos, self) {
             Some((item, end)) => (pos, end, Some(item)),
             None => (pos.clone(), pos, None),
         };
@@ -427,68 +427,69 @@ impl<T: Item> SumTree<T> {
         }
 
         let mut pos = D::zero(cx);
-        return match Self::find_recurse::<_, _, false>(cx, target, bias, &mut pos, self) {
+        return match Self::find_iterate::<_, _, false>(cx, target, bias, &mut pos, self) {
             Some((item, end)) => (pos, end, Some(item)),
             None => (pos.clone(), pos, None),
         };
     }
 
-    fn find_recurse<'tree, 'a, D, Target, const EXACT: bool>(
+    fn find_iterate<'tree, 'a, D, Target, const EXACT: bool>(
         cx: <T::Summary as Summary>::Context<'a>,
         target: &Target,
         bias: Bias,
         position: &mut D,
-        this: &'tree SumTree<T>,
+        mut this: &'tree SumTree<T>,
     ) -> Option<(&'tree T, D)>
     where
         D: Dimension<'tree, T::Summary>,
         Target: SeekTarget<'tree, T::Summary, D>,
     {
-        match &*this.0 {
-            Node::Internal {
-                child_summaries,
-                child_trees,
-                ..
-            } => {
-                for (child_tree, child_summary) in child_trees.iter().zip(child_summaries) {
-                    let child_end = position.clone().with_added_summary(child_summary, cx);
-
-                    let comparison = target.cmp(&child_end, cx);
-                    let target_in_child = comparison == Ordering::Less
-                        || (comparison == Ordering::Equal && bias == Bias::Left);
-                    if target_in_child {
-                        return Self::find_recurse::<D, Target, EXACT>(
-                            cx, target, bias, position, child_tree,
-                        );
+        'iterate: loop {
+            match &*this.0 {
+                Node::Internal {
+                    child_summaries,
+                    child_trees,
+                    ..
+                } => {
+                    for (child_tree, child_summary) in child_trees.iter().zip(child_summaries) {
+                        let child_end = position.clone().with_added_summary(child_summary, cx);
+
+                        let comparison = target.cmp(&child_end, cx);
+                        let target_in_child = comparison == Ordering::Less
+                            || (comparison == Ordering::Equal && bias == Bias::Left);
+                        if target_in_child {
+                            this = child_tree;
+                            continue 'iterate;
+                        }
+                        *position = child_end;
                     }
-                    *position = child_end;
                 }
-            }
-            Node::Leaf {
-                items,
-                item_summaries,
-                ..
-            } => {
-                for (item, item_summary) in items.iter().zip(item_summaries) {
-                    let mut child_end = position.clone();
-                    child_end.add_summary(item_summary, cx);
+                Node::Leaf {
+                    items,
+                    item_summaries,
+                    ..
+                } => {
+                    for (item, item_summary) in items.iter().zip(item_summaries) {
+                        let mut child_end = position.clone();
+                        child_end.add_summary(item_summary, cx);
+
+                        let comparison = target.cmp(&child_end, cx);
+                        let entry_found = if EXACT {
+                            comparison == Ordering::Equal
+                        } else {
+                            comparison == Ordering::Less
+                                || (comparison == Ordering::Equal && bias == Bias::Left)
+                        };
+                        if entry_found {
+                            return Some((item, child_end));
+                        }
 
-                    let comparison = target.cmp(&child_end, cx);
-                    let entry_found = if EXACT {
-                        comparison == Ordering::Equal
-                    } else {
-                        comparison == Ordering::Less
-                            || (comparison == Ordering::Equal && bias == Bias::Left)
-                    };
-                    if entry_found {
-                        return Some((item, child_end));
+                        *position = child_end;
                     }
-
-                    *position = child_end;
                 }
             }
+            return None;
         }
-        None
     }
 
     /// A more efficient version of `Cursor::new()` + `Cursor::seek()` + `Cursor::item()`
@@ -511,75 +512,72 @@ impl<T: Item> SumTree<T> {
         }
 
         let mut pos = D::zero(cx);
-        return match Self::find_recurse_with_prev::<_, _, false>(
-            cx, target, bias, &mut pos, self, None,
-        ) {
+        return match Self::find_with_prev_iterate::<_, _, false>(cx, target, bias, &mut pos, self) {
             Some((prev, item, end)) => (pos, end, Some((prev, item))),
             None => (pos.clone(), pos, None),
         };
     }
 
-    fn find_recurse_with_prev<'tree, 'a, D, Target, const EXACT: bool>(
+    fn find_with_prev_iterate<'tree, 'a, D, Target, const EXACT: bool>(
         cx: <T::Summary as Summary>::Context<'a>,
         target: &Target,
         bias: Bias,
         position: &mut D,
-        this: &'tree SumTree<T>,
-        prev: Option<&'tree T>,
+        mut this: &'tree SumTree<T>,
     ) -> Option<(Option<&'tree T>, &'tree T, D)>
     where
         D: Dimension<'tree, T::Summary>,
         Target: SeekTarget<'tree, T::Summary, D>,
     {
-        match &*this.0 {
-            Node::Internal {
-                child_summaries,
-                child_trees,
-                ..
-            } => {
-                let mut prev = prev;
-                for (child_tree, child_summary) in child_trees.iter().zip(child_summaries) {
-                    let child_end = position.clone().with_added_summary(child_summary, cx);
-
-                    let comparison = target.cmp(&child_end, cx);
-                    let target_in_child = comparison == Ordering::Less
-                        || (comparison == Ordering::Equal && bias == Bias::Left);
-                    if target_in_child {
-                        return Self::find_recurse_with_prev::<D, Target, EXACT>(
-                            cx, target, bias, position, child_tree, prev,
-                        );
+        let mut prev = None;
+        'iterate: loop {
+            match &*this.0 {
+                Node::Internal {
+                    child_summaries,
+                    child_trees,
+                    ..
+                } => {
+                    for (child_tree, child_summary) in child_trees.iter().zip(child_summaries) {
+                        let child_end = position.clone().with_added_summary(child_summary, cx);
+
+                        let comparison = target.cmp(&child_end, cx);
+                        let target_in_child = comparison == Ordering::Less
+                            || (comparison == Ordering::Equal && bias == Bias::Left);
+                        if target_in_child {
+                            this = child_tree;
+                            continue 'iterate;
+                        }
+                        prev = child_tree.last();
+                        *position = child_end;
                     }
-                    prev = child_tree.last();
-                    *position = child_end;
                 }
-            }
-            Node::Leaf {
-                items,
-                item_summaries,
-                ..
-            } => {
-                let mut prev = prev;
-                for (item, item_summary) in items.iter().zip(item_summaries) {
-                    let mut child_end = position.clone();
-                    child_end.add_summary(item_summary, cx);
-
-                    let comparison = target.cmp(&child_end, cx);
-                    let entry_found = if EXACT {
-                        comparison == Ordering::Equal
-                    } else {
-                        comparison == Ordering::Less
-                            || (comparison == Ordering::Equal && bias == Bias::Left)
-                    };
-                    if entry_found {
-                        return Some((prev, item, child_end));
-                    }
+                Node::Leaf {
+                    items,
+                    item_summaries,
+                    ..
+                } => {
+                    for (item, item_summary) in items.iter().zip(item_summaries) {
+                        let mut child_end = position.clone();
+                        child_end.add_summary(item_summary, cx);
+
+                        let comparison = target.cmp(&child_end, cx);
+                        let entry_found = if EXACT {
+                            comparison == Ordering::Equal
+                        } else {
+                            comparison == Ordering::Less
+                                || (comparison == Ordering::Equal && bias == Bias::Left)
+                        };
+                        if entry_found {
+                            return Some((prev, item, child_end));
+                        }
 
-                    prev = Some(item);
-                    *position = child_end;
+                        prev = Some(item);
+                        *position = child_end;
+                    }
                 }
             }
+            return None;
         }
-        None
     }
 
     pub fn cursor<'a, 'b, D>(