Limit collated data size in collator

crypto/block/block.cpp

@@ -701,12 +701,19 @@ int BlockLimits::classify_lt(ton::LogicalTime lt) const {
   return lt_delta.classify(lt - start_lt);
 }
 
-int BlockLimits::classify(td::uint64 size, td::uint64 gas, ton::LogicalTime lt) const {
-  return std::max(std::max(classify_size(size), classify_gas(gas)), classify_lt(lt));
+int BlockLimits::classify_collated_data_size(td::uint64 size) const {
+  return bytes.classify(size);  // TODO: Maybe separate limits in config
 }
 
-bool BlockLimits::fits(unsigned cls, td::uint64 size, td::uint64 gas_value, ton::LogicalTime lt) const {
-  return bytes.fits(cls, size) && gas.fits(cls, gas_value) && lt_delta.fits(cls, lt - start_lt);
+int BlockLimits::classify(td::uint64 size, td::uint64 gas, ton::LogicalTime lt, td::uint64 collated_size) const {
+  return std::max(
+      {classify_size(size), classify_gas(gas), classify_lt(lt), classify_collated_data_size(collated_size)});
+}
+
+bool BlockLimits::fits(unsigned cls, td::uint64 size, td::uint64 gas_value, ton::LogicalTime lt,
+                       td::uint64 collated_size) const {
+  return bytes.fits(cls, size) && gas.fits(cls, gas_value) && lt_delta.fits(cls, lt - start_lt) &&
+         bytes.fits(cls, collated_size);
 }
 
 td::uint64 BlockLimitStatus::estimate_block_size(const vm::NewCellStorageStat::Stat* extra) const {
@@ -719,20 +726,22 @@ td::uint64 BlockLimitStatus::estimate_block_size(const vm::NewCellStorageStat::S
 }
 
 int BlockLimitStatus::classify() const {
-  return limits.classify(estimate_block_size(), gas_used, cur_lt);
+  return limits.classify(estimate_block_size(), gas_used, cur_lt, collated_data_stat.estimate_proof_size());
 }
 
 bool BlockLimitStatus::fits(unsigned cls) const {
   return cls >= ParamLimits::limits_cnt ||
          (limits.gas.fits(cls, gas_used) && limits.lt_delta.fits(cls, cur_lt - limits.start_lt) &&
-          limits.bytes.fits(cls, estimate_block_size()));
+          limits.bytes.fits(cls, estimate_block_size()) &&
+          limits.bytes.fits(cls, collated_data_stat.estimate_proof_size()));
 }
 
 bool BlockLimitStatus::would_fit(unsigned cls, ton::LogicalTime end_lt, td::uint64 more_gas,
                                  const vm::NewCellStorageStat::Stat* extra) const {
   return cls >= ParamLimits::limits_cnt || (limits.gas.fits(cls, gas_used + more_gas) &&
                                             limits.lt_delta.fits(cls, std::max(cur_lt, end_lt) - limits.start_lt) &&
-                                            limits.bytes.fits(cls, estimate_block_size(extra)));
+                                            limits.bytes.fits(cls, estimate_block_size(extra)) &&
+                                            limits.bytes.fits(cls, collated_data_stat.estimate_proof_size()));
 }
 
 // SETS: account_dict, shard_libraries_, mc_state_extra

crypto/block/block.h

@@ -252,8 +252,9 @@ struct BlockLimits {
   int classify_size(td::uint64 size) const;
   int classify_gas(td::uint64 gas) const;
   int classify_lt(ton::LogicalTime lt) const;
-  int classify(td::uint64 size, td::uint64 gas, ton::LogicalTime lt) const;
-  bool fits(unsigned cls, td::uint64 size, td::uint64 gas, ton::LogicalTime lt) const;
+  int classify_collated_data_size(td::uint64 size) const;
+  int classify(td::uint64 size, td::uint64 gas, ton::LogicalTime lt, td::uint64 collated_size) const;
+  bool fits(unsigned cls, td::uint64 size, td::uint64 gas, ton::LogicalTime lt, td::uint64 collated_size) const;
 };
 
 struct BlockLimitStatus {
@@ -262,6 +263,7 @@ struct BlockLimitStatus {
   td::uint64 gas_used{};
   vm::NewCellStorageStat st_stat;
   unsigned accounts{}, transactions{}, extra_out_msgs{};
+  vm::ProofStorageStat collated_data_stat;
   BlockLimitStatus(const BlockLimits& limits_, ton::LogicalTime lt = 0)
       : limits(limits_), cur_lt(std::max(limits_.start_lt, lt)) {
   }
@@ -271,6 +273,7 @@ struct BlockLimitStatus {
     transactions = accounts = 0;
     gas_used = 0;
     extra_out_msgs = 0;
+    collated_data_stat = {};
   }
   td::uint64 estimate_block_size(const vm::NewCellStorageStat::Stat* extra = nullptr) const;
   int classify() const;

crypto/vm/boc.cpp

@@ -1214,4 +1214,35 @@ 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) {
+    return;
+  }
+  if (status == c_prunned) {
+    proof_size_ -= estimate_prunned_size();
+  }
+  status = c_loaded;
+  proof_size_ += estimate_serialized_size(cell);
+  for (unsigned i = 0; i < cell->size_refs(); ++i) {
+    auto& child_status = cells_[cell->get_ref(i)->get_hash()];
+    if (child_status == c_none) {
+      child_status = c_prunned;
+      proof_size_ += estimate_prunned_size();
+    }
+  }
+}
+
+td::uint64 ProofStorageStat::estimate_proof_size() const {
+  return proof_size_;
+}
+
 }  // namespace vm

crypto/vm/boc.h

@@ -41,6 +41,7 @@ class NewCellStorageStat {
     Stat(td::uint64 cells_, td::uint64 bits_, td::uint64 internal_refs_ = 0, td::uint64 external_refs_ = 0)
         : cells(cells_), bits(bits_), internal_refs(internal_refs_), external_refs(external_refs_) {
     }
+    Stat(const Stat&) = default;
     td::uint64 cells{0};
     td::uint64 bits{0};
     td::uint64 internal_refs{0};
@@ -160,6 +161,18 @@ struct VmStorageStat {
   }
 };
 
+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
+  };
+  std::map<vm::Cell::Hash, CellStatus> cells_;
+  td::uint64 proof_size_ = 0;
+};
+
 struct CellSerializationInfo {
   bool special;
   Cell::LevelMask level_mask;

crypto/vm/cells/CellUsageTree.cpp

@@ -22,12 +22,12 @@ namespace vm {
 //
 // CellUsageTree::NodePtr
 //
-bool CellUsageTree::NodePtr::on_load() const {
+bool CellUsageTree::NodePtr::on_load(const td::Ref<vm::DataCell>& cell) const {
   auto tree = tree_weak_.lock();
   if (!tree) {
     return false;
   }
-  tree->on_load(node_id_);
+  tree->on_load(node_id_, cell);
   return true;
 }
 
@@ -111,8 +111,14 @@ void CellUsageTree::set_use_mark_for_is_loaded(bool use_mark) {
   use_mark_ = use_mark;
 }
 
-void CellUsageTree::on_load(NodeId node_id) {
+void CellUsageTree::on_load(NodeId node_id, const td::Ref<vm::DataCell>& cell) {
+  if (nodes_[node_id].is_loaded) {
+    return;
+  }
   nodes_[node_id].is_loaded = true;
+  if (cell_load_callback_) {
+    cell_load_callback_(cell);
+  }
 }
 
 CellUsageTree::NodeId CellUsageTree::create_child(NodeId node_id, unsigned ref_id) {

crypto/vm/cells/CellUsageTree.h

@@ -22,8 +22,12 @@
 
 #include "td/utils/int_types.h"
 #include "td/utils/logging.h"
+#include <functional>
 
 namespace vm {
+
+class DataCell;
+
 class CellUsageTree : public std::enable_shared_from_this<CellUsageTree> {
  public:
   using NodeId = td::uint32;
@@ -38,7 +42,7 @@ class CellUsageTree : public std::enable_shared_from_this<CellUsageTree> {
       return node_id_ == 0 || tree_weak_.expired();
     }
 
-    bool on_load() const;
+    bool on_load(const td::Ref<vm::DataCell>& cell) const;
     NodePtr create_child(unsigned ref_id) const;
     bool mark_path(CellUsageTree* master_tree) const;
     bool is_from_tree(const CellUsageTree* master_tree) const;
@@ -59,6 +63,10 @@ class CellUsageTree : public std::enable_shared_from_this<CellUsageTree> {
   void set_use_mark_for_is_loaded(bool use_mark = true);
   NodeId create_child(NodeId node_id, unsigned ref_id);
 
+  void set_cell_load_callback(std::function<void(const td::Ref<vm::DataCell>&)> f) {
+    cell_load_callback_ = std::move(f);
+  }
+
  private:
   struct Node {
     bool is_loaded{false};
@@ -68,8 +76,9 @@ class CellUsageTree : public std::enable_shared_from_this<CellUsageTree> {
   };
   bool use_mark_{false};
   std::vector<Node> nodes_{2};
+  std::function<void(const td::Ref<vm::DataCell>&)> cell_load_callback_;
 
-  void on_load(NodeId node_id);
+  void on_load(NodeId node_id, const td::Ref<vm::DataCell>& cell);
   NodeId create_node(NodeId parent);
 };
 }  // namespace vm

crypto/vm/cells/MerkleProof.h

@@ -66,6 +66,10 @@ class MerkleProofBuilder {
   td::Result<Ref<Cell>> extract_proof() const;
   bool extract_proof_to(Ref<Cell> &proof_root) const;
   td::Result<td::BufferSlice> extract_proof_boc() const;
+
+  void set_cell_load_callback(std::function<void(const td::Ref<vm::DataCell>&)> f) {
+    usage_tree->set_cell_load_callback(std::move(f));
+  }
 };
 
 }  // namespace vm

crypto/vm/cells/UsageCell.h

@@ -39,7 +39,7 @@ class UsageCell : public Cell {
   // load interface
   td::Result<LoadedCell> load_cell() const override {
     TRY_RESULT(loaded_cell, cell_->load_cell());
-    if (tree_node_.on_load()) {
+    if (tree_node_.on_load(loaded_cell.data_cell)) {
       CHECK(loaded_cell.tree_node.empty());
       loaded_cell.tree_node = tree_node_;
     }