update Patricia.new for optional IPv6 support

IPv4 is still the default and will likely remain so
indefinitely.

4 files changed, 103 insertions, 24 deletions

diff --git a/README b/README
index 3b39e23..6c8bbaf 100644
--- a/README
+++ b/README
@@ -60,9 +60,11 @@ new:
 
       pt = Patricia.new
 
-    This is the class' constructor - it returns a Patricia object.  For
-    now, the constructor takes no arguments, and defaults to creating a
-    tree which uses AF_INET IPv4 address and mask values as keys.
+    This is the class' constructor - it returns a Patricia object for
+    handling IPv4 addresses.  To handle IPv6 addresses (only),
+    Patricia.new may be called with an additional argument:
+
+      pt = Patricia.new(:AF_INET6)
 
     The Patricia object will be destroyed automatically when there are
     no longer any references to it.
@@ -74,7 +76,8 @@ add:
     specification in canonical form, e.g. ``10.0.0.0/8'', where the
     number after the slash represents the number of bits in the
     netmask. If no mask width is specified, the longest possible mask
-    is assumed, i.e. 32 bits for AF_INET addresses.
+    is assumed, i.e. 32 bits for AF_INET addresses and 128 bits for
+    AF_INET6 addresses.
 
     The second argument, user_data, is optional. If supplied, it
     should be a STRING object specifying the user data that will be 
@@ -89,6 +92,12 @@ add:
     On success, this method returns the object of the Patricia Trie 
     node.
 
+family:
+
+    Returns either :AF_INET or :AF_INET6 symbol depending on how the
+    object was initialized.  A Patricia object may only handle IPv4 or
+    IPv6 addresses.
+
 add_node:    An alias of add.
 
 search_best:
@@ -103,7 +112,7 @@ search_best:
     canonical form, e.g. ``10.0.0.0/8'', where the number after the
     slash represents the number of bits in the netmask. If no mask
     width value is specified, the longest mask is assumed, i.e. 32
-    bits for AF_INET addresses.
+    bits for AF_INET addresses and 128 bits for AF_INET6 addresses.
 
     If a matching node is found in the Patricia Trie, this method
     returns the object of the node. This method returns nil on
diff --git a/ext/rpatricia/rpatricia.c b/ext/rpatricia/rpatricia.c
index 17534c6..f8ce0b5 100644
--- a/ext/rpatricia/rpatricia.c
+++ b/ext/rpatricia/rpatricia.c
@@ -9,6 +9,7 @@
 #include <assert.h>
 
 static VALUE cPatricia, cNode;
+static VALUE sym_AF_INET, sym_AF_INET6;
 
 static void dummy(void) {}
 
@@ -266,25 +267,42 @@ p_tree_mark (void *ptr)
 {
   patricia_tree_t *tree = ptr;
 
-  patricia_process(tree, p_tree_mark_each);
+  if (tree)
+    patricia_process(tree, p_tree_mark_each);
 }
 
 static void
 p_tree_free (void *ptr)
 {
   patricia_tree_t *tree = ptr;
-
   /* no need to explicitly free each node->data, GC will do it for us */
-  Destroy_Patricia(tree, NULL);
+  if (tree)
+    Destroy_Patricia(tree, NULL);
 }
 
 static VALUE
 p_alloc(VALUE klass)
 {
-  patricia_tree_t *tree;
-  tree = New_Patricia(32); /* assuming only IPv4 */
+  return Data_Wrap_Struct(klass, p_tree_mark, p_tree_free, NULL);
+}
+
+static VALUE
+p_init(int argc, VALUE *argv, VALUE self)
+{
+  VALUE family;
+  int maxbits;
+
+  rb_scan_args(argc, argv, "01", &family);
+
+  if (NIL_P(family) || family == sym_AF_INET)
+    maxbits = 32;
+  else if (family == sym_AF_INET6)
+    maxbits = 128;
+  else
+    rb_raise(rb_eArgError, "unknown family (must be :AF_INET6 or :AF_INET)");
 
-  return Data_Wrap_Struct(klass, p_tree_mark, p_tree_free, tree);
+  DATA_PTR(self) = New_Patricia(maxbits);
+  return self;
 }
 
 static VALUE
@@ -299,7 +317,8 @@ p_init_copy(VALUE self, VALUE orig)
     prefix_t prefix;
     VALUE user_data;
 
-    Data_Get_Struct(self, patricia_tree_t, tree);
+    DATA_PTR(self) = tree = New_Patricia(orig_tree->maxbits);
+
     PATRICIA_WALK(orig_tree->head, orig_node) {
       node = patricia_lookup(tree, orig_node->prefix);
       assert(node->prefix == orig_node->prefix);
@@ -312,20 +331,39 @@ p_init_copy(VALUE self, VALUE orig)
   }
 }
 
+static VALUE
+p_family(VALUE self)
+{
+  patricia_tree_t *tree;
+
+  Data_Get_Struct(self, patricia_tree_t, tree);
+
+  switch (tree->maxbits) {
+  case 32: return sym_AF_INET;
+  case 128: return sym_AF_INET6;
+  }
+  assert(0 && "unknown maxbits, corrupt tree");
+  return Qnil;
+}
+
 void
 Init_rpatricia (void)
 {
   cPatricia = rb_define_class("Patricia", rb_cObject);
   cNode = rb_define_class_under(cPatricia, "Node", rb_cObject);
+  sym_AF_INET = ID2SYM(rb_intern("AF_INET"));
+  sym_AF_INET6 = ID2SYM(rb_intern("AF_INET6"));
 
   /* allocate new Patricia object, called before initialize  */
   rb_define_alloc_func(cPatricia, p_alloc);
+  rb_define_private_method(cPatricia, "initialize", p_init, -1);
   rb_define_method(cPatricia, "initialize_copy", p_init_copy, 1);
 
   /*---------- methods to tree ----------*/
   /* add string */
   rb_define_method(cPatricia, "add", p_add, -1);
   rb_define_method(cPatricia, "add_node", p_add, -1);
+  rb_define_method(cPatricia, "family", p_family, 0);
 
   /* match prefix */
   rb_define_method(cPatricia, "match_best", p_match, 1); 
@@ -356,6 +394,4 @@ Init_rpatricia (void)
   rb_define_method(cNode, "network", p_network, 0);
   rb_define_method(cNode, "prefix", p_prefix, 0);
   rb_define_method(cNode, "prefixlen", p_prefixlen, 0);
-  //  rb_define_method(cPatricia, "family", p_family, 0); 
-
 }
diff --git a/test/test_duplicate.rb b/test/test_duplicate.rb
index c50a211..c6ddc1c 100644
--- a/test/test_duplicate.rb
+++ b/test/test_duplicate.rb
@@ -8,21 +8,42 @@ class TestDuplicate < Test::Unit::TestCase
     t = Patricia.new
     t.add('127.0.0.0/8', tmp)
     t2 = t.dup
+    assert_equal :AF_INET, t2.family
     assert_equal 1, t2.num_nodes
     assert_equal tmp.object_id, t2.match_best('127.0.0.1').data.object_id
     t2.add('10.0.0.0/8', zz = [])
     assert_equal 2, t2.num_nodes
     assert_equal 1, t.num_nodes
 
-    oldout = $stdout
-    begin
-      $stdout = stringio = StringIO.new
-      t2.show_nodes
-      puts "--"
-      t.show_nodes
-    ensure
-      $stdout = oldout
-    end
-    p stringio.string
+    tio = StringIO.new
+    t.show_nodes(tio)
+    assert_equal "node: 127.0.0.0/8\n", tio.string
+
+    t2io = StringIO.new
+    t2.show_nodes(t2io)
+    assert_equal("node: 10.0.0.0/8\nnode: 127.0.0.0/8\n", t2io.string)
+  end
+
+  def test_dup_ipv6
+    tmp = {}
+    t = Patricia.new :AF_INET6
+    t.add('1234:4321::/32', tmp)
+    t.add('2600:0102:a100::/43', tmp)
+    t2 = t.dup
+    assert_equal :AF_INET6, t2.family
+    assert_equal 2, t2.num_nodes
+    t2.add('::1/128', zz = [])
+    assert_equal 3, t2.num_nodes
+    assert_equal 2, t.num_nodes
+
+    tio = StringIO.new
+    t.show_nodes(tio)
+    expect = "node: 1234:4321::/32\nnode: 2600:102:a100::/43\n"
+    assert_equal expect, tio.string
+
+    t2io = StringIO.new
+    t2.show_nodes(t2io)
+    expect = "node: ::1/128\nnode: 1234:4321::/32\nnode: 2600:102:a100::/43\n"
+    assert_equal expect, t2io.string
   end
 end
diff --git a/test/test_ipv6.rb b/test/test_ipv6.rb
new file mode 100644
index 0000000..7bf95c2
--- /dev/null
+++ b/test/test_ipv6.rb
@@ -0,0 +1,13 @@
+require 'test/unit'
+require 'rpatricia'
+
+class TestIPv6 < Test::Unit::TestCase
+  def test_ipv6
+    t = Patricia.new :AF_INET6
+    assert_equal :AF_INET6, t.family
+    node = t.add('1234:567::/35', 'hello world')
+    assert_kind_of Patricia::Node, node
+    assert_equal 'hello world', node.data
+    t.match_best('1234:567::/128').data
+  end
+end