/* Released Into the Public Domain By: Roy Keene, 2012 */
 
 #include <assert.h>
 #include <stdlib.h>
 #include <limits.h>
 
 #ifndef HAT_DEFAULT_PAGESIZE
 #define HAT_DEFAULT_PAGESIZE 4096
 #endif
 
 struct hat {
     unsigned int baselen;
     unsigned int numleaves;
     void ***leaves;
 };
 
 /* Compute the leaf index a given index is under */
 static unsigned int hat_compute_leaf(unsigned int baselen, unsigned int idx) {
     unsigned int leaf;
     unsigned int base;
 
     /* This computes the leaf node using integer addition and multiplication */
     for (base = leaf = 0; base <= idx; leaf++) {
         base += baselen;
         baselen *= 2;
     }
     leaf--;
 
     return(leaf);
 }
 
 /* Compute the number of items in a given leaf index */
 static unsigned int hat_compute_numitems_leaf(unsigned int baselen, unsigned int leaf) {
     unsigned int numitemsthisleaf;
 
     if (leaf == 0) {
         numitemsthisleaf = baselen;
     } else {
         numitemsthisleaf = (baselen * (2 << (leaf - 1)));
     }
 
     return(numitemsthisleaf);
 }
 
 /* Returns the leaf index and item index for a given linear index */
 static void hat_get_idx(unsigned int baselen, unsigned int idx, unsigned int *leaf_p, unsigned int *leaf_idx_p) {
     unsigned int leaf, leaf_idx;
 
     assert(leaf_p != NULL);
     assert(leaf_idx_p != NULL);
 
     leaf = hat_compute_leaf(baselen, idx);
 
     if (leaf == 0) {
         leaf_idx = idx;
     } else {
         leaf_idx = idx - (baselen * ((2 << (leaf - 1)) - 1));
     }
 
     assert(leaf <= idx);
     assert(leaf_idx <= idx);
 
     *leaf_p = leaf;
     *leaf_idx_p = leaf_idx;
 
     return;
 }
 
 /* Grows the structure such that it will be have "newnumleaves" number of leaf indexes */
 static void hat_alloc_leaves(struct hat *hat, unsigned int newnumleaves) {
     unsigned int oldnumleaves, numitemsthisleaf, idx;
     size_t leaves_new_size;
     struct hat tmp_hat;
 
     assert(hat != NULL);
     assert(newnumleaves >= hat->numleaves);
 
     oldnumleaves = hat->numleaves;
     leaves_new_size = sizeof(*(hat->leaves)) * newnumleaves;
 
     if (hat->leaves == NULL) {
         tmp_hat.leaves = malloc(leaves_new_size);
     } else {
         tmp_hat.leaves = realloc(hat->leaves, leaves_new_size);
     }
 
     if (tmp_hat.leaves == NULL) {
         return;
     }
 
     hat->leaves = tmp_hat.leaves;
 
     for (idx = oldnumleaves; idx < newnumleaves; idx++) {
         tmp_hat.leaves[idx] = NULL;
     }
 
     for (idx = oldnumleaves; idx < newnumleaves; idx++) {
         numitemsthisleaf = hat_compute_numitems_leaf(hat->baselen, idx);
 
         tmp_hat.leaves[idx] = malloc(sizeof(**(tmp_hat.leaves)) * numitemsthisleaf);
 
         if (tmp_hat.leaves[idx] == NULL) {
             for (idx = oldnumleaves; idx < newnumleaves; idx++) {
                 if (tmp_hat.leaves[idx] != NULL) {
                     free(tmp_hat.leaves[idx]);
                 }
             }
 
             return;
         }
     }
 
     hat->numleaves = newnumleaves;
 
     return;
 }
 
 /* Determine system page size */
 #ifdef HAT_FEATURE_USE_GETPAGESIZE
 #include <unistd.h>
 #endif
 static unsigned long hat_getpagesize(void) {
 #ifdef HAT_FEATURE_USE_GETPAGESIZE
     static long pagesize = -1;
 
     if (pagesize != -1) {
         return(pagesize);
     }
 
     pagesize = sysconf(_SC_PAGESIZE);
 
     if (pagesize > 0) {
         return(pagesize);
     }
 #endif
 
     return(HAT_DEFAULT_PAGESIZE);
 }
 
 /* Determine optimal baselen size for a given page size */
 static unsigned int hat_get_baselen(void) {
     unsigned long pagesize;
     unsigned int baselen;
 
     /* Get OS page size */
     pagesize = hat_getpagesize();
 
     /* Compute the number of entries that will fit in a page */
     baselen = pagesize / sizeof(void *);
 
     /* Leave some room in the page for "buddy info" */
     baselen -= 2;
 
     return(baselen);
 }
 
 static unsigned int hat_get_baselen_from_max(unsigned int maxidx) {
     unsigned long best_baselen_size, best_baselen;
     unsigned long pagesize, baselen_size;
     unsigned int baselen;
     unsigned int leaf;
 
     /* Get OS page size */
     pagesize = hat_getpagesize();
     pagesize -= sizeof(void *) * 2;
 
     /*
      * Really we are finding the start of the block we hope not to have to
      * allocate, rather than the end of the current block
      */
     maxidx++;
 
     best_baselen_size = 0;
     for (leaf = 2; leaf < 64; leaf++) {
         baselen = (maxidx / ((2 << (leaf - 1)) - 1));
         baselen_size = baselen * sizeof(void *);
         baselen_size = baselen_size % pagesize;
 
         /* Don't create too large of a baselen to inhibit growth rate (and waste) for the average case */
         if (baselen > 16384) {
             continue;
         }
 
         if (baselen < 2) {
             break;
         }
 
         if (baselen_size == 0) {
             baselen_size = pagesize;
         }
 
         if (baselen_size >= best_baselen_size) {
             best_baselen_size = baselen_size;
             best_baselen = baselen;
         }
     }
 
     if (best_baselen_size != 0) {
         return(best_baselen + 1);
     }
 
     return(hat_get_baselen());
 }
 
 /* Initialize a "hat" structure */
 int hat_init(struct hat *hat, unsigned int baselen, unsigned int numleaves, unsigned int maxidx) {
     assert(hat != NULL);
 
     if (baselen == 0) {
         if (maxidx == 0) {
             baselen = hat_get_baselen();
         } else {
             baselen = hat_get_baselen_from_max(maxidx);
         }
     }
 
     if (numleaves == 0) {
         numleaves = 1;
     }
 
     hat->baselen = baselen;
     hat->numleaves = 0;
     hat->leaves = NULL;
 
     hat_alloc_leaves(hat, numleaves);
 
     return(0);
 }
 
 /* De-initialize and free a "hat" structure */
 void hat_free(struct hat *hat) {
     unsigned int leaf;
 
     assert(hat != NULL);
 
     if (hat->leaves == NULL) {
         assert(hat->numleaves == 0);
 
         return;
     }
 
     for (leaf = 0; leaf < hat->numleaves; leaf++) {
         assert(hat->leaves[leaf] != NULL);
 
         free(hat->leaves[leaf]);
     }
 
     free(hat->leaves);
 
     return;
 }
 
 /* Set an item */
 int hat_set(struct hat *hat, unsigned int idx, void *value) {
     unsigned int leaf, leaf_idx;
 
     assert(hat != NULL);
 
     hat_get_idx(hat->baselen, idx, &leaf, &leaf_idx);
 
     if (leaf >= hat->numleaves) {
         hat_alloc_leaves(hat, leaf + 1);
 
         if (leaf >= hat->numleaves) {
             return(-1);
         }
     }
 
     hat->leaves[leaf][leaf_idx] = value;
 
     return(0);
 }
 
 /* Get an item */
 int hat_get(struct hat *hat, unsigned int idx, void **value_p) {
     unsigned int leaf, leaf_idx;
 
     assert(hat != NULL);
     assert(value_p != NULL);
 
     hat_get_idx(hat->baselen, idx, &leaf, &leaf_idx);
 
     if (leaf >= hat->numleaves) {
         return(-1);
     }
 
     *value_p = hat->leaves[leaf][leaf_idx];
 
     return(0);
 }
 
 #ifdef test_driver
 /* Test and example code on how to use this implementation */
 
 #include <stdio.h>
 
 int main(int argc, char **argv) {
     struct hat hat;
     unsigned int idx;
     char *word;
     int chk_rv;
 
     chk_rv = hat_init(&hat, 1022, 1, 0);
     if (chk_rv != 0) {
         fprintf(stderr, "hat_init() failed\n");
 
         return(1);
     }
 
     /* Test immediate destruction */
     hat_free(&hat);
 
     /* Re-create for further testing */
     hat_init(&hat, 0, 1, 268435455);
 
     /* Test setting 5000 initial values */
     for (idx = 0; idx < 5000; idx++) {
         chk_rv = hat_set(&hat, idx, NULL);
 
         if (chk_rv != 0) {
             fprintf(stderr, "hat_set() failed\n");
 
             return(1);
         }
     }
 
     /* Test setting to specific values and then retreiving them */
     hat_set(&hat, 0, "This");
     hat_set(&hat, 1, "is");
     hat_set(&hat, 2, "a");
     hat_set(&hat, 3, "test");
     hat_set(&hat, 4, "bob");
     hat_set(&hat, 10, "ten");
     hat_set(&hat, 20, "twenty");
     hat_set(&hat, 40, "fourty");
     hat_set(&hat, 80, "eighty");
     hat_set(&hat, 160, "one hundred sixty");
     hat_set(&hat, 320, "three hundred twenty");
     hat_set(&hat, 321, "three hundred twenty-one");
     hat_set(&hat, 640, "six hundred fourty");
     chk_rv = hat_set(&hat, 268435455, "big");
     if (chk_rv != 0) {
         fprintf(stderr, "hat_set(..., 268435455, ...) failed\n");
     }
 
     for (idx = 0; idx <= 641; idx++) {
         if (idx == 5) idx = 10;
         if (idx == 11) idx = 20;
         if (idx == 21) idx = 40;
         if (idx == 41) idx = 80;
         if (idx == 81) idx = 160;
         if (idx == 161) idx = 320;
         if (idx == 322) idx = 640;
         if (idx == 641) idx = 268435455;
 
         chk_rv = hat_get(&hat, idx, (void **) &word);
         if (chk_rv != 0) {
             fprintf(stderr, "hat_get() failed\n");
 
             continue;
         }
 
         printf("hat[%u] = %s\n", idx, word);
     }
 
     /* Test destruction */
     hat_free(&hat);
 
     return(0);
 
     /* NOTREACHED */
     argc = argc;
     argv = argv;
 }
 #endif