New account storage stat for accelerator

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