Feature: Count sigops for Transaction

bitcoin/src/blockdata/script/borrowed.rs

@@ -9,6 +9,7 @@ use core::ops::{Index, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, Ran
 use hashes::Hash;
 use secp256k1::{Secp256k1, Verification};
 
+use super::PushBytes;
 use crate::blockdata::opcodes::all::*;
 use crate::blockdata::opcodes::{self, Opcode};
 use crate::blockdata::script::witness_version::WitnessVersion;
@@ -189,6 +190,24 @@ impl Script {
             && self.0[24] == OP_CHECKSIG.to_u8()
     }
 
+    /// Checks whether a script is push only.
+    #[inline]
+    pub fn is_push_only(&self) -> bool {
+        for inst in self.instructions() {
+            match inst {
+                Err(_) => return false,
+                Ok(Instruction::Op(op)) => match op.to_u8() {
+                    // From Bitcoin Core:
+                    // if (opcode > OP_16 (0x60)) return false
+                    0..=0x60 => {}
+                    _ => return false,
+                },
+                Ok(Instruction::PushBytes(_)) => {}
+            }
+        }
+        true
+    }
+
     /// Checks whether a script pubkey is a P2PK output.
     ///
     /// You can obtain the public key, if its valid,
@@ -481,6 +500,27 @@ impl Script {
         }
     }
 
+    /// Iterates the script to find the last pushdata.
+    ///
+    /// Returns `None` if the instruction is an opcode or if the script is empty.
+    pub(crate) fn last_pushdata(&self) -> Option<Push> {
+        match self.instructions().last() {
+            Some(Ok(Instruction::PushBytes(bytes))) => Some(Push::Data(bytes)),
+            // OP_16 (0x60) and lower are considered "pushes" by Bitcoin Core.
+            Some(Ok(Instruction::Op(op))) if op.to_u8() <= 0x60 => {
+                // OP value as i8 - OP_RESERVED (0x50)
+                // OP_NEGATIVE1 (0x4f) is now -1 and the push nums are 1 to 16
+                let num = op.to_u8() as i8 - 0x50;
+                if num != 0 {
+                    Some(Push::Num(num))
+                } else {
+                    Some(Push::Reserved)
+                }
+            }
+            _ => None,
+        }
+    }
+
     /// Converts a [`Box<Script>`](Box) into a [`ScriptBuf`] without copying or allocating.
     #[must_use = "`self` will be dropped if the result is not used"]
     pub fn into_script_buf(self: Box<Self>) -> ScriptBuf {
@@ -494,6 +534,12 @@ impl Script {
     }
 }
 
+pub(crate) enum Push<'a> {
+    Data(&'a PushBytes),
+    Num(i8),
+    Reserved,
+}
+
 /// Iterator over bytes of a script
 pub struct Bytes<'a>(core::iter::Copied<core::slice::Iter<'a, u8>>);
 

bitcoin/src/blockdata/transaction.rs

@@ -21,15 +21,17 @@ use internals::write_err;
 use super::Weight;
 use crate::blockdata::locktime::absolute::{self, Height, Time};
 use crate::blockdata::locktime::relative;
-use crate::blockdata::script::ScriptBuf;
+use crate::blockdata::script::{self, Script, ScriptBuf};
 use crate::blockdata::witness::Witness;
 #[cfg(feature = "bitcoinconsensus")]
 pub use crate::consensus::validation::TxVerifyError;
 use crate::consensus::{encode, Decodable, Encodable};
+use crate::constants::WITNESS_SCALE_FACTOR;
 use crate::hash_types::{Txid, Wtxid};
 use crate::internal_macros::impl_consensus_encoding;
 use crate::parse::impl_parse_str_from_int_infallible;
 use crate::prelude::*;
+use crate::script::Push;
 #[cfg(doc)]
 use crate::sighash::{EcdsaSighashType, TapSighashType};
 use crate::string::FromHexStr;
@@ -833,6 +835,99 @@ impl Transaction {
     pub fn script_pubkey_lens(&self) -> impl Iterator<Item = usize> + '_ {
         self.output.iter().map(|txout| txout.script_pubkey.len())
     }
+
+    /// Get the sigop count for legacy transactions
+    fn count_sigops_legacy(&self) -> usize {
+        let mut n = 0;
+        for input in &self.input {
+            n += input.script_sig.count_sigops_legacy();
+        }
+        for output in &self.output {
+            n += output.script_pubkey.count_sigops_legacy();
+        }
+        n
+    }
+
+    fn count_p2sh_sigops<S>(&self, mut spent: Option<&mut S>) -> usize
+    where
+        S: FnMut(&OutPoint) -> Option<TxOut>,
+    {
+        if self.is_coinbase() {
+            return 0;
+        }
+        let mut n = 0;
+        for input in &self.input {
+            if let Some(prevout) = spent.as_mut().and_then(|s| s(&input.previous_output)) {
+                if prevout.script_pubkey.is_p2sh() {
+                    if let Some(Ok(script::Instruction::PushBytes(redeem))) =
+                        input.script_sig.instructions().last()
+                    {
+                        let script = Script::from_bytes(redeem.as_bytes());
+                        n += script.count_sigops();
+                    }
+                }
+            }
+        }
+        n
+    }
+
+    fn count_witness_sigops<S>(&self, mut spent: Option<&mut S>) -> usize
+    where
+        S: FnMut(&OutPoint) -> Option<TxOut>,
+    {
+        let mut n = 0;
+
+        #[inline]
+        fn count_with_prevout(prevout: TxOut, script_sig: &Script, witness: &Witness) -> usize {
+            let mut n = 0;
+
+            let script = if prevout.script_pubkey.is_witness_program() {
+                &prevout.script_pubkey
+            } else if prevout.script_pubkey.is_p2sh() && script_sig.is_push_only() {
+                if let Some(Push::Data(push_bytes)) = script_sig.last_pushdata() {
+                    Script::from_bytes(push_bytes.as_bytes())
+                } else {
+                    return 0;
+                }
+            } else {
+                return 0;
+            };
+
+            if script.is_v0_p2wsh() {
+                n += witness.sig_ops(script);
+            } else if script.is_v0_p2wpkh() {
+                n += 1;
+            }
+            n
+        }
+
+        for input in &self.input {
+            if let Some(Some(prevout)) = spent.as_mut().map(|s| s(&input.previous_output)) {
+                n += count_with_prevout(prevout, &input.script_sig, &input.witness);
+            }
+        }
+
+        n
+    }
+
+    /// Get the sigop cost for this transaction.
+    ///
+    /// The `spent` parameter is an optional closure/function that takes in an [`OutPoint`] and
+    /// returns a [`TxOut`]. Without access to the previous [`TxOut`], any sigops in redeemScripts,
+    /// witnessScripts, and P2WPKH sigops will not be counted.
+    pub fn get_sigop_cost<S, F>(&self, mut spent: Option<S>) -> usize
+    where
+        S: FnMut(&OutPoint) -> Option<TxOut>,
+    {
+        let mut n_sigop_cost = self.count_sigops_legacy() * WITNESS_SCALE_FACTOR;
+        if self.is_coinbase() {
+            return n_sigop_cost;
+        }
+        n_sigop_cost += self.count_p2sh_sigops(spent.as_mut()) * WITNESS_SCALE_FACTOR;
+        n_sigop_cost += self.count_witness_sigops(spent.as_mut());
+
+        n_sigop_cost
+    }
 }
 
 impl_consensus_encoding!(TxOut, value, script_pubkey);

bitcoin/src/blockdata/witness.rs

@@ -390,6 +390,26 @@ impl Witness {
             .and_then(|script_pos_from_last| self.nth(len - script_pos_from_last))
             .map(Script::from_bytes)
     }
+
+    /// Get sigops for the Witness
+    ///
+    /// witness_version is the raw opcode. OP_0 is 0, OP_1 is 81, etc.
+    pub(crate) fn sig_ops(&self, prevout_script: &Script) -> usize {
+        let bytes = prevout_script.as_bytes();
+        match (bytes.first(), bytes.len()) {
+            // Segwit v0, P2WPKH
+            (Some(0), 22) => 1,
+            // Segwit v0, P2WSH
+            (Some(0), 34) => {
+                if let Some(n) = self.last().map(Script::from_bytes).map(|s| s.count_sigops()) {
+                    n
+                } else {
+                    0
+                }
+            }
+            _ => 0,
+        }
+    }
 }
 
 impl Index<usize> for Witness {