Merge pull request #381 from scratchcpp/fix_if_and_loop_skipping

Fix skipping of if statements and loops

Author: adazem009

src/engine/virtualmachine_p.cpp

@@ -236,25 +236,40 @@ unsigned int *VirtualMachinePrivate::run(unsigned int *pos, bool reset)
     checkpoint = pos - 1;
     DISPATCH();
 
-do_if:
+do_if : {
     if (!READ_LAST_REG()->toBool()) {
         unsigned int ifCounter = 1;
-        while ((*pos != OP_ELSE && *pos != OP_ENDIF) || (ifCounter > 0)) {
+        while (!((*pos == OP_ELSE && ifCounter == 1) || (*pos == OP_ENDIF && ifCounter == 0))) {
             pos += instruction_arg_count[*pos++];
 
-            if (*pos == OP_IF)
+            if ((*pos == OP_IF) || (*pos == OP_FOREVER_LOOP) || (*pos == OP_REPEAT_LOOP) || (*pos == OP_UNTIL_LOOP))
                 ifCounter++;
-            else if ((*pos == OP_ELSE) || (*pos == OP_ENDIF))
+            else if ((*pos == OP_ENDIF) || (*pos == OP_LOOP_END)) {
+                assert(ifCounter > 0);
                 ifCounter--;
+            }
         }
     }
     FREE_REGS(1);
     DISPATCH();
+}
 
-do_else:
-    while (*pos != OP_ENDIF)
+do_else : {
+    unsigned int ifCounter = 1;
+    while (!(*pos == OP_ENDIF && ifCounter == 0)) {
         pos += instruction_arg_count[*pos++];
 
+        if ((*pos == OP_IF) || (*pos == OP_FOREVER_LOOP) || (*pos == OP_REPEAT_LOOP) || (*pos == OP_UNTIL_LOOP))
+            ifCounter++;
+        else if ((*pos == OP_ENDIF) || (*pos == OP_LOOP_END)) {
+            assert(ifCounter > 0);
+            ifCounter--;
+        }
+
+        assert(!(*pos == OP_ELSE && ifCounter == 1));
+    }
+}
+
 do_endif:
     DISPATCH();
 
@@ -275,10 +290,12 @@ unsigned int *VirtualMachinePrivate::run(unsigned int *pos, bool reset)
         while ((*loopEnd != OP_LOOP_END) || (loopCounter > 0)) {
             loopEnd += instruction_arg_count[*loopEnd++];
 
-            if ((*loopEnd == OP_FOREVER_LOOP) || (*loopEnd == OP_REPEAT_LOOP) || (*loopEnd == OP_UNTIL_LOOP))
+            if ((*loopEnd == OP_IF) || (*loopEnd == OP_FOREVER_LOOP) || (*loopEnd == OP_REPEAT_LOOP) || (*loopEnd == OP_UNTIL_LOOP))
                 loopCounter++;
-            else if (*loopEnd == OP_LOOP_END)
+            else if ((*loopEnd == OP_ENDIF) || (*loopEnd == OP_LOOP_END)) {
+                assert(loopCounter > 0);
                 loopCounter--;
+            }
         }
         pos = loopEnd;
     } else {
@@ -317,8 +334,17 @@ do_repeat_loop_index1 : {
         pos = loopStart;
     } else {
         pos = loopStart;
-        while (*pos != OP_LOOP_END)
+        unsigned int loopCounter = 1;
+        while ((*pos != OP_LOOP_END) || (loopCounter > 0)) {
             pos += instruction_arg_count[*pos++];
+
+            if ((*pos == OP_IF) || (*pos == OP_FOREVER_LOOP) || (*pos == OP_REPEAT_LOOP) || (*pos == OP_UNTIL_LOOP))
+                loopCounter++;
+            else if ((*pos == OP_ENDIF) || (*pos == OP_LOOP_END)) {
+                assert(loopCounter > 0);
+                loopCounter--;
+            }
+        }
     }
     FREE_REGS(1);
     DISPATCH();

test/virtual_machine/virtual_machine_test.cpp

@@ -1626,3 +1626,36 @@ TEST(VirtualMachineTest, NoCrashInNestedIfStatementsWithLoop)
     vm.run();
     ASSERT_EQ(vm.registerCount(), 0);
 }
+
+TEST(VirtualMachineTest, NoCrashInNestedIfStatementsWithLoopAndIfElse)
+{
+    // Regtest for #378
+    static unsigned int bytecode[] = { OP_START, OP_NULL, OP_IF, OP_NULL, OP_REPEAT_LOOP, OP_NULL, OP_IF, OP_ELSE, OP_ENDIF, OP_LOOP_END, OP_ENDIF, OP_HALT };
+
+    VirtualMachine vm(nullptr, nullptr, nullptr);
+    vm.setBytecode(bytecode);
+    vm.run();
+    ASSERT_EQ(vm.registerCount(), 0);
+}
+
+TEST(VirtualMachineTest, NoCrashInNestedLoopsInRepeatUntilLoop)
+{
+    // Regtest for #379
+    static unsigned int bytecode[] = { OP_START, OP_NULL, OP_NOT, OP_IF, OP_ELSE, OP_NULL, OP_REPEAT_LOOP, OP_NULL, OP_IF, OP_ENDIF, OP_LOOP_END, OP_ENDIF, OP_HALT };
+
+    VirtualMachine vm(nullptr, nullptr, nullptr);
+    vm.setBytecode(bytecode);
+    vm.run();
+    ASSERT_EQ(vm.registerCount(), 0);
+}
+
+TEST(VirtualMachineTest, NoCrashInNestedLoopsInIfElseStatements)
+{
+    // Regtest for #380
+    static unsigned int bytecode[] = { OP_START, OP_UNTIL_LOOP, OP_NULL, OP_NOT, OP_BEGIN_UNTIL_LOOP, OP_NULL, OP_REPEAT_LOOP, OP_LOOP_END, OP_LOOP_END, OP_HALT };
+
+    VirtualMachine vm(nullptr, nullptr, nullptr);
+    vm.setBytecode(bytecode);
+    vm.run();
+    ASSERT_EQ(vm.registerCount(), 0);
+}