Replace recursive tree traversal

.gitignore

@@ -1 +1,2 @@
-/target
+*sqlite
+target

banyan/Cargo.toml

@@ -19,6 +19,7 @@ lru = "0.6.4"
 maplit = "1.0.2"
 rand = "0.8.3"
 salsa20 = "0.7.2"
+smallvec = "1.6.1"
 tracing = "0.1.23"
 zstd = "0.6.0"
 

banyan/src/forest/mod.rs

@@ -15,6 +15,8 @@ mod read;
 mod stream;
 mod write;
 
+pub use read::ForestIter;
+
 pub type FutureResult<'a, T> = BoxFuture<'a, Result<T>>;
 
 #[derive(Debug, Clone)]

banyan/src/forest/read.rs

@@ -1,15 +1,9 @@
 use super::{BranchCache, Config, CryptoConfig, FilteredChunk, Forest, TreeTypes};
 use crate::{
-    index::deserialize_compressed,
-    index::zip_with_offset,
-    index::Branch,
-    index::BranchIndex,
-    index::CompactSeq,
-    index::Index,
-    index::IndexRef,
-    index::Leaf,
-    index::LeafIndex,
-    index::NodeInfo,
+    index::{
+        deserialize_compressed, zip_with_offset, Branch, BranchIndex, CompactSeq, Index, IndexRef,
+        Leaf, LeafIndex, NodeInfo,
+    },
     query::Query,
     store::ReadOnlyStore,
     util::{BoxedIter, IterExt},
@@ -19,7 +13,148 @@ use anyhow::{anyhow, Result};
 use core::fmt::Debug;
 use futures::{prelude::*, stream::BoxStream};
 use libipld::cbor::DagCbor;
+use smallvec::{smallvec, SmallVec};
+
 use std::{iter, sync::Arc, time::Instant};
+pub struct ForestIter<T: TreeTypes, V, R, Q: Query<T>, F> {
+    pub forest: Forest<T, V, R>,
+    pub offset: u64,
+    pub query: Q,
+    pub mk_extra: F,
+    pub index_stack: SmallVec<[Arc<Index<T>>; 64]>,
+    pub pos_stack: SmallVec<[(usize, SmallVec<[bool; 64]>); 32]>,
+}
+
+impl<T: TreeTypes, V, R, Q, E, F> Iterator for ForestIter<T, V, R, Q, F>
+where
+    T: TreeTypes + 'static,
+    V: DagCbor + Clone + Send + Sync + Debug + 'static,
+    R: ReadOnlyStore<T::Link> + Clone + Send + Sync + 'static,
+    Q: Query<T> + Clone + Send + 'static,
+    E: Send + 'static,
+    F: Fn(IndexRef<T>) -> E + Send + Sync + 'static,
+{
+    type Item = Result<FilteredChunk<T, V, E>>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let res: FilteredChunk<T, V, E> = 'outer: loop {
+            let head = match self.index_stack.last() {
+                Some(i) => i,
+                // Nothing to do ..
+                _ => return None,
+            };
+
+            let (pos, matching) = self.pos_stack.last_mut().expect("not empty");
+
+            // Branch is exhausted: Ascend.
+            if *pos == matching.len() {
+                // tracing::trace!("branch exhausted");
+
+                // Ascend to parent's node
+                self.index_stack.pop().expect("not empty");
+                self.pos_stack.pop();
+
+                // increase last stack ptr, if there is still something left to
+                // traverse
+                if !self.index_stack.is_empty() {
+                    let last = self.pos_stack.last_mut().expect("not empty");
+                    last.0 += 1;
+                }
+                continue;
+            }
+
+            match self.forest.load_node(head) {
+                Ok(NodeInfo::Branch(index, branch)) => {
+                    // we hit this branch node for the first time. Apply the
+                    // query on its children and store it
+                    if *pos == 0 {
+                        let mut q_matching = smallvec![true; index.summaries.len()];
+                        self.query.intersecting(self.offset, index, &mut q_matching);
+                        debug_assert_eq!(branch.children.len(), q_matching.len());
+                        let _ = std::mem::replace(matching, q_matching);
+                    }
+
+                    while !matching[*pos] {
+                        self.offset += branch.children[*pos].count();
+                        *pos += 1;
+
+                        // Branch exhausted, ascend and continue
+                        if *pos == branch.children.len() {
+                            // Ascend to parent's node
+                            self.index_stack.pop().expect("not empty");
+                            self.pos_stack.pop();
+                            let last = self.pos_stack.last_mut().expect("not empty");
+                            last.0 += 1;
+                            continue 'outer;
+                        }
+                    }
+
+                    // Descend into next child
+                    self.index_stack
+                        // TODO: clone :-( ?
+                        .push(Arc::new(branch.children[*pos].clone()));
+                    let new_vec: SmallVec<[_; 64]> = smallvec![matching[*pos]];
+                    self.pos_stack.push((0, new_vec));
+                    continue;
+                }
+
+                Ok(NodeInfo::Leaf(index, leaf)) => {
+                    let chunk = {
+                        let mut matching: SmallVec<[_; 32]> = smallvec![true; index.keys.len()];
+                        self.query.containing(self.offset, index, &mut matching);
+                        let keys = index.select_keys(&matching);
+                        let elems: Vec<V> = match leaf.as_ref().select(&matching) {
+                            Ok(i) => i,
+                            Err(e) => return Some(Err(e)),
+                        };
+                        let pairs = keys
+                            .zip(elems)
+                            .map(|((o, k), v)| (o + self.offset, k, v))
+                            .collect::<Vec<_>>();
+                        FilteredChunk {
+                            range: self.offset..self.offset + index.keys.count(),
+                            data: pairs,
+                            extra: (self.mk_extra)(IndexRef::Leaf(index)),
+                        }
+                    };
+                    self.offset += index.keys.count();
+
+                    // Ascend to parent's node
+                    self.index_stack.pop().expect("not empty");
+                    self.pos_stack.pop();
+                    let last = self.pos_stack.last_mut().expect("not empty");
+                    last.0 += 1;
+
+                    break chunk;
+                }
+
+                // even for purged leafs and branches or ignored chunks,
+                // produce a placeholder.
+                //
+                // the caller can find out if we skipped purged parts of the
+                // tree by using an appropriate mk_extra fn, or check
+                // `data.len()`.
+                Ok(_) => {
+                    // Ascend to parent's node
+                    let index = self.index_stack.pop().expect("not empty");
+                    self.pos_stack.pop();
+                    let last = self.pos_stack.last_mut().expect("Index stack not empty");
+                    last.0 += 1;
+                    self.offset += index.count();
+
+                    let placeholder: FilteredChunk<T, V, E> = FilteredChunk {
+                        range: self.offset..self.offset + index.count(),
+                        data: Vec::new(),
+                        extra: (self.mk_extra)(index.as_index_ref()),
+                    };
+                    break placeholder;
+                }
+                Err(e) => return Some(Err(e)),
+            };
+        };
+        Some(Ok(res))
+    }
+}
 
 /// basic random access append only tree
 impl<T, V, R> Forest<T, V, R>
@@ -305,8 +440,7 @@ where
         query: Q,
         index: Arc<Index<T>>,
     ) -> BoxedIter<'static, Result<(u64, T::Key, V)>> {
-        self.clone()
-            .iter_filtered_chunked0(offset, query, index, &|_| {})
+        self.traverse0(offset, query, index, &|_| {})
             .map(|res| match res {
                 Ok(chunk) => chunk.data.into_iter().map(Ok).left_iter(),
                 Err(cause) => iter::once(Err(cause)).right_iter(),
@@ -315,84 +449,6 @@ where
             .boxed()
     }
 
-    pub(crate) fn iter_filtered_chunked0<
-        Q: Query<T> + Clone + Send + 'static,
-        E: Send + 'static,
-        F: Fn(IndexRef<T>) -> E + Send + Sync + 'static,
-    >(
-        &self,
-        offset: u64,
-        query: Q,
-        index: Arc<Index<T>>,
-        mk_extra: &'static F,
-    ) -> BoxedIter<'static, Result<FilteredChunk<T, V, E>>> {
-        let this = self.clone();
-        let inner = || {
-            Ok(match self.load_node(&index)? {
-                NodeInfo::Leaf(index, node) => {
-                    // todo: don't get the node here, since we might not need it
-                    let mut matching = vec![true; index.keys.len()];
-                    query.containing(offset, index, &mut matching);
-                    let keys = index.select_keys(&matching);
-                    let elems: Vec<V> = node.as_ref().select(&matching)?;
-                    let pairs = keys
-                        .zip(elems)
-                        .map(|((o, k), v)| (o + offset, k, v))
-                        .collect::<Vec<_>>();
-                    let chunk = FilteredChunk {
-                        range: offset..offset + index.keys.count(),
-                        data: pairs,
-                        extra: mk_extra(IndexRef::Leaf(index)),
-                    };
-                    iter::once(Ok(chunk)).left_iter().left_iter()
-                }
-                NodeInfo::Branch(index, node) => {
-                    // todo: don't get the node here, since we might not need it
-                    let mut matching = vec![true; index.summaries.len()];
-                    query.intersecting(offset, index, &mut matching);
-                    let offsets = zip_with_offset(node.children.to_vec(), offset);
-                    let iter = matching.into_iter().zip(offsets).map(
-                        move |(is_matching, (child, offset))| {
-                            if is_matching {
-                                this.iter_filtered_chunked0(
-                                    offset,
-                                    query.clone(),
-                                    Arc::new(child),
-                                    mk_extra,
-                                )
-                                .left_iter()
-                            } else {
-                                let placeholder = FilteredChunk {
-                                    range: offset..offset + child.count(),
-                                    data: Vec::new(),
-                                    extra: mk_extra(child.as_index_ref()),
-                                };
-                                iter::once(Ok(placeholder)).right_iter()
-                            }
-                        },
-                    );
-                    iter.flatten().right_iter().left_iter()
-                }
-                NodeInfo::PurgedBranch(_) | NodeInfo::PurgedLeaf(_) => {
-                    // even for purged leafs and branches, produce a placeholder.
-                    //
-                    // the caller can find out if we skipped purged parts of the tree by
-                    // using an appropriate mk_extra fn.
-                    let placeholder = FilteredChunk {
-                        range: offset..offset + index.count(),
-                        data: Vec::new(),
-                        extra: mk_extra(index.as_index_ref()),
-                    };
-                    iter::once(Ok(placeholder)).right_iter()
-                }
-            })
-        };
-        match inner() {
-            Ok(iter) => iter.boxed(),
-            Err(cause) => iter::once(Err(cause)).boxed(),
-        }
-    }
-
     pub(crate) fn stream_filtered_chunked_reverse0<
         Q: Query<T> + Clone + Send + 'static,
         E: Send + 'static,
@@ -475,53 +531,25 @@ where
 
     pub(crate) fn dump0(&self, index: &Index<T>, prefix: &str) -> Result<()> {
         match self.load_node(index)? {
-            NodeInfo::Leaf(index, _) => {
-                println!(
-                    "{}Leaf(count={}, value_bytes={}, sealed={}, link={})",
-                    prefix,
-                    index.keys.count(),
-                    index.value_bytes,
-                    index.sealed,
-                    index
-                        .link
-                        .map(|x| format!("{}", x))
-                        .unwrap_or_else(|| "".to_string())
-                );
-            }
-
             NodeInfo::Branch(index, branch) => {
                 println!(
-                    "{}Branch(count={}, key_bytes={}, value_bytes={}, sealed={}, link={})",
-                    prefix,
+                    "Branch(count={}, key_bytes={}, value_bytes={}, sealed={}, link={}, children={})",
                     index.count,
                     index.key_bytes,
                     index.value_bytes,
                     index.sealed,
                     index
                         .link
                         .map(|x| format!("{}", x))
-                        .unwrap_or_else(|| "".to_string())
-                );
+                        .unwrap_or_else(|| "".to_string()),
+                    branch.children.len()
+            );
                 let prefix = prefix.to_string() + "  ";
                 for x in branch.children.iter() {
                     self.dump0(x, &prefix)?;
                 }
             }
-            NodeInfo::PurgedBranch(index) => {
-                println!(
-                    "{}PurgedBranch(count={}, key_bytes={}, value_bytes={}, sealed={})",
-                    prefix, index.count, index.key_bytes, index.value_bytes, index.sealed,
-                );
-            }
-            NodeInfo::PurgedLeaf(index) => {
-                println!(
-                    "{}PurgedLeaf(count={}, key_bytes={}, sealed={})",
-                    prefix,
-                    index.keys.count(),
-                    index.value_bytes,
-                    index.sealed,
-                );
-            }
+            x => println!("{}{}", prefix, x),
         };
         Ok(())
     }

banyan/src/forest/stream.rs

@@ -112,7 +112,7 @@ impl<
                 let offset = offset.clone();
                 let iter = forest
                     .clone()
-                    .iter_filtered_chunked0(0, query, index, mk_extra)
+                    .traverse0(0, query, index, mk_extra)
                     .take_while(move |result| {
                         if let Ok(chunk) = result {
                             // update the offset