optimized converting from blockstate to blockkind

10 files changed, 60 insertions, 30 deletions

diff --git a/azalea-block/azalea-block-macros/src/lib.rs b/azalea-block/azalea-block-macros/src/lib.rs
index 47d12345..cecb2fb9 100644
--- a/azalea-block/azalea-block-macros/src/lib.rs
+++ b/azalea-block/azalea-block-macros/src/lib.rs
@@ -151,15 +151,19 @@ pub fn make_block_states(input: TokenStream) -> TokenStream {
     let mut block_structs = quote! {};
 
     let mut from_state_to_block_match = quote! {};
-    let mut from_registry_block_to_block_match = quote! {};
-    let mut from_registry_block_to_blockstate_match = quote! {};
-    let mut from_registry_block_to_blockstates_match = quote! {};
+    let mut from_kind_to_block_match = quote! {};
+    let mut from_kind_to_state_match = quote! {};
+    let mut from_kind_to_states_match = quote! {};
+
+    let mut from_state_to_kind_table = quote! {};
 
     // keys are enum names like Waterlogged
     let mut properties_to_state_ids = HashMap::<String, Vec<PropertyVariantData>>::new();
 
+    let last_block_kind_id = u32::try_from(input.blocks.blocks.len() - 1).unwrap();
+
     let mut state_id: BlockStateIntegerRepr = 0;
-    for block in &input.blocks.blocks {
+    for (block_kind_id, block) in input.blocks.blocks.iter().enumerate() {
         let block_property_names = &block
             .properties_and_defaults
             .iter()
@@ -423,15 +427,19 @@ pub fn make_block_states(input: TokenStream) -> TokenStream {
             }
         });
 
-        from_registry_block_to_block_match.extend(quote! {
+        from_kind_to_block_match.extend(quote! {
             BlockKind::#block_name_pascal_case => Box::new(#block_struct_name::default()),
         });
-        from_registry_block_to_blockstate_match.extend(quote! {
+        from_kind_to_state_match.extend(quote! {
             BlockKind::#block_name_pascal_case => BlockState::new_const(#default_state_id),
         });
-        from_registry_block_to_blockstates_match.extend(quote! {
+        from_kind_to_states_match.extend(quote! {
             BlockKind::#block_name_pascal_case => BlockStates::from(#first_state_id..=#last_state_id),
         });
+        for _ in first_state_id..=last_state_id {
+            let block_kind_id = block_kind_id as u32;
+            from_state_to_kind_table.extend(quote! { #block_kind_id, });
+        }
 
         let mut property_map_inner = quote! {};
         let mut get_property_match_inner = quote! {};
@@ -512,7 +520,7 @@ pub fn make_block_states(input: TokenStream) -> TokenStream {
                 fn as_block_state(&self) -> BlockState {
                     #as_block_state
                 }
-                fn as_registry_block(&self) -> BlockKind {
+                fn as_block_kind(&self) -> BlockKind {
                     BlockKind::#block_name_pascal_case
                 }
 
@@ -566,6 +574,17 @@ pub fn make_block_states(input: TokenStream) -> TokenStream {
             pub fn property<P: Property>(self) -> Option<P::Value> {
                 P::try_from_block_state(self)
             }
+
+            pub fn as_block_kind(self) -> BlockKind {
+                static TABLE: &[u32] = &[
+                    #from_state_to_kind_table
+                ];
+                const _: () = assert!(BlockKind::is_valid_id(#last_block_kind_id));
+
+                // SAFETY: the table was constructed from trusted values, and
+                // we just checked that the highest one must be valid
+                unsafe { BlockKind::from_u32(TABLE[self.id() as usize]).unwrap_unchecked() }
+            }
         }
     };
 
@@ -597,7 +616,7 @@ pub fn make_block_states(input: TokenStream) -> TokenStream {
             impl From<BlockKind> for Box<dyn BlockTrait> {
                 fn from(block: BlockKind) -> Self {
                     match block {
-                        #from_registry_block_to_block_match
+                        #from_kind_to_block_match
                         _ => unreachable!("There should always be a block struct for every BlockKind variant")
                     }
                 }
@@ -605,7 +624,7 @@ pub fn make_block_states(input: TokenStream) -> TokenStream {
             impl From<BlockKind> for BlockState {
                 fn from(block: BlockKind) -> Self {
                     match block {
-                        #from_registry_block_to_blockstate_match
+                        #from_kind_to_state_match
                         _ => unreachable!("There should always be a block state for every BlockKind variant")
                     }
                 }
@@ -613,7 +632,7 @@ pub fn make_block_states(input: TokenStream) -> TokenStream {
             impl From<BlockKind> for BlockStates {
                 fn from(block: BlockKind) -> Self {
                     match block {
-                        #from_registry_block_to_blockstates_match
+                        #from_kind_to_states_match
                         _ => unreachable!("There should always be a block state for every BlockKind variant")
                     }
                 }
diff --git a/azalea-block/src/block_state.rs b/azalea-block/src/block_state.rs
index 9a3eece8..3806c915 100644
--- a/azalea-block/src/block_state.rs
+++ b/azalea-block/src/block_state.rs
@@ -1,5 +1,6 @@
 use std::{
     fmt::{self, Debug},
+    hint::assert_unchecked,
     io::{self, Cursor, Write},
 };
 
@@ -68,6 +69,12 @@ impl BlockState {
     /// hard-code them or store them in databases.
     #[inline]
     pub const fn id(&self) -> BlockStateIntegerRepr {
+        // this unsafe assert allows us to skip bounds checks if the array has at least
+        // MAX_STATE items.
+        // SAFETY: BlockState::id is private and is always checked with is_valid_state
+        // before being constructed.
+        unsafe { assert_unchecked(Self::is_valid_state(self.id)) };
+
         self.id
     }
 }
@@ -97,13 +104,12 @@ 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 {
-            Err(())
+    fn try_from(id: u16) -> Result<Self, Self::Error> {
+        let id = id as BlockStateIntegerRepr;
+        if !Self::is_valid_state(id) {
+            return Err(());
         }
+        Ok(BlockState { id })
     }
 }
 impl From<BlockState> for u32 {
@@ -137,8 +143,8 @@ impl Debug for BlockState {
 }
 
 impl From<BlockState> for BlockKind {
-    fn from(value: BlockState) -> Self {
-        Box::<dyn BlockTrait>::from(value).as_registry_block()
+    fn from(block_state: BlockState) -> Self {
+        block_state.as_block_kind()
     }
 }
 
diff --git a/azalea-block/src/generated.rs b/azalea-block/src/generated.rs
index 85a911c1..62b580d7 100644
--- a/azalea-block/src/generated.rs
+++ b/azalea-block/src/generated.rs
@@ -3,6 +3,7 @@
 use std::{fmt::Debug, str::FromStr};
 
 use azalea_block_macros::make_block_states;
+use azalea_registry::{Registry, builtin::BlockKind};
 
 use crate::{
     BlockBehavior, BlockState, BlockStates, BlockTrait, InvalidPropertyError, Property,
diff --git a/azalea-block/src/lib.rs b/azalea-block/src/lib.rs
index 60849b83..9e2a3c5c 100644
--- a/azalea-block/src/lib.rs
+++ b/azalea-block/src/lib.rs
@@ -31,7 +31,12 @@ pub trait BlockTrait: Debug + Any {
     ///
     /// This is a lossy conversion, as [`BlockKind`] doesn't contain any state
     /// data.
-    fn as_registry_block(&self) -> BlockKind;
+    fn as_block_kind(&self) -> BlockKind;
+    #[deprecated = "renamed to as_block_kind"]
+    #[doc(hidden)]
+    fn as_registry_block(&self) -> BlockKind {
+        self.as_block_kind()
+    }
 
     /// Returns a map of property names on this block to their values as
     /// strings.
diff --git a/azalea-client/src/plugins/mining.rs b/azalea-client/src/plugins/mining.rs
index 480ba46b..bcbad747 100644
--- a/azalea-client/src/plugins/mining.rs
+++ b/azalea-client/src/plugins/mining.rs
@@ -512,7 +512,7 @@ pub fn handle_finish_mining_block_observer(
         return;
     };
 
-    let registry_block = Box::<dyn BlockTrait>::from(block_state).as_registry_block();
+    let registry_block = block_state.as_block_kind();
     if !can_use_game_master_blocks(abilities, permission_level)
         && matches!(
             registry_block,
diff --git a/azalea-entity/src/mining.rs b/azalea-entity/src/mining.rs
index bc2e6c46..19346ba9 100644
--- a/azalea-entity/src/mining.rs
+++ b/azalea-entity/src/mining.rs
@@ -35,7 +35,7 @@ pub fn get_mine_progress(
     };
 
     let base_destroy_speed = destroy_speed(
-        block.as_registry_block(),
+        block.as_block_kind(),
         held_item,
         fluid_on_eyes,
         physics,
@@ -52,7 +52,7 @@ fn has_correct_tool_for_drops(block: &dyn BlockTrait, item: &ItemStack) -> bool
     let Some(tool) = item.get_component::<Tool>() else {
         return false;
     };
-    let registry_block = block.as_registry_block();
+    let registry_block = block.as_block_kind();
     for rule in &tool.rules {
         if let Some(correct) = rule.correct_for_drops
             && rule.blocks.contains(registry_block)
diff --git a/azalea-entity/src/plugin/mod.rs b/azalea-entity/src/plugin/mod.rs
index bedb03b0..6035f674 100644
--- a/azalea-entity/src/plugin/mod.rs
+++ b/azalea-entity/src/plugin/mod.rs
@@ -144,7 +144,7 @@ pub fn update_on_climbable(
         let block_pos = BlockPos::from(position);
         let block_state_at_feet = world.get_block_state(block_pos).unwrap_or_default();
         let block_at_feet = Box::<dyn BlockTrait>::from(block_state_at_feet);
-        let registry_block_at_feet = block_at_feet.as_registry_block();
+        let registry_block_at_feet = block_at_feet.as_block_kind();
 
         **on_climbable = tags::blocks::CLIMBABLE.contains(&registry_block_at_feet)
             || (tags::blocks::TRAPDOORS.contains(&registry_block_at_feet)
@@ -167,7 +167,7 @@ fn is_trapdoor_usable_as_ladder(
 
     // block below must be a ladder
     let block_below = world.get_block_state(block_pos.down(1)).unwrap_or_default();
-    let registry_block_below = Box::<dyn BlockTrait>::from(block_below).as_registry_block();
+    let registry_block_below = block_below.as_block_kind();
     if registry_block_below != BlockKind::Ladder {
         return false;
     }
diff --git a/azalea-registry/azalea-registry-macros/src/lib.rs b/azalea-registry/azalea-registry-macros/src/lib.rs
index a8c6fe18..5cb35982 100644
--- a/azalea-registry/azalea-registry-macros/src/lib.rs
+++ b/azalea-registry/azalea-registry-macros/src/lib.rs
@@ -104,7 +104,7 @@ pub fn registry(input: TokenStream) -> TokenStream {
             }
 
             #[inline]
-            pub fn is_valid_id(id: u32) -> bool {
+            pub const fn is_valid_id(id: u32) -> bool {
                 id <= #max_id
             }
         }
diff --git a/azalea-world/src/heightmap.rs b/azalea-world/src/heightmap.rs
index bae11cea..6d99ba9d 100644
--- a/azalea-world/src/heightmap.rs
+++ b/azalea-world/src/heightmap.rs
@@ -1,4 +1,4 @@
-use azalea_block::{BlockState, BlockTrait};
+use azalea_block::BlockState;
 use azalea_core::{heightmap_kind::HeightmapKind as HeightmapKind_, math, position::ChunkBlockPos};
 use azalea_registry::tags::blocks::LEAVES;
 use tracing::warn;
@@ -31,8 +31,7 @@ fn motion_blocking(block_state: BlockState) -> bool {
 }
 
 pub fn is_heightmap_opaque(heightmap: HeightmapKind_, block_state: BlockState) -> bool {
-    let block = Box::<dyn BlockTrait>::from(block_state);
-    let registry_block = block.as_registry_block();
+    let registry_block = block_state.as_block_kind();
     type K = HeightmapKind_;
     match heightmap {
         K::WorldSurfaceWg => !block_state.is_air(),
diff --git a/azalea/src/auto_tool.rs b/azalea/src/auto_tool.rs
index 5fe0f529..4f59d4a9 100644
--- a/azalea/src/auto_tool.rs
+++ b/azalea/src/auto_tool.rs
@@ -81,7 +81,7 @@ pub fn accurate_best_tool_in_hotbar_for_block(
     let mut best_slot = None;
 
     let block = Box::<dyn BlockTrait>::from(block);
-    let registry_block = block.as_registry_block();
+    let registry_block = block.as_block_kind();
 
     if matches!(registry_block, BlockKind::Water | BlockKind::Lava) {
         // can't mine fluids