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

iota_genesis_builder/stardust/migration/
migration.rs
