From fb396ee2c48885091d7517d2c6d571931dfef27e Mon Sep 17 00:00:00 2001
From: Aaron Mahan <aaron@forerunnergames.com>
Date: Wed, 5 Oct 2016 01:37:39 -0800
Subject: [PATCH] Improve ImmutableMultimapSerializer.

- Add support for serializing ImmutableSetMultimap, which is missing.

- Add explicit support for serializing ImmutableListMultimap
  (ImmutableMap delegates to ImmutableListMultimap in most cases, so it
  was implicitly supported).

- Add tests for ImmutableSetMultimap & ImmutableListMultimap, and
  improve existing tests.
---
 .../guava/ImmutableMultimapSerializer.java    |  54 +++++--
 .../ImmutableMultimapSerializerTest.java      | 135 ++++++++++++++++--
 2 files changed, 164 insertions(+), 25 deletions(-)

diff --git a/src/main/java/de/javakaffee/kryoserializers/guava/ImmutableMultimapSerializer.java b/src/main/java/de/javakaffee/kryoserializers/guava/ImmutableMultimapSerializer.java
index 08cfbc24..013f14a1 100644
--- a/src/main/java/de/javakaffee/kryoserializers/guava/ImmutableMultimapSerializer.java
+++ b/src/main/java/de/javakaffee/kryoserializers/guava/ImmutableMultimapSerializer.java
@@ -4,9 +4,13 @@
 import com.esotericsoftware.kryo.Serializer;
 import com.esotericsoftware.kryo.io.Input;
 import com.esotericsoftware.kryo.io.Output;
+
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSetMultimap;
 
 import java.util.HashMap;
 import java.util.List;
@@ -32,10 +36,20 @@ public void write(Kryo kryo, Output output, ImmutableMultimap<Object, Object> im
 
     @Override
     public ImmutableMultimap<Object, Object> read(Kryo kryo, Input input, Class<ImmutableMultimap<Object, Object>> type) {
-        Map map = kryo.readObject(input, ImmutableMap.class);
+        final Map map = kryo.readObject(input, ImmutableMap.class);
+        final Set<Map.Entry<Object, List<? extends Object>>> entries = map.entrySet();
+
+        final ImmutableMultimap.Builder builder;
+        if (type.equals (ImmutableListMultimap.class)) {
+            builder = ImmutableMultimap.builder();
+        }
+        else if (type.equals (ImmutableSetMultimap.class)) {
+            builder = ImmutableSetMultimap.builder();
+        }
+        else {
+            builder = ImmutableMultimap.builder();
+        }
 
-        Set<Map.Entry<Object, List<? extends Object>>> entries = map.entrySet();
-        ImmutableMultimap.Builder<Object, Object> builder = ImmutableMultimap.builder();
         for (Map.Entry<Object, List<? extends Object>> entry : entries) {
             builder.putAll(entry.getKey(), entry.getValue());
         }
@@ -50,24 +64,42 @@ public ImmutableMultimap<Object, Object> read(Kryo kryo, Input input, Class<Immu
      * @param kryo the {@link Kryo} instance to set the serializer on
      */
     public static void registerSerializers(final Kryo kryo) {
+        // ImmutableMap is used by ImmutableMultimap. However,
+        // we already have a separate serializer class for ImmutableMap,
+        // ImmutableMapSerializer. If it is not already being used, register it.
+        Serializer immutableMapSerializer = kryo.getSerializer(ImmutableMap.class);
+        if (!(immutableMapSerializer instanceof ImmutableMapSerializer)) {
+            ImmutableMapSerializer.registerSerializers(kryo);
+        }
 
+        // ImmutableList is used by ImmutableListMultimap. However,
+        // we already have a separate serializer class for ImmutableList,
+        // ImmutableListSerializer. If it is not already being used, register it.
         Serializer immutableListSerializer = kryo.getSerializer(ImmutableList.class);
         if (!(immutableListSerializer instanceof ImmutableListSerializer)) {
             ImmutableListSerializer.registerSerializers(kryo);
         }
 
-        Serializer immutableMapSerializer = kryo.getSerializer(ImmutableMap.class);
-        if (!(immutableMapSerializer instanceof ImmutableMapSerializer)) {
-            ImmutableMapSerializer.registerSerializers(kryo);
+        // ImmutableSet is used by ImmutableSetMultimap. However,
+        // we already have a separate serializer class for ImmutableSet,
+        // ImmutableSetSerializer. If it is not already being used, register it.
+        Serializer immutableSetSerializer = kryo.getSerializer(ImmutableSet.class);
+        if (!(immutableSetSerializer instanceof ImmutableSetSerializer)) {
+            ImmutableSetSerializer.registerSerializers(kryo);
         }
 
         final ImmutableMultimapSerializer serializer = new ImmutableMultimapSerializer();
 
-        kryo.register(ImmutableMultimap.class, serializer);
-        kryo.register(ImmutableMultimap.of().getClass(), serializer);
-        Object o = new Object();
-
-        kryo.register(ImmutableMultimap.of(o, o).getClass(), serializer);
+        // ImmutableMultimap (abstract class)
+        //  +- EmptyImmutableListMultimap
+        //  +- ImmutableListMultimap
+        //  +- EmptyImmutableSetMultimap
+        //  +- ImmutableSetMultimap
 
+        kryo.register(ImmutableMultimap.class, serializer);
+        kryo.register(ImmutableListMultimap.of().getClass(), serializer);
+        kryo.register(ImmutableListMultimap.of("A", "B").getClass(), serializer);
+        kryo.register(ImmutableSetMultimap.of().getClass(), serializer);
+        kryo.register(ImmutableSetMultimap.of("A", "B").getClass(), serializer);
     }
 }
diff --git a/src/test/java/de/javakaffee/kryoserializers/guava/ImmutableMultimapSerializerTest.java b/src/test/java/de/javakaffee/kryoserializers/guava/ImmutableMultimapSerializerTest.java
index 3a41ec30..2612e94f 100644
--- a/src/test/java/de/javakaffee/kryoserializers/guava/ImmutableMultimapSerializerTest.java
+++ b/src/test/java/de/javakaffee/kryoserializers/guava/ImmutableMultimapSerializerTest.java
@@ -1,10 +1,12 @@
 package de.javakaffee.kryoserializers.guava;
 
 import com.esotericsoftware.kryo.Kryo;
+import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMultimap;
 
 import static org.testng.Assert.*;
 
+import com.google.common.collect.ImmutableSetMultimap;
 import org.testng.annotations.BeforeTest;
 import org.testng.annotations.Test;
 
@@ -23,43 +25,57 @@ public void setUp() throws Exception {
     }
 
     @Test
-    public void testEmpty() {
+    public void testRegularEmpty() {
         final ImmutableMultimap<?, ?> obj = ImmutableMultimap.of();
         final byte[] serialized = serialize(_kryo, obj);
         final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableMultimap.class);
         assertTrue(deserialized.isEmpty());
-        assertEquals(deserialized.size(), obj.size());
+        assertEquals(deserialized, obj);
     }
 
     @Test
-    public void testImmutableListSerializerAlreadyRegistered() {
-        ImmutableListSerializer.registerSerializers(_kryo);
-        final ImmutableMultimap<?, ?> obj = ImmutableMultimap.of();
+    public void testImmutableListMultimapEmpty() {
+        final ImmutableMultimap<?, ?> obj = ImmutableListMultimap.of();
         final byte[] serialized = serialize(_kryo, obj);
-        final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableMultimap.class);
+        final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableListMultimap.class);
         assertTrue(deserialized.isEmpty());
-        assertEquals(deserialized.size(), obj.size());
+        assertEquals(deserialized, obj);
     }
 
     @Test
-    public void testImmutableMapSerializerAlreadyRegistered() {
-        ImmutableMapSerializer.registerSerializers(_kryo);
-        final ImmutableMultimap<?, ?> obj = ImmutableMultimap.of();
+    public void testImmutableSetMultimapEmpty() {
+        final ImmutableMultimap<?, ?> obj = ImmutableSetMultimap.of();
         final byte[] serialized = serialize(_kryo, obj);
-        final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableMultimap.class);
+        final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableSetMultimap.class);
         assertTrue(deserialized.isEmpty());
-        assertEquals(deserialized.size(), obj.size());
+        assertEquals(deserialized, obj);
     }
 
-
     @Test
-    public void testSingleton() {
+    public void testRegularSingleton() {
         final ImmutableMultimap<?, ?> obj = ImmutableMultimap.of(3, "k");
         final byte[] serialized = serialize(_kryo, obj);
         final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableMultimap.class);
         assertEquals(deserialized, obj);
     }
 
+    @Test
+    public void testImmutableListMultimapSingleton() {
+        final ImmutableMultimap<?, ?> obj = ImmutableListMultimap.of(3, "k");
+        final byte[] serialized = serialize(_kryo, obj);
+        final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableListMultimap.class);
+        assertEquals(deserialized, obj);
+    }
+
+    @Test
+    public void testImmutableSetMultimapSingleton() {
+        final ImmutableMultimap<?, ?> obj = ImmutableSetMultimap.of(3, "k");
+        final byte[] serialized = serialize(_kryo, obj);
+        final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableSetMultimap.class);
+        assertEquals(deserialized.getClass(), obj.getClass());
+        assertEquals(deserialized, obj);
+    }
+
     @Test
     public void testRegular() {
         final ImmutableMultimap<?, ?> obj = ImmutableMultimap.of(3, "k", 5, "r", 6, "y");
@@ -68,6 +84,23 @@ public void testRegular() {
         assertEquals(deserialized, obj);
     }
 
+    @Test
+    public void testImmutableListMultimap() {
+        final ImmutableMultimap<?, ?> obj = ImmutableListMultimap.of(3, "k", 5, "r", 6, "y");
+        final byte[] serialized = serialize(_kryo, obj);
+        final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableListMultimap.class);
+        assertEquals(deserialized, obj);
+    }
+
+    @Test
+    public void testImmutableSetMultimap() {
+        final ImmutableMultimap<?, ?> obj = ImmutableSetMultimap.of(3, "k", 5, "r", 6, "y");
+        final byte[] serialized = serialize(_kryo, obj);
+        final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableSetMultimap.class);
+        assertEquals(deserialized.getClass(), obj.getClass());
+        assertEquals(deserialized, obj);
+    }
+
     @Test
     public void testRegularMultipleElementsPerKey() {
         final ImmutableMultimap<?, ?> obj = ImmutableMultimap.of(3, "k", 3, "r", 4, "y", 4, "o");
@@ -76,6 +109,52 @@ public void testRegularMultipleElementsPerKey() {
         assertEquals(deserialized, obj);
     }
 
+    @Test
+    public void testImmutableListMultimapMultipleElementsPerKey() {
+        final ImmutableMultimap<?, ?> obj = ImmutableListMultimap.of(3, "k", 3, "r", 4, "y", 4, "o");
+        final byte[] serialized = serialize(_kryo, obj);
+        final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableListMultimap.class);
+        assertEquals(deserialized, obj);
+    }
+
+    @Test
+    public void testImmutableSetMultimapMultipleElementsPerKey() {
+        final ImmutableMultimap<?, ?> obj = ImmutableSetMultimap.of(3, "k", 3, "r", 4, "y", 4, "o");
+        final byte[] serialized = serialize(_kryo, obj);
+        final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableSetMultimap.class);
+        assertEquals(deserialized, obj);
+    }
+
+    @Test
+    public void testImmutableMapSerializerAlreadyRegistered() {
+        ImmutableMapSerializer.registerSerializers(_kryo);
+        final ImmutableMultimap<?, ?> obj = ImmutableMultimap.of();
+        final byte[] serialized = serialize(_kryo, obj);
+        final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableMultimap.class);
+        assertTrue(deserialized.isEmpty());
+        assertEquals(deserialized, obj);
+    }
+
+    @Test
+    public void testImmutableListSerializerAlreadyRegistered() {
+        ImmutableListSerializer.registerSerializers(_kryo);
+        final ImmutableMultimap<?, ?> obj = ImmutableListMultimap.of();
+        final byte[] serialized = serialize(_kryo, obj);
+        final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableListMultimap.class);
+        assertTrue(deserialized.isEmpty());
+        assertEquals(deserialized, obj);
+    }
+
+    @Test
+    public void testImmutableSetSerializerAlreadyRegistered() {
+        ImmutableSetSerializer.registerSerializers(_kryo);
+        final ImmutableMultimap<?, ?> obj = ImmutableSetMultimap.of();
+        final byte[] serialized = serialize(_kryo, obj);
+        final ImmutableMultimap<?, ?> deserialized = deserialize(_kryo, serialized, ImmutableSetMultimap.class);
+        assertTrue(deserialized.isEmpty());
+        assertEquals(deserialized, obj);
+    }
+
     // Kryo#copy tests
 
     @Test
@@ -98,4 +177,32 @@ public void testCopyRegular() {
         final ImmutableMultimap<?, ?> copied = _kryo.copy(obj);
         assertSame(copied, obj);
     }
+
+    @Test
+    public void testCopyImmutableListMultimap() {
+        final ImmutableMultimap<?, ?> obj = ImmutableListMultimap.of(1, "k", 2, "r", 3, "y");
+        final ImmutableMultimap<?, ?> copied = _kryo.copy(obj);
+        assertSame(copied, obj);
+    }
+
+    @Test
+    public void testCopyImmutableListMultimapEmpty() {
+        final ImmutableMultimap<?, ?> obj = ImmutableListMultimap.of();
+        final ImmutableMultimap<?, ?> copied = _kryo.copy(obj);
+        assertSame(copied, obj);
+    }
+
+    @Test
+    public void testCopyImmutableSetMultimap() {
+        final ImmutableMultimap<?, ?> obj = ImmutableSetMultimap.of(1, "k", 2, "r", 3, "y");
+        final ImmutableMultimap<?, ?> copied = _kryo.copy(obj);
+        assertSame(copied, obj);
+    }
+
+    @Test
+    public void testCopyImmutableSetMultimapEmpty() {
+        final ImmutableMultimap<?, ?> obj = ImmutableSetMultimap.of();
+        final ImmutableMultimap<?, ?> copied = _kryo.copy(obj);
+        assertSame(copied, obj);
+    }
 }
\ No newline at end of file