R2Code: Why Your LLM Knows What Code to Write But Not Which Requirement It Satisfies

# R2Code framework (reconstructed from arXiv:2604.22432 architecture description)
# The key innovation is generating code and trace links simultaneously,
# then using the LLM to verify its own trace links

import json
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class Requirement:
    id: str         # "REQ-001"
    text: str       # full requirement text
    priority: str   # "shall" (mandatory) | "should" (preferred) | "may" (optional)

@dataclass
class TraceLink:
    requirement_id: str    # "REQ-001"
    code_unit: str         # "AuthenticationService.authenticate"
    confidence: float      # 0.0-1.0, from reflection verification
    verified: bool = False

@dataclass
class CodeUnit:
    name: str
    code: str
    trace_links: list[str]  # list of requirement IDs this unit satisfies

# Phase 1: Parse requirements into structured representation
def parse_requirements(spec_text: str, llm) -> list[Requirement]:
    """
    LLM extracts structured requirements from natural language spec.
    Assigns IDs, identifies priority (shall/should/may), extracts entities.
    ← Why structured output? Downstream coverage checking needs to enumerate
      requirements as discrete items. Free-text doesn't support this.
    """
    prompt = f"""
    Parse the following requirements specification into structured requirements.
    For each requirement, extract:
    - id: unique identifier (REQ-001, REQ-002, etc.)
    - text: verbatim requirement text
    - priority: "shall" | "should" | "may"

    Requirements specification:
    {spec_text}

    Return JSON: [{{"id": str, "text": str, "priority": str}}]
    """
    response = llm.generate(prompt, format="json")
    return [Requirement(**r) for r in json.loads(response)]

# Phase 2: Generate code with embedded trace annotations
def generate_traced_code(requirements: list[Requirement], llm) -> list[CodeUnit]:
    """
    Generate implementation code where each code unit is annotated
    with the requirements it satisfies.

    ← THIS is the key design decision: trace links are NOT inferred after
      code generation. They are generated SIMULTANEOUSLY with the code.
      This means the model is explicitly asked to reason about traceability
      during code generation, not as an afterthought.

    The @satisfies annotation format is chosen for:
    - Parsability: easily extracted by regex or AST parser
    - Visibility: reviewers can see trace links directly in code
    - Maintainability: trace links co-locate with the code they describe
    """
    req_list = "\n".join([f"{r.id}: {r.text}" for r in requirements])

    prompt = f"""
    Implement all of the following requirements in Python.
    For each class and method, add a @satisfies annotation listing which
    requirements it implements. A single method may satisfy multiple requirements.
    A requirement may be satisfied by multiple methods.

    Requirements:
    {req_list}

    Format for annotations:
    # @satisfies: REQ-001, REQ-003  (for class-level)
    # @satisfies: REQ-001           (for method-level)

    Generate complete, runnable Python code with all annotations.
    """
    raw_code = llm.generate(prompt)
    return parse_code_units(raw_code)  # extract CodeUnits from generated code

# Phase 3: Self-reflection verification loop
def reflect_and_verify(
    requirements: list[Requirement],
    code_units: list[CodeUnit],
    llm,
    max_iterations: int = 3,
) -> tuple[list[CodeUnit], list[TraceLink]]:
    """
    The self-reflection loop: verify trace links and regenerate as needed.

    Three verification passes per iteration:
    1. Coverage: every requirement has at least one satisfying code unit
    2. Spurious links: every asserted trace link is actually correct
    3. Missing links: code units that should satisfy a requirement but don't

    ← Why use the LLM for verification? Static analysis can check syntax.
      Only the LLM can check semantic correctness: "does this code actually
      implement this requirement?"

    ← Why iterate? Initial generation may produce incomplete trace annotations.
      The reflection loop is the error-correction mechanism.
    """
    trace_links = extract_trace_links(code_units)
    req_map = {r.id: r for r in requirements}

    for iteration in range(max_iterations):
        issues_found = []

        # Check 1: Coverage (every requirement has implementing code)
        covered_reqs = {link.requirement_id for link in trace_links}
        missing_coverage = [r for r in requirements if r.id not in covered_reqs]

        for req in missing_coverage:
            issues_found.append({
                "type": "missing_coverage",
                "requirement": req,
                "message": f"No code unit satisfies {req.id}: {req.text}"
            })

        # Check 2: Spurious link verification
        for link in trace_links:
            req = req_map.get(link.requirement_id)
            code_unit = get_code_unit(code_units, link.code_unit)
            if req and code_unit:
                # ← LLM verifies semantic alignment between requirement and code
                # This is the core of self-reflection: the model critiques its own output
                verification_prompt = f"""
                Does this code unit satisfy the stated requirement?

                Requirement {req.id}: {req.text}

                Code unit ({code_unit.name}):
                {code_unit.code}

                Answer: YES or NO, then explain in one sentence.
                """
                verdict = llm.generate(verification_prompt)
                if "NO" in verdict.upper():
                    issues_found.append({
                        "type": "spurious_link",
                        "link": link,
                        "reason": verdict,
                    })
                    link.confidence = 0.0
                else:
                    link.confidence = 1.0
                    link.verified = True

        if not issues_found:
            break  # ← Early termination: all links verified

        # Selective regeneration: only fix broken code units
        # ← THIS is the trick: don't regenerate everything, only what's broken
        # This preserves correct code and reduces token cost per iteration
        code_units = selective_regenerate(
            code_units, issues_found, requirements, llm
        )
        trace_links = extract_trace_links(code_units)

    return code_units, trace_links

# Traceability matrix output and evaluation metrics
# This is what R2Code produces and how it is measured

import json
from collections import defaultdict

def build_traceability_matrix(
    requirements: list[Requirement],
    trace_links: list[TraceLink],
    code_units: list[CodeUnit],
) -> dict:
    """
    Build bidirectional traceability matrix from verified trace links.

    Bidirectional traceability:
    - Forward: requirement → {code units that implement it}
    - Backward: code unit → {requirements it implements}

    ← Why bidirectional? Both directions are needed in practice:
      Forward: when requirement changes, find all code to update
      Backward: when reviewing code, verify it has a traced requirement (no orphan code)
    """
    forward_trace = defaultdict(list)  # req_id → [code_units]
    backward_trace = defaultdict(list) # code_unit → [req_ids]

    for link in trace_links:
        if link.verified or link.confidence > 0.5:
            forward_trace[link.requirement_id].append(link.code_unit)
            backward_trace[link.code_unit].append(link.requirement_id)

    # Identify coverage gaps
    gaps = [r for r in requirements if r.id not in forward_trace]
    orphans = [cu for cu in code_units if cu.name not in backward_trace]

    return {
        "forward": dict(forward_trace),
        "backward": dict(backward_trace),
        "coverage": {
            "total_requirements": len(requirements),
            "covered": len(forward_trace),
            "gaps": [(r.id, r.text[:60]) for r in gaps],  # uncovered requirements
            "orphans": [cu.name for cu in orphans],        # code with no requirement
        }
    }

def evaluate_traceability(
    predicted_matrix: dict,
    ground_truth_matrix: dict,
    metric: str = "f1",
) -> dict:
    """
    Standard evaluation metrics for requirements traceability.

    Precision: of all predicted trace links, what fraction are correct?
    Recall: of all true trace links, what fraction were predicted?
    F1: harmonic mean of precision and recall
    F2: recall-weighted F1 (used in safety-critical contexts where missing
        a trace link is more costly than a spurious one)

    ← F2 is often preferred in traceability because missing a link
      (failing to trace code to a requirement) is more dangerous than
      a false link (incorrectly claiming code satisfies a requirement).
      In safety-critical systems, the former can mean unverified compliance.
    """
    predicted_links = set()
    for req_id, code_units in predicted_matrix["forward"].items():
        for cu in code_units:
            predicted_links.add((req_id, cu))

    ground_truth_links = set()
    for req_id, code_units in ground_truth_matrix["forward"].items():
        for cu in code_units:
            ground_truth_links.add((req_id, cu))

    true_positives = len(predicted_links & ground_truth_links)
    false_positives = len(predicted_links - ground_truth_links)
    false_negatives = len(ground_truth_links - predicted_links)

    precision = true_positives / (true_positives + false_positives + 1e-9)
    recall = true_positives / (true_positives + false_negatives + 1e-9)
    f1 = 2 * precision * recall / (precision + recall + 1e-9)
    # ← F2: weights recall twice as much as precision
    f2 = (1 + 4) * precision * recall / (4 * precision + recall + 1e-9)

    return {"precision": precision, "recall": recall, "f1": f1, "f2": f2}

# R2Code benchmark results (from arXiv:2604.22432):
# Dataset         Precision    Recall    F1      F2
# iTrust          0.84         0.79      0.81    0.80
# EasyClinic      0.81         0.82      0.81    0.81
# CoCoTrace       0.76         0.75      0.75    0.75
#
# Prior SOTA (TraceLLM, arXiv:2602.01253, on same datasets):
# Achieved strong F2 scores on aerospace/healthcare datasets
# but was evaluated on different artifact types (requirements → design)

REQ-001: The system shall authenticate users via valid email and password combination.
REQ-002: Authentication sessions shall expire after 15 minutes of inactivity.
REQ-003: Failed login attempts shall be logged with timestamp, user email, and IP address.
REQ-004: After 5 consecutive failed login attempts, the account shall be locked.
REQ-005: Password reset shall require verification via email link valid for 60 minutes.

[
  {"id": "REQ-001", "text": "...", "priority": "shall", "entities": ["user", "email", "password"]},
  {"id": "REQ-002", "text": "...", "priority": "shall", "entities": ["session", "timeout", "15min"]},
  {"id": "REQ-003", "text": "...", "priority": "shall", "entities": ["log", "timestamp", "email", "IP"]},
  {"id": "REQ-004", "text": "...", "priority": "shall", "entities": ["lockout", "5 attempts"]},
  {"id": "REQ-005", "text": "...", "priority": "shall", "entities": ["password_reset", "email", "60min"]}
]
Dependency graph: REQ-003 depends on REQ-001; REQ-004 depends on REQ-003

# @satisfies: REQ-001, REQ-004
class AuthenticationService:
    # @satisfies: REQ-001
    def authenticate(self, email: str, password: str) -> AuthResult:
        ...

    # @satisfies: REQ-002
    def check_session_timeout(self, session_id: str) -> bool:
        ...

    # @satisfies: REQ-003
    def log_failed_attempt(self, email: str, ip_address: str) -> None:
        ...

    # @satisfies: REQ-004
    def enforce_lockout_policy(self, email: str) -> None:
        ...

    # @satisfies: REQ-005
    def initiate_password_reset(self, email: str) -> str:
        ...

REQ-001: [authenticate]
REQ-002: [check_session_timeout]
REQ-003: [log_failed_attempt]
REQ-004: [enforce_lockout_policy]  ← also listed at class level but not REQ-001 class annotation
REQ-005: [initiate_password_reset]

Coverage check:
  All 5 requirements covered ✓

Spurious link check:
  (REQ-001, AuthenticationService class) → verify: "Does the class implement auth?"
  LLM verdict: YES ✓
  (REQ-004, enforce_lockout_policy) → verify: "Does this implement 5-attempt lockout?"
  LLM inspects code: missing the count-and-compare logic
  Verdict: NO — code calls lockout but doesn't count attempts
  → Issue: spurious link at method level

  (REQ-004, authenticate) → not annotated but LLM detects: authenticate calls
  enforce_lockout_policy which is REQ-004, but authenticate itself should
  also be annotated as satisfying REQ-004 (it's where the attempt is counted)
  → Issue: missing link

Selective regeneration:
  Regenerate: enforce_lockout_policy (add attempt counting logic)
  Add missing annotation: authenticate @satisfies: REQ-001, REQ-004

Coverage check: ✓ (all 5 requirements)
Spurious link check: all links verified ✓
No further issues found.
Terminate at iteration 2.

Final trace matrix:
  REQ-001: [authenticate, AuthenticationService]
  REQ-002: [check_session_timeout]
  REQ-003: [log_failed_attempt]
  REQ-004: [enforce_lockout_policy, authenticate]
  REQ-005: [initiate_password_reset]

Traceability coverage:
  5/5 requirements covered (100%)
  0 orphan code units
  0 uncovered requirements

Token usage:
  Phase 1 (parsing):         ~600 tokens
  Phase 2 (generation):     ~3,200 tokens
  Phase 3 iter 1:           ~2,100 tokens (verification of 5 links + 2 regenerations)
  Phase 3 iter 2:           ~1,400 tokens (verification only)
  Total:                    ~7,300 tokens

vs. baseline (generate code, then separately recover trace links):
  Code generation:           ~3,200 tokens
  Trace recovery (separate): ~4,500 tokens (for all pairwise comparisons)
  Total:                     ~7,700 tokens
  ← R2Code is slightly cheaper AND produces better trace links
     because simultaneous generation avoids the pairwise comparison explosion