New account storage stat for accelerator

crypto/block/account-storage-stat.cpp

@@ -53,51 +53,83 @@ AccountStorageStat& AccountStorageStat::operator=(AccountStorageStat&& other) {
   return *this;
 }
 
-td::Result<AccountStorageStat::CellInfo> AccountStorageStat::add_root(const Ref<vm::Cell>& cell) {
-  roots_.push_back(cell);
-  return add_cell(cell);
-}
-
-td::Status AccountStorageStat::remove_root(const Ref<vm::Cell>& cell) {
-  auto it = std::find_if(roots_.begin(), roots_.end(),
-                         [&](const Ref<vm::Cell>& c) { return c->get_hash() == cell->get_hash(); });
-  if (it == roots_.end()) {
-    return td::Status::Error(PSTRING() << "no such root " << cell->get_hash().to_hex());
-  }
-  roots_.erase(it);
-  return remove_cell(cell);
-}
-
 td::Result<AccountStorageStat::CellInfo> AccountStorageStat::replace_roots(std::vector<Ref<vm::Cell>> new_roots) {
-  std::vector<Ref<vm::Cell>> old_roots = roots_;
+  std::erase_if(new_roots, [](const Ref<vm::Cell>& c) { return c.is_null(); });
+  auto cmp = [](const Ref<vm::Cell>& c1, const Ref<vm::Cell>& c2) { return c1->get_hash() < c2->get_hash(); };
+  std::sort(new_roots.begin(), new_roots.end(), cmp);
+  std::sort(roots_.begin(), roots_.end(), cmp);
+  std::vector<Ref<vm::Cell>> to_add, to_del;
+  std::set_difference(new_roots.begin(), new_roots.end(), roots_.begin(), roots_.end(), std::back_inserter(to_add),
+                      cmp);
+  std::set_difference(roots_.begin(), roots_.end(), new_roots.begin(), new_roots.end(), std::back_inserter(to_del),
+                      cmp);
+
   td::uint32 max_merkle_depth = 0;
-  for (const Ref<vm::Cell>& root : new_roots) {
+  for (const Ref<vm::Cell>& root : to_add) {
-    if (root.is_null()) {
+    TRY_RESULT(info, add_cell(root));
-      continue;
-    }
-    TRY_RESULT(info, add_root(root));
     max_merkle_depth = std::max(max_merkle_depth, info.max_merkle_depth);
   }
-  for (const Ref<vm::Cell>& root : old_roots) {
+  for (const Ref<vm::Cell>& root : to_del) {
-    TRY_STATUS(remove_root(root));
+    TRY_STATUS(remove_cell(root));
   }
+
+  roots_ = std::move(new_roots);
+  td::Status S = td::Status::OK();
+  cache_.for_each([&](Entry& e) {
+    if (S.is_ok()) {
+      S = commit_entry(e);
+    }
+  });
+  TRY_STATUS(std::move(S));
   return CellInfo{max_merkle_depth};
 }
 
+void AccountStorageStat::add_hint(const td::HashSet<vm::CellHash>& hint) {
+  td::HashSet<vm::CellHash> visited;
+  std::function<void(const Ref<vm::Cell>&, bool)> dfs = [&](const Ref<vm::Cell>& cell, bool is_root) {
+    if (!visited.insert(cell->get_hash()).second) {
+      return;
+    }
+    Entry& e = get_entry(cell);
+    e.exists = e.exists_known = true;
+    if (is_root) {
+      fetch_entry(e).ignore();
+      if (e.max_merkle_depth && e.max_merkle_depth.value() != 0) {
+        return;
+      }
+    }
+    if (hint.contains(cell->get_hash())) {
+      bool spec;
+      vm::CellSlice cs = vm::load_cell_slice_special(cell, spec);
+      for (unsigned i = 0; i < cs.size_refs(); ++i) {
+        dfs(cs.prefetch_ref(i), false);
+      }
+    }
+  };
+  for (const Ref<vm::Cell>& root : roots_) {
+    dfs(root, true);
+  }
+}
+
 td::Result<AccountStorageStat::CellInfo> AccountStorageStat::add_cell(const Ref<vm::Cell>& cell) {
   Entry& e = get_entry(cell);
-  ++e.refcnt;
+  if (!e.exists_known || e.refcnt_diff < 0) {
-  if (e.refcnt == 0) {
+    TRY_STATUS(fetch_entry(e));
-    return td::Status::Error(PSTRING() << "cell " << cell->get_hash().to_hex() << ": refcnt overflow");
   }
-  if (e.refcnt != 1) {
+  ++e.refcnt_diff;
-    update_dict(e);
+  if (e.exists || e.refcnt_diff > 1 || (e.refcnt && e.refcnt.value() + e.refcnt_diff != 1)) {
-    return CellInfo{e.max_merkle_depth};
+    if (!e.max_merkle_depth) {
+      TRY_STATUS(fetch_entry(e));
+      if (!e.max_merkle_depth) {
+        return td::Status::Error(PSTRING() << "unexpected unknown Merkle depth of cell " << cell->get_hash());
+      }
+    }
+    return CellInfo{e.max_merkle_depth.value()};
   }
+
   td::uint32 max_merkle_depth = 0;
-  vm::CellSlice cs{vm::NoVm{}, cell};
+  bool spec;
-  ++total_cells_;
+  vm::CellSlice cs = vm::load_cell_slice_special(cell, spec);
-  total_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);
@@ -109,64 +141,99 @@ 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;
-  update_dict(e2);
   return CellInfo{max_merkle_depth};
 }
 
 td::Status AccountStorageStat::remove_cell(const Ref<vm::Cell>& cell) {
   Entry& e = get_entry(cell);
-  if (e.refcnt == 0) {
+  if (!e.exists_known) {
-    return td::Status::Error(PSTRING() << "cell " << cell->get_hash().to_hex() << " is not in the dict");
+    TRY_STATUS(fetch_entry(e));
   }
-  --e.refcnt;
+  if (!e.exists) {
-  update_dict(e);
+    return td::Status::Error(PSTRING() << "Failed to remove cell " << cell->get_hash().to_hex()
-  if (e.refcnt != 0) {
+                                       << " : does not exist in the dict");
+  }
+  --e.refcnt_diff;
+  if (e.refcnt_diff < 0 && !e.refcnt) {
+    TRY_STATUS(fetch_entry(e));
+  }
+  if (e.refcnt.value() + e.refcnt_diff != 0) {
     return td::Status::OK();
   }
-  vm::CellSlice cs{vm::NoVm{}, std::move(cell)};
+  bool spec;
-  if (total_cells_ == 0 || total_bits_ < cs.size()) {
+  vm::CellSlice cs = vm::load_cell_slice_special(cell, spec);
-    return td::Status::Error("total_cell/total_bits becomes negative");
-  }
-  --total_cells_;
-  total_bits_ -= cs.size();
   for (unsigned i = 0; i < cs.size_refs(); ++i) {
     TRY_STATUS(remove_cell(cs.prefetch_ref(i)));
   }
   return td::Status::OK();
 }
 
-bool AccountStorageStat::Entry::serialize(vm::CellBuilder& cb) const {
-  return cb.store_long_bool(refcnt, 32) && cb.store_long_bool(max_merkle_depth, 2);
-}
-
-void AccountStorageStat::Entry::fetch(Ref<vm::CellSlice> cs) {
-  if (cs.is_null()) {
-    refcnt = max_merkle_depth = 0;
-  } else {
-    refcnt = (td::uint32)cs.write().fetch_ulong(32);
-    max_merkle_depth = (td::uint32)cs.write().fetch_ulong(2);
-  }
-}
-
 AccountStorageStat::Entry& AccountStorageStat::get_entry(const Ref<vm::Cell>& cell) {
   return cache_.apply(cell->get_hash().as_slice(), [&](Entry& e) {
     if (e.inited) {
       return;
     }
     e.inited = true;
-    e.hash = cell->get_hash();
+    e.cell = cell;
-    e.fetch(dict_.lookup(e.hash.as_bitslice()));
   });
 }
 
-void AccountStorageStat::update_dict(const Entry& e) {
+td::Status AccountStorageStat::fetch_entry(Entry& e) {
-  if (e.refcnt == 0) {
+  if (e.exists_known && e.refcnt && (!e.exists || e.max_merkle_depth)) {
-    dict_.lookup_delete(e.hash.as_bitslice());
+    return td::Status::OK();
-  } else {
-    vm::CellBuilder cb;
-    CHECK(e.serialize(cb));
-    dict_.set_builder(e.hash.as_bitslice(), cb);
   }
+  auto cs = dict_.lookup(e.cell->get_hash().as_bitslice());
+  if (cs.is_null()) {
+    e.exists = false;
+    e.refcnt = 0;
+  } else {
+    if (cs->size_ext() != 32 + 2) {
+      return td::Status::Error(PSTRING() << "invalid record for cell " << e.cell->get_hash().to_hex());
+    }
+    e.exists = true;
+    e.refcnt = (td::uint32)cs.write().fetch_ulong(32);
+    e.max_merkle_depth = (td::uint32)cs.write().fetch_ulong(2);
+    if (e.refcnt.value() == 0) {
+      return td::Status::Error(PSTRING() << "invalid refcnt=0 for cell " << e.cell->get_hash().to_hex());
+    }
+  }
+  e.exists_known = true;
+  return td::Status::OK();
+}
+
+td::Status AccountStorageStat::commit_entry(Entry& e) {
+  if (e.refcnt_diff == 0) {
+    return td::Status::OK();
+  }
+  TRY_STATUS(fetch_entry(e));
+  e.refcnt.value() += e.refcnt_diff;
+  e.refcnt_diff = 0;
+  bool spec;
+  if (e.refcnt.value() == 0) {
+    --total_cells_;
+    total_bits_ -= vm::load_cell_slice_special(e.cell, spec).size();
+    e.exists = false;
+    if (dict_.lookup_delete(e.cell->get_hash().as_bitslice()).is_null()) {
+      return td::Status::Error(PSTRING() << "Failed to delete entry " << e.cell->get_hash().to_hex());
+    }
+  } else {
+    if (!e.exists) {
+      ++total_cells_;
+      total_bits_ += vm::load_cell_slice_special(e.cell, spec).size();
+    }
+    e.exists = true;
+    if (!e.max_merkle_depth) {
+      return td::Status::Error(PSTRING() << "Failed to store entry " << e.cell->get_hash().to_hex()
+                                         << " : unknown merkle depth");
+    }
+    vm::CellBuilder cb;
+    dict_.set_builder(e.cell->get_hash().as_bitslice(), cb);
+    CHECK(cb.store_long_bool(e.refcnt.value(), 32) && cb.store_long_bool(e.max_merkle_depth.value(), 2));
+    if (!dict_.set_builder(e.cell->get_hash().as_bitslice(), cb)) {
+      return td::Status::Error(PSTRING() << "Failed to store entry " << e.cell->get_hash().to_hex());
+    }
+  }
+  return td::Status::OK();
 }
 
 }  // namespace block

crypto/block/account-storage-stat.h

@@ -58,9 +58,8 @@ class AccountStorageStat {
     return dict_.is_empty() ? td::Bits256::zero() : td::Bits256{dict_.get_root_cell()->get_hash().bits()};
   }
 
-  td::Result<CellInfo> add_root(const Ref<vm::Cell> &cell);
+  td::Result<CellInfo> replace_roots(std::vector<Ref<vm::Cell>> hint);
-  td::Status remove_root(const Ref<vm::Cell> &cell);
+  void add_hint(const td::HashSet<vm::CellHash> &visited);
-  td::Result<CellInfo> replace_roots(std::vector<Ref<vm::Cell>> new_roots);
 
  private:
   vm::Dictionary dict_;
@@ -72,15 +71,14 @@ class AccountStorageStat {
 
   struct Entry {
     bool inited = false;
-    vm::Cell::Hash hash;
+    Ref<vm::Cell> cell;
-    td::uint32 refcnt = 0;
+    bool exists_known = false;
-    td::uint32 max_merkle_depth = 0;
+    bool exists = false;
-
+    td::optional<td::uint32> refcnt, max_merkle_depth;
-    void fetch(Ref<vm::CellSlice> cs);
+    td::int32 refcnt_diff = 0;
-    bool serialize(vm::CellBuilder &cb) const;
 
     vm::Cell::Hash key() const {
-      return hash;
+      return cell->get_hash();
     }
     bool operator<(const Entry &other) const {
       return key() < other.key();
@@ -111,7 +109,8 @@ class AccountStorageStat {
   vm::CellHashTable<Entry> cache_;
 
   Entry &get_entry(const Ref<vm::Cell> &cell);
-  void update_dict(const Entry &e);
+  td::Status fetch_entry(Entry &e);
+  td::Status commit_entry(Entry &e);
 
   static constexpr td::uint32 MERKLE_DEPTH_LIMIT = 3;
 };

crypto/block/block.tlb

@@ -852,6 +852,7 @@ top_block_descr#d5 proof_for:BlockIdExt signatures:(Maybe ^BlockSignatures)
 //  COLLATED DATA
 //
 top_block_descr_set#4ac789f3 collection:(HashmapE 96 ^TopBlockDescr) = TopBlockDescrSet;
+account_storage_dict_proof#37c1e3fc proof:^Cell = AccountStorageDictProof;
 
 //
 //  VALIDATOR MISBEHAVIOR COMPLAINTS

crypto/block/transaction.cpp

@@ -262,6 +262,7 @@ bool Account::unpack_storage_info(vm::CellSlice& cs) {
   } else {
     storage_dict_hash = {};
   }
+  orig_storage_dict_hash = storage_dict_hash;
   if (info.due_payment->prefetch_ulong(1) == 1) {
     vm::CellSlice& cs2 = info.due_payment.write();
     cs2.advance(1);
@@ -310,8 +311,8 @@ bool Account::unpack_state(vm::CellSlice& cs) {
     tock = z & 1;
     LOG(DEBUG) << "tick=" << tick << ", tock=" << tock;
   }
-  code = state.code->prefetch_ref();
+  code = orig_code = state.code->prefetch_ref();
-  data = state.data->prefetch_ref();
+  data = orig_data = state.data->prefetch_ref();
   library = orig_library = state.library->prefetch_ref();
   return true;
 }
@@ -531,7 +532,7 @@ bool Account::init_new(ton::UnixTime now) {
   now_ = now;
   last_paid = 0;
   storage_used = {};
-  storage_dict_hash = {};
+  orig_storage_dict_hash = storage_dict_hash = {};
   due_payment = td::zero_refint();
   balance.set_zero();
   if (my_addr_exact.is_null()) {
@@ -562,6 +563,113 @@ bool Account::init_new(ton::UnixTime now) {
   return true;
 }
 
+/**
+ * Removes extra currencies dict from AccountStorage.
+ *
+ * This is used for computing account storage stats.
+ *
+ * @param storage_cs AccountStorage as CellSlice.
+ *
+ * @returns AccountStorage without extra currencies as CellSlice.
+ */
+static td::Ref<vm::CellSlice> storage_without_extra_currencies(td::Ref<vm::CellSlice> storage_cs) {
+  block::gen::AccountStorage::Record rec;
+  if (!block::gen::csr_unpack(storage_cs, rec)) {
+    LOG(ERROR) << "failed to unpack AccountStorage";
+    return {};
+  }
+  if (rec.balance->size_refs() > 0) {
+    block::gen::CurrencyCollection::Record balance;
+    if (!block::gen::csr_unpack(rec.balance, balance)) {
+      LOG(ERROR) << "failed to unpack AccountStorage";
+      return {};
+    }
+    balance.other = vm::CellBuilder{}.store_zeroes(1).as_cellslice_ref();
+    if (!block::gen::csr_pack(rec.balance, balance)) {
+      LOG(ERROR) << "failed to pack AccountStorage";
+      return {};
+    }
+  }
+  td::Ref<vm::CellSlice> result;
+  if (!block::gen::csr_pack(result, rec)) {
+    LOG(ERROR) << "failed to pack AccountStorage";
+    return {};
+  }
+  return result;
+}
+
+/**
+ * Computes storage dict of the account from scratch.
+ * This requires storage_dict_hash to be set, as it guarantees that the stored storage_used was computed recently
+ * (in older versions it included extra currency balance, in newer versions it does not).
+ *
+ * @returns Root of the dictionary, or Error
+ */
+td::Result<Ref<vm::Cell>> Account::compute_account_storage_dict() const {
+  if (storage.is_null()) {
+    return td::Status::Error("cannot compute storage dict: empty storage");
+  }
+  if (!storage_dict_hash) {
+    return td::Status::Error("cannot compute storage dict: storage_dict_hash is not set");
+  }
+  AccountStorageStat stat;
+  auto storage_for_stat = storage_without_extra_currencies(storage);
+  if (storage_for_stat.is_null()) {
+    return td::Status::Error("cannot compute storage dict: invalid storage");
+  }
+  TRY_STATUS(stat.replace_roots(storage_for_stat->prefetch_all_refs()).move_as_status());
+  // Root of AccountStorage is not counted in AccountStorageStat
+  td::uint64 expected_cells = stat.get_total_cells() + 1;
+  td::uint64 expected_bits = stat.get_total_bits() + storage->size();
+  if (expected_cells != storage_used.cells || expected_bits != storage_used.bits) {
+    return td::Status::Error(PSTRING() << "invalid storage_used: computed cells=" << expected_cells
+                                       << " bits=" << expected_bits << ", found cells" << storage_used.cells
+                                       << " bits=" << storage_used.bits);
+  }
+  if (storage_dict_hash.value() != stat.get_dict_hash()) {
+    return td::Status::Error(PSTRING() << "invalid storage dict hash: computed " << stat.get_dict_hash().to_hex()
+                                       << ", found " << storage_dict_hash.value().to_hex());
+  }
+  return stat.get_dict_root();
+}
+
+/**
+ * Initializes account_storage_stat of the account using the existing dict_root.
+ * This is not strictly necessary, as the storage stat is recomputed in Transaction.
+ * However, it can be used to optimize cell usage.
+ * This requires storage_dict_hash to be set, as it guarantees that the stored storage_used was computed recently
+ * (in older versions it included extra currency balance, in newer versions it does not).
+ *
+ * @param dict_root Root of the storage dictionary.
+ *
+ * @returns Status of the operation.
+ */
+td::Status Account::init_account_storage_stat(Ref<vm::Cell> dict_root) {
+  if (storage.is_null()) {
+    if (dict_root.not_null()) {
+      return td::Status::Error("storage is null, but dict_root is not null");
+    }
+    account_storage_stat = {};
+    return td::Status::OK();
+  }
+  if (!storage_dict_hash) {
+    return td::Status::Error("cannot init storage dict: storage_dict_hash is not set");
+  }
+  // Root of AccountStorage is not counted in AccountStorageStat
+  if (storage_used.cells < 1 || storage_used.bits < storage->size()) {
+    return td::Status::Error(PSTRING() << "storage_used is too small: cells=" << storage_used.cells
+                                       << " bits=" << storage_used.bits << " storage_root_bits=" << storage->size());
+  }
+  AccountStorageStat new_stat(std::move(dict_root), storage->prefetch_all_refs(), storage_used.cells - 1,
+                              storage_used.bits - storage->size());
+  if (storage_dict_hash.value() != new_stat.get_dict_hash()) {
+    return td::Status::Error(PSTRING() << "invalid storage dict hash: computed " << new_stat.get_dict_hash().to_hex()
+                                       << ", found " << storage_dict_hash.value().to_hex());
+  }
+  account_storage_stat = std::move(new_stat);
+  return td::Status::OK();
+}
+
 /**
  * Resets the split depth of the account.
  *
@@ -1814,6 +1922,7 @@ bool Transaction::prepare_compute_phase(const ComputePhaseConfig& cfg) {
                << cfg.flat_gas_limit << "]; remaining balance=" << balance.to_str();
     CHECK(td::sgn(balance.grams) >= 0);
   }
+  cp.vm_loaded_cells = vm.extract_loaded_cells();
   return true;
 }
 
@@ -2956,6 +3065,9 @@ td::Status Transaction::check_state_limits(const SizeLimitsConfig& size_limits,
   {
     TD_PERF_COUNTER(transaction_storage_stat_a);
     td::Timer timer;
+    if (update_storage_stat && compute_phase) {
+      storage_stat.add_hint(compute_phase->vm_loaded_cells);
+    }
     TRY_RESULT(info, storage_stat.replace_roots({new_code, new_data, new_library}));
     if (info.max_merkle_depth > max_allowed_merkle_depth) {
       return td::Status::Error("too big Merkle depth");
@@ -3131,41 +3243,6 @@ bool Account::store_acc_status(vm::CellBuilder& cb, int acc_status) const {
   return cb.store_long_bool(v, 2);
 }
 
-/**
- * Removes extra currencies dict from AccountStorage.
- *
- * This is used for computing account storage stats.
- *
- * @param storage_cs AccountStorage as CellSlice.
- *
- * @returns AccountStorage without extra currencies as CellSlice.
- */
-static td::Ref<vm::CellSlice> storage_without_extra_currencies(td::Ref<vm::CellSlice> storage_cs) {
-  block::gen::AccountStorage::Record rec;
-  if (!block::gen::csr_unpack(storage_cs, rec)) {
-    LOG(ERROR) << "failed to unpack AccountStorage";
-    return {};
-  }
-  if (rec.balance->size_refs() > 0) {
-    block::gen::CurrencyCollection::Record balance;
-    if (!block::gen::csr_unpack(rec.balance, balance)) {
-      LOG(ERROR) << "failed to unpack AccountStorage";
-      return {};
-    }
-    balance.other = vm::CellBuilder{}.store_zeroes(1).as_cellslice_ref();
-    if (!block::gen::csr_pack(rec.balance, balance)) {
-      LOG(ERROR) << "failed to pack AccountStorage";
-      return {};
-    }
-  }
-  td::Ref<vm::CellSlice> result;
-  if (!block::gen::csr_pack(result, rec)) {
-    LOG(ERROR) << "failed to pack AccountStorage";
-    return {};
-  }
-  return result;
-}
-
 namespace transaction {
 /**
  * Computes the new state of the account.
@@ -3284,6 +3361,9 @@ bool Transaction::compute_state(const SerializeConfig& cfg) {
     }
     AccountStorageStat& stats = new_account_storage_stat.value_force();
     // Don't check Merkle depth and size here - they were checked in check_state_limits
+    if (compute_phase) {
+      stats.add_hint(compute_phase->vm_loaded_cells);
+    }
     auto S = stats.replace_roots(new_storage->prefetch_all_refs()).move_as_status();
     if (S.is_error()) {
       LOG(ERROR) << "Cannot recompute storage stats for account " << account.addr.to_hex() << ": " << S.move_as_error();
@@ -3306,6 +3386,26 @@ bool Transaction::compute_state(const SerializeConfig& cfg) {
   if (!cfg.store_storage_dict_hash) {
     new_storage_dict_hash = {};
   }
+  if (false) {
+    vm::CellStorageStat control_stats;
+    control_stats.add_used_storage(new_storage);
+    if (control_stats.bits != new_storage_used.bits || control_stats.cells != new_storage_used.cells) {
+      LOG(ERROR) << " [ QQQQQQ 1 Wrong storage stat " << account.workchain << ":" << account.addr.to_hex() << " "
+                 << start_lt << " : " << new_storage_used.cells << "," << new_storage_used.bits
+                 << " != " << control_stats.cells << "," << control_stats.bits << " ] ";
+      return false;
+    }
+    AccountStorageStat control_stats_2;
+    control_stats_2.replace_roots(new_storage->prefetch_all_refs());
+    if (control_stats_2.get_total_bits() + new_storage->size() != new_storage_used.bits ||
+        control_stats_2.get_total_cells() + 1 != new_storage_used.cells) {
+      LOG(ERROR) << " [ QQQQQQ 2 Wrong storage stat " << account.workchain << ":" << account.addr.to_hex() << " "
+                 << start_lt << " : " << new_storage_used.cells << "," << new_storage_used.bits
+                 << " != " << control_stats_2.get_total_cells() + 1 << ","
+                 << control_stats_2.get_total_bits() + new_storage->size() << " ] ";
+      return false;
+    }
+  }
 
   CHECK(cb.store_long_bool(1, 1)                                 // account$1
         && cb.append_cellslice_bool(account.my_addr)             // addr:MsgAddressInt

crypto/block/transaction.h

@@ -208,6 +208,7 @@ struct ComputePhase {
   Ref<vm::Cell> actions;
   std::string vm_log;
   td::optional<td::uint64> precompiled_gas_usage;
+  td::HashSet<vm::CellHash> vm_loaded_cells;
 };
 
 struct ActionPhase {
@@ -272,6 +273,7 @@ struct Account {
   ton::UnixTime last_paid;
   StorageUsed storage_used;
   td::optional<td::Bits256> storage_dict_hash;
+  td::optional<td::Bits256> orig_storage_dict_hash;
   td::optional<AccountStorageStat> account_storage_stat;
 
   block::CurrencyCollection balance;
@@ -281,7 +283,8 @@ struct Account {
   Ref<vm::CellSlice> storage;      // AccountStorage
   Ref<vm::CellSlice> inner_state;  // StateInit
   ton::Bits256 state_hash;         // hash of StateInit for frozen accounts
-  Ref<vm::Cell> code, data, library, orig_library;
+  Ref<vm::Cell> code, data, library;
+  Ref<vm::Cell> orig_code, orig_data, orig_library;
   std::vector<LtCellRef> transactions;
   Account() = default;
   Account(ton::WorkchainId wc, td::ConstBitPtr _addr) : workchain(wc), addr(_addr) {
@@ -295,6 +298,8 @@ struct Account {
   bool set_address(ton::WorkchainId wc, td::ConstBitPtr new_addr);
   bool unpack(Ref<vm::CellSlice> account, ton::UnixTime now, bool special);
   bool init_new(ton::UnixTime now);
+  td::Result<Ref<vm::Cell>> compute_account_storage_dict() const;
+  td::Status init_account_storage_stat(Ref<vm::Cell> dict_root);
   bool deactivate();
   bool recompute_tmp_addr(Ref<vm::CellSlice>& tmp_addr, int split_depth, td::ConstBitPtr orig_addr_rewrite) const;
   td::RefInt256 compute_storage_fees(ton::UnixTime now, const std::vector<block::StoragePrices>& pricing) const;

crypto/vm/boc.cpp

@@ -1259,14 +1259,6 @@ bool VmStorageStat::add_storage(const CellSlice& cs) {
   return true;
 }
 
-static td::uint64 estimate_prunned_size() {
-  return 41;
-}
-
-static td::uint64 estimate_serialized_size(const Ref<DataCell>& cell) {
-  return cell->get_serialized_size() + cell->size_refs() * 3 + 3;
-}
-
 void ProofStorageStat::add_cell(const Ref<DataCell>& cell) {
   auto& status = cells_[cell->get_hash()];
   if (status == c_loaded) {
@@ -1290,4 +1282,17 @@ td::uint64 ProofStorageStat::estimate_proof_size() const {
   return proof_size_;
 }
 
+ProofStorageStat::CellStatus ProofStorageStat::get_cell_status(const Cell::Hash& hash) const {
+  auto it = cells_.find(hash);
+  return it == cells_.end() ? c_none : it->second;
+}
+
+td::uint64 ProofStorageStat::estimate_prunned_size() {
+  return 41;
+}
+
+td::uint64 ProofStorageStat::estimate_serialized_size(const Ref<DataCell>& cell) {
+  return cell->get_serialized_size() + cell->size_refs() * 3 + 3;
+}
+
 }  // namespace vm

crypto/vm/boc.h

@@ -167,11 +167,17 @@ class ProofStorageStat {
  public:
   void add_cell(const Ref<DataCell>& cell);
   td::uint64 estimate_proof_size() const;
- private:
+
   enum CellStatus {
     c_none = 0, c_prunned = 1, c_loaded = 2
   };
-  td::HashMap<vm::Cell::Hash, CellStatus> cells_;
+  CellStatus get_cell_status(const Cell::Hash& hash) const;
+
+  static td::uint64 estimate_prunned_size();
+  static td::uint64 estimate_serialized_size(const Ref<DataCell>& cell);
+
+ private:
+  td::HashMap<Cell::Hash, CellStatus> cells_;
   td::uint64 proof_size_ = 0;
 };
 

crypto/vm/cells/CellUsageTree.cpp

@@ -112,7 +112,7 @@ void CellUsageTree::set_use_mark_for_is_loaded(bool use_mark) {
 }
 
 void CellUsageTree::on_load(NodeId node_id, const td::Ref<vm::DataCell>& cell) {
-  if (nodes_[node_id].is_loaded) {
+  if (ignore_loads_ || nodes_[node_id].is_loaded) {
     return;
   }
   nodes_[node_id].is_loaded = true;

crypto/vm/cells/CellUsageTree.h

@@ -66,6 +66,9 @@ class CellUsageTree : public std::enable_shared_from_this<CellUsageTree> {
   void set_cell_load_callback(std::function<void(const td::Ref<vm::DataCell>&)> f) {
     cell_load_callback_ = std::move(f);
   }
+  void set_ignore_loads(bool value) {
+    ignore_loads_ = value;
+  }
 
  private:
   struct Node {
@@ -80,5 +83,6 @@ class CellUsageTree : public std::enable_shared_from_this<CellUsageTree> {
 
   void on_load(NodeId node_id, const td::Ref<vm::DataCell>& cell);
   NodeId create_node(NodeId parent);
+  bool ignore_loads_ = false;
 };
 }  // namespace vm

crypto/vm/vm.h

@@ -19,6 +19,7 @@
 #pragma once
 
 #include "common/refcnt.hpp"
+#include "td/utils/HashMap.h"
 #include "vm/cellslice.h"
 #include "vm/stack.hpp"
 #include "vm/vmstate.h"
@@ -424,6 +425,10 @@ class VmState final : public VmStateInterface {
     }
   }
 
+  td::HashSet<CellHash> extract_loaded_cells() {
+    return std::move(loaded_cells);
+  }
+
  private:
   void init_cregs(bool same_c3 = false, bool push_0 = true);
   int run_inner();

validator/impl/collator-impl.h

@@ -220,6 +220,15 @@ class Collator final : public td::actor::Actor {
   std::vector<vm::MerkleProofBuilder> neighbor_proof_builders_;
   std::vector<Ref<vm::Cell>> collated_roots_;
 
+  struct AccountStorageDict {
+    bool inited = false;
+    vm::MerkleProofBuilder mpb;
+    Ref<vm::Cell> proof_root;
+    size_t proof_size_estimate = 0;
+    bool add_to_collated_data = false;
+  };
+  std::map<td::Bits256, AccountStorageDict> account_storage_dicts_;
+
   std::unique_ptr<ton::BlockCandidate> block_candidate;
 
   std::unique_ptr<vm::AugmentedDictionary> dispatch_queue_, old_dispatch_queue_;
@@ -245,6 +254,7 @@ class Collator final : public td::actor::Actor {
   block::Account* lookup_account(td::ConstBitPtr addr) const;
   std::unique_ptr<block::Account> make_account_from(td::ConstBitPtr addr, Ref<vm::CellSlice> account,
                                                     bool force_create);
+  bool init_account_storage_dict(block::Account& account);
   td::Result<block::Account*> make_account(td::ConstBitPtr addr, bool force_create = false);
   td::actor::ActorId<Collator> get_self() {
     return actor_id(this);
@@ -344,6 +354,7 @@ class Collator final : public td::actor::Actor {
   bool register_dispatch_queue_op(bool force = false);
   bool update_account_dict_estimation(const block::transaction::Transaction& trans);
   bool update_min_mc_seqno(ton::BlockSeqno some_mc_seqno);
+  bool process_account_storage_dict(const block::Account& account);
   bool combine_account_transactions();
   bool update_public_libraries();
   bool update_account_public_libraries(Ref<vm::Cell> orig_libs, Ref<vm::Cell> final_libs, const td::Bits256& addr);

validator/impl/collator.cpp

@@ -2614,9 +2614,60 @@ std::unique_ptr<block::Account> Collator::make_account_from(td::ConstBitPtr addr
     return nullptr;
   }
   ptr->block_lt = start_lt;
+  if (!init_account_storage_dict(*ptr)) {
+    return nullptr;
+  }
   return ptr;
 }
 
+/**
+ * If full collated data is enabled, initialize account storage dict and prepare MerkleProofBuilder for it
+ *
+ * @param account Account to initialize storage dict for
+ * @return True on success, False on failure
+ */
+bool Collator::init_account_storage_dict(block::Account& account) {
+  if (!full_collated_data_ || is_masterchain() || !account.storage_dict_hash || account.storage.is_null()) {
+    return true;
+  }
+  if (account.storage_used.cells < 10) {  // TODO: some other threshold?
+    return true;
+  }
+  td::Bits256 storage_dict_hash = account.storage_dict_hash.value();
+  if (storage_dict_hash.is_zero()) {
+    return true;
+  }
+  AccountStorageDict& dict = account_storage_dicts_[storage_dict_hash];
+  if (!dict.inited) {
+    dict.inited = true;
+    // don't mark cells in account state as loaded during compute_account_storage_dict
+    state_usage_tree_->set_ignore_loads(true);
+    auto res = account.compute_account_storage_dict();
+    state_usage_tree_->set_ignore_loads(false);
+    if (res.is_error()) {
+      return fatal_error(res.move_as_error_prefix(PSTRING() << "Failed to init account storage dict for "
+                                                            << account.addr.to_hex() << ": "));
+    }
+    if (res.ok().is_null()) {  // Impossible if storage_dict_hash is not zero
+      return fatal_error(PSTRING() << "Failed to init account storage dict for " << account.addr.to_hex()
+                                   << ": dict is empty");
+    }
+    dict.mpb = vm::MerkleProofBuilder(res.move_as_ok());
+    dict.mpb.set_cell_load_callback([&](const td::Ref<vm::DataCell>& cell) {
+      if (block_limit_status_) {
+        block_limit_status_->collated_data_stat.add_cell(cell);
+      }
+    });
+  }
+  auto S = account.init_account_storage_stat(dict.mpb.root());
+  if (S.is_error()) {
+    return fatal_error(S.move_as_error_prefix(PSTRING() << "Failed to init account storage dict for "
+                                                        << account.addr.to_hex() << ": "));
+  }
+  return true;
+}
+
+
 /**
  * Looks up an account in the Collator's account map.
  *
@@ -2666,12 +2717,99 @@ td::Result<block::Account*> Collator::make_account(td::ConstBitPtr addr, bool fo
   return ins.first->second.get();
 }
 
+/**
+ * Decides whether to include storage dict proof to collated data for this account or not.
+ *
+ * @param account Account object
+ *
+ * @returns True if the operation is successful, false otherwise.
+ */
+bool Collator::process_account_storage_dict(const block::Account& account) {
+  if (!account.orig_storage_dict_hash) {
+    return true;
+  }
+  td::Bits256 storage_dict_hash = account.orig_storage_dict_hash.value();
+  auto it = account_storage_dicts_.find(storage_dict_hash);
+  if (it == account_storage_dicts_.end()) {
+    return true;
+  }
+  CHECK(full_collated_data_ && !is_masterchain());
+  AccountStorageDict& dict = it->second;
+  if (dict.add_to_collated_data) {
+    LOG(DEBUG) << "Storage dict proof of account " << account.addr.to_hex() << " : already included";
+    return true;
+  }
+
+  td::HashSet<vm::CellHash> visited;
+  bool calculate_proof_size_diff = true;
+  td::int64 proof_size_diff = 0;
+  std::function<void(const Ref<vm::Cell>&)> dfs = [&](const Ref<vm::Cell>& cell) {
+    if (cell.is_null() || !visited.emplace(cell->get_hash()).second) {
+      return;
+    }
+    auto loaded_cell = cell->load_cell().move_as_ok();
+    if (calculate_proof_size_diff) {
+      switch (block_limit_status_->collated_data_stat.get_cell_status(cell->get_hash())) {
+        case vm::ProofStorageStat::c_none:
+          proof_size_diff += vm::ProofStorageStat::estimate_serialized_size(loaded_cell.data_cell);
+          break;
+        case vm::ProofStorageStat::c_prunned:
+          proof_size_diff -= vm::ProofStorageStat::estimate_prunned_size();
+          proof_size_diff += vm::ProofStorageStat::estimate_serialized_size(loaded_cell.data_cell);
+          break;
+        case vm::ProofStorageStat::c_loaded:
+          break;
+      }
+    }
+    vm::CellSlice cs{std::move(loaded_cell)};
+    for (unsigned i = 0; i < cs.size_refs(); ++i) {
+      dfs(cs.prefetch_ref(i));
+    }
+  };
+
+  // Visit all cells in the original account storage to calculate collated data increase
+  state_usage_tree_->set_ignore_loads(true);
+  dfs(account.orig_code);
+  dfs(account.orig_data);
+  dfs(account.orig_library);
+  state_usage_tree_->set_ignore_loads(false);
+
+  if (proof_size_diff > (td::int64)dict.proof_size_estimate) {
+    LOG(DEBUG) << "Storage dict proof of account " << account.addr.to_hex()
+               << " : account_proof_size=" << proof_size_diff << ", dict_proof_size=" << dict.proof_size_estimate
+               << ", include dict in collated data";
+    dict.add_to_collated_data = true;
+  } else {
+    LOG(DEBUG) << "Storage dict proof of account " << account.addr.to_hex()
+               << " : account_proof_size=" << proof_size_diff << ", dict_proof_size=" << dict.proof_size_estimate
+               << ", DO NOT include dict in collated data";
+    // Include account storage in collated data
+    calculate_proof_size_diff = false;
+    visited.clear();
+    dfs(account.orig_code);
+    dfs(account.orig_data);
+    dfs(account.orig_library);
+  }
+
+  return true;
+}
+
 /**
  * Combines account transactions and updates the ShardAccountBlocks and ShardAccounts.
  *
  * @returns True if the operation is successful, false otherwise.
  */
 bool Collator::combine_account_transactions() {
+  for (auto& [hash, dict] : account_storage_dicts_) {
+    auto res = dict.mpb.extract_proof();
+    if (res.is_error()) {
+      return fatal_error(res.move_as_error_prefix(PSTRING() << "Failed to generate proof for account storage dict "
+                                                            << hash.to_hex() << ": "));
+    }
+    dict.proof_root = res.move_as_ok();
+    dict.proof_size_estimate = vm::std_boc_serialize(dict.proof_root, 31).move_as_ok().size();
+  }
+
   vm::AugmentedDictionary dict{256, block::tlb::aug_ShardAccountBlocks};
   for (auto& z : accounts) {
     block::Account& acc = *(z.second);
@@ -2755,6 +2893,9 @@ bool Collator::combine_account_transactions() {
           }
         }
       }
+      if (!process_account_storage_dict(acc)) {
+        return false;
+      }
     } else {
       if (acc.total_state->get_hash() != acc.orig_total_state->get_hash()) {
         return fatal_error(std::string{"total state of account "} + z.first.to_hex() +
@@ -5885,6 +6026,19 @@ bool Collator::create_collated_data() {
   for (auto& p : proofs) {
     collated_roots_.push_back(std::move(p.second));
   }
+
+  // 5. Proofs for account storage dicts
+  for (auto& [_, dict] : account_storage_dicts_) {
+    if (!dict.add_to_collated_data) {
+      continue;
+    }
+    CHECK(dict.proof_root.not_null());
+    // account_storage_dict_proof#37c1e3fc proof:^Cell = AccountStorageDictProof;
+    collated_roots_.push_back(vm::CellBuilder()
+                                  .store_long(block::gen::AccountStorageDictProof::cons_tag[0], 32)
+                                  .store_ref(dict.proof_root)
+                                  .finalize_novm());
+  }
   return true;
 }
 

validator/impl/validate-query.cpp

@@ -633,6 +633,23 @@ bool ValidateQuery::extract_collated_data_from(Ref<vm::Cell> croot, int idx) {
     top_shard_descr_dict_ = std::make_unique<vm::Dictionary>(cs.prefetch_ref(), 96);
     return true;
   }
+  if (block::gen::t_AccountStorageDictProof.has_valid_tag(cs)) {
+    if (!block::gen::t_AccountStorageDictProof.validate_upto(10000, cs)) {
+      return reject_query("invalid AccountStorageDictProof");
+    }
+    // account_storage_dict_proof#37c1e3fc proof:^Cell = AccountStorageDictProof;
+    Ref<vm::Cell> proof = cs.prefetch_ref();
+    auto virt_root = vm::MerkleProof::virtualize(proof, 1);
+    if (virt_root.is_null()) {
+      return reject_query("invalid Merkle proof in AccountStorageDictProof");
+    }
+    LOG(DEBUG) << "collated datum # " << idx << " is an AccountStorageDictProof with hash "
+               << virt_root->get_hash().to_hex();
+    if (!virt_account_storage_dicts_.emplace(virt_root->get_hash().bits(), virt_root).second) {
+      return reject_query("duplicate AccountStorageDictProof");
+    }
+    return true;
+  }
   LOG(WARNING) << "collated datum # " << idx << " has unknown type (magic " << cs.prefetch_ulong(32) << "), ignoring";
   return true;
 }
@@ -5253,6 +5270,18 @@ std::unique_ptr<block::Account> ValidateQuery::unpack_account(td::ConstBitPtr ad
                            << " does not really belong to current shard");
     return {};
   }
+  if (new_acc->storage_dict_hash) {
+    auto it = virt_account_storage_dicts_.find(new_acc->storage_dict_hash.value());
+    if (it != virt_account_storage_dicts_.end()) {
+      LOG(DEBUG) << "Using account storage dict proof for account " << addr.to_hex(256)
+                 << ", hash=" << it->second->get_hash().to_hex();
+      auto S = new_acc->init_account_storage_stat(it->second);
+      if (S.is_error()) {
+        reject_query(PSTRING() << "Failed to init account storage stat for account " << addr.to_hex(256), std::move(S));
+        return {};
+      }
+    }
+  }
   return new_acc;
 }
 

validator/impl/validate-query.hpp

@@ -170,6 +170,7 @@ class ValidateQuery : public td::actor::Actor {
   std::vector<Ref<vm::Cell>> collated_roots_;
   std::map<RootHash, Ref<vm::Cell>> virt_roots_;
   std::unique_ptr<vm::Dictionary> top_shard_descr_dict_;
+  std::map<td::Bits256, Ref<vm::Cell>> virt_account_storage_dicts_;
 
   Ref<vm::CellSlice> shard_hashes_;              // from McBlockExtra
   Ref<vm::CellSlice> blk_config_params_;         // from McBlockExtra