Skip to content

Commit 252908a

Browse files
authored
Merge pull request scala/scala#10632 from lrytz/t12921
Fix RedBlackTree.doFrom
2 parents 5416735 + fc420ca commit 252908a

File tree

1 file changed

+19
-4
lines changed

1 file changed

+19
-4
lines changed

library/src/scala/collection/immutable/RedBlackTree.scala

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,21 @@ import scala.runtime.Statics.releaseFence
2626
* optimizations behind a reasonably clean API.
2727
*/
2828
private[collection] object RedBlackTree {
29+
def validate[A](tree: Tree[A, _])(implicit ordering: Ordering[A]): tree.type = {
30+
def impl(tree: Tree[A, _], keyProp: A => Boolean): Int = {
31+
assert(keyProp(tree.key), s"key check failed: $tree")
32+
if (tree.isRed) {
33+
assert(tree.left == null || tree.left.isBlack, s"red-red left $tree")
34+
assert(tree.right == null || tree.right.isBlack, s"red-red right $tree")
35+
}
36+
val leftBlacks = if (tree.left == null) 0 else impl(tree.left, k => keyProp(k) && ordering.compare(k, tree.key) < 0)
37+
val rightBlacks = if (tree.right == null) 0 else impl(tree.right, k => keyProp(k) && ordering.compare(k, tree.key) > 0)
38+
assert(leftBlacks == rightBlacks, s"not balanced: $tree")
39+
leftBlacks + (if (tree.isBlack) 1 else 0)
40+
}
41+
if (tree != null) impl(tree, _ => true)
42+
tree
43+
}
2944

3045
def isEmpty(tree: Tree[_, _]): Boolean = tree eq null
3146

@@ -437,23 +452,23 @@ private[collection] object RedBlackTree {
437452
if (ordering.lt(tree.key, from)) return doFrom(tree.right, from)
438453
val newLeft = doFrom(tree.left, from)
439454
if (newLeft eq tree.left) tree
440-
else if (newLeft eq null) upd(tree.right, tree.key, tree.value, overwrite = false)
455+
else if (newLeft eq null) maybeBlacken(upd(tree.right, tree.key, tree.value, overwrite = false))
441456
else join(newLeft, tree.key, tree.value, tree.right)
442457
}
443458
private[this] def doTo[A, B](tree: Tree[A, B], to: A)(implicit ordering: Ordering[A]): Tree[A, B] = {
444459
if (tree eq null) return null
445460
if (ordering.lt(to, tree.key)) return doTo(tree.left, to)
446461
val newRight = doTo(tree.right, to)
447462
if (newRight eq tree.right) tree
448-
else if (newRight eq null) upd(tree.left, tree.key, tree.value, overwrite = false)
449-
else join (tree.left, tree.key, tree.value, newRight)
463+
else if (newRight eq null) maybeBlacken(upd(tree.left, tree.key, tree.value, overwrite = false))
464+
else join(tree.left, tree.key, tree.value, newRight)
450465
}
451466
private[this] def doUntil[A, B](tree: Tree[A, B], until: A)(implicit ordering: Ordering[A]): Tree[A, B] = {
452467
if (tree eq null) return null
453468
if (ordering.lteq(until, tree.key)) return doUntil(tree.left, until)
454469
val newRight = doUntil(tree.right, until)
455470
if (newRight eq tree.right) tree
456-
else if (newRight eq null) upd(tree.left, tree.key, tree.value, overwrite = false)
471+
else if (newRight eq null) maybeBlacken(upd(tree.left, tree.key, tree.value, overwrite = false))
457472
else join(tree.left, tree.key, tree.value, newRight)
458473
}
459474

0 commit comments

Comments
 (0)