Improve binary json

src/main/kotlin/ru/dbotthepony/mc/otm/core/util/BinaryJson.kt

@@ -1,5 +1,6 @@
 package ru.dbotthepony.mc.otm.core.util
 
+import com.google.common.collect.ImmutableList
 import com.google.gson.JsonArray
 import com.google.gson.JsonElement
 import com.google.gson.JsonNull
@@ -8,16 +9,186 @@ import com.google.gson.JsonParseException
 import com.google.gson.JsonPrimitive
 import com.google.gson.JsonSyntaxException
 import net.minecraft.nbt.NbtAccounter
+import ru.dbotthepony.mc.otm.core.collect.stream
 import java.io.InputStream
 import java.io.OutputStream
+import java.util.function.Predicate
 
-private const val TYPE_NULL = 0x01
+private enum class BinaryElementType(private val predicate: Predicate<JsonElement>) : Predicate<JsonElement> by predicate {
-private const val TYPE_DOUBLE = 0x02
+	NULL({ it is JsonNull }) {
-private const val TYPE_BOOLEAN = 0x03
+		override fun write(stream: OutputStream, element: JsonElement) {
-private const val TYPE_INT = 0x04
+		}
-private const val TYPE_STRING = 0x05
+
-private const val TYPE_ARRAY = 0x06
+		override fun read(stream: InputStream, accounter: NbtAccounter): JsonElement {
-private const val TYPE_OBJECT = 0x07
+			return JsonNull.INSTANCE
+		}
+	},
+
+	DOUBLE({ it is JsonPrimitive && it.isNumber && (it.asNumber is Double || it.asNumber is Float) }) {
+		override fun write(stream: OutputStream, element: JsonElement) {
+			stream.writeDouble(element.asDouble)
+		}
+
+		override fun read(stream: InputStream, accounter: NbtAccounter): JsonElement {
+			accounter.accountBytes(8L)
+			return JsonPrimitive(stream.readDouble())
+		}
+	},
+
+	BOOLEAN({ it is JsonPrimitive && it.isBoolean }) {
+		override fun write(stream: OutputStream, element: JsonElement) {
+			stream.write(if (element.asBoolean) 1 else 0)
+		}
+
+		override fun read(stream: InputStream, accounter: NbtAccounter): JsonElement {
+			accounter.accountBytes(1L)
+			return JsonPrimitive(stream.read() > 0)
+		}
+	},
+
+	INT({ it is JsonPrimitive && it.isNumber }) {
+		override fun write(stream: OutputStream, element: JsonElement) {
+			val it = element.asNumber
+
+			if (it is Int || it is Long || it is Short || it is Byte)
+				stream.writeVarLongLE(element.asLong)
+			else
+				throw IllegalArgumentException("Unknown number type: ${it::class.qualifiedName}")
+		}
+
+		override fun read(stream: InputStream, accounter: NbtAccounter): JsonElement {
+			return JsonPrimitive(stream.readVarLongLE(accounter))
+		}
+	},
+
+	STRING({ it is JsonPrimitive && it.isString }) {
+		override fun write(stream: OutputStream, element: JsonElement) {
+			stream.writeBinaryString(element.asString)
+		}
+
+		override fun read(stream: InputStream, accounter: NbtAccounter): JsonElement {
+			return JsonPrimitive(stream.readBinaryString(accounter))
+		}
+	},
+
+	ARRAY({ it is JsonArray }) {
+		override fun write(stream: OutputStream, element: JsonElement) {
+			writeArray(stream, element, OutputStream::writeBinaryJson)
+		}
+
+		override fun read(stream: InputStream, accounter: NbtAccounter): JsonElement {
+			return readArray(stream, accounter, InputStream::readBinaryJson)
+		}
+	},
+
+	OBJECT({ it is JsonObject }) {
+		override fun write(stream: OutputStream, element: JsonElement) {
+			element as JsonObject
+			stream.writeVarIntLE(element.size())
+
+			for ((k, v) in element.entrySet()) {
+				stream.writeBinaryString(k)
+				stream.writeBinaryJson(v)
+			}
+		}
+
+		override fun read(stream: InputStream, accounter: NbtAccounter): JsonElement {
+			val count = stream.readVarIntLE(accounter)
+
+			if (count == 0) return JsonObject()
+			if (count < 0) throw JsonSyntaxException("Tried to read json object with $count elements in it")
+
+			val build = JsonObject()
+
+			for (i in 0 until count) {
+				val key = try {
+					stream.readBinaryString(accounter)
+				} catch(err: Throwable) {
+					throw JsonSyntaxException("Reading json object at $i", err)
+				}
+
+				try {
+					build.add(key, stream.readBinaryJson(accounter))
+				} catch(err: Throwable) {
+					throw JsonSyntaxException("Reading json object at $i with name $key", err)
+				}
+			}
+
+			return build
+		}
+	},
+
+	DOUBLE_ARRAY({ it is JsonArray && it.stream().allMatch { DOUBLE.test(it) } }) {
+		override fun write(stream: OutputStream, element: JsonElement) {
+			writeArray(stream, element, DOUBLE::write)
+		}
+
+		override fun read(stream: InputStream, accounter: NbtAccounter): JsonElement {
+			return readArray(stream, accounter, DOUBLE::read)
+		}
+	},
+
+	INT_ARRAY({ it is JsonArray && it.stream().allMatch { INT.test(it) } }) {
+		override fun write(stream: OutputStream, element: JsonElement) {
+			writeArray(stream, element, INT::write)
+		}
+
+		override fun read(stream: InputStream, accounter: NbtAccounter): JsonElement {
+			return readArray(stream, accounter, INT::read)
+		}
+	},
+
+	BOOLEAN_ARRAY({ it is JsonArray && it.stream().allMatch { BOOLEAN.test(it) } }) {
+		override fun write(stream: OutputStream, element: JsonElement) {
+			writeArray(stream, element, BOOLEAN::write)
+		}
+
+		override fun read(stream: InputStream, accounter: NbtAccounter): JsonElement {
+			return readArray(stream, accounter, BOOLEAN::read)
+		}
+	},
+
+	NULL_ARRAY({ it is JsonArray && it.stream().allMatch { NULL.test(it) } }) {
+		override fun write(stream: OutputStream, element: JsonElement) {
+			writeArray(stream, element, NULL::write)
+		}
+
+		override fun read(stream: InputStream, accounter: NbtAccounter): JsonElement {
+			return readArray(stream, accounter, NULL::read)
+		}
+	};
+
+	protected fun writeArray(stream: OutputStream, element: JsonElement, writer: (OutputStream, JsonElement) -> Unit) {
+		element as JsonArray
+		stream.writeVarIntLE(element.size())
+
+		for (v in element) {
+			writer.invoke(stream, v)
+		}
+	}
+
+	protected fun readArray(stream: InputStream, accounter: NbtAccounter, reader: (InputStream, NbtAccounter) -> JsonElement): JsonArray {
+		val count = stream.readVarIntLE(accounter)
+
+		if (count == 0) return JsonArray()
+		if (count < 0) throw JsonSyntaxException("Tried to read json array with $count elements in it")
+
+		return JsonArray(count).also {
+			for (i in 0 until count) it.add(reader.invoke(stream, accounter))
+		}
+	}
+
+	abstract fun write(stream: OutputStream, element: JsonElement)
+	abstract fun read(stream: InputStream, accounter: NbtAccounter): JsonElement
+
+	companion object {
+		val cached: ImmutableList<BinaryElementType> = ImmutableList.of(
+			NULL, DOUBLE, BOOLEAN, INT, STRING, OBJECT,
+			DOUBLE_ARRAY, INT_ARRAY, BOOLEAN_ARRAY, ARRAY)
+
+		val ordered: ImmutableList<BinaryElementType> = ImmutableList.copyOf(values())
+	}
+}
 
 /**
  * Writes binary json to stream in Starbound Object Notation format
@@ -25,48 +196,15 @@ private const val TYPE_OBJECT = 0x07
  * just copy pasted this code from my another project because i was lazy
  */
 fun OutputStream.writeBinaryJson(element: JsonElement) {
-	if (element is JsonObject) {
+	for (writer in BinaryElementType.cached) {
-		write(TYPE_OBJECT)
+		if (writer.test(element)) {
-		writeVarIntLE(element.size())
+			write(writer.ordinal + 1)
-
+			writer.write(this, element)
-		for ((k, v) in element.entrySet()) {
+			return
-			writeBinaryString(k)
-			writeBinaryJson(v)
 		}
-	} else if (element is JsonArray) {
-		write(TYPE_ARRAY)
-		writeVarIntLE(element.size())
-
-		for (v in element) {
-			writeBinaryJson(v)
-		}
-	} else if (element is JsonPrimitive) {
-		if (element.isNumber) {
-			val num = element.asNumber
-
-			if (num is Int || num is Long || num is Short || num is Byte) {
-				write(TYPE_INT)
-				writeVarLongLE(num.toLong())
-			} else if (num is Float || num is Double) {
-				write(TYPE_DOUBLE)
-				writeDouble(num.toDouble())
-			} else {
-				throw IllegalArgumentException("Unknown number type: ${num::class.qualifiedName}")
-			}
-		} else if (element.isString) {
-			write(TYPE_STRING)
-			writeBinaryString(element.asString)
-		} else if (element.isBoolean) {
-			write(TYPE_BOOLEAN)
-			write(if (element.asBoolean) 1 else 0)
-		} else {
-			write(TYPE_NULL)
-		}
-	} else if (element is JsonNull) {
-		write(TYPE_NULL)
-	} else {
-		throw IllegalArgumentException("Unknown element type: ${element::class.qualifiedName}")
 	}
+
+	throw IllegalArgumentException("Unknown element type: ${element::class.qualifiedName}")
 }
 
 /**
@@ -76,54 +214,7 @@ fun OutputStream.writeBinaryJson(element: JsonElement) {
  */
 fun InputStream.readBinaryJson(sizeLimit: NbtAccounter = NbtAccounter(1L shl 18 /* 256 KiB */)): JsonElement {
 	sizeLimit.accountBytes(1L)
-
+	val id = read() - 1
-	return when (val id = read()) {
+	val reader = BinaryElementType.ordered.getOrNull(id) ?: throw JsonParseException("Unknown element type ${id + 1}")
-		TYPE_NULL -> JsonNull.INSTANCE
+	return reader.read(this, sizeLimit)
-		TYPE_DOUBLE -> {
-			sizeLimit.accountBytes(8L)
-			JsonPrimitive(readDouble())
-		}
-		TYPE_BOOLEAN -> {
-			sizeLimit.accountBytes(1L)
-			JsonPrimitive(read() > 1)
-		}
-		TYPE_INT -> JsonPrimitive(readVarLongLE(sizeLimit))
-		TYPE_STRING -> JsonPrimitive(readBinaryString(sizeLimit))
-		TYPE_ARRAY -> {
-			val values = readVarIntLE(sizeLimit)
-
-			if (values == 0) return JsonArray()
-			if (values < 0) throw JsonSyntaxException("Tried to read json array with $values elements in it")
-
-			val build = JsonArray(values)
-			for (i in 0 until values) build.add(readBinaryJson(sizeLimit))
-			return build
-		}
-		TYPE_OBJECT -> {
-			val values = readVarIntLE(sizeLimit)
-			if (values == 0) return JsonObject()
-			if (values < 0) throw JsonSyntaxException("Tried to read json object with $values elements in it")
-
-			val build = JsonObject()
-
-			for (i in 0 until values) {
-				val key: String
-
-				try {
-					key = readBinaryString(sizeLimit)
-				} catch(err: Throwable) {
-					throw JsonSyntaxException("Reading json object at $i", err)
-				}
-
-				try {
-					build.add(key, readBinaryJson(sizeLimit))
-				} catch(err: Throwable) {
-					throw JsonSyntaxException("Reading json object at $i with name $key", err)
-				}
-			}
-
-			return build
-		}
-		else -> throw JsonParseException("Unknown element type $id")
-	}
 }