Parse and generate impls for enums with #[asn(enumerated)] attributes #…

…11
kellerkindt · Apr 28, 2020 · cce4c27 · cce4c27
1 parent 12bdfc4
commit cce4c27
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 21 deletions.
diff --git a/asn1rs-macros/src/ast/mod.rs b/asn1rs-macros/src/ast/mod.rs
@@ -1,7 +1,7 @@
 mod range;
 mod tag;
 
-use asn1rs_model::model::{Definition, Field, Model, Range, Tag, Type};
+use asn1rs_model::model::{Definition, Enumerated, Field, Model, Range, Tag, Type};
 use proc_macro::TokenStream;
 use quote::quote;
 use range::MaybeRanged;
@@ -35,8 +35,14 @@ pub(crate) fn parse(attr: TokenStream, item: TokenStream) -> TokenStream {
 
     let mut additional_impl: Vec<TokenStream2> = Vec::default();
 
+    let mut model: Model<asn1rs_model::model::Asn> = Model {
+        name: "__proc_macro".to_string(),
+        imports: vec![],
+        definitions: vec![],
+    };
+
     let item = match item {
-        Item::Struct(mut strct) => {
+        Item::Struct(mut strct) if asn_type_decl == "sequence" => {
             let mut fields = Vec::new();
             for field in strct.fields.iter_mut() {
                 let mut removed = None;
@@ -82,25 +88,47 @@ pub(crate) fn parse(attr: TokenStream, item: TokenStream) -> TokenStream {
             println!("---------- parsed");
             let definition = Definition(strct.ident.to_string(), Type::Sequence(fields).untagged());
             println!("{:#?}", definition);
-            let model: Model<asn1rs_model::model::Asn> = Model {
-                name: "__proc_macro".to_string(),
-                imports: vec![],
-                definitions: vec![definition],
-            };
-            let model_rust = model.to_rust();
-
-            use asn1rs_model::gen::rust::walker::AsnDefWalker;
-            let stringified = AsnDefWalker::stringify(&model_rust);
-            additional_impl.push(TokenStream2::from_str(&stringified).unwrap());
+            model.definitions.push(definition);
 
             println!("---------- output");
             let st = Item::Struct(strct.clone());
             println!("{}", TokenStream::from(quote! {#st}).to_string());
+
             Item::Struct(strct)
         }
+        Item::Enum(enm) if asn_type_decl == "enumerated" => {
+            let plain_enum = enm.variants.iter().all(|v| v.fields.is_empty());
+            let variants = enm
+                .variants
+                .iter()
+                .map(|v| v.ident.to_string())
+                .collect::<Vec<_>>();
+            if plain_enum {
+                // TODO extensible
+                // TODO tags
+                let enumerated = Enumerated::from_names(variants.into_iter());
+                model.definitions.push(Definition(
+                    enm.ident.to_string(),
+                    Type::Enumerated(enumerated).untagged(),
+                ));
+            } else {
+                // data enum
+                panic!("ENUMERATED does not allow data carried on Variants. Consider type CHOICE");
+            }
+
+            Item::Enum(enm)
+        }
         item => item,
     };
 
+    if !model.definitions.is_empty() {
+        let model_rust = model.to_rust();
+
+        use asn1rs_model::gen::rust::walker::AsnDefWalker;
+        let stringified = AsnDefWalker::stringify(&model_rust);
+        additional_impl.push(TokenStream2::from_str(&stringified).unwrap());
+    }
+
     let result = TokenStream::from(quote! {
         #item
         #(#additional_impl)*

diff --git a/asn1rs-model/src/gen/rust/walker.rs b/asn1rs-model/src/gen/rust/walker.rs
@@ -25,7 +25,12 @@ impl AsnDefWalker {
                     self.write_type_declaration(scope, &name, &field, r#type);
                 }
             }
-            Rust::Enum(_) => {}
+            Rust::Enum(_enm) => {
+                scope.raw(&format!(
+                    "type AsnDef{} = {}Enumerated<{}>;",
+                    name, CRATE_SYN_PREFIX, name
+                ));
+            }
             Rust::DataEnum(_) => {}
             Rust::TupleStruct(_) => {}
         }
@@ -204,7 +209,7 @@ impl AsnDefWalker {
 
     fn write_enumerated_constraint(&self, scope: &mut Scope, name: &str, enumerated: &PlainEnum) {
         let mut imp = Impl::new(name);
-        imp.impl_trait(format!("{}::enumerated::Constraint", CRATE_SYN_PREFIX));
+        imp.impl_trait(format!("{}enumerated::Constraint", CRATE_SYN_PREFIX));
 
         imp.new_fn("to_choice_index")
             .arg_ref_self()
@@ -233,7 +238,7 @@ impl AsnDefWalker {
             scope,
             imp,
             &[
-                format!("const NAME: &str = \"{}\";", name),
+                format!("const NAME: &'static str = \"{}\";", name),
                 format!("const VARIANT_COUNT: usize = {};", enumerated.len()),
                 format!(
                     "const STD_VARIANT_COUNT: usize = {};",

diff --git a/asn1rs-model/src/model/mod.rs b/asn1rs-model/src/model/mod.rs
@@ -819,8 +819,7 @@ pub struct Enumerated {
 }
 
 impl Enumerated {
-    #[cfg(test)]
-    pub(crate) fn from_names<'a, 'b: 'a>(variants: impl Iterator<Item = &'a &'b str>) -> Self {
+    pub fn from_names<I: ToString>(variants: impl Iterator<Item = I>) -> Self {
         Self {
             variants: variants
                 .map(|name| EnumeratedVariant::from_name(name))
@@ -921,8 +920,7 @@ impl<S: ToString> From<S> for EnumeratedVariant {
 }
 
 impl EnumeratedVariant {
-    #[cfg(test)]
-    pub(crate) fn from_name(name: &str) -> Self {
+    pub fn from_name<I: ToString>(name: I) -> Self {
         Self {
             name: name.to_string(),
             number: None,
@@ -1158,8 +1156,7 @@ pub(crate) mod tests {
         assert_eq!(
             Definition(
                 "Neither".into(),
-                Type::Enumerated(Enumerated::from_names(["ABC".into(), "DEF".into()].iter()))
-                    .untagged()
+                Type::Enumerated(Enumerated::from_names(["ABC", "DEF"].iter())).untagged()
             ),
             model.definitions[2]
         );

diff --git a/tests/basic_proc_macro_attribute.rs b/tests/basic_proc_macro_attribute.rs
@@ -68,3 +68,29 @@ fn test_deserialize_with_uper() {
         p
     );
 }
+
+#[asn(enumerated)]
+#[derive(Debug, PartialOrd, PartialEq)]
+pub enum Topping {
+    NotPineapple,
+    EvenLessPineapple,
+    NoPineappleAtAll,
+}
+
+#[test]
+fn topping_test_serialize_with_uper() {
+    let mut uper = UperWriter::default();
+    uper.write(&Topping::NotPineapple).unwrap();
+    uper.write(&Topping::EvenLessPineapple).unwrap();
+    uper.write(&Topping::NoPineappleAtAll).unwrap();
+    assert_eq!(&[0x00 | 0x40 >> 2 | 0x80 >> 4], uper.byte_content());
+    assert_eq!(6, uper.bit_len());
+}
+
+#[test]
+fn topping_test_deserialize_with_uper() {
+    let mut uper = UperReader::from_bits(vec![0x00_u8 | 0x40 >> 2 | 0x80 >> 4], 6);
+    assert_eq!(Topping::NotPineapple, uper.read::<Topping>().unwrap());
+    assert_eq!(Topping::EvenLessPineapple, uper.read::<Topping>().unwrap());
+    assert_eq!(Topping::NoPineappleAtAll, uper.read::<Topping>().unwrap());
+}