From c8a5059e999b126f45d7aea2f68896c57071d797 Mon Sep 17 00:00:00 2001
From: Vladislav Oleshko <vlad@dragonflydb.io>
Date: Tue, 16 Jul 2024 21:21:18 +0300
Subject: [PATCH 1/4] chore: Make KeyIndex iterable

---
 src/facade/facade_types.h            |  4 ++
 src/server/acl/validator.cc          | 22 ++------
 src/server/main_service.cc           | 28 ++--------
 src/server/multi_command_squasher.cc | 35 ++++--------
 src/server/transaction.cc            | 80 ++++++++++++----------------
 src/server/tx_base.h                 | 49 ++++++++++++-----
 6 files changed, 94 insertions(+), 124 deletions(-)

diff --git a/src/facade/facade_types.h b/src/facade/facade_types.h
index 877047941068..5b8e0eff1c30 100644
--- a/src/facade/facade_types.h
+++ b/src/facade/facade_types.h
@@ -91,6 +91,10 @@ struct ArgRange {
     return Range().second;
   }
 
+  std::string_view operator[](size_t idx) const {
+    return std::visit([idx](const auto& span) { return facade::ToSV(span[idx]); }, span);
+  }
+
   std::variant<CmdArgList, ArgSlice, OwnedArgSlice> span;
 };
 struct ConnectionStats {
diff --git a/src/server/acl/validator.cc b/src/server/acl/validator.cc
index 9482230e5f28..9dcb1dba689b 100644
--- a/src/server/acl/validator.cc
+++ b/src/server/acl/validator.cc
@@ -76,23 +76,11 @@ namespace dfly::acl {
 
   bool keys_allowed = true;
   if (!keys.all_keys && id.first_key_pos() != 0 && (is_read_command || is_write_command)) {
-    const auto keys_index = DetermineKeys(&id, tail_args).value();
-    const size_t end = keys_index.end;
-    if (keys_index.bonus) {
-      auto target = facade::ToSV(tail_args[*keys_index.bonus]);
-      if (!iterate_globs(target)) {
-        keys_allowed = false;
-      }
-    }
-    if (keys_allowed) {
-      for (size_t i = keys_index.start; i < end; i += keys_index.step) {
-        auto target = facade::ToSV(tail_args[i]);
-        if (!iterate_globs(target)) {
-          keys_allowed = false;
-          break;
-        }
-      }
-    }
+    auto keys_index = DetermineKeys(&id, tail_args);
+    DCHECK(keys_index);
+
+    for (std::string_view key : keys_index->Range(tail_args))
+      keys_allowed &= iterate_globs(key);
   }
 
   return {keys_allowed, AclLog::Reason::KEY};
diff --git a/src/server/main_service.cc b/src/server/main_service.cc
index 702ae5de21c0..ac5318b39abf 100644
--- a/src/server/main_service.cc
+++ b/src/server/main_service.cc
@@ -710,7 +710,7 @@ Transaction::MultiMode DeduceExecMode(ExecEvalState state,
         StoredCmd cmd = scmd;
         cmd.Fill(&arg_vec);
         auto keys = DetermineKeys(scmd.Cid(), absl::MakeSpan(arg_vec));
-        transactional |= (keys && keys.value().num_args() > 0);
+        transactional |= (keys && keys.value().Size() > 0);
       } else {
         transactional |= scmd.Cid()->IsTransactional();
       }
@@ -935,10 +935,8 @@ optional<ErrorReply> Service::CheckKeysOwnership(const CommandId* cid, CmdArgLis
   optional<cluster::SlotId> keys_slot;
   bool cross_slot = false;
   // Iterate keys and check to which slot they belong.
-  for (unsigned i = key_index.start; i < key_index.end; i += key_index.step) {
-    string_view key = ArgS(args, i);
-    cluster::SlotId slot = cluster::KeySlot(key);
-    if (keys_slot && slot != *keys_slot) {
+  for (string_view key : key_index.Range(args)) {
+    if (cluster::SlotId slot = cluster::KeySlot(key); keys_slot && slot != *keys_slot) {
       cross_slot = true;  // keys belong to different slots
       break;
     } else {
@@ -983,18 +981,7 @@ optional<ErrorReply> CheckKeysDeclared(const ConnectionState::ScriptInfo& eval_i
   // TODO: Switch to transaction internal locked keys once single hop multi transactions are merged
   // const auto& locked_keys = trans->GetMultiKeys();
   const auto& locked_tags = eval_info.lock_tags;
-
-  const auto& key_index = *key_index_res;
-  for (unsigned i = key_index.start; i < key_index.end; ++i) {
-    string_view key = ArgS(args, i);
-    LockTag tag{key};
-    if (!locked_tags.contains(tag)) {
-      return ErrorReply(absl::StrCat(kUndeclaredKeyErr, ", key: ", key));
-    }
-  }
-
-  if (key_index.bonus) {
-    string_view key = ArgS(args, *key_index.bonus);
+  for (string_view key : key_index_res->Range(args)) {
     if (!locked_tags.contains(LockTag{key})) {
       return ErrorReply(absl::StrCat(kUndeclaredKeyErr, ", key: ", key));
     }
@@ -2113,13 +2100,8 @@ template <typename F> void IterateAllKeys(ConnectionState::ExecInfo* exec_info,
     if (!key_res.ok())
       continue;
 
-    auto key_index = key_res.value();
-
-    for (unsigned i = key_index.start; i < key_index.end; i += key_index.step)
+    for (unsigned i : key_res->Range())
       f(arg_vec[i]);
-
-    if (key_index.bonus)
-      f(arg_vec[*key_index.bonus]);
   }
 }
 
diff --git a/src/server/multi_command_squasher.cc b/src/server/multi_command_squasher.cc
index b9eeb6b545bb..ecd16e61acd1 100644
--- a/src/server/multi_command_squasher.cc
+++ b/src/server/multi_command_squasher.cc
@@ -13,6 +13,7 @@
 #include "server/conn_context.h"
 #include "server/engine_shard_set.h"
 #include "server/transaction.h"
+#include "server/tx_base.h"
 
 namespace dfly {
 
@@ -22,14 +23,6 @@ using namespace util;
 
 namespace {
 
-template <typename F> void IterateKeys(CmdArgList args, KeyIndex keys, F&& f) {
-  for (unsigned i = keys.start; i < keys.end; i += keys.step)
-    f(args[i]);
-
-  if (keys.bonus)
-    f(args[*keys.bonus]);
-}
-
 void CheckConnStateClean(const ConnectionState& state) {
   DCHECK_EQ(state.exec_info.state, ConnectionState::ExecInfo::EXEC_INACTIVE);
   DCHECK(state.exec_info.body.empty());
@@ -90,29 +83,21 @@ MultiCommandSquasher::SquashResult MultiCommandSquasher::TrySquash(StoredCmd* cm
   auto keys = DetermineKeys(cmd->Cid(), args);
   if (!keys.ok())
     return SquashResult::ERROR;
+  if (keys->Size() == 0)
+    return SquashResult::NOT_SQUASHED;
 
   // Check if all commands belong to one shard
-  bool found_more = false;
   cluster::UniqueSlotChecker slot_checker;
   ShardId last_sid = kInvalidSid;
-  IterateKeys(args, *keys, [&last_sid, &found_more, &slot_checker](MutableSlice key) {
-    if (found_more)
-      return;
-
-    string_view key_sv = facade::ToSV(key);
 
-    slot_checker.Add(key_sv);
-
-    ShardId sid = Shard(key_sv, shard_set->size());
-    if (last_sid == kInvalidSid || last_sid == sid) {
+  for (string_view key : keys->Range(args)) {
+    slot_checker.Add(key);
+    ShardId sid = Shard(key, shard_set->size());
+    if (last_sid == kInvalidSid || last_sid == sid)
       last_sid = sid;
-      return;
-    }
-    found_more = true;
-  });
-
-  if (found_more || last_sid == kInvalidSid)
-    return SquashResult::NOT_SQUASHED;
+    else
+      return SquashResult::NOT_SQUASHED;  // at least two shards
+  }
 
   auto& sinfo = PrepareShardInfo(last_sid, slot_checker.GetUniqueSlotId());
 
diff --git a/src/server/transaction.cc b/src/server/transaction.cc
index 64fa7b48d475..dfb058060a33 100644
--- a/src/server/transaction.cc
+++ b/src/server/transaction.cc
@@ -190,21 +190,17 @@ void Transaction::BuildShardIndex(const KeyIndex& key_index, std::vector<PerShar
     }
   };
 
-  if (key_index.bonus) {
-    DCHECK(key_index.step == 1);
-    string_view key = ArgS(full_args_, *key_index.bonus);
-    unique_slot_checker_.Add(key);
-    uint32_t sid = Shard(key, shard_data_.size());
-    add(sid, *key_index.bonus, *key_index.bonus + 1);
-  }
+  // Because of the way we iterate in InitShardData
+  DCHECK(!key_index.bonus || key_index.step == 1);
 
-  for (unsigned i = key_index.start; i < key_index.end; i += key_index.step) {
+  for (unsigned i : key_index.Range()) {
     string_view key = ArgS(full_args_, i);
     unique_slot_checker_.Add(key);
-    uint32_t sid = Shard(key, shard_data_.size());
-    shard_index[sid].key_step = key_index.step;
+    ShardId sid = Shard(key, shard_data_.size());
 
-    add(sid, i, i + key_index.step);
+    unsigned step = i == key_index.bonus ? 1 : key_index.step;
+    shard_index[sid].key_step = step;
+    add(sid, i, step);
   }
 }
 
@@ -235,11 +231,9 @@ void Transaction::InitShardData(absl::Span<const PerShardCache> shard_index, siz
     unique_shard_cnt_++;
     unique_shard_id_ = i;
 
-    for (size_t j = 0; j < src.slices.size(); ++j) {
-      IndexSlice slice = src.slices[j];
-      args_slices_.push_back(slice);
-      for (uint32_t k = slice.first; k < slice.second; k += src.key_step) {
-        string_view key = ArgS(full_args_, k);
+    for (const auto [start, end] : src.slices) {
+      args_slices_.emplace_back(start, end);
+      for (string_view key : KeyIndex(start, end, src.key_step).Range(full_args_)) {
         kv_fp_.push_back(LockTag(key).Fingerprint());
         sd.fp_count++;
       }
@@ -267,10 +261,8 @@ void Transaction::StoreKeysInArgs(const KeyIndex& key_index) {
 
   // even for a single key we may have multiple arguments per key (MSET).
   args_slices_.emplace_back(key_index.start, key_index.end);
-  for (unsigned j = key_index.start; j < key_index.end; j += key_index.step) {
-    string_view key = ArgS(full_args_, j);
+  for (string_view key : key_index.Range(full_args_))
     kv_fp_.push_back(LockTag(key).Fingerprint());
-  }
 }
 
 void Transaction::InitByKeys(const KeyIndex& key_index) {
@@ -284,14 +276,14 @@ void Transaction::InitByKeys(const KeyIndex& key_index) {
   // Stub transactions always operate only on single shard.
   bool is_stub = multi_ && multi_->role == SQUASHED_STUB;
 
-  if ((key_index.HasSingleKey() && !IsAtomicMulti()) || is_stub) {
+  if ((key_index.Size() == 1 && !IsAtomicMulti()) || is_stub) {
     DCHECK(!IsActiveMulti() || multi_->mode == NON_ATOMIC);
 
     // We don't have to split the arguments by shards, so we can copy them directly.
     StoreKeysInArgs(key_index);
 
     unique_shard_cnt_ = 1;
-    string_view akey = ArgS(full_args_, key_index.start);
+    string_view akey = *key_index.Range(full_args_).begin();
     if (is_stub)  // stub transactions don't migrate
       DCHECK_EQ(unique_shard_id_, Shard(akey, shard_set->size()));
     else {
@@ -317,7 +309,7 @@ void Transaction::InitByKeys(const KeyIndex& key_index) {
   BuildShardIndex(key_index, &shard_index);
 
   // Initialize shard data based on distributed arguments.
-  InitShardData(shard_index, key_index.num_args());
+  InitShardData(shard_index, key_index.Size());
 
   DCHECK(!multi_ || multi_->mode != LOCK_AHEAD || !multi_->tag_fps.empty());
 
@@ -428,7 +420,7 @@ void Transaction::StartMultiLockedAhead(DbIndex dbid, CmdArgList keys, bool skip
   PrepareMultiFps(keys);
 
   InitBase(dbid, keys);
-  InitByKeys(KeyIndex::Range(0, keys.size()));
+  InitByKeys(KeyIndex(0, keys.size()));
 
   if (!skip_scheduling)
     ScheduleInternal();
@@ -1418,23 +1410,24 @@ bool Transaction::CanRunInlined() const {
 }
 
 OpResult<KeyIndex> DetermineKeys(const CommandId* cid, CmdArgList args) {
-  KeyIndex key_index;
-
   if (cid->opt_mask() & (CO::GLOBAL_TRANS | CO::NO_KEY_TRANSACTIONAL))
-    return key_index;
+    return KeyIndex{};
 
   int num_custom_keys = -1;
 
-  if (cid->opt_mask() & CO::VARIADIC_KEYS) {
+  unsigned start = 0, end = 0, step = 0;
+  std::optional<unsigned> bonus = std::nullopt;
+
+  if (cid->opt_mask() & CO::VARIADIC_KEYS) {  // number of keys is not trivially deducable
     // ZUNION/INTER <num_keys> <key1> [<key2> ...]
     // EVAL <script> <num_keys>
     // XREAD ... STREAMS ...
-    if (args.size() < 2) {
+    if (args.size() < 2)
       return OpStatus::SYNTAX_ERR;
-    }
 
     string_view name{cid->name()};
 
+    // Determine based on STREAMS argument position
     if (name == "XREAD" || name == "XREADGROUP") {
       for (size_t i = 0; i < args.size(); ++i) {
         string_view arg = ArgS(args, i);
@@ -1443,24 +1436,20 @@ OpResult<KeyIndex> DetermineKeys(const CommandId* cid, CmdArgList args) {
           if (left < 2 || left % 2 != 0)
             return OpStatus::SYNTAX_ERR;
 
-          key_index.start = i + 1;
-          key_index.end = key_index.start + (left / 2);
-          key_index.step = 1;
-
-          return key_index;
+          return KeyIndex(i + 1, i + 1 + (left / 2));
         }
       }
       return OpStatus::SYNTAX_ERR;
     }
 
     if (absl::EndsWith(name, "STORE"))
-      key_index.bonus = 0;  // Z<xxx>STORE <key> commands
+      bonus = 0;  // Z<xxx>STORE <key> commands
 
     unsigned num_keys_index;
     if (absl::StartsWith(name, "EVAL"))
       num_keys_index = 1;
     else
-      num_keys_index = key_index.bonus ? *key_index.bonus + 1 : 0;
+      num_keys_index = bonus ? *bonus + 1 : 0;
 
     string_view num = ArgS(args, num_keys_index);
     if (!absl::SimpleAtoi(num, &num_custom_keys) || num_custom_keys < 0)
@@ -1477,22 +1466,22 @@ OpResult<KeyIndex> DetermineKeys(const CommandId* cid, CmdArgList args) {
   }
 
   if (cid->first_key_pos() > 0) {
-    key_index.start = cid->first_key_pos() - 1;
+    start = cid->first_key_pos() - 1;
     int last = cid->last_key_pos();
 
     if (num_custom_keys >= 0) {
-      key_index.end = key_index.start + num_custom_keys;
+      end = start + num_custom_keys;
     } else {
-      key_index.end = last > 0 ? last : (int(args.size()) + last + 1);
+      end = last > 0 ? last : (int(args.size()) + last + 1);
     }
     if (cid->opt_mask() & CO::INTERLEAVED_KEYS) {
       if (cid->name() == "JSON.MSET") {
-        key_index.step = 3;
+        step = 3;
       } else {
-        key_index.step = 2;
+        step = 2;
       }
     } else {
-      key_index.step = 1;
+      step = 1;
     }
 
     if (cid->opt_mask() & CO::STORE_LAST_KEY) {
@@ -1502,17 +1491,16 @@ OpResult<KeyIndex> DetermineKeys(const CommandId* cid, CmdArgList args) {
         // key member radius .. STORE destkey
         string_view opt = ArgS(args, args.size() - 2);
         if (absl::EqualsIgnoreCase(opt, "STORE") || absl::EqualsIgnoreCase(opt, "STOREDIST")) {
-          key_index.bonus = args.size() - 1;
+          bonus = args.size() - 1;
         }
       }
     }
 
-    return key_index;
+    return KeyIndex{start, end, step, bonus};
   }
 
   LOG(FATAL) << "TBD: Not supported " << cid->name();
-
-  return key_index;
+  return {};
 }
 
 std::vector<Transaction::PerShardCache>& Transaction::TLTmpSpace::GetShardIndex(unsigned size) {
diff --git a/src/server/tx_base.h b/src/server/tx_base.h
index 03aab1a5f080..581a3cc40d71 100644
--- a/src/server/tx_base.h
+++ b/src/server/tx_base.h
@@ -8,6 +8,7 @@
 
 #include <optional>
 
+#include "base/iterator.h"
 #include "src/facade/facade_types.h"
 
 namespace dfly {
@@ -33,28 +34,50 @@ struct KeyLockArgs {
 
 // Describes key indices.
 struct KeyIndex {
-  unsigned start;
-  unsigned end;   // does not include this index (open limit).
-  unsigned step;  // 1 for commands like mget. 2 for commands like mset.
+  KeyIndex(unsigned start = 0, unsigned end = 0, unsigned step = 1,
+           std::optional<unsigned> bonus = std::nullopt)
+      : start(start), end(end), step(step), bonus(bonus) {
+  }
+
+  using iterator_category = std::forward_iterator_tag;
+  using value_type = unsigned;
+  using difference_type = std::ptrdiff_t;
+  using pointer = value_type;
+  using reference = value_type;
 
-  // if index is non-zero then adds another key index (usually 0).
-  // relevant for for commands like ZUNIONSTORE/ZINTERSTORE for destination key.
-  std::optional<uint16_t> bonus{};
+  unsigned operator*() {
+    if (bonus)
+      return *bonus;
+    return start;
+  }
 
-  KeyIndex(unsigned s = 0, unsigned e = 0, unsigned step = 0) : start(s), end(e), step(step) {
+  KeyIndex& operator++() {
+    if (bonus)
+      bonus.reset();
+    else
+      start = std::max(end, start + step);
+    return *this;
   }
 
-  static KeyIndex Range(unsigned start, unsigned end, unsigned step = 1) {
-    return KeyIndex{start, end, step};
+  bool operator!=(const KeyIndex& ki) const {
+    return std::tie(start, end, step, bonus) != std::tie(ki.start, ki.end, ki.step, ki.bonus);
   }
 
-  bool HasSingleKey() const {
-    return !bonus && (start + step >= end);
+  unsigned Size() const {
+    return (end - start) / step + unsigned(bonus.has_value());
   }
 
-  unsigned num_args() const {
-    return end - start + bool(bonus);
+  auto Range() const {
+    return base::it::Range(*this, KeyIndex{end, end, step, std::nullopt});
   }
+
+  auto Range(facade::ArgRange args) const {
+    return base::it::Transform([args](unsigned idx) { return args[idx]; }, Range());
+  }
+
+ public:
+  unsigned start, end, step;      // [start, end) with step
+  std::optional<unsigned> bonus;  // destination key, for example for commands that end with STORE
 };
 
 struct DbContext {

From 3d27f5ff93bf542a38e82f8b4c49c9489224dbb4 Mon Sep 17 00:00:00 2001
From: Vladislav Oleshko <vlad@dragonflydb.io>
Date: Tue, 16 Jul 2024 23:15:28 +0300
Subject: [PATCH 2/4] chore: fixes

Signed-off-by: Vladislav Oleshko <vlad@dragonflydb.io>
---
 src/server/main_service.cc           |  2 +-
 src/server/multi_command_squasher.cc |  2 +-
 src/server/transaction.cc            | 31 ++++++++++++++--------------
 src/server/tx_base.cc                | 18 ++++++++++++++++
 src/server/tx_base.h                 | 24 +++++----------------
 5 files changed, 41 insertions(+), 36 deletions(-)

diff --git a/src/server/main_service.cc b/src/server/main_service.cc
index df110e92101e..cdfaf56d45a5 100644
--- a/src/server/main_service.cc
+++ b/src/server/main_service.cc
@@ -714,7 +714,7 @@ Transaction::MultiMode DeduceExecMode(ExecEvalState state,
         StoredCmd cmd = scmd;
         cmd.Fill(&arg_vec);
         auto keys = DetermineKeys(scmd.Cid(), absl::MakeSpan(arg_vec));
-        transactional |= (keys && keys.value().Size() > 0);
+        transactional |= (keys && keys.value().NumArgs() > 0);
       } else {
         transactional |= scmd.Cid()->IsTransactional();
       }
diff --git a/src/server/multi_command_squasher.cc b/src/server/multi_command_squasher.cc
index b3d2bcec220f..80d518fbf15f 100644
--- a/src/server/multi_command_squasher.cc
+++ b/src/server/multi_command_squasher.cc
@@ -83,7 +83,7 @@ MultiCommandSquasher::SquashResult MultiCommandSquasher::TrySquash(StoredCmd* cm
   auto keys = DetermineKeys(cmd->Cid(), args);
   if (!keys.ok())
     return SquashResult::ERROR;
-  if (keys->Size() == 0)
+  if (keys->NumArgs() == 0)
     return SquashResult::NOT_SQUASHED;
 
   // Check if all commands belong to one shard
diff --git a/src/server/transaction.cc b/src/server/transaction.cc
index 4ddef621d45b..d2841f86a5f5 100644
--- a/src/server/transaction.cc
+++ b/src/server/transaction.cc
@@ -190,28 +190,26 @@ void Transaction::InitGlobal() {
 }
 
 void Transaction::BuildShardIndex(const KeyIndex& key_index, std::vector<PerShardCache>* out) {
-  auto& shard_index = *out;
-
-  auto add = [&shard_index](uint32_t sid, uint32_t b, uint32_t e) {
-    auto& slices = shard_index[sid].slices;
-    if (!slices.empty() && slices.back().second == b) {
-      slices.back().second = e;
-    } else {
-      slices.emplace_back(b, e);
-    }
-  };
-
   // Because of the way we iterate in InitShardData
   DCHECK(!key_index.bonus || key_index.step == 1);
 
+  auto& shard_index = *out;
   for (unsigned i : key_index.Range()) {
     string_view key = ArgS(full_args_, i);
     unique_slot_checker_.Add(key);
     ShardId sid = Shard(key, shard_data_.size());
 
+    VLOG(0) << i << " " << key;
+
     unsigned step = i == key_index.bonus ? 1 : key_index.step;
+
     shard_index[sid].key_step = step;
-    add(sid, i, step);
+    auto& slices = shard_index[sid].slices;
+    if (!slices.empty() && slices.back().second == i) {
+      slices.back().second = i + step;
+    } else {
+      slices.emplace_back(i, i + step);
+    }
   }
 }
 
@@ -242,7 +240,7 @@ void Transaction::InitShardData(absl::Span<const PerShardCache> shard_index, siz
     unique_shard_cnt_++;
     unique_shard_id_ = i;
 
-    for (const auto [start, end] : src.slices) {
+    for (const auto& [start, end] : src.slices) {
       args_slices_.emplace_back(start, end);
       for (string_view key : KeyIndex(start, end, src.key_step).Range(full_args_)) {
         kv_fp_.push_back(LockTag(key).Fingerprint());
@@ -282,12 +280,15 @@ void Transaction::InitByKeys(const KeyIndex& key_index) {
     return;
   }
 
+  VLOG(0) << "InitByKeys " << key_index.start << " " << key_index.end << " " << key_index.step
+          << " -> " << key_index.NumArgs();
+
   DCHECK_LT(key_index.start, full_args_.size());
 
   // Stub transactions always operate only on single shard.
   bool is_stub = multi_ && multi_->role == SQUASHED_STUB;
 
-  if ((key_index.Size() == 1 && !IsAtomicMulti()) || is_stub) {
+  if ((key_index.NumArgs() == 1 && !IsAtomicMulti()) || is_stub) {
     DCHECK(!IsActiveMulti() || multi_->mode == NON_ATOMIC);
 
     // We don't have to split the arguments by shards, so we can copy them directly.
@@ -320,7 +321,7 @@ void Transaction::InitByKeys(const KeyIndex& key_index) {
   BuildShardIndex(key_index, &shard_index);
 
   // Initialize shard data based on distributed arguments.
-  InitShardData(shard_index, key_index.Size());
+  InitShardData(shard_index, key_index.NumArgs());
 
   DCHECK(!multi_ || multi_->mode != LOCK_AHEAD || !multi_->tag_fps.empty());
 
diff --git a/src/server/tx_base.cc b/src/server/tx_base.cc
index 08081df7b091..370ebe933506 100644
--- a/src/server/tx_base.cc
+++ b/src/server/tx_base.cc
@@ -17,6 +17,24 @@ namespace dfly {
 using namespace std;
 using Payload = journal::Entry::Payload;
 
+unsigned KeyIndex::operator*() const {
+  if (bonus)
+    return *bonus;
+  return start;
+}
+
+KeyIndex& KeyIndex::operator++() {
+  if (bonus)
+    bonus.reset();
+  else
+    start = std::min(end, start + step);
+  return *this;
+}
+
+bool KeyIndex::operator!=(const KeyIndex& ki) const {
+  return std::tie(start, end, step, bonus) != std::tie(ki.start, ki.end, ki.step, ki.bonus);
+}
+
 DbSlice& DbContext::GetDbSlice(ShardId shard_id) const {
   return ns->GetDbSlice(shard_id);
 }
diff --git a/src/server/tx_base.h b/src/server/tx_base.h
index b2eb43050cd5..aa8db41fee77 100644
--- a/src/server/tx_base.h
+++ b/src/server/tx_base.h
@@ -46,26 +46,12 @@ struct KeyIndex {
   using pointer = value_type;
   using reference = value_type;
 
-  unsigned operator*() {
-    if (bonus)
-      return *bonus;
-    return start;
-  }
-
-  KeyIndex& operator++() {
-    if (bonus)
-      bonus.reset();
-    else
-      start = std::max(end, start + step);
-    return *this;
-  }
-
-  bool operator!=(const KeyIndex& ki) const {
-    return std::tie(start, end, step, bonus) != std::tie(ki.start, ki.end, ki.step, ki.bonus);
-  }
+  unsigned operator*() const;
+  KeyIndex& operator++();
+  bool operator!=(const KeyIndex& ki) const;
 
-  unsigned Size() const {
-    return (end - start) / step + unsigned(bonus.has_value());
+  unsigned NumArgs() const {
+    return (end - start) + unsigned(bonus.has_value());
   }
 
   auto Range() const {

From 41bac7aa08564e7cd2c20f762af5a53a535993e6 Mon Sep 17 00:00:00 2001
From: Vladislav Oleshko <vlad@dragonflydb.io>
Date: Wed, 17 Jul 2024 12:24:27 +0300
Subject: [PATCH 3/4] chore; add test

Signed-off-by: Vladislav Oleshko <vlad@dragonflydb.io>
---
 src/server/cluster/cluster_family_test.cc | 1 +
 src/server/transaction.cc                 | 9 +--------
 2 files changed, 2 insertions(+), 8 deletions(-)

diff --git a/src/server/cluster/cluster_family_test.cc b/src/server/cluster/cluster_family_test.cc
index 11224d035e36..d523bb5ae35c 100644
--- a/src/server/cluster/cluster_family_test.cc
+++ b/src/server/cluster/cluster_family_test.cc
@@ -686,6 +686,7 @@ TEST_F(ClusterFamilyTest, ClusterCrossSlot) {
 
   EXPECT_THAT(Run({"MSET", "key", "value", "key2", "value2"}), ErrArg("CROSSSLOT"));
   EXPECT_THAT(Run({"MGET", "key", "key2"}), ErrArg("CROSSSLOT"));
+  EXPECT_THAT(Run({"ZINTERSTORE", "key", "2", "key1", "key2"}), ErrArg("CROSSSLOT"));
 
   EXPECT_EQ(Run({"MSET", "key{tag}", "value", "key2{tag}", "value2"}), "OK");
   EXPECT_THAT(Run({"MGET", "key{tag}", "key2{tag}"}), RespArray(ElementsAre("value", "value2")));
diff --git a/src/server/transaction.cc b/src/server/transaction.cc
index d2841f86a5f5..dfdf69d231b2 100644
--- a/src/server/transaction.cc
+++ b/src/server/transaction.cc
@@ -199,11 +199,7 @@ void Transaction::BuildShardIndex(const KeyIndex& key_index, std::vector<PerShar
     unique_slot_checker_.Add(key);
     ShardId sid = Shard(key, shard_data_.size());
 
-    VLOG(0) << i << " " << key;
-
-    unsigned step = i == key_index.bonus ? 1 : key_index.step;
-
-    shard_index[sid].key_step = step;
+    unsigned step = shard_index[sid].key_step = key_index.bonus ? 1 : key_index.step;
     auto& slices = shard_index[sid].slices;
     if (!slices.empty() && slices.back().second == i) {
       slices.back().second = i + step;
@@ -280,9 +276,6 @@ void Transaction::InitByKeys(const KeyIndex& key_index) {
     return;
   }
 
-  VLOG(0) << "InitByKeys " << key_index.start << " " << key_index.end << " " << key_index.step
-          << " -> " << key_index.NumArgs();
-
   DCHECK_LT(key_index.start, full_args_.size());
 
   // Stub transactions always operate only on single shard.

From 5226f972993bbd3076a569ef6843ba94d29824da Mon Sep 17 00:00:00 2001
From: Vladislav Oleshko <vlad@dragonflydb.io>
Date: Fri, 19 Jul 2024 13:43:23 +0300
Subject: [PATCH 4/4] fix: fixes

---
 src/server/transaction.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/server/transaction.cc b/src/server/transaction.cc
index dfdf69d231b2..150187403c7f 100644
--- a/src/server/transaction.cc
+++ b/src/server/transaction.cc
@@ -199,7 +199,8 @@ void Transaction::BuildShardIndex(const KeyIndex& key_index, std::vector<PerShar
     unique_slot_checker_.Add(key);
     ShardId sid = Shard(key, shard_data_.size());
 
-    unsigned step = shard_index[sid].key_step = key_index.bonus ? 1 : key_index.step;
+    unsigned step = key_index.bonus ? 1 : key_index.step;
+    shard_index[sid].key_step = step;
     auto& slices = shard_index[sid].slices;
     if (!slices.empty() && slices.back().second == i) {
       slices.back().second = i + step;