Add support for custom attributes

Signed-off-by: Martin Kröning <martin.kroening@eonerc.rwth-aachen.de>

bindgen-cli/Cargo.toml

@@ -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/options.rs

@@ -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,85 @@ 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<TokenStream> {
+            if self.kind.map(|kind| kind == info.kind).unwrap_or(true) &&
+               self.regex_set.matches(info.name)
+            {
+                return self
+                    .attributes
+                    .iter()
+                    .map(|s| s.parse().unwrap())
+                    .collect();
+            }
+            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-integration/Cargo.toml

@@ -9,6 +9,8 @@ build = "build.rs"
 [build-dependencies]
 bindgen = { path = "../bindgen", features = ["experimental"] }
 cc = "1.0"
+proc-macro2 = { version = "1", default-features = false }
+quote = { version = "1", default-features = false }
 
 [features]
 static = ["bindgen/static"]

bindgen-integration/build.rs

@@ -1,9 +1,13 @@
 extern crate bindgen;
+extern crate proc_macro2;
+extern crate quote;
 
 use bindgen::callbacks::{
     DeriveInfo, IntKind, MacroParsingBehavior, ParseCallbacks,
 };
 use bindgen::{Builder, EnumVariation, Formatter};
+use proc_macro2::TokenStream;
+use quote::quote;
 use std::collections::HashSet;
 use std::env;
 use std::path::PathBuf;
@@ -133,6 +137,15 @@ impl ParseCallbacks for MacroCallback {
             vec![]
         }
     }
+
+    // Test the "custom derives" capability.
+    fn add_attributes(&self, info: &bindgen::callbacks::AttributeInfo<'_>) -> Vec<TokenStream> {
+        if info.name == "Test" {
+            vec![quote!(#[cfg_attr(test, derive(PartialOrd))])]
+        } else {
+            vec![]
+        }
+    }
 }
 
 impl Drop for MacroCallback {

bindgen-integration/src/lib.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/callbacks.rs

@@ -1,5 +1,7 @@
 //! A public API for more fine-grained customization of bindgen behavior.
 
+use proc_macro2::TokenStream;
+
 pub use crate::ir::analysis::DeriveTrait;
 pub use crate::ir::derive::CanDerive as ImplementsTrait;
 pub use crate::ir::enum_ty::{EnumVariantCustomBehavior, EnumVariantValue};
@@ -129,6 +131,14 @@ pub trait ParseCallbacks: fmt::Debug {
         vec![]
     }
 
+    /// Provide a list of custom attributes.
+    ///
+    /// If no additional attributes are wanted, this function should return an
+    /// empty `Vec`.
+    fn add_attributes(&self, _info: &AttributeInfo<'_>) -> Vec<TokenStream> {
+        vec![]
+    }
+
     /// Process a source code comment.
     fn process_comment(&self, _comment: &str) -> Option<String> {
         None
@@ -167,6 +177,17 @@ pub struct DeriveInfo<'a> {
     pub kind: TypeKind,
 }
 
+/// Relevant information about a type to which new attributes will be added using
+/// [`ParseCallbacks::add_attributes`].
+#[derive(Debug)]
+#[non_exhaustive]
+pub struct AttributeInfo<'a> {
+    /// The name of the type.
+    pub name: &'a str,
+    /// The kind of the type.
+    pub kind: TypeKind,
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 /// The kind of the current type.
 pub enum TypeKind {

bindgen/codegen/mod.rs

@@ -20,7 +20,7 @@ use self::struct_layout::StructLayoutTracker;
 
 use super::BindgenOptions;
 
-use crate::callbacks::{DeriveInfo, FieldInfo, TypeKind as DeriveTypeKind};
+use crate::callbacks::{AttributeInfo, DeriveInfo, FieldInfo, TypeKind as DeriveTypeKind};
 use crate::codegen::error::Error;
 use crate::ir::analysis::{HasVtable, Sizedness};
 use crate::ir::annotations::{
@@ -1047,6 +1047,15 @@ impl CodeGenerator for Type {
                             .extend(custom_derives.iter().map(|s| s.as_str()));
                         attributes.push(attributes::derives(&derives));
 
+                        let custom_attributes =
+                            ctx.options().all_callbacks(|cb| {
+                                cb.add_attributes(&AttributeInfo {
+                                    name: &name,
+                                    kind: DeriveTypeKind::Struct,
+                                })
+                            });
+                        attributes.extend(custom_attributes);
+
                         quote! {
                             #( #attributes )*
                             pub struct #rust_name
@@ -2377,6 +2386,18 @@ impl CodeGenerator for CompInfo {
             attributes.push(attributes::derives(&derives))
         }
 
+        let custom_attributes = ctx.options().all_callbacks(|cb| {
+            cb.add_attributes(&AttributeInfo {
+                name: &canonical_name,
+                kind: if is_rust_union {
+                    DeriveTypeKind::Union
+                } else {
+                    DeriveTypeKind::Struct
+                },
+            })
+        });
+        attributes.extend(custom_attributes);
+
         if item.must_use(ctx) {
             attributes.push(attributes::must_use());
         }
@@ -3568,6 +3589,16 @@ impl CodeGenerator for Enum {
             // In most cases this will be a no-op, since custom_derives will be empty.
             derives.extend(custom_derives.iter().map(|s| s.as_str()));
 
+            // The custom attribute callback may return a list of attributes;
+            // add them to the end of the list.
+            let custom_attributes = ctx.options().all_callbacks(|cb| {
+                cb.add_attributes(&AttributeInfo {
+                    name: &name,
+                    kind: DeriveTypeKind::Enum,
+                })
+            });
+            attrs.extend(custom_attributes);
+
             attrs.push(attributes::derives(&derives));
         }