feat(analytics): consider expiring and locked outputs

src/analytics/ledger/active_addresses.rs

@@ -1,4 +1,4 @@
-// Copyright 2023 IOTA Stiftung
+// Copyright 2024 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
 use std::collections::HashSet;
@@ -43,15 +43,15 @@ impl IntervalAnalytics for AddressActivityMeasurement {
 impl Analytics for AddressActivityAnalytics {
     type Measurement = AddressActivityMeasurement;
 
-    fn handle_transaction(&mut self, consumed: &[LedgerSpent], created: &[LedgerOutput], ctx: &dyn AnalyticsContext) {
+    fn handle_transaction(&mut self, consumed: &[LedgerSpent], created: &[LedgerOutput], _ctx: &dyn AnalyticsContext) {
         for output in consumed {
             if let Some(a) = output.owning_address() {
                 self.addresses.insert(*a);
             }
         }
 
         for output in created {
-            if let Some(a) = output.output.owning_address(ctx.at().milestone_timestamp) {
+            if let Some(a) = output.owning_address() {
                 self.addresses.insert(*a);
             }
         }

src/analytics/ledger/address_balance.rs

@@ -1,11 +1,11 @@
-// Copyright 2023 IOTA Stiftung
+// Copyright 2024 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 
 use super::*;
 use crate::model::{
-    payload::milestone::MilestoneTimestamp,
+    payload::{milestone::MilestoneTimestamp, transaction::output::OutputId},
     utxo::{Address, TokenAmount},
 };
 
@@ -24,10 +24,12 @@ pub(crate) struct DistributionStat {
     pub(crate) total_amount: TokenAmount,
 }
 
-/// Computes the number of addresses the currently hold a balance.
-#[derive(Serialize, Deserialize)]
+/// Computes the number of addresses that currently hold a balance.
+#[derive(Serialize, Deserialize, Default)]
 pub(crate) struct AddressBalancesAnalytics {
     balances: HashMap<Address, TokenAmount>,
+    expiring: BTreeMap<(MilestoneTimestamp, OutputId), (Address, Address, TokenAmount)>,
+    locked: BTreeMap<(MilestoneTimestamp, OutputId), (Address, TokenAmount)>,
 }
 
 impl AddressBalancesAnalytics {
@@ -36,49 +38,118 @@ impl AddressBalancesAnalytics {
         unspent_outputs: impl IntoIterator<Item = &'a LedgerOutput>,
         milestone_timestamp: MilestoneTimestamp,
     ) -> Self {
-        let mut balances = HashMap::new();
-        for output in unspent_outputs {
-            if let Some(&a) = output.output.owning_address(milestone_timestamp) {
-                *balances.entry(a).or_default() += output.amount();
+        let mut res = AddressBalancesAnalytics::default();
+        for created in unspent_outputs {
+            res.handle_created(created, milestone_timestamp);
+        }
+        res
+    }
+
+    fn handle_created(&mut self, created: &LedgerOutput, milestone_timestamp: MilestoneTimestamp) {
+        if let Some(&owning_address) = created.owning_address() {
+            if let Some(expiration) = created.output.expiration() {
+                // If the output is expired already, add the value to the return address
+                if milestone_timestamp >= expiration.timestamp {
+                    *self.balances.entry(expiration.return_address).or_default() += created.amount();
+                } else {
+                    // Otherwise, add it to the set of expiring values to be handled later
+                    *self.balances.entry(owning_address).or_default() += created.amount();
+                    self.expiring.insert(
+                        (expiration.timestamp, created.output_id),
+                        (owning_address, expiration.return_address, created.amount()),
+                    );
+                }
+            } else if let Some(timelock) = created.output.timelock() {
+                // If the output is unlocked, add the value to the address
+                if milestone_timestamp >= timelock.timestamp {
+                    *self.balances.entry(owning_address).or_default() += created.amount();
+                } else {
+                    // Otherwise, add it to the set of locked values to be handled later
+                    self.locked.insert(
+                        (timelock.timestamp, created.output_id),
+                        (owning_address, created.amount()),
+                    );
+                }
+            } else {
+                *self.balances.entry(owning_address).or_default() += created.amount();
+            }
+        }
+    }
+
+    fn handle_consumed(&mut self, consumed: &LedgerSpent, milestone_timestamp: MilestoneTimestamp) {
+        if let Some(&owning_address) = consumed.output.owning_address() {
+            if let Some(expiration) = consumed.output.output.expiration() {
+                // No longer need to handle the expiration
+                self.expiring.remove(&(expiration.timestamp, consumed.output_id()));
+                // If the output is past the expiration time, remove the value from the return address
+                if milestone_timestamp >= expiration.timestamp {
+                    *self.balances.entry(expiration.return_address).or_default() -= consumed.amount();
+                // Otherwise, remove it from the original address
+                } else {
+                    *self.balances.entry(owning_address).or_default() -= consumed.amount();
+                }
+            } else if let Some(timelock) = consumed.output.output.timelock() {
+                // No longer need to handle the lock
+                self.locked.remove(&(timelock.timestamp, consumed.output_id()));
+                *self.balances.entry(owning_address).or_default() -= consumed.amount();
+            } else {
+                *self.balances.entry(owning_address).or_default() -= consumed.amount();
+            }
+        }
+    }
+
+    fn handle_expired(&mut self, milestone_timestamp: MilestoneTimestamp) {
+        while let Some((address, return_address, amount)) = self.expiring.first_entry().and_then(|entry| {
+            if milestone_timestamp >= entry.key().0 {
+                Some(entry.remove())
+            } else {
+                None
             }
+        }) {
+            *self.balances.entry(address).or_default() -= amount;
+            *self.balances.entry(return_address).or_default() += amount;
+        }
+    }
+
+    fn handle_locked(&mut self, milestone_timestamp: MilestoneTimestamp) {
+        while let Some((address, amount)) = self.locked.first_entry().and_then(|entry| {
+            if milestone_timestamp >= entry.key().0 {
+                Some(entry.remove())
+            } else {
+                None
+            }
+        }) {
+            *self.balances.entry(address).or_default() += amount;
         }
-        Self { balances }
     }
 }
 
 impl Analytics for AddressBalancesAnalytics {
     type Measurement = AddressBalanceMeasurement;
 
     fn handle_transaction(&mut self, consumed: &[LedgerSpent], created: &[LedgerOutput], ctx: &dyn AnalyticsContext) {
-        for output in consumed {
-            if let Some(a) = output.owning_address() {
-                // All inputs should be present in `addresses`. If not, we skip it's value.
-                if let Some(amount) = self.balances.get_mut(a) {
-                    *amount -= output.amount();
-                    if amount.0 == 0 {
-                        self.balances.remove(a);
-                    }
-                }
-            }
+        // Handle consumed outputs first, as they can remove entries from expiration/locked
+        for consumed in consumed {
+            self.handle_consumed(consumed, ctx.at().milestone_timestamp);
         }
-
-        for output in created {
-            if let Some(&a) = output.output.owning_address(ctx.at().milestone_timestamp) {
-                // All inputs should be present in `addresses`. If not, we skip it's value.
-                *self.balances.entry(a).or_default() += output.amount();
-            }
+        // Handle any expiring or unlocking outputs for this milestone
+        self.handle_expired(ctx.at().milestone_timestamp);
+        self.handle_locked(ctx.at().milestone_timestamp);
+        // Finally, handle the created transactions, which can insert new expiration/locked records for the future
+        for created in created {
+            self.handle_created(created, ctx.at().milestone_timestamp);
         }
     }
 
     fn take_measurement(&mut self, ctx: &dyn AnalyticsContext) -> Self::Measurement {
         let bucket_max = ctx.protocol_params().token_supply.ilog10() as usize + 1;
         let mut token_distribution = vec![DistributionStat::default(); bucket_max];
 
-        for amount in self.balances.values() {
+        for balance in self.balances.values() {
             // Balances are partitioned into ranges defined by: [10^index..10^(index+1)).
-            let index = amount.0.ilog10() as usize;
+            let index = balance.ilog10() as usize;
             token_distribution[index].address_count += 1;
-            token_distribution[index].total_amount += *amount;
+            token_distribution[index].total_amount += *balance;
         }
         AddressBalanceMeasurement {
             address_with_balance_count: self.balances.len(),

src/bin/inx-chronicle/cli/analytics.rs

@@ -1,4 +1,4 @@
-// Copyright 2023 IOTA Stiftung
+// Copyright 2024 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
 use std::collections::HashSet;
@@ -23,7 +23,9 @@ use tracing::{debug, info};
 
 use crate::config::ChronicleConfig;
 
-/// This command accepts both milestone index and date ranges. The following rules apply:
+/// This command accepts both milestone index and date ranges.
+///
+/// The following rules apply:
 ///
 /// - If both milestone and date are specified, the date will be used for interval analytics
 /// while the milestone will be used for per-milestone analytics.

src/model/block/payload/transaction/mod.rs

@@ -1,4 +1,4 @@
-// Copyright 2022 IOTA Stiftung
+// Copyright 2024 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
 //! Module containing types related to transactions.
@@ -17,7 +17,7 @@ pub mod output;
 pub mod unlock;
 
 /// Uniquely identifies a transaction.
-#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Hash)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Hash, PartialOrd, Ord)]
 #[serde(transparent)]
 pub struct TransactionId(#[serde(with = "bytify")] pub [u8; Self::LENGTH]);
 

src/model/block/payload/transaction/output/mod.rs

@@ -1,4 +1,4 @@
-// Copyright 2022 IOTA Stiftung
+// Copyright 2024 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
 //! Module containing the [`Output`] types.
@@ -23,6 +23,7 @@ use mongodb::bson::{doc, Bson};
 use packable::PackableExt;
 use serde::{Deserialize, Serialize};
 
+use self::unlock_condition::{ExpirationUnlockCondition, TimelockUnlockCondition};
 pub use self::{
     address::{Address, AliasAddress, Ed25519Address, NftAddress},
     alias::{AliasId, AliasOutput},
@@ -54,15 +55,17 @@ use crate::model::{
     derive_more::AddAssign,
     derive_more::SubAssign,
     derive_more::Sum,
+    derive_more::Deref,
 )]
+#[repr(transparent)]
 pub struct TokenAmount(#[serde(with = "stringify")] pub u64);
 
 /// The index of an output within a transaction.
 pub type OutputIndex = u16;
 
 /// An id which uniquely identifies an output. It is computed from the corresponding [`TransactionId`], as well as the
 /// [`OutputIndex`].
-#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Hash)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Hash, PartialOrd, Ord)]
 pub struct OutputId {
     /// The transaction id part of the [`OutputId`].
     pub transaction_id: TransactionId,
@@ -182,6 +185,36 @@ impl Output {
         })
     }
 
+    /// Returns the expiration unlock condition, if there is one.
+    pub fn expiration(&self) -> Option<&ExpirationUnlockCondition> {
+        match self {
+            Self::Basic(BasicOutput {
+                expiration_unlock_condition,
+                ..
+            })
+            | Self::Nft(NftOutput {
+                expiration_unlock_condition,
+                ..
+            }) => expiration_unlock_condition.as_ref(),
+            _ => None,
+        }
+    }
+
+    /// Returns the timelock unlock condition, if there is one.
+    pub fn timelock(&self) -> Option<&TimelockUnlockCondition> {
+        match self {
+            Self::Basic(BasicOutput {
+                timelock_unlock_condition,
+                ..
+            })
+            | Self::Nft(NftOutput {
+                timelock_unlock_condition,
+                ..
+            }) => timelock_unlock_condition.as_ref(),
+            _ => None,
+        }
+    }
+
     /// Returns the amount associated with an output.
     pub fn amount(&self) -> TokenAmount {
         match self {

src/model/block/payload/transaction/output/unlock_condition/timelock.rs

@@ -1,4 +1,4 @@
-// Copyright 2022 IOTA Stiftung
+// Copyright 2024 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
 #![allow(missing_docs)]
@@ -13,7 +13,7 @@ use crate::model::tangle::MilestoneTimestamp;
 /// Defines a unix timestamp until which the output can not be unlocked.
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct TimelockUnlockCondition {
-    timestamp: MilestoneTimestamp,
+    pub(crate) timestamp: MilestoneTimestamp,
 }
 
 impl<T: Borrow<iota::TimelockUnlockCondition>> From<T> for TimelockUnlockCondition {