implement search_coverage in c, fork from @Habbie, which had rather o…

…dd changelog and conflicting rebase, so decided to make a clean patch

README.rst

@@ -85,6 +85,11 @@ A simple example that demonstrates most of the features: ::
  # that contains the search term (inverse routing-style lookup)
  rnode = rtree.search_worst("10.123.45.6")
 
+ # Covered search will return all prefixes inside the given
+ # search term, as a list (including the search term itself,
+ # if present in the tree)
+ rnodes = rtree.search_covered("10.123.0.0/16")
+
  # There are a couple of implicit members of a RadixNode:
  print rnode.network # -> "10.0.0.0"
  print rnode.prefix # -> "10.0.0.0/8"

radix/__init__.py

@@ -17,6 +17,7 @@ def __init__(self):
  self.search_exact = self._radix.search_exact
  self.search_best = self._radix.search_best
  self.search_worst = self._radix.search_worst
+ self.search_covered = self._radix.search_covered
  self.nodes = self._radix.nodes
  self.prefixes = self._radix.prefixes
 

radix/_radix.c

@@ -508,6 +508,50 @@ Radix_search_worst(RadixObject *self, PyObject *args, PyObject *kw_args)
  return (PyObject *)node_obj;
 }
 
+PyDoc_STRVAR(Radix_search_covered_doc,
+"Radix.search_covered(network[, masklen][, packed] -> None\n\
+\n\
+FIXME");
+
+static PyObject *
+Radix_search_covered(RadixObject *self, PyObject *args, PyObject *kw_args)
+{
+ radix_node_t *node;
+ radix_node_t *node_iter;
+ prefix_t *prefix;
+ PyObject *ret;
+
+ static char *keywords[] = { "network", "masklen", "packed", NULL };
+
+ char *addr = NULL, *packed = NULL;
+ long prefixlen = -1;
+ int packlen = -1;
+
+ if (!PyArg_ParseTupleAndKeywords(args, kw_args, "|sls#:search_covered", keywords, &addr, &prefixlen, &packed, &packlen))
+ return NULL;
+
+ if ((prefix = args_to_prefix(addr, packed, packlen, prefixlen)) == NULL)
+ return NULL;
+
+ if ((ret = PyList_New(0)) == NULL)
+ return NULL;
+
+ if ((node = radix_search_node(PICKRT(prefix, self), prefix)) == NULL) {
+ Deref_Prefix(prefix);
+ Py_INCREF(Py_None);
+ return ret;
+ }
+
+ RADIX_WALK(node, node_iter) {
+ if (node_iter->data != NULL) {
+ PyList_Append(ret, ((RadixNodeObject *)node_iter->data));
+ }
+ } RADIX_WALK_END;
+
+ Deref_Prefix(prefix);
+ return (ret);
+}
+
 
 PyDoc_STRVAR(Radix_nodes_doc,
 "Radix.nodes(prefix) -> List of RadixNode\n\
@@ -589,6 +633,7 @@ static PyMethodDef Radix_methods[] = {
  {"search_exact",(PyCFunction)Radix_search_exact,METH_VARARGS|METH_KEYWORDS, Radix_search_exact_doc },
  {"search_best", (PyCFunction)Radix_search_best, METH_VARARGS|METH_KEYWORDS, Radix_search_best_doc },
  {"search_worst", (PyCFunction)Radix_search_worst, METH_VARARGS|METH_KEYWORDS, Radix_search_worst_doc },
+ {"search_covered", (PyCFunction)Radix_search_covered, METH_VARARGS|METH_KEYWORDS, Radix_search_covered_doc },
  {"nodes", (PyCFunction)Radix_nodes, METH_VARARGS, Radix_nodes_doc },
  {"prefixes", (PyCFunction)Radix_prefixes, METH_VARARGS, Radix_prefixes_doc },
  {NULL, NULL} /* sentinel */

radix/_radix/radix.c

@@ -275,6 +275,38 @@ radix_node_t
 }
 
 
+/* seach for a node without checking if it is a "real" node */
+radix_node_t
+*radix_search_node(radix_tree_t *radix, prefix_t *prefix)
+{
+ radix_node_t *node;
+ u_char *addr;
+ u_int bitlen;
+
+ if (radix->head == NULL)
+ return (NULL);
+
+ node = radix->head;
+ addr = prefix_touchar(prefix);
+ bitlen = prefix->bitlen;
+
+ while (node->bit < bitlen) {
+ if (BIT_TEST(addr[node->bit >> 3], 0x80 >> (node->bit & 0x07)))
+ node = node->r;
+ else
+ node = node->l;
+
+ if (node == NULL)
+ return (NULL);
+ }
+
+ if (comp_with_mask(prefix_tochar(node->prefix), prefix_tochar(prefix), bitlen))
+ return (node);
+
+ return (NULL);
+}
+
+
 /* if inclusive != 0, "best" may be the given prefix itself */
 static radix_node_t
 *radix_search_best2(radix_tree_t *radix, prefix_t *prefix, int inclusive)

radix/_radix/radix.h

@@ -121,6 +121,7 @@ radix_tree_t *New_Radix(void);
 void Destroy_Radix(radix_tree_t *radix, rdx_cb_t func, void *cbctx);
 radix_node_t *radix_lookup(radix_tree_t *radix, prefix_t *prefix);
 void radix_remove(radix_tree_t *radix, radix_node_t *node);
+radix_node_t *radix_search_node(radix_tree_t *radix, prefix_t *prefix);
 radix_node_t *radix_search_exact(radix_tree_t *radix, prefix_t *prefix);
 radix_node_t *radix_search_best(radix_tree_t *radix, prefix_t *prefix);
 radix_node_t *radix_search_worst(radix_tree_t *radix, prefix_t *prefix);

radix/radix.py

@@ -324,6 +324,34 @@ def search_worst(self, prefix):
  return node
  return None
 
+ def search_covered(self, prefix):
+ results = []
+ if self.head is None:
+ return results
+ node = self.head
+ addr = prefix.addr
+ bitlen = prefix.bitlen
+
+ while node.bitlen < bitlen:
+ if self._addr_test(addr, node.bitlen):
+ node = node.right
+ else:
+ node = node.left
+ if node is None:
+ return results
+
+ stack = [node]
+ while stack:
+ node = stack.pop()
+ if self._prefix_match(node._prefix, prefix, prefix.bitlen):
+ results.append(node)
+ if node.right:
+ stack.append(node.right)
+ if node.left:
+ stack.append(node.left)
+
+ return results
+
  def _prefix_match(self, left, right, bitlen):
  l = left.addr
  r = right.addr
@@ -446,6 +474,13 @@ def search_worst(self, network=None, masklen=None, packed=None):
  else:
  return None
 
+ def search_covered(self, network=None, masklen=None, packed=None):
+ prefix = RadixPrefix(network, masklen, packed)
+ if prefix.family == AF_INET:
+ return self._tree4.search_covered(prefix)
+ else:
+ return self._tree6.search_covered(prefix)
+
  def _iter(self, node):
  stack = []
  while node is not None:

tests/test_regression.py

@@ -405,6 +405,33 @@ def test_25_search_default(self):
  '0.0.0.0/0')
 
 
+ def test_26_search_covered(self):
+ tree = radix.Radix()
+ tree.add('10.0.0.0/8')
+ tree.add('10.0.0.0/13')
+ tree.add('10.0.0.0/31')
+ tree.add('11.0.0.0/16')
+ self.assertEquals(
+ [n.prefix for n in tree.search_covered('11.0.0.0/8')],
+ ['11.0.0.0/16'])
+ self.assertEquals(
+ [n.prefix for n in tree.search_covered('10.0.0.0/9')],
+ ['10.0.0.0/13', '10.0.0.0/31'])
+ self.assertEquals(
+ [n.prefix for n in tree.search_covered('10.0.0.0/8')],
+ ['10.0.0.0/8', '10.0.0.0/13', '10.0.0.0/31'])
+ self.assertEquals(
+ [n.prefix for n in tree.search_covered('11.0.0.0/8')],
+ ['11.0.0.0/16'])
+ self.assertEquals(
+ [n.prefix for n in tree.search_covered('21.0.0.0/8')],
+ [])
+ self.assertEquals(
+ [n.prefix for n in tree.search_covered('0.0.0.0/0')],
+ ['10.0.0.0/8', '10.0.0.0/13', '10.0.0.0/31', '11.0.0.0/16'])
+
+
+
 def main():
  unittest.main()