//! A reader holds a program to the profile's bounds as it reads it (extensions/quantum-wasmq.md
//! §2.3, §4.8): a short program that would need more memory, stack or time to read than a host
//! has is refused, with little held, on every host. This file counts what is allocated.

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};

struct Counting;
static NOW: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, l: Layout) -> *mut u8 {
        let p = unsafe { System.alloc(l) };
        if !p.is_null() {
            PEAK.fetch_max(NOW.fetch_add(l.size(), Relaxed) + l.size(), Relaxed);
        }
        p
    }
    unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
        unsafe { System.dealloc(p, l) };
        NOW.fetch_sub(l.size(), Relaxed);
    }
    unsafe fn realloc(&self, p: *mut u8, l: Layout, n: usize) -> *mut u8 {
        let q = unsafe { System.realloc(p, l, n) };
        if !q.is_null() {
            if n > l.size() {
                PEAK.fetch_max(NOW.fetch_add(n - l.size(), Relaxed) + n - l.size(), Relaxed);
            } else {
                NOW.fetch_sub(l.size() - n, Relaxed);
            }
        }
        q
    }
}

#[global_allocator]
static COUNTING: Counting = Counting;

/// The most a refused program may make a reader hold, beyond what it held before: no more
/// than a program at the bounds does (some 232 MiB for two million gates: its tokens, its
/// gates and its steps).
const HELD: usize = 256 << 20;

/// One reading at a time: the counts are the process's, and two tests reading at once would
/// each count the other's.
static ONE_AT_A_TIME: std::sync::Mutex<()> = std::sync::Mutex::new(());

/// What reading `src` held at most, and its refusal's code and message.
fn read(src: &str) -> (usize, Option<(&'static str, String)>) {
    let _one = ONE_AT_A_TIME.lock().unwrap_or_else(|e| e.into_inner());
    let before = NOW.load(Relaxed);
    PEAK.store(before, Relaxed);
    let refused = wai_wasmq::envelope::read(src, &Default::default()).err().map(|e| (e.code, e.message));
    (PEAK.load(Relaxed) - before, refused)
}

const HEAD: &str = "OPENQASM 3.0;\ninclude \"stdgates.inc\";\nextern f(bit[1]) -> bit[1];\nqubit[1] q;\nbit[1] d;\nd = f(d);\n";

/// Programs of a few hundred bytes that would each need more than a host has, and the
/// words of the bound each is refused by.
fn cases() -> Vec<(&'static str, String, &'static str)> {
    let deep = |open: &str, close: &str| format!("{HEAD}int x = {}1{};\n", open.repeat(100_000), close.repeat(100_000));
    let wide = |body: &str| format!("OPENQASM 3.0;\ninclude \"stdgates.inc\";\nextern f(bit[1]) -> bit[1];\nqubit[32768] q;\nbit[32768] c;\nbit[1] d;\nd = f(d);\n{body}");
    vec![
        ("a register of 2^30 bits", format!("{HEAD}bit[1073741824] c;\n"), "bits; the program is held"),
        ("bits past the program's", format!("{HEAD}bit[1048576] c;\n"), "bits; the program is held"),
        ("calls whose bits add up past what a reader holds", "OPENQASM 3.0;\ninclude \"stdgates.inc\";\nextern f(bit[1048575]) -> bit[1];\nqubit[1] q;\nbit[1048575] c;\nbit[1] d;\nfor uint i in [0:63] { d = f(c); }\n".into(), "the calls of externs"),
        ("steps past the bound, from resets of every qubit", wide("for uint i in [0:9999] { reset q; }\n"), "steps it is held to"),
        ("steps past the bound, from measurements of every qubit", wide("for uint i in [0:9999] { c = measure q; }\n"), "steps it is held to"),
        ("steps past the bound, from gates on every qubit", wide("for uint i in [0:9999] { h q; }\n"), "gates"),
        ("physical qubits past what the method holds", "OPENQASM 3.0;\ninclude \"stdgates.inc\";\nextern f(bit[1]) -> bit[1];\nbit[1] d;\nd = f(d);\nh $40000;\n".into(), "qubits"),
        ("qubits past what the method holds", "OPENQASM 3.0;\ninclude \"stdgates.inc\";\nextern f(bit[1]) -> bit[1];\nqubit[32769] q;\nbit[1] d;\nd = f(d);\n".into(), "qubits; the program is held"),
        ("a register past any reader's", "OPENQASM 3.0;\ninclude \"stdgates.inc\";\nextern f(bit[1]) -> bit[1];\nqubit[4294967295] q;\nbit[1] d;\nd = f(d);\n".into(), "qubits"),
        ("an array of 2^26 elements", format!("{HEAD}array[int[8], 1024, 65536] a;\n"), "an array of"),
        ("an array of a thousand registers of a million bits", format!("{HEAD}array[bit[1048575], 1024] a;\n"), "an array of"),
        ("reading that does not finish its work", format!("{HEAD}bit[1048575] a;\nfor uint i in [0:199] {{ a = ~a; }}\n"), "units of work"),
        ("brackets nested 100,000 deep", deep("(", ")"), "nests deeper"),
        ("unary operators nested 100,000 deep", format!("{HEAD}int x = {}1;\n", "-".repeat(100_000)), "nests deeper"),
        ("statements nested 50,000 deep", format!("{HEAD}{}x q[0];\n", "if (d == 1) ".repeat(50_000)), "nests deeper"),
        ("blocks nested 100,000 deep", format!("{HEAD}{}x q[0];{}\n", "{ ".repeat(100_000), " }".repeat(100_000)), "nests deeper"),
        ("a program of 2 MiB", format!("{HEAD}{}", "x q[0];\n".repeat(1 << 18)), "a reader takes at most"),
    ]
}

#[test]
fn a_program_past_the_bounds_is_refused_with_little_held() {
    let mut wrong = Vec::new();
    for (what, src, why) in cases() {
        let (held, refused) = read(&src);
        match refused {
            Some(("W016", message)) if message.contains(why) && held <= HELD => {}
            other => wrong.push(format!("{what}: refused as {other:?}, holding {} MiB, where it is refused for `{why}`", held >> 20)),
        }
    }
    assert!(wrong.is_empty(), "{wrong:#?}");
}

#[test]
fn a_program_at_the_bounds_is_read() {
    // The deepest nesting the reader takes, and a program at the bounds of qubits, bits and
    // steps: read, on every host (on wasm32 too, with its smaller stack).
    // Inside the reader's bound on nesting in the text (`qasm::parse::MAX_NESTING`, `Q006`),
    // with room for the statement and what holds the expression; past it, refused.
    let depth = wai_quantum::qasm::parse::MAX_NESTING - 8;
    let parens = format!("{HEAD}int x = {}1{};\n", "(".repeat(depth), ")".repeat(depth));
    assert!(read(&parens).1.is_none(), "{depth} deep");
    let unary = format!("{HEAD}int x = {}1;\n", "-".repeat(depth));
    assert!(read(&unary).1.is_none(), "{depth} unary operators");
    let ifs = format!("{HEAD}{}x q[0];\n", "if (d == 1) ".repeat(depth));
    assert!(read(&ifs).1.is_none(), "{depth} deep through statements");
    let past = format!("{HEAD}{}x q[0];\n", "if (d == 1) ".repeat(2 * wai_quantum::qasm::parse::MAX_NESTING));
    assert!(read(&past).1.is_some_and(|(_, m)| m.contains("Q006")), "past it");
    let wide = "OPENQASM 3.0;\ninclude \"stdgates.inc\";\nextern f(bit[1]) -> bit[1];\nqubit[32768] q;\nbit[1048575] c;\nbit[1] d;\nd = f(d);\nh q[0];\nc[0] = measure q[0];\n";
    let (held, refused) = read(wide);
    assert!(refused.is_none(), "the most qubits and bits: {refused:?}");
    assert!(held < 256 << 20, "{} MiB", held >> 20);
}