add size method to NinePMessage

1 files changed, 19 insertions, 5 deletions

diff --git a/src/main/kotlin/NinePMessage.kt b/src/main/kotlin/NinePMessage.kt
index 9edf361..1d0530a 100644
--- a/src/main/kotlin/NinePMessage.kt
+++ b/src/main/kotlin/NinePMessage.kt
@@ -15,20 +15,27 @@ import kotlin.math.pow
  */
 class NinePMessage(val type: NinePMessageType, val tag: UShort, val fieldNames: List<String>, val fieldValuesInt: Map<String, SizedMessageField>, val fieldValuesStr: Map<String, String>) {
     /**
+     * Intersection between [fieldNames] and [fieldValuesInt]. In other words: the amount of integer fields that are
+     * going to be used when writing the message.
+     */
+    private val insecInts = fieldNames.intersect(fieldValuesInt.keys)
+    /**
+     * Intersection between [fieldNames] and [fieldValuesStr]. In other words: the amount of string fields that are
+     * going to be used when writing the message.
+     */
+    private val insecStrs = fieldNames.intersect(fieldValuesStr.keys)
+
+    /**
      * Send the message using the given networking API.
      *
      * @param npt The networking API.
      * @throws IllegalArgumentException if [fieldNames], [fieldValuesInt], and [fieldValuesStr] are incoherent.
      */
     fun write(npt: NetworkPacketTransporter) {
-        // shorthands
-        val insecInts = fieldNames.intersect(fieldValuesInt.keys)
-        val insecStrs = fieldNames.intersect(fieldValuesStr.keys)
         // check that names in fieldNames exist as keys in either fieldValuesInt or fieldValuesStr but not both
         require(insecInts.size == fieldNames.size - insecStrs.size)
 
-        val totalSize = 4 + 1 + 2 + insecInts.sumOf { fieldValuesInt[it]!!.size } + insecStrs.sumOf { 2 + fieldValuesStr[it]!!.length }
-
+        val totalSize = size()
         writeMessageSizeType(npt, totalSize, type)
         writeInteger(npt, SizedMessageField(2, tag.toInt().toBigInteger()))
         for (field in fieldNames) {
@@ -79,4 +86,11 @@ class NinePMessage(val type: NinePMessageType, val tag: UShort, val fieldNames:
         writeInteger(npt, SizedMessageField(4, size.toBigInteger()))
         writeInteger(npt, SizedMessageField(1, type.value.toInt().toBigInteger()))
     }
+
+    /**
+     * Calculate the expected size of the message.
+     */
+    fun size(): Int {
+        return 4 + 1 + 2 + this.insecInts.sumOf { this.fieldValuesInt[it]!!.size } + this.insecStrs.sumOf { 2 + this.fieldValuesStr[it]!!.length }
+    }
 }
\ No newline at end of file