fix: fallible systems using extra arguments

ten3roberts · Nov 30, 2024 · af418c9 · af418c9
1 parent 4ffc4d9
commit af418c9
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 7 deletions.
diff --git a/flax-derive/src/system.rs b/flax-derive/src/system.rs
@@ -150,17 +150,27 @@ pub(crate) fn system_impl(
                 .map(|i| format_ident!("__extra_arg_{i}"))
                 .collect_vec();
 
-            let call = if *ret == ReturnType::Default {
-                quote! { #call_sig(#(#query_idents),* #(,#with_adapters #item_names)* ) }
+            let (call_expr, ret_sig, ret) = if *ret == ReturnType::Default {
+                (
+                    quote! { #call_sig(#(#query_idents),* #(,#with_adapters #item_names)* ) },
+                    quote! { () },
+                    quote! {},
+                )
             } else {
-                quote! { #call_sig(#(#query_idents),* #(,#with_adapters #item_names)* )?; }
+                (
+                    quote! { #call_sig(#(#query_idents),* #(,#with_adapters #item_names)* )?; },
+                    quote! { #crate_name::__internal::anyhow::Result<()> },
+                    quote! { Ok(()) },
+                )
             };
 
             quote! {
-                build(|#(mut #item_names: #items_types,)* mut main_query: #crate_name::QueryBorrow<'_, _, _>| {
+                build(|#(mut #item_names: #items_types,)* mut main_query: #crate_name::QueryBorrow<'_, _, _>| -> #ret_sig {
                     for (#(#query_idents,)*) in &mut main_query {
-                        #call
+                        #call_expr
                     }
+
+                    #ret
                 })
             }
         }

diff --git a/src/lib.rs b/src/lib.rs
@@ -284,3 +284,8 @@ pub(crate) use vtable::ComponentVTable;
 #[doc(inline)]
 #[cfg(feature = "derive")]
 pub use flax_derive::*;
+
+#[doc(hidden)]
+pub mod __internal {
+    pub use anyhow;
+}
diff --git a/tests/system_macro.rs b/tests/system_macro.rs
@@ -47,9 +47,12 @@ fn system_macro() {
         eprintln!("{a} {b} {c_renamed} {d:?}");
     }
 
-    #[system(par)]
-    fn fallible(a: &mut i32) -> anyhow::Result<()> {
+    #[system(filter(a().with()), with_cmd_mut)]
+    fn fallible(a: &mut i32, _cmd: &mut CommandBuffer) -> anyhow::Result<()> {
+        (anyhow::Ok(()))?;
+
         let _ = a;
+
         Ok(())
     }