sorted_set.go

sorted_set.go - Overview

This file implements a sorted set data structure using a B-tree for efficient storage and retrieval of members with associated scores.

Detailed Documentation

type Item struct

type Item struct {
	btree.Item
	Score  float64
	Member string
}

• Purpose: Represents a member of a sorted set.
  • Score: The score of the member.
  • Member: The member itself (a string).

func (a *Item) Less(b btree.Item) bool

func (a *Item) Less(b btree.Item) bool {
	other := b.(*Item)
	if a.Score != other.Score {
		return a.Score < other.Score
	}
	return a.Member < other.Member
}

• Purpose: Compares two Items based on their scores and members.
  • Parameters:
    • b: The item to compare with.
  • Returns: True if the item a is less than item b, false otherwise.

type Set struct

type Set struct {
	// tree is a btree that stores Items.
	tree *btree.BTree
	// memberMap is a map that stores members and their scores.
	memberMap map[string]float64
}

• Purpose: Represents a sorted set data structure.
  • tree: A B-tree to store the items (members and scores).
  • memberMap: A map to store the members and their associated scores for quick lookup.

func New() *Set

func New() *Set {
	return &Set{
		tree:      btree.New(2),
		memberMap: make(map[string]float64),
	}
}

• Purpose: Creates a new, empty sorted set.
  • Returns: A pointer to the newly created sorted set.

func FromObject(obj *object.Obj) (value *Set, err []byte)

func FromObject(obj *object.Obj) (value *Set, err []byte) {
	if err := object.AssertType(obj.Type, object.ObjTypeSortedSet); err != nil {
		return nil, err
	}
	value, ok := obj.Value.(*Set)
	if !ok {
		return nil, diceerrors.NewErrWithMessage("Invalid sorted set object")
	}
	return value, nil
}

• Purpose: Converts a generic object to a sorted set.
  • Parameters:
    • obj: The object to convert.
  • Returns:
    • value: A pointer to the sorted set.
    • err: An error if the object is not a sorted set.

func (ss *Set) Upsert(score float64, member string) bool

func (ss *Set) Upsert(score float64, member string) bool {
	existingScore, exists := ss.memberMap[member]

	if exists {
		oldItem := &Item{Score: existingScore, Member: member}
		ss.tree.Delete(oldItem)
	}

	item := &Item{Score: score, Member: member}
	ss.tree.ReplaceOrInsert(item)
	ss.memberMap[member] = score

	return !exists
}

• Purpose: Adds or updates a member with a score in the sorted set.
  • Parameters:
    • score: The score of the member.
    • member: The member to add or update.
  • Returns: True if the member was added, false if it was updated.

func (ss *Set) RankWithScore(member string, reverse bool) (rank int64, score float64)

func (ss *Set) RankWithScore(member string, reverse bool) (rank int64, score float64) {
	score, exists := ss.memberMap[member]
	if !exists {
		return -1, 0
	}

	rank = int64(0)
	ss.tree.Ascend(func(item btree.Item) bool {
		if item.(*Item).Member == member {
			return false
		}
		rank++
		return true
	})

	if reverse {
		rank = int64(len(ss.memberMap)) - rank - 1
	}

	return
}

• Purpose: Returns the rank and score of a member in the sorted set.
  • Parameters:
    • member: The member to get the rank and score for.
    • reverse: Whether to return the rank in reverse order.
  • Returns:
    • rank: The rank of the member.
    • score: The score of the member.

func (ss *Set) Remove(member string) bool

func (ss *Set) Remove(member string) bool {
	score, exists := ss.memberMap[member]
	if !exists {
		return false
	}

	item := &Item{Score