fix(db): add unlock_condition to id_index (#402)

* fix(db): add `unlock_condition` to `id_index` * Rename to `UnlockConditionType`
iotaledger · Jul 11, 2022 · e0145b3 · e0145b3
1 parent 250ccbf
commit e0145b3
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 15 deletions.
diff --git 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
@@ -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
@@ -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())),
         }
     }
 }