lib/uint64set: reduce memory allocations in Set.AppendTo

This commit is contained in:
Aliaksandr Valialkin 2020-01-17 22:32:30 +02:00
parent f9289b804a
commit bc0f897fcb
2 changed files with 56 additions and 4 deletions

View file

@ -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)

View file

@ -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)