Instantiator.v3

// Copyright 2020 Ben L. Titzer. All rights reserved.
// See LICENSE for details of Apache 2.0 license.

// Extension point for import processing. An {ImportProcessor} lookups or materializes
// an {Exportable} for a given name and import kind. It can use the {binder} as a
// callback to access parts of the module that need polymorphic substitution.
class ImportProcessor(name: string) {
	def preprocess(error: ErrorGen, module: Module, imports: Array<Exportable>) { }
	def processFunction(error: ErrorGen, name: string, decl: FuncDecl) -> Exportable { return null; }
	def processTable(error: ErrorGen, name: string, decl: TableDecl) -> Exportable { return null; }
	def processMemory(error: ErrorGen, name: string, decl: MemoryDecl) -> Exportable { return null; }
	def processGlobal(error: ErrorGen, name: string, decl: GlobalDecl) -> Exportable { return null; }
	def processTag(error: ErrorGen, name: string, decl: TagDecl) -> Exportable { return null; }
	def postprocess(error: ErrorGen, instance: Instance) { }
}

// Creates {Instance} objects, given a module and a list of imports.
class Instantiator(extensions: Extension.set, module: Module, var imports: Array<Exportable>, error: ErrorGen) { // TODO: s/error/err
	def instance = Instance.new(module, imports);
	def processors = Vector<ImportProcessor>.new();
	var tiering: ExecutionStrategy;
	var binder: (Decl, Exportable) -> Exportable;		// last binding step
	var trap_reason: TrapReason;
	var import_pos: u31;
	var import_info: ImportInfo;

	private var partially_instantiated = false;
	private var fully_instantiated = false;

	def run() -> Instance {
		return runPartial(null);
	}

	// Initialize the instance with the given processors to populate imports.
	// `deferred_imports` contains the import indices that should be skipped during first pass
	// initialization. It represents the imports from preloaded wasm modules that are dynamically linked.
	//
	// For single module instantiation, `deferred_imports` is null or empty.
	// For dynamic linking with multiple wasm modules, this method is called twice with the same set of `deferred_imports`.
	// The instantiator keeps track of partial/completed instantiation state.
	def runPartial(deferred_imports: Array<int>) -> Instance {
		if (fully_instantiated) return instance;
		//if (!partially_instantiated) {
		if (tiering != null) tiering.onInstantiateStart(this);
		if (module.imports.length > 0) {
			if (imports == null) imports = Array.new(module.imports.length);
			else if (imports.length < module.imports.length) {
				return fail(Strings.format2("expected %d imports, got %d", module.imports.length, imports.length));
			}
		}
		for (i < processors.length) {
			processors[i].preprocess(error, module, imports);
		}

		var needSigIds = false;
		for (j < module.decls.length) {
			if (error.error()) return null;
			match (module.decls[j]) {
				Sig(index) => {
					instance.heaptypes[index] = module.heaptypes[index];
				}
				Struct(index) => {
					instance.heaptypes[index] = module.heaptypes[index];
				}
				Array(index) => {
					instance.heaptypes[index] = module.heaptypes[index];
				}
				Func(index, sig_index) => {
					var decl = module.functions[index];
					if ((import_info = decl.imp) != null) {
						if (partially_instantiated != checkDeferredImport(decl.imp.import_index, deferred_imports)) continue;
						var r = imports[decl.imp.import_index];
						instance.functions[index] = importFunction(decl, r);
					} else if (!partially_instantiated) {
						var f = WasmFunction.new(instance, mapFuncDecl(index));
						instance.functions[index] = f;
					}
				}
				Table(index) => {
					var decl = module.tables[index];
					if ((import_info = decl.imp) != null) {
						if (partially_instantiated != checkDeferredImport(decl.imp.import_index, deferred_imports)) continue;
						var r = imports[decl.imp.import_index];
						var t =  importTable(decl, r);
						instance.tables[index] = t;
						if (t != null && t.funcs != null) needSigIds = true;
					} else if (!partially_instantiated) {
						var t = Table.new(decl.elemtype, decl);
						if (t.oom) return fail("out of memory allocating table");
						instance.tables[index] = t;
						if (t.funcs != null) needSigIds = true;
					}
				}
				Memory(index) => {
					var decl = module.memories[index];
					if ((import_info = decl.imp) != null) {
						if (partially_instantiated != checkDeferredImport(decl.imp.import_index, deferred_imports)) continue;
						var r = imports[decl.imp.import_index];
						instance.memories[index] = importMemory(decl, r);
					} else if (!partially_instantiated) {
						var m = Target.newMemory(decl);
						if (m.oom) return fail("out of memory allocating memory");
						instance.memories[index] = m;
					}
				}
				Global(index) => {
					var decl = module.globals[index];
					if ((import_info = decl.imp) != null) {
						if (partially_instantiated != checkDeferredImport(decl.imp.import_index, deferred_imports)) continue;
						var r = imports[decl.imp.import_index];
						instance.globals[index] = importGlobal(decl, r);
					} else if (!partially_instantiated) {
						var g = Global.new(decl.valtype, decl);
						instance.globals[index] = g;
						g.set(instance.evalInitExpr(decl.init));
					}
				}
				Tag(index) => {
					var decl = module.tags[index];
					if ((import_info = decl.imp) != null) {
						if (partially_instantiated != checkDeferredImport(decl.imp.import_index, deferred_imports)) continue;
						var r = imports[decl.imp.import_index];
						instance.tags[index] = importTag(decl, r);
					} else if (!partially_instantiated) {
						instance.tags[index] = Tag.new(getSigDecl(decl.sig_index), decl);
					}
				}
				Continuation(index) => {
					// TODO-SS: implement
				}
			}
		}
		if (error.error()) return null;
		import_info = null;
		if (needSigIds) {
			for (i < instance.heaptypes.length) {
				match (instance.heaptypes[i]) {
					x: SigDecl => instance.sig_ids[i] = Canon.sigId(x);
					_ => instance.sig_ids[i] = -1;
				}
			}
		}

		// Initialize tables that have default element values.
		for (i < instance.tables.length) {
			var t = instance.tables[i];
			var decl = instance.module.tables[i];
			var is_deferred = if(t == null, true,
				decl.imp != null && checkDeferredImport(decl.imp.import_index, deferred_imports));
			// deferred imports are only loaded during finalization, regular imports are only
			// loaded on first pass.
			if (partially_instantiated != is_deferred || !t.decl.has_default_elem) continue;
			var val = instance.evalInitExpr(t.decl.default_elem);
			t.fill(0, val, u32.view(t.elems.length));
		}
		// Load element segments.
		for (i < module.elems.length) {
			var e = module.elems[i];
			match (e.mode) {
				Passive => ;
				Active(index, offset) => {
					var t = instance.tables[index];
					var decl = instance.module.tables[index];
					var is_deferred = if(t == null, true,
						decl.imp != null && checkDeferredImport(decl.imp.import_index, deferred_imports));
					// deferred imports are only loaded during finalization, regular imports are only
					// loaded on first pass.
					if (partially_instantiated != is_deferred) continue;
					var evals = instance.getElems(i);
					loadElems(index, offset, evals);
					instance.dropped_elems[i] = true;
					if (error.error()) return null;
				}
				Declarative => {
					instance.dropped_elems[i] = true;
				}
			}
		}
		// Load data segments.
		for (i < module.data.length) {
			var d = module.data[i];
			match (d.mode) {
				Passive => ;
				Active(index, offset) => {
					var mem = instance.memories[index];
					var decl = instance.module.memories[index];
					var is_deferred = if (mem == null, true,
						decl.imp != null && checkDeferredImport(decl.imp.import_index, deferred_imports));
					// deferred imports are only loaded during finalization, regular imports are only
					// loaded on first pass.
					if (partially_instantiated != is_deferred) continue;
					loadData(index, offset, d);
					instance.dropped_data[i] = true;
					if (error.error()) return null;
				}
				Declarative => {
					instance.dropped_data[i] = true;
				}
			}
		}

		// Organize exports
		var exports = instance.exports;
		for (i < module.exports.length) {
			exports[i] = getDecl(module.exports[i].1);
		}

		var has_deferred_imports = deferred_imports != null && deferred_imports.length > 0;
		if (!partially_instantiated) {
			partially_instantiated = true;
			fully_instantiated = !has_deferred_imports;
		} else {
			fully_instantiated = true;
		}

		if (fully_instantiated) {
			for (i < processors.length) {
				processors[i].postprocess(error, instance);
			}
			if (tiering != null) tiering.onInstantiateFinish(this, error);
		}
		return instance;
	}
	def importFunction(decl: FuncDecl, r: Exportable) -> Function {
		if (r == null) r = processImport(decl, ImportProcessor.processFunction(_, error, _, decl));
		if (binder != null) r = binder(decl, r);
		match (r) {
			func: Function => {
				var expected_sig = getSigDecl(decl.sig_index);
				if (func.sig != expected_sig) {
					var buf = StringBuilder.new().puts("expected function with signature ");
					if (expected_sig == null) buf.puts("(null)");
					else expected_sig.render(buf);
					buf.puts(", got ");
					func.sig.render(buf);
					if (func.sig.isAssignableSig(expected_sig)) {
						// structurally equivalent, but not allowed in Wasm's canonicalization--add context
						buf.put2(", structurally equivalent (canon #%d, #%d)", expected_sig.canon_id, func.sig.canon_id);
					}
					fail(buf.toString());
				}
				return func;
			}
			null => fail("function import not found");
			_ => fail("expected function import");
		}
		return null;
	}
	def importTable(decl: TableDecl, r: Exportable) -> Table {
		if (r == null) r = processImport(decl, ImportProcessor.processTable(_, error, _, decl));
		if (binder != null) r = binder(decl, r);
		match (r) {
			table: Table => {
				var f = table.decl.size;
				if (!checkSizeConstraint(f.is64, u64.!(table.elems.length), f.maximum, decl.size)) {
					fail("table limits mismatch");
				}
				var expected_type = decl.elemtype;
				var got_type = table.decl.elemtype;
				if (got_type != expected_type) {  // TODO: proper equality check
					fail("table element type mismatch");
				}
				return table;
			}
		}
		fail("expected table import");
		return null;
	}
	def importMemory(decl: MemoryDecl, r: Exportable) -> Memory {
		if (r == null) r = processImport(decl, ImportProcessor.processMemory(_, error, _, decl));
		if (binder != null) r = binder(decl, r);
		match (r) {
			memory: Memory => {
				var f = memory.decl.size;
				if (!checkSizeConstraint(f.is64, memory.num_pages, f.maximum, decl.size)) {
					fail("memory limits mismatch");
				}
				if (memory.decl.shared != decl.shared) {
					fail("memory sharing mismatch");
				}
				return memory;
			}
		}
		fail("expected memory import");
		return null;
	}
	def importTag(decl: TagDecl, r: Exportable) -> Tag {
		if (r == null) r = processImport(decl, ImportProcessor.processTag(_, error, _, decl));
		if (binder != null) r = binder(decl, r);
		match (r) {
			tag: Tag => {
				var expected_sig = getSigDecl(decl.sig_index);
				if (tag.sig != expected_sig) {
					var buf = StringBuilder.new().puts("expected tag with signature ");
					if (expected_sig == null) buf.puts("(null)");
					else expected_sig.render(buf);
					buf.puts(", got ");
					tag.sig.render(buf);
					if (tag.sig.isAssignableSig(expected_sig)) {
						// structurally equivalent, but not allowed in Wasm's canonicalization--add context
						buf.put2(", structurally equivalent (canon #%d, #%d)", expected_sig.canon_id, tag.sig.canon_id);
					}
					fail(buf.toString());
				}
				return tag;
			}
			null => fail("tag import not found");
			_ => fail("expected tag import");
		}
		return null;
	}
	def importGlobal(decl: GlobalDecl, r: Exportable) -> Global {
		if (r == null) r = processImport(decl, ImportProcessor.processGlobal(_, error, _, decl));
		if (binder != null) r = binder(decl, r);
		match (r) {
			global: Global => {
				if (decl.mutable != global.decl.mutable) {
					fail("global immutability mismatch");
				}
				var expected_type = decl.valtype;
				var got_type = global.valtype;
				if (decl.mutable) {
					if (got_type != expected_type) fail("global type mismatch");
				} else {
					if (!ValueTypes.isAssignable(got_type, expected_type)) 	fail("global type mismatch");
				}
				return global;
			}
		}
		fail("expected global import");
		return null;
	}
	def processImport(decl: Decl, f: (ImportProcessor, string) -> Exportable) -> Exportable {
		var i = decl.imp;
		import_pos = i.import_index;
		var modname = i.module_name, fieldname = i.field_name;
		for (i < processors.length) {
			var p = processors[i];
			if (Strings.equal(modname, p.name)) {
				var r = f(p, fieldname);
				if (r != null) {
					imports[decl.imp.import_index] = r;
					return r;
				}
			}
		}
		return null;
	}

	def getDecl(d: Decl) -> Exportable {
		match (d) {
			x: FuncDecl => return instance.functions[x.func_index];
			x: TableDecl => return instance.tables[x.table_index];
			x: MemoryDecl => return instance.memories[x.memory_index];
			x: GlobalDecl => return instance.globals[x.global_index];
			x: TagDecl => return instance.tags[x.tag_index];
		}
		return null;
	}
	def getSigDecl(sig_index: int) -> SigDecl {
		if (u32.view(sig_index) >= instance.heaptypes.length) return null;
		return SigDecl.!(instance.heaptypes[sig_index]);
	}

	// Substitution utilities.
	def mapFuncDecl(func_index: int) -> FuncDecl {
		var orig = module.functions[func_index];
		return orig;
	}

	private def checkDeferredImport(imp_index: int, deferred_imports: Array<int>) -> bool {
		if (deferred_imports == null || deferred_imports.length == 0) return false;
		for (i < deferred_imports.length) {
			if (deferred_imports[i] == imp_index) return true;
		}
		return false;
	}
	private def fail(msg: string) -> Instance {
		if (import_info != null) {
			var buf = StringBuilder.new()
				.put3("import #%d(\"%s\".\"%s\"): ", import_pos, import_info.module_name, import_info.field_name)
				.puts(msg);
			msg = buf.toString();
		}
		error.abs(import_pos).set(msg);
		return null;
	}
	private def checkSizeConstraint(is64: bool, size: u64, fmax: Max, t: SizeConstraint) -> bool {
		if (size < t.initial) return false;
		if (is64 != t.is64) return false;
		match (t.maximum) {
			None => return true;
			Set(tmax_val) => match (fmax) {
				None => return false;
				Set(fmax_val) => return fmax_val <= tmax_val;
			}
		}
	}
	private def loadElems(index: int, offset: InitExpr, evals: Array<Value>) {
		var dst_val = instance.evalInitExpr(offset);
		var dst_offset = if(Value.I64.?(dst_val), Value.I64.!(dst_val).val, Value.I32.!(dst_val).val);
		var table = instance.tables[index];
		var r = table.copyE(instance, dst_offset, evals, 0, u32.view(evals.length));
		if (r != TrapReason.NONE) {
			fail("out of bounds in initialization");
			trap_reason = r;
		}
	}
	private def loadData(index: int, offset: InitExpr, ddecl: DataDecl) {
		var dst_val = instance.evalInitExpr(offset);
		var dst_offset = if(Value.I64.?(dst_val), Value.I64.!(dst_val).val, Value.I32.!(dst_val).val);
		var memory = instance.memories[index];
		var r = memory.copyIn(dst_offset, if(ddecl != null, ddecl.data), 0, u32.view(ddecl.data.length));
		if (r != TrapReason.NONE) {
			fail("out of bounds in initialization");
			trap_reason = r;
		}
	}
}