Add mutate method to jsonpath (#2606)

* Add mutate method to jsonpath

Plug it in into OpToggle op.

---------

Signed-off-by: Roman Gershman <roman@dragonflydb.io>

src/core/json/jsonpath_test.cc

@@ -267,4 +267,45 @@ TEST_F(JsonPathTest, Wildcard) {
   ASSERT_THAT(arr, ElementsAre(1, 2, 3));
 }
 
+TEST_F(JsonPathTest, Mutate) {
+  JsonType json = JsonFromString(R"([1, 2, 3, 5, 6])", pmr::get_default_resource()).value();
+  ASSERT_EQ(0, Parse("$[*]"));
+  Path path = driver_.TakePath();
+  MutateCallback cb = [&](optional<string_view>, JsonType* val) {
+    int intval = val->as<int>();
+    *val = intval + 1;
+    return false;
+  };
+
+  MutatePath(path, cb, &json);
+  vector<int> arr;
+  for (auto& el : json.array_range()) {
+    arr.push_back(el.as<int>());
+  }
+  ASSERT_THAT(arr, ElementsAre(2, 3, 4, 6, 7));
+
+  json = JsonFromString(R"(
+    {"a":[7], "inner": {"a": {"bool": true, "c": 42}}}
+  )",
+                        pmr::get_default_resource());
+  ASSERT_EQ(0, Parse("$..a.*"));
+  path = driver_.TakePath();
+  MutatePath(
+      path,
+      [&](optional<string_view> key, JsonType* val) {
+        if (val->is_int64() && !key) {  // array element
+          *val = 42;
+          return false;
+        }
+        if (val->is_bool()) {
+          *val = false;
+          return false;
+        }
+        return true;
+      },
+      &json);
+
+  ASSERT_EQ(R"({"a":[42],"inner":{"a":{"bool":false}}})", json.to_string());
+}
+
 }  // namespace dfly::json

src/core/json/path.cc

@@ -102,6 +102,7 @@ class Dfs {
 
   // TODO: for some operations we need to know the type of mismatches.
   void Traverse(absl::Span<const PathSegment> path, const JsonType& json, const Cb& callback);
+  void Mutate(absl::Span<const PathSegment> path, const MutateCallback& callback, JsonType* json);
 
   unsigned matches() const {
     return matches_;
@@ -113,11 +114,21 @@ class Dfs {
   nonstd::expected<void, MatchStatus> PerformStep(const PathSegment& segment, const JsonType& node,
                                                   const Cb& callback);
 
+  nonstd::expected<void, MatchStatus> MutateStep(const PathSegment& segment,
+                                                 const MutateCallback& cb, JsonType* node);
+
+  void Mutate(const PathSegment& segment, const MutateCallback& callback, JsonType* node);
+
   void DoCall(const Cb& callback, optional<string_view> key, const JsonType& node) {
     ++matches_;
     callback(key, node);
   }
 
+  bool Mutate(const MutateCallback& callback, optional<string_view> key, JsonType* node) {
+    ++matches_;
+    return callback(key, node);
+  }
+
   unsigned matches_ = 0;
 };
 
@@ -253,6 +264,44 @@ void Dfs::Traverse(absl::Span<const PathSegment> path, const JsonType& root, con
   } while (!stack.empty());
 }
 
+void Dfs::Mutate(absl::Span<const PathSegment> path, const MutateCallback& callback,
+                 JsonType* json) {
+  DCHECK(!path.empty());
+  if (path.size() == 1) {
+    MutateStep(path[0], callback, json);
+    return;
+  }
+
+  using Item = DfsItem<false>;
+  vector<Item> stack;
+  stack.emplace_back(json);
+
+  do {
+    unsigned segment_index = stack.back().segment_idx();
+    const auto& path_segment = path[segment_index];
+
+    // init or advance the current object
+    Item::AdvanceResult res = stack.back().Advance(path_segment);
+    if (res && res->first != nullptr) {
+      JsonType* next = res->first;
+      DVLOG(2) << "Handling now " << next->type() << " " << next->to_string();
+
+      // We descent only if next is object or an array.
+      if (IsRecursive(next->type())) {
+        unsigned next_seg_id = res->second;
+
+        if (next_seg_id + 1 < path.size()) {
+          stack.emplace_back(next, next_seg_id);
+        } else {
+          MutateStep(path[next_seg_id], callback, next);
+        }
+      }
+    } else {
+      stack.pop_back();
+    }
+  } while (!stack.empty());
+}
+
 auto Dfs::PerformStep(const PathSegment& segment, const JsonType& node, const Cb& callback)
     -> nonstd::expected<void, MatchStatus> {
   switch (segment.type()) {
@@ -290,6 +339,50 @@ auto Dfs::PerformStep(const PathSegment& segment, const JsonType& node, const Cb
   return {};
 }
 
+auto Dfs::MutateStep(const PathSegment& segment, const MutateCallback& cb, JsonType* node)
+    -> nonstd::expected<void, MatchStatus> {
+  switch (segment.type()) {
+    case SegmentType::IDENTIFIER: {
+      if (!node->is_object())
+        return make_unexpected(MISMATCH);
+
+      auto it = node->find(segment.identifier());
+      if (it != node->object_range().end()) {
+        if (Mutate(cb, it->key(), &it->value())) {
+          node->erase(it);
+        }
+      }
+    } break;
+    case SegmentType::INDEX: {
+      if (!node->is_array())
+        return make_unexpected(MISMATCH);
+      if (segment.index() >= node->size()) {
+        return make_unexpected(OUT_OF_BOUNDS);
+      }
+      if (Mutate(cb, nullopt, &node[segment.index()])) {
+        node->erase(node->array_range().begin() + segment.index());
+      }
+    } break;
+
+    case SegmentType::DESCENT:
+    case SegmentType::WILDCARD: {
+      if (node->is_object()) {
+        auto it = node->object_range().begin();
+        while (it != node->object_range().end()) {
+          it = Mutate(cb, it->key(), &it->value()) ? node->erase(it) : it + 1;
+        }
+      } else if (node->is_array()) {
+        auto it = node->array_range().begin();
+        while (it != node->array_range().end()) {
+          it = Mutate(cb, nullopt, &*it) ? node->erase(it) : it + 1;
+        }
+      }
+    } break;
+  }
+
+  return {};
+}
+
 }  // namespace
 
 void EvaluatePath(const Path& path, const JsonType& json, PathCallback callback) {
@@ -298,4 +391,10 @@ void EvaluatePath(const Path& path, const JsonType& json, PathCallback callback)
   Dfs().Traverse(path, json, std::move(callback));
 }
 
+void MutatePath(const Path& path, MutateCallback callback, JsonType* json) {
+  if (path.empty())
+    return;
+  Dfs().Mutate(path, callback, json);
+}
+
 }  // namespace dfly::json

src/core/json/path.h

@@ -56,6 +56,10 @@ using Path = std::vector<PathSegment>;
 // The second argument is a json value of either object fields or array elements.
 using PathCallback = absl::FunctionRef<void(std::optional<std::string_view>, const JsonType&)>;
 
+// Returns true if the entry should be deleted, false otherwise.
+using MutateCallback = absl::FunctionRef<bool(std::optional<std::string_view>, JsonType*)>;
+
 void EvaluatePath(const Path& path, const JsonType& json, PathCallback callback);
+void MutatePath(const Path& path, MutateCallback callback, JsonType* json);
 
 }  // namespace dfly::json

src/server/json_family.cc

@@ -578,23 +578,49 @@ OpResult<vector<OptSizeT>> OpArrLen(const OpArgs& op_args, string_view key, Json
   return vec;
 }
 
-OpResult<vector<OptBool>> OpToggle(const OpArgs& op_args, string_view key, string_view path) {
+OpResult<vector<OptBool>> OpToggle(const OpArgs& op_args, string_view key, string_view path,
+                                   JsonPathV2 expression) {
   vector<OptBool> vec;
-  auto cb = [&vec](const auto&, JsonType& val) {
-    if (val.is_bool()) {
-      bool current_val = val.as_bool() ^ true;
-      val = current_val;
-      vec.emplace_back(current_val);
-    } else {
-      vec.emplace_back(nullopt);
+  if (std::holds_alternative<json::Path>(expression)) {
+    auto it_res = op_args.shard->db_slice().FindMutable(op_args.db_cntx, key, OBJ_JSON);
+    if (!it_res.ok()) {
+      return it_res.status();
     }
-  };
 
-  OpStatus status = UpdateEntry(op_args, key, path, cb);
-  if (status != OpStatus::OK) {
-    return status;
+    PrimeValue& pv = it_res->it->second;
+
+    op_args.shard->search_indices()->RemoveDoc(key, op_args.db_cntx, pv);
+
+    const json::Path& expr = std::get<json::Path>(expression);
+    auto cb = [&vec](optional<string_view>, JsonType* val) {
+      if (val->is_bool()) {
+        bool next_val = val->as_bool() ^ true;
+        *val = next_val;
+        vec.emplace_back(next_val);
+      } else {
+        vec.emplace_back(nullopt);
+      }
+      return false;
+    };
+    json::MutatePath(expr, std::move(cb), pv.GetJson());
+    it_res->post_updater.Run();
+    op_args.shard->search_indices()->AddDoc(key, op_args.db_cntx, pv);
+  } else {
+    auto cb = [&vec](const auto&, JsonType& val) {
+      if (val.is_bool()) {
+        bool current_val = val.as_bool() ^ true;
+        val = current_val;
+        vec.emplace_back(current_val);
+      } else {
+        vec.emplace_back(nullopt);
+      }
+    };
+
+    OpStatus status = UpdateEntry(op_args, key, path, cb);
+    if (status != OpStatus::OK) {
+      return status;
+    }
   }
-
   return vec;
 }
 
@@ -1698,8 +1724,10 @@ void JsonFamily::Toggle(CmdArgList args, ConnectionContext* cntx) {
   string_view key = ArgS(args, 0);
   string_view path = ArgS(args, 1);
 
+  JsonPathV2 expression = PARSE_PATHV2(path);
+
   auto cb = [&](Transaction* t, EngineShard* shard) {
-    return OpToggle(t->GetOpArgs(shard), key, path);
+    return OpToggle(t->GetOpArgs(shard), key, path, std::move(expression));
   };
 
   Transaction* trans = cntx->transaction;