summaryrefslogtreecommitdiff
path: root/src/library/scala/collection/mutable/TreeSet.scala
blob: ada6f145ad42adb463cc1917857c543c1a5bc1bc (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
189
/*                     __                                               *\
**     ________ ___   / /  ___     Scala API                            **
**    / __/ __// _ | / /  / _ |    (c) 2003-2013, LAMP/EPFL             **
**  __\ \/ /__/ __ |/ /__/ __ |    http://scala-lang.org/               **
** /____/\___/_/ |_/____/_/ | |                                         **
**                          |/                                          **
\*                                                                      */

package scala
package collection
package mutable

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

/**
 * @define Coll `mutable.TreeSet`
 * @define coll mutable tree set
 * @factoryInfo
 *   Companion object of TreeSet providing factory related utilities.
 *
 * @author Lucien Pereira
 *
 */
object TreeSet extends MutableSortedSetFactory[TreeSet] {
  /**
   *  The empty set of this type
   */
  def empty[A](implicit ordering: Ordering[A]) = new TreeSet[A]()

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

/**
 * A mutable sorted set 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 set.
 *
 * @author Rui Gonçalves
 * @version 2.12
 * @since 2.10
 *
 * @define Coll mutable.TreeSet
 * @define coll mutable tree set
 */
// Original API designed in part by Lucien Pereira
@SerialVersionUID(-3642111301929493640L)
sealed class TreeSet[A] private (tree: RB.Tree[A, Null])(implicit val ordering: Ordering[A])
  extends AbstractSortedSet[A]
  with SortedSet[A]
  with SetLike[A, TreeSet[A]]
  with SortedSetLike[A, TreeSet[A]]
  with Serializable {

  if (ordering eq null)
    throw new NullPointerException("ordering must not be null")

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

  override def empty = TreeSet.empty
  override protected[this] def newBuilder = TreeSet.newBuilder[A]

  /**
   * Creates a ranged projection of this set. Any mutations in the ranged projection affect will update the original set
   * and vice versa.
   *
   * Only keys between this projection's key range will ever appear as elements of this set, independently of whether
   * the elements are added through the original set or through this view. That means that if one inserts an element in
   * a view whose key is outside the view's bounds, calls to `contains` will _not_ consider the newly added element.
   * Mutations are always reflected in the original set, 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]): TreeSet[A] = new TreeSetView(from, until)

  def -=(key: A): this.type = { RB.delete(tree, key); this }
  def +=(elem: A): this.type = { RB.insert(tree, elem, null); this }

  def contains(elem: A) = RB.contains(tree, elem)

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

  override def size = RB.size(tree)
  override def isEmpty = RB.isEmpty(tree)

  override def head = RB.minKey(tree).get
  override def headOption = RB.minKey(tree)
  override def last = RB.maxKey(tree).get
  override def lastOption = RB.maxKey(tree)

  override def foreach[U](f: A => U): Unit = RB.foreachKey(tree, f)
  override def clear(): Unit = RB.clear(tree)

  override def stringPrefix = "TreeSet"

  /**
   * A ranged projection of a [[TreeSet]]. Mutations on this set affect the original set and vice versa.
   *
   * Only keys between this projection's key range will ever appear as elements of this set, independently of whether
   * the elements are added through the original set or through this view. That means that if one inserts an element in
   * a view whose key is outside the view's bounds, calls to `contains` will _not_ consider the newly added element.
   * Mutations are always reflected in the original set, 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(7087824939194006086L)
  private[this] final class TreeSetView(from: Option[A], until: Option[A]) extends TreeSet[A](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]): TreeSet[A] =
      new TreeSetView(pickLowerBound(from), pickUpperBound(until))

    override def contains(key: A) = isInsideViewBounds(key) && RB.contains(tree, key)

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

    override def size = iterator.length
    override def isEmpty = !iterator.hasNext

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

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

    // Using the iterator should be efficient enough; if performance is deemed a problem later, a specialized
    // `foreachKey(f, from, until)` method 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 => U): Unit = iterator.foreach(f)

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