From 5f8864e251453231fd91298c7c757da2b6566afc Mon Sep 17 00:00:00 2001
From: Erin <erin@erindesu.cz>
Date: Sun, 4 Feb 2024 02:21:14 +0100
Subject: [PATCH] Fixed stack allocation for debug

---
 hbxrt/src/main.rs |  2 +-
 hbxrt/src/mem.rs  | 16 ++++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/hbxrt/src/main.rs b/hbxrt/src/main.rs
index ae86ae6..e7963f1 100644
--- a/hbxrt/src/main.rs
+++ b/hbxrt/src/main.rs
@@ -31,7 +31,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     }
 
     // Allocate stack
-    let mut stack = Box::new(MaybeUninit::<[u8; 1024 * 1024 * 2]>::uninit());
+    let mut stack = unsafe { mem::alloc_stack() };
     eprintln!("[I] Stack allocated at {:p}", stack.as_ptr());
 
     // Load program
diff --git a/hbxrt/src/mem.rs b/hbxrt/src/mem.rs
index 4fae0ce..323e925 100644
--- a/hbxrt/src/mem.rs
+++ b/hbxrt/src/mem.rs
@@ -1,3 +1,5 @@
+use std::alloc::Layout;
+
 use hbvm::mem::{Address, LoadError, Memory, StoreError};
 
 pub struct HostMemory;
@@ -29,3 +31,17 @@ impl Memory for HostMemory {
         unsafe { core::ptr::read(addr.get() as *const T) }
     }
 }
+
+const STACK_SIZE: usize = 2; // MiB
+type Stack = [u8; 1024 * 1024 * STACK_SIZE];
+
+/// Allocate stack of size [`STACK_SIZE`] MiB
+pub unsafe fn alloc_stack() -> Box<Stack> {
+    let layout = Layout::new::<Stack>();
+    let ptr = unsafe { std::alloc::alloc(layout) };
+    if ptr.is_null() {
+        std::alloc::handle_alloc_error(layout);
+    }
+
+    unsafe { Box::from_raw(ptr.cast()) }
+}