MOD-11277 implement RLookupRow::write_key_by_name porting C RLookup_WriteKeyByName function

JonasKruckenberg · DarthB · JonasKruckenberg · commit 03aeae4def20 · 2025-10-02T09:25:46.000+02:00
Co-authored-by: Tim Janus &lt;tim@janus.rs&gt;
diff --git a/src/redisearch_rs/rlookup/src/lib.rs b/src/redisearch_rs/rlookup/src/lib.rs
@@ -15,6 +15,7 @@ mod row;
 
 pub use bindings::IndexSpecCache;
 pub use lookup::{
-    RLookup, RLookupKey, RLookupKeyFlag, RLookupKeyFlags, RLookupOption, RLookupOptions,
+    Cursor, CursorMut, RLookup, RLookupKey, RLookupKeyFlag, RLookupKeyFlags, RLookupOption,
+    RLookupOptions,
 };
 pub use row::RLookupRow;
diff --git a/src/redisearch_rs/rlookup/src/lookup.rs b/src/redisearch_rs/rlookup/src/lookup.rs
@@ -230,13 +230,13 @@ struct KeyList<'a> {
     rowlen: u32,
 }
 
-/// A cursor over a [`KeyList`].
+/// A cursor over an [`RLookup`] key list.
 pub struct Cursor<'list, 'a> {
     _rlookup: &'list KeyList<'a>,
     current: Option<NonNull<RLookupKey<'a>>>,
 }
 
-/// A cursor over a [`KeyList`] with editing operations.
+/// A cursor over an [`RLookup`] key list with editing operations.
 pub struct CursorMut<'list, 'a> {
     _rlookup: &'list mut KeyList<'a>,
     current: Option<NonNull<RLookupKey<'a>>>,
@@ -721,7 +721,7 @@ impl Drop for KeyList<'_> {
 // ===== impl Cursor =====
 
 impl<'list, 'a> Cursor<'list, 'a> {
-    /// Move the cursor to the next [`RLookupKey`] in the [`KeyList`].
+    /// Move the cursor to the next [`RLookupKey`] in the key list.
     ///
     /// Note that contrary to [`Self::next`] this **does not** skip over hidden keys.
     pub fn move_next(&mut self) {
@@ -761,7 +761,7 @@ impl<'list, 'a> Cursor<'list, 'a> {
 impl<'list, 'a> Iterator for Cursor<'list, 'a> {
     type Item = &'list RLookupKey<'a>;
 
-    /// Advances the [`Cursor`] to the next [`RLookupKey`] in the [`KeyList`] and returns it.
+    /// Advances the [`Cursor`] to the next [`RLookupKey`] in the key list and returns it.
     ///
     /// This will automatically skip over any keys with the [`RLookupKeyFlag::Hidden`] flag.
     fn next(&mut self) -> Option<Self::Item> {
@@ -879,7 +879,7 @@ impl<'list, 'a> CursorMut<'list, 'a> {
 impl<'list, 'a> Iterator for CursorMut<'list, 'a> {
     type Item = Pin<&'list mut RLookupKey<'a>>;
 
-    /// Advances the [`CursorMut`] to the next [`RLookupKey`] in the [`KeyList`] and returns it.
+    /// Advances the [`CursorMut`] to the next [`RLookupKey`] in the key list and returns it.
     ///
     /// This will automatically skip over any keys with the [`RLookupKeyFlag::Hidden`] flag.
     fn next(&mut self) -> Option<Self::Item> {
@@ -938,6 +938,12 @@ impl<'a> RLookup<'a> {
         self.index_spec_cache = spcache;
     }
 
+    /// Find a [`RLookupKey`] in this `RLookup`'s key list by its `name`
+    /// and return a [`Cursor`] pointing to the key if found.
+    pub fn find_by_name(&self, name: &CStr) -> Option<Cursor<'_, 'a>> {
+        self.keys.find_by_name(name)
+    }
+
     // ===== Get key for reading (create only if in schema and sortable) =====
 
     /// Gets a key by its name from the lookup table, if not found it uses the schema as a fallback to search the key.
@@ -955,7 +961,8 @@ impl<'a> RLookup<'a> {
 
         let available = self.keys.find_by_name(&name).is_some();
         if available {
-            // FIXME: Duplication because of borrow-checker false positive. Duplication means performance implications.
+            // FIXME: We cannot use let-some above because of a borrow-checker false positive.
+            // This duplication might have performance implications.
             // See <https://github.com/rust-lang/rust/issues/54663>
             return self.keys.find_by_name(&name).unwrap().into_current();
         }
diff --git a/src/redisearch_rs/rlookup/src/row.rs b/src/redisearch_rs/rlookup/src/row.rs
@@ -7,11 +7,11 @@
  * GNU Affero General Public License v3 (AGPLv3).
 */
 
+use crate::{RLookup, RLookupKey, RLookupKeyFlag, RLookupKeyFlags};
 use sorting_vector::RSSortingVector;
+use std::{borrow::Cow, ffi::CStr};
 use value::RSValueTrait;
 
-use crate::{RLookupKey, RLookupKeyFlag};
-
 /// Row data for a lookup key. This abstracts the question of if the data comes from a borrowed [RSSortingVector]
 /// or from dynamic values stored in the row during processing.
 ///
@@ -127,6 +127,26 @@ impl<'a, T: RSValueTrait> RLookupRow<'a, T> {
         prev
     }
 
+    /// Write a value to the lookup table *by-name*. This is useful for 'dynamic' keys
+    /// for which it is not necessary to use the boilerplate of getting an explicit
+    /// key.
+    pub fn write_key_by_name(
+        &mut self,
+        rlookup: &mut RLookup<'a>,
+        name: impl Into<Cow<'a, CStr>>,
+        val: T,
+    ) {
+        let name = name.into();
+        let key = if let Some(cursor) = rlookup.find_by_name(&name) {
+            cursor.into_current().expect("the cursor returned by `Keys::find_by_name` must have a current key. This is a bug!")
+        } else {
+            rlookup
+                .get_key_write(name, RLookupKeyFlags::empty())
+                .expect("`RLookup::get_key_write` must never return None for non-existent keys. This is a bug!")
+        };
+        self.write_key(key, val);
+    }
+
     /// Wipes the row, retaining its memory but decrementing the ref count of any included instance of `T`.
     /// This does not free all the memory consumed by the row, but simply resets
     /// the row data (preserving any caches) so that it may be refilled.
diff --git a/src/redisearch_rs/rlookup/tests/row.rs b/src/redisearch_rs/rlookup/tests/row.rs
@@ -7,11 +7,14 @@
  * GNU Affero General Public License v3 (AGPLv3).
 */
 
-use rlookup::{RLookupKey, RLookupKeyFlag, RLookupKeyFlags, RLookupRow};
+use rlookup::{RLookup, RLookupKey, RLookupKeyFlag, RLookupKeyFlags, RLookupRow};
 use sorting_vector::RSSortingVector;
 use std::{
     ffi::CString,
+    mem::offset_of,
     ops::{Deref, DerefMut},
+    ptr::NonNull,
+    sync::atomic::{AtomicUsize, Ordering},
 };
 use value::{RSValueMock, RSValueTrait};
 
@@ -388,6 +391,125 @@ fn test_rlookup_get_item_priority_dynamic_over_static() {
     assert_eq!(result.unwrap().as_str(), Some("dynamic_value"));
 }
 
+#[test]
+fn test_write_key_by_name_new_key() {
+    // Test case: name is not yet part of the lookup and gets created
+    let mut lookup = RLookup::new();
+    let mut row = RLookupRow::new();
+
+    let key_name = CString::new("new_key").unwrap();
+    let value = RSValueMock::create_string("test_value".to_string());
+
+    // Initially, row should be empty
+    assert_eq!(row.len(), 0);
+
+    // Write the key
+    row.write_key_by_name(&mut lookup, key_name.to_owned(), value.clone());
+
+    // Verify we can find the key by name
+    let cursor = lookup.find_by_name(&key_name);
+    assert!(cursor.is_some());
+
+    // Verify the rlookup row is in correct state
+    assert_eq!(row.len(), 1);
+    assert!(row.dyn_values()[0].is_some());
+    assert_eq!(
+        row.dyn_values()[0].as_ref().unwrap().as_str(),
+        Some("test_value")
+    );
+}
+
+#[test]
+fn test_write_key_by_name_existing_key_overwrite() {
+    // Test case: name is part of the lookup and its value gets overwritten
+    let mut lookup = RLookup::new();
+    let mut row = RLookupRow::new();
+
+    let key_name = CString::new("existing_key").unwrap();
+    let initial_value = RSValueMock::create_string("initial_value".to_string());
+    let new_value = RSValueMock::create_string("new_value".to_string());
+
+    // Write initial value
+    row.write_key_by_name(&mut lookup, key_name.to_owned(), initial_value.clone());
+
+    // Verify initial state
+    let cursor = lookup.find_by_name(&key_name).unwrap();
+    assert!(cursor.into_current().is_some());
+    assert_eq!(row.len(), 1);
+    assert_eq!(
+        row.dyn_values()[0].as_ref().unwrap().as_str(),
+        initial_value.as_str()
+    );
+
+    // Overwrite with new value - key count should not increase
+    row.write_key_by_name(&mut lookup, key_name.to_owned(), new_value.clone());
+
+    let cursor = lookup.find_by_name(&key_name).unwrap();
+    assert!(cursor.into_current().is_some());
+    assert_eq!(row.len(), 1);
+    assert_eq!(
+        row.dyn_values()[0].as_ref().unwrap().as_str(),
+        new_value.as_str()
+    );
+}
+
+#[test]
+fn test_write_multiple_different_keys() {
+    // Test case: writing multiple different keys
+    let mut lookup = RLookup::new();
+    let mut row = RLookupRow::new();
+
+    let key1_name = CString::new("key1").unwrap();
+    let key2_name = CString::new("key2").unwrap();
+    let key3_name = CString::new("key3").unwrap();
+
+    let value1 = RSValueMock::create_string("value1".to_string());
+    let value2 = RSValueMock::create_string("value2".to_string());
+    let value3 = RSValueMock::create_string("value3".to_string());
+
+    // Write multiple keys
+    row.write_key_by_name(&mut lookup, key1_name.to_owned(), value1.clone());
+    row.write_key_by_name(&mut lookup, key2_name.to_owned(), value2.clone());
+    row.write_key_by_name(&mut lookup, key3_name.to_owned(), value3.clone());
+
+    // Verify all keys were added
+    assert_eq!(row.len(), 3);
+
+    for (key_name, value) in [
+        (&key1_name, value1),
+        (&key2_name, value2),
+        (&key3_name, value3),
+    ] {
+        let cursor = lookup.find_by_name(key_name);
+        let key = cursor.unwrap().into_current().unwrap();
+        assert!(row.dyn_values()[key.dstidx as usize].is_some());
+        assert_eq!(
+            row.dyn_values()[key.dstidx as usize]
+                .as_ref()
+                .unwrap()
+                .as_str(),
+            value.as_str(),
+        );
+    }
+}
+
+/// Mock implementation of `IndexSpecCache_Decref` from spec.h for testing purposes
+#[unsafe(no_mangle)]
+extern "C" fn IndexSpecCache_Decref(spcache: Option<NonNull<ffi::IndexSpecCache>>) {
+    let spcache = spcache.expect("`spcache` must not be null");
+    let refcount = unsafe {
+        spcache
+            .byte_add(offset_of!(ffi::IndexSpecCache, refcount))
+            .cast::<usize>()
+    };
+
+    let refcount = unsafe { AtomicUsize::from_ptr(refcount.as_ptr()) };
+
+    if refcount.fetch_sub(1, Ordering::Relaxed) == 1 {
+        drop(unsafe { Box::from_raw(spcache.as_ptr()) });
+    }
+}
+
 fn create_test_key(dstidx: u16, svidx: u16, flags: RLookupKeyFlags) -> RLookupKey<'static> {
     let str = format!("mock_key_{}_{}", dstidx, svidx);
     let cstring = CString::new(str).unwrap();
