Merge pull request #15 from tcharding/01-11-combine

Improve code for the Combiner role

src/macros.rs

@@ -1,18 +1,24 @@
 // SPDX-License-Identifier: CC0-1.0
 
-/// Combines the given `$thing` field of two types by setting `self.thing` to be
-/// the value in `other.thing` iff `self.thing` is currently empty.
+/// Combines two `Option<Foo>` fields.
 ///
+/// Sets `self.thing` to be `Some(other.thing)` iff `self.thing` is `None`.
 /// If `self.thing` already contains a value then this macro does nothing.
-#[allow(unused_macros)]
-macro_rules! combine {
+macro_rules! combine_option {
     ($thing:ident, $slf:ident, $other:ident) => {
         if let (&None, Some($thing)) = (&$slf.$thing, $other.$thing) {
             $slf.$thing = Some($thing);
         }
     };
 }
 
+/// Combines to `BTreeMap` fields by extending the map in `self.thing`.
+macro_rules! combine_map {
+    ($thing:ident, $slf:ident, $other:ident) => {
+        $slf.$thing.extend($other.$thing)
+    };
+}
+
 // Implements our Serialize/Deserialize traits using bitcoin consensus serialization.
 macro_rules! impl_psbt_de_serialize {
     ($thing:ty) => {

src/v0/error.rs

@@ -6,8 +6,7 @@ use core::fmt;
 
 use bitcoin::{sighash, FeeRate, Transaction};
 
-use crate::error::{write_err, InconsistentKeySourcesError};
-use crate::prelude::Box;
+use crate::error::write_err;
 use crate::v0::map::{global, input, output};
 use crate::v0::Psbt;
 
@@ -303,50 +302,6 @@ impl From<IndexOutOfBoundsError> for SignError {
     fn from(e: IndexOutOfBoundsError) -> Self { SignError::IndexOutOfBounds(e) }
 }
 
-/// Error combining two PSBTs.
-#[derive(Debug, Clone, PartialEq, Eq)]
-#[non_exhaustive]
-pub enum CombineError {
-    /// Attempting to combine with a PSBT describing a different unsigned transaction.
-    UnexpectedUnsignedTx {
-        /// Expected transaction.
-        expected: Box<Transaction>,
-        /// Actual transaction.
-        actual: Box<Transaction>,
-    },
-    /// Global extended public key has inconsistent key sources.
-    InconsistentKeySources(InconsistentKeySourcesError),
-}
-
-impl fmt::Display for CombineError {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        use CombineError::*;
-
-        match *self {
-            UnexpectedUnsignedTx { ref expected, ref actual } =>
-                write!(f, "combine, transaction differs from actual {:?} {:?}", expected, actual),
-            InconsistentKeySources(ref e) =>
-                write_err!(f, "combine with inconsistent key sources"; e),
-        }
-    }
-}
-
-#[cfg(feature = "std")]
-impl std::error::Error for CombineError {
-    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
-        use CombineError::*;
-
-        match *self {
-            UnexpectedUnsignedTx { .. } => None,
-            InconsistentKeySources(ref e) => Some(e),
-        }
-    }
-}
-
-impl From<InconsistentKeySourcesError> for CombineError {
-    fn from(e: InconsistentKeySourcesError) -> Self { Self::InconsistentKeySources(e) }
-}
-
 /// Unsigned transaction checks error.
 #[derive(Debug, Clone, PartialEq, Eq)]
 #[non_exhaustive]

src/v0/map/global.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: CC0-1.0
 
 use core::convert::TryFrom;
-use core::{cmp, fmt};
+use core::fmt;
 
 use bitcoin::bip32::{self, ChildNumber, DerivationPath, Fingerprint, KeySource, Xpub};
 // TODO: This should be exposed like this in rust-bitcoin.
@@ -18,7 +18,7 @@ use crate::consts::{
 use crate::error::{write_err, InconsistentKeySourcesError};
 use crate::io::{self, Cursor, Read};
 use crate::prelude::*;
-use crate::v0::error::{CombineError, UnsignedTxChecksError};
+use crate::v0::error::UnsignedTxChecksError;
 use crate::v0::map::Map;
 use crate::version::Version;
 use crate::{consts, raw, serialize, V0};
@@ -234,18 +234,20 @@ impl Global {
     /// In accordance with BIP 174 this function is commutative i.e., `A.combine(B) == B.combine(A)`
     pub fn combine(&mut self, other: Self) -> Result<(), CombineError> {
         if self.unsigned_tx != other.unsigned_tx {
-            return Err(CombineError::UnexpectedUnsignedTx {
-                expected: Box::new(self.unsigned_tx.clone()),
-                actual: Box::new(other.unsigned_tx),
+            return Err(CombineError::UnsignedTxMismatch {
+                this: Box::new(self.unsigned_tx.clone()),
+                that: Box::new(other.unsigned_tx),
             });
         }
 
+        // Combining different versions of PSBT without explicit conversion is out of scope.
+        if self.version != other.version {
+            return Err(CombineError::VersionMismatch { this: self.version, that: other.version });
+        }
+
         // BIP 174: The Combiner must remove any duplicate key-value pairs, in accordance with
         //          the specification. It can pick arbitrarily when conflicts occur.
 
-        // Keeping the highest version
-        self.version = cmp::max(self.version, other.version);
-
         // Merging xpubs
         for (xpub, (fingerprint1, derivation1)) in other.xpubs {
             match self.xpubs.entry(xpub) {
@@ -281,8 +283,9 @@ impl Global {
             }
         }
 
-        self.proprietaries.extend(other.proprietaries);
-        self.unknowns.extend(other.unknowns);
+        combine_map!(proprietaries, self, other);
+        combine_map!(unknowns, self, other);
+
         Ok(())
     }
 }
@@ -511,3 +514,59 @@ impl From<consensus::Error> for InsertPairError {
 impl From<bip32::Error> for InsertPairError {
     fn from(e: bip32::Error) -> Self { Self::Bip32(e) }
 }
+
+/// Error combining two global maps.
+#[derive(Debug, Clone, PartialEq, Eq)]
+#[non_exhaustive]
+pub enum CombineError {
+    /// The unsigned transactions are not the same.
+    UnsignedTxMismatch {
+        /// Attempted to combine a PBST with `this` transaction.
+        this: Box<Transaction>,
+        /// Into a PBST with `that` transaction.
+        that: Box<Transaction>,
+    },
+    ///
+    VersionMismatch {
+        /// Attempted to combine a PBST with `this` version.
+        this: Version,
+        /// Into a PBST with `that` version.
+        that: Version,
+    },
+    /// Xpubs have inconsistent key sources.
+    InconsistentKeySources(InconsistentKeySourcesError),
+}
+
+impl fmt::Display for CombineError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use CombineError::*;
+
+        match *self {
+            UnsignedTxMismatch { ref this, ref that } => write!(
+                f,
+                "combine two PSBTs with different unsigned transactions: {:?} {:?}",
+                this, that
+            ),
+            VersionMismatch { ref this, ref that } =>
+                write!(f, "combine two PSBTs with different versions: {:?} {:?}", this, that),
+            InconsistentKeySources(ref e) =>
+                write_err!(f, "combine with inconsistent key sources"; e),
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for CombineError {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        use CombineError::*;
+
+        match *self {
+            InconsistentKeySources(ref e) => Some(e),
+            UnsignedTxMismatch { .. } | VersionMismatch { .. } => None,
+        }
+    }
+}
+
+impl From<InconsistentKeySourcesError> for CombineError {
+    fn from(e: InconsistentKeySourcesError) -> Self { Self::InconsistentKeySources(e) }
+}

src/v0/map/input.rs

@@ -278,32 +278,32 @@ impl Input {
 
     /// Combines this [`Input`] with `other` `Input` (as described by BIP 174).
     pub fn combine(&mut self, other: Self) {
-        combine!(non_witness_utxo, self, other);
+        combine_option!(non_witness_utxo, self, other);
 
         if let (&None, Some(witness_utxo)) = (&self.witness_utxo, other.witness_utxo) {
             self.witness_utxo = Some(witness_utxo);
             self.non_witness_utxo = None; // Clear out any non-witness UTXO when we set a witness one
         }
 
-        self.partial_sigs.extend(other.partial_sigs);
-        self.bip32_derivations.extend(other.bip32_derivations);
-        self.ripemd160_preimages.extend(other.ripemd160_preimages);
-        self.sha256_preimages.extend(other.sha256_preimages);
-        self.hash160_preimages.extend(other.hash160_preimages);
-        self.hash256_preimages.extend(other.hash256_preimages);
-        self.tap_script_sigs.extend(other.tap_script_sigs);
-        self.tap_scripts.extend(other.tap_scripts);
-        self.tap_key_origins.extend(other.tap_key_origins);
-        self.proprietaries.extend(other.proprietaries);
-        self.unknowns.extend(other.unknowns);
-
-        combine!(redeem_script, self, other);
-        combine!(witness_script, self, other);
-        combine!(final_script_sig, self, other);
-        combine!(final_script_witness, self, other);
-        combine!(tap_key_sig, self, other);
-        combine!(tap_internal_key, self, other);
-        combine!(tap_merkle_root, self, other);
+        combine_map!(partial_sigs, self, other);
+        // TODO: Why do we not combine sighash_type?
+        combine_option!(redeem_script, self, other);
+        combine_option!(witness_script, self, other);
+        combine_map!(bip32_derivations, self, other);
+        combine_option!(final_script_sig, self, other);
+        combine_option!(final_script_witness, self, other);
+        combine_map!(ripemd160_preimages, self, other);
+        combine_map!(sha256_preimages, self, other);
+        combine_map!(hash160_preimages, self, other);
+        combine_map!(hash256_preimages, self, other);
+        combine_option!(tap_key_sig, self, other);
+        combine_map!(tap_script_sigs, self, other);
+        combine_map!(tap_scripts, self, other);
+        combine_map!(tap_key_origins, self, other);
+        combine_option!(tap_internal_key, self, other);
+        combine_option!(tap_merkle_root, self, other);
+        combine_map!(proprietaries, self, other);
+        combine_map!(unknowns, self, other);
     }
 }
 

src/v0/map/output.rs

@@ -123,15 +123,14 @@ impl Output {
 
     /// Combines this [`Output`] with `other` `Output` (as described by BIP 174).
     pub fn combine(&mut self, other: Self) {
-        self.bip32_derivations.extend(other.bip32_derivations);
-        self.proprietaries.extend(other.proprietaries);
-        self.unknowns.extend(other.unknowns);
-        self.tap_key_origins.extend(other.tap_key_origins);
-
-        combine!(redeem_script, self, other);
-        combine!(witness_script, self, other);
-        combine!(tap_internal_key, self, other);
-        combine!(tap_tree, self, other);
+        combine_option!(redeem_script, self, other);
+        combine_option!(witness_script, self, other);
+        combine_map!(bip32_derivations, self, other);
+        combine_option!(tap_internal_key, self, other);
+        combine_option!(tap_tree, self, other);
+        combine_map!(tap_key_origins, self, other);
+        combine_map!(proprietaries, self, other);
+        combine_map!(unknowns, self, other);
     }
 }
 

src/v0/mod.rs

@@ -26,18 +26,25 @@ use bitcoin::{ecdsa, Amount, ScriptBuf};
 
 use crate::error::{write_err, FeeError, FundingUtxoError};
 use crate::prelude::*;
-use crate::v0::map::Map;
+use crate::v0::map::{global, Map};
 
 #[rustfmt::skip]                // Keep pubic re-exports separate
 pub use self::{
-    error::{IndexOutOfBoundsError, SignerChecksError, SignError, CombineError, UnsignedTxChecksError, DeserializePsbtError},
+    error::{IndexOutOfBoundsError, SignerChecksError, SignError, UnsignedTxChecksError, DeserializePsbtError},
     map::{Input, Output, Global},
 };
 
 #[rustfmt::skip]                // Keep public re-exports separate.
 #[cfg(feature = "base64")]
 pub use self::display_from_str::PsbtParseError;
 
+/// Combines these two PSBTs as described by BIP-174.
+///
+/// This function is commutative `combine(this, that) = combine(that, this)`.
+pub fn combine(this: Psbt, that: Psbt) -> Result<Psbt, global::CombineError> {
+    this.combine_with(that)
+}
+
 /// A Partially Signed Transaction.
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -167,21 +174,23 @@ impl Psbt {
         Ok(psbt)
     }
 
-    /// Combines this [`Psbt`] with `other` PSBT as described by BIP 174.
+    /// Combines this [`Psbt`] with `other` PSBT as described by BIP-174.
     ///
-    /// In accordance with BIP 174 this function is commutative i.e., `A.combine(B) == B.combine(A)`
-    pub fn combine(&mut self, other: Self) -> Result<(), CombineError> {
+    /// This function is commutative `A.combine_with(B) == B.combine_with(A)`
+    pub fn combine_with(mut self, other: Self) -> Result<Psbt, global::CombineError> {
+        // Errors if this PSBT and other do not contain the same unsigned transaction.
         self.global.combine(other.global)?;
 
+        // Ok to zip because inputs vectors are always equal to the unsigned tx input vector and we
+        // checked already that both PSBTs have the same unsigned tx. Some for outputs below.
         for (self_input, other_input) in self.inputs.iter_mut().zip(other.inputs.into_iter()) {
             self_input.combine(other_input);
         }
-
         for (self_output, other_output) in self.outputs.iter_mut().zip(other.outputs.into_iter()) {
             self_output.combine(other_output);
         }
 
-        Ok(())
+        Ok(self)
     }
 
     /// Returns `Ok` if PSBT is
@@ -1284,27 +1293,33 @@ mod tests {
 
     // PSBTs taken from BIP 174 test vectors.
     #[test]
-    fn combine_psbts() {
-        let mut psbt1 = hex_psbt(include_str!("../../tests/data/psbt1.hex")).unwrap();
+    fn combine_works() {
+        let psbt1 = hex_psbt(include_str!("../../tests/data/psbt1.hex")).unwrap();
         let psbt2 = hex_psbt(include_str!("../../tests/data/psbt2.hex")).unwrap();
-        let psbt_combined = hex_psbt(include_str!("../../tests/data/psbt2.hex")).unwrap();
+        let want = hex_psbt(include_str!("../../tests/data/psbt2.hex")).unwrap();
 
-        psbt1.combine(psbt2).expect("psbt combine to succeed");
-        assert_eq!(psbt1, psbt_combined);
+        let got = psbt1.clone().combine_with(psbt2.clone()).expect("failed to combine_with");
+        assert_eq!(got, want);
+
+        // Sanity check, combine works as well.
+        let got = combine(psbt1, psbt2).expect("failed to combine");
+        assert_eq!(got, want);
     }
 
     #[test]
-    fn combine_psbts_commutative() {
-        let mut psbt1 = hex_psbt(include_str!("../../tests/data/psbt1.hex")).unwrap();
-        let mut psbt2 = hex_psbt(include_str!("../../tests/data/psbt2.hex")).unwrap();
-
-        let psbt1_clone = psbt1.clone();
-        let psbt2_clone = psbt2.clone();
+    fn combine_is_commutative() {
+        let psbt1 = hex_psbt(include_str!("../../tests/data/psbt1.hex")).unwrap();
+        let psbt2 = hex_psbt(include_str!("../../tests/data/psbt2.hex")).unwrap();
 
-        psbt1.combine(psbt2_clone).expect("psbt1 combine to succeed");
-        psbt2.combine(psbt1_clone).expect("psbt2 combine to succeed");
+        // `combine_with` consumes.
+        let a = psbt1.clone().combine_with(psbt2.clone()).expect("failed to combine 1 with 2");
+        let b = psbt2.clone().combine_with(psbt1.clone()).expect("failed to combine 2 with 1");
+        assert_eq!(a, b);
 
-        assert_eq!(psbt1, psbt2);
+        // Sanity check, combine works as well.
+        let a = combine(psbt1.clone(), psbt2.clone());
+        let b = combine(psbt2, psbt1);
+        assert_eq!(a, b);
     }
 
     #[cfg(feature = "rand-std")]