Temporary allocator for `context`

7 files changed, 144 insertions, 13 deletions

diff --git a/core/fmt/fmt.odin b/core/fmt/fmt.odin
index 459d6e2ce..b8a0f9a24 100644
--- a/core/fmt/fmt.odin
+++ b/core/fmt/fmt.odin
@@ -159,6 +159,24 @@ aprintf :: proc(fmt: string, args: ..any) -> string {
 }
 
 
+// tprint* procedures return a string that was allocated with the current context's temporary allocator
+tprint :: proc(args: ..any) -> string {
+	buf := String_Buffer(make([dynamic]byte, context.temp_allocator));
+	sbprint(&buf, ..args);
+	return to_string(buf);
+}
+tprintln :: proc(args: ..any) -> string {
+	buf := String_Buffer(make([dynamic]byte, context.temp_allocator));
+	sbprintln(&buf, ..args);
+	return to_string(buf);
+}
+tprintf :: proc(fmt: string, args: ..any) -> string {
+	buf := String_Buffer(make([dynamic]byte, context.temp_allocator));
+	sbprintf(&buf, fmt, ..args);
+	return to_string(buf);
+}
+
+
 // bprint* procedures return a string using a buffer from an array
 bprint :: proc(buf: []byte, args: ..any) -> string {
 	sb := string_buffer_from_slice(buf[0:len(buf)]);
diff --git a/core/mem/alloc.odin b/core/mem/alloc.odin
index 7dee561e2..b63317815 100644
--- a/core/mem/alloc.odin
+++ b/core/mem/alloc.odin
@@ -159,6 +159,84 @@ nil_allocator :: proc() -> Allocator {
 	};
 }
 
+Scratch_Allocator :: struct {
+	data:     []byte,
+	curr_offset: int,
+	prev_offset: int,
+	backup_allocator: Allocator,
+}
+
+scratch_allocator_init :: proc(scratch: ^Scratch_Allocator, data: []byte, backup_allocator := context.allocator) {
+	scratch.data = data;
+	scratch.curr_offset = 0;
+	scratch.prev_offset = 0;
+	scratch.backup_allocator = backup_allocator;
+}
+
+scratch_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode,
+                               size, alignment: int,
+                               old_memory: rawptr, old_size: int, flags: u64 = 0, loc := #caller_location) -> rawptr {
+
+	scratch := (^Scratch_Allocator)(allocator_data);
+
+	switch mode {
+	case Allocator_Mode.Alloc:
+		switch {
+		case scratch.curr_offset+size <= len(scratch.data):
+			offset := align_forward_uintptr(uintptr(scratch.curr_offset), uintptr(alignment));
+			ptr := &scratch.data[offset];
+			zero(ptr, size);
+			scratch.prev_offset = int(offset);
+			scratch.curr_offset = int(offset) + size;
+			return ptr;
+		case size <= len(scratch.data):
+			offset := align_forward_uintptr(uintptr(0), uintptr(alignment));
+			ptr := &scratch.data[offset];
+			zero(ptr, size);
+			scratch.prev_offset = int(offset);
+			scratch.curr_offset = int(offset) + size;
+			return ptr;
+		}
+		// TODO(bill): Should leaks be notified about? Should probably use a logging system that is built into the context system
+		a := scratch.backup_allocator;
+		if a.procedure == nil {
+			a = context.allocator;
+		}
+		return alloc(size, alignment, a, loc);
+
+	case Allocator_Mode.Free:
+		last_ptr := rawptr(&scratch.data[scratch.prev_offset]);
+		if old_memory == last_ptr {
+			size := scratch.curr_offset - scratch.prev_offset;
+			scratch.curr_offset = scratch.prev_offset;
+			zero(last_ptr, size);
+			return nil;
+		}
+		// NOTE(bill): It's scratch memory, don't worry about freeing
+
+	case Allocator_Mode.Free_All:
+		scratch.curr_offset = 0;
+		scratch.prev_offset = 0;
+
+	case Allocator_Mode.Resize:
+		last_ptr := rawptr(&scratch.data[scratch.prev_offset]);
+		if old_memory == last_ptr && len(scratch.data)-scratch.prev_offset >= size {
+			scratch.curr_offset = scratch.prev_offset+size;
+			return old_memory;
+		}
+		return scratch_allocator_proc(allocator_data, Allocator_Mode.Alloc, size, alignment, old_memory, old_size, flags, loc);
+	}
+
+	return nil;
+}
+
+scratch_allocator :: proc(scratch: ^Scratch_Allocator) -> Allocator {
+	return Allocator{
+		procedure = scratch_allocator_proc,
+		data = scratch,
+	};
+}
+
 
 
 
diff --git a/core/mem/mem.odin b/core/mem/mem.odin
index 51d66aa7f..283d3e29c 100644
--- a/core/mem/mem.odin
+++ b/core/mem/mem.odin
@@ -10,8 +10,9 @@ foreign _ {
 swap :: proc[swap16, swap32, swap64];
 
 
-set :: proc "contextless" (data: rawptr, value: i32, len: int) -> rawptr {
+set :: proc "contextless" (data: rawptr, value: byte, len: int) -> rawptr {
 	if data == nil do return nil;
+	if len < 0 do return data;
 	foreign _ {
 		when size_of(rawptr) == 8 {
 			@(link_name="llvm.memset.p0i8.i64")
@@ -138,6 +139,16 @@ align_forward :: proc(ptr: rawptr, align: uintptr) -> rawptr {
 	return rawptr(p);
 }
 
+align_forward_uintptr :: proc(ptr, align: uintptr) -> uintptr {
+	assert(is_power_of_two(align));
+
+	a := uintptr(align);
+	p := uintptr(ptr);
+	modulo := p & (a-1);
+	if modulo != 0 do p += a - modulo;
+	return uintptr(p);
+}
+
 
 
 AllocationHeader :: struct {size: int};
@@ -181,7 +192,7 @@ Arena :: struct {
 	temp_count: int,
 }
 
-ArenaTempMemory :: struct {
+Arena_Temp_Memory :: struct {
 	arena:          ^Arena,
 	original_count: int,
 }
@@ -249,7 +260,7 @@ arena_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode,
 
 	case Free:
 		// NOTE(bill): Free all at once
-		// Use ArenaTempMemory if you want to free a block
+		// Use Arena_Temp_Memory if you want to free a block
 
 	case Free_All:
 		(^Raw_Slice)(&arena.memory).len = 0;
@@ -261,15 +272,15 @@ arena_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode,
 	return nil;
 }
 
-begin_arena_temp_memory :: proc(a: ^Arena) -> ArenaTempMemory {
-	tmp: ArenaTempMemory;
+begin_arena_temp_memory :: proc(a: ^Arena) -> Arena_Temp_Memory {
+	tmp: Arena_Temp_Memory;
 	tmp.arena = a;
 	tmp.original_count = len(a.memory);
 	a.temp_count += 1;
 	return tmp;
 }
 
-end_arena_temp_memory :: proc(using tmp: ArenaTempMemory) {
+end_arena_temp_memory :: proc(using tmp: Arena_Temp_Memory) {
 	assert(len(arena.memory) >= original_count);
 	assert(arena.temp_count > 0);
 	(^Raw_Dynamic_Array)(&arena.memory).len = original_count;
diff --git a/core/runtime/core.odin b/core/runtime/core.odin
index 7bf3a9826..9b3e96d56 100644
--- a/core/runtime/core.odin
+++ b/core/runtime/core.odin
@@ -184,7 +184,8 @@ Source_Code_Location :: struct {
 }
 
 Context :: struct {
-	allocator:  mem.Allocator,
+	allocator:      mem.Allocator,
+	temp_allocator: mem.Allocator,
 	thread_id:  int,
 
 	user_data:  any,
@@ -194,6 +195,8 @@ Context :: struct {
 	derived:    any, // May be used for derived data types
 }
 
+global_scratch_allocator_data: mem.Scratch_Allocator;
+
 
 
 
@@ -315,9 +318,15 @@ __init_context :: proc "contextless" (c: ^Context) {
 	if c == nil do return;
 
 	c.allocator = os.heap_allocator();
+	c.temp_allocator = mem.scratch_allocator(&global_scratch_allocator_data);
 	c.thread_id = os.current_thread_id();
 }
 
+@(builtin)
+init_global_temporary_allocator :: proc(data: []byte, backup_allocator := context.allocator) {
+	mem.scratch_allocator_init(&global_scratch_allocator_data, data, backup_allocator);
+}
+
 
 
 @(builtin)
diff --git a/examples/demo/demo.odin b/examples/demo/demo.odin
index f16571d56..ed26b9664 100644
--- a/examples/demo/demo.odin
+++ b/examples/demo/demo.odin
@@ -489,12 +489,15 @@ threading_example :: proc() {
 	when os.OS == "windows" {
 		fmt.println("# threading_example");
 
-		unordered_remove :: proc(array: ^[dynamic]$T, index: int, loc := #caller_location) {
+		unordered_remove :: proc(array: ^$D/[dynamic]$T, index: int, loc := #caller_location) {
 			runtime.bounds_check_error_loc(loc, index, len(array));
-			array[index] = array[len(array)-1];
+			n := len(array)-1;
+			if index != n {
+				array[index] = array[n];
+			}
 			pop(array);
 		}
-		ordered_remove :: proc(array: ^[dynamic]$T, index: int, loc := #caller_location) {
+		ordered_remove :: proc(array: ^$D/[dynamic]$T, index: int, loc := #caller_location) {
 			runtime.bounds_check_error_loc(loc, index, len(array));
 			copy(array[index:], array[index+1:]);
 			pop(array);
diff --git a/src/checker.cpp b/src/checker.cpp
index b0d6d72aa..2035ded75 100644
--- a/src/checker.cpp
+++ b/src/checker.cpp
@@ -1309,6 +1309,7 @@ void generate_minimum_dependency_set(Checker *c, Entity *start) {
 
 		str_lit("args__"),
 		str_lit("type_table"),
+		str_lit("global_scratch_allocator"),
 
 		str_lit("Type_Info"),
 		str_lit("Source_Code_Location"),
diff --git a/src/ir.cpp b/src/ir.cpp
index daa0770d2..53391c236 100644
--- a/src/ir.cpp
+++ b/src/ir.cpp
@@ -1658,11 +1658,22 @@ irValue *ir_find_or_generate_context_ptr(irProcedure *proc) {
 	ir_push_context_onto_stack(proc, c);
 	ir_emit_store(proc, c, ir_emit_load(proc, proc->module->global_default_context));
 
-	irValue *ep = ir_emit_struct_ep(proc, c, 0);
+
+#if 1
 	Array<irValue *> args = {};
-	irValue *v = ir_emit_package_call(proc, "os", "heap_allocator", args);
-	ir_emit_store(proc, ep, v);
+	ir_emit_store(proc, ir_emit_struct_ep(proc, c, 0), ir_emit_package_call(proc, "os", "heap_allocator", args));
+	ir_emit_store(proc, ir_emit_struct_ep(proc, c, 2), ir_emit_package_call(proc, "os", "current_thread_id", args));
 
+	array_init(&args, heap_allocator(), 1);
+	AstPackage *rt_pkg = get_core_package(proc->module->info, str_lit("runtime"));
+	Entity *e = scope_lookup_current(rt_pkg->scope, str_lit("global_scratch_allocator_data"));
+	irValue **found = map_get(&proc->module->values, hash_entity(e));
+	GB_ASSERT_MSG(found != nullptr, "%.*s", LIT(e->token.string));
+	args[0] = *found;
+	ir_emit_store(proc, ir_emit_struct_ep(proc, c, 1), ir_emit_package_call(proc, "mem", "scratch_allocator", args));
+#else
+	ir_emit_init_context(proc, c);
+#endif
 	return c;
 }