MeiliSearch/milli/src/search/new/query_graph.rs

use std::fmt::Debug;

use heed::RoTxn;
use roaring::RoaringBitmap;

use super::db_cache::DatabaseCache;
use super::query_term::{self, LocatedQueryTerm, QueryTerm, WordDerivations};
use crate::{Index, Result};

#[derive(Debug, Clone)]
pub enum QueryNode {
    Term(LocatedQueryTerm),
    Deleted,
    Start,
    End,
}

#[derive(Debug, Clone)]
pub struct Edges {
    // TODO: use a tiny bitset instead, something like a simple Vec<u8> where most queries will see a vector of one element
    pub predecessors: RoaringBitmap,
    pub successors: RoaringBitmap,
}

#[derive(Debug, Clone)]
pub struct QueryGraph {
    pub root_node: u32,
    pub end_node: u32,
    pub nodes: Vec<QueryNode>,
    pub edges: Vec<Edges>,
}

fn _assert_sizes() {
    // TODO: QueryNodes are too big now, 184B is an unreasonable size
    let _: [u8; 184] = [0; std::mem::size_of::<QueryNode>()];
    let _: [u8; 48] = [0; std::mem::size_of::<Edges>()];
}

impl Default for QueryGraph {
    /// Create a new QueryGraph with two disconnected nodes: the root and end nodes.
    fn default() -> Self {
        let nodes = vec![QueryNode::Start, QueryNode::End];
        let edges = vec![
            Edges { predecessors: RoaringBitmap::new(), successors: RoaringBitmap::new() },
            Edges { predecessors: RoaringBitmap::new(), successors: RoaringBitmap::new() },
        ];

        Self { root_node: 0, end_node: 1, nodes, edges }
    }
}

impl QueryGraph {
    fn connect_to_node(&mut self, from_nodes: &[u32], to_node: u32) {
        for &from_node in from_nodes {
            self.edges[from_node as usize].successors.insert(to_node);
            self.edges[to_node as usize].predecessors.insert(from_node);
        }
    }
    fn add_node(&mut self, from_nodes: &[u32], node: QueryNode) -> u32 {
        let new_node_idx = self.nodes.len() as u32;
        self.nodes.push(node);
        self.edges.push(Edges {
            predecessors: from_nodes.iter().collect(),
            successors: RoaringBitmap::new(),
        });
        for from_node in from_nodes {
            self.edges[*from_node as usize].successors.insert(new_node_idx);
        }
        new_node_idx
    }
}

impl QueryGraph {
    // TODO: return the list of all matching words here as well
    pub fn from_query<'transaction>(
        index: &Index,
        txn: &RoTxn,
        _db_cache: &mut DatabaseCache<'transaction>,
        terms: Vec<LocatedQueryTerm>,
    ) -> Result<QueryGraph> {
        // TODO: maybe empty nodes should not be removed here, to compute
        // the score of the `words` ranking rule correctly
        // it is very easy to traverse the graph and remove afterwards anyway
        // Still, I'm keeping this here as a demo
        let mut empty_nodes = vec![];

        let word_set = index.words_fst(txn)?;
        let mut graph = QueryGraph::default();

        let (mut prev2, mut prev1, mut prev0): (Vec<u32>, Vec<u32>, Vec<u32>) =
            (vec![], vec![], vec![graph.root_node]);

        // TODO: split words / synonyms
        for length in 1..=terms.len() {
            let query = &terms[..length];

            let term0 = query.last().unwrap();

            let mut new_nodes = vec![];
            let new_node_idx = graph.add_node(&prev0, QueryNode::Term(term0.clone()));
            new_nodes.push(new_node_idx);
            if term0.is_empty() {
                empty_nodes.push(new_node_idx);
            }

            if !prev1.is_empty() {
                if let Some((ngram2_str, ngram2_pos)) =
                    query_term::ngram2(&query[length - 2], &query[length - 1])
                {
                    if word_set.contains(ngram2_str.as_bytes()) {
                        let ngram2 = LocatedQueryTerm {
                            value: QueryTerm::Word {
                                derivations: WordDerivations {
                                    original: ngram2_str.clone(),
                                    // TODO: could add a typo if it's an ngram?
                                    zero_typo: vec![ngram2_str],
                                    one_typo: vec![],
                                    two_typos: vec![],
                                    use_prefix_db: false,
                                    synonyms: vec![],  // TODO: ngram synonyms
                                    split_words: None, // TODO: maybe ngram split words?
                                },
                            },
                            positions: ngram2_pos,
                        };
                        let ngram2_idx = graph.add_node(&prev1, QueryNode::Term(ngram2));
                        new_nodes.push(ngram2_idx);
                    }
                }
            }
            if !prev2.is_empty() {
                if let Some((ngram3_str, ngram3_pos)) =
                    query_term::ngram3(&query[length - 3], &query[length - 2], &query[length - 1])
                {
                    if word_set.contains(ngram3_str.as_bytes()) {
                        let ngram3 = LocatedQueryTerm {
                            value: QueryTerm::Word {
                                derivations: WordDerivations {
                                    original: ngram3_str.clone(),
                                    // TODO: could add a typo if it's an ngram?
                                    zero_typo: vec![ngram3_str],
                                    one_typo: vec![],
                                    two_typos: vec![],
                                    use_prefix_db: false,
                                    synonyms: vec![], // TODO: ngram synonyms
                                    split_words: None, // TODO: maybe ngram split words?
                                                      // would be nice for typos like su nflower
                                },
                            },
                            positions: ngram3_pos,
                        };
                        let ngram3_idx = graph.add_node(&prev2, QueryNode::Term(ngram3));
                        new_nodes.push(ngram3_idx);
                    }
                }
            }
            (prev0, prev1, prev2) = (new_nodes, prev0, prev1);
        }
        graph.connect_to_node(&prev0, graph.end_node);

        graph.remove_nodes_keep_edges(&empty_nodes);

        Ok(graph)
    }
    pub fn remove_nodes(&mut self, nodes: &[u32]) {
        for &node in nodes {
            self.nodes[node as usize] = QueryNode::Deleted;
            let edges = self.edges[node as usize].clone();
            for pred in edges.predecessors.iter() {
                self.edges[pred as usize].successors.remove(node);
            }
            for succ in edges.successors {
                self.edges[succ as usize].predecessors.remove(node);
            }
            self.edges[node as usize] =
                Edges { predecessors: RoaringBitmap::new(), successors: RoaringBitmap::new() };
        }
    }
    pub fn remove_nodes_keep_edges(&mut self, nodes: &[u32]) {
        for &node in nodes {
            self.nodes[node as usize] = QueryNode::Deleted;
            let edges = self.edges[node as usize].clone();
            for pred in edges.predecessors.iter() {
                self.edges[pred as usize].successors.remove(node);
                self.edges[pred as usize].successors |= &edges.successors;
            }
            for succ in edges.successors {
                self.edges[succ as usize].predecessors.remove(node);
                self.edges[succ as usize].predecessors |= &edges.predecessors;
            }
            self.edges[node as usize] =
                Edges { predecessors: RoaringBitmap::new(), successors: RoaringBitmap::new() };
        }
    }
    pub fn remove_words_at_position(&mut self, position: i8) -> bool {
        let mut nodes_to_remove_keeping_edges = vec![];
        for (node_idx, node) in self.nodes.iter().enumerate() {
            let node_idx = node_idx as u32;
            let QueryNode::Term(LocatedQueryTerm { value: _, positions }) = node else { continue };
            if positions.start() == &position {
                nodes_to_remove_keeping_edges.push(node_idx);
            }
        }

        self.remove_nodes_keep_edges(&nodes_to_remove_keeping_edges);

        self.simplify();
        !nodes_to_remove_keeping_edges.is_empty()
    }

    fn simplify(&mut self) {
        loop {
            let mut nodes_to_remove = vec![];
            for (node_idx, node) in self.nodes.iter().enumerate() {
                if (!matches!(node, QueryNode::End | QueryNode::Deleted)
                    && self.edges[node_idx].successors.is_empty())
                    || (!matches!(node, QueryNode::Start | QueryNode::Deleted)
                        && self.edges[node_idx].predecessors.is_empty())
                {
                    nodes_to_remove.push(node_idx as u32);
                }
            }
            if nodes_to_remove.is_empty() {
                break;
            } else {
                self.remove_nodes(&nodes_to_remove);
            }
        }
    }
}