4 files changed, 56 insertions, 26 deletions

diff --git a/macros/src/codegen/idle.rs b/macros/src/codegen/idle.rs
index c0eecbca..1f05d129 100644
--- a/macros/src/codegen/idle.rs
+++ b/macros/src/codegen/idle.rs
@@ -20,8 +20,6 @@ pub fn codegen(
     Vec<TokenStream2>,
     // user_idle
     Option<TokenStream2>,
-    // call_idle
-    TokenStream2,
 ) {
     if let Some(idle) = &app.idle {
         let mut mod_app = vec![];
@@ -60,22 +58,10 @@ pub fn codegen(
             }
         ));
 
-        let call_idle = quote!(#name(
-            #name::Context::new()
-        ));
-
-        (mod_app, root_idle, user_idle, call_idle)
+        (mod_app, root_idle, user_idle)
     } else {
         // TODO: No idle defined, check for 0-priority tasks and generate an executor if needed
-        unimplemented!();
 
-        // (
-        //     vec![],
-        //     vec![],
-        //     None,
-        //     quote!(loop {
-        //         rtic::export::nop()
-        //     }),
-        // )
+        (vec![], vec![], None)
     }
 }
diff --git a/macros/src/codegen/init.rs b/macros/src/codegen/init.rs
index 2aa8fb31..3b2bcd47 100644
--- a/macros/src/codegen/init.rs
+++ b/macros/src/codegen/init.rs
@@ -18,8 +18,6 @@ type CodegenResult = (
     Vec<TokenStream2>,
     // user_init -- the `#[init]` function written by the user
     TokenStream2,
-    // call_init -- the call to the user `#[init]`
-    TokenStream2,
 );
 
 /// Generates support code for `#[init]` functions
@@ -63,6 +61,7 @@ pub fn codegen(app: &App, analysis: &Analysis) -> CodegenResult {
             )
         })
         .collect();
+
     root_init.push(quote! {
         struct #shared {
             #(#shared_resources)*
@@ -97,11 +96,7 @@ pub fn codegen(app: &App, analysis: &Analysis) -> CodegenResult {
         mod_app = Some(constructor);
     }
 
-    let call_init = quote! {
-        let (shared_resources, local_resources) = #name(#name::Context::new(core.into()));
-    };
-
     root_init.push(module::codegen(Context::Init, app, analysis));
 
-    (mod_app, root_init, user_init, call_init)
+    (mod_app, root_init, user_init)
 }
diff --git a/macros/src/codegen/main.rs b/macros/src/codegen/main.rs
new file mode 100644
index 00000000..90f09ae0
--- /dev/null
+++ b/macros/src/codegen/main.rs
@@ -0,0 +1,51 @@
+use crate::{analyze::Analysis, codegen::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);
+
+    let pre_init_stmts = pre_init::codegen(app, analysis);
+
+    let post_init_stmts = post_init::codegen(app, analysis);
+
+    let call_idle = if let Some(idle) = &app.idle {
+        let name = &idle.name;
+        quote!(#name(#name::Context::new()))
+    } else {
+        // TODO: No idle defined, check for 0-priority tasks and generate an executor if needed
+
+        quote!(loop {
+            rtic::export::nop()
+        })
+    };
+
+    let main = util::suffixed("main");
+    let init_name = &app.init.name;
+    quote!(
+        #[doc(hidden)]
+        #[no_mangle]
+        unsafe extern "C" fn #main() -> ! {
+            #(#assertion_stmts)*
+
+            #(#pre_init_stmts)*
+
+            #[inline(never)]
+            fn __rtic_init_resources<F>(f: F) where F: FnOnce() {
+                f();
+            }
+
+            // 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()));
+
+                #(#post_init_stmts)*
+            });
+
+            #call_idle
+        }
+    )
+}
diff --git a/macros/src/codegen/post_init.rs b/macros/src/codegen/post_init.rs
index e8183b93..c4e53837 100644
--- a/macros/src/codegen/post_init.rs
+++ b/macros/src/codegen/post_init.rs
@@ -1,9 +1,7 @@
-use crate::syntax::ast::App;
+use crate::{analyze::Analysis, codegen::util, syntax::ast::App};
 use proc_macro2::TokenStream as TokenStream2;
 use quote::quote;
 
-use crate::{analyze::Analysis, codegen::util};
-
 /// Generates code that runs after `#[init]` returns
 pub fn codegen(app: &App, analysis: &Analysis) -> Vec<TokenStream2> {
     let mut stmts = vec![];