Skip to content

Commit fe03cae

Browse files
huonwhuonw
committed
std::trie: use macros to share code between the iterator implementations.
1 parent f07c74d commit fe03cae

File tree

1 file changed

+126
-129
lines changed

1 file changed

+126
-129
lines changed

‎src/libstd/trie.rs‎

Lines changed: 126 additions & 129 deletions
Original file line numberDiff line numberDiff line change
@@ -141,40 +141,95 @@ impl<T> TrieMap<T> {
141141
remaining_max: self.length
142142
}
143143
}
144+
}
144145

146+
// FIXME #5846 we want to be able to choose between &x and &mut x
147+
// (with many different `x`) below, so we need to optionally pass mut
148+
// as a tt, but the only thing we can do with a `tt` is pass them to
149+
// other macros, so this takes the `& <mutability> <operand>` token
150+
// sequence and forces their evalutation as an expression. (see also
151+
// `item!` below.)
152+
macro_rules! addr { ($e:expr) => { $e } }
153+
154+
macro_rules! bound {
155+
($iterator_name:ident,
156+
// the current treemap
157+
self = $this:expr,
158+
// the key to look for
159+
key = $key:expr,
160+
// are we looking at the upper bound?
161+
is_upper = $upper:expr,
162+
163+
// method names for slicing/iterating.
164+
slice_from = $slice_from:ident,
165+
iter = $iter:ident,
166+
167+
// see the comment on `addr!`, this is just an optional mut, but
168+
// there's no 0-or-1 repeats yet.
169+
mutability = $($mut_:tt)*) => {
170+
{
171+
// # For `mut`
172+
// We need an unsafe pointer here because we are borrowing
173+
// mutable references to the internals of each of these
174+
// mutable nodes, while still using the outer node.
175+
//
176+
// However, we're allowed to flaunt rustc like this because we
177+
// never actually modify the "shape" of the nodes. The only
178+
// place that mutation is can actually occur is of the actual
179+
// values of the TrieMap (as the return value of the
180+
// iterator), i.e. we can never cause a deallocation of any
181+
// TrieNodes so the raw pointer is always valid.
182+
//
183+
// # For non-`mut`
184+
// We like sharing code so much that even a little unsafe won't
185+
// stop us.
186+
let this = $this;
187+
let mut node = addr!(& $($mut_)* this.root as * $($mut_)* TrieNode<T>);
188+
189+
let key = $key;
190+
191+
let mut idx = 0;
192+
let mut it = $iterator_name {
193+
stack: ~[],
194+
remaining_min: 0,
195+
remaining_max: this.length
196+
};
197+
// this addr is necessary for the `Internal` pattern.
198+
addr!(loop {
199+
let children = unsafe {addr!(& $($mut_)* (*node).children)};
200+
let child_id = chunk(key, idx);
201+
match children[child_id] {
202+
Internal(ref $($mut_)* n) => {
203+
node = addr!(& $($mut_)* **n as * $($mut_)* TrieNode<T>);
204+
}
205+
External(stored, _) => {
206+
if stored < key || ($upper && stored == key) {
207+
it.stack.push(children.$slice_from(child_id + 1).$iter());
208+
} else {
209+
it.stack.push(children.$slice_from(child_id).$iter());
210+
}
211+
return it;
212+
}
213+
Nothing => {
214+
it.stack.push(children.$slice_from(child_id + 1).$iter());
215+
return it
216+
}
217+
}
218+
it.stack.push(children.$slice_from(child_id + 1).$iter());
219+
idx += 1;
220+
})
221+
}
222+
}
223+
}
224+
225+
impl<T> TrieMap<T> {
145226
// If `upper` is true then returns upper_bound else returns lower_bound.
146227
#[inline]
147228
fn bound<'a>(&'a self, key: uint, upper: bool) -> TrieMapIterator<'a, T> {
148-
let mut node: &'a TrieNode<T> = &self.root;
149-
let mut idx = 0;
150-
let mut it = TrieMapIterator {
151-
stack: ~[],
152-
remaining_min: 0,
153-
remaining_max: self.length
154-
};
155-
loop {
156-
let children = &node.children;
157-
let child_id = chunk(key, idx);
158-
match children[child_id] {
159-
Internal(ref n) => {
160-
node = &**n;
161-
it.stack.push(children.slice_from(child_id + 1).iter());
162-
}
163-
External(stored, _) => {
164-
if stored < key || (upper && stored == key) {
165-
it.stack.push(children.slice_from(child_id + 1).iter());
166-
} else {
167-
it.stack.push(children.slice_from(child_id).iter());
168-
}
169-
return it;
170-
}
171-
Nothing => {
172-
it.stack.push(children.slice_from(child_id + 1).iter());
173-
return it
174-
}
175-
}
176-
idx += 1;
177-
}
229+
bound!(TrieMapIterator, self = self,
230+
key = key, is_upper = upper,
231+
slice_from = slice_from, iter = iter,
232+
mutability = )
178233
}
179234

180235
/// Get an iterator pointing to the first key-value pair whose key is not less than `key`.
@@ -191,47 +246,10 @@ impl<T> TrieMap<T> {
191246
// If `upper` is true then returns upper_bound else returns lower_bound.
192247
#[inline]
193248
fn mut_bound<'a>(&'a mut self, key: uint, upper: bool) -> TrieMapMutIterator<'a, T> {
194-
// we need an unsafe pointer here because we are borrowing
195-
// references to the internals of each of these
196-
// nodes.
197-
//
198-
// However, we're allowed to flaunt rustc like this because we
199-
// never actually modify the "shape" of the nodes. The only
200-
// place that mutation is can actually occur is of the actual
201-
// values of the TrieMap (as the return value of the
202-
// iterator), i.e. we can never cause a deallocation of any
203-
// TrieNodes so this pointer is always valid.
204-
let mut node = &mut self.root as *mut TrieNode<T>;
205-
206-
let mut idx = 0;
207-
let mut it = TrieMapMutIterator {
208-
stack: ~[],
209-
remaining_min: 0,
210-
remaining_max: self.length
211-
};
212-
loop {
213-
let children = unsafe {&mut (*node).children};
214-
let child_id = chunk(key, idx);
215-
match children[child_id] {
216-
Internal(ref mut n) => {
217-
node = &mut **n as *mut TrieNode<T>;
218-
}
219-
External(stored, _) => {
220-
if stored < key || (upper && stored == key) {
221-
it.stack.push(children.mut_slice_from(child_id + 1).mut_iter());
222-
} else {
223-
it.stack.push(children.mut_slice_from(child_id).mut_iter());
224-
}
225-
return it;
226-
}
227-
Nothing => {
228-
it.stack.push(children.mut_slice_from(child_id + 1).mut_iter());
229-
return it
230-
}
231-
}
232-
it.stack.push(children.mut_slice_from(child_id + 1).mut_iter());
233-
idx += 1;
234-
}
249+
bound!(TrieMapMutIterator, self = self,
250+
key = key, is_upper = upper,
251+
slice_from = mut_slice_from, iter = mut_iter,
252+
mutability = mut)
235253
}
236254

237255
/// Get an iterator pointing to the first key-value pair whose key is not less than `key`.
@@ -464,39 +482,6 @@ pub struct TrieMapIterator<'a, T> {
464482
priv remaining_max: uint
465483
}
466484

467-
impl<'a, T> Iterator<(uint, &'a T)> for TrieMapIterator<'a, T> {
468-
fn next(&mut self) -> Option<(uint, &'a T)> {
469-
while !self.stack.is_empty() {
470-
match self.stack[self.stack.len() - 1].next() {
471-
None => {
472-
self.stack.pop();
473-
}
474-
Some(ref child) => {
475-
match **child {
476-
Internal(ref node) => {
477-
self.stack.push(node.children.iter());
478-
}
479-
External(key, ref value) => {
480-
self.remaining_max -= 1;
481-
if self.remaining_min > 0 {
482-
self.remaining_min -= 1;
483-
}
484-
return Some((key, value));
485-
}
486-
Nothing => {}
487-
}
488-
}
489-
}
490-
}
491-
return None;
492-
}
493-
494-
#[inline]
495-
fn size_hint(&self) -> (uint, Option<uint>) {
496-
(self.remaining_min, Some(self.remaining_max))
497-
}
498-
}
499-
500485
/// Forward iterator over the key-value pairs of a map, with the
501486
/// values being mutable.
502487
pub struct TrieMapMutIterator<'a, T> {
@@ -505,39 +490,51 @@ pub struct TrieMapMutIterator<'a, T> {
505490
priv remaining_max: uint
506491
}
507492

508-
impl<'a, T> Iterator<(uint, &'a mut T)> for TrieMapMutIterator<'a, T> {
509-
fn next(&mut self) -> Option<(uint, &'a mut T)> {
510-
while !self.stack.is_empty() {
511-
match self.stack[self.stack.len() - 1].next() {
512-
None => {
513-
self.stack.pop();
514-
}
515-
Some(child) => {
516-
match *child {
517-
Internal(ref mut node) => {
518-
self.stack.push(node.children.mut_iter());
519-
}
520-
External(key, ref mut value) => {
521-
self.remaining_max -= 1;
522-
if self.remaining_min > 0 {
523-
self.remaining_min -= 1;
493+
// FIXME #5846: see `addr!` above.
494+
macro_rules! item { ($i:item) => {$i}}
495+
496+
macro_rules! iterator_impl {
497+
($name:ident,
498+
iter = $iter:ident,
499+
mutability = $($mut_:tt)*) => {
500+
item!(impl<'a, T> Iterator<(uint, &'a $($mut_)* T)> for $name<'a, T> {
501+
fn next(&mut self) -> Option<(uint, &'a $($mut_)* T)> {
502+
while !self.stack.is_empty() {
503+
match self.stack[self.stack.len() - 1].next() {
504+
None => {
505+
self.stack.pop();
506+
}
507+
Some(child) => {
508+
addr!(match *child {
509+
Internal(ref $($mut_)* node) => {
510+
self.stack.push(node.children.$iter());
511+
}
512+
External(key, ref $($mut_)* value) => {
513+
self.remaining_max -= 1;
514+
if self.remaining_min > 0 {
515+
self.remaining_min -= 1;
516+
}
517+
return Some((key, value));
518+
}
519+
Nothing => {}
520+
})
524521
}
525-
return Some((key, value));
526522
}
527-
Nothing => {}
528523
}
524+
return None;
529525
}
530-
}
531-
}
532-
return None;
533-
}
534526

535-
#[inline]
536-
fn size_hint(&self) -> (uint, Option<uint>) {
537-
(self.remaining_min, Some(self.remaining_max))
527+
#[inline]
528+
fn size_hint(&self) -> (uint, Option<uint>) {
529+
(self.remaining_min, Some(self.remaining_max))
530+
}
531+
})
538532
}
539533
}
540534

535+
iterator_impl! { TrieMapIterator, iter = iter, mutability = }
536+
iterator_impl! { TrieMapMutIterator, iter = mut_iter, mutability = mut }
537+
541538
/// Forward iterator over a set
542539
pub struct TrieSetIterator<'a> {
543540
priv iter: TrieMapIterator<'a, ()>

0 commit comments

Comments
 (0)