Add support for declarative CHOICE types #11

asn1rs-macros/src/ast/mod.rs

@@ -1,7 +1,9 @@
 mod range;
 mod tag;
 
-use asn1rs_model::model::{Definition, Enumerated, Field, Model, Range, Tag, Type};
+use asn1rs_model::model::{
+    Choice, ChoiceVariant, Definition, Enumerated, Field, Model, Range, Tag, Type,
+};
 use proc_macro::TokenStream;
 use quote::quote;
 use range::MaybeRanged;
@@ -63,29 +65,18 @@ pub(crate) fn parse(attr: TokenStream, item: TokenStream) -> TokenStream {
                 if let Some(removed) = removed {
                     match removed.parse_args::<Asn>() {
                         Ok(asn) => {
-                            let parsed = asn1rs_model::model::Asn {
-                                tag: asn.tag,
-                                r#type: match asn.r#type {
-                                    Some(some) => {
-                                        if let Type::TypeReference(_) = some {
-                                            let ty = field.ty.clone();
-                                            Type::TypeReference(quote! { #ty }.to_string())
-                                        } else {
-                                            some
-                                        }
-                                    }
+                            fields.push(Field {
+                                name: field.ident.as_ref().map(ToString::to_string).unwrap(),
+                                optional: asn.optional,
+                                role: match into_asn(&field.ty, asn) {
+                                    Some(asn) => asn,
                                     None => {
                                         return TokenStream::from(
                                             syn::Error::new(field.span(), "Missing ASN-Type")
                                                 .to_compile_error(),
                                         );
                                     }
                                 },
-                            };
-                            fields.push(Field {
-                                name: field.ident.as_ref().map(ToString::to_string).unwrap(),
-                                role: parsed,
-                                optional: asn.optional,
                             });
                         }
                         Err(e) => return TokenStream::from(e.to_compile_error()),
@@ -125,6 +116,74 @@ pub(crate) fn parse(attr: TokenStream, item: TokenStream) -> TokenStream {
 
             Item::Enum(enm)
         }
+        Item::Enum(mut enm) if asn_type_decl == "choice" => {
+            let data_enum = enm.variants.iter().all(|v| !v.fields.is_empty());
+            let variants = enm
+                .variants
+                .iter_mut()
+                .map(|v| {
+                    if v.fields.len() != 1 {
+                        panic!("Variants of CHOICE have to have exactly one field");
+                    }
+                    let mut attr = None;
+                    'inner: for i in 0..v.attrs.len() {
+                        if v.attrs[i]
+                            .path
+                            .segments
+                            .first()
+                            .unwrap()
+                            .ident
+                            .to_string()
+                            .eq("asn")
+                        {
+                            attr = Some(v.attrs.remove(i));
+                            break 'inner;
+                        }
+                    }
+                    let attr = attr.expect("Missing #[asn(..)] attribute");
+
+                    match attr.parse_args::<Asn>() {
+                        Ok(asn) => match into_asn(&v.fields.iter().next().unwrap().ty, asn) {
+                            Some(asn) => {
+                                let name = v.ident.to_string();
+                                Ok(ChoiceVariant {
+                                    name,
+                                    tag: asn.tag,
+                                    r#type: asn.r#type,
+                                })
+                            }
+                            None => Err(TokenStream::from(
+                                syn::Error::new(v.span(), "Missing ASN-Type").to_compile_error(),
+                            )),
+                        },
+                        Err(e) => Err(TokenStream::from(e.to_compile_error())),
+                    }
+                })
+                .collect::<Vec<_>>();
+
+            if data_enum {
+                // TODO extensible
+                // TODO tags
+                let choice = Choice::from_variants({
+                    let mut new = Vec::with_capacity(variants.len());
+                    for var in variants {
+                        new.push(match var {
+                            Ok(variant) => variant,
+                            Err(e) => return e,
+                        });
+                    }
+                    new
+                });
+                model.definitions.push(Definition(
+                    enm.ident.to_string(),
+                    Type::Choice(choice).untagged(),
+                ));
+            } else {
+                // mixed case
+                panic!("CHOICE does not allow any Variant to not have data attached!");
+            }
+            Item::Enum(enm)
+        }
         item => item,
     };
 
@@ -146,6 +205,22 @@ pub(crate) fn parse(attr: TokenStream, item: TokenStream) -> TokenStream {
     result
 }
 
+fn into_asn(ty: &syn::Type, asn: Asn) -> Option<asn1rs_model::model::Asn> {
+    Some(asn1rs_model::model::Asn {
+        tag: asn.tag,
+        r#type: match asn.r#type {
+            Some(some) => {
+                if let Type::TypeReference(_) = some {
+                    Type::TypeReference(quote! { #ty }.to_string())
+                } else {
+                    some
+                }
+            }
+            None => return None,
+        },
+    })
+}
+
 #[derive(Debug, Default)]
 struct Asn {
     r#type: Option<Type>,

asn1rs-model/src/gen/rust/walker.rs

@@ -1,6 +1,6 @@
 use crate::gen::rust::GeneratorSupplement;
 use crate::gen::RustCodeGenerator;
-use crate::model::rust::PlainEnum;
+use crate::model::rust::{DataEnum, PlainEnum};
 use crate::model::{Definition, Model, Range, Rust, RustType};
 use codegen::{Block, Impl, Scope};
 use std::fmt::Display;
@@ -31,8 +31,16 @@ impl AsnDefWalker {
                     name, CRATE_SYN_PREFIX, name
                 ));
             }
-            Rust::DataEnum(_) => {}
-            Rust::TupleStruct(_) => {}
+            Rust::DataEnum(enm) => {
+                scope.raw(&format!(
+                    "type AsnDef{} = {}Choice<{}>;",
+                    name, CRATE_SYN_PREFIX, name
+                ));
+                for (field, r#type) in enm.variants() {
+                    self.write_type_declaration(scope, &name, &field, r#type);
+                }
+            }
+            Rust::TupleStruct(_) => unimplemented!("TupleStruct in Walker::write_type_definitions"),
         }
     }
 
@@ -87,8 +95,8 @@ impl AsnDefWalker {
             Rust::Enum(plain) => {
                 self.write_enumerated_constraint(scope, &name, plain);
             }
-            Rust::DataEnum(_) => {}
-            Rust::TupleStruct(_) => {}
+            Rust::DataEnum(data) => self.write_choice_constraint(scope, &name, data),
+            Rust::TupleStruct(_) => unimplemented!("TupleStruct for Walker::write_constraints"),
         }
     }
 
@@ -251,6 +259,71 @@ impl AsnDefWalker {
         );
     }
 
+    fn write_choice_constraint(&self, scope: &mut Scope, name: &str, choice: &DataEnum) {
+        let mut imp = Impl::new(name);
+        imp.impl_trait(format!("{}choice::Constraint", CRATE_SYN_PREFIX));
+
+        imp.new_fn("to_choice_index")
+            .arg_ref_self()
+            .ret("usize")
+            .push_block({
+                let mut match_block = Block::new("match self");
+                for (index, (variant, _type)) in choice.variants().enumerate() {
+                    match_block.line(format!("Self::{}(_) => {},", variant, index));
+                }
+                match_block
+            });
+
+        imp.new_fn("write_content")
+            .generic(&format!("W: {}Writer", CRATE_SYN_PREFIX))
+            .arg_ref_self()
+            .arg("writer", "&mut W")
+            .ret("Result<(), W::Error>")
+            .push_block({
+                let mut match_block = Block::new("match self");
+                for (variant, _type) in choice.variants() {
+                    let combined = Self::combined_field_type_name(name, variant);
+                    match_block.line(format!(
+                        "Self::{}(c) => AsnDef{}::write_value(writer, c),",
+                        variant, combined
+                    ));
+                }
+                match_block
+            });
+
+        imp.new_fn("read_content")
+            .generic(&format!("R: {}Reader", CRATE_SYN_PREFIX))
+            .arg("index", "usize")
+            .arg("reader", "&mut R")
+            .ret("Result<Option<Self>, R::Error>")
+            .push_block({
+                let mut match_block = Block::new("match index");
+                for (index, (variant, _type)) in choice.variants().enumerate() {
+                    let combined = Self::combined_field_type_name(name, variant);
+                    match_block.line(format!(
+                        "{} => Ok(Some(Self::{}(AsnDef{}::read_value(reader)?))),",
+                        index, variant, combined
+                    ));
+                }
+                match_block.line("_ => Ok(None),");
+                match_block
+            });
+
+        Self::insert_consts(
+            scope,
+            imp,
+            &[
+                format!("const NAME: &'static str = \"{}\";", name),
+                format!("const VARIANT_COUNT: usize = {};", choice.len()),
+                format!(
+                    "const STD_VARIANT_COUNT: usize = {};",
+                    choice.last_standard_index().unwrap_or_else(|| choice.len())
+                ),
+                format!("const EXTENSIBLE: bool = {};", choice.is_extensible()),
+            ],
+        );
+    }
+
     fn write_integer_constraint_type<T: Display>(
         scope: &mut Scope,
         name: &str,

asn1rs-model/src/model/mod.rs

@@ -695,7 +695,6 @@ pub struct Choice {
 }
 
 impl Choice {
-    #[cfg(test)]
     pub fn from_variants(variants: Vec<ChoiceVariant>) -> Self {
         Self {
             variants,
@@ -774,9 +773,9 @@ impl TryFrom<&mut Peekable<IntoIter<Token>>> for Choice {
 
 #[derive(Debug, Clone, PartialOrd, PartialEq)]
 pub struct ChoiceVariant {
-    pub(crate) name: String,
-    pub(crate) tag: Option<Tag>,
-    pub(crate) r#type: Type,
+    pub name: String,
+    pub tag: Option<Tag>,
+    pub r#type: Type,
 }
 
 impl ChoiceVariant {

src/io/uper.rs

@@ -96,6 +96,7 @@ pub trait Reader {
         }
     }
 
+    /// Range is inclusive
     fn read_int(&mut self, range: (i64, i64)) -> Result<i64, Error> {
         let (lower, upper) = range;
         let leading_zeros = ((upper - lower) as u64).leading_zeros();
@@ -233,6 +234,7 @@ pub trait Writer {
         }
     }
 
+    /// Range is inclusive
     fn write_int(&mut self, value: i64, range: (i64, i64)) -> Result<(), Error> {
         let (lower, upper) = range;
         let value = {

src/syn/choice.rs

@@ -0,0 +1,44 @@
+use crate::syn::{ReadableType, Reader, WritableType, Writer};
+use core::marker::PhantomData;
+
+pub struct Choice<C: Constraint>(PhantomData<C>);
+
+impl<C: Constraint> Default for Choice<C> {
+    fn default() -> Self {
+        Self(Default::default())
+    }
+}
+
+pub trait Constraint: Sized {
+    const NAME: &'static str;
+    const VARIANT_COUNT: usize;
+    const STD_VARIANT_COUNT: usize;
+    const EXTENSIBLE: bool = false;
+
+    fn to_choice_index(&self) -> usize;
+
+    fn write_content<W: Writer>(&self, writer: &mut W) -> Result<(), W::Error>;
+
+    fn read_content<R: Reader>(index: usize, reader: &mut R) -> Result<Option<Self>, R::Error>;
+}
+
+impl<C: Constraint> WritableType for Choice<C> {
+    type Type = C;
+
+    #[inline]
+    fn write_value<W: Writer>(
+        writer: &mut W,
+        value: &Self::Type,
+    ) -> Result<(), <W as Writer>::Error> {
+        writer.write_choice(value)
+    }
+}
+
+impl<C: Constraint> ReadableType for Choice<C> {
+    type Type = C;
+
+    #[inline]
+    fn read_value<R: Reader>(reader: &mut R) -> Result<Self::Type, <R as Reader>::Error> {
+        reader.read_choice::<Self::Type>()
+    }
+}

src/syn/io/println.rs

@@ -50,6 +50,26 @@ impl Writer for PrintlnWriter {
         })
     }
 
+    fn write_choice<C: choice::Constraint>(&mut self, choice: &C) -> Result<(), Self::Error> {
+        self.indented_println(&format!("Write choice {}", C::NAME));
+        self.with_increased_indentation(|w| {
+            if C::EXTENSIBLE {
+                w.indented_println("extensible");
+            } else {
+                w.indented_println("normal");
+            }
+            w.with_increased_indentation(|w| {
+                w.indented_println(&format!(
+                    "choice_index {}/{}/{}",
+                    choice.to_choice_index(),
+                    C::STD_VARIANT_COUNT,
+                    C::VARIANT_COUNT
+                ));
+                choice.write_content(w)
+            })
+        })
+    }
+
     fn write_opt<T: WritableType>(&mut self, value: Option<&T::Type>) -> Result<(), Self::Error> {
         self.indented_println("Writing OPTIONAL");
         self.with_increased_indentation(|w| {