Module.v3

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

// Representation of a Wasm module, including its various index spaces of declarations.
// Keeps the original order of declarations in order to perform polymorphic
// import binding.
class Module(filename: string) {
	var name: string;		// can be manually set
	def decls = Vector<Decl2>.new();
	def imports = Vector<Decl>.new(); // all imports
	def heaptypes = Vector<HeapTypeDecl>.new();
	def functions = Vector<FuncDecl>.new();
	def tables = Vector<TableDecl>.new();
	def memories = Vector<MemoryDecl>.new();
	def globals = Vector<GlobalDecl>.new();
	def tags = Vector<TagDecl>.new();		// exception-handling extension
	def exports = Vector<(string, Decl)>.new();
	def elems = Vector<ElemDecl>.new();
	def data = Vector<DataDecl>.new();
	def custom_sections = Vector<CustomSection>.new();
	var probes: Array<Array<Probe>>;
	var dyn_probes: Vector<(int, int, Probe)>;
	var outline: ModuleOutline;

	var start_function = -1;
	var explicit_data_count = -1;
	var target_module: TargetModule;
	var new_funcs: int;

	var names: NameSection;

	// Add a new declaration to this module. Automatically manages the various
	// index spaces and updates any space-specific indexes.
	def addDecl(d: Decl) {
		match (d) {
			x: SigDecl => {
				var index = heaptypes.length;
				decls.put(Decl2.Sig(x.heaptype_index = index));
				heaptypes.put(x);
			}
			x: StructDecl => {
				var index = heaptypes.length;
				decls.put(Decl2.Struct(x.heaptype_index = index));
				heaptypes.put(x);
			}
			x: ArrayDecl => {
				var index = heaptypes.length;
				decls.put(Decl2.Array(x.heaptype_index = index));
				heaptypes.put(x);
			}
			x: FuncDecl => {
				var index = functions.length;
				decls.put(Decl2.Func(x.func_index = index, x.sig_index));
				functions.put(x);
				if (u32.view(x.sig_index) < heaptypes.length) {
					var d = heaptypes[x.sig_index];
					if (SigDecl.?(d)) x.sig = SigDecl.!(d);
				}
			}
			x: TableDecl => {
				var index = tables.length;
				decls.put(Decl2.Table(x.table_index = index));
				tables.put(x);
			}
			x: MemoryDecl => {
				var index = memories.length;
				decls.put(Decl2.Memory(x.memory_index = index));
				memories.put(x);
			}
			x: GlobalDecl => {
				var index = globals.length;
				decls.put(Decl2.Global(x.global_index = index));
				globals.put(x);
			}
			x: TagDecl => {
				var index = tags.length;
				decls.put(Decl2.Tag(x.tag_index = index));
				tags.put(x);
				if (u32.view(x.sig_index) < heaptypes.length) {
					var d = heaptypes[x.sig_index];
					if (SigDecl.?(d)) x.fields = SigDecl.!(d).params;
				}
			}
			x: ContDecl => {
				var index = heaptypes.length;
				decls.put(Decl2.Continuation(index));
				heaptypes.put(x);
			}
		}
	}
	// Add a new import declaration to this module. Adds this declaration to {imports} as well.
	def addImport(module_name: string, field_name: string, d: Decl) {
		if (d == null) return;
		d.imp = ImportInfo.new(module_name, field_name, u31.!(imports.length));
		addDecl(d);
		imports.put(d);
	}
}

// A pair of a name and a module.
type NamedModule(path: string, name: string, module: Module);

// For imported quantities, the module name, field name, index, and args.
class ImportInfo(module_name: string, field_name: string, import_index: u31) {
}

// Superclass of all declared and importable/exportable declarations.
class Decl {
	var imp: ImportInfo;
	def imported() -> bool { return imp != null; }
}

def NO_EX_HANDLERS: Array<ExHandlerEntry> = [];
def NO_DEST_MAP: Array<int> = [];
def NO_DESTS: Array<TargetHandlerDest> = [];
def NO_PC_MAP: Array<(int, Array<int>)> = [];
class FuncHandlerInfo {
	var ex_handlers = NO_EX_HANDLERS;
	var suspend_handlers = NO_EX_HANDLERS;
	var switch_handlers = NO_EX_HANDLERS;
	var handler_dest_map: Array<int> = NO_DEST_MAP;
	var handler_dests: Array<TargetHandlerDest> = NO_DESTS;
	var handler_pc_index_map = NO_PC_MAP;

	def get_handler_dest(index: int) -> TargetHandlerDest {
		// XXX: flatten this lookup with an {Array.new(n_handlers)} where each elem is its target dest?
		return handler_dests[handler_dest_map[index]];
	}
}

// Function declaration, including signature and code.
class FuncDecl(sig_index: int) extends Decl {
	var sig: SigDecl;
	var func_index = -1;
	var reffed: bool;
	var entry_probed: bool;
	var num_locals: u16;
	def var cur_bytecode: Array<byte>;		// current (potentially instrumented) bytecode
	def var orig_bytecode: Array<byte>;		// unmodified original bytecode
	var sidetable = Sidetables.NO_SIDETABLE;	// sidetable, including control transfers
	var cbd_sidetable: Array<u8>;			// CBD u8 sidetable
	var frame_var_tags: Array<byte>;		// value tags for frame variables
	var target_code: TargetCode;
	var tierup_trigger: int = int.max;
	var handlers = FuncHandlerInfo.new();

	def render(names: NameSection, buf: StringBuilder) -> StringBuilder {
		var name = if (names != null, names.getFuncName(func_index));
		if (name != null) return buf.putsq(name);
		else return buf.put1("#%d", func_index);
	}
	def setOrigCode(code: Array<byte>) -> this {
		cur_bytecode = orig_bytecode = code;
		var tc: TargetCode;
		var tr: TargetCode;
		target_code = tc; // reset target code as well
		sidetable = Sidetables.NO_SIDETABLE;
		cbd_sidetable = null;
	}
	def activateProbingAt(pc: int, probe_byte: byte) {
		if (pc == 0) return void(entry_probed = true); // special case for function entry
		// "orig" will become a copy of the original code, to allow in-place modification of old code
		if (cur_bytecode == orig_bytecode) orig_bytecode = Arrays.dup(orig_bytecode);
		cur_bytecode[pc] = probe_byte;
	}
	def deactivateProbingAt(pc: int) {
		if (pc == 0) return void(entry_probed = false);
		if (cur_bytecode == orig_bytecode) return;
		cur_bytecode[pc] = orig_bytecode[pc];
	}
	def reset() -> this {
		if (cur_bytecode == orig_bytecode) return;
		ArrayUtil.copyInto(cur_bytecode, 0, orig_bytecode, 0, orig_bytecode.length);
		orig_bytecode = cur_bytecode;
	}
	def dup(new_sig: SigDecl) -> FuncDecl {
		var n = FuncDecl.new(this.sig_index);
		n.sig = new_sig;
		n.func_index = this.func_index;
		n.reffed = this.reffed;
		n.cur_bytecode = this.cur_bytecode;
		n.orig_bytecode = this.orig_bytecode;
		n.sidetable = this.sidetable;
		n.num_locals = this.num_locals;
		n.target_code = this.target_code;
		return n;
	}
	def findExHandler(instance: Instance, tag: Tag, throw_pc: int) -> ExHandler {
		return findHandler("findExHandler", handlers.ex_handlers, instance, tag, throw_pc);
	}
	def findSuspensionHandler(instance: Instance, tag: Tag, resume_pc: int) -> ExHandler {
		return findHandler("suspend", handlers.suspend_handlers, instance, tag, resume_pc);
	}
	def findSwitchHandler(instance: Instance, tag: Tag, resume_pc: int) -> ExHandler {
		return findHandler("switch", handlers.switch_handlers, instance, tag, resume_pc);
	}
	def findHandler(name: string, handlers: Array<ExHandlerEntry>, instance: Instance, tag: Tag, throw_pc: int) -> ExHandler {
		var i = 0;
		if (Trace.exception) {
			Trace.OUT.puts(name);
			Trace.OUT.put3("(func=%q, tag=%d, throw_pc=%d)",
			this.render(instance.module.names, _), tag.decl.tag_index, throw_pc).ln();
		}
		while (i < handlers.length) { // XXX: speed this up with a binary search
			var e = handlers[i];
			if (Trace.exception) Trace.OUT.put3("  entry[%d...%d] tag=%d", e.start, e.end, e.tag).ln();
			if (throw_pc < e.start || throw_pc >= e.end) {
				// not in range of the handler
				if (Trace.exception) Trace.OUT.put2("RANGE (%d, %d)\n", e.start, e.end).ln();
			} else if (e.tag == Modules.EX_TAG_DELEGATE) {
				// DELEGATE restarts the search as if it thrown at a different PC
				i = 0;
				throw_pc = ExHandlerInfo.Direct.!(e.info).handler_pc;
				if (Trace.exception) Trace.OUT.put1("  delegate: throw_pc=%d", throw_pc).ln();
				continue;
			} else if (e.tag == Modules.EX_TAG_ALL || instance.tags[e.tag] == tag) {
				match (e.info) {
					Direct(handler_pc, ex_slot, val_stack_top, sidetable_pos) => {
						if (Trace.exception) Trace.OUT.put3("  match (direct): handler_pc=%d, ex_slot=%d, vsp=%d",
							handler_pc, ex_slot, val_stack_top).ln();
						return ExHandler(e.index, handler_pc, false, ex_slot, val_stack_top, sidetable_pos);
					}
					Sidetable(push_exnref, sidetable_entry) => {
						var entry = sidetable.getCatchEntry(sidetable_entry);
						if (Trace.exception) {
							Trace.OUT.put2("  match (sidetable): push_exnref=%z, sidetable_entry=%d, ",
								push_exnref, sidetable_entry);
							Trace.OUT.put2("handler_pc=%d, vsp=%d", entry.handler_pc, entry.val_stack_top).ln();
						}
						return ExHandler(e.index, entry.handler_pc, push_exnref, -1, entry.val_stack_top, entry.sidetable_pos);
					}
				}
			}
			i++;
		}
		if (Trace.exception) Trace.OUT.puts("  no handler found").ln();
		return ExHandler(-1, -1, false, -1, -1, -1);
	}
	def num_slots() -> int {
		return if(frame_var_tags != null, frame_var_tags.length) + num_locals;
	}
}

// Both tables and elements have an index type, initial size, and optional maximum size.
type SizeConstraint(is64: bool, initial: u64, maximum: Max) #unboxed {
	def indexType() -> ValueType { return if(is64, ValueType.I64, ValueType.I32); }
}

// Table declaration, including element type, limits.
class TableDecl(elemtype: ValueType.Ref, size: SizeConstraint) extends Decl {
	var table_index = -1;
	var has_default_elem: bool;
	var default_elem: InitExpr;
}

// Memory declaration, including limits, shared attribute, and index type.
class MemoryDecl(size: SizeConstraint, shared: bool, log2_pageSize: u5) extends Decl {
	var memory_index = -1;
}

// Global variable declaration, including type and mutability.
class GlobalDecl(valtype: ValueType, mutable: bool, init: InitExpr) extends Decl {
	var global_index = -1;
}

// Tag declaration. (ext:exception-handling)
class TagDecl(sig_index: int) extends Decl {
	var fields: Array<ValueType>;
	var tag_index = -1;
}

// An element declaration that can occur in the elements section.
class ElemDecl(elemtype: ValueType, mode: SegmentMode, details: ElemDetails) {
	var elem_index: int;
}
type ElemDetails {
	case FuncRefs(vals: Array<int>);
	case Exprs(vals: Array<InitExpr>); /* ext:reference-types */
	def length() -> int {
		match (this) {
			FuncRefs(vals) => return vals.length;
			Exprs(vals) => return vals.length;
		}
	}
}
type SegmentMode {
	case Passive;
	case Active(index: int, offset: InitExpr);
	case Declarative;
}

// A data segment declaration that can occur in the data section.
class DataDecl(mode: SegmentMode, data: Array<byte>) {
}

// An uninterpreted custom section.
class CustomSection(name: string, payload: Array<byte>) {
}

// A special kind of expression that can be used in initialization.
type InitExpr {
	case I32(val: i32);
	case I64(val: i64);
	case F32(val: u32);
	case F64(val: u64);
	case V128(low: u64, high: u64);
	case FuncRefNull;
	case ExternRefNull;
	case RefNull(ht: HeapType);
	case Global(global_index: int, decl: GlobalDecl);
	case FuncRef(func_index: int, decl: FuncDecl);
	case Const(val: Value);
	case I31(val: InitExpr);
	case Array(t: HeapType.Array, len: InitExpr, elem: InitExpr);
	case FixedArray(t: HeapType.Array, vals: Array<InitExpr>);
	case Struct(t: HeapType.Struct, vals: Array<InitExpr>);
	case ArrayNewData(t: HeapType.Array, data_index: int, offset: InitExpr, len: InitExpr);
	case ArrayNewElem(t: HeapType.Array, elem_index: int, offset: InitExpr, len: InitExpr);
	case I32_ADD(a: InitExpr, b: InitExpr);
	case I32_SUB(a: InitExpr, b: InitExpr);
	case I32_MUL(a: InitExpr, b: InitExpr);
	case I64_ADD(a: InitExpr, b: InitExpr);
	case I64_SUB(a: InitExpr, b: InitExpr);
	case I64_MUL(a: InitExpr, b: InitExpr);
}

// Optional maximum for a table or memory.
type Max {
	case None;
	case Set(max: u64);

	def min(that: u64) -> u64 {
		match (this) {
			None => return that;
			Set(max) => return if(max < that, max, that);
		}
	}
	def check(val: u64) -> bool {
		match (this) {
			None => return true;
			Set(max) => return val <= max;
		}
	}
	def render(buf: StringBuilder) -> StringBuilder {
		match (this) {
			None => buf.puts("None");
			Set(max) => buf.putd(max);
		}
		return buf;
	}
}

// Globals associated with modules.
component Modules {
	def EX_TAG_ALL = -1;
	def EX_TAG_DELEGATE = -2;
}

type Decl2 {
	case Sig(index: int);
	case Struct(index: int);
	case Array(index: int);
	case Func(index: int, sig_index: int);
	case Table(index: int);
	case Memory(index: int);
	case Global(index: int);
	case Tag(index: int);
	case Continuation(index: int);
}

// Stores handler information for efficient runtime lookup.
type ExHandlerEntry(index: int, tag: int, start: int, end: int, info: ExHandlerInfo) #unboxed;
type ExHandlerInfo {
	case Direct(handler_pc: int, ex_slot: int, val_stack_top: int, sidetable_pos: int);
	case Sidetable(push_exnref: bool, sidetable_entry: int);
}
type ExHandler(index: int, handler_pc: int, push_exnref: bool, ex_slot: int, val_stack_top: int, sidetable_pos: int) #unboxed { }

// An outline stores the byte-offsets of various declarations in the original module bytes.
class ModuleOutline {
	var orig_bytes: Range<byte>;

	def imports = SectionOutline.new();
	def heaptypes = SectionOutline.new();
	def functions = SectionOutline.new();
	def tables = SectionOutline.new();
	def memories = SectionOutline.new();
	def globals = SectionOutline.new();
	def tags = SectionOutline.new();
	def exports = SectionOutline.new();
	def elements = SectionOutline.new();
	def data = SectionOutline.new();
	def custom_sections = Vector<(u64, u64)>.new();
}
// For sections that have repeated entries, this helper simplifies getting a range of
// a specific entry.
class SectionOutline {
	var start: u64;
	def offsets = Vector<int>.new(); // < 0 indicates an imported definition
	var end: u64;
	def range(orig: Range<byte>) -> Range<byte> {
		return orig[start ... end];
	}
	def get(orig: Range<byte>, i: int) -> Range<byte> {
		if (i < 0 || i >= offsets.length) System.error("BoundsCheckException", "");
		var start = offsets[i];
		if (start < 0) return null;  // indicates an imported definition
		if (i == offsets.length - 1) return orig[start ... end];
		else return orig[start ... offsets[i+1]];
	}
	def set(start: u64, size: u32) -> Vector<int> {
		this.start = start;
		this.end = start + size;
		return offsets;
	}
}