starry-kernel 0.9.0

use core::mem::{align_of, size_of};

use ax_runtime::hal::{
    cpu::{UserAccessError, UserAtomicError, UserAtomicU32Op},
    time::monotonic_time,
};
use linux_raw_sys::general::{
    FUTEX_CLOCK_REALTIME, FUTEX_CMP_REQUEUE, FUTEX_OP_ADD, FUTEX_OP_ANDN, FUTEX_OP_CMP_EQ,
    FUTEX_OP_CMP_GE, FUTEX_OP_CMP_GT, FUTEX_OP_CMP_LE, FUTEX_OP_CMP_LT, FUTEX_OP_CMP_NE,
    FUTEX_OP_OPARG_SHIFT, FUTEX_OP_OR, FUTEX_OP_SET, FUTEX_OP_XOR, FUTEX_REQUEUE, FUTEX_WAIT,
    FUTEX_WAIT_BITSET, FUTEX_WAKE, FUTEX_WAKE_BITSET, FUTEX_WAKE_OP, robust_list_head, timespec,
};

use crate::{
    StarryError, StarryResult,
    mm::{
        VmMutPtr, VmPtr, atomic_update_user_u32_nofault, check_access, fault_in_user_u32_read,
        fault_in_user_u32_write, read_user_u32_nofault,
    },
    task::{
        FutexAccessError, FutexContext, FutexKeyMode, FutexWaitError, TidNumber, UserTaskRef,
        get_user_task_by_number,
    },
    time::{ClockDeadline, TimeValueLike},
};

const FUTEX_PRIVATE_FLAG: u32 = 128;
const FUTEX_COMMAND_MASK: u32 = FUTEX_PRIVATE_FLAG - 1;
const SUPPORTED_FLAGS: u32 = FUTEX_PRIVATE_FLAG | FUTEX_CLOCK_REALTIME;

#[derive(Clone, Copy, PartialEq, Eq)]
enum FutexCommand {
    Wait,
    Wake,
    WaitBitset,
    WakeBitset,
    Requeue,
    CmpRequeue,
    WakeOp,
}

struct ParsedFutexOp {
    command: FutexCommand,
    key_mode: FutexKeyMode,
    clock_realtime: bool,
}

#[derive(Clone, Copy)]
struct ParsedFutexWakeOp {
    operation: UserAtomicU32Op,
    argument: u32,
    comparison: u32,
    comparison_argument: i32,
}

fn assert_non_negative_i32(value: u32) -> crate::StarryResult<u32> {
    if (value as i32) < 0 {
        Err(StarryError::InvalidInput)
    } else {
        Ok(value)
    }
}

fn validate_futex_word(current: &UserTaskRef, uaddr: *const u32) -> crate::StarryResult<()> {
    if !uaddr.addr().is_multiple_of(align_of::<u32>()) {
        return Err(StarryError::InvalidInput);
    }
    loop {
        match futex_read_user_nofault(uaddr) {
            Ok(_) => return Ok(()),
            Err(FutexAccessError::UserFault) => {
                fault_in_user_u32_read(current, uaddr)?;
                crate::task::yield_now();
            }
            Err(FutexAccessError::Retry) => crate::task::yield_now(),
            Err(FutexAccessError::Operation(error)) => return Err(error.into()),
        }
    }
}

fn validate_futex_wake_key(
    current: &UserTaskRef,
    uaddr: *const u32,
    key_mode: FutexKeyMode,
) -> crate::StarryResult<()> {
    if matches!(key_mode, FutexKeyMode::Private) {
        // Linux private futex keys contain only current->mm and the virtual
        // address. Wake therefore performs access_ok() but never dereferences
        // the futex word or resolves its VMA.
        return check_access(uaddr.addr(), size_of::<u32>()).map_err(Into::into);
    }

    validate_futex_word(current, uaddr)
}

fn sign_extend_12(value: u32) -> i32 {
    ((value << 20) as i32) >> 20
}

fn futex_wake_op_arg(raw_op: u32, encoded_op: u32) -> i32 {
    let mut oparg = sign_extend_12((encoded_op >> 12) & 0xfff);
    if raw_op & FUTEX_OP_OPARG_SHIFT != 0 {
        oparg = (1u32 << ((oparg & 31) as u32)) as i32;
    }
    oparg
}

#[cfg(all(test, not(axtest)))]
fn apply_futex_wake_op(old_value: u32, raw_op: u32, oparg: i32) -> StarryResult<u32> {
    let op = raw_op & !FUTEX_OP_OPARG_SHIFT;
    let new_value = match op {
        FUTEX_OP_SET => oparg as u32,
        FUTEX_OP_ADD => (old_value as i32).wrapping_add(oparg) as u32,
        FUTEX_OP_OR => old_value | oparg as u32,
        FUTEX_OP_ANDN => old_value & !(oparg as u32),
        FUTEX_OP_XOR => old_value ^ oparg as u32,
        _ => return Err(StarryError::Unsupported),
    };
    Ok(new_value)
}

fn compare_futex_wake_op(old_value: u32, raw_cmp: u32, cmparg: i32) -> StarryResult<bool> {
    let old_value = old_value as i32;
    let matched = match raw_cmp {
        FUTEX_OP_CMP_EQ => old_value == cmparg,
        FUTEX_OP_CMP_NE => old_value != cmparg,
        FUTEX_OP_CMP_LT => old_value < cmparg,
        FUTEX_OP_CMP_LE => old_value <= cmparg,
        FUTEX_OP_CMP_GT => old_value > cmparg,
        FUTEX_OP_CMP_GE => old_value >= cmparg,
        _ => return Err(StarryError::Unsupported),
    };
    Ok(matched)
}

fn parse_futex_wake_op(encoded_op: u32) -> crate::StarryResult<ParsedFutexWakeOp> {
    let raw_op = (encoded_op >> 28) & 0xf;
    let raw_cmp = (encoded_op >> 24) & 0xf;
    let oparg = futex_wake_op_arg(raw_op, encoded_op);
    let cmparg = sign_extend_12(encoded_op & 0xfff);

    let operation = match raw_op & !FUTEX_OP_OPARG_SHIFT {
        FUTEX_OP_SET => UserAtomicU32Op::Set,
        FUTEX_OP_ADD => UserAtomicU32Op::Add,
        FUTEX_OP_OR => UserAtomicU32Op::Or,
        FUTEX_OP_ANDN => UserAtomicU32Op::AndNot,
        FUTEX_OP_XOR => UserAtomicU32Op::Xor,
        _ => return Err(crate::StarryError::Unsupported),
    };
    match raw_cmp {
        FUTEX_OP_CMP_EQ | FUTEX_OP_CMP_NE | FUTEX_OP_CMP_LT | FUTEX_OP_CMP_LE | FUTEX_OP_CMP_GT
        | FUTEX_OP_CMP_GE => {}
        _ => return Err(crate::StarryError::Unsupported),
    }
    Ok(ParsedFutexWakeOp {
        operation,
        argument: oparg as u32,
        comparison: raw_cmp,
        comparison_argument: cmparg,
    })
}

fn futex_atomic_op_in_user_nofault(
    uaddr: *mut u32,
    operation: ParsedFutexWakeOp,
) -> Result<bool, FutexAccessError> {
    let old_value = atomic_update_user_u32_nofault(uaddr, operation.operation, operation.argument)
        .map_err(|error| match error {
            UserAtomicError::Fault => FutexAccessError::UserFault,
            UserAtomicError::Retry => FutexAccessError::Retry,
        })?;
    compare_futex_wake_op(
        old_value,
        operation.comparison,
        operation.comparison_argument,
    )
    .map_err(Into::into)
}

fn futex_read_user_nofault(uaddr: *const u32) -> Result<u32, FutexAccessError> {
    read_user_u32_nofault(uaddr).map_err(|error| match error {
        UserAccessError::Fault => FutexAccessError::UserFault,
    })
}

fn apply_wake_op_without_waiters(
    current: &UserTaskRef,
    uaddr: *mut u32,
    operation: ParsedFutexWakeOp,
) -> crate::StarryResult<()> {
    loop {
        match futex_atomic_op_in_user_nofault(uaddr, operation) {
            Ok(_) => return Ok(()),
            Err(FutexAccessError::UserFault) => {
                fault_in_user_u32_write(current, uaddr)?;
                crate::task::yield_now();
            }
            Err(FutexAccessError::Retry) => crate::task::yield_now(),
            Err(FutexAccessError::Operation(error)) => return Err(error.into()),
        }
    }
}

fn parse_futex_op(futex_op: u32) -> StarryResult<ParsedFutexOp> {
    let flags = futex_op & !FUTEX_COMMAND_MASK;
    if flags & !SUPPORTED_FLAGS != 0 {
        return Err(StarryError::InvalidInput);
    }

    let command = match futex_op & FUTEX_COMMAND_MASK {
        FUTEX_WAIT => FutexCommand::Wait,
        FUTEX_WAKE => FutexCommand::Wake,
        FUTEX_WAIT_BITSET => FutexCommand::WaitBitset,
        FUTEX_WAKE_BITSET => FutexCommand::WakeBitset,
        FUTEX_REQUEUE => FutexCommand::Requeue,
        FUTEX_CMP_REQUEUE => FutexCommand::CmpRequeue,
        FUTEX_WAKE_OP => FutexCommand::WakeOp,
        _ => return Err(StarryError::Unsupported),
    };

    let clock_realtime = flags & FUTEX_CLOCK_REALTIME != 0;
    if clock_realtime && command != FutexCommand::WaitBitset {
        return Err(StarryError::Unsupported);
    }

    let key_mode = if flags & FUTEX_PRIVATE_FLAG != 0 {
        FutexKeyMode::Private
    } else {
        FutexKeyMode::Auto
    };

    Ok(ParsedFutexOp {
        command,
        key_mode,
        clock_realtime,
    })
}

fn futex_wait_timeout(
    current: &UserTaskRef,
    op: &ParsedFutexOp,
    timeout: *const timespec,
) -> crate::StarryResult<Option<ClockDeadline>> {
    let Some(ts) = timeout.nullable() else {
        return Ok(None);
    };

    let timeout = unsafe { ts.vm_read_uninit(current)?.assume_init() }.try_into_time_value()?;
    // Resolve relative waits once, before any user-fault or notification retry.
    let deadline = if op.command == FutexCommand::Wait {
        ClockDeadline::Monotonic(monotonic_time().saturating_add(timeout))
    } else if op.clock_realtime {
        ClockDeadline::Realtime(timeout)
    } else {
        ClockDeadline::Monotonic(timeout)
    };
    Ok(Some(deadline))
}

fn complete_futex_wake(count: usize) -> crate::StarryResult<isize> {
    // Waker publication makes the target runnable and lets ax-task set the
    // owner CPU's sticky preemption state. Mirroring Linux wake_q completion,
    // the futex syscall must not add a second, unconditional scheduling point.
    Ok(count as isize)
}

pub fn sys_futex(
    current: &UserTaskRef,
    uaddr: *const u32,
    futex_op: u32,
    value: u32,
    timeout: *const timespec,
    uaddr2: *mut u32,
    value3: u32,
) -> StarryResult<isize> {
    debug!(
        "sys_futex <= uaddr: {uaddr:?}, futex_op: {futex_op}, value: {value}, uaddr2: {uaddr2:?}, \
         value3: {value3}",
    );

    let op = parse_futex_op(futex_op)?;
    if !uaddr.addr().is_multiple_of(align_of::<u32>()) {
        return Err(StarryError::InvalidInput);
    }
    if matches!(
        op.command,
        FutexCommand::WaitBitset | FutexCommand::WakeBitset
    ) && value3 == 0
    {
        return Err(StarryError::InvalidInput);
    }

    match op.command {
        FutexCommand::Wait | FutexCommand::WaitBitset => {
            let deadline = futex_wait_timeout(current, &op, timeout)?;

            let bitset = if op.command == FutexCommand::WaitBitset {
                value3
            } else {
                u32::MAX
            };
            let context = FutexContext::new(current);

            loop {
                let futex = context.resolve(uaddr.addr(), op.key_mode);
                match futex.wait_nofault_until(context.task(), bitset, deadline, || {
                    futex_read_user_nofault(uaddr).map(|observed| observed == value)
                }) {
                    Ok(true) => break,
                    Ok(false) => return Err(crate::StarryError::WouldBlock),
                    Err(FutexWaitError::SchedulerNotification) => continue,
                    Err(FutexWaitError::Access(FutexAccessError::UserFault)) => {
                        fault_in_user_u32_read(current, uaddr)?;
                        crate::task::yield_now();
                    }
                    Err(FutexWaitError::Access(FutexAccessError::Retry)) => {
                        crate::task::yield_now()
                    }
                    Err(FutexWaitError::Access(FutexAccessError::Operation(error))) => {
                        return Err(error.into());
                    }
                }
            }

            Ok(0)
        }
        FutexCommand::Wake | FutexCommand::WakeBitset => {
            let wake_count = assert_non_negative_i32(value)? as usize;
            validate_futex_wake_key(current, uaddr, op.key_mode)?;

            let futex = FutexContext::new(current).resolve(uaddr.addr(), op.key_mode);
            let bitset = if op.command == FutexCommand::WakeBitset {
                value3
            } else {
                u32::MAX
            };
            let count = futex.wake(wake_count, bitset);
            complete_futex_wake(count)
        }
        FutexCommand::Requeue | FutexCommand::CmpRequeue => {
            let wake_count = assert_non_negative_i32(value)? as usize;
            let requeue_count = assert_non_negative_i32(timeout.addr() as u32)? as usize;
            if op.command == FutexCommand::Requeue {
                validate_futex_word(current, uaddr)?;
            }
            validate_futex_word(current, uaddr2)?;
            let context = FutexContext::new(current);

            let count = loop {
                let (source_futex, target_futex) =
                    context.resolve_pair(uaddr.addr(), uaddr2.addr(), op.key_mode);
                match source_futex.requeue_to(
                    &target_futex,
                    wake_count,
                    u32::MAX,
                    requeue_count,
                    || {
                        if op.command == FutexCommand::CmpRequeue {
                            futex_read_user_nofault(uaddr).map(|observed| observed == value3)
                        } else {
                            Ok(true)
                        }
                    },
                ) {
                    Ok(Some(count)) => break count,
                    Ok(None) => return Err(crate::StarryError::WouldBlock),
                    Err(FutexAccessError::UserFault) => {
                        fault_in_user_u32_read(current, uaddr)?;
                        crate::task::yield_now();
                    }
                    Err(FutexAccessError::Retry) => crate::task::yield_now(),
                    Err(FutexAccessError::Operation(error)) => return Err(error.into()),
                }
            };

            complete_futex_wake(count)
        }
        FutexCommand::WakeOp => {
            let wake_count = value as usize;
            let wake2_count = timeout.addr();
            validate_futex_word(current, uaddr)?;
            if !uaddr2.addr().is_multiple_of(align_of::<u32>()) {
                return Err(crate::StarryError::InvalidInput);
            }
            let wake_operation = parse_futex_wake_op(value3)?;

            let count = if wake_count == 0 && wake2_count == 0 {
                // No waiter state can change when both wake limits are zero.
                // The user RMW is already atomic, so taking futex table locks
                // would add PI contention without protecting any kernel data.
                apply_wake_op_without_waiters(current, uaddr2, wake_operation)?;
                0
            } else {
                let context = FutexContext::new(current);
                loop {
                    // Shared keys depend on the current VMA backing and must be
                    // recomputed after fault-in, matching Linux futex retry.
                    let (source, target) =
                        context.resolve_pair(uaddr.addr(), uaddr2.addr(), op.key_mode);
                    match source.wake_op(wake_count, &target, wake2_count, || {
                        futex_atomic_op_in_user_nofault(uaddr2, wake_operation)
                    }) {
                        Ok(count) => break count,
                        Err(FutexAccessError::UserFault) => {
                            fault_in_user_u32_write(current, uaddr2)?;
                            crate::task::yield_now();
                        }
                        Err(FutexAccessError::Retry) => crate::task::yield_now(),
                        Err(FutexAccessError::Operation(error)) => return Err(error.into()),
                    }
                }
            };

            complete_futex_wake(count)
        }
    }
}

pub fn sys_get_robust_list(
    current: &UserTaskRef,
    tid: u32,
    head: *mut *const robust_list_head,
    size: *mut usize,
) -> StarryResult<isize> {
    let task = if tid == 0 {
        current.clone()
    } else {
        get_user_task_by_number(TidNumber::try_from(tid)?)?
    };
    head.vm_write(current, task.as_thread().robust_list_head() as _)?;
    size.vm_write(current, size_of::<robust_list_head>())?;

    Ok(0)
}

pub fn sys_set_robust_list(
    current: &crate::task::UserTaskRef,
    head: *const robust_list_head,
    size: usize,
) -> crate::StarryResult<isize> {
    if size != size_of::<robust_list_head>() {
        return Err(StarryError::InvalidInput);
    }
    current.as_thread().set_robust_list_head(head.addr());

    Ok(0)
}

#[cfg(all(test, not(axtest)))]
fn futex_op_and_compare_rules_hold_for_test() -> bool {
    // sign_extend_12: sign-extends a 12-bit value.
    assert!(sign_extend_12(0x000) == 0);
    assert!(sign_extend_12(0x7FF) == 2047); // max positive
    assert!(sign_extend_12(0x800) == -2048); // min negative
    assert!(sign_extend_12(0xFFF) == -1);

    // futex_wake_op_arg: extracts oparg from encoded_op, optionally shifts.
    let raw_op_set = FUTEX_OP_SET;
    let encoded_no_shift = (5u32) << 12; // oparg=5, no shift
    assert!(futex_wake_op_arg(raw_op_set, encoded_no_shift) == 5);

    let raw_op_shift = FUTEX_OP_SET | FUTEX_OP_OPARG_SHIFT;
    let encoded_shift = (3u32) << 12; // oparg=3, shift by 3
    assert!(futex_wake_op_arg(raw_op_shift, encoded_shift) == 8); // 1 << 3 = 8

    // apply_futex_wake_op: applies the operation to old_value.
    assert!(apply_futex_wake_op(10, FUTEX_OP_SET, 42).unwrap() == 42);
    assert!(apply_futex_wake_op(10, FUTEX_OP_ADD, 5).unwrap() == 15);
    assert!(apply_futex_wake_op(0b1100, FUTEX_OP_OR, 0b1010).unwrap() == 0b1110);
    assert!(apply_futex_wake_op(0xFF, FUTEX_OP_ANDN, 0x0F).unwrap() == 0xF0);
    assert!(apply_futex_wake_op(0xAA, FUTEX_OP_XOR, 0xFF).unwrap() == 0x55);
    assert!(apply_futex_wake_op(0, 0xFFFF, 0).is_err()); // unsupported op

    // compare_futex_wake_op: compares old_value with cmparg.
    assert!(compare_futex_wake_op(5, FUTEX_OP_CMP_EQ, 5).unwrap());
    assert!(!compare_futex_wake_op(5, FUTEX_OP_CMP_EQ, 6).unwrap());
    assert!(compare_futex_wake_op(5, FUTEX_OP_CMP_NE, 6).unwrap());
    assert!(compare_futex_wake_op(5, FUTEX_OP_CMP_LT, 10).unwrap());
    assert!(compare_futex_wake_op(5, FUTEX_OP_CMP_LE, 5).unwrap());
    assert!(compare_futex_wake_op(5, FUTEX_OP_CMP_GT, 3).unwrap());
    assert!(compare_futex_wake_op(5, FUTEX_OP_CMP_GE, 5).unwrap());
    assert!(compare_futex_wake_op(0, 0xFFFF, 0).is_err()); // unsupported cmp

    true
}

#[cfg(all(test, axtest))]
fn futex_wake_completion_is_scheduler_driven_for_test() -> bool {
    crate::task::reset_yield_now_calls_for_test();
    let result = complete_futex_wake(1);
    matches!(result, Ok(1)) && crate::task::yield_now_calls_for_test() == 0
}

#[cfg(test)]
mod tests {
    #[cfg(not(axtest))]
    #[test]
    fn futex_op_and_compare_rules_hold() {
        assert!(super::futex_op_and_compare_rules_hold_for_test());
    }

    #[cfg(axtest)]
    #[axtest::axtest]
    fn futex_wake_completion_is_scheduler_driven() {
        assert!(super::futex_wake_completion_is_scheduler_driven_for_test());
    }
}