|
| 1 | +use core::fmt; |
| 2 | +use core::mem::{self, PinMut}; |
| 3 | + |
| 4 | +use Future; |
| 5 | +use super::{Context, Poll, TaskObj, SpawnErrorKind, SpawnObjError}; |
| 6 | + |
| 7 | +/// A custom trait object for polling tasks, roughly akin to |
| 8 | +/// `Box<Future<Output = ()>>`. |
| 9 | +/// Contrary to `TaskObj`, `LocalTaskObj` does not have a `Send` bound |
| 10 | +pub struct LocalTaskObj { |
| 11 | + ptr: *mut (), |
| 12 | + poll_fn: unsafe fn(*mut (), &mut Context) -> Poll<()>, |
| 13 | + drop_fn: unsafe fn(*mut ()), |
| 14 | +} |
| 15 | + |
| 16 | +impl fmt::Debug for LocalTaskObj { |
| 17 | + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { |
| 18 | + f.debug_struct("LocalTaskObj") |
| 19 | + .finish() |
| 20 | + } |
| 21 | +} |
| 22 | + |
| 23 | +/// A custom implementation of a task trait object for `LocalTaskObj`, providing |
| 24 | +/// a hand-rolled vtable. |
| 25 | +/// |
| 26 | +/// This custom representation is typically used only in `no_std` contexts, |
| 27 | +/// where the default `Box`-based implementation is not available. |
| 28 | +/// |
| 29 | +/// The implementor must guarantee that it is safe to call `poll` repeatedly (in |
| 30 | +/// a non-concurrent fashion) with the result of `into_raw` until `drop` is |
| 31 | +/// called. |
| 32 | +pub unsafe trait UnsafeLocalTask: 'static { |
| 33 | + /// Convert a owned instance into a (conceptually owned) void pointer. |
| 34 | + fn into_raw(self) -> *mut (); |
| 35 | + |
| 36 | + /// Poll the task represented by the given void pointer. |
| 37 | + /// |
| 38 | + /// # Safety |
| 39 | + /// |
| 40 | + /// The trait implementor must guarantee that it is safe to repeatedly call |
| 41 | + /// `poll` with the result of `into_raw` until `drop` is called; such calls |
| 42 | + /// are not, however, allowed to race with each other or with calls to `drop`. |
| 43 | + unsafe fn poll(task: *mut (), cx: &mut Context) -> Poll<()>; |
| 44 | + |
| 45 | + /// Drops the task represented by the given void pointer. |
| 46 | + /// |
| 47 | + /// # Safety |
| 48 | + /// |
| 49 | + /// The trait implementor must guarantee that it is safe to call this |
| 50 | + /// function once per `into_raw` invocation; that call cannot race with |
| 51 | + /// other calls to `drop` or `poll`. |
| 52 | + unsafe fn drop(task: *mut ()); |
| 53 | +} |
| 54 | + |
| 55 | +impl LocalTaskObj { |
| 56 | + /// Create a `LocalTaskObj` from a custom trait object representation. |
| 57 | + #[inline] |
| 58 | + pub fn new<T: UnsafeLocalTask>(t: T) -> LocalTaskObj { |
| 59 | + LocalTaskObj { |
| 60 | + ptr: t.into_raw(), |
| 61 | + poll_fn: T::poll, |
| 62 | + drop_fn: T::drop, |
| 63 | + } |
| 64 | + } |
| 65 | + |
| 66 | + /// Converts the `LocalTaskObj` into a `TaskObj` |
| 67 | + /// To make this operation safe one has to ensure that the `LocalTaskObj` |
| 68 | + /// implements `Send`. |
| 69 | + pub unsafe fn as_task_obj(self) -> TaskObj { |
| 70 | + // Safety: Both structs have the same memory layout |
| 71 | + mem::transmute::<LocalTaskObj, TaskObj>(self) |
| 72 | + } |
| 73 | +} |
| 74 | + |
| 75 | +impl From<TaskObj> for LocalTaskObj { |
| 76 | + fn from(task: TaskObj) -> LocalTaskObj { |
| 77 | + unsafe { |
| 78 | + // Safety: Both structs have the same memory layout |
| 79 | + mem::transmute::<TaskObj, LocalTaskObj>(task) |
| 80 | + } |
| 81 | + } |
| 82 | +} |
| 83 | + |
| 84 | +impl Future for LocalTaskObj { |
| 85 | + type Output = (); |
| 86 | + |
| 87 | + #[inline] |
| 88 | + fn poll(self: PinMut<Self>, cx: &mut Context) -> Poll<()> { |
| 89 | + unsafe { |
| 90 | + (self.poll_fn)(self.ptr, cx) |
| 91 | + } |
| 92 | + } |
| 93 | +} |
| 94 | + |
| 95 | +impl Drop for LocalTaskObj { |
| 96 | + fn drop(&mut self) { |
| 97 | + unsafe { |
| 98 | + (self.drop_fn)(self.ptr) |
| 99 | + } |
| 100 | + } |
| 101 | +} |
| 102 | + |
| 103 | +/// The result of a failed spawn |
| 104 | +#[derive(Debug)] |
| 105 | +pub struct SpawnLocalObjError { |
| 106 | + /// The kind of error |
| 107 | + pub kind: SpawnErrorKind, |
| 108 | + |
| 109 | + /// The task for which spawning was attempted |
| 110 | + pub task: LocalTaskObj, |
| 111 | +} |
| 112 | + |
| 113 | +impl SpawnLocalObjError { |
| 114 | + /// Converts the `SpawnLocalObjError` into a `SpawnObjError` |
| 115 | + /// To make this operation safe one has to ensure that the `LocalTaskObj` |
| 116 | + /// stored inside implements `Send`. |
| 117 | + pub unsafe fn as_spawn_obj_error(self) -> SpawnObjError { |
| 118 | + // Safety: Both structs have the same memory layout |
| 119 | + mem::transmute::<SpawnLocalObjError, SpawnObjError>(self) |
| 120 | + } |
| 121 | +} |
| 122 | + |
| 123 | +impl From<SpawnObjError> for SpawnLocalObjError { |
| 124 | + fn from(error: SpawnObjError) -> SpawnLocalObjError { |
| 125 | + unsafe { |
| 126 | + // Safety: Both structs have the same memory layout |
| 127 | + mem::transmute::<SpawnObjError, SpawnLocalObjError>(error) |
| 128 | + } |
| 129 | + } |
| 130 | +} |
| 131 | + |
| 132 | +if_std! { |
| 133 | + use std::boxed::{Box, PinBox}; |
| 134 | + |
| 135 | + unsafe impl<F: Future<Output = ()> + 'static> UnsafeLocalTask for PinBox<F> { |
| 136 | + fn into_raw(self) -> *mut () { |
| 137 | + PinBox::into_raw(self) as *mut () |
| 138 | + } |
| 139 | + |
| 140 | + unsafe fn poll(task: *mut (), cx: &mut Context) -> Poll<()> { |
| 141 | + let ptr = task as *mut F; |
| 142 | + let pin: PinMut<F> = PinMut::new_unchecked(&mut *ptr); |
| 143 | + pin.poll(cx) |
| 144 | + } |
| 145 | + |
| 146 | + unsafe fn drop(task: *mut ()) { |
| 147 | + drop(PinBox::from_raw(task as *mut F)) |
| 148 | + } |
| 149 | + } |
| 150 | + |
| 151 | + impl<F: Future<Output = ()> + 'static> From<PinBox<F>> for LocalTaskObj { |
| 152 | + fn from(boxed: PinBox<F>) -> Self { |
| 153 | + LocalTaskObj::new(boxed) |
| 154 | + } |
| 155 | + } |
| 156 | + |
| 157 | + impl<F: Future<Output = ()> + 'static> From<Box<F>> for LocalTaskObj { |
| 158 | + fn from(boxed: Box<F>) -> Self { |
| 159 | + LocalTaskObj::new(PinBox::from(boxed)) |
| 160 | + } |
| 161 | + } |
| 162 | +} |
0 commit comments