From 7c87cea09c7676dcb85db16f99362f5db8456ff9 Mon Sep 17 00:00:00 2001
From: Kevin Conaway <kevin.conaway@datadoghq.com>
Date: Mon, 30 Mar 2020 11:54:39 -0400
Subject: [PATCH] Use stack allocated iterators in each container type to avoid
 duplicating iterator logic

---
 arraycontainer.go  |  6 ++++--
 bitmapcontainer.go |  9 +++++----
 roaring.go         |  7 +++----
 runcontainer.go    | 21 +++------------------
 4 files changed, 15 insertions(+), 28 deletions(-)

diff --git a/arraycontainer.go b/arraycontainer.go
index 8001b994..eb124f3b 100644
--- a/arraycontainer.go
+++ b/arraycontainer.go
@@ -25,8 +25,10 @@ func (ac *arrayContainer) fillLeastSignificant16bits(x []uint32, i int, mask uin
 }
 
 func (ac *arrayContainer) iterate(cb func(x uint16) bool) bool {
-	for i := 0; i < len(ac.content); i++ {
-		if !cb(ac.content[i]) {
+	iterator := shortIterator{ac.content, 0}
+
+	for iterator.hasNext() {
+		if !cb(iterator.next()) {
 			return false
 		}
 	}
diff --git a/bitmapcontainer.go b/bitmapcontainer.go
index ea419a14..cd259fd2 100644
--- a/bitmapcontainer.go
+++ b/bitmapcontainer.go
@@ -97,13 +97,14 @@ func (bc *bitmapContainer) maximum() uint16 {
 }
 
 func (bc *bitmapContainer) iterate(cb func(x uint16) bool) bool {
-	for i := bc.NextSetBit(0); i >= 0; {
-		j := i
-		i = bc.NextSetBit(i + 1)
-		if !cb(uint16(j)) {
+	iterator := bitmapContainerShortIterator{bc, bc.NextSetBit(0)}
+
+	for iterator.hasNext() {
+		if !cb(iterator.next()) {
 			return false
 		}
 	}
+
 	return true
 }
 
diff --git a/roaring.go b/roaring.go
index a1685447..ed75d58b 100644
--- a/roaring.go
+++ b/roaring.go
@@ -421,10 +421,8 @@ func (rb *Bitmap) String() string {
 // The iteration results are undefined if the bitmap is modified (e.g., with Add or Remove).
 // There is no guarantee as to what order the values will be iterated
 func (rb *Bitmap) Iterate(cb func(x uint32) bool) {
-	var hs uint32
 	for i := 0; i < rb.highlowcontainer.size(); i++ {
-		hs = uint32(rb.highlowcontainer.getKeyAtIndex(i)) << 16
-
+		hs := uint32(rb.highlowcontainer.getKeyAtIndex(i)) << 16
 		c := rb.highlowcontainer.getContainerAtIndex(i)
 
 		var shouldContinue bool
@@ -443,8 +441,9 @@ func (rb *Bitmap) Iterate(cb func(x uint32) bool) {
 				return cb(uint32(x) | hs)
 			})
 		}
+
 		if !shouldContinue {
-			return
+			break
 		}
 	}
 }
diff --git a/runcontainer.go b/runcontainer.go
index 4da49c16..5a0f985f 100644
--- a/runcontainer.go
+++ b/runcontainer.go
@@ -1163,25 +1163,10 @@ func (rc *runContainer16) newRunIterator16() *runIterator16 {
 }
 
 func (rc *runContainer16) iterate(cb func(x uint16) bool) bool {
-	curIndex := int64(0)
-	curPosInIndex := uint16(0)
+	iterator := runIterator16{rc, 0, 0}
 
-	hasNext := func() bool {
-		return int64(len(rc.iv)) > curIndex+1 ||
-			(int64(len(rc.iv)) == curIndex+1 && rc.iv[curIndex].length >= curPosInIndex)
-	}
-
-	for hasNext() {
-		next := rc.iv[curIndex].start + curPosInIndex
-
-		if curPosInIndex == rc.iv[curIndex].length {
-			curPosInIndex = 0
-			curIndex++
-		} else {
-			curPosInIndex++
-		}
-
-		if !cb(next) {
+	for iterator.hasNext() {
+		if !cb(iterator.next()) {
 			return false
 		}
 	}