Add some wrappers around SPLAY_* for the common map-from-string-to-X case.

It will probably be less blindingly fast than using SPLAY_* directly, but
only slightly so.


svn:r1306

4 files changed, 322 insertions, 2 deletions

diff --git a/src/common/test.h b/src/common/test.h
index d30a458e0..d2b449db4 100644
--- a/src/common/test.h
+++ b/src/common/test.h
@@ -72,7 +72,7 @@ extern int have_failed;
 
 #define test_streq(expr1, expr2)                                \
   STMT_BEGIN                                                    \
-    char *v1=(expr1), *v2=(expr2);                              \
+    const char *v1=(expr1), *v2=(expr2);                        \
     if(!strcmp(v1,v2)) { printf("."); } else {                  \
     have_failed = 1;                                            \
     printf("\nFile %s: line %d (%s): Assertion failed: (%s==%s)\n"\
@@ -87,7 +87,7 @@ extern int have_failed;
 
 #define test_strneq(expr1, expr2)                               \
   STMT_BEGIN                                                    \
-    char *v1=(expr1), *v2=(expr2);                              \
+    const char *v1=(expr1), *v2=(expr2);                        \
     if(strcmp(v1,v2)) { printf("."); } else {                   \
     have_failed = 1;                                            \
     printf("\nFile %s: line %d (%s): Assertion failed: (%s!=%s)\n"\
diff --git a/src/common/util.c b/src/common/util.c
index b95b23d8f..101ac0b1d 100644
--- a/src/common/util.c
+++ b/src/common/util.c
@@ -3,6 +3,7 @@
 /* $Id$ */
 
 #include "../or/or.h"
+#include "../or/tree.h"
 
 #ifdef HAVE_UNAME
 #include <sys/utsname.h>
@@ -156,6 +157,229 @@ void *smartlist_choose(smartlist_t *sl) {
 }
 
 /*
+ * Splay-tree implementation of string-to-void* map
+ */
+struct strmap_entry_t {
+  SPLAY_ENTRY(strmap_entry_t) node;
+  char *key;
+  void *val;
+};
+
+struct strmap_t {
+  SPLAY_HEAD(strmap_tree, strmap_entry_t) head;
+};
+
+static int compare_strmap_entries(struct strmap_entry_t *a,
+				 struct strmap_entry_t *b)
+{
+  return strcmp(a->key, b->key);
+}
+
+SPLAY_PROTOTYPE(strmap_tree, strmap_entry_t, node, compare_strmap_entries);
+SPLAY_GENERATE(strmap_tree, strmap_entry_t, node, compare_strmap_entries);
+
+/* Create a new empty map from strings to void*'s.
+ */
+strmap_t* strmap_new(void)
+{
+  strmap_t *result;
+  result = tor_malloc(sizeof(strmap_t));
+  SPLAY_INIT(&result->head);
+  return result;
+}
+
+/* Set the current value for <key> with <val>.  Returns the previous
+ * value for <key> if one was set, or NULL if one was not.
+ *
+ * This function makes a copy of 'key' if necessary, but not of 'val'.
+ */
+void* strmap_set(strmap_t *map, const char *key, void *val)
+{
+  strmap_entry_t *resolve;
+  strmap_entry_t search;
+  void *oldval;
+  assert(map && key && val);
+  search.key = (char*)key;
+  resolve = SPLAY_FIND(strmap_tree, &map->head, &search);
+  if (resolve) {
+    oldval = resolve->val;
+    resolve->val = val;
+    return oldval;
+  } else {
+    resolve = tor_malloc_zero(sizeof(strmap_entry_t));
+    resolve->key = tor_strdup(key);
+    resolve->val = val;
+    SPLAY_INSERT(strmap_tree, &map->head, resolve);
+    return NULL;
+  }
+}
+
+/* Return the current value associated with <key>, or NULL if no
+ * value is set.
+ */
+void* strmap_get(strmap_t *map, const char *key)
+{
+  strmap_entry_t *resolve;
+  strmap_entry_t search;
+  assert(map && key);
+  search.key = (char*)key;
+  resolve = SPLAY_FIND(strmap_tree, &map->head, &search);
+  if (resolve) {
+    return resolve->val;
+  } else {
+    return NULL;
+  }
+}
+
+/* Remove the value currently associated with <key> from the map.
+ * Return the value if one was set, or NULL if there was no entry for
+ * <key>.
+ *
+ * Note: you must free any storage associated with the returned value.
+ */
+void* strmap_remove(strmap_t *map, const char *key)
+{
+  strmap_entry_t *resolve;
+  strmap_entry_t search;
+  void *oldval;
+  assert(map && key);
+  search.key = (char*)key;
+  resolve = SPLAY_FIND(strmap_tree, &map->head, &search);
+  if (resolve) {
+    oldval = resolve->val;
+    SPLAY_REMOVE(strmap_tree, &map->head, resolve);
+    tor_free(resolve->key);
+    tor_free(resolve);
+    return oldval;
+  } else {
+    return NULL;
+  }
+}
+
+/* Invoke fn() on every entry of the map, in order.  For every entry,
+ * fn() is invoked with that entry's key, that entry's value, and the
+ * value of <data> supplied to strmap_foreach.  fn() must return a new
+ * (possibly unmodified) value for each entry: if fn() returns NULL, the
+ * entry is removed.
+ *
+ * Example:
+ *   static void* upcase_and_remove_empty_vals(const char *key, void *val,
+ *                                             void* data) {
+ *     char *cp = (char*)val;
+ *     if (!*cp) {  // val is an empty string.
+ *       free(val);
+ *       return NULL;
+ *     } else {
+ *       for (; *cp; cp++)
+ *         *cp = toupper(*cp);
+ *       }
+ *       return val;
+ *     }
+ *   }
+ *
+ *   ...
+ *
+ *   strmap_foreach(map, upcase_and_remove_empty_vals, NULL);
+ */
+void strmap_foreach(strmap_t *map,
+		    void* (*fn)(const char *key, void *val, void *data),
+		    void *data)
+{
+  strmap_entry_t *ptr, *next;
+  assert(map && fn);
+  for (ptr = SPLAY_MIN(strmap_tree, &map->head); ptr != NULL; ptr = next) {
+    /* This remove-in-place usage is specifically blessed in tree(3). */
+    next = SPLAY_NEXT(strmap_tree, &map->head, ptr);
+    ptr->val = fn(ptr->key, ptr->val, data);
+    if (!ptr->val) {
+      SPLAY_REMOVE(strmap_tree, &map->head, ptr);
+      tor_free(ptr->key);
+      tor_free(ptr);
+    }
+  }
+}
+
+/* return an 'iterator' pointer to the front of a map.
+ *
+ * Iterator example:
+ *
+ * // uppercase values in "map", removing empty values.
+ *
+ * strmap_iterator_t *iter;
+ * const char *key;
+ * void *val;
+ * char *cp;
+ *
+ * for (iter = strmap_iter_init(map); !strmap_iter_done(iter); ) {
+ *    strmap_iter_get(iter, &key, &val);
+ *    cp = (char*)val;
+ *    if (!*cp) {
+ *       iter = strmap_iter_next_rmv(iter);
+ *       free(val);
+ *    } else {
+ *       for(;*cp;cp++) *cp = toupper(*cp);
+ *       iter = strmap_iter_next(iter);
+ *    }
+ * }
+ *
+ */
+strmap_iter_t *strmap_iter_init(strmap_t *map)
+{
+  assert(map);
+  return SPLAY_MIN(strmap_tree, &map->head);
+}
+/* Advance the iterator 'iter' for map a single step to the next entry.
+ */
+strmap_iter_t *strmap_iter_next(strmap_t *map, strmap_iter_t *iter)
+{
+  assert(map && iter);
+  return SPLAY_NEXT(strmap_tree, &map->head, iter);
+}
+/* Advance the iterator 'iter' a single step to the next entry, removing
+ * the current entry.
+ */
+strmap_iter_t *strmap_iter_next_rmv(strmap_t *map, strmap_iter_t *iter)
+{
+  strmap_iter_t *next;
+  assert(map && iter);
+  next = SPLAY_NEXT(strmap_tree, &map->head, iter);
+  SPLAY_REMOVE(strmap_tree, &map->head, iter);
+  tor_free(iter->key);
+  tor_free(iter);
+  return next;
+}
+/* Set *keyp and *valp to the current entry pointed to by iter.
+ */
+void strmap_iter_get(strmap_iter_t *iter, const char **keyp, void **valp)
+{
+  assert(iter && keyp && valp);
+  *keyp = iter->key;
+  *valp = iter->val;
+}
+/* Return true iff iter has advanced past the last entry of map.
+ */
+int strmap_iter_done(strmap_iter_t *iter)
+{
+  return iter == NULL;
+}
+/* Remove all entries from <map>, and deallocate storage for those entries.
+ * If free_val is provided, it is invoked on every value in <map>.
+ */
+void strmap_free(strmap_t *map, void (*free_val)(void*))
+{
+  strmap_entry_t *ent, *next;
+  for (ent = SPLAY_MIN(strmap_tree, &map->head); ent != NULL; ent = next) {
+    next = SPLAY_NEXT(strmap_tree, &map->head, ent);
+    SPLAY_REMOVE(strmap_tree, &map->head, ent);
+    tor_free(ent->key);
+    if (free_val)
+      tor_free(ent->val);
+  }
+  assert(SPLAY_EMPTY(&map->head));
+  tor_free(map);
+}
+
+/*
  *    String manipulation
  */
 
diff --git a/src/common/util.h b/src/common/util.h
index f9b2b90eb..5bdcaed96 100644
--- a/src/common/util.h
+++ b/src/common/util.h
@@ -58,6 +58,26 @@ void smartlist_intersect(smartlist_t *sl1, smartlist_t *sl2);
 void smartlist_subtract(smartlist_t *sl1, smartlist_t *sl2);
 void *smartlist_choose(smartlist_t *sl);
 
+/* Map from const char * to void*. Implemented with a splay tree. */
+typedef struct strmap_t strmap_t;
+typedef struct strmap_entry_t strmap_entry_t;
+typedef struct strmap_entry_t strmap_iter_t;
+strmap_t* strmap_new(void);
+void* strmap_set(strmap_t *map, const char *key, void *val);
+void* strmap_get(strmap_t *map, const char *key);
+void* strmap_remove(strmap_t *map, const char *key);
+void strmap_foreach(strmap_t *map,
+		    void* (*fn)(const char *key, void *val, void *data),
+		    void *data);
+void strmap_free(strmap_t *map, void (*free_val)(void*));
+
+strmap_iter_t *strmap_iter_init(strmap_t *map);
+strmap_iter_t *strmap_iter_next(strmap_t *map, strmap_iter_t *iter);
+strmap_iter_t *strmap_iter_next_rmv(strmap_t *map, strmap_iter_t *iter);
+void strmap_iter_get(strmap_iter_t *iter, const char **keyp, void **valp);
+
+int strmap_iter_done(strmap_iter_t *iter);
+
 const char *eat_whitespace(const char *s);
 const char *eat_whitespace_no_nl(const char *s);
 const char *find_whitespace(const char *s);
diff --git a/src/or/test.c b/src/or/test.c
index a3ab992ae..b4810d2e4 100644
--- a/src/or/test.c
+++ b/src/or/test.c
@@ -466,6 +466,81 @@ test_util() {
   test_eq((time_t) 1076393695UL, tor_timegm(&a_time));
 }
 
+static void* _squareAndRemoveK4(const char *key, void*val, void *data)
+{
+  int *ip = (int*)data;
+  int v;
+  if (strcmp(key,"K4") == 0) {
+    ++(*ip);
+    return NULL;
+  }
+  v = (int)val;
+  return (void*)(v*v);
+}
+
+void test_strmap() {
+  strmap_t *map;
+  strmap_iter_t *iter;
+  const char *k;
+  void *v;
+  int count;
+
+  map = strmap_new();
+  v = strmap_set(map, "K1", (void*)99);
+  test_eq(v, NULL);
+  v = strmap_set(map, "K2", (void*)101);
+  test_eq(v, NULL);
+  v = strmap_set(map, "K1", (void*)100);
+  test_eq(v, (void*)99);
+  test_eq(strmap_get(map,"K1"), (void*)100);
+  test_eq(strmap_get(map,"K2"), (void*)101);
+  test_eq(strmap_get(map,"K-not-there"), NULL);
+
+  v = strmap_remove(map,"K2");
+  test_eq(v, (void*)101);
+  test_eq(strmap_get(map,"K2"), NULL);
+  test_eq(strmap_remove(map,"K2"), NULL);
+
+  strmap_set(map, "K2", (void*)101);
+  strmap_set(map, "K3", (void*)102);
+  strmap_set(map, "K4", (void*)103);
+  strmap_set(map, "K5", (void*)104);
+  strmap_set(map, "K6", (void*)105);
+
+  count = 0;
+  strmap_foreach(map, _squareAndRemoveK4, &count);
+  test_eq(count, 1);
+  test_eq(strmap_get(map, "K4"), NULL);
+  test_eq(strmap_get(map, "K1"), (void*)10000);
+  test_eq(strmap_get(map, "K6"), (void*)11025);
+
+  iter = strmap_iter_init(map);
+  strmap_iter_get(iter,&k,&v);
+  test_streq(k, "K1");
+  test_eq(v, (void*)10000);
+  iter = strmap_iter_next(map,iter);
+  strmap_iter_get(iter,&k,&v);
+  test_streq(k, "K2");
+  test_eq(v, (void*)10201);
+  iter = strmap_iter_next_rmv(map,iter);
+  strmap_iter_get(iter,&k,&v);
+  test_streq(k, "K3");
+  test_eq(v, (void*)10404);
+  iter = strmap_iter_next(map,iter); /* K5 */
+  test_assert(!strmap_iter_done(iter));
+  iter = strmap_iter_next(map,iter); /* K6 */
+  test_assert(!strmap_iter_done(iter));
+  iter = strmap_iter_next(map,iter); /* done */
+  test_assert(strmap_iter_done(iter));
+
+  /* Make sure we removed K2, but not the others. */
+  test_eq(strmap_get(map, "K2"), NULL);
+  test_eq(strmap_get(map, "K5"), (void*)10816);
+
+  /* Clean up after ourselves. */
+  strmap_free(map, NULL);
+}
+
 void test_onion() {
 #if 0
   char **names;
@@ -711,6 +786,7 @@ main(int c, char**v){
   test_crypto_dh();
   puts("\n========================= Util ============================");
   test_util();
+  test_strmap();
   puts("\n========================= Onion Skins =====================");
   test_onion();
   test_onion_handshake();