1 files changed, 39 insertions, 19 deletions

diff --git a/src/std/hash_map.c b/src/std/hash_map.c
index 61030da..ab580ba 100644
--- a/src/std/hash_map.c
+++ b/src/std/hash_map.c
@@ -9,6 +9,13 @@ Basic hash map implementation. TODO: Improve performance
 
 #define HASH_MAP_INITIAL_SIZE 8
 
+/* Structure:
+HashMapBucketNode *next;
+usize hash;
+u32 key_size;
+u8[..] key;
+u8[..] value;
+*/
 typedef struct HashMapBucketNode HashMapBucketNode;
 
 typedef struct {
@@ -22,37 +29,47 @@ static void bucket_node_set_next(HashMapBucketNode *self, HashMapBucketNode *nex
 
 static HashMapBucketNode* bucket_node_get_next(HashMapBucketNode *self) {
     HashMapBucketNode *next;
-    am_memcpy(&next, self, sizeof(HashMapBucketNode*));
+    am_memcpy(&next, self, sizeof(next));
     return next;
 }
 
+static void bucket_node_set_hash(HashMapBucketNode *self, usize hash) {
+    am_memcpy((char*)self + sizeof(HashMapBucketNode*), &hash, sizeof(hash));
+}
+
+static usize bucket_node_get_hash(HashMapBucketNode *self) {
+    usize hash;
+    am_memcpy(&hash, (char*)self + sizeof(HashMapBucketNode*), sizeof(hash));
+    return hash;
+}
+
 static void bucket_node_set_key(HashMapBucketNode *self, BufferView key) {
     u32 key_size;
     key_size = (u32)key.size;
-    am_memcpy((char*)self + sizeof(HashMapBucketNode*), &key_size, sizeof(u32));
-    am_memcpy((char*)self + sizeof(HashMapBucketNode*) + sizeof(u32), key.data, key_size);
+    am_memcpy((char*)self + sizeof(HashMapBucketNode*) + sizeof(usize), &key_size, sizeof(u32));
+    am_memcpy((char*)self + sizeof(HashMapBucketNode*) + sizeof(usize) + sizeof(u32), key.data, key_size);
 }
 
 static BufferView bucket_node_get_key(HashMapBucketNode *self) {
     BufferView key;
     u32 key_size;
-    am_memcpy(&key_size, (char*)self + sizeof(HashMapBucketNode*), sizeof(u32));
+    am_memcpy(&key_size, (char*)self + sizeof(HashMapBucketNode*) + sizeof(usize), sizeof(u32));
     key.size = key_size;
-    key.data = (char*)self + sizeof(HashMapBucketNode*) + sizeof(u32);
+    key.data = (char*)self + sizeof(HashMapBucketNode*) + sizeof(usize) + sizeof(u32);
     return key;
 }
 
 static void bucket_node_set_value(HashMapBucketNode *self, void *value, usize value_type_size) {
     u32 key_size;
-    am_memcpy(&key_size, (char*)self + sizeof(HashMapBucketNode*), sizeof(key_size));
-    am_memcpy((char*)self + sizeof(HashMapBucketNode*) + sizeof(u32) + key_size, value, value_type_size);
+    am_memcpy(&key_size, (char*)self + sizeof(HashMapBucketNode*) + sizeof(usize), sizeof(key_size));
+    am_memcpy((char*)self + sizeof(HashMapBucketNode*) + sizeof(usize) + sizeof(u32) + key_size, value, value_type_size);
 }
 
 static void* bucket_node_get_value(HashMapBucketNode *self) {
     u32 key_size;
     void *value;
-    am_memcpy(&key_size, (char*)self + sizeof(HashMapBucketNode*), sizeof(key_size));
-    value = (char*)self + sizeof(HashMapBucketNode*) + sizeof(u32) + key_size;
+    am_memcpy(&key_size, (char*)self + sizeof(HashMapBucketNode*) + sizeof(usize), sizeof(key_size));
+    value = (char*)self + sizeof(HashMapBucketNode*) + sizeof(usize) + sizeof(u32) + key_size;
     return value;
 }
 
@@ -72,12 +89,13 @@ int hash_map_init(HashMap *self, ScopedAllocator *allocator, usize value_type_si
     return 0;
 }
 
-static CHECK_RESULT int hash_map_bucket_add(HashMap *self, HashMapBucket *bucket, BufferView key, void *value) {
+static CHECK_RESULT int hash_map_bucket_add(HashMap *self, HashMapBucket *bucket, BufferView key, void *value, usize hash) {
     HashMapBucketNode *new_bucket_node;
     return_if_error(scoped_allocator_alloc(self->allocator, 
-                                           sizeof(HashMapBucketNode*) + sizeof(u32) + key.size + self->value_type_size,
+                                           sizeof(HashMapBucketNode*) + sizeof(hash) + sizeof(u32) + key.size + self->value_type_size,
                                            (void**)&new_bucket_node));
     bucket_node_set_next(new_bucket_node, bucket->start);
+    bucket_node_set_hash(new_bucket_node, hash);
     bucket_node_set_key(new_bucket_node, key);
     bucket_node_set_value(new_bucket_node, value, self->value_type_size);
     bucket->start = new_bucket_node;
@@ -102,11 +120,9 @@ static void hash_map_reorder_nodes(HashMap *self, usize end_index) {
         prev_bucket_node = NULL;
         bucket_node = bucket->start;
         while(bucket_node) {
-            BufferView bucket_key;
             usize bucket_index;
-
-            bucket_key = bucket_node_get_key(bucket_node);
-            bucket_index = self->hash_func((const u8*)bucket_key.data, bucket_key.size) % bucket_size;
+            bucket_index = bucket_node_get_hash(bucket_node) % bucket_size;
+            
             if(bucket_index != index) {
                 /* Add node to new bucket */
                 HashMapBucketNode *moved_node;
@@ -145,6 +161,7 @@ static CHECK_RESULT int hash_map_increase_buckets(HashMap *self) {
 }
 
 int hash_map_insert(HashMap *self, BufferView key, void *value) {
+    usize hash;
     usize bucket_index;
     usize bucket_size;
     HashMapBucket *bucket;
@@ -155,9 +172,10 @@ int hash_map_insert(HashMap *self, BufferView key, void *value) {
         bucket_size = buffer_get_size(&self->buckets, HashMapBucket);
     }
 
-    bucket_index = self->hash_func((const u8*)key.data, key.size) % bucket_size;
+    hash = self->hash_func((const u8*)key.data, key.size);
+    bucket_index = hash % bucket_size;
     bucket = ((HashMapBucket*)self->buckets.data) + bucket_index;
-    return_if_error(hash_map_bucket_add(self, bucket, key, value));
+    return_if_error(hash_map_bucket_add(self, bucket, key, value, hash));
     ++self->num_elements;
     return 0;
 }
@@ -165,17 +183,19 @@ int hash_map_insert(HashMap *self, BufferView key, void *value) {
 bool hash_map_get(HashMap *self, BufferView key, void *value) {
     usize bucket_size;
     usize bucket_index;
+    usize hash;
     HashMapBucket *bucket;
     HashMapBucketNode *bucket_node;
 
     bucket_size = buffer_get_size(&self->buckets, HashMapBucket);
-    bucket_index = self->hash_func((const u8*)key.data, key.size) % bucket_size;
+    hash = self->hash_func((const u8*)key.data, key.size);
+    bucket_index = hash % bucket_size;
 
     bucket = ((HashMapBucket*)self->buckets.data) + bucket_index;
     for(bucket_node = bucket->start; bucket_node; bucket_node = bucket_node_get_next(bucket_node)) {
         BufferView bucket_key;
         bucket_key = bucket_node_get_key(bucket_node);
-        if(self->compare_func(&key, &bucket_key) == 0) {
+        if(hash == bucket_node_get_hash(bucket_node) && self->compare_func(&key, &bucket_key) == 0) {
             am_memcpy(value, bucket_node_get_value(bucket_node), self->value_type_size);
             return bool_true;
         }