Get coverage working better, add --show-locs flag to aptos move disas…

…semble

crates/aptos/src/move_tool/bytecode.rs

@@ -27,7 +27,7 @@ use move_disassembler::disassembler::{Disassembler, DisassemblerOptions};
 use move_ir_types::location::Spanned;
 use std::{
     fs,
-    path::{Path, PathBuf},
+    path::{Component, Path, PathBuf},
     process::Command,
 };
 
@@ -63,6 +63,9 @@ pub struct BytecodeCommand {
     #[clap(long)]
     pub is_script: bool,
 
+    #[clap(long)]
+    pub show_locs: bool,
+
     #[clap(flatten)]
     input: BytecodeCommandInput,
 
@@ -203,18 +206,51 @@ impl BytecodeCommand {
 
     fn disassemble(&self, bytecode_path: &Path) -> Result<String, CliError> {
         let bytecode_bytes = read_from_file(bytecode_path)?;
-        let move_path = bytecode_path.with_extension(MOVE_EXTENSION);
-        let source_map_path = bytecode_path.with_extension(SOURCE_MAP_EXTENSION);
 
-        let source = fs::read_to_string(move_path).ok();
-        let source_map = source_map_from_file(&source_map_path).ok();
+        let source = {
+            let move_path = bytecode_path.with_extension(MOVE_EXTENSION);
+            if let Ok(source) = fs::read_to_string(move_path.clone()) {
+                Some(source)
+            } else {
+                let move_path = move_path
+                    .components()
+                    .map(|elt| {
+                        if elt.as_os_str() == "bytecode_modules" {
+                            Component::Normal("sources".as_ref())
+                        } else {
+                            elt.clone()
+                        }
+                    })
+                    .collect::<PathBuf>();
+                fs::read_to_string(move_path).ok()
+            }
+        };
+        let source_map = {
+            let source_map_path = bytecode_path.with_extension(SOURCE_MAP_EXTENSION);
+            if let Ok(source_map) = source_map_from_file(&source_map_path) {
+                Some(source_map)
+            } else {
+                let source_map_path = source_map_path
+                    .components()
+                    .map(|elt| {
+                        if elt.as_os_str() == "bytecode_modules" {
+                            Component::Normal("source_maps".as_ref())
+                        } else {
+                            elt.clone()
+                        }
+                    })
+                    .collect::<PathBuf>();
+                source_map_from_file(&source_map_path).ok()
+            }
+        };
 
         let disassembler_options = DisassemblerOptions {
             print_code: true,
             only_externally_visible: false,
             print_basic_blocks: true,
             print_locals: true,
             print_bytecode_stats: false,
+            show_locs: self.show_locs,
         };
         let no_loc = Spanned::unsafe_no_loc(()).loc;
         let module: CompiledModule;
@@ -232,6 +268,11 @@ impl BytecodeCommand {
         let mut source_mapping = if let Some(s) = source_map {
             SourceMapping::new(s, bytecode)
         } else {
+            if self.show_locs {
+                return Err(CliError::UnexpectedError(
+                    "Missing source_map with option show-locs".to_string(),
+                ));
+            }
             SourceMapping::new_from_view(bytecode, no_loc)
                 .context("Unable to build dummy source mapping")?
         };

crates/aptos/src/move_tool/coverage.rs

@@ -124,7 +124,17 @@ impl CliCommand<()> for SourceCoverage {
             }) => (module, source_map),
             _ => panic!("Should all be modules"),
         };
-        let source_coverage = SourceCoverageBuilder::new(module, &coverage_map, source_map);
+        let packages: Vec<_> = package
+            .root_modules()
+            .map(|unit| match &unit.unit {
+                CompiledUnit::Module(NamedCompiledModule {
+                    module, source_map, ..
+                }) => (module, source_map),
+                _ => panic!("Should all be modules"),
+            })
+            .collect();
+        let source_coverage =
+            SourceCoverageBuilder::new(module, &coverage_map, source_map, packages);
         let source_coverage = source_coverage.compute_source_coverage(source_path);
         let output_result =
             source_coverage.output_source_coverage(&mut std::io::stdout(), self.color, self.tag);

third_party/move/move-compiler-v2/src/lib.rs

@@ -504,14 +504,7 @@ pub fn bytecode_pipeline(env: &GlobalEnv) -> FunctionTargetPipeline {
 pub fn disassemble_compiled_units(units: &[CompiledUnit]) -> anyhow::Result<String> {
     let disassembled_units: anyhow::Result<Vec<_>> = units
         .iter()
-        .map(|unit| {
-            let view = match unit {
-                CompiledUnit::Module(module) => BinaryIndexedView::Module(&module.module),
-                CompiledUnit::Script(script) => BinaryIndexedView::Script(&script.script),
-            };
-            Disassembler::from_view(view, location::Loc::new(FileHash::empty(), 0, 0))
-                .and_then(|d| d.disassemble())
-        })
+        .map(|unit| Disassembler::from_unit(unit).disassemble())
         .collect();
     Ok(disassembled_units?.concat())
 }

third_party/move/move-ir/types/src/location.rs

@@ -36,6 +36,7 @@ pub struct Loc {
 
 impl Loc {
     pub fn new(file_hash: FileHash, start: ByteIndex, end: ByteIndex) -> Loc {
+        assert!(start <= end);
         Loc {
             file_hash,
             start,
@@ -61,6 +62,94 @@ impl Loc {
             end: self.end as usize,
         }
     }
+
+    pub fn contains(&self, other: &Self) -> bool {
+        if self.file_hash != other.file_hash {
+            false
+        } else {
+            self.start <= other.start && other.end <= self.end
+        }
+    }
+
+    pub fn overlaps(&self, other: &Self) -> bool {
+        if self.file_hash != other.file_hash {
+            false
+        } else {
+            // [a, b] overlaps? [c, d]
+            // a <= b
+            // c <= d
+            other.start <= self.end && self.start <= other.end
+            // c <= b
+            // a <= d
+            // One of these:
+            // [a <= (c <= b] <= d)
+            // (c <= [a <= b] <= d)
+            // (c <= [a <= d) <= b]
+        }
+    }
+
+    pub fn overlaps_or_abuts(&self, other: &Self) -> bool {
+        if self.file_hash != other.file_hash {
+            false
+        } else {
+            // [a, b] overlaps? [c, d]
+            // a <= b
+            // c <= d
+            other.start <= self.end + 1 && self.start <= other.end + 1
+            // c <= b + 1
+            // a <= d + 1
+            // One of these:
+            // a <= c <= b+1 <= d+1
+            // c <= a <= b+1 <= d+1
+            // c <= a <= d+1 <= b+1
+        }
+    }
+
+    pub fn try_merge(&mut self, other: &Self) -> bool {
+        if self.overlaps_or_abuts(other) {
+            self.start = std::cmp::min(self.start, other.start);
+            self.end = std::cmp::max(self.end, other.end);
+            true
+        } else {
+            false
+        }
+    }
+
+    // if other overlaps with this, then subtract it out.
+    pub fn subtract(self, other: &Self) -> Vec<Loc> {
+        if !self.overlaps(other) {
+            vec![self]
+        } else {
+            if other.start <= self.start {
+                if self.end <= other.end {
+                    vec![]
+                } else {
+                    vec![Loc {
+                        start: other.end + 1,
+                        ..self
+                    }]
+                }
+            } else {
+                if self.end <= other.end {
+                    vec![Loc {
+                        end: other.start - 1,
+                        ..self
+                    }]
+                } else {
+                    vec![
+                        Loc {
+                            end: other.start - 1,
+                            ..self
+                        },
+                        Loc {
+                            start: other.end + 1,
+                            ..self
+                        },
+                    ]
+                }
+            }
+        }
+    }
 }
 
 impl PartialOrd for Loc {

third_party/move/move-model/src/model.rs

@@ -3337,6 +3337,7 @@ impl<'env> ModuleEnv<'env> {
             print_basic_blocks: true,
             print_locals: true,
             print_bytecode_stats: false,
+            show_locs: false,
         });
         Some(
             disas

third_party/move/tools/move-cli/src/base/coverage.rs

@@ -79,8 +79,17 @@ impl Coverage {
                     }) => (module, source_map),
                     _ => panic!("Should all be modules"),
                 };
+                let packages: Vec<_> = package
+                    .root_modules()
+                    .map(|unit| match &unit.unit {
+                        CompiledUnit::Module(NamedCompiledModule {
+                            module, source_map, ..
+                        }) => (module, source_map),
+                        _ => panic!("Should all be modules"),
+                    })
+                    .collect();
                 let source_coverage_builder =
-                    SourceCoverageBuilder::new(module, &coverage_map, source_map);
+                    SourceCoverageBuilder::new(module, &coverage_map, source_map, packages);
                 let source_coverage = source_coverage_builder.compute_source_coverage(source_path);
                 source_coverage
                     .output_source_coverage(&mut std::io::stdout(), self.color, self.tag)