From f53a9551f4b462692424a668857b8e9422a2511b Mon Sep 17 00:00:00 2001
From: Luis Michaelis <me@lmichaelis.de>
Date: Fri, 9 Aug 2024 20:52:54 +0200
Subject: [PATCH] [#95] fix(DaedalusVm): fixup the stack after function calls

This prevents issues when functions have too many or too few items on the stack.
---
 include/zenkit/DaedalusVm.hh | 17 ++++------
 src/DaedalusVm.cc            | 60 +++++++++++++++++++++++++++---------
 2 files changed, 52 insertions(+), 25 deletions(-)

diff --git a/include/zenkit/DaedalusVm.hh b/include/zenkit/DaedalusVm.hh
index ede9dd66..28adb66e 100644
--- a/include/zenkit/DaedalusVm.hh
+++ b/include/zenkit/DaedalusVm.hh
@@ -82,6 +82,7 @@ namespace zenkit {
 	struct DaedalusCallStackFrame {
 		DaedalusSymbol const* function;
 		std::uint32_t program_counter;
+		std::uint32_t stack_ptr;
 		std::shared_ptr<DaedalusInstance> context;
 	};
 
@@ -153,20 +154,14 @@ namespace zenkit {
 			unsafe_call(sym);
 
 			if constexpr (std::is_same_v<R, IgnoreReturnValue>) {
-				// clear the stack
-				_m_stack_ptr = 0;
+				if (sym->has_return()) {
+					// Make sure to still pop the return value off of the stack.
+					_m_stack_ptr--;
+				}
 
 				return {};
 			} else if constexpr (!std::is_same_v<R, void>) {
-				auto ret = pop_call_return_value<R>();
-
-				// clear the stack
-				_m_stack_ptr = 0;
-
-				return ret;
-			} else {
-				// clear the stack
-				_m_stack_ptr = 0;
+				return pop_call_return_value<R>();
 			}
 		}
 
diff --git a/src/DaedalusVm.cc b/src/DaedalusVm.cc
index 1b8c4e91..2168d37a 100644
--- a/src/DaedalusVm.cc
+++ b/src/DaedalusVm.cc
@@ -36,19 +36,31 @@ namespace zenkit {
 		~StackGuard() {
 			if (_m_inhibited) return;
 
-			switch (_m_type) {
+			StackGuard::fix(_m_machine, _m_type);
+		}
+
+		/// \brief Inhibits this guard.
+		///
+		/// Calling this function will cause the guard to disengage and thus not push a
+		/// value onto the stack.
+		void inhibit() {
+			_m_inhibited = true;
+		}
+
+		static void fix(DaedalusVm* vm, DaedalusDataType type) {
+			switch (type) {
 			case DaedalusDataType::FLOAT:
-				_m_machine->push_float(0);
+				vm->push_float(0);
 				break;
 			case DaedalusDataType::INT:
 			case DaedalusDataType::FUNCTION:
-				_m_machine->push_int(0);
+				vm->push_int(0);
 				break;
 			case DaedalusDataType::STRING:
-				_m_machine->push_string("");
+				vm->push_string("");
 				break;
 			case DaedalusDataType::INSTANCE:
-				_m_machine->push_instance(nullptr);
+				vm->push_instance(nullptr);
 				break;
 			case DaedalusDataType::VOID:
 			case DaedalusDataType::CLASS:
@@ -57,14 +69,6 @@ namespace zenkit {
 			}
 		}
 
-		/// \brief Inhibits this guard.
-		///
-		/// Calling this function will cause the guard to disengage and thus not push a
-		/// value onto the stack.
-		void inhibit() {
-			_m_inhibited = true;
-		}
-
 	private:
 		DaedalusDataType _m_type;
 		DaedalusVm* _m_machine;
@@ -449,11 +453,39 @@ namespace zenkit {
 	}
 
 	void DaedalusVm::push_call(DaedalusSymbol const* sym) {
-		_m_call_stack.push({sym, _m_pc, _m_instance});
+		auto var_count = this->find_parameters_for_function(sym).size();
+		_m_call_stack.push({sym, _m_pc, _m_stack_ptr - static_cast<uint32_t>(var_count), _m_instance});
 	}
 
 	void DaedalusVm::pop_call() {
 		auto const& call = _m_call_stack.top();
+
+		// First, try to fix up the stack.
+		if (!call.function->has_return()) {
+			// No special logic needed, there are supposed to be no more stack frames for
+			// this function, so just reset the stack for the caller.
+			_m_stack_ptr = call.stack_ptr;
+		} else {
+			auto remaining_locals = _m_stack_ptr - call.stack_ptr;
+			if (remaining_locals == 0) {
+				// The function is supposed to have a return value, but it does not seem to have one.
+				// To fix this, we just insert a default value (either 0, "" or NULL).
+				StackGuard::fix(this, call.function->rtype());
+			} else if (remaining_locals > 1) {
+				// Now we have too many items left on the stack. Remove all of them, except the topmost one,
+				// since that one is supposed to be the return value of the function.
+				DaedalusStackFrame frame = _m_stack[--_m_stack_ptr];
+				_m_stack_ptr = call.stack_ptr;
+				_m_stack[_m_stack_ptr++] = std::move(frame);
+			}
+			// else {
+			//     We have exactly one value to be returned (as indicated by the symbol's return type).
+			//     That means, that the compiler did not mess up the stack management, so we can just
+			//     return that value.
+			// }
+		}
+
+		// Second, reset PC and context, then remove the call stack frame
 		_m_pc = call.program_counter;
 		_m_instance = call.context;
 		_m_call_stack.pop();