blvm-consensus 0.1.82

//! Signature operation counting functions
//!
//! Implements consensus's sigop counting for block validation.
//! Sigops are counted to enforce MAX_BLOCK_SIGOPS_COST limit (80,000).
//!
//! Reference: consensus `tx_verify.cpp` and `script.cpp`

use crate::error::Result;
use crate::opcodes::*;
use crate::segwit::Witness;
use crate::types::*;
use crate::utxo_overlay::UtxoLookup;
#[cfg(feature = "spec-lock")]
use blvm_spec_lock::spec_locked;

/// Maximum number of public keys in a multisig (for sigop counting)
/// This is used when we can't accurately determine the number from the script
const MAX_PUBKEYS_PER_MULTISIG: u32 = 20;

/// Witness scale factor for sigop cost calculation
/// Legacy sigops count as 4x their actual number in sigop cost
const WITNESS_SCALE_FACTOR: u64 = 4;

/// Count sigops in a script (legacy counting)
///
/// Counts OP_CHECKSIG, OP_CHECKSIGVERIFY, OP_CHECKMULTISIG, OP_CHECKMULTISIGVERIFY.
/// Matches consensus's CScript::GetSigOpCount(bool fAccurate).
///
/// Uses GetOp-style iteration: each call to the loop body reads one opcode and
/// advances past any associated push data, exactly like consensus's GetOp().
///
/// # Arguments
/// * `script` - Script to count sigops in
/// * `accurate` - If true, use OP_1-OP_16 before OP_CHECKMULTISIG to determine key count
///
/// # Returns
/// Number of sigops in the script
#[cfg_attr(feature = "spec-lock", spec_locked("5.2.2", "CountSigOpsInScript"))]
pub fn count_sigops_in_script(script: &ByteString, accurate: bool) -> u32 {
    let mut count = 0u32;
    let mut last_opcode: Option<u8> = None;
    let mut i = 0;

    while i < script.len() {
        let opcode = script[i];

        // Skip past push data (matches consensus's GetOp)
        if opcode > 0 && opcode < OP_PUSHDATA1 {
            // Direct push: opcode IS the length (1-75 bytes)
            let len = opcode as usize;
            last_opcode = Some(opcode);
            i += 1 + len;
            continue;
        } else if opcode == OP_PUSHDATA1 {
            // OP_PUSHDATA1: next byte is length
            if i + 1 >= script.len() {
                break;
            }
            let len = script[i + 1] as usize;
            last_opcode = Some(opcode);
            i += 2 + len;
            continue;
        } else if opcode == OP_PUSHDATA2 {
            // OP_PUSHDATA2: next 2 bytes (little-endian) are length
            if i + 2 >= script.len() {
                break;
            }
            let len = u16::from_le_bytes([script[i + 1], script[i + 2]]) as usize;
            last_opcode = Some(opcode);
            i += 3 + len;
            continue;
        } else if opcode == OP_PUSHDATA4 {
            // OP_PUSHDATA4: next 4 bytes (little-endian) are length
            if i + 4 >= script.len() {
                break;
            }
            let len =
                u32::from_le_bytes([script[i + 1], script[i + 2], script[i + 3], script[i + 4]])
                    as usize;
            last_opcode = Some(opcode);
            i += 5 + len;
            continue;
        }

        // OP_CHECKSIG and OP_CHECKSIGVERIFY count as 1 sigop each
        if opcode == OP_CHECKSIG || opcode == OP_CHECKSIGVERIFY {
            count = count.saturating_add(1);
        }
        // OP_CHECKMULTISIG and OP_CHECKMULTISIGVERIFY count as multiple sigops
        else if opcode == OP_CHECKMULTISIG || opcode == OP_CHECKMULTISIGVERIFY {
            if accurate {
                // If accurate mode and previous opcode is OP_1-OP_16, use that number
                if let Some(prev_op) = last_opcode {
                    if (OP_1..=OP_16).contains(&prev_op) {
                        // OP_1 = 0x51, OP_16 = 0x60
                        // Decode: OP_N = N - 0x50
                        let n = (prev_op - OP_1 + 1) as u32;
                        count = count.saturating_add(n);
                    } else {
                        count = count.saturating_add(MAX_PUBKEYS_PER_MULTISIG);
                    }
                } else {
                    count = count.saturating_add(MAX_PUBKEYS_PER_MULTISIG);
                }
            } else {
                // Not accurate: assume maximum
                count = count.saturating_add(MAX_PUBKEYS_PER_MULTISIG);
            }
        }

        last_opcode = Some(opcode);
        i += 1;
    }

    count
}

/// Count sigops in a tapscript per BIP 342 (`CountTapscriptSigOps`).
///
/// Used for **per-tapscript validation weight** during script execution. This must **not** be folded
/// into block `MAX_BLOCK_SIGOPS_COST`: witness sigops only count witness **v0**; v1
/// (Taproot) contributes **0** to that total.
#[cfg_attr(feature = "spec-lock", spec_locked("11.2.8", "CountTapscriptSigOps"))]
pub fn count_tapscript_sigops(script: &ByteString) -> u32 {
    let mut count = 0u32;
    let mut i = 0;

    while i < script.len() {
        let opcode = script[i];

        if opcode > 0 && opcode < OP_PUSHDATA1 {
            let len = opcode as usize;
            i += 1 + len;
            continue;
        } else if opcode == OP_PUSHDATA1 {
            if i + 1 >= script.len() {
                break;
            }
            let len = script[i + 1] as usize;
            i += 2 + len;
            continue;
        } else if opcode == OP_PUSHDATA2 {
            if i + 2 >= script.len() {
                break;
            }
            let len = u16::from_le_bytes([script[i + 1], script[i + 2]]) as usize;
            i += 3 + len;
            continue;
        } else if opcode == OP_PUSHDATA4 {
            if i + 4 >= script.len() {
                break;
            }
            let len =
                u32::from_le_bytes([script[i + 1], script[i + 2], script[i + 3], script[i + 4]])
                    as usize;
            i += 5 + len;
            continue;
        }

        if opcode == OP_CHECKSIG || opcode == OP_CHECKSIGVERIFY || opcode == OP_CHECKSIGADD {
            count = count.saturating_add(1);
        }
        i += 1;
    }
    count
}

/// Orange Paper 5.2.1: P2SH scriptPubKey pattern (consensus `IsP2SH`-style predicate).
///
/// Structural invariant: if the function returns `true`, the script has exactly 23 bytes
/// (OP_HASH160 + 0x14 push + 20-byte hash + OP_EQUAL).  When length ≠ 23 the first AND
/// condition short-circuits to `false`, so `result == true` implies `script.len() == 23`.
#[cfg_attr(feature = "spec-lock", spec_locked("5.2.1", "IsP2SH"))]
#[cfg_attr(feature = "spec-lock", blvm_spec_lock::ensures(result == false || script.len() == 23))]
pub fn is_pay_to_script_hash(script: &[u8]) -> bool {
    script.len() == 23
        && script[0] == OP_HASH160  // OP_HASH160
        && script[1] == 0x14  // Push 20 bytes
        && script[22] == OP_EQUAL // OP_EQUAL
}

/// An opcode is a push while extracting the P2SH redeem for sigop counting.
///
/// The bound is `OP_16` (`0x60`). `OP_1NEGATE` (`0x4f`) and `OP_RESERVED` (`0x50`)
/// are pushes. An opcode above `OP_16` drops the input's P2SH and nested-witness
/// sigop contribution to zero.
#[cfg_attr(feature = "spec-lock", spec_locked("5.2.2", "F_SigopCountPushOpcode"))]
#[inline]
fn sigop_count_push_opcode(op: u8) -> bool {
    op <= OP_16
}

/// Last immediate of a push-only scriptSig, for P2SH and nested witness sigop counts.
///
/// A truncated push or any opcode above `OP_16` yields none. `OP_1NEGATE`,
/// `OP_RESERVED`, and `OP_1`..=`OP_16` are pushes whose immediate data is empty;
/// a later real push still supplies the redeem script.
fn extract_redeem_script_from_scriptsig(script_sig: &ByteString) -> Option<ByteString> {
    let mut i = 0;
    let mut last_data: Option<ByteString> = None;

    while i < script_sig.len() {
        let opcode = script_sig[i];
        if !sigop_count_push_opcode(opcode) {
            return None;
        }

        if opcode <= OP_PUSHDATA4 {
            let (len, advance) = if opcode < OP_PUSHDATA1 {
                (opcode as usize, 1)
            } else if opcode == OP_PUSHDATA1 {
                if i + 1 >= script_sig.len() {
                    return None;
                }
                (script_sig[i + 1] as usize, 2)
            } else if opcode == OP_PUSHDATA2 {
                if i + 2 >= script_sig.len() {
                    return None;
                }
                (
                    u16::from_le_bytes([script_sig[i + 1], script_sig[i + 2]]) as usize,
                    3,
                )
            } else {
                if i + 4 >= script_sig.len() {
                    return None;
                }
                (
                    u32::from_le_bytes([
                        script_sig[i + 1],
                        script_sig[i + 2],
                        script_sig[i + 3],
                        script_sig[i + 4],
                    ]) as usize,
                    5,
                )
            };

            if i + advance + len > script_sig.len() {
                return None;
            }

            last_data = Some(script_sig[i + advance..i + advance + len].to_vec());
            i += advance + len;
        } else {
            last_data = Some(Vec::new());
            i += 1;
        }
    }

    last_data
}

/// Get legacy sigop count from transaction
///
/// Counts sigops in scriptSig and scriptPubKey of all inputs and outputs.
/// Matches consensus's GetLegacySigOpCount().
///
/// # Arguments
/// * `tx` - Transaction to count sigops in
///
/// # Returns
/// Total number of legacy sigops
#[cfg_attr(feature = "spec-lock", spec_locked("5.2.2", "GetLegacySigOpCount"))]
pub fn get_legacy_sigop_count(tx: &Transaction) -> u32 {
    let mut count = 0u32;

    // Count sigops in all input scriptSigs
    for input in &tx.inputs {
        count = count.saturating_add(count_sigops_in_script(&input.script_sig, false));
    }

    // Count sigops in all output scriptPubKeys
    for output in &tx.outputs {
        count = count.saturating_add(count_sigops_in_script(&output.script_pubkey, false));
    }

    count
}

/// Get P2SH sigop count from transaction
///
/// Counts sigops in P2SH redeem scripts. Only counts sigops for outputs
/// that are P2SH (Pay-to-Script-Hash).
/// Matches consensus's GetP2SHSigOpCount().
///
/// # Arguments
/// * `tx` - Transaction to count sigops in
/// * `utxo_lookup` - UTXO lookup (UtxoSet or UtxoOverlay)
///
/// # Returns
/// Total number of P2SH sigops
#[cfg_attr(feature = "spec-lock", spec_locked("5.2.2", "GetP2SHSigOpCount"))]
pub fn get_p2sh_sigop_count<U: UtxoLookup>(tx: &Transaction, utxo_lookup: &U) -> Result<u32> {
    // Coinbase transactions have no P2SH sigops
    use crate::transaction::is_coinbase;
    if is_coinbase(tx) {
        return Ok(0);
    }

    let mut count = 0u32;

    for input in &tx.inputs {
        // Get the UTXO (scriptPubKey) for this input
        if let Some(utxo) = utxo_lookup.get(&input.prevout) {
            // Check if this is a P2SH output
            if is_pay_to_script_hash(utxo.script_pubkey.as_ref()) {
                // Extract redeem script from scriptSig
                if let Some(redeem_script) = extract_redeem_script_from_scriptsig(&input.script_sig)
                {
                    // Count sigops in redeem script (use accurate counting)
                    count = count.saturating_add(count_sigops_in_script(&redeem_script, true));
                }
            }
        }
    }

    Ok(count)
}

/// True for a version-0 witness program of 20 or 32 bytes.
fn is_version_zero_witness_program(script: &[u8]) -> bool {
    (script.len() == 22 && script[0] == OP_0 && script[1] == 0x14)
        || (script.len() == 34 && script[0] == OP_0 && script[1] == 0x20)
}

/// Sigops for one version-0 witness program.
///
/// A 20-byte program costs 1 even when the witness stack is empty. A 32-byte program
/// costs the accurate sigop count of its witness script, and costs 0 when that stack
/// is empty. Any other script costs 0. This cost is not scaled by four.
fn witness_program_sigops(script: &[u8], witness: Option<&Witness>) -> u64 {
    if script.len() == 22 && script[0] == OP_0 && script[1] == 0x14 {
        return 1;
    }
    if script.len() == 34 && script[0] == OP_0 && script[1] == 0x20 {
        let Some(witness) = witness.filter(|w| !w.is_empty()) else {
            return 0;
        };
        let Some(witness_script) = witness.last() else {
            return 0;
        };
        return count_sigops_in_script(witness_script, true) as u64;
    }
    0
}

/// Witness sigops for one input.
///
/// Native version-0 programs are counted from the scriptPubKey. When P2SH is enabled,
/// a redeem script that is itself a version-0 program is counted the same way. Version-1
/// programs, including those wrapped in P2SH, add nothing: their limit is the tapscript
/// validation weight, not the block sigop cost.
fn input_witness_sigops(
    script_pubkey: &[u8],
    script_sig: &ByteString,
    witness: Option<&Witness>,
    flags: u32,
) -> u64 {
    if flags & 0x800 == 0 {
        return 0;
    }
    if is_version_zero_witness_program(script_pubkey) {
        return witness_program_sigops(script_pubkey, witness);
    }
    if flags & 0x01 != 0 && is_pay_to_script_hash(script_pubkey) {
        if let Some(redeem) = extract_redeem_script_from_scriptsig(script_sig) {
            return witness_program_sigops(&redeem, witness);
        }
    }
    0
}

/// Count witness sigops in transaction
///
/// Native version-0 programs and P2SH-wrapped version-0 programs. A 20-byte program
/// costs 1. A 32-byte program costs the sigops in the witness script. Version 1 adds 0.
///
/// # Arguments
/// * `tx` - Transaction
/// * `witnesses` - Witness data for each input (slice of Witness vectors)
/// * `utxo_lookup` - UTXO lookup (UtxoSet or UtxoOverlay)
/// * `flags` - Script verification flags
///
/// # Returns
/// Number of witness sigops
#[cfg_attr(feature = "spec-lock", spec_locked("11.1", "CountWitnessSigOps"))]
pub(crate) fn count_witness_sigops<U: UtxoLookup>(
    tx: &Transaction,
    witnesses: &[Witness],
    utxo_lookup: &U,
    flags: u32,
) -> Result<u64> {
    use crate::transaction::is_coinbase;

    if is_coinbase(tx) {
        return Ok(0);
    }

    let mut count = 0u64;
    for (i, input) in tx.inputs.iter().enumerate() {
        if let Some(utxo) = utxo_lookup.get(&input.prevout) {
            count = count.saturating_add(input_witness_sigops(
                utxo.script_pubkey.as_ref(),
                &input.script_sig,
                witnesses.get(i),
                flags,
            ));
        }
    }

    Ok(count)
}

/// Legacy sigop count with accurate OP_CHECKMULTISIG (OP_1..OP_16 = 1..16, else 20).
/// Used for BIP54 per-tx 2500 limit (accurate sigop count with scriptSig).
#[cfg_attr(feature = "spec-lock", spec_locked("5.2.2", "GetLegacySigOpCount"))]
pub fn get_legacy_sigop_count_accurate(tx: &Transaction) -> u32 {
    let mut count = 0u32;
    for input in &tx.inputs {
        count = count.saturating_add(count_sigops_in_script(&input.script_sig, true));
    }
    for output in &tx.outputs {
        count = count.saturating_add(count_sigops_in_script(&output.script_pubkey, true));
    }
    count
}

/// Get total transaction sigop count (BIP54 limit).
///
/// Sum of legacy + P2SH + witness sigop counts (same accounting as BIP16).
/// Used to enforce per-transaction limit of 2500 sigops after BIP54 activation.
pub fn get_transaction_sigop_count<U: UtxoLookup>(
    tx: &Transaction,
    utxo_lookup: &U,
    witnesses: Option<&[Witness]>,
    flags: u32,
) -> Result<u64> {
    let legacy = get_legacy_sigop_count(tx) as u64;
    let p2sh = get_p2sh_sigop_count(tx, utxo_lookup)? as u64;
    let witness = witnesses
        .map(|w| count_witness_sigops(tx, w, utxo_lookup, flags))
        .unwrap_or(Ok(0))?;
    Ok(legacy.saturating_add(p2sh).saturating_add(witness))
}

/// BIP54 per-tx sigop count: legacy (accurate) + P2SH + witness.
/// BIP54 sigop check (accurate GetSigOpCount with scriptSig for legacy).
pub fn get_transaction_sigop_count_for_bip54<U: UtxoLookup>(
    tx: &Transaction,
    utxo_lookup: &U,
    witnesses: Option<&[Witness]>,
    flags: u32,
) -> Result<u64> {
    let legacy = get_legacy_sigop_count_accurate(tx) as u64;
    let p2sh = get_p2sh_sigop_count(tx, utxo_lookup)? as u64;
    let witness = witnesses
        .map(|w| count_witness_sigops(tx, w, utxo_lookup, flags))
        .unwrap_or(Ok(0))?;
    Ok(legacy.saturating_add(p2sh).saturating_add(witness))
}

/// Get total transaction sigop cost
///
/// Calculates total sigop cost for a transaction, including:
/// - Legacy sigops × 4 (witness scale factor)
/// - P2SH sigops × 4 (if P2SH enabled)
/// - Witness sigops (actual count, not scaled)
///
/// Matches consensus's GetTransactionSigOpCost().
///
/// # Arguments
/// * `tx` - Transaction to count sigops in
/// * `utxo_set` - UTXO set to lookup inputs
/// * `witness` - Witness data for this transaction (one Witness per input)
/// * `flags` - Script verification flags
///
/// # Returns
/// Total sigop cost
#[cfg_attr(feature = "spec-lock", spec_locked("5.2.2", "GetTransactionSigOpCost"))]
pub fn get_transaction_sigop_cost<U: UtxoLookup>(
    tx: &Transaction,
    utxo_lookup: &U,
    witness: Option<&Witness>,
    flags: u32,
) -> Result<u64> {
    let witness_slices = witness.map(std::slice::from_ref);
    get_transaction_sigop_cost_with_witness_slices(tx, utxo_lookup, witness_slices, flags)
}

/// Same as get_transaction_sigop_cost but accepts pre-fetched UTXOs in input order.
/// Avoids redundant overlay lookups when caller already has UTXO data.
#[cfg_attr(feature = "spec-lock", spec_locked("5.2.2", "GetTransactionSigOpCost"))]
pub fn get_transaction_sigop_cost_with_utxos(
    tx: &Transaction,
    utxos: &[Option<&UTXO>],
    witnesses: Option<&[Witness]>,
    flags: u32,
) -> Result<u64> {
    let legacy_count = get_legacy_sigop_count(tx) as u64;
    let mut total_cost = legacy_count.saturating_mul(WITNESS_SCALE_FACTOR);

    use crate::transaction::is_coinbase;
    if is_coinbase(tx) {
        return Ok(total_cost);
    }

    if (flags & 0x01) != 0 {
        let mut p2sh_count = 0u32;
        for (input, utxo_opt) in tx.inputs.iter().zip(utxos.iter()) {
            if let Some(utxo) = utxo_opt {
                if is_pay_to_script_hash(utxo.script_pubkey.as_ref()) {
                    if let Some(redeem_script) =
                        extract_redeem_script_from_scriptsig(&input.script_sig)
                    {
                        p2sh_count =
                            p2sh_count.saturating_add(count_sigops_in_script(&redeem_script, true));
                    }
                }
            }
        }
        total_cost = total_cost
            .saturating_add(p2sh_count.saturating_mul(WITNESS_SCALE_FACTOR as u32) as u64);
    }

    if let Some(witnesses) = witnesses {
        for (i, (input, utxo_opt)) in tx.inputs.iter().zip(utxos.iter()).enumerate() {
            if let Some(utxo) = utxo_opt {
                total_cost = total_cost.saturating_add(input_witness_sigops(
                    utxo.script_pubkey.as_ref(),
                    &input.script_sig,
                    witnesses.get(i),
                    flags,
                ));
            }
        }
    }

    Ok(total_cost)
}

/// Same as get_transaction_sigop_cost but accepts per-input witness slices directly.
/// Avoids flattening witness data in block validation hot path.
#[cfg_attr(feature = "spec-lock", spec_locked("5.2.2", "GetTransactionSigOpCost"))]
pub fn get_transaction_sigop_cost_with_witness_slices<U: UtxoLookup>(
    tx: &Transaction,
    utxo_lookup: &U,
    witnesses: Option<&[Witness]>,
    flags: u32,
) -> Result<u64> {
    // Legacy sigops × witness scale factor
    let legacy_count = get_legacy_sigop_count(tx) as u64;
    let mut total_cost = legacy_count.saturating_mul(WITNESS_SCALE_FACTOR);

    use crate::transaction::is_coinbase;
    if is_coinbase(tx) {
        return Ok(total_cost);
    }

    // P2SH sigops × witness scale factor (if P2SH enabled)
    if (flags & 0x01) != 0 {
        // SCRIPT_VERIFY_P2SH flag enabled
        let p2sh_count = get_p2sh_sigop_count(tx, utxo_lookup)? as u64;
        total_cost = total_cost.saturating_add(p2sh_count.saturating_mul(WITNESS_SCALE_FACTOR));
    }

    // Witness sigops (actual count, not scaled)
    if let Some(witnesses) = witnesses {
        let witness_count = count_witness_sigops(tx, witnesses, utxo_lookup, flags)?;
        total_cost = total_cost.saturating_add(witness_count);
    }

    Ok(total_cost)
}

#[cfg(test)]
#[path = "sigop_tests.rs"]
mod tests;