Cranelift: add debug tag infrastructure.

cranelift/codegen/src/cursor.rs

@@ -2,7 +2,11 @@
 //!
 //! This module defines cursor data types that can be used for inserting instructions.
 
-use crate::ir;
+use crate::{
+    inst_predicates::has_lowering_side_effect,
+    ir::{self},
+};
+use alloc::vec::Vec;
 
 /// The possible positions of a cursor.
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
@@ -34,6 +38,9 @@ pub trait Cursor {
     /// Set the source location that should be assigned to new instructions.
     fn set_srcloc(&mut self, srcloc: ir::SourceLoc);
 
+    /// Set the debug tags that should be assigned to new side-effecting instructions.
+    fn set_debug_tags(&mut self, tags: Vec<ir::DebugTag>);
+
     /// Borrow a reference to the function layout that this cursor is navigating.
     fn layout(&self) -> &ir::Layout;
 
@@ -61,6 +68,30 @@ pub trait Cursor {
         self
     }
 
+    /// Exchange this cursor for one with a set debug tag list.
+    ///
+    /// These tags will be attached to all newly inserted
+    /// side-effecting instructions.
+    ///
+    /// This is intended to be used as a builder method:
+    ///
+    /// ```
+    /// # use cranelift_codegen::ir::{Function, Block, DebugTag};
+    /// # use cranelift_codegen::cursor::{Cursor, FuncCursor};
+    /// fn edit_func(func: &mut Function, tags: Vec<DebugTag>) {
+    ///     let mut pos = FuncCursor::new(func).with_debug_tags(tags);
+    ///
+    ///     // Use `pos`...
+    /// }
+    /// ```
+    fn with_debug_tags(mut self, tags: Vec<ir::DebugTag>) -> Self
+    where
+        Self: Sized,
+    {
+        self.set_debug_tags(tags);
+        self
+    }
+
     /// Rebuild this cursor positioned at `pos`.
     fn at_position(mut self, pos: CursorPosition) -> Self
     where
@@ -617,6 +648,7 @@ pub trait Cursor {
 pub struct FuncCursor<'f> {
     pos: CursorPosition,
     srcloc: ir::SourceLoc,
+    debug_tags: Vec<ir::DebugTag>,
 
     /// The referenced function.
     pub func: &'f mut ir::Function,
@@ -628,6 +660,7 @@ impl<'f> FuncCursor<'f> {
         Self {
             pos: CursorPosition::Nowhere,
             srcloc: Default::default(),
+            debug_tags: vec![],
             func,
         }
     }
@@ -661,6 +694,10 @@ impl<'f> Cursor for FuncCursor<'f> {
         self.srcloc = srcloc;
     }
 
+    fn set_debug_tags(&mut self, tags: Vec<ir::DebugTag>) {
+        self.debug_tags = tags;
+    }
+
     fn layout(&self) -> &ir::Layout {
         &self.func.layout
     }
@@ -684,6 +721,9 @@ impl<'c, 'f> ir::InstInserterBase<'c> for &'c mut FuncCursor<'f> {
         if !self.srcloc.is_default() {
             self.func.set_srcloc(inst, self.srcloc);
         }
+        if has_lowering_side_effect(self.func, inst) && !self.debug_tags.is_empty() {
+            self.func.debug_tags.set(inst, self.debug_tags.clone());
+        }
         &mut self.func.dfg
     }
 }

cranelift/codegen/src/inline.rs

@@ -20,7 +20,7 @@
 //! Cranelift the body of the callee that is to be inlined.
 
 use crate::cursor::{Cursor as _, FuncCursor};
-use crate::ir::{self, ExceptionTableData, ExceptionTableItem, InstBuilder as _};
+use crate::ir::{self, DebugTag, ExceptionTableData, ExceptionTableItem, InstBuilder as _};
 use crate::result::CodegenResult;
 use crate::trace;
 use crate::traversals::Dfs;
@@ -366,6 +366,10 @@ fn inline_one(
     // callee.
     let mut last_inlined_block = inline_block_layout(func, call_block, callee, &entity_map);
 
+    // Get a copy of debug tags on the call instruction; these are
+    // prepended to debug tags on inlined instructions.
+    let call_debug_tags = func.debug_tags.get(call_inst).to_vec();
+
     // Translate each instruction from the callee into the caller,
     // appending them to their associated block in the caller.
     //
@@ -403,6 +407,29 @@ fn inline_one(
             let inlined_inst = func.dfg.make_inst(inlined_inst_data);
             func.layout.append_inst(inlined_inst, inlined_block);
 
+            // Copy over debug tags, translating referenced entities
+            // as appropriate.
+            let debug_tags = callee.debug_tags.get(callee_inst);
+            // If there are tags on the inlined instruction, we always
+            // add tags, and we prepend any tags from the call
+            // instruction; but we don't add tags if only the callsite
+            // had them (this would otherwise mean that every single
+            // instruction in an inlined function body would get
+            // tags).
+            if !debug_tags.is_empty() {
+                let tags = call_debug_tags
+                    .iter()
+                    .cloned()
+                    .chain(debug_tags.iter().map(|tag| match *tag {
+                        DebugTag::User(value) => DebugTag::User(value),
+                        DebugTag::StackSlot(slot) => {
+                            DebugTag::StackSlot(entity_map.inlined_stack_slot(slot))
+                        }
+                    }))
+                    .collect::<SmallVec<[_; 4]>>();
+                func.debug_tags.set(inlined_inst, tags);
+            }
+
             let opcode = callee.dfg.insts[callee_inst].opcode();
             if opcode.is_return() {
                 // Instructions that return do not define any values, so we

cranelift/codegen/src/ir/debug_tags.rs

@@ -0,0 +1,127 @@
+//! Debug tag storage.
+//!
+//! Cranelift permits the embedder to place "debug tags" on
+//! instructions in CLIF. These tags are sequences of items of various
+//! kinds, with no other meaning imposed by Cranelift. They are passed
+//! through to metadata provided alongside the compilation result.
+//!
+//! When Cranelift inlines a function, it will prepend any tags from
+//! the call instruction at the inlining callsite to tags on all
+//! inlined instructions.
+//!
+//! These tags can be used, for example, to identify stackslots that
+//! store user state, or to denote positions in user source. In
+//! general, the intent is to allow perfect reconstruction of original
+//! (source-level) program state in an instrumentation-based
+//! debug-info scheme, as long as the instruction(s) on which these
+//! tags are attached are preserved. This will be the case for any
+//! instructions with side-effects.
+//!
+//! A few answers to design questions that lead to this design:
+//!
+//! - Why not use the SourceLoc mechanism? Debug tags are richer than
+//!   that infrastructure because they preserve inlining location and
+//!   are interleaved properly with any other tags describing the
+//!   frame.
+//! - Why not attach debug tags only to special sequence-point
+//!   instructions? This is driven by inlining: we should have the
+//!   semantic information about a callsite attached directly to the
+//!   call and observe it there, not have a magic "look backward to
+//!   find a sequence point" behavior in the inliner.
+//!
+//! In other words, the needs of preserving "virtual" frames across an
+//! inlining transform drive this design.
+
+use crate::ir::{Inst, StackSlot};
+use alloc::collections::BTreeMap;
+use alloc::vec::Vec;
+use core::ops::Range;
+
+/// Debug tags for instructions.
+#[derive(Clone, PartialEq, Hash, Default)]
+#[cfg_attr(
+    feature = "enable-serde",
+    derive(serde_derive::Serialize, serde_derive::Deserialize)
+)]
+pub struct DebugTags {
+    /// Pool of tags, referred to by `insts` below.
+    tags: Vec<DebugTag>,
+
+    /// Per-instruction range for its list of tags in the tag pool (if
+    /// any).
+    ///
+    /// Note: we don't use `PackedOption` and `EntityList` here
+    /// because the values that we are storing are not entities.
+    insts: BTreeMap<Inst, Range<u32>>,
+}
+
+/// One debug tag.
+#[derive(Clone, Debug, PartialEq, Hash)]
+#[cfg_attr(
+    feature = "enable-serde",
+    derive(serde_derive::Serialize, serde_derive::Deserialize)
+)]
+pub enum DebugTag {
+    /// User-specified `u32` value, opaque to Cranelift.
+    User(u32),
+
+    /// A stack slot reference.
+    StackSlot(StackSlot),
+}
+
+impl DebugTags {
+    /// Set the tags on an instruction, overwriting existing tag list.
+    pub fn set(&mut self, inst: Inst, tags: impl IntoIterator<Item = DebugTag>) {
+        let start = u32::try_from(self.tags.len()).unwrap();
+        self.tags.extend(tags);
+        let end = u32::try_from(self.tags.len()).unwrap();
+        if end > start {
+            self.insts.insert(inst, start..end);
+        } else {
+            self.insts.remove(&inst);
+        }
+    }
+
+    /// Get the tags associated with an instruction.
+    pub fn get(&self, inst: Inst) -> &[DebugTag] {
+        if let Some(range) = self.insts.get(&inst) {
+            let start = usize::try_from(range.start).unwrap();
+            let end = usize::try_from(range.end).unwrap();
+            &self.tags[start..end]
+        } else {
+            &[]
+        }
+    }
+
+    /// Clone the tags from one instruction to another.
+    ///
+    /// This clone is cheap (references the same underlying storage)
+    /// because the tag lists are immutable.
+    pub fn clone_tags(&mut self, from: Inst, to: Inst) {
+        if let Some(range) = self.insts.get(&from).cloned() {
+            self.insts.insert(to, range);
+        }
+    }
+
+    /// Are any debug tags present?
+    ///
+    /// This is used for adjusting margins when pretty-printing CLIF.
+    pub fn is_empty(&self) -> bool {
+        self.insts.is_empty()
+    }
+
+    /// Clear all tags.
+    pub fn clear(&mut self) {
+        self.insts.clear();
+        self.tags.clear();
+    }
+}
+
+impl core::fmt::Display for DebugTag {
+    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
+        match self {
+            DebugTag::User(value) => write!(f, "{value}"),
+            DebugTag::StackSlot(slot) => write!(f, "{slot}"),
+        }
+    }
+}

cranelift/codegen/src/ir/function.rs

@@ -5,6 +5,7 @@
 
 use crate::HashMap;
 use crate::entity::{PrimaryMap, SecondaryMap};
+use crate::ir::DebugTags;
 use crate::ir::{
     self, Block, DataFlowGraph, DynamicStackSlot, DynamicStackSlotData, DynamicStackSlots,
     DynamicType, ExtFuncData, FuncRef, GlobalValue, GlobalValueData, Inst, JumpTable,
@@ -190,6 +191,15 @@ pub struct FunctionStencil {
     /// interpreted by Cranelift, only preserved.
     pub srclocs: SourceLocs,
 
+    /// Opaque debug-info tags on instructions.
+    ///
+    /// These tags are not interpreted by Cranelift, and are passed
+    /// through to compilation-result metadata. The only semantic
+    /// structure that Cranelift imposes is that when inlining, it
+    /// prepends the callsite call instruction's tags to the tags on
+    /// inlined instructions.
+    pub debug_tags: DebugTags,
+
     /// An optional global value which represents an expression evaluating to
     /// the stack limit for this function. This `GlobalValue` will be
     /// interpreted in the prologue, if necessary, to insert a stack check to
@@ -209,6 +219,7 @@ impl FunctionStencil {
         self.dfg.clear();
         self.layout.clear();
         self.srclocs.clear();
+        self.debug_tags.clear();
         self.stack_limit = None;
     }
 
@@ -408,6 +419,7 @@ impl Function {
                 layout: Layout::new(),
                 srclocs: SecondaryMap::new(),
                 stack_limit: None,
+                debug_tags: DebugTags::default(),
             },
             params: FunctionParameters::new(),
         }

cranelift/codegen/src/ir/layout.rs

@@ -756,7 +756,7 @@ mod tests {
     use super::*;
     use crate::cursor::{Cursor, CursorPosition};
     use crate::entity::EntityRef;
-    use crate::ir::{Block, Inst, SourceLoc};
+    use crate::ir::{Block, DebugTag, Inst, SourceLoc};
     use alloc::vec::Vec;
     use core::cmp::Ordering;
 
@@ -795,6 +795,10 @@ mod tests {
             unimplemented!()
         }
 
+        fn set_debug_tags(&mut self, _tags: Vec<DebugTag>) {
+            unimplemented!()
+        }
+
         fn layout(&self) -> &Layout {
             self.layout
         }

cranelift/codegen/src/ir/mod.rs

@@ -4,6 +4,7 @@ mod atomic_rmw_op;
 mod builder;
 pub mod condcodes;
 pub mod constant;
+mod debug_tags;
 pub mod dfg;
 pub mod dynamic_type;
 pub mod entities;
@@ -36,6 +37,7 @@ pub use crate::ir::builder::{
     InsertBuilder, InstBuilder, InstBuilderBase, InstInserterBase, ReplaceBuilder,
 };
 pub use crate::ir::constant::{ConstantData, ConstantPool};
+pub use crate::ir::debug_tags::{DebugTag, DebugTags};
 pub use crate::ir::dfg::{BlockData, DataFlowGraph, ValueDef};
 pub use crate::ir::dynamic_type::{DynamicTypeData, DynamicTypes, dynamic_to_fixed};
 pub use crate::ir::entities::{