Merge rust-bitcoin/rust-bitcoin#3164: Extension traits for `ScriptBuf`

2bb90b8203 Introduce two extensions traits for ScriptBuf (Tobin C. Harding)
ae0a5bd64a Run cargo fmt (Tobin C. Harding)
3fdc574851 Add temporary script buf modules (Tobin C. Harding)
4ff5d6886b Add private ScriptBufAsVec type (Tobin C. Harding)
c81fb93359 Make push_slice_no_opt pub(crate) (Tobin C. Harding)
1001a33f19 Add second ScriptBuf impl block (Tobin C. Harding)
3625d74e8b Make pub in crate functions pub crate (Tobin C. Harding)
b368384317 Separate ScriptBuf POD methods (Tobin C. Harding)

Pull request description:

  Similar to #3155 but for `ScriptBuf`, however it is a little more involved.

  Note:
  - the change to use `impl` syntax (and addition of #3179)
  - mad trickery of `ScriptBufAsVec` (props to Kix)
  - widening of scope of private functions

  Onward and upward!

ACKs for top commit:
  Kixunil:
    ACK 2bb90b8203
  apoelstra:
    ACK 2bb90b8203 successfully ran local tests

Tree-SHA512: 7209d8dc436e52b23e1dbfd9db8432df225ebdb701f465e4d1b55328e22988c98a0f28efdf2a8b3edbafc754354d718ab36bd2f5b1621d12e061b2dadaf49a05

bitcoin/examples/ecdsa-psbt.rs

@@ -37,6 +37,7 @@ use bitcoin::bip32::{ChildNumber, DerivationPath, Fingerprint, IntoDerivationPat
 use bitcoin::consensus::encode;
 use bitcoin::locktime::absolute;
 use bitcoin::psbt::{self, Input, Psbt, PsbtSighashType};
+use bitcoin::script::ScriptBufExt as _;
 use bitcoin::secp256k1::{Secp256k1, Signing, Verification};
 use bitcoin::{
     transaction, Address, Amount, CompressedPublicKey, Network, OutPoint, ScriptBuf, Sequence,

bitcoin/examples/taproot-psbt.rs

@@ -84,7 +84,7 @@ use bitcoin::consensus::encode;
 use bitcoin::key::{TapTweak, XOnlyPublicKey};
 use bitcoin::opcodes::all::{OP_CHECKSIG, OP_CLTV, OP_DROP};
 use bitcoin::psbt::{self, Input, Output, Psbt, PsbtSighashType};
-use bitcoin::script::ScriptExt as _;
+use bitcoin::script::{ScriptBufExt as _, ScriptExt as _};
 use bitcoin::secp256k1::Secp256k1;
 use bitcoin::sighash::{self, SighashCache, TapSighash, TapSighashType};
 use bitcoin::taproot::{self, LeafVersion, TapLeafHash, TaprootBuilder, TaprootSpendInfo};

bitcoin/src/address/mod.rs

@@ -894,6 +894,7 @@ mod tests {
     use super::*;
     use crate::network::params;
     use crate::network::Network::{Bitcoin, Testnet};
+    use crate::script::ScriptBufExt as _;
 
     fn roundtrips(addr: &Address, network: Network) {
         assert_eq!(

bitcoin/src/blockdata/script/builder.rs

@@ -7,7 +7,7 @@ use crate::locktime::absolute;
 use crate::opcodes::all::*;
 use crate::opcodes::{self, Opcode};
 use crate::prelude::Vec;
-use crate::script::{ScriptExt as _, ScriptExtPriv as _};
+use crate::script::{ScriptBufExt as _, ScriptExt as _, ScriptExtPriv as _};
 use crate::Sequence;
 
 /// An Object which can be used to construct a script piece by piece.

bitcoin/src/blockdata/script/instruction.rs

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: CC0-1.0
 
-use super::{read_uint_iter, Error, PushBytes, Script, ScriptBuf, UintError};
+use super::{
+    read_uint_iter, Error, PushBytes, Script, ScriptBuf, ScriptBufExtPriv as _, UintError,
+};
 use crate::opcodes::{self, Opcode};
 
 /// A "parsed opcode" which allows iterating over a [`Script`] in a more sensible way.

bitcoin/src/blockdata/script/owned.rs

@@ -64,20 +64,6 @@ impl ScriptBuf {
     /// Returns a mutable reference to unsized script.
     pub fn as_mut_script(&mut self) -> &mut Script { Script::from_bytes_mut(&mut self.0) }
 
-    /// Creates a new script builder
-    pub fn builder() -> Builder { Builder::new() }
-
-    /// Generates OP_RETURN-type of scriptPubkey for the given data.
-    pub fn new_op_return<T: AsRef<PushBytes>>(data: T) -> Self {
-        Builder::new().push_opcode(OP_RETURN).push_slice(data).into_script()
-    }
-
-    /// Creates a [`ScriptBuf`] from a hex string.
-    pub fn from_hex(s: &str) -> Result<Self, hex::HexToBytesError> {
-        let v = Vec::from_hex(s)?;
-        Ok(ScriptBuf::from_bytes(v))
-    }
-
     /// Converts byte vector into script.
     ///
     /// This method doesn't (re)allocate.
@@ -88,63 +74,55 @@ impl ScriptBuf {
     /// This method doesn't (re)allocate.
     pub fn into_bytes(self) -> Vec<u8> { self.0 }
 
+    /// Converts this `ScriptBuf` into a [boxed](Box) [`Script`].
+    ///
+    /// This method reallocates if the capacity is greater than length of the script but should not
+    /// when they are equal. If you know beforehand that you need to create a script of exact size
+    /// use [`reserve_exact`](Self::reserve_exact) before adding data to the script so that the
+    /// reallocation can be avoided.
+    #[must_use = "`self` will be dropped if the result is not used"]
+    #[inline]
+    pub fn into_boxed_script(self) -> Box<Script> {
+        // Copied from PathBuf::into_boxed_path
+        let rw = Box::into_raw(self.0.into_boxed_slice()) as *mut Script;
+        unsafe { Box::from_raw(rw) }
+    }
+}
+
+crate::internal_macros::define_extension_trait! {
+    /// Extension functionality for the [`ScriptBuf`] type.
+    pub trait ScriptBufExt impl for ScriptBuf {
+        /// Creates a new script builder
+        fn builder() -> Builder { Builder::new() }
+
+        /// Generates OP_RETURN-type of scriptPubkey for the given data.
+        fn new_op_return<T: AsRef<PushBytes>>(data: T) -> Self {
+            Builder::new().push_opcode(OP_RETURN).push_slice(data).into_script()
+        }
+
+        /// Creates a [`ScriptBuf`] from a hex string.
+        fn from_hex(s: &str) -> Result<ScriptBuf, hex::HexToBytesError> {
+            let v = Vec::from_hex(s)?;
+            Ok(ScriptBuf::from_bytes(v))
+        }
+
         /// Adds a single opcode to the script.
-    pub fn push_opcode(&mut self, data: Opcode) { self.0.push(data.to_u8()); }
+        fn push_opcode(&mut self, data: Opcode) { self.as_byte_vec().push(data.to_u8()); }
 
         /// Adds instructions to push some arbitrary data onto the stack.
-    pub fn push_slice<T: AsRef<PushBytes>>(&mut self, data: T) {
+        fn push_slice<T: AsRef<PushBytes>>(&mut self, data: T) {
             let data = data.as_ref();
-        self.reserve(Self::reserved_len_for_slice(data.len()));
+            self.reserve(ScriptBuf::reserved_len_for_slice(data.len()));
             self.push_slice_no_opt(data);
         }
 
-    /// Pushes the slice without reserving
-    fn push_slice_no_opt(&mut self, data: &PushBytes) {
-        // Start with a PUSH opcode
-        match data.len().to_u64() {
-            n if n < opcodes::Ordinary::OP_PUSHDATA1 as u64 => {
-                self.0.push(n as u8);
-            }
-            n if n < 0x100 => {
-                self.0.push(opcodes::Ordinary::OP_PUSHDATA1.to_u8());
-                self.0.push(n as u8);
-            }
-            n if n < 0x10000 => {
-                self.0.push(opcodes::Ordinary::OP_PUSHDATA2.to_u8());
-                self.0.push((n % 0x100) as u8);
-                self.0.push((n / 0x100) as u8);
-            }
-            // `PushBytes` enforces len < 0x100000000
-            n => {
-                self.0.push(opcodes::Ordinary::OP_PUSHDATA4.to_u8());
-                self.0.push((n % 0x100) as u8);
-                self.0.push(((n / 0x100) % 0x100) as u8);
-                self.0.push(((n / 0x10000) % 0x100) as u8);
-                self.0.push((n / 0x1000000) as u8);
-            }
-        }
-        // Then push the raw bytes
-        self.0.extend_from_slice(data.as_bytes());
-    }
-
-    /// Computes the sum of `len` and the length of an appropriate push opcode.
-    pub(in crate::blockdata::script) fn reserved_len_for_slice(len: usize) -> usize {
-        len + match len {
-            0..=0x4b => 1,
-            0x4c..=0xff => 2,
-            0x100..=0xffff => 3,
-            // we don't care about oversized, the other fn will panic anyway
-            _ => 5,
-        }
-    }
-
         /// Add a single instruction to the script.
         ///
         /// # Panics
         ///
         /// The method panics if the instruction is a data push with length greater or equal to
         /// 0x100000000.
-    pub fn push_instruction(&mut self, instruction: Instruction<'_>) {
+        fn push_instruction(&mut self, instruction: Instruction<'_>) {
             match instruction {
                 Instruction::Op(opcode) => self.push_opcode(opcode),
                 Instruction::PushBytes(bytes) => self.push_slice(bytes),
@@ -152,7 +130,7 @@ impl ScriptBuf {
         }
 
         /// Like push_instruction, but avoids calling `reserve` to not re-check the length.
-    pub fn push_instruction_no_opt(&mut self, instruction: Instruction<'_>) {
+        fn push_instruction_no_opt(&mut self, instruction: Instruction<'_>) {
             match instruction {
                 Instruction::Op(opcode) => self.push_opcode(opcode),
                 Instruction::PushBytes(bytes) => self.push_slice_no_opt(bytes),
@@ -173,34 +151,74 @@ impl ScriptBuf {
         /// This function needs to iterate over the script to find the last instruction. Prefer
         /// `Builder` if you're creating the script from scratch or if you want to push `OP_VERIFY`
         /// multiple times.
-    pub fn scan_and_push_verify(&mut self) { self.push_verify(self.last_opcode()); }
+        fn scan_and_push_verify(&mut self) { self.push_verify(self.last_opcode()); }
+    }
+}
+
+crate::internal_macros::define_extension_trait! {
+    pub(crate) trait ScriptBufExtPriv impl for ScriptBuf {
+        /// Pretends to convert `&mut ScriptBuf` to `&mut Vec<u8>` so that it can be modified.
+        ///
+        /// Note: if the returned value leaks the original `ScriptBuf` will become empty.
+        fn as_byte_vec(&mut self) -> ScriptBufAsVec<'_> {
+            let vec = core::mem::take(self).into_bytes();
+            ScriptBufAsVec(self, vec)
+        }
+
+        /// Pushes the slice without reserving
+        fn push_slice_no_opt(&mut self, data: &PushBytes) {
+            let mut this = self.as_byte_vec();
+            // Start with a PUSH opcode
+            match data.len().to_u64() {
+                n if n < opcodes::Ordinary::OP_PUSHDATA1 as u64 => {
+                    this.push(n as u8);
+                }
+                n if n < 0x100 => {
+                    this.push(opcodes::Ordinary::OP_PUSHDATA1.to_u8());
+                    this.push(n as u8);
+                }
+                n if n < 0x10000 => {
+                    this.push(opcodes::Ordinary::OP_PUSHDATA2.to_u8());
+                    this.push((n % 0x100) as u8);
+                    this.push((n / 0x100) as u8);
+                }
+                // `PushBytes` enforces len < 0x100000000
+                n => {
+                    this.push(opcodes::Ordinary::OP_PUSHDATA4.to_u8());
+                    this.push((n % 0x100) as u8);
+                    this.push(((n / 0x100) % 0x100) as u8);
+                    this.push(((n / 0x10000) % 0x100) as u8);
+                    this.push((n / 0x1000000) as u8);
+                }
+            }
+            // Then push the raw bytes
+            this.extend_from_slice(data.as_bytes());
+        }
+
+        /// Computes the sum of `len` and the length of an appropriate push opcode.
+        fn reserved_len_for_slice(len: usize) -> usize {
+            len + match len {
+                0..=0x4b => 1,
+                0x4c..=0xff => 2,
+                0x100..=0xffff => 3,
+                // we don't care about oversized, the other fn will panic anyway
+                _ => 5,
+            }
+        }
 
         /// Adds an `OP_VERIFY` to the script or changes the most-recently-added opcode to `VERIFY`
         /// alternative.
         ///
         /// See the public fn [`Self::scan_and_push_verify`] to learn more.
-    pub(in crate::blockdata::script) fn push_verify(&mut self, last_opcode: Option<Opcode>) {
+        fn push_verify(&mut self, last_opcode: Option<Opcode>) {
             match opcode_to_verify(last_opcode) {
                 Some(opcode) => {
-                self.0.pop();
+                    self.as_byte_vec().pop();
                     self.push_opcode(opcode);
                 }
                 None => self.push_opcode(OP_VERIFY),
             }
         }
-
-    /// Converts this `ScriptBuf` into a [boxed](Box) [`Script`].
-    ///
-    /// This method reallocates if the capacity is greater than length of the script but should not
-    /// when they are equal. If you know beforehand that you need to create a script of exact size
-    /// use [`reserve_exact`](Self::reserve_exact) before adding data to the script so that the
-    /// reallocation can be avoided.
-    #[must_use = "`self` will be dropped if the result is not used"]
-    #[inline]
-    pub fn into_boxed_script(self) -> Box<Script> {
-        // Copied from PathBuf::into_boxed_path
-        let rw = Box::into_raw(self.0.into_boxed_slice()) as *mut Script;
-        unsafe { Box::from_raw(rw) }
     }
 }
 
@@ -250,3 +268,27 @@ impl<'a> Extend<Instruction<'a>> for ScriptBuf {
         }
     }
 }
+
+/// Pretends that this is a mutable reference to [`ScriptBuf`]'s internal buffer.
+///
+/// In reality the backing `Vec<u8>` is swapped with an empty one and this is holding both the
+/// reference and the vec. The vec is put back when this drops so it also covers paics. (But not
+/// leaks, which is OK since we never leak.)
+pub(crate) struct ScriptBufAsVec<'a>(&'a mut ScriptBuf, Vec<u8>);
+
+impl<'a> core::ops::Deref for ScriptBufAsVec<'a> {
+    type Target = Vec<u8>;
+
+    fn deref(&self) -> &Self::Target { &self.1 }
+}
+
+impl<'a> core::ops::DerefMut for ScriptBufAsVec<'a> {
+    fn deref_mut(&mut self) -> &mut Self::Target { &mut self.1 }
+}
+
+impl<'a> Drop for ScriptBufAsVec<'a> {
+    fn drop(&mut self) {
+        let vec = core::mem::take(&mut self.1);
+        *(self.0) = ScriptBuf::from_bytes(vec);
+    }
+}

bitcoin/src/crypto/ecdsa.rs

@@ -13,6 +13,8 @@ use io::Write;
 
 use crate::prelude::{DisplayHex, Vec};
 use crate::script::PushBytes;
+#[cfg(doc)]
+use crate::script::ScriptBufExt as _;
 use crate::sighash::{EcdsaSighashType, NonStandardSighashTypeError};
 
 const MAX_SIG_LEN: usize = 73;

bitcoin/src/crypto/sighash.rs

@@ -1470,6 +1470,7 @@ mod tests {
     use super::*;
     use crate::consensus::deserialize;
     use crate::locktime::absolute;
+    use crate::script::ScriptBufExt as _;
 
     extern crate serde_json;
 

bitcoin/src/internal_macros.rs

@@ -251,7 +251,7 @@ macro_rules! define_extension_trait {
     ($(#[$($trait_attrs:tt)*])* $trait_vis:vis trait $trait_name:ident impl for $ty:ident {
         $(
             $(#[$($fn_attrs:tt)*])*
-            fn $fn:ident$(<$($gen:ident: $gent:ident),*>)?($($params:tt)*) $( -> $ret:ty )? $body:block
+            fn $fn:ident$(<$($gen:ident: $gent:path),*>)?($($params:tt)*) $( -> $ret:ty )? $body:block
         )*
     }) => {
         $(#[$($trait_attrs)*])* $trait_vis trait $trait_name {

bitcoin/src/psbt/mod.rs

@@ -1216,7 +1216,7 @@ mod tests {
     use crate::locktime::absolute;
     use crate::network::NetworkKind;
     use crate::psbt::serialize::{Deserialize, Serialize};
-    use crate::script::ScriptBuf;
+    use crate::script::{ScriptBuf, ScriptBufExt as _};
     use crate::transaction::{self, OutPoint, TxIn};
     use crate::witness::Witness;
     use crate::Sequence;

bitcoin/src/psbt/serialize.rs

@@ -383,6 +383,7 @@ fn key_source_len(key_source: &KeySource) -> usize { 4 + 4 * (key_source.1).as_r
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::script::ScriptBufExt as _;
 
     // Composes tree matching a given depth map, filled with dumb script leafs,
     // each of which consists of a single push-int op code, with int value

bitcoin/src/taproot/mod.rs

@@ -1544,6 +1544,7 @@ mod test {
     use secp256k1::VerifyOnly;
 
     use super::*;
+    use crate::script::ScriptBufExt as _;
     use crate::sighash::{TapSighash, TapSighashTag};
     use crate::{Address, KnownHrp};
     extern crate serde_json;

bitcoin/tests/bip_174.rs

@@ -9,7 +9,7 @@ use bitcoin::consensus::encode::{deserialize, serialize_hex};
 use bitcoin::hex::FromHex;
 use bitcoin::opcodes::OP_0;
 use bitcoin::psbt::{Psbt, PsbtSighashType};
-use bitcoin::script::PushBytes;
+use bitcoin::script::{PushBytes, ScriptBufExt as _};
 use bitcoin::secp256k1::Secp256k1;
 use bitcoin::{
     absolute, script, transaction, Amount, Denomination, NetworkKind, OutPoint, PrivateKey,

bitcoin/tests/serde.rs

@@ -33,6 +33,7 @@ use bitcoin::hex::FromHex;
 use bitcoin::locktime::{absolute, relative};
 use bitcoin::psbt::raw::{self, Key, Pair, ProprietaryKey};
 use bitcoin::psbt::{Input, Output, Psbt, PsbtSighashType};
+use bitcoin::script::ScriptBufExt as _;
 use bitcoin::sighash::{EcdsaSighashType, TapSighashType};
 use bitcoin::taproot::{self, ControlBlock, LeafVersion, TapTree, TaprootBuilder};
 use bitcoin::witness::Witness;