feat: add basic loop overloading

bolt/codegen.py

@@ -961,11 +961,29 @@ def while_statement(
     ) -> Generator[AstNode, Optional[List[str]], Optional[List[str]]]:
         acc.statement(f"while True:")
         with acc.block():
-            condition = yield from visit_single(node.arguments[0], required=True)
-            acc.statement(f"if not {condition}:")
+            acc.statement(f"with _bolt_runtime.scope() as _bolt_condition_commands:")
             with acc.block():
-                acc.statement("break")
-            yield from visit_body(cast(AstRoot, node.arguments[1]), acc)
+                condition = yield from visit_single(node.arguments[0], required=True)
+
+            again = acc.make_variable()
+            loop = acc.helper("loop", condition)
+            acc.statement(f"with {loop} as (_bolt_loop_overridden, {again}):")
+            with acc.block():
+                acc.statement("_bolt_runtime.commands.extend(_bolt_condition_commands)")
+
+                acc.statement("if not _bolt_loop_overridden:")
+                with acc.block():
+                    acc.statement(f"{condition} = bool({condition})")
+
+                with acc.if_statement(condition):
+                    yield from visit_body(cast(AstRoot, node.arguments[1]), acc)
+
+                    acc.statement(f"if {again}:", lineno=node.arguments[0])
+                    with acc.block():
+                        acc.statement("continue")
+
+            acc.statement("break")
+
         return []
 
     @rule(AstCommand, identifier="for:target:in:iterable:body")

bolt/helpers.py

@@ -3,11 +3,11 @@
 ]
 
 
+from contextlib import contextmanager
 from dataclasses import dataclass, replace
 from functools import partial, wraps
 from importlib import import_module
-from types import TracebackType
-from typing import Any, Callable, ContextManager, Dict, Optional, Type
+from typing import Any, Callable, Dict
 from uuid import UUID
 
 from mecha import (
@@ -59,7 +59,8 @@ def get_bolt_helpers() -> Dict[str, Any]:
         "operator_not": operator_not,
         "operator_in": operator_in,
         "operator_not_in": operator_not_in,
-        "branch": BranchDriver,
+        "branch": branch_driver,
+        "loop": loop_driver,
         "get_dup": get_dup,
         "get_rebind": get_rebind,
         "get_attribute_handler": AttributeHandler,
@@ -124,32 +125,24 @@ def operator_not_in(item: Any, container: Any):
     return operator_not(operator_in(item, container))
 
 
-class BranchDriver:
-    obj: Any
-    context_manager: Optional[ContextManager[Any]]
+@contextmanager
+@internal
+def branch_driver(obj: Any):
+    if func := getattr(type(obj), "__branch__", None):
+        with func(obj) as condition:
+            yield condition
+    else:
+        yield obj
 
-    @internal
-    def __init__(self, obj: Any):
-        self.obj = obj
-        self.context_manager = None
-        if func := getattr(type(obj), "__branch__", None):
-            self.context_manager = func(obj)
 
-    @internal
-    def __enter__(self) -> Any:
-        if self.context_manager is not None:
-            return self.context_manager.__enter__()
-        return self.obj
-
-    @internal
-    def __exit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc: Optional[BaseException],
-        traceback: Optional[TracebackType],
-    ) -> Optional[bool]:
-        if self.context_manager is not None:
-            return self.context_manager.__exit__(exc_type, exc, traceback)
+@contextmanager
+@internal
+def loop_driver(obj: Any):
+    if func := getattr(type(obj), "__loop__", None):
+        with func(obj) as cont:
+            yield True, cont
+    else:
+        yield False, True
 
 
 @internal

examples/bolt_loop/beet.yml

@@ -0,0 +1,6 @@
+require:
+  - bolt
+data_pack:
+  load: "src"
+pipeline:
+  - mecha

examples/bolt_loop/src/data/demo/functions/foo.mcfunction

@@ -0,0 +1,90 @@
+from contextlib import contextmanager
+
+
+class LoopAgainAtRuntime:
+    def __init__(self, name: str):
+        self.name = name
+
+    def __bool__(self):
+        function self.name
+        return False
+
+
+class Tmp:
+    def __init__(self, name: str = None):
+        if name is None:
+            name = ctx.generate.format("tmp{incr}")
+        self.name = name
+
+    def __dup__(self):
+        result = Tmp()
+        result = self
+        return result
+
+    def __rebind__(self, rhs: object):
+        if self is rhs:
+            return self
+        if isinstance(rhs, Tmp):
+            scoreboard players operation self.name global = rhs.name global
+        else:
+            scoreboard players set self.name global int(rhs)
+        return self
+
+    @contextmanager
+    def __branch__(self):
+        unless score self.name global matches 0:
+            yield True
+
+    @contextmanager
+    def __loop__(self):
+        name = ctx.generate.format(~/ + "/loop{incr}")
+        execute function name:
+            yield LoopAgainAtRuntime(name)
+
+    def __not__(self):
+        result = Tmp()
+        result = 1
+        unless score self.name global matches 0:
+            result = 0
+        return result
+
+    def __eq__(self, rhs: object):
+        result = Tmp()
+        result = 0
+        if isinstance(rhs, Tmp):
+            if score self.name global = rhs.name global:
+                result = 1
+        else:
+            if score self.name global matches int(rhs):
+                result = 1
+        return result
+
+    def __iadd__(self, rhs: object):
+        if isinstance(rhs, Tmp):
+            scoreboard players operation self.name global += rhs.name global
+        else:
+            scoreboard players add self.name global int(rhs)
+        return self
+
+    def __str__(self):
+        return self.name
+
+
+def display(value: object):
+    if isinstance(value, Tmp):
+        tellraw @a {"score": {"name": value.name, "objective": "global"}}
+    else:
+        tellraw @a {"text": f"{value}"}
+
+
+def f(i):
+    while not i == 4:
+        display(i)
+        i += 1
+
+a = 0
+f(a)
+
+b = Tmp("b")
+b = 0
+f(b)

examples/bolt_loop/src/data/demo/functions/init.mcfunction

@@ -0,0 +1,5 @@
+append function_tag minecraft:load {
+    "values": [__name__]
+}
+
+scoreboard objectives add global dummy

tests/snapshots/bolt__parse_138__1.txt

@@ -1,24 +1,33 @@
-_bolt_lineno = [1, 12], [1, 2]
-_bolt_helper_operator_not = _bolt_runtime.helpers['operator_not']
+_bolt_lineno = [1, 18, 25], [1, 2, 1]
+_bolt_helper_loop = _bolt_runtime.helpers['loop']
 _bolt_helper_branch = _bolt_runtime.helpers['branch']
+_bolt_helper_operator_not = _bolt_runtime.helpers['operator_not']
 _bolt_helper_children = _bolt_runtime.helpers['children']
 _bolt_helper_replace = _bolt_runtime.helpers['replace']
-with _bolt_runtime.scope() as _bolt_var2:
+with _bolt_runtime.scope() as _bolt_var3:
     while True:
-        _bolt_var0 = True
-        if not _bolt_var0:
-            break
-        _bolt_var1 = 'hello'
-        _bolt_var1_inverse = _bolt_helper_operator_not(_bolt_var1)
-        with _bolt_helper_branch(_bolt_var1) as _bolt_condition:
-            if _bolt_condition:
-                pass
-        with _bolt_helper_branch(_bolt_var1_inverse) as _bolt_condition:
-            if _bolt_condition:
-                break
-_bolt_var3 = _bolt_helper_replace(_bolt_refs[0], commands=_bolt_helper_children(_bolt_var2))
+        with _bolt_runtime.scope() as _bolt_condition_commands:
+            _bolt_var0 = True
+        with _bolt_helper_loop(_bolt_var0) as (_bolt_loop_overridden, _bolt_var1):
+            _bolt_runtime.commands.extend(_bolt_condition_commands)
+            if not _bolt_loop_overridden:
+                _bolt_var0 = bool(_bolt_var0)
+            with _bolt_helper_branch(_bolt_var0) as _bolt_condition:
+                if _bolt_condition:
+                    _bolt_var2 = 'hello'
+                    _bolt_var2_inverse = _bolt_helper_operator_not(_bolt_var2)
+                    with _bolt_helper_branch(_bolt_var2) as _bolt_condition:
+                        if _bolt_condition:
+                            pass
+                    with _bolt_helper_branch(_bolt_var2_inverse) as _bolt_condition:
+                        if _bolt_condition:
+                            break
+                    if _bolt_var1:
+                        continue
+        break
+_bolt_var4 = _bolt_helper_replace(_bolt_refs[0], commands=_bolt_helper_children(_bolt_var3))
 ---
-output = _bolt_var3
+output = _bolt_var4
 ---
 _bolt_refs[0]
 <class 'mecha.ast.AstRoot'>

tests/snapshots/bolt__parse_41__1.txt

@@ -1,15 +1,25 @@
 _bolt_lineno = [1], [1]
+_bolt_helper_loop = _bolt_runtime.helpers['loop']
+_bolt_helper_branch = _bolt_runtime.helpers['branch']
 _bolt_helper_children = _bolt_runtime.helpers['children']
 _bolt_helper_replace = _bolt_runtime.helpers['replace']
-with _bolt_runtime.scope() as _bolt_var1:
+with _bolt_runtime.scope() as _bolt_var2:
     while True:
-        _bolt_var0 = True
-        if not _bolt_var0:
-            break
-        _bolt_runtime.commands.extend(_bolt_refs[0].commands)
-_bolt_var2 = _bolt_helper_replace(_bolt_refs[1], commands=_bolt_helper_children(_bolt_var1))
+        with _bolt_runtime.scope() as _bolt_condition_commands:
+            _bolt_var0 = True
+        with _bolt_helper_loop(_bolt_var0) as (_bolt_loop_overridden, _bolt_var1):
+            _bolt_runtime.commands.extend(_bolt_condition_commands)
+            if not _bolt_loop_overridden:
+                _bolt_var0 = bool(_bolt_var0)
+            with _bolt_helper_branch(_bolt_var0) as _bolt_condition:
+                if _bolt_condition:
+                    _bolt_runtime.commands.extend(_bolt_refs[0].commands)
+                    if _bolt_var1:
+                        continue
+        break
+_bolt_var3 = _bolt_helper_replace(_bolt_refs[1], commands=_bolt_helper_children(_bolt_var2))
 ---
-output = _bolt_var2
+output = _bolt_var3
 ---
 _bolt_refs[0]
 <class 'mecha.ast.AstRoot'>

tests/snapshots/bolt__parse_44__1.txt

@@ -1,22 +1,40 @@
-_bolt_lineno = [1], [1]
+_bolt_lineno = [1, 18, 31], [1, 2, 4]
+_bolt_helper_loop = _bolt_runtime.helpers['loop']
+_bolt_helper_branch = _bolt_runtime.helpers['branch']
 _bolt_helper_children = _bolt_runtime.helpers['children']
 _bolt_helper_replace = _bolt_runtime.helpers['replace']
-with _bolt_runtime.scope() as _bolt_var3:
-    with _bolt_runtime.push_nesting('execute:subcommand'), _bolt_runtime.push_nesting('execute:if:score:target:targetObjective:matches:range:subcommand', *_bolt_refs[4:7]), _bolt_runtime.push_nesting('execute:commands'), _bolt_runtime.scope() as _bolt_var2:
+with _bolt_runtime.scope() as _bolt_var5:
+    with _bolt_runtime.push_nesting('execute:subcommand'), _bolt_runtime.push_nesting('execute:if:score:target:targetObjective:matches:range:subcommand', *_bolt_refs[4:7]), _bolt_runtime.push_nesting('execute:commands'), _bolt_runtime.scope() as _bolt_var4:
         while True:
-            _bolt_var0 = True
-            if not _bolt_var0:
-                break
-            _bolt_runtime.commands.extend(_bolt_refs[0].commands)
+            with _bolt_runtime.scope() as _bolt_condition_commands:
+                _bolt_var0 = True
+            with _bolt_helper_loop(_bolt_var0) as (_bolt_loop_overridden, _bolt_var1):
+                _bolt_runtime.commands.extend(_bolt_condition_commands)
+                if not _bolt_loop_overridden:
+                    _bolt_var0 = bool(_bolt_var0)
+                with _bolt_helper_branch(_bolt_var0) as _bolt_condition:
+                    if _bolt_condition:
+                        _bolt_runtime.commands.extend(_bolt_refs[0].commands)
+                        if _bolt_var1:
+                            continue
+            break
         while True:
-            _bolt_var1 = True
-            if not _bolt_var1:
-                break
-            _bolt_runtime.commands.extend(_bolt_refs[1].commands)
-    _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[8], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[7], arguments=_bolt_helper_children([*_bolt_refs[4:7], _bolt_helper_replace(_bolt_refs[3], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[2], commands=_bolt_helper_children(_bolt_var2))]))]))])))
-_bolt_var4 = _bolt_helper_replace(_bolt_refs[9], commands=_bolt_helper_children(_bolt_var3))
+            with _bolt_runtime.scope() as _bolt_condition_commands:
+                _bolt_var2 = True
+            with _bolt_helper_loop(_bolt_var2) as (_bolt_loop_overridden, _bolt_var3):
+                _bolt_runtime.commands.extend(_bolt_condition_commands)
+                if not _bolt_loop_overridden:
+                    _bolt_var2 = bool(_bolt_var2)
+                with _bolt_helper_branch(_bolt_var2) as _bolt_condition:
+                    if _bolt_condition:
+                        _bolt_runtime.commands.extend(_bolt_refs[1].commands)
+                        if _bolt_var3:
+                            continue
+            break
+    _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[8], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[7], arguments=_bolt_helper_children([*_bolt_refs[4:7], _bolt_helper_replace(_bolt_refs[3], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[2], commands=_bolt_helper_children(_bolt_var4))]))]))])))
+_bolt_var6 = _bolt_helper_replace(_bolt_refs[9], commands=_bolt_helper_children(_bolt_var5))
 ---
-output = _bolt_var4
+output = _bolt_var6
 ---
 _bolt_refs[0]
 <class 'mecha.ast.AstRoot'>

tests/snapshots/bolt__parse_45__1.txt

@@ -1,17 +1,27 @@
-_bolt_lineno = [1], [1]
+_bolt_lineno = [1, 18], [1, 2]
+_bolt_helper_loop = _bolt_runtime.helpers['loop']
+_bolt_helper_branch = _bolt_runtime.helpers['branch']
 _bolt_helper_children = _bolt_runtime.helpers['children']
 _bolt_helper_replace = _bolt_runtime.helpers['replace']
-with _bolt_runtime.scope() as _bolt_var2:
-    with _bolt_runtime.push_nesting('execute:subcommand'), _bolt_runtime.push_nesting('execute:if:score:target:targetObjective:matches:range:subcommand', *_bolt_refs[3:6]), _bolt_runtime.push_nesting('execute:commands'), _bolt_runtime.scope() as _bolt_var1:
+with _bolt_runtime.scope() as _bolt_var3:
+    with _bolt_runtime.push_nesting('execute:subcommand'), _bolt_runtime.push_nesting('execute:if:score:target:targetObjective:matches:range:subcommand', *_bolt_refs[3:6]), _bolt_runtime.push_nesting('execute:commands'), _bolt_runtime.scope() as _bolt_var2:
         while True:
-            _bolt_var0 = True
-            if not _bolt_var0:
-                break
-            _bolt_runtime.commands.extend(_bolt_refs[0].commands)
-    _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[7], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[6], arguments=_bolt_helper_children([*_bolt_refs[3:6], _bolt_helper_replace(_bolt_refs[2], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[1], commands=_bolt_helper_children(_bolt_var1))]))]))])))
-_bolt_var3 = _bolt_helper_replace(_bolt_refs[8], commands=_bolt_helper_children(_bolt_var2))
+            with _bolt_runtime.scope() as _bolt_condition_commands:
+                _bolt_var0 = True
+            with _bolt_helper_loop(_bolt_var0) as (_bolt_loop_overridden, _bolt_var1):
+                _bolt_runtime.commands.extend(_bolt_condition_commands)
+                if not _bolt_loop_overridden:
+                    _bolt_var0 = bool(_bolt_var0)
+                with _bolt_helper_branch(_bolt_var0) as _bolt_condition:
+                    if _bolt_condition:
+                        _bolt_runtime.commands.extend(_bolt_refs[0].commands)
+                        if _bolt_var1:
+                            continue
+            break
+    _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[7], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[6], arguments=_bolt_helper_children([*_bolt_refs[3:6], _bolt_helper_replace(_bolt_refs[2], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[1], commands=_bolt_helper_children(_bolt_var2))]))]))])))
+_bolt_var4 = _bolt_helper_replace(_bolt_refs[8], commands=_bolt_helper_children(_bolt_var3))
 ---
-output = _bolt_var3
+output = _bolt_var4
 ---
 _bolt_refs[0]
 <class 'mecha.ast.AstRoot'>