1use 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#[derive(Clone)]
21pub struct JITEngine {
22 inner: Arc<Mutex<JITEngineInner>>,
23 target: Arc<Target>,
25 engine_id: EngineId,
26}
27
28impl JITEngine {
29 #[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 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 fn target(&self) -> &Target {
83 &self.target
84 }
85
86 fn register_signature(&self, func_type: &FunctionType) -> VMSharedSignatureIndex {
88 let compiler = self.inner();
89 compiler.signatures().register(func_type)
90 }
91
92 fn lookup_signature(&self, sig: VMSharedSignatureIndex) -> Option<FunctionType> {
94 let compiler = self.inner();
95 compiler.signatures().lookup(sig)
96 }
97
98 fn validate(&self, binary: &[u8]) -> Result<(), CompileError> {
100 self.inner().validate(binary)
101 }
102
103 #[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 #[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 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
140pub struct JITEngineInner {
142 #[cfg(feature = "compiler")]
144 compiler: Option<Box<dyn Compiler>>,
145 features: Features,
147 code_memory: Vec<CodeMemory>,
150 signatures: SignatureRegistry,
153}
154
155impl JITEngineInner {
156 #[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 #[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 #[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 pub fn features(&self) -> &Features {
182 &self.features
183 }
184
185 #[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 pub(crate) fn publish_compiled_code(&mut self) {
280 self.code_memory.last_mut().unwrap().publish();
281 }
282
283 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 pub fn signatures(&self) -> &SignatureRegistry {
298 &self.signatures
299 }
300}