From e0145b376ee12cdae792af62283e9c2e669804d7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jochen=20G=C3=B6rtler?= <grtlr@users.noreply.github.com>
Date: Mon, 11 Jul 2022 08:34:23 +0200
Subject: [PATCH] fix(db): add `unlock_condition` to `id_index` (#402)

* fix(db): add `unlock_condition` to `id_index`

* Rename to `UnlockConditionType`
---
 src/db/collections/ledger_update.rs           | 13 ++++--
 src/types/stardust/block/output/mod.rs        |  5 +--
 .../stardust/block/output/unlock_condition.rs | 44 +++++++++++++++----
 3 files changed, 47 insertions(+), 15 deletions(-)

diff --git a/src/db/collections/ledger_update.rs b/src/db/collections/ledger_update.rs
index f1656d7a7..f597baa6d 100644
--- a/src/db/collections/ledger_update.rs
+++ b/src/db/collections/ledger_update.rs
@@ -14,7 +14,7 @@ use crate::{
     db::MongoDb,
     types::{
         ledger::{MilestoneIndexTimestamp, OutputWithMetadata},
-        stardust::block::{Address, OutputId},
+        stardust::block::{Address, OutputId, UnlockConditionType},
         tangle::MilestoneIndex,
     },
 };
@@ -25,6 +25,9 @@ struct LedgerUpdateDocument {
     address: Address,
     output_id: OutputId,
     at: MilestoneIndexTimestamp,
+    #[serde(with = "crate::types::util::stringify")]
+    amount: u64,
+    unlock_condition_type: UnlockConditionType,
     is_spent: bool,
 }
 
@@ -111,7 +114,7 @@ impl MongoDb {
         collection
             .create_index(
                 IndexModel::builder()
-                    .keys(doc! { "output_id": 1, "is_spent": 1 })
+                    .keys(doc! { "output_id": 1, "is_spent": 1, "unlock_condition": 1 })
                     .options(
                         IndexOptions::builder()
                             // An output can be spent and unspent only once.
@@ -140,12 +143,14 @@ impl MongoDb {
             let is_spent = metadata.spent.is_some();
 
             // Ledger updates
-            for owner in output.owning_addresses() {
+            for (address, unlock_condition) in output.owning_addresses() {
                 let ledger_update_document = LedgerUpdateDocument {
-                    address: owner,
+                    address,
                     output_id: metadata.output_id,
                     at,
                     is_spent,
+                    amount: output.amount(),
+                    unlock_condition_type: unlock_condition,
                 };
 
                 let _ = self
diff --git a/src/types/stardust/block/output/mod.rs b/src/types/stardust/block/output/mod.rs
index 1a10f4a18..a85f339bb 100644
--- a/src/types/stardust/block/output/mod.rs
+++ b/src/types/stardust/block/output/mod.rs
@@ -26,7 +26,7 @@ pub use self::{
     native_token::{NativeToken, TokenScheme},
     nft::{NftId, NftOutput},
     treasury::TreasuryOutput,
-    unlock_condition::UnlockCondition,
+    unlock_condition::{UnlockCondition, UnlockConditionType},
 };
 use super::Address;
 use crate::types::stardust::block::TransactionId;
@@ -89,7 +89,7 @@ pub enum Output {
 }
 
 impl Output {
-    pub fn owning_addresses(&self) -> Vec<Address> {
+    pub fn owning_addresses(&self) -> Vec<(Address, UnlockConditionType)> {
         match self {
             Self::Treasury(_) => Vec::new(),
             Self::Basic(BasicOutput { unlock_conditions, .. })
@@ -98,7 +98,6 @@ impl Output {
             | Self::Foundry(FoundryOutput { unlock_conditions, .. }) => unlock_conditions
                 .iter()
                 .filter_map(UnlockCondition::owning_address)
-                .cloned()
                 .collect(),
         }
     }
diff --git a/src/types/stardust/block/output/unlock_condition.rs b/src/types/stardust/block/output/unlock_condition.rs
index 954032efe..beaa9994a 100644
--- a/src/types/stardust/block/output/unlock_condition.rs
+++ b/src/types/stardust/block/output/unlock_condition.rs
@@ -35,16 +35,44 @@ pub enum UnlockCondition {
     },
 }
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case", tag = "kind")]
+pub enum UnlockConditionType {
+    Address,
+    StorageDepositReturn { amount: u64 },
+    Timelock { timestamp: u32 },
+    Expiration { timestamp: u32 },
+    StateControllerAddress,
+    GovernorAddress,
+    ImmutableAliasAddress,
+}
+
+impl From<&UnlockCondition> for UnlockConditionType {
+    fn from(value: &UnlockCondition) -> Self {
+        match *value {
+            UnlockCondition::Address { .. } => UnlockConditionType::Address,
+            UnlockCondition::StorageDepositReturn { amount, .. } => {
+                UnlockConditionType::StorageDepositReturn { amount }
+            }
+            UnlockCondition::Timelock { timestamp } => UnlockConditionType::Timelock { timestamp },
+            UnlockCondition::Expiration { timestamp, .. } => UnlockConditionType::Expiration { timestamp },
+            UnlockCondition::StateControllerAddress { .. } => UnlockConditionType::StateControllerAddress,
+            UnlockCondition::GovernorAddress { .. } => UnlockConditionType::GovernorAddress,
+            UnlockCondition::ImmutableAliasAddress { .. } => UnlockConditionType::ImmutableAliasAddress,
+        }
+    }
+}
+
 impl UnlockCondition {
-    pub fn owning_address(&self) -> Option<&Address> {
-        match self {
-            Self::Address { address } => Some(address),
-            Self::StorageDepositReturn { return_address, .. } => Some(return_address),
+    pub fn owning_address(&self) -> Option<(Address, UnlockConditionType)> {
+        match *self {
+            Self::Address { address } => Some((address, self.into())),
+            Self::StorageDepositReturn { return_address, .. } => Some((return_address, self.into())),
             Self::Timelock { .. } => None,
-            Self::Expiration { return_address, .. } => Some(return_address),
-            Self::StateControllerAddress { address } => Some(address),
-            Self::GovernorAddress { address } => Some(address),
-            Self::ImmutableAliasAddress { address } => Some(address),
+            Self::Expiration { return_address, .. } => Some((return_address, self.into())),
+            Self::StateControllerAddress { address } => Some((address, self.into())),
+            Self::GovernorAddress { address } => Some((address, self.into())),
+            Self::ImmutableAliasAddress { address } => Some((address, self.into())),
         }
     }
 }