From ab4d5d72ebaebea8a33c0944631df111cc7646cd Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Fri, 17 Jan 2020 22:32:30 +0200
Subject: [PATCH] lib/uint64set: reduce memory allocations in Set.AppendTo

---
 lib/uint64set/uint64set.go      | 28 ++++++++++++++++++++++++----
 lib/uint64set/uint64set_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 4 deletions(-)

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)