From 778de3cfe0aa3498627f2ac027929ad9d74e1d61 Mon Sep 17 00:00:00 2001
From: Vincent Ollivier <v@vinc.cc>
Date: Sat, 5 Nov 2022 09:40:02 +0100
Subject: [PATCH] Add compilation option to set memory size (#433)

---
 Makefile             |  4 +++-
 src/sys/allocator.rs | 14 ++++++++++----
 src/sys/process.rs   |  7 ++++++-
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/Makefile b/Makefile
index e68f9f43b..c24099c32 100644
--- a/Makefile
+++ b/Makefile
@@ -8,6 +8,7 @@ setup:
 	cargo install bootimage
 
 # Compilation options
+memory = 32
 output = video# video, serial
 keyboard = qwerty# qwerty, azerty, dvorak
 mode = release
@@ -19,6 +20,7 @@ kvm = false
 pcap = false
 
 export MOROS_KEYBOARD = $(keyboard)
+export MOROS_MEMORY = $(memory)
 
 # Build userspace binaries
 user-nasm:
@@ -56,7 +58,7 @@ image: $(img)
 	dd conv=notrunc if=$(bin) of=$(img)
 
 
-qemu-opts = -m 32 -drive file=$(img),format=raw \
+qemu-opts = -m $(memory) -drive file=$(img),format=raw \
 			 -audiodev $(audio),id=a0 -machine pcspk-audiodev=a0 \
 			 -netdev user,id=e0,hostfwd=tcp::8080-:80 -device $(nic),netdev=e0
 ifeq ($(kvm),true)
diff --git a/src/sys/allocator.rs b/src/sys/allocator.rs
index 0dc404ba1..e7c68cad9 100644
--- a/src/sys/allocator.rs
+++ b/src/sys/allocator.rs
@@ -12,15 +12,21 @@ use x86_64::structures::paging::mapper::MapToError;
 use x86_64::structures::paging::{FrameAllocator, Mapper, Page, PageTableFlags, Size4KiB};
 use x86_64::VirtAddr;
 
-pub const HEAP_START: usize = 0x4444_4444_0000;
+pub const HEAP_START: u64 = 0x4444_4444_0000;
 
 #[global_allocator]
 static ALLOCATOR: LockedHeap = LockedHeap::empty();
 
+fn max_memory() -> u64 {
+    option_env!("MOROS_MEMORY").unwrap_or("32").parse::<u64>().unwrap() << 20 // MB
+}
+
 pub fn init_heap(mapper: &mut impl Mapper<Size4KiB>, frame_allocator: &mut impl FrameAllocator<Size4KiB>) -> Result<(), MapToError<Size4KiB>> {
-    // Use half of the memory for the heap, caped to 16 MB because the allocator is too slow
-    let heap_size = cmp::min(sys::mem::memory_size() / 2, 16 << 20);
-    let heap_start = VirtAddr::new(HEAP_START as u64);
+    // Use half of the memory for the heap caped to 16MB by default because the
+    // allocator is slow.
+    let heap_size = cmp::min(sys::mem::memory_size(), max_memory()) / 2;
+    let heap_start = VirtAddr::new(HEAP_START);
+    sys::process::init_process_addr(HEAP_START + heap_size);
 
     let pages = {
         let heap_end = heap_start + heap_size - 1u64;
diff --git a/src/sys/process.rs b/src/sys/process.rs
index be0bafe56..dbce359ed 100644
--- a/src/sys/process.rs
+++ b/src/sys/process.rs
@@ -200,7 +200,12 @@ use crate::sys::gdt::GDT;
 use core::sync::atomic::AtomicU64;
 use x86_64::VirtAddr;
 
-static CODE_ADDR: AtomicU64 = AtomicU64::new((sys::allocator::HEAP_START as u64) + (16 << 20));
+static CODE_ADDR: AtomicU64 = AtomicU64::new(0);
+
+// Called during kernel heap initialization
+pub fn init_process_addr(addr: u64) {
+    sys::process::CODE_ADDR.store(addr, Ordering::SeqCst);
+}
 
 #[repr(align(8), C)]
 #[derive(Debug, Clone, Copy, Default)]