Support for iterating over the leaves in a qp-trie

The iterator object records a path through the trie, in a similar
manner to the existing dns_rbtnodechain.

lib/dns/include/dns/qp.h

@@ -154,10 +154,7 @@ typedef union dns_qpreadable {
 #define dns_qpreader(qpr) ((qpr).qp)
 
 /*%
- * A trie lookup key is a small array, allocated on the stack during trie
- * searches. Keys are usually created on demand from DNS names using
- * `dns_qpkey_fromname()`, but in principle you can define your own
- * functions to convert other types to trie lookup keys.
+ * The maximum size of a key is also the maximum depth of a trie.
  *
  * A domain name can be up to 255 bytes. When converted to a key, each
  * character in the name corresponds to one byte in the key if it is a
@@ -165,7 +162,29 @@ typedef union dns_qpreadable {
  * using two bytes in the key. So we allow keys to be up to 512 bytes.
  * (The actual max is (255 - 5) * 2 + 6 == 506)
  */
-typedef uint8_t dns_qpkey_t[512];
+#define DNS_QP_MAXKEY 512
+
+/*%
+ * A trie lookup key is a small array, allocated on the stack during trie
+ * searches. Keys are usually created on demand from DNS names using
+ * `dns_qpkey_fromname()`, but in principle you can define your own
+ * functions to convert other types to trie lookup keys.
+ */
+typedef uint8_t dns_qpkey_t[DNS_QP_MAXKEY];
+
+/*%
+ * A trie iterator describes a path through the trie from the root to
+ * a leaf node, for use with `dns_qpiter_init()` and `dns_qpiter_next()`.
+ */
+typedef struct dns_qpiter {
+	unsigned int	magic;
+	dns_qpreader_t *qp;
+	uint16_t	sp;
+	struct __attribute__((__packed__)) {
+		uint32_t ref;
+		uint8_t	 more;
+	} stack[DNS_QP_MAXKEY];
+} dns_qpiter_t;
 
 /*%
  * These leaf methods allow the qp-trie code to call back to the code
@@ -376,16 +395,13 @@ dns_qpmulti_memusage(dns_qpmulti_t *multi);
 /*
  * XXXFANF todo, based on what we discover BIND needs
  *
- * fancy searches: longest match, lexicographic predecessor,
- * etc.
+ * fancy searches: longest match, lexicographic predecessor
+ * (for NSEC), successor (for modification-safe iteration), etc.
  *
  * do we need specific lookup functions to find out if the
  * returned value is readonly or mutable?
  *
  * richer modification such as dns_qp_replace{key,name}
- *
- * iteration - probably best to put an explicit stack in the iterator,
- * cf. rbtnodechain
  */
 
 size_t
@@ -484,6 +500,48 @@ dns_qp_deletename(dns_qp_t *qp, const dns_name_t *name);
  * \li  ISC_R_SUCCESS if the leaf was deleted from the trie
  */
 
+void
+dns_qpiter_init(dns_qpreadable_t qpr, dns_qpiter_t *qpi);
+/*%<
+ * Initialize an iterator
+ *
+ * SAFETY NOTE: If `qpr` is a `dns_qp_t`, it is not safe to modify the
+ * trie during iteration. If `qpr` is a `dns_qpread_t` or `dns_qpsnap_t`
+ * then (like any other read-only access) modifications will not affect
+ * iteration.
+ *
+ * Requires:
+ * \li  `qp` is a pointer to a valid qp-trie
+ * \li  `qpi` is a pointer to a qp iterator
+ */
+
+isc_result_t
+dns_qpiter_next(dns_qpiter_t *qpi, void **pval_r, uint32_t *ival_r);
+/*%<
+ * Get the next leaf object of a trie in lexicographic order of its keys.
+ *
+ * NOTE: see the safety note under `dns_qpiter_init()`.
+ *
+ * For example,
+ *
+ *	dns_qpiter_t qpi;
+ *	void *pval;
+ *	uint32_t ival;
+ *	dns_qpiter_init(qp, &qpi);
+ *	while (dns_qpiter_next(&qpi, &pval, &ival)) {
+ *		// do something with pval and ival
+ *	}
+ *
+ * Requires:
+ * \li  `qpi` is a pointer to a valid qp iterator
+ * \li  `pval_r != NULL`
+ * \li  `ival_r != NULL`
+ *
+ * Returns:
+ * \li  ISC_R_SUCCESS if a leaf was found and pval_r and ival_r were set
+ * \li  ISC_R_NOMORE otherwise
+ */
+
 /***********************************************************************
  *
  *  functions - transactions

lib/dns/qp.c

@@ -1681,6 +1681,72 @@ dns_qp_deletename(dns_qp_t *qp, const dns_name_t *name) {
 	return (dns_qp_deletekey(qp, key, keylen));
 }
 
+/***********************************************************************
+ *
+ *  iterate
+ */
+
+void
+dns_qpiter_init(dns_qpreadable_t qpr, dns_qpiter_t *qpi) {
+	dns_qpreader_t *qp = dns_qpreader(qpr);
+	REQUIRE(QP_VALID(qp));
+	REQUIRE(qpi != NULL);
+	qpi->magic = QPITER_MAGIC;
+	qpi->qp = qp;
+	qpi->sp = 0;
+	qpi->stack[qpi->sp].ref = qp->root_ref;
+	qpi->stack[qpi->sp].more = 0;
+}
+
+/*
+ * note: this function can go wrong when the iterator refers to
+ * a mutable view of the trie which is altered while iterating
+ */
+isc_result_t
+dns_qpiter_next(dns_qpiter_t *qpi, void **pval_r, uint32_t *ival_r) {
+	REQUIRE(QPITER_VALID(qpi));
+	REQUIRE(QP_VALID(qpi->qp));
+	REQUIRE(pval_r != NULL);
+	REQUIRE(ival_r != NULL);
+
+	dns_qpreader_t *qp = qpi->qp;
+
+	if (qpi->stack[qpi->sp].ref == INVALID_REF) {
+		INSIST(qpi->sp == 0);
+		qpi->magic = 0;
+		return (ISC_R_NOMORE);
+	}
+
+	/* push branch nodes onto the stack until we reach a leaf */
+	for (;;) {
+		qp_node_t *n = ref_ptr(qp, qpi->stack[qpi->sp].ref);
+		if (node_tag(n) == LEAF_TAG) {
+			*pval_r = leaf_pval(n);
+			*ival_r = leaf_ival(n);
+			break;
+		}
+		qpi->sp++;
+		INSIST(qpi->sp < DNS_QP_MAXKEY);
+		qpi->stack[qpi->sp].ref = branch_twigs_ref(n);
+		qpi->stack[qpi->sp].more = branch_twigs_size(n) - 1;
+	}
+
+	/* pop the stack until we find a twig with a successor */
+	while (qpi->sp > 0 && qpi->stack[qpi->sp].more == 0) {
+		qpi->sp--;
+	}
+
+	/* move across to the next twig */
+	if (qpi->stack[qpi->sp].more > 0) {
+		qpi->stack[qpi->sp].more--;
+		qpi->stack[qpi->sp].ref++;
+	} else {
+		INSIST(qpi->sp == 0);
+		qpi->stack[qpi->sp].ref = INVALID_REF;
+	}
+	return (ISC_R_SUCCESS);
+}
+
 /***********************************************************************
  *
  *  search

lib/dns/qp_p.h

@@ -379,11 +379,13 @@ ref_ptr(dns_qpreadable_t qpr, qp_ref_t ref) {
  */
 
 #define QP_MAGIC       ISC_MAGIC('t', 'r', 'i', 'e')
+#define QPITER_MAGIC   ISC_MAGIC('q', 'p', 'i', 't')
 #define QPMULTI_MAGIC  ISC_MAGIC('q', 'p', 'm', 'v')
 #define QPREADER_MAGIC ISC_MAGIC('q', 'p', 'r', 'x')
 
-#define QP_VALID(qp)	  ISC_MAGIC_VALID(qp, QP_MAGIC)
-#define QPMULTI_VALID(qp) ISC_MAGIC_VALID(qp, QPMULTI_MAGIC)
+#define QP_VALID(p)	 ISC_MAGIC_VALID(p, QP_MAGIC)
+#define QPITER_VALID(p)	 ISC_MAGIC_VALID(p, QPITER_MAGIC)
+#define QPMULTI_VALID(p) ISC_MAGIC_VALID(p, QPMULTI_MAGIC)
 
 /*
  * Polymorphic initialization of the `dns_qpreader_t` prefix.

tests/dns/qp_test.c

@@ -22,6 +22,9 @@
 #define UNIT_TESTING
 #include <cmocka.h>
 
+#include <isc/qsbr.h>
+#include <isc/random.h>
+#include <isc/refcount.h>
 #include <isc/result.h>
 #include <isc/string.h>
 #include <isc/util.h>
@@ -29,6 +32,8 @@
 #include <dns/name.h>
 #include <dns/qp.h>
 
+#include "qp_p.h"
+
 #include <tests/dns.h>
 #include <tests/qp.h>
 
@@ -129,9 +134,92 @@ ISC_RUN_TEST_IMPL(qpkey_sort) {
 	}
 }
 
+#define ITER_ITEMS 100
+
+static uint32_t
+check_leaf(void *uctx, void *pval, uint32_t ival) {
+	uint32_t *items = uctx;
+	assert_in_range(ival, 1, ITER_ITEMS - 1);
+	assert_ptr_equal(items + ival, pval);
+	return (1);
+}
+
+static size_t
+qpiter_makekey(dns_qpkey_t key, void *uctx, void *pval, uint32_t ival) {
+	check_leaf(uctx, pval, ival);
+
+	char str[8];
+	snprintf(str, sizeof(str), "%03u", ival);
+
+	size_t i = 0;
+	while (str[i] != '\0') {
+		key[i] = str[i] - '0' + SHIFT_BITMAP;
+		i++;
+	}
+	key[i++] = SHIFT_NOBYTE;
+
+	return (i);
+}
+
+static void
+getname(void *uctx, char *buf, size_t size) {
+	strlcpy(buf, "test", size);
+	UNUSED(uctx);
+	UNUSED(size);
+}
+
+const struct dns_qpmethods qpiter_methods = {
+	check_leaf,
+	check_leaf,
+	qpiter_makekey,
+	getname,
+};
+
+ISC_RUN_TEST_IMPL(qpiter) {
+	dns_qp_t *qp = NULL;
+	uint32_t item[ITER_ITEMS] = { 0 };
+
+	dns_qp_create(mctx, &qpiter_methods, item, &qp);
+	for (size_t tests = 0; tests < 1234; tests++) {
+		uint32_t ival = isc_random_uniform(ITER_ITEMS - 1) + 1;
+		void *pval = &item[ival];
+		item[ival] = ival;
+
+		/* randomly insert or remove */
+		dns_qpkey_t key;
+		size_t len = qpiter_makekey(key, item, pval, ival);
+		if (dns_qp_insert(qp, pval, ival) == ISC_R_EXISTS) {
+			dns_qp_deletekey(qp, key, len);
+			item[ival] = 0;
+		}
+
+		/* check that we see only valid items in the correct order */
+		uint32_t prev = 0;
+		dns_qpiter_t qpi;
+		dns_qpiter_init(qp, &qpi);
+		while (dns_qpiter_next(&qpi, &pval, &ival) == ISC_R_SUCCESS) {
+			assert_in_range(ival, prev + 1, ITER_ITEMS - 1);
+			assert_int_equal(ival, item[ival]);
+			assert_ptr_equal(pval, &item[ival]);
+			item[ival] = ~ival;
+			prev = ival;
+		}
+
+		/* ensure we saw every item */
+		for (ival = 0; ival < ITER_ITEMS; ival++) {
+			if (item[ival] != 0) {
+				assert_int_equal(item[ival], ~ival);
+				item[ival] = ival;
+			}
+		}
+	}
+	dns_qp_destroy(&qp);
+}
+
 ISC_TEST_LIST_START
 ISC_TEST_ENTRY(qpkey_name)
 ISC_TEST_ENTRY(qpkey_sort)
+ISC_TEST_ENTRY(qpiter)
 ISC_TEST_LIST_END
 
 ISC_TEST_MAIN