sonolus.py 0.3.4__py3-none-any.whl → 0.4.1__py3-none-any.whl

sonolus/backend/utils.py

@@ -23,29 +23,33 @@ def get_tree_from_file(file: str | Path) -> ast.Module:
 class FindFunction(ast.NodeVisitor):
     def __init__(self, line):
         self.line = line
-        self.node: ast.FunctionDef | None = None
+        self.results: list[ast.FunctionDef | ast.Lambda] = []
 
     def visit_FunctionDef(self, node: ast.FunctionDef):
-        if node.lineno == self.line or (
-            node.decorator_list and (node.decorator_list[-1].end_lineno <= self.line <= node.lineno)
-        ):
-            self.node = node
-        else:
-            self.generic_visit(node)
+        self.results.append(node)
+        self.generic_visit(node)
 
     def visit_Lambda(self, node: ast.Lambda):
-        if node.lineno == self.line:
-            if self.node is not None:
-                raise ValueError("Multiple functions defined on the same line are not supported")
-            self.node = node
-        else:
-            self.generic_visit(node)
+        self.results.append(node)
+        self.generic_visit(node)
 
 
-def find_function(tree: ast.Module, line: int):
-    visitor = FindFunction(line)
+@cache
+def get_functions(tree: ast.Module) -> list[ast.FunctionDef | ast.Lambda]:
+    visitor = FindFunction(0)
     visitor.visit(tree)
-    return visitor.node
+    return visitor.results
+
+
+def find_function(tree: ast.Module, line: int):
+    for node in get_functions(tree):
+        if node.lineno == line or (
+            isinstance(node, ast.FunctionDef)
+            and node.decorator_list
+            and (node.decorator_list[-1].end_lineno <= line <= node.lineno)
+        ):
+            return node
+    raise ValueError("Function not found")
 
 
 class ScanWrites(ast.NodeVisitor):
@@ -61,3 +65,27 @@ def scan_writes(node: ast.AST) -> set[str]:
     visitor = ScanWrites()
     visitor.visit(node)
     return set(visitor.writes)
+
+
+class HasDirectYield(ast.NodeVisitor):
+    def __init__(self):
+        self.started = False
+        self.has_yield = False
+
+    def visit_Yield(self, node: ast.Yield):
+        self.has_yield = True
+
+    def visit_YieldFrom(self, node: ast.YieldFrom):
+        self.has_yield = True
+
+    def visit_FunctionDef(self, node: ast.FunctionDef):
+        if self.started:
+            return
+        self.started = True
+        self.generic_visit(node)
+
+
+def has_yield(node: ast.AST) -> bool:
+    visitor = HasDirectYield()
+    visitor.visit(node)
+    return visitor.has_yield

sonolus/backend/visitor.py

@@ -1,26 +1,31 @@
 # ruff: noqa: N802
+from __future__ import annotations
+
 import ast
 import builtins
 import functools
 import inspect
-from collections.abc import Callable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from inspect import ismethod
 from types import FunctionType, MethodType, MethodWrapperType
-from typing import Any, Never, Self
+from typing import Any, Never
 
 from sonolus.backend.excepthook import install_excepthook
-from sonolus.backend.utils import get_function, scan_writes
+from sonolus.backend.utils import get_function, has_yield, scan_writes
 from sonolus.script.debug import assert_true
 from sonolus.script.internal.builtin_impls import BUILTIN_IMPLS, _bool, _float, _int, _len
 from sonolus.script.internal.constant import ConstantValue
-from sonolus.script.internal.context import Context, EmptyBinding, Scope, ValueBinding, ctx, set_ctx
+from sonolus.script.internal.context import Context, EmptyBinding, Scope, ValueBinding, ctx, set_ctx, using_ctx
 from sonolus.script.internal.descriptor import SonolusDescriptor
 from sonolus.script.internal.error import CompilationError
-from sonolus.script.internal.impl import validate_value
+from sonolus.script.internal.impl import meta_fn, validate_value
+from sonolus.script.internal.transient import TransientValue
 from sonolus.script.internal.tuple_impl import TupleImpl
 from sonolus.script.internal.value import Value
 from sonolus.script.iterator import SonolusIterator
+from sonolus.script.maybe import Maybe
 from sonolus.script.num import Num, _is_num
+from sonolus.script.record import Record
 
 _compiler_internal_ = True
 
@@ -31,6 +36,30 @@ def compile_and_call[**P, R](fn: Callable[P, R], /, *args: P.args, **kwargs: P.k
     return validate_value(generate_fn_impl(fn)(*args, **kwargs))
 
 
+def compile_and_call_at_definition[**P, R](fn: Callable[P, R], /, *args: P.args, **kwargs: P.kwargs) -> R:
+    if not ctx():
+        return fn(*args, **kwargs)
+    source_file, node = get_function(fn)
+    location_args = {
+        "lineno": node.lineno,
+        "col_offset": node.col_offset,
+        "end_lineno": node.lineno,
+        "end_col_offset": node.col_offset + 1,
+    }
+    expr = ast.Expression(
+        body=ast.Call(
+            func=ast.Name(id="fn", ctx=ast.Load(), **location_args),
+            args=[],
+            keywords=[],
+            **location_args,
+        )
+    )
+    return eval(
+        compile(expr, filename=source_file, mode="eval"),
+        {"fn": lambda: compile_and_call(fn, *args, **kwargs), "_filter_traceback_": True, "_traceback_root_": True},
+    )
+
+
 def generate_fn_impl(fn: Callable):
     install_excepthook()
     match fn:
@@ -189,15 +218,18 @@ class Visitor(ast.NodeVisitor):
         Context | list[Context]
     ]  # Contexts at loop heads, from outer to inner. Contains a list for unrolled (tuple) loops
     break_ctxs: list[list[Context]]  # Contexts at break statements, from outer to inner
-    active_ctx: Context | None  # The active context for use in nested functions=
-    parent: Self | None  # The parent visitor for use in nested functions
+    yield_ctxs: list[Context]  # Contexts at yield statements, which will branch to the exit
+    resume_ctxs: list[Context]  # Contexts after yield statements
+    active_ctx: Context | None  # The active context for use in nested functions
+    parent: Visitor | None  # The parent visitor for use in nested functions
+    used_parent_binding_values: dict[str, Value]  # Values of parent bindings used in this function
 
     def __init__(
         self,
         source_file: str,
         bound_args: inspect.BoundArguments,
         global_vars: dict[str, Any],
-        parent: Self | None = None,
+        parent: Visitor | None = None,
     ):
         self.source_file = source_file
         self.globals = global_vars
@@ -206,12 +238,16 @@ class Visitor(ast.NodeVisitor):
         self.return_ctxs = []
         self.loop_head_ctxs = []
         self.break_ctxs = []
+        self.yield_ctxs = []
+        self.resume_ctxs = []
         self.active_ctx = None
         self.parent = parent
+        self.used_parent_binding_values = {}
 
     def run(self, node):
         before_ctx = ctx()
-        set_ctx(before_ctx.branch_with_scope(None, Scope()))
+        start_ctx = before_ctx.branch_with_scope(None, Scope())
+        set_ctx(start_ctx)
         for name, value in self.bound_args.arguments.items():
             ctx().scope.set_value(name, validate_value(value))
         match node:
@@ -224,8 +260,75 @@ class Visitor(ast.NodeVisitor):
             case ast.Lambda(body=body):
                 result = self.visit(body)
                 ctx().scope.set_value("$return", result)
+            case ast.GeneratorExp(elt=elt, generators=generators):
+                first_generator = generators[0]
+                iterable = self.visit(first_generator.iter)
+                if isinstance(iterable, TupleImpl):
+                    initial_iterator = iterable
+                else:
+                    if not hasattr(iterable, "__iter__"):
+                        raise TypeError(f"Object of type '{type(iterable).__name__}' is not iterable")
+                    initial_iterator = self.handle_call(first_generator.iter, iterable.__iter__)
+                    if not isinstance(initial_iterator, SonolusIterator):
+                        raise ValueError("Unsupported iterator")
+                # The initial iterator is evaluated eagerly in Python
+                before_ctx = ctx().branch_with_scope(None, before_ctx.scope.copy())
+                start_ctx = before_ctx.branch_with_scope(None, Scope())
+                set_ctx(start_ctx)
+                self.construct_genexpr(generators, elt, initial_iterator)
+                ctx().scope.set_value("$return", validate_value(None))
             case _:
                 raise NotImplementedError("Unsupported syntax")
+        if has_yield(node) or isinstance(node, ast.GeneratorExp):
+            return_ctx = Context.meet([*self.return_ctxs, ctx()])
+            result_binding = return_ctx.scope.get_binding("$return")
+            if not isinstance(result_binding, ValueBinding):
+                raise ValueError("Function has conflicting return values")
+            if not result_binding.value._is_py_() and result_binding.value._as_py_() is not None:
+                raise ValueError("Generator function return statements must return None")
+            with using_ctx(start_ctx):
+                state_var = Num._alloc_()
+                is_present_var = Num._alloc_()
+            with using_ctx(before_ctx):
+                state_var._set_(0)
+            with using_ctx(return_ctx):
+                state_var._set_(len(self.return_ctxs) + 1)
+                is_present_var._set_(0)
+            del before_ctx.outgoing[None]  # Unlink the state machine body from the call site
+            entry = before_ctx.new_empty_disconnected()
+            entry.test = state_var.ir()
+            for i, tgt in enumerate([start_ctx, *self.resume_ctxs]):
+                entry.outgoing[i] = tgt
+            entry.outgoing[None] = return_ctx
+            yield_between_ctxs = []
+            for i, out in enumerate(self.yield_ctxs, start=1):
+                between = out.branch(None)
+                with using_ctx(between):
+                    state_var._set_(i)
+                yield_between_ctxs.append(between)
+            if yield_between_ctxs:
+                yield_merge_ctx = Context.meet(yield_between_ctxs)
+            else:
+                yield_merge_ctx = before_ctx.new_empty_disconnected()
+                # Making it default to a number for convenience when used with stuff like min, etc.
+                yield_merge_ctx.scope.set_value("$yield", validate_value(0))
+            yield_binding = yield_merge_ctx.scope.get_binding("$yield")
+            if not isinstance(yield_binding, ValueBinding):
+                raise ValueError("Function has conflicting yield values")
+            with using_ctx(yield_merge_ctx):
+                is_present_var._set_(1)
+            next_result_ctx = Context.meet([yield_merge_ctx, return_ctx])
+            set_ctx(before_ctx)
+            return_test = Num._alloc_()
+            next_result_ctx.test = return_test.ir()
+            return Generator(
+                return_test,
+                entry,
+                next_result_ctx,
+                Maybe(present=is_present_var, value=yield_binding.value),
+                self.used_parent_binding_values,
+                self,
+            )
         after_ctx = Context.meet([*self.return_ctxs, ctx()])
         self.active_ctx = after_ctx
         result_binding = after_ctx.scope.get_binding("$return")
@@ -234,6 +337,69 @@ class Visitor(ast.NodeVisitor):
         set_ctx(after_ctx.branch_with_scope(None, before_ctx.scope.copy()))
         return result_binding.value
 
+    def construct_genexpr(
+        self, generators: Iterable[ast.comprehension], elt: ast.expr, initial_iterator: Value | None = None
+    ):
+        if not generators:
+            # Note that there may effectively be multiple yields in an expression since
+            # tuples are unrolled.
+            value = self.visit(elt)
+            ctx().scope.set_value("$yield", validate_value(value))
+            self.yield_ctxs.append(ctx())
+            resume_ctx = ctx().new_disconnected()
+            self.resume_ctxs.append(resume_ctx)
+            set_ctx(resume_ctx)
+            return
+        generator, *others = generators
+        if initial_iterator is not None:
+            iterable = initial_iterator
+        else:
+            iterable = self.visit(generator.iter)
+        if isinstance(iterable, TupleImpl):
+            for value in iterable.value:
+                set_ctx(ctx().branch(None))
+                self.handle_assign(generator.target, validate_value(value))
+                self.construct_genexpr(others, elt)
+        else:
+            if initial_iterator is not None:
+                iterator = initial_iterator
+            else:
+                if not hasattr(iterable, "__iter__"):
+                    raise TypeError(f"Object of type '{type(iterable).__name__}' is not iterable")
+                iterator = self.handle_call(generator.iter, iterable.__iter__)
+            if not isinstance(iterator, SonolusIterator):
+                raise ValueError("Unsupported iterator")
+            header_ctx = ctx().branch(None)
+            set_ctx(header_ctx)
+            next_value = self.handle_call(generator.iter, iterator.next)
+            if not isinstance(next_value, Maybe):
+                raise ValueError("Iterator next must return a Maybe")
+            if next_value._present._is_py_() and not next_value._present._as_py_():
+                # This will never run
+                return
+            ctx().test = next_value._present.ir()
+            body_ctx = ctx().branch(None)
+            else_ctx = ctx().branch(0)
+            set_ctx(body_ctx)
+            self.handle_assign(generator.target, next_value._value)
+            for if_expr in generator.ifs:
+                test = self.convert_to_boolean_num(if_expr, self.visit(if_expr))
+                if test._is_py_():
+                    if test._as_py_():
+                        continue
+                    else:
+                        ctx().outgoing[None] = header_ctx
+                        set_ctx(ctx().into_dead())
+                else:
+                    if_then_ctx = ctx().branch(None)
+                    if_else_ctx = ctx().branch(0)
+                    ctx().test = test.ir()
+                    if_else_ctx.outgoing[None] = header_ctx
+                    set_ctx(if_then_ctx)
+            self.construct_genexpr(others, elt)
+            ctx().outgoing[None] = header_ctx
+            set_ctx(else_ctx)
+
     def visit(self, node):
         """Visit a node."""
         # We want this here so this is filtered out of tracebacks
@@ -351,6 +517,8 @@ class Visitor(ast.NodeVisitor):
             if break_ctxs:
                 set_ctx(Context.meet([*break_ctxs, ctx()]))
             return
+        if not hasattr(iterable, "__iter__"):
+            raise TypeError(f"Object of type '{type(iterable).__name__}' is not iterable")
         iterator = self.handle_call(node, iterable.__iter__)
         if not isinstance(iterator, SonolusIterator):
             raise ValueError("Unsupported iterator")
@@ -359,8 +527,10 @@ class Visitor(ast.NodeVisitor):
         self.loop_head_ctxs.append(header_ctx)
         self.break_ctxs.append([])
         set_ctx(header_ctx)
-        has_next = self.convert_to_boolean_num(node, self.handle_call(node, iterator.has_next))
-        if has_next._is_py_() and not has_next._as_py_():
+        next_value = self.handle_call(node, iterator.next)
+        if not isinstance(next_value, Maybe):
+            raise ValueError("Iterator next must return a Maybe")
+        if next_value._present._is_py_() and not next_value._present._as_py_():
             # The loop will never run, continue after evaluating the condition
             self.loop_head_ctxs.pop()
             self.break_ctxs.pop()
@@ -369,12 +539,12 @@ class Visitor(ast.NodeVisitor):
                     break
                 self.visit(stmt)
             return
-        ctx().test = has_next.ir()
+        ctx().test = next_value._present.ir()
         body_ctx = ctx().branch(None)
         else_ctx = ctx().branch(0)
 
         set_ctx(body_ctx)
-        self.handle_assign(node.target, self.handle_call(node, iterator.next))
+        self.handle_assign(node.target, next_value._value)
         for stmt in node.body:
             if not ctx().live:
                 break
@@ -811,16 +981,50 @@ class Visitor(ast.NodeVisitor):
         raise NotImplementedError("Dict comprehensions are not supported")
 
     def visit_GeneratorExp(self, node):
-        raise NotImplementedError("Generator expressions are not supported")
+        self.active_ctx = ctx()
+        return Visitor(self.source_file, inspect.Signature([]).bind(), self.globals, self).run(node)
 
     def visit_Await(self, node):
         raise NotImplementedError("Await expressions are not supported")
 
     def visit_Yield(self, node):
-        raise NotImplementedError("Yield expressions are not supported")
+        value = self.visit(node.value) if node.value else validate_value(None)
+        ctx().scope.set_value("$yield", value)
+        self.yield_ctxs.append(ctx())
+        resume_ctx = ctx().new_disconnected()
+        self.resume_ctxs.append(resume_ctx)
+        set_ctx(resume_ctx)
+        return validate_value(None)  # send() is unsupported, so yield returns None
 
     def visit_YieldFrom(self, node):
-        raise NotImplementedError("Yield from expressions are not supported")
+        value = self.visit(node.value)
+        if isinstance(value, TupleImpl):
+            for entry in value.value:
+                ctx().scope.set_value("$yield", validate_value(entry))
+                self.yield_ctxs.append(ctx())
+                resume_ctx = ctx().new_disconnected()
+                self.resume_ctxs.append(resume_ctx)
+                set_ctx(resume_ctx)
+            return validate_value(None)
+        if not hasattr(value, "__iter__"):
+            raise TypeError(f"Object of type '{type(value).__name__}' is not iterable")
+        iterator = self.handle_call(node, value.__iter__)
+        if not isinstance(iterator, SonolusIterator):
+            raise ValueError("Expected a SonolusIterator")
+        header = ctx().branch(None)
+        set_ctx(header)
+        result = self.handle_call(node, iterator.next)
+        nothing_branch = ctx().branch(0)
+        some_branch = ctx().branch(None)
+        ctx().test = result._present.ir()
+        set_ctx(some_branch)
+        ctx().scope.set_value("$yield", result._value)
+        self.yield_ctxs.append(ctx())
+        resume_ctx = ctx().new_disconnected()
+        self.resume_ctxs.append(resume_ctx)
+        resume_ctx.outgoing[None] = header
+        set_ctx(nothing_branch)
+        return validate_value(None)
 
     def _has_real_method(self, obj: Value, method_name: str) -> bool:
         return hasattr(obj, method_name) and not isinstance(getattr(obj, method_name), MethodWrapperType)
@@ -929,11 +1133,16 @@ class Visitor(ast.NodeVisitor):
         raise NotImplementedError("Starred expressions are not supported")
 
     def visit_Name(self, node):
+        if node.id in self.used_parent_binding_values:
+            return self.used_parent_binding_values[node.id]
         self.active_ctx = ctx()
         v = self
         while v:
             if not isinstance(v.active_ctx.scope.get_binding(node.id), EmptyBinding):
-                return v.active_ctx.scope.get_value(node.id)
+                result = v.active_ctx.scope.get_value(node.id)
+                if v is not self:
+                    self.used_parent_binding_values[node.id] = result
+                return result
             v = v.parent
         if node.id in self.globals:
             value = self.globals[node.id]
@@ -1141,6 +1350,9 @@ class Visitor(ast.NodeVisitor):
             if length._is_py_():
                 return Num._accept_(length._as_py_() > 0)
             return length > Num._accept_(0)
+        if isinstance(value, Record):
+            return Num._accept_(1)
+        # Not allowing other types to default to truthy for now in case there's any edge cases.
         raise TypeError(f"Converting {type(value).__name__} to bool is not supported")
 
     def arguments_to_signature(self, arguments: ast.arguments) -> inspect.Signature:
@@ -1269,3 +1481,62 @@ class ReportingErrorsAtNode:
 
         if exc_value is not None:
             self.compiler.raise_exception_at_node(self.node, exc_value)
+
+
+class Generator(TransientValue, SonolusIterator):
+    def __init__(
+        self,
+        return_test: Num,
+        entry: Context,
+        exit_: Context,
+        value: Maybe,
+        used_bindings: dict[str, Value],
+        parent: Visitor,
+    ):
+        self.i = 0
+        self.return_test = return_test
+        self.entry = entry
+        self.exit = exit_
+        self.value = value
+        self.used_bindings = used_bindings
+        self.parent = parent
+
+    @meta_fn
+    def next(self):
+        self._validate_bindings()
+        self.return_test._set_(self.i)
+        after_ctx = ctx().new_disconnected()
+        ctx().outgoing[None] = self.entry
+        self.exit.outgoing[self.i] = after_ctx
+        self.i += 1
+        set_ctx(after_ctx)
+        return self.value
+
+    def _validate_bindings(self):
+        for key, value in self.used_bindings.items():
+            v = self.parent
+            while v:
+                if not isinstance(v.active_ctx.scope.get_binding(key), EmptyBinding):
+                    result = v.active_ctx.scope.get_value(key)
+                    if result is not value:
+                        raise ValueError(f"Binding '{key}' has been modified since the generator was created")
+                v = v.parent
+
+    def __iter__(self):
+        return self
+
+    @classmethod
+    def _accepts_(cls, value: Any) -> bool:
+        return isinstance(value, cls)
+
+    @classmethod
+    def _accept_(cls, value: Any) -> Generator:
+        if not cls._accepts_(value):
+            raise TypeError(f"Cannot accept value of type {type(value).__name__} as {cls.__name__}")
+        return value
+
+    def _is_py_(self) -> bool:
+        return False
+
+    def _as_py_(self) -> Any:
+        raise NotImplementedError

sonolus/build/cli.py

@@ -10,10 +10,12 @@ import sys
 from pathlib import Path
 from time import perf_counter
 
+from sonolus.backend.excepthook import print_simple_traceback
 from sonolus.backend.optimize.optimize import FAST_PASSES, MINIMAL_PASSES, STANDARD_PASSES
-from sonolus.build.engine import no_gil, package_engine
+from sonolus.build.engine import no_gil, package_engine, validate_engine
 from sonolus.build.level import package_level_data
 from sonolus.build.project import build_project_to_collection, get_project_schema
+from sonolus.script.internal.error import CompilationError
 from sonolus.script.project import BuildConfig, Project
 
 
@@ -81,7 +83,11 @@ def build_project(project: Project, build_dir: Path, config: BuildConfig):
         level_path.write_bytes(package_level_data(level.data))
 
 
-def build_collection(project: Project, build_dir: Path, config: BuildConfig):
+def validate_project(project: Project, config: BuildConfig):
+    validate_engine(project.engine.data, config)
+
+
+def build_collection(project: Project, build_dir: Path, config: BuildConfig | None):
     site_dir = build_dir / "site"
     shutil.rmtree(site_dir, ignore_errors=True)
     site_dir.mkdir(parents=True, exist_ok=True)
@@ -204,6 +210,15 @@ def main():
         help="Module path (e.g., 'module.name'). If omitted, will auto-detect if only one module exists.",
     )
 
+    check_parser = subparsers.add_parser("check")
+    check_parser.add_argument(
+        "module",
+        type=str,
+        nargs="?",
+        help="Module path (e.g., 'module.name'). If omitted, will auto-detect if only one module exists.",
+    )
+    add_common_arguments(check_parser)
+
     args = parser.parse_args()
 
     if not args.module:
@@ -220,24 +235,37 @@ def main():
     if hasattr(sys, "_jit") and sys._jit.is_enabled():
         print("Python JIT is enabled")
 
+    start_time = perf_counter()
     project = import_project(args.module)
+    end_time = perf_counter()
     if project is None:
         sys.exit(1)
+    print(f"Project imported in {end_time - start_time:.2f}s")
 
-    if args.command == "build":
-        build_dir = Path(args.build_dir)
-        start_time = perf_counter()
-        config = get_config(args)
-        build_project(project, build_dir, config)
-        end_time = perf_counter()
-        print(f"Project built successfully to '{build_dir.resolve()}' in {end_time - start_time:.2f}s")
-    elif args.command == "dev":
-        build_dir = Path(args.build_dir)
-        start_time = perf_counter()
-        config = get_config(args)
-        build_collection(project, build_dir, config)
-        end_time = perf_counter()
-        print(f"Build finished in {end_time - start_time:.2f}s")
-        run_server(build_dir / "site", port=args.port)
-    elif args.command == "schema":
-        print(json.dumps(get_project_schema(project), indent=2))
+    try:
+        if args.command == "build":
+            build_dir = Path(args.build_dir)
+            start_time = perf_counter()
+            config = get_config(args)
+            build_project(project, build_dir, config)
+            end_time = perf_counter()
+            print(f"Project built successfully to '{build_dir.resolve()}' in {end_time - start_time:.2f}s")
+        elif args.command == "dev":
+            build_dir = Path(args.build_dir)
+            start_time = perf_counter()
+            config = get_config(args)
+            build_collection(project, build_dir, config)
+            end_time = perf_counter()
+            print(f"Build finished in {end_time - start_time:.2f}s")
+            run_server(build_dir / "site", port=args.port)
+        elif args.command == "schema":
+            print(json.dumps(get_project_schema(project), indent=2))
+        elif args.command == "check":
+            start_time = perf_counter()
+            config = get_config(args)
+            validate_project(project, config)
+            end_time = perf_counter()
+            print(f"Project validation completed successfully in {end_time - start_time:.2f}s")
+    except CompilationError:
+        exc_info = sys.exc_info()
+        print_simple_traceback(*exc_info)

sonolus/build/compile.py

@@ -7,8 +7,8 @@ from sonolus.backend.mode import Mode
 from sonolus.backend.ops import Op
 from sonolus.backend.optimize.flow import BasicBlock
 from sonolus.backend.optimize.optimize import STANDARD_PASSES
-from sonolus.backend.optimize.passes import CompilerPass, run_passes
-from sonolus.backend.visitor import compile_and_call
+from sonolus.backend.optimize.passes import CompilerPass, OptimizerConfig, run_passes
+from sonolus.backend.visitor import compile_and_call_at_definition
 from sonolus.build.node import OutputNodeGenerator
 from sonolus.script.archetype import _BaseArchetype
 from sonolus.script.internal.callbacks import CallbackInfo
@@ -31,6 +31,7 @@ def compile_mode(
     global_callbacks: list[tuple[CallbackInfo, Callable]] | None,
     passes: Sequence[CompilerPass] | None = None,
     thread_pool: Executor | None = None,
+    validate_only: bool = False,
 ) -> dict:
     if passes is None:
         passes = STANDARD_PASSES
@@ -55,7 +56,13 @@ def compile_mode(
         - (cb_info.name, {"index": node_index, "order": cb_order}) for archetype callbacks.
         """
         cfg = callback_to_cfg(global_state, cb, cb_info.name, arch)
-        cfg = run_passes(cfg, passes)
+        if validate_only:
+            if arch is not None:
+                cb_order = getattr(cb, "_callback_order_", 0)
+                return cb_info.name, {"index": 0, "order": cb_order}
+            else:
+                return cb_info.name, 0
+        cfg = run_passes(cfg, passes, OptimizerConfig(mode=mode, callback=cb_info.name))
         node = cfg_to_engine_node(cfg)
         node_index = nodes.add(node)
 
@@ -137,9 +144,9 @@ def callback_to_cfg(
     context = Context(global_state, callback_state)
     with using_ctx(context):
         if archetype is not None:
-            result = compile_and_call(callback, archetype._for_compilation())
+            result = compile_and_call_at_definition(callback, archetype._for_compilation())
         else:
-            result = compile_and_call(callback)
+            result = compile_and_call_at_definition(callback)
         if _is_num(result):
             ctx().add_statements(IRInstr(Op.Break, [IRConst(1), result.ir()]))
     return context_to_cfg(context)