From 439238ac4a5721f81703c266cc463bfdb4914fa6 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Tue, 28 May 2024 10:14:58 -0400
Subject: [PATCH] libia2: Implement ia2_thread_begin

---
 runtime/libia2/threads.c | 31 ++++++++++++++++++++++++++-----
 1 file changed, 26 insertions(+), 5 deletions(-)

diff --git a/runtime/libia2/threads.c b/runtime/libia2/threads.c
index f0cf6398d..e0624c4e4 100644
--- a/runtime/libia2/threads.c
+++ b/runtime/libia2/threads.c
@@ -33,9 +33,9 @@ void *ia2_thread_begin(void *arg) {
   size_t tag = ia2_get_tag();
   void **new_sp_addr = ia2_stackptr_for_tag(tag);
 
-#if LIBIA2_X86_64
   /* Switch to the stack for this compartment, then call `fn(data)`. */
   void *result;
+#if LIBIA2_X86_64
   __asm__ volatile(
       /* clang-format off */
       // Copy stack pointer to rdi.
@@ -61,12 +61,33 @@ void *ia2_thread_begin(void *arg) {
       : "=a"(result)
       : [fn] "r"(fn), [data] "r"(data), [new_sp_addr] "r"(new_sp_addr)
       : "rdi");
-  /* clang-format on */
-  return result;
+      /* clang-format on */
 #elif LIBIA2_AARCH64
-#warning "libia2 does not implement ia2_thread_begin yet"
-  __builtin_trap();
+#warning "ia2_thread_begin does not align the stack correctly"
+  __asm__ volatile(
+        // Copy stack pointer to x10
+        "mov x10, sp\n"
+        // Load the stack pointer for this compartment's stack
+        "ldr x0, [%[new_sp_addr]]\n"
+        "mov sp, x0\n"
+        // Push the old stack pointer
+        "str x10, [sp, #-8]!\n"
+        // Load argument
+        "ldr x0, [%[data]]\n"
+        // Call fn(data)
+        "blr %[fn]\n"
+        // x0 now contains ret value
+        "mov %[result], x0\n"
+        // Pop the old stack pointer
+        "ldr x10, [sp], #8\n"
+        // Switch stacks back
+        "mov sp, x10\n"
+      : [result] "=r"(result)
+      : [fn] "r"(fn), [data] "r"(&data), [new_sp_addr] "r"(new_sp_addr)
+      : "x10");
 #endif
+
+  return result;
 }
 
 int __real_pthread_create(pthread_t *restrict thread,