diff --git a/data/src/partial_vec.rs b/data/src/partial_vec.rs index e5ed03c174d..911883826fa 100644 --- a/data/src/partial_vec.rs +++ b/data/src/partial_vec.rs @@ -159,6 +159,34 @@ impl PartialVec { self.entries[idx + 1].start } + /// Retrieve a reference to the element at the given index, if defined. + pub fn get(&self, idx: usize) -> Option<&T> { + let entry_idx = self.entries.partition_point(|entry| entry.end() <= idx); + + if entry_idx >= self.entries.len() { + // All entries are before the requested index. + return None; + } + + let entry = &self.entries[entry_idx]; + let local_idx = idx.checked_sub(entry.start)?; + entry.data.get(local_idx) + } + + /// Retrieve a mutable reference to the element at the given index, if defined. + pub fn get_mut(&mut self, idx: usize) -> Option<&mut T> { + let entry_idx = self.entries.partition_point(|entry| entry.end() <= idx); + + if entry_idx >= self.entries.len() { + // All entries are before the requested index. + return None; + } + + let entry = &mut self.entries[entry_idx]; + let local_idx = idx.checked_sub(entry.start)?; + entry.data.get_mut(local_idx) + } + /// Define a range within the partial vector. /// /// Existing data is overwritten if it overlaps the newly defined range. diff --git a/data/src/partial_vec/tests.rs b/data/src/partial_vec/tests.rs index 9cb7f3883e5..ec42aed28b6 100644 --- a/data/src/partial_vec/tests.rs +++ b/data/src/partial_vec/tests.rs @@ -4,6 +4,8 @@ //! Tests for [`PartialVec`] +use std::collections::BTreeMap; + use proptest::arbitrary::any; use proptest::collection::vec; use proptest::prelude::Strategy; @@ -45,6 +47,53 @@ proptest! { added_so_far += chunk_len; } } + + /// `get` returns the correct value for any index in defined ranges. + #[test] + fn get_returns_defined_data( + entries in vec((..100usize, data_vec()), ..16), + query_indices in vec(..200usize, ..32) + ) { + let mut vec = PartialVec::empty(); + let mut reference = BTreeMap::new(); + + // Build the PartialVec and a reference map (later writes overwrite earlier ones) + for (offset, data) in entries { + for (i, &val) in data.iter().enumerate() { + reference.insert(offset.saturating_add(i), val); + } + + vec.define(offset, data); + } + + // Getting should match the reference for all queried indices + for query_idx in query_indices { + assert_eq!(vec.get(query_idx), reference.get(&query_idx)); + } + } + + /// `get_mut` allows modifying values and subsequent `get` reflects the change. + #[test] + fn get_mut_modifies_data( + offset in 0usize..1000, + data in vec(any::(), 1..64), + local_idx in any::(), + new_value in any::() + ) { + let mut vec = PartialVec::empty(); + vec.define(offset, data.clone()); + + let local_idx = local_idx.index(data.len()); + let idx = offset + local_idx; + + // Modify via get_mut + if let Some(elem) = vec.get_mut(idx) { + *elem = new_value; + } + + // Verify the modification + assert_eq!(vec.get(idx), Some(&new_value)); + } } /// Ensure defining continuous ranges can be recovered.