Просмотр исходного кода

util/set: add more functionality to IntSet (#16640)

Joe Tsai 7 месяцев назад
Родитель
Сommit
1cc842b389
2 измененных файлов с 32 добавлено и 0 удалено
  1. 26 0
      util/set/intset.go
  2. 6 0
      util/set/intset_test.go

+ 26 - 0
util/set/intset.go

@@ -28,6 +28,15 @@ type IntSet[T constraints.Integer] struct {
 	extraLen int
 }
 
+// IntsOf constructs an [IntSet] with the provided elements.
+func IntsOf[T constraints.Integer](slice ...T) IntSet[T] {
+	var s IntSet[T]
+	for _, e := range slice {
+		s.Add(e)
+	}
+	return s
+}
+
 // Values returns an iterator over the elements of the set.
 // The iterator will yield the elements in no particular order.
 func (s IntSet[T]) Values() iter.Seq[T] {
@@ -111,6 +120,23 @@ func (s *IntSet[T]) Delete(e T) {
 	}
 }
 
+// DeleteSeq deletes the values in seq from the set.
+func (s *IntSet[T]) DeleteSeq(seq iter.Seq[T]) {
+	for e := range seq {
+		s.Delete(e)
+	}
+}
+
+// Equal reports whether s is equal to other.
+func (s IntSet[T]) Equal(other IntSet[T]) bool {
+	for hi, bits := range s.extra {
+		if other.extra[hi] != bits {
+			return false
+		}
+	}
+	return s.extraLen == other.extraLen && s.bits == other.bits
+}
+
 // Clone returns a copy of s that doesn't alias the original.
 func (s IntSet[T]) Clone() IntSet[T] {
 	return IntSet[T]{

+ 6 - 0
util/set/intset_test.go

@@ -47,6 +47,9 @@ func TestIntSet(t *testing.T) {
 		deleteInt(t, ss, &si, math.MinInt64)
 		deleteInt(t, ss, &si, math.MaxInt64)
 		intValues(t, ss, si)
+		if !si.Equal(IntsOf(ss.Slice()...)) {
+			t.Errorf("{%v}.Equal({%v}) = false, want true", si, ss)
+		}
 	})
 
 	t.Run("Uint64", func(t *testing.T) {
@@ -80,6 +83,9 @@ func TestIntSet(t *testing.T) {
 		intValues(t, ss, si)
 		deleteInt(t, ss, &si, math.MaxInt64)
 		intValues(t, ss, si)
+		if !si.Equal(IntsOf(ss.Slice()...)) {
+			t.Errorf("{%v}.Equal({%v}) = false, want true", si, ss)
+		}
 	})
 }