Skip to content

Commit

Permalink
Auto merge of #75480 - ssomers:btree_check_invariant, r=Mark-Simulacrum
Browse files Browse the repository at this point in the history
BTreeMap: check some invariants in unit tests
  • Loading branch information
bors committed Aug 19, 2020
2 parents e7f6ed1 + 8c4641b commit 443e177
Show file tree
Hide file tree
Showing 3 changed files with 302 additions and 41 deletions.
35 changes: 3 additions & 32 deletions library/alloc/src/collections/btree/map.rs
Expand Up @@ -1236,10 +1236,10 @@ impl<K: Ord, V> BTreeMap<K, V> {
right_root.fix_left_border();

if left_root.height() < right_root.height() {
self.recalc_length();
self.length = left_root.node_as_ref().calc_length();
right.length = total_num - self.len();
} else {
right.recalc_length();
right.length = right_root.node_as_ref().calc_length();
self.length = total_num - right.len();
}

Expand Down Expand Up @@ -1283,42 +1283,13 @@ impl<K: Ord, V> BTreeMap<K, V> {
{
DrainFilter { pred, inner: self.drain_filter_inner() }
}

pub(super) fn drain_filter_inner(&mut self) -> DrainFilterInner<'_, K, V> {
let root_node = self.root.as_mut().map(|r| r.node_as_mut());
let front = root_node.map(|rn| rn.first_leaf_edge());
DrainFilterInner { length: &mut self.length, cur_leaf_edge: front }
}

/// Calculates the number of elements if it is incorrect.
fn recalc_length(&mut self) {
fn dfs<'a, K, V>(node: NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal>) -> usize
where
K: 'a,
V: 'a,
{
let mut res = node.len();

if let Internal(node) = node.force() {
let mut edge = node.first_edge();
loop {
res += dfs(edge.reborrow().descend());
match edge.right_kv() {
Ok(right_kv) => {
edge = right_kv.right_edge();
}
Err(_) => {
break;
}
}
}
}

res
}

self.length = dfs(self.root.as_ref().unwrap().node_as_ref());
}

/// Creates a consuming iterator visiting all the keys, in sorted order.
/// The map cannot be used after calling this.
/// The iterator element type is `K`.
Expand Down

0 comments on commit 443e177

Please sign in to comment.