RegionIO/internal/world/encode.go

package world

import (
	"math/bits"
	"sync"
	"sync/atomic"

	"regionio/internal/protocol"
)

// World vertical geometry for the overworld dimension type.
const (
	MinY         = -64
	WorldHeight  = 384
	SectionCount = WorldHeight / 16 // 24 sections
	sectionVol   = 16 * 16 * 16     // 4096 blocks
)

// Common block-state network IDs (from the generated block report).
// Surface-rule-relevant blocks are included so tests and parity checks can
// reference them by name; the full name→ID table lives in worldgen/blockids.go.
const (
	StateAir     uint16 = 0
	StateStone   uint16 = 1
	StateGrass   uint16 = 9
	StateDirt    uint16 = 10
	StateBedrock uint16 = 85
	StateWater   uint16 = 86
	StateLava    uint16 = 102
	StateSand    uint16 = 118
	StateGravel  uint16 = 124
	StateOakLog  uint16 = 137
	StateOakLeaf uint16 = 279

	// Surface-rule blocks (IDs captured from blocks.json 26.1.2).
	StateCoarseDirt   uint16 = 11
	StatePodzol       uint16 = 13
	StateRedSand      uint16 = 123
	StateSandstone    uint16 = 578
	StateSnow         uint16 = 6919 // snow, layers=1
	StateSnowBlock    uint16 = 6928
	StateIce          uint16 = 6927
	StateGlowstone    uint16 = 7016
	StateMycelium     uint16 = 8919
	StateTerracotta   uint16 = 12912
	StateRedSandstone uint16 = 13247
	StateCalcite      uint16 = 24687
	StatePowderSnow   uint16 = 24689
)

// BiomePlains is the network ID (registry index) of minecraft:plains.
const BiomePlains uint16 = 40

// totalBlockStates is one past the largest block-state ID; it sets the
// direct-palette bit width.
const totalBlockStates = 29873

// Biome-cell geometry for the overworld. A biome cell is biomeCellSize³ blocks
// (4×4×4), so each 16-block chunk section holds biomeCellsPerSection biome
// cells. totalBiomes is the size of the synchronized biome registry and sets
// the biome direct-palette bit width.
const (
	biomeCellSize        = 4
	biomeCellsXZ         = 16 / biomeCellSize                         // 4
	biomeCellsPerSection = biomeCellsXZ * biomeCellsXZ * biomeCellsXZ // 64
	totalBiomes          = 65                                         // synced minecraft:worldgen/biome registry size

	// maxBiomeLinearBits is the widest indirect (linear) biome palette the client
	// will read. Vanilla's SECTION_BIOMES strategy switches on the bit count with
	// `tableswitch {0..3}`: 0 is single-valued, 1-3 are linear, and everything
	// else falls through to the global palette — there is no hashmap tier for
	// biomes, unlike block states. Writing a linear palette at 4+ bits makes the
	// client read the container as global: it consumes no palette prefix and
	// re-reads the long array at bitsFor(totalBiomes), so the rest of the chunk
	// payload is misaligned.
	maxBiomeLinearBits = 3
)

// Chunk is a 16xWorldHeightx16 column of block states. Each section may carry a
// per-cell biome array (4×4×4); when biomes[si] is nil the section falls back to
// the column-wide biome field (used by flat/simple generators).
type Chunk struct {
	mu         sync.RWMutex
	revision   atomic.Uint64
	X, Z       int32
	sections   [SectionCount]*[sectionVol]uint16
	biomes     [SectionCount]*[biomeCellsPerSection]uint16
	skyLight   [SectionCount]*[2048]byte
	blockLight [SectionCount]*[2048]byte
	lightReady bool
	// lightValidated is runtime-only. Persisted arrays are ready to read but are
	// reconciled with current neighbor blocks once after entering a live cache.
	lightValidated bool
	biome          uint16 // fallback uniform biome when biomes[si] is nil
}

// NewChunk returns an empty (all-air) chunk at (x, z) with the given biome.
func NewChunk(x, z int32, biome uint16) *Chunk {
	return &Chunk{X: x, Z: z, biome: biome}
}

// blockIndex maps local coordinates to the YZX-ordered section array index.
func blockIndex(lx, ly, lz int) int { return (ly&15)<<8 | (lz&15)<<4 | (lx & 15) }

// section returns section i, allocating it on first write.
func (c *Chunk) section(i int) *[sectionVol]uint16 {
	if c.sections[i] == nil {
		c.sections[i] = new([sectionVol]uint16)
	}
	return c.sections[i]
}

// GetBlock returns the block state at local (lx, lz) and world height y, or
// StateAir if the section is empty or y is out of range.
func (c *Chunk) GetBlock(lx, y, lz int) uint16 {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.getBlock(lx, y, lz)
}

// getBlock is the lock-free form used while operating on a private snapshot.
func (c *Chunk) getBlock(lx, y, lz int) uint16 {
	si := (y - MinY) >> 4
	if si < 0 || si >= SectionCount {
		return StateAir
	}
	s := c.sections[si]
	if s == nil {
		return StateAir
	}
	return s[blockIndex(lx, y, lz)]
}

// GetBiome returns the biome of the 4×4×4 cell containing block (lx, y, lz). It
// mirrors GetBlock: the per-section biome array if present, else the column's
// uniform fallback biome. Needed for on-disk chunk serialization.
func (c *Chunk) GetBiome(lx, y, lz int) uint16 {
	c.mu.RLock()
	defer c.mu.RUnlock()
	si := (y - MinY) >> 4
	if si < 0 || si >= SectionCount {
		return c.biome
	}
	b := c.biomes[si]
	if b == nil {
		return c.biome
	}
	return b[biomeIndex(lx, y, lz)]
}

// SetBlock sets the block at local (lx, lz) and absolute world height y.
func (c *Chunk) SetBlock(lx, y, lz int, state uint16) {
	_, changed := c.setBlock(lx, y, lz, state)
	if changed {
		c.mu.Lock()
		c.lightReady = false
		c.lightValidated = false
		c.mu.Unlock()
	}
}

// setBlockRaw writes to an unpublished chunk during generation. Generators call
// it only after their parallel sampling phases have joined and before the chunk
// enters Cache, avoiding a mutex operation for every solid terrain block.
func (c *Chunk) setBlockRaw(lx, y, lz int, state uint16) {
	si := (y - MinY) >> 4
	if si < 0 || si >= SectionCount {
		return
	}
	c.section(si)[blockIndex(lx, y, lz)] = state
}

// setBlock changes a block and returns the resulting revision. The changed flag
// lets the cache avoid dirtying a chunk for a no-op client prediction.
func (c *Chunk) setBlock(lx, y, lz int, state uint16) (revision uint64, changed bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	si := (y - MinY) >> 4
	if si < 0 || si >= SectionCount {
		return c.revision.Load(), false
	}
	idx := blockIndex(lx, y, lz)
	s := c.section(si)
	if s[idx] == state {
		return c.revision.Load(), false
	}
	s[idx] = state
	return c.revision.Add(1), true
}

// biomeIndex maps a block within a section to its YZX-ordered 4×4×4 biome cell.
// Coordinates are folded into 0..15 (block coords) then divided to cell coords.
func biomeIndex(lx, ly, lz int) int {
	bx := (lx & 15) / biomeCellSize
	by := (ly & 15) / biomeCellSize
	bz := (lz & 15) / biomeCellSize
	return by<<(biomeCellsXZBits*2) | bz<<biomeCellsXZBits | bx
}

// biomeCellsXZBits is log2(biomeCellsXZ) for the YZX index assembly.
const biomeCellsXZBits = 2 // biomeCellsXZ=4 → 2 bits

// SetBiome sets the biome for the 4×4×4 cell containing block (lx, y, lz). The
// section's per-cell biome array is allocated lazily on first write. Any block
// in the cell shares its biome, matching the 4-block resolution vanilla uses.
func (c *Chunk) SetBiome(lx, y, lz int, biome uint16) {
	c.mu.Lock()
	defer c.mu.Unlock()
	si := (y - MinY) >> 4
	if si < 0 || si >= SectionCount {
		return
	}
	if c.biomes[si] == nil {
		cells := new([biomeCellsPerSection]uint16)
		for i := range cells {
			cells[i] = c.biome
		}
		c.biomes[si] = cells
	}
	idx := biomeIndex(lx, y, lz)
	if c.biomes[si][idx] != biome {
		c.biomes[si][idx] = biome
		c.revision.Add(1)
	}
}

// Encode serializes the level_chunk_with_light body for this chunk.
func (c *Chunk) Encode() []byte {
	snapshot, _ := c.snapshot()
	return snapshot.encode()
}

// snapshot returns a detached, immutable copy and the revision it represents.
// Section arrays are copied so encoding and persistence never race block edits.
func (c *Chunk) snapshot() (*Chunk, uint64) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	revision := c.revision.Load()
	clone := &Chunk{X: c.X, Z: c.Z, biome: c.biome}
	clone.revision.Store(revision)
	for i := 0; i < SectionCount; i++ {
		if c.sections[i] != nil {
			section := *c.sections[i]
			clone.sections[i] = &section
		}
		if c.biomes[i] != nil {
			biomes := *c.biomes[i]
			clone.biomes[i] = &biomes
		}
		if c.skyLight[i] != nil {
			light := *c.skyLight[i]
			clone.skyLight[i] = &light
		}
		if c.blockLight[i] != nil {
			light := *c.blockLight[i]
			clone.blockLight[i] = &light
		}
	}
	clone.lightReady = c.lightReady
	clone.lightValidated = c.lightValidated
	return clone, revision
}

// currentRevision returns the latest mutation revision.
func (c *Chunk) currentRevision() uint64 {
	return c.revision.Load()
}

// encode serializes a detached snapshot.
func (c *Chunk) encode() []byte {
	w := protocol.NewWriter(8192)
	w.Int32(c.X).Int32(c.Z)
	c.writeHeightmaps(w)

	// Section data is length-prefixed.
	sec := protocol.NewWriter(4096)
	for i := 0; i < SectionCount; i++ {
		c.writeSection(sec, i)
	}
	w.VarInt(int32(sec.Len()))
	w.Raw(sec.Bytes())

	w.VarInt(0) // block entity count
	c.writeLight(w)
	return w.Bytes()
}

// Heightmap.Types ids whose Usage is CLIENT. Vanilla sends exactly these three
// and no others.
const (
	hmWorldSurface           = 1
	hmMotionBlocking         = 4
	hmMotionBlockingNoLeaves = 5
)

// writeHeightmaps emits the three heightmaps the client is sent.
//
// They are not the same map, which is what this code used to assume. Each stops
// at a different block: WORLD_SURFACE at the first thing that is not air,
// MOTION_BLOCKING at the first thing that blocks movement or holds fluid, and
// MOTION_BLOCKING_NO_LEAVES at the first such thing that is not leaves. A
// flower, a torch, a sapling or a tree canopy separates them — the client uses
// MOTION_BLOCKING to place rain and snow particles and to decide where a
// fishing bobber lands, so a canopy reported as solid ground rains indoors.
func (c *Chunk) writeHeightmaps(w *protocol.Writer) {
	surface, motion, motionNoLeaves := c.heightmaps()
	w.VarInt(3)
	for _, hm := range [...]struct {
		id     int32
		values [256]uint16
	}{
		{hmWorldSurface, surface},
		{hmMotionBlocking, motion},
		{hmMotionBlockingNoLeaves, motionNoLeaves},
	} {
		packed := packHeightmap(hm.values)
		w.VarInt(hm.id)
		w.VarInt(int32(len(packed)))
		for _, v := range packed {
			w.Int64(int64(v))
		}
	}
}

// heightmaps walks every column once from the top down, recording the first
// block that satisfies each predicate. The stored value is one above the
// matching block, relative to the world floor — what Heightmap.setHeight
// writes — so 0 means the column has no matching block at all.
func (c *Chunk) heightmaps() (surface, motion, motionNoLeaves [256]uint16) {
	for lx := 0; lx < 16; lx++ {
		for lz := 0; lz < 16; lz++ {
			i := lz*16 + lx
			var haveSurface, haveMotion, haveNoLeaves bool
			for si := SectionCount - 1; si >= 0; si-- {
				s := c.sections[si]
				if s == nil {
					continue
				}
				for ly := 15; ly >= 0; ly-- {
					y := MinY + si*16 + ly
					state := s[blockIndex(lx, y, lz)]
					if state == StateAir {
						continue
					}
					h := uint16(y + 1 - MinY)
					if !haveSurface {
						surface[i], haveSurface = h, true
					}
					if !haveMotion && blocksMotionOrFluid(state) {
						motion[i], haveMotion = h, true
					}
					if !haveNoLeaves && blocksMotionNoLeaves(state) {
						motionNoLeaves[i], haveNoLeaves = h, true
					}
					if haveSurface && haveMotion && haveNoLeaves {
						break
					}
				}
				if haveSurface && haveMotion && haveNoLeaves {
					break
				}
			}
		}
	}
	return surface, motion, motionNoLeaves
}

// columnHeights returns the WORLD_SURFACE heightmap on its own, for the
// on-disk Heightmaps tag and the parity test.
func (c *Chunk) columnHeights() [256]uint16 {
	surface, _, _ := c.heightmaps()
	return surface
}

// packHeightmap packs 256 column heights at 9 bits each, 7 values per long,
// without spanning longs (37 longs).
func packHeightmap(h [256]uint16) []uint64 {
	const bpe = 9
	const perLong = 64 / bpe // 7
	out := make([]uint64, (256+perLong-1)/perLong)
	for i, v := range h {
		out[i/perLong] |= uint64(v&0x1FF) << uint((i%perLong)*bpe)
	}
	return out
}

// writeSection emits one chunk section: block count, block paletted container,
// then the biome paletted container (per-cell 4×4×4, or single-valued for legacy
// generators that only set a column-wide biome).
func (c *Chunk) writeSection(w *protocol.Writer, i int) {
	s := c.sections[i]
	if s == nil {
		w.Uint16(0) // non-air block count
		w.Uint16(0) // reserved 2-byte field (always 0 in vanilla)
		writeSingleValued(w, uint32(StateAir))
	} else {
		w.Uint16(uint16(nonAirCount(s)))
		w.Uint16(0) // reserved 2-byte field
		writeBlockPalette(w, s)
	}
	// Biome container: per-cell palette when present, else the uniform fallback.
	if b := c.biomes[i]; b != nil {
		writeBiomePalette(w, b)
	} else {
		writeSingleValued(w, uint32(c.biome))
	}
}

func nonAirCount(s *[sectionVol]uint16) int {
	n := 0
	for _, v := range s {
		if v != StateAir {
			n++
		}
	}
	return n
}

// writeSingleValued writes a bits-per-entry-0 paletted container (no data).
func writeSingleValued(w *protocol.Writer, value uint32) {
	w.Byte(0)
	w.VarInt(int32(value))
}

// writeBlockPalette writes a block-state paletted container, choosing the
// single-valued, indirect, or direct encoding as appropriate.
func writeBlockPalette(w *protocol.Writer, s *[sectionVol]uint16) {
	palette, indexOf := buildPalette(s[:])
	if len(palette) == 1 {
		writeSingleValued(w, uint32(palette[0]))
		return
	}

	bpe := bitsFor(len(palette))
	if bpe < 4 {
		bpe = 4 // minimum for the indirect block format
	}
	if bpe > 8 {
		writeDirect(w, s)
		return
	}

	w.Byte(byte(bpe))
	w.VarInt(int32(len(palette)))
	for _, st := range palette {
		w.VarInt(int32(st))
	}
	writePackedIndices(w, bpe, sectionVol, func(i int) uint32 {
		return uint32(indexOf[s[i]])
	})
}

// writeBiomePalette writes a biome paletted container over the 64 cells of a
// section. It mirrors writeBlockPalette but with biome-specific thresholds: the
// indirect palette allows a minimum of 1 bit per entry (vs 4 for blocks), and
// the direct form is used once the palette bit width exceeds the biome
// registry width.
func writeBiomePalette(w *protocol.Writer, s *[biomeCellsPerSection]uint16) {
	palette, indexOf := buildPalette(s[:])
	if len(palette) == 1 {
		writeSingleValued(w, uint32(palette[0]))
		return
	}

	bpe := bitsFor(len(palette))
	if bpe < 1 {
		bpe = 1 // minimum for the indirect biome format
	}
	if bpe > maxBiomeLinearBits {
		writeBiomeDirect(w, s)
		return
	}

	w.Byte(byte(bpe))
	w.VarInt(int32(len(palette)))
	for _, st := range palette {
		w.VarInt(int32(st))
	}
	writePackedIndices(w, bpe, biomeCellsPerSection, func(i int) uint32 {
		return uint32(indexOf[s[i]])
	})
}

// writeBiomeDirect writes a direct (palette-less) biome container of registry
// IDs, sized to the full biome registry width.
func writeBiomeDirect(w *protocol.Writer, s *[biomeCellsPerSection]uint16) {
	bpe := bitsFor(totalBiomes)
	w.Byte(byte(bpe))
	writePackedIndices(w, bpe, biomeCellsPerSection, func(i int) uint32 {
		return uint32(s[i])
	})
}

// writeDirect writes a direct (palette-less) container of global state IDs.
func writeDirect(w *protocol.Writer, s *[sectionVol]uint16) {
	bpe := bitsFor(totalBlockStates)
	w.Byte(byte(bpe))
	writePackedIndices(w, bpe, sectionVol, func(i int) uint32 {
		return uint32(s[i])
	})
}

// writePackedIndices emits the long-array data: count entries of bpe bits each,
// packed perLong=64/bpe values per long, never spanning a long boundary. The
// long count is NOT length-prefixed; the client derives it from bpe.
func writePackedIndices(w *protocol.Writer, bpe, count int, value func(i int) uint32) {
	perLong := 64 / bpe
	numLongs := (count + perLong - 1) / perLong

	mask := uint64(1)<<uint(bpe) - 1
	for l := 0; l < numLongs; l++ {
		var packed uint64
		for j := 0; j < perLong; j++ {
			idx := l*perLong + j
			if idx >= count {
				break
			}
			packed |= (uint64(value(idx)) & mask) << uint(j*bpe)
		}
		w.Int64(int64(packed))
	}
}

// buildPalette returns the distinct values in s and a value->index map. It
// takes a slice so the same routine serves block sections (sectionVol entries)
// and biome cells (biomeCellsPerSection entries); callers pass array[:] in.
func buildPalette(s []uint16) ([]uint16, map[uint16]int) {
	indexOf := make(map[uint16]int)
	var palette []uint16
	for _, v := range s {
		if _, ok := indexOf[v]; !ok {
			indexOf[v] = len(palette)
			palette = append(palette, v)
		}
	}
	return palette, indexOf
}

// bitsFor returns the bits needed to index n distinct values (min 1).
func bitsFor(n int) int {
	if n <= 1 {
		return 0
	}
	return bits.Len(uint(n - 1))
}