f convert already_forwarded_htlcs to a hashmap

Author: valentinewallace

lightning/src/ln/channelmanager.rs

@@ -18468,17 +18468,24 @@ impl<
 		// edge or (b) removed from the outbound edge via claim. If it's in neither of these states, we
 		// infer that it was removed from the outbound edge via fail, and fail it backwards to ensure
 		// that it is handled.
-		let mut already_forwarded_htlcs = Vec::new();
-		let prune_forwarded_htlc =
-			|already_forwarded_htlcs: &mut Vec<(PaymentHash, HTLCPreviousHopData, u64)>,
-			 prev_hop: &HTLCPreviousHopData| {
-				if let Some(idx) = already_forwarded_htlcs.iter().position(|(_, htlc, _)| {
-					prev_hop.htlc_id == htlc.htlc_id
-						&& prev_hop.prev_outbound_scid_alias == htlc.prev_outbound_scid_alias
-				}) {
-					already_forwarded_htlcs.swap_remove(idx);
+		let mut already_forwarded_htlcs: HashMap<
+			ChannelId,
+			Vec<(PaymentHash, HTLCPreviousHopData, u64)>,
+		> = new_hash_map();
+		let prune_forwarded_htlc = |already_forwarded_htlcs: &mut HashMap<
+			ChannelId,
+			Vec<(PaymentHash, HTLCPreviousHopData, u64)>,
+		>,
+		                            prev_hop: &HTLCPreviousHopData| {
+			if let hash_map::Entry::Occupied(mut entry) =
+				already_forwarded_htlcs.entry(prev_hop.channel_id)
+			{
+				entry.get_mut().retain(|(_, htlc, _)| prev_hop.htlc_id != htlc.htlc_id);
+				if entry.get().is_empty() {
+					entry.remove();
 				}
-			};
+			}
+		};
 		{
 			// If we're tracking pending payments, ensure we haven't lost any by looking at the
 			// ChannelMonitor data for any channels for which we do not have authorative state
@@ -18523,11 +18530,23 @@ impl<
 											hop_data,
 											outbound_amt_msat,
 										} => {
-											already_forwarded_htlcs.push((
-												payment_hash,
-												hop_data,
-												outbound_amt_msat,
-											));
+											match already_forwarded_htlcs.entry(hop_data.channel_id)
+											{
+												hash_map::Entry::Occupied(mut entry) => {
+													entry.get_mut().push((
+														payment_hash,
+														hop_data,
+														outbound_amt_msat,
+													));
+												},
+												hash_map::Entry::Vacant(entry) => {
+													entry.insert(vec![(
+														payment_hash,
+														hop_data,
+														outbound_amt_msat,
+													)]);
+												},
+											}
 										},
 										InboundUpdateAdd::Legacy => {
 											return Err(DecodeError::InvalidValue)
@@ -19296,44 +19315,50 @@ impl<
 
 		let mut processed_claims: HashSet<Vec<MPPClaimHTLCSource>> = new_hash_set();
 		for (channel_id, monitor) in args.channel_monitors.iter() {
+			let counterparty_node_id = monitor.get_counterparty_node_id();
 			for (payment_hash, (payment_preimage, payment_claims)) in monitor.get_stored_preimages()
 			{
 				// If we have unresolved inbound committed HTLCs that were already forwarded to the
 				// outbound edge and removed via claim, we need to make sure to claim them backwards via
 				// adding them to `pending_claims_to_replay`.
-				for (hash, hop_data, outbound_amt_msat) in
-					mem::take(&mut already_forwarded_htlcs).drain(..)
+				if let hash_map::Entry::Occupied(mut entry) =
+					already_forwarded_htlcs.entry(*channel_id)
 				{
-					if hash != payment_hash {
-						already_forwarded_htlcs.push((hash, hop_data, outbound_amt_msat));
-						continue;
-					}
-					let new_pending_claim = !pending_claims_to_replay.iter().any(|(src, _, _, _, _, _, _)| {
-						matches!(src, HTLCSource::PreviousHopData(hop) if hop.htlc_id == hop_data.htlc_id && hop.prev_outbound_scid_alias == hop_data.prev_outbound_scid_alias)
-					});
-					if new_pending_claim {
-						let counterparty_node_id = monitor.get_counterparty_node_id();
-						let is_channel_closed = channel_manager
-							.per_peer_state
-							.read()
-							.unwrap()
-							.get(&counterparty_node_id)
-							.map_or(true, |peer_state_mtx| {
-								!peer_state_mtx
-									.lock()
-									.unwrap()
-									.channel_by_id
-									.contains_key(channel_id)
+					for (hash, hop_data, outbound_amt_msat) in mem::take(entry.get_mut()) {
+						if hash != payment_hash {
+							entry.get_mut().push((hash, hop_data, outbound_amt_msat));
+							continue;
+						}
+						let new_pending_claim =
+							!pending_claims_to_replay.iter().any(|(src, _, _, _, _, _, _)| {
+								matches!(src, HTLCSource::PreviousHopData(hop) if hop.htlc_id == hop_data.htlc_id && hop.channel_id == hop_data.channel_id)
 							});
-						pending_claims_to_replay.push((
-							HTLCSource::PreviousHopData(hop_data),
-							payment_preimage,
-							outbound_amt_msat,
-							is_channel_closed,
-							counterparty_node_id,
-							monitor.get_funding_txo(),
-							*channel_id,
-						));
+						if new_pending_claim {
+							let is_channel_closed = channel_manager
+								.per_peer_state
+								.read()
+								.unwrap()
+								.get(&counterparty_node_id)
+								.map_or(true, |peer_state_mtx| {
+									!peer_state_mtx
+										.lock()
+										.unwrap()
+										.channel_by_id
+										.contains_key(channel_id)
+								});
+							pending_claims_to_replay.push((
+								HTLCSource::PreviousHopData(hop_data),
+								payment_preimage,
+								outbound_amt_msat,
+								is_channel_closed,
+								counterparty_node_id,
+								monitor.get_funding_txo(),
+								*channel_id,
+							));
+						}
+					}
+					if entry.get().is_empty() {
+						entry.remove();
 					}
 				}
 				if !payment_claims.is_empty() {
@@ -19577,7 +19602,7 @@ impl<
 			channel_manager
 				.fail_htlc_backwards_internal(&source, &hash, &reason, receiver, ev_action);
 		}
-		for (hash, htlc, _) in already_forwarded_htlcs {
+		for (hash, htlc, _) in already_forwarded_htlcs.into_values().flatten() {
 			let channel_id = htlc.channel_id;
 			let node_id = htlc.counterparty_node_id;
 			let source = HTLCSource::PreviousHopData(htlc);