Snap for 10453563 from 9d69f56f4e1fde5260b6d129c57b30af15037a10 to mainline-media-swcodec-releaseaml_swc_341711000aml_swc_341619000aml_swc_341513600aml_swc_341312300aml_swc_341312020aml_swc_341111000aml_swc_341011020aml_swc_340922010android14-mainline-media-swcodec-release

Change-Id: I9834050046331edaf2685f471a22b2fb6e18ed5c

4 files changed, 160 insertions, 49 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 4a868bf..dd689ec 100644
--- a/src/lib.rs
+++ b/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;
diff --git a/src/raw.rs b/src/raw.rs
index 9b5854c..bb031da 100644
--- a/src/raw.rs
+++ b/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.
@@ -34,6 +34,12 @@ pub(crate) struct TaskVTable {
 
     /// Creates a new waker associated with the task.
     pub(crate) clone_waker: unsafe fn(ptr: *const ()) -> RawWaker,
+
+    /// The memory layout of the task. This information enables
+    /// debuggers to decode raw task memory blobs. Do not remove
+    /// the field, even if it appears to be unused.
+    #[allow(unused)]
+    pub(crate) layout_info: &'static Option<TaskLayout>,
 }
 
 /// Memory layout of a task.
@@ -45,7 +51,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 +86,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 +136,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.
@@ -120,6 +161,7 @@ where
                     destroy: Self::destroy,
                     run: Self::run,
                     clone_waker: Self::clone_waker,
+                    layout_info: &Self::TASK_LAYOUT,
                 },
             });
 
@@ -149,32 +191,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(),
         }
     }
 
diff --git a/src/task.rs b/src/task.rs
index fff918c..8ecd746 100644
--- a/src/task.rs
+++ b/src/task.rs
@@ -99,6 +99,7 @@ impl<T> Task<T> {
     /// # Examples
     ///
     /// ```
+    /// # if cfg!(miri) { return; } // Miri does not support epoll
     /// use smol::{future, Executor, Timer};
     /// use std::thread;
     /// use std::time::Duration;
@@ -395,6 +396,19 @@ impl<T> Task<T> {
         let header = ptr as *const Header;
         unsafe { &*header }
     }
+
+    /// Returns `true` if the current task is finished.
+    ///
+    /// Note that in a multithreaded environment, this task can change finish immediately after calling this function.
+    pub fn is_finished(&self) -> bool {
+        let ptr = self.ptr.as_ptr();
+        let header = ptr as *const Header;
+
+        unsafe {
+            let state = (*header).state.load(Ordering::Acquire);
+            state & (CLOSED | COMPLETED) != 0
+        }
+    }
 }
 
 impl<T> Drop for Task<T> {
@@ -472,6 +486,7 @@ impl<T> FallibleTask<T> {
     /// # Examples
     ///
     /// ```
+    /// # if cfg!(miri) { return; } // Miri does not support epoll
     /// use smol::{future, Executor, Timer};
     /// use std::thread;
     /// use std::time::Duration;
diff --git a/src/utils.rs b/src/utils.rs
index cb9b65e..189e9af 100644
--- a/src/utils.rs
+++ b/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
+    }
 }