Single-pass compiler: Introduce CompilerPlugin protocol, CompilerHooks, and context objects

[masenf]

Summary

Introduce the foundational types for the new single-pass compiler: CompilerPlugin (protocol), CompilerHooks (dispatcher), BaseContext, PageContext, and CompileContext. These are the building blocks that all subsequent compiler work depends on.

This is part of the 2a. Single Pass Compiler roadmap item. The goal is to replace the current multi-pass compilation approach (where the component tree is walked multiple times to extract imports, hooks, custom code, dynamic imports, refs, and app wrap components separately) with a single tree walk using a plugin architecture.

Background & Motivation

Today, compiling a single page involves at least 6 separate recursive tree walks over the same component tree:

1. _get_all_imports() — collects imports from every component
2. _get_all_hooks() — collects React hooks
3. _get_all_custom_code() — collects custom JS code snippets
4. _get_all_dynamic_imports() — collects dynamic import statements
5. _get_all_refs() — collects refs
6. _get_all_app_wrap_components() — collects app wrapper components

Each of these methods in reflex/components/component.py follows the same pattern: call the non-recursive version on self, then recursively call on all children, merging results. This is wasteful — a single walk could collect all information in one pass.

Additionally, compilation logic is currently spread across reflex/app.py (the _compile() method, ~400 lines), reflex/compiler/compiler.py, and reflex/compiler/utils.py, making it hard to extend or adapt for different targets.

Design

A working demo of the proposed architecture is available as a starting point (see the CompilerPlugin protocol, CompilerHooks, BaseContext, PageContext, CompileContext classes in the demo code shared in the project brief). The key ideas:

CompilerPlugin Protocol

class CompilerPlugin(Protocol):
    async def eval_page(self, page_fn, /, **kwargs) -> PageContext | None: ...
    async def compile_page(self, page_ctx, /, **kwargs) -> None: ...
    async def compile_component(self, comp, /, **kwargs) -> AsyncGenerator[...]: ...

• eval_page: Evaluate a page function into a PageContext. Returns None if the plugin doesn't handle this page.
• compile_page: Perform a transformation on the PageContext during page compilation.
• compile_component: Two-phase async generator for component transformation. Pre-yield (descending): sees component before children are visited. Post-yield (ascending): runs after children are visited, receives (component, children) tuple.

CompilerHooks Dispatcher

Holds a tuple[CompilerPlugin, ...] and dispatches hook calls to all registered plugins in order. Plugin ordering controls priority.

Context Objects

• BaseContext: Async context manager that sets itself as a ContextVar, allowing any code in the call stack to access the current context via ContextClass.get().
• PageContext: Holds per-page compilation state — name, root_component, imports, module_code, etc.
• CompileContext: Holds the full compilation state — list of pages, hooks, compiled results.

Acceptance Criteria

• CompilerPlugin protocol is defined in a new module (e.g. reflex/compiler/plugins.py or reflex/compiler/plugin.py)
• CompilerHooks dataclass is implemented with _dispatch, eval_page, compile_page, and compile_component methods
• BaseContext is implemented with ContextVar-based async context management, get() classmethod, and ensure_context_attached() guard
• PageContext is implemented with fields for: name, route, root_component, imports (list of ParsedImportDict), module_code (set of str), hooks (dict), dynamic_imports (set), refs (dict), app_wrap_components (dict)
• CompileContext is implemented with fields for: pages, hooks (CompilerHooks), compiled_pages dict, and a compile() method that orchestrates page compilation
• All new types have docstrings and type annotations
• Unit tests cover: plugin dispatch ordering, context var lifecycle (enter/exit/get), PageContext field accumulation

Key Files to Understand

• reflex/components/component.py — BaseComponent._get_all_* methods (lines ~1590-2002) show the current tree-walking pattern
• reflex/compiler/compiler.py — _compile_page(), compile_stateful_components(), _get_shared_components_recursive()
• reflex/app.py — App._compile() method (~lines 1151-1549) is the main compilation orchestrator

Notes

• The context objects use ContextVar so that plugins running during compilation can access PageContext.get() or CompileContext.get() from anywhere in the call stack without explicit parameter passing.
• Plugin ordering matters: e.g., a FooPlugin that transforms components must run before ConsolidateImportsPlugin so the transformed components' imports are captured.
• This issue only introduces the types and basic tests. Wiring them into the actual compilation pipeline is a separate issue.

[linear]

ENG-9142

[masenf]

Inspiration

"""Demo for compiler plugins and hooks."""

import contextlib
import dataclasses
import inspect
from collections.abc import AsyncGenerator, Callable
from contextvars import ContextVar, Token
from typing import Any, ClassVar, Protocol, Self

import reflex as rx
from reflex.components.component import BaseComponent, Component
from reflex.components.radix.themes.typography.text import Text
from reflex.utils.format import json_dumps
from reflex.utils.imports import ParsedImportDict, collapse_imports, merge_imports

ComponentAndChildrenTuple = tuple[BaseComponent, list[BaseComponent]]
PageFunction = Callable[[], Component]


class CompilerPlugin(Protocol):
    """Protocol for compiler plugins that hook into different stages of compilation."""

    async def eval_page(
        self, page_fn: PageFunction, /, **kwargs
    ) -> "PageContext | None":
        """Evaluate a page function and return a PageContext.

        Args:
            page_fn: The page function being processed.
            kwargs: Reserved for future use.

        Returns:
            A PageContext capable of compiling the page, or None if the plugin does not handle this page.
        """

    async def compile_page(self, page_ctx: "PageContext", /, **kwargs) -> None:
        """Perform a transformation on the PageContext during page compilation.

        Args:
            page_ctx: The page context being processed.
            kwargs: Reserved for future use.
        """

    async def compile_component(
        self,
        comp: BaseComponent,
        /,
        **kwargs,
    ) -> AsyncGenerator[
        BaseComponent | ComponentAndChildrenTuple | None, ComponentAndChildrenTuple
    ]:
        """Transform a component during compilation.

        Args:
            comp: The original component being processed.
            kwargs: Reserved for future use.

        Receives:
            Tuple of the transformed component and its transformed children.

        Yields:
            - A single transformed component: children will be considered unmodified and applied to the transformed component.
            - A tuple of (component, children): subsequent plugin invocations will receive the updated component and children.
            - None to keep the original component and children.
        """
        yield


@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
class CompilerHooks:
    """Container for compiler plugins that provides methods to dispatch hook calls to the plugins."""

    plugins: tuple[CompilerPlugin, ...]

    async def _dispatch(
        self, hook_name: str, *args, **kwargs
    ) -> AsyncGenerator[Any, None]:
        """Dispatch a hook call to all plugins.

        Args:
            hook_name: The name of the hook to dispatch.
            args: Positional arguments to pass to the hook.
            kwargs: Keyword arguments to pass to the hook.

        Returns:
            An async generator that yields the result of each plugin's hook.
        """
        default_impl = getattr(CompilerPlugin, hook_name)
        return (
            (
                await result
                if inspect.isawaitable(result := hook(*args, **kwargs))
                else result
            )
            for hook in (getattr(plugin, hook_name, None) for plugin in self.plugins)
            if hook is not None and hook is not default_impl
        )

    async def eval_page(
        self, page_fn: PageFunction, /, **kwargs
    ) -> "AsyncGenerator[PageContext | None]":
        """Dispatch the eval_page hook to each plugin.

        The caller should take the first non-None PageContext then aclose the generator.

        Args:
            page_fn: The page function being processed.
            kwargs: Reserved for future use.

        Returns:
            An async generator that yields a PageContext capable of compiling the page,
            or None for each plugin that does not handle the given page function.
        """
        return await self._dispatch("eval_page", page_fn, **kwargs)

    async def compile_page(
        self, page_ctx: "PageContext", /, **kwargs
    ) -> AsyncGenerator[None]:
        """Dispatch the compile_page hook to each plugin.

        Args:
            page_ctx: The page context being processed.
            kwargs: Reserved for future use.

        Returns:
            An async generator that yields the result of each plugin's compile_page hook.
        """
        return await self._dispatch("compile_page", page_ctx, **kwargs)

    async def compile_component(
        self, comp: BaseComponent, /, **kwargs
    ) -> AsyncGenerator[
        AsyncGenerator[
            BaseComponent | ComponentAndChildrenTuple | None, ComponentAndChildrenTuple
        ]
    ]:
        """Dispatch the compile_component hook to each plugin.

        Args:
            comp: The original component being processed.
            kwargs: Reserved for future use.

        Returns:
            An async generator that yields each plugin's compile_component generator.
        """
        return await self._dispatch("compile_component", comp, **kwargs)


@dataclasses.dataclass(slots=True, kw_only=True)
class BaseContext:
    """Base context class that acts as an async context manager to set the context var."""

    __context_var__: ClassVar[ContextVar[Self]]

    _attached_context_token: Token[Self] | None = dataclasses.field(
        default=None, init=False, repr=False
    )

    @classmethod
    def __init_subclass__(cls, **kwargs):
        """Initialize the context variable for the subclass."""
        super().__init_subclass__(**kwargs)
        cls.__context_var__ = ContextVar(cls.__name__)

    @classmethod
    def get(cls) -> Self:
        """Get the context from the context variable.

        Returns:
            The context instance.
        """
        return cls.__context_var__.get()

    async def __aenter__(self) -> Self:
        """Enter the context.

        Returns:
            This context instance.
        """
        if self._attached_context_token is not None:
            msg = "Context is already attached, cannot enter context manager."
            raise RuntimeError(msg)
        self._attached_context_token = self.__context_var__.set(self)
        return self

    async def __aexit__(self, *exc_info):
        """Exit the context."""
        if self._attached_context_token is not None:
            self.__context_var__.reset(self._attached_context_token)
        self._attached_context_token = None

    def ensure_context_attached(self):
        """Ensure that the context is attached to the current context variable.

        Raises:
            RuntimeError: If the context is not attached.
        """
        if self._attached_context_token is None:
            msg = f"{type(self).__name__} must be entered with 'async with' before calling this method."
            raise RuntimeError(msg)


@dataclasses.dataclass(slots=True, kw_only=True)
class PageContext(BaseContext):
    """Context for a single page during compilation."""

    name: str
    root_component: BaseComponent
    imports: list[ParsedImportDict] = dataclasses.field(default_factory=list)
    module_code: set[str] = dataclasses.field(default_factory=set)

    @classmethod
    def from_page_fn(cls, page_fn: Callable[[], Component]) -> "PageContext":
        """Create a PageContext from a page function.

        Args:
            page_fn: The page function to create the context from.

        Returns:
            A PageContext with the name of the page function and the evaluated root component.
        """
        return cls(name=page_fn.__name__, root_component=page_fn())


@dataclasses.dataclass(slots=True, kw_only=True)
class CompileContext(BaseContext):
    """Context for the entire compilation process."""

    pages: list[PageFunction]
    hooks: CompilerHooks

    compiled_pages: dict[str, PageContext] = dataclasses.field(default_factory=dict)

    async def compile(self):
        """Compile the app by dispatching hooks for each page and gathering their contexts."""
        self.ensure_context_attached()
        for page_fn in self.pages:
            async with contextlib.aclosing(
                await self.hooks.eval_page(page_fn)
            ) as page_ctx_gen:
                async for page_ctx in page_ctx_gen:
                    if page_ctx is not None:
                        self.compiled_pages[page_ctx.name] = page_ctx
                        break
                else:
                    msg = f"No plugin was able to evaluate page function {page_fn.__name__}."
                    raise RuntimeError(msg)
            async with (
                page_ctx as page_ctx,
                contextlib.aclosing(
                    await self.hooks.compile_page(page_ctx)
                ) as compile_page_gen,
            ):
                async for _ in compile_page_gen:
                    pass


class DefaultPagePlugin(CompilerPlugin):
    """Plugin that evaluates and compiles a page."""

    async def eval_page(self, page_fn: PageFunction) -> PageContext:
        """Evaluate the page function.

        Args:
            page_fn: The page function being processed.

        Returns:
            A PageContext with the evaluated root component.
        """
        return PageContext.from_page_fn(page_fn)

    async def compile_page(self, page_ctx: PageContext) -> None:
        """Compile the page.

        Args:
            page_ctx: The page context being processed.
        """
        page_ctx.root_component = await self._compile_component_recursive(
            page_ctx.root_component
        )

    async def _compile_component_recursive(self, comp: BaseComponent) -> BaseComponent:
        """Compile a component by dispatching the compile_component hook to each plugin.

        This method recursively walks the component tree, dispatching the
        compile_component hook for each component. Plugins can transform the
        component and its children on the way down the tree, and then receive
        the transformed component and children on the way back up the tree to
        apply further transformations.

        Args:
            comp: The component to compile.

        Returns:
            The compiled, possibly transformed component.
        """
        gens = []
        children = comp.children
        async with contextlib.aclosing(
            await CompileContext.get().hooks.compile_component(comp)
        ) as compile_component_gen:
            async for gen in compile_component_gen:
                gens.append(gen)
                if (replaced := await anext(gen)) is not None:
                    if isinstance(replaced, tuple):
                        comp, children = replaced
                    else:
                        comp = replaced
        # Recursively walk the children with the same hook transformations.
        children = [
            await self._compile_component_recursive(child) for child in children
        ]
        # Send the transformed component and children back to the hook generators to
        # allow plugins to apply additional transformations on the way back up the tree.
        for gen in gens:
            with contextlib.suppress(StopAsyncIteration):
                if (
                    replaced := await gen.asend((  # codespell:ignore asend
                        comp,
                        children,
                    ))
                ) is not None:
                    if isinstance(replaced, tuple):
                        comp = replaced[0]
                        if replaced[1] is not children:
                            # If the plugin changed the children, we need to traverse back down to incorporate those changes.
                            children = [
                                await self._compile_component_recursive(child)
                                for child in replaced[1]
                            ]
                    else:
                        comp = replaced
        # Apply the transformed children to the (potentially) transformed component before returning.
        comp.children = children
        return comp


class ConsolidateImportsPlugin(CompilerPlugin):
    """Plugin that collapses and merges page imports into a single import dict."""

    async def compile_component(
        self,
        comp: BaseComponent,
    ) -> AsyncGenerator[None, ComponentAndChildrenTuple]:
        """Gather imports from the compiled component.

        Args:
            comp: The original component being processed.

        Yields:
            None, no transformation is done by this plugin.
        """
        # Get imports from the potentially updated component on the way back up the tree.
        comp, _ = yield
        if isinstance(comp, Component) and (imports := comp._get_imports()):
            page_ctx = PageContext.get()
            page_ctx.imports.append(imports)

    async def compile_page(self, page_ctx: PageContext) -> None:
        """Consolidate the page imports.

        This runs _after_ the page has already been compiled by
        DefaultPagePlugin, so we can be sure we've gathered all imports from the
        page context.

        Args:
            page_ctx: The page context being processed.
        """
        page_ctx.imports = [collapse_imports(merge_imports(*page_ctx.imports))]


class ConsolidateModuleCodePlugin(CompilerPlugin):
    """Plugin that gathers custom module code from components."""

    async def compile_component(
        self,
        comp: BaseComponent,
    ) -> AsyncGenerator[None, ComponentAndChildrenTuple]:
        """Gather custom module code from the compiled component.

        Args:
            comp: The original component being processed.

        Yields:
            None, no transformation is done by this plugin.
        """
        # Get custom code from the potentially updated component on the way back up the tree.
        comp, _ = yield
        if isinstance(comp, Component) and (module_code := comp._get_custom_code()):
            page_ctx = PageContext.get()
            page_ctx.module_code.add(module_code)


class FooPlugin(CompilerPlugin):
    """Example plugin that transforms text components to have red color and large font size."""

    async def compile_component(
        self,
        comp: BaseComponent,
    ) -> AsyncGenerator[ComponentAndChildrenTuple | None, ComponentAndChildrenTuple]:
        """Transform text components.

        Args:
            comp: The component to transform.

        Yields:
            The transformed component, or None to keep the original.
        """
        comp, children = yield
        if isinstance(comp, Text) and comp.style.get("color") is None:
            yield (
                rx.badge(color_scheme="amber"),
                [rx.text(*children, color="red", font_size="4em")],
            )
        else:
            yield


async def main():
    """Example usage of the compiler with plugins."""

    class MyWeirdComponent(rx.Fragment):
        def _get_custom_code(self) -> str | None:
            return "console.log('This is some custom code for MyWeirdComponent');"

    def page() -> Component:
        return rx.vstack(
            rx.text("Hello, World!"),
            rx.button("Click me", on_click=rx.console_log("clicked")),
            MyWeirdComponent.create(),
        )

    def page2() -> Component:
        return rx.container(
            rx.vstack(
                rx.text("About this app"),
                rx.text("This is a demo of Reflex."),
            )
        )

    async with CompileContext(
        pages=[page, page2],
        hooks=CompilerHooks(
            # Order is important! If FooPlugin follows ConsolidateImportsPlugin,
            # the transformed components won't be able to contribute their
            # imports to the page context.
            plugins=(
                DefaultPagePlugin(),
                FooPlugin(),
                ConsolidateImportsPlugin(),
                ConsolidateModuleCodePlugin(),
            ),
        ),
    ) as compile_ctx:
        await compile_ctx.compile()
    print(json_dumps(compile_ctx.compiled_pages, indent=2))  # noqa: T201


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())