[fixes] Fix raw module data & add env vars

aptos-labs · Nov 12, 2024 · e425538 · e425538
1 parent a7fa2c3
commit e425538
Show file tree

Hide file tree

Showing 8 changed files with 231 additions and 194 deletions.
diff --git a/aptos-move/aptos-vm-environment/src/environment.rs b/aptos-move/aptos-vm-environment/src/environment.rs
@@ -15,7 +15,8 @@ use aptos_native_interface::SafeNativeBuilder;
 use aptos_types::{
     chain_id::ChainId,
     on_chain_config::{
-        ConfigurationResource, Features, OnChainConfig, TimedFeatures, TimedFeaturesBuilder,
+        ConfigurationResource, FeatureFlag, Features, OnChainConfig, TimedFeatures,
+        TimedFeaturesBuilder,
     },
     state_store::StateView,
 };
@@ -175,8 +176,13 @@ impl Environment {
     ) -> Self {
         // We compute and store a hash of configs in order to distinguish different environments.
         let mut sha3_256 = Sha3_256::new();
-        let features =
+        let mut features =
             fetch_config_and_update_hash::<Features>(&mut sha3_256, state_view).unwrap_or_default();
+        if std::env::var("USE_LOADER_V2").is_ok() {
+            features.enable(FeatureFlag::ENABLE_LOADER_V2);
+        } else {
+            features.disable(FeatureFlag::ENABLE_LOADER_V2);
+        }
 
         // If no chain ID is in storage, we assume we are in a testing environment.
         let chain_id = fetch_config_and_update_hash::<ChainId>(&mut sha3_256, state_view)

diff --git a/aptos-move/aptos-vm/src/move_vm_ext/vm.rs b/aptos-move/aptos-vm/src/move_vm_ext/vm.rs
@@ -7,7 +7,7 @@ use aptos_gas_schedule::{MiscGasParameters, NativeGasParameters, LATEST_GAS_FEAT
 use aptos_native_interface::SafeNativeBuilder;
 use aptos_types::{
     chain_id::ChainId,
-    on_chain_config::{Features, TimedFeaturesBuilder},
+    on_chain_config::{FeatureFlag, Features, TimedFeaturesBuilder},
     transaction::user_transaction_context::UserTransactionContext,
 };
 use aptos_vm_environment::{
@@ -30,7 +30,13 @@ pub struct GenesisRuntimeBuilder {
 impl GenesisRuntimeBuilder {
     /// Returns a builder, capable of creating VM and runtime environment to run genesis.
     pub fn new(chain_id: ChainId) -> Self {
-        let features = Features::default();
+        let mut features = Features::default();
+        if std::env::var("USE_LOADER_V2").is_ok() {
+            features.enable(FeatureFlag::ENABLE_LOADER_V2);
+        } else {
+            features.disable(FeatureFlag::ENABLE_LOADER_V2);
+        }
+
         let timed_features = TimedFeaturesBuilder::enable_all().build();
 
         let vm_config =

diff --git a/crates/transaction-generator-lib/src/args.rs b/crates/transaction-generator-lib/src/args.rs
@@ -323,9 +323,13 @@ impl TransactionTypeArg {
             },
             TransactionTypeArg::DeserializeU256 => call_custom_module(EntryPoints::DeserializeU256),
             TransactionTypeArg::SimpleScript => call_custom_module(EntryPoints::SimpleScript),
-            TransactionTypeArg::ChainDependencies => call_custom_module(EntryPoints::ChainDependencies),
+            TransactionTypeArg::ChainDependencies => {
+                call_custom_module(EntryPoints::ChainDependencies)
+            },
             TransactionTypeArg::ChainFriends => call_custom_module(EntryPoints::ChainFriends),
-            TransactionTypeArg::StarDependencies => call_custom_module(EntryPoints::StarDependencies),
+            TransactionTypeArg::StarDependencies => {
+                call_custom_module(EntryPoints::StarDependencies)
+            },
         }
     }
 

diff --git a/crates/transaction-generator-lib/src/publishing/module_simple.rs b/crates/transaction-generator-lib/src/publishing/module_simple.rs
@@ -395,9 +395,9 @@ impl EntryPoints {
             EntryPoints::IncGlobalMilestoneAggV2 { .. }
             | EntryPoints::CreateGlobalMilestoneAggV2 { .. } => "counter_with_milestone",
             EntryPoints::DeserializeU256 => "bcs_stream",
-            EntryPoints::ChainDependencies => "module_loading_chain_dependencies",
-            EntryPoints::ChainFriends => "module_loading_chain_friends",
-            EntryPoints::StarDependencies => "module_loading_star_dependencies",
+            EntryPoints::ChainDependencies => "chain_dependencies",
+            EntryPoints::ChainFriends => "chain_friends",
+            EntryPoints::StarDependencies => "star_dependencies",
         }
     }
 
@@ -741,7 +741,11 @@ impl EntryPoints {
                     ],
                 )
             },
-            EntryPoints::ChainDependencies | EntryPoints::ChainFriends | EntryPoints::StarDependencies => get_payload_void(module_id, ident_str!("run").to_owned()),
+            EntryPoints::ChainDependencies
+            | EntryPoints::ChainFriends
+            | EntryPoints::StarDependencies => {
+                get_payload_void(module_id, ident_str!("run").to_owned())
+            },
         }
     }
 
@@ -850,7 +854,9 @@ impl EntryPoints {
             EntryPoints::DeserializeU256 => AutomaticArgs::None,
             EntryPoints::IncGlobalMilestoneAggV2 { .. } => AutomaticArgs::None,
             EntryPoints::CreateGlobalMilestoneAggV2 { .. } => AutomaticArgs::Signer,
-            EntryPoints::ChainDependencies | EntryPoints::ChainFriends | EntryPoints::StarDependencies => AutomaticArgs::None,
+            EntryPoints::ChainDependencies
+            | EntryPoints::ChainFriends
+            | EntryPoints::StarDependencies => AutomaticArgs::None,
         }
     }
 }