diff --git a/rust/flatbuffers/src/vector.rs b/rust/flatbuffers/src/vector.rs index 2078a76df0..17a3a1698e 100644 --- a/rust/flatbuffers/src/vector.rs +++ b/rust/flatbuffers/src/vector.rs @@ -107,6 +107,20 @@ impl<'a, T: Follow<'a> + 'a> Vector<'a, T> { key: K, f: fn(&>::Inner, &K) -> Ordering, ) -> Option { + 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( + &self, + key: K, + f: fn(&>::Inner, &K) -> Ordering, + ) -> Option { 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 { diff --git a/tests/rust_usage_test/tests/integration_test.rs b/tests/rust_usage_test/tests/integration_test.rs index 38928d156e..c0fdab8178 100644 --- a/tests/rust_usage_test/tests/integration_test.rs +++ b/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::(&[]); + 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.