bsp_define 0.1.54

//! 定义多核启动相关结构

use core::{
    mem::offset_of,
    sync::atomic::{AtomicUsize, Ordering},
};

/// 系统最大支持 cpu 数量
pub const NR_CPUS: usize = 16;

/// psci 兼容性检测
#[derive(Clone, Copy)]
pub enum PsciCompatible {
    /// arm,psci
    Psci0_1,
    /// arm,psci-0.2
    Psci0_2,
    /// arm,psci-1.0
    Psci1_0,
}

/// psci 调用方式
#[derive(Clone, Copy)]
pub enum PsciEnableMethod {
    /// hvc
    Hvc,
    /// smc
    Smc,
}

/// 多核启动定义
#[derive(Clone, Copy)]
pub struct SmpDefine {
    compatible: PsciCompatible,
    enable_method: PsciEnableMethod,
    cpu_on: usize,
    hwid: [usize; NR_CPUS],
    nr_cpus: usize,
}

impl SmpDefine {
    /// 创建多核定义
    pub const fn create(
        compatible: PsciCompatible,
        enable_method: PsciEnableMethod,
        cpu_on: usize,
        hwid: [usize; NR_CPUS],
        nr_cpus: usize,
    ) -> Self {
        Self { compatible, enable_method, cpu_on, hwid, nr_cpus }
    }

    /// psci兼容性
    #[inline(always)]
    pub fn compatible(&self) -> PsciCompatible {
        self.compatible
    }

    /// psci调用方式
    #[inline(always)]
    pub fn method(&self) -> PsciEnableMethod {
        self.enable_method
    }

    /// `psci_0.1`版本特有`cpu_on function id`
    #[inline(always)]
    pub fn cpu_on_id(&self) -> usize {
        self.cpu_on
    }

    /// cpu id 对应的 hw id
    #[inline(always)]
    pub fn hwid(&self, cpu: usize) -> usize {
        self.hwid[cpu]
    }

    /// 总的 cpu 数量
    #[inline(always)]
    pub fn nr_cpus(&self) -> usize {
        self.nr_cpus
    }
}

/// cpumask
#[derive(Default)]
pub struct Cpumask {
    map: AtomicUsize,
}

impl Cpumask {
    /// 构造 cpumask
    pub const fn new() -> Self {
        Self { map: AtomicUsize::new(0) }
    }

    /// 拷贝一份cpumask
    pub fn copy(&self) -> Self {
        let val = self.map.load(Ordering::Relaxed);
        Self { map: AtomicUsize::new(val) }
    }

    /// 设置一个 cpu bit
    ///
    /// # Panics
    /// 如果设置 cpu 超过最大 cpu 限制将会 panic.
    #[inline(always)]
    pub fn set(&self, cpu: usize) {
        assert!(cpu < NR_CPUS);
        self.map.fetch_or(1 << cpu, Ordering::Relaxed);
    }

    /// 清除一个 cpu bit
    ///
    /// # Panics
    /// 如果清除 cpu 超过最大 cpu 限制将会 panic.
    #[inline(always)]
    pub fn clear(&self, cpu: usize) {
        assert!(cpu < NR_CPUS);
        self.map.fetch_and(!(1 << cpu), Ordering::Relaxed);
    }

    /// cpu 数量
    #[inline(always)]
    pub fn nr_cpus(&self) -> u32 {
        self.map.load(Ordering::Relaxed).count_ones()
    }

    /// 测试 cpu bit 是否设置
    #[inline(always)]
    pub fn test(&self, cpu: usize) -> bool {
        if cpu >= NR_CPUS {
            return false;
        }
        (self.map.load(Ordering::Relaxed) & (1 << cpu)) != 0
    }

    /// 测试并且清零 cpu bit
    #[inline(always)]
    pub fn test_and_clear_bit(&self, cpu: usize) -> bool {
        if cpu >= NR_CPUS {
            return false;
        }
        let old = self.map.load(Ordering::Relaxed);
        self.clear(cpu);
        (old & (1 << cpu)) != 0
    }
}

/// cpumask 迭代器
pub struct CpumaskIter<'a> {
    pos: usize,
    map: &'a Cpumask,
}

impl<'a> CpumaskIter<'a> {
    fn new(map: &'a Cpumask) -> Self {
        Self { pos: 0, map }
    }
}

impl Cpumask {
    /// 创建迭代器
    pub fn iter(&self) -> CpumaskIter<'_> {
        CpumaskIter::new(self)
    }
}

impl<'a> Iterator for CpumaskIter<'a> {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        for i in self.pos..NR_CPUS {
            if self.map.map.load(Ordering::Relaxed) & (1 << i) != 0 {
                self.pos = i + 1;
                return Some(i);
            }
        }
        None
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test() {
        let cpu = Cpumask::new();
        cpu.set(1);
        cpu.set(3);
        cpu.set(7);

        assert_eq!(cpu.nr_cpus(), 3);

        let check = [1, 3, 7];
        let mut j = 0;
        for i in cpu.iter() {
            assert_eq!(i, check[j]);
            j += 1;
        }

        cpu.set(10);
        cpu.clear(3);
        let check = [1, 7, 10];
        let mut j = 0;
        for i in cpu.iter() {
            assert_eq!(i, check[j]);
            j += 1;
        }
    }
}

/// 从核task数据
#[derive(Default)]
#[repr(C)]
pub struct SecondaryData {
    stack: AtomicUsize,
    task: AtomicUsize,
}

impl SecondaryData {
    /// 构造
    pub const fn new() -> Self {
        Self { stack: AtomicUsize::new(0), task: AtomicUsize::new(0) }
    }

    /// 设置stack
    pub fn set_stack(&self, base: usize) {
        self.stack.store(base, Ordering::Relaxed);
    }

    /// 设置task
    pub fn set_task(&self, task: usize) {
        self.task.store(task, Ordering::Relaxed);
    }
}

/// `CPU_BOOT_STACK`
pub const CPU_BOOT_STACK: usize = offset_of!(SecondaryData, stack);
/// `CPU_BOOT_TASK`
pub const CPU_BOOT_TASK: usize = offset_of!(SecondaryData, task);