Improve the design of context switch

src/fiber/context/aarch64.cr

@@ -3,20 +3,29 @@
 class Fiber
   # :nodoc:
   def makecontext(stack_ptr, fiber_main) : Nil
-    # in ARMv8, the context switch push/pop 12 registers and 8 FPU registers,
-    # and one more to store the argument of `fiber_main` (+ alignment), we thus
-    # reserve space for 22 pointers:
-    @context.stack_top = (stack_ptr - 22).as(Void*)
-    @context.resumable = 1
+    # Initial Stack
+    #
+    # +-------------------------+
+    # |   dummy return address  |
+    # +-------------------------+
+    # |      entry function     |
+    # +-------------------------+
+    # |      fiber address      |
+    # +-------------------------+
+    # |     helper function     | ---> load first argument and jump to entry function
+    # +-------------------------+
 
-    stack_ptr[-2] = self.as(Void*)      # x0 (r0): puts `self` as first argument for `fiber_main`
-    stack_ptr[-14] = fiber_main.pointer # x30 (lr): initial `resume` will `ret` to this address
+    @context.stack_top = (stack_ptr - 3).as(Void*)
+    @context.resumable = 1
+    stack_ptr[0] = Pointer(Void).null
+    stack_ptr[-1] = fiber_main.pointer
+    stack_ptr[-2] = self.as(Void*)
+    stack_ptr[-3] = (->Fiber.load_first_argument).pointer
   end
 
-  # :nodoc:
   @[NoInline]
   @[Naked]
-  def self.swapcontext(current_context, new_context) : Nil
+  private def self.suspend_context(current_context, new_context, resume_func)
     # adapted from https://github.com/ldc-developers/druntime/blob/ldc/src/core/threadasm.S
     #
     # preserve/restore AAPCS64 registers:
@@ -26,17 +35,19 @@ class Fiber
     # x0        self argument (initial call)
     # d8-d15    5.1.2 says callee only must save bottom 64-bits (the "d" regs)
     #
-    # ARM assembly requires integer literals to be moved to a register before
-    # being stored at an address; we use x19 as a scratch register that will be
-    # overwritten by the new context.
-    #
-    # AArch64 assembly also requires a register to load/store the stack top
-    # pointer. We use x19 as a scratch register again.
+    # Stack information:
     #
-    # Eventually reset LR to zero to avoid the ARM unwinder to mistake the
-    # context switch as a regular call.
+    # |           :          |
+    # +----------------------+
+    # |        x19-x30       |
+    # +----------------------+
+    # |         d8-d15       |
+    # +----------------------+
+    # |    resume_context    | <--- stack top
+    # +----------------------+
+
     asm("
-      stp     d15, d14, [sp, #-22*8]!
+      stp     d15, d14, [sp, #-20*8]!
       stp     d13, d12, [sp, #2*8]
       stp     d11, d10, [sp, #4*8]
       stp     d9,  d8,  [sp, #6*8]
@@ -46,7 +57,9 @@ class Fiber
       stp     x24, x23, [sp, #14*8]
       stp     x22, x21, [sp, #16*8]
       stp     x20, x19, [sp, #18*8]
-      stp     x0,  x1,  [sp, #20*8] // push 1st argument (+ alignment)
+
+      // push resume_context address
+      str     $2, [sp, #-8]!
 
       mov     x19, sp               // current_context.stack_top = sp
       str     x19, [$0, #0]
@@ -58,7 +71,14 @@ class Fiber
       ldr     x19, [$1, #0]         // sp = new_context.stack_top (x19)
       mov     sp, x19
 
-      ldp     x0,  x1,  [sp, #20*8] // pop 1st argument (+ alignement)
+      ldr     x30, [sp], #8
+      " :: "r"(current_context), "r"(new_context), "r"(resume_func))
+  end
+
+  @[NoInline]
+  @[Naked]
+  private def self.resume_context
+    asm("
       ldp     x20, x19, [sp, #18*8]
       ldp     x22, x21, [sp, #16*8]
       ldp     x24, x23, [sp, #14*8]
@@ -68,12 +88,39 @@ class Fiber
       ldp     d9,  d8,  [sp, #6*8]
       ldp     d11, d10, [sp, #4*8]
       ldp     d13, d12, [sp, #2*8]
-      ldp     d15, d14, [sp], #22*8
+      ldp     d15, d14, [sp], #20*8
 
       // avoid a stack corruption that will confuse the unwinder
       mov     x16, x30 // save lr
       mov     x30, #0  // reset lr
       br      x16      // jump to new pc value
-      " :: "r"(current_context), "r"(new_context))
+      ")
+  end
+
+  # :nodoc:
+  def self.swapcontext(current_context, new_context) : Nil
+    suspend_context current_context, new_context, (->resume_context).pointer
+  end
+
+  @[NoInline]
+  @[Naked]
+  protected def self.load_first_argument
+    # Stack requirement
+    #
+    # |            :           |
+    # |            :           |
+    # +------------------------+
+    # |  next return address   | ---> for lr register
+    # +------------------------+
+    # |  target function addr  | ---> for pc register
+    # +------------------------+
+    # |     first argument     | ---> for x0 register
+    # +------------------------+
+
+    asm("
+      ldp     x16, x30, [sp, #8]
+      ldr     x0, [sp], #24
+      br      x16
+      ")
   end
 end

src/fiber/context/arm.cr

@@ -2,71 +2,130 @@
 
 class Fiber
   # :nodoc:
-  def makecontext(stack_ptr, fiber_main) : Void*
-    # in ARMv6 / ARVMv7, the context switch push/pop 8 registers, add one more
-    # to store the argument of `fiber_main`, and 8 64-bit FPU registers if a FPU
-    # is present, we thus reserve space for 9 or 25 pointers:
-    {% if flag?(:armhf) %}
-      @context.stack_top = (stack_ptr - 25).as(Void*)
-    {% else %}
-      @context.stack_top = (stack_ptr - 9).as(Void*)
-    {% end %}
-    @context.resumable = 1
+  def makecontext(stack_ptr, fiber_main) : Nil
+    # Initial Stack
+    #
+    # +-------------------------+
+    # |   dummy return address  |
+    # +-------------------------+
+    # |      entry function     |
+    # +-------------------------+
+    # |      fiber address      |
+    # +-------------------------+
+    # |     helper function     | ---> load first argument and jump to entry function
+    # +-------------------------+
 
-    stack_ptr[0] = fiber_main.pointer # lr: initial `resume` will `ret` to this address
-    stack_ptr[-9] = self.as(Void*)    # r0: puts `self` as first argument for `fiber_main`
+    @context.stack_top = (stack_ptr - 3).as(Void*)
+    @context.resumable = 1
+    stack_ptr[0] = Pointer(Void).null
+    stack_ptr[-1] = fiber_main.pointer
+    stack_ptr[-2] = self.as(Void*)
+    stack_ptr[-3] = (->Fiber.load_first_argument).pointer
   end
 
-  # :nodoc:
   @[NoInline]
   @[Naked]
-  def self.swapcontext(current_context, new_context) : Nil
+  private def self.suspend_context(current_context, new_context, resume_func)
     # ARM assembly requires integer literals to be moved to a register before
     # being stored at an address; we use r4 as a scratch register that will be
     # overwritten by the new context.
     #
-    # Eventually reset LR to zero to avoid the ARM unwinder to mistake the
-    # context switch as a regular call.
+    # Stack Information
+    #
+    # |           :           |
+    # +-----------------------+
+    # |      link register    |
+    # +-----------------------+
+    # | callee-saved register |
+    # +-----------------------+
+    # |     resume_context    |
+    # +-----------------------+
 
     {% if flag?(:armhf) %}
       asm("
         // declare the presence of a conservative FPU to the ASM compiler
         .fpu vfp
 
-        stmdb  sp!, {r0, r4-r11, lr}  // push 1st argument + callee-saved registers
+        stmdb  sp!, {r4-r11, lr}      // push callee-saved registers + return address
         vstmdb sp!, {d8-d15}          // push FPU registers
+
+        // store resume_context address
+        stmdb  sp!, {$2}
+
         str    sp, [$0, #0]           // current_context.stack_top = sp
         mov    r4, #1                 // current_context.resumable = 1
         str    r4, [$0, #4]
 
         mov    r4, #0                 // new_context.resumable = 0
         str    r4, [$1, #4]
         ldr    sp, [$1, #0]           // sp = new_context.stack_top
-        vldmia sp!, {d8-d15}          // pop FPU registers
-        ldmia  sp!, {r0, r4-r11, lr}  // pop 1st argument + calleed-saved registers
 
-        // avoid a stack corruption that will confuse the unwinder
-        mov    r1, lr
-        mov    lr, #0
-        mov    pc, r1
-        " :: "r"(current_context), "r"(new_context))
+        ldmia  sp!, {pc}
+        " :: "r"(current_context), "r"(new_context), "r"(resume_func))
     {% elsif flag?(:arm) %}
       asm("
-        stmdb  sp!, {r0, r4-r11, lr}  // push 1st argument + calleed-saved registers
+        stmdb  sp!, {r4-r11, lr}      // push calleed-saved registers + return address
+
+        // store resume_context address
+        stmdb  sp!, {$2}
+
         str    sp, [$0, #0]           // current_context.stack_top = sp
         mov    r4, #1                 // current_context.resumable = 1
         str    r4, [$0, #4]
 
         mov    r4, #0                 // new_context.resumable = 0
         str    r4, [$1, #4]
         ldr    sp, [$1, #0]           // sp = new_context.stack_top
-        ldmia  sp!, {r0, r4-r11, lr}  // pop 1st argument + calleed-saved registers
 
+        ldmia  sp, {pc}
+        " :: "r"(current_context), "r"(new_context), "r"(resume_func))
+    {% end %}
+  end
+
+  @[NoInline]
+  @[Naked]
+  private def self.resume_context
+    {% if flag?(:armhf) %}
+      asm("
+        // avoid a stack corruption that will confuse the unwinder
+        mov    lr, #0
+
+        vldmia sp!, {d8-d15}          // pop FPU registers
+        ldmia  sp!, {r4-r11, pc}      // pop calleed-saved registers and return
+        ")
+    {% elsif flag?(:arm) %}
+      asm("
         // avoid a stack corruption that will confuse the unwinder
-        mov    r1, lr
         mov    lr, #0
-        mov    pc, r1
-        " :: "r"(current_context), "r"(new_context))
+
+        ldmia  sp!, {r4-r11, pc}      // pop calleed-saved registers and return
+        ")
     {% end %}
   end
+
+  # :nodoc:
+  def self.swapcontext(current_context, new_context) : Nil
+    suspend_context current_context, new_context, (->resume_context).pointer
+  end
+
+  @[NoInline]
+  @[Naked]
+  protected def self.load_first_argument
+    # Stack requirement
+    #
+    # |            :           |
+    # |            :           |
+    # +------------------------+
+    # |  next return address   | ---> for lr register
+    # +------------------------+
+    # |  target function addr  | ---> for pc register
+    # +------------------------+
+    # |     first argument     | ---> for r0 register
+    # +------------------------+
+
+    asm("
+        ldmia sp!, {r0, r4, lr}
+        mov pc, r4
+        ")
+  end
 end

src/fiber/context/i686.cr

@@ -2,36 +2,77 @@
 
 class Fiber
   # :nodoc:
-  def makecontext(stack_ptr, fiber_main)
-    # in IA32 (x86), the context switch push/pop 4 registers, and we need two
-    # more to store the argument for `fiber_main` and keep the stack aligned on
-    # 16 bytes, we thus reserve space for 6 pointers:
-    @context.stack_top = (stack_ptr - 6).as(Void*)
-    @context.resumable = 1
+  def makecontext(stack_ptr, fiber_main) : Nil
+    # Initial Stack
+    #
+    # +-----------------------+
+    # |      fiber address    |
+    # +-----------------------+
+    # |     dummy address     |
+    # +-----------------------+
+    # |     entry function    |
+    # +-----------------------+
 
-    stack_ptr[0] = self.as(Void*)      # first argument passed on the stack
-    stack_ptr[-1] = Pointer(Void).null # empty space to keep the stack alignment (16 bytes)
-    stack_ptr[-2] = fiber_main.pointer # initial `resume` will `ret` to this address
+    @context.stack_top = (stack_ptr - 2).as(Void*)
+    @context.resumable = 1
+    stack_ptr[0] = self.as(Void*)
+    stack_ptr[-1] = Pointer(Void).null
+    stack_ptr[-2] = fiber_main.pointer
   end
 
-  # :nodoc:
   @[NoInline]
   @[Naked]
-  def self.swapcontext(current_context, new_context) : Nil
+  private def self.suspend_context(current_context, new_context, resume_func)
+    # Stack Information
+    #
+    # |            :             |
+    # +--------------------------+
+    # |   callee-saved register  |
+    # +--------------------------+
+    # |       resume context     | <--- stack top
+    # +--------------------------+
+
     asm("
-      pushl %edi        // push 1st argument (because of initial resume)
       pushl %ebx        // push callee-saved registers on the stack
       pushl %ebp
       pushl %esi
+      pushl $2          // push resume_context function_pointer
       movl %esp, 0($0)  // current_context.stack_top = %esp
       movl $$1, 4($0)   // current_context.resumable = 1
 
       movl $$0, 4($1)   // new_context.resumable = 0
       movl 0($1), %esp  // %esp = new_context.stack_top
-      popl %esi         // pop callee-saved registers from the stack
-      popl %ebp
-      popl %ebx
-      popl %edi         // pop first argument (for initial resume)
-      " :: "r"(current_context), "r"(new_context))
+      " :: "r"(current_context), "r"(new_context), "r"(resume_func))
+  end
+
+  @[NoInline]
+  @[Naked]
+  private def self.resume_context
+    asm("
+        popl %esi
+        popl %ebp
+        popl %ebx
+        ")
+  end
+
+  # :nodoc:
+  def self.swapcontext(current_context, new_context) : Nil
+    suspend_context current_context, new_context, (->resume_context).pointer
+  end
+
+  @[NoInline]
+  @[Naked]
+  protected def self.load_first_argument
+    # Stack requirement
+    #
+    # |            :           |
+    # |            :           |
+    # +------------------------+
+    # |  target function addr  | ---> for pc register
+    # +------------------------+
+    # |     first argument     | ---> for edi register
+    # +------------------------+
+
+    asm("popl %edi")
   end
 end