From 993bca8096d9bcc4282f2f5ba389a943e300c48a Mon Sep 17 00:00:00 2001 From: statxc Date: Fri, 6 Mar 2026 21:20:12 +0000 Subject: [PATCH 1/2] feat: add lookup_index_by_key to Rust Vector for index-based binary search --- rust/flatbuffers/src/vector.rs | 36 +++++ .../rust_usage_test/tests/integration_test.rs | 153 ++++++++++++++++++ 2 files changed, 189 insertions(+) diff --git a/rust/flatbuffers/src/vector.rs b/rust/flatbuffers/src/vector.rs index 2078a76df0..83bb399a2f 100644 --- a/rust/flatbuffers/src/vector.rs +++ b/rust/flatbuffers/src/vector.rs @@ -132,6 +132,42 @@ impl<'a, T: Follow<'a> + 'a> Vector<'a, T> { None } + /// 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( + &self, + key: K, + f: fn(&>::Inner, &K) -> Ordering, + ) -> Option { + if self.is_empty() { + return None; + } + + let mut left: usize = 0; + let mut right = self.len() - 1; + + while left <= right { + let mid = (left + right) / 2; + let value = self.get(mid); + match f(&value, &key) { + Ordering::Equal => return Some(mid), + Ordering::Less => left = mid + 1, + Ordering::Greater => { + if mid == 0 { + return None; + } + right = mid - 1; + } + } + } + + None + } + #[inline(always)] pub fn iter(&self) -> VectorIter<'a, T> { VectorIter::from_vector(*self) diff --git a/tests/rust_usage_test/tests/integration_test.rs b/tests/rust_usage_test/tests/integration_test.rs index d97bf88cff..929994cc5b 100644 --- a/tests/rust_usage_test/tests/integration_test.rs +++ b/tests/rust_usage_test/tests/integration_test.rs @@ -3224,4 +3224,157 @@ fn test_shared_strings() { assert_eq!(string_vector.get(1), "foo"); } +#[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::(&[]); + 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()); + } +} + } From 946077cb9b88d8d77ee8a091077d105a0c108b28 Mon Sep 17 00:00:00 2001 From: statxc Date: Sat, 7 Mar 2026 04:28:53 +0000 Subject: [PATCH 2/2] fix: remove duplicated code --- rust/flatbuffers/src/vector.rs | 24 +----------------------- 1 file changed, 1 insertion(+), 23 deletions(-) diff --git a/rust/flatbuffers/src/vector.rs b/rust/flatbuffers/src/vector.rs index 83bb399a2f..17a3a1698e 100644 --- a/rust/flatbuffers/src/vector.rs +++ b/rust/flatbuffers/src/vector.rs @@ -107,29 +107,7 @@ impl<'a, T: Follow<'a> + 'a> Vector<'a, T> { key: K, f: fn(&>::Inner, &K) -> Ordering, ) -> Option { - if self.is_empty() { - return None; - } - - let mut left: usize = 0; - let mut right = self.len() - 1; - - while left <= right { - let mid = (left + right) / 2; - let value = self.get(mid); - match f(&value, &key) { - Ordering::Equal => return Some(value), - Ordering::Less => left = mid + 1, - Ordering::Greater => { - if mid == 0 { - return None; - } - right = mid - 1; - } - } - } - - None + self.lookup_index_by_key(key, f).map(|idx| self.get(idx)) } /// Binary search by key, returning the index of the matching element.