Merge pull request scala#7193 from odd/vectormap-tombstones

Add tombstones to VectorMap.

Author: szeiger

src/library/scala/collection/immutable/VectorMap.scala

@@ -14,8 +14,7 @@ package scala
 package collection
 package immutable
 
-import scala.collection.mutable.{Builder, ImmutableBuilder}
-import scala.annotation.unchecked.{uncheckedVariance => uV}
+import scala.annotation.tailrec
 
 /** This class implements immutable maps using a vector/map-based data structure, which preserves insertion order.
   *
@@ -26,140 +25,194 @@ import scala.annotation.unchecked.{uncheckedVariance => uV}
   *  @tparam V      the type of the values associated with the keys in this vector map.
   *
   * @author Matthew de Detrich
+  * @author Odd Möller
   * @version 2.13
   * @since 2.13
   * @define coll immutable vector map
   * @define Coll `immutable.VectorMap`
   */
-final class VectorMap[K, +V] private[immutable] (
-    private[immutable] val fields: Vector[K],
-    private[immutable] val underlying: Map[K, (Int, V)])
-    extends AbstractMap[K, V]
+final class VectorMap[K, +V] private (
+    private[immutable] val fields: Vector[Any],
+    private[immutable] val underlying: Map[K, (Int, V)], dummy: Boolean)
+  extends AbstractMap[K, V]
     with SeqMap[K, V]
     with MapOps[K, V, VectorMap, VectorMap[K, V]]
     with StrictOptimizedIterableOps[(K, V), Iterable, VectorMap[K, V]] {
+
+  import VectorMap._
+
   override protected[this] def className: String = "VectorMap"
 
+  private[immutable] def this(fields: Vector[K], underlying: Map[K, (Int, V)]) = {
+    this(fields, underlying, false)
+  }
+
+  override val size = underlying.size
+
+  override def knownSize: Int = size
+
+  override def isEmpty: Boolean = size == 0
+
   def updated[V1 >: V](key: K, value: V1): VectorMap[K, V1] = {
     underlying.get(key) match {
-      case Some(oldIndexWithValue) =>
-        new VectorMap(fields,
-          underlying.updated(key, (oldIndexWithValue._1, value)))
+      case Some((slot, _)) =>
+        new VectorMap(fields, underlying.updated[(Int, V1)](key, (slot, value)), false)
       case None =>
-        new VectorMap(
-          fields :+ key,
-          underlying.updated(key, (fields.length, value)))
+        new VectorMap(fields :+ key, underlying.updated[(Int, V1)](key, (fields.length, value)), false)
     }
   }
 
   override def withDefault[V1 >: V](d: K => V1): Map.WithDefault[K, V1] =
     new Map.WithDefault(this, d)
 
-  override def withDefaultValue[V1 >: V](d: V1): Map.WithDefault[K, V1] = new Map.WithDefault[K, V1](this, _ => d)
+  override def withDefaultValue[V1 >: V](d: V1): Map.WithDefault[K, V1] =
+    new Map.WithDefault[K, V1](this, _ => d)
+
+  def get(key: K): Option[V] = underlying.get(key) match {
+    case Some(v) => Some(v._2)
+    case None    => None
+  }
+
+  @tailrec
+  private def field(slot: Int): (Int, K) = {
+    fields(slot) match {
+      case Tombstone.Kinless =>
+        (-1, null.asInstanceOf[K])
+      case Tombstone.NextOfKin(distance) =>
+        field(slot + distance)
+      case k: K =>
+        (slot, k)
+    }
+  }
 
   def iterator: Iterator[(K, V)] = new AbstractIterator[(K, V)] {
-    private val fieldsIterator = fields.iterator
+    private[this] val fieldsLength = fields.length
+    private[this] var slot = -1
+    private[this] var key: K = null.asInstanceOf[K]
+
+    private[this] def advance(): Unit = {
+      val nextSlot = slot + 1
+      if (nextSlot >= fieldsLength) {
+        slot = fieldsLength
+        key = null.asInstanceOf[K]
+      } else {
+        field(nextSlot) match {
+          case (-1, _) ⇒
+            slot = fieldsLength
+            key = null.asInstanceOf[K]
+          case (s, k) ⇒
+            slot = s
+            key = k
+        }
+      }
+    }
+
+    advance()
 
-    override def hasNext: Boolean = fieldsIterator.hasNext
+    override def hasNext: Boolean = slot < fieldsLength
 
     override def next(): (K, V) = {
-      val field = fieldsIterator.next()
-      (field, underlying(field)._2)
+      if (!hasNext) throw new NoSuchElementException("next called on depleted iterator")
+      val result = (key, underlying(key)._2)
+      advance()
+      result
     }
   }
 
-  def get(key: K): Option[V] = underlying.get(key) match {
-    case Some(v) => Some(v._2)
-    case None => None
-  }
-
   def remove(key: K): VectorMap[K, V] = {
-    underlying.get(key) match {
-      case Some((index, _)) =>
-        val finalUnderlying = (underlying - key).map{ case (k, (currentIndex,v)) =>
-          if (currentIndex > index)
-            (k, (currentIndex - 1, v))
-          else
-            (k, (currentIndex, v))
-        }
-        new VectorMap(fields.patch(index, Nil, 1), finalUnderlying)
-      case _ =>
-        this
+    if (isEmpty) empty
+    else {
+      var fs = fields
+      val sz = fs.size
+      underlying.get(key) match {
+        case Some((slot, _)) =>
+          val s = field(slot)._1
+          // Calculate distance to next of kin
+          val d =
+            if (s < sz - 1) fs(s + 1) match {
+              case Tombstone.Kinless => 0
+              case Tombstone.NextOfKin(d) => d + 1
+              case _ => 1
+            } else 0
+          fs = fs.updated(s, Tombstone(d))
+          if (s > 0) {
+            // Adjust distance to next of kin for all preceding tombstones
+            var t = s - 1
+            var prev = fs(t)
+            while (t >= 0 && prev.isInstanceOf[Tombstone]) {
+              fs = prev match {
+                case Tombstone.Kinless => throw new IllegalStateException("kinless tombstone found in prefix: " + key)
+                case Tombstone.NextOfKin(_) if d == 0 => fs.updated(t, Tombstone.Kinless)
+                case Tombstone.NextOfKin(d) => fs.updated(t, Tombstone(d + 1))
+                case _ => fs
+              }
+              t -= 1
+              if (t >= 0) prev = fs(t)
+            }
+          }
+          new VectorMap(fs, underlying - key, false)
+        case _ =>
+          this
+      }
     }
   }
 
   override def mapFactory: MapFactory[VectorMap] = VectorMap
 
-  override def size: Int = fields.size
-
-  override def knownSize: Int = fields.size
-
-  override def isEmpty: Boolean = fields.isEmpty
-
-  override final def contains(key: K): Boolean = underlying.contains(key)
+  override def contains(key: K): Boolean = underlying.contains(key)
 
   override def head: (K, V) = iterator.next()
 
   override def last: (K, V) = {
-    val last = fields.last
+    val last = fields
+      .reverseIterator
+      .find(!_.isInstanceOf[Tombstone])
+      .get
+      .asInstanceOf[K]
     (last, underlying(last)._2)
   }
 
   override def lastOption: Option[(K, V)] = {
-    fields.lastOption match {
-      case Some(last) => Some((last, underlying(last)._2))
-      case None => None
-    }
+    fields
+      .reverseIterator
+      .find(!_.isInstanceOf[Tombstone])
+      .map { f ⇒
+        val last = f.asInstanceOf[K]
+        (last, underlying(last)._2)
+      }
   }
 
   override def tail: VectorMap[K, V] = {
-    new VectorMap(fields.tail, underlying.remove(fields.head))
+    val (slot, key) = field(0)
+    new VectorMap(fields.drop(slot + 1), underlying - key, false)
   }
 
   override def init: VectorMap[K, V] = {
-    new VectorMap(fields.init, underlying.remove(fields.last))
+    val (slot, key) = field(size - 1)
+    new VectorMap(fields.dropRight(size - 1 - slot + 1), underlying - key, false)
   }
 
-  // Only care about content, not ordering for equality
-  override def equals(that: Any): Boolean =
-    that match {
-      case map: Map[K, V] =>
-        (this eq map) ||
-          (this.size == map.size) && {
-            try {
-              var i = 0
-              val _size = size
-              while (i < _size) {
-                val k = fields(i)
-
-                map.get(k) match {
-                  case Some(value) =>
-                    if (!(value == underlying(k)._2)) {
-                      return false
-                    }
-                  case None =>
-                    return false
-                }
-                i += 1
-              }
-              true
-            } catch {
-              case _: ClassCastException => false
-            }
-          }
-      case _ => super.equals(that)
-    }
-
-  override def foreach[U](f: ((K, V)) => U): Unit = iterator.foreach(f)
-
-  override def keys: Iterable[K] = fields
+  def keyIterator: Iterator[K] = iterator.map(_._1)
+  override def keys: Vector[K] = keyIterator.toVector
 
   override def values: Iterable[V] = new Iterable[V] {
-    override def iterator: Iterator[V] = fields.iterator.map(underlying(_)._2)
+    override def iterator: Iterator[V] = keyIterator.map(underlying(_)._2)
   }
 }
 
 object VectorMap extends MapFactory[VectorMap] {
+  private[VectorMap] sealed trait Tombstone
+  private[VectorMap] object Tombstone {
+    final case object Kinless extends Tombstone {
+      override def toString = "⤞"
+    }
+    final case class NextOfKin private (distance: Int) extends Tombstone {
+      override def toString = "⥅" + distance
+    }
+    def apply(distance: Int): Tombstone =
+      if (distance <= 0) Kinless
+      else NextOfKin(distance)
+  }
 
   def empty[K, V]: VectorMap[K, V] =
     new VectorMap[K, V](
@@ -172,8 +225,8 @@ object VectorMap extends MapFactory[VectorMap] {
       case _                   => (newBuilder[K, V] ++= it).result()
     }
 
-  def newBuilder[K, V]: Builder[(K, V), VectorMap[K, V]] =
-    new ImmutableBuilder[(K, V), VectorMap[K, V]](empty) {
+  def newBuilder[K, V]: mutable.Builder[(K, V), VectorMap[K, V]] =
+    new mutable.ImmutableBuilder[(K, V), VectorMap[K, V]](empty) {
       def addOne(elem: (K, V)): this.type = { elems = elems + elem; this }
     }
 

test/scalacheck/scala/collection/immutable/VectorMapProperties.scala

@@ -14,24 +14,28 @@ object VectorMapProperties extends Properties("immutable.VectorMap") {
   property("internal underlying index match") = forAll { m: Map[K, V] =>
     !m.isEmpty ==> {
       val vm = VectorMap.from(m)
-      val last = vm.fields.last
-      vm.underlying(last)._1 == vm.size - 1
+      val last = vm.keys.last
+      vm.fields(vm.underlying(last)._1) == last
     }
   }
 
   property("internal underlying and field length") = forAll { m: Map[K, V] => {
       val vm = VectorMap.from(m)
-      vm.underlying.size == vm.fields.length
+      vm.underlying.size == vm.keys.length
     }
   }
 
-  property("internal underlying has consistent index") = forAll { m: Map[K, V] =>
+  property("internal underlying and index are consistent after removal") = forAll { (m: Map[K, V]) =>
     m.size >= 3 ==> {
       val v = Vector.from(m)
       val random = v(new scala.util.Random().nextInt(v.size))
       val vm = VectorMap.from(v)
       val removed = vm - random._1
-      removed.underlying.toList.map{case (k, v) => v._1}.sorted.sliding(2).forall(x => x(1) - x(0) == 1 )
+      removed.underlying.forall { case (k, (s, v)) => removed.fields(s) == k }
+      removed.fields.zipWithIndex.forall {
+        case (k: K, s) => removed.underlying(k)._1 == s
+        case _ => true
+      }
     }
   }
 }