preprocess_slice_pdg.py

#!/usr/bin/env python3
"""
preprocess_slice_pdg.py — PDG backward-slice graphs from Devign LLVM IR.

§12 experiment: §11 (DFG-only backward slice) produced 56.64% — below the §7
baseline of 58.00%. Root cause: guard conditions are control dependence, not data
dependence. The `icmp`+`br` nodes that determine whether a dangerous call is safe
sit in a CFG predecessor block and have no DFG edge into the sink. DFG-only slicing
produces identical subgraphs for guarded (safe) and unguarded (vulnerable) code.

Fix: Program Dependence Graph (PDG) slice = DFG + control dependence.

Algorithm (fixed-point):
  1. Seed with dangerous sink nodes (same as §11 preprocess_slice.py).
  2. DFG backward BFS from all visited nodes.
  3. For each newly visited instruction node, add the terminator (br/switch) of
     each CFG predecessor block to the slice.
  4. Repeat until no new nodes are added.

In LLVM IR, `br i1 %cmp ...` has a DFG edge from `%cmp` (VK_INSTRUCTION operand).
Adding the `br` terminator automatically pulls in the `icmp` guard and its operands
via the next DFG BFS iteration — no special casing needed.

Output: data/{train,valid,test}_slice_pdg_graphs.pkl
  Same format as _slice_graphs.pkl — drop-in for train_slice_pdg.py.

Usage:
    python preprocess_slice_pdg.py --subset 200 --workers 1   # smoke test
    python preprocess_slice_pdg.py                             # full Devign
    python preprocess_slice_pdg.py --workers 8
"""

import argparse
import ctypes
import json
import pickle
import random
import re
import sys
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

import numpy as np
import llvmlite.binding as llvm

HERE = Path(__file__).parent
DATA = HERE / "data"

sys.path.insert(0, str(HERE))
from preprocess import compile_to_ir, download_devign

# ---------------------------------------------------------------------------
# Opcode vocabulary (identical to preprocess_instr.py — 110 entries)
# ---------------------------------------------------------------------------

OPCODE_VOCAB: dict[str, int] = {
    "add": 2,  "sub": 3,  "mul": 4,  "udiv": 5,  "sdiv": 6,
    "urem": 7, "srem": 8, "shl": 9,  "lshr": 10, "ashr": 11,
    "and": 12, "or": 13,  "xor": 14,
    "fadd": 15, "fsub": 16, "fmul": 17, "fdiv": 18, "frem": 19,
    "fneg": 20, "extractelement": 21, "insertelement": 22, "shufflevector": 23,
    "alloca": 26, "load": 27, "store": 28, "getelementptr": 29,
    "fence": 30, "cmpxchg": 31, "atomicrmw": 32,
    "br": 36, "switch": 37, "ret": 38, "invoke": 39,
    "resume": 40, "unreachable": 41, "indirectbr": 42, "callbr": 43,
    "icmp": 46, "fcmp": 47,
    "trunc": 48, "zext": 49, "sext": 50, "fptrunc": 51, "fpext": 52,
    "fptoui": 53, "fptosi": 54, "uitofp": 55, "sitofp": 56,
    "ptrtoint": 57, "inttoptr": 58, "bitcast": 59, "addrspacecast": 60,
    "phi": 61, "select": 62, "call": 63, "extractvalue": 64,
    "insertvalue": 65, "va_arg": 66, "landingpad": 67, "freeze": 68,
}
VOCAB_SIZE = 110

IDX_CONTEXT   = 0
IDX_ARGUMENT  = 1
IDX_MOCK      = 75
IDX_CONST_INT = 76
IDX_CONST_FP  = 77
IDX_UNDEF     = 78
IDX_UNKNOWN   = 79

_ICMP_PRED_RE = re.compile(r'\bicmp\s+(\w+)\b')
_FCMP_PRED_RE = re.compile(r'\bfcmp\s+(\w+)\b')

_ICMP_PRED_IDS: dict[str, int] = {
    "eq": 80,  "ne": 81,
    "slt": 82, "sle": 83, "sgt": 84, "sge": 85,
    "ult": 86, "ule": 87, "ugt": 88, "uge": 89,
}
_FCMP_PRED_IDS: dict[str, int] = {
    "false": 90, "oeq": 91, "ogt": 92, "oge": 93,
    "olt":  94,  "ole": 95, "one": 96, "ord": 97,
    "uno":  98,  "ueq": 99, "ugt": 100, "uge": 101,
    "ult":  102, "ule": 103, "une": 104, "true": 105,
}

VK_ARGUMENT     = 0
VK_BASIC_BLOCK  = 1
VK_FUNCTION     = 5
VK_GLOBAL_VAR   = 8
VK_UNDEF        = 14
VK_CONSTANT_INT = 18
VK_CONSTANT_FP  = 19
VK_INSTRUCTION  = 24
VK_POISON       = 25


def _instr_node_id(instr) -> int:
    op = instr.opcode
    if op == "icmp":
        m = _ICMP_PRED_RE.search(str(instr))
        if m:
            return _ICMP_PRED_IDS.get(m.group(1), IDX_UNKNOWN)
        return 46
    if op == "fcmp":
        m = _FCMP_PRED_RE.search(str(instr))
        if m:
            return _FCMP_PRED_IDS.get(m.group(1), IDX_UNKNOWN)
        return 47
    return OPCODE_VOCAB.get(op, IDX_UNKNOWN)


def _ptr_id(v) -> int:
    return ctypes.cast(v._ptr, ctypes.c_void_p).value


# ---------------------------------------------------------------------------
# Dangerous sink patterns (identical to preprocess_slice.py)
# ---------------------------------------------------------------------------

DANGEROUS_SINKS = frozenset({
    "strcpy", "strncpy", "strcat", "strncat",
    "memcpy", "memmove", "memset", "bcopy",
    "sprintf", "snprintf", "vsprintf", "vsnprintf",
    "gets", "fgets", "scanf", "sscanf", "fscanf",
    "read", "recv", "recvfrom", "pread",
    "malloc", "calloc", "realloc", "free", "xmalloc", "xrealloc",
    "printf", "fprintf", "syslog", "err", "warn",
    # Integer conversion — unchecked return used as array index / size is a
    # common vulnerability pattern (scar_atoi in scarnet, CWE-190/191).
    "atoi", "atol", "atoll", "atof",
    "strtol", "strtoul", "strtoll", "strtoull", "strtod",
})

# Functions whose return value is user-controlled / network-facing input.
# A mock node for any of these in the slice means external data reaches the sink.
INPUT_SOURCES = frozenset({
    "read", "recv", "recvfrom", "pread",
    "fgets", "fread", "getline", "getdelim",
    "scanf", "sscanf", "fscanf",
    "gets",
})

_SINK_SUFFIXES = tuple(DANGEROUS_SINKS)


def _normalize_sink_name(name: str) -> str:
    """Strip compiler-added decorations to recover the base function name.

    Clang's FORTIFY_SOURCE replaces e.g. memcpy with __memcpy_chk — the
    resulting IR call is no longer in DANGEROUS_SINKS by exact match.
    Strip leading underscores and trailing _chk / _chk_warn suffixes so
    __memcpy_chk normalizes to memcpy and is recognized.
    """
    name = name.lstrip("_")
    for suffix in ("_chk_warn", "_chk"):
        if name.endswith(suffix):
            name = name[:-len(suffix)]
            break
    return name


def _is_dangerous(name: str) -> bool:
    name = name.lstrip("@")
    if name in DANGEROUS_SINKS:
        return True
    # Normalize __foo_chk / __foo_chk_warn → foo before matching
    norm = _normalize_sink_name(name)
    if norm != name and norm in DANGEROUS_SINKS:
        return True
    for s in _SINK_SUFFIXES:
        if name.endswith(s) or name.endswith("_" + s):
            return True
    # LLVM memory intrinsics: e.g. llvm.memcpy.p0i8.p0i8.i64 / llvm.memcpy.p0.p0.i64
    for s in ("memcpy", "memmove", "memset", "bcopy"):
        if name.startswith(f"llvm.{s}."):
            return True
    return False


def _canonical_name(name: str) -> str:
    """Map IR callee name (including LLVM intrinsics) to canonical sink name."""
    name = name.lstrip("@")
    if name in DANGEROUS_SINKS:
        return name
    for s in ("memcpy", "memmove", "memset", "bcopy"):
        if name.startswith(f"llvm.{s}."):
            return s
    for s in _SINK_SUFFIXES:
        if name.endswith(s) or name.endswith("_" + s):
            return s
    return name


# ---------------------------------------------------------------------------
# PDG backward slice extractor
# ---------------------------------------------------------------------------

_CONSTANT_IDS = frozenset({IDX_CONST_INT, IDX_CONST_FP, IDX_UNDEF, IDX_CONTEXT})


def _extract_slice_pdg(x, edge_index, edge_type, mock_names,
                       instr_to_block, block_preds, block_last_instr):
    """
    PDG backward slice: DFG backward BFS + control dependence (block terminators).

    Fixed-point loop:
      1. DFG backward BFS from all visited nodes.
      2. For each newly visited instruction node, add the br/switch terminator
         of every CFG predecessor block (these have DFG edges from their icmp
         condition, so the guard comparison is pulled in on the next BFS pass).
      3. Repeat until stable.

    Returns None if no dangerous sinks found (caller falls back to full graph).
    """
    E = edge_index.shape[1] if edge_index.ndim == 2 and edge_index.shape[1] > 0 else 0

    fwd_dfg = defaultdict(list)
    rev_dfg = defaultdict(list)
    for i in range(E):
        if int(edge_type[i]) == 1:
            s, d = int(edge_index[0, i]), int(edge_index[1, i])
            fwd_dfg[s].append(d)
            rev_dfg[d].append(s)

    # Sink type 1: dangerous call sites
    dangerous_mocks = {nid for nid, nm in mock_names.items() if _is_dangerous(nm)}
    sink_ids:    set[int]       = set()
    sink_to_fn: dict[int, str] = {}   # old_node_id → dangerous function name
    for mid in dangerous_mocks:
        for consumer in fwd_dfg[mid]:
            if int(x[consumer, 0]) == 63:
                sink_ids.add(consumer)
                sink_to_fn[consumer] = _canonical_name(mock_names[mid])

    # Sink type 2: GEP or VLA alloca with non-constant operand
    # alloca(non-const) = variable-length array; same structural pattern as GEP
    for i in range(E):
        if int(edge_type[i]) == 1:
            s, d = int(edge_index[0, i]), int(edge_index[1, i])
            if int(x[d, 0]) in (29, 26) and int(x[s, 0]) not in _CONSTANT_IDS:
                sink_ids.add(d)

    if not sink_ids:
        return None

    visited      = set(sink_ids)
    ctrl_checked = set()

    changed = True
    while changed:
        changed = False

        # DFG backward BFS
        frontier = list(visited)
        while frontier:
            nxt = []
            for node in frontier:
                for pred in rev_dfg[node]:
                    if pred not in visited and pred != 0:
                        visited.add(pred)
                        nxt.append(pred)
                        changed = True
            frontier = nxt

        # Control dependence: add predecessor-block terminators for new nodes
        new_nodes = visited - ctrl_checked
        ctrl_checked |= new_nodes
        for node in new_nodes:
            block_id = instr_to_block.get(node)
            if block_id is None:
                continue
            for pred_block in block_preds.get(block_id, []):
                term_id = block_last_instr.get(pred_block)
                if term_id is not None and term_id not in visited and term_id != 0:
                    visited.add(term_id)
                    changed = True

    slice_nodes = sorted(visited)
    slice_size  = len(slice_nodes) + 1
    old_to_new  = {old: new + 1 for new, old in enumerate(slice_nodes)}

    if slice_size < 2:
        return None

    new_x = np.zeros((slice_size, 1), dtype=np.int64)
    new_x[0, 0] = IDX_CONTEXT
    for new_id, old_id in enumerate(slice_nodes, start=1):
        new_x[new_id, 0] = int(x[old_id, 0])

    # Map sink function names to new node indices
    sink_fn_names = {old_to_new[old_id]: fn
                     for old_id, fn in sink_to_fn.items()
                     if old_id in old_to_new}

    new_src, new_dst, new_et = [], [], []
    for i in range(E):
        et = int(edge_type[i])
        if et == 2:
            continue
        s, d = int(edge_index[0, i]), int(edge_index[1, i])
        if s in old_to_new and d in old_to_new:
            new_src.append(old_to_new[s])
            new_dst.append(old_to_new[d])
            new_et.append(et)

    for new_id in range(1, slice_size):
        new_src.extend([new_id, 0])
        new_dst.extend([0, new_id])
        new_et.extend([2, 2])

    new_edge_index = (np.array([new_src, new_dst], dtype=np.int64)
                      if new_src else np.zeros((2, 0), dtype=np.int64))
    new_edge_type  = (np.array(new_et, dtype=np.int64)
                      if new_et  else np.zeros(0, dtype=np.int64))

    # Collect input-source mock nodes that landed in the slice
    source_fn_names = {old_to_new[nid]: _canonical_name(mock_names[nid])
                       for nid, nm in mock_names.items()
                       if _canonical_name(nm) in INPUT_SOURCES
                       and nid in old_to_new}

    return {"x": new_x, "edge_index": new_edge_index, "edge_type": new_edge_type,
            "sink_fn_names": sink_fn_names, "source_fn_names": source_fn_names,
            "_sliced": True, "_n_sinks": len(sink_ids)}


# ---------------------------------------------------------------------------
# Graph builder — 5-pass algorithm + PDG slice extraction
# ---------------------------------------------------------------------------

def ir_to_graph_slice_pdg(ir_text, fn_name: str | None = None):
    """
    Build instruction-level graph then extract PDG backward slice.

    Additions over ir_to_graph_slice (§11):
    - Pass 1: instr_to_block tracks node_id → block ptr_id for instructions
    - Pass 2: block_preds and block_last_instr built alongside CFG edges
    - Calls _extract_slice_pdg() with control-dependence support

    fn_name: if given, select that specific function from a multi-function
             module. If None, picks the last non-declaration (single-function
             mode, original behaviour).

    Returns None if parsing fails or result has < 2 nodes.
    Caller adds 'y' and 'idx'.
    """
    try:
        mod = llvm.parse_assembly(ir_text)
    except Exception:
        return None

    target_fn = None
    for fn in mod.functions:
        if fn.is_declaration:
            continue
        if fn_name is None:
            target_fn = fn          # last non-declaration (original behaviour)
        elif fn.name == fn_name:
            target_fn = fn
            break
    if target_fn is None:
        return None

    # -- Pass 1: allocate nodes + track instruction→block membership ----------
    node_opcodes   = []
    ptr_to_id      = {}
    instr_to_block = {}   # node_id → block ptr_id (instructions only)
    node_counter   = 0

    node_opcodes.append(IDX_CONTEXT)
    node_counter = 1

    for arg in target_fn.arguments:
        ptr_to_id[_ptr_id(arg)] = node_counter
        node_opcodes.append(IDX_ARGUMENT)
        node_counter += 1

    block_first_instr = {}
    for block in target_fn.blocks:
        bpid = _ptr_id(block)
        first_in_block = True
        for instr in block.instructions:
            ipid = _ptr_id(instr)
            if first_in_block:
                block_first_instr[bpid] = node_counter
                first_in_block = False
            ptr_to_id[ipid]            = node_counter
            instr_to_block[node_counter] = bpid
            node_opcodes.append(_instr_node_id(instr))
            node_counter += 1

    if node_counter < 2:
        return None

    edges_src  = []
    edges_dst  = []
    edges_type = []

    # -- Pass 2: CFG edges + predecessor/terminator maps ----------------------
    block_preds      = defaultdict(list)
    block_last_instr = {}

    for block in target_fn.blocks:
        bpid    = _ptr_id(block)
        prev_id = None
        instrs  = list(block.instructions)
        for instr in instrs:
            cur_id = ptr_to_id[_ptr_id(instr)]
            if prev_id is not None:
                edges_src.append(prev_id)
                edges_dst.append(cur_id)
                edges_type.append(0)
            prev_id = cur_id
        if instrs:
            block_last_instr[bpid] = ptr_to_id[_ptr_id(instrs[-1])]
            terminator = instrs[-1]
            term_id    = ptr_to_id[_ptr_id(terminator)]
            for op in terminator.operands:
                if op.value_kind == VK_BASIC_BLOCK:
                    succ_bpid  = _ptr_id(op)
                    succ_first = block_first_instr.get(succ_bpid)
                    if succ_first is not None:
                        edges_src.append(term_id)
                        edges_dst.append(succ_first)
                        edges_type.append(0)
                    block_preds[succ_bpid].append(bpid)

    # -- Pass 3: DFG edges + mock name tracking --------------------------------
    constant_cache = {}
    mock_cache     = {}
    mock_names     = {}

    for block in target_fn.blocks:
        for instr in block.instructions:
            dst_id = ptr_to_id[_ptr_id(instr)]
            for op in instr.operands:
                vk = op.value_kind

                if vk == VK_INSTRUCTION or vk == VK_ARGUMENT:
                    src_id = ptr_to_id.get(_ptr_id(op))
                    if src_id is not None:
                        edges_src.append(src_id)
                        edges_dst.append(dst_id)
                        edges_type.append(1)

                elif vk == VK_CONSTANT_INT:
                    opid = _ptr_id(op)
                    if opid not in constant_cache:
                        constant_cache[opid] = node_counter
                        node_opcodes.append(IDX_CONST_INT)
                        node_counter += 1
                    edges_src.append(constant_cache[opid])
                    edges_dst.append(dst_id)
                    edges_type.append(1)

                elif vk == VK_CONSTANT_FP:
                    opid = _ptr_id(op)
                    if opid not in constant_cache:
                        constant_cache[opid] = node_counter
                        node_opcodes.append(IDX_CONST_FP)
                        node_counter += 1
                    edges_src.append(constant_cache[opid])
                    edges_dst.append(dst_id)
                    edges_type.append(1)

                elif vk in (VK_GLOBAL_VAR, VK_FUNCTION):
                    name = op.name
                    if name not in mock_cache:
                        mock_cache[name]     = node_counter
                        mock_names[node_counter] = name
                        node_opcodes.append(IDX_MOCK)
                        node_counter += 1
                    edges_src.append(mock_cache[name])
                    edges_dst.append(dst_id)
                    edges_type.append(1)

                elif vk in (VK_UNDEF, VK_POISON):
                    opid = _ptr_id(op)
                    if opid not in constant_cache:
                        constant_cache[opid] = node_counter
                        node_opcodes.append(IDX_UNDEF)
                        node_counter += 1
                    edges_src.append(constant_cache[opid])
                    edges_dst.append(dst_id)
                    edges_type.append(1)

    # -- Pass 4: global context edges (type 2) — bidirectional ----------------
    for i in range(1, node_counter):
        edges_src.extend([i, 0])
        edges_dst.extend([0, i])
        edges_type.extend([2, 2])

    x          = np.array(node_opcodes, dtype=np.int64).reshape(-1, 1)
    edge_index = (np.array([edges_src, edges_dst], dtype=np.int64)
                  if edges_src else np.zeros((2, 0), dtype=np.int64))
    edge_type  = (np.array(edges_type, dtype=np.int64)
                  if edges_type else np.zeros(0, dtype=np.int64))

    g = _extract_slice_pdg(x, edge_index, edge_type, mock_names,
                            instr_to_block, block_preds, block_last_instr)
    if g is None:
        g = {"x": x, "edge_index": edge_index, "edge_type": edge_type,
             "sink_fn_names": {}, "source_fn_names": {}, "_sliced": False, "_n_sinks": 0}

    return g


# ---------------------------------------------------------------------------
# Per-item processing
# ---------------------------------------------------------------------------

def process_item_slice_pdg(item):
    ir = compile_to_ir(item["func"])
    if ir is None:
        return None
    g = ir_to_graph_slice_pdg(ir)
    if g is None:
        return None
    g["y"]   = int(item["target"])
    g["idx"] = item.get("idx", 0)
    return g


def process_split_slice_pdg(jsonl_path, subset, workers, seed=42):
    with open(jsonl_path) as f:
        items = [json.loads(l) for l in f]

    rng = random.Random(seed)
    if subset:
        vuln  = [x for x in items if x["target"] == 1]
        fixed = [x for x in items if x["target"] == 0]
        rng.shuffle(vuln); rng.shuffle(fixed)
        items = vuln[:subset // 2] + fixed[:subset // 2]
    else:
        rng.shuffle(items)

    graphs, ok, fail = [], 0, 0
    total = len(items)
    print(f"  Processing {total} functions with {workers} workers ...")

    if workers == 1:
        for i, item in enumerate(items, 1):
            g = process_item_slice_pdg(item)
            if g:
                graphs.append(g); ok += 1
            else:
                fail += 1
            if i % 500 == 0:
                print(f"    {i}/{total}  ok={ok}  failed={fail}")
    else:
        with ProcessPoolExecutor(max_workers=workers) as pool:
            futs = {pool.submit(process_item_slice_pdg, it): it for it in items}
            for i, fut in enumerate(as_completed(futs), 1):
                g = fut.result()
                if g:
                    graphs.append(g); ok += 1
                else:
                    fail += 1
                if i % 500 == 0:
                    print(f"    {i}/{total}  ok={ok}  failed={fail}")

    attrition = fail / total * 100 if total > 0 else 0
    print(f"  Done: {ok} graphs built, {fail} failed ({attrition:.0f}% attrition)")

    node_counts = [g["x"].shape[0] for g in graphs]
    n_sliced    = sum(1 for g in graphs if g.get("_sliced", False))
    n_fallback  = ok - n_sliced
    if node_counts:
        print(f"  Slice stats: mean={np.mean(node_counts):.0f} nodes  "
              f"median={int(np.median(node_counts))}  max={max(node_counts)}")
        print(f"  Sliced: {n_sliced}/{ok} ({100*n_sliced/ok:.0f}%)  "
              f"Fallback (no sinks): {n_fallback}/{ok} ({100*n_fallback/ok:.0f}%)")

    for g in graphs:
        g.pop("_sliced", None)
        g.pop("_n_sinks", None)

    return graphs


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--subset",        type=int,  default=None)
    ap.add_argument("--workers",       type=int,  default=4)
    ap.add_argument("--seed",          type=int,  default=42)
    ap.add_argument("--skip-download", action="store_true")
    args = ap.parse_args()

    DATA.mkdir(parents=True, exist_ok=True)

    if not args.skip_download:
        missing = any(not (DATA / f"{s}.jsonl").exists()
                      for s in ["train", "valid", "test"])
        if missing:
            print("\n-- Download --------------------------------------------------")
            download_devign()
        else:
            print("  data/*.jsonl present, skipping download.")

    for split in ["train", "valid", "test"]:
        src = DATA / f"{split}.jsonl"
        dst = DATA / f"{split}_slice_pdg_graphs.pkl"
        if not src.exists():
            print(f"Missing {src} -- run preprocess.py or drop --skip-download.")
            sys.exit(1)
        print(f"\n-- {split} ---------------------------------------------------")
        graphs = process_split_slice_pdg(src, subset=args.subset,
                                         workers=args.workers, seed=args.seed)
        with open(dst, "wb") as f:
            pickle.dump(graphs, f)
        print(f"  Saved {len(graphs)} graphs -> {dst}")

    print("\nDone. Run train_slice_pdg.py next.\n")


if __name__ == "__main__":
    main()