From f16ed256eaf90fb0fed1e795f6c62cd356180422 Mon Sep 17 00:00:00 2001
From: flysand7 <thebumboni@gmail.com>
Date: Wed, 11 Sep 2024 08:00:27 +1100
Subject: [PATCH] [mem]: Fix handling of default resize to check alignment

---
 base/runtime/internal.odin | 11 +++++++----
 core/mem/alloc.odin        |  2 +-
 core/mem/mem.odin          | 11 +++++++++++
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/base/runtime/internal.odin b/base/runtime/internal.odin
index ff60cf547..a0bda9d40 100644
--- a/base/runtime/internal.odin
+++ b/base/runtime/internal.odin
@@ -118,16 +118,15 @@ mem_copy_non_overlapping :: proc "contextless" (dst, src: rawptr, len: int) -> r
 DEFAULT_ALIGNMENT :: 2*align_of(rawptr)
 
 mem_alloc_bytes :: #force_inline proc(size: int, alignment: int = DEFAULT_ALIGNMENT, allocator := context.allocator, loc := #caller_location) -> ([]byte, Allocator_Error) {
-	if size == 0 {
-		return nil, nil
-	}
-	if allocator.procedure == nil {
+	assert(is_power_of_two_int(alignment), "Alignment must be a power of two", loc)
+	if size == 0 || allocator.procedure == nil{
 		return nil, nil
 	}
 	return allocator.procedure(allocator.data, .Alloc, size, alignment, nil, 0, loc)
 }
 
 mem_alloc :: #force_inline proc(size: int, alignment: int = DEFAULT_ALIGNMENT, allocator := context.allocator, loc := #caller_location) -> ([]byte, Allocator_Error) {
+	assert(is_power_of_two_int(alignment), "Alignment must be a power of two", loc)
 	if size == 0 || allocator.procedure == nil {
 		return nil, nil
 	}
@@ -135,6 +134,7 @@ mem_alloc :: #force_inline proc(size: int, alignment: int = DEFAULT_ALIGNMENT, a
 }
 
 mem_alloc_non_zeroed :: #force_inline proc(size: int, alignment: int = DEFAULT_ALIGNMENT, allocator := context.allocator, loc := #caller_location) -> ([]byte, Allocator_Error) {
+	assert(is_power_of_two_int(alignment), "Alignment must be a power of two", loc)
 	if size == 0 || allocator.procedure == nil {
 		return nil, nil
 	}
@@ -174,6 +174,7 @@ mem_free_all :: #force_inline proc(allocator := context.allocator, loc := #calle
 }
 
 _mem_resize :: #force_inline proc(ptr: rawptr, old_size, new_size: int, alignment: int = DEFAULT_ALIGNMENT, allocator := context.allocator, should_zero: bool, loc := #caller_location) -> (data: []byte, err: Allocator_Error) {
+	assert(is_power_of_two_int(alignment), "Alignment must be a power of two", loc)
 	if allocator.procedure == nil {
 		return nil, nil
 	}
@@ -215,9 +216,11 @@ _mem_resize :: #force_inline proc(ptr: rawptr, old_size, new_size: int, alignmen
 }
 
 mem_resize :: proc(ptr: rawptr, old_size, new_size: int, alignment: int = DEFAULT_ALIGNMENT, allocator := context.allocator, loc := #caller_location) -> (data: []byte, err: Allocator_Error) {
+	assert(is_power_of_two_int(alignment), "Alignment must be a power of two", loc)
 	return _mem_resize(ptr, old_size, new_size, alignment, allocator, true, loc)
 }
 non_zero_mem_resize :: proc(ptr: rawptr, old_size, new_size: int, alignment: int = DEFAULT_ALIGNMENT, allocator := context.allocator, loc := #caller_location) -> (data: []byte, err: Allocator_Error) {
+	assert(is_power_of_two_int(alignment), "Alignment must be a power of two", loc)
 	return _mem_resize(ptr, old_size, new_size, alignment, allocator, false, loc)
 }
 
diff --git a/core/mem/alloc.odin b/core/mem/alloc.odin
index 1ede92837..5f65e9ebc 100644
--- a/core/mem/alloc.odin
+++ b/core/mem/alloc.odin
@@ -1096,7 +1096,7 @@ _default_resize_bytes_align :: #force_inline proc(
 		err := free_bytes(old_data, allocator, loc)
 		return nil, err
 	}
-	if new_size == old_size {
+	if new_size == old_size && is_aligned(old_memory, alignment) {
 		return old_data, .None
 	}
 	new_memory : []byte
diff --git a/core/mem/mem.odin b/core/mem/mem.odin
index 0554cee23..b57b18ffc 100644
--- a/core/mem/mem.odin
+++ b/core/mem/mem.odin
@@ -456,6 +456,17 @@ is_power_of_two :: proc "contextless" (x: uintptr) -> bool {
 	return (x & (x-1)) == 0
 }
 
+/*
+Check if a pointer is aligned.
+
+This procedure checks whether a pointer `x` is aligned to a boundary specified
+by `align`, and returns `true` if the pointer is aligned, and false otherwise.
+*/
+is_aligned :: proc "contextless" (x: rawptr, align: int) -> bool {
+	p := uintptr(x)
+	return (p & (1<<uintptr(align) - 1)) == 0
+}
+
 /*
 Align uintptr forward.