From 85d96d35a857a45ce70f45d5996b5b30fc3ff1c2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Lecrenier?= <loic.lecrenier@me.com>
Date: Wed, 3 May 2023 13:39:19 +0200
Subject: [PATCH] Highlight ngram matches as well

---
 milli/src/search/new/mod.rs         |  4 ++--
 milli/src/search/new/query_graph.rs | 14 ++++++++++----
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/milli/src/search/new/mod.rs b/milli/src/search/new/mod.rs
index 7e8426bf9..cbc085b12 100644
--- a/milli/src/search/new/mod.rs
+++ b/milli/src/search/new/mod.rs
@@ -397,8 +397,8 @@ pub fn execute_search(
         None
     };
     let bucket_sort_output = if let Some(query_terms) = query_terms {
-        let graph = QueryGraph::from_query(ctx, &query_terms)?;
-        located_query_terms = Some(query_terms);
+        let (graph, new_located_query_terms) = QueryGraph::from_query(ctx, &query_terms)?;
+        located_query_terms = Some(new_located_query_terms);
 
         let ranking_rules = get_ranking_rules_for_query_graph_search(
             ctx,
diff --git a/milli/src/search/new/query_graph.rs b/milli/src/search/new/query_graph.rs
index 0e7d5a7f3..dc25d1bc3 100644
--- a/milli/src/search/new/query_graph.rs
+++ b/milli/src/search/new/query_graph.rs
@@ -88,12 +88,15 @@ pub struct QueryGraph {
 }
 
 impl QueryGraph {
-    /// Build the query graph from the parsed user search query.
+    /// Build the query graph from the parsed user search query, return an updated list of the located query terms
+    /// which contains ngrams.
     pub fn from_query(
         ctx: &mut SearchContext,
         // NOTE: the terms here must be consecutive
         terms: &[LocatedQueryTerm],
-    ) -> Result<QueryGraph> {
+    ) -> Result<(QueryGraph, Vec<LocatedQueryTerm>)> {
+        let mut new_located_query_terms = terms.to_vec();
+
         let nbr_typos = number_of_typos_allowed(ctx)?;
 
         let mut nodes_data: Vec<QueryNodeData> = vec![QueryNodeData::Start, QueryNodeData::End];
@@ -107,10 +110,11 @@ impl QueryGraph {
         let original_terms_len = terms.len();
         for term_idx in 0..original_terms_len {
             let mut new_nodes = vec![];
+
             let new_node_idx = add_node(
                 &mut nodes_data,
                 QueryNodeData::Term(LocatedQueryTermSubset {
-                    term_subset: QueryTermSubset::full(Interned::from_raw(term_idx as u16)),
+                    term_subset: QueryTermSubset::full(terms[term_idx].value),
                     positions: terms[term_idx].positions.clone(),
                     term_ids: term_idx as u8..=term_idx as u8,
                 }),
@@ -121,6 +125,7 @@ impl QueryGraph {
                 if let Some(ngram) =
                     query_term::make_ngram(ctx, &terms[term_idx - 1..=term_idx], &nbr_typos)?
                 {
+                    new_located_query_terms.push(ngram.clone());
                     let ngram_idx = add_node(
                         &mut nodes_data,
                         QueryNodeData::Term(LocatedQueryTermSubset {
@@ -136,6 +141,7 @@ impl QueryGraph {
                 if let Some(ngram) =
                     query_term::make_ngram(ctx, &terms[term_idx - 2..=term_idx], &nbr_typos)?
                 {
+                    new_located_query_terms.push(ngram.clone());
                     let ngram_idx = add_node(
                         &mut nodes_data,
                         QueryNodeData::Term(LocatedQueryTermSubset {
@@ -167,7 +173,7 @@ impl QueryGraph {
         let mut graph = QueryGraph { root_node, end_node, nodes };
         graph.build_initial_edges();
 
-        Ok(graph)
+        Ok((graph, new_located_query_terms))
     }
 
     /// Remove the given nodes, connecting all their predecessors to all their successors.