pw_kernel/arch/riscv/disable_interrupts_atomic.rs

// Copyright 2025 The Pigweed Authors
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not
// use this file except in compliance with the License. You may obtain a copy of
// the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations under
// the License.

use core::cell::UnsafeCell;
use core::sync::atomic::Ordering;

use pw_atomic::{
    Atomic, AtomicAdd, AtomicCompareExchange, AtomicFalse, AtomicLoad, AtomicNew, AtomicStore,
    AtomicSub, AtomicZero,
};

// TODO: once the interrupt controller work lands, make use of it here.
unsafe fn disable_interrupts() -> usize {
    let mut val: usize;
    unsafe { core::arch::asm!("csrrw {val}, mstatus, x0", val = out(reg) val) };
    val
}

unsafe fn enabled_interrupts(val: usize) {
    unsafe { core::arch::asm!("csrrw x0, mstatus, {val}", val = in(reg) val) };
}

macro_rules! impl_for_scalar {
    ($atomic_type:ty, $primitive_type:ty) => {
        impl AtomicLoad<$primitive_type> for $atomic_type {
            fn load(&self, _ordering: Ordering) -> $primitive_type {
                unsafe { *self.value.get() }
            }
        }

        impl AtomicStore<$primitive_type> for $atomic_type {
            fn store(&self, val: $primitive_type, _ordering: Ordering) {
                unsafe {
                    *self.value.get() = val;
                }
            }
        }

        impl AtomicCompareExchange<$primitive_type> for $atomic_type {
            fn compare_exchange(
                &self,
                current: $primitive_type,
                new: $primitive_type,
                _success: Ordering,
                _failure: Ordering,
            ) -> Result<$primitive_type, $primitive_type> {
                unsafe {
                    let i = disable_interrupts();
                    let ptr = self.value.get();
                    if *ptr == current {
                        *ptr = new;
                        enabled_interrupts(i);
                        Ok(current)
                    } else {
                        enabled_interrupts(i);
                        Err(current)
                    }
                }
            }
        }

        impl AtomicNew<$primitive_type> for $atomic_type {
            fn new(val: $primitive_type) -> Self {
                Self {
                    value: UnsafeCell::new(val),
                }
            }
        }

        impl Atomic<$primitive_type> for $atomic_type {}
    };
}

macro_rules! impl_for_numeric {
    ($atomic_type:ty, $primitive_type:ty) => {
        impl_for_scalar!($atomic_type, $primitive_type);

        impl AtomicAdd<$primitive_type> for $atomic_type {
            fn fetch_add(&self, val: $primitive_type, _ordering: Ordering) -> $primitive_type {
                unsafe {
                    let i = disable_interrupts();
                    let ptr = self.value.get();
                    let current = *ptr;
                    // Wrapping semantics per
                    // https://doc.rust-lang.org/std/sync/atomic/struct.AtomicUsize.html#method.fetch_add
                    *ptr = current.wrapping_add(val);
                    enabled_interrupts(i);
                    current
                }
            }
        }

        impl AtomicSub<$primitive_type> for $atomic_type {
            fn fetch_sub(&self, val: $primitive_type, _ordering: Ordering) -> $primitive_type {
                unsafe {
                    let i = disable_interrupts();
                    let ptr = self.value.get();
                    let current = *ptr;
                    // Wrapping semantics per
                    // https://doc.rust-lang.org/std/sync/atomic/struct.AtomicUsize.html#method.fetch_sub
                    *ptr = current.wrapping_sub(val);
                    enabled_interrupts(i);
                    current
                }
            }
        }

        impl AtomicZero for $atomic_type {
            const ZERO: Self = Self {
                value: UnsafeCell::new(0),
            };
        }
    };
}

pub struct AtomicBool {
    value: UnsafeCell<bool>,
}

impl AtomicBool {
    pub const fn new(val: bool) -> Self {
        Self {
            value: UnsafeCell::new(val),
        }
    }
}

impl AtomicFalse for AtomicBool {
    const FALSE: Self = Self::new(false);
}
impl_for_scalar!(AtomicBool, bool);
impl pw_atomic::AtomicBool for AtomicBool {}

// SAFETY: Atomicity is guaranteed by disabling interrupts on a UP system.
// MP systems are not supported.
unsafe impl Sync for AtomicBool {}

pub struct AtomicUsize {
    value: UnsafeCell<usize>,
}

impl_for_numeric!(AtomicUsize, usize);
impl pw_atomic::AtomicUsize for AtomicUsize {}

// SAFETY: Atomicity is guaranteed by disabling interrupts on a UP system.
// MP systems are not supported.
unsafe impl Sync for AtomicUsize {}