MachInst backend: pass through SourceLoc information.

This change adds SourceLoc information per instruction in a `VCode<Inst>`
container, and keeps this information up-to-date across register allocation
and branch reordering. The information is initially collected during
instruction lowering, eventually collected on the MachSection, and finally
provided to the environment that wraps the codegen crate for wasmtime.

This PR is based on top of bytecodealliance#1570 and bytecodealliance#1571 (part of a series fixing tests).

This PR depends on wasmtime/regalloc.rs#50, a change to the register
allocator to provide instruction-granularity info on the rewritten
instruction stream (rather than block-granularity).

With the prior PRs applied as well, quite a few more unit tests pass;
the exclusion list in bytecodealliance#1526 should be updated if this PR lands first.

cranelift/codegen/src/machinst/lower.rs

@@ -371,6 +371,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
                             "lower_branch_group: targets = {:?} branches = {:?}",
                             targets, branches
                         );
+                        self.vcode.set_srcloc(Some(self.srcloc(branches[0])));
                         backend.lower_branch_group(
                             &mut self,
                             &branches[..],
@@ -387,6 +388,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
                     let num_uses = self.num_uses[inst];
                     let side_effect = has_side_effect(self.f, inst);
                     if side_effect || num_uses > 0 {
+                        self.vcode.set_srcloc(Some(self.srcloc(inst)));
                         backend.lower(&mut self, inst);
                         self.vcode.end_ir_inst();
                     } else {
@@ -415,6 +417,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
                     "lower_branch_group: targets = {:?} branches = {:?}",
                     targets, branches
                 );
+                self.vcode.set_srcloc(Some(self.srcloc(branches[0])));
                 backend.lower_branch_group(&mut self, &branches[..], &targets[..], fallthrough);
                 self.vcode.end_ir_inst();
                 branches.clear();

cranelift/codegen/src/machinst/mod.rs

@@ -97,7 +97,6 @@
 //! ```
 
 use crate::binemit::{CodeInfo, CodeOffset};
-use crate::entity::SecondaryMap;
 use crate::ir::condcodes::IntCC;
 use crate::ir::{Function, Type};
 use crate::result::CodegenResult;

cranelift/codegen/src/machinst/sections.rs

@@ -74,6 +74,15 @@ impl MachSections {
         sink.end_codegen();
     }
 
+    /// Get a list of source location mapping tuples.
+    pub fn get_srclocs(&self) -> Vec<MachSrcLoc> {
+        let mut ret = vec![];
+        for section in &self.sections {
+            ret.extend(section.srclocs.clone().into_iter());
+        }
+        ret
+    }
+
     /// Get the total required size for these sections.
     pub fn total_size(&self) -> CodeOffset {
         if self.sections.len() == 0 {
@@ -143,6 +152,12 @@ pub trait MachSectionOutput {
     /// Add a call return address record at the current offset.
     fn add_call_site(&mut self, loc: SourceLoc, opcode: Opcode);
 
+    /// Start the output for the given source-location at the current offset.
+    fn start_srcloc(&mut self, loc: SourceLoc);
+
+    /// End the output for the previously-given source-location at the current offset.
+    fn end_srcloc(&mut self);
+
     /// Align up to the given alignment.
     fn align_to(&mut self, align_to: CodeOffset) {
         assert!(align_to.is_power_of_two());
@@ -168,8 +183,13 @@ pub struct MachSection {
     pub relocs: Vec<MachReloc>,
     /// Any trap records referring to this section.
     pub traps: Vec<MachTrap>,
-    /// Any call site record referring to this section.
+    /// Any call site records referring to this section.
     pub call_sites: Vec<MachCallSite>,
+    /// Any source location mappings referring to this section.
+    pub srclocs: Vec<MachSrcLoc>,
+    /// The curren source location in progress (after `start_srcloc()` and before `end_srcloc()`).
+    /// This is a (start_offset, src_loc) tuple.
+    pub cur_srcloc: Option<(CodeOffset, SourceLoc)>,
 }
 
 impl MachSection {
@@ -182,6 +202,8 @@ impl MachSection {
             relocs: vec![],
             traps: vec![],
             call_sites: vec![],
+            srclocs: vec![],
+            cur_srcloc: None,
         }
     }
 
@@ -266,6 +288,22 @@ impl MachSectionOutput for MachSection {
             opcode,
         });
     }
+
+    fn start_srcloc(&mut self, loc: SourceLoc) {
+        self.cur_srcloc = Some((self.cur_offset_from_start(), loc));
+    }
+
+    fn end_srcloc(&mut self) {
+        if let Some((start, loc)) = self.cur_srcloc.take() {
+            let end = self.cur_offset_from_start();
+            // Skip zero-length extends.
+            if end > start {
+                self.srclocs.push(MachSrcLoc { start, end, loc });
+            }
+        } else {
+            panic!("end_srcloc() called without start_srcloc()");
+        }
+    }
 }
 
 /// A MachSectionOutput implementation that records only size.
@@ -315,6 +353,10 @@ impl MachSectionOutput for MachSectionSize {
     fn add_trap(&mut self, _: SourceLoc, _: TrapCode) {}
 
     fn add_call_site(&mut self, _: SourceLoc, _: Opcode) {}
+
+    fn start_srcloc(&mut self, _: SourceLoc) {}
+
+    fn end_srcloc(&mut self) {}
 }
 
 /// A relocation resulting from a compilation.
@@ -352,3 +394,16 @@ pub struct MachCallSite {
     /// The call's opcode.
     pub opcode: Opcode,
 }
+
+/// A source-location mapping resulting from a compilation.
+#[derive(Clone, Debug)]
+pub struct MachSrcLoc {
+    /// The start of the region of code corresponding to a source location.
+    /// This is absolute, not relative to the start of the section.
+    pub start: CodeOffset,
+    /// The end of the  region of code corresponding to a source location.
+    /// This is absolute, not relative to the start of the section.
+    pub end: CodeOffset,
+    /// The source location.
+    pub loc: SourceLoc,
+}

cranelift/codegen/src/machinst/vcode.rs

@@ -17,7 +17,9 @@
 //! See the main module comment in `mod.rs` for more details on the VCode-based
 //! backend pipeline.
 
+use crate::entity::SecondaryMap;
 use crate::ir;
+use crate::ir::SourceLoc;
 use crate::machinst::*;
 use crate::settings;
 
@@ -59,6 +61,10 @@ pub struct VCode<I: VCodeInst> {
     /// Lowered machine instructions in order corresponding to the original IR.
     insts: Vec<I>,
 
+    /// Source locations for each instruction. (`SourceLoc` is a `u32`, so it is
+    /// reasonable to keep one of these per instruction.)
+    srclocs: Vec<Option<SourceLoc>>,
+
     /// Entry block.
     entry: BlockIndex,
 
@@ -115,13 +121,16 @@ pub struct VCodeBuilder<I: VCodeInst> {
 
     /// Current basic block instructions, in reverse order (because blocks are
     /// built bottom-to-top).
-    bb_insns: SmallVec<[I; 32]>,
+    bb_insns: SmallVec<[(I, Option<SourceLoc>); 32]>,
 
     /// Current IR-inst instructions, in forward order.
-    ir_inst_insns: SmallVec<[I; 4]>,
+    ir_inst_insns: SmallVec<[(I, Option<SourceLoc>); 4]>,
 
     /// Start of succs for the current block in the concatenated succs list.
     succ_start: usize,
+
+    /// Current source location.
+    cur_srcloc: Option<SourceLoc>,
 }
 
 impl<I: VCodeInst> VCodeBuilder<I> {
@@ -133,6 +142,7 @@ impl<I: VCodeInst> VCodeBuilder<I> {
             bb_insns: SmallVec::new(),
             ir_inst_insns: SmallVec::new(),
             succ_start: 0,
+            cur_srcloc: None,
         }
     }
 
@@ -179,8 +189,8 @@ impl<I: VCodeInst> VCodeBuilder<I> {
     /// End the current IR instruction. Must be called after pushing any
     /// instructions and prior to ending the basic block.
     pub fn end_ir_inst(&mut self) {
-        while let Some(i) = self.ir_inst_insns.pop() {
-            self.bb_insns.push(i);
+        while let Some(pair) = self.ir_inst_insns.pop() {
+            self.bb_insns.push(pair);
         }
     }
 
@@ -191,8 +201,9 @@ impl<I: VCodeInst> VCodeBuilder<I> {
         let block_num = self.vcode.block_ranges.len() as BlockIndex;
         // Push the instructions.
         let start_idx = self.vcode.insts.len() as InsnIndex;
-        while let Some(i) = self.bb_insns.pop() {
+        while let Some((i, loc)) = self.bb_insns.pop() {
             self.vcode.insts.push(i);
+            self.vcode.srclocs.push(loc);
         }
         let end_idx = self.vcode.insts.len() as InsnIndex;
         // Add the instruction index range to the list of blocks.
@@ -224,7 +235,12 @@ impl<I: VCodeInst> VCodeBuilder<I> {
                 }
             }
         }
-        self.ir_inst_insns.push(insn);
+        self.ir_inst_insns.push((insn, self.cur_srcloc));
+    }
+
+    /// Set the current source location.
+    pub fn set_srcloc(&mut self, srcloc: Option<SourceLoc>) {
+        self.cur_srcloc = srcloc;
     }
 
     /// Build the final VCode.
@@ -286,6 +302,7 @@ impl<I: VCodeInst> VCode<I> {
             liveouts: abi.liveouts(),
             vreg_types: vec![],
             insts: vec![],
+            srclocs: vec![],
             entry: 0,
             block_ranges: vec![],
             block_succ_range: vec![],
@@ -349,14 +366,18 @@ impl<I: VCodeInst> VCode<I> {
             block_ranges(result.target_map.elems(), result.insns.len());
         let mut final_insns = vec![];
         let mut final_block_ranges = vec![(0, 0); self.num_blocks()];
+        let mut final_srclocs = vec![];
 
         for block in &self.final_block_order {
             let (start, end) = block_ranges[*block as usize];
             let final_start = final_insns.len() as InsnIndex;
 
             if *block == self.entry {
                 // Start with the prologue.
-                final_insns.extend(self.abi.gen_prologue().into_iter());
+                let prologue = self.abi.gen_prologue();
+                let len = prologue.len();
+                final_insns.extend(prologue.into_iter());
+                final_srclocs.extend(iter::repeat(None).take(len));
             }
 
             for i in start..end {
@@ -368,21 +389,34 @@ impl<I: VCodeInst> VCode<I> {
                     continue;
                 }
 
+                // Is there a srcloc associated with this insn? Look it up based on original
+                // instruction index (if new insn corresponds to some original insn, i.e., is not
+                // an inserted load/spill/move).
+                let srcloc: Option<SourceLoc> = result.orig_insn_map[InstIx::new(i as u32)]
+                    .and_then(|orig_idx| self.srclocs[orig_idx.get() as usize]);
+
                 // Whenever encountering a return instruction, replace it
                 // with the epilogue.
                 let is_ret = insn.is_term() == MachTerminator::Ret;
                 if is_ret {
-                    final_insns.extend(self.abi.gen_epilogue().into_iter());
+                    let epilogue = self.abi.gen_epilogue();
+                    let len = epilogue.len();
+                    final_insns.extend(epilogue.into_iter());
+                    final_srclocs.extend(iter::repeat(srcloc).take(len));
                 } else {
                     final_insns.push(insn.clone());
+                    final_srclocs.push(srcloc);
                 }
             }
 
             let final_end = final_insns.len() as InsnIndex;
             final_block_ranges[*block as usize] = (final_start, final_end);
         }
 
+        debug_assert!(final_insns.len() == final_srclocs.len());
+
         self.insts = final_insns;
+        self.srclocs = final_srclocs;
         self.block_ranges = final_block_ranges;
     }
 
@@ -512,6 +546,7 @@ impl<I: VCodeInst> VCode<I> {
         let code_section = sections.get_section(code_idx);
 
         let flags = self.abi.flags();
+        let mut cur_srcloc: Option<SourceLoc> = None;
         for &block in &self.final_block_order {
             let new_offset = I::align_basic_block(code_section.cur_offset_from_start());
             while new_offset > code_section.cur_offset_from_start() {
@@ -523,8 +558,24 @@ impl<I: VCodeInst> VCode<I> {
 
             let (start, end) = self.block_ranges[block as usize];
             for iix in start..end {
+                let srcloc = self.srclocs[iix as usize];
+                if srcloc != cur_srcloc {
+                    if cur_srcloc.is_some() {
+                        code_section.end_srcloc();
+                    }
+                    if srcloc.is_some() {
+                        code_section.start_srcloc(srcloc.unwrap());
+                    }
+                    cur_srcloc = srcloc;
+                }
+
                 self.insts[iix as usize].emit(code_section, flags);
             }
+
+            if cur_srcloc.is_some() {
+                code_section.end_srcloc();
+                cur_srcloc = None;
+            }
         }
 
         sections

crates/environ/src/cranelift.rs

@@ -8,6 +8,7 @@ use crate::compilation::{
 use crate::func_environ::{get_func_name, FuncEnvironment};
 use crate::{CacheConfig, FunctionBodyData, ModuleLocal, ModuleTranslation, Tunables};
 use cranelift_codegen::ir::{self, ExternalName};
+use cranelift_codegen::machinst::sections::MachSrcLoc;
 use cranelift_codegen::print_errors::pretty_error;
 use cranelift_codegen::{binemit, isa, Context};
 use cranelift_entity::PrimaryMap;
@@ -121,13 +122,24 @@ fn get_function_address_map<'data>(
 ) -> FunctionAddressMap {
     let mut instructions = Vec::new();
 
-    let func = &context.func;
-    let mut blocks = func.layout.blocks().collect::<Vec<_>>();
-    blocks.sort_by_key(|block| func.offsets[*block]); // Ensure inst offsets always increase
+    if let Some(ref mcr) = &context.mach_compile_result {
+        // New-style backend: we have a `MachCompileResult` that will give us `MachSrcLoc` mapping
+        // tuples.
+        let mut srclocs = mcr.sections.get_srclocs();
+        srclocs.sort_by_key(|mach_src_loc| mach_src_loc.start);
+        for MachSrcLoc { start, end, loc } in srclocs.into_iter() {
+            instructions.push(InstructionAddressMap {
+                srcloc: loc,
+                code_offset: start as usize,
+                code_len: (end - start) as usize,
+            });
+        }
+    } else {
+        // Old-style backend: we need to traverse the instruction/encoding info in the function.
+        let func = &context.func;
+        let mut blocks = func.layout.blocks().collect::<Vec<_>>();
+        blocks.sort_by_key(|block| func.offsets[*block]); // Ensure inst offsets always increase
 
-    // FIXME(#1523): New backend does not support debug info or instruction-address mapping
-    // yet.
-    if !isa.get_mach_backend().is_some() {
         let encinfo = isa.encoding_info();
         for block in blocks {
             for (offset, inst, size) in func.inst_offsets(block, &encinfo) {