Skip to content

Commit 117b1bc

Browse files
philrhcPhil Cummins
philrhc
and
Phil Cummins
authored
adds stubwasi.rs (#111)
* adds stubwasi.rs * stubwasi.link_stub_modules returns Result * fix lib.rs --------- Co-authored-by: Phil Cummins <philip.cummins@bsc.es>
1 parent 0a440e3 commit 117b1bc

File tree

2 files changed

+131
-124
lines changed

2 files changed

+131
-124
lines changed

‎src/lib.rs‎

Lines changed: 2 additions & 124 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,6 @@ use {
1919
str,
2020
},
2121
summary::{Escape, Locations, Summary},
22-
wasm_convert::IntoValType,
23-
wasm_encoder::{
24-
CodeSection, ExportKind, ExportSection, Function, FunctionSection, Instruction, Module,
25-
TypeSection,
26-
},
27-
wasmparser::{FuncType, Parser, Payload, TypeRef},
2822
wasmtime::{
2923
component::{Component, Instance, Linker, ResourceTable, ResourceType},
3024
Config, Engine, Store,
@@ -43,6 +37,7 @@ pub mod command;
4337
mod prelink;
4438
#[cfg(feature = "pyo3")]
4539
mod python;
40+
mod stubwasi;
4641
mod summary;
4742
#[cfg(test)]
4843
mod test;
@@ -287,16 +282,12 @@ pub async fn componentize(
287282
// Link all the libraries (including any native extensions) into a single component.
288283
let mut linker = wit_component::Linker::default().validate(true);
289284

290-
let mut wasi_imports = HashMap::new();
291285
for Library {
292286
name,
293287
module,
294288
dl_openable,
295289
} in &libraries
296290
{
297-
if stub_wasi {
298-
add_wasi_imports(module, &mut wasi_imports)?;
299-
}
300291
linker = linker.library(name, module, *dl_openable)?;
301292
}
302293

@@ -311,51 +302,7 @@ pub async fn componentize(
311302
let component = linker.encode()?;
312303

313304
let stubbed_component = if stub_wasi {
314-
// When `stub_wasi` is `true`, we apply the pre-initialization snapshot to an alternate version of the
315-
// component -- one where the WASI imports have been stubbed out.
316-
317-
let mut linker = wit_component::Linker::default().validate(true);
318-
319-
for Library {
320-
name,
321-
module,
322-
dl_openable,
323-
} in &libraries
324-
{
325-
linker = linker.library(name, module, *dl_openable)?;
326-
}
327-
328-
for (module, imports) in &wasi_imports {
329-
linker = linker.adapter(module, &make_stub_adapter(module, imports))?;
330-
}
331-
332-
let component = linker.encode()?;
333-
334-
// As of this writing, `wit_component::Linker` generates a component such that the first module is the
335-
// `main` one, followed by any adapters, followed by any libraries, followed by the `init` module, which is
336-
// finally followed by any shim modules. Given that the stubbed component may contain more adapters than
337-
// the non-stubbed version, we need to tell `component-init` how to translate module indexes from the
338-
// former to the latter.
339-
//
340-
// TODO: this is pretty fragile in that it could silently break if `wit_component::Linker`'s implementation
341-
// changes. Can we make it more robust?
342-
343-
let old_adapter_count = 1;
344-
let new_adapter_count = u32::try_from(wasi_imports.len()).unwrap();
345-
assert!(new_adapter_count >= old_adapter_count);
346-
347-
Some((component, move |index: u32| {
348-
if index == 0 {
349-
// `main` module
350-
0
351-
} else if index <= new_adapter_count {
352-
// adapter module
353-
old_adapter_count
354-
} else {
355-
// one of the other kinds of module
356-
index + old_adapter_count - new_adapter_count
357-
}
358-
}))
305+
stubwasi::link_stub_modules(libraries)?
359306
} else {
360307
None
361308
};
@@ -694,72 +641,3 @@ fn add_wasi_and_stubs(
694641

695642
Ok(())
696643
}
697-
698-
fn add_wasi_imports<'a>(
699-
module: &'a [u8],
700-
imports: &mut HashMap<&'a str, HashMap<&'a str, FuncType>>,
701-
) -> Result<()> {
702-
let mut types = Vec::new();
703-
for payload in Parser::new(0).parse_all(module) {
704-
match payload? {
705-
Payload::TypeSection(reader) => {
706-
types = reader
707-
.into_iter_err_on_gc_types()
708-
.collect::<Result<Vec<_>, _>>()?;
709-
}
710-
711-
Payload::ImportSection(reader) => {
712-
for import in reader {
713-
let import = import?;
714-
715-
if import.module == "wasi_snapshot_preview1"
716-
|| import.module.starts_with("wasi:")
717-
{
718-
if let TypeRef::Func(ty) = import.ty {
719-
imports
720-
.entry(import.module)
721-
.or_default()
722-
.insert(import.name, types[usize::try_from(ty).unwrap()].clone());
723-
} else {
724-
bail!("encountered non-function import from WASI namespace")
725-
}
726-
}
727-
}
728-
break;
729-
}
730-
731-
_ => {}
732-
}
733-
}
734-
735-
Ok(())
736-
}
737-
738-
fn make_stub_adapter(_module: &str, stubs: &HashMap<&str, FuncType>) -> Vec<u8> {
739-
let mut types = TypeSection::new();
740-
let mut functions = FunctionSection::new();
741-
let mut exports = ExportSection::new();
742-
let mut code = CodeSection::new();
743-
744-
for (index, (name, ty)) in stubs.iter().enumerate() {
745-
let index = u32::try_from(index).unwrap();
746-
types.ty().function(
747-
ty.params().iter().map(|&v| IntoValType(v).into()),
748-
ty.results().iter().map(|&v| IntoValType(v).into()),
749-
);
750-
functions.function(index);
751-
exports.export(name, ExportKind::Func, index);
752-
let mut function = Function::new([]);
753-
function.instruction(&Instruction::Unreachable);
754-
function.instruction(&Instruction::End);
755-
code.function(&function);
756-
}
757-
758-
let mut module = Module::new();
759-
module.section(&types);
760-
module.section(&functions);
761-
module.section(&exports);
762-
module.section(&code);
763-
764-
module.finish()
765-
}

‎src/stubwasi.rs‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
use std::collections::HashMap;
2+
3+
use anyhow::{bail, Error};
4+
use wasm_convert::IntoValType;
5+
use wasm_encoder::{
6+
CodeSection, ExportKind, ExportSection, Function, FunctionSection, Instruction as Ins, Module,
7+
TypeSection,
8+
};
9+
use wasmparser::{FuncType, Parser, Payload, TypeRef};
10+
11+
use crate::Library;
12+
13+
pub fn link_stub_modules(
14+
libraries: Vec<Library>,
15+
) -> Result<Option<(Vec<u8>, impl Fn(u32) -> u32)>, Error> {
16+
let mut wasi_imports = HashMap::new();
17+
let mut linker = wit_component::Linker::default().validate(true);
18+
19+
for Library {
20+
name,
21+
module,
22+
dl_openable,
23+
} in &libraries
24+
{
25+
add_wasi_imports(module, &mut wasi_imports)?;
26+
linker = linker.library(name, module, *dl_openable)?;
27+
}
28+
29+
for (module, imports) in &wasi_imports {
30+
linker = linker.adapter(module, &make_stub_adapter(module, imports))?;
31+
}
32+
33+
let component = linker.encode()?;
34+
35+
// As of this writing, `wit_component::Linker` generates a component such that the first module is the
36+
// `main` one, followed by any adapters, followed by any libraries, followed by the `init` module, which is
37+
// finally followed by any shim modules. Given that the stubbed component may contain more adapters than
38+
// the non-stubbed version, we need to tell `component-init` how to translate module indexes from the
39+
// former to the latter.
40+
//
41+
// TODO: this is pretty fragile in that it could silently break if `wit_component::Linker`'s implementation
42+
// changes. Can we make it more robust?
43+
44+
let old_adapter_count = 1;
45+
let new_adapter_count = u32::try_from(wasi_imports.len())?;
46+
assert!(new_adapter_count >= old_adapter_count);
47+
48+
Ok(Some((component, move |index: u32| {
49+
if index == 0 {
50+
// `main` module
51+
0
52+
} else if index <= new_adapter_count {
53+
// adapter module
54+
old_adapter_count
55+
} else {
56+
// one of the other kinds of module
57+
index + old_adapter_count - new_adapter_count
58+
}
59+
})))
60+
}
61+
62+
fn add_wasi_imports<'a>(
63+
module: &'a [u8],
64+
imports: &mut HashMap<&'a str, HashMap<&'a str, FuncType>>,
65+
) -> Result<(), Error> {
66+
let mut types = Vec::new();
67+
for payload in Parser::new(0).parse_all(module) {
68+
match payload? {
69+
Payload::TypeSection(reader) => {
70+
types = reader
71+
.into_iter_err_on_gc_types()
72+
.collect::<Result<Vec<_>, _>>()?;
73+
}
74+
75+
Payload::ImportSection(reader) => {
76+
for import in reader {
77+
let import = import?;
78+
79+
if import.module == "wasi_snapshot_preview1"
80+
|| import.module.starts_with("wasi:")
81+
{
82+
if let TypeRef::Func(ty) = import.ty {
83+
imports
84+
.entry(import.module)
85+
.or_default()
86+
.insert(import.name, types[usize::try_from(ty).unwrap()].clone());
87+
} else {
88+
bail!("encountered non-function import from WASI namespace")
89+
}
90+
}
91+
}
92+
break;
93+
}
94+
95+
_ => {}
96+
}
97+
}
98+
99+
Ok(())
100+
}
101+
102+
fn make_stub_adapter(_module: &str, stubs: &HashMap<&str, FuncType>) -> Vec<u8> {
103+
let mut types = TypeSection::new();
104+
let mut functions = FunctionSection::new();
105+
let mut exports = ExportSection::new();
106+
let mut code = CodeSection::new();
107+
108+
for (index, (name, ty)) in stubs.iter().enumerate() {
109+
let index = u32::try_from(index).unwrap();
110+
types.ty().function(
111+
ty.params().iter().map(|&v| IntoValType(v).into()),
112+
ty.results().iter().map(|&v| IntoValType(v).into()),
113+
);
114+
functions.function(index);
115+
exports.export(name, ExportKind::Func, index);
116+
let mut function = Function::new([]);
117+
function.instruction(&Ins::Unreachable);
118+
function.instruction(&Ins::End);
119+
code.function(&function);
120+
}
121+
122+
let mut module = Module::new();
123+
module.section(&types);
124+
module.section(&functions);
125+
module.section(&exports);
126+
module.section(&code);
127+
128+
module.finish()
129+
}

0 commit comments

Comments
 (0)