Remove a clone of all the non-dust htlcs when rebuilding a commitment tx

We leverage the `Borrow` trait to allow passing both exclusive and
shared references to `CommitmentTransaction::internal_build_outputs`.
This allows us to avoid cloning all the htlcs when a
`CommitmentTransaction::verify` call recreates the
`BuiltCommitmentTransaction`.

The iterator trait used in
`CommitmentTransaction::internal_build_outputs` allows us to handle both
collections of pointers, and a pointer to a collection without requiring
any memory allocations.

We also add extensive documentation of the inner mechanics of
`CommitmentTransaction`.

Author: tankyleo

lightning/src/ln/chan_utils.rs

@@ -40,6 +40,7 @@ use bitcoin::secp256k1::{Secp256k1, ecdsa::Signature, Message};
 use bitcoin::{secp256k1, Sequence, Witness};
 
 use crate::io;
+use core::borrow::Borrow;
 use core::cmp;
 use crate::util::transaction_utils::sort_outputs;
 use crate::ln::channel::{INITIAL_COMMITMENT_NUMBER, ANCHOR_OUTPUT_VALUE_SATOSHI};
@@ -1463,7 +1464,7 @@ impl CommitmentTransaction {
 		let keys = TxCreationKeys::from_channel_static_keys(per_commitment_point, channel_parameters.broadcaster_pubkeys(), channel_parameters.countersignatory_pubkeys(), secp_ctx);
 
 		// Sort outputs and populate output indices
-		let (outputs, nondust_htlcs) = Self::internal_build_outputs(&keys, to_broadcaster_value_sat, to_countersignatory_value_sat, nondust_htlcs, channel_parameters).unwrap();
+		let (outputs, nondust_htlcs) = Self::build_outputs_and_htlcs(&keys, to_broadcaster_value_sat, to_countersignatory_value_sat, nondust_htlcs, channel_parameters).unwrap();
 
 		let (obscured_commitment_transaction_number, txins) = Self::internal_build_inputs(commitment_number, channel_parameters);
 		let transaction = Self::make_transaction(obscured_commitment_transaction_number, txins, outputs);
@@ -1492,11 +1493,17 @@ impl CommitmentTransaction {
 		self
 	}
 
+	// This function
+	// - Builds the inputs of the commitment transaction.
+	// - Calls `internal_build_outputs` to get a `Vec<(TxOut, Option<&HTLCOutputInCommitment>)>` sorted according to the LN specification,
+	//   and as they'd appear on the commitment transaction.
+	// - Maps this vector to a `Vec<TxOut>`, keeping the original order of the outputs.
+	// - Then uses the vectors of inputs and outputs to build the commitment transaction.
 	fn internal_rebuild_transaction(&self, keys: &TxCreationKeys, channel_parameters: &DirectedChannelTransactionParameters) -> Result<BuiltCommitmentTransaction, ()> {
 		let (obscured_commitment_transaction_number, txins) = Self::internal_build_inputs(self.commitment_number, channel_parameters);
 
-		let mut nondust_htlcs = self.htlcs.clone();
-		let (outputs, _) = Self::internal_build_outputs(keys, self.to_broadcaster_value_sat, self.to_countersignatory_value_sat, nondust_htlcs.iter_mut().collect(), channel_parameters)?;
+		let txout_htlc_pairs = Self::internal_build_outputs(keys, self.to_broadcaster_value_sat, self.to_countersignatory_value_sat, self.htlcs().iter(), channel_parameters)?;
+		let outputs = txout_htlc_pairs.into_iter().map(|(output, _)| output).collect();
 
 		let transaction = Self::make_transaction(obscured_commitment_transaction_number, txins, outputs);
 		let txid = transaction.compute_txid();
@@ -1516,17 +1523,58 @@ impl CommitmentTransaction {
 		}
 	}
 
-	// This is used in two cases:
-	// - initial sorting of outputs / HTLCs in the constructor, in which case T is auxiliary data the
-	//   caller needs to have sorted together with the HTLCs so it can keep track of the output index
-	// - building of a bitcoin transaction during a verify() call, in which case T is just ()
-	fn internal_build_outputs(keys: &TxCreationKeys, to_broadcaster_value_sat: Amount, to_countersignatory_value_sat: Amount, nondust_htlcs: Vec<&mut HTLCOutputInCommitment>, channel_parameters: &DirectedChannelTransactionParameters) -> Result<(Vec<TxOut>, Vec<HTLCOutputInCommitment>), ()> {
+	// This function
+	// - Calls `internal_build_outputs` to get a `Vec<(TxOut, Option<&mut HTLCOutputInCommitment>)>` sorted according to the LN specification,
+	//   and as they'd appear on the commitment transaction.
+	// - Then iterates over this vector of tuples:
+	//   - populates *all* the output indices of the exclusive references in `nondust_htlcs`,
+	//   - creates a vector of `HTLCOutputInCommitment` sorted by their output index,
+	//   - creates a vector of `TxOut`, sorted according to the LN Spec, see `sort_outputs`.
+	// - Returns the two vectors: the vector of `TxOut`, and the vector of `HTLCOutputInCommitment`:
+	//   - `Vec<TxOut>` gets moved to the `Transaction` cached by `CommitmentTransaction`,
+	//   - `Vec<HTLCOutputInCommitment>` gets moved to and cached by `CommitmentTransaction`.
+	fn build_outputs_and_htlcs(keys: &TxCreationKeys, to_broadcaster_value_sat: Amount, to_countersignatory_value_sat: Amount, nondust_htlcs: Vec<&mut HTLCOutputInCommitment>, channel_parameters: &DirectedChannelTransactionParameters) -> Result<(Vec<TxOut>, Vec<HTLCOutputInCommitment>), ()> {
+		let mut htlcs: Vec<HTLCOutputInCommitment> = Vec::with_capacity(nondust_htlcs.len());
+		let mut txout_htlc_pairs: Vec<(TxOut, Option<&mut HTLCOutputInCommitment>)> = Self::internal_build_outputs(keys, to_broadcaster_value_sat, to_countersignatory_value_sat, nondust_htlcs.into_iter(), channel_parameters)?;
+		let mut outputs: Vec<TxOut> = Vec::with_capacity(txout_htlc_pairs.len());
+		for (idx, out) in txout_htlc_pairs.drain(..).enumerate() {
+			if let Some(htlc) = out.1 {
+				// Remember `htlc` is a `&mut HTLCOutputInCommitment`, so this populates the HTLC exclusive references passed to `CommitmentTransaction::new`
+				htlc.transaction_output_index = Some(idx as u32);
+				// We also clone the htlcs to cache them in `CommitmentTransaction`
+				htlcs.push(htlc.clone());
+			}
+			outputs.push(out.0);
+		}
+		Ok((outputs, htlcs))
+	}
+
+	// This function
+	// - Takes an iterator that yields objects of type `T` that can yield a `&HTLCOutputInCommitment`; `T: Borrow<HTLCOutputInCommitment>`.
+	//   For our use-cases, these are `&mut HTLCOutputInCommitment` (used by `CommitmentTransaction::new`), and `&HTLCOutputInCommitment` (used by `CommitmentTransaction::internal_rebuild_transaction`).
+	// - Builds a `Vec<(TxOut, Option<T>)>`, containing all the outputs of the commitment transaction.
+	//   For the `Option` in the tuple, an HTLC output sets the option to `Some(T)`, and all the other outputs set the option to `None`.
+	// - Passes this vector to `sort_outputs` to sort all the tuples in the vector according to the LN specification.
+	// - Then returns this sorted vector to the caller.
+	fn internal_build_outputs<T: Borrow<HTLCOutputInCommitment>>(keys: &TxCreationKeys, to_broadcaster_value_sat: Amount, to_countersignatory_value_sat: Amount, nondust_htlcs: impl Iterator<Item = T>, channel_parameters: &DirectedChannelTransactionParameters) -> Result<Vec<(TxOut, Option<T>)>, ()> {
 		let countersignatory_pubkeys = channel_parameters.countersignatory_pubkeys();
 		let broadcaster_funding_key = &channel_parameters.broadcaster_pubkeys().funding_pubkey;
 		let countersignatory_funding_key = &countersignatory_pubkeys.funding_pubkey;
 		let contest_delay = channel_parameters.contest_delay();
 
-		let mut txouts: Vec<(TxOut, Option<&mut HTLCOutputInCommitment>)> = Vec::new();
+		// Assumes the max number of outputs on the transaction is equal to n_htlcs + to_local + to_remote + local_anchor + remote_anchor
+		let mut txouts: Vec<(TxOut, Option<T>)> = Vec::with_capacity(nondust_htlcs.size_hint().0 + 4);
+
+		for htlc in nondust_htlcs {
+			let script = get_htlc_redeemscript(htlc.borrow(), &channel_parameters.channel_type_features(), &keys);
+			let txout = TxOut {
+				script_pubkey: script.to_p2wsh(),
+				value: htlc.borrow().to_bitcoin_amount(),
+			};
+			txouts.push((txout, Some(htlc)));
+		}
+		// We avoid relying on `ExactSizeIterator` here
+		let nondust_htlcs_is_empty = txouts.is_empty();
 
 		if to_countersignatory_value_sat > Amount::ZERO {
 			let script = if channel_parameters.channel_type_features().supports_anchors_zero_fee_htlc_tx() {
@@ -1559,7 +1607,7 @@ impl CommitmentTransaction {
 		}
 
 		if channel_parameters.channel_type_features().supports_anchors_zero_fee_htlc_tx() {
-			if to_broadcaster_value_sat > Amount::ZERO || !nondust_htlcs.is_empty() {
+			if to_broadcaster_value_sat > Amount::ZERO || !nondust_htlcs_is_empty {
 				let anchor_script = get_anchor_redeemscript(broadcaster_funding_key);
 				txouts.push((
 					TxOut {
@@ -1570,7 +1618,7 @@ impl CommitmentTransaction {
 				));
 			}
 
-			if to_countersignatory_value_sat > Amount::ZERO || !nondust_htlcs.is_empty() {
+			if to_countersignatory_value_sat > Amount::ZERO || !nondust_htlcs_is_empty {
 				let anchor_script = get_anchor_redeemscript(countersignatory_funding_key);
 				txouts.push((
 					TxOut {
@@ -1582,41 +1630,23 @@ impl CommitmentTransaction {
 			}
 		}
 
-		let mut htlcs = Vec::with_capacity(nondust_htlcs.len());
-		for htlc in nondust_htlcs {
-			let script = get_htlc_redeemscript(&htlc, &channel_parameters.channel_type_features(), &keys);
-			let txout = TxOut {
-				script_pubkey: script.to_p2wsh(),
-				value: htlc.to_bitcoin_amount(),
-			};
-			txouts.push((txout, Some(htlc)));
-		}
-
 		// Sort output in BIP-69 order (amount, scriptPubkey).  Tie-breaks based on HTLC
 		// CLTV expiration height.
 		sort_outputs(&mut txouts, |a, b| {
-			if let &Some(ref a_htlcout) = a {
-				if let &Some(ref b_htlcout) = b {
-					a_htlcout.cltv_expiry.cmp(&b_htlcout.cltv_expiry)
+			if let Some(a_htlcout) = a {
+				if let Some(b_htlcout) = b {
+					a_htlcout.borrow().cltv_expiry.cmp(&b_htlcout.borrow().cltv_expiry)
 						// Note that due to hash collisions, we have to have a fallback comparison
 						// here for fuzzing mode (otherwise at least chanmon_fail_consistency
 						// may fail)!
-						.then(a_htlcout.payment_hash.0.cmp(&b_htlcout.payment_hash.0))
+						.then(a_htlcout.borrow().payment_hash.0.cmp(&b_htlcout.borrow().payment_hash.0))
 				// For non-HTLC outputs, if they're copying our SPK we don't really care if we
 				// close the channel due to mismatches - they're doing something dumb:
 				} else { cmp::Ordering::Equal }
 			} else { cmp::Ordering::Equal }
 		});
 
-		let mut outputs = Vec::with_capacity(txouts.len());
-		for (idx, out) in txouts.drain(..).enumerate() {
-			if let Some(htlc) = out.1 {
-				htlc.transaction_output_index = Some(idx as u32);
-				htlcs.push(htlc.clone());
-			}
-			outputs.push(out.0);
-		}
-		Ok((outputs, htlcs))
+		Ok(txouts)
 	}
 
 	fn internal_build_inputs(commitment_number: u64, channel_parameters: &DirectedChannelTransactionParameters) -> (u64, Vec<TxIn>) {