Browse Source

[runtime] Add builtin `shrink` for dynamic arrays and maps

Asks the allocator to shrink the backing allocation to the current __length__, or a capacity
of the user's choosing.

Returns `(did_shrink: bool, err: mem.Allocator_Error)`.

```
shrink(&array) // shrinks to len(array)
shrink(&array, N) // shrink to N capacity

shrink(&map) // shrinks down to len(map)
shrink(&map, N) // shrink to N capacity
```
Tetralux 3 years ago
parent
commit
fa2296a124

+ 71 - 2
core/runtime/core_builtin.odin

@@ -129,6 +129,9 @@ reserve :: proc{reserve_dynamic_array, reserve_map}
 @builtin
 resize :: proc{resize_dynamic_array}
 
+// Shrinks the capacity of a dynamic array or map down to the current length, or the given capacity.
+@builtin
+shrink :: proc{shrink_dynamic_array, shrink_map}
 
 @builtin
 free :: proc{mem_free}
@@ -284,10 +287,28 @@ clear_map :: proc "contextless" (m: ^$T/map[$K]$V) {
 }
 
 @builtin
-reserve_map :: proc(m: ^$T/map[$K]$V, capacity: int) {
+reserve_map :: proc(m: ^$T/map[$K]$V, capacity: int, loc := #caller_location) {
+	if m != nil {
+		__dynamic_map_reserve(__get_map_header(m), capacity, loc)
+	}
+}
+
+/*
+	Shrinks the capacity of a map down to the current length, or the given capacity.
+
+	If `new_cap` is -1, then `len(m)` is used.
+
+	Returns false if `cap(m) < new_cap`, or the allocator report failure.
+
+	If `len(m) < new_cap`, then `len(m)` will be left unchanged.
+*/
+@builtin
+shrink_map :: proc(m: ^$T/map[$K]$V, new_cap := -1, loc := #caller_location) -> (did_shrink: bool) {
 	if m != nil {
-		__dynamic_map_reserve(__get_map_header(m), capacity)
+		new_cap := new_cap if new_cap != -1 else len(m)
+		return __dynamic_map_shrink(__get_map_header(m), new_cap, loc)
 	}
+	return
 }
 
 // The delete_key built-in procedure deletes the element with the specified key (m[key]) from the map.
@@ -536,6 +557,54 @@ resize_dynamic_array :: proc(array: ^$T/[dynamic]$E, length: int, loc := #caller
 	return true
 }
 
+/*
+	Shrinks the capacity of a dynamic array down to the current length, or the given capacity.
+
+	If `new_cap` is -1, then `len(array)` is used.
+
+	Returns false if `cap(array) < new_cap`, or the allocator report failure.
+
+	If `len(array) < new_cap`, then `len(array)` will be left unchanged.
+*/
+shrink_dynamic_array :: proc(array: ^$T/[dynamic]$E, new_cap := -1, loc := #caller_location) -> (did_shrink: bool) {
+	if array == nil {
+		return
+	}
+	a := (^Raw_Dynamic_Array)(array)
+
+	new_cap := new_cap if new_cap != -1 else a.len
+
+	if new_cap > a.cap {
+		return
+	}
+
+	if a.allocator.procedure == nil {
+		a.allocator = context.allocator
+	}
+	assert(a.allocator.procedure != nil)
+
+	old_size := a.cap * size_of(E)
+	new_size := new_cap * size_of(E)
+
+	new_data, err := a.allocator.procedure(
+		a.allocator.data,
+		.Resize,
+		new_size,
+		align_of(E),
+		a.data,
+		old_size,
+		loc,
+	)
+	if err != nil {
+		return
+	}
+
+	a.data = raw_data(new_data)
+	a.len = min(new_cap, a.len)
+	a.cap = new_cap
+	return true
+}
+
 @builtin
 map_insert :: proc(m: ^$T/map[$K]$V, key: K, value: V, loc := #caller_location) -> (ptr: ^V) {
 	key, value := key, value

+ 29 - 0
core/runtime/dynamic_array_internal.odin

@@ -41,6 +41,35 @@ __dynamic_array_reserve :: proc(array_: rawptr, elem_size, elem_align: int, cap:
 	return false
 }
 
+__dynamic_array_shrink :: proc(array_: rawptr, elem_size, elem_align: int, new_cap: int, loc := #caller_location) -> (did_shrink: bool) {
+	array := (^Raw_Dynamic_Array)(array_)
+
+	// NOTE(tetra, 2020-01-26): We set the allocator before earlying-out below, because user code is usually written
+	// assuming that appending/reserving will set the allocator, if it is not already set.
+	if array.allocator.procedure == nil {
+		array.allocator = context.allocator
+	}
+	assert(array.allocator.procedure != nil)
+
+	if new_cap > array.cap {
+		return
+	}
+
+	old_size  := array.cap * elem_size
+	new_size  := new_cap * elem_size
+	allocator := array.allocator
+
+	new_data, err := allocator.procedure(allocator.data, .Resize, new_size, elem_align, array.data, old_size, loc)
+	if err != nil {
+		return
+	}
+
+	array.data = raw_data(new_data)
+	array.len = min(new_cap, array.len)
+	array.cap = new_cap
+	return true
+}
+
 __dynamic_array_resize :: proc(array_: rawptr, elem_size, elem_align: int, len: int, loc := #caller_location) -> bool {
 	array := (^Raw_Dynamic_Array)(array_)
 

+ 10 - 0
core/runtime/dynamic_map_internal.odin

@@ -239,6 +239,16 @@ __dynamic_map_reserve :: proc(using header: Map_Header, cap: int, loc := #caller
 	}
 }
 
+__dynamic_map_shrink :: proc(using header: Map_Header, cap: int, loc := #caller_location) -> (did_shrink: bool) {
+	c := context
+	if m.entries.allocator.procedure != nil {
+		c.allocator = m.entries.allocator
+	}
+	context = c
+
+	return __dynamic_array_shrink(&m.entries, entry_size, entry_align, cap, loc)
+}
+
 __dynamic_map_rehash :: proc(using header: Map_Header, new_count: int, loc := #caller_location) {
 	#force_inline __dynamic_map_reserve(header, new_count, loc)
 }