From 68b9c9b81f5a2d274c67d5b0455eac8b83ddde99 Mon Sep 17 00:00:00 2001
From: Sayan Nandan <nandansayan@outlook.com>
Date: Wed, 8 Sep 2021 04:40:30 -0700
Subject: [PATCH] Add listmap compatibility with `exists`

---
 server/src/actions/exists.rs   | 46 +++++++++++++++++++++-------------
 server/src/kvengine/listmap.rs | 12 +++++++++
 server/src/kvengine/mod.rs     | 23 +++++++++++++++++
 3 files changed, 63 insertions(+), 18 deletions(-)

diff --git a/server/src/actions/exists.rs b/server/src/actions/exists.rs
index 0bbe0e4c..f51413a5 100644
--- a/server/src/actions/exists.rs
+++ b/server/src/actions/exists.rs
@@ -27,7 +27,9 @@
 //! # `EXISTS` queries
 //! This module provides functions to work with `EXISTS` queries
 
+use crate::corestore::table::DataModel;
 use crate::dbnet::connection::prelude::*;
+use crate::kvengine::KVTable;
 use crate::queryengine::ActionIter;
 use crate::util::compiler;
 
@@ -36,24 +38,32 @@ action!(
     fn exists(handle: &Corestore, con: &'a mut T, act: ActionIter<'a>) {
         err_if_len_is!(act, con, eq 0);
         let mut how_many_of_them_exist = 0usize;
-        let kve = kve!(con, handle);
-        let encoding_is_okay = if kve.needs_key_encoding() {
-            true
-        } else {
-            let encoder = kve.get_key_encoder();
-            act.as_ref().all(|k| encoder.is_ok(k))
-        };
-        if compiler::likely(encoding_is_okay) {
-            {
-                act.for_each(|key| {
-                    if kve.exists_unchecked(key) {
-                        how_many_of_them_exist += 1;
-                    }
-                });
-            }
-            con.write_response(how_many_of_them_exist).await?;
-        } else {
-            conwrite!(con, groups::ENCODING_ERROR)?;
+        macro_rules! exists {
+            ($engine:expr) => {{
+                let encoding_is_okay = if $engine.kve_key_encoded() {
+                    let encoder = $engine.kve_get_key_encoder();
+                    act.as_ref().all(|k| encoder.is_ok(k))
+                } else {
+                    true
+                };
+                if compiler::likely(encoding_is_okay) {
+                    act.for_each(|key| {
+                        if $engine.kve_exists(key) {
+                            how_many_of_them_exist += 1;
+                        }
+                    });
+                    conwrite!(con, how_many_of_them_exist)?;
+                } else {
+                    compiler::cold_err(conwrite!(con, groups::ENCODING_ERROR))?;
+                }
+            }};
+        }
+        let tbl = get_tbl!(handle, con);
+        match tbl.get_model_ref() {
+            DataModel::KV(kve) => exists!(kve),
+            DataModel::KVExtListmap(kve) => exists!(kve),
+            #[allow(unreachable_patterns)]
+            _ => conwrite!(con, groups::WRONG_MODEL)?,
         }
         Ok(())
     }
diff --git a/server/src/kvengine/listmap.rs b/server/src/kvengine/listmap.rs
index 4c40e6c3..6ab17064 100644
--- a/server/src/kvengine/listmap.rs
+++ b/server/src/kvengine/listmap.rs
@@ -141,4 +141,16 @@ impl<'a> KVTable<'a, Coremap<Data, RwLock<Vec<Data>>>> for KVEListMap {
     {
         self.base.true_if_removed(input)
     }
+    fn kve_exists<Q: ?Sized + Eq + Hash>(&self, input: &Q) -> bool
+    where
+        Data: Borrow<Q>,
+    {
+        self.base.contains_key(input)
+    }
+    fn kve_keylen<Q: ?Sized + Eq + Hash>(&self, input: &Q) -> Option<usize>
+    where
+        Data: Borrow<Q>,
+    {
+        self.base.get(input).map(|v| v.key().len())
+    }
 }
diff --git a/server/src/kvengine/mod.rs b/server/src/kvengine/mod.rs
index 6cb948c6..f63bd833 100644
--- a/server/src/kvengine/mod.rs
+++ b/server/src/kvengine/mod.rs
@@ -53,15 +53,26 @@ pub trait KVTable<'a, T> {
     fn kve_payload_encoded(&self) -> bool;
     /// Get a reference to the inner table for a given KVE Table
     fn kve_inner_ref(&'a self) -> &'a T;
+    /// Remove a key from the KVE
     fn kve_remove<Q: ?Sized + Eq + Hash>(&self, input: &Q) -> bool
     where
         Data: Borrow<Q>;
+    /// Get the key encoder
     fn kve_get_key_encoder(&self) -> SingleEncoder {
         s_encoder_booled!(self.kve_key_encoded())
     }
+    /// Get the payload encoder
     fn kve_get_payload_encoder(&self) -> SingleEncoder {
         s_encoder_booled!(self.kve_payload_encoded())
     }
+    /// Check if the KVE contains a certain key
+    fn kve_exists<Q: ?Sized + Eq + Hash>(&self, input: &Q) -> bool
+    where
+        Data: Borrow<Q>;
+    /// Get the length of a certain key in the KVE
+    fn kve_keylen<Q: ?Sized + Eq + Hash>(&self, input: &Q) -> Option<usize>
+    where
+        Data: Borrow<Q>;
 }
 
 impl<'a> KVTable<'a, Coremap<Data, Data>> for KVEngine {
@@ -86,6 +97,18 @@ impl<'a> KVTable<'a, Coremap<Data, Data>> for KVEngine {
     {
         self.table.true_if_removed(input)
     }
+    fn kve_exists<Q: ?Sized + Eq + Hash>(&self, input: &Q) -> bool
+    where
+        Data: Borrow<Q>,
+    {
+        self.table.contains_key(input)
+    }
+    fn kve_keylen<Q: ?Sized + Eq + Hash>(&self, input: &Q) -> Option<usize>
+    where
+        Data: Borrow<Q>,
+    {
+        self.table.get(input).map(|v| v.key().len())
+    }
 }
 
 /// An arbitrary unicode/binary _double encoder_ for two byte slice inputs