#include <assert.h> // assert
#include <stdlib.h> // calloc, free

/** A prefix tree
 * An efficient key-value storage with string keys
 *
 * There exist two type of nodes:
 * 1. node->symbol == 0 -> node where a key ended, has node->value
 * 2. node->symbol != 0 -> node representing a character in a key, has node->child
 *
 * Each node potentially has a node->next, which represents a node in the same
 * key position but with a different symbol
 */
struct Trie {
    char symbol;
    struct Trie *next;
    union {
        void *value;
        struct Trie *child;
    };
};

/** Create a new Trie struct
 * This function assumes NULL == 0, but that's probably true in every machine
 * I'm not sure it is true in C++ but eh
 */
struct Trie *trie_new() {
    return calloc(1, sizeof(struct Trie));
}

/** Walk through a Trie and call a function on every value
 * This *must* be used to free values
 *
 * Example:
 *   struct Trie *trie = trie_new();
 *
 *   int *x = malloc(sizeof *x);
 *   *x = 1;
 *   trie_set(trie, "a", x);
 *
 *   int *y = malloc(sizeof *y);
 *   *y = 2;
 *   trie_set(trie, "b", y);
 *
 *   assert(trie_get(trie, "a") == x);
 *   assert(trie_get(trie, "b") == y);
 *
 *   trie_foreach_value(struct Trie *root, free);
 *   trie_free(trie);
 */
void trie_foreach_value(struct Trie *root, void (*f)(void*)) {
    if (root == NULL) return;
    trie_foreach_value(root->next, f);
    if (root->symbol != 0) trie_foreach_value(root->child, f);
    else f(root->value);
}

/** Free a Trie
 * This function does not try to free values, for that utilize
 * trie_foreach_value with a custom free function
 */
void trie_free(struct Trie *root) {
    if (root == NULL) return;
    trie_free(root->next);
    if (root->symbol != 0) trie_free(root->child);
    free(root);
}

/** Set a key to a value in a Trie
 * This function requires that `value != NULL` so that missing keys can return
 * NULL in `trie_get`
 *
 * It also requires that the key is non-empty (i.e. strlen(key) != 0) for ease
 * of implementation
 */
void trie_set(struct Trie *root, const char *key, void *value) {
    assert(value != NULL);
    assert(*key != 0);

    struct Trie *conductor = root;

    /* For each character of the key */
    for (; *key != 0; ++key) {
        /* Find a suitable node at this level */
        for (;; conductor = conductor->next) {
            /* If there already is a node at this level with the key's current
             * character, just use that and move on to the next level */
            if (conductor->symbol == *key) {
                conductor = conductor->child;
                break;
            }

            /* If we reach this case, it means there are no more nodes in the
             * current level */
            if (conductor->next == NULL) {
                /* So we create a new node and set it to be the last node's
                 * successor */
                struct Trie *next = trie_new();
                next->symbol = *key;
                conductor->next = next;
                conductor = next;

                /* Then we move on to the next level as usual */
                struct Trie *child = trie_new();
                conductor->child = child;
                conductor = child;
                break;
            }
        }
    }

    /* When we're here, we're at the last character of the key and so we just
     * set the current node's value */
    conductor->value = value;
}

/** Get a key's value in a Trie
 * This function returns NULL if a key is not present
 */
void *trie_get(struct Trie *root, const char *key) {
    assert(*key != 0);

    if (root == NULL) return NULL;
    struct Trie *conductor = root;

    for (; *key != 0; ++key) {
        for (; conductor->symbol != *key; conductor = conductor->next) {
            if (conductor->next == NULL) return NULL;
        }

        if (conductor->child == NULL) return NULL;
        conductor = conductor->child;
    }

    return conductor->value;
}