1 files changed, 53 insertions, 49 deletions

diff --git a/macros/src/codegen/init.rs b/macros/src/codegen/init.rs
index fa273fee..8942439b 100644
--- a/macros/src/codegen/init.rs
+++ b/macros/src/codegen/init.rs
@@ -1,5 +1,5 @@
 use proc_macro2::TokenStream as TokenStream2;
-use quote::quote;
+use quote::{format_ident, quote};
 use rtic_syntax::{ast::App, Context};
 
 use crate::{
@@ -10,12 +10,11 @@ use crate::{
 
 /// Generates support code for `#[init]` functions
 pub fn codegen(
-    core: u8,
     app: &App,
     analysis: &Analysis,
     extra: &Extra,
 ) -> (
-    // const_app_idle -- the `${init}Resources` constructor
+    // mod_app_idle -- the `${init}Resources` constructor
     Option<TokenStream2>,
     // root_init -- items that must be placed in the root of the crate:
     // - the `${init}Locals` struct
@@ -25,58 +24,55 @@ pub fn codegen(
     Vec<TokenStream2>,
     // user_init -- the `#[init]` function written by the user
     Option<TokenStream2>,
+    // user_init_imports -- the imports for `#[init]` functio written by the user
+    Vec<TokenStream2>,
     // call_init -- the call to the user `#[init]` if there's one
     Option<TokenStream2>,
 ) {
-    if let Some(init) = app.inits.get(&core) {
-        let cfg_core = util::cfg_core(core, app.args.cores);
+    if app.inits.len() > 0 {
+        let init = &app.inits.first().unwrap();
         let mut needs_lt = false;
         let name = &init.name;
 
         let mut root_init = vec![];
 
-        let ret = {
-            let late_fields = analysis
-                .late_resources
-                .get(&core)
-                .map(|resources| {
-                    resources
-                        .iter()
-                        .map(|name| {
-                            let ty = &app.late_resources[name].ty;
-                            let cfgs = &app.late_resources[name].cfgs;
-
-                            quote!(
-                            #(#cfgs)*
-                            pub #name: #ty
-                            )
-                        })
-                        .collect::<Vec<_>>()
+        let late_fields = analysis
+            .late_resources
+            .iter()
+            .flat_map(|resources| {
+                resources.iter().map(|name| {
+                    let ty = &app.late_resources[name].ty;
+                    let cfgs = &app.late_resources[name].cfgs;
+
+                    quote!(
+                        #(#cfgs)*
+                        pub #name: #ty
+                    )
                 })
-                .unwrap_or(vec![]);
+            })
+            .collect::<Vec<_>>();
 
-            if !late_fields.is_empty() {
-                let late_resources = util::late_resources_ident(&name);
+        let mut user_init_imports = vec![];
+        let late_resources = util::late_resources_ident(&name);
 
-                root_init.push(quote!(
-                    /// Resources initialized at runtime
-                    #cfg_core
-                    #[allow(non_snake_case)]
-                    pub struct #late_resources {
-                        #(#late_fields),*
-                    }
-                ));
-
-                Some(quote!(-> #name::LateResources))
-            } else {
-                None
+        root_init.push(quote!(
+            /// Resources initialized at runtime
+            #[allow(non_snake_case)]
+            pub struct #late_resources {
+                #(#late_fields),*
             }
-        };
+        ));
+
+        let name_late = format_ident!("{}LateResources", name);
+        user_init_imports.push(quote!(
+                #[allow(non_snake_case)]
+                use super::#name_late;
+        ));
 
         let mut locals_pat = None;
         let mut locals_new = None;
         if !init.locals.is_empty() {
-            let (struct_, pat) = locals::codegen(Context::Init(core), &init.locals, core, app);
+            let (struct_, pat) = locals::codegen(Context::Init, &init.locals, app);
 
             locals_new = Some(quote!(#name::Locals::new()));
             locals_pat = Some(pat);
@@ -86,25 +82,33 @@ pub fn codegen(
         let context = &init.context;
         let attrs = &init.attrs;
         let stmts = &init.stmts;
-        let section = util::link_section("text", core);
         let locals_pat = locals_pat.iter();
         let user_init = Some(quote!(
             #(#attrs)*
-            #cfg_core
             #[allow(non_snake_case)]
-            #section
-            fn #name(#(#locals_pat,)* #context: #name::Context) #ret {
+            fn #name(#(#locals_pat,)* #context: #name::Context) -> #name::LateResources {
                 #(#stmts)*
             }
         ));
+        user_init_imports.push(quote!(
+                #(#attrs)*
+                #[allow(non_snake_case)]
+                use super::#name;
+        ));
 
-        let mut const_app = None;
+        let mut mod_app = None;
         if !init.args.resources.is_empty() {
             let (item, constructor) =
-                resources_struct::codegen(Context::Init(core), 0, &mut needs_lt, app, analysis);
+                resources_struct::codegen(Context::Init, 0, &mut needs_lt, app, analysis);
 
             root_init.push(item);
-            const_app = Some(constructor);
+            mod_app = Some(constructor);
+
+            let name_late = format_ident!("{}Resources", name);
+            user_init_imports.push(quote!(
+                    #[allow(non_snake_case)]
+                    use super::#name_late;
+            ));
         }
 
         let locals_new = locals_new.iter();
@@ -112,10 +116,10 @@ pub fn codegen(
             quote!(let late = crate::#name(#(#locals_new,)* #name::Context::new(core.into()));),
         );
 
-        root_init.push(module::codegen(Context::Init(core), needs_lt, app, extra));
+        root_init.push(module::codegen(Context::Init, needs_lt, app, extra));
 
-        (const_app, root_init, user_init, call_init)
+        (mod_app, root_init, user_init, user_init_imports, call_init)
     } else {
-        (None, vec![], None, None)
+        (None, vec![], None, vec![], None)
     }
 }