From 9d5462dcb5b4b4601bdf2e628b9d80fb74000cb2 Mon Sep 17 00:00:00 2001
From: GethDW <33738921+GethDW@users.noreply.github.com>
Date: Tue, 4 Oct 2022 03:57:53 +0100
Subject: std: fix memory leak in ArrayHashMap (#13001)

---
 lib/std/array_hash_map.zig | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/lib/std/array_hash_map.zig b/lib/std/array_hash_map.zig
index 031a9fab5d..a502c5fa3f 100644
--- a/lib/std/array_hash_map.zig
+++ b/lib/std/array_hash_map.zig
@@ -773,9 +773,9 @@ pub fn ArrayHashMapUnmanaged(
                 }
             }
 
+            try self.entries.ensureTotalCapacity(allocator, new_capacity);
             const new_bit_index = try IndexHeader.findBitIndex(new_capacity);
             const new_header = try IndexHeader.alloc(allocator, new_bit_index);
-            try self.entries.ensureTotalCapacity(allocator, new_capacity);
 
             if (self.index_header) |old_header| old_header.free(allocator);
             self.insertAllEntriesIntoNewHeader(if (store_hash) {} else ctx, new_header);
@@ -2042,6 +2042,19 @@ test "ensure capacity" {
     try testing.expect(initial_capacity == map.capacity());
 }
 
+test "ensure capacity leak" {
+    try testing.checkAllAllocationFailures(std.testing.allocator, struct {
+        pub fn f(allocator: Allocator) !void {
+            var map = AutoArrayHashMap(i32, i32).init(allocator);
+            defer map.deinit();
+
+            var i: i32 = 0;
+            // put more than `linear_scan_max` in so index_header gets allocated.
+            while (i <= 20) : (i += 1) try map.put(i, i);
+        }
+    }.f, .{});
+}
+
 test "big map" {
     var map = AutoArrayHashMap(i32, i32).init(std.testing.allocator);
     defer map.deinit();
-- 
cgit v1.2.3