Expand sorted_{,multi}set a little

bazel-contrib · Jan 25, 2024 · e3e6b43 · e3e6b43
1 parent 7aafc05
commit e3e6b43
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 11 deletions.
diff --git a/java/gazelle/private/sorted_multiset/multiset.go b/java/gazelle/private/sorted_multiset/multiset.go
@@ -33,7 +33,15 @@ func (s *SortedMultiSet[K, V]) Keys() []K {
 	return s.keys.SortedSlice()
 }
 
-func (s *SortedMultiSet[K, V]) Values(key K) []V {
+func (s *SortedMultiSet[K, V]) Values(key K) *sorted_set.SortedSet[V] {
+	if s == nil {
+		return sorted_set.NewSortedSet[V](nil)
+	}
+
+	return s.ms[key]
+}
+
+func (s *SortedMultiSet[K, V]) SortedValues(key K) []V {
 	if s == nil {
 		return nil
 	}

diff --git a/java/gazelle/private/sorted_multiset/multiset_test.go b/java/gazelle/private/sorted_multiset/multiset_test.go
@@ -15,7 +15,7 @@ func TestNil(t *testing.T) {
 		t.Errorf("want len(nil.Keys()) == 0, got %d", gotKeysLen)
 	}
 
-	gotValues := s.Values("foo")
+	gotValues := s.SortedValues("foo")
 	if gotValuesLen := len(gotValues); gotValuesLen != 0 {
 		t.Errorf("want len(nil.Values(\"foo\")) == 0, got %d", gotValuesLen)
 	}
@@ -30,7 +30,12 @@ func TestMultiSet(t *testing.T) {
 			t.Errorf("want no keys, got %d: %v", gotKeysLen, gotKeys)
 		}
 
-		gotValues := s.Values("tasty")
+		gotValuesSlice := s.SortedValues("tasty")
+		if gotValuesLen := len(gotValuesSlice); gotValuesLen != 0 {
+			t.Errorf("want no values for tasty, got %d: %v", gotValuesLen, gotValuesSlice)
+		}
+
+		gotValues := s.Values("tasty").SortedSlice()
 		if gotValuesLen := len(gotValues); gotValuesLen != 0 {
 			t.Errorf("want no values for tasty, got %d: %v", gotValuesLen, gotValues)
 		}
@@ -45,12 +50,12 @@ func TestMultiSet(t *testing.T) {
 		}
 
 		wantTastyValues := []string{"hummus"}
-		gotTastyValues := s.Values("tasty")
+		gotTastyValues := s.SortedValues("tasty")
 		if !reflect.DeepEqual(wantTastyValues, gotTastyValues) {
 			t.Errorf("want tasty values %v got %v", wantTastyValues, gotTastyValues)
 		}
 
-		gotBadValues := s.Values("bad")
+		gotBadValues := s.SortedValues("bad")
 		if gotBadValuesLen := len(gotBadValues); gotBadValuesLen != 0 {
 			t.Errorf("want no values for tasty, got %d: %v", gotBadValuesLen, gotBadValues)
 		}
@@ -65,12 +70,12 @@ func TestMultiSet(t *testing.T) {
 		}
 
 		wantTastyValues := []string{"cheese", "hummus"}
-		gotTastyValues := s.Values("tasty")
+		gotTastyValues := s.SortedValues("tasty")
 		if !reflect.DeepEqual(wantTastyValues, gotTastyValues) {
 			t.Errorf("want tasty values %v got %v", wantTastyValues, gotTastyValues)
 		}
 
-		gotBadValues := s.Values("bad")
+		gotBadValues := s.SortedValues("bad")
 		if gotBadValuesLen := len(gotBadValues); gotBadValuesLen != 0 {
 			t.Errorf("want no values for tasty, got %d: %v", gotBadValuesLen, gotBadValues)
 		}
@@ -85,12 +90,12 @@ func TestMultiSet(t *testing.T) {
 		}
 
 		wantTastyValues := []string{"cheese", "hummus"}
-		gotTastyValues := s.Values("tasty")
+		gotTastyValues := s.SortedValues("tasty")
 		if !reflect.DeepEqual(wantTastyValues, gotTastyValues) {
 			t.Errorf("want tasty values %v got %v", wantTastyValues, gotTastyValues)
 		}
 
-		gotBadValues := s.Values("bad")
+		gotBadValues := s.SortedValues("bad")
 		if gotBadValuesLen := len(gotBadValues); gotBadValuesLen != 0 {
 			t.Errorf("want no values for tasty, got %d: %v", gotBadValuesLen, gotBadValues)
 		}
@@ -105,13 +110,18 @@ func TestMultiSet(t *testing.T) {
 		}
 
 		wantTastyValues := []string{"cheese", "hummus"}
-		gotTastyValues := s.Values("tasty")
+		gotTastyValues := s.SortedValues("tasty")
 		if !reflect.DeepEqual(wantTastyValues, gotTastyValues) {
 			t.Errorf("want tasty values %v got %v", wantTastyValues, gotTastyValues)
 		}
 
+		gotTastySortedSetValues := s.Values("tasty")
+		if !reflect.DeepEqual(wantTastyValues, gotTastySortedSetValues.SortedSlice()) {
+			t.Errorf("want tasty values %v got %v", wantTastyValues, gotTastyValues)
+		}
+
 		wantBadValues := []string{"soil"}
-		gotBadValues := s.Values("bad")
+		gotBadValues := s.SortedValues("bad")
 		if !reflect.DeepEqual(wantBadValues, gotBadValues) {
 			t.Errorf("want bad values %v got %v", wantBadValues, gotBadValues)
 		}

diff --git a/java/gazelle/private/sorted_set/btreeset.go b/java/gazelle/private/sorted_set/btreeset.go
@@ -66,6 +66,9 @@ func (s *SortedSet[T]) Add(v T) {
 
 // AddAll adds all of the elements from the other SortedSet to this one.
 func (s *SortedSet[T]) AddAll(other *SortedSet[T]) {
+	if other == nil {
+		return
+	}
 	other.tree.Ascend(func(v T) bool {
 		s.Add(v)
 		return true

diff --git a/java/gazelle/private/sorted_set/btreeset_test.go b/java/gazelle/private/sorted_set/btreeset_test.go
@@ -251,6 +251,13 @@ func TestAddAll(t *testing.T) {
 	if !reflect.DeepEqual(want, got) {
 		t.Errorf("want %v got %v", want, got)
 	}
+
+	s.AddAll(nil)
+
+	got = s.SortedSlice()
+	if !reflect.DeepEqual(want, got) {
+		t.Errorf("want %v got %v", want, got)
+	}
 }
 
 func TestFilter(t *testing.T) {