Single file NBT reader/writer

NBT.java

/*
This is free and unencumbered software released into the public domain.

Anyone is free to copy, modify, publish, use, compile, sell, or
distribute this software, either in source code form or as a compiled
binary, for any purpose, commercial or non-commercial, and by any
means.

In jurisdictions that recognize copyright laws, the author or authors
of this software dedicate any and all copyright interest in the
software to the public domain. We make this dedication for the benefit
of the public at large and to the detriment of our heirs and
successors. We intend this dedication to be an overt act of
relinquishment in perpetuity of all present and future rights to this
software under copyright law.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.

For more information, please refer to <http://unlicense.org>

NBT.java by Jan Polák, 2019
 */

package com.darkyen.minecraft.util;

import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInput;
import java.io.DataInputStream;
import java.io.DataOutput;
import java.io.DataOutputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.Base64;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Objects;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

/**
 * Very light NBT reader/writer implementation.
 */
public final class NBT {

	private static final Logger LOG = Logger.getLogger("NBT");

	public enum Type {
		/** No name, marks end of compound tags. */
		END(null),
		BYTE(null),
		SHORT(null),
		INT(null),
		LONG(null),
		FLOAT(null),
		DOUBLE(null),

		BYTE_ARRAY(byte[].class),
		STRING(String.class),
		LIST(Tag[].class),
		COMPOUND(Map/*<String, Tag>*/.class),
		INT_ARRAY(int[].class),
		LONG_ARRAY(long[].class);

		@Nullable
		final Class<?> compoundType;

		public static final Type[] TYPES = values();

		Type(@Nullable Class<?> compoundType) {
			this.compoundType = compoundType;
		}
	}

	private static final Tag END_TAG = new Tag(Type.END, 0, null);

	@SuppressWarnings("WeakerAccess")
	public static final class Tag {

		/** Type of this tag. Inspect before calling value() accessors. */
		@NotNull
		public final Type type;

		private final long scalar;
		@Nullable
		private final Object compound;

		Tag(@NotNull Type type, long scalar, @Nullable Object compound) {
			assert type != Type.END || (scalar == 0 && compound == null);
			assert (type.compoundType == null ? compound == null : type.compoundType.isInstance(compound));
			this.type = type;
			this.scalar = scalar;
			this.compound = compound;
		}

		public Tag(byte value) {
			this.type = Type.BYTE;
			this.scalar = value;
			this.compound = null;
		}

		public Tag(short value) {
			this.type = Type.SHORT;
			this.scalar = value;
			this.compound = null;
		}

		public Tag(int value) {
			this.type = Type.INT;
			this.scalar = value;
			this.compound = null;
		}

		public Tag(long value) {
			this.type = Type.LONG;
			this.scalar = value;
			this.compound = null;
		}

		public Tag(float value) {
			this.type = Type.FLOAT;
			this.scalar = Float.floatToRawIntBits(value);
			this.compound = null;
		}

		public Tag(double value) {
			this.type = Type.DOUBLE;
			this.scalar = Double.doubleToRawLongBits(value);
			this.compound = null;
		}

		public Tag(@NotNull byte[] value) {
			this.type = Type.BYTE_ARRAY;
			this.scalar = 0;
			this.compound = value;
		}

		public Tag(@NotNull int[] value) {
			this.type = Type.INT_ARRAY;
			this.scalar = 0;
			this.compound = value;
		}

		public Tag(@NotNull long[] value) {
			this.type = Type.LONG_ARRAY;
			this.scalar = 0;
			this.compound = value;
		}

		public Tag(@NotNull String value) {
			this.type = Type.STRING;
			this.scalar = 0;
			this.compound = value;
		}

		public Tag(@NotNull Map<String, Tag> value) {
			this.type = Type.COMPOUND;
			this.scalar = 0;
			this.compound = value;
		}

		public Tag(@NotNull String compoundKey, @NotNull Tag compoundValue) {
			this.type = Type.COMPOUND;
			this.scalar = 0;
			this.compound = Collections.singletonMap(compoundKey, compoundValue);
		}

		public Tag(@NotNull String compoundKey1, @NotNull Tag compoundValue1, @NotNull Object...compoundKeyValue) {
			this.type = Type.COMPOUND;
			this.scalar = 0;
			final HashMap<String, Tag> compound = new HashMap<>();
			compound.put(compoundKey1, compoundValue1);
			assert compoundKeyValue.length % 2 == 0;
			for (int i = 0; i < compoundKeyValue.length; i += 2) {
				compound.put((String)compoundKeyValue[i], (Tag)compoundKeyValue[i+1]);
			}
			this.compound = compound;
		}

		public Tag(@NotNull Tag...value) {
			if (value.length > 1) {
				final Type type = value[0].type;
				for (int i = 1; i < value.length; i++) {
					assert value[i].type == type : "All types in the list must be the same";
				}
			}
			this.type = Type.LIST;
			this.scalar = 0;
			this.compound = value;
		}

		public byte valueByte() {
			assert type == Type.BYTE;
			return (byte)scalar;
		}

		public short valueShort() {
			assert type == Type.SHORT;
			return (short)scalar;
		}

		public int valueInt() {
			assert type == Type.INT;
			return (int)scalar;
		}

		public long valueLong() {
			assert type == Type.LONG;
			return scalar;
		}

		public float valueFloat() {
			assert type == Type.FLOAT;
			return Float.intBitsToFloat((int)scalar);
		}

		public double valueDouble() {
			assert type == Type.DOUBLE;
			return Double.longBitsToDouble(scalar);
		}

		@NotNull
		public byte[] valueByteArray() {
			assert type == Type.BYTE_ARRAY;
			//noinspection ConstantConditions
			return (byte[])compound;
		}

		@NotNull
		public int[] valueIntArray() {
			assert type == Type.INT_ARRAY;
			//noinspection ConstantConditions
			return (int[])compound;
		}

		@NotNull
		public long[] valueLongArray() {
			assert type == Type.LONG_ARRAY;
			//noinspection ConstantConditions
			return (long[])compound;
		}

		@NotNull
		public String valueString() {
			assert type == Type.STRING;
			//noinspection ConstantConditions
			return (String)compound;
		}

		@SuppressWarnings("unchecked")
		@NotNull
		public Map<String, Tag> valueCompound() {
			assert type == Type.COMPOUND;
			//noinspection ConstantConditions
			return (Map<String, Tag>) compound;
		}

		@NotNull
		public Tag[] valueList() {
			assert type == Type.LIST;
			//noinspection ConstantConditions
			return (Tag[])compound;
		}

		@Override
		public boolean equals(Object o) {
			if (this == o) return true;
			if (o == null || getClass() != o.getClass()) return false;

			Tag tag = (Tag) o;
			if (type != tag.type) return false;

			switch (type) {
				case END:
					return true;
				case BYTE:
					return ((scalar ^ tag.scalar) & 0xFF) == 0;
				case SHORT:
					return ((scalar ^ tag.scalar) & 0xFFFF) == 0;
				case INT:
				case FLOAT:
					return ((scalar ^ tag.scalar) & 0xFFFF_FFFFL) == 0L;
				case LONG:
				case DOUBLE:
					return scalar == tag.scalar;
				case BYTE_ARRAY:
					return Arrays.equals((byte[])compound, (byte[])tag.compound);
				case INT_ARRAY:
					return Arrays.equals((int[])compound, (int[])tag.compound);
				case LONG_ARRAY:
					return Arrays.equals((long[])compound, (long[])tag.compound);
				case STRING:
				case COMPOUND:
					return Objects.equals(compound, tag.compound);
				case LIST:
					return Arrays.equals((Tag[])compound, (Tag[])tag.compound);
			}

			return false;
		}

		@Override
		public int hashCode() {
			int result = 31 * type.hashCode();
			switch (type) {
				case END:
					break;
				case BYTE:
				case SHORT:
				case INT:
				case LONG:
				case FLOAT:
				case DOUBLE:
					result += scalar;
					break;
				case BYTE_ARRAY:
					result += Arrays.hashCode(valueByteArray());
					break;
				case INT_ARRAY:
					result += Arrays.hashCode(valueIntArray());
					break;
				case LONG_ARRAY:
					result += Arrays.hashCode(valueLongArray());
					break;
				case STRING:
				case COMPOUND:
					result += compound != null ? compound.hashCode() : 0;
					break;
				case LIST:
					result += Arrays.hashCode(valueList());
					break;
			}
			return result;
		}

		@Override
		public String toString() {
			final StringBuilder sb = new StringBuilder();
			toString(sb);
			return sb.toString();
		}

		private void toString(StringBuilder sb) {
			switch (type) {
				case END:
					sb.append("END");
					break;
				case BYTE:
					sb.append(valueByte()).append('b');
					break;
				case SHORT:
					sb.append(valueShort()).append('s');
					break;
				case INT:
					sb.append(valueInt());
					break;
				case LONG:
					sb.append(valueLong()).append('l');
					break;
				case FLOAT:
					sb.append(valueFloat()).append('f');
					break;
				case DOUBLE:
					sb.append(valueDouble()).append('d');
					break;
				case BYTE_ARRAY:
					sb.append("[B;");
					final @NotNull byte[] bytes = valueByteArray();
					if (bytes.length > 0) {
						sb.append(bytes[0]);
						for (int i = 1; i < bytes.length; i++) {
							sb.append(',').append(bytes[i]);
						}
					}
					sb.append(']');
					break;
				case STRING:
					final String string = valueString();
					sb.append('"');
					for (int i = 0; i < string.length(); i++) {
						final char c = string.charAt(i);
						if (c == '"') {
							sb.append("\"");
						} else {
							sb.append(c);
						}
					}
					sb.append('"');
					break;
				case LIST:
					sb.append("[");
					final @NotNull Tag[] tags = valueList();
					if (tags.length > 0) {
						tags[0].toString(sb);
						for (int i = 1; i < tags.length; i++) {
							sb.append(',');
							tags[i].toString(sb);
						}
					}
					sb.append(']');
					break;
				case COMPOUND:
					sb.append("{");
					boolean first = true;
					for (Map.Entry<String, Tag> entry : valueCompound().entrySet()) {
						if (first) {
							first = false;
						} else {
							sb.append(',');
						}
						sb.append(entry.getKey()).append(':');
						entry.getValue().toString(sb);
					}
					sb.append('}');
					break;
				case INT_ARRAY:
					sb.append("[I;");
					final @NotNull int[] ints = valueIntArray();
					if (ints.length > 0) {
						sb.append(ints[0]);
						for (int i = 1; i < ints.length; i++) {
							sb.append(',').append(ints[i]);
						}
					}
					sb.append(']');
					break;
				case LONG_ARRAY:
					sb.append("[B;");
					final @NotNull long[] longs = valueLongArray();
					if (longs.length > 0) {
						sb.append(longs[0]);
						for (int i = 1; i < longs.length; i++) {
							sb.append(',').append(longs[i]);
						}
					}
					sb.append(']');
					break;
			}
		}

		@Nullable
		public static Tag fromBase64(@NotNull String base64, boolean gzipped) {
			try {
				final byte[] baseBytes = Base64.getDecoder().decode(base64);
				final DataInputStream dataInputStream;
				if (gzipped) {
					dataInputStream = new DataInputStream(new GZIPInputStream(new ByteArrayInputStream(baseBytes)));
				} else {
					dataInputStream = new DataInputStream(new ByteArrayInputStream(baseBytes));
				}
				return NBT.read(dataInputStream);
			} catch (IllegalArgumentException | IOException e) {
				LOG.log(Level.WARNING, "Failed to decode", e);
				return null;
			}
		}

		@Nullable
		public String toBase64(boolean gzipped) {
			try {
				final ByteArrayOutputStream byteOut = new ByteArrayOutputStream();
				try (final DataOutputStream dataOutputStream = new DataOutputStream(gzipped ? new GZIPOutputStream(byteOut) : byteOut)) {
					NBT.write(this, dataOutputStream);
				}
				return Base64.getEncoder().encodeToString(byteOut.toByteArray());
			} catch (IOException e) {
				LOG.log(Level.WARNING, "Failed to decode: "+this, e);
				return null;
			}
		}
	}

	@Nullable
	public static Tag read(DataInput in) throws IOException {
		final Type type = readType(in);
		if (type == null) {
			return null;
		} else if (type == Type.END) {
			return END_TAG;
		}

		final int tagLength = in.readUnsignedShort();
		in.skipBytes(tagLength);

		final Object compound = readCompound(type, in);
		if (compound == COMPOUND_ERROR) {
			return null;
		}
		final long scalar = readScalar(type, in);

		return new Tag(type, scalar, compound);
	}

	@Nullable
	private static Type readType(DataInput in) throws IOException {
		final int tag = in.readByte() & 0xFF;
		if (tag >= Type.TYPES.length) {
			LOG.log(Level.WARNING, "Invalid type: "+tag);
			return null;
		}

		return Type.TYPES[tag];
	}

	private static long readScalar(@NotNull Type type, DataInput in) throws IOException {
		switch (type) {
			case BYTE:
				return in.readByte();
			case SHORT:
				return in.readShort();
			case INT:
			case FLOAT:
				return in.readInt();
			case LONG:
			case DOUBLE:
				return in.readLong();
			default:
				return 0;
		}
	}

	private static final Object COMPOUND_ERROR = new Object();

	private static Object readCompound(@NotNull Type type, DataInput in) throws IOException {
		switch (type) {
			case BYTE_ARRAY: {
				final int length = in.readInt();
				final byte[] bytes = new byte[length];
				in.readFully(bytes);
				return bytes;
			}
			case STRING:
				return in.readUTF();
			case LIST: {
				final Type elementType = readType(in);
				if (elementType == null) {
					return COMPOUND_ERROR;
				}
				final int length = in.readInt();
				final Tag[] tags = new Tag[length];
				for (int i = 0; i < length; i++) {
					final Object compound = readCompound(elementType, in);
					if (compound == COMPOUND_ERROR) {
						return COMPOUND_ERROR;
					}
					final long scalar = readScalar(elementType, in);

					tags[i] = new Tag(elementType, scalar, compound);
				}
				return tags;
			}
			case COMPOUND: {
				final LinkedHashMap<String, Tag> tags = new LinkedHashMap<>();
				while (true) {
					final Type nestedType = readType(in);
					if (nestedType == null) {
						return COMPOUND_ERROR;
					} else if (nestedType == Type.END) {
						return tags;
					}
					final String name = in.readUTF();

					final Object compound = readCompound(nestedType, in);
					if (compound == COMPOUND_ERROR) {
						return COMPOUND_ERROR;
					}
					final long scalar = readScalar(nestedType, in);
					tags.put(name, new Tag(nestedType, scalar, compound));
				}
			}
			case INT_ARRAY: {
				final int length = in.readInt();
				final int[] ints = new int[length];
				for (int i = 0; i < length; i++) {
					ints[i] = in.readInt();
				}
				return ints;
			}
			case LONG_ARRAY: {
				final int length = in.readInt();
				final long[] longs = new long[length];
				for (int i = 0; i < length; i++) {
					longs[i] = in.readLong();
				}
				return longs;
			}
			default:
				return null;
		}
	}

	public static void write(@NotNull Tag tag, @NotNull DataOutput out) throws IOException {
		final Type type = tag.type;
		out.writeByte(type.ordinal());
		if (type == Type.END) {
			return;
		}

		out.writeShort(0); // Tag name
		writePayload(tag, out);
	}

	private static void writePayload(@NotNull Tag tag, @NotNull DataOutput out) throws IOException {
		switch (tag.type) {
			case BYTE:
				out.writeByte((byte)tag.scalar);
				break;
			case SHORT:
				out.writeShort((short)tag.scalar);
				break;
			case INT:
			case FLOAT:
				out.writeInt((int)tag.scalar);
				break;
			case LONG:
			case DOUBLE:
				out.writeLong(tag.scalar);
				break;
			case BYTE_ARRAY: {
				final @NotNull byte[] b = tag.valueByteArray();
				out.writeInt(b.length);
				out.write(b);
				break;
			}
			case STRING:
				out.writeUTF(tag.valueString());
				break;
			case LIST: {
				final @NotNull Tag[] tagList = tag.valueList();
				out.writeByte((tagList.length <= 0 ? Type.END : tagList[0].type).ordinal());
				out.writeInt(tagList.length);
				for (Tag listTag : tagList) {
					writePayload(listTag, out);
				}
				break;
			}
			case COMPOUND: {
				final Map<String, Tag> compound = tag.valueCompound();
				for (Map.Entry<String, Tag> entry : compound.entrySet()) {
					final Tag entryTag = entry.getValue();
					final String entryName = entry.getKey();

					out.writeByte(entryTag.type.ordinal());
					out.writeUTF(entryName);
					writePayload(entryTag, out);
				}
				out.writeByte(Type.END.ordinal());
				break;
			}
			case INT_ARRAY: {
				final @NotNull int[] ints = tag.valueIntArray();
				out.writeInt(ints.length);
				for (int i : ints) {
					out.writeInt(i);
				}
				break;
			}
			case LONG_ARRAY: {
				final @NotNull long[] longs = tag.valueLongArray();
				out.writeInt(longs.length);
				for (long i : longs) {
					out.writeLong(i);
				}
				break;
			}
		}
	}

}