expose sorted sets builder

2 files changed, 37 insertions, 4 deletions

diff --git a/sets.go b/sets.go
index c8ce39f..b41bd37 100644
--- a/sets.go
+++ b/sets.go
@@ -194,13 +194,13 @@ func (itr *SortedSetIterator[T]) Next() (val T, ok bool) {
 	return
 }
 
-// Next moves the iterator to the previous value.
+// Prev moves the iterator to the previous value.
 func (itr *SortedSetIterator[T]) Prev() (val T, ok bool) {
 	val, _, ok = itr.mi.Prev()
 	return
 }
 
-// Next moves the iterator to the given value.
+// Seek moves the iterator to the given value.
 //
 // If the value does not exist then the next value is used. If no more keys exist
 // then the iterator is marked as done.
@@ -209,11 +209,12 @@ func (itr *SortedSetIterator[T]) Seek(val T) {
 }
 
 type SortedSetBuilder[T any] struct {
-	s SortedSet[T]
+	s *SortedSet[T]
 }
 
 func NewSortedSetBuilder[T any](comparer Comparer[T]) *SortedSetBuilder[T] {
-	return &SortedSetBuilder[T]{s: NewSortedSet(comparer)}
+	s := NewSortedSet(comparer)
+	return &SortedSetBuilder[T]{s: &s}
 }
 
 func (s SortedSetBuilder[T]) Set(val T) {
@@ -231,3 +232,12 @@ func (s SortedSetBuilder[T]) Has(val T) bool {
 func (s SortedSetBuilder[T]) Len() int {
 	return s.s.Len()
 }
+
+// SortedSet returns the current copy of the set.
+// The builder should not be used again after the list after this call.
+func (s SortedSetBuilder[T]) SortedSet() SortedSet[T] {
+	assert(s.s != nil, "immutable.SortedSetBuilder.SortedSet(): duplicate call to fetch sorted set")
+	set := s.s
+	s.s = nil
+	return *set
+}
diff --git a/sets_test.go b/sets_test.go
index ca3aa3d..6612cba 100644
--- a/sets_test.go
+++ b/sets_test.go
@@ -101,3 +101,26 @@ func TestSortedSetsDelete(t *testing.T) {
 		t.Fatalf("Unexpected set element after delete")
 	}
 }
+
+func TestSortedSetBuilder(t *testing.T) {
+	b := NewSortedSetBuilder[string](nil)
+	b.Set("test3")
+	b.Set("test1")
+	b.Set("test2")
+
+	s := b.SortedSet()
+	items := s.Items()
+
+	if len(items) != 3 {
+		t.Fatalf("Set has wrong number of items")
+	}
+	if items[0] != "test1" {
+		t.Fatalf("First item incorrectly sorted")
+	}
+	if items[1] != "test2" {
+		t.Fatalf("Second item incorrectly sorted")
+	}
+	if items[2] != "test3" {
+		t.Fatalf("Third item incorrectly sorted")
+	}
+}