diff --git a/c2rust-postprocess/README.md b/c2rust-postprocess/README.md index 1b90e42e08..2819f56eb2 100644 --- a/c2rust-postprocess/README.md +++ b/c2rust-postprocess/README.md @@ -22,6 +22,12 @@ attempts to produce fully safe and idiomatic Rust; not a replacement for these e - `c2rust-postprocess path/to/transpiled_rust.rs`, or - `uv run postprocess path/to/transpiled_rust.rs` +Up to four functions in each Rust file are processed concurrently by default. +Use `-j 8` / `--jobs 8` to change the limit, or `-j 1` to process serially. +Each function's trimming and comment transfer still run in order. Files and +transform passes also run in order, and rewrites are applied and checked only +after all functions in a file finish. Existing cached responses are reused. + ## Excluding/Filtering `c2rust-postprocess` has a few ways to filter/exclude the function identifiers that are processed. diff --git a/c2rust-postprocess/postprocess/__init__.py b/c2rust-postprocess/postprocess/__init__.py index 1f6753af5b..13ad3dd608 100644 --- a/c2rust-postprocess/postprocess/__init__.py +++ b/c2rust-postprocess/postprocess/__init__.py @@ -3,6 +3,7 @@ """ import argparse +import asyncio import logging import os from argparse import BooleanOptionalAction @@ -20,7 +21,7 @@ from postprocess.models.mock import MockGenerativeModel from postprocess.transforms import get_transform_by_id from postprocess.transforms.base import TransformError, TransformResult -from postprocess.utils import existing_file +from postprocess.utils import existing_file, positive_int from postprocess.validate import BaselineError, make_validator DEFAULT_LLM_MODEL = "gpt-5.6-luna" @@ -37,6 +38,14 @@ def build_arg_parser() -> argparse.ArgumentParser: help="Path to Rust source file referenced by Cargo.toml", ) + parser.add_argument( + "-j", + "--jobs", + type=positive_int, + default=4, + help="Maximum concurrent function transformations per file (default: 4)", + ) + parser.add_argument( "--log-level", type=str, @@ -170,22 +179,36 @@ def get_model(model_id: str) -> AbstractGenerativeModel: def main(argv: Sequence[str] | None = None): try: - parser = build_arg_parser() - args = parser.parse_args(argv) + return asyncio.run(_main(argv)) + except BaselineError as error: + logging.error(error) + return 1 + except TransformError as error: + logging.exception(f"Aborting at first transform failure: {error}") + return 1 + except KeyboardInterrupt: + logging.warning("Interrupted by user, terminating...") + return 130 # 128 + SIGINT(2) - logging.basicConfig(level=logging.getLevelName(args.log_level.upper())) - if args.cache_dir is not None: - cache = DirectoryCache(args.cache_dir) - else: - cache = getattr(DirectoryCache, args.cache_scope)() - if args.update_cache and args.prune_cache_days > 0: - cache.prune(args.prune_cache_days) - if not args.update_cache: - cache = FrozenCache(cache) +async def _main(argv: Sequence[str] | None = None) -> int: + parser = build_arg_parser() + args = parser.parse_args(argv) - model = get_model(args.llm_model) + logging.basicConfig(level=logging.getLevelName(args.log_level.upper())) + if args.cache_dir is not None: + cache = DirectoryCache(args.cache_dir) + else: + cache = getattr(DirectoryCache, args.cache_scope)() + if args.update_cache and args.prune_cache_days > 0: + cache.prune(args.prune_cache_days) + if not args.update_cache: + cache = FrozenCache(cache) + + model = get_model(args.llm_model) + + try: # 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( @@ -214,7 +237,7 @@ def main(argv: Sequence[str] | None = None): ) for transform in transforms: result.extend( - transform.apply_dir( + await transform.apply_dir( root_rust_source_file=args.root_rust_source_file, exclude_list=IdentifierExcludeList(src_path=args.exclude_file), ident_filter=args.ident_filter, @@ -222,6 +245,7 @@ def main(argv: Sequence[str] | None = None): keep_going=args.on_error != "abort", failure_log_level=failure_log_level, validator=validator, + jobs=args.jobs, ) ) @@ -237,12 +261,5 @@ def main(argv: Sequence[str] | None = None): return 1 return 0 - except BaselineError as error: - logging.error(error) - return 1 - except TransformError as error: - logging.exception(f"Aborting at first transform failure: {error}") - return 1 - except KeyboardInterrupt: - logging.warning("Interrupted by user, terminating...") - return 130 # 128 + SIGINT(2) + finally: + await model.aclose() diff --git a/c2rust-postprocess/postprocess/models/base.py b/c2rust-postprocess/postprocess/models/base.py index 4473476b24..f4e63e4745 100644 --- a/c2rust-postprocess/postprocess/models/base.py +++ b/c2rust-postprocess/postprocess/models/base.py @@ -16,10 +16,14 @@ def id(self) -> str: return self._id @abstractmethod - def generate_with_tools( + async def generate_with_tools( self, messages: list[dict[str, Any]], tools: Iterable[Callable[..., Any]] = (), max_tool_loops: int = 5, ) -> str | None: pass + + async def aclose(self) -> None: + """Release any resources held by the model client.""" + return None diff --git a/c2rust-postprocess/postprocess/models/gemini.py b/c2rust-postprocess/postprocess/models/gemini.py index 092e8d22bb..ebf908ba02 100644 --- a/c2rust-postprocess/postprocess/models/gemini.py +++ b/c2rust-postprocess/postprocess/models/gemini.py @@ -21,7 +21,7 @@ def __init__( super().__init__(id) self.client = genai.Client(api_key=api_key) - def generate_with_tools( + async def generate_with_tools( self, messages: list[dict[str, Any]], tools: Iterable[Callable[..., Any]] = (), @@ -38,12 +38,18 @@ def generate_with_tools( ), ) - response = self.client.models.generate_content( + response = await self.client.aio.models.generate_content( model=self._id, contents=contents, config=config ) return response.text + async def aclose(self) -> None: + try: + await self.client.aio.aclose() + finally: + self.client.close() + def _convert_messages(self, messages: list[dict[str, Any]]) -> list[types.Content]: """ Converts standard list of dicts to Google GenAI 'Content' objects. diff --git a/c2rust-postprocess/postprocess/models/gpt.py b/c2rust-postprocess/postprocess/models/gpt.py index eb0c8904a4..4e8bfe0074 100644 --- a/c2rust-postprocess/postprocess/models/gpt.py +++ b/c2rust-postprocess/postprocess/models/gpt.py @@ -1,7 +1,7 @@ from collections.abc import Callable, Iterable from typing import Any -from openai import OpenAI +from openai import AsyncOpenAI from postprocess.models import AbstractGenerativeModel @@ -14,9 +14,9 @@ def __init__( base_url: str | None = None, ): super().__init__(id) - self.client = OpenAI(api_key=api_key, base_url=base_url) + self.client = AsyncOpenAI(api_key=api_key, base_url=base_url) - def generate_with_tools( + async def generate_with_tools( self, messages: list[dict[str, Any]], tools: Iterable[Callable[..., Any]] = (), @@ -25,10 +25,13 @@ def generate_with_tools( # TODO: implement tool calling support assert not tools, "Tool calling not yet implemented for GPTModel" - response = self.client.responses.create( + response = await self.client.responses.create( model=self.id, input=messages[0]["content"], max_tool_calls=max_tool_loops, ) return response.output_text + + async def aclose(self) -> None: + await self.client.close() diff --git a/c2rust-postprocess/postprocess/models/mock.py b/c2rust-postprocess/postprocess/models/mock.py index 9095ccd8c2..1cc7b172c7 100644 --- a/c2rust-postprocess/postprocess/models/mock.py +++ b/c2rust-postprocess/postprocess/models/mock.py @@ -13,7 +13,7 @@ class MockGenerativeModel(AbstractGenerativeModel): def __init__(self): super().__init__(id="mock-llm") - def generate_with_tools( + async def generate_with_tools( self, messages: list[dict[str, Any]], tools: Iterable[Callable[..., Any]] = (), diff --git a/c2rust-postprocess/postprocess/transforms/base.py b/c2rust-postprocess/postprocess/transforms/base.py index b6e21ecd8a..91139d5d08 100644 --- a/c2rust-postprocess/postprocess/transforms/base.py +++ b/c2rust-postprocess/postprocess/transforms/base.py @@ -1,3 +1,4 @@ +import asyncio import logging import re from collections.abc import Callable @@ -62,7 +63,7 @@ def __init__( def system_instruction(self) -> str: return self._system_instruction - def apply_ident( + async def apply_ident( self, rust_source_file: Path, rust_definition: str, @@ -74,7 +75,7 @@ def apply_ident( Apply the transform to one Rust definition and commit it through merge_rust. """ - new_definition = self.try_apply_ident( + new_definition = await self.try_apply_ident( rust_source_file=rust_source_file, rust_definition=rust_definition, c_definition=c_definition, @@ -99,7 +100,7 @@ def apply_ident( logging.info(f"{self.__class__.__name__}: Transformed Rust fn {identifier}") return new_definition - def try_apply_ident( + async def try_apply_ident( self, rust_source_file: Path, rust_definition: str, @@ -115,7 +116,7 @@ def try_apply_ident( "or override apply_ident directly" ) - def generate( + async def generate( self, identifier: str, messages: list[dict[str, Any]], @@ -155,7 +156,7 @@ def generate( for attempt in range(self.max_attempts): try: - response = self.model.generate_with_tools(messages) + response = await self.model.generate_with_tools(messages) if response is None: raise TransformError(f"model returned no response for {identifier}") result = validate(response) @@ -181,7 +182,7 @@ def generate( f"for {identifier}: {error}" ) from error - def apply_dir( + async def apply_dir( self, root_rust_source_file: Path, exclude_list: IdentifierExcludeList, @@ -190,6 +191,7 @@ def apply_dir( keep_going: bool = False, failure_log_level: int = logging.ERROR, validator: BatchValidator | None = None, + jobs: int = 1, ) -> TransformResult: """ Run `self.apply_file` on each `*.rs` in `dir` @@ -200,13 +202,13 @@ def apply_dir( result = TransformResult() root_dir = root_rust_source_file.parent c_decls_json_suffix = ".c_decls.json" - for c_decls_path in root_dir.glob(f"**/*{c_decls_json_suffix}"): + for c_decls_path in sorted(root_dir.glob(f"**/*{c_decls_json_suffix}")): rs_path = c_decls_path.with_name( c_decls_path.name.removesuffix(c_decls_json_suffix) + ".rs" ) assert rs_path.exists() result.extend( - self.apply_file( + await self.apply_file( rust_source_file=rs_path, exclude_list=exclude_list, ident_filter=ident_filter, @@ -214,11 +216,12 @@ def apply_dir( keep_going=keep_going, failure_log_level=failure_log_level, validator=validator, + jobs=jobs, ) ) return result - def apply_file( + async def apply_file( self, rust_source_file: Path, exclude_list: IdentifierExcludeList, @@ -227,7 +230,10 @@ def apply_file( keep_going: bool = False, failure_log_level: int = logging.ERROR, validator: BatchValidator | None = None, + jobs: int = 1, ) -> TransformResult: + if jobs < 1: + raise ValueError("jobs must be at least 1") ident_regex = re.compile(ident_filter) if ident_filter else None result = TransformResult() @@ -237,58 +243,95 @@ def apply_file( logging.info(f"Loaded {len(rust_definitions)} Rust definitions") logging.info(f"Loaded {len(c_definitions)} C definitions") - # Collect candidates without touching the file; application and - # validation happen per batch below. - candidates: list[Candidate] = [] - for identifier, rust_definition in rust_definitions.items(): - if exclude_list.contains(path=rust_source_file, identifier=identifier): - logging.info( - f"Skipping Rust fn {identifier} in {rust_source_file} " - f"due to exclude file {exclude_list.src_path}" - ) - continue - - if ident_regex and not ident_regex.search(identifier): - logging.info( - f"Skipping Rust fn {identifier} in {rust_source_file} " - f"due to ident filter {ident_filter}" - ) - continue - - if identifier not in c_definitions: - logging.warning(f"No corresponding C definition found for {identifier}") - continue - - c_definition = c_definitions[identifier] - - highlighted_c_definition = get_highlighted_c(c_definition.effective) - logging.debug( - f"C function {identifier} definition:\n{highlighted_c_definition}\n" - ) - - try: - new_definition = self.apply_ident( - rust_source_file=rust_source_file, - rust_definition=rust_definition, - c_definition=c_definition, - identifier=identifier, - update_rust=False, - ) - except TransformError as error: - if not keep_going: - raise - logging.log( - failure_log_level, - f"Transform failed for {identifier} in {rust_source_file}: {error}", + # Each worker completes a function's trim/retry sequence before taking + # another one. Cache operations stay on this event loop; no file writes + # or cargo checks run until all workers finish. + definitions = iter(rust_definitions.items()) + new_definitions: dict[str, str] = {} + failures: dict[str, TransformFailure] = {} + + async def worker() -> None: + for identifier, rust_definition in definitions: + if exclude_list.contains(path=rust_source_file, identifier=identifier): + logging.info( + f"Skipping Rust fn {identifier} in {rust_source_file} " + f"due to exclude file {exclude_list.src_path}" + ) + continue + + if ident_regex and not ident_regex.search(identifier): + logging.info( + f"Skipping Rust fn {identifier} in {rust_source_file} " + f"due to ident filter {ident_filter}" + ) + continue + + if identifier not in c_definitions: + logging.warning( + f"No corresponding C definition found for {identifier}" + ) + continue + + c_definition = c_definitions[identifier] + + highlighted_c_definition = get_highlighted_c(c_definition.effective) + logging.debug( + f"C function {identifier} definition:\n{highlighted_c_definition}\n" ) - result.failed.append( - (rust_source_file, identifier, "failed to transform") - ) - continue - if new_definition is None: + try: + new_definition = await self.apply_ident( + rust_source_file=rust_source_file, + rust_definition=rust_definition, + c_definition=c_definition, + identifier=identifier, + update_rust=False, + ) + except TransformError as error: + if not keep_going: + raise + logging.log( + failure_log_level, + f"Transform failed for {identifier} " + f"in {rust_source_file}: {error}", + ) + failures[identifier] = ( + rust_source_file, + identifier, + "failed to transform", + ) + continue + + if new_definition is None: + continue + + new_definitions[identifier] = new_definition + + workers = [ + asyncio.create_task(worker()) + for _ in range(min(jobs, len(rust_definitions))) + ] + try: + await asyncio.gather(*workers) + except BaseException: + # Abort and interruption must cancel pending model calls before + # returning or closing the shared client. + for task in workers: + task.cancel() + await asyncio.gather(*workers, return_exceptions=True) + raise + + # Apply and report in source order, independent of response timing. + result.failed = [ + failures[identifier] + for identifier in rust_definitions + if identifier in failures + ] + candidates: list[Candidate] = [] + for identifier in rust_definitions: + if identifier not in new_definitions: continue - + new_definition = new_definitions[identifier] candidates.append( Candidate( identifier=identifier, @@ -319,7 +362,7 @@ def apply_file( f"Validating {len(candidates)} rewrite(s) in {rust_source_file} " "with cargo check" ) - _, rejected = validator.validate(candidates) + _, rejected = await validator.validate(candidates) for candidate, error in rejected: candidate.invalidate() result.failed.append( diff --git a/c2rust-postprocess/postprocess/transforms/comments.py b/c2rust-postprocess/postprocess/transforms/comments.py index 8e6f1030f4..564e92e551 100644 --- a/c2rust-postprocess/postprocess/transforms/comments.py +++ b/c2rust-postprocess/postprocess/transforms/comments.py @@ -104,7 +104,7 @@ def __init__(self, cache: AbstractCache, model: AbstractGenerativeModel): super().__init__(SYSTEM_INSTRUCTION, cache, model) self.trim_transform = TrimTransform(cache, model) - def try_apply_ident( + async def try_apply_ident( self, rust_source_file: Path, rust_definition: str, @@ -124,7 +124,7 @@ def try_apply_ident( logging.info(f"Skipping C function without comments: {identifier}") return None - match self.trim_transform.apply_ident( + match await self.trim_transform.apply_ident( rust_source_file=rust_source_file, rust_definition=rust_definition, c_definition=c_definition, @@ -205,7 +205,7 @@ def validate(response: str) -> str: return rust_fn - rust_fn = self.generate(identifier, messages, validate) + rust_fn = await self.generate(identifier, messages, validate) if rust_fn is None: return None diff --git a/c2rust-postprocess/postprocess/transforms/trim.py b/c2rust-postprocess/postprocess/transforms/trim.py index 64226c27f2..0c70344f05 100644 --- a/c2rust-postprocess/postprocess/transforms/trim.py +++ b/c2rust-postprocess/postprocess/transforms/trim.py @@ -61,7 +61,7 @@ class TrimTransform(AbstractTransform): def __init__(self, cache: AbstractCache, model: AbstractGenerativeModel): super().__init__(SYSTEM_INSTRUCTION, cache, model) - def apply_ident( + async def apply_ident( self, rust_source_file: Path, rust_definition: str, @@ -122,7 +122,7 @@ def validate(response: str) -> str: return trimmed try: - return self.generate(identifier, messages, validate) + return await self.generate(identifier, messages, validate) except TransformError as error: # Trimming is best-effort; callers fall back to the untrimmed input. logging.warning(f"{self.__class__.__name__}: {error}") diff --git a/c2rust-postprocess/postprocess/utils.py b/c2rust-postprocess/postprocess/utils.py index b577fb187d..73555be7cb 100644 --- a/c2rust-postprocess/postprocess/utils.py +++ b/c2rust-postprocess/postprocess/utils.py @@ -48,6 +48,13 @@ def existing_file(value: str) -> Path: raise argparse.ArgumentTypeError(f"{value!r} is not a readable file") +def positive_int(value: str) -> int: + number = int(value) + if number < 1: + raise argparse.ArgumentTypeError("must be at least 1") + return number + + # TODO: test def get_compile_commands(compile_commands_path: Path) -> list[dict[str, Any]]: try: diff --git a/c2rust-postprocess/postprocess/validate.py b/c2rust-postprocess/postprocess/validate.py index 61192b9e45..238afaae2b 100644 --- a/c2rust-postprocess/postprocess/validate.py +++ b/c2rust-postprocess/postprocess/validate.py @@ -1,5 +1,6 @@ """Transactional validation of applied rewrites via `cargo check`.""" +import asyncio import json import logging import subprocess @@ -35,7 +36,7 @@ class BatchValidator: def __init__(self, check: Callable[[], str | None]): self._check = check - def validate( + async def validate( self, candidates: Sequence[Candidate] ) -> tuple[list[Candidate], list[tuple[Candidate, str]]]: """ @@ -56,7 +57,11 @@ def validate( try: for candidate in candidates: candidate.apply() + # SIGINT cancels the main asyncio task. Observe cancellation from + # synchronous merges/checks while rollback is still possible. + await asyncio.sleep(0) error = self._check() + await asyncio.sleep(0) ok = error is None finally: # `finally` rather than `except Exception` so KeyboardInterrupt @@ -75,8 +80,8 @@ def validate( # candidates, so interacting candidates are isolated correctly. logging.info(f"Check failed for batch of {len(candidates)}; bisecting") mid = len(candidates) // 2 - accepted, rejected = self.validate(candidates[:mid]) - right_accepted, right_rejected = self.validate(candidates[mid:]) + accepted, rejected = await self.validate(candidates[:mid]) + right_accepted, right_rejected = await self.validate(candidates[mid:]) return accepted + right_accepted, rejected + right_rejected diff --git a/c2rust-postprocess/tests/test_apply_file.py b/c2rust-postprocess/tests/test_apply_file.py index eaa985a983..8f52f684f6 100644 --- a/c2rust-postprocess/tests/test_apply_file.py +++ b/c2rust-postprocess/tests/test_apply_file.py @@ -1,3 +1,4 @@ +import asyncio import json from pathlib import Path from typing import Any @@ -115,11 +116,13 @@ def fake_update(*, root_rust_source_file, identifier, new_definition): def test_accepted_batch_is_applied(patched_io, rust_file: Path) -> None: cache, transform = make_comments_transform() - result = transform.apply_file( - rust_source_file=rust_file, - exclude_list=IdentifierExcludeList(None), - keep_going=True, - validator=BatchValidator(lambda: None), + result = asyncio.run( + transform.apply_file( + rust_source_file=rust_file, + exclude_list=IdentifierExcludeList(None), + keep_going=True, + validator=BatchValidator(lambda: None), + ) ) assert result.failed == [] @@ -132,11 +135,13 @@ def test_rejection_counts_failure_and_invalidates_final_response( ) -> None: cache, transform = make_comments_transform() - result = transform.apply_file( - rust_source_file=rust_file, - exclude_list=IdentifierExcludeList(None), - keep_going=True, - validator=BatchValidator(lambda: "type error"), + result = asyncio.run( + transform.apply_file( + rust_source_file=rust_file, + exclude_list=IdentifierExcludeList(None), + keep_going=True, + validator=BatchValidator(lambda: "type error"), + ) ) assert result.failed == [(rust_file, "enabled", "rejected by cargo check")] @@ -150,11 +155,13 @@ def test_rejection_raises_without_keep_going(patched_io, rust_file: Path) -> Non cache, transform = make_comments_transform() with pytest.raises(TransformError, match="cargo check rejected"): - transform.apply_file( - rust_source_file=rust_file, - exclude_list=IdentifierExcludeList(None), - keep_going=False, - validator=BatchValidator(lambda: "type error"), + asyncio.run( + transform.apply_file( + rust_source_file=rust_file, + exclude_list=IdentifierExcludeList(None), + keep_going=False, + validator=BatchValidator(lambda: "type error"), + ) ) assert rust_file.read_text() == "original\n" @@ -167,12 +174,14 @@ def test_no_update_rust_is_purely_generative(patched_io, rust_file: Path) -> Non def exploding_check() -> str | None: raise AssertionError("validator must not run") - result = transform.apply_file( - rust_source_file=rust_file, - exclude_list=IdentifierExcludeList(None), - update_rust=False, - keep_going=True, - validator=BatchValidator(exploding_check), + result = asyncio.run( + transform.apply_file( + rust_source_file=rust_file, + exclude_list=IdentifierExcludeList(None), + update_rust=False, + keep_going=True, + validator=BatchValidator(exploding_check), + ) ) assert result.failures == 0 @@ -182,11 +191,13 @@ def exploding_check() -> str | None: def test_without_validator_candidates_are_applied(patched_io, rust_file: Path) -> None: cache, transform = make_comments_transform() - result = transform.apply_file( - rust_source_file=rust_file, - exclude_list=IdentifierExcludeList(None), - keep_going=True, - validator=None, + result = asyncio.run( + transform.apply_file( + rust_source_file=rust_file, + exclude_list=IdentifierExcludeList(None), + keep_going=True, + validator=None, + ) ) assert result.failures == 0 @@ -215,7 +226,7 @@ def __init__(self, rewrites: dict[str, str]): super().__init__("", CannedCache({}), MockGenerativeModel()) self.rewrites = rewrites - def try_apply_ident( + async def try_apply_ident( self, rust_source_file: Path, rust_definition: str, @@ -268,11 +279,13 @@ def test_end_to_end_type_invalid_rewrite_is_isolated(tmp_path: Path) -> None: checker = CargoChecker(manifest) assert checker() is None # baseline - result = transform.apply_file( - rust_source_file=lib_rs, - exclude_list=IdentifierExcludeList(None), - keep_going=True, - validator=BatchValidator(checker), + result = asyncio.run( + transform.apply_file( + rust_source_file=lib_rs, + exclude_list=IdentifierExcludeList(None), + keep_going=True, + validator=BatchValidator(checker), + ) ) content = lib_rs.read_text() diff --git a/c2rust-postprocess/tests/test_cli.py b/c2rust-postprocess/tests/test_cli.py new file mode 100644 index 0000000000..f6e257fe4c --- /dev/null +++ b/c2rust-postprocess/tests/test_cli.py @@ -0,0 +1,60 @@ +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +import postprocess +from postprocess.transforms.base import TransformError, TransformResult +from postprocess.validate import BaselineError + + +@pytest.mark.parametrize("value", ["0", "-1", "many"]) +def test_jobs_rejects_invalid_values(tmp_path: Path, value: str) -> None: + source = tmp_path / "lib.rs" + source.touch() + with pytest.raises(SystemExit) as error: + postprocess.build_arg_parser().parse_args([str(source), "-j", value]) + assert error.value.code == 2 + + +@pytest.mark.parametrize("option", ["-j", "--jobs"]) +def test_jobs_option_and_default(tmp_path: Path, option: str) -> None: + source = tmp_path / "lib.rs" + source.touch() + parser = postprocess.build_arg_parser() + assert parser.parse_args([str(source)]).jobs == 4 + assert parser.parse_args([str(source), option, "8"]).jobs == 8 + + +@pytest.mark.parametrize("failure", [None, "baseline", "transform"]) +def test_main_forwards_jobs_and_closes_model(monkeypatch, tmp_path, failure) -> None: + source = tmp_path / "lib.rs" + source.touch() + close = AsyncMock() + model = SimpleNamespace(aclose=close) + monkeypatch.setattr(postprocess, "get_model", lambda _: model) + + apply = AsyncMock(return_value=TransformResult()) + if failure == "transform": + apply.side_effect = TransformError("rejected") + transform = SimpleNamespace(apply_dir=apply) + monkeypatch.setattr(postprocess, "get_transform_by_id", lambda *a, **kw: transform) + + validator = Mock(return_value=None) + if failure == "baseline": + validator.side_effect = BaselineError("broken baseline") + monkeypatch.setattr(postprocess, "make_validator", validator) + + result = postprocess.main( + [str(source), "--cache-dir", str(tmp_path / "cache"), "-j", "8"] + ) + + assert result == (1 if failure else 0) + if failure == "baseline": + apply.assert_not_awaited() + else: + apply.assert_awaited_once() + assert apply.await_args is not None + assert apply.await_args.kwargs["jobs"] == 8 + close.assert_awaited_once_with() diff --git a/c2rust-postprocess/tests/test_comments_transform.py b/c2rust-postprocess/tests/test_comments_transform.py index ee594a0e71..d522980394 100644 --- a/c2rust-postprocess/tests/test_comments_transform.py +++ b/c2rust-postprocess/tests/test_comments_transform.py @@ -1,3 +1,4 @@ +import asyncio from pathlib import Path from typing import Any @@ -79,12 +80,14 @@ def test_directive_line_comment_survives_preprocessed_check() -> None: 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, + asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition=rust_definition, + c_definition=c_definition, + identifier="enabled", + update_rust=False, + ) ) transforms = [transform for transform, _ in cache.lookups] @@ -128,12 +131,14 @@ def fake_update(*, root_rust_source_file, identifier, new_definition): monkeypatch.setattr(base, "update_rust_definition", fake_update) - transform.apply_ident( - rust_source_file=Path("unused.rs"), - rust_definition=RUST_DEFINITION_NO_COMMENTS, - c_definition=C_DEFINITION_BODY_COMMENT, - identifier="f", - update_rust=True, + asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition=RUST_DEFINITION_NO_COMMENTS, + c_definition=C_DEFINITION_BODY_COMMENT, + identifier="f", + update_rust=True, + ) ) assert "/// @note" not in merged["f"] @@ -150,12 +155,14 @@ def test_syntactically_invalid_response_is_rejected() -> None: transform = CommentsTransform(cache=cache, model=MockGenerativeModel()) with pytest.raises(TransformError, match="not syntactically valid"): - transform.apply_ident( - rust_source_file=Path("unused.rs"), - rust_definition=RUST_DEFINITION_NO_COMMENTS, - c_definition=C_DEFINITION_BODY_COMMENT, - identifier="f", - update_rust=False, + asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition=RUST_DEFINITION_NO_COMMENTS, + c_definition=C_DEFINITION_BODY_COMMENT, + identifier="f", + update_rust=False, + ) ) @@ -169,12 +176,14 @@ def test_response_that_changes_non_comment_code_is_rejected() -> None: transform = CommentsTransform(cache=cache, model=MockGenerativeModel()) with pytest.raises(TransformError, match="non-comment Rust code changed"): - transform.apply_ident( - rust_source_file=Path("unused.rs"), - rust_definition=RUST_DEFINITION_NO_COMMENTS, - c_definition=C_DEFINITION_BODY_COMMENT, - identifier="f", - update_rust=False, + asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition=RUST_DEFINITION_NO_COMMENTS, + c_definition=C_DEFINITION_BODY_COMMENT, + identifier="f", + update_rust=False, + ) ) @@ -191,12 +200,14 @@ def test_response_may_reformat_non_comment_code() -> None: cache = StaticCache(response) transform = CommentsTransform(cache=cache, model=MockGenerativeModel()) - result = transform.apply_ident( - rust_source_file=Path("unused.rs"), - rust_definition=RUST_DEFINITION_NO_COMMENTS, - c_definition=C_DEFINITION_BODY_COMMENT, - identifier="f", - update_rust=False, + result = asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition=RUST_DEFINITION_NO_COMMENTS, + c_definition=C_DEFINITION_BODY_COMMENT, + identifier="f", + update_rust=False, + ) ) assert result is not None @@ -246,7 +257,7 @@ def __init__(self, responses: list[str]): self.responses = responses self.calls = 0 - def generate_with_tools(self, messages, tools=(), max_tool_loops=5): + async def generate_with_tools(self, messages, tools=(), max_tool_loops=5): self.calls += 1 return self.responses.pop(0) @@ -273,12 +284,14 @@ def apply_to_body_comment_fn( cache: AbstractCache, model: AbstractGenerativeModel ) -> str | None: transform = CommentsTransform(cache=cache, model=model) - return transform.apply_ident( - rust_source_file=Path("unused.rs"), - rust_definition=RUST_DEFINITION_NO_COMMENTS, - c_definition=C_DEFINITION_BODY_COMMENT, - identifier="f", - update_rust=False, + return asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition=RUST_DEFINITION_NO_COMMENTS, + c_definition=C_DEFINITION_BODY_COMMENT, + identifier="f", + update_rust=False, + ) ) diff --git a/c2rust-postprocess/tests/test_models.py b/c2rust-postprocess/tests/test_models.py new file mode 100644 index 0000000000..2416f7f3c1 --- /dev/null +++ b/c2rust-postprocess/tests/test_models.py @@ -0,0 +1,92 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from google.genai import types + +from postprocess.models import gemini, gpt + + +def test_gpt_awaits_response_and_closes_client(monkeypatch) -> None: + create = AsyncMock(return_value=SimpleNamespace(output_text="rewritten function")) + close = AsyncMock() + client = SimpleNamespace(responses=SimpleNamespace(create=create), close=close) + constructor = Mock(return_value=client) + monkeypatch.setattr(gpt, "AsyncOpenAI", constructor) + + async def run() -> str: + model = gpt.GPTModel( + id="test/model", api_key="test-key", base_url="https://example.test/v1" + ) + try: + return await model.generate_with_tools( + [{"role": "user", "content": "function prompt"}], max_tool_loops=2 + ) + finally: + await model.aclose() + + assert asyncio.run(run()) == "rewritten function" + constructor.assert_called_once_with( + api_key="test-key", base_url="https://example.test/v1" + ) + create.assert_awaited_once_with( + model="test/model", input="function prompt", max_tool_calls=2 + ) + close.assert_awaited_once_with() + + +@pytest.mark.parametrize("response_text", ["rewritten function", None]) +def test_gemini_awaits_response_and_closes_both_clients( + monkeypatch, response_text: str | None +) -> None: + generate = AsyncMock(return_value=SimpleNamespace(text=response_text)) + async_close = AsyncMock() + sync_close = Mock() + client = SimpleNamespace( + aio=SimpleNamespace( + models=SimpleNamespace(generate_content=generate), aclose=async_close + ), + close=sync_close, + ) + constructor = Mock(return_value=client) + monkeypatch.setattr(gemini.genai, "Client", constructor) + + def validate(response: str) -> str: + return response + + async def run() -> str | None: + model = gemini.GoogleGenerativeModel(id="gemini-test", api_key="test-key") + try: + return await model.generate_with_tools( + [ + {"role": "user", "content": "function prompt"}, + {"role": "assistant", "content": "previous response"}, + ], + tools=[validate], + max_tool_loops=2, + ) + finally: + await model.aclose() + + assert asyncio.run(run()) == response_text + constructor.assert_called_once_with(api_key="test-key") + generate.assert_awaited_once_with( + model="gemini-test", + contents=[ + types.Content( + role="user", parts=[types.Part.from_text(text="function prompt")] + ), + types.Content( + role="model", parts=[types.Part.from_text(text="previous response")] + ), + ], + config=types.GenerateContentConfig( + tools=[validate], + automatic_function_calling=types.AutomaticFunctionCallingConfig( + disable=False, maximum_remote_calls=2 + ), + ), + ) + async_close.assert_awaited_once_with() + sync_close.assert_called_once_with() diff --git a/c2rust-postprocess/tests/test_parallel.py b/c2rust-postprocess/tests/test_parallel.py new file mode 100644 index 0000000000..baa0c01ad2 --- /dev/null +++ b/c2rust-postprocess/tests/test_parallel.py @@ -0,0 +1,252 @@ +import asyncio +from collections.abc import Awaitable, Callable +from pathlib import Path + +import pytest +from test_apply_file import CannedCache + +from postprocess.cache import DirectoryCache, FrozenCache +from postprocess.definitions import CDefinition +from postprocess.exclude_list import IdentifierExcludeList +from postprocess.models.base import AbstractGenerativeModel +from postprocess.models.mock import MockGenerativeModel +from postprocess.transforms import base +from postprocess.transforms.base import AbstractTransform, TransformError + +IDENTIFIERS = ["first", "second", "third", "fourth"] + + +class AsyncTransform(AbstractTransform): + def __init__(self, generate: Callable[[str], Awaitable[str | None]]): + super().__init__("", CannedCache({}), MockGenerativeModel()) + self.generate_definition = generate + + async def try_apply_ident( + self, + rust_source_file: Path, + rust_definition: str, + c_definition: CDefinition, + identifier: str, + ) -> str | None: + return await self.generate_definition(identifier) + + +@pytest.fixture +def source(monkeypatch, tmp_path: Path) -> tuple[Path, list[str]]: + path = tmp_path / "lib.rs" + path.write_text("original\n") + writes: list[str] = [] + monkeypatch.setattr( + base, + "get_rust_definitions", + lambda path: {identifier: "original" for identifier in IDENTIFIERS}, + ) + monkeypatch.setattr( + base, + "get_c_definitions", + lambda path: { + identifier: CDefinition( + definition="void f(void) {}", preprocessed_definition=None + ) + for identifier in IDENTIFIERS + }, + ) + + def update(*, root_rust_source_file, identifier, new_definition): + writes.append(identifier) + root_rust_source_file.write_text(new_definition) + + monkeypatch.setattr(base, "update_rust_definition", update) + return path, writes + + +@pytest.mark.parametrize("jobs", [1, 2]) +def test_jobs_bounds_concurrent_transformations(source, jobs: int) -> None: + path, writes = source + + async def run() -> None: + started: asyncio.Queue[str] = asyncio.Queue() + release: asyncio.Queue[None] = asyncio.Queue() + active = 0 + peak = 0 + + async def generate(identifier: str) -> str: + nonlocal active, peak + active += 1 + peak = max(peak, active) + started.put_nowait(identifier) + try: + await release.get() + return identifier + finally: + active -= 1 + + task = asyncio.create_task( + AsyncTransform(generate).apply_file( + path, IdentifierExcludeList(None), jobs=jobs + ) + ) + for _ in range(jobs): + await asyncio.wait_for(started.get(), timeout=2) + assert active == jobs + assert writes == [] + for _ in IDENTIFIERS: + release.put_nowait(None) + result = await asyncio.wait_for(task, timeout=2) + assert result.failures == 0 + assert peak == jobs + assert active == 0 + + asyncio.run(run()) + assert writes == IDENTIFIERS + + +def test_out_of_order_responses_are_applied_in_source_order(source) -> None: + path, writes = source + + async def run() -> None: + started: asyncio.Queue[str] = asyncio.Queue() + completed: asyncio.Queue[str] = asyncio.Queue() + release = {identifier: asyncio.Event() for identifier in IDENTIFIERS} + + async def generate(identifier: str) -> str: + started.put_nowait(identifier) + await release[identifier].wait() + completed.put_nowait(identifier) + return identifier + + task = asyncio.create_task( + AsyncTransform(generate).apply_file( + path, IdentifierExcludeList(None), jobs=4 + ) + ) + for _ in IDENTIFIERS: + await asyncio.wait_for(started.get(), timeout=2) + for identifier in reversed(IDENTIFIERS): + assert writes == [] + release[identifier].set() + assert await asyncio.wait_for(completed.get(), timeout=2) == identifier + await asyncio.wait_for(task, timeout=2) + + asyncio.run(run()) + assert writes == IDENTIFIERS + + +def test_keep_going_applies_successes_and_reports_failure(source) -> None: + path, writes = source + + async def generate(identifier: str) -> str: + await asyncio.sleep(0) + if identifier == "second": + raise TransformError("rejected") + return identifier + + async def run(): + return await asyncio.wait_for( + AsyncTransform(generate).apply_file( + path, IdentifierExcludeList(None), jobs=2, keep_going=True + ), + timeout=2, + ) + + result = asyncio.run(run()) + assert result.failed == [(path, "second", "failed to transform")] + assert writes == ["first", "third", "fourth"] + + +def test_abort_cancels_pending_transformations_before_writing(source) -> None: + path, writes = source + + async def run() -> None: + first_started = asyncio.Event() + cancelled = asyncio.Event() + pending = asyncio.Event() + active: set[str] = set() + + async def generate(identifier: str) -> str: + active.add(identifier) + try: + if identifier == "second": + await first_started.wait() + raise TransformError("rejected") + if identifier == "first": + first_started.set() + await pending.wait() + return identifier + except asyncio.CancelledError: + if identifier == "first": + cancelled.set() + raise + finally: + active.remove(identifier) + + with pytest.raises(TransformError, match="rejected"): + await asyncio.wait_for( + AsyncTransform(generate).apply_file( + path, IdentifierExcludeList(None), jobs=2 + ), + timeout=2, + ) + assert cancelled.is_set() + assert not active + assert writes == [] + assert path.read_text() == "original\n" + + asyncio.run(run()) + + +def test_parallel_responses_are_reused_by_serial_frozen_run( + source, tmp_path: Path, monkeypatch +) -> None: + path, writes = source + + class CountingModel(AbstractGenerativeModel): + def __init__(self): + super().__init__("test-model") + self.calls = 0 + + async def generate_with_tools(self, messages, tools=(), max_tool_loops=5): + self.calls += 1 + await asyncio.sleep(0) + return messages[0]["content"] + + class CachedTransform(AbstractTransform): + async def try_apply_ident( + self, rust_source_file, rust_definition, c_definition, identifier + ): + return await self.generate( + identifier, + [{"role": "user", "content": identifier}], + lambda response: response, + ) + + cache = DirectoryCache(tmp_path / "cache") + model = CountingModel() + result = asyncio.run( + CachedTransform("", cache, model).apply_file( + path, IdentifierExcludeList(None), jobs=4 + ) + ) + assert result.failures == 0 + assert writes == IDENTIFIERS + assert model.calls == len(IDENTIFIERS) + assert len(list(cache.path.glob("*/*/*/metadata.toml"))) == len(IDENTIFIERS) + first_output = path.read_text() + + path.write_text("original\n") + writes.clear() + frozen = FrozenCache(DirectoryCache(cache.path)) + + def unexpected_update(**kwargs): + raise AssertionError("cache hits must not update responses") + + monkeypatch.setattr(frozen, "update", unexpected_update) + result = asyncio.run( + CachedTransform("", frozen, model).apply_file( + path, IdentifierExcludeList(None), jobs=1 + ) + ) + assert result.failures == 0 + assert writes == IDENTIFIERS + assert path.read_text() == first_output + assert model.calls == len(IDENTIFIERS) diff --git a/c2rust-postprocess/tests/test_trim_transform.py b/c2rust-postprocess/tests/test_trim_transform.py index 22e4f1439b..114e93832d 100644 --- a/c2rust-postprocess/tests/test_trim_transform.py +++ b/c2rust-postprocess/tests/test_trim_transform.py @@ -1,3 +1,4 @@ +import asyncio from pathlib import Path from typing import Any @@ -109,12 +110,14 @@ def invalidate(self, *, transform: str, identifier: str) -> None: def apply_trim(c_definition: CDefinition, cache: AbstractCache) -> str | None: transform = TrimTransform(cache=cache, model=MockGenerativeModel()) - return transform.apply_ident( - rust_source_file=Path("unused.rs"), - rust_definition="fn f() {}", - c_definition=c_definition, - identifier="f", - update_rust=False, + return asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition="fn f() {}", + c_definition=c_definition, + identifier="f", + update_rust=False, + ) ) diff --git a/c2rust-postprocess/tests/test_validate.py b/c2rust-postprocess/tests/test_validate.py index e7653bc109..2fa8839770 100644 --- a/c2rust-postprocess/tests/test_validate.py +++ b/c2rust-postprocess/tests/test_validate.py @@ -1,3 +1,4 @@ +import asyncio import shutil import textwrap from pathlib import Path @@ -47,7 +48,7 @@ def test_valid_batch_is_retained_with_one_check(source: Path) -> None: check = ContentCheck(source) candidates = [append_candidate(source, "a"), append_candidate(source, "b")] - accepted, rejected = BatchValidator(check).validate(candidates) + accepted, rejected = asyncio.run(BatchValidator(check).validate(candidates)) assert accepted == candidates assert rejected == [] @@ -64,7 +65,7 @@ def test_mixed_batch_retains_only_valid_candidates(source: Path) -> None: append_candidate(source, "d"), ] - accepted, rejected = BatchValidator(check).validate(candidates) + accepted, rejected = asyncio.run(BatchValidator(check).validate(candidates)) assert [c.identifier for c in accepted] == ["a", "c", "d"] assert [(c.identifier, error) for c, error in rejected] == [("BAD", "found BAD")] @@ -72,8 +73,8 @@ def test_mixed_batch_retains_only_valid_candidates(source: Path) -> None: def test_single_rejected_candidate_restores_baseline_exactly(source: Path) -> None: - accepted, rejected = BatchValidator(ContentCheck(source)).validate( - [append_candidate(source, "BAD")] + accepted, rejected = asyncio.run( + BatchValidator(ContentCheck(source)).validate([append_candidate(source, "BAD")]) ) assert accepted == [] @@ -90,8 +91,10 @@ def check() -> str | None: return "x and y conflict" return None - accepted, rejected = BatchValidator(check).validate( - [append_candidate(source, "x"), append_candidate(source, "y")] + accepted, rejected = asyncio.run( + BatchValidator(check).validate( + [append_candidate(source, "x"), append_candidate(source, "y")] + ) ) assert [c.identifier for c in accepted] == ["x"] @@ -104,7 +107,7 @@ def check() -> str | None: raise RuntimeError("cargo crashed") with pytest.raises(RuntimeError, match="cargo crashed"): - BatchValidator(check).validate([append_candidate(source, "a")]) + asyncio.run(BatchValidator(check).validate([append_candidate(source, "a")])) assert source.read_text() == "baseline\n" @@ -117,13 +120,47 @@ def broken_apply() -> None: identifier="a", files=(source,), apply=broken_apply, invalidate=lambda: None ) with pytest.raises(RuntimeError, match="merge failed"): - BatchValidator(ContentCheck(source)).validate([candidate]) + asyncio.run(BatchValidator(ContentCheck(source)).validate([candidate])) assert source.read_text() == "baseline\n" +@pytest.mark.parametrize("cancel_during", ["apply", "check"]) +def test_pending_cancellation_restores_file(source: Path, cancel_during: str) -> None: + check_calls = 0 + + async def run() -> None: + task = asyncio.current_task() + assert task is not None + + def apply() -> None: + source.write_text("rewritten\n") + if cancel_during == "apply": + task.cancel() + + def check() -> str | None: + nonlocal check_calls + check_calls += 1 + assert source.read_text() == "rewritten\n" + if cancel_during == "check": + task.cancel() + return None + + candidate = Candidate( + identifier="f", files=(source,), apply=apply, invalidate=lambda: None + ) + await BatchValidator(check).validate([candidate]) + + # asyncio.run handles SIGINT by cancelling the main task; synchronous + # subprocess calls can return normally with that cancellation pending. + with pytest.raises(asyncio.CancelledError): + asyncio.run(run()) + assert source.read_text() == "baseline\n" + assert check_calls == (1 if cancel_during == "check" else 0) + + def test_empty_batch_runs_no_check(source: Path) -> None: check = ContentCheck(source) - assert BatchValidator(check).validate([]) == ([], []) + assert asyncio.run(BatchValidator(check).validate([])) == ([], []) assert check.calls == 0