summaryrefslogtreecommitdiff
path: root/src/library/scala/collection/mutable/TreeMap.scala
blob: 14ae7c9c8cc8e05fe95032ec899acb2971a032f6 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
package scala
package collection
package mutable

import scala.collection.generic._
import scala.collection.mutable.{RedBlackTree => RB}

/**
 * $factoryInfo
 *
 * @define Coll mutable.TreeMap
 * @define coll mutable tree map
 */
object TreeMap extends MutableSortedMapFactory[TreeMap] {

  def empty[A, B](implicit ord: Ordering[A]) = new TreeMap[A, B]()(ord)

  /** $sortedMapCanBuildFromInfo */
  implicit def canBuildFrom[A, B](implicit ord: Ordering[A]): CanBuildFrom[Coll, (A, B), TreeMap[A, B]] =
    new SortedMapCanBuildFrom[A, B]
}

/**
 * A mutable sorted map implemented using a mutable red-black tree as underlying data structure.
 *
 * @param ordering the implicit ordering used to compare objects of type `A`.
 * @tparam A the type of the keys contained in this tree map.
 * @tparam B the type of the values associated with the keys.
 *
 * @author Rui Gonçalves
 * @version 2.12
 * @since 2.12
 *
 * @define Coll mutable.TreeMap
 * @define coll mutable tree map
 */
@SerialVersionUID(-2558985573956740112L)
sealed class TreeMap[A, B] private (tree: RB.Tree[A, B])(implicit val ordering: Ordering[A])
  extends AbstractSortedMap[A, B]
  with SortedMap[A, B]
  with MapLike[A, B, TreeMap[A, B]]
  with SortedMapLike[A, B, TreeMap[A, B]]
  with Serializable {

  /**
   * Creates an empty `TreeMap`.
   * @param ord the implicit ordering used to compare objects of type `A`.
   * @return an empty `TreeMap`.
   */
  def this()(implicit ord: Ordering[A]) = this(RB.Tree.empty)(ord)

  override def empty = TreeMap.empty
  override protected[this] def newBuilder = TreeMap.newBuilder[A, B]

  /**
   * Creates a ranged projection of this map. Any mutations in the ranged projection will update the original map and
   * vice versa.
   *
   * Only entries with keys between this projection's key range will ever appear as elements of this map, independently
   * of whether the entries are added through the original map or through this view. That means that if one inserts a
   * key-value in a view whose key is outside the view's bounds, calls to `get` or `contains` will _not_ consider the
   * newly added entry. Mutations are always reflected in the original map, though.
   *
   * @param from the lower bound (inclusive) of this projection wrapped in a `Some`, or `None` if there is no lower
   *             bound.
   * @param until the upper bound (exclusive) of this projection wrapped in a `Some`, or `None` if there is no upper
   *              bound.
   */
  def rangeImpl(from: Option[A], until: Option[A]): TreeMap[A, B] = new TreeMapView(from, until)

  def -=(key: A): this.type = { RB.delete(tree, key); this }
  def +=(kv: (A, B)): this.type = { RB.insert(tree, kv._1, kv._2); this }

  def get(key: A) = RB.get(tree, key)

  def iterator = RB.iterator(tree)
  def iteratorFrom(start: A) = RB.iterator(tree, Some(start))
  def keysIteratorFrom(start: A) = RB.keysIterator(tree, Some(start))
  def valuesIteratorFrom(start: A) = RB.valuesIterator(tree, Some(start))

  override def size = RB.size(tree)
  override def isEmpty = RB.isEmpty(tree)
  override def contains(key: A) = RB.contains(tree, key)

  override def head = RB.min(tree).get
  override def headOption = RB.min(tree)
  override def last = RB.max(tree).get
  override def lastOption = RB.max(tree)

  override def keysIterator = RB.keysIterator(tree)
  override def valuesIterator = RB.valuesIterator(tree)

  override def foreach[U](f: ((A, B)) => U): Unit = RB.foreach(tree, f)
  override def transform(f: (A, B) => B) = { RB.transform(tree, f); this }
  override def clear(): Unit = RB.clear(tree)

  override def stringPrefix = "TreeMap"

  /**
   * A ranged projection of a [[TreeMap]]. Mutations on this map affect the original map and vice versa.
   *
   * Only entries with keys between this projection's key range will ever appear as elements of this map, independently
   * of whether the entries are added through the original map or through this view. That means that if one inserts a
   * key-value in a view whose key is outside the view's bounds, calls to `get` or `contains` will _not_ consider the
   * newly added entry. Mutations are always reflected in the original map, though.
   *
   * @param from the lower bound (inclusive) of this projection wrapped in a `Some`, or `None` if there is no lower
   *             bound.
   * @param until the upper bound (exclusive) of this projection wrapped in a `Some`, or `None` if there is no upper
   *              bound.
   */
  @SerialVersionUID(2219159283273389116L)
  private[this] final class TreeMapView(from: Option[A], until: Option[A]) extends TreeMap[A, B](tree) {

    /**
     * Given a possible new lower bound, chooses and returns the most constraining one (the maximum).
     */
    private[this] def pickLowerBound(newFrom: Option[A]): Option[A] = (from, newFrom) match {
      case (Some(fr), Some(newFr)) => Some(ordering.max(fr, newFr))
      case (None, _) => newFrom
      case _ => from
    }

    /**
     * Given a possible new upper bound, chooses and returns the most constraining one (the minimum).
     */
    private[this] def pickUpperBound(newUntil: Option[A]): Option[A] = (until, newUntil) match {
      case (Some(unt), Some(newUnt)) => Some(ordering.min(unt, newUnt))
      case (None, _) => newUntil
      case _ => until
    }

    /**
     * Returns true if the argument is inside the view bounds (between `from` and `until`).
     */
    private[this] def isInsideViewBounds(key: A): Boolean = {
      val afterFrom = from.isEmpty || ordering.compare(from.get, key) <= 0
      val beforeUntil = until.isEmpty || ordering.compare(key, until.get) < 0
      afterFrom && beforeUntil
    }

    override def rangeImpl(from: Option[A], until: Option[A]): TreeMap[A, B] =
      new TreeMapView(pickLowerBound(from), pickUpperBound(until))

    override def get(key: A) = if (isInsideViewBounds(key)) RB.get(tree, key) else None

    override def iterator = RB.iterator(tree, from, until)
    override def iteratorFrom(start: A) = RB.iterator(tree, pickLowerBound(Some(start)), until)
    override def keysIteratorFrom(start: A) = RB.keysIterator(tree, pickLowerBound(Some(start)), until)
    override def valuesIteratorFrom(start: A) = RB.valuesIterator(tree, pickLowerBound(Some(start)), until)

    override def size = iterator.length
    override def isEmpty = !iterator.hasNext
    override def contains(key: A) = isInsideViewBounds(key) && RB.contains(tree, key)

    override def head = headOption.get
    override def headOption = {
      val entry = if (from.isDefined) RB.minAfter(tree, from.get) else RB.min(tree)
      (entry, until) match {
        case (Some(e), Some(unt)) if ordering.compare(e._1, unt) >= 0 => None
        case _ => entry
      }
    }

    override def last = lastOption.get
    override def lastOption = {
      val entry = if (until.isDefined) RB.maxBefore(tree, until.get) else RB.max(tree)
      (entry, from) match {
        case (Some(e), Some(fr)) if ordering.compare(e._1, fr) < 0 => None
        case _ => entry
      }
    }

    // Using the iterator should be efficient enough; if performance is deemed a problem later, specialized
    // `foreach(f, from, until)` and `transform(f, from, until)` methods can be created in `RedBlackTree`. See
    // https://github.com/scala/scala/pull/4608#discussion_r34307985 for a discussion about this.
    override def foreach[U](f: ((A, B)) => U): Unit = iterator.foreach(f)
    override def transform(f: (A, B) => B) = {
      iterator.foreach { case (key, value) => update(key, f(key, value)) }
      this
    }

    override def valuesIterator: Iterator[B] = RB.valuesIterator(tree, from, until)
    override def keysIterator: Iterator[A] = RB.keysIterator(tree, from, until)

    override def clone() = super.clone().rangeImpl(from, until)
  }
}