1 files changed, 16 insertions, 10 deletions

diff --git a/src/program.c b/src/program.c
index 128f0f9..4c0f1a2 100644
--- a/src/program.c
+++ b/src/program.c
@@ -78,7 +78,7 @@ void amal_program_deinit(amal_program *self) {
         buffer_deinit(&self->data);
 }
 
-int amal_program_register_extern_func(amal_program *self, BufferView name, void *func_ptr, int args_byte_size) {
+int amal_program_register_extern_func(amal_program *self, BufferView name, void *func_ptr, u32 args_byte_size) {
     ProgramExternFunc extern_func;
     extern_func.func = func_ptr;
     extern_func.args_byte_size = args_byte_size;
@@ -112,15 +112,16 @@ static void amal_program_get_header_extern_function_by_index(amal_program *self,
     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;
+        BytecodeHeaderExternFunction *extern_func = (BytecodeHeaderExternFunction*)extern_funcs_start;
         extern_funcs_start += sizeof(BytecodeHeaderExternFunction);
-        name_len = *extern_funcs_start;
-        extern_funcs_start += 1 + name_len + 1; /* +1 for skipping length byte and the null-terminated character */
+        extern_funcs_start += extern_func->name_len + 1; /* +1 for skipping the null-terminated character */
     }
+
     am_memcpy(&result->extern_func, extern_funcs_start, sizeof(result->extern_func));
-    result->name.size = *(extern_funcs_start + sizeof(result->extern_func));
-    result->name.data = (const char*)extern_funcs_start + sizeof(result->extern_func) + sizeof(u8);
+    result->name.size = result->extern_func.name_len;
+    result->name.data = (const char*)extern_funcs_start + sizeof(result->extern_func);
 }
 
 
@@ -133,10 +134,15 @@ static CHECK_RESULT int amal_program_get_extern_func_by_index(amal_program *self
         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 != (int)extern_func_get_total_param_size(&extern_func.extern_func)) {
-        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, (int)extern_func_get_total_param_size(&extern_func.extern_func));
+    if(extern_func.extern_func.flags & FUNC_FLAG_VARARGS) {
+        if(extern_func_get_total_param_size(&extern_func.extern_func) < result->args_byte_size) {
+            amal_log_error("Extern function %.*s was registered to take at least %u byte(s), but the program says it takes at least %u byte(s)",
+                extern_func.name.size, extern_func.name.data, result->args_byte_size, extern_func_get_total_param_size(&extern_func.extern_func));
+            return AMAL_PROGRAM_NO_SUCH_EXTERNAL_FUNCTION;
+        }
+    } else if(extern_func_get_total_param_size(&extern_func.extern_func) != result->args_byte_size) {
+        amal_log_error("Extern function %.*s was registered to take %u byte(s), but the program says it takes %u byte(s)",
+            extern_func.name.size, extern_func.name.data, result->args_byte_size, extern_func_get_total_param_size(&extern_func.extern_func));
         return AMAL_PROGRAM_NO_SUCH_EXTERNAL_FUNCTION;
     }
     return 0;