path: root/ss/compile.cpp

#include "ms.h"
#include <algorithm>

namespace ms {

static inline const std::string &unwrap_id(const Value *id) {
	return *static_cast<const std::string *>(id->p1);
}

// Helpers for list; node is either a PAIR or NIL
#define IS_END(node) \
	(node->is_a(Value::NIL))
#define CHECK_LIST(node) { \
	if (!node->is_a(Value::PAIR)) {\
		compile_error("expected pair"); \
		return -1; \
	} \
}
#define CHECK_END(node) { \
	if (!IS_END(node)) { \
		compile_error("expected end of token"); \
		return -1; \
	} \
}
#define CHECK_NOTEND(node) { \
	CHECK_LIST(node); \
	if (IS_END(node)) { \
		compile_error("unexpected end of token"); \
		return -1; \
	} \
}
#define CHECK(node, t) { \
	CHECK_NOTEND(node); \
	if (!car(node)->is_a(t)) { \
		compile_error("unexpected token"); \
		return -1; \
	} \
}
#define BEGINS_WITH_KEYWORD(node, s) \
	node->is_a(Value::PAIR) && \
	(car(node)->is_a(Value::SYMBOL) && \
	    *static_cast<const std::string *>(car(node)->p1) == s)

#define CONSUME_KEYWORD(node, s) { \
	CHECK(node, Value::SYMBOL); \
	if (*static_cast<const std::string *>(car(node)->p1) != s) { \
		compile_error("unexpected id: expected " s); \
		return -1; \
	} \
	node = cdr(node); \
}
#define CONSUME(node) __extension__({ \
	CHECK_LIST(node); \
	const Value *ret = car(node); \
	node = cdr(node); \
	ret; \
})
#define CONSUME_ID(node) __extension__({ \
	CHECK(node, Value::SYMBOL); \
	CONSUME(node); \
})

class Compiler {

public:
	Compiler(const Compiler *parent) :
		parent(parent), variables(), seq(nullptr), label_next(0), num_args(0),
		seq_shield(get_memory()) {
		seq = static_cast<Assembly *>(seq_shield.new_object(Value::ASM));
		seq->p1 = new std::vector<Instruction>();
		seq->p2 = nullptr;
	}

	int feed_toplevel_sequence(const Value *node) {
		while (!IS_END(node)) {
			if (feed_toplevel(CONSUME(node)) < 0)
				return -1;
			if (!IS_END(node))
				emit(Instruction::POP);
		}
		return 0;
	}

	Assembly *result() {
		if (finish() < 0)
			return nullptr;
		return seq;
	}

private:

	int feed_toplevel(const Value *node) {
		if (BEGINS_WITH_KEYWORD(node, "define"))
			return feed_define(node);
		else if (BEGINS_WITH_KEYWORD(node, "define-cfunc"))
			return feed_define_cfunc(node);
		else if (BEGINS_WITH_KEYWORD(node, "load"))
			return -1; // Not supported
		else
			return feed_exp(node);
	}

	int feed_define(const Value *node) {
		CONSUME_KEYWORD(node, "define");

		const Value *id;
		if (car(node)->is_a(Value::PAIR)) {
			// (define (Id Id* [. Id]) Body)
			const Value *id_and_args = CONSUME(node);
			id = CONSUME_ID(id_and_args);
			emit_lambda(id_and_args, node);
		}
		else {
			// (define Id Exp)
			id = CONSUME_ID(node);
			if (feed_exp(CONSUME(node)) < 0)
				return -1;
			CHECK_END(node);
		}
		int n = bind_variable(id);
		// At toplevel, define behaves like set! if id is bound
		if (n < 0 && parent)
			return -1;
		emit(Instruction::SETVAR, id);
		emit(Instruction::PUSH, id);
		return 0;
	}

	int feed_define_cfunc(const Value *node) {
		// (define-cfunc Id)
		CONSUME_KEYWORD(node, "define-cfunc");
		const Value *id = CONSUME_ID(node);
		CHECK_END(node);

		int n = bind_variable(id);
		if (n < 0 && parent)
			return -1;
		emit(Instruction::PUSH, get_cfunc(unwrap_id(id)));
		emit(Instruction::SETVAR, id);
		emit(Instruction::PUSH, id);
		return 0;
	}

	int feed_exp(const Value *node) {
		switch (node->type) {
		case Value::INTEGER:
		case Value::BOOLEAN:
		case Value::STRING:
		case Value::NIL:
			emit(Instruction::PUSH, node);
			return 0;
		case Value::SYMBOL:
			// TODO: better way
			if (unwrap_id(node) == ":callcc") {
				emit(Instruction::PUSH, unshield().new_object(Value::CALLCC));
				return 0;
			}
			if (check_variable(node) < 0)
				return -1;
			emit(Instruction::GETVAR, node);
			return 0;
		default:
			if (node->type != Value::PAIR) {
				compile_error("expected list");
				return -1;
			}
			if (BEGINS_WITH_KEYWORD(node, "lambda"))
				return feed_lambda(node);
			else if (BEGINS_WITH_KEYWORD(node, "quote"))
				return feed_quote(node);
			else if (BEGINS_WITH_KEYWORD(node, "set!"))
				return feed_set(node);
			else if (BEGINS_WITH_KEYWORD(node, "let"))
				return feed_let(node);
			else if (BEGINS_WITH_KEYWORD(node, "let*"))
				return feed_letstar(node);
			else if (BEGINS_WITH_KEYWORD(node, "letrec"))
				return feed_letrec(node);
			else if (BEGINS_WITH_KEYWORD(node, "if"))
				return feed_if(node);
			else if (BEGINS_WITH_KEYWORD(node, "cond"))
				return feed_cond(node);
			else if (BEGINS_WITH_KEYWORD(node, "and"))
				return feed_and(node);
			else if (BEGINS_WITH_KEYWORD(node, "or"))
				return feed_or(node);
			else if (BEGINS_WITH_KEYWORD(node, "begin"))
				return feed_begin(node);
			else if (BEGINS_WITH_KEYWORD(node, "do"))
				return feed_do(node);
			else
				return feed_apply(node);
		}
	}

	void emit_lambda(const Value *args, const Value *body) {
		emit(Instruction::LAMBDA, unshield().new_pair(args, body));
	}

	int feed_lambda(const Value *node) {
		Memory::Shield shield(get_memory());

		CONSUME_KEYWORD(node, "lambda");
		const Value *args;
		CHECK_LIST(node);
		if (car(node)->is_a(Value::SYMBOL)) {
			// (lambda Id Body)
			args = shield.new_pair((Value *)CONSUME_ID(node), shield.new_nil());
		} else {
			// (lambda (Id* [Id . Id]) Body)
			// check car(node) is a pair or nil
			args = CONSUME(node);
		}
		emit_lambda(args, node);
		return 0;
	}

	int feed_body(const Value *node) {
		if (IS_END(node)) {
			emit(Instruction::PUSH, unshield().new_nil());
			return 0;
		}
		bool found_exp = false;
		while (!IS_END(node)) {
			const Value *item = CONSUME(node);
			if (BEGINS_WITH_KEYWORD(item, "define")) {
				if (found_exp) {
					compile_error("Define after Exp in Body");
					return -1;
				}
				if (feed_define(item) < 0)
					return -1;
			} else {
				found_exp = true;
				if (feed_exp(item) < 0)
					return -1;
			}
			if (!IS_END(node))
				emit(Instruction::POP);
		}
		return 0;
	}

	int feed_apply(const Value *node) {
		if (feed_exp(CONSUME(node)) < 0)
			return -1;

		uintptr_t args = 0;
		while (!node->is_a(Value::NIL)) {
			if (!node->is_a(Value::PAIR)) {
				compile_error("unexpected DOT");
				return -1;
			}
			args++;
			if (feed_exp(CONSUME(node)) < 0)
				return -1;
		}
		CHECK_END(node);
		emit(Instruction::APPLY, args);
		return 0;
	}

	int feed_quote(const Value *node) {
		CONSUME_KEYWORD(node, "quote");
		const Value *ret = CONSUME(node);
		CHECK_END(node);
		emit(Instruction::PUSH, ret);
		return 0;
	}

	int feed_set(const Value *node) {
		CONSUME_KEYWORD(node, "set!");
		const Value *id = CONSUME_ID(node);
		if (feed_exp(CONSUME(node)) < 0)
			return -1;
		CHECK_END(node);
		if (check_variable(id) < 0)
			return -1;
		emit(Instruction::DUP);
		emit(Instruction::SETVAR, id);
		return 0;
	}

	int parse_bindings(Memory::Shield &shield, const Value *bindings, Value **pargs, Value **pexprs) {
		Value *pair = const_cast<Value *>(CONSUME(bindings));
		Value *bid = const_cast<Value *>(CONSUME(pair)),
		      *expr = const_cast<Value *>(CONSUME(pair));
		CHECK_END(pair);

		Value *nil = shield.new_nil();
		Value *args = shield.new_pair(bid, nil),
		      *exprs = shield.new_pair(expr, nil),
		      *args_prev = args, *expr_prev = exprs;
		while (!IS_END(bindings)) {
			if (!bindings->is_a(Value::PAIR)) {
				compile_error("unexpected DOT in Bindings");
				return -1;
			}
			pair = const_cast<Value *>(CONSUME(bindings));
			bid = const_cast<Value *>(CONSUME(pair));
			expr = const_cast<Value *>(CONSUME(pair));
			CHECK_END(pair);
			args_prev = static_cast<Value *>(args_prev->p2 = shield.new_pair(bid, nil));
			expr_prev = static_cast<Value *>(expr_prev->p2 = shield.new_pair(expr, nil));
		}
		*pargs = args;
		*pexprs = exprs;
		return 0;
	}

	int feed_let(const Value *node) {
		Memory::Shield shield(get_memory());

		CONSUME_KEYWORD(node, "let");
		// (let ((a b) (c d)) body)
		//   => ((lambda (a c) body) b d)
		const Value *id = CONSUME(node),
		      *bindings;
		if (id->is_a(Value::SYMBOL)) {
			if (bind_variable(id) < 0)
				return -1;
			bindings = CONSUME(node);
		} else {
			bindings = id;
			id = nullptr;
		}

		if (IS_END(bindings)) {
			emit_lambda(bindings, node);
			if (id) {
				emit(Instruction::DUP);
				emit(Instruction::SETVAR, id);
			}
			emit(Instruction::APPLY, (uintptr_t)0);
			return 0;
		}

		Value *args, *exprs;
		if (parse_bindings(shield, bindings, &args, &exprs) < 0)
			return -1;
		emit_lambda(args, node);
		if (id) {
			emit(Instruction::DUP);
			emit(Instruction::SETVAR, id);
		}
		uintptr_t argc = 0;
		do {
			argc++;
			if (feed_exp((Value *)exprs->p1) < 0)
				return -1;
		} while ((exprs = (Value *)exprs->p2)->is_a(Value::PAIR));
		emit(Instruction::APPLY, argc);
		return 0;
	}

	int feed_letstar(const Value *node) {
		Memory::Shield shield(get_memory());

		CONSUME_KEYWORD(node, "let*");
		// (let* ((a b) (c d)) body)
		//   => ((lambda (a) ((lambda (c) body) d)) b)
		const Value *bindings = CONSUME(node),
		      *body = node;

		Value *nil = shield.new_nil();
		const Value *args = nil, *exprs = nil;
		while (!IS_END(bindings)) {
			if (!bindings->is_a(Value::PAIR)) {
				compile_error("unexpected DOT in Bindings");
				return -1;
			}
			const Value *pair = CONSUME(bindings);
			args = shield.new_pair(CONSUME(pair), args);
			exprs = shield.new_pair(CONSUME(pair), exprs);
			CHECK_END(pair);
		}

		const Value *lambda = shield.new_symbol("lambda");
		while (!IS_END(args)) {
			const Value *la = shield.new_pair(CONSUME(args), nil);
			const Value *lm = shield.new_pair(lambda, shield.new_pair(la, body));
			body = shield.new_pair(shield.new_pair(lm, shield.new_pair(CONSUME(exprs), nil)), nil);
		}

		CHECK_LIST(body);
		return feed_apply(car(body));
	}

	int feed_letrec(const Value *node) {
		Memory::Shield shield(get_memory());

		CONSUME_KEYWORD(node, "letrec");
		// (letrec ((a b) (c d)) body)
		//   => ((lambda () (define a b) (define c d) body))
		const Value *bindings = CONSUME(node),
		      *body = node;

		Value *nil = shield.new_nil();
		const Value *args = nil, *exprs = nil;
		while (!IS_END(bindings)) {
			if (!bindings->is_a(Value::PAIR)) {
				compile_error("unexpected DOT in Bindings");
				return -1;
			}
			const Value *pair = CONSUME(bindings);
			args = shield.new_pair(CONSUME(pair), args);
			exprs = shield.new_pair(CONSUME(pair), exprs);
			CHECK_END(pair);
		}

		const Value *define = shield.new_symbol("define");
		while (!IS_END(args)) {
			const Value *def = shield.new_pair(define,
						    shield.new_pair(CONSUME(args),
								    shield.new_pair(CONSUME(exprs), nil)));
			body = shield.new_pair(def, body);
		}

		emit_lambda(nil, body);
		emit(Instruction::APPLY);
		return 0;
	}

	int feed_if(const Value *node) {
		CONSUME_KEYWORD(node, "if");
		const Value *ccond = CONSUME(node);
		const Value *cthen = CONSUME(node);

		if (feed_exp(ccond) < 0)
			return -1;
		int label_unless = new_label(),
		    label_end = new_label();
		emit(Instruction::JUMPUNLESS, (uintptr_t)label_unless);
		if (feed_exp(cthen) < 0)
			return -1;
		emit(Instruction::JUMP, (uintptr_t)label_end);
		emit(Instruction::NOP, (uintptr_t)label_unless);
		if (!IS_END(node)) {
			const Value *celse = CONSUME(node);
			if (feed_exp(celse) < 0)
				return -1;
			CHECK_END(node);
		}
		emit(Instruction::NOP, (uintptr_t)label_end);
		return 0;
	}

	int feed_cond(const Value *node) {
		CONSUME_KEYWORD(node, "cond");
		CHECK_NOTEND(node);
		int label_next = new_label(),
		    label_end = new_label();
		while (!IS_END(node)) {
			emit(Instruction::NOP, (uintptr_t)label_next);

			const Value *cl = CONSUME(node);
			if (BEGINS_WITH_KEYWORD(cl, "else")) {
				CONSUME_KEYWORD(cl, "else");
				if (feed_sequence(cl) < 0)
					return -1;
			}
			else {
				if (feed_exp(CONSUME(cl)) < 0)
					return -1;
				label_next = new_label();
				emit(Instruction::JUMPUNLESS, (uintptr_t)label_next);
				if (feed_sequence(cl) < 0)
					return -1;
				emit(Instruction::JUMP, (uintptr_t)label_end);
			}
		}
		emit(Instruction::NOP, (uintptr_t)label_end);
		return 0;
	}

	int feed_and(const Value *node) {
		CONSUME_KEYWORD(node, "and");
		int label = new_label();
		while (!IS_END(node)) {
			if (feed_exp(CONSUME(node)) < 0)
				return -1;
			if (!IS_END(node))
				emit(Instruction::JUMPUNLESS, (uintptr_t)label);
		}
		emit(Instruction::NOP, (uintptr_t)label);
		return 0;
	}

	int feed_or(const Value *node) {
		CONSUME_KEYWORD(node, "or");
		int label = new_label();
		while (!IS_END(node)) {
			if (feed_exp(CONSUME(node)) < 0)
				return -1;
			if (!IS_END(node))
				emit(Instruction::JUMPIF, (uintptr_t)label);
		}
		emit(Instruction::NOP, (uintptr_t)label);
		return 0;
	}

	int feed_begin(const Value *node) {
		CONSUME_KEYWORD(node, "begin");
		return feed_sequence(node);
	}

	int feed_do(const Value *node) {
		Memory::Shield shield(get_memory());

		CONSUME_KEYWORD(node, "do");
		// (do
		//   ((Id1 Def1 Step1) (Id2 Def2 Step2))
		//   (Test Exp*)
		//   Body)
		// =>
		// (x=(lambda (:$ Id1 Id2)
		//   (if Test (begin Exp*) (begin Body (:$ Step1 Step2))))
		//   x Def1 Def2)
		const Value *vars = CONSUME(node),
		      *test_and_exps = CONSUME(node),
		      *test = CONSUME(test_and_exps),
		      *exps = test_and_exps,
		      *body = node;

		std::vector<const Value *> defs;
		const Value *nil = shield.new_nil(),
		      *dollar = shield.new_symbol(":$"),
		      *lambda_params = nil, *lambda_args = nil;
		while (!IS_END(vars)) {
			const Value *var = CONSUME(vars),
			      *id = CONSUME(var),
			      *def = CONSUME(var),
			      *step = id;
			if (!IS_END(var))
				step = CONSUME(var);
			CHECK_END(var);

			lambda_params = shield.new_pair(id, lambda_params);
			lambda_args = shield.new_pair(step, lambda_args);
			defs.push_back(def);
		}
		const Value *if_then =
			shield.new_pair(shield.new_symbol("begin"), exps);
		lambda_args = shield.new_pair(dollar, lambda_args);
		const Value *if_else =
			shield.new_pair(shield.new_symbol("begin"),
				 shield.new_pair(shield.new_pair(shield.new_symbol("begin"), body),
					  shield.new_pair(shield.new_pair(dollar, lambda_args), nil)));
		const Value *lambda_body =
			shield.new_pair(shield.new_symbol("if"),
				 shield.new_pair(test,
					  shield.new_pair(if_then,
						   shield.new_pair(if_else, nil))));

		lambda_params = shield.new_pair(dollar, lambda_params);
		// ((lambda ...) DUP ...defs)
		emit_lambda(lambda_params, shield.new_pair(lambda_body, nil));
		//dump(lambda_body);
		emit(Instruction::DUP);
		for (size_t i = 0; i < defs.size(); i++)
			if (feed_exp(defs[i]) < 0)
				return -1;
		emit(Instruction::APPLY, defs.size() + 1);

		return 0;
	}

	int feed_sequence(const Value *node) {
		Memory::Shield shield(get_memory());

		if (IS_END(node)) {
			emit(Instruction::PUSH, shield.new_nil());
			return 0;
		}
		while (!IS_END(node)) {
			const Value *q = CONSUME(node);
			if (feed_exp(q) < 0)
				return -1;
			if (!IS_END(node))
				emit(Instruction::POP);
		}
		return 0;
	}

	int feed_arguments(const Value *node) {
		while (!IS_END(node)) {
			if (node->is_a(Value::SYMBOL)) {
				// (arg . varargs)
				if (bind_variable(node) < 0)
					return -1;
				num_args = -num_args - 1;
				break;
			}
			CHECK(node, Value::SYMBOL);
			if (bind_variable(CONSUME_ID(node)) < 0)
				return -1;
			num_args++;
		}
		return 0;
	}

	void compile_error(const std::string &msg) const {
		std::cerr << "compile: " << msg << "\n";
	}

	void compile_error(const std::string &msg, const std::string &msg2) const {
		std::cerr << "compile: " << msg << msg2 << "\n";
	}

	int bind_variable(const Value *id) {
		const std::string &str = unwrap_id(id);
		const std::vector<std::string>::const_reverse_iterator it
			= std::find(variables.rbegin(), variables.rend(), str);
		if (it == variables.rend()) {
			int n = variables.size();
			variables.push_back(str);
			return n;
		}
		return -1;
	}

	int check_variable(const Value *id) const {
		const std::string &str = unwrap_id(id);
		const std::vector<std::string>::const_reverse_iterator it
			= std::find(variables.rbegin(), variables.rend(), str);
		if (it == variables.rend()) {
			if (!parent) {
				compile_error("undefined variable reference: ", str);
				return -1;
			}
			return parent->check_variable(id);
		}
		return 0;
	}

	void emit(Instruction::Type op, uintptr_t arg) {
		Instruction i(op, arg);
		seq->seq()->push_back(i);
	}

	void emit(Instruction::Type op, const Value *arg) {
		emit(op, (uintptr_t)arg);
		if (arg) // Conservative
			seq_shield.add((Value *)arg);
	}

	void emit(Instruction::Type op) {
		emit(op, (uintptr_t)0);
	}

	int new_label() {
		return label_next++;
	}

	int finish() {
		std::vector<size_t> labels(label_next);
		bool *tailcallable = (bool *)std::malloc(seq->seq()->size() * sizeof(bool));
		if (!tailcallable)
			return -1;
		std::fill(tailcallable, tailcallable + sizeof(bool) * seq->seq()->size(), false);
		int ret = -1;

		for (size_t pos = 0; pos < seq->seq()->size(); pos++) {
			Instruction *i = &seq->seq()->at(seq->seq()->size() - pos - 1);
			switch (i->type) {
			case Instruction::LAMBDA: {
				const Value *memo = (const Value *)i->arg;
				Compiler sub(this);
				if (sub.feed_arguments((const Value *)memo->p1) < 0)
					goto out;
				if (sub.feed_body((const Value *)memo->p2) < 0)
					goto out;
				Assembly *assm = sub.result();
				if (!assm)
					goto out;
				i->arg = (uintptr_t)assm;
				break;
			}
			case Instruction::NOP: {
				tailcallable[pos] = (pos == 0 || tailcallable[pos - 1]);

				labels[i->arg] = pos;
				break;
			}
			case Instruction::JUMP: {
				tailcallable[pos] = (pos == 0 || tailcallable[pos - labels[i->arg]]);
			} /* fall through */
			case Instruction::JUMPIF:
			case Instruction::JUMPUNLESS: {
				// NOP always comes after JUMP*
				i->arg = pos - labels[i->arg];
				break;
			}
			case Instruction::GETVAR:
			case Instruction::SETVAR: {
				const std::string &str = unwrap_id(reinterpret_cast<const Value *>(i->arg));
				uintptr_t offset = 0;
				const Compiler *cmpl = this;
				while (cmpl) {
					const std::vector<std::string>::const_iterator it
						= std::find(cmpl->variables.begin(), cmpl->variables.end(), str);
					if (it == cmpl->variables.end()) {
						offset += cmpl->variables.size();
						cmpl = cmpl->parent;
						continue;
					}
					i->arg = offset + std::distance(cmpl->variables.begin(), it);
					break;
				}
				break;
			}
			case Instruction::RET: {
				tailcallable[pos] = true;
				break;
			}
			case Instruction::APPLY: {
				if (pos == 0 || tailcallable[pos - 1])
					i->type = Instruction::APPLY_TAILCALL;
				break;
			}
			default:
				break;
			}
		}

		emit(Instruction::RET, (uintptr_t)0);
		seq->set_num_variables(variables.size(), num_args);
		ret = 0;
out:
		std::free(tailcallable);
		return ret;
	}

	const Value *toplevel;
	const Compiler *parent;
	std::vector<std::string> variables;
	Assembly *seq;
	int label_next;
	int num_args;
	Memory::Shield seq_shield;
};

Assembly *compile(const Value *node) {
	Compiler compiler(nullptr);
	if (compiler.feed_toplevel_sequence(node) < 0)
		return nullptr;
	return compiler.result();
}

}