iota_mainnet_unlocks/
store.rs

1// Copyright (c) 2025 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::BTreeMap;
5
6use anyhow::{Context, Result};
7use chrono::{DateTime, Utc};
8use csv::ReaderBuilder;
9use serde::{Deserialize, Serialize};
10
11use crate::aggregated_data::AGGREGATED_DATA_CSV;
12
13/// File name of the mainnet unlock data.
14pub const INPUT_FILE: &str = "aggregated_mainnet_unlocks.csv";
15
16/// Represents a single entry in the store.
17/// It defines how many tokens still remain locked at a specific point in time.
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct StillLockedEntry {
20    /// UTC timestamp at which the tokens are still locked.
21    pub timestamp: DateTime<Utc>,
22    /// Total locked amount (nano-units) still locked at the timestamp.
23    pub amount_still_locked: u64,
24}
25
26/// In-memory store holding the aggregated token unlock data.
27#[derive(Debug, Clone)]
28pub struct MainnetUnlocksStore {
29    // Each entry represents the total number of tokens still locked at the specific point in time.
30    entries: BTreeMap<DateTime<Utc>, StillLockedEntry>,
31}
32
33impl MainnetUnlocksStore {
34    /// Creates a new store with the aggregated unlock data for mainnet.
35    pub fn new() -> Result<Self> {
36        Self::from_csv_str(AGGREGATED_DATA_CSV)
37    }
38
39    /// Parses the given CSV string into a `MainnetUnlocksStore`.
40    fn from_csv_str(csv_str: &str) -> Result<Self> {
41        let mut rdr = ReaderBuilder::new()
42            .has_headers(true)
43            .from_reader(csv_str.as_bytes());
44
45        let mut map = BTreeMap::new();
46
47        for result in rdr.deserialize() {
48            let entry: StillLockedEntry =
49                result.context("failed to deserialize CSV row into StillLockedEntry")?;
50
51            if let Some(old_entry) = map.insert(entry.timestamp, entry) {
52                return Err(anyhow::anyhow!(
53                    "duplicate entry found for timestamp: {}",
54                    old_entry.timestamp
55                ));
56            }
57        }
58
59        Ok(Self { entries: map })
60    }
61
62    /// Returns the total amount of tokens (in nano-units) that are still locked
63    /// at the given timestamp.
64    pub fn still_locked_tokens(&self, date_time: DateTime<Utc>) -> u64 {
65        self.entries
66            .range(..=date_time)
67            .next_back()
68            .map(|(_, entry)| entry.amount_still_locked)
69            .unwrap_or_else(|| {
70                // No earlier entries exist: use first available as retroactively valid
71                self.entries
72                    .iter()
73                    .next()
74                    .map(|(_, e)| e.amount_still_locked)
75                    .unwrap_or(0)
76            })
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use chrono::{TimeZone, Utc};
83
84    use super::*;
85
86    #[test]
87    fn new() {
88        let store = MainnetUnlocksStore::new().unwrap();
89        assert_eq!(store.entries.len(), 105);
90    }
91
92    fn store(csv: &str) -> MainnetUnlocksStore {
93        MainnetUnlocksStore::from_csv_str(csv).unwrap()
94    }
95
96    #[test]
97    fn test_no_entries() {
98        let store = store("timestamp,amount_still_locked\n");
99        assert_eq!(store.still_locked_tokens(Utc::now()), 0);
100    }
101
102    #[test]
103    fn test_single_entry() {
104        let csv = r#"
105            timestamp,amount_still_locked
106            2000-01-01T00:00:00Z,999
107        "#;
108        let store = store(csv.trim());
109
110        let before = Utc.with_ymd_and_hms(1999, 12, 31, 23, 59, 59).unwrap();
111        let exact = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
112        let after = Utc.with_ymd_and_hms(2001, 1, 1, 0, 0, 0).unwrap();
113
114        assert_eq!(store.still_locked_tokens(before), 999);
115        assert_eq!(store.still_locked_tokens(exact), 999);
116        assert_eq!(store.still_locked_tokens(after), 999);
117    }
118
119    #[test]
120    fn test_multiple_entries() {
121        let csv = r#"
122            timestamp,amount_still_locked
123            2023-01-01T00:00:00Z,300
124            2024-01-01T00:00:00Z,200
125            2025-01-01T00:00:00Z,100
126        "#;
127        let store = store(csv.trim());
128
129        let t0 = Utc.with_ymd_and_hms(2022, 12, 31, 0, 0, 0).unwrap();
130        let t0_between = Utc.with_ymd_and_hms(2023, 6, 1, 0, 0, 0).unwrap();
131        let t1 = Utc.with_ymd_and_hms(2023, 1, 1, 0, 0, 0).unwrap();
132        let t2 = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
133        let t3 = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
134        let t4 = Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap();
135
136        assert_eq!(store.still_locked_tokens(t0), 300);
137        assert_eq!(store.still_locked_tokens(t1), 300);
138        assert_eq!(store.still_locked_tokens(t0_between), 300);
139        assert_eq!(store.still_locked_tokens(t2), 200);
140        assert_eq!(store.still_locked_tokens(t3), 100);
141        assert_eq!(store.still_locked_tokens(t4), 100);
142    }
143
144    #[test]
145    fn test_zero_at_latest_entry() {
146        let csv = r#"
147            timestamp,amount_still_locked
148            2023-01-01T00:00:00Z,1000
149            2025-01-01T00:00:00Z,0
150        "#;
151        let store = store(csv.trim());
152
153        let t1 = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
154        let t2 = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
155        let t3 = Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap();
156
157        assert_eq!(store.still_locked_tokens(t1), 1000);
158        assert_eq!(store.still_locked_tokens(t2), 0);
159        assert_eq!(store.still_locked_tokens(t3), 0);
160    }
161
162    #[test]
163    fn test_gap_between_entries() {
164        let csv = r#"
165            timestamp,amount_still_locked
166            2020-01-01T00:00:00Z,1000
167            2030-01-01T00:00:00Z,100
168        "#;
169        let store = store(csv.trim());
170
171        let t_before = Utc.with_ymd_and_hms(2019, 1, 1, 0, 0, 0).unwrap();
172        let t_mid = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
173        let t_after = Utc.with_ymd_and_hms(2040, 1, 1, 0, 0, 0).unwrap();
174
175        assert_eq!(store.still_locked_tokens(t_before), 1000);
176        assert_eq!(store.still_locked_tokens(t_mid), 1000);
177        assert_eq!(store.still_locked_tokens(t_after), 100);
178    }
179
180    #[test]
181    fn test_dense_entry() {
182        let csv = r#"
183            timestamp,amount_still_locked
184            2023-10-01T00:00:00Z,300
185            2023-10-15T00:00:00Z,200
186            2023-11-01T00:00:00Z,100
187        "#;
188        let store = store(csv.trim());
189
190        let t_exact_mid = Utc.with_ymd_and_hms(2023, 10, 15, 0, 0, 0).unwrap();
191        let t_between = Utc.with_ymd_and_hms(2023, 10, 20, 0, 0, 0).unwrap();
192
193        assert_eq!(store.still_locked_tokens(t_exact_mid), 200);
194        assert_eq!(store.still_locked_tokens(t_between), 200);
195    }
196
197    #[test]
198    fn test_first_entry_is_retrospective() {
199        let csv = r#"
200            timestamp,amount_still_locked
201            2022-06-01T00:00:00Z,888
202        "#;
203        let store = store(csv.trim());
204
205        let far_before = Utc.with_ymd_and_hms(2000, 1, 1, 0, 0, 0).unwrap();
206        let just_before = Utc.with_ymd_and_hms(2022, 5, 31, 23, 59, 59).unwrap();
207        let exact = Utc.with_ymd_and_hms(2022, 6, 1, 0, 0, 0).unwrap();
208        let just_after = Utc.with_ymd_and_hms(2022, 6, 1, 0, 0, 1).unwrap();
209
210        assert_eq!(store.still_locked_tokens(far_before), 888);
211        assert_eq!(store.still_locked_tokens(just_before), 888);
212        assert_eq!(store.still_locked_tokens(exact), 888);
213        assert_eq!(store.still_locked_tokens(just_after), 888);
214    }
215
216    #[test]
217    fn test_unsorted_input() {
218        let csv = r#"
219            timestamp,amount_still_locked
220            2023-11-01T00:00:00Z,100
221            2023-01-01T00:00:00Z,300
222            2023-10-01T00:00:00Z,200
223        "#;
224        let store = store(csv.trim());
225
226        let query_before = Utc.with_ymd_and_hms(2023, 6, 1, 0, 0, 0).unwrap();
227        assert_eq!(store.still_locked_tokens(query_before), 300);
228
229        let query_exact = Utc.with_ymd_and_hms(2023, 10, 1, 0, 0, 0).unwrap();
230        assert_eq!(store.still_locked_tokens(query_exact), 200);
231    }
232}