fix(rivetc.codegen): fix break with defer statement

Author: StunxFS

lib/core/src/test_runner.ri renamed to lib/core/src/TestRunner.ri

@@ -14,28 +14,28 @@ extern (Rivet) {
     static testRunner: TestRunner;
 }
 
-public func bold(msg: string) -> string {
+func bold(msg: string) -> string {
     if !supportColor {
         return msg;
     }
     return "\x1b[1m{}\x1b[22m".fmt(msg);
 }
 
-public func green(msg: string) -> string {
+func green(msg: string) -> string {
     if !supportColor {
         return msg;
     }
     return "\x1b[32m{}\x1b[39m".fmt(msg);
 }
 
-public func red(msg: string) -> string {
+func red(msg: string) -> string {
     if !supportColor {
         return msg;
     }
     return "\x1b[31m{}\x1b[39m".fmt(msg);
 }
 
-public func yellow(msg: string) -> string {
+func yellow(msg: string) -> string {
     if !supportColor {
         return msg;
     }
@@ -75,7 +75,7 @@ struct Test {
     mut err: Throwable;
 }
 
-public struct TestRunner {
+struct TestRunner {
     tests: []mut Test;
     mut ok_tests: uint64;
     mut fail_tests: uint64;

rivetc/src/ast.py

@@ -312,14 +312,15 @@ def __init__(self, scope, lefts, right, pos):
         self.pos = pos
 
 class WhileStmt:
-    def __init__(self, cond, continue_expr, stmt, else_stmt, is_inf, pos):
+    def __init__(self, scope, cond, continue_expr, stmt, else_stmt, is_inf, pos):
         self.cond = cond
         self.is_inf = is_inf
         self.continue_expr = continue_expr
         self.has_continue_expr = not isinstance(continue_expr, EmptyExpr)
         self.stmt = stmt
         self.else_stmt = else_stmt
         self.has_else_stmt = else_stmt != None
+        self.scope = scope
         self.pos = pos
 
 class ForStmt:

rivetc/src/codegen/__init__.py

@@ -139,6 +139,7 @@ def __init__(self, comp):
 
         self.loop_entry_label = ""
         self.loop_exit_label = ""
+        self.loop_scope = None
         self.while_continue_expr = None
 
     def gen_source_files(self, source_files):
@@ -472,6 +473,8 @@ def gen_stmts(self, stmts):
 
     def gen_stmt(self, stmt):
         if isinstance(stmt, ast.ForStmt):
+            old_loop_scope = self.loop_scope
+            self.loop_scope = stmt.scope
             old_while_continue_expr = self.while_continue_expr
             old_entry_label = self.loop_entry_label
             old_exit_label = self.loop_exit_label
@@ -537,8 +540,11 @@ def gen_stmt(self, stmt):
             self.cur_fn.add_label(self.loop_exit_label)
             self.loop_entry_label = old_entry_label
             self.loop_exit_label = old_exit_label
+            self.loop_scope = old_loop_scope
             self.while_continue_expr = old_while_continue_expr
         elif isinstance(stmt, ast.WhileStmt):
+            old_loop_scope = self.loop_scope
+            self.loop_scope = stmt.scope
             old_while_continue_expr = self.while_continue_expr
             old_entry_label = self.loop_entry_label
             old_exit_label = self.loop_exit_label
@@ -587,6 +593,7 @@ def gen_stmt(self, stmt):
             self.cur_fn.add_label(self.loop_exit_label)
             self.loop_entry_label = old_entry_label
             self.loop_exit_label = old_exit_label
+            self.loop_scope = old_loop_scope
             self.while_continue_expr = old_while_continue_expr
         elif isinstance(stmt, ast.StaticDeclStmt):
             if len(stmt.lefts) == 1:
@@ -2331,7 +2338,8 @@ def gen_expr(self, expr, custom_tmp = None):
                 self.cur_fn.add_br(self.loop_entry_label)
             else:
                 self.gen_defer_stmts(
-                    scope = expr.scope, run_defer_previous = True
+                    scope = expr.scope, run_defer_previous = True,
+                    scope_limit = self.loop_scope
                 )
                 self.cur_fn.add_br(self.loop_exit_label)
             return ir.Skip()
@@ -2467,7 +2475,7 @@ def gen_defer_stmt_vars(self, defer_stmts):
 
     def gen_defer_stmts(
         self, gen_errdefer = False, last_ret_was_err = None, scope = None,
-        run_defer_previous = False
+        run_defer_previous = False, scope_limit = None
     ):
         for i in range(len(self.cur_fn_defer_stmts) - 1, -1, -1):
             defer_stmt = self.cur_fn_defer_stmts[i]
@@ -2477,11 +2485,15 @@ def gen_defer_stmts(
             ):
                 continue
             if defer_stmt.mode == ast.DeferMode.ERROR and not gen_errdefer:
-                # Should be run only when an error occurs
+                # should be run only when an error occurs
                 continue
             if defer_stmt.mode == ast.DeferMode.SUCCESS and gen_errdefer:
-                # Should not be run in case an error occurs
+                # should not be run in case an error occurs
                 continue
+            if scope_limit != None and defer_stmt.scope.start == scope_limit.start:
+                # do not generate `defer` that are in a scope greater than
+                # the established limit, this is used by `break`
+                break
             defer_start = self.cur_fn.local_name()
             defer_end = self.cur_fn.local_name()
             self.cur_fn.add_comment(

rivetc/src/parser.py

@@ -576,7 +576,7 @@ def parse_stmt(self):
             if self.accept(Kind.KwElse):
                 else_stmt = self.parse_stmt()
             return ast.WhileStmt(
-                cond, continue_expr, stmt, else_stmt, is_inf, pos
+                self.scope, cond, continue_expr, stmt, else_stmt, is_inf, pos
             )
         elif self.accept(Kind.KwFor):
             pos = self.prev_tok.pos

tests/valid/src/defer_stmt.ri

@@ -111,3 +111,16 @@ test "`defer(success)` with throwed error" {
     throw_and_defer_error(err_defer) catch {};
     @assert(err_defer.i == 4);
 }
+
+test "`defer` with `break`" {
+    mut a := 0;
+    mut i := 0;
+    defer @assert(a == 12);
+    defer a += 2;
+    while i < 20 : i += 1 {
+        defer a = 10;
+        if i == 10 {
+            break;
+        }
+    }
+}