Codegen: Align representation of stackmap with SpiderMonkey

This commit aligns the representation of stackmaps to be the same
as Spidermonkey's by:
 * Reversing the order of the bitmap from low addresses to high addresses
 * Including incoming stack arguments
 * Excluding outgoing stack arguments

Additionally, some accessor functions were added to allow Spidermonkey
to access the internals of the bitmap.

cranelift/codegen/src/binemit/stackmap.rs

@@ -6,10 +6,19 @@ use alloc::vec::Vec;
 type Num = u32;
 const NUM_BITS: usize = core::mem::size_of::<Num>() * 8;
 
-/// Wrapper class for longer bit vectors that cannot be represented by a single BitSet.
+/// A stack map is a bitmap with one bit per machine word on the stack. Stack
+/// maps are created at `safepoint` instructions and record all live reference
+/// values that are on the stack. All slot kinds, except `OutgoingArg` are
+/// captured in a stack map. The `OutgoingArg`'s will be captured in the callee
+/// function as `IncomingArg`'s.
+///
+/// The first value in the bitmap is of the lowest addressed slot on the stack.
+/// As all stacks in Isa's supported by Cranelift grow down, this means that
+/// first value is of the top of the stack and values proceed down the stack.
 #[derive(Clone, Debug)]
 pub struct Stackmap {
     bitmap: Vec<BitSet<Num>>,
+    mapped_words: u32,
 }
 
 impl Stackmap {
@@ -27,32 +36,37 @@ impl Stackmap {
         for val in args {
             if let Some(value_loc) = loc.get(*val) {
                 match *value_loc {
-                    ir::ValueLoc::Stack(stack_slot) => live_ref_in_stack_slot.insert(stack_slot),
-                    _ => false,
-                };
+                    ir::ValueLoc::Stack(stack_slot) => {
+                        live_ref_in_stack_slot.insert(stack_slot);
+                    }
+                    _ => {}
+                }
             }
         }
 
-        // SpiderMonkey stackmap structure:
-        // <trap reg dump> + <general spill> + <frame> + <inbound args>
-        // Bit vector goes from lower addresses to higher addresses.
-
-        // TODO: Get trap register layout from Spidermonkey and prepend to bitvector below.
         let stack = &func.stack_slots;
-        let frame_size = stack.frame_size.unwrap();
-        let word_size = ir::stackslot::StackSize::from(isa.pointer_bytes());
-        let num_words = (frame_size / word_size) as usize;
-        let mut vec = alloc::vec::Vec::with_capacity(num_words);
+        let info = func.stack_slots.layout_info.unwrap();
 
+        // Refer to the doc comment for `Stackmap` above to understand the
+        // bitmap representation used here.
+        let map_size = (info.frame_size + info.inbound_args_size) as usize;
+        let word_size = isa.pointer_bytes() as usize;
+        let num_words = map_size / word_size;
+
+        let mut vec = alloc::vec::Vec::with_capacity(num_words);
         vec.resize(num_words, false);
 
-        // Frame (includes spills and inbound args).
         for (ss, ssd) in stack.iter() {
-            if live_ref_in_stack_slot.contains(&ss) {
-                // Assumption: greater magnitude of offset imply higher address.
-                let index = (((ssd.offset.unwrap().abs() as u32) - ssd.size) / word_size) as usize;
-                vec[index] = true;
+            if !live_ref_in_stack_slot.contains(&ss)
+                || ssd.kind == ir::stackslot::StackSlotKind::OutgoingArg
+            {
+                continue;
             }
+
+            debug_assert!(ssd.size as usize == word_size);
+            let bytes_from_bottom = info.frame_size as i32 + ssd.offset.unwrap();
+            let words_from_bottom = (bytes_from_bottom as usize) / word_size;
+            vec[words_from_bottom] = true;
         }
 
         Self::from_slice(&vec)
@@ -73,7 +87,10 @@ impl Stackmap {
             }
             bitmap.push(BitSet(curr_word));
         }
-        Self { bitmap }
+        Self {
+            mapped_words: len as u32,
+            bitmap,
+        }
     }
 
     /// Returns a specified bit.
@@ -83,6 +100,16 @@ impl Stackmap {
         let word_offset = (bit_index % NUM_BITS) as u8;
         self.bitmap[word_index].contains(word_offset)
     }
+
+    /// Returns the raw bitmap that represents this stack map.
+    pub fn as_slice(&self) -> &[BitSet<u32>] {
+        &self.bitmap
+    }
+
+    /// Returns the number of words represented by this stack map.
+    pub fn mapped_words(&self) -> u32 {
+        self.mapped_words
+    }
 }
 
 #[cfg(test)]

cranelift/codegen/src/ir/mod.rs

@@ -53,7 +53,7 @@ pub use crate::ir::libcall::{get_libcall_funcref, get_probestack_funcref, LibCal
 pub use crate::ir::memflags::MemFlags;
 pub use crate::ir::progpoint::{ExpandedProgramPoint, ProgramOrder, ProgramPoint};
 pub use crate::ir::sourceloc::SourceLoc;
-pub use crate::ir::stackslot::{StackSlotData, StackSlotKind, StackSlots};
+pub use crate::ir::stackslot::{StackLayoutInfo, StackSlotData, StackSlotKind, StackSlots};
 pub use crate::ir::table::TableData;
 pub use crate::ir::trapcode::TrapCode;
 pub use crate::ir::types::Type;

cranelift/codegen/src/ir/stackslot.rs

@@ -162,6 +162,23 @@ impl fmt::Display for StackSlotData {
     }
 }
 
+/// Stack frame layout information.
+///
+/// This is computed by the `layout_stack()` method.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
+pub struct StackLayoutInfo {
+    /// The total size of the stack frame.
+    ///
+    /// This is the distance from the stack pointer in the current function to the stack pointer in
+    /// the calling function, so it includes a pushed return address as well as space for outgoing
+    /// call arguments.
+    pub frame_size: StackSize,
+
+    /// The total size of the stack frame for inbound arguments pushed by the caller.
+    pub inbound_args_size: StackSize,
+}
+
 /// Stack frame manager.
 ///
 /// Keep track of all the stack slots used by a function.
@@ -177,14 +194,8 @@ pub struct StackSlots {
     /// All the emergency slots.
     emergency: Vec<StackSlot>,
 
-    /// The total size of the stack frame.
-    ///
-    /// This is the distance from the stack pointer in the current function to the stack pointer in
-    /// the calling function, so it includes a pushed return address as well as space for outgoing
-    /// call arguments.
-    ///
-    /// This is computed by the `layout()` method.
-    pub frame_size: Option<StackSize>,
+    /// Layout information computed from `layout_stack`.
+    pub layout_info: Option<StackLayoutInfo>,
 }
 
 /// Stack slot manager functions that behave mostly like an entity map.
@@ -195,7 +206,7 @@ impl StackSlots {
             slots: PrimaryMap::new(),
             outgoing: Vec::new(),
             emergency: Vec::new(),
-            frame_size: None,
+            layout_info: None,
         }
     }
 
@@ -204,7 +215,7 @@ impl StackSlots {
         self.slots.clear();
         self.outgoing.clear();
         self.emergency.clear();
-        self.frame_size = None;
+        self.layout_info = None;
     }
 
     /// Allocate a new stack slot.

cranelift/codegen/src/isa/stack.rs

@@ -36,8 +36,9 @@ impl StackRef {
     /// Get a reference to `ss` using the stack pointer as a base.
     pub fn sp(ss: StackSlot, frame: &StackSlots) -> Self {
         let size = frame
-            .frame_size
-            .expect("Stack layout must be computed before referencing stack slots");
+            .layout_info
+            .expect("Stack layout must be computed before referencing stack slots")
+            .frame_size;
         let slot = &frame[ss];
         let offset = if slot.kind == StackSlotKind::OutgoingArg {
             // Outgoing argument slots have offsets relative to our stack pointer.

cranelift/codegen/src/stack_layout.rs

@@ -1,7 +1,7 @@
 //! Computing stack layout.
 
 use crate::ir::stackslot::{StackOffset, StackSize, StackSlotKind};
-use crate::ir::StackSlots;
+use crate::ir::{StackLayoutInfo, StackSlots};
 use crate::result::{CodegenError, CodegenResult};
 use core::cmp::{max, min};
 
@@ -44,6 +44,7 @@ pub fn layout_stack(
     // require the stack to be aligned.
 
     let mut incoming_min = 0;
+    let mut incoming_max = 0;
     let mut outgoing_max = 0;
     let mut min_align = alignment;
     let mut must_align = is_leaf;
@@ -56,6 +57,7 @@ pub fn layout_stack(
         match slot.kind {
             StackSlotKind::IncomingArg => {
                 incoming_min = min(incoming_min, slot.offset.unwrap());
+                incoming_max = max(incoming_max, slot.offset.unwrap() + slot.size as i32);
             }
             StackSlotKind::OutgoingArg => {
                 let offset = slot
@@ -119,8 +121,14 @@ pub fn layout_stack(
         offset &= -(alignment as StackOffset);
     }
 
+    // Set the computed layout information for the frame
     let frame_size = (offset as StackSize).wrapping_neg();
-    frame.frame_size = Some(frame_size);
+    let inbound_args_size = incoming_max as u32;
+    frame.layout_info = Some(StackLayoutInfo {
+        frame_size,
+        inbound_args_size,
+    });
+
     Ok(frame_size)
 }
 

cranelift/filetests/src/test_binemit.rs

@@ -142,7 +142,10 @@ impl SubTest for TestBinEmit {
             .values()
             .map(|slot| slot.offset.unwrap())
             .min();
-        func.stack_slots.frame_size = min_offset.map(|off| (-off) as u32);
+        func.stack_slots.layout_info = min_offset.map(|off| ir::StackLayoutInfo {
+            frame_size: (-off) as u32,
+            inbound_args_size: 0,
+        });
 
         let opt_level = isa.flags().opt_level();
 

cranelift/wasm/src/code_translator.rs

@@ -1151,7 +1151,6 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             segment,
             table: table_index,
         } => {
-            // The WebAssembly MVP only supports one table and we assume it here.
             let table = state.get_table(builder.func, *table_index, environ)?;
             let len = state.pop1();
             let src = state.pop1();