VM: make the callstack more compact

mrunix00 · Sep 11, 2024 · d63e242 · d63e242
1 parent 2e043d7
commit d63e242
Show file tree

Hide file tree

Showing 7 changed files with 115 additions and 101 deletions.
diff --git a/examples/fib.hal b/examples/fib.hal
@@ -8,18 +8,18 @@ global_pointers: 0
     PrintTopStackI64;
     Exit;
 }
-:1 { args: 1 ptr_args: 0 locals: 0 local_pointers: 0 } {
-    LoadArgI64 $0;
+:1 { args: 1 ptr_args: 0 locals: 1 local_pointers: 0 } {
+    LoadLocalI64 $0;
     PushI64 2;
     LessThanI64;
     JumpIfFalse #6;
-    LoadArgI64 $0;
+    LoadLocalI64 $0;
     ReturnI64;
-    LoadArgI64 $0;
+    LoadLocalI64 $0;
     PushI64 1;
     SubI64;
     Call :1;
-    LoadArgI64 $0;
+    LoadLocalI64 $0;
     PushI64 2;
     SubI64;
     Call :1;

diff --git a/include/assembler/lexer.h b/include/assembler/lexer.h
@@ -12,7 +12,7 @@ typedef enum
 	TOKEN_PTR_ARGS,
 	TOKEN_LOCALS,
 	TOKEN_LOCAL_POINTERS,
-	TOKEN_LoadArgI64,
+	TOKEN_LoadLocalI64,
 	TOKEN_PushI64,
 	TOKEN_LessThanI64,
 	TOKEN_GreaterThanI64,

diff --git a/include/hal64.h b/include/hal64.h
@@ -6,7 +6,7 @@
 typedef enum
 {
 	OP_NOOP = 0,
-	OP_LOAD_ARG_I64,
+	OP_LOAD_LOCAL_I64,
 	OP_PUSH_I64,
 	OP_LESS_THAN_I64,
 	OP_GREATER_THAN_I64,
@@ -39,7 +39,9 @@ typedef struct
 	size_t ptr_args_count;
 	size_t locals_count;
 	size_t local_pointers_count;
+	size_t stack_depth;
 	size_t instructions_count;
+	size_t stack_frame_size;
 } Function;
 
 typedef struct
@@ -52,29 +54,13 @@ typedef struct
 
 typedef struct
 {
-	uint64_t *stack;
-	size_t stack_size;
-	size_t stack_pointer;
-} OperandsStack;
-
-typedef struct {
-	uint64_t *args;
+	uint64_t *call_stack;
+	uint64_t *operands_stack;
 	uint64_t *locals;
-	size_t return_function;
-	size_t return_instruction;
-} StackFrame;
-
-typedef struct
-{
-	StackFrame *stack;
-	size_t stack_size;
-	size_t stack_pointer;
-} CallStack;
-
-typedef struct
-{
-	OperandsStack stack;
-	CallStack call_stack;
+	size_t call_stack_size;
+	size_t call_stack_capacity;
+	size_t operands_stack_size;
+	size_t operands_stack_capacity;
 } VM;
 
 Program init_program(void);

diff --git a/src/assembler/assemble.c b/src/assembler/assemble.c
@@ -176,8 +176,8 @@ read_instruction(Instruction *instruction)
 	switch (token.type) {
 	case TOKEN_CLOSE_BRACE:
 		return HAL64_END_OF_BODY;
-	case TOKEN_LoadArgI64:
-		instruction->op = OP_LOAD_ARG_I64;
+	case TOKEN_LoadLocalI64:
+		instruction->op = OP_LOAD_LOCAL_I64;
 		token = read_index();
 		instruction->data = strtoll(token.value, NULL, 10);
 		if (instruction->data == LONG_MAX) {

diff --git a/src/hal64.c b/src/hal64.c
@@ -51,8 +51,8 @@ instruction_as_string(Instruction instruction, char *string, size_t max_length)
 	case OP_NOOP:
 		snprintf(string, max_length, "NOOP");
 		break;
-	case OP_LOAD_ARG_I64:
-		snprintf(string, max_length, "LOAD_ARG_I64 $%zu", instruction.data);
+	case OP_LOAD_LOCAL_I64:
+		snprintf(string, max_length, "LOAD_LOCAL_I64 $%zu", instruction.data);
 		break;
 	case OP_PUSH_I64:
 		snprintf(string, max_length, "PUSH_I64 %zu", instruction.data);
@@ -94,14 +94,16 @@ print_program(Program program)
 }
 
 Program
-init_program(void) {
+init_program(void)
+{
 	Program program;
 	memset(&program, 0, sizeof(Program));
 	return program;
 }
 
 Function
-init_function() {
+init_function()
+{
 	Function function;
 	memset(&function, 0, sizeof(Function));
 	return function;
@@ -110,6 +112,7 @@ init_function() {
 void
 emit_function(Program *program, Function function)
 {
+	uint64_t i, depth;
 	if (function.id >= program->functions_count) {
 		program->functions_count = function.id + 1;
 		program->functions = safe_realloc(program->functions, (program->functions_count + 1) * sizeof(Function));
@@ -120,7 +123,8 @@ emit_function(Program *program, Function function)
 void
 emit_instruction(Function *function, Instruction instruction)
 {
-	function->instructions = safe_realloc(function->instructions, (function->instructions_count + 1) * sizeof(Instruction));
+	function->instructions =
+		safe_realloc(function->instructions, (function->instructions_count + 1) * sizeof(Instruction));
 	function->instructions[function->instructions_count] = instruction;
 	function->instructions_count++;
 }
diff --git a/src/lexer.l b/src/lexer.l
@@ -22,7 +22,7 @@
 "ptr_args"              { return TOKEN_PTR_ARGS; }
 "locals"                { return TOKEN_LOCALS; }
 "local_pointers"        { return TOKEN_LOCAL_POINTERS; }
-"LoadArgI64"            { return TOKEN_LoadArgI64; }
+"LoadLocalI64"          { return TOKEN_LoadLocalI64; }
 "PushI64"               { return TOKEN_PushI64; }
 "LessThanI64"           { return TOKEN_LessThanI64; }
 "GreaterThanI64"        { return TOKEN_GreaterThanI64; }

diff --git a/src/vm.c b/src/vm.c
@@ -7,102 +7,120 @@ VM
 init_vm(void)
 {
 	VM vm;
-	vm.stack.stack_size = 1024;
-	vm.stack.stack = safe_malloc(vm.stack.stack_size * sizeof(uint64_t));
-	vm.stack.stack_pointer = 0;
-
-	vm.call_stack.stack_size = 1024;
-	vm.call_stack.stack = safe_malloc(vm.call_stack.stack_size * sizeof(StackFrame));
-	vm.call_stack.stack_pointer = 0;
+	vm.call_stack_size = 0;
+	vm.operands_stack_size = 0;
+	vm.call_stack_capacity = 1024;
+	vm.operands_stack_capacity = 1024;
+	vm.call_stack = safe_malloc(vm.call_stack_capacity * sizeof(uint64_t));
+	vm.operands_stack = safe_malloc(vm.operands_stack_capacity * sizeof(uint64_t));
 	return vm;
 }
 
 void
 free_vm(VM vm)
 {
-	free(vm.stack.stack);
-	free(vm.call_stack.stack);
+	free(vm.call_stack);
 }
 
 static void
 push_stack(VM *vm, uint64_t value)
 {
-	if (vm->stack.stack_pointer >= vm->stack.stack_size) {
-		vm->stack.stack_size *= 2;
-		vm->stack.stack = safe_realloc(vm->stack.stack, vm->stack.stack_size * sizeof(uint64_t));
-	}
-	vm->stack.stack[vm->stack.stack_pointer++] = value;
+	vm->operands_stack[vm->operands_stack_size++] = value;
 }
 
 static uint64_t
 pop_stack(VM *vm)
 {
-#ifdef HAL64_DEBUG
-	if (vm->stack.stack_pointer == 0) {
-		fprintf(stderr, "Stack underflow\n");
-		exit(EXIT_FAILURE);
-	}
-#endif
-	return vm->stack.stack[--vm->stack.stack_pointer];
+	return vm->operands_stack[--vm->operands_stack_size];
 }
 
 static uint64_t
 top_stack(VM *vm)
 {
-#ifdef HAL64_DEBUG
-	if (vm->stack.stack_pointer == 0) {
-		fprintf(stderr, "Stack underflow\n");
-		exit(EXIT_FAILURE);
-	}
-#endif
-	return vm->stack.stack[vm->stack.stack_pointer - 1];
+	return vm->operands_stack[vm->operands_stack_size - 1];
 }
 
-static StackFrame *
-top_stack_frame(VM *vm)
+static size_t
+get_stack_frame_size(VM *vm)
 {
-#ifdef HAL64_DEBUG
-	if (vm->call_stack.stack_pointer == 0) {
-		fprintf(stderr, "Stack underflow\n");
-		exit(EXIT_FAILURE);
-	}
-#endif
-	return vm->call_stack.stack + vm->call_stack.stack_pointer;
+	return vm->call_stack[vm->call_stack_size - 1];
 }
 
 static void
 pop_stack_frame(VM *vm)
 {
-#ifdef HAL64_DEBUG
-	if (vm->call_stack.stack_pointer == 0) {
-		fprintf(stderr, "Stack underflow\n");
-		exit(EXIT_FAILURE);
-	}
-#endif
-	free(top_stack_frame(vm)->locals);
-	free(top_stack_frame(vm)->args);
-	vm->call_stack.stack_pointer--;
+	vm->call_stack_size -= get_stack_frame_size(vm);
+	vm->locals = vm->call_stack + vm->call_stack_size - get_stack_frame_size(vm);
 }
 
 static void
 call_function(VM *vm, const Program *program, size_t current_function, size_t current_instruction, size_t next_function)
 {
-	StackFrame frame;
-	size_t i;
+	uint64_t i;
+	Function function = program->functions[next_function];
+
+	if (vm->call_stack_size + function.stack_frame_size >= vm->call_stack_capacity) {
+		vm->call_stack_capacity *= 2;
+		vm->call_stack = safe_realloc(vm->call_stack, vm->call_stack_capacity * sizeof(uint64_t));
+	}
+	if (vm->operands_stack_size + function.stack_depth >= vm->operands_stack_capacity) {
+		vm->operands_stack_capacity *= 2;
+		vm->operands_stack = safe_realloc(vm->operands_stack, vm->operands_stack_capacity * sizeof(uint64_t));
+	}
 
-	frame.return_function = current_function;
-	frame.return_instruction = current_instruction;
-	frame.locals = safe_malloc(program->functions[next_function].locals_count * sizeof(uint64_t));
-	frame.args = safe_malloc(program->functions[next_function].args_count * sizeof(uint64_t));
+	vm->locals = vm->call_stack + vm->call_stack_size;
+	vm->call_stack_size += function.stack_frame_size;
+	vm->call_stack[vm->call_stack_size - 1] = function.stack_frame_size;
+	vm->call_stack[vm->call_stack_size - 2] = current_instruction;
+	vm->call_stack[vm->call_stack_size - 3] = current_function;
 
-	for (i = program->functions[next_function].args_count - 1; i != -1; i--)
-		frame.args[i] = pop_stack(vm);
+	for (i = function.args_count - 1; i != -1; i--)
+		vm->locals[i] = pop_stack(vm);
+}
 
-	if (vm->call_stack.stack_pointer >= vm->call_stack.stack_size) {
-		vm->call_stack.stack_size *= 2;
-		vm->call_stack.stack = safe_realloc(vm->call_stack.stack, vm->call_stack.stack_size * sizeof(StackFrame));
+int64_t
+get_stack_depth(Program *program, Instruction inst)
+{
+	switch (inst.op) {
+	case OP_PUSH_I64:
+	case OP_LOAD_LOCAL_I64:
+		return 1;
+	case OP_ADD_I64:
+	case OP_SUB_I64:
+	case OP_MUL_I64:
+	case OP_DIV_I64:
+	case OP_MOD_I64:
+	case OP_LESS_THAN_I64:
+	case OP_GREATER_THAN_I64:
+	case OP_EQUALS_I64:
+	case OP_NOT_EQUALS_I64:
+	case OP_RETURN_I64:
+		return -1;
+	case OP_CALL:
+		return -program->functions[inst.data].args_count + 1;
+	default:
+		return 0;
+	}
+}
+
+static void
+prepare_program(Program *program)
+{
+	Function *function;
+	size_t i, j, depth;
+	for (i = 0; i < program->functions_count; i++) {
+		function = &program->functions[i];
+		function->stack_depth = 0;
+		for (j = 0; j < function->instructions_count; j++) {
+			depth = get_stack_depth(program, function->instructions[j]);
+			if (depth > 0)
+				function->stack_depth += depth;
+		}
+		function->stack_frame_size =
+			function->locals_count +
+				function->local_pointers_count
+				+ 3;
 	}
-	vm->call_stack.stack[++vm->call_stack.stack_pointer] = frame;
 }
 
 void
@@ -114,14 +132,19 @@ execute_program(Program program)
 	Instruction *instr;
 	uint64_t a, b;
 
-	vm.call_stack.stack[0].locals = safe_malloc(func->locals_count * sizeof(uint64_t));
+	prepare_program(&program);
+	vm.call_stack_size = func->stack_frame_size;
+	vm.locals = vm.call_stack;
+	vm.call_stack[vm.call_stack_size - 1] = vm.call_stack_size;
+	vm.call_stack[vm.call_stack_size - 2] = 0;
+	vm.call_stack[vm.call_stack_size - 3] = 0;
 	for (instr = func->instructions;; instr++) {
 		switch (instr->op) {
 		case OP_PUSH_I64:
 			push_stack(&vm, instr->data);
 			break;
-		case OP_LOAD_ARG_I64:
-			push_stack(&vm, top_stack_frame(&vm)->args[instr->data]);
+		case OP_LOAD_LOCAL_I64:
+			push_stack(&vm, vm.locals[instr->data]);
 			break;
 		case OP_ADD_I64:
 			push_stack(&vm, pop_stack(&vm) + pop_stack(&vm));
@@ -177,10 +200,11 @@ execute_program(Program program)
 			func = program.functions + instr->data;
 			instr = func->instructions - 1;
 			break;
-		case OP_RETURN_I64:
-			func = program.functions + top_stack_frame(&vm)->return_function;
-			instr = func->instructions + top_stack_frame(&vm)->return_instruction;
+		case OP_RETURN_I64: {
+			func = program.functions + vm.call_stack[vm.call_stack_size - 3];
+			instr = func->instructions + vm.call_stack[vm.call_stack_size - 2];
 			pop_stack_frame(&vm);
+		}
 			break;
 		default:
 			instruction_as_string(*instr, buff, 256);