#![no_main]

use libfuzzer_sys::arbitrary::{self, Result, Unstructured};

use libfuzzer_sys::fuzz_target;

use std::sync::atomic::AtomicUsize;

use std::sync::atomic::Ordering::SeqCst;

use std::sync::{Mutex, Once};

use wasmtime_fuzzing::generators::{Config, DiffValue, DiffValueType, SingleInstModule};

use wasmtime_fuzzing::oracles::diff_wasmtime::WasmtimeInstance;

use wasmtime_fuzzing::oracles::engine::{build_allowed_env_list, parse_env_list};

use wasmtime_fuzzing::oracles::{DiffEqResult, differential, engine, log_wasm};

// Upper limit on the number of invocations for each WebAssembly function

// executed by this fuzz target.

const NUM_INVOCATIONS: usize = 5;

// Only run once when the fuzz target loads.

static SETUP: Once = Once::new();

// Environment-specified configuration for controlling the kinds of engines and

// modules used by this fuzz target. E.g.:

// - ALLOWED_ENGINES=wasmi,spec cargo +nightly fuzz run ...

// - ALLOWED_ENGINES=-v8 cargo +nightly fuzz run ...

// - ALLOWED_MODULES=single-inst cargo +nightly fuzz run ...

static ALLOWED_ENGINES: Mutex<Vec<Option<&str>>> = Mutex::new(vec![]);

static ALLOWED_MODULES: Mutex<Vec<Option<&str>>> = Mutex::new(vec![]);

// Statistics about what's actually getting executed during fuzzing

static STATS: RuntimeStats = RuntimeStats::new();

fuzz_target!(|data: &[u8]| {

    SETUP.call_once(|| {

        // To avoid a uncaught `SIGSEGV` due to signal handlers; see comments on

        // `setup_ocaml_runtime`.

        engine::setup_engine_runtimes();

        // Retrieve the configuration for this fuzz target from `ALLOWED_*`

        // environment variables.

        let allowed_engines = build_allowed_env_list(

            parse_env_list("ALLOWED_ENGINES"),

            &["wasmtime", "wasmi", "spec", "v8", "winch", "pulley"],

        );

        let allowed_modules = build_allowed_env_list(

            parse_env_list("ALLOWED_MODULES"),

            &["wasm-smith", "single-inst"],

        );

        *ALLOWED_ENGINES.lock().unwrap() = allowed_engines;

        *ALLOWED_MODULES.lock().unwrap() = allowed_modules;

    });

    // Errors in `run` have to do with not enough input in `data`, which we

    // ignore here since it doesn't affect how we'd like to fuzz.

    let _ = execute_one(&data);

});

fn execute_one(data: &[u8]) -> Result<()> {

    wasmtime_fuzzing::init_fuzzing();

    STATS.bump_attempts();

    let mut u = Unstructured::new(data);

    // Generate a Wasmtime and module configuration and update its settings

    // initially to be suitable for differential execution where the generated

    // wasm will behave the same in two different engines. This will get further

    // refined below.

    let mut config: Config = u.arbitrary()?;

    config.set_differential_config();

    let allowed_engines = ALLOWED_ENGINES.lock().unwrap();

    let allowed_modules = ALLOWED_MODULES.lock().unwrap();

    // Choose an engine that Wasmtime will be differentially executed against.

    // The chosen engine is then created, which might update `config`, and

    // returned as a trait object.

    let lhs = match *u.choose(&allowed_engines)? {

        Some(engine) => engine,

        None => {

            log::debug!("test case uses a runtime-disabled engine");

            return Ok(());

        }

    };

    log::trace!("Building LHS engine");

    let mut lhs = match engine::build(&mut u, lhs, &mut config)? {

        Some(engine) => engine,

        // The chosen engine does not have support compiled into the fuzzer,

        // discard this test case.

        None => return Ok(()),

    };

    log::debug!("lhs engine: {}", lhs.name());

    // Using the now-legalized module configuration generate the Wasm module;

    // this is specified by either the ALLOWED_MODULES environment variable or a

    // random selection between wasm-smith and single-inst.

    let build_wasm_smith_module = |u: &mut Unstructured, config: &Config| -> Result<_> {

        log::debug!("build wasm-smith with {config:?}");

        STATS.wasm_smith_modules.fetch_add(1, SeqCst);

        let module = config.generate(u, Some(1000))?;

        Ok(module.to_bytes())

    };

    let build_single_inst_module = |u: &mut Unstructured, config: &Config| -> Result<_> {

        log::debug!("build single-inst with {config:?}");

        STATS.single_instruction_modules.fetch_add(1, SeqCst);

        let module = SingleInstModule::new(u, &config.module_config)?;

        Ok(module.to_bytes())

    };

    if allowed_modules.is_empty() {

        panic!("unable to generate a module to fuzz against; check `ALLOWED_MODULES`")

    }

    let wasm = match *u.choose(&allowed_modules)? {

        Some("wasm-smith") => build_wasm_smith_module(&mut u, &config)?,

        Some("single-inst") => build_single_inst_module(&mut u, &config)?,

        None => {

            log::debug!("test case uses a runtime-disabled module strategy");

            return Ok(());

        }

        _ => unreachable!(),

    };

    log_wasm(&wasm);

    // Instantiate the generated wasm file in the chosen differential engine.

    let lhs_instance = lhs.instantiate(&wasm);

    STATS.bump_engine(lhs.name());

    // Always use Wasmtime as the second engine to instantiate within.

    log::debug!("Building RHS Wasmtime");

    let rhs_store = config.to_store();

    let rhs_module = wasmtime::Module::new(rhs_store.engine(), &wasm).unwrap();

    let rhs_instance = WasmtimeInstance::new(rhs_store, rhs_module);

    let (mut lhs_instance, mut rhs_instance) =

        match DiffEqResult::new(&*lhs, lhs_instance, rhs_instance) {

            // Both sides successful, continue below to invoking exports.

            DiffEqResult::Success(l, r) => (l, r),

            // Both sides failed, or computation has diverged. In both cases this

            // test case is done.

            DiffEqResult::Poisoned | DiffEqResult::Failed => return Ok(()),

        };

    // Call each exported function with different sets of arguments.

    'outer: for (name, signature) in rhs_instance.exported_functions() {

        let mut invocations = 0;

        loop {

            let arguments = match signature

                .params()

                .map(|ty| {

                    let ty = ty

                        .try_into()

                        .map_err(|_| arbitrary::Error::IncorrectFormat)?;

                    DiffValue::arbitrary_of_type(&mut u, ty)

                })

                .collect::<Result<Vec<_>>>()

            {

                Ok(args) => args,

                // This function signature isn't compatible with differential

                // fuzzing yet, try the next exported function in the meantime.

                Err(_) => continue 'outer,

            };

            let result_tys = match signature

                .results()

                .map(|ty| {

                    DiffValueType::try_from(ty).map_err(|_| arbitrary::Error::IncorrectFormat)

                })

                .collect::<Result<Vec<_>>>()

            {

                Ok(tys) => tys,

                // This function signature isn't compatible with differential

                // fuzzing yet, try the next exported function in the meantime.

                Err(_) => continue 'outer,

            };

            let ok = differential(

                lhs_instance.as_mut(),

                lhs.as_ref(),

                &mut rhs_instance,

                &name,

                &arguments,

                &result_tys,

            )

            .expect("failed to run differential evaluation");

            invocations += 1;

            STATS.total_invocations.fetch_add(1, SeqCst);

            // If this differential execution has resulted in the two instances

            // diverging in state we can't keep executing so don't execute any

            // more functions.

            if !ok {

                break 'outer;

            }

            // We evaluate the same function with different arguments until we

            // Hit a predetermined limit or we run out of unstructured data--it

            // does not make sense to re-evaluate the same arguments over and

            // over.

            if invocations > NUM_INVOCATIONS || u.is_empty() {

                break;

            }

        }

    }

    STATS.successes.fetch_add(1, SeqCst);

    Ok(())

}

#[derive(Default)]

struct RuntimeStats {

    /// Total number of fuzz inputs processed

    attempts: AtomicUsize,

    /// Number of times we've invoked engines

    total_invocations: AtomicUsize,

    /// Number of times a fuzz input finished all the way to the end without any

    /// sort of error (including `Arbitrary` errors)

    successes: AtomicUsize,

    // Counters for which engine was chosen

    wasmi: AtomicUsize,

    v8: AtomicUsize,

    spec: AtomicUsize,

    wasmtime: AtomicUsize,

    winch: AtomicUsize,

    pulley: AtomicUsize,

    // Counters for which style of module is chosen

    wasm_smith_modules: AtomicUsize,

    single_instruction_modules: AtomicUsize,

}

impl RuntimeStats {

    const fn new() -> RuntimeStats {

        RuntimeStats {

            attempts: AtomicUsize::new(0),

            total_invocations: AtomicUsize::new(0),

            successes: AtomicUsize::new(0),

            wasmi: AtomicUsize::new(0),

            v8: AtomicUsize::new(0),

            spec: AtomicUsize::new(0),

            wasmtime: AtomicUsize::new(0),

            winch: AtomicUsize::new(0),

            pulley: AtomicUsize::new(0),

            wasm_smith_modules: AtomicUsize::new(0),

            single_instruction_modules: AtomicUsize::new(0),

        }

    }

    fn bump_attempts(&self) {

        let attempts = self.attempts.fetch_add(1, SeqCst);

        if attempts == 0 || attempts % 1_000 != 0 {

            return;

        }

        let successes = self.successes.load(SeqCst);

        println!(

            "=== Execution rate ({} successes / {} attempted modules): {:.02}% ===",

            successes,

            attempts,

            successes as f64 / attempts as f64 * 100f64,

        );

        let v8 = self.v8.load(SeqCst);

        let spec = self.spec.load(SeqCst);

        let wasmi = self.wasmi.load(SeqCst);

        let wasmtime = self.wasmtime.load(SeqCst);

        let winch = self.winch.load(SeqCst);

        let pulley = self.pulley.load(SeqCst);

        let total = v8 + spec + wasmi + wasmtime + winch + pulley;

        println!(

            "\twasmi: {:.02}%, spec: {:.02}%, wasmtime: {:.02}%, v8: {:.02}%, \

             winch: {:.02}, \

             pulley: {:.02}%",

            wasmi as f64 / total as f64 * 100f64,

            spec as f64 / total as f64 * 100f64,

            wasmtime as f64 / total as f64 * 100f64,

            v8 as f64 / total as f64 * 100f64,

            winch as f64 / total as f64 * 100f64,

            pulley as f64 / total as f64 * 100f64,

        );

        let wasm_smith = self.wasm_smith_modules.load(SeqCst);

        let single_inst = self.single_instruction_modules.load(SeqCst);

        let total = wasm_smith + single_inst;

        println!(

            "\twasm-smith: {:.02}%, single-inst: {:.02}%",

            wasm_smith as f64 / total as f64 * 100f64,

            single_inst as f64 / total as f64 * 100f64,

        );

    }

    fn bump_engine(&self, name: &str) {

        match name {

            "wasmi" => self.wasmi.fetch_add(1, SeqCst),

            "wasmtime" => self.wasmtime.fetch_add(1, SeqCst),

            "spec" => self.spec.fetch_add(1, SeqCst),

            "v8" => self.v8.fetch_add(1, SeqCst),

            "winch" => self.winch.fetch_add(1, SeqCst),

            "pulley" => self.pulley.fetch_add(1, SeqCst),

            _ => return,

        };

    }

}