Whamm.v3

// Copyright 2024 Wizard authors. All rights reserved.
// See LICENSE for details of Apache 2.0 license.

// Represents a declared parameter to a Whamm probe call.
type WhammParam {
	case DynamicLoc;
	case FrameAccessor;
	case Pc;
	case Func;
	case Imm(orig: Token, i: u31);
	case Arg(orig: Token, i: u31);
	case Local(orig: Token, i: u31);
	case Call(target: Token, params: Array<WhammParam>);

	def render(buf: StringBuilder) -> StringBuilder {
		match (this) {
			DynamicLoc => return buf.puts("loc");
			FrameAccessor => return buf.puts("frame");
			Pc => return buf.puts("pc");
			Func => return buf.puts("fid");
			Imm(orig, i) => return buf.put1("imm%d", i);
			Arg(orig, i) => return buf.put1("arg%d", i);
			Local(orig, i) => return buf.put1("local%d", i);
			Call(target, params) => {
				buf.puts(target.image);
				buf.putc('(');
				Trace.renderCspRange(buf, params, WhammParam.render);
				buf.putc(')');
				return buf;
			}
		}
	}
}

// {WhammParam}s are converted to {WhammArg}s when they are bound at a match site. After a match
// is determined, some {WhammParam}s are turned into constants (e.g. func, pc), while others (e.g.
// local slot) have their types determined by the match location. Similarly, calls will be
// evaluated at the match site and replaced with their results.
type WhammArg {
	case FrameAccessor;
	case Val(v: Value);
	case Operand(t: ValueType, i: int);
	case Local(t: ValueType, i: int);
}

class WhammPredicate(call: WhammParam.Call) { }

// Utilities associated with Whamm probes.
component Whamm {
	def NO_PARAMS: Array<WhammParam> = [];

	// Parse a string representing an opcode match with optional parameters.
	def parseOpcodePattern(r: TextReader) -> (Opcode, Array<WhammParam>, WhammPredicate) {
		if (r.optN("wasm:") >= 0) {
			if (r.optN("opcode:") >= 0) {
				var op = Opcodes.parseName(r.data[r.pos ...]);
				if (op != Opcode.INVALID) {
					r.advance(op.mnemonic.length);
					r.skipWhitespace();
					var pred = if(r.char == '/', parsePredicate(r));
					r.skipWhitespace();
					var params = if(r.char == '(', parseParams(r), NO_PARAMS);
					return (op, params, pred);
				}
			}
		}
		return (Opcode.INVALID, NO_PARAMS, null);
	}
	// Parse a string representing Whamm parameters, enclosed in parentheses.
	// Any parse errors are left in the supplied {TextReader}.
	def parseParams(r: TextReader) -> Array<WhammParam> {
		var params = Vector<WhammParam>.new();
		r.req1('(');
		while (r.ok) {
			if (r.opt1(')') > 0) break;
			var p = parseParam0(r);
			if (!r.ok) break;
			params.put(p);
			if (r.opt1(')') > 0) break;
			r.req1(',');
		}
		return params.extract();
	}
	// Parse a string representation of a predicate, which is restricted to be
	// a call, thus "/$func(args)/".
	def parsePredicate(r: TextReader) -> WhammPredicate {
		r.req1('/');
		r.skipWhitespace();
		match (parseParam0(r)) {
			x: WhammParam.Call => {
				var pred = WhammPredicate.new(x);
				r.skipWhitespace();
				r.req1('/');
				return pred;
			}
			_ => {
				r.fail("expected call in predicate");
				return null;
			}
		}
	}
}

// A probe that adapts a Wasm function to be called by the engine-internal probing mechanism.
class WhammProbe(func: Function, sig: Array<WhammArg>) extends Probe {
	var trampoline: TargetCode;
	// properties set by the spc to make inlining optimization decisions.
	var inline_heuristic_checked = false;
	var spc_inline_func = false;
	var spc_swap_instance = false;
	var spc_swap_membase = false;

	private def args = if(sig.length == 0, Values.NONE, Array<Value>.new(sig.length));

	def fire(loc: DynamicLoc) -> Resumption {
		for (i < sig.length) {
			var v: Value;
			match (sig[i]) {
				FrameAccessor => v = Value.Ref(loc.frame.getFrameAccessor().getMetaRef());
				Val(val) => v = val;
				Operand(t, i) => v = loc.frame.getFrameAccessor().getOperand(i);
				Local(t, i) => v = loc.frame.getFrameAccessor().getLocal(i);
			}
			args[i] = v;
		}
		match (Execute.call(func, args)) {
			Throw => System.error("whamm", "probe threw exception");
			_ => ;
		}
		return Resumption.Continue;
	}
}

def parseParam0(r: TextReader) -> WhammParam {
	var i = r.star_rel(0, isAlphaOrUnderscore);
	var id = r.data[r.pos ... i];
	if (Ranges.equal("arg", id)) return parseUint(r, i, WhammParam.Arg);
	if (Ranges.equal("imm", id)) return parseUint(r, i, WhammParam.Imm);
	if (Ranges.equal("local", id)) return parseUint(r, i, WhammParam.Local);

	i = r.star_rel(i - r.pos, isIdentChar);

	if (i == r.pos) {
		r.fail("expected identifier");
		return WhammParam.DynamicLoc;
	}

	var token = r.readToken(i - r.pos);
	if (r.char == '(') {
		var params = Whamm.parseParams(r); // TODO: don't allow nested calls.
		return WhammParam.Call(token, params);
	}

	if (Strings.equal("pc", token.image)) return WhammParam.Pc;
	if (Strings.equal("fid", token.image)) return WhammParam.Func;
	if (Strings.equal("frame", token.image)) return WhammParam.FrameAccessor;

	r.setFirstError(token.beginLine, token.beginColumn, Strings.format1("unresolved identifier: \"%s\"", token.image));
	return WhammParam.DynamicLoc;
}
def parseUint<T>(r: TextReader, i: int, f: (Token, u31) -> T) -> T {
	var t = Ints.parsePosDecimal(r.data, i);
	if (t.0 > 0) {
		var token = r.readToken(i + t.0 - r.pos);
		return f(token, u31.!(t.1));
	}
	r.failRel(i, "expected positive integer");
	var d: T;
	return d;
}
def isAlphaOrUnderscore(ch: byte) -> bool {
	return (ch >= 'a' && ch <= 'z')
		|| (ch >= 'A' && ch <= 'Z')
		|| ch == '_';
}
def isIdentChar(ch: byte) -> bool {
	return (ch >= 'a' && ch <= 'z')
		|| (ch >= 'A' && ch <= 'Z')
		|| (ch >= '0' && ch <= '9')
		|| ch == '_' || ch == '$';
}

def debug(r: TextReader, where: string) {
	Trace.OUT.puts(where).ln();
	r.renderCurrentLineWithCaret(Trace.OUT, r.pos);
	Trace.OUT.ln();
}