Minimum code for proving the context switch assebly on Windows

closer_to_crystal_impl.cr

STACK_SIZE         = 1024 * 1024 * 8
PUSH_POP_REG_COUNT = 7

class FiberPoc
  @stack : Array(UInt8)
  property rsp : UInt8* = Pointer(UInt8).null
  property f = Proc(Void).new { }

  def initialize
    @stack = Array.new(STACK_SIZE, 0_u8) # OK, as long as we don't push/pop -> reallocate
  end

  def stack
    @stack
  end

  def run
    @f.call
  end

  def makecontext(proc : Proc)
    @f = proc
     # we pop 7 registers before we return

    s_ptr = @stack.to_unsafe
    s_start = s_ptr + (STACK_SIZE - 32)

    # Our entry
    @rsp = (s_start - 7*8)
    s_start_ptr = s_start.as(UInt64*)
    s_start_ptr.value = ->(f : FiberPoc){f.run}.pointer.address

    # First parameter
    first_param = (s_start - 8).as(UInt64*)
    first_param.value = self.as(Void*).address

    p "fiber_address: #{self.as(Void*)}"
  end
end

class Scheduler
  @fibers = Array(FiberPoc).new(3)
  @current = 0
  @running = 1

  def initialize
    @fibers = Array(FiberPoc).new
    (0..2).each do |i|
      @fibers << FiberPoc.new
    end
    # base thread
    @fibers[0].makecontext(->{})
  end

  def spawn(fiber_no : UInt64, f : Proc)
    p "spawn"
    @running += 1
    @fibers[fiber_no].makecontext(f)
  end

  def run
    self.yield_control
  end

  def yield_control
    p "yield running: #{@running}, current: #{@current}"
    if @running == 2
      # do nothing
    elsif @running == 3
      if @current == 0
        p "old: 0, new: 1, from: #{@fibers[0].rsp}, to: #{@fibers[1].rsp}"
        @current = 1
        Scheduler.swapcontext(@fibers[0].rsp, @fibers[1].rsp)
      elsif @current == 1
        p "old: 1, new: 2, from: #{@fibers[1].rsp}, to: #{@fibers[2].rsp}"
        @current = 2
        Scheduler.swapcontext(@fibers[1].rsp, @fibers[2].rsp)
      elsif @current == 2
        p "old: 2, new: 1, from: #{@fibers[2].rsp}, to: #{@fibers[1].rsp}"
        @current = 1
        Scheduler.swapcontext(@fibers[2].rsp, @fibers[1].rsp)
      end
    end
  end

  def ret
    p "ret"
    @running -= 1
    if @running == 1
      p "Done. Exiting."
      Process.exit(0)
    end
    self.yield_control
  end

  @[NoInline]
  @[Naked]
  def self.swapcontext(current : UInt8*, to : UInt8*) : Nil
    asm("
        pushq %rdi
        pushq %rbx
        pushq %rbp
        pushq %r12
        pushq %r13
        pushq %r14
        pushq %r15
        movq %rsp, ($0)

        movq $1, %rsp
        popq %r15
        popq %r14
        popq %r13
        popq %r12
        popq %rbp
        popq %rbx
        popq %rdi"
              :: "r"(current), "r"(to))
  end
end

fun test
    puts "Hello"
end

class Executor
  @scheduler = Scheduler.new

  def main
    # hello1 = ->{
    #   (0..10).each do |i|
    #     puts("Fiber1: #{i}")
    #     @scheduler.yield_control
    #   end
    #   @scheduler.ret
    # }
    # hello2 = ->{
    #   (0..10).each do |i|
    #     puts("Fiber2: #{i}")
    #     @scheduler.yield_control
    #   end
    #   @scheduler.ret
    # }

      hello1 = ->{
      puts("Fiber2: HELLO")
        @scheduler.yield_control
        puts("Fiber2: HELLO2")
        @scheduler.ret
    }
    hello2 = ->{
        puts("Fiber2: HELLO")
        @scheduler.yield_control
        puts("Fiber2: HELLO2")
      @scheduler.ret
    }

    @scheduler.spawn(1, hello1)
    @scheduler.spawn(2, hello2)
    @scheduler.run
  end
end

executor = Executor.new
executor.main

crystal_ctx_switch_win_poc.cr

STACK_SIZE = 100

class FiberPoc
    @stack : Array(UInt8)
    def initialize
        @stack = Array.new(STACK_SIZE, 0_u8)
    end
    def stack
        @stack
    end

    def makecontext (f :  Proc)
        p f
        s_ptr = @stack.to_unsafe
        s_start = s_ptr + (STACK_SIZE - 32)
        s_start_ptr = s_start.as(UInt64*)
        s_start_ptr.value = ->start(Void*, Void*).pointer.as(UInt64*).value
        # # ===== UNNCOMMENT FOR DEBUGGING =====
        # p_ptr = ->start(Void*, Void*).pointer.as(UInt64*).address
        # p_ptr = Pointer(Void).new(p_ptr)
        # proc = Proc(Void*, Void*, Void).new(p_ptr, Pointer(Void).null)
        # proc.call(f.pointer, f.closure_data)
        # p @stack
    end
end

fun start(p_ptr : Void*, p_closure : Void*)
    fn = Proc(Void).new(p_ptr, p_closure)
    fn.call
end


class Scheduler
    @fibers : Array(FiberPoc)
    @current = 0
    @running = 0

    def initialize
        @fibers = Array(FiberPoc).new
        (0..1).each do |i|
           @fibers << FiberPoc.new
        end
    end

    def spawn (fiber_no : UInt64, f : Proc)
        @running += 1
        @fibers[fiber_no].makecontext(f)
    end

    def yield_control
        new_fiber : UInt64
        old_fiber : UInt64
        if @running == 2
            if @current = 0
                #switch(0, 1)
                p "old: 0, new: 1"
            else
                p "new: 0, old: 1"
            end
        end
    end

    def ret
        @running -= 1
        if @running == 0
            p "Done. Exiting."
            Process.exit(0)
        end
        self.yield_control
    end
end

fun switch(old : Void*, new : Void*)
    p "Switching"
end

class Executor
    @scheduler = Scheduler.new

    def main
        hello1 = ->{
            (0..10).each do |i|
            puts("Fiber1: {}", i)
            @scheduler.yield_control
            end
            @scheduler.ret
        }
        hello2 = ->{
            (0..10).each do |i|
            puts("Fiber2: {}", i)
            @scheduler.yield_control
            end
            @scheduler.ret
        }
        
        @scheduler.spawn(0, hello1)
        @scheduler.spawn(0, hello2)
    end

end

executor = Executor.new
executor.main()

dev.cr

STACK_SIZE = 1024 * 1024 * 8
# give the OS some time to kill the process if it's bad
SLEEP_INTERVAL = 0.5

# Defined here as well just for easier editing. The two things that are interesting to play with is the `swapcontext`
# and the `makecontext` functions
class Scheduler
  @[NoInline]
  @[Naked]
  def self.swapcontext(current : Pointer(Pointer(Void)), to : Pointer(Void)) : Nil
    
     {% if flag?(:win32) %}
      asm("
        pushq %rcx
        pushq %gs: 0x10
        pushq %gs:0x08
        pushq %rdi
        pushq %rbx
        pushq %rbp
        pushq %rsi
        pushq %r12
        pushq %r13
        pushq %r14
        pushq %r15
        movq %rsp, ($0)
        movq ($1), %rsp
        popq %r15
        popq %r14
        popq %r13
        popq %r12
        popq %rsi
        popq %rbp
        popq %rbx
        popq %rdi
        popq %gs:0x08
        popq %gs:010
        popq %rcx
        "
              :
              : "r"(current), "r"(to)
              :
              : "volatile", "alignstack"
              )
    {% else %}
      asm("
        pushq %rdi
        pushq %rbx
        pushq %rbp
        pushq %r12
        pushq %r13
        pushq %r14
        pushq %r15
        leaq -0x08(%rsp), %rsp # testing 16 byte align, just remove, remember to change @rsp = (s_start - 7*8)
        movq %rsp, ($0)

        movq $1, %rsp
        leaq 0x08(%rsp), %rsp # testing 16 byte align, just remove, remember to change @rsp = (s_start - 7*8)
        popq %r15
        popq %r14
        popq %r13
        popq %r12
        popq %rbp
        popq %rbx
        popq %rdi
        "
              :: "r"(current), "r"(to))
    {% end %}
  end
end

class FiberPoc
  {% if flag?(:win32) %}
  def makecontext(proc : Proc)
    @f = proc

    s_ptr = @stack.to_unsafe
    s_start = s_ptr + (STACK_SIZE - 32) # note that this is different from the original crystal impl (uses offset of 16 IRC)
    puts s_start.address
    # Our entry
    # 8 registers + 2 qwords for NT_TIB + 1 parameter 
    @rsp = (s_start - 11*8).as(Void*) 
    s_start_ptr = s_start.as(UInt64*)
    s_start_ptr.value = ->(f : FiberPoc) { f.run }.pointer.address

    # https://en.wikipedia.org/wiki/Win32_Thread_Information_Block

    # stack end - "stack limit" - low address
    stack_limit = (s_start - 24).as(UInt64*)
    stack_limit.value = s_ptr.address

    # stack top - "stack base" - high address
    stack_base = (s_start - 16).as(UInt64*)
    stack_base.value = (s_ptr + STACK_SIZE).address

    # First parameter
    first_param = (s_start - 8).as(UInt64*)
    first_param.value = self.as(Void*).address
  end
  {% else %}
  def makecontext(proc : Proc)
    @f = proc

    
    s_ptr = @stack.to_unsafe
    s_start = s_ptr + (STACK_SIZE - 32) # note that this is different from the original crystal impl (uses offset of 16 IRC)
    # Our entry
    # we pop 7 registers before we return
    @rsp = (s_start - 8*8).as(Void*)
    s_start_ptr = s_start.as(UInt64*)
    s_start_ptr.value = ->(f : FiberPoc) { f.run }.pointer.address

    # First parameter
    first_param = (s_start - 8).as(UInt64*)
    first_param.value = self.as(Void*).address
  end
  {% end %}
  

  property stack : Array(UInt8)
  property rsp : Void* = Pointer(Void).null
  property f = Proc(Void).new { }

  def initialize
    @stack = Array.new(STACK_SIZE, 0_u8) # OK, as long as we don't push/pop -> reallocate
  end

  def run
    @f.call
  end
end

class Scheduler
  @fibers = Array(FiberPoc).new(3)
  @current = 0
  @running = 1
  @finished : Int32 = 0 # just a hack since we only have two executing fibers, this is the one that finished first

  def initialize
    @fibers = Array(FiberPoc).new
    (0..2).each do |i|
      @fibers << FiberPoc.new
    end
    # base thread
    @fibers[0].makecontext(->{})
  end

  def spawn(fiber_no : UInt64, f : Proc)
    puts "Queuing fiber #{fiber_no}"
    @running += 1
    @fibers[fiber_no].makecontext(f)
  end

  def run
    self.yield_control
  end

  def yield_control
    # if we only have two threads, one is finished so we only swap if it just finished to continoue on the last
    if @running == 2
      if @current == @finished
        next_ctx = @finished == 2 ? 1 : 2
        current = @fibers[next_ctx]
        @current = next_ctx
        Scheduler.swapcontext(pointerof(current.@rsp), current.@rsp)
      end
    elsif @running == 3
      if @current == 0
        current = @fibers[0]
        # p "old: 0, new: 1, from: #{@fibers[0].rsp}, to: #{@fibers[1].rsp}"
        @current = 1
        Scheduler.swapcontext(pointerof(current.@rsp), @fibers[1].@rsp)
      elsif @current == 1
        current = @fibers[1]
        # p "old: 1, new: 2, from: #{@fibers[1].rsp}, to: #{@fibers[2].rsp}"
        @current = 2
        Scheduler.swapcontext(pointerof(current.@rsp), @fibers[2].@rsp)
      elsif @current == 2
        # p "old: 2, new: 1, from: #{@fibers[2].rsp}, to: #{@fibers[1].rsp}"
        @current = 1
        current = @fibers[2]
        Scheduler.swapcontext(pointerof(current.@rsp), @fibers[1].@rsp)
      end
    end
  end

  def ret
    puts "ret"
    @finished = @current
    @running -= 1
    if @running == 1
      puts "Done. Exiting."
      Process.exit(0)
    end
    self.yield_control
  end
end

class Executor
  @scheduler = Scheduler.new

  def main
    hello1 = ->{
      (0..5).each do |i|
        puts("Fiber1: #{i}")
        sleep SLEEP_INTERVAL
        @scheduler.yield_control
      end
      @scheduler.ret
    }
    hello2 = ->{
      (0..10).each do |i|
        puts("Fiber2: #{i}")
        sleep SLEEP_INTERVAL
        @scheduler.yield_control
      end
      @scheduler.ret
    }

    @scheduler.spawn(1, hello1)
    @scheduler.spawn(2, hello2)
    @scheduler.run
  end
end

executor = Executor.new
executor.main

first_switch_working.cr

STACK_SIZE = 1024 * 1024 * 8

class FiberPoc
  @stack : Array(UInt8)
  property rsp : UInt64 = 0

  def initialize
    @stack = Array.new(STACK_SIZE, 0_u8)
  end

  def stack
    @stack
  end

  def makecontext(f : Proc)
    p f
    s_ptr = @stack.to_unsafe
    s_start = s_ptr + (STACK_SIZE - 32)
    p s_start


     @rsp = (s_start - 8*8).address # we pop 8 registers before we return
    
    # Our entry
    s_start_ptr = s_start.as(UInt64*)
    s_start_ptr.value = ->start(Void*, Void*).pointer.as(UInt64*).address

    # First parameter
    first_param = (s_start - 8).as(UInt64*)
    first_param.value = f.pointer.address
    

    # Second parameter
    second_param = (s_start - 16).as(UInt64*)
    second_param.value = f.closure_data.address

    p f.pointer
    p f.closure_data
    
    # # ===== UNNCOMMENT FOR DEBUGGING =====
    # p_ptr = ->start(Void*, Void*).pointer.as(UInt64*).address
    # p_ptr = Pointer(Void).new(p_ptr)
    # proc = Proc(Void*, Void*, Void).new(p_ptr, Pointer(Void).null)
    # proc.call(f.pointer, f.closure_data)
    # p @stack
  end
end

fun start(p_ptr : Void*, p_closure : Void*)
p p_ptr
p p_closure
   fn = Proc(Void).new(p_ptr, p_closure)
   fn.call
  puts "hello from start"
end

class Scheduler
  @fibers = Array(FiberPoc).new(3)
  @current = 0
  @running = 1

  def initialize
    @fibers = Array(FiberPoc).new
    (0..2).each do |i|
      @fibers << FiberPoc.new
    end
    @fibers[0].makecontext(->{});
  end

  def spawn(fiber_no : UInt64, f : Proc)
    p "spawn"
    @running += 1
    p @running
    @fibers[fiber_no].makecontext(f)
  end

  def run
    self.yield_control
  end

  def yield_control
    p "yield running: #{@running}, current: #{@current}"
    if @running == 2
        # do nothing
    elsif @running == 3 
    if @current == 0
        p "old: 0, new: 1, from: #{@fibers[0].rsp}, to: #{@fibers[1].rsp}"
      @current = 1
      Scheduler.swapcontext(@fibers[0].rsp, @fibers[1].rsp)
    elsif @current == 1
      p "old: 1, new: 2, from: #{@fibers[1].rsp}, to: #{@fibers[2].rsp}"
      @current = 2
      Scheduler.swapcontext(@fibers[1].rsp, @fibers[2].rsp)
    elsif @current == 2
      p "old: 2, new: 1, from: #{@fibers[2].rsp}, to: #{@fibers[1].rsp}"
      @current = 1
      Scheduler.swapcontext(@fibers[2].rsp, @fibers[1].rsp)
    end
    end
    
  end

  def ret
    p "ret"
    @running -= 1
    if @running == 1
      p "Done. Exiting."
      Process.exit(0)
    end
    self.yield_control
  end

  @[NoInline]
  @[Naked]
  def self.swapcontext(current : UInt64, to : UInt64) : Nil
    # # ===== UNNCOMMENT FOR DEBUGGING =====
    # p_ptr = Pointer(Void*).new(to)
    # p p_ptr
    # p_ptr = p_ptr.value
    # proc = Proc(Void*, Void*, Void).new(p_ptr, Pointer(Void).null)
    # proc.call(Pointer(Void).null, Pointer(Void).null)
    # asm("
    #     movq $0, %rsp
    #    "
    # :: "r"(to))

    asm("
        # Needed to run the closure
        pushq %rdi
        pushq %rsi
        
        # Registers
        pushq %rbx
        pushq %rbp
        pushq %r12
        pushq %r13
        pushq %r14
        pushq %r15
        
        movq %rsp, $0
        movq $1, %rsp
        
        # Registers
        popq %r15
        popq %r14
        popq %r13
        popq %r12
        popq %rbp
        popq %rbx
        
        # Needed to run the closure
        popq %rsi
        popq %rdi
        "
              :
              : "r"(current), "r"(to)
              :
              : "volatile"
    )
  end
end

class Executor
  @scheduler = Scheduler.new

  def main
    hello1 = ->{
      (0..10).each do |i|
        puts("Fiber1: #{i}")
        @scheduler.yield_control
      end
      @scheduler.ret
    }
    hello2 = ->{
      (0..10).each do |i|
        puts("Fiber2: #{i}")
        @scheduler.yield_control
      end
      @scheduler.ret
    }

    @scheduler.spawn(1, hello1)
    @scheduler.spawn(2, hello2)
    @scheduler.run
  end
end

executor = Executor.new
executor.main