From ac81cde84e7851960b02537b7208d16ee1ddf5b7 Mon Sep 17 00:00:00 2001
From: Syrus Akbary <me@syrusakbary.com>
Date: Mon, 25 Mar 2024 18:31:57 +0100
Subject: [PATCH] Fix singlepass panic. Partially fix #4519

---
 lib/compiler-singlepass/src/codegen.rs       | 2 +-
 lib/compiler-singlepass/src/machine.rs       | 2 +-
 lib/compiler-singlepass/src/machine_arm64.rs | 6 ++++--
 lib/compiler-singlepass/src/machine_x64.rs   | 6 ++++--
 4 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/lib/compiler-singlepass/src/codegen.rs b/lib/compiler-singlepass/src/codegen.rs
index 932f38b6600..ab7a783cf4e 100644
--- a/lib/compiler-singlepass/src/codegen.rs
+++ b/lib/compiler-singlepass/src/codegen.rs
@@ -6681,7 +6681,7 @@ impl<'a, M: Machine> FuncGen<'a, M> {
         let address_map =
             get_function_address_map(self.machine.instructions_address_map(), data, body_len);
         let traps = self.machine.collect_trap_information();
-        let mut body = self.machine.assembler_finalize();
+        let mut body = self.machine.assembler_finalize()?;
         body.shrink_to_fit();
 
         Ok((
diff --git a/lib/compiler-singlepass/src/machine.rs b/lib/compiler-singlepass/src/machine.rs
index d192bc0a1ba..d4602ae77fe 100644
--- a/lib/compiler-singlepass/src/machine.rs
+++ b/lib/compiler-singlepass/src/machine.rs
@@ -233,7 +233,7 @@ pub trait Machine {
     fn new_machine_state(&self) -> MachineState;
 
     /// Finalize the assembler
-    fn assembler_finalize(self) -> Vec<u8>;
+    fn assembler_finalize(self) -> Result<Vec<u8>, CompileError>;
 
     /// get_offset of Assembler
     fn get_offset(&self) -> Offset;
diff --git a/lib/compiler-singlepass/src/machine_arm64.rs b/lib/compiler-singlepass/src/machine_arm64.rs
index ee93c992894..f5e2b0c1b18 100644
--- a/lib/compiler-singlepass/src/machine_arm64.rs
+++ b/lib/compiler-singlepass/src/machine_arm64.rs
@@ -2290,8 +2290,10 @@ impl Machine for MachineARM64 {
     }
 
     // assembler finalize
-    fn assembler_finalize(self) -> Vec<u8> {
-        self.assembler.finalize().unwrap()
+    fn assembler_finalize(self) -> Result<Vec<u8>, CompileError> {
+        self.assembler.finalize().map_err(|e| {
+            CompileError::Codegen(format!("Assembler failed finalization with: {:?}", e))
+        })
     }
 
     fn get_offset(&self) -> Offset {
diff --git a/lib/compiler-singlepass/src/machine_x64.rs b/lib/compiler-singlepass/src/machine_x64.rs
index a10e096a4b5..5bb3655d374 100644
--- a/lib/compiler-singlepass/src/machine_x64.rs
+++ b/lib/compiler-singlepass/src/machine_x64.rs
@@ -2499,8 +2499,10 @@ impl Machine for MachineX86_64 {
     }
 
     // assembler finalize
-    fn assembler_finalize(self) -> Vec<u8> {
-        self.assembler.finalize().unwrap()
+    fn assembler_finalize(self) -> Result<Vec<u8>, CompileError> {
+        self.assembler.finalize().map_err(|e| {
+            CompileError::Codegen(format!("Assembler failed finalization with: {:?}", e))
+        })
     }
 
     fn get_offset(&self) -> Offset {