iota_genesis_builder/stardust/migration/
migration.rs

1// Copyright (c) 2024 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4//! Contains the logic for the migration process.
5
6use std::{
7    cmp::Reverse,
8    collections::{HashMap, HashSet},
9    io::{BufWriter, prelude::Write},
10};
11
12use anyhow::Result;
13use iota_move_build::CompiledPackage;
14use iota_protocol_config::ProtocolVersion;
15use iota_sdk::types::block::output::{FoundryOutput, Output, OutputId};
16use iota_types::{
17    IOTA_FRAMEWORK_PACKAGE_ID, IOTA_SYSTEM_PACKAGE_ID, MOVE_STDLIB_PACKAGE_ID, STARDUST_PACKAGE_ID,
18    balance::Balance,
19    base_types::{IotaAddress, ObjectID, TxContext},
20    epoch_data::EpochData,
21    object::Object,
22    stardust::coin_type::CoinType,
23    timelock::timelock::{self, TimeLock, is_timelocked_balance},
24};
25use move_binary_format::file_format_common::VERSION_MAX;
26use tracing::info;
27
28use crate::stardust::{
29    migration::{
30        MigrationTargetNetwork,
31        executor::Executor,
32        verification::{created_objects::CreatedObjects, verify_outputs},
33    },
34    native_token::package_data::NativeTokenPackageData,
35    process_outputs::process_outputs_for_iota,
36    types::{
37        address_swap_map::AddressSwapMap, address_swap_split_map::AddressSwapSplitMap,
38        output_header::OutputHeader,
39    },
40};
41
42/// We fix the protocol version used in the migration.
43pub const MIGRATION_PROTOCOL_VERSION: u64 = 1;
44
45/// The dependencies of the generated packages for native tokens.
46pub const PACKAGE_DEPS: [ObjectID; 4] = [
47    MOVE_STDLIB_PACKAGE_ID,
48    IOTA_FRAMEWORK_PACKAGE_ID,
49    IOTA_SYSTEM_PACKAGE_ID,
50    STARDUST_PACKAGE_ID,
51];
52
53pub(crate) const NATIVE_TOKEN_BAG_KEY_TYPE: &str = "0x01::ascii::String";
54
55/// An alias for representing the timestamp of a Timelock
56pub type ExpirationTimestamp = u64;
57
58/// The orchestrator of the migration process.
59///
60/// It is run by providing an [`Iterator`] of stardust UTXOs, and holds an inner
61/// executor and in-memory object storage for their conversion into objects.
62///
63/// It guarantees the following:
64///
65/// * That foundry UTXOs are sorted by `(milestone_timestamp, output_id)`.
66/// * That the foundry packages and total supplies are created first
67/// * That all other outputs are created in a second iteration over the original
68///   UTXOs.
69/// * That the resulting ledger state is valid.
70///
71/// The migration process results in the generation of a snapshot file with the
72/// generated objects serialized.
73pub struct Migration {
74    target_milestone_timestamp_sec: u32,
75    total_supply: u64,
76    executor: Executor,
77    pub(super) output_objects_map: HashMap<OutputId, CreatedObjects>,
78    /// The coin type to use in order to migrate outputs. Can only be equal to
79    /// `Iota` at the moment. Is fixed for the entire migration process.
80    coin_type: CoinType,
81    address_swap_map: AddressSwapMap,
82}
83
84impl Migration {
85    /// Try to setup the migration process by creating the inner executor
86    /// and bootstrapping the in-memory storage.
87    pub fn new(
88        target_milestone_timestamp_sec: u32,
89        total_supply: u64,
90        target_network: MigrationTargetNetwork,
91        coin_type: CoinType,
92        address_swap_map: AddressSwapMap,
93    ) -> Result<Self> {
94        let executor = Executor::new(
95            ProtocolVersion::new(MIGRATION_PROTOCOL_VERSION),
96            target_network,
97            coin_type,
98        )?;
99        Ok(Self {
100            target_milestone_timestamp_sec,
101            total_supply,
102            executor,
103            output_objects_map: Default::default(),
104            coin_type,
105            address_swap_map,
106        })
107    }
108
109    /// Run all stages of the migration except snapshot migration.
110    /// Factored out to facilitate testing.
111    ///
112    /// See also `Self::run`.
113    pub(crate) fn run_migration(
114        &mut self,
115        outputs: impl IntoIterator<Item = (OutputHeader, Output)>,
116    ) -> Result<()> {
117        let (mut foundries, mut outputs) = outputs.into_iter().fold(
118            (Vec::new(), Vec::new()),
119            |(mut foundries, mut outputs), (header, output)| {
120                if let Output::Foundry(foundry) = output {
121                    foundries.push((header, foundry));
122                } else {
123                    outputs.push((header, output));
124                }
125                (foundries, outputs)
126            },
127        );
128        // We sort the outputs to make sure the order of outputs up to
129        // a certain milestone timestamp remains the same between runs.
130        //
131        // This guarantees that fresh ids created through the transaction
132        // context will also map to the same objects between runs.
133        outputs.sort_by_key(|(header, _)| (header.ms_timestamp(), header.output_id()));
134        foundries.sort_by_key(|(header, _)| (header.ms_timestamp(), header.output_id()));
135        info!("Migrating foundries...");
136        self.migrate_foundries(&foundries)?;
137        info!("Migrating the rest of outputs...");
138        self.migrate_outputs(&outputs)?;
139        let outputs = outputs
140            .into_iter()
141            .chain(foundries.into_iter().map(|(h, f)| (h, Output::Foundry(f))))
142            .collect::<Vec<_>>();
143        info!("Verifying ledger state...");
144        self.verify_ledger_state(&outputs)?;
145        self.address_swap_map.verify_all_addresses_swapped()?;
146        Ok(())
147    }
148
149    /// Run all stages of the migration.
150    ///
151    /// * Generate and build the foundry packages
152    /// * Create the foundry packages, and associated objects.
153    /// * Create all other objects.
154    /// * Validate the resulting object-based ledger state.
155    /// * Create the snapshot file.
156    pub fn run(
157        mut self,
158        outputs: impl IntoIterator<Item = (OutputHeader, Output)>,
159        writer: impl Write,
160    ) -> Result<()> {
161        info!("Starting the migration...");
162        self.run_migration(outputs)?;
163        info!("Migration ended.");
164        info!("Writing snapshot file...");
165        create_snapshot(self.into_objects(), writer)?;
166        info!("Snapshot file written.");
167        Ok(())
168    }
169
170    /// Run all stages of the migration coming from a Hornet snapshot with IOTA
171    /// coin type.
172    pub fn run_for_iota<'a>(
173        self,
174        target_milestone_timestamp: u32,
175        swap_split_map: AddressSwapSplitMap,
176        outputs: impl Iterator<Item = Result<(OutputHeader, Output)>> + 'a,
177        writer: impl Write,
178    ) -> Result<()> {
179        itertools::process_results(
180            process_outputs_for_iota(target_milestone_timestamp, swap_split_map, outputs),
181            |outputs| self.run(outputs, writer),
182        )?
183    }
184
185    /// The migration objects.
186    ///
187    /// The system packages and underlying `init` objects
188    /// are filtered out because they will be generated
189    /// in the genesis process.
190    fn into_objects(self) -> Vec<Object> {
191        self.executor.into_objects()
192    }
193
194    /// Create the packages, and associated objects representing foundry
195    /// outputs.
196    fn migrate_foundries<'a>(
197        &mut self,
198        foundries: impl IntoIterator<Item = &'a (OutputHeader, FoundryOutput)>,
199    ) -> Result<()> {
200        let compiled = foundries
201            .into_iter()
202            .map(|(header, output)| {
203                let pkg = generate_package(output)?;
204                Ok((header, output, pkg))
205            })
206            .collect::<Result<Vec<_>>>()?;
207        self.output_objects_map
208            .extend(self.executor.create_foundries(compiled.into_iter())?);
209        Ok(())
210    }
211
212    /// Create objects for all outputs except for foundry outputs.
213    fn migrate_outputs<'a>(
214        &mut self,
215        outputs: impl IntoIterator<Item = &'a (OutputHeader, Output)>,
216    ) -> Result<()> {
217        for (header, output) in outputs {
218            let created = match output {
219                Output::Alias(alias) => self.executor.create_alias_objects(
220                    header,
221                    alias,
222                    self.coin_type,
223                    &mut self.address_swap_map,
224                )?,
225                Output::Nft(nft) => self.executor.create_nft_objects(
226                    header,
227                    nft,
228                    self.coin_type,
229                    &mut self.address_swap_map,
230                )?,
231                Output::Basic(basic) => {
232                    // All timelocked vested rewards(basic outputs with the specific ID format)
233                    // should be migrated as TimeLock<Balance<IOTA>> objects.
234                    if timelock::is_timelocked_vested_reward(
235                        header.output_id(),
236                        basic,
237                        self.target_milestone_timestamp_sec,
238                    ) {
239                        self.executor.create_timelock_object(
240                            header.output_id(),
241                            basic,
242                            self.target_milestone_timestamp_sec,
243                            &mut self.address_swap_map,
244                        )?
245                    } else {
246                        self.executor.create_basic_objects(
247                            header,
248                            basic,
249                            self.target_milestone_timestamp_sec,
250                            &self.coin_type,
251                            &mut self.address_swap_map,
252                        )?
253                    }
254                }
255                Output::Treasury(_) | Output::Foundry(_) => continue,
256            };
257            self.output_objects_map.insert(header.output_id(), created);
258        }
259        Ok(())
260    }
261
262    /// Verify the ledger state represented by the objects in
263    /// [`InMemoryStorage`](iota_types::in_memory_storage::InMemoryStorage).
264    pub fn verify_ledger_state<'a>(
265        &self,
266        outputs: impl IntoIterator<Item = &'a (OutputHeader, Output)>,
267    ) -> Result<()> {
268        verify_outputs(
269            outputs,
270            &self.output_objects_map,
271            self.executor.native_tokens(),
272            self.target_milestone_timestamp_sec,
273            self.total_supply,
274            self.executor.store(),
275            &self.address_swap_map,
276        )?;
277        Ok(())
278    }
279
280    /// Consumes the `Migration` and returns the underlying `Executor` and
281    /// created objects map, so tests can continue to work in the same
282    /// environment as the migration.
283    #[cfg(test)]
284    pub(super) fn into_parts(self) -> (Executor, HashMap<OutputId, CreatedObjects>) {
285        (self.executor, self.output_objects_map)
286    }
287}
288
289/// All the objects created during the migration.
290///
291/// Internally it maintains indexes of [`TimeLock`] and
292/// [`iota_types::gas_coin::GasCoin`] objects grouped by their owners to
293/// accommodate queries of this sort.
294#[derive(Debug, Clone, Default)]
295pub struct MigrationObjects {
296    inner: Vec<Object>,
297    owner_timelock: HashMap<IotaAddress, Vec<usize>>,
298    owner_gas_coin: HashMap<IotaAddress, Vec<usize>>,
299}
300
301impl Extend<Object> for MigrationObjects {
302    fn extend<T: IntoIterator<Item = Object>>(&mut self, iter: T) {
303        let last_current_ix = self.inner.len();
304        self.inner.extend(iter);
305        for (i, tag, object) in self.inner[last_current_ix..]
306            .iter()
307            .zip(last_current_ix..)
308            .filter_map(|(object, i)| {
309                let tag = object.struct_tag()?;
310                Some((i, tag, object))
311            })
312        {
313            let owner_object_map = if is_timelocked_balance(&tag) {
314                &mut self.owner_timelock
315            } else if object.is_gas_coin() {
316                &mut self.owner_gas_coin
317            } else {
318                continue;
319            };
320            let owner = object
321                .owner
322                .get_owner_address()
323                .expect("timelocks should have an address owner");
324            owner_object_map
325                .entry(owner)
326                .and_modify(|object_ixs| object_ixs.push(i))
327                .or_insert_with(|| vec![i]);
328        }
329    }
330}
331
332impl MigrationObjects {
333    pub fn new(objects: Vec<Object>) -> Self {
334        let mut migration_objects = Self::default();
335        migration_objects.extend(objects);
336        migration_objects
337    }
338
339    /// Evict the objects with the specified ids
340    pub fn evict(&mut self, objects: impl IntoIterator<Item = ObjectID>) {
341        let eviction_set = objects.into_iter().collect::<HashSet<_>>();
342        let inner = std::mem::take(&mut self.inner);
343        self.inner = inner
344            .into_iter()
345            .filter(|object| !eviction_set.contains(&object.id()))
346            .collect();
347    }
348
349    /// Take the inner migration objects.
350    ///
351    /// This follows the semantics of [`std::mem::take`].
352    pub fn take_objects(&mut self) -> Vec<Object> {
353        std::mem::take(&mut self.inner)
354    }
355
356    /// Checks if inner is empty.
357    pub fn is_empty(&self) -> bool {
358        self.inner.is_empty()
359    }
360
361    /// Get [`TimeLock`] objects created during the migration together with
362    /// their expiration timestamp.
363    ///
364    /// The query is filtered by the object owner.
365    ///
366    /// The returned objects are ordered by expiration timestamp, in descending
367    /// order.
368    pub fn get_sorted_timelocks_and_expiration_by_owner(
369        &self,
370        address: IotaAddress,
371    ) -> Option<Vec<(&Object, ExpirationTimestamp)>> {
372        self.get_timelocks_and_expiration_by_owner(address)
373            .map(|mut timelocks| {
374                timelocks.sort_by_key(|&(_, timestamp)| Reverse(timestamp));
375                timelocks
376            })
377    }
378
379    /// Get [`TimeLock`] objects created during the migration together with
380    /// their expiration timestamp.
381    ///
382    /// The query is filtered by the object owner.
383    pub fn get_timelocks_and_expiration_by_owner(
384        &self,
385        address: IotaAddress,
386    ) -> Option<Vec<(&Object, ExpirationTimestamp)>> {
387        Some(
388            self.owner_timelock
389                .get(&address)?
390                .iter()
391                .map(|i| {
392                    (
393                        &self.inner[*i],
394                        self.inner[*i]
395                            .to_rust::<TimeLock<Balance>>()
396                            .expect("this should be a TimeLock object")
397                            .expiration_timestamp_ms(),
398                    )
399                })
400                .collect(),
401        )
402    }
403
404    /// Get [`iota_types::gas_coin::GasCoin`] objects created during the
405    /// migration.
406    ///
407    /// The query is filtered by the object owner.
408    pub fn get_gas_coins_by_owner(&self, address: IotaAddress) -> Option<Vec<&Object>> {
409        Some(
410            self.owner_gas_coin
411                .get(&address)?
412                .iter()
413                .map(|i| &self.inner[*i])
414                .collect(),
415        )
416    }
417}
418
419// Build a `CompiledPackage` from a given `FoundryOutput`.
420fn generate_package(foundry: &FoundryOutput) -> Result<CompiledPackage> {
421    let native_token_data = NativeTokenPackageData::try_from(foundry)?;
422    crate::stardust::native_token::package_builder::build_and_compile(native_token_data)
423}
424
425/// Serialize the objects stored in [`InMemoryStorage`] into a file using
426/// [`bcs`] encoding.
427fn create_snapshot(ledger: Vec<Object>, writer: impl Write) -> Result<()> {
428    let mut writer = BufWriter::new(writer);
429    writer.write_all(&bcs::to_bytes(&ledger)?)?;
430    Ok(writer.flush()?)
431}
432
433/// Get the bytes of all bytecode modules (not including direct or transitive
434/// dependencies) of [`CompiledPackage`].
435pub(super) fn package_module_bytes(pkg: &CompiledPackage) -> Result<Vec<Vec<u8>>> {
436    pkg.get_modules()
437        .map(|module| {
438            let mut buf = Vec::new();
439            module.serialize_with_version(VERSION_MAX, &mut buf)?;
440            Ok(buf)
441        })
442        .collect::<Result<_>>()
443}
444
445/// Create a [`TxContext]` that remains the same across invocations.
446pub(super) fn create_migration_context(
447    coin_type: &CoinType,
448    target_network: MigrationTargetNetwork,
449) -> TxContext {
450    TxContext::new(
451        &IotaAddress::default(),
452        &target_network.migration_transaction_digest(coin_type),
453        &EpochData::new_genesis(0),
454    )
455}
456
457#[cfg(test)]
458mod tests {
459    use iota_protocol_config::ProtocolConfig;
460    use iota_types::{
461        balance::Balance,
462        base_types::SequenceNumber,
463        gas_coin::GasCoin,
464        id::UID,
465        object::{Data, Owner},
466        timelock::timelock::{TimeLock, to_genesis_object},
467    };
468
469    use super::*;
470
471    #[test]
472    fn migration_objects_get_timelocks() {
473        let owner = IotaAddress::random_for_testing_only();
474        let address = IotaAddress::random_for_testing_only();
475        let tx_context = TxContext::random_for_testing_only();
476        let expected_timelocks = (0..4)
477            .map(|_| TimeLock::new(UID::new(ObjectID::random()), Balance::new(0), 0, None))
478            .map(|timelock| {
479                to_genesis_object(
480                    timelock,
481                    owner,
482                    &ProtocolConfig::get_for_min_version(),
483                    &tx_context,
484                    SequenceNumber::MIN,
485                )
486                .unwrap()
487            })
488            .collect::<Vec<_>>();
489        let non_matching_timelocks = (0..8)
490            .map(|_| TimeLock::new(UID::new(ObjectID::random()), Balance::new(0), 0, None))
491            .map(|timelock| {
492                to_genesis_object(
493                    timelock,
494                    address,
495                    &ProtocolConfig::get_for_min_version(),
496                    &tx_context,
497                    SequenceNumber::MIN,
498                )
499                .unwrap()
500            });
501        let non_matching_objects = (0..8)
502            .map(|_| GasCoin::new_for_testing(0).to_object(SequenceNumber::MIN))
503            .map(|move_object| {
504                Object::new_from_genesis(
505                    Data::Move(move_object),
506                    Owner::AddressOwner(address),
507                    tx_context.digest(),
508                )
509            });
510        let migration_objects = MigrationObjects::new(
511            non_matching_objects
512                .chain(non_matching_timelocks)
513                .chain(expected_timelocks.clone())
514                .collect(),
515        );
516        let matching_objects = migration_objects
517            .get_timelocks_and_expiration_by_owner(owner)
518            .unwrap();
519        assert_eq!(
520            expected_timelocks,
521            matching_objects
522                .into_iter()
523                .map(|(timelock, _)| timelock.clone())
524                .collect::<Vec<Object>>()
525        );
526    }
527
528    #[test]
529    fn migration_objects_get_gas_coins() {
530        let owner = IotaAddress::random_for_testing_only();
531        let address = IotaAddress::random_for_testing_only();
532        let tx_context = TxContext::random_for_testing_only();
533        let non_matching_timelocks = (0..8)
534            .map(|_| TimeLock::new(UID::new(ObjectID::random()), Balance::new(0), 0, None))
535            .map(|timelock| {
536                to_genesis_object(
537                    timelock,
538                    address,
539                    &ProtocolConfig::get_for_min_version(),
540                    &tx_context,
541                    SequenceNumber::MIN,
542                )
543                .unwrap()
544            });
545        let expected_gas_coins = (0..8)
546            .map(|_| GasCoin::new_for_testing(0).to_object(SequenceNumber::MIN))
547            .map(|move_object| {
548                Object::new_from_genesis(
549                    Data::Move(move_object),
550                    Owner::AddressOwner(owner),
551                    tx_context.digest(),
552                )
553            })
554            .collect::<Vec<_>>();
555        let migration_objects = MigrationObjects::new(
556            non_matching_timelocks
557                .chain(expected_gas_coins.clone())
558                .collect(),
559        );
560        let matching_objects = migration_objects.get_gas_coins_by_owner(owner).unwrap();
561        assert_eq!(
562            expected_gas_coins,
563            matching_objects
564                .into_iter()
565                .cloned()
566                .collect::<Vec<Object>>()
567        );
568    }
569}
iota_genesis_builder/stardust/migration/migration.rs

iota_genesis_builder/stardust/migration/
migration.rs
