Xdsl opt main

class xDSLOptMain(CommandLineTool):
    available_passes: dict[str, Callable[[], type[ModulePass]]]
    """
    A mapping from pass names to functions that apply the pass to a ModuleOp.
    """

    available_targets: dict[
        str, Callable[[ModuleOp, IO[str]], None] | Callable[[], type[Target]]
    ]
    """
    A mapping from target names to either:
    - Old-style: a function `(ModuleOp, IO[str]) -> None` (deprecated soon)
    - New-style: a factory `() -> type[Target]`
    """

    pipeline: PassPipeline
    """ The pass-pipeline to be applied. """

    version_format: ClassVar[str] = "xdsl-opt built from xdsl version {xdsl_version}"
    """
    Format string for the message printed by `--version`. `{xdsl_version}` is
    replaced by the installed xDSL version.
    """

    def __init__(
        self,
        description: str = "xDSL modular optimizer driver",
        args: Sequence[str] | None = None,
    ):
        self.available_frontends = {}
        self.available_passes = {}
        self.available_targets = {}

        self.ctx = Context()
        self.register_all_dialects()
        self.register_all_frontends()
        self.register_all_passes()
        self.register_all_targets()

        # arg handling
        arg_parser = argparse.ArgumentParser(description=description)
        self.register_all_arguments(arg_parser)
        self.args = arg_parser.parse_args(args=args)

        target_spec = parse_spec(self.args.target)
        if target_spec.name not in self.available_targets:
            arg_parser.error(
                f"argument -t/--target: invalid choice: '{target_spec.name}' "
                f"(choose from {', '.join(self.available_targets)})"
            )

        # Validate frontend spec
        if self.args.frontend:
            frontend_spec = parse_spec(self.args.frontend)
            if frontend_spec.name not in self.available_frontends:
                arg_parser.error(
                    f"argument -f/--frontend: invalid choice: '{frontend_spec.name}' "
                    f"(choose from {', '.join(self.available_frontends)})"
                )

        self.ctx.allow_unregistered = self.args.allow_unregistered_dialect

        if self.args.disable_verify:
            Attribute.__post_init__ = _empty_post_init

        if self.args.syntax_highlight:
            Diagnostic.colored = True

        self.setup_pipeline()

    def run(self):
        """
        Executes the different steps.
        """
        chunks, file_extension = self.prepare_input()
        output_stream = self.prepare_output()
        try:
            for i, (chunk, offset) in enumerate(chunks):
                try:
                    if i > 0:
                        output_stream.write("// -----\n")
                    module = self.parse_chunk(chunk, file_extension, offset)

                    if module is not None:
                        self.apply_passes(module)
                        output_stream.write(self.output_resulting_program(module))
                    output_stream.flush()
                except ParseError as e:
                    e.span = dataclasses.replace(e.span, line_offset=offset)
                    if self.args.parsing_diagnostics:
                        print(e)
                    else:
                        raise
                except DiagnosticException as e:
                    if self.args.verify_diagnostics:
                        print(e)
                        # __notes__ only in Python 3.11 and above
                        if hasattr(e, "__notes__"):
                            for e in getattr(e, "__notes__"):
                                print(e)
                    else:
                        raise
                finally:
                    chunk.close()
        except ShrinkException:
            assert self.args.shrink
            print("Success, can shrink")
            # Exit with value 0 to let shrinkray know that it can shrink
            exit(0)
        finally:
            if output_stream is not sys.stdout:
                output_stream.close()
        if self.args.shrink:
            print("Failure, can't shrink")
            # Exit with non-0 value to let shrinkray know that it cannot shrink
            exit(1)

    def register_all_arguments(self, arg_parser: argparse.ArgumentParser):
        """
        Registers all the command line arguments that are used by this tool.

        Add other/additional arguments by overloading this function.
        """
        super().register_all_arguments(arg_parser)

        target_names = ",".join(self.available_targets)
        arg_parser.add_argument(
            "-t",
            "--target",
            type=str,
            required=False,
            help=f"Target to use for output. Available targets are: {target_names}",
            default="mlir",
        )

        arg_parser.add_argument(
            "-o", "--output-file", type=str, required=False, help="path to output file"
        )

        pass_names = ",".join([name for name in self.available_passes])
        arg_parser.add_argument(
            "-p",
            "--passes",
            required=False,
            help=f"Delimited list of passes. Available passes are: {pass_names}",
            type=str,
            default="",
        )

        arg_parser.add_argument(
            "--print-between-passes",
            default=False,
            action="store_true",
            help="Print the IR between each pass",
        )

        arg_parser.add_argument(
            "--time-passes",
            default=False,
            action="store_true",
            help="Print timing information for each pass",
        )

        arg_parser.add_argument(
            "--verify-diagnostics",
            default=False,
            action="store_true",
            help="Prints the content of a triggered "
            "verifier exception and exits with code 0",
        )

        arg_parser.add_argument(
            "--parsing-diagnostics",
            default=False,
            action="store_true",
            help="Prints the content of a triggered "
            "parsing exception and exits with code 0",
        )

        arg_parser.add_argument(
            "--split-input-file",
            default=False,
            action="store_true",
            help="Split the input file into pieces and process each chunk "
            "independently by using `// -----`",
        )

        arg_parser.add_argument(
            "--print-op-generic",
            default=False,
            action="store_true",
            help="Print operations with the generic format",
        )

        arg_parser.add_argument(
            "--print-no-properties",
            default=False,
            action="store_true",
            help="Print properties as if they were attributes for retrocompatibility.",
        )

        arg_parser.add_argument(
            "--print-debuginfo",
            default=False,
            action="store_true",
            help="Print operations with debug info annotation, such as location.",
        )

        arg_parser.add_argument(
            "--syntax-highlight",
            default=False,
            action="store_true",
            help="Enable printing with syntax highlighting on the terminal.",
        )

        arg_parser.add_argument(
            "-v",
            "--version",
            action=VersionAction,
            version_format=self.version_format,
        )

        arg_parser.add_argument(
            "--shrink",
            default=False,
            action="store_true",
            help="Return success on exit if ShrinkException was raised.",
        )

    def register_pass(
        self, pass_name: str, pass_factory: Callable[[], type[ModulePass]]
    ):
        self.available_passes[pass_name] = pass_factory

    def register_all_passes(self):
        """
        Register all passes that can be used.

        Add other/additional passes by overloading this function.
        """
        multiverse = Universe.get_multiverse()
        for pass_name, pass_factory in multiverse.all_passes.items():
            self.register_pass(pass_name, pass_factory)

    def register_all_targets(self):
        """
        Register all targets that can be used.

        Add other/additional targets by overloading this function.
        """
        multiverse = Universe.get_multiverse()
        for target_name, target_factory in multiverse.all_targets.items():
            if target_name not in self.available_targets:
                self.available_targets[target_name] = target_factory

    def setup_pipeline(self):
        """
        Creates a pipeline that consists of all the passes specified.

        Fails, if not all passes are registered.
        """

        start_timer = time.perf_counter()

        def callback(
            previous_pass: ModulePass | None,
            module: ModuleOp,
            next_pass: ModulePass | None,
        ) -> None:
            nonlocal start_timer
            if self.args.time_passes and previous_pass is not None:
                end_timer = time.perf_counter()
                duration = end_timer - start_timer
                print(f"Pass {previous_pass.name} took {duration} seconds")

            if not self.args.disable_verify:
                module.verify()
            if previous_pass and next_pass and self.args.print_between_passes:
                print(f"IR after {previous_pass.name}:")
                printer = Printer(stream=sys.stdout)
                printer.print_op(module)
                print("\n\n\n")

            if self.args.time_passes and next_pass is not None:
                start_timer = time.perf_counter()

        self.pipeline = PassPipeline.parse_spec(
            self.available_passes,
            self.args.passes,
            callback,
        )

    def prepare_input(self) -> tuple[list[tuple[IO[bytes], int]], str]:
        """
        Prepare input by eventually splitting it in chunks. If not set, the parser
        registered for this file extension is used.
        """

        # when using the split input flag, program is split into multiple chunks
        # it's used for split input file

        chunks: list[tuple[IO[bytes], int]] = []
        f, file_extension = self.get_input_stream()
        chunks = [(f, 0)]
        if self.args.split_input_file:
            chunks_bytes = f.read().split(b"// -----")
            chunks_off = accumulate(
                [0, *[chunk.count(b"\n") for chunk in chunks_bytes[:-1]]]
            )
            chunks = [
                (BytesIO(chunk), off)
                for chunk, off in zip(chunks_bytes, chunks_off, strict=True)
            ]
            f.close()

        if file_extension not in self.available_frontends:
            for chunk, _ in chunks:
                chunk.close()
            raise ValueError(f"Unrecognized file extension '{file_extension}'")

        return chunks, file_extension

    def prepare_output(self) -> IO[str]:
        if self.args.output_file is None:
            return sys.stdout
        else:
            return open(self.args.output_file, "w")

    def apply_passes(self, prog: ModuleOp):
        """Apply passes in order."""
        self.pipeline.apply(self.ctx, prog)

    def output_resulting_program(self, prog: ModuleOp) -> str:
        """Get the resulting program."""
        output = StringIO()

        spec = parse_spec(self.args.target)
        if spec.name not in self.available_targets:
            raise ValueError(
                f"Unknown target '{spec.name}'. "
                f"Available targets: {list(self.available_targets)}"
            )

        target_entry = self.available_targets[spec.name]
        sig = inspect.signature(target_entry)

        if not sig.parameters:
            factory = cast(Callable[[], type[Target]], target_entry)
            target = factory().from_spec(spec)
            if spec.name == "mlir":
                target = dataclasses.replace(
                    target,
                    **{
                        "print_generic_format": self.args.print_op_generic,
                        "print_properties_as_attributes": self.args.print_no_properties,
                        "print_debuginfo": self.args.print_debuginfo,
                        "syntax_highlight": self.args.syntax_highlight,
                    },
                )
            target.emit(self.ctx, prog, output)
        else:
            legacy = cast(Callable[[ModuleOp, IO[str]], None], target_entry)
            legacy(prog, output)

        return output.getvalue()

class VersionAction(argparse.Action):
    version_format: str

    def __init__(
        self,
        *args: Any,
        version_format: str = xDSLOptMain.version_format,
        **kwargs: Any,
    ):
        super().__init__(nargs=0, *args, **kwargs)
        self.version_format = version_format

    def __call__(
        self,
        parser: argparse.ArgumentParser,
        namespace: argparse.Namespace,
        values: Any,
        option_string: str | None = None,
    ) -> None:
        print(self.version_format.format(xdsl_version=version("xdsl")))
        parser.exit()