"""Build compact tree-sitter def-use/reachability features for TRAIN + VALID.""" from __future__ import annotations import argparse import json import signal from collections import Counter, deque from multiprocessing import Pool from pathlib import Path import pandas as pd from src.cpg.defuse import build_defuse_graph from src.cpg.trace_miner import _extract_param_names, _known_dangerous_call_names, is_source_candidate TIMEOUT_S = 15 TYPE_NAMES = ("call_expression", "assignment_expression", "declaration", "init_declarator", "return_statement", "if_statement", "for_statement", "while_statement", "switch_statement", "expression_statement", "binary_expression") DANGEROUS_CALL_NAMES = _known_dangerous_call_names(None) class RowTimeout(Exception): pass def _alarm(signum, frame): raise RowTimeout() def feature_row(task): rec_id, pair_id, target, code = task signal.signal(signal.SIGALRM, _alarm) signal.alarm(TIMEOUT_S) try: graph = build_defuse_graph(code) params = _extract_param_names(code) nodes = [s.node_id for s in graph.statements] node_set = set(nodes) out = {n: [] for n in nodes} indegree = Counter() for a, b in graph.edges: if a in node_set and b in node_set: out[a].append(b) indegree[b] += 1 sources = {n for n in nodes if is_source_candidate(graph, n, params)} distances = {n: 0 for n in sources} queue = deque(sources) while queue: node = queue.popleft() for nxt in out[node]: if nxt not in distances: distances[nxt] = distances[node] + 1 queue.append(nxt) dangerous_sinks = set() call_count = 0 for stmt in graph.statements: if stmt.node_type == "call_expression": call_count += 1 if graph.uses.get(stmt.node_id, set()) & DANGEROUS_CALL_NAMES: dangerous_sinks.add(stmt.node_id) reachable_dangerous = [distances[n] for n in dangerous_sinks if n in distances] n_nodes, n_edges = len(nodes), len(graph.edges) degrees_out = [len(out[n]) for n in nodes] degrees_in = [indegree[n] for n in nodes] n_def, n_use = sum(bool(v) for v in graph.defines.values()), sum(bool(v) for v in graph.uses.values()) features = { "n_statements": n_nodes, "n_defuse_edges": n_edges, "edges_per_statement": n_edges / max(1, n_nodes), "n_sources": len(sources), "source_fraction": len(sources) / max(1, n_nodes), "n_reachable_statements": len(distances), "reachable_fraction": len(distances) / max(1, n_nodes), "n_calls": call_count, "n_dangerous_call_candidates": len(dangerous_sinks), "n_reachable_dangerous_calls": len(reachable_dangerous), "dangerous_reachability_fraction": len(reachable_dangerous) / max(1, len(dangerous_sinks)), "has_source_to_dangerous_call": int(bool(reachable_dangerous)), "shortest_source_dangerous_distance": min(reachable_dangerous) if reachable_dangerous else 0, "mean_source_dangerous_distance": sum(reachable_dangerous) / max(1, len(reachable_dangerous)), "n_defining_statements": n_def, "n_using_statements": n_use, "defined_fraction": n_def / max(1, n_nodes), "used_fraction": n_use / max(1, n_nodes), "max_in_degree": max(degrees_in, default=0), "mean_in_degree": sum(degrees_in) / max(1, n_nodes), "max_out_degree": max(degrees_out, default=0), "mean_out_degree": sum(degrees_out) / max(1, n_nodes), "n_control_dependent": len(graph.control_parent), } type_counts = Counter(s.node_type for s in graph.statements) for name in TYPE_NAMES: features[f"type_{name}"] = type_counts[name] return {"sample_id": rec_id, "pair_id": pair_id, "target": target, "status": "ok", "features": features} except RowTimeout: return {"sample_id": rec_id, "pair_id": pair_id, "target": target, "status": "timeout"} except Exception as exc: return {"sample_id": rec_id, "pair_id": pair_id, "target": target, "status": "error", "error": f"{type(exc).__name__}: {exc}"} finally: signal.alarm(0) def main(): p = argparse.ArgumentParser(description=__doc__) p.add_argument("--corpus", type=Path, default=Path("data/processed/corpus.parquet")) p.add_argument("--train", type=Path, default=Path("data/splits/chronological/train.json")) p.add_argument("--valid", type=Path, default=Path("data/splits/chronological/valid.json")) p.add_argument("--output", type=Path, default=Path("results/phase7/graph_feature_table.jsonl")) p.add_argument("--workers", type=int, default=8) a = p.parse_args() train_ids, valid_ids = set(json.loads(a.train.read_text())), set(json.loads(a.valid.read_text())) if train_ids & valid_ids: raise ValueError("TRAIN/VALID pair overlap") df = pd.read_parquet(a.corpus, columns=["sample_id", "pair_id", "target", "func_vuln", "func_patch"]) df = df[df.pair_id.isin(train_ids | valid_ids)] tasks = [] for row in df.to_dict("records"): code = row["func_vuln"] if int(row["target"]) == 1 else row["func_patch"] tasks.append((row["sample_id"], row["pair_id"], int(row["target"]), str(code))) a.output.parent.mkdir(parents=True, exist_ok=True) results = [] with Pool(max(1, a.workers)) as pool: for i, rec in enumerate(pool.imap_unordered(feature_row, tasks, chunksize=5), 1): results.append(rec) if i % 500 == 0: print(f"built {i}/{len(tasks)}", flush=True) results.sort(key=lambda x: (x["pair_id"], x["target"])) with a.output.open("w") as f: for rec in results: f.write(json.dumps(rec, sort_keys=True) + "\n") counts = Counter(r["status"] for r in results) report = {"n_rows": len(results), "n_pairs": len(train_ids | valid_ids), "train_pairs": len(train_ids), "valid_pairs": len(valid_ids), "statuses": dict(counts), "timeout_seconds": TIMEOUT_S, "features_per_row": len(next((r["features"] for r in results if r["status"] == "ok"), {})), "test_read": False} report_path = a.output.with_suffix(".report.json") report_path.write_text(json.dumps(report, indent=2) + "\n") print(json.dumps(report, indent=2)) if counts["timeout"] + counts["error"] > 0.02 * len(results): raise RuntimeError("graph feature failure rate exceeds 2%; do not train classifier") if __name__ == "__main__": main()