diff --git a/llvm/utils/lit/lit/ShellEnvironment.py b/llvm/utils/lit/lit/ShellEnvironment.py index 1945865b19199..e2cf602de96ff 100644 --- a/llvm/utils/lit/lit/ShellEnvironment.py +++ b/llvm/utils/lit/lit/ShellEnvironment.py @@ -1,7 +1,11 @@ +from __future__ import annotations + +import io import os import platform import subprocess import tempfile +from typing import BinaryIO, Protocol, TextIO import lit.util from lit.ShCommands import GlobItem @@ -206,6 +210,67 @@ def processRedirects(cmd, stdin_source, cmd_shenv, opened_files): return std_fds +def as_binary_reader(stream: None | int | io.TextIOBase | BinaryIO) -> BinaryIO: + """Adapts a builtin's stdin source into a binary, read()-able stream. + + Args: + stream: Standard input source from pipeline dispatch. Supported types: + - None or subprocess sentinel (PIPE, DEVNULL, STDOUT) for no input + - Binary stream (BytesIO, temporary/spooled file, or 'rb' mode) + - Text stream (text-mode '<' redirect or upstream universal_newlines pipe) + + Returns: + A binary, read()-able stream. Text streams are unwrapped to their + underlying buffer so an in-process builtin sees the same raw bytes a + child process would. + """ + if stream is None or isinstance(stream, int): + # No real input to read. + return io.BytesIO(b"") + if isinstance(stream, io.TextIOBase): + buffer = getattr(stream, "buffer", None) + if buffer is not None: + return buffer + data = stream.read() + return io.BytesIO(data.encode() if isinstance(data, str) else data) + # Already a binary reader. + assert hasattr(stream, "read"), f"expected a binary reader, got {type(stream)!r}" + return stream + + +class BinaryFileWriter: + """Writes bytes straight to a file descriptor, matching Popen's behavior. + + Writes directly to the underlying file descriptor with os.write, bypassing + the wrapped file object's text-mode buffering and newline translation, so + the exact bytes passed in reach the fd unchanged. + """ + + # TODO: Replace __slots__ with @dataclass(slots=True) + # once the minimum Python version is bumped to 3.10 + # https://github.com/llvm/llvm-project/issues/200531 + __slots__ = ("fd",) + + def __init__(self, fileobj: TextIO) -> None: + # Owned by opened_files, never closed here + self.fd = fileobj.fileno() + + def write(self, data: bytes) -> int: + return os.write(self.fd, data) + + +def binary_fd(fileobj: TextIO) -> BinaryFileWriter: + """Wraps a redirect file object as a byte-exact writer for in-process builtins.""" + return BinaryFileWriter(fileobj) + + +class ByteWriter(Protocol): + """Structural type for objects that support writing a bytes buffer.""" + + def write(self, data: bytes) -> object: + ... + + def expand_glob(arg, cwd): if isinstance(arg, GlobItem): return sorted(arg.resolve(cwd)) diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py index 97cb54c63be64..681a56d9b45db 100644 --- a/llvm/utils/lit/lit/TestRunner.py +++ b/llvm/utils/lit/lit/TestRunner.py @@ -1,6 +1,7 @@ from __future__ import annotations import enum +import io import os import pathlib import re @@ -11,17 +12,23 @@ import tempfile import threading import traceback +from typing import IO, BinaryIO, Callable, List, TextIO, Union import lit.InprocBuiltins as InprocBuiltins import lit.ShUtil as ShUtil import lit.Test as Test import lit.util +import lit.builtin_commands.cat as builtin_cat +import lit.builtin_commands.diff as builtin_diff from lit.BooleanExpression import BooleanExpression from lit.ShCommands import Command from lit.ShellEnvironment import ( + ByteWriter, InternalShellError, ShellCommandResult, ShellEnvironment, + as_binary_reader, + binary_fd, expand_glob, expand_glob_expressions, kAvoidDevNull, @@ -212,6 +219,299 @@ def _replaceReadFile(match): return arguments +class PipeIOConfig: + """Configuration for how one pipeline stage's output is routed, whether + to the next stage, a redirect file, or captured for the caller, and the + working directory the stage runs in. + + Attributes: + out_sink: Where a non-last stage's output goes if the next stage + reads it. None for the last stage or a real redirect. + redirect_out: Binary writer for a '>'/'>>' target, else None. + capture_out: True if this stage's output should be captured and + returned to the caller instead of piped or redirected. + err_sink: Binary writer for a '2>file' target, else None. + cwd: The shell environment's current working directory. + merge_err: True for '2>&1', stderr writes to the out target instead. + """ + + # TODO: Replace __slots__ with @dataclass(slots=True) + # once the minimum Python version is bumped to 3.10 + # https://github.com/llvm/llvm-project/issues/200531 + __slots__ = ( + "out_sink", + "redirect_out", + "capture_out", + "err_sink", + "cwd", + "merge_err", + ) + + def __init__( + self, + out_sink: IO[bytes] | None, + redirect_out: ByteWriter | None, + capture_out: bool, + err_sink: ByteWriter | None, + cwd: str, + merge_err: bool = False, + ) -> None: + self.out_sink = out_sink + self.redirect_out = redirect_out + self.capture_out = capture_out + self.err_sink = err_sink + self.cwd = cwd + self.merge_err = merge_err + + +# A builtin's run(argv, stdin, stdout, stderr, cwd) entry point, e.g. +# builtin_cat.run or builtin_diff.run. stdout/stderr accept anything with a +# byte-oriented write() method: a real BytesIO or the duck-typed ByteWriter. +RunFn = Callable[ + [ + List[str], + BinaryIO, + Union[IO[bytes], ByteWriter, None], + Union[IO[bytes], ByteWriter, None], + str, + ], + int, +] + + +class InProcessPipe: + """Popen-compatible shim for an I/O-heavy builtin command. + + Runs the builtin (e.g., cat or diff) in-process instead of spawning + it. Anything that isn't a lit builtin still spawns a real subprocess. + + Exposes the same interface as subprocess.Popen, so it's a drop-in + replacement wherever a real subprocess is expected. + + Attributes: + returncode: The builtin's exit code. + stdout: The captured stdout if this stage's output was captured, + else None. + stderr: The captured stderr. + """ + + # TODO: Replace __slots__ with @dataclass(slots=True) + # once the minimum Python version is bumped to 3.10 + # https://github.com/llvm/llvm-project/issues/200531 + __slots__ = ("returncode", "stdout", "stderr", "_out", "_err") + + def __init__( + self, + run_fn: RunFn, + args: List[str], + stdin: None | int | io.TextIOBase | BinaryIO, + stage_io: PipeIOConfig, + ) -> None: + """Runs run_fn synchronously and captures the result. + + Args: + run_fn: The builtin's run function, e.g. cat.run or diff.run. + args: The argv to pass to run_fn, args[0] is the command name. + stdin: This stage's input stream. + stage_io: Where this stage's output and error streams go, and + its cwd. + """ + in_stream = as_binary_reader(stdin) + out_buf = io.BytesIO() if stage_io.capture_out else None + out_target = ( + out_buf + if stage_io.capture_out + else (stage_io.redirect_out or stage_io.out_sink) + ) + if stage_io.merge_err: + err_target, err_buf = out_target, None + elif stage_io.err_sink is not None: + err_target, err_buf = stage_io.err_sink, None + else: + err_buf = io.BytesIO() + err_target = err_buf + + self.returncode = run_fn(args, in_stream, out_target, err_target, stage_io.cwd) + + self._out = out_buf.getvalue() if out_buf is not None else b"" + self._err = err_buf.getvalue() if err_buf is not None else b"" + self.stdout = io.BytesIO(self._out) if stage_io.capture_out else None + self.stderr = io.BytesIO(self._err) + + def communicate(self) -> tuple[bytes, bytes]: + return (self._out, self._err) + + def wait(self) -> int: + return self.returncode + + def poll(self) -> int: + return self.returncode + + def kill(self) -> None: + pass + + def terminate(self) -> None: + pass + + +def _resolve_redirect_out(stdout: int | TextIO) -> ByteWriter | None: + """Wraps a '>' or '>>' redirect target as a binary writer. + + Args: + stdout: This stage's stdout target, a subprocess sentinel + (PIPE/STDOUT) or an open file object for a real redirect. + + Returns: + A binary writer wrapping the redirect file. None if stdout is a + PIPE or STDOUT sentinel instead of a real redirect file, since + that case doesn't need this sink. + """ + if isinstance(stdout, int): + return None + return binary_fd(stdout) + + +def _should_capture(is_last: bool, stdout: int | TextIO) -> bool: + """Whether this stage's output should be captured for the caller. + + Args: + is_last: True for the pipeline's final stage, the only stage allowed + to capture instead of feeding the next stage or a redirect. + stdout: The stage's stdout target. Capture only applies when + this is subprocess.PIPE, ruling out a real redirect file. + + Returns: + True only for the last stage of a piped command reading from a + PIPE, the one stage that reports output back to the caller + instead of writing to a temp file or redirect. + """ + return is_last and stdout == subprocess.PIPE + + +def _make_out_sink(stdout: int | TextIO, is_last: bool) -> IO[bytes] | None: + """Builds the output sink for a non-last stage feeding a pipe. + + Args: + stdout: The stage's stdout target. Only subprocess.PIPE gets a + sink here, a real redirect target writes there directly. + is_last: True for the pipeline's final stage, which streams to + the caller instead of spooling for a downstream stage. + + Returns: + A SpooledTemporaryFile, which stays in memory unless its size or + a downstream call to .fileno() forces it to disk. None for the + last stage or a non-pipe target, which don't need this sink. + """ + if stdout == subprocess.PIPE and not is_last: + return tempfile.SpooledTemporaryFile(max_size=1 << 20) + return None + + +def _resolve_err_sink(stderr: int | TextIO, merge_err: bool) -> ByteWriter | None: + """Wraps a '2>file' redirect target as a binary writer. + + Args: + stderr: This stage's stderr target. Only a real redirect file + needs wrapping here, PIPE/STDOUT are handled elsewhere. + merge_err: Whether stderr is merged into stdout (2>&1). + + Returns: + A binary writer wrapping the redirect file. None when stderr is + merged into stdout or is a PIPE/STDOUT sentinel, since both cases + are captured into the result instead of written through a sink + here. + """ + if merge_err or isinstance(stderr, int): + return None + return binary_fd(stderr) + + +def _should_run_inproc( + builtin_fn: RunFn | None, + not_crash: bool, + cmd_shenv: ShellEnvironment, + shenv: ShellEnvironment, +) -> bool: + """Whether this pipeline stage can run in-process instead of spawning. + + Args: + builtin_fn: The builtin's run function for this command (e.g. + builtin_cat.run), or None if the command has no in-process + implementation. + not_crash: True if the command is wrapped in 'not --crash'. + cmd_shenv: The environment this specific command runs under. + shenv: The pipeline's shared environment. + + Returns: + False if builtin_fn is None, if not_crash is set, or if cmd_shenv + is not shenv (a per-command 'env' only takes effect on a spawned + child, not an in-process call sharing our environment). True + otherwise. + """ + return builtin_fn is not None and not not_crash and cmd_shenv is shenv + + +def _run_inproc_stage( + args: List[str], + builtin_fn: RunFn, + stdin: None | int | io.TextIOBase | BinaryIO, + stdout: int | TextIO, + stderr: int | TextIO, + cmd_shenv: ShellEnvironment, + is_last: bool, + named_temp_files: List[str], +) -> tuple[InProcessPipe, IO[bytes] | int]: + """Runs one in-process pipeline stage and returns its default_stdin. + + Args: + args: The command's argument vector. + builtin_fn: The builtin's run function to execute (e.g., + cat.run or diff.run). + stdin: Current stage's input. None or a placeholder int means no + real input, otherwise an open binary or text stream to read + from. + stdout: Target for the stage's stdout, either subprocess.PIPE or + an open redirect file. Whether PIPE feeds the next stage or + is captured depends on is_last. + stderr: Target for the stage's stderr. Set to subprocess.STDOUT + when merged into stdout. + cmd_shenv: Shell environment the stage runs under. + is_last: True for the final stage in the pipeline. Together with + stdout being subprocess.PIPE, this determines whether output + is captured instead of feeding a downstream stage. + named_temp_files: Caller-owned cleanup list. Any temp file created + to stand in for /dev/null is appended here. + + Returns: + The InProcessPipe wrapping the stage, and the stdin the next stage + should read from (the spooled output, or subprocess.PIPE). + """ + if kAvoidDevNull: + for arg_idx, arg in enumerate(args): + if isinstance(arg, str) and kDevNull in arg: + devnull = tempfile.NamedTemporaryFile(delete=False) + devnull.close() + named_temp_files.append(devnull.name) + args[arg_idx] = arg.replace(kDevNull, devnull.name) + args = expand_glob_expressions(args, cmd_shenv.cwd) + merge_err = stderr == subprocess.STDOUT + stage_io = PipeIOConfig( + out_sink=_make_out_sink(stdout, is_last), + redirect_out=_resolve_redirect_out(stdout), + capture_out=_should_capture(is_last, stdout), + err_sink=_resolve_err_sink(stderr, merge_err), + cwd=cmd_shenv.cwd, + merge_err=merge_err, + ) + proc = InProcessPipe(builtin_fn, args, stdin, stage_io) + if stage_io.out_sink is not None: + stage_io.out_sink.seek(0) + default_stdin = stage_io.out_sink + else: + default_stdin = subprocess.PIPE + return proc, default_stdin + + def _executeShCmd(cmd, shenv, results, timeoutHelper): if timeoutHelper.timeoutReached(): # Prevent further recursion if the timeout has been hit @@ -268,6 +568,10 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper): "umask": InprocBuiltins.executeBuiltinUmask, ":": InprocBuiltins.executeBuiltinColon, } + pipeline_builtins = { + "cat": builtin_cat.run, + "diff": builtin_diff.run, + } # To avoid deadlock, we use a single stderr stream for piped # output. This is null until we have seen some output using # stderr. @@ -362,13 +666,14 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper): results.append(result) return result.exitCode - # Resolve any out-of-process builtin command before adding back 'not' - # commands. - if args[0] in builtin_commands: + builtin_fn = pipeline_builtins.get(args[0]) + use_inproc = _should_run_inproc(builtin_fn, not_crash, cmd_shenv, shenv) + if not use_inproc and args[0] in builtin_commands: args.insert(0, sys.executable) cmd_shenv.env["PYTHONPATH"] = os.path.dirname(os.path.abspath(__file__)) args[1] = os.path.join(builtin_commands_dir, args[1] + ".py") + # We had to search through the 'not' commands to find all the 'env' # commands and any other in-process builtin command. We don't want to # reimplement 'not' and its '--crash' here, so just push all 'not' @@ -395,6 +700,21 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper): j, default_stdin, cmd_shenv, opened_files ) + if use_inproc: + proc, default_stdin = _run_inproc_stage( + args, + builtin_fn, + stdin, + stdout, + stderr, + cmd_shenv, + is_last=j is cmd.commands[-1], + named_temp_files=named_temp_files, + ) + procs.append(proc) + proc_not_counts.append(not_count) + proc_not_fail_if_crash.append(False) + continue # If stderr wants to come from stdout, but stdout isn't a pipe, then put # stderr on a pipe and treat it as stdout. if stderr == subprocess.STDOUT and stdout != subprocess.PIPE: diff --git a/llvm/utils/lit/lit/builtin_commands/diff.py b/llvm/utils/lit/lit/builtin_commands/diff.py index 9d1a398106664..3a85d2210705f 100644 --- a/llvm/utils/lit/lit/builtin_commands/diff.py +++ b/llvm/utils/lit/lit/builtin_commands/diff.py @@ -92,11 +92,7 @@ def compareTwoBinaryFiles(flags, filepaths, filelines, stdout): ) for diff in diffs: - stdout.write( - diff.decode(errors="backslashreplace").encode( - locale.getpreferredencoding(False) - ) - ) + stdout.write(diff.decode(errors="backslashreplace").encode("utf-8")) exitCode = 1 return exitCode @@ -153,9 +149,7 @@ def printDirVsFile(dir_path, file_path, stdout): msg = "File %s is a directory while file %s is a regular file" else: msg = "File %s is a directory while file %s is a regular empty file" - stdout.write( - (msg % (dir_path, file_path) + "\n").encode(locale.getpreferredencoding(False)) - ) + stdout.write((msg % (dir_path, file_path) + "\n").encode("utf-8")) def printFileVsDir(file_path, dir_path, stdout): @@ -163,16 +157,12 @@ def printFileVsDir(file_path, dir_path, stdout): msg = "File %s is a regular file while file %s is a directory" else: msg = "File %s is a regular empty file while file %s is a directory" - stdout.write( - (msg % (file_path, dir_path) + "\n").encode(locale.getpreferredencoding(False)) - ) + stdout.write((msg % (file_path, dir_path) + "\n").encode("utf-8")) def printOnlyIn(basedir, path, name, stdout): stdout.write( - ("Only in %s: %s\n" % (os.path.join(basedir, path), name)).encode( - locale.getpreferredencoding(False) - ) + ("Only in %s: %s\n" % (os.path.join(basedir, path), name)).encode("utf-8") )