chore(cranelift): unify rayon and non-rayon code with closure (#5896)

Author: marxin

examples/throw_exception.rs

@@ -122,7 +122,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let instance = Instance::new(&mut store, &module, &import_object)?;
 
     // Here we go.
-    let f: TypedFunction<(), ()> = instance.exports.get_function("f")?.typed(&mut store)?;
+    let f: TypedFunction<(), ()> = instance.exports.get_function("f")?.typed(&store)?;
 
     println!("Calling `f` function...");
     let result = f.call(&mut store);

lib/compiler-cranelift/src/compiler.rs

@@ -118,133 +118,9 @@ impl CraneliftCompiler {
             }
         };
 
-        #[cfg_attr(not(feature = "unwind"), allow(unused_mut))]
-        let mut custom_sections = PrimaryMap::new();
-
-        #[cfg(not(feature = "rayon"))]
-        let mut func_translator = FuncTranslator::new();
-        #[cfg(not(feature = "rayon"))]
-        #[cfg_attr(not(feature = "unwind"), allow(unused_variables))]
-        let (functions, fdes): (Vec<CompiledFunction>, Vec<_>) = function_body_inputs
-            .iter()
-            .collect::<Vec<(LocalFunctionIndex, &FunctionBodyData<'_>)>>()
-            .into_iter()
-            .map(|(i, input)| {
-                let func_index = module.func_index(i);
-                let mut context = Context::new();
-                let mut func_env = FuncEnvironment::new(
-                    isa.frontend_config(),
-                    module,
-                    &signatures,
-                    &memory_styles,
-                    table_styles,
-                );
-                context.func.name = match get_function_name(&mut context.func, func_index) {
-                    ExternalName::User(nameref) => {
-                        if context.func.params.user_named_funcs().is_valid(nameref) {
-                            let name = &context.func.params.user_named_funcs()[nameref];
-                            UserFuncName::User(name.clone())
-                        } else {
-                            UserFuncName::default()
-                        }
-                    }
-                    ExternalName::TestCase(testcase) => UserFuncName::Testcase(testcase),
-                    _ => UserFuncName::default(),
-                };
-                context.func.signature = signatures[module.functions[func_index]].clone();
-                // if generate_debug_info {
-                //     context.func.collect_debug_info();
-                // }
-                let mut reader =
-                    MiddlewareBinaryReader::new_with_offset(input.data, input.module_offset);
-                reader.set_middleware_chain(
-                    self.config
-                        .middlewares
-                        .generate_function_middleware_chain(i),
-                );
-
-                func_translator.translate(
-                    module_translation_state,
-                    &mut reader,
-                    &mut context.func,
-                    &mut func_env,
-                    i,
-                )?;
-
-                let mut code_buf: Vec<u8> = Vec::new();
-                let mut ctrl_plane = Default::default();
-                let func_name_map = context.func.params.user_named_funcs().clone();
-                let result = context
-                    .compile(&*isa, &mut ctrl_plane)
-                    .map_err(|error| CompileError::Codegen(error.inner.to_string()))?;
-                code_buf.extend_from_slice(result.code_buffer());
-
-                let func_relocs = result
-                    .buffer
-                    .relocs()
-                    .into_iter()
-                    .map(|r| mach_reloc_to_reloc(module, &func_name_map, r))
-                    .collect::<Vec<_>>();
-
-                let traps = result
-                    .buffer
-                    .traps()
-                    .into_iter()
-                    .map(mach_trap_to_trap)
-                    .collect::<Vec<_>>();
-
-                let (unwind_info, fde) = match compiled_function_unwind_info(&*isa, &context)? {
-                    #[cfg(feature = "unwind")]
-                    CraneliftUnwindInfo::Fde(fde) => {
-                        if dwarf_frametable.is_some() {
-                            let fde = fde.to_fde(Address::Symbol {
-                                // The symbol is the kind of relocation.
-                                // "0" is used for functions
-                                symbol: WriterRelocate::FUNCTION_SYMBOL,
-                                // We use the addend as a way to specify the
-                                // function index
-                                addend: i.index() as _,
-                            });
-                            // The unwind information is inserted into the dwarf section
-                            (Some(CompiledFunctionUnwindInfo::Dwarf), Some(fde))
-                        } else {
-                            (None, None)
-                        }
-                    }
-                    #[cfg(feature = "unwind")]
-                    other => (other.maybe_into_to_windows_unwind(), None),
-
-                    // This is a bit hacky, but necessary since gimli is not
-                    // available when the "unwind" feature is disabled.
-                    #[cfg(not(feature = "unwind"))]
-                    other => (other.maybe_into_to_windows_unwind(), None::<()>),
-                };
-
-                let range = reader.range();
-                let address_map = get_function_address_map(&context, range, code_buf.len());
-
-                Ok((
-                    CompiledFunction {
-                        body: FunctionBody {
-                            body: code_buf,
-                            unwind_info,
-                        },
-                        relocations: func_relocs,
-                        frame_info: CompiledFunctionFrameInfo { address_map, traps },
-                    },
-                    fde,
-                ))
-            })
-            .collect::<Result<Vec<_>, CompileError>>()?
-            .into_iter()
-            .unzip();
-        #[cfg(feature = "rayon")]
-        #[cfg_attr(not(feature = "unwind"), allow(unused_variables))]
-        let (functions, fdes): (Vec<CompiledFunction>, Vec<_>) = function_body_inputs
-            .iter()
-            .collect::<Vec<(LocalFunctionIndex, &FunctionBodyData<'_>)>>()
-            .par_iter()
-            .map_init(FuncTranslator::new, |func_translator, (i, input)| {
+        let compile_function =
+            |func_translator: &mut FuncTranslator,
+             (i, input): (&LocalFunctionIndex, &FunctionBodyData)| {
                 let func_index = module.func_index(*i);
                 let mut context = Context::new();
                 let mut func_env = FuncEnvironment::new(
@@ -368,6 +244,31 @@ impl CraneliftCompiler {
                     },
                     fde,
                 ))
+            };
+
+        #[cfg_attr(not(feature = "unwind"), allow(unused_mut))]
+        let mut custom_sections = PrimaryMap::new();
+
+        #[cfg(not(feature = "rayon"))]
+        let mut func_translator = FuncTranslator::new();
+        #[cfg(not(feature = "rayon"))]
+        #[cfg_attr(not(feature = "unwind"), allow(unused_variables))]
+        let (functions, fdes): (Vec<CompiledFunction>, Vec<_>) = function_body_inputs
+            .iter()
+            .collect::<Vec<(LocalFunctionIndex, &FunctionBodyData<'_>)>>()
+            .into_iter()
+            .map(|(i, input)| compile_function(&mut func_translator, (&i, input)))
+            .collect::<Result<Vec<_>, CompileError>>()?
+            .into_iter()
+            .unzip();
+        #[cfg(feature = "rayon")]
+        #[cfg_attr(not(feature = "unwind"), allow(unused_variables))]
+        let (functions, fdes): (Vec<CompiledFunction>, Vec<_>) = function_body_inputs
+            .iter()
+            .collect::<Vec<(LocalFunctionIndex, &FunctionBodyData<'_>)>>()
+            .par_iter()
+            .map_init(FuncTranslator::new, |func_translator, &(i, input)| {
+                compile_function(func_translator, (&i, input))
             })
             .collect::<Result<Vec<_>, CompileError>>()?
             .into_iter()
@@ -402,7 +303,7 @@ impl CraneliftCompiler {
             .map(|sig| make_trampoline_function_call(&self.config().callbacks, &*isa, &mut cx, sig))
             .collect::<Result<Vec<FunctionBody>, CompileError>>()?
             .into_iter()
-            .collect::<PrimaryMap<SignatureIndex, FunctionBody>>();
+            .collect();
         #[cfg(feature = "rayon")]
         let function_call_trampolines = module
             .signatures
@@ -414,7 +315,7 @@ impl CraneliftCompiler {
             })
             .collect::<Result<Vec<FunctionBody>, CompileError>>()?
             .into_iter()
-            .collect::<PrimaryMap<SignatureIndex, FunctionBody>>();
+            .collect();
 
         use wasmer_types::VMOffsets;
         let offsets = VMOffsets::new_for_trampolines(frontend_config.pointer_bytes());
@@ -437,7 +338,7 @@ impl CraneliftCompiler {
             })
             .collect::<Result<Vec<_>, CompileError>>()?
             .into_iter()
-            .collect::<PrimaryMap<FunctionIndex, FunctionBody>>();
+            .collect();
         #[cfg(feature = "rayon")]
         let dynamic_function_trampolines = module
             .imported_function_types()
@@ -454,7 +355,7 @@ impl CraneliftCompiler {
             })
             .collect::<Result<Vec<_>, CompileError>>()?
             .into_iter()
-            .collect::<PrimaryMap<FunctionIndex, FunctionBody>>();
+            .collect();
 
         let got = wasmer_compiler::types::function::GOT::empty();