kdtree 0.1 → 0.3

data/.gitignore

@@ -0,0 +1,4 @@
+*.gem
+Gemfile.lock
+lib/*.so
+tmp

data/.travis.yml

@@ -0,0 +1,9 @@
+language: ruby
+rvm:
+  - 1.8.7
+  - 1.9.2
+  - 1.9.3
+  - rbx-18mode
+  - rbx-19mode
+#  - ruby-head
+  - ree

data/Gemfile

@@ -0,0 +1,2 @@
+source "http://rubygems.org"
+gemspec

data/LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2009 Adam Doppelt
+Copyright (c) 2012 Adam Doppelt
 
 Permission is hereby granted, free of charge, to any person obtaining
 a copy of this software and associated documentation files (the

data/README.md

@@ -0,0 +1,70 @@
+## Kdtree
+
+A kd tree is a data structure that recursively partitions the world in order to rapidly answer nearest neighbor queries. A generic kd tree can support any number of dimensions, and can return either the nearest neighbor or a set of N nearest neighbors.
+
+This gem is a blazingly fast, native, 2d kdtree. It's specifically built to find the nearest neighbor when searching millions of points. It's used in production at Urbanspoon and several other companies.
+
+The first version of this gem was released back in 2009. See the original [blog post](http://gurge.com/2009/10/22/ruby-nearest-neighbor-fast-kdtree-gem/) for the full story. Wikipedia has a great [article on kdtrees](http://en.wikipedia.org/wiki/K-d_tree).
+
+### Usage
+
+Usage is very simple:
+
+* **Kdtree.new(points)** - construct a new tree. Each point should be of the form `[x, y, id]`, where `x/y` are floats and `id` is an int. Not a string, not an object, just an int.
+* **kd.nearest(x, y)** - find the nearest point. Returns an id.
+* **kd.nearestk(x, y, k)** - find the nearest `k` points. Returns an array of ids.
+
+Also, I made it possible to **persist** the tree to disk and load it later. That way you can calculate the tree offline and load it quickly at some future point. Loading a persisted tree w/ 1 millions points takes less than a second, as opposed to the 3.5 second startup time shown above. For example:
+
+```ruby
+File.open("treefile", "w") { |f| kd.persist(f) }
+... later ...
+kd2 = File.open("treefile") { |f| Kdtree.new(f) }
+```
+
+### Performance
+
+Kdtree is fast. How fast? Using a tree with 1 million points on my i5 2.8ghz:
+
+```
+build               3.5s
+nearest point       0.000003s
+nearest 5 points    0.000004s
+nearest 50 points   0.000014s
+nearest 255 points  0.000063s
+```
+
+### Limitations
+
+* No **editing** allowed! Once you construct a tree you’re stuck with it.
+* The tree is stored in **one big memory block**, 20 bytes per point. A tree with one million points will allocate a single 19mb block to store its nodes.
+* Persisted trees are **architecture dependent**, and may not work across different machines due to endian issues.
+* nearestk is limited to **255 results**
+
+### Contributors
+
+Since this gem was originally released, several folks have contributed important patches:
+
+* @antifuchs (thread safety)
+* @evanphx (native cleanups, perf)
+* @ghazel (C89 compliance)
+* @mcerna (1.9 compat)
+
+### Changelog
+
+#### 0.3 (in progress, unreleased)
+
+* Ruby 1.9.x compatibility (@mcerna and others)
+* renamed KDTree to the more idiomatic Kdtree
+* use IO methods directly instead of rooting around in rb_io
+* thread safe, no more statics (@antifuchs)
+* C90 compliance, no warnings (@ghazel)
+* native cleanups (@evanphx)
+
+#### 0.2
+
+skipped this version to prevent confusion with other flavors of the gem
+
+#### 0.1
+
+* Original release

data/Rakefile

@@ -0,0 +1,41 @@
+require "bundler/setup"
+require "rake/extensiontask"
+require "rake/testtask"
+
+# load the spec, we use it below
+spec = Gem::Specification.load("kdtree.gemspec")
+
+#
+# gem
+#
+
+task :build do
+  system "gem build --quiet kdtree.gemspec"
+end
+
+task :install => :build do
+  system "sudo gem install --quiet kdtree-#{spec.version}.gem"
+end
+
+task :release => :build do
+  system "git tag -a #{spec.version} -m 'Tagging #{spec.version}'"
+  system "git push --tags"
+  system "gem push kdtree-#{spec.version}.gem"
+end
+
+#
+# rake-compiler
+#
+
+Rake::ExtensionTask.new("kdtree", spec)
+
+
+#
+# testing
+#
+
+Rake::TestTask.new(:test) do |test|
+  test.libs << "test"
+end
+task :test => :compile
+task :default => :test

data/ext/{extconf.rb → kdtree/extconf.rb}

@@ -1,3 +1,3 @@
-require 'mkmf'
+require "mkmf"
 
 create_makefile("kdtree")

data/ext/{kdtree.c → kdtree/kdtree.c}

@@ -1,15 +1,10 @@
 #include "ruby.h"
-#include "rubyio.h"
-#include "version.h"
-
-#ifndef HAVE_RB_IO_T
-#define rb_io_t OpenFile
-#endif
 
 //
 // interface
 //
 
+// the tree itself
 typedef struct kdtree_data
 {
     int root;
@@ -17,6 +12,7 @@ typedef struct kdtree_data
     struct kdtree_node *nodes;
 } kdtree_data;
 
+// a node in the tree
 typedef struct kdtree_node
 {
     float x, y;
@@ -25,10 +21,18 @@ typedef struct kdtree_node
     int right;
 } kdtree_node;
 
+// a result node from kdtree_nearestk0
+typedef struct kresult {
+    int index;
+    float distance;
+} kresult;
+
+// helper macro for digging out our struct
 #define KDTREEP \
     struct kdtree_data *kdtreep; \
     Data_Get_Struct(kdtree, struct kdtree_data, kdtreep);
 
+// kdtree public methods
 static VALUE kdtree_alloc(VALUE klass);
 static void kdtree_free(struct kdtree_data *kdtreep);
 static VALUE kdtree_initialize(VALUE kdtree, VALUE points);
@@ -37,13 +41,20 @@ static VALUE kdtree_nearestk(VALUE kdtree, VALUE x, VALUE y, VALUE k);
 static VALUE kdtree_persist(VALUE kdtree, VALUE io);
 static VALUE kdtree_to_s(VALUE kdtree);
 
-// helpers
+// kdtree helpers
 static int kdtree_build(struct kdtree_data *kdtreep, int min, int max, int depth);
-static void kdtree_nearest0(struct kdtree_data *kdtreep, int i, float x, float y, int depth);
-static void kdtree_nearestk0(struct kdtree_data *kdtreep, int i, float x, float y, int k, int depth);
+static void kdtree_nearest0(struct kdtree_data *kdtreep, int i, float x, float y, int depth, int *n_index, float *n_dist);
+static void kdtree_nearestk0(struct kdtree_data *kdtreep, int i, float x, float y, int k, int depth, kresult *k_list, int *k_len, float *k_dist);
+
+// io helpers
+static void read_all(VALUE io, void *buf, int len);
+static void write_all(VALUE io, const void *buf, int len);
 
 #define KDTREE_MAGIC "KdTr"
 
+// ids
+static ID id_read, id_write, id_binmode;
+
 //
 // implementation
 //
@@ -63,24 +74,12 @@ static void kdtree_free(struct kdtree_data *kdtreep)
     }
 }
 
-static void read_all(struct rb_io_t *fptr, char *buf, int len)
-{
-    while (len > 0) {
-        int n = rb_io_fread(buf, len, fptr->f);
-        if (n == 0) {
-            rb_eof_error();
-        }
-        buf += n;
-        len -= n;
-    }
-}
-
 /*
  * call-seq:
- *   KDTree.new(points)    => kdtree
- *   KDTree.new(io)        => kdtree
+ *   Kdtree.new(points)    => kdtree
+ *   Kdtree.new(io)        => kdtree
  *
- * Returns a new <code>KDTree</code>. To construct a tree, pass an array of
+ * Returns a new <code>Kdtree</code>. To construct a tree, pass an array of
  * <i>points</i>. Each point should be an array of the form <code>[x, y,
  * id]</code>, where <i>x</i> and <i>y</i> are floats and <i>id</i> is an
  * integer. The <i>id</i> is arbitrary and will be returned to you whenever you
@@ -90,7 +89,7 @@ static void read_all(struct rb_io_t *fptr, char *buf, int len)
  *   points = []
  *   points << [47.6, -122.3, 1] # Seattle
  *   points << [40.7, -74.0, 2]  # New York
- *   kd = KDTree.new(points)
+ *   kd = Kdtree.new(points)
  *
  * Alternately, you can pass in an <i>IO</i> object containing a persisted
  * kdtree. This makes it possible to build the tree in advance, persist it, and
@@ -103,104 +102,97 @@ static VALUE kdtree_initialize(VALUE kdtree, VALUE arg)
     if (TYPE(arg) == T_ARRAY) {
         // init from array of pints
         VALUE points = arg;
+        int i;
         kdtreep->len = RARRAY_LEN(points);
         kdtreep->nodes = ALLOC_N(struct kdtree_node, kdtreep->len);
 
-        int i;
         for (i = 0; i < RARRAY_LEN(points); ++i) {
             struct kdtree_node *n = kdtreep->nodes + i;
-            
-            VALUE ptr = RARRAY_PTR(points)[i];
+
+            VALUE ptr = rb_ary_entry(points, i);
             VALUE v = rb_check_array_type(ptr);
             if (NIL_P(v) || RARRAY_LEN(v) != 3) {
                 continue;
             }
-            VALUE *a = RARRAY_PTR(ptr);
-            n->x = NUM2DBL(a[0]);
-            n->y = NUM2DBL(a[1]);        
-            n->id = NUM2INT(a[2]);
+            n->x = NUM2DBL(rb_ary_entry(v, 0));
+            n->y = NUM2DBL(rb_ary_entry(v, 1));
+            n->id = NUM2INT(rb_ary_entry(v, 2));
         }
 
         // now build the tree
         kdtreep->root = kdtree_build(kdtreep, 0, kdtreep->len, 0);
     } else if (rb_respond_to(arg, rb_intern("read"))) {
         VALUE io = arg;
-        if (rb_respond_to(io, rb_intern("binmode"))) {
-            rb_funcall2(io, rb_intern("binmode"), 0, 0);
+        char buf[4];
+        if (rb_respond_to(io, id_binmode)) {
+            rb_funcall(io, id_binmode, 0);
         }
 
-        struct rb_io_t *fptr = RFILE(rb_io_taint_check(io))->fptr;
-        rb_io_check_readable(fptr);
-
         // check magic
-        char buf[4];
-        read_all(fptr, buf, 4);
+        read_all(io, buf, 4);
         if (memcmp(KDTREE_MAGIC, buf, 4) != 0) {
             rb_raise(rb_eRuntimeError, "wrong magic number in kdtree file");
         }
-        
+
         // read start of the struct
-        read_all(fptr, (char *)kdtreep, sizeof(struct kdtree_data) - sizeof(struct kdtree_node *));
+        read_all(io, kdtreep, sizeof(struct kdtree_data) - sizeof(struct kdtree_node *));
+
         // read the nodes
         kdtreep->nodes = ALLOC_N(struct kdtree_node, kdtreep->len);
-        read_all(fptr, (char *)kdtreep->nodes, sizeof(struct kdtree_node) * kdtreep->len);
+        read_all(io, kdtreep->nodes, sizeof(struct kdtree_node) * kdtreep->len);
     } else {
-        rb_raise(rb_eTypeError, "array or IO required to init KDTree");
+        rb_raise(rb_eTypeError, "array or IO required to init Kdtree");
     }
-    
+
     return kdtree;
 }
 
 static int comparex(const void *pa, const void *pb)
 {
     float a = ((const struct kdtree_node*)pa)->x;
-    float b = ((const struct kdtree_node*)pb)->x;    
+    float b = ((const struct kdtree_node*)pb)->x;
     return (a < b) ? -1 : ((a > b) ? 1 : 0);
 }
 
 static int comparey(const void *pa, const void *pb)
 {
     float a = ((const struct kdtree_node*)pa)->y;
-    float b = ((const struct kdtree_node*)pb)->y;    
+    float b = ((const struct kdtree_node*)pb)->y;
     return (a < b) ? -1 : ((a > b) ? 1 : 0);
 }
 
 static int kdtree_build(struct kdtree_data *kdtreep, int min, int max, int depth)
 {
+    int(*compar)(const void *, const void *);
+    struct kdtree_node *m;
+    int median;
     if (max <= min) {
         return -1;
     }
 
     // sort nodes from min to max
-    int(*compar)(const void *, const void *) = (depth % 2) ? comparex : comparey;
+    compar = (depth % 2) ? comparex : comparey;
     qsort(kdtreep->nodes + min, max - min, sizeof(struct kdtree_node), compar);
 
-    int median = (min + max) / 2;
-    struct kdtree_node *m = kdtreep->nodes + median;
+    median = (min + max) / 2;
+    m = kdtreep->nodes + median;
     m->left = kdtree_build(kdtreep, min, median, depth + 1);
     m->right = kdtree_build(kdtreep, median + 1, max, depth + 1);
     return median;
 }
 
-//
-// nearest
-//
-
-static int n_index;
-static float n_dist;
-
 /*
  * call-seq:
  *   kd.nearest(x, y)    => id
  *
  * Finds the point closest to <i>x</i>, <i>y</i> and returns the id for that
  * point. Returns -1 if the tree is empty.
- * 
+ *
  *   points = []
  *   points << [47.6, -122.3, 1] # Seattle
  *   points << [40.7, -74.0, 2]  # New York
- *   kd = KDTree.new(points)
- *   
+ *   kd = Kdtree.new(points)
+ *
  *   # which city is closest to Portland?
  *   kd.nearest(45.5, -122.8) #=> 1
  *   # which city is closest to Boston?
@@ -208,52 +200,59 @@ static float n_dist;
  */
 static VALUE kdtree_nearest(VALUE kdtree, VALUE x, VALUE y)
 {
+    int n_index;
+    float n_dist;
     KDTREEP;
 
     n_index = -1;
     n_dist = INT_MAX;
-    kdtree_nearest0(kdtreep, kdtreep->root, NUM2DBL(x), NUM2DBL(y), 0);
+
+    kdtree_nearest0(kdtreep, kdtreep->root, NUM2DBL(x), NUM2DBL(y), 0, &n_index, &n_dist);
     if (n_index == -1) {
         return -1;
     }
     return INT2NUM((kdtreep->nodes + n_index)->id);
 }
 
-static void kdtree_nearest0(struct kdtree_data *kdtreep, int i, float x, float y, int depth)
+static void kdtree_nearest0(struct kdtree_data *kdtreep, int i, float x, float y, int depth, int *n_index, float *n_dist)
 {
+    struct kdtree_node *n;
+    float ad;
+    int near, far;
+    float dx;
+
     if (i == -1) {
         return;
     }
-    
-    struct kdtree_node *n = kdtreep->nodes + i;
 
-    float ad = (depth % 2) ? (x - n->x) : (y - n->y);
+    n = kdtreep->nodes + i;
+
+    ad = (depth % 2) ? (x - n->x) : (y - n->y);
 
     //
     // recurse near, and perhaps far as well
     //
-    
-    int near, far;
+
     if (ad <= 0) {
         near = n->left; far = n->right;
     } else {
         near = n->right; far = n->left;
     }
-    kdtree_nearest0(kdtreep, near,  x, y, depth + 1);    
-    if (ad * ad < n_dist) {
-        kdtree_nearest0(kdtreep, far, x, y, depth + 1);
+    kdtree_nearest0(kdtreep, near,  x, y, depth + 1, n_index, n_dist);
+    if (ad * ad < *n_dist) {
+        kdtree_nearest0(kdtreep, far, x, y, depth + 1, n_index, n_dist);
     }
 
     //
     // do we beat the old distance?
     //
-    
-    float dx = (x - n->x) * (x - n->x);
-    if (dx < n_dist) {
+
+    dx = (x - n->x) * (x - n->x);
+    if (dx < *n_dist) {
         float d = dx + ((y - n->y) * (y - n->y));
-        if (d < n_dist) {
-            n_index = i;
-            n_dist = d;
+        if (d < *n_dist) {
+            *n_index = i;
+            *n_dist = d;
         }
     }
 }
@@ -264,15 +263,6 @@ static void kdtree_nearest0(struct kdtree_data *kdtreep, int i, float x, float y
 
 #define MAX_K 255
 
-typedef struct kresult {
-    int index;
-    float distance;
-} kresult;
-// note I leave an extra slot here at the end because of the way our binary insert works
-static struct kresult k_list[MAX_K + 1];
-static int k_len;
-static float k_dist;
-
 /*
  * call-seq:
  *   kd.nearestk(x, y, k)    => array
@@ -280,77 +270,84 @@ static float k_dist;
  * Finds the <i>k</i> points closest to <i>x</i>, <i>y</i>. Returns an array of
  * ids, sorted by distance. Returns an empty array if the tree is empty. Note
  * that <i>k</i> is capped at 255.
- * 
+ *
  *   points = []
  *   points << [47.6, -122.3, 1] # Seattle
  *   points << [45.5, -122.8, 2] # Portland
  *   points << [40.7, -74.0,  3] # New York
- *   kd = KDTree.new(points)
- *   
+ *   kd = Kdtree.new(points)
+ *
  *   # which two cities are closest to San Francisco?
  *   kd.nearest(34.1, -118.2) #=> [2, 1]
  */
 static VALUE kdtree_nearestk(VALUE kdtree, VALUE x, VALUE y, VALUE k)
 {
+    // note I leave an extra slot here at the end because of the way our binary insert works
+    kresult k_list[MAX_K + 1];
+    int k_len = 0;
+    float k_dist = INT_MAX;
+    int ki = NUM2INT(k);
+    VALUE ary;
+    int i;
     KDTREEP;
 
-    k_len = 0;
-    k_dist = INT_MAX;
-
-    int ki = NUM2INT(k);
     if (ki < 1) {
         ki = 1;
     } else if (ki > MAX_K) {
         ki = MAX_K;
     }
-    kdtree_nearestk0(kdtreep, kdtreep->root, NUM2DBL(x), NUM2DBL(y), ki, 0);
+    kdtree_nearestk0(kdtreep, kdtreep->root, NUM2DBL(x), NUM2DBL(y), ki, 0, k_list, &k_len, &k_dist);
 
     // convert result to ruby array
-    VALUE ary = rb_ary_new();
-    int i;
+    ary = rb_ary_new();
     for (i = 0; i < k_len; ++i) {
         rb_ary_push(ary, INT2NUM(kdtreep->nodes[k_list[i].index].id));
     }
     return ary;
 }
 
-static void kdtree_nearestk0(struct kdtree_data *kdtreep, int i, float x, float y, int k, int depth)
+static void kdtree_nearestk0(struct kdtree_data *kdtreep, int i, float x, float y, int k, int depth, kresult *k_list, int *k_len, float *k_dist)
 {
+    struct kdtree_node *n;
+    float ad;
+    int near, far;
+    float dx;
+    int lo, hi;
+
     if (i == -1) {
         return;
     }
-    
-    struct kdtree_node *n = kdtreep->nodes + i;
 
-    float ad = (depth % 2) ? (x - n->x) : (y - n->y);
+    n = kdtreep->nodes + i;
+
+    ad = (depth % 2) ? (x - n->x) : (y - n->y);
 
     //
     // recurse near, and then perhaps far as well
     //
-    
-    int near, far;
+
     if (ad <= 0) {
         near = n->left; far = n->right;
     } else {
         near = n->right; far = n->left;
     }
-    kdtree_nearestk0(kdtreep, near,  x, y, k, depth + 1);
-    if (ad * ad < k_dist) {
-        kdtree_nearestk0(kdtreep, far, x, y, k, depth + 1);
+    kdtree_nearestk0(kdtreep, near,  x, y, k, depth + 1, k_list, k_len, k_dist);
+    if (ad * ad < *k_dist) {
+        kdtree_nearestk0(kdtreep, far, x, y, k, depth + 1, k_list, k_len, k_dist);
     }
 
     //
     // do we beat the old distance?
     //
-    
-    float dx = (x - n->x) * (x - n->x);
-    if (dx < k_dist) {
+
+    dx = (x - n->x) * (x - n->x);
+    if (dx < *k_dist) {
         float d = dx + ((y - n->y) * (y - n->y));
-        if (d < k_dist) {
+        if (d < *k_dist) {
             //
             // find spot to insert
             //
-            int lo = 0, hi = k_len;
+            lo = 0, hi = *k_len;
             while (lo < hi) {
                 int mid = (lo + hi) / 2;
                 if (k_list[mid].distance < d) {
@@ -363,21 +360,21 @@ static void kdtree_nearestk0(struct kdtree_data *kdtreep, int i, float x, float
             //
             // insert
             //
-            
-            memmove(k_list + lo + 1, k_list + lo, (k_len - lo) * sizeof(struct kresult));
+
+            memmove(k_list + lo + 1, k_list + lo, (*k_len - lo) * sizeof(struct kresult));
             k_list[lo].index = i;
             k_list[lo].distance = d;
 
             //
             // adjust len/dist if necessary
             //
-            
-            if (k_len < k) {
-                ++k_len;
+
+            if (*k_len < k) {
+                ++(*k_len);
             } else {
-                k_dist = k_list[k - 1].distance;
+                *k_dist = k_list[k - 1].distance;
             }
-         }
+        }
     }
 }
 
@@ -386,43 +383,42 @@ static void kdtree_nearestk0(struct kdtree_data *kdtreep, int i, float x, float
  *   kd.persist(io)
  *
  * Writes the tree out to <i>io</i> so you can quickly load it later with
- * KDTree.new. This avoids the startup cost of initializing a tree. Apart from a
+ * Kdtree.new. This avoids the startup cost of initializing a tree. Apart from a
  * small header, the size of the file is proportional to the number of points,
  * requiring 20 bytes per point.
  *
  * This file is <b>NOT PORTABLE</b> across different architectures due to endian
  * issues.
- * 
+ *
  *   points = []
  *   points << [47.6, -122.3, 1] # Seattle
  *   points << [45.5, -122.8, 2] # Portland
  *   points << [40.7, -74.0,  3] # New York
- *   kd = KDTree.new(points)
+ *   kd = Kdtree.new(points)
  *
  *   # persist the tree to disk
  *   File.open("treefile", "w") { |f| kd.persist(f) }
  *
  *   ...
- *   
+ *
  *   # later, read the tree from disk
- *   kd2 = File.open("treefile") { |f| KDTree.new(f) }
+ *   kd2 = File.open("treefile") { |f| Kdtree.new(f) }
  */
 static VALUE kdtree_persist(VALUE kdtree, VALUE io)
 {
+    VALUE str;
     KDTREEP;
-    
+
     if (!rb_respond_to(io, rb_intern("write"))) {
         rb_raise(rb_eTypeError, "instance of IO needed");
     }
-    if (rb_respond_to(io, rb_intern("binmode"))) {
-        rb_funcall2(io, rb_intern("binmode"), 0, 0);
+    if (rb_respond_to(io, id_binmode)) {
+        rb_funcall(io, id_binmode, 0);
     }
 
-    VALUE str = rb_str_buf_new(0);
-    rb_str_buf_cat(str, KDTREE_MAGIC, 4);
-    rb_str_buf_cat(str, (char*)kdtreep, sizeof(struct kdtree_data) - sizeof(struct kdtree_node *));
-    rb_str_buf_cat(str, (char*)kdtreep->nodes, sizeof(struct kdtree_node) * kdtreep->len);
-    rb_io_write(io, str);
+    write_all(io, KDTREE_MAGIC, 4);
+    write_all(io, kdtreep, sizeof(struct kdtree_data) - sizeof(struct kdtree_node *));
+    write_all(io, kdtreep->nodes, sizeof(struct kdtree_node) * kdtreep->len);
     return io;
 }
 
@@ -434,22 +430,40 @@ static VALUE kdtree_persist(VALUE kdtree, VALUE io)
  */
 static VALUE kdtree_to_s(VALUE kdtree)
 {
+    char buf[256];
     KDTREEP;
 
-    char buf[256];
     sprintf(buf, "#<%s:%p nodes=%d>", rb_obj_classname(kdtree), (void*)kdtree, kdtreep->len);
     return rb_str_new(buf, strlen(buf));
 }
 
+//
+// io helpers
+//
+
+static void read_all(VALUE io, void *buf, int len)
+{
+    VALUE string = rb_funcall(io, id_read, 1, INT2NUM(len));
+    if (NIL_P(string) || RSTRING_LEN(string) != len) {
+        rb_raise(rb_eEOFError, "end of file reached");
+    }
+    memcpy(buf, RSTRING_PTR(string), len);
+}
+
+static void write_all(VALUE io, const void *buf, int len)
+{
+    rb_funcall(io, id_write, 1, rb_str_new(buf, len));
+}
+
 //
 // entry point
 //
 
 /*
- *  KDTree is an insanely fast data structure for finding the nearest
+ *  Kdtree is an insanely fast data structure for finding the nearest
  *  neighbor(s) to a given point. This implementation only supports 2d
  *  points. Also, it only supports static points - there is no way to edit the
- *  tree after it has been initialized. KDTree should scale to millions of
+ *  tree after it has been initialized. Kdtree should scale to millions of
  *  points, though it's only been tested with around 1 million.
  *
  *  Once the tree is constructed, it can be searched with nearest and nearestk.
@@ -462,8 +476,8 @@ static VALUE kdtree_to_s(VALUE kdtree)
  *   points << [47.6, -122.3, 1] # Seattle
  *   points << [45.5, -122.8, 2] # Portland
  *   points << [40.7, -74.0,  3] # New York
- *   kd = KDTree.new(points)
- *   
+ *   kd = Kdtree.new(points)
+ *
  *   # which city is closest to San Francisco?
  *   kd.nearest(34.1, -118.2) #=> 2
  *   # which two cities are closest to San Francisco?
@@ -477,12 +491,17 @@ void Init_kdtree()
 {
     static VALUE clazz;
 
-    clazz = rb_define_class("KDTree", rb_cObject);
-    
-    rb_define_alloc_func(clazz, kdtree_alloc);    
+    clazz = rb_define_class("Kdtree", rb_cObject);
+
+    rb_define_alloc_func(clazz, kdtree_alloc);
     rb_define_method(clazz, "initialize", kdtree_initialize, 1);
     rb_define_method(clazz, "nearest", kdtree_nearest, 2);
     rb_define_method(clazz, "nearestk", kdtree_nearestk, 3);
-    rb_define_method(clazz, "persist", kdtree_persist, 1);    
-    rb_define_method(clazz, "to_s", kdtree_to_s, 0);        
+    rb_define_method(clazz, "persist", kdtree_persist, 1);
+    rb_define_method(clazz, "to_s", kdtree_to_s, 0);
+
+    // function ids
+    id_binmode = rb_intern("binmode");
+    id_read = rb_intern("read");
+    id_write = rb_intern("write");
 }

data/kdtree.gemspec

@@ -0,0 +1,22 @@
+Gem::Specification.new do |s|
+  s.name        = "kdtree"
+  s.version     = "0.3"
+
+  s.authors     = ["Adam Doppelt"]
+  s.email       = ["amd@gurge.com"]
+  s.homepage    = "http://github.com/gurgeous/kdtree"
+  s.summary     = "Blazingly fast, native 2d kdtree."
+  s.description = <<EOF
+A kdtree is a data structure that makes it possible to quickly solve
+the nearest neighbor problem. This is a native 2d kdtree suitable for
+production use with millions of points.
+EOF
+
+  s.rubyforge_project = "kdtree"
+  s.add_development_dependency "rake-compiler"
+
+  s.files      = `git ls-files`.split("\n")
+  s.test_files = `git ls-files -- {test,spec,features}/*`.split("\n")
+  s.extensions = ["ext/kdtree/extconf.rb"]
+  s.require_paths = ["lib"]
+end

data/lib/kdtree.rb

@@ -0,0 +1 @@
+require "kdtree.so"

data/test/test_kdtree.rb

@@ -0,0 +1,150 @@
+require "benchmark"
+require "kdtree"
+require "tempfile"
+require "test/unit"
+
+#
+# create a tree
+#
+
+class KdtreeTest < Test::Unit::TestCase
+  TMP = "#{Dir.tmpdir}/kdtree_test"
+
+  def setup
+    @points = (0...1000).map { |i| [rand_coord, rand_coord, i] }
+    @kdtree = Kdtree.new(@points)
+  end
+
+  def teardown
+    File.unlink(TMP) if File.exists?(TMP)
+  end
+
+  def test_nearest
+    100.times do
+      pt = [rand_coord, rand_coord]
+
+      # kdtree search
+      id = @kdtree.nearest(pt[0], pt[1])
+      kdpt = @points[id]
+
+      # slow search
+      sortpt = @points.sort_by { |i| distance(i, pt) }.first
+
+      # assert
+      kdd = distance(kdpt, pt)
+      sortd = distance(sortpt, pt)
+      assert((kdd - sortd).abs < 0.0000001, "kdtree didn't return the closest result")
+    end
+  end
+
+  def test_nearestk
+    100.times do
+      pt = [rand_coord, rand_coord]
+
+      # kdtree search
+      list = @kdtree.nearestk(pt[0], pt[1], 5)
+      kdpt = @points[list.last]
+
+      # slow search
+      sortpt = @points.sort_by { |i| distance(i, pt) }[list.length - 1]
+
+      # assert
+      kdd = distance(kdpt, pt)
+      sortd = distance(sortpt, pt)
+      assert((kdd - sortd).abs < 0.0000001, "kdtree didn't return the closest result")
+    end
+  end
+
+  def test_persist
+    # write
+    File.open(TMP, "w") { |f| @kdtree.persist(f) }
+    # read
+    kdtree2 = File.open(TMP, "r") { |f| Kdtree.new(f) }
+
+    # now test some random points
+    100.times do
+      pt = [rand_coord, rand_coord]
+      id1 = @kdtree.nearest(*pt)
+      id2 = kdtree2.nearest(*pt)
+      assert(id1 == id2, "kdtree2 differed from kdtree")
+    end
+  end
+
+  def test_bad_magic
+    File.open(TMP, "w") { |f| f.puts "That ain't right" }
+    assert_raise RuntimeError do
+      File.open(TMP, "r") { |f| Kdtree.new(f) }
+    end
+  end
+
+  def test_eof
+    File.open(TMP, "w") { |f| @kdtree.persist(f) }
+    bytes = File.read(TMP)
+
+    [2, 10, 100].each do |len|
+      File.open(TMP, "w") { |f| f.write(bytes[0, len]) }
+      assert_raise EOFError do
+        File.open(TMP, "r") { |f| Kdtree.new(f) }
+      end
+    end
+  end
+
+  def dont_test_speed
+    sizes = [1, 100, 1000, 10000, 100000, 1000000]
+    ks = [1, 5, 50, 255]
+    sizes.each do |s|
+      points = (0...s).map { |i| [rand_coord, rand_coord, i] }
+
+      # build
+      Benchmark.bm(17) do |bm|
+        kdtree = nil
+        bm.report "build" do
+          kdtree = Kdtree.new(points)
+        end
+        bm.report "persist" do
+          File.open(TMP, "w") { |f| kdtree.persist(f) }
+        end
+        bm.report "read" do
+          File.open(TMP, "r") { |f| Kdtree.new(f) }
+        end
+
+        ks.each do |k|
+          bm.report "100 queries (#{k})" do
+            total = count = 0
+            100.times do
+              tm = Time.now
+              if k == 1
+                kdtree.nearest(rand_coord, rand_coord)
+              else
+                kdtree.nearestk(rand_coord, rand_coord, k)
+              end
+            end
+          end
+        end
+      end
+      puts
+    end
+  end
+
+  protected
+
+  def distance(a, b)
+    x, y = a[0] - b[0], a[1] - b[1]
+    x * x + y * y
+  end
+
+  def rand_coord
+    rand(0) * 10 - 5
+  end
+end
+
+# running dont_test_speed on my i5 2.8ghz:
+#
+#                         user     system      total        real
+# build               3.350000   0.020000   3.370000 (  3.520528)
+# persist             0.150000   0.020000   0.170000 (  0.301963)
+# read                0.280000   0.000000   0.280000 (  0.432676)
+# 100 queries (1)     0.000000   0.000000   0.000000 (  0.000319)
+# 100 queries (5)     0.000000   0.000000   0.000000 (  0.000412)
+# 100 queries (50)    0.000000   0.000000   0.000000 (  0.001417)
+# 100 queries (255)   0.000000   0.000000   0.000000 (  0.006268)

metadata

@@ -1,61 +1,85 @@
---- !ruby/object:Gem::Specification 
+--- !ruby/object:Gem::Specification
 name: kdtree
-version: !ruby/object:Gem::Version 
-  version: "0.1"
+version: !ruby/object:Gem::Version
+  version: '0.3'
+  prerelease: 
 platform: ruby
-authors: 
+authors:
 - Adam Doppelt
 autorequire: 
 bindir: bin
 cert_chain: []
+date: 2012-10-17 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: rake-compiler
+  requirement: !ruby/object:Gem::Requirement
+    none: false
+    requirements:
+    - - ! '>='
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    none: false
+    requirements:
+    - - ! '>='
+      - !ruby/object:Gem::Version
+        version: '0'
+description: ! 'A kdtree is a data structure that makes it possible to quickly solve
 
-date: 2010-01-21 00:00:00 -08:00
-default_executable: 
-dependencies: []
+  the nearest neighbor problem. This is a native 2d kdtree suitable for
 
-description: 
-email: amd@gurge.com
-executables: []
+  production use with millions of points.
 
-extensions: 
-- ext/extconf.rb
+'
+email:
+- amd@gurge.com
+executables: []
+extensions:
+- ext/kdtree/extconf.rb
 extra_rdoc_files: []
-
-files: 
-- ext/extconf.rb
-- ext/kdtree.c
+files:
+- .gitignore
+- .travis.yml
+- Gemfile
 - LICENSE
-- test/test.rb
-has_rdoc: true
-homepage: 
+- README.md
+- Rakefile
+- ext/kdtree/extconf.rb
+- ext/kdtree/kdtree.c
+- kdtree.gemspec
+- lib/kdtree.rb
+- test/test_kdtree.rb
+homepage: http://github.com/gurgeous/kdtree
 licenses: []
-
 post_install_message: 
-rdoc_options: 
-- --exclude
-- test
-- --exclude
-- extconf
-require_paths: 
-- .
-required_ruby_version: !ruby/object:Gem::Requirement 
-  requirements: 
-  - - ">="
-    - !ruby/object:Gem::Version 
-      version: 1.8.5
-  version: 
-required_rubygems_version: !ruby/object:Gem::Requirement 
-  requirements: 
-  - - ">="
-    - !ruby/object:Gem::Version 
-      version: "0"
-  version: 
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  none: false
+  requirements:
+  - - ! '>='
+    - !ruby/object:Gem::Version
+      version: '0'
+      segments:
+      - 0
+      hash: -3094601017742930682
+required_rubygems_version: !ruby/object:Gem::Requirement
+  none: false
+  requirements:
+  - - ! '>='
+    - !ruby/object:Gem::Version
+      version: '0'
+      segments:
+      - 0
+      hash: -3094601017742930682
 requirements: []
-
-rubyforge_project: 
-rubygems_version: 1.3.5
+rubyforge_project: kdtree
+rubygems_version: 1.8.21
 signing_key: 
 specification_version: 3
-summary: Blazingly fast 2d kdtree.
-test_files: 
-- test/test.rb
+summary: Blazingly fast, native 2d kdtree.
+test_files: []

data/test/test.rb

@@ -1,138 +0,0 @@
-require "#{File.expand_path(File.dirname(__FILE__))}/../ext/kdtree.o"
-require "test/unit"
-require "tempfile"
-
-#
-# create a tree
-#
-
-class KDTreeTest < Test::Unit::TestCase
-  TMP = "#{Dir.tmpdir}/kdtree_test"
-  
-  def test_nearest
-    setup_tree(1000)
-    100.times do
-      pt = [rand_coord, rand_coord]
-
-      # kdtree search
-      id = @kdtree.nearest(pt[0], pt[1])
-      kdpt = @points[id]
-
-      # slow search
-      sortpt = @points.sort_by { |i| distance(i, pt) }.first
-
-      # assert
-      kdd = distance(kdpt, pt)
-      sortd = distance(sortpt, pt)
-      assert((kdd - sortd).abs < 0.0000001, "kdtree didn't return the closest result")
-    end
-  end
-
-  def test_nearestk
-    setup_tree(1000)    
-    100.times do
-      pt = [rand_coord, rand_coord]
-
-      # kdtree search
-      list = @kdtree.nearestk(pt[0], pt[1], 5)
-      kdpt = @points[list.last]
-
-      # slow search
-      sortpt = @points.sort_by { |i| distance(i, pt) }[list.length - 1]
-
-      # assert
-      kdd = distance(kdpt, pt)
-      sortd = distance(sortpt, pt)
-      assert((kdd - sortd).abs < 0.0000001, "kdtree didn't return the closest result")
-    end
-  end
-
-  def test_persist
-    setup_tree(1000)
-
-    begin
-      # write
-      File.open(TMP, "w") { |f| @kdtree.persist(f) }
-      # read
-      kdtree2 = File.open(TMP, "r") { |f| KDTree.new(f) }
-
-      # now test some random points
-      100.times do
-        pt = [rand_coord, rand_coord]
-        id1 = @kdtree.nearest(*pt)
-        id2 = kdtree2.nearest(*pt)
-        assert(id1 == id2, "kdtree2 differed from kdtree")
-      end
-    ensure
-      File.unlink(TMP)
-    end
-
-    # now test magic problems
-    begin
-      File.open(TMP, "w") { |f| f.puts "That ain't right" }
-      assert_raise RuntimeError do
-        File.open(TMP, "r") { |f| KDTree.new(f) }
-      end
-    ensure
-      File.unlink(TMP)
-    end
-  end
-
-  def dont_test_speed
-    printf("\n")
-    sizes = [1, 100, 1000, 10000, 100000, 1000000]
-    ks = [1, 5, 50, 255]
-    sizes.each do |s|
-      points = (0...s).map { |i| [rand_coord, rand_coord, i] }
-
-      # build
-      tm = Time.now
-      kdtree = KDTree.new(points)
-      printf "build %d took %.6fs\n", s, Time.now - tm
-      
-      begin
-        # write
-        tm = Time.now
-        File.open(TMP, "w") { |f| kdtree.persist(f) }
-        printf "write %d took %.6fs\n", s, Time.now - tm
-        # read
-        tm = Time.now
-        File.open(TMP, "r") { |f| KDTree.new(f) }
-        printf "read  %d took %.6fs\n", s, Time.now - tm        
-      ensure
-        File.unlink(TMP)
-      end
-      
-      ks.each do |k|
-        total = count = 0
-        100.times do
-          tm = Time.now
-          if k == 1
-            kdtree.nearest(rand_coord, rand_coord)
-          else
-            kdtree.nearestk(rand_coord, rand_coord, k)
-          end
-          total += Time.now - tm
-          count += 1
-        end
-        printf "avg query time = %.6fs [%d/%d]\n", total / count, s, k
-      end
-    end
-  end
-
-  protected
-  
-  def setup_tree(len)
-    @points = (0...len).map { |i| [rand_coord, rand_coord, i] }
-    @kdtree = KDTree.new(@points)
-  end
-  
-  def distance(a, b)
-    x, y = a[0] - b[0], a[1] - b[1]
-    x * x + y * y
-  end
-
-  def rand_coord
-    rand(0) * 10 - 5
-  end
-end