What Eric and Zakir said

lib/blacklist.c

@@ -125,5 +125,20 @@ int blacklist_init_from_files(char *whitelist_filename, char *blacklist_filename
 	uint64_t allowed = blacklist_count_allowed();
 	log_debug("blacklist", "%lu addresses allowed to be scanned (%0.0f%% of address space)", 
 			  allowed, allowed*100./((long long int)1 << 32));
+
+	/*
+	// test
+	log_debug("blacklist", "testing started");
+	uint64_t count = constraint_count_ips(constraint, ADDR_ALLOWED);
+	for (unsigned int i=0; i < count; i++) {
+		int ip = constraint_lookup_index(constraint, i, ADDR_ALLOWED);
+		if ((i & 0xFFFFFF) == 0)
+			log_info("blacklist", "%x", i & 0xFF000000);
+		if (constraint_lookup_ip(constraint, ip) != ADDR_ALLOWED) {
+			log_error("blacklist", "test failed for index %d", i);
+		}
+	}
+	log_debug("blacklist", "testing complete");
+	*/
 	return 0;
 }

lib/constraint.c

@@ -50,6 +50,7 @@ typedef struct node {
 	struct node *l;
 	struct node *r;
 	value_t value;
+	uint64_t count;
 } node_t;
 
 // As an optimization, we precompute lookups for every prefix of this
@@ -60,6 +61,8 @@ struct _constraint {
 	node_t *root;   // root node of the tree
 	node_t **radix; // array of nodes for every RADIX_LENGTH prefix
 	int optimized;  // is radix populated and up-to-date?
+	int painted;    // have we precomputed counts for each node?
+	value_t paint_value; // value for which we precomputed counts
 };
 
 // Tree operations respect the invariant that every node that isn't a
@@ -176,7 +179,7 @@ static int _lookup_ip(node_t *root, uint32_t address)
 
 // Return the value pertaining to an address.
 // (Note: address must be in host byte order.)
-int constraint_lookup_ip(constraint_t *con, uint32_t address)
+value_t constraint_lookup_ip(constraint_t *con, uint32_t address)
 {
 	assert(con);
 	if (con->optimized) {
@@ -193,27 +196,91 @@ int constraint_lookup_ip(constraint_t *con, uint32_t address)
 	}
 }
 
+// Return the nth painted IP address.
+static int _lookup_index(node_t *root, uint64_t n)
+{
+	assert(root);
+	node_t *node = root;
+	uint32_t ip = 0;
+	uint32_t mask = 0x80000000;
+	for (;;) {
+		if (IS_LEAF(node)) {
+			return ip | n;
+		}
+		if (n < node->l->count) {
+			node = node->l;
+		} else {
+			n -= node->l->count;
+			node = node->r;
+			ip |= mask;
+		}
+		mask >>= 1;
+	}
+}
+
+// For a given value, return the IP address with zero-based index n.
+// (i.e., if there are three addresses with value 0xFF, looking up index 1
+// will return the second one).
+// Note that the tree must have been previously painted with this value.
+uint32_t constraint_lookup_index(constraint_t *con, uint64_t index, value_t value)
+{
+	assert(con);
+	if (!con->painted || con->paint_value != value) {
+		constraint_paint_value(con, value);
+	}
+	if (con->optimized) {
+		// TK TK TK
+	}
+
+	assert(index < con->root->count);
+	return _lookup_index(con->root, index);
+}
+
+
 // Implement count_ips by recursing on halves of the tree.  Size represents
 // the number of addresses in a prefix at the current level of the tree.
-static uint64_t _count_ips_recurse(node_t *node, value_t value, uint64_t size)
+// If paint is specified, each node will have its count set to the number of
+// leaves under it set to value.
+static uint64_t _count_ips_recurse(node_t *node, value_t value, uint64_t size, int paint)
 {
 	assert(node);
+	uint64_t n;
 	if (IS_LEAF(node)) {
 		if (node->value == value) {
-			return size;
+			n = size;
 		} else {
-			return 0;
+			n = 0;
 		}
+	} else {
+		n = _count_ips_recurse(node->l, value, size >> 1, paint) +
+			_count_ips_recurse(node->r, value, size >> 1, paint);
 	}
-	return _count_ips_recurse(node->l, value, size >> 1) +
-		_count_ips_recurse(node->r, value, size >> 1);
+	if (paint) {
+		node->count = n;
+	}
+	return n;
+}
+
+// For each node, precompute the count of leaves beneath it set to value.
+// Note that the tree can be painted for only one value at a time.
+void constraint_paint_value(constraint_t *con, value_t value)
+{
+	assert(con);
+	log_info("constraint", "Painting value %lu", value);
+	_count_ips_recurse(con->root, value, (uint64_t)1 << 32, 1);
+	con->painted = 1;
+	con->paint_value = value;
 }
 
 // Return the number of addresses that have a given value.
 uint64_t constraint_count_ips(constraint_t *con, value_t value)
 {
 	assert(con);
-	return _count_ips_recurse(con->root, value, (uint64_t)1 << 32);
+	if (con->painted && con->paint_value == value) {
+		return con->root->count;
+	} else {
+		return _count_ips_recurse(con->root, value, (uint64_t)1 << 32, 0);
+	}
 }
 
 // Initialize the tree.

lib/constraint.h

@@ -2,13 +2,15 @@
 #define _CONSTRAINT_H
 
 typedef struct _constraint constraint_t;
-typedef int value_t;
+typedef unsigned int value_t;
 
 constraint_t* constraint_init(value_t value);
 void constraint_free(constraint_t *con);
 void constraint_set(constraint_t *con, uint32_t prefix, int len, value_t value);
 void constraint_optimize(constraint_t *con);
-int constraint_lookup_ip(constraint_t *con, uint32_t address);
+value_t constraint_lookup_ip(constraint_t *con, uint32_t address);
 uint64_t constraint_count_ips(constraint_t *con, value_t value);
+uint32_t constraint_lookup_index(constraint_t *con, uint64_t index, value_t value);
+void constraint_paint_value(constraint_t *con, value_t value);
 
 #endif //_CONSTRAINT_H