diff --git a/c2rust-postprocess/postprocess/__init__.py b/c2rust-postprocess/postprocess/__init__.py index 98e1d915c0..27d8f96ab3 100644 --- a/c2rust-postprocess/postprocess/__init__.py +++ b/c2rust-postprocess/postprocess/__init__.py @@ -21,6 +21,7 @@ from postprocess.validate import BaselineError, make_validator DEFAULT_LLM_MODEL = "gemini-3.5-flash" +DEFAULT_TRANSFORMS = ("comments", "asserts", "formatting") def build_arg_parser() -> argparse.ArgumentParser: @@ -125,10 +126,11 @@ def build_arg_parser() -> argparse.ArgumentParser: type=str, required=False, action="append", - default=["comments"], + default=None, help=( "Transform to apply; pass multiple times to apply multiple transforms " - "in sorted order (default: comments)" + "in the order provided; duplicate transforms are ignored " + f"(default: {', '.join(DEFAULT_TRANSFORMS)})" ), ) @@ -186,13 +188,16 @@ def main(argv: Sequence[str] | None = None): model = get_model(args.llm_model) - # sort transform IDs to transforms always run in the same order to - # maximize cache hits even if the user passed them in a different order - transform_ids = sorted( - transform_id.strip() - for transform_id in set(args.transform) - if transform_id.strip() + # De-duplicate transform IDs while preserving their first occurrence. + transform_args = args.transform or DEFAULT_TRANSFORMS + transform_ids = list( + dict.fromkeys( + transform_id.strip() + for transform_id in transform_args + if transform_id.strip() + ) ) + transforms = [ get_transform_by_id( transform_id, diff --git a/c2rust-postprocess/postprocess/models/gpt.py b/c2rust-postprocess/postprocess/models/gpt.py index eb0c8904a4..d37a5f1798 100644 --- a/c2rust-postprocess/postprocess/models/gpt.py +++ b/c2rust-postprocess/postprocess/models/gpt.py @@ -1,11 +1,25 @@ +import inspect +import json from collections.abc import Callable, Iterable -from typing import Any +from typing import Any, Protocol, cast from openai import OpenAI +from openai.types.responses import ( + FunctionToolParam, + ResponseFunctionToolCall, + ResponseInputParam, +) +from openai.types.responses.response_input_param import FunctionCallOutput from postprocess.models import AbstractGenerativeModel +class NamedCallable(Protocol): + __name__: str + + def __call__(self, *args: Any, **kwargs: Any) -> Any: ... + + class GPTModel(AbstractGenerativeModel): def __init__( self, @@ -22,13 +36,100 @@ def generate_with_tools( tools: Iterable[Callable[..., Any]] = (), max_tool_loops: int = 5, ) -> str: - # TODO: implement tool calling support - assert not tools, "Tool calling not yet implemented for GPTModel" + tools = [self._named_tool(tool) for tool in tools] + tool_schemas = [self._tool_schema(tool) for tool in tools] + tool_by_name = {tool.__name__: tool for tool in tools} + + if tool_schemas: + response = self.client.responses.create( + model=self.id, + input=messages[0]["content"], + max_tool_calls=max_tool_loops, + tools=tool_schemas, + ) + else: + response = self.client.responses.create( + model=self.id, + input=messages[0]["content"], + max_tool_calls=max_tool_loops, + ) - response = self.client.responses.create( - model=self.id, - input=messages[0]["content"], - max_tool_calls=max_tool_loops, - ) + for _ in range(max_tool_loops): + tool_calls = [ + cast(ResponseFunctionToolCall, item) + for item in response.output + if getattr(item, "type", None) == "function_call" + ] + if not tool_calls: + return response.output_text + + tool_outputs: ResponseInputParam = [ + FunctionCallOutput( + type="function_call_output", + call_id=tool_call.call_id, + output=self._call_tool(tool_call, tool_by_name), + ) + for tool_call in tool_calls + ] + response = self.client.responses.create( + model=self.id, + input=tool_outputs, + previous_response_id=response.id, + max_tool_calls=max_tool_loops, + tools=tool_schemas, + ) return response.output_text + + def _named_tool(self, tool: Callable[..., Any]) -> NamedCallable: + if not hasattr(tool, "__name__"): + raise TypeError(f"Tool must be a named function: {tool!r}") + return cast(NamedCallable, tool) + + def _tool_schema(self, tool: NamedCallable) -> FunctionToolParam: + signature = inspect.signature(tool) + properties: dict[str, object] = {} + required: list[str] = [] + for name, parameter in signature.parameters.items(): + properties[name] = { + "type": self._json_schema_type(parameter.annotation), + } + if parameter.default is inspect.Parameter.empty: + required.append(name) + + return { + "type": "function", + "name": tool.__name__, + "description": inspect.getdoc(tool) or f"Call `{tool.__name__}`.", + "parameters": { + "type": "object", + "properties": properties, + "required": required, + "additionalProperties": False, + }, + "strict": False, + } + + def _json_schema_type(self, annotation: Any) -> str: + if annotation is bool: + return "boolean" + if annotation is int: + return "integer" + if annotation is float: + return "number" + return "string" + + def _call_tool( + self, + tool_call: ResponseFunctionToolCall, + tool_by_name: dict[str, NamedCallable], + ) -> str: + if tool_call.name not in tool_by_name: + raise ValueError(f"Unknown tool call: {tool_call.name}") + + arguments = json.loads(tool_call.arguments or "{}") + if not isinstance(arguments, dict): + raise ValueError(f"Tool call arguments must be an object: {arguments}") + + result = tool_by_name[tool_call.name](**arguments) + return result if isinstance(result, str) else json.dumps(result) diff --git a/c2rust-postprocess/postprocess/transforms/__init__.py b/c2rust-postprocess/postprocess/transforms/__init__.py index 8cc791ed13..2ba21fddff 100644 --- a/c2rust-postprocess/postprocess/transforms/__init__.py +++ b/c2rust-postprocess/postprocess/transforms/__init__.py @@ -1,7 +1,9 @@ from postprocess.cache import AbstractCache from postprocess.models import AbstractGenerativeModel +from postprocess.transforms.asserts import AssertsTransform from postprocess.transforms.base import AbstractTransform from postprocess.transforms.comments import CommentsTransform +from postprocess.transforms.formatting import FormattingTransform def get_transform_by_id( @@ -13,5 +15,9 @@ def get_transform_by_id( match id.lower(): case "comments": return CommentsTransform(cache=cache, model=model) + case "asserts": + return AssertsTransform(cache=cache, model=model) + case "formatting": + return FormattingTransform(cache=cache, model=model) case _: raise ValueError(f"Unsupported transform: {id}") diff --git a/c2rust-postprocess/postprocess/transforms/asserts.py b/c2rust-postprocess/postprocess/transforms/asserts.py new file mode 100644 index 0000000000..5dda0a2519 --- /dev/null +++ b/c2rust-postprocess/postprocess/transforms/asserts.py @@ -0,0 +1,133 @@ +import logging +from collections.abc import Callable +from pathlib import Path +from textwrap import dedent + +from postprocess.cache import AbstractCache +from postprocess.definitions import CDefinition +from postprocess.models import AbstractGenerativeModel +from postprocess.transforms.base import AbstractTransform, TransformError +from postprocess.utils import remove_backticks + +SYSTEM_INSTRUCTION = ( + "You are a helpful assistant that rewrites c2rust-transpiled assert patterns " + "into idiomatic Rust assert! macros." +) + + +class AssertsTransformPrompt: + c_function: str + rust_function: str + prompt_text: str + identifier: str + + __slots__ = ("c_function", "rust_function", "prompt_text", "identifier") + + def __init__( + self, c_function: str, rust_function: str, prompt_text: str, identifier: str + ): + self.c_function = c_function + self.rust_function = rust_function + self.prompt_text = prompt_text + self.identifier = identifier + + def __str__(self) -> str: + return ( + self.prompt_text + + "\n\n" + + "C function:\n```c\n" + + self.c_function + + "\n```\n\n" + + "Rust function:\n```rust\n" + + self.rust_function + + "\n```\n" + ) + + +class AssertsTransform(AbstractTransform): + def __init__(self, cache: AbstractCache, model: AbstractGenerativeModel): + super().__init__(SYSTEM_INSTRUCTION, cache, model) + + @staticmethod + def get_validation_fn(expected_assert_count: int) -> Callable[[str], str]: + def validate_response(rust_fn: str) -> str: + rust_fn = remove_backticks(rust_fn) + + if "__assert_fail(" in rust_fn: + return ( + "FAILURE: Rust function still contains __assert_fail. " + "Rewrite those into assert! calls. " + "Reply with the full Rust function definition only; " + "say nothing else." + ) + + actual_assert_count = rust_fn.count("assert!(") + if actual_assert_count < expected_assert_count: + return ( + "FAILURE: Missing rewritten assert! calls. " + f"Expected at least {expected_assert_count}, " + f"got {actual_assert_count}. " + "Reply with the full Rust function definition only; " + "say nothing else." + ) + + return "SUCCESS: Asserts transformed correctly!" + + return validate_response + + def try_apply_ident( + self, + rust_source_file: Path, + rust_definition: str, + c_definition: CDefinition, + identifier: str, + ) -> str | None: + _ = rust_source_file + expected_assert_count = rust_definition.count("__assert_fail(") + if expected_assert_count == 0: + logging.info( + f"{self.__class__.__name__}: " + f"Skipping function without transpiled asserts: {identifier}" + ) + return + + prompt_text = """ + Rewrite the Rust function below by replacing transpiled C assert-macro + expansions (which call __assert_fail) with idiomatic Rust assert! calls. + + Requirements: + - Preserve function behavior. + - Preserve formatting and indentation. + - Keep all non-assert logic unchanged. + - Return the full Rust function definition only; say nothing else. + """ + prompt_text = dedent(prompt_text).strip() + + prompt = AssertsTransformPrompt( + c_function=c_definition.effective, + rust_function=rust_definition, + prompt_text=prompt_text, + identifier=identifier, + ) + + messages = [{"role": "user", "content": str(prompt)}] + validate_response = self.get_validation_fn(expected_assert_count) + + def check(response: str) -> str: + if response.strip() == "": + raise TransformError("model returned an empty response") + + validation_result = validate_response(response) + if not validation_result.startswith("SUCCESS"): + raise TransformError( + f"model response for {identifier} failed validation: " + f"{validation_result}\nResponse was:\n{response}" + ) + return remove_backticks(response) + + return self.generate( + identifier, + messages, + check, + tools=[validate_response], + ) diff --git a/c2rust-postprocess/postprocess/transforms/base.py b/c2rust-postprocess/postprocess/transforms/base.py index 1175a36381..3f267cb11e 100644 --- a/c2rust-postprocess/postprocess/transforms/base.py +++ b/c2rust-postprocess/postprocess/transforms/base.py @@ -1,6 +1,6 @@ import logging import re -from collections.abc import Callable +from collections.abc import Callable, Iterable from dataclasses import dataclass, field from functools import partial from pathlib import Path @@ -119,6 +119,7 @@ def generate( identifier: str, messages: list[dict[str, Any]], check: Callable[[str], str], + tools: Iterable[Callable[..., Any]] = (), ) -> str | None: """ Model call with caching, validation, and retries. `check` returns the @@ -154,7 +155,7 @@ def generate( for attempt in range(self.max_attempts): try: - response = self.model.generate_with_tools(messages) + response = self.model.generate_with_tools(messages, tools=tools) if response is None: raise TransformError(f"model returned no response for {identifier}") result = check(response) diff --git a/c2rust-postprocess/postprocess/transforms/comments.py b/c2rust-postprocess/postprocess/transforms/comments.py index 27687cd5a2..19370776c3 100644 --- a/c2rust-postprocess/postprocess/transforms/comments.py +++ b/c2rust-postprocess/postprocess/transforms/comments.py @@ -40,10 +40,10 @@ def __str__(self) -> str: + "\n```\n\n" + "C function:\n```c\n" + self.c_function - + "```\n\n" + + "\n```\n\n" + "Rust function:\n```rust\n" + self.rust_function - + "```\n" + + "\n```\n" ) diff --git a/c2rust-postprocess/postprocess/transforms/formatting.py b/c2rust-postprocess/postprocess/transforms/formatting.py new file mode 100644 index 0000000000..5e1bba744b --- /dev/null +++ b/c2rust-postprocess/postprocess/transforms/formatting.py @@ -0,0 +1,161 @@ +import logging +import re +from collections.abc import Callable +from pathlib import Path +from textwrap import dedent + +from postprocess.cache import AbstractCache +from postprocess.definitions import CDefinition +from postprocess.models import AbstractGenerativeModel +from postprocess.transforms.base import AbstractTransform, TransformError +from postprocess.utils import remove_backticks + +SYSTEM_INSTRUCTION = ( + "You are a helpful assistant that conservatively reformats c2rust-transpiled " + "Rust functions for compactness while preserving idiomatic Rust formatting." +) + + +class FormattingTransformPrompt: + c_function: str + rust_function: str + prompt_text: str + identifier: str + + __slots__ = ("c_function", "rust_function", "prompt_text", "identifier") + + def __init__( + self, c_function: str, rust_function: str, prompt_text: str, identifier: str + ): + self.c_function = c_function + self.rust_function = rust_function + self.prompt_text = prompt_text + self.identifier = identifier + + def __str__(self) -> str: + return ( + self.prompt_text + + "\n\n" + + "C function:\n```c\n" + + self.c_function + + "\n```\n\n" + + "Rust function:\n```rust\n" + + self.rust_function + + "\n```\n" + ) + + +class FormattingTransform(AbstractTransform): + def __init__(self, cache: AbstractCache, model: AbstractGenerativeModel): + super().__init__(SYSTEM_INSTRUCTION, cache, model) + + @staticmethod + def should_attempt_formatting(c_function: str, rust_function: str) -> bool: + c_line_count = len([line for line in c_function.splitlines() if line.strip()]) + rust_line_count = len( + [line for line in rust_function.splitlines() if line.strip()] + ) + + if rust_line_count <= max(c_line_count * 2, c_line_count + 20): + return False + + compactable_item_lines = 0 + for line in rust_function.splitlines(): + stripped = line.strip() + if len(stripped) <= 48 and re.fullmatch(r"[^,{};]+,", stripped): + compactable_item_lines += 1 + + return " = [" in rust_function and compactable_item_lines >= 16 + + @staticmethod + def get_validation_fn(identifier: str) -> Callable[[str], str]: + def validate_response(rust_fn: str) -> str: + rust_fn = remove_backticks(rust_fn).strip() + + if not rust_fn: + return ( + "FAILURE: Empty response. Reply with the full Rust function " + "definition only; say nothing else." + ) + + if "```" in rust_fn: + return ( + "FAILURE: Response contains Markdown code fences. Reply with the " + "full Rust function definition only; say nothing else." + ) + + if f"fn {identifier}" not in rust_fn: + return ( + f"FAILURE: Response does not contain function `{identifier}`. " + "Reply with the full Rust function definition only; say nothing else." # noqa: E501 + ) + + return "SUCCESS: Function formatted correctly!" + + return validate_response + + def try_apply_ident( + self, + rust_source_file: Path, + rust_definition: str, + c_definition: CDefinition, + identifier: str, + ) -> str | None: + _ = rust_source_file + if not self.should_attempt_formatting(c_definition.effective, rust_definition): + logging.info( + f"{self.__class__.__name__}: " + f"Skipping function without obvious compactness issue: {identifier}" + ) + return + + prompt_text = """ + Reformat the Rust function below only where the transpiled formatting is + needlessly verbose compared with the corresponding C function. + + Most Rust functions should stay exactly as rustfmt would format them. Make + changes only for mechanically expanded tables, arrays, lookup data, or similar + data-heavy structures where the Rust version is much longer than the C version + because rustfmt placed one small element per line. + + Requirements: + - Preserve Rust syntax, behavior, attributes, signature, names, types, + expressions, comments, and control flow. + - Do not try to make ordinary Rust statements imitate C brace or indentation + style. Keep ordinary code idiomatic for Rust. + - For compacted data structures, take formatting clues from the C version: + group comparable numbers of elements per line, keep related comments near the + same data, and preserve useful visual structure. + - Add #[rustfmt::skip] to the function if needed so rustfmt will not expand the + compacted data structure again. + - If there is no clear table/array/data-structure compactness problem, return + the original Rust function unchanged. + - Return the full Rust function definition only; say nothing else. + """ + prompt_text = dedent(prompt_text).strip() + + prompt = FormattingTransformPrompt( + c_function=c_definition.effective, + rust_function=rust_definition, + prompt_text=prompt_text, + identifier=identifier, + ) + + messages = [{"role": "user", "content": str(prompt)}] + validate_response = self.get_validation_fn(identifier) + + def check(response: str) -> str: + validation_result = validate_response(response) + if not validation_result.startswith("SUCCESS"): + raise TransformError( + f"model response for {identifier} failed validation: " + f"{validation_result}\nResponse was:\n{response}" + ) + return remove_backticks(response) + + return self.generate( + identifier, + messages, + check, + tools=[validate_response], + )