From 5692982dd10eee7a1f99a8b4e3e630af6ea91342 Mon Sep 17 00:00:00 2001
From: dudaodong <lanliddd.2007@163.com>
Date: Tue, 29 Nov 2022 23:59:22 +0800
Subject: [PATCH] feat: add AddIfNotExistBy function for set

---
 datastructure/set/set.go      | 15 +++++++++++++++
 datastructure/set/set_test.go | 21 +++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/datastructure/set/set.go b/datastructure/set/set.go
index 8b5576b..243029f 100644
--- a/datastructure/set/set.go
+++ b/datastructure/set/set.go
@@ -30,6 +30,21 @@ func (s Set[T]) AddIfNotExist(value T) bool {
 	return false
 }
 
+// AddIfNotExistBy checks if item exists in the set and pass the `checker` function
+// it adds the item to set and returns true if it does not exists in the set and
+// function `checker` returns true, or else it does nothing and returns false.
+func (s Set[T]) AddIfNotExistBy(item T, checker func(element T) bool) bool {
+	if !s.Contain(item) {
+		if checker(item) {
+			if _, ok := s[item]; !ok {
+				s[item] = struct{}{}
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // Contain checks if set contains value or not
 func (s Set[T]) Contain(value T) bool {
 	_, ok := s[value]
diff --git a/datastructure/set/set_test.go b/datastructure/set/set_test.go
index b8067fc..691f203 100644
--- a/datastructure/set/set_test.go
+++ b/datastructure/set/set_test.go
@@ -28,6 +28,27 @@ func TestSet_AddIfNotExist(t *testing.T) {
 	assert.Equal(NewSet(1, 2, 3, 4), set)
 }
 
+func TestSet_AddIfNotExistBy(t *testing.T) {
+	assert := internal.NewAssert(t, "TestSet_AddIfNotExistBy")
+
+	set := NewSet[int]()
+	set.Add(1, 2)
+
+	ok := set.AddIfNotExistBy(3, func(val int) bool {
+		return val%2 != 0
+	})
+
+	notOk := set.AddIfNotExistBy(4, func(val int) bool {
+		return val%2 != 0
+	})
+
+	assert.Equal(true, ok)
+	assert.Equal(false, notOk)
+
+	assert.Equal(true, set.Contain(3))
+	assert.Equal(false, set.Contain(4))
+}
+
 func TestSet_Contain(t *testing.T) {
 	assert := internal.NewAssert(t, "TestSet_Contain")