Add support for custom attributes

Signed-off-by: Martin Kröning <[email protected]>

Author: mkroening

CHANGELOG.md

@@ -235,6 +235,7 @@
 - Add option to use DST structs for flexible arrays (--flexarray-dst, #2772).
 - Add option to dynamically load variables (#2812).
 - Add option in CLI to use rustified non-exhaustive enums (--rustified-non-exhaustive-enum, #2847).
+- Add support for custom attributes (--with-attribute-custom, #2866)
 ## Changed
 - Remove which and lazy-static dependencies (#2809, #2817).
 - Generate compile-time layout tests (#2787).

Cargo.lock

@@ -25,6 +25,7 @@ clap = { version = "4", features = ["derive"] }
 clap_complete = "4"
 env_logger = { version = "0.10.0", optional = true }
 log = { version = "0.4", optional = true }
+proc-macro2 = { version = "1", default-features = false }
 shlex = "1"
 
 [features]

bindgen-cli/Cargo.toml

@@ -6,10 +6,12 @@ use bindgen::{
 };
 use clap::error::{Error, ErrorKind};
 use clap::{CommandFactory, Parser};
+use proc_macro2::TokenStream;
 use std::fs::File;
 use std::io;
 use std::path::{Path, PathBuf};
 use std::process::exit;
+use std::str::FromStr;
 
 fn rust_target_help() -> String {
     format!(
@@ -87,6 +89,43 @@ fn parse_custom_derive(
     Ok((derives, regex.to_owned()))
 }
 
+fn parse_custom_attribute(
+    custom_attribute: &str,
+) -> Result<(Vec<String>, String), Error> {
+    let mut brace_level = 0;
+    let (regex, attributes) = custom_attribute
+        .rsplit_once(|c| {
+            match c {
+                ']' => brace_level += 1,
+                '[' => brace_level -= 1,
+                _ => {}
+            }
+            c == '=' && brace_level == 0
+        })
+        .ok_or_else(|| Error::raw(ErrorKind::InvalidValue, "Missing `=`"))?;
+
+    let mut brace_level = 0;
+    let attributes = attributes
+        .split(|c| {
+            match c {
+                ']' => brace_level += 1,
+                '[' => brace_level -= 1,
+                _ => {}
+            }
+            c == ',' && brace_level == 0
+        })
+        .map(|s| s.to_owned())
+        .collect::<Vec<_>>();
+
+    for attribute in &attributes {
+        if let Err(err) = TokenStream::from_str(attribute) {
+            return Err(Error::raw(ErrorKind::InvalidValue, err));
+        }
+    }
+
+    Ok((attributes, regex.to_owned()))
+}
+
 #[derive(Parser, Debug)]
 #[clap(
     about = "Generates Rust bindings from C/C++ headers.",
@@ -424,6 +463,18 @@ struct BindgenCommand {
     /// Derive custom traits on a `union`. The CUSTOM value must be of the shape REGEX=DERIVE where DERIVE is a coma-separated list of derive macros.
     #[arg(long, value_name = "CUSTOM", value_parser = parse_custom_derive)]
     with_derive_custom_union: Vec<(Vec<String>, String)>,
+    /// Add custom attributes on any kind of type. The CUSTOM value must be of the shape REGEX=ATTRIBUTE where ATTRIBUTE is a coma-separated list of attributes.
+    #[arg(long, value_name = "CUSTOM", value_parser = parse_custom_attribute)]
+    with_attribute_custom: Vec<(Vec<String>, String)>,
+    /// Add custom attributes on a `struct`. The CUSTOM value must be of the shape REGEX=ATTRIBUTE where ATTRIBUTE is a coma-separated list of attributes.
+    #[arg(long, value_name = "CUSTOM", value_parser = parse_custom_attribute)]
+    with_attribute_custom_struct: Vec<(Vec<String>, String)>,
+    /// Add custom attributes on an `enum. The CUSTOM value must be of the shape REGEX=ATTRIBUTE where ATTRIBUTE is a coma-separated list of attributes.
+    #[arg(long, value_name = "CUSTOM", value_parser = parse_custom_attribute)]
+    with_attribute_custom_enum: Vec<(Vec<String>, String)>,
+    /// Add custom attributes on a `union`. The CUSTOM value must be of the shape REGEX=ATTRIBUTE where ATTRIBUTE is a coma-separated list of attributes.
+    #[arg(long, value_name = "CUSTOM", value_parser = parse_custom_attribute)]
+    with_attribute_custom_union: Vec<(Vec<String>, String)>,
     /// Generate wrappers for `static` and `static inline` functions.
     #[arg(long, requires = "experimental")]
     wrap_static_fns: bool,
@@ -574,6 +625,10 @@ where
         with_derive_custom_struct,
         with_derive_custom_enum,
         with_derive_custom_union,
+        with_attribute_custom,
+        with_attribute_custom_struct,
+        with_attribute_custom_enum,
+        with_attribute_custom_union,
         wrap_static_fns,
         wrap_static_fns_path,
         wrap_static_fns_suffix,
@@ -1130,6 +1185,82 @@ where
         }
     }
 
+    #[derive(Debug)]
+    struct CustomAttributeCallback {
+        attributes: Vec<String>,
+        kind: Option<TypeKind>,
+        regex_set: bindgen::RegexSet,
+    }
+
+    impl bindgen::callbacks::ParseCallbacks for CustomAttributeCallback {
+        fn cli_args(&self) -> Vec<String> {
+            let mut args = vec![];
+
+            let flag = match &self.kind {
+                None => "--with-attribute-custom",
+                Some(TypeKind::Struct) => "--with-attribute-custom-struct",
+                Some(TypeKind::Enum) => "--with-attribute-custom-enum",
+                Some(TypeKind::Union) => "--with-attribute-custom-union",
+            };
+
+            let attributes = self.attributes.join(",");
+
+            for item in self.regex_set.get_items() {
+                args.extend_from_slice(&[
+                    flag.to_owned(),
+                    format!("{}={}", item, attributes),
+                ]);
+            }
+
+            args
+        }
+
+        fn add_attributes(
+            &self,
+            info: &bindgen::callbacks::AttributeInfo<'_>,
+        ) -> Vec<String> {
+            if self.kind.map(|kind| kind == info.kind).unwrap_or(true) &&
+                self.regex_set.matches(info.name)
+            {
+                return self.attributes.clone();
+            }
+            vec![]
+        }
+    }
+
+    for (custom_attributes, kind, name) in [
+        (with_attribute_custom, None, "--with-attribute-custom"),
+        (
+            with_attribute_custom_struct,
+            Some(TypeKind::Struct),
+            "--with-attribute-custom-struct",
+        ),
+        (
+            with_attribute_custom_enum,
+            Some(TypeKind::Enum),
+            "--with-attribute-custom-enum",
+        ),
+        (
+            with_attribute_custom_union,
+            Some(TypeKind::Union),
+            "--with-attribute-custom-union",
+        ),
+    ] {
+        let name = emit_diagnostics.then_some(name);
+        for (attributes, regex) in custom_attributes {
+            let mut regex_set = RegexSet::new();
+            regex_set.insert(regex);
+            regex_set.build_with_diagnostics(false, name);
+
+            builder =
+                builder.parse_callbacks(Box::new(CustomAttributeCallback {
+                    attributes,
+                    kind,
+                    regex_set,
+                }));
+        }
+    }
+
     if wrap_static_fns {
         builder = builder.wrap_static_fns(true);
     }

bindgen-cli/options.rs

@@ -133,6 +133,18 @@ impl ParseCallbacks for MacroCallback {
             vec![]
         }
     }
+
+    // Test the "custom attributes" capability.
+    fn add_attributes(
+        &self,
+        info: &bindgen::callbacks::AttributeInfo<'_>,
+    ) -> Vec<String> {
+        if info.name == "Test" {
+            vec!["#[cfg_attr(test, derive(PartialOrd))]".into()]
+        } else {
+            vec![]
+        }
+    }
 }
 
 impl Drop for MacroCallback {

bindgen-integration/build.rs

@@ -297,6 +297,15 @@ fn test_custom_derive() {
     assert!(!(test1 > test2));
 }
 
+#[test]
+fn test_custom_attributes() {
+    // The `add_attributes` callback should have added `#[cfg_attr(test, derive(PartialOrd))])`
+    // to the `Test` struct. If it didn't, this will fail to compile.
+    let test1 = unsafe { bindings::Test::new(5) };
+    let test2 = unsafe { bindings::Test::new(6) };
+    assert!(test1 < test2);
+}
+
 #[test]
 fn test_wrap_static_fns() {
     // GH-1090: https://github.com/rust-lang/rust-bindgen/issues/1090

bindgen-integration/src/lib.rs

@@ -10,6 +10,7 @@ clap = { version = "4", features = ["derive"] }
 clap_complete = "4"
 shlex = "1"
 prettyplease = { version = "0.2.7", features = ["verbatim"] }
+proc-macro2 = { version = "1", default-features = false }
 syn = { version = "2.0" }
 tempfile = "3"
 similar = { version = "2.2.1", features = ["inline"] }

bindgen-tests/Cargo.toml

@@ -0,0 +1,14 @@
+// bindgen-flags: --default-enum-style rust --default-non-copy-union-style manually_drop --no-default=".*" --no-hash=".*" --no-partialeq=".*" --no-debug=".*" --no-copy=".*" --with-attribute-custom="foo_[^e].*=#[doc(hidden)]" --with-attribute-custom-struct="foo.*=#[derive(Default)]" --with-attribute-custom-enum="foo.*=#[cfg_attr(test, derive(PartialOrd, Copy))]" --with-attribute-custom-union="foo.*=#[derive(Clone)],#[derive(Copy)]"
+struct foo_struct {
+    int inner;
+};
+enum foo_enum {
+    inner = 0
+};
+union foo_union {
+    int fst;
+    float snd;
+};
+struct non_matching {
+    int inner;
+};