Merge pull request #903 from S714726/SI-5906

SI-5906 Search for sorted sequences

4 files changed, 193 insertions, 0 deletions

diff --git a/src/library/scala/collection/Searching.scala b/src/library/scala/collection/Searching.scala
new file mode 100644
index 0000000000..c1f7f4cae6
--- /dev/null
+++ b/src/library/scala/collection/Searching.scala
@@ -0,0 +1,116 @@
+/*                     __                                               *\
+**     ________ ___   / /  ___     Scala API                            **
+**    / __/ __// _ | / /  / _ |    (c) 2003-2012, LAMP/EPFL             **
+**  __\ \/ /__/ __ |/ /__/ __ |    http://scala-lang.org/               **
+** /____/\___/_/ |_/____/_/ | |                                         **
+**                          |/                                          **
+\*                                                                      */
+
+package scala.collection
+
+import scala.annotation.tailrec
+import scala.collection.generic.IsSeqLike
+import scala.math.Ordering
+
+/** A collection of wrappers that provide sequence classes with search functionality.
+  *
+  * Example usage:
+  * {{{
+  *    import scala.collection.Searching._
+  *    val l = List(1, 2, 3, 4, 5)
+  *    l.search(3)
+  *    // == Found(2)
+  * }}}
+  */
+object Searching {
+  sealed abstract class SearchResult {
+    def insertionPoint: Int
+  }
+
+  case class Found(foundIndex: Int) extends SearchResult {
+    override def insertionPoint = foundIndex
+  }
+  case class InsertionPoint(insertionPoint: Int) extends SearchResult
+
+  class SearchImpl[A, Repr](val coll: SeqLike[A, Repr]) {
+    /** Search the sorted sequence for a specific element. If the sequence is an
+      * `IndexedSeq`, a binary search is used. Otherwise, a linear search is used.
+      *
+      * The sequence should be sorted with the same `Ordering` before calling; otherwise,
+      * the results are undefined.
+      *
+      * @see [[scala.math.IndexedSeq]]
+      * @see [[scala.math.Ordering]]
+      * @see [[scala.collection.SeqLike]], method `sorted`
+      *
+      * @param elem the element to find.
+      * @param ord  the ordering to be used to compare elements.
+      *
+      * @return a `Found` value containing the index corresponding to the element in the
+      *         sequence, or the `InsertionPoint` where the element would be inserted if
+      *         the element is not in the sequence.
+      */
+    final def search[B >: A](elem: B)(implicit ord: Ordering[B]): SearchResult = 
+      coll match {
+        case _: IndexedSeq[A] => binarySearch(elem, -1, coll.length)(ord)
+        case _ => linearSearch(coll.view, elem, 0)(ord)
+      }
+
+    /** Search within an interval in the sorted sequence for a specific element. If the
+      * sequence is an IndexedSeq, a binary search is used. Otherwise, a linear search
+      * is used.
+      *
+      * The sequence should be sorted with the same `Ordering` before calling; otherwise,
+      * the results are undefined.
+      *
+      * @see [[scala.math.IndexedSeq]]
+      * @see [[scala.math.Ordering]]
+      * @see [[scala.collection.SeqLike]], method `sorted`
+      *
+      * @param elem the element to find.
+      * @param from the index where the search starts.
+      * @param to   the index following where the search ends.
+      * @param ord  the ordering to be used to compare elements.
+      *
+      * @return a `Found` value containing the index corresponding to the element in the
+      *         sequence, or the `InsertionPoint` where the element would be inserted if
+      *         the element is not in the sequence.
+      */
+    final def search[B >: A](elem: B, from: Int, to: Int)
+    (implicit ord: Ordering[B]): SearchResult = 
+      coll match {
+        case _: IndexedSeq[A] => binarySearch(elem, from-1, to)(ord)
+        case _ => linearSearch(coll.view(from, to), elem, from)(ord)
+      }
+
+    @tailrec
+    private def binarySearch[B >: A](elem: B, from: Int, to: Int)
+    (implicit ord: Ordering[B]): SearchResult = {
+      if ((to-from) == 1) InsertionPoint(from) else {
+        val idx = from+(to-from)/2
+        math.signum(ord.compare(elem, coll(idx))) match {
+          case -1 => binarySearch(elem, from, idx)(ord)
+          case  1 => binarySearch(elem, idx, to)(ord)
+          case  _ => Found(idx)
+        }
+      }
+    }
+
+    private def linearSearch[B >: A](c: SeqView[A, Repr], elem: B, offset: Int)
+    (implicit ord: Ordering[B]): SearchResult = {
+      var idx = offset
+      val it = c.iterator
+      while (it.hasNext) {
+        val cur = it.next()
+        if (ord.equiv(elem, cur)) return Found(idx)
+        else if (ord.lt(elem, cur)) return InsertionPoint(idx-1)
+        idx += 1
+      }
+      InsertionPoint(idx)
+    }
+
+  }
+
+  implicit def search[Repr, A](coll: Repr)
+  (implicit fr: IsSeqLike[Repr]): SearchImpl[fr.A, Repr] = new SearchImpl(fr.conversion(coll))
+}
diff --git a/src/library/scala/collection/generic/IsSeqLike.scala b/src/library/scala/collection/generic/IsSeqLike.scala
new file mode 100644
index 0000000000..8eac025ed6
--- /dev/null
+++ b/src/library/scala/collection/generic/IsSeqLike.scala
@@ -0,0 +1,57 @@
+/*                     __                                               *\
+**     ________ ___   / /  ___     Scala API                            **
+**    / __/ __// _ | / /  / _ |    (c) 2003-2012, LAMP/EPFL             **
+**  __\ \/ /__/ __ |/ /__/ __ |    http://scala-lang.org/               **
+** /____/\___/_/ |_/____/_/ | |                                         **
+**                          |/                                          **
+\*                                                                      */
+
+package scala.collection
+package generic
+
+/** Type class witnessing that a collection representation type `Repr` has
+  * elements of type `A` and has a conversion to `SeqLike[A, Repr]`.
+  *
+  * This type enables simple enrichment of `Seq`s with extension methods which
+  * can make full use of the mechanics of the Scala collections framework in
+  * their implementation.
+  *
+  * Example usage:
+  * {{{
+  *    class FilterMapImpl[A, Repr](val r: SeqLike[A, Repr]) {
+  *      final def filterMap[B, That](f: A => Option[B])(implicit cbf: CanBuildFrom[Repr, B, That]): That =
+  *        r.flatMap(f(_))
+  *    }
+  *    implicit def filterMap[Repr, A](r: Repr)(implicit fr: IsSeqLike[Repr]): FilterMapImpl[fr.A,Repr] =
+  *      new FilterMapImpl(fr.conversion(r))
+  *
+  *    val l = List(1, 2, 3, 4, 5)
+  *    List(1, 2, 3, 4, 5) filterMap (i => if(i % 2 == 0) Some(i) else None)
+  *    // == List(2, 4)
+  * }}}
+  *
+  * @see [[scala.collection.generic.Seq]]
+  * @see [[scala.collection.generic.IsTraversableLike]]
+  */
+trait IsSeqLike[Repr] {
+  /** The type of elements we can traverse over. */
+  type A
+  /** A conversion from the representation type `Repr` to a `SeqLike[A,Repr]`. */
+  val conversion: Repr => SeqLike[A, Repr]
+}
+
+object IsSeqLike {
+  import language.higherKinds
+
+  implicit val stringRepr: IsSeqLike[String] { type A = Char } =
+    new IsSeqLike[String] {
+      type A = Char
+      val conversion = implicitly[String => SeqLike[Char, String]]
+    }
+
+  implicit def seqLikeRepr[C[_], A0](implicit conv: C[A0] => SeqLike[A0,C[A0]]): IsSeqLike[C[A0]] { type A = A0 } = 
+    new IsSeqLike[C[A0]] {
+      type A = A0
+      val conversion = conv
+    }
+}
diff --git a/test/files/run/search.check b/test/files/run/search.check
new file mode 100644
index 0000000000..a885696509
--- /dev/null
+++ b/test/files/run/search.check
@@ -0,0 +1,6 @@
+Found(2)
+Found(4)
+InsertionPoint(9)
+Found(2)
+Found(4)
+InsertionPoint(9)
diff --git a/test/files/run/search.scala b/test/files/run/search.scala
new file mode 100644
index 0000000000..ed7fed54a7
--- /dev/null
+++ b/test/files/run/search.scala
@@ -0,0 +1,14 @@
+object Test extends App {
+  import scala.collection.{LinearSeq, IndexedSeq}
+  import scala.collection.Searching.search
+
+  val ls = LinearSeq(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 13)
+  println(ls.search(3))
+  println(ls.search(5, 3, 8))
+  println(ls.search(12))
+
+  val is = IndexedSeq(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 13)
+  println(is.search(3))
+  println(is.search(5, 3, 8))
+  println(is.search(12))
+}