From 9ca8866d05ef69b945f29524c0f034b5c85ccbe4 Mon Sep 17 00:00:00 2001
From: Robin Krahl <robin@nitrokey.com>
Date: Mon, 25 Mar 2024 08:56:23 +0100
Subject: [PATCH] derive: Allow multiple extensions attributes for
 ExtensionDispatch

This makes it easier to select the supported extensions using
conditional compilation (cfg_attr).
---
 derive/src/extension_dispatch.rs | 31 +++++++++++++++----------------
 derive/src/util.rs               | 10 +++++++---
 2 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/derive/src/extension_dispatch.rs b/derive/src/extension_dispatch.rs
index 235ff58b6aa..54d0e2a03af 100644
--- a/derive/src/extension_dispatch.rs
+++ b/derive/src/extension_dispatch.rs
@@ -152,13 +152,14 @@ impl ExtensionAttrs {
     fn new(input: &DeriveInput) -> Result<Self> {
         let mut extensions: HashMap<Ident, Path> = Default::default();
 
-        let attr = util::require_attr(input, &input.attrs, "extensions")?;
-        attr.parse_nested_meta(|meta| {
-            let ident = meta.path.require_ident()?;
-            let s: LitStr = meta.value()?.parse()?;
-            extensions.insert(ident.to_owned(), s.parse()?);
-            Ok(())
-        })?;
+        for attr in util::get_attrs(&input.attrs, "extensions") {
+            attr.parse_nested_meta(|meta| {
+                let ident = meta.path.require_ident()?;
+                let s: LitStr = meta.value()?.parse()?;
+                extensions.insert(ident.to_owned(), s.parse()?);
+                Ok(())
+            })?;
+        }
 
         Ok(Self { extensions })
     }
@@ -173,21 +174,19 @@ struct Backend {
 }
 
 impl Backend {
-    fn new(i: usize, field: &Field, extensions: &HashMap<Ident, Path>) -> Result<Self> {
+    fn new(i: usize, field: &Field, extension_types: &HashMap<Ident, Path>) -> Result<Self> {
         let ident = field.ident.clone().ok_or_else(|| {
             Error::new_spanned(
                 field,
                 "ExtensionDispatch can only be derived for a struct with named fields",
             )
         })?;
-        let extensions = if let Some(attr) = util::get_attr(&field.attrs, "extensions")? {
-            attr.parse_args_with(Punctuated::<LitStr, Token![,]>::parse_terminated)?
-                .into_iter()
-                .map(|s| Extension::new(&s, extensions))
-                .collect::<Result<_>>()?
-        } else {
-            Default::default()
-        };
+        let mut extensions = Vec::new();
+        for attr in util::get_attrs(&field.attrs, "extensions") {
+            for s in attr.parse_args_with(Punctuated::<LitStr, Token![,]>::parse_terminated)? {
+                extensions.push(Extension::new(&s, extension_types)?);
+            }
+        }
         Ok(Self {
             id: util::to_camelcase(&ident),
             field: ident,
diff --git a/derive/src/util.rs b/derive/src/util.rs
index 9d1e8b2cf25..24ba5963b6c 100644
--- a/derive/src/util.rs
+++ b/derive/src/util.rs
@@ -1,8 +1,12 @@
 use quote::ToTokens;
 use syn::{Attribute, Error, Ident, Result};
 
-pub fn get_attr<'a>(attrs: &'a [Attribute], name: &str) -> Result<Option<&'a Attribute>> {
-    let mut attrs = attrs.iter().filter(|attr| attr.path().is_ident(name));
+pub fn get_attrs<'a>(attrs: &'a [Attribute], name: &'a str) -> impl Iterator<Item = &'a Attribute> {
+    attrs.iter().filter(|attr| attr.path().is_ident(name))
+}
+
+pub fn get_attr<'a>(attrs: &'a [Attribute], name: &'a str) -> Result<Option<&'a Attribute>> {
+    let mut attrs = get_attrs(attrs, name);
     let first = attrs.next();
     if let Some(next) = attrs.next() {
         Err(Error::new_spanned(
@@ -17,7 +21,7 @@ pub fn get_attr<'a>(attrs: &'a [Attribute], name: &str) -> Result<Option<&'a Att
 pub fn require_attr<'a>(
     span: &dyn ToTokens,
     attrs: &'a [Attribute],
-    name: &str,
+    name: &'a str,
 ) -> Result<&'a Attribute> {
     get_attr(attrs, name)?
         .ok_or_else(|| Error::new_spanned(span, format!("missing #[{}(...)] attribute", name)))