Merge pull request #499 from scratchcpp/fix_vm_bugs

Fix some VM bugs

Author: adazem009

src/engine/virtualmachine_p.cpp

@@ -121,6 +121,7 @@ VirtualMachinePrivate::VirtualMachinePrivate(VirtualMachine *vm, Target *target,
     regs = regsVector.data();
     loops.reserve(256);
     callTree.reserve(1024);
+    loopCountTree.reserve(1024);
 
     if (!rng)
         rng = RandomGenerator::instance().get();
@@ -233,10 +234,19 @@ unsigned int *VirtualMachinePrivate::run(unsigned int *pos, bool reset)
         warp = oldState.second;
         callTree.pop_back();
         procedureArgTree.pop_back();
+
         if (procedureArgTree.empty())
             procedureArgs = nullptr;
         else
             procedureArgs = &procedureArgTree.back();
+
+        int loopsToRemove = loops.size() - loopCountTree.back();
+        assert(loopsToRemove >= 0);
+        loopCountTree.pop_back();
+
+        for (int i = 0; i < loopsToRemove; i++)
+            loops.pop_back();
+
         DISPATCH();
     }
 
@@ -781,6 +791,14 @@ unsigned int *VirtualMachinePrivate::run(unsigned int *pos, bool reset)
             DISPATCH(); // this avoids freeing registers after "stopping" a warp script
         }
         FREE_REGS(ret);
+
+        if (atEnd) {
+            if (regCount > 0)
+                std::cout << "warning: VM: " << regCount << " registers were leaked by the script; this is most likely a bug in the VM or in the compiler" << std::endl;
+
+            return pos;
+        }
+
         DISPATCH();
     }
 
@@ -799,6 +817,7 @@ unsigned int *VirtualMachinePrivate::run(unsigned int *pos, bool reset)
             callTree.push_back({ ++pos, warp });
             procedureArgs = nextProcedureArgs;
             nextProcedureArgs = nullptr;
+            loopCountTree.push_back(loops.size());
             pos = procedurePos;
         } else
             pos++;

src/engine/virtualmachine_p.h

@@ -50,6 +50,7 @@ struct VirtualMachinePrivate
         std::vector<std::vector<Value>> procedureArgTree;
         std::vector<Value> *procedureArgs = nullptr;
         std::vector<Value> *nextProcedureArgs = nullptr;
+        std::vector<unsigned int> loopCountTree;
         bool noBreak = true;
         bool warp = false;
         bool stop = false;

test/virtual_machine/virtual_machine_test.cpp

@@ -1773,3 +1773,44 @@ TEST(VirtualMachineTest, NoCrashWhenReadingProcedureArgsAfterStopping)
     ASSERT_EQ(vm.registerCount(), 0);
     ASSERT_TRUE(vm.atEnd());
 }
+
+unsigned int killTest(VirtualMachine *vm)
+{
+    vm->kill();
+    return 1;
+}
+
+TEST(VirtualMachineTest, StopAfterKilled)
+{
+    // Regtest for #497
+    static unsigned int bytecode[] = { OP_START, OP_NULL, OP_EXEC, 0, OP_NULL, OP_NULL, OP_HALT };
+    static BlockFunc functions[] = { &killTest };
+
+    VirtualMachine vm;
+    vm.setBytecode(bytecode);
+    vm.setFunctions(functions);
+
+    vm.run();
+    ASSERT_TRUE(vm.atEnd());
+    ASSERT_EQ(vm.registerCount(), 0);
+}
+
+TEST(VirtualMachineTest, NoCrashAfterCallingProcedureFromLoop)
+{
+    // Regtest for #498
+    static unsigned int bytecode[] = { OP_START, OP_CONST, 0, OP_REPEAT_LOOP, OP_INIT_PROCEDURE, OP_CONST, 1, OP_ADD_ARG, OP_CALL_PROCEDURE, 0, OP_LOOP_END, OP_HALT };
+    static unsigned int procedure[] = { OP_START, OP_WARP, OP_CONST, 0, OP_REPEAT_LOOP, OP_READ_ARG, 0, OP_PRINT, OP_HALT, OP_LOOP_END, OP_HALT };
+    static unsigned int *procedures[] = { procedure };
+    static Value constValues[] = { 2, "test" };
+
+    VirtualMachine vm;
+    vm.setBytecode(bytecode);
+    vm.setProcedures(procedures);
+    vm.setConstValues(constValues);
+
+    ::testing::internal::CaptureStdout();
+    vm.run();
+    ASSERT_TRUE(vm.atEnd());
+    ASSERT_EQ(vm.registerCount(), 0);
+    ASSERT_EQ(::testing::internal::GetCapturedStdout(), "test\ntest\n");
+}