starry-kernel 0.8.1

use core::mem::align_of;

use ax_runtime::hal::{
    cpu::{UserAccessError, UserAtomicError, UserAtomicU32Op},
    time::{TimeValue, monotonic_time, wall_time},
};
use ax_task::current;
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 starry_vm::{VmMutPtr, VmPtr};

use crate::{
    StarryError, StarryResult,
    mm::{
        atomic_update_user_u32_nofault, fault_in_user_u32_read, fault_in_user_u32_write,
        read_user_u32_nofault,
    },
    task::{
        AsThread, FutexAccessError, FutexKey, FutexKeyMode, TidNumber, futex_table_for,
        get_user_task_by_number, retry_futex_nofault,
    },
    time::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,
}

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

fn validate_futex_word(uaddr: *const u32) -> StarryResult<()> {
    if !uaddr.addr().is_multiple_of(align_of::<u32>()) {
        return Err(StarryError::InvalidInput);
    }
    uaddr.vm_read()?;
    Ok(())
}

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(test)]
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)
}

struct ParsedFutexWakeOp {
    operation: UserAtomicU32Op,
    argument: u32,
    raw_cmp: u32,
    cmparg: i32,
}

fn parse_futex_wake_op(encoded_op: u32) -> StarryResult<ParsedFutexWakeOp> {
    let raw_op = (encoded_op >> 28) & 0xf;
    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(StarryError::Unsupported),
    };
    let raw_cmp = (encoded_op >> 24) & 0xf;
    if !matches!(
        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(StarryError::Unsupported);
    }
    Ok(ParsedFutexWakeOp {
        operation,
        argument: futex_wake_op_arg(raw_op, encoded_op) as u32,
        raw_cmp,
        cmparg: sign_extend_12(encoded_op & 0xfff),
    })
}

fn futex_atomic_op_in_user_nofault(
    uaddr: *mut u32,
    operation: &ParsedFutexWakeOp,
) -> Result<bool, FutexAccessError> {
    if !uaddr.addr().is_multiple_of(align_of::<u32>()) {
        return Err(FutexAccessError::Operation(StarryError::InvalidInput));
    }

    let old_value =
        match atomic_update_user_u32_nofault(uaddr, operation.operation, operation.argument) {
            Ok(value) => value,
            Err(UserAtomicError::Fault) => return Err(FutexAccessError::Fault),
            Err(UserAtomicError::Retry) => return Err(FutexAccessError::Retry),
        };
    compare_futex_wake_op(old_value, operation.raw_cmp, operation.cmparg)
        .map_err(FutexAccessError::Operation)
}

fn futex_value_matches_nofault(uaddr: *const u32, expected: u32) -> Result<bool, FutexAccessError> {
    match read_user_u32_nofault(uaddr) {
        Ok(value) => Ok(value == expected),
        Err(UserAccessError::Fault) => Err(FutexAccessError::Fault),
    }
}

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::WakeOp {
        return Err(StarryError::Unsupported);
    }
    if clock_realtime && !matches!(command, FutexCommand::Wait | FutexCommand::WaitBitset) {
        return Err(StarryError::InvalidInput);
    }

    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(
    op: &ParsedFutexOp,
    timeout: *const timespec,
) -> StarryResult<Option<TimeValue>> {
    let Some(ts) = timeout.nullable() else {
        return Ok(None);
    };

    let timeout = unsafe { ts.vm_read_uninit()?.assume_init() }.try_into_time_value()?;
    // FUTEX_WAIT keeps the traditional relative timeout. FUTEX_WAIT_BITSET
    // uses an absolute deadline on the selected clock.
    if op.command == FutexCommand::Wait {
        return Ok(Some(timeout));
    }

    let now = if op.clock_realtime {
        wall_time()
    } else {
        monotonic_time()
    };

    Ok(Some(timeout.saturating_sub(now)))
}

pub fn sys_futex(
    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);
    }

    let key = FutexKey::new_current(uaddr.addr(), op.key_mode);

    let futex_table = futex_table_for(&key);

    match op.command {
        FutexCommand::Wait | FutexCommand::WaitBitset => {
            // Fast path
            if uaddr.vm_read()? != value {
                return Err(StarryError::WouldBlock);
            }

            let timeout = futex_wait_timeout(&op, timeout)?;

            let futex = futex_table.get_or_insert(&key);
            let cleanup = futex_table.cleanup_for(&key);

            let bitset = if op.command == FutexCommand::WaitBitset {
                value3
            } else {
                u32::MAX
            };

            let waited = retry_futex_nofault(
                || {
                    futex.wq.wait_if_with_cleanup_nofault(
                        bitset,
                        timeout,
                        Some(cleanup.clone()),
                        || futex_value_matches_nofault(uaddr, value),
                    )
                },
                || fault_in_user_u32_read(uaddr),
            )?;
            if !waited {
                return Err(StarryError::WouldBlock);
            }

            Ok(0)
        }
        FutexCommand::Wake | FutexCommand::WakeBitset => {
            let wake_count = assert_non_negative_i32(value)? as usize;
            validate_futex_word(uaddr)?;

            let futex = futex_table.get(&key);
            let mut count = 0;
            if let Some(futex) = futex {
                let bitset = if op.command == FutexCommand::WakeBitset {
                    value3
                } else {
                    u32::MAX
                };
                count = futex.wq.wake(wake_count, bitset);
            }
            ax_task::yield_now();
            Ok(count as _)
        }
        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(uaddr)?;
            }
            validate_futex_word(uaddr2)?;

            let key2 = FutexKey::new_current(uaddr2.addr(), op.key_mode);
            let table2 = futex_table_for(&key2);
            let target = table2.get_or_insert(&key2);
            let target_cleanup = table2.cleanup_for(&key2);

            let count = retry_futex_nofault(
                || {
                    let Some(source) = futex_table.get(&key) else {
                        return if op.command == FutexCommand::CmpRequeue {
                            futex_value_matches_nofault(uaddr, value3)
                                .map(|matches| if matches { Some(0) } else { None })
                        } else {
                            Ok(Some(0))
                        };
                    };

                    source.wq.wake_requeue_if(
                        wake_count,
                        u32::MAX,
                        requeue_count,
                        target_cleanup.clone(),
                        &target.wq,
                        || {
                            if op.command == FutexCommand::CmpRequeue {
                                futex_value_matches_nofault(uaddr, value3)
                            } else {
                                Ok(true)
                            }
                        },
                    )
                },
                || fault_in_user_u32_read(uaddr),
            )?;

            let Some(count) = count else {
                return Err(StarryError::WouldBlock);
            };

            if count > 0 {
                ax_task::yield_now();
            }
            Ok(count as _)
        }
        FutexCommand::WakeOp => {
            let wake_count = value as usize;
            let wake2_count = timeout.addr();
            let operation = parse_futex_wake_op(value3)?;
            validate_futex_word(uaddr)?;

            let key2 = FutexKey::new_current(uaddr2.addr(), op.key_mode);
            let table2 = futex_table_for(&key2);

            let count = retry_futex_nofault(
                || {
                    let source = futex_table.get_or_insert(&key);
                    let target = table2.get_or_insert(&key2);
                    source.wq.wake_op(wake_count, &target.wq, wake2_count, || {
                        futex_atomic_op_in_user_nofault(uaddr2, &operation)
                    })
                },
                || fault_in_user_u32_write(uaddr2),
            )?;

            if count > 0 {
                ax_task::yield_now();
            }
            Ok(count as _)
        }
    }
}

pub fn sys_get_robust_list(
    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(task.as_thread().robust_list_head() as _)?;
    size.vm_write(size_of::<robust_list_head>())?;

    Ok(0)
}

pub fn sys_set_robust_list(head: *const robust_list_head, size: usize) -> 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(test)]
pub(crate) 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() == true);
    assert!(compare_futex_wake_op(5, FUTEX_OP_CMP_EQ, 6).unwrap() == false);
    assert!(compare_futex_wake_op(5, FUTEX_OP_CMP_NE, 6).unwrap() == true);
    assert!(compare_futex_wake_op(5, FUTEX_OP_CMP_LT, 10).unwrap() == true);
    assert!(compare_futex_wake_op(5, FUTEX_OP_CMP_LE, 5).unwrap() == true);
    assert!(compare_futex_wake_op(5, FUTEX_OP_CMP_GT, 3).unwrap() == true);
    assert!(compare_futex_wake_op(5, FUTEX_OP_CMP_GE, 5).unwrap() == true);
    assert!(compare_futex_wake_op(0, 0xFFFF, 0).is_err()); // unsupported cmp

    true
}