Trial: Use a sparse array representation for Tree-address maps.

Make use of the fact that every tree has a unique id, so we can use that
as an index into a sparse array.

Author: odersky

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

@@ -8,6 +8,7 @@ import dotty.tools.tasty.TastyBuffer
 import TastyBuffer.{Addr, NoAddr, AddrWidth}
 
 import util.Util.bestFit
+import util.SparseIntArray
 import config.Printers.pickling
 import ast.untpd.Tree
 
@@ -22,16 +23,27 @@ class TreeBuffer extends TastyBuffer(50000) {
 
   /** 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); currentAddr
+    case null =>
+      treeAddrs.put(tree, currentAddr)
+      treeAddrs2(tree.uniqueId) = currentAddr.index
+      currentAddr
     case addr: Addr => addr
   }
 
-  def addrOfTree(tree: Tree): Addr = treeAddrs.get(tree) match {
-    case null => NoAddr
-    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
 
   private def offset(i: Int): Addr = Addr(offsets(i))
 
@@ -164,6 +176,7 @@ class TreeBuffer extends TastyBuffer(50000) {
         case addr: Addr => treeAddrs.put(tree, adjusted(addr))
       }
     }
+    treeAddrs2.transform((id, addr) => adjusted(Addr(addr)).index)
   }
 
   /** Final assembly, involving the following steps:

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

@@ -0,0 +1,228 @@
+package dotty.tools.dotc
+package util
+
+import java.util.NoSuchElementException
+
+class SparseIntArray:
+  import SparseIntArray._
+
+  private var siz: Int = 0
+  private var root: Node = LeafNode()
+
+  private def grow() =
+    val newRoot = InnerNode(root.level + 1)
+    newRoot.elems(0) = root
+    root = newRoot
+
+  private def capacity: Int = root.elemSize * NodeSize
+
+  def size = siz
+
+  def contains(index: Int): Boolean =
+    0 <= index && index < capacity && root.contains(index)
+
+  def apply(index: Int): Value =
+    require(index >= 0)
+    if index >= capacity then throw NoSuchElementException()
+    root.apply(index)
+
+  def update(index: Int, value: Value): Unit =
+    require(index >= 0)
+    while capacity <= index do grow()
+    if !root.update(index, value) then siz += 1
+
+  /** Remove element at `index` if it is present
+   *  @return element was present
+   */
+  def remove(index: Int): Boolean =
+    require(index >= 0)
+    index < capacity && {
+      val result = root.remove(index)
+      if result then siz -= 1
+      result
+    }
+
+  def foreachBinding(op: (Int, Value) => Unit): Unit =
+    root.foreachBinding(op, 0)
+
+  def transform(op: (Int, Value) => Value): Unit =
+    root.transform(op, 0)
+
+  /** Access to some info about low-level representation */
+  def repr: Repr = root
+
+  override def toString =
+    val b = StringBuilder() ++= "SparseIntArray("
+    var first = true
+    foreachBinding { (idx, elem) =>
+      if first then first = false else b ++= ", "
+      b ++= s"$idx -> $elem"
+    }
+    b ++= ")"
+    b.toString
+
+object SparseIntArray:
+  type Value = Int
+
+  private inline val NodeSizeLog = 5
+  private inline val NodeSize = 1 << NodeSizeLog
+
+  /** The exposed representation. Should be used just for nodeCount and
+   *  low-level toString.
+   */
+  abstract class Repr:
+    def nodeCount: Int
+
+  private abstract class Node(val level: Int) extends Repr:
+    private[SparseIntArray] def elemShift = level * NodeSizeLog
+    private[SparseIntArray] def elemSize = 1 << elemShift
+    private[SparseIntArray] def elemMask = elemSize - 1
+    def contains(index: Int): Boolean
+    def apply(index: Int): Value
+    def update(index: Int, value: Value): Boolean
+    def remove(index: Int): Boolean
+    def isEmpty: Boolean
+    def foreachBinding(op: (Int, Value) => Unit, offset: Int): Unit
+    def transform(op: (Int, Value) => Value, offset: Int): Unit
+    def nodeCount: Int
+  end Node
+
+  private class LeafNode extends Node(0):
+    private val elems = new Array[Value](NodeSize)
+    private var present: Int = 0
+
+    def contains(index: Int): Boolean =
+      (present & (1 << index)) != 0
+
+    def apply(index: Int) =
+      if !contains(index) then throw NoSuchElementException()
+      elems(index)
+
+    def update(index: Int, value: Value): Boolean =
+      elems(index) = value
+      val result = contains(index)
+      present = present | (1 << index)
+      result
+
+    def remove(index: Int): Boolean =
+      val result = contains(index)
+      present = present & ~(1 << index)
+      result
+
+    def isEmpty = present == 0
+
+    def foreachBinding(op: (Int, Value) => Unit, offset: Int): Unit =
+      var i = 0
+      while i < NodeSize do
+        if contains(i) then op(offset + i, elems(i))
+        i += 1
+
+    def transform(op: (Int, Value) => Value, offset: Int): Unit =
+      var i = 0
+      while i < NodeSize do
+        if contains(i) then elems(i) = op(offset + i, elems(i))
+        i += 1
+
+    def nodeCount = 1
+
+    override def toString =
+      elems
+        .zipWithIndex
+        .filter((elem, idx) => contains(idx))
+        .map((elem, idx) => s"$idx -> $elem").mkString(s"0#(", ", ", ")")
+  end LeafNode
+
+  private class InnerNode(level: Int) extends Node(level):
+    private[SparseIntArray] val elems = new Array[Node](NodeSize)
+    private var empty: Boolean = true
+
+    def contains(index: Int): Boolean =
+      val elem = elems(index >>> elemShift)
+      elem != null && elem.contains(index & elemMask)
+
+    def apply(index: Int): Value =
+      val elem = elems(index >>> elemShift)
+      if elem == null then throw NoSuchElementException()
+      elem.apply(index & elemMask)
+
+    def update(index: Int, value: Value): Boolean =
+      empty = false
+      var elem = elems(index >>> elemShift)
+      if elem == null then
+        elem = newNode(level - 1)
+        elems(index >>> elemShift) = elem
+      elem.update(index & elemMask, value)
+
+    def remove(index: Int): Boolean =
+      val elem = elems(index >>> elemShift)
+      if elem == null then false
+      else
+        val result = elem.remove(index & elemMask)
+        if elem.isEmpty then
+          elems(index >>> elemShift) = null
+          var i = 0
+          while i < NodeSize && elems(i) == null do i += 1
+          if i == NodeSize then empty = true
+        result
+
+    def isEmpty = empty
+
+    def foreachBinding(op: (Int, Value) => Unit, offset: Int): Unit =
+      var i = 0
+      while i < NodeSize do
+        if elems(i) != null then
+          elems(i).foreachBinding(op, offset + i * elemSize)
+        i += 1
+
+    def transform(op: (Int, Value) => Value, offset: Int): Unit =
+      var i = 0
+      while i < NodeSize do
+        if elems(i) != null then
+          elems(i).transform(op, offset + i * elemSize)
+        i += 1
+
+    def nodeCount =
+      1 + elems.filter(_ != null).map(_.nodeCount).sum
+
+    override def toString =
+      elems
+        .zipWithIndex
+        .filter((elem, idx) => elem != null)
+        .map((elem, idx) => s"$idx -> $elem").mkString(s"$level#(", ", ", ")")
+  end InnerNode
+
+  private def newNode(level: Int): Node =
+    if level == 0 then LeafNode() else InnerNode(level)
+
+end SparseIntArray
+
+@main def Test =
+  val a = SparseIntArray()
+  println(s"a = $a")
+  a(1) = 22
+  println(s"a = $a")
+  a(222) = 33
+  println(s"a = $a")
+  a(55555) = 44
+  println(s"a = $a")
+  assert(a.size == 3, a)
+  assert(a.contains(1), a)
+  assert(a.contains(222), a)
+  assert(a.contains(55555), a)
+  assert(!a.contains(2))
+  assert(!a.contains(20000000))
+  a(222) = 44
+  assert(a.size == 3)
+  assert(a(1) == 22)
+  assert(a(222) == 44)
+  assert(a(55555) == 44)
+  assert(a.remove(1))
+  println(s"a = $a")
+  assert(a(222) == 44, a)
+  assert(a.remove(55555))
+  assert(a(222) == 44, a)
+  assert(a.size == 1)
+  assert(!a.contains(1))
+  assert(!a.remove(55555))
+  assert(a.remove(222))
+  assert(a.size == 0)