fix: prevent interpreter crash when using labeled break/continue in n…

…ested loops
ArgonLang · Oct 5, 2023 · 7bf9184 · 7bf9184
1 parent 212c24f
commit 7bf9184
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 11 deletions.
diff --git a/argon/lang/compiler.cpp b/argon/lang/compiler.cpp
@@ -1258,7 +1258,9 @@ void Compiler::CompileJump(const parser::Unary *jump) {
     if (jump->value != nullptr)
         label = (String *) (((const Unary *) jump->value)->value);
 
-    if ((jb = this->unit_->FindLoop(label)) == nullptr) {
+    unsigned short pops;
+
+    if ((jb = this->unit_->FindLoop(label, pops)) == nullptr) {
         ErrorFormat(kCompilerError[0], "unknown loop label, the loop '%s' cannot be %s",
                     ARGON_RAW_STRING((String *) ((const Unary *) jump->value)->value),
                     jump->token_type == scanner::TokenType::KW_BREAK ? "breaked" : "continued");
@@ -1269,19 +1271,19 @@ void Compiler::CompileJump(const parser::Unary *jump) {
     dst = jb->end;
 
     if (jump->token_type == scanner::TokenType::KW_BREAK) {
-        for (auto i = 0; i < jb->pops; i++)
+        for (auto i = 0; i < pops; i++)
             this->unit_->Emit(vm::OpCode::POP, nullptr);
 
         // Don't decrease the stack size
-        this->unit_->IncrementStack(jb->pops);
+        this->unit_->IncrementStack(pops);
     } else if (jump->token_type == scanner::TokenType::KW_CONTINUE) {
-        if (jb != this->unit_->jstack) {
-            for (auto i = 0; i < jb->pops; i++)
-                this->unit_->Emit(vm::OpCode::POP, nullptr);
+        pops -= jb->pops;
 
-            // Don't decrease the stack size
-            this->unit_->IncrementStack(jb->pops);
-        }
+        for (auto i = 0; i < pops; i++)
+            this->unit_->Emit(vm::OpCode::POP, nullptr);
+
+        // Don't decrease the stack size
+        this->unit_->IncrementStack(pops);
 
         dst = jb->start;
     }

diff --git a/argon/lang/translation_unit.cpp b/argon/lang/translation_unit.cpp
@@ -235,8 +235,12 @@ JBlock *TranslationUnit::JBNew(BasicBlock *start, BasicBlock *end, unsigned shor
     return jb;
 }
 
-JBlock *TranslationUnit::FindLoop(String *label) const {
+JBlock *TranslationUnit::FindLoop(String *label, unsigned short &pops) const {
+    pops = 0;
+
     for (JBlock *block = this->jstack; block != nullptr; block = block->prev) {
+        pops += block->pops;
+
         if (block->type != JBlockType::LOOP)
             continue;
 

diff --git a/argon/lang/translation_unit.h b/argon/lang/translation_unit.h
@@ -78,7 +78,7 @@ namespace argon::lang {
 
         JBlock *JBNew(BasicBlock *start, BasicBlock *end, unsigned short pops);
 
-        JBlock *FindLoop(vm::datatype::String *label) const;
+        JBlock *FindLoop(vm::datatype::String *label, unsigned short &pops) const;
 
         unsigned int ComputeAssemblyLength(unsigned int *out_linfo_sz) const;