Specialize SimpleIdentitySet with 3 elements

Author: odersky

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

@@ -37,7 +37,7 @@ object SimpleIdentitySet {
   private class Set1[+Elem <: AnyRef](x0: AnyRef) extends SimpleIdentitySet[Elem] {
     def size = 1
     def + [E >: Elem <: AnyRef](x: E): SimpleIdentitySet[E] =
-      new Set2[E](x0, x)
+      if (contains(x)) this else new Set2[E](x0, x)
     def - [E >: Elem <: AnyRef](x: E): SimpleIdentitySet[Elem] =
       if (x `eq` x0) empty else this
     def contains[E >: Elem <: AnyRef](x: E): Boolean = x `eq` x0
@@ -46,13 +46,8 @@ object SimpleIdentitySet {
 
   private class Set2[+Elem <: AnyRef](x0: AnyRef, x1: AnyRef) extends SimpleIdentitySet[Elem] {
     def size = 2
-    def + [E >: Elem <: AnyRef](x: E): SimpleIdentitySet[E] = {
-      val xs = new Array[AnyRef](3)
-      xs(0) = x0
-      xs(1) = x1
-      xs(2) = x
-      new SetN[E](xs)
-    }
+    def + [E >: Elem <: AnyRef](x: E): SimpleIdentitySet[E] =
+      if (contains(x)) this else new Set3(x0, x1, x)
     def - [E >: Elem <: AnyRef](x: E): SimpleIdentitySet[Elem] =
       if (x `eq` x0) new Set1(x1)
       else if (x `eq` x1) new Set1(x0)
@@ -61,6 +56,29 @@ object SimpleIdentitySet {
     def foreach(f: Elem => Unit): Unit = { f(x0.asInstanceOf[Elem]); f(x1.asInstanceOf[Elem]) }
   }
 
+  private class Set3[+Elem <: AnyRef](x0: AnyRef, x1: AnyRef, x2: AnyRef) extends SimpleIdentitySet[Elem] {
+    def size = 3
+    def + [E >: Elem <: AnyRef](x: E): SimpleIdentitySet[E] =
+      if (contains(this)) this
+      else {
+        val xs = new Array[AnyRef](4)
+        xs(0) = x0
+        xs(1) = x1
+        xs(2) = x2
+        xs(3) = x
+        new SetN[E](xs)
+      }
+    def - [E >: Elem <: AnyRef](x: E): SimpleIdentitySet[Elem] =
+      if (x `eq` x0) new Set2(x1, x2)
+      else if (x `eq` x1) new Set2(x0, x2)
+      else if (x `eq` x2) new Set2(x0, x1)
+      else this
+    def contains[E >: Elem <: AnyRef](x: E): Boolean = (x `eq` x0) || (x `eq` x1) || (x `eq` x2)
+    def foreach(f: Elem => Unit): Unit = {
+      f(x0.asInstanceOf[Elem]); f(x1.asInstanceOf[Elem]); f(x2.asInstanceOf[Elem])
+    }
+  }
+
   private class SetN[+Elem <: AnyRef](xs: Array[AnyRef]) extends SimpleIdentitySet[Elem] {
     def size = xs.length
     def + [E >: Elem <: AnyRef](x: E): SimpleIdentitySet[E] = {
@@ -73,10 +91,11 @@ object SimpleIdentitySet {
       var i = 0
       while (i < size && (xs(i) `ne` x)) i += 1
       if (i == size) this
-      else if (size == 3)
-        if (i == 0) new Set2(xs(1), xs(2))
-        else if (i == 1) new Set2(xs(0), xs(2))
-        else new Set2(xs(0), xs(1))
+      else if (size == 4)
+        if (i == 0) new Set3(xs(1), xs(2), xs(3))
+        else if (i == 1) new Set3(xs(0), xs(2), xs(3))
+        else if (i == 2) new Set3(xs(0), xs(1), xs(3))
+        else new Set3(xs(0), xs(1), xs(2))
       else {
         val xs1 = new Array[AnyRef](size - 1)
         Array.copy(xs, 0, xs1, 0, i)