simln-lib: add unique index to htlcs in ChannelState

Author: elnosh

simln-lib/src/sim_node.rs

@@ -155,8 +155,11 @@ macro_rules! fail_forwarding_inequality {
 #[derive(Clone)]
 struct ChannelState {
     local_balance_msat: u64,
-    in_flight: HashMap<PaymentHash, Htlc>,
+    /// Maps payment hash to htlc and index that it was added at.
+    in_flight: HashMap<PaymentHash, (Htlc, u64)>,
     policy: ChannelPolicy,
+    /// Tracks unique identifier for htlcs proposed by this node (sent in the outgoing direction).
+    index: u64,
 }
 
 impl ChannelState {
@@ -168,12 +171,13 @@ impl ChannelState {
             local_balance_msat,
             in_flight: HashMap::new(),
             policy,
+            index: 0,
         }
     }
 
     /// Returns the sum of all the *in flight outgoing* HTLCs on the channel.
     fn in_flight_total(&self) -> u64 {
-        self.in_flight.values().map(|h| h.amount_msat).sum()
+        self.in_flight.values().map(|h| h.0.amount_msat).sum()
     }
 
     /// Checks whether the proposed HTLC abides by the channel policy advertised for using this channel as the
@@ -232,21 +236,28 @@ impl ChannelState {
     ///
     /// Note: MPP payments are not currently supported, so this function will fail if a duplicate payment hash is
     /// reported.
-    fn add_outgoing_htlc(&mut self, hash: PaymentHash, htlc: Htlc) -> ForwardResult {
+    fn add_outgoing_htlc(
+        &mut self,
+        hash: PaymentHash,
+        htlc: Htlc,
+    ) -> Result<Result<u64, ForwardingError>, CriticalError> {
         if let Err(fwd_err) = self.check_outgoing_addition(&htlc) {
             return Ok(Err(fwd_err));
         }
 
         if self.in_flight.contains_key(&hash) {
             return Err(CriticalError::PaymentHashExists(hash));
         }
+        let index = self.index;
+        self.index += 1;
+
         self.local_balance_msat -= htlc.amount_msat;
-        self.in_flight.insert(hash, htlc);
-        Ok(Ok(()))
+        self.in_flight.insert(hash, (htlc, index));
+        Ok(Ok(index))
     }
 
     /// Removes the HTLC from our set of outgoing in-flight HTLCs, failing if the payment hash is not found.
-    fn remove_outgoing_htlc(&mut self, hash: &PaymentHash) -> Result<Htlc, CriticalError> {
+    fn remove_outgoing_htlc(&mut self, hash: &PaymentHash) -> Result<(Htlc, u64), CriticalError> {
         self.in_flight
             .remove(hash)
             .ok_or(CriticalError::PaymentHashNotFound(*hash))
@@ -354,19 +365,14 @@ impl SimulatedChannel {
         sending_node: &PublicKey,
         hash: PaymentHash,
         htlc: Htlc,
-    ) -> ForwardResult {
+    ) -> Result<Result<u64, ForwardingError>, CriticalError> {
         if htlc.amount_msat == 0 {
             return Err(CriticalError::ZeroAmountHtlc);
         }
 
-        if let Err(fwd_err) = self
-            .get_node_mut(sending_node)?
-            .add_outgoing_htlc(hash, htlc)?
-        {
-            return Ok(Err(fwd_err));
-        }
         self.sanity_check()?;
-        Ok(Ok(()))
+        self.get_node_mut(sending_node)?
+            .add_outgoing_htlc(hash, htlc)
     }
 
     /// Performs a sanity check on the total balances in a channel. Note that we do not currently include on-chain
@@ -395,12 +401,14 @@ impl SimulatedChannel {
         sending_node: &PublicKey,
         hash: &PaymentHash,
         success: bool,
-    ) -> Result<(), CriticalError> {
+    ) -> Result<(Htlc, u64), CriticalError> {
         let htlc = self
             .get_node_mut(sending_node)?
             .remove_outgoing_htlc(hash)?;
-        self.settle_htlc(sending_node, htlc.amount_msat, success)?;
-        self.sanity_check()
+        self.settle_htlc(sending_node, htlc.0.amount_msat, success)?;
+        self.sanity_check()?;
+
+        Ok(htlc)
     }
 
     /// Updates the local balance of each node in the channel once a htlc has been resolved, pushing funds to the