Merge builtins.wasm and builtins.wasi into one function

doc/manual/source/protocols/wasm.md

@@ -1,36 +1,57 @@
 # Wasm Host Interface
 
-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).
+Nix provides a builtin for calling WebAssembly modules: `builtins.wasm`. This allows 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:
+The `builtins.wasm` builtin takes two arguments:
+1. A configuration attribute set with the following attributes:
+   - `path` - Path to the WebAssembly module (required)
+   - `function` - Name of the Wasm function to call (required for non-WASI modules, not allowed for WASI modules)
+2. The argument value to pass to the function
 
-- **`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.
+WASI mode is automatically detected by checking if the module imports from `wasi_snapshot_preview1`. There are two calling conventions:
+
+- **Non-WASI mode** (no WASI imports) calls the Wasm export specified by `function` directly. The function receives its input as a `ValueId` parameter and returns a `ValueId`.
+- **WASI mode** (when the module imports from `wasi_snapshot_preview1`) runs the 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 Points
 
-### `builtins.wasm` Entry Point
+### Non-WASI Mode
+
+Non-WASI mode is used when the module does **not** import from `wasi_snapshot_preview1`.
 
-Usage: `builtins.wasm <module> <function-name> <arg>`
+Usage:
+```nix
+builtins.wasm {
+  path = <module>;
+  function = <function-name>;
+} <arg>
+```
 
-Every Wasm module used with `builtins.wasm` must export:
+Every Wasm module used in non-WASI mode 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`).
+- The entry point function, whose name is specified by the `function` attribute. It takes a single `ValueId` and returns a single `ValueId` (i.e. it has type `fn(arg: u32) -> u32`).
 
-### `builtins.wasi` Entry Point
+### WASI Mode
 
-Usage: `builtins.wasi <module> <arg>`
+WASI mode is automatically used when the module imports a `wasi_snapshot_preview1` function.
+
+Usage:
+```nix
+builtins.wasm {
+  path = <module>;
+} <arg>
+```
 
-Every WASI module used with `builtins.wasi` must export:
+Every WASI module 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.
 
@@ -42,7 +63,7 @@ Standard output and standard error from the WASI module are captured and emitted
 
 ## Host Functions
 
-All host functions are imported from the `env` module provided by `builtins.wasm` and `builtins.wasi`.
+All host functions are imported from the `env` module.
 
 ### Error Handling
 
@@ -327,16 +348,16 @@ Creates a lazy or partially applied function application.
 
 **Returns:** Value ID of the unevaluated application
 
-### Returning Results (`builtins.wasi` only)
+### Returning Results (WASI mode 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`.
+Returns a result value to the Nix evaluator from a WASI module. This function is only available in WASI mode.
 
 **Parameters:**
-- `value` - ID of the Nix value to return as the result of the `builtins.wasi` call
+- `value` - ID of the Nix value to return as the result of the `builtins.wasm` 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.
+**Note:** Calling this function immediately terminates the WASI module's execution. The module must call `return_to_nix` before finishing; otherwise, an error is raised.
 
 ### File I/O
 

src/libexpr/primops/wasm.cc

@@ -84,17 +84,31 @@ struct NixWasmInstancePre
 {
     Engine & engine;
     SourcePath wasmPath;
+    bool useWasi;
     InstancePre instancePre;
 
-    NixWasmInstancePre(SourcePath _wasmPath, bool useWasi = false)
+    NixWasmInstancePre(SourcePath _wasmPath)
         : engine(getEngine())
         , wasmPath(_wasmPath)
+        , useWasi(false)
         , instancePre(({
+            // Compile the module
+            auto module = unwrap(Module::compile(engine, string2span(wasmPath.readFile())));
+
+            // Auto-detect WASI by checking for wasi_snapshot_preview1 imports.
+            for (const auto & ref : module.imports())
+                if (const_cast<std::decay_t<decltype(ref)> &>(ref).module() == "wasi_snapshot_preview1") {
+                    useWasi = true;
+                    break;
+                }
+
+            // Create linker with appropriate WASI support
             Linker linker(engine);
             if (useWasi)
                 unwrap(linker.define_wasi());
             regFuns(linker, useWasi);
-            unwrap(instantiate_pre(linker, unwrap(Module::compile(engine, string2span(wasmPath.readFile())))));
+
+            unwrap(instantiate_pre(linker, module));
         }))
     {
     }
@@ -522,61 +536,24 @@ static void regFuns(Linker & linker, bool useWasi)
     }
 }
 
-static NixWasmInstance instantiateWasm(EvalState & state, const SourcePath & wasmPath, bool useWasi)
+static NixWasmInstance instantiateWasm(EvalState & state, const SourcePath & wasmPath)
 {
     // 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;
+    static boost::concurrent_flat_map<SourcePath, std::shared_ptr<NixWasmInstancePre>> instancesPre;
 
     std::shared_ptr<NixWasmInstancePre> instancePre;
 
     instancesPre.try_emplace_and_cvisit(
-        {wasmPath, useWasi},
+        wasmPath,
         nullptr,
-        [&](auto & i) { instancePre = i.second = std::make_shared<NixWasmInstancePre>(wasmPath, useWasi); },
+        [&](auto & i) { instancePre = i.second = std::make_shared<NixWasmInstancePre>(wasmPath); },
         [&](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 {
-        auto instance = instantiateWasm(state, wasmPath, false);
-
-        debug("calling wasm module");
-
-        // FIXME: use the "start" function if present.
-        instance.runFunction("nix_wasm_init_v1", {});
-
-        auto res = instance.runFunction(functionName, {(int32_t) instance.addValue(args[2])});
-        if (res.size() != 1)
-            throw Error("Wasm function '%s' from '%s' did not return exactly one value", functionName, wasmPath);
-        if (res[0].kind() != ValKind::I32)
-            throw Error("Wasm function '%s' from '%s' did not return an i32 value", functionName, wasmPath);
-        auto & vRes = instance.getValue(res[0].i32());
-        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_wasm(
-    {.name = "__wasm",
-     .args = {"wasm", "entry", "arg"},
-     .doc = R"(
-      Call a Wasm function with the specified argument.
-     )",
-     .fun = prim_wasm,
-     .experimentalFeature = Xp::WasmBuiltin});
-
 /**
  * Callback for WASI stdout/stderr writes. It splits the output into lines and logs each line separately.
  */
@@ -606,54 +583,121 @@ struct WasiLogger
     }
 };
 
-static void prim_wasi(EvalState & state, const PosIdx pos, Value ** args, Value & v)
+static void prim_wasm(EvalState & state, const PosIdx pos, Value ** args, Value & v)
 {
-    auto wasmPath = state.realisePath(pos, *args[0]);
-    auto functionName = "_start";
+    state.forceAttrs(*args[0], pos, "while evaluating the first argument to `builtins.wasm`");
+
+    // Extract 'path' attribute
+    auto pathAttr = args[0]->attrs()->get(state.symbols.create("path"));
+    if (!pathAttr)
+        throw Error("missing required 'path' attribute in first argument to `builtins.wasm`");
+    auto wasmPath = state.realisePath(pos, *pathAttr->value);
+
+    // Check for unknown attributes
+    for (auto & attr : *args[0]->attrs()) {
+        auto name = state.symbols[attr.name];
+        if (name != "path" && name != "function")
+            throw Error("unknown attribute '%s' in first argument to `builtins.wasm`", name);
+    }
+
+    // Second argument is the value to pass to the function
+    auto argValue = args[1];
 
     try {
-        auto instance = instantiateWasm(state, wasmPath, true);
+        auto instance = instantiateWasm(state, wasmPath);
+
+        // Extract 'function' attribute (optional for wasi, required for non-wasi)
+        std::string functionName;
+        auto functionAttr = args[0]->attrs()->get(state.symbols.create("function"));
+        if (instance.pre->useWasi) {
+            functionName = "_start";
+            if (functionAttr)
+                throw Error("'function' attribute is not allowed for WASI modules");
+        } else {
+            if (!functionAttr)
+                throw Error(
+                    "missing required 'function' attribute in first argument to `builtins.wasm` for non-WASI modules");
+            functionName = std::string(
+                state.forceStringNoCtx(*functionAttr->value, pos, "while evaluating the 'function' attribute"));
+        }
 
         debug("calling wasm module");
 
-        auto argId = instance.addValue(args[1]);
+        auto argId = instance.addValue(argValue);
 
-        WasiLogger logger{instance};
+        if (instance.pre->useWasi) {
+            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;
-        };
+            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)));
+            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 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;
+            auto & vRes = instance.getValue(instance.resultId);
+            state.forceValue(vRes, pos);
+            v = vRes;
+        } else {
+            // FIXME: use the "start" function if present.
+            instance.runFunction("nix_wasm_init_v1", {});
+
+            auto res = instance.runFunction(functionName, {(int32_t) argId});
+            if (res.size() != 1)
+                throw Error("Wasm function '%s' from '%s' did not return exactly one value", functionName, wasmPath);
+            if (res[0].kind() != ValKind::I32)
+                throw Error("Wasm function '%s' from '%s' did not return an i32 value", functionName, wasmPath);
+            auto & vRes = instance.getValue(res[0].i32());
+            state.forceValue(vRes, pos);
+            v = vRes;
+        }
     } catch (Error & e) {
-        e.addTrace(state.positions[pos], "while executing the Wasm function '%s' from '%s'", functionName, wasmPath);
+        e.addTrace(state.positions[pos], "while executing the Wasm function from '%s'", wasmPath);
         throw;
     }
 }
 
-static RegisterPrimOp primop_wasi(
-    {.name = "__wasi",
-     .args = {"wasi", "arg"},
+static RegisterPrimOp primop_wasm(
+    {.name = "__wasm",
+     .args = {"config", "arg"},
      .doc = R"(
-      Call a WASI function with the specified argument.
+      Call a Wasm function with the specified argument.
+
+      The first argument must be an attribute set with the following attributes:
+      - `path`: Path to the Wasm module (required)
+      - `function`: Function name to call (required for non-WASI modules, not allowed for WASI modules)
+
+      The second argument is the value to pass to the function.
+
+      WASI mode is automatically enabled if the module imports from `wasi_snapshot_preview1`.
+
+      Example (non-WASI):
+      ```nix
+      builtins.wasm {
+        path = ./foo.wasm;
+        function = "fib";
+      } 33
+      ```
+
+      Example (WASI):
+      ```nix
+      builtins.wasm {
+        path = ./bar.wasm;
+      } { x = 42; }
+      ```
      )",
-     .fun = prim_wasi,
+     .fun = prim_wasm,
      .experimentalFeature = Xp::WasmBuiltin});
 
 } // namespace nix