feat: error when using unsupported statements in memo

bolt/parse.py

@@ -16,6 +16,7 @@
     "AssignmentStatementParser",
     "EscapeAnalysisParser",
     "EscapeAnalysisResolver",
+    "MemoRootParser",
     "MemoHandler",
     "parse_decorator",
     "AssignmentTargetParser",
@@ -320,7 +321,7 @@ def get_bolt_parsers(
                 ]
             )
         ),
-        "bolt:memo_root": EscapeAnalysisParser(delegate("nested_root")),
+        "bolt:memo_root": MemoRootParser(EscapeAnalysisParser(delegate("nested_root"))),
         "bolt:del_target": parse_del_target,
         "bolt:interpolation": BuiltinCallRestriction(
             PrimaryParser(delegate("bolt:identifier"), truncate=True),
@@ -608,6 +609,7 @@ def create_bolt_root_parser(parser: Parser, macro_handler: "MacroHandler"):
     parser = LexicalScopeConstraint(
         parser,
         type=FunctionScope,
+        flags={"memo": False},
         command_identifiers={
             "return",
             "return:value",
@@ -952,6 +954,17 @@ def resolve(self, node: AstRoot) -> AstRoot:
         return node
 
 
+@dataclass
+class MemoRootParser:
+    """Parse memo root."""
+
+    parser: Parser
+
+    def __call__(self, stream: TokenStream) -> Any:
+        with stream.provide(loop=False, memo=True):
+            return self.parser(stream)
+
+
 @dataclass
 class MemoHandler:
     """Handle memo."""
@@ -2021,6 +2034,8 @@ def resolve_deferred(self, node: AstRoot, stream: TokenStream) -> AstRoot:
                 self.macro_handler.flush_pending_macros(stream)
 
                 deferred_stream = deferred_root.stream
+                deferred_stream.data["loop"] = False
+                deferred_stream.data["memo"] = False
                 deferred_stream.data["local_spec"] = False
                 deferred_stream.data["spec"] = get_stream_spec(stream)
                 deferred_stream.data["macro_scope"] = get_stream_macro_scope(stream)
@@ -2058,20 +2073,26 @@ class LexicalScopeConstraint:
 
     parser: Parser
     type: Type[LexicalScope]
+    flags: Dict[str, bool]
     command_identifiers: Set[str]
 
     def __call__(self, stream: TokenStream) -> AstRoot:
         node = self.parser(stream)
 
         lexical_scope = get_stream_lexical_scope(stream)
-        if isinstance(lexical_scope, self.type):
+        if isinstance(lexical_scope, self.type) and all(
+            stream.data.get(flag, False) == enabled
+            for flag, enabled in self.flags.items()
+        ):
             return node
 
         if isinstance(node, AstRoot):
             for command in node.commands:
                 if command.identifier in self.command_identifiers:
                     name, _, _ = command.identifier.partition(":")
-                    exc = InvalidSyntax(f'Can only use "{name}" in functions.')
+                    exc = InvalidSyntax(
+                        f'Statement "{name}" is not allowed in this context.'
+                    )
                     raise set_location(exc, command)
 
         return node

tests/resources/bolt_examples.mcfunction

@@ -1273,3 +1273,15 @@ memo thing:
     thing += 1
 thing *= 7
 say thing
+###
+for i in range(10):
+    def f():
+        break
+###
+for i in range(10):
+    memo:
+        break
+###
+def f():
+    memo:
+        return 9999

tests/snapshots/bolt__parse_119__0.txt

@@ -1,5 +1,5 @@
 if True:
-#>ERROR  Can only use "yield" in functions.
+#>ERROR  Statement "yield" is not allowed in this context.
 # line 2, column 5
 #      1 |  if True:
 #      2 |      yield

tests/snapshots/bolt__parse_260__0.txt

@@ -1,5 +1,5 @@
 class Foo:
-#>ERROR  Can only use "return" in functions.
+#>ERROR  Statement "return" is not allowed in this context.
 # line 2, column 5
 #      1 |  class Foo:
 #      2 |      return

tests/snapshots/bolt__parse_261__0.txt

@@ -1,6 +1,6 @@
 def f():
     class Foo:
-#>ERROR  Can only use "return" in functions.
+#>ERROR  Statement "return" is not allowed in this context.
 # line 3, column 9
 #      2 |      class Foo:
 #      3 |          return

tests/snapshots/bolt__parse_347__0.txt

@@ -0,0 +1,8 @@
+for i in range(10):
+    def f():
+#>ERROR  Can only use "break" in loops.
+# line 3, column 9
+#      2 |      def f():
+#      3 |          break
+#        :          ^^^^^
+        break

tests/snapshots/bolt__parse_348__0.txt

@@ -0,0 +1,8 @@
+for i in range(10):
+    memo:
+#>ERROR  Can only use "break" in loops.
+# line 3, column 9
+#      2 |      memo:
+#      3 |          break
+#        :          ^^^^^
+        break

tests/snapshots/bolt__parse_349__0.txt

@@ -0,0 +1,8 @@
+def f():
+    memo:
+#>ERROR  Statement "return" is not allowed in this context.
+# line 3, column 9
+#      2 |      memo:
+#      3 |          return 9999
+#        :          ^^^^^^^^^^^
+        return 9999

tests/snapshots/bolt__parse_85__0.txt

@@ -1,4 +1,4 @@
-#>ERROR  Can only use "return" in functions.
+#>ERROR  Statement "return" is not allowed in this context.
 # line 1, column 1
 #      1 |  return "hello"
 #        :  ^^^^^^^^^^^^^^