From 2352e266f356a0a84ed0a38c7eaeed1f0d7798b7 Mon Sep 17 00:00:00 2001
From: Akie Mejia <rfmejia@gmail.com>
Date: Thu, 6 May 2021 00:36:04 -0700
Subject: [PATCH] Implement `zDiff` and `zDiffWithScores` API (#367)

---
 .../main/scala/zio/redis/api/SortedSets.scala | 53 +++++++++++++++
 .../test/scala/zio/redis/SortedSetsSpec.scala | 65 +++++++++++++++++++
 2 files changed, 118 insertions(+)

diff --git a/redis/src/main/scala/zio/redis/api/SortedSets.scala b/redis/src/main/scala/zio/redis/api/SortedSets.scala
index 359560aef..69092034f 100644
--- a/redis/src/main/scala/zio/redis/api/SortedSets.scala
+++ b/redis/src/main/scala/zio/redis/api/SortedSets.scala
@@ -146,6 +146,58 @@ trait SortedSets {
     command.run((key, range))
   }
 
+  /**
+   * Subtract multiple sorted sets and return members.
+   *
+   * @param inputKeysNum Number of input keys
+   * @param key Key of a sorted set
+   * @param keys Keys of other sorted sets
+   * @return Chunk of differences between the first and successive input sorted sets.
+   */
+  final def zDiff[K: Schema, M: Schema](
+    inputKeysNum: Long,
+    key: K,
+    keys: K*
+  ): ZIO[RedisExecutor, RedisError, Chunk[M]] = {
+    val command =
+      RedisCommand(
+        ZDiff,
+        Tuple2(
+          LongInput,
+          NonEmptyList(ArbitraryInput[K]())
+        ),
+        ChunkOutput(ArbitraryOutput[M]())
+      )
+    command.run((inputKeysNum, (key, keys.toList)))
+  }
+
+  /**
+   * Subtract multiple sorted sets and return members and their associated score.
+   *
+   * @param inputKeysNum Number of input keys
+   * @param key Key of a sorted set
+   * @param keys Keys of other sorted sets
+   * @return Chunk of differences and scores between the first and successive input sorted sets.
+   */
+  final def zDiffWithScores[K: Schema, M: Schema](
+    inputKeysNum: Long,
+    key: K,
+    keys: K*
+  ): ZIO[RedisExecutor, RedisError, Chunk[MemberScore[M]]] = {
+    val command =
+      RedisCommand(
+        ZDiff,
+        Tuple3(
+          LongInput,
+          NonEmptyList(ArbitraryInput[K]()),
+          ArbitraryInput[String]()
+        ),
+        ChunkTuple2Output(ArbitraryOutput[M](), DoubleOutput)
+          .map(_.map { case (m, s) => MemberScore(s, m) })
+      )
+    command.run((inputKeysNum, (key, keys.toList), WithScores.stringify))
+  }
+
   /**
    * Subtract multiple sorted sets and store the resulting sorted set in a destination key.
    *
@@ -810,6 +862,7 @@ private[redis] object SortedSets {
   final val ZAdd             = "ZADD"
   final val ZCard            = "ZCARD"
   final val ZCount           = "ZCOUNT"
+  final val ZDiff            = "ZDIFF"
   final val ZDiffStore       = "ZDIFFSTORE"
   final val ZIncrBy          = "ZINCRBY"
   final val ZInter           = "ZINTER"
diff --git a/redis/src/test/scala/zio/redis/SortedSetsSpec.scala b/redis/src/test/scala/zio/redis/SortedSetsSpec.scala
index bd0b09d91..d418abc3f 100644
--- a/redis/src/test/scala/zio/redis/SortedSetsSpec.scala
+++ b/redis/src/test/scala/zio/redis/SortedSetsSpec.scala
@@ -213,6 +213,71 @@ trait SortedSetsSpec extends BaseSpec {
           } yield assert(count)(equalTo(0L))
         }
       ),
+      suite("zDiff")(
+        testM("empty sets") {
+          for {
+            key1 <- uuid
+            key2 <- uuid
+            key3 <- uuid
+            diff <- zDiff[String, String](3, key1, key2, key3)
+          } yield assert(diff)(isEmpty)
+        },
+        testM("non-empty set with empty set") {
+          for {
+            key1 <- uuid
+            key2 <- uuid
+            _ <- zAdd(key1)(
+                   MemberScore(1d, "a"),
+                   MemberScore(2d, "b")
+                 )
+            diff <- zDiff[String, String](2, key1, key2)
+          } yield assert(diff)(hasSameElements(Chunk("a", "b")))
+        },
+        testM("non-empty sets") {
+          for {
+            key1 <- uuid
+            key2 <- uuid
+            _ <- zAdd(key1)(
+                   MemberScore(1d, "a"),
+                   MemberScore(2d, "b"),
+                   MemberScore(3d, "c")
+                 )
+            _    <- zAdd(key2)(MemberScore(1d, "a"), MemberScore(2d, "b"))
+            diff <- zDiff[String, String](2, key1, key2)
+          } yield assert(diff)(hasSameElements(Chunk("c")))
+        }
+      ),
+      suite("zDiffWithScores")(
+        testM("empty sets") {
+          for {
+            key1 <- uuid
+            key2 <- uuid
+            key3 <- uuid
+            diff <- zDiffWithScores[String, String](3, key1, key2, key3)
+          } yield assert(diff)(isEmpty)
+        },
+        testM("non-empty set with empty set") {
+          for {
+            key1 <- uuid
+            key2 <- uuid
+            _    <- zAdd(key1)(MemberScore(1d, "a"), MemberScore(2d, "b"))
+            diff <- zDiffWithScores[String, String](2, key1, key2)
+          } yield assert(diff)(hasSameElements(Chunk(MemberScore(1d, "a"), MemberScore(2d, "b"))))
+        },
+        testM("non-empty sets") {
+          for {
+            key1 <- uuid
+            key2 <- uuid
+            _ <- zAdd(key1)(
+                   MemberScore(1d, "a"),
+                   MemberScore(2d, "b"),
+                   MemberScore(3d, "c")
+                 )
+            _    <- zAdd(key2)(MemberScore(1d, "a"), MemberScore(2d, "b"))
+            diff <- zDiffWithScores[String, String](2, key1, key2)
+          } yield assert(diff)(hasSameElements(Chunk(MemberScore(3d, "c"))))
+        }
+      ),
       suite("zDiffStore")(
         testM("empty sets") {
           for {