Extend the trie_walk_init api + test

The trie_walk_init() function now supports also searching whole trie
subnet and all successor subnets (in lexicographic order). This behavior
can be accomplished by setting @net, and @include_successors to subnet,
and non-zero respectivelly.

filter/data.h

@@ -127,7 +127,7 @@ void *trie_add_prefix(struct f_trie *t, const net_addr *n, uint l, uint h);
 int trie_match_net(const struct f_trie *t, const net_addr *n);
 int trie_match_longest_ip4(const struct f_trie *t, const net_addr_ip4 *net, net_addr_ip4 *dst, ip4_addr *found0);
 int trie_match_longest_ip6(const struct f_trie *t, const net_addr_ip6 *net, net_addr_ip6 *dst, ip6_addr *found0);
-void trie_walk_init(struct f_trie_walk_state *s, const struct f_trie *t, const net_addr *from);
+int trie_walk_init(struct f_trie_walk_state *s, const struct f_trie *t, const net_addr *from, u8 include_successors);
 int trie_walk_next(struct f_trie_walk_state *s, net_addr *net);
 int trie_same(const struct f_trie *t1, const struct f_trie *t2);
 void trie_format(const struct f_trie *t, buffer *buf);
@@ -179,14 +179,17 @@ trie_match_next_longest_ip6(net_addr_ip6 *n, ip6_addr *found)
 
 #define TRIE_WALK_TO_ROOT_END })
 
-
-#define TRIE_WALK(trie, net, from) ({				\
+#define TRIE_WALK2(trie, net, from, include) ({			\
   net_addr net;							\
   struct f_trie_walk_state tws_;				\
-  trie_walk_init(&tws_, trie, from);				\
+  trie_walk_init(&tws_, trie, from, include);			\
   while (trie_walk_next(&tws_, &net))
 
-#define TRIE_WALK_END })
+#define TRIE_WALK2_END })
+
+#define TRIE_WALK(trie, net, from) TRIE_WALK2(trie, net, from, 0) \
+
+#define TRIE_WALK_END TRIE_WALK2_END
 
 
 #define F_CMP_ERROR 999

filter/trie.c

@@ -790,14 +790,22 @@ done:
  * @s: walk state
  * @t: trie
  * @net: optional subnet for walk
+ * @include_successors: optional flag for continue walking beyond subnet @net
  *
  * Initialize walk state for subsequent walk through nodes of the trie @t by
  * trie_walk_next(). The argument @net allows to restrict walk to given subnet,
  * otherwise full walk over all nodes is used. This is done by finding node at
- * or below @net and starting position in it.
+ * or below @net and starting position in it. The argument @include_successors
+ * removes the restriction for all subnets lexicographically succeeding the
+ * @net. In case of @net search fail the walk state starting position points to
+ * the nearest parent node availible. If you use @net and @include_successors,
+ * beware that the trie_walk_next() could return a net preceding the one
+ * specified in @net.
+ *
+ * If desired start position node was found in trie, 1 is returned, 0 otherwise.
  */
-void
-trie_walk_init(struct f_trie_walk_state *s, const struct f_trie *t, const net_addr *net)
+int
+trie_walk_init(struct f_trie_walk_state *s, const struct f_trie *t, const net_addr *net, u8 include_successors)
 {
   *s = (struct f_trie_walk_state) {
     .ipv4 = t->ipv4,
@@ -809,7 +817,7 @@ trie_walk_init(struct f_trie_walk_state *s, const struct f_trie *t, const net_ad
   };
 
   if (!net)
-    return;
+    return 1;
 
   /* We want to find node of level at least plen */
   int plen = ROUND_DOWN_POW2(net->pxlen, TRIE_STEP);
@@ -840,16 +848,42 @@ trie_walk_init(struct f_trie_walk_state *s, const struct f_trie *t, const net_ad
 	s->accept_length = net->pxlen;
       }
 
-      s->stack[0] = n;
-      return;
+      /* Set as root node the only searched subnet */
+      if (!include_successors)
+	s->stack[0] = n;
+      /* Save the last node on the stack otherwise */
+      else
+      {
+	/* Found prefect match, no advancing */
+	s->stack[s->stack_pos] = n;
+	/* Search whole trie except skipped parts */
+	s->start_pos = 1;
+      }
+
+      return 1;
     }
 
+    /* We store node in stack before moving on */
+    if (include_successors)
+      s->stack[s->stack_pos++] = n;
+
     /* Choose child */
     n = GET_CHILD(n, v4, GET_NET_BITS(net, v4, nlen, TRIE_STEP));
   }
 
-  s->stack[0] = NULL;
-  return;
+  /* We do not override the trie root in case of inclusive search */
+  if (!include_successors)
+    s->stack[0] = NULL;
+
+  /* Be careful about underflow */
+  else if (s->stack_pos > 0)
+    s->stack_pos--;
+
+  /* Search whole trie except skipped parts */
+  if (include_successors)
+    s->start_pos = 1;
+
+  return 0;
 }
 
 #define GET_ACCEPT_BIT(N,X,B) ((X) ? ip4_getbit((N)->v4.accept, (B)) : ip6_getbit((N)->v6.accept, (B)))

filter/trie_test.c

@@ -880,6 +880,137 @@ t_trie_walk_to_root(void)
   return 1;
 }
 
+static int
+t_trie_walk_inclusive(void)
+{
+  bt_bird_init();
+  bt_config_parse(BT_CONFIG_SIMPLE);
+
+  for (int round = 0; round < TESTS_NUM*8; round++)
+  {
+    int level = round / TESTS_NUM;
+    int v6 = level % 2;
+    int num = PREFIXES_NUM * (int[]){1, 10, 100, 1000}[level / 2];
+    int pos = 0, end = 0;
+    list *prefixes = make_random_prefix_list(num, v6, 1);
+    struct f_trie *trie = make_trie_from_prefix_list(prefixes);
+    struct f_prefix *pxset = malloc((num + 1) * sizeof(struct f_prefix));
+
+    struct f_prefix_node *n;
+    WALK_LIST(n, *prefixes)
+      pxset[pos++] = n->prefix;
+    memset(&pxset[pos], 0, sizeof (struct f_prefix));
+
+    qsort(pxset, num, sizeof(struct f_prefix), compare_prefixes);
+
+    /* // print sorted prefixes
+    bt_debug("sorted prefixes\n");
+    for (struct f_prefix *px = pxset; px < pxset + num; px++)
+    {
+      char buf[64];
+      bt_format_net(buf, 64, &px->net);
+      bt_debug("%s{%d,%d}\n", buf, px->lo, px->hi);
+    }
+    */
+ 
+    /* Full walk */
+    bt_debug("Full walk inclusive (round %d, %d nets)\n", round, num);
+
+    pos = 0;
+    uint pxc = 0;
+    TRIE_WALK2(trie, net, NULL, 1)
+    {
+      log_networks(&net, &pxset[pos].net);
+      bt_assert(net_equal(&net, &pxset[pos].net));
+
+      /* Skip possible duplicates */
+      while (net_equal(&pxset[pos].net, &pxset[pos + 1].net))
+	pos++;
+
+      pos++;
+      pxc++;
+    }
+    TRIE_WALK2_END;
+
+    bt_assert(pos == num);
+    bt_assert(pxc == trie->prefix_count);
+    bt_debug("Full walk inclusive done\n");
+
+
+    /* Prepare net for subnet walk - start with random prefix */
+    pos = bt_random() % num;
+    end = pos + (int[]){2, 2, 3, 4}[level / 2];
+    end = MIN(end, num);
+
+    struct f_prefix from = pxset[pos];
+
+    /* Find a common superprefix to several subsequent prefixes */
+    for (; pos < end; pos++)
+    {
+      if (net_equal(&from.net, &pxset[pos].net))
+	continue;
+
+      int common = !v6 ?
+	ip4_pxlen(net4_prefix(&from.net), net4_prefix(&pxset[pos].net)) :
+	ip6_pxlen(net6_prefix(&from.net), net6_prefix(&pxset[pos].net));
+      from.net.pxlen = MIN(from.net.pxlen, common);
+
+      if (!v6)
+	((net_addr_ip4 *) &from.net)->prefix =
+	  ip4_and(net4_prefix(&from.net), net4_prefix(&pxset[pos].net));
+      else
+	((net_addr_ip6 *) &from.net)->prefix =
+	  ip6_and(net6_prefix(&from.net), net6_prefix(&pxset[pos].net));
+    }
+
+    /* Fix irrelevant bits */
+    if (!v6)
+      ((net_addr_ip4 *) &from.net)->prefix =
+	ip4_and(net4_prefix(&from.net), ip4_mkmask(net4_pxlen(&from.net)));
+    else
+      ((net_addr_ip6 *) &from.net)->prefix =
+	ip6_and(net6_prefix(&from.net), ip6_mkmask(net6_pxlen(&from.net)));
+
+
+    /* Find initial position for final prefix */
+    for (pos = 0; pos < num; pos++)
+      if (compare_prefixes(&pxset[pos], &from) >= 0)
+	break;
+
+      /* Account for subnets before searched net from */
+    for (; pos < num; pos++)
+      if (net_compare(&pxset[pos].net, &from.net) >= 0)
+	 break;
+
+    int p0 = pos;
+    char buf0[64];
+    bt_format_net(buf0, 64, &from.net);
+    bt_debug("Subnet walk inclusive for %s (round %d, %d nets)\n", buf0, round, num);
+
+    /* Subnet walk */
+    TRIE_WALK2(trie, net, &from.net, 1)
+    {
+      log_networks(&net, &pxset[pos].net);
+      bt_assert(net_compare(&net, &pxset[pos].net) >= 0);
+
+      /* Skip possible duplicates */
+      while (net_equal(&pxset[pos].net, &pxset[pos + 1].net))
+	pos++;
+
+      pos++;
+    }
+    TRIE_WALK2_END;
+
+    bt_assert(pos == num);
+    bt_debug("Subnet walk done inclusive for %s (found %d nets)\n", buf0, pos - p0);
+
+    tmp_flush();
+  }
+
+  bt_bird_cleanup();
+  return 1;
+}
+
 int
 main(int argc, char *argv[])
 {
@@ -891,6 +1022,7 @@ main(int argc, char *argv[])
   bt_test_suite(t_trie_same, "A trie filled forward should be same with a trie filled backward.");
   bt_test_suite(t_trie_walk, "Testing TRIE_WALK() on random tries");
   bt_test_suite(t_trie_walk_to_root, "Testing TRIE_WALK_TO_ROOT() on random tries");
+  bt_test_suite(t_trie_walk_inclusive, "Testing TRIE_WALK2() on random tries");
 
   // bt_test_suite(t_bench_trie_datasets_subset, "Benchmark tries from datasets by random subset of nets");
   // bt_test_suite(t_bench_trie_datasets_random, "Benchmark tries from datasets by generated addresses");

nest/rt-table.c

@@ -2057,7 +2057,7 @@ rt_table_export_start(struct rt_exporter *re, struct rt_export_request *req)
       {
 	hook->walk_state = mb_allocz(p, sizeof (struct f_trie_walk_state));
 	hook->walk_lock = rt_lock_trie(tab);
-	trie_walk_init(hook->walk_state, tab->trie, req->addr);
+	trie_walk_init(hook->walk_state, tab->trie, req->addr, 0);
 	hook->event = ev_new_init(p, rt_feed_by_trie, hook);
 	break;
       }