Fix for order of operations issues with enum and function typedef code

7 files changed, 131 insertions, 69 deletions

diff --git a/compile.c b/compile.c
index 0acbceb6..acd08f3a 100644
--- a/compile.c
+++ b/compile.c
@@ -264,6 +264,8 @@ CORD compile_statement(env_t *env, ast_t *ast)
         }
         if (when->else_body) {
             code = CORD_all(code, "default: {\n", compile_statement(env, when->else_body), "\nbreak;\n}");
+        } else {
+            code = CORD_all(code, "default: errx(1, \"Invalid tag!\");\n");
         }
         code = CORD_all(code, "\n}\n}");
         return code;
@@ -615,7 +617,7 @@ CORD compile_statement(env_t *env, ast_t *ast)
         env->code->funcs = CORD_all(env->code->funcs, code, " ", body, "\n");
 
         if (fndef->cache && fndef->cache->tag == Int && Match(fndef->cache, Int)->i > 0) {
-            const char *arg_type_name = heap_strf("%s$args", CORD_to_const_char_star(name));
+            const char *arg_type_name = heap_strf("%s$args", Match(fndef->name, Var)->name);
             ast_t *args_def = FakeAST(StructDef, .name=arg_type_name, .fields=fndef->args);
             prebind_statement(env, args_def);
             bind_statement(env, args_def);
@@ -633,8 +635,10 @@ CORD compile_statement(env_t *env, ast_t *ast)
                                     ") Table$remove(&cache, NULL, table_type);\n");
             }
 
-            CORD arg_typedef = compile_struct_header(env, args_def);
+            CORD arg_typedef = compile_struct_typedef(env, args_def);
             env->code->local_typedefs = CORD_all(env->code->local_typedefs, arg_typedef);
+            env->code->staticdefs = CORD_all(env->code->staticdefs,
+                                             "extern const TypeInfo ", namespace_prefix(env->libname, env->namespace), arg_type_name, ";\n");
             CORD wrapper = CORD_all(
                 is_private ? CORD_EMPTY : "public ", ret_type_code, " ", name, arg_signature, "{\n"
                 "static table_t cache = {};\n",
@@ -1787,7 +1791,7 @@ CORD compile(env_t *env, ast_t *ast)
             code = CORD_all(code, compile_type(arg_type), " $", arg->name, ", ");
         }
 
-        CORD args_typedef = compile_statement_header(env, ast);
+        CORD args_typedef = compile_statement_typedefs(env, ast);
         env->code->local_typedefs = CORD_all(env->code->local_typedefs, args_typedef);
 
         table_t *closed_vars = get_closed_vars(env, ast);
@@ -2243,12 +2247,12 @@ void compile_namespace(env_t *env, const char *ns_name, ast_t *block)
     }
 }
 
-CORD compile_namespace_headers(env_t *env, const char *ns_name, ast_t *block)
+CORD compile_namespace_definitions(env_t *env, const char *ns_name, ast_t *block)
 {
     env_t *ns_env = namespace_env(env, ns_name);
     CORD header = CORD_EMPTY;
     for (ast_list_t *stmt = block ? Match(block, Block)->statements : NULL; stmt; stmt = stmt->next) {
-        header = CORD_all(header, compile_statement_header(ns_env, stmt->ast));
+        header = CORD_all(header, compile_statement_definitions(ns_env, stmt->ast));
     }
     return header;
 }
@@ -2514,17 +2518,83 @@ CORD compile_file(env_t *env, ast_t *ast)
     );
 }
 
-CORD compile_statement_header(env_t *env, ast_t *ast)
+CORD compile_statement_imports(env_t *env, ast_t *ast)
 {
     switch (ast->tag) {
     case DocTest: {
         auto test = Match(ast, DocTest);
-        return compile_statement_header(env, test->expr);
+        return compile_statement_imports(env, test->expr);
+    }
+    case Declare: {
+        auto decl = Match(ast, Declare);
+        if (decl->value->tag == Use || decl->value->tag == Import)
+            return compile_statement_imports(env, decl->value);
+        return CORD_EMPTY;
+    }
+    case Import: {
+        const char *path = Match(ast, Import)->path;
+        return CORD_all("#include \"", path, ".tm.h\"\n");
+    }
+    case Use: {
+        const char *name = Match(ast, Use)->name;
+        if (strncmp(name, "-l", 2) == 0)
+            return CORD_EMPTY;
+        else
+            return CORD_all("#include <tomo/lib", name, ".h>\n");
+    }
+    default: return CORD_EMPTY;
+    }
+}
+
+CORD compile_statement_typedefs(env_t *env, ast_t *ast)
+{
+    switch (ast->tag) {
+    case DocTest: {
+        auto test = Match(ast, DocTest);
+        return compile_statement_typedefs(env, test->expr);
+    }
+    case StructDef: {
+        return compile_struct_typedef(env, ast);
+    }
+    case EnumDef: {
+        return compile_enum_typedef(env, ast);
+    }
+    case LangDef: {
+        auto def = Match(ast, LangDef);
+        return CORD_all("typedef CORD ", namespace_prefix(env->libname, env->namespace), def->name, "_t;\n");
+    }
+    case Lambda: {
+        auto lambda = Match(ast, Lambda);
+        table_t *closed_vars = get_closed_vars(env, ast);
+        if (Table$length(*closed_vars) == 0)
+            return CORD_EMPTY;
+
+        CORD def = "typedef struct {";
+        for (int64_t i = 1; i <= Table$length(*closed_vars); i++) {
+            struct { const char *name; binding_t *b; } *entry = Table$entry(*closed_vars, i);
+            if (entry->b->type->tag == ModuleType)
+                continue;
+            def = CORD_all(def, compile_declaration(entry->b->type, entry->name), "; ");
+        }
+        CORD name = CORD_asprintf("%rlambda$%ld", namespace_prefix(env->libname, env->namespace), lambda->id);
+        return CORD_all(def, "} ", name, "$userdata_t;");
+    }
+    default:
+        return CORD_EMPTY;
+    }
+}
+
+CORD compile_statement_definitions(env_t *env, ast_t *ast)
+{
+    switch (ast->tag) {
+    case DocTest: {
+        auto test = Match(ast, DocTest);
+        return compile_statement_definitions(env, test->expr);
     }
     case Declare: {
         auto decl = Match(ast, Declare);
         if (decl->value->tag == Use || decl->value->tag == Import) {
-            return compile_statement_header(env, decl->value);
+            return compile_statement_definitions(env, decl->value);
         }
         type_t *t = get_type(env, decl->value);
         assert(t->tag != ModuleType);
@@ -2535,7 +2605,7 @@ CORD compile_statement_header(env_t *env, ast_t *ast)
         if (!is_constant(env, decl->value))
             code_err(decl->value, "This value is not a valid constant initializer.");
 
-        CORD code = (decl->value->tag == Use || decl->value->tag == Import) ? compile_statement_header(env, decl->value) : CORD_EMPTY;
+        CORD code = (decl->value->tag == Use || decl->value->tag == Import) ? compile_statement_definitions(env, decl->value) : CORD_EMPTY;
         if (is_private) {
             return code;
         } else {
@@ -2544,18 +2614,21 @@ CORD compile_statement_header(env_t *env, ast_t *ast)
         }
     }
     case StructDef: {
-        return compile_struct_header(env, ast);
+        auto def = Match(ast, StructDef);
+        CORD full_name = CORD_cat(namespace_prefix(env->libname, env->namespace), def->name);
+        return CORD_all(
+            "extern const TypeInfo ", full_name, ";\n",
+            compile_namespace_definitions(env, def->name, def->namespace));
     }
     case EnumDef: {
-        return compile_enum_header(env, ast);
+        return compile_enum_declarations(env, ast);
     }
     case LangDef: {
         auto def = Match(ast, LangDef);
+        CORD full_name = CORD_cat(namespace_prefix(env->libname, env->namespace), def->name);
         return CORD_all(
-            "typedef CORD ", namespace_prefix(env->libname, env->namespace), def->name, "_t;\n",
-            "extern const TypeInfo ", namespace_prefix(env->libname, env->namespace), def->name, ";\n");
-        compile_namespace(env, def->name, def->namespace);
-        return CORD_EMPTY;
+            "extern const TypeInfo ", full_name, ";\n",
+            compile_namespace_definitions(env, def->name, def->namespace));
     }
     case FunctionDef: {
         auto fndef = Match(ast, FunctionDef);
@@ -2572,25 +2645,7 @@ CORD compile_statement_header(env_t *env, ast_t *ast)
 
         type_t *ret_t = fndef->ret_type ? parse_type_ast(env, fndef->ret_type) : Type(VoidType);
         CORD ret_type_code = compile_type(ret_t);
-        CORD header = CORD_all(ret_type_code, " ", CORD_cat(namespace_prefix(env->libname, env->namespace), decl_name), arg_signature, ";\n");
-        return header;
-    }
-    case Lambda: {
-        auto lambda = Match(ast, Lambda);
-        CORD name = CORD_asprintf("%rlambda$%ld", namespace_prefix(env->libname, env->namespace), lambda->id);
-        table_t *closed_vars = get_closed_vars(env, ast);
-        if (Table$length(*closed_vars) == 0) {
-            return CORD_EMPTY;
-        } else {
-            CORD def = "typedef struct {";
-            for (int64_t i = 1; i <= Table$length(*closed_vars); i++) {
-                struct { const char *name; binding_t *b; } *entry = Table$entry(*closed_vars, i);
-                if (entry->b->type->tag == ModuleType)
-                    continue;
-                def = CORD_all(def, compile_declaration(entry->b->type, entry->name), "; ");
-            }
-            return CORD_all(def, "} ", name, "$userdata_t;");
-        }
+        return CORD_all(ret_type_code, " ", namespace_prefix(env->libname, env->namespace), decl_name, arg_signature, ";\n");
     }
     case Extern: {
         auto ext = Match(ast, Extern);
@@ -2610,17 +2665,6 @@ CORD compile_statement_header(env_t *env, ast_t *ast)
         }
         return CORD_all("extern ", decl, ";\n");
     }
-    case Import: {
-        const char *path = Match(ast, Import)->path;
-        return CORD_all("#include \"", path, ".tm.h\"\n");
-    }
-    case Use: {
-        const char *name = Match(ast, Use)->name;
-        if (strncmp(name, "-l", 2) == 0)
-            return CORD_EMPTY;
-        else
-            return CORD_all("#include <tomo/lib", name, ".h>\n");
-    }
     default:
         return CORD_EMPTY;
     }
@@ -2633,7 +2677,13 @@ CORD compile_header(env_t *env, ast_t *ast)
                   "#include <tomo/tomo.h>\n";
 
     for (ast_list_t *stmt = Match(ast, Block)->statements; stmt; stmt = stmt->next)
-        header = CORD_all(header, compile_statement_header(env, stmt->ast));
+        header = CORD_all(header, compile_statement_imports(env, stmt->ast));
+
+    for (ast_list_t *stmt = Match(ast, Block)->statements; stmt; stmt = stmt->next)
+        header = CORD_all(header, compile_statement_typedefs(env, stmt->ast));
+
+    for (ast_list_t *stmt = Match(ast, Block)->statements; stmt; stmt = stmt->next)
+        header = CORD_all(header, compile_statement_definitions(env, stmt->ast));
 
     return header;
 }
diff --git a/compile.h b/compile.h
index 62e74729..fa18ee89 100644
--- a/compile.h
+++ b/compile.h
@@ -16,9 +16,10 @@ CORD compile_declaration(type_t *t, const char *name);
 CORD compile_type(type_t *t);
 CORD compile(env_t *env, ast_t *ast);
 void compile_namespace(env_t *env, const char *ns_name, ast_t *block);
-CORD compile_namespace_headers(env_t *env, const char *ns_name, ast_t *block);
+CORD compile_namespace_definitions(env_t *env, const char *ns_name, ast_t *block);
 CORD compile_statement(env_t *env, ast_t *ast);
-CORD compile_statement_header(env_t *env, ast_t *ast);
+CORD compile_statement_typedefs(env_t *env, ast_t *ast);
+CORD compile_statement_definitions(env_t *env, ast_t *ast);
 CORD compile_type_info(env_t *env, type_t *t);
 CORD compile_cli_arg_call(env_t *env, CORD fn_name, type_t *fn_type);
 
diff --git a/enums.c b/enums.c
index 4e632216..9b73f3f7 100644
--- a/enums.c
+++ b/enums.c
@@ -200,12 +200,11 @@ void compile_enum_def(env_t *env, ast_t *ast)
     compile_namespace(env, def->name, def->namespace);
 }
 
-CORD compile_enum_header(env_t *env, ast_t *ast)
+CORD compile_enum_typedef(env_t *env, ast_t *ast)
 {
     auto def = Match(ast, EnumDef);
     CORD full_name = CORD_cat(namespace_prefix(env->libname, env->namespace), def->name);
-    CORD header = CORD_all("typedef struct ", full_name, "_s ", full_name, "_t;\n",
-                           "extern const TypeInfo ", full_name, ";\n");
+    CORD all_defs = CORD_all("typedef struct ", full_name, "_s ", full_name, "_t;\n");
     CORD enum_def = CORD_all("struct ", full_name, "_s {\n"
                              "\tenum {");
     for (tag_ast_t *tag = def->tags; tag; tag = tag->next) {
@@ -214,11 +213,24 @@ CORD compile_enum_header(env_t *env, ast_t *ast)
     }
     enum_def = CORD_cat(enum_def, "} $tag;\n"
                         "union {\n");
-    CORD constructors = CORD_EMPTY;
     for (tag_ast_t *tag = def->tags; tag; tag = tag->next) {
-        CORD field_header = compile_struct_header(env, WrapAST(ast, StructDef, .name=CORD_to_const_char_star(CORD_all(def->name, "$", tag->name)), .fields=tag->fields));
-        header = CORD_all(header, field_header);
+        CORD field_def = compile_struct_typedef(env, WrapAST(ast, StructDef, .name=CORD_to_const_char_star(CORD_all(def->name, "$", tag->name)), .fields=tag->fields));
+        all_defs = CORD_all(all_defs, field_def);
         enum_def = CORD_all(enum_def, full_name, "$", tag->name, "_t ", tag->name, ";\n");
+    }
+    enum_def = CORD_cat(enum_def, "};\n};\n");
+    all_defs = CORD_all(all_defs, enum_def);
+    return all_defs;
+}
+
+CORD compile_enum_declarations(env_t *env, ast_t *ast)
+{
+    auto def = Match(ast, EnumDef);
+    CORD full_name = CORD_cat(namespace_prefix(env->libname, env->namespace), def->name);
+    CORD all_defs = CORD_all("extern const TypeInfo ", full_name, ";\n");
+    for (tag_ast_t *tag = def->tags; tag; tag = tag->next) {
+        all_defs = CORD_all(all_defs,
+                            "extern const TypeInfo ", namespace_prefix(env->libname, env->namespace), def->name, "$", tag->name, ";\n");
         if (tag->fields) { // Constructor macros:
             CORD arg_sig = CORD_EMPTY;
             for (arg_ast_t *field = tag->fields; field; field = field->next) {
@@ -228,12 +240,10 @@ CORD compile_enum_header(env_t *env, ast_t *ast)
             }
             if (arg_sig == CORD_EMPTY) arg_sig = "void";
             CORD constructor_def = CORD_all(full_name, "_t ", full_name, "$tagged$", tag->name, "(", arg_sig, ");\n");
-            constructors = CORD_cat(constructors, constructor_def);
+            all_defs = CORD_cat(all_defs, constructor_def);
         }
     }
-    enum_def = CORD_cat(enum_def, "};\n};\n");
-    header = CORD_all(header, enum_def, constructors);
-    return CORD_all(header, compile_namespace_headers(env, def->name, def->namespace));
+    return CORD_all(all_defs, compile_namespace_definitions(env, def->name, def->namespace));
 }
 
 // vim: ts=4 sw=0 et cino=L2,l1,(0,W4,m1,\:0
diff --git a/enums.h b/enums.h
index 68b03ec9..603c17bb 100644
--- a/enums.h
+++ b/enums.h
@@ -8,6 +8,7 @@
 #include "environment.h"
 
 void compile_enum_def(env_t *env, ast_t *ast);
-CORD compile_enum_header(env_t *env, ast_t *ast);
+CORD compile_enum_typedef(env_t *env, ast_t *ast);
+CORD compile_enum_declarations(env_t *env, ast_t *ast);
 
 // vim: ts=4 sw=0 et cino=L2,l1,(0,W4,m1,\:0
diff --git a/structs.c b/structs.c
index 22965b72..80049cdc 100644
--- a/structs.c
+++ b/structs.c
@@ -165,11 +165,11 @@ void compile_struct_def(env_t *env, ast_t *ast)
     compile_namespace(env, def->name, def->namespace);
 }
 
-CORD compile_struct_header(env_t *env, ast_t *ast)
+CORD compile_struct_typedef(env_t *env, ast_t *ast)
 {
     auto def = Match(ast, StructDef);
     CORD full_name = CORD_cat(namespace_prefix(env->libname, env->namespace), def->name);
-    CORD header = CORD_all("typedef struct ", full_name, "_s ", full_name, "_t;\n");
+    CORD code = CORD_all("typedef struct ", full_name, "_s ", full_name, "_t;\n");
 
     CORD struct_code = CORD_all("struct ", full_name, "_s {\n");
     for (arg_ast_t *field = def->fields; field; field = field->next) {
@@ -179,10 +179,8 @@ CORD compile_struct_header(env_t *env, ast_t *ast)
                      CORD_cmp(type_code, "Bool_t") ? "" : ":1");
     }
     struct_code = CORD_all(struct_code, "};\n");
-    header = CORD_all(header, struct_code);
-    header = CORD_all(header, "extern const TypeInfo ", full_name, ";\n");
-    header = CORD_all(header, compile_namespace_headers(env, def->name, def->namespace));
-    return header;
+    code = CORD_all(code, struct_code);
+    return code;
 }
 
 // vim: ts=4 sw=0 et cino=L2,l1,(0,W4,m1,\:0
diff --git a/structs.h b/structs.h
index 166a382b..0f7b0ad8 100644
--- a/structs.h
+++ b/structs.h
@@ -8,6 +8,6 @@
 #include "environment.h"
 
 void compile_struct_def(env_t *env, ast_t *ast);
-CORD compile_struct_header(env_t *env, ast_t *ast);
+CORD compile_struct_typedef(env_t *env, ast_t *ast);
 
 // vim: ts=4 sw=0 et cino=L2,l1,(0,W4,m1,\:0
diff --git a/tomo.c b/tomo.c
index f6030577..631192e4 100644
--- a/tomo.c
+++ b/tomo.c
@@ -209,9 +209,11 @@ int main(int argc, char *argv[])
             if (!ast) errx(1, "Could not parse file %s", filename);
             env->namespace = new(namespace_t, .name=file_base_name(filename));
             for (ast_list_t *stmt = Match(ast, Block)->statements; stmt; stmt = stmt->next) {
-                if (stmt->ast->tag == Import || (stmt->ast->tag == Declare && Match(stmt->ast, Declare)->value->tag == Import))
-                    continue;
-                CORD h = compile_statement_header(env, stmt->ast);
+                CORD h = compile_statement_typedefs(env, stmt->ast);
+                if (h) CORD_put(h, header_prog);
+            }
+            for (ast_list_t *stmt = Match(ast, Block)->statements; stmt; stmt = stmt->next) {
+                CORD h = compile_statement_definitions(env, stmt->ast);
                 if (h) CORD_put(h, header_prog);
             }
         }