From 0fe213765b89f18e4b70abf1dbbdfd6e1485fc9e Mon Sep 17 00:00:00 2001
From: Magnus Norddahl <dpjudas@users.noreply.github.com>
Date: Sun, 18 Nov 2018 13:49:19 +0100
Subject: [PATCH] - only allocate stack space for vmframe and call returns when
 we need them

---
 src/scripting/vm/jit.cpp      | 42 +++++++++++++++++++++++++++++------
 src/scripting/vm/jit_call.cpp |  8 ++++---
 src/scripting/vm/jitintern.h  | 13 +++++++++--
 3 files changed, 51 insertions(+), 12 deletions(-)

diff --git a/src/scripting/vm/jit.cpp b/src/scripting/vm/jit.cpp
index b8b7240a1..e05a4881f 100644
--- a/src/scripting/vm/jit.cpp
+++ b/src/scripting/vm/jit.cpp
@@ -436,6 +436,37 @@ void JitCompiler::BindLabels()
 	cc.setCursor(cursor);
 }
 
+void JitCompiler::CheckVMFrame()
+{
+	if (!vmframeAllocated)
+	{
+		auto cursor = cc.getCursor();
+		cc.setCursor(vmframeCursor);
+
+		auto vmstack = cc.newStack(sfunc->StackSize, 16, "vmstack");
+		vmframe = cc.newIntPtr("vmframe");
+		cc.lea(vmframe, vmstack);
+
+		cc.setCursor(cursor);
+		vmframeAllocated = true;
+	}
+}
+
+asmjit::X86Gp JitCompiler::GetCallReturns()
+{
+	if (!callReturnsAllocated)
+	{
+		auto cursor = cc.getCursor();
+		cc.setCursor(callReturnsCursor);
+		auto stackalloc = cc.newStack(sizeof(VMReturn) * MAX_RETURNS, alignof(VMReturn), "stackalloc");
+		callReturns = cc.newIntPtr("callReturns");
+		cc.lea(callReturns, stackalloc);
+		cc.setCursor(cursor);
+		callReturnsAllocated = true;
+	}
+	return callReturns;
+}
+
 void JitCompiler::Setup()
 {
 	using namespace asmjit;
@@ -466,9 +497,7 @@ void JitCompiler::Setup()
 	cc.setArg(3, ret);
 	cc.setArg(4, numret);
 
-	auto stackalloc = cc.newStack(sizeof(VMReturn) * MAX_RETURNS, alignof(VMReturn), "stackalloc");
-	callReturns = cc.newIntPtr("callReturns");
-	cc.lea(callReturns, stackalloc);
+	callReturnsCursor = cc.getCursor();
 
 	konstd = sfunc->KonstD;
 	konstf = sfunc->KonstF;
@@ -492,8 +521,6 @@ void JitCompiler::SetupFrame()
 	offsetD = offsetA + (int)(sfunc->NumRegA * sizeof(void*));
 	offsetExtra = (offsetD + (int)(sfunc->NumRegD * sizeof(int32_t)) + 15) & ~15;
 
-	vmframe = cc.newIntPtr("vmframe");
-
 	if (sfunc->SpecialInits.Size() == 0 && sfunc->NumRegS == 0)
 	{
 		SetupSimpleFrame();
@@ -510,8 +537,7 @@ void JitCompiler::SetupSimpleFrame()
 
 	// This is a simple frame with no constructors or destructors. Allocate it on the stack ourselves.
 
-	auto vmstack = cc.newStack(sfunc->StackSize, 16, "vmstack");
-	cc.lea(vmframe, vmstack);
+	vmframeCursor = cc.getCursor();
 
 	int argsPos = 0;
 	int regd = 0, regf = 0, rega = 0;
@@ -592,8 +618,10 @@ void JitCompiler::SetupFullVMFrame()
 	allocFrame->setArg(1, args);
 	allocFrame->setArg(2, numargs);
 
+	vmframe = cc.newIntPtr("vmframe");
 	cc.mov(vmframe, x86::ptr(stack)); // stack->Blocks
 	cc.mov(vmframe, x86::ptr(vmframe, VMFrameStack::OffsetLastFrame())); // Blocks->LastFrame
+	vmframeAllocated = true;
 
 	for (int i = 0; i < sfunc->NumRegD; i++)
 		cc.mov(regD[i], x86::dword_ptr(vmframe, offsetD + i * sizeof(int32_t)));
diff --git a/src/scripting/vm/jit_call.cpp b/src/scripting/vm/jit_call.cpp
index 56147c902..92f596693 100644
--- a/src/scripting/vm/jit_call.cpp
+++ b/src/scripting/vm/jit_call.cpp
@@ -71,6 +71,8 @@ void JitCompiler::EmitDoCall(asmjit::X86Gp vmfunc, VMFunction *target)
 		simpleFrameTarget = starget->SpecialInits.Size() == 0 && starget->NumRegS == 0;
 	}
 
+	CheckVMFrame();
+
 	int numparams = StoreCallParams(simpleFrameTarget);
 	if (numparams != B)
 		I_FatalError("OP_CALL parameter count does not match the number of preceding OP_PARAM instructions");
@@ -106,7 +108,7 @@ void JitCompiler::EmitScriptCall(asmjit::X86Gp vmfunc, asmjit::X86Gp paramsptr)
 	call->setArg(0, vmfunc);
 	call->setArg(1, paramsptr);
 	call->setArg(2, Imm(B));
-	call->setArg(3, callReturns);
+	call->setArg(3, GetCallReturns());
 	call->setArg(4, Imm(C));
 }
 
@@ -346,7 +348,7 @@ void JitCompiler::FillReturns(const VMOP *retval, int numret)
 			break;
 		}
 
-		cc.mov(x86::ptr(callReturns, i * sizeof(VMReturn) + myoffsetof(VMReturn, Location)), regPtr);
-		cc.mov(x86::byte_ptr(callReturns, i * sizeof(VMReturn) + myoffsetof(VMReturn, RegType)), type);
+		cc.mov(x86::ptr(GetCallReturns(), i * sizeof(VMReturn) + myoffsetof(VMReturn, Location)), regPtr);
+		cc.mov(x86::byte_ptr(GetCallReturns(), i * sizeof(VMReturn) + myoffsetof(VMReturn, RegType)), type);
 	}
 }
diff --git a/src/scripting/vm/jitintern.h b/src/scripting/vm/jitintern.h
index d417a100a..59004cc39 100644
--- a/src/scripting/vm/jitintern.h
+++ b/src/scripting/vm/jitintern.h
@@ -237,9 +237,18 @@ private:
 	int offsetA;
 	int offsetD;
 	int offsetExtra;
-	asmjit::X86Gp vmframe;
-	int NumParam = 0; // Actually part of vmframe (f->NumParam), but nobody seems to read that?
+
 	TArray<const VMOP *> ParamOpcodes;
+
+	void CheckVMFrame();
+	asmjit::X86Gp GetCallReturns();
+
+	bool vmframeAllocated = false;
+	asmjit::CBNode *vmframeCursor = nullptr;
+	asmjit::X86Gp vmframe;
+
+	bool callReturnsAllocated = false;
+	asmjit::CBNode *callReturnsCursor = nullptr;
 	asmjit::X86Gp callReturns;
 
 	const int *konstd;