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}