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
diff --git a/rust/flatbuffers/src/vector.rs b/rust/flatbuffers/src/vector.rs
index 2078a76..17a3a16 100644
--- a/rust/flatbuffers/src/vector.rs
+++ b/rust/flatbuffers/src/vector.rs
@@ -107,6 +107,20 @@
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 @@
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 {
diff --git a/tests/rust_usage_test/tests/integration_test.rs b/tests/rust_usage_test/tests/integration_test.rs
index 38928d1..c0fdab8 100644
--- a/tests/rust_usage_test/tests/integration_test.rs
+++ b/tests/rust_usage_test/tests/integration_test.rs
@@ -3225,6 +3225,156 @@
}
#[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.
