[wasm] vm.c: stop unwinding to main for every vm_exec call by setjmp

the original rb_wasm_setjmp implementation always unwinds to the root
call frame to have setjmp compatible interface, and simulate sjlj's
undefined behavior. Therefore, every vm_exec call unwinds to main, and
a deep call stack makes setjmp call very expensive. The following
snippet from optcarrot takes 5s even though it takes less than 0.3s on
native.

```
[0x0, 0x4, 0x8, 0xc].map do |attr|
  (0..7).map do |j|
    (0...0x10000).map do |i|
      clr = i[15 - j] * 2 + i[7 - j]
      clr != 0 ? attr | clr : 0
    end
  end
end
```

This patch adds a WASI specialized vm_exec which uses lightweight
try-catch API without unwinding to the root frame. After this patch, the
above snippet takes only 0.5s.

4 files changed, 194 insertions, 0 deletions

diff --git a/vm.c b/vm.c
index 4cede5e247..0bc9784046 100644
--- a/vm.c
+++ b/vm.c
@@ -2196,6 +2196,85 @@ static inline VALUE
 vm_exec_handle_exception(rb_execution_context_t *ec, enum ruby_tag_type state,
                          VALUE errinfo, VALUE *initial);
 
+// for non-Emscripten Wasm build, use vm_exec with optimized setjmp for runtime performance
+#if defined(__wasm__) && !defined(__EMSCRIPTEN__)
+
+struct rb_vm_exec_context {
+    rb_execution_context_t *ec;
+    struct rb_vm_tag *tag;
+    VALUE initial;
+    VALUE result;
+    enum ruby_tag_type state;
+    bool mjit_enable_p;
+};
+
+static void
+vm_exec_enter_vm_loop(rb_execution_context_t *ec, struct rb_vm_exec_context *ctx,
+                      struct rb_vm_tag *_tag, bool skip_first_ex_handle) {
+    if (skip_first_ex_handle) {
+        goto vm_loop_start;
+    }
+
+    ctx->result = ec->errinfo;
+    rb_ec_raised_reset(ec, RAISED_STACKOVERFLOW | RAISED_NOMEMORY);
+    while ((ctx->result = vm_exec_handle_exception(ec, ctx->state, ctx->result, &ctx->initial)) == Qundef) {
+        /* caught a jump, exec the handler */
+        ctx->result = vm_exec_core(ec, ctx->initial);
+    vm_loop_start:
+        VM_ASSERT(ec->tag == _tag);
+        /* when caught `throw`, `tag.state` is set. */
+        if ((ctx->state = _tag->state) == TAG_NONE) break;
+        _tag->state = TAG_NONE;
+    }
+}
+
+static void
+vm_exec_bottom_main(void *context)
+{
+    struct rb_vm_exec_context *ctx = (struct rb_vm_exec_context *)context;
+
+    ctx->state = TAG_NONE;
+    if (!ctx->mjit_enable_p || (ctx->result = mjit_exec(ctx->ec)) == Qundef) {
+        ctx->result = vm_exec_core(ctx->ec, ctx->initial);
+    }
+    vm_exec_enter_vm_loop(ctx->ec, ctx, ctx->tag, true);
+}
+
+static void
+vm_exec_bottom_rescue(void *context)
+{
+    struct rb_vm_exec_context *ctx = (struct rb_vm_exec_context *)context;
+    ctx->state = rb_ec_tag_state(ctx->ec);
+    vm_exec_enter_vm_loop(ctx->ec, ctx, ctx->tag, false);
+}
+
+VALUE
+vm_exec(rb_execution_context_t *ec, bool mjit_enable_p)
+{
+    struct rb_vm_exec_context ctx = {
+        .ec = ec,
+        .initial = 0, .result = Qundef,
+        .mjit_enable_p = mjit_enable_p,
+    };
+    struct rb_wasm_try_catch try_catch;
+
+    EC_PUSH_TAG(ec);
+
+    _tag.retval = Qnil;
+    ctx.tag = &_tag;
+
+    EC_REPUSH_TAG();
+
+    rb_wasm_try_catch_init(&try_catch, vm_exec_bottom_main, vm_exec_bottom_rescue, &ctx);
+
+    rb_wasm_try_catch_loop_run(&try_catch, &_tag.buf);
+
+    EC_POP_TAG();
+    return ctx.result;
+}
+
+#else
+
 VALUE
 vm_exec(rb_execution_context_t *ec, bool mjit_enable_p)
 {
@@ -2228,6 +2307,7 @@ vm_exec(rb_execution_context_t *ec, bool mjit_enable_p)
     EC_POP_TAG();
     return result;
 }
+#endif
 
 static inline VALUE
 vm_exec_handle_exception(rb_execution_context_t *ec, enum ruby_tag_type state,
diff --git a/wasm/asyncify.h b/wasm/asyncify.h
index 834fc8b570..49eb125593 100644
--- a/wasm/asyncify.h
+++ b/wasm/asyncify.h
@@ -3,8 +3,18 @@
 
 __attribute__((import_module("asyncify"), import_name("start_unwind")))
 void asyncify_start_unwind(void *buf);
+#define asyncify_start_unwind(buf) do {  \
+    extern void *rb_asyncify_unwind_buf; \
+    rb_asyncify_unwind_buf = (buf);      \
+    asyncify_start_unwind((buf));        \
+  } while (0)
 __attribute__((import_module("asyncify"), import_name("stop_unwind")))
 void asyncify_stop_unwind(void);
+#define asyncify_stop_unwind() do {      \
+    extern void *rb_asyncify_unwind_buf; \
+    rb_asyncify_unwind_buf = NULL;       \
+    asyncify_stop_unwind();              \
+  } while (0)
 __attribute__((import_module("asyncify"), import_name("start_rewind")))
 void asyncify_start_rewind(void *buf);
 __attribute__((import_module("asyncify"), import_name("stop_rewind")))
diff --git a/wasm/setjmp.c b/wasm/setjmp.c
index 3f017e202a..1e69cb1c64 100644
--- a/wasm/setjmp.c
+++ b/wasm/setjmp.c
@@ -57,6 +57,7 @@ async_buf_init(struct __rb_wasm_asyncify_jmp_buf* buf)
 
 // Global unwinding/rewinding jmpbuf state
 static rb_wasm_jmp_buf *_rb_wasm_active_jmpbuf;
+void *rb_asyncify_unwind_buf;
 
 __attribute__((noinline))
 int
@@ -106,6 +107,75 @@ _rb_wasm_longjmp(rb_wasm_jmp_buf* env, int value)
     asyncify_start_unwind(&env->longjmp_buf);
 }
 
+
+enum try_catch_phase {
+  TRY_CATCH_PHASE_MAIN   = 0,
+  TRY_CATCH_PHASE_RESCUE = 1,
+};
+
+void
+rb_wasm_try_catch_init(struct rb_wasm_try_catch *try_catch,
+                       rb_wasm_try_catch_func_t try_f,
+                       rb_wasm_try_catch_func_t catch_f,
+                       void *context)
+{
+    try_catch->state = TRY_CATCH_PHASE_MAIN;
+    try_catch->try_f = try_f;
+    try_catch->catch_f = catch_f;
+    try_catch->context = context;
+}
+
+// NOTE: This function is not processed by Asyncify due to a call of asyncify_stop_rewind
+void
+rb_wasm_try_catch_loop_run(struct rb_wasm_try_catch *try_catch, rb_wasm_jmp_buf *target)
+{
+    extern void *rb_asyncify_unwind_buf;
+    extern rb_wasm_jmp_buf *_rb_wasm_active_jmpbuf;
+
+    target->state = JMP_BUF_STATE_CAPTURED;
+
+    switch ((enum try_catch_phase)try_catch->state) {
+    case TRY_CATCH_PHASE_MAIN: {
+        // may unwind
+        try_catch->try_f(try_catch->context);
+        break;
+    }
+    case TRY_CATCH_PHASE_RESCUE: {
+        if (try_catch->catch_f) {
+            // may unwind
+            try_catch->catch_f(try_catch->context);
+        }
+        break;
+    }
+    }
+
+    while (1) {
+        // catch longjmp with target jmp_buf
+        if (rb_asyncify_unwind_buf && _rb_wasm_active_jmpbuf == target) {
+            // do similar steps setjmp does when JMP_BUF_STATE_RETURNING
+
+            // stop unwinding
+            // (but call stop_rewind to update the asyncify state to "normal" from "unwind")
+            asyncify_stop_rewind();
+            // clear the active jmpbuf because it's already stopped
+            _rb_wasm_active_jmpbuf = NULL;
+            // reset jmpbuf state to be able to unwind again
+            target->state = JMP_BUF_STATE_CAPTURED;
+            // move to catch loop phase
+            try_catch->state = TRY_CATCH_PHASE_RESCUE;
+            if (try_catch->catch_f) {
+                try_catch->catch_f(try_catch->context);
+            }
+            continue;
+        } else if (rb_asyncify_unwind_buf /* unrelated unwind */) {
+            return;
+        }
+        // no unwind, then exit
+        break;
+    }
+    return;
+}
+
 void *
 rb_wasm_handle_jmp_unwind(void)
 {
diff --git a/wasm/setjmp.h b/wasm/setjmp.h
index 30ea23ca12..65e35c03b3 100644
--- a/wasm/setjmp.h
+++ b/wasm/setjmp.h
@@ -58,4 +58,38 @@ typedef rb_wasm_jmp_buf jmp_buf;
 #define setjmp(env) rb_wasm_setjmp(env)
 #define longjmp(env, payload) rb_wasm_longjmp(env, payload)
 
+
+typedef void (*rb_wasm_try_catch_func_t)(void *ctx);
+
+struct rb_wasm_try_catch {
+    rb_wasm_try_catch_func_t try_f;
+    rb_wasm_try_catch_func_t catch_f;
+    void *context;
+    int state;
+};
+
+//
+// Lightweight try-catch API without unwinding to root frame.
+//
+
+void
+rb_wasm_try_catch_init(struct rb_wasm_try_catch *try_catch,
+                       rb_wasm_try_catch_func_t try_f,
+                       rb_wasm_try_catch_func_t catch_f,
+                       void *context);
+
+// Run, catch longjmp thrown by run, and re-catch longjmp thrown by catch, ...
+//
+// 1. run try_f of try_catch struct
+// 2. catch longjmps with the given target jmp_buf or exit
+// 3. run catch_f if not NULL, otherwise exit
+// 4. catch longjmps with the given target jmp_buf or exit
+// 5. repeat from step 3
+//
+// NOTICE: This API assumes that all longjmp targeting the given jmp_buf are NOT called
+// after the function that called this function has exited.
+//
+void
+rb_wasm_try_catch_loop_run(struct rb_wasm_try_catch *try_catch, rb_wasm_jmp_buf *target);
+
 #endif