Use WeakHashSet instead of HashSet for hash-consing types

This mimics what Scala 2 has been doing for a long time now and serves
the same purpose: it considerably reduces peak memory usage when
compiling some projects, for example previously compiling the Scalatest
tests required a heap of at least 11 GB, but now it fits in about 4 GB.

This required changing the implementation of WeakHashSet to have
overridable `hash` and `isEqual` methods just like HashSet, it also
required making various private methods protected since NamedTypeUniques
and AppliedUniques contain an inlined implementation of `put`.

This commit also changes the default load factor of a WeakHashSet from
0.75 to 0.5 to match the load factor we use for HashSets, though note that
Scala 2 has always been using 0.75.

For a history of the usage of WeakHashSet in Scala 2 see:
- scala/scala#247
- scala/scala#2605
- scala/scala#2901

Author: smarter

compiler/src/dotty/tools/dotc/core/Contexts.scala

@@ -559,7 +559,7 @@ object Contexts {
     def platform: Platform                 = base.platform
     def pendingUnderlying: util.HashSet[Type]      = base.pendingUnderlying
     def uniqueNamedTypes: Uniques.NamedTypeUniques = base.uniqueNamedTypes
-    def uniques: util.HashSet[Type]                = base.uniques
+    def uniques: util.WeakHashSet[Type]            = base.uniques
 
     def initialize()(using Context): Unit = base.initialize()
   }

compiler/src/dotty/tools/dotc/core/Uniques.scala

@@ -4,9 +4,11 @@ package core
 import Types._, Contexts._, util.Stats._, Hashable._, Names._
 import config.Config
 import Decorators._
-import util.{HashSet, Stats}
+import util.{WeakHashSet, Stats}
+import WeakHashSet.Entry
+import scala.annotation.tailrec
 
-class Uniques extends HashSet[Type](Config.initialUniquesCapacity):
+class Uniques extends WeakHashSet[Type](Config.initialUniquesCapacity):
   override def hash(x: Type): Int = x.hash
   override def isEqual(x: Type, y: Type) = x.eql(y)
 
@@ -32,7 +34,7 @@ object Uniques:
     if tp.hash == NotCached then tp
     else ctx.uniques.put(tp).asInstanceOf[T]
 
-  final class NamedTypeUniques extends HashSet[NamedType](Config.initialUniquesCapacity * 4) with Hashable:
+  final class NamedTypeUniques extends WeakHashSet[NamedType](Config.initialUniquesCapacity * 4) with Hashable:
     override def hash(x: NamedType): Int = x.hash
 
     def enterIfNew(prefix: Type, designator: Designator, isTerm: Boolean)(using Context): NamedType =
@@ -43,17 +45,25 @@ object Uniques:
         else new CachedTypeRef(prefix, designator, h)
       if h == NotCached then newType
       else
+        // Inlined from WeakHashSet#put
         Stats.record(statsItem("put"))
-        var idx = index(h)
-        var e = entryAt(idx)
-        while e != null do
-          if (e.prefix eq prefix) && (e.designator eq designator) && (e.isTerm == isTerm) then return e
-          idx = nextIndex(idx)
-          e = entryAt(idx)
-        addEntryAt(idx, newType)
+        removeStaleEntries()
+        val bucket = index(h)
+        val oldHead = table(bucket)
+
+        @tailrec
+        def linkedListLoop(entry: Entry[NamedType]): NamedType = entry match
+          case null                    => addEntryAt(bucket, newType, h, oldHead)
+          case _                       =>
+            val e = entry.get
+            if e != null && (e.prefix eq prefix) && (e.designator eq designator) && (e.isTerm == isTerm) then e
+            else linkedListLoop(entry.tail)
+
+        linkedListLoop(oldHead)
+      end if
   end NamedTypeUniques
 
-  final class AppliedUniques extends HashSet[AppliedType](Config.initialUniquesCapacity * 2) with Hashable:
+  final class AppliedUniques extends WeakHashSet[AppliedType](Config.initialUniquesCapacity * 2) with Hashable:
     override def hash(x: AppliedType): Int = x.hash
 
     def enterIfNew(tycon: Type, args: List[Type]): AppliedType =
@@ -62,13 +72,21 @@ object Uniques:
       if monitored then recordCaching(h, classOf[CachedAppliedType])
       if h == NotCached then newType
       else
+        // Inlined from WeakHashSet#put
         Stats.record(statsItem("put"))
-        var idx = index(h)
-        var e = entryAt(idx)
-        while e != null do
-          if (e.tycon eq tycon) && e.args.eqElements(args) then return e
-          idx = nextIndex(idx)
-          e = entryAt(idx)
-        addEntryAt(idx, newType)
+        removeStaleEntries()
+        val bucket = index(h)
+        val oldHead = table(bucket)
+
+        @tailrec
+        def linkedListLoop(entry: Entry[AppliedType]): AppliedType = entry match
+          case null                    => addEntryAt(bucket, newType, h, oldHead)
+          case _                       =>
+            val e = entry.get
+            if e != null && (e.tycon eq tycon) && e.args.eqElements(args) then e
+            else linkedListLoop(entry.tail)
+
+        linkedListLoop(oldHead)
+      end if
   end AppliedUniques
 end Uniques

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

@@ -17,25 +17,23 @@ import scala.collection.mutable
  * This set implementation is not in general thread safe without external concurrency control. However it behaves
  * properly when GC concurrently collects elements in this set.
  */
-final class WeakHashSet[A <: AnyRef](initialCapacity: Int, loadFactor: Double) extends MutableSet[A] {
+abstract class WeakHashSet[A <: AnyRef](initialCapacity: Int = 8, loadFactor: Double = 0.5) extends MutableSet[A] {
 
   import WeakHashSet._
 
-  def this() = this(initialCapacity = WeakHashSet.defaultInitialCapacity, loadFactor = WeakHashSet.defaultLoadFactor)
-
   type This = WeakHashSet[A]
 
   /**
    * queue of Entries that hold elements scheduled for GC
    * the removeStaleEntries() method works through the queue to remove
    * stale entries from the table
    */
-  private val queue = new ReferenceQueue[A]
+  protected val queue = new ReferenceQueue[A]
 
   /**
    * the number of elements in this set
    */
-  private var count = 0
+  protected var count = 0
 
   /**
    * from a specified initial capacity compute the capacity we'll use as being the next
@@ -52,33 +50,20 @@ final class WeakHashSet[A <: AnyRef](initialCapacity: Int, loadFactor: Double) e
   /**
    * the underlying table of entries which is an array of Entry linked lists
    */
-  private var table = new Array[Entry[A]](computeCapacity)
+  protected var table = new Array[Entry[A]](computeCapacity)
 
   /**
    * the limit at which we'll increase the size of the hash table
    */
-  private var threshold = computeThreshold
+  protected var threshold = computeThreshold
 
   private def computeThreshold: Int = (table.size * loadFactor).ceil.toInt
 
-  /**
-   * find the bucket associated with an element's hash code
-   */
-  private def bucketFor(hash: Int): Int = {
-    // spread the bits around to try to avoid accidental collisions using the
-    // same algorithm as java.util.HashMap
-    var h = hash
-    h ^= h >>> 20 ^ h >>> 12
-    h ^= h >>> 7 ^ h >>> 4
-
-    // this is finding h % table.length, but takes advantage of the
-    // fact that table length is a power of 2,
-    // if you don't do bit flipping in your head, if table.length
-    // is binary 100000.. (with n 0s) then table.length - 1
-    // is 1111.. with n 1's.
-    // In other words this masks on the last n bits in the hash
-    h & (table.length - 1)
-  }
+  protected def hash(key: A): Int
+  protected def isEqual(x: A, y: A): Boolean = x.equals(y)
+
+  /** Turn hashcode `x` into a table index */
+  protected def index(x: Int): Int = x & (table.length - 1)
 
   /**
    * remove a single entry from a linked list in a given bucket
@@ -95,14 +80,14 @@ final class WeakHashSet[A <: AnyRef](initialCapacity: Int, loadFactor: Double) e
   /**
    * remove entries associated with elements that have been gc'ed
    */
-  private def removeStaleEntries(): Unit = {
+  protected def removeStaleEntries(): Unit = {
     def poll(): Entry[A] = queue.poll().asInstanceOf[Entry[A]]
 
     @tailrec
     def queueLoop(): Unit = {
       val stale = poll()
       if (stale != null) {
-        val bucket = bucketFor(stale.hash)
+        val bucket = index(stale.hash)
 
         @tailrec
         def linkedListLoop(prevEntry: Entry[A], entry: Entry[A]): Unit = if (stale eq entry) remove(bucket, prevEntry, entry)
@@ -120,7 +105,7 @@ final class WeakHashSet[A <: AnyRef](initialCapacity: Int, loadFactor: Double) e
   /**
    * Double the size of the internal table
    */
-  private def resize(): Unit = {
+  protected def resize(): Unit = {
     Stats.record(statsItem("resize"))
     val oldTable = table
     table = new Array[Entry[A]](oldTable.size * 2)
@@ -132,7 +117,7 @@ final class WeakHashSet[A <: AnyRef](initialCapacity: Int, loadFactor: Double) e
       def linkedListLoop(entry: Entry[A]): Unit = entry match {
         case null => ()
         case _ =>
-          val bucket = bucketFor(entry.hash)
+          val bucket = index(entry.hash)
           val oldNext = entry.tail
           entry.tail = table(bucket)
           table(bucket) = entry
@@ -150,43 +135,43 @@ final class WeakHashSet[A <: AnyRef](initialCapacity: Int, loadFactor: Double) e
     case _    =>
       Stats.record(statsItem("lookup"))
       removeStaleEntries()
-      val hash = elem.hashCode
-      val bucket = bucketFor(hash)
+      val bucket = index(hash(elem))
 
       @tailrec
       def linkedListLoop(entry: Entry[A]): A = entry match {
         case null                    => null.asInstanceOf[A]
         case _                       =>
           val entryElem = entry.get
-          if (elem.equals(entryElem)) entryElem
+          if (isEqual(elem, entryElem)) entryElem
           else linkedListLoop(entry.tail)
       }
 
       linkedListLoop(table(bucket))
   }
 
+  protected def addEntryAt(bucket: Int, elem: A, elemHash: Int, oldHead: Entry[A]): A = {
+    Stats.record(statsItem("addEntryAt"))
+    table(bucket) = new Entry(elem, elemHash, oldHead, queue)
+    count += 1
+    if (count > threshold) resize()
+    elem
+  }
+
   def put(elem: A): A = elem match {
     case null => throw new NullPointerException("WeakHashSet cannot hold nulls")
     case _    =>
       Stats.record(statsItem("put"))
       removeStaleEntries()
-      val hash = elem.hashCode
-      val bucket = bucketFor(hash)
+      val h = hash(elem)
+      val bucket = index(h)
       val oldHead = table(bucket)
 
-      def add() = {
-        table(bucket) = new Entry(elem, hash, oldHead, queue)
-        count += 1
-        if (count > threshold) resize()
-        elem
-      }
-
       @tailrec
       def linkedListLoop(entry: Entry[A]): A = entry match {
-        case null                    => add()
+        case null                    => addEntryAt(bucket, elem, h, oldHead)
         case _                       =>
           val entryElem = entry.get
-          if (elem.equals(entryElem)) entryElem
+          if (isEqual(elem, entryElem)) entryElem
           else linkedListLoop(entry.tail)
       }
 
@@ -200,14 +185,14 @@ final class WeakHashSet[A <: AnyRef](initialCapacity: Int, loadFactor: Double) e
     case _ =>
       Stats.record(statsItem("-="))
       removeStaleEntries()
-      val bucket = bucketFor(elem.hashCode)
+      val bucket = index(hash(elem))
 
 
 
       @tailrec
       def linkedListLoop(prevEntry: Entry[A], entry: Entry[A]): Unit = entry match {
         case null => ()
-        case _ if elem.equals(entry.get) => remove(bucket, prevEntry, entry)
+        case _ if isEqual(elem, entry.get) => remove(bucket, prevEntry, entry)
         case _ => linkedListLoop(entry, entry.tail)
       }
 
@@ -307,9 +292,9 @@ final class WeakHashSet[A <: AnyRef](initialCapacity: Int, loadFactor: Double) e
           assert(entry.get != null, s"$entry had a null value indicated that gc activity was happening during diagnostic validation or that a null value was inserted")
           computedCount += 1
           val cachedHash = entry.hash
-          val realHash = entry.get.hashCode
+          val realHash = hash(entry.get)
           assert(cachedHash == realHash, s"for $entry cached hash was $cachedHash but should have been $realHash")
-          val computedBucket = bucketFor(realHash)
+          val computedBucket = index(realHash)
           assert(computedBucket == bucket, s"for $entry the computed bucket was $computedBucket but should have been $bucket")
 
           entry = entry.tail
@@ -355,11 +340,6 @@ object WeakHashSet {
    * A single entry in a WeakHashSet. It's a WeakReference plus a cached hash code and
    * a link to the next Entry in the same bucket
    */
-  private class Entry[A](@constructorOnly element: A, val hash:Int, var tail: Entry[A], @constructorOnly queue: ReferenceQueue[A]) extends WeakReference[A](element, queue)
-
-  private final val defaultInitialCapacity = 16
-  private final val defaultLoadFactor = .75
+  class Entry[A](@constructorOnly element: A, val hash:Int, var tail: Entry[A], @constructorOnly queue: ReferenceQueue[A]) extends WeakReference[A](element, queue)
 
-  def apply[A <: AnyRef](initialCapacity: Int = defaultInitialCapacity, loadFactor: Double = defaultLoadFactor): WeakHashSet[A] =
-    new WeakHashSet(initialCapacity, loadFactor)
 }