From 2c422bebd336783ef2842f45c0bfd7efeaa2aabb Mon Sep 17 00:00:00 2001
From: Andrius Dagys <andrius.dagys@r3cev.com>
Date: Wed, 11 May 2016 14:18:09 +0100
Subject: [PATCH] Added noneOrSingle extension method, which returns a single
 element, null if no elements found and throws if more than one element found

---
 core/src/main/kotlin/core/Utils.kt            | 27 ++++++++++++
 .../core/node/services/NotaryService.kt       | 10 ++---
 .../utilities/CollectionExtensionTests.kt     | 42 +++++++++++++++++++
 3 files changed, 74 insertions(+), 5 deletions(-)
 create mode 100644 src/test/kotlin/core/utilities/CollectionExtensionTests.kt

diff --git a/core/src/main/kotlin/core/Utils.kt b/core/src/main/kotlin/core/Utils.kt
index a093c9db98..c0a09cf49d 100644
--- a/core/src/main/kotlin/core/Utils.kt
+++ b/core/src/main/kotlin/core/Utils.kt
@@ -83,6 +83,33 @@ fun <T> List<T>.indexOfOrThrow(item: T): Int {
     return i
 }
 
+/**
+ * Returns the single element matching the given [predicate], or `null` if element was not found,
+ * or throws if more than one element was found.
+ */
+fun <T> Iterable<T>.noneOrSingle(predicate: (T) -> Boolean): T? {
+    var single: T? = null
+    for (element in this) {
+        if (predicate(element)) {
+            if (single == null) {
+                single = element
+            } else throw IllegalArgumentException("Collection contains more than one matching element.")
+        }
+    }
+    return single
+}
+
+/** Returns single element, or `null` if element was not found, or throws if more than one element was found. */
+fun <T> Iterable<T>.noneOrSingle(): T? {
+    var single: T? = null
+    for (element in this) {
+        if (single == null) {
+            single = element
+        } else throw IllegalArgumentException("Collection contains more than one matching element.")
+    }
+    return single
+}
+
 // An alias that can sometimes make code clearer to read.
 val RunOnCallerThread = MoreExecutors.directExecutor()
 
diff --git a/src/main/kotlin/core/node/services/NotaryService.kt b/src/main/kotlin/core/node/services/NotaryService.kt
index 05e7461012..c34e9a6860 100644
--- a/src/main/kotlin/core/node/services/NotaryService.kt
+++ b/src/main/kotlin/core/node/services/NotaryService.kt
@@ -7,6 +7,7 @@ import core.crypto.DigitalSignature
 import core.crypto.SignedData
 import core.crypto.signWithECDSA
 import core.messaging.MessagingService
+import core.noneOrSingle
 import core.serialization.SerializedBytes
 import core.serialization.deserialize
 import core.serialization.serialize
@@ -64,12 +65,11 @@ class NotaryService(net: MessagingService,
     }
 
     private fun validateTimestamp(tx: WireTransaction) {
-        // Need to have at most one timestamp command
-        val timestampCmds = tx.commands.filter { it.value is TimestampCommand }
-        if (timestampCmds.count() > 1)
+        val timestampCmd = try {
+            tx.commands.noneOrSingle { it.value is TimestampCommand } ?: return
+        } catch (e: IllegalArgumentException) {
             throw NotaryException(NotaryError.MoreThanOneTimestamp())
-
-        val timestampCmd = timestampCmds.singleOrNull() ?: return
+        }
         if (!timestampCmd.signers.contains(identity.owningKey))
             throw NotaryException(NotaryError.NotForMe())
         if (!timestampChecker.isValid(timestampCmd.value as TimestampCommand))
diff --git a/src/test/kotlin/core/utilities/CollectionExtensionTests.kt b/src/test/kotlin/core/utilities/CollectionExtensionTests.kt
new file mode 100644
index 0000000000..f40ac1fb22
--- /dev/null
+++ b/src/test/kotlin/core/utilities/CollectionExtensionTests.kt
@@ -0,0 +1,42 @@
+package core.utilities
+
+import core.indexOfOrThrow
+import core.noneOrSingle
+import org.junit.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+
+class CollectionExtensionTests {
+    @Test
+    fun `noneOrSingle returns a single item`() {
+        val collection = listOf(1)
+        assertEquals(collection.noneOrSingle(), 1)
+        assertEquals(collection.noneOrSingle { it == 1 }, 1)
+    }
+
+    @Test
+    fun `noneOrSingle returns null if item not found`() {
+        val collection = emptyList<Int>()
+        assertEquals(collection.noneOrSingle(), null)
+    }
+
+    @Test
+    fun `noneOrSingle throws if more than one item found`() {
+        val collection = listOf(1, 2)
+        assertFailsWith<IllegalArgumentException> { collection.noneOrSingle() }
+        assertFailsWith<IllegalArgumentException> { collection.noneOrSingle { it > 0 } }
+    }
+
+    @Test
+    fun `indexOfOrThrow returns index of the given item`() {
+        val collection = listOf(1, 2)
+        assertEquals(collection.indexOfOrThrow(1), 0)
+        assertEquals(collection.indexOfOrThrow(2), 1)
+    }
+
+    @Test
+    fun `indexOfOrThrow throws if the given item is not found`() {
+        val collection = listOf(1)
+        assertFailsWith<IllegalArgumentException> { collection.indexOfOrThrow(2) }
+    }
+}
\ No newline at end of file