feat: add lookup_index_by_key to Rust Vector for index-based search (#8959)

* feat: add lookup_index_by_key to Rust Vector for index-based binary search

* fix: remove duplicated code

rust/flatbuffers/src/vector.rs

@@ -107,6 +107,20 @@ impl<'a, T: Follow<'a> + 'a> Vector<'a, T> {
         key: K,
         f: fn(&<T as Follow<'a>>::Inner, &K) -> Ordering,
     ) -> Option<T::Inner> {
+        self.lookup_index_by_key(key, f).map(|idx| self.get(idx))
+    }
+
+    /// Binary search by key, returning the index of the matching element.
+    ///
+    /// This is similar to `lookup_by_key`, but returns the index of the found
+    /// element rather than the element itself. This is useful when you need
+    /// to reference elements by their position in the vector.
+    #[inline(always)]
+    pub fn lookup_index_by_key<K: Ord>(
+        &self,
+        key: K,
+        f: fn(&<T as Follow<'a>>::Inner, &K) -> Ordering,
+    ) -> Option<usize> {
         if self.is_empty() {
             return None;
         }
@@ -118,7 +132,7 @@ impl<'a, T: Follow<'a> + 'a> Vector<'a, T> {
             let mid = (left + right) / 2;
             let value = self.get(mid);
             match f(&value, &key) {
-                Ordering::Equal => return Some(value),
+                Ordering::Equal => return Some(mid),
                 Ordering::Less => left = mid + 1,
                 Ordering::Greater => {
                     if mid == 0 {

tests/rust_usage_test/tests/integration_test.rs

@@ -3225,6 +3225,156 @@ fn test_shared_strings() {
 }
 
 #[test]
+fn lookup_index_by_key_returns_correct_index() {
+    let b = &mut flatbuffers::FlatBufferBuilder::new();
+    // Abilities are sorted by id (the key field).
+    let v = b.create_vector(&[
+        my_game::example::Ability::new(1, 10),
+        my_game::example::Ability::new(3, 30),
+        my_game::example::Ability::new(5, 50),
+    ]);
+    let name = b.create_string("test");
+    let mon = my_game::example::Monster::create(b, &my_game::example::MonsterArgs {
+        name: Some(name),
+        testarrayofsortedstruct: Some(v),
+        ..Default::default()
+    });
+    my_game::example::finish_monster_buffer(b, mon);
+    let buf = b.finished_data();
+    let mon = my_game::example::root_as_monster(buf).unwrap();
+    let abilities = mon.testarrayofsortedstruct().unwrap();
+
+    // Lookup each element and verify the returned index.
+    assert_eq!(
+        abilities.lookup_index_by_key(1u32, |a, key| a.key_compare_with_value(*key)),
+        Some(0)
+    );
+    assert_eq!(
+        abilities.lookup_index_by_key(3u32, |a, key| a.key_compare_with_value(*key)),
+        Some(1)
+    );
+    assert_eq!(
+        abilities.lookup_index_by_key(5u32, |a, key| a.key_compare_with_value(*key)),
+        Some(2)
+    );
+}
+
+#[test]
+fn lookup_index_by_key_returns_none_for_missing_key() {
+    let b = &mut flatbuffers::FlatBufferBuilder::new();
+    let v = b.create_vector(&[
+        my_game::example::Ability::new(1, 10),
+        my_game::example::Ability::new(3, 30),
+        my_game::example::Ability::new(5, 50),
+    ]);
+    let name = b.create_string("test");
+    let mon = my_game::example::Monster::create(b, &my_game::example::MonsterArgs {
+        name: Some(name),
+        testarrayofsortedstruct: Some(v),
+        ..Default::default()
+    });
+    my_game::example::finish_monster_buffer(b, mon);
+    let buf = b.finished_data();
+    let mon = my_game::example::root_as_monster(buf).unwrap();
+    let abilities = mon.testarrayofsortedstruct().unwrap();
+
+    // Keys that do not exist in the vector.
+    assert_eq!(
+        abilities.lookup_index_by_key(0u32, |a, key| a.key_compare_with_value(*key)),
+        None
+    );
+    assert_eq!(
+        abilities.lookup_index_by_key(2u32, |a, key| a.key_compare_with_value(*key)),
+        None
+    );
+    assert_eq!(
+        abilities.lookup_index_by_key(99u32, |a, key| a.key_compare_with_value(*key)),
+        None
+    );
+}
+
+#[test]
+fn lookup_index_by_key_on_empty_vector() {
+    let b = &mut flatbuffers::FlatBufferBuilder::new();
+    let v = b.create_vector::<my_game::example::Ability>(&[]);
+    let name = b.create_string("test");
+    let mon = my_game::example::Monster::create(b, &my_game::example::MonsterArgs {
+        name: Some(name),
+        testarrayofsortedstruct: Some(v),
+        ..Default::default()
+    });
+    my_game::example::finish_monster_buffer(b, mon);
+    let buf = b.finished_data();
+    let mon = my_game::example::root_as_monster(buf).unwrap();
+    let abilities = mon.testarrayofsortedstruct().unwrap();
+
+    assert_eq!(
+        abilities.lookup_index_by_key(1u32, |a, key| a.key_compare_with_value(*key)),
+        None
+    );
+}
+
+#[test]
+fn lookup_index_by_key_single_element() {
+    let b = &mut flatbuffers::FlatBufferBuilder::new();
+    let v = b.create_vector(&[my_game::example::Ability::new(42, 100)]);
+    let name = b.create_string("test");
+    let mon = my_game::example::Monster::create(b, &my_game::example::MonsterArgs {
+        name: Some(name),
+        testarrayofsortedstruct: Some(v),
+        ..Default::default()
+    });
+    my_game::example::finish_monster_buffer(b, mon);
+    let buf = b.finished_data();
+    let mon = my_game::example::root_as_monster(buf).unwrap();
+    let abilities = mon.testarrayofsortedstruct().unwrap();
+
+    assert_eq!(
+        abilities.lookup_index_by_key(42u32, |a, key| a.key_compare_with_value(*key)),
+        Some(0)
+    );
+    assert_eq!(
+        abilities.lookup_index_by_key(1u32, |a, key| a.key_compare_with_value(*key)),
+        None
+    );
+}
+
+#[test]
+fn lookup_index_by_key_consistent_with_lookup_by_key() {
+    let b = &mut flatbuffers::FlatBufferBuilder::new();
+    let v = b.create_vector(&[
+        my_game::example::Ability::new(2, 20),
+        my_game::example::Ability::new(4, 40),
+        my_game::example::Ability::new(6, 60),
+        my_game::example::Ability::new(8, 80),
+        my_game::example::Ability::new(10, 100),
+    ]);
+    let name = b.create_string("test");
+    let mon = my_game::example::Monster::create(b, &my_game::example::MonsterArgs {
+        name: Some(name),
+        testarrayofsortedstruct: Some(v),
+        ..Default::default()
+    });
+    my_game::example::finish_monster_buffer(b, mon);
+    let buf = b.finished_data();
+    let mon = my_game::example::root_as_monster(buf).unwrap();
+    let abilities = mon.testarrayofsortedstruct().unwrap();
+
+    // For every key that exists, lookup_index_by_key should return an index
+    // whose element matches lookup_by_key.
+    for key in &[2u32, 4, 6, 8, 10] {
+        let obj = abilities.lookup_by_key(*key, |a, k| a.key_compare_with_value(*k));
+        let idx = abilities.lookup_index_by_key(*key, |a, k| a.key_compare_with_value(*k));
+        assert!(obj.is_some());
+        assert!(idx.is_some());
+        assert_eq!(abilities.get(idx.unwrap()).id(), obj.unwrap().id());
+    }
+
+    // For keys that don't exist, both should return None.
+    for key in &[0u32, 1, 3, 5, 7, 9, 11] {
+        assert!(abilities.lookup_by_key(*key, |a, k| a.key_compare_with_value(*k)).is_none());
+        assert!(abilities.lookup_index_by_key(*key, |a, k| a.key_compare_with_value(*k)).is_none());
+    }
 fn test_shared_strings_pool_deduplication() {
     // Verifies that create_shared_string correctly deduplicates across many
     // unique strings and that the resulting buffer contains valid data.