sonolus.py 0.1.8__py3-none-any.whl → 0.2.0__py3-none-any.whl

sonolus/backend/finalize.py

@@ -42,6 +42,7 @@ def cfg_to_engine_node(entry: BasicBlock):
                     for cond, target in targets.items():
                         if cond is None:
                             default = block_indexes[target]
+                            continue
                         args.append(ConstantNode(value=cond))
                         args.append(ConstantNode(value=block_indexes[target]))
                     args.append(ConstantNode(value=default))
@@ -55,7 +56,7 @@ def ir_to_engine_node(stmt) -> EngineNode:
     match stmt:
         case int() | float():
             return ConstantNode(value=float(stmt))
-        case IRConst(value=value):
+        case IRConst(value=int(value) | float(value)):
             return ConstantNode(value=value)
         case IRPureInstr(op=op, args=args) | IRInstr(op=op, args=args):
             return FunctionNode(func=op, args=[ir_to_engine_node(arg) for arg in args])

sonolus/backend/optimize/constant_evaluation.py

@@ -315,8 +315,8 @@ class SparseConditionalConstantPropagation(CompilerPass):
                             return 1
                         return functools.reduce(operator.pow, args)
                     case Op.Log:
-                        assert len(args) == 2
-                        return math.log(args[0], args[1])
+                        assert len(args) == 1
+                        return math.log(args[0])
                     case Op.Ceil:
                         assert len(args) == 1
                         return math.ceil(args[0])

sonolus/backend/optimize/copy_coalesce.py

@@ -39,20 +39,29 @@ class CopyCoalesce(CompilerPass):
         interference = self.get_interference(entry)
         copies = self.get_copies(entry)
 
-        mapping = {}
+        mapping: dict[TempBlock, set[TempBlock]] = {}
 
         for target, sources in copies.items():
             for source in sources:
-                if source in mapping:
-                    continue
                 if source in interference.get(target, set()):
                     continue
-                mapping[source] = mapping.get(target, target)
+                combined_mapping = mapping.get(target, {target}) | mapping.get(source, {source})
                 combined_interference = interference.get(target, set()) | interference.get(source, set())
-                interference[source] = combined_interference
-                interference[target] = combined_interference
+                for place in combined_mapping:
+                    mapping[place] = combined_mapping
+                    interference[place] = combined_interference
+                for place in combined_interference:
+                    interference[place].update(combined_mapping)
+
+        canonical_mapping = {}
+        for place, group in mapping.items():
+            if place in canonical_mapping:
+                continue
+            canonical = min(group)
+            for member in group:
+                canonical_mapping[member] = canonical
 
-        return mapping
+        return canonical_mapping
 
     def get_interference(self, entry: BasicBlock) -> dict[TempBlock, set[TempBlock]]:
         result = {}

sonolus/backend/optimize/optimize.py

@@ -12,13 +12,21 @@ from sonolus.backend.optimize.passes import run_passes
 from sonolus.backend.optimize.simplify import CoalesceFlow, NormalizeSwitch, RewriteToSwitch
 from sonolus.backend.optimize.ssa import FromSSA, ToSSA
 
-MINIMAL_PASSES = [
+MINIMAL_PASSES = (
     CoalesceFlow(),
     UnreachableCodeElimination(),
     AllocateBasic(),
-]
+)
 
-STANDARD_PASSES = [
+FAST_PASSES = (
+    CoalesceFlow(),
+    UnreachableCodeElimination(),
+    AdvancedDeadCodeElimination(),
+    CoalesceFlow(),
+    Allocate(),
+)
+
+STANDARD_PASSES = (
     CoalesceFlow(),
     UnreachableCodeElimination(),
     DeadCodeElimination(),
@@ -37,7 +45,7 @@ STANDARD_PASSES = [
     CoalesceFlow(),
     NormalizeSwitch(),
     Allocate(),
-]
+)
 
 
 def optimize_and_allocate(cfg: BasicBlock):

sonolus/backend/optimize/passes.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from collections import deque
+from collections.abc import Sequence
 
 from sonolus.backend.optimize.flow import BasicBlock
 
@@ -35,7 +36,7 @@ class CompilerPass(ABC):
         pass
 
 
-def run_passes(entry: BasicBlock, passes: list[CompilerPass]) -> BasicBlock:
+def run_passes(entry: BasicBlock, passes: Sequence[CompilerPass]) -> BasicBlock:
     active_passes = set()
     queue = deque(passes)
     while queue:

sonolus/backend/place.py

@@ -1,5 +1,5 @@
 from collections.abc import Iterator
-from typing import NamedTuple, Self
+from typing import Self
 
 from sonolus.backend.blocks import Block
 
@@ -8,9 +8,13 @@ type BlockValue = Block | int | TempBlock | Place
 type IndexValue = int | Place
 
 
-class TempBlock(NamedTuple):
-    name: str
-    size: int = 1
+class TempBlock:
+    __slots__ = ("__hash", "name", "size")
+
+    def __init__(self, name: str, size: int = 1):
+        self.name = name
+        self.size = size
+        self.__hash = hash(name)  # Precompute hash based on name alone
 
     def __repr__(self):
         return f"TempBlock(name={self.name!r}, size={self.size!r})"
@@ -28,14 +32,38 @@ class TempBlock(NamedTuple):
     def __eq__(self, other):
         return isinstance(other, TempBlock) and self.name == other.name and self.size == other.size
 
+    def __lt__(self, other):
+        if not isinstance(other, TempBlock):
+            return NotImplemented
+        return str(self) < str(other)
+
+    def __le__(self, other):
+        if not isinstance(other, TempBlock):
+            return NotImplemented
+        return str(self) <= str(other)
+
+    def __gt__(self, other):
+        if not isinstance(other, TempBlock):
+            return NotImplemented
+        return str(self) > str(other)
+
+    def __ge__(self, other):
+        if not isinstance(other, TempBlock):
+            return NotImplemented
+        return str(self) >= str(other)
+
     def __hash__(self):
-        return hash(self.name)  # Typically will be unique by name alone
+        return self.__hash
 
 
-class BlockPlace(NamedTuple):
-    block: BlockValue
-    index: IndexValue = 0
-    offset: int = 0
+class BlockPlace:
+    __slots__ = ("__hash", "block", "index", "offset")
+
+    def __init__(self, block: BlockValue, index: IndexValue = 0, offset: int = 0):
+        self.block = block
+        self.index = index
+        self.offset = offset
+        self.__hash = hash((block, index, offset))
 
     def __repr__(self):
         return f"BlockPlace(block={self.block!r}, index={self.index!r}, offset={self.offset!r})"
@@ -61,13 +89,42 @@ class BlockPlace(NamedTuple):
             and self.offset == other.offset
         )
 
+    def __lt__(self, other):
+        if not isinstance(other, BlockPlace):
+            return NotImplemented
+        return str(self) < str(other)
+
+    def __le__(self, other):
+        if not isinstance(other, BlockPlace):
+            return NotImplemented
+        return str(self) <= str(other)
+
+    def __gt__(self, other):
+        if not isinstance(other, BlockPlace):
+            return NotImplemented
+        return str(self) > str(other)
+
+    def __ge__(self, other):
+        if not isinstance(other, BlockPlace):
+            return NotImplemented
+        return str(self) >= str(other)
+
     def __hash__(self):
-        return hash((self.block, self.index, self.offset))
+        return self.__hash
+
+    def __iter__(self):
+        yield self.block
+        yield self.index
+        yield self.offset
 
 
-class SSAPlace(NamedTuple):
-    name: str
-    num: int
+class SSAPlace:
+    __slots__ = ("__hash", "name", "num")
+
+    def __init__(self, name: str, num: int):
+        self.name = name
+        self.num = num
+        self.__hash = hash((name, num))
 
     def __repr__(self):
         return f"SSAPlace(name={self.name!r}, num={self.num!r})"
@@ -78,5 +135,29 @@ class SSAPlace(NamedTuple):
     def __eq__(self, other):
         return isinstance(other, SSAPlace) and self.name == other.name and self.num == other.num
 
+    def __lt__(self, other):
+        if not isinstance(other, SSAPlace):
+            return NotImplemented
+        return str(self) < str(other)
+
+    def __le__(self, other):
+        if not isinstance(other, SSAPlace):
+            return NotImplemented
+        return str(self) <= str(other)
+
+    def __gt__(self, other):
+        if not isinstance(other, SSAPlace):
+            return NotImplemented
+        return str(self) > str(other)
+
+    def __ge__(self, other):
+        if not isinstance(other, SSAPlace):
+            return NotImplemented
+        return str(self) >= str(other)
+
     def __hash__(self):
-        return hash((self.name, self.num))
+        return self.__hash
+
+    def __iter__(self):
+        yield self.name
+        yield self.num

sonolus/backend/visitor.py

@@ -172,7 +172,9 @@ class Visitor(ast.NodeVisitor):
     bound_args: inspect.BoundArguments
     used_names: dict[str, int]
     return_ctxs: list[Context]  # Contexts at return statements, which will branch to the exit
-    loop_head_ctxs: list[Context]  # Contexts at loop heads, from outer to inner
+    loop_head_ctxs: list[
+        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
@@ -203,6 +205,8 @@ class Visitor(ast.NodeVisitor):
             case ast.FunctionDef(body=body):
                 ctx().scope.set_value("$return", validate_value(None))
                 for stmt in body:
+                    if not ctx().live:
+                        break
                     self.visit(stmt)
             case ast.Lambda(body=body):
                 result = self.visit(body)
@@ -318,11 +322,21 @@ class Visitor(ast.NodeVisitor):
         iterable = self.visit(node.iter)
         if isinstance(iterable, TupleImpl):
             # Unroll the loop
+            break_ctxs = []
             for value in iterable.value:
                 set_ctx(ctx().branch(None))
+                self.loop_head_ctxs.append([])
+                self.break_ctxs.append([])
                 self.handle_assign(node.target, validate_value(value))
                 for stmt in node.body:
+                    if not ctx().live:
+                        break
                     self.visit(stmt)
+                continue_ctxs = [*self.loop_head_ctxs.pop(), ctx()]
+                break_ctxs.extend(self.break_ctxs.pop())
+                set_ctx(Context.meet(continue_ctxs))
+            if break_ctxs:
+                set_ctx(Context.meet([*break_ctxs, ctx()]))
             return
         iterator = self.handle_call(node, iterable.__iter__)
         if not isinstance(iterator, SonolusIterator):
@@ -335,7 +349,11 @@ class Visitor(ast.NodeVisitor):
         has_next = self.ensure_boolean_num(self.handle_call(node, iterator.has_next))
         if has_next._is_py_() and not has_next._as_py_():
             # The loop will never run, continue after evaluating the condition
+            self.loop_head_ctxs.pop()
+            self.break_ctxs.pop()
             for stmt in node.orelse:
+                if not ctx().live:
+                    break
                 self.visit(stmt)
             return
         ctx().test = has_next.ir()
@@ -345,16 +363,21 @@ class Visitor(ast.NodeVisitor):
         set_ctx(body_ctx)
         self.handle_assign(node.target, self.handle_call(node, iterator.next))
         for stmt in node.body:
+            if not ctx().live:
+                break
             self.visit(stmt)
         ctx().branch_to_loop_header(header_ctx)
 
+        self.loop_head_ctxs.pop()
+        break_ctxs = self.break_ctxs.pop()
+
         set_ctx(else_ctx)
         for stmt in node.orelse:
+            if not ctx().live:
+                break
             self.visit(stmt)
         else_end_ctx = ctx()
 
-        self.loop_head_ctxs.pop()
-        break_ctxs = self.break_ctxs.pop()
         after_ctx = Context.meet([else_end_ctx, *break_ctxs])
         set_ctx(after_ctx)
 
@@ -365,27 +388,55 @@ class Visitor(ast.NodeVisitor):
         self.break_ctxs.append([])
         set_ctx(header_ctx)
         test = self.ensure_boolean_num(self.visit(node.test))
-        if test._is_py_() and not test._as_py_():
-            # The loop will never run, continue after evaluating the condition
-            for stmt in node.orelse:
-                self.visit(stmt)
-            return
+        if test._is_py_():
+            if test._as_py_():
+                # The loop will run until a break / return
+                body_ctx = ctx().branch(None)
+                set_ctx(body_ctx)
+                for stmt in node.body:
+                    if not ctx().live:
+                        break
+                    self.visit(stmt)
+                ctx().branch_to_loop_header(header_ctx)
+
+                self.loop_head_ctxs.pop()
+                break_ctxs = self.break_ctxs.pop()
+
+                # Skip the else block
+
+                after_ctx = Context.meet([ctx().into_dead(), *break_ctxs])
+                set_ctx(after_ctx)
+                return
+            else:
+                # The loop will never run, continue after evaluating the condition
+                self.loop_head_ctxs.pop()
+                self.break_ctxs.pop()
+                for stmt in node.orelse:
+                    if not ctx().live:
+                        break
+                    self.visit(stmt)
+                return
         ctx().test = test.ir()
         body_ctx = ctx().branch(None)
         else_ctx = ctx().branch(0)
 
         set_ctx(body_ctx)
         for stmt in node.body:
+            if not ctx().live:
+                break
             self.visit(stmt)
         ctx().branch_to_loop_header(header_ctx)
 
+        self.loop_head_ctxs.pop()
+        break_ctxs = self.break_ctxs.pop()
+
         set_ctx(else_ctx)
         for stmt in node.orelse:
+            if not ctx().live:
+                break
             self.visit(stmt)
         else_end_ctx = ctx()
 
-        self.loop_head_ctxs.pop()
-        break_ctxs = self.break_ctxs.pop()
         after_ctx = Context.meet([else_end_ctx, *break_ctxs])
         set_ctx(after_ctx)
 
@@ -395,9 +446,13 @@ class Visitor(ast.NodeVisitor):
         if test._is_py_():
             if test._as_py_():
                 for stmt in node.body:
+                    if not ctx().live:
+                        break
                     self.visit(stmt)
             else:
                 for stmt in node.orelse:
+                    if not ctx().live:
+                        break
                     self.visit(stmt)
             return
 
@@ -408,11 +463,15 @@ class Visitor(ast.NodeVisitor):
 
         set_ctx(true_ctx)
         for stmt in node.body:
+            if not ctx().live:
+                break
             self.visit(stmt)
         true_end_ctx = ctx()
 
         set_ctx(false_ctx)
         for stmt in node.orelse:
+            if not ctx().live:
+                break
             self.visit(stmt)
         false_end_ctx = ctx()
 
@@ -439,6 +498,8 @@ class Visitor(ast.NodeVisitor):
             if guard._is_py_():
                 if guard._as_py_():
                     for stmt in case.body:
+                        if not ctx().live:
+                            break
                         self.visit(stmt)
                     end_ctxs.append(ctx())
                 else:
@@ -450,6 +511,8 @@ class Visitor(ast.NodeVisitor):
                 guard_false_ctx = ctx().branch(0)
                 set_ctx(guard_true_ctx)
                 for stmt in case.body:
+                    if not ctx().live:
+                        break
                     self.visit(stmt)
                 end_ctxs.append(ctx())
                 false_ctx = Context.meet([false_ctx, guard_false_ctx])
@@ -610,7 +673,11 @@ class Visitor(ast.NodeVisitor):
         set_ctx(ctx().into_dead())
 
     def visit_Continue(self, node):
-        ctx().branch_to_loop_header(self.loop_head_ctxs[-1])
+        loop_head = self.loop_head_ctxs[-1]
+        if isinstance(loop_head, list):
+            loop_head.append(ctx())
+        else:
+            ctx().branch_to_loop_header(loop_head)
         set_ctx(ctx().into_dead())
 
     def visit_BoolOp(self, node) -> Value:
@@ -959,7 +1026,11 @@ class Visitor(ast.NodeVisitor):
             if isinstance(target, ConstantValue):
                 # Unwrap so we can access fields
                 target = target._as_py_()
-            descriptor = type(target).__dict__.get(key)
+            descriptor = None
+            for cls in type.mro(type(target)):
+                descriptor = cls.__dict__.get(key, None)
+                if descriptor is not None:
+                    break
             match descriptor:
                 case property(fget=getter):
                     return self.handle_call(node, getter, target)

sonolus/build/cli.py

@@ -10,10 +10,11 @@ import sys
 from pathlib import Path
 from time import perf_counter
 
+from sonolus.backend.optimize.optimize import FAST_PASSES, MINIMAL_PASSES, STANDARD_PASSES
 from sonolus.build.engine import package_engine
 from sonolus.build.level import package_level_data
 from sonolus.build.project import build_project_to_collection, get_project_schema
-from sonolus.script.project import Project
+from sonolus.script.project import BuildConfig, Project
 
 
 def find_default_module() -> str | None:
@@ -63,29 +64,29 @@ def import_project(module_path: str) -> Project | None:
         return project
     except Exception as e:
         print(f"Error: Failed to import project: {e}")
-        return None
+        raise e from None
 
 
-def build_project(project: Project, build_dir: Path):
+def build_project(project: Project, build_dir: Path, config: BuildConfig):
     dist_dir = build_dir / "dist"
     levels_dir = dist_dir / "levels"
     shutil.rmtree(dist_dir, ignore_errors=True)
     dist_dir.mkdir(parents=True, exist_ok=True)
     levels_dir.mkdir(parents=True, exist_ok=True)
 
-    package_engine(project.engine.data).write(dist_dir / "engine")
+    package_engine(project.engine.data, config).write(dist_dir / "engine")
 
     for level in project.levels:
         level_path = levels_dir / level.name
         level_path.write_bytes(package_level_data(level.data))
 
 
-def build_collection(project: Project, build_dir: Path):
+def build_collection(project: Project, build_dir: Path, config: BuildConfig):
     site_dir = build_dir / "site"
     shutil.rmtree(site_dir, ignore_errors=True)
     site_dir.mkdir(parents=True, exist_ok=True)
 
-    collection = build_project_to_collection(project)
+    collection = build_project_to_collection(project, config)
     collection.write(site_dir)
 
 
@@ -125,10 +126,55 @@ def run_server(base_dir: Path, port: int = 8000):
             httpd.shutdown()
 
 
+def get_config(args: argparse.Namespace) -> BuildConfig:
+    if hasattr(args, "optimize_minimal") and args.optimize_minimal:
+        optimization_passes = MINIMAL_PASSES
+    elif hasattr(args, "optimize_fast") and args.optimize_fast:
+        optimization_passes = FAST_PASSES
+    elif hasattr(args, "optimize_standard") and args.optimize_standard:
+        optimization_passes = STANDARD_PASSES
+    else:
+        optimization_passes = FAST_PASSES if args.command == "dev" else STANDARD_PASSES
+
+    if any(hasattr(args, attr) and getattr(args, attr) for attr in ["play", "watch", "preview", "tutorial"]):
+        build_play = hasattr(args, "play") and args.play
+        build_watch = hasattr(args, "watch") and args.watch
+        build_preview = hasattr(args, "preview") and args.preview
+        build_tutorial = hasattr(args, "tutorial") and args.tutorial
+    else:
+        build_play = build_watch = build_preview = build_tutorial = True
+
+    return BuildConfig(
+        passes=optimization_passes,
+        build_play=build_play,
+        build_watch=build_watch,
+        build_preview=build_preview,
+        build_tutorial=build_tutorial,
+    )
+
+
 def main():
     parser = argparse.ArgumentParser(description="Sonolus project build and development tools")
     subparsers = parser.add_subparsers(dest="command", required=True)
 
+    def add_common_arguments(parser):
+        optimization_group = parser.add_mutually_exclusive_group()
+        optimization_group.add_argument(
+            "-o0", "--optimize-minimal", action="store_true", help="Use minimal optimization passes"
+        )
+        optimization_group.add_argument(
+            "-o1", "--optimize-fast", action="store_true", help="Use fast optimization passes"
+        )
+        optimization_group.add_argument(
+            "-o2", "--optimize-standard", action="store_true", help="Use standard optimization passes"
+        )
+
+        build_components = parser.add_argument_group("build components")
+        build_components.add_argument("--play", action="store_true", help="Build play component")
+        build_components.add_argument("--watch", action="store_true", help="Build watch component")
+        build_components.add_argument("--preview", action="store_true", help="Build preview component")
+        build_components.add_argument("--tutorial", action="store_true", help="Build tutorial component")
+
     build_parser = subparsers.add_parser("build")
     build_parser.add_argument(
         "module",
@@ -137,6 +183,7 @@ def main():
         help="Module path (e.g., 'module.name'). If omitted, will auto-detect if only one module exists.",
     )
     build_parser.add_argument("--build-dir", type=str, default="./build")
+    add_common_arguments(build_parser)
 
     dev_parser = subparsers.add_parser("dev")
     dev_parser.add_argument(
@@ -147,6 +194,7 @@ def main():
     )
     dev_parser.add_argument("--build-dir", type=str, default="./build")
     dev_parser.add_argument("--port", type=int, default=8000)
+    add_common_arguments(dev_parser)
 
     schema_parser = subparsers.add_parser("schema")
     schema_parser.add_argument(
@@ -161,7 +209,8 @@ def main():
     if not args.module:
         default_module = find_default_module()
         if default_module:
-            print(f"Using auto-detected module: {default_module}")
+            if args.command != "schema":
+                print(f"Using auto-detected module: {default_module}")
             args.module = default_module
         else:
             parser.error("Module argument is required when multiple or no modules are found")
@@ -173,13 +222,15 @@ def main():
     if args.command == "build":
         build_dir = Path(args.build_dir)
         start_time = perf_counter()
-        build_project(project, build_dir)
+        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()
-        build_collection(project, build_dir)
+        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)