alex/beta-minecraft-rust

Minecraft Beta 1.7.3 Rewrite
root / crates / common / src / world / chunk.rs

use std::cmp::Reverse;
use std::ops::Range;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::mpsc::{self, Receiver, TryRecvError};
use std::sync::{Arc, LazyLock, Mutex, PoisonError};

use palettevec::PaletteVec;
use palettevec::index_buffer::AlignedIndexBuffer;
use palettevec::palette::vec::VecPalette;
use rustc_hash::FxHashMap;

use super::{Block, TileEntity};

pub const WIDTH: usize = 16;
pub const HEIGHT: usize = 128;
pub const VOLUME: usize = WIDTH * WIDTH * HEIGHT;

/// Highest value of the height map, which never reaches [`HEIGHT`].
const MAX_HEIGHT: usize = HEIGHT - 1;

#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub struct ChunkPos {
    pub x: i32,
    pub z: i32,
}

impl ChunkPos {
    pub const fn new(x: i32, z: i32) -> Self {
        Self { x, z }
    }

    pub const fn containing(block_x: i32, block_z: i32) -> Self {
        Self::new(block_x >> 4, block_z >> 4)
    }

    pub const fn offset(self, dx: i32, dz: i32) -> Self {
        Self::new(self.x + dx, self.z + dz)
    }

    pub const fn min_block_x(self) -> i32 {
        self.x << 4
    }

    pub const fn min_block_z(self) -> i32 {
        self.z << 4
    }
}

/// Index into chunk arrays, using the original `x << 11 | z << 7 | y` layout.
pub const fn index(x: usize, y: usize, z: usize) -> usize {
    x << 11 | z << 7 | y
}

const SECTIONS: usize = HEIGHT / 16;
const SECTION_VOLUME: usize = VOLUME / SECTIONS;
const SECTION_BYTES: usize = SECTION_VOLUME / 2;

/// The 16³ section of an [`index`] and the position within it, `x << 8 | z << 4 | y & 15`.
const fn split(index: usize) -> (usize, usize) {
    (index >> 4 & (SECTIONS - 1), (index >> 3 & !15) | index & 15)
}

type NibbleSection = Arc<[u8; SECTION_BYTES]>;

/// A shared section for each value, for the sections that hold nothing else.
static UNIFORM: LazyLock<[NibbleSection; 16]> =
    LazyLock::new(|| std::array::from_fn(|value| Arc::new([value as u8 * 0x11; SECTION_BYTES])));

fn nibble_section(bytes: [u8; SECTION_BYTES]) -> NibbleSection {
    match uniform(&bytes) {
        Some(value) => UNIFORM[usize::from(value)].clone(),
        None => Arc::new(bytes),
    }
}

/// The value of every nibble, if they are all the same.
fn uniform(bytes: &[u8; SECTION_BYTES]) -> Option<u8> {
    let value = bytes[0] & 0xF;
    bytes.iter().all(|&b| b == value * 0x11).then_some(value)
}

/// 4-bit values like `NibbleArray`, packed per 16³ section, low nibble first. Sections that hold a
/// single value share one array, and clones share all sections until one of them is written.
#[derive(Clone)]
pub struct Nibbles([NibbleSection; SECTIONS]);

impl Nibbles {
    #[inline(always)]
    pub fn get(&self, index: usize) -> u8 {
        (self.byte(index >> 1) >> ((index & 1) << 2)) & 0xF
    }

    pub fn set(&mut self, index: usize, value: u8) {
        let value = value & 0xF;
        if self.get(index) == value {
            return;
        }
        set_nibble(
            Arc::make_mut(&mut self.0[index >> 4 & (SECTIONS - 1)]),
            index,
            value,
        );
    }

    /// Byte `index / 2` of the original array.
    #[inline(always)]
    fn byte(&self, byte: usize) -> u8 {
        self.0[byte >> 3 & (SECTIONS - 1)][local_byte(byte)]
    }

    /// Overwrites the original array's bytes from `start`.
    fn write_bytes(&mut self, mut start: usize, mut data: &[u8]) {
        // Runs of 8 bytes, half a column, stay within one section
        while !data.is_empty() {
            let len = (8 - (start & 7)).min(data.len());
            let local = local_byte(start);
            Arc::make_mut(&mut self.0[start >> 3 & (SECTIONS - 1)])[local..local + len]
                .copy_from_slice(&data[..len]);
            (start, data) = (start + len, &data[len..]);
        }
    }

    /// The original array.
    pub fn to_bytes(&self) -> Vec<u8> {
        (0..VOLUME / 2).map(|byte| self.byte(byte)).collect()
    }

    /// Whether both share all sections.
    pub fn ptr_eq(&self, other: &Self) -> bool {
        self.0.iter().zip(&other.0).all(|(a, b)| Arc::ptr_eq(a, b))
    }

    /// Shares the sections that hold a single value again, of those given as bits of their
    /// section y.
    pub fn compact(&mut self, sections: u8) {
        for (y, section) in self.0.iter_mut().enumerate() {
            if sections >> y & 1 != 0
                && !Arc::ptr_eq(section, &UNIFORM[usize::from(section[0] & 0xF)])
                && let Some(value) = uniform(section)
            {
                *section = UNIFORM[usize::from(value)].clone();
            }
        }
    }

    /// Reads the original array, or `None` if the length is wrong.
    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
        let bytes: &[u8; VOLUME / 2] = bytes.try_into().ok()?;
        Some(Self(std::array::from_fn(|section| {
            nibble_section(std::array::from_fn(|local| {
                bytes[(local & !7) << 3 | section << 3 | local & 7]
            }))
        })))
    }
}

/// Sets the nibble of an [`index`] within its section.
fn set_nibble(bytes: &mut [u8; SECTION_BYTES], index: usize, value: u8) {
    let (byte, shift) = (&mut bytes[local_byte(index >> 1)], (index & 1) << 2);
    *byte = (*byte & !(0xF << shift)) | (value << shift);
}

/// The byte within its section of byte `index / 2` of the original array.
const fn local_byte(byte: usize) -> usize {
    (byte >> 3 & (SECTION_BYTES - 8)) | byte & 7
}

impl Default for Nibbles {
    /// All zeros.
    fn default() -> Self {
        Self(std::array::from_fn(|_| UNIFORM[0].clone()))
    }
}

type BlockSection = PaletteVec<Block, VecPalette<Block>, AlignedIndexBuffer>;

#[derive(Clone)]
enum BlockStorage {
    Dense(Arc<[Block; VOLUME]>),
    Packed(Arc<[BlockSection; SECTIONS]>),
}

/// Block ids, either in the original array or palette compressed per 16³ section, which takes a
/// fraction of the memory but is slower to access. Clones share them until one of them is written.
#[derive(Clone)]
pub struct Blocks(BlockStorage);

impl Blocks {
    #[inline(always)]
    pub fn get(&self, index: usize) -> Block {
        match &self.0 {
            BlockStorage::Dense(blocks) => blocks[index],
            BlockStorage::Packed(sections) => packed_block(sections, index),
        }
    }

    fn replace(&mut self, index: usize, block: Block) -> Block {
        match &mut self.0 {
            BlockStorage::Dense(blocks) => {
                std::mem::replace(&mut Arc::make_mut(blocks)[index], block)
            }
            BlockStorage::Packed(sections) => {
                let (section, local) = split(index);
                let section = &mut Arc::make_mut(sections)[section];
                let old = section.get(local).copied().unwrap_or_default();
                section.set(local, &block);
                old
            }
        }
    }

    /// The ids in the order of the original array.
    pub fn ids(&self) -> impl Iterator<Item = u8> + '_ {
        (0..VOLUME).map(|index| self.get(index).0)
    }

    /// The original array, shared unless the ids are packed.
    pub fn dense(&self) -> Arc<[Block; VOLUME]> {
        match &self.0 {
            BlockStorage::Dense(blocks) => blocks.clone(),
            BlockStorage::Packed(_) => self.unpacked(),
        }
    }

    #[cold]
    #[inline(never)]
    fn unpacked(&self) -> Arc<[Block; VOLUME]> {
        let mut blocks = Arc::new([Block::AIR; VOLUME]);
        let dense = Arc::get_mut(&mut blocks).unwrap_or_else(|| unreachable!());
        dense
            .iter_mut()
            .zip(self.ids())
            .for_each(|(block, id)| *block = Block(id));
        blocks
    }

    pub fn is_packed(&self) -> bool {
        matches!(self.0, BlockStorage::Packed(_))
    }

    /// Whether both share the same ids.
    pub fn ptr_eq(&self, other: &Self) -> bool {
        match (&self.0, &other.0) {
            (BlockStorage::Dense(a), BlockStorage::Dense(b)) => Arc::ptr_eq(a, b),
            (BlockStorage::Packed(a), BlockStorage::Packed(b)) => Arc::ptr_eq(a, b),
            _ => false,
        }
    }

    /// A copy in the other form: packed if these are not, and the other way round.
    pub fn converted(&self) -> Self {
        let BlockStorage::Dense(blocks) = &self.0 else {
            return Self(BlockStorage::Dense(self.unpacked()));
        };
        Self(BlockStorage::Packed(Arc::new(std::array::from_fn(
            |section| {
                let block = |local: usize| blocks[(local & !15) << 3 | section << 4 | local & 15];
                let first = block(0);
                if (1..SECTION_VOLUME).all(|local| block(local) == first) {
                    return BlockSection::filled(first, SECTION_VOLUME);
                }
                let mut packed = BlockSection::new();
                (0..SECTION_VOLUME).for_each(|local| packed.push(block(local)));
                packed
            },
        ))))
    }
}

/// Kept out of line, so reading unpacked ids stays as cheap as indexing an array.
#[inline(never)]
fn packed_block(sections: &[BlockSection; SECTIONS], index: usize) -> Block {
    let (section, local) = split(index);
    sections[section].get(local).copied().unwrap_or_default()
}

impl From<Box<[Block; VOLUME]>> for Blocks {
    fn from(blocks: Box<[Block; VOLUME]>) -> Self {
        Self(BlockStorage::Dense(blocks.into()))
    }
}

/// Highest share of chunks whose block ids stay unpacked for being accessed the most. Measured
/// CPU time keeps falling up to this cap (by about a quarter against packing every settled
/// chunk), while the extra memory stays a few MB.
pub const HOT_SHARE: f64 = 0.25;
/// Ticks between rankings by heat, which also halve it.
const REPACK_PERIOD: u32 = 20;
/// Heat of a block write, which costs packed ids about as much as four reads.
const WRITE_HEAT: u32 = 4;

/// Recent accesses of a chunk's block ids, weighted by what each would cost on packed ids, in
/// block reads. On its own cache line, so counting never slows readers of the chunk on other
/// threads.
#[derive(Default)]
#[repr(align(64))]
pub struct Heat(AtomicU32);

impl Heat {
    #[inline(always)]
    pub fn add(&self, reads: u32) {
        // Not a read-modify-write: it is counted on the owner's thread, and a rare lost update
        // from another one does not matter
        let heat = self.0.load(Ordering::Relaxed);
        self.0.store(heat.saturating_add(reads), Ordering::Relaxed);
    }

    pub fn get(&self) -> u32 {
        self.0.load(Ordering::Relaxed)
    }

    /// Halves the heat, so it follows recent accesses.
    fn cool(&self) {
        self.0.store(self.get() / 2, Ordering::Relaxed);
    }
}

impl Clone for Heat {
    fn clone(&self) -> Self {
        Self(AtomicU32::new(self.get()))
    }
}

/// Chunks with the block ids they were given and the ids to take instead.
type Repacked = Vec<(ChunkPos, Blocks, Blocks)>;

/// Keeps the block ids of the hottest chunks unpacked and packs those of the other settled ones,
/// converting them on a worker.
#[derive(Default)]
pub struct Packer {
    running: Option<Mutex<Receiver<Repacked>>>,
    countdown: u32,
}

impl Packer {
    /// Call once per tick with the loaded chunks, each with whether it is settled, that is no
    /// longer written in bulk. Returns finished conversions for [`Chunk::adopt_blocks`]. Every
    /// [`REPACK_PERIOD`] ticks it ranks the chunks by heat, keeps the hottest [`HOT_SHARE`]
    /// unpacked, and starts converting the settled chunks that are in the other form.
    pub fn tick<'a>(
        &mut self,
        chunks: impl FnOnce() -> Vec<(ChunkPos, &'a Chunk, bool)>,
    ) -> Repacked {
        self.countdown = self.countdown.saturating_sub(1);
        if let Some(running) = &mut self.running {
            match running
                .get_mut()
                .unwrap_or_else(PoisonError::into_inner)
                .try_recv()
            {
                Ok(repacked) => {
                    self.running = None;
                    return repacked;
                }
                Err(TryRecvError::Empty) => return Vec::new(),
                Err(TryRecvError::Disconnected) => self.running = None,
            }
        }
        if self.countdown > 0 {
            return Vec::new();
        }
        self.countdown = REPACK_PERIOD;
        let mut chunks = chunks();
        let hot = (chunks.len() as f64 * HOT_SHARE) as usize;
        if hot > 0 {
            chunks.select_nth_unstable_by_key(hot - 1, |(_, chunk, _)| Reverse(chunk.heat.get()));
        }
        let due: Vec<_> = (chunks.into_iter().enumerate())
            .filter_map(|(rank, (pos, chunk, settled))| {
                chunk.heat.cool();
                let unpacked = rank < hot;
                (settled && chunk.blocks.is_packed() == unpacked)
                    .then(|| (pos, chunk.blocks.clone()))
            })
            .collect();
        if !due.is_empty() {
            let (sender, receiver) = mpsc::sync_channel(1);
            rayon::spawn(move || {
                let repacked = due.into_iter().map(|(pos, blocks)| {
                    let converted = blocks.converted();
                    (pos, blocks, converted)
                });
                drop(sender.send(repacked.collect()));
            });
            self.running = Some(Mutex::new(receiver));
        }
        Vec::new()
    }
}

/// Cheap to clone: the block and nibble arrays are shared until written.
#[derive(Clone)]
pub struct Chunk {
    pub pos: ChunkPos,
    blocks: Blocks,
    pub metadata: Nibbles,
    pub sky_light: Nibbles,
    pub block_light: Nibbles,
    height_map: [u8; WIDTH * WIDTH],
    /// Keyed by [`index`].
    pub tile_entities: FxHashMap<u16, TileEntity>,
    pub populated: bool,
    /// False in the Nether (`WorldProvider.hasNoSky`): no sky light at all.
    pub has_sky: bool,
    pub heat: Heat,
}

impl std::fmt::Debug for Chunk {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Chunk")
            .field("pos", &self.pos)
            .finish_non_exhaustive()
    }
}

impl Chunk {
    pub fn new(pos: ChunkPos, blocks: Box<[Block; VOLUME]>) -> Self {
        Self {
            pos,
            blocks: Blocks::from(blocks),
            metadata: Nibbles::default(),
            sky_light: Nibbles::default(),
            block_light: Nibbles::default(),
            height_map: [0; WIDTH * WIDTH],
            tile_entities: FxHashMap::default(),
            populated: false,
            has_sky: true,
            heat: Heat::default(),
        }
    }

    /// Drops the sky light, for dimensions without a sky.
    pub fn without_sky(mut self) -> Self {
        self.has_sky = false;
        self.sky_light = Nibbles::default();
        self
    }

    /// A chunk read from a save. Without a valid height map and sky light both are regenerated.
    pub fn from_saved(
        pos: ChunkPos,
        blocks: Box<[Block; VOLUME]>,
        metadata: Option<Nibbles>,
        sky_light: Option<Nibbles>,
        block_light: Option<Nibbles>,
        height_map: Option<[u8; WIDTH * WIDTH]>,
    ) -> Self {
        let mut chunk = Self::new(pos, blocks);
        chunk.metadata = metadata.unwrap_or_default();
        chunk.block_light = block_light.unwrap_or_default();
        match (sky_light, height_map) {
            (Some(sky_light), Some(height_map)) => {
                chunk.sky_light = sky_light;
                chunk.height_map = height_map;
            }
            _ => chunk.generate_sky_light(),
        }
        chunk
    }

    /// The raw height map bytes, indexed `z << 4 | x`.
    pub fn height_map(&self) -> &[u8; WIDTH * WIDTH] {
        &self.height_map
    }

    pub fn blocks(&self) -> &Blocks {
        &self.blocks
    }

    /// Takes `converted` as the block ids if they are still those it was converted from.
    pub fn adopt_blocks(&mut self, from: &Blocks, converted: Blocks) {
        if self.blocks.ptr_eq(from) {
            self.blocks = converted;
        }
    }

    #[inline]
    pub fn block(&self, x: usize, y: usize, z: usize) -> Block {
        self.heat.add(1);
        self.blocks.get(index(x, y, z))
    }

    /// Lowest y whose column above is fully transparent to light, capped at 127 like the original.
    pub fn height(&self, x: usize, z: usize) -> usize {
        usize::from(self.height_map[z << 4 | x])
    }

    /// `Chunk.canBlockSeeTheSky`.
    pub fn sees_sky(&self, x: usize, y: usize, z: usize) -> bool {
        y >= self.height(x, z)
    }

    fn scan_height(&self, x: usize, z: usize, from: usize) -> usize {
        let column = index(x, 0, z);
        let mut height = from;
        while height > 0 && self.blocks.get(column + height - 1).light_opacity() == 0 {
            height -= 1;
        }
        height
    }

    /// Port of `Chunk.generateSkylightMap` without the neighbor checks, which the
    /// light engine performs instead.
    pub fn generate_sky_light(&mut self) {
        // `ChunkProviderHell` never generates it, so Nether height maps stay zero
        if !self.has_sky {
            return;
        }
        self.generate_height_map();
        // Section by section from the top, each column carrying its light down
        let mut lights = [15_i32; WIDTH * WIDTH];
        for section in (0..SECTIONS).rev() {
            if lights.iter().all(|&light| light <= 0) {
                break;
            }
            let bytes = Arc::make_mut(&mut self.sky_light.0[section]);
            for (column, light) in columns(0..WIDTH, 0, 0..WIDTH).zip(&mut lights) {
                for y in (section * 16..section * 16 + 16).rev().filter(|&y| y > 0) {
                    if *light <= 0 {
                        break;
                    }
                    *light -= i32::from(self.blocks.get(column + y).light_opacity());
                    if *light > 0 {
                        set_nibble(bytes, column + y, *light as u8);
                    }
                }
            }
        }
        self.sky_light.compact(u8::MAX);
    }

    fn generate_height_map(&mut self) {
        for x in 0..WIDTH {
            for z in 0..WIDTH {
                self.height_map[z << 4 | x] = self.scan_height(x, z, MAX_HEIGHT) as u8;
            }
        }
    }

    /// Writes a block id and returns the previous one, the first step of `Chunk.setBlockID`.
    /// Follow with [`Self::finish_block_change`] once the old block has been notified.
    pub fn swap_block(&mut self, x: usize, y: usize, z: usize, block: Block) -> Block {
        self.heat.add(WRITE_HEAT);
        self.blocks.replace(index(x, y, z), block)
    }

    /// Second step of `Chunk.setBlockID`: stores the metadata (`None` for `setBlockID`, which
    /// clears it) and updates the height map and sky light of the column. Returns the y range
    /// whose sky light was rewritten.
    pub fn finish_block_change(
        &mut self,
        x: usize,
        y: usize,
        z: usize,
        metadata: Option<u8>,
    ) -> Option<Range<usize>> {
        self.metadata.set(index(x, y, z), metadata.unwrap_or(0));
        // Only `setBlockIDWithMetadata` skips the column without a sky
        if !self.has_sky && metadata.is_some() {
            return None;
        }
        let height = self.height(x, z);
        if self.blocks.get(index(x, y, z)).light_opacity() != 0 {
            (y >= height).then(|| self.relight_column(x, y + 1, z))?
        } else if y + 1 == height {
            self.relight_column(x, y, z)
        } else {
            None
        }
    }

    /// Port of `Chunk.relightBlock`.
    fn relight_column(&mut self, x: usize, y: usize, z: usize) -> Option<Range<usize>> {
        let old = self.height(x, z);
        let mut height = self.scan_height(x, z, y.max(old));
        if height == old {
            return None;
        }
        self.height_map[z << 4 | x] = height as u8;
        let column = index(x, 0, z);
        let lit = if height < old { 15 } else { 0 };
        for y in old.min(height)..old.max(height) {
            self.sky_light.set(column + y, lit);
        }

        let top = height.max(old);
        let mut light = 15_i32;
        while height > 0 && light > 0 {
            height -= 1;
            light =
                (light - i32::from(self.blocks.get(column + height).light_opacity().max(1))).max(0);
            self.sky_light.set(column + height, light as u8);
        }
        // A raised column was darkened from its old height up, which needs settling too
        Some(height.min(old)..top)
    }

    /// `Chunk.getChunkData`: appends the ids, metadata, block light and sky light of a box,
    /// column by column. `y` must start and end on even heights.
    pub fn write_region(
        &self,
        out: &mut Vec<u8>,
        xs: Range<usize>,
        ys: Range<usize>,
        zs: Range<usize>,
    ) {
        let columns = columns(xs, ys.start, zs);
        for start in columns.clone() {
            out.extend((start..start + ys.len()).map(|index| self.blocks.get(index).0));
        }
        for nibbles in [&self.metadata, &self.block_light, &self.sky_light] {
            for start in columns.clone() {
                out.extend((start / 2..(start + ys.len()) / 2).map(|byte| nibbles.byte(byte)));
            }
        }
    }

    /// `Chunk.setChunkData`: reads a box written by [`Self::write_region`] and rebuilds the
    /// height map. Returns the number of bytes consumed.
    pub fn read_region(
        &mut self,
        data: &[u8],
        xs: Range<usize>,
        ys: Range<usize>,
        zs: Range<usize>,
    ) -> usize {
        let whole = xs.len() == WIDTH && ys.len() == HEIGHT && zs.len() == WIDTH;
        let columns = columns(xs, ys.start, zs);
        let mut offset = 0;
        // A whole chunk arrives in the order of the original array
        if let Some(ids) = data.first_chunk::<VOLUME>().filter(|_| whole) {
            self.blocks = Blocks::from(Box::new(ids.map(Block)));
            offset = VOLUME;
        } else {
            self.heat
                .add((columns.clone().count() * ys.len()) as u32 * WRITE_HEAT);
            for start in columns.clone() {
                for (index, &id) in (start..start + ys.len()).zip(&data[offset..]) {
                    self.blocks.replace(index, Block(id));
                }
                offset += ys.len();
            }
        }
        for nibbles in [
            &mut self.metadata,
            &mut self.block_light,
            &mut self.sky_light,
        ] {
            for start in columns.clone() {
                let len = ys.len() / 2;
                nibbles.write_bytes(start / 2, &data[offset..offset + len]);
                offset += len;
            }
            nibbles.compact(u8::MAX);
        }
        self.generate_height_map();
        // Like the original renderer, tile entities only exist for their blocks
        let blocks = &self.blocks;
        self.tile_entities
            .retain(|&index, _| blocks.get(usize::from(index)).is_container());
        offset
    }

    pub fn tile_entity(&self, x: usize, y: usize, z: usize) -> Option<&TileEntity> {
        self.tile_entities.get(&(index(x, y, z) as u16))
    }

    pub fn tile_entity_mut(&mut self, x: usize, y: usize, z: usize) -> Option<&mut TileEntity> {
        self.tile_entities.get_mut(&(index(x, y, z) as u16))
    }
}

/// The [`index`] of `y` in each column of a box, x-major.
fn columns(xs: Range<usize>, y: usize, zs: Range<usize>) -> impl Iterator<Item = usize> + Clone {
    xs.flat_map(move |x| zs.clone().map(move |z| index(x, y, z)))
}

fn zeroed_box<T: Copy + Default, const N: usize>() -> Box<[T; N]> {
    vec![T::default(); N]
        .into_boxed_slice()
        .try_into()
        .unwrap_or_else(|_| unreachable!())
}

pub fn empty_blocks() -> Box<[Block; VOLUME]> {
    zeroed_box()
}