mirror of
https://github.com/revng/revng
synced 2026-06-21 14:07:57 +00:00
Giacomo Vercesi
49b2c523ab
Add the `revng merge llvm` command which allows linking multiple LLVM IR modules together into a single module.
97 lines
3.6 KiB
Python
97 lines
3.6 KiB
Python
#
|
|
# This file is distributed under the MIT License. See LICENSE.md for details.
|
|
#
|
|
|
|
import tarfile
|
|
from argparse import ArgumentParser, RawDescriptionHelpFormatter
|
|
from base64 import b64decode
|
|
from contextlib import suppress
|
|
from pathlib import Path
|
|
from shutil import copyfileobj
|
|
from tempfile import TemporaryDirectory, TemporaryFile
|
|
from typing import IO
|
|
|
|
import yaml
|
|
|
|
from revng.internal.cli.commands_registry import Command, CommandsRegistry, Options
|
|
from revng.internal.cli.support import build_command_with_loads, file_wrapper, try_run
|
|
from revng.support import TarDictionary
|
|
|
|
|
|
def read_single_file(input_: IO[bytes]) -> list[bytes]:
|
|
# If we encounter a non-seekable file (e.g. stdin) first copy it to a
|
|
# temporary file and then re-run this function
|
|
if not input_.seekable():
|
|
with TemporaryFile() as temp_file:
|
|
copyfileobj(input_, temp_file)
|
|
temp_file.seek(0)
|
|
return read_single_file(temp_file)
|
|
|
|
# If here, the file is seekable, test the various formats by trying to read
|
|
# them and moving onto the next format if the previous throws an error
|
|
|
|
# Test for YAML
|
|
with suppress(yaml.YAMLError):
|
|
data = yaml.load(input_, Loader=yaml.CSafeLoader)
|
|
return [b64decode(v) for v in data.values()]
|
|
|
|
# Test for tar
|
|
input_.seek(0)
|
|
with suppress(tarfile.ReadError):
|
|
return list(TarDictionary(input_).values())
|
|
|
|
# No good format found, assume it's a single plain file
|
|
input_.seek(0)
|
|
return [input_.read()]
|
|
|
|
|
|
class MergeLLVMIRCommand(Command):
|
|
def __init__(self):
|
|
super().__init__(("merge", "llvm"), "Merge multiple LLVM IR files into a single one")
|
|
|
|
def register_arguments(self, parser: ArgumentParser):
|
|
parser.formatter_class = RawDescriptionHelpFormatter
|
|
parser.description = """Links multiple LLVM IR files into a single one.
|
|
This program works differently depending on the number of inputs:
|
|
* If one input is passed it is interpreted as either:
|
|
* a YAML dict with base64-encoded values
|
|
* a tar where each file is a bitcode file
|
|
* a single LLVM module
|
|
* If no input is passed, treat stdin as a single input
|
|
* If multiple input files are passed, treat all of them as bitcode files
|
|
"""
|
|
parser.add_argument("-o", "--output", help="Where to output the file, stdout if omitted")
|
|
parser.add_argument("inputs", nargs="*", help="Input file(s), omit to use stdin")
|
|
|
|
def run(self, options: Options):
|
|
args = options.parsed_args
|
|
if len(args.inputs) >= 2:
|
|
return self.run_merge(args.inputs, options)
|
|
|
|
input_file = args.inputs[0] if len(args.inputs) == 1 else "-"
|
|
with file_wrapper(input_file, "rb") as input_:
|
|
data = read_single_file(input_)
|
|
|
|
with TemporaryDirectory(prefix="tmp.revng-merge-llvm.") as temp_dir:
|
|
temp_dir_path = Path(temp_dir)
|
|
input_paths = []
|
|
for index, value in enumerate(data):
|
|
element_path = temp_dir_path / str(index)
|
|
element_path.write_bytes(value)
|
|
input_paths.append(str(element_path.resolve()))
|
|
|
|
return self.run_merge(input_paths, options)
|
|
|
|
def run_merge(self, inputs: list[str], options: Options):
|
|
with file_wrapper(options.parsed_args.output, "wb") as output_file:
|
|
return try_run(
|
|
build_command_with_loads("merge-llvm", inputs, options),
|
|
options,
|
|
stdout=output_file,
|
|
)
|
|
|
|
|
|
def setup(commands_registry: CommandsRegistry):
|
|
commands_registry.define_namespace(("merge",), "Tools to merge function-wise object together")
|
|
commands_registry.register_command(MergeLLVMIRCommand())
|