iota_single_node_benchmark/tx_generator/
package_publish_tx_generator.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{collections::BTreeMap, fs, path::PathBuf};
6
7use iota_move_build::{BuildConfig, CompiledPackage};
8use iota_test_transaction_builder::{PublishData, TestTransactionBuilder};
9use iota_types::{
10    base_types::ObjectID,
11    transaction::{DEFAULT_VALIDATOR_GAS_PRICE, Transaction},
12};
13use move_package::source_package::manifest_parser::parse_move_manifest_from_file;
14use move_symbol_pool::Symbol;
15use serde::{Deserialize, Serialize};
16use tracing::info;
17
18use crate::{
19    benchmark_context::BenchmarkContext, mock_account::Account, tx_generator::TxGenerator,
20};
21
22pub struct PackagePublishTxGenerator {
23    compiled_package: CompiledPackage,
24}
25
26impl PackagePublishTxGenerator {
27    pub async fn new(ctx: &mut BenchmarkContext, manifest_path: PathBuf) -> Self {
28        let manifest = load_manifest_json(&manifest_path);
29        let dir = manifest_path.parent().unwrap();
30        let PackageDependencyManifest {
31            dependencies,
32            root_package,
33        } = manifest;
34        let mut dep_map = BTreeMap::new();
35        for dependency in dependencies {
36            let Package {
37                name,
38                path,
39                is_source_code,
40            } = dependency;
41
42            info!("Publishing dependent package {}", name);
43            let target_path = dir.join(&path);
44            let module_bytes = if is_source_code {
45                let compiled_package = BuildConfig::new_for_testing_replace_addresses(vec![(
46                    name.clone(),
47                    ObjectID::ZERO,
48                )])
49                .build(&target_path)
50                .unwrap();
51                compiled_package.get_package_bytes(false)
52            } else {
53                let toml = parse_move_manifest_from_file(&target_path.join("Move.toml")).unwrap();
54                let package_name = toml.package.name.as_str();
55                let module_dir = target_path
56                    .join("build")
57                    .join(package_name)
58                    .join("bytecode_modules");
59                let mut all_bytes = Vec::new();
60                info!("Loading module bytes from {:?}", module_dir);
61                for entry in fs::read_dir(module_dir).unwrap() {
62                    let entry = entry.unwrap();
63                    let file_path = entry.path();
64                    if file_path.extension().and_then(|s| s.to_str()) == Some("mv") {
65                        let contents = fs::read(file_path).unwrap();
66                        all_bytes.push(contents);
67                    }
68                }
69                all_bytes
70            };
71            let package_id = ctx
72                .publish_package(PublishData::ModuleBytes(module_bytes))
73                .await
74                .0;
75            info!("Published dependent package {}", package_id);
76            dep_map.insert(Symbol::from(name), package_id);
77        }
78
79        let Package {
80            name,
81            path,
82            is_source_code,
83        } = root_package;
84
85        info!("Compiling root package {}", name);
86        assert!(
87            is_source_code,
88            "Only support building root package from source code"
89        );
90
91        let target_path = dir.join(path);
92        let published_deps = dep_map.clone();
93
94        dep_map.insert(Symbol::from(name), ObjectID::ZERO);
95        let mut compiled_package = BuildConfig::new_for_testing_replace_addresses(
96            dep_map.into_iter().map(|(k, v)| (k.to_string(), v)),
97        )
98        .build(&target_path)
99        .unwrap();
100
101        compiled_package.dependency_ids.published = published_deps;
102        Self { compiled_package }
103    }
104}
105
106impl TxGenerator for PackagePublishTxGenerator {
107    fn generate_tx(&self, account: Account) -> Transaction {
108        TestTransactionBuilder::new(
109            account.sender,
110            account.gas_objects[0],
111            DEFAULT_VALIDATOR_GAS_PRICE,
112        )
113        .publish_with_data(PublishData::CompiledPackage(self.compiled_package.clone()))
114        .build_and_sign(account.keypair.as_ref())
115    }
116
117    fn name(&self) -> &'static str {
118        "PackagePublishTxGenerator"
119    }
120}
121
122#[derive(Serialize, Deserialize, Debug)]
123struct PackageDependencyManifest {
124    dependencies: Vec<Package>,
125    root_package: Package,
126}
127
128#[derive(Serialize, Deserialize, Debug)]
129struct Package {
130    name: String,
131    path: PathBuf,
132    is_source_code: bool,
133}
134
135fn load_manifest_json(file_path: &PathBuf) -> PackageDependencyManifest {
136    let data = fs::read_to_string(file_path)
137        .unwrap_or_else(|_| panic!("Unable to read file at: {:?}", file_path));
138    let parsed_data: PackageDependencyManifest = serde_json::from_str(&data)
139        .unwrap_or_else(|_| panic!("Unable to parse json from file at: {:?}", file_path));
140
141    parsed_data
142}