Add builtins.wasi

doc/manual/source/protocols/wasm.md

@@ -1,25 +1,48 @@
-# `builtins.wasm` Host Interface
+# Wasm Host Interface
 
-The `builtins.wasm` function allows Nix expressions to call WebAssembly functions. This enables extending Nix with custom functionality written in languages that compile to WebAssembly (such as Rust).
+Nix provides two builtins for calling WebAssembly modules: `builtins.wasm` and `builtins.wasi`. These allow extending Nix with custom functionality written in languages that compile to WebAssembly (such as Rust).
 
 ## Overview
 
 WebAssembly modules can interact with Nix values through a host interface that provides functions for creating and inspecting Nix values. The WASM module receives Nix values as opaque `ValueId` handles and uses host functions to work with them.
 
+There are two calling conventions:
+
+- **`builtins.wasm`** calls a named Wasm export directly. The function receives its input as a `ValueId` parameter and returns a `ValueId`.
+- **`builtins.wasi`** runs a WASI module's `_start` entry point. The input `ValueId` is passed as a command-line argument (`argv[1]`), and the result is returned by calling the `return_to_nix` host function.
+
 ## Value IDs
 
 Nix values are represented in Wasm code as a `u32` referred to below as a `ValueId`. These are opaque handles that reference values managed by the Nix evaluator. Value ID 0 is reserved to represent a missing attribute lookup result.
 
-## Entry Point
+## Entry Points
+
+### `builtins.wasm` Entry Point
 
-Every Wasm module must export:
+Usage: `builtins.wasm <module> <function-name> <arg>`
+
+Every Wasm module used with `builtins.wasm` must export:
 - A `memory` object that the host can use to read/write data.
 - `nix_wasm_init_v1()`, a function that is called once when the module is instantiated.
 - The entry point of the Wasm function, its name corresponding to the second argument to `builtins.wasm`. It takes a single `ValueId` and returns a single `ValueId` (i.e. it has type `fn(arg: u32) -> u32`).
 
+### `builtins.wasi` Entry Point
+
+Usage: `builtins.wasi <module> <arg>`
+
+Every WASI module used with `builtins.wasi` must export:
+- A `memory` object that the host can use to read/write data.
+- `_start()`, the standard WASI entry point. This function takes no parameters.
+
+The input value is passed as a command-line argument: `argv[1]` is set to the decimal representation of the `ValueId` of the input value.
+
+To return a result to Nix, the module must call the `return_to_nix` host function (see below) with the `ValueId` of the result. If `_start` finishes without calling `return_to_nix`, an error is raised.
+
+Standard output and standard error from the WASI module are captured and emitted as Nix warnings (one warning per line).
+
 ## Host Functions
 
-All host functions are imported from the `env` module provided by `builtins.wasm`.
+All host functions are imported from the `env` module provided by `builtins.wasm` and `builtins.wasi`.
 
 ### Error Handling
 
@@ -304,6 +327,17 @@ Creates a lazy or partially applied function application.
 
 **Returns:** Value ID of the unevaluated application
 
+### Returning Results (`builtins.wasi` only)
+
+#### `return_to_nix(value: ValueId)`
+
+Returns a result value to the Nix evaluator from a WASI module. This function is only available in modules called via `builtins.wasi`.
+
+**Parameters:**
+- `value` - ID of the Nix value to return as the result of the `builtins.wasi` call
+
+**Note:** Calling this function immediately terminates the WASI module's execution. The module's `_start` function must call `return_to_nix` before finishing; otherwise, an error is raised.
+
 ### File I/O
 
 #### `read_file(path: ValueId, ptr: u32, len: u32) -> u32`

src/libexpr/primops/wasm.cc

@@ -78,20 +78,22 @@ TrapResult<InstancePre> instantiate_pre(Linker & linker, const Module & m)
     return InstancePre(instance_pre);
 }
 
-void regFuns(Linker & linker);
+static void regFuns(Linker & linker, bool useWasi);
 
 struct NixWasmInstancePre
 {
     Engine & engine;
     SourcePath wasmPath;
     InstancePre instancePre;
 
-    NixWasmInstancePre(SourcePath _wasmPath)
+    NixWasmInstancePre(SourcePath _wasmPath, bool useWasi = false)
         : engine(getEngine())
         , wasmPath(_wasmPath)
         , instancePre(({
             Linker linker(engine);
-            regFuns(linker);
+            if (useWasi)
+                unwrap(linker.define_wasi());
+            regFuns(linker, useWasi);
             unwrap(instantiate_pre(linker, unwrap(Module::compile(engine, string2span(wasmPath.readFile())))));
         }))
     {
@@ -112,13 +114,18 @@ struct NixWasmInstance
 
     std::optional<std::string> functionName;
 
+    ValueId resultId = 0;
+
+    std::string logPrefix;
+
     NixWasmInstance(EvalState & _state, ref<NixWasmInstancePre> _pre)
         : state(_state)
         , pre(_pre)
         , wasmStore(pre->engine)
         , wasmCtx(wasmStore)
         , instance(unwrap(pre->instancePre.instantiate(wasmCtx)))
         , memory_(getExport<Memory>("memory"))
+        , logPrefix(pre->wasmPath.baseName())
     {
         wasmCtx.set_data(this);
 
@@ -177,14 +184,18 @@ struct NixWasmInstance
 
     std::monostate warn(uint32_t ptr, uint32_t len)
     {
-        nix::warn(
-            "'%s' function '%s': %s",
-            pre->wasmPath,
-            functionName.value_or("<unknown>"),
-            span2string(memory().subspan(ptr, len)));
+        doWarn(span2string(memory().subspan(ptr, len)));
         return {};
     }
 
+    void doWarn(std::string_view s)
+    {
+        if (functionName)
+            nix::warn("'%s' function '%s': %s", logPrefix, functionName.value_or("<unknown>"), s);
+        else
+            nix::warn("'%s': %s", logPrefix, s);
+    }
+
     uint32_t get_type(ValueId valueId)
     {
         auto & value = getValue(valueId);
@@ -475,7 +486,7 @@ static void regFun(Linker & linker, std::string_view name, R (NixWasmInstance::*
     }));
 }
 
-void regFuns(Linker & linker)
+static void regFuns(Linker & linker, bool useWasi)
 {
     regFun(linker, "panic", &NixWasmInstance::panic);
     regFun(linker, "warn", &NixWasmInstance::warn);
@@ -500,32 +511,46 @@ void regFuns(Linker & linker)
     regFun(linker, "call_function", &NixWasmInstance::call_function);
     regFun(linker, "make_app", &NixWasmInstance::make_app);
     regFun(linker, "read_file", &NixWasmInstance::read_file);
+
+    if (useWasi) {
+        unwrap(linker.func_wrap(
+            "env", "return_to_nix", [](Caller caller, ValueId resultId) -> Result<std::monostate, Trap> {
+                auto instance = std::any_cast<NixWasmInstance *>(caller.context().get_data());
+                instance->resultId = resultId;
+                return Trap("return_to_nix");
+            }));
+    }
 }
 
-void prim_wasm(EvalState & state, const PosIdx pos, Value ** args, Value & v)
+static NixWasmInstance instantiateWasm(EvalState & state, const SourcePath & wasmPath, bool useWasi)
+{
+    // FIXME: make this a weak Boehm GC pointer so that it can be freed during GC.
+    // FIXME: move to EvalState?
+    // Note: InstancePre in Rust is Send+Sync so it should be safe to share between threads.
+    static boost::concurrent_flat_map<std::pair<SourcePath, bool>, std::shared_ptr<NixWasmInstancePre>> instancesPre;
+
+    std::shared_ptr<NixWasmInstancePre> instancePre;
+
+    instancesPre.try_emplace_and_cvisit(
+        {wasmPath, useWasi},
+        nullptr,
+        [&](auto & i) { instancePre = i.second = std::make_shared<NixWasmInstancePre>(wasmPath, useWasi); },
+        [&](auto & i) { instancePre = i.second; });
+
+    return NixWasmInstance{state, ref(instancePre)};
+}
+
+static void prim_wasm(EvalState & state, const PosIdx pos, Value ** args, Value & v)
 {
     auto wasmPath = state.realisePath(pos, *args[0]);
     std::string functionName =
         std::string(state.forceStringNoCtx(*args[1], pos, "while evaluating the second argument of `builtins.wasm`"));
 
     try {
-        // FIXME: make this a weak Boehm GC pointer so that it can be freed during GC.
-        // FIXME: move to EvalState?
-        // Note: InstancePre in Rust is Send+Sync so it should be safe to share between threads.
-        static boost::concurrent_flat_map<SourcePath, std::shared_ptr<NixWasmInstancePre>> instancesPre;
-
-        std::shared_ptr<NixWasmInstancePre> instancePre;
-
-        instancesPre.try_emplace_and_cvisit(
-            wasmPath,
-            nullptr,
-            [&](auto & i) { instancePre = i.second = std::make_shared<NixWasmInstancePre>(wasmPath); },
-            [&](auto & i) { instancePre = i.second; });
+        auto instance = instantiateWasm(state, wasmPath, false);
 
         debug("calling wasm module");
 
-        NixWasmInstance instance{state, ref(instancePre)};
-
         // FIXME: use the "start" function if present.
         instance.runFunction("nix_wasm_init_v1", {});
 
@@ -552,4 +577,83 @@ static RegisterPrimOp primop_wasm(
      .fun = prim_wasm,
      .experimentalFeature = Xp::WasmBuiltin});
 
+/**
+ * Callback for WASI stdout/stderr writes. It splits the output into lines and logs each line separately.
+ */
+struct WasiLogger
+{
+    NixWasmInstance & instance;
+
+    std::string data;
+
+    ~WasiLogger()
+    {
+        if (!data.empty())
+            instance.doWarn(data);
+    }
+
+    void operator()(std::string_view s)
+    {
+        data.append(s);
+
+        while (true) {
+            auto pos = data.find('\n');
+            if (pos == std::string_view::npos)
+                break;
+            instance.doWarn(data.substr(0, pos));
+            data.erase(0, pos + 1);
+        }
+    }
+};
+
+static void prim_wasi(EvalState & state, const PosIdx pos, Value ** args, Value & v)
+{
+    auto wasmPath = state.realisePath(pos, *args[0]);
+    auto functionName = "_start";
+
+    try {
+        auto instance = instantiateWasm(state, wasmPath, true);
+
+        debug("calling wasm module");
+
+        auto argId = instance.addValue(args[1]);
+
+        WasiLogger logger{instance};
+
+        auto loggerTrampoline = [](void * data, const unsigned char * buf, size_t len) -> ptrdiff_t {
+            auto logger = static_cast<WasiLogger *>(data);
+            (*logger)(std::string_view((const char *) buf, len));
+            return len;
+        };
+
+        WasiConfig wasiConfig;
+        wasi_config_set_stdout_custom(wasiConfig.capi(), loggerTrampoline, &logger, nullptr);
+        wasi_config_set_stderr_custom(wasiConfig.capi(), loggerTrampoline, &logger, nullptr);
+        wasiConfig.argv({"wasi", std::to_string(argId)});
+        unwrap(instance.wasmStore.context().set_wasi(std::move(wasiConfig)));
+
+        auto res = instance.getExport<Func>(functionName).call(instance.wasmCtx, {});
+        if (!instance.resultId) {
+            unwrap(std::move(res));
+            throw Error("Wasm function '%s' from '%s' finished without returning a value", functionName, wasmPath);
+        }
+
+        auto & vRes = instance.getValue(instance.resultId);
+        state.forceValue(vRes, pos);
+        v = vRes;
+    } catch (Error & e) {
+        e.addTrace(state.positions[pos], "while executing the Wasm function '%s' from '%s'", functionName, wasmPath);
+        throw;
+    }
+}
+
+static RegisterPrimOp primop_wasi(
+    {.name = "wasi",
+     .args = {"wasi", "arg"},
+     .doc = R"(
+      Call a WASI function with the specified argument.
+     )",
+     .fun = prim_wasi,
+     .experimentalFeature = Xp::WasmBuiltin});
+
 } // namespace nix