feat: optimize list lrange method

xgzlucario · Sep 8, 2024 · 5a5e387 · 5a5e387
1 parent 997d63f
commit 5a5e387
Show file tree

Hide file tree

Showing 5 changed files with 86 additions and 77 deletions.
diff --git a/command.go b/command.go
@@ -350,15 +350,7 @@ func lrangeCommand(writer *RESPWriter, args []RESP) {
 		writer.WriteError(err)
 		return
 	}
-
-	if stop == -1 {
-		stop = ls.Size()
-	} else if stop < ls.Size() {
-		stop++ // range 1 3 means range[1,3]
-	}
-	start = min(start, stop)
-
-	writer.WriteArrayHead(stop - start)
+	writer.WriteArrayHead(ls.RangeCount(start, stop))
 	ls.Range(start, stop, func(data []byte) {
 		writer.WriteBulk(data)
 	})

diff --git a/command_test.go b/command_test.go
@@ -199,31 +199,45 @@ func testCommand(t *testing.T, rdb *redis.Client, sleepFn func(time.Duration)) {
 
 	t.Run("list", func(t *testing.T) {
 		// lpush
-		n, _ := rdb.LPush(ctx, "list", "a", "b", "c").Result()
+		n, _ := rdb.LPush(ctx, "list", "3", "2", "1").Result()
 		assert.Equal(n, int64(3))
 
 		// rpush
-		n, _ = rdb.RPush(ctx, "list", "d", "e", "f").Result()
+		n, _ = rdb.RPush(ctx, "list", "4", "5", "6").Result()
 		assert.Equal(n, int64(6))
 
+		// list: [1,2,3,4,5,6]
 		// lrange
 		res, _ := rdb.LRange(ctx, "list", 0, -1).Result()
-		assert.Equal(res, []string{"c", "b", "a", "d", "e", "f"})
-
+		assert.Equal(res, []string{"1", "2", "3", "4", "5", "6"})
+		res, _ = rdb.LRange(ctx, "list", -100, 100).Result()
+		assert.Equal(res, []string{"1", "2", "3", "4", "5", "6"})
 		res, _ = rdb.LRange(ctx, "list", 1, 3).Result()
-		assert.Equal(res, []string{"b", "a", "d"})
-
-		res, err := rdb.LRange(ctx, "list", 3, 2).Result()
-		assert.Equal(len(res), 0)
-		assert.Nil(err)
+		assert.Equal(res, []string{"2", "3", "4"})
+		res, _ = rdb.LRange(ctx, "list", 3, 3).Result()
+		assert.Equal(res, []string{"4"})
+		res, _ = rdb.LRange(ctx, "list", -5, 2).Result()
+		assert.Equal(res, []string{"2", "3"})
+
+		// revrange not support
+		res, _ = rdb.LRange(ctx, "list", -1, -3).Result()
+		assert.Equal(res, []string{})
+		res, _ = rdb.LRange(ctx, "list", -1, 2).Result()
+		assert.Equal(res, []string{})
+		res, _ = rdb.LRange(ctx, "list", 3, 2).Result()
+		assert.Equal(res, []string{})
+		res, _ = rdb.LRange(ctx, "list", 99, 100).Result()
+		assert.Equal(res, []string{})
+		res, _ = rdb.LRange(ctx, "list", -100, -99).Result()
+		assert.Equal(res, []string{})
 
 		// lpop
 		val, _ := rdb.LPop(ctx, "list").Result()
-		assert.Equal(val, "c")
+		assert.Equal(val, "1")
 
 		// rpop
 		val, _ = rdb.RPop(ctx, "list").Result()
-		assert.Equal(val, "f")
+		assert.Equal(val, "6")
 
 		// pop nil
 		{
@@ -237,7 +251,7 @@ func testCommand(t *testing.T, rdb *redis.Client, sleepFn func(time.Duration)) {
 		// error wrong type
 		rdb.Set(ctx, "key", "value", 0)
 
-		_, err = rdb.LPush(ctx, "key", "1").Result()
+		_, err := rdb.LPush(ctx, "key", "1").Result()
 		assert.Equal(err.Error(), errWrongType.Error())
 
 		_, err = rdb.RPush(ctx, "key", "1").Result()
@@ -380,6 +394,9 @@ func testCommand(t *testing.T, rdb *redis.Client, sleepFn func(time.Duration)) {
 
 		_, err = rdb.ZRem(ctx, "key", "member1").Result()
 		assert.Equal(err.Error(), errWrongType.Error())
+
+		_, err = rdb.ZPopMin(ctx, "key").Result()
+		assert.Equal(err.Error(), errWrongType.Error())
 	})
 
 	t.Run("eval", func(t *testing.T) {

diff --git a/internal/list/bench_test.go b/internal/list/bench_test.go
@@ -38,11 +38,4 @@ func BenchmarkList(b *testing.B) {
 			ls.Range(0, -1, func([]byte) {})
 		}
 	})
-	b.Run("revrange", func(b *testing.B) {
-		ls := genList(0, 100)
-		b.ResetTimer()
-		for i := 0; i < b.N; i++ {
-			ls.RevRange(0, -1, func([]byte) {})
-		}
-	})
 }
diff --git a/internal/list/list.go b/internal/list/list.go
@@ -87,15 +87,36 @@ func (ls *QuickList) free(n *Node) {
 
 func (ls *QuickList) Size() int { return ls.size }
 
-func (ls *QuickList) Range(start, stop int, f func(data []byte)) {
-	count := stop - start
+func (ls *QuickList) RangeCount(start, stop int) int {
+	if start < 0 {
+		start += ls.Size()
+	}
+	if stop < 0 {
+		stop += ls.Size()
+	}
+	start = max(0, start)
+	stop = min(ls.Size(), stop)
+
+	if start <= stop {
+		return min(ls.Size(), stop-start+1)
+	}
+	return 0
+}
+
+func (ls *QuickList) Range(start, stop int, fn func(data []byte)) {
+	count := ls.RangeCount(start, stop)
+	if count == 0 {
+		return
+	}
+	if start < 0 {
+		start += ls.Size()
+	}
 
 	lp := ls.head
 	for lp != nil && start > lp.Size() {
 		start -= lp.Size()
 		lp = lp.next
 	}
-
 	it := lp.Iterator().SeekFirst()
 	for range start {
 		it.Next()
@@ -109,32 +130,6 @@ func (ls *QuickList) Range(start, stop int, f func(data []byte)) {
 			lp = lp.next
 			it = lp.Iterator().SeekFirst()
 		}
-		f(it.Next())
-	}
-}
-
-func (ls *QuickList) RevRange(start, stop int, f func(data []byte)) {
-	count := stop - start
-
-	lp := ls.tail
-	for lp != nil && start > lp.Size() {
-		start -= lp.Size()
-		lp = lp.prev
-	}
-
-	it := lp.Iterator().SeekLast()
-	for range start {
-		it.Prev()
-	}
-
-	for range count {
-		if it.IsFirst() {
-			if lp.prev == nil {
-				return
-			}
-			lp = lp.prev
-			it = lp.Iterator().SeekLast()
-		}
-		f(it.Prev())
+		fn(it.Next())
 	}
 }
diff --git a/internal/list/list_test.go b/internal/list/list_test.go
@@ -2,6 +2,7 @@ package list
 
 import (
 	"slices"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -79,40 +80,51 @@ func TestList(t *testing.T) {
 	})
 
 	t.Run("range", func(t *testing.T) {
-		ls := genList(0, N)
-		i := 0
-		ls.Range(0, N, func(data []byte) {
-			assert.Equal(string(data), genKey(i))
-			i++
-		})
-		assert.Equal(i, N)
+		ls := New()
+		// [0, 1, 2, 3, 4]
+		for i := range 5 {
+			ls.RPush(strconv.Itoa(i))
+		}
 
-		for _, start := range []int{100, 1000, 5000} {
-			i = 0
-			ls.Range(start, start+100, func(data []byte) {
-				assert.Equal(string(data), genKey(start+i))
-				i++
+		rangeFn := func(start, stop int) (res []string) {
+			ls.Range(start, stop, func(data []byte) {
+				res = append(res, string(data))
 			})
-			assert.Equal(i, 100)
+			assert.Equal(len(res), ls.RangeCount(start, stop))
+			return
 		}
+
+		assert.Equal(rangeFn(0, -1), []string{"0", "1", "2", "3", "4"})
+		assert.Equal(rangeFn(-100, 100), []string{"0", "1", "2", "3", "4"})
+		assert.Equal(rangeFn(1, 3), []string{"1", "2", "3"})
+		assert.Equal(rangeFn(-3, -1), []string{"2", "3", "4"})
+		assert.Equal(rangeFn(3, 3), []string{"3"})
+		assert.Equal(rangeFn(-3, 2), []string{"2"})
+
+		// empty
+		var nilStrings []string
+		assert.Equal(rangeFn(99, 100), nilStrings)
+		assert.Equal(rangeFn(-100, -99), nilStrings)
+		assert.Equal(rangeFn(-1, -3), nilStrings)
+		assert.Equal(rangeFn(3, 2), nilStrings)
 	})
 
-	t.Run("revrange", func(t *testing.T) {
+	t.Run("range2", func(t *testing.T) {
 		ls := genList(0, N)
 		i := 0
-		ls.RevRange(0, N, func(data []byte) {
-			assert.Equal(string(data), genKey(N-i-1))
+		ls.Range(0, N, func(data []byte) {
+			assert.Equal(string(data), genKey(i))
 			i++
 		})
 		assert.Equal(i, N)
 
 		for _, start := range []int{100, 1000, 5000} {
 			i = 0
-			ls.RevRange(start, start+100, func(data []byte) {
-				assert.Equal(string(data), genKey(N-start-i-1))
+			ls.Range(start, start+100, func(data []byte) {
+				assert.Equal(string(data), genKey(start+i))
 				i++
 			})
-			assert.Equal(i, 100)
+			assert.Equal(i, 101)
 		}
 	})
 }