From 94032f6817a9f715adac53316aa551dd3b6b3901 Mon Sep 17 00:00:00 2001
From: Nadav Ivgi <nadav@shesek.info>
Date: Mon, 23 Mar 2020 19:30:13 +0200
Subject: [PATCH] Add Block::get_size() and Block::get_weight()

---
 src/blockdata/block.rs | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/blockdata/block.rs b/src/blockdata/block.rs
index 059ab9ae..4b272183 100644
--- a/src/blockdata/block.rs
+++ b/src/blockdata/block.rs
@@ -29,7 +29,8 @@ use util::uint::Uint256;
 use consensus::encode::Encodable;
 use network::constants::Network;
 use blockdata::transaction::Transaction;
-use blockdata::constants::max_target;
+use blockdata::constants::{max_target, WITNESS_SCALE_FACTOR};
+use VarInt;
 
 /// A block header, which contains all the block's information except
 /// the actual transactions
@@ -124,6 +125,21 @@ impl Block {
         );
         bitcoin_merkle_root(hashes).into()
     }
+
+    /// Get the size of the block
+    pub fn get_size(&self) -> usize {
+        // The size of the header + the size of the varint with the tx count + the txs themselves
+        let base_size = 80 + VarInt(self.txdata.len() as u64).len();
+        let txs_size: usize = self.txdata.iter().map(Transaction::get_size).sum();
+        base_size + txs_size
+    }
+
+    /// Get the weight of the block
+    pub fn get_weight(&self) -> usize {
+        let base_weight = WITNESS_SCALE_FACTOR * (80 + VarInt(self.txdata.len() as u64).len());
+        let txs_weight: usize = self.txdata.iter().map(Transaction::get_weight).sum();
+        base_weight + txs_weight
+    }
 }
 
 impl BlockHeader {
@@ -240,6 +256,9 @@ mod tests {
         assert_eq!(real_decode.header.nonce, 2067413810);
         // [test] TODO: check the transaction data
 
+        assert_eq!(real_decode.get_size(), some_block.len());
+        assert_eq!(real_decode.get_weight(), some_block.len() * 4);
+
         // should be also ok for a non-witness block as commitment is optional in that case
         assert!(real_decode.check_witness_commitment());
 
@@ -267,6 +286,9 @@ mod tests {
         assert_eq!(real_decode.header.nonce, 1879759182);
         // [test] TODO: check the transaction data
 
+        assert_eq!(real_decode.get_size(), segwit_block.len());
+        assert_eq!(real_decode.get_weight(), 17168);
+
         assert!(real_decode.check_witness_commitment());
 
         assert_eq!(serialize(&real_decode), segwit_block);