Stack Limit as an Argument Purpose (#372)

* Initial approach.

* Move stack_limit check before opening the frame

* Account for GPRs and frame pointer in stack check

* Check stack_limit example.

* Remove stack_limit attribute code.

Amends #359

* fmt

lib/codegen/src/context.rs

@@ -297,10 +297,6 @@ impl Context {
 
     /// Insert prologue and epilogues after computing the stack frame layout.
     pub fn prologue_epilogue(&mut self, isa: &TargetIsa) -> CodegenResult<()> {
-        assert!(
-            self.func.stack_limit.is_none(),
-            "stack_limit isn't implemented yet"
-        );
         isa.prologue_epilogue(&mut self.func)?;
         self.verify_if(isa)?;
         self.verify_locations_if(isa)?;

lib/codegen/src/ir/extfunc.rs

@@ -293,10 +293,25 @@ pub enum ArgumentPurpose {
     /// This is a special-purpose argument used to identify the calling convention expected by the
     /// caller in an indirect call. The callee can verify that the expected signature ID matches.
     SignatureId,
+
+    /// A stack limit pointer.
+    ///
+    /// This is a pointer to a stack limit. It is used to check the current stack pointer
+    /// against. Can only appear once in a signature.
+    StackLimit,
 }
 
 /// Text format names of the `ArgumentPurpose` variants.
-static PURPOSE_NAMES: [&str; 7] = ["normal", "sret", "link", "fp", "csr", "vmctx", "sigid"];
+static PURPOSE_NAMES: [&str; 8] = [
+    "normal",
+    "sret",
+    "link",
+    "fp",
+    "csr",
+    "vmctx",
+    "sigid",
+    "stack_limit",
+];
 
 impl fmt::Display for ArgumentPurpose {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -315,6 +330,7 @@ impl FromStr for ArgumentPurpose {
             "csr" => Ok(ArgumentPurpose::CalleeSaved),
             "vmctx" => Ok(ArgumentPurpose::VMContext),
             "sigid" => Ok(ArgumentPurpose::SignatureId),
+            "stack_limit" => Ok(ArgumentPurpose::StackLimit),
             _ => Err(()),
         }
     }
@@ -370,6 +386,8 @@ mod tests {
             ArgumentPurpose::FramePointer,
             ArgumentPurpose::CalleeSaved,
             ArgumentPurpose::VMContext,
+            ArgumentPurpose::SignatureId,
+            ArgumentPurpose::StackLimit,
         ];
         for (&e, &n) in all_purpose.iter().zip(PURPOSE_NAMES.iter()) {
             assert_eq!(e.to_string(), n);

lib/codegen/src/ir/function.rs

@@ -32,10 +32,6 @@ pub struct Function {
     /// Stack slots allocated in this function.
     pub stack_slots: StackSlots,
 
-    /// If not `None`, represents the address that the stack pointer should
-    /// be checked against.
-    pub stack_limit: Option<ir::GlobalValue>,
-
     /// Global values referenced.
     pub global_values: PrimaryMap<ir::GlobalValue, ir::GlobalValueData>,
 
@@ -82,7 +78,6 @@ impl Function {
             name,
             signature: sig,
             stack_slots: StackSlots::new(),
-            stack_limit: None,
             global_values: PrimaryMap::new(),
             heaps: PrimaryMap::new(),
             tables: PrimaryMap::new(),
@@ -133,15 +128,6 @@ impl Function {
         self.stack_slots.push(data)
     }
 
-    /// Sets the stack limit for the function.
-    ///
-    /// Returns previous one if any.
-    pub fn set_stack_limit(&mut self, stack_limit: Option<GlobalValue>) -> Option<GlobalValue> {
-        let prev = self.stack_limit.take();
-        self.stack_limit = stack_limit;
-        prev
-    }
-
     /// Adds a signature which can later be used to declare an external function import.
     pub fn import_signature(&mut self, signature: Signature) -> SigRef {
         self.dfg.signatures.push(signature)

lib/codegen/src/isa/x86/abi.rs

@@ -427,6 +427,21 @@ fn insert_common_prologue(
     csrs: &RegisterSet,
     isa: &TargetIsa,
 ) {
+    if stack_size > 0 {
+        // Check if there is a special stack limit parameter. If so insert stack check.
+        if let Some(stack_limit_arg) = pos.func.special_param(ArgumentPurpose::StackLimit) {
+            // Total stack size is the size of all stack area used by the function, including
+            // pushed CSRs, frame pointer.
+            // Also, the size of a return address, implicitly pushed by a x86 `call` instruction, also
+            // should be accounted for.
+            // TODO: Check if the function body actually contains a `call` instruction.
+            let word_size = isa.pointer_bytes();
+            let total_stack_size = (csrs.iter(GPR).len() + 1 + 1) as i64 * word_size as i64;
+
+            insert_stack_check(pos, total_stack_size, stack_limit_arg);
+        }
+    }
+
     // Append param to entry EBB
     let ebb = pos.current_ebb().expect("missing ebb under cursor");
     let fp = pos.func.dfg.append_ebb_param(ebb, reg_type);
@@ -493,6 +508,29 @@ fn insert_common_prologue(
     }
 }
 
+/// Insert a check that generates a trap if the stack pointer goes
+/// below a value in `stack_limit_arg`.
+fn insert_stack_check(pos: &mut EncCursor, stack_size: i64, stack_limit_arg: ir::Value) {
+    use ir::condcodes::IntCC;
+
+    // Copy `stack_limit_arg` into a %rax and use it for calculating
+    // a SP threshold.
+    let stack_limit_copy = pos.ins().copy(stack_limit_arg);
+    pos.func.locations[stack_limit_copy] = ir::ValueLoc::Reg(RU::rax as RegUnit);
+    let sp_threshold = pos.ins().iadd_imm(stack_limit_copy, stack_size);
+    pos.func.locations[sp_threshold] = ir::ValueLoc::Reg(RU::rax as RegUnit);
+
+    // If the stack pointer currently reaches the SP threshold or below it then after opening
+    // the current stack frame, the current stack pointer will reach the limit.
+    let cflags = pos.ins().ifcmp_sp(sp_threshold);
+    pos.func.locations[cflags] = ir::ValueLoc::Reg(RU::rflags as RegUnit);
+    pos.ins().trapif(
+        IntCC::UnsignedGreaterThanOrEqual,
+        cflags,
+        ir::TrapCode::StackOverflow,
+    );
+}
+
 /// Find all `return` instructions and insert epilogues before them.
 fn insert_common_epilogues(
     pos: &mut EncCursor,

lib/codegen/src/legalizer/boundary.rs

@@ -73,6 +73,7 @@ fn legalize_entry_params(func: &mut Function, entry: Ebb) {
     let mut has_link = false;
     let mut has_vmctx = false;
     let mut has_sigid = false;
+    let mut has_stack_limit = false;
 
     // Insert position for argument conversion code.
     // We want to insert instructions before the first instruction in the entry block.
@@ -111,6 +112,10 @@ fn legalize_entry_params(func: &mut Function, entry: Ebb) {
                     debug_assert!(!has_sigid, "Multiple sigid arguments found");
                     has_sigid = true;
                 }
+                ArgumentPurpose::StackLimit => {
+                    debug_assert!(!has_stack_limit, "Multiple stack_limit arguments found");
+                    has_stack_limit = true;
+                }
                 _ => panic!("Unexpected special-purpose arg {}", abi_type),
             }
             abi_arg += 1;
@@ -167,6 +172,10 @@ fn legalize_entry_params(func: &mut Function, entry: Ebb) {
                 debug_assert!(!has_sigid, "Multiple sigid parameters found");
                 has_sigid = true;
             }
+            ArgumentPurpose::StackLimit => {
+                debug_assert!(!has_stack_limit, "Multiple stack_limit parameters found");
+                has_stack_limit = true;
+            }
         }
 
         // Just create entry block values to match here. We will use them in `handle_return_abi()`

lib/codegen/src/write.rs

@@ -136,11 +136,6 @@ fn write_preamble(
         writeln!(w, "    {} = {}", jt, jt_data)?;
     }
 
-    if let Some(stack_limit) = func.stack_limit {
-        any = true;
-        writeln!(w, "    stack_limit = {}", stack_limit)?;
-    }
-
     Ok(any)
 }