Clean up and correct nearest-neighbor algorithm.

4 files changed, 41 insertions, 27 deletions

diff --git a/Makefile b/Makefile
index 9cc99c6..3814c14 100644
--- a/Makefile
+++ b/Makefile
@@ -10,7 +10,7 @@
 #####################################################################
 
 CC ?= gcc
-CFLAGS ?= -std=c99 -pipe -O2 -flto -Werror -Wall -Wpedantic -Wextra -Wno-sign-compare -Wno-unused-parameter -Wunreachable-code -Wshadow -Wpointer-arith -Wwrite-strings -Wcast-align -Wstrict-prototypes
+CFLAGS ?= -std=c99 -pipe -g -O3 -flto -Werror -Wall -Wpedantic -Wextra -Wno-sign-compare -Wno-unused-parameter -Wunreachable-code -Wshadow -Wpointer-arith -Wwrite-strings -Wcast-align -Wstrict-prototypes
 LDFLAGS ?= -Wl,-O1,--sort-common,--as-needed,-z,relro
 LIBS ?= -lm -lpng
 RM ?= rm -f
diff --git a/kd-forest.c b/kd-forest.c
index be1b574..759c41e 100644
--- a/kd-forest.c
+++ b/kd-forest.c
@@ -137,44 +137,58 @@ kd_build_tree(kd_node_t **buffers[KD_BUFSIZE], size_t size)
 }
 
 static double
-kd_distance_sq(const kd_node_t *a, const kd_node_t *b)
+kd_distance_sq(double a[KD_DIMEN], double b[KD_DIMEN])
 {
   double result = 0.0;
   for (int i = 0; i < KD_DIMEN; ++i) {
-    double d = a->coords[i] - b->coords[i];
+    double d = a[i] - b[i];
     result += d*d;
   }
   return result;
 }
 
 static void
-kd_find_nearest_recursive(kd_node_t *root, const kd_node_t *target, kd_node_t **best, double *limit, unsigned int coord)
+kd_find_nearest_recursive(kd_node_t *node, double target[KD_DIMEN], double closest[KD_DIMEN], kd_node_t **best, double *limit, unsigned int coord)
 {
-  double dist = target->coords[coord] - root->coords[coord];
-  double dist_sq = dist*dist;
-
-  if (!root->removed) {
-    double root_dist_sq = kd_distance_sq(root, target);
-    if (root_dist_sq < *limit) {
-      *best = root;
-      *limit = root_dist_sq;
+  if (!node->removed) {
+    double node_dist_sq = kd_distance_sq(node->coords, target);
+    if (node_dist_sq < *limit) {
+      *limit = node_dist_sq;
+      *best = node;
     }
   }
 
-  coord = (coord + 1)%KD_DIMEN;
+  kd_node_t *first;
+  kd_node_t *second;
+  if (target[coord] < node->coords[coord]) {
+    first = node->left;
+    second = node->right;
+  } else {
+    first = node->right;
+    second = node->left;
+  }
+
+  unsigned int next = (coord + 1)%KD_DIMEN;
 
-  if (root->left && (dist < 0.0 || dist_sq < *limit)) {
-    kd_find_nearest_recursive(root->left, target, best, limit, coord);
+  if (first) {
+    kd_find_nearest_recursive(first, target, closest, best, limit, next);
   }
-  if (root->right && (dist > 0.0 || dist_sq < *limit)) {
-    kd_find_nearest_recursive(root->right, target, best, limit, coord);
+
+  if (second) {
+    double new_closest[KD_DIMEN];
+    memcpy(new_closest, closest, sizeof(new_closest));
+    new_closest[coord] = node->coords[coord];
+
+    if (kd_distance_sq(new_closest, target) < *limit) {
+      kd_find_nearest_recursive(second, target, new_closest, best, limit, next);
+    }
   }
 }
 
 static void
-kd_find_nearest(kd_node_t *root, const kd_node_t *target, kd_node_t **best, double *limit)
+kd_find_nearest(kd_node_t *root, double target[KD_DIMEN], kd_node_t **best, double *limit)
 {
-  kd_find_nearest_recursive(root, target, best, limit, 0);
+  kd_find_nearest_recursive(root, target, target, best, limit, 0);
 }
 
 void
@@ -275,7 +289,7 @@ kdf_remove(kd_forest_t *kdf, kd_node_t *node)
 }
 
 kd_node_t *
-kdf_find_nearest(kd_forest_t *kdf, const kd_node_t *target)
+kdf_find_nearest(kd_forest_t *kdf, double target[KD_DIMEN])
 {
   double limit = INFINITY;
   kd_node_t *best = NULL;
diff --git a/kd-forest.h b/kd-forest.h
index 4cf5750..3607903 100644
--- a/kd-forest.h
+++ b/kd-forest.h
@@ -51,6 +51,6 @@ void kdf_init(kd_forest_t *kdf);
 void kdf_destroy(kd_forest_t *kdf);
 void kdf_insert(kd_forest_t *kdf, kd_node_t *node);
 void kdf_remove(kd_forest_t *kdf, kd_node_t *node);
-kd_node_t *kdf_find_nearest(kd_forest_t *kdf, const kd_node_t *target);
+kd_node_t *kdf_find_nearest(kd_forest_t *kdf, double target[KD_DIMEN]);
 
 #endif // KD_FOREST_H
diff --git a/main.c b/main.c
index 7558f96..c2b6c55 100644
--- a/main.c
+++ b/main.c
@@ -483,16 +483,16 @@ generate_bitmap(const state_t *state)
 
       uint32_t color = state->colors[j];
 
-      kd_node_t target;
+      double target[KD_DIMEN];
       switch (state->options->color_space) {
       case COLOR_SPACE_RGB:
-        color_set_RGB(target.coords, color);
+        color_set_RGB(target, color);
         break;
       case COLOR_SPACE_LAB:
-        color_set_Lab(target.coords, color);
+        color_set_Lab(target, color);
         break;
       case COLOR_SPACE_LUV:
-        color_set_Luv(target.coords, color);
+        color_set_Luv(target, color);
         break;
       }
 
@@ -501,11 +501,11 @@ generate_bitmap(const state_t *state)
         // First node goes in the center
         new_node = nodes + state->size/2 + state->width/2;
       } else {
-        kd_node_t *nearest = kdf_find_nearest(&kdf, &target);
+        kd_node_t *nearest = kdf_find_nearest(&kdf, target);
         new_node = next_neighbor(state, nearest);
       }
 
-      memcpy(new_node->coords, target.coords, sizeof(target.coords));
+      memcpy(new_node->coords, target, sizeof(target));
       kdf_insert(&kdf, new_node);
       remove_non_boundary(state, &kdf, new_node);