Refactor azalea-protocol (#190)

* start updating to 1.21.4

* fix block codegen and stop using block data from burger

* rename packet related modules and structs to be simpler

* ItemSlot -> ItemStack for more consistency with mojmap

* .get() -> .into_packet()

* simplify declare_state_packets by removing packet ids

* rename read_from and write_into to azalea_read and azalea_write

* rename McBufReadable and McBufWritable to AzaleaRead and AzaleaWrite

* McBuf -> AzBuf

* remove most uses of into_variant

* update codegen and use resourcelocation names for packets

* implement #[limit(i)] attribute for AzBuf derive macro

* fixes for 1.21.4

* fix examples

* update some physics code and fix ChatType

* remove unused imports in codegen

* re-add some things to migrate.py and update +mc version numbers automatically

* downgrade to 1.21.3 lol

7 files changed, 176 insertions, 271 deletions

diff --git a/codegen/lib/code/blocks.py b/codegen/lib/code/blocks.py
index cda95339..2733093b 100755
--- a/codegen/lib/code/blocks.py
+++ b/codegen/lib/code/blocks.py
@@ -12,7 +12,7 @@ BLOCKS_RS_DIR = get_dir_location('../azalea-block/src/generated.rs')
 # - Block: Has properties and states.
 
 
-def generate_blocks(blocks_burger: dict, blocks_report: dict, pixlyzer_block_datas: dict, ordered_blocks: list[str], mappings: Mappings):
+def generate_blocks(blocks_report: dict, pixlyzer_block_datas: dict, ordered_blocks: list[str]):
     with open(BLOCKS_RS_DIR, 'r') as f:
         existing_code = f.read().splitlines()
 
@@ -25,25 +25,14 @@ def generate_blocks(blocks_burger: dict, blocks_report: dict, pixlyzer_block_dat
     # This dict looks like { 'FloweringAzaleaLeavesDistance': 'distance' }
     property_struct_names_to_names = {}
     for block_id in ordered_blocks:
-        block_data_burger = blocks_burger[block_id]
         block_data_report = blocks_report[f'minecraft:{block_id}']
 
         block_properties = {}
-        for property_name in list(block_data_report.get('properties', {}).keys()):
-            property_burger = None
-            for property in block_data_burger.get('states', []):
-                if property['name'] == property_name:
-                    property_burger = property
-                    break
-
-            property_variants = block_data_report['properties'][property_name]
-
-            if property_burger is None:
-                print(
-                    f'Warning: The reports have states for a block, but Burger doesn\'t! (missing "{property_name}")', block_data_burger)
+        for property_id in list(block_data_report.get('properties', {}).keys()):
+            property_variants = block_data_report['properties'][property_id]
 
             property_struct_name = get_property_struct_name(
-                property_burger, block_data_burger, property_variants, mappings)
+                block_id, property_id, property_variants)
 
             if property_struct_name in properties:
                 if not properties[property_struct_name] == property_variants:
@@ -55,7 +44,7 @@ def generate_blocks(blocks_burger: dict, blocks_report: dict, pixlyzer_block_dat
 
             block_properties[property_struct_name] = property_variants
 
-            property_struct_names_to_names[property_struct_name] = property_name
+            property_struct_names_to_names[property_struct_name] = property_id
 
         properties.update(block_properties)
 
@@ -67,7 +56,7 @@ def generate_blocks(blocks_burger: dict, blocks_report: dict, pixlyzer_block_dat
         #     Wall,
         #     Ceiling,
         # },
-        property_name = property_struct_names_to_names[property_struct_name]
+        property_id = property_struct_names_to_names[property_struct_name]
 
         # if the only variants are true and false, we make it unit struct with a boolean instead of an enum
         if property_variants == ['true', 'false']:
@@ -79,38 +68,28 @@ def generate_blocks(blocks_burger: dict, blocks_report: dict, pixlyzer_block_dat
             property_shape_code += '        }'
 
         new_make_block_states_macro_code.append(
-            f'        "{property_name}" => {property_shape_code},')
+            f'        "{property_id}" => {property_shape_code},')
 
     new_make_block_states_macro_code.append('    },')
 
     # Block codegen
     new_make_block_states_macro_code.append('    Blocks => {')
     for block_id in ordered_blocks:
-        block_data_burger = blocks_burger[block_id]
         block_data_report = blocks_report['minecraft:' + block_id]
         block_data_pixlyzer = pixlyzer_block_datas.get(f'minecraft:{block_id}', {})
 
-        block_properties = block_data_burger.get('states', [])
-        block_properties_burger = block_data_burger.get('states', [])
-
         default_property_variants: dict[str, str] = {}
         for state in block_data_report['states']:
             if state.get('default'):
                 default_property_variants = state.get('properties', {})
 
         properties_code = '{'
-        for property_name in list(block_data_report.get('properties', {}).keys()):
-            property_burger = None
-            for property in block_data_burger.get('states', []):
-                if property['name'] == property_name:
-                    property_burger = property
-                    break
-
-            property_default = default_property_variants.get(property_name)
-            property_variants = block_data_report['properties'][property_name]
+        for property_id in list(block_data_report.get('properties', {}).keys()):
+            property_default = default_property_variants.get(property_id)
+            property_variants = block_data_report['properties'][property_id]
 
             property_struct_name = get_property_struct_name(
-                property_burger, block_data_burger, property_variants, mappings)
+                block_id, property_id, property_variants)
 
             is_boolean_property = property_variants == ['true', 'false']
 
@@ -123,7 +102,7 @@ def generate_blocks(blocks_burger: dict, blocks_report: dict, pixlyzer_block_dat
 
             assert property_default is not None
 
-            this_property_code = f'"{property_name}": {property_default_type}'
+            this_property_code = f'"{property_id}": {property_default_type}'
 
             properties_code += f'\n            {this_property_code},'
         # if there's nothing inside the properties, keep it in one line
@@ -177,8 +156,7 @@ def generate_blocks(blocks_burger: dict, blocks_report: dict, pixlyzer_block_dat
     with open(BLOCKS_RS_DIR, 'w') as f:
         f.write('\n'.join(new_code))
 
-
-def get_property_struct_name(property: Optional[dict], block_data_burger: dict, property_variants: list[str], mappings: Mappings) -> str:
+def get_property_struct_name(block_id: str, property_id: str, property_variants: list[str]) -> str:
     # these are hardcoded because otherwise they cause conflicts
     # some names inspired by https://github.com/feather-rs/feather/blob/main/feather/blocks/src/generated/table.rs
     if property_variants == ['north', 'east', 'south', 'west', 'up', 'down']:
@@ -205,34 +183,34 @@ def get_property_struct_name(property: Optional[dict], block_data_burger: dict,
         return 'VaultState'
     if 'harp' in property_variants and 'didgeridoo' in property_variants:
         return 'Sound'
-
-    if property is None:
-        return ''.join(map(to_camel_case, property_variants))
-
-    if property_variants == ['true', 'false']:
-        # booleans are weird, so just return the string name minecraft uses
-        return to_camel_case(property['name'])
-
-    for class_name in [block_data_burger['class']] + block_data_burger['super']:
-        property_name = mappings.get_field(
-            class_name, property['field_name'])
-        if property_name:
-            break
-    if property_name is None:
-        if 'declared_in' in property:
-            property_name = mappings.get_field(
-                property['declared_in'], property['field_name'])
-    if property_name is None:
-        property_name = property['name']
-    assert property_name
-    property_name = to_camel_case(property_name.lower())
-    if property['type'] == 'int':
-        property_name = to_camel_case(
-            block_data_burger['text_id']) + property_name
-
-    # if property_variants == ['none', 'low', 'tall']:
-
+    if is_list_of_string_integers(property_variants):
+        # if the values are all integers, then prepend the block name
+        return to_camel_case(block_id) + to_camel_case(property_id)
     if property_variants == ['up', 'side', 'none']:
-        property_name = 'Wire' + to_camel_case(property_name)
-
-    return property_name
+        return 'Wire' + to_camel_case(property_id)
+    if property_variants == ['none', 'low', 'tall']:
+        return 'Wall' + to_camel_case(property_id)
+
+    return to_camel_case(property_id)
+
+def is_list_of_string_integers(l: list[str]) -> bool:
+    return all(map(str.isdigit, l))
+
+def get_ordered_blocks(registries_report: dict[str, dict]) -> list[str]:
+    '''
+    Returns a list of block ids (like ['air', 'stone', ...]) ordered by their protocol id.
+    '''
+    blocks_registry = registries_report['minecraft:block']
+
+    blocks_to_ids = {} 
+    for block_id, value in blocks_registry['entries'].items():
+        prefix = 'minecraft:'
+        assert block_id.startswith(prefix)
+        block_id = block_id[len(prefix):]
+        protocol_id = value['protocol_id']
+        blocks_to_ids[block_id] = protocol_id
+    
+    ordered_blocks = []
+    for block_id in sorted(blocks_to_ids, key=blocks_to_ids.get):
+        ordered_blocks.append(block_id)
+    return ordered_blocks
diff --git a/codegen/lib/code/entity.py b/codegen/lib/code/entity.py
index 5f0bc3d9..a380d3d9 100644
--- a/codegen/lib/code/entity.py
+++ b/codegen/lib/code/entity.py
@@ -111,7 +111,7 @@ use azalea_core::{
     direction::Direction,
     position::{BlockPos, Vec3},
 };
-use azalea_inventory::ItemSlot;
+use azalea_inventory::ItemStack;
 use bevy_ecs::{bundle::Bundle, component::Component};
 use derive_more::{Deref, DerefMut};
 use thiserror::Error;
@@ -428,7 +428,7 @@ impl From<EntityDataValue> for UpdateMetadataError {
                         elif type_name == 'OptionalUnsignedInt':
                             default = f'OptionalUnsignedInt(Some({default}))' if default != 'Empty' else 'OptionalUnsignedInt(None)'
                         elif type_name == 'ItemStack':
-                            default = f'ItemSlot::Present({default})' if default != 'Empty' else 'ItemSlot::Empty'
+                            default = f'ItemStack::Present({default})' if default != 'Empty' else 'ItemStack::Empty'
                         elif type_name == 'BlockState':
                             default = f'{default}' if default != 'Empty' else 'azalea_block::BlockState::AIR'
                         elif type_name == 'OptionalBlockState':
diff --git a/codegen/lib/code/packet.py b/codegen/lib/code/packet.py
index b34d3455..71bc3a70 100755
--- a/codegen/lib/code/packet.py
+++ b/codegen/lib/code/packet.py
@@ -9,179 +9,80 @@ import re
 def make_packet_mod_rs_line(packet_id: int, packet_class_name: str):
     return f'        {padded_hex(packet_id)}: {to_snake_case(packet_class_name)}::{to_camel_case(packet_class_name)},'
 
-
-def fix_state(state: str):
-    return {'PLAY': 'game'}.get(state, state.lower())
-
-
-def generate_packet(burger_packets, mappings: Mappings, target_packet_id, target_packet_direction, target_packet_state):
-    for packet in burger_packets.values():
-        if packet['id'] != target_packet_id:
-            continue
-
-        direction = packet['direction'].lower()  # serverbound or clientbound
-        state = fix_state(packet['state'])
-
-        if state != target_packet_state or direction != target_packet_direction:
-            continue
-
-        generated_packet_code = []
-        uses = set()
-        extra_code = []
-
-        packet_derive_name = f'{to_camel_case(direction)}{to_camel_case(state)}Packet'
-
-        generated_packet_code.append(
-            f'#[derive(Clone, Debug, McBuf, {packet_derive_name})]')
-        uses.add(f'azalea_protocol_macros::{packet_derive_name}')
-        uses.add(f'azalea_buf::McBuf')
-
-        obfuscated_class_name = packet['class'].split('.')[0]
-        class_name = mappings.get_class(
-            obfuscated_class_name).split('.')[-1]
-        if '$' in class_name:
-            class_name, extra_part = class_name.split('$')
-            if class_name.endswith('Packet'):
-                class_name = class_name[:-
-                                        len('Packet')] + extra_part + 'Packet'
-
-        generated_packet_code.append(
-            f'pub struct {to_camel_case(class_name)} {{')
-
-        # call burger_instruction_to_code for each instruction
-        i = -1
-        instructions = packet.get('instructions', [])
-        while (i + 1) < len(instructions):
-            i += 1
-
-            if instructions[i]['operation'] == 'write':
-                skip = burger_instruction_to_code(
-                    instructions, i, generated_packet_code, mappings, obfuscated_class_name, uses, extra_code)
-                if skip:
-                    i += skip
-            else:
-                generated_packet_code.append(f'// TODO: {instructions[i]}')
-
-        generated_packet_code.append('}')
-
-        if uses:
-            # empty line before the `use` statements
-            generated_packet_code.insert(0, '')
-        for use in uses:
-            generated_packet_code.insert(0, f'use {use};')
-        for line in extra_code:
-            generated_packet_code.append(line)
-
-        print(generated_packet_code)
-        write_packet_file(state, to_snake_case(class_name),
-                          '\n'.join(generated_packet_code))
-        print()
-
+MOJMAP_TO_AZALEA_STATE_NAME_MAPPING = {
+    # shorter name, i like it more
+    'configuration': 'config',
+    # in the files mojang calls the directory "game" so we do that too
+    'play': 'game'
+}
+AZALEA_TO_MOJMAP_STATE_NAME_MAPPING = {v: k for k, v in MOJMAP_TO_AZALEA_STATE_NAME_MAPPING.items()}
+
+def generate_packet(packets_report, packet_name, direction, state):
+    mojmap_state = AZALEA_TO_MOJMAP_STATE_NAME_MAPPING.get(state, state)
+    _packet_report = packets_report[mojmap_state][direction]['minecraft:' + packet_name]
+
+    code = []
+    uses = set()
+
+    packet_derive_name = f'{to_camel_case(direction)}{to_camel_case(state)}Packet'
+
+    packet_struct_name = to_camel_case(f'{direction}_{packet_name}')
+    packet_module_name = f'{direction[0]}_{packet_name}'
+
+    code.append(f'use azalea_buf::AzBuf;')
+    code.append(f'use azalea_protocol_macros::{packet_derive_name};')
+    code.append('')
+   
+    code.append(
+        f'#[derive(Clone, Debug, AzBuf, {packet_derive_name})]')
+    code.append(
+        f'pub struct {packet_struct_name} {{')
+    code.append('    TODO')
+    code.append('}')
+
+    print(code)
+    write_packet_file(state, packet_module_name, '\n'.join(code))
+
+    # this won't handle writing to the packets/{state}/mod.rs file since we'd need to know the full packet list
+
+def set_packets(packets_report):
+    for mojmap_state in packets_report:
+        state = MOJMAP_TO_AZALEA_STATE_NAME_MAPPING.get(mojmap_state, mojmap_state)
         mod_rs_dir = get_dir_location(
             f'../azalea-protocol/src/packets/{state}/mod.rs')
-        with open(mod_rs_dir, 'r') as f:
-            mod_rs = f.read().splitlines()
-
-        pub_mod_line = f'pub mod {to_snake_case(class_name)};'
-        if pub_mod_line not in mod_rs:
-            mod_rs.insert(0, pub_mod_line)
-            packet_mod_rs_line = make_packet_mod_rs_line(
-                packet['id'], class_name)
-
-            in_serverbound = False
-            in_clientbound = False
-            for i, line in enumerate(mod_rs):
-                if line.strip() == 'Serverbound => {':
-                    in_serverbound = True
-                    continue
-                elif line.strip() == 'Clientbound => {':
-                    in_clientbound = True
-                    continue
-                elif line.strip() in ('}', '},'):
-                    if (in_serverbound and direction == 'serverbound') or (in_clientbound and direction == 'clientbound'):
-                        mod_rs.insert(i, packet_mod_rs_line)
-                        break
-                    in_serverbound = in_clientbound = False
-                    continue
-
-                if line.strip() == '' or line.strip().startswith('//') or (not in_serverbound and direction == 'serverbound') or (not in_clientbound and direction == 'clientbound'):
-                    continue
-
-                line_packet_id_hex = line.strip().split(':')[0]
-                assert line_packet_id_hex.startswith('0x')
-                line_packet_id = int(line_packet_id_hex[2:], 16)
-                if line_packet_id > packet['id']:
-                    mod_rs.insert(i, packet_mod_rs_line)
-                    break
-
-            with open(mod_rs_dir, 'w') as f:
-                f.write('\n'.join(mod_rs))
-
-
-def set_packets(packet_ids: list[int], packet_class_names: list[str], direction: str, state: str):
-    assert len(packet_ids) == len(packet_class_names)
-
-    # ids are repeated
-    assert len(packet_ids) == len(set(packet_ids))
-
-    # sort the packets by id
-    packet_ids, packet_class_names = [list(x) for x in zip(
-        *sorted(zip(packet_ids, packet_class_names), key=lambda pair: pair[0]))]  # type: ignore
-
-    mod_rs_dir = get_dir_location(
-        f'../azalea-protocol/src/packets/{state}/mod.rs')
-    with open(mod_rs_dir, 'r') as f:
-        mod_rs = f.read().splitlines()
-    new_mod_rs = []
-
-    required_modules = []
-
-    ignore_lines = False
-
-    for line in mod_rs:
-        if line.strip() == 'Serverbound => {':
-            new_mod_rs.append(line)
-            if direction == 'serverbound':
-                ignore_lines = True
-                for packet_id, packet_class_name in zip(packet_ids, packet_class_names):
-                    new_mod_rs.append(
-                        make_packet_mod_rs_line(packet_id, packet_class_name)
-                    )
-                    required_modules.append(packet_class_name)
-            else:
-                ignore_lines = False
-            continue
-        elif line.strip() == 'Clientbound => {':
-            new_mod_rs.append(line)
-            if direction == 'clientbound':
-                ignore_lines = True
-                for packet_id, packet_class_name in zip(packet_ids, packet_class_names):
-                    new_mod_rs.append(
-                        make_packet_mod_rs_line(packet_id, packet_class_name)
-                    )
-                    required_modules.append(packet_class_name)
-            else:
-                ignore_lines = False
-            continue
-        elif line.strip() in ('}', '},'):
-            ignore_lines = False
-        elif line.strip().startswith('pub mod '):
-            continue
-
-        if not ignore_lines:
-            new_mod_rs.append(line)
-            # 0x00: clientbound_status_response_packet::ClientboundStatusResponsePacket,
-            if line.strip().startswith('0x'):
-                required_modules.append(
-                    line.strip().split(':')[1].split('::')[0].strip())
-
-    for i, required_module in enumerate(required_modules):
-        if required_module not in mod_rs:
-            new_mod_rs.insert(i, f'pub mod {required_module};')
-
-    with open(mod_rs_dir, 'w') as f:
-        f.write('\n'.join(new_mod_rs))
 
+        serverbound_packets = packet_direction_report_to_packet_names(packets_report[mojmap_state]['serverbound'])
+        clientbound_packets = packet_direction_report_to_packet_names(packets_report[mojmap_state].get('clientbound', {}))
+
+        code = []
+        code.append('// NOTE: This file is generated automatically by codegen/packet.py.')
+        code.append("// Don't edit it directly!")
+        code.append('')
+        code.append('use azalea_protocol_macros::declare_state_packets;')
+        code.append('')
+        code.append(f'declare_state_packets!({to_camel_case(state)}Packet,')
+        code.append('    Clientbound => [')
+        for packet_name in clientbound_packets:
+            code.append(f'        {packet_name},')
+        code.append('    ],')
+        code.append('    Serverbound => [')
+        for packet_name in serverbound_packets:
+            code.append(f'        {packet_name},')
+        code.append('    ]')
+        code.append(');')
+        code.append('')
+
+        with open(mod_rs_dir, 'w') as f:
+            f.write('\n'.join(code))
+
+def packet_direction_report_to_packet_names(report):
+    name_to_id = {}
+    for resource_location, packet in report.items():
+        packet_id = packet['protocol_id']
+        name_to_id[resource_location.split(':')[-1]] = packet_id
+    
+    names_sorted = [name for name in sorted(name_to_id, key=lambda x: name_to_id[x])]
+    return names_sorted
 
 def get_packets(direction: str, state: str):
     mod_rs_dir = get_dir_location(
diff --git a/codegen/lib/code/shapes.py b/codegen/lib/code/shapes.py
index 12b333d1..c4964fa6 100755
--- a/codegen/lib/code/shapes.py
+++ b/codegen/lib/code/shapes.py
@@ -5,11 +5,11 @@ COLLISION_BLOCKS_RS_DIR = get_dir_location(
     '../azalea-physics/src/collision/blocks.rs')
 
 
-def generate_block_shapes(blocks_pixlyzer: dict, shapes: dict, aabbs: dict, block_states_report, block_datas_burger, mappings: Mappings):
+def generate_block_shapes(blocks_pixlyzer: dict, shapes: dict, aabbs: dict, block_states_report):
     blocks, shapes = simplify_shapes(blocks_pixlyzer, shapes, aabbs)
 
     code = generate_block_shapes_code(
-        blocks, shapes, block_states_report, block_datas_burger, mappings)
+        blocks, shapes, block_states_report)
     with open(COLLISION_BLOCKS_RS_DIR, 'w') as f:
         f.write(code)
 
@@ -63,7 +63,7 @@ def simplify_shapes(blocks: dict, shapes: dict, aabbs: dict):
     return new_blocks, new_shapes
 
 
-def generate_block_shapes_code(blocks: dict, shapes: dict, block_states_report, block_datas_burger, mappings: Mappings):
+def generate_block_shapes_code(blocks: dict, shapes: dict, block_states_report):
     # look at __cache__/generator-mod-*/blockCollisionShapes.json for format of blocks and shapes
 
     generated_shape_code = ''
@@ -75,7 +75,9 @@ def generate_block_shapes_code(blocks: dict, shapes: dict, block_states_report,
     empty_shapes = []
     full_shapes = []
 
-    block_state_ids_to_shape_ids = []
+    # the index into this list is the block state id
+    shapes_map = []
+
     for block_id, shape_ids in blocks.items():
         if isinstance(shape_ids, int):
             shape_ids = [shape_ids]
@@ -89,17 +91,19 @@ def generate_block_shapes_code(blocks: dict, shapes: dict, block_states_report,
             elif shape_id == 1 :
                 full_shapes.append(block_state_id)
 
-            block_state_ids_to_shape_ids.append((block_state_id, shape_id))
-
+            while len(shapes_map) <= block_state_id:
+                # default to shape 1 for missing shapes (full block)
+                shapes_map.append(1)
+            shapes_map[block_state_id] = shape_id
 
-    generated_map_code = f'static SHAPES_MAP: [&LazyLock<VoxelShape>; {len(block_state_ids_to_shape_ids)}] = ['
+    
 
-    block_state_ids_to_shape_ids = sorted(block_state_ids_to_shape_ids, key=lambda x: x[0])
+    generated_map_code = f'static SHAPES_MAP: [&LazyLock<VoxelShape>; {len(shapes_map)}] = ['
 
     empty_shape_match_code = convert_ints_to_rust_ranges(empty_shapes)
     block_shape_match_code = convert_ints_to_rust_ranges(full_shapes)
 
-    for block_state_id, shape_id in block_state_ids_to_shape_ids:
+    for block_state_id, shape_id in enumerate(shapes_map):
         generated_map_code += f'&SHAPE{shape_id},\n'
     generated_map_code += '];'
 
@@ -109,7 +113,7 @@ def generate_block_shapes_code(blocks: dict, shapes: dict, block_states_report,
     return f'''
 //! Autogenerated block collisions for every block
 
-// This file is generated from codegen/lib/code/block_shapes.py. If you want to
+// This file is generated from codegen/lib/code/shapes.py. If you want to
 // modify it, change that file.
 
 #![allow(clippy::explicit_auto_deref)]
diff --git a/codegen/lib/code/utils.py b/codegen/lib/code/utils.py
index 29adc247..1a87b7f6 100755
--- a/codegen/lib/code/utils.py
+++ b/codegen/lib/code/utils.py
@@ -131,7 +131,7 @@ def burger_type_to_rust_type(burger_type, field_name: Optional[str] = None, inst
             field_type_rs = to_camel_case(
                 enum_name.split('.')[-1].split('$')[-1])
             extra_code.append('')
-            extra_code.append(f'#[derive(McBuf, Clone, Copy, Debug)]')
+            extra_code.append(f'#[derive(AzBuf, Clone, Copy, Debug)]')
             extra_code.append(f'pub enum {field_type_rs} {{')
             for index, variant in enumerate(enum_variants):
                 extra_code.append(
@@ -166,8 +166,8 @@ def burger_type_to_rust_type(burger_type, field_name: Optional[str] = None, inst
     return field_type_rs, is_var, uses, extra_code
 
 
-def write_packet_file(state, packet_name_snake_case, code):
-    with open(get_dir_location(f'../azalea-protocol/src/packets/{state}/{packet_name_snake_case}.rs'), 'w') as f:
+def write_packet_file(state, packet_module_name, code):
+    with open(get_dir_location(f'../azalea-protocol/src/packets/{state}/{packet_module_name}.rs'), 'w') as f:
         f.write(code)
 
 
diff --git a/codegen/lib/code/version.py b/codegen/lib/code/version.py
index c2053da3..d4a37232 100755
--- a/codegen/lib/code/version.py
+++ b/codegen/lib/code/version.py
@@ -31,7 +31,28 @@ def set_version_id(version_id: str) -> None:
 
     with open(README_DIR, 'wb') as f:
         f.write(readme_text.encode())
-
+    
+    # update the version in all Cargo.toml files
+    # version = "0.10.3+mc1.21.1"
+    for root, _, files in os.walk(get_dir_location('..')):
+        for file in files:
+            if file == 'Cargo.toml':
+                with open(os.path.join(root, file), 'r') as f:
+                    cargo_toml = f.read().splitlines()
+                for i, line in enumerate(cargo_toml):
+                    if line.strip().startswith('version = '):
+                        replaced = re.sub(r'\+mc[^"]+?"', f'+mc{version_id}"', line)
+                        cargo_toml[i] = replaced
+                        break
+                else:
+                    # didn't have a version line
+                    continue
+                if cargo_toml[-1] != '':
+                    # make sure there's always a trailing newline
+                    cargo_toml.append('')
+                with open(os.path.join(root, file), 'w') as f:
+                    f.write('\n'.join(cargo_toml))
+    print('Updated version in README.md and Cargo.toml files')
 
 def get_protocol_version() -> str:
     # azalea-protocol/src/packets/mod.rs
@@ -49,7 +70,7 @@ def set_protocol_version(protocol_version: str) -> None:
     with open(get_dir_location('../azalea-protocol/src/packets/mod.rs'), 'r') as f:
         mod_rs = f.read().splitlines()
     for i, line in enumerate(mod_rs):
-        if line.strip().startswith('pub const PROTOCOL_VERSION'):
+        if line.strip().startswith('pub const PROTOCOL_VERSION:'):
             mod_rs[i] = f'pub const PROTOCOL_VERSION: i32 = {protocol_version};'
             break
     else:
@@ -58,3 +79,16 @@ def set_protocol_version(protocol_version: str) -> None:
 
     with open(get_dir_location('../azalea-protocol/src/packets/mod.rs'), 'w') as f:
         f.write('\n'.join(mod_rs))
+def set_version_name(version_name: str) -> None:
+    with open(get_dir_location('../azalea-protocol/src/packets/mod.rs'), 'r') as f:
+        mod_rs = f.read().splitlines()
+    for i, line in enumerate(mod_rs):
+        if line.strip().startswith('pub const VERSION_NAME:'):
+            mod_rs[i] = f'pub const VERSION_NAME: &str = "{version_name}";'
+            break
+    else:
+        raise Exception(
+            'Could not find version name in azalea-protocol/src/packets/mod.rs')
+
+    with open(get_dir_location('../azalea-protocol/src/packets/mod.rs'), 'w') as f:
+        f.write('\n'.join(mod_rs))
diff --git a/codegen/lib/extract.py b/codegen/lib/extract.py
index da69b753..4119ad55 100755
--- a/codegen/lib/extract.py
+++ b/codegen/lib/extract.py
@@ -1,7 +1,7 @@
 # Extracting data from the Minecraft jars
 
 from typing import TYPE_CHECKING
-from lib.download import get_mappings_for_version, get_server_jar, get_burger, get_client_jar, get_pixlyzer, get_yarn_data, get_fabric_api_versions, get_fabric_loader_versions
+from lib.download import get_server_jar, get_burger, get_client_jar, get_pixlyzer, get_yarn_data, get_fabric_api_versions, get_fabric_loader_versions
 from lib.utils import get_dir_location, to_camel_case, upper_first_letter
 from zipfile import ZipFile
 import subprocess
@@ -23,17 +23,16 @@ def generate_data_from_server_jar(version_id: str):
 
 
 def get_block_states_report(version_id: str):
-    generate_data_from_server_jar(version_id)
-    with open(get_dir_location(f'__cache__/generated-{version_id}/reports/blocks.json'), 'r') as f:
-        return json.load(f)
-
-
+    return get_report(version_id, 'blocks')
 def get_registries_report(version_id: str):
+    return get_report(version_id, 'registries')
+def get_packets_report(version_id: str):
+    return get_report(version_id, 'packets')
+def get_report(version_id: str, name: str):
     generate_data_from_server_jar(version_id)
-    with open(get_dir_location(f'__cache__/generated-{version_id}/reports/registries.json'), 'r') as f:
+    with open(get_dir_location(f'__cache__/generated-{version_id}/reports/{name}.json'), 'r') as f:
         return json.load(f)
 
-
 def get_registry_tags(version_id: str, name: str):
     generate_data_from_server_jar(version_id)
     tags_directory = get_dir_location(f'__cache__/generated-{version_id}/data/minecraft/tags/{name}')
@@ -50,17 +49,6 @@ def get_registry_tags(version_id: str, name: str):
                 tags[relative_path[:-5]] = json.load(f)
     return tags
 
-
-def get_block_states_burger(version_id: str):
-    burger_data = get_burger_data_for_version(version_id)
-    return burger_data[0]['blocks']['block']
-
-
-def get_ordered_blocks_burger(version_id: str):
-    burger_data = get_burger_data_for_version(version_id)
-    return burger_data[0]['blocks']['ordered_blocks']
-
-
 python_command = None