Merge pull request #4194 from Feoramund/update-bit-array

Update `bit_array`

3 files changed, 307 insertions, 17 deletions

diff --git a/core/container/bit_array/bit_array.odin b/core/container/bit_array/bit_array.odin
index b53bacda7..9a76dc78f 100644
--- a/core/container/bit_array/bit_array.odin
+++ b/core/container/bit_array/bit_array.odin
@@ -1,5 +1,6 @@
 package container_dynamic_bit_array
 
+import "base:builtin"
 import "base:intrinsics"
 import "core:mem"
 
@@ -18,7 +19,7 @@ NUM_BITS :: 64
 Bit_Array :: struct {
 	bits:         [dynamic]u64,
 	bias:         int,
-	max_index:    int,
+	length:       int,
 	free_pointer: bool,
 }
 
@@ -52,9 +53,9 @@ Returns:
 */
 iterate_by_all :: proc (it: ^Bit_Array_Iterator) -> (set: bool, index: int, ok: bool) {
 	index = it.word_idx * NUM_BITS + int(it.bit_idx) + it.array.bias
-	if index > it.array.max_index { return false, 0, false }
+	if index >= it.array.length + it.array.bias { return false, 0, false }
 
-	word := it.array.bits[it.word_idx] if len(it.array.bits) > it.word_idx else 0
+	word := it.array.bits[it.word_idx] if builtin.len(it.array.bits) > it.word_idx else 0
 	set = (word >> it.bit_idx & 1) == 1
 
 	it.bit_idx += 1
@@ -106,22 +107,22 @@ Returns:
 */
 @(private="file")
 iterate_internal_ :: proc (it: ^Bit_Array_Iterator, $ITERATE_SET_BITS: bool) -> (index: int, ok: bool) {
-	word := it.array.bits[it.word_idx] if len(it.array.bits) > it.word_idx else 0
+	word := it.array.bits[it.word_idx] if builtin.len(it.array.bits) > it.word_idx else 0
 	when ! ITERATE_SET_BITS { word = ~word }
 
 	// If the word is empty or we have already gone over all the bits in it,
 	// b.bit_idx is greater than the index of any set bit in the word,
 	// meaning that word >> b.bit_idx == 0.
-	for it.word_idx < len(it.array.bits) && word >> it.bit_idx == 0 {
+	for it.word_idx < builtin.len(it.array.bits) && word >> it.bit_idx == 0 {
 		it.word_idx += 1
 		it.bit_idx = 0
-		word = it.array.bits[it.word_idx] if len(it.array.bits) > it.word_idx else 0
+		word = it.array.bits[it.word_idx] if builtin.len(it.array.bits) > it.word_idx else 0
 		when ! ITERATE_SET_BITS { word = ~word }
 	}
 
 	// If we are iterating the set bits, reaching the end of the array means we have no more bits to check
 	when ITERATE_SET_BITS {
-		if it.word_idx >= len(it.array.bits) {
+		if it.word_idx >= builtin.len(it.array.bits) {
 			return 0, false
 		}
 	}
@@ -135,7 +136,7 @@ iterate_internal_ :: proc (it: ^Bit_Array_Iterator, $ITERATE_SET_BITS: bool) ->
 		it.bit_idx = 0
 		it.word_idx += 1
 	}
-	return index, index <= it.array.max_index
+	return index, index < it.array.length + it.array.bias
 }
 /*
 Gets the state of a bit in the bit-array
@@ -160,7 +161,7 @@ get :: proc(ba: ^Bit_Array, #any_int index: uint) -> (res: bool, ok: bool) #opti
 		If we `get` a bit that doesn't fit in the Bit Array, it's naturally `false`.
 		This early-out prevents unnecessary resizing.
 	*/
-	if leg_index + 1 > len(ba.bits) { return false, true }
+	if leg_index + 1 > builtin.len(ba.bits) { return false, true }
 
 	val := u64(1 << uint(bit_index))
 	res = ba.bits[leg_index] & val == val
@@ -208,7 +209,7 @@ set :: proc(ba: ^Bit_Array, #any_int index: uint, set_to: bool = true, allocator
 
 	resize_if_needed(ba, leg_index) or_return
 
-	ba.max_index = max(idx, ba.max_index)
+	ba.length = max(1 + idx, ba.length)
 
 	if set_to {
 		ba.bits[leg_index] |=  1 << uint(bit_index)
@@ -261,6 +262,9 @@ unsafe_unset :: proc(b: ^Bit_Array, bit: int) #no_bounds_check {
 /*
 A helper function to create a Bit Array with optional bias, in case your smallest index is non-zero (including negative).
 
+The range of bits created by this procedure is `min_index..<max_index`, and the
+array will be able to expand beyond `max_index` if needed.
+
 *Allocates (`new(Bit_Array) & make(ba.bits)`)*
 
 Inputs:
@@ -275,7 +279,7 @@ create :: proc(max_index: int, min_index: int = 0, allocator := context.allocato
 	context.allocator = allocator
 	size_in_bits := max_index - min_index
 
-	if size_in_bits < 1 { return {}, false }
+	if size_in_bits < 0 { return {}, false }
 
 	legs := size_in_bits >> INDEX_SHIFT
 	if size_in_bits & INDEX_MASK > 0 {legs+=1}
@@ -284,7 +288,7 @@ create :: proc(max_index: int, min_index: int = 0, allocator := context.allocato
 	res = new(Bit_Array)
 	res.bits         = bits
 	res.bias         = min_index
-	res.max_index    = max_index
+	res.length       = max_index - min_index
 	res.free_pointer = true
 	return
 }
@@ -299,6 +303,48 @@ clear :: proc(ba: ^Bit_Array) {
 	mem.zero_slice(ba.bits[:])
 }
 /*
+Gets the length of set and unset valid bits in the Bit_Array.
+
+Inputs:
+- ba: The target Bit_Array
+
+Returns:
+- length: The length of valid bits.
+*/
+len :: proc(ba: ^Bit_Array) -> (length: int) {
+	if ba == nil { return }
+	return ba.length
+}
+/*
+Shrinks the Bit_Array's backing storage to the smallest possible size.
+
+Inputs:
+- ba: The target Bit_Array
+*/
+shrink :: proc(ba: ^Bit_Array) #no_bounds_check {
+	if ba == nil { return }
+	legs_needed := builtin.len(ba.bits)
+	for i := legs_needed - 1; i >= 0; i -= 1 {
+		if ba.bits[i] == 0 {
+			legs_needed -= 1
+		} else {
+			break
+		}
+	}
+	if legs_needed == builtin.len(ba.bits) {
+		return
+	}
+	ba.length = 0
+	if legs_needed > 0 {
+		if legs_needed > 1 {
+			ba.length = (legs_needed - 1) * NUM_BITS
+		}
+		ba.length += NUM_BITS - int(intrinsics.count_leading_zeros(ba.bits[legs_needed - 1]))
+	}
+	resize(&ba.bits, legs_needed)
+	builtin.shrink(&ba.bits)
+}
+/*
 Deallocates the Bit_Array and its backing storage
 
 Inputs:
@@ -321,8 +367,8 @@ resize_if_needed :: proc(ba: ^Bit_Array, legs: int, allocator := context.allocat
 
 	context.allocator = allocator
 
-	if legs + 1 > len(ba.bits) {
+	if legs + 1 > builtin.len(ba.bits) {
 		resize(&ba.bits, legs + 1)
 	}
-	return len(ba.bits) > legs
+	return builtin.len(ba.bits) > legs
 }
diff --git a/core/flags/internal_assignment.odin b/core/flags/internal_assignment.odin
index bb49977eb..be3997ef1 100644
--- a/core/flags/internal_assignment.odin
+++ b/core/flags/internal_assignment.odin
@@ -14,10 +14,10 @@ import "core:reflect"
 // positionals first before adding it to a fallback field.
 @(optimization_mode="favor_size")
 push_positional :: #force_no_inline proc (model: ^$T, parser: ^Parser, arg: string) -> (error: Error) {
-	if bit_array.get(&parser.filled_pos, parser.filled_pos.max_index) {
-		// The max index is set, which means we're out of space.
+	if set, valid_index := bit_array.get(&parser.filled_pos, parser.filled_pos.length - 1); set || !valid_index {
+		// The index below the last one is either set or invalid, which means we're out of space.
 		// Add one free bit by setting the index above to false.
-		bit_array.set(&parser.filled_pos, 1 + parser.filled_pos.max_index, false)
+		bit_array.set(&parser.filled_pos, parser.filled_pos.length, false)
 	}
 
 	pos: int = ---
diff --git a/tests/core/container/test_core_bit_array.odin b/tests/core/container/test_core_bit_array.odin
new file mode 100644
index 000000000..8283d73f9
--- /dev/null
+++ b/tests/core/container/test_core_bit_array.odin
@@ -0,0 +1,244 @@
+package test_core_container
+
+import "core:container/bit_array"
+import "core:log"
+import "core:math/rand"
+import "core:slice"
+import "core:testing"
+
+ELEM_BIT_SIZE :: 8 * size_of(u64)
+
+@test
+test_bit_array_bias :: proc(t: ^testing.T) {
+	for bias in -ELEM_BIT_SIZE..=ELEM_BIT_SIZE {
+		M :: 19
+		list := []int{0,1,3,5,7,11,13,17,M}
+
+		ba := bit_array.create(M + bias, bias)
+		defer bit_array.destroy(ba)
+
+		for v, i in list {
+			list[i] = v + bias
+		}
+
+		for i in list {
+			bit_array.set(ba, i)
+			testing.expectf(t, bit_array.get(ba, i),
+				"Expected bit_array<length: %i, bias: %i>[%i] to be true",
+				ba.length, ba.bias, i)
+		}
+
+		seen: [dynamic]int
+		defer delete(seen)
+
+		iter := bit_array.make_iterator(ba)
+		for i in bit_array.iterate_by_set(&iter) {
+			append(&seen, i)
+		}
+
+		testing.expectf(t, slice.equal(list, seen[:]),
+			"Expected bit_array<length: %i, bias: %i> to be: %v, got %v",
+			ba.length, ba.bias, list, seen)
+	}
+}
+
+@test
+test_bit_array_empty_iteration :: proc(t: ^testing.T) {
+	ba: ^bit_array.Bit_Array = &{}
+	defer bit_array.destroy(ba)
+
+	for x in 0..=1 {
+		if x == 1 {
+			// Run the same tests with a created bit_array.
+			ba = bit_array.create(0,0)
+		}
+
+		iter := bit_array.make_iterator(ba)
+		for v, i in bit_array.iterate_by_all(&iter) {
+			log.errorf("Empty bit array had iterable: %v, %i", v, i)
+		}
+
+		iter = bit_array.make_iterator(ba)
+		for i in bit_array.iterate_by_unset(&iter) {
+			log.errorf("Empty bit array had iterable: %v", i)
+		}
+	}
+}
+
+@test
+test_bit_array_biased_max_index :: proc(t: ^testing.T) {
+	for bias in -ELEM_BIT_SIZE..=ELEM_BIT_SIZE {
+		for max_index in 1+bias..<ELEM_BIT_SIZE {
+			length := max_index - bias
+			ba := bit_array.create(max_index, bias)
+			defer bit_array.destroy(ba)
+
+			bit_array.set(ba, max_index - 1)
+
+			expected := max_index - bias
+			testing.expectf(t, ba.length == expected,
+				"Expected bit_array<max_index: %i, bias: %i> length to be: %i, got %i",
+				max_index, bias, expected, ba.length)
+
+			list := make([]int, length)
+			defer delete(list)
+			for i in 0..<len(list) {
+				list[i] = i + bias
+			}
+
+			seen: [dynamic]int
+			defer delete(seen)
+
+			iter := bit_array.make_iterator(ba)
+			for _, i in bit_array.iterate_by_all(&iter) {
+				append(&seen, i)
+			}
+			testing.expectf(t, slice.equal(list[:], seen[:]),
+				"Expected bit_array<max_index: %i, bias: %i> to contain: %v, got %v",
+				max_index, bias, list, seen)
+		}
+	}
+}
+
+@test
+test_bit_array_shrink :: proc(t: ^testing.T) {
+	for bias in -ELEM_BIT_SIZE..=ELEM_BIT_SIZE {
+		ba := bit_array.create(bias, bias)
+		defer bit_array.destroy(ba)
+
+		N :: 3*ELEM_BIT_SIZE
+
+		for i in 0..=N {
+			biased_i := bias + i
+			bit_array.set(ba, biased_i)
+
+			testing.expectf(t, bit_array.get(ba, biased_i),
+				"Expected bit_array<bias: %i>[%i] to be true",
+				ba.bias, biased_i)
+			testing.expectf(t, ba.length == 1 + i,
+				"Expected bit_array<bias: %i> length to be %i, got %i",
+				ba.bias, 1 + i, ba.length)
+
+			legs := 1 + i / ELEM_BIT_SIZE
+
+			testing.expectf(t, len(ba.bits) == legs,
+				"Expected bit_array<bias: %i> to have %i legs with index %i set, had %i legs",
+				ba.bias, legs, biased_i, len(ba.bits))
+
+			bit_array.unset(ba, biased_i)
+
+			if i >= ELEM_BIT_SIZE {
+				// Test shrinking arrays with bits set across two legs.
+				bit_array.set(ba, bias)
+				bit_array.shrink(ba)
+
+				testing.expectf(t, ba.length == 1,
+					"Expected bit_array<bias: %i> length to be 1 after >1 leg shrink, got %i",
+					ba.bias, ba.length)
+				testing.expectf(t, len(ba.bits) == 1,
+					"Expected bit_array<bias: %i> to have one leg after >1 leg shrink, had %i",
+					ba.bias, len(ba.bits))
+
+				bit_array.unset(ba, bias)
+			}
+
+			bit_array.shrink(ba)
+
+			testing.expectf(t, ba.length == 0,
+				"Expected bit_array<bias: %i> length to be zero after final shrink, got %i",
+				ba.bias, ba.length)
+			testing.expectf(t, len(ba.bits) == 0,
+				"Expected bit_array<bias: %i> to have zero legs with index %i set after final shrink, had %i",
+				ba.bias, biased_i, len(ba.bits))
+		}
+	}
+}
+
+@test
+test_bit_array :: proc(t: ^testing.T) {
+	ba := bit_array.create(0, 0)
+	defer bit_array.destroy(ba)
+
+	list_set: [dynamic]int
+	seen_set: [dynamic]int
+	list_unset: [dynamic]int
+	seen_unset: [dynamic]int
+	defer {
+		delete(list_set)
+		delete(seen_set)
+		delete(list_unset)
+		delete(seen_unset)
+	}
+
+	// Setup bits.
+	MAX_INDEX :: 1+16*ELEM_BIT_SIZE
+	for i in 0..=MAX_INDEX {
+		append(&list_unset, i)
+	}
+	for i in 1..=16 {
+		for j in -1..=1 {
+			n := ELEM_BIT_SIZE * i + j
+			bit_array.set(ba, n)
+			append(&list_set, n)
+		}
+	}
+	#reverse for i in list_set {
+		ordered_remove(&list_unset, i)
+	}
+
+	// Test iteration.
+	iter := bit_array.make_iterator(ba)
+	for i in bit_array.iterate_by_set(&iter) {
+		append(&seen_set, i)
+	}
+	testing.expectf(t, slice.equal(list_set[:], seen_set[:]),
+		"Expected set bit_array to be: %v, got %v",
+		list_set, seen_set)
+
+	iter = bit_array.make_iterator(ba)
+	for i in bit_array.iterate_by_unset(&iter) {
+		append(&seen_unset, i)
+	}
+	testing.expectf(t, slice.equal(list_unset[:], seen_unset[:]),
+		"Expected unset bit_array to be: %v, got %v",
+		list_unset, seen_unset)
+
+	// Test getting.
+	for i in list_set {
+		testing.expectf(t, bit_array.get(ba, i),
+			"Expected index %i to be true, got false",
+			i)
+	}
+	for i in list_unset {
+		testing.expectf(t, bit_array.get(ba, i) == false,
+			"Expected index %i to be false, got true",
+			i)
+	}
+
+	// Test flipping bits.
+	rand.shuffle(list_set[:])
+	rand.shuffle(list_unset[:])
+
+	for i in list_set {
+		bit_array.unset(ba, i)
+		testing.expectf(t, bit_array.get(ba, i) == false,
+			"Expected index %i to be false after unsetting, got true",
+			i)
+	}
+
+	for i in list_unset {
+		bit_array.set(ba, i)
+		testing.expectf(t, bit_array.get(ba, i),
+			"Expected index %i to be true after setting, got false",
+			i)
+	}
+
+	// Test clearing.
+	bit_array.clear(ba)
+	iter = bit_array.make_iterator(ba)
+	for i in 0..=MAX_INDEX {
+		testing.expectf(t, bit_array.get(ba, i) == false,
+			"Expected index %i to be false after clearing, got true",
+			i)
+	}
+}