Fix calculating storage dict, simplify storage stat code (#1643)

SpyCheese · web-flow · commit 5a86f0ac2d52 · 2025-05-09T19:12:21.000+03:00
diff --git a/crypto/block/account-storage-stat.cpp b/crypto/block/account-storage-stat.cpp
@@ -63,7 +63,7 @@ td::Status AccountStorageStat::replace_roots(std::vector<Ref<vm::Cell>> new_root
   for (const Ref<vm::Cell>& root : to_add) {
     TRY_RESULT(info, add_cell(root));
     if (check_merkle_depth && info.max_merkle_depth > MAX_MERKLE_DEPTH) {
-      return td::Status::Error("too big Merkle depth");
+      return td::Status::Error(errorcode_limits_exceeded, "too big Merkle depth");
     }
   }
   for (const Ref<vm::Cell>& root : to_del) {
@@ -92,10 +92,10 @@ void AccountStorageStat::add_hint(const td::HashSet<vm::CellHash>& hint) {
         return;
       }
     }
+    e.max_merkle_depth = 0;
     if (hint.contains(cell->get_hash())) {
       bool spec;
       vm::CellSlice cs = vm::load_cell_slice_special(cell, spec);
-      e.size_bits = cs.size();
       for (unsigned i = 0; i < cs.size_refs(); ++i) {
         dfs(cs.prefetch_ref(i), false);
       }
@@ -108,11 +108,11 @@ void AccountStorageStat::add_hint(const td::HashSet<vm::CellHash>& hint) {
 
 td::Result<AccountStorageStat::CellInfo> AccountStorageStat::add_cell(const Ref<vm::Cell>& cell) {
   Entry& e = get_entry(cell);
-  if (!e.exists_known || e.refcnt_diff < 0) {
+  if (!e.exists_known) {
     TRY_STATUS(fetch_from_dict(e));
   }
   ++e.refcnt_diff;
-  if (e.exists || e.refcnt_diff > 1 || (e.refcnt && e.refcnt.value() + e.refcnt_diff != 1)) {
+  if (e.exists || e.refcnt_diff > 1) {
     if (!e.max_merkle_depth) {
       TRY_STATUS(fetch_from_dict(e));
       if (!e.max_merkle_depth) {
@@ -125,7 +125,6 @@ td::Result<AccountStorageStat::CellInfo> AccountStorageStat::add_cell(const Ref<
   td::uint32 max_merkle_depth = 0;
   bool spec;
   vm::CellSlice cs = vm::load_cell_slice_special(cell, spec);
-  e.size_bits = cs.size();
   for (unsigned i = 0; i < cs.size_refs(); ++i) {
     TRY_RESULT(info, add_cell(cs.prefetch_ref(i)));
     max_merkle_depth = std::max(max_merkle_depth, info.max_merkle_depth);
@@ -137,6 +136,8 @@ td::Result<AccountStorageStat::CellInfo> AccountStorageStat::add_cell(const Ref<
   max_merkle_depth = std::min(max_merkle_depth, MERKLE_DEPTH_LIMIT);
   Entry& e2 = get_entry(cell);
   e2.max_merkle_depth = max_merkle_depth;
+  ++total_cells_;
+  total_bits_ += cs.size();
   return CellInfo{max_merkle_depth};
 }
 
@@ -158,16 +159,28 @@ td::Status AccountStorageStat::remove_cell(const Ref<vm::Cell>& cell) {
   }
   bool spec;
   vm::CellSlice cs = vm::load_cell_slice_special(cell, spec);
-  e.size_bits = cs.size();
   for (unsigned i = 0; i < cs.size_refs(); ++i) {
     TRY_STATUS(remove_cell(cs.prefetch_ref(i)));
   }
+  --total_cells_;
+  total_bits_ -= cs.size();
   return td::Status::OK();
 }
 
 td::Result<Ref<vm::Cell>> AccountStorageStat::get_dict_root() {
   if (!dict_up_to_date_) {
     std::vector<std::pair<td::ConstBitPtr, Ref<vm::CellBuilder>>> values;
+    if (parent_ && !parent_->dict_up_to_date_) {
+      for (auto& [_, ep] : parent_->cache_) {
+        if (ep.dict_refcnt_diff == 0) {
+          continue;
+        }
+        Entry& e = cache_[ep.hash];
+        if (!e.inited) {
+          e = ep;
+        }
+      }
+    }
     for (auto& [_, e] : cache_) {
       if (e.dict_refcnt_diff == 0) {
         continue;
@@ -259,25 +272,9 @@ td::Status AccountStorageStat::finalize_entry(Entry& e) {
   e.refcnt.value() += e.refcnt_diff;
   e.dict_refcnt_diff += e.refcnt_diff;
   e.refcnt_diff = 0;
-  if (e.refcnt.value() == 0) {
-    if (!e.size_bits) {
-      return td::Status::Error(PSTRING() << "Failed to store entry " << e.hash.to_hex() << " : unknown cell bits");
-    }
-    --total_cells_;
-    total_bits_ -= e.size_bits.value();
-    e.exists = false;
-  } else {
-    if (!e.exists) {
-      if (!e.size_bits) {
-        return td::Status::Error(PSTRING() << "Failed to store entry " << e.hash.to_hex() << " : unknown cell bits");
-      }
-      ++total_cells_;
-      total_bits_ += e.size_bits.value();
-    }
-    e.exists = true;
-    if (!e.max_merkle_depth) {
-      return td::Status::Error(PSTRING() << "Failed to store entry " << e.hash.to_hex() << " : unknown merkle depth");
-    }
+  e.exists = (e.refcnt.value() != 0);
+  if (e.exists && !e.max_merkle_depth) {
+    return td::Status::Error(PSTRING() << "Error on entry " << e.hash.to_hex() << " : unknown merkle depth");
   }
   return td::Status::OK();
 }
diff --git a/crypto/block/account-storage-stat.h b/crypto/block/account-storage-stat.h
@@ -59,6 +59,8 @@ class AccountStorageStat {
 
   void apply_child_stat(AccountStorageStat &&child);
 
+  static constexpr int errorcode_limits_exceeded = 999;
+
  private:
   vm::Dictionary dict_;
   bool dict_up_to_date_ = true;
@@ -76,7 +78,6 @@ class AccountStorageStat {
   struct Entry {
     bool inited = false;
     vm::CellHash hash;
-    td::optional<unsigned> size_bits;
     bool exists_known = false;
     bool exists = false;
     td::optional<td::uint32> refcnt, max_merkle_depth;
diff --git a/crypto/block/transaction.cpp b/crypto/block/transaction.cpp
@@ -1930,12 +1930,17 @@ bool Transaction::prepare_action_phase(const ActionPhaseConfig& cfg) {
   ap.action_fine = td::zero_refint();
 
   td::Ref<vm::Cell> old_code = new_code, old_data = new_data, old_library = new_library;
-  auto enforce_state_limits = [&]() {
+  // 1 - ok, 0 - limits exceeded, -1 - fatal error
+  auto enforce_state_limits = [&]() -> int {
     if (account.is_special) {
-      return true;
+      return 1;
     }
     auto S = check_state_limits(cfg.size_limits);
     if (S.is_error()) {
+      if (S.code() != AccountStorageStat::errorcode_limits_exceeded) {
+        LOG(ERROR) << "Account storage stat error: " << S.move_as_error();
+        return -1;
+      }
       // Rollback changes to state, fail action phase
       LOG(INFO) << "Account state size exceeded limits: " << S.move_as_error();
       new_account_storage_stat = {};
@@ -1944,9 +1949,9 @@ bool Transaction::prepare_action_phase(const ActionPhaseConfig& cfg) {
       new_library = old_library;
       ap.result_code = 50;
       ap.state_exceeds_limits = true;
-      return false;
+      return 0;
     }
-    return true;
+    return 1;
   };
 
   int n = 0;
@@ -2057,7 +2062,9 @@ bool Transaction::prepare_action_phase(const ActionPhaseConfig& cfg) {
       }
       LOG(DEBUG) << "invalid action " << ap.result_arg << " in action list: error code " << ap.result_code;
       // This is required here because changes to libraries are applied even if action phase fails
-      enforce_state_limits();
+      if (enforce_state_limits() == -1) {
+        return false;
+      }
       if (cfg.action_fine_enabled) {
         ap.action_fine = std::min(ap.action_fine, balance.grams);
         ap.total_action_fees = ap.action_fine;
@@ -2079,7 +2086,11 @@ bool Transaction::prepare_action_phase(const ActionPhaseConfig& cfg) {
     new_code = ap.new_code;
   }
   new_data = compute_phase->new_data;  // tentative persistent data update applied
-  if (!enforce_state_limits()) {
+  int res = enforce_state_limits();
+  if (res == -1) {
+    return false;
+  }
+  if (res == 0) {
     if (cfg.extra_currency_v2) {
       end_lt = ap.end_lt = start_lt + 1;
       if (cfg.action_fine_enabled) {
@@ -3054,7 +3065,8 @@ static td::uint32 get_public_libraries_diff_count(const td::Ref<vm::Cell>& old_l
  *
  * @returns A `td::Status` indicating the result of the check.
  *          - If the state limits are within the allowed range, returns OK.
- *          - If the state limits exceed the maximum allowed range, returns an error.
+ *          - If the state limits exceed the maximum allowed range, returns an error with AccountStorageStat::errorcode_limits_exceeded code.
+ *          - If an error occurred during storage stat calculation, returns other error.
  */
 td::Status Transaction::check_state_limits(const SizeLimitsConfig& size_limits, bool update_storage_stat) {
   auto cell_equal = [](const td::Ref<vm::Cell>& a, const td::Ref<vm::Cell>& b) -> bool {
@@ -3079,15 +3091,17 @@ td::Status Transaction::check_state_limits(const SizeLimitsConfig& size_limits,
 
   if (storage_stat.get_total_cells() > size_limits.max_acc_state_cells ||
       storage_stat.get_total_bits() > size_limits.max_acc_state_bits) {
-    return td::Status::Error(PSTRING() << "account state is too big: cells=" << storage_stat.get_total_cells()
+    return td::Status::Error(AccountStorageStat::errorcode_limits_exceeded,
+                             PSTRING() << "account state is too big: cells=" << storage_stat.get_total_cells()
                                        << ", bits=" << storage_stat.get_total_bits()
                                        << " (max cells=" << size_limits.max_acc_state_cells
                                        << ", max bits=" << size_limits.max_acc_state_bits << ")");
   }
   if (account.is_masterchain() && !cell_equal(account.library, new_library)) {
     auto libraries_count = get_public_libraries_count(new_library);
     if (libraries_count > size_limits.max_acc_public_libraries) {
-      return td::Status::Error(PSTRING() << "too many public libraries: " << libraries_count << " (max "
+      return td::Status::Error(AccountStorageStat::errorcode_limits_exceeded,
+                               PSTRING() << "too many public libraries: " << libraries_count << " (max "
                                          << size_limits.max_acc_public_libraries << ")");
     }
   }
