1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

#[cfg(features = "std")]
use std::collections::{HashSet as Set};
#[cfg(not(features = "std"))]
use std::collections::{BTreeSet as Set};
use std::vec::Vec;

use parity_wasm::elements;

#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Copy, Clone, Debug)]
pub enum Symbol {
	Type(usize),
	Import(usize),
	Global(usize),
	Function(usize),
	Export(usize),
}

pub fn resolve_function(module: &elements::Module, index: u32) -> Symbol {
	let mut functions = 0;
	if let Some(import_section) = module.import_section() {
		for (item_index, item) in import_section.entries().iter().enumerate() {
			if let &elements::External::Function(_) = item.external() {
				if functions == index {
					return Symbol::Import(item_index as usize);
				}
				functions += 1;
			}
		}
	}

	Symbol::Function(index as usize - functions as usize)
}

pub fn resolve_global(module: &elements::Module, index: u32) -> Symbol {
	let mut globals = 0;
	if let Some(import_section) = module.import_section() {
		for (item_index, item) in import_section.entries().iter().enumerate() {
			if let &elements::External::Global(_) = item.external() {
				if globals == index {
					return Symbol::Import(item_index as usize);
				}
				globals += 1;
			}
		}
	}

	Symbol::Global(index as usize - globals as usize)
}

pub fn push_code_symbols(module: &elements::Module, instructions: &[elements::Instruction], dest: &mut Vec<Symbol>) {
	use parity_wasm::elements::Instruction::*;

	for instruction in instructions {
		match instruction {
			&Call(idx) => {
				dest.push(resolve_function(module, idx));
			},
			&CallIndirect(idx, _) => {
				dest.push(Symbol::Type(idx as usize));
			},
			&GetGlobal(idx) | &SetGlobal(idx) => {
				dest.push(resolve_global(module, idx))
			},
			_ => { },
		}
	}
}

pub fn expand_symbols(module: &elements::Module, set: &mut Set<Symbol>) {
	use self::Symbol::*;

	// symbols that were already processed
	let mut stop: Set<Symbol> = Set::new();
	let mut fringe = set.iter().cloned().collect::<Vec<Symbol>>();
	loop {
		let next = match fringe.pop() {
			Some(s) if stop.contains(&s) => { continue; }
			Some(s) => s,
			_ => { break; }
		};
		trace!("Processing symbol {:?}", next);

		match next {
			Export(idx) => {
				let entry = &module.export_section().expect("Export section to exist").entries()[idx];
				match entry.internal() {
					&elements::Internal::Function(func_idx) => {
						let symbol = resolve_function(module, func_idx);
						if !stop.contains(&symbol) {
							fringe.push(symbol);
						}
						set.insert(symbol);
					},
					&elements::Internal::Global(global_idx) => {
						let symbol = resolve_global(module, global_idx);
						if !stop.contains(&symbol) {
							fringe.push(symbol);
						}
						set.insert(symbol);
					},
					_ => {}
				}
			},
			Import(idx) => {
				let entry = &module.import_section().expect("Import section to exist").entries()[idx];
				if let &elements::External::Function(type_idx) = entry.external() {
					let type_symbol = Symbol::Type(type_idx as usize);
					if !stop.contains(&type_symbol) {
						fringe.push(type_symbol);
					}
					set.insert(type_symbol);
				}
			},
			Function(idx) => {
				let body = &module.code_section().expect("Code section to exist").bodies()[idx];
				let mut code_symbols = Vec::new();
				push_code_symbols(module, body.code().elements(), &mut code_symbols);
				for symbol in code_symbols.drain(..) {
					if !stop.contains(&symbol) {
						fringe.push(symbol);
					}
					set.insert(symbol);
				}

				let signature = &module.function_section().expect("Functions section to exist").entries()[idx];
				let type_symbol = Symbol::Type(signature.type_ref() as usize);
				if !stop.contains(&type_symbol) {
					fringe.push(type_symbol);
				}
				set.insert(type_symbol);
			},
			Global(idx) => {
				let entry = &module.global_section().expect("Global section to exist").entries()[idx];
				let mut code_symbols = Vec::new();
				push_code_symbols(module, entry.init_expr().code(), &mut code_symbols);
				for symbol in code_symbols.drain(..) {
					if !stop.contains(&symbol) {
						fringe.push(symbol);
					}
					set.insert(symbol);
				}
			}
			_ => {}
		}

		stop.insert(next);
	}
}