accept `init: fn() -> init::LateResources`

3 files changed, 175 insertions, 35 deletions

diff --git a/macros/src/check.rs b/macros/src/check.rs
index 0bc31c5f..045d152f 100644
--- a/macros/src/check.rs
+++ b/macros/src/check.rs
@@ -35,17 +35,20 @@ pub fn app(app: &App) -> parse::Result<()> {
         }
     }
 
-    // Check that all late resources have been initialized in `#[init]`
-    for res in app
-        .resources
-        .iter()
-        .filter_map(|(name, res)| if res.expr.is_none() { Some(name) } else { None })
-    {
-        if app.init.assigns.iter().all(|assign| assign.left != *res) {
-            return Err(parse::Error::new(
-                res.span(),
-                "late resources MUST be initialized at the end of `init`",
-            ));
+    // Check that all late resources have been initialized in `#[init]` if `init` has signature
+    // `fn()`
+    if !app.init.returns_late_resources {
+        for res in app
+            .resources
+            .iter()
+            .filter_map(|(name, res)| if res.expr.is_none() { Some(name) } else { None })
+        {
+            if app.init.assigns.iter().all(|assign| assign.left != *res) {
+                return Err(parse::Error::new(
+                    res.span(),
+                    "late resources MUST be initialized at the end of `init`",
+                ));
+            }
         }
     }
 
@@ -112,11 +115,19 @@ pub fn app(app: &App) -> parse::Result<()> {
         }
     }
 
-    // Check that `init` contains no early returns *if* late resources exist
+    // Check that `init` contains no early returns *if* late resources exist and `init` signature is
+    // `fn()`
     if app.resources.values().any(|res| res.expr.is_none()) {
-        for stmt in &app.init.stmts {
-            noreturn_stmt(stmt)?;
+        if !app.init.returns_late_resources {
+            for stmt in &app.init.stmts {
+                noreturn_stmt(stmt)?;
+            }
         }
+    } else if app.init.returns_late_resources {
+        return Err(parse::Error::new(
+            Span::call_site(),
+            "`init` signature must be `[unsafe] fn()` if there are no late resources",
+        ));
     }
 
     Ok(())
diff --git a/macros/src/codegen.rs b/macros/src/codegen.rs
index 6c1baeca..c7adbd67 100644
--- a/macros/src/codegen.rs
+++ b/macros/src/codegen.rs
@@ -94,7 +94,7 @@ pub fn app(app: &App, analysis: &Analysis) -> TokenStream {
 
     let (dispatchers_data, dispatchers) = dispatchers(&mut ctxt, &app, analysis);
 
-    let init_fn = init(&mut ctxt, &app, analysis);
+    let (init_fn, has_late_resources) = init(&mut ctxt, &app, analysis);
     let init_arg = if cfg!(feature = "timer-queue") {
         quote!(rtfm::Peripherals {
             CBP: p.CBP,
@@ -123,6 +123,30 @@ pub fn app(app: &App, analysis: &Analysis) -> TokenStream {
         })
     };
 
+    let init = &ctxt.init;
+    let init_phase = if has_late_resources {
+        let assigns = app
+            .resources
+            .iter()
+            .filter_map(|(name, res)| {
+                if res.expr.is_none() {
+                    let alias = &ctxt.statics[name];
+
+                    Some(quote!(#alias.set(res.#name);))
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>();
+
+        quote!(
+            let res = #init(#init_arg);
+            #(#assigns)*
+        )
+    } else {
+        quote!(#init(#init_arg);)
+    };
+
     let post_init = post_init(&ctxt, &app, analysis);
 
     let (idle_fn, idle_expr) = idle(&mut ctxt, &app, analysis);
@@ -147,7 +171,6 @@ pub fn app(app: &App, analysis: &Analysis) -> TokenStream {
     let assertions = assertions(app, analysis);
 
     let main = mk_ident(None);
-    let init = &ctxt.init;
     quote!(
         #resources
 
@@ -185,7 +208,7 @@ pub fn app(app: &App, analysis: &Analysis) -> TokenStream {
 
             #pre_init
 
-            #init(#init_arg);
+            #init_phase
 
             #post_init
 
@@ -290,10 +313,11 @@ fn resources(ctxt: &mut Context, app: &App, analysis: &Analysis) -> proc_macro2:
     quote!(#(#items)*)
 }
 
-fn init(ctxt: &mut Context, app: &App, analysis: &Analysis) -> proc_macro2::TokenStream {
+fn init(ctxt: &mut Context, app: &App, analysis: &Analysis) -> (proc_macro2::TokenStream, bool) {
     let attrs = &app.init.attrs;
     let locals = mk_locals(&app.init.statics, true);
     let stmts = &app.init.stmts;
+    // TODO remove in v0.5.x
     let assigns = app
         .init
         .assigns
@@ -334,12 +358,47 @@ fn init(ctxt: &mut Context, app: &App, analysis: &Analysis) -> proc_macro2::Toke
         analysis,
     );
 
+    let (late_resources, late_resources_ident, ret) = if app.init.returns_late_resources {
+        // create `LateResources` struct in the root of the crate
+        let ident = mk_ident(None);
+
+        let fields = app
+            .resources
+            .iter()
+            .filter_map(|(name, res)| {
+                if res.expr.is_none() {
+                    let ty = &res.ty;
+                    Some(quote!(pub #name: #ty))
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>();
+
+        let late_resources = quote!(
+            #[allow(non_snake_case)]
+            pub struct #ident {
+                #(#fields),*
+            }
+        );
+
+        (
+            Some(late_resources),
+            Some(ident),
+            Some(quote!(-> init::LateResources)),
+        )
+    } else {
+        (None, None, None)
+    };
+    let has_late_resources = late_resources.is_some();
+
     let module = module(
         ctxt,
         Kind::Init,
         !app.init.args.schedule.is_empty(),
         !app.init.args.spawn.is_empty(),
         app,
+        late_resources_ident,
     );
 
     #[cfg(feature = "timer-queue")]
@@ -365,25 +424,30 @@ fn init(ctxt: &mut Context, app: &App, analysis: &Analysis) -> proc_macro2::Toke
     let unsafety = &app.init.unsafety;
     let device = &app.args.device;
     let init = &ctxt.init;
-    quote!(
-        #module
+    (
+        quote!(
+            #late_resources
 
-        #(#attrs)*
-        #unsafety fn #init(mut core: rtfm::Peripherals) {
-            #(#locals)*
+            #module
 
-            #baseline_let
+            #(#attrs)*
+            #unsafety fn #init(mut core: rtfm::Peripherals) #ret {
+                #(#locals)*
 
-            #prelude
+                #baseline_let
 
-            let mut device = unsafe { #device::Peripherals::steal() };
+                #prelude
 
-            #start_let
+                let mut device = unsafe { #device::Peripherals::steal() };
 
-            #(#stmts)*
+                #start_let
 
-            #(#assigns)*
-        }
+                #(#stmts)*
+
+                #(#assigns)*
+            }
+        ),
+        has_late_resources,
     )
 }
 
@@ -440,6 +504,7 @@ fn module(
     schedule: bool,
     spawn: bool,
     app: &App,
+    late_resources: Option<Ident>,
 ) -> proc_macro2::TokenStream {
     let mut items = vec![];
     let mut fields = vec![];
@@ -572,6 +637,12 @@ fn module(
         Kind::Task(_) => "Software task",
     };
 
+    if let Some(late_resources) = late_resources {
+        items.push(quote!(
+            pub use super::#late_resources as LateResources;
+        ));
+    }
+
     quote!(
         #root
 
@@ -950,6 +1021,7 @@ fn idle(
             !idle.args.schedule.is_empty(),
             !idle.args.spawn.is_empty(),
             app,
+            None,
         );
 
         let unsafety = &idle.unsafety;
@@ -1004,6 +1076,7 @@ fn exceptions(ctxt: &mut Context, app: &App, analysis: &Analysis) -> Vec<proc_ma
                 !exception.args.schedule.is_empty(),
                 !exception.args.spawn.is_empty(),
                 app,
+                None,
             );
 
             #[cfg(feature = "timer-queue")]
@@ -1083,6 +1156,7 @@ fn interrupts(
             !interrupt.args.schedule.is_empty(),
             !interrupt.args.spawn.is_empty(),
             app,
+            None,
         ));
 
         #[cfg(feature = "timer-queue")]
@@ -1245,6 +1319,7 @@ fn tasks(ctxt: &mut Context, app: &App, analysis: &Analysis) -> proc_macro2::Tok
             !task.args.schedule.is_empty(),
             !task.args.spawn.is_empty(),
             app,
+            None,
         ));
 
         let attrs = &task.attrs;
diff --git a/macros/src/syntax.rs b/macros/src/syntax.rs
index 85f3caaa..ad7d8bde 100644
--- a/macros/src/syntax.rs
+++ b/macros/src/syntax.rs
@@ -596,6 +596,7 @@ impl Parse for InitArgs {
     }
 }
 
+// TODO remove in v0.5.x
 pub struct Assign {
     pub attrs: Vec<Attribute>,
     pub left: Ident,
@@ -608,32 +609,83 @@ pub struct Init {
     pub unsafety: Option<Token![unsafe]>,
     pub statics: HashMap<Ident, Static>,
     pub stmts: Vec<Stmt>,
+    // TODO remove in v0.5.x
     pub assigns: Vec<Assign>,
+    pub returns_late_resources: bool,
 }
 
 impl Init {
     fn check(args: InitArgs, item: ItemFn) -> parse::Result<Self> {
-        let valid_signature = item.vis == Visibility::Inherited
+        let mut valid_signature = item.vis == Visibility::Inherited
             && item.constness.is_none()
             && item.asyncness.is_none()
             && item.abi.is_none()
             && item.decl.generics.params.is_empty()
             && item.decl.generics.where_clause.is_none()
             && item.decl.inputs.is_empty()
-            && item.decl.variadic.is_none()
-            && is_unit(&item.decl.output);
+            && item.decl.variadic.is_none();
+
+        let returns_late_resources = match &item.decl.output {
+            ReturnType::Default => false,
+            ReturnType::Type(_, ty) => {
+                match &**ty {
+                    Type::Tuple(t) => {
+                        if t.elems.is_empty() {
+                            // -> ()
+                            true
+                        } else {
+                            valid_signature = false;
+
+                            false // don't care
+                        }
+                    }
+
+                    Type::Path(p) => {
+                        let mut segments = p.path.segments.iter();
+                        if p.qself.is_none()
+                            && p.path.leading_colon.is_none()
+                            && p.path.segments.len() == 2
+                            && segments.next().map(|s| {
+                                s.arguments == PathArguments::None && s.ident.to_string() == "init"
+                            }) == Some(true)
+                            && segments.next().map(|s| {
+                                s.arguments == PathArguments::None
+                                    && s.ident.to_string() == "LateResources"
+                            }) == Some(true)
+                        {
+                            // -> init::LateResources
+                            true
+                        } else {
+                            valid_signature = false;
+
+                            false // don't care
+                        }
+                    }
+
+                    _ => {
+                        valid_signature = false;
+
+                        false // don't care
+                    }
+                }
+            }
+        };
 
         let span = item.span();
 
         if !valid_signature {
             return Err(parse::Error::new(
                 span,
-                "`init` must have type signature `[unsafe] fn()`",
+                "`init` must have type signature `[unsafe] fn() [-> init::LateResources]`",
             ));
         }
 
         let (statics, stmts) = extract_statics(item.block.stmts);
-        let (stmts, assigns) = extract_assignments(stmts);
+        let (stmts, assigns) = if returns_late_resources {
+            (stmts, vec![])
+        } else {
+            extract_assignments(stmts)
+        };
 
         Ok(Init {
             args,
@@ -642,6 +694,7 @@ impl Init {
             statics: Static::parse(statics)?,
             stmts,
             assigns,
+            returns_late_resources,
         })
     }
 }
@@ -1207,6 +1260,7 @@ fn extract_statics(stmts: Vec<Stmt>) -> (Statics, Vec<Stmt>) {
     (statics, stmts)
 }
 
+// TODO remove in v0.5.x
 fn extract_assignments(stmts: Vec<Stmt>) -> (Vec<Stmt>, Vec<Assign>) {
     let mut istmts = stmts.into_iter().rev();