Merge pull request #1469 from ap29600/bit_iterator

add `iterator` to `core:container/bit_array`

1 files changed, 111 insertions, 1 deletions

diff --git a/core/container/bit_array/bit_array.odin b/core/container/bit_array/bit_array.odin
index 0fa80c623..5eebe1bcb 100644
--- a/core/container/bit_array/bit_array.odin
+++ b/core/container/bit_array/bit_array.odin
@@ -11,13 +11,121 @@ INDEX_SHIFT :: 6
 @(private="file")
 INDEX_MASK  :: 63
 
+@(private="file")
+NUM_BITS :: 64
+
 Bit_Array :: struct {
 	bits: [dynamic]u64,
 	bias: int,
+	max_index: int,
+}
+
+Bit_Array_Iterator :: struct {
+	array: ^Bit_Array,
+	word_idx: int,
+	bit_idx: uint,
+}
+
+/*
+	In:
+		- ba:   ^Bit_Array - the array to iterate over
+
+	Out:
+		- it:   ^Bit_Array_Iterator - the iterator that holds iteration state
+*/
+make_iterator :: proc (ba: ^Bit_Array) -> (it: Bit_Array_Iterator) {
+	return Bit_Array_Iterator { array = ba }
+}
+
+/*
+	In:
+		- it:    ^Bit_Array_Iterator - the iterator struct that holds the state.
+
+	Out:
+		- set:    bool - the state of the bit at `index`
+		- index:  int - the next bit of the Bit_Array referenced by `it`.
+		- ok:	  bool - `true` if the iterator returned a valid index,
+			  `false` if there were no more bits
+*/
+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 }
+
+	word := it.array.bits[it.word_idx] if len(it.array.bits) > it.word_idx else 0
+	set = (word >> it.bit_idx & 1) == 1
+
+	it.bit_idx += 1
+	if it.bit_idx >= NUM_BITS {
+		it.bit_idx = 0
+		it.word_idx += 1
+	}
+
+	return set, index, true
 }
 
 /*
 	In:
+		- it:     ^Bit_Array_Iterator - the iterator struct that holds the state.
+
+	Out:
+		- index:  int - the next set bit of the Bit_Array referenced by `it`.
+		- ok:	  bool - `true` if the iterator returned a valid index,
+			  `false` if there were no more bits set
+*/
+iterate_by_set :: proc (it: ^Bit_Array_Iterator) -> (index: int, ok: bool) {
+	return iterate_internal_(it, true)
+}
+
+/*
+	In:
+		- it:	  ^Bit_Array_Iterator - the iterator struct that holds the state.
+
+	Out:
+		- index:  int - the next unset bit of the Bit_Array referenced by `it`.
+		- ok:	  bool - `true` if the iterator returned a valid index,
+			  `false` if there were no more unset bits
+*/
+iterate_by_unset:: proc (it: ^Bit_Array_Iterator) -> (index: int, ok: bool) {
+	return iterate_internal_(it, false)
+}
+
+@(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
+	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 {
+		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
+		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) {
+			return 0, false
+		}
+	}
+
+	// reaching here means that the word has some set bits
+	it.bit_idx += uint(intrinsics.count_trailing_zeros(word >> it.bit_idx))
+	index = it.word_idx * NUM_BITS + int(it.bit_idx) + it.array.bias
+
+	it.bit_idx += 1
+	if it.bit_idx >= NUM_BITS {
+		it.bit_idx = 0
+		it.word_idx += 1
+	}
+	return index, index <= it.array.max_index
+}
+
+
+/*
+	In:
 		- ba:    ^Bit_Array - a pointer to the Bit Array
 		- index: The bit index. Can be an enum member.
 
@@ -70,6 +178,7 @@ set :: proc(ba: ^Bit_Array, #any_int index: uint, allocator := context.allocator
 
 	resize_if_needed(ba, leg_index) or_return
 
+	ba.max_index = max(idx, ba.max_index)
 	ba.bits[leg_index] |= 1 << uint(bit_index)
 	return true
 }
@@ -87,6 +196,7 @@ create :: proc(max_index: int, min_index := 0, allocator := context.allocator) -
 
 	res = Bit_Array{
 		bias = min_index,
+		max_index = max_index,
 	}
 	return res, resize_if_needed(&res, legs)
 }
@@ -121,4 +231,4 @@ resize_if_needed :: proc(ba: ^Bit_Array, legs: int, allocator := context.allocat
 		resize(&ba.bits, legs + 1)
 	}
 	return len(ba.bits) > legs
-}
\ No newline at end of file
+}