Fixing source map locations when peephole optimizations are applied.

third_party/move/move-compiler-v2/src/file_format_generator/function_generator.rs

@@ -159,8 +159,15 @@ impl<'a> FunctionGenerator<'a> {
                     .get_extension::<Options>()
                     .expect("Options is available");
                 if options.experiment_on(Experiment::PEEPHOLE_OPTIMIZATION) {
-                    // TODO: fix source mapping (#14167)
-                    peephole_optimizer::run(&mut code);
+                    let transformed_code_chunk = peephole_optimizer::optimize(&code.code);
+                    // Fix the source map for the optimized code.
+                    fun_gen
+                        .gen
+                        .source_map
+                        .remap_code_map(def_idx, transformed_code_chunk.original_offsets)
+                        .expect(SOURCE_MAP_OK);
+                    // Replace the code with the optimized one.
+                    code.code = transformed_code_chunk.code;
                 }
             } else {
                 // Write the spec block table back to the environment.

third_party/move/move-compiler-v2/src/file_format_generator/peephole_optimizer.rs

@@ -12,17 +12,18 @@ pub mod reducible_pairs;
 use inefficient_loads::InefficientLoads;
 use move_binary_format::{
     control_flow_graph::{ControlFlowGraph, VMControlFlowGraph},
-    file_format::{Bytecode, CodeOffset, CodeUnit},
+    file_format::{Bytecode, CodeOffset},
 };
-use optimizers::{BasicBlockOptimizer, WindowProcessor};
+use optimizers::{BasicBlockOptimizer, TransformedCodeChunk, WindowProcessor};
 use reducible_pairs::ReduciblePairs;
-use std::{collections::BTreeMap, mem};
+use std::collections::BTreeMap;
 
 /// Pre-requisite: `code` should not have spec block associations.
 /// Run peephole optimizers on the given `code`, possibly modifying it.
-pub fn run(code: &mut CodeUnit) {
-    let original_code = mem::take(&mut code.code);
-    code.code = BasicBlockOptimizerPipeline::default().optimize(original_code);
+/// Returns the optimized code, along with mapping to original offsets
+/// in `code`.
+pub fn optimize(code: &[Bytecode]) -> TransformedCodeChunk {
+    BasicBlockOptimizerPipeline::default().optimize(code)
 }
 
 /// A pipeline of basic block optimizers.
@@ -44,23 +45,24 @@ impl BasicBlockOptimizerPipeline {
 
     /// Run the basic block optimization pipeline on the given `code`,
     /// returning new (possibly optimized) code.
-    pub fn optimize(&self, mut code: Vec<Bytecode>) -> Vec<Bytecode> {
-        let mut cfg = VMControlFlowGraph::new(&code);
+    pub fn optimize(&self, code: &[Bytecode]) -> TransformedCodeChunk {
+        let mut code_chunk = TransformedCodeChunk::make_from(code);
+        let mut cfg = VMControlFlowGraph::new(&code_chunk.code);
         loop {
-            let optimized_blocks = self.get_optimized_blocks(&code, &cfg);
-            let optimized_code = Self::flatten_blocks(optimized_blocks);
-            let optimized_cfg = VMControlFlowGraph::new(&optimized_code);
+            let optimized_blocks = self.get_optimized_blocks(&code_chunk.code, &cfg);
+            let optimized_code_chunk = Self::flatten_blocks(optimized_blocks);
+            let optimized_cfg = VMControlFlowGraph::new(&optimized_code_chunk.code);
             if optimized_cfg.num_blocks() == cfg.num_blocks() {
                 // Proxy for convergence of basic block optimizations.
                 // This is okay for peephole optimizations that merge basic blocks.
                 // But may need to revisit if we have peephole optimizations that can
                 // split a basic block.
-                return optimized_code;
+                return optimized_code_chunk;
             } else {
                 // Number of basic blocks changed, re-run the basic-block
                 // optimization pipeline again on the new basic blocks.
                 cfg = optimized_cfg;
-                code = optimized_code;
+                code_chunk = optimized_code_chunk.remap(code_chunk.original_offsets);
             }
         }
     }
@@ -71,29 +73,33 @@ impl BasicBlockOptimizerPipeline {
         &self,
         code: &[Bytecode],
         cfg: &VMControlFlowGraph,
-    ) -> BTreeMap<CodeOffset, Vec<Bytecode>> {
+    ) -> BTreeMap<CodeOffset, TransformedCodeChunk> {
         let mut optimized_blocks = BTreeMap::new();
         for block_id in cfg.blocks() {
             let start = cfg.block_start(block_id);
             let end = cfg.block_end(block_id); // `end` is inclusive
-            let mut block = code[start as usize..=end as usize].to_vec();
+            let mut block = TransformedCodeChunk::make_from(&code[start as usize..=end as usize]);
             for bb_optimizer in self.optimizers.iter() {
-                block = bb_optimizer.optimize(&block);
+                block = bb_optimizer
+                    .optimize(&block.code)
+                    .remap(block.original_offsets);
             }
             optimized_blocks.insert(start, block);
         }
         optimized_blocks
     }
 
     /// Flatten the individually optimized basic blocks into a single code vector.
-    fn flatten_blocks(optimized_blocks: BTreeMap<CodeOffset, Vec<Bytecode>>) -> Vec<Bytecode> {
-        let mut optimized_code = vec![];
+    fn flatten_blocks(
+        optimized_blocks: BTreeMap<CodeOffset, TransformedCodeChunk>,
+    ) -> TransformedCodeChunk {
+        let mut optimized_code = TransformedCodeChunk::empty();
         let mut block_mapping = BTreeMap::new();
-        for (offset, mut block) in optimized_blocks {
-            block_mapping.insert(offset, optimized_code.len() as CodeOffset);
-            optimized_code.append(&mut block);
+        for (offset, block) in optimized_blocks {
+            block_mapping.insert(offset, optimized_code.code.len() as CodeOffset);
+            optimized_code.extend(block, offset);
         }
-        Self::remap_branch_targets(&mut optimized_code, &block_mapping);
+        Self::remap_branch_targets(&mut optimized_code.code, &block_mapping);
         optimized_code
     }
 

third_party/move/move-compiler-v2/src/file_format_generator/peephole_optimizer/inefficient_loads.rs

@@ -24,8 +24,11 @@
 //! cannot use it without a subsequent store.
 //! So, skipping the store to `u` is safe.
 
-use crate::file_format_generator::peephole_optimizer::optimizers::WindowOptimizer;
-use move_binary_format::file_format::Bytecode;
+use crate::file_format_generator::peephole_optimizer::optimizers::{
+    TransformedCodeChunk, WindowOptimizer,
+};
+use move_binary_format::file_format::{Bytecode, CodeOffset};
+use std::iter;
 
 /// An optimizer for inefficient loads.
 pub struct InefficientLoads;
@@ -37,7 +40,7 @@ impl InefficientLoads {
 }
 
 impl WindowOptimizer for InefficientLoads {
-    fn optimize_window(&self, window: &[Bytecode]) -> Option<(Vec<Bytecode>, usize)> {
+    fn optimize_window(&self, window: &[Bytecode]) -> Option<(TransformedCodeChunk, usize)> {
         use Bytecode::*;
         if window.len() < Self::MIN_WINDOW_SIZE {
             return None;
@@ -61,8 +64,12 @@ impl WindowOptimizer for InefficientLoads {
                     // We have reached the end of the pattern (point 4 in the module documentation).
                     let sequence = &window[2..index + 2];
                     let load_constant = &window[0..1];
+                    let transformed_code = [sequence, load_constant].concat();
+                    let original_offsets = (2..(index + 2) as CodeOffset)
+                        .chain(iter::once(0))
+                        .collect::<Vec<_>>();
                     return Some((
-                        [sequence, load_constant].concat(),
+                        TransformedCodeChunk::new(transformed_code, original_offsets),
                         index + Self::MIN_WINDOW_SIZE,
                     ));
                 },

third_party/move/move-compiler-v2/src/file_format_generator/peephole_optimizer/optimizers.rs

@@ -3,33 +3,91 @@
 
 //! This module contains setup for basic block peephole optimizers.
 
-use move_binary_format::file_format::Bytecode;
+use move_binary_format::file_format::{Bytecode, CodeOffset};
+
+/// A contiguous chunk of bytecode that may have been transformed from some
+/// other "original" contiguous chunk of bytecode.
+pub struct TransformedCodeChunk {
+    /// The transformed bytecode.
+    pub code: Vec<Bytecode>,
+    /// Mapping to the original offsets.
+    /// The instruction in `code[i]` corresponds to the instruction at
+    /// `original_offsets[i]` in the original bytecode.
+    pub original_offsets: Vec<CodeOffset>,
+}
+
+impl TransformedCodeChunk {
+    /// Create an instance of `TransformedCodeChunk` from the given `code`
+    /// and `original_offsets`.
+    pub fn new(code: Vec<Bytecode>, original_offsets: Vec<CodeOffset>) -> Self {
+        debug_assert_eq!(code.len(), original_offsets.len());
+        Self {
+            code,
+            original_offsets,
+        }
+    }
+
+    /// Create an empty chunk.
+    pub fn empty() -> Self {
+        Self::new(vec![], vec![])
+    }
+
+    /// Extend this chunk with another `other` chunk.
+    /// The `original_offsets` for the `other` chunk are incremented by `adjust`.
+    pub fn extend(&mut self, other: TransformedCodeChunk, adjust: CodeOffset) {
+        self.code.extend(other.code);
+        self.original_offsets
+            .extend(other.original_offsets.into_iter().map(|off| off + adjust));
+    }
+
+    /// Make a new chunk from the given `code`.
+    pub fn make_from(code: &[Bytecode]) -> Self {
+        Self::new(code.to_vec(), Vec::from_iter(0..code.len() as CodeOffset))
+    }
+
+    /// Remap the original offsets using the given `previous_offsets`.
+    pub fn remap(self, previous_offsets: Vec<CodeOffset>) -> Self {
+        Self::new(
+            self.code,
+            self.original_offsets
+                .into_iter()
+                .map(|off| previous_offsets[off as usize])
+                .collect(),
+        )
+    }
+}
 
 /// A basic block optimizer that optimizes a basic block of bytecode.
 pub trait BasicBlockOptimizer {
-    /// Given a basic `block`, return its optimized version.
-    fn optimize(&self, block: &[Bytecode]) -> Vec<Bytecode>;
+    /// Given a basic `block`, return its optimized version [*].
+    ///
+    /// [*] The optimized version returned via `TransformedCodeChunk` maintains a
+    /// mapping to the original offsets in `block`.
+    fn optimize(&self, block: &[Bytecode]) -> TransformedCodeChunk;
 }
 
 /// An optimizer for a window of bytecode within a basic block.
 /// The window is always a suffix of a basic block.
 pub trait WindowOptimizer {
     /// Given a `window` of bytecode, return a tuple containing:
-    ///   1. an optimized version of a non-empty prefix of the `window`.
+    ///   1. an optimized version of a non-empty prefix of the `window`. [*]
     ///   2. size of this prefix (should be non-zero).
     /// If `None` is returned, the `window` is not optimized.
-    fn optimize_window(&self, window: &[Bytecode]) -> Option<(Vec<Bytecode>, usize)>;
+    ///
+    /// [*] When an optimized version is returned, the corresponding `TransformedCodeChunk`
+    /// maintains a mapping to the original offsets of `window`.
+    fn optimize_window(&self, window: &[Bytecode]) -> Option<(TransformedCodeChunk, usize)>;
 }
 
 /// A processor to perform window optimizations of a particular style on a basic block.
 pub struct WindowProcessor<T: WindowOptimizer>(T);
 
 impl<T: WindowOptimizer> BasicBlockOptimizer for WindowProcessor<T> {
-    fn optimize(&self, block: &[Bytecode]) -> Vec<Bytecode> {
-        let mut old_block = block.to_vec();
+    fn optimize(&self, block: &[Bytecode]) -> TransformedCodeChunk {
+        let mut old_block = TransformedCodeChunk::make_from(block);
         // Run single passes until code stops changing.
-        while let Some(new_block) = self.optimize_single_pass(&old_block) {
-            old_block = new_block;
+        while let Some(new_chunk) = self.optimize_single_pass(&old_block.code) {
+            old_block = new_chunk.remap(old_block.original_offsets);
         }
         old_block
     }
@@ -43,19 +101,22 @@ impl<T: WindowOptimizer> WindowProcessor<T> {
 
     /// Run a single pass of the window peephole optimization on the given basic `block`.
     /// If the block cannot be optimized further, return `None`.
-    fn optimize_single_pass(&self, block: &[Bytecode]) -> Option<Vec<Bytecode>> {
+    fn optimize_single_pass(&self, block: &[Bytecode]) -> Option<TransformedCodeChunk> {
         let mut changed = false;
-        let mut new_block: Vec<Bytecode> = vec![];
+        let mut new_block = TransformedCodeChunk::empty();
         let mut left = 0;
         while left < block.len() {
             let window = &block[left..];
             if let Some((optimized_window, consumed)) = self.0.optimize_window(window) {
                 debug_assert!(consumed != 0);
-                new_block.extend(optimized_window);
+                new_block.extend(optimized_window, left as CodeOffset);
                 left += consumed;
                 changed = true;
             } else {
-                new_block.push(block[left].clone());
+                new_block.extend(
+                    TransformedCodeChunk::make_from(&block[left..left + 1]),
+                    left as CodeOffset,
+                );
                 left += 1;
             }
         }

third_party/move/move-compiler-v2/src/file_format_generator/peephole_optimizer/reducible_pairs.rs

@@ -39,7 +39,9 @@
 //! Finally, note that fixed window optimizations are performed on windows within a basic
 //! block, not spanning across multiple basic blocks.
 
-use crate::file_format_generator::peephole_optimizer::optimizers::WindowOptimizer;
+use crate::file_format_generator::peephole_optimizer::optimizers::{
+    TransformedCodeChunk, WindowOptimizer,
+};
 use move_binary_format::file_format::Bytecode;
 
 pub struct ReduciblePairs;
@@ -49,7 +51,7 @@ impl ReduciblePairs {
 }
 
 impl WindowOptimizer for ReduciblePairs {
-    fn optimize_window(&self, window: &[Bytecode]) -> Option<(Vec<Bytecode>, usize)> {
+    fn optimize_window(&self, window: &[Bytecode]) -> Option<(TransformedCodeChunk, usize)> {
         use Bytecode::*;
         if window.len() < Self::WINDOW_SIZE {
             return None;
@@ -59,13 +61,17 @@ impl WindowOptimizer for ReduciblePairs {
             (StLoc(u), MoveLoc(v)) | (CopyLoc(u), StLoc(v)) | (MoveLoc(u), StLoc(v))
                 if *u == *v =>
             {
-                vec![]
+                TransformedCodeChunk::new(vec![], vec![])
             },
-            (CopyLoc(_), Pop) => vec![],
-            (LdTrue, BrTrue(target)) | (LdFalse, BrFalse(target)) => vec![Branch(*target)],
-            (LdTrue, BrFalse(_)) | (LdFalse, BrTrue(_)) => vec![],
-            (Not, BrFalse(target)) => vec![BrTrue(*target)],
-            (Not, BrTrue(target)) => vec![BrFalse(*target)],
+            (CopyLoc(_), Pop) => TransformedCodeChunk::new(vec![], vec![]),
+            (LdTrue, BrTrue(target)) | (LdFalse, BrFalse(target)) => {
+                TransformedCodeChunk::new(vec![Branch(*target)], vec![0])
+            },
+            (LdTrue, BrFalse(_)) | (LdFalse, BrTrue(_)) => {
+                TransformedCodeChunk::new(vec![], vec![])
+            },
+            (Not, BrFalse(target)) => TransformedCodeChunk::new(vec![BrTrue(*target)], vec![0]),
+            (Not, BrTrue(target)) => TransformedCodeChunk::new(vec![BrFalse(*target)], vec![0]),
             _ => return None,
         };
         Some((optimized, Self::WINDOW_SIZE))

third_party/move/move-compiler-v2/tests/reference-safety/v1-borrow-tests/writeref_borrow_invalid.old.exp

@@ -4,7 +4,7 @@
 Diagnostics:
 bug: bytecode verification failed with unexpected status code `WRITEREF_EXISTS_BORROW_ERROR`. This is a compiler bug, consider reporting it.
 Error message: none
-   ┌─ tests/reference-safety/v1-borrow-tests/writeref_borrow_invalid.move:10:18
+   ┌─ tests/reference-safety/v1-borrow-tests/writeref_borrow_invalid.move:10:9
    │
 10 │         *g_mut = G { v: 0 };
-   │                  ^^^^^^^^^^
+   │         ^^^^^^^^^^^^^^^^^^^

third_party/move/move-ir-compiler/move-bytecode-source-map/src/source_map.rs

@@ -202,6 +202,29 @@ impl FunctionSourceMap {
         };
     }
 
+    /// Remap the code map based on the given `remap`.
+    /// If `remap[i] == j`, then the code location associated with code offset `j`
+    /// will now be associated with code offset `i`.
+    pub fn remap_code_map(&mut self, remap: Vec<CodeOffset>) {
+        let remapped_locations = remap
+            .iter()
+            .map(|old_offset| self.get_code_location(*old_offset))
+            .collect::<Vec<_>>();
+        self.code_map.clear();
+        let mut prev_loc = None;
+        for (new_offset, loc) in remapped_locations.into_iter().enumerate() {
+            if prev_loc == loc {
+                // optimization: if a series of instructions map to the same location, we only need
+                // to add code mapping for the first one.
+                continue;
+            }
+            if let Some(loc) = loc {
+                self.add_code_mapping(new_offset as CodeOffset, loc);
+                prev_loc = Some(loc);
+            }
+        }
+    }
+
     /// Record the code offset for an Nop label
     pub fn add_nop_mapping(&mut self, label: NopLabel, offset: CodeOffset) {
         assert!(self.nops.insert(label, offset).is_none())
@@ -362,6 +385,22 @@ impl SourceMap {
         Ok(())
     }
 
+    /// Remap the code map for the function given by `fdef_idx`, according to `remap`.
+    /// If `remap[i] == j`, then the code location associated with code offset `j`
+    /// will now be associated with code offset `i`.
+    pub fn remap_code_map(
+        &mut self,
+        fdef_idx: FunctionDefinitionIndex,
+        remap: Vec<CodeOffset>,
+    ) -> Result<()> {
+        let func_entry = self
+            .function_map
+            .get_mut(&fdef_idx.0)
+            .ok_or_else(|| format_err!("Tried to remap code map of undefined function index"))?;
+        func_entry.remap_code_map(remap);
+        Ok(())
+    }
+
     pub fn add_nop_mapping(
         &mut self,
         fdef_idx: FunctionDefinitionIndex,