MeiliSearch/milli/src/search/criteria/typo.rs

use std::{borrow::Cow, mem::take};

use anyhow::bail;
use roaring::RoaringBitmap;

use crate::Index;
use crate::search::query_tree::{Operation, Query, QueryKind};
use crate::search::word_typos;
use super::{Candidates, Criterion};

// FIXME we must stop when the number of typos is equal to
// the maximum number of typos for this query tree.
const MAX_NUM_TYPOS: u8 = 8;

pub struct Typo<'t> {
    index: &'t Index,
    rtxn: &'t heed::RoTxn<'t>,
    words_fst: fst::Set<Cow<'t, [u8]>>,
    query_tree: Option<Operation>,
    number_typos: u8,
    candidates: Candidates,
    parent: Option<Box<dyn Criterion>>,
}

impl<'t> Typo<'t> {
    pub fn initial(
        index: &'t Index,
        rtxn: &'t heed::RoTxn<'t>,
        query_tree: Option<Operation>,
        candidates: Option<RoaringBitmap>,
    ) -> anyhow::Result<Self> where Self: Sized
    {
        Ok(Typo {
            index,
            rtxn,
            words_fst: index.words_fst(rtxn)?,
            query_tree,
            number_typos: 0,
            candidates: candidates.map_or_else(Candidates::default, Candidates::Allowed),
            parent: None,
        })
    }

    pub fn new(
        index: &'t Index,
        rtxn: &'t heed::RoTxn<'t>,
        parent: Box<dyn Criterion>,
    ) -> anyhow::Result<Self> where Self: Sized
    {
        Ok(Typo {
            index,
            rtxn,
            words_fst: index.words_fst(rtxn)?,
            query_tree: None,
            number_typos: 0,
            candidates: Candidates::default(),
            parent: Some(parent),
        })
    }
}

impl<'t> Criterion for Typo<'t> {
    fn next(&mut self) -> anyhow::Result<Option<(Option<Operation>, RoaringBitmap)>> {
        use Candidates::{Allowed, Forbidden};
        while self.number_typos < MAX_NUM_TYPOS {
            match (&mut self.query_tree, &mut self.candidates) {
                (_, Allowed(candidates)) if candidates.is_empty() => {
                    self.query_tree = None;
                    self.candidates = Candidates::default();
                },
                (Some(query_tree), Allowed(candidates)) => {
                    let new_query_tree = alterate_query_tree(&self.words_fst, query_tree.clone(), self.number_typos)?;
                    let mut new_candidates = resolve_candidates(&self.index, &self.rtxn, &new_query_tree, self.number_typos)?;
                    new_candidates.intersect_with(&candidates);
                    candidates.difference_with(&new_candidates);
                    self.number_typos += 1;

                    return Ok(Some((Some(new_query_tree), new_candidates)));
                },
                (Some(query_tree), Forbidden(candidates)) => {
                    let new_query_tree = alterate_query_tree(&self.words_fst, query_tree.clone(), self.number_typos)?;
                    let mut new_candidates = resolve_candidates(&self.index, &self.rtxn, &new_query_tree, self.number_typos)?;
                    new_candidates.difference_with(&candidates);
                    candidates.union_with(&new_candidates);
                    self.number_typos += 1;

                    return Ok(Some((Some(new_query_tree), new_candidates)));
                },
                (None, Allowed(_)) => {
                    return Ok(Some((None, take(&mut self.candidates).into_inner())));
                },
                (None, Forbidden(_)) => {
                    match self.parent.as_mut() {
                        Some(parent) => {
                            match parent.next()? {
                                Some((query_tree, candidates)) => {
                                    self.query_tree = query_tree;
                                    self.candidates = Candidates::Allowed(candidates);
                                },
                                None => return Ok(None),
                            }
                        },
                        None => return Ok(None),
                    }
                },
            }
        }

        Ok(None)
    }
}

/// Modify the query tree by replacing every tolerant query by an Or operation
/// containing all of the corresponding exact words in the words FST. Each tolerant
/// query will only be replaced by exact query with up to `number_typos` maximum typos.
fn alterate_query_tree(
    words_fst: &fst::Set<Cow<[u8]>>,
    mut query_tree: Operation,
    number_typos: u8,
) -> anyhow::Result<Operation>
{
    fn recurse(
        words_fst: &fst::Set<Cow<[u8]>>,
        operation: &mut Operation,
        number_typos: u8,
    ) -> anyhow::Result<()>
    {
        use Operation::{And, Consecutive, Or};

        match operation {
            And(ops) | Consecutive(ops) | Or(_, ops) => {
                ops.iter_mut().try_for_each(|op| recurse(words_fst, op, number_typos))
            },
            Operation::Query(q) => {
                if let QueryKind::Tolerant { typo, word } = &q.kind {
                    let typo = *typo.min(&number_typos);
                    let words = word_typos(word, q.prefix, typo, words_fst)?;

                    let queries = words.into_iter().map(|(word, _typo)| {
                        Operation::Query(Query {
                            prefix: false,
                            kind: QueryKind::Exact { original_typo: typo, word },
                        })
                    }).collect();

                    *operation = Operation::or(false, queries);
                }

                Ok(())
            },
        }
    }

    recurse(words_fst, &mut query_tree, number_typos)?;
    Ok(query_tree)
}

fn resolve_candidates(
    index: &Index,
    rtxn: &heed::RoTxn,
    query_tree: &Operation,
    number_typos: u8,
) -> anyhow::Result<RoaringBitmap>
{
    // FIXME add a cache
    // FIXME keep the cache between typos iterations
    // cache: HashMap<(&Operation, u8), RoaringBitmap>,

    fn resolve_operation(
        index: &Index,
        rtxn: &heed::RoTxn,
        query_tree: &Operation,
        number_typos: u8,
    ) -> anyhow::Result<RoaringBitmap>
    {
        use Operation::{And, Consecutive, Or, Query};

        match query_tree {
            And(ops) => {
                mdfs(index, rtxn, ops, number_typos)
            },
            Consecutive(ops) => {
                let mut candidates = RoaringBitmap::new();
                let mut first_loop = true;
                for slice in ops.windows(2) {
                    match (&slice[0], &slice[1]) {
                        (Operation::Query(left), Operation::Query(right)) => {
                            let key_pair = &(left.kind.word(), right.kind.word(), 1);
                            match index.word_pair_proximity_docids.get(rtxn, key_pair)? {
                                Some(pair_docids) => {
                                    if first_loop {
                                        candidates = pair_docids;
                                        first_loop = false;
                                    } else {
                                        candidates.intersect_with(&pair_docids)
                                    }
                                },
                                None => return Ok(RoaringBitmap::new()),
                            }

                        },
                        _ => bail!("invalid consecutive query type"),
                    }
                }
                Ok(candidates)
            },
            Or(_, ops) => {
                let mut candidates = RoaringBitmap::new();
                for op in ops {
                    let docids = resolve_operation(index, rtxn, op, number_typos)?;
                    candidates.union_with(&docids);
                }
                Ok(candidates)
            },
            Query(q) => if q.kind.typo() == number_typos {
                let word = q.kind.word();
                Ok(index.word_docids.get(rtxn, word)?.unwrap_or_default())
            } else {
                Ok(RoaringBitmap::new())
            },
        }
    }

    /// FIXME Make this function generic and mutualize it between Typo and proximity criterion
    fn mdfs(
        index: &Index,
        rtxn: &heed::RoTxn,
        branches: &[Operation],
        mana: u8,
    ) -> anyhow::Result<RoaringBitmap>
    {
        match branches.split_first() {
            Some((head, [])) => resolve_operation(index, rtxn, head, mana),
            Some((head, tail)) => {
                let mut candidates = RoaringBitmap::new();

                for m in 0..=mana {
                    let mut head_candidates = resolve_operation(index, rtxn, head, m)?;
                    if !head_candidates.is_empty() {
                        let tail_candidates = mdfs(index, rtxn, tail, mana - m)?;
                        head_candidates.intersect_with(&tail_candidates);
                        candidates.union_with(&head_candidates);
                    }
                }

                Ok(candidates)
            },
            None => Ok(RoaringBitmap::new()),
        }
    }

    resolve_operation(index, rtxn, query_tree, number_typos)
}