diff --git a/src/array_hash_set.zig b/src/array_hash_set.zig index 59bc4b7..8cf906e 100644 --- a/src/array_hash_set.zig +++ b/src/array_hash_set.zig @@ -94,9 +94,16 @@ pub fn ArraySet(comptime E: type) type { } pub fn add(self: *Self, allocator: Allocator, element: E) Allocator.Error!bool { - const prevCount = self.unmanaged.count(); - try self.unmanaged.put(allocator, element, {}); - return prevCount != self.unmanaged.count(); + const result = try self.unmanaged.getOrPut(allocator, element); + return !result.found_existing; + } + + /// Adds a single element to the set. Asserts that there is enough capacity. + /// A bool is returned indicating if the element was actually added + /// if not already known. + pub fn addAssumeCapacity(self: *Self, element: E) bool { + const result = self.unmanaged.getOrPutAssumeCapacity(element); + return !result.found_existing; } /// Appends all elements from the provided slice, and may allocate. @@ -830,3 +837,28 @@ test "sizeOf matches" { try expectEqual(expectedByteSize, @sizeOf(std.AutoArrayHashMapUnmanaged(u32, void))); try expectEqual(expectedByteSize, @sizeOf(ArraySet(u32))); } + +test "add boolean" { + var A: ArraySet(u32) = .empty; + defer A.deinit(testing.allocator); + + const five_inserted = try A.add(testing.allocator, 5); + try expectEqual(five_inserted, true); + const six_inserted = try A.add(testing.allocator, 6); + try expectEqual(six_inserted, true); + + const five_again_inserted = try A.add(testing.allocator, 5); + try expectEqual(five_again_inserted, false); +} + +test "addAssumeCapacity" { + var A = try ArraySet(u32).initCapacity(testing.allocator, 2); + defer A.deinit(testing.allocator); + + try expectEqual(true, A.addAssumeCapacity(5)); + try expectEqual(true, A.addAssumeCapacity(6)); + try expectEqual(false, A.addAssumeCapacity(5)); + try expectEqual(2, A.cardinality()); + try expect(A.contains(5)); + try expect(A.contains(6)); +} diff --git a/src/hash_set.zig b/src/hash_set.zig index 435651e..ea627f6 100644 --- a/src/hash_set.zig +++ b/src/hash_set.zig @@ -113,18 +113,16 @@ pub fn HashSetWithContext(comptime E: type, comptime Context: type, comptime max } pub fn add(self: *Self, allocator: Allocator, element: E) Allocator.Error!bool { - const prevCount = self.unmanaged.count(); - try self.unmanaged.put(allocator, element, {}); - return prevCount != self.unmanaged.count(); + const result = try self.unmanaged.getOrPut(allocator, element); + return !result.found_existing; } /// Adds a single element to the set. Asserts that there is enough capacity. /// A bool is returned indicating if the element was actually added /// if not already known. pub fn addAssumeCapacity(self: *Self, element: E) bool { - const prevCount = self.unmanaged.count(); - self.unmanaged.putAssumeCapacity(element, {}); - return prevCount != self.unmanaged.count(); + const result = self.unmanaged.getOrPutAssumeCapacity(element); + return !result.found_existing; } /// Appends all elements from the provided set, and may allocate. @@ -1001,3 +999,30 @@ test "custom hash function string usage" { _ = try A.add(testing.allocator, "Hello\t"); try expectEqual(7, A.cardinality()); } + +test "add boolean" { + var A: HashSet(u32) = .empty; + defer A.deinit(testing.allocator); + + const five_inserted = try A.add(testing.allocator, 5); + try expectEqual(five_inserted, true); + const six_inserted = try A.add(testing.allocator, 6); + try expectEqual(six_inserted, true); + + const five_again_inserted = try A.add(testing.allocator, 5); + try expectEqual(five_again_inserted, false); +} + +test "addAssumeCapacity" { + var A: HashSet(u32) = try .initCapacity(testing.allocator, 3); + defer A.deinit(testing.allocator); + + try expect(A.addAssumeCapacity(1)); + try expect(A.addAssumeCapacity(2)); + try expect(A.addAssumeCapacity(3)); + + try expect(!A.addAssumeCapacity(1)); + + try expectEqual(3, A.cardinality()); + try expect(A.capacity() >= 3); +}