From 9a09ebb1b92f33913bdcc16e9d978fb6e5a820f1 Mon Sep 17 00:00:00 2001
From: Zander Khan <git@zander.xyz>
Date: Sun, 17 Jan 2021 14:41:20 +0000
Subject: [PATCH] Replace `shrink` with `shrinkAndFree` and
 `shrinkRetainingCapacity`

---
 lib/std/priority_dequeue.zig | 53 ++++++++++++++++++++++++++++++++++--
 lib/std/priority_queue.zig   | 53 ++++++++++++++++++++++++++++++++++--
 2 files changed, 100 insertions(+), 6 deletions(-)

diff --git a/lib/std/priority_dequeue.zig b/lib/std/priority_dequeue.zig
index 9655458890..bb119438d2 100644
--- a/lib/std/priority_dequeue.zig
+++ b/lib/std/priority_dequeue.zig
@@ -382,9 +382,29 @@ pub fn PriorityDequeue(comptime T: type) type {
             self.len = new_len;
         }
 
-        pub fn shrink(self: *Self, new_len: usize) void {
-            // TODO take advantage of the new realloc semantics
-            assert(new_len <= self.len);
+        /// Reduce allocated capacity to `new_len`.
+        pub fn shrinkAndFree(self: *Self, new_len: usize) void {
+            assert(new_len <= self.items.len);
+
+            // Cannot shrink to smaller than the current queue size without invalidating the heap property
+            assert(new_len >= self.len);
+
+            self.items = self.allocator.realloc(self.items[0..], new_len) catch |e| switch (e) {
+                error.OutOfMemory => { // no problem, capacity is still correct then.
+                    self.items.len = new_len;
+                    return;
+                },
+            };
+            self.len = new_len;
+        }
+
+        /// Reduce length to `new_len`.
+        pub fn shrinkRetainingCapacity(self: *Self, new_len: usize) void {
+            assert(new_len <= self.items.len);
+
+            // Cannot shrink to smaller than the current queue size without invalidating the heap property
+            assert(new_len >= self.len);
+
             self.len = new_len;
         }
 
@@ -798,6 +818,33 @@ test "std.PriorityDequeue: iterator while empty" {
     expectEqual(it.next(), null);
 }
 
+test "std.PriorityDequeue: shrinkRetainingCapacity and shrinkAndFree" {
+    var queue = PDQ.init(testing.allocator, lessThanComparison);
+    defer queue.deinit();
+
+    try queue.ensureCapacity(4);
+    expect(queue.capacity() >= 4);
+
+    try queue.add(1);
+    try queue.add(2);
+    try queue.add(3);
+    expect(queue.capacity() >= 4);
+    expectEqual(@as(usize, 3), queue.len);
+
+    queue.shrinkRetainingCapacity(3);
+    expect(queue.capacity() >= 4);
+    expectEqual(@as(usize, 3), queue.len);
+
+    queue.shrinkAndFree(3);
+    expectEqual(@as(usize, 3), queue.capacity());
+    expectEqual(@as(usize, 3), queue.len);
+
+    expectEqual(@as(u32, 3), queue.removeMax());
+    expectEqual(@as(u32, 2), queue.removeMax());
+    expectEqual(@as(u32, 1), queue.removeMax());
+    expect(queue.removeMaxOrNull() == null);
+}
+
 test "std.PriorityDequeue: fuzz testing min" {
     var prng = std.rand.DefaultPrng.init(0x12345678);
 
diff --git a/lib/std/priority_queue.zig b/lib/std/priority_queue.zig
index dc3070d1b3..844d37580c 100644
--- a/lib/std/priority_queue.zig
+++ b/lib/std/priority_queue.zig
@@ -192,9 +192,29 @@ pub fn PriorityQueue(comptime T: type) type {
             self.len = new_len;
         }
 
-        pub fn shrink(self: *Self, new_len: usize) void {
-            // TODO take advantage of the new realloc semantics
-            assert(new_len <= self.len);
+        /// Reduce allocated capacity to `new_len`.
+        pub fn shrinkAndFree(self: *Self, new_len: usize) void {
+            assert(new_len <= self.items.len);
+
+            // Cannot shrink to smaller than the current queue size without invalidating the heap property
+            assert(new_len >= self.len);
+
+            self.items = self.allocator.realloc(self.items[0..], new_len) catch |e| switch (e) {
+                error.OutOfMemory => { // no problem, capacity is still correct then.
+                    self.items.len = new_len;
+                    return;
+                },
+            };
+            self.len = new_len;
+        }
+
+        /// Reduce length to `new_len`.
+        pub fn shrinkRetainingCapacity(self: *Self, new_len: usize) void {
+            assert(new_len <= self.items.len);
+
+            // Cannot shrink to smaller than the current queue size without invalidating the heap property
+            assert(new_len >= self.len);
+
             self.len = new_len;
         }
 
@@ -477,6 +497,33 @@ test "std.PriorityQueue: iterator while empty" {
     expectEqual(it.next(), null);
 }
 
+test "std.PriorityQueue: shrinkRetainingCapacity and shrinkAndFree" {
+    var queue = PQ.init(testing.allocator, lessThan);
+    defer queue.deinit();
+
+    try queue.ensureCapacity(4);
+    expect(queue.capacity() >= 4);
+
+    try queue.add(1);
+    try queue.add(2);
+    try queue.add(3);
+    expect(queue.capacity() >= 4);
+    expectEqual(@as(usize, 3), queue.len);
+
+    queue.shrinkRetainingCapacity(3);
+    expect(queue.capacity() >= 4);
+    expectEqual(@as(usize, 3), queue.len);
+
+    queue.shrinkAndFree(3);
+    expectEqual(@as(usize, 3), queue.capacity());
+    expectEqual(@as(usize, 3), queue.len);
+
+    expectEqual(@as(u32, 1), queue.remove());
+    expectEqual(@as(u32, 2), queue.remove());
+    expectEqual(@as(u32, 3), queue.remove());
+    expect(queue.removeOrNull() == null);
+}
+
 test "std.PriorityQueue: update min heap" {
     var queue = PQ.init(testing.allocator, lessThan);
     defer queue.deinit();