bl/common/utils/zset/zset_test.go

//go:build !go1.18
// +build !go1.18

package zset

import (
	"math/rand"
	"reflect"
	"strconv"
	"testing"
)

type TestRank struct {
	member string
	score  int
}

func (a TestRank) Key() string {
	return a.member
}

func (a TestRank) Less(than Item) bool {
	return a.score < than.(TestRank).score
}

// perm returns a random permutation of n Int items in the range [0, n).
func perm(n int) (out []TestRank) {
	out = make([]TestRank, 0, n)
	for _, v := range rand.Perm(n) {
		out = append(out, TestRank{
			member: strconv.Itoa(v),
			score:  v,
		})
	}
	return
}

// rang returns an ordered list of Int items in the range [0, n).
func rang(n int) (out []TestRank) {
	for i := 0; i < n; i++ {
		out = append(out, TestRank{
			member: strconv.Itoa(i),
			score:  i,
		})
	}
	return
}

func revrang(n int, count int) (out []TestRank) {
	for i := n - 1; i >= n-count; i-- {
		out = append(out, TestRank{
			member: strconv.Itoa(i),
			score:  i,
		})
	}
	return
}

func TestZSetRank(t *testing.T) {
	const listSize = 10000
	zs := New()
	for i := 0; i < 10; i++ {
		for _, v := range perm(listSize) {
			zs.Add(v.member, v)
		}
		for _, v := range perm(listSize) {
			if zs.Rank(v.Key(), false) != v.score+1 {
				t.Error("rank error")
			}
			if zs.Rank(v.Key(), true) != listSize-v.score {
				t.Error("rank error")
			}
		}

		var r []Item
		zs.Range(0, 1, false, func(item Item, _ int) bool {
			r = append(r, item)
			return true
		})
		if !reflect.DeepEqual(r, rang(2)) {
			t.Error("range error")
		}

		r = r[:0]
		zs.RangeByScore(func(i Item) bool {
			return i.(TestRank).score >= 0
		}, func(i Item) bool {
			return i.(TestRank).score <= 1
		}, false, func(item Item, rank int) bool {
			r = append(r, item)
			return true
		})
		if !reflect.DeepEqual(r, rang(2)) {
			t.Error("RangeItem error", r, rang(2))
		}

		r = r[:0]
		zs.Range(0, 1, true, func(item Item, _ int) bool {
			r = append(r, item)
			return true
		})
		if !reflect.DeepEqual(r, revrang(listSize, 2)) {
			t.Error("range error")
		}

		for i := 0; i < listSize/2; i++ {
			zs.Remove(strconv.Itoa(i))
		}
		for i := listSize + 1; i < listSize; i++ {
			if r := zs.Rank(strconv.Itoa(i), false); r != i-listSize/2 {
				t.Error("rank failed")
			}
		}
	}
}

func TestRangeItem(t *testing.T) {
	zs := New()
	zs.RangeByScore(nil, nil, false, func(i Item, rank int) bool {
		return true
	})

	for _, i := range perm(10) {
		zs.Add(i.member, i)
	}

	var r []Item
	zs.RangeByScore(nil, nil, false, func(i Item, rank int) bool {
		r = append(r, i)
		return true
	})
	if !reflect.DeepEqual(r, rang(10)) {
		t.Error("RangeItem error", r, rang(10))
	}

	r = r[:0]
	zs.RangeByScore(func(i Item) bool {
		return i.(TestRank).score >= 3
	}, func(i Item) bool {
		return i.(TestRank).score <= 5
	}, false, func(i Item, rank int) bool {
		r = append(r, i)
		return true
	})
	var expect []Item
	for i := 3; i <= 5; i++ {
		expect = append(expect, TestRank{
			member: strconv.Itoa(i),
			score:  i,
		})
	}
	if !reflect.DeepEqual(r, expect) {
		t.Error("RangeItem error", r, expect)
	}

	r = r[:0]
	zs.RangeByScore(func(i Item) bool {
		return i.(TestRank).score >= 3
	}, func(i Item) bool {
		return i.(TestRank).score <= 5
	}, true, func(i Item, rank int) bool {
		r = append(r, i)
		return true
	})
	expect = expect[:0]
	for i := 5; i >= 3; i-- {
		expect = append(expect, TestRank{
			member: strconv.Itoa(i),
			score:  i,
		})
	}
	if !reflect.DeepEqual(r, expect) {
		t.Error("RangeItem error", r, expect)
	}
}

const benchmarkListSize = 10000

func BenchmarkAdd(b *testing.B) {
	b.StopTimer()
	insertP := perm(benchmarkListSize)
	b.StartTimer()
	i := 0
	for i < b.N {
		tr := New()
		for _, item := range insertP {
			tr.Add(item.member, item)
			i++
			if i >= b.N {
				return
			}
		}
	}
}

func BenchmarkAddIncrease(b *testing.B) {
	b.StopTimer()
	insertP := rang(benchmarkListSize)
	b.StartTimer()
	i := 0
	for i < b.N {
		tr := New()
		for _, item := range insertP {
			tr.Add(item.Key(), item)
			i++
			if i >= b.N {
				return
			}
		}
	}
}

func BenchmarkAddDecrease(b *testing.B) {
	b.StopTimer()
	insertP := revrang(benchmarkListSize, benchmarkListSize)
	b.StartTimer()
	i := 0
	for i < b.N {
		tr := New()
		for _, item := range insertP {
			tr.Add(item.member, item)
			i++
			if i >= b.N {
				return
			}
		}
	}
}

func BenchmarkRemoveAdd(b *testing.B) {
	b.StopTimer()
	insertP := perm(benchmarkListSize)
	tr := New()
	for _, item := range insertP {
		tr.Add(item.member, item)
	}
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		tr.Remove(insertP[i%benchmarkListSize].Key())
		item := insertP[i%benchmarkListSize]
		tr.Add(item.member, item)
	}
}

func BenchmarkRemove(b *testing.B) {
	b.StopTimer()
	insertP := perm(benchmarkListSize)
	removeP := perm(benchmarkListSize)
	b.StartTimer()
	i := 0
	for i < b.N {
		b.StopTimer()
		tr := New()
		for _, v := range insertP {
			tr.Add(v.member, v)
		}
		b.StartTimer()
		for _, item := range removeP {
			tr.Remove(item.Key())
			i++
			if i >= b.N {
				return
			}
		}
		if tr.Length() > 0 {
			b.Error(tr.Length())
		}
	}
}

func BenchmarkRank(b *testing.B) {
	b.StopTimer()
	insertP := perm(benchmarkListSize)
	tr := New()
	for _, v := range insertP {
		tr.Add(v.member, v)
	}
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		tr.Rank(insertP[i%benchmarkListSize].Key(), true)
	}
}

func BenchmarkRange(b *testing.B) {
	insertP := perm(benchmarkListSize)
	tr := New()
	for _, item := range insertP {
		tr.Add(item.member, item)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		tr.Range(0, 100, true, func(i Item, rank int) bool {
			return true
		})
	}
}

func BenchmarkRangeIterator(b *testing.B) {
	insertP := perm(benchmarkListSize)
	tr := New()
	for _, item := range insertP {
		tr.Add(item.member, item)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		it := tr.RangeIterator(0, 100, true)
		for ; it.Valid(); it.Next() {
		}
	}
}

func BenchmarkRangeItem(b *testing.B) {
	insertP := perm(benchmarkListSize)
	tr := New()
	for _, item := range insertP {
		tr.Add(item.member, item)
	}
	minScore, maxScore := 0, 100
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		tr.RangeByScore(func(i Item) bool {
			return i.(TestRank).score >= minScore
		}, func(i Item) bool {
			return i.(TestRank).score <= maxScore
		}, true, func(i Item, rank int) bool {
			return true
		})
	}
}