[move] Added bytecode map generation during bytecode disassembly (Mys…

…tenLabs#20797)

## Description 

This PR adds support for generating "bytecode maps" (source maps for
disassembled bytecode)

## Test plan 

A new test has been added and all other test have to pass

crates/sui-framework-tests/tests/move_tests.rs

@@ -58,7 +58,7 @@ pub(crate) fn tests(path: &Path) -> datatest_stable::Result<()> {
     testing_config.filter = std::env::var("FILTER").ok().map(|s| s.to_string());
 
     assert_eq!(
-        run_move_unit_tests(path, move_config, Some(testing_config), false).unwrap(),
+        run_move_unit_tests(path, move_config, Some(testing_config), false, false).unwrap(),
         UnitTestResult::Success
     );
 

crates/sui-move/Cargo.toml

@@ -24,6 +24,7 @@ move-disassembler.workspace = true
 move-ir-types.workspace = true
 move-package.workspace = true
 move-prover.workspace = true
+move-bytecode-source-map.workspace = true
 move-unit-test.workspace = true
 telemetry-subscribers.workspace = true
 tokio = { workspace = true, features = ["full"] }

crates/sui-move/src/disassemble.rs

@@ -3,6 +3,7 @@
 
 use clap::Parser;
 use move_binary_format::CompiledModule;
+use move_bytecode_source_map::utils::serialize_to_json_string;
 use move_cli::base;
 use move_disassembler::disassembler::Disassembler;
 use move_ir_types::location::Spanned;
@@ -25,6 +26,10 @@ pub struct Disassemble {
 
     #[clap(short = 'i', long = "interactive")]
     interactive: bool,
+
+    /// Print the "bytecode map" (source map for disassembled bytecode)
+    #[clap(long = "bytecode-map")]
+    pub bytecode_map: bool,
 }
 
 impl Disassemble {
@@ -47,6 +52,7 @@ impl Disassemble {
                 package_name: None,
                 module_or_script_name: module_name,
                 debug: self.debug,
+                bytecode_map: self.bytecode_map,
             }
             .execute(package_path, build_config)?;
             return Ok(());
@@ -69,7 +75,11 @@ impl Disassemble {
             println!("{module:#?}");
         } else {
             let d = Disassembler::from_module(&module, Spanned::unsafe_no_loc(()).loc)?;
-            println!("{}", d.disassemble()?);
+            let (disassemble_string, bcode_map) = d.disassemble_with_source_map()?;
+            if self.bytecode_map {
+                println!("{}", serialize_to_json_string(&bcode_map)?);
+            }
+            println!("{}", disassemble_string);
         }
 
         Ok(())

crates/sui-move/src/unit_test.rs

@@ -42,6 +42,8 @@ impl Test {
                 "The --coverage flag is currently supported only in debug builds. Please build the Sui CLI from source in debug mode."
             ));
         }
+        // save disassembly if trace execution is enabled
+        let save_disassembly = self.test.trace_execution.is_some();
         // find manifest file directory from a given path or (if missing) from current dir
         let rerooted_path = base::reroot_path(path)?;
         let unit_test_config = self.test.unit_test_config();
@@ -50,6 +52,7 @@ impl Test {
             build_config,
             Some(unit_test_config),
             compute_coverage,
+            save_disassembly,
         )
     }
 }
@@ -71,6 +74,7 @@ pub fn run_move_unit_tests(
     build_config: BuildConfig,
     config: Option<UnitTestingConfig>,
     compute_coverage: bool,
+    save_disassembly: bool,
 ) -> anyhow::Result<UnitTestResult> {
     // bind the extension hook if it has not yet been done
     Lazy::force(&SET_EXTENSION_HOOK);
@@ -91,6 +95,7 @@ pub fn run_move_unit_tests(
         ),
         Some(initial_cost_schedule_for_unit_tests()),
         compute_coverage,
+        save_disassembly,
         &mut std::io::stdout(),
     );
     result.map(|(test_result, warning_diags)| {

external-crates/move/crates/move-bytecode-source-map/src/source_map.rs

@@ -712,4 +712,68 @@ impl SourceMap {
 
         Ok(empty_source_map)
     }
+
+    pub fn replace_file_hashes(&mut self, file_hash: FileHash) {
+        self.definition_location = Loc::new(
+            file_hash,
+            self.definition_location.start(),
+            self.definition_location.end(),
+        );
+        for (_, struct_map) in self.struct_map.iter_mut() {
+            struct_map.definition_location = Loc::new(
+                file_hash,
+                struct_map.definition_location.start(),
+                struct_map.definition_location.end(),
+            );
+            for (_, loc) in struct_map.type_parameters.iter_mut() {
+                *loc = Loc::new(file_hash, loc.start(), loc.end());
+            }
+            for loc in struct_map.fields.iter_mut() {
+                *loc = Loc::new(file_hash, loc.start(), loc.end());
+            }
+        }
+        for (_, enum_map) in self.enum_map.iter_mut() {
+            enum_map.definition_location = Loc::new(
+                file_hash,
+                enum_map.definition_location.start(),
+                enum_map.definition_location.end(),
+            );
+            for (_, loc) in enum_map.type_parameters.iter_mut() {
+                *loc = Loc::new(file_hash, loc.start(), loc.end());
+            }
+            for ((_, loc), field_locations) in enum_map.variants.iter_mut() {
+                *loc = Loc::new(file_hash, loc.start(), loc.end());
+                for field_loc in field_locations.iter_mut() {
+                    *field_loc = Loc::new(file_hash, field_loc.start(), field_loc.end());
+                }
+            }
+        }
+        for (_, function_map) in self.function_map.iter_mut() {
+            function_map.location = Loc::new(
+                file_hash,
+                function_map.location.start(),
+                function_map.location.end(),
+            );
+            function_map.definition_location = Loc::new(
+                file_hash,
+                function_map.definition_location.start(),
+                function_map.definition_location.end(),
+            );
+            for (_, loc) in function_map.type_parameters.iter_mut() {
+                *loc = Loc::new(file_hash, loc.start(), loc.end());
+            }
+            for (_, loc) in function_map.parameters.iter_mut() {
+                *loc = Loc::new(file_hash, loc.start(), loc.end());
+            }
+            for loc in function_map.returns.iter_mut() {
+                *loc = Loc::new(file_hash, loc.start(), loc.end());
+            }
+            for (_, loc) in function_map.locals.iter_mut() {
+                *loc = Loc::new(file_hash, loc.start(), loc.end());
+            }
+            for (_, loc) in function_map.code_map.iter_mut() {
+                *loc = Loc::new(file_hash, loc.start(), loc.end());
+            }
+        }
+    }
 }

external-crates/move/crates/move-bytecode-source-map/src/utils.rs

@@ -27,13 +27,16 @@ pub fn source_map_from_file(file_path: &Path) -> Result<SourceMap> {
         .map_err(|_| format_err!("Error deserializing into source map"))
 }
 
+pub fn serialize_to_json_string(map: &SourceMap) -> Result<String> {
+    serde_json::to_string_pretty(map).map_err(|e| format_err!("Error serializing to json: {}", e))
+}
+
 pub fn serialize_to_json(map: &SourceMap) -> Result<Vec<u8>> {
     serde_json::to_vec(map).map_err(|e| format_err!("Error serializing to json: {}", e))
 }
 
 pub fn serialize_to_json_file(map: &SourceMap, file_path: &Path) -> Result<()> {
-    let json = serde_json::to_string_pretty(map)
-        .map_err(|e| format_err!("Error serializing to json: {}", e))?;
+    let json = serialize_to_json_string(map)?;
     let mut f =
         std::fs::File::create(file_path).map_err(|e| format_err!("Error creating file: {}", e))?;
     f.write_all(json.as_bytes())

external-crates/move/crates/move-cli/Cargo.toml

@@ -22,6 +22,7 @@ toml_edit.workspace = true
 
 bcs.workspace = true
 
+move-bytecode-source-map.workspace = true
 move-bytecode-verifier.workspace = true
 move-disassembler.workspace = true
 move-docgen.workspace = true
@@ -40,6 +41,7 @@ move-vm-test-utils.workspace = true
 move-binary-format.workspace = true
 move-package.workspace = true
 move-prover.workspace = true
+move-symbol-pool.workspace = true
 move-unit-test.workspace = true
 move-bytecode-viewer.workspace = true
 

external-crates/move/crates/move-cli/src/base/disassemble.rs

@@ -3,6 +3,7 @@
 
 use super::reroot_path;
 use clap::*;
+use move_bytecode_source_map::utils::serialize_to_json_string;
 use move_compiler::compiled_unit::NamedCompiledModule;
 use move_disassembler::disassembler::Disassembler;
 use move_package::{compilation::compiled_package::CompiledUnitWithSource, BuildConfig};
@@ -24,6 +25,9 @@ pub struct Disassemble {
     #[clap(long = "Xdebug")]
     /// Also print the raw disassembly using Rust's Debug output, at the end.
     pub debug: bool,
+    #[clap(long = "bytecode-map")]
+    /// Print the "bytecode map" (source map for disassembled bytecode)
+    pub bytecode_map: bool,
 }
 
 impl Disassemble {
@@ -34,6 +38,7 @@ impl Disassemble {
             package_name,
             module_or_script_name,
             debug,
+            bytecode_map,
         } = self;
         // Make sure the package is built
         let package = config.compile_package(&rerooted_path, &mut Vec::new())?;
@@ -66,7 +71,12 @@ impl Disassemble {
                         source_path,
                     )
                 } else {
-                    println!("{}", Disassembler::from_unit(&unit.unit).disassemble()?);
+                    let d = Disassembler::from_unit(&unit.unit);
+                    let (disassemble_string, bcode_map) = d.disassemble_with_source_map()?;
+                    if bytecode_map {
+                        println!("{}", serialize_to_json_string(&bcode_map)?);
+                    }
+                    println!("{}", disassemble_string);
                     if debug {
                         println!("\n{:#?}", &unit.unit.module)
                     }

external-crates/move/crates/move-cli/src/base/test.rs

@@ -84,13 +84,16 @@ impl Test {
     ) -> anyhow::Result<()> {
         let rerooted_path = reroot_path(path)?;
         let compute_coverage = self.compute_coverage;
+        // save disassembly if trace execution is enabled
+        let save_disassembly = self.trace_execution.is_some();
         let result = run_move_unit_tests(
             &rerooted_path,
             config,
             self.unit_test_config(),
             natives,
             cost_table,
             compute_coverage,
+            save_disassembly,
             &mut std::io::stdout(),
         )?;
 
@@ -143,11 +146,13 @@ pub fn run_move_unit_tests<W: Write + Send>(
     natives: Vec<NativeFunctionRecord>,
     cost_table: Option<CostTable>,
     compute_coverage: bool,
+    save_disassembly: bool,
     writer: &mut W,
 ) -> Result<(UnitTestResult, Option<Diagnostics>)> {
     let mut test_plan = None;
     build_config.test_mode = true;
     build_config.dev_mode = true;
+    build_config.save_disassembly = save_disassembly;
 
     // Build the resolution graph (resolution graph diagnostics are only needed for CLI commands so
     // ignore them by passing a vector as the writer)

external-crates/move/crates/move-cli/tests/build_tests/disassemble_module_source_map/Move.toml

@@ -0,0 +1,3 @@
+[package]
+name = "Test"
+edition = "2024.beta"