Add files using upload-large-folder tool

5bfc214 verified about 1 month ago

4.44 kB

"""Phase 2 per-sample annotation (spec §4.2-§4.4): ties statement_mapper,
defuse, and trace_miner together to produce, for one corpus row, the L1/L2/L3
localization targets and trace ground truth.
Uses the tree-sitter fallback (defuse.py) by default — not real Joern — since
running Joern per-sample across the full ~13.6k-pair corpus is a real,
multi-hour, CPU-parallelized batch job (spec §4: "budget accordingly...
target 16 workers... resumable") that hasn't been executed yet (see
DECISIONS.md). `annotate_sample(..., use_joern=True)` is wired for when that
batch job runs; this module's job is to define *what* an annotation record
looks like and compute it correctly, independent of which backend produced
the underlying graph.
"""
from __future__ import annotations

from dataclasses import asdict, dataclass, field

from src.cpg.defuse import backward_slice, build_defuse_graph
from src.cpg.statement_mapper import map_patch_lines_to_statements_detailed, statement_lines
from src.cpg.trace_miner import mine_trace


@dataclass
class SampleAnnotation:
    sample_id: str
    l1_lines: list[int]
    l2_lines: list[int]
    l3_lines: list[int]  # whole function (degenerate upper bound, spec §4.2)
    trace_gt: list[list[str]]  # stringified node ids (JSON-safe), spec §4.3
    trace_unverifiable: bool
    trace_reason: str
    annotation_backend: str  # "tree_sitter" | "joern"
    trace_lengths: list[int] = field(default_factory=list)  # for the trace-length distribution report
    used_fallback_sink_search: bool = False  # True if L1 was empty and the old dangerous-call search ran
    n_dangerous_call_sinks: int = 0  # enrichment count, spec amendment §1.3-a
    cwe: str | None = None  # carried through for the per-CWE trace-coverage table
    l1_mapping_reasons: list[str] = field(default_factory=list)  # spec amendment §2: per-line reason log


def annotate_sample(
    sample_id: str,
    func_vuln: str,
    patch_lines_vuln: list[int],
    cwe: str | None,
    func_len_lines: int | None = None,
) -> SampleAnnotation:
    """spec §4.2-§4.3 for a single sample. Only meaningful for target=1 rows
    (patch_lines_vuln is empty for target=0 benign rows by construction —
    spec §3.2 — so L1/L2/trace are empty/unverifiable for those, correctly)."""
    if not patch_lines_vuln:
        n_lines = func_len_lines or len(func_vuln.splitlines())
        return SampleAnnotation(
            sample_id=sample_id, l1_lines=[], l2_lines=[], l3_lines=list(range(1, n_lines + 1)),
            trace_gt=[], trace_unverifiable=True, trace_reason="target_0_no_patch_lines",
            annotation_backend="tree_sitter", cwe=cwe,
        )

    detailed = map_patch_lines_to_statements_detailed(func_vuln, patch_lines_vuln)
    l1_mapping_reasons = [reason for _stmt, reason in detailed]
    l1_statements = [stmt for stmt, _reason in detailed if stmt is not None]
    l1_lines = sorted(statement_lines(l1_statements))

    graph = build_defuse_graph(func_vuln)
    l1_ids = {s.node_id for s in l1_statements}
    l2_ids = backward_slice(graph, l1_ids, max_depth=3)
    l2_statement_objs = [s for s in graph.statements if s.node_id in l2_ids]
    l2_lines = sorted(statement_lines(l2_statement_objs))

    n_lines = func_len_lines or len(func_vuln.splitlines())
    l3_lines = list(range(1, n_lines + 1))

    # Trace amendment: sinks = L1 statements themselves, not a dangerous-call
    # table lookup (see DECISIONS.md "Phase 2 trace-mining redefinition").
    trace_result = mine_trace(func_vuln, cwe, l1_stmt_ids=sorted(l1_ids), graph=graph)
    # stmt ids are (start_byte, end_byte) tuples — stringify for JSON safety
    trace_gt = [[f"{a}-{b}" for (a, b) in path] for path in trace_result.trace_gt]

    return SampleAnnotation(
        sample_id=sample_id,
        l1_lines=l1_lines,
        l2_lines=sorted(set(l1_lines) | set(l2_lines)),
        l3_lines=l3_lines,
        trace_gt=trace_gt,
        trace_unverifiable=trace_result.trace_unverifiable,
        trace_reason=trace_result.reason,
        annotation_backend="tree_sitter",
        trace_lengths=[len(p) for p in trace_result.trace_gt],
        used_fallback_sink_search=trace_result.used_fallback_sink_search,
        n_dangerous_call_sinks=len(trace_result.dangerous_call_sinks),
        cwe=cwe,
        l1_mapping_reasons=l1_mapping_reasons,
    )


def annotation_to_dict(ann: SampleAnnotation) -> dict:
    return asdict(ann)

	"""Phase 2 per-sample annotation (spec §4.2-§4.4): ties statement_mapper,
	defuse, and trace_miner together to produce, for one corpus row, the L1/L2/L3
	localization targets and trace ground truth.

	Uses the tree-sitter fallback (defuse.py) by default — not real Joern — since
	running Joern per-sample across the full ~13.6k-pair corpus is a real,
	multi-hour, CPU-parallelized batch job (spec §4: "budget accordingly...
	target 16 workers... resumable") that hasn't been executed yet (see
	DECISIONS.md). `annotate_sample(..., use_joern=True)` is wired for when that
	batch job runs; this module's job is to define what an annotation record
	looks like and compute it correctly, independent of which backend produced
	the underlying graph.
	"""
	from __future__ import annotations

	from dataclasses import asdict, dataclass, field

	from src.cpg.defuse import backward_slice, build_defuse_graph
	from src.cpg.statement_mapper import map_patch_lines_to_statements_detailed, statement_lines
	from src.cpg.trace_miner import mine_trace


	@dataclass
	class SampleAnnotation:
	sample_id: str
	l1_lines: list[int]
	l2_lines: list[int]
	l3_lines: list[int] # whole function (degenerate upper bound, spec §4.2)
	trace_gt: list[list[str]] # stringified node ids (JSON-safe), spec §4.3
	trace_unverifiable: bool
	trace_reason: str
	annotation_backend: str # "tree_sitter" \| "joern"
	trace_lengths: list[int] = field(default_factory=list) # for the trace-length distribution report
	used_fallback_sink_search: bool = False # True if L1 was empty and the old dangerous-call search ran
	n_dangerous_call_sinks: int = 0 # enrichment count, spec amendment §1.3-a
	cwe: str \| None = None # carried through for the per-CWE trace-coverage table
	l1_mapping_reasons: list[str] = field(default_factory=list) # spec amendment §2: per-line reason log


	def annotate_sample(
	sample_id: str,
	func_vuln: str,
	patch_lines_vuln: list[int],
	cwe: str \| None,
	func_len_lines: int \| None = None,
	) -> SampleAnnotation:
	"""spec §4.2-§4.3 for a single sample. Only meaningful for target=1 rows
	(patch_lines_vuln is empty for target=0 benign rows by construction —
	spec §3.2 — so L1/L2/trace are empty/unverifiable for those, correctly)."""
	if not patch_lines_vuln:
	n_lines = func_len_lines or len(func_vuln.splitlines())
	return SampleAnnotation(
	sample_id=sample_id, l1_lines=[], l2_lines=[], l3_lines=list(range(1, n_lines + 1)),
	trace_gt=[], trace_unverifiable=True, trace_reason="target_0_no_patch_lines",
	annotation_backend="tree_sitter", cwe=cwe,
	)

	detailed = map_patch_lines_to_statements_detailed(func_vuln, patch_lines_vuln)
	l1_mapping_reasons = [reason for _stmt, reason in detailed]
	l1_statements = [stmt for stmt, _reason in detailed if stmt is not None]
	l1_lines = sorted(statement_lines(l1_statements))

	graph = build_defuse_graph(func_vuln)
	l1_ids = {s.node_id for s in l1_statements}
	l2_ids = backward_slice(graph, l1_ids, max_depth=3)
	l2_statement_objs = [s for s in graph.statements if s.node_id in l2_ids]
	l2_lines = sorted(statement_lines(l2_statement_objs))

	n_lines = func_len_lines or len(func_vuln.splitlines())
	l3_lines = list(range(1, n_lines + 1))

	# Trace amendment: sinks = L1 statements themselves, not a dangerous-call
	# table lookup (see DECISIONS.md "Phase 2 trace-mining redefinition").
	trace_result = mine_trace(func_vuln, cwe, l1_stmt_ids=sorted(l1_ids), graph=graph)
	# stmt ids are (start_byte, end_byte) tuples — stringify for JSON safety
	trace_gt = [[f"{a}-{b}" for (a, b) in path] for path in trace_result.trace_gt]

	return SampleAnnotation(
	sample_id=sample_id,
	l1_lines=l1_lines,
	l2_lines=sorted(set(l1_lines) \| set(l2_lines)),
	l3_lines=l3_lines,
	trace_gt=trace_gt,
	trace_unverifiable=trace_result.trace_unverifiable,
	trace_reason=trace_result.reason,
	annotation_backend="tree_sitter",
	trace_lengths=[len(p) for p in trace_result.trace_gt],
	used_fallback_sink_search=trace_result.used_fallback_sink_search,
	n_dangerous_call_sinks=len(trace_result.dangerous_call_sinks),
	cwe=cwe,
	l1_mapping_reasons=l1_mapping_reasons,
	)


	def annotation_to_dict(ann: SampleAnnotation) -> dict:
	return asdict(ann)