Add skip attribute for filtering codegen

Author: robertwayne

rust/bufferfish/src/compiler.rs

@@ -18,6 +18,18 @@ use syn::{
 /// Requires Rust types to be annotated with `#[derive(Encode)]` and/or
 /// `#[derive(Decode]` macros.
 pub fn generate(src_path: &str, output_dst: &str) -> io::Result<()> {
+    generate_with_target(src_path, output_dst, None)
+}
+
+/// Generate a TypeScript file at `output_dst` from a directory of Rust source
+/// files at `src_path`, filtering by target (e.g., "client" or "server").
+///
+/// Items with `skip = "target"` will be excluded from generation.
+pub fn generate_with_target(
+    src_path: &str,
+    output_dst: &str,
+    target: Option<&str>,
+) -> io::Result<()> {
     let mut files = Vec::new();
 
     fn visit_dirs(dir: &std::path::Path, files: &mut Vec<String>) -> io::Result<()> {
@@ -38,7 +50,7 @@ pub fn generate(src_path: &str, output_dst: &str) -> io::Result<()> {
     visit_dirs(std::path::Path::new(src_path), &mut files)?;
 
     let mut output = String::new();
-    generate_output_string(files, &mut output)?;
+    generate_output_string_with_target(files, &mut output, target)?;
     write_typescript_file(output_dst, &output)?;
 
     Ok(())
@@ -134,7 +146,11 @@ fn get_items_implementing_encode(items: Vec<Item>) -> (Vec<ItemStruct>, Vec<Item
     (structs, enums)
 }
 
-fn generate_output_string(input: Vec<String>, output: &mut String) -> Result<(), std::io::Error> {
+fn generate_output_string_with_target(
+    input: Vec<String>,
+    output: &mut String,
+    target: Option<&str>,
+) -> Result<(), std::io::Error> {
     output.push_str("/* AUTOGENERATED BUFFERFISH FILE, DO NOT EDIT */\n");
     output.push_str("import { Bufferfish } from 'bufferfish'\n");
 
@@ -148,16 +164,38 @@ fn generate_output_string(input: Vec<String>, output: &mut String) -> Result<(),
         let (structs, enums) = get_items_implementing_encode(items);
 
         for item in &structs {
-            if let Some(packet_id) = get_packet_id(&item.attrs)
+            let (packet_id, skip_attr) = get_packet_attributes(&item.attrs);
+            if should_skip_for_target(&skip_attr, target) {
+                continue;
+            }
+
+            if let Some(packet_id) = packet_id
                 && let Some(enum_name) = extract_enum_name_from_packet_id(&packet_id)
                 && !packet_id_enum_names.contains(&enum_name)
             {
                 packet_id_enum_names.push(enum_name);
             }
         }
 
-        all_structs.extend(structs);
-        all_enums.extend(enums);
+        // Filter structs and enums based on target
+        let filtered_structs: Vec<_> = structs
+            .into_iter()
+            .filter(|item| {
+                let (_, skip_attr) = get_packet_attributes(&item.attrs);
+                !should_skip_for_target(&skip_attr, target)
+            })
+            .collect();
+
+        let filtered_enums: Vec<_> = enums
+            .into_iter()
+            .filter(|item| {
+                let skip_attr = get_enum_skip_attribute(&item.attrs);
+                !should_skip_for_target(&skip_attr, target)
+            })
+            .collect();
+
+        all_structs.extend(filtered_structs);
+        all_enums.extend(filtered_enums);
     }
 
     for item in &all_enums {
@@ -184,7 +222,44 @@ fn generate_output_string(input: Vec<String>, output: &mut String) -> Result<(),
     Ok(())
 }
 
-fn generate_typescript_enum_encoders(item: ItemEnum, output: &mut String) {
+fn get_packet_attributes(attrs: &[Attribute]) -> (Option<String>, Option<String>) {
+    for attr in attrs {
+        if attr.path().is_ident("bufferfish")
+            && let Meta::List(list) = &attr.meta
+        {
+            let tokens = list.tokens.to_string();
+            let parts: Vec<&str> = tokens.split(',').map(|s| s.trim()).collect();
+
+            let mut packet_id = None;
+            let mut skip_target = None;
+
+            for part in parts {
+                if part.contains("skip") && part.contains('=') {
+                    // Parse skip = "target"
+                    if let Some(target) = part.split('=').nth(1) {
+                        skip_target = Some(target.trim().trim_matches('"').to_string());
+                    }
+                } else if !part.contains('=') {
+                    // This is the packet ID
+                    packet_id = Some(part.replace(" ", "").replace("::", "."));
+                }
+            }
+
+            return (packet_id, skip_target);
+        }
+    }
+    (None, None)
+}
+
+fn should_skip_for_target(skip_attr: &Option<String>, target: Option<&str>) -> bool {
+    if let (Some(skip_target), Some(current_target)) = (skip_attr, target) {
+        skip_target == current_target
+    } else {
+        false
+    }
+}
+
+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());
 
@@ -212,29 +287,7 @@ fn generate_typescript_enum_encoders(item: ItemEnum, output: &mut String) {
     output.push_str("}\n");
 }
 
-/// Extract the PacketID from struct attributes and format it for TypeScript
-fn get_packet_id(attrs: &[Attribute]) -> Option<String> {
-    for attr in attrs {
-        if attr.path().is_ident("bufferfish")
-            && let Meta::List(list) = &attr.meta
-        {
-            let tokens = list.tokens.to_string();
-            let cleaned = tokens.trim().replace(" ", "").replace("::", ".");
-
-            return Some(cleaned);
-        }
-    }
-    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 a TypeScript encoder function for a specific packet ID enum
-/// This is used specifically for enum types that are used as packet IDs
-fn generate_typescript_packet_id_encoder(item: ItemEnum, output: &mut String) {
+fn generate_typescript_enum_encoders(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());
 
@@ -268,7 +321,7 @@ fn generate_typescript_struct_encoders(
     packet_id_enums: &[String],
 ) {
     let struct_name = item.ident.to_string();
-    let packet_id = get_packet_id(&item.attrs);
+    let (packet_id, _) = get_packet_attributes(&item.attrs);
 
     match &item.fields {
         Fields::Named(fields_named) => {
@@ -687,6 +740,28 @@ fn get_repr_type(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())
+}
+
+fn get_enum_skip_attribute(attrs: &[Attribute]) -> Option<String> {
+    for attr in attrs {
+        if attr.path().is_ident("bufferfish")
+            && let Meta::List(list) = &attr.meta
+        {
+            let tokens = list.tokens.to_string();
+            if tokens.contains("skip")
+                && tokens.contains('=')
+                && let Some(target) = tokens.split('=').nth(1)
+            {
+                return Some(target.trim().trim_matches('"').to_string());
+            }
+        }
+    }
+    None
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -784,7 +859,8 @@ export function encodeUnknownPacket(bf: Bufferfish, value: UnknownPacket): void
         let mut packet_id_enum_names = Vec::new();
 
         for item in &structs {
-            if let Some(packet_id) = get_packet_id(&item.attrs)
+            let (packet_id, _) = get_packet_attributes(&item.attrs);
+            if let Some(packet_id) = packet_id
                 && let Some(enum_name) = extract_enum_name_from_packet_id(&packet_id)
                 && !packet_id_enum_names.contains(&enum_name)
             {
@@ -820,4 +896,52 @@ export function encodeUnknownPacket(bf: Bufferfish, value: UnknownPacket): void
             panic!("Output does not match expected output");
         }
     }
+
+    #[test]
+    fn test_skip_attribute_parsing() {
+        let test_file = r#"
+#[derive(Encode)]
+#[bufferfish(PacketId::Join, skip = "client")]
+pub struct ClientOnlyPacket {
+    pub data: String,
+}
+
+#[derive(Encode)]
+#[bufferfish(PacketId::Leave, skip = "server")]
+pub struct ServerOnlyPacket {
+    pub data: String,
+}
+
+#[derive(Encode)]
+#[bufferfish(PacketId::Both)]
+pub struct BothPacket {
+    pub data: String,
+}
+    "#;
+
+        let items = syn::parse_file(test_file)
+            .unwrap_or_else(|e| panic!("Failed to parse: {e}"))
+            .items;
+        let (structs, _) = get_items_implementing_encode(items);
+
+        let client_structs: Vec<_> = structs
+            .iter()
+            .filter(|item| {
+                let (_, skip_attr) = get_packet_attributes(&item.attrs);
+                !should_skip_for_target(&skip_attr, Some("client"))
+            })
+            .collect();
+
+        assert_eq!(client_structs.len(), 2);
+
+        let server_structs: Vec<_> = structs
+            .iter()
+            .filter(|item| {
+                let (_, skip_attr) = get_packet_attributes(&item.attrs);
+                !should_skip_for_target(&skip_attr, Some("server"))
+            })
+            .collect();
+
+        assert_eq!(server_structs.len(), 2);
+    }
 }