use super::ref_types_module;
use super::skip_pooling_allocator_tests;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst};
use std::sync::Arc;
use wasmtime::*;

struct SetFlagOnDrop(Arc<AtomicBool>);

impl Drop for SetFlagOnDrop {
    fn drop(&mut self) {
        self.0.store(true, SeqCst);
    }
}

#[test]
#[cfg_attr(miri, ignore)]
fn smoke_test_gc() -> anyhow::Result<()> {
    smoke_test_gc_impl(false)
}

#[test]
#[cfg_attr(miri, ignore)]
fn smoke_test_gc_epochs() -> anyhow::Result<()> {
    smoke_test_gc_impl(true)
}

fn smoke_test_gc_impl(use_epochs: bool) -> anyhow::Result<()> {
    let (mut store, module) = ref_types_module(
        use_epochs,
        r#"
            (module
                (import "" "" (func $do_gc))
                (func $recursive (export "func") (param i32 externref) (result externref)
                    local.get 0
                    i32.eqz
                    if (result externref)
                        call $do_gc
                        local.get 1
                    else
                        local.get 0
                        i32.const 1
                        i32.sub
                        local.get 1
                        call $recursive
                    end
                )
            )
        "#,
    )?;

    let do_gc = Func::wrap(&mut store, |mut caller: Caller<'_, _>| {
        // Do a GC with `externref`s on the stack in Wasm frames.
        caller.gc();
    });
    let instance = Instance::new(&mut store, &module, &[do_gc.into()])?;
    let func = instance.get_func(&mut store, "func").unwrap();

    let inner_dropped = Arc::new(AtomicBool::new(false));
    let r = ExternRef::new(SetFlagOnDrop(inner_dropped.clone()));
    {
        let args = [Val::I32(5), Val::ExternRef(Some(r.clone()))];
        func.call(&mut store, &args, &mut [Val::I32(0)])?;
    }

    // Still held alive by the `VMExternRefActivationsTable` (potentially in
    // multiple slots within the table) and by this `r` local.
    assert!(r.strong_count() >= 2);

    // Doing a GC should see that there aren't any `externref`s on the stack in
    // Wasm frames anymore.
    store.gc();
    assert_eq!(r.strong_count(), 1);

    // Dropping `r` should drop the inner `SetFlagOnDrop` value.
    drop(r);
    assert!(inner_dropped.load(SeqCst));

    Ok(())
}

#[test]
#[cfg_attr(miri, ignore)]
fn wasm_dropping_refs() -> anyhow::Result<()> {
    let (mut store, module) = ref_types_module(
        false,
        r#"
            (module
                (func (export "drop_ref") (param externref)
                    nop
                )
            )
        "#,
    )?;

    let instance = Instance::new(&mut store, &module, &[])?;
    let drop_ref = instance.get_func(&mut store, "drop_ref").unwrap();

    let num_refs_dropped = Arc::new(AtomicUsize::new(0));

    // NB: 4096 is greater than the initial `VMExternRefActivationsTable`
    // capacity, so this will trigger at least one GC.
    for _ in 0..4096 {
        let r = ExternRef::new(CountDrops(num_refs_dropped.clone()));
        let args = [Val::ExternRef(Some(r))];
        drop_ref.call(&mut store, &args, &mut [])?;
    }

    assert!(num_refs_dropped.load(SeqCst) > 0);

    // And after doing a final GC, all the refs should have been dropped.
    store.gc();
    assert_eq!(num_refs_dropped.load(SeqCst), 4096);

    return Ok(());

    struct CountDrops(Arc<AtomicUsize>);

    impl Drop for CountDrops {
        fn drop(&mut self) {
            self.0.fetch_add(1, SeqCst);
        }
    }
}

#[test]
#[cfg_attr(miri, ignore)]
fn many_live_refs() -> anyhow::Result<()> {
    let mut wat = r#"
        (module
            ;; Make new `externref`s.
            (import "" "make_ref" (func $make_ref (result externref)))

            ;; Observe an `externref` so it is kept live.
            (import "" "observe_ref" (func $observe_ref (param externref)))

            (func (export "many_live_refs")
    "#
    .to_string();

    // This is more than the initial `VMExternRefActivationsTable` capacity, so
    // it will need to allocate additional bump chunks.
    const NUM_LIVE_REFS: usize = 1024;

    // Push `externref`s onto the stack.
    for _ in 0..NUM_LIVE_REFS {
        wat.push_str("(call $make_ref)\n");
    }

    // Pop `externref`s from the stack. Because we pass each of them to a
    // function call here, they are all live references for the duration of
    // their lifetimes.
    for _ in 0..NUM_LIVE_REFS {
        wat.push_str("(call $observe_ref)\n");
    }

    wat.push_str(
        "
            ) ;; func
        ) ;; module
        ",
    );

    let (mut store, module) = ref_types_module(false, &wat)?;

    let live_refs = Arc::new(AtomicUsize::new(0));

    let make_ref = Func::wrap(&mut store, {
        let live_refs = live_refs.clone();
        move || Some(ExternRef::new(CountLiveRefs::new(live_refs.clone())))
    });

    let observe_ref = Func::wrap(&mut store, |r: Option<ExternRef>| {
        let r = r.unwrap();
        let r = r.data().downcast_ref::<CountLiveRefs>().unwrap();
        assert!(r.live_refs.load(SeqCst) > 0);
    });

    let instance = Instance::new(&mut store, &module, &[make_ref.into(), observe_ref.into()])?;
    let many_live_refs = instance.get_func(&mut store, "many_live_refs").unwrap();

    many_live_refs.call(&mut store, &[], &mut [])?;

    store.gc();
    assert_eq!(live_refs.load(SeqCst), 0);

    return Ok(());

    struct CountLiveRefs {
        live_refs: Arc<AtomicUsize>,
    }

    impl CountLiveRefs {
        fn new(live_refs: Arc<AtomicUsize>) -> Self {
            live_refs.fetch_add(1, SeqCst);
            Self { live_refs }
        }
    }

    impl Drop for CountLiveRefs {
        fn drop(&mut self) {
            self.live_refs.fetch_sub(1, SeqCst);
        }
    }
}

#[test]
#[cfg_attr(miri, ignore)]
fn drop_externref_via_table_set() -> anyhow::Result<()> {
    let (mut store, module) = ref_types_module(
        false,
        r#"
            (module
                (table $t 1 externref)

                (func (export "table-set") (param externref)
                  (table.set $t (i32.const 0) (local.get 0))
                )
            )
        "#,
    )?;

    let instance = Instance::new(&mut store, &module, &[])?;
    let table_set = instance.get_func(&mut store, "table-set").unwrap();

    let foo_is_dropped = Arc::new(AtomicBool::new(false));
    let bar_is_dropped = Arc::new(AtomicBool::new(false));

    let foo = ExternRef::new(SetFlagOnDrop(foo_is_dropped.clone()));
    let bar = ExternRef::new(SetFlagOnDrop(bar_is_dropped.clone()));

    {
        let args = vec![Val::ExternRef(Some(foo))];
        table_set.call(&mut store, &args, &mut [])?;
    }
    store.gc();
    assert!(!foo_is_dropped.load(SeqCst));
    assert!(!bar_is_dropped.load(SeqCst));

    {
        let args = vec![Val::ExternRef(Some(bar))];
        table_set.call(&mut store, &args, &mut [])?;
    }
    store.gc();
    assert!(foo_is_dropped.load(SeqCst));
    assert!(!bar_is_dropped.load(SeqCst));

    table_set.call(&mut store, &[Val::ExternRef(None)], &mut [])?;
    assert!(foo_is_dropped.load(SeqCst));
    assert!(bar_is_dropped.load(SeqCst));

    Ok(())
}

#[test]
#[cfg_attr(miri, ignore)]
fn global_drops_externref() -> anyhow::Result<()> {
    test_engine(&Engine::default())?;

    if !skip_pooling_allocator_tests() {
        test_engine(&Engine::new(
            Config::new().allocation_strategy(InstanceAllocationStrategy::pooling()),
        )?)?;
    }

    return Ok(());

    fn test_engine(engine: &Engine) -> anyhow::Result<()> {
        let mut store = Store::new(&engine, ());
        let flag = Arc::new(AtomicBool::new(false));
        let externref = ExternRef::new(SetFlagOnDrop(flag.clone()));
        Global::new(
            &mut store,
            GlobalType::new(ValType::ExternRef, Mutability::Const),
            externref.into(),
        )?;
        drop(store);
        assert!(flag.load(SeqCst));

        let mut store = Store::new(&engine, ());
        let module = Module::new(
            &engine,
            r#"
                (module
                    (global (mut externref) (ref.null extern))

                    (func (export "run") (param externref)
                        local.get 0
                        global.set 0
                    )
                )
            "#,
        )?;
        let instance = Instance::new(&mut store, &module, &[])?;
        let run = instance.get_typed_func::<Option<ExternRef>, ()>(&mut store, "run")?;
        let flag = Arc::new(AtomicBool::new(false));
        let externref = ExternRef::new(SetFlagOnDrop(flag.clone()));
        run.call(&mut store, Some(externref))?;
        drop(store);
        assert!(flag.load(SeqCst));
        Ok(())
    }
}

#[test]
#[cfg_attr(miri, ignore)]
fn table_drops_externref() -> anyhow::Result<()> {
    test_engine(&Engine::default())?;

    if !skip_pooling_allocator_tests() {
        test_engine(&Engine::new(
            Config::new().allocation_strategy(InstanceAllocationStrategy::pooling()),
        )?)?;
    }

    return Ok(());

    fn test_engine(engine: &Engine) -> anyhow::Result<()> {
        let mut store = Store::new(&engine, ());
        let flag = Arc::new(AtomicBool::new(false));
        let externref = ExternRef::new(SetFlagOnDrop(flag.clone()));
        Table::new(
            &mut store,
            TableType::new(ValType::ExternRef, 1, None),
            externref.into(),
        )?;
        drop(store);
        assert!(flag.load(SeqCst));

        let mut store = Store::new(&engine, ());
        let module = Module::new(
            &engine,
            r#"
            (module
                (table 1 externref)

                (func (export "run") (param externref)
                    i32.const 0
                    local.get 0
                    table.set 0
                )
            )
        "#,
        )?;
        let instance = Instance::new(&mut store, &module, &[])?;
        let run = instance.get_typed_func::<Option<ExternRef>, ()>(&mut store, "run")?;
        let flag = Arc::new(AtomicBool::new(false));
        let externref = ExternRef::new(SetFlagOnDrop(flag.clone()));
        run.call(&mut store, Some(externref))?;
        drop(store);
        assert!(flag.load(SeqCst));
        Ok(())
    }
}

#[test]
#[cfg_attr(miri, ignore)]
fn gee_i_sure_hope_refcounting_is_atomic() -> anyhow::Result<()> {
    let mut config = Config::new();
    config.wasm_reference_types(true);
    config.epoch_interruption(true);
    let engine = Engine::new(&config)?;
    let mut store = Store::new(&engine, ());
    let module = Module::new(
        &engine,
        r#"
            (module
                (global (mut externref) (ref.null extern))
                (table 1 externref)

                (func (export "run") (param externref)
                    local.get 0
                    global.set 0
                    i32.const 0
                    local.get 0
                    table.set 0
                    loop
                        global.get 0
                        global.set 0

                        i32.const 0
                        i32.const 0
                        table.get
                        table.set

                        local.get 0
                        call $f

                        br 0
                    end
                )

                (func $f (param externref))
            )
        "#,
    )?;

    let instance = Instance::new(&mut store, &module, &[])?;
    let run = instance.get_typed_func::<Option<ExternRef>, ()>(&mut store, "run")?;

    let flag = Arc::new(AtomicBool::new(false));
    let externref = ExternRef::new(SetFlagOnDrop(flag.clone()));
    let externref2 = externref.clone();

    let child = std::thread::spawn(move || run.call(&mut store, Some(externref2)));

    for _ in 0..10000 {
        drop(externref.clone());
    }
    engine.increment_epoch();

    assert!(child.join().unwrap().is_err());
    assert!(!flag.load(SeqCst));
    assert_eq!(externref.strong_count(), 1);
    drop(externref);
    assert!(flag.load(SeqCst));

    Ok(())
}

#[test]
fn global_init_no_leak() -> anyhow::Result<()> {
    let (mut store, module) = ref_types_module(
        false,
        r#"
            (module
                (import "" "" (global externref))
                (global externref (global.get 0))
            )
        "#,
    )?;

    let externref = ExternRef::new(());
    let global = Global::new(
        &mut store,
        GlobalType::new(ValType::ExternRef, Mutability::Const),
        externref.clone().into(),
    )?;
    Instance::new(&mut store, &module, &[global.into()])?;
    drop(store);
    assert_eq!(externref.strong_count(), 1);

    Ok(())
}

#[test]
#[cfg_attr(miri, ignore)]
fn no_gc_middle_of_args() -> anyhow::Result<()> {
    let (mut store, module) = ref_types_module(
        false,
        r#"
            (module
                (import "" "return_some" (func $return (result externref externref externref)))
                (import "" "take_some" (func $take (param externref externref externref)))
                (func (export "run")
                    (local i32)
                    i32.const 1000
                    local.set 0
                    loop
                        call $return
                        call $take
                        local.get 0
                        i32.const -1
                        i32.add
                        local.tee 0
                        br_if 0
                    end
                )
            )
        "#,
    )?;

    let mut linker = Linker::new(store.engine());
    linker.func_wrap("", "return_some", || {
        (
            Some(ExternRef::new("a".to_string())),
            Some(ExternRef::new("b".to_string())),
            Some(ExternRef::new("c".to_string())),
        )
    })?;
    linker.func_wrap(
        "",
        "take_some",
        |a: Option<ExternRef>, b: Option<ExternRef>, c: Option<ExternRef>| {
            let a = a.unwrap();
            let b = b.unwrap();
            let c = c.unwrap();
            assert_eq!(a.data().downcast_ref::<String>().unwrap(), "a");
            assert_eq!(b.data().downcast_ref::<String>().unwrap(), "b");
            assert_eq!(c.data().downcast_ref::<String>().unwrap(), "c");
        },
    )?;

    let instance = linker.instantiate(&mut store, &module)?;
    let func = instance.get_typed_func::<(), ()>(&mut store, "run")?;
    func.call(&mut store, ())?;

    Ok(())
}

#[test]
#[cfg_attr(any(
    miri,
    // TODO(6530): s390x doesn't support tail calls yet.
    target_arch = "s390x"
), ignore)]
fn gc_and_tail_calls_and_stack_arguments() -> anyhow::Result<()> {
    // Test that GC refs in tail-calls' stack arguments get properly accounted
    // for in stack maps.
    //
    // What we do _not_ want to happen is for tail callers to be responsible for
    // including stack arguments in their stack maps (and therefore whether or
    // not they get marked at runtime). If that was the case, then we could have
    // the following scenario:
    //
    // * `f` calls `g` without any stack arguments,
    // * `g` tail calls `h` with GC ref stack arguments,
    // * and then `h` triggers a GC.
    //
    // Because `g`, who is responsible for including the GC refs in its stack
    // map in this hypothetical scenario, is no longer on the stack, we never
    // see its stack map, and therefore never mark the GC refs, and then we
    // collect them too early, and then we can get user-after-free bugs. Not
    // good! Note also that `f`, which is the frame that `h` will return to,
    // _cannot_ be responsible for including these stack arguments in its stack
    // map, because it has no idea what frame will be returning to it, and it
    // could be any number of different functions using that frame for long (and
    // indirect!) tail-call chains.
    //
    // In Cranelift we avoid this scenario because stack arguments are eagerly
    // loaded into virtual registers, and then when we insert a GC safe point,
    // we spill these virtual registers to the callee stack frame, and the stack
    // map includes entries for these stack slots.
    //
    // Nonetheless, this test exercises the above scenario just in case we do
    // something in the future like lazily load stack arguments into virtual
    // registers, to make sure that everything shows up in stack maps like they
    // are supposed to.

    let (mut store, module) = ref_types_module(
        false,
        r#"
            (module
                (import "" "make_some" (func $make (result externref externref externref)))
                (import "" "take_some" (func $take (param externref externref externref)))
                (import "" "gc" (func $gc))

                (func $stack_args (param externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref externref)
                  call $gc
                  ;; Make sure all these GC refs are live, so that they need to
                  ;; be put into the stack map.
                  local.get 0
                  local.get 1
                  local.get 2
                  call $take
                  local.get 3
                  local.get 4
                  local.get 5
                  call $take
                  local.get 6
                  local.get 7
                  local.get 8
                  call $take
                  local.get 9
                  local.get 10
                  local.get 11
                  call $take
                  local.get 12
                  local.get 13
                  local.get 14
                  call $take
                  local.get 15
                  local.get 16
                  local.get 17
                  call $take
                  local.get 18
                  local.get 19
                  local.get 20
                  call $take
                  local.get 21
                  local.get 22
                  local.get 23
                  call $take
                  local.get 24
                  local.get 25
                  local.get 26
                  call $take
                  local.get 27
                  local.get 28
                  local.get 29
                  call $take
                )

                (func $no_stack_args
                  call $make
                  call $make
                  call $make
                  call $make
                  call $make
                  call $make
                  call $make
                  call $make
                  call $make
                  call $make
                  return_call $stack_args
                )

                (func (export "run")
                    (local i32)
                    i32.const 1000
                    local.set 0
                    loop
                        call $no_stack_args
                        local.get 0
                        i32.const -1
                        i32.add
                        local.tee 0
                        br_if 0
                    end
                )
            )
        "#,
    )?;

    let mut linker = Linker::new(store.engine());
    linker.func_wrap("", "make_some", || {
        (
            Some(ExternRef::new("a".to_string())),
            Some(ExternRef::new("b".to_string())),
            Some(ExternRef::new("c".to_string())),
        )
    })?;
    linker.func_wrap(
        "",
        "take_some",
        |a: Option<ExternRef>, b: Option<ExternRef>, c: Option<ExternRef>| {
            let a = a.unwrap();
            let b = b.unwrap();
            let c = c.unwrap();
            assert_eq!(a.data().downcast_ref::<String>().unwrap(), "a");
            assert_eq!(b.data().downcast_ref::<String>().unwrap(), "b");
            assert_eq!(c.data().downcast_ref::<String>().unwrap(), "c");
        },
    )?;
    linker.func_wrap("", "gc", |mut caller: Caller<()>| {
        caller.gc();
    })?;

    let instance = linker.instantiate(&mut store, &module)?;
    let func = instance.get_typed_func::<(), ()>(&mut store, "run")?;
    func.call(&mut store, ())?;

    Ok(())
}