wasmer_engine_jit/
engine.rs

1//! JIT compilation.
2
3use crate::{CodeMemory, JITArtifact};
4use std::sync::{Arc, Mutex};
5#[cfg(feature = "compiler")]
6use wasmer_compiler::Compiler;
7use wasmer_compiler::{
8    CompileError, CustomSection, CustomSectionProtection, FunctionBody, SectionIndex, Target,
9};
10use wasmer_engine::{Artifact, DeserializeError, Engine, EngineId, FunctionExtent, Tunables};
11use wasmer_types::entity::PrimaryMap;
12use wasmer_types::Features;
13use wasmer_types::{FunctionIndex, FunctionType, LocalFunctionIndex, SignatureIndex};
14use wasmer_vm::{
15    FunctionBodyPtr, ModuleInfo, SectionBodyPtr, SignatureRegistry, VMFunctionBody,
16    VMSharedSignatureIndex, VMTrampoline,
17};
18
19/// A WebAssembly `JIT` Engine.
20#[derive(Clone)]
21pub struct JITEngine {
22    inner: Arc<Mutex<JITEngineInner>>,
23    /// The target for the compiler
24    target: Arc<Target>,
25    engine_id: EngineId,
26}
27
28impl JITEngine {
29    /// Create a new `JITEngine` with the given config
30    #[cfg(feature = "compiler")]
31    pub fn new(compiler: Box<dyn Compiler>, target: Target, features: Features) -> Self {
32        Self {
33            inner: Arc::new(Mutex::new(JITEngineInner {
34                compiler: Some(compiler),
35                code_memory: vec![],
36                signatures: SignatureRegistry::new(),
37                features,
38            })),
39            target: Arc::new(target),
40            engine_id: EngineId::default(),
41        }
42    }
43
44    /// Create a headless `JITEngine`
45    ///
46    /// A headless engine is an engine without any compiler attached.
47    /// This is useful for assuring a minimal runtime for running
48    /// WebAssembly modules.
49    ///
50    /// For example, for running in IoT devices where compilers are very
51    /// expensive, or also to optimize startup speed.
52    ///
53    /// # Important
54    ///
55    /// Headless engines can't compile or validate any modules,
56    /// they just take already processed Modules (via `Module::serialize`).
57    pub fn headless() -> Self {
58        Self {
59            inner: Arc::new(Mutex::new(JITEngineInner {
60                #[cfg(feature = "compiler")]
61                compiler: None,
62                code_memory: vec![],
63                signatures: SignatureRegistry::new(),
64                features: Features::default(),
65            })),
66            target: Arc::new(Target::default()),
67            engine_id: EngineId::default(),
68        }
69    }
70
71    pub(crate) fn inner(&self) -> std::sync::MutexGuard<'_, JITEngineInner> {
72        self.inner.lock().unwrap()
73    }
74
75    pub(crate) fn inner_mut(&self) -> std::sync::MutexGuard<'_, JITEngineInner> {
76        self.inner.lock().unwrap()
77    }
78}
79
80impl Engine for JITEngine {
81    /// The target
82    fn target(&self) -> &Target {
83        &self.target
84    }
85
86    /// Register a signature
87    fn register_signature(&self, func_type: &FunctionType) -> VMSharedSignatureIndex {
88        let compiler = self.inner();
89        compiler.signatures().register(func_type)
90    }
91
92    /// Lookup a signature
93    fn lookup_signature(&self, sig: VMSharedSignatureIndex) -> Option<FunctionType> {
94        let compiler = self.inner();
95        compiler.signatures().lookup(sig)
96    }
97
98    /// Validates a WebAssembly module
99    fn validate(&self, binary: &[u8]) -> Result<(), CompileError> {
100        self.inner().validate(binary)
101    }
102
103    /// Compile a WebAssembly binary
104    #[cfg(feature = "compiler")]
105    fn compile(
106        &self,
107        binary: &[u8],
108        tunables: &dyn Tunables,
109    ) -> Result<Arc<dyn Artifact>, CompileError> {
110        Ok(Arc::new(JITArtifact::new(&self, binary, tunables)?))
111    }
112
113    /// Compile a WebAssembly binary
114    #[cfg(not(feature = "compiler"))]
115    fn compile(
116        &self,
117        _binary: &[u8],
118        _tunables: &dyn Tunables,
119    ) -> Result<Arc<dyn Artifact>, CompileError> {
120        Err(CompileError::Codegen(
121            "The JITEngine is operating in headless mode, so it can not compile Modules."
122                .to_string(),
123        ))
124    }
125
126    /// Deserializes a WebAssembly module
127    unsafe fn deserialize(&self, bytes: &[u8]) -> Result<Arc<dyn Artifact>, DeserializeError> {
128        Ok(Arc::new(JITArtifact::deserialize(&self, &bytes)?))
129    }
130
131    fn id(&self) -> &EngineId {
132        &self.engine_id
133    }
134
135    fn cloned(&self) -> Arc<dyn Engine + Send + Sync> {
136        Arc::new(self.clone())
137    }
138}
139
140/// The inner contents of `JITEngine`
141pub struct JITEngineInner {
142    /// The compiler
143    #[cfg(feature = "compiler")]
144    compiler: Option<Box<dyn Compiler>>,
145    /// The features to compile the Wasm module with
146    features: Features,
147    /// The code memory is responsible of publishing the compiled
148    /// functions to memory.
149    code_memory: Vec<CodeMemory>,
150    /// The signature registry is used mainly to operate with trampolines
151    /// performantly.
152    signatures: SignatureRegistry,
153}
154
155impl JITEngineInner {
156    /// Gets the compiler associated to this engine.
157    #[cfg(feature = "compiler")]
158    pub fn compiler(&self) -> Result<&dyn Compiler, CompileError> {
159        if self.compiler.is_none() {
160            return Err(CompileError::Codegen("The JITEngine is operating in headless mode, so it can only execute already compiled Modules.".to_string()));
161        }
162        Ok(&**self.compiler.as_ref().unwrap())
163    }
164
165    /// Validate the module
166    #[cfg(feature = "compiler")]
167    pub fn validate<'data>(&self, data: &'data [u8]) -> Result<(), CompileError> {
168        self.compiler()?.validate_module(self.features(), data)
169    }
170
171    /// Validate the module
172    #[cfg(not(feature = "compiler"))]
173    pub fn validate<'data>(&self, _data: &'data [u8]) -> Result<(), CompileError> {
174        Err(CompileError::Validate(
175            "The JITEngine is not compiled with compiler support, which is required for validating"
176                .to_string(),
177        ))
178    }
179
180    /// The Wasm features
181    pub fn features(&self) -> &Features {
182        &self.features
183    }
184
185    /// Allocate compiled functions into memory
186    #[allow(clippy::type_complexity)]
187    pub(crate) fn allocate(
188        &mut self,
189        _module: &ModuleInfo,
190        functions: &PrimaryMap<LocalFunctionIndex, FunctionBody>,
191        function_call_trampolines: &PrimaryMap<SignatureIndex, FunctionBody>,
192        dynamic_function_trampolines: &PrimaryMap<FunctionIndex, FunctionBody>,
193        custom_sections: &PrimaryMap<SectionIndex, CustomSection>,
194    ) -> Result<
195        (
196            PrimaryMap<LocalFunctionIndex, FunctionExtent>,
197            PrimaryMap<SignatureIndex, VMTrampoline>,
198            PrimaryMap<FunctionIndex, FunctionBodyPtr>,
199            PrimaryMap<SectionIndex, SectionBodyPtr>,
200        ),
201        CompileError,
202    > {
203        let function_bodies = functions
204            .values()
205            .chain(function_call_trampolines.values())
206            .chain(dynamic_function_trampolines.values())
207            .collect::<Vec<_>>();
208        let (executable_sections, data_sections): (Vec<_>, _) = custom_sections
209            .values()
210            .partition(|section| section.protection == CustomSectionProtection::ReadExecute);
211        self.code_memory.push(CodeMemory::new());
212
213        let (mut allocated_functions, allocated_executable_sections, allocated_data_sections) =
214            self.code_memory
215                .last_mut()
216                .unwrap()
217                .allocate(
218                    function_bodies.as_slice(),
219                    executable_sections.as_slice(),
220                    data_sections.as_slice(),
221                )
222                .map_err(|message| {
223                    CompileError::Resource(format!(
224                        "failed to allocate memory for functions: {}",
225                        message
226                    ))
227                })?;
228
229        let allocated_functions_result = allocated_functions
230            .drain(0..functions.len())
231            .map(|slice| FunctionExtent {
232                ptr: FunctionBodyPtr(slice.as_ptr()),
233                length: slice.len(),
234            })
235            .collect::<PrimaryMap<LocalFunctionIndex, _>>();
236
237        let mut allocated_function_call_trampolines: PrimaryMap<SignatureIndex, VMTrampoline> =
238            PrimaryMap::new();
239        for ptr in allocated_functions
240            .drain(0..function_call_trampolines.len())
241            .map(|slice| slice.as_ptr())
242        {
243            let trampoline =
244                unsafe { std::mem::transmute::<*const VMFunctionBody, VMTrampoline>(ptr) };
245            allocated_function_call_trampolines.push(trampoline);
246        }
247
248        let allocated_dynamic_function_trampolines = allocated_functions
249            .drain(..)
250            .map(|slice| FunctionBodyPtr(slice.as_ptr()))
251            .collect::<PrimaryMap<FunctionIndex, _>>();
252
253        let mut exec_iter = allocated_executable_sections.iter();
254        let mut data_iter = allocated_data_sections.iter();
255        let allocated_custom_sections = custom_sections
256            .iter()
257            .map(|(_, section)| {
258                SectionBodyPtr(
259                    if section.protection == CustomSectionProtection::ReadExecute {
260                        exec_iter.next()
261                    } else {
262                        data_iter.next()
263                    }
264                    .unwrap()
265                    .as_ptr(),
266                )
267            })
268            .collect::<PrimaryMap<SectionIndex, _>>();
269
270        Ok((
271            allocated_functions_result,
272            allocated_function_call_trampolines,
273            allocated_dynamic_function_trampolines,
274            allocated_custom_sections,
275        ))
276    }
277
278    /// Make memory containing compiled code executable.
279    pub(crate) fn publish_compiled_code(&mut self) {
280        self.code_memory.last_mut().unwrap().publish();
281    }
282
283    /// Register DWARF-type exception handling information associated with the code.
284    pub(crate) fn publish_eh_frame(&mut self, eh_frame: Option<&[u8]>) -> Result<(), CompileError> {
285        self.code_memory
286            .last_mut()
287            .unwrap()
288            .unwind_registry_mut()
289            .publish(eh_frame)
290            .map_err(|e| {
291                CompileError::Resource(format!("Error while publishing the unwind code: {}", e))
292            })?;
293        Ok(())
294    }
295
296    /// Shared signature registry.
297    pub fn signatures(&self) -> &SignatureRegistry {
298        &self.signatures
299    }
300}