8000 Custom coded version of range/from/to/until. · scala/scala@7824dbd · GitHub
[go: up one dir, main page]
Skip to content

Commit 7824dbd

Browse files
erikrozendaalerikrozendaal
committed
Custom coded version of range/from/to/until.
This avoids unnecessary allocation of Option and Function objects, mostly helping performance of small trees.
1 parent 00b5cb8 commit 7824dbd

File tree

4 files changed

+59
-42
lines changed

4 files changed

+59
-42
lines changed

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

Lines changed: 38 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -45,11 +45,16 @@ object RedBlackTree {
4545
def count(tree: Tree[_, _]) = if (tree eq null) 0 else tree.count
4646
def update[A, B, B1 >: B](tree: Tree[A, B], k: A, v: B1)(implicit ordering: Ordering[A]): Tree[A, B1] = blacken(upd(tree, k, v))
4747
def delete[A, B](tree: Tree[A, B], k: A)(implicit ordering: Ordering[A]): Tree[A, B] = blacken(del(tree, k))
48-
def range[A, B](tree: Tree[A, B], low: Option[A], lowInclusive: Boolean, high: Option[A], highInclusive: Boolean)(implicit ordering: Ordering[A]): Tree[A, B] = {
49-
val after: Option[A => Boolean] = low.map(key => if (lowInclusive) ordering.lt(_, key) else ordering.lteq(_, key))
50-
val before: Option[A => Boolean] = high.map(key => if (highInclusive) ordering.lt(key, _) else ordering.lteq(key, _))
51-
blacken(rng(tree, after, before))
48+
def rangeImpl[A: Ordering, B](tree: Tree[A, B], from: Option[A], until: Option[A]): Tree[A, B] = (from, until) match {
49+
case (Some(from), Some(until)) => this.range(tree, from, until)
50+
case (Some(from), None) => this.from(tree, from)
51+
case (None, Some(until)) => this.until(tree, until)
52+
case (None, None) => tree
5253
}
54+
def range[A: Ordering, B](tree: Tree[A, B], from: A, until: A): Tree[A, B] = blacken(doRange(tree, from, until))
55+
def from[A: Ordering, B](tree: Tree[A, B], from: A): Tree[A, B] = blacken(doFrom(tree, from))
56+
def to[A: Ordering, B](tree: Tree[A, B], to: A): Tree[A, B] = blacken(doTo(tree, to))
57+
def until[A: Ordering, B](tree: Tree[A, B], key: A): Tree[A, B] = blacken(doUntil(tree, key))
5358

5459
def smallest[A, B](tree: Tree[A, B]): Tree[A, B] = {
5560
if (tree eq null) throw new NoSuchElementException("empty map")
@@ -202,13 +207,36 @@ object RedBlackTree {
202207
else append(tree.left, tree.right)
203208
}
204209

205-
private[this] def rng[A, B](tree: Tree[A, B], after: Option[A => Boolean], before: Option[A => Boolean])(implicit ordering: Ordering[A]): Tree[A, B] = {
210+
private[this] def doFrom[A, B](tree: Tree[A, B], from: A)(implicit ordering: Ordering[A]): Tree[A, B] = {
206211
if (tree eq null) return null
207-
if (after == None && before == None) return tree
208-
if (after != None && after.get(tree.key)) return rng(tree.right, after, before);
209-
if (before != None && before.get(tree.key)) return rng(tree.left, after, before);
210-
val newLeft = rng(tree.left, after, None)
211-
val newRight = rng(tree.right, None, before)
212+
if (ordering.lt(tree.key, from)) return doFrom(tree.right, from)
213+
val newLeft = doFrom(tree.left, from)
214+
if (newLeft eq tree.left) tree
215+
else if (newLeft eq null) upd(tree.right, tree.key, tree.value)
216+
else rebalance(tree, newLeft, tree.right)
217+
}
218+
private[this] def doTo[A, B](tree: Tree[A, B], to: A)(implicit ordering: Ordering[A]): Tree[A, B] = {
219+
if (tree eq null) return null
220+
if (ordering.lt(to, tree.key)) return doTo(tree.left, to)
221+
val newRight = doTo(tree.right, to)
222+
if (newRight eq tree.right) tree
223+
else if (newRight eq null) upd(tree.left, tree.key, tree.value)
224+
else rebalance(tree, tree.left, newRight)
225+
}
226+
private[this] def doUntil[A, B](tree: Tree[A, B], until: A)(implicit ordering: Ordering[A]): Tree[A, B] = {
227+
if (tree eq null) return null
228+
if (ordering.lteq(until, tree.key)) return doUntil(tree.left, until)
229+
val newRight = doUntil(tree.right, until)
230+
if (newRight eq tree.right) tree
231+
else if (newRight eq null) upd(tree.left, tree.key, tree.value)
232+
else rebalance(tree, tree.left, newRight)
233+
}
234+
private[this] def doRange[A, B](tree: Tree[A, B], from: A, until: A)(implicit ordering: Ordering[A]): Tree[A, B] = {
235+
if (tree eq null) return null
236+
if (ordering.lt(tree.key, from)) return doRange(tree.right, from, until);
237+
if (ordering.lteq(until, tree.key)) return doRange(tree.left, from, until);
238+
val newLeft = doFrom(tree.left, from)
239+
val newRight = doUntil(tree.right, until)
212240
if ((newLeft eq tree.left) && (newRight eq tree.right)) tree
213241
else if (newLeft eq null) upd(newRight, tree.key, tree.value);
214242
else if (newRight eq null) upd(newLeft, tree.key, tree.value);

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

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -62,14 +62,11 @@ class TreeMap[A, +B] private (tree: RB.Tree[A, B])(implicit val ordering: Orderi
6262

6363
def this()(implicit ordering: Ordering[A]) = this(null)(ordering)
6464

65-
override def rangeImpl(from: Option[A], until: Option[A]): TreeMap[A, B] = {
66-
val ntree = RB.range(tree, from, true, until, false)
67-
new TreeMap[A, B](ntree)
68-
}
69-
override def to(to: A): TreeMap[A, B] = {
70-
val ntree = RB.range(tree, None, true, Some(to), true)
71-
new TreeMap[A, B](ntree)
72-
}
65+
override def rangeImpl(from: Option[A], until: Option[A]): TreeMap[A, B] = new TreeMap[A, B](RB.rangeImpl(tree, from, until))
66+
override def range(from: A, until: A): TreeMap[A, B] = new TreeMap[A, B](RB.range(tree, from, until))
67+
override def from(from: A): TreeMap[A, B] = new TreeMap[A, B](RB.from(tree, from))
68+
override def to(to: A): TreeMap[A, B] = new TreeMap[A, B](RB.to(tree, to))
69+
override def until(until: A): TreeMap[A, B] = new TreeMap[A, B](RB.until(tree, until))
7370

7471
override def firstKey = RB.smallest(tree).key
7572
override def lastKey = RB.greatest(tree).key

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

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -150,14 +150,12 @@ class TreeSet[A] private (tree: RB.Tree[A, Unit])(implicit val ordering: Orderin
150150

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

153-
override def rangeImpl(from: Option[A], until: Option[A]): TreeSet[A] = {
154-
val ntree = RB.range(tree, from, true, until, false)
155-
newSet(ntree)
156-
}
157-
override def to(to: A): TreeSet[A] = {
158-
val ntree = RB.range(tree, None, true, Some(to), true)
159-
newSet(ntree)
160-
}
153+
override def rangeImpl(from: Option[A], until: Option[A]): TreeSet[A] = newSet(RB.rangeImpl(tree, from, until))
154+
override def range(from: A, until: A): TreeSet[A] = newSet(RB.range(tree, from, until))
155+
override def from(from: A): TreeSet[A] = newSet(RB.from(tree, from))
156+
override def to(to: A): TreeSet[A] = newSet(RB.to(tree, to))
157+
override def until(until: A): TreeSet[A] = newSet(RB.until(tree, until))
158+
161159
override def firstKey = head
162160
override def lastKey = last
163161
}

‎test/files/scalacheck/redblacktree.scala

Lines changed: 10 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -174,39 +174,33 @@ package scala.collection.immutable.redblacktree {
174174
object TestRange extends RedBlackTreeTest with RedBlackTreeInvariants {
175175
import RB._
176176

177-
override type ModifyParm = (Option[Int], Boolean, Option[Int], Boolean)
177+
override type ModifyParm = (Option[Int], Option[Int])
178178
override def genParm(tree: Tree[String, Int]): Gen[ModifyParm] = for {
179179
from <- choose(0, iterator(tree).size)
180-
fromInclusive <- oneOf(false, true)
181180
to <- choose(0, iterator(tree).size) suchThat (from <=)
182-
toInclusive <- oneOf(false, true)
183181
optionalFrom <- oneOf(Some(from), None, Some(from)) // Double Some(n) to get around a bug
184182
optionalTo <- oneOf(Some(to), None, Some(to)) // Double Some(n) to get around a bug
185-
} yield (optionalFrom, fromInclusive, optionalTo, toInclusive)
183+
} yield (optionalFrom, optionalTo)
186184

187185
override def modify(tree: Tree[String, Int], parm: ModifyParm): Tree[String, Int] = {
188186
val from = parm._1 flatMap (nodeAt(tree, _) map (_._1))
189-
val to = parm._3 flatMap (nodeAt(tree, _) map (_._1))
190-
range(tree, from, parm._2, to, parm._4)
187+
val to = parm._2 flatMap (nodeAt(tree, _) map (_._1))
188+
rangeImpl(tree, from, to)
191189
}
192190

193191
property("range boundaries respected") = forAll(genInput) { case (tree, parm, newTree) =>
194192
val from = parm._1 flatMap (nodeAt(tree, _) map (_._1))
195-
val fromPredicate: String => String => Boolean = if (parm._2) (_ <=) else (_ <)
196-
val to = parm._3 flatMap (nodeAt(tree, _) map (_._1))
197-
val toPredicate: String => String => Boolean = if (parm._4) (_ >=) else (_ >)
198-
("lower boundary" |: (from forall ( key => keysIterator(newTree) forall fromPredicate(key)))) &&
199-
("upper boundary" |: (to forall ( key => keysIterator(newTree) forall toPredicate(key))))
193+
val to = parm._2 flatMap (nodeAt(tree, _) map (_._1))
194+
("lower boundary" |: (from forall ( key => keysIterator(newTree) forall (key <=)))) &&
195+
("upper boundary" |: (to forall ( key => keysIterator(newTree) forall (key >))))
200196
}
201197

202198
property("range returns all elements") = forAll(genInput) { case (tree, parm, newTree) =>
203199
val from = parm._1 flatMap (nodeAt(tree, _) map (_._1))
204-
val fromPredicate: String => String => Boolean = if (parm._2) (_ >=) else (_ >)
205-
val to = parm._3 flatMap (nodeAt(tree, _) map (_._1))
206-
val toPredicate: String => String => Boolean = if (parm._4) (_ <=) else (_ <)
200+
val to = parm._2 flatMap (nodeAt(tree, _) map (_._1))
207201
val filteredTree = (keysIterator(tree)
208-
.filter(key => from forall fromPredicate(key))
209-
.filter(key => to forall toPredicate(key))
202+
.filter(key => from forall (key >=))
203+
.filter(key => to forall (key <))
210204
.toList)
211205
filteredTree == keysIterator(newTree).toList
212206
}

0 commit comments

Comments
 (0)
0