C++とアセンブラによるコルーチンの試験的実装

mycoroutine.cpp

#include <array>
#include <chrono>
#include <cstring>
#include <iostream>
#include <thread>

using namespace std;

struct CoroContext {
  uint64_t rip, rsp, rbp, rbx, r12, r13, r14, r15;
};

struct alignas(16) CoroEnv {
  CoroContext run_ctx;
  CoroContext task_ctx;
  std::array<uint64_t, 512 - 16> task_stack;
};

struct CoroStatus {
  int64_t task_exitcode;
  enum Status {
    kNotInitialized = 0,
    kRunnable = 1,
    kFinished = 2,
  } status;
};

extern "C" {
using TaskFunc = int64_t (CoroEnv*, int64_t);
void ExitTask();
CoroStatus RunCoro(CoroEnv*);
void Yield(CoroEnv*);
}

void InitCoro(CoroEnv* env, TaskFunc* f, int64_t arg) {
  memset(env, 0, sizeof(CoroEnv));
  uint64_t* stk = env->task_stack.end() - 3;
  stk[0] = reinterpret_cast<uint64_t>(&ExitTask);
  stk[1] = reinterpret_cast<uint64_t>(env);
  env->task_ctx.rsp = reinterpret_cast<uint64_t>(stk);
  env->task_ctx.rip = reinterpret_cast<uint64_t>(f);
  env->task_ctx.r12 = arg;
}

int64_t LongTask(CoroEnv* env, int64_t arg) {
  cout << "task " << arg << ": task phase 1 ..." << flush;
  this_thread::sleep_for(chrono::milliseconds(500));
  cout << "done" << endl;

  Yield(env);

  cout << "task " << arg << ": task phase 2 ..." << flush;
  this_thread::sleep_for(chrono::milliseconds(500));
  cout << "done" << endl;

  Yield(env);

  cout << "task " << arg << ": task phase 3 ..." << flush;
  this_thread::sleep_for(chrono::milliseconds(500));
  cout << "done" << endl;

  return 42 + arg;
}

int main() {
  alignas(16) CoroEnv env;
  alignas(16) CoroEnv env2;
  InitCoro(&env, LongTask, 1);
  InitCoro(&env2, LongTask, 2);

  while (true) {
    auto r = RunCoro(&env);
    if (r.status == CoroStatus::kFinished) {
      cout << "LongTask finished with code=" << r.task_exitcode << endl;
    } else if (r.status == CoroStatus::kNotInitialized) {
      cout << "env is not initialized" << endl;
    }
    r = RunCoro(&env2);
    if (r.status == CoroStatus::kFinished) {
      cout << "LongTask finished with code=" << r.task_exitcode << endl;
    } else if (r.status == CoroStatus::kNotInitialized) {
      cout << "env2 is not initialized" << endl;
      return 0;
    }
    cout << "LongTask continues..." << endl;
  }
}

mycoroutine_helper.s

.intel_syntax noprefix
.code64

.section .text

.global ExitTask
ExitTask:
    mov rdi, [rsp]
    mov edx, 2                     # status = kFinished
    mov qword ptr [rdi + 0x40], 0  # env->task_ctx.rip = NULL
    jmp SuspendTask

.global RunCoro
RunCoro: # CoroStatus RunCoro(CoroEnv* env);
    mov rdx, [rdi + 0x40]  # env->task_ctx.rip
    test rdx, rdx
    jz RunCoroError

    # run_ctx への保存
    mov rax, [rsp]
    mov [rdi + 0x00], rax  # env->run_ctx.rip = RunCoro's ret addr
    lea rax, [rsp + 8]
    mov [rdi + 0x08], rax  # env->run_ctx.rsp = rsp when RunCoro returned

    mov [rdi + 0x10], rbp  # env->run_ctx.rbp = rbp
    mov [rdi + 0x18], rbx
    mov [rdi + 0x20], r12
    mov [rdi + 0x28], r13
    mov [rdi + 0x30], r14
    mov [rdi + 0x38], r15

    # task_ctx からの復帰
    mov rsp, [rdi + 0x48]  # rsp = env->task_ctx.rsp
    mov rbp, [rdi + 0x50]
    mov rbx, [rdi + 0x58]
    mov r12, [rdi + 0x60]
    mov r13, [rdi + 0x68]
    mov r14, [rdi + 0x70]
    mov r15, [rdi + 0x78]

    mov rsi, r12
    jmp rdx  # タスクへジャンプ

RunCoroError:
    # here rdx == 0 => status = kNotInitialized
    ret

.global Yield
Yield: # void Yield(CoroEnv* env);
    # task_ctx への保存
    mov rax, [rsp]
    mov [rdi + 0x40], rax  # env->task_ctx.rip = Yeild's ret addr
    lea rax, [rsp + 8]
    mov [rdi + 0x48], rax  # env->task_ctx.rsp = rsp when Yeild returned

    mov [rdi + 0x50], rbp  # env->task_ctx.rbp = rbp
    mov [rdi + 0x58], rbx
    mov [rdi + 0x60], r12
    mov [rdi + 0x68], r13
    mov [rdi + 0x70], r14
    mov [rdi + 0x78], r15
    mov edx, 1             # status = kRunnable

SuspendTask:
    # run_ctx からの復帰
    mov rsp, [rdi + 0x08]  # rsp = env->run_ctx.rsp
    mov rbp, [rdi + 0x10]
    mov rbx, [rdi + 0x18]
    mov r12, [rdi + 0x20]
    mov r13, [rdi + 0x28]
    mov r14, [rdi + 0x30]
    mov r15, [rdi + 0x38]

    mov rcx, [rdi + 0x00]  # env->run_ctx.rip
    jmp rcx  # RunCoro の次の行までジャンプ

output.log

$ g++ mycoroutine.cpp mycoroutine_helper.s
$ ./a.out
task 1: task phase 1 ...done
task 2: task phase 1 ...done
LongTask continues...
task 1: task phase 2 ...done
task 2: task phase 2 ...done
LongTask continues...
task 1: task phase 3 ...done
LongTask finished with code=43
task 2: task phase 3 ...done
LongTask finished with code=44
LongTask continues...
env is not initialized
env2 is not initialized
$