Make RTIC 2 work on stable by using `main`'s stack as an allocator

1 files changed, 29 insertions, 6 deletions

diff --git a/rtic-macros/src/codegen/main.rs b/rtic-macros/src/codegen/main.rs
index 3848ab0..5612796 100644
--- a/rtic-macros/src/codegen/main.rs
+++ b/rtic-macros/src/codegen/main.rs
@@ -1,9 +1,12 @@
-use crate::{analyze::Analysis, codegen::util, syntax::ast::App};
+use super::{assertions, post_init, pre_init};
+use crate::{
+    analyze::Analysis,
+    codegen::{bindings, util},
+    syntax::ast::App,
+};
 use proc_macro2::TokenStream as TokenStream2;
 use quote::quote;
 
-use super::{assertions, post_init, pre_init};
-
 /// Generates code for `fn main`
 pub fn codegen(app: &App, analysis: &Analysis) -> TokenStream2 {
     let assertion_stmts = assertions::codegen(app, analysis);
@@ -19,12 +22,26 @@ pub fn codegen(app: &App, analysis: &Analysis) -> TokenStream2 {
         let dispatcher = util::zero_prio_dispatcher_ident();
         quote!(#dispatcher();)
     } else {
-        quote!(loop {
-        })
+        quote!(loop {})
     };
 
+    let mut executor_allocations = Vec::new();
+
+    for (name, _) in app.software_tasks.iter() {
+        let exec_name = util::internal_task_ident(name, "EXEC");
+        let new_n_args = util::new_n_args_ident(app.software_tasks[name].inputs.len());
+
+        executor_allocations.push(quote!(
+            let executor = ::core::mem::ManuallyDrop::new(rtic::export::executor::AsyncTaskExecutor::#new_n_args(#name));
+            executors_size += ::core::mem::size_of_val(&executor);
+            #exec_name.set_in_main(&executor);
+        ));
+    }
+
     let main = util::suffixed("main");
     let init_name = &app.init.name;
+    let msp_check = bindings::check_stack_overflow_before_init(app, analysis);
+
     quote!(
         #[doc(hidden)]
         #[no_mangle]
@@ -38,9 +55,15 @@ pub fn codegen(app: &App, analysis: &Analysis) -> TokenStream2 {
                 f();
             }
 
+            // Generate allocations for async executors.
+            let mut executors_size = 0;
+            #(#executor_allocations)*
+
+            #(#msp_check)*
+
             // Wrap late_init_stmts in a function to ensure that stack space is reclaimed.
             __rtic_init_resources(||{
-                let (shared_resources, local_resources) = #init_name(#init_name::Context::new(core.into()));
+                let (shared_resources, local_resources) = #init_name(#init_name::Context::new(core.into(), executors_size));
 
                 #(#post_init_stmts)*
             });