Add Config::fuel_cost to customize fuel cost func

crates/cranelift/src/compiler.rs

@@ -32,7 +32,7 @@ use std::mem;
 use std::sync::{Arc, Mutex};
 use wasmparser::{FuncValidatorAllocations, FunctionBody};
 use wasmtime_environ::{
-    AddressMapSection, CacheStore, CompileError, FilePos, FlagValue, FunctionBodyData,
+    AddressMapSection, CacheStore, CompileError, FilePos, FlagValue, FuelCost, FunctionBodyData,
     FunctionInfo, InstructionAddressMap, Module, ModuleTranslation, ModuleTypes, PtrSize,
     StackMapInformation, Trampoline, TrapCode, TrapEncodingBuilder, TrapInformation, Tunables,
     VMOffsets,
@@ -188,6 +188,7 @@ impl wasmtime_environ::Compiler for Compiler {
         func_index: DefinedFuncIndex,
         input: FunctionBodyData<'_>,
         tunables: &Tunables,
+        fuel_cost: &FuelCost,
         types: &ModuleTypes,
     ) -> Result<Box<dyn Any + Send>, CompileError> {
         let isa = &*self.isa;
@@ -211,7 +212,7 @@ impl wasmtime_environ::Compiler for Compiler {
             context.func.collect_debug_info();
         }
 
-        let mut func_env = FuncEnvironment::new(isa, translation, types, tunables);
+        let mut func_env = FuncEnvironment::new(isa, translation, types, tunables, fuel_cost);
 
         // The `stack_limit` global value below is the implementation of stack
         // overflow checks in Wasmtime.

crates/cranelift/src/func_environ.rs

@@ -16,8 +16,8 @@ use std::convert::TryFrom;
 use std::mem;
 use wasmparser::Operator;
 use wasmtime_environ::{
-    BuiltinFunctionIndex, MemoryPlan, MemoryStyle, Module, ModuleTranslation, ModuleTypes, PtrSize,
-    TableStyle, Tunables, VMOffsets, WASM_PAGE_SIZE,
+    BuiltinFunctionIndex, FuelCost, MemoryPlan, MemoryStyle, Module, ModuleTranslation,
+    ModuleTypes, PtrSize, TableStyle, Tunables, VMOffsets, WASM_PAGE_SIZE,
 };
 use wasmtime_environ::{FUNCREF_INIT_BIT, FUNCREF_MASK};
 
@@ -147,6 +147,8 @@ pub struct FuncEnvironment<'module_environment> {
     epoch_ptr_var: cranelift_frontend::Variable,
 
     fuel_consumed: i64,
+
+    fuel_cost: &'module_environment FuelCost,
 }
 
 impl<'module_environment> FuncEnvironment<'module_environment> {
@@ -155,6 +157,7 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         translation: &'module_environment ModuleTranslation<'module_environment>,
         types: &'module_environment ModuleTypes,
         tunables: &'module_environment Tunables,
+        fuel_cost: &'module_environment FuelCost,
     ) -> Self {
         let builtin_function_signatures = BuiltinFunctionSignatures::new(
             isa.pointer_type(),
@@ -182,6 +185,7 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
             // Start with at least one fuel being consumed because even empty
             // functions should consume at least some fuel.
             fuel_consumed: 1,
+            fuel_cost,
         }
     }
 
@@ -385,23 +389,7 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
             return;
         }
 
-        self.fuel_consumed += match op {
-            // Nop and drop generate no code, so don't consume fuel for them.
-            Operator::Nop | Operator::Drop => 0,
-
-            // Control flow may create branches, but is generally cheap and
-            // free, so don't consume fuel. Note the lack of `if` since some
-            // cost is incurred with the conditional check.
-            Operator::Block { .. }
-            | Operator::Loop { .. }
-            | Operator::Unreachable
-            | Operator::Return
-            | Operator::Else
-            | Operator::End => 0,
-
-            // everything else, just call it one operation.
-            _ => 1,
-        };
+        self.fuel_consumed = self.fuel_consumed.saturating_add((self.fuel_cost)(op));
 
         match op {
             // Exiting a function (via a return or unreachable) or otherwise

crates/environ/src/compilation.rs

@@ -91,6 +91,9 @@ pub trait CacheStore: Send + Sync + std::fmt::Debug {
     fn insert(&self, key: &[u8], value: Vec<u8>) -> bool;
 }
 
+///
+pub type FuelCost = dyn Fn(&wasmparser::Operator<'_>) -> i64 + Send + Sync;
+
 /// Abstract trait representing the ability to create a `Compiler` below.
 ///
 /// This is used in Wasmtime to separate compiler implementations, currently
@@ -169,6 +172,7 @@ pub trait Compiler: Send + Sync {
         index: DefinedFuncIndex,
         data: FunctionBodyData<'_>,
         tunables: &Tunables,
+        fuel_cost: &FuelCost,
         types: &ModuleTypes,
     ) -> Result<Box<dyn Any + Send>, CompileError>;
 

crates/wasmtime/src/config.rs

@@ -10,7 +10,7 @@ use std::sync::Arc;
 use wasmparser::WasmFeatures;
 #[cfg(feature = "cache")]
 use wasmtime_cache::CacheConfig;
-use wasmtime_environ::Tunables;
+use wasmtime_environ::{FuelCost, Tunables};
 use wasmtime_jit::{JitDumpAgent, NullProfilerAgent, ProfilingAgent, VTuneAgent};
 use wasmtime_runtime::{InstanceAllocator, OnDemandInstanceAllocator, RuntimeMemoryCreator};
 
@@ -92,6 +92,7 @@ pub struct Config {
     profiling_strategy: ProfilingStrategy,
 
     pub(crate) tunables: Tunables,
+    pub(crate) fuel_cost: Option<Arc<FuelCost>>,
     #[cfg(feature = "cache")]
     pub(crate) cache_config: CacheConfig,
     pub(crate) mem_creator: Option<Arc<dyn RuntimeMemoryCreator>>,
@@ -169,6 +170,7 @@ impl Config {
     pub fn new() -> Self {
         let mut ret = Self {
             tunables: Tunables::default(),
+            fuel_cost: None,
             #[cfg(compiler)]
             compiler_config: CompilerConfig::default(),
             #[cfg(feature = "cache")]
@@ -461,6 +463,25 @@ impl Config {
         self
     }
 
+    /// Customizes the cost function used to calculate fuel usage.
+    ///
+    /// The passed function takes a [`WasmOpcode`] and returns a value denoting
+    /// how much fuel is used when that opcode is present in a wasm module.
+    /// Note that there is no inherent meaning behind this integer; other than
+    /// triggering the out-of-fuel behavior when fuel reaches zero, what "one
+    /// fuel point" means is determined by the user.
+    pub fn fuel_cost(
+        &mut self,
+        f: impl Fn(WasmOpcode) -> u64 + Send + Sync + 'static,
+    ) -> &mut Self {
+        self.fuel_cost = Some(Arc::new(move |op| {
+            f(WasmOpcode::from_operator(op))
+                .try_into()
+                .unwrap_or(i64::MAX)
+        }));
+        self
+    }
+
     /// Enables epoch-based interruption.
     ///
     /// When executing code in async mode, we sometimes want to
@@ -1440,6 +1461,27 @@ impl Config {
         self
     }
 
+    pub(crate) fn get_fuel_cost(&self) -> &FuelCost {
+        use wasmparser::Operator;
+        self.fuel_cost.as_deref().unwrap_or(&|op| match op {
+            // Nop and drop generate no code, so don't consume fuel for them.
+            Operator::Nop | Operator::Drop => 0,
+
+            // Control flow may create branches, but is generally cheap and
+            // free, so don't consume fuel. Note the lack of `if` since some
+            // cost is incurred with the conditional check.
+            Operator::Block { .. }
+            | Operator::Loop { .. }
+            | Operator::Unreachable
+            | Operator::Return
+            | Operator::Else
+            | Operator::End => 0,
+
+            // everything else, just call it one operation.
+            _ => 1,
+        })
+    }
+
     pub(crate) fn validate(&self) -> Result<()> {
         if self.features.reference_types && !self.features.bulk_memory {
             bail!("feature 'reference_types' requires 'bulk_memory' to be enabled");
@@ -1689,3 +1731,26 @@ pub enum WasmBacktraceDetails {
     /// `WASMTIME_BACKTRACE_DETAILS` environment variable.
     Environment,
 }
+
+macro_rules! define_opcode_enum {
+    ($( @$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident)*) => {
+        /// The different operations the [WASM specification] defines. Note that this enum is not
+        /// stable and is subject to gain more variants as the WASM spec evolves.
+        ///
+        /// [WASM specification]: https://webassembly.github.io/spec/core/binary/instructions.html
+        #[derive(Copy, Clone)]
+        #[allow(missing_docs)]
+        pub enum WasmOpcode {
+            $( $op, )*
+        }
+        impl WasmOpcode {
+            fn from_operator(op: &wasmparser::Operator<'_>) -> Self {
+                use wasmparser::Operator::*;
+                match op {
+                    $( $op { .. } => Self::$op, )*
+                }
+            }
+        }
+    };
+}
+wasmparser::for_each_operator!(define_opcode_enum);

crates/wasmtime/src/module.rs

@@ -367,6 +367,7 @@ impl Module {
         types: &ModuleTypes,
     ) -> Result<(MmapVec, Option<CompiledModuleInfo>)> {
         let tunables = &engine.config().tunables;
+        let fuel_cost = engine.config().get_fuel_cost();
         let functions = mem::take(&mut translation.function_body_inputs);
         let functions = functions.into_iter().collect::<Vec<_>>();
         let compiler = engine.compiler();
@@ -377,8 +378,14 @@ impl Module {
             || -> Result<_> {
                 let funcs = engine.run_maybe_parallel(functions, |(index, func)| {
                     let offset = func.body.range().start;
-                    let result =
-                        compiler.compile_function(&translation, index, func, tunables, types);
+                    let result = compiler.compile_function(
+                        &translation,
+                        index,
+                        func,
+                        tunables,
+                        fuel_cost,
+                        types,
+                    );
                     result.with_context(|| {
                         let index = translation.module.func_index(index);
                         let name = match translation.debuginfo.name_section.func_names.get(&index) {