Skip to main content

codama_attributes/codama_directives/
enum_discriminator_directive.rs

1use crate::{
2    utils::{FromMeta, SetOnce},
3    Attribute, CodamaAttribute, CodamaDirective,
4};
5use codama_errors::CodamaError;
6use codama_nodes::{
7    CamelCaseString, InstructionArgumentNode, NestedTypeNode, Node, NumberFormat::U8,
8    NumberTypeNode, StructFieldTypeNode, TypeNode,
9};
10use codama_syn_helpers::{extensions::*, Meta};
11
12#[derive(Debug, PartialEq, Default, Clone)]
13pub struct EnumDiscriminatorDirective {
14    pub name: Option<CamelCaseString>,
15    pub size: Option<NestedTypeNode<NumberTypeNode>>,
16}
17
18impl EnumDiscriminatorDirective {
19    pub fn parse(meta: &Meta) -> syn::Result<Self> {
20        let pl = meta
21            .assert_directive("enum_discriminator")?
22            .as_path_list()?;
23
24        let mut name = SetOnce::<CamelCaseString>::new("name");
25        let mut size: SetOnce<NestedTypeNode<NumberTypeNode>> =
26            SetOnce::<NestedTypeNode<NumberTypeNode>>::new("size");
27        pl.each(|ref meta| match meta.path_str().as_str() {
28            "name" => name.set(meta.as_value()?.as_expr()?.as_string()?.into(), meta),
29            "size" => {
30                let node = TypeNode::from_meta(meta.as_value()?)?;
31                match NestedTypeNode::<NumberTypeNode>::try_from(node) {
32                    Ok(node) => size.set(node, meta),
33                    _ => Err(meta.error("size must be a NumberTypeNode")),
34                }
35            }
36            _ => Err(meta.error("unrecognized attribute")),
37        })?;
38
39        let directive = EnumDiscriminatorDirective {
40            name: name.option(),
41            size: size.option(),
42        };
43
44        if directive.name.is_none() && directive.size.is_none() {
45            return Err(meta.error("enum_discriminator must specify at least one of: name, size"));
46        }
47
48        Ok(directive)
49    }
50}
51
52impl<'a> TryFrom<&'a CodamaAttribute<'a>> for &'a EnumDiscriminatorDirective {
53    type Error = CodamaError;
54
55    fn try_from(attribute: &'a CodamaAttribute) -> Result<Self, Self::Error> {
56        match attribute.directive.as_ref() {
57            CodamaDirective::EnumDiscriminator(ref a) => Ok(a),
58            _ => Err(CodamaError::InvalidCodamaDirective {
59                expected: "enum_discriminator".to_string(),
60                actual: attribute.directive.name().to_string(),
61            }),
62        }
63    }
64}
65
66impl<'a> TryFrom<&'a Attribute<'a>> for &'a EnumDiscriminatorDirective {
67    type Error = CodamaError;
68
69    fn try_from(attribute: &'a Attribute) -> Result<Self, Self::Error> {
70        <&CodamaAttribute>::try_from(attribute)?.try_into()
71    }
72}
73
74impl From<&Option<Node>> for EnumDiscriminatorDirective {
75    fn from(node: &Option<Node>) -> Self {
76        let size = match node {
77            Some(Node::DefinedType(defined_type)) => match defined_type.r#type.as_ref() {
78                TypeNode::Enum(data) => Some(data.size.clone()),
79                _ => None,
80            },
81            _ => None,
82        };
83        EnumDiscriminatorDirective {
84            size,
85            ..EnumDiscriminatorDirective::default()
86        }
87    }
88}
89
90impl From<&EnumDiscriminatorDirective> for StructFieldTypeNode {
91    fn from(directive: &EnumDiscriminatorDirective) -> Self {
92        StructFieldTypeNode::new(
93            directive.name.clone().unwrap_or("discriminator".into()),
94            directive
95                .size
96                .clone()
97                .unwrap_or(NumberTypeNode::le(U8).into()),
98        )
99    }
100}
101
102impl From<&EnumDiscriminatorDirective> for InstructionArgumentNode {
103    fn from(directive: &EnumDiscriminatorDirective) -> Self {
104        StructFieldTypeNode::from(directive).into()
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111    use codama_nodes::NumberFormat::{U16, U32};
112
113    #[test]
114    fn enum_discriminator_with_name() {
115        let meta: Meta = syn::parse_quote! { enum_discriminator(name = "banana") };
116        let directive = EnumDiscriminatorDirective::parse(&meta).unwrap();
117        assert_eq!(
118            directive,
119            EnumDiscriminatorDirective {
120                name: Some("banana".into()),
121                size: None,
122            }
123        );
124    }
125
126    #[test]
127    fn enum_discriminator_with_size() {
128        let meta: Meta = syn::parse_quote! { enum_discriminator(size = number(u32)) };
129        let directive = EnumDiscriminatorDirective::parse(&meta).unwrap();
130        assert_eq!(
131            directive,
132            EnumDiscriminatorDirective {
133                name: None,
134                size: Some(NumberTypeNode::le(U32).into()),
135            }
136        );
137    }
138
139    #[test]
140    fn enum_discriminator_with_name_and_size() {
141        let meta: Meta =
142            syn::parse_quote! { enum_discriminator(name = "banana", size = number(u16)) };
143        let directive = EnumDiscriminatorDirective::parse(&meta).unwrap();
144        assert_eq!(
145            directive,
146            EnumDiscriminatorDirective {
147                name: Some("banana".into()),
148                size: Some(NumberTypeNode::le(U16).into()),
149            }
150        );
151    }
152
153    #[test]
154    fn empty_enum_discriminator() {
155        let meta: Meta = syn::parse_quote! { enum_discriminator() };
156        let error = EnumDiscriminatorDirective::parse(&meta).unwrap_err();
157        assert_eq!(
158            error.to_string(),
159            "enum_discriminator must specify at least one of: name, size"
160        );
161    }
162}