diff --git a/lib/uint64set/uint64set.go b/lib/uint64set/uint64set.go index ef8f898492..ad826a2442 100644 --- a/lib/uint64set/uint64set.go +++ b/lib/uint64set/uint64set.go @@ -660,14 +660,18 @@ func (b *bucket16) delFromSmallPool(x uint16) bool { func (b *bucket16) appendTo(dst []uint64, hi uint32, hi16 uint16) []uint64 { hi64 := uint64(hi)<<32 | uint64(hi16)<<16 if b.bits == nil { - a := b.smallPool[:b.smallPoolLen] - if len(a) > 1 { - sort.Slice(a, func(i, j int) bool { return a[i] < a[j] }) + // Use uint16Sorter instead of sort.Slice here in order to reduce memory allocations. + a := uint16SorterPool.Get().(*uint16Sorter) + *a = uint16Sorter(b.smallPool[:b.smallPoolLen]) + if len(*a) > 1 && !sort.IsSorted(a) { + sort.Sort(a) } - for _, v := range a { + for _, v := range *a { x := hi64 | uint64(v) dst = append(dst, x) } + *a = nil + uint16SorterPool.Put(a) return dst } var wordNum uint64 @@ -691,6 +695,22 @@ func (b *bucket16) appendTo(dst []uint64, hi uint32, hi16 uint16) []uint64 { return dst } +var uint16SorterPool = &sync.Pool{ + New: func() interface{} { + return &uint16Sorter{} + }, +} + +type uint16Sorter []uint16 + +func (s uint16Sorter) Len() int { return len(s) } +func (s uint16Sorter) Less(i, j int) bool { + return s[i] < s[j] +} +func (s uint16Sorter) Swap(i, j int) { + s[i], s[j] = s[j], s[i] +} + func getWordNumBitMask(x uint16) (uint16, uint64) { wordNum := x / 64 bitMask := uint64(1) << (x & 63) diff --git a/lib/uint64set/uint64set_test.go b/lib/uint64set/uint64set_test.go index e8ac8ebd1e..bbb9675bcc 100644 --- a/lib/uint64set/uint64set_test.go +++ b/lib/uint64set/uint64set_test.go @@ -240,6 +240,38 @@ func testSetBasicOps(t *testing.T, itemsCount int) { } } + // Verify UnionMayOwn + { + const unionOffset = 12345 + var s1, s2 Set + for i := 0; i < itemsCount; i++ { + s1.Add(uint64(i) + offset) + s2.Add(uint64(i) + offset + unionOffset) + } + s1.UnionMayOwn(&s2) + expectedLen := 2 * itemsCount + if itemsCount > unionOffset { + expectedLen = itemsCount + unionOffset + } + if n := s1.Len(); n != expectedLen { + t.Fatalf("unexpected s1.Len() after union; got %d; want %d", n, expectedLen) + } + + // Verify union on empty set. + var s3 Set + expectedLen = s1.Len() + s3.UnionMayOwn(&s1) + if n := s3.Len(); n != expectedLen { + t.Fatalf("unexpected s3.Len() after union with empty set; got %d; want %d", n, expectedLen) + } + var s4 Set + expectedLen = s3.Len() + s3.UnionMayOwn(&s4) + if n := s3.Len(); n != expectedLen { + t.Fatalf("unexpected s3.Len() after union with empty set; got %d; want %d", n, expectedLen) + } + } + // Verify intersect { // Verify s1.Intersect(s2) and s2.Intersect(s1)