Refactor azalea-registry (#294)

* move registries in azalea-registry into separate modules

* rename Item and Block to ItemKind and BlockKind

* remove 'extra' registries from azalea-registry

* hide deprecated items from docs

* use DamageKindKey instead of Identifier when parsing registries

* store tag entries as a Vec instead of a HashSet

* sort tag values by protocol id

* update changelog

5 files changed, 143 insertions, 49 deletions

diff --git a/codegen/lib/code/blocks.py b/codegen/lib/code/blocks.py
index 9f783690..84d9fffb 100644
--- a/codegen/lib/code/blocks.py
+++ b/codegen/lib/code/blocks.py
@@ -6,7 +6,7 @@ BLOCKS_RS_DIR = get_dir_location("../azalea-block/src/generated.rs")
 # - Property: A property of a block, like "direction"
 # - Variant: A potential state of a property, like "up"
 # - State: A possible state of a block, a combination of variants
-# - Block: Has properties and states.
+# - BlockKind: Has properties and states.
 
 
 def generate_blocks(
@@ -85,7 +85,7 @@ def generate_blocks(
 
     new_make_block_states_macro_code.append("    },")
 
-    # Block codegen
+    # BlockKind codegen
     new_make_block_states_macro_code.append("    Blocks => {")
     for block_id in ordered_blocks:
         block_data_report = blocks_report["minecraft:" + block_id]
diff --git a/codegen/lib/code/data_components.py b/codegen/lib/code/data_components.py
index d4c13118..38060662 100644
--- a/codegen/lib/code/data_components.py
+++ b/codegen/lib/code/data_components.py
@@ -47,7 +47,7 @@ def generate(version_id: str):
 
 def get_expected_variants(version_id: str):
     expected_variants = []
-    registries = lib.extract.get_registries_report(version_id)
+    registries = lib.extract.get_builtin_registries_report(version_id)
 
     registry = registries["minecraft:data_component_type"]
     registry_entries = sorted(
@@ -182,7 +182,7 @@ use std::collections::HashMap;
 use azalea_chat::translatable_component::TranslatableComponent;
 use azalea_core::attribute_modifier_operation::AttributeModifierOperation;
 use azalea_registry::{
-    Attribute, Block, DataRegistry, EntityKind, HolderSet, Item, MobEffect, SoundEvent,
+    Attribute, BlockKind, DataRegistry, EntityKind, HolderSet, ItemKind, MobEffect, SoundEvent,
 };
 use simdnbt::owned::NbtCompound;
 
@@ -207,7 +207,7 @@ use crate::{
                 component_value
             )
 
-    registries = lib.extract.get_registries_report(version_id)
+    registries = lib.extract.get_builtin_registries_report(version_id)
     item_resource_id_to_protocol_id = {}
     item_resource_ids = [None] * len(registries["minecraft:item"]["entries"])
     for item_resource_id, item_data in registries["minecraft:item"]["entries"].items():
@@ -358,7 +358,7 @@ use crate::{
                     list(python_value.values())[0], target_rust_type
                 )
             elif target_rust_type == "ItemStack":
-                item_rust_value = python_to_rust_value(python_value["id"], "Item")
+                item_rust_value = python_to_rust_value(python_value["id"], "ItemKind")
                 count = python_value["count"]
                 if count == 1:
                     return f"ItemStack::from({item_rust_value})"
@@ -430,7 +430,7 @@ use crate::{
             elif target_rust_type == "DamageType":
                 # TODO: this is intentionally incorrect, see the comment in
                 # azalea-registry/src/data.rs to see how to fix this properly
-                return "DamageType::Registry(azalea_registry::DamageKind::new_raw(0))"
+                return "DamageType::Registry(azalea_registry::data::DamageKind::new_raw(0))"
             else:
                 # enum variant
                 return f"{target_rust_type}::{lib.utils.to_camel_case(python_value.split(':')[-1])}"
@@ -464,9 +464,9 @@ use crate::{
                     tag_name = lib.utils.to_snake_case(v.split(":")[-1]).upper()
                     if inner_type == "EntityKind":
                         tag_module = "entities"
-                    elif inner_type == "Item":
+                    elif inner_type == "ItemKind":
                         tag_module = "items"
-                    elif inner_type == "Block":
+                    elif inner_type == "BlockKind":
                         tag_module = "blocks"
                     else:
                         tag_module = "FIXME_UNKNOWN_MODULE"
@@ -561,7 +561,9 @@ use crate::{
             if len(values_set) == 1:
                 # always returns the same value
                 code.append(f"impl DefaultableComponent for {component_struct_name} {{")
-                code.append("    fn default_for_item(_item: Item) -> Option<Self> {")
+                code.append(
+                    "    fn default_for_item(_item: ItemKind) -> Option<Self> {"
+                )
                 value = next(iter(values_set))
                 code.append(f"        Some({transform_value_fn(value)})")
                 code.append("    }")
@@ -587,7 +589,7 @@ use crate::{
             code.append(static_def_line)
 
             code.append(f"impl DefaultableComponent for {component_struct_name} {{")
-            code.append("    fn default_for_item(item: Item) -> Option<Self> {")
+            code.append("    fn default_for_item(item: ItemKind) -> Option<Self> {")
             code.append(f"        let value = {static_values_name}[item as usize];")
             if none_value_is_used:
                 code.append(f"        if value == {none_value} {{")
@@ -599,18 +601,22 @@ use crate::{
         elif includes_every_item_but_mostly_same_values:
             code.append(f"impl DefaultableComponent for {component_struct_name} {{")
             if default_values_count_except_most_common > 0:
-                code.append("    fn default_for_item(item: Item) -> Option<Self> {")
+                code.append("    fn default_for_item(item: ItemKind) -> Option<Self> {")
                 code.append("        let value = match item {")
                 for item_resource_id, value in item_defaults.items():
                     if value == most_common_default_value:
                         continue
                     item_variant_name = lib.utils.to_camel_case(item_resource_id)
-                    code.append(f"            Item::{item_variant_name} => {value},")
+                    code.append(
+                        f"            ItemKind::{item_variant_name} => {value},"
+                    )
                 code.append(f"            _ => {most_common_default_value},")
                 code.append("        };")
                 code.append(f"        Some({transform_value_fn('value')})")
             else:
-                code.append("    fn default_for_item(_item: Item) -> Option<Self> {")
+                code.append(
+                    "    fn default_for_item(_item: ItemKind) -> Option<Self> {"
+                )
                 code.append(
                     f"        Some({transform_value_fn(most_common_default_value)})"
                 )
@@ -618,11 +624,11 @@ use crate::{
             code.append("}")
         else:
             code.append(f"impl DefaultableComponent for {component_struct_name} {{")
-            code.append("    fn default_for_item(item: Item) -> Option<Self> {")
+            code.append("    fn default_for_item(item: ItemKind) -> Option<Self> {")
             code.append("        let value = match item {")
             for item_resource_id, value in item_defaults.items():
                 item_variant_name = lib.utils.to_camel_case(item_resource_id)
-                code.append(f"            Item::{item_variant_name} => {value},")
+                code.append(f"            ItemKind::{item_variant_name} => {value},")
             code.append("            _ => return None,")
             code.append("        };")
             code.append(f"        Some({transform_value_fn('value')})")
diff --git a/codegen/lib/code/entity.py b/codegen/lib/code/entity.py
index e9da0404..2fc66b0f 100644
--- a/codegen/lib/code/entity.py
+++ b/codegen/lib/code/entity.py
@@ -114,7 +114,7 @@ use azalea_core::{
     position::{BlockPos, Vec3f32},
 };
 use azalea_inventory::{ItemStack, components};
-use azalea_registry::DataRegistry;
+use azalea_registry::{DataRegistry, builtin::EntityKind};
 use bevy_ecs::{bundle::Bundle, component::Component};
 use derive_more::{Deref, DerefMut};
 use thiserror::Error;
@@ -456,15 +456,15 @@ impl From<EntityDataValue> for UpdateMetadataError {
                             default = "simdnbt::owned::NbtCompound::default()"
                         # elif type_name == 'CatVariant':
                         #     # TODO: the default should be Tabby but we don't have a way to get that from here
-                        #     default = 'azalea_registry::CatVariant::new_raw(0)'
+                        #     default = 'azalea_registry::data::CatVariant::new_raw(0)'
                         # elif type_name == 'PaintingVariant':
-                        #     default = 'azalea_registry::PaintingVariant::Kebab'
+                        #     default = 'azalea_registry::data::PaintingVariant::Kebab'
                         # elif type_name == 'FrogVariant':
-                        #     default = 'azalea_registry::FrogVariant::Temperate'
+                        #     default = 'azalea_registry::data::FrogVariant::Temperate'
                         elif type_name.endswith("Variant"):
-                            default = f"azalea_registry::{type_name}::new_raw(0)"
+                            default = f"azalea_registry::data::{type_name}::new_raw(0)"
                         elif type_name == "VillagerData":
-                            default = "VillagerData { kind: azalea_registry::VillagerKind::Plains, profession: azalea_registry::VillagerProfession::None, level: 0 }"
+                            default = "VillagerData { kind: azalea_registry::builtin::VillagerKind::Plains, profession: azalea_registry::builtin::VillagerProfession::None, level: 0 }"
                         else:
                             default = (
                                 f"{type_name}::default()"
@@ -581,7 +581,7 @@ impl From<EntityDataValue> for UpdateMetadataError {
     code.append(
         """pub fn apply_metadata(
     entity: &mut bevy_ecs::system::EntityCommands,
-    entity_kind: azalea_registry::EntityKind,
+    entity_kind: EntityKind,
     items: Vec<EntityDataItem>,
 ) -> Result<(), UpdateMetadataError> {
     match entity_kind {"""
@@ -591,7 +591,7 @@ impl From<EntityDataValue> for UpdateMetadataError {
             # not actually an entity
             continue
         struct_name: str = upper_first_letter(to_camel_case(entity_id))
-        code.append(f"        azalea_registry::EntityKind::{struct_name} => {{")
+        code.append(f"        EntityKind::{struct_name} => {{")
         code.append("            for d in items {")
         code.append(f"                {struct_name}::apply_metadata(entity, d)?;")
         code.append("            }")
@@ -601,15 +601,15 @@ impl From<EntityDataValue> for UpdateMetadataError {
     code.append("}")
     code.append("")
 
-    # pub fn apply_default_metadata(entity: &mut bevy_ecs::system::EntityCommands, kind: azalea_registry::EntityKind) {
+    # pub fn apply_default_metadata(entity: &mut bevy_ecs::system::EntityCommands, kind: EntityKind) {
     #     match kind {
-    #         azalea_registry::EntityKind::AreaEffectCloud => {
+    #         EntityKind::AreaEffectCloud => {
     #             entity.insert(AreaEffectCloudMetadataBundle::default());
     #         }
     #     }
     # }
     code.append(
-        "pub fn apply_default_metadata(entity: &mut bevy_ecs::system::EntityCommands, kind: azalea_registry::EntityKind) {"
+        "pub fn apply_default_metadata(entity: &mut bevy_ecs::system::EntityCommands, kind: EntityKind) {"
     )
     code.append("    match kind {")
     for entity_id in burger_entity_metadata:
@@ -617,7 +617,7 @@ impl From<EntityDataValue> for UpdateMetadataError {
             # not actually an entity
             continue
         struct_name: str = upper_first_letter(to_camel_case(entity_id))
-        code.append(f"        azalea_registry::EntityKind::{struct_name} => {{")
+        code.append(f"        EntityKind::{struct_name} => {{")
         code.append(
             f"            entity.insert({struct_name}MetadataBundle::default());"
         )
diff --git a/codegen/lib/code/registry.py b/codegen/lib/code/registry.py
index 84d613ba..a01dc845 100644
--- a/codegen/lib/code/registry.py
+++ b/codegen/lib/code/registry.py
@@ -1,16 +1,17 @@
 from lib.utils import get_dir_location, to_camel_case
 
-REGISTRIES_DIR = get_dir_location("../azalea-registry/src/lib.rs")
+BUILTIN_REGISTRIES_DIR = get_dir_location("../azalea-registry/src/builtin.rs")
+DATA_REGISTRIES_DIR = get_dir_location("../azalea-registry/src/data.rs")
 
 
-def generate_registries(registries: dict):
-    with open(REGISTRIES_DIR, "r") as f:
+def generate_builtin_registries(registries: dict):
+    with open(BUILTIN_REGISTRIES_DIR, "r") as f:
         code = f.read().split("\n")
 
     existing_registry_enum_names = set()
 
     for registry_name, registry in registries.items():
-        # registry!(Block, {
+        # registry!(BlockKind, {
         #     Air => "minecraft:air",
         #     Stone => "minecraft:stone"
         # });
@@ -68,18 +69,91 @@ def generate_registries(registries: dict):
         else:
             i += 1
 
-    with open(REGISTRIES_DIR, "w") as f:
+    with open(BUILTIN_REGISTRIES_DIR, "w") as f:
+        f.write("\n".join(code))
+
+
+# data_registries looks like { "enchantment": [ "aqua_affinity", ... ] }
+def generate_data_registries(data_registries: dict):
+    with open(DATA_REGISTRIES_DIR, "r") as f:
+        code = f.read().split("\n")
+
+    existing_registry_struct_names = set()
+    for registry_name, registry_entries in data_registries.items():
+        registry_enum_name = registry_name_to_enum_name(registry_name.split("/")[-1])
+        existing_registry_struct_names.add(registry_enum_name)
+
+    # delete the unused data registries
+    i = 0
+    while i < len(code):
+        if code[i] == "data_registry! {":
+            i += 1
+            struct_name = code[i].split(" ")[0]
+            if struct_name not in existing_registry_struct_names:
+                print("removing data registry", struct_name)
+                i -= 1
+                while code[i] != "}":
+                    code.pop(i)
+                code.pop(i)
+                # close the data_registry! block
+                code.pop(i)
+        else:
+            i += 1
+
+    for registry_name, registry_entries in data_registries.items():
+        # data_registry! {
+        #     Enchantment => "enchantment",
+        #     enum EnchantmentKey {
+        #         AquaAffinity => "minecraft:aqua_affinity",
+        #     }
+        # }
+
+        registry_enum_name = registry_name_to_enum_name(registry_name.split("/")[-1])
+
+        registry_code = []
+        registry_code.append(f'{registry_enum_name} => "{registry_name}",')
+        registry_code.append(f"enum {registry_enum_name}Key {{")
+        registry_entries.sort()
+        for variant_name in registry_entries:
+            variant_struct_name = to_camel_case(variant_name.split(":")[-1])
+            registry_code.append(f'    {variant_struct_name} => "{variant_name}",')
+        registry_code.append("}")
+
+        # when we find a "data_registry! {" line, find the next line that starts
+        # with "enum <name>" and replace that until we find a line that's "}"
+        found = False
+        in_registry_macro = False
+        for i, line in enumerate(list(code)):
+            if not in_registry_macro and line == "data_registry! {":
+                in_registry_macro = True
+            elif in_registry_macro and line == registry_code[1]:
+                # found it, now delete until we get to "}"
+                while code[i] != "}":
+                    code.pop(i)
+                code[i] = "\n".join(registry_code[1:])
+                found = True
+                break
+        if not found:
+            code.append("data_registry! {")
+            code.append("\n".join(registry_code))
+            code.append("}")
+            code.append("")
+
+    with open(DATA_REGISTRIES_DIR, "w") as f:
         f.write("\n".join(code))
 
 
 def registry_name_to_enum_name(registry_name: str) -> str:
     registry_name = registry_name.split(":")[-1]
 
-    if registry_name.endswith("_type"):
+    if registry_name == "block_type":
+        # avoid conflicting with BlockKind
+        registry_name = "abstract_block_kind"
+    elif registry_name.endswith("_type"):
         # change _type to _kind because that's Rustier (and because _type
         # is a reserved keyword)
         registry_name = registry_name[:-5] + "_kind"
-    elif registry_name in {"menu"}:
+    elif registry_name in {"menu", "block", "item"}:
         registry_name += "_kind"
 
     return to_camel_case(registry_name)
diff --git a/codegen/lib/code/tags.py b/codegen/lib/code/tags.py
index d3e561df..ff5a4061 100644
--- a/codegen/lib/code/tags.py
+++ b/codegen/lib/code/tags.py
@@ -1,34 +1,48 @@
+from lib.code.registry import registry_name_to_enum_name
 from lib.utils import to_snake_case, upper_first_letter, get_dir_location, to_camel_case
 
-REGISTRIES_DIR = get_dir_location("../azalea-registry/src/tags")
+TAGS_DIR = get_dir_location("../azalea-registry/src/tags")
 
 
-def generate_tags(registries: dict, file_name: str, struct_name: str):
-    tags_dir = f"{REGISTRIES_DIR}/{file_name}.rs"
+def generate_tags(registries: dict, tags: dict, file_name: str, registry_name: str):
+    struct_name = registry_name_to_enum_name(registry_name)
 
-    generated = f"""// This file was @generated by codegen/lib/code/tags.py, don't edit it manually!
+    tags_dir = f"{TAGS_DIR}/{file_name}.rs"
 
-use std::{{collections::HashSet, sync::LazyLock}};
+    generated = f"""// This file was @generated by codegen/lib/code/tags.py, don't edit it manually!
 
+use std::sync::LazyLock;
 
-use crate::{struct_name};
+use crate::{{builtin::{struct_name}, tags::RegistryTag}};
 
 """
 
-    for tag_name, tag in sorted(registries.items(), key=lambda x: x[0]):
-        tag_name = tag_name.replace("/", "_")
-        static_set_name = to_snake_case(tag_name).upper()
-        generated += f"pub static {static_set_name}: LazyLock<HashSet<{struct_name}>> = LazyLock::new(|| HashSet::from_iter(["
+    protocol_ids = {}
+    for k, v in registries["minecraft:" + registry_name]["entries"].items():
+        protocol_ids[k.split(":")[1]] = v["protocol_id"]
 
+    for tag_name, tag in sorted(tags.items(), key=lambda x: x[0]):
+        entries = []
         queue = tag["values"].copy()
         while queue != []:
-            item = queue.pop(0)
-            namespace, item_name = item.split(":")
+            ident = queue.pop(0)
+            namespace, entry_name = ident.split(":")
             if namespace[0] == "#":
-                queue += registries[item_name]["values"]
+                queue += tags[entry_name]["values"]
                 continue
+            entries.append(entry_name)
+
+        tag_name = tag_name.replace("/", "_")
+        static_set_name = to_snake_case(tag_name).upper()
+        generated += f"pub static {static_set_name}: LazyLock<RegistryTag<{struct_name}>> = LazyLock::new(|| RegistryTag::new(vec!["
+
+        # this is important because we binary search registries in some cases
+        # and they need to be sorted by their rust Ord order
+        entries.sort(key=lambda e: protocol_ids[e])
+
+        for entry_name in entries:
             generated += (
-                f"{struct_name}::{upper_first_letter(to_camel_case(item_name))},\n"
+                f"{struct_name}::{upper_first_letter(to_camel_case(entry_name))},\n"
             )
         generated += "]));\n"