Simplify amortize_fee, correct scheduler's amortizing of branch fees

processor/src/multisigs/scheduler.rs

@@ -423,29 +423,37 @@ impl<N: Network> Scheduler<N> {
       None => return,
     };
 
-    // Amortize the fee amongst all payments
+    // Amortize the fee amongst all payments underneath this branch
-    // While some networks, like Ethereum, may have some payments take notably more gas, those
+    {
-    // payments will have their own gas deducted when they're created. The difference in output
+      let mut to_amortize = actual - expected;
-    // value present here is solely the cost of the branch, which is used for all of these
+      // If the payments are worth less than this fee we need to amortize, return, dropping them
-    // payments, regardless of how much they'll end up costing
+      if payments.iter().map(|payment| payment.amount).sum::<u64>() < to_amortize {
-    let diff = actual - expected;
+        return;
+      }
+      while to_amortize != 0 {
         let payments_len = u64::try_from(payments.len()).unwrap();
-    let per_payment = diff / payments_len;
+        let per_payment = to_amortize / payments_len;
-    // The above division isn't perfect
+        let mut overage = to_amortize % payments_len;
-    let mut remainder = diff - (per_payment * payments_len);
 
         for payment in payments.iter_mut() {
-      // TODO: This usage of saturating_sub is invalid as we *need* to subtract this value
+          let to_subtract = per_payment + overage;
-      payment.amount = payment.amount.saturating_sub(per_payment + remainder);
+          // Only subtract the overage once
-      // Only subtract the remainder once
+          overage = 0;
-      remainder = 0;
+
+          let subtractable = payment.amount.min(to_subtract);
+          to_amortize -= subtractable;
+          payment.amount -= subtractable;
+        }
+      }
     }
 
     // Drop payments now below the dust threshold
     let payments =
-      payments.drain(..).filter(|payment| payment.amount >= N::DUST).collect::<Vec<_>>();
+      payments.into_iter().filter(|payment| payment.amount >= N::DUST).collect::<Vec<_>>();
     // Sanity check this was done properly
     assert!(actual >= payments.iter().map(|payment| payment.amount).sum::<u64>());
+
+    // If there's no payments left, return
     if payments.is_empty() {
       return;
     }

processor/src/networks/mod.rs

@@ -26,7 +26,7 @@ pub mod monero;
 #[cfg(feature = "monero")]
 pub use monero::Monero;
 
-use crate::{Payment, Plan};
+use crate::Plan;
 
 #[derive(Clone, Copy, Error, Debug)]
 pub enum NetworkError {
@@ -343,7 +343,7 @@ pub trait Network: 'static + Send + Sync + Clone + PartialEq + Eq + Debug {
     // Sanity check this has at least one output planned
     assert!((!plan.payments.is_empty()) || plan.change.is_some());
 
-    let Some(fee) = self.needed_fee(block_number, &plan, fee_rate).await? else {
+    let Some(tx_fee) = self.needed_fee(block_number, &plan, fee_rate).await? else {
       // This Plan is not fulfillable
       // TODO: Have Plan explicitly distinguish payments and branches in two separate Vecs?
       return Ok(PreparedSend {
@@ -358,28 +358,11 @@ pub trait Network: 'static + Send + Sync + Clone + PartialEq + Eq + Debug {
       });
     };
 
-    let (post_fee_branches, mut operating_costs) = {
-      pub struct AmortizeFeeRes {
-        post_fee_branches: Vec<PostFeeBranch>,
-        operating_costs: u64,
-      }
-
     // Amortize a fee over the plan's payments
-      fn amortize_fee<N: Network>(
+    let (post_fee_branches, mut operating_costs) = (|| {
-        plan: &mut Plan<N>,
+      // If we're creating a change output, letting us recoup coins, amortize the operating costs
-        operating_costs: u64,
-        tx_fee: u64,
-      ) -> AmortizeFeeRes {
-        let total_fee = {
-          let mut total_fee = tx_fee;
-          // Since we're creating a change output, letting us recoup coins, amortize the operating
-          // costs
       // as well
-          if plan.change.is_some() {
+      let total_fee = tx_fee + if plan.change.is_some() { operating_costs } else { 0 };
-            total_fee += operating_costs;
-          }
-          total_fee
-        };
 
       let original_outputs = plan.payments.iter().map(|payment| payment.amount).sum::<u64>();
       // If this isn't enough for the total fee, drop and move on
@@ -391,55 +374,48 @@ pub trait Network: 'static + Send + Sync + Clone + PartialEq + Eq + Debug {
           // Yet decrease by the payments we managed to drop
           remaining_operating_costs = remaining_operating_costs.saturating_sub(original_outputs);
         }
-          return AmortizeFeeRes {
+        return (drop_branches(&plan), remaining_operating_costs);
-            post_fee_branches: drop_branches(plan),
-            operating_costs: remaining_operating_costs,
-          };
       }
 
+      let initial_payment_amounts =
+        plan.payments.iter().map(|payment| payment.amount).collect::<Vec<_>>();
+
       // Amortize the transaction fee across outputs
-        let mut payments_len = u64::try_from(plan.payments.len()).unwrap();
+      let mut remaining_fee = total_fee;
-        // Use a formula which will round up
+      // Run as many times as needed until we can successfully subtract this fee
-        let per_output_fee = |payments| (total_fee + (payments - 1)) / payments;
+      while remaining_fee != 0 {
+        // This shouldn't be a / by 0 as these payments have enough value to cover the fee
+        let this_iter_fee = remaining_fee / u64::try_from(plan.payments.len()).unwrap();
+        let mut overage = remaining_fee % u64::try_from(plan.payments.len()).unwrap();
+        for payment in &mut plan.payments {
+          let this_payment_fee = this_iter_fee + overage;
+          // Only subtract the overage once
+          overage = 0;
 
-        let post_fee = |payment: &Payment<N>, per_output_fee| {
+          let subtractable = payment.amount.min(this_payment_fee);
-          let mut post_fee = payment.amount.checked_sub(per_output_fee);
+          remaining_fee -= subtractable;
-          // If this is under our dust threshold, drop it
+          payment.amount -= subtractable;
-          if let Some(amount) = post_fee {
-            if amount < N::DUST {
-              post_fee = None;
         }
       }
-          post_fee
-        };
 
-        // If we drop outputs for being less than the fee, we won't successfully reduce the amount
+      // If any payment is now below the dust threshold, set its value to 0 so it'll be dropped
-        // spent (dropping a 800 output due to a 1000 fee leaves 200 we still have to deduct)
+      for payment in &mut plan.payments {
-        // Do initial runs until the amount of output we will drop is known
+        if payment.amount < Self::DUST {
-        while {
+          payment.amount = 0;
-          let last = payments_len;
+        }
-          payments_len = u64::try_from(
+      }
-            plan
-              .payments
-              .iter()
-              .filter(|payment| post_fee(payment, per_output_fee(payments_len)).is_some())
-              .count(),
-          )
-          .unwrap();
-          last != payments_len
-        } {}
 
-        // Now that we know how many outputs will survive, calculate the actual per_output_fee
+      // Note the branch outputs' new values
-        let per_output_fee = per_output_fee(payments_len);
       let mut branch_outputs = vec![];
-        for payment in plan.payments.iter_mut() {
+      for (initial_amount, payment) in initial_payment_amounts.into_iter().zip(&plan.payments) {
-          let post_fee = post_fee(payment, per_output_fee);
+        if payment.address == Self::branch_address(plan.key) {
-          // Note the branch output, if this is one
+          branch_outputs.push(PostFeeBranch {
-          if payment.address == N::branch_address(plan.key) {
+            expected: initial_amount,
-            branch_outputs.push(PostFeeBranch { expected: payment.amount, actual: post_fee });
+            actual: if payment.amount == 0 { None } else { Some(payment.amount) },
+          });
         }
-          payment.amount = post_fee.unwrap_or(0);
       }
+
       // Drop payments now worth 0
       plan.payments = plan.payments.drain(..).filter(|payment| payment.amount != 0).collect();
 
@@ -447,9 +423,9 @@ pub trait Network: 'static + Send + Sync + Clone + PartialEq + Eq + Debug {
       let new_outputs = plan.payments.iter().map(|payment| payment.amount).sum::<u64>();
       assert!((new_outputs + total_fee) <= original_outputs);
 
-        AmortizeFeeRes {
+      (
-          post_fee_branches: branch_outputs,
+        branch_outputs,
-          operating_costs: if plan.change.is_none() {
+        if plan.change.is_none() {
           // If the change is None, this had no effect on the operating costs
           operating_costs
         } else {
@@ -457,21 +433,15 @@ pub trait Network: 'static + Send + Sync + Clone + PartialEq + Eq + Debug {
           // recouped
           0
         },
-        }
+      )
-      }
+    })();
-
-      let AmortizeFeeRes { post_fee_branches, operating_costs } =
-        amortize_fee(&mut plan, operating_costs, fee);
-
-      (post_fee_branches, operating_costs)
-    };
 
     let Some(tx) = self.signable_transaction(block_number, &plan, fee_rate).await? else {
       panic!(
         "{}. post-amortization plan: {:?}, successfully amoritized fee: {}",
         "signable_transaction returned None for a TX we prior successfully calculated the fee for",
         &plan,
-        fee,
+        tx_fee,
       )
     };
 
@@ -485,7 +455,7 @@ pub trait Network: 'static + Send + Sync + Clone + PartialEq + Eq + Debug {
       // to bare it)
       // This call wants the actual value, post-amortization over outputs, and since Plan is
       // unaware of the fee, has to manually adjust
-      let on_chain_expected_change = plan.expected_change() - fee;
+      let on_chain_expected_change = plan.expected_change() - tx_fee;
       // If the change value is less than the dust threshold, it becomes an operating cost
       // This may be slightly inaccurate as dropping payments may reduce the fee, raising the
       // change above dust