foundry

//! A tiny "libc" for user programs: syscall wrappers, print!/println!,
//! eprintln!, a small heap (so `Vec` and `String` work), the `_start` entry
//! point and a panic handler.
//!
//! Each user program includes it with
//!     #[macro_use] #[path = "ulib.rs"] mod ulib;
//! and defines `fn main() -> i64` (the exit code).
//!
//! User programs are built for the `x86_64-unknown-none` target, which has
//! SSE/floating point disabled. That matters: the kernel doesn't save SSE
//! registers on a context switch.

#![allow(dead_code, unused_macros)]

extern crate alloc;

use core::alloc::{GlobalAlloc, Layout};
use core::arch::asm;
use core::cell::UnsafeCell;
use core::fmt;

pub const SYS_READ: u64 = 0;
pub const SYS_WRITE: u64 = 1;
pub const SYS_OPEN: u64 = 2;
pub const SYS_CLOSE: u64 = 3;
pub const SYS_STAT: u64 = 4;
pub const SYS_FSTAT: u64 = 5;
pub const SYS_LSEEK: u64 = 8;
pub const SYS_IOCTL: u64 = 16;
pub const SYS_PIPE: u64 = 22;
pub const SYS_SCHED_YIELD: u64 = 24;
pub const SYS_NANOSLEEP: u64 = 35;
pub const SYS_GETPID: u64 = 39;
pub const SYS_EXIT: u64 = 60;
pub const SYS_WAIT4: u64 = 61;
pub const SYS_MMAP: u64 = 9;
pub const SYS_MPROTECT: u64 = 10;
pub const SYS_MUNMAP: u64 = 11;
pub const SYS_BRK: u64 = 12;
pub const SYS_RT_SIGACTION: u64 = 13;
pub const SYS_RT_SIGPROCMASK: u64 = 14;
pub const SYS_RT_SIGRETURN: u64 = 15;
pub const SYS_DUP: u64 = 32;
pub const SYS_DUP2: u64 = 33;
pub const SYS_PAUSE: u64 = 34;
pub const SYS_FORK: u64 = 57;
pub const SYS_VFORK: u64 = 58;
pub const SYS_EXECVE: u64 = 59;
pub const SYS_KILL: u64 = 62;
pub const SYS_FCNTL: u64 = 72;
pub const SYS_SETPGID: u64 = 109;
pub const SYS_GETPPID: u64 = 110;
pub const SYS_GETPGRP: u64 = 111;
pub const SYS_EXIT_GROUP: u64 = 231;
pub const SYS_PIPE2: u64 = 293;
pub const SYS_TRACE_ME: u64 = 505;
pub const SYS_UNAME: u64 = 63;
pub const SYS_GETCWD: u64 = 79;
pub const SYS_CHDIR: u64 = 80;
pub const SYS_MKDIR: u64 = 83;
pub const SYS_RMDIR: u64 = 84;
pub const SYS_UNLINK: u64 = 87;
pub const SYS_READLINK: u64 = 89;
pub const SYS_GETDENTS64: u64 = 217;
pub const SYS_SPAWN: u64 = 500;
pub const SYS_MOUNTS: u64 = 501;
pub const SYS_SYSINFO: u64 = 502;
pub const SYS_SOCKET: u64 = 41;
pub const SYS_CONNECT: u64 = 42;
pub const SYS_ACCEPT: u64 = 43;
pub const SYS_SENDTO: u64 = 44;
pub const SYS_RECVFROM: u64 = 45;
pub const SYS_BIND: u64 = 49;
pub const SYS_LISTEN: u64 = 50;
pub const SYS_SETSOCKOPT: u64 = 54;
pub const SYS_CLOCK_GETTIME: u64 = 228;
pub const SYS_RESOLVE: u64 = 503;
pub const SYS_NETINFO: u64 = 504;
pub const SYS_GETRANDOM: u64 = 318;

pub const O_CREAT: u64 = 0o100;
pub const O_TRUNC: u64 = 0o1000;
pub const O_APPEND: u64 = 0o2000;
pub const TTY_SET_RAW: u64 = 1;
// /dev/fb0 ioctls (see the table in the kernel's syscall.rs).
pub const FB_GET_INFO: u64 = 0x4600;
pub const FB_FLUSH: u64 = 0x46f0;
pub const FB_FILL: u64 = 0x46f1;
pub const FB_SET_CONSOLE: u64 = 0x46f2;
pub const WNOHANG: u64 = 1;
pub const WUNTRACED: u64 = 2;
pub const WCONTINUED: u64 = 8;
pub const O_CLOEXEC: u64 = 0o2000000;
pub const F_GETFD: u64 = 1;
pub const F_SETFD: u64 = 2;
pub const FD_CLOEXEC: u64 = 1;
pub const TIOCGPGRP: u64 = 0x540F;
pub const TIOCSPGRP: u64 = 0x5410;
pub const SEEK_SET: u64 = 0;
pub const SEEK_CUR: u64 = 1;
pub const SEEK_END: u64 = 2;
pub const SPAWN_TRACE: u64 = 1;

pub const EPERM: i64 = 1;
pub const ENOENT: i64 = 2;
pub const ESRCH: i64 = 3;
pub const EINTR: i64 = 4;
pub const EAGAIN: i64 = 11;
pub const EBADF: i64 = 9;
pub const ECHILD: i64 = 10;
pub const EFAULT: i64 = 14;
pub const EBUSY: i64 = 16;
pub const EEXIST: i64 = 17;
pub const ENOTDIR: i64 = 20;
pub const EISDIR: i64 = 21;
pub const ESPIPE: i64 = 29;
pub const ENOTEMPTY: i64 = 39;
pub const ETIMEDOUT: i64 = 110;

/// Same layout as the kernel's `fs::file::Stat`.
#[repr(C)]
#[derive(Default, Clone, Copy)]
pub struct Stat {
    pub kind: u64, // KIND_*
    pub size: u64,
    pub mode: u64, // permission bits like 0o644; 0 = the file system has none
}
pub const KIND_FILE: u64 = 1;
pub const KIND_DIR: u64 = 2;
pub const KIND_CONSOLE: u64 = 3;
pub const KIND_PIPE: u64 = 4;
pub const KIND_SYMLINK: u64 = 5;
pub const KIND_DEVICE: u64 = 7; // /dev/null, /dev/fb0, ...

/// Same layout as the kernel's `syscall::SpawnOptions`.
#[repr(C)]
pub struct SpawnOptions {
    pub stdin: i64, // my fd to give the child as its fd 0 (-1 = my fd 0)
    pub stdout: i64,
    pub stderr: i64,
    pub flags: u64, // SPAWN_TRACE
}

impl SpawnOptions {
    pub const INHERIT: SpawnOptions = SpawnOptions { stdin: -1, stdout: -1, stderr: -1, flags: 0 };
}

/// Raw system call: number in rax, args in rdi/rsi/rdx, result in rax.
/// The `syscall` instruction overwrites rcx and r11.
pub fn syscall3(nr: u64, a1: u64, a2: u64, a3: u64) -> i64 {
    let ret: i64;
    unsafe {
        asm!(
            "syscall",
            inlateout("rax") nr as i64 => ret,
            in("rdi") a1,
            in("rsi") a2,
            in("rdx") a3,
            lateout("rcx") _,
            lateout("r11") _,
            options(nostack),
        );
    }
    ret
}

/// 5-argument version (4th argument goes in r10, 5th in r8).
pub fn syscall5(nr: u64, a1: u64, a2: u64, a3: u64, a4: u64, a5: u64) -> i64 {
    syscall6(nr, a1, a2, a3, a4, a5, 0)
}

/// 6-argument version (6th in r9). mmap is the only one that needs it.
pub fn syscall6(nr: u64, a1: u64, a2: u64, a3: u64, a4: u64, a5: u64, a6: u64) -> i64 {
    let ret: i64;
    unsafe {
        asm!(
            "syscall",
            inlateout("rax") nr as i64 => ret,
            in("rdi") a1,
            in("rsi") a2,
            in("rdx") a3,
            in("r10") a4,
            in("r8") a5,
            in("r9") a6,
            lateout("rcx") _,
            lateout("r11") _,
            options(nostack),
        );
    }
    ret
}

pub fn read(fd: i64, buf: &mut [u8]) -> i64 {
    syscall3(SYS_READ, fd as u64, buf.as_mut_ptr() as u64, buf.len() as u64)
}
pub fn write(fd: i64, buf: &[u8]) -> i64 {
    syscall3(SYS_WRITE, fd as u64, buf.as_ptr() as u64, buf.len() as u64)
}
pub fn open(path: &str, flags: u64) -> i64 {
    syscall3(SYS_OPEN, path.as_ptr() as u64, path.len() as u64, flags)
}
pub fn close(fd: i64) -> i64 {
    syscall3(SYS_CLOSE, fd as u64, 0, 0)
}
pub fn stat(path: &str, st: &mut Stat) -> i64 {
    syscall3(SYS_STAT, path.as_ptr() as u64, path.len() as u64, st as *mut Stat as u64)
}
pub fn fstat(fd: i64, st: &mut Stat) -> i64 {
    syscall3(SYS_FSTAT, fd as u64, st as *mut Stat as u64, 0)
}
pub fn lseek(fd: i64, offset: i64, whence: u64) -> i64 {
    syscall3(SYS_LSEEK, fd as u64, offset as u64, whence)
}
/// Returns (read end, write end), or the negative errno.
pub fn pipe() -> Result<(i64, i64), i64> {
    pipe2(0)
}
/// pipe with flags (O_CLOEXEC: both ends close on exec).
pub fn pipe2(flags: u64) -> Result<(i64, i64), i64> {
    let mut fds = [0u32; 2];
    match syscall3(SYS_PIPE2, fds.as_mut_ptr() as u64, flags, 0) {
        0 => Ok((fds[0] as i64, fds[1] as i64)),
        e => Err(e),
    }
}
pub fn sched_yield() {
    syscall3(SYS_SCHED_YIELD, 0, 0, 0);
}
pub fn sleep_ms(ms: u64) {
    nanosleep_us(ms * 1000);
}
/// nanosleep(*timespec, *remaining): 0, or -EINTR if a signal woke us.
pub fn nanosleep_us(us: u64) -> i64 {
    let ts = [us / 1_000_000, (us % 1_000_000) * 1000];
    syscall3(SYS_NANOSLEEP, ts.as_ptr() as u64, 0, 0)
}
pub fn getpid() -> i64 {
    syscall3(SYS_GETPID, 0, 0, 0)
}
pub fn exit(code: i64) -> ! {
    syscall3(SYS_EXIT_GROUP, code as u64, 0, 0);
    loop {} // not reached
}

// --- Processes (fork, exec, wait), stage 14 --------------------------------

/// fork: 0 in the child, the child's pid in the parent (or -errno).
pub fn fork() -> i64 {
    syscall3(SYS_FORK, 0, 0, 0)
}
/// vfork: like fork, but the parent sleeps until the child calls execve or
/// exits. (Here the child gets a copy-on-write copy anyway, so it's safe to
/// do more than exec in it; on Linux it borrows the parent's memory.)
pub fn vfork() -> i64 {
    syscall3(SYS_VFORK, 0, 0, 0)
}
pub fn getppid() -> i64 {
    syscall3(SYS_GETPPID, 0, 0, 0)
}
pub fn getpgrp() -> i64 {
    syscall3(SYS_GETPGRP, 0, 0, 0)
}
/// Put process `pid` (0 = me) into group `pgid` (0 = a new group named
/// after the process).
pub fn setpgid(pid: i64, pgid: i64) -> i64 {
    syscall3(SYS_SETPGID, pid as u64, pgid as u64, 0)
}
/// Which process group gets the console's Ctrl+C.
pub fn tcgetpgrp(fd: i64) -> i64 {
    let mut pgid = 0u32;
    match ioctl(fd, TIOCGPGRP, &mut pgid as *mut u32 as u64) {
        0 => pgid as i64,
        e => e,
    }
}
pub fn tcsetpgrp(fd: i64, pgid: i64) -> i64 {
    let pgid = pgid as u32;
    ioctl(fd, TIOCSPGRP, &pgid as *const u32 as u64)
}
/// Print my system calls from now on (inherited across fork and exec).
pub fn trace_me(on: bool) -> i64 {
    syscall3(SYS_TRACE_ME, on as u64, 0, 0)
}
pub fn dup(fd: i64) -> i64 {
    syscall3(SYS_DUP, fd as u64, 0, 0)
}
pub fn dup2(old: i64, new: i64) -> i64 {
    syscall3(SYS_DUP2, old as u64, new as u64, 0)
}
pub fn fcntl(fd: i64, cmd: u64, arg: u64) -> i64 {
    syscall3(SYS_FCNTL, fd as u64, cmd, arg)
}

/// A list of C strings (NUL-terminated) plus the NULL-terminated array of
/// pointers to them that execve wants.
pub struct CStrings {
    data: alloc::vec::Vec<u8>,
    ends: alloc::vec::Vec<usize>,
    ptrs: alloc::vec::Vec<u64>,
}

impl CStrings {
    pub fn new<'a>(items: impl IntoIterator<Item = &'a str>) -> CStrings {
        let mut data = alloc::vec::Vec::new();
        let mut ends = alloc::vec::Vec::new();
        for item in items {
            data.extend_from_slice(item.as_bytes());
            data.push(0);
            ends.push(data.len());
        }
        CStrings { data, ends, ptrs: alloc::vec::Vec::new() }
    }
    /// The `char **`: pointers into `data`, then NULL.
    pub fn as_ptr(&mut self) -> *const u64 {
        let base = self.data.as_ptr() as u64;
        let mut start = 0;
        self.ptrs.clear();
        for &end in &self.ends {
            self.ptrs.push(base + start as u64);
            start = end;
        }
        self.ptrs.push(0);
        self.ptrs.as_ptr()
    }
}

/// execve(path, argv, envp): only returns on failure (-errno).
pub fn execve(path: &str, argv: &[&str], envp: &[&str]) -> i64 {
    let mut cpath = alloc::vec::Vec::from(path.as_bytes());
    cpath.push(0);
    let mut args = CStrings::new(argv.iter().copied());
    let mut env = CStrings::new(envp.iter().copied());
    syscall3(SYS_EXECVE, cpath.as_ptr() as u64, args.as_ptr() as u64, env.as_ptr() as u64)
}

/// The same environment I got: for passing on to execve.
pub fn environ() -> alloc::vec::Vec<&'static str> {
    (0..).map_while(env).collect()
}

/// wait4(pid, &status, options): returns (pid, status), or -errno. pid -1 =
/// any child; WNOHANG and nothing finished: pid 0.
pub fn wait4(pid: i64, options: u64) -> Result<(i64, i32), i64> {
    let mut status = 0i32;
    match syscall5(SYS_WAIT4, pid as u64, &mut status as *mut i32 as u64, options, 0, 0) {
        n if n < 0 => Err(n),
        n => Ok((n, status)),
    }
}

// The status macros (C's WIFEXITED etc.): wait4's `status` packs how the
// child ended into one int -- exit code << 8, or the signal number, or
// 0x7f for "stopped".
pub fn wifexited(status: i32) -> bool {
    status & 0x7f == 0
}
pub fn wexitstatus(status: i32) -> i32 {
    (status >> 8) & 0xff
}
pub fn wifsignaled(status: i32) -> bool {
    (status & 0x7f) + 1 >= 2 && status & 0x7f != 0x7f
}
pub fn wtermsig(status: i32) -> i32 {
    status & 0x7f
}
pub fn wcoredump(status: i32) -> bool {
    status & 0x80 != 0
}
pub fn wifstopped(status: i32) -> bool {
    status & 0xff == 0x7f
}
pub fn wstopsig(status: i32) -> i32 {
    wexitstatus(status)
}
pub fn wifcontinued(status: i32) -> bool {
    status == 0xffff
}
/// A shell-style exit code: the code itself, or 128 + the signal number.
pub fn exit_code(status: i32) -> i64 {
    if wifsignaled(status) { 128 + wtermsig(status) as i64 } else { wexitstatus(status) as i64 }
}

/// Wait for child `pid` to end; returns its exit code (139 = killed by
/// SIGSEGV: 128 + 11), or -errno.
pub fn wait(pid: i64) -> i64 {
    loop {
        match wait4(pid, 0) {
            Ok((_, status)) => return exit_code(status),
            Err(e) if e == -EINTR => continue,
            Err(e) => return e,
        }
    }
}
/// Like `wait`, but don't block: -EAGAIN while the child is still running.
pub fn try_wait(pid: i64) -> i64 {
    match wait4(pid, WNOHANG) {
        Ok((0, _)) => -EAGAIN,
        Ok((_, status)) => exit_code(status),
        Err(e) => e,
    }
}

// --- Signals -----------------------------------------------------------------

pub const SIGHUP: i32 = 1;
pub const SIGINT: i32 = 2;
pub const SIGQUIT: i32 = 3;
pub const SIGILL: i32 = 4;
pub const SIGABRT: i32 = 6;
pub const SIGBUS: i32 = 7;
pub const SIGFPE: i32 = 8;
pub const SIGKILL: i32 = 9;
pub const SIGUSR1: i32 = 10;
pub const SIGSEGV: i32 = 11;
pub const SIGUSR2: i32 = 12;
pub const SIGPIPE: i32 = 13;
pub const SIGALRM: i32 = 14;
pub const SIGTERM: i32 = 15;
pub const SIGCHLD: i32 = 17;
pub const SIGCONT: i32 = 18;
pub const SIGSTOP: i32 = 19;
pub const SIGTSTP: i32 = 20;
pub const SIGTTIN: i32 = 21;
pub const SIGTTOU: i32 = 22;

/// Signal names, index = number ("" = unused).
pub const SIGNAL_NAMES: [&str; 32] = [
    "", "HUP", "INT", "QUIT", "ILL", "TRAP", "ABRT", "BUS", "FPE", "KILL", "USR1", "SEGV", "USR2", "PIPE", "ALRM", "TERM",
    "STKFLT", "CHLD", "CONT", "STOP", "TSTP", "TTIN", "TTOU", "URG", "XCPU", "XFSZ", "VTALRM", "PROF", "WINCH", "IO", "PWR",
    "SYS",
];
/// "INT" or "SIGINT" or "2" -> 2.
pub fn signal_number(name: &str) -> Option<i32> {
    if let Ok(n) = name.parse::<i32>() {
        return (1..64).contains(&n).then_some(n);
    }
    let name = name.strip_prefix("SIG").unwrap_or(name);
    SIGNAL_NAMES.iter().position(|&n| !n.is_empty() && n.eq_ignore_ascii_case(name)).map(|i| i as i32)
}

pub const SIG_DFL: usize = 0;
pub const SIG_IGN: usize = 1;
pub const SA_SIGINFO: u64 = 4;
pub const SA_RESTORER: u64 = 0x0400_0000;
pub const SA_RESTART: u64 = 0x1000_0000;
pub const SA_NODEFER: u64 = 0x4000_0000;
pub const SIG_BLOCK: u64 = 0;
pub const SIG_UNBLOCK: u64 = 1;
pub const SIG_SETMASK: u64 = 2;

/// The kernel's `struct sigaction` (not quite libc's: no sa_sigaction
/// union, the mask is 64 bits).
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct SigAction {
    pub handler: usize, // fn(i32) or fn(i32, *SigInfo, *ucontext); or SIG_DFL / SIG_IGN
    pub flags: u64,
    pub restorer: usize,
    pub mask: u64,
}

/// The first fields of Linux's `siginfo_t` (it's 128 bytes in all).
#[repr(C)]
pub struct SigInfo {
    pub signo: i32,
    pub errno: i32,
    pub code: i32,
    _pad: i32,
    /// SIGSEGV/SIGBUS: the faulting address. SIGCHLD/kill: the sender's
    /// pid (low 32 bits) and uid (high 32 bits).
    pub addr_or_pid: u64,
    /// SIGCHLD: the exit code (or the signal that ended/stopped the child).
    pub status: i32,
}

// After a handler returns it "returns" into this: the restorer, which asks
// the kernel to put everything back as it was before the signal (from the
// frame the kernel saved on our stack). Linux's libc does exactly this.
core::arch::global_asm!(
    ".global __restore_rt",
    "__restore_rt:",
    "mov rax, 15", // rt_sigreturn
    "syscall",
);
unsafe extern "C" {
    fn __restore_rt();
}

/// Install `handler` for `sig` (plain `fn(i32)`). Returns the old action.
pub fn signal(sig: i32, handler: extern "C" fn(i32)) -> i64 {
    sigaction_raw(sig, handler as usize, 0, 0)
}
/// SA_SIGINFO handler: also gets the SigInfo (fault address, sender pid...).
pub fn sigaction_info(sig: i32, handler: extern "C" fn(i32, *const SigInfo, *mut u8), flags: u64) -> i64 {
    sigaction_raw(sig, handler as usize, flags | SA_SIGINFO, 0)
}
/// Set SIG_DFL / SIG_IGN / a handler address, with flags and a mask of
/// signals to block while it runs.
pub fn sigaction_raw(sig: i32, handler: usize, flags: u64, mask: u64) -> i64 {
    let act = SigAction { handler, flags: flags | SA_RESTORER, restorer: __restore_rt as usize, mask };
    syscall5(SYS_RT_SIGACTION, sig as u64, &act as *const _ as u64, 0, 8, 0)
}
/// Change the blocked set; returns the old one (or -errno).
pub fn sigprocmask(how: u64, set: u64) -> i64 {
    let mut old = 0u64;
    match syscall5(SYS_RT_SIGPROCMASK, how, &set as *const u64 as u64, &mut old as *mut u64 as u64, 8, 0) {
        0 => old as i64,
        e => e,
    }
}
pub fn sigbit(sig: i32) -> u64 {
    1 << (sig - 1)
}
/// kill(pid, sig): pid > 0 one process, 0 my group, -pgid a group, -1 all.
pub fn kill(pid: i64, sig: i32) -> i64 {
    syscall3(SYS_KILL, pid as u64, sig as u64, 0)
}
/// Sleep until a signal handler has run (always returns -EINTR).
pub fn pause() -> i64 {
    syscall3(SYS_PAUSE, 0, 0, 0)
}

// --- Memory: mmap and brk ----------------------------------------------------

pub const PROT_NONE: u64 = 0;
pub const PROT_READ: u64 = 1;
pub const PROT_WRITE: u64 = 2;
pub const PROT_EXEC: u64 = 4;
pub const MAP_SHARED: u64 = 0x01;
pub const MAP_PRIVATE: u64 = 0x02;
pub const MAP_FIXED: u64 = 0x10;
pub const MAP_ANONYMOUS: u64 = 0x20;

/// mmap: the address of the new mapping, or -errno. fd -1 = anonymous.
pub fn mmap(addr: u64, len: u64, prot: u64, flags: u64, fd: i64, offset: u64) -> i64 {
    syscall6(SYS_MMAP, addr, len, prot, flags, fd as u64, offset)
}
pub fn munmap(addr: u64, len: u64) -> i64 {
    syscall3(SYS_MUNMAP, addr, len, 0)
}
pub fn mprotect(addr: u64, len: u64, prot: u64) -> i64 {
    syscall3(SYS_MPROTECT, addr, len, prot)
}
/// brk: move the end of the heap; returns the new end (brk(0) = current).
pub fn brk(addr: u64) -> u64 {
    syscall3(SYS_BRK, addr, 0, 0) as u64
}

/// Console raw mode on/off (keys arrive one by one, no echo). Returns the
/// old mode, or -ENOTTY if `fd` isn't the console.
pub fn set_raw(fd: i64, raw: bool) -> i64 {
    syscall3(SYS_IOCTL, fd as u64, TTY_SET_RAW, raw as u64)
}
/// Linux's `struct utsname`: sysname, nodename, release, version, machine,
/// domainname. Returns them as strings.
pub fn uname(buf: &mut [u8; 6 * 65]) -> [&str; 6] {
    syscall3(SYS_UNAME, buf.as_mut_ptr() as u64, 0, 0);
    core::array::from_fn(|i| {
        let field = &buf[i * 65..(i + 1) * 65];
        let len = field.iter().position(|&b| b == 0).unwrap_or(65);
        core::str::from_utf8(&field[..len]).unwrap_or("?")
    })
}
/// Fill `buf` with random bytes from the kernel (virtio-rng). Linux 318.
pub fn getrandom(buf: &mut [u8]) -> i64 {
    syscall3(SYS_GETRANDOM, buf.as_mut_ptr() as u64, buf.len() as u64, 0)
}
/// A raw ioctl: what `request` means depends on the device.
pub fn ioctl(fd: i64, request: u64, arg: u64) -> i64 {
    syscall3(SYS_IOCTL, fd as u64, request, arg)
}
/// The mount table as text ("path fstype" lines); returns the length.
pub fn mounts(buf: &mut [u8]) -> i64 {
    syscall3(SYS_MOUNTS, buf.as_mut_ptr() as u64, buf.len() as u64, 0)
}
/// What `sysinfo` reports (same layout as the kernel's syscall::SysInfo).
/// `ticks[cpu]` counts that CPU's ~10 ms timer ticks, `busy_ticks[cpu]` the
/// ones where it was running a real thread rather than idling: sample twice
/// and subtract to get a CPU's load.
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct SysInfo {
    pub cpus: u64,
    pub this_cpu: u64,
    pub threads: u64,
    pub processes: u64,
    pub frames_in_use: u64,
    pub uptime_ticks: u64,
    pub ticks: [u64; 16],
    pub busy_ticks: [u64; 16],
}
/// Ask the kernel for CPU, thread, process and memory counts.
pub fn sysinfo() -> SysInfo {
    let mut info = SysInfo::default();
    syscall3(SYS_SYSINFO, &mut info as *mut SysInfo as u64, 0, 0);
    info
}
/// Current directory into `buf`; returns its length (or -errno). The
/// kernel, like Linux, also writes a NUL and counts it; we don't.
pub fn getcwd(buf: &mut [u8]) -> i64 {
    let n = syscall3(SYS_GETCWD, buf.as_mut_ptr() as u64, buf.len() as u64, 0);
    if n > 0 { n - 1 } else { n }
}
pub fn chdir(path: &str) -> i64 {
    syscall3(SYS_CHDIR, path.as_ptr() as u64, path.len() as u64, 0)
}
pub fn mkdir(path: &str) -> i64 {
    syscall3(SYS_MKDIR, path.as_ptr() as u64, path.len() as u64, 0)
}
pub fn rmdir(path: &str) -> i64 {
    syscall3(SYS_RMDIR, path.as_ptr() as u64, path.len() as u64, 0)
}
/// The target of a symlink, into `buf`; returns its length or -errno.
pub fn readlink(path: &str, buf: &mut [u8]) -> i64 {
    syscall5(SYS_READLINK, path.as_ptr() as u64, path.len() as u64, buf.as_mut_ptr() as u64, buf.len() as u64, 0)
}
pub fn unlink(path: &str) -> i64 {
    syscall3(SYS_UNLINK, path.as_ptr() as u64, path.len() as u64, 0)
}
pub fn getdents64(fd: i64, buf: &mut [u8]) -> i64 {
    syscall3(SYS_GETDENTS64, fd as u64, buf.as_mut_ptr() as u64, buf.len() as u64)
}
/// Start program `path` with arguments `argv` (NUL-separated, e.g.
/// "child\0hi"). Returns the child's pid (or -errno).
pub fn spawn(path: &str, argv: &str, opts: &SpawnOptions) -> i64 {
    syscall5(
        SYS_SPAWN,
        path.as_ptr() as u64,
        path.len() as u64,
        argv.as_ptr() as u64,
        argv.len() as u64,
        opts as *const SpawnOptions as u64,
    )
}

/// Walk the `linux_dirent64` records getdents64 put in `buf[..len]`,
/// calling `f(name, d_type)` for each (d_type: 4 = directory, 8 = file).
pub fn for_each_dirent(buf: &[u8], len: usize, mut f: impl FnMut(&str, u8)) {
    let mut pos = 0;
    while pos + 19 < len {
        let reclen = u16::from_le_bytes([buf[pos + 16], buf[pos + 17]]) as usize;
        let d_type = buf[pos + 18];
        let name = &buf[pos + 19..pos + reclen];
        let name = &name[..name.iter().position(|&b| b == 0).unwrap_or(name.len())];
        f(core::str::from_utf8(name).unwrap_or("?"), d_type);
        pos += reclen;
    }
}

// --- program arguments, environment, auxiliary vector -------------------------
//
// The kernel starts a program the way Linux does: rsp points at
//     argc, argv[0..argc], NULL, envp[0..], NULL, auxv pairs..., AT_NULL 0
// and the strings themselves sit higher up on the stack.

static mut ARGC: usize = 0;
static mut ARGV: *const *const u8 = core::ptr::null();
static mut ENVP: *const *const u8 = core::ptr::null();
static mut AUXV: *const [u64; 2] = core::ptr::null();

/// A C string as &str (up to its NUL).
pub unsafe fn cstr(ptr: *const u8) -> &'static str {
    unsafe {
        let mut len = 0;
        while *ptr.add(len) != 0 {
            len += 1;
        }
        core::str::from_utf8(core::slice::from_raw_parts(ptr, len)).unwrap_or("?")
    }
}

/// The program's arguments; `arg(0)` is the program name.
pub fn argc() -> usize {
    unsafe { ARGC }
}
pub fn arg(i: usize) -> Option<&'static str> {
    unsafe { if i < ARGC { Some(cstr(*ARGV.add(i))) } else { None } }
}
/// The i-th environment string ("NAME=value").
pub fn env(i: usize) -> Option<&'static str> {
    unsafe {
        let p = *ENVP.add(i);
        if p.is_null() { None } else { Some(cstr(p)) }
    }
}
/// getenv("PATH") -> the value after "PATH=".
pub fn getenv(name: &str) -> Option<&'static str> {
    (0..).map_while(env).find_map(|e| e.strip_prefix(name)?.strip_prefix('='))
}
pub const AT_PHDR: u64 = 3;
pub const AT_PHENT: u64 = 4;
pub const AT_PHNUM: u64 = 5;
pub const AT_PAGESZ: u64 = 6;
pub const AT_ENTRY: u64 = 9;
pub const AT_RANDOM: u64 = 25;
pub const AT_EXECFN: u64 = 31;
/// An auxiliary vector entry (AT_PAGESZ -> 4096 ...), if the kernel gave it.
pub fn getauxval(kind: u64) -> Option<u64> {
    unsafe {
        let mut p = AUXV;
        while (*p)[0] != 0 {
            if (*p)[0] == kind {
                return Some((*p)[1]);
            }
            p = p.add(1);
        }
        None
    }
}

/// Write all of `buf` (pipes and files may take less than all at once).
pub fn write_all(fd: i64, mut buf: &[u8]) -> i64 {
    while !buf.is_empty() {
        let n = write(fd, buf);
        if n <= 0 {
            return n;
        }
        buf = &buf[n as usize..];
    }
    0
}

/// Call `f` for each line read from `fd` (without the '\n'). Lines longer
/// than 1 KiB come in pieces. Returns 0, or -errno if a read fails.
pub fn for_each_line(fd: i64, mut f: impl FnMut(&[u8])) -> i64 {
    let mut line = [0u8; 1024];
    let mut len = 0;
    let mut buf = [0u8; 512];
    loop {
        let n = read(fd, &mut buf);
        if n < 0 {
            return n;
        }
        if n == 0 {
            if len > 0 {
                f(&line[..len]); // last line without a newline
            }
            return 0;
        }
        for &b in &buf[..n as usize] {
            if b == b'\n' || len == line.len() {
                f(&line[..len]);
                len = 0;
                if b == b'\n' {
                    continue;
                }
            }
            line[len] = b;
            len += 1;
        }
    }
}

/// `ls`, `cat` etc. print errors like "cat: x: No such file or directory".
pub fn strerror(errno: i64) -> &'static str {
    match errno.abs() {
        1 => "Operation not permitted",
        2 => "No such file or directory",
        3 => "No such process",
        4 => "Interrupted system call",
        5 => "I/O error",
        7 => "Argument list too long",
        8 => "Exec format error",
        9 => "Bad file descriptor",
        10 => "No child processes",
        11 => "Try again",
        12 => "Out of memory",
        14 => "Bad address",
        16 => "Device or resource busy",
        19 => "No such device",
        17 => "File exists",
        20 => "Not a directory",
        21 => "Is a directory",
        22 => "Invalid argument",
        24 => "Too many open files",
        25 => "Not a terminal",
        27 => "File too large",
        28 => "No space left on device",
        29 => "Illegal seek",
        30 => "Read-only file system",
        32 => "Broken pipe",
        34 => "Result too large",
        36 => "File name too long",
        38 => "Function not implemented",
        39 => "Directory not empty",
        40 => "Too many levels of symbolic links",
        95 => "Operation not supported",
        88 => "Socket operation on non-socket",
        89 => "Destination address required",
        90 => "Message too long",
        93 => "Protocol not supported",
        97 => "Address family not supported",
        98 => "Address already in use",
        100 => "Network is down",
        101 => "Network is unreachable",
        104 => "Connection reset by peer",
        106 => "Already connected",
        107 => "Not connected",
        110 => "Connection timed out",
        111 => "Connection refused",
        _ => "Unknown error",
    }
}

// --- Networking ------------------------------------------------------------------
//
// BSD sockets, the API every OS copied from 1983's 4.2BSD Unix: `socket()`
// makes an fd, `connect()` dials, read/write move data, `close()` hangs up.
// Servers `bind()` a port, `listen()`, and `accept()` each caller.

pub const AF_INET: u64 = 2;
pub const SOCK_STREAM: u64 = 1; // TCP
pub const SOCK_DGRAM: u64 = 2; // UDP (or ICMP with IPPROTO_ICMP)
pub const IPPROTO_ICMP: u64 = 1;

/// An IPv4 address that prints as "10.0.2.15".
#[derive(Clone, Copy, PartialEq, Eq, Default)]
pub struct Ip(pub [u8; 4]);

impl fmt::Display for Ip {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let [a, b, c, d] = self.0;
        write!(f, "{}.{}.{}.{}", a, b, c, d)
    }
}

/// "10.0.2.2" -> Ip, or None if it isn't a dotted-quad address.
pub fn parse_ip(s: &str) -> Option<Ip> {
    let mut ip = [0u8; 4];
    let mut parts = s.split('.');
    for b in ip.iter_mut() {
        *b = parts.next()?.parse().ok()?;
    }
    if parts.next().is_some() { None } else { Some(Ip(ip)) }
}

/// Linux's `struct sockaddr_in` (port and address big-endian).
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct SockAddrIn {
    pub family: u16,
    pub port_be: u16,
    pub addr: [u8; 4],
    pub zero: [u8; 8],
}

impl SockAddrIn {
    pub fn new(ip: Ip, port: u16) -> SockAddrIn {
        SockAddrIn { family: AF_INET as u16, port_be: port.to_be(), addr: ip.0, zero: [0; 8] }
    }
    pub fn ip(&self) -> Ip {
        Ip(self.addr)
    }
    pub fn port(&self) -> u16 {
        u16::from_be(self.port_be)
    }
}

pub fn socket(ty: u64, protocol: u64) -> i64 {
    syscall3(SYS_SOCKET, AF_INET, ty, protocol)
}
pub fn connect(fd: i64, ip: Ip, port: u16) -> i64 {
    let sa = SockAddrIn::new(ip, port);
    syscall3(SYS_CONNECT, fd as u64, &sa as *const _ as u64, 16)
}
pub fn bind(fd: i64, port: u16) -> i64 {
    let sa = SockAddrIn::new(Ip([0; 4]), port);
    syscall3(SYS_BIND, fd as u64, &sa as *const _ as u64, 16)
}
pub fn listen(fd: i64, backlog: u64) -> i64 {
    syscall3(SYS_LISTEN, fd as u64, backlog, 0)
}
/// Wait for a connection: (new fd or -errno, caller's address).
pub fn accept(fd: i64) -> (i64, SockAddrIn) {
    let mut sa = SockAddrIn::default();
    let fd = syscall3(SYS_ACCEPT, fd as u64, &mut sa as *mut _ as u64, 0);
    (fd, sa)
}
pub fn sendto(fd: i64, buf: &[u8], ip: Ip, port: u16) -> i64 {
    let sa = SockAddrIn::new(ip, port);
    syscall5(SYS_SENDTO, fd as u64, buf.as_ptr() as u64, buf.len() as u64, 0, &sa as *const _ as u64)
}
/// Receive one datagram (or some stream bytes): (length or -errno, sender).
pub fn recvfrom(fd: i64, buf: &mut [u8]) -> (i64, SockAddrIn) {
    let mut sa = SockAddrIn::default();
    let n = syscall5(SYS_RECVFROM, fd as u64, buf.as_mut_ptr() as u64, buf.len() as u64, 0, &mut sa as *mut _ as u64);
    (n, sa)
}
/// SO_RCVTIMEO: recv/accept give up with -EAGAIN after `ms` (0 = never).
pub fn set_recv_timeout(fd: i64, ms: u64) -> i64 {
    let tv = [ms / 1000, (ms % 1000) * 1000]; // struct timeval
    syscall5(SYS_SETSOCKOPT, fd as u64, 1, 20, tv.as_ptr() as u64, 16)
}
/// Microseconds since boot (clock_gettime(CLOCK_MONOTONIC)).
pub fn now_us() -> u64 {
    let mut ts = [0u64; 2];
    syscall3(SYS_CLOCK_GETTIME, 1, ts.as_mut_ptr() as u64, 0);
    ts[0] * 1_000_000 + ts[1] / 1000
}
/// Look up a host name with DNS (or parse it if it's already an address).
pub fn resolve(name: &str) -> Result<Ip, i64> {
    if let Some(ip) = parse_ip(name) {
        return Ok(ip);
    }
    let mut out = [[0u8; 4]; 1];
    match syscall5(SYS_RESOLVE, name.as_ptr() as u64, name.len() as u64, out.as_mut_ptr() as u64, 1, 0) {
        n if n < 0 => Err(n),
        0 => Err(-2), // no A record: ENOENT
        _ => Ok(Ip(out[0])),
    }
}
/// All IPv4 addresses for `name` (up to `out.len()`); count or -errno.
pub fn resolve_all(name: &str, out: &mut [[u8; 4]]) -> i64 {
    syscall5(SYS_RESOLVE, name.as_ptr() as u64, name.len() as u64, out.as_mut_ptr() as u64, out.len() as u64, 0)
}

/// What `ifconfig` shows; same layout as the kernel's `net::NetInfo`.
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct NetInfo {
    pub present: u8,
    pub configured: u8,
    pub mac: [u8; 6],
    pub address: [u8; 4],
    pub prefix_len: u8,
    pub irq: u8,
    pub _pad: [u8; 2],
    pub gateway: [u8; 4],
    pub dns: [u8; 4],
    pub rx_packets: u64,
    pub rx_bytes: u64,
    pub tx_packets: u64,
    pub tx_bytes: u64,
}
pub fn netinfo() -> NetInfo {
    let mut info = NetInfo::default();
    syscall3(SYS_NETINFO, &mut info as *mut NetInfo as u64, 0, 0);
    info
}

pub struct Stdout;

impl fmt::Write for Stdout {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        write(1, s.as_bytes());
        Ok(())
    }
}

/// Collects the pieces of one `print!`/`println!` so the whole line goes to
/// the kernel in a single `write`. With several CPUs, two programs printing
/// at the same moment would otherwise have their pieces mixed up mid-line.
/// Flushes when full and when dropped (at the end of the macro).
pub struct LineWriter {
    fd: i64,
    buf: [u8; 256],
    len: usize,
}

impl LineWriter {
    pub fn new(fd: i64) -> LineWriter {
        LineWriter { fd, buf: [0; 256], len: 0 }
    }
    pub fn flush(&mut self) {
        if self.len > 0 {
            write_all(self.fd, &self.buf[..self.len]);
            self.len = 0;
        }
    }
}

impl fmt::Write for LineWriter {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        if self.len + s.len() > self.buf.len() {
            self.flush();
        }
        if s.len() > self.buf.len() {
            write_all(self.fd, s.as_bytes()); // too big to buffer: send as is
        } else {
            self.buf[self.len..self.len + s.len()].copy_from_slice(s.as_bytes());
            self.len += s.len();
        }
        Ok(())
    }
}

impl Drop for LineWriter {
    fn drop(&mut self) {
        self.flush();
    }
}

macro_rules! print {
    ($($arg:tt)*) => {{
        use core::fmt::Write;
        let mut out = $crate::ulib::LineWriter::new(1);
        let _ = write!(out, $($arg)*);
    }};
}

pub struct Stderr;

impl fmt::Write for Stderr {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        write(2, s.as_bytes());
        Ok(())
    }
}

/// Write formatted text to any fd: `fprint!(fd, "...")`.
pub struct Fd(pub i64);

impl fmt::Write for Fd {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        write_all(self.0, s.as_bytes());
        Ok(())
    }
}

macro_rules! eprintln {
    ($($arg:tt)*) => {{
        use core::fmt::Write;
        let mut out = $crate::ulib::LineWriter::new(2);
        let _ = writeln!(out, $($arg)*);
    }};
}

macro_rules! println {
    ($($arg:tt)*) => {{
        use core::fmt::Write;
        let mut out = $crate::ulib::LineWriter::new(1);
        let _ = writeln!(out, $($arg)*);
    }};
}

// The entry point. The kernel jumps here with rsp at `argc` (see above);
// hand that pointer to Rust (rdi = first argument), keeping rsp 16-byte
// aligned the way a `call` expects.
core::arch::global_asm!(
    ".global _start",
    "_start:",
    "mov rdi, rsp",
    "and rsp, -16",
    "call start_rust",
    "ud2",
);

/// Find argv/envp/auxv on the initial stack, run main, exit with its result.
#[unsafe(no_mangle)]
extern "C" fn start_rust(sp: *const usize) -> ! {
    unsafe {
        ARGC = *sp;
        ARGV = sp.add(1) as *const *const u8;
        ENVP = ARGV.add(ARGC + 1);
        let mut p = ENVP;
        while !(*p).is_null() {
            p = p.add(1);
        }
        AUXV = p.add(1) as *const [u64; 2];
    }
    exit(crate::main())
}

#[panic_handler]
fn panic(info: &core::panic::PanicInfo) -> ! {
    println!("[user panic] {}", info);
    exit(101)
}

// --- heap ----------------------------------------------------------------
//
// A tiny allocator so programs can use `Vec`, `String`, `format!` (with
// `extern crate alloc;`). Memory comes from `brk`: the kernel's "program
// break" is the end of a heap area right after the program, and moving it
// up gives us more (zeroed, demand-paged) memory.
// Same idea as the kernel's fixed-size-block allocator: sizes are rounded up
// to a power of two (16 B .. 1 MiB) and freed blocks go on one free list
// per size, to be reused by the next allocation of that size. Programs may
// have several threads (stage 15), so a tiny spinlock guards the heap; a
// thread that finds it taken yields its CPU instead of spinning long.

const CLASSES: usize = 17; // 16 << 16 = 1 MiB

struct Heap {
    start: UnsafeCell<u64>, // where the heap began (0 = not set up yet)
    used: UnsafeCell<u64>,  // bump pointer: end of what's handed out
    end: UnsafeCell<u64>,   // the current break
    free: UnsafeCell<[usize; CLASSES]>, // free list heads (addresses, 0 = empty)
    lock: AtomicBool,
}

unsafe impl Sync for Heap {} // the fields are only touched under `lock`

impl Heap {
    fn lock(&self) {
        while self.lock.swap(true, Ordering::Acquire) {
            sched_yield();
        }
    }
    fn unlock(&self) {
        self.lock.store(false, Ordering::Release);
    }
}

#[global_allocator]
static HEAP: Heap = Heap {
    start: UnsafeCell::new(0),
    used: UnsafeCell::new(0),
    end: UnsafeCell::new(0),
    free: UnsafeCell::new([0; CLASSES]),
    lock: AtomicBool::new(false),
};

/// Size class for a layout: 0 = 16 bytes, 1 = 32, ...
fn class_of(layout: Layout) -> usize {
    let size = layout.size().max(layout.align()).max(16).next_power_of_two();
    size.trailing_zeros() as usize - 4
}

unsafe impl GlobalAlloc for Heap {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        self.lock();
        let p = unsafe { self.alloc_locked(layout) };
        self.unlock();
        p
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let class = class_of(layout);
        self.lock();
        unsafe {
            let free = &mut *self.free.get();
            *(ptr as *mut usize) = free[class];
            free[class] = ptr as usize;
        }
        self.unlock();
    }
}

impl Heap {
    unsafe fn alloc_locked(&self, layout: Layout) -> *mut u8 {
        let class = class_of(layout);
        if class >= CLASSES {
            return core::ptr::null_mut(); // too big: "out of memory"
        }
        unsafe {
            let free = &mut *self.free.get();
            if free[class] != 0 {
                // Pop a freed block; its first 8 bytes hold the next one.
                let block = free[class] as *mut u8;
                free[class] = *(block as *const usize);
                return block;
            }
            let (used, end) = (&mut *self.used.get(), &mut *self.end.get());
            if *self.start.get() == 0 {
                *self.start.get() = brk(0);
                *used = *self.start.get();
                *end = *used;
            }
            // Carve a new block. Blocks are aligned to their size, which
            // covers any alignment up to that size.
            let size = 16u64 << class;
            let start = (*used).next_multiple_of(size);
            if start + size > *end {
                // Grow the heap (at least 64 KiB at a time).
                let want = (start + size).max(*end + 64 * 1024);
                let got = brk(want);
                if got < start + size {
                    return core::ptr::null_mut(); // the kernel said no
                }
                *end = got;
            }
            *used = start + size;
            start as *mut u8
        }
    }
}

// --- Stage 15: what a libc's plumbing needs ----------------------------------

pub const SYS_POLL: u64 = 7;
pub const SYS_PREAD64: u64 = 17;
pub const SYS_PWRITE64: u64 = 18;
pub const SYS_READV: u64 = 19;
pub const SYS_WRITEV: u64 = 20;
pub const SYS_CLONE: u64 = 56;
pub const SYS_EXIT_THREAD: u64 = 60; // Linux's "exit": just this thread
pub const SYS_SETSID: u64 = 112;
pub const SYS_ARCH_PRCTL: u64 = 158;
pub const SYS_GETTID: u64 = 186;
pub const SYS_FUTEX: u64 = 202;
pub const SYS_EPOLL_WAIT: u64 = 232;
pub const SYS_EPOLL_CTL: u64 = 233;
pub const SYS_TGKILL: u64 = 234;
pub const SYS_EPOLL_CREATE1: u64 = 291;

pub const O_NOCTTY: u64 = 0o400;
pub const O_NONBLOCK: u64 = 0o4000;
pub const O_RDWR: u64 = 2;
pub const F_GETFL: u64 = 3;
pub const F_SETFL: u64 = 4;
pub const EINPROGRESS: i64 = 115;

/// Turn O_NONBLOCK on or off for `fd` (fcntl F_GETFL + F_SETFL).
pub fn set_nonblocking(fd: i64, on: bool) -> i64 {
    let flags = fcntl(fd, F_GETFL, 0);
    if flags < 0 {
        return flags;
    }
    let flags = if on { flags as u64 | O_NONBLOCK } else { flags as u64 & !O_NONBLOCK };
    fcntl(fd, F_SETFL, flags)
}

pub fn pread(fd: i64, buf: &mut [u8], offset: u64) -> i64 {
    syscall6(SYS_PREAD64, fd as u64, buf.as_mut_ptr() as u64, buf.len() as u64, offset, 0, 0)
}
pub fn pwrite(fd: i64, buf: &[u8], offset: u64) -> i64 {
    syscall6(SYS_PWRITE64, fd as u64, buf.as_ptr() as u64, buf.len() as u64, offset, 0, 0)
}
/// writev: several buffers in one call (`struct iovec {base, len}`).
pub fn writev(fd: i64, bufs: &[&[u8]]) -> i64 {
    let iov: alloc::vec::Vec<[u64; 2]> = bufs.iter().map(|b| [b.as_ptr() as u64, b.len() as u64]).collect();
    syscall3(SYS_WRITEV, fd as u64, iov.as_ptr() as u64, iov.len() as u64)
}

/// Linux 228: 0 = CLOCK_REALTIME (wall clock), 1 = CLOCK_MONOTONIC.
pub fn clock_gettime(clock: u64) -> (u64, u64) {
    let mut ts = [0u64; 2];
    syscall3(SYS_CLOCK_GETTIME, clock, ts.as_mut_ptr() as u64, 0);
    (ts[0], ts[1])
}

// poll / epoll (see the kernel's poll.rs).
pub const POLLIN: u16 = 0x1;
pub const POLLOUT: u16 = 0x4;
pub const POLLERR: u16 = 0x8;
pub const POLLHUP: u16 = 0x10;
pub const POLLNVAL: u16 = 0x20;

/// `struct pollfd`.
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct PollFd {
    pub fd: i32,
    pub events: u16,
    pub revents: u16,
}
/// Wait until one of `fds` is ready (timeout_ms -1 = forever); returns
/// how many are, 0 on timeout.
pub fn poll(fds: &mut [PollFd], timeout_ms: i64) -> i64 {
    syscall3(SYS_POLL, fds.as_mut_ptr() as u64, fds.len() as u64, timeout_ms as u64)
}

pub const EPOLL_CTL_ADD: u64 = 1;
pub const EPOLL_CTL_DEL: u64 = 2;
pub const EPOLL_CTL_MOD: u64 = 3;
pub const EPOLLIN: u32 = 0x1;
pub const EPOLLOUT: u32 = 0x4;

/// `struct epoll_event` is *packed* on x86_64: 4 + 8 = 12 bytes.
#[repr(C, packed)]
#[derive(Clone, Copy, Default)]
pub struct EpollEvent {
    pub events: u32,
    pub data: u64,
}
pub fn epoll_create1(flags: u64) -> i64 {
    syscall3(SYS_EPOLL_CREATE1, flags, 0, 0)
}
pub fn epoll_ctl(epfd: i64, op: u64, fd: i64, events: u32, data: u64) -> i64 {
    let ev = EpollEvent { events, data };
    syscall5(SYS_EPOLL_CTL, epfd as u64, op, fd as u64, &ev as *const EpollEvent as u64, 0)
}
pub fn epoll_wait(epfd: i64, out: &mut [EpollEvent], timeout_ms: i64) -> i64 {
    syscall5(SYS_EPOLL_WAIT, epfd as u64, out.as_mut_ptr() as u64, out.len() as u64, timeout_ms as u64, 0)
}

// Sessions and terminals.
pub fn setsid() -> i64 {
    syscall3(SYS_SETSID, 0, 0, 0)
}
pub const TCGETS: u64 = 0x5401;
pub const TCSETS: u64 = 0x5402;
pub const TIOCSCTTY: u64 = 0x540E;
pub const TIOCGPTN: u64 = 0x8004_5430;
pub const TIOCSPTLCK: u64 = 0x4004_5431;
pub const ICANON: u32 = 0x2;
pub const ECHO: u32 = 0x8;
pub const ISIG: u32 = 0x1;
pub const VMIN: usize = 6;
pub const VTIME: usize = 5;

/// Linux's `struct termios` (36 bytes): the terminal's settings.
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct Termios {
    pub iflag: u32,
    pub oflag: u32,
    pub cflag: u32,
    pub lflag: u32,
    pub line: u8,
    pub cc: [u8; 19],
}
pub fn tcgetattr(fd: i64) -> Result<Termios, i64> {
    let mut t = Termios::default();
    match ioctl(fd, TCGETS, &mut t as *mut Termios as u64) {
        0 => Ok(t),
        e => Err(e),
    }
}
pub fn tcsetattr(fd: i64, t: &Termios) -> i64 {
    ioctl(fd, TCSETS, t as *const Termios as u64)
}
/// Raw mode, like libc's cfmakeraw: no line editing, echo or signals,
/// no character conversions; read returns as soon as 1 byte is there.
pub fn cfmakeraw(t: &mut Termios) {
    t.iflag &= !(0x1 | 0x2 | 0x8 | 0x20 | 0x40 | 0x80 | 0x100 | 0x400); // IGNBRK BRKINT PARMRK ISTRIP INLCR IGNCR ICRNL IXON
    t.oflag &= !0x1; // OPOST
    t.lflag &= !(ECHO | 0x40 | ICANON | ISIG | 0x8000); // ECHO ECHONL ICANON ISIG IEXTEN
    t.cflag = (t.cflag & !0x130) | 0x30; // CS8, no parity
    t.cc[VMIN] = 1;
    t.cc[VTIME] = 0;
}

/// Open a new pseudo-terminal: returns (master fd, slave number N); the
/// slave is /dev/pts/N. Like posix_openpt + grantpt + unlockpt.
pub fn openpty() -> Result<(i64, u32), i64> {
    let master = open("/dev/ptmx", O_RDWR | O_NOCTTY);
    if master < 0 {
        return Err(master);
    }
    let mut n = 0u32;
    let unlock = 0u32;
    let r = ioctl(master, TIOCGPTN, &mut n as *mut u32 as u64);
    let r = if r == 0 { ioctl(master, TIOCSPTLCK, &unlock as *const u32 as u64) } else { r };
    if r != 0 {
        close(master);
        return Err(r);
    }
    Ok((master, n))
}

// Threads: clone + futex + TLS.
pub const FUTEX_WAIT: u64 = 0;
pub const FUTEX_WAKE: u64 = 1;
pub const FUTEX_PRIVATE_FLAG: u64 = 128;
pub const CLONE_VM: u64 = 0x100;
pub const CLONE_FS: u64 = 0x200;
pub const CLONE_FILES: u64 = 0x400;
pub const CLONE_SIGHAND: u64 = 0x800;
pub const CLONE_THREAD: u64 = 0x10000;
pub const CLONE_SYSVSEM: u64 = 0x40000;
pub const CLONE_SETTLS: u64 = 0x80000;
pub const CLONE_PARENT_SETTID: u64 = 0x100000;
pub const CLONE_CHILD_CLEARTID: u64 = 0x200000;
pub const ARCH_SET_FS: u64 = 0x1002;
pub const ARCH_GET_FS: u64 = 0x1003;

use core::sync::atomic::{AtomicBool, AtomicU32, Ordering};

pub fn gettid() -> i64 {
    syscall3(SYS_GETTID, 0, 0, 0)
}
pub fn tgkill(tgid: i64, tid: i64, sig: i32) -> i64 {
    syscall3(SYS_TGKILL, tgid as u64, tid as u64, sig as u64)
}
/// Sleep while `*word == expected` (until futex_wake, or `timeout_us`).
pub fn futex_wait(word: &AtomicU32, expected: u32, timeout_us: Option<u64>) -> i64 {
    let ts = timeout_us.map(|us| [us / 1_000_000, (us % 1_000_000) * 1000]);
    let ts_ptr = ts.as_ref().map_or(0, |t| t.as_ptr() as u64);
    syscall6(SYS_FUTEX, word.as_ptr() as u64, FUTEX_WAIT | FUTEX_PRIVATE_FLAG, expected as u64, ts_ptr, 0, 0)
}
/// Wake up to `n` threads sleeping on `word`; returns how many woke.
pub fn futex_wake(word: &AtomicU32, n: u32) -> i64 {
    syscall6(SYS_FUTEX, word.as_ptr() as u64, FUTEX_WAKE | FUTEX_PRIVATE_FLAG, n as u64, 0, 0, 0)
}

/// A mutex built on a futex (Ulrich Drepper's "Futexes Are Tricky",
/// mutex #2). The word is 0 = unlocked, 1 = locked, 2 = locked and
/// someone may be sleeping. Taking a free lock is one atomic instruction,
/// no syscall; only contention goes to the kernel.
pub struct FutexMutex {
    state: AtomicU32,
}
impl FutexMutex {
    pub const fn new() -> FutexMutex {
        FutexMutex { state: AtomicU32::new(0) }
    }
    pub fn lock(&self) {
        if self.state.compare_exchange(0, 1, Ordering::Acquire, Ordering::Relaxed).is_ok() {
            return; // fast path: it was free
        }
        // Slow path: mark "contended" (2) and sleep until it's free. The
        // kernel only puts us to sleep if the word is *still* 2, so an
        // unlock between our swap and the futex call can't be missed.
        while self.state.swap(2, Ordering::Acquire) != 0 {
            futex_wait(&self.state, 2, None);
        }
    }
    pub fn unlock(&self) {
        // 1 -> 0: nobody waits. 2 -> 0: wake one sleeper.
        if self.state.swap(0, Ordering::Release) == 2 {
            futex_wake(&self.state, 1);
        }
    }
}

/// Thread-local storage, the x86_64 way: each thread's FS register base
/// points at its own block, whose first word points at itself (so code can
/// find the block with `mov rax, fs:[0]`). We keep one u64 per thread in
/// the second word.
#[repr(C)]
pub struct TlsBlock {
    pub self_ptr: u64,
    pub value: u64,
}
pub fn arch_prctl(code: u64, addr: u64) -> i64 {
    syscall3(SYS_ARCH_PRCTL, code, addr, 0)
}
/// This thread's TLS value (fs:[8]).
pub fn tls_value() -> u64 {
    let v: u64;
    unsafe { asm!("mov {}, fs:[8]", out(reg) v, options(nostack, readonly)) };
    v
}

/// A running thread (see `spawn_thread`).
pub struct Thread {
    pub tid: i64,
    /// The kernel writes our tid here (CLONE_PARENT_SETTID) and zeroes it +
    /// futex-wakes it when the thread exits (CLONE_CHILD_CLEARTID): that's
    /// how `join` knows it's done, exactly like pthread_join.
    tid_word: alloc::boxed::Box<AtomicU32>,
    stack: u64,
    tls: alloc::boxed::Box<TlsBlock>,
}

const THREAD_STACK: u64 = 64 * 1024;
type ThreadFn = alloc::boxed::Box<dyn FnOnce()>;

/// The new thread starts here, on its own stack, with the boxed closure.
extern "C" fn thread_entry(arg: u64) -> ! {
    let f = unsafe { alloc::boxed::Box::from_raw(arg as *mut ThreadFn) };
    f();
    // exit (60), not exit_group: only this thread ends.
    syscall3(SYS_EXIT_THREAD, 0, 0, 0);
    loop {}
}

/// Start `f` in a new thread of this process (same memory, same fds),
/// with `tls_value` in its TLS block. Like pthread_create.
pub fn spawn_thread(tls_value: u64, f: impl FnOnce() + Send + 'static) -> Result<Thread, i64> {
    let stack = mmap(0, THREAD_STACK, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if stack < 0 {
        return Err(stack);
    }
    let stack = stack as u64;
    let boxed: alloc::boxed::Box<ThreadFn> = alloc::boxed::Box::new(alloc::boxed::Box::new(f));
    let arg = alloc::boxed::Box::into_raw(boxed) as u64;
    let mut tls = alloc::boxed::Box::new(TlsBlock { self_ptr: 0, value: tls_value });
    tls.self_ptr = &*tls as *const TlsBlock as u64;
    let tid_word = alloc::boxed::Box::new(AtomicU32::new(0));
    // The child's stack starts with [entry, arg] for the asm below.
    let sp = stack + THREAD_STACK - 16;
    unsafe {
        *(sp as *mut u64) = thread_entry as usize as u64;
        *((sp + 8) as *mut u64) = arg;
    }
    let flags = CLONE_VM | CLONE_FS | CLONE_FILES | CLONE_SIGHAND | CLONE_THREAD | CLONE_SYSVSEM
        | CLONE_SETTLS | CLONE_PARENT_SETTID | CLONE_CHILD_CLEARTID;
    let tid_ptr = tid_word.as_ptr() as u64;
    let ret: i64;
    // clone returns twice: the parent gets the tid, the child 0 -- on its
    // new stack, so it must not touch the parent's stack frame. It pops
    // [entry, arg] and calls entry(arg) straight from assembly (after the
    // two pops rsp is 16-byte aligned, as `call` wants).
    unsafe {
        asm!(
            "syscall",
            "test rax, rax",
            "jnz 2f",
            "pop rax",
            "pop rdi",
            "call rax",
            "ud2",
            "2:",
            inlateout("rax") SYS_CLONE as i64 => ret,
            in("rdi") flags,
            in("rsi") sp,
            in("rdx") tid_ptr,
            in("r10") tid_ptr,
            in("r8") tls.self_ptr,
            lateout("rcx") _,
            lateout("r11") _,
            options(nostack)
        );
    }
    if ret < 0 {
        unsafe { drop(alloc::boxed::Box::from_raw(arg as *mut ThreadFn)) };
        munmap(stack, THREAD_STACK);
        return Err(ret);
    }
    Ok(Thread { tid: ret, tid_word, stack, tls })
}

impl Thread {
    /// Wait for the thread to finish, then free its stack.
    pub fn join(self) {
        loop {
            let tid = self.tid_word.load(Ordering::Acquire);
            if tid == 0 {
                break; // the kernel cleared it: the thread is gone
            }
            futex_wait(&self.tid_word, tid, None);
        }
        munmap(self.stack, THREAD_STACK);
        drop(self.tls);
    }
}

// --- Self-checking test programs ---------------------------------------------

/// Counts failed checks for the test programs (forktest, sigtest, ...):
/// each prints "ok" or "FAILED" per check and exits with the failures.
pub struct Checks {
    pub name: &'static str,
    pub failed: i64,
}

impl Checks {
    pub fn new(name: &'static str) -> Checks {
        Checks { name, failed: 0 }
    }
    pub fn check(&mut self, ok: bool, what: &str) {
        if ok {
            println!("[{}] ok: {}", self.name, what);
        } else {
            println!("[{}] FAILED: {}", self.name, what);
            self.failed += 1;
        }
    }
    /// Print the summary; the result is the exit code.
    pub fn done(self) -> i64 {
        if self.failed == 0 {
            println!("[{}] all checks passed", self.name);
        } else {
            println!("[{}] {} check(s) failed", self.name, self.failed);
        }
        self.failed
    }
}