postprocess: include C definitions before and after preprocessing in prompts

Read both forms from the structured *.c_decls.json instead of running a clang subprocess per snippet. The CommentsTransform prompt includes the preprocessed text only when it differs from the original, so prompts (and llm-cache keys) are unchanged for directive-free functions.

Comment gating and response validation use comments from the preprocessed text plus comments on original preprocessor directive lines, so inactive-region comments are not transferred while directive-line comments are not dropped just because clang omits them from directives-only output.

Author: thedataking

c2rust-postprocess/postprocess/definitions/__init__.py

@@ -2,6 +2,7 @@
 import logging
 import subprocess
 from collections.abc import Generator, Iterable
+from dataclasses import dataclass
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 from typing import Any
@@ -233,12 +234,37 @@ def get_rust_definitions(root_rust_source_file: Path) -> dict[str, str]:
     return json.loads(result.stdout)
 
 
-def get_c_definitions(root_rust_source_file: Path) -> dict[str, str]:
+@dataclass(frozen=True)
+class CDefinition:
+    definition: str
+    preprocessed_definition: str | None
+
+    @property
+    def effective(self) -> str:
+        """Preprocessed text when available, original otherwise."""
+        return self.preprocessed_definition or self.definition
+
+    @property
+    def was_changed_by_preprocessing(self) -> bool:
+        return (
+            self.preprocessed_definition is not None
+            and self.preprocessed_definition != self.definition
+        )
+
+
+def get_c_definitions(root_rust_source_file: Path) -> dict[str, CDefinition]:
     c_defs_json = root_rust_source_file.with_suffix(".c_decls.json")
 
     logging.debug(f"Loading C definitions from {c_defs_json}")
     try:
-        return json.loads(c_defs_json.read_text())
+        c_decls = json.loads(c_defs_json.read_text())
+        return {
+            identifier: CDefinition(
+                definition=decl["definition"],
+                preprocessed_definition=decl.get("preprocessed_definition"),
+            )
+            for identifier, decl in c_decls["definitions"].items()
+        }
     except OSError as e:
         e.add_note(f"C definitions JSON file not found: {c_defs_json}")
         raise e

c2rust-postprocess/postprocess/transforms/base.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 
 from postprocess.definitions import (
+    CDefinition,
     get_c_definitions,
     get_rust_definitions,
 )
@@ -32,7 +33,7 @@ def apply_ident(
         self,
         rust_source_file: Path,
         rust_definition: str,
-        c_definition: str,
+        c_definition: CDefinition,
         identifier: str,
         update_rust: bool = True,
     ) -> None:
@@ -114,7 +115,7 @@ def apply_file(
 
             c_definition = c_definitions[identifier]
 
-            highlighted_c_definition = get_highlighted_c(c_definition)
+            highlighted_c_definition = get_highlighted_c(c_definition.effective)
             logging.debug(
                 f"C function {identifier} definition:\n{highlighted_c_definition}\n"
             )

c2rust-postprocess/postprocess/transforms/comments.py

@@ -1,9 +1,11 @@
 import logging
+from collections import Counter
 from pathlib import Path
 from textwrap import dedent
 
 from postprocess.cache import AbstractCache
 from postprocess.definitions import (
+    CDefinition,
     get_c_comments,
     get_rust_comments,
     update_rust_definition,
@@ -20,31 +22,74 @@
 
 class CommentsTransformPrompt:
     c_function: str
+    c_function_preprocessed: str | None
     rust_function: str
     prompt_text: str
     identifier: str
 
-    __slots__ = ("c_function", "rust_function", "prompt_text", "identifier")
+    __slots__ = (
+        "c_function",
+        "c_function_preprocessed",
+        "rust_function",
+        "prompt_text",
+        "identifier",
+    )
 
     def __init__(
-        self, c_function: str, rust_function: str, prompt_text: str, identifier: str
+        self,
+        c_function: str,
+        c_function_preprocessed: str | None,
+        rust_function: str,
+        prompt_text: str,
+        identifier: str,
     ):
         self.c_function = c_function
+        self.c_function_preprocessed = c_function_preprocessed
         self.rust_function = rust_function
         self.prompt_text = prompt_text
         self.identifier = identifier
 
     def __str__(self) -> str:
-        return (
+        prompt = (
             self.prompt_text
             + "\n\n"
             + "C function:\n```c\n"
             + self.c_function
             + "```\n\n"
-            + "Rust function:\n```rust\n"
-            + self.rust_function
-            + "```\n"
         )
+        if self.c_function_preprocessed is not None:
+            prompt += (
+                "The same C function after preprocessing; comments that are"
+                " absent here may have been in inactive preprocessor regions"
+                " and must not be transferred. Comments on preprocessor"
+                " directive lines in the original C function may be absent"
+                " here even when they annotate active code:\n```c\n"
+                + self.c_function_preprocessed
+                + "```\n\n"
+            )
+        prompt += "Rust function:\n```rust\n" + self.rust_function + "```\n"
+        return prompt
+
+
+def _get_directive_line_comments(code: str) -> list[str]:
+    return get_c_comments(
+        "\n".join(line for line in code.splitlines() if line.lstrip().startswith("#"))
+    )
+
+
+def _get_transferable_c_comments(c_definition: CDefinition) -> list[str]:
+    if c_definition.preprocessed_definition is None:
+        return get_c_comments(c_definition.definition)
+
+    wanted_comments = Counter(get_c_comments(c_definition.preprocessed_definition))
+    wanted_comments.update(_get_directive_line_comments(c_definition.definition))
+
+    comments = []
+    for comment in get_c_comments(c_definition.definition):
+        if wanted_comments[comment] > 0:
+            comments.append(comment)
+            wanted_comments[comment] -= 1
+    return comments
 
 
 class CommentsTransform(AbstractTransform):
@@ -57,7 +102,7 @@ def apply_ident(
         self,
         rust_source_file: Path,
         rust_definition: str,
-        c_definition: str,
+        c_definition: CDefinition,
         identifier: str,
         update_rust: bool = True,
     ) -> None:
@@ -70,8 +115,7 @@ def apply_ident(
             )
             return
 
-        # Skip functions without comments in C definition
-        c_comments = get_c_comments(c_definition)
+        c_comments = _get_transferable_c_comments(c_definition)
         if not c_comments:
             logging.info(f"Skipping C function without comments: {identifier}")
             return
@@ -86,7 +130,14 @@ def apply_ident(
         prompt_text = dedent(prompt_text).strip()
 
         prompt = CommentsTransformPrompt(
-            c_function=c_definition,
+            c_function=c_definition.definition,
+            # Only include the preprocessed text when it differs, so prompts
+            # (and thus cache keys) are unchanged for directive-free functions.
+            c_function_preprocessed=(
+                c_definition.preprocessed_definition
+                if c_definition.was_changed_by_preprocessing
+                else None
+            ),
             rust_function=rust_definition,
             prompt_text=prompt_text,
             identifier=identifier,
@@ -121,7 +172,6 @@ def apply_ident(
 
         rust_fn = remove_backticks(response)
 
-        c_comments = get_c_comments(prompt.c_function)
         logging.debug(f"{c_comments=}")
 
         rust_comments = get_rust_comments(rust_fn)

c2rust-postprocess/tests/examples/qsort.c_decls.json

@@ -1,5 +1 @@
-{
-    "swap": "void swap(int* a, int* b)\n{\n    int t = *a;\n    *a = *b;\n    *b = t;\n}",
-    "quickSort": "void quickSort(int arr[], int low, int high)\n{\n    if (low < high) {\n        /* pi is the partitioning index; arr[pi] is now at the right place */\n        int pi = partition(arr, low, high);\n\n        /* Recursively sort elements before and after partition */\n        quickSort(arr, low, pi - 1);\n        quickSort(arr, pi + 1, high);\n    }\n}",
-    "partition": "/* \n * Lomuto Partition Scheme:\n * Partitions the array so that elements < pivot are on the left, \n * and elements >= pivot are on the right.\n */\nint partition (int arr[], int low, int high)\n{\n    // Partition the subarray around the last element as pivot and return pivot's final index.\n    int pivot = arr[high];\n    int i = low - 1;\n\n    for (int j = low; j <= high - 1; j++) {\n        if (arr[j] <= pivot) {\n            i++;\n            // Move elements <= pivot into the left partition.\n            swap(&arr[i], &arr[j]);\n        }\n    }\n    // Place pivot just after the final element of the left partition.\n    swap(&arr[i + 1], &arr[high]);\n    return i + 1;\n}"
-}
+{"definitions":{"swap":{"definition":"void swap(int* a, int* b)\n{\n    int t = *a;\n    *a = *b;\n    *b = t;\n}","preprocessed_definition":"void swap(int* a, int* b)\n{\n    int t = *a;\n    *a = *b;\n    *b = t;\n}"},"partition":{"definition":"/* \n * Lomuto Partition Scheme:\n * Partitions the array so that elements < pivot are on the left, \n * and elements >= pivot are on the right.\n */\nint partition (int arr[], int low, int high)\n{\n    // Partition the subarray around the last element as pivot and return pivot's final index.\n    int pivot = arr[high];\n    int i = low - 1;\n\n    for (int j = low; j <= high - 1; j++) {\n        if (arr[j] <= pivot) {\n            i++;\n            // Move elements <= pivot into the left partition.\n            swap(&arr[i], &arr[j]);\n        }\n    }\n    // Place pivot just after the final element of the left partition.\n    swap(&arr[i + 1], &arr[high]);\n    return i + 1;\n}","preprocessed_definition":"/* \n * Lomuto Partition Scheme:\n * Partitions the array so that elements < pivot are on the left, \n * and elements >= pivot are on the right.\n */\nint partition (int arr[], int low, int high)\n{\n    // Partition the subarray around the last element as pivot and return pivot's final index.\n    int pivot = arr[high];\n    int i = low - 1;\n\n    for (int j = low; j <= high - 1; j++) {\n        if (arr[j] <= pivot) {\n            i++;\n            // Move elements <= pivot into the left partition.\n            swap(&arr[i], &arr[j]);\n        }\n    }\n    // Place pivot just after the final element of the left partition.\n    swap(&arr[i + 1], &arr[high]);\n    return i + 1;\n}"},"quickSort":{"definition":"void quickSort(int arr[], int low, int high)\n{\n    if (low < high) {\n        /* pi is the partitioning index; arr[pi] is now at the right place */\n        int pi = partition(arr, low, high);\n\n        /* Recursively sort elements before and after partition */\n        quickSort(arr, low, pi - 1);\n        quickSort(arr, pi + 1, high);\n    }\n}","preprocessed_definition":"void quickSort(int arr[], int low, int high)\n{\n    if (low < high) {\n        /* pi is the partitioning index; arr[pi] is now at the right place */\n        int pi = partition(arr, low, high);\n\n        /* Recursively sort elements before and after partition */\n        quickSort(arr, low, pi - 1);\n        quickSort(arr, pi + 1, high);\n    }\n}"}}}

c2rust-postprocess/tests/test_comments_transform.py

@@ -0,0 +1,82 @@
+from pathlib import Path
+from typing import Any
+
+from postprocess.cache import AbstractCache
+from postprocess.definitions import CDefinition
+from postprocess.models.mock import MockGenerativeModel
+from postprocess.transforms.comments import CommentsTransform
+
+
+class StaticCache(AbstractCache):
+    def __init__(self, response: str):
+        super().__init__(Path())
+        self.response = response
+        self.lookups = 0
+        self.messages: list[list[dict[str, Any]]] = []
+
+    def lookup(
+        self,
+        *,
+        transform: str,
+        identifier: str,
+        model: str,
+        messages: list[dict[str, Any]],
+    ) -> str | None:
+        self.lookups += 1
+        self.messages.append(messages)
+        return self.response
+
+    def update(
+        self,
+        *,
+        transform: str,
+        identifier: str,
+        model: str,
+        messages: list[dict[str, Any]],
+        response: str,
+    ) -> None:
+        raise AssertionError("cached response should not be updated")
+
+
+def test_directive_line_comment_survives_preprocessed_check() -> None:
+    c_definition = CDefinition(
+        definition="""\
+int enabled(void) {
+#ifdef FEATURE
+    return 1;
+#endif /* enabled path */
+}
+""",
+        preprocessed_definition="""\
+int enabled(void) {
+
+    return 1;
+
+}
+""",
+    )
+    rust_definition = """\
+pub unsafe extern "C" fn enabled() -> libc::c_int {
+    return 1 as libc::c_int;
+}
+"""
+    response = """\
+/// enabled path
+pub unsafe extern "C" fn enabled() -> libc::c_int {
+    return 1 as libc::c_int;
+}
+"""
+    cache = StaticCache(response)
+    transform = CommentsTransform(cache=cache, model=MockGenerativeModel())
+
+    transform.apply_ident(
+        rust_source_file=Path("unused.rs"),
+        rust_definition=rust_definition,
+        c_definition=c_definition,
+        identifier="enabled",
+        update_rust=False,
+    )
+
+    assert cache.lookups == 1
+    prompt = cache.messages[0][0]["content"]
+    assert "preprocessor directive lines" in prompt