1 files changed, 22 insertions, 7 deletions

diff --git a/src/std/hash_map.c b/src/std/hash_map.c
index 1ffca90..febb3fd 100644
--- a/src/std/hash_map.c
+++ b/src/std/hash_map.c
@@ -16,6 +16,7 @@ typedef struct {
 } HashMapBucket;
 
 static void bucket_node_set_next(HashMapBucketNode *self, HashMapBucketNode *next) {
+    assert(next != self);
     am_memcpy(self, &next, sizeof(next));
 }
 
@@ -83,7 +84,7 @@ static CHECK_RESULT int hash_map_bucket_add(HashMap *self, HashMapBucket *bucket
     return 0;
 }
 
-static void hash_map_reorder_buckets(HashMap *self, usize end_index) {
+static void hash_map_reorder_nodes(HashMap *self, usize end_index) {
     HashMapBucket *bucket;
     HashMapBucket *bucket_end;
     usize bucket_size;
@@ -97,29 +98,43 @@ static void hash_map_reorder_buckets(HashMap *self, usize end_index) {
     for(; bucket != bucket_end; ++bucket, ++index) {
         HashMapBucketNode *prev_bucket_node;
         HashMapBucketNode *bucket_node;
+
         prev_bucket_node = NULL;
-        for(bucket_node = bucket->start; bucket_node; bucket_node_get_next(bucket_node)) {
+        bucket_node = bucket->start;
+        while(bucket_node) {
             BufferView bucket_key;
             usize bucket_index;
 
             /* Node between current and previous has been removed, chain previous and current */
             if(prev_bucket_node && !bucket_node_get_next(prev_bucket_node))
                 bucket_node_set_next(prev_bucket_node, bucket_node);
+            
+            if(!prev_bucket_node)
+                bucket->start = bucket_node;
 
             bucket_key = bucket_node_get_key(bucket_node);
             bucket_index = self->hash_func((const u8*)bucket_key.data, bucket_key.size) % bucket_size;
             if(bucket_index != index) {
                 /* Add node to new bucket */
+                HashMapBucketNode *moved_node;
                 HashMapBucket *new_bucket;
                 new_bucket = ((HashMapBucket*)self->buckets.data) + bucket_index;
-                bucket_node_set_next(bucket_node, new_bucket->start);
-                new_bucket->start = bucket_node;
+                moved_node = bucket_node;
+                bucket_node = bucket_node_get_next(bucket_node);
+                bucket_node_set_next(moved_node, new_bucket->start);
+                new_bucket->start = moved_node;
                 /* Mark bucket node to be removed from previous bucket */
-                bucket_node_set_next(prev_bucket_node, NULL);
+                if(prev_bucket_node)
+                    bucket_node_set_next(prev_bucket_node, NULL);
             } else {
                 prev_bucket_node = bucket_node;
+                bucket_node = bucket_node_get_next(bucket_node);
             }
         }
+
+        /* All nodes removed in bucket */
+        if(!prev_bucket_node)
+            bucket->start = NULL;
     }
 }
 
@@ -129,8 +144,8 @@ static CHECK_RESULT int hash_map_increase_buckets(HashMap *self) {
     prev_size = buffer_get_size(&self->buckets, HashMapBucket);
     bytes_to_append = prev_size * 1.5;
     return_if_error(buffer_append(&self->buckets, NULL, bytes_to_append));
-    am_memset(self->buckets.data + prev_size, 0, prev_size + bytes_to_append);
-    hash_map_reorder_buckets(self, prev_size);
+    am_memset(((HashMapBucket*)self->buckets.data) + prev_size, 0, bytes_to_append);
+    hash_map_reorder_nodes(self, prev_size);
     return 0;
 }