neo-decompiler 0.11.0

use super::*;
use crate::decompiler::cfg::ssa::DominanceInfo;
use crate::decompiler::cfg::ssa::SsaForm;
use crate::decompiler::cfg::ssa::{SsaBlock, SsaExpr, SsaStmt, SsaVariable};
use crate::decompiler::cfg::{BasicBlock, BlockId, Cfg, EdgeKind, Terminator};
use crate::decompiler::ir::{BinOp, ControlFlow, Literal, Stmt};

fn v(base: &str, n: usize) -> SsaVariable {
    SsaVariable::new(base.to_string(), n)
}

/// Build a diamond: BB0 branches to BB1/BB2, both jump to BB3 (merge/ret).
fn diamond_cfg() -> Cfg {
    let mut cfg = Cfg::new();
    cfg.add_block(BasicBlock::new(
        BlockId(0),
        0,
        1,
        0..1,
        Terminator::Branch {
            then_target: BlockId(1),
            else_target: BlockId(2),
        },
    ));
    cfg.add_block(BasicBlock::new(
        BlockId(1),
        1,
        2,
        1..2,
        Terminator::Jump { target: BlockId(3) },
    ));
    cfg.add_block(BasicBlock::new(
        BlockId(2),
        2,
        3,
        2..3,
        Terminator::Jump { target: BlockId(3) },
    ));
    cfg.add_block(BasicBlock::new(BlockId(3), 3, 4, 3..4, Terminator::Return));
    cfg.add_edge(BlockId(0), BlockId(1), EdgeKind::ConditionalTrue);
    cfg.add_edge(BlockId(0), BlockId(2), EdgeKind::ConditionalFalse);
    cfg.add_edge(BlockId(1), BlockId(3), EdgeKind::Unconditional);
    cfg.add_edge(BlockId(2), BlockId(3), EdgeKind::Unconditional);
    cfg
}

fn block_with(stmts: Vec<SsaStmt>) -> SsaBlock {
    let mut b = SsaBlock::new();
    for s in stmts {
        b.add_stmt(s);
    }
    b
}

#[test]
fn structures_a_diamond_into_an_if_else() {
    let cfg = diamond_cfg();
    let mut blocks = std::collections::BTreeMap::new();
    // BB0: condition def  b0_0 = (loc0 < 1)
    blocks.insert(
        BlockId(0),
        block_with(vec![SsaStmt::assign(
            v("b0", 0),
            SsaExpr::binary(
                BinOp::Lt,
                SsaExpr::var(v("loc0", 0)),
                SsaExpr::lit(Literal::Int(1)),
            ),
        )]),
    );
    // BB1 (then): b1_0 = 10
    blocks.insert(
        BlockId(1),
        block_with(vec![SsaStmt::assign(
            v("b1", 0),
            SsaExpr::lit(Literal::Int(10)),
        )]),
    );
    // BB2 (else): b2_0 = 20
    blocks.insert(
        BlockId(2),
        block_with(vec![SsaStmt::assign(
            v("b2", 0),
            SsaExpr::lit(Literal::Int(20)),
        )]),
    );
    blocks.insert(BlockId(3), SsaBlock::new());

    let ssa = SsaForm {
        cfg,
        dominance: DominanceInfo::new(),
        blocks,
        definitions: std::collections::BTreeMap::new(),
        uses: std::collections::BTreeMap::new(),
    };

    let structured = structure(&ssa);

    // Expect: condition assign, then an If ControlFlow with both branches.
    let has_if = structured
        .stmts
        .iter()
        .any(|s| matches!(s, Stmt::ControlFlow(cf) if matches!(**cf, ControlFlow::If { .. })));
    assert!(
        has_if,
        "expected an If ControlFlow; got {:?}",
        structured.stmts
    );

    let if_cf = structured
        .stmts
        .iter()
        .rev()
        .find_map(|s| match s {
            Stmt::ControlFlow(cf) => Some(cf),
            _ => None,
        })
        .expect("an If node");
    let ControlFlow::If {
        then_branch,
        else_branch,
        ..
    } = if_cf.as_ref()
    else {
        panic!("expected If, got {if_cf:?}");
    };
    assert!(!then_branch.is_empty(), "then-branch should carry BB1 body");
    assert!(
        else_branch.is_some(),
        "an if-else diamond should yield an else branch"
    );
}

/// The if-condition must inline the comparison that drives the branch, not
/// render the bare reaching-definition variable. When BB0's last def is
/// `t = (loc0 < 1)`, the condition must be `(loc0 < 1)` and the def must
/// NOT be duplicated as a body statement.
#[test]
fn inlines_branch_comparison_condition_and_does_not_duplicate_it() {
    let mut cfg = Cfg::new();
    cfg.add_block(BasicBlock::new(
        BlockId(0),
        0,
        1,
        0..1,
        Terminator::Branch {
            then_target: BlockId(1),
            else_target: BlockId(2),
        },
    ));
    cfg.add_block(BasicBlock::new(
        BlockId(1),
        1,
        2,
        1..2,
        Terminator::Jump { target: BlockId(3) },
    ));
    cfg.add_block(BasicBlock::new(
        BlockId(2),
        2,
        3,
        2..3,
        Terminator::Jump { target: BlockId(3) },
    ));
    cfg.add_block(BasicBlock::new(BlockId(3), 3, 4, 3..4, Terminator::Return));
    cfg.add_edge(BlockId(0), BlockId(1), EdgeKind::ConditionalTrue);
    cfg.add_edge(BlockId(0), BlockId(2), EdgeKind::ConditionalFalse);
    cfg.add_edge(BlockId(1), BlockId(3), EdgeKind::Unconditional);
    cfg.add_edge(BlockId(2), BlockId(3), EdgeKind::Unconditional);

    let mut blocks = std::collections::BTreeMap::new();
    // BB0: only the comparison def — it IS the branch condition.
    blocks.insert(
        BlockId(0),
        block_with(vec![SsaStmt::assign(
            v("t", 0),
            SsaExpr::binary(
                BinOp::Lt,
                SsaExpr::var(v("loc0", 0)),
                SsaExpr::lit(Literal::Int(1)),
            ),
        )]),
    );
    blocks.insert(
        BlockId(1),
        block_with(vec![SsaStmt::assign(
            v("b1", 0),
            SsaExpr::lit(Literal::Int(10)),
        )]),
    );
    blocks.insert(
        BlockId(2),
        block_with(vec![SsaStmt::assign(
            v("b2", 0),
            SsaExpr::lit(Literal::Int(20)),
        )]),
    );
    blocks.insert(BlockId(3), SsaBlock::new());

    let ssa = SsaForm {
        cfg,
        dominance: DominanceInfo::new(),
        blocks,
        definitions: std::collections::BTreeMap::new(),
        uses: std::collections::BTreeMap::new(),
    };

    let structured = structure(&ssa);
    let rendered = crate::decompiler::ir::render_block(&structured, 0);

    // The condition must be the inlined comparison (versioned SSA name `loc0_0`,
    // double parens are a renderer quirk around a parenthesised Binary).
    assert!(
        rendered.contains("loc0 < 1"),
        "branch condition should inline the comparison; got:\n{rendered}"
    );
    // And it must NOT render the bare reaching-definition variable as the
    // condition.
    assert!(
        !rendered.contains("if (t_0)") && !rendered.contains("if (t)"),
        "branch condition should not be the bare t_0; got:\n{rendered}"
    );
    // The comparison def must not be duplicated as a body assignment.
    assert!(
        !rendered.contains("t_0 ="),
        "the comparison def must be consumed by the condition, not emitted in the body; got:\n{rendered}"
    );
}

#[test]
fn straight_line_cfg_emits_flat_block() {
    let mut cfg = Cfg::new();
    cfg.add_block(BasicBlock::new(BlockId(0), 0, 1, 0..2, Terminator::Return));
    let mut blocks = std::collections::BTreeMap::new();
    blocks.insert(
        BlockId(0),
        block_with(vec![
            SsaStmt::assign(v("b0", 0), SsaExpr::lit(Literal::Int(1))),
            SsaStmt::assign(v("b0", 1), SsaExpr::lit(Literal::Int(2))),
        ]),
    );
    let ssa = SsaForm {
        cfg,
        dominance: DominanceInfo::new(),
        blocks,
        definitions: std::collections::BTreeMap::new(),
        uses: std::collections::BTreeMap::new(),
    };
    let structured = structure(&ssa);
    let assigns = structured
        .stmts
        .iter()
        .filter(|s| matches!(s, Stmt::Assign { .. }))
        .count();
    // One assignment emitted directly; the last assignment becomes the
    // return value (the block's terminator is Return).
    assert_eq!(
        assigns,
        1,
        "first assignment emitted, last becomes return value:\n{}",
        crate::decompiler::ir::render_block(&structured, 0)
    );
    assert!(matches!(structured.stmts[0], Stmt::Assign { .. }));
    // The return statement should carry the last assignment's expression.
    assert!(matches!(structured.stmts[1], Stmt::Return(Some(_))));
}

/// A while loop: BB0 (header) branches to BB1 (body) / BB2 (exit); BB1
/// jumps back to BB0. dominance(BB0 ≥ BB1) makes BB0 a loop header.
#[test]
fn structures_a_back_edge_into_a_while_loop() {
    let mut cfg = Cfg::new();
    cfg.add_block(BasicBlock::new(
        BlockId(0),
        0,
        1,
        0..1,
        Terminator::Branch {
            then_target: BlockId(1),
            else_target: BlockId(2),
        },
    ));
    cfg.add_block(BasicBlock::new(
        BlockId(1),
        1,
        2,
        1..2,
        Terminator::Jump { target: BlockId(0) },
    ));
    cfg.add_block(BasicBlock::new(BlockId(2), 2, 3, 2..3, Terminator::Return));
    cfg.add_edge(BlockId(0), BlockId(1), EdgeKind::ConditionalTrue);
    cfg.add_edge(BlockId(0), BlockId(2), EdgeKind::ConditionalFalse);
    cfg.add_edge(BlockId(1), BlockId(0), EdgeKind::Unconditional);
    let dominance = crate::decompiler::cfg::ssa::compute(&cfg);

    let mut blocks = std::collections::BTreeMap::new();
    // header condition def: b0_0 = (loc0 < 3)
    blocks.insert(
        BlockId(0),
        block_with(vec![SsaStmt::assign(
            v("t", 0),
            SsaExpr::binary(
                BinOp::Lt,
                SsaExpr::var(v("loc0", 0)),
                SsaExpr::lit(Literal::Int(3)),
            ),
        )]),
    );
    // body: b1_0 = 1
    blocks.insert(
        BlockId(1),
        block_with(vec![SsaStmt::assign(
            v("t", 1),
            SsaExpr::lit(Literal::Int(1)),
        )]),
    );
    blocks.insert(BlockId(2), SsaBlock::new());

    let ssa = SsaForm {
        cfg,
        dominance,
        blocks,
        definitions: std::collections::BTreeMap::new(),
        uses: std::collections::BTreeMap::new(),
    };

    let structured = structure(&ssa);
    let has_while = structured
        .stmts
        .iter()
        .any(|s| matches!(s, Stmt::ControlFlow(cf) if matches!(**cf, ControlFlow::While { .. })));
    assert!(
        has_while,
        "a back-edge branch should structure as a While; got {:?}",
        structured.stmts
    );
}

/// A try/catch: TryEntry{body, catch, finally=None}; body and catch both
/// reach an EndTry whose continuation is the post-try block.
#[test]
fn structures_a_try_entry_into_try_catch() {
    let mut cfg = Cfg::new();
    cfg.add_block(BasicBlock::new(
        BlockId(0),
        0,
        1,
        0..1,
        Terminator::TryEntry {
            body_target: BlockId(1),
            catch_target: Some(BlockId(2)),
            finally_target: None,
        },
    ));
    cfg.add_block(BasicBlock::new(
        BlockId(1),
        1,
        2,
        1..2,
        Terminator::Jump { target: BlockId(3) },
    ));
    cfg.add_block(BasicBlock::new(
        BlockId(2),
        2,
        3,
        2..3,
        Terminator::Jump { target: BlockId(3) },
    ));
    cfg.add_block(BasicBlock::new(
        BlockId(3),
        3,
        4,
        3..4,
        Terminator::EndTry {
            continuation: BlockId(4),
        },
    ));
    cfg.add_block(BasicBlock::new(BlockId(4), 4, 5, 4..5, Terminator::Return));
    cfg.add_edge(BlockId(0), BlockId(1), EdgeKind::Unconditional);
    cfg.add_edge(BlockId(0), BlockId(2), EdgeKind::Unconditional);
    cfg.add_edge(BlockId(1), BlockId(3), EdgeKind::Unconditional);
    cfg.add_edge(BlockId(2), BlockId(3), EdgeKind::Unconditional);
    let dominance = crate::decompiler::cfg::ssa::compute(&cfg);

    let mut blocks = std::collections::BTreeMap::new();
    blocks.insert(
        BlockId(1),
        block_with(vec![SsaStmt::assign(
            v("t", 0),
            SsaExpr::lit(Literal::Int(1)),
        )]),
    );
    blocks.insert(
        BlockId(2),
        block_with(vec![SsaStmt::assign(
            v("t", 1),
            SsaExpr::lit(Literal::Int(2)),
        )]),
    );
    blocks.insert(BlockId(0), SsaBlock::new());
    blocks.insert(BlockId(3), SsaBlock::new());
    blocks.insert(BlockId(4), SsaBlock::new());

    let ssa = SsaForm {
        cfg,
        dominance,
        blocks,
        definitions: std::collections::BTreeMap::new(),
        uses: std::collections::BTreeMap::new(),
    };

    let structured = structure(&ssa);
    let has_try = structured.stmts.iter().any(
        |s| matches!(s, Stmt::ControlFlow(cf) if matches!(**cf, ControlFlow::TryCatch { .. })),
    );
    assert!(
        has_try,
        "a TryEntry should structure as TryCatch; got {:?}",
        structured.stmts
    );
}

/// A do-while: BB0 (body entry, falls through to the latch) is the loop
/// header; BB1 (latch) tests the condition and branches back to BB0 or out
/// to BB2. BB0 dominates BB1, so BB0 is a loop header whose terminator is
/// not a Branch → do-while.
#[test]
fn structures_a_bottom_tested_loop_into_do_while() {
    let mut cfg = Cfg::new();
    cfg.add_block(BasicBlock::new(
        BlockId(0),
        0,
        1,
        0..1,
        Terminator::Fallthrough { target: BlockId(1) },
    ));
    cfg.add_block(BasicBlock::new(
        BlockId(1),
        1,
        2,
        1..2,
        Terminator::Branch {
            then_target: BlockId(0),
            else_target: BlockId(2),
        },
    ));
    cfg.add_block(BasicBlock::new(BlockId(2), 2, 3, 2..3, Terminator::Return));
    cfg.add_edge(BlockId(0), BlockId(1), EdgeKind::Unconditional);
    cfg.add_edge(BlockId(1), BlockId(0), EdgeKind::ConditionalTrue);
    cfg.add_edge(BlockId(1), BlockId(2), EdgeKind::ConditionalFalse);
    let dominance = crate::decompiler::cfg::ssa::compute(&cfg);

    let mut blocks = std::collections::BTreeMap::new();
    // body: b0_0 = step()
    blocks.insert(
        BlockId(0),
        block_with(vec![SsaStmt::assign(
            v("t", 0),
            SsaExpr::call("step".to_string(), vec![]),
        )]),
    );
    // latch condition: b1_0 = (loc0 < 3)
    blocks.insert(
        BlockId(1),
        block_with(vec![SsaStmt::assign(
            v("t", 1),
            SsaExpr::binary(
                BinOp::Lt,
                SsaExpr::var(v("loc0", 0)),
                SsaExpr::lit(Literal::Int(3)),
            ),
        )]),
    );
    blocks.insert(BlockId(2), SsaBlock::new());

    let ssa = SsaForm {
        cfg,
        dominance,
        blocks,
        definitions: std::collections::BTreeMap::new(),
        uses: std::collections::BTreeMap::new(),
    };

    let structured = structure(&ssa);
    let has_dowhile = structured
        .stmts
        .iter()
        .any(|s| matches!(s, Stmt::ControlFlow(cf) if matches!(**cf, ControlFlow::DoWhile { .. })));
    assert!(
        has_dowhile,
        "a bottom-tested loop should structure as DoWhile; got {:?}",
        structured.stmts
    );
}

/// A switch: an equality cascade on one scrutinee. B0 compares `loc0 == 0`
/// → case0 body, else B1; B1 compares `loc0 == 1` → case1 body, else B2
/// (default); all bodies join at the merge B5.
#[test]
fn structures_an_equality_cascade_into_a_switch() {
    let mut cfg = Cfg::new();
    // B0: loc0 == 0 ? case0 : B1
    cfg.add_block(BasicBlock::new(
        BlockId(0),
        0,
        1,
        0..1,
        Terminator::Branch {
            then_target: BlockId(3),
            else_target: BlockId(1),
        },
    ));
    // B1: loc0 == 1 ? case1 : default(B2)
    cfg.add_block(BasicBlock::new(
        BlockId(1),
        1,
        2,
        1..2,
        Terminator::Branch {
            then_target: BlockId(4),
            else_target: BlockId(2),
        },
    ));
    // B2 (default) -> merge
    cfg.add_block(BasicBlock::new(
        BlockId(2),
        2,
        3,
        2..3,
        Terminator::Jump { target: BlockId(5) },
    ));
    // B3 (case0 body) -> merge
    cfg.add_block(BasicBlock::new(
        BlockId(3),
        3,
        4,
        3..4,
        Terminator::Jump { target: BlockId(5) },
    ));
    // B4 (case1 body) -> merge
    cfg.add_block(BasicBlock::new(
        BlockId(4),
        4,
        5,
        4..5,
        Terminator::Jump { target: BlockId(5) },
    ));
    // B5 (merge)
    cfg.add_block(BasicBlock::new(BlockId(5), 5, 6, 5..6, Terminator::Return));
    cfg.add_edge(BlockId(0), BlockId(3), EdgeKind::ConditionalTrue);
    cfg.add_edge(BlockId(0), BlockId(1), EdgeKind::ConditionalFalse);
    cfg.add_edge(BlockId(1), BlockId(4), EdgeKind::ConditionalTrue);
    cfg.add_edge(BlockId(1), BlockId(2), EdgeKind::ConditionalFalse);
    cfg.add_edge(BlockId(2), BlockId(5), EdgeKind::Unconditional);
    cfg.add_edge(BlockId(3), BlockId(5), EdgeKind::Unconditional);
    cfg.add_edge(BlockId(4), BlockId(5), EdgeKind::Unconditional);
    let dominance = crate::decompiler::cfg::ssa::compute(&cfg);

    let mut blocks = std::collections::BTreeMap::new();
    // B0: t0 = (loc0 == 0)
    blocks.insert(
        BlockId(0),
        block_with(vec![SsaStmt::assign(
            v("t", 0),
            SsaExpr::binary(
                BinOp::Eq,
                SsaExpr::var(v("loc0", 0)),
                SsaExpr::lit(Literal::Int(0)),
            ),
        )]),
    );
    // B1: t1 = (loc0 == 1)
    blocks.insert(
        BlockId(1),
        block_with(vec![SsaStmt::assign(
            v("t", 1),
            SsaExpr::binary(
                BinOp::Eq,
                SsaExpr::var(v("loc0", 1)),
                SsaExpr::lit(Literal::Int(1)),
            ),
        )]),
    );
    blocks.insert(
        BlockId(2),
        block_with(vec![SsaStmt::assign(
            v("loc0", 4),
            SsaExpr::lit(Literal::Int(12)),
        )]),
    );
    blocks.insert(
        BlockId(3),
        block_with(vec![SsaStmt::assign(
            v("loc0", 2),
            SsaExpr::lit(Literal::Int(10)),
        )]),
    );
    blocks.insert(
        BlockId(4),
        block_with(vec![SsaStmt::assign(
            v("loc0", 3),
            SsaExpr::lit(Literal::Int(11)),
        )]),
    );
    blocks.insert(BlockId(5), SsaBlock::new());

    let ssa = SsaForm {
        cfg,
        dominance,
        blocks,
        definitions: std::collections::BTreeMap::new(),
        uses: std::collections::BTreeMap::new(),
    };

    let structured = structure(&ssa);
    let has_switch = structured
        .stmts
        .iter()
        .any(|s| matches!(s, Stmt::ControlFlow(cf) if matches!(**cf, ControlFlow::Switch { .. })));
    assert!(
        has_switch,
        "an equality cascade on one scrutinee should structure as a Switch; got {:?}",
        structured.stmts
    );
}