[Swift] Add arm64 calling conventions

The Swift ABI repurposes three callee-saved registers for implicit
parameters (self, error, and async context). Supporting the various
combinations of these requires registering several different calling
conventions.

The demangler is taught to explicitly apply these calling conventions to
functions that need them.

3 files changed, 199 insertions, 3 deletions

diff --git a/arch/arm64/arch_arm64.cpp b/arch/arm64/arch_arm64.cpp
index b71e79e8..a56ed35b 100644
--- a/arch/arm64/arch_arm64.cpp
+++ b/arch/arm64/arch_arm64.cpp
@@ -2765,6 +2765,115 @@ class AppleArm64SystemCallConvention : public CallingConvention
 	virtual bool IsEligibleForHeuristics() override { return false; }
 };
 
+
+// Swift calling convention for ARM64.
+//
+// The Swift ABI repurposes three callee-saved registers for implicit parameters:
+//   x20 - self (swiftself): context/self parameter, passed as the last integer argument
+//   x21 - error (swifterror): caller initializes to zero, callee sets to error pointer on throw
+//   x22 - async context (swiftasync): implicit async context for async functions
+//
+// Each combination of these three properties requires a distinct calling convention,
+// since it changes which registers are arguments, callee-saved, or implicitly defined.
+// The naming scheme is "swift" with optional suffixes: "-self", "-throws", "-async".
+class SwiftArm64CallingConvention : public CallingConvention
+{
+	bool m_hasSelf;
+	bool m_throws;
+	bool m_isAsync;
+
+public:
+	SwiftArm64CallingConvention(Architecture* arch, const string& name, bool hasSelf, bool throws, bool isAsync)
+	    : CallingConvention(arch, name), m_hasSelf(hasSelf), m_throws(throws), m_isAsync(isAsync)
+	{
+	}
+
+	virtual vector<uint32_t> GetIntegerArgumentRegisters() override
+	{
+		// Self goes first because our function types list self as the first
+		// parameter. Parameters are assigned to registers sequentially, so
+		// this ensures self -> x20 and remaining args -> x0-x7.
+		// Async context and error go last as implicit trailing registers.
+		vector<uint32_t> regs;
+		if (m_hasSelf)
+			regs.push_back(REG_X20);
+		regs.insert(regs.end(), {REG_X0, REG_X1, REG_X2, REG_X3, REG_X4, REG_X5, REG_X6, REG_X7});
+		if (m_isAsync)
+			regs.push_back(REG_X22);
+		if (m_throws)
+			regs.push_back(REG_X21);
+		return regs;
+	}
+
+	virtual vector<uint32_t> GetFloatArgumentRegisters() override
+	{
+		return vector<uint32_t> {REG_V0, REG_V1, REG_V2, REG_V3, REG_V4, REG_V5, REG_V6, REG_V7};
+	}
+
+	virtual vector<uint32_t> GetCallerSavedRegisters() override
+	{
+		vector<uint32_t> regs {REG_X0, REG_X1, REG_X2, REG_X3, REG_X4, REG_X5, REG_X6, REG_X7, REG_X8,
+		    REG_X9, REG_X10, REG_X11, REG_X12, REG_X13, REG_X14, REG_X15, REG_X16, REG_X17, REG_X18,
+		    REG_X30, REG_V0, REG_V1, REG_V2, REG_V3, REG_V4, REG_V5, REG_V6, REG_V7, REG_V16, REG_V17,
+		    REG_V18, REG_V19, REG_V20, REG_V21, REG_V22, REG_V23, REG_V24, REG_V25, REG_V26, REG_V27,
+		    REG_V28, REG_V29, REG_V30, REG_V31};
+		// When used as special registers, they are no longer callee-saved.
+		if (m_hasSelf)
+			regs.push_back(REG_X20);
+		if (m_throws)
+			regs.push_back(REG_X21);
+		if (m_isAsync)
+			regs.push_back(REG_X22);
+		return regs;
+	}
+
+	virtual vector<uint32_t> GetCalleeSavedRegisters() override
+	{
+		vector<uint32_t> regs {REG_X19, REG_X23, REG_X24, REG_X25, REG_X26, REG_X27, REG_X28, REG_X29};
+		// Only include x20/x21/x22 as callee-saved when they are NOT repurposed.
+		if (!m_hasSelf)
+			regs.push_back(REG_X20);
+		if (!m_throws)
+			regs.push_back(REG_X21);
+		if (!m_isAsync)
+			regs.push_back(REG_X22);
+		return regs;
+	}
+
+	virtual vector<uint32_t> GetImplicitlyDefinedRegisters() override
+	{
+		vector<uint32_t> regs;
+		// Throwing functions implicitly define x21 (the error register) on return.
+		if (m_throws)
+			regs.push_back(REG_X21);
+		return regs;
+	}
+
+	virtual uint32_t GetIntegerReturnValueRegister() override { return REG_X0; }
+
+	virtual uint32_t GetFloatReturnValueRegister() override { return REG_V0; }
+
+	virtual vector<uint32_t> GetRequiredArgumentRegisters() override
+	{
+		vector<uint32_t> regs;
+		if (m_hasSelf)
+			regs.push_back(REG_X20);
+		if (m_throws)
+			regs.push_back(REG_X21);
+		if (m_isAsync)
+			regs.push_back(REG_X22);
+		return regs;
+	}
+
+	virtual bool AreArgumentRegistersUsedForVarArgs() override { return false; }
+
+	virtual bool IsEligibleForHeuristics() override
+	{
+		return m_hasSelf || m_throws || m_isAsync;
+	}
+};
+
+
 #define PAGE(x)        (uint32_t)((x) >> 12)
 #define PAGE_OFF(x)    (uint32_t)((x)&0xfff)
 #define PAGE_NO_OFF(x) (uint32_t)((x)&0xFFFFF000)
@@ -3550,6 +3659,20 @@ extern "C"
 		conv = new AppleArm64CallingConvention(arm64);
 		arm64->RegisterCallingConvention(conv);
 
+		// Register Swift calling conventions (all combinations of self/throws/async).
+		for (int swiftFlags = 0; swiftFlags < 8; swiftFlags++)
+		{
+			bool hasSelf = (swiftFlags & 1) != 0;
+			bool throws  = (swiftFlags & 2) != 0;
+			bool isAsync = (swiftFlags & 4) != 0;
+			string name = "swift";
+			if (hasSelf) name += "-self";
+			if (throws)  name += "-throws";
+			if (isAsync) name += "-async";
+			conv = new SwiftArm64CallingConvention(arm64, name, hasSelf, throws, isAsync);
+			arm64->RegisterCallingConvention(conv);
+		}
+
 		for (uint32_t i = REG_X0; i <= REG_X28; i++)
 		{
 			if (i == REG_X16 || i == REG_X17 || i == REG_X18) // reserved by os.
diff --git a/plugins/workflow_swift/src/demangler/function_type.rs b/plugins/workflow_swift/src/demangler/function_type.rs
index 3313143e..908f5940 100644
--- a/plugins/workflow_swift/src/demangler/function_type.rs
+++ b/plugins/workflow_swift/src/demangler/function_type.rs
@@ -10,6 +10,43 @@ use swift_demangler::{
 
 use super::type_reconstruction::{make_named_type_ref, TypeRefExt};
 
+/// Architecture-specific register assignments for Swift implicit parameters.
+struct PlatformAbi {
+    arch: CoreArchitecture,
+    error_reg: &'static str,
+    async_context_reg: &'static str,
+}
+
+impl PlatformAbi {
+    fn for_arch(arch: &CoreArchitecture) -> Option<Self> {
+        match arch.name().as_ref() {
+            "aarch64" => Some(Self {
+                arch: *arch,
+                error_reg: "x21",
+                async_context_reg: "x22",
+            }),
+            _ => None,
+        }
+    }
+
+    fn error_location(&self) -> Option<Variable> {
+        self.register_variable(self.error_reg)
+    }
+
+    fn async_context_location(&self) -> Option<Variable> {
+        self.register_variable(self.async_context_reg)
+    }
+
+    fn register_variable(&self, name: &str) -> Option<Variable> {
+        let reg = self.arch.register_by_name(name)?;
+        Some(Variable::new(
+            VariableSourceType::RegisterVariableSourceType,
+            0,
+            reg.id().0 as i64,
+        ))
+    }
+}
+
 /// Swift calling convention builder.
 ///
 /// Tracks which implicit parameters (self, error, async context) are present,
@@ -17,6 +54,7 @@ use super::type_reconstruction::{make_named_type_ref, TypeRefExt};
 /// architecture-specific calling convention when building the final function type.
 struct CallingConvention {
     arch: CoreArchitecture,
+    abi: Option<PlatformAbi>,
     flags: u8,
     leading_params: Vec<FunctionParameter>,
     trailing_params: Vec<FunctionParameter>,
@@ -30,6 +68,7 @@ impl CallingConvention {
     fn for_arch(arch: &CoreArchitecture) -> Self {
         Self {
             arch: *arch,
+            abi: PlatformAbi::for_arch(arch),
             flags: 0,
             leading_params: Vec::new(),
             trailing_params: Vec::new(),
@@ -94,6 +133,8 @@ impl CallingConvention {
     /// Prepends `self` before `params` and appends error/async context after,
     /// then looks up the appropriate calling convention by name on the architecture.
     fn build_type(self, ret_type: &Type, params: Vec<FunctionParameter>) -> Ref<Type> {
+        let cc_name = self.cc_name();
+
         let mut all_params = self.leading_params;
         all_params.extend(params);
         all_params.extend(self.trailing_params);
@@ -103,7 +144,7 @@ impl CallingConvention {
             all_params.push(FunctionParameter {
                 ty: error_ty.into(),
                 name: "error".to_string(),
-                location: None,
+                location: self.abi.as_ref().and_then(|a| a.error_location()),
             });
         }
 
@@ -112,11 +153,29 @@ impl CallingConvention {
             all_params.push(FunctionParameter {
                 ty: ptr_ty.into(),
                 name: "asyncContext".to_string(),
-                location: None,
+                location: self.abi.as_ref().and_then(|a| a.async_context_location()),
             });
         }
 
-        Type::function(ret_type, all_params, false)
+        if let Some(cc) = self.arch.calling_convention_by_name(&cc_name) {
+            Type::function_with_opts(ret_type, &all_params, false, cc, Conf::new(0, 0))
+        } else {
+            Type::function(ret_type, all_params, false)
+        }
+    }
+
+    fn cc_name(&self) -> String {
+        let mut name = String::from("swift");
+        if self.flags & Self::SELF != 0 {
+            name.push_str("-self");
+        }
+        if self.flags & Self::THROWS != 0 {
+            name.push_str("-throws");
+        }
+        if self.flags & Self::ASYNC != 0 {
+            name.push_str("-async");
+        }
+        name
     }
 }
 
diff --git a/rust/src/architecture.rs b/rust/src/architecture.rs
index 6b644e88..e854d53b 100644
--- a/rust/src/architecture.rs
+++ b/rust/src/architecture.rs
@@ -1262,6 +1262,20 @@ pub trait ArchitectureExt: Architecture {
         }
     }
 
+    fn calling_convention_by_name(&self, name: &str) -> Option<Ref<CoreCallingConvention>> {
+        let name = name.to_cstr();
+        unsafe {
+            let result = NonNull::new(BNGetArchitectureCallingConventionByName(
+                self.as_ref().handle,
+                name.as_ptr(),
+            ))?;
+            Some(CoreCallingConvention::ref_from_raw(
+                result.as_ptr(),
+                self.as_ref().handle(),
+            ))
+        }
+    }
+
     fn calling_conventions(&self) -> Array<CoreCallingConvention> {
         unsafe {
             let mut count = 0;