Implement interrupting wasm code, reimplement stack overflow

Cargo.toml

@@ -39,6 +39,7 @@ file-per-thread-logger = "0.1.1"
 wat = "1.0.10"
 libc = "0.2.60"
 rayon = "1.2.1"
+humantime = "1.3.0"
 
 [dev-dependencies]
 filecheck = "0.5.0"

cranelift/codegen/src/ir/entities.rs

@@ -392,6 +392,8 @@ pub enum AnyEntity {
     Heap(Heap),
     /// A table.
     Table(Table),
+    /// A function's stack limit
+    StackLimit,
 }
 
 impl fmt::Display for AnyEntity {
@@ -409,6 +411,7 @@ impl fmt::Display for AnyEntity {
             Self::SigRef(r) => r.fmt(f),
             Self::Heap(r) => r.fmt(f),
             Self::Table(r) => r.fmt(f),
+            Self::StackLimit => write!(f, "stack_limit"),
         }
     }
 }

cranelift/codegen/src/ir/function.rs

@@ -95,6 +95,13 @@ pub struct Function {
     ///
     /// This is used for some ABIs to generate unwind information.
     pub epilogues_start: Vec<Inst>,
+
+    /// 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
+    /// ensure that a trap happens if the stack pointer goes below the
+    /// threshold specified here.
+    pub stack_limit: Option<ir::GlobalValue>,
 }
 
 impl Function {
@@ -119,6 +126,7 @@ impl Function {
             srclocs: SecondaryMap::new(),
             prologue_end: None,
             epilogues_start: Vec::new(),
+            stack_limit: None,
         }
     }
 
@@ -140,6 +148,7 @@ impl Function {
         self.srclocs.clear();
         self.prologue_end = None;
         self.epilogues_start.clear();
+        self.stack_limit = None;
     }
 
     /// Create a new empty, anonymous function with a Fast calling convention.

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

@@ -685,21 +685,32 @@ fn insert_common_prologue(
     fpr_slot: Option<&StackSlot>,
     isa: &dyn 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.
-            // If any FPR are present, count them as well as necessary alignment space.
-            // TODO: Check if the function body actually contains a `call` instruction.
-            let mut total_stack_size =
-                (csrs.iter(GPR).len() + 1 + 1) as i64 * (isa.pointer_bytes() as isize) as i64;
-
-            total_stack_size += csrs.iter(FPR).len() as i64 * types::F64X2.bytes() as i64;
-
-            insert_stack_check(pos, total_stack_size, stack_limit_arg);
+    // If this is a leaf function with zero stack, then there's no need to
+    // insert a stack check since it can't overflow anything and
+    // forward-progress is guarantee so long as loop are handled anyway.
+    //
+    // If this has a stack size it could stack overflow, or if it isn't a leaf
+    // it could be part of a long call chain which we need to check anyway.
+    //
+    // First we look for the stack limit as a special argument to the function,
+    // and failing that we see if a custom stack limit factory has been provided
+    // which will be used to likely calculate the stack limit from the arguments
+    // or perhaps constants.
+    if stack_size > 0 || !pos.func.is_leaf() {
+        let scratch = ir::ValueLoc::Reg(RU::rax as RegUnit);
+        let stack_limit_arg = match pos.func.special_param(ArgumentPurpose::StackLimit) {
+            Some(arg) => {
+                let copy = pos.ins().copy(arg);
+                pos.func.locations[copy] = scratch;
+                Some(copy)
+            }
+            None => pos
+                .func
+                .stack_limit
+                .map(|gv| interpret_gv(pos, gv, scratch)),
+        };
+        if let Some(stack_limit_arg) = stack_limit_arg {
+            insert_stack_check(pos, stack_size, stack_limit_arg);
         }
     }
 
@@ -811,16 +822,76 @@ fn insert_common_prologue(
     );
 }
 
+/// Inserts code necessary to calculate `gv`.
+///
+/// Note that this is typically done with `ins().global_value(...)` but that
+/// requires legalization to run to encode it, and we're running super late
+/// here in the backend where legalization isn't possible. To get around this
+/// we manually interpret the `gv` specified and do register allocation for
+/// intermediate values.
+///
+/// This is an incomplete implementation of loading `GlobalValue` values to get
+/// compared to the stack pointer, but currently it serves enough functionality
+/// to get this implemented in `wasmtime` itself. This'll likely get expanded a
+/// bit over time!
+fn interpret_gv(pos: &mut EncCursor, gv: ir::GlobalValue, scratch: ir::ValueLoc) -> ir::Value {
+    match pos.func.global_values[gv] {
+        ir::GlobalValueData::VMContext => pos
+            .func
+            .special_param(ir::ArgumentPurpose::VMContext)
+            .expect("no vmcontext parameter found"),
+        ir::GlobalValueData::Load {
+            base,
+            offset,
+            global_type,
+            readonly: _,
+        } => {
+            let base = interpret_gv(pos, base, scratch);
+            let ret = pos
+                .ins()
+                .load(global_type, ir::MemFlags::trusted(), base, offset);
+            pos.func.locations[ret] = scratch;
+            return ret;
+        }
+        ref other => panic!("global value for stack limit not supported: {}", other),
+    }
+}
+
 /// 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 crate::ir::condcodes::IntCC;
 
+    // Our stack pointer, after subtracting `stack_size`, must not be below
+    // `stack_limit_arg`. To do this we're going to add `stack_size` to
+    // `stack_limit_arg` and see if the stack pointer is below that. The
+    // `stack_size + stack_limit_arg` computation might overflow, however, due
+    // to how stack limits may be loaded and set externally to trigger a trap.
+    //
+    // To handle this we'll need an extra comparison to see if the stack
+    // pointer is already below `stack_limit_arg`. Most of the time this
+    // isn't necessary though since the stack limit which triggers a trap is
+    // likely a sentinel somewhere around `usize::max_value()`. In that case
+    // only conditionally emit this pre-flight check. That way most functions
+    // only have the one comparison, but are also guaranteed that if we add
+    // `stack_size` to `stack_limit_arg` is won't overflow.
+    //
+    // This does mean that code generators which use this stack check
+    // functionality need to ensure that values stored into the stack limit
+    // will never overflow if this threshold is added.
+    if stack_size >= 32 * 1024 {
+        let cflags = pos.ins().ifcmp_sp(stack_limit_arg);
+        pos.func.locations[cflags] = ir::ValueLoc::Reg(RU::rflags as RegUnit);
+        pos.ins().trapif(
+            IntCC::UnsignedGreaterThanOrEqual,
+            cflags,
+            ir::TrapCode::StackOverflow,
+        );
+    }
+
     // 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);
+    let sp_threshold = pos.ins().iadd_imm(stack_limit_arg, 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

cranelift/codegen/src/write.rs

@@ -107,6 +107,11 @@ pub trait FuncWriter {
             self.write_entity_definition(w, func, cref.into(), cval)?;
         }
 
+        if let Some(limit) = func.stack_limit {
+            any = true;
+            self.write_entity_definition(w, func, AnyEntity::StackLimit, &limit)?;
+        }
+
         Ok(any)
     }
 

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

@@ -1,6 +1,7 @@
 test compile
 set opt_level=speed_and_size
 set is_pic
+set enable_probestack=false
 target x86_64 haswell
 
 ; An empty function.
@@ -244,7 +245,7 @@ block0(v0: i64):
 ; nextln: 
 ; nextln: block0(v0: i64 [%rdi], v4: i64 [%rbp]):
 ; nextln:     v1 = copy v0
-; nextln:     v2 = iadd_imm v1, 16
+; nextln:     v2 = iadd_imm v1, 176
 ; nextln:     v3 = ifcmp_sp v2
 ; nextln:     trapif uge v3, stk_ovf
 ; nextln:     x86_push v4
@@ -254,3 +255,60 @@ block0(v0: i64):
 ; nextln:     v5 = x86_pop.i64
 ; nextln:     return v5
 ; nextln: }
+
+function %big_stack_limit(i64 stack_limit) {
+    ss0 = explicit_slot 40000
+block0(v0: i64):
+    return
+}
+
+; check: function %big_stack_limit(i64 stack_limit [%rdi], i64 fp [%rbp]) -> i64 fp [%rbp] fast {
+; nextln:     ss0 = explicit_slot 40000, offset -40016
+; nextln:     ss1 = incoming_arg 16, offset -16
+; nextln: 
+; nextln: block0(v0: i64 [%rdi], v5: i64 [%rbp]):
+; nextln:     v1 = copy v0
+; nextln:     v2 = ifcmp_sp v1
+; nextln:     trapif uge v2, stk_ovf
+; nextln:     v3 = iadd_imm v1, 0x9c40
+; nextln:     v4 = ifcmp_sp v3
+; nextln:     trapif uge v4, stk_ovf
+; nextln:     x86_push v5
+; nextln:     copy_special %rsp -> %rbp
+; nextln:     adjust_sp_down_imm 0x9c40
+; nextln:     adjust_sp_up_imm 0x9c40
+; nextln:     v6 = x86_pop.i64
+; nextln:     return v6
+; nextln: }
+
+function %limit_preamble(i64 vmctx) {
+    gv0 = vmctx
+    gv1 = load.i64 notrap aligned gv0
+    gv2 = load.i64 notrap aligned gv1+4
+    stack_limit = gv2
+    ss0 = explicit_slot 20
+block0(v0: i64):
+    return
+}
+
+; check: function %limit_preamble(i64 vmctx [%rdi], i64 fp [%rbp]) -> i64 fp [%rbp] fast {
+; nextln:     ss0 = explicit_slot 20, offset -36
+; nextln:     ss1 = incoming_arg 16, offset -16
+; nextln:     gv0 = vmctx
+; nextln:     gv1 = load.i64 notrap aligned gv0
+; nextln:     gv2 = load.i64 notrap aligned gv1+4
+; nextln:     stack_limit = gv2
+; nextln: 
+; nextln: block0(v0: i64 [%rdi], v5: i64 [%rbp]):
+; nextln:     v1 = load.i64 notrap aligned v0
+; nextln:     v2 = load.i64 notrap aligned v1+4
+; nextln:     v3 = iadd_imm v2, 32
+; nextln:     v4 = ifcmp_sp v3
+; nextln:     trapif uge v4, stk_ovf
+; nextln:     x86_push v5
+; nextln:     copy_special %rsp -> %rbp
+; nextln:     adjust_sp_down_imm 32
+; nextln:     adjust_sp_up_imm 32
+; nextln:     v6 = x86_pop.i64
+; nextln:     return v6
+; nextln: }

cranelift/reader/src/parser.rs

@@ -358,6 +358,15 @@ impl<'a> Context<'a> {
         Ok(())
     }
 
+    // Configure the stack limit of the current function.
+    fn add_stack_limit(&mut self, limit: GlobalValue, loc: Location) -> ParseResult<()> {
+        if self.function.stack_limit.is_some() {
+            return err!(loc, "stack limit defined twice");
+        }
+        self.function.stack_limit = Some(limit);
+        Ok(())
+    }
+
     // Resolve a reference to a constant.
     fn check_constant(&self, c: Constant, loc: Location) -> ParseResult<()> {
         if !self.map.contains_constant(c) {
@@ -598,6 +607,15 @@ impl<'a> Parser<'a> {
         err!(self.loc, "expected constant number: const«n»")
     }
 
+    // Match and consume a stack limit token
+    fn match_stack_limit(&mut self) -> ParseResult<()> {
+        if let Some(Token::Identifier("stack_limit")) = self.token() {
+            self.consume();
+            return Ok(());
+        }
+        err!(self.loc, "expected identifier: stack_limit")
+    }
+
     // Match and consume a block reference.
     fn match_block(&mut self, err_msg: &str) -> ParseResult<Block> {
         if let Some(Token::Block(block)) = self.token() {
@@ -1455,6 +1473,7 @@ 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<()> {
@@ -1503,6 +1522,11 @@ impl<'a> Parser<'a> {
                     self.parse_constant_decl()
                         .and_then(|(c, v)| ctx.add_constant(c, v, self.loc))
                 }
+                Some(Token::Identifier("stack_limit")) => {
+                    self.start_gathering_comments();
+                    self.parse_stack_limit_decl()
+                        .and_then(|gv| ctx.add_stack_limit(gv, self.loc))
+                }
                 // More to come..
                 _ => return Ok(()),
             }?;
@@ -1907,6 +1931,28 @@ impl<'a> Parser<'a> {
         Ok((name, data))
     }
 
+    // Parse a stack limit decl
+    //
+    // stack-limit-decl ::= * StackLimit "=" GlobalValue(gv)
+    fn parse_stack_limit_decl(&mut self) -> ParseResult<GlobalValue> {
+        self.match_stack_limit()?;
+        self.match_token(Token::Equal, "expected '=' in stack limit decl")?;
+        let limit = match self.token() {
+            Some(Token::GlobalValue(base_num)) => match GlobalValue::with_number(base_num) {
+                Some(gv) => gv,
+                None => return err!(self.loc, "invalid global value number for stack limit"),
+            },
+            _ => return err!(self.loc, "expected global value"),
+        };
+        self.consume();
+
+        // Collect any trailing comments.
+        self.token();
+        self.claim_gathered_comments(AnyEntity::StackLimit);
+
+        Ok(limit)
+    }
+
     // Parse a function body, add contents to `ctx`.
     //
     // function-body ::= * { extended-basic-block }

cranelift/reader/src/run_command.rs

@@ -12,7 +12,8 @@ use std::fmt::{Display, Formatter, Result};
 
 /// A run command appearing in a test file.
 ///
-/// For parsing, see [Parser::parse_run_command].
+/// For parsing, see
+/// [Parser::parse_run_command](crate::parser::Parser::parse_run_command).
 #[derive(PartialEq, Debug)]
 pub enum RunCommand {
     /// Invoke a function and print its result.
@@ -66,6 +67,8 @@ impl Display for Invocation {
 
 /// Represent a data value. Where [Value] is an SSA reference, [DataValue] is the type + value
 /// that would be referred to by a [Value].
+///
+/// [Value]: cranelift_codegen::ir::Value
 #[allow(missing_docs)]
 #[derive(Clone, Debug, PartialEq)]
 pub enum DataValue {