From bfc5a0a312a468f2cfd693b308dfad4de3f7ab13 Mon Sep 17 00:00:00 2001
From: Lucas <x@rvcas.dev>
Date: Fri, 16 Sep 2022 16:53:45 -0400
Subject: [PATCH] feat: Provide access to all assets at a tx out (#180)

---
 pallas-codec/src/utils.rs     | 10 +++++-
 pallas-traverse/src/lib.rs    | 27 +++++++++++++-
 pallas-traverse/src/output.rs | 67 +++++++++++++++++++++++++++++++++--
 3 files changed, 100 insertions(+), 4 deletions(-)

diff --git a/pallas-codec/src/utils.rs b/pallas-codec/src/utils.rs
index bd6d453..722cb08 100644
--- a/pallas-codec/src/utils.rs
+++ b/pallas-codec/src/utils.rs
@@ -1,6 +1,6 @@
 use minicbor::{data::Tag, Decode, Encode};
 use serde::{Deserialize, Serialize};
-use std::ops::Deref;
+use std::{fmt, ops::Deref};
 
 /// Utility for skipping parts of the CBOR payload, use only for debugging
 #[derive(Debug, PartialEq, PartialOrd, Eq, Ord)]
@@ -778,6 +778,14 @@ impl From<Bytes> for String {
     }
 }
 
+impl fmt::Display for Bytes {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let bytes: Vec<u8> = self.clone().into();
+
+        f.write_str(&hex::encode(bytes))
+    }
+}
+
 #[derive(
     Serialize, Deserialize, Clone, Copy, Encode, Decode, Debug, PartialEq, Eq, PartialOrd, Ord,
 )]
diff --git a/pallas-traverse/src/lib.rs b/pallas-traverse/src/lib.rs
index cc356df..bd87342 100644
--- a/pallas-traverse/src/lib.rs
+++ b/pallas-traverse/src/lib.rs
@@ -7,7 +7,11 @@ use thiserror::Error;
 
 use pallas_codec::utils::KeepRaw;
 use pallas_crypto::hash::Hash;
-use pallas_primitives::{alonzo, babbage, byron};
+use pallas_primitives::{
+    alonzo,
+    babbage::{self, AssetName, PolicyId},
+    byron,
+};
 
 pub mod block;
 pub mod cert;
@@ -132,6 +136,27 @@ pub enum MultiEraSigners<'b> {
 #[derive(Debug, Clone)]
 pub struct OutputRef(Hash<32>, u64);
 
+#[derive(Debug, Clone)]
+pub struct Asset {
+    pub subject: Subject,
+    pub quantity: u64,
+}
+
+#[derive(Debug, Clone)]
+pub enum Subject {
+    Lovelace,
+    NativeAsset(PolicyId, AssetName),
+}
+
+impl ToString for Subject {
+    fn to_string(&self) -> String {
+        match self {
+            Self::Lovelace => String::from("lovelace"),
+            Self::NativeAsset(p, n) => format!("{p}.{n}"),
+        }
+    }
+}
+
 #[derive(Debug, Error)]
 pub enum Error {
     #[error("Invalid CBOR structure: {0}")]
diff --git a/pallas-traverse/src/output.rs b/pallas-traverse/src/output.rs
index 51a239c..483d58c 100644
--- a/pallas-traverse/src/output.rs
+++ b/pallas-traverse/src/output.rs
@@ -4,11 +4,11 @@ use pallas_addresses::{Address, ByronAddress, Error as AddressError};
 use pallas_codec::minicbor;
 use pallas_primitives::{
     alonzo,
-    babbage::{self, DatumOption, ScriptRef},
+    babbage::{self, Coin, DatumOption, ScriptRef},
     byron,
 };
 
-use crate::{Era, MultiEraOutput};
+use crate::{Asset, Era, MultiEraOutput, Subject};
 
 impl<'b> MultiEraOutput<'b> {
     pub fn from_byron(output: &'b byron::TxOut) -> Self {
@@ -77,6 +77,50 @@ impl<'b> MultiEraOutput<'b> {
         }
     }
 
+    pub fn assets(&self) -> Vec<Asset> {
+        let mut assets = Vec::new();
+
+        match self {
+            MultiEraOutput::Byron(x) => {
+                push_lovelace(&mut assets, x.amount);
+            }
+            MultiEraOutput::Babbage(x) => match x.deref().deref() {
+                babbage::TransactionOutput::Legacy(x) => match &x.amount {
+                    babbage::Value::Coin(c) => {
+                        push_lovelace(&mut assets, *c);
+                    }
+                    babbage::Value::Multiasset(c, multi_asset) => {
+                        push_lovelace(&mut assets, *c);
+
+                        push_native_asset(&mut assets, multi_asset);
+                    }
+                },
+                babbage::TransactionOutput::PostAlonzo(x) => match &x.value {
+                    babbage::Value::Coin(c) => {
+                        push_lovelace(&mut assets, *c);
+                    }
+                    babbage::Value::Multiasset(c, multi_asset) => {
+                        push_lovelace(&mut assets, *c);
+
+                        push_native_asset(&mut assets, multi_asset);
+                    }
+                },
+            },
+            MultiEraOutput::AlonzoCompatible(x) => match &x.amount {
+                alonzo::Value::Coin(c) => {
+                    push_lovelace(&mut assets, *c);
+                }
+                alonzo::Value::Multiasset(c, multi_asset) => {
+                    push_lovelace(&mut assets, *c);
+
+                    push_native_asset(&mut assets, multi_asset);
+                }
+            },
+        };
+
+        assets
+    }
+
     pub fn as_babbage(&self) -> Option<&babbage::TransactionOutput> {
         match self {
             MultiEraOutput::AlonzoCompatible(_) => None,
@@ -131,6 +175,24 @@ impl<'b> MultiEraOutput<'b> {
     }
 }
 
+fn push_lovelace(assets: &mut Vec<Asset>, quantity: u64) {
+    assets.push(Asset {
+        subject: Subject::Lovelace,
+        quantity,
+    })
+}
+
+fn push_native_asset(assets: &mut Vec<Asset>, multi_asset: &alonzo::Multiasset<Coin>) {
+    for (policy_id, names) in multi_asset.iter() {
+        for (asset_name, quantity) in names.iter() {
+            assets.push(Asset {
+                subject: Subject::NativeAsset(*policy_id, asset_name.clone()),
+                quantity: *quantity,
+            });
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::MultiEraBlock;
@@ -143,6 +205,7 @@ mod tests {
 
         for tx in block.txs() {
             for output in tx.outputs() {
+                assert_ne!(output.assets()[0].quantity, 0);
                 assert_ne!(output.ada_amount(), 0);
                 assert!(matches!(output.address(), Ok(_)));
             }