From 5238298ac473d4c16b18828109e0965739e1d253 Mon Sep 17 00:00:00 2001
From: Sayan Nandan <nandansayan@outlook.com>
Date: Tue, 7 Sep 2021 08:01:40 -0700
Subject: [PATCH] Add `VALUEAT` and `LIMIT` subactions to `lget`

---
 server/src/actions/lists/mod.rs  | 64 +++++++++++++++++++++++++++++---
 server/src/kvengine/listmap.rs   |  3 ++
 server/src/protocol/responses.rs |  2 +
 3 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/server/src/actions/lists/mod.rs b/server/src/actions/lists/mod.rs
index fb66505f..71854a42 100644
--- a/server/src/actions/lists/mod.rs
+++ b/server/src/actions/lists/mod.rs
@@ -29,9 +29,12 @@ use crate::corestore::Data;
 use crate::dbnet::connection::prelude::*;
 use crate::kvengine::listmap::LockedVec;
 use crate::kvengine::KVTable;
+use crate::resp::writer;
 use crate::resp::writer::TypedArrayWriter;
 
 const LEN: &[u8] = "LEN".as_bytes();
+const LIMIT: &[u8] = "LIMIT".as_bytes();
+const VALUEAT: &[u8] = "VALUEAT".as_bytes();
 
 macro_rules! listmap {
     ($tbl:expr, $con:expr) => {
@@ -42,6 +45,17 @@ macro_rules! listmap {
     };
 }
 
+macro_rules! writelist {
+    ($con:expr, $listmap:expr, $items:expr) => {
+        let mut typed_array_writer =
+            unsafe { TypedArrayWriter::new($con, $listmap.get_payload_tsymbol(), $items.len()) }
+                .await?;
+        for item in $items {
+            typed_array_writer.write_element(item).await?;
+        }
+    };
+}
+
 action! {
     /// Handle an `LSET` query for the list model
     /// Syntax: `LSET <listname> <values ...>`
@@ -80,6 +94,14 @@ action! {
         // get the list name
         let listname = unsafe { act.next_unchecked() };
         // now let us see what we need to do
+        macro_rules! get_numeric_count {
+            () => {
+                match unsafe { String::from_utf8_lossy(act.next_unchecked()) }.parse::<usize>() {
+                    Ok(int) => int,
+                    Err(_) => return conwrite!(con, groups::WRONGTYPE_ERR),
+                }
+            };
+        }
         match act.next_uppercase().as_ref() {
             None => {
                 // just return everything in the list
@@ -88,12 +110,7 @@ action! {
                 } else {
                     return conwrite!(con, groups::NIL);
                 };
-                let mut typed_array_writer = unsafe {
-                    TypedArrayWriter::new(con, listmap.get_payload_tsymbol(), items.len())
-                }.await?;
-                for item in items {
-                    typed_array_writer.write_element(item).await?;
-                }
+                writelist!(con, listmap, items);
             }
             Some(subaction) => {
                 match subaction.as_ref() {
@@ -104,6 +121,41 @@ action! {
                             conwrite!(con, groups::NIL)?;
                         }
                     }
+                    LIMIT => {
+                        err_if_len_is!(act, con, not 1);
+                        let count = get_numeric_count!();
+                        let items = if let Some(keys) = listmap.get_cloned(listname, count) {
+                            keys
+                        } else {
+                            return conwrite!(con, groups::NIL);
+                        };
+                        writelist!(con, listmap, items);
+                    }
+                    VALUEAT => {
+                        err_if_len_is!(act, con, not 1);
+                        let idx = get_numeric_count!();
+                        let maybe_value = listmap.get(listname).map(|list| {
+                            let readlist = list.read();
+                            let get = readlist.get(idx).cloned();
+                            get
+                        });
+                        match maybe_value {
+                            Some(Some(value)) => {
+                                unsafe {
+                                    // tsymbol is verified
+                                    writer::write_raw_mono(con, listmap.get_payload_tsymbol(), &value).await?;
+                                }
+                            },
+                            Some(None) => {
+                                // bad index
+                                conwrite!(con, groups::LISTMAP_BAD_INDEX)?;
+                            },
+                            None => {
+                                // not found
+                                conwrite!(con, groups::NIL)?;
+                            }
+                        }
+                    }
                     _ => conwrite!(con, groups::UNKNOWN_ACTION)?,
                 }
             }
diff --git a/server/src/kvengine/listmap.rs b/server/src/kvengine/listmap.rs
index fa920cf6..6e0eb1d1 100644
--- a/server/src/kvengine/listmap.rs
+++ b/server/src/kvengine/listmap.rs
@@ -91,6 +91,9 @@ impl KVEListMap {
         pub fn {borrow: Data} get(self: &Self, key: &Q) -> Option<Vecref<'_>> {
             self.base.get(key)
         }
+        pub fn {borrow: Data} get_cloned(self: &Self, key: &Q, count: usize) -> Option<Vec<Data>> {
+            self.base.get(key).map(|v| v.read().iter().take(count).cloned().collect())
+        }
     }
     /// Create and add a new list to the map
     pub fn add_list(&self, listname: Data) -> Option<bool> {
diff --git a/server/src/protocol/responses.rs b/server/src/protocol/responses.rs
index 47c3073d..8e437854 100644
--- a/server/src/protocol/responses.rs
+++ b/server/src/protocol/responses.rs
@@ -101,6 +101,8 @@ pub mod groups {
     pub const KEYSPACE_NOT_EMPTY: &[u8] = "!18\nkeyspace-not-empty\n".as_bytes();
     /// Bad type supplied in a DDL query for the key
     pub const BAD_TYPE_FOR_KEY: &[u8] = "!16\nbad-type-for-key\n".as_bytes();
+    /// The index for the provided list was non-existent
+    pub const LISTMAP_BAD_INDEX: &[u8] = "!9\nbad-list-index\n".as_bytes();
 }
 
 pub mod full_responses {