util: implement iterator for hash tables and sets (#4639)

* Implement iterators for hash tables.
* Implement iterator over set.
* Constify get-size() of hash tables.
* Mention iterators in documentation.

DEVELOPERS.md

@@ -203,6 +203,9 @@ rz_core_wrap.cxx:32103:61: error: assigning to 'RzDebugReasonType' from incompat
 int sum = 0; // set sum to 0
 ```
 
+* If you want to iterate over values of your struct, implement `RzIterator *mystruct_as_iter()` and `RzIterator *mystruct_as_iter_mut()` for them.
+  See `rz_iterator.h` for details about the iterator.
+
 * If you need bitmaps, do not shift and OR the bits manually on `ut32`. Use bit vectors from `rz_bitvector.h` instead.
 
 ### Shell Scripts

librz/include/rz_util/ht_inc.h

@@ -4,6 +4,8 @@
 // SPDX-FileCopyrightText: 2024 pelijah
 // SPDX-License-Identifier: BSD-3-Clause
 
+#include <rz_util/rz_iterator.h>
+
 #ifndef HT_TYPE
 #error HT_TYPE should be defined before including this header
 #endif
@@ -166,6 +168,22 @@ typedef struct Ht_(t) {
 }
 HtName_(Ht);
 
+typedef struct Ht_(iter_mut_t) {
+	HtName_(Ht) *ht; ///< The hash table to iterate over.
+	ut32 ti; ///< Table index
+	ut32 bi; ///< Bucket index
+	HT_(Kv) *kv; ///< Current Key-Value-pair.
+}
+HT_(IterMutState);
+
+typedef struct Ht_(iter_t) {
+	const HtName_(Ht) *ht; ///< The hash table to iterate over.
+	ut32 ti; ///< Table index
+	ut32 bi; ///< Bucket index
+	const HT_(Kv) *kv; ///< Current Key-Value-pair.
+}
+HT_(IterState);
+
 // Create a new Ht with the provided Options
 RZ_API RZ_OWN HtName_(Ht) *Ht_(new_opt)(RZ_NONNULL HT_(Options) *opt);
 // Create a new Ht with the provided Options and initial size
@@ -190,6 +208,21 @@ RZ_API VALUE_TYPE Ht_(find)(RZ_NONNULL HtName_(Ht) *ht, const KEY_TYPE key, RZ_N
 // NOTE: cb can delete the current element, but it should be avoided
 RZ_API void Ht_(foreach)(RZ_NONNULL HtName_(Ht) *ht, RZ_NONNULL HT_(ForeachCallback) cb, RZ_NULLABLE void *user);
 
+RZ_API ut32 Ht_(size)(const RZ_NONNULL HtName_(Ht) *ht);
+
 RZ_API RZ_BORROW HT_(Kv) *Ht_(find_kv)(RZ_NONNULL HtName_(Ht) *ht, const KEY_TYPE key, RZ_NULLABLE bool *found);
 RZ_API bool Ht_(insert_kv)(RZ_NONNULL HtName_(Ht) *ht, RZ_NONNULL HT_(Kv) *kv, bool update);
 RZ_API HtRetCode Ht_(insert_kv_ex)(RZ_NONNULL HtName_(Ht) *ht, RZ_NONNULL HT_(Kv) *kv, bool update, RZ_OUT RZ_NULLABLE HT_(Kv) **out_kv);
+
+RZ_API RZ_OWN HT_(IterMutState) *Ht_(new_iter_mut_state)(RZ_NONNULL HtName_(Ht) *ht);
+RZ_API RZ_OWN HT_(IterState) *Ht_(new_iter_state)(const RZ_NONNULL HtName_(Ht) *ht);
+RZ_API void Ht_(free_iter_mut_state)(RZ_NULLABLE HT_(IterMutState) *state);
+RZ_API void Ht_(free_iter_state)(RZ_NULLABLE HT_(IterState) *state);
+
+RZ_API RZ_BORROW VALUE_TYPE *Ht_(iter_next_mut)(RzIterator *it);
+RZ_API const VALUE_TYPE *Ht_(iter_next)(RzIterator *it);
+RZ_API const KEY_TYPE *Ht_(iter_next_key)(RzIterator *it);
+
+RZ_API RZ_OWN RzIterator /* <HtName_(Ht)> */ *Ht_(as_iter_mut)(RZ_NONNULL HtName_(Ht) *ht);
+RZ_API RZ_OWN RzIterator /* <HtName_(Ht)> */ *Ht_(as_iter)(const RZ_NONNULL HtName_(Ht) *ht);
+RZ_API RZ_OWN RzIterator /* <HtName_(Ht)> */ *Ht_(as_iter_keys)(const RZ_NONNULL HtName_(Ht) *ht);

librz/include/rz_util/rz_set.h

@@ -20,7 +20,9 @@ RZ_API void rz_set_s_free(RZ_NULLABLE RzSetS *set);
 RZ_API void rz_set_s_add(RZ_NONNULL RzSetS *set, const char *str);
 RZ_API bool rz_set_s_contains(RZ_NONNULL RzSetS *set, const char *str);
 RZ_API void rz_set_s_delete(RZ_NONNULL RzSetS *set, const char *str);
+RZ_API ut32 rz_set_s_size(const RZ_NONNULL RzSetS *set);
 RZ_API RZ_OWN RzPVector /*<char *>*/ *rz_set_s_to_vector(RZ_NONNULL RzSetS *set);
+RZ_API RzIterator /* <RzSetS> */ *rz_set_s_as_iter(const RZ_NONNULL RzSetS *set);
 
 typedef HtUP RzSetU;
 
@@ -29,6 +31,8 @@ RZ_API void rz_set_u_free(RZ_NULLABLE RzSetU *set);
 RZ_API void rz_set_u_add(RZ_NONNULL RzSetU *set, ut64 u);
 RZ_API bool rz_set_u_contains(RZ_NONNULL RzSetU *set, ut64 u);
 RZ_API void rz_set_u_delete(RZ_NONNULL RzSetU *set, ut64 u);
+RZ_API ut32 rz_set_u_size(const RZ_NONNULL RzSetU *set);
+RZ_API RzIterator /* <RzSetU> */ *rz_set_u_as_iter(const RZ_NONNULL RzSetU *set);
 
 #ifdef __cplusplus
 }

librz/util/ht/ht_inc.c

@@ -4,7 +4,9 @@
 // SPDX-FileCopyrightText: 2024 pelijah
 // SPDX-License-Identifier: BSD-3-Clause
 
+#include <rz_util/rz_log.h>
 #include <rz_util/rz_assert.h>
+#include <rz_util/rz_iterator.h>
 #include <rz_util/rz_str.h>
 
 #define LOAD_FACTOR     1
@@ -512,3 +514,193 @@ RZ_API void Ht_(foreach)(RZ_NONNULL HtName_(Ht) *ht, RZ_NONNULL HT_(ForeachCallb
 		}
 	}
 }
+
+/**
+ * \brief Returns the number of elements stored in the hash map \p ht.
+ *
+ * \param ht The hash map.
+ *
+ * \return The number of elements saved in the hash map.
+ */
+RZ_API ut32 Ht_(size)(const RZ_NONNULL HtName_(Ht) *ht) {
+	rz_return_val_if_fail(ht, 0);
+	return ht->count;
+}
+
+/**
+ * \brief Advances an RzIterator over a hashtable to the next value and returns it.
+ *
+ * \param it The next value or NULL if iteration terminated. The value is mutable.
+ */
+RZ_API RZ_BORROW VALUE_TYPE *Ht_(iter_next_mut)(RzIterator *it) {
+	rz_return_val_if_fail(it, NULL);
+
+	HT_(IterMutState) *state = it->u;
+	if (state->ti >= state->ht->size) {
+		// Iteration is done. No elements left to select.
+		return NULL;
+	}
+	// Iterate over tables until a table with an element is found.
+	for (; state->ti < state->ht->size; state->ti++) {
+		if (state->ht->table[state->ti].count == 0) {
+			// Table has no elements. Check next table.
+			continue;
+		}
+		if (state->bi < state->ht->table[state->ti].count) {
+			// Table has elements, select the element.
+			state->kv = &state->ht->table[state->ti].arr[state->bi];
+			// For the next iteration, increment bucket index to the following element.
+			state->bi++;
+			return &state->kv->value;
+		}
+		// Reset bucket index to first bucket.
+		state->bi = 0;
+		// Go to next table
+	}
+	// Iteration is done. No elements left to select.
+	return NULL;
+}
+
+/**
+ * \brief Advances an RzIterator over a hash table to the next value and returns it.
+ *
+ * \param it The next value as immutable or NULL if iteration terminated.
+ */
+RZ_API const VALUE_TYPE *Ht_(iter_next)(RzIterator *it) {
+	rz_return_val_if_fail(it, NULL);
+
+	HT_(IterState) *state = it->u;
+	if (state->ti >= state->ht->size) {
+		// Iteration is done. No elements left to select.
+		return NULL;
+	}
+	// Iterate over tables until a table with an element is found.
+	for (; state->ti < state->ht->size; state->ti++) {
+		if (state->ht->table[state->ti].count == 0) {
+			// Table has no elements. Check next table.
+			continue;
+		}
+		if (state->bi < state->ht->table[state->ti].count) {
+			// Table has elements, select the element.
+			state->kv = &state->ht->table[state->ti].arr[state->bi];
+			// For the next iteration, increment bucket index to the following element.
+			state->bi++;
+			return (const VALUE_TYPE *)&state->kv->value;
+		}
+		// Reset bucket index to first bucket.
+		state->bi = 0;
+		// Go to next table
+	}
+	// Iteration is done. No elements left to select.
+	return NULL;
+}
+
+/**
+ * \brief Advances an RzIterator over a hash table to the next key in
+ * and returns it.
+ *
+ * \param it The next key as immutable or NULL if iteration terminated.
+ */
+RZ_API const KEY_TYPE *Ht_(iter_next_key)(RzIterator *it) {
+	rz_return_val_if_fail(it, NULL);
+
+	HT_(IterState) *state = it->u;
+	if (state->ti >= state->ht->size) {
+		// Iteration is done. No elements left to select.
+		return NULL;
+	}
+	// Iterate over tables until a table with an element is found.
+	for (; state->ti < state->ht->size; state->ti++) {
+		if (state->ht->table[state->ti].count == 0) {
+			// Table has no elements. Check next table.
+			continue;
+		}
+		if (state->bi < state->ht->table[state->ti].count) {
+			// Table has elements, select the element.
+			state->kv = &state->ht->table[state->ti].arr[state->bi];
+			// For the next iteration, increment bucket index to the following element.
+			state->bi++;
+			return (const KEY_TYPE *)&state->kv->key;
+		}
+		// Reset bucket index to first bucket.
+		state->bi = 0;
+		// Go to next table
+	}
+	// Iteration is done. No elements left to select.
+	return NULL;
+}
+
+RZ_API RZ_OWN HT_(IterMutState) *Ht_(new_iter_mut_state)(RZ_NONNULL HtName_(Ht) *ht) {
+	rz_return_val_if_fail(ht, NULL);
+	HT_(IterMutState) *state = RZ_NEW0(HT_(IterMutState));
+	rz_return_val_if_fail(state, NULL);
+	state->ht = ht;
+	return state;
+}
+
+RZ_API RZ_OWN HT_(IterState) *Ht_(new_iter_state)(const RZ_NONNULL HtName_(Ht) *ht) {
+	rz_return_val_if_fail(ht, NULL);
+	HT_(IterState) *state = RZ_NEW0(HT_(IterState));
+	rz_return_val_if_fail(state, NULL);
+	state->ht = ht;
+	return state;
+}
+
+RZ_API void Ht_(free_iter_mut_state)(RZ_NULLABLE HT_(IterMutState) *state) {
+	free(state);
+}
+
+RZ_API void Ht_(free_iter_state)(RZ_NULLABLE HT_(IterState) *state) {
+	free(state);
+}
+
+/**
+ * \brief Returns an iterator over the hash table \p ht. The iterator yields mutable values.
+ *
+ * \param ht The hash table to create the iterator for.
+ *
+ * \return The iterator over the hash table values or NULL in case of failure.
+ */
+RZ_API RZ_OWN RzIterator /* <HtName_(Ht)> */ *Ht_(as_iter_mut)(RZ_NONNULL HtName_(Ht) *ht) {
+	rz_return_val_if_fail(ht, NULL);
+	HT_(IterMutState) *state = Ht_(new_iter_mut_state)(ht);
+	if (!state) {
+		RZ_LOG_ERROR("Could not allocate a new ht_iter state.\n");
+		return NULL;
+	}
+
+	RzIterator *iter = rz_iterator_new((rz_iterator_next_cb)Ht_(iter_next_mut), NULL, (rz_iterator_free_cb)Ht_(free_iter_mut_state), state);
+	return iter;
+}
+
+/**
+ * \brief Returns an iterator over the hash table \p ht. The iterator yields immutable values.
+ *
+ * \param ht The hash table to create the iterator for.
+ *
+ * \return The iterator over the hash table values or NULL in case of failure.
+ */
+RZ_API RZ_OWN RzIterator /* <HtName_(Ht)> */ *Ht_(as_iter)(const RZ_NONNULL HtName_(Ht) *ht) {
+	rz_return_val_if_fail(ht, NULL);
+	HT_(IterState) *state = Ht_(new_iter_state)(ht);
+	rz_return_val_if_fail(state, NULL);
+
+	RzIterator *iter = rz_iterator_new((rz_iterator_next_cb)Ht_(iter_next), NULL, (rz_iterator_free_cb)Ht_(free_iter_state), state);
+	return iter;
+}
+
+/**
+ * \brief Returns an iterator over the hash table \p ht. The iterator yields immutable keys.
+ *
+ * \param ht The hash table to create the iterator for.
+ *
+ * \return The iterator over the hash table keys or NULL in case of failure.
+ */
+RZ_API RZ_OWN RzIterator /* <HtName_(Ht)> */ *Ht_(as_iter_keys)(const RZ_NONNULL HtName_(Ht) *ht) {
+	rz_return_val_if_fail(ht, NULL);
+	HT_(IterState) *state = Ht_(new_iter_state)(ht);
+	rz_return_val_if_fail(state, NULL);
+
+	RzIterator *iter = rz_iterator_new((rz_iterator_next_cb)Ht_(iter_next_key), NULL, (rz_iterator_free_cb)Ht_(free_iter_state), state);
+	return iter;
+}

librz/util/set.c

@@ -37,6 +37,16 @@ RZ_API void rz_set_s_delete(RZ_NONNULL RzSetS *set, const char *str) {
 	ht_sp_delete(set, str);
 }
 
+/**
+ * \brief Returns an iterator over the \p set with immutable elements.
+ *
+ * \return Iterator yielding immutable elements.
+ */
+RZ_API RzIterator /* <RzSetS> */ *rz_set_s_as_iter(const RZ_NONNULL RzSetS *set) {
+	rz_return_val_if_fail(set, NULL);
+	return ht_sp_as_iter_keys((const HtSP *)set);
+}
+
 static bool push_to_pvector(void *user, const char *k, RZ_UNUSED const void *v) {
 	RzPVector *vec = (RzPVector *)user;
 	return !!rz_pvector_push(vec, (void *)k);
@@ -65,6 +75,14 @@ RZ_API void rz_set_s_free(RZ_NULLABLE RzSetS *set) {
 	ht_sp_free((HtSP *)set);
 }
 
+/**
+ * \brief Return number of elements saved in the set.
+ */
+RZ_API ut32 rz_set_s_size(const RZ_NONNULL RzSetS *set) {
+	rz_return_val_if_fail(set, 0);
+	return ht_sp_size((HtSP *)set);
+}
+
 /**
  * \brief Create a new hash set with ut64 as elements.
  */
@@ -99,3 +117,21 @@ RZ_API void rz_set_u_delete(RZ_NONNULL RzSetU *set, ut64 u) {
 RZ_API void rz_set_u_free(RZ_NULLABLE RzSetU *set) {
 	ht_up_free((HtUP *)set);
 }
+
+/**
+ * \brief Return number of elements saved in the set.
+ */
+RZ_API ut32 rz_set_u_size(const RZ_NONNULL RzSetU *set) {
+	rz_return_val_if_fail(set, 0);
+	return ht_up_size((HtUP *)set);
+}
+
+/**
+ * \brief Returns an iterator over the \p set with immutable elements.
+ *
+ * \return Iterator yielding immutable elements.
+ */
+RZ_API RzIterator /* <RzSetU> */ *rz_set_u_as_iter(const RZ_NONNULL RzSetU *set) {
+	rz_return_val_if_fail(set, NULL);
+	return ht_up_as_iter_keys((const HtUP *)set);
+}

test/unit/meson.build

@@ -54,6 +54,7 @@ if get_option('enable_tests')
     'graph',
     'hash',
     'hex',
+    'ht',
     'id_storage',
     'idpool',
     'idstorage',
@@ -81,7 +82,6 @@ if get_option('enable_tests')
     'rz_test',
     'sdb_array',
     'sdb_diff',
-    'sdb_hash',
     'sdb_sdb',
     'sdb_util',
     'serialize_analysis',