1 files changed, 70 insertions, 0 deletions
diff --git a/third_party/rust/wasm-smith/src/core/terminate.rs b/third_party/rust/wasm-smith/src/core/terminate.rs
new file mode 100644
index 0000000000..adcfeed54f
--- /dev/null
+++ b/third_party/rust/wasm-smith/src/core/terminate.rs
@@ -0,0 +1,70 @@
+use super::*;
+use std::mem;
+use wasm_encoder::BlockType;
+
+impl Module {
+    /// Ensure that all of this Wasm module's functions will terminate when
+    /// executed.
+    ///
+    /// This adds a new mutable, exported global to the module to keep track of
+    /// how much "fuel" is left. Fuel is decremented at the head of each loop
+    /// and function. When fuel reaches zero, a trap is raised.
+    ///
+    /// The index of the fuel global is returned, so that you may control how
+    /// much fuel the module is given.
+    pub fn ensure_termination(&mut self, default_fuel: u32) -> u32 {
+        let fuel_global = self.globals.len() as u32;
+        self.globals.push(GlobalType {
+            val_type: ValType::I32,
+            mutable: true,
+        });
+        self.defined_globals.push((
+            fuel_global,
+            GlobalInitExpr::ConstExpr(ConstExpr::i32_const(default_fuel as i32)),
+        ));
+
+        for code in &mut self.code {
+            let check_fuel = |insts: &mut Vec<Instruction>| {
+                // if fuel == 0 { trap }
+                insts.push(Instruction::GlobalGet(fuel_global));
+                insts.push(Instruction::I32Eqz);
+                insts.push(Instruction::If(BlockType::Empty));
+                insts.push(Instruction::Unreachable);
+                insts.push(Instruction::End);
+
+                // fuel -= 1
+                insts.push(Instruction::GlobalGet(fuel_global));
+                insts.push(Instruction::I32Const(1));
+                insts.push(Instruction::I32Sub);
+                insts.push(Instruction::GlobalSet(fuel_global));
+            };
+
+            let instrs = match &mut code.instructions {
+                Instructions::Generated(list) => list,
+                // only present on modules contained within
+                // `MaybeInvalidModule`, which doesn't expose its internal
+                // `Module`.
+                Instructions::Arbitrary(_) => unreachable!(),
+            };
+            let mut new_insts = Vec::with_capacity(instrs.len() * 2);
+
+            // Check fuel at the start of functions to deal with infinite
+            // recursion.
+            check_fuel(&mut new_insts);
+
+            for inst in mem::replace(instrs, vec![]) {
+                let is_loop = matches!(&inst, Instruction::Loop(_));
+                new_insts.push(inst);
+
+                // Check fuel at loop heads to deal with infinite loops.
+                if is_loop {
+                    check_fuel(&mut new_insts);
+                }
+            }
+
+            *instrs = new_insts;
+        }
+
+        fuel_global
+    }
+}