add local variables and (let) form

8 files changed, 145 insertions, 20 deletions

diff --git a/com.c b/com.c
index 22c811d..b2f1f09 100644
--- a/com.c
+++ b/com.c
@@ -32,6 +32,7 @@ static int stack_effect_of(Op opcode) {
 	switch (opcode) {
 	case OP_LOADK:
 	case OP_GETGLOBAL:
+	case OP_GETLOCAL:
 	case OP_NIL:
 	case OP_TRUE:
 	case OP_FALSE:
@@ -55,11 +56,12 @@ static int stack_effect_of(Op opcode) {
 	case OP_MOD:
 		return -1;
 
+	// these ones depend on their argument. handle them specifically
 	case OP_CALL:
-		// it depends. handle that one specifically
+	case OP_ENDSCOPE:
 		return 0;
 	default:
-		abort();
+		ERROR("unknown stack effect of opcode %d",opcode); 
 	}
 }
 
@@ -102,12 +104,19 @@ static size_t compile_constant(Compiler *C, Val v) {
 	return ix;
 }
 
+// len is 1 + number of args
 static void compile_call_instr(Compiler *C, uint8_t len) {
 	compile_opcode(C, OP_CALL);
 	compile_byte(C, len);
 	C->stack_cur -= len;
 }
 
+static void compile_endscope_instr(Compiler *C, uint8_t nlocals) {
+	compile_opcode(C, OP_ENDSCOPE);
+	compile_byte(C, nlocals);
+	C->stack_cur -= nlocals;
+}
+
 static size_t placeholder(Compiler *C) {
 	size_t old_ix = BYTECODE(C).len;
 	compile_byte(C, 0x00);
@@ -119,20 +128,11 @@ static void patch(Compiler *C, size_t addr, uint16_t val) {
 	BYTECODE(C).d[addr+1] = (val & 0xff00) >> 8;
 }
 
-
-
-
+// ----
 
 static void compile_node(Compiler *C, AstNode a);
 
 typedef void (*form_compiler)(Compiler *C, AstVec l, Op op);
-typedef struct {
-	char *name;
-	int min_params;
-	bool ellipsis;
-	form_compiler action;
-	Op op;
-} BuiltinForm;
 
 
 void single_form(Compiler *C, AstVec l, Op op) {
@@ -211,7 +211,7 @@ void arith_form(Compiler *C, AstVec l, Op op) {
 	compile_opcode(C, op);
 }
 
-void fn_form(Compiler *C, AstVec l, Op op) {
+void fn_form(Compiler *C, AstVec l, Op _) {
 	// (fn (arg arg arg) body ...)
 	CHECK(l.vals[1].ty == AST_LIST, "fn's first argument must be list");
 	AstVec arglist = l.vals[1].as.list;
@@ -232,8 +232,70 @@ void fn_form(Compiler *C, AstVec l, Op op) {
 	compile_byte(C, compile_constant(C, VAL_OBJ(func)));
 }
 
+static void begin_scope(Compiler *C) {
+	Scope *sc = malloc(sizeof(Scope));
+	CHECK(sc != NULL, "memory fail");
+	memset(sc, 0, sizeof(Scope));
+	sc->outer = C->scope;
+	C->scope = sc;
+}
+static void end_scope(Compiler *C) {
+	Scope *sc = C->scope;
+	CHECK(sc != NULL, "attempt to end nonexistent scope");
+	C->scope = sc->outer;
+	// printf("ending scope with %d locals, named: \n", sc->nlocals);
+	// for (int i = 0; i < sc->nlocals; i++) {
+	// 	Local loc = sc->locals[i];
+	// 	printf("\t%3d  %s\n",loc.slot, loc.name);
+	// }
+
+	compile_endscope_instr(C, sc->nlocals);
+
+	free(sc);
+}
+static void declare_local(Compiler *C, char *name) {
+	Scope *sc = C->scope;
+	CHECK(sc != NULL, "can't declare local outside of scope");
+	Local *l = &sc->locals[sc->nlocals++];
+	l->name = name;
+	// -1 because local is expected to be already on the stack
+	// ie sitting just below where stack_cur points
+	l->slot = C->stack_cur - 1;
+}
+
+void let_form(Compiler *C, AstVec l, Op _) {
+	CHECK(l.vals[1].ty == AST_LIST, "let's first argument must be list");
+	AstVec bindlist = l.vals[1].as.list;
+	CHECK(bindlist.len % 2 == 0, "unmatched binding in let");
+	int nbinds = bindlist.len / 2;
+
+	begin_scope(C);
+	for (int i = 0; i < nbinds; i++) {
+		int ix = i * 2;
+		AstNode name = bindlist.vals[ix];
+		AstNode expr = bindlist.vals[ix+1];
+		CHECK(name.ty == AST_IDENT, "binding name must be identifier");
+		compile_node(C, expr);
+		declare_local(C, name.as.str);
+	}
+	for (int i = 2; i < l.len - 1; i++) {
+		compile_node(C, l.vals[i]);
+		compile_opcode(C, OP_DROP);
+	}
+	compile_node(C, l.vals[l.len-1]);
+	end_scope(C);
+}
+
 
 
+typedef struct {
+	char *name;
+	int min_params;
+	bool ellipsis;
+	form_compiler action;
+	Op op;
+} BuiltinForm;
+
 static BuiltinForm builtin_forms[] = {
 	{ "puts",  1, false, single_form, OP_PUTS },
 	{ "print", 1, false, single_form, OP_PRINT },
@@ -242,6 +304,7 @@ static BuiltinForm builtin_forms[] = {
 	{ "if",    3, false, if_form, 0 },
 	{ "while", 2,  true, while_form, 0 },
 	{ "fn",    2,  true, fn_form, 0 },
+	{ "let",   2,  true, let_form, 0 },
 #define ARITH_OP(str, op) \
 	{ str, 2, false, arith_form, op },
 	ARITH_OP("+", OP_ADD)
@@ -279,10 +342,26 @@ static void compile_node(Compiler *C, AstNode a) {
 			}
 		}
 		if (!found_builtin) {
-			// read global variable
-			ObjString *o = objstring_copy_cstr(C->S, a.as.str);
-			compile_opcode(C, OP_GETGLOBAL);
-			compile_byte(C, compile_constant(C, VAL_OBJ(o)));
+			// read local or global variable
+			bool found_local = false;
+			if (C->scope != NULL) {
+				Scope *sc = C->scope;
+				for (int i = 0; i < sc->nlocals; i++) {
+					Local loc = sc->locals[i];
+					if (0 == strcmp(ident, loc.name)) {
+						compile_opcode(C, OP_GETLOCAL);
+						compile_byte(C, loc.slot);
+						found_local = true;
+						break;
+					}
+				}
+			}
+			if (!found_local) {
+				// read global
+				ObjString *o = objstring_copy_cstr(C->S, a.as.str);
+				compile_opcode(C, OP_GETGLOBAL);
+				compile_byte(C, compile_constant(C, VAL_OBJ(o)));
+			}
 		}
 		break;
 	case AST_NUM:
diff --git a/com.h b/com.h
index 26b7494..c623d49 100644
--- a/com.h
+++ b/com.h
@@ -7,10 +7,24 @@ typedef struct _compiler Compiler;
 #include "chunk.h"
 
 
+typedef struct {
+	char *name;
+	uint8_t slot;
+} Local;
+
+#define MAX_LOCALS 256
+typedef struct _scope Scope;
+struct _scope {
+	Scope *outer;
+	uint8_t nlocals;
+	Local locals[MAX_LOCALS];
+};
+
 struct _compiler {
 	State *S;
 	Chunk *ch;
 	int stack_cur;
+	Scope *scope;
 };
 Compiler compiler_new(Compiler *outer, Chunk *ch);
 
diff --git a/dis.c b/dis.c
index 0b1d322..6827be0 100644
--- a/dis.c
+++ b/dis.c
@@ -39,6 +39,11 @@ static size_t disasm_instr_h(Chunk *ch, size_t ip, int depth) {
 		print_const(ch, ix);
 		break;
 	}
+	case OP_GETLOCAL: {
+		uint8_t ix = ch->bc.d[ip++];
+		printf("getlocal #%d\n",ix);
+		break;
+	}
 	case OP_SETGLOBAL: {
 		uint8_t ix = ch->bc.d[ip++];
 		printf("setglobal #%d\t; ",ix);
@@ -50,6 +55,11 @@ static size_t disasm_instr_h(Chunk *ch, size_t ip, int depth) {
 		printf("call #%hhu\n",nargs);
 		break;
 	}
+	case OP_ENDSCOPE: {
+		uint8_t nlocals = ch->bc.d[ip++];
+		printf("endscope #%hhu\n",nlocals);
+		break;
+	}
 	#define RSHORT() (uint16_t)( ch->bc.d[ip-2] | ch->bc.d[ip-1] << 8 )
 	case OP_SKIP: {
 		ip += 2;
diff --git a/tests/vars1.bth b/tests/vars1.bth
index 718b52b..6e91bf0 100644
--- a/tests/vars1.bth
+++ b/tests/vars1.bth
@@ -1,2 +1,2 @@
 (+ 2 (let (a 10)
-          (* a a)))
+               (* a a)))
diff --git a/tests/vars2.bth b/tests/vars2.bth
new file mode 100644
index 0000000..3474d54
--- /dev/null
+++ b/tests/vars2.bth
@@ -0,0 +1,4 @@
+(+ 2 (* 6 (let (a 10
+				b 20
+				c 30)
+               (* b b))))
diff --git a/tests/vars2.out b/tests/vars2.out
new file mode 100644
index 0000000..65b7a7f
--- /dev/null
+++ b/tests/vars2.out
@@ -0,0 +1 @@
+2402
diff --git a/vm.c b/vm.c
index 4e412c1..b5b66d3 100644
--- a/vm.c
+++ b/vm.c
@@ -111,6 +111,12 @@ int runvm(State *S) {
 			break;
 		}
 
+		case OP_GETLOCAL: {
+			uint8_t lidx = RBYTE();
+			PUSH(th->stack[lidx]);
+			break;
+		}
+
 #define BINARY_OP(opcode, OP, RET_TYPE) \
 		case opcode: { \
 			Val b = POP(); \
@@ -148,8 +154,8 @@ int runvm(State *S) {
 				goto done;
 			}
 			PUSH(VAL_NUM(fmod(AS_NUM(a), AS_NUM(b))));
-			} 
 			break;
+		} 
 
 
 		case OP_NIL: PUSH(VAL_NIL); break;
@@ -173,8 +179,16 @@ int runvm(State *S) {
 
 			th->ch = &func->ch;
 			th->ip = 0;
+			break;
+		}
+
+		case OP_ENDSCOPE: {
+			uint8_t nlocals = RBYTE();
+			Val retval = POP();
+			th->sp -= nlocals;
+			PUSH(retval);
+			break;
 		}
-		break;
 
 		case OP_RET: {
 			StackFrame *sf = &th->rstack[--th->rsp]; 
diff --git a/vm.h b/vm.h
index 74fe506..0d1782d 100644
--- a/vm.h
+++ b/vm.h
@@ -50,6 +50,8 @@ typedef enum {
 	OP_GETGLOBAL,
 	OP_SETGLOBAL,
 
+	OP_GETLOCAL,
+
 	OP_TRUE,
 	OP_FALSE,
 	OP_NIL,
@@ -59,6 +61,7 @@ typedef enum {
 	OP_REDO,
 
 	OP_CALL,
+	OP_ENDSCOPE,
 } Op;
 
 int runvm(State *S);