wasmer_engine_jit/
artifact.rs

1//! Define `JITArtifact` to allow compiling and instantiating to be
2//! done as separate steps.
3
4use crate::engine::{JITEngine, JITEngineInner};
5use crate::link::link_module;
6#[cfg(feature = "compiler")]
7use crate::serialize::SerializableCompilation;
8use crate::serialize::SerializableModule;
9use std::sync::{Arc, Mutex};
10use wasmer_compiler::{CompileError, Features, Triple};
11#[cfg(feature = "compiler")]
12use wasmer_compiler::{CompileModuleInfo, ModuleEnvironment};
13use wasmer_engine::{
14    register_frame_info, Artifact, DeserializeError, FunctionExtent, GlobalFrameInfoRegistration,
15    SerializeError,
16};
17#[cfg(feature = "compiler")]
18use wasmer_engine::{Engine, SerializableFunctionFrameInfo, Tunables};
19use wasmer_types::entity::{BoxedSlice, PrimaryMap};
20use wasmer_types::{
21    FunctionIndex, LocalFunctionIndex, MemoryIndex, OwnedDataInitializer, SignatureIndex,
22    TableIndex,
23};
24use wasmer_vm::{
25    FunctionBodyPtr, MemoryStyle, ModuleInfo, TableStyle, VMSharedSignatureIndex, VMTrampoline,
26};
27
28/// A compiled wasm module, ready to be instantiated.
29pub struct JITArtifact {
30    serializable: SerializableModule,
31    finished_functions: BoxedSlice<LocalFunctionIndex, FunctionBodyPtr>,
32    finished_function_call_trampolines: BoxedSlice<SignatureIndex, VMTrampoline>,
33    finished_dynamic_function_trampolines: BoxedSlice<FunctionIndex, FunctionBodyPtr>,
34    signatures: BoxedSlice<SignatureIndex, VMSharedSignatureIndex>,
35    frame_info_registration: Mutex<Option<GlobalFrameInfoRegistration>>,
36    finished_function_lengths: BoxedSlice<LocalFunctionIndex, usize>,
37}
38
39impl JITArtifact {
40    const MAGIC_HEADER: &'static [u8] = b"\0wasmer-jit";
41
42    /// Check if the provided bytes look like a serialized `JITArtifact`.
43    pub fn is_deserializable(bytes: &[u8]) -> bool {
44        bytes.starts_with(Self::MAGIC_HEADER)
45    }
46
47    /// Compile a data buffer into a `JITArtifact`, which may then be instantiated.
48    #[cfg(feature = "compiler")]
49    pub fn new(
50        jit: &JITEngine,
51        data: &[u8],
52        tunables: &dyn Tunables,
53    ) -> Result<Self, CompileError> {
54        let environ = ModuleEnvironment::new();
55        let mut inner_jit = jit.inner_mut();
56        let features = inner_jit.features();
57
58        let translation = environ.translate(data).map_err(CompileError::Wasm)?;
59
60        let memory_styles: PrimaryMap<MemoryIndex, MemoryStyle> = translation
61            .module
62            .memories
63            .values()
64            .map(|memory_type| tunables.memory_style(memory_type))
65            .collect();
66        let table_styles: PrimaryMap<TableIndex, TableStyle> = translation
67            .module
68            .tables
69            .values()
70            .map(|table_type| tunables.table_style(table_type))
71            .collect();
72
73        let mut compile_info = CompileModuleInfo {
74            module: Arc::new(translation.module),
75            features: features.clone(),
76            memory_styles,
77            table_styles,
78        };
79
80        let compiler = inner_jit.compiler()?;
81
82        // Compile the Module
83        let compilation = compiler.compile_module(
84            &jit.target(),
85            &mut compile_info,
86            // SAFETY: Calling `unwrap` is correct since
87            // `environ.translate()` above will write some data into
88            // `module_translation_state`.
89            translation.module_translation_state.as_ref().unwrap(),
90            translation.function_body_inputs,
91        )?;
92        let function_call_trampolines = compilation.get_function_call_trampolines();
93        let dynamic_function_trampolines = compilation.get_dynamic_function_trampolines();
94
95        let data_initializers = translation
96            .data_initializers
97            .iter()
98            .map(OwnedDataInitializer::new)
99            .collect::<Vec<_>>()
100            .into_boxed_slice();
101
102        let frame_infos = compilation
103            .get_frame_info()
104            .values()
105            .map(|frame_info| SerializableFunctionFrameInfo::Processed(frame_info.clone()))
106            .collect::<PrimaryMap<LocalFunctionIndex, _>>();
107
108        let serializable_compilation = SerializableCompilation {
109            function_bodies: compilation.get_function_bodies(),
110            function_relocations: compilation.get_relocations(),
111            function_jt_offsets: compilation.get_jt_offsets(),
112            function_frame_info: frame_infos,
113            function_call_trampolines,
114            dynamic_function_trampolines,
115            custom_sections: compilation.get_custom_sections(),
116            custom_section_relocations: compilation.get_custom_section_relocations(),
117            debug: compilation.get_debug(),
118        };
119        let serializable = SerializableModule {
120            compilation: serializable_compilation,
121            compile_info,
122            data_initializers,
123        };
124        Self::from_parts(&mut inner_jit, serializable)
125    }
126
127    /// Compile a data buffer into a `JITArtifact`, which may then be instantiated.
128    #[cfg(not(feature = "compiler"))]
129    pub fn new(_jit: &JITEngine, _data: &[u8]) -> Result<Self, CompileError> {
130        Err(CompileError::Codegen(
131            "Compilation is not enabled in the engine".to_string(),
132        ))
133    }
134
135    /// Deserialize a JITArtifact
136    pub fn deserialize(jit: &JITEngine, bytes: &[u8]) -> Result<Self, DeserializeError> {
137        if !Self::is_deserializable(bytes) {
138            return Err(DeserializeError::Incompatible(
139                "The provided bytes are not wasmer-jit".to_string(),
140            ));
141        }
142
143        let inner_bytes = &bytes[Self::MAGIC_HEADER.len()..];
144
145        // let r = flexbuffers::Reader::get_root(bytes).map_err(|e| DeserializeError::CorruptedBinary(format!("{:?}", e)))?;
146        // let serializable = SerializableModule::deserialize(r).map_err(|e| DeserializeError::CorruptedBinary(format!("{:?}", e)))?;
147
148        let serializable: SerializableModule = bincode::deserialize(inner_bytes)
149            .map_err(|e| DeserializeError::CorruptedBinary(format!("{:?}", e)))?;
150
151        Self::from_parts(&mut jit.inner_mut(), serializable).map_err(DeserializeError::Compiler)
152    }
153
154    /// Construct a `JITArtifact` from component parts.
155    pub fn from_parts(
156        inner_jit: &mut JITEngineInner,
157        serializable: SerializableModule,
158    ) -> Result<Self, CompileError> {
159        let (
160            finished_functions,
161            finished_function_call_trampolines,
162            finished_dynamic_function_trampolines,
163            custom_sections,
164        ) = inner_jit.allocate(
165            &serializable.compile_info.module,
166            &serializable.compilation.function_bodies,
167            &serializable.compilation.function_call_trampolines,
168            &serializable.compilation.dynamic_function_trampolines,
169            &serializable.compilation.custom_sections,
170        )?;
171
172        link_module(
173            &serializable.compile_info.module,
174            &finished_functions,
175            &serializable.compilation.function_jt_offsets,
176            serializable.compilation.function_relocations.clone(),
177            &custom_sections,
178            &serializable.compilation.custom_section_relocations,
179        );
180
181        // Compute indices into the shared signature table.
182        let signatures = {
183            let signature_registry = inner_jit.signatures();
184            serializable
185                .compile_info
186                .module
187                .signatures
188                .values()
189                .map(|sig| signature_registry.register(sig))
190                .collect::<PrimaryMap<_, _>>()
191        };
192
193        let eh_frame = match &serializable.compilation.debug {
194            Some(debug) => {
195                let eh_frame_section_size = serializable.compilation.custom_sections
196                    [debug.eh_frame]
197                    .bytes
198                    .len();
199                let eh_frame_section_pointer = custom_sections[debug.eh_frame];
200                Some(unsafe {
201                    std::slice::from_raw_parts(*eh_frame_section_pointer, eh_frame_section_size)
202                })
203            }
204            None => None,
205        };
206        // Make all code compiled thus far executable.
207        inner_jit.publish_compiled_code();
208
209        inner_jit.publish_eh_frame(eh_frame)?;
210
211        let finished_function_lengths = finished_functions
212            .values()
213            .map(|extent| extent.length)
214            .collect::<PrimaryMap<LocalFunctionIndex, usize>>()
215            .into_boxed_slice();
216        let finished_functions = finished_functions
217            .values()
218            .map(|extent| extent.ptr)
219            .collect::<PrimaryMap<LocalFunctionIndex, FunctionBodyPtr>>()
220            .into_boxed_slice();
221        let finished_function_call_trampolines =
222            finished_function_call_trampolines.into_boxed_slice();
223        let finished_dynamic_function_trampolines =
224            finished_dynamic_function_trampolines.into_boxed_slice();
225        let signatures = signatures.into_boxed_slice();
226
227        Ok(Self {
228            serializable,
229            finished_functions,
230            finished_function_call_trampolines,
231            finished_dynamic_function_trampolines,
232            signatures,
233            frame_info_registration: Mutex::new(None),
234            finished_function_lengths,
235        })
236    }
237
238    /// Get the default extension when serializing this artifact
239    pub fn get_default_extension(_triple: &Triple) -> &'static str {
240        // `.wjit` is the default extension for all the triples
241        "wjit"
242    }
243}
244
245impl Artifact for JITArtifact {
246    fn module(&self) -> Arc<ModuleInfo> {
247        self.serializable.compile_info.module.clone()
248    }
249
250    fn module_ref(&self) -> &ModuleInfo {
251        &self.serializable.compile_info.module
252    }
253
254    fn module_mut(&mut self) -> Option<&mut ModuleInfo> {
255        Arc::get_mut(&mut self.serializable.compile_info.module)
256    }
257
258    fn register_frame_info(&self) {
259        let mut info = self.frame_info_registration.lock().unwrap();
260
261        if info.is_some() {
262            return;
263        }
264
265        let finished_function_extents = self
266            .finished_functions
267            .values()
268            .copied()
269            .zip(self.finished_function_lengths.values().copied())
270            .map(|(ptr, length)| FunctionExtent { ptr, length })
271            .collect::<PrimaryMap<LocalFunctionIndex, _>>()
272            .into_boxed_slice();
273
274        let frame_infos = &self.serializable.compilation.function_frame_info;
275        *info = register_frame_info(
276            self.serializable.compile_info.module.clone(),
277            &finished_function_extents,
278            frame_infos.clone(),
279        );
280    }
281
282    fn features(&self) -> &Features {
283        &self.serializable.compile_info.features
284    }
285
286    fn data_initializers(&self) -> &[OwnedDataInitializer] {
287        &*self.serializable.data_initializers
288    }
289
290    fn memory_styles(&self) -> &PrimaryMap<MemoryIndex, MemoryStyle> {
291        &self.serializable.compile_info.memory_styles
292    }
293
294    fn table_styles(&self) -> &PrimaryMap<TableIndex, TableStyle> {
295        &self.serializable.compile_info.table_styles
296    }
297
298    fn finished_functions(&self) -> &BoxedSlice<LocalFunctionIndex, FunctionBodyPtr> {
299        &self.finished_functions
300    }
301
302    fn finished_function_call_trampolines(&self) -> &BoxedSlice<SignatureIndex, VMTrampoline> {
303        &self.finished_function_call_trampolines
304    }
305
306    fn finished_dynamic_function_trampolines(&self) -> &BoxedSlice<FunctionIndex, FunctionBodyPtr> {
307        &self.finished_dynamic_function_trampolines
308    }
309
310    fn signatures(&self) -> &BoxedSlice<SignatureIndex, VMSharedSignatureIndex> {
311        &self.signatures
312    }
313
314    fn serialize(&self) -> Result<Vec<u8>, SerializeError> {
315        // let mut s = flexbuffers::FlexbufferSerializer::new();
316        // self.serializable.serialize(&mut s).map_err(|e| SerializeError::Generic(format!("{:?}", e)));
317        // Ok(s.take_buffer())
318        let bytes = bincode::serialize(&self.serializable)
319            .map_err(|e| SerializeError::Generic(format!("{:?}", e)))?;
320
321        // Prepend the header.
322        let mut serialized = Self::MAGIC_HEADER.to_vec();
323        serialized.extend(bytes);
324        Ok(serialized)
325    }
326}
wasmer_engine_jit/artifact.rs

wasmer_engine_jit/
artifact.rs
