make BlockState a u16 and add a BlockStateIntegerRepr type

9 files changed, 121 insertions, 46 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 0f935ae8..d42f4594 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -125,6 +125,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c1fd03a028ef38ba2276dce7e33fcd6369c158a1bca17946c4b1b701891c1ff7"
 
 [[package]]
+name = "arrayvec"
+version = "0.7.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
+
+[[package]]
 name = "as-any"
 version = "0.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -227,6 +233,7 @@ dependencies = [
  "futures",
  "futures-lite",
  "nohash-hasher",
+ "num-format",
  "num-traits",
  "parking_lot",
  "priority-queue",
@@ -2058,6 +2065,16 @@ dependencies = [
 ]
 
 [[package]]
+name = "num-format"
+version = "0.4.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a652d9771a63711fd3c3deb670acfbe5c30a4072e664d7a3bf5a9e1056ac72c3"
+dependencies = [
+ "arrayvec",
+ "itoa",
+]
+
+[[package]]
 name = "num-integer"
 version = "0.1.46"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/azalea-block/azalea-block-macros/src/lib.rs b/azalea-block/azalea-block-macros/src/lib.rs
index de13ff6e..9374527f 100755
--- a/azalea-block/azalea-block-macros/src/lib.rs
+++ b/azalea-block/azalea-block-macros/src/lib.rs
@@ -19,6 +19,8 @@ use syn::{
 };
 use utils::{combinations_of, to_pascal_case};
 
+// must be the same as the type in `azalea-block/src/lib.rs`
+type BlockStateIntegerRepr = u16;
 enum PropertyType {
     /// `Axis { X, Y, Z }`
     Enum {
@@ -275,7 +277,7 @@ impl Parse for MakeBlockStates {
 }
 
 struct PropertyVariantData {
-    pub block_state_ids: Vec<u32>,
+    pub block_state_ids: Vec<BlockStateIntegerRepr>,
     pub ident: Ident,
     pub is_enum: bool,
 }
@@ -288,7 +290,7 @@ pub fn make_block_states(input: TokenStream) -> TokenStream {
     let mut properties_map = HashMap::new();
     let mut property_struct_names_to_names = HashMap::new();
 
-    let mut state_id: u32 = 0;
+    let mut state_id: BlockStateIntegerRepr = 0;
 
     for property in &input.property_definitions.properties {
         let property_struct_name: Ident;
@@ -339,8 +341,8 @@ pub fn make_block_states(input: TokenStream) -> TokenStream {
                         #property_enum_variants
                     }
 
-                    impl From<u32> for #property_struct_name {
-                        fn from(value: u32) -> Self {
+                    impl From<crate::BlockStateIntegerRepr> for #property_struct_name {
+                        fn from(value: crate::BlockStateIntegerRepr) -> Self {
                             match value {
                                 #property_from_number_variants
                                 _ => panic!("Invalid property value: {}", value),
@@ -358,8 +360,8 @@ pub fn make_block_states(input: TokenStream) -> TokenStream {
                     #[derive(Debug, Clone, Copy, PartialEq, Eq)]
                     pub struct #property_struct_name(pub bool);
 
-                    impl From<u32> for #property_struct_name {
-                        fn from(value: u32) -> Self {
+                    impl From<crate::BlockStateIntegerRepr> for #property_struct_name {
+                        fn from(value: crate::BlockStateIntegerRepr) -> Self {
                             match value {
                                 0 => Self(false),
                                 1 => Self(true),
@@ -583,7 +585,7 @@ pub fn make_block_states(input: TokenStream) -> TokenStream {
         //     }
         // }
         let mut from_state_to_block_inner = quote! {};
-        let mut division = 1u32;
+        let mut division: BlockStateIntegerRepr = 1;
         for i in (0..properties_with_name.len()).rev() {
             let PropertyWithNameAndDefault {
                 property_type: property_struct_name_ident,
@@ -593,7 +595,7 @@ pub fn make_block_states(input: TokenStream) -> TokenStream {
             } = &properties_with_name[i];
 
             let property_variants = &block_properties_vec[i];
-            let property_variants_count = property_variants.len() as u32;
+            let property_variants_count = property_variants.len() as crate::BlockStateIntegerRepr;
             let conversion_code = {
                 if &property_value_type.to_string() == "bool" {
                     assert_eq!(property_variants_count, 2);
@@ -695,7 +697,7 @@ pub fn make_block_states(input: TokenStream) -> TokenStream {
     let mut generated = quote! {
         impl BlockState {
             /// The highest possible block state ID.
-            pub const MAX_STATE: u32 = #last_state_id;
+            pub const MAX_STATE: crate::BlockStateIntegerRepr = #last_state_id;
 
             /// Get a property from this block state. Will be `None` if the block can't have the property.
             ///
diff --git a/azalea-block/src/lib.rs b/azalea-block/src/lib.rs
index 3de0aa5a..65dd581e 100755
--- a/azalea-block/src/lib.rs
+++ b/azalea-block/src/lib.rs
@@ -40,22 +40,34 @@ pub trait Property {
     fn try_from_block_state(state: BlockState) -> Option<Self::Value>;
 }
 
+/// The type that's used internally to represent a block state ID.
+///
+/// This should be either `u16` or `u32`. If you choose to modify it, you must
+/// also change it in `azalea-block-macros/src/lib.rs`.
+///
+/// This does not affect protocol serialization, it just allows you to make the
+/// internal type smaller if you want.
+pub type BlockStateIntegerRepr = u16;
+
 /// A representation of a state a block can be in.
 ///
 /// For example, a stone block only has one state but each possible stair
 /// rotation is a different state.
+///
+/// Note that this type is internally either a `u16` or `u32`, depending on
+/// [`BlockStateIntegerRepr`].
 #[derive(Copy, Clone, PartialEq, Eq, Default, Hash)]
 pub struct BlockState {
     /// The protocol ID for the block state. IDs may change every
     /// version, so you shouldn't hard-code them or store them in databases.
-    pub id: u32,
+    pub id: BlockStateIntegerRepr,
 }
 
 impl BlockState {
     pub const AIR: BlockState = BlockState { id: 0 };
 
     #[inline]
-    pub fn is_valid_state(state_id: u32) -> bool {
+    pub fn is_valid_state(state_id: BlockStateIntegerRepr) -> bool {
         state_id <= Self::MAX_STATE
     }
 
@@ -70,8 +82,22 @@ impl BlockState {
 impl TryFrom<u32> for BlockState {
     type Error = ();
 
-    /// Safely converts a state id to a block state.
+    /// Safely converts a u32 state id to a block state.
     fn try_from(state_id: u32) -> Result<Self, Self::Error> {
+        let state_id = state_id as BlockStateIntegerRepr;
+        if Self::is_valid_state(state_id) {
+            Ok(BlockState { id: state_id })
+        } else {
+            Err(())
+        }
+    }
+}
+impl TryFrom<u16> for BlockState {
+    type Error = ();
+
+    /// Safely converts a u16 state id to a block state.
+    fn try_from(state_id: u16) -> Result<Self, Self::Error> {
+        let state_id = state_id as BlockStateIntegerRepr;
         if Self::is_valid_state(state_id) {
             Ok(BlockState { id: state_id })
         } else {
@@ -90,7 +116,7 @@ impl AzaleaRead for BlockState {
 }
 impl AzaleaWrite for BlockState {
     fn azalea_write(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
-        u32::azalea_write_var(&self.id, buf)
+        u32::azalea_write_var(&(self.id as u32), buf)
     }
 }
 
@@ -158,12 +184,16 @@ impl From<FluidState> for BlockState {
             azalea_registry::Fluid::Empty => BlockState::AIR,
             azalea_registry::Fluid::Water | azalea_registry::Fluid::FlowingWater => {
                 BlockState::from(crate::blocks::Water {
-                    level: crate::properties::WaterLevel::from(state.height as u32),
+                    level: crate::properties::WaterLevel::from(
+                        state.height as BlockStateIntegerRepr,
+                    ),
                 })
             }
             azalea_registry::Fluid::Lava | azalea_registry::Fluid::FlowingLava => {
                 BlockState::from(crate::blocks::Lava {
-                    level: crate::properties::LavaLevel::from(state.height as u32),
+                    level: crate::properties::LavaLevel::from(
+                        state.height as BlockStateIntegerRepr,
+                    ),
                 })
             }
         }
@@ -182,7 +212,10 @@ mod tests {
 
     #[test]
     fn test_from_u32() {
-        assert_eq!(BlockState::try_from(0).unwrap(), BlockState::AIR);
+        assert_eq!(
+            BlockState::try_from(0 as BlockStateIntegerRepr).unwrap(),
+            BlockState::AIR
+        );
 
         assert!(BlockState::try_from(BlockState::MAX_STATE).is_ok());
         assert!(BlockState::try_from(BlockState::MAX_STATE + 1).is_err());
diff --git a/azalea-block/src/range.rs b/azalea-block/src/range.rs
index 9b520d49..779d805d 100644
--- a/azalea-block/src/range.rs
+++ b/azalea-block/src/range.rs
@@ -3,15 +3,15 @@ use std::{
     ops::{Add, RangeInclusive},
 };
 
-use crate::BlockState;
+use crate::{BlockState, BlockStateIntegerRepr};
 
 #[derive(Debug, Clone)]
 pub struct BlockStates {
     pub set: HashSet<BlockState>,
 }
 
-impl From<RangeInclusive<u32>> for BlockStates {
-    fn from(range: RangeInclusive<u32>) -> Self {
+impl From<RangeInclusive<BlockStateIntegerRepr>> for BlockStates {
+    fn from(range: RangeInclusive<BlockStateIntegerRepr>) -> Self {
         let mut set = HashSet::with_capacity((range.end() - range.start() + 1) as usize);
         for id in range {
             set.insert(BlockState { id });
diff --git a/azalea-world/src/chunk_storage.rs b/azalea-world/src/chunk_storage.rs
index 7eb10471..4ea54628 100755
--- a/azalea-world/src/chunk_storage.rs
+++ b/azalea-world/src/chunk_storage.rs
@@ -7,7 +7,7 @@ use std::{
     sync::{Arc, Weak},
 };
 
-use azalea_block::BlockState;
+use azalea_block::{BlockState, BlockStateIntegerRepr};
 use azalea_buf::{AzaleaRead, AzaleaWrite, BufReadError};
 use azalea_core::position::{BlockPos, ChunkBlockPos, ChunkPos, ChunkSectionBlockPos};
 use nohash_hasher::IntMap;
@@ -450,7 +450,7 @@ impl AzaleaRead for Section {
         let states = PalettedContainer::read_with_type(buf, &PalettedContainerKind::BlockStates)?;
 
         for i in 0..states.storage.size() {
-            if !BlockState::is_valid_state(states.storage.get(i) as u32) {
+            if !BlockState::is_valid_state(states.storage.get(i) as BlockStateIntegerRepr) {
                 return Err(BufReadError::Custom(format!(
                     "Invalid block state {} (index {i}) found in section.",
                     states.storage.get(i)
diff --git a/azalea-world/src/palette.rs b/azalea-world/src/palette.rs
index 24dbe47a..269e443b 100755
--- a/azalea-world/src/palette.rs
+++ b/azalea-world/src/palette.rs
@@ -1,5 +1,6 @@
 use std::io::{Cursor, Write};
 
+use azalea_block::BlockStateIntegerRepr;
 use azalea_buf::{AzaleaRead, AzaleaReadVar, AzaleaWrite, AzaleaWriteVar, BufReadError};
 use azalea_core::math;
 use tracing::warn;
@@ -110,7 +111,7 @@ impl PalettedContainer {
     /// This function panics if the index is greater than or equal to the number
     /// of things in the storage. (So for block states, it must be less than
     /// 4096).
-    pub fn get_at_index(&self, index: usize) -> u32 {
+    pub fn get_at_index(&self, index: usize) -> BlockStateIntegerRepr {
         // first get the palette id
         let paletted_value = self.storage.get(index);
         // and then get the value from that id
@@ -118,28 +119,35 @@ impl PalettedContainer {
     }
 
     /// Returns the value at the given coordinates.
-    pub fn get(&self, x: usize, y: usize, z: usize) -> u32 {
+    pub fn get(&self, x: usize, y: usize, z: usize) -> BlockStateIntegerRepr {
         // let paletted_value = self.storage.get(self.get_index(x, y, z));
         // self.palette.value_for(paletted_value as usize)
         self.get_at_index(self.index_from_coords(x, y, z))
     }
 
     /// Sets the id at the given coordinates and return the previous id
-    pub fn get_and_set(&mut self, x: usize, y: usize, z: usize, value: u32) -> u32 {
+    pub fn get_and_set(
+        &mut self,
+        x: usize,
+        y: usize,
+        z: usize,
+        value: BlockStateIntegerRepr,
+    ) -> BlockStateIntegerRepr {
         let paletted_value = self.id_for(value);
         self.storage
-            .get_and_set(self.index_from_coords(x, y, z), paletted_value as u64) as u32
+            .get_and_set(self.index_from_coords(x, y, z), paletted_value as u64)
+            as BlockStateIntegerRepr
     }
 
     /// Sets the id at the given index and return the previous id. You probably
     /// want `.set` instead.
-    pub fn set_at_index(&mut self, index: usize, value: u32) {
+    pub fn set_at_index(&mut self, index: usize, value: BlockStateIntegerRepr) {
         let paletted_value = self.id_for(value);
         self.storage.set(index, paletted_value as u64);
     }
 
     /// Sets the id at the given coordinates and return the previous id
-    pub fn set(&mut self, x: usize, y: usize, z: usize, value: u32) {
+    pub fn set(&mut self, x: usize, y: usize, z: usize, value: BlockStateIntegerRepr) {
         self.set_at_index(self.index_from_coords(x, y, z), value);
     }
 
@@ -168,7 +176,7 @@ impl PalettedContainer {
         }
     }
 
-    fn on_resize(&mut self, bits_per_entry: u8, value: u32) -> usize {
+    fn on_resize(&mut self, bits_per_entry: u8, value: BlockStateIntegerRepr) -> usize {
         // in vanilla this is always true, but it's sometimes false in purpur servers
         // assert!(bits_per_entry <= 5, "bits_per_entry must be <= 5");
         let mut new_data = self.create_or_reuse_data(bits_per_entry);
@@ -185,7 +193,7 @@ impl PalettedContainer {
         }
     }
 
-    pub fn id_for(&mut self, value: u32) -> usize {
+    pub fn id_for(&mut self, value: BlockStateIntegerRepr) -> usize {
         match &mut self.palette {
             Palette::SingleValue(v) => {
                 if *v != value {
@@ -245,21 +253,21 @@ pub enum PaletteKind {
 #[derive(Clone, Debug)]
 pub enum Palette {
     /// ID of the corresponding entry in its global palette
-    SingleValue(u32),
+    SingleValue(BlockStateIntegerRepr),
     // in vanilla this keeps a `size` field that might be less than the length, but i'm not sure
     // it's actually needed?
-    Linear(Vec<u32>),
-    Hashmap(Vec<u32>),
+    Linear(Vec<BlockStateIntegerRepr>),
+    Hashmap(Vec<BlockStateIntegerRepr>),
     Global,
 }
 
 impl Palette {
-    pub fn value_for(&self, id: usize) -> u32 {
+    pub fn value_for(&self, id: usize) -> BlockStateIntegerRepr {
         match self {
             Palette::SingleValue(v) => *v,
             Palette::Linear(v) => v[id],
             Palette::Hashmap(v) => v.get(id).copied().unwrap_or_default(),
-            Palette::Global => id as u32,
+            Palette::Global => id as BlockStateIntegerRepr,
         }
     }
 }
@@ -301,9 +309,17 @@ impl PaletteKind {
 
     pub fn read(&self, buf: &mut Cursor<&[u8]>) -> Result<Palette, BufReadError> {
         Ok(match self {
-            PaletteKind::SingleValue => Palette::SingleValue(u32::azalea_read_var(buf)?),
-            PaletteKind::Linear => Palette::Linear(Vec::<u32>::azalea_read_var(buf)?),
-            PaletteKind::Hashmap => Palette::Hashmap(Vec::<u32>::azalea_read_var(buf)?),
+            // since they're read as varints it's actually fine to just use BlockStateIntegerRepr
+            // instead of the correct type (u32)
+            PaletteKind::SingleValue => {
+                Palette::SingleValue(BlockStateIntegerRepr::azalea_read_var(buf)?)
+            }
+            PaletteKind::Linear => {
+                Palette::Linear(Vec::<BlockStateIntegerRepr>::azalea_read_var(buf)?)
+            }
+            PaletteKind::Hashmap => {
+                Palette::Hashmap(Vec::<BlockStateIntegerRepr>::azalea_read_var(buf)?)
+            }
             PaletteKind::Global => Palette::Global,
         })
     }
diff --git a/azalea/Cargo.toml b/azalea/Cargo.toml
index 59f253bd..d1bf4884 100644
--- a/azalea/Cargo.toml
+++ b/azalea/Cargo.toml
@@ -35,6 +35,7 @@ derive_more = { workspace = true, features = ["deref", "deref_mut"] }
 futures = { workspace = true }
 futures-lite = { workspace = true }
 nohash-hasher = { workspace = true }
+num-format = "0.4.4"
 num-traits = { workspace = true }
 parking_lot = { workspace = true }
 priority-queue = { workspace = true }
diff --git a/azalea/src/pathfinder/astar.rs b/azalea/src/pathfinder/astar.rs
index 2442adb6..6573c883 100644
--- a/azalea/src/pathfinder/astar.rs
+++ b/azalea/src/pathfinder/astar.rs
@@ -5,6 +5,7 @@ use std::{
     time::{Duration, Instant},
 };
 
+use num_format::ToFormattedString;
 use priority_queue::PriorityQueue;
 use rustc_hash::FxHashMap;
 use tracing::{debug, trace, warn};
@@ -131,6 +132,12 @@ where
 
     let best_path = determine_best_path(&best_paths, &start);
 
+    debug!(
+        "A* ran at {} nodes per second",
+        ((num_nodes as f64 / start_time.elapsed().as_secs_f64()) as u64)
+            .to_formatted_string(&num_format::Locale::en)
+    );
+
     Path {
         movements: reconstruct_path(nodes, best_path),
         partial: true,
diff --git a/azalea/src/pathfinder/mining.rs b/azalea/src/pathfinder/mining.rs
index 049e974a..62963306 100644
--- a/azalea/src/pathfinder/mining.rs
+++ b/azalea/src/pathfinder/mining.rs
@@ -1,7 +1,6 @@
 use std::{cell::UnsafeCell, ops::RangeInclusive};
 
-use azalea_block::{properties::Waterlogged, BlockState, BlockStates};
-use azalea_core::bitset::BitSet;
+use azalea_block::{properties::Waterlogged, BlockState, BlockStateIntegerRepr, BlockStates};
 use azalea_inventory::Menu;
 use nohash_hasher::IntMap;
 
@@ -9,11 +8,11 @@ use super::costs::BLOCK_BREAK_ADDITIONAL_PENALTY;
 use crate::auto_tool::best_tool_in_hotbar_for_block;
 
 pub struct MiningCache {
-    block_state_id_costs: UnsafeCell<IntMap<u32, f32>>,
+    block_state_id_costs: UnsafeCell<IntMap<BlockStateIntegerRepr, f32>>,
     inventory_menu: Option<Menu>,
 
-    water_block_state_range: RangeInclusive<u32>,
-    lava_block_state_range: RangeInclusive<u32>,
+    water_block_state_range: RangeInclusive<BlockStateIntegerRepr>,
+    lava_block_state_range: RangeInclusive<BlockStateIntegerRepr>,
 
     falling_blocks: Vec<BlockState>,
 }
@@ -23,16 +22,16 @@ impl MiningCache {
         let water_block_states = BlockStates::from(azalea_registry::Block::Water);
         let lava_block_states = BlockStates::from(azalea_registry::Block::Lava);
 
-        let mut water_block_state_range_min = u32::MAX;
-        let mut water_block_state_range_max = u32::MIN;
+        let mut water_block_state_range_min = BlockStateIntegerRepr::MAX;
+        let mut water_block_state_range_max = BlockStateIntegerRepr::MIN;
         for state in water_block_states {
             water_block_state_range_min = water_block_state_range_min.min(state.id);
             water_block_state_range_max = water_block_state_range_max.max(state.id);
         }
         let water_block_state_range = water_block_state_range_min..=water_block_state_range_max;
 
-        let mut lava_block_state_range_min = u32::MAX;
-        let mut lava_block_state_range_max = u32::MIN;
+        let mut lava_block_state_range_min = BlockStateIntegerRepr::MAX;
+        let mut lava_block_state_range_max = BlockStateIntegerRepr::MIN;
         for state in lava_block_states {
             lava_block_state_range_min = lava_block_state_range_min.min(state.id);
             lava_block_state_range_max = lava_block_state_range_max.max(state.id);