From 2ec38841890113d12251c7122d1d58d85ece30a6 Mon Sep 17 00:00:00 2001
From: wheremyfoodat <44909372+wheremyfoodat@users.noreply.github.com>
Date: Fri, 7 Jul 2023 15:00:18 +0300
Subject: [PATCH] Proper stack management

---
 include/memory.hpp       |  2 +-
 src/core/loader/elf.cpp  |  7 +++++++
 src/core/loader/ncsd.cpp | 21 +++++++++++++++++++++
 src/core/memory.cpp      | 26 +++++++++++++++++++-------
 4 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/include/memory.hpp b/include/memory.hpp
index 5beeaa47..40349b46 100644
--- a/include/memory.hpp
+++ b/include/memory.hpp
@@ -29,7 +29,6 @@ namespace VirtualAddrs {
 		// Stack for main ARM11 thread.
 		// Typically 0x4000 bytes, determined by exheader
 		StackTop = 0x10000000,
-		StackBottom = 0x0FFFC000,
 		DefaultStackSize = 0x4000,
 
 		NormalHeapStart = 0x08000000,
@@ -248,4 +247,5 @@ public:
 	u32 getUsedUserMem() { return usedUserMemory; }
 
 	void setVRAM(u8* pointer) { vram = pointer; }
+	bool allocateMainThreadStack(u32 size);
 };
\ No newline at end of file
diff --git a/src/core/loader/elf.cpp b/src/core/loader/elf.cpp
index 54128c64..62802953 100644
--- a/src/core/loader/elf.cpp
+++ b/src/core/loader/elf.cpp
@@ -14,6 +14,13 @@ std::optional<u32> Memory::loadELF(std::ifstream& file) {
 		return std::nullopt;
 	}
 
+    // Allocate stack space. For ELFs we use the default stack size, which is 16KB
+	if (!allocateMainThreadStack(VirtualAddrs::DefaultStackSize)) {
+		// Should be unreachable
+		printf("Failed to allocate stack space for ELF file\n");
+		return std::nullopt;
+	}
+
     auto segNum = reader.segments.size();
     printf("Number of segments: %d\n", segNum);
     printf(" #  Perms       Vaddr           File Size       Mem Size\n");
diff --git a/src/core/loader/ncsd.cpp b/src/core/loader/ncsd.cpp
index 21737ce2..2655cf72 100644
--- a/src/core/loader/ncsd.cpp
+++ b/src/core/loader/ncsd.cpp
@@ -9,6 +9,27 @@ bool Memory::mapCXI(NCSD& ncsd, NCCH& cxi) {
 	printf("Text address = %08X, size = %08X\n", cxi.text.address, cxi.text.size);
 	printf("Rodata address = %08X, size = %08X\n", cxi.rodata.address, cxi.rodata.size);
 	printf("Data address = %08X, size = %08X\n", cxi.data.address, cxi.data.size);
+	printf("Stack size: %08X\n", cxi.stackSize);
+
+	if (!isAligned(cxi.stackSize)) {
+		Helpers::warn("CXI has a suspicious stack size of %08X which is not a multiple of 4KB", cxi.stackSize);
+	}
+
+	// Round up the size of the CXI stack size to a page (4KB) boundary, as the OS can only allocate memory this way
+	u32 stackSize = (cxi.stackSize + pageSize - 1) & -pageSize;
+
+	if (stackSize > 512_KB) {
+		// TODO: Figure out the actual max stack size
+		Helpers::warn("CXI stack size is %08X which seems way too big. Clamping to 512KB", stackSize);
+		stackSize = 512_KB;
+	}
+
+	// Allocate stack
+	if (!allocateMainThreadStack(stackSize)) {
+		// Should be unreachable
+		printf("Failed to allocate stack for CXI partition. Requested stack size: %08X\n", stackSize);
+		return false;
+	}
 
 	// Map code file to memory
 	auto& code = cxi.codeFile;
diff --git a/src/core/memory.cpp b/src/core/memory.cpp
index dfa155a2..abea7606 100644
--- a/src/core/memory.cpp
+++ b/src/core/memory.cpp
@@ -28,16 +28,16 @@ void Memory::reset() {
 		writeTable[i] = 0;
 	}
 
-	// Map stack pages as R/W
-	// We have 16KB for the stack, so we allocate the last 16KB of APPLICATION FCRAM for the stack
-	u32 basePaddrForStack = FCRAM_APPLICATION_SIZE - VirtualAddrs::DefaultStackSize;
-	allocateMemory(VirtualAddrs::StackBottom, basePaddrForStack, VirtualAddrs::DefaultStackSize, true);
+	// Map (32 * 4) KB of FCRAM before the stack for the TLS of each thread
+	std::optional<u32> tlsBaseOpt = findPaddr(32 * 4_KB);
+	if (!tlsBaseOpt.has_value()) {  // Should be unreachable but still good to have
+		Helpers::panic("Failed to allocate memory for thread-local storage");
+	}
 
-	// And map (4 * 32)KB of FCRAM before the stack for the TLS of each thread
-	u32 basePaddrForTLS = basePaddrForStack;
+	u32 basePaddrForTLS = tlsBaseOpt.value();
 	for (int i = 0; i < appResourceLimits.maxThreads; i++) {
 		u32 vaddr = VirtualAddrs::TLSBase + i * VirtualAddrs::TLSSize;
-		basePaddrForTLS -= VirtualAddrs::TLSSize;
+		basePaddrForTLS += VirtualAddrs::TLSSize;
 		allocateMemory(vaddr, basePaddrForTLS, VirtualAddrs::TLSSize, true);
 	}
 
@@ -59,6 +59,18 @@ void Memory::reset() {
 	}
 }
 
+bool Memory::allocateMainThreadStack(u32 size) {
+	// Map stack pages as R/W
+	std::optional<u32> basePaddr = findPaddr(size);
+	if (!basePaddr.has_value()) {  // Should also be unreachable but still good to have
+		return false;
+	}
+
+	const u32 stackBottom = VirtualAddrs::StackTop - size;
+	std::optional<u32> result = allocateMemory(stackBottom, basePaddr.value(), size, true); // Should never be nullopt
+	return result.has_value();
+}
+
 u8 Memory::read8(u32 vaddr) {
 	const u32 page = vaddr >> pageShift;
 	const u32 offset = vaddr & pageMask;