optimized converting from blockstate to blockkind

4 files changed, 51 insertions, 20 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.