Add ChannelContext::for_each_{inbound, outbound}

These functions sort included from excluded htlcs based on their state,
and run the passed closure on each sorted htlc.

Author: tankyleo

lightning/src/ln/channel.rs

@@ -240,17 +240,6 @@ impl InboundHTLCState {
 			_ => None,
 		}
 	}
-
-	// Determines whether a HTLC is included in a commitment transaction, either as a dust or non-dust HTLC
-	fn included_in_commitment(&self, generated_by_local: bool) -> bool {
-		match self {
-			InboundHTLCState::RemoteAnnounced(_) => !generated_by_local,
-			InboundHTLCState::AwaitingRemoteRevokeToAnnounce(_) => !generated_by_local,
-			InboundHTLCState::AwaitingAnnouncedRemoteRevoke(_) => true,
-			InboundHTLCState::Committed => true,
-			InboundHTLCState::LocalRemoved(_) => !generated_by_local,
-		}
-	}
 }
 
 struct InboundHTLCOutput {
@@ -353,17 +342,6 @@ impl OutboundHTLCState {
 			_ => None,
 		}
 	}
-
-	// Determines whether a HTLC is included in a commitment transaction, either as a dust or non-dust HTLC
-	fn included_in_commitment(&self, generated_by_local: bool) -> bool {
-		match self {
-			OutboundHTLCState::LocalAnnounced(_) => generated_by_local,
-			OutboundHTLCState::Committed => true,
-			OutboundHTLCState::RemoteRemoved(_) => generated_by_local,
-			OutboundHTLCState::AwaitingRemoteRevokeToRemove(_) => generated_by_local,
-			OutboundHTLCState::AwaitingRemovedRemoteRevoke(_) => false,
-		}
-	}
 }
 
 #[derive(Clone)]
@@ -974,10 +952,10 @@ struct HTLCStats {
 }
 
 /// A struct gathering data on a commitment, either local or remote.
-struct CommitmentData<'a> {
+struct CommitmentData {
 	tx: CommitmentTransaction,
 	stats: CommitmentStats,
-	htlcs_included: Vec<(HTLCOutputInCommitment, Option<&'a HTLCSource>)>, // the list of HTLCs (dust HTLCs *included*) which were not ignored when building the transaction
+	htlcs_included: Vec<(HTLCOutputInCommitment, Option<Box<HTLCSource>>)>, // the list of HTLCs (dust HTLCs *included*) which were not ignored when building the transaction
 	outbound_htlc_preimages: Vec<PaymentPreimage>, // preimages for successful offered HTLCs since last commitment
 	inbound_htlc_preimages: Vec<PaymentPreimage>, // preimages for successful received HTLCs since last commitment
 }
@@ -3604,7 +3582,6 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 			}
 			bitcoin_tx.txid
 		};
-		let mut htlcs_cloned: Vec<_> = commitment_data.htlcs_included.iter().map(|htlc| (htlc.0.clone(), htlc.1.map(|h| h.clone()))).collect();
 
 		// If our counterparty updated the channel fee in this commitment transaction, check that
 		// they can actually afford the new fee now.
@@ -3655,10 +3632,10 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 			separate_nondust_htlc_sources = rand_val % 2 == 0;
 		}
 
-		let mut nondust_htlc_sources = Vec::with_capacity(htlcs_cloned.len());
-		let mut htlcs_and_sigs = Vec::with_capacity(htlcs_cloned.len());
+		let mut nondust_htlc_sources = Vec::with_capacity(commitment_data.htlcs_included.len());
+		let mut htlcs_and_sigs = Vec::with_capacity(commitment_data.htlcs_included.len());
 		let holder_keys = commitment_data.tx.trust().keys();
-		for (idx, (htlc, mut source_opt)) in htlcs_cloned.drain(..).enumerate() {
+		for (idx, (htlc, mut source_opt)) in commitment_data.htlcs_included.into_iter().enumerate() {
 			if let Some(_) = htlc.transaction_output_index {
 				let htlc_tx = chan_utils::build_htlc_transaction(&commitment_txid, commitment_data.tx.feerate_per_kw(),
 					funding.get_counterparty_selected_contest_delay().unwrap(), &htlc, &self.channel_type,
@@ -3674,14 +3651,14 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 					return Err(ChannelError::close("Invalid HTLC tx signature from peer".to_owned()));
 				}
 				if !separate_nondust_htlc_sources {
-					htlcs_and_sigs.push((htlc, Some(msg.htlc_signatures[idx]), source_opt.take()));
+					htlcs_and_sigs.push((htlc, Some(msg.htlc_signatures[idx]), source_opt.take().map(|source| *source)));
 				}
 			} else {
-				htlcs_and_sigs.push((htlc, None, source_opt.take()));
+				htlcs_and_sigs.push((htlc, None, source_opt.take().map(|source| *source)));
 			}
 			if separate_nondust_htlc_sources {
 				if let Some(source) = source_opt.take() {
-					nondust_htlc_sources.push(source);
+					nondust_htlc_sources.push(*source);
 				}
 			}
 			debug_assert!(source_opt.is_none(), "HTLCSource should have been put somewhere");
@@ -3705,6 +3682,36 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 		})
 	}
 
+	fn for_each_inbound<F>(&self, generated_by_local: bool, mut for_each: F)
+		where F: FnMut(bool, &InboundHTLCOutput)
+	{
+		self.pending_inbound_htlcs.iter().for_each(|htlc| {
+			let included = match htlc.state {
+				InboundHTLCState::RemoteAnnounced(_) => !generated_by_local,
+				InboundHTLCState::AwaitingRemoteRevokeToAnnounce(_) => !generated_by_local,
+				InboundHTLCState::AwaitingAnnouncedRemoteRevoke(_) => true,
+				InboundHTLCState::Committed => true,
+				InboundHTLCState::LocalRemoved(_) => !generated_by_local,
+			};
+			for_each(included, htlc);
+		})
+	}
+
+	fn for_each_outbound<F>(&self, generated_by_local: bool, mut for_each: F)
+		where F: FnMut(bool, &OutboundHTLCOutput),
+	{
+		self.pending_outbound_htlcs.iter().for_each(|htlc| {
+			let included = match htlc.state {
+				OutboundHTLCState::LocalAnnounced(_) => generated_by_local,
+				OutboundHTLCState::Committed => true,
+				OutboundHTLCState::RemoteRemoved(_) => generated_by_local,
+				OutboundHTLCState::AwaitingRemoteRevokeToRemove(_) => generated_by_local,
+				OutboundHTLCState::AwaitingRemovedRemoteRevoke(_) => false,
+			};
+			for_each(included, htlc);
+		})
+	}
+
 	/// Generates stats on a potential commitment transaction build, without actually building the
 	/// commitment transaction. See `build_commitment_transaction` for further docs.
 	#[inline]
@@ -3745,37 +3752,34 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 			}
 		}
 
-		for ref htlc in self.pending_inbound_htlcs.iter() {
-			if htlc.state.included_in_commitment(generated_by_local) {
+		let for_each = |included: bool, htlc: &InboundHTLCOutput| {
+			if included {
 				count_nondust_htlc!(htlc, false);
 				remote_htlc_total_msat += htlc.amount_msat;
 			} else {
 				log_trace!(logger, "   ...not counting inbound HTLC {} (hash {}) with value {} due to state ({})", htlc.htlc_id, htlc.payment_hash, htlc.amount_msat, htlc.state.as_str());
-				match htlc.state {
-					InboundHTLCState::LocalRemoved(InboundHTLCRemovalReason::Fulfill(_preimage)) => {
-						value_to_self_msat_offset += htlc.amount_msat as i64;
-					},
-					_ => {},
+				if let InboundHTLCState::LocalRemoved(InboundHTLCRemovalReason::Fulfill(_preimage)) = htlc.state {
+					value_to_self_msat_offset += htlc.amount_msat as i64;
 				}
 			}
-		}
+		};
+		self.for_each_inbound(generated_by_local, for_each);
 
-		for ref htlc in self.pending_outbound_htlcs.iter() {
-			if htlc.state.included_in_commitment(generated_by_local) {
+		let for_each = |included: bool, htlc: &OutboundHTLCOutput| {
+			if included {
 				count_nondust_htlc!(htlc, true);
 				local_htlc_total_msat += htlc.amount_msat;
 			} else {
 				log_trace!(logger, "   ...not counting outbound HTLC {} (hash {}) with value {} due to state ({})", htlc.htlc_id, htlc.payment_hash, htlc.amount_msat, htlc.state.as_str());
-				match htlc.state {
-					OutboundHTLCState::AwaitingRemoteRevokeToRemove(OutboundHTLCOutcome::Success(_)) |
+				if let OutboundHTLCState::AwaitingRemoteRevokeToRemove(OutboundHTLCOutcome::Success(_)) |
 					OutboundHTLCState::AwaitingRemovedRemoteRevoke(OutboundHTLCOutcome::Success(_)) |
-					OutboundHTLCState::RemoteRemoved(OutboundHTLCOutcome::Success(_)) => {
-						value_to_self_msat_offset -= htlc.amount_msat as i64;
-					},
-					_ => {},
+					OutboundHTLCState::RemoteRemoved(OutboundHTLCOutcome::Success(_))
+				= htlc.state {
+					value_to_self_msat_offset -= htlc.amount_msat as i64;
 				}
 			}
-		}
+		};
+		self.for_each_outbound(generated_by_local, for_each);
 
 		// # Panics
 		//
@@ -3844,9 +3848,9 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 			remote_balance_before_fee_anchors_msat
 		} = stats;
 
-		let mut included_dust_htlcs: Vec<(HTLCOutputInCommitment, Option<&HTLCSource>)> = Vec::new();
+		let mut included_dust_htlcs: Vec<(HTLCOutputInCommitment, Option<Box<HTLCSource>>)> = Vec::new();
 		let num_htlcs = self.pending_inbound_htlcs.len() + self.pending_outbound_htlcs.len();
-		let mut included_non_dust_htlcs: Vec<(HTLCOutputInCommitment, Option<&HTLCSource>)> = Vec::with_capacity(num_htlcs);
+		let mut included_non_dust_htlcs: Vec<(HTLCOutputInCommitment, Option<Box<HTLCSource>>)> = Vec::with_capacity(num_htlcs);
 
 		log_trace!(logger, "Building commitment transaction number {} (really {} xor {}) for channel {} for {}, generated by {} with fee {}...",
 			commitment_number, (INITIAL_COMMITMENT_NUMBER - commitment_number),
@@ -3880,30 +3884,32 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider {
 		}
 
 		let mut inbound_htlc_preimages: Vec<PaymentPreimage> = Vec::new();
+		let mut outbound_htlc_preimages: Vec<PaymentPreimage> = Vec::new();
 
-		for ref htlc in self.pending_inbound_htlcs.iter() {
-			if htlc.state.included_in_commitment(generated_by_local) {
+		let for_each = |included: bool, htlc: &InboundHTLCOutput| {
+			if included {
 				add_htlc_output!(htlc, false, None);
 			} else {
 				log_trace!(logger, "   ...not including inbound HTLC {} (hash {}) with value {} due to state ({})", htlc.htlc_id, htlc.payment_hash, htlc.amount_msat, htlc.state.as_str());
 				if let Some(preimage) = htlc.state.preimage() {
 					inbound_htlc_preimages.push(preimage);
 				}
 			}
-		}
-
-		let mut outbound_htlc_preimages: Vec<PaymentPreimage> = Vec::new();
+		};
+		self.for_each_inbound(generated_by_local, for_each);
 
-		for ref htlc in self.pending_outbound_htlcs.iter() {
+		let for_each = |included: bool, htlc: &OutboundHTLCOutput| {
 			if let Some(preimage) = htlc.state.preimage() {
 				outbound_htlc_preimages.push(preimage);
 			}
-			if htlc.state.included_in_commitment(generated_by_local) {
-				add_htlc_output!(htlc, true, Some(&htlc.source));
+			if included {
+				// We box the source here because borrowed data cannot escape a closure
+				add_htlc_output!(htlc, true, Some(Box::new(htlc.source.clone())));
 			} else {
 				log_trace!(logger, "   ...not including outbound HTLC {} (hash {}) with value {} due to state ({})", htlc.htlc_id, htlc.payment_hash, htlc.amount_msat, htlc.state.as_str());
 			}
-		}
+		};
+		self.for_each_outbound(generated_by_local, for_each);
 
 		// We MUST use saturating subs here, as the funder's balance is not guaranteed to be greater
 		// than or equal to the sum of `total_fee_sat` and `total_anchors_sat`.
@@ -8814,11 +8820,9 @@ impl<SP: Deref> FundedChannel<SP> where
 		}
 		self.context.resend_order = RAACommitmentOrder::RevokeAndACKFirst;
 
-		let (mut htlcs_ref, counterparty_commitment_tx) =
+		let (htlcs, counterparty_commitment_tx) =
 			self.build_commitment_no_state_update(logger);
 		let counterparty_commitment_txid = counterparty_commitment_tx.trust().txid();
-		let htlcs: Vec<(HTLCOutputInCommitment, Option<Box<HTLCSource>>)> =
-			htlcs_ref.drain(..).map(|(htlc, htlc_source)| (htlc, htlc_source.map(|source_ref| Box::new(source_ref.clone())))).collect();
 
 		if self.context.announcement_sigs_state == AnnouncementSigsState::MessageSent {
 			self.context.announcement_sigs_state = AnnouncementSigsState::Committed;
@@ -8845,7 +8849,7 @@ impl<SP: Deref> FundedChannel<SP> where
 	}
 
 	fn build_commitment_no_state_update<L: Deref>(&self, logger: &L)
-	-> (Vec<(HTLCOutputInCommitment, Option<&HTLCSource>)>, CommitmentTransaction)
+	-> (Vec<(HTLCOutputInCommitment, Option<Box<HTLCSource>>)>, CommitmentTransaction)
 	where L::Target: Logger
 	{
 		let commitment_data = self.context.build_commitment_transaction(&self.funding,