diff --git a/src/libstd/treemap.rs b/src/libstd/treemap.rs index d8deea60725..d3583828f9a 100644 --- a/src/libstd/treemap.rs +++ b/src/libstd/treemap.rs @@ -576,63 +576,62 @@ pure fn each_reverse(node: &r/Option<~TreeNode>, } // Remove left horizontal link by rotating right -fn skew(mut node: ~TreeNode) -> ~TreeNode { +fn skew(node: &mut ~TreeNode) { if node.left.map_default(false, |x| x.level == node.level) { let mut save = node.left.swap_unwrap(); node.left <-> save.right; // save.right now None - save.right = Some(node); - save - } else { - node // nothing to do + *node <-> save; + node.right = Some(save); } } // Remove dual horizontal link by rotating left and increasing level of // the parent -fn split(mut node: ~TreeNode) -> ~TreeNode { +fn split(node: &mut ~TreeNode) { if node.right.map_default(false, |x| x.right.map_default(false, |y| y.level == node.level)) { let mut save = node.right.swap_unwrap(); node.right <-> save.left; // save.left now None - save.left = Some(node); save.level += 1; - save - } else { - node // nothing to do + *node <-> save; + node.left = Some(save); } } fn insert(node: &mut Option<~TreeNode>, key: K, value: V) -> bool { - if node.is_none() { - *node = Some(~TreeNode::new(key, value)); - true - } else { - let mut save = node.swap_unwrap(); + match *node { + Some(ref mut save) => { if key < save.key { let inserted = insert(&mut save.left, key, value); - *node = Some(split(skew(save))); // re-balance, if necessary + skew(save); + split(save); inserted } else if save.key < key { let inserted = insert(&mut save.right, key, value); - *node = Some(split(skew(save))); // re-balance, if necessary + skew(save); + split(save); inserted } else { save.key = key; save.value = value; - *node = Some(save); false } + } + None => { + *node = Some(~TreeNode::new(key, value)); + true + } } } fn remove(node: &mut Option<~TreeNode>, key: &K) -> bool { - fn heir_swap(node: &mut TreeNode, + fn heir_swap(node: &mut ~TreeNode, child: &mut Option<~TreeNode>) { // *could* be done without recursion, but it won't borrow check do child.mutate |mut child| { if child.right.is_some() { - heir_swap(&mut *node, &mut child.right); + heir_swap(node, &mut child.right); } else { node.key <-> child.key; node.value <-> child.value; @@ -641,15 +640,15 @@ fn remove(node: &mut Option<~TreeNode>, key: &K) -> bool { } } - if node.is_none() { + match *node { + None => { return false // bottom of tree - } else { - let mut save = node.swap_unwrap(); - - let removed = if save.key < *key { - remove(&mut save.right, key) + } + Some(ref mut save) => { + let (removed, this) = if save.key < *key { + (remove(&mut save.right, key), false) } else if *key < save.key { - remove(&mut save.left, key) + (remove(&mut save.left, key), false) } else { if save.left.is_some() { if save.right.is_some() { @@ -663,16 +662,22 @@ fn remove(node: &mut Option<~TreeNode>, key: &K) -> bool { save.left = Some(left); remove(&mut save.left, key); } else { - save = save.left.swap_unwrap(); + *save = save.left.swap_unwrap(); } + (true, false) } else if save.right.is_some() { - save = save.right.swap_unwrap(); + *save = save.right.swap_unwrap(); + (true, false) } else { - return true // leaf + (true, true) } - true }; + if this { + *node = None; + return true; + } + let left_level = save.left.map_default(0, |x| x.level); let right_level = save.right.map_default(0, |x| x.level); @@ -684,19 +689,28 @@ fn remove(node: &mut Option<~TreeNode>, key: &K) -> bool { do save.right.mutate |mut x| { x.level = save.level; x } } - save = skew(save); + skew(save); - do save.right.mutate |mut right| { - right = skew(right); - right.right.mutate(skew); - right + match save.right { + Some(ref mut right) => { + skew(right); + match right.right { + Some(ref mut x) => { skew(x) }, + None => () + } + } + None => () + } + + split(save); + match save.right { + Some(ref mut x) => { split(x) }, + None => () } - save = split(save); - save.right.mutate(split); } - *node = Some(save); removed + } } }