//! unroll_const_loops on the loop shape of a `for` over a small array with calls in its body
//! (sandbox reconcile's `for (shape, applied, wanted, solid) in [(..), (..)]`): the counter steps in
//! the block after the header test, a second test there leaves for the header's exit block (the
//! loaded Option's tag), the body has a `continue` (two latches) and an invoked call with an unwind
//! edge. The loop must be fully unrolled, and the RIR before and after must agree on an evaluator
//! that records every call (callee argument) in order. The body reads the array (a stack slot of
//! borrows: another slot's address is stored in it) at `slot + k * 8`: unrolled, those reads forward
//! and the borrowed slot stops escaping, the gain that lets the general rule unroll it. The same loop
//! over memory behind a pointer, or over a slot of plain values (box3d's scatter_bodies), stays rolled.
use relax_backend::regalloc::reachable;
use relax_backend::rir::*;

const U64: IntTy = IntTy { bits: 64, signed: false };

fn push(f: &mut Func, b: u32, i: Inst) {
    f.blocks[b as usize].insts.push(i);
    f.blocks[b as usize].pos.push(1);
}

/// Integer evaluator: a Host call appends its first argument to the trace. Returns (result, trace).
fn eval(f: &Func, args: &[u64]) -> (Vec<u64>, Vec<u64>) {
    let mut r = vec![0u64; f.vregs.len()];
    for (p, a) in f.params.iter().zip(args) {
        r[p.0 as usize] = *a;
    }
    let mut trace = Vec::new();
    let mut mem: std::collections::HashMap<u64, u64> = std::collections::HashMap::new();
    let mut b = 0usize;
    for _ in 0..100_000 {
        for i in &f.blocks[b].insts {
            match *i {
                Inst::Iconst(d, k) => r[d.0 as usize] = k as u64,
                Inst::SlotAddr(d, s) => r[d.0 as usize] = 0x1000 + 0x100 * s as u64,
                Inst::Store(_, base, off, src) => {
                    mem.insert(r[base.0 as usize].wrapping_add(off as i64 as u64), r[src.0 as usize]);
                }
                Inst::Load(_, d, base, off) => r[d.0 as usize] = mem.get(&r[base.0 as usize].wrapping_add(off as i64 as u64)).copied().unwrap_or(0),
                Inst::Mov(d, s) => r[d.0 as usize] = r[s.0 as usize],
                Inst::IBin(IOp::Add, _, d, a, c) => r[d.0 as usize] = r[a.0 as usize].wrapping_add(r[c.0 as usize]),
                Inst::IBinI(IOp::Add, _, d, a, k) => r[d.0 as usize] = r[a.0 as usize].wrapping_add(k as u64),
                Inst::IBinI(IOp::And, _, d, a, k) => r[d.0 as usize] = r[a.0 as usize] & k as u64,
                Inst::IBinI(IOp::Mul, _, d, a, k) => r[d.0 as usize] = r[a.0 as usize].wrapping_mul(k as u64),
                Inst::ICmpI(Cond::Eq, _, d, a, k) => r[d.0 as usize] = (r[a.0 as usize] == k as u64) as u64,
                Inst::ICmpI(Cond::Ne, _, d, a, k) => r[d.0 as usize] = (r[a.0 as usize] != k as u64) as u64,
                Inst::ICmpI(Cond::Lt, false, d, a, k) => r[d.0 as usize] = (r[a.0 as usize] < k as u64) as u64,
                Inst::Call(Callee::Host(_), ref a, _) => trace.push(r[a[0].0 as usize]),
                ref other => panic!("evaluator: {:?}", other),
            }
        }
        match f.blocks[b].term {
            Term::Jump(t) => b = t as usize,
            Term::Branch(c, t, e) => b = if r[c.0 as usize] != 0 { t as usize } else { e as usize },
            Term::Invoke(t, _) => b = t as usize,
            Term::Ret(ref v) => return (v.iter().map(|x| r[x.0 as usize]).collect(), trace),
            ref t => panic!("evaluator: {:?}", t),
        }
    }
    panic!("runaway")
}

/// b0: c = 0, acc = 0 -> b1
/// b1 (header): done = c == 3 -> b7 : b2
/// b2: k = c; c = c + 1; tag = (seed & 1) == 0 || k != 2 (the "loaded Some"); -> b3 : b7
/// b3: skip = (k + seed) & 1 -> b1 (continue: latch 1) : b4
/// b4: call host(k * 10 + seed + base + k * 8) -> invoke b5, unwind b8 (base: the array slot's
///     address, or the parameter `p` for `Base::Pointer`)
/// b5: acc = acc + k -> b1 (latch 2)
/// b7: ret acc
/// b8: (cleanup) ret acc
#[derive(Clone, Copy, PartialEq)]
enum Base {
    /// a slot holding another slot's address (an array of borrows)
    Borrows,
    /// a slot holding a plain value
    Values,
    /// a pointer parameter
    Pointer,
}

fn array_loop(kind: Base) -> Func {
    let mut f = Func::new("arr".into(), 0);
    let seed = f.vreg(Cls::I);
    f.params.push(seed);
    let p = f.vreg(Cls::I);
    f.params.push(p);
    let base = f.vreg(Cls::I);
    let (sc, ad) = (f.vreg(Cls::I), f.vreg(Cls::I));
    f.rets.push(Cls::I);
    for _ in 0..9 {
        f.block();
    }
    let v: Vec<VReg> = (0..12).map(|_| f.vreg(Cls::I)).collect();
    let (c, acc, done, k, t0, t1, tag, s1, skip, arg, t2, t3) = (v[0], v[1], v[2], v[3], v[4], v[5], v[6], v[7], v[8], v[9], v[10], v[11]);
    push(&mut f, 0, Inst::Iconst(c, 0));
    push(&mut f, 0, Inst::Iconst(acc, 0));
    if kind == Base::Pointer {
        push(&mut f, 0, Inst::Mov(base, p));
    } else {
        let s = f.slot(32, 8);
        push(&mut f, 0, Inst::SlotAddr(base, s));
        let x = f.vreg(Cls::I);
        if kind == Base::Borrows {
            let s2 = f.slot(8, 8);
            push(&mut f, 0, Inst::SlotAddr(x, s2));
        } else {
            push(&mut f, 0, Inst::Iconst(x, 5));
        }
        push(&mut f, 0, Inst::Store(Mem::Int(8, false), base, 0, x));
    }
    f.blocks[0].term = Term::Jump(1);
    push(&mut f, 1, Inst::ICmpI(Cond::Eq, false, done, c, 3));
    f.blocks[1].term = Term::Branch(done, 7, 2);
    push(&mut f, 2, Inst::Mov(k, c));
    push(&mut f, 2, Inst::IBinI(IOp::Add, U64, c, c, 1));
    push(&mut f, 2, Inst::IBinI(IOp::And, U64, t0, seed, 1));
    push(&mut f, 2, Inst::ICmpI(Cond::Eq, false, t1, t0, 0));
    push(&mut f, 2, Inst::ICmpI(Cond::Ne, false, t2, k, 2));
    push(&mut f, 2, Inst::IBin(IOp::Add, U64, t3, t1, t2));
    push(&mut f, 2, Inst::ICmpI(Cond::Ne, false, tag, t3, 0));
    f.blocks[2].term = Term::Branch(tag, 3, 7);
    push(&mut f, 3, Inst::IBin(IOp::Add, U64, s1, k, seed));
    push(&mut f, 3, Inst::IBinI(IOp::And, U64, skip, s1, 1));
    f.blocks[3].term = Term::Branch(skip, 1, 4);
    push(&mut f, 4, Inst::IBinI(IOp::Mul, U64, arg, k, 10));
    push(&mut f, 4, Inst::IBin(IOp::Add, U64, arg, arg, seed));
    push(&mut f, 4, Inst::IBinI(IOp::Mul, U64, sc, k, 8));
    push(&mut f, 4, Inst::IBin(IOp::Add, U64, ad, base, sc));
    push(&mut f, 4, Inst::IBin(IOp::Add, U64, arg, arg, ad));
    push(&mut f, 4, Inst::Call(Callee::Host(Sym::new(7)), vec![arg], vec![]));
    f.blocks[4].term = Term::Invoke(5, 8);
    push(&mut f, 5, Inst::IBin(IOp::Add, U64, acc, acc, k));
    f.blocks[5].term = Term::Jump(1);
    f.blocks[6].term = Term::Unreachable;
    f.blocks[7].term = Term::Ret(vec![acc]);
    f.blocks[8].term = Term::Ret(vec![acc]);
    for b in &mut f.blocks {
        b.pos = vec![1; b.insts.len()];
    }
    f
}

#[test]
fn a_loop_with_calls_two_latches_and_a_split_exit_test_unrolls() {
    let f = array_loop(Base::Borrows);
    f.verify().unwrap();
    let o = relax_session::Options::default();
    let mut g = f.clone();
    assert!(relax_codegen::loops::unroll_const_loops(&o, &mut g), "not unrolled");
    g.verify().unwrap();
    for seed in 0..8u64 {
        assert_eq!(eval(&f, &[seed, 0x5000]), eval(&g, &[seed, 0x5000]), "seed {}", seed);
    }
    // folded, the original loop is unreachable: every call left is in a copy, one per trip
    relax_backend::passes::simplify(&o, &mut g, None, true, true);
    g.verify().unwrap();
    let reach = reachable(&g);
    let preds = relax_backend::regalloc::preds_of(&g, &reach);
    assert!(relax_backend::regalloc::natural_loops(&g, &reach, &preds, false).is_empty(), "a loop is left:\n{}", g.dump());
    let calls = g.blocks.iter().enumerate().filter(|(bi, b)| reach[*bi] && b.insts.iter().any(|i| matches!(i, Inst::Call(..)))).count();
    assert_eq!(calls, 3, "one call per trip:\n{}", g.dump());
    for seed in 0..8u64 {
        assert_eq!(eval(&f, &[seed, 0x5000]), eval(&g, &[seed, 0x5000]), "seed {} after simplify", seed);
    }
}

#[test]
fn the_same_loop_over_a_pointer_or_a_slot_of_values_stays_rolled() {
    for kind in [Base::Pointer, Base::Values] {
        let f = array_loop(kind);
        f.verify().unwrap();
        let mut g = f.clone();
        assert!(!relax_codegen::loops::unroll_const_loops(&relax_session::Options::default(), &mut g), "unrolled:\n{}", g.dump());
    }
}

/// A small counted loop in front of the array loop: unrolled first, its copies add vregs the
/// array loop's gain check then scans (once an out-of-range index that dropped the O3 build). The
/// result must agree (the array loop's own copies stay unentered: the pass keeps the entries of its
/// start).
#[test]
fn the_gain_check_after_an_earlier_unrolled_loop_in_the_same_pass() {
    let mut f = array_loop(Base::Borrows);
    let (j, t, s) = (f.vreg(Cls::I), f.vreg(Cls::I), f.vreg(Cls::I));
    let pre = f.block();
    push(&mut f, 0, Inst::Iconst(j, 0));
    push(&mut f, 0, Inst::Iconst(s, 0));
    f.blocks[0].term = Term::Jump(pre);
    push(&mut f, pre, Inst::IBin(IOp::Add, U64, s, s, j));
    push(&mut f, pre, Inst::IBinI(IOp::Add, U64, j, j, 1));
    push(&mut f, pre, Inst::ICmpI(Cond::Lt, false, t, j, 2));
    f.blocks[pre as usize].term = Term::Branch(t, pre, 1);
    f.verify().unwrap();
    let o = relax_session::Options::default();
    let mut g = f.clone();
    assert!(relax_codegen::loops::unroll_const_loops(&o, &mut g), "not unrolled");
    g.verify().unwrap();
    for seed in 0..8u64 {
        assert_eq!(eval(&f, &[seed, 0x5000]), eval(&g, &[seed, 0x5000]), "seed {}", seed);
    }
}

/// camera_path_limit's slab test: `for a in 0..3` over a local array of plain values (`let d =
/// [dir.x, dir.y, dir.z]`), read at `slot + a * 8` and nowhere else but fixed-offset stores, whose two
/// latches also leave the loop (`continue 'entities` from the test that ends each iteration):
/// b0: s = slot; s[0], s[8], s[16] = seed, seed + 1, seed * 3; c = 0; acc = 0 -> b1
/// b1 (header): c < 3 -> b2 : b5
/// b2: k = c; c = c + 1; x = s[k * 8]; odd = x & 1 -> b6 : b3
/// b3: acc = acc + x; stop = acc == 7 -> b4 (leaves: latch 1) : b1
/// b6: acc = acc + 1; stop2 = acc == 5 -> b4 (latch 2) : b1
/// b4: ret acc + 1000;  b5: ret acc
/// Unrolled (the counter indexes the array: the copies read it at constant offsets), and the result
/// agrees; with MAKEPAD_RUN_OPT=-unrollearly it stays rolled.
fn slab_loop() -> Func {
    let mut f = Func::new("slab".into(), 0);
    let seed = f.vreg(Cls::I);
    f.params.push(seed);
    f.rets.push(Cls::I);
    for _ in 0..7 {
        f.block();
    }
    let v: Vec<VReg> = (0..16).map(|_| f.vreg(Cls::I)).collect();
    let (s, x1, x3, c, acc, lt, k, sh, ad, x, odd, stop, stop2, r4) = (v[0], v[1], v[2], v[3], v[4], v[5], v[6], v[7], v[8], v[9], v[10], v[11], v[12], v[13]);
    let sl = f.slot(24, 8);
    push(&mut f, 0, Inst::SlotAddr(s, sl));
    push(&mut f, 0, Inst::Store(Mem::Int(8, false), s, 0, seed));
    push(&mut f, 0, Inst::IBinI(IOp::Add, U64, x1, seed, 1));
    push(&mut f, 0, Inst::Store(Mem::Int(8, false), s, 8, x1));
    push(&mut f, 0, Inst::IBinI(IOp::Mul, U64, x3, seed, 3));
    push(&mut f, 0, Inst::Store(Mem::Int(8, false), s, 16, x3));
    push(&mut f, 0, Inst::Iconst(c, 0));
    push(&mut f, 0, Inst::Iconst(acc, 0));
    f.blocks[0].term = Term::Jump(1);
    push(&mut f, 1, Inst::ICmpI(Cond::Lt, false, lt, c, 3));
    f.blocks[1].term = Term::Branch(lt, 2, 5);
    push(&mut f, 2, Inst::Mov(k, c));
    push(&mut f, 2, Inst::IBinI(IOp::Add, U64, c, c, 1));
    push(&mut f, 2, Inst::IBinI(IOp::Mul, U64, sh, k, 8));
    push(&mut f, 2, Inst::IBin(IOp::Add, U64, ad, s, sh));
    push(&mut f, 2, Inst::Load(Mem::Int(8, false), x, ad, 0));
    push(&mut f, 2, Inst::IBinI(IOp::And, U64, odd, x, 1));
    f.blocks[2].term = Term::Branch(odd, 6, 3);
    push(&mut f, 3, Inst::IBin(IOp::Add, U64, acc, acc, x));
    push(&mut f, 3, Inst::ICmpI(Cond::Eq, false, stop, acc, 7));
    f.blocks[3].term = Term::Branch(stop, 4, 1);
    push(&mut f, 6, Inst::IBinI(IOp::Add, U64, acc, acc, 1));
    push(&mut f, 6, Inst::ICmpI(Cond::Eq, false, stop2, acc, 5));
    f.blocks[6].term = Term::Branch(stop2, 4, 1);
    push(&mut f, 4, Inst::IBinI(IOp::Add, U64, r4, acc, 1000));
    f.blocks[4].term = Term::Ret(vec![r4]);
    f.blocks[5].term = Term::Ret(vec![acc]);
    for b in &mut f.blocks {
        b.pos = vec![1; b.insts.len()];
    }
    f
}

#[test]
fn a_short_loop_over_a_local_array_of_values_whose_latches_leave_unrolls() {
    let f = slab_loop();
    f.verify().unwrap();
    let o = relax_session::Options::default();
    let mut g = f.clone();
    assert!(relax_codegen::loops::unroll_const_loops(&o, &mut g), "not unrolled:\n{}", f.dump());
    g.verify().unwrap();
    for seed in 0..16u64 {
        assert_eq!(eval(&f, &[seed]), eval(&g, &[seed]), "seed {}", seed);
    }
    relax_backend::passes::simplify(&o, &mut g, None, true, true);
    g.verify().unwrap();
    let reach = reachable(&g);
    let preds = relax_backend::regalloc::preds_of(&g, &reach);
    assert!(relax_backend::regalloc::natural_loops(&g, &reach, &preds, false).is_empty(), "a loop is left:\n{}", g.dump());
    for seed in 0..16u64 {
        assert_eq!(eval(&f, &[seed]), eval(&g, &[seed]), "seed {} after simplify", seed);
    }
    let mut off = relax_session::Options::default();
    off.opt = vec!["-unrollearly".to_string()];
    let mut h = f.clone();
    assert!(!relax_codegen::loops::unroll_const_loops(&off, &mut h), "unrolled with -unrollearly");
}