Add custom serialiser for NonEmptySet

core/src/main/kotlin/com/r3corda/core/serialization/Kryo.kt

@@ -16,6 +16,8 @@ import com.r3corda.core.crypto.generateKeyPair
 import com.r3corda.core.crypto.sha256
 import com.r3corda.core.node.AttachmentsClassLoader
 import com.r3corda.core.node.services.AttachmentStorage
+import com.r3corda.core.utilities.NonEmptySet
+import com.r3corda.core.utilities.NonEmptySetSerializer
 import de.javakaffee.kryoserializers.ArraysAsListSerializer
 import net.i2p.crypto.eddsa.EdDSAPrivateKey
 import net.i2p.crypto.eddsa.EdDSAPublicKey
@@ -350,6 +352,9 @@ fun createKryo(k: Kryo = Kryo()): Kryo {
         register(Issued::class.java, ImmutableClassSerializer(Issued::class))
         register(TransactionState::class.java, ImmutableClassSerializer(TransactionState::class))
 
+        // This ensures a NonEmptySetSerializer is constructed with an initial value.
+        register(NonEmptySet::class.java, NonEmptySetSerializer)
+        
         noReferencesWithin<WireTransaction>()
     }
 }

core/src/main/kotlin/com/r3corda/core/utilities/NonEmptySet.kt

@@ -1,12 +1,20 @@
 package com.r3corda.core.utilities
 
+import com.esotericsoftware.kryo.Kryo
+import com.esotericsoftware.kryo.Serializer
+import com.esotericsoftware.kryo.io.Input
+import com.esotericsoftware.kryo.io.Output
+import java.util.*
+
 /**
  * A set which is constrained to ensure it can never be empty. An initial value must be provided at
  * construction, and attempting to remove the last element will cause an IllegalStateException.
+ * The underlying set is exposed for Kryo to access, but should not be accessed directly.
  */
-class NonEmptySet<T>(initial: T, private val set: MutableSet<T> = mutableSetOf()) : MutableSet<T> {
+class NonEmptySet<T>(initial: T) : MutableSet<T> {
+    private val set: MutableSet<T> = HashSet<T>()
+
     init {
-        require (set.isEmpty()) { "Provided set must be empty." }
         set.add(initial)
     }
 
@@ -80,4 +88,28 @@ fun <T> nonEmptySetOf(initial: T, vararg elements: T): NonEmptySet<T> {
     // We add the first element twice, but it's a set, so who cares
     set.addAll(elements)
     return set
+}
+
+/**
+ * Custom serializer which understands it has to read in an item before
+ * trying to construct the set.
+ */
+object NonEmptySetSerializer : Serializer<NonEmptySet<Any>>() {
+    override fun write(kryo: Kryo, output: Output, obj: NonEmptySet<Any>) {
+        // Write out the contents as normal
+        output.writeInt(obj.size)
+        obj.forEach { kryo.writeClassAndObject(output, it) }
+    }
+
+    override fun read(kryo: Kryo, input: Input, type: Class<NonEmptySet<Any>>): NonEmptySet<Any> {
+        val size = input.readInt()
+        require(size >= 1) { "Size is positive" }
+        // TODO: Is there an upper limit we can apply to how big one of these could be?
+        val first = kryo.readClassAndObject(input)
+        // Read the first item and use it to construct the NonEmptySet
+        val set = NonEmptySet(first)
+        // Read in the rest of the set
+        for (i in 2..size) { set.add(kryo.readClassAndObject(input)) }
+        return set
+    }
 }

core/src/test/kotlin/com/r3corda/core/utilities/NonEmptySetTest.kt

@@ -8,6 +8,8 @@ import com.google.common.collect.testing.testers.CollectionAddAllTester
 import com.google.common.collect.testing.testers.CollectionClearTester
 import com.google.common.collect.testing.testers.CollectionRemoveAllTester
 import com.google.common.collect.testing.testers.CollectionRetainAllTester
+import com.r3corda.core.serialization.deserialize
+import com.r3corda.core.serialization.serialize
 import junit.framework.TestSuite
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -17,7 +19,8 @@ import kotlin.test.assertEquals
 @RunWith(Suite::class)
 @Suite.SuiteClasses(
         NonEmptySetTest.Guava::class,
-        NonEmptySetTest.Remove::class
+        NonEmptySetTest.Remove::class,
+        NonEmptySetTest.Serializer::class
 )
 class NonEmptySetTest {
     /**
@@ -93,6 +96,20 @@ class NonEmptySetTest {
             }
         }
     }
+
+    /**
+     * Test serialization/deserialization.
+     */
+    class Serializer {
+        @Test
+        fun `serialize deserialize`() {
+            val expected: NonEmptySet<Int> = nonEmptySetOf(-17, 22, 17)
+            val serialized = expected.serialize().bits
+            val actual = serialized.deserialize<NonEmptySet<Int>>()
+
+            assertEquals(expected, actual)
+        }
+    }
 }
 
 /**