Chunk Rewrite With Paletted Containers (#91)

The current approach to managing chunk data is misconceived. This new approach uses genuine paletted containers and does not suffer from complexities caused by caching. As a result, memory usage (according to htop) in the terrain example with render distance = 32 has gone from 785 megs to 137 megs. That's 17.4% of the memory it used to use. Terrain generation speed was not affected.

build/block.rs

@@ -453,11 +453,6 @@ pub fn build() -> anyhow::Result<TokenStream> {
                 }
             }
 
-            pub(crate) const fn from_raw_unchecked(id: u16) -> Self {
-                debug_assert!(Self::from_raw(id).is_some());
-                Self(id)
-            }
-
             /// Returns the [`BlockKind`] of this block state.
             pub const fn to_kind(self) -> BlockKind {
                 match self.0 {

examples/combat.rs

@@ -88,9 +88,8 @@ impl Config for Game {
         let (_, world) = server.worlds.insert(DimensionId::default(), ());
         server.state = Some(server.player_lists.insert(()).0);
 
-        let dim = server.shared.dimension(DimensionId::default());
+        let min_y = world.chunks.min_y();
-        let min_y = dim.min_y;
+        let height = world.chunks.height();
-        let height = dim.height as usize;
 
         // Create circular arena.
         let size = 2;

examples/conway.rs

@@ -279,7 +279,7 @@ impl Config for Game {
             mem::swap(&mut server.state.board, &mut server.state.board_buf);
         }
 
-        let min_y = server.shared.dimensions().next().unwrap().1.min_y;
+        let min_y = world.chunks.min_y();
 
         for chunk_x in 0..Integer::div_ceil(&SIZE_X, &16) {
             for chunk_z in 0..Integer::div_ceil(&SIZE_Z, &16) {

examples/terrain.rs

@@ -203,8 +203,8 @@ impl Config for Game {
 
                     // Add grass
                     for y in (0..chunk.height()).rev() {
-                        if chunk.get_block_state(x, y, z).is_air()
+                        if chunk.block_state(x, y, z).is_air()
-                            && chunk.get_block_state(x, y - 1, z) == BlockState::GRASS_BLOCK
+                            && chunk.block_state(x, y - 1, z) == BlockState::GRASS_BLOCK
                         {
                             let density = fbm(
                                 &self.grass_noise,
@@ -215,7 +215,7 @@ impl Config for Game {
                             );
 
                             if density > 0.55 {
-                                if density > 0.7 && chunk.get_block_state(x, y + 1, z).is_air() {
+                                if density > 0.7 && chunk.block_state(x, y + 1, z).is_air() {
                                     let upper = BlockState::TALL_GRASS
                                         .set(PropName::Half, PropValue::Upper);
                                     let lower = BlockState::TALL_GRASS

src/chunk/paletted_container.rs

@@ -0,0 +1,310 @@
+use std::array;
+use std::io::Write;
+
+use arrayvec::ArrayVec;
+
+use crate::chunk::{compact_u64s_len, encode_compact_u64s};
+use crate::protocol::{Encode, VarInt};
+use crate::util::log2_ceil;
+
+/// `HALF_LEN` must be equal to `ceil(LEN / 2)`.
+#[derive(Clone, Debug)]
+pub enum PalettedContainer<T: PalettedContainerElement, const LEN: usize, const HALF_LEN: usize> {
+    Single(T),
+    Indirect(Box<Indirect<T, LEN, HALF_LEN>>),
+    Direct(Box<[T; LEN]>),
+}
+
+pub trait PalettedContainerElement: Copy + Eq + Default {
+    /// The minimum number of bits required to represent all instances of the
+    /// element type. If `N` is the total number of possible values, then
+    /// `DIRECT_BITS` is `ceil(log2(N))`.
+    const DIRECT_BITS: usize;
+    /// The maximum number of bits per element allowed in the indirect
+    /// representation while encoding. Any higher than this will force
+    /// conversion to the direct representation while encoding.
+    const MAX_INDIRECT_BITS: usize;
+    /// The minimum number of bits used to represent the element type in the
+    /// indirect representation while encoding. If the bits per index is lower,
+    /// it will be rounded up to this.
+    const MIN_INDIRECT_BITS: usize;
+    /// Converts the element type to bits. The output must be less than two to
+    /// the power of `DIRECT_BITS`.
+    fn to_bits(self) -> u64;
+}
+
+#[derive(Clone, Debug)]
+pub struct Indirect<T: PalettedContainerElement, const LEN: usize, const HALF_LEN: usize> {
+    /// Each element is a unique instance of `T`.
+    palette: ArrayVec<T, 16>,
+    /// Each half-byte is an index into `palette`.
+    indices: [u8; HALF_LEN],
+}
+
+impl<T: PalettedContainerElement, const LEN: usize, const HALF_LEN: usize>
+    PalettedContainer<T, LEN, HALF_LEN>
+{
+    pub fn new() -> Self {
+        assert_eq!(num::Integer::div_ceil(&LEN, &2), HALF_LEN);
+        assert_ne!(LEN, 0);
+
+        Self::Single(T::default())
+    }
+
+    pub fn fill(&mut self, val: T) {
+        *self = Self::Single(val)
+    }
+
+    pub fn get(&self, idx: usize) -> T {
+        self.check_oob(idx);
+
+        match self {
+            Self::Single(elem) => *elem,
+            Self::Indirect(ind) => ind.get(idx),
+            Self::Direct(elems) => elems[idx],
+        }
+    }
+
+    pub fn set(&mut self, idx: usize, val: T) -> T {
+        self.check_oob(idx);
+
+        match self {
+            Self::Single(old_val) => {
+                if *old_val == val {
+                    *old_val
+                } else {
+                    // Upgrade to indirect.
+                    let old = *old_val;
+                    let mut ind = Box::new(Indirect {
+                        palette: ArrayVec::from_iter([old, val]),
+                        // All indices are initialized to index 0 (the old element).
+                        indices: [0; HALF_LEN],
+                    });
+
+                    ind.indices[idx / 2] = 1 << (idx % 2 * 4);
+                    *self = Self::Indirect(ind);
+                    old
+                }
+            }
+            Self::Indirect(ind) => {
+                if let Some(old) = ind.set(idx, val) {
+                    old
+                } else {
+                    // Upgrade to direct.
+                    *self = Self::Direct(Box::new(array::from_fn(|i| ind.get(i))));
+                    self.set(idx, val)
+                }
+            }
+            Self::Direct(vals) => {
+                let old = vals[idx];
+                vals[idx] = val;
+                old
+            }
+        }
+    }
+
+    pub fn optimize(&mut self) {
+        match self {
+            Self::Single(_) => {}
+            Self::Indirect(ind) => {
+                let mut new_ind = Indirect {
+                    palette: ArrayVec::new(),
+                    indices: [0; HALF_LEN],
+                };
+
+                for i in 0..LEN {
+                    new_ind.set(i, ind.get(i));
+                }
+
+                if new_ind.palette.len() == 1 {
+                    *self = Self::Single(new_ind.palette[0]);
+                } else {
+                    **ind = new_ind;
+                }
+            }
+            Self::Direct(dir) => {
+                let mut ind = Indirect {
+                    palette: ArrayVec::new(),
+                    indices: [0; HALF_LEN],
+                };
+
+                for (i, val) in dir.iter().cloned().enumerate() {
+                    if ind.set(i, val).is_none() {
+                        return;
+                    }
+                }
+
+                *self = if ind.palette.len() == 1 {
+                    Self::Single(ind.palette[0])
+                } else {
+                    Self::Indirect(Box::new(ind))
+                };
+            }
+        }
+    }
+
+    #[inline]
+    fn check_oob(&self, idx: usize) {
+        assert!(
+            idx < LEN,
+            "index {idx} is out of bounds in paletted container of length {LEN}"
+        );
+    }
+}
+
+impl<T: PalettedContainerElement, const LEN: usize, const HALF_LEN: usize> Default
+    for PalettedContainer<T, LEN, HALF_LEN>
+{
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<T: PalettedContainerElement, const LEN: usize, const HALF_LEN: usize>
+    Indirect<T, LEN, HALF_LEN>
+{
+    pub fn get(&self, idx: usize) -> T {
+        let palette_idx = self.indices[idx / 2] >> (idx % 2 * 4) & 0b1111;
+        self.palette[palette_idx as usize]
+    }
+
+    pub fn set(&mut self, idx: usize, val: T) -> Option<T> {
+        let palette_idx = if let Some(i) = self.palette.iter().position(|v| *v == val) {
+            i
+        } else {
+            self.palette.try_push(val).ok()?;
+            self.palette.len() - 1
+        };
+
+        let old_val = self.get(idx);
+        let u8 = &mut self.indices[idx / 2];
+        let shift = idx % 2 * 4;
+        *u8 = (*u8 & !(0b1111 << shift)) | ((palette_idx as u8) << shift);
+        Some(old_val)
+    }
+}
+
+/// Encodes the paletted container in the format that Minecraft expects.
+impl<T: PalettedContainerElement, const LEN: usize, const HALF_LEN: usize> Encode
+    for PalettedContainer<T, LEN, HALF_LEN>
+{
+    fn encode(&self, w: &mut impl Write) -> anyhow::Result<()> {
+        assert!(T::DIRECT_BITS <= 64);
+        assert!(T::MAX_INDIRECT_BITS <= 64);
+        assert!(T::MIN_INDIRECT_BITS <= T::MAX_INDIRECT_BITS);
+        assert!(T::MIN_INDIRECT_BITS <= 4);
+
+        match self {
+            Self::Single(val) => {
+                // Bits per entry
+                0_u8.encode(w)?;
+
+                // Palette
+                VarInt(val.to_bits() as i32).encode(w)?;
+
+                // Number of longs
+                VarInt(0).encode(w)?;
+            }
+            Self::Indirect(ind) => {
+                let bits_per_entry = T::MIN_INDIRECT_BITS.max(log2_ceil(ind.palette.len()));
+
+                // TODO: if bits_per_entry > MAX_INDIRECT_BITS, encode as direct.
+                debug_assert!(bits_per_entry <= T::MAX_INDIRECT_BITS);
+
+                // Bits per entry
+                (bits_per_entry as u8).encode(w)?;
+
+                // Palette len
+                VarInt(ind.palette.len() as i32).encode(w)?;
+                // Palette
+                for val in &ind.palette {
+                    VarInt(val.to_bits() as i32).encode(w)?;
+                }
+
+                // Number of longs in data array.
+                VarInt(compact_u64s_len(LEN, bits_per_entry) as _).encode(w)?;
+                // Data array
+                encode_compact_u64s(
+                    w,
+                    ind.indices
+                        .iter()
+                        .cloned()
+                        .flat_map(|byte| [byte & 0b1111, byte >> 4])
+                        .map(u64::from)
+                        .take(LEN),
+                    bits_per_entry,
+                )?;
+            }
+            Self::Direct(dir) => {
+                // Bits per entry
+                (T::DIRECT_BITS as u8).encode(w)?;
+
+                // Number of longs in data array.
+                VarInt(compact_u64s_len(LEN, T::DIRECT_BITS) as _).encode(w)?;
+                // Data array
+                encode_compact_u64s(w, dir.iter().map(|v| v.to_bits()), T::DIRECT_BITS)?;
+            }
+        }
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rand::Rng;
+
+    use super::*;
+
+    fn check<T: PalettedContainerElement, const LEN: usize, const HALF_LEN: usize>(
+        p: &PalettedContainer<T, LEN, HALF_LEN>,
+        s: &[T],
+    ) -> bool {
+        assert_eq!(s.len(), LEN);
+        (0..LEN).all(|i| p.get(i) == s[i])
+    }
+
+    impl PalettedContainerElement for u32 {
+        const DIRECT_BITS: usize = 0;
+        const MAX_INDIRECT_BITS: usize = 0;
+        const MIN_INDIRECT_BITS: usize = 0;
+
+        fn to_bits(self) -> u64 {
+            self.into()
+        }
+    }
+
+    #[test]
+    fn random_assignments() {
+        const LEN: usize = 100;
+        let range = 0..64;
+
+        let mut rng = rand::thread_rng();
+
+        for _ in 0..20 {
+            let mut p = PalettedContainer::<u32, LEN, { LEN / 2 }>::new();
+
+            let init = rng.gen_range(range.clone());
+
+            p.fill(init);
+            let mut a = [init; LEN];
+
+            assert!(check(&p, &a));
+
+            let mut rng = rand::thread_rng();
+
+            for _ in 0..LEN * 10 {
+                let idx = rng.gen_range(0..LEN);
+                let val = rng.gen_range(range.clone());
+
+                assert_eq!(p.get(idx), p.set(idx, val));
+                assert_eq!(val, p.get(idx));
+                a[idx] = val;
+
+                p.optimize();
+
+                assert!(check(&p, &a));
+            }
+        }
+    }
+}

src/client.rs

@@ -1252,7 +1252,6 @@ impl<C: Config> Client<C> {
             if let Some(chunk) = world.chunks.get(pos) {
                 if self.loaded_chunks.insert(pos) {
                     self.send_packet(chunk.chunk_data_packet(pos));
-                    chunk.block_change_packets(pos, dimension.min_y, |pkt| self.send_packet(pkt));
                 }
             }
         }

src/config.rs

@@ -2,7 +2,6 @@
 
 use std::borrow::Cow;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
-use std::panic::{RefUnwindSafe, UnwindSafe};
 
 use async_trait::async_trait;
 use serde::Serialize;
@@ -24,7 +23,7 @@ use crate::{Ticks, STANDARD_TPS};
 /// [async_trait]: https://docs.rs/async-trait/latest/async_trait/
 #[async_trait]
 #[allow(unused_variables)]
-pub trait Config: Sized + Send + Sync + UnwindSafe + RefUnwindSafe + 'static {
+pub trait Config: Sized + Send + Sync + 'static {
     /// Custom state to store with the [`Server`].
     type ServerState: Send + Sync;
     /// Custom state to store with every [`Client`](crate::client::Client).
@@ -306,3 +305,33 @@ pub struct PlayerSampleEntry<'a> {
     /// The player UUID.
     pub id: Uuid,
 }
+
+/// A minimal `Config` implementation for testing purposes.
+#[cfg(test)]
+pub(crate) struct MockConfig<S = (), Cl = (), E = (), W = (), Ch = (), P = ()> {
+    _marker: std::marker::PhantomData<(S, Cl, E, W, Ch, P)>,
+}
+
+#[cfg(test)]
+impl<S, Cl, E, W, Ch, P> Config for MockConfig<S, Cl, E, W, Ch, P>
+where
+    S: Send + Sync + 'static,
+    Cl: Default + Send + Sync + 'static,
+    E: Send + Sync + 'static,
+    W: Send + Sync + 'static,
+    Ch: Send + Sync + 'static,
+    P: Send + Sync + 'static,
+{
+    type ServerState = S;
+    type ClientState = Cl;
+    type EntityState = E;
+    type WorldState = W;
+    type ChunkState = Ch;
+    type PlayerListState = P;
+
+    fn max_connections(&self) -> usize {
+        64
+    }
+
+    fn update(&self, _server: &mut Server<Self>) {}
+}

src/entity.rs

@@ -821,25 +821,9 @@ mod tests {
     use uuid::Uuid;
 
     use super::{Entities, EntityId, EntityKind};
-    use crate::config::Config;
-    use crate::server::Server;
     use crate::slab_versioned::Key;
 
-    /// Created for the sole purpose of use during unit tests.
+    type MockConfig = crate::config::MockConfig<(), (), u8>;
-    struct MockConfig;
-    impl Config for MockConfig {
-        type ServerState = ();
-        type ClientState = ();
-        type EntityState = u8; // Just for identification purposes
-        type WorldState = ();
-        type ChunkState = ();
-        type PlayerListState = ();
-
-        fn max_connections(&self) -> usize {
-            10
-        }
-        fn update(&self, _server: &mut Server<Self>) {}
-    }
 
     #[test]
     fn entities_has_valid_new_state() {

src/server.rs

@@ -440,12 +440,6 @@ fn do_update_loop<C: Config>(server: &mut Server<C>) -> ShutdownResult {
         shared.config().update(server);
 
         server.worlds.par_iter_mut().for_each(|(id, world)| {
-            // Chunks created this tick can have their changes applied immediately because
-            // they have not been observed by clients yet. Clients will not have to be sent
-            // the block change packet in this case, since the changes are applied before we
-            // update clients.
-            world.chunks.update_created_this_tick();
-
             world.spatial_index.update(&server.entities, id);
         });
 

src/util.rs

@@ -127,6 +127,17 @@ pub fn ray_box_intersect(ro: Vec3<f64>, rd: Vec3<f64>, bb: Aabb<f64>) -> Option<
     }
 }
 
+/// Calculates the log base 2 rounded up.
+pub(crate) const fn log2_ceil(n: usize) -> usize {
+    debug_assert!(n != 0);
+
+    // TODO: replace with `n.wrapping_next_power_of_two().trailing_zeros()`.
+    match n.checked_next_power_of_two() {
+        Some(n) => n.trailing_zeros() as usize,
+        None => 0_u64.trailing_zeros() as usize,
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use approx::assert_relative_eq;

src/world.rs

@@ -44,12 +44,18 @@ impl<C: Config> Worlds<C> {
 
     /// Creates a new world on the server with the provided dimension. A
     /// reference to the world along with its ID is returned.
-    pub fn insert(&mut self, dim: DimensionId, state: C::WorldState) -> (WorldId, &mut World<C>) {
+    pub fn insert(
+        &mut self,
+        dimension: DimensionId,
+        state: C::WorldState,
+    ) -> (WorldId, &mut World<C>) {
+        let dim = self.shared.dimension(dimension);
+
         let (id, world) = self.slab.insert(World {
             state,
             spatial_index: SpatialIndex::new(),
-            chunks: Chunks::new(self.shared.clone(), dim),
+            chunks: Chunks::new(dim.height, dim.min_y),
-            meta: WorldMeta { dimension: dim },
+            meta: WorldMeta { dimension },
         });
 
         (WorldId(id), world)