Make StackFuture::from fail statically if the future doesn't fit

Cargo.toml

@@ -4,7 +4,7 @@
 [package]
 name = "stackfuture"
 description = "StackFuture is a wrapper around futures that stores the wrapped future in space provided by the caller."
-version = "0.2.0"
+version = "0.3.0"
 edition = "2021"
 license = "MIT"
 authors = ["Microsoft"]

src/lib.rs

@@ -72,45 +72,6 @@ use alloc::boxed::Box;
 /// it also respects any alignment requirements for the wrapped future. Note that the
 /// wrapped future's alignment must be less than or equal to that of the overall
 /// `StackFuture` struct.
-///
-/// The following example would panic because the future is too large:
-/// ```should_panic
-/// # use stackfuture::*;
-/// fn large_stack_future() -> StackFuture<'static, u8, { 4 }> {
-///     StackFuture::from(async {
-///         let mut buf = [0u8; 256];
-///         fill_buf(&mut buf).await;
-///         buf[0]
-///     })
-/// }
-///
-/// async fn fill_buf(buf: &mut [u8]) {
-///     buf[0] = 42;
-/// }
-///
-/// let future = large_stack_future();
-/// ```
-///
-/// The following example would panic because the alignment requirements can't be met:
-/// ```should_panic
-/// # use stackfuture::*;
-/// #[repr(align(4096))]
-/// struct LargeAlignment(i32);
-///
-/// fn large_alignment_future() -> StackFuture<'static, i32, { 8192 }> {
-///     StackFuture::from(async {
-///         let mut buf = LargeAlignment(0);
-///         fill_buf(&mut buf).await;
-///         buf.0
-///     })
-/// }
-///
-/// async fn fill_buf(buf: &mut LargeAlignment) {
-///     buf.0 = 42;
-/// }
-///
-/// let future = large_alignment_future();
-/// ```
 #[repr(C)] // Ensures the data first does not have any padding before it in the struct
 pub struct StackFuture<'a, T, const STACK_SIZE: usize> {
     /// An array of bytes that is used to store the wrapped future.
@@ -135,10 +96,40 @@ impl<'a, T, const STACK_SIZE: usize> StackFuture<'a, T, { STACK_SIZE }> {
     ///
     /// See the documentation on [`StackFuture`] for examples of how to use this.
     ///
-    /// # Panics
+    /// The size and alignment requirements are statically checked, so it is a compiler error
+    /// to use this with a future that does not fit within the [`StackFuture`]'s size and
+    /// alignment requirements.
+    ///
+    /// The following example illustrates a compile error for a future that is too large.
+    /// ```compile_fail
+    /// # use stackfuture::StackFuture;
+    /// // Fails because the future contains a large array and is therefore too big to fit in
+    /// // a 16-byte `StackFuture`.
+    /// let f = StackFuture::<_, { 16 }>::from(async {
+    ///     let x = [0u8; 4096];
+    ///     async {}.await;
+    ///     println!("{}", x.len());
+    /// });
+    /// # #[cfg(miri)] break rust; // FIXME: miri doesn't detect this breakage for some reason...
+    /// ```
     ///
-    /// Panics if the requested `StackFuture` is not large enough to hold `future` or we cannot
-    /// satisfy the alignment requirements for `F`.
+    /// The example below illustrates a compiler error for a future whose alignment is too large.
+    /// ```compile_fail
+    /// # use stackfuture::StackFuture;
+    ///
+    /// #[derive(Debug)]
+    /// #[repr(align(256))]
+    /// struct BigAlignment(usize);
+    ///
+    /// // Fails because the future contains a large array and is therefore too big to fit in
+    /// // a 16-byte `StackFuture`.
+    /// let f = StackFuture::<_, { 16 }>::from(async {
+    ///     let x = BigAlignment(42);
+    ///     async {}.await;
+    ///     println!("{x:?}");
+    /// });
+    /// # #[cfg(miri)] break rust; // FIXME: miri doesn't detect this breakage for some reason...
+    /// ```
     pub fn from<F>(future: F) -> Self
     where
         F: Future<Output = T> + Send + 'a, // the bounds here should match those in the _phantom field
@@ -152,6 +143,15 @@ impl<'a, T, const STACK_SIZE: usize> StackFuture<'a, T, { STACK_SIZE }> {
         // However, libcore provides a blanket `impl<T> From<T> for T`, and since `StackFuture: Future`,
         // both impls end up being applicable to do `From<StackFuture> for StackFuture`.
 
+        // Statically assert that `F` meets all the size and alignment requirements
+        #[allow(clippy::let_unit_value)]
+        let _ = AssertFits::<F, STACK_SIZE>::ASSERT;
+
+        // Since we have the static assert above, we know this will not fail. This means we could
+        // use `unwrap_unchecked` here. The extra `unsafe` code probably isn't worth it since very
+        // likely the compiler will be able to make that optimization for us, this comment is here
+        // as a reminder that we can change it if for some reason this ends up being a performance
+        // issue in the future.
         Self::try_from(future).unwrap_or_else(|f| {
             match (Self::has_alignment_for_val(&f), Self::has_space_for_val(&f)) {
                 (false, false) => panic!(
@@ -274,29 +274,29 @@ impl<'a, T, const STACK_SIZE: usize> StackFuture<'a, T, { STACK_SIZE }> {
     }
 
     /// Computes how much space is required to store a value of type `F`
-    fn required_space<F>() -> usize {
+    const fn required_space<F>() -> usize {
         mem::size_of::<F>()
     }
 
     /// Determines whether this `StackFuture` can hold a value of type `F`
-    pub fn has_space_for<F>() -> bool {
+    pub const fn has_space_for<F>() -> bool {
         Self::required_space::<F>() <= STACK_SIZE
     }
 
     /// Determines whether this `StackFuture` can hold the referenced value
-    pub fn has_space_for_val<F>(_: &F) -> bool {
+    pub const fn has_space_for_val<F>(_: &F) -> bool {
         Self::has_space_for::<F>()
     }
 
     /// Determines whether this `StackFuture`'s alignment is compatible with the
     /// type `F`.
-    pub fn has_alignment_for<F>() -> bool {
+    pub const fn has_alignment_for<F>() -> bool {
         mem::align_of::<F>() <= mem::align_of::<Self>()
     }
 
     /// Determines whether this `StackFuture`'s alignment is compatible with the
     /// referenced value.
-    pub fn has_alignment_for_val<F>(_: &F) -> bool {
+    pub const fn has_alignment_for_val<F>(_: &F) -> bool {
         Self::has_alignment_for::<F>()
     }
 }
@@ -324,5 +324,19 @@ impl<'a, T, const STACK_SIZE: usize> Drop for StackFuture<'a, T, { STACK_SIZE }>
     }
 }
 
+struct AssertFits<F, const STACK_SIZE: usize>(PhantomData<F>);
+
+impl<F, const STACK_SIZE: usize> AssertFits<F, STACK_SIZE> {
+    const ASSERT: () = {
+        if !StackFuture::<F, STACK_SIZE>::has_space_for::<F>() {
+            panic!("F is too large");
+        }
+
+        if !StackFuture::<F, STACK_SIZE>::has_alignment_for::<F>() {
+            panic!("F has incompatible alignment");
+        }
+    };
+}
+
 #[cfg(test)]
 mod tests;

src/tests.rs

@@ -80,6 +80,21 @@ fn destructor_runs() {
     assert!(destructed);
 }
 
+#[test]
+fn test_size_failure() {
+    async fn fill_buf(buf: &mut [u8]) {
+        buf[0] = 42;
+    }
+
+    let f = async {
+        let mut buf = [0u8; 256];
+        fill_buf(&mut buf).await;
+        buf[0]
+    };
+
+    assert!(StackFuture::<_, 4>::try_from(f).is_err());
+}
+
 #[test]
 fn test_alignment() {
     // A test to make sure we store the wrapped future with the correct alignment
@@ -99,7 +114,6 @@ fn test_alignment() {
 }
 
 #[test]
-#[should_panic]
 fn test_alignment_failure() {
     // A test to make sure we store the wrapped future with the correct alignment
 
@@ -113,7 +127,7 @@ fn test_alignment_failure() {
             Poll::Pending
         }
     }
-    StackFuture::<'_, _, 1016>::from(BigAlignment(42));
+    assert!(StackFuture::<'_, _, 1016>::try_from(BigAlignment(42)).is_err());
 }
 
 #[cfg(feature = "alloc")]