Evaluate the layouts for the tasks at compile time (#30)

Author: notgull

.clippy.toml

@@ -1 +1 @@
-msrv = "1.39"
+msrv = "1.47"

.github/workflows/ci.yml

@@ -45,7 +45,7 @@ jobs:
       matrix:
         # When updating this, the reminder to update the minimum supported
         # Rust version in Cargo.toml and .clippy.toml.
-        rust: ['1.39']
+        rust: ['1.47']
     steps:
       - uses: actions/checkout@v3
       - name: Install Rust

Cargo.toml

@@ -6,7 +6,7 @@ name = "async-task"
 version = "4.2.0"
 authors = ["Stjepan Glavina <[email protected]>"]
 edition = "2018"
-rust-version = "1.39"
+rust-version = "1.47"
 license = "Apache-2.0 OR MIT"
 repository = "https://github.com/smol-rs/async-task"
 description = "Task abstraction for building executors"

src/lib.rs

@@ -74,6 +74,17 @@
 
 extern crate alloc;
 
+/// We can't use `?` in const contexts yet, so this macro acts
+/// as a workaround.
+macro_rules! leap {
+    ($x: expr) => {{
+        match ($x) {
+            Some(val) => val,
+            None => return None,
+        }
+    }};
+}
+
 mod header;
 mod raw;
 mod runnable;

src/raw.rs

@@ -1,4 +1,4 @@
-use alloc::alloc::Layout;
+use alloc::alloc::Layout as StdLayout;
 use core::cell::UnsafeCell;
 use core::future::Future;
 use core::mem::{self, ManuallyDrop};
@@ -9,7 +9,7 @@ use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
 
 use crate::header::Header;
 use crate::state::*;
-use crate::utils::{abort, abort_on_panic, extend};
+use crate::utils::{abort, abort_on_panic, max, Layout};
 use crate::Runnable;
 
 /// The vtable for a task.
@@ -45,7 +45,7 @@ pub(crate) struct TaskVTable {
 #[derive(Clone, Copy)]
 pub(crate) struct TaskLayout {
     /// Memory layout of the whole task.
-    pub(crate) layout: Layout,
+    pub(crate) layout: StdLayout,
 
     /// Offset into the task at which the schedule function is stored.
     pub(crate) offset_s: usize,
@@ -80,6 +80,39 @@ impl<F, T, S> Clone for RawTask<F, T, S> {
     }
 }
 
+impl<F, T, S> RawTask<F, T, S> {
+    const TASK_LAYOUT: Option<TaskLayout> = Self::eval_task_layout();
+
+    /// Computes the memory layout for a task.
+    #[inline]
+    const fn eval_task_layout() -> Option<TaskLayout> {
+        // Compute the layouts for `Header`, `S`, `F`, and `T`.
+        let layout_header = Layout::new::<Header>();
+        let layout_s = Layout::new::<S>();
+        let layout_f = Layout::new::<F>();
+        let layout_r = Layout::new::<T>();
+
+        // Compute the layout for `union { F, T }`.
+        let size_union = max(layout_f.size(), layout_r.size());
+        let align_union = max(layout_f.align(), layout_r.align());
+        let layout_union = Layout::from_size_align(size_union, align_union);
+
+        // Compute the layout for `Header` followed `S` and `union { F, T }`.
+        let layout = layout_header;
+        let (layout, offset_s) = leap!(layout.extend(layout_s));
+        let (layout, offset_union) = leap!(layout.extend(layout_union));
+        let offset_f = offset_union;
+        let offset_r = offset_union;
+
+        Some(TaskLayout {
+            layout: unsafe { layout.into_std() },
+            offset_s,
+            offset_f,
+            offset_r,
+        })
+    }
+}
+
 impl<F, T, S> RawTask<F, T, S>
 where
     F: Future<Output = T>,
@@ -97,7 +130,9 @@ where
     /// It is assumed that initially only the `Runnable` and the `Task` exist.
     pub(crate) fn allocate(future: F, schedule: S) -> NonNull<()> {
         // Compute the layout of the task for allocation. Abort if the computation fails.
-        let task_layout = abort_on_panic(|| Self::task_layout());
+        //
+        // n.b. notgull: task_layout now automatically aborts instead of panicking
+        let task_layout = Self::task_layout();
 
         unsafe {
             // Allocate enough space for the entire task.
@@ -149,32 +184,12 @@ where
         }
     }
 
-    /// Returns the memory layout for a task.
+    /// Returns the layout of the task.
     #[inline]
     fn task_layout() -> TaskLayout {
-        // Compute the layouts for `Header`, `S`, `F`, and `T`.
-        let layout_header = Layout::new::<Header>();
-        let layout_s = Layout::new::<S>();
-        let layout_f = Layout::new::<F>();
-        let layout_r = Layout::new::<T>();
-
-        // Compute the layout for `union { F, T }`.
-        let size_union = layout_f.size().max(layout_r.size());
-        let align_union = layout_f.align().max(layout_r.align());
-        let layout_union = unsafe { Layout::from_size_align_unchecked(size_union, align_union) };
-
-        // Compute the layout for `Header` followed `S` and `union { F, T }`.
-        let layout = layout_header;
-        let (layout, offset_s) = extend(layout, layout_s);
-        let (layout, offset_union) = extend(layout, layout_union);
-        let offset_f = offset_union;
-        let offset_r = offset_union;
-
-        TaskLayout {
-            layout,
-            offset_s,
-            offset_f,
-            offset_r,
+        match Self::TASK_LAYOUT {
+            Some(tl) => tl,
+            None => abort(),
         }
     }
 

src/utils.rs

@@ -1,4 +1,4 @@
-use core::alloc::Layout;
+use core::alloc::Layout as StdLayout;
 use core::mem;
 
 /// Aborts the process.
@@ -36,29 +36,92 @@ pub(crate) fn abort_on_panic<T>(f: impl FnOnce() -> T) -> T {
     t
 }
 
-/// Returns the layout for `a` followed by `b` and the offset of `b`.
-///
-/// This function was adapted from the currently unstable `Layout::extend()`:
-/// https://doc.rust-lang.org/nightly/std/alloc/struct.Layout.html#method.extend
-#[inline]
-pub(crate) fn extend(a: Layout, b: Layout) -> (Layout, usize) {
-    let new_align = a.align().max(b.align());
-    let pad = padding_needed_for(a, b.align());
+/// A version of `alloc::alloc::Layout` that can be used in the const
+/// position.
+#[derive(Clone, Copy, Debug)]
+pub(crate) struct Layout {
+    size: usize,
+    align: usize,
+}
 
-    let offset = a.size().checked_add(pad).unwrap();
-    let new_size = offset.checked_add(b.size()).unwrap();
+impl Layout {
+    /// Creates a new `Layout` with the given size and alignment.
+    #[inline]
+    pub(crate) const fn from_size_align(size: usize, align: usize) -> Self {
+        Self { size, align }
+    }
+
+    /// Creates a new `Layout` for the given sized type.
+    #[inline]
+    pub(crate) const fn new<T>() -> Self {
+        Self::from_size_align(mem::size_of::<T>(), mem::align_of::<T>())
+    }
 
-    let layout = Layout::from_size_align(new_size, new_align).unwrap();
-    (layout, offset)
+    /// Convert this into the standard library's layout type.
+    ///
+    /// # Safety
+    ///
+    /// - `align` must be non-zero and a power of two.
+    /// - When rounded up to the nearest multiple of `align`, the size
+    ///   must not overflow.
+    #[inline]
+    pub(crate) const unsafe fn into_std(self) -> StdLayout {
+        StdLayout::from_size_align_unchecked(self.size, self.align)
+    }
+
+    /// Get the alignment of this layout.
+    #[inline]
+    pub(crate) const fn align(&self) -> usize {
+        self.align
+    }
+
+    /// Get the size of this layout.
+    #[inline]
+    pub(crate) const fn size(&self) -> usize {
+        self.size
+    }
+
+    /// Returns the layout for `a` followed by `b` and the offset of `b`.
+    ///
+    /// This function was adapted from the currently unstable `Layout::extend()`:
+    /// https://doc.rust-lang.org/nightly/std/alloc/struct.Layout.html#method.extend
+    #[inline]
+    pub(crate) const fn extend(self, other: Layout) -> Option<(Layout, usize)> {
+        let new_align = max(self.align(), other.align());
+        let pad = self.padding_needed_for(other.align());
+
+        let offset = leap!(self.size().checked_add(pad));
+        let new_size = leap!(offset.checked_add(other.size()));
+
+        // return None if any of the following are true:
+        // - align is 0 (implied false by is_power_of_two())
+        // - align is not a power of 2
+        // - size rounded up to align overflows
+        if !new_align.is_power_of_two() || new_size > core::usize::MAX - (new_align - 1) {
+            return None;
+        }
+
+        let layout = Layout::from_size_align(new_size, new_align);
+        Some((layout, offset))
+    }
+
+    /// Returns the padding after `layout` that aligns the following address to `align`.
+    ///
+    /// This function was adapted from the currently unstable `Layout::padding_needed_for()`:
+    /// https://doc.rust-lang.org/nightly/std/alloc/struct.Layout.html#method.padding_needed_for
+    #[inline]
+    pub(crate) const fn padding_needed_for(self, align: usize) -> usize {
+        let len = self.size();
+        let len_rounded_up = len.wrapping_add(align).wrapping_sub(1) & !align.wrapping_sub(1);
+        len_rounded_up.wrapping_sub(len)
+    }
 }
 
-/// Returns the padding after `layout` that aligns the following address to `align`.
-///
-/// This function was adapted from the currently unstable `Layout::padding_needed_for()`:
-/// https://doc.rust-lang.org/nightly/std/alloc/struct.Layout.html#method.padding_needed_for
 #[inline]
-pub(crate) fn padding_needed_for(layout: Layout, align: usize) -> usize {
-    let len = layout.size();
-    let len_rounded_up = len.wrapping_add(align).wrapping_sub(1) & !align.wrapping_sub(1);
-    len_rounded_up.wrapping_sub(len)
+pub(crate) const fn max(left: usize, right: usize) -> usize {
+    if left > right {
+        left
+    } else {
+        right
+    }
 }