1 files changed, 202 insertions, 36 deletions

diff --git a/src/program.c b/src/program.c
index eef49b6..a0a2407 100644
--- a/src/program.c
+++ b/src/program.c
@@ -1,4 +1,5 @@
 #include "../include/program.h"
+
 #include "../include/std/mem.h"
 #include "../include/std/hash.h"
 #include "../include/std/alloc.h"
@@ -29,6 +30,15 @@ typedef struct {
     NumberUnion value;
 } Number;
 
+static int hash_map_compare_u64(const void *a, const void *b) {
+    return *(u64*)a == *(u64*)b;
+}
+
+static usize hash_u64(const u8 *data, usize size) {
+    (void)size;
+    return *(u64*)data;
+}
+
 int amal_program_init(amal_program *self) {
     ignore_result_int(buffer_init(&self->data, NULL));
     self->string_indices                = NULL;
@@ -38,16 +48,20 @@ int amal_program_init(amal_program *self) {
     self->extern_funcs_start            = NULL;
     self->exported_funcs                = NULL;
     self->exported_funcs_end            = NULL;
+    self->imports_start                 = NULL;
     self->read_index                    = 0;
-    self->main_func_instruction_offset  = ~0U;
+    self->main_func_instruction_offset  = ~(u32)0U;
     self->num_intermediates             = 0;
     self->num_strings                   = 0;
     self->num_functions                 = 0;
     self->num_extern_functions          = 0;
     self->num_exported_functions        = 0;
+    self->num_imports                   = 0;
+    self->return_value_index            = 0;
 
     cleanup_if_error(arena_allocator_init(&self->allocator));
     cleanup_if_error(hash_map_init(&self->extern_funcs_map, &self->allocator, sizeof(ProgramExternFunc), hash_map_compare_string, amal_hash_string));
+    cleanup_if_error(hash_map_init(&self->deferred_func_calls, &self->allocator, sizeof(Buffer), hash_map_compare_u64, hash_u64));
     cleanup_if_error(buffer_append_header(&self->data));
     return 0;
 
@@ -60,6 +74,8 @@ void amal_program_deinit(amal_program *self) {
     arena_allocator_deinit(&self->allocator);
     am_free(self->extern_func_indices);
     am_free(self->string_indices);
+    self->extern_func_indices = NULL;
+    self->string_indices = NULL;
     if(self->data.data)
         buffer_deinit(&self->data);
 }
@@ -75,34 +91,46 @@ int amal_program_add_extern_func(amal_program *self, BufferView name, void *func
     return hash_map_insert(&self->extern_funcs_map, name, &extern_func);
 }
 
-static CHECK_RESULT int amal_program_get_extern_func_by_index(amal_program *self, u16 index, ProgramExternFunc *result) {
-    u8 *extern_func_ptr;
-    u8 num_args;
-    u8 func_name_len;
-    BufferView func_name;
-
-    if(index >= self->num_extern_functions) {
-        result->func = NULL;
-        result->args_byte_size = 0;
-        amal_log_error("Extern func index index %ld is out of range (%ld)", index, self->num_extern_functions);
-        return AMAL_PROGRAM_INSTRUCTION_INVALID_EXTERN_FUNC_INDEX;
+static u8* amal_program_get_extern_funcs_start_by_import_index(amal_program *self, u8 import_index) {
+    BytecodeHeaderImport *header_import = (BytecodeHeaderImport*)self->imports_start;
+    header_import += import_index;
+    return (u8*)self->data.data + header_import->extern_function_index;
+}
+
+typedef struct {
+    u8 num_params;
+    BufferView name;
+} BytecodeHeaderExternFunction;
+
+/* TODO: Optimize this */
+static void amal_program_get_header_extern_function_by_index(amal_program *self, u8 import_index, u16 index, BytecodeHeaderExternFunction *result) {
+    u32 i;
+    u8 *extern_funcs_start = amal_program_get_extern_funcs_start_by_import_index(self, import_index);
+    extern_funcs_start += sizeof(u16) + sizeof(u32);
+    for(i = 0; i < (u32)index; ++i) {
+        u8 name_len = *(extern_funcs_start + sizeof(u8));
+        /* +1 for the null-terminated character */
+        extern_funcs_start += sizeof(u8) + sizeof(u8) + name_len + 1;
     }
+    result->num_params = extern_funcs_start[0];
+    result->name.size = extern_funcs_start[1];
+    result->name.data = (const char*)extern_funcs_start + sizeof(u8) + sizeof(u8);
+}
 
-    extern_func_ptr = self->extern_funcs_start + self->extern_func_indices[index];
 
-    am_memcpy(&num_args, extern_func_ptr, sizeof(num_args));
-    am_memcpy(&func_name_len, extern_func_ptr + sizeof(num_args), sizeof(func_name_len));
-    func_name.size = func_name_len;
-    func_name.data = (const char*)(extern_func_ptr + sizeof(num_args) + sizeof(func_name_len));
+static CHECK_RESULT int amal_program_get_extern_func_by_index(amal_program *self, u8 import_index, u16 index, ProgramExternFunc *result) {
+    BytecodeHeaderExternFunction extern_func;
+    amal_program_get_header_extern_function_by_index(self, import_index, index, &extern_func);
     
-    if(!hash_map_get(&self->extern_funcs_map, func_name, result)) {
-        amal_log_error("No such extern function: %.*s", func_name.size, func_name.data);
+    if(!hash_map_get(&self->extern_funcs_map, extern_func.name, result)) {
+        amal_log_error("No such extern function: %.*s", extern_func.name.size, extern_func.name.data);
         return AMAL_PROGRAM_NO_SUCH_EXTERNAL_FUNCTION;
     }
     
     /* TODO: This assumes all arguments are of size sizeof(isize) */
-    if(result->args_byte_size != -1 && result->args_byte_size != num_args * (int)sizeof(isize)) {
-        amal_log_error("Extern function %.*s was registered to take %d byte(s), but the program says it takes %d byte(s)", func_name.size, func_name.data, result->args_byte_size, num_args * sizeof(isize));
+    if(result->args_byte_size != -1 && result->args_byte_size != extern_func.num_params * (int)sizeof(isize)) {
+        amal_log_error("Extern function %.*s was registered to take %d byte(s), but the program says it takes %d byte(s)",
+            extern_func.name.size, extern_func.name.data, result->args_byte_size, extern_func.num_params * sizeof(isize));
         return AMAL_PROGRAM_NO_SUCH_EXTERNAL_FUNCTION;
     }
     return 0;
@@ -125,7 +153,7 @@ static CHECK_RESULT int amal_program_set_exported_function_instruction_offset_ad
         num_args = self->exported_funcs[sizeof(instruction_offset)];
         func_name_size = self->exported_funcs[sizeof(instruction_offset) + sizeof(num_args)];
         self->exported_funcs += sizeof(instruction_offset) + sizeof(num_args) + sizeof(func_name_size);
-        if(self->main_func_instruction_offset == ~0U && func_name_size == 4 && am_memeql(self->exported_funcs, "main", 4))
+        if(self->main_func_instruction_offset == ~(u32)0U && func_name_size == 4 && am_memeql(self->exported_funcs, "main", 4))
             self->main_func_instruction_offset = instruction_offset;
 
         /* +1 to skip null-termination character */
@@ -191,10 +219,25 @@ static CHECK_RESULT bool amal_program_read_advance(amal_program *self, void *out
     return bool_true;
 }
 
+static CHECK_RESULT int amal_program_advance_section_magic_number(amal_program *self) {
+    u32 magic_number;
+    if(bytes_left_to_read(self) < sizeof(u32))
+        return AMAL_PROGRAM_SECTION_ERROR;
+    
+    am_memcpy(&magic_number, self->data.data + self->read_index, sizeof(magic_number));
+    if(magic_number != AMAL_BYTECODE_SECTION_MAGIC_NUMBER)
+        return AMAL_PROGRAM_SECTION_ERROR;
+    
+    self->read_index += sizeof(magic_number);
+    return 0;
+}
+
 static CHECK_RESULT int amal_program_read_intermediates(amal_program *self) {
     u32 intermediates_size;
     /*u32 read_end;*/
 
+    return_if_error(amal_program_advance_section_magic_number(self));
+
     if(bytes_left_to_read(self) < sizeof(intermediates_size)) {
         amal_log_error("Not enough space in program to intermediates size");
         return AMAL_PROGRAM_INVALID_INTERMEDIATES;
@@ -219,6 +262,8 @@ static CHECK_RESULT int amal_program_read_strings(amal_program *self) {
     u32 strings_size;
     u32 *string_index_ptr;
 
+    return_if_error(amal_program_advance_section_magic_number(self));
+
     if(!amal_program_read_advance(self, &self->num_strings, sizeof(u16)))
         return AMAL_PROGRAM_INVALID_STRINGS;
 
@@ -262,8 +307,18 @@ static CHECK_RESULT int amal_program_read_strings(amal_program *self) {
 }
 
 static CHECK_RESULT int amal_program_read_functions(amal_program *self) {
+    u32 funcs_size;
+
+    return_if_error(amal_program_advance_section_magic_number(self));
+
     if(!amal_program_read_advance(self, &self->num_functions, sizeof(u16)))
         return AMAL_PROGRAM_INVALID_FUNCTIONS;
+
+    if(!amal_program_read_advance(self, &funcs_size, sizeof(funcs_size)) || bytes_left_to_read(self) < funcs_size)
+        return AMAL_PROGRAM_INVALID_FUNCTIONS;
+
+    self->funcs_start = (u8*)(self->data.data + self->read_index);
+    self->read_index += funcs_size;
     return AMAL_PROGRAM_OK;
 }
 
@@ -271,6 +326,8 @@ static CHECK_RESULT int amal_program_read_external_functions(amal_program *self)
     u32 extern_funcs_size;
     u32 *extern_func_index_ptr;
 
+    return_if_error(amal_program_advance_section_magic_number(self));
+
     if(!amal_program_read_advance(self, &self->num_extern_functions, sizeof(u16)))
         return AMAL_PROGRAM_INVALID_EXTERNAL_FUNCTIONS;
 
@@ -291,17 +348,17 @@ static CHECK_RESULT int amal_program_read_external_functions(amal_program *self)
         const u32 read_end = read_start + extern_funcs_size;
         self->extern_funcs_start = (u8*)(self->data.data + self->read_index);
         while(self->read_index < read_end) {
-            u8 num_args;
+            u8 num_params;
             u8 func_name_size;
 
-            if(bytes_left_to_read(self) < sizeof(num_args) + sizeof(func_name_size))
+            if(bytes_left_to_read(self) < sizeof(num_params) + sizeof(func_name_size))
                 return AMAL_PROGRAM_INVALID_EXTERNAL_FUNCTIONS;
 
             *extern_func_index_ptr = self->read_index - read_start;
             ++extern_func_index_ptr;
-            num_args = self->data.data[self->read_index];
-            func_name_size = self->data.data[self->read_index + sizeof(num_args)];
-            self->read_index += sizeof(num_args) + sizeof(func_name_size);
+            num_params = self->data.data[self->read_index];
+            func_name_size = self->data.data[self->read_index + sizeof(num_params)];
+            self->read_index += sizeof(num_params) + sizeof(func_name_size);
 
             /* +1 to skip null-termination character */
             if(bytes_left_to_read(self) < func_name_size + 1U)
@@ -318,6 +375,8 @@ static CHECK_RESULT int amal_program_read_external_functions(amal_program *self)
 static CHECK_RESULT int amal_program_read_exported_functions(amal_program *self) {
     u32 export_funcs_size;
 
+    return_if_error(amal_program_advance_section_magic_number(self));
+
     if(!amal_program_read_advance(self, &self->num_exported_functions, sizeof(u16)))
         return AMAL_PROGRAM_INVALID_EXPORTED_FUNCTIONS;
 
@@ -337,6 +396,22 @@ static CHECK_RESULT int amal_program_read_exported_functions(amal_program *self)
     return AMAL_PROGRAM_OK;
 }
 
+static CHECK_RESULT int amal_program_read_imports(amal_program *self) {
+    u32 imports_size;
+
+    return_if_error(amal_program_advance_section_magic_number(self));
+
+    if(!amal_program_read_advance(self, &self->num_imports, sizeof(u8)))
+        return AMAL_PROGRAM_INVALID_IMPORTS;
+
+    if(!amal_program_read_advance(self, &imports_size, sizeof(imports_size)) || bytes_left_to_read(self) < imports_size)
+        return AMAL_PROGRAM_INVALID_IMPORTS;
+
+    self->imports_start = (u8*)(self->data.data + self->read_index);
+    self->read_index += imports_size;
+    return AMAL_PROGRAM_OK;
+}
+
 static CHECK_RESULT int amal_program_get_intermediate_by_index(amal_program *self, u16 index, Number *result) {
     if(index >= self->num_intermediates)
         return AMAL_PROGRAM_INSTRUCTION_INVALID_INTERMEDIATE_INDEX;
@@ -360,6 +435,48 @@ static CHECK_RESULT int amal_program_get_data_by_index(amal_program *self, u16 i
     return 0;
 }
 
+static u8* amal_program_get_funcs_start_by_import_index(amal_program *self, u8 import_index) {
+    BytecodeHeaderImport *header_import = (BytecodeHeaderImport*)self->imports_start;
+    header_import += import_index;
+    return (u8*)self->data.data + header_import->function_index;
+}
+
+static BytecodeHeaderFunction* amal_program_get_header_function_by_index(amal_program *self, u8 import_index, u16 index) {
+    u8 *funcs_start = amal_program_get_funcs_start_by_import_index(self, import_index);
+    BytecodeHeaderFunction *header_func = (BytecodeHeaderFunction*)(funcs_start + sizeof(u16) + sizeof(u32));
+    return header_func + index;
+}
+
+static u64 deferred_func_call_get_key(amal_program *self, u8 import_index, u16 func_index) {
+    BytecodeHeaderImport *header_import = (BytecodeHeaderImport*)self->imports_start;
+    header_import += import_index;
+    return ((u64)func_index << 32) | (u64)header_import->function_index;
+}
+
+static CHECK_RESULT int resolve_deferred_func_calls(amal_program *self, amal_executor *executor, u16 func_index) {
+    u64 key = deferred_func_call_get_key(self, 0, func_index);
+    BufferView key_mem = create_buffer_view((char*)&key, sizeof(key));
+    u32 current_code_offset = amal_exec_get_code_offset(executor);
+
+    Buffer* deferred_func_call_list;
+    if(!hash_map_get_ref(&self->deferred_func_calls, key_mem, (void**)&deferred_func_call_list))
+        return 0;
+    
+    {
+        u32 *code_offset = buffer_begin(deferred_func_call_list);
+        u32 *code_offset_end = buffer_end(deferred_func_call_list);
+        for(; code_offset != code_offset_end; ++code_offset) {
+            amal_exec_call_overwrite(executor, *code_offset, current_code_offset - *code_offset);
+        }
+        return buffer_set_capacity(deferred_func_call_list, 0);
+    }
+}
+
+static void header_func_set_offset(amal_program *self, u16 func_index, u32 code_offset) {
+    BytecodeHeaderFunction *header_func = ((BytecodeHeaderFunction*)self->funcs_start) + func_index;
+    header_func->func_offset = code_offset;
+}
+
 static CHECK_RESULT int amal_program_read_instructions(amal_program *self, amal_executor *executor) {
     u32 instructions_size;
     u32 read_start;
@@ -370,6 +487,9 @@ static CHECK_RESULT int amal_program_read_instructions(amal_program *self, amal_
     inside_func = bool_false;
     (void)inside_func;
     func_counter = 0;
+    self->return_value_index = 0;
+
+    return_if_error(amal_program_advance_section_magic_number(self));
 
     if(!amal_program_read_advance(self, &instructions_size, sizeof(instructions_size)))
         return AMAL_PROGRAM_INVALID_INSTRUCTIONS_SIZE;
@@ -484,14 +604,51 @@ static CHECK_RESULT int amal_program_read_instructions(amal_program *self, amal_
                 self->read_index += 2;
                 break;
             }
+            case AMAL_OP_PUSH_RET: {
+                /* TODO: Validate return value index doesn't go out of bounds? */
+                self->return_values_stack[self->return_value_index++] = self->data.data[self->read_index];
+                self->read_index += 1;
+                break;
+            }
             case AMAL_OP_CALL: {
+                u8 import_index;
                 u16 func_index;
                 u8 num_args;
+                BytecodeHeaderFunction *func_def;
                 i8 dst_reg;
-                am_memcpy(&func_index, self->data.data + self->read_index, sizeof(func_index));
-                num_args = self->data.data[self->read_index + sizeof(func_index)];
-                dst_reg = self->data.data[self->read_index + sizeof(func_index) + sizeof(num_args)];
-                return_if_error(amal_exec_call(executor, func_index, num_args, dst_reg));
+
+                am_memcpy(&import_index, self->data.data + self->read_index, sizeof(import_index));
+                am_memcpy(&func_index, self->data.data + self->read_index + sizeof(import_index), sizeof(func_index));
+                am_memcpy(&num_args, self->data.data + self->read_index + sizeof(import_index) + sizeof(func_index), sizeof(num_args));
+                func_def = amal_program_get_header_function_by_index(self, import_index, func_index);
+                assert(func_def->num_return_types == 1 && "TODO: Support 0 and more than 1 return values");
+                assert(self->return_value_index == 1);
+                dst_reg = self->return_values_stack[0];
+                self->return_value_index -= func_def->num_return_types;
+
+                if((char*)func_def < self->data.data + self->read_index) {
+                    return_if_error(amal_exec_call(executor, func_def->func_offset, num_args, dst_reg));
+                } else {
+                    /*
+                        The code for the function has not been generated yet (the function is defined after the current location).
+                        Make a dummy call and replace the call target after the function has been generated
+                    */
+                    u64 key = deferred_func_call_get_key(self, import_index, func_index);
+                    BufferView key_mem = create_buffer_view((char*)&key, sizeof(key));
+                    u32 code_offset = amal_exec_get_code_offset(executor);
+
+                    Buffer* deferred_func_call_list;
+                    if(hash_map_get_ref(&self->deferred_func_calls, key_mem, (void**)&deferred_func_call_list))
+                        return_if_error(buffer_append(deferred_func_call_list, &code_offset, sizeof(code_offset)));
+                    else {
+                        Buffer new_deferred_call_list;
+                        return_if_error(buffer_init(&new_deferred_call_list, &self->allocator));
+                        return_if_error(buffer_append(&new_deferred_call_list, &code_offset, sizeof(code_offset)));
+                        return_if_error(hash_map_insert(&self->deferred_func_calls, key_mem, &new_deferred_call_list));
+                    }
+                    /* Dummy call to offset 0, offset will be replace later when the target function hits AMAL_OP_FUNC_START */
+                    return_if_error(amal_exec_call(executor, 0, num_args, dst_reg));
+                }
                 self->read_index += 4;
                 break;
             }
@@ -500,16 +657,21 @@ static CHECK_RESULT int amal_program_read_instructions(amal_program *self, amal_
                 self->read_index += 2;
                 break;
             case AMAL_OP_CALLE: {
+                u8 import_index;
                 u16 extern_func_index;
                 u8 num_args;
                 i8 dst_reg;
-                am_memcpy(&extern_func_index, self->data.data + self->read_index, sizeof(extern_func_index));
-                num_args = self->data.data[self->read_index + sizeof(extern_func_index)];
-                dst_reg = self->data.data[self->read_index + sizeof(extern_func_index) + sizeof(num_args)];
+
+                am_memcpy(&import_index, self->data.data + self->read_index, sizeof(import_index));
+                am_memcpy(&extern_func_index, self->data.data + self->read_index + sizeof(import_index), sizeof(extern_func_index));
+                am_memcpy(&num_args, self->data.data + self->read_index + sizeof(import_index) + sizeof(extern_func_index), sizeof(num_args));
+                assert(self->return_value_index == 1 && "TODO: Support extern functions that don't return any value");
+                dst_reg = self->return_values_stack[0];
+                self->return_value_index = 0;
 
                 {
                     ProgramExternFunc extern_func;
-                    return_if_error(amal_program_get_extern_func_by_index(self, extern_func_index, &extern_func));
+                    return_if_error(amal_program_get_extern_func_by_index(self, import_index, extern_func_index, &extern_func));
                     return_if_error(amal_exec_calle(executor, extern_func.func, num_args, dst_reg));
                 }
                 self->read_index += 4;
@@ -549,6 +711,9 @@ static CHECK_RESULT int amal_program_read_instructions(amal_program *self, amal_
                 assert(!inside_func);
                 inside_func = bool_true;
                 assert(func_counter < self->num_functions);
+
+                header_func_set_offset(self, func_counter, amal_exec_get_code_offset(executor));
+                return_if_error(resolve_deferred_func_calls(self, executor, func_counter));
                 ++func_counter;
 
                 func_flags = self->data.data[self->read_index];
@@ -595,9 +760,10 @@ int amal_program_run(amal_program *self) {
         cleanup_if_error(amal_program_read_functions(self));
         cleanup_if_error(amal_program_read_external_functions(self));
         cleanup_if_error(amal_program_read_exported_functions(self));
+        cleanup_if_error(amal_program_read_imports(self));
         cleanup_if_error(amal_program_read_instructions(self, executor));
     }
-    if(self->main_func_instruction_offset == ~0U) {
+    if(self->main_func_instruction_offset == ~(u32)0U) {
         amal_log_error("The program is missing a main function");
         result = AMAL_PROGRAM_NO_MAIN_FUNC;
         goto cleanup;