Skip to content

Commit 7f10286

Browse files
edolstraedolstra
authored
Merge pull request #405 from DeterminateSystems/wasm-improvements
builtins.wasm: Support WAT, add a test
2 parents 2f1ca41 + 52eb36a commit 7f10286

5 files changed

Lines changed: 124 additions & 39 deletions

File tree

‎src/libexpr/primops/wasm.cc‎

Lines changed: 71 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -82,34 +82,44 @@ static void regFuns(Linker & linker, bool useWasi);
8282

8383
struct NixWasmInstancePre
8484
{
85-
Engine & engine;
86-
SourcePath wasmPath;
87-
bool useWasi;
85+
Engine & engine = getEngine();
86+
std::string name;
87+
bool useWasi = false;
8888
InstancePre instancePre;
8989

90-
NixWasmInstancePre(SourcePath _wasmPath)
91-
: engine(getEngine())
92-
, wasmPath(_wasmPath)
93-
, useWasi(false)
94-
, instancePre(({
95-
// Compile the module
96-
auto module = unwrap(Module::compile(engine, string2span(wasmPath.readFile())));
90+
InstancePre compile(std::span<uint8_t> bytes)
91+
{
92+
// Compile the module
93+
auto module = unwrap(Module::compile(engine, bytes));
94+
95+
// Auto-detect WASI by checking for wasi_snapshot_preview1 imports.
96+
for (const auto & ref : module.imports())
97+
if (const_cast<std::decay_t<decltype(ref)> &>(ref).module() == "wasi_snapshot_preview1") {
98+
useWasi = true;
99+
break;
100+
}
97101

98-
// Auto-detect WASI by checking for wasi_snapshot_preview1 imports.
99-
for (const auto & ref : module.imports())
100-
if (const_cast<std::decay_t<decltype(ref)> &>(ref).module() == "wasi_snapshot_preview1") {
101-
useWasi = true;
102-
break;
103-
}
102+
// Create linker with appropriate WASI support
103+
Linker linker(engine);
104+
if (useWasi)
105+
unwrap(linker.define_wasi());
106+
regFuns(linker, useWasi);
104107

105-
// Create linker with appropriate WASI support
106-
Linker linker(engine);
107-
if (useWasi)
108-
unwrap(linker.define_wasi());
109-
regFuns(linker, useWasi);
108+
return unwrap(instantiate_pre(linker, module));
109+
}
110110

111-
unwrap(instantiate_pre(linker, module));
112-
}))
111+
NixWasmInstancePre(SourcePath wasmPath)
112+
: name(wasmPath.baseName())
113+
, instancePre(compile(string2span(wasmPath.readFile())))
114+
{
115+
}
116+
117+
NixWasmInstancePre(std::string_view wat)
118+
: name("<inline wat>")
119+
, instancePre([&] {
120+
auto wasm = unwrap(wat2wasm(wat));
121+
return compile(std::span<uint8_t>(wasm));
122+
}())
113123
{
114124
}
115125
};
@@ -139,7 +149,7 @@ struct NixWasmInstance
139149
, wasmCtx(wasmStore)
140150
, instance(unwrap(pre->instancePre.instantiate(wasmCtx)))
141151
, memory_(getExport<Memory>("memory"))
142-
, logPrefix(pre->wasmPath.baseName())
152+
, logPrefix(pre->name)
143153
{
144154
wasmCtx.set_data(this);
145155

@@ -173,10 +183,10 @@ struct NixWasmInstance
173183
{
174184
auto ext = instance.get(wasmCtx, name);
175185
if (!ext)
176-
throw Error("Wasm module '%s' does not export '%s'", pre->wasmPath, name);
186+
throw Error("Wasm module '%s' does not export '%s'", pre->name, name);
177187
auto res = std::get_if<T>(&*ext);
178188
if (!res)
179-
throw Error("export '%s' of Wasm module '%s' does not have the right type", name, pre->wasmPath);
189+
throw Error("export '%s' of Wasm module '%s' does not have the right type", name, pre->name);
180190
return *res;
181191
}
182192

@@ -593,24 +603,30 @@ static void prim_wasm(EvalState & state, const PosIdx pos, Value ** args, Value
593603
{
594604
state.forceAttrs(*args[0], pos, "while evaluating the first argument to `builtins.wasm`");
595605

596-
// Extract 'path' attribute
597-
auto pathAttr = args[0]->attrs()->get(state.symbols.create("path"));
598-
if (!pathAttr)
599-
throw Error("missing required 'path' attribute in first argument to `builtins.wasm`");
600-
auto wasmPath = state.realisePath(pos, *pathAttr->value);
601-
602606
// Check for unknown attributes
603607
for (auto & attr : *args[0]->attrs()) {
604608
auto name = state.symbols[attr.name];
605-
if (name != "path" && name != "function")
609+
if (name != "path" && name != "wat" && name != "function")
606610
throw Error("unknown attribute '%s' in first argument to `builtins.wasm`", name);
607611
}
608612

613+
auto pathAttr = args[0]->attrs()->get(state.symbols.create("path"));
614+
auto watAttr = args[0]->attrs()->get(state.symbols.create("wat"));
615+
616+
if (pathAttr && watAttr)
617+
throw Error("'path' and 'wat' are mutually exclusive in first argument to `builtins.wasm`");
618+
if (!pathAttr && !watAttr)
619+
throw Error("missing required 'path' or 'wat' attribute in first argument to `builtins.wasm`");
620+
609621
// Second argument is the value to pass to the function
610622
auto argValue = args[1];
611623

612624
try {
613-
auto instance = instantiateWasm(state, wasmPath);
625+
auto instance = pathAttr ? instantiateWasm(state, state.realisePath(pos, *pathAttr->value))
626+
: NixWasmInstance{
627+
state,
628+
make_ref<NixWasmInstancePre>(state.forceStringNoCtx(
629+
*watAttr->value, pos, "while evaluating the 'wat' attribute"))};
614630

615631
// Extract 'function' attribute (optional for wasi, required for non-wasi)
616632
std::string functionName;
@@ -649,7 +665,10 @@ static void prim_wasm(EvalState & state, const PosIdx pos, Value ** args, Value
649665
auto res = instance.getExport<Func>(functionName).call(instance.wasmCtx, {});
650666
if (!instance.resultId) {
651667
unwrap(std::move(res));
652-
throw Error("Wasm function '%s' from '%s' finished without returning a value", functionName, wasmPath);
668+
throw Error(
669+
"Wasm function '%s' from '%s' finished without returning a value",
670+
functionName,
671+
instance.pre->name);
653672
}
654673

655674
auto & vRes = instance.getValue(instance.resultId);
@@ -661,15 +680,17 @@ static void prim_wasm(EvalState & state, const PosIdx pos, Value ** args, Value
661680

662681
auto res = instance.runFunction(functionName, {(int32_t) argId});
663682
if (res.size() != 1)
664-
throw Error("Wasm function '%s' from '%s' did not return exactly one value", functionName, wasmPath);
683+
throw Error(
684+
"Wasm function '%s' from '%s' did not return exactly one value", functionName, instance.pre->name);
665685
if (res[0].kind() != ValKind::I32)
666-
throw Error("Wasm function '%s' from '%s' did not return an i32 value", functionName, wasmPath);
686+
throw Error(
687+
"Wasm function '%s' from '%s' did not return an i32 value", functionName, instance.pre->name);
667688
auto & vRes = instance.getValue(res[0].i32());
668689
state.forceValue(vRes, pos);
669690
v = vRes;
670691
}
671692
} catch (Error & e) {
672-
e.addTrace(state.positions[pos], "while executing the Wasm function from '%s'", wasmPath);
693+
e.addTrace(state.positions[pos], "while executing a Wasm module");
673694
throw;
674695
}
675696
}
@@ -681,9 +702,12 @@ static RegisterPrimOp primop_wasm(
681702
Call a Wasm function with the specified argument.
682703
683704
The first argument must be an attribute set with the following attributes:
684-
- `path`: Path to the Wasm module (required)
705+
- `path`: Path to the Wasm module (mutually exclusive with `wat`)
706+
- `wat`: WebAssembly Text format source as a string (mutually exclusive with `path`)
685707
- `function`: Function name to call (required for non-WASI modules, not allowed for WASI modules)
686708
709+
Exactly one of `path` or `wat` must be specified.
710+
687711
The second argument is the value to pass to the function.
688712
689713
WASI mode is automatically enabled if the module imports from `wasi_snapshot_preview1`.
@@ -696,6 +720,14 @@ static RegisterPrimOp primop_wasm(
696720
} 33
697721
```
698722
723+
Example (reading from a WAT file):
724+
```nix
725+
builtins.wasm {
726+
wat = builtins.readFile ./fib.wat;
727+
function = "fib";
728+
} 10
729+
```
730+
699731
Example (WASI):
700732
```nix
701733
builtins.wasm {

‎tests/functional/fib.wasm‎

160 Bytes
Binary file not shown.

‎tests/functional/fib.wat‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
(module
2+
;; Import host functions from the Nix env module
3+
(import "env" "get_int" (func $get_int (param i32) (result i64)))
4+
(import "env" "make_int" (func $make_int (param i64) (result i32)))
5+
6+
;; The host requires an exported memory to read/write data
7+
(memory (export "memory") 1)
8+
9+
;; Called once when the module is instantiated; nothing to initialize here
10+
(func (export "nix_wasm_init_v1"))
11+
12+
;; Pure wasm: compute fib(n) recursively
13+
;; fib(0) = 1
14+
;; fib(1) = 1
15+
;; fib(n) = fib(n-1) + fib(n-2)
16+
(func $fib (param $n i64) (result i64)
17+
(if (i64.le_s (local.get $n) (i64.const 1))
18+
(then (return (i64.const 1)))
19+
)
20+
(i64.add
21+
(call $fib (i64.sub (local.get $n) (i64.const 1)))
22+
(call $fib (i64.sub (local.get $n) (i64.const 2)))
23+
)
24+
)
25+
26+
;; Entry point: receives ValueId of the input integer, returns ValueId of fib(n)
27+
;; Type: fn(arg: u32) -> u32 (ValueId = u32)
28+
(func (export "fib") (param $arg i32) (result i32)
29+
(call $make_int
30+
(call $fib
31+
(call $get_int (local.get $arg))
32+
)
33+
)
34+
)
35+
)

‎tests/functional/meson.build‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,7 @@ suites = [
175175
'help.sh',
176176
'symlinks.sh',
177177
'external-builders.sh',
178+
'wasm.sh',
178179
],
179180
'workdir' : meson.current_source_dir(),
180181
},

‎tests/functional/wasm.sh‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
#!/usr/bin/env bash
2+
3+
source common.sh
4+
5+
if [[ $(nix eval --extra-experimental-features wasm-builtin --expr 'builtins ? wasm') = false ]]; then
6+
skipTest "builtins.wasm not available"
7+
fi
8+
9+
# Test running a WebAssembly module in text format (WAT).
10+
[[ $(nix eval --json --impure \
11+
--extra-experimental-features wasm-builtin \
12+
--expr "builtins.wasm { wat = builtins.readFile ./fib.wat; function = \"fib\"; } 40") = 165580141 ]]
13+
14+
# Test running a WebAssembly module in binary format (.wasm).
15+
[[ $(nix eval --json --impure \
16+
--extra-experimental-features wasm-builtin \
17+
--expr "builtins.wasm { path = ./fib.wasm; function = \"fib\"; } 40") = 165580141 ]]

0 commit comments

Comments
 (0)