Implement symbolic labels, optimize worker with select

Author: object-Object

python/src/mlogv32/preprocessor/app.py

@@ -25,11 +25,12 @@
 from .filters import FILTERS, ram_var
 from .models import BuildConfig
 from .parser import (
-    DirectiveError,
-    Statement,
+    MlogError,
     check_unsaved_variables,
+    count_statements,
     iter_labels,
     parse_mlog,
+    replace_symbolic_labels,
 )
 from .types import ConfigArgs, ConfigsYaml, Labels
 
@@ -101,11 +102,11 @@ def build(
     size: Annotated[int | None, Option("-s", "--size")] = None,
     output: Annotated[Path | None, Option("-o", "--output")] = None,
     bin_path: Annotated[Path | None, Option("--bin")] = None,
-    include_all: Annotated[bool, Option("--all")] = False,
-    include_cpu: Annotated[bool, Option("--cpu")] = False,
-    include_peripherals: Annotated[bool, Option("--peripherals")] = False,
-    include_memory: Annotated[bool, Option("--memory")] = False,
-    include_debugger: Annotated[bool, Option("--debugger")] = False,
+    include_all: Annotated[bool, Option("-A", "--all")] = False,
+    include_cpu: Annotated[bool, Option("-C", "--cpu")] = False,
+    include_peripherals: Annotated[bool, Option("-P", "--peripherals")] = False,
+    include_memory: Annotated[bool, Option("-M", "--memory")] = False,
+    include_debugger: Annotated[bool, Option("-D", "--debugger")] = False,
     include_keyboard: Annotated[bool, Option("--keyboard/--no-keyboard")] = True,
 ):
     """Generate a CPU schematic.
@@ -260,20 +261,24 @@ def _render_template(
 
     worker_ast = parse_mlog(worker_code)
 
-    check_unsaved_variables(worker_ast)
+    try:
+        check_unsaved_variables(worker_ast)
+        worker_labels = dict(iter_labels(worker_ast))
+        worker_code = replace_symbolic_labels(worker_code, worker_ast, worker_labels)
+    except MlogError as e:
+        e.add_note(f"{worker_output}:{e.token.line}")
+        raise
+
+    # hack
+    worker_output.write_text(worker_code, "utf-8")
+
     print(
         f"""\
 Worker:
-  Instructions: {len(list(n for n in worker_ast if isinstance(n, Statement)))} / 1000
+  Instructions: {count_statements(worker_ast)} / 1000
   Bytes: {len(worker_code.encode())} / {1024 * 100}"""
     )
 
-    try:
-        worker_labels = dict(iter_labels(worker_ast))
-    except DirectiveError as e:
-        e.add_note(f"{worker_output}:{e.token.line}")
-        raise
-
     # preprocess controller
 
     controller_code, _, _ = _render_template(
@@ -291,7 +296,7 @@ def _render_template(
     print(
         f"""\
 Controller:
-  Instructions: {len(list(n for n in controller_ast if isinstance(n, Statement)))} / 1000
+  Instructions: {count_statements(controller_ast)} / 1000
   Bytes: {len(controller_code.encode())} / {1024 * 100}"""
     )
 

python/src/mlogv32/preprocessor/parser/__init__.py

@@ -1,19 +1,25 @@
 __all__ = [
-    "DirectiveError",
     "Directive",
+    "DirectiveError",
     "Label",
+    "MlogError",
     "Statement",
     "check_unsaved_variables",
+    "count_statements",
     "iter_labels",
     "parse_mlog",
+    "replace_symbolic_labels",
 ]
 
 from .mlog import (
     Directive,
     DirectiveError,
     Label,
+    MlogError,
     Statement,
     check_unsaved_variables,
+    count_statements,
     iter_labels,
     parse_mlog,
+    replace_symbolic_labels,
 )

python/src/mlogv32/preprocessor/parser/mlog.lark

@@ -4,13 +4,15 @@ start: (line? _NEWLINE)*
      | statement
      | directive
 
-statement: TOKEN (TOKEN | LABEL | STRING)*
+statement: TOKEN (TOKEN | LABEL_REF | LABEL | STRING)*
 
 directive: _DIRECTIVE /assert_counter|(start|end)_(fetch|assert_length)|(push|pop)_saved/ TOKEN*
 
 STRING.2: /"[^\n]*"/
 
-LABEL.1: /[^\n #\t;]+:/
+LABEL.1: TOKEN ":"
+
+LABEL_REF.1: "%" TOKEN "%"
 
 TOKEN: /[^\n #\t;]+/
 

python/src/mlogv32/preprocessor/parser/mlog.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from importlib import resources
-from typing import TYPE_CHECKING, Callable, Iterable, Iterator
+from typing import TYPE_CHECKING, Callable, Iterator
 
 from lark import Lark, Token, Transformer
 
@@ -34,6 +34,8 @@ class Directive:
 
 type ASTNode = Label | Statement | Directive
 
+type AST = list[ASTNode]
+
 
 @dataclass
 @v_args(inline=True)
@@ -65,12 +67,16 @@ def parse_mlog(text: str):
     return MlogTransformer().transform(tree)
 
 
-class DirectiveError(AssertionError):
+class MlogError(AssertionError):
     def __init__(self, message: str, token: Token, *args: object) -> None:
         super().__init__(message, *args)
         self.token = token
 
 
+class DirectiveError(MlogError):
+    pass
+
+
 def expect_int(
     n: Token,
     predicate: Callable[[int], bool] | None = None,
@@ -85,7 +91,7 @@ def expect_int(
     raise DirectiveError(f"Invalid {description}: {result}", n)
 
 
-def iter_labels(ast: Iterable[ASTNode]) -> Iterator[tuple[str, int]]:
+def iter_labels(ast: AST) -> Iterator[tuple[str, int]]:
     counter = 0
     length_assertion = None
     for node in ast:
@@ -139,7 +145,7 @@ def iter_labels(ast: Iterable[ASTNode]) -> Iterator[tuple[str, int]]:
         )
 
 
-def check_unsaved_variables(ast: Iterable[ASTNode]):
+def check_unsaved_variables(ast: AST):
     saved_variables = {"@counter"}
     warned_variables = set[str]()
     state = "init"
@@ -194,3 +200,64 @@ def check_unsaved_variables(ast: Iterable[ASTNode]):
             warned_variables.add(var)
 
     print(f"Saved variable count: {len(saved_variables - {'@counter'})}")
+
+
+def count_statements(ast: AST):
+    n = 0
+    for node in ast:
+        if isinstance(node, Statement):
+            n += 1
+    return n
+
+
+def replace_symbolic_labels(
+    text: str,
+    ast: AST,
+    labels: dict[str, int],
+    *,
+    check: bool = True,
+) -> str:
+    parts = list[str]()
+
+    start_pos = 0
+
+    for token in _iter_label_refs(ast):
+        label = token[1:-1]
+        if label not in labels:
+            raise MlogError(f"Unknown symbolic label: {label}", token)
+
+        if token.start_pos is None or token.end_pos is None:
+            raise MlogError(f"Internal error: invalid token: {token}", token)
+
+        parts += [
+            text[start_pos : token.start_pos],
+            str(labels[label]),
+        ]
+
+        start_pos = token.end_pos
+
+    parts += text[start_pos:]
+
+    new_text = "".join(parts)
+
+    # sanity checks
+    if check:
+        new_ast = parse_mlog(new_text)
+        if count_statements(ast) != count_statements(new_ast):
+            raise AssertionError(
+                "Source transformation failed: mismatched statement count"
+            )
+        if labels != dict(iter_labels(new_ast)):
+            raise AssertionError("Source transformation failed: mismatched labels")
+
+    return new_text
+
+
+def _iter_label_refs(ast: AST) -> Iterator[Token]:
+    for node in ast:
+        if not isinstance(node, Statement):
+            continue
+        for token in node.args:
+            if token.type != "LABEL_REF":
+                continue
+            yield token