From 12db59090621feb2999b43c47baa182b8741bc92 Mon Sep 17 00:00:00 2001
From: telezhnaya <olyatelezhnaya@gmail.com>
Date: Fri, 8 Mar 2024 16:01:08 +0000
Subject: [PATCH] fix max seat number computation

It's still not 100% accurate, but it would give the result closer to the truth
---
 src/common.rs                   | 15 +++++++++++++++
 src/proposals/mod.rs            |  6 +-----
 src/validators/block_id/mod.rs  |  6 +-----
 src/validators/epoch_id/next.rs |  6 +-----
 4 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/src/common.rs b/src/common.rs
index 14f1bcd..9c090fa 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -1,5 +1,7 @@
 use color_eyre::eyre::{Context, ContextCompat};
 use near_cli_rs::types::near_token::NearToken;
+use near_jsonrpc_client::methods::EXPERIMENTAL_protocol_config::RpcProtocolConfigResponse;
+use near_primitives::types::NumSeats;
 use num_rational::Rational32;
 
 /// This implementation is ported from near-api-js:
@@ -16,6 +18,19 @@ pub fn find_seat_price(
     find_seat_price_for_protocol_after_49(stakes, max_number_of_seats, minimum_stake_ratio)
 }
 
+/// With statelessnet feature enabled, we need to take into account `num_chunk_only_producer_seats`.
+pub fn find_max_number_of_seats(protocol_config: &RpcProtocolConfigResponse) -> NumSeats {
+    let seats_before_statelessnet = protocol_config.num_block_producer_seats
+        + protocol_config
+            .avg_hidden_validator_seats_per_shard
+            .iter()
+            .sum::<u64>();
+    std::cmp::max(
+        seats_before_statelessnet,
+        protocol_config.num_chunk_only_producer_seats,
+    )
+}
+
 /// This implementation is ported from near-api-js:
 /// https://github.com/near/near-api-js/blob/bdbf839f54fbc399d7299da8cf9966bbdc426238/packages/utils/src/validators.ts#L24-L50
 fn find_seat_price_for_protocol_before_49(
diff --git a/src/proposals/mod.rs b/src/proposals/mod.rs
index 196b3e4..3fea1d1 100644
--- a/src/proposals/mod.rs
+++ b/src/proposals/mod.rs
@@ -133,11 +133,7 @@ pub fn display_proposals_info(
         })
         .wrap_err("Failed to get protocol config.")?;
 
-    let max_number_of_seats = protocol_config.num_block_producer_seats
-        + protocol_config
-            .avg_hidden_validator_seats_per_shard
-            .iter()
-            .sum::<u64>();
+    let max_number_of_seats = crate::common::find_max_number_of_seats(&protocol_config);
 
     let expected_seat_price = crate::common::find_seat_price(
         combine_validators_and_proposals_table
diff --git a/src/validators/block_id/mod.rs b/src/validators/block_id/mod.rs
index 7e887d4..80940e4 100644
--- a/src/validators/block_id/mod.rs
+++ b/src/validators/block_id/mod.rs
@@ -83,11 +83,7 @@ pub fn display_current_validators_info(
         .blocking_call(&RpcProtocolConfigRequest { block_reference })
         .wrap_err("Failed to get protocol config.")?;
 
-    let max_number_of_seats = protocol_config.num_block_producer_seats
-        + protocol_config
-            .avg_hidden_validator_seats_per_shard
-            .iter()
-            .sum::<u64>();
+    let max_number_of_seats = crate::common::find_max_number_of_seats(&protocol_config);
     eprintln!(
         "Validators (total: {}, seat price: {})",
         current_validators.len(),
diff --git a/src/validators/epoch_id/next.rs b/src/validators/epoch_id/next.rs
index 729caf6..9c8bf39 100644
--- a/src/validators/epoch_id/next.rs
+++ b/src/validators/epoch_id/next.rs
@@ -65,11 +65,7 @@ fn display_next_validators_info(
         })
         .wrap_err("Failed to get protocol config.")?;
 
-    let max_number_of_seats = protocol_config.num_block_producer_seats
-        + protocol_config
-            .avg_hidden_validator_seats_per_shard
-            .iter()
-            .sum::<u64>();
+    let max_number_of_seats = crate::common::find_max_number_of_seats(&protocol_config);
     let seat_price = crate::common::find_seat_price(
         next_validators
             .iter()