diff --git a/src/inclusive_map.rs b/src/inclusive_map.rs index ae9c89e..c3d375b 100644 --- a/src/inclusive_map.rs +++ b/src/inclusive_map.rs @@ -189,6 +189,12 @@ where self.get_key_value(key).map(|(_range, value)| value) } + /// Returns a mutable reference to the value corresponding to the given key, + /// if the key is covered by any range in the map. + pub fn get_mut(&mut self, key: &K) -> Option<&mut V>{ + self.get_key_value_mut(key).map(|(_, value)| value) + } + /// Returns the range-value pair (as a pair of references) corresponding /// to the given key, if the key is covered by any range in the map. pub fn get_key_value(&self, key: &K) -> Option<(&RangeInclusive, &V)> { @@ -208,6 +214,25 @@ where .map(|(range_start_wrapper, value)| (&range_start_wrapper.range, value)) } + /// Returns the range-value pair (immutable key reference and mutable value + /// reference) corresponding to the given key, if the key is covered by any range in the map. + pub fn get_key_value_mut(&mut self, key: &K) -> Option<(&RangeInclusive, &mut V)> { + use core::ops::Bound; + + // The only stored range that could contain the given key is the + // last stored range whose start is less than or equal to this key. + let key_as_start = RangeInclusiveStartWrapper::new(key.clone()..=key.clone()); + self.btm + .range_mut((Bound::Unbounded, Bound::Included(key_as_start))) + .next_back() + .filter(|(range_start_wrapper, _value)| { + // Does the only candidate range contain + // the requested key? + range_start_wrapper.contains(key) + }) + .map(|(range_start_wrapper, value)| (&range_start_wrapper.range, value)) + } + /// Returns `true` if any range in the map covers the specified key. pub fn contains_key(&self, key: &K) -> bool { self.get(key).is_some() @@ -1348,6 +1373,17 @@ mod tests { assert_eq!(range_map.get(&51), None); } + #[test] + fn get_mut(){ + let mut range_map: RangeInclusiveMap = RangeInclusiveMap::new(); + range_map.insert(0..=50, false); + if let Some(value) = range_map.get_mut(&50){ + *value = true; + } + assert_eq!(range_map.get(&50), Some(&true)); + assert_eq!(range_map.get(&51), None); + } + #[test] fn get_key_value() { let mut range_map: RangeInclusiveMap = RangeInclusiveMap::new(); @@ -1356,6 +1392,16 @@ mod tests { assert_eq!(range_map.get_key_value(&51), None); } + #[test] + fn get_key_value_mut(){ + let mut range_map: RangeInclusiveMap = RangeInclusiveMap::new(); + range_map.insert(0..=50, false); + if let Some((_, value)) = range_map.get_key_value_mut(&50){ + *value = true + } + assert_eq!(range_map.get_key_value(&50), Some((&(0..=50), &true))); + assert_eq!(range_map.get_key_value(&51), None); + } // // Removal tests // diff --git a/src/map.rs b/src/map.rs index 7296c5b..a25e6df 100644 --- a/src/map.rs +++ b/src/map.rs @@ -158,6 +158,12 @@ where self.get_key_value(key).map(|(_range, value)| value) } + /// Returns a mutable reference to the value corresponding to the given key, + /// if the key is covered by any range in the map. + pub fn get_mut(&mut self, key: &K) -> Option<&mut V>{ + self.get_key_value_mut(key).map(|(_, value)| value) + } + /// Returns the range-value pair (as a pair of references) corresponding /// to the given key, if the key is covered by any range in the map. pub fn get_key_value(&self, key: &K) -> Option<(&Range, &V)> { @@ -175,6 +181,23 @@ where .map(|(start_wrapper, value)| (&start_wrapper.end_wrapper.range, value)) } + /// Returns the range-value pair (immutable key reference and mutable value + /// reference) corresponding to the given key, if the key is covered by any range in the map. + pub fn get_key_value_mut(&mut self, key: &K) -> Option<(&Range, &mut V)> { + // The only stored range that could contain the given key is the + // last stored range whose start is less than or equal to this key. + let key_as_start = RangeStartWrapper::new(key.clone()..key.clone()); + self.btm + .range_mut((Bound::Unbounded, Bound::Included(key_as_start))) + .next_back() + .filter(|(start_wrapper, _value)| { + // Does the only candidate range contain + // the requested key? + start_wrapper.end_wrapper.range.contains(key) + }) + .map(|(start_wrapper, value)| (&start_wrapper.end_wrapper.range, value)) + } + /// Returns `true` if any range in the map covers the specified key. pub fn contains_key(&self, key: &K) -> bool { self.get(key).is_some() @@ -1232,6 +1255,18 @@ mod tests { assert_eq!(range_map.get(&50), None); } + + #[test] + fn get_mut(){ + let mut range_map: RangeMap = RangeMap::new(); + range_map.insert(0..50, false); + if let Some(value) = range_map.get_mut(&49){ + *value = true; + } + assert_eq!(range_map.get(&49), Some(&true)); + assert_eq!(range_map.get(&50), None); + } + #[test] fn get_key_value() { let mut range_map: RangeMap = RangeMap::new(); @@ -1240,6 +1275,17 @@ mod tests { assert_eq!(range_map.get_key_value(&50), None); } + #[test] + fn get_key_value_mut(){ + let mut range_map: RangeMap = RangeMap::new(); + range_map.insert(0..50, false); + if let Some((_, value)) = range_map.get_key_value_mut(&49){ + *value = true + } + assert_eq!(range_map.get_key_value(&49), Some((&(0..50), &true))); + assert_eq!(range_map.get_key_value(&50), None); + } + // // Removal tests //