Kai Zhang

"""
Unit tests for cfg_features.py — all new CFG feature computations.

These tests use plain Python data structures (index-based adjacency lists)
and require no Binary Ninja dependency.
"""
import pytest

from redb.extractors.decompiler.bninja.analysis.cfg_features import (
    bfs_order,
    bfs_max_depth,
    count_back_edges,
    compute_topology_hash,
    compute_md_index_topdown,
    compute_md_index_bottomup,
    compute_prime_product,
    build_block_features,
    compute_cfg_feature_tlsh,
    compute_wl_minhash,
    pack_adjacency,
    LLIL_OP_CATEGORIES,
    CAT_ARITHMETIC,
    CAT_LOGIC,
    CAT_CALL,
    CAT_MEMORY,
    NUM_WL_MINHASH_PERMS,
)


# ===================================================================
# Helper: common graph topologies
# ===================================================================

def _linear_chain(n):
    """0 -> 1 -> 2 -> ... -> (n-1)"""
    return [[i + 1] if i < n - 1 else [] for i in range(n)]


def _diamond():
    """
    0 -> 1, 0 -> 2, 1 -> 3, 2 -> 3
    (classic if/else diamond)
    """
    return [[1, 2], [3], [3], []]


def _predecessors_from_successors(successors, n):
    preds = [[] for _ in range(n)]
    for src, targets in enumerate(successors):
        for tgt in targets:
            preds[tgt].append(src)
    return preds


# ===================================================================
# TestBfsOrder
# ===================================================================

class TestBfsOrder:
    def test_empty_graph(self):
        assert bfs_order([], 0) == []

    def test_single_node(self):
        assert bfs_order([[]], 1) == [0]

    def test_linear_chain(self):
        succs = _linear_chain(4)
        assert bfs_order(succs, 4) == [0, 1, 2, 3]

    def test_diamond(self):
        succs = _diamond()
        order = bfs_order(succs, 4)
        assert order[0] == 0
        assert order[-1] == 3
        assert set(order) == {0, 1, 2, 3}

    def test_unreachable_nodes(self):
        # 0 -> 1, node 2 is unreachable
        succs = [[1], [], []]
        order = bfs_order(succs, 3)
        assert order[:2] == [0, 1]
        assert 2 in order  # unreachable appended

    def test_all_nodes_visited(self):
        succs = _diamond()
        order = bfs_order(succs, 4)
        assert len(order) == 4


# ===================================================================
# TestBfsMaxDepth
# ===================================================================

class TestBfsMaxDepth:
    def test_empty_graph(self):
        assert bfs_max_depth([], 0) == 0

    def test_single_block(self):
        assert bfs_max_depth([[]], 1) == 0

    def test_linear_chain(self):
        succs = _linear_chain(5)
        assert bfs_max_depth(succs, 5) == 4

    def test_diamond(self):
        succs = _diamond()
        assert bfs_max_depth(succs, 4) == 2

    def test_wide_graph(self):
        # 0 -> 1, 0 -> 2, 0 -> 3 (all at depth 1)
        succs = [[1, 2, 3], [], [], []]
        assert bfs_max_depth(succs, 4) == 1


# ===================================================================
# TestCountBackEdges
# ===================================================================

class TestCountBackEdges:
    def test_empty_graph(self):
        assert count_back_edges([], 0) == 0

    def test_no_loops(self):
        succs = _linear_chain(3)
        assert count_back_edges(succs, 3) == 0

    def test_single_loop(self):
        # 0 -> 1 -> 2 -> 0 (one back edge: 2->0)
        succs = [[1], [2], [0]]
        assert count_back_edges(succs, 3) == 1

    def test_nested_loops(self):
        # 0 -> 1 -> 2 -> 1 (inner), 2 -> 3 -> 0 (outer)
        succs = [[1], [2], [1, 3], [0]]
        assert count_back_edges(succs, 4) == 2

    def test_self_loop(self):
        # 0 -> 0 (self-loop)
        succs = [[0]]
        assert count_back_edges(succs, 1) == 1

    def test_diamond_no_loops(self):
        succs = _diamond()
        assert count_back_edges(succs, 4) == 0

    def test_single_node_no_loop(self):
        succs = [[]]
        assert count_back_edges(succs, 1) == 0


# ===================================================================
# TestTopologyHash
# ===================================================================

class TestTopologyHash:
    def test_same_graph_same_hash(self):
        succs = _diamond()
        bfs = bfs_order(succs, 4)
        h1 = compute_topology_hash(succs, bfs, 4)
        h2 = compute_topology_hash(succs, bfs, 4)
        assert h1 == h2

    def test_different_graphs_different_hash(self):
        succs1 = _linear_chain(3)
        bfs1 = bfs_order(succs1, 3)
        h1 = compute_topology_hash(succs1, bfs1, 3)

        succs2 = _diamond()
        bfs2 = bfs_order(succs2, 4)
        h2 = compute_topology_hash(succs2, bfs2, 4)

        assert h1 != h2

    def test_returns_16_bytes(self):
        succs = _diamond()
        bfs = bfs_order(succs, 4)
        h = compute_topology_hash(succs, bfs, 4)
        assert isinstance(h, bytes)
        assert len(h) == 16

    def test_isomorphic_graphs_same_hash(self):
        # Graph A: 0->1, 0->2, 1->3, 2->3 (diamond with successors [1,2])
        succs_a = [[1, 2], [3], [3], []]
        # Graph B: same structure but successors listed as [2,1]
        # BFS from 0 will visit them in different order, but after remapping
        # the canonical form should be identical for isomorphic graphs
        succs_b = [[2, 1], [3], [3], []]

        bfs_a = bfs_order(succs_a, 4)
        bfs_b = bfs_order(succs_b, 4)

        h_a = compute_topology_hash(succs_a, bfs_a, 4)
        h_b = compute_topology_hash(succs_b, bfs_b, 4)
        assert h_a == h_b

    def test_empty_graph(self):
        h = compute_topology_hash([], [], 0)
        assert h == b'\x00' * 16

    def test_single_node(self):
        succs = [[]]
        bfs = bfs_order(succs, 1)
        h = compute_topology_hash(succs, bfs, 1)
        assert isinstance(h, bytes)
        assert len(h) == 16


# ===================================================================
# TestMdIndex
# ===================================================================

class TestMdIndex:
    def test_single_block_topdown(self):
        succs = [[]]
        preds = [[]]
        bfs = [0]
        result = compute_md_index_topdown(succs, preds, bfs)
        assert isinstance(result, int)
        assert result > 0

    def test_single_block_bottomup(self):
        succs = [[]]
        preds = [[]]
        result = compute_md_index_bottomup(succs, preds, 1)
        assert isinstance(result, int)
        assert result > 0

    def test_linear_chain_topdown_vs_bottomup(self):
        succs = _linear_chain(4)
        preds = _predecessors_from_successors(succs, 4)
        bfs = bfs_order(succs, 4)
        td = compute_md_index_topdown(succs, preds, bfs)
        bu = compute_md_index_bottomup(succs, preds, 4)
        # Top-down and bottom-up should be different for a linear chain
        # (entry has in_deg=0, exit has out_deg=0, so the sequences differ)
        assert td != bu

    def test_deterministic(self):
        succs = _diamond()
        preds = _predecessors_from_successors(succs, 4)
        bfs = bfs_order(succs, 4)
        td1 = compute_md_index_topdown(succs, preds, bfs)
        td2 = compute_md_index_topdown(succs, preds, bfs)
        assert td1 == td2

    def test_different_graphs_different_index(self):
        succs1 = _linear_chain(3)
        preds1 = _predecessors_from_successors(succs1, 3)
        bfs1 = bfs_order(succs1, 3)
        td1 = compute_md_index_topdown(succs1, preds1, bfs1)

        succs2 = _diamond()
        preds2 = _predecessors_from_successors(succs2, 4)
        bfs2 = bfs_order(succs2, 4)
        td2 = compute_md_index_topdown(succs2, preds2, bfs2)

        assert td1 != td2

    def test_topdown_empty(self):
        assert compute_md_index_topdown([], [], []) == 0

    def test_bottomup_empty(self):
        assert compute_md_index_bottomup([], [], 0) == 0


# ===================================================================
# TestPrimeProduct
# ===================================================================

class TestPrimeProduct:
    def test_empty(self):
        assert compute_prime_product([]) == 0

    def test_known_sequence(self):
        # Use actual LLIL enum values from conftest_binja_stubs:
        # LLIL_NOP=0 -> prime 1, LLIL_LOAD=4 -> prime 5
        from redb.extractors.decompiler.bninja.analysis.cfg_features import LLIL_OP_PRIMES
        nop_val = 0   # LLIL_NOP
        load_val = 4  # LLIL_LOAD
        expected = LLIL_OP_PRIMES.get(nop_val, 1) * LLIL_OP_PRIMES.get(load_val, 1)
        result = compute_prime_product([nop_val, load_val])
        assert result == expected

    def test_order_independence(self):
        # LLIL_LOAD=4, LLIL_STORE=5, LLIL_ADD=13
        ops_a = [4, 5, 13]
        ops_b = [13, 4, 5]
        assert compute_prime_product(ops_a) == compute_prime_product(ops_b)

    def test_unknown_ops_map_to_1(self):
        # Unknown ops get prime 1, so they don't change the product
        result_known = compute_prime_product([4])  # LLIL_LOAD -> 5
        result_with_unknown = compute_prime_product([4, 9999])  # LOAD * unknown(1)
        assert result_known == result_with_unknown

    def test_mod_2_64(self):
        # Product should be mod 2^64
        result = compute_prime_product([4] * 1000)  # LLIL_LOAD
        assert 0 <= result < 2**64

    def test_single_op(self):
        # LLIL_STORE=5 -> prime 7
        assert compute_prime_product([5]) == 7


# ===================================================================
# TestBuildBlockFeatures
# ===================================================================

class TestBuildBlockFeatures:
    def test_empty_llil(self):
        succs = [[1], []]
        features = build_block_features([[], []], succs, 2)
        assert len(features) == 2
        # All zeros except successor_count
        assert features[0] == [0, 0, 0, 0, 0, 0, 0, 1]  # 1 successor
        assert features[1] == [0, 0, 0, 0, 0, 0, 0, 0]  # 0 successors

    def test_correct_categorization(self):
        # Set up categories for testing
        import redb.extractors.decompiler.bninja.analysis.cfg_features as cf
        old_cats = cf.LLIL_OP_CATEGORIES.copy()
        cf.LLIL_OP_CATEGORIES.update({
            100: CAT_ARITHMETIC,
            101: CAT_ARITHMETIC,
            200: CAT_LOGIC,
            300: CAT_CALL,
            400: CAT_MEMORY,
        })
        try:
            block_ops = [[100, 101, 200, 300, 400]]
            succs = [[]]
            features = build_block_features(block_ops, succs, 1)
            assert features[0][0] == 5   # instr_count
            assert features[0][1] == 2   # arithmetic
            assert features[0][2] == 1   # logic
            assert features[0][4] == 1   # call
            assert features[0][6] == 1   # memory
        finally:
            cf.LLIL_OP_CATEGORIES.clear()
            cf.LLIL_OP_CATEGORIES.update(old_cats)

    def test_cap_at_65535(self):
        # More than 65535 ops in one block
        huge_ops = [0] * 70000  # NOP x 70000
        succs = [[]]
        features = build_block_features([huge_ops], succs, 1)
        assert features[0][0] == 65535  # capped

    def test_missing_block_ops(self):
        # block_llil_ops shorter than n
        succs = [[1], [2], []]
        features = build_block_features([[1, 2]], succs, 3)
        assert len(features) == 3
        # Block 1 and 2 get empty ops since block_llil_ops only has 1 entry
        assert features[1] == [0, 0, 0, 0, 0, 0, 0, 1]
        assert features[2] == [0, 0, 0, 0, 0, 0, 0, 0]


# ===================================================================
# TestCfgFeatureTlsh
# ===================================================================

class TestCfgFeatureTlsh:
    def test_too_few_blocks_returns_none(self):
        # 5 blocks = 5 * 9 bytes = 45 < 50
        bb_features = [[10, 1, 0, 2, 0, 1, 1, 2]] * 5
        bfs = list(range(5))
        result = compute_cfg_feature_tlsh(bb_features, bfs)
        assert result is None

    def test_uniform_data_returns_none(self):
        # 7 identical blocks — TLSH returns TNULL for low-entropy input
        bb_features = [[10, 1, 0, 2, 0, 1, 1, 2]] * 7
        bfs = list(range(7))
        result = compute_cfg_feature_tlsh(bb_features, bfs)
        assert result is None

    def test_varied_data_returns_string(self):
        # 20 blocks with varied features — enough entropy for TLSH
        bb_features = [
            [i * 7 + 3, (i * 13) % 50, (i * 17) % 30, (i * 23) % 40,
             (i * 11) % 20, (i * 7) % 25, (i * 19) % 35, (i * 3) % 10]
            for i in range(20)
        ]
        bfs = list(range(20))
        result = compute_cfg_feature_tlsh(bb_features, bfs)
        assert isinstance(result, str)
        assert len(result) > 0
        assert result.startswith("T1")


# ===================================================================
# TestWlMinhash
# ===================================================================

class TestWlMinhash:
    def test_empty_function(self):
        result = compute_wl_minhash([], [], [], 0)
        assert result == [255] * NUM_WL_MINHASH_PERMS

    def test_returns_128_elements(self):
        succs = _diamond()
        preds = _predecessors_from_successors(succs, 4)
        bb_feats = [[5, 1, 0, 2, 0, 1, 1, 2]] * 4
        result = compute_wl_minhash(succs, preds, bb_feats, 4)
        assert len(result) == 128

    def test_all_uint8(self):
        succs = _linear_chain(3)
        preds = _predecessors_from_successors(succs, 3)
        bb_feats = [[3, 1, 0, 1, 0, 0, 1, 1]] * 3
        result = compute_wl_minhash(succs, preds, bb_feats, 3)
        assert all(0 <= v <= 255 for v in result)

    def test_identical_graphs_same_signature(self):
        succs = _diamond()
        preds = _predecessors_from_successors(succs, 4)
        bb_feats = [[5, 1, 0, 2, 0, 1, 1, 2]] * 4
        sig1 = compute_wl_minhash(succs, preds, bb_feats, 4)
        sig2 = compute_wl_minhash(succs, preds, bb_feats, 4)
        assert sig1 == sig2

    def test_different_graphs_different_signatures(self):
        # Graph 1: linear chain
        succs1 = _linear_chain(4)
        preds1 = _predecessors_from_successors(succs1, 4)
        bb_feats1 = [[5, 1, 0, 2, 0, 1, 1, i] for i in range(4)]
        sig1 = compute_wl_minhash(succs1, preds1, bb_feats1, 4)

        # Graph 2: diamond
        succs2 = _diamond()
        preds2 = _predecessors_from_successors(succs2, 4)
        bb_feats2 = [[10, 3, 2, 1, 0, 0, 0, i] for i in range(4)]
        sig2 = compute_wl_minhash(succs2, preds2, bb_feats2, 4)

        assert sig1 != sig2

    def test_single_node(self):
        succs = [[]]
        preds = [[]]
        bb_feats = [[1, 0, 0, 0, 0, 0, 0, 0]]
        result = compute_wl_minhash(succs, preds, bb_feats, 1)
        assert len(result) == 128


# ===================================================================
# TestPackAdjacency
# ===================================================================

class TestPackAdjacency:
    def test_empty(self):
        assert pack_adjacency([]) == []

    def test_single_edge(self):
        succs = [[1], []]
        edges = pack_adjacency(succs)
        assert len(edges) == 1
        assert edges[0] == (0 << 16) | 1

    def test_correct_packing(self):
        succs = _diamond()
        edges = pack_adjacency(succs)
        assert len(edges) == 4
        # 0->1, 0->2, 1->3, 2->3
        expected = {
            (0 << 16) | 1,
            (0 << 16) | 2,
            (1 << 16) | 3,
            (2 << 16) | 3,
        }
        assert set(edges) == expected

    def test_roundtrip(self):
        """Unpack edges and verify source/target pairs."""
        succs = [[1, 2], [3], [3], []]
        edges = pack_adjacency(succs)
        unpacked = [(e >> 16, e & 0xFFFF) for e in edges]
        expected = [(0, 1), (0, 2), (1, 3), (2, 3)]
        assert sorted(unpacked) == sorted(expected)

    def test_large_index_filtered(self):
        # Create a successor list where index >= 65536
        succs = [[] for _ in range(65537)]
        succs[0] = [65536]  # target is exactly 65536 — should be filtered
        edges = pack_adjacency(succs)
        assert len(edges) == 0

    def test_max_valid_index(self):
        # Index 65535 is the maximum valid
        succs = [[] for _ in range(65536)]
        succs[0] = [65535]
        edges = pack_adjacency(succs)
        assert len(edges) == 1
        assert edges[0] == (0 << 16) | 65535