keep track of stack position while compiling

6 files changed, 82 insertions, 27 deletions

diff --git a/chunk.h b/chunk.h
index ba01614..2dad595 100644
--- a/chunk.h
+++ b/chunk.h
@@ -19,7 +19,4 @@ struct _chunk {
 	} consts;
 };
 Chunk chunk_new(State *S);
-size_t compile_opcode(Compiler *C, uint8_t byte);
-size_t compile_constant(Compiler *C, Val v);
-
 #endif
diff --git a/com.c b/com.c
index 9a90a7f..22c811d 100644
--- a/com.c
+++ b/com.c
@@ -9,20 +9,8 @@
 #include "read.h"
 #include "util.h"
 
-#define BYTECODE(C) (C->ch->bc)
-
-static size_t placeholder(Compiler *C) {
-	size_t old_ix = BYTECODE(C).len;
-	compile_opcode(C, 0x00);
-	compile_opcode(C, 0x00);
-	return old_ix;
-}
-
-static void patch(Compiler *C, size_t addr, uint16_t val) {
-	BYTECODE(C).d[addr] = val & 0xff;
-	BYTECODE(C).d[addr+1] = (val & 0xff00) >> 8;
-}
 
+#define BYTECODE(C) (C->ch->bc)
 
 Chunk chunk_new(State *S) {
 	return (Chunk){ 0 };
@@ -32,11 +20,52 @@ Compiler compiler_new(Compiler *outer, Chunk *ch) {
 	return (Compiler){
 		.S = outer->S,
 		.ch = ch,
+		.stack_cur = 0,
 	};
 }
 
+static void compile_byte(Compiler *C, uint8_t byte);
+static void compile_opcode(Compiler *C, Op opcode);
+static size_t compile_constant(Compiler *C, Val v);
+
+static int stack_effect_of(Op opcode) {
+	switch (opcode) {
+	case OP_LOADK:
+	case OP_GETGLOBAL:
+	case OP_NIL:
+	case OP_TRUE:
+	case OP_FALSE:
+		return 1;
+	case OP_PUTS:
+	case OP_PRINT:
+	case OP_SKIP:
+	case OP_REDO:
+	case OP_SETGLOBAL:
+	case OP_RET:
+	case OP_HALT:
+		return 0;
+	case OP_DROP:
+	case OP_0BRANCH:
+	case OP_ADD:
+	case OP_SUB:
+	case OP_MUL:
+	case OP_DIV:
+	case OP_CMP:
+	case OP_EQU:
+	case OP_MOD:
+		return -1;
+
+	case OP_CALL:
+		// it depends. handle that one specifically
+		return 0;
+	default:
+		abort();
+	}
+}
+
 
-size_t compile_opcode(Compiler *C, uint8_t byte) {
+// compile that byte directly
+static void compile_byte(Compiler *C, uint8_t byte) {
 	Chunk *ch = C->ch;
 	if (ch->bc.len == ch->bc.cap) {
 		size_t newsz = (ch->bc.cap == 0 ? 8 : ch->bc.cap * 2);
@@ -46,9 +75,17 @@ size_t compile_opcode(Compiler *C, uint8_t byte) {
 	size_t ix = ch->bc.len;
 	ch->bc.d[ix] = byte;
 	ch->bc.len ++;
-	return ix;
 }
-size_t compile_constant(Compiler *C, Val v) {
+// compile an opcode, keeping track of its stack effect
+static void compile_opcode(Compiler *C, Op opcode) {
+	int stack_effect = stack_effect_of(opcode);
+	C->stack_cur += stack_effect;
+	compile_byte(C, opcode);
+}
+// add a new constant to the constant table, and return its index
+// (but return the index of any existing identical constant instead of
+//  inserting a duplicate)
+static size_t compile_constant(Compiler *C, Val v) {
 	Chunk *ch = C->ch;
 	for (int i = 0; i < ch->consts.len; i ++)
 		if (val_equal(v, ch->consts.d[i])) return i;
@@ -65,6 +102,24 @@ size_t compile_constant(Compiler *C, Val v) {
 	return ix;
 }
 
+static void compile_call_instr(Compiler *C, uint8_t len) {
+	compile_opcode(C, OP_CALL);
+	compile_byte(C, len);
+	C->stack_cur -= len;
+}
+
+static size_t placeholder(Compiler *C) {
+	size_t old_ix = BYTECODE(C).len;
+	compile_byte(C, 0x00);
+	compile_byte(C, 0x00);
+	return old_ix;
+}
+static void patch(Compiler *C, size_t addr, uint16_t val) {
+	BYTECODE(C).d[addr] = val & 0xff;
+	BYTECODE(C).d[addr+1] = (val & 0xff00) >> 8;
+}
+
+
 
 
 
@@ -91,7 +146,7 @@ void set_form(Compiler *C, AstVec l, Op _) {
 	ObjString *o = objstring_copy_cstr(C->S, ident.as.str);
 	compile_node(C, l.vals[2]);
 	compile_opcode(C, OP_SETGLOBAL);
-	compile_opcode(C, compile_constant(C, VAL_OBJ(o)));
+	compile_byte(C, compile_constant(C, VAL_OBJ(o)));
 }
 
 void do_form(Compiler *C, AstVec l, Op _) {
@@ -174,7 +229,7 @@ void fn_form(Compiler *C, AstVec l, Op op) {
 	compile_opcode(&subcompiler, OP_RET);
 
 	compile_opcode(C, OP_LOADK);
-	compile_opcode(C, compile_constant(C, VAL_OBJ(func)));
+	compile_byte(C, compile_constant(C, VAL_OBJ(func)));
 }
 
 
@@ -227,17 +282,17 @@ static void compile_node(Compiler *C, AstNode a) {
 			// read global variable
 			ObjString *o = objstring_copy_cstr(C->S, a.as.str);
 			compile_opcode(C, OP_GETGLOBAL);
-			compile_opcode(C, compile_constant(C, VAL_OBJ(o)));
+			compile_byte(C, compile_constant(C, VAL_OBJ(o)));
 		}
 		break;
 	case AST_NUM:
 		compile_opcode(C, OP_LOADK);
-		compile_opcode(C, compile_constant(C, VAL_NUM(a.as.num)));
+		compile_byte(C, compile_constant(C, VAL_NUM(a.as.num)));
 		break;
 	case AST_STRING: {
 		ObjString *o = objstring_copy_cstr(C->S, a.as.str);
 		compile_opcode(C, OP_LOADK);
-		compile_opcode(C, compile_constant(C, VAL_OBJ(o)));
+		compile_byte(C, compile_constant(C, VAL_OBJ(o)));
 		break;
 	}
 	case AST_LIST: {
@@ -277,8 +332,7 @@ static void compile_node(Compiler *C, AstNode a) {
 			for (int i = 0; i < l.len; i++) {
 				compile_node(C, l.vals[i]);
 			}
-			compile_opcode(C, OP_CALL);
-			compile_opcode(C, l.len);
+			compile_call_instr(C, l.len);
 		}
 
 		break;
diff --git a/com.h b/com.h
index f08769e..26b7494 100644
--- a/com.h
+++ b/com.h
@@ -10,6 +10,7 @@ typedef struct _compiler Compiler;
 struct _compiler {
 	State *S;
 	Chunk *ch;
+	int stack_cur;
 };
 Compiler compiler_new(Compiler *outer, Chunk *ch);
 
diff --git a/dis.c b/dis.c
index 0b22ac3..0b1d322 100644
--- a/dis.c
+++ b/dis.c
@@ -24,7 +24,7 @@ size_t disasm_instr(Chunk *ch, size_t ip) {
 static size_t disasm_instr_h(Chunk *ch, size_t ip, int depth) {
 	size_t orig_ip = ip;
 	uint8_t instr = ch->bc.d[ip];
-	printf("%*s%04zd\t",depth,"",ip);
+	printf("%*s%04zd\t%3d ",depth,"",ip,instr);
 	ip ++;
 	switch (instr) {
 	case OP_LOADK: {
diff --git a/tests/vars1.bth b/tests/vars1.bth
new file mode 100644
index 0000000..718b52b
--- /dev/null
+++ b/tests/vars1.bth
@@ -0,0 +1,2 @@
+(+ 2 (let (a 10)
+          (* a a)))
diff --git a/tests/vars1.out b/tests/vars1.out
new file mode 100644
index 0000000..257e563
--- /dev/null
+++ b/tests/vars1.out
@@ -0,0 +1 @@
+102