Skip to content

Commit 5b7dfe7

Browse files
progvalprogval
committed
Add ManagedMap::range.
1 parent 2da4c59 commit 5b7dfe7

File tree

1 file changed

+266
-2
lines changed

1 file changed

+266
-2
lines changed

‎src/map.rs

Lines changed: 266 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,14 @@ use core::borrow::Borrow;
66
#[cfg(feature = "std")]
77
use std::collections::BTreeMap;
88
#[cfg(feature = "std")]
9-
use std::collections::btree_map::{Iter as BTreeIter, IterMut as BTreeIterMut};
9+
use std::collections::btree_map::{Iter as BTreeIter, IterMut as BTreeIterMut, Range as BTreeRange};
10+
#[cfg(feature = "std")]
11+
use std::collections::Bound::{Included, Excluded};
1012
#[cfg(all(feature = "alloc", not(feature = "std")))]
1113
use alloc::btree_map::BTreeMap;
1214
#[cfg(all(feature = "alloc", not(feature = "std")))]
13-
use alloc::btree_map::{Iter as BTreeIter, IterMut as BTreeIterMut};
15+
use alloc::btree_map::{Iter as BTreeIter, IterMut as BTreeIterMut, Range as BTreeRange};
16+
use core::ops::Bound::{Included,Excluded};
1417

1518
/// A managed map.
1619
///
@@ -90,6 +93,92 @@ impl<T> Into<Option<T>> for RevOption<T> {
9093
}
9194
}
9295

96+
pub enum Range<'a, K: 'a, V: 'a> {
97+
/// Borrowed variant.
98+
Borrowed(&'a [Option<(K, V)>], usize),
99+
/// Owned variant, only available with the `std` or `alloc` feature enabled.
100+
#[cfg(any(feature = "std", feature = "alloc"))]
101+
Owned(BTreeRange<'a, K, V>),
102+
}
103+
104+
impl<'a, K: 'a, V: 'a> Iterator for Range<'a, K, V> {
105+
type Item = (&'a K, &'a V);
106+
107+
fn next(&mut self) -> Option<Self::Item> {
108+
match *self {
109+
Range::Borrowed(ref slice, ref mut index) => {
110+
*index += 1;
111+
if *index-1 >= slice.len() {
112+
None
113+
} else {
114+
match slice[*index-1] {
115+
None => None,
116+
Some((ref k, ref v)) => Some((k, v))
117+
}
118+
}
119+
},
120+
#[cfg(any(feature = "std", feature = "alloc"))]
121+
Range::Owned(ref mut range) => range.next(),
122+
}
123+
}
124+
}
125+
126+
fn binary_search_by_key_range<K, V, Q>(slice: &[Option<(K, V)>], key_begin: &Q, key_end: &Q) -> Result<(usize, usize), ()>
127+
where K: Ord + Borrow<Q>, Q: Ord + ?Sized
128+
{
129+
if slice.len() == 0 {
130+
return Err(())
131+
}
132+
let (mut left, mut right) = (0, slice.len()-1);
133+
134+
//let key = |entry: Option<(K, V)>| { entry.as_ref().map(|&(ref key, _)| key.borrow()) };
135+
macro_rules! key {
136+
( $e:expr) => { $e.as_ref().map(|&(ref key, _)| key.borrow()) }
137+
}
138+
139+
// We cannot use slice.binary_search_by_key instead of each of the
140+
// two loops here, because we need the left-most match (for begin) and
141+
// the right-most match (for end), and the stdlib does not provide
142+
// any of these guarantees.
143+
144+
// Find the beginning
145+
while left < right {
146+
let middle = left + (right-left)/2;
147+
if key!(slice[middle]) < Some(key_begin) {
148+
left = middle+1;
149+
} else if middle > 0 && key!(slice[middle-1]) >= Some(key_begin) {
150+
right = middle-1;
151+
} else {
152+
left = middle;
153+
break;
154+
}
155+
}
156+
let begin = left;
157+
if key!(slice[begin]) < Some(key_begin) {
158+
return Err(())
159+
}
160+
161+
// Find the ending
162+
right = slice.len(); // no need to reset left
163+
while left < right {
164+
let middle = left + (right-left+1)/2;
165+
if key!(slice[middle-1]) >= Some(key_end) {
166+
right = middle-1;
167+
} else if middle < slice.len() && key!(slice[middle]) < Some(key_end) {
168+
left = middle+1;
169+
} else {
170+
right = middle;
171+
break;
172+
}
173+
}
174+
let end = right;
175+
if end == 0 || key!(slice[end-1]) > Some(key_end) {
176+
return Err(())
177+
}
178+
179+
Ok((begin, end))
180+
}
181+
93182
fn binary_search_by_key<K, V, Q>(slice: &[Option<(K, V)>], key: &Q) -> Result<usize, usize>
94183
where K: Ord + Borrow<Q>, Q: Ord + ?Sized
95184
{
@@ -155,6 +244,23 @@ impl<'a, K: Ord + 'a, V: 'a> ManagedMap<'a, K, V> {
155244
}
156245
}
157246

247+
pub fn range<'b, Q>(&'b self, key_begin: &Q, key_end: &Q) -> Range<'a, K, V>
248+
where K: Borrow<Q>, Q: Ord + ?Sized, 'b: 'a
249+
{
250+
match self {
251+
&ManagedMap::Borrowed(ref pairs) => {
252+
match binary_search_by_key_range(&pairs[0..self.len()], key_begin, key_end) {
253+
Ok((begin, end)) => Range::Borrowed(&pairs[begin..end], 0),
254+
Err(()) => Range::Borrowed(&[], 0),
255+
}
256+
},
257+
#[cfg(any(feature = "std", feature = "alloc"))]
258+
&ManagedMap::Owned(ref map) => {
259+
Range::Owned(map.range((Included(key_begin), Excluded(key_end))))
260+
},
261+
}
262+
}
263+
158264
pub fn insert(&mut self, key: K, new_value: V) -> Result<Option<V>, (K, V)> {
159265
match self {
160266
&mut ManagedMap::Borrowed(ref mut pairs) if pairs.len() < 1 =>
@@ -379,6 +485,164 @@ mod test {
379485
assert_eq!(map.get("q"), None);
380486
}
381487

488+
#[test]
489+
fn test_get_none_empty() {
490+
let mut pairs = all_pairs_empty();
491+
let map = ManagedMap::Borrowed(&mut pairs);
492+
assert_eq!(map.len(), 0);
493+
assert!(map.is_empty());
494+
assert_eq!(map.get("q"), None);
495+
}
496+
497+
#[test]
498+
fn test_range_full() {
499+
let mut pairs = all_pairs_full();
500+
let map = ManagedMap::Borrowed(&mut pairs);
501+
assert_eq!(map.len(), 4);
502+
503+
let mut range = map.range("0", "a");
504+
assert_eq!(range.next(), None);
505+
let mut range = map.range("0", "b");
506+
assert_eq!(range.next(), Some((&"a", &1)));
507+
assert_eq!(range.next(), None);
508+
let mut range = map.range("0", "c");
509+
assert_eq!(range.next(), Some((&"a", &1)));
510+
assert_eq!(range.next(), Some((&"b", &2)));
511+
assert_eq!(range.next(), None);
512+
let mut range = map.range("0", "d");
513+
assert_eq!(range.next(), Some((&"a", &1)));
514+
assert_eq!(range.next(), Some((&"b", &2)));
515+
assert_eq!(range.next(), Some((&"c", &3)));
516+
assert_eq!(range.next(), None);
517+
let mut range = map.range("0", "e");
518+
assert_eq!(range.next(), Some((&"a", &1)));
519+
assert_eq!(range.next(), Some((&"b", &2)));
520+
assert_eq!(range.next(), Some((&"c", &3)));
521+
assert_eq!(range.next(), Some((&"d", &4)));
522+
assert_eq!(range.next(), None);
523+
524+
let mut range = map.range("a", "a");
525+
assert_eq!(range.next(), None);
526+
let mut range = map.range("a", "b");
527+
assert_eq!(range.next(), Some((&"a", &1)));
528+
assert_eq!(range.next(), None);
529+
let mut range = map.range("a", "c");
530+
assert_eq!(range.next(), Some((&"a", &1)));
531+
assert_eq!(range.next(), Some((&"b", &2)));
532+
assert_eq!(range.next(), None);
533+
let mut range = map.range("a", "d");
534+
assert_eq!(range.next(), Some((&"a", &1)));
535+
assert_eq!(range.next(), Some((&"b", &2)));
536+
assert_eq!(range.next(), Some((&"c", &3)));
537+
assert_eq!(range.next(), None);
538+
let mut range = map.range("a", "e");
539+
assert_eq!(range.next(), Some((&"a", &1)));
540+
assert_eq!(range.next(), Some((&"b", &2)));
541+
assert_eq!(range.next(), Some((&"c", &3)));
542+
assert_eq!(range.next(), Some((&"d", &4)));
543+
assert_eq!(range.next(), None);
544+
545+
let mut range = map.range("b", "a");
546+
assert_eq!(range.next(), None);
547+
let mut range = map.range("b", "b");
548+
assert_eq!(range.next(), None);
549+
let mut range = map.range("b", "c");
550+
assert_eq!(range.next(), Some((&"b", &2)));
551+
assert_eq!(range.next(), None);
552+
let mut range = map.range("b", "d");
553+
assert_eq!(range.next(), Some((&"b", &2)));
554+
assert_eq!(range.next(), Some((&"c", &3)));
555+
assert_eq!(range.next(), None);
556+
let mut range = map.range("b", "e");
557+
assert_eq!(range.next(), Some((&"b", &2)));
558+
assert_eq!(range.next(), Some((&"c", &3)));
559+
assert_eq!(range.next(), Some((&"d", &4)));
560+
assert_eq!(range.next(), None);
561+
562+
let mut range = map.range("c", "a");
563+
assert_eq!(range.next(), None);
564+
let mut range = map.range("c", "b");
565+
assert_eq!(range.next(), None);
566+
let mut range = map.range("c", "c");
567+
assert_eq!(range.next(), None);
568+
let mut range = map.range("c", "d");
569+
assert_eq!(range.next(), Some((&"c", &3)));
570+
assert_eq!(range.next(), None);
571+
let mut range = map.range("c", "e");
572+
assert_eq!(range.next(), Some((&"c", &3)));
573+
assert_eq!(range.next(), Some((&"d", &4)));
574+
assert_eq!(range.next(), None);
575+
576+
let mut range = map.range("d", "a");
577+
assert_eq!(range.next(), None);
578+
let mut range = map.range("d", "b");
579+
assert_eq!(range.next(), None);
580+
let mut range = map.range("d", "c");
581+
assert_eq!(range.next(), None);
582+
let mut range = map.range("d", "d");
583+
assert_eq!(range.next(), None);
584+
let mut range = map.range("d", "e");
585+
assert_eq!(range.next(), Some((&"d", &4)));
586+
assert_eq!(range.next(), None);
587+
588+
let mut range = map.range("e", "a");
589+
assert_eq!(range.next(), None);
590+
let mut range = map.range("e", "b");
591+
assert_eq!(range.next(), None);
592+
let mut range = map.range("e", "c");
593+
assert_eq!(range.next(), None);
594+
let mut range = map.range("e", "d");
595+
assert_eq!(range.next(), None);
596+
let mut range = map.range("e", "e");
597+
assert_eq!(range.next(), None);
598+
}
599+
600+
#[test]
601+
fn test_range_one_pair() {
602+
let mut pairs = one_pair_full();
603+
let map = ManagedMap::Borrowed(&mut pairs);
604+
assert_eq!(map.len(), 1);
605+
606+
let mut range = map.range("0", "a");
607+
assert_eq!(range.next(), None);
608+
let mut range = map.range("0", "b");
609+
assert_eq!(range.next(), Some((&"a", &1)));
610+
assert_eq!(range.next(), None);
611+
let mut range = map.range("0", "c");
612+
assert_eq!(range.next(), Some((&"a", &1)));
613+
assert_eq!(range.next(), None);
614+
615+
let mut range = map.range("a", "a");
616+
assert_eq!(range.next(), None);
617+
let mut range = map.range("a", "b");
618+
assert_eq!(range.next(), Some((&"a", &1)));
619+
assert_eq!(range.next(), None);
620+
let mut range = map.range("a", "c");
621+
assert_eq!(range.next(), Some((&"a", &1)));
622+
assert_eq!(range.next(), None);
623+
624+
let mut range = map.range("b", "a");
625+
assert_eq!(range.next(), None);
626+
let mut range = map.range("b", "b");
627+
assert_eq!(range.next(), None);
628+
let mut range = map.range("b", "c");
629+
assert_eq!(range.next(), None);
630+
}
631+
632+
#[test]
633+
fn test_range_empty() {
634+
let mut pairs = all_pairs_empty();
635+
let map = ManagedMap::Borrowed(&mut pairs);
636+
assert_eq!(map.len(), 0);
637+
638+
let mut range = map.range("b", "a");
639+
assert_eq!(range.next(), None);
640+
let mut range = map.range("b", "b");
641+
assert_eq!(range.next(), None);
642+
let mut range = map.range("b", "c");
643+
assert_eq!(range.next(), None);
644+
}
645+
382646
#[test]
383647
fn test_get_mut_some() {
384648
let mut pairs = all_pairs_full();

0 commit comments

Comments
 (0)