Refactor preemption-related code in OSTD

ostd/src/task/preempt/cpu_local.rs

@@ -0,0 +1,61 @@
+// SPDX-License-Identifier: MPL-2.0
+
+//! This module maintains preemption-related information for the curren task
+//! on a CPU with a single 32-bit, CPU-local integer value.
+//!
+//! * Bits from 0 to 30 represents an unsigned counter called `guard_count`,
+//!   which is the number of `DisablePreemptGuard` instances held by the
+//!   current CPU;
+//! * Bit 31 is set to `!need_preempt`, where `need_preempt` is a boolean value
+//!   that will be set by the scheduler when it decides that the current task
+//!   _needs_ to be preempted.
+//!
+//! Thus, the current task on a CPU _should_ be preempted if and only if this
+//! integer is equal to zero.
+//!
+//! The initial value of this integer is equal to `1 << 31`.
+//!
+//! This module provides a set of functions to access and manipulate
+//! `guard_count` and `need_preempt`.
+
+use crate::cpu_local_cell;
+
+/// Returns whether the current task _should_ be preempted or not.
+///
+/// `should_preempt() == need_preempt() && get_guard_count() == 0`.
+pub(in crate::task) fn should_preempt() -> bool {
+    PREEMPT_INFO.load() == 0
+}
+
+#[allow(dead_code)]
+pub(in crate::task) fn need_preempt() -> bool {
+    PREEMPT_INFO.load() & NEED_PREEMPT_MASK == 0
+}
+
+pub(in crate::task) fn set_need_preempt() {
+    PREEMPT_INFO.bitand_assign(!NEED_PREEMPT_MASK);
+}
+
+pub(in crate::task) fn clear_need_preempt() {
+    PREEMPT_INFO.bitor_assign(NEED_PREEMPT_MASK);
+}
+
+pub(in crate::task) fn get_guard_count() -> u32 {
+    PREEMPT_INFO.load() & GUARD_COUNT_MASK
+}
+
+pub(in crate::task) fn inc_guard_count() {
+    PREEMPT_INFO.add_assign(1);
+}
+
+pub(in crate::task) fn dec_guard_count() {
+    debug_assert!(get_guard_count() > 0);
+    PREEMPT_INFO.sub_assign(1);
+}
+
+cpu_local_cell! {
+    static PREEMPT_INFO: u32 = NEED_PREEMPT_MASK;
+}
+
+const NEED_PREEMPT_MASK: u32 = 1 << 31;
+const GUARD_COUNT_MASK: u32 = (1 << 31) - 1;

ostd/src/task/preempt/guard.rs

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: MPL-2.0
+
+/// A guard for disable preempt.
+#[clippy::has_significant_drop]
+#[must_use]
+pub struct DisablePreemptGuard {
+    // This private field prevents user from constructing values of this type directly.
+    _private: (),
+}
+
+impl !Send for DisablePreemptGuard {}
+
+impl DisablePreemptGuard {
+    fn new() -> Self {
+        super::cpu_local::inc_guard_count();
+        Self { _private: () }
+    }
+
+    /// Transfer this guard to a new guard.
+    /// This guard must be dropped after this function.
+    pub fn transfer_to(&self) -> Self {
+        disable_preempt()
+    }
+}
+
+impl Drop for DisablePreemptGuard {
+    fn drop(&mut self) {
+        super::cpu_local::dec_guard_count();
+    }
+}
+
+/// Disables preemption.
+pub fn disable_preempt() -> DisablePreemptGuard {
+    DisablePreemptGuard::new()
+}

ostd/src/task/preempt/mod.rs

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: MPL-2.0
+
+pub(super) mod cpu_local;
+mod guard;
+
+pub use self::guard::{disable_preempt, DisablePreemptGuard};