Fix out-of-band allocations in dynamic arenas

2 files changed, 16 insertions, 15 deletions

diff --git a/core/mem/allocators.odin b/core/mem/allocators.odin
index 59fe72fbb..5b0b178f8 100644
--- a/core/mem/allocators.odin
+++ b/core/mem/allocators.odin
@@ -1819,18 +1819,18 @@ memory region.
 */
 @(require_results)
 dynamic_arena_alloc_bytes_non_zeroed :: proc(a: ^Dynamic_Arena, size: int, loc := #caller_location) -> ([]byte, Allocator_Error) {
-	n := align_formula(size, a.alignment)
-	if n > a.block_size {
-		return nil, .Invalid_Argument
-	}
-	if n >= a.out_band_size {
-		assert(a.block_allocator.procedure != nil, "Backing block allocator must be initialized", loc=loc)
-		memory, err := alloc_bytes_non_zeroed(a.block_size, a.alignment, a.block_allocator, loc)
+	if size >= a.out_band_size {
+		assert(a.out_band_allocations.allocator.procedure != nil, "Backing array allocator must be initialized", loc=loc)
+		memory, err := alloc_bytes_non_zeroed(size, a.alignment, a.out_band_allocations.allocator, loc)
 		if memory != nil {
 			append(&a.out_band_allocations, raw_data(memory), loc = loc)
 		}
 		return memory, err
 	}
+	n := align_formula(size, a.alignment)
+	if n > a.block_size {
+		return nil, .Invalid_Argument
+	}
 	if a.bytes_left < n {
 		err := _dynamic_arena_cycle_new_block(a, loc)
 		if err != nil {
@@ -1867,7 +1867,7 @@ dynamic_arena_reset :: proc(a: ^Dynamic_Arena, loc := #caller_location) {
 	}
 	clear(&a.used_blocks)
 	for allocation in a.out_band_allocations {
-		free(allocation, a.block_allocator, loc=loc)
+		free(allocation, a.out_band_allocations.allocator, loc=loc)
 	}
 	clear(&a.out_band_allocations)
 	a.bytes_left = 0 // Make new allocations call `_dynamic_arena_cycle_new_block` again.
diff --git a/tests/core/mem/test_mem_dynamic_pool.odin b/tests/core/mem/test_mem_dynamic_pool.odin
index fa204d3b1..82cee89af 100644
--- a/tests/core/mem/test_mem_dynamic_pool.odin
+++ b/tests/core/mem/test_mem_dynamic_pool.odin
@@ -44,11 +44,11 @@ expect_pool_allocation :: proc(t: ^testing.T, expected_used_bytes, num_bytes, al
 	testing.expect(t, pool.used_blocks == nil)
 }
 
-expect_pool_allocation_out_of_band :: proc(t: ^testing.T, num_bytes, out_band_size: int) {
+expect_pool_allocation_out_of_band :: proc(t: ^testing.T, num_bytes, block_size, out_band_size: int) {
 	testing.expect(t, num_bytes >= out_band_size, "Sanity check failed, your test call is flawed! Make sure that num_bytes >= out_band_size!")
 
 	pool: mem.Dynamic_Pool
-	mem.dynamic_pool_init(&pool, out_band_size = out_band_size)
+	mem.dynamic_pool_init(&pool, block_size = block_size, out_band_size = out_band_size)
 	pool_allocator := mem.dynamic_pool_allocator(&pool)
 
 	element, err := mem.alloc(num_bytes, allocator = pool_allocator)
@@ -69,14 +69,15 @@ test_dynamic_pool_alloc_aligned :: proc(t: ^testing.T) {
 
 @(test)
 test_dynamic_pool_alloc_unaligned :: proc(t: ^testing.T) {
-	expect_pool_allocation(t, expected_used_bytes =   8,   num_bytes=1, alignment=8)
-	expect_pool_allocation(t, expected_used_bytes =   16,  num_bytes=9, alignment=8)
+	expect_pool_allocation(t, expected_used_bytes = 8,  num_bytes = 1, alignment = 8)
+	expect_pool_allocation(t, expected_used_bytes = 16, num_bytes = 9, alignment = 8)
 }
 
 @(test)
 test_dynamic_pool_alloc_out_of_band :: proc(t: ^testing.T) {
-	expect_pool_allocation_out_of_band(t, num_bytes = 128, out_band_size = 128)
-	expect_pool_allocation_out_of_band(t, num_bytes = 129, out_band_size = 128)
+	expect_pool_allocation_out_of_band(t, num_bytes = 128, block_size = 512, out_band_size = 128)
+	expect_pool_allocation_out_of_band(t, num_bytes = 129, block_size = 512, out_band_size = 128)
+	expect_pool_allocation_out_of_band(t, num_bytes = 513, block_size = 512, out_band_size = 128)
 }
 
 @(test)
@@ -98,4 +99,4 @@ intentionally_leaky_test :: proc(t: ^testing.T) {
 leak_verifier :: proc(t: ^testing.T, ta: ^mem.Tracking_Allocator) {
 	testing.expect_value(t, len(ta.allocation_map), 1)
 	testing.expect_value(t, len(ta.bad_free_array), 1)
-}
\ No newline at end of file
+}