From 2e93be3a4c8f59e3856c22152c3771054577b618 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Sun, 5 Feb 2023 17:24:02 +0000
Subject: [PATCH] Add create_wrapper_function helper

---
 src/allocator.rs | 75 +++++++++---------------------------------------
 src/common.rs    | 38 ++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 61 deletions(-)

diff --git a/src/allocator.rs b/src/allocator.rs
index 8508227179a..1c73957ca57 100644
--- a/src/allocator.rs
+++ b/src/allocator.rs
@@ -70,37 +70,13 @@ fn codegen_inner(
             params: arg_tys.iter().cloned().map(AbiParam::new).collect(),
             returns: output.into_iter().map(AbiParam::new).collect(),
         };
-
-        let caller_name = format!("__rust_{}", method.name);
-        let callee_name = kind.fn_name(method.name);
-
-        let func_id = module.declare_function(&caller_name, Linkage::Export, &sig).unwrap();
-
-        let callee_func_id = module.declare_function(&callee_name, Linkage::Import, &sig).unwrap();
-
-        let mut ctx = Context::new();
-        ctx.func.signature = sig.clone();
-        {
-            let mut func_ctx = FunctionBuilderContext::new();
-            let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
-
-            let block = bcx.create_block();
-            bcx.switch_to_block(block);
-            let args = arg_tys
-                .into_iter()
-                .map(|ty| bcx.append_block_param(block, ty))
-                .collect::<Vec<Value>>();
-
-            let callee_func_ref = module.declare_func_in_func(callee_func_id, &mut bcx.func);
-            let call_inst = bcx.ins().call(callee_func_ref, &args);
-            let results = bcx.inst_results(call_inst).to_vec(); // Clone to prevent borrow error
-
-            bcx.ins().return_(&results);
-            bcx.seal_all_blocks();
-            bcx.finalize();
-        }
-        module.define_function(func_id, &mut ctx).unwrap();
-        unwind_context.add_function(func_id, &ctx, module.isa());
+        crate::common::create_wrapper_function(
+            module,
+            unwind_context,
+            sig,
+            &format!("__rust_{}", method.name),
+            &kind.fn_name(method.name),
+        );
     }
 
     let sig = Signature {
@@ -108,36 +84,13 @@ fn codegen_inner(
         params: vec![AbiParam::new(usize_ty), AbiParam::new(usize_ty)],
         returns: vec![],
     };
-
-    let callee_name = alloc_error_handler_kind.fn_name(sym::oom);
-
-    let func_id =
-        module.declare_function("__rust_alloc_error_handler", Linkage::Export, &sig).unwrap();
-
-    let callee_func_id = module.declare_function(&callee_name, Linkage::Import, &sig).unwrap();
-
-    let mut ctx = Context::new();
-    ctx.func.signature = sig;
-    {
-        let mut func_ctx = FunctionBuilderContext::new();
-        let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
-
-        let block = bcx.create_block();
-        bcx.switch_to_block(block);
-        let args = (&[usize_ty, usize_ty])
-            .iter()
-            .map(|&ty| bcx.append_block_param(block, ty))
-            .collect::<Vec<Value>>();
-
-        let callee_func_ref = module.declare_func_in_func(callee_func_id, &mut bcx.func);
-        bcx.ins().call(callee_func_ref, &args);
-
-        bcx.ins().trap(TrapCode::UnreachableCodeReached);
-        bcx.seal_all_blocks();
-        bcx.finalize();
-    }
-    module.define_function(func_id, &mut ctx).unwrap();
-    unwind_context.add_function(func_id, &ctx, module.isa());
+    crate::common::create_wrapper_function(
+        module,
+        unwind_context,
+        sig,
+        "__rust_alloc_error_handler",
+        &alloc_error_handler_kind.fn_name(sym::oom),
+    );
 
     let data_id = module.declare_data(OomStrategy::SYMBOL, Linkage::Export, false, false).unwrap();
     let mut data_ctx = DataContext::new();
diff --git a/src/common.rs b/src/common.rs
index f41af3a9e63..16ae526090b 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -254,6 +254,44 @@ pub(crate) fn type_sign(ty: Ty<'_>) -> bool {
     }
 }
 
+pub(crate) fn create_wrapper_function(
+    module: &mut impl Module,
+    unwind_context: &mut UnwindContext,
+    sig: Signature,
+    wrapper_name: &str,
+    callee_name: &str,
+) {
+    let wrapper_func_id = module.declare_function(wrapper_name, Linkage::Export, &sig).unwrap();
+    let callee_func_id = module.declare_function(callee_name, Linkage::Import, &sig).unwrap();
+
+    let mut ctx = Context::new();
+    ctx.func.signature = sig;
+    {
+        let mut func_ctx = FunctionBuilderContext::new();
+        let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
+
+        let block = bcx.create_block();
+        bcx.switch_to_block(block);
+        let func = &mut bcx.func.stencil;
+        let args = func
+            .signature
+            .params
+            .iter()
+            .map(|param| func.dfg.append_block_param(block, param.value_type))
+            .collect::<Vec<Value>>();
+
+        let callee_func_ref = module.declare_func_in_func(callee_func_id, &mut bcx.func);
+        let call_inst = bcx.ins().call(callee_func_ref, &args);
+        let results = bcx.inst_results(call_inst).to_vec(); // Clone to prevent borrow error
+
+        bcx.ins().return_(&results);
+        bcx.seal_all_blocks();
+        bcx.finalize();
+    }
+    module.define_function(wrapper_func_id, &mut ctx).unwrap();
+    unwind_context.add_function(wrapper_func_id, &ctx, module.isa());
+}
+
 pub(crate) struct FunctionCx<'m, 'clif, 'tcx: 'm> {
     pub(crate) cx: &'clif mut crate::CodegenCx,
     pub(crate) module: &'m mut dyn Module,