Add logical CPU ID newtype, switch to u32.

src/arch/x86_64/gdt.rs

@@ -3,6 +3,7 @@
 use core::convert::TryInto;
 use core::mem;
 
+use crate::LogicalCpuId;
 use crate::paging::{RmmA, RmmArch};
 use crate::percpu::PercpuBlock;
 
@@ -144,7 +145,7 @@ unsafe fn load_segments() {
 
 /// Initialize GDT and PCR.
 #[cold]
-pub unsafe fn init_paging(stack_offset: usize, cpu_id: usize) {
+pub unsafe fn init_paging(stack_offset: usize, cpu_id: LogicalCpuId) {
     let pcr_frame = crate::memory::allocate_frames(1).expect("failed to allocate PCR");
     let pcr = &mut *(RmmA::phys_to_virt(pcr_frame.start_address()).data() as *mut ProcessorControlRegion);
 

src/arch/x86_64/idt.rs

@@ -8,7 +8,7 @@ use alloc::collections::BTreeMap;
 use x86::segmentation::Descriptor as X86IdtEntry;
 use x86::dtables::{self, DescriptorTablePointer};
 
-use crate::interrupt::*;
+use crate::{interrupt::*, LogicalCpuId};
 use crate::ipi::IpiKind;
 
 use spin::RwLock;
@@ -68,10 +68,10 @@ impl Idt {
 static mut INIT_BSP_IDT: Idt = Idt::new();
 
 // TODO: VecMap?
-pub static IDTS: RwLock<Option<BTreeMap<usize, &'static mut Idt>>> = RwLock::new(None);
+pub static IDTS: RwLock<Option<BTreeMap<LogicalCpuId, &'static mut Idt>>> = RwLock::new(None);
 
 #[inline]
-pub fn is_reserved(cpu_id: usize, index: u8) -> bool {
+pub fn is_reserved(cpu_id: LogicalCpuId, index: u8) -> bool {
     let byte_index = index / 64;
     let bit = index % 64;
 
@@ -79,7 +79,7 @@ pub fn is_reserved(cpu_id: usize, index: u8) -> bool {
 }
 
 #[inline]
-pub fn set_reserved(cpu_id: usize, index: u8, reserved: bool) {
+pub fn set_reserved(cpu_id: LogicalCpuId, index: u8, reserved: bool) {
     let byte_index = index / 64;
     let bit = index % 64;
 
@@ -97,7 +97,7 @@ pub fn allocate_interrupt() -> Option<NonZeroU8> {
     None
 }
 
-pub fn available_irqs_iter(cpu_id: usize) -> impl Iterator<Item = u8> + 'static {
+pub fn available_irqs_iter(cpu_id: LogicalCpuId) -> impl Iterator<Item = u8> + 'static {
     (32..=254).filter(move |&index| !is_reserved(cpu_id, index))
 }
 
@@ -123,7 +123,7 @@ const fn new_idt_reservations() -> [AtomicU64; 4] {
 }
 
 /// Initialize the IDT for a
-pub unsafe fn init_paging_post_heap(is_bsp: bool, cpu_id: usize) {
+pub unsafe fn init_paging_post_heap(is_bsp: bool, cpu_id: LogicalCpuId) {
     let mut idts_guard = IDTS.write();
     let idts_btree = idts_guard.get_or_insert_with(BTreeMap::new);
 

src/arch/x86_64/misc.rs

@@ -1,9 +1,10 @@
 use x86::controlregs::Cr4;
 use x86::cpuid::ExtendedFeatures;
 
+use crate::LogicalCpuId;
 use crate::cpuid::cpuid_always;
 
-pub unsafe fn init(cpu_id: usize) {
+pub unsafe fn init(cpu_id: LogicalCpuId) {
     let has_ext_feat = |feat: fn(ExtendedFeatures) -> bool| {
         cpuid_always()
             .get_extended_feature_info()
@@ -33,6 +34,6 @@ pub unsafe fn init(cpu_id: usize) {
     }
 
     if let Some(feats) = cpuid_always().get_extended_processor_and_feature_identifiers() && feats.has_rdtscp() {
-        x86::msr::wrmsr(x86::msr::IA32_TSC_AUX, cpu_id as u64);
+        x86::msr::wrmsr(x86::msr::IA32_TSC_AUX, cpu_id.get().into());
     }
 }

src/arch/x86_64/rmm.rs

@@ -22,6 +22,8 @@ use rmm::{
 };
 use spin::Mutex;
 
+use crate::LogicalCpuId;
+
 use super::CurrentRmmArch as RmmA;
 
 // Keep synced with OsMemoryKind in bootloader
@@ -275,14 +277,14 @@ impl KernelMapper {
 
         prev_count > 0
     }
-    pub unsafe fn lock_for_manual_mapper(current_processor: usize, mapper: crate::paging::PageMapper) -> Self {
-        let ro = Self::lock_inner(current_processor);
+    pub unsafe fn lock_for_manual_mapper(current_processor: LogicalCpuId, mapper: crate::paging::PageMapper) -> Self {
+        let ro = Self::lock_inner(current_processor.get() as usize);
         Self {
             mapper,
             ro,
         }
     }
-    pub fn lock_manually(current_processor: usize) -> Self {
+    pub fn lock_manually(current_processor: LogicalCpuId) -> Self {
         unsafe { Self::lock_for_manual_mapper(current_processor, PageMapper::current(TableKind::Kernel, FRAME_ALLOCATOR)) }
     }
     pub fn lock() -> Self {

src/arch/x86_64/start.rs

@@ -4,9 +4,9 @@
 /// defined in other files inside of the `arch` module
 
 use core::slice;
-use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering, AtomicU32};
 
-use crate::{allocator, memory};
+use crate::{allocator, memory, LogicalCpuId};
 #[cfg(feature = "acpi")]
 use crate::acpi;
 use crate::arch::pti;
@@ -28,7 +28,10 @@ static DATA_TEST_NONZERO: usize = usize::max_value();
 
 pub static KERNEL_BASE: AtomicUsize = AtomicUsize::new(0);
 pub static KERNEL_SIZE: AtomicUsize = AtomicUsize::new(0);
-pub static CPU_COUNT: AtomicUsize = AtomicUsize::new(0);
+
+// TODO: This probably shouldn't be an atomic. Only the BSP starts APs.
+pub static CPU_COUNT: AtomicU32 = AtomicU32::new(0);
+
 pub static AP_READY: AtomicBool = AtomicBool::new(false);
 static BSP_READY: AtomicBool = AtomicBool::new(false);
 
@@ -129,7 +132,7 @@ pub unsafe extern fn kstart(args_ptr: *const KernelArgs) -> ! {
         paging::init();
 
         // Set up GDT after paging with TLS
-        gdt::init_paging(args.stack_base as usize + args.stack_size as usize, 0);
+        gdt::init_paging(args.stack_base as usize + args.stack_size as usize, LogicalCpuId::BSP);
 
         // Set up IDT
         idt::init_paging_bsp();
@@ -149,13 +152,13 @@ pub unsafe extern fn kstart(args_ptr: *const KernelArgs) -> ! {
         #[cfg(feature = "graphical_debug")]
         graphical_debug::init_heap();
 
-        idt::init_paging_post_heap(true, 0);
+        idt::init_paging_post_heap(true, LogicalCpuId::BSP);
 
         // Activate memory logging
         log::init();
 
         // Initialize miscellaneous processor features
-        misc::init(0);
+        misc::init(LogicalCpuId::BSP);
 
         // Initialize devices
         device::init();
@@ -196,7 +199,9 @@ pub unsafe extern fn kstart(args_ptr: *const KernelArgs) -> ! {
 
 #[repr(packed)]
 pub struct KernelArgsAp {
+    // TODO: u32?
     cpu_id: u64,
+
     page_table: u64,
     stack_start: u64,
     stack_end: u64,
@@ -206,7 +211,7 @@ pub struct KernelArgsAp {
 pub unsafe extern fn kstart_ap(args_ptr: *const KernelArgsAp) -> ! {
     let cpu_id = {
         let args = &*args_ptr;
-        let cpu_id = args.cpu_id as usize;
+        let cpu_id = LogicalCpuId::new(args.cpu_id as u32);
         let bsp_table = args.page_table as usize;
         let _stack_start = args.stack_start as usize;
         let stack_end = args.stack_end as usize;

src/context/context.rs

@@ -10,6 +10,7 @@ use alloc::{
 };
 use spin::RwLock;
 
+use crate::LogicalCpuId;
 use crate::arch::{interrupt::InterruptStack, paging::PAGE_SIZE};
 use crate::common::aligned_box::AlignedBox;
 use crate::common::unique::Unique;
@@ -144,9 +145,9 @@ pub struct ContextSnapshot {
     pub status: Status,
     pub status_reason: &'static str,
     pub running: bool,
-    pub cpu_id: Option<usize>,
+    pub cpu_id: Option<LogicalCpuId>,
     pub cpu_time: u128,
-    pub sched_affinity: Option<usize>,
+    pub sched_affinity: Option<LogicalCpuId>,
     pub syscall: Option<(usize, usize, usize, usize, usize, usize)>,
     // Clone fields
     //TODO: is there a faster way than allocation?
@@ -230,7 +231,7 @@ pub struct Context {
     /// Context running or not
     pub running: bool,
     /// Current CPU ID
-    pub cpu_id: Option<usize>,
+    pub cpu_id: Option<LogicalCpuId>,
     /// Time this context was switched to
     pub switch_time: u128,
     /// Amount of CPU time used
@@ -238,7 +239,7 @@ pub struct Context {
     /// Scheduler CPU affinity. If set, [`cpu_id`] can except [`None`] never be anything else than
     /// this value.
     // TODO: bitmask (selection of multiple allowed CPUs)?
-    pub sched_affinity: Option<usize>,
+    pub sched_affinity: Option<LogicalCpuId>,
     /// Current system call
     pub syscall: Option<(usize, usize, usize, usize, usize, usize)>,
     /// Head buffer to use when system call buffers are not page aligned

src/context/list.rs

@@ -60,7 +60,7 @@ impl ContextList {
         // Zero is not a valid context ID, therefore add 1.
         //
         // FIXME: Ensure the number of CPUs can't switch between new_context calls.
-        let min = crate::cpu_count() + 1;
+        let min = crate::cpu_count() as usize + 1;
 
         self.next_id = core::cmp::max(self.next_id, min);
 

src/context/mod.rs

@@ -63,7 +63,7 @@ pub use self::arch::empty_cr3;
 
 pub fn init() {
     let mut contexts = contexts_mut();
-    let id = ContextId::from(crate::cpu_id() + 1);
+    let id = ContextId::from(crate::cpu_id().get() as usize + 1);
     let context_lock = contexts.insert_context_raw(id).expect("could not initialize first context");
     let mut context = context_lock.write();
     context.sched_affinity = Some(crate::cpu_id());

src/context/switch.rs

@@ -8,14 +8,14 @@ use spin::{RwLock, RwLockWriteGuard};
 
 use crate::context::signal::signal_handler;
 use crate::context::{arch, contexts, Context};
-use crate::interrupt;
+use crate::{interrupt, LogicalCpuId};
 use crate::percpu::PercpuBlock;
 use crate::ptrace;
 use crate::time;
 
 use super::ContextId;
 
-unsafe fn update_runnable(context: &mut Context, cpu_id: usize) -> bool {
+unsafe fn update_runnable(context: &mut Context, cpu_id: LogicalCpuId) -> bool {
     // Ignore already running contexts
     if context.running {
         return false;

src/lib.rs

@@ -65,7 +65,7 @@ extern crate alloc;
 #[macro_use]
 extern crate bitflags;
 
-use core::sync::atomic::{AtomicUsize, Ordering};
+use core::sync::atomic::{AtomicU32, Ordering};
 
 use crate::scheme::SchemeNamespace;
 
@@ -144,16 +144,16 @@ static ALLOCATOR: allocator::Allocator = allocator::Allocator;
 
 /// Get the current CPU's scheduling ID
 #[inline(always)]
-pub fn cpu_id() -> usize {
+pub fn cpu_id() -> LogicalCpuId {
     crate::percpu::PercpuBlock::current().cpu_id
 }
 
 /// The count of all CPUs that can have work scheduled
-static CPU_COUNT : AtomicUsize = AtomicUsize::new(0);
+static CPU_COUNT: AtomicU32 = AtomicU32::new(0);
 
 /// Get the number of CPUs currently active
 #[inline(always)]
-pub fn cpu_count() -> usize {
+pub fn cpu_count() -> u32 {
     CPU_COUNT.load(Ordering::Relaxed)
 }
 
@@ -175,14 +175,14 @@ pub struct Bootstrap {
 static BOOTSTRAP: spin::Once<Bootstrap> = spin::Once::new();
 
 /// This is the kernel entry point for the primary CPU. The arch crate is responsible for calling this
-pub fn kmain(cpus: usize, bootstrap: Bootstrap) -> ! {
-    CPU_COUNT.store(cpus, Ordering::SeqCst);
+pub fn kmain(cpu_count: u32, bootstrap: Bootstrap) -> ! {
+    CPU_COUNT.store(cpu_count, Ordering::SeqCst);
 
     //Initialize the first context, stored in kernel/src/context/mod.rs
     context::init();
 
     let pid = syscall::getpid();
-    info!("BSP: {:?} {}", pid, cpus);
+    info!("BSP: {:?} {}", pid, cpu_count);
     info!("Env: {:?}", ::core::str::from_utf8(bootstrap.env));
 
     BOOTSTRAP.call_once(|| bootstrap);
@@ -215,12 +215,12 @@ pub fn kmain(cpus: usize, bootstrap: Bootstrap) -> ! {
 
 /// This is the main kernel entry point for secondary CPUs
 #[allow(unreachable_code, unused_variables)]
-pub fn kmain_ap(id: usize) -> ! {
+pub fn kmain_ap(cpu_id: LogicalCpuId) -> ! {
     if cfg!(feature = "multi_core") {
         context::init();
 
         let pid = syscall::getpid();
-        info!("AP {}: {:?}", id, pid);
+        info!("AP {}: {:?}", cpu_id, pid);
 
         loop {
             unsafe {
@@ -234,7 +234,7 @@ pub fn kmain_ap(id: usize) -> ! {
             }
         }
     } else {
-        info!("AP {}: Disabled", id);
+        info!("AP {}: Disabled", cpu_id);
 
         loop {
             unsafe {
@@ -284,3 +284,26 @@ macro_rules! linker_offsets(
 pub mod kernel_executable_offsets {
     linker_offsets!(__text_start, __text_end, __rodata_start, __rodata_end, __data_start, __data_end, __bss_start, __bss_end, __usercopy_start, __usercopy_end);
 }
+
+/// A unique number used internally by the kernel to identify CPUs.
+///
+/// This is usually but not necessarily the same as the APIC ID.
+
+// TODO: Differentiate between logical CPU IDs and hardware CPU IDs (e.g. APIC IDs)
+#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
+// TODO: NonMaxUsize?
+// TODO: Optimize away this type if not cfg!(feature = "multi_core")
+pub struct LogicalCpuId(u32);
+
+impl LogicalCpuId {
+    pub const BSP: Self = Self::new(0);
+
+    pub const fn new(inner: u32) -> Self { Self(inner) }
+    pub const fn get(self) -> u32 { self.0 }
+}
+
+impl core::fmt::Display for LogicalCpuId {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        write!(f, "[logical cpu #{}]", self.0)
+    }
+}

src/percpu.rs

@@ -1,10 +1,10 @@
+use crate::LogicalCpuId;
 use crate::context::switch::ContextSwitchPercpu;
 
 /// The percpu block, that stored all percpu variables.
 pub struct PercpuBlock {
     /// A unique immutable number that identifies the current CPU - used for scheduling
-    // TODO: Differentiate between logical CPU IDs and hardware CPU IDs (e.g. APIC IDs)
-    pub cpu_id: usize,
+    pub cpu_id: LogicalCpuId,
 
     /// Context management
     pub switch_internals: ContextSwitchPercpu,

src/scheme/irq.rs

@@ -10,7 +10,7 @@ use spin::{Mutex, RwLock};
 
 use crate::arch::interrupt::{available_irqs_iter, bsp_apic_id, is_reserved, set_reserved};
 
-use crate::event;
+use crate::{event, LogicalCpuId};
 use crate::interrupt::irq::acknowledge;
 use crate::scheme::{AtomicSchemeId, SchemeId};
 use crate::syscall::data::Stat;
@@ -127,10 +127,10 @@ impl IrqScheme {
                 return Err(Error::new(EINVAL));
             }
             if flags & O_STAT == 0 {
-                if is_reserved(usize::from(cpu_id), irq_to_vector(irq_number)) {
+                if is_reserved(LogicalCpuId::new(cpu_id.into()), irq_to_vector(irq_number)) {
                     return Err(Error::new(EEXIST));
                 }
-                set_reserved(usize::from(cpu_id), irq_to_vector(irq_number), true);
+                set_reserved(LogicalCpuId::new(cpu_id.into()), irq_to_vector(irq_number), true);
             }
             Handle::Irq { ack: AtomicUsize::new(0), irq: irq_number }
         } else {
@@ -188,7 +188,7 @@ impl Scheme for IrqScheme {
                     let mut data = String::new();
                     use core::fmt::Write;
 
-                    for vector in available_irqs_iter(cpu_id.into()) {
+                    for vector in available_irqs_iter(LogicalCpuId::new(cpu_id.into())) {
                         let irq = vector_to_irq(vector);
                         if Some(u32::from(cpu_id)) == bsp_apic_id() && irq < BASE_IRQ_COUNT {
                             continue;
@@ -253,7 +253,7 @@ impl Scheme for IrqScheme {
 
         if let &Handle::Irq { irq: handle_irq, .. } = handle {
             if handle_irq > BASE_IRQ_COUNT {
-                set_reserved(0, irq_to_vector(handle_irq), false);
+                set_reserved(LogicalCpuId::BSP, irq_to_vector(handle_irq), false);
             }
         }
         Ok(0)

src/scheme/proc.rs

@@ -13,7 +13,7 @@ use crate::{
         flag::*,
         scheme::{CallerCtx, Scheme},
         self, usercopy::{UserSliceWo, UserSliceRo},
-    },
+    }, LogicalCpuId,
 };
 
 use alloc::{
@@ -862,7 +862,8 @@ impl KernelScheme for ProcScheme {
                 Ok(mem::size_of::<usize>())
             }
             Operation::SchedAffinity => {
-                buf.write_usize(context::contexts().get(info.pid).ok_or(Error::new(EBADFD))?.read().sched_affinity.map_or(usize::MAX, |a| a % crate::cpu_count()))?;
+                let id = context::contexts().get(info.pid).ok_or(Error::new(EBADFD))?.read().sched_affinity.map_or(u32::MAX, |a| a.get() % crate::cpu_count());
+                buf.write_usize(id as usize)?;
                 Ok(mem::size_of::<usize>())
             }
             // TODO: Replace write() with SYS_DUP_FORWARD.
@@ -1067,8 +1068,16 @@ impl KernelScheme for ProcScheme {
             }
             // TODO: Deduplicate code.
             Operation::SchedAffinity => {
-                let val = buf.read_usize()?;
-                context::contexts().get(info.pid).ok_or(Error::new(EBADFD))?.write().sched_affinity = if val == usize::MAX { None } else { Some(val % crate::cpu_count()) };
+                // TODO: read_u32
+                let val = u32::try_from(buf.read_usize()?).map_err(|_| Error::new(EINVAL))?;
+
+                context::contexts().get(info.pid)
+                    .ok_or(Error::new(EBADFD))?.write()
+                    .sched_affinity = if val == u32::MAX {
+                        None
+                    } else {
+                        Some(LogicalCpuId::new(val % crate::cpu_count()))
+                    };
                 Ok(mem::size_of::<usize>())
             }