From fb7110f1f824b3eb13602e4ffc0f3590b336b2e6 Mon Sep 17 00:00:00 2001
From: ethanoroshiba <ethan@astria.org>
Date: Fri, 27 Sep 2024 12:10:49 -0500
Subject: [PATCH] Requested changes

---
 .../src/primitive/v1/asset/denom.rs           | 87 +++++++++++++++++++
 crates/astria-sequencer/src/accounts/query.rs |  7 +-
 2 files changed, 92 insertions(+), 2 deletions(-)

diff --git a/crates/astria-core/src/primitive/v1/asset/denom.rs b/crates/astria-core/src/primitive/v1/asset/denom.rs
index 91dd06bf24..03cb146499 100644
--- a/crates/astria-core/src/primitive/v1/asset/denom.rs
+++ b/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::{
diff --git a/crates/astria-sequencer/src/accounts/query.rs b/crates/astria-sequencer/src/accounts/query.rs
index e65223841a..75ce03afce 100644
--- a/crates/astria-sequencer/src/accounts/query.rs
+++ b/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,
@@ -228,13 +231,13 @@ fn compare_asset_balances(lhs: &AssetBalance, rhs: &AssetBalance) -> std::cmp::O
     use std::cmp::Ordering;
 
     match lhs.balance.cmp(&rhs.balance) {
-        // Denoms should never have the same display name
-        Ordering::Equal => lhs.denom.to_string().cmp(&rhs.denom.to_string()),
+        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() {