1 files changed, 39 insertions, 20 deletions

diff --git a/lib/std/buf_set.zig b/lib/std/buf_set.zig
index 5e09aa6bcb..c2a1f9bcb9 100644
--- a/lib/std/buf_set.zig
+++ b/lib/std/buf_set.zig
@@ -9,50 +9,69 @@ const mem = @import("mem.zig");
 const Allocator = mem.Allocator;
 const testing = std.testing;
 
+/// A BufSet is a set of strings.  The BufSet duplicates
+/// strings internally, and never takes ownership of strings
+/// which are passed to it.
 pub const BufSet = struct {
     hash_map: BufSetHashMap,
 
     const BufSetHashMap = StringHashMap(void);
+    pub const Iterator = BufSetHashMap.KeyIterator;
 
+    /// Create a BufSet using an allocator.  The allocator will
+    /// be used internally for both backing allocations and
+    /// string duplication.
     pub fn init(a: *Allocator) BufSet {
         var self = BufSet{ .hash_map = BufSetHashMap.init(a) };
         return self;
     }
 
+    /// Free a BufSet along with all stored keys.
     pub fn deinit(self: *BufSet) void {
-        var it = self.hash_map.iterator();
-        while (it.next()) |entry| {
-            self.free(entry.key);
+        var it = self.hash_map.keyIterator();
+        while (it.next()) |key_ptr| {
+            self.free(key_ptr.*);
         }
         self.hash_map.deinit();
         self.* = undefined;
     }
 
-    pub fn put(self: *BufSet, key: []const u8) !void {
-        if (self.hash_map.get(key) == null) {
-            const key_copy = try self.copy(key);
-            errdefer self.free(key_copy);
-            try self.hash_map.put(key_copy, {});
+    /// Insert an item into the BufSet.  The item will be
+    /// copied, so the caller may delete or reuse the
+    /// passed string immediately.
+    pub fn insert(self: *BufSet, value: []const u8) !void {
+        const gop = try self.hash_map.getOrPut(value);
+        if (!gop.found_existing) {
+            gop.key_ptr.* = self.copy(value) catch |err| {
+                _ = self.hash_map.remove(value);
+                return err;
+            };
         }
     }
 
-    pub fn exists(self: BufSet, key: []const u8) bool {
-        return self.hash_map.get(key) != null;
+    /// Check if the set contains an item matching the passed string
+    pub fn contains(self: BufSet, value: []const u8) bool {
+        return self.hash_map.contains(value);
     }
 
-    pub fn delete(self: *BufSet, key: []const u8) void {
-        const entry = self.hash_map.remove(key) orelse return;
-        self.free(entry.key);
+    /// Remove an item from the set.
+    pub fn remove(self: *BufSet, value: []const u8) void {
+        const kv = self.hash_map.fetchRemove(value) orelse return;
+        self.free(kv.key);
     }
 
+    /// Returns the number of items stored in the set
     pub fn count(self: *const BufSet) usize {
         return self.hash_map.count();
     }
 
-    pub fn iterator(self: *const BufSet) BufSetHashMap.Iterator {
-        return self.hash_map.iterator();
+    /// Returns an iterator over the items stored in the set.
+    /// Iteration order is arbitrary.
+    pub fn iterator(self: *const BufSet) Iterator {
+        return self.hash_map.keyIterator();
     }
 
+    /// Get the allocator used by this set
     pub fn allocator(self: *const BufSet) *Allocator {
         return self.hash_map.allocator;
     }
@@ -72,12 +91,12 @@ test "BufSet" {
     var bufset = BufSet.init(std.testing.allocator);
     defer bufset.deinit();
 
-    try bufset.put("x");
+    try bufset.insert("x");
     try testing.expect(bufset.count() == 1);
-    bufset.delete("x");
+    bufset.remove("x");
     try testing.expect(bufset.count() == 0);
 
-    try bufset.put("x");
-    try bufset.put("y");
-    try bufset.put("z");
+    try bufset.insert("x");
+    try bufset.insert("y");
+    try bufset.insert("z");
 }