Support swapped word pairs in new proximity ranking rule impl

milli/src/search/new/db_cache.rs

@@ -11,6 +11,8 @@ pub struct DatabaseCache<'transaction> {
     pub word_pair_proximity_docids: FxHashMap<(u8, String, String), Option<&'transaction [u8]>>,
     pub word_prefix_pair_proximity_docids:
         FxHashMap<(u8, String, String), Option<&'transaction [u8]>>,
+    pub prefix_word_pair_proximity_docids:
+        FxHashMap<(u8, String, String), Option<&'transaction [u8]>>,
     pub word_docids: FxHashMap<String, Option<&'transaction [u8]>>,
     pub exact_word_docids: FxHashMap<String, Option<&'transaction [u8]>>,
     pub word_prefix_docids: FxHashMap<String, Option<&'transaction [u8]>>,
@@ -115,4 +117,25 @@ impl<'transaction> DatabaseCache<'transaction> {
             }
         }
     }
+    pub fn get_prefix_word_pair_proximity_docids(
+        &mut self,
+        index: &Index,
+        txn: &'transaction RoTxn,
+        word1: &str,
+        prefix2: &str,
+        proximity: u8,
+    ) -> Result<Option<&'transaction [u8]>> {
+        let key = (proximity, prefix2.to_owned(), word1.to_owned());
+        match self.prefix_word_pair_proximity_docids.entry(key) {
+            Entry::Occupied(bitmap_ptr) => Ok(*bitmap_ptr.get()),
+            Entry::Vacant(entry) => {
+                let bitmap_ptr = index
+                    .prefix_word_pair_proximity_docids
+                    .remap_data_type::<ByteSlice>()
+                    .get(txn, &(proximity, prefix2, word1))?;
+                entry.insert(bitmap_ptr);
+                Ok(bitmap_ptr)
+            }
+        }
+    }
 }

milli/src/search/new/ranking_rule_graph/proximity/compute_docids.rs

@@ -20,6 +20,11 @@ pub fn compute_docids<'transaction>(
             }
             WordPair::WordPrefix { left, right_prefix } => db_cache
                 .get_word_prefix_pair_proximity_docids(index, txn, left, right_prefix, *proximity),
+            WordPair::WordsSwapped { left, right } => {
+                db_cache.get_word_pair_proximity_docids(index, txn, left, right, *proximity)
+            }
+            WordPair::WordPrefixSwapped { left, right_prefix } => db_cache
+                .get_prefix_word_pair_proximity_docids(index, txn, left, right_prefix, *proximity),
         }?;
         let bitmap =
             bytes.map(CboRoaringBitmapCodec::deserialize_from).transpose()?.unwrap_or_default();

milli/src/search/new/ranking_rule_graph/proximity/mod.rs

@@ -13,7 +13,9 @@ use crate::{Index, Result};
 pub enum WordPair {
     // TODO: add WordsSwapped and WordPrefixSwapped case
     Words { left: String, right: String },
+    WordsSwapped { left: String, right: String },
     WordPrefix { left: String, right_prefix: String },
+    WordPrefixSwapped { left: String, right_prefix: String },
 }
 
 pub struct ProximityEdge {