|
| 1 | +use std::{collections::BTreeMap, ops::Bound}; |
| 2 | + |
| 3 | +pub(crate) trait Neighbors<K, V> |
| 4 | +where |
| 5 | + K: Ord + Eq, |
| 6 | +{ |
| 7 | + fn get_neighbors(&self, key: &K) -> (Option<&V>, Option<&V>); |
| 8 | + |
| 9 | + fn lo_bound(&self, bound: Bound<&K>) -> Option<(&K, &V)>; |
| 10 | + fn up_bound(&self, bound: Bound<&K>) -> Option<(&K, &V)>; |
| 11 | +} |
| 12 | + |
| 13 | +impl<K, V> Neighbors<K, V> for BTreeMap<K, V> |
| 14 | +where |
| 15 | + K: Ord + Eq, |
| 16 | +{ |
| 17 | + fn get_neighbors(&self, key: &K) -> (Option<&V>, Option<&V>) { |
| 18 | + // NOTE (@Techassi): These functions might get added to the standard |
| 19 | + // library at some point. If that's the case, we can use the ones |
| 20 | + // provided by the standard lib. |
| 21 | + // See: https://github.com/rust-lang/rust/issues/107540 |
| 22 | + match ( |
| 23 | + self.lo_bound(Bound::Excluded(key)), |
| 24 | + self.up_bound(Bound::Excluded(key)), |
| 25 | + ) { |
| 26 | + (Some((k, v)), None) => { |
| 27 | + if key > k { |
| 28 | + (Some(v), None) |
| 29 | + } else { |
| 30 | + (self.lo_bound(Bound::Excluded(k)).map(|(_, v)| v), None) |
| 31 | + } |
| 32 | + } |
| 33 | + (None, Some((k, v))) => { |
| 34 | + if key < k { |
| 35 | + (None, Some(v)) |
| 36 | + } else { |
| 37 | + (None, self.up_bound(Bound::Excluded(k)).map(|(_, v)| v)) |
| 38 | + } |
| 39 | + } |
| 40 | + (Some((_, lo)), Some((_, up))) => (Some(lo), Some(up)), |
| 41 | + (None, None) => unreachable!(), |
| 42 | + } |
| 43 | + } |
| 44 | + |
| 45 | + fn lo_bound(&self, bound: Bound<&K>) -> Option<(&K, &V)> { |
| 46 | + self.range((Bound::Unbounded, bound)).next_back() |
| 47 | + } |
| 48 | + |
| 49 | + fn up_bound(&self, bound: Bound<&K>) -> Option<(&K, &V)> { |
| 50 | + self.range((bound, Bound::Unbounded)).next() |
| 51 | + } |
| 52 | +} |
| 53 | + |
| 54 | +#[cfg(test)] |
| 55 | +mod test { |
| 56 | + use super::*; |
| 57 | + use rstest::rstest; |
| 58 | + |
| 59 | + #[rstest] |
| 60 | + #[case(0, (None, Some(&"test1")))] |
| 61 | + #[case(1, (None, Some(&"test3")))] |
| 62 | + #[case(2, (Some(&"test1"), Some(&"test3")))] |
| 63 | + #[case(3, (Some(&"test1"), None))] |
| 64 | + #[case(4, (Some(&"test3"), None))] |
| 65 | + fn test(#[case] key: i32, #[case] expected: (Option<&&str>, Option<&&str>)) { |
| 66 | + let map = BTreeMap::from([(1, "test1"), (3, "test3")]); |
| 67 | + let neigbors = map.get_neighbors(&key); |
| 68 | + |
| 69 | + assert_eq!(neigbors, expected); |
| 70 | + } |
| 71 | +} |
0 commit comments