Source code for guppylang.optimizer

"""Optimization configuration for Guppy compilation.

Guppy applies a predefined set of optimization passes when compiling a program
using the TKET compiler.

These passes clean up artifacts introduced by the compiler and may simplify both
classical and quantum operations when calling ``compile()``,
``compile_function()``, or ``emulator()``.

Use ``with_opt_level()`` before compiling or creating an emulator to select a
different :py:class:`OptimizationLevel`. The method can be chained before
compiling or creating an emulator.


Choosing an optimization level
------------------------------

Pass a member of :py:class:`OptimizationLevel` to ``with_opt_level()``:

.. code-block:: python

    from guppylang import OptimizationLevel, guppy
    from guppylang.std.builtins import output
    from guppylang.std.quantum import h, measure, qubit

    @guppy
    def main() -> None:
        q = qubit()
        h(q)
        h(q)
        if measure(q):
            output("result", 2 + 2)
        else:
            output("result", 3 + 3)

    # Classical optimization will keep the self-inverse Hadamard gates.
    package = main.with_opt_level(OptimizationLevel.Classical).compile()

The available levels are:

* :py:attr:`OptimizationLevel.Default` applies Guppy's standard optimization
  level. This may include both classical and quantum optimizations that do
  not alter the program's gateset. Calling ``main.compile()`` or
  ``main.emulator(...)`` directly uses this level.
* :py:attr:`OptimizationLevel.Classical` restricts optimization to classical
  operations. The program will execute the same quantum operations as the original
  source, but may have a simplified control flow structure.
* :py:attr:`OptimizationLevel.Minimal` applies only structural rewrites needed
  to produce executable output. This is useful for low-level program analysis or
  when more control over the optimization passes is desired.

See :py:class:`OptimizationLevel` for more details.

Note that gate rebasing or other program transformations may still be performed
further down the compilation pipeline where required. For example, emulators may
require a specific gateset when targeting a particular architecture.

``with_minimal_opt()`` is shorthand for selecting :py:attr:`OptimizationLevel.Minimal`.
It disables optional optimizations on the program.

.. code-block:: python

    emulator = main.with_minimal_opt().emulator(n_qubits=1)


Running custom passes
---------------------

Use :py:meth:`OptimizerInstance.with_optimization` to append any HUGR
``ComposablePass`` to an optimization pipeline. For example, the following
starts with minimal optimization and then runs tket's function-inlining pass:

.. code-block:: python

    from tket.passes import InlineFunctions

    # Apply a tket pass to inline Guppy functions
    package = main.with_minimal_opt().with_optimization(InlineFunctions()).compile()

    package = (
        main.with_minimal_opt().with_optimization(passes.InlineFunctions())
        .compile()
    )

Multiple custom passes can be added by chaining ``with_optimization()`` calls.
They run in the order they are added, after the passes supplied by the selected
optimization level:

.. code-block:: python

    package = (
        main.with_opt_level(OptimizationLevel.Classical)
        .with_optimization(first_pass)
        .with_optimization(second_pass)
        .compile()
    )
"""

from __future__ import annotations

import functools
from dataclasses import dataclass, field
from enum import Enum
from typing import (
    TYPE_CHECKING,
    Generic,
    ParamSpec,
    TypeVar,
)

if TYPE_CHECKING:
    from collections.abc import Sequence

    from hugr.package import Package
    from hugr.passes.composable import ComposablePass

    from guppylang.defs import GuppyFunctionDefinition
    from guppylang.emulator import EmulatorBuilder, EmulatorInstance, Platform

__all__ = (
    "OptimizationLevel",
    "OptimizerInstance",
)

P = ParamSpec("P")
Out = TypeVar("Out")




class OptimizationLevel(Enum):
    """Optimization level used when compiling a Guppy program."""

    Default = "default"
    """
    Guppy's standard optimization level.

    This may include both classical and quantum optimizations that do not alter
    the program's gateset. Calling ``main.compile()`` or ``main.emulator(...)``
    directly uses this level.

    Currently, this applies pytket's `RemoveRedundancies` after the
    optimizations in :py:attr:`OptimizationLevel.Classical`. This may be
    modified in future versions.
    """

    Classical = "classical"
    """
    Restricts optimization to classical operations.

    The program will execute the same quantum operations as the original source,
    but may have a simplified control flow structure.

    Currently, this runs tket's `Normalize
    <https://quantinuum.github.io/tket2/generated/tket.passes.Normalize.html#tket.passes.Normalize>`_
    pass to simplify classical control flow and remove redundant classical
    operations. This set may be modified in future versions.
    """

    Minimal = "minimal"
    """
    Applies only structural rewrites needed to produce executable output.

    This is useful for low-level program analysis or when more control over
    the optimization passes is desired.
    """



    def passes(self) -> list[ComposablePass]:
        """Return the list of HUGR passes ran by this optimization level."""
        match self:
            case OptimizationLevel.Default:
                # The pytket dependency could be bypassed by using the json
                # encoding of the passes rather than the pytket objects
                # themselves.
                from pytket.passes import RemoveRedundancies
                from tket import passes

                return [passes.Normalize(), passes.PytketHugrPass(RemoveRedundancies())]
            case OptimizationLevel.Classical:
                from tket import passes

                return [passes.Normalize()]
            case OptimizationLevel.Minimal:
                return []




def _apply_passes(package: Package, passes: Sequence[ComposablePass]) -> Package:
    if not passes:
        return package

    # Compose the passes to trigger any cross-pass optimizations that may be possible.
    composed = functools.reduce(lambda x, y: x.then(y), passes)

    for module in package.modules:
        composed.run(module, inplace=True)

    return package




@dataclass(frozen=True)
class OptimizerInstance(Generic[P, Out]):
    """Builder used to configure optimizations for compiling a Guppy program.

    Obtained by calling :py:meth:`GuppyFunctionDefinition.with_opt_level` or
    :py:meth:`GuppyFunctionDefinition.with_minimal_opt`.

    See :py:mod:`guppylang.optimizer` for usage examples.
    """

    definition: GuppyFunctionDefinition[P, Out]
    passes: list[ComposablePass] = field(default_factory=list)



    def with_optimization(
        self, optimization: ComposablePass
    ) -> OptimizerInstance[P, Out]:
        """Add an additional optimization pass to run while compiling the program."""
        return OptimizerInstance(self.definition, [*self.passes, optimization])




    def emulator(
        self,
        n_qubits: int | None = None,
        builder: EmulatorBuilder | None = None,
        libs: list[Package] | None = None,
        platform: Platform = "helios",
    ) -> EmulatorInstance:
        """Compile this function for emulation with the configured optimizations."""
        return self.definition._emulator(
            self.compile_function(), n_qubits, builder, libs, platform
        )


    def compile(self) -> Package:
        """Compile an execution entrypoint with the configured optimizations.

        Alias for :py:meth:`compile_entrypoint`.
        """
        return self.compile_entrypoint()



    def compile_entrypoint(self) -> Package:
        """Compile an entrypoint with the configured optimizations."""
        return _apply_passes(self.definition._compile_entrypoint(), self.passes)




    def compile_function(self) -> Package:
        """Compile a function with the configured optimizations."""
        return _apply_passes(self.definition._compile_function(), self.passes)