onion-mc/analysis/client-perf/harness/b14/nbt.py

"""Minimal read-only NBT / Anvil region reader (stdlib only) for B14 site finding.

Enough of the format to read level.dat and 1.18+ chunk biomes + heightmaps:
https://minecraft.wiki/w/NBT_format, /Region_file_format, /Chunk_format.
"""
import gzip
import io
import math
import struct
import zlib

END, BYTE, SHORT, INT, LONG, FLOAT, DOUBLE, BYTE_ARRAY, STRING, LIST, COMPOUND, \
    INT_ARRAY, LONG_ARRAY = range(13)


class _Reader:
    def __init__(self, data):
        self.b, self.i = data, 0

    def take(self, n):
        s = self.b[self.i:self.i + n]
        if len(s) != n:
            raise ValueError("truncated NBT")
        self.i += n
        return s

    def unpack(self, fmt):
        size = struct.calcsize(fmt)
        return struct.unpack(fmt, self.take(size))

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

    def payload(self, t):
        if t == BYTE:
            return self.unpack(">b")[0]
        if t == SHORT:
            return self.unpack(">h")[0]
        if t == INT:
            return self.unpack(">i")[0]
        if t == LONG:
            return self.unpack(">q")[0]
        if t == FLOAT:
            return self.unpack(">f")[0]
        if t == DOUBLE:
            return self.unpack(">d")[0]
        if t == BYTE_ARRAY:
            (n,) = self.unpack(">i")
            return self.take(n)
        if t == STRING:
            return self.string()
        if t == LIST:
            et, n = self.unpack(">bi")
            return [self.payload(et) for _ in range(max(0, n))]
        if t == COMPOUND:
            out = {}
            while True:
                (ct,) = self.unpack(">b")
                if ct == END:
                    return out
                name = self.string()
                out[name] = self.payload(ct)
        if t == INT_ARRAY:
            (n,) = self.unpack(">i")
            return list(struct.unpack(f">{n}i", self.take(4 * n)))
        if t == LONG_ARRAY:
            (n,) = self.unpack(">i")
            return list(struct.unpack(f">{n}q", self.take(8 * n)))
        raise ValueError(f"bad NBT tag type {t}")


def loads(data):
    """Parse an uncompressed NBT blob -> (root name, root compound)."""
    r = _Reader(data)
    (t,) = r.unpack(">b")
    if t != COMPOUND:
        raise ValueError("NBT root is not a compound")
    name = r.string()
    return name, r.payload(COMPOUND)


def load_gzip(data):
    """level.dat style (gzip) NBT."""
    return loads(gzip.decompress(data))


def region_chunks(data):
    """Yield (local_x, local_z, chunk_compound) for every chunk in an .mca blob."""
    if len(data) < 8192:
        return
    for idx in range(1024):
        off = int.from_bytes(data[idx * 4:idx * 4 + 3], "big")
        if off == 0:
            continue
        start = off * 4096
        if start + 5 > len(data):
            continue
        length = int.from_bytes(data[start:start + 4], "big")
        comp = data[start + 4]
        blob = data[start + 5:start + 4 + length]
        try:
            if comp == 2:
                raw = zlib.decompress(blob)
            elif comp == 1:
                raw = gzip.GzipFile(fileobj=io.BytesIO(blob)).read()
            elif comp == 3:
                raw = blob
            else:  # 4 = lz4 (not in vanilla defaults) / 128+ = external .mcc
                continue
            _, chunk = loads(raw)
        except (ValueError, zlib.error, OSError, struct.error):
            continue
        yield idx % 32, idx // 32, chunk


def _unsigned(v):
    return v & 0xFFFFFFFFFFFFFFFF


def unpack_packed(longs, bits, count):
    """1.16+ packed array: entries never span longs."""
    per = 64 // bits
    mask = (1 << bits) - 1
    out = []
    for i in range(count):
        word = _unsigned(longs[i // per])
        out.append((word >> ((i % per) * bits)) & mask)
    return out


def section_biome_grid(section):
    """64 biome ids (index = (y4*4 + z4)*4 + x4) for one chunk section, or None."""
    b = section.get("biomes")
    if not b or not b.get("palette"):
        return None
    pal = b["palette"]
    if len(pal) == 1 or "data" not in b:
        return [pal[0]] * 64
    bits = max(1, math.ceil(math.log2(len(pal))))
    return [pal[i] if i < len(pal) else pal[0] for i in unpack_packed(b["data"], bits, 64)]


def heightmap(chunk, kind="MOTION_BLOCKING_NO_LEAVES", min_y=-64, height=384):
    """256 world-y values (first free block above the top matching block), idx = z*16+x."""
    hm = (chunk.get("Heightmaps") or {}).get(kind)
    if not hm:
        return None
    bits = math.ceil(math.log2(height + 1))
    return [v + min_y for v in unpack_packed(hm, bits, 256)]


def biome_at(chunk, x, y, z):
    """Biome id at block (x, y, z) in this chunk (x, z may be world coords)."""
    sy = y >> 4
    for s in chunk.get("sections", []):
        if s.get("Y") == sy:
            grid = section_biome_grid(s)
            if grid is None:
                return None
            return grid[(((y & 15) >> 2) * 4 + ((z & 15) >> 2)) * 4 + ((x & 15) >> 2)]
    return None