Refactor main.c to take configuration on the command line instead of #defines.

2 files changed, 358 insertions, 140 deletions

diff --git a/Makefile b/Makefile
index b55c57c..a4f8280 100644
--- a/Makefile
+++ b/Makefile
@@ -19,7 +19,7 @@ kd-forest: kd-forest.c kd-forest.h util.c util.h color.c color.h main.c
 	$(CC) $(CFLAGS) $(LDFLAGS) kd-forest.c util.c color.c main.c $(LIBS) -o kd-forest
 
 kd-forest.png: kd-forest
-	./kd-forest
+	./kd-forest -b 24 -s -c Lab -o kd-forest.png
 
 image: kd-forest.png
 
diff --git a/main.c b/main.c
index b78ddf7..ca13889 100644
--- a/main.c
+++ b/main.c
@@ -16,6 +16,7 @@
 #include <math.h>
 #include <png.h>
 #include <setjmp.h>
+#include <stdarg.h>
 #include <stdio.h>
 #include <stdbool.h>
 #include <stdlib.h>
@@ -24,21 +25,214 @@
 #  include <unistd.h>
 #endif
 
-// Number of trailing zero bits on each color chanel, set to zero for all
-// 24-bit images
-#define BIT_DEPTH 24
-// Whether to sort by hue
-#define HUE_SORT true
+// Possible generation modes
+typedef enum {
+  MODE_HUE_SORT,
+  MODE_RANDOM,
+} mode_t;
+
+// Possible color spaces
+typedef enum {
+  COLOR_SPACE_RGB,
+  COLOR_SPACE_LAB,
+  COLOR_SPACE_LUV,
+} color_space_t;
+
+// Command-line options
+typedef struct {
+  unsigned int bit_depth;
+  mode_t mode;
+  color_space_t color_space;
+  const char *filename;
+  bool help;
+} options_t;
+
+static bool parse_options(options_t *options, int argc, char *argv[]);
+static void print_usage(FILE *file, const char *command);
+
+// All-encompasing state struct
+typedef struct {
+  const options_t *options;
+  unsigned int width;
+  unsigned int height;
+  size_t size;
+  uint32_t *colors;
+  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 destroy_state(state_t *state);
+
+// Entry point
+int
+main(int argc, char *argv[])
+{
+  options_t options;
+  if (!parse_options(&options, argc, argv)) {
+    fprintf(stderr, "\n");
+    print_usage(stderr, argv[0]);
+    return EXIT_FAILURE;
+  }
+
+  if (options.help) {
+    print_usage(stdout, argv[0]);
+    return EXIT_SUCCESS;
+  }
+
+  // For consistent images
+  srand(0);
+
+  state_t state;
+  init_state(&state, &options);
+  generate_image(&state);
+  destroy_state(&state);
+  return EXIT_SUCCESS;
+}
+
+static bool
+parse_arg(int argc, char *argv[],
+          const char *short_form, const char *long_form,
+          const char **value, int *i, bool *error)
+{
+  size_t short_len = strlen(short_form);
+  size_t long_len = strlen(long_form);
+
+  const char *actual_form;
+  const char *arg = argv[*i];
+  const char *candidate = NULL;
+
+  if (strncmp(arg, short_form, short_len) == 0) {
+    actual_form = short_form;
+    if (strlen(arg) > short_len) {
+      candidate = arg + short_len;
+    }
+  } else if (strncmp(arg, long_form, long_len) == 0) {
+    actual_form = long_form;
+    if (strlen(arg) > long_len) {
+      if (arg[long_len] == '=') {
+        candidate = arg + long_len + 1;
+      } else {
+        return false;
+      }
+    }
+  } else {
+    return false;
+  }
+
+  if (value) {
+    if (candidate) {
+      *value = candidate;
+    } else if (*i < argc - 1) {
+      ++*i;
+      *value = argv[*i];
+    } else {
+      fprintf(stderr, "Expected a value for %s\n", arg);
+      *error = true;
+      return false;
+    }
+  } else if (candidate) {
+    fprintf(stderr, "Unexpected value for %s: `%s'\n",
+            actual_form, candidate);
+    *error = true;
+    return false;
+  }
+
+  return true;
+}
+
+static bool
+str_to_uint(const char *str, unsigned int *value)
+{
+  char *endptr;
+  long result = strtol(str, &endptr, 10);
+  if (*str == '\0' || *endptr != '\0') {
+    return false;
+  }
+  if (result < 0 || result > UINT_MAX) {
+    return false;
+  }
+
+  *value = result;
+  return true;
+}
+
+static void
+print_usage(FILE *file, const char *command)
+{
+  size_t length = strlen(command);
+  char whitespace[length];
+  memset(whitespace, ' ', length);
+
+  fprintf(file, "Usage:\n");
+  fprintf(file, "  $ %s [-b|--bit-depth DEPTH]\n", command);
+  fprintf(file, "    %s [-s|--hue-sort] [-r|--random]\n", whitespace);
+  fprintf(file, "    %s [-c|--color-space RGB|Lab|Luv]\n", whitespace);
+  fprintf(file, "    %s [-o|--output PATH]\n", whitespace);
+  fprintf(file, "    %s [-h|--help]\n", whitespace);
+  fprintf(file, "\n");
+  fprintf(file, "  -b, --bit-depth DEPTH:  Use all DEPTH-bit colors (default: 24)\n");
+  fprintf(file, "  -s, --hue-sort:         Sort colors by hue first (default)\n");
+  fprintf(file, "  -r, --random:           Randomize colors first\n");
+  fprintf(file, "  -c, --color-space RGB|Lab|Luv:\n");
+  fprintf(file, "                          Use the given color space (default: Lab)\n");
+  fprintf(file, "  -o, --output PATH:      Output a PNG file at PATH (default: kd-forest.png)\n");
+  fprintf(file, "  -h, --help:             Show this message\n");
+}
+
+static bool
+parse_options(options_t *options, int argc, char *argv[])
+{
+  // Set defaults
+  options->bit_depth = 24;
+  options->mode = MODE_HUE_SORT;
+  options->color_space = COLOR_SPACE_LAB;
+  options->filename = "kd-forest.png";
+  options->help = false;
+
+  bool result = true;
+
+  for (int i = 1; i < argc; ++i) {
+    const char *value;
+    bool error = false;
+
+    if (parse_arg(argc, argv, "-b", "--bit-depth", &value, &i, &error)) {
+      if (!str_to_uint(value, &options->bit_depth)
+          || options->bit_depth <= 0
+          || options->bit_depth > 24) {
+        fprintf(stderr, "Invalid bit depth: `%s'\n", value);
+        error = true;
+      }
+    } else if (parse_arg(argc, argv, "-s", "--hue-sort", NULL, &i, &error)) {
+      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, "-o", "--output", &value, &i, &error)) {
+      options->filename = value;
+    } else if (parse_arg(argc, argv, "-c", "--color-space", &value, &i, &error)) {
+      if (strcmp(value, "RGB") == 0) {
+        options->color_space = COLOR_SPACE_RGB;
+      } else if (strcmp(value, "Lab") == 0) {
+        options->color_space = COLOR_SPACE_LAB;
+      } else if (strcmp(value, "Luv") == 0) {
+        options->color_space = COLOR_SPACE_LUV;
+      } else {
+        fprintf(stderr, "Invalid color space: `%s'\n", value);
+      }
+    } else if (parse_arg(argc, argv, "-h", "--help", NULL, &i, &error)) {
+      options->help = true;
+    } else if (!error) {
+      fprintf(stderr, "Unexpected argument `%s'\n", argv[i]);
+      error = true;
+    }
 
-// Which color space to use
-#define USE_RGB false
-#define USE_LAB true
-#define USE_LUV false
+    if (error) {
+      result = false;
+    }
+  }
 
-// Computed constants
-static const unsigned int WIDTH = 1U << (BIT_DEPTH + 1)/2; // Round up
-static const unsigned int HEIGHT = 1U << (BIT_DEPTH)/2;    // Round down
-static const unsigned int SIZE = 1U << BIT_DEPTH;
+  return result;
+}
 
 static unsigned int
 rand_in(unsigned int range)
@@ -52,20 +246,105 @@ rand_in(unsigned int range)
   return res%range;
 }
 
+static uint32_t *
+create_colors(const state_t *state)
+{
+  const unsigned int bit_depth = state->options->bit_depth;
+
+  // From least to most perceptually important
+  const unsigned int bskip = 1U << (24 - bit_depth)/3;
+  const unsigned int rskip = 1U << (24 - bit_depth + 1)/3;
+  const unsigned int gskip = 1U << (24 - bit_depth + 2)/3;
+
+  uint32_t *colors = xmalloc(state->size*sizeof(uint32_t));
+  for (unsigned int b = 0, i = 0; b < 0x100; b += bskip) {
+    for (unsigned int g = 0; g < 0x100; g += gskip) {
+      for (unsigned int r = 0; r < 0x100; r += rskip, ++i) {
+        colors[i] = (r << 16) | (g << 8) | b;
+      }
+    }
+  }
+
+  switch (state->options->mode) {
+  case MODE_HUE_SORT:
+    qsort(colors, state->size, sizeof(uint32_t), color_comparator);
+    break;
+
+  case MODE_RANDOM:
+    // Fisher-Yates shuffle
+    for (unsigned int i = state->size; i-- > 0;) {
+      unsigned int j = rand_in(i + 1);
+      uint32_t temp = colors[i];
+      colors[i] = colors[j];
+      colors[j] = temp;
+    }
+    break;
+  }
+
+  return colors;
+}
+
+static png_byte **
+create_bitmap(const state_t *state)
+{
+  png_byte **rows = xmalloc(state->height*sizeof(png_byte *));
+  const size_t row_size = 3*state->width*sizeof(png_byte);
+  for (unsigned int i = 0; i < state->height; ++i) {
+    rows[i] = xmalloc(row_size);
+    memset(rows[i], 0, row_size);
+  }
+  return rows;
+}
+
+static void
+init_state(state_t *state, const options_t *options)
+{
+  state->options = options;
+  state->width = 1U << (options->bit_depth + 1)/2; // Round up
+  state->height = 1U << options->bit_depth/2; // Round down
+  state->size = (size_t)state->width*state->height;
+
+  printf("Generating a %u-bit, %ux%u image (%zu pixels)\n",
+         options->bit_depth, state->width, state->height, state->size);
+
+  state->colors = create_colors(state);
+  state->bitmap = create_bitmap(state);
+}
+
+static void generate_bitmap(const state_t *state);
+static void write_png(const state_t *state, const char *filename);
+
+static void
+generate_image(const state_t *state)
+{
+  generate_bitmap(state);
+  write_png(state, state->options->filename);
+}
+
+static void
+destroy_state(state_t *state)
+{
+  for (unsigned int i = 0; i < state->height; ++i) {
+    free(state->bitmap[i]);
+  }
+  free(state->bitmap);
+  free(state->colors);
+}
+
 static kd_node_t *
-try_neighbor(kd_node_t *node, int dx, int dy)
+try_neighbor(const state_t *state, kd_node_t *node, int dx, int dy)
 {
   if (dx < 0 && node->x < -dx) {
     return NULL;
-  } else if (dx > 0 && node->x + dx >= WIDTH) {
+  } else if (dx > 0 && node->x + dx >= state->width) {
     return NULL;
   } else if (dy < 0 && node->y < -dy) {
     return NULL;
-  } else if (dy > 0 && node->y + dy >= HEIGHT) {
+  } else if (dy > 0 && node->y + dy >= state->height) {
     return NULL;
   }
 
-  return node + (int)WIDTH*dy + dx;
+  return node + (int)state->width*dy + dx;
 }
 
 // Star pattern:
@@ -84,12 +363,12 @@ static int neighbor_order[][2] = {
 };
 
 static kd_node_t *
-next_neighbor(kd_node_t *node)
+next_neighbor(const state_t *state, kd_node_t *node)
 {
   unsigned int first = rand_in(8);
   for (unsigned int i = first; i < first + 8; ++i) {
     int *delta = neighbor_order[i%8];
-    kd_node_t *neighbor = try_neighbor(node, delta[0], delta[1]);
+    kd_node_t *neighbor = try_neighbor(state, node, delta[0], delta[1]);
     if (neighbor && !neighbor->added) {
       return neighbor;
     }
@@ -99,84 +378,61 @@ next_neighbor(kd_node_t *node)
 }
 
 static void
-remove_if_surrounded(kd_forest_t *kdf, kd_node_t *node)
+remove_if_surrounded(const state_t *state, kd_forest_t *kdf, kd_node_t *node)
 {
-  if (node->added && !node->removed && next_neighbor(node) == NULL) {
+  if (node->added && !node->removed && next_neighbor(state, node) == NULL) {
     kdf_remove(kdf, node);
   }
 }
 
 static void
-remove_non_boundary(kd_forest_t *kdf, kd_node_t *node)
+remove_non_boundary(const state_t *state, kd_forest_t *kdf, kd_node_t *node)
 {
   for (int dy = -1; dy <= 1; ++dy) {
     for (int dx = -1; dx <= 1; ++dx) {
-      kd_node_t *neighbor = try_neighbor(node, dx, dy);
+      kd_node_t *neighbor = try_neighbor(state, node, dx, dy);
       if (neighbor) {
-        remove_if_surrounded(kdf, neighbor);
+        remove_if_surrounded(state, kdf, neighbor);
       }
     }
   }
 }
 
-static uint32_t *
-create_colors(void)
+static void
+print_progress(const char *format, ...)
 {
-  // From least to most perceptually important
-  const unsigned int bskip = 1U << (24 - BIT_DEPTH)/3;
-  const unsigned int rskip = 1U << (24 - BIT_DEPTH + 1)/3;
-  const unsigned int gskip = 1U << (24 - BIT_DEPTH + 2)/3;
-
-  uint32_t *colors = xmalloc(SIZE*sizeof(uint32_t));
-  for (unsigned int b = 0, i = 0; b < 0x100; b += bskip) {
-    for (unsigned int g = 0; g < 0x100; g += gskip) {
-      for (unsigned int r = 0; r < 0x100; r += rskip, ++i) {
-        colors[i] = (r << 16) | (g << 8) | b;
-      }
-    }
-  }
-
-  if (HUE_SORT) {
-    qsort(colors, SIZE, sizeof(uint32_t), color_comparator);
-  } else {
-    // Fisher-Yates shuffle
-    for (unsigned int i = SIZE; i-- > 0;) {
-      unsigned int j = rand_in(i + 1);
-      uint32_t temp = colors[i];
-      colors[i] = colors[j];
-      colors[j] = temp;
-    }
+#if __unix__
+  static bool tty_checked = false;
+  static bool tty = false;
+  if (!tty_checked) {
+    tty = isatty(STDOUT_FILENO);
+    tty_checked = true;
   }
+  const char *clear_line = tty ? "\033[2K\r" : "";
+  const char *new_line = tty ? "" : "\n";
+#else
+  const char *clear_line = "";
+  const char *new_line = "\n";
+#endif
 
-  return colors;
-}
+  printf("%s", clear_line);
 
-static kd_node_t *
-create_kd_nodes(void)
-{
-  return xmalloc(SIZE*sizeof(kd_node_t));
-}
+  va_list args;
+  va_start(args, format);
+  vprintf(format, args);
+  va_end(args);
 
-static png_byte **
-create_bitmap(void)
-{
-  png_byte **rows = xmalloc(HEIGHT*sizeof(png_byte *));
-  const size_t row_size = 3*WIDTH*sizeof(png_byte);
-  for (unsigned int i = 0; i < HEIGHT; ++i) {
-    rows[i] = xmalloc(row_size);
-    memset(rows[i], 0, row_size);
-  }
-  return rows;
+  printf("%s", new_line);
+  fflush(stdout);
 }
 
 static void
-generate_image(const uint32_t *colors, kd_node_t *nodes,
-               unsigned int initial_x, unsigned int initial_y,
-               png_byte **bitmap)
+generate_bitmap(const state_t *state)
 {
-  for (unsigned int y = 0, i = 0; y < HEIGHT; ++y) {
-    for (unsigned int x = 0; x < WIDTH; ++x, ++i) {
-      kd_node_init(nodes + y*WIDTH + x, x, y);
+  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);
     }
   }
 
@@ -184,71 +440,65 @@ generate_image(const uint32_t *colors, kd_node_t *nodes,
   kd_forest_t kdf;
   kdf_init(&kdf);
 
-#if __unix__
-  bool tty = isatty(1);
-  const char *clear_line = tty ? "\033[2K\r" : "";
-  const char *new_line = tty ? "" : "\n";
-#else
-  const char *clear_line = "";
-  const char *new_line = "\n";
-#endif
-
   size_t max_size = 0;
 
   // Do multiple passes to get rid of artifacts in HUE_SORT mode
-  for (unsigned int i = 1, progress = 0; i <= BIT_DEPTH; ++i) {
+  unsigned int bit_depth = state->options->bit_depth;
+  for (unsigned int i = 1, progress = 0; i <= bit_depth; ++i) {
     unsigned int stripe = 1 << i;
 
-    for (unsigned int j = stripe/2 - 1; j < SIZE; j += stripe, ++progress) {
-      if (progress%WIDTH == 0) {
-        printf("%s%.2f%%\t| boundary size: %zu\t| max boundary size: %zu%s",
-               clear_line, 100.0*progress/SIZE, kdf.size, max_size, new_line);
-        fflush(stdout);
+    for (unsigned int j = stripe/2 - 1; j < state->size; j += stripe, ++progress) {
+      if (progress%state->width == 0) {
+        print_progress("%.2f%%\t| boundary size: %zu\t| max boundary size: %zu",
+                       100.0*progress/state->size, kdf.size, max_size);
       }
 
-      uint32_t color = colors[j];
+      uint32_t color = state->colors[j];
 
       kd_node_t target;
-#if USE_RGB
-      color_set_RGB(target.coords, color);
-#elif USE_LAB
-      color_set_Lab(target.coords, color);
-#elif USE_LUV
-      color_set_Luv(target.coords, color);
-#else
-#  error "Pick one!"
-#endif
+      switch (state->options->color_space) {
+      case COLOR_SPACE_RGB:
+        color_set_RGB(target.coords, color);
+        break;
+      case COLOR_SPACE_LAB:
+        color_set_Lab(target.coords, color);
+        break;
+      case COLOR_SPACE_LUV:
+        color_set_Luv(target.coords, color);
+        break;
+      }
 
       kd_node_t *new_node;
       if (j == 0) {
         // First node goes in the center
-        new_node = nodes + WIDTH*initial_y + initial_x;
+        new_node = nodes + state->size/2 + state->width/2;
       } else {
         kd_node_t *nearest = kdf_find_nearest(&kdf, &target);
-        new_node = next_neighbor(nearest);
+        new_node = next_neighbor(state, nearest);
       }
 
       memcpy(new_node->coords, target.coords, sizeof(target.coords));
       kdf_insert(&kdf, new_node);
-      remove_non_boundary(&kdf, new_node);
+      remove_non_boundary(state, &kdf, new_node);
 
       if (kdf.size > max_size) {
         max_size = kdf.size;
       }
 
-      png_byte *pixel = bitmap[new_node->y] + 3*new_node->x;
+      png_byte *pixel = state->bitmap[new_node->y] + 3*new_node->x;
       color_unpack(pixel, color);
     }
   }
 
-  printf("%s%.2f%%\t| boundary size: 0\t| max boundary size: %zu\n",
-         clear_line, 100.0, max_size);
+  print_progress("%.2f%%\t| boundary size: 0\t| max boundary size: %zu\n",
+                 100.0, max_size);
 
   kdf_destroy(&kdf);
+  free(nodes);
 }
 
 static void
-write_png(const char *filename, png_byte **bitmap)
+write_png(const state_t *state, const char *filename)
 {
   FILE *file = fopen(filename, "wb");
   if (!file) {
@@ -272,44 +522,12 @@ write_png(const char *filename, png_byte **bitmap)
   }
 
   png_init_io(png_ptr, file);
-  png_set_IHDR(png_ptr, info_ptr, WIDTH, HEIGHT, 8,
+  png_set_IHDR(png_ptr, info_ptr, state->width, state->height, 8,
                PNG_COLOR_TYPE_RGB, PNG_INTERLACE_ADAM7,
                PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
   png_set_sRGB_gAMA_and_cHRM(png_ptr, info_ptr, PNG_sRGB_INTENT_ABSOLUTE);
-  png_set_rows(png_ptr, info_ptr, bitmap);
+  png_set_rows(png_ptr, info_ptr, state->bitmap);
   png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, NULL);
   png_destroy_write_struct(&png_ptr, &info_ptr);
   fclose(file);
 }
-
-int
-main(void)
-{
-  printf("Generating a %ux%u image (%u pixels)\n", WIDTH, HEIGHT, SIZE);
-
-  // For consistent images
-  srand(0);
-
-  // Generate all the colors
-  uint32_t *colors = create_colors();
-  // Make a pool of potential k-d nodes
-  kd_node_t *nodes = create_kd_nodes();
-
-  // Allocate the bitmap
-  png_byte **bitmap = create_bitmap();
-
-  // Generate the image
-  generate_image(colors, nodes, WIDTH/2, HEIGHT/2, bitmap);
-
-  // Write out the image
-  write_png("kd-forest.png", bitmap);
-
-  // Clean up
-  for (unsigned int i = 0; i < HEIGHT; ++i) {
-    free(bitmap[i]);
-  }
-  free(bitmap);
-  free(nodes);
-  free(colors);
-  return 0;
-}