Skip to content

Commit ad36554

Browse files
author
Phil Cummins
committed
adds stubwasi.rs
1 parent f19009e commit ad36554

File tree

2 files changed

+132
-124
lines changed

2 files changed

+132
-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;
@@ -282,16 +277,12 @@ pub async fn componentize(
282277
// Link all the libraries (including any native extensions) into a single component.
283278
let mut linker = wit_component::Linker::default().validate(true);
284279

285-
let mut wasi_imports = HashMap::new();
286280
for Library {
287281
name,
288282
module,
289283
dl_openable,
290284
} in &libraries
291285
{
292-
if stub_wasi {
293-
add_wasi_imports(module, &mut wasi_imports)?;
294-
}
295286
linker = linker.library(name, module, *dl_openable)?;
296287
}
297288

@@ -306,51 +297,7 @@ pub async fn componentize(
306297
let component = linker.encode()?;
307298

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

673620
Ok(())
674621
}
675-
676-
fn add_wasi_imports<'a>(
677-
module: &'a [u8],
678-
imports: &mut HashMap<&'a str, HashMap<&'a str, FuncType>>,
679-
) -> Result<()> {
680-
let mut types = Vec::new();
681-
for payload in Parser::new(0).parse_all(module) {
682-
match payload? {
683-
Payload::TypeSection(reader) => {
684-
types = reader
685-
.into_iter_err_on_gc_types()
686-
.collect::<Result<Vec<_>, _>>()?;
687-
}
688-
689-
Payload::ImportSection(reader) => {
690-
for import in reader {
691-
let import = import?;
692-
693-
if import.module == "wasi_snapshot_preview1"
694-
|| import.module.starts_with("wasi:")
695-
{
696-
if let TypeRef::Func(ty) = import.ty {
697-
imports
698-
.entry(import.module)
699-
.or_default()
700-
.insert(import.name, types[usize::try_from(ty).unwrap()].clone());
701-
} else {
702-
bail!("encountered non-function import from WASI namespace")
703-
}
704-
}
705-
}
706-
break;
707-
}
708-
709-
_ => {}
710-
}
711-
}
712-
713-
Ok(())
714-
}
715-
716-
fn make_stub_adapter(_module: &str, stubs: &HashMap<&str, FuncType>) -> Vec<u8> {
717-
let mut types = TypeSection::new();
718-
let mut functions = FunctionSection::new();
719-
let mut exports = ExportSection::new();
720-
let mut code = CodeSection::new();
721-
722-
for (index, (name, ty)) in stubs.iter().enumerate() {
723-
let index = u32::try_from(index).unwrap();
724-
types.ty().function(
725-
ty.params().iter().map(|&v| IntoValType(v).into()),
726-
ty.results().iter().map(|&v| IntoValType(v).into()),
727-
);
728-
functions.function(index);
729-
exports.export(name, ExportKind::Func, index);
730-
let mut function = Function::new([]);
731-
function.instruction(&Instruction::Unreachable);
732-
function.instruction(&Instruction::End);
733-
code.function(&function);
734-
}
735-
736-
let mut module = Module::new();
737-
module.section(&types);
738-
module.section(&functions);
739-
module.section(&exports);
740-
module.section(&code);
741-
742-
module.finish()
743-
}

‎src/stubwasi.rs‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
use std::collections::HashMap;
2+
3+
use anyhow::bail;
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(libraries: Vec<Library>) -> Option<(Vec<u8>, impl Fn(u32) -> u32)> {
14+
let mut wasi_imports = HashMap::new();
15+
let mut linker = wit_component::Linker::default().validate(true);
16+
17+
for Library {
18+
name,
19+
module,
20+
dl_openable,
21+
} in &libraries
22+
{
23+
add_wasi_imports(module, &mut wasi_imports).unwrap();
24+
linker = linker.library(name, module, *dl_openable).unwrap();
25+
}
26+
27+
for (module, imports) in &wasi_imports {
28+
linker = linker
29+
.adapter(module, &make_stub_adapter(module, imports))
30+
.unwrap();
31+
}
32+
33+
let component = linker.encode().unwrap();
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()).unwrap();
46+
assert!(new_adapter_count >= old_adapter_count);
47+
48+
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<(), anyhow::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+
.unwrap();
74+
}
75+
76+
Payload::ImportSection(reader) => {
77+
for import in reader {
78+
let import = import?;
79+
80+
if import.module == "wasi_snapshot_preview1"
81+
|| import.module.starts_with("wasi:")
82+
{
83+
if let TypeRef::Func(ty) = import.ty {
84+
imports
85+
.entry(import.module)
86+
.or_default()
87+
.insert(import.name, types[usize::try_from(ty).unwrap()].clone());
88+
} else {
89+
bail!("encountered non-function import from WASI namespace")
90+
}
91+
}
92+
}
93+
break;
94+
}
95+
96+
_ => {}
97+
}
98+
}
99+
100+
Ok(())
101+
}
102+
103+
fn make_stub_adapter(_module: &str, stubs: &HashMap<&str, FuncType>) -> Vec<u8> {
104+
let mut types = TypeSection::new();
105+
let mut functions = FunctionSection::new();
106+
let mut exports = ExportSection::new();
107+
let mut code = CodeSection::new();
108+
109+
for (index, (name, ty)) in stubs.iter().enumerate() {
110+
let index = u32::try_from(index).unwrap();
111+
types.ty().function(
112+
ty.params().iter().map(|&v| IntoValType(v).into()),
113+
ty.results().iter().map(|&v| IntoValType(v).into()),
114+
);
115+
functions.function(index);
116+
exports.export(name, ExportKind::Func, index);
117+
let mut function = Function::new([]);
118+
function.instruction(&Ins::Unreachable);
119+
function.instruction(&Ins::End);
120+
code.function(&function);
121+
}
122+
123+
let mut module = Module::new();
124+
module.section(&types);
125+
module.section(&functions);
126+
module.section(&exports);
127+
module.section(&code);
128+
129+
module.finish()
130+
}

0 commit comments

Comments
 (0)