| ... | ... | @@ -118,7 +118,7 @@ pub fn HashMap(comptime K: type, comptime V: type, comptime hash: fn (key: K) u3 |
| 118 | 118 | }; |
| 119 | 119 | } |
| 120 | 120 | self.incrementModificationCount(); |
| 121 | | try self.ensureCapacity(); |
| 121 | try self.autoCapacity(); |
| 122 | 122 | const put_result = self.internalPut(key); |
| 123 | 123 | assert(put_result.old_kv == null); |
| 124 | 124 | return GetOrPutResult{ |
| ... | ... | @@ -135,15 +135,37 @@ pub fn HashMap(comptime K: type, comptime V: type, comptime hash: fn (key: K) u3 |
| 135 | 135 | return res.kv; |
| 136 | 136 | } |
| 137 | 137 | |
| 138 | | fn ensureCapacity(self: *Self) !void { |
| 139 | | if (self.entries.len == 0) { |
| 140 | | return self.initCapacity(16); |
| 138 | fn optimizedCapacity(expected_count: usize) usize { |
| 139 | // ensure that the hash map will be at most 60% full if |
| 140 | // expected_count items are put into it |
| 141 | var optimized_capacity = expected_count * 5 / 3; |
| 142 | // round capacity to the next power of two |
| 143 | const pow = math.log2_int_ceil(usize, optimized_capacity); |
| 144 | return math.pow(usize, 2, pow); |
| 145 | } |
| 146 | |
| 147 | /// Increases capacity so that the hash map will be at most |
| 148 | /// 60% full when expected_count items are put into it |
| 149 | pub fn ensureCapacity(self: *Self, expected_count: usize) !void { |
| 150 | const optimized_capacity = optimizedCapacity(expected_count); |
| 151 | return self.ensureCapacityExact(optimized_capacity); |
| 152 | } |
| 153 | |
| 154 | /// Sets the capacity to the new capacity if the new |
| 155 | /// capacity is greater than the current capacity. |
| 156 | /// New capacity must be a power of two. |
| 157 | fn ensureCapacityExact(self: *Self, new_capacity: usize) !void { |
| 158 | const is_power_of_two = new_capacity & (new_capacity-1) == 0; |
| 159 | assert(is_power_of_two); |
| 160 | |
| 161 | if (new_capacity <= self.entries.len) { |
| 162 | return; |
| 141 | 163 | } |
| 142 | 164 | |
| 143 | | // if we get too full (60%), double the capacity |
| 144 | | if (self.size * 5 >= self.entries.len * 3) { |
| 145 | | const old_entries = self.entries; |
| 146 | | try self.initCapacity(self.entries.len * 2); |
| 165 | const old_entries = self.entries; |
| 166 | try self.initCapacity(new_capacity); |
| 167 | self.incrementModificationCount(); |
| 168 | if (old_entries.len > 0) { |
| 147 | 169 | // dump all of the old elements into the new table |
| 148 | 170 | for (old_entries) |*old_entry| { |
| 149 | 171 | if (old_entry.used) { |
| ... | ... | @@ -156,8 +178,13 @@ pub fn HashMap(comptime K: type, comptime V: type, comptime hash: fn (key: K) u3 |
| 156 | 178 | |
| 157 | 179 | /// Returns the kv pair that was already there. |
| 158 | 180 | pub fn put(self: *Self, key: K, value: V) !?KV { |
| 181 | try self.autoCapacity(); |
| 182 | return putAssumeCapacity(self, key, value); |
| 183 | } |
| 184 | |
| 185 | pub fn putAssumeCapacity(self: *Self, key: K, value: V) ?KV { |
| 186 | assert(self.count() < self.entries.len); |
| 159 | 187 | self.incrementModificationCount(); |
| 160 | | try self.ensureCapacity(); |
| 161 | 188 | |
| 162 | 189 | const put_result = self.internalPut(key); |
| 163 | 190 | put_result.new_entry.kv.value = value; |
| ... | ... | @@ -227,6 +254,16 @@ pub fn HashMap(comptime K: type, comptime V: type, comptime hash: fn (key: K) u3 |
| 227 | 254 | return other; |
| 228 | 255 | } |
| 229 | 256 | |
| 257 | fn autoCapacity(self: *Self) !void { |
| 258 | if (self.entries.len == 0) { |
| 259 | return self.ensureCapacityExact(16); |
| 260 | } |
| 261 | // if we get too full (60%), double the capacity |
| 262 | if (self.size * 5 >= self.entries.len * 3) { |
| 263 | return self.ensureCapacityExact(self.entries.len * 2); |
| 264 | } |
| 265 | } |
| 266 | |
| 230 | 267 | fn initCapacity(hm: *Self, capacity: usize) !void { |
| 231 | 268 | hm.entries = try hm.allocator.alloc(Entry, capacity); |
| 232 | 269 | hm.size = 0; |
| ... | ... | @@ -427,6 +464,24 @@ test "iterator hash map" { |
| 427 | 464 | testing.expect(entry.value == values[0]); |
| 428 | 465 | } |
| 429 | 466 | |
| 467 | test "ensure capacity" { |
| 468 | var direct_allocator = std.heap.DirectAllocator.init(); |
| 469 | defer direct_allocator.deinit(); |
| 470 | |
| 471 | var map = AutoHashMap(i32, i32).init(&direct_allocator.allocator); |
| 472 | defer map.deinit(); |
| 473 | |
| 474 | try map.ensureCapacity(20); |
| 475 | const initialCapacity = map.entries.len; |
| 476 | testing.expect(initialCapacity >= 20); |
| 477 | var i : i32 = 0; |
| 478 | while (i < 20) : (i += 1) { |
| 479 | testing.expect(map.putAssumeCapacity(i, i+10) == null); |
| 480 | } |
| 481 | // shouldn't resize from putAssumeCapacity |
| 482 | testing.expect(initialCapacity == map.entries.len); |
| 483 | } |
| 484 | |
| 430 | 485 | pub fn getHashPtrAddrFn(comptime K: type) (fn (K) u32) { |
| 431 | 486 | return struct { |
| 432 | 487 | fn hash(key: K) u32 { |
Andrew Kelley <andrew@ziglang.org> 2019-05-04 14:13:18-04:00
GitHub <noreply@github.com> 2019-05-04 14:13:18-04:00