Tighten type of getAllEquivalenceClasses to use immutable collections…

…, and use a Multimap internally which makes the logic a little easier to follow.

http://errorprone.info/bugpattern/MutableMethodReturnType

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=192208829

src/com/google/javascript/jscomp/graph/StandardUnionFind.java

@@ -17,16 +17,18 @@
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.Iterators.filter;
+import static com.google.common.collect.Multimaps.asMap;
 
 import com.google.common.annotations.GwtCompatible;
 import com.google.common.base.Objects;
 import com.google.common.base.Predicate;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.MultimapBuilder;
+import com.google.common.collect.SetMultimap;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.io.Serializable;
 import java.util.AbstractSet;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
@@ -119,17 +121,15 @@ public Set<E> elements() {
   }
 
   @Override
-  public Collection<Set<E>> allEquivalenceClasses() {
-    Map<Node<E>, ImmutableSet.Builder<E>> groupsTmp = new LinkedHashMap<>();
+  public ImmutableList<ImmutableSet<E>> allEquivalenceClasses() {
+    SetMultimap<Node<E>, E> groupsTmp =
+        MultimapBuilder.linkedHashKeys().linkedHashSetValues().build();
     for (Node<E> elem : elmap.values()) {
-      Node<E> root = findRoot(elem);
-      ImmutableSet.Builder<E> builder =
-          groupsTmp.computeIfAbsent(root, (Node<E> k) -> ImmutableSet.builder());
-      builder.add(elem.element);
+      groupsTmp.put(findRoot(elem), elem.element);
     }
-    ImmutableList.Builder<Set<E>> result = ImmutableList.builder();
-    for (ImmutableSet.Builder<E> group : groupsTmp.values()) {
-      result.add(group.build());
+    ImmutableList.Builder<ImmutableSet<E>> result = ImmutableList.builder();
+    for (Set<E> group : asMap(groupsTmp).values()) {
+      result.add(ImmutableSet.copyOf(group));
     }
     return result.build();
   }

src/com/google/javascript/jscomp/graph/UnionFind.java

@@ -16,8 +16,9 @@
 package com.google.javascript.jscomp.graph;
 
 import com.google.common.annotations.GwtCompatible;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
-import java.util.Collection;
 import java.util.Set;
 
 /**
@@ -70,11 +71,11 @@ public interface UnionFind<E> {
   public Set<E> elements();
 
   /**
-   * Returns an immutable collection containing all equivalence classes.  The
-   * returned collection represents a snapshot of the current state and will not
-   * reflect changes made to the data structure.
+   * Returns an immutable collection containing all equivalence classes. The returned collection
+   * represents a snapshot of the current state and will not reflect changes made to the data
+   * structure.
    */
-  public Collection<Set<E>> allEquivalenceClasses();
+  public ImmutableList<ImmutableSet<E>> allEquivalenceClasses();
 
   /**
    * Returns the elements in the same equivalence class as {@code value}.

test/com/google/javascript/jscomp/graph/StandardUnionFindTest.java

@@ -18,13 +18,11 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
-
-import junit.framework.TestCase;
-
-import java.util.Collection;
 import java.util.Iterator;
 import java.util.Set;
+import junit.framework.TestCase;
 
 /**
  * Unit test for the {@link StandardUnionFind} data structure.
@@ -93,7 +91,7 @@ public void testAllEquivalenceClasses() {
     union.union("D", "E");
     union.union("F", "F");
 
-    Collection<Set<String>> classes = union.allEquivalenceClasses();
+    ImmutableList<ImmutableSet<String>> classes = union.allEquivalenceClasses();
     assertThat(classes).containsExactly(
         ImmutableSet.of("A", "B", "C"), ImmutableSet.of("D", "E"), ImmutableSet.of("F"));
   }