Add get_or_create method for hash map, more efficient than checking and then creating

3 files changed, 52 insertions, 5 deletions

diff --git a/README.md b/README.md
index c9c3a82..9f9d49f 100644
--- a/README.md
+++ b/README.md
@@ -2,4 +2,5 @@ A tiny scripting language that is designed to be a replacement for small shell/p
 Written in ANSI C to allow embedding everywhere and a WTFPL license that allows it to be used anywhere without any restrictions.
 # TODO
 * Remove dependency on `gc`. Write our own gc instead.
-* Implement big int, which numbers should automatically switch to when they are too large to fit into double.
\ No newline at end of file
+* Implement big int, which numbers should automatically switch to when they are too large to fit into double.
+* Use a list instead of a hash map for variable lookup, since the list will be small most of the time.
\ No newline at end of file
diff --git a/include/std/hash_map.h b/include/std/hash_map.h
index 986576d..b430d22 100644
--- a/include/std/hash_map.h
+++ b/include/std/hash_map.h
@@ -17,7 +17,9 @@ void tsl_hash_map_init(TslHashMap *self);
 void tsl_hash_map_deinit(TslHashMap *self);
 
 int tsl_hash_map_insert(TslHashMap *self, const TslStringView *key, const void *data, size_t size, TslHashFunc hash_func);
-/* Get a reference to the value by key @key */
+/* Get a reference to the value by key @key, or NULL if it doesn't exist */
 void* tsl_hash_map_get(TslHashMap *self, const TslStringView *key, TslHashFunc hash_func);
+/* Get a reference to the value by key @key, or insert a new value for the key with a size of @size */
+int tsl_hash_map_get_or_create(TslHashMap *self, const TslStringView *key, size_t size, TslHashFunc hash_func, void **output);
 
 #endif /* TSL_HASH_MAP_H */
diff --git a/src/std/hash_map.c b/src/std/hash_map.c
index 80fd059..44d36db 100644
--- a/src/std/hash_map.c
+++ b/src/std/hash_map.c
@@ -45,9 +45,16 @@ void tsl_hash_map_deinit(TslHashMap *self) {
     free(self->buckets_data);
 }
 
-static int tsl_hash_map_append_bucket(TslHashMapNode **head_node, uint64_t hash, const TslStringView *key, size_t size, const void *data) {
+/* TODO: Remove if (data) and if (output) */
+static int tsl_hash_map_append_bucket(TslHashMapNode **head_node, uint64_t hash, const TslStringView *key, size_t size, const void *data, void **output) {
     TslHashMapNode *next_node;
+    /*
+        TODO: Instead of allocating space for the data, allow the user to pass a pointer in the insert
+        method and use that directly.
+    */
     uint8_t *node_data = malloc(sizeof(hash) + sizeof(key->size) + key->size + sizeof(size) + size);
+    if(output)
+        *output = node_data;
     if(!node_data) {
         fprintf(stderr, "Error: hash map append failed. Reason: out of memory\n");
         return 0;
@@ -56,6 +63,8 @@ static int tsl_hash_map_append_bucket(TslHashMapNode **head_node, uint64_t hash,
     next_node = malloc(sizeof(TslHashMapNode));
     if(!next_node) {
         free(node_data);
+        if(output)
+            *output = NULL;
         fprintf(stderr, "Error: hash map append failed. Reason: out of memory\n");
         return 0;
     }
@@ -64,7 +73,8 @@ static int tsl_hash_map_append_bucket(TslHashMapNode **head_node, uint64_t hash,
     memcpy(node_data + sizeof(hash), &key->size, sizeof(key->size));
     memcpy(node_data + sizeof(hash) + sizeof(key->size), key->data, key->size);
     memcpy(node_data + sizeof(hash) + sizeof(key->size) + key->size, &size, sizeof(size));
-    memcpy(node_data + sizeof(hash) + sizeof(key->size) + key->size + sizeof(size), data, size);
+    if(data)
+        memcpy(node_data + sizeof(hash) + sizeof(key->size) + key->size + sizeof(size), data, size);
 
     next_node->data = node_data;
     if(*head_node) {
@@ -78,6 +88,7 @@ static int tsl_hash_map_append_bucket(TslHashMapNode **head_node, uint64_t hash,
 }
 
 static size_t tsl_hash_map_get_index(TslHashMap *self, uint64_t hash) {
+    assert((self->buckets_capacity >> 3) > 0);
     /* >>3 = 8 = sizeof(TslHashMapNode*) */
     return hash & ((self->buckets_capacity >> 3) - 1);
 }
@@ -176,7 +187,7 @@ int tsl_hash_map_insert(TslHashMap *self, const TslStringView *key, const void *
 
     index = tsl_hash_map_get_index(self, hash);
     bucket = (TslHashMapNode**)self->buckets_data + index;
-    return tsl_hash_map_append_bucket(bucket, hash, key, size, data);
+    return tsl_hash_map_append_bucket(bucket, hash, key, size, data, NULL);
 }
 
 void* tsl_hash_map_get(TslHashMap *self, const TslStringView *key, TslHashFunc hash_func) {
@@ -202,3 +213,36 @@ void* tsl_hash_map_get(TslHashMap *self, const TslStringView *key, TslHashFunc h
 
     return NULL;
 }
+
+int tsl_hash_map_get_or_create(TslHashMap *self, const TslStringView *key, size_t size, TslHashFunc hash_func, void **output) {
+    uint64_t hash;
+    size_t index;
+    TslHashMapNode **bucket;
+    TslHashMapNode *node;
+    assert(size > 0);
+    if(self->buckets_capacity == 0) {
+        if(!tsl_hash_map_ensure_bucket_capacity_for_one_new_item(self)) {
+            *output = NULL;
+            return 0;
+        }
+    }
+
+    hash = hash_func(key->data, key->size);
+    index = tsl_hash_map_get_index(self, hash);
+    bucket = (TslHashMapNode**)self->buckets_data + index;
+    node = *bucket;
+    while(node) {
+        uint64_t node_hash;
+        TslStringView node_key;
+        size_t node_size;
+        uint8_t *node_data;
+        hash_map_node_get(node, &node_hash, &node_key, &node_size, &node_data);
+        if(hash == node_hash && key->size == node_key.size && memcmp(key->data, node_key.data, node_key.size) == 0) {
+            *output = node_data;
+            return 1;
+        }
+        node = node->next;
+    }
+
+    return tsl_hash_map_append_bucket(bucket, hash, key, size, NULL, output);
+}