From 45a5f30cdd66f3889b03172bd83ba296bd6a48bf Mon Sep 17 00:00:00 2001
From: SALEH <salehwithDM@gmail.com>
Date: Thu, 13 Oct 2022 03:21:27 +0100
Subject: [PATCH] feat(server): Implement SMISMEMBER command (#377)

* feat(server): Implement SMISMEMBER command (#361)

Signed-off-by: ATM SALEH <saleh.cse08@gmail.com>
Signed-off-by: Roman Gershman <roman@dragonflydb.io>
Co-authored-by: Roman Gershman <roman@dragonflydb.io>
---
 CONTRIBUTORS.md               |  3 ++-
 src/server/set_family.cc      | 42 ++++++++++++++++++++++++++++++++++-
 src/server/set_family_test.cc | 30 +++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
index 31c4429cb..cfa4cc78a 100644
--- a/CONTRIBUTORS.md
+++ b/CONTRIBUTORS.md
@@ -8,4 +8,5 @@
 * **[Ryan Russell](https://github.com/ryanrussell)**
   * Docs & Code Readability
 * **[Ali-Akber Saifee](https://github.com/alisaifee)**
-* **[Elle Y](https://github.com/inohime)**
\ No newline at end of file
+* **[Elle Y](https://github.com/inohime)**
+* **[ATM SALEH](https://github.com/ATM-SALEH)**
diff --git a/src/server/set_family.cc b/src/server/set_family.cc
index a643319a5..16776609a 100644
--- a/src/server/set_family.cc
+++ b/src/server/set_family.cc
@@ -32,7 +32,6 @@ using absl::GetFlag;
 using ResultStringVec = vector<OpResult<StringVec>>;
 using ResultSetView = OpResult<absl::flat_hash_set<std::string_view>>;
 using SvArray = vector<std::string_view>;
-
 namespace {
 
 constexpr uint32_t kMaxIntSetEntries = 256;
@@ -283,6 +282,15 @@ bool IsInSet(const DbContext& db_context, const SetType& st, string_view member)
   }
 }
 
+void FindInSet(StringVec& memberships,
+		const DbContext& db_context, const SetType& st, 
+		const vector<string_view>& members) {
+  for (const auto& member : members) {
+    bool status = IsInSet(db_context, st, member);
+    memberships.emplace_back(to_string(status));
+  }
+}
+
 // Removes arg from result.
 void DiffStrSet(const DbContext& db_context, const SetType& st,
                 absl::flat_hash_set<string>* result) {
@@ -1020,6 +1028,37 @@ void SIsMember(CmdArgList args, ConnectionContext* cntx) {
   }
 }
 
+void SMIsMember(CmdArgList args, ConnectionContext* cntx) {
+  string_view key = ArgS(args, 1);
+
+  vector<string_view> vals(args.size() - 2);
+  for (size_t i = 2; i < args.size(); ++i) {
+    vals[i - 2] = ArgS(args, i);
+  }
+
+  StringVec memberships;
+  memberships.reserve(vals.size());
+
+  auto cb = [&](Transaction* t, EngineShard* shard) {
+    OpResult<PrimeIterator> find_res = shard->db_slice().Find(t->db_context(), key, OBJ_SET);
+    if (find_res) {
+      SetType st{find_res.value()->second.RObjPtr(), find_res.value()->second.Encoding()};
+      FindInSet(memberships, t->db_context(), st, vals);
+      return OpStatus::OK;
+    }
+    return find_res.status();
+  };
+
+  OpResult<void> result = cntx->transaction->ScheduleSingleHop(std::move(cb));
+  if (result == OpStatus::KEY_NOTFOUND) {
+    memberships.assign(vals.size(), "0");
+    return (*cntx)->SendStringArr(memberships);
+  } else if (result == OpStatus::OK) {
+    return (*cntx)->SendStringArr(memberships);
+  }
+  (*cntx)->SendError(result.status());
+}
+
 void SMove(CmdArgList args, ConnectionContext* cntx) {
   string_view src = ArgS(args, 1);
   string_view dest = ArgS(args, 2);
@@ -1464,6 +1503,7 @@ void SetFamily::Register(CommandRegistry* registry) {
             << CI{"SINTERSTORE", CO::WRITE | CO::DENYOOM, -3, 1, -1, 1}.HFUNC(SInterStore)
             << CI{"SMEMBERS", CO::READONLY, 2, 1, 1, 1}.HFUNC(SMembers)
             << CI{"SISMEMBER", CO::FAST | CO::READONLY, 3, 1, 1, 1}.HFUNC(SIsMember)
+            << CI{"SMISMEMBER", CO::READONLY, -3, 1, 1, 1}.HFUNC(SMIsMember)
             << CI{"SMOVE", CO::FAST | CO::WRITE, 4, 1, 2, 1}.HFUNC(SMove)
             << CI{"SREM", CO::WRITE | CO::FAST | CO::DENYOOM, -3, 1, 1, 1}.HFUNC(SRem)
             << CI{"SCARD", CO::READONLY | CO::FAST, 2, 1, 1, 1}.HFUNC(SCard)
diff --git a/src/server/set_family_test.cc b/src/server/set_family_test.cc
index 5d468cf88..dcbcc3bd4 100644
--- a/src/server/set_family_test.cc
+++ b/src/server/set_family_test.cc
@@ -132,6 +132,36 @@ TEST_F(SetFamilyTest, SPop) {
   EXPECT_THAT(resp.GetVec(), IsSubsetOf({"a", "b", "c"}));
 }
 
+TEST_F(SetFamilyTest, SMIsMember) {
+  Run({"sadd", "foo", "a"});
+  Run({"sadd", "foo", "b"});
+
+  auto resp = Run({"smismember", "foo"});
+  EXPECT_THAT(resp, ErrArg("wrong number of arguments"));
+  
+  resp = Run({"smismember", "foo1", "a", "b"});
+  ASSERT_THAT(resp, ArgType(RespExpr::ARRAY));
+  EXPECT_THAT(resp.GetVec(), ElementsAre("0", "0"));
+  
+  resp = Run({"smismember", "foo", "a", "c"});
+  ASSERT_THAT(resp, ArgType(RespExpr::ARRAY));
+  EXPECT_THAT(resp.GetVec(), ElementsAre("1", "0"));
+  
+  resp = Run({"smismember", "foo", "a", "b"});
+  ASSERT_THAT(resp, ArgType(RespExpr::ARRAY));
+  EXPECT_THAT(resp.GetVec(), ElementsAre("1", "1"));
+  
+  resp = Run({"smismember", "foo", "d", "e"});
+  ASSERT_THAT(resp, ArgType(RespExpr::ARRAY));
+  EXPECT_THAT(resp.GetVec(), ElementsAre("0", "0"));
+  
+  resp = Run({"smismember", "foo", "b"});
+  EXPECT_THAT(resp, "1");
+  
+  resp = Run({"smismember", "foo", "x"});
+  EXPECT_THAT(resp, "0");
+}
+
 TEST_F(SetFamilyTest, Empty) {
   auto resp = Run({"smembers", "x"});
   ASSERT_THAT(resp, ArrLen(0));