t_zset.go

package main

import (
	"bytes"
	"fmt"
	"github.com/pengdafu/redis-golang/dict"
	"github.com/pengdafu/redis-golang/sds"
	"github.com/pengdafu/redis-golang/util"
	"github.com/pengdafu/redis-golang/ziplist"
	"math"
	"strconv"
	"unsafe"
)

type zskiplist struct {
	header, tail *zskiplistNode
	length       int64
	level        int
}
type zskiplistNode struct {
	ele      *sds.SDS
	score    float64
	backward *zskiplistNode
	level    []*zskiplistLevel
}
type zskiplistLevel struct {
	forward *zskiplistNode
	span    int64
}

type zset struct {
	dict *dict.Dict
	zsl  *zskiplist
}

func zaddCommand(c *Client) {
	zaddGenericCommand(c, zaddNone)
}

func zincrbyCommand(c *Client) {
	zaddGenericCommand(c, zaddIncr)
}

const (
	zaddNone    = 0
	zaddIncr    = 1 << 0
	zaddNx      = 1 << 1
	zaddXX      = 1 << 2
	zaddNop     = 1 << 3
	zaddNan     = 1 << 4
	zaddAdded   = 1 << 5
	zaddUpdated = 1 << 6
	zaddCh      = 1 << 16
)

func zaddGenericCommand(c *Client, flags int) {
	nanErr := "resulting score is not number (NaN)"
	key := c.argv[1]
	var zobj *robj
	var ele sds.SDS

	var score float64
	var scores []float64
	var added int
	var updated int
	var processed int

	scoreIdx := 2
	for scoreIdx < c.argc {
		opt := (*sds.SDS)(c.argv[scoreIdx].ptr).BufData(0)
		if util.StrCaseCmp(opt, "nx") {
			flags |= zaddNx
		} else if util.StrCaseCmp(opt, "xx") {
			flags |= zaddXX
		} else if util.StrCaseCmp(opt, "ch") {
			flags |= zaddCh
		} else if util.StrCaseCmp(opt, "incr") {
			flags |= zaddIncr
		} else {
			break
		}
		scoreIdx++
	}

	incr := (flags & zaddIncr) > 0
	nx := (flags & zaddNx) > 0
	xx := (flags & zaddXX) > 0
	ch := (flags & zaddCh) > 0
	elements := c.argc - scoreIdx
	if elements%2 != 0 || elements == 0 {
		addReply(c, shared.syntaxErr)
		return
	}

	elements /= 2

	if nx && xx {
		addReplyError(c, "XX and NX options at the same time are not compatible")
		return
	}
	if incr && elements > 1 {
		addReplyError(c,
			"INCR option supports a single increment-element pair")
		return
	}

	scores = make([]float64, elements)
	for j := 0; j < elements; j++ {
		if !getDoubleFromObjectOrReply(c, c.argv[scoreIdx+j*2], &scores[j], "") {
			goto cleanUp
		}
	}

	zobj = c.db.lookupKeyWrite(key)
	if zobj == nil {
		if xx {
			goto replyToClient
		}
		if server.zsetMaxZipListEntries == 0 ||
			server.zsetMaxZipListValue < sds.Len(*(*sds.SDS)(c.argv[scoreIdx+1].ptr)) {
			zobj = createZsetObject()
		} else {
			zobj = createZsetZiplistObject()
		}
		c.db.dbAdd(key, zobj)
	} else {
		if zobj.checkType(c, ObjSet) {
			goto cleanUp
		}
	}

	for j := 0; j < elements; j++ {
		var newscore float64
		score = scores[j]
		retFlags := flags

		ele = *(*sds.SDS)(c.argv[scoreIdx+1+j*2].ptr)
		retval := zsetAdd(zobj, score, ele, &retFlags, &newscore)
		if !retval {
			addReplyError(c, nanErr)
			goto cleanUp
		}
		if retFlags&zaddAdded > 0 {
			added++
		}
		if retFlags&zaddUpdated > 0 {
			updated++
		}
		if retFlags&zaddNop == 0 {
			processed++
		}
		score = newscore
	}
	server.dirty++

replyToClient:
	if incr {
		if processed > 0 {
			addReplyDouble(c, score)
		} else {
			addReplyNull(c)
		}
	} else {
		num := added
		if ch {
			num += updated
		}
		addReplyLongLong(c, num)
	}
cleanUp:
	if added > 0 || updated > 0 {
		signalModifiedKey(c, c.db, key)
		cmd := "zincr"
		if !incr {
			cmd = "zadd"
		}
		notifyKeySpaceEvent(notifyZset, cmd, key, c.db.id)
	}
	return
}

func zsetAdd(zobj *robj, score float64, ele sds.SDS, flags *int, newscore *float64) bool {
	incr := (*flags & zaddIncr) > 0
	nx := (*flags & zaddNx) > 0
	xx := (*flags & zaddXX) > 0

	*flags = 0
	var curScore float64
	if math.IsNaN(score) {
		*flags = zaddNan
		return false
	}

	if zobj.getEncoding() == ObjEncodingZipList {
		var eptr []byte
		zl := *(*[]byte)(zobj.ptr)
		if eptr = zzlFind(*(*[]byte)(zobj.ptr), ele, &curScore); eptr != nil {
			if nx {
				*flags |= zaddNop
				return true
			}

			if incr {
				score += curScore
				if math.IsNaN(score) {
					*flags |= zaddNan
					return false
				}
				if newscore != nil {
					*newscore = score
				}
			}
			if score != curScore {
				zl = zzlDelete(zl, eptr)
				zl = zzlInsert(zl, ele, score)
				*flags |= zaddUpdated
				zobj.ptr = unsafe.Pointer(&zl)
			}
			return true
		} else if !xx {
			if zzlLength(zl)+1 > server.zsetMaxZipListEntries ||
				sds.Len(ele) > server.zsetMaxZipListValue ||
				!ziplist.SafeToAdd(zl, sds.Len(ele)) {
				zsetConvert(zobj, ObjEncodingSkipList)
			} else {
				zl = zzlInsert(zl, ele, score)
				zobj.ptr = unsafe.Pointer(&zl)
				if newscore != nil {
					*newscore = score
				}
				*flags |= zaddAdded
				return true
			}
		} else {
			*flags |= zaddNop
			return true
		}
	}

	if zobj.getEncoding() == ObjEncodingSkipList {

	}
}

func zremCommand(c *Client) {

}

func zcountCommand(c *Client) {

}

func zcardCommand(c *Client) {

}

func zscoreCommand(c *Client) {

}

func zsetConvert(zobj *robj, encoding int) {

}

// zsl : zset skiplist

const (
	zskiplistMaxLevel = 32
	zskiplistP        = 0.25
)

func zslCreate() *zskiplist {
	zsl := new(zskiplist)
	zsl.level = 1
	zsl.length = 0
	zsl.header = zslCreateNode(zskiplistMaxLevel, 0, nil)
	for j := 0; j < zskiplistMaxLevel; j++ {
		zsl.header.level[j].forward = nil
		zsl.header.level[j].span = 0
	}
	zsl.header.backward = nil
	zsl.tail = nil
	return zsl
}

func zslCreateNode(level int, score float64, ele *sds.SDS) *zskiplistNode {
	zn := new(zskiplistNode)
	zn.level = make([]*zskiplistLevel, level)
	for i := 0; i < level; i++ {
		zn.level[i] = new(zskiplistLevel)
	}
	zn.score = score
	zn.ele = ele
	return zn
}

// zzl : zset ziplist
func zzlFind(zl []byte, ele sds.SDS, score *float64) []byte {
	eptr := ziplist.Index(zl, 0)
	var sptr []byte

	for eptr != nil {
		sptr = ziplist.Next(zl, eptr)

		if ziplist.Compare(eptr, ele.BufData(0), sds.Len(ele)) {
			if score != nil {
				*score = zzlGetScore(sptr)
			}
			return eptr
		}

		eptr = ziplist.Next(zl, sptr)
	}
	return nil
}

func zzlGetScore(sptr []byte) float64 {
	var vstr []byte
	var vlen int
	var vlong int64

	ziplist.Get(sptr, &vstr, &vlen, &vlong)

	if vstr != nil {
		score, err := strconv.ParseFloat(string(vstr[:vlen]), 64)
		if err != nil {
			return 0
		}
		return score
	} else {
		return float64(vlong)
	}
}

func zzlDelete(zl []byte, eptr []byte) []byte {
	p := eptr
	zl = ziplist.Delete(zl, &p)
	zl = ziplist.Delete(zl, &p)
	return zl
}

func zzlInsertAt(zl []byte, eptr []byte, ele sds.SDS, score float64) []byte {
	scoreBuf := fmt.Sprintf("%v", score)
	if eptr == nil {
		zl = ziplist.Push(zl, ele.BufData(0), ziplist.Tail)
		zl = ziplist.Push(zl, []byte(scoreBuf), ziplist.Tail)
	} else {
		offset := int(uintptr(unsafe.Pointer(&eptr[0])) - uintptr(unsafe.Pointer(&zl[0])))
		zl = ziplist.Insert(zl, eptr, ele.BufData(0))
		eptr = zl[offset:]

		sptr := ziplist.Next(zl, eptr)
		zl = ziplist.Insert(zl, sptr, []byte(scoreBuf))
	}
	return zl
}

func zzlInsert(zl []byte, ele sds.SDS, score float64) []byte {
	eptr := ziplist.Index(zl, 0)
	var sptr []byte

	for eptr != nil {
		sptr = ziplist.Next(zl, eptr)

		s := zzlGetScore(sptr)

		if s > score {
			zl = zzlInsertAt(zl, eptr, ele, score)
			break
		} else if score == s {
			if zzlCompareElements(eptr, ele.BufData(0), sds.Len(ele)) > 0 {
				zl = zzlInsertAt(zl, eptr, ele, score)
				break
			}
		}
		eptr = ziplist.Next(zl, sptr)
	}

	if eptr == nil {
		zl = zzlInsertAt(zl, nil, ele, score)
	}
	return zl
}

func zzlCompareElements(eptr, cstr []byte, clen int) int {
	var vstr []byte
	var vlen int
	var vlong int64
	ziplist.Get(eptr, &vstr, &vlen, &vlong)

	if vstr == nil {
		vbuf := fmt.Sprintf("%v", vlong)
		vstr = []byte(vbuf)
		vlen = len(vstr)
	}

	minLen := clen
	if vlen < clen {
		minLen = vlen
	}

	cmp := bytes.Compare(vstr[:minLen], cstr[:minLen])
	if cmp == 0 {
		return vlen - clen
	}
	return cmp
}

func zzlLength(zl []byte) int {
	return ziplist.Len(zl) / 2
}