add setting locals

4 files changed, 30 insertions, 4 deletions

diff --git a/com.c b/com.c
index a3cc5a8..66d3941 100644
--- a/com.c
+++ b/com.c
@@ -43,6 +43,7 @@ static int stack_effect_of(Op opcode) {
 	case OP_SKIP:
 	case OP_REDO:
 	case OP_SETGLOBAL:
+	case OP_SETLOCAL:
 	case OP_RET:
 	case OP_HALT:
 		return 0;
@@ -141,13 +142,25 @@ void single_form(Compiler *C, AstVec l, Op op) {
 	compile_opcode(C, op);
 }
 
+static Local *locate_local(Compiler *C, char *name);
+
 void set_form(Compiler *C, AstVec l, Op _) {
 	AstNode ident = l.vals[1];
 	CHECK(ident.ty == AST_IDENT, "set's first argument must be identifier");
-	ObjString *o = objstring_copy_cstr(C->S, ident.as.str);
-	compile_node(C, l.vals[2]);
-	compile_opcode(C, OP_SETGLOBAL);
-	compile_byte(C, compile_constant(C, VAL_OBJ(o)));
+	char *name = ident.as.str;
+
+	Local *loc = locate_local(C, name);
+	if (loc != NULL) {
+		compile_node(C, l.vals[2]);
+		compile_opcode(C, OP_SETLOCAL);
+		compile_byte(C, loc->slot);
+	} else {
+		// write global
+		ObjString *o = objstring_copy_cstr(C->S, name);
+		compile_node(C, l.vals[2]);
+		compile_opcode(C, OP_SETGLOBAL);
+		compile_byte(C, compile_constant(C, VAL_OBJ(o)));
+	}
 }
 
 void do_form(Compiler *C, AstVec l, Op _) {
diff --git a/dis.c b/dis.c
index 6827be0..193f410 100644
--- a/dis.c
+++ b/dis.c
@@ -44,6 +44,11 @@ static size_t disasm_instr_h(Chunk *ch, size_t ip, int depth) {
 		printf("getlocal #%d\n",ix);
 		break;
 	}
+	case OP_SETLOCAL: {
+		uint8_t ix = ch->bc.d[ip++];
+		printf("setlocal #%d\n",ix);
+		break;
+	}
 	case OP_SETGLOBAL: {
 		uint8_t ix = ch->bc.d[ip++];
 		printf("setglobal #%d\t; ",ix);
diff --git a/vm.c b/vm.c
index 39bf65d..9a69472 100644
--- a/vm.c
+++ b/vm.c
@@ -45,6 +45,7 @@ int runvm(State *S) {
 		if (S->do_trace) {
 			printf("\t[%lu + %lu] : ",th->fp, th->sp);
 			for (int i = 0; i < th->sp; i++) {
+				printf("(%d) ",i);
 				print_val(th->stack[th->fp + i]);
 				printf(" ; ");
 			}
@@ -124,6 +125,12 @@ int runvm(State *S) {
 			PUSH(th->stack[th->fp + lidx]);
 			break;
 		}
+		case OP_SETLOCAL: {
+			uint8_t lidx = RBYTE();
+			Val v = PEEK();
+			th->stack[th->fp + lidx] = v;
+			break;
+		}
 
 #define BINARY_OP(opcode, OP, RET_TYPE) \
 		case opcode: { \
diff --git a/vm.h b/vm.h
index 39fc0bb..d9fcb50 100644
--- a/vm.h
+++ b/vm.h
@@ -52,6 +52,7 @@ typedef enum {
 	OP_SETGLOBAL,
 
 	OP_GETLOCAL,
+	OP_SETLOCAL,
 
 	OP_TRUE,
 	OP_FALSE,