build: Add inner function to extract struct from messages

Signed-off-by: Patrick José Pereira <[email protected]>
bluerobotics · Apr 20, 2024 · 81dbee0 · 81dbee0
1 parent 93d38ed
commit 81dbee0
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 1 deletion.
diff --git a/build/binder.rs b/build/binder.rs
@@ -46,7 +46,7 @@ pub fn generate<W: Write>(modules: Vec<String>, out: &mut W) {
         }
     }).collect::<Vec<TokenStream>>();
 
-    let deserialize_tokens = modules.into_iter().map(|module| {
+    let deserialize_tokens = modules.clone().into_iter().map(|module| {
         let pascal_case_ident = quote::format_ident!("{}", module.to_case(Case::Pascal));
         let module_ident = quote::format_ident!("{module}");
         quote! {
@@ -56,10 +56,30 @@ pub fn generate<W: Write>(modules: Vec<String>, out: &mut W) {
         }
     }).collect::<Vec<TokenStream>>();
 
+    let enum_tokens_inner = modules
+        .clone()
+        .into_iter()
+        .map(|module| {
+            let pascal_case_ident = quote::format_ident!("{}", module.to_case(Case::Pascal));
+
+            quote! {
+                Self::#pascal_case_ident(inner_enum) => inner_enum.inner(),
+            }
+        })
+        .collect::<Vec<TokenStream>>();
+
     let enum_ident = quote! {
         pub enum Messages {
             #(#enum_tokens)*
         }
+
+        impl Messages {
+            pub fn inner<T: 'static>(&self) -> Option<&T> {
+                match self {
+                    #(#enum_tokens_inner)*
+                }
+            }
+        }
     };
 
     let try_from_ident = quote! {

diff --git a/build/parser.rs b/build/parser.rs
@@ -436,12 +436,29 @@ pub fn generate<R: Read, W: Write>(input: &mut R, output_rust: &mut W) {
             quote!(#pascal_message_name(#pascal_struct_name),)
         })
         .collect::<Vec<TokenStream>>();
+
+    let message_enums_inner = messages
+        .iter()
+        .map(|(name, _message)| {
+            let pascal_message_name = ident!(name.to_case(Case::Pascal));
+            quote!(Self::#pascal_message_name(inner_struct) => (inner_struct as &dyn std::any::Any).downcast_ref::<T>(),)
+        })
+        .collect::<Vec<TokenStream>>();
+
     let message_enums = quote! {
         #[derive(Debug, Clone, PartialEq)]
         #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
         pub enum Messages {
             #(#message_enums)*
         }
+
+        impl Messages {
+            pub fn inner<T: 'static>(&self) -> Option<&T> {
+                match self {
+                    #(#message_enums_inner)*
+                }
+            }
+        }
     };
 
     let message_tokens = messages