refactor project structure

datastruct/sortedset/sortedset.go

@@ -0,0 +1,226 @@
+package sortedset
+
+import (
+	"strconv"
+)
+
+type SortedSet struct {
+	dict     map[string]*Element
+	skiplist *skiplist
+}
+
+func Make() *SortedSet {
+	return &SortedSet{
+		dict:     make(map[string]*Element),
+		skiplist: makeSkiplist(),
+	}
+}
+
+/*
+ * return: has inserted new node
+ */
+func (sortedSet *SortedSet) Add(member string, score float64) bool {
+	element, ok := sortedSet.dict[member]
+	sortedSet.dict[member] = &Element{
+		Member: member,
+		Score:  score,
+	}
+	if ok {
+		if score != element.Score {
+			sortedSet.skiplist.remove(member, score)
+			sortedSet.skiplist.insert(member, score)
+		}
+		return false
+	} else {
+		sortedSet.skiplist.insert(member, score)
+		return true
+	}
+}
+
+func (sortedSet *SortedSet) Len() int64 {
+	return int64(len(sortedSet.dict))
+}
+
+func (sortedSet *SortedSet) Get(member string) (element *Element, ok bool) {
+	element, ok = sortedSet.dict[member]
+	if !ok {
+		return nil, false
+	}
+	return element, true
+}
+
+func (sortedSet *SortedSet) Remove(member string) bool {
+	v, ok := sortedSet.dict[member]
+	if ok {
+		sortedSet.skiplist.remove(member, v.Score)
+		delete(sortedSet.dict, member)
+		return true
+	}
+	return false
+}
+
+/**
+ * get 0-based rank
+ */
+func (sortedSet *SortedSet) GetRank(member string, desc bool) (rank int64) {
+	element, ok := sortedSet.dict[member]
+	if !ok {
+		return -1
+	}
+	r := sortedSet.skiplist.getRank(member, element.Score)
+	if desc {
+		r = sortedSet.skiplist.length - r
+	} else {
+		r--
+	}
+	return r
+}
+
+/**
+ * traverse [start, stop), 0-based rank
+ */
+func (sortedSet *SortedSet) ForEach(start int64, stop int64, desc bool, consumer func(element *Element) bool) {
+	size := int64(sortedSet.Len())
+	if start < 0 || start >= size {
+		panic("illegal start " + strconv.FormatInt(start, 10))
+	}
+	if stop < start || stop > size {
+		panic("illegal end " + strconv.FormatInt(stop, 10))
+	}
+
+	// find start node
+	var node *Node
+	if desc {
+		node = sortedSet.skiplist.tail
+		if start > 0 {
+			node = sortedSet.skiplist.getByRank(int64(size - start))
+		}
+	} else {
+		node = sortedSet.skiplist.header.level[0].forward
+		if start > 0 {
+			node = sortedSet.skiplist.getByRank(int64(start + 1))
+		}
+	}
+
+	sliceSize := int(stop - start)
+	for i := 0; i < sliceSize; i++ {
+		if !consumer(&node.Element) {
+			break
+		}
+		if desc {
+			node = node.backward
+		} else {
+			node = node.level[0].forward
+		}
+	}
+}
+
+/**
+ * return [start, stop), 0-based rank
+ * assert start in [0, size), stop in [start, size]
+ */
+func (sortedSet *SortedSet) Range(start int64, stop int64, desc bool) []*Element {
+	sliceSize := int(stop - start)
+	slice := make([]*Element, sliceSize)
+	i := 0
+	sortedSet.ForEach(start, stop, desc, func(element *Element) bool {
+		slice[i] = element
+		i++
+		return true
+	})
+	return slice
+}
+
+func (sortedSet *SortedSet) Count(min *ScoreBorder, max *ScoreBorder) int64 {
+	var i int64 = 0
+	// ascending order
+	sortedSet.ForEach(0, sortedSet.Len(), false, func(element *Element) bool {
+		gtMin := min.less(element.Score) // greater than min
+		if !gtMin {
+			// has not into range, continue foreach
+			return true
+		}
+		ltMax := max.greater(element.Score) // less than max
+		if !ltMax {
+			// break through score border, break foreach
+			return false
+		}
+		// gtMin && ltMax
+		i++
+		return true
+	})
+	return i
+}
+
+func (sortedSet *SortedSet) ForEachByScore(min *ScoreBorder, max *ScoreBorder, offset int64, limit int64, desc bool, consumer func(element *Element) bool) {
+	// find start node
+	var node *Node
+	if desc {
+		node = sortedSet.skiplist.getLastInScoreRange(min, max)
+	} else {
+		node = sortedSet.skiplist.getFirstInScoreRange(min, max)
+	}
+
+	for node != nil && offset > 0 {
+		if desc {
+			node = node.backward
+		} else {
+			node = node.level[0].forward
+		}
+		offset--
+	}
+
+	// A negative limit returns all elements from the offset
+	for i := 0; (i < int(limit) || limit < 0) && node != nil; i++ {
+		if !consumer(&node.Element) {
+			break
+		}
+		if desc {
+			node = node.backward
+		} else {
+			node = node.level[0].forward
+		}
+		if node == nil {
+			break
+		}
+		gtMin := min.less(node.Element.Score) // greater than min
+		ltMax := max.greater(node.Element.Score)
+		if !gtMin || !ltMax {
+			break // break through score border
+		}
+	}
+}
+
+/*
+ * param limit: <0 means no limit
+ */
+func (sortedSet *SortedSet) RangeByScore(min *ScoreBorder, max *ScoreBorder, offset int64, limit int64, desc bool) []*Element {
+	if limit == 0 || offset < 0 {
+		return make([]*Element, 0)
+	}
+	slice := make([]*Element, 0)
+	sortedSet.ForEachByScore(min, max, offset, limit, desc, func(element *Element) bool {
+		slice = append(slice, element)
+		return true
+	})
+	return slice
+}
+
+func (sortedSet *SortedSet) RemoveByScore(min *ScoreBorder, max *ScoreBorder) int64 {
+	removed := sortedSet.skiplist.RemoveRangeByScore(min, max)
+	for _, element := range removed {
+		delete(sortedSet.dict, element.Member)
+	}
+	return int64(len(removed))
+}
+
+/*
+ * 0-based rank, [start, stop)
+ */
+func (sortedSet *SortedSet) RemoveByRank(start int64, stop int64) int64 {
+	removed := sortedSet.skiplist.RemoveRangeByRank(start+1, stop+1)
+	for _, element := range removed {
+		delete(sortedSet.dict, element.Member)
+	}
+	return int64(len(removed))
+}