sonolus.py 0.8.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

sonolus/backend/ir.py

@@ -5,12 +5,25 @@ type IRExpr = IRConst | IRPureInstr | IRGet
 type IRStmt = IRExpr | IRInstr | IRSet
 
 
+_IR_CONST_CACHE_START = -5
+_IR_CONST_CACHE_STOP = 257
+
+
 class IRConst:
+    __slots__ = ("value",)
+
     value: float | int
 
+    def __new__(cls, value):
+        if float(value).is_integer():
+            int_value = int(value)
+            if _IR_CONST_CACHE_START <= int_value < _IR_CONST_CACHE_STOP:
+                return _IR_CONST_CACHE[int_value - _IR_CONST_CACHE_START]
+            else:
+                return _create_raw_const(int_value)
+        return super().__new__(cls)
+
     def __init__(self, value: float):
-        if isinstance(value, bool):
-            value = int(value)
         self.value = value
 
     def __repr__(self):
@@ -26,7 +39,18 @@ class IRConst:
         return hash(self.value)
 
 
+def _create_raw_const(value: float | int) -> IRConst:
+    result = object.__new__(IRConst)
+    result.value = value
+    return result
+
+
+_IR_CONST_CACHE = tuple(_create_raw_const(i) for i in range(_IR_CONST_CACHE_START, _IR_CONST_CACHE_STOP))
+
+
 class IRPureInstr:
+    __slots__ = ("args", "array_defs", "defs", "is_array_init", "live", "op", "uses", "visited")
+
     op: Op
     args: list[IRExpr]
 
@@ -49,6 +73,8 @@ class IRPureInstr:
 
 
 class IRInstr:
+    __slots__ = ("args", "array_defs", "defs", "is_array_init", "live", "op", "uses", "visited")
+
     op: Op
     args: list[IRExpr]
 
@@ -70,6 +96,8 @@ class IRInstr:
 
 
 class IRGet:
+    __slots__ = ("array_defs", "defs", "is_array_init", "live", "place", "uses", "visited")
+
     place: Place
 
     def __init__(self, place: Place):
@@ -89,6 +117,8 @@ class IRGet:
 
 
 class IRSet:
+    __slots__ = ("array_defs", "defs", "is_array_init", "live", "place", "uses", "value", "visited")
+
     place: Place
     value: IRExpr | IRInstr
 

sonolus/backend/optimize/copy_coalesce.py

@@ -69,10 +69,14 @@ class CopyCoalesce(CompilerPass):
     def get_interference(self, entry: BasicBlock) -> dict[TempBlock, set[TempBlock]]:
         result = {}
         for block in traverse_cfg_preorder(entry):
-            for stmt in [*block.statements, block.test]:
+            for stmt in block.statements:
                 live = {p for p in get_live(stmt) if isinstance(p, TempBlock) and p.size == 1}
                 for place in live:
                     result.setdefault(place, set()).update(live - {place})
+            if not isinstance(block.test, IRConst):
+                live = {p for p in get_live(block.test) if isinstance(p, TempBlock) and p.size == 1}
+                for place in live:
+                    result.setdefault(place, set()).update(live - {place})
         return result
 
     def get_copies(self, entry: BasicBlock) -> dict[TempBlock, set[TempBlock]]:

sonolus/backend/optimize/flow.py

@@ -140,6 +140,20 @@ def cfg_to_mermaid(entry: BasicBlock):
     return f"graph\n{body}"
 
 
+def hash_cfg(entry: BasicBlock) -> int:
+    # Note: not stable between Python runs due to hash randomization
+    block_indexes = {block: i for i, block in enumerate(traverse_cfg_preorder(entry))}
+
+    h = 0
+    for block, index in block_indexes.items():
+        outgoing = [(edge.cond, block_indexes.get(edge.dst, -1)) for edge in block.outgoing]
+        h = hash(
+            (h, index, tuple(block.statements), block.test, tuple(sorted(outgoing, key=lambda x: (x[0] is None, x[0]))))
+        )
+
+    return h
+
+
 def cfg_to_text(entry: BasicBlock) -> str:
     def indent(iterable, prefix="  "):
         for line in iterable:

sonolus/backend/optimize/inlining.py

@@ -50,7 +50,7 @@ class InlineVars(CompilerPass):
             # Update the definition if it's a Get from another SSAPlace until we reach a definition that is not a Get
             while defn and isinstance(defn, IRGet) and isinstance(defn.place, SSAPlace):
                 canonical_definitions[p] = defn
-                defn = definitions.get(defn.place, None)  # Can be None if it's a phi
+                defn = definitions.get(defn.place)  # Can be None if it's a phi
             canonical_defn = canonical_definitions[p]
             if (
                 use_counts.get(p, 0) > 0

sonolus/backend/optimize/liveness.py

@@ -31,8 +31,9 @@ class LivenessAnalysis(CompilerPass):
                 statement.defs = self.get_defs(statement)
                 statement.is_array_init = False  # True if this may be the first assignment to an array
                 statement.array_defs = self.get_array_defs(statement)
-            block.test.live = set()
-            block.test.uses = self.get_uses(block.test, set())
+            if not isinstance(block.test, IRConst):
+                block.test.live = set()
+                block.test.uses = self.get_uses(block.test, set())
         self.preprocess_arrays(entry)
 
     def process(self, entry: BasicBlock):
@@ -72,8 +73,9 @@ class LivenessAnalysis(CompilerPass):
             for place in block.live_out
             if not (isinstance(place, TempBlock) and place.size > 1 and place not in block.array_defs_out)
         }
-        block.test.live.update(live)
-        live.update(block.test.uses)
+        if not isinstance(block.test, IRConst):
+            block.test.live.update(live)
+            live.update(block.test.uses)
         for statement in reversed(block.statements):
             statement.live.update(live)
             if self.can_skip(statement, live):

sonolus/build/cli.py

@@ -1,17 +1,16 @@
 import argparse
-import contextlib
-import http.server
 import importlib
 import json
 import shutil
-import socket
-import socketserver
 import sys
 from pathlib import Path
 from time import perf_counter
+from types import ModuleType
 
 from sonolus.backend.excepthook import print_simple_traceback
 from sonolus.backend.optimize.optimize import FAST_PASSES, MINIMAL_PASSES, STANDARD_PASSES
+from sonolus.build.compile import CompileCache
+from sonolus.build.dev_server import run_server
 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
@@ -35,8 +34,10 @@ def find_default_module() -> str | None:
     return potential_modules[0] if len(potential_modules) == 1 else None
 
 
-def import_project(module_path: str) -> Project | None:
+def import_project(module_path: str) -> tuple[Project, ModuleType, set[str]] | tuple[None, None, None]:
     try:
+        initial_modules = set(sys.modules)
+
         current_dir = Path.cwd()
         if current_dir not in sys.path:
             sys.path.insert(0, str(current_dir))
@@ -44,8 +45,8 @@ def import_project(module_path: str) -> Project | None:
         project = None
 
         try:
-            module = importlib.import_module(module_path)
-            project = getattr(module, "project", None)
+            project_module = importlib.import_module(module_path)
+            project = getattr(project_module, "project", None)
         except ImportError as e:
             if not str(e).endswith(f"'{module_path}'"):
                 # It's an error from the module itself
@@ -61,9 +62,9 @@ def import_project(module_path: str) -> Project | None:
 
         if project is None:
             print(f"Error: No Project instance found in module {module_path} or {module_path}.project")
-            return None
+            return None, None, None
 
-        return project
+        return project, project_module, initial_modules
     except Exception as e:
         print(f"Error: Failed to import project: {e}")
         raise e from None
@@ -87,51 +88,15 @@ def validate_project(project: Project, config: BuildConfig):
     validate_engine(project.engine.data, config)
 
 
-def build_collection(project: Project, build_dir: Path, config: BuildConfig | None):
+def build_collection(project: Project, build_dir: Path, config: BuildConfig | None, cache: CompileCache | None = None):
     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, config)
+    collection = build_project_to_collection(project, config, cache=cache)
     collection.write(site_dir)
 
 
-def get_local_ips():
-    hostname = socket.gethostname()
-    local_ips = []
-
-    with contextlib.suppress(socket.gaierror):
-        local_ips.append(socket.gethostbyname(socket.getfqdn()))
-
-    try:
-        for info in socket.getaddrinfo(hostname, None):
-            ip = info[4][0]
-            if not ip.startswith("127.") and ":" not in ip:
-                local_ips.append(ip)
-    except socket.gaierror:
-        pass
-
-    return sorted(set(local_ips))
-
-
-def run_server(base_dir: Path, port: int = 8000):
-    class DirectoryHandler(http.server.SimpleHTTPRequestHandler):
-        def __init__(self, *args, **kwargs):
-            super().__init__(*args, directory=str(base_dir), **kwargs)
-
-    with socketserver.TCPServer(("", port), DirectoryHandler) as httpd:
-        local_ips = get_local_ips()
-        print(f"Server started on port {port}")
-        print("Available on:")
-        for ip in local_ips:
-            print(f"  http://{ip}:{port}")
-        try:
-            httpd.serve_forever()
-        except KeyboardInterrupt:
-            print("\nStopping server...")
-            httpd.shutdown()
-
-
 def get_config(args: argparse.Namespace) -> BuildConfig:
     if hasattr(args, "optimize_minimal") and args.optimize_minimal:
         optimization_passes = MINIMAL_PASSES
@@ -236,7 +201,7 @@ def main():
         print("Python JIT is enabled")
 
     start_time = perf_counter()
-    project = import_project(args.module)
+    project, project_module, core_module_names = import_project(args.module)
     end_time = perf_counter()
     if project is None:
         sys.exit(1)
@@ -254,10 +219,13 @@ def main():
             build_dir = Path(args.build_dir)
             start_time = perf_counter()
             config = get_config(args)
-            build_collection(project, build_dir, config)
+            cache = CompileCache()
+            build_collection(project, build_dir, config, cache=cache)
             end_time = perf_counter()
             print(f"Build finished in {end_time - start_time:.2f}s")
-            run_server(build_dir / "site", port=args.port)
+            run_server(
+                build_dir / "site", args.port, project_module.__name__, core_module_names, build_dir, config, cache
+            )
         elif args.command == "schema":
             print(json.dumps(get_project_schema(project), indent=2))
         elif args.command == "check":

sonolus/build/compile.py

@@ -1,11 +1,13 @@
 from collections.abc import Callable, Sequence
 from concurrent.futures import Executor, Future
+from threading import Event, Lock
 
 from sonolus.backend.finalize import cfg_to_engine_node
 from sonolus.backend.ir import IRConst, IRInstr
 from sonolus.backend.mode import Mode
+from sonolus.backend.node import EngineNode
 from sonolus.backend.ops import Op
-from sonolus.backend.optimize.flow import BasicBlock
+from sonolus.backend.optimize.flow import BasicBlock, hash_cfg
 from sonolus.backend.optimize.optimize import STANDARD_PASSES
 from sonolus.backend.optimize.passes import CompilerPass, OptimizerConfig, run_passes
 from sonolus.backend.visitor import compile_and_call_at_definition
@@ -25,6 +27,31 @@ from sonolus.script.internal.error import CompilationError
 from sonolus.script.num import _is_num
 
 
+class CompileCache:
+    def __init__(self):
+        self._cache = {}
+        self._lock = Lock()
+        self._event = Event()
+
+    def get(self, entry_hash: int) -> EngineNode | None:
+        with self._lock:
+            if entry_hash not in self._cache:
+                self._cache[entry_hash] = None  # Mark as being compiled
+                return None
+        while True:
+            with self._lock:
+                node = self._cache[entry_hash]
+                if node is not None:
+                    return node
+            self._event.wait()
+
+    def set(self, entry_hash: int, node: EngineNode) -> None:
+        with self._lock:
+            self._cache[entry_hash] = node
+            self._event.set()
+            self._event.clear()
+
+
 def compile_mode(
     mode: Mode,
     rom: ReadOnlyMemory,
@@ -33,6 +60,7 @@ def compile_mode(
     passes: Sequence[CompilerPass] | None = None,
     thread_pool: Executor | None = None,
     validate_only: bool = False,
+    cache: CompileCache | None = None,
 ) -> dict:
     if passes is None:
         passes = STANDARD_PASSES
@@ -63,9 +91,19 @@ def compile_mode(
                 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)
+
+        if cache is not None:
+            cfg_hash = hash_cfg(cfg)
+            node = cache.get(cfg_hash)
+            if node is None:
+                optimized_cfg = run_passes(cfg, passes, OptimizerConfig(mode=mode, callback=cb_info.name))
+                node = cfg_to_engine_node(optimized_cfg)
+                cache.set(cfg_hash, node)
+            node_index = nodes.add(node)
+        else:
+            optimized_cfg = run_passes(cfg, passes, OptimizerConfig(mode=mode, callback=cb_info.name))
+            node = cfg_to_engine_node(optimized_cfg)
+            node_index = nodes.add(node)
 
         if arch is not None:
             cb_order = getattr(cb, "_callback_order_", 0)

sonolus/build/dev_server.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+import argparse
+import contextlib
+import http.server
+import importlib
+import queue
+import shlex
+import socket
+import socketserver
+import sys
+import textwrap
+import threading
+import traceback
+from dataclasses import dataclass
+from pathlib import Path
+from time import perf_counter
+from typing import TYPE_CHECKING, Protocol
+
+from sonolus.backend.excepthook import print_simple_traceback
+from sonolus.backend.utils import get_function, get_functions, get_tree_from_file
+from sonolus.build.compile import CompileCache
+from sonolus.script.internal.error import CompilationError
+
+if TYPE_CHECKING:
+    from sonolus.script.project import BuildConfig
+
+HELP_TEXT = """
+[r]ebuild
+[q]uit
+""".strip()
+
+HELP_TEXT = textwrap.dedent(HELP_TEXT)
+
+
+class Command(Protocol):
+    def execute(
+        self,
+        project_module_name: str,
+        core_module_names: set[str],
+        build_dir: Path,
+        config: BuildConfig,
+        cache: CompileCache,
+    ) -> None: ...
+
+
+@dataclass
+class RebuildCommand:
+    def execute(
+        self,
+        project_module_name: str,
+        core_module_names: set[str],
+        build_dir: Path,
+        config: BuildConfig,
+        cache: CompileCache,
+    ):
+        from sonolus.build.cli import build_collection
+
+        for module_name in tuple(sys.modules):
+            if module_name not in core_module_names:
+                del sys.modules[module_name]
+
+        try:
+            project_module = importlib.import_module(project_module_name)
+        except Exception:
+            print(traceback.format_exc())
+            return
+
+        get_function.cache_clear()
+        get_tree_from_file.cache_clear()
+        get_functions.cache_clear()
+        print("Rebuilding...")
+        try:
+            start_time = perf_counter()
+            build_collection(project_module.project, build_dir, config, cache=cache)
+            end_time = perf_counter()
+            print(f"Rebuild completed in {end_time - start_time:.2f} seconds")
+        except CompilationError:
+            exc_info = sys.exc_info()
+            print_simple_traceback(*exc_info)
+
+
+@dataclass
+class ExitCommand:
+    def execute(
+        self,
+        project_module_name: str,
+        core_module_names: set[str],
+        build_dir: Path,
+        config: BuildConfig,
+        cache: CompileCache,
+    ):
+        print("Exiting...")
+        sys.exit(0)
+
+
+def parse_dev_command(command_line: str) -> Command | None:
+    parser = argparse.ArgumentParser(prog="", add_help=False, exit_on_error=False)
+    subparsers = parser.add_subparsers(dest="cmd")
+
+    subparsers.add_parser("rebuild", aliases=["r"])
+    subparsers.add_parser("quit", aliases=["q"])
+
+    try:
+        args = parser.parse_args(shlex.split(command_line))
+        if args.cmd in {"rebuild", "r"}:
+            return RebuildCommand()
+        elif args.cmd in {"quit", "q"}:
+            return ExitCommand()
+        return None
+    except argparse.ArgumentError:
+        return None
+
+
+def command_input_thread(command_queue: queue.Queue, stop_event: threading.Event, prompt_event: threading.Event):
+    print(f"\nAvailable commands:\n{HELP_TEXT}")
+
+    while not stop_event.is_set():
+        try:
+            prompt_event.wait()
+            prompt_event.clear()
+
+            if stop_event.is_set():
+                break
+
+            print("\n> ", end="", flush=True)
+            command_line = input()
+            if command_line.strip():
+                cmd = parse_dev_command(command_line.strip())
+                if cmd:
+                    command_queue.put(cmd)
+                    if isinstance(cmd, ExitCommand):
+                        break
+                else:
+                    print(f"Unknown command. Available commands:\n{HELP_TEXT}")
+                    # Show prompt again
+                    prompt_event.set()
+        except EOFError:
+            break
+        except Exception as e:
+            print(f"Error reading command: {e}")
+            break
+
+
+def get_local_ips():
+    hostname = socket.gethostname()
+    local_ips = []
+
+    with contextlib.suppress(socket.gaierror):
+        local_ips.append(socket.gethostbyname(socket.getfqdn()))
+
+    try:
+        for info in socket.getaddrinfo(hostname, None):
+            ip = info[4][0]
+            if not ip.startswith("127.") and ":" not in ip:
+                local_ips.append(ip)
+    except socket.gaierror:
+        pass
+
+    return sorted(set(local_ips))
+
+
+def run_server(
+    base_dir: Path,
+    port: int,
+    project_module_name: str | None,
+    core_module_names: set[str] | None,
+    build_dir: Path,
+    config: BuildConfig,
+    cache: CompileCache,
+):
+    interactive = project_module_name is not None and core_module_names is not None
+
+    class DirectoryHandler(http.server.SimpleHTTPRequestHandler):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, directory=str(base_dir), **kwargs)
+
+        def log_message(self, fmt, *args):
+            sys.stdout.write("\r\033[K")  # Clear line
+            sys.stdout.write(f"{self.address_string()} - - [{self.log_date_time_string()}] {fmt % args}\n")
+            if interactive:
+                sys.stdout.write("> ")
+            sys.stdout.flush()
+
+    with socketserver.TCPServer(("", port), DirectoryHandler) as httpd:
+        local_ips = get_local_ips()
+        print(f"Server started on port {port}")
+        print("Available on:")
+        for ip in local_ips:
+            print(f"  http://{ip}:{port}")
+
+        if interactive:
+            threading.Thread(target=httpd.serve_forever, daemon=True).start()
+
+            command_queue = queue.Queue()
+            stop_event = threading.Event()
+            prompt_event = threading.Event()
+            input_thread = threading.Thread(
+                target=command_input_thread, args=(command_queue, stop_event, prompt_event), daemon=True
+            )
+            input_thread.start()
+
+            prompt_event.set()
+
+            try:
+                while True:
+                    try:
+                        cmd = command_queue.get(timeout=0.5)
+                        cmd.execute(project_module_name, core_module_names, build_dir, config, cache)
+                        prompt_event.set()
+                    except queue.Empty:
+                        continue
+            except KeyboardInterrupt:
+                print("\nStopping server...")
+                sys.exit(0)
+            finally:
+                httpd.shutdown()
+                stop_event.set()
+                prompt_event.set()
+                input_thread.join()
+        else:
+            httpd.serve_forever()