use generic version of btree (#3404)

This commit is contained in:
Denis Mishin 2022-06-06 14:31:05 -04:00 committed by GitHub
parent 63673363c5
commit f7b6ed0ad4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -2,51 +2,47 @@ package sets
import "github.com/google/btree" import "github.com/google/btree"
type stringItem string func lessFn(a, b string) bool { return a < b }
func (item stringItem) Less(than btree.Item) bool {
return item < than.(stringItem)
}
// A SortedString is a set of strings with sorted iteration. // A SortedString is a set of strings with sorted iteration.
type SortedString struct { type SortedString struct {
b *btree.BTree b *btree.BTreeG[string]
} }
// NewSortedString creates a new sorted string set. // NewSortedString creates a new sorted string set.
func NewSortedString() *SortedString { func NewSortedString() *SortedString {
return &SortedString{ return &SortedString{
b: btree.New(8), b: btree.NewG[string](8, lessFn),
} }
} }
// Add adds a string to the set. // Add adds a string to the set.
func (s *SortedString) Add(elements ...string) { func (s *SortedString) Add(elements ...string) {
for _, element := range elements { for _, element := range elements {
s.b.ReplaceOrInsert(stringItem(element)) s.b.ReplaceOrInsert(element)
} }
} }
// Clear clears the set. // Clear clears the set.
func (s *SortedString) Clear() { func (s *SortedString) Clear() {
s.b = btree.New(8) s.b = btree.NewG[string](8, lessFn)
} }
// Delete deletes an element from the set. // Delete deletes an element from the set.
func (s *SortedString) Delete(element string) { func (s *SortedString) Delete(element string) {
s.b.Delete(stringItem(element)) s.b.Delete(element)
} }
// ForEach iterates over the set in ascending order. // ForEach iterates over the set in ascending order.
func (s *SortedString) ForEach(callback func(element string) bool) { func (s *SortedString) ForEach(callback func(element string) bool) {
s.b.Ascend(func(i btree.Item) bool { s.b.Ascend(func(item string) bool {
return callback(string(i.(stringItem))) return callback(item)
}) })
} }
// Has returns true if the elment is in the set. // Has returns true if the elment is in the set.
func (s *SortedString) Has(element string) bool { func (s *SortedString) Has(element string) bool {
return s.b.Has(stringItem(element)) return s.b.Has(element)
} }
// Size returns the size of the set. // Size returns the size of the set.
@ -57,8 +53,8 @@ func (s *SortedString) Size() int {
// ToSlice returns a slice of all the elements in the set. // ToSlice returns a slice of all the elements in the set.
func (s *SortedString) ToSlice() []string { func (s *SortedString) ToSlice() []string {
arr := make([]string, 0, s.Size()) arr := make([]string, 0, s.Size())
s.b.Ascend(func(i btree.Item) bool { s.b.Ascend(func(item string) bool {
arr = append(arr, string(i.(stringItem))) arr = append(arr, item)
return true return true
}) })
return arr return arr