Use only new SparseInArray scheme for tree addresses

Author: odersky

compiler/src/dotty/tools/dotc/core/tasty/TreeBuffer.scala

@@ -21,29 +21,20 @@ class TreeBuffer extends TastyBuffer(50000) {
   private var delta: Array[Int] = _
   private var numOffsets = 0
 
-  /** A map from trees to the address at which a tree is pickled. */
-  private val treeAddrs = new java.util.IdentityHashMap[Tree, Any] // really: Addr | Null
-  private val treeAddrs2 = SparseIntArray()
-
-  def registerTreeAddr(tree: Tree): Addr = treeAddrs.get(tree) match {
-    case null =>
-      treeAddrs.put(tree, currentAddr)
-      treeAddrs2(tree.uniqueId) = currentAddr.index
+  /** A map from tree unique ids to the address index at which a tree is pickled. */
+  private val addrOfTree = SparseIntArray()
+
+  def registerTreeAddr(tree: Tree): Addr =
+    val id = tree.uniqueId
+    if addrOfTree.contains(id) then Addr(addrOfTree(id))
+    else
+      addrOfTree(tree.uniqueId) = currentAddr.index
       currentAddr
-    case addr: Addr => addr
-  }
 
   def addrOfTree(tree: Tree): Addr =
-    val old = treeAddrs.get(tree) match {
-      case null => NoAddr
-      case addr: Addr => addr
-    }
-    val now =
-      val idx = tree.uniqueId
-      if treeAddrs2.contains(idx) then Addr(treeAddrs2(idx))
-      else NoAddr
-    assert(old == now)
-    now
+    val idx = tree.uniqueId
+    if addrOfTree.contains(idx) then Addr(addrOfTree(idx))
+    else NoAddr
 
   private def offset(i: Int): Addr = Addr(offsets(i))
 
@@ -168,16 +159,8 @@ class TreeBuffer extends TastyBuffer(50000) {
     wasted
   }
 
-  def adjustTreeAddrs(): Unit = {
-    val it = treeAddrs.keySet.iterator
-    while (it.hasNext) {
-      val tree = it.next
-      treeAddrs.get(tree) match {
-        case addr: Addr => treeAddrs.put(tree, adjusted(addr))
-      }
-    }
-    treeAddrs2.transform((id, addr) => adjusted(Addr(addr)).index)
-  }
+  def adjustTreeAddrs(): Unit =
+    addrOfTree.transform((id, addr) => adjusted(Addr(addr)).index)
 
   /** Final assembly, involving the following steps:
    *   - compute deltas

compiler/src/dotty/tools/dotc/util/SparseIntArray.scala

@@ -42,9 +42,16 @@ class SparseIntArray:
       result
     }
 
+  /** All defined indices in an iterator */
+  def keysIterator: Iterator[Int] = root.keysIterator(0)
+
+  /** Perform operation for each key/value pair */
   def foreachBinding(op: (Int, Value) => Unit): Unit =
     root.foreachBinding(op, 0)
 
+  /** Transform each defined value with transformation `op`.
+   *  The transformation gets the element index and value as parameters.
+   */
   def transform(op: (Int, Value) => Value): Unit =
     root.transform(op, 0)
 
@@ -82,6 +89,7 @@ object SparseIntArray:
     def update(index: Int, value: Value): Boolean
     def remove(index: Int): Boolean
     def isEmpty: Boolean
+    def keysIterator(offset: Int): Iterator[Int]
     def foreachBinding(op: (Int, Value) => Unit, offset: Int): Unit
     def transform(op: (Int, Value) => Value, offset: Int): Unit
     def nodeCount: Int
@@ -111,6 +119,17 @@ object SparseIntArray:
 
     def isEmpty = present == 0
 
+    private def skipUndefined(i: Int): Int =
+      if i < NodeSize && !contains(i) then skipUndefined(i + 1) else i
+
+    def keysIterator(offset: Int) = new Iterator[Int]:
+      private var curIdx = skipUndefined(0)
+      def hasNext = curIdx < NodeSize
+      def next(): Int =
+        val result = curIdx + offset
+        curIdx = skipUndefined(curIdx + 1)
+        result
+
     def foreachBinding(op: (Int, Value) => Unit, offset: Int): Unit =
       var i = 0
       while i < NodeSize do
@@ -167,6 +186,20 @@ object SparseIntArray:
 
     def isEmpty = empty
 
+    private def skipUndefined(i: Int): Int =
+      if i < NodeSize && elems(i) == null then skipUndefined(i + 1) else i
+
+    def keysIterator(offset: Int) = new Iterator[Value]:
+      private var curIdx = skipUndefined(0)
+      private var elemIt = Iterator.empty[Int]
+      def hasNext = elemIt.hasNext || curIdx < NodeSize
+      def next(): Value =
+        if elemIt.hasNext then elemIt.next()
+        else
+          elemIt = elems(curIdx).keysIterator(offset + curIdx * elemSize)
+          curIdx = skipUndefined(curIdx + 1)
+          elemIt.next()
+
     def foreachBinding(op: (Int, Value) => Unit, offset: Int): Unit =
       var i = 0
       while i < NodeSize do
@@ -205,6 +238,7 @@ end SparseIntArray
   println(s"a = $a")
   a(55555) = 44
   println(s"a = $a")
+  println(s"iterator of a yields ${a.keysIterator.toList}")
   assert(a.size == 3, a)
   assert(a.contains(1), a)
   assert(a.contains(222), a)