From d7308c361d8fdea44d0420e924aa5ba37b60b34e Mon Sep 17 00:00:00 2001
From: losfair <zhy20000919@hotmail.com>
Date: Tue, 19 Nov 2019 01:25:01 +0800
Subject: [PATCH] Fix call_indirect on imported functions.

---
 lib/singlepass-backend/src/codegen_x64.rs     | 23 ++++--
 lib/singlepass-backend/src/emitter_x64.rs     | 11 +++
 .../src/translator_aarch64.rs                 | 72 +++++++++++++++++++
 3 files changed, 100 insertions(+), 6 deletions(-)

diff --git a/lib/singlepass-backend/src/codegen_x64.rs b/lib/singlepass-backend/src/codegen_x64.rs
index 58ca6a111..cd2a729cb 100644
--- a/lib/singlepass-backend/src/codegen_x64.rs
+++ b/lib/singlepass-backend/src/codegen_x64.rs
@@ -472,12 +472,15 @@ impl ModuleCodeGenerator<X64FunctionCode, X64ExecutionContext, CodegenError>
     for X64ModuleCodeGenerator
 {
     fn new() -> X64ModuleCodeGenerator {
+        let mut a = Assembler::new().unwrap();
+        a.notify_begin();
+
         X64ModuleCodeGenerator {
             functions: vec![],
             signatures: None,
             function_signatures: None,
             function_labels: Some(HashMap::new()),
-            assembler: Some(Assembler::new().unwrap()),
+            assembler: Some(a),
             func_import_count: 0,
             config: None,
         }
@@ -556,7 +559,7 @@ impl ModuleCodeGenerator<X64FunctionCode, X64ExecutionContext, CodegenError>
         mut self,
         _: &ModuleInfo,
     ) -> Result<(X64ExecutionContext, Box<dyn CacheGen>), CodegenError> {
-        let (assembler, function_labels, breakpoints) = match self.functions.last_mut() {
+        let (mut assembler, function_labels, breakpoints) = match self.functions.last_mut() {
             Some(x) => (
                 x.assembler.take().unwrap(),
                 x.function_labels.take().unwrap(),
@@ -569,6 +572,7 @@ impl ModuleCodeGenerator<X64FunctionCode, X64ExecutionContext, CodegenError>
             ),
         };
 
+        assembler.notify_end();
         let total_size = assembler.get_offset().0;
         let _output = assembler.finalize().unwrap();
         let mut output = CodeMemory::new(_output.len());
@@ -5139,10 +5143,17 @@ impl FunctionCodeGenerator<CodegenError> for X64FunctionCode {
                     a,
                     &mut self.machine,
                     |a| {
-                        a.emit_call_location(Location::Memory(
-                            GPR::RAX,
-                            (vm::Anyfunc::offset_func() as usize) as i32,
-                        ));
+                        if a.arch_requires_indirect_call_trampoline() {
+                            a.arch_emit_indirect_call_with_trampoline(Location::Memory(
+                                GPR::RAX,
+                                (vm::Anyfunc::offset_func() as usize) as i32,
+                            ));
+                        } else {
+                            a.emit_call_location(Location::Memory(
+                                GPR::RAX,
+                                (vm::Anyfunc::offset_func() as usize) as i32,
+                            ));
+                        }
                     },
                     params.iter().map(|x| *x),
                     Some((&mut self.fsm, &mut self.control_stack)),
diff --git a/lib/singlepass-backend/src/emitter_x64.rs b/lib/singlepass-backend/src/emitter_x64.rs
index f426b6d02..5a7855435 100644
--- a/lib/singlepass-backend/src/emitter_x64.rs
+++ b/lib/singlepass-backend/src/emitter_x64.rs
@@ -271,6 +271,17 @@ pub trait Emitter {
     fn arch_supports_canonicalize_nan(&self) -> bool {
         true
     }
+
+    fn arch_requires_indirect_call_trampoline(&self) -> bool {
+        false
+    }
+
+    fn arch_emit_indirect_call_with_trampoline(&mut self, _loc: Location) {
+        unimplemented!()
+    }
+
+    fn notify_begin(&mut self) {}
+    fn notify_end(&mut self) {}
 }
 
 fn _dummy(a: &mut Assembler) {
diff --git a/lib/singlepass-backend/src/translator_aarch64.rs b/lib/singlepass-backend/src/translator_aarch64.rs
index d6361fc1a..a26e2073b 100644
--- a/lib/singlepass-backend/src/translator_aarch64.rs
+++ b/lib/singlepass-backend/src/translator_aarch64.rs
@@ -1702,6 +1702,78 @@ impl Emitter for Assembler {
     fn arch_supports_canonicalize_nan(&self) -> bool {
         false
     }
+
+    fn arch_requires_indirect_call_trampoline(&self) -> bool {
+        true
+    }
+
+    fn arch_emit_indirect_call_with_trampoline(&mut self, loc: Location) {
+        match loc {
+            Location::GPR(x) => {
+                dynasm!(self
+                    // Push return address.
+                    ; sub x_rsp, x_rsp, 8
+                    ; adr x_tmp1, >done
+                    ; str x_tmp1, [x_rsp]
+                );
+                dynasm!(self
+                    ; adr x_tmp1, >program_end
+                    ; cmp X(map_gpr(x).x()), x_tmp1
+                    ; b.ge >external
+                    ; adr x_tmp1, <program_begin
+                    ; cmp X(map_gpr(x).x()), x_tmp1
+                    ; b.lt >external
+                    ; br X(map_gpr(x).x())
+                    ; external:
+                );
+                self.emit_homomorphic_host_redirection(x);
+                dynasm!(self ; done: );
+            }
+            Location::Memory(base, disp) => {
+                if disp >= 0 {
+                    dynasm!(self ; b >after ; disp: ; .dword disp ; after: ; ldr w_tmp3, <disp ; add x_tmp3, x_tmp3, X(map_gpr(base).x()));
+                } else {
+                    dynasm!(self ; b >after ; disp: ; .dword -disp ; after: ; ldr w_tmp3, <disp ; sub x_tmp3, X(map_gpr(base).x()), x_tmp3);
+                }
+                dynasm!(self
+                    // Push return address.
+                    ; sub x_rsp, x_rsp, 8
+                    ; adr x_tmp1, >done
+                    ; str x_tmp1, [x_rsp]
+
+                    // Read memory.
+                    ; ldr X(map_gpr(GPR::RAX).x()), [x_tmp3]
+                );
+                dynasm!(self
+                    ; adr x_tmp1, >program_end
+                    ; cmp X(map_gpr(GPR::RAX).x()), x_tmp1
+                    ; b.ge >external
+                    ; adr x_tmp1, <program_begin
+                    ; cmp X(map_gpr(GPR::RAX).x()), x_tmp1
+                    ; b.lt >external
+                    ; br X(map_gpr(GPR::RAX).x())
+                    ; external:
+                );
+                self.emit_homomorphic_host_redirection(GPR::RAX);
+                dynasm!(self ; done: );
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    fn notify_begin(&mut self) {
+        dynasm!(
+            self
+            ; program_begin:
+        );
+    }
+
+    fn notify_end(&mut self) {
+        dynasm!(
+            self
+            ; program_end:
+        );
+    }
 }
 
 fn emit_clz_variant(