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

cranelift/docs/ir.rst

@@ -546,17 +546,6 @@ instructions before instruction selection::
 When Cranelift code is running in a sandbox, it can also be necessary to include
 stack overflow checks in the prologue.
 
-.. inst:: stack_limit = GV
-
-    Set the stack limit of the current function.
-
-    If set, in the preamble read the stack limit from ``GV`` and compare it to the stack pointer. If
-    the stack pointer has reached or exceeded the limit, generate a trap with a
-    ``stk_ovf`` code.
-
-    Setting `stack_limit` is an alternative way to detect stack overflow, when using
-    a calling convention that doesn't perform stack probes.
-
 Global values
 -------------
 

cranelift/filetests/isa/x86/prologue-epilogue.clif

@@ -228,3 +228,28 @@ ebb4:
 ; check: function %divert
 ; check: regmove v5, %rcx -> %rbx
 ; check: [Op1popq#58,%rbx]                   v15 = x86_pop.i64
+
+; Stack limit checking
+
+function %stack_limit(i64 stack_limit) {
+    ss0 = explicit_slot 168
+ebb0(v0: i64):
+    return
+}
+
+; check: function %stack_limit(i64 stack_limit [%rdi], i64 fp [%rbp]) -> i64 fp [%rbp] fast {
+; nextln:     ss0 = explicit_slot 168, offset -184
+; nextln:     ss1 = incoming_arg 16, offset -16
+; nextln: 
+; nextln: ebb0(v0: i64 [%rdi], v4: i64 [%rbp]):
+; nextln:     v1 = copy v0
+; nextln:     v2 = iadd_imm v1, 16
+; nextln:     v3 = ifcmp_sp v2
+; nextln:     trapif uge v3, stk_ovf
+; nextln:     x86_push v4
+; nextln:     copy_special %rsp -> %rbp
+; nextln:     adjust_sp_down_imm 176
+; nextln:     adjust_sp_up_imm 176
+; nextln:     v5 = x86_pop.i64
+; nextln:     return v5
+; nextln: }

cranelift/filetests/parser/preamble.clif

@@ -1,18 +0,0 @@
-test cat
-
-; Verify parsing of stack_limit.
-function %minimal(i64 vmctx) {
-gv0 = vmctx
-; Stack limit
-stack_limit = gv0
-
-ebb0:
-    trap user0
-}
-; sameln: function %minimal(i64 vmctx) fast {
-; nextln: gv0 = vmctx
-; nextln: stack_limit = gv0
-; nextln: 
-; nextln: ebb0:
-; nextln:     trap user0
-; nextln: }

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)
 }
 

lib/reader/src/parser.rs

@@ -274,15 +274,6 @@ impl<'a> Context<'a> {
         }
     }
 
-    // Assign the global for the stack limit.
-    fn set_stack_limit(&mut self, gv: GlobalValue, loc: Location) -> ParseResult<()> {
-        if self.function.set_stack_limit(Some(gv)).is_some() {
-            err!(loc, "multiple stack_limit declarations")
-        } else {
-            Ok(())
-        }
-    }
-
     // Allocate a new EBB.
     fn add_ebb(&mut self, ebb: Ebb, loc: Location) -> ParseResult<Ebb> {
         self.map.def_ebb(ebb, loc)?;
@@ -1032,7 +1023,6 @@ impl<'a> Parser<'a> {
     //                   * function-decl
     //                   * signature-decl
     //                   * jump-table-decl
-    //                   * stack-limit-decl
     //
     // The parsed decls are added to `ctx` rather than returned.
     fn parse_preamble(&mut self, ctx: &mut Context) -> ParseResult<()> {
@@ -1074,9 +1064,6 @@ impl<'a> Parser<'a> {
                     self.parse_jump_table_decl()
                         .and_then(|(jt, dat)| ctx.add_jt(jt, dat, self.loc))
                 }
-                Some(Token::Identifier("stack_limit")) => self
-                    .parse_stack_limit_decl()
-                    .and_then(|gv| ctx.set_stack_limit(gv, self.loc)),
                 // More to come..
                 _ => return Ok(()),
             }?;
@@ -1414,15 +1401,6 @@ impl<'a> Parser<'a> {
         }
     }
 
-    /// stack-limit-decl ::= "stack_limit" "=" GlobalValue(gv)
-    fn parse_stack_limit_decl(&mut self) -> ParseResult<GlobalValue> {
-        self.consume();
-        self.match_token(Token::Equal, "expected '=' in stack limit declaration")?;
-        let gv = self.match_gv("expected global value")?;
-
-        Ok(gv)
-    }
-
     // Parse a function body, add contents to `ctx`.
     //
     // function-body ::= * { extended-basic-block }