Context

class FuncInfo(NamedTuple):
    """Information about a decorated function being generated into IR."""

    file: str
    """The path of the file containing the function."""

    globals: dict[str, Any]
    """The globals defined in that file up to the point of function definition."""

    ast: ast.FunctionDef
    """The Python AST representation of the function."""

@dataclass
class PyASTContext:
    """Encapsulate the mapping between Python and IR types and operations."""

    type_registry: TypeRegistry = field(default_factory=TypeRegistry)
    """Mappings between source code and IR type."""

    function_registry: FunctionRegistry = field(default_factory=FunctionRegistry)
    """Mappings between functions and their operation types."""

    literal_registry: LiteralRegistry = field(default_factory=LiteralRegistry)
    """Mappings between literal types and their operation constructors."""

    post_transforms: list[ModulePass] = field(
        default_factory=lambda: [FrontendDesymrefyPass()]
    )
    """An ordered list of passes to apply to the built module."""

    post_callback: PassPipelineCallbackType | None = default_pipeline_callback
    """Callback to run between post transforms."""

    ir_context: Context = field(
        default_factory=lambda: Context(allow_unregistered=True)
    )
    """The xDSL context to use when applying transformations to the built module."""

    function_definition_constructor: FunctionDefinitionConstructor = build_func
    """Construct the operation for each function."""

    return_constructor: ReturnConstructor = build_return
    """Construct return operations, optionally omitting the terminator."""

    def register_type(
        self,
        source_type: TypeForm[Any],
        ir_type: TypeAttribute,
    ) -> None:
        """Associate a type in the source code with its type in the IR."""
        self.type_registry.insert(source_type, ir_type)

    def register_function(
        self, function: Callable[..., Any], ir_constructor: Callable[..., Operation]
    ) -> None:
        """Associate a method on an object in the source code with its IR implementation."""
        self.function_registry.insert(function, ir_constructor)

    def register_literal(
        self, value_type: type[object], ir_constructor: Callable[[Any], Operation]
    ) -> None:
        """Associate a Python literal type with an IR constructor."""
        self.literal_registry.insert(value_type, ir_constructor)

    def register_function_definition(
        self, ir_constructor: FunctionDefinitionConstructor
    ) -> None:
        """Set the definition constructor, receiving name, type, body, and location."""
        self.function_definition_constructor = ir_constructor

    def register_return(self, ir_constructor: ReturnConstructor) -> None:
        """Set the return constructor, receiving expression results and location."""
        self.return_constructor = ir_constructor

    def register_post_transform(self, transform: ModulePass) -> None:
        """Add a module pass to be run on the generated IR."""
        self.post_transforms.append(transform)

    def register_dialect(self, dialect: Dialect) -> None:
        """Add a dialect to the context used for transformation."""
        self.ir_context.load_dialect(dialect)

    @property
    def pass_pipeline(self) -> PassPipeline:
        """Get a pass pipeline from the context state."""
        return PassPipeline(tuple(self.post_transforms), self.post_callback)

    @classmethod
    def _get_func_info(cls, func: Callable[P, R]) -> FuncInfo:
        """Get information about the decorated function."""
        func_file = func.__code__.co_filename
        func_globals = func.__globals__

        # Remove leading indentation from the source code to avoid parsing errors
        source = getsource(func.__code__)
        dedented = textwrap.dedent(source)
        indent = len(source.splitlines()[0]) - len(dedented.splitlines()[0])

        # Retrieve the AST for the function body, without the decorator
        func_ast = ast.parse(dedented).body[0]
        assert isinstance(func_ast, ast.FunctionDef)

        # Count positions from the start of the file, not of the function source
        ast.increment_lineno(func_ast, func.__code__.co_firstlineno - 1)
        for node in ast.walk(func_ast):
            if isinstance(node, ast.expr | ast.stmt | ast.arg | ast.keyword):
                node.col_offset += indent
                if node.end_col_offset is not None:
                    node.end_col_offset += indent
        assert func_ast.name == func.__name__
        assert len(func_ast.decorator_list) == 1
        func_ast.decorator_list = []

        # Return the information about the function
        return FuncInfo(func_file, func_globals, func_ast)

    @classmethod
    def _get_wrapped_program(
        cls, func: Callable[P, R], builder: PyASTBuilder
    ) -> PyASTProgram[P, R]:
        """Return a PyAST program for this function with the builder."""
        program = PyASTProgram[P, R](
            name=func.__name__,
            func=func,
            _builder=builder,
        )
        functools.update_wrapper(program, func)
        assert program.__doc__ == func.__doc__
        return program

    def parse_program(self, func: Callable[P, R]) -> PyASTProgram[P, R]:
        """Get a program wrapper by decorating a function."""
        func_file, func_globals, func_ast = self._get_func_info(func)
        builder = PyASTBuilder(
            type_registry=self.type_registry,
            function_registry=self.function_registry,
            literal_registry=self.literal_registry,
            file=func_file,
            globals=func_globals,
            function_ast=func_ast,
            build_context=self.ir_context,
            post_transforms=self.pass_pipeline,
            function_definition_constructor=self.function_definition_constructor,
            return_constructor=self.return_constructor,
        )
        return self._get_wrapped_program(func, builder)