Skip to content

Commit

Permalink
Add ZREMOVE, fix RB tree deletion
Browse files Browse the repository at this point in the history
Signed-off-by: Devansh Singh <[email protected]>
  • Loading branch information
Devansh3712 committed Mar 30, 2024
1 parent 73c13b5 commit 8e24e67
Show file tree
Hide file tree
Showing 4 changed files with 31 additions and 1 deletion.
3 changes: 3 additions & 0 deletions server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ const (
// Sorted set commands
CMD_ZADD = "ZADD"
CMD_ZCARD = "ZCARD"
CMD_ZREMOVE = "ZREMOVE"
CMD_ZMEMBERS = "ZMEMBERS"
)

Expand Down Expand Up @@ -147,6 +148,8 @@ func (s *Server) HandleCommand() {
s.zAdd(cmd)
case CMD_ZCARD:
s.zCard(cmd)
case CMD_ZREMOVE:
s.zRemove(cmd)
case CMD_ZMEMBERS:
s.zMembers(cmd)
default:
Expand Down
12 changes: 12 additions & 0 deletions server/zset.go
Original file line number Diff line number Diff line change
Expand Up @@ -40,3 +40,15 @@ func (s *Server) zCard(cmd Command) {
}
cmd.write(strconv.Itoa(size))
}

func (s *Server) zRemove(cmd Command) {
if len(cmd.Args) < 2 {
cmd.error(ErrNotEnoughArgs)
return
}
err := s.DB.ZRemove(cmd.Args[0], cmd.Args[1])
if err != nil {
cmd.error(err)
return
}
}
11 changes: 11 additions & 0 deletions store/zset.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,3 +35,14 @@ func (s *Store) ZCard(set string) (int, error) {
}
return value.Size(), nil
}

func (s *Store) ZRemove(set, key string) error {
s.Mutex.RLock()
defer s.Mutex.RUnlock()

value, ok := s.ZSets[set]
if !ok {
return ErrSetNotExists
}
return value.Remove(key)
}
6 changes: 5 additions & 1 deletion zset/rbtree.go
Original file line number Diff line number Diff line change
Expand Up @@ -271,14 +271,18 @@ func (t *RBTree) delete(value string) {
} else {
temp = t.successor(node)
}

var tchild *Node
if temp.Left != nil {
tchild = temp.Left
} else {
tchild = temp.Right
}
// If temp is leaf node, then tchild will be nil
if tchild != nil {
tchild.Parent = temp.Parent
}

tchild.Parent = temp.Parent
if temp.Parent == nil {
t.Root = tchild
} else if temp == temp.Parent.Left {
Expand Down

0 comments on commit 8e24e67

Please sign in to comment.