fix #498: Store a tree of running loop count

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();
     }
 
@@ -807,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

@@ -1794,3 +1794,23 @@ TEST(VirtualMachineTest, StopAfterKilled)
     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");
+}