1 files changed, 39 insertions, 2 deletions
diff --git a/codegen/lib/code/registry.py b/codegen/lib/code/registry.py
index c22eefe9..4df82c09 100755
--- a/codegen/lib/code/registry.py
+++ b/codegen/lib/code/registry.py
@@ -9,8 +9,45 @@ REGISTRIES_DIR = get_dir_location('../azalea-registry/src/lib.rs')
 def generate_registries(registries: dict):
     code = []
 
-    code.append('use azalea_registry_macros::registry;')
-    code.append('')
+    code.append('''// 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),
+        }
+    }
+}
+''')
 
     for registry_name, registry in registries.items():
         # registry!(Block, {