```
feat: 更新战斗系统模型结构和Redis消息处理 - 引入gredis依赖用于Redis消息处理 - 将战斗相关的枚举和结构体从info包迁移到model包 - 更新战斗结束原因、攻击值等类型的引用路径 - 添加新的zset工具包到工作区 - 修改Redis消息处理逻辑以正确解析gredis.Message类型 - 在战斗控制器中统一使用model包下的类型定义
This commit is contained in:
昔念
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/gogf/gf/v2/database/gredis"
|
||||
"github.com/gogf/gf/v2/errors/gerror"
|
||||
"github.com/gogf/gf/v2/frame/g"
|
||||
"github.com/gogf/gf/v2/text/gstr"
|
||||
@@ -158,17 +159,32 @@ func ListenFunc(ctx g.Ctx) {
|
||||
|
||||
// 处理消息（保留原有业务逻辑）
|
||||
if data != nil {
|
||||
dataMap := data.MapStrStr()
|
||||
if dataMap["Kind"] == "subscribe" {
|
||||
dataMap, ok := data.Interface().(*gredis.Message)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if dataMap["Channel"] == subscribeTopic {
|
||||
Logger.Debug(ctx, "执行函数", "payload", dataMap["Payload"])
|
||||
err := RunFunc(ctx, dataMap["Payload"])
|
||||
// if dataMap. == "subscribe" {
|
||||
// continue
|
||||
// }
|
||||
if dataMap.Channel == subscribeTopic {
|
||||
Logger.Debug(ctx, "执行函数", "payload", dataMap.Payload)
|
||||
err := RunFunc(ctx, dataMap.Payload)
|
||||
if err != nil {
|
||||
Logger.Error(ctx, "执行函数失败", "payload", dataMap["Payload"], "error", err)
|
||||
Logger.Error(ctx, "执行函数失败", "payload", dataMap.Payload, "error", err)
|
||||
}
|
||||
}
|
||||
if dataMap.Channel == "sun:join:2458" {
|
||||
|
||||
println("收到sun:join:2458", dataMap.Payload)
|
||||
//universalClient, _ := g.Redis("cool").Client().(goredis.UniversalClient)
|
||||
zs.Add(1001, User{ID: 1001, JoinTime: time.Now().Unix(), Score: 1000})
|
||||
if zs.Length() > 2 {
|
||||
zs.FindNext(func(i User) bool { return i.ID > 1000 })
|
||||
//找到上一个,如果区间分数少于一定,
|
||||
//直接进行匹配
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
25
common/cool/user.go
Normal file
25
common/cool/user.go
Normal file
@@ -0,0 +1,25 @@
|
||||
package cool
|
||||
|
||||
import "github.com/liwnn/zset"
|
||||
|
||||
type User struct {
|
||||
JoinTime int64
|
||||
ID int
|
||||
Score int
|
||||
}
|
||||
|
||||
func (u User) Key() uint32 {
|
||||
return uint32(u.ID)
|
||||
}
|
||||
|
||||
// 如果分数不对的话,就按时间排序
|
||||
func (u User) Less(than User) bool {
|
||||
if u.Score == than.Score {
|
||||
return u.JoinTime < than.JoinTime
|
||||
}
|
||||
return u.Score < than.Score
|
||||
}
|
||||
|
||||
var zs = zset.New[uint32, User](func(a, b User) bool {
|
||||
return a.Less(b)
|
||||
})
|
||||
177
common/utils/zset/README.md
Normal file
177
common/utils/zset/README.md
Normal file
@@ -0,0 +1,177 @@
|
||||
# ZSet
|
||||
This Go package provides an implementation of sorted set in redis.
|
||||
|
||||
## Usage (go < 1.18)
|
||||
All you have to do is to implement a comparison `function Less(Item) bool` and a `function Key() string` for your Item which will be store in the zset, here are some examples.
|
||||
``` go
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/liwnn/zset"
|
||||
)
|
||||
|
||||
type User struct {
|
||||
Name string
|
||||
Score int
|
||||
}
|
||||
|
||||
func (u User) Key() string {
|
||||
return u.Name
|
||||
}
|
||||
|
||||
func (u User) Less(than zset.Item) bool {
|
||||
if u.Score == than.(User).Score {
|
||||
return u.Name < than.(User).Name
|
||||
}
|
||||
return u.Score < than.(User).Score
|
||||
}
|
||||
|
||||
func main() {
|
||||
zs := zset.New()
|
||||
|
||||
// Add
|
||||
zs.Add("Hurst", User{Name: "Hurst", Score: 88})
|
||||
zs.Add("Peek", User{Name: "Peek", Score: 100})
|
||||
zs.Add("Beaty", User{Name: "Beaty", Score: 66})
|
||||
|
||||
// Rank
|
||||
rank := zs.Rank("Hurst", true)
|
||||
fmt.Printf("Hurst's rank is %v\n", rank) // expected 2
|
||||
|
||||
// Range
|
||||
fmt.Println()
|
||||
fmt.Println("Range[0,3]:")
|
||||
zs.Range(0, 3, true, func(v zset.Item, rank int) bool {
|
||||
fmt.Printf("%v's rank is %v\n", v.(User).Key(), rank)
|
||||
return true
|
||||
})
|
||||
|
||||
// Range with Iterator
|
||||
fmt.Println()
|
||||
fmt.Println("Range[0,3] with Iterator:")
|
||||
for it := zs.RangeIterator(0, 3, true); it.Valid(); it.Next() {
|
||||
fmt.Printf("Ite: %v's rank is %v\n", it.Item().(User).Key(), it.Rank())
|
||||
}
|
||||
|
||||
// Range by score [88, 100]
|
||||
fmt.Println()
|
||||
fmt.Println("RangeByScore[88,100]:")
|
||||
zs.RangeByScore(func(i zset.Item) bool {
|
||||
return i.(User).Score >= 88
|
||||
}, func(i zset.Item) bool {
|
||||
return i.(User).Score <= 100
|
||||
}, true, func(i zset.Item, rank int) bool {
|
||||
fmt.Printf("%v's score[%v] rank is %v\n", i.(User).Key(), i.(User).Score, rank)
|
||||
return true
|
||||
})
|
||||
|
||||
// Remove
|
||||
zs.Remove("Peek")
|
||||
|
||||
// Rank
|
||||
fmt.Println()
|
||||
fmt.Println("After remove Peek:")
|
||||
rank = zs.Rank("Hurst", true)
|
||||
fmt.Printf("Hurst's rank is %v\n", rank) // expected 1
|
||||
}
|
||||
```
|
||||
Output:
|
||||
```
|
||||
Hurst's rank is 2
|
||||
|
||||
Range[0,3]:
|
||||
Peek's rank is 1
|
||||
Hurst's rank is 2
|
||||
Beaty's rank is 3
|
||||
|
||||
Range[0,3] with Iterator:
|
||||
Ite: Peek's rank is 1
|
||||
Ite: Hurst's rank is 2
|
||||
Ite: Beaty's rank is 3
|
||||
|
||||
RangeByScore[88,100]:
|
||||
Peek's score[100] rank is 1
|
||||
Hurst's score[88] rank is 2
|
||||
|
||||
After remove Peek:
|
||||
Hurst's rank is 1
|
||||
```
|
||||
## Usage (go >= 1.18)
|
||||
``` go
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/liwnn/zset"
|
||||
)
|
||||
|
||||
type User struct {
|
||||
Name string
|
||||
Score int
|
||||
}
|
||||
|
||||
func (u User) Key() string {
|
||||
return u.Name
|
||||
}
|
||||
|
||||
func (u User) Less(than User) bool {
|
||||
if u.Score == than.Score {
|
||||
return u.Name < than.Name
|
||||
}
|
||||
return u.Score < than.Score
|
||||
}
|
||||
|
||||
func main() {
|
||||
zs := zset.New[string, User](func(a, b User) bool {
|
||||
return a.Less(b)
|
||||
})
|
||||
|
||||
// Add
|
||||
zs.Add("Hurst", User{Name: "Hurst", Score: 88})
|
||||
zs.Add("Peek", User{Name: "Peek", Score: 100})
|
||||
zs.Add("Beaty", User{Name: "Beaty", Score: 66})
|
||||
|
||||
// Rank
|
||||
rank := zs.Rank("Hurst", true)
|
||||
fmt.Printf("Hurst's rank is %v\n", rank) // expected 2
|
||||
|
||||
// Range
|
||||
fmt.Println()
|
||||
fmt.Println("Range[0,3]:")
|
||||
zs.Range(0, 3, true, func(v User, rank int) bool {
|
||||
fmt.Printf("%v's rank is %v\n", v.Key(), rank)
|
||||
return true
|
||||
})
|
||||
|
||||
// Range with Iterator
|
||||
fmt.Println()
|
||||
fmt.Println("Range[0,3] with Iterator:")
|
||||
for it := zs.RangeIterator(0, 3, true); it.Valid(); it.Next() {
|
||||
fmt.Printf("Ite: %v's rank is %v\n", it.Item().Key(), it.Rank())
|
||||
}
|
||||
|
||||
// Range by score [88, 100]
|
||||
fmt.Println()
|
||||
fmt.Println("RangeByScore[88,100]:")
|
||||
zs.RangeByScore(func(i User) bool {
|
||||
return i.Score >= 88
|
||||
}, func(i User) bool {
|
||||
return i.Score <= 100
|
||||
}, true, func(i User, rank int) bool {
|
||||
fmt.Printf("%v's score[%v] rank is %v\n", i.Key(), i.Score, rank)
|
||||
return true
|
||||
})
|
||||
|
||||
// Remove
|
||||
zs.Remove("Peek")
|
||||
|
||||
// Rank
|
||||
fmt.Println()
|
||||
fmt.Println("After remove Peek:")
|
||||
rank = zs.Rank("Hurst", true)
|
||||
fmt.Printf("Hurst's rank is %v\n", rank) // expected 1
|
||||
}
|
||||
```
|
||||
3
common/utils/zset/go.mod
Normal file
3
common/utils/zset/go.mod
Normal file
@@ -0,0 +1,3 @@
|
||||
module github.com/liwnn/zset
|
||||
|
||||
go 1.18
|
||||
540
common/utils/zset/zset.go
Normal file
540
common/utils/zset/zset.go
Normal file
@@ -0,0 +1,540 @@
|
||||
//go:build !go1.18
|
||||
// +build !go1.18
|
||||
|
||||
// Package zset implements sorted set similar to redis zset.
|
||||
package zset
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"strconv"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultMaxLevel = 32 // (1/p)^MaxLevel >= maxNode
|
||||
DefaultP = 0.25 // SkipList P = 1/4
|
||||
|
||||
DefaultFreeListSize = 32
|
||||
)
|
||||
|
||||
// Item represents a single object in the set.
|
||||
type Item interface {
|
||||
// Less must provide a strict weak ordering
|
||||
Less(Item) bool
|
||||
}
|
||||
|
||||
// ItemIterator allows callers of Range* to iterate of the zset.
|
||||
// When this function returns false, iteration will stop.
|
||||
type ItemIterator func(i Item, rank int) bool
|
||||
|
||||
type skipListLevel struct {
|
||||
forward *node
|
||||
span int
|
||||
}
|
||||
|
||||
// node is an element of a skip list
|
||||
type node struct {
|
||||
item Item
|
||||
backward *node
|
||||
level []skipListLevel
|
||||
}
|
||||
|
||||
// FreeList represents a free list of set node.
|
||||
type FreeList struct {
|
||||
freelist []*node
|
||||
}
|
||||
|
||||
// NewFreeList creates a new free list.
|
||||
func NewFreeList(size int) *FreeList {
|
||||
return &FreeList{freelist: make([]*node, 0, size)}
|
||||
}
|
||||
|
||||
func (f *FreeList) newNode(lvl int) (n *node) {
|
||||
if len(f.freelist) == 0 {
|
||||
n = new(node)
|
||||
n.level = make([]skipListLevel, lvl)
|
||||
return
|
||||
}
|
||||
index := len(f.freelist) - 1
|
||||
n = f.freelist[index]
|
||||
f.freelist[index] = nil
|
||||
f.freelist = f.freelist[:index]
|
||||
|
||||
if cap(n.level) < lvl {
|
||||
n.level = make([]skipListLevel, lvl)
|
||||
} else {
|
||||
n.level = n.level[:lvl]
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (f *FreeList) freeNode(n *node) (out bool) {
|
||||
// for gc
|
||||
n.item = nil
|
||||
for j := 0; j < len(n.level); j++ {
|
||||
n.level[j] = skipListLevel{}
|
||||
}
|
||||
|
||||
if len(f.freelist) < cap(f.freelist) {
|
||||
f.freelist = append(f.freelist, n)
|
||||
out = true
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// skipList represents a skip list
|
||||
type skipList struct {
|
||||
header, tail *node
|
||||
length int
|
||||
level int // current level count
|
||||
maxLevel int
|
||||
freelist *FreeList
|
||||
random *rand.Rand
|
||||
}
|
||||
|
||||
// newSkipList creates a skip list
|
||||
func newSkipList(maxLevel int) *skipList {
|
||||
if maxLevel < DefaultMaxLevel {
|
||||
panic("maxLevel must < 32")
|
||||
}
|
||||
return &skipList{
|
||||
level: 1,
|
||||
header: &node{
|
||||
level: make([]skipListLevel, maxLevel),
|
||||
},
|
||||
maxLevel: maxLevel,
|
||||
freelist: NewFreeList(DefaultFreeListSize),
|
||||
random: rand.New(rand.NewSource(time.Now().UnixNano())),
|
||||
}
|
||||
}
|
||||
|
||||
// insert an item into the SkipList.
|
||||
func (sl *skipList) insert(item Item) *node {
|
||||
var update [DefaultMaxLevel]*node // [0...list.maxLevel)
|
||||
var rank [DefaultMaxLevel]int
|
||||
x := sl.header
|
||||
for i := sl.level - 1; i >= 0; i-- {
|
||||
if i == sl.level-1 {
|
||||
rank[i] = 0
|
||||
} else {
|
||||
rank[i] = rank[i+1]
|
||||
}
|
||||
for y := x.level[i].forward; y != nil && y.item.Less(item); y = x.level[i].forward {
|
||||
rank[i] += x.level[i].span
|
||||
x = y
|
||||
}
|
||||
update[i] = x
|
||||
}
|
||||
|
||||
lvl := sl.randomLevel()
|
||||
if lvl > sl.level {
|
||||
for i := sl.level; i < lvl; i++ {
|
||||
rank[i] = 0
|
||||
update[i] = sl.header
|
||||
update[i].level[i].span = sl.length
|
||||
}
|
||||
sl.level = lvl
|
||||
}
|
||||
|
||||
x = sl.freelist.newNode(lvl)
|
||||
x.item = item
|
||||
for i := 0; i < lvl; i++ {
|
||||
x.level[i].forward = update[i].level[i].forward
|
||||
update[i].level[i].forward = x
|
||||
|
||||
x.level[i].span = update[i].level[i].span - (rank[0] - rank[i])
|
||||
update[i].level[i].span = (rank[0] - rank[i]) + 1
|
||||
}
|
||||
|
||||
// increment span for untouched levels
|
||||
for i := lvl; i < sl.level; i++ {
|
||||
update[i].level[i].span++
|
||||
}
|
||||
|
||||
if update[0] == sl.header {
|
||||
x.backward = nil
|
||||
} else {
|
||||
x.backward = update[0]
|
||||
}
|
||||
if x.level[0].forward == nil {
|
||||
sl.tail = x
|
||||
} else {
|
||||
x.level[0].forward.backward = x
|
||||
}
|
||||
sl.length++
|
||||
return x
|
||||
}
|
||||
|
||||
// delete element
|
||||
func (sl *skipList) delete(n *node) Item {
|
||||
var preAlloc [DefaultMaxLevel]*node // [0...list.maxLevel)
|
||||
update := preAlloc[:sl.maxLevel]
|
||||
x := sl.header
|
||||
for i := sl.level - 1; i >= 0; i-- {
|
||||
for y := x.level[i].forward; y != nil && y.item.Less(n.item); y = x.level[i].forward {
|
||||
x = y
|
||||
}
|
||||
update[i] = x
|
||||
}
|
||||
x = x.level[0].forward
|
||||
if x != nil && !n.item.Less(x.item) {
|
||||
for i := 0; i < sl.level; i++ {
|
||||
if update[i].level[i].forward == x {
|
||||
update[i].level[i].span += x.level[i].span - 1
|
||||
update[i].level[i].forward = x.level[i].forward
|
||||
} else {
|
||||
update[i].level[i].span--
|
||||
}
|
||||
}
|
||||
for sl.level > 1 && sl.header.level[sl.level-1].forward == nil {
|
||||
sl.level--
|
||||
}
|
||||
if x.level[0].forward == nil {
|
||||
sl.tail = x.backward
|
||||
} else {
|
||||
x.level[0].forward.backward = x.backward
|
||||
}
|
||||
removeItem := x.item
|
||||
sl.freelist.freeNode(x)
|
||||
sl.length--
|
||||
return removeItem
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (sl *skipList) updateItem(node *node, item Item) bool {
|
||||
if (node.level[0].forward == nil || !node.level[0].forward.item.Less(item)) &&
|
||||
(node.backward == nil || !item.Less(node.backward.item)) {
|
||||
node.item = item
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// getRank find the rank for an element.
|
||||
// Returns 0 when the element cannot be found, rank otherwise.
|
||||
// Note that the rank is 1-based
|
||||
func (sl *skipList) getRank(item Item) int {
|
||||
var rank int
|
||||
x := sl.header
|
||||
for i := sl.level - 1; i >= 0; i-- {
|
||||
for y := x.level[i].forward; y != nil && !item.Less(y.item); y = x.level[i].forward {
|
||||
rank += x.level[i].span
|
||||
x = y
|
||||
}
|
||||
if x.item != nil && !x.item.Less(item) {
|
||||
return rank
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (sl *skipList) randomLevel() int {
|
||||
lvl := 1
|
||||
for lvl < sl.maxLevel && float32(sl.random.Uint32()&0xFFFF) < DefaultP*0xFFFF {
|
||||
lvl++
|
||||
}
|
||||
return lvl
|
||||
}
|
||||
|
||||
// Finds an element by its rank. The rank argument needs to be 1-based.
|
||||
func (sl *skipList) getNodeByRank(rank int) *node {
|
||||
var traversed int
|
||||
x := sl.header
|
||||
for i := sl.level - 1; i >= 0; i-- {
|
||||
for x.level[i].forward != nil && traversed+x.level[i].span <= rank {
|
||||
traversed += x.level[i].span
|
||||
x = x.level[i].forward
|
||||
}
|
||||
if traversed == rank {
|
||||
return x
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (sl *skipList) getMinNode() *node {
|
||||
return sl.header.level[0].forward
|
||||
}
|
||||
|
||||
func (sl *skipList) getMaxNode() *node {
|
||||
return sl.tail
|
||||
}
|
||||
|
||||
// return the first node greater and the node's 1-based rank.
|
||||
func (sl *skipList) findNext(greater func(i Item) bool) (*node, int) {
|
||||
x := sl.header
|
||||
var rank int
|
||||
for i := sl.level - 1; i >= 0; i-- {
|
||||
for y := x.level[i].forward; y != nil && !greater(y.item); y = x.level[i].forward {
|
||||
rank += x.level[i].span
|
||||
x = y
|
||||
}
|
||||
}
|
||||
return x.level[0].forward, rank + x.level[0].span
|
||||
}
|
||||
|
||||
// return the first node less and the node's 1-based rank.
|
||||
func (sl *skipList) findPrev(less func(i Item) bool) (*node, int) {
|
||||
var rank int
|
||||
x := sl.header
|
||||
for i := sl.level - 1; i >= 0; i-- {
|
||||
for y := x.level[i].forward; y != nil && less(y.item); y = x.level[i].forward {
|
||||
rank += x.level[i].span
|
||||
x = y
|
||||
}
|
||||
}
|
||||
return x, rank
|
||||
}
|
||||
|
||||
// ZSet set
|
||||
type ZSet struct {
|
||||
dict map[string]*node
|
||||
sl *skipList
|
||||
}
|
||||
|
||||
// New creates a new ZSet.
|
||||
func New() *ZSet {
|
||||
return &ZSet{
|
||||
dict: make(map[string]*node),
|
||||
sl: newSkipList(DefaultMaxLevel),
|
||||
}
|
||||
}
|
||||
|
||||
// Add a new element or update the score of an existing element. If an item already
|
||||
// exist, the removed item is returned. Otherwise, nil is returned.
|
||||
func (zs *ZSet) Add(key string, item Item) (removeItem Item) {
|
||||
if node := zs.dict[key]; node != nil {
|
||||
// if the node after update, would be still exactly at the same position,
|
||||
// we can just update item.
|
||||
if zs.sl.updateItem(node, item) {
|
||||
return
|
||||
}
|
||||
removeItem = zs.sl.delete(node)
|
||||
}
|
||||
zs.dict[key] = zs.sl.insert(item)
|
||||
return
|
||||
}
|
||||
|
||||
// Remove the element 'ele' from the sorted set,
|
||||
// return true if the element existed and was deleted, false otherwise
|
||||
func (zs *ZSet) Remove(key string) (removeItem Item) {
|
||||
node := zs.dict[key]
|
||||
if node == nil {
|
||||
return nil
|
||||
}
|
||||
removeItem = zs.sl.delete(node)
|
||||
delete(zs.dict, key)
|
||||
return
|
||||
}
|
||||
|
||||
// Rank return 1-based rank or 0 if not exist
|
||||
func (zs *ZSet) Rank(key string, reverse bool) int {
|
||||
node := zs.dict[key]
|
||||
if node != nil {
|
||||
rank := zs.sl.getRank(node.item)
|
||||
if rank > 0 {
|
||||
if reverse {
|
||||
return zs.sl.length - rank + 1
|
||||
}
|
||||
return rank
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (zs *ZSet) FindNext(iGreaterThan func(i Item) bool) (v Item, rank int) {
|
||||
n, rank := zs.sl.findNext(iGreaterThan)
|
||||
if n == nil {
|
||||
return
|
||||
}
|
||||
return n.item, rank
|
||||
}
|
||||
|
||||
func (zs *ZSet) FindPrev(iLessThan func(i Item) bool) (v Item, rank int) {
|
||||
n, rank := zs.sl.findPrev(iLessThan)
|
||||
if n == nil {
|
||||
return
|
||||
}
|
||||
return n.item, rank
|
||||
}
|
||||
|
||||
// RangeByScore calls the iterator for every value within the range [min, max],
|
||||
// until iterator return false. If min is nil, it represents negative infinity.
|
||||
// If max is nil, it represents positive infinity.
|
||||
func (zs *ZSet) RangeByScore(min, max func(i Item) bool, reverse bool, iterator ItemIterator) {
|
||||
llen := zs.sl.length
|
||||
var minNode, maxNode *node
|
||||
var minRank, maxRank int
|
||||
if min == nil {
|
||||
minNode = zs.sl.getMinNode()
|
||||
minRank = 1
|
||||
} else {
|
||||
minNode, minRank = zs.sl.findNext(min)
|
||||
}
|
||||
if minNode == nil {
|
||||
return
|
||||
}
|
||||
if max == nil {
|
||||
maxNode = zs.sl.getMaxNode()
|
||||
maxRank = llen
|
||||
} else {
|
||||
maxNode, maxRank = zs.sl.findPrev(max)
|
||||
}
|
||||
if maxNode == nil {
|
||||
return
|
||||
}
|
||||
if reverse {
|
||||
n := maxNode
|
||||
for i := maxRank; i >= minRank; i-- {
|
||||
if iterator(n.item, llen-i+1) {
|
||||
n = n.backward
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
} else {
|
||||
n := minNode
|
||||
for i := minRank; i <= maxRank; i++ {
|
||||
if iterator(n.item, i) {
|
||||
n = n.level[0].forward
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Range calls the iterator for every value with in index range [start, end],
|
||||
// until iterator return false. The <start> and <stop> arguments represent
|
||||
// zero-based indexes.
|
||||
func (zs *ZSet) Range(start, end int, reverse bool, iterator ItemIterator) {
|
||||
llen := zs.sl.length
|
||||
if start < 0 {
|
||||
start = llen + start
|
||||
}
|
||||
if end < 0 {
|
||||
end = llen + end
|
||||
}
|
||||
if start < 0 {
|
||||
start = 0
|
||||
}
|
||||
if start > end || start >= llen {
|
||||
return
|
||||
}
|
||||
if end >= llen {
|
||||
end = llen - 1
|
||||
}
|
||||
|
||||
rangeLen := end - start + 1
|
||||
if reverse {
|
||||
ln := zs.sl.getNodeByRank(llen - start)
|
||||
for i := 1; i <= rangeLen; i++ {
|
||||
if iterator(ln.item, start+i) {
|
||||
ln = ln.backward
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
} else {
|
||||
ln := zs.sl.getNodeByRank(start + 1)
|
||||
for i := 1; i <= rangeLen; i++ {
|
||||
if iterator(ln.item, start+i) {
|
||||
ln = ln.level[0].forward
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type RangeIterator struct {
|
||||
node *node
|
||||
start, end, cur int
|
||||
reverse bool
|
||||
}
|
||||
|
||||
func (r *RangeIterator) Len() int {
|
||||
return r.end - r.start + 1
|
||||
}
|
||||
|
||||
func (r *RangeIterator) Valid() bool {
|
||||
return r.cur <= r.end
|
||||
}
|
||||
|
||||
func (r *RangeIterator) Next() {
|
||||
if r.reverse {
|
||||
r.node = r.node.backward
|
||||
} else {
|
||||
r.node = r.node.level[0].forward
|
||||
}
|
||||
r.cur++
|
||||
}
|
||||
|
||||
func (r *RangeIterator) Item() Item {
|
||||
return r.node.item
|
||||
}
|
||||
|
||||
func (r *RangeIterator) Rank() int {
|
||||
return r.cur + 1
|
||||
}
|
||||
|
||||
// RangeIterator return iterator for visit elements in [start, end].
|
||||
// It is slower than Range.
|
||||
func (zs *ZSet) RangeIterator(start, end int, reverse bool) RangeIterator {
|
||||
llen := zs.sl.length
|
||||
if start < 0 {
|
||||
start = llen + start
|
||||
}
|
||||
if end < 0 {
|
||||
end = llen + end
|
||||
}
|
||||
if start < 0 {
|
||||
start = 0
|
||||
}
|
||||
|
||||
if start > end || start >= llen {
|
||||
return RangeIterator{end: -1}
|
||||
}
|
||||
|
||||
if end >= llen {
|
||||
end = llen - 1
|
||||
}
|
||||
|
||||
var n *node
|
||||
if reverse {
|
||||
n = zs.sl.getNodeByRank(llen - start)
|
||||
} else {
|
||||
n = zs.sl.getNodeByRank(start + 1)
|
||||
}
|
||||
return RangeIterator{
|
||||
start: start,
|
||||
cur: start,
|
||||
end: end,
|
||||
node: n,
|
||||
reverse: reverse,
|
||||
}
|
||||
}
|
||||
|
||||
// Get return Item in dict.
|
||||
func (zs *ZSet) Get(key string) Item {
|
||||
if node, ok := zs.dict[key]; ok {
|
||||
return node.item
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Length return the element count
|
||||
func (zs *ZSet) Length() int {
|
||||
return zs.sl.length
|
||||
}
|
||||
|
||||
type Int int
|
||||
|
||||
func (a Int) Key() string {
|
||||
return strconv.Itoa(int(a))
|
||||
}
|
||||
|
||||
func (a Int) Less(b Item) bool {
|
||||
return a < b.(Int)
|
||||
}
|
||||
529
common/utils/zset/zset_generic.go
Normal file
529
common/utils/zset/zset_generic.go
Normal file
@@ -0,0 +1,529 @@
|
||||
//go:build go1.18
|
||||
|
||||
// Package zset implements sorted set similar to redis zset.
|
||||
package zset
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultMaxLevel = 32 // (1/p)^MaxLevel >= maxNode
|
||||
DefaultP = 0.25 // SkipList P = 1/4
|
||||
|
||||
DefaultFreeListSize = 32
|
||||
)
|
||||
|
||||
// ItemIterator allows callers of Range* to iterate of the zset.
|
||||
// When this function returns false, iteration will stop.
|
||||
type ItemIterator[T any] func(i T, rank int) bool
|
||||
|
||||
type skipListLevel[T any] struct {
|
||||
forward *node[T]
|
||||
span int
|
||||
}
|
||||
|
||||
// node is an element of a skip list
|
||||
type node[T any] struct {
|
||||
item T
|
||||
backward *node[T]
|
||||
level []skipListLevel[T]
|
||||
}
|
||||
|
||||
// FreeList represents a free list of set node.
|
||||
type FreeList[T any] struct {
|
||||
freelist []*node[T]
|
||||
}
|
||||
|
||||
// NewFreeList creates a new free list.
|
||||
func NewFreeList[T any](size int) *FreeList[T] {
|
||||
return &FreeList[T]{freelist: make([]*node[T], 0, size)}
|
||||
}
|
||||
|
||||
func (f *FreeList[T]) newNode(lvl int) (n *node[T]) {
|
||||
if len(f.freelist) == 0 {
|
||||
n = new(node[T])
|
||||
n.level = make([]skipListLevel[T], lvl)
|
||||
return
|
||||
}
|
||||
index := len(f.freelist) - 1
|
||||
n = f.freelist[index]
|
||||
f.freelist[index] = nil
|
||||
f.freelist = f.freelist[:index]
|
||||
|
||||
if cap(n.level) < lvl {
|
||||
n.level = make([]skipListLevel[T], lvl)
|
||||
} else {
|
||||
n.level = n.level[:lvl]
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (f *FreeList[T]) freeNode(n *node[T]) (out bool) {
|
||||
// for gc
|
||||
var zero T
|
||||
n.item = zero
|
||||
for j := 0; j < len(n.level); j++ {
|
||||
n.level[j] = skipListLevel[T]{}
|
||||
}
|
||||
|
||||
if len(f.freelist) < cap(f.freelist) {
|
||||
f.freelist = append(f.freelist, n)
|
||||
out = true
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// skipList represents a skip list
|
||||
type skipList[T any] struct {
|
||||
header, tail *node[T]
|
||||
length int
|
||||
level int // current level count
|
||||
maxLevel int
|
||||
freelist *FreeList[T]
|
||||
random *rand.Rand
|
||||
less LessFunc[T]
|
||||
}
|
||||
|
||||
// newSkipList creates a skip list
|
||||
func newSkipList[T any](maxLevel int, less LessFunc[T]) *skipList[T] {
|
||||
if maxLevel < DefaultMaxLevel {
|
||||
panic("maxLevel must < 32")
|
||||
}
|
||||
return &skipList[T]{
|
||||
level: 1,
|
||||
header: &node[T]{
|
||||
level: make([]skipListLevel[T], maxLevel),
|
||||
},
|
||||
maxLevel: maxLevel,
|
||||
freelist: NewFreeList[T](DefaultFreeListSize),
|
||||
random: rand.New(rand.NewSource(time.Now().UnixNano())),
|
||||
less: less,
|
||||
}
|
||||
}
|
||||
|
||||
// insert an item into the SkipList.
|
||||
func (sl *skipList[T]) insert(item T) *node[T] {
|
||||
var update [DefaultMaxLevel]*node[T] // [0...list.maxLevel)
|
||||
var rank [DefaultMaxLevel]int
|
||||
x := sl.header
|
||||
for i := sl.level - 1; i >= 0; i-- {
|
||||
if i == sl.level-1 {
|
||||
rank[i] = 0
|
||||
} else {
|
||||
rank[i] = rank[i+1]
|
||||
}
|
||||
for y := x.level[i].forward; y != nil && sl.less(y.item, item); y = x.level[i].forward {
|
||||
rank[i] += x.level[i].span
|
||||
x = y
|
||||
}
|
||||
update[i] = x
|
||||
}
|
||||
|
||||
lvl := sl.randomLevel()
|
||||
if lvl > sl.level {
|
||||
for i := sl.level; i < lvl; i++ {
|
||||
rank[i] = 0
|
||||
update[i] = sl.header
|
||||
update[i].level[i].span = sl.length
|
||||
}
|
||||
sl.level = lvl
|
||||
}
|
||||
|
||||
x = sl.freelist.newNode(lvl)
|
||||
x.item = item
|
||||
for i := 0; i < lvl; i++ {
|
||||
x.level[i].forward = update[i].level[i].forward
|
||||
update[i].level[i].forward = x
|
||||
|
||||
x.level[i].span = update[i].level[i].span - (rank[0] - rank[i])
|
||||
update[i].level[i].span = (rank[0] - rank[i]) + 1
|
||||
}
|
||||
|
||||
// increment span for untouched levels
|
||||
for i := lvl; i < sl.level; i++ {
|
||||
update[i].level[i].span++
|
||||
}
|
||||
|
||||
if update[0] == sl.header {
|
||||
x.backward = nil
|
||||
} else {
|
||||
x.backward = update[0]
|
||||
}
|
||||
if x.level[0].forward == nil {
|
||||
sl.tail = x
|
||||
} else {
|
||||
x.level[0].forward.backward = x
|
||||
}
|
||||
sl.length++
|
||||
return x
|
||||
}
|
||||
|
||||
// delete element
|
||||
func (sl *skipList[T]) delete(n *node[T]) (_ T) {
|
||||
var preAlloc [DefaultMaxLevel]*node[T] // [0...list.maxLevel)
|
||||
update := preAlloc[:sl.maxLevel]
|
||||
x := sl.header
|
||||
for i := sl.level - 1; i >= 0; i-- {
|
||||
for y := x.level[i].forward; y != nil && sl.less(y.item, n.item); y = x.level[i].forward {
|
||||
x = y
|
||||
}
|
||||
update[i] = x
|
||||
}
|
||||
x = x.level[0].forward
|
||||
if x != nil && !sl.less(n.item, x.item) {
|
||||
for i := 0; i < sl.level; i++ {
|
||||
if update[i].level[i].forward == x {
|
||||
update[i].level[i].span += x.level[i].span - 1
|
||||
update[i].level[i].forward = x.level[i].forward
|
||||
} else {
|
||||
update[i].level[i].span--
|
||||
}
|
||||
}
|
||||
for sl.level > 1 && sl.header.level[sl.level-1].forward == nil {
|
||||
sl.level--
|
||||
}
|
||||
if x.level[0].forward == nil {
|
||||
sl.tail = x.backward
|
||||
} else {
|
||||
x.level[0].forward.backward = x.backward
|
||||
}
|
||||
removeItem := x.item
|
||||
sl.freelist.freeNode(x)
|
||||
sl.length--
|
||||
return removeItem
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (sl *skipList[T]) updateItem(node *node[T], item T) bool {
|
||||
if (node.level[0].forward == nil || !sl.less(node.level[0].forward.item, item)) &&
|
||||
(node.backward == nil || !sl.less(item, node.backward.item)) {
|
||||
node.item = item
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// getRank find the rank for an element.
|
||||
// Returns 0 when the element cannot be found, rank otherwise.
|
||||
// Note that the rank is 1-based
|
||||
func (sl *skipList[T]) getRank(item T) int {
|
||||
var rank int
|
||||
x := sl.header
|
||||
for i := sl.level - 1; i >= 0; i-- {
|
||||
for y := x.level[i].forward; y != nil && !sl.less(item, y.item); y = x.level[i].forward {
|
||||
rank += x.level[i].span
|
||||
x = y
|
||||
}
|
||||
if x != sl.header && !sl.less(x.item, item) {
|
||||
return rank
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (sl *skipList[T]) randomLevel() int {
|
||||
lvl := 1
|
||||
for lvl < sl.maxLevel && float32(sl.random.Uint32()&0xFFFF) < DefaultP*0xFFFF {
|
||||
lvl++
|
||||
}
|
||||
return lvl
|
||||
}
|
||||
|
||||
// Finds an element by its rank. The rank argument needs to be 1-based.
|
||||
func (sl *skipList[T]) getNodeByRank(rank int) *node[T] {
|
||||
var traversed int
|
||||
x := sl.header
|
||||
for i := sl.level - 1; i >= 0; i-- {
|
||||
for x.level[i].forward != nil && traversed+x.level[i].span <= rank {
|
||||
traversed += x.level[i].span
|
||||
x = x.level[i].forward
|
||||
}
|
||||
if traversed == rank {
|
||||
return x
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (sl *skipList[T]) getMinNode() *node[T] {
|
||||
return sl.header.level[0].forward
|
||||
}
|
||||
|
||||
func (sl *skipList[T]) getMaxNode() *node[T] {
|
||||
return sl.tail
|
||||
}
|
||||
|
||||
// return the first node greater and the node's 1-based rank.
|
||||
func (sl *skipList[T]) findNext(greater func(i T) bool) (*node[T], int) {
|
||||
x := sl.header
|
||||
var rank int
|
||||
for i := sl.level - 1; i >= 0; i-- {
|
||||
for y := x.level[i].forward; y != nil && !greater(y.item); y = x.level[i].forward {
|
||||
rank += x.level[i].span
|
||||
x = y
|
||||
}
|
||||
}
|
||||
return x.level[0].forward, rank + x.level[0].span
|
||||
}
|
||||
|
||||
// return the first node less and the node's 1-based rank.
|
||||
func (sl *skipList[T]) findPrev(less func(i T) bool) (*node[T], int) {
|
||||
var rank int
|
||||
x := sl.header
|
||||
for i := sl.level - 1; i >= 0; i-- {
|
||||
for y := x.level[i].forward; y != nil && less(y.item); y = x.level[i].forward {
|
||||
rank += x.level[i].span
|
||||
x = y
|
||||
}
|
||||
}
|
||||
return x, rank
|
||||
}
|
||||
|
||||
// ZSet set
|
||||
type ZSet[K comparable, T any] struct {
|
||||
dict map[K]*node[T]
|
||||
sl *skipList[T]
|
||||
}
|
||||
|
||||
// LessFunc determines how to order a type 'T'. It should implement a strict
|
||||
// ordering, and should return true if within that ordering, 'a' < 'b'.
|
||||
type LessFunc[T any] func(a, b T) bool
|
||||
|
||||
// New creates a new ZSet.
|
||||
func New[K comparable, T any](less LessFunc[T]) *ZSet[K, T] {
|
||||
return &ZSet[K, T]{
|
||||
dict: make(map[K]*node[T]),
|
||||
sl: newSkipList[T](DefaultMaxLevel, less),
|
||||
}
|
||||
}
|
||||
|
||||
// Add a new element or update the score of an existing element. If an item already
|
||||
// exist, the removed item is returned. Otherwise, nil is returned.
|
||||
func (zs *ZSet[K, T]) Add(key K, item T) (removeItem T) {
|
||||
if node := zs.dict[key]; node != nil {
|
||||
// if the node after update, would be still exactly at the same position,
|
||||
// we can just update item.
|
||||
if zs.sl.updateItem(node, item) {
|
||||
return
|
||||
}
|
||||
removeItem = zs.sl.delete(node)
|
||||
}
|
||||
zs.dict[key] = zs.sl.insert(item)
|
||||
return
|
||||
}
|
||||
|
||||
// Remove the element 'ele' from the sorted set,
|
||||
// return true if the element existed and was deleted, false otherwise
|
||||
func (zs *ZSet[K, T]) Remove(key K) (removeItem T) {
|
||||
node := zs.dict[key]
|
||||
if node == nil {
|
||||
return
|
||||
}
|
||||
removeItem = zs.sl.delete(node)
|
||||
delete(zs.dict, key)
|
||||
return
|
||||
}
|
||||
|
||||
// Rank return 1-based rank or 0 if not exist
|
||||
func (zs *ZSet[K, T]) Rank(key K, reverse bool) int {
|
||||
node := zs.dict[key]
|
||||
if node != nil {
|
||||
rank := zs.sl.getRank(node.item)
|
||||
if rank > 0 {
|
||||
if reverse {
|
||||
return zs.sl.length - rank + 1
|
||||
}
|
||||
return rank
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (zs *ZSet[K, T]) FindNext(iGreaterThan func(i T) bool) (v T, rank int) {
|
||||
n, rank := zs.sl.findNext(iGreaterThan)
|
||||
if n == nil {
|
||||
return
|
||||
}
|
||||
return n.item, rank
|
||||
}
|
||||
|
||||
func (zs *ZSet[K, T]) FindPrev(iLessThan func(i T) bool) (v T, rank int) {
|
||||
n, rank := zs.sl.findPrev(iLessThan)
|
||||
if n == nil {
|
||||
return
|
||||
}
|
||||
return n.item, rank
|
||||
}
|
||||
|
||||
// RangeByScore calls the iterator for every value within the range [min, max],
|
||||
// until iterator return false. If min is nil, it represents negative infinity.
|
||||
// If max is nil, it represents positive infinity.
|
||||
func (zs *ZSet[K, T]) RangeByScore(min, max func(i T) bool, reverse bool, iterator ItemIterator[T]) {
|
||||
llen := zs.sl.length
|
||||
var minNode, maxNode *node[T]
|
||||
var minRank, maxRank int
|
||||
if min == nil {
|
||||
minNode = zs.sl.getMinNode()
|
||||
minRank = 1
|
||||
} else {
|
||||
minNode, minRank = zs.sl.findNext(min)
|
||||
}
|
||||
if minNode == nil {
|
||||
return
|
||||
}
|
||||
if max == nil {
|
||||
maxNode = zs.sl.getMaxNode()
|
||||
maxRank = llen
|
||||
} else {
|
||||
maxNode, maxRank = zs.sl.findPrev(max)
|
||||
}
|
||||
if maxNode == nil {
|
||||
return
|
||||
}
|
||||
if reverse {
|
||||
n := maxNode
|
||||
for i := maxRank; i >= minRank; i-- {
|
||||
if iterator(n.item, llen-i+1) {
|
||||
n = n.backward
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
} else {
|
||||
n := minNode
|
||||
for i := minRank; i <= maxRank; i++ {
|
||||
if iterator(n.item, i) {
|
||||
n = n.level[0].forward
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Range calls the iterator for every value with in index range [start, end],
|
||||
// until iterator return false. The <start> and <stop> arguments represent
|
||||
// zero-based indexes.
|
||||
func (zs *ZSet[K, T]) Range(start, end int, reverse bool, iterator ItemIterator[T]) {
|
||||
llen := zs.sl.length
|
||||
if start < 0 {
|
||||
start = llen + start
|
||||
}
|
||||
if end < 0 {
|
||||
end = llen + end
|
||||
}
|
||||
if start < 0 {
|
||||
start = 0
|
||||
}
|
||||
if start > end || start >= llen {
|
||||
return
|
||||
}
|
||||
if end >= llen {
|
||||
end = llen - 1
|
||||
}
|
||||
|
||||
rangeLen := end - start + 1
|
||||
if reverse {
|
||||
ln := zs.sl.getNodeByRank(llen - start)
|
||||
for i := 1; i <= rangeLen; i++ {
|
||||
if iterator(ln.item, start+i) {
|
||||
ln = ln.backward
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
} else {
|
||||
ln := zs.sl.getNodeByRank(start + 1)
|
||||
for i := 1; i <= rangeLen; i++ {
|
||||
if iterator(ln.item, start+i) {
|
||||
ln = ln.level[0].forward
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type RangeIterator[T any] struct {
|
||||
node *node[T]
|
||||
start, end, cur int
|
||||
reverse bool
|
||||
}
|
||||
|
||||
func (r *RangeIterator[T]) Len() int {
|
||||
return r.end - r.start + 1
|
||||
}
|
||||
|
||||
func (r *RangeIterator[T]) Valid() bool {
|
||||
return r.cur <= r.end
|
||||
}
|
||||
|
||||
func (r *RangeIterator[T]) Next() {
|
||||
if r.reverse {
|
||||
r.node = r.node.backward
|
||||
} else {
|
||||
r.node = r.node.level[0].forward
|
||||
}
|
||||
r.cur++
|
||||
}
|
||||
|
||||
func (r *RangeIterator[T]) Item() T {
|
||||
return r.node.item
|
||||
}
|
||||
|
||||
func (r *RangeIterator[T]) Rank() int {
|
||||
return r.cur + 1
|
||||
}
|
||||
|
||||
// RangeIterator return iterator for visit elements in [start, end].
|
||||
// It is slower than Range.
|
||||
func (zs *ZSet[K, T]) RangeIterator(start, end int, reverse bool) RangeIterator[T] {
|
||||
llen := zs.sl.length
|
||||
if start < 0 {
|
||||
start = llen + start
|
||||
}
|
||||
if end < 0 {
|
||||
end = llen + end
|
||||
}
|
||||
if start < 0 {
|
||||
start = 0
|
||||
}
|
||||
|
||||
if start > end || start >= llen {
|
||||
return RangeIterator[T]{end: -1}
|
||||
}
|
||||
|
||||
if end >= llen {
|
||||
end = llen - 1
|
||||
}
|
||||
|
||||
var n *node[T]
|
||||
if reverse {
|
||||
n = zs.sl.getNodeByRank(llen - start)
|
||||
} else {
|
||||
n = zs.sl.getNodeByRank(start + 1)
|
||||
}
|
||||
return RangeIterator[T]{
|
||||
start: start,
|
||||
cur: start,
|
||||
end: end,
|
||||
node: n,
|
||||
reverse: reverse,
|
||||
}
|
||||
}
|
||||
|
||||
// Get return Item in dict.
|
||||
func (zs *ZSet[K, T]) Get(key K) (item T, found bool) {
|
||||
if n, ok := zs.dict[key]; ok {
|
||||
return n.item, ok
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Length return the element count
|
||||
func (zs *ZSet[K, T]) Length() int {
|
||||
return zs.sl.length
|
||||
}
|
||||
342
common/utils/zset/zset_generic_test.go
Normal file
342
common/utils/zset/zset_generic_test.go
Normal file
@@ -0,0 +1,342 @@
|
||||
//go:build go1.18
|
||||
|
||||
package zset
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"reflect"
|
||||
"strconv"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type TestRank struct {
|
||||
member string
|
||||
score int
|
||||
}
|
||||
|
||||
// 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[string, TestRank](func(a, b TestRank) bool {
|
||||
return a.score < b.score
|
||||
})
|
||||
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.member, false) != v.score+1 {
|
||||
t.Error("rank error")
|
||||
}
|
||||
if zs.Rank(v.member, true) != listSize-v.score {
|
||||
t.Error("rank error")
|
||||
}
|
||||
}
|
||||
|
||||
var r []TestRank
|
||||
zs.Range(0, 1, false, func(item TestRank, _ 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 TestRank) bool {
|
||||
return i.score >= 0
|
||||
}, func(i TestRank) bool {
|
||||
return i.score <= 1
|
||||
}, false, func(item TestRank, 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 TestRank, _ 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[string, TestRank](func(a, b TestRank) bool {
|
||||
return a.score < b.score
|
||||
})
|
||||
zs.RangeByScore(nil, nil, false, func(i TestRank, rank int) bool {
|
||||
return true
|
||||
})
|
||||
|
||||
for _, i := range perm(10) {
|
||||
zs.Add(i.member, i)
|
||||
}
|
||||
|
||||
var r []TestRank
|
||||
zs.RangeByScore(nil, nil, false, func(i TestRank, 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 TestRank) bool {
|
||||
return i.score >= 3
|
||||
}, func(i TestRank) bool {
|
||||
return i.score <= 5
|
||||
}, false, func(i TestRank, rank int) bool {
|
||||
r = append(r, i)
|
||||
return true
|
||||
})
|
||||
var expect []TestRank
|
||||
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 TestRank) bool {
|
||||
return i.score >= 3
|
||||
}, func(i TestRank) bool {
|
||||
return i.score <= 5
|
||||
}, true, func(i TestRank, 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[string, TestRank](func(a, b TestRank) bool {
|
||||
return a.score < b.score
|
||||
})
|
||||
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[string, TestRank](func(a, b TestRank) bool {
|
||||
return a.score < b.score
|
||||
})
|
||||
for _, item := range insertP {
|
||||
tr.Add(item.member, 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[string, TestRank](func(a, b TestRank) bool {
|
||||
return a.score < b.score
|
||||
})
|
||||
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[string, TestRank](func(a, b TestRank) bool {
|
||||
return a.score < b.score
|
||||
})
|
||||
for _, item := range insertP {
|
||||
tr.Add(item.member, item)
|
||||
}
|
||||
b.StartTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
tr.Remove(insertP[i%benchmarkListSize].member)
|
||||
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[string, TestRank](func(a, b TestRank) bool {
|
||||
return a.score < b.score
|
||||
})
|
||||
for _, item := range insertP {
|
||||
tr.Add(item.member, item)
|
||||
}
|
||||
b.StartTimer()
|
||||
for _, item := range removeP {
|
||||
tr.Remove(item.member)
|
||||
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[string, TestRank](func(a, b TestRank) bool {
|
||||
return a.score < b.score
|
||||
})
|
||||
for _, item := range insertP {
|
||||
tr.Add(item.member, item)
|
||||
}
|
||||
b.StartTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
tr.Rank(insertP[i%benchmarkListSize].member, true)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkRange(b *testing.B) {
|
||||
insertP := perm(benchmarkListSize)
|
||||
tr := New[string, TestRank](func(a, b TestRank) bool {
|
||||
return a.score < b.score
|
||||
})
|
||||
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 TestRank, rank int) bool {
|
||||
return true
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkRangeIterator(b *testing.B) {
|
||||
insertP := perm(benchmarkListSize)
|
||||
tr := New[string, TestRank](func(a, b TestRank) bool {
|
||||
return a.score < b.score
|
||||
})
|
||||
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[string, TestRank](func(a, b TestRank) bool {
|
||||
return a.score < b.score
|
||||
})
|
||||
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 TestRank) bool {
|
||||
return i.score >= minScore
|
||||
}, func(i TestRank) bool {
|
||||
return i.score <= maxScore
|
||||
}, true, func(i TestRank, rank int) bool {
|
||||
return true
|
||||
})
|
||||
}
|
||||
}
|
||||
329
common/utils/zset/zset_test.go
Normal file
329
common/utils/zset/zset_test.go
Normal file
@@ -0,0 +1,329 @@
|
||||
//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
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user