shrinkray 0.0.0__py3-none-any.whl → 25.12.26__py3-none-any.whl

shrinkray/subprocess/worker.py

@@ -0,0 +1,491 @@
+"""Worker subprocess that runs the reducer with trio and communicates via JSON protocol."""
+
+import os
+import sys
+import time
+import traceback
+from typing import Any, Protocol
+
+import trio
+from binaryornot.helpers import is_binary_string
+
+from shrinkray.problem import InvalidInitialExample
+from shrinkray.subprocess.protocol import (
+    PassStatsData,
+    ProgressUpdate,
+    Request,
+    Response,
+    deserialize,
+    serialize,
+)
+
+
+class InputStream(Protocol):
+    """Protocol for async input streams."""
+
+    def __aiter__(self) -> "InputStream": ...
+    async def __anext__(self) -> bytes | bytearray: ...
+
+
+class OutputStream(Protocol):
+    """Protocol for output streams."""
+
+    async def send(self, data: bytes) -> None: ...
+
+
+class StdoutStream:
+    """Wrapper around sys.stdout for the OutputStream protocol."""
+
+    async def send(self, data: bytes) -> None:
+        sys.stdout.write(data.decode("utf-8"))
+        sys.stdout.flush()
+
+
+class ReducerWorker:
+    """Runs the reducer in a subprocess with JSON protocol communication."""
+
+    def __init__(
+        self,
+        input_stream: InputStream | None = None,
+        output_stream: OutputStream | None = None,
+    ):
+        self.running = False
+        self.reducer = None
+        self.problem = None
+        self.state = None
+        self._cancel_scope: trio.CancelScope | None = None
+        # Parallelism tracking
+        self._parallel_samples = 0
+        self._parallel_total = 0
+        # I/O streams - None means use stdin/stdout
+        self._input_stream = input_stream
+        self._output_stream = output_stream
+
+    async def emit(self, msg: Response | ProgressUpdate) -> None:
+        """Write a message to the output stream."""
+        line = serialize(msg) + "\n"
+        if self._output_stream is not None:
+            await self._output_stream.send(line.encode("utf-8"))
+        else:
+            sys.stdout.write(line)
+            sys.stdout.flush()
+
+    async def read_commands(
+        self,
+        input_stream: InputStream | None = None,
+        task_status: trio.TaskStatus[None] = trio.TASK_STATUS_IGNORED,
+    ) -> None:
+        """Read commands from input stream and dispatch them."""
+        task_status.started()
+
+        # Use provided stream, or instance stream, or default to stdin
+        stream: InputStream
+        if input_stream is not None:
+            stream = input_stream
+        elif self._input_stream is not None:
+            stream = self._input_stream
+        else:
+            stream = trio.lowlevel.FdStream(os.dup(sys.stdin.fileno()))
+
+        buffer = b""
+        async for chunk in stream:
+            buffer += chunk
+            while b"\n" in buffer:
+                line, buffer = buffer.split(b"\n", 1)
+                if line:
+                    await self.handle_line(line.decode("utf-8"))
+
+    async def handle_line(self, line: str) -> None:
+        """Handle a single command line."""
+        try:
+            request = deserialize(line)
+            if not isinstance(request, Request):
+                await self.emit(Response(id="", error="Expected a request"))
+                return
+            response = await self.handle_command(request)
+            await self.emit(response)
+        except Exception as e:
+            traceback.print_exc()
+            await self.emit(Response(id="", error=str(e)))
+
+    async def handle_command(self, request: Request) -> Response:
+        """Handle a command request and return a response."""
+        match request.command:
+            case "start":
+                return await self._handle_start(request.id, request.params)
+            case "status":
+                return self._handle_status(request.id)
+            case "cancel":
+                return self._handle_cancel(request.id)
+            case "disable_pass":
+                return self._handle_disable_pass(request.id, request.params)
+            case "enable_pass":
+                return self._handle_enable_pass(request.id, request.params)
+            case "skip_pass":
+                return self._handle_skip_pass(request.id)
+            case _:
+                return Response(id=request.id, error=f"Unknown command: {request.command}")
+
+    async def _handle_start(self, request_id: str, params: dict) -> Response:
+        """Start the reduction process."""
+        if self.running:
+            return Response(id=request_id, error="Already running")
+
+        try:
+            await self._start_reduction(params)
+            return Response(id=request_id, result={"status": "started"})
+        except* InvalidInitialExample as excs:
+            assert len(excs.exceptions) == 1
+            (e,) = excs.exceptions
+            # Build a detailed error message for invalid initial examples
+            if self.state is not None:
+                error_message = await self.state.build_error_message(e)
+            else:
+                error_message = str(e)
+        except* Exception as e:
+            traceback.print_exc()
+            error_message = str(e.exceptions[0])
+        return Response(id=request_id, error=error_message)
+
+    async def _start_reduction(self, params: dict) -> None:
+        """Initialize and start the reduction."""
+        from shrinkray.cli import InputType
+        from shrinkray.passes.clangdelta import (
+            C_FILE_EXTENSIONS,
+            ClangDelta,
+            find_clang_delta,
+        )
+        from shrinkray.state import (
+            ShrinkRayDirectoryState,
+            ShrinkRayStateSingleFile,
+        )
+        from shrinkray.work import Volume
+
+        filename = params["file_path"]
+        test = params["test"]
+        parallelism = params.get("parallelism", os.cpu_count() or 1)
+        timeout = params.get("timeout", 1.0)
+        seed = params.get("seed", 0)
+        input_type = InputType[params.get("input_type", "all")]
+        in_place = params.get("in_place", False)
+        formatter = params.get("formatter", "default")
+        volume = Volume[params.get("volume", "normal")]
+        no_clang_delta = params.get("no_clang_delta", False)
+        clang_delta_path = params.get("clang_delta", "")
+        trivial_is_error = params.get("trivial_is_error", True)
+
+        clang_delta_executable = None
+        if os.path.splitext(filename)[1] in C_FILE_EXTENSIONS and not no_clang_delta:
+            if not clang_delta_path:
+                clang_delta_path = find_clang_delta()
+            if clang_delta_path:
+                clang_delta_executable = ClangDelta(clang_delta_path)
+
+        state_kwargs: dict[str, Any] = dict(
+            input_type=input_type,
+            in_place=in_place,
+            test=test,
+            timeout=timeout,
+            base=os.path.basename(filename),
+            parallelism=parallelism,
+            filename=filename,
+            formatter=formatter,
+            trivial_is_error=trivial_is_error,
+            seed=seed,
+            volume=volume,
+            clang_delta_executable=clang_delta_executable,
+        )
+
+        if os.path.isdir(filename):
+            files = [os.path.join(d, f) for d, _, fs in os.walk(filename) for f in fs]
+            initial = {}
+            for f in files:
+                with open(f, "rb") as i:
+                    initial[os.path.relpath(f, filename)] = i.read()
+            self.state = ShrinkRayDirectoryState(initial=initial, **state_kwargs)
+        else:
+            with open(filename, "rb") as reader:
+                initial = reader.read()
+            self.state = ShrinkRayStateSingleFile(initial=initial, **state_kwargs)
+
+        self.problem = self.state.problem
+        self.reducer = self.state.reducer
+
+        # Validate initial example before starting - this will raise
+        # InvalidInitialExample if the initial test case fails
+        await self.problem.setup()
+
+        self.running = True
+
+    def _handle_status(self, request_id: str) -> Response:
+        """Get current status."""
+        if not self.running or self.problem is None:
+            return Response(id=request_id, result={"running": False})
+
+        stats = self.problem.stats
+        return Response(
+            id=request_id,
+            result={
+                "running": True,
+                "status": self.reducer.status if self.reducer else "",
+                "size": stats.current_test_case_size,
+                "original_size": stats.initial_test_case_size,
+                "calls": stats.calls,
+                "reductions": stats.reductions,
+            },
+        )
+
+    def _handle_cancel(self, request_id: str) -> Response:
+        """Cancel the reduction."""
+        if self._cancel_scope is not None:
+            self._cancel_scope.cancel()
+        self.running = False
+        return Response(id=request_id, result={"status": "cancelled"})
+
+    def _get_known_pass_names(self) -> set[str]:
+        """Get the set of known pass names from pass stats."""
+        if self.reducer is None or self.reducer.pass_stats is None:
+            return set()
+        return set(self.reducer.pass_stats._stats.keys())
+
+    def _handle_disable_pass(self, request_id: str, params: dict) -> Response:
+        """Disable a reduction pass by name."""
+        pass_name = params.get("pass_name", "")
+        if not pass_name:
+            return Response(id=request_id, error="pass_name is required")
+
+        known_passes = self._get_known_pass_names()
+        if known_passes and pass_name not in known_passes:
+            return Response(
+                id=request_id,
+                error=f"Unknown pass '{pass_name}'. Known passes: {sorted(known_passes)}",
+            )
+
+        if self.reducer is not None and hasattr(self.reducer, "disable_pass"):
+            self.reducer.disable_pass(pass_name)
+            return Response(id=request_id, result={"status": "disabled", "pass_name": pass_name})
+        return Response(id=request_id, error="Reducer does not support pass control")
+
+    def _handle_enable_pass(self, request_id: str, params: dict) -> Response:
+        """Enable a previously disabled reduction pass."""
+        pass_name = params.get("pass_name", "")
+        if not pass_name:
+            return Response(id=request_id, error="pass_name is required")
+
+        known_passes = self._get_known_pass_names()
+        if known_passes and pass_name not in known_passes:
+            return Response(
+                id=request_id,
+                error=f"Unknown pass '{pass_name}'. Known passes: {sorted(known_passes)}",
+            )
+
+        if self.reducer is not None and hasattr(self.reducer, "enable_pass"):
+            self.reducer.enable_pass(pass_name)
+            return Response(id=request_id, result={"status": "enabled", "pass_name": pass_name})
+        return Response(id=request_id, error="Reducer does not support pass control")
+
+    def _handle_skip_pass(self, request_id: str) -> Response:
+        """Skip the currently running pass."""
+        if self.reducer is not None and hasattr(self.reducer, "skip_current_pass"):
+            self.reducer.skip_current_pass()
+            return Response(id=request_id, result={"status": "skipped"})
+        return Response(id=request_id, error="Reducer does not support pass control")
+
+    def _get_content_preview(self) -> tuple[str, bool]:
+        """Get a preview of the current test case content."""
+        if self.problem is None:
+            return "", False
+
+        test_case = self.problem.current_test_case
+
+        # Handle directory mode
+        if isinstance(test_case, dict):
+            lines = []
+            for name, content in sorted(test_case.items()):
+                size = len(content)
+                lines.append(f"{name}: {size} bytes")
+            return "\n".join(lines[:50]), False
+
+        # Handle single file mode
+        hex_mode = is_binary_string(
+            test_case[:1024] if len(test_case) > 1024 else test_case
+        )
+
+        if hex_mode:
+            # Show hex dump for binary files
+            preview_bytes = test_case[:512]
+            lines = []
+            for i in range(0, len(preview_bytes), 16):
+                chunk = preview_bytes[i : i + 16]
+                hex_part = " ".join(f"{b:02x}" for b in chunk)
+                ascii_part = "".join(chr(b) if 32 <= b < 127 else "." for b in chunk)
+                lines.append(f"{i:08x}  {hex_part:<48}  {ascii_part}")
+            return "\n".join(lines), True
+        else:
+            # Show text for text files
+            # Send up to 100KB of content - the TUI will handle display truncation
+            try:
+                text = test_case.decode("utf-8", errors="replace")
+                # Truncate by character count to handle files with very long lines
+                max_chars = 100_000
+                if len(text) > max_chars:
+                    text = text[:max_chars]
+                return text, False
+            except Exception:
+                return "", True
+
+    async def _build_progress_update(self) -> ProgressUpdate | None:
+        """Build a progress update from current state."""
+        if self.problem is None:
+            return None
+
+        stats = self.problem.stats
+        content_preview, hex_mode = self._get_content_preview()
+
+        # Get parallel workers count and track average
+        parallel_workers = 0
+        if self.state is not None and hasattr(self.state, "parallel_tasks_running"):
+            parallel_workers = self.state.parallel_tasks_running
+            self._parallel_samples += 1
+            self._parallel_total += parallel_workers
+
+        # Calculate parallelism stats
+        average_parallelism = 0.0
+        effective_parallelism = 0.0
+        if self._parallel_samples > 0:
+            average_parallelism = self._parallel_total / self._parallel_samples
+            wasteage = (
+                stats.wasted_interesting_calls / stats.calls if stats.calls > 0 else 0.0
+            )
+            effective_parallelism = average_parallelism * (1.0 - wasteage)
+
+        # Collect pass statistics in run order (only those with test evaluations)
+        pass_stats_list = []
+        current_pass_name = ""
+        if self.reducer is not None:
+            # Get the currently running pass name
+            current_pass = self.reducer.current_reduction_pass
+            if current_pass is not None:
+                current_pass_name = getattr(current_pass, "__name__", "")
+
+            # Get all stats in the order they were first run
+            pass_stats = self.reducer.pass_stats
+            if pass_stats is not None:
+                all_stats = pass_stats.get_stats_in_order()
+
+                # Only include passes that have made at least one test evaluation
+                pass_stats_list = [
+                    PassStatsData(
+                        pass_name=ps.pass_name,
+                        bytes_deleted=ps.bytes_deleted,
+                        run_count=ps.run_count,
+                        test_evaluations=ps.test_evaluations,
+                        successful_reductions=ps.successful_reductions,
+                        success_rate=ps.success_rate,
+                    )
+                    for ps in all_stats
+                    if ps.test_evaluations > 0
+                ]
+
+        # Get disabled passes
+        if self.reducer is not None and hasattr(self.reducer, "disabled_passes"):
+            disabled_passes = list(self.reducer.disabled_passes)
+        else:
+            disabled_passes = []
+
+        return ProgressUpdate(
+            status=self.reducer.status if self.reducer else "",
+            size=stats.current_test_case_size,
+            original_size=stats.initial_test_case_size,
+            calls=stats.calls,
+            reductions=stats.reductions,
+            interesting_calls=stats.interesting_calls,
+            wasted_calls=stats.wasted_interesting_calls,
+            runtime=time.time() - stats.start_time,
+            parallel_workers=parallel_workers,
+            average_parallelism=average_parallelism,
+            effective_parallelism=effective_parallelism,
+            time_since_last_reduction=stats.time_since_last_reduction(),
+            content_preview=content_preview,
+            hex_mode=hex_mode,
+            pass_stats=pass_stats_list,
+            current_pass_name=current_pass_name,
+            disabled_passes=disabled_passes,
+        )
+
+    async def emit_progress_updates(self) -> None:
+        """Periodically emit progress updates."""
+        # Emit initial progress update immediately
+        update = await self._build_progress_update()
+        if update is not None:
+            await self.emit(update)
+
+        while self.running:
+            await trio.sleep(0.1)
+            update = await self._build_progress_update()
+            if update is not None:
+                await self.emit(update)
+
+    async def run_reducer(self) -> None:
+        """Run the reducer."""
+        if self.reducer is None:
+            return
+
+        try:
+            with trio.CancelScope() as scope:
+                self._cancel_scope = scope
+                await self.reducer.run()
+
+            # Check for trivial result after successful completion
+            if self.state is not None and self.problem is not None:
+                trivial_error = self.state.check_trivial_result(self.problem)
+                if trivial_error:
+                    await self.emit(Response(id="", error=trivial_error))
+        except* InvalidInitialExample as excs:
+            assert len(excs.exceptions) == 1
+            (e,) = excs.exceptions
+            # Build a detailed error message for invalid initial examples
+            if self.state is not None:
+                error_message = await self.state.build_error_message(e)
+            else:
+                error_message = str(e)
+            await self.emit(Response(id="", error=error_message))
+        except* Exception as e:
+            # Catch any other exception during reduction and emit as error
+            traceback.print_exc()
+            await self.emit(Response(id="", error=str(e.exceptions[0])))
+        finally:
+            self._cancel_scope = None
+            self.running = False
+
+    async def run(self) -> None:
+        """Main entry point for the worker."""
+        async with trio.open_nursery() as nursery:
+            await nursery.start(self.read_commands)
+
+            # Wait for start command
+            while not self.running:
+                await trio.sleep(0.01)
+
+            # Start progress updates and reducer
+            nursery.start_soon(self.emit_progress_updates)
+            await self.run_reducer()
+
+            # Emit final progress update before completion
+            final_update = await self._build_progress_update()
+            if final_update is not None:
+                await self.emit(final_update)
+
+            # Signal completion
+            await self.emit(Response(id="", result={"status": "completed"}))
+            nursery.cancel_scope.cancel()
+
+
+def main() -> None:
+    """Entry point for the worker subprocess."""
+    worker = ReducerWorker()
+    trio.run(worker.run)
+
+
+if __name__ == "__main__":
+    main()