From 429c23d5f4de5274cbcc22e72f6a2e3a00069817 Mon Sep 17 00:00:00 2001 From: Piotr Czarnecki Date: Tue, 13 Jan 2015 21:55:44 +0100 Subject: [PATCH] Implement range and range_mut for BTree Simplify BTree's iterators, too. --- src/libcollections/btree/map.rs | 467 ++++++++++++++++++++++--------- src/libcollections/btree/node.rs | 292 ++++++++++++++----- src/libcollections/btree/set.rs | 46 +++ src/libcollections/lib.rs | 11 +- src/libstd/collections/mod.rs | 1 + 5 files changed, 614 insertions(+), 203 deletions(-) diff --git a/src/libcollections/btree/map.rs b/src/libcollections/btree/map.rs index 3ac6b2775bf..c56592177b4 100644 --- a/src/libcollections/btree/map.rs +++ b/src/libcollections/btree/map.rs @@ -27,6 +27,7 @@ use core::iter::{Map, FromIterator}; use core::ops::{Index, IndexMut}; use core::{iter, fmt, mem}; +use Bound::{self, Included, Excluded, Unbounded}; use ring_buf::RingBuf; @@ -37,8 +38,6 @@ use super::node::{Traversal, MutTraversal, MoveTraversal}; use super::node::{self, Node, Found, GoDown}; -// FIXME(conventions): implement bounded iterators - /// A map based on a B-Tree. /// /// B-Trees represent a fundamental compromise between cache-efficiency and actually minimizing @@ -92,9 +91,7 @@ pub struct BTreeMap { /// An abstract base over-which all other BTree iterators are built. struct AbsIter { - lca: T, - left: RingBuf, - right: RingBuf, + traversals: RingBuf, size: uint, } @@ -128,6 +125,16 @@ pub struct Values<'a, K: 'a, V: 'a> { inner: Map<(&'a K, &'a V), &'a V, Iter<'a, K, V>, fn((&'a K, &'a V)) -> &'a V> } +/// An iterator over a sub-range of BTreeMap's entries. +pub struct Range<'a, K: 'a, V: 'a> { + inner: AbsIter> +} + +/// A mutable iterator over a sub-range of BTreeMap's entries. +pub struct RangeMut<'a, K: 'a, V: 'a> { + inner: AbsIter> +} + /// A view into a single entry in a map, which may either be vacant or occupied. #[unstable = "precise API still under development"] pub enum Entry<'a, K:'a, V:'a> { @@ -924,74 +931,45 @@ fn traverse(node: Node) -> MoveTraversal { } /// Represents an operation to perform inside the following iterator methods. -/// This is necessary to use in `next` because we want to modify self.left inside -/// a match that borrows it. Similarly, in `next_back` for self.right. Instead, we use this -/// enum to note what we want to do, and do it after the match. +/// This is necessary to use in `next` because we want to modify `self.traversals` inside +/// a match that borrows it. Similarly in `next_back`. Instead, we use this enum to note +/// what we want to do, and do it after the match. enum StackOp { Push(T), Pop, } - impl Iterator for AbsIter where T: DoubleEndedIterator> + Traverse, { type Item = (K, V); - // This function is pretty long, but only because there's a lot of cases to consider. - // Our iterator represents two search paths, left and right, to the smallest and largest - // elements we have yet to yield. lca represents the least common ancestor of these two paths, - // above-which we never walk, since everything outside it has already been consumed (or was - // never in the range to iterate). - // - // Note that the design of these iterators permits an *arbitrary* initial pair of min and max, - // making these arbitrary sub-range iterators. However the logic to construct these paths - // efficiently is fairly involved, so this is a FIXME. The sub-range iterators also wouldn't be - // able to accurately predict size, so those iterators can't implement ExactSizeIterator. + // Our iterator represents a queue of all ancestors of elements we have + // yet to yield, from smallest to largest. Note that the design of these + // iterators permits an *arbitrary* initial pair of min and max, making + // these arbitrary sub-range iterators. fn next(&mut self) -> Option<(K, V)> { loop { - // We want the smallest element, so try to get the top of the left stack - let op = match self.left.back_mut() { - // The left stack is empty, so try to get the next element of the two paths - // LCAs (the left search path is currently a subpath of the right one) - None => match self.lca.next() { - // The lca has been exhausted, walk further down the right path - None => match self.right.pop_front() { - // The right path is exhausted, so we're done - None => return None, - // The right path had something, make that the new LCA - // and restart the whole process - Some(right) => { - self.lca = right; - continue; - } - }, - // The lca yielded an edge, make that the new head of the left path - Some(Edge(next)) => Push(Traverse::traverse(next)), - // The lca yielded an entry, so yield that - Some(Elem(k, v)) => { - self.size -= 1; - return Some((k, v)) - } - }, - // The left stack wasn't empty, so continue along the node in its head + // We want the smallest element, so try to get the back of the queue + let op = match self.traversals.back_mut() { + None => return None, + // The queue wasn't empty, so continue along the node in its head Some(iter) => match iter.next() { - // The head of the left path is empty, so Pop it off and restart the process + // The head is empty, so Pop it off and continue the process None => Pop, - // The head of the left path yielded an edge, so make that the new head - // of the left path + // The head yielded an edge, so make that the new head Some(Edge(next)) => Push(Traverse::traverse(next)), - // The head of the left path yielded entry, so yield that - Some(Elem(k, v)) => { + // The head yielded an entry, so yield that + Some(Elem(kv)) => { self.size -= 1; - return Some((k, v)) + return Some(kv) } } }; - // Handle any operation on the left stack as necessary + // Handle any operation as necessary, without a conflicting borrow of the queue match op { - Push(item) => { self.left.push_back(item); }, - Pop => { self.left.pop_back(); }, + Push(item) => { self.traversals.push_back(item); }, + Pop => { self.traversals.pop_back(); }, } } } @@ -1005,36 +983,24 @@ impl DoubleEndedIterator for AbsIter where T: DoubleEndedIterator> + Traverse, { // next_back is totally symmetric to next + #[inline] fn next_back(&mut self) -> Option<(K, V)> { loop { - let op = match self.right.back_mut() { - None => match self.lca.next_back() { - None => match self.left.pop_front() { - None => return None, - Some(left) => { - self.lca = left; - continue; - } - }, - Some(Edge(next)) => Push(Traverse::traverse(next)), - Some(Elem(k, v)) => { - self.size -= 1; - return Some((k, v)) - } - }, + let op = match self.traversals.front_mut() { + None => return None, Some(iter) => match iter.next_back() { None => Pop, Some(Edge(next)) => Push(Traverse::traverse(next)), - Some(Elem(k, v)) => { + Some(Elem(kv)) => { self.size -= 1; - return Some((k, v)) + return Some(kv) } } }; match op { - Push(item) => { self.right.push_back(item); }, - Pop => { self.right.pop_back(); } + Push(item) => { self.traversals.push_front(item); }, + Pop => { self.traversals.pop_front(); } } } } @@ -1111,6 +1077,24 @@ fn next_back(&mut self) -> Option<(&'a V)> { self.inner.next_back() } #[stable] impl<'a, K, V> ExactSizeIterator for Values<'a, K, V> {} +impl<'a, K, V> Iterator for Range<'a, K, V> { + type Item = (&'a K, &'a V); + + fn next(&mut self) -> Option<(&'a K, &'a V)> { self.inner.next() } +} +impl<'a, K, V> DoubleEndedIterator for Range<'a, K, V> { + fn next_back(&mut self) -> Option<(&'a K, &'a V)> { self.inner.next_back() } +} + +impl<'a, K, V> Iterator for RangeMut<'a, K, V> { + type Item = (&'a K, &'a mut V); + + fn next(&mut self) -> Option<(&'a K, &'a mut V)> { self.inner.next() } +} +impl<'a, K, V> DoubleEndedIterator for RangeMut<'a, K, V> { + fn next_back(&mut self) -> Option<(&'a K, &'a mut V)> { self.inner.next_back() } +} + impl<'a, K: Ord, V> Entry<'a, K, V> { #[unstable = "matches collection reform v2 specification, waiting for dust to settle"] /// Returns a mutable reference to the entry if occupied, or the VacantEntry if vacant @@ -1188,11 +1172,12 @@ impl BTreeMap { #[stable] pub fn iter(&self) -> Iter { let len = self.len(); + // NB. The initial capacity for ringbuf is large enough to avoid reallocs in many cases. + let mut lca = RingBuf::new(); + lca.push_back(Traverse::traverse(&self.root)); Iter { inner: AbsIter { - lca: Traverse::traverse(&self.root), - left: RingBuf::new(), - right: RingBuf::new(), + traversals: lca, size: len, } } @@ -1220,11 +1205,11 @@ pub fn iter(&self) -> Iter { #[stable] pub fn iter_mut(&mut self) -> IterMut { let len = self.len(); + let mut lca = RingBuf::new(); + lca.push_back(Traverse::traverse(&mut self.root)); IterMut { inner: AbsIter { - lca: Traverse::traverse(&mut self.root), - left: RingBuf::new(), - right: RingBuf::new(), + traversals: lca, size: len, } } @@ -1249,11 +1234,11 @@ pub fn iter_mut(&mut self) -> IterMut { #[stable] pub fn into_iter(self) -> IntoIter { let len = self.len(); + let mut lca = RingBuf::new(); + lca.push_back(Traverse::traverse(self.root)); IntoIter { inner: AbsIter { - lca: Traverse::traverse(self.root), - left: RingBuf::new(), - right: RingBuf::new(), + traversals: lca, size: len, } } @@ -1334,7 +1319,189 @@ pub fn len(&self) -> uint { self.length } pub fn is_empty(&self) -> bool { self.len() == 0 } } +macro_rules! range_impl { + ($root:expr, $min:expr, $max:expr, $as_slices_internal:ident, $iter:ident, $Range:ident, + $edges:ident, [$($mutability:ident)*]) => ( + { + // A deque that encodes two search paths containing (left-to-right): + // a series of truncated-from-the-left iterators, the LCA's doubly-truncated iterator, + // and a series of truncated-from-the-right iterators. + let mut traversals = RingBuf::new(); + let (root, min, max) = ($root, $min, $max); + + let mut leftmost = None; + let mut rightmost = None; + + match (&min, &max) { + (&Unbounded, &Unbounded) => { + traversals.push_back(Traverse::traverse(root)) + } + (&Unbounded, &Included(_)) | (&Unbounded, &Excluded(_)) => { + rightmost = Some(root); + } + (&Included(_), &Unbounded) | (&Excluded(_), &Unbounded) => { + leftmost = Some(root); + } + (&Included(min_key), &Included(max_key)) + | (&Included(min_key), &Excluded(max_key)) + | (&Excluded(min_key), &Included(max_key)) + | (&Excluded(min_key), &Excluded(max_key)) => { + // lca represents the Lowest Common Ancestor, above which we never + // walk, since everything else is outside the range to iterate. + // ___________________ + // |__0_|_80_|_85_|_90_| (root) + // | | | | | + // | + // v + // ___________________ + // |__5_|_15_|_30_|_73_| + // | | | | | + // | + // v + // ___________________ + // |_33_|_58_|_63_|_68_| lca for the range [41, 65] + // | |\___|___/| | iterator at traversals[2] + // | | + // | v + // v rightmost + // leftmost + let mut is_leaf = root.is_leaf(); + let mut lca = root.$as_slices_internal(); + loop { + let slice = lca.slice_from(min_key).slice_to(max_key); + if let [ref $($mutability)* edge] = slice.edges { + // Follow the only edge that leads the node that covers the range. + is_leaf = edge.is_leaf(); + lca = edge.$as_slices_internal(); + } else { + let mut iter = slice.$iter(); + if is_leaf { + leftmost = None; + rightmost = None; + } else { + // Only change the state of nodes with edges. + leftmost = iter.next_edge_item(); + rightmost = iter.next_edge_item_back(); + } + traversals.push_back(iter); + break; + } + } + } + } + // Keep narrowing the range by going down. + // ___________________ + // |_38_|_43_|_48_|_53_| + // | |____|____|____/ iterator at traversals[1] + // | + // v + // ___________________ + // |_39_|_40_|_41_|_42_| (leaf, the last leftmost) + // \_________| iterator at traversals[0] + match min { + Included(key) | Excluded(key) => + while let Some(left) = leftmost { + let is_leaf = left.is_leaf(); + let mut iter = left.$as_slices_internal().slice_from(key).$iter(); + leftmost = if is_leaf { + None + } else { + // Only change the state of nodes with edges. + iter.next_edge_item() + }; + traversals.push_back(iter); + }, + _ => {} + } + // If the leftmost iterator starts with an element, then it was an exact match. + if let (Excluded(_), Some(leftmost_iter)) = (min, traversals.back_mut()) { + // Drop this excluded element. `next_kv_item` has no effect when + // the next item is an edge. + leftmost_iter.next_kv_item(); + } + + // The code for the right side is similar. + match max { + Included(key) | Excluded(key) => + while let Some(right) = rightmost { + let is_leaf = right.is_leaf(); + let mut iter = right.$as_slices_internal().slice_to(key).$iter(); + rightmost = if is_leaf { + None + } else { + iter.next_edge_item_back() + }; + traversals.push_front(iter); + }, + _ => {} + } + if let (Excluded(_), Some(rightmost_iter)) = (max, traversals.front_mut()) { + rightmost_iter.next_kv_item_back(); + } + + $Range { + inner: AbsIter { + traversals: traversals, + size: 0, // unused + } + } + } + ) +} + impl BTreeMap { + /// Constructs a double-ended iterator over a sub-range of elements in the map, starting + /// at min, and ending at max. If min is `Unbounded`, then it will be treated as "negative + /// infinity", and if max is `Unbounded`, then it will be treated as "positive infinity". + /// Thus range(Unbounded, Unbounded) will yield the whole collection. + /// + /// # Examples + /// + /// ``` + /// use std::collections::BTreeMap; + /// use std::collections::Bound::{Included, Unbounded}; + /// + /// let mut map = BTreeMap::new(); + /// map.insert(3u, "a"); + /// map.insert(5u, "b"); + /// map.insert(8u, "c"); + /// for (&key, &value) in map.range(Included(&4), Included(&8)) { + /// println!("{}: {}", key, value); + /// } + /// assert_eq!(Some((&5u, &"b")), map.range(Included(&4), Unbounded).next()); + /// ``` + #[unstable = "matches collection reform specification, waiting for dust to settle"] + pub fn range<'a>(&'a self, min: Bound<&K>, max: Bound<&K>) -> Range<'a, K, V> { + range_impl!(&self.root, min, max, as_slices_internal, iter, Range, edges, []) + } + + /// Constructs a mutable double-ended iterator over a sub-range of elements in the map, starting + /// at min, and ending at max. If min is `Unbounded`, then it will be treated as "negative + /// infinity", and if max is `Unbounded`, then it will be treated as "positive infinity". + /// Thus range(Unbounded, Unbounded) will yield the whole collection. + /// + /// # Examples + /// + /// ``` + /// use std::collections::BTreeMap; + /// use std::collections::Bound::{Included, Excluded}; + /// + /// let mut map: BTreeMap<&str, i32> = ["Alice", "Bob", "Carol", "Cheryl"].iter() + /// .map(|&s| (s, 0)) + /// .collect(); + /// for (_, balance) in map.range_mut(Included(&"B"), Excluded(&"Cheryl")) { + /// *balance += 100; + /// } + /// for (name, balance) in map.iter() { + /// println!("{} => {}", name, balance); + /// } + /// ``` + #[unstable = "matches collection reform specification, waiting for dust to settle"] + pub fn range_mut<'a>(&'a mut self, min: Bound<&K>, max: Bound<&K>) -> RangeMut<'a, K, V> { + range_impl!(&mut self.root, min, max, as_slices_internal_mut, iter_mut, RangeMut, + edges_mut, [mut]) + } + /// Gets the given key's corresponding entry in the map for in-place manipulation. /// /// # Examples @@ -1410,8 +1577,10 @@ pub fn entry<'a>(&'a mut self, mut key: K) -> Entry<'a, K, V> { #[cfg(test)] mod test { use prelude::*; + use std::iter::range_inclusive; use super::{BTreeMap, Occupied, Vacant}; + use Bound::{self, Included, Excluded, Unbounded}; #[test] fn test_basic_large() { @@ -1481,28 +1650,7 @@ fn test_iter() { // Forwards let mut map: BTreeMap = range(0, size).map(|i| (i, i)).collect(); - { - let mut iter = map.iter(); - for i in range(0, size) { - assert_eq!(iter.size_hint(), (size - i, Some(size - i))); - assert_eq!(iter.next().unwrap(), (&i, &i)); - } - assert_eq!(iter.size_hint(), (0, Some(0))); - assert_eq!(iter.next(), None); - } - - { - let mut iter = map.iter_mut(); - for i in range(0, size) { - assert_eq!(iter.size_hint(), (size - i, Some(size - i))); - assert_eq!(iter.next().unwrap(), (&i, &mut (i + 0))); - } - assert_eq!(iter.size_hint(), (0, Some(0))); - assert_eq!(iter.next(), None); - } - - { - let mut iter = map.into_iter(); + fn test(size: uint, mut iter: T) where T: Iterator { for i in range(0, size) { assert_eq!(iter.size_hint(), (size - i, Some(size - i))); assert_eq!(iter.next().unwrap(), (i, i)); @@ -1510,7 +1658,9 @@ fn test_iter() { assert_eq!(iter.size_hint(), (0, Some(0))); assert_eq!(iter.next(), None); } - + test(size, map.iter().map(|(&k, &v)| (k, v))); + test(size, map.iter_mut().map(|(&k, &mut v)| (k, v))); + test(size, map.into_iter()); } #[test] @@ -1520,28 +1670,7 @@ fn test_iter_rev() { // Forwards let mut map: BTreeMap = range(0, size).map(|i| (i, i)).collect(); - { - let mut iter = map.iter().rev(); - for i in range(0, size) { - assert_eq!(iter.size_hint(), (size - i, Some(size - i))); - assert_eq!(iter.next().unwrap(), (&(size - i - 1), &(size - i - 1))); - } - assert_eq!(iter.size_hint(), (0, Some(0))); - assert_eq!(iter.next(), None); - } - - { - let mut iter = map.iter_mut().rev(); - for i in range(0, size) { - assert_eq!(iter.size_hint(), (size - i, Some(size - i))); - assert_eq!(iter.next().unwrap(), (&(size - i - 1), &mut(size - i - 1))); - } - assert_eq!(iter.size_hint(), (0, Some(0))); - assert_eq!(iter.next(), None); - } - - { - let mut iter = map.into_iter().rev(); + fn test(size: uint, mut iter: T) where T: Iterator { for i in range(0, size) { assert_eq!(iter.size_hint(), (size - i, Some(size - i))); assert_eq!(iter.next().unwrap(), (size - i - 1, size - i - 1)); @@ -1549,7 +1678,93 @@ fn test_iter_rev() { assert_eq!(iter.size_hint(), (0, Some(0))); assert_eq!(iter.next(), None); } + test(size, map.iter().rev().map(|(&k, &v)| (k, v))); + test(size, map.iter_mut().rev().map(|(&k, &mut v)| (k, v))); + test(size, map.into_iter().rev()); + } + #[test] + fn test_iter_mixed() { + let size = 10000u; + + // Forwards + let mut map: BTreeMap = range(0, size).map(|i| (i, i)).collect(); + + fn test(size: uint, mut iter: T) + where T: Iterator + DoubleEndedIterator { + for i in range(0, size / 4) { + assert_eq!(iter.size_hint(), (size - i * 2, Some(size - i * 2))); + assert_eq!(iter.next().unwrap(), (i, i)); + assert_eq!(iter.next_back().unwrap(), (size - i - 1, size - i - 1)); + } + for i in range(size / 4, size * 3 / 4) { + assert_eq!(iter.size_hint(), (size * 3 / 4 - i, Some(size * 3 / 4 - i))); + assert_eq!(iter.next().unwrap(), (i, i)); + } + assert_eq!(iter.size_hint(), (0, Some(0))); + assert_eq!(iter.next(), None); + } + test(size, map.iter().map(|(&k, &v)| (k, v))); + test(size, map.iter_mut().map(|(&k, &mut v)| (k, v))); + test(size, map.into_iter()); + } + + #[test] + fn test_range_small() { + let size = 5u; + + // Forwards + let map: BTreeMap = range(0, size).map(|i| (i, i)).collect(); + + let mut j = 0u; + for ((&k, &v), i) in map.range(Included(&2), Unbounded).zip(range(2u, size)) { + assert_eq!(k, i); + assert_eq!(v, i); + j += 1; + } + assert_eq!(j, size - 2); + } + + #[test] + fn test_range_1000() { + let size = 1000u; + let map: BTreeMap = range(0, size).map(|i| (i, i)).collect(); + + fn test(map: &BTreeMap, size: uint, min: Bound<&uint>, max: Bound<&uint>) { + let mut kvs = map.range(min, max).map(|(&k, &v)| (k, v)); + let mut pairs = range(0, size).map(|i| (i, i)); + + for (kv, pair) in kvs.by_ref().zip(pairs.by_ref()) { + assert_eq!(kv, pair); + } + assert_eq!(kvs.next(), None); + assert_eq!(pairs.next(), None); + } + test(&map, size, Included(&0), Excluded(&size)); + test(&map, size, Unbounded, Excluded(&size)); + test(&map, size, Included(&0), Included(&(size - 1))); + test(&map, size, Unbounded, Included(&(size - 1))); + test(&map, size, Included(&0), Unbounded); + test(&map, size, Unbounded, Unbounded); + } + + #[test] + fn test_range() { + let size = 200u; + let map: BTreeMap = range(0, size).map(|i| (i, i)).collect(); + + for i in range(0, size) { + for j in range(i, size) { + let mut kvs = map.range(Included(&i), Included(&j)).map(|(&k, &v)| (k, v)); + let mut pairs = range_inclusive(i, j).map(|i| (i, i)); + + for (kv, pair) in kvs.by_ref().zip(pairs.by_ref()) { + assert_eq!(kv, pair); + } + assert_eq!(kvs.next(), None); + assert_eq!(pairs.next(), None); + } + } } #[test] diff --git a/src/libcollections/btree/node.rs b/src/libcollections/btree/node.rs index 7d5422290e2..afce5f7dbda 100644 --- a/src/libcollections/btree/node.rs +++ b/src/libcollections/btree/node.rs @@ -77,6 +77,24 @@ pub struct Node { _capacity: uint, } +struct NodeSlice<'a, K: 'a, V: 'a> { + keys: &'a [K], + vals: &'a [V], + pub edges: &'a [Node], + head_is_edge: bool, + tail_is_edge: bool, + has_edges: bool, +} + +struct MutNodeSlice<'a, K: 'a, V: 'a> { + keys: &'a [K], + vals: &'a mut [V], + pub edges: &'a mut [Node], + head_is_edge: bool, + tail_is_edge: bool, + has_edges: bool, +} + /// Rounds up to a multiple of a power of two. Returns the closest multiple /// of `target_alignment` that is higher or equal to `unrounded`. /// @@ -342,7 +360,8 @@ pub fn as_slices_mut<'a>(&'a mut self) -> (&'a mut [K], &'a mut [V]) { } #[inline] - pub fn as_slices_internal<'a>(&'a self) -> (&'a [K], &'a [V], &'a [Node]) { + pub fn as_slices_internal<'b>(&'b self) -> NodeSlice<'b, K, V> { + let is_leaf = self.is_leaf(); let (keys, vals) = self.as_slices(); let edges: &[_] = if self.is_leaf() { &[] @@ -354,12 +373,18 @@ pub fn as_slices_internal<'a>(&'a self) -> (&'a [K], &'a [V], &'a [Node]) }) } }; - (keys, vals, edges) + NodeSlice { + keys: keys, + vals: vals, + edges: edges, + head_is_edge: true, + tail_is_edge: true, + has_edges: !is_leaf, + } } #[inline] - pub fn as_slices_internal_mut<'a>(&'a mut self) -> (&'a mut [K], &'a mut [V], - &'a mut [Node]) { + pub fn as_slices_internal_mut<'b>(&'b mut self) -> MutNodeSlice<'b, K, V> { unsafe { mem::transmute(self.as_slices_internal()) } } @@ -385,12 +410,12 @@ pub fn vals_mut<'a>(&'a mut self) -> &'a mut [V] { #[inline] pub fn edges<'a>(&'a self) -> &'a [Node] { - self.as_slices_internal().2 + self.as_slices_internal().edges } #[inline] pub fn edges_mut<'a>(&'a mut self) -> &'a mut [Node] { - self.as_slices_internal_mut().2 + self.as_slices_internal_mut().edges } } @@ -522,30 +547,11 @@ pub fn search>>(node: NodeRef, key: // FIXME(Gankro): Tune when to search linear or binary based on B (and maybe K/V). // For the B configured as of this writing (B = 6), binary search was *significantly* // worse for uints. - let (found, index) = node.search_linear(key); - if found { - Found(Handle { - node: node, - index: index - }) - } else { - GoDown(Handle { - node: node, - index: index - }) + match node.as_slices_internal().search_linear(key) { + (index, true) => Found(Handle { node: node, index: index }), + (index, false) => GoDown(Handle { node: node, index: index }), } } - - fn search_linear(&self, key: &Q) -> (bool, uint) where Q: BorrowFrom + Ord { - for (i, k) in self.keys().iter().enumerate() { - match key.cmp(BorrowFrom::borrow_from(k)) { - Greater => {}, - Equal => return (true, i), - Less => return (false, i), - } - } - (false, self.len()) - } } // Public interface @@ -1043,31 +1049,11 @@ pub fn kv_handle(&mut self, index: uint) -> Handle<&mut Node, handle::KV, } pub fn iter<'a>(&'a self) -> Traversal<'a, K, V> { - let is_leaf = self.is_leaf(); - let (keys, vals, edges) = self.as_slices_internal(); - Traversal { - inner: ElemsAndEdges( - keys.iter().zip(vals.iter()), - edges.iter() - ), - head_is_edge: true, - tail_is_edge: true, - has_edges: !is_leaf, - } + self.as_slices_internal().iter() } pub fn iter_mut<'a>(&'a mut self) -> MutTraversal<'a, K, V> { - let is_leaf = self.is_leaf(); - let (keys, vals, edges) = self.as_slices_internal_mut(); - MutTraversal { - inner: ElemsAndEdges( - keys.iter().zip(vals.iter_mut()), - edges.iter_mut() - ), - head_is_edge: true, - tail_is_edge: true, - has_edges: !is_leaf, - } + self.as_slices_internal_mut().iter_mut() } pub fn into_iter(self) -> MoveTraversal { @@ -1311,12 +1297,15 @@ fn min_load_from_capacity(cap: uint) -> uint { /// A trait for pairs of `Iterator`s, one over edges and the other over key/value pairs. This is /// necessary, as the `MoveTraversalImpl` needs to have a destructor that deallocates the `Node`, /// and a pair of `Iterator`s would require two independent destructors. -trait TraversalImpl { - fn next_kv(&mut self) -> Option<(K, V)>; - fn next_kv_back(&mut self) -> Option<(K, V)>; +trait TraversalImpl { + type Item; + type Edge; - fn next_edge(&mut self) -> Option; - fn next_edge_back(&mut self) -> Option; + fn next_kv(&mut self) -> Option; + fn next_kv_back(&mut self) -> Option; + + fn next_edge(&mut self) -> Option; + fn next_edge_back(&mut self) -> Option; } /// A `TraversalImpl` that actually is backed by two iterators. This works in the non-moving case, @@ -1324,9 +1313,11 @@ trait TraversalImpl { struct ElemsAndEdges(Elems, Edges); impl - TraversalImpl for ElemsAndEdges + TraversalImpl for ElemsAndEdges where Elems : Iterator, Edges : Iterator { + type Item = (K, V); + type Edge = E; fn next_kv(&mut self) -> Option<(K, V)> { self.0.next() } fn next_kv_back(&mut self) -> Option<(K, V)> { self.0.next_back() } @@ -1347,7 +1338,10 @@ struct MoveTraversalImpl { is_leaf: bool } -impl TraversalImpl> for MoveTraversalImpl { +impl TraversalImpl for MoveTraversalImpl { + type Item = (K, V); + type Edge = Node; + fn next_kv(&mut self) -> Option<(K, V)> { match (self.keys.next(), self.vals.next()) { (Some(k), Some(v)) => Some((k, v)), @@ -1398,9 +1392,12 @@ struct AbsTraversal { has_edges: bool, } -/// A single atomic step in a traversal. Either an element is visited, or an edge is followed +/// A single atomic step in a traversal. pub enum TraversalItem { - Elem(K, V), + /// An element is visited. This isn't written as `Elem(K, V)` just because `opt.map(Elem)` + /// requires the function to take a single argument. (Enum constructors are functions.) + Elem((K, V)), + /// An edge is followed. Edge(E), } @@ -1417,32 +1414,175 @@ pub enum TraversalItem { /// An owning traversal over a node's entries and edges pub type MoveTraversal = AbsTraversal>; -#[old_impl_check] -impl> Iterator for AbsTraversal { + +impl Iterator for AbsTraversal + where Impl: TraversalImpl { type Item = TraversalItem; fn next(&mut self) -> Option> { - let head_is_edge = self.head_is_edge; - self.head_is_edge = !head_is_edge; - - if head_is_edge && self.has_edges { - self.inner.next_edge().map(|node| Edge(node)) - } else { - self.inner.next_kv().map(|(k, v)| Elem(k, v)) - } + self.next_edge_item().map(Edge).or_else(|| + self.next_kv_item().map(Elem) + ) } } -#[old_impl_check] -impl> DoubleEndedIterator for AbsTraversal { +impl DoubleEndedIterator for AbsTraversal + where Impl: TraversalImpl { fn next_back(&mut self) -> Option> { - let tail_is_edge = self.tail_is_edge; - self.tail_is_edge = !tail_is_edge; + self.next_edge_item_back().map(Edge).or_else(|| + self.next_kv_item_back().map(Elem) + ) + } +} - if tail_is_edge && self.has_edges { - self.inner.next_edge_back().map(|node| Edge(node)) +impl AbsTraversal + where Impl: TraversalImpl { + /// Advances the iterator and returns the item if it's an edge. Returns None + /// and does nothing if the first item is not an edge. + pub fn next_edge_item(&mut self) -> Option { + // NB. `&& self.has_edges` might be redundant in this condition. + let edge = if self.head_is_edge && self.has_edges { + self.inner.next_edge() } else { - self.inner.next_kv_back().map(|(k, v)| Elem(k, v)) + None + }; + self.head_is_edge = false; + edge + } + + /// Advances the iterator and returns the item if it's an edge. Returns None + /// and does nothing if the last item is not an edge. + pub fn next_edge_item_back(&mut self) -> Option { + let edge = if self.tail_is_edge && self.has_edges { + self.inner.next_edge_back() + } else { + None + }; + self.tail_is_edge = false; + edge + } + + /// Advances the iterator and returns the item if it's a key-value pair. Returns None + /// and does nothing if the first item is not a key-value pair. + pub fn next_kv_item(&mut self) -> Option<(K, V)> { + if !self.head_is_edge { + self.head_is_edge = true; + self.inner.next_kv() + } else { + None + } + } + + /// Advances the iterator and returns the item if it's a key-value pair. Returns None + /// and does nothing if the last item is not a key-value pair. + pub fn next_kv_item_back(&mut self) -> Option<(K, V)> { + if !self.tail_is_edge { + self.tail_is_edge = true; + self.inner.next_kv_back() + } else { + None } } } + +macro_rules! node_slice_impl { + ($NodeSlice:ident, $Traversal:ident, + $as_slices_internal:ident, $slice_from:ident, $slice_to:ident, $iter:ident) => { + impl<'a, K: Ord + 'a, V: 'a> $NodeSlice<'a, K, V> { + /// Performs linear search in a slice. Returns a tuple of (index, is_exact_match). + fn search_linear(&self, key: &Q) -> (uint, bool) + where Q: BorrowFrom + Ord { + for (i, k) in self.keys.iter().enumerate() { + match key.cmp(BorrowFrom::borrow_from(k)) { + Greater => {}, + Equal => return (i, true), + Less => return (i, false), + } + } + (self.keys.len(), false) + } + + /// Returns a sub-slice with elements starting with `min_key`. + pub fn slice_from(self, min_key: &K) -> $NodeSlice<'a, K, V> { + // _______________ + // |_1_|_3_|_5_|_7_| + // | | | | | + // 0 0 1 1 2 2 3 3 4 index + // | | | | | + // \___|___|___|___/ slice_from(&0); pos = 0 + // \___|___|___/ slice_from(&2); pos = 1 + // |___|___|___/ slice_from(&3); pos = 1; result.head_is_edge = false + // \___|___/ slice_from(&4); pos = 2 + // \___/ slice_from(&6); pos = 3 + // \|/ slice_from(&999); pos = 4 + let (pos, pos_is_kv) = self.search_linear(min_key); + $NodeSlice { + has_edges: self.has_edges, + edges: if !self.has_edges { + self.edges + } else { + self.edges.$slice_from(pos) + }, + keys: self.keys.slice_from(pos), + vals: self.vals.$slice_from(pos), + head_is_edge: !pos_is_kv, + tail_is_edge: self.tail_is_edge, + } + } + + /// Returns a sub-slice with elements up to and including `max_key`. + pub fn slice_to(self, max_key: &K) -> $NodeSlice<'a, K, V> { + // _______________ + // |_1_|_3_|_5_|_7_| + // | | | | | + // 0 0 1 1 2 2 3 3 4 index + // | | | | | + //\|/ | | | | slice_to(&0); pos = 0 + // \___/ | | | slice_to(&2); pos = 1 + // \___|___| | | slice_to(&3); pos = 1; result.tail_is_edge = false + // \___|___/ | | slice_to(&4); pos = 2 + // \___|___|___/ | slice_to(&6); pos = 3 + // \___|___|___|___/ slice_to(&999); pos = 4 + let (pos, pos_is_kv) = self.search_linear(max_key); + let pos = pos + if pos_is_kv { 1 } else { 0 }; + $NodeSlice { + has_edges: self.has_edges, + edges: if !self.has_edges { + self.edges + } else { + self.edges.$slice_to(pos + 1) + }, + keys: self.keys.slice_to(pos), + vals: self.vals.$slice_to(pos), + head_is_edge: self.head_is_edge, + tail_is_edge: !pos_is_kv, + } + } + } + + impl<'a, K: 'a, V: 'a> $NodeSlice<'a, K, V> { + /// Returns an iterator over key/value pairs and edges in a slice. + #[inline] + pub fn $iter(self) -> $Traversal<'a, K, V> { + let mut edges = self.edges.$iter(); + // Skip edges at both ends, if excluded. + if !self.head_is_edge { edges.next(); } + if !self.tail_is_edge { edges.next_back(); } + // The key iterator is always immutable. + $Traversal { + inner: ElemsAndEdges( + self.keys.iter().zip(self.vals.$iter()), + edges + ), + head_is_edge: self.head_is_edge, + tail_is_edge: self.tail_is_edge, + has_edges: self.has_edges, + } + } + } + } +} + +node_slice_impl!(NodeSlice, Traversal, as_slices_internal, slice_from, slice_to, iter); +node_slice_impl!(MutNodeSlice, MutTraversal, as_slices_internal_mut, slice_from_mut, + slice_to_mut, iter_mut); diff --git a/src/libcollections/btree/set.rs b/src/libcollections/btree/set.rs index 6e048e0e83c..4d71f9dbea8 100644 --- a/src/libcollections/btree/set.rs +++ b/src/libcollections/btree/set.rs @@ -25,6 +25,7 @@ use core::ops::{BitOr, BitAnd, BitXor, Sub}; use btree_map::{BTreeMap, Keys}; +use Bound; // FIXME(conventions): implement bounded iterators @@ -50,6 +51,11 @@ pub struct IntoIter { iter: Map<(T, ()), T, ::btree_map::IntoIter, fn((T, ())) -> T> } +/// An iterator over a sub-range of BTreeSet's items. +pub struct Range<'a, T: 'a> { + iter: Map<(&'a T, &'a ()), &'a T, ::btree_map::Range<'a, T, ()>, fn((&'a T, &'a ())) -> &'a T> +} + /// A lazy iterator producing elements in the set difference (in-order). #[stable] pub struct Difference<'a, T:'a> { @@ -145,6 +151,36 @@ fn first((a, _): (A, B)) -> A { a } } } +impl BTreeSet { + /// Constructs a double-ended iterator over a sub-range of elements in the set, starting + /// at min, and ending at max. If min is `Unbounded`, then it will be treated as "negative + /// infinity", and if max is `Unbounded`, then it will be treated as "positive infinity". + /// Thus range(Unbounded, Unbounded) will yield the whole collection. + /// + /// # Examples + /// + /// ``` + /// use std::collections::BTreeSet; + /// use std::collections::Bound::{Included, Unbounded}; + /// + /// let mut set = BTreeSet::new(); + /// set.insert(3u); + /// set.insert(5u); + /// set.insert(8u); + /// for &elem in set.range(Included(&4), Included(&8)) { + /// println!("{}", elem); + /// } + /// assert_eq!(Some(&5u), set.range(Included(&4), Unbounded).next()); + /// ``` + #[unstable = "matches collection reform specification, waiting for dust to settle"] + pub fn range<'a>(&'a self, min: Bound<&T>, max: Bound<&T>) -> Range<'a, T> { + fn first((a, _): (A, B)) -> A { a } + let first: fn((&'a T, &'a ())) -> &'a T = first; // coerce to fn pointer + + Range { iter: self.map.range(min, max).map(first) } + } +} + impl BTreeSet { /// Visits the values representing the difference, in ascending order. /// @@ -598,6 +634,16 @@ fn next_back(&mut self) -> Option { self.iter.next_back() } #[stable] impl ExactSizeIterator for IntoIter {} + +impl<'a, T> Iterator for Range<'a, T> { + type Item = &'a T; + + fn next(&mut self) -> Option<&'a T> { self.iter.next() } +} +impl<'a, T> DoubleEndedIterator for Range<'a, T> { + fn next_back(&mut self) -> Option<&'a T> { self.iter.next_back() } +} + /// Compare `x` and `y`, but return `short` if x is None and `long` if y is None fn cmp_opt(x: Option<&T>, y: Option<&T>, short: Ordering, long: Ordering) -> Ordering { diff --git a/src/libcollections/lib.rs b/src/libcollections/lib.rs index 797042c3216..10ea0421eee 100644 --- a/src/libcollections/lib.rs +++ b/src/libcollections/lib.rs @@ -26,7 +26,6 @@ #![feature(unsafe_destructor, slicing_syntax)] #![feature(box_syntax)] #![feature(unboxed_closures)] -#![feature(old_impl_check)] #![allow(unknown_features)] #![feature(int_uint)] #![allow(unstable)] #![no_std] @@ -142,3 +141,13 @@ mod prelude { pub use string::{String, ToString}; pub use vec::Vec; } + +/// An endpoint of a range of keys. +pub enum Bound { + /// An inclusive bound. + Included(T), + /// An exclusive bound. + Excluded(T), + /// An infinite endpoint. Indicates that there is no bound in this direction. + Unbounded, +} diff --git a/src/libstd/collections/mod.rs b/src/libstd/collections/mod.rs index 71ab89027ff..f085fd259de 100644 --- a/src/libstd/collections/mod.rs +++ b/src/libstd/collections/mod.rs @@ -311,6 +311,7 @@ #![stable] +pub use core_collections::Bound; pub use core_collections::{BinaryHeap, Bitv, BitvSet, BTreeMap, BTreeSet}; pub use core_collections::{DList, RingBuf, VecMap};