Add TreeSetSerializer, closes #139 (Kryo gets ClassCastException when deserializing TreeSet with Comparator)

git-svn-id: http://kryo.googlecode.com/svn/trunk@416 a9660ef0-a895-11de-8ba5-91705c1537ed

Author: [email protected]

src/com/esotericsoftware/kryo/Kryo.java

@@ -18,6 +18,7 @@
 import java.util.Map;
 import java.util.TimeZone;
 import java.util.TreeMap;
+import java.util.TreeSet;
 
 import org.objenesis.instantiator.ObjectInstantiator;
 import org.objenesis.strategy.InstantiatorStrategy;
@@ -68,6 +69,7 @@
 import com.esotericsoftware.kryo.serializers.DefaultSerializers.StringSerializer;
 import com.esotericsoftware.kryo.serializers.DefaultSerializers.TimeZoneSerializer;
 import com.esotericsoftware.kryo.serializers.DefaultSerializers.TreeMapSerializer;
+import com.esotericsoftware.kryo.serializers.DefaultSerializers.TreeSetSerializer;
 import com.esotericsoftware.kryo.serializers.DefaultSerializers.VoidSerializer;
 import com.esotericsoftware.kryo.serializers.FieldSerializer;
 import com.esotericsoftware.kryo.serializers.MapSerializer;
@@ -180,6 +182,7 @@ public Kryo (ClassResolver classResolver, ReferenceResolver referenceResolver, S
 		addDefaultSerializer(Collections.singletonList(null).getClass(), CollectionsSingletonListSerializer.class);
 		addDefaultSerializer(Collections.singletonMap(null, null).getClass(), CollectionsSingletonMapSerializer.class);
 		addDefaultSerializer(Collections.singleton(null).getClass(), CollectionsSingletonSetSerializer.class);
+		addDefaultSerializer(TreeSet.class, TreeSetSerializer.class);
 		addDefaultSerializer(Collection.class, CollectionSerializer.class);
 		addDefaultSerializer(TreeMap.class, TreeMapSerializer.class);
 		addDefaultSerializer(Map.class, MapSerializer.class);

src/com/esotericsoftware/kryo/serializers/DefaultSerializers.java

@@ -4,6 +4,7 @@
 import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.util.Calendar;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.Currency;
@@ -16,6 +17,7 @@
 import java.util.Set;
 import java.util.TimeZone;
 import java.util.TreeMap;
+import java.util.TreeSet;
 
 import com.esotericsoftware.kryo.Kryo;
 import com.esotericsoftware.kryo.KryoException;
@@ -556,4 +558,20 @@ protected Map createCopy (Kryo kryo, Map original) {
 			return new TreeMap(((TreeMap)original).comparator());
 		}
 	}
+
+	static public class TreeSetSerializer extends CollectionSerializer {
+		public void write (Kryo kryo, Output output, Collection collection) {
+			TreeSet treeSet = (TreeSet)collection;
+			kryo.writeClassAndObject(output, treeSet.comparator());
+			super.write(kryo, output, collection);
+		}
+
+		protected TreeSet create (Kryo kryo, Input input, Class<Collection> type) {
+			return new TreeSet((Comparator)kryo.readClassAndObject(input));
+		}
+
+		protected TreeSet createCopy (Kryo kryo, Collection original) {
+			return new TreeSet(((TreeSet)original).comparator());
+		}
+	}
 }

test/com/esotericsoftware/kryo/CollectionSerializerTest.java

@@ -3,8 +3,12 @@
 
 import java.util.ArrayList;
 import java.util.LinkedList;
+import java.util.TreeMap;
+import java.util.TreeSet;
 import java.util.concurrent.CopyOnWriteArrayList;
 
+import com.esotericsoftware.kryo.MapSerializerTest.KeyComparator;
+import com.esotericsoftware.kryo.MapSerializerTest.KeyThatIsntComparable;
 import com.esotericsoftware.kryo.serializers.CollectionSerializer;
 import com.esotericsoftware.kryo.serializers.DefaultSerializers.StringSerializer;
 
@@ -37,5 +41,19 @@ public void testCollections () {
 		roundTrip(8, 8, list("1", "2", "3"));
 		serializer.setElementsCanBeNull(false);
 		roundTrip(8, 8, list("1", "2", "3"));
+
+		kryo.register(TreeSet.class);
+		TreeSet set = new TreeSet();
+		set.add("1");
+		set.add("2");
+		roundTrip(9, 9, set);
+
+		kryo.register(KeyThatIsntComparable.class);
+		kryo.register(KeyComparator.class);
+		set = new TreeSet(new KeyComparator());
+		set.add(new KeyThatIsntComparable("1"));
+		set.add(new KeyThatIsntComparable("2"));
+		roundTrip(9, 9, set);
 	}
+
 }

test/com/esotericsoftware/kryo/MapSerializerTest.java

@@ -147,5 +147,10 @@ public int compare (KeyThatIsntComparable o1, KeyThatIsntComparable o2) {
 
 	static public class KeyThatIsntComparable {
 		public String value;
+		public KeyThatIsntComparable () {
+		}
+		public KeyThatIsntComparable (String value) {
+			this.value = value;
+		}
 	}
 }