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,19 @@ 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(_)) => return false,
+                Ok(Instruction::PushBytes(_)) => {}
+            }
+        }
+        true
+    }
+
     /// Checks whether a script pubkey is a P2PK output.
     ///
     /// You can obtain the public key, if its valid,
@@ -489,6 +503,16 @@ impl Script {
         }
     }
 
+    /// Iterates the script to find the last pushdata.
+    ///
+    /// Returns `None` is the instruction is an opcode or if the script is empty.
+    pub(crate) fn last_pushdata(&self) -> Option<&PushBytes> {
+        match self.instructions().last() {
+            Some(Ok(Instruction::PushBytes(bytes))) => Some(bytes),
+            _ => 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 {

bitcoin/src/blockdata/transaction.rs

@@ -21,11 +21,12 @@ 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;
@@ -833,6 +834,98 @@ 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) -> Result<usize, script::Error> {
+        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()?;
+        }
+        Ok(n)
+    }
+
+    fn count_p2sh_sigops<S>(&self, mut spent: Option<&mut S>) -> Result<usize, script::Error>
+    where
+        S: FnMut(&OutPoint) -> Option<TxOut>,
+    {
+        if self.is_coinbase() {
+            return Ok(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()?;
+                    }
+                }
+            }
+        }
+        Ok(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()
+                && script_sig.len() > 0
+            {
+                Script::from_bytes(script_sig.last_pushdata().unwrap().as_bytes())
+            } 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>) -> Result<usize, script::Error>
+    where
+        S: FnMut(&OutPoint) -> Option<TxOut>,
+    {
+        let mut n_sigop_cost = self.count_sigops_legacy()? * WITNESS_SCALE_FACTOR;
+        if self.is_coinbase() {
+            return Ok(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());
+
+        Ok(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.get(0), bytes.len()) {
+            // Segwit v0, P2WPKH
+            (Some(0), 22) => 1,
+            // Segwit v0, P2WSH
+            (Some(0), 34) => {
+                if let Some(Ok(n)) = self.last().map(Script::from_bytes).map(|s| s.count_sigops()) {
+                    n
+                } else {
+                    0
+                }
+            }
+            _ => 0,
+        }
+    }
 }
 
 impl Index<usize> for Witness {