axtask/
task_ext.rs

1//! User-defined task extended data.
2
3use core::alloc::Layout;
4use core::mem::{align_of, size_of};
5
6#[unsafe(no_mangle)]
7#[linkage = "weak"]
8static __AX_TASK_EXT_SIZE: usize = 0;
9
10#[unsafe(no_mangle)]
11#[linkage = "weak"]
12static __AX_TASK_EXT_ALIGN: usize = 0;
13
14/// A wrapper of pointer to the task extended data.
15pub(crate) struct AxTaskExt {
16    ptr: *mut u8,
17}
18
19impl AxTaskExt {
20    /// Returns the expected size of the task extended structure.
21    pub fn size() -> usize {
22        unsafe extern "C" {
23            static __AX_TASK_EXT_SIZE: usize;
24        }
25        unsafe { __AX_TASK_EXT_SIZE }
26    }
27
28    /// Returns the expected alignment of the task extended structure.
29    pub fn align() -> usize {
30        unsafe extern "C" {
31            static __AX_TASK_EXT_ALIGN: usize;
32        }
33        unsafe { __AX_TASK_EXT_ALIGN }
34    }
35
36    /// Construct an empty task extended structure that contains no data
37    /// (zero size).
38    pub const fn empty() -> Self {
39        Self {
40            ptr: core::ptr::null_mut(),
41        }
42    }
43
44    /// Returns `true` if the task extended structure is empty.
45    pub const fn is_empty(&self) -> bool {
46        self.ptr.is_null()
47    }
48
49    /// Allocates the space for the task extended data, but does not
50    /// initialize the data.
51    pub unsafe fn uninited() -> Self {
52        let size = Self::size();
53        let align = Self::align();
54        let ptr = if size == 0 {
55            core::ptr::null_mut()
56        } else {
57            let layout = Layout::from_size_align(size, align).unwrap();
58            unsafe { alloc::alloc::alloc(layout) }
59        };
60        Self { ptr }
61    }
62
63    /// Gets the raw pointer to the task extended data.
64    pub const fn as_ptr(&self) -> *mut u8 {
65        self.ptr
66    }
67
68    /// Write the given object to the task extended data.
69    ///
70    /// Returns [`None`] if the data size is zero, otherwise returns a mutable
71    /// reference to the content.
72    ///
73    /// # Panics
74    ///
75    /// Panics If the sizes and alignments of the two object do not match.
76    pub fn write<T: Sized>(&mut self, data: T) -> Option<&mut T> {
77        let data_size = size_of::<T>();
78        let data_align = align_of::<T>();
79        if data_size != Self::size() {
80            panic!("size mismatch: {} != {}", data_size, Self::size());
81        }
82        if data_align != Self::align() {
83            panic!("align mismatch: {} != {}", data_align, Self::align());
84        }
85
86        if self.ptr.is_null() {
87            *self = unsafe { Self::uninited() };
88        }
89        if data_size > 0 {
90            let ptr = self.ptr as *mut T;
91            assert!(!ptr.is_null());
92            unsafe {
93                ptr.write(data);
94                Some(&mut *ptr)
95            }
96        } else {
97            None
98        }
99    }
100}
101
102impl Drop for AxTaskExt {
103    fn drop(&mut self) {
104        if !self.ptr.is_null() {
105            let layout = Layout::from_size_align(Self::size(), 0x10).unwrap();
106            unsafe { alloc::alloc::dealloc(self.ptr, layout) };
107        }
108    }
109}
110
111/// A trait to convert [`TaskInner::task_ext_ptr`] to the reference of the
112/// concrete type.
113///
114/// [`TaskInner::task_ext_ptr`]: crate::TaskInner::task_ext_ptr
115pub trait TaskExtRef<T: Sized> {
116    /// Get a reference to the task extended data.
117    fn task_ext(&self) -> &T;
118}
119
120/// A trait to convert [`TaskInner::task_ext_ptr`] to the mutable reference of
121/// the concrete type.
122///
123/// [`TaskInner::task_ext_ptr`]: crate::TaskInner::task_ext_ptr
124pub trait TaskExtMut<T: Sized> {
125    /// Get a mutable reference to the task extended data.
126    fn task_ext_mut(&mut self) -> &mut T;
127}
128
129/// Define the task extended data.
130///
131/// It automatically implements [`TaskExtRef`] and [`TaskExtMut`] for
132/// [`TaskInner`].
133///
134/// # Example
135///
136/// ```
137/// # #![allow(non_local_definitions)]
138/// use axtask::{def_task_ext, TaskExtRef, TaskInner};
139///
140/// pub struct TaskExtImpl {
141///    proc_id: usize,
142/// }
143///
144/// def_task_ext!(TaskExtImpl);
145///
146/// axtask::init_scheduler();
147///
148/// let mut inner = TaskInner::new(|| {},  "".into(), 0x1000);
149/// assert!(inner.init_task_ext(TaskExtImpl { proc_id: 233 }).is_some());
150/// // cannot initialize twice
151/// assert!(inner.init_task_ext(TaskExtImpl { proc_id: 0xdead }).is_none());
152///
153/// let task = axtask::spawn_task(inner);
154/// assert_eq!(task.task_ext().proc_id, 233);
155/// ```
156///
157/// [`TaskInner`]: crate::TaskInner
158#[macro_export]
159macro_rules! def_task_ext {
160    ($task_ext_struct:ty) => {
161        #[unsafe(no_mangle)]
162        static __AX_TASK_EXT_SIZE: usize = ::core::mem::size_of::<$task_ext_struct>();
163
164        #[unsafe(no_mangle)]
165        static __AX_TASK_EXT_ALIGN: usize = ::core::mem::align_of::<$task_ext_struct>();
166
167        impl $crate::TaskExtRef<$task_ext_struct> for $crate::TaskInner {
168            fn task_ext(&self) -> &$task_ext_struct {
169                unsafe {
170                    let ptr = self.task_ext_ptr() as *const $task_ext_struct;
171                    assert!(!ptr.is_null());
172                    &*ptr
173                }
174            }
175        }
176
177        impl $crate::TaskExtMut<$task_ext_struct> for $crate::TaskInner {
178            fn task_ext_mut(&mut self) -> &mut $task_ext_struct {
179                unsafe {
180                    let ptr = self.task_ext_ptr() as *mut $task_ext_struct;
181                    assert!(!ptr.is_null());
182                    &mut *ptr
183                }
184            }
185        }
186    };
187}
axtask/task_ext.rs

axtask/
task_ext.rs
