Set up initial project structure

Also add bare-bones test which can be run like so:

```
    uvx pytest tests/test_postprocess.py
```

as long as `uvx` is in path.

Author: thedataking

.gitignore

@@ -10,6 +10,8 @@
 **/tests/
 /build
 *.pyc
+**/__pycache__
+*.egg-info/
 .vagrant
 **/compile_commands.json
 .python-version

c2rust-postprocess/README.md

@@ -0,0 +1,43 @@
+# LLM-based postprocessing of c2rust transpiler output
+
+This is currently a prototype effort to gauge the extent to which LLMs can 
+accelerate the types of translation and migration that help move C code to Rust.
+
+# Prerequisites
+
+- Python 3.12 or later
+- `uv` in path
+- A valid `GEMINI_API_KEY` set
+- A transpiled codebase with a correct `compile_commands.json`
+
+# Running
+
+- `c2rust-postprocess`, or
+- `uv run postprocess`
+
+# Testing
+
+## Test prerequisites
+
+- `bear` and `c2rust` in path 
+
+```
+   uv run pytest -v
+   uv run pytest -v tests/test_utils.py # filter tests to run
+```
+
+## Misc
+
+- `uv run ruff check --fix .` to format & lint
+
+# TODOs
+
+- testable prototype
+  - pluggable support for getting definitions
+  - gemini api support
+  - filtering by file and function name
+  - file-based caching of model responses
+  - openai model support
+  - antropic model support
+  - openrouter API support?
+

c2rust-postprocess/c2rust-postprocess

@@ -0,0 +1,2 @@
+#!/bin/sh
+uv run postproc

c2rust-postprocess/postprocess/__init__.py

@@ -0,0 +1,66 @@
+"""
+c2rust-postprocess: Transfer comments from C functions to Rust functions using LLMs.
+"""
+
+
+from pathlib import Path
+from typing import Any
+from postprocess.definitions import get_function_span_pairs
+from postprocess.utils import read_chunk, get_rust_files, get_compile_commands
+
+# TODO: could also include
+# - validation function to check result
+# - list of comments to check for
+class CommentTransferPrompt:
+    c_function: str
+    rust_function: str
+    prompt_text: str
+
+    __slots__ = ("c_function", "rust_function", "prompt_text")
+
+    def __init__(self, c_function: str, rust_function: str, prompt_text: str) -> None:
+        self.c_function = c_function
+        self.rust_function = rust_function
+        self.prompt_text = prompt_text
+
+    def __str__(self) -> str:
+        raise NotImplementedError("String conversion not implemented yet.")
+        
+    
+
+def generate_prompts(compile_commands: dict[str, Any], rust_file: Path) -> list[CommentTransferPrompt]:
+    pairs = get_function_span_pairs(compile_commands, rust_file)
+
+    prompts = []
+
+    for rust_fn, c_fn in pairs:
+        c_def = read_chunk(c_fn['file'], c_fn['start_byte'], c_fn['end_byte'])
+        # TODO: log on verbose level
+        # print(f"C function {c_fn['name']} definition:\n{c_def}\n")
+
+        rust_def = read_chunk(rust_fn['file'], rust_fn['start_byte'], rust_fn['end_byte'])
+        # TODO: log on verbose level
+        # print(f"Rust function {rust_fn['name']} definition:\n{rust_def}\n")
+        
+        # TODO: make this function take a model and get prompt from model
+        prompt_text = f"""Transfer the comments from the following C function to the corresponding Rust function."""
+        prompt = CommentTransferPrompt(
+            c_function=c_def,
+            rust_function=rust_def,
+            prompt_text=prompt_text
+        )
+
+        prompts.append(prompt)
+
+    return prompts
+
+
+def transfer_comments(compile_commands_path: Path):
+    rust_sources = get_rust_files(compile_commands_path.parent)
+    
+    compile_commands  = get_compile_commands(compile_commands_path)
+    for rust_file in rust_sources:
+        prompts = generate_prompts(compile_commands, rust_file)
+        # Call LLM API with the prompt and get the response
+        # Process the response to extract comments and associate them with Rust functions
+        pass

c2rust-postprocess/postprocess/__main__.py

@@ -0,0 +1,46 @@
+import argparse
+import logging
+import sys
+from collections.abc import Sequence
+
+from postprocess.utils import existing_file
+
+
+def build_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Transfer C function comments to Rust using LLMs.",
+    )
+    parser.add_argument(
+        "compile_commands",
+        type=existing_file,
+        help="Path to compile_commands.json.",
+    )
+
+    parser.add_argument(
+        "--log-level",
+        type=str,
+        required=False,
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        help="Logging level (default: INFO)",
+    )
+
+    return parser
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    parser = build_arg_parser()
+    args = parser.parse_args(argv)
+
+    logging.basicConfig(level=logging.getLevelName(args.log_level.upper()))
+
+    _compile_commands = args.compile_commands
+    # TODO: implement post-processing logic using compile_commands
+
+
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

c2rust-postprocess/postprocess/definitions/__init__.py

@@ -0,0 +1,105 @@
+import logging
+from pathlib import Path
+from typing import Any
+
+import tree_sitter_rust as tsrust
+from tree_sitter import Language, Parser
+
+
+def get_c_sourcefile(compile_commands, rustfile: Path) -> Path | None:
+    c_file_guesses = [rustfile.with_suffix(".c"), rustfile.with_suffix(".C")]
+
+    files = [Path(d["file"]) for d in compile_commands]
+
+    for guess in c_file_guesses:
+        if guess in files:
+            return guess
+
+    return None
+
+
+def get_rust_function_spans(rustfile: Path) -> list[dict[str, Any]]:
+    LANGUAGE = Language(tsrust.language())
+    parser = Parser(LANGUAGE)
+
+    if not rustfile.exists():
+        raise FileNotFoundError(f"{rustfile} does not exist")
+    if not rustfile.is_file():
+        raise NotADirectoryError(f"{rustfile} is not a file")
+
+    try:
+        with open(rustfile, "rb") as rust_source:
+            source_bytes = rust_source.read()
+    except OSError as exc:
+        logging.error(f"Failed to read Rust file {rustfile}: {exc}")
+        return []
+
+    tree = parser.parse(source_bytes)
+
+    functions = []
+
+    for node in tree.root_node.children:
+        if node.type == 'function_item':
+            name_node = node.child_by_field_name('name')
+            func_name = (source_bytes[
+                name_node.start_byte: # type: ignore
+                name_node.end_byte    # type: ignore
+            ].decode('utf-8'))
+
+            functions.append({
+                "name": func_name,
+                "start_line": node.start_point[0] + 1, # 0-indexed
+                "end_line": node.end_point[0] + 1, # 0-indexed
+                "start_byte": node.start_byte,
+                "end_byte": node.end_byte
+            })
+
+    return functions
+
+
+def get_c_functions_spans(compile_commands: dict[str, Any], c_file: Path):
+    from .clang import get_c_ast_as_json, get_functions_from_clang_ast
+    cmd = (c for c in compile_commands if c["file"] == str(c_file))
+    entry = next(cmd, None)
+
+    assert entry is not None, f"No compile command entry for {c_file}"
+
+    c_fn_asts = get_functions_from_clang_ast(get_c_ast_as_json(entry))
+
+    # print(json.dumps(c_fn_asts, indent=4))
+
+    functions = []
+    for fn in c_fn_asts:
+        loc = fn["loc"]
+        if "line" in loc and "col" in loc:
+            functions.append({
+                "name": fn["name"],
+                "start_line": loc["line"],
+                "start_byte": fn["range"]["begin"]["offset"],
+                "end_line": fn["range"]["end"]["line"],
+                "end_byte": fn["range"]["end"]["offset"],
+        })
+
+    return functions
+
+
+def get_function_span_pairs(compile_commands: dict[str, Any], rustfile: Path) -> list[tuple[dict[str, Any], dict[str, Any]]]:
+    """Get pairs of Rust and C function spans for the given Rust file."""
+
+    rust_fn_spans = get_rust_function_spans(rustfile)   
+    c_file = get_c_sourcefile(compile_commands, rustfile)
+    if not c_file:
+        raise FileNotFoundError(f"No corresponding C source file found for {rustfile}")
+
+    c_fn_spans = get_c_functions_spans(compile_commands, c_file)
+
+    # TODO: handle cases where ordering or counts differ
+    # A reasonable assumption is that we can still pair functions by name
+    # which means that this tool needs to run fairly soon after transpilation
+    assert len(c_fn_spans) == len(rust_fn_spans), "Mismatched number of functions between Rust and C source files"
+    for rust_fn, c_fn in zip(rust_fn_spans, c_fn_spans):
+        assert rust_fn['name'] == c_fn['name']
+        rust_fn['file'] = rustfile
+        c_fn['file'] = c_file
+
+    return list(zip(rust_fn_spans, c_fn_spans))

c2rust-postprocess/postprocess/definitions/clang/__init__.py

@@ -0,0 +1,68 @@
+
+import json
+import subprocess
+from typing import Any
+
+import jq
+
+
+def get_functions_from_clang_ast(ast: dict[str, Any]) -> list[dict[str, Any]]:
+    """
+    Extract function declarations from the Clang AST JSON.
+    Args:
+        ast (dict): The AST JSON as a dictionary.
+    Returns:
+        list[dict]: A list of dictionaries, each representing a function declaration.
+    """
+    query = jq.compile(
+        '.inner[] | select(.kind =="FunctionDecl") | {name: .name, loc: .loc, range: .range}'
+    )
+    return query.transform(ast, multiple_output=True)
+
+
+def get_c_ast_as_json(entry: dict[str, Any]) -> dict[str, Any]:
+    """
+    Get AST as JSON for a translation unit identified by compile commands entry.
+    """
+    source_file = entry["file"]
+
+    cmd = entry["arguments"]
+    cmd[0] = "clang"  # make sure we use clang
+    # drop the last four elements which are the output options
+    cmd = cmd[:-4]  # TODO: validate that these are the output options
+    # add the necessary flags to dump the AST as JSON
+    cmd += [
+        "-fsyntax-only",
+        "-Xclang",
+        "-ast-dump=json",
+        "-fparse-all-comments",  # NOTE: Clang AST only includes doc comments
+        source_file,
+    ]
+    try:
+        # cwd to the directory from the compile_commands.json entry to make sure
+        # relative paths in the command work correctly
+        result = subprocess.run(
+            cmd, capture_output=True, text=True, check=True, cwd=entry["directory"]
+        )
+        return json.loads(result.stdout)
+    except subprocess.CalledProcessError as e:
+        print(f"Error running clang on {source_file}: {e.stderr}")
+        raise
+
+
+def is_entry_from_c_file(entry: dict[str, Any], c_file: str) -> bool:
+    """
+    Check if the entry is from the specified C file.
+    """
+    loc = entry["loc"]
+    if "file" in loc:
+        return loc["file"] == c_file
+    elif "spellingLoc" in loc and "includedFrom" in loc["spellingLoc"]:
+        return loc["spellingLoc"]["includedFrom"]["file"] == c_file
+    elif "expansionLoc" in loc and "includedFrom" in loc["expansionLoc"]:
+        return loc["expansionLoc"]["includedFrom"]["file"] == c_file
+    if "includedFrom" not in loc:
+        return (
+            True  # entry was parsed from c_file so by default it is from that file
+        )
+    return False

c2rust-postprocess/postprocess/models/__init__.py

@@ -0,0 +1,48 @@
+from abc import ABC, abstractmethod
+from typing import List, Dict, Callable, Optional, Any
+
+
+class AbstractGenAIModel(ABC):
+    """
+    Abstract base class for LLM clients using Native Function Calling.
+    """
+
+    def __init__(self, model_id: str, **kwargs: Any):
+        self._model_id = model_id
+        self._config = kwargs
+
+    @property
+    def model_id(self) -> str:
+        return self._model_id
+
+    @abstractmethod
+    async def agenerate_with_tools(
+        self, 
+        messages: List[Dict[str, Any]], 
+        tools: Optional[List[Callable]] = None,
+        max_tool_loops: int = 5
+    ) -> Any:
+        """
+        Generate a response using native automatic function calling.
+        
+        Args:
+            messages: Chat history.
+            tools: List of Python functions (the concrete class handles schema conversion).
+            max_tool_loops: Maximum number of times the model can call tools 
+                            consecutively (prevent infinite loops).
+        
+        Returns:
+            The final natural language response from the model.
+        """
+        pass
+
+    def generate_with_tools(
+        self, 
+        messages: List[Dict[str, Any]], 
+        tools: Optional[List[Callable]] = None,
+        max_tool_loops: int = 5
+    ) -> str:
+        """Synchronous wrapper for agenerate_response."""
+        import asyncio
+        return asyncio.run(self.agenerate_with_tools(messages, tools, max_tool_loops))
+