generate_aggregated_data/
generate_aggregated_data.rs

1// Copyright (c) 2025 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{
5    collections::BTreeMap,
6    path::{Path, PathBuf},
7    process::Command,
8};
9
10use anyhow::{Context, Result};
11use chrono::{DateTime, Utc};
12use csv::{Reader, Writer};
13use iota_mainnet_unlocks::store::{INPUT_FILE as OUTPUT_FILE, StillLockedEntry};
14use regex::Regex;
15use tempfile::{TempDir, tempdir};
16
17// Folders in the raw data repository that contain the CSV files.
18const FOLDERS: &[&str] = &[
19    "Assembly_IF_Members",
20    "Assembly_Investors",
21    "IOTA_Airdrop",
22    "IOTA_Foundation",
23    "New_Investors",
24    "TEA",
25    "Treasury_DAO",
26    "UAE",
27];
28
29/// Clones the repository containing raw data into a temporary directory.
30fn clone_repo(tmp_dir: &TempDir) -> Result<PathBuf> {
31    let repo_path = tmp_dir.path().join("new_supply");
32
33    let status = Command::new("git")
34        .args([
35            "clone",
36            "--depth",
37            "1",
38            "https://github.com/iotaledger/new_supply.git",
39        ])
40        .arg(&repo_path)
41        .status()
42        .context("failed to execute `git clone`")?;
43
44    if !status.success() {
45        anyhow::bail!("`git clone` failed with exit status: {}", status);
46    }
47
48    Ok(repo_path)
49}
50
51/// Reads and aggregates the CSV unlock data from the cloned repository.
52/// Returns a BTreeMap keyed by unlock date (as a String) with the aggregated
53/// token amount (in nano-units).
54fn aggregate_unlocks(repo_path: &Path) -> Result<BTreeMap<String, u64>> {
55    let mut locked_by_date: BTreeMap<String, u64> = BTreeMap::new();
56
57    for folder in FOLDERS {
58        let csv_path = repo_path.join(folder).join("summary.csv");
59        println!("Processing file: {:?}", csv_path);
60
61        let mut rdr = Reader::from_path(&csv_path)
62            .with_context(|| format!("failed to open CSV file: {csv_path:?}"))?;
63
64        // Iterate over CSV records (header is skipped automatically).
65        for result in rdr.records() {
66            let record = result?;
67            if record.len() < 2 {
68                return Err(anyhow::anyhow!("invalid record: {record:?}"));
69            }
70
71            let tokens_str = record.get(0).unwrap().trim();
72            let unlock_date = record.get(1).unwrap().trim().to_string();
73
74            // Convert token amount to a u64 and then to nano-units.
75            let tokens: u64 = tokens_str
76                .parse()
77                .with_context(|| format!("invalid token amount: {tokens_str}"))?;
78            let nanos = tokens * 1000;
79
80            *locked_by_date.entry(unlock_date).or_insert(0) += nanos;
81        }
82    }
83    Ok(locked_by_date)
84}
85
86/// Converts a raw unlock date string into ISO 8601 format.
87/// It removes a trailing " ([+0-9]+ UTC)" suffix, replaces the first space with
88/// "T", and appends "Z".
89fn format_date(ts: &str, re: &Regex) -> Result<DateTime<Utc>> {
90    let cleaned = re.replace(ts, "");
91    let iso = if let Some(space_index) = cleaned.find(' ') {
92        let mut s = cleaned.to_string();
93        s.replace_range(space_index..=space_index, "T");
94        s.push('Z');
95        s
96    } else {
97        format!("{cleaned}Z")
98    };
99
100    Ok(DateTime::parse_from_rfc3339(&iso)
101        .context(format!("failed to parse timestamp: {iso}",))?
102        .with_timezone(&Utc))
103}
104
105/// Writes the aggregated unlock data into a CSV file.
106fn write_output_csv(output_file: &PathBuf, entries: &[StillLockedEntry]) -> Result<()> {
107    let mut wtr = Writer::from_path(output_file).with_context(|| {
108        format!(
109            "failed to create output CSV file: {}",
110            output_file.display()
111        )
112    })?;
113    for entry in entries {
114        wtr.serialize(entry)?;
115    }
116    wtr.flush()?;
117    Ok(())
118}
119
120fn main() -> Result<()> {
121    // Clone the repository containing raw data.
122    let tmp_dir = tempdir()?;
123    let repo_path = clone_repo(&tmp_dir)?;
124
125    let crate_dir = env!("CARGO_MANIFEST_DIR");
126    let output_file = PathBuf::from(crate_dir).join("data").join(OUTPUT_FILE);
127
128    // Aggregate unlock data from CSV files.
129    let locked_by_date = aggregate_unlocks(&repo_path)?;
130
131    if locked_by_date.is_empty() {
132        println!("No data found – writing empty CSV.");
133        write_output_csv(&output_file, &[])?;
134        return Ok(());
135    }
136
137    // Compute the total locked tokens.
138    let total_locked: u64 = locked_by_date.values().sum();
139
140    // Prepare to transform each entry into an output record.
141    let re = Regex::new(r" [\+0-9]+ UTC")?;
142    let mut cumulative_unlocked = 0;
143    let mut output_entries = Vec::new();
144
145    // Process unlock dates in order.
146    for (ts, &unlocked) in &locked_by_date {
147        cumulative_unlocked += unlocked;
148        let still_locked = total_locked - cumulative_unlocked;
149        let iso_ts = format_date(ts, &re)?;
150        output_entries.push(StillLockedEntry {
151            timestamp: iso_ts,
152            amount_still_locked: still_locked,
153        });
154    }
155
156    // Write the aggregated data to a CSV file.
157    write_output_csv(&output_file, &output_entries)?;
158    println!("Done: {}", output_file.display());
159
160    Ok(())
161}
generate_aggregated_data/generate_aggregated_data.rs

generate_aggregated_data/
generate_aggregated_data.rs
