From e49434bbb2767f8426dc09d9be5a8e9366e39815 Mon Sep 17 00:00:00 2001
From: Pablo Sichert <mail@pablosichert.com>
Date: Wed, 2 Oct 2024 17:03:26 +0200
Subject: [PATCH] Implement `RustCodeGenerator::add_local_attr` (#86)

---
 asn1rs-model/src/generate/rust.rs | 129 ++++++++++++++++++++++++++++++
 1 file changed, 129 insertions(+)

diff --git a/asn1rs-model/src/generate/rust.rs b/asn1rs-model/src/generate/rust.rs
index fba3492..1769ded 100644
--- a/asn1rs-model/src/generate/rust.rs
+++ b/asn1rs-model/src/generate/rust.rs
@@ -9,6 +9,7 @@ use codegen::Impl;
 use codegen::Scope;
 use codegen::Struct;
 use std::borrow::Cow;
+use std::collections::HashMap;
 use std::convert::Infallible;
 use std::fmt::Display;
 
@@ -36,6 +37,7 @@ pub trait GeneratorSupplement<T> {
 pub struct RustCodeGenerator {
     models: Vec<Model<Rust>>,
     global_derives: Vec<String>,
+    local_attrs: HashMap<String, Vec<String>>,
     direct_field_access: bool,
     getter_and_setter: bool,
 }
@@ -53,6 +55,7 @@ impl Default for RustCodeGenerator {
         RustCodeGenerator {
             models: Default::default(),
             global_derives: Vec::default(),
+            local_attrs: HashMap::new(),
             direct_field_access: true,
             getter_and_setter: false,
         }
@@ -90,6 +93,18 @@ impl RustCodeGenerator {
         self
     }
 
+    pub fn add_local_attr<N: Into<String>, I: Into<String>>(&mut self, name: N, attr: I) {
+        self.local_attrs
+            .entry(name.into())
+            .or_default()
+            .push(attr.into());
+    }
+
+    pub fn without_additional_local_attrs(mut self) -> Self {
+        self.local_attrs.clear();
+        self
+    }
+
     pub const fn fields_are_pub(&self) -> bool {
         self.direct_field_access
     }
@@ -898,6 +913,11 @@ impl RustCodeGenerator {
         self.global_derives.iter().for_each(|derive| {
             str_ct.derive(derive);
         });
+        if let Some(local_attrs) = self.local_attrs.get(name) {
+            local_attrs.iter().for_each(|attr| {
+                str_ct.attr(attr);
+            });
+        }
         str_ct
     }
 
@@ -915,6 +935,11 @@ impl RustCodeGenerator {
         self.global_derives.iter().for_each(|derive| {
             en_m.derive(derive);
         });
+        if let Some(local_attrs) = self.local_attrs.get(name) {
+            local_attrs.iter().for_each(|attr| {
+                en_m.r#macro(&format!("#[{attr}]")); // Workaround for missing `.attr` for enums in codegen
+            });
+        }
         en_m
     }
 }
@@ -1010,4 +1035,108 @@ pub(crate) mod tests {
             &file_content,
         );
     }
+
+    #[test]
+    pub fn test_struct_local_attr() {
+        let model = Model::try_from(Tokenizer::default().parse(
+            r#"Test DEFINITIONS AUTOMATIC TAGS ::=
+            BEGIN
+
+            MyStruct ::= SEQUENCE {
+                myField BOOLEAN
+            }
+
+            END
+        "#,
+        ))
+        .unwrap()
+        .try_resolve()
+        .unwrap()
+        .to_rust();
+
+        let mut generator = RustCodeGenerator::from(model).without_additional_global_derives();
+        generator.add_local_attr("MyStruct", "my_attr");
+        let (_file_name, file_content) = generator
+            .to_string_without_generators()
+            .into_iter()
+            .next()
+            .unwrap();
+
+        assert_starts_with_lines(
+            r#"
+            use asn1rs::prelude::*;
+
+            #[asn(sequence)]
+            #[derive(Default, Debug, Clone, PartialEq, Hash)]
+            #[my_attr]
+            pub struct MyStruct {
+                #[asn(boolean)] pub my_field: bool,
+            }
+
+            impl MyStruct {
+            }
+        "#,
+            &file_content,
+        );
+    }
+
+    #[test]
+    pub fn test_enum_local_attr() {
+        let model = Model::try_from(Tokenizer::default().parse(
+            r#"Test DEFINITIONS AUTOMATIC TAGS ::=
+            BEGIN
+            MyEnum ::= ENUMERATED {
+                a,
+                b
+            }
+            END
+        "#,
+        ))
+        .unwrap()
+        .try_resolve()
+        .unwrap()
+        .to_rust();
+        let mut generator = RustCodeGenerator::from(model).without_additional_global_derives();
+        generator.add_local_attr("MyEnum", "my_attr");
+        let (_file_name, file_content) = generator
+            .to_string_without_generators()
+            .into_iter()
+            .next()
+            .unwrap();
+
+        assert_starts_with_lines(
+            r#"
+            use asn1rs::prelude::*;
+            #[asn(enumerated)]
+            #[derive(Debug, Clone, PartialEq, Hash, Copy, PartialOrd, Eq, Default)]
+            #[my_attr]
+            pub enum MyEnum {
+                #[default] A,
+                B,
+            }
+            impl MyEnum {
+                pub fn variant(index: usize) -> Option<Self> {
+                    match index {
+                        0 => Some(MyEnum::A),
+                        1 => Some(MyEnum::B),
+                        _ => None,
+                    }
+                }
+                pub const fn variants() -> [Self; 2] {
+                    [
+                        MyEnum::A,
+                        MyEnum::B,
+                    ]
+                }
+                pub fn value_index(self) -> usize {
+                    match self {
+                        MyEnum::A => 0,
+                        MyEnum::B => 1,
+                    }
+                }
+            }
+        "#,
+            &file_content,
+        );
+    }
 }