cobblemon-survival-pack/tools/anvil.py

"""Read Minecraft 1.21 worlds: NBT, region (.mca) files and chunk data.

Only reading, and only what the world tools need: block and biome palettes,
heightmaps and level.dat. Long arrays come back as numpy int64 arrays.
"""
import gzip
import struct
import zlib
from pathlib import Path

import numpy as np

MIN_Y = -64


class _Reader:
    __slots__ = ("buf", "pos")

    def __init__(self, buf):
        self.buf = buf
        self.pos = 0

    def take(self, n):
        start = self.pos
        self.pos += n
        return self.buf[start:self.pos]

    def unpack(self, fmt, size):
        value = struct.unpack_from(fmt, self.buf, self.pos)[0]
        self.pos += size
        return value

    def string(self):
        length = self.unpack(">H", 2)
        return bytes(self.take(length)).decode("utf-8", "replace")

    def payload(self, tag):
        if tag == 1:
            return self.unpack(">b", 1)
        if tag == 2:
            return self.unpack(">h", 2)
        if tag == 3:
            return self.unpack(">i", 4)
        if tag == 4:
            return self.unpack(">q", 8)
        if tag == 5:
            return self.unpack(">f", 4)
        if tag == 6:
            return self.unpack(">d", 8)
        if tag == 7:
            return np.frombuffer(self.take(self.unpack(">i", 4)), dtype=np.int8)
        if tag == 8:
            return self.string()
        if tag == 9:
            inner, length = self.unpack(">b", 1), self.unpack(">i", 4)
            return [self.payload(inner) for _ in range(length)]
        if tag == 10:
            out = {}
            while (inner := self.unpack(">b", 1)) != 0:
                name = self.string()
                out[name] = self.payload(inner)
            return out
        if tag == 11:
            return np.frombuffer(self.take(4 * self.unpack(">i", 4)), dtype=">i4")
        if tag == 12:
            return np.frombuffer(self.take(8 * self.unpack(">i", 4)), dtype=">i8").astype(np.int64)
        raise ValueError(f"unknown NBT tag {tag}")


def parse_nbt(data):
    reader = _Reader(memoryview(data))
    tag = reader.unpack(">b", 1)
    reader.string()
    return reader.payload(tag)


def read_level_dat(world):
    return parse_nbt(gzip.decompress(Path(world, "level.dat").read_bytes()))["Data"]


def region_chunks(path):
    """Yield (chunk_x, chunk_z, root_compound) for every chunk stored in a region file."""
    data = Path(path).read_bytes()
    if len(data) < 8192:
        return
    _, rx, rz, _ = Path(path).name.split(".")
    rx, rz = int(rx), int(rz)
    for i in range(1024):
        loc = struct.unpack_from(">I", data, i * 4)[0]
        offset, sectors = loc >> 8, loc & 0xFF
        if offset == 0 or sectors == 0:
            continue
        start = offset * 4096
        length, compression = struct.unpack_from(">IB", data, start)
        payload = data[start + 5:start + 4 + length]
        if compression == 2:
            payload = zlib.decompress(payload)
        elif compression == 1:
            payload = gzip.decompress(payload)
        elif compression != 3:
            continue  # external (.mcc) or LZ4 chunks: not produced by vanilla defaults
        yield rx * 32 + (i % 32), rz * 32 + (i // 32), parse_nbt(payload)


def unpack(longs, bits, count):
    """Unpack a 1.16+ packed long array (entries never span two longs)."""
    if bits == 0 or longs is None or len(longs) == 0:
        return np.zeros(count, dtype=np.int64)
    per_long = 64 // bits
    shifts = np.arange(per_long, dtype=np.uint64) * np.uint64(bits)
    values = (longs.astype(np.uint64)[:, None] >> shifts[None, :]) & np.uint64((1 << bits) - 1)
    return values.reshape(-1)[:count].astype(np.int64)


def heightmap(chunk, name):
    """16x16 array (z, x) of the first free y above the column, in world coordinates."""
    longs = chunk.get("Heightmaps", {}).get(name)
    if longs is None:
        return None
    return unpack(longs, 9, 256).reshape(16, 16) + MIN_Y


def section_blocks(section):
    """(palette names, 4096 indices in y,z,x order) for one chunk section."""
    states = section.get("block_states", {})
    palette = [entry["Name"] for entry in states.get("palette", [])]
    bits = max(4, (len(palette) - 1).bit_length())
    indices = unpack(states.get("data"), bits, 4096) if len(palette) > 1 else np.zeros(4096, dtype=np.int64)
    return palette, indices


def section_biomes(section):
    """(palette, 64 indices in y,z,x order) for one chunk section."""
    biomes = section.get("biomes", {})
    palette = list(biomes.get("palette", []))
    bits = (len(palette) - 1).bit_length()
    indices = unpack(biomes.get("data"), bits, 64) if len(palette) > 1 else np.zeros(64, dtype=np.int64)
    return palette, indices