From 6dc8543a1e9ed03395e84c5290830f35b75c7e55 Mon Sep 17 00:00:00 2001
From: iko1 <me@remotecpp.dev>
Date: Wed, 22 Feb 2023 00:13:15 +0200
Subject: [PATCH] fix(server): fix json.arrtrim implementation (#844)

Signed-off-by: iko1 <me@remotecpp.dev>
---
 src/server/json_family.cc      | 29 ++++++++++++++++-------------
 src/server/json_family_test.cc | 25 +++++++++++++++++++------
 2 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/src/server/json_family.cc b/src/server/json_family.cc
index a11222b0d393..1061a87a1bb5 100644
--- a/src/server/json_family.cc
+++ b/src/server/json_family.cc
@@ -677,11 +677,16 @@ OpResult<vector<OptSizeT>> OpArrTrim(const OpArgs& op_args, string_view key, str
                                      int start_index, int stop_index) {
   vector<OptSizeT> vec;
   auto cb = [&](const string& path, JsonType& val) {
-    if (!val.is_array() || val.empty()) {
+    if (!val.is_array()) {
       vec.emplace_back(nullopt);
       return;
     }
 
+    if (val.empty()) {
+      vec.emplace_back(0);
+      return;
+    }
+
     size_t trim_start_index;
     if (start_index < 0) {
       trim_start_index = 0;
@@ -689,28 +694,26 @@ OpResult<vector<OptSizeT>> OpArrTrim(const OpArgs& op_args, string_view key, str
       trim_start_index = start_index;
     }
 
-    size_t trim_stop_index;
+    size_t trim_end_index;
     if ((size_t)stop_index >= val.size()) {
-      trim_stop_index = val.size();
+      trim_end_index = val.size();
     } else {
-      trim_stop_index = stop_index;
+      trim_end_index = stop_index;
     }
 
-    if (trim_start_index >= val.size() || trim_start_index > trim_stop_index) {
+    if (trim_start_index >= val.size() || trim_start_index > trim_end_index) {
       val.erase(val.array_range().begin(), val.array_range().end());
-      vec.emplace_back(val.size());
+      vec.emplace_back(0);
       return;
     }
 
-    auto it = std::next(val.array_range().begin(), trim_start_index);
-    while (it != val.array_range().end()) {
-      if (trim_start_index++ == trim_stop_index) {
-        break;
-      }
-
-      it = val.erase(it);
+    auto trim_start_it = std::next(val.array_range().begin(), trim_start_index);
+    auto trim_end_it = val.array_range().end();
+    if (trim_end_index < val.size()) {
+      trim_end_it = std::next(val.array_range().begin(), trim_end_index + 1);
     }
 
+    val = json_array<JsonType>(trim_start_it, trim_end_it);
     vec.emplace_back(val.size());
   };
 
diff --git a/src/server/json_family_test.cc b/src/server/json_family_test.cc
index 9fb2734d54a8..8bbddc4a32d6 100644
--- a/src/server/json_family_test.cc
+++ b/src/server/json_family_test.cc
@@ -715,10 +715,10 @@ TEST_F(JsonFamilyTest, ArrTrim) {
 
   resp = Run({"JSON.ARRTRIM", "json", "$[*]", "0", "1"});
   ASSERT_EQ(RespExpr::ARRAY, resp.type);
-  EXPECT_THAT(resp.GetVec(), ElementsAre(ArgType(RespExpr::NIL), IntArg(0), IntArg(1), IntArg(2)));
+  EXPECT_THAT(resp.GetVec(), ElementsAre(IntArg(0), IntArg(1), IntArg(2), IntArg(2)));
 
   resp = Run({"JSON.GET", "json"});
-  EXPECT_EQ(resp, R"([[],[],["b"],["b","c"]])");
+  EXPECT_EQ(resp, R"([[],["a"],["a","b"],["a","b"]])");
 
   json = R"(
     {"a":[], "nested": {"a": [1,4]}}
@@ -729,10 +729,10 @@ TEST_F(JsonFamilyTest, ArrTrim) {
 
   resp = Run({"JSON.ARRTRIM", "json", "$..a", "0", "1"});
   ASSERT_EQ(RespExpr::ARRAY, resp.type);
-  EXPECT_THAT(resp.GetVec(), ElementsAre(ArgType(RespExpr::NIL), IntArg(1)));
+  EXPECT_THAT(resp.GetVec(), ElementsAre(IntArg(0), IntArg(2)));
 
   resp = Run({"JSON.GET", "json"});
-  EXPECT_EQ(resp, R"({"a":[],"nested":{"a":[4]}})");
+  EXPECT_EQ(resp, R"({"a":[],"nested":{"a":[1,4]}})");
 
   json = R"(
     {"a":[1,2,3,2], "nested": {"a": false}}
@@ -743,10 +743,23 @@ TEST_F(JsonFamilyTest, ArrTrim) {
 
   resp = Run({"JSON.ARRTRIM", "json", "$..a", "1", "2"});
   ASSERT_EQ(RespExpr::ARRAY, resp.type);
-  EXPECT_THAT(resp.GetVec(), ElementsAre(IntArg(3), ArgType(RespExpr::NIL)));
+  EXPECT_THAT(resp.GetVec(), ElementsAre(IntArg(2), ArgType(RespExpr::NIL)));
+
+  resp = Run({"JSON.GET", "json"});
+  EXPECT_EQ(resp, R"({"a":[2,3],"nested":{"a":false}})");
+
+  json = R"(
+    [1,2,3,4,5,6,7]
+  )";
+
+  resp = Run({"JSON.SET", "json", "$", json});
+  ASSERT_THAT(resp, "OK");
+
+  resp = Run({"JSON.ARRTRIM", "json", "$", "2", "3"});
+  EXPECT_THAT(resp, IntArg(2));
 
   resp = Run({"JSON.GET", "json"});
-  EXPECT_EQ(resp, R"({"a":[1,3,2],"nested":{"a":false}})");
+  EXPECT_EQ(resp, R"([3,4])");
 }
 
 TEST_F(JsonFamilyTest, ArrInsert) {