Skip to content

Commit 33f4a2b

Browse files
mkeskellsmkeskells
committed
mutable RedBlackTree to reduce garbage
1 parent d03e693 commit 33f4a2b

File tree

3 files changed

+311
-46
lines changed

3 files changed

+311
-46
lines changed

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

Lines changed: 271 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-
17+
import scala.annotation.meta.{getter, setter}
1918
import scala.annotation.tailrec
20-
import scala.annotation.meta.getter
21-
22-
import java.lang.{Integer, String}
2319

2420
/** An object containing the RedBlack tree implementation used by for `TreeMaps` and `TreeSets`.
2521
*
@@ -45,6 +41,143 @@ private[collection] object NewRedBlackTree {
4541
else if (cmp > 0) lookup(tree.right, x)
4642
else tree
4743
}
44+
private[immutable] abstract class Helper[A] {
45+
def beforePublish[B](tree: Tree[A, B]): Tree[A, B] = {
46+
if (tree eq null) tree
47+
else {
48+
val res = tree.mutableBlack.makeImmutable
49+
VM.releaseFence()
50+
res
51+
}
52+
}
53+
/** Create a new balanced tree where `newLeft` replaces `tree.left`.
54+
* tree and newLeft are never null */
55+
protected[this] def mutableBalanceLeft[A, B, B1 >: B](tree: Tree[A, B], newLeft: Tree[A, B1]): Tree[A, B1] = {
56+
// Parameter trees
57+
// tree | newLeft
58+
// -- KV R | nl.L nl.KV nl.R
59+
// | nl.R.L nl.R.KV nl.R.R
60+
if (tree.left eq newLeft) tree
61+
else if (newLeft.isRed) {
62+
val newLeft_left = newLeft.left
63+
val newLeft_right = newLeft.right
64+
if (isRedTree(newLeft_left)) {
65+
// RED
66+
// black(nl.L) nl.KV black
67+
// nl.R KV R
68+
val resultLeft = newLeft_left.mutableBlack
69+
val resultRight = tree.mutableBlackWithLeft(newLeft_right)
70+
71+
newLeft.mutableWithLeftRight(resultLeft, resultRight)
72+
} else if (isRedTree(newLeft_right)) {
73+
// RED
74+
// black nl.R.KV black
75+
// nl.L nl.KV nl.R.L nl.R.R KV R
76+
77+
val newLeft_right_right = newLeft_right.right
78+
79+
val resultLeft = newLeft.mutableBlackWithRight(newLeft_right.left)
80+
val resultRight = tree.mutableBlackWithLeft(newLeft_right_right)
81+
82+
newLeft_right.mutableWithLeftRight(resultLeft, resultRight)
83+
} else {
84+
// tree
85+
// newLeft KV R
86+
tree.mutableWithLeft(newLeft)
87+
}
88+
} else {
89+
// tree
90+
// newLeft KV R
91+
tree.mutableWithLeft(newLeft)
92+
}
93+
}
94+
/** Create a new balanced tree where `newRight` replaces `tree.right`.
95+
* tree and newRight are never null */
96+
protected[this] def mutableBalanceRight[A, B, B1 >: B](tree: Tree[A, B], newRight: Tree[A, B1]): Tree[A, B1] = {
97+
// Parameter trees
98+
// tree | newRight
99+
// L KV -- | nr.L nr.KV nr.R
100+
// | nr.L.L nr.L.KV nr.L.R
101+
if (tree.right eq newRight) tree
102+
else if (newRight.isRed) {
103+
val newRight_left = newRight.left
104+
if (isRedTree(newRight_left)) {
105+
// RED
106+
// black nr.L.KV black
107+
// L KV nr.L.L nr.L.R nr.KV nr.R
108+
109+
val resultLeft = tree.mutableBlackWithRight(newRight_left.left)
110+
val resultRight = newRight.mutableBlackWithLeft(newRight_left.right)
111+
112+
newRight_left.mutableWithLeftRight(resultLeft, resultRight)
113+
114+
} else {
115+
val newRight_right = newRight.right
116+
if (isRedTree(newRight_right)) {
117+
// RED
118+
// black nr.KV black(nr.R)
119+
// L KV nr.L
120+
121+
val resultLeft = tree.mutableBlackWithRight(newRight_left)
122+
val resultRight = newRight_right.mutableBlack
123+
124+
newRight.mutableWithLeftRight(resultLeft, resultRight)
125+
} else {
126+
// tree
127+
// L KV newRight
128+
tree.mutableWithRight(newRight)
129+
}
130+
}
131+
} else {
132+
// tree
133+
// L KV newRight
134+
tree.mutableWithRight(newRight)
135+
}
136+
}
137+
}
138+
private[immutable] class SetHelper[A](implicit ordering: Ordering[A]) extends Helper[A] {
139+
def addMutable(tree: Tree[A, Any], k: A): Tree[A, Any] = {
140+
mutableUpd(tree, k)
141+
}
142+
143+
protected[this] def mutableUpd(tree: Tree[A, Any], k: A): Tree[A, Any] =
144+
if (tree eq null) {
145+
mutableRedTree(k, (), null, null)
146+
} else if (k.asInstanceOf[AnyRef] eq tree.key.asInstanceOf[AnyRef]) {
147+
tree
148+
} else {
149+
val cmp = ordering.compare(k, tree.key)
150+
if (cmp < 0)
151+
mutableBalanceLeft(tree, mutableUpd(tree.left, k))
152+
else if (cmp > 0)
153+
mutableBalanceRight(tree, mutableUpd(tree.right, k))
154+
else tree
155+
}
156+
}
157+
private[immutable] class MapHelper[A, B](implicit ordering: Ordering[A]) extends Helper[A] {
158+
159+
def addMutable(tree: Tree[A, B], k: A, v: B): Tree[A, B] = {
160+
mutableUpd(tree, k, v)
161+
}
162+
163+
protected[this] def mutableUpd[B1 >: B](tree: Tree[A, B], k: A, v: B1): Tree[A, B1] =
164+
if (tree eq null) {
165+
mutableRedTree(k, v, null, null)
166+
} else if (k.asInstanceOf[AnyRef] eq tree.key.asInstanceOf[AnyRef]) {
167+
if (v.asInstanceOf[AnyRef] ne tree.value.asInstanceOf[AnyRef])
168+
tree.mutableWithKV(tree.key, v)
169+
else tree
170+
} else {
171+
val cmp = ordering.compare(k, tree.key)
172+
if (cmp < 0)
173+
mutableBalanceLeft(tree, mutableUpd(tree.left, k, v))
174+
else if (cmp > 0)
175+
mutableBalanceRight(tree, mutableUpd(tree.right, k, v))
176+
else if (v.asInstanceOf[AnyRef] ne tree.value.asInstanceOf[AnyRef])
177+
tree.mutableWithKV(tree.key, v)
178+
else tree
179+
}
180+
}
48181

49182
def count(tree: Tree[_, _]) = if (tree eq null) 0 else tree.count
50183
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))
@@ -176,8 +309,8 @@ private[collection] object NewRedBlackTree {
176309

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

179-
@`inline` private[this] def isRedTree(tree: Tree[_, _]) = tree.isInstanceOf[RedTree[_, _]]
180-
@`inline` private[this] def isBlackTree(tree: Tree[_, _]) = tree.isInstanceOf[BlackTree[_, _]]
312+
@`inline` private[this] def isRedTree(tree: Tree[_, _]) = (tree ne null) && tree.isRed
313+
@`inline` private[this] def isBlackTree(tree: Tree[_, _]) = (tree ne null) && tree.isBlack
181314

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

@@ -231,7 +364,6 @@ private[collection] object NewRedBlackTree {
231364
}
232365
}
233366
}
234-
235367
/** Create a new balanced tree where `newRight` replaces `tree.right`. */
236368
private[this] def balanceRight[A, B1](tree: Tree[A, B1], newRight: Tree[A, B1]): Tree[A, B1] = {
237369
// Parameter trees
@@ -275,6 +407,10 @@ private[collection] object NewRedBlackTree {
275407

276408
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) {
277409
RedTree(k, v, null, null)
410+
} else if (k.asInstanceOf[AnyRef] eq tree.key.asInstanceOf[AnyRef]) {
411+
if (overwrite && (v.asInstanceOf[AnyRef] ne tree.value.asInstanceOf[AnyRef]))
412+
mkTree(isBlackTree(tree), tree.key, v, tree.left, tree.right)
413+
else tree
278414
} else {
279415
val cmp = ordering.compare(k, tree.key)
280416
if (cmp < 0)
@@ -372,41 +508,135 @@ private[collection] object NewRedBlackTree {
372508
*
373509
* An alternative is to implement the these classes using plain old Java code...
374510
*/
375-
sealed abstract class Tree[A, +B](
376-
@(`inline` @getter) final val key: A,
377-
@(`inline` @getter) final val value: B,
378-
@(`inline` @getter) final val left: Tree[A, B],
379-
@(`inline` @getter) final val right: Tree[A, B]) extends Serializable
511+
final class Tree[A, +B](
512+
@(inline @getter @setter) private var _key: A,
513+
@(inline @getter @setter) private var _value: AnyRef,
514+
@(inline @getter @setter) private var _left: Tree[A, _],
515+
@(inline @getter @setter) private var _right: Tree[A, _],
516+
@(inline @getter @setter) private var _count: Int)
380517
{
381-
@(`inline` @getter) final val count: Int = 1 + NewRedBlackTree.count(left) + NewRedBlackTree.count(right)
382-
def black: Tree[A, B]
383-
def red: Tree[A, B]
384-
}
385-
final class RedTree[A, +B](key: A,
386-
value: B,
387-
left: Tree[A, B],
388-
right: Tree[A, B]) extends Tree[A, B](key, value, left, right) {
389-
override def black: Tree[A, B] = BlackTree(key, value, left, right)
390-
override def red: Tree[A, B] = this
391-
override def toString: String = "RedTree(" + key + ", " + value + ", " + left + ", " + right + ")"
392-
}
393-
final class BlackTree[A, +B](key: A,
394-
value: B,
395-
left: Tree[A, B],
396-
right: Tree[A, B]) extends Tree[A, B](key, value, left, right) {
397-
override def black: Tree[A, B] = this
398-
override def red: Tree[A, B] = RedTree(key, value, left, right)
399-
override def toString: String = "BlackTree(" + key + ", " + value + ", " + left + ", " + right + ")"
400-
}
401-
402-
object RedTree {
403-
@`inline` def apply[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = new RedTree(key, value, left, right)
404-
def unapply[A, B](t: RedTree[A, B]) = Some((t.key, t.value, t.left, t.right))
405-
}
406-
object BlackTree {
407-
@`inline` def apply[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = new BlackTree(key, value, left, right)
408-
def unapply[A, B](t: BlackTree[A, B]) = Some((t.key, t.value, t.left, t.right))
518+
// read only APIs
519+
@`inline` final def count = _count & 0x7FFFFFFF
520+
@`inline` final def key = _key
521+
@`inline` final def value = _value.asInstanceOf[B]
522+
@`inline` final def left = _left.asInstanceOf[Tree[A, B]]
523+
@`inline` final def right = _right.asInstanceOf[Tree[A, B]]
524+
@`inline` final def isBlack = _count < 0
525+
@`inline` final def isRed = _count >= 0
526+
527+
override def toString: String = s"${if(isRed) "RedTree" else "BlackTree"}($key, $value, $left, $right)"
528+
529+
@`inline` private def initialCount = _count & 0x80000000
530+
@`inline` def initialBlackCount = 0x80000000
531+
@`inline` def initialRedCount = 0
532+
533+
//mutable APIs
534+
@`inline` def mutable = (_count & 0x7FFFFFFF) == 0
535+
def makeImmutable: Tree[A, B] = {
536+
def makeImmutableImpl() = {
537+
if (mutable) {
538+
var size = 1
539+
if (_left ne null) {
540+
_left.makeImmutable
541+
size += _left.count
542+
}
543+
if (_right ne null) {
544+
_right.makeImmutable
545+
size += _right.count
546+
}
547+
_count |= size //retains colour
548+
}
549+
this
550+
}
551+
makeImmutableImpl()
552+
//TODO call releaseFence after backported
553+
this
554+
}
555+
556+
def mutableBlack: Tree[A, B] = {
557+
if (isBlack) this
558+
else if (mutable) {
559+
_count = initialBlackCount
560+
this
561+
}
562+
else new Tree(_key, _value, _left, _right, initialBlackCount)
563+
}
564+
def mutableRed: Tree[A, B] = {
565+
if (isRed) this
566+
else if (mutable) {
567+
_count = initialRedCount
568+
this
569+
}
570+
else new Tree(_key, _value, _left, _right, initialRedCount)
571+
}
572+
def mutableWithKV[B1 >: B](key: A, value: B1): Tree[A, B1] = {
573+
if (mutable) {
574+
_key = key
575+
_value = value.asInstanceOf[AnyRef]
576+
this
577+
} else new Tree(key, value.asInstanceOf[AnyRef], _left, _right, _count)
578+
}
579+
def mutableWithLeft[B1 >: B](newLeft: Tree[A, B1]): Tree[A, B1] = {
580+
if (mutable) {
581+
_left = newLeft
582+
this
583+
} else new Tree(_key, _value, newLeft, _right, initialCount)
584+
}
585+
def mutableWithRight[B1 >: B](newRight: Tree[A, B1]): Tree[A, B1] = {
586+
if (mutable) {
587+
_right = newRight
588+
this
589+
} else new Tree(_key, _value, _left, newRight, initialCount)
590+
}
591+
def mutableWithLeftRight[B1 >: B](newLeft: Tree[A, B1], newRight: Tree[A, B1]): Tree[A, B1] = {
592+
if (mutable) {
593+
_left = newLeft
594+
_right = newRight
595+
this
596+
} else new Tree(_key, _value, newLeft, newRight, initialCount)
597+
}
598+
def mutableBlackWithLeft[B1 >: B](newLeft: Tree[A, B1]): Tree[A, B1] = {
599+
if (mutable) {
600+
_count = initialBlackCount
601+
_left = newLeft
602+
this
603+
} else new Tree(_key, _value, newLeft, _right, initialBlackCount)
604+
}
605+
def mutableBlackWithRight[B1 >: B](newRight: Tree[A, B1]): Tree[A, B1] = {
606+
if (mutable) {
607+
_count = initialBlackCount
608+
_right = newRight
609+
this
610+
} else new Tree(_key, _value, _left, newRight, initialBlackCount)
611+
}
612+
// def mutableBlackWithLeftRight[B1 >: B](newLeft: Tree[A, B1], newRight: Tree[A, B1]): Tree[A, B1] = {
613+
// if (mutable) {
614+
// _count = initialBlackCount
615+
// _left = newLeft
616+
// _right = newRight
617+
// this
618+
// } else new Tree(_key, _value, newLeft, newRight, initialBlackCount)
619+
// }
620+
//immutable APIs
621+
def black: Tree[A, B] = mutableBlack.makeImmutable
622+
def red: Tree[A, B] = mutableRed.makeImmutable
623+
def withKV[B1 >: B](key: A, value: B1): Tree[A, B1] = mutableWithKV(key,value).makeImmutable
624+
def withLeft[B1 >: B](newLeft: Tree[A, B1]): Tree[A, B1] = mutableWithLeft(newLeft).makeImmutable
625+
def withRight[B1 >: B](newRight: Tree[A, B1]): Tree[A, B1] = mutableWithRight(newRight).makeImmutable
626+
def withLeftRight[B1 >: B](newLeft: Tree[A, B1], newRight: Tree[A, B1]): Tree[A, B1] = mutableWithLeftRight(newLeft, newRight).makeImmutable
409627
}
628+
@`inline` def initialBlackCount = 0x80000000
629+
@`inline` def initialRedCount = 0
630+
631+
@`inline` def mutableRedTree[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = new Tree[A,B](key, value.asInstanceOf[AnyRef], left, right, initialRedCount)
632+
@`inline` def mutableBlackTree[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = new Tree[A,B](key, value.asInstanceOf[AnyRef], left, right, initialBlackCount)
633+
634+
/** create a new immutable red tree.
635+
* left and right may be null
636+
* TODO - make the callers able to cope with mutability
637+
*/
638+
@`inline` def RedTree[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]): Tree[A, B] = mutableRedTree(key, value, left, right).makeImmutable
639+
@`inline` def BlackTree[A, B](key: A, value: B, left: Tree[A, B], right: Tree[A, B]): Tree[A, B] = mutableBlackTree(key, value, left, right).makeImmutable
410640

411641
private[this] abstract class TreeIterator[A, B, R](root: Tree[A, B], start: Option[A])(protected implicit val ordering: Ordering[A]) extends Iterator[R] {
412642
protected[this] def nextResult(tree: Tree[A, B]): R

0 commit comments

Comments
 (0)