Fiber switching for WebAssembly

.github/workflows/wasm32.yml

@@ -18,6 +18,9 @@ jobs:
         with:
           wasmtime-version: "2.0.0"
 
+      - name: Install Binaryen
+        run: apt-get update && apt-get install -y binaryen
+
       - name: Install LLVM 13
         run: |
           apt-get update
@@ -35,6 +38,9 @@ jobs:
           tar -f wasm32-wasi-libs.tar.gz -C wasm32-wasi-libs -xz
           rm wasm32-wasi-libs.tar.gz
 
+      - name: Build fresh compiler
+        run: make crystal # TODO: Remove this after next version update
+
       - name: Build spec/wasm32_std_spec.cr
         run: bin/crystal build spec/wasm32_std_spec.cr -o wasm32_std_spec.wasm --target wasm32-wasi -Duse_pcre
         env:

src/compiler/crystal/compiler.cr

@@ -254,7 +254,8 @@ module Crystal
 
     private def codegen(program, node : ASTNode, sources, output_filename)
       llvm_modules = @progress_tracker.stage("Codegen (crystal)") do
-        program.codegen node, debug: debug, single_module: @single_module || @release || @cross_compile || !@emit_targets.none?
+        # TODO: wasm32 shouldn't require single_module. But not using it somehow messes up with imports from named modules.
+        program.codegen node, debug: debug, single_module: @single_module || @release || @cross_compile || !@emit_targets.none? || program.has_flag? "wasm32"
       end
 
       output_dir = CacheDir.instance.directory_for(sources)
@@ -315,15 +316,16 @@ module Crystal
         target_machine.emit_obj_to_file llvm_mod, object_name
       end
 
-      _, command, args = linker_command(program, [object_name], output_filename, nil)
-      print_command(command, args)
+      linker_commands(program, [object_name], output_filename, nil).each do |(_, command, args)|
+        print_command(command, args)
+      end
     end
 
     private def print_command(command, args)
       stdout.puts command.sub(%("${@}"), args && Process.quote(args))
     end
 
-    private def linker_command(program : Program, object_names, output_filename, output_dir, expand = false)
+    private def linker_commands(program : Program, object_names, output_filename, output_dir, expand = false)
       if program.has_flag? "msvc"
         lib_flags = program.lib_flags
         # Execute and expand `subcommands`.
@@ -379,15 +381,25 @@ module Crystal
           cmd = "#{cl} #{Process.quote_windows("@" + args_filename)}"
         end
 
-        {cl, cmd, nil}
+        [{cl, cmd, nil}]
       elsif program.has_flag? "wasm32"
         link_flags = @link_flags || ""
-        {"wasm-ld", %(wasm-ld "${@}" -o #{Process.quote_posix(output_filename)} #{link_flags} -lc #{program.lib_flags}), object_names}
+        link_flags += " --compress-relocations" if @release
+        link_flags += " --strip-all" if @debug.none?
+
+        opt_flags = "--asyncify"
+        opt_flags += " -Os --all-features" if @release
+        opt_flags += " -g" unless @debug.none?
+        output = Process.quote_posix(output_filename)
+        [
+          {"wasm-ld", %(wasm-ld "${@}" -o #{output} #{link_flags} -lc #{program.lib_flags}), object_names},
+          {"wasm-opt", %(wasm-opt #{output} -o #{output} #{opt_flags}), nil},
+        ]
       else
         link_flags = @link_flags || ""
         link_flags += " -rdynamic"
 
-        {CC, %(#{CC} "${@}" -o #{Process.quote_posix(output_filename)} #{link_flags} #{program.lib_flags}), object_names}
+        [{CC, %(#{CC} "${@}" -o #{Process.quote_posix(output_filename)} #{link_flags} #{program.lib_flags}), object_names}]
       end
     end
 
@@ -419,7 +431,9 @@ module Crystal
 
       @progress_tracker.stage("Codegen (linking)") do
         Dir.cd(output_dir) do
-          run_linker *linker_command(program, object_names, output_filename, output_dir, expand: true)
+          linker_commands(program, object_names, output_filename, output_dir, expand: true).each do |command|
+            run_linker *command
+          end
         end
       end
 

src/crystal/scheduler.cr

@@ -175,8 +175,11 @@ class Crystal::Scheduler
   end
 
   protected def yield : Nil
-    # TODO: Fiber switching and libevent for wasm32
-    {% unless flag?(:wasm32) %}
+    {% if flag?(:wasm32) %}
+      # TODO: event loop for wasm32
+      enqueue @current
+      reschedule
+    {% else %}
       sleep(0.seconds)
     {% end %}
   end

src/crystal/system/wasi/main.cr

@@ -29,5 +29,7 @@ end
 
 # `__main_argc_argv` is called by wasi-libc's `__main_void` with the program arguments.
 fun __main_argc_argv(argc : Int32, argv : UInt8**) : Int32
-  main(argc, argv)
+  Fiber.wrap_with_fibers do
+    main(argc, argv)
+  end
 end

src/fiber.cr

@@ -141,7 +141,7 @@ class Fiber
   end
 
   # :nodoc:
-  def run
+  def run : Nil
     GC.unlock_read
     @proc.call
   rescue ex

src/fiber/context/wasm32.cr

@@ -1,15 +1,189 @@
 {% skip_file unless flag?(:wasm32) %}
 
+# WebAssembly as built by LLVM has 3 stacks:
+# - A value stack for low level operations. Wasm is a stack machine. For example: to write data to
+#   memory you first push the target address, then the value, then perform the 'store' instruction.
+# - A stack with private frames. Each function has access to a stack frame accessed with the local.get
+#   and local.set instructions. These frames are private and can't be manipulated outside of the function.
+# - A shadow stack in the main memory. It is controlled by a stack_pointer global and grows down. This
+#   is not a feature of WebAssembly, just a convention that LLVM uses.
+#
+# There is a proposal to add stack switching to WebAssembly, but as of this writing it is not standardized
+# and isn't implemented by any runtime. See details at https://github.com/WebAssembly/stack-switching.
+#
+# Here we use an alternative implementation on top of Binaryen's Asyncify pass. It is a static transformation
+# on the WebAssembly file that must be performed after the Crystal program is compiled. It rewrites every
+# relevant function so that it writes its stack frame to memory and rewinds to the caller, as well as
+# recovering the stack from memory. It is all controlled by a global state variable.
+# Please read more details at https://github.com/WebAssembly/binaryen/blob/main/src/passes/Asyncify.cpp.
+#
+# Stack switching is accomplished by the cooperation between two functions: Fiber.swapcontext and
+# Fiber.wrap_with_fibers. Before they are explained, the memory layout must be understood:
+#
+# +--------+--------+------------------+----------------
+# | Unused | Data   | Main stack       | Heap ...
+# +--------+--------+------------------+----------------
+# |        |        |                  |
+# 0    __data_start |              __heap_base
+#              __data_end        __stack_pointer
+#
+# The first 1024 bytes of memory are unused (but they are usable, including the address 0!). Follows
+# the static data like constant strings. This would be read only data on other platforms, but here it
+# is writable. The special symbols `__data_start` and `__data_end` are generated by LLVM to mark this
+# region. Then follows the main stack. A `__heap_base` symbol points to the bottom of this stack. The
+# global `__stack_pointer` starts at this position and is moved by the program during execution.
+# After that, the rest of the memory follows, used mainly by malloc. It is important to note that
+# during execution everything is just memory, there is no real difference between those sections.
+#
+# For our implementation of stack switching we manage this memory layout on every stack:
+#
+# +--------+---------------------------+------------------------+
+# | Header | Asyncify buffer =>        |               <= Stack |
+# +--------+---------------------------+------------------------+
+#
+# The header stores 4 pointers as metadata:
+#  - [0] => A function pointer to fiber_main
+#  - [1] => A pointer to the Fiber instance
+#  - [2] => The current position on the asyncify buffer (starts after the header, grows up)
+#  - [4] => The end position on the asyncify buffer (the current stack top)
+# On the main stack all of them are guaranteed to start null.
+#
+# Stack switching should happens as follows:
+#
+# 1. CrystalMain should call Fiber.wrap_with_fibers, passing a proc to the main function.
+# 2. Fiber.wrap_with_fibers will call this proc immediately.
+# 3. At some point a new Fiber will be created and Fiber.swapcontext will be called. It will:
+#    a. store the new context at Fiber.@@next_context.
+#    b. store the current stack_pointer global on context.stack_top.
+#    c. update the current position of the asyncify buffer to be just after the header
+#    d. update the end position  of the asyncify buffer to be the current stack top
+#    e. mark Fiber.manipulating_stack_with_asyncify = true
+#    f. begin stack unwinding with LibAsyncify.start_unwind() and returns.
+# 4. As a consequence of the Asyncify transformation, all functions behave differently and instead
+#    of executing, they will write their local stack to the Asyncify buffer and return.
+# 5. At some point execution will arrive at Fiber.wrap_with_fibers again. We know that we are
+#    unwinding as Fiber.manipulating_stack_with_asyncify? is marked. This means we have to either
+#    start a new fiber or rewind into a previously running fiber. If there is a asyncify buffer, then
+#    setup the rewinding process. Then call into the fiber main function. If it's null, then this is
+#    the main fiber, just call the original block.
+
+lib LibC
+  $__data_end : UInt8
+end
+
+@[Link(wasm_import_module: "asyncify")]
+lib LibAsyncify
+  struct Data
+    current_location : Void*
+    end_location : Void*
+  end
+
+  fun start_unwind(data : Data*)
+  fun stop_unwind
+  fun start_rewind(data : Data*)
+  fun stop_rewind
+end
+
+private def get_stack_pointer
+  stack_pointer = uninitialized Void*
+  asm("
+    .globaltype __stack_pointer, i32
+    global.get __stack_pointer
+    local.set $0
+  " : "=r"(stack_pointer))
+
+  stack_pointer
+end
+
+private def set_stack_pointer(stack_pointer)
+  asm("
+    .globaltype __stack_pointer, i32
+    local.get $0
+    global.set __stack_pointer
+  " :: "r"(stack_pointer))
+end
+
+private def get_main_stack_low
+  pointerof(LibC.__data_end).as(Void*)
+end
+
 class Fiber
   # :nodoc:
-  def makecontext(stack_ptr, fiber_main)
-    # TODO: Implement this using Binaryen Asyncify
+  class_property next_context : Context*?
+
+  # :nodoc:
+  class_property? manipulating_stack_with_asyncify = false
+
+  struct Context
+    property stack_low : Void* = get_main_stack_low
+  end
+
+  # :nodoc:
+  def makecontext(stack_ptr : Void**, fiber_main : Fiber ->)
+    @context.stack_top = stack_ptr.as(Void*)
+    @context.stack_low = (stack_ptr.as(UInt8*) - StackPool::STACK_SIZE + 32).as(Void*)
+    @context.resumable = 1
+
+    ctx_data_ptr = @context.stack_low.as(Void**)
+    ctx_data_ptr[0] = fiber_main.pointer
+    ctx_data_ptr[1] = self.as(Void*)
+    ctx_data_ptr[2] = Pointer(Void).null
+    ctx_data_ptr[3] = Pointer(Void).null
   end
 
   # :nodoc:
   @[NoInline]
-  @[Naked]
   def self.swapcontext(current_context, new_context) : Nil
-    # TODO: Implement this using Binaryen Asyncify
+    if Fiber.manipulating_stack_with_asyncify?
+      Fiber.manipulating_stack_with_asyncify = false
+      LibAsyncify.stop_rewind
+      return
+    end
+
+    new_context.value.resumable = 0
+    current_context.value.resumable = 1
+    Fiber.next_context = new_context
+
+    current_context.value.stack_top = get_stack_pointer
+
+    ctx_data_ptr = current_context.value.stack_low.as(Void**)
+    ctx_data_ptr[2] = (ctx_data_ptr + 4).as(Void*)
+    ctx_data_ptr[3] = current_context.value.stack_top
+
+    asyncify_data_ptr = (ctx_data_ptr + 2).as(LibAsyncify::Data*)
+    Fiber.manipulating_stack_with_asyncify = true
+    LibAsyncify.start_unwind(asyncify_data_ptr)
+  end
+
+  # :nodoc:
+  @[NoInline]
+  def self.wrap_with_fibers(&block : -> T) : T forall T
+    result = block.call
+
+    while Fiber.manipulating_stack_with_asyncify?
+      Fiber.manipulating_stack_with_asyncify = false
+      LibAsyncify.stop_unwind
+
+      next_context = Fiber.next_context.not_nil!
+      ctx_data_ptr = next_context.value.stack_low.as(Void**)
+
+      set_stack_pointer next_context.value.stack_top
+
+      asyncify_data_ptr = (ctx_data_ptr + 2).as(LibAsyncify::Data*)
+      unless asyncify_data_ptr.value.current_location == Pointer(Void).null
+        Fiber.manipulating_stack_with_asyncify = true
+        LibAsyncify.start_rewind(asyncify_data_ptr)
+      end
+
+      if ctx_data_ptr[0].null?
+        result = block.call
+      else
+        fiber_main = Proc(Fiber, Void).new(ctx_data_ptr[0], Pointer(Void).null)
+        fiber = ctx_data_ptr[1].as(Fiber)
+        fiber_main.call(fiber)
+      end
+    end
+
+    result
   end
 end