Add `mem.Allocator_Query_Info` and `mem.query_info`

5 files changed, 86 insertions, 14 deletions

diff --git a/core/mem/alloc.odin b/core/mem/alloc.odin
index 0ed691dd6..ea870cf7d 100644
--- a/core/mem/alloc.odin
+++ b/core/mem/alloc.odin
@@ -19,6 +19,15 @@ Allocator_Mode_Set :: runtime.Allocator_Mode_Set;
 Allocator_Mode_Set :: distinct bit_set[Allocator_Mode];
 */
 
+Allocator_Query_Info :: runtime.Allocator_Query_Info;
+/*
+Allocator_Query_Info :: struct {
+	pointer:   Maybe(rawptr),
+	size:      Maybe(int),
+	alignment: Maybe(int),
+}
+*/
+
 Allocator_Proc :: runtime.Allocator_Proc;
 /*
 Allocator_Proc :: #type proc(allocator_data: rawptr, mode: Allocator_Mode,
@@ -69,17 +78,22 @@ resize :: inline proc(ptr: rawptr, old_size, new_size: int, alignment: int = DEF
 	return allocator.procedure(allocator.data, Allocator_Mode.Resize, new_size, alignment, ptr, old_size, 0, loc);
 }
 
-query_features :: proc(allocator: Allocator, loc := #caller_location) -> Allocator_Mode_Set {
+query_features :: proc(allocator: Allocator, loc := #caller_location) -> (set: Allocator_Mode_Set) {
 	if allocator.procedure != nil {
-		set: Allocator_Mode_Set;
-		res := allocator.procedure(allocator.data, Allocator_Mode.Query_Features, 0, 0, &set, 0, 0, loc);
-		if res == &set {
-			return set;
-		}
+		allocator.procedure(allocator.data, Allocator_Mode.Query_Features, 0, 0, &set, 0, 0, loc);
+		return set;
 	}
 	return nil;
 }
 
+query_info :: proc(pointer: rawptr, allocator: Allocator, loc := #caller_location) -> (props: Allocator_Query_Info) {
+	props.pointer = pointer;
+	if allocator.procedure != nil {
+		allocator.procedure(allocator.data, Allocator_Mode.Query_Info, 0, 0, &props, 0, 0, loc);
+	}
+	return;
+}
+
 
 
 delete_string :: proc(str: string, allocator := context.allocator, loc := #caller_location) {
diff --git a/core/mem/allocators.odin b/core/mem/allocators.odin
index 8c27f6887..1a4871732 100644
--- a/core/mem/allocators.odin
+++ b/core/mem/allocators.odin
@@ -81,6 +81,9 @@ arena_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode,
 			set^ = {.Alloc, .Free_All, .Resize, .Query_Features};
 		}
 		return set;
+
+	case .Query_Info:
+		return nil;
 	}
 
 	return nil;
@@ -211,6 +214,9 @@ scratch_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode,
 			set^ = {.Alloc, .Free, .Free_All, .Resize, .Query_Features};
 		}
 		return set;
+
+	case .Query_Info:
+		return nil;
 	}
 
 	return nil;
@@ -369,6 +375,8 @@ stack_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode,
 			set^ = {.Alloc, .Free, .Free_All, .Resize, .Query_Features};
 		}
 		return set;
+	case .Query_Info:
+		return nil;
 	}
 
 	return nil;
@@ -496,6 +504,9 @@ small_stack_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode,
 			set^ = {.Alloc, .Free, .Free_All, .Resize, .Query_Features};
 		}
 		return set;
+
+	case .Query_Info:
+		return nil;
 	}
 
 	return nil;
@@ -551,9 +562,18 @@ dynamic_pool_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode
 	case .Query_Features:
 		set := (^Allocator_Mode_Set)(old_memory);
 		if set != nil {
-			set^ = {.Alloc, .Free_All, .Resize, .Query_Features};
+			set^ = {.Alloc, .Free_All, .Resize, .Query_Features, .Query_Info};
 		}
 		return set;
+
+	case .Query_Info:
+		info := (^Allocator_Query_Info)(old_memory);
+		if info != nil && info.pointer != nil {
+			info.size = pool.block_size;
+			info.alignment = pool.alignment;
+			return info;
+		}
+		return nil;
 	}
 	return nil;
 }
@@ -696,6 +716,9 @@ panic_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode,
 			set^ = {.Query_Features};
 		}
 		return set;
+
+	case .Query_Info:
+		return nil;
 	}
 
 	return nil;
@@ -750,6 +773,9 @@ alloca_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode,
 			set^ = {.Alloc, .Resize, .Query_Features};
 		}
 		return set;
+
+	case .Query_Info:
+		return nil;
 	}
 	return nil;
 }
@@ -766,8 +792,9 @@ alloca_allocator :: proc() -> Allocator {
 
 
 Tracking_Allocator_Entry :: struct {
-	memory:   rawptr,
-	size:     int,
+	memory:    rawptr,
+	size:      int,
+	alignment: int,
 	location: runtime.Source_Code_Location,
 }
 Tracking_Allocator :: struct {
@@ -794,6 +821,20 @@ tracking_allocator :: proc(data: ^Tracking_Allocator) -> Allocator {
 
 tracking_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 {
 	data := (^Tracking_Allocator)(allocator_data);
+	if mode == .Query_Info {
+		info := (^Allocator_Query_Info)(old_memory);
+		if info != nil && info.pointer != nil {
+			if entry, ok := data.allocation_map[info.pointer]; ok {
+				info.size = entry.size;
+				info.alignment = entry.alignment;
+				return info;
+			}
+			info.pointer = nil;
+		}
+
+		return nil;
+	}
+
 	result := data.backing.procedure(data.backing.data, mode, size, alignment, old_memory, old_size, flags, loc);
 
 	if data.allocation_map.allocator.procedure == nil {
@@ -805,6 +846,7 @@ tracking_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode, si
 		data.allocation_map[result] = Tracking_Allocator_Entry{
 			memory = result,
 			size = size,
+			alignment = alignment,
 			location = loc,
 		};
 	case .Free:
@@ -816,6 +858,7 @@ tracking_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode, si
 		data.allocation_map[result] = Tracking_Allocator_Entry{
 			memory = result,
 			size = size,
+			alignment = alignment,
 			location = loc,
 		};
 
@@ -827,9 +870,12 @@ tracking_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode, si
 	case .Query_Features:
 		set := (^Allocator_Mode_Set)(old_memory);
 		if set != nil {
-			set^ = {.Alloc, .Free, .Free_All, .Resize, .Query_Features};
+			set^ = {.Alloc, .Free, .Free_All, .Resize, .Query_Features, .Query_Info};
 		}
 		return set;
+
+	case .Query_Info:
+		unreachable();
 	}
 
 	return result;
diff --git a/core/os/os.odin b/core/os/os.odin
index 9342355cf..10613e513 100644
--- a/core/os/os.odin
+++ b/core/os/os.odin
@@ -201,6 +201,9 @@ heap_allocator_proc :: proc(allocator_data: rawptr, mode: mem.Allocator_Mode,
 			set^ = {.Alloc, .Free, .Resize, .Query_Features};
 		}
 		return set;
+
+	case .Query_Info:
+		return nil;
 	}
 
 	return nil;
diff --git a/core/runtime/core.odin b/core/runtime/core.odin
index d3522c268..4e871351b 100644
--- a/core/runtime/core.odin
+++ b/core/runtime/core.odin
@@ -234,6 +234,9 @@ args__: []cstring;
 
 // IMPORTANT NOTE(bill): Must be in this order (as the compiler relies upon it)
 
+@builtin
+Maybe :: union(T: typeid) #maybe {T};
+
 
 Source_Code_Location :: struct {
 	file_path:    string,
@@ -252,10 +255,17 @@ Allocator_Mode :: enum byte {
 	Free_All,
 	Resize,
 	Query_Features,
+	Query_Info,
 }
 
 Allocator_Mode_Set :: distinct bit_set[Allocator_Mode];
 
+Allocator_Query_Info :: struct {
+	pointer:   rawptr,
+	size:      Maybe(int),
+	alignment: Maybe(int),
+}
+
 Allocator_Proc :: #type proc(allocator_data: rawptr, mode: Allocator_Mode,
                              size, alignment: int,
                              old_memory: rawptr, old_size: int, flags: u64 = 0, location: Source_Code_Location = #caller_location) -> rawptr;
@@ -310,10 +320,6 @@ Context :: struct {
 }
 
 
-@builtin
-Maybe :: union(T: typeid) #maybe {T};
-
-
 
 @thread_local global_default_temp_allocator_data: Default_Temp_Allocator;
 
diff --git a/core/runtime/default_allocators.odin b/core/runtime/default_allocators.odin
index 3643fccdf..14f4edcda 100644
--- a/core/runtime/default_allocators.odin
+++ b/core/runtime/default_allocators.odin
@@ -139,6 +139,9 @@ default_temp_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode
 			set^ = {.Alloc, .Free, .Free_All, .Resize, .Query_Features};
 		}
 		return set;
+
+	case .Query_Info:
+		return nil;
 	}
 
 	return nil;