pkiter: add iterator for pkarr and pkbktarr

1 files changed, 106 insertions, 4 deletions

diff --git a/pkbktarr.h b/pkbktarr.h
index 29613ae..4f146e4 100644
--- a/pkbktarr.h
+++ b/pkbktarr.h
@@ -53,6 +53,7 @@ enum PK_BKT_ARR_HANDLE_VALIDATION pk_bkt_arr_handle_validate(struct pk_bkt_arr *
 
 void pk_bkt_arr_init(struct pk_bkt_arr *bkt_arr, unsigned long stride, unsigned long alignment, struct pk_bkt_arr_handle limits, struct pk_membucket *bkt_buckets, struct pk_membucket *bkt_data);
 void pk_bkt_arr_clear(struct pk_bkt_arr *bkt_arr);
+void pk_bkt_arr_reserve(struct pk_bkt_arr *bkt_arr, size_t count);
 struct pk_bkt_arr_handle pk_bkt_arr_find_first_handle(struct pk_bkt_arr *bkt_arr, pk_bkt_arr_compare_fn fn, void *user_data, const void *user_obj_data);
 void pk_bkt_arr_iterate(struct pk_bkt_arr *bkt_arr, pk_bkt_arr_iterate_fn fn, void *user_data);
 void pk_bkt_arr_teardown(struct pk_bkt_arr *bkt_arr);
@@ -62,6 +63,14 @@ int pk_bkt_arr_handle_compare(struct pk_bkt_arr_handle lhs, struct pk_bkt_arr_ha
 struct pk_bkt_arr_handle pk_bkt_arr_handle_increment(struct pk_bkt_arr *arr, struct pk_bkt_arr_handle h);
 struct pk_bkt_arr_handle pk_bkt_arr_handle_decrement(struct pk_bkt_arr *arr, struct pk_bkt_arr_handle h);
 
+#ifdef PK_IMPL_ITER
+#include "pkiter.h" /*deleteme*/
+bool pk_bkt_arr_iter_begin(struct pk_bkt_arr *arr, struct pk_iter *it);
+bool pk_bkt_arr_iter_end(struct pk_bkt_arr *arr, struct pk_iter *it);
+bool pk_bkt_arr_iter_increment(struct pk_bkt_arr *arr, struct pk_iter *it);
+bool pk_bkt_arr_iter_decrement(struct pk_bkt_arr *arr, struct pk_iter *it);
+#endif
+
 #if defined (__cplusplus)
 #include "pktmpln.h" /*deleteme*/
 #include <assert.h>
@@ -90,7 +99,7 @@ T &pk_bkt_arr_t<T>::operator[](struct pk_bkt_arr_handle handle) {
 	assert(this->bucketed_data != nullptr);
 	assert(handle.b <= this->limits.b);
 	assert(handle.i <= this->limits.i);
-	assert(handle.b != this->head_r.b || handle.i < this->head_r.i);
+	assert(handle.b != this->head_r.b || handle.i <= this->head_r.i);
 	T** two_star_programmer = reinterpret_cast<T**>(this->bucketed_data);
 	return two_star_programmer[handle.b][handle.i];
 }
@@ -108,9 +117,10 @@ enum PK_BKT_ARR_HANDLE_VALIDATION pk_bkt_arr_handle_validate(struct pk_bkt_arr *
 	if (handle.b >= bkt_arr->reserved_buckets || handle.b >= bkt_arr->limits.b) {
 		ret |= PK_BKT_ARR_HANDLE_VALIDATION_BUCKET_INDEX_TOO_HIGH;
 	}
-	if (handle.b == bkt_arr->reserved_buckets-1 && handle.i > bkt_arr->head_r.i) {
+	if (handle.i >= bkt_arr->limits.i) {
 		ret |= PK_BKT_ARR_HANDLE_VALIDATION_ITEM_INDEX_TOO_HIGH;
-	} else if (handle.b < bkt_arr->reserved_buckets-1 && handle.i > bkt_arr->limits.i) {
+	}
+	if (handle.b == bkt_arr->head_r.b && handle.i > bkt_arr->head_r.i) {
 		ret |= PK_BKT_ARR_HANDLE_VALIDATION_ITEM_INDEX_TOO_HIGH;
 	}
 	return (enum PK_BKT_ARR_HANDLE_VALIDATION)ret;
@@ -149,6 +159,26 @@ void pk_bkt_arr_clear(struct pk_bkt_arr *bkt_arr) {
 	}
 }
 
+void pk_bkt_arr_reserve(struct pk_bkt_arr *bkt_arr, size_t count) {
+	size_t bucket_count = count / bkt_arr->limits.i;
+	if (bkt_arr->reserved_buckets >= bucket_count) return;
+	unsigned long long *new_idx_unused = (unsigned long long *)pk_new(sizeof(unsigned long long) * bucket_count, alignof(uint64_t), bkt_arr->bkt_buckets);
+	void **new_bucketed_data = (void **)pk_new(sizeof(void *) * bucket_count, alignof(void *), bkt_arr->bkt_buckets);
+	if (bkt_arr->reserved_buckets > 0) {
+		memcpy(new_idx_unused, bkt_arr->idx_unused, sizeof(unsigned long long) * bkt_arr->reserved_buckets);
+		memcpy(new_bucketed_data, bkt_arr->bucketed_data, sizeof(void *) * bkt_arr->reserved_buckets);
+		pk_delete_bkt(bkt_arr->bucketed_data, sizeof(void *) * bkt_arr->reserved_buckets, bkt_arr->bkt_buckets);
+		pk_delete_bkt(bkt_arr->idx_unused, sizeof(unsigned long long) * bkt_arr->reserved_buckets, bkt_arr->bkt_buckets);
+	}
+	for (unsigned int i = bkt_arr->reserved_buckets; i < bucket_count; ++i) {
+		new_idx_unused[i] = PK_BKT_ARR_ALL_UNUSED_VAL;
+		new_bucketed_data[i] = pk_new(bkt_arr->stride * bkt_arr->limits.i, bkt_arr->alignment, bkt_arr->bkt_data);
+	}
+	bkt_arr->idx_unused = new_idx_unused;
+	bkt_arr->bucketed_data = new_bucketed_data;
+	bkt_arr->reserved_buckets = bucket_count;
+}
+
 struct pk_bkt_arr_handle pk_bkt_arr_find_first_handle(struct pk_bkt_arr *bkt_arr, pk_bkt_arr_compare_fn fn, void *user_data, const void *user_obj_data) {
 	assert(bkt_arr != NULL);
 	assert(fn != NULL);
@@ -179,7 +209,7 @@ void pk_bkt_arr_iterate(struct pk_bkt_arr *bkt_arr, pk_bkt_arr_iterate_fn fn, vo
 	unsigned int b, i, ii;
 	for (b = 0; b < bkt_arr->reserved_buckets; ++b) {
 		char *arr = ((char**)(bkt_arr->bucketed_data))[b];
-		ii = b == bkt_arr->reserved_buckets-1 ? bkt_arr->head_r.i : bkt_arr->limits.i;
+		ii = b == bkt_arr->head_r.b ? bkt_arr->head_r.i : bkt_arr->limits.i;
 		for (i = 0; i < ii; ++i) {
 			if (PK_HAS_FLAG(bkt_arr->idx_unused[b], 1ull << i)) {
 				continue;
@@ -306,4 +336,76 @@ struct pk_bkt_arr_handle pk_bkt_arr_handle_decrement(struct pk_bkt_arr *arr, str
 	return h;
 }
 
+#ifdef PK_IMPL_ITER
+
+bool pk_bkt_arr_iter_begin(struct pk_bkt_arr *arr, struct pk_iter *it) {
+	it->data = nullptr;
+	it->id.bkt.b = 0;
+	it->id.bkt.i = 0;
+	if ((arr->idx_unused[it->id.bkt.b] & (1ull << it->id.bkt.i)) != 0) return false;
+	it->data = arr->bucketed_data[0];
+	return true;
+}
+
+bool pk_bkt_arr_iter_end(struct pk_bkt_arr *arr, struct pk_iter *it) {
+	it->data = nullptr;
+	it->id.bkt.b = 0;
+	it->id.bkt.i = 0;
+	if (arr->head_r.b == 0 && arr->head_r.i == 0) return false;
+	do {
+		struct pk_bkt_arr_handle handle = arr->head_r;
+		for (;;) {
+			if ((arr->idx_unused[handle.b] & (1ull << handle.i)) == 0) break;
+			if (handle.b == 0 && handle.i == 0) return false;
+			handle = pk_bkt_arr_handle_decrement(arr, handle);
+		}
+		it->id.bkt.b = handle.b;
+		it->id.bkt.i = handle.i;
+		break;
+	} while (true);
+	if (arr->bucketed_data != nullptr && arr->bucketed_data[it->id.bkt.b] != nullptr) {
+		it->data = (char*)(arr->bucketed_data[it->id.bkt.b]) + (arr->stride * it->id.bkt.i);
+		return true;
+	}
+	return false;
+}
+
+bool pk_bkt_arr_iter_increment(struct pk_bkt_arr *arr, struct pk_iter *it) {
+	struct pk_bkt_arr_handle handle = {
+		.b = it->id.bkt.b,
+		.i = it->id.bkt.i,
+	};
+	if (it->id.bkt.b == arr->limits.b-1 && it->id.bkt.i == arr->limits.i-1) return false;
+	for (;;) {
+		handle = pk_bkt_arr_handle_increment(arr, handle);
+		if ((arr->idx_unused[handle.b] & (1ull << handle.i)) == 0) break;
+	}
+	it->id.bkt.b = handle.b;
+	it->id.bkt.i = handle.i;
+	if (handle.b >= arr->reserved_buckets) return false;
+	if ((arr->idx_unused[it->id.bkt.b] & (1ull << it->id.bkt.i)) != 0) return false;
+	it->data = (char*)(arr->bucketed_data[it->id.bkt.b]) + (arr->stride * it->id.bkt.i);
+	return true;
+}
+
+bool pk_bkt_arr_iter_decrement(struct pk_bkt_arr *arr, struct pk_iter *it) {
+	struct pk_bkt_arr_handle handle = {
+		.b = it->id.bkt.b,
+		.i = it->id.bkt.i,
+	};
+	for (;;) {
+		handle = pk_bkt_arr_handle_decrement(arr, handle);
+		if ((arr->idx_unused[handle.b] & (1ull << handle.i)) == 0) break;
+		if (handle.b == 0 && handle.i == 0) break;
+	}
+	if (it->id.bkt.b == handle.b && it->id.bkt.i == handle.i) return false;
+	it->id.bkt.b = handle.b;
+	it->id.bkt.i = handle.i;
+	if ((arr->idx_unused[it->id.bkt.b] & (1ull << it->id.bkt.i)) != 0) return false;
+	it->data = ((char*)(arr->bucketed_data[it->id.bkt.b])) + (arr->stride * it->id.bkt.i);
+	return true;
+}
+
+#endif /* PK_IMPL_ITER */
+
 #endif /* PK_IMPL_BKTARR */