rust-kernel

//! 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)]

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_SLEEP_MS: u64 = 35;
pub const SYS_GETPID: u64 = 39;
pub const SYS_EXIT: u64 = 60;
pub const SYS_WAIT: u64 = 61;
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 SEEK_SET: u64 = 0;
pub const SEEK_CUR: u64 = 1;
pub const SEEK_END: u64 = 2;
pub const SPAWN_TRACE: u64 = 1;

pub const ENOENT: i64 = 2;
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 {
    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,
            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> {
    let mut fds = [0u32; 2];
    match syscall3(SYS_PIPE, fds.as_mut_ptr() as u64, 0, 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) {
    syscall3(SYS_SLEEP_MS, ms, 0, 0);
}
pub fn getpid() -> i64 {
    syscall3(SYS_GETPID, 0, 0, 0)
}
pub fn exit(code: i64) -> ! {
    syscall3(SYS_EXIT, code as u64, 0, 0);
    loop {} // not reached
}
/// Wait for child `pid` to end; returns its exit code (139 = killed).
pub fn wait(pid: i64) -> i64 {
    syscall3(SYS_WAIT, pid as u64, 0, 0)
}
/// Like `wait`, but don't block: -EAGAIN while the child is still running.
pub fn try_wait(pid: i64) -> i64 {
    syscall3(SYS_WAIT, pid as u64, WNOHANG, 0)
}
/// 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).
pub fn getcwd(buf: &mut [u8]) -> i64 {
    syscall3(SYS_GETCWD, buf.as_mut_ptr() as u64, buf.len() as u64, 0)
}
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 ---------------------------------------------------

/// `argv` as the kernel passes it: (pointer, length) pairs.
static mut ARGC: usize = 0;
static mut ARGV: *const [usize; 2] = core::ptr::null();

/// 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 {
            return None;
        }
        let [ptr, len] = *ARGV.add(i);
        core::str::from_utf8(core::slice::from_raw_parts(ptr as *const u8, len)).ok()
    }
}

/// 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() {
        2 => "No such file or directory",
        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 kernel jumps here (see process.rs) with rdi = argc, rsi = argv.
/// Save the arguments, run main, exit with its result.
#[unsafe(no_mangle)]
pub extern "C" fn _start(argc: usize, argv: *const [usize; 2]) -> ! {
    unsafe {
        ARGC = argc;
        ARGV = argv;
    }
    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 a 64 KiB array in .bss, which
// the ELF loader maps as zeroed pages -- no `brk`/`mmap` syscall needed.
// Same idea as the kernel's fixed-size-block allocator: sizes are rounded up
// to a power of two (16 B .. 32 KiB) and freed blocks go on one free list
// per size, to be reused by the next allocation of that size. Programs are
// single-threaded, so no lock.

const HEAP_SIZE: usize = 64 * 1024;
const CLASSES: usize = 12; // 16 << 11 = 32 KiB

#[repr(C, align(4096))]
struct Heap {
    memory: UnsafeCell<[u8; HEAP_SIZE]>,
    used: UnsafeCell<usize>,               // bump pointer: bytes handed out so far
    free: UnsafeCell<[usize; CLASSES]>,    // free list heads (addresses, 0 = empty)
}

unsafe impl Sync for Heap {} // single-threaded programs

#[global_allocator]
static HEAP: Heap = Heap {
    memory: UnsafeCell::new([0; HEAP_SIZE]),
    used: UnsafeCell::new(0),
    free: UnsafeCell::new([0; CLASSES]),
};

/// 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 {
        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;
            }
            // Carve a new block. Blocks are aligned to their size, which
            // covers any alignment up to that size.
            let size = 16 << class;
            let used = &mut *self.used.get();
            let start = (*used).next_multiple_of(size);
            if start + size > HEAP_SIZE {
                return core::ptr::null_mut();
            }
            *used = start + size;
            (self.memory.get() as *mut u8).add(start)
        }
    }

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