postprocess: run function transformations asynchronously with bounded jobs

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.

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,14 +237,15 @@ 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,
                     update_rust=args.update_rust,
                     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()

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

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.

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()

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]] = (),