Ensure that the set returned by ImmutableMap<K, V>.keySet() is serializable when K is serializable, and similarly for values().

Set<T> should be serializable when T is serializable but that is not always the case for the set returned by ImmutableMap.keySet() due to a reference from the returned set back to the original map. When serializing this set, the original map is serialized is well. This change changes this so that only the keys are serialized.

RELNOTES=`collect`: Ensure that the set returned by `ImmutableMap<K, V>.keySet()` is serializable when `K` is serializable (and similarly for `values()`).

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

Author: adob

android/guava-tests/test/com/google/common/collect/ImmutableMapTest.java

@@ -17,6 +17,7 @@
 package com.google.common.collect;
 
 import static com.google.common.testing.SerializableTester.reserialize;
+import static com.google.common.truth.Truth.assertThat;
 
 import com.google.common.annotations.GwtCompatible;
 import com.google.common.annotations.GwtIncompatible;
@@ -45,12 +46,15 @@
 import com.google.common.testing.EqualsTester;
 import com.google.common.testing.NullPointerTester;
 import com.google.common.testing.SerializableTester;
+import java.io.ByteArrayOutputStream;
+import java.io.ObjectOutputStream;
 import java.io.Serializable;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Set;
 import junit.framework.Test;
 import junit.framework.TestCase;
 import junit.framework.TestSuite;
@@ -768,6 +772,58 @@ public void testViewSerialization() {
     assertTrue(reserializedValues instanceof ImmutableCollection);
   }
 
+  @GwtIncompatible // SerializableTester
+  public void testKeySetIsSerializable_regularImmutableMap() {
+    class NonSerializableClass {}
+
+    Map<String, NonSerializableClass> map =
+        RegularImmutableMap.create(1, new Object[] {"one", new NonSerializableClass()});
+    Set<String> set = map.keySet();
+
+    LenientSerializableTester.reserializeAndAssertLenient(set);
+  }
+
+  @GwtIncompatible // SerializableTester
+  public void testValuesCollectionIsSerializable_regularImmutableMap() {
+    class NonSerializableClass {}
+
+    Map<NonSerializableClass, String> map =
+        RegularImmutableMap.create(1, new Object[] {new NonSerializableClass(), "value"});
+    Collection<String> collection = map.values();
+
+    LenientSerializableTester.reserializeAndAssertElementsEqual(collection);
+  }
+
+  // TODO: Re-enable this test after moving to new serialization format in ImmutableMap.
+  @GwtIncompatible // SerializableTester
+  @SuppressWarnings("unchecked")
+  public void ignore_testSerializationNoDuplication_regularImmutableMap() throws Exception {
+    // Tests that searializing a map, its keySet, and values only writes the underlying data once.
+
+    Object[] entries = new Object[2000];
+    for (int i = 0; i < entries.length; i++) {
+      entries[i] = i;
+    }
+
+    ImmutableMap<Integer, Integer> map = RegularImmutableMap.create(entries.length / 2, entries);
+    Set<Integer> keySet = map.keySet();
+    Collection<Integer> values = map.values();
+
+    ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+    ObjectOutputStream oos = new ObjectOutputStream(bytes);
+    oos.writeObject(map);
+    oos.flush();
+
+    int mapSize = bytes.size();
+    oos.writeObject(keySet);
+    oos.writeObject(values);
+    oos.close();
+
+    int finalSize = bytes.size();
+
+    assertThat(finalSize - mapSize).isLessThan(100);
+  }
+
   public void testEquals() {
     new EqualsTester()
         .addEqualityGroup(ImmutableMap.of(), ImmutableMap.builder().build())

android/guava-tests/test/com/google/common/collect/LenientSerializableTester.java

@@ -24,6 +24,7 @@
 import com.google.common.annotations.GwtIncompatible;
 import com.google.common.testing.SerializableTester;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import java.util.Collection;
 import java.util.Set;
 
 /**
@@ -60,5 +61,14 @@ static <E> Multiset<E> reserializeAndAssertLenient(Multiset<E> original) {
     return copy;
   }
 
+  @CanIgnoreReturnValue
+  @GwtIncompatible // SerializableTester
+  static <E> Collection<E> reserializeAndAssertElementsEqual(Collection<E> original) {
+    Collection<E> copy = reserialize(original);
+    assertTrue(Iterables.elementsEqual(original, copy));
+    assertTrue(copy instanceof ImmutableCollection);
+    return copy;
+  }
+
   private LenientSerializableTester() {}
 }

android/guava/src/com/google/common/collect/ImmutableBiMap.java

@@ -351,22 +351,21 @@ public V forcePut(K key, V value) {
    * <p>Since the bimap is immutable, ImmutableBiMap doesn't require special logic for keeping the
    * bimap and its inverse in sync during serialization, the way AbstractBiMap does.
    */
-  private static class SerializedForm extends ImmutableMap.SerializedForm {
-    SerializedForm(ImmutableBiMap<?, ?> bimap) {
+  private static class SerializedForm<K, V> extends ImmutableMap.SerializedForm<K, V> {
+    SerializedForm(ImmutableBiMap<K, V> bimap) {
       super(bimap);
     }
 
     @Override
-    Object readResolve() {
-      Builder<Object, Object> builder = new Builder<>();
-      return createMap(builder);
+    Builder<K, V> makeBuilder(int size) {
+      return new Builder<>(size);
     }
 
     private static final long serialVersionUID = 0;
   }
 
   @Override
   Object writeReplace() {
-    return new SerializedForm(this);
+    return new SerializedForm<>(this);
   }
 }

android/guava/src/com/google/common/collect/ImmutableMap.java

@@ -706,37 +706,85 @@ public String toString() {
    * reconstructed using public factory methods. This ensures that the implementation types remain
    * as implementation details.
    */
-  static class SerializedForm implements Serializable {
-    private final Object[] keys;
-    private final Object[] values;
-
-    SerializedForm(ImmutableMap<?, ?> map) {
-      keys = new Object[map.size()];
-      values = new Object[map.size()];
-      int i = 0;
-      for (Entry<?, ?> entry : map.entrySet()) {
-        keys[i] = entry.getKey();
-        values[i] = entry.getValue();
-        i++;
+  static class SerializedForm<K, V> implements Serializable {
+    // This object retains references to collections returned by keySet() and value(). This saves
+    // bytes when the both the map and its keySet or value collection are written to the same
+    // instance of ObjectOutputStream.
+
+    // TODO(b/160980469): remove support for the old serialization format after some time
+    private static final boolean USE_LEGACY_SERIALIZATION = true;
+
+    private final Object keys;
+    private final Object values;
+
+    SerializedForm(ImmutableMap<K, V> map) {
+      if (USE_LEGACY_SERIALIZATION) {
+        Object[] keys = new Object[map.size()];
+        Object[] values = new Object[map.size()];
+        int i = 0;
+        for (Entry<?, ?> entry : map.entrySet()) {
+          keys[i] = entry.getKey();
+          values[i] = entry.getValue();
+          i++;
+        }
+        this.keys = keys;
+        this.values = values;
+        return;
       }
+      this.keys = map.keySet();
+      this.values = map.values();
     }
 
-    Object readResolve() {
-      Builder<Object, Object> builder = new Builder<>(keys.length);
-      return createMap(builder);
+    @SuppressWarnings("unchecked")
+    final Object readResolve() {
+      if (!(this.keys instanceof ImmutableSet)) {
+        return legacyReadResolve();
+      }
+
+      ImmutableSet<K> keySet = (ImmutableSet<K>) this.keys;
+      ImmutableCollection<V> values = (ImmutableCollection<V>) this.values;
+
+      Builder<K, V> builder = makeBuilder(keySet.size());
+
+      UnmodifiableIterator<K> keyIter = keySet.iterator();
+      UnmodifiableIterator<V> valueIter = values.iterator();
+
+      while (keyIter.hasNext()) {
+        builder.put(keyIter.next(), valueIter.next());
+      }
+
+      return builder.build();
     }
 
-    Object createMap(Builder<Object, Object> builder) {
+    @SuppressWarnings("unchecked")
+    final Object legacyReadResolve() {
+      K[] keys = (K[]) this.keys;
+      V[] values = (V[]) this.values;
+
+      Builder<K, V> builder = makeBuilder(keys.length);
+
       for (int i = 0; i < keys.length; i++) {
         builder.put(keys[i], values[i]);
       }
       return builder.build();
     }
 
+    /**
+     * Returns a builder that builds the unserialized type. Subclasses should override this method.
+     */
+    Builder<K, V> makeBuilder(int size) {
+      return new Builder<>(size);
+    }
+
     private static final long serialVersionUID = 0;
   }
 
+  /**
+   * Returns a serializable form of this object. Non-public subclasses should not override this
+   * method. Publicly-accessible subclasses must override this method and should return a subclass
+   * of SerializedForm whose readResolve() method returns objects of the subclass type.
+   */
   Object writeReplace() {
-    return new SerializedForm(this);
+    return new SerializedForm<>(this);
   }
 }

android/guava/src/com/google/common/collect/ImmutableSortedMap.java

@@ -881,27 +881,25 @@ public ImmutableSortedSet<K> descendingKeySet() {
    * are reconstructed using public factory methods. This ensures that the implementation types
    * remain as implementation details.
    */
-  private static class SerializedForm extends ImmutableMap.SerializedForm {
-    private final Comparator<Object> comparator;
+  private static class SerializedForm<K, V> extends ImmutableMap.SerializedForm<K, V> {
+    private final Comparator<? super K> comparator;
 
-    @SuppressWarnings("unchecked")
-    SerializedForm(ImmutableSortedMap<?, ?> sortedMap) {
+    SerializedForm(ImmutableSortedMap<K, V> sortedMap) {
       super(sortedMap);
-      comparator = (Comparator<Object>) sortedMap.comparator();
+      comparator = sortedMap.comparator();
     }
 
     @Override
-    Object readResolve() {
-      Builder<Object, Object> builder = new Builder<>(comparator);
-      return createMap(builder);
+    Builder<K, V> makeBuilder(int size) {
+      return new Builder<>(comparator);
     }
 
     private static final long serialVersionUID = 0;
   }
 
   @Override
   Object writeReplace() {
-    return new SerializedForm(this);
+    return new SerializedForm<>(this);
   }
 
   // This class is never actually serialized directly, but we have to make the