1 files changed, 26 insertions, 47 deletions

diff --git a/codegen/lib/code/registry.py b/codegen/lib/code/registry.py
index 1e9d9f43..a67b5e4d 100755
--- a/codegen/lib/code/registry.py
+++ b/codegen/lib/code/registry.py
@@ -7,49 +7,8 @@ REGISTRIES_DIR = get_dir_location('../azalea-registry/src/lib.rs')
 
 
 def generate_registries(registries: dict):
-    code = []
-
-    code.append('''#![doc = include_str!("../README.md")]
-
-// This file is automatically generated in codegen/lib/code/registry.py
-
-use azalea_buf::{BufReadError, McBufReadable, McBufVarReadable, McBufVarWritable, McBufWritable};
-use azalea_registry_macros::registry;
-use std::io::{Cursor, Write};
-
-pub trait Registry
-where
-    Self: Sized,
-{
-    fn from_u32(value: u32) -> Option<Self>;
-    fn to_u32(&self) -> u32;
-}
-
-/// A registry that might not be present. This is transmitted as a single
-/// varint in the protocol.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub struct OptionalRegistry<T: Registry>(Option<T>);
-
-impl<T: Registry> McBufReadable for OptionalRegistry<T> {
-    fn read_from(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
-        Ok(OptionalRegistry(match u32::var_read_from(buf)? {
-            0 => None,
-            value => Some(
-                T::from_u32(value - 1)
-                    .ok_or(BufReadError::UnexpectedEnumVariant { id: value as i32 })?,
-            ),
-        }))
-    }
-}
-impl<T: Registry> McBufWritable for OptionalRegistry<T> {
-    fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
-        match &self.0 {
-            None => 0u32.var_write_into(buf),
-            Some(value) => (value.to_u32() + 1).var_write_into(buf),
-        }
-    }
-}
-''')
+    with open(REGISTRIES_DIR, 'r') as f:
+        code = f.read().split('\n')
 
     for registry_name, registry in registries.items():
         # registry!(Block, {
@@ -64,15 +23,35 @@ impl<T: Registry> McBufWritable for OptionalRegistry<T> {
 
         registry_struct_name = to_camel_case(registry_name.split(':')[1])
 
-        code.append(f'registry!({registry_struct_name}, {{')
+        registry_code = []
+        registry_code.append(f'enum {registry_struct_name} {{')
         registry_entries = sorted(
             registry['entries'].items(), key=lambda x: x[1]['protocol_id'])
         for variant_name, _variant in registry_entries:
             variant_struct_name = to_camel_case(
                 variant_name.split(':')[1])
-            code.append(f'\t{variant_struct_name} => "{variant_name}",')
-        code.append('});')
-        code.append('')
+            registry_code.append(f'\t{variant_struct_name} => "{variant_name}",')
+        registry_code.append('}')
+
+        # when we find a "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 == "registry! {":
+                in_registry_macro = True
+            elif in_registry_macro and line == registry_code[0]:
+                # found it, now delete until we get to "}"
+                while code[i] != '}':
+                    code.pop(i)
+                code[i] = '\n'.join(registry_code)
+                found = True
+                break
+        if not found:
+            code.append('registry! {')
+            code.append('\n'.join(registry_code))
+            code.append('}')
+            code.append('')
 
     with open(REGISTRIES_DIR, 'w') as f:
         f.write('\n'.join(code))