IterFuncByRankRange now behaves similar to GetByRankRange

- refactor sanitizeIndexes code
- refactor findNodeByRank code
- add test

sortedset.go

@@ -387,16 +387,8 @@ func (this *SortedSet) GetByScoreRange(start SCORE, end SCORE, options *GetBySco
 	return nodes
 }
 
-// Get nodes within specific rank range [start, end]
-// Note that the rank is 1-based integer. Rank 1 means the first node; Rank -1 means the last node;
-//
-// If start is greater than end, the returned array is in reserved order
-// If remove is true, the returned nodes are removed
-//
-// Time complexity of this method is : O(log(N))
-func (this *SortedSet) GetByRankRange(start int, end int, remove bool) []*SortedSetNode {
-
-	/* Sanitize indexes. */
+// sanitizeIndexes return start, end, and reverse flag
+func (this *SortedSet) sanitizeIndexes(start int, end int) (int, int, bool) {
 	if start < 0 {
 		start = int(this.length) + start + 1
 	}
@@ -414,12 +406,11 @@ func (this *SortedSet) GetByRankRange(start int, end int, remove bool) []*Sorted
 	if reverse { // swap start and end
 		start, end = end, start
 	}
+	return start, end, reverse
+}
 
-	var update [SKIPLIST_MAXLEVEL]*SortedSetNode
-	var nodes []*SortedSetNode
-	var traversed int = 0
-
-	x := this.header
+func (this *SortedSet) findNodeByRank(start int, remove bool) (traversed int, x *SortedSetNode, update [SKIPLIST_MAXLEVEL]*SortedSetNode) {
+	x = this.header
 	for i := this.level - 1; i >= 0; i-- {
 		for x.level[i].forward != nil &&
 			traversed+int(x.level[i].span) < start {
@@ -434,6 +425,22 @@ func (this *SortedSet) GetByRankRange(start int, end int, remove bool) []*Sorted
 			}
 		}
 	}
+	return
+}
+
+// Get nodes within specific rank range [start, end]
+// Note that the rank is 1-based integer. Rank 1 means the first node; Rank -1 means the last node;
+//
+// If start is greater than end, the returned array is in reserved order
+// If remove is true, the returned nodes are removed
+//
+// Time complexity of this method is : O(log(N))
+func (this *SortedSet) GetByRankRange(start int, end int, remove bool) []*SortedSetNode {
+	start, end, reverse := this.sanitizeIndexes(start, end)
+
+	var nodes []*SortedSetNode
+
+	traversed, x, update := this.findNodeByRank(start, remove)
 
 	traversed++
 	x = x.level[0].forward
@@ -509,19 +516,40 @@ func (this *SortedSet) FindRank(key string) int {
 	return 0
 }
 
-// IterFunc apply fn to node until fn return false or end of list
-func (this *SortedSet) IterFunc(fn func(i int, key string, value interface{}) bool) {
-	end := int(this.length)
-	var traversed int
-	x := this.header.level[0].forward
+// IterFuncByRankRange apply fn to node within specific rank range [start, end]
+// or until fn return false
+//
+// Note that the rank is 1-based integer. Rank 1 means the first node; Rank -1 means the last node;
+// If start is greater than end, apply fn in reserved order
+// If fn is nil, this function return without doing anything
+func (this *SortedSet) IterFuncByRankRange(start int, end int, fn func(key string, value interface{}) bool) {
+	if fn == nil {
+		return
+	}
+
+	start, end, reverse := this.sanitizeIndexes(start, end)
+	traversed, x, _ := this.findNodeByRank(start, false)
+	var nodes []*SortedSetNode
+
+	x = x.level[0].forward
 	for x != nil && traversed < end {
 		next := x.level[0].forward
 
-		if !fn(traversed, x.key, x.Value) {
+		if reverse {
+			nodes = append(nodes, x)
+		} else if !fn(x.key, x.Value) {
 			return
 		}
 
 		traversed++
 		x = next
 	}
+
+	if reverse {
+		for i := len(nodes) - 1; i >= 0; i-- {
+			if !fn(nodes[i].key, nodes[i].Value) {
+				return
+			}
+		}
+	}
 }

sortedset_test.go

@@ -16,6 +16,28 @@ func checkOrder(t *testing.T, nodes []*SortedSetNode, expectedOrder []string) {
 	}
 }
 
+func checkIterByRankRange(t *testing.T, sortedset *SortedSet, start int, end int, expectedOrder []string) {
+	var keys []string
+	sortedset.IterFuncByRankRange(start, end, func(key string, _ interface{}) bool {
+		keys = append(keys, key)
+		return true
+	})
+	if len(expectedOrder) != len(keys) {
+		t.Errorf("keys does not contain %d elements", len(expectedOrder))
+	}
+	for i := 0; i < len(expectedOrder); i++ {
+		if keys[i] != expectedOrder[i] {
+			t.Errorf("keys[%d] is %q, but the expected key is %q", i, keys[i], expectedOrder[i])
+		}
+	}
+}
+
+func checkRankRangeIterAndOrder(t *testing.T, sortedset *SortedSet, start int, end int, remove bool, expectedOrder []string) {
+	checkIterByRankRange(t, sortedset, start, end, expectedOrder)
+	nodes := sortedset.GetByRankRange(start, end, remove)
+	checkOrder(t, nodes, expectedOrder)
+}
+
 func TestCase1(t *testing.T) {
 	sortedset := New()
 
@@ -43,16 +65,13 @@ func TestCase1(t *testing.T) {
 	}
 
 	// get all nodes since the first one to last one
-	nodes := sortedset.GetByRankRange(1, -1, false)
-	checkOrder(t, nodes, []string{"d", "h", "a", "e", "f", "g", "c"})
+	checkRankRangeIterAndOrder(t, sortedset, 1, -1, false, []string{"d", "h", "a", "e", "f", "g", "c"})
 
 	// get & remove the 2nd/3rd nodes in reserve order
-	nodes = sortedset.GetByRankRange(-2, -3, true)
-	checkOrder(t, nodes, []string{"g", "f"})
+	checkRankRangeIterAndOrder(t, sortedset, -2, -3, true, []string{"g", "f"})
 
 	// get all nodes since the last one to first one
-	nodes = sortedset.GetByRankRange(-1, 1, false)
-	checkOrder(t, nodes, []string{"c", "e", "a", "h", "d"})
+	checkRankRangeIterAndOrder(t, sortedset, -1, 1, false, []string{"c", "e", "a", "h", "d"})
 
 }