OpenE2K
/
rust
2
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

extra::treemap: share code between the mutable and immutable iterators. 

Yay for macros.
This commit is contained in:
Huon Wilson 2014-01-06 01:53:27 +11:00
parent e33fcf9da9
commit a9a348f2fb
1 changed files with 213 additions and 279 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

492
src/libextra/treemap.rs
View File

@ -169,8 +169,57 @@ impl<K: TotalOrd, V> TreeMap<K, V> {
}
/// Get a lazy iterator that consumes the treemap.
pub fn move_iter(self) -> TreeMapMoveIterator<K, V> {
let TreeMap { root: root, length: length } = self;
let stk = match root {
None => ~[],
Some(~tn) => ~[tn]
};
TreeMapMoveIterator {
stack: stk,
remaining: length
}
}
}
// range iterators.
macro_rules! bound_setup {
// initialiser of the iterator to manipulate
($iter:expr,
// whether we are looking for the lower or upper bound.
$is_lower_bound:expr) => {
{
let mut iter = $iter;
loop {
if !iter.node.is_null() {
let node_k = unsafe {&(*iter.node).key};
match k.cmp(node_k) {
Less => iter.traverse_left(),
Greater => iter.traverse_right(),
Equal => {
if $is_lower_bound {
iter.traverse_complete();
return iter;
} else {
iter.traverse_right()
}
}
}
} else {
iter.traverse_complete();
return iter;
}
}
}
}
}
impl<K: TotalOrd, V> TreeMap<K, V> {
/// Get a lazy iterator that should be initialized using
/// `iter_traverse_left`/`iter_traverse_right`/`iter_traverse_complete`.
/// `traverse_left`/`traverse_right`/`traverse_complete`.
fn iter_for_traversal<'a>(&'a self) -> TreeMapIterator<'a, K, V> {
TreeMapIterator {
stack: ~[],
@ -183,49 +232,17 @@ impl<K: TotalOrd, V> TreeMap<K, V> {
/// Return a lazy iterator to the first key-value pair whose key is not less than `k`
/// If all keys in map are less than `k` an empty iterator is returned.
pub fn lower_bound<'a>(&'a self, k: &K) -> TreeMapIterator<'a, K, V> {
let mut iter: TreeMapIterator<'a, K, V> = self.iter_for_traversal();
loop {
match iter.node {
Some(r) => {
match k.cmp(&r.key) {
Less => iter_traverse_left(&mut iter),
Greater => iter_traverse_right(&mut iter),
Equal => {
iter_traverse_complete(&mut iter);
return iter;
}
}
}
None => {
iter_traverse_complete(&mut iter);
return iter;
}
}
}
bound_setup!(self.iter_for_traversal(), true)
}
/// Return a lazy iterator to the first key-value pair whose key is greater than `k`
/// If all keys in map are not greater than `k` an empty iterator is returned.
pub fn upper_bound<'a>(&'a self, k: &K) -> TreeMapIterator<'a, K, V> {
let mut iter: TreeMapIterator<'a, K, V> = self.iter_for_traversal();
loop {
match iter.node {
Some(r) => {
match k.cmp(&r.key) {
Less => iter_traverse_left(&mut iter),
Greater => iter_traverse_right(&mut iter),
Equal => iter_traverse_right(&mut iter)
}
}
None => {
iter_traverse_complete(&mut iter);
return iter;
}
}
}
bound_setup!(self.iter_for_traversal(), false)
}
/// Get a lazy iterator that should be initialized using
/// `mut_iter_traverse_left`/`mut_iter_traverse_right`/`mut_iter_traverse_complete`.
/// `traverse_left`/`traverse_right`/`traverse_complete`.
fn mut_iter_for_traversal<'a>(&'a mut self) -> TreeMapMutIterator<'a, K, V> {
TreeMapMutIterator {
stack: ~[],
@ -241,23 +258,7 @@ impl<K: TotalOrd, V> TreeMap<K, V> {
/// If all keys in map are less than `k` an empty iterator is
/// returned.
pub fn mut_lower_bound<'a>(&'a mut self, k: &K) -> TreeMapMutIterator<'a, K, V> {
let mut iter = self.mut_iter_for_traversal();
loop {
if !iter.node.is_null() {
let node_k = unsafe {&(*iter.node).key};
match k.cmp(node_k) {
Less => mut_iter_traverse_left(&mut iter),
Greater => mut_iter_traverse_right(&mut iter),
Equal => {
mut_iter_traverse_complete(&mut iter);
return iter;
}
}
} else {
mut_iter_traverse_complete(&mut iter);
return iter;
}
}
bound_setup!(self.mut_iter_for_traversal(), true)
}
/// Return a lazy iterator to the first key-value pair (with the
@ -266,160 +267,26 @@ impl<K: TotalOrd, V> TreeMap<K, V> {
/// If all keys in map are not greater than `k` an empty iterator
/// is returned.
pub fn mut_upper_bound<'a>(&'a mut self, k: &K) -> TreeMapMutIterator<'a, K, V> {
let mut iter = self.mut_iter_for_traversal();
loop {
if !iter.node.is_null() {
let node_k = unsafe {&(*iter.node).key};
match k.cmp(node_k) {
Less => mut_iter_traverse_left(&mut iter),
Greater => mut_iter_traverse_right(&mut iter),
Equal => mut_iter_traverse_right(&mut iter)
}
} else {
mut_iter_traverse_complete(&mut iter);
return iter;
}
}
}
/// Get a lazy iterator that consumes the treemap.
pub fn move_iter(self) -> TreeMapMoveIterator<K, V> {
let TreeMap { root: root, length: length } = self;
let stk = match root {
None => ~[],
Some(~tn) => ~[tn]
};
TreeMapMoveIterator {
stack: stk,
remaining: length
}
bound_setup!(self.mut_iter_for_traversal(), false)
}
}
/// Lazy forward iterator over a map
pub struct TreeMapIterator<'a, K, V> {
priv stack: ~[&'a TreeNode<K, V>],
priv node: Option<&'a TreeNode<K, V>>,
// See the comment on TreeMapMutIterator; this is just to allow
// code-sharing (for this immutable-values iterator it *could* very
// well be Option<&'a TreeNode<K,V>>).
priv node: *TreeNode<K, V>,
priv remaining_min: uint,
priv remaining_max: uint
}
fn deref<'a, K, V>(node: &'a Option<~TreeNode<K, V>>) -> Option<&'a TreeNode<K, V>> {
match *node {
Some(ref n) => {
let n: &TreeNode<K, V> = *n;
Some(n)
}
None => None
}
}
impl<'a, K, V> TreeMapIterator<'a, K, V> {
#[inline(always)]
fn next_(&mut self, forward: bool) -> Option<(&'a K, &'a V)> {
while !self.stack.is_empty() || self.node.is_some() {
match self.node {
Some(x) => {
self.stack.push(x);
self.node = deref(if forward { &x.left } else { &x.right });
}
None => {
let res = self.stack.pop();
self.node = deref(if forward { &res.right } else { &res.left });
self.remaining_max -= 1;
if self.remaining_min > 0 {
self.remaining_min -= 1;
}
return Some((&res.key, &res.value));
}
}
}
None
}
}
impl<'a, K, V> Iterator<(&'a K, &'a V)> for TreeMapIterator<'a, K, V> {
/// Advance the iterator to the next node (in order) and return a
/// tuple with a reference to the key and value. If there are no
/// more nodes, return `None`.
fn next(&mut self) -> Option<(&'a K, &'a V)> {
self.next_(true)
}
#[inline]
fn size_hint(&self) -> (uint, Option<uint>) {
(self.remaining_min, Some(self.remaining_max))
}
}
/// Lazy backward iterator over a map
pub struct TreeMapRevIterator<'a, K, V> {
priv iter: TreeMapIterator<'a, K, V>,
}
impl<'a, K, V> Iterator<(&'a K, &'a V)> for TreeMapRevIterator<'a, K, V> {
/// Advance the iterator to the next node (in order) and return a
/// tuple with a reference to the key and value. If there are no
/// more nodes, return `None`.
fn next(&mut self) -> Option<(&'a K, &'a V)> {
self.iter.next_(false)
}
#[inline]
fn size_hint(&self) -> (uint, Option<uint>) {
self.iter.size_hint()
}
}
/// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to
/// initialize TreeMapIterator pointing to element inside tree structure.
///
/// They should be used in following manner:
/// - create iterator using TreeMap::iter_for_traversal
/// - find required node using `iter_traverse_left`/`iter_traverse_right`
/// (current node is `TreeMapIterator::node` field)
/// - complete initialization with `iter_traverse_complete`
#[inline]
fn iter_traverse_left<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) {
let node = it.node.unwrap();
it.stack.push(node);
it.node = deref(&node.left);
}
#[inline]
fn iter_traverse_right<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) {
it.node = deref(&it.node.get_ref().right);
}
/// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to
/// initialize TreeMapIterator pointing to element inside tree structure.
///
/// Completes traversal. Should be called before using iterator.
/// Iteration will start from `self.node`.
/// If `self.node` is None iteration will start from last node from which we
/// traversed left.
#[inline]
fn iter_traverse_complete<'a, K, V>(it: &mut TreeMapIterator<'a, K, V>) {
match it.node {
Some(n) => {
it.stack.push(n);
it.node = None;
}
None => ()
}
}
fn mut_deref<K, V>(x: &mut Option<~TreeNode<K, V>>) -> *mut TreeNode<K, V> {
match *x {
Some(ref mut n) => {
let n: &mut TreeNode<K, V> = *n;
n as *mut TreeNode<K, V>
}
None => ptr::mut_null()
}
}
/// Lazy forward iterator over a map that allows for the mutation of
/// the values.
pub struct TreeMapMutIterator<'a, K, V> {
@ -448,102 +315,169 @@ pub struct TreeMapMutIterator<'a, K, V> {
priv remaining_max: uint
}
impl<'a, K, V> TreeMapMutIterator<'a, K, V> {
#[inline(always)]
fn next_(&mut self, forward: bool) -> Option<(&'a K, &'a mut V)> {
while !self.stack.is_empty() || !self.node.is_null() {
let node = self.node;
if !node.is_null() {
let node = unsafe {&mut *node};
{
let next_node = if forward { &mut node.left } else { &mut node.right };
self.node = mut_deref(next_node);
}
self.stack.push(node);
} else {
let node = self.stack.pop();
self.node = mut_deref(if forward { &mut node.right } else { &mut node.left });
self.remaining_max -= 1;
if self.remaining_min > 0 {
self.remaining_min -= 1;
}
return Some((&node.key, &mut node.value));
}
}
None
}
}
impl<'a, K, V> Iterator<(&'a K, &'a mut V)> for TreeMapMutIterator<'a, K, V> {
/// Advance the iterator to the next node (in order) and return a
/// tuple with a reference to the key and value. If there are no
/// more nodes, return `None`.
fn next(&mut self) -> Option<(&'a K, &'a mut V)> {
self.next_(true)
}
#[inline]
fn size_hint(&self) -> (uint, Option<uint>) {
(self.remaining_min, Some(self.remaining_max))
}
}
/// Lazy backward iterator over a map
pub struct TreeMapMutRevIterator<'a, K, V> {
priv iter: TreeMapMutIterator<'a, K, V>,
}
impl<'a, K, V> Iterator<(&'a K, &'a mut V)> for TreeMapMutRevIterator<'a, K, V> {
/// Advance the iterator to the next node (in order) and return a
/// tuple with a reference to the key and value. If there are no
/// more nodes, return `None`.
fn next(&mut self) -> Option<(&'a K, &'a mut V)> {
self.iter.next_(false)
}
#[inline]
fn size_hint(&self) -> (uint, Option<uint>) {
self.iter.size_hint()
}
}
// FIXME #5846 we want to be able to choose between &x and &mut x
// (with many different `x`) below, so we need to optionally pass mut
// as a tt, but the only thing we can do with a `tt` is pass them to
// other macros, so this takes the `& <mutability> <operand>` token
// sequence and forces their evalutation as an expression.
macro_rules! addr { ($e:expr) => { $e }}
/// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to
/// initialize TreeMapMutIterator pointing to element inside tree structure.
///
/// They should be used in following manner:
/// - create iterator using TreeMap::iter_for_traversal
/// - find required node using `iter_traverse_left`/`iter_traverse_right`
/// (current node is `TreeMapMutIterator::node` field)
/// - complete initialization with `iter_traverse_complete`
#[inline]
fn mut_iter_traverse_left<'a, K, V>(it: &mut TreeMapMutIterator<'a, K, V>) {
// guaranteed to be non-null
let node = unsafe {&mut *it.node};
it.node = mut_deref(&mut node.left);
it.stack.push(node);
}
macro_rules! define_iterator {
($name:ident,
$rev_name:ident,
// the type of the values of the treemap in the return value of
// the iterator (i.e. &V or &mut V). This is non-hygienic in the
// name of the lifetime.
value_type = $value_type:ty,
#[inline]
fn mut_iter_traverse_right<'a, K, V>(it: &mut TreeMapMutIterator<'a, K, V>) {
// guaranteed to be non-null
let node = unsafe {&mut *it.node};
it.node = mut_deref(&mut node.right);
}
// the function to go from &m Option<~TreeNode> to *m TreeNode
deref = $deref:ident,
/// iter_traverse_left, iter_traverse_right and iter_traverse_complete are used to
/// initialize TreeMapMutIterator pointing to element inside tree structure.
///
/// Completes traversal. Should be called before using iterator.
/// Iteration will start from `self.node`.
/// If `self.node` is None iteration will start from last node from which we
/// traversed left.
#[inline]
fn mut_iter_traverse_complete<'a, K, V>(it: &mut TreeMapMutIterator<'a, K, V>) {
if !it.node.is_null() {
unsafe {
it.stack.push(&mut *it.node);
// see comment on `addr!`, this is just an optional `mut`, but
// there's no support for 0-or-1 repeats.
addr_mut = $($addr_mut:tt)*
) => {
// private methods on the forward iterator
impl<'a, K, V> $name<'a, K, V> {
#[inline(always)]
fn next_(&mut self, forward: bool) -> Option<(&'a K, $value_type)> {
while !self.stack.is_empty() || !self.node.is_null() {
if !self.node.is_null() {
let node = unsafe {addr!(& $($addr_mut)* *self.node)};
{
let next_node = if forward {
addr!(& $($addr_mut)* node.left)
} else {
addr!(& $($addr_mut)* node.right)
};
self.node = $deref(next_node);
}
self.stack.push(node);
} else {
let node = self.stack.pop();
let next_node = if forward {
addr!(& $($addr_mut)* node.right)
} else {
addr!(& $($addr_mut)* node.left)
};
self.node = $deref(next_node);
self.remaining_max -= 1;
if self.remaining_min > 0 {
self.remaining_min -= 1;
}
return Some((&node.key, addr!(& $($addr_mut)* node.value)));
}
}
None
}
/// traverse_left, traverse_right and traverse_complete are
/// used to initialize TreeMapIterator/TreeMapMutIterator
/// pointing to element inside tree structure.
///
/// They should be used in following manner:
/// - create iterator using TreeMap::[mut_]iter_for_traversal
/// - find required node using `traverse_left`/`traverse_right`
/// (current node is `TreeMapIterator::node` field)
/// - complete initialization with `traverse_complete`
///
/// After this, iteration will start from `self.node`. If
/// `self.node` is None iteration will start from last
/// node from which we traversed left.
#[inline]
fn traverse_left(&mut self) {
let node = unsafe {addr!(& $($addr_mut)* *self.node)};
self.node = $deref(addr!(& $($addr_mut)* node.left));
self.stack.push(node);
}
#[inline]
fn traverse_right(&mut self) {
let node = unsafe {addr!(& $($addr_mut)* *self.node)};
self.node = $deref(addr!(& $($addr_mut)* node.right));
}
#[inline]
fn traverse_complete(&mut self) {
if !self.node.is_null() {
unsafe {
self.stack.push(addr!(& $($addr_mut)* *self.node));
}
self.node = ptr::RawPtr::null();
}
}
}
it.node = ptr::mut_null();
// the forward Iterator impl.
impl<'a, K, V> Iterator<(&'a K, $value_type)> for $name<'a, K, V> {
/// Advance the iterator to the next node (in order) and return a
/// tuple with a reference to the key and value. If there are no
/// more nodes, return `None`.
fn next(&mut self) -> Option<(&'a K, $value_type)> {
self.next_(true)
}
#[inline]
fn size_hint(&self) -> (uint, Option<uint>) {
(self.remaining_min, Some(self.remaining_max))
}
}
// the reverse Iterator impl.
impl<'a, K, V> Iterator<(&'a K, $value_type)> for $rev_name<'a, K, V> {
fn next(&mut self) -> Option<(&'a K, $value_type)> {
self.iter.next_(false)
}
#[inline]
fn size_hint(&self) -> (uint, Option<uint>) {
self.iter.size_hint()
}
}
}
} // end of define_iterator
define_iterator! {
TreeMapIterator,
TreeMapRevIterator,
value_type = &'a V,
deref = deref,
// immutable, so no mut
addr_mut =
}
define_iterator! {
TreeMapMutIterator,
TreeMapMutRevIterator,
value_type = &'a mut V,
deref = mut_deref,
addr_mut = mut
}
fn deref<'a, K, V>(node: &'a Option<~TreeNode<K, V>>) -> *TreeNode<K, V> {
match *node {
Some(ref n) => {
let n: &TreeNode<K, V> = *n;
n as *TreeNode<K, V>
}
None => ptr::null()
}
}
fn mut_deref<K, V>(x: &mut Option<~TreeNode<K, V>>) -> *mut TreeNode<K, V> {
match *x {
Some(ref mut n) => {
let n: &mut TreeNode<K, V> = *n;
n as *mut TreeNode<K, V>
}
None => ptr::mut_null()
}
}