From d6490c9babe8878c44927db5f0d9fe9a4633d6ad Mon Sep 17 00:00:00 2001
From: Christoffer Lerno <christoffer@aegik.com>
Date: Wed, 31 Jul 2024 18:03:52 +0200
Subject: [PATCH] Make it possible to set max limit for memory pages.

---
 src/build/build_options.c |  9 +++++++++
 src/main.c                | 13 ++++++++++++-
 src/utils/lib.h           |  2 +-
 src/utils/malloc.c        |  9 +++++----
 src/utils/vmem.c          |  7 +++++++
 src/utils/vmem.h          |  1 +
 6 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/src/build/build_options.c b/src/build/build_options.c
index c20baf31f..673488ddc 100644
--- a/src/build/build_options.c
+++ b/src/build/build_options.c
@@ -89,6 +89,7 @@ static void usage(void)
 	PRINTF("  --template <template>      - Select template for 'init': \"exe\", \"static-lib\", \"dynamic-lib\" or a path.");
 	PRINTF("  --about                    - Prints a short description of C3.");
 	PRINTF("  --symtab <value>           - Sets the preferred symtab size.");
+	PRINTF("  --max-mem <value>          - Sets the preferred max memory size.");
 	PRINTF("  -V --version               - Print version information.");
 	PRINTF("  -E                         - Lex only.");
 	PRINTF("  -P                         - Only parse and output the AST as JSON.");
@@ -652,6 +653,14 @@ static void parse_option(BuildOptions *options)
 				options->backend = BACKEND_TB;
 				return;
 			}
+			if (match_longopt("max-mem"))
+			{
+				if (at_end() || next_is_opt()) error_exit("error: --max-mem needs a valid integer.");
+				const char *maxmem_string = next_arg();
+				int maxmem = atoi(maxmem_string);
+				if (maxmem < 128) PRINTF("Expected a valid positive integer >= 128.");
+				return;
+			}
 			if (match_longopt("symtab"))
 			{
 				if (at_end() || next_is_opt()) error_exit("error: --symtab needs a valid integer.");
diff --git a/src/main.c b/src/main.c
index 50d390911..89edb9335 100644
--- a/src/main.c
+++ b/src/main.c
@@ -38,8 +38,19 @@ int main_real(int argc, const char *argv[])
 		return result;
 	}
 
+	// Memory must be handled here
+	long long max_mem = 0;
+	for (int i = 0; i < argc; i++)
+	{
+		if (str_eq(argv[i], "--max-mem") && i < argc - 1)
+		{
+			max_mem = atoll(argv[i + 1]);
+			if (max_mem) max_mem = next_highest_power_of_2(max_mem);
+			break;
+		}
+	}
 	// First setup memory
-	memory_init();
+	memory_init(max_mem);
 
 	// Parse arguments.
 	BuildOptions build_options = parse_arguments(argc, argv);
diff --git a/src/utils/lib.h b/src/utils/lib.h
index 7652f7490..fefcef1f4 100644
--- a/src/utils/lib.h
+++ b/src/utils/lib.h
@@ -91,7 +91,7 @@ const char *execute_cmd(const char *cmd, bool ignore_failure);
 bool execute_cmd_failable(const char *cmd, const char **result);
 void *cmalloc(size_t size);
 void *ccalloc(size_t size, size_t elements);
-void memory_init(void);
+void memory_init(size_t max_mem);
 void memory_release();
 
 #define ptrid(ptr_) ((((uintptr_t)(ptr_)) - arena_zero) / 16)
diff --git a/src/utils/malloc.c b/src/utils/malloc.c
index 04e30d2f0..eb3243f2d 100644
--- a/src/utils/malloc.c
+++ b/src/utils/malloc.c
@@ -14,8 +14,9 @@ static Vmem arena;
 uintptr_t arena_zero;
 static Vmem char_arena;
 
-void memory_init(void)
+void memory_init(size_t max_mem)
 {
+	if (max_mem) vmem_set_max_limit(max_mem);
 	vmem_init(&arena, 2048);
 	vmem_init(&char_arena, 512);
 	allocations_done = 0;
@@ -66,9 +67,9 @@ void run_arena_allocator_tests(void)
 {
 	printf("Begin arena allocator testing.\n");
 	bool was_init = arena.ptr != NULL;
-	if (!was_init) memory_init();
+	if (!was_init) memory_init(0);
 	memory_release();
-	memory_init();
+	memory_init(0);
 	ASSERT(calloc_arena(10) != calloc_arena(10), "Expected different values...");
 	printf("-- Tested basic allocation - OK.\n");
 	ASSERT(arena.allocated == 48, "Expected allocations rounded to next 16 bytes");
@@ -79,7 +80,7 @@ void run_arena_allocator_tests(void)
 	free_arena();
 	ASSERT(arena.allocated == 0, "Arena not freed?");
 	printf("-- Test freeing arena - OK.\n");
-	if (was_init) memory_init();
+	if (was_init) memory_init(0);
 }
 
 void *cmalloc(size_t size)
diff --git a/src/utils/vmem.c b/src/utils/vmem.c
index e3ec5538e..a7d41d16e 100644
--- a/src/utils/vmem.c
+++ b/src/utils/vmem.c
@@ -59,8 +59,15 @@ static inline void* mmap_allocate(Vmem *vmem, size_t to_allocate)
 	return ptr;
 }
 
+size_t max = 0x10000000;
+void vmem_set_max_limit(size_t size_in_mb)
+{
+	max = size_in_mb;
+}
+
 void vmem_init(Vmem *vmem, size_t size_in_mb)
 {
+	if (size_in_mb > max) size_in_mb = max;
 	mmap_init(vmem, 1024 * 1024 * size_in_mb);
 }
 
diff --git a/src/utils/vmem.h b/src/utils/vmem.h
index be4143932..3b1fc6522 100644
--- a/src/utils/vmem.h
+++ b/src/utils/vmem.h
@@ -19,3 +19,4 @@ typedef struct
 void vmem_init(Vmem *vmem, size_t size_in_mb);
 void *vmem_alloc(Vmem *vmem, size_t alloc);
 void vmem_free(Vmem *vmem);
+void vmem_set_max_limit(size_t size_in_mb);
\ No newline at end of file