feat(sequencer): make ABCI response for account balances deterministic

crates/astria-core/src/primitive/v1/asset/denom.rs

@@ -147,6 +147,29 @@ impl FromStr for Denom {
     }
 }
 
+impl PartialOrd for Denom {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for Denom {
+    /// If the denoms are the same type, returns their comparison. Otherwise, returns IBC prefixed
+    /// denoms as less than trace prefixed denoms.
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        match self {
+            Self::TracePrefixed(lhs) => match other {
+                Self::TracePrefixed(rhs) => lhs.cmp(rhs),
+                Self::IbcPrefixed(_) => std::cmp::Ordering::Greater,
+            },
+            Self::IbcPrefixed(lhs) => match other {
+                Self::IbcPrefixed(rhs) => lhs.cmp(rhs),
+                Self::TracePrefixed(_) => std::cmp::Ordering::Less,
+            },
+        }
+    }
+}
+
 #[derive(Debug, thiserror::Error)]
 #[error(transparent)]
 pub struct ParseDenomError(ParseDenomErrorKind);
@@ -284,6 +307,21 @@ impl TracePrefixed {
     }
 }
 
+impl PartialOrd for TracePrefixed {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for TracePrefixed {
+    /// Returns trace comparison if not equal, otherwise returns base denom comparison.
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.trace
+            .cmp(&other.trace)
+            .then_with(|| self.base_denom.cmp(&other.base_denom))
+    }
+}
+
 #[derive(Debug, Clone, Hash, PartialEq, Eq)]
 struct TraceSegments {
     inner: VecDeque<PortAndChannel>,
@@ -354,6 +392,28 @@ impl FromStr for TraceSegments {
         Ok(parsed_segments)
     }
 }
+
+impl PartialOrd for TraceSegments {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for TraceSegments {
+    /// Returns the first non-equal comparison between the two trace segments. If one doesn't exist,
+    /// returns the shortest, and if they are equal, returns equal.
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.inner
+            .iter()
+            .zip(other.inner.iter())
+            .find_map(|(self_segment, other_segment)| {
+                Some(self_segment.cmp(other_segment))
+                    .filter(|&cmp| cmp != std::cmp::Ordering::Equal)
+            })
+            .unwrap_or(self.inner.len().cmp(&other.inner.len()))
+    }
+}
+
 #[derive(Debug, Clone, Hash, PartialEq, Eq)]
 pub struct PortAndChannel {
     port: String,
@@ -372,6 +432,21 @@ impl PortAndChannel {
     }
 }
 
+impl PartialOrd for PortAndChannel {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for PortAndChannel {
+    /// Returns port comparison if not equal, otherwise returns channel comparison.
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.port
+            .cmp(&other.port)
+            .then_with(|| self.channel.cmp(&other.channel))
+    }
+}
+
 impl std::fmt::Display for TracePrefixed {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         for segment in &self.trace.inner {
@@ -549,6 +624,18 @@ impl FromStr for IbcPrefixed {
     }
 }
 
+impl PartialOrd for IbcPrefixed {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for IbcPrefixed {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.id.cmp(&other.id)
+    }
+}
+
 #[cfg(feature = "serde")]
 mod serde_impl {
     use serde::{

crates/astria-sequencer/src/accounts/query.rs

@@ -77,6 +77,9 @@ async fn get_trace_prefixed_account_balances<S: StateRead>(
     stream.try_collect::<Vec<_>>().await
 }
 
+/// Returns a list of `AssetBalance`s for the provided address. `AssetBalance`s are sorted first in
+/// descending order by balance, then in ascending order by denom. IBC prefixed denoms are treated
+/// as less than trace prefixed denoms.
 pub(crate) async fn balance_request(
     storage: Storage,
     request: request::Query,
@@ -88,7 +91,7 @@ pub(crate) async fn balance_request(
         Err(err_rsp) => return err_rsp,
     };
 
-    let balances = match get_trace_prefixed_account_balances(&snapshot, address).await {
+    let mut balances = match get_trace_prefixed_account_balances(&snapshot, address).await {
         Ok(balance) => balance,
         Err(err) => {
             return response::Query {
@@ -100,6 +103,11 @@ pub(crate) async fn balance_request(
             };
         }
     };
+
+    // Unstable sort does not allocate auxillary memory and is typically faster. Custom sorting
+    // function is deterministic and only allocates in the case that two values are equal.
+    balances.sort_unstable_by(compare_asset_balances);
+
     let payload = BalanceResponse {
         height: height.value(),
         balances,
@@ -215,3 +223,84 @@ async fn preprocess_request(
     };
     Ok((address, snapshot, height))
 }
+
+/// Custom deterministic sorting function for `AssetBalance` that sorts by balance in descending
+/// order and then by denom in ascending order. This function should never return `Ordering::Equal`,
+/// as calling `sort_unstable_by` on equal values has a non-deterministic result.
+fn compare_asset_balances(lhs: &AssetBalance, rhs: &AssetBalance) -> std::cmp::Ordering {
+    use std::cmp::Ordering;
+
+    match lhs.balance.cmp(&rhs.balance) {
+        Ordering::Equal => lhs.denom.cmp(&rhs.denom),
+        Ordering::Less => Ordering::Greater,
+        Ordering::Greater => Ordering::Less,
+    }
+}
+
+#[cfg(test)]
+mod test {
+    #[test]
+    fn compare_asset_balances_sorts_asset_balances_as_expected() {
+        use astria_core::{
+            primitive::v1::asset::Denom,
+            protocol::account::v1alpha1::AssetBalance,
+        };
+
+        let trace_denom_1 = "first/new/asset".parse::<Denom>().unwrap();
+        let trace_denom_2 = "second/new/asset".parse::<Denom>().unwrap();
+        let ibc_denom_1 = format!("ibc/{}", hex::encode([0u8; 32]))
+            .parse::<Denom>()
+            .unwrap();
+        let ibc_denom_2 = format!("ibc/{}", hex::encode([1u8; 32]))
+            .parse::<Denom>()
+            .unwrap();
+
+        let mut balances = vec![
+            AssetBalance {
+                denom: trace_denom_1.clone(),
+                balance: 1,
+            },
+            AssetBalance {
+                denom: trace_denom_1.clone(),
+                balance: 2,
+            },
+            AssetBalance {
+                denom: trace_denom_2.clone(),
+                balance: 2,
+            },
+            AssetBalance {
+                denom: trace_denom_2.clone(),
+                balance: 3,
+            },
+            AssetBalance {
+                denom: ibc_denom_1.clone(),
+                balance: 3,
+            },
+            AssetBalance {
+                denom: ibc_denom_1.clone(),
+                balance: 4,
+            },
+            AssetBalance {
+                denom: ibc_denom_2.clone(),
+                balance: 4,
+            },
+            AssetBalance {
+                denom: ibc_denom_2.clone(),
+                balance: 5,
+            },
+        ];
+
+        let original_balances = balances.clone();
+
+        balances.sort_unstable_by(super::compare_asset_balances);
+
+        assert_eq!(balances[0], original_balances[7]);
+        assert_eq!(balances[1], original_balances[5]);
+        assert_eq!(balances[2], original_balances[6]);
+        assert_eq!(balances[3], original_balances[4]);
+        assert_eq!(balances[4], original_balances[3]);
+        assert_eq!(balances[5], original_balances[1]);
+        assert_eq!(balances[6], original_balances[2]);
+        assert_eq!(balances[7], original_balances[0]);
+    }
+}