//! The differential's first half: modules from wasm-smith, and this checker's verdicts.
//!
//!     cargo run --release --example differential -- --out DIR [--count N] [--mutants K] [--seed HEX]
//!
//! Writes `DIR/<i>.wasm` for i in 0..N (10,000 by default), K mutated copies of each
//! (`DIR/<i>-<k>.wasm`, 2 by default), and `DIR/verdicts.txt`, one line a module:
//! `<name> <caps> <integer verdict> | <float verdict>`. `differential.mjs DIR` then runs
//! check-module.mjs, the second checker, on every module at both levels under the same
//! caps, and names every verdict that differs.
//!
//! wasm-smith writes valid modules, or modules invalid in the ways it is told to be, so its
//! modules alone test how the two validators accept far more than how they refuse. Each
//! mutated copy changes one or two bytes of its module: a byte anywhere after the preamble
//! replaced, inserted or removed, or a byte of the code section replaced by a one-byte
//! opcode (`select`, `drop`, `end`, `ref.is_null`, ...). Most no longer decode or validate
//! (R2); the rest are other modules, which both checkers must give the same verdict.
//!
//! Each module is a function of the seed and its index alone. Its bytes come from a
//! splitmix64 stream seeded by both, read through wasm-smith's own `Config` and `Module`:
//!
//! - a quarter take wasm-smith's configuration as it comes, threads, the compact import
//!   section and its other proposals included, so many are outside WebAssembly 3.0 (R2);
//! - a quarter are shaped toward the profile, each knob with its own chance: the exports
//!   of minimal-int (§9.1) in name and type, no imports, one memory with a maximum, at most
//!   one table, no start export;
//! - half are shaped to it, their exports minimal-int's or, for some, a planted module's
//!   that differ in one export, and at most one knob turned the other way (imports, a
//!   second memory, a start function, two tables, tags), so that most reach
//!   the code section, where floats, relaxed vector instructions, exceptions, GC, tail calls
//!   and reference types are each on or off; their verdicts spread over held and every rule;
//! - one in eight may hold an invalid function body, one in four over-long LEB128s, and one
//!   in ten is checked under lower caps (one below its length, few functions, few pages).
//!
//! The seed is committed below; a run with another (`--seed`) is a new corpus.

use std::fmt::Write as _;
use std::path::PathBuf;

use arbitrary::{Arbitrary, Unstructured};
use wai_module_check::{check, Caps, Level};
use wasm_smith::{Config, Module};

/// The committed seed of the differential.
const SEED: u64 = 0x6465_742d_6d6f_6431;

/// The exports a profile-shaped module is given: minimal-int's, in name and type.
const INTERFACE: &[u8] = include_bytes!("../planted/minimal-int.wasm");

/// Other exports a module shaped to the profile may be given, each a planted module whose
/// exports are minimal-int's but for one: a fixed memory (held), and a 64-bit memory, a
/// memory with no maximum, an i64 parameter, a seventh export and an exported global.
const INTERFACES: [&[u8]; 6] = [
    include_bytes!("../planted/fixed-memory.wasm"),
    include_bytes!("../planted/memory64.wasm"),
    include_bytes!("../planted/no-maximum.wasm"),
    include_bytes!("../planted/i64-export.wasm"),
    include_bytes!("../planted/extra-export.wasm"),
    include_bytes!("../planted/exported-global.wasm"),
];

struct SplitMix(u64);

impl SplitMix {
    fn next(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^ (z >> 31)
    }

    fn bytes(&mut self, n: usize) -> Vec<u8> {
        let mut out = Vec::with_capacity(n + 8);
        while out.len() < n {
            out.extend_from_slice(&self.next().to_le_bytes());
        }
        out.truncate(n);
        out
    }

    /// True with chance `num` in `den`.
    fn chance(&mut self, num: u64, den: u64) -> bool {
        self.next() % den < num
    }
}

fn shaped(cfg: &mut Config, r: &mut SplitMix) {
    if r.chance(3, 4) {
        cfg.exports = Some(INTERFACE.to_vec());
    }
    if r.chance(3, 4) {
        cfg.max_imports = 0;
    }
    if r.chance(3, 4) {
        cfg.min_memories = 1;
        cfg.max_memories = 1;
        cfg.memory_max_size_required = true;
        cfg.max_memory32_bytes = 1 << (16 + r.next() % 8);
    }
    if r.chance(1, 2) {
        cfg.memory64_enabled = false;
    }
    if r.chance(3, 4) {
        cfg.max_tables = cfg.max_tables.min(1);
        cfg.table_max_size_required = true;
    }
    if r.chance(3, 4) {
        cfg.allow_start_export = false;
    }
    cfg.threads_enabled = r.chance(1, 10);
    cfg.compact_imports_enabled = r.chance(1, 10);
    cfg.allow_floats = r.chance(1, 2);
    cfg.relaxed_simd_enabled = cfg.simd_enabled && r.chance(1, 4);
    cfg.exceptions_enabled = r.chance(1, 4);
    cfg.tail_call_enabled = r.chance(1, 4);
    cfg.gc_enabled = cfg.reference_types_enabled && r.chance(1, 4);
}

fn strict(cfg: &mut Config, r: &mut SplitMix) {
    // The exports bring the module's one memory; wasm-smith adds none of its own.
    cfg.exports = Some(match r.next() % 16 {
        k @ 0..=5 => INTERFACES[k as usize].to_vec(),
        _ => INTERFACE.to_vec(),
    });
    cfg.max_imports = 0;
    cfg.min_memories = 0;
    cfg.max_memories = 0;
    cfg.memory64_enabled = false;
    cfg.max_tables = if r.chance(1, 2) { 0 } else { cfg.max_tables.min(1) };
    cfg.table_max_size_required = true;
    cfg.allow_start_export = false;
    cfg.max_tags = 0;
    cfg.threads_enabled = false;
    cfg.compact_imports_enabled = false;
    match r.next() % 10 {
        0 => cfg.max_imports = 2,
        1 => {
            cfg.min_memories = 1;
            cfg.max_memories = 1;
        }
        2 => cfg.allow_start_export = true,
        3 => cfg.max_tables = 2,
        4 => cfg.max_tags = 2,
        _ => {}
    }
    if r.chance(1, 2) {
        cfg.reference_types_enabled = false;
        cfg.max_tables = cfg.max_tables.min(1);
    }
    cfg.allow_floats = r.chance(1, 2);
    cfg.relaxed_simd_enabled = cfg.simd_enabled && r.chance(1, 3);
    cfg.exceptions_enabled = r.chance(1, 4);
    cfg.tail_call_enabled = r.chance(1, 4);
    cfg.gc_enabled = cfg.reference_types_enabled && r.chance(1, 4);
}

fn generate(seed: u64, index: u64) -> (Vec<u8>, Caps) {
    for attempt in 0u64.. {
        let mut r = SplitMix(seed ^ index.wrapping_mul(0xA24B_AED4_963E_E407) ^ attempt.wrapping_mul(0x9FB2_1C65_1E98_DF25));
        let shape = r.next() % 4;
        let len = 1024 + (r.next() % 32_768) as usize;
        let data = r.bytes(len);
        let mut u = Unstructured::new(&data);
        let Ok(mut cfg) = Config::arbitrary(&mut u) else { continue };
        match shape {
            0 => {
                // Two more proposals outside WebAssembly 3.0, which wasm-smith's own choice
                // leaves off.
                cfg.custom_page_sizes_enabled = r.chance(1, 8);
                cfg.wide_arithmetic_enabled = r.chance(1, 8);
            }
            1 => shaped(&mut cfg, &mut r),
            _ => strict(&mut cfg, &mut r),
        }
        cfg.allow_invalid_funcs = r.chance(1, 8);
        cfg.min_uleb_size = if r.chance(1, 4) { (1 + r.next() % 5) as u8 } else { 1 };
        cfg.generate_custom_sections = r.chance(1, 4);
        let Ok(module) = Module::new(cfg, &mut u) else { continue };
        let bytes = module.to_bytes();
        let mut caps = Caps::PROFILE;
        if r.chance(1, 10) {
            match r.next() % 3 {
                0 => caps.module_bytes = (bytes.len() as u64).saturating_sub(r.next() % 2),
                1 => caps.functions = r.next() % 8,
                _ => caps.memory_pages = r.next() % 4,
            }
        }
        return (bytes, caps);
    }
    unreachable!()
}

/// One-byte opcodes a mutation writes into a function body.
const OPCODES: [u8; 16] = [
    0x00, 0x01, 0x0b, 0x0f, 0x1a, 0x1b, 0x45, 0x50, 0x6a, 0x71, 0xa7, 0xac, 0xd1, 0xd3, 0xd4, 0x0d,
];

/// The content of the code section, if the module has one: its start and end.
fn code_section(bytes: &[u8]) -> Option<(usize, usize)> {
    let leb = |at: &mut usize| -> Option<usize> {
        let (mut v, mut shift) = (0usize, 0);
        loop {
            let c = *bytes.get(*at)?;
            *at += 1;
            v |= usize::from(c & 0x7f) << shift;
            if c & 0x80 == 0 {
                return Some(v);
            }
            shift += 7;
            if shift > 28 {
                return None;
            }
        }
    };
    let mut at = 8;
    while at < bytes.len() {
        let id = bytes[at];
        at += 1;
        let size = leb(&mut at)?;
        if id == 10 {
            return Some((at, (at + size).min(bytes.len())));
        }
        at += size;
    }
    None
}

/// A copy of `bytes` with one or two edits.
fn mutate(bytes: &[u8], r: &mut SplitMix) -> Vec<u8> {
    let mut out = bytes.to_vec();
    for _ in 0..1 + r.next() % 2 {
        if out.len() <= 8 {
            break;
        }
        let anywhere = 8 + (r.next() as usize) % (out.len() - 8);
        match r.next() % 8 {
            0..=3 => out[anywhere] = r.next() as u8,
            4 | 5 => match code_section(&out) {
                Some((start, end)) if end > start => {
                    let at = start + (r.next() as usize) % (end - start);
                    out[at] = OPCODES[(r.next() % OPCODES.len() as u64) as usize];
                }
                _ => out[anywhere] = r.next() as u8,
            },
            6 => out.insert(anywhere, r.next() as u8),
            _ => {
                out.remove(anywhere);
            }
        }
    }
    out
}

fn caps_text(c: &Caps) -> String {
    let mut parts = Vec::new();
    if c.module_bytes != Caps::PROFILE.module_bytes {
        parts.push(format!("module_bytes={}", c.module_bytes));
    }
    if c.functions != Caps::PROFILE.functions {
        parts.push(format!("functions={}", c.functions));
    }
    if c.memory_pages != Caps::PROFILE.memory_pages {
        parts.push(format!("memory_pages={}", c.memory_pages));
    }
    if parts.is_empty() {
        "-".into()
    } else {
        parts.join(",")
    }
}

fn main() {
    let (mut out, mut count, mut mutants, mut seed) = (None, 10_000u64, 2u64, SEED);
    let mut args = std::env::args().skip(1);
    while let Some(a) = args.next() {
        match a.as_str() {
            "--out" => out = args.next().map(PathBuf::from),
            "--count" => count = args.next().and_then(|n| n.parse().ok()).expect("--count N"),
            "--mutants" => mutants = args.next().and_then(|n| n.parse().ok()).expect("--mutants K"),
            "--seed" => {
                seed = args
                    .next()
                    .and_then(|s| u64::from_str_radix(s.trim_start_matches("0x"), 16).ok())
                    .expect("--seed HEX")
            }
            _ => panic!("usage: differential --out DIR [--count N] [--mutants K] [--seed HEX]"),
        }
    }
    let out = out.expect("--out DIR");
    std::fs::create_dir_all(&out).expect("the output directory");
    let mut lines = String::new();
    let mut write = |name: String, bytes: &[u8], caps: &Caps| {
        std::fs::write(out.join(format!("{name}.wasm")), bytes).expect("a module");
        let int = check(bytes, Level::Integer, caps);
        let float = check(bytes, Level::Float, caps);
        writeln!(lines, "{name} {} {int} | {float}", caps_text(caps)).unwrap();
    };
    for i in 0..count {
        let (bytes, caps) = generate(seed, i);
        write(i.to_string(), &bytes, &caps);
        let mut r = SplitMix(seed ^ i.wrapping_mul(0xD6E8_FEB8_6659_FD93) ^ 0x6d75_7461_6e74);
        for k in 0..mutants {
            write(format!("{i}-{k}"), &mutate(&bytes, &mut r), &caps);
        }
    }
    std::fs::write(out.join("verdicts.txt"), lines).expect("verdicts.txt");
    println!(
        "differential: {count} modules and {} mutated copies from seed {seed:#018x} in {}",
        count * mutants,
        out.display()
    );
}