cpu-local 0.4.2

use core::cell::Cell;

use super::*;

pub(super) const CURRENT_MODEL: ArchitectureCurrentModel = ArchitectureCurrentModel {
    linux_current: CurrentContextSource::ArchitectureRegister,
    unikernel_tls: CurrentContextSource::RuntimeAnchor,
};

pub(super) struct Backend;

impl ArchitectureRegisterBackend for Backend {}

std::thread_local! {
    static CPU_BASE: Cell<usize> = const { Cell::new(0) };
    static ARCHITECTURE_CURRENT: Cell<usize> = const { Cell::new(0) };
    static KERNEL_TLS: Cell<usize> = const { Cell::new(0) };
    static CPU_BASE_READS: Cell<usize> = const { Cell::new(0) };
    static CURRENT_CONTEXT_READS: Cell<usize> = const { Cell::new(0) };
    static BINDING_OBSERVATIONS: Cell<usize> = const { Cell::new(0) };
    static INITIALIZED_AREA_VALIDATIONS: Cell<usize> = const { Cell::new(0) };
    #[cfg(test)]
    static MIGRATION_TARGET: Cell<usize> = const { Cell::new(0) };
}

pub(super) fn validate_environment() -> Result<(), CpuLocalError> {
    Ok(())
}

pub(super) unsafe fn install_cpu_base(area_base: usize, boot_context: usize) {
    CPU_BASE.set(area_base);
    ARCHITECTURE_CURRENT.set(if cfg!(kernel_tls) { 0 } else { boot_context });
}

pub(super) unsafe fn read_cpu_base() -> Result<usize, CpuLocalError> {
    CPU_BASE_READS.set(CPU_BASE_READS.get().wrapping_add(1));
    Ok(CPU_BASE.get())
}

pub(super) unsafe fn read_current_context(area_base: usize) -> usize {
    CURRENT_CONTEXT_READS.set(CURRENT_CONTEXT_READS.get().wrapping_add(1));
    #[cfg(test)]
    MIGRATION_TARGET.with(|target| {
        let target = target.replace(0);
        if target != 0 {
            CPU_BASE.set(target);
        }
    });
    if cfg!(kernel_tls) {
        // SAFETY: the shared caller validated the sampled shutdown-lifetime
        // CPU area before asking the host backend for its selected source.
        return unsafe { area_runtime_anchor(area_base) }.current_context_raw();
    }
    ARCHITECTURE_CURRENT
        .with(|current| {
            let current = current.get();
            (current != 0).then_some(current)
        })
        .unwrap_or(0)
}

pub(super) unsafe fn write_current_context(value: usize) {
    ARCHITECTURE_CURRENT.set(value);
}

#[cfg(kernel_tls)]
pub(super) unsafe fn read_kernel_tls() -> usize {
    KERNEL_TLS.get()
}

#[cfg(kernel_tls)]
pub(super) unsafe fn write_kernel_tls(value: usize) {
    KERNEL_TLS.set(value);
}

pub(super) fn reset_register_read_counts() {
    CPU_BASE_READS.set(0);
    CURRENT_CONTEXT_READS.set(0);
    BINDING_OBSERVATIONS.set(0);
    INITIALIZED_AREA_VALIDATIONS.set(0);
}

pub(super) fn register_read_counts() -> super::host_test::RegisterReadCounts {
    super::host_test::RegisterReadCounts {
        cpu_base: CPU_BASE_READS.get(),
        current_context: CURRENT_CONTEXT_READS.get(),
        binding_observations: BINDING_OBSERVATIONS.get(),
        initialized_area_validations: INITIALIZED_AREA_VALIDATIONS.get(),
    }
}

pub(super) fn record_binding_observation() {
    BINDING_OBSERVATIONS.set(BINDING_OBSERVATIONS.get().wrapping_add(1));
}

pub(super) fn record_initialized_area_validation() {
    INITIALIZED_AREA_VALIDATIONS.set(INITIALIZED_AREA_VALIDATIONS.get().wrapping_add(1));
}

#[cfg(test)]
pub(super) fn migrate_on_next_current_read(area_base: usize) {
    MIGRATION_TARGET.set(area_base);
}

#[cfg(all(test, not(kernel_tls)))]
pub(super) fn set_architecture_current(current: usize) {
    ARCHITECTURE_CURRENT.set(current);
}

unsafe fn area_runtime_anchor(area_base: usize) -> &'static crate::CpuRuntimeAnchor {
    // SAFETY: forwarded caller contract supplies a validated CPU-area base.
    unsafe {
        &*((area_base + crate::CPU_AREA_RUNTIME_ANCHOR_OFFSET) as *const crate::CpuRuntimeAnchor)
    }
}