fix(zscript): stack offset issue when 'return'ing inside a 'loop()'

ZQuestClassic · Apr 5, 2024 · dc6c350 · dc6c350
1 parent 1f010f2
commit dc6c350
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 5 deletions.
diff --git a/src/parser/BuildVisitors.cpp b/src/parser/BuildVisitors.cpp
@@ -82,6 +82,16 @@ class MiniStackMgr
 			--q;
 		return true;
 	}
+	uint pop_all()
+	{
+		uint sz = peekinds.size();
+		peekinds.clear();
+		return sz;
+	}
+	uint count()
+	{
+		return peekinds.size();
+	}
 };
 #define VISIT_PUSH(node, ind) \
 do { \
@@ -1160,7 +1170,7 @@ void BuildOpcodes::caseStmtRangeLoop(ASTStmtRangeLoop &host, void *param)
 	addOpcode(new ONoOp(loopstart));
 
 	//run the inside of the loop.
-	push_break(loopend, arrayRefs.size());
+	push_break(loopend, arrayRefs.size(), mgr.count());
 	push_cont(loopcont, arrayRefs.size());
 
 	targ_sz = commentTarget();
@@ -1348,6 +1358,8 @@ void BuildOpcodes::caseStmtRangeLoop(ASTStmtRangeLoop &host, void *param)
 
 	scope = scope->getParent();
 
+	addOpcode(new OPopArgsRegister(new VarArgument(NUL), new LiteralArgument(mgr.pop_all())));
+
 	if(host.hasElse())
 	{
 		addOpcode(new ONoOp(elselabel));
@@ -1357,8 +1369,6 @@ void BuildOpcodes::caseStmtRangeLoop(ASTStmtRangeLoop &host, void *param)
 	}
 
 	addOpcode(new ONoOp(loopend));
-	while(mgr.pop())
-		addOpcode(new OPopRegister(new VarArgument(NUL)));
 }
 
 void BuildOpcodes::caseStmtWhile(ASTStmtWhile &host, void *param)
@@ -1511,6 +1521,8 @@ void BuildOpcodes::caseStmtDo(ASTStmtDo &host, void *param)
 void BuildOpcodes::caseStmtReturn(ASTStmtReturn&, void*)
 {
 	deallocateRefsUntilCount(0);
+	if(uint pops = scope_pops_count())
+		pop_params(pops);
 	addOpcode(new OGotoImmediate(new LabelArgument(returnlabelid)));
 	commentBack("return (Void)");
 }
@@ -1523,6 +1535,8 @@ void BuildOpcodes::caseStmtReturnVal(ASTStmtReturnVal &host, void *param)
 	INITC_INIT();
 	INITC_DEALLOC();
 	deallocateRefsUntilCount(0);
+	if(uint pops = scope_pops_count())
+		pop_params(pops);
 	addOpcode(new OGotoImmediate(new LabelArgument(returnlabelid)));
 	commentStartEnd(targ_sz,"return");
 }
@@ -1538,6 +1552,8 @@ void BuildOpcodes::caseStmtBreak(ASTStmtBreak &host, void *)
 	int32_t refcount = breakRefCounts.at(breakRefCounts.size()-host.breakCount);
 	int32_t breaklabel = breaklabelids.at(breaklabelids.size()-host.breakCount);
 	deallocateRefsUntilCount(refcount);
+	if(uint pops = scope_pops_back(host.breakCount))
+		pop_params(pops);
 	addOpcode(new OGotoImmediate(new LabelArgument(breaklabel)));
 	commentBack(fmt::format("break {};",host.breakCount));
 	inc_break(host.breakCount);
@@ -1555,6 +1571,8 @@ void BuildOpcodes::caseStmtContinue(ASTStmtContinue &host, void *)
 	int32_t refcount = continueRefCounts.at(continueRefCounts.size()-host.contCount);
 	int32_t contlabel = continuelabelids.at(continuelabelids.size()-host.contCount);
 	deallocateRefsUntilCount(refcount);
+	if(uint pops = scope_pops_back(host.contCount-1))
+		pop_params(pops);
 	addOpcode(new OGotoImmediate(new LabelArgument(contlabel)));
 	commentBack(fmt::format("continue {};",host.contCount));
 	inc_cont(host.contCount);
@@ -4150,6 +4168,12 @@ optional<int> BuildOpcodes::eatSetCompare()
 	return nullopt;
 }
 
+void BuildOpcodes::pop_params(uint count)
+{
+	if(count)
+		addOpcode(new OPopArgsRegister(new VarArgument(NUL), new LiteralArgument(count)));
+}
+
 /////////////////////////////////////////////////////////////////////////////////
 // LValBOHelper
 

diff --git a/src/parser/BuildVisitors.h b/src/parser/BuildVisitors.h
@@ -139,6 +139,7 @@ namespace ZScript
 		std::vector<uint> break_to_counts;
 		std::vector<uint> continue_past_counts;
 		std::vector<uint> continue_to_counts;
+		std::vector<uint> scope_allocations;
 		std::list<int32_t> arrayRefs;
 		// Stack of opcode targets. Only the latest is used.
 		std::vector<std::vector<std::shared_ptr<Opcode>>*> opcodeTargets;
@@ -166,13 +167,14 @@ namespace ZScript
 					++continue_to_counts[q];
 			}
 		}
-		void push_break(int32_t id, int32_t count)
+		void push_break(int32_t id, int32_t count, uint scope_pops = 0)
 		{
 			++break_depth;
 			breaklabelids.push_back(id);
 			breakRefCounts.push_back(count);
 			break_past_counts.push_back(0);
 			break_to_counts.push_back(0);
+			scope_allocations.emplace_back(scope_pops);
 		}
 		void push_cont(int32_t id, int32_t count)
 		{
@@ -189,6 +191,7 @@ namespace ZScript
 			breakRefCounts.pop_back();
 			break_past_counts.pop_back();
 			break_to_counts.pop_back();
+			scope_allocations.pop_back();
 		}
 		void pop_cont()
 		{
@@ -198,9 +201,26 @@ namespace ZScript
 			continue_past_counts.pop_back();
 			continue_to_counts.pop_back();
 		}
+		uint scope_pops_back(uint break_count)
+		{
+			uint ret = 0;
+			for(auto it = scope_allocations.rbegin(); break_count && it != scope_allocations.rend(); ++it)
+			{
+				ret += *it;
+				--break_count;
+			}
+			return ret;
+		}
+		uint scope_pops_count()
+		{
+			uint ret = 0;
+			for(uint v : scope_allocations)
+				ret += v;
+			return ret;
+		}
 
 		// Helper Functions.
-
+		void pop_params(uint count);
 		// For when ASTDataDecl is for a single variable.
 		void buildVariable(ASTDataDecl& host, OpcodeContext& context);
 		// For when ASTDataDecl is an initialized array.