authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2019-05-04 14:13:18-04:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2019-05-04 14:13:18-04:00
log6f0aa801c80a434b7cc931e879d234a8de501e38
tree9c04d6cfeb4685c0d7ff71a3c005e41e1566154b
parent21c8d57fca6bad5c9573679c2f7b40e584bd597c
parentcf8dde2d686199474847c6c4e342dc1ac46a435b
signature Signed by PGP key 4AEE18F83AFDEB23

Merge pull request #2404 from squeek502/hash-map-ensure-cap

std.HashMap: add public ensureCapacity fn

1 files changed, 64 insertions(+), 9 deletions(-)

std/hash_map.zig+64-9
...@@ -118,7 +118,7 @@ pub fn HashMap(comptime K: type, comptime V: type, comptime hash: fn (key: K) u3...@@ -118,7 +118,7 @@ pub fn HashMap(comptime K: type, comptime V: type, comptime hash: fn (key: K) u3
118 };118 };
119 }119 }
120 self.incrementModificationCount();120 self.incrementModificationCount();
121 try self.ensureCapacity();121 try self.autoCapacity();
122 const put_result = self.internalPut(key);122 const put_result = self.internalPut(key);
123 assert(put_result.old_kv == null);123 assert(put_result.old_kv == null);
124 return GetOrPutResult{124 return GetOrPutResult{
...@@ -135,15 +135,37 @@ pub fn HashMap(comptime K: type, comptime V: type, comptime hash: fn (key: K) u3...@@ -135,15 +135,37 @@ pub fn HashMap(comptime K: type, comptime V: type, comptime hash: fn (key: K) u3
135 return res.kv;135 return res.kv;
136 }136 }
137137
138 fn ensureCapacity(self: *Self) !void {138 fn optimizedCapacity(expected_count: usize) usize {
139 if (self.entries.len == 0) {139 // ensure that the hash map will be at most 60% full if
140 return self.initCapacity(16);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 }
142164
143 // if we get too full (60%), double the capacity165 const old_entries = self.entries;
144 if (self.size * 5 >= self.entries.len * 3) {166 try self.initCapacity(new_capacity);
145 const old_entries = self.entries;167 self.incrementModificationCount();
146 try self.initCapacity(self.entries.len * 2);168 if (old_entries.len > 0) {
147 // dump all of the old elements into the new table169 // dump all of the old elements into the new table
148 for (old_entries) |*old_entry| {170 for (old_entries) |*old_entry| {
149 if (old_entry.used) {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,8 +178,13 @@ pub fn HashMap(comptime K: type, comptime V: type, comptime hash: fn (key: K) u3
156178
157 /// Returns the kv pair that was already there.179 /// Returns the kv pair that was already there.
158 pub fn put(self: *Self, key: K, value: V) !?KV {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 self.incrementModificationCount();187 self.incrementModificationCount();
160 try self.ensureCapacity();
161188
162 const put_result = self.internalPut(key);189 const put_result = self.internalPut(key);
163 put_result.new_entry.kv.value = value;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,6 +254,16 @@ pub fn HashMap(comptime K: type, comptime V: type, comptime hash: fn (key: K) u3
227 return other;254 return other;
228 }255 }
229256
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 fn initCapacity(hm: *Self, capacity: usize) !void {267 fn initCapacity(hm: *Self, capacity: usize) !void {
231 hm.entries = try hm.allocator.alloc(Entry, capacity);268 hm.entries = try hm.allocator.alloc(Entry, capacity);
232 hm.size = 0;269 hm.size = 0;
...@@ -427,6 +464,24 @@ test "iterator hash map" {...@@ -427,6 +464,24 @@ test "iterator hash map" {
427 testing.expect(entry.value == values[0]);464 testing.expect(entry.value == values[0]);
428}465}
429466
467test "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
430pub fn getHashPtrAddrFn(comptime K: type) (fn (K) u32) {485pub fn getHashPtrAddrFn(comptime K: type) (fn (K) u32) {
431 return struct {486 return struct {
432 fn hash(key: K) u32 {487 fn hash(key: K) u32 {