Add ZDIFF command

Signed-off-by: Devansh Singh <[email protected]>

server/server.go

@@ -36,6 +36,7 @@ const (
 	// Sorted set commands
 	CMD_ZADD     = "ZADD"
 	CMD_ZCARD    = "ZCARD"
+	CMD_ZDIFF    = "ZDIFF"
 	CMD_ZREMOVE  = "ZREMOVE"
 	CMD_ZMEMBERS = "ZMEMBERS"
 )
@@ -148,6 +149,8 @@ func (s *Server) HandleCommand() {
 			s.zAdd(cmd)
 		case CMD_ZCARD:
 			s.zCard(cmd)
+		case CMD_ZDIFF:
+			s.zDiff(cmd)
 		case CMD_ZREMOVE:
 			s.zRemove(cmd)
 		case CMD_ZMEMBERS:

server/zset.go

@@ -52,3 +52,18 @@ func (s *Server) zRemove(cmd Command) {
 		return
 	}
 }
+
+func (s *Server) zDiff(cmd Command) {
+	if len(cmd.Args) < 2 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
+	elements, err := s.DB.ZDiff(cmd.Args[0], cmd.Args[1])
+	if err != nil {
+		cmd.error(err)
+		return
+	}
+	for index, element := range elements {
+		cmd.write(fmt.Sprintf("%d) %s", index+1, element))
+	}
+}

store/zset.go

@@ -1,6 +1,10 @@
 package store
 
-import "github.com/Devansh3712/tandb/zset"
+import (
+	"fmt"
+
+	"github.com/Devansh3712/tandb/zset"
+)
 
 func (s *Store) ZAdd(set, key string) {
 	s.Mutex.Lock()
@@ -46,3 +50,20 @@ func (s *Store) ZRemove(set, key string) error {
 	}
 	return value.Remove(key)
 }
+
+func (s *Store) ZDiff(s1, s2 string) ([]string, error) {
+	s.Mutex.RLock()
+	defer s.Mutex.RUnlock()
+
+	v1, ok := s.ZSets[s1]
+	if !ok {
+		return nil, fmt.Errorf("the set %s does not exist", s1)
+	}
+	v2, ok := s.ZSets[s2]
+	if !ok {
+		return nil, fmt.Errorf("the set %s does not exist", s2)
+	}
+
+	diff := v1.Difference(v2)
+	return diff.Members(), nil
+}