Skip to content

Commit 3880eb0

Browse files
committed
Clone stdlib-bootstrapped/src/scala/collection/mutable/TreeMap.scala
1 parent 302be2f commit 3880eb0

File tree

1 file changed

+257
-0
lines changed

1 file changed

+257
-0
lines changed
Lines changed: 257 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,257 @@
1+
/*
2+
* Scala (https://www.scala-lang.org)
3+
*
4+
* Copyright EPFL and Lightbend, Inc.
5+
*
6+
* Licensed under Apache License 2.0
7+
* (http://www.apache.org/licenses/LICENSE-2.0).
8+
*
9+
* See the NOTICE file distributed with this work for
10+
* additional information regarding copyright ownership.
11+
*/
12+
13+
package scala
14+
package collection
15+
package mutable
16+
17+
import scala.collection.Stepper.EfficientSplit
18+
import scala.collection.generic.DefaultSerializable
19+
import scala.collection.mutable.{RedBlackTree => RB}
20+
21+
/**
22+
* A mutable sorted map implemented using a mutable red-black tree as underlying data structure.
23+
*
24+
* @param ordering the implicit ordering used to compare objects of type `A`.
25+
* @tparam K the type of the keys contained in this tree map.
26+
* @tparam V the type of the values associated with the keys.
27+
*
28+
* @define Coll mutable.TreeMap
29+
* @define coll mutable tree map
30+
*/
31+
sealed class TreeMap[K, V] private (tree: RB.Tree[K, V])(implicit val ordering: Ordering[K])
32+
extends AbstractMap[K, V]
33+
with SortedMap[K, V]
34+
with SortedMapOps[K, V, TreeMap, TreeMap[K, V]]
35+
with StrictOptimizedIterableOps[(K, V), Iterable, TreeMap[K, V]]
36+
with StrictOptimizedMapOps[K, V, Map, TreeMap[K, V]]
37+
with StrictOptimizedSortedMapOps[K, V, TreeMap, TreeMap[K, V]]
38+
with SortedMapFactoryDefaults[K, V, TreeMap, Iterable, Map]
39+
with DefaultSerializable {
40+
41+
override def sortedMapFactory = TreeMap
42+
43+
/**
44+
* Creates an empty `TreeMap`.
45+
* @param ord the implicit ordering used to compare objects of type `K`.
46+
* @return an empty `TreeMap`.
47+
*/
48+
def this()(implicit ord: Ordering[K]) = this(RB.Tree.empty)(ord)
49+
50+
def iterator: Iterator[(K, V)] = {
51+
if (isEmpty) Iterator.empty
52+
else RB.iterator(tree)
53+
}
54+
55+
override def keysIterator: Iterator[K] = {
56+
if (isEmpty) Iterator.empty
57+
else RB.keysIterator(tree, None)
58+
}
59+
60+
override def valuesIterator: Iterator[V] = {
61+
if (isEmpty) Iterator.empty
62+
else RB.valuesIterator(tree, None)
63+
}
64+
65+
def keysIteratorFrom(start: K): Iterator[K] = {
66+
if (isEmpty) Iterator.empty
67+
else RB.keysIterator(tree, Some(start))
68+
}
69+
70+
def iteratorFrom(start: K): Iterator[(K, V)] = {
71+
if (isEmpty) Iterator.empty
72+
else RB.iterator(tree, Some(start))
73+
}
74+
75+
override def valuesIteratorFrom(start: K): Iterator[V] = {
76+
if (isEmpty) Iterator.empty
77+
else RB.valuesIterator(tree, Some(start))
78+
}
79+
80+
override def stepper[S <: Stepper[_]](implicit shape: StepperShape[(K, V), S]): S with EfficientSplit =
81+
shape.parUnbox(
82+
scala.collection.convert.impl.AnyBinaryTreeStepper.from[(K, V), RB.Node[K, V]](
83+
size, tree.root, _.left, _.right, x => (x.key, x.value)
84+
)
85+
)
86+
87+
override def keyStepper[S <: Stepper[_]](implicit shape: StepperShape[K, S]): S with EfficientSplit = {
88+
import scala.collection.convert.impl._
89+
type T = RB.Node[K, V]
90+
val s = shape.shape match {
91+
case StepperShape.IntShape => IntBinaryTreeStepper.from[T] (size, tree.root, _.left, _.right, _.key.asInstanceOf[Int])
92+
case StepperShape.LongShape => LongBinaryTreeStepper.from[T] (size, tree.root, _.left, _.right, _.key.asInstanceOf[Long])
93+
case StepperShape.DoubleShape => DoubleBinaryTreeStepper.from[T](size, tree.root, _.left, _.right, _.key.asInstanceOf[Double])
94+
case _ => shape.parUnbox(AnyBinaryTreeStepper.from[K, T](size, tree.root, _.left, _.right, _.key))
95+
}
96+
s.asInstanceOf[S with EfficientSplit]
97+
}
98+
99+
override def valueStepper[S <: Stepper[_]](implicit shape: StepperShape[V, S]): S with EfficientSplit = {
100+
import scala.collection.convert.impl._
101+
type T = RB.Node[K, V]
102+
val s = shape.shape match {
103+
case StepperShape.IntShape => IntBinaryTreeStepper.from[T] (size, tree.root, _.left, _.right, _.value.asInstanceOf[Int])
104+
case StepperShape.LongShape => LongBinaryTreeStepper.from[T] (size, tree.root, _.left, _.right, _.value.asInstanceOf[Long])
105+
case StepperShape.DoubleShape => DoubleBinaryTreeStepper.from[T] (size, tree.root, _.left, _.right, _.value.asInstanceOf[Double])
106+
case _ => shape.parUnbox(AnyBinaryTreeStepper.from[V, T] (size, tree.root, _.left, _.right, _.value))
107+
}
108+
s.asInstanceOf[S with EfficientSplit]
109+
}
110+
111+
def addOne(elem: (K, V)): this.type = { RB.insert(tree, elem._1, elem._2); this }
112+
113+
def subtractOne(elem: K): this.type = { RB.delete(tree, elem); this }
114+
115+
override def clear(): Unit = RB.clear(tree)
116+
117+
def get(key: K): Option[V] = RB.get(tree, key)
118+
119+
/**
120+
* Creates a ranged projection of this map. Any mutations in the ranged projection will update the original map and
121+
* vice versa.
122+
*
123+
* Only entries with keys between this projection's key range will ever appear as elements of this map, independently
124+
* of whether the entries are added through the original map or through this view. That means that if one inserts a
125+
* key-value in a view whose key is outside the view's bounds, calls to `get` or `contains` will _not_ consider the
126+
* newly added entry. Mutations are always reflected in the original map, though.
127+
*
128+
* @param from the lower bound (inclusive) of this projection wrapped in a `Some`, or `None` if there is no lower
129+
* bound.
130+
* @param until the upper bound (exclusive) of this projection wrapped in a `Some`, or `None` if there is no upper
131+
* bound.
132+
*/
133+
def rangeImpl(from: Option[K], until: Option[K]): TreeMap[K, V] = new TreeMapProjection(from, until)
134+
135+
override def foreach[U](f: ((K, V)) => U): Unit = RB.foreach(tree, f)
136+
override def foreachEntry[U](f: (K, V) => U): Unit = RB.foreachEntry(tree, f)
137+
138+
override def size: Int = RB.size(tree)
139+
override def knownSize: Int = size
140+
override def isEmpty: Boolean = RB.isEmpty(tree)
141+
142+
override def contains(key: K): Boolean = RB.contains(tree, key)
143+
144+
override def head: (K, V) = RB.min(tree).get
145+
146+
override def last: (K, V) = RB.max(tree).get
147+
148+
override def minAfter(key: K): Option[(K, V)] = RB.minAfter(tree, key)
149+
150+
override def maxBefore(key: K): Option[(K, V)] = RB.maxBefore(tree, key)
151+
152+
override protected[this] def className: String = "TreeMap"
153+
154+
155+
/**
156+
* A ranged projection of a [[TreeMap]]. Mutations on this map affect the original map and vice versa.
157+
*
158+
* Only entries with keys between this projection's key range will ever appear as elements of this map, independently
159+
* of whether the entries are added through the original map or through this view. That means that if one inserts a
160+
* key-value in a view whose key is outside the view's bounds, calls to `get` or `contains` will _not_ consider the
161+
* newly added entry. Mutations are always reflected in the original map, though.
162+
*
163+
* @param from the lower bound (inclusive) of this projection wrapped in a `Some`, or `None` if there is no lower
164+
* bound.
165+
* @param until the upper bound (exclusive) of this projection wrapped in a `Some`, or `None` if there is no upper
166+
* bound.
167+
*/
168+
private[this] final class TreeMapProjection(from: Option[K], until: Option[K]) extends TreeMap[K, V](tree) {
169+
170+
/**
171+
* Given a possible new lower bound, chooses and returns the most constraining one (the maximum).
172+
*/
173+
private[this] def pickLowerBound(newFrom: Option[K]): Option[K] = (from, newFrom) match {
174+
case (Some(fr), Some(newFr)) => Some(ordering.max(fr, newFr))
175+
case (None, _) => newFrom
176+
case _ => from
177+
}
178+
179+
/**
180+
* Given a possible new upper bound, chooses and returns the most constraining one (the minimum).
181+
*/
182+
private[this] def pickUpperBound(newUntil: Option[K]): Option[K] = (until, newUntil) match {
183+
case (Some(unt), Some(newUnt)) => Some(ordering.min(unt, newUnt))
184+
case (None, _) => newUntil
185+
case _ => until
186+
}
187+
188+
/**
189+
* Returns true if the argument is inside the view bounds (between `from` and `until`).
190+
*/
191+
private[this] def isInsideViewBounds(key: K): Boolean = {
192+
val afterFrom = from.isEmpty || ordering.compare(from.get, key) <= 0
193+
val beforeUntil = until.isEmpty || ordering.compare(key, until.get) < 0
194+
afterFrom && beforeUntil
195+
}
196+
197+
override def rangeImpl(from: Option[K], until: Option[K]): TreeMap[K, V] =
198+
new TreeMapProjection(pickLowerBound(from), pickUpperBound(until))
199+
200+
override def get(key: K) = if (isInsideViewBounds(key)) RB.get(tree, key) else None
201+
202+
override def iterator = if (RB.size(tree) == 0) Iterator.empty else RB.iterator(tree, from, until)
203+
override def keysIterator: Iterator[K] = if (RB.size(tree) == 0) Iterator.empty else RB.keysIterator(tree, from, until)
204+
override def valuesIterator: Iterator[V] = if (RB.size(tree) == 0) Iterator.empty else RB.valuesIterator(tree, from, until)
205+
override def keysIteratorFrom(start: K) = if (RB.size(tree) == 0) Iterator.empty else RB.keysIterator(tree, pickLowerBound(Some(start)), until)
206+
override def iteratorFrom(start: K) = if (RB.size(tree) == 0) Iterator.empty else RB.iterator(tree, pickLowerBound(Some(start)), until)
207+
override def valuesIteratorFrom(start: K) = if (RB.size(tree) == 0) Iterator.empty else RB.valuesIterator(tree, pickLowerBound(Some(start)), until)
208+
override def size = if (RB.size(tree) == 0) 0 else iterator.length
209+
override def knownSize: Int = if (RB.size(tree) == 0) 0 else -1
210+
override def isEmpty = RB.size(tree) == 0 || !iterator.hasNext
211+
override def contains(key: K) = isInsideViewBounds(key) && RB.contains(tree, key)
212+
213+
override def head = headOption.get
214+
override def headOption = {
215+
val entry = if (from.isDefined) RB.minAfter(tree, from.get) else RB.min(tree)
216+
(entry, until) match {
217+
case (Some(e), Some(unt)) if ordering.compare(e._1, unt) >= 0 => None
218+
case _ => entry
219+
}
220+
}
221+
222+
override def last = lastOption.get
223+
override def lastOption = {
224+
val entry = if (until.isDefined) RB.maxBefore(tree, until.get) else RB.max(tree)
225+
(entry, from) match {
226+
case (Some(e), Some(fr)) if ordering.compare(e._1, fr) < 0 => None
227+
case _ => entry
228+
}
229+
}
230+
231+
// Using the iterator should be efficient enough; if performance is deemed a problem later, specialized
232+
// `foreach(f, from, until)` and `transform(f, from, until)` methods can be created in `RedBlackTree`. See
233+
// https://github.com/scala/scala/pull/4608#discussion_r34307985 for a discussion about this.
234+
override def foreach[U](f: ((K, V)) => U): Unit = iterator.foreach(f)
235+
236+
override def clone() = super.clone().rangeImpl(from, until)
237+
}
238+
239+
}
240+
241+
/**
242+
* $factoryInfo
243+
*
244+
* @define Coll mutable.TreeMap
245+
* @define coll mutable tree map
246+
*/
247+
@SerialVersionUID(3L)
248+
object TreeMap extends SortedMapFactory[TreeMap] {
249+
250+
def from[K : Ordering, V](it: IterableOnce[(K, V)]): TreeMap[K, V] =
251+
Growable.from(empty[K, V], it)
252+
253+
def empty[K : Ordering, V]: TreeMap[K, V] = new TreeMap[K, V]()
254+
255+
def newBuilder[K: Ordering, V]: Builder[(K, V), TreeMap[K, V]] = new GrowableBuilder(empty[K, V])
256+
257+
}

0 commit comments

Comments
 (0)