Homogenize mutable symbol maps

Use IdentityHashMap instead of mutable.HashMap to make fast equality check

Author: nicolasstucki

compiler/sjs/backend/sjs/JSPrimitives.scala

@@ -59,7 +59,7 @@ class JSPrimitives(ctx: Context) extends DottyPrimitives(ctx) {
   /** Initialize the primitive map */
   private def initJSPrimitives(implicit ctx: Context): Map[Symbol, Int] = {
 
-    val primitives = new mutable.HashMap[Symbol, Int]()
+    val primitives = newMutableSymbolMap[Int]
 
     // !!! Code duplicate with DottyPrimitives
     /** Add a primitive operation to the map */

compiler/src/dotty/tools/backend/jvm/GenBCode.scala

@@ -40,7 +40,7 @@ class GenBCode extends Phase {
   private val entryPoints = new mutable.HashSet[Symbol]()
   def registerEntryPoint(sym: Symbol) = entryPoints += sym
 
-  private val superCallsMap = new mutable.HashMap[Symbol, Set[ClassSymbol]]()
+  private val superCallsMap = newMutableSymbolMap[Set[ClassSymbol]]
   def registerSuperCall(sym: Symbol, calls: ClassSymbol) = {
     val old = superCallsMap.getOrElse(sym, Set.empty)
     superCallsMap.put(sym, old + calls)

compiler/src/dotty/tools/backend/jvm/LabelDefs.scala

@@ -76,7 +76,7 @@ class LabelDefs extends MiniPhase {
             case t: DefDef =>
               assert(t.symbol is Label)
               EmptyTree
-            case _ => if (labelDefs.nonEmpty) super.transform(tree) else tree
+            case _ => if (!labelDefs.isEmpty) super.transform(tree) else tree
           }
         }
       }
@@ -86,9 +86,9 @@ class LabelDefs extends MiniPhase {
     }
   }
 
-  private def collectLabelDefs(tree: Tree)(implicit ctx: Context): mutable.HashMap[Symbol, DefDef] = {
+  private def collectLabelDefs(tree: Tree)(implicit ctx: Context): MutableSymbolMap[DefDef] = {
     // labelSymbol -> Defining tree
-    val labelDefs = new mutable.HashMap[Symbol, DefDef]()
+    val labelDefs = newMutableSymbolMap[DefDef]
     new TreeTraverser {
       override def traverse(tree: Tree)(implicit ctx: Context): Unit = tree match {
         case t: DefDef =>

compiler/src/dotty/tools/backend/jvm/scalaPrimitives.scala

@@ -128,7 +128,7 @@ class DottyPrimitives(ctx: Context) {
     implicit val ctx = this.ctx
 
     import core.Symbols.defn
-    val primitives = new mutable.HashMap[Symbol, Int]()
+    val primitives = core.Symbols.newMutableSymbolMap[Int]
 
     /** Add a primitive operation to the map */
     def addPrimitive(s: Symbol, code: Int): Unit = {

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

@@ -23,10 +23,8 @@ object Comments {
     * docstrings via `Symbol` and expanding templates
     */
   class ContextDocstrings {
-    import scala.collection.mutable
 
-    private[this] val _docstrings: mutable.Map[Symbol, Comment] =
-      mutable.Map.empty
+    private[this] val _docstrings: MutableSymbolMap[Comment] = newMutableSymbolMap
 
     val templateExpander = new CommentExpander
 
@@ -35,7 +33,7 @@ object Comments {
     def docstring(sym: Symbol): Option[Comment] = _docstrings.get(sym)
 
     def addDocstring(sym: Symbol, doc: Option[Comment]): Unit =
-      doc.map(d => _docstrings += (sym -> d))
+      doc.map(d => _docstrings.put(sym, d))
   }
 
   /** A `Comment` contains the unformatted docstring as well as a position
@@ -182,7 +180,7 @@ object Comments {
     protected def superComment(sym: Symbol)(implicit ctx: Context): Option[String] =
       allInheritedOverriddenSymbols(sym).iterator map (x => cookedDocComment(x)) find (_ != "")
 
-    private val cookedDocComments = mutable.HashMap[Symbol, String]()
+    private val cookedDocComments = newMutableSymbolMap[String]
 
     /** The raw doc comment of symbol `sym`, minus usecase and define sections, augmented by
      *  missing sections of an inherited doc comment.

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

@@ -678,4 +678,70 @@ object Symbols {
   def currentClass(implicit ctx: Context): ClassSymbol = ctx.owner.enclosingClass.asClass
 
   @sharable var stubs: List[Symbol] = Nil // diagnostic only
+
+  /* Mutable map from symbols any T */
+  class MutableSymbolMap[T](private[Symbols] val value: java.util.IdentityHashMap[Symbol, T]) extends AnyVal {
+
+    def apply(sym: Symbol): T = value.get(sym)
+
+    def get(sym: Symbol): Option[T] = Option(value.get(sym))
+
+    def getOrElse[U >: T](sym: Symbol, default: => U): U = {
+      val v = value.get(sym)
+      if (v != null) v else default
+    }
+
+    def getOrElseUpdate(sym: Symbol, op: => T): T = {
+      val v = value.get(sym)
+      if (v != null) v
+      else {
+        val v = op
+        assert(v != null)
+        value.put(sym, v)
+        v
+      }
+    }
+
+    def update(sym: Symbol, x: T): Unit = {
+      assert(x != null)
+      value.put(sym, x)
+    }
+    def put(sym: Symbol, x: T): T = {
+      assert(x != null)
+      value.put(sym, x)
+    }
+
+    def -=(sym: Symbol): Unit = value.remove(sym)
+    def remove(sym: Symbol): Option[T] = Option(value.remove(sym))
+
+    def contains(sym: Symbol): Boolean = value.containsKey(sym)
+
+    def isEmpty: Boolean = value.isEmpty
+
+    def clear(): Unit = value.clear()
+
+    def filter(p: ((Symbol, T)) => Boolean): Map[Symbol, T] = {
+      import scala.collection.JavaConversions._
+      value.toMap.filter(p)
+    }
+
+    def iterator: Iterator[(Symbol, T)] = {
+      import scala.collection.JavaConversions._
+      value.iterator
+    }
+
+    def keysIterator: Iterator[Symbol] = {
+      import scala.collection.JavaConversions._
+      value.keySet().iterator
+    }
+
+    def toMap: Map[Symbol, T] = {
+      import scala.collection.JavaConversions._
+      value.toMap
+    }
+  }
+
+  @inline def newMutableSymbolMap[T]: MutableSymbolMap[T] =
+    new MutableSymbolMap(new java.util.IdentityHashMap[Symbol, T]())
+
 }

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

@@ -23,8 +23,8 @@ class TreePickler(pickler: TastyPickler) {
   import pickler.nameBuffer.nameIndex
   import ast.tpd._
 
-  private val symRefs = new mutable.HashMap[Symbol, Addr]
-  private val forwardSymRefs = new mutable.HashMap[Symbol, List[Addr]]
+  private val symRefs = Symbols.newMutableSymbolMap[Addr]
+  private val forwardSymRefs = Symbols.newMutableSymbolMap[List[Addr]]
   private val pickledTypes = new java.util.IdentityHashMap[Type, Any] // Value type is really Addr, but that's not compatible with null
 
   private def withLength(op: => Unit) = {
@@ -586,15 +586,15 @@ class TreePickler(pickler: TastyPickler) {
 
   def pickle(trees: List[Tree])(implicit ctx: Context) = {
     trees.foreach(tree => if (!tree.isEmpty) pickleTree(tree))
-    def missing = forwardSymRefs.keySet.toList.map(_.showLocated)
+    def missing = forwardSymRefs.keysIterator.map(_.showLocated).toList
     assert(forwardSymRefs.isEmpty, i"unresolved symbols: $missing%, % when pickling ${ctx.source}")
   }
 
   def compactify() = {
     buf.compactify()
 
-    def updateMapWithDeltas[T](mp: collection.mutable.Map[T, Addr]) =
-      for (key <- mp.keysIterator.toBuffer[T]) mp(key) = adjusted(mp(key))
+    def updateMapWithDeltas(mp: MutableSymbolMap[Addr]) =
+      for (key <- mp.keysIterator.toBuffer[Symbol]) mp(key) = adjusted(mp(key))
 
     updateMapWithDeltas(symRefs)
   }

compiler/src/dotty/tools/dotc/transform/Bridges.scala

@@ -33,7 +33,7 @@ class Bridges(root: ClassSymbol)(implicit ctx: Context) {
   private[this] var toBeRemoved = immutable.Set[Symbol]()
   private val bridges = mutable.ListBuffer[Tree]()
   private val bridgesScope = newScope
-  private val bridgeTarget = mutable.HashMap[Symbol, Symbol]()
+  private val bridgeTarget = newMutableSymbolMap[Symbol]
 
   /** Add a bridge between `member` and `other`, where `member` overrides `other`
    *  before erasure, if the following conditions are satisfied.

compiler/src/dotty/tools/dotc/transform/ExtensionMethods.scala

@@ -139,7 +139,7 @@ class ExtensionMethods extends MiniPhase with DenotTransformer with FullParamete
     extensionMeth
   }
 
-  private val extensionDefs = mutable.Map[Symbol, mutable.ListBuffer[Tree]]()
+  private val extensionDefs = newMutableSymbolMap[mutable.ListBuffer[Tree]]
   // TODO: this is state and should be per-run
   // todo: check that when transformation finished map is empty
 

compiler/src/dotty/tools/dotc/transform/NonLocalReturns.scala

@@ -33,7 +33,7 @@ class NonLocalReturns extends MiniPhase {
     defn.NonLocalReturnControlType.appliedTo(argtype)
 
   /** A hashmap from method symbols to non-local return keys */
-  private val nonLocalReturnKeys = mutable.Map[Symbol, TermSymbol]()
+  private val nonLocalReturnKeys = newMutableSymbolMap[TermSymbol]
 
   /** Return non-local return key for given method */
   private def nonLocalReturnKey(meth: Symbol)(implicit ctx: Context) =