Skip to content

Commit 8ffaf9c

Browse files
mkeskellsmkeskells
committed
mutable RedBlackTree
1 parent 7868448 commit 8ffaf9c

File tree

2 files changed

+318
-41
lines changed

2 files changed

+318
-41
lines changed

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

Lines changed: 308 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,8 @@ package scala
1414
package collection
1515
package immutable
1616

17-
import collection.Iterator
18-
19-
import scala.annotation.tailrec
2017
import scala.annotation.meta.getter
21-
22-
import java.lang.{Integer, String}
18+
import scala.annotation.tailrec
2319

2420
/** An object containing the RedBlack tree implementation used by for `TreeMaps` and `TreeSets`.
2521
*
@@ -48,6 +44,10 @@ private[collection] object RedBlackTree {
4844

4945
def count(tree: Tree[_, _]) = if (tree eq null) 0 else tree.count
5046
def update[A: Ordering, B, B1 >: B](tree: Tree[A, B], k: A, v: B1, overwrite: Boolean): Tree[A, B1] = blacken(upd(tree, k, v, overwrite))
47+
def mutableUpdate[A: Ordering, B, B1 >: B](tree: Tree[A, B], k: A, v: B1, overwrite: Boolean): Tree[A, B1] = mutableUpd(tree, k, v, overwrite)
48+
def afterMutableUpdate[A: Ordering, B, B1 >: B](tree: Tree[A, B]): Tree[A, B1] =
49+
if (tree eq null) null
50+
else blacken(tree) //blacken calls black calls .makeImmutable
5151
def delete[A: Ordering, B](tree: Tree[A, B], k: A): Tree[A, B] = blacken(del(tree, k))
5252
def rangeImpl[A: Ordering, B](tree: Tree[A, B], from: Option[A], until: Option[A]): Tree[A, B] = (from, until) match {
5353
case (Some(from), Some(until)) => this.range(tree, from, until)
@@ -157,8 +157,8 @@ private[collection] object RedBlackTree {
157157

158158
def isBlack(tree: Tree[_, _]) = (tree eq null) || isBlackTree(tree)
159159

160-
@`inline` private[this] def isRedTree(tree: Tree[_, _]) = tree.isInstanceOf[RedTree[_, _]]
161-
@`inline` private[this] def isBlackTree(tree: Tree[_, _]) = tree.isInstanceOf[BlackTree[_, _]]
160+
@`inline` private[this] def isRedTree(tree: Tree[_, _]) = (tree ne null) && tree.isRed
161+
@`inline` private[this] def isBlackTree(tree: Tree[_, _]) = (tree ne null) && tree.isBlack
162162

163163
private[this] def blacken[A, B](t: Tree[A, B]): Tree[A, B] = if (t eq null) null else t.black
164164

@@ -186,6 +186,166 @@ private[collection] object RedBlackTree {
186186
else
187187
mkTree(isBlack, x, xv, a, r)
188188
}
189+
190+
/**
191+
* balances a tree with a newLeft
192+
* Origin tree
193+
* tree
194+
* -- KV R newLeft
195+
* nl.L nl.KV nl.R
196+
* nl.R.L nl.R.KV nl.R.R
197+
*/
198+
private[this] def balanceLeft1[A, B, B1 >: B](tree_isBlack: Boolean, tree_key: A, tree_value: B, newLeft: Tree[A, B1], tree_right: Tree[A, B1]): Tree[A, B1] = {
199+
if (isRedTree(newLeft)) {
200+
if (isRedTree(newLeft.left)) {
201+
// produce tree
202+
// RED
203+
// black(nl.L) nl.KV black
204+
// nl.R KV R
205+
RedTree(newLeft.key, newLeft.value, newLeft.left.black, BlackTree(tree_key, tree_value, newLeft.right, tree_right))
206+
} else if (isRedTree(newLeft.right)) {
207+
// produce tree
208+
// RED
209+
// black nl.R.KV black
210+
// nl.L nl.KV nl.R.L nl.R.R KV R
211+
RedTree(newLeft.right.key, newLeft.right.value, BlackTree(newLeft.key, newLeft.value, newLeft.left, newLeft.right.left), BlackTree(tree_key, tree_value, newLeft.right.right, tree_right))
212+
} else {
213+
// tree
214+
// newLeft KV R
215+
mkTree(tree_isBlack, tree_key, tree_value, newLeft, tree_right)
216+
}
217+
} else {
218+
// tree
219+
// newLeft KV R
220+
mkTree(tree_isBlack, tree_key, tree_value, newLeft, tree_right)
221+
}
222+
}
223+
/**
224+
* balances a tree with a newRight
225+
* Origin tree
226+
* tree
227+
* L KV -- newRight
228+
* nr.L nr.KV nr.R
229+
* nr.L.L nr.L.KV nr.L.R
230+
*/
231+
private[this] def balanceRight1[A, B, B1 >: B](tree_isBlack: Boolean, tree_key: A, tree_value: B, tree_left: Tree[A, B1], newRight: Tree[A, B1]): Tree[A, B1] = {
232+
if (isRedTree(newRight)) {
233+
if (isRedTree(newRight.left)) {
234+
// produce
235+
// RED
236+
// black nr.L.KV black
237+
// L KV nr.L.L nr.L.R nr.KV nr.R
238+
RedTree(newRight.left.key, newRight.left.value, BlackTree(tree_key, tree_value, tree_left, newRight.left.left), BlackTree(newRight.key, newRight.value, newRight.left.right, newRight.right))
239+
} else if (isRedTree(newRight.right)) {
240+
// produce
241+
// RED
242+
// black nr.KV black(nr.R)
243+
// L KV nr.L
244+
RedTree(newRight.key, newRight.value, BlackTree(tree_key, tree_value, tree_left, newRight.left), newRight.right.black)
245+
} else {
246+
// tree
247+
// L KV newRight
248+
mkTree(tree_isBlack, tree_key, tree_value, tree_left, newRight)
249+
}
250+
} else {
251+
// tree
252+
// L KV newRight
253+
mkTree(tree_isBlack, tree_key, tree_value, tree_left, newRight)
254+
}
255+
}
256+
/**
257+
* balances a tree with a newLeft
258+
* Origin tree
259+
* tree
260+
* -- KV R newLeft
261+
* nl.L nl.KV nl.R
262+
* nl.R.L nl.R.KV nl.R.R
263+
*/
264+
private[this] def balanceLeft[A, B, B1 >: B](tree: Tree[A, B], newLeft: Tree[A, B1]): Tree[A, B1] = {
265+
// if (tree.left eq newLeft) tree else
266+
if (isRedTree(newLeft)) {
267+
val newLeft_left = newLeft.left
268+
val newLeft_right = newLeft.right
269+
if (isRedTree(newLeft_left)) {
270+
// produce tree
271+
// RED
272+
// black(nl.L) nl.KV black
273+
// nl.R KV R
274+
val resultLeft = newLeft_left.mutableBlack
275+
val resultRight = tree.mutableBlackWithLeft(newLeft_right)
276+
277+
newLeft.mutableWithLeftRight(resultLeft, resultRight)
278+
} else if (isRedTree(newLeft_right)) {
279+
// produce tree
280+
// RED
281+
// black nl.R.KV black
282+
// nl.L nl.KV nl.R.L nl.R.R KV R
283+
284+
val newLeft_right_right = newLeft_right.right
285+
286+
val resultLeft = newLeft.mutableBlackWithRight(newLeft_right.left)
287+
val resultRight = tree.mutableBlackWithLeft(newLeft_right_right)
288+
289+
newLeft_right.mutableWithLeftRight(resultLeft, resultRight)
290+
} else {
291+
// tree
292+
// newLeft KV R
293+
tree.mutableWithLeft(newLeft)
294+
}
295+
} else {
296+
// tree
297+
// newLeft KV R
298+
tree.mutableWithLeft(newLeft)
299+
}
300+
}
301+
/**
302+
* balances a tree with a newRight
303+
* Origin tree
304+
* tree
305+
* L KV -- newRight
306+
* nr.L nr.KV nr.R
307+
* nr.L.L nr.L.KV nr.L.R
308+
*/
309+
private[this] def balanceRight[A, B, B1 >: B](tree: Tree[A, B], newRight: Tree[A, B1]): Tree[A, B1] = {
310+
// if (tree.right eq newRight) tree else
311+
if (isRedTree(newRight)) {
312+
val newRight_left = newRight.left
313+
if (isRedTree(newRight_left)) {
314+
// produce
315+
// RED
316+
// black nr.L.KV black
317+
// L KV nr.L.L nr.L.R nr.KV nr.R
318+
319+
320+
val resultLeft = tree.mutableBlackWithRight(newRight_left.left)
321+
val resultRight = newRight.mutableBlackWithLeft(newRight_left.right)
322+
323+
newRight_left.mutableWithLeftRight(resultLeft, resultRight)
324+
325+
} else {
326+
val newRight_right = newRight.right
327+
if (isRedTree(newRight_right)) {
328+
// produce
329+
// RED
330+
// black nr.KV black(nr.R)
331+
// L KV nr.L
332+
333+
val resultLeft = tree.mutableBlackWithRight(newRight_left)
334+
val resultRight = newRight_right.mutableBlack
335+
336+
newRight.mutableWithLeftRight(resultLeft, resultRight)
337+
} else {
338+
// tree
339+
// L KV newRight
340+
tree.mutableWithRight(newRight)
341+
}
342+
}
343+
} else {
344+
// tree
345+
// L KV newRight
346+
tree.mutableWithRight(newRight)
347+
}
348+
}
189349
private[this] def upd[A, B, B1 >: B](tree: Tree[A, B], k: A, v: B1, overwrite: Boolean)(implicit ordering: Ordering[A]): Tree[A, B1] = if (tree eq null) {
190350
RedTree(k, v, null, null)
191351
} else {
@@ -195,6 +355,15 @@ private[collection] object RedBlackTree {
195355
else if (overwrite || k != tree.key) mkTree(isBlackTree(tree), tree.key, v, tree.left, tree.right)
196356
else tree
197357
}
358+
private[this] def mutableUpd[A, B, B1 >: B](tree: Tree[A, B], k: A, v: B1, overwrite: Boolean)(implicit ordering: Ordering[A]): Tree[A, B1] = if (tree eq null) {
359+
mutableRedTree(k, v, null, null)
360+
} else {
361+
val cmp = ordering.compare(k, tree.key)
362+
if (cmp < 0) balanceLeft(tree, mutableUpd(tree.left, k, v, overwrite))
363+
else if (cmp > 0) balanceRight(tree, mutableUpd(tree.right, k, v, overwrite))
364+
else if (overwrite || k != tree.key) tree.mutableWithkV(k,v) //*** should that be tree.key ???
365+
else tree
366+
}
198367
private[this] def updNth[A, B, B1 >: B](tree: Tree[A, B], idx: Int, k: A, v: B1, overwrite: Boolean): Tree[A, B1] = if (tree eq null) {
199368
RedTree(k, v, null, null)
200369
} else {
@@ -281,41 +450,139 @@ private[collection] object RedBlackTree {
281450
*
282451
* An alternative is to implement the these classes using plain old Java code...
283452
*/
284-
sealed abstract class Tree[A, +B](
285-
@(`inline` @getter) final val key: A,
286-
@(`inline` @getter) final val value: B,
287-
@(`inline` @getter) final val left: Tree[A, B],
288-
@(`inline` @getter) final val right: Tree[A, B])
453+
final class Tree[A, +B](
454+
@(inline @getter) private var _key: A,
455+
@(inline @getter) private var _value: AnyRef,
456+
@(inline @getter) private var _left: Tree[A, _],
457+
@(inline @getter) private var _right: Tree[A, _],
458+
@(inline @getter) private var _count: Int)
289459
{
290-
@(`inline` @getter) final val count: Int = 1 + RedBlackTree.count(left) + RedBlackTree.count(right)
291-
def black: Tree[A, B]
292-
def red: Tree[A, B]
293-
}
294-
final class RedTree[A, +B](key: A,
295-
value: B,
296-
left: Tree[A, B],
297-
right: Tree[A, B]) extends Tree[A, B](key, value, left, right) {
298-
override def black: Tree[A, B] = BlackTree(key, value, left, right)
299-
override def red: Tree[A, B] = this
300-
override def toString: String = "RedTree(" + key + ", " + value + ", " + left + ", " + right + ")"
301-
}
302-
final class BlackTree[A, +B](key: A,
303-
value: B,
304-
left: Tree[A, B],
305-
right: Tree[A, B]) extends Tree[A, B](key, value, left, right) {
306-
override def black: Tree[A, B] = this
307-
override def red: Tree[A, B] = RedTree(key, value, left, right)
308-
override def toString: String = "BlackTree(" + key + ", " + value + ", " + left + ", " + right + ")"
309-
}
310-
311-
object RedTree {
312-
@`inline` def apply[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = new RedTree(key, value, left, right)
313-
def unapply[A, B](t: RedTree[A, B]) = Some((t.key, t.value, t.left, t.right))
314-
}
315-
object BlackTree {
316-
@`inline` def apply[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = new BlackTree(key, value, left, right)
317-
def unapply[A, B](t: BlackTree[A, B]) = Some((t.key, t.value, t.left, t.right))
318-
}
460+
// read only APIs
461+
@`inline` final def count = _count & 0x7FFFFFFF
462+
@`inline` final def key = _key
463+
@`inline` final def value = _value.asInstanceOf[B]
464+
@`inline` final def left = _left.asInstanceOf[Tree[A, B]]
465+
@`inline` final def right = _right.asInstanceOf[Tree[A, B]]
466+
@`inline` final def isBlack = _count < 0
467+
@`inline` final def isRed = _count >= 0
468+
469+
override def toString: String = s"${if(isRed) "RedTree" else "BlackTree"}($key, $value, $left, $right)"
470+
471+
@`inline` private def initialCount = if (isBlack) initialBlackCount else initialRedCount
472+
//mutable APIs
473+
@`inline` private def mutable = count == 0
474+
def makeImmutable: Tree[A, B] = {
475+
if (mutable) {
476+
var size = 1
477+
if (_left ne null) {
478+
_left.makeImmutable
479+
size += _left.count
480+
}
481+
if (_right ne null) {
482+
_right.makeImmutable
483+
size += _right.count
484+
}
485+
_count |= size //retains colour
486+
}
487+
this
488+
}
489+
490+
def mutableBlack: Tree[A, B] = {
491+
if (isBlack) this
492+
else if (mutable) {
493+
_count = initialBlackCount;
494+
this
495+
}
496+
else new Tree(_key, _value, _left, _right, initialBlackCount)
497+
}
498+
def mutableRed: Tree[A, B] = {
499+
if (isRed) this
500+
else if (mutable) {
501+
_count = initialRedCount;
502+
this
503+
}
504+
else new Tree(_key, _value, _left, _right, initialRedCount)
505+
}
506+
def mutableWithkV[B1 >: B](key: A, value: B1): Tree[A, B1] = {
507+
if (mutable) {
508+
_key = key
509+
_value = value.asInstanceOf[AnyRef]
510+
this
511+
} else new Tree(key, value.asInstanceOf[AnyRef], _left, _right, _count)
512+
}
513+
def mutableWithLeft[B1 >: B](newLeft: Tree[A, B1]): Tree[A, B1] = {
514+
if (mutable) {
515+
_left = newLeft
516+
this
517+
} else new Tree(_key, _value, newLeft, _right, initialCount)
518+
}
519+
def mutableWithRight[B1 >: B](newRight: Tree[A, B1]): Tree[A, B1] = {
520+
if (mutable) {
521+
_right = newRight
522+
this
523+
} else new Tree(_key, _value, _left, newRight, initialCount)
524+
}
525+
def mutableWithLeftRight[B1 >: B](newLeft: Tree[A, B1], newRight: Tree[A, B1]): Tree[A, B1] = {
526+
if (mutable) {
527+
_left = newLeft
528+
_right = newRight
529+
this
530+
} else new Tree(_key, _value, newLeft, newRight, initialCount)
531+
}
532+
def mutableBlackWithLeft[B1 >: B](newLeft: Tree[A, B1]): Tree[A, B1] = {
533+
if (mutable) {
534+
_count = initialBlackCount
535+
_left = newLeft
536+
this
537+
} else new Tree(_key, _value, newLeft, _right, initialBlackCount)
538+
}
539+
def mutableBlackWithRight[B1 >: B](newRight: Tree[A, B1]): Tree[A, B1] = {
540+
if (mutable) {
541+
_count = initialBlackCount
542+
_right = newRight
543+
this
544+
} else new Tree(_key, _value, _left, newRight, initialBlackCount)
545+
}
546+
def mutableBlackWithLeftRight[B1 >: B](newLeft: Tree[A, B1], newRight: Tree[A, B1]): Tree[A, B1] = {
547+
if (mutable) {
548+
_count = initialBlackCount
549+
_left = newLeft
550+
_right = newRight
551+
this
552+
} else new Tree(_key, _value, newLeft, newRight, initialBlackCount)
553+
}
554+
//immutable APIs
555+
def black: Tree[A, B] = mutableBlack.makeImmutable
556+
def red: Tree[A, B] = mutableRed.makeImmutable
557+
def withkV[B1 >: B](key: A, value: B1): Tree[A, B1] = mutableWithkV(key,value).makeImmutable
558+
def withLeft[B1 >: B](newLeft: Tree[A, B1]): Tree[A, B1] = mutableWithLeft(newLeft).makeImmutable
559+
def withRight[B1 >: B](newRight: Tree[A, B1]): Tree[A, B1] = mutableWithRight(newRight).makeImmutable
560+
def withLeftRight[B1 >: B](newLeft: Tree[A, B1], newRight: Tree[A, B1]): Tree[A, B1] = mutableWithLeftRight(newLeft, newRight).makeImmutable
561+
}
562+
//
563+
// final class RedTree[A, +B](key: A,
564+
// value: B,
565+
// left: Tree[A, B],
566+
// right: Tree[A, B]) extends Tree[A, B](key, value, left, right) {
567+
// override def black: Tree[A, B] = BlackTree(key, value, left, right)
568+
// override def red: Tree[A, B] = this
569+
// override def toString: String = "RedTree(" + key + ", " + value + ", " + left + ", " + right + ")"
570+
// }
571+
// final class BlackTree[A, +B](key: A,
572+
// value: B,
573+
// left: Tree[A, B],
574+
// right: Tree[A, B]) extends Tree[A, B](key, value, left, right) {
575+
// override def black: Tree[A, B] = this
576+
// override def red: Tree[A, B] = RedTree(key, value, left, right)
577+
// override def toString: String = "BlackTree(" + key + ", " + value + ", " + left + ", " + right + ")"
578+
// }
579+
@inline def initialBlackCount = Int.MinValue
580+
@inline def initialRedCount = 0
581+
582+
@`inline` def mutableRedTree[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = new Tree(key, value.asInstanceOf[AnyRef], left, right, initialRedCount)
583+
@`inline` def mutableBlackTree[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = new Tree(key, value.asInstanceOf[AnyRef], left, right, initialBlackCount)
584+
@`inline` def RedTree[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = mutableRedTree(key, value, left, right).makeImmutable
585+
@`inline` def BlackTree[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = mutableBlackTree(key, value, left, right).makeImmutable
319586

320587
private[this] abstract class TreeIterator[A, B, R](root: Tree[A, B], start: Option[A])(implicit ordering: Ordering[A]) extends AbstractIterator[R] {
321588
protected[this] def nextResult(tree: Tree[A, B]): R

0 commit comments

Comments
 (0)