Allow `-> impl Future<Output = ()>` in #[task]

1 files changed, 45 insertions, 2 deletions

diff --git a/embassy-executor/src/lib.rs b/embassy-executor/src/lib.rs
index dfe420bab..70abfcc3a 100644
--- a/embassy-executor/src/lib.rs
+++ b/embassy-executor/src/lib.rs
@@ -65,8 +65,17 @@ pub mod _export {
 
     use crate::raw::TaskPool;
 
+    trait TaskReturnValue {}
+    impl TaskReturnValue for () {}
+    impl TaskReturnValue for Never {}
+
+    #[diagnostic::on_unimplemented(
+        message = "task function futures must resolve to `()`",
+        note = "use `async fn` or change the return type to `impl Future<Output = ()>`"
+    )]
+    #[allow(private_bounds)]
     pub trait TaskFn<Args>: Copy {
-        type Fut: Future + 'static;
+        type Fut: Future<Output: TaskReturnValue> + 'static;
     }
 
     macro_rules! task_fn_impl {
@@ -74,7 +83,7 @@ pub mod _export {
             impl<F, Fut, $($Tn,)*> TaskFn<($($Tn,)*)> for F
             where
                 F: Copy + FnOnce($($Tn,)*) -> Fut,
-                Fut: Future + 'static,
+                Fut: Future<Output: TaskReturnValue> + 'static,
             {
                 type Fut = Fut;
             }
@@ -205,4 +214,38 @@ pub mod _export {
         Align268435456: 268435456,
         Align536870912: 536870912,
     );
+
+    #[allow(dead_code)]
+    trait HasOutput {
+        type Output;
+    }
+
+    impl<O> HasOutput for fn() -> O {
+        type Output = O;
+    }
+
+    #[allow(dead_code)]
+    type Never = <fn() -> ! as HasOutput>::Output;
+}
+
+/// Implementation details for embassy macros.
+/// Do not use. Used for macros and HALs only. Not covered by semver guarantees.
+#[doc(hidden)]
+#[cfg(feature = "nightly")]
+pub mod _export {
+    pub trait TaskReturnValue {}
+    impl TaskReturnValue for () {}
+    impl TaskReturnValue for Never {}
+
+    #[allow(dead_code)]
+    trait HasOutput {
+        type Output;
+    }
+
+    impl<O> HasOutput for fn() -> O {
+        type Output = O;
+    }
+
+    #[allow(dead_code)]
+    type Never = <fn() -> ! as HasOutput>::Output;
 }