Use template argument for list subtype, extract further information from chunks

9 files changed, 108 insertions, 32 deletions

diff --git a/src/Buffer.hpp b/src/Buffer.hpp
index 9af622d..c7fa9c7 100644
--- a/src/Buffer.hpp
+++ b/src/Buffer.hpp
@@ -27,6 +27,7 @@
 #pragma once
 
 #include <cstdint>
+#include <cstring>
 #include <stdexcept>
 #include <string>
 
@@ -58,6 +59,15 @@ public:
 		len -= n;
 		return std::string(reinterpret_cast<const char *>(data - n), n);
 	}
+
+	void getData(void *buf, size_t n) {
+		if (n > len)
+			throw std::runtime_error("Buffer::get(): buffer underrun");
+
+		data += n;
+		len -= n;
+		std::memcpy(buf, data - n, n);
+	}
 };
 
 }
diff --git a/src/MinedMap.cpp b/src/MinedMap.cpp
index 97074b2..8759d9d 100644
--- a/src/MinedMap.cpp
+++ b/src/MinedMap.cpp
@@ -25,6 +25,7 @@
 
 
 #include "World/Region.hpp"
+#include "NBT/ListTag.hpp"
 
 #include <iostream>
 #include <unistd.h>
@@ -42,8 +43,12 @@ int main(int argc, char *argv[]) {
 
 	const World::Chunk &chunk = *region(0, 0);
 
-	for (auto &entry : std::dynamic_pointer_cast<const NBT::CompoundTag>(chunk->get("Level"))->getValues()) {
-		std::cout << entry.first << " " << entry.second->getType() << std::endl;
+	for (auto &entry : chunk.getSections()) {
+		std::cout << entry->getType() << ":" << std::endl;
+
+		for (auto &entry2 : *entry) {
+			std::cout << entry2.first << " " << entry2.second->getType() << std::endl;
+		}
 	}
 
 	return 0;
diff --git a/src/NBT/ByteArrayTag.hpp b/src/NBT/ByteArrayTag.hpp
index d5a6333..ee0d7d1 100644
--- a/src/NBT/ByteArrayTag.hpp
+++ b/src/NBT/ByteArrayTag.hpp
@@ -47,9 +47,7 @@ private:
 		len |= uint32_t(buffer->get());
 
 		value.resize(len);
-
-		for (uint32_t i = 0; i < len; i++)
-			value[i] = buffer->get();
+		buffer->getData(value.data(), len);
 	}
 
 public:
diff --git a/src/NBT/CompoundTag.hpp b/src/NBT/CompoundTag.hpp
index 18de1c6..28efc41 100644
--- a/src/NBT/CompoundTag.hpp
+++ b/src/NBT/CompoundTag.hpp
@@ -35,19 +35,17 @@
 namespace MinedMap {
 namespace NBT {
 
-class CompoundTag : public Tag {
+class CompoundTag : public Tag, public std::unordered_map<std::string, std::shared_ptr<const Tag>> {
 private:
 	friend class Tag;
 
-	std::unordered_map<std::string, std::shared_ptr<const Tag>> values;
-
 	CompoundTag(Buffer *buffer) {
 		while (true) {
 			std::pair<std::string, std::shared_ptr<const Tag>> v = Tag::readNamedTag(buffer);
 			if (v.second->getType() == Type::End)
 				break;
 
-			values.insert(std::move(v));
+			insert(std::move(v));
 		}
 	}
 
@@ -56,20 +54,18 @@ public:
 		return Type::Compound;
 	}
 
-	const std::unordered_map<std::string, std::shared_ptr<const Tag>> & getValues() const {
-		return values;
-	}
-
-	const std::shared_ptr<const Tag> & get(const std::string &key) const {
-		return values.at(key);
+	template<typename T>
+	std::shared_ptr<const T> get(const std::string &key) const {
+		return std::dynamic_pointer_cast<const T>(at(key));
 	}
 
-	template<typename... Args> const std::shared_ptr<const Tag> & get(const std::string &key, const Args &...args) const {
-		std::shared_ptr<const CompoundTag> tag = std::dynamic_pointer_cast<const CompoundTag>(get(key));
+	template<typename T, typename... Args>
+		std::shared_ptr<const T> get(const std::string &key, const Args &...args) const {
+		std::shared_ptr<const CompoundTag> tag = get<CompoundTag>(key);
 		if (!tag)
-			return std::shared_ptr<Tag>(nullptr);
+			return std::shared_ptr<const T>();
 
-		tag->get(args...);
+		return tag->get<T>(args...);
 	}
 };
 
diff --git a/src/NBT/ListTag.hpp b/src/NBT/ListTag.hpp
index 67ba619..1a48062 100644
--- a/src/NBT/ListTag.hpp
+++ b/src/NBT/ListTag.hpp
@@ -34,30 +34,39 @@
 namespace MinedMap {
 namespace NBT {
 
-class ListTag : public Tag {
-private:
-	friend class Tag;
 
-	std::vector<std::shared_ptr<const Tag>> value;
+class ListTagBase : public Tag {
+public:
+	virtual Type getType() const {
+		return Type::List;
+	}
 
+	virtual Type getSubtype() const = 0;
+};
 
-	ListTag(Buffer *buffer) {
-		Type type = static_cast<Type>(buffer->get());
 
+template<typename T>
+class ListTag : public ListTagBase, public std::vector<std::shared_ptr<const T>> {
+private:
+	friend class Tag;
+
+	Type type;
+
+	ListTag(Type type0, Buffer *buffer) : type(type0) {
 		uint32_t len = uint32_t(buffer->get()) << 24;
 		len |= uint32_t(buffer->get()) << 16;
 		len |= uint32_t(buffer->get()) << 8;
 		len |= uint32_t(buffer->get());
 
-		value.resize(len);
+		this->resize(len);
 
 		for (uint32_t i = 0; i < len; i++)
-			value[i] = Tag::readTag(type, buffer);
+			(*this)[i] = std::static_pointer_cast<const T>(Tag::readTag(type, buffer));
 	}
 
 public:
-	virtual Type getType() const {
-		return Type::List;
+	virtual Type getSubtype() const {
+		return type;
 	}
 };
 
diff --git a/src/NBT/Tag.cpp b/src/NBT/Tag.cpp
index a2d57e3..1829dc5 100644
--- a/src/NBT/Tag.cpp
+++ b/src/NBT/Tag.cpp
@@ -76,7 +76,7 @@ std::shared_ptr<const Tag> Tag::readTag(Type type, Buffer *buffer) {
 		return std::shared_ptr<StringTag>(new StringTag(buffer));
 
 	case Type::List:
-		return std::shared_ptr<ListTag>(new ListTag(buffer));
+		return readList(buffer);
 
 	case Type::Compound:
 		return std::shared_ptr<CompoundTag>(new CompoundTag(buffer));
@@ -85,7 +85,52 @@ std::shared_ptr<const Tag> Tag::readTag(Type type, Buffer *buffer) {
 		return std::shared_ptr<IntArrayTag>(new IntArrayTag(buffer));
 
 	default:
-		throw std::runtime_error("Tag::read: unknown tag type");
+		throw std::runtime_error("Tag::readTag: unknown tag type");
+	}
+}
+
+std::shared_ptr<const Tag> Tag::readList(Buffer *buffer) {
+	Type type = static_cast<Type>(buffer->get());
+
+	switch (type) {
+	case Type::End:
+		return std::shared_ptr<Tag>(new ListTag<EndTag>(type, buffer));
+
+	case Type::Byte:
+		return std::shared_ptr<Tag>(new ListTag<ByteTag>(type, buffer));
+
+	case Type::Short:
+		return std::shared_ptr<Tag>(new ListTag<ShortTag>(type, buffer));
+
+	case Type::Int:
+		return std::shared_ptr<Tag>(new ListTag<IntTag>(type, buffer));
+
+	case Type::Long:
+		return std::shared_ptr<Tag>(new ListTag<LongTag>(type, buffer));
+
+	case Type::Float:
+		return std::shared_ptr<Tag>(new ListTag<FloatTag>(type, buffer));
+
+	case Type::Double:
+		return std::shared_ptr<Tag>(new ListTag<DoubleTag>(type, buffer));
+
+	case Type::ByteArray:
+		return std::shared_ptr<Tag>(new ListTag<ByteArrayTag>(type, buffer));
+
+	case Type::String:
+		return std::shared_ptr<Tag>(new ListTag<StringTag>(type, buffer));
+
+	case Type::List:
+		return std::shared_ptr<Tag>(new ListTag<ListTagBase>(type, buffer));
+
+	case Type::Compound:
+		return std::shared_ptr<Tag>(new ListTag<CompoundTag>(type, buffer));
+
+	case Type::IntArray:
+		return std::shared_ptr<Tag>(new ListTag<IntArrayTag>(type, buffer));
+
+	default:
+		throw std::runtime_error("Tag::readList: unknown tag type");
 	}
 }
 
diff --git a/src/NBT/Tag.hpp b/src/NBT/Tag.hpp
index 077c508..a3c0c76 100644
--- a/src/NBT/Tag.hpp
+++ b/src/NBT/Tag.hpp
@@ -37,6 +37,9 @@ namespace MinedMap {
 namespace NBT {
 
 class Tag {
+private:
+	static std::shared_ptr<const Tag> readList(Buffer *buffer);
+
 public:
 	enum class Type {
 		End = 0,
@@ -54,7 +57,6 @@ public:
 	};
 
 	static std::shared_ptr<const Tag> readTag(Type type, Buffer *buffer);
-
 	static std::pair<std::string, std::shared_ptr<const Tag>> readNamedTag(Buffer *buffer);
 
 	virtual Type getType() const = 0;
diff --git a/src/World/Chunk.cpp b/src/World/Chunk.cpp
index f3a140d..fb46466 100644
--- a/src/World/Chunk.cpp
+++ b/src/World/Chunk.cpp
@@ -95,6 +95,10 @@ Chunk::Chunk(uint8_t *buffer, size_t buflen) {
 
 	if (!(*this) || tag.first != "")
 		throw std::invalid_argument("invalid root tag");
+
+	sections = (*this)->get<const NBT::ListTag<NBT::CompoundTag>>("Level", "Sections");
+	if (!sections)
+		throw std::invalid_argument("no sections found");
 }
 
 }
diff --git a/src/World/Chunk.hpp b/src/World/Chunk.hpp
index 7c485eb..3ed6cf4 100644
--- a/src/World/Chunk.hpp
+++ b/src/World/Chunk.hpp
@@ -29,6 +29,7 @@
 
 #include "../UniqueCPtr.hpp"
 #include "../NBT/CompoundTag.hpp"
+#include "../NBT/ListTag.hpp"
 
 #include <cstdint>
 #include <tuple>
@@ -44,8 +45,14 @@ public:
 private:
 	static std::pair<UniqueCPtr<uint8_t[]>, size_t> inflateChunk(uint8_t *data, size_t len);
 
+	std::shared_ptr<const NBT::ListTag<NBT::CompoundTag>> sections;
+
 public:
 	Chunk(uint8_t *buffer, size_t buflen);
+
+	const NBT::ListTag<NBT::CompoundTag> & getSections() const {
+		return *sections;
+	}
 };
 
 }