Support average selection.

6 files changed, 244 insertions, 67 deletions

diff --git a/Makefile b/Makefile
index b611bdf..0e1ba9d 100644
--- a/Makefile
+++ b/Makefile
@@ -10,17 +10,17 @@
 #####################################################################
 
 CC = gcc
-CFLAGS = -std=c99 -D_POSIX_C_SOURCE=200809L -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
+CFLAGS = -std=c99 -D_POSIX_C_SOURCE=200809L -pipe -g -O3 -flto -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
 
-HEADERS = color.h kd-forest.h options.h util.h
+DEPS = Makefile color.h kd-forest.h options.h util.h
 
 kd-forest: color.o kd-forest.o main.o options.o util.o
 	$(CC) $(CFLAGS) $(LDFLAGS) $^ $(LIBS) -o kd-forest
 
-%.o: %.c $(HEADERS)
+%.o: %.c $(DEPS)
 	$(CC) $(CFLAGS) -c $< -o $@
 
 image: kd-forest.png
@@ -32,10 +32,10 @@ anim: kd-forest.mkv
 
 kd-forest.mkv: kd-forest
 	$(RM) kd-forest.mkv
-	$(RM) -r frames
-	mkdir -p frames
-	./kd-forest -b 20 -s -c Lab -a -o frames
-	ffmpeg -r 60 -i frames/%04d.png -c:v libx264 -preset veryslow -qp 0 kd-forest.mkv
+	mkdir /tmp/kd-frames
+	./kd-forest -b 21 -s -l mean -c Lab -a -o /tmp/kd-frames
+	ffmpeg -r 60 -i /tmp/kd-frames/%04d.png -c:v libx264 -preset veryslow -qp 0 kd-forest.mkv
+	$(RM) -r /tmp/kd-frames
 
 clean:
 	$(RM) *.o
diff --git a/kd-forest.c b/kd-forest.c
index 3716537..3840a61 100644
--- a/kd-forest.c
+++ b/kd-forest.c
@@ -16,29 +16,49 @@
 #include <stdlib.h>
 #include <string.h>
 
-void
-kd_node_init(kd_node_t *node, unsigned int x, unsigned int y)
+kd_node_t *
+new_kd_node(double coords[KD_DIMEN], unsigned int x, unsigned int y)
 {
+  kd_node_t *node = xmalloc(sizeof(kd_node_t));
+  memcpy(node->coords, coords, sizeof(node->coords));
   node->left = node->right = NULL;
   node->x = x;
   node->y = y;
-  node->added = node->removed = false;
+  node->removed = false;
+  return node;
+}
+
+static void
+kd_destroy(kd_node_t *node)
+{
+  if (node) {
+    kd_destroy(node->left);
+    kd_destroy(node->right);
+    free(node);
+  }
 }
 
 static size_t
 kd_collect_nodes(kd_node_t *root, kd_node_t **buffer, bool include_removed)
 {
   size_t count = 0;
+
   if (include_removed || !root->removed) {
     buffer[0] = root;
     ++count;
   }
+
   if (root->left) {
     count += kd_collect_nodes(root->left, buffer + count, include_removed);
   }
   if (root->right) {
     count += kd_collect_nodes(root->right, buffer + count, include_removed);
   }
+
+  if (root->removed && !include_removed) {
+    free(root);
+  }
+
   return count;
 }
 
@@ -202,6 +222,10 @@ kdf_init(kd_forest_t *kdf)
 void
 kdf_destroy(kd_forest_t *kdf)
 {
+  for (unsigned int i = 0; i < kdf->roots_size; ++i) {
+    kd_destroy(kdf->roots[i]);
+  }
+
   free(kdf->roots);
 }
 
@@ -274,8 +298,6 @@ kdf_balance(kd_forest_t *kdf, kd_node_t *new_node, bool force)
 void
 kdf_insert(kd_forest_t *kdf, kd_node_t *node)
 {
-  node->added = true;
-
   // If half or more of the nodes are removed, force a complete rebalance
   bool force = (kdf->size_est + 1) >= 2*(kdf->size + 1);
   kdf_balance(kdf, node, force);
diff --git a/kd-forest.h b/kd-forest.h
index 10756a2..3651bfe 100644
--- a/kd-forest.h
+++ b/kd-forest.h
@@ -26,14 +26,14 @@ typedef struct kd_node_t {
   struct kd_node_t *left, *right;
   // Used to keep track of which sub-tree a node is in during construction
   bool is_left;
-  // State flags
-  bool added, removed;
+  // Weak delete support
+  bool removed;
 
   // Corresponding image position for this node
   unsigned int x, y;
 } kd_node_t;
 
-void kd_node_init(kd_node_t *node, unsigned int x, unsigned int y);
+kd_node_t *new_kd_node(double coords[KD_DIMEN], unsigned int x, unsigned int y);
 
 // A forest of k-d trees
 typedef struct {
diff --git a/main.c b/main.c
index eaa7f01..bf4fe75 100644
--- a/main.c
+++ b/main.c
@@ -25,6 +25,14 @@
 #  include <unistd.h>
 #endif
 
+// A single pixel in all its glory
+typedef struct {
+  double value[KD_DIMEN];
+  kd_node_t *node;
+  unsigned int x, y;
+  bool filled;
+} pixel_t;
+
 // All-encompasing state struct
 typedef struct {
   const options_t *options;
@@ -32,11 +40,12 @@ typedef struct {
   unsigned int height;
   size_t size;
   uint32_t *colors;
+  pixel_t *pixels;
   png_byte **bitmap;
 } state_t;
 
 static void init_state(state_t *state, const options_t *options);
-static void generate_image(const state_t *state);
+static void generate_image(state_t *state);
 static void destroy_state(state_t *state);
 
 // Entry point
@@ -100,6 +109,22 @@ create_colors(const state_t *state)
   return colors;
 }
 
+static pixel_t *
+create_pixels(const state_t *state)
+{
+  pixel_t *pixels = xmalloc(state->size*sizeof(pixel_t));
+  for (unsigned int y = 0, i = 0; y < state->height; ++y) {
+    for (unsigned int x = 0; x < state->width; ++x, ++i) {
+      pixel_t *pixel = pixels + i;
+      pixel->node = NULL;
+      pixel->x = x;
+      pixel->y = y;
+      pixel->filled = false;
+    }
+  }
+  return pixels;
+}
+
 static png_byte **
 create_bitmap(const state_t *state)
 {
@@ -126,14 +151,15 @@ init_state(state_t *state, const options_t *options)
          options->bit_depth, state->width, state->height, state->size);
 
   state->colors = create_colors(state);
+  state->pixels = create_pixels(state);
   state->bitmap = create_bitmap(state);
 }
 
-static void generate_bitmap(const state_t *state);
+static void generate_bitmap(state_t *state);
 static void write_png(const state_t *state, const char *filename);
 
 static void
-generate_image(const state_t *state)
+generate_image(state_t *state)
 {
   generate_bitmap(state);
 
@@ -149,68 +175,185 @@ destroy_state(state_t *state)
     free(state->bitmap[i]);
   }
   free(state->bitmap);
+  free(state->pixels);
   free(state->colors);
 }
 
-static kd_node_t *
-try_neighbor(const state_t *state, kd_node_t *node, int dx, int dy)
+static pixel_t *
+get_pixel(const state_t *state, unsigned int x, unsigned int y)
 {
-  if (dx < 0 && node->x < -dx) {
+  return state->pixels + state->width*y + x;
+}
+
+static pixel_t *
+try_neighbor(const state_t *state, pixel_t *pixel, int dx, int dy)
+{
+  if (dx < 0 && pixel->x < -dx) {
     return NULL;
-  } else if (dx > 0 && node->x + dx >= state->width) {
+  } else if (dx > 0 && pixel->x + dx >= state->width) {
     return NULL;
-  } else if (dy < 0 && node->y < -dy) {
+  } else if (dy < 0 && pixel->y < -dy) {
     return NULL;
-  } else if (dy > 0 && node->y + dy >= state->height) {
+  } else if (dy > 0 && pixel->y + dy >= state->height) {
     return NULL;
   }
 
-  return node + (int)state->width*dy + dx;
+  return pixel + (int)state->width*dy + dx;
 }
 
-static kd_node_t *
-next_neighbor(const state_t *state, kd_node_t *node)
+static unsigned int
+get_all_neighbors(const state_t *state, pixel_t *pixel, pixel_t *neighbors[8])
 {
-  kd_node_t *neighbors[8];
   unsigned int size = 0;
+
   for (int dy = -1; dy <= 1; ++dy) {
     for (int dx = -1; dx <= 1; ++dx) {
       if (dx == 0 && dy == 0) {
         continue;
       }
 
-      kd_node_t *neighbor = try_neighbor(state, node, dx, dy);
-      if (neighbor && !neighbor->added) {
+      pixel_t *neighbor = try_neighbor(state, pixel, dx, dy);
+      if (neighbor) {
         neighbors[size++] = neighbor;
       }
     }
   }
 
-  if (size == 0) {
-    return NULL;
+  return size;
+}
+
+static unsigned int
+get_neighbors(const state_t *state, pixel_t *pixel, pixel_t *neighbors[8], bool filled)
+{
+  unsigned int size = 0;
+
+  for (int dy = -1; dy <= 1; ++dy) {
+    for (int dx = -1; dx <= 1; ++dx) {
+      if (dx == 0 && dy == 0) {
+        continue;
+      }
+
+      pixel_t *neighbor = try_neighbor(state, pixel, dx, dy);
+      if (neighbor && neighbor->filled == filled) {
+        neighbors[size++] = neighbor;
+      }
+    }
   }
 
+  return size;
+}
+
+static pixel_t *
+select_empty_neighbor(const state_t *state, pixel_t *pixel)
+{
+  pixel_t *neighbors[8];
+  unsigned int size = get_neighbors(state, pixel, neighbors, false);
   return neighbors[xrand(size)];
 }
 
-static void
-remove_if_surrounded(const state_t *state, kd_forest_t *kdf, kd_node_t *node)
+static pixel_t *
+find_next_pixel(const state_t *state, const kd_forest_t *kdf, const double target[KD_DIMEN])
 {
-  if (node->added && !node->removed && next_neighbor(state, node) == NULL) {
-    kdf_remove(kdf, node);
+  kd_node_t *nearest = kdf_find_nearest(kdf, target);
+  pixel_t *pixel = get_pixel(state, nearest->x, nearest->y);
+
+  if (state->options->selection == SELECTION_MIN) {
+    pixel = select_empty_neighbor(state, pixel);
   }
+
+  return pixel;
 }
 
 static void
-remove_non_boundary(const state_t *state, kd_forest_t *kdf, kd_node_t *node)
+ensure_pixel_removed(kd_forest_t *kdf, pixel_t *pixel)
+{
+  if (pixel->node) {
+    kdf_remove(kdf, pixel->node);
+    pixel->node = NULL;
+  }
+}
+
+static bool
+has_empty_neighbors(const state_t *state, pixel_t *pixel)
 {
   for (int dy = -1; dy <= 1; ++dy) {
     for (int dx = -1; dx <= 1; ++dx) {
-      kd_node_t *neighbor = try_neighbor(state, node, dx, dy);
-      if (neighbor) {
-        remove_if_surrounded(state, kdf, neighbor);
+      if (dx == 0 && dy == 0) {
+        continue;
+      }
+
+      pixel_t *neighbor = try_neighbor(state, pixel, dx, dy);
+      if (neighbor && !neighbor->filled) {
+        return true;
+      }
+    }
+  }
+
+  return false;
+}
+
+static void
+insert_new_pixel_min(state_t *state, kd_forest_t *kdf, pixel_t *pixel)
+{
+  pixel->filled = true;
+
+  if (has_empty_neighbors(state, pixel)) {
+    pixel->node = new_kd_node(pixel->value, pixel->x, pixel->y);
+    kdf_insert(kdf, pixel->node);
+  }
+
+  pixel_t *neighbors[8];
+  unsigned int size = get_all_neighbors(state, pixel, neighbors);
+  for (unsigned int i = 0; i < size; ++i) {
+    pixel_t *neighbor = neighbors[i];
+    if (!has_empty_neighbors(state, neighbor)) {
+      ensure_pixel_removed(kdf, neighbor);
+    }
+  }
+}
+
+static void
+insert_new_pixel_mean(state_t *state, kd_forest_t *kdf, pixel_t *pixel)
+{
+  pixel->filled = true;
+  ensure_pixel_removed(kdf, pixel);
+
+  pixel_t *neighbors[8];
+  unsigned int size = get_neighbors(state, pixel, neighbors, false);
+  for (unsigned int i = 0; i < size; ++i) {
+    pixel_t *neighbor = neighbors[i];
+
+    double value[KD_DIMEN] = { 0.0 };
+
+    pixel_t *filled[8];
+    unsigned int nfilled = get_neighbors(state, neighbor, filled, true);
+    for (unsigned int j = 0; j < nfilled; ++j) {
+      for (unsigned int k = 0; k < KD_DIMEN; ++k) {
+        value[k] += filled[j]->value[k];
       }
     }
+
+    for (unsigned int j = 0; j < KD_DIMEN; ++j) {
+      value[j] /= nfilled;
+    }
+
+    ensure_pixel_removed(kdf, neighbor);
+    neighbor->node = new_kd_node(value, neighbor->x, neighbor->y);
+    kdf_insert(kdf, neighbor->node);
+  }
+}
+
+static void
+insert_new_pixel(state_t *state, kd_forest_t *kdf, pixel_t *pixel)
+{
+  switch (state->options->selection) {
+  case SELECTION_MIN:
+    insert_new_pixel_min(state, kdf, pixel);
+    break;
+
+  case SELECTION_MEAN:
+    insert_new_pixel_mean(state, kdf, pixel);
+    break;
   }
 }
 
@@ -243,15 +386,8 @@ print_progress(const char *format, ...)
 }
 
 static void
-generate_bitmap(const state_t *state)
+generate_bitmap(state_t *state)
 {
-  kd_node_t *nodes = xmalloc(state->size*sizeof(kd_node_t));
-  for (unsigned int y = 0, i = 0; y < state->height; ++y) {
-    for (unsigned int x = 0; x < state->width; ++x, ++i) {
-      kd_node_init(nodes + y*state->width + x, x, y);
-    }
-  }
-
   // Make the forest
   kd_forest_t kdf;
   kdf_init(&kdf);
@@ -294,25 +430,22 @@ generate_bitmap(const state_t *state)
         break;
       }
 
-      kd_node_t *new_node;
+      pixel_t *pixel;
       if (j == 0) {
         // First node goes in the center
-        new_node = nodes + state->size/2 + state->width/2;
+        pixel = get_pixel(state, state->width/2, state->height/2);
       } else {
-        kd_node_t *nearest = kdf_find_nearest(&kdf, target);
-        new_node = next_neighbor(state, nearest);
+        pixel = find_next_pixel(state, &kdf, target);
       }
 
-      memcpy(new_node->coords, target, sizeof(target));
-      kdf_insert(&kdf, new_node);
-      remove_non_boundary(state, &kdf, new_node);
-
+      memcpy(pixel->value, target, sizeof(target));
+      insert_new_pixel(state, &kdf, pixel);
       if (kdf.size > max_size) {
         max_size = kdf.size;
       }
 
-      png_byte *pixel = state->bitmap[new_node->y] + 3*new_node->x;
-      color_unpack(pixel, color);
+      png_byte *png_pixel = state->bitmap[pixel->y] + 3*pixel->x;
+      color_unpack(png_pixel, color);
     }
   }
 
@@ -335,11 +468,10 @@ generate_bitmap(const state_t *state)
 #endif
   }
 
-  print_progress("%.2f%%\t| boundary size: 0\t| max boundary size: %zu\n",
-                 100.0, max_size);
+  print_progress("%.2f%%\t| boundary size: %zu\t| max boundary size: %zu\n",
+                 100.0, kdf.size, max_size);
 
   kdf_destroy(&kdf);
-  free(nodes);
 }
 
 static void
diff --git a/options.c b/options.c
index 04f1f82..06deddb 100644
--- a/options.c
+++ b/options.c
@@ -226,6 +226,7 @@ print_usage(FILE *file, const char *command)
   usage("Usage:\n");
   usage("  !$! *%s* [-b|--bit-depth @DEPTH@]\n", command);
   usage("    %s [-s|--hue-sort] [-r|--random]\n", whitespace);
+  usage("    %s [-l|--selection @min@|@mean@]\n", whitespace);
   usage("    %s [-c|--color-space @RGB@|@Lab@|@Luv@]\n", whitespace);
   usage("    %s [-a|--animate]\n", whitespace);
   usage("    %s [-o|--output @PATH@]\n", whitespace);
@@ -238,6 +239,10 @@ print_usage(FILE *file, const char *command)
   usage("          Sort colors by hue first (!default!)\n");
   usage("  -r, --random:\n");
   usage("          Randomize colors first\n\n");
+  usage("  -l, --selection @min@|@mean@:\n");
+  usage("          Specify the selection mode (!default!: @min@)\n\n");
+  usage("          @min@:  Pick the pixel with the closest neighboring pixel\n");
+  usage("          @mean@: Pick the pixel with the closest average of all its neighbors\n\n");
   usage("  -c, --color-space @RGB@|@Lab@|@Luv@:\n");
   usage("          Use the given color space (!default!: @Lab@)\n\n");
   usage("  -a, --animate:\n");
@@ -259,10 +264,11 @@ parse_options(options_t *options, int argc, char *argv[])
   // Set defaults
   options->bit_depth = 24;
   options->mode = MODE_HUE_SORT;
-  options->seed = 0;
+  options->selection = SELECTION_MIN;
   options->color_space = COLOR_SPACE_LAB;
   options->animate = false;
   options->filename = NULL;
+  options->seed = 0;
   options->help = false;
 
   bool result = true;
@@ -282,15 +288,19 @@ parse_options(options_t *options, int argc, char *argv[])
       options->mode = MODE_HUE_SORT;
     } else if (parse_arg(argc, argv, "-r", "--random", NULL, &i, &error)) {
       options->mode = MODE_RANDOM;
-    } else if (parse_arg(argc, argv, "-e", "--seed", &value, &i, &error)) {
-      if (!str_to_uint(value, &options->seed)) {
-        fprintf(stderr, "Invalid random seed: `%s'\n", value);
-        error = true;
-      }
     } else if (parse_arg(argc, argv, "-a", "--animate", NULL, &i, &error)) {
       options->animate = true;
     } else if (parse_arg(argc, argv, "-o", "--output", &value, &i, &error)) {
       options->filename = value;
+    } else if (parse_arg(argc, argv, "-l", "--selection", &value, &i, &error)) {
+      if (strcmp(value, "min") == 0) {
+        options->selection = SELECTION_MIN;
+      } else if (strcmp(value, "mean") == 0) {
+        options->selection = SELECTION_MEAN;
+      } else {
+        fprintf(stderr, "Invalid selection mode: `%s'\n", value);
+        error = true;
+      }
     } else if (parse_arg(argc, argv, "-c", "--color-space", &value, &i, &error)) {
       if (strcmp(value, "RGB") == 0) {
         options->color_space = COLOR_SPACE_RGB;
@@ -300,6 +310,12 @@ parse_options(options_t *options, int argc, char *argv[])
         options->color_space = COLOR_SPACE_LUV;
       } else {
         fprintf(stderr, "Invalid color space: `%s'\n", value);
+        error = true;
+      }
+    } else if (parse_arg(argc, argv, "-e", "--seed", &value, &i, &error)) {
+      if (!str_to_uint(value, &options->seed)) {
+        fprintf(stderr, "Invalid random seed: `%s'\n", value);
+        error = true;
       }
     } else if (parse_arg(argc, argv, "-h", "--help", NULL, &i, &error)) {
       options->help = true;
diff --git a/options.h b/options.h
index 6f1b2f8..fe996f7 100644
--- a/options.h
+++ b/options.h
@@ -21,6 +21,12 @@ typedef enum {
   MODE_RANDOM,
 } mode_t;
 
+// Possible pixel selection modes
+typedef enum {
+  SELECTION_MIN,
+  SELECTION_MEAN,
+} selection_t;
+
 // Possible color spaces
 typedef enum {
   COLOR_SPACE_RGB,
@@ -32,6 +38,7 @@ typedef enum {
 typedef struct {
   unsigned int bit_depth;
   mode_t mode;
+  selection_t selection;
   color_space_t color_space;
   bool animate;
   const char *filename;