Encoder packet ID generation uses enum names

Author: robertwayne

rust/bufferfish/src/compiler.rs

@@ -138,23 +138,46 @@ fn generate_output_string(input: Vec<String>, output: &mut String) -> Result<(),
     output.push_str("/* AUTOGENERATED BUFFERFISH FILE, DO NOT EDIT */\n");
     output.push_str("import { Bufferfish } from 'bufferfish'\n");
 
-    for path in input {
-        let items = parse_rust_source_file(&path)?;
+    let mut packet_id_enum_names = Vec::new();
+    let mut all_structs = Vec::new();
+    let mut all_enums = Vec::new();
+
+    for path in &input {
+        let items = parse_rust_source_file(path)?;
         let (structs, enums) = get_items_implementing_encode(items);
 
-        for item in &enums {
-            generate_typescript_enum_defs(item.clone(), output);
-            generate_typescript_enum_decoders(item.clone(), output);
-            generate_typescript_enum_encoders(item.clone(), output);
+        // Find struct attributes with #[bufferfish(...)] to identify packet ID enums
+        for item in &structs {
+            if let Some(packet_id) = get_packet_id(&item.attrs) {
+                if let Some(enum_name) = extract_enum_name_from_packet_id(&packet_id) {
+                    if !packet_id_enum_names.contains(&enum_name) {
+                        packet_id_enum_names.push(enum_name);
+                    }
+                }
+            }
         }
 
-        for item in &structs {
-            generate_typescript_struct_defs(item.clone(), output);
-            generate_typescript_struct_decoders(item.clone(), output);
-            generate_typescript_struct_encoders(item.clone(), output);
+        all_structs.extend(structs);
+        all_enums.extend(enums);
+    }
+
+    for item in &all_enums {
+        generate_typescript_enum_defs(item.clone(), output);
+        generate_typescript_enum_decoders(item.clone(), output);
+        generate_typescript_enum_encoders(item.clone(), output);
+
+        let enum_name = item.ident.to_string();
+        if packet_id_enum_names.contains(&enum_name) {
+            generate_typescript_packet_id_encoder(item.clone(), output);
         }
     }
 
+    for item in &all_structs {
+        generate_typescript_struct_defs(item.clone(), output);
+        generate_typescript_struct_decoders(item.clone(), output);
+        generate_typescript_struct_encoders(item.clone(), output, &packet_id_enum_names);
+    }
+
     Ok(())
 }
 
@@ -192,7 +215,6 @@ fn get_packet_id(attrs: &[Attribute]) -> Option<String> {
         if attr.path().is_ident("bufferfish") {
             if let Meta::List(list) = &attr.meta {
                 let tokens = list.tokens.to_string();
-
                 let cleaned = tokens.trim().replace(" ", "").replace("::", ".");
 
                 return Some(cleaned);
@@ -202,8 +224,46 @@ fn get_packet_id(attrs: &[Attribute]) -> Option<String> {
     None
 }
 
+/// Extract the enum name from a packet ID reference like "EnumName.Variant"
+fn extract_enum_name_from_packet_id(id: &str) -> Option<String> {
+    id.split('.').next().map(|s| s.to_string())
+}
+
 /// Generate TypeScript encoders for structs including PacketID if specified
-fn generate_typescript_struct_encoders(item: ItemStruct, output: &mut String) {
+/// Generate a TypeScript encoder function for a specific packet ID enum
+fn generate_typescript_packet_id_encoder(item: ItemEnum, output: &mut String) {
+    let enum_name = item.ident.to_string();
+    let repr_type = get_repr_type(&item.attrs).unwrap_or("u8".to_string());
+
+    let write_fn = match repr_type.as_str() {
+        "u8" => "writeUint8",
+        "u16" => "writeUint16",
+        "u32" => "writeUint32",
+        "u64" => "writeUint64",
+        "u128" => "writeUint128",
+        "i8" => "writeInt8",
+        "i16" => "writeInt16",
+        "i32" => "writeInt32",
+        "i64" => "writeInt64",
+        "i128" => "writeInt128",
+        _ => panic!("Unsupported repr type"),
+    };
+
+    output.push_str(
+        format!(
+            "\nexport function encode{enum_name}(bf: Bufferfish, value: {enum_name}): void {{\n",
+        )
+        .as_str(),
+    );
+    output.push_str(format!("    bf.{write_fn}(value)\n").as_str());
+    output.push_str("}\n");
+}
+
+fn generate_typescript_struct_encoders(
+    item: ItemStruct,
+    output: &mut String,
+    packet_id_enums: &[String],
+) {
     let struct_name = item.ident.to_string();
     let packet_id = get_packet_id(&item.attrs);
 
@@ -223,10 +283,13 @@ fn generate_typescript_struct_encoders(item: ItemStruct, output: &mut String) {
             );
 
             if let Some(id) = &packet_id {
-                output.push_str(format!("    encodePacketId(bf, {id})\n").as_str());
+                if let Some(enum_name) = extract_enum_name_from_packet_id(id) {
+                    if packet_id_enums.contains(&enum_name) {
+                        output.push_str(format!("    encode{enum_name}(bf, {id})\n").as_str());
+                    }
+                }
             }
 
-            // Write struct fields
             for field in &fields_named.named {
                 if let Some(field_name) = &field.ident {
                     let field_ts_name = snake_to_camel_case(field_name.to_string());
@@ -260,10 +323,13 @@ fn generate_typescript_struct_encoders(item: ItemStruct, output: &mut String) {
             );
 
             if let Some(id) = &packet_id {
-                output.push_str(format!("    encodePacketId(bf, {id})\n").as_str());
+                if let Some(enum_name) = extract_enum_name_from_packet_id(id) {
+                    if packet_id_enums.contains(&enum_name) {
+                        output.push_str(format!("    encode{enum_name}(bf, {id})\n").as_str());
+                    }
+                }
             }
 
-            // Write tuple fields
             for (i, field) in fields_unnamed.unnamed.iter().enumerate() {
                 output.push_str(
                     format!(
@@ -285,7 +351,13 @@ fn generate_typescript_struct_encoders(item: ItemStruct, output: &mut String) {
                     )
                     .as_str(),
                 );
-                output.push_str(format!("    encodePacketId(bf, {id})\n").as_str());
+
+                if let Some(enum_name) = extract_enum_name_from_packet_id(&id) {
+                    if packet_id_enums.contains(&enum_name) {
+                        output.push_str(format!("    encode{enum_name}(bf, {id})\n").as_str());
+                    }
+                }
+
                 output.push_str("}\n");
             }
         }
@@ -511,7 +583,7 @@ fn generate_typescript_struct_decoders(item: ItemStruct, lines: &mut String) {
                 }
             }
 
-            lines.push_str("    };\n");
+            lines.push_str("    }\n");
             lines.push_str("}\n");
         }
         Fields::Unnamed(fields_unnamed) => {
@@ -616,7 +688,7 @@ mod tests {
 
     #[test]
     fn test_ts_generation() {
-        let input = r#"
+        let test_file = r#"
 #[derive(Encode)]
 #[repr(u16)]
 pub enum PacketId {
@@ -658,6 +730,10 @@ export function encodePacketId(bf: Bufferfish, value: PacketId): void {
     bf.writeUint16(value)
 }
 
+export function encodePacketId(bf: Bufferfish, value: PacketId): void {
+    bf.writeUint16(value)
+}
+
 export interface JoinPacket {
     id: number
     username: string
@@ -667,7 +743,7 @@ export function decodeJoinPacket(bf: Bufferfish): JoinPacket {
     return {
         id: bf.readUint8() as number,
         username: bf.readString() as string,
-    };
+    }
 }
 
 export function encodeJoinPacket(bf: Bufferfish, value: JoinPacket): void {
@@ -696,25 +772,48 @@ export function encodeUnknownPacket(bf: Bufferfish, value: UnknownPacket): void
 }
 "#;
 
+        let items = syn::parse_file(test_file)
+            .unwrap_or_else(|e| panic!("Failed to parse: {e}"))
+            .items;
+        let (structs, enums) = get_items_implementing_encode(items);
+
         let mut output = String::new();
         output.push_str("/* AUTOGENERATED BUFFERFISH FILE, DO NOT EDIT */\n");
         output.push_str("import { Bufferfish } from 'bufferfish'\n");
 
-        let syntax_tree = syn::parse_file(input).unwrap();
-        let (structs, enums) = get_items_implementing_encode(syntax_tree.items);
+        let mut packet_id_enum_names = Vec::new();
 
-        for item in enums {
+        for item in &structs {
+            if let Some(packet_id) = get_packet_id(&item.attrs) {
+                if let Some(enum_name) = extract_enum_name_from_packet_id(&packet_id) {
+                    if !packet_id_enum_names.contains(&enum_name) {
+                        packet_id_enum_names.push(enum_name);
+                    }
+                }
+            }
+        }
+
+        for item in &enums {
             generate_typescript_enum_defs(item.clone(), &mut output);
             generate_typescript_enum_decoders(item.clone(), &mut output);
             generate_typescript_enum_encoders(item.clone(), &mut output);
+
+            let enum_name = item.ident.to_string();
+            if packet_id_enum_names.contains(&enum_name) {
+                generate_typescript_packet_id_encoder(item.clone(), &mut output);
+            }
         }
 
-        for item in structs {
+        for item in &structs {
             generate_typescript_struct_defs(item.clone(), &mut output);
             generate_typescript_struct_decoders(item.clone(), &mut output);
-            generate_typescript_struct_encoders(item.clone(), &mut output);
+            generate_typescript_struct_encoders(item.clone(), &mut output, &packet_id_enum_names);
         }
 
-        assert_eq!(output, expected_output);
+        if output != expected_output {
+            println!("Expected:\n{expected_output}");
+            println!("Got:\n{output}");
+            panic!("Output does not match expected output");
+        }
     }
 }