Re-implement set based algorithm for attribute criterion

milli/src/search/criteria/attribute.rs

@@ -1,7 +1,7 @@
-use std::borrow::Cow;
 use std::cmp::{self, Ordering};
 use std::collections::binary_heap::PeekMut;
 use std::collections::{btree_map, BTreeMap, BinaryHeap, HashMap};
+use std::iter::Peekable;
 use std::mem::take;
 
 use roaring::RoaringBitmap;
@@ -17,10 +17,6 @@ use crate::{Result, TreeLevel};
 /// We chose the LCM of all numbers between 1 and 10 as the multiplier (https://en.wikipedia.org/wiki/Least_common_multiple).
 const LCM_10_FIRST_NUMBERS: u32 = 2520;
 
-/// To compute the interval size of a level,
-/// we use 4 as the exponentiation base and the level as the exponent.
-const LEVEL_EXPONENTIATION_BASE: u32 = 4;
-
 /// Threshold on the number of candidates that will make
 /// the system to choose between one algorithm or another.
 const CANDIDATES_THRESHOLD: u64 = 1000;
@@ -32,7 +28,8 @@ pub struct Attribute<'t> {
     state: Option<(Operation, FlattenedQueryTree, RoaringBitmap)>,
     bucket_candidates: RoaringBitmap,
     parent: Box<dyn Criterion + 't>,
-    current_buckets: Option<btree_map::IntoIter<u64, RoaringBitmap>>,
+    linear_buckets: Option<btree_map::IntoIter<u64, RoaringBitmap>>,
+    set_buckets: Option<BinaryHeap<Branch<'t>>>,
 }
 
 impl<'t> Attribute<'t> {
@@ -42,7 +39,8 @@ impl<'t> Attribute<'t> {
             state: None,
             bucket_candidates: RoaringBitmap::new(),
             parent,
-            current_buckets: None,
+            linear_buckets: None,
+            set_buckets: None,
         }
     }
 }
@@ -67,19 +65,19 @@ impl<'t> Criterion for Attribute<'t> {
                 }
                 Some((query_tree, flattened_query_tree, mut allowed_candidates)) => {
                     let found_candidates = if allowed_candidates.len() < CANDIDATES_THRESHOLD {
-                        let current_buckets = match self.current_buckets.as_mut() {
+                        let linear_buckets = match self.linear_buckets.as_mut() {
-                            Some(current_buckets) => current_buckets,
+                            Some(linear_buckets) => linear_buckets,
                             None => {
-                                let new_buckets = linear_compute_candidates(
+                                let new_buckets = initialize_linear_buckets(
                                     self.ctx,
                                     &flattened_query_tree,
                                     &allowed_candidates,
                                 )?;
-                                self.current_buckets.get_or_insert(new_buckets.into_iter())
+                                self.linear_buckets.get_or_insert(new_buckets.into_iter())
                             }
                         };
 
-                        match current_buckets.next() {
+                        match linear_buckets.next() {
                             Some((_score, candidates)) => candidates,
                             None => {
                                 return Ok(Some(CriterionResult {
@@ -91,13 +89,21 @@ impl<'t> Criterion for Attribute<'t> {
                             }
                         }
                     } else {
-                        match set_compute_candidates(
+                        let mut set_buckets = match self.set_buckets.as_mut() {
-                            self.ctx,
+                            Some(set_buckets) => set_buckets,
-                            &flattened_query_tree,
+                            None => {
-                            &allowed_candidates,
+                                let new_buckets = initialize_set_buckets(
-                            params.wdcache,
+                                    self.ctx,
-                        )? {
+                                    &flattened_query_tree,
-                            Some(candidates) => candidates,
+                                    &allowed_candidates,
+                                    params.wdcache,
+                                )?;
+                                self.set_buckets.get_or_insert(new_buckets)
+                            }
+                        };
+
+                        match set_compute_candidates(&mut set_buckets, &allowed_candidates)? {
+                            Some((_score, candidates)) => candidates,
                             None => {
                                 return Ok(Some(CriterionResult {
                                     query_tree: Some(query_tree),
@@ -148,7 +154,7 @@ impl<'t> Criterion for Attribute<'t> {
                         }
 
                         self.state = Some((query_tree, flattened_query_tree, candidates));
-                        self.current_buckets = None;
+                        self.linear_buckets = None;
                     }
                     Some(CriterionResult {
                         query_tree: None,
@@ -170,142 +176,52 @@ impl<'t> Criterion for Attribute<'t> {
     }
 }
 
-/// WordLevelIterator is an pseudo-Iterator over intervals of word-position for one word,
-/// it will begin at the first non-empty interval and will return every interval without
-/// jumping over empty intervals.
-struct WordLevelIterator<'t, 'q> {
-    inner: Box<
-        dyn Iterator<Item = heed::Result<((&'t str, TreeLevel, u32, u32), RoaringBitmap)>> + 't,
-    >,
-    level: TreeLevel,
-    interval_size: u32,
-    word: Cow<'q, str>,
-    in_prefix_cache: bool,
-    inner_next: Option<(u32, u32, RoaringBitmap)>,
-    current_interval: Option<(u32, u32)>,
-}
-
-impl<'t, 'q> WordLevelIterator<'t, 'q> {
-    fn new(
-        ctx: &'t dyn Context<'t>,
-        word: Cow<'q, str>,
-        in_prefix_cache: bool,
-    ) -> heed::Result<Option<Self>> {
-        match ctx.word_position_last_level(&word, in_prefix_cache)? {
-            Some(_) => {
-                // HOTFIX Meilisearch#1707: it is better to only iterate over level 0 for performances reasons.
-                let level = TreeLevel::min_value();
-                let interval_size = LEVEL_EXPONENTIATION_BASE.pow(Into::<u8>::into(level) as u32);
-                let inner =
-                    ctx.word_position_iterator(&word, level, in_prefix_cache, None, None)?;
-                Ok(Some(Self {
-                    inner,
-                    level,
-                    interval_size,
-                    word,
-                    in_prefix_cache,
-                    inner_next: None,
-                    current_interval: None,
-                }))
-            }
-            None => Ok(None),
-        }
-    }
-
-    fn dig(
-        &self,
-        ctx: &'t dyn Context<'t>,
-        level: &TreeLevel,
-        left_interval: Option<u32>,
-    ) -> heed::Result<Self> {
-        let level = *level.min(&self.level);
-        let interval_size = LEVEL_EXPONENTIATION_BASE.pow(Into::<u8>::into(level) as u32);
-        let word = self.word.clone();
-        let in_prefix_cache = self.in_prefix_cache;
-        let inner =
-            ctx.word_position_iterator(&word, level, in_prefix_cache, left_interval, None)?;
-
-        Ok(Self {
-            inner,
-            level,
-            interval_size,
-            word,
-            in_prefix_cache,
-            inner_next: None,
-            current_interval: None,
-        })
-    }
-
-    fn next(&mut self) -> heed::Result<Option<(u32, u32, RoaringBitmap)>> {
-        fn is_next_interval(last_right: u32, next_left: u32) -> bool {
-            last_right + 1 == next_left
-        }
-
-        let inner_next = match self.inner_next.take() {
-            Some(inner_next) => Some(inner_next),
-            None => self
-                .inner
-                .next()
-                .transpose()?
-                .map(|((_, _, left, right), docids)| (left, right, docids)),
-        };
-
-        match inner_next {
-            Some((left, right, docids)) => match self.current_interval {
-                Some((last_left, last_right)) if !is_next_interval(last_right, left) => {
-                    let blank_left = last_left + self.interval_size;
-                    let blank_right = last_right + self.interval_size;
-                    self.current_interval = Some((blank_left, blank_right));
-                    self.inner_next = Some((left, right, docids));
-                    Ok(Some((blank_left, blank_right, RoaringBitmap::new())))
-                }
-                _ => {
-                    self.current_interval = Some((left, right));
-                    Ok(Some((left, right, docids)))
-                }
-            },
-            None => Ok(None),
-        }
-    }
-}
-
 /// QueryLevelIterator is an pseudo-Iterator for a Query,
 /// It contains WordLevelIterators and is chainned with other QueryLevelIterator.
-struct QueryLevelIterator<'t, 'q> {
+struct QueryLevelIterator<'t> {
-    parent: Option<Box<QueryLevelIterator<'t, 'q>>>,
+    inner: Vec<
-    inner: Vec<WordLevelIterator<'t, 'q>>,
+        Peekable<
-    level: TreeLevel,
+            Box<
-    accumulator: Vec<Option<(u32, u32, RoaringBitmap)>>,
+                dyn Iterator<Item = heed::Result<((&'t str, TreeLevel, u32, u32), RoaringBitmap)>>
-    parent_accumulator: Vec<Option<(u32, u32, RoaringBitmap)>>,
+                    + 't,
-    interval_to_skip: usize,
+            >,
+        >,
+    >,
 }
 
-impl<'t, 'q> QueryLevelIterator<'t, 'q> {
+impl<'t> QueryLevelIterator<'t> {
     fn new(
         ctx: &'t dyn Context<'t>,
-        queries: &'q [Query],
+        queries: &[Query],
         wdcache: &mut WordDerivationsCache,
-    ) -> Result<Option<Self>> {
+    ) -> Result<Self> {
         let mut inner = Vec::with_capacity(queries.len());
         for query in queries {
+            let in_prefix_cache = query.prefix && ctx.in_prefix_cache(query.kind.word());
             match &query.kind {
                 QueryKind::Exact { word, .. } => {
-                    if !query.prefix || ctx.in_prefix_cache(&word) {
+                    if !query.prefix || in_prefix_cache {
-                        let word = Cow::Borrowed(query.kind.word());
+                        let iter = ctx.word_position_iterator(
-                        if let Some(word_level_iterator) =
+                            query.kind.word(),
-                            WordLevelIterator::new(ctx, word, query.prefix)?
+                            TreeLevel::min_value(),
-                        {
+                            in_prefix_cache,
-                            inner.push(word_level_iterator);
+                            None,
-                        }
+                            None,
+                        )?;
+
+                        inner.push(iter.peekable());
                     } else {
                         for (word, _) in word_derivations(&word, true, 0, ctx.words_fst(), wdcache)?
                         {
-                            let word = Cow::Owned(word.to_owned());
+                            let iter = ctx.word_position_iterator(
-                            if let Some(word_level_iterator) =
+                                &word,
-                                WordLevelIterator::new(ctx, word, false)?
+                                TreeLevel::min_value(),
-                            {
+                                in_prefix_cache,
-                                inner.push(word_level_iterator);
+                                None,
-                            }
+                                None,
+                            )?;
+
+                            inner.push(iter.peekable());
                         }
                     }
                 }
@@ -313,360 +229,255 @@ impl<'t, 'q> QueryLevelIterator<'t, 'q> {
                     for (word, _) in
                         word_derivations(&word, query.prefix, *typo, ctx.words_fst(), wdcache)?
                     {
-                        let word = Cow::Owned(word.to_owned());
+                        let iter = ctx.word_position_iterator(
-                        if let Some(word_level_iterator) = WordLevelIterator::new(ctx, word, false)?
+                            &word,
-                        {
+                            TreeLevel::min_value(),
-                            inner.push(word_level_iterator);
+                            in_prefix_cache,
-                        }
+                            None,
+                            None,
+                        )?;
+
+                        inner.push(iter.peekable());
                     }
                 }
-            }
+            };
         }
 
-        let highest = inner.iter().max_by_key(|wli| wli.level).map(|wli| wli.level);
+        Ok(Self { inner })
-        match highest {
-            Some(level) => Ok(Some(Self {
-                parent: None,
-                inner,
-                level,
-                accumulator: vec![],
-                parent_accumulator: vec![],
-                interval_to_skip: 0,
-            })),
-            None => Ok(None),
-        }
-    }
-
-    fn parent(&mut self, parent: QueryLevelIterator<'t, 'q>) -> &Self {
-        self.parent = Some(Box::new(parent));
-        self
-    }
-
-    /// create a new QueryLevelIterator with a lower level than the current one.
-    fn dig(&self, ctx: &'t dyn Context<'t>) -> heed::Result<Self> {
-        let (level, parent) = match &self.parent {
-            Some(parent) => {
-                let parent = parent.dig(ctx)?;
-                (parent.level.min(self.level), Some(Box::new(parent)))
-            }
-            None => (self.level.saturating_sub(1), None),
-        };
-
-        let left_interval = self
-            .accumulator
-            .get(self.interval_to_skip)
-            .map(|opt| opt.as_ref().map(|(left, _, _)| *left))
-            .flatten();
-        let mut inner = Vec::with_capacity(self.inner.len());
-        for word_level_iterator in self.inner.iter() {
-            inner.push(word_level_iterator.dig(ctx, &level, left_interval)?);
-        }
-
-        Ok(Self {
-            parent,
-            inner,
-            level,
-            accumulator: vec![],
-            parent_accumulator: vec![],
-            interval_to_skip: 0,
-        })
-    }
-
-    fn inner_next(&mut self, level: TreeLevel) -> heed::Result<Option<(u32, u32, RoaringBitmap)>> {
-        let mut accumulated: Option<(u32, u32, RoaringBitmap)> = None;
-        let u8_level = Into::<u8>::into(level);
-        let interval_size = LEVEL_EXPONENTIATION_BASE.pow(u8_level as u32);
-        for wli in self.inner.iter_mut() {
-            let wli_u8_level = Into::<u8>::into(wli.level);
-            let accumulated_count = LEVEL_EXPONENTIATION_BASE.pow((u8_level - wli_u8_level) as u32);
-            for _ in 0..accumulated_count {
-                if let Some((next_left, _, next_docids)) = wli.next()? {
-                    accumulated = match accumulated.take() {
-                        Some((acc_left, acc_right, mut acc_docids)) => {
-                            acc_docids |= next_docids;
-                            Some((acc_left, acc_right, acc_docids))
-                        }
-                        None => Some((next_left, next_left + interval_size, next_docids)),
-                    };
-                }
-            }
-        }
-
-        Ok(accumulated)
-    }
-
-    /// return the next meta-interval created from inner WordLevelIterators,
-    /// and from eventual chainned QueryLevelIterator.
-    fn next(
-        &mut self,
-        allowed_candidates: &RoaringBitmap,
-        tree_level: TreeLevel,
-    ) -> heed::Result<Option<(u32, u32, RoaringBitmap)>> {
-        let parent_result = match self.parent.as_mut() {
-            Some(parent) => Some(parent.next(allowed_candidates, tree_level)?),
-            None => None,
-        };
-
-        match parent_result {
-            Some(parent_next) => {
-                let inner_next = self.inner_next(tree_level)?;
-                self.interval_to_skip += interval_to_skip(
-                    &self.parent_accumulator,
-                    &self.accumulator,
-                    self.interval_to_skip,
-                    allowed_candidates,
-                );
-                self.accumulator.push(inner_next);
-                self.parent_accumulator.push(parent_next);
-                let mut merged_interval: Option<(u32, u32, RoaringBitmap)> = None;
-
-                for current in self
-                    .accumulator
-                    .iter()
-                    .rev()
-                    .zip(self.parent_accumulator.iter())
-                    .skip(self.interval_to_skip)
-                {
-                    if let (Some((left_a, right_a, a)), Some((left_b, right_b, b))) = current {
-                        match merged_interval.as_mut() {
-                            Some((_, _, merged_docids)) => *merged_docids |= a & b,
-                            None => {
-                                merged_interval = Some((left_a + left_b, right_a + right_b, a & b))
-                            }
-                        }
-                    }
-                }
-                Ok(merged_interval)
-            }
-            None => {
-                let level = self.level;
-                match self.inner_next(level)? {
-                    Some((left, right, mut candidates)) => {
-                        self.accumulator = vec![Some((left, right, RoaringBitmap::new()))];
-                        candidates &= allowed_candidates;
-                        Ok(Some((left, right, candidates)))
-                    }
-                    None => {
-                        self.accumulator = vec![None];
-                        Ok(None)
-                    }
-                }
-            }
-        }
     }
 }
 
-/// Count the number of interval that can be skiped when we make the cross-intersections
+impl<'t> Iterator for QueryLevelIterator<'t> {
-/// in order to compute the next meta-interval.
+    type Item = heed::Result<(u32, RoaringBitmap)>;
-/// A pair of intervals is skiped when both intervals doesn't contain any allowed docids.
+
-fn interval_to_skip(
+    fn next(&mut self) -> Option<Self::Item> {
-    parent_accumulator: &[Option<(u32, u32, RoaringBitmap)>],
+        // sort inner words from the closest next position to the more far next position.
-    current_accumulator: &[Option<(u32, u32, RoaringBitmap)>],
+        let expected_pos = self
-    already_skiped: usize,
+            .inner
-    allowed_candidates: &RoaringBitmap,
+            .iter_mut()
-) -> usize {
+            .filter_map(|wli| match wli.peek() {
-    parent_accumulator
+                Some(Ok(((_, _, pos, _), _))) => Some(*pos),
-        .iter()
+                _ => None,
-        .zip(current_accumulator.iter())
+            })
-        .skip(already_skiped)
+            .min()?;
-        .take_while(|(parent, current)| {
+
-            let skip_parent = parent.as_ref().map_or(true, |(_, _, docids)| docids.is_empty());
+        let mut candidates = None;
-            let skip_current = current
+        for wli in self.inner.iter_mut() {
-                .as_ref()
+            if let Some(Ok(((_, _, pos, _), _))) = wli.peek() {
-                .map_or(true, |(_, _, docids)| docids.is_disjoint(allowed_candidates));
+                if *pos > expected_pos {
-            skip_parent && skip_current
+                    continue;
-        })
+                }
-        .count()
+            }
+
+            match wli.next() {
+                Some(Ok((_, docids))) => {
+                    candidates = match candidates.take() {
+                        Some(candidates) => Some(candidates | docids),
+                        None => Some(docids),
+                    }
+                }
+                Some(Err(e)) => return Some(Err(e)),
+                None => continue,
+            }
+        }
+
+        candidates.map(|candidates| Ok((expected_pos, candidates)))
+    }
 }
 
 /// A Branch is represent a possible alternative of the original query and is build with the Query Tree,
 /// This branch allows us to iterate over meta-interval of position and to dig in it if it contains interesting candidates.
-struct Branch<'t, 'q> {
+struct Branch<'t> {
-    query_level_iterator: QueryLevelIterator<'t, 'q>,
+    query_level_iterator: Vec<(u32, RoaringBitmap, Peekable<QueryLevelIterator<'t>>)>,
-    last_result: (u32, u32, RoaringBitmap),
+    last_result: (u32, RoaringBitmap),
-    tree_level: TreeLevel,
     branch_size: u32,
 }
 
-impl<'t, 'q> Branch<'t, 'q> {
+impl<'t> Branch<'t> {
+    fn new(
+        ctx: &'t dyn Context<'t>,
+        flatten_branch: &[Vec<Query>],
+        wdcache: &mut WordDerivationsCache,
+        allowed_candidates: &RoaringBitmap,
+    ) -> Result<Self> {
+        let mut query_level_iterator = Vec::new();
+        for queries in flatten_branch {
+            let mut qli = QueryLevelIterator::new(ctx, queries, wdcache)?.peekable();
+            let (pos, docids) = qli.next().transpose()?.unwrap_or((0, RoaringBitmap::new()));
+            query_level_iterator.push((pos, docids & allowed_candidates, qli));
+        }
+
+        let mut branch = Self {
+            query_level_iterator,
+            last_result: (0, RoaringBitmap::new()),
+            branch_size: flatten_branch.len() as u32,
+        };
+
+        branch.update_last_result();
+
+        Ok(branch)
+    }
+
     /// return the next meta-interval of the branch,
     /// and update inner interval in order to be ranked by the BinaryHeap.
     fn next(&mut self, allowed_candidates: &RoaringBitmap) -> heed::Result<bool> {
-        let tree_level = self.query_level_iterator.level;
+        // update the first query.
-        match self.query_level_iterator.next(allowed_candidates, tree_level)? {
+        let index = self.lowest_iterator_index();
-            Some(last_result) => {
+        match self.query_level_iterator.get_mut(index) {
-                self.last_result = last_result;
+            Some((cur_pos, cur_docids, qli)) => match qli.next().transpose()? {
-                self.tree_level = tree_level;
+                Some((next_pos, next_docids)) => {
-                Ok(true)
+                    *cur_pos = next_pos;
-            }
+                    *cur_docids |= next_docids & allowed_candidates;
-            None => Ok(false),
+                }
+                None => return Ok(false),
+            },
+            None => return Ok(false),
         }
+
+        self.update_last_result();
+
+        Ok(true)
     }
 
-    /// make the current Branch iterate over smaller intervals.
+    fn lowest_iterator_index(&mut self) -> usize {
-    fn dig(&mut self, ctx: &'t dyn Context<'t>) -> heed::Result<()> {
+        let (index, _) = self
-        self.query_level_iterator = self.query_level_iterator.dig(ctx)?;
+            .query_level_iterator
-        Ok(())
+            .iter_mut()
+            .map(|(pos, docids, qli)| {
+                if docids.is_empty() {
+                    0
+                } else {
+                    qli.peek()
+                        .map(|result| {
+                            result.as_ref().map(|(next_pos, _)| *next_pos - *pos).unwrap_or(0)
+                        })
+                        .unwrap_or(u32::MAX)
+                }
+            })
+            .enumerate()
+            .min_by_key(|(_, diff)| *diff)
+            .unwrap_or((0, 0));
+
+        index
     }
 
-    /// because next() method could be time consuming,
+    fn update_last_result(&mut self) {
-    /// update inner interval in order to be ranked by the binary_heap without computing it,
+        let mut result_pos = 0;
-    /// the next() method should be called when the real interval is needed.
+        let mut result_docids = None;
-    fn lazy_next(&mut self) {
-        let u8_level = Into::<u8>::into(self.tree_level);
-        let interval_size = LEVEL_EXPONENTIATION_BASE.pow(u8_level as u32);
-        let (left, right, _) = self.last_result;
 
-        self.last_result = (left + interval_size, right + interval_size, RoaringBitmap::new());
+        for (pos, docids, _qli) in self.query_level_iterator.iter() {
+            result_pos += pos;
+            result_docids = result_docids
+                .take()
+                .map_or_else(|| Some(docids.clone()), |candidates| Some(candidates & docids));
+        }
+
+        // remove last result docids from inner iterators
+        if let Some(docids) = result_docids.as_ref() {
+            for (_, query_docids, _) in self.query_level_iterator.iter_mut() {
+                *query_docids -= docids;
+            }
+        }
+
+        self.last_result = (result_pos, result_docids.unwrap_or_default());
     }
 
     /// return the score of the current inner interval.
     fn compute_rank(&self) -> u32 {
-        // we compute a rank from the left interval.
+        // we compute a rank from the position.
-        let (left, _, _) = self.last_result;
+        let (pos, _) = self.last_result;
-        left.saturating_sub((0..self.branch_size).sum()) * LCM_10_FIRST_NUMBERS / self.branch_size
+        pos.saturating_sub((0..self.branch_size).sum()) * LCM_10_FIRST_NUMBERS / self.branch_size
     }
 
     fn cmp(&self, other: &Self) -> Ordering {
         let self_rank = self.compute_rank();
         let other_rank = other.compute_rank();
-        let left_cmp = self_rank.cmp(&other_rank);
-        // on level: lower is better,
-        // we want to dig faster into levels on interesting branches.
-        let level_cmp = self.tree_level.cmp(&other.tree_level);
 
-        left_cmp.then(level_cmp).then(self.last_result.2.len().cmp(&other.last_result.2.len()))
+        // lower rank is better, and because BinaryHeap give the higher ranked branch, we reverse it.
+        self_rank.cmp(&other_rank).reverse()
     }
 }
 
-impl<'t, 'q> Ord for Branch<'t, 'q> {
+impl<'t> Ord for Branch<'t> {
     fn cmp(&self, other: &Self) -> Ordering {
         self.cmp(other)
     }
 }
 
-impl<'t, 'q> PartialOrd for Branch<'t, 'q> {
+impl<'t> PartialOrd for Branch<'t> {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
         Some(self.cmp(other))
     }
 }
 
-impl<'t, 'q> PartialEq for Branch<'t, 'q> {
+impl<'t> PartialEq for Branch<'t> {
     fn eq(&self, other: &Self) -> bool {
         self.cmp(other) == Ordering::Equal
     }
 }
 
-impl<'t, 'q> Eq for Branch<'t, 'q> {}
+impl<'t> Eq for Branch<'t> {}
 
-fn initialize_query_level_iterators<'t, 'q>(
+fn initialize_set_buckets<'t>(
-    ctx: &'t dyn Context<'t>,
-    branches: &'q FlattenedQueryTree,
-    allowed_candidates: &RoaringBitmap,
-    wdcache: &mut WordDerivationsCache,
-) -> Result<BinaryHeap<Branch<'t, 'q>>> {
-    let mut positions = BinaryHeap::with_capacity(branches.len());
-    for branch in branches {
-        let mut branch_positions = Vec::with_capacity(branch.len());
-        for queries in branch {
-            match QueryLevelIterator::new(ctx, queries, wdcache)? {
-                Some(qli) => branch_positions.push(qli),
-                None => {
-                    // the branch seems to be invalid, so we skip it.
-                    branch_positions.clear();
-                    break;
-                }
-            }
-        }
-        // QueryLevelIterator need to be sorted by level and folded in descending order.
-        branch_positions.sort_unstable_by_key(|qli| qli.level);
-        let folded_query_level_iterators =
-            branch_positions.into_iter().fold(None, |fold: Option<QueryLevelIterator>, mut qli| {
-                match fold {
-                    Some(fold) => {
-                        qli.parent(fold);
-                        Some(qli)
-                    }
-                    None => Some(qli),
-                }
-            });
-
-        if let Some(mut folded_query_level_iterators) = folded_query_level_iterators {
-            let tree_level = folded_query_level_iterators.level;
-            let last_result = folded_query_level_iterators.next(allowed_candidates, tree_level)?;
-            if let Some(last_result) = last_result {
-                let branch = Branch {
-                    last_result,
-                    tree_level,
-                    query_level_iterator: folded_query_level_iterators,
-                    branch_size: branch.len() as u32,
-                };
-                positions.push(branch);
-            }
-        }
-    }
-
-    Ok(positions)
-}
-
-fn set_compute_candidates<'t>(
     ctx: &'t dyn Context<'t>,
     branches: &FlattenedQueryTree,
     allowed_candidates: &RoaringBitmap,
     wdcache: &mut WordDerivationsCache,
-) -> Result<Option<RoaringBitmap>> {
+) -> Result<BinaryHeap<Branch<'t>>> {
-    let mut branches_heap =
+    let mut heap = BinaryHeap::new();
-        initialize_query_level_iterators(ctx, branches, allowed_candidates, wdcache)?;
+    for flatten_branch in branches {
-    let lowest_level = TreeLevel::min_value();
+        let branch = Branch::new(ctx, flatten_branch, wdcache, allowed_candidates)?;
+        heap.push(branch);
+    }
+
+    Ok(heap)
+}
+
+fn set_compute_candidates(
+    branches_heap: &mut BinaryHeap<Branch>,
+    allowed_candidates: &RoaringBitmap,
+) -> Result<Option<(u32, RoaringBitmap)>> {
     let mut final_candidates: Option<(u32, RoaringBitmap)> = None;
     let mut allowed_candidates = allowed_candidates.clone();
 
     while let Some(mut branch) = branches_heap.peek_mut() {
-        let is_lowest_level = branch.tree_level == lowest_level;
-        let branch_rank = branch.compute_rank();
         // if current is worst than best we break to return
         // candidates that correspond to the best rank
+        let branch_rank = branch.compute_rank();
         if let Some((best_rank, _)) = final_candidates {
             if branch_rank > best_rank {
                 break;
             }
         }
-        let _left = branch.last_result.0;
+
-        let candidates = take(&mut branch.last_result.2);
+        let candidates = take(&mut branch.last_result.1);
         if candidates.is_empty() {
             // we don't have candidates, get next interval.
             if !branch.next(&allowed_candidates)? {
                 PeekMut::pop(branch);
             }
-        } else if is_lowest_level {
+        } else {
-            // we have candidates, but we can't dig deeper.
             allowed_candidates -= &candidates;
             final_candidates = match final_candidates.take() {
                 // we add current candidates to best candidates
                 Some((best_rank, mut best_candidates)) => {
                     best_candidates |= candidates;
-                    branch.lazy_next();
+                    branch.next(&allowed_candidates)?;
                     Some((best_rank, best_candidates))
                 }
                 // we take current candidates as best candidates
                 None => {
-                    branch.lazy_next();
+                    branch.next(&allowed_candidates)?;
                     Some((branch_rank, candidates))
                 }
             };
-        } else {
-            // we have candidates, lets dig deeper in levels.
-            branch.dig(ctx)?;
-            if !branch.next(&allowed_candidates)? {
-                PeekMut::pop(branch);
-            }
         }
     }
 
-    Ok(final_candidates.map(|(_rank, candidates)| candidates))
+    Ok(final_candidates)
 }
 
-fn linear_compute_candidates(
+fn initialize_linear_buckets(
     ctx: &dyn Context,
     branches: &FlattenedQueryTree,
     allowed_candidates: &RoaringBitmap,