From b48a95e87930fa51ef6fb47ad08a95d3159d9bac Mon Sep 17 00:00:00 2001
From: Emil Fresk <emil.fresk@gmail.com>
Date: Fri, 5 Aug 2022 08:59:16 +0200
Subject: [PATCH] Fix codegen when having executor at multiple priorities

The codegen generated code for all executors in all
dispatchers, which caused some weird bugs.
Also the definition of an executor was not generated
globally, this caused use after free errors when having
multiple priority levels.
---
 examples/async-task-multiple-prios.rs | 73 +++++++++++++++++++++++++++
 macros/src/check.rs                   |  6 ++-
 macros/src/codegen/dispatchers.rs     | 47 ++++++++++-------
 macros/src/codegen/idle.rs            |  3 ++
 4 files changed, 109 insertions(+), 20 deletions(-)
 create mode 100644 examples/async-task-multiple-prios.rs

diff --git a/examples/async-task-multiple-prios.rs b/examples/async-task-multiple-prios.rs
new file mode 100644
index 0000000000..a90c11bd98
--- /dev/null
+++ b/examples/async-task-multiple-prios.rs
@@ -0,0 +1,73 @@
+#![no_main]
+#![no_std]
+#![feature(type_alias_impl_trait)]
+
+use panic_semihosting as _;
+
+// NOTES:
+//
+// - Async tasks cannot have `#[lock_free]` resources, as they can interleve and each async
+//   task can have a mutable reference stored.
+// - Spawning an async task equates to it being polled once.
+
+#[rtic::app(device = lm3s6965, dispatchers = [SSI0, UART0], peripherals = true)]
+mod app {
+    use cortex_m_semihosting::{debug, hprintln};
+    use systick_monotonic::*;
+
+    #[shared]
+    struct Shared {}
+
+    #[local]
+    struct Local {}
+
+    #[monotonic(binds = SysTick, default = true)]
+    type MyMono = Systick<100>;
+
+    #[init]
+    fn init(cx: init::Context) -> (Shared, Local, init::Monotonics) {
+        hprintln!("init").unwrap();
+
+        normal_task::spawn().ok();
+        async_task::spawn().ok();
+        normal_task2::spawn().ok();
+        async_task2::spawn().ok();
+
+        (
+            Shared {},
+            Local {},
+            init::Monotonics(Systick::new(cx.core.SYST, 12_000_000)),
+        )
+    }
+
+    #[idle]
+    fn idle(_: idle::Context) -> ! {
+        // debug::exit(debug::EXIT_SUCCESS);
+        loop {
+            // hprintln!("idle");
+            cortex_m::asm::wfi(); // put the MCU in sleep mode until interrupt occurs
+        }
+    }
+
+    #[task(priority = 1)]
+    fn normal_task(_cx: normal_task::Context) {
+        hprintln!("hello from normal 1").ok();
+    }
+
+    #[task(priority = 1)]
+    async fn async_task(_cx: async_task::Context) {
+        hprintln!("hello from async 1").ok();
+
+        debug::exit(debug::EXIT_SUCCESS);
+    }
+
+    #[task(priority = 2)]
+    fn normal_task2(_cx: normal_task2::Context) {
+        hprintln!("hello from normal 2").ok();
+    }
+
+    #[task(priority = 2)]
+    async fn async_task2(_cx: async_task2::Context) {
+        hprintln!("hello from async 2").ok();
+    }
+}
diff --git a/macros/src/check.rs b/macros/src/check.rs
index 374fcedd09..4b360738f5 100644
--- a/macros/src/check.rs
+++ b/macros/src/check.rs
@@ -40,7 +40,11 @@ pub fn app(app: &App, _analysis: &Analysis) -> parse::Result<Extra> {
         })
         .collect::<HashSet<_>>();
 
-    let need = priorities.len();
+    let need = priorities
+        .iter()
+        // Only count if not 0
+        .filter_map(|prio| if *prio > 0 { Some(prio) } else { None })
+        .count();
     let given = app.args.extern_interrupts.len();
     if need > given {
         let s = {
diff --git a/macros/src/codegen/dispatchers.rs b/macros/src/codegen/dispatchers.rs
index 9552451c75..cc9944025e 100644
--- a/macros/src/codegen/dispatchers.rs
+++ b/macros/src/codegen/dispatchers.rs
@@ -10,6 +10,21 @@ pub fn codegen(app: &App, analysis: &Analysis, extra: &Extra) -> Vec<TokenStream
 
     let interrupts = &analysis.interrupts;
 
+    // Generate executor definition in global scope
+    for (name, task) in app.software_tasks.iter() {
+        if task.is_async {
+            let type_name = util::internal_task_ident(name, "F");
+            let exec_name = util::internal_task_ident(name, "EXEC");
+
+            items.push(quote!(
+                type #type_name = impl core::future::Future + 'static;
+                static #exec_name:
+                    rtic::RacyCell<rtic::export::executor::AsyncTaskExecutor<#type_name>> =
+                        rtic::RacyCell::new(rtic::export::executor::AsyncTaskExecutor::new());
+            ));
+        }
+    }
+
     for (&level, channel) in &analysis.channels {
         let mut stmts = vec![];
 
@@ -123,24 +138,17 @@ pub fn codegen(app: &App, analysis: &Analysis, extra: &Extra) -> Vec<TokenStream
             })
             .collect::<Vec<_>>();
 
-        for (name, task) in app.software_tasks.iter() {
-            if task.is_async {
-                let type_name = util::internal_task_ident(name, "F");
-                let exec_name = util::internal_task_ident(name, "EXEC");
-
-                stmts.push(quote!(
-                    type #type_name = impl core::future::Future + 'static;
-                    static #exec_name:
-                        rtic::RacyCell<rtic::export::executor::AsyncTaskExecutor<#type_name>> =
-                            rtic::RacyCell::new(rtic::export::executor::AsyncTaskExecutor::new());
-                ));
-            }
-        }
-
-        let n_executors = app
-            .software_tasks
+        let n_executors = channel
+            .tasks
             .iter()
-            .map(|(_, task)| if task.is_async { 1 } else { 0 })
+            .map(|name| {
+                let task = &app.software_tasks[name];
+                if task.is_async {
+                    1
+                } else {
+                    0
+                }
+            })
             .sum::<usize>()
             .max(1);
 
@@ -165,9 +173,10 @@ pub fn codegen(app: &App, analysis: &Analysis, extra: &Extra) -> Vec<TokenStream
             }
         ));
 
-        for (name, _task) in app.software_tasks.iter().filter_map(|(name, task)| {
+        for name in channel.tasks.iter().filter_map(|name| {
+            let task = &app.software_tasks[name];
             if task.is_async {
-                Some((name, task))
+                Some(name)
             } else {
                 None
             }
diff --git a/macros/src/codegen/idle.rs b/macros/src/codegen/idle.rs
index 83b85d7ba6..ea23f814af 100644
--- a/macros/src/codegen/idle.rs
+++ b/macros/src/codegen/idle.rs
@@ -80,6 +80,9 @@ pub fn codegen(
 
         (mod_app, root_idle, user_idle, call_idle)
     } else {
+        // TODO: No idle defined, check for 0-priority tasks and generate an executor if needed
+
+        //
         (
             vec![],
             vec![],