From 2cb811b42f6d0589defd8aadd1c2d19b260407df Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Wed, 25 Sep 2019 00:34:09 +0300
Subject: [PATCH] lib/uint64set: optimize Set.AppendTo

---
 lib/uint64set/uint64set.go      | 17 +++++++++++++----
 lib/uint64set/uint64set_test.go | 14 +++++++-------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/lib/uint64set/uint64set.go b/lib/uint64set/uint64set.go
index e357c6669..5132c6e27 100644
--- a/lib/uint64set/uint64set.go
+++ b/lib/uint64set/uint64set.go
@@ -1,6 +1,7 @@
 package uint64set
 
 import (
+	"math/bits"
 	"sort"
 )
 
@@ -299,11 +300,19 @@ func (b *bucket16) appendTo(dst []uint64, hi uint32, hi16 uint16) []uint64 {
 	hi64 := uint64(hi)<<32 | uint64(hi16)<<16
 	var wordNum uint64
 	for _, word := range b.bits {
-		for bitNum := uint64(0); bitNum < 64; bitNum++ {
-			if word&(uint64(1)<<bitNum) != 0 {
-				x := hi64 | uint64(wordNum)*64 | bitNum
-				dst = append(dst, x)
+		if word == 0 {
+			wordNum++
+			continue
+		}
+		x64 := hi64 | (wordNum * 64)
+		for {
+			tzn := uint64(bits.TrailingZeros64(word))
+			if tzn >= 64 {
+				break
 			}
+			word &^= uint64(1) << tzn
+			x := x64 | tzn
+			dst = append(dst, x)
 		}
 		wordNum++
 	}
diff --git a/lib/uint64set/uint64set_test.go b/lib/uint64set/uint64set_test.go
index a6a60d275..2b6733b60 100644
--- a/lib/uint64set/uint64set_test.go
+++ b/lib/uint64set/uint64set_test.go
@@ -48,14 +48,14 @@ func testSetBasicOps(t *testing.T, itemsCount int) {
 	// Verify Has on existing bits
 	for i := 0; i < itemsCount; i++ {
 		if !s.Has(uint64(i) + offset) {
-			t.Fatalf("missing bit %d", i)
+			t.Fatalf("missing bit %d", uint64(i)+offset)
 		}
 	}
 
 	// Verify Has on missing bits
 	for i := itemsCount; i < 2*itemsCount; i++ {
 		if s.Has(uint64(i) + offset) {
-			t.Fatalf("unexpected bit found: %d", i)
+			t.Fatalf("unexpected bit found: %d", uint64(i)+offset)
 		}
 	}
 
@@ -66,7 +66,7 @@ func testSetBasicOps(t *testing.T, itemsCount int) {
 	}
 	for i := 0; i < itemsCount; i++ {
 		if !sCopy.Has(uint64(i) + offset) {
-			t.Fatalf("missing bit %d on sCopy", i)
+			t.Fatalf("missing bit %d on sCopy", uint64(i)+offset)
 		}
 	}
 
@@ -84,7 +84,7 @@ func testSetBasicOps(t *testing.T, itemsCount int) {
 	}
 	for i := 0; i < itemsCount; i++ {
 		if !m[uint64(i)+offset] {
-			t.Fatalf("missing bit %d in the exported bits; array:\n%d", i, a)
+			t.Fatalf("missing bit %d in the exported bits; array:\n%d", uint64(i)+offset, a)
 		}
 	}
 
@@ -106,11 +106,11 @@ func testSetBasicOps(t *testing.T, itemsCount int) {
 	for i := 0; i < itemsCount; i++ {
 		if i >= itemsCount/2 && i < itemsCount-itemsCount/4 {
 			if m[uint64(i)+offset] {
-				t.Fatalf("unexpected bit found after deleting: %d", i)
+				t.Fatalf("unexpected bit found after deleting: %d", uint64(i)+offset)
 			}
 		} else {
 			if !m[uint64(i)+offset] {
-				t.Fatalf("missing bit %d in the exported bits after deleting", i)
+				t.Fatalf("missing bit %d in the exported bits after deleting", uint64(i)+offset)
 			}
 		}
 	}
@@ -131,7 +131,7 @@ func testSetBasicOps(t *testing.T, itemsCount int) {
 	}
 	for i := 0; i < itemsCount; i++ {
 		if !sCopy.Has(uint64(i) + offset) {
-			t.Fatalf("missing bit %d on sCopy", i)
+			t.Fatalf("missing bit %d on sCopy", uint64(i)+offset)
 		}
 	}
 }