wasm_instrument/
export_globals.rs

1use alloc::{format, vec::Vec};
2use parity_wasm::elements;
3
4/// Export all declared mutable globals as `prefix_index`.
5///
6/// This will export all internal mutable globals under the name of
7/// concat(`prefix`, `"_"`, `i`) where i is the index inside the range of
8/// [0..total number of internal mutable globals].
9pub fn export_mutable_globals(module: &mut elements::Module, prefix: &str) {
10	let exports = global_section(module)
11		.map(|section| {
12			section
13				.entries()
14				.iter()
15				.enumerate()
16				.filter_map(
17					|(index, global)| {
18						if global.global_type().is_mutable() {
19							Some(index)
20						} else {
21							None
22						}
23					},
24				)
25				.collect::<Vec<_>>()
26		})
27		.unwrap_or_default();
28
29	if module.export_section().is_none() {
30		module
31			.sections_mut()
32			.push(elements::Section::Export(elements::ExportSection::default()));
33	}
34
35	for (symbol_index, export) in exports.into_iter().enumerate() {
36		let new_entry = elements::ExportEntry::new(
37			format!("{}_{}", prefix, symbol_index),
38			elements::Internal::Global(
39				(module.import_count(elements::ImportCountType::Global) + export) as _,
40			),
41		);
42		export_section(module)
43			.expect("added above if does not exists")
44			.entries_mut()
45			.push(new_entry);
46	}
47}
48
49fn export_section(module: &mut elements::Module) -> Option<&mut elements::ExportSection> {
50	for section in module.sections_mut() {
51		if let elements::Section::Export(sect) = section {
52			return Some(sect)
53		}
54	}
55	None
56}
57
58fn global_section(module: &mut elements::Module) -> Option<&mut elements::GlobalSection> {
59	for section in module.sections_mut() {
60		if let elements::Section::Global(sect) = section {
61			return Some(sect)
62		}
63	}
64	None
65}
66
67#[cfg(test)]
68mod tests {
69
70	use super::export_mutable_globals;
71	use parity_wasm::elements;
72
73	fn parse_wat(source: &str) -> elements::Module {
74		let module_bytes = wat::parse_str(source).unwrap();
75		wasmparser::validate(&module_bytes).unwrap();
76		elements::deserialize_buffer(module_bytes.as_ref()).expect("failed to parse module")
77	}
78
79	macro_rules! test_export_global {
80		(name = $name:ident; input = $input:expr; expected = $expected:expr) => {
81			#[test]
82			fn $name() {
83				let mut input_module = parse_wat($input);
84				let expected_module = parse_wat($expected);
85
86				export_mutable_globals(&mut input_module, "exported_internal_global");
87
88				let actual_bytes = elements::serialize(input_module)
89					.expect("injected module must have a function body");
90
91				let expected_bytes = elements::serialize(expected_module)
92					.expect("injected module must have a function body");
93
94				let actual_wat = wasmprinter::print_bytes(actual_bytes).unwrap();
95				let expected_wat = wasmprinter::print_bytes(expected_bytes).unwrap();
96
97				if actual_wat != expected_wat {
98					for diff in diff::lines(&expected_wat, &actual_wat) {
99						match diff {
100							diff::Result::Left(l) => println!("-{}", l),
101							diff::Result::Both(l, _) => println!(" {}", l),
102							diff::Result::Right(r) => println!("+{}", r),
103						}
104					}
105					panic!()
106				}
107			}
108		};
109	}
110
111	test_export_global! {
112		name = simple;
113		input = r#"
114		(module
115			(global (;0;) (mut i32) (i32.const 1))
116			(global (;1;) (mut i32) (i32.const 0)))
117		"#;
118		expected = r#"
119		(module
120			(global (;0;) (mut i32) (i32.const 1))
121			(global (;1;) (mut i32) (i32.const 0))
122			(export "exported_internal_global_0" (global 0))
123			(export "exported_internal_global_1" (global 1)))
124		"#
125	}
126
127	test_export_global! {
128		name = with_import;
129		input = r#"
130		(module
131			(import "env" "global" (global $global i64))
132			(global (;0;) (mut i32) (i32.const 1))
133			(global (;1;) (mut i32) (i32.const 0)))
134		"#;
135		expected = r#"
136		(module
137			(import "env" "global" (global $global i64))
138			(global (;0;) (mut i32) (i32.const 1))
139			(global (;1;) (mut i32) (i32.const 0))
140			(export "exported_internal_global_0" (global 1))
141			(export "exported_internal_global_1" (global 2)))
142		"#
143	}
144
145	test_export_global! {
146		name = with_import_and_some_are_immutable;
147		input = r#"
148		(module
149			(import "env" "global" (global $global i64))
150			(global (;0;) i32 (i32.const 1))
151			(global (;1;) (mut i32) (i32.const 0)))
152		"#;
153		expected = r#"
154		(module
155			(import "env" "global" (global $global i64))
156			(global (;0;) i32 (i32.const 1))
157			(global (;1;) (mut i32) (i32.const 0))
158			(export "exported_internal_global_0" (global 2)))
159		"#
160	}
161}
wasm_instrument/export_globals.rs

wasm_instrument/
export_globals.rs
