Merge pull request scala#4717 from Ichoran/issue/9347

SI-9347  Efficient head/tail, if possible, for immutable maps & sets

src/library/scala/collection/BitSetLike.scala

@@ -204,6 +204,27 @@ trait BitSetLike[+This <: BitSetLike[This] with SortedSet[Int]] extends SortedSe
   def subsetOf(other: BitSet): Boolean =
     (0 until nwords) forall (idx => (this.word(idx) & ~ other.word(idx)) == 0L)
 
+  override def head: Int = {
+    val n = nwords
+    var i = 0
+    while (i < n) {
+      val wi = word(i)
+      if (wi != 0L) return WordLength*i + java.lang.Long.numberOfTrailingZeros(wi)
+      i += 1
+    }
+    throw new NoSuchElementException("Empty BitSet")
+  }
+
+  override def last: Int = {
+    var i = nwords - 1
+    while (i >= 0) {
+      val wi = word(i)
+      if (wi != 0L) return WordLength*i + 63 - java.lang.Long.numberOfLeadingZeros(wi)
+      i += 1
+    }
+    throw new NoSuchElementException("Empty BitSet")
+  }
+
   override def addString(sb: StringBuilder, start: String, sep: String, end: String) = {
     sb append start
     var pre = ""

src/library/scala/collection/immutable/BitSet.scala

@@ -103,13 +103,20 @@ object BitSet extends BitSetFactory[BitSet] {
     else new BitSetN(elems)
   }
 
+  @SerialVersionUID(2260107458435649300L)
   class BitSet1(val elems: Long) extends BitSet {
     protected def nwords = 1
     protected def word(idx: Int) = if (idx == 0) elems else 0L
     protected def updateWord(idx: Int, w: Long): BitSet =
       if (idx == 0) new BitSet1(w)
       else if (idx == 1) new BitSet2(elems, w)
       else fromBitMaskNoCopy(updateArray(Array(elems), idx, w))
+    override def head: Int =
+      if (elems == 0L) throw new NoSuchElementException("Empty BitSet")
+      else java.lang.Long.numberOfTrailingZeros(elems)
+    override def tail: BitSet =
+      if (elems == 0L) throw new NoSuchElementException("Empty BitSet")
+      else new BitSet1(elems - java.lang.Long.lowestOneBit(elems))
   }
 
   class BitSet2(val elems0: Long, elems1: Long) extends BitSet {
@@ -119,6 +126,18 @@ object BitSet extends BitSetFactory[BitSet] {
       if (idx == 0) new BitSet2(w, elems1)
       else if (idx == 1) new BitSet2(elems0, w)
       else fromBitMaskNoCopy(updateArray(Array(elems0, elems1), idx, w))
+    override def head: Int =
+      if (elems0 == 0L) {
+        if (elems1 == 0) throw new NoSuchElementException("Empty BitSet")
+        64 + java.lang.Long.numberOfTrailingZeros(elems1)
+      }
+      else java.lang.Long.numberOfTrailingZeros(elems0)
+    override def tail: BitSet =
+      if (elems0 == 0L) {
+        if (elems1 == 0L) throw new NoSuchElementException("Empty BitSet")
+        new BitSet2(elems0, elems1 - java.lang.Long.lowestOneBit(elems1))
+      }
+      else new BitSet2(elems0 - java.lang.Long.lowestOneBit(elems0), elems1)
   }
 
   /** The implementing class for bit sets with elements >= 128 (exceeding
@@ -131,5 +150,15 @@ object BitSet extends BitSetFactory[BitSet] {
     protected def nwords = elems.length
     protected def word(idx: Int) = if (idx < nwords) elems(idx) else 0L
     protected def updateWord(idx: Int, w: Long): BitSet = fromBitMaskNoCopy(updateArray(elems, idx, w))
+    override def tail: BitSet = {
+      val n = nwords
+      var i = 0
+      while (i < n) {
+        val wi = word(i)
+        if (wi != 0L) return fromBitMaskNoCopy(updateArray(elems, i, wi - java.lang.Long.lowestOneBit(wi)))
+        i += 1
+      }
+      throw new NoSuchElementException("Empty BitSet")
+    }
   }
 }

src/library/scala/collection/immutable/HashMap.scala

@@ -65,6 +65,8 @@ class HashMap[A, +B] extends AbstractMap[A, B]
   def - (key: A): HashMap[A, B] =
     removed0(key, computeHash(key), 0)
 
+  override def tail: HashMap[A, B] = this - head._1
+
   override def filter(p: ((A, B)) => Boolean) = {
     val buffer = new Array[HashMap[A, B]](bufferSize(size))
     nullToEmpty(filter0(p, false, 0, buffer, 0))
@@ -156,7 +158,10 @@ object HashMap extends ImmutableMapFactory[HashMap] with BitOperations.Int {
   implicit def canBuildFrom[A, B]: CanBuildFrom[Coll, (A, B), HashMap[A, B]] = new MapCanBuildFrom[A, B]
   def empty[A, B]: HashMap[A, B] = EmptyHashMap.asInstanceOf[HashMap[A, B]]
 
-  private object EmptyHashMap extends HashMap[Any, Nothing] { }
+  private object EmptyHashMap extends HashMap[Any, Nothing] { 
+    override def head: (Any, Nothing) = throw new NoSuchElementException("Empty Map")
+    override def tail: HashMap[Any, Nothing] = throw new NoSuchElementException("Empty Map")
+  }
 
   // utility method to create a HashTrieMap from two leaf HashMaps (HashMap1 or HashMapCollision1) with non-colliding hash code)
   private def makeHashTrieMap[A, B](hash0:Int, elem0:HashMap[A, B], hash1:Int, elem1:HashMap[A, B], level:Int, size:Int) : HashTrieMap[A, B] = {

src/library/scala/collection/immutable/HashSet.scala

@@ -162,6 +162,8 @@ class HashSet[A] extends AbstractSet[A]
   def - (e: A): HashSet[A] =
     nullToEmpty(removed0(e, computeHash(e), 0))
 
+  override def tail: HashSet[A] = this - head
+
   override def filter(p: A => Boolean) = {
     val buffer = new Array[HashSet[A]](bufferSize(size))
     nullToEmpty(filter0(p, false, 0, buffer, 0))
@@ -213,7 +215,10 @@ object HashSet extends ImmutableSetFactory[HashSet] {
   /** $setCanBuildFromInfo */
   implicit def canBuildFrom[A]: CanBuildFrom[Coll, A, HashSet[A]] = setCanBuildFrom[A]
 
-  private object EmptyHashSet extends HashSet[Any] { }
+  private object EmptyHashSet extends HashSet[Any] {
+    override def head: Any = throw new NoSuchElementException("Empty Set")
+    override def tail: HashSet[Any] = throw new NoSuchElementException("Empty Set")
+  }
   private[collection] def emptyInstance: HashSet[Any] = EmptyHashSet
 
   // utility method to create a HashTrieSet from two leaf HashSets (HashSet1 or HashSetCollision1) with non-colliding hash code)

src/library/scala/collection/immutable/ListMap.scala

@@ -20,6 +20,9 @@ import scala.annotation.tailrec
  *  @see [[http://docs.scala-lang.org/overviews/collections/concrete-immutable-collection-classes.html#list_maps "Scala's Collection Library overview"]]
  *  section on `List Maps` for more information.
  *
+ *  Note that `ListMap` is built in reverse order to canonical traversal order (traversal order is oldest first).
+ *  Thus, `head` and `tail` are O(n).  To rapidly partition a `ListMap` into elements, use `last` and `init` instead.  These are O(1).
+ *
  *  @define Coll immutable.ListMap
  *  @define coll immutable list map
  */
@@ -33,6 +36,8 @@ object ListMap extends ImmutableMapFactory[ListMap] {
   private object EmptyListMap extends ListMap[Any, Nothing] {
     override def apply(key: Any) = throw new NoSuchElementException("key not found: " + key)
     override def contains(key: Any) = false
+    override def last: (Any, Nothing) = throw new NoSuchElementException("Empty ListMap")
+    override def init: ListMap[Any, Nothing] = throw new NoSuchElementException("Empty ListMap")
   }
 }
 
@@ -216,5 +221,9 @@ extends AbstractMap[A, B]
         remove0(k, cur.next, cur::acc)
 
     override protected def next: ListMap[A, B1] = ListMap.this
+
+    override def last: (A, B1) = (key, value)
+
+    override def init: ListMap[A, B1] = next
   }
 }

src/library/scala/collection/immutable/Set.scala

@@ -103,10 +103,11 @@ object Set extends ImmutableSetFactory[Set] {
       if (p(elem1)) Some(elem1)
       else None
     }
+    override def head: A = elem1
+    override def tail: Set[A] = Set.empty
     // Why is Set1 non-final?  Need to fix that!
     @deprecatedOverriding("This immutable set should do nothing on toSet but cast itself to a Set with a wider element type.", "2.11.8")
     override def toSet[B >: A]: Set[B] = this.asInstanceOf[Set1[B]]
-
   }
 
   /** An optimized representation for immutable sets of size 2 */
@@ -138,6 +139,8 @@ object Set extends ImmutableSetFactory[Set] {
       else if (p(elem2)) Some(elem2)
       else None
     }
+    override def head: A = elem1
+    override def tail: Set[A] = new Set1(elem2)
     // Why is Set2 non-final?  Need to fix that!
     @deprecatedOverriding("This immutable set should do nothing on toSet but cast itself to a Set with a wider element type.", "2.11.8")
     override def toSet[B >: A]: Set[B] = this.asInstanceOf[Set2[B]]
@@ -174,6 +177,8 @@ object Set extends ImmutableSetFactory[Set] {
       else if (p(elem3)) Some(elem3)
       else None
     }
+    override def head: A = elem1
+    override def tail: Set[A] = new Set2(elem2, elem3)
     // Why is Set3 non-final?  Need to fix that!
     @deprecatedOverriding("This immutable set should do nothing on toSet but cast itself to a Set with a wider element type.", "2.11.8")
     override def toSet[B >: A]: Set[B] = this.asInstanceOf[Set3[B]]
@@ -212,6 +217,8 @@ object Set extends ImmutableSetFactory[Set] {
       else if (p(elem4)) Some(elem4)
       else None
     }
+    override def head: A = elem1
+    override def tail: Set[A] = new Set3(elem2, elem3, elem4)
     // Why is Set4 non-final?  Need to fix that!
     @deprecatedOverriding("This immutable set should do nothing on toSet but cast itself to a Set with a wider element type.", "2.11.8")
     override def toSet[B >: A]: Set[B] = this.asInstanceOf[Set4[B]]