Refactor a bit in order to catch issue with ConvertDefs creating

infinite loops

6 files changed, 79 insertions, 69 deletions

diff --git a/src/compile/files.c b/src/compile/files.c
index a6af2300..c250e6cc 100644
--- a/src/compile/files.c
+++ b/src/compile/files.c
@@ -95,13 +95,13 @@ static Text_t compile_top_level_code(env_t *env, ast_t *ast) {
         return compile_function(env, name_code, ast, &env->code->staticdefs);
     }
     case ConvertDef: {
-        type_t *type = get_function_def_type(env, ast);
-        const char *name = get_type_name(Match(type, FunctionType)->ret);
+        type_t *type = get_function_return_type(env, ast);
+        const char *name = get_type_name(type);
         if (!name)
             code_err(ast,
                      "Conversions are only supported for text, struct, and enum "
                      "types, not ",
-                     type_to_str(Match(type, FunctionType)->ret));
+                     type_to_str(type));
         Text_t name_code =
             namespace_name(env, env->namespace, Texts(name, "$", get_line_number(ast->file, ast->start)));
         return compile_function(env, name_code, ast, &env->code->staticdefs);
diff --git a/src/compile/functions.c b/src/compile/functions.c
index 6caefa8b..f04a3b59 100644
--- a/src/compile/functions.c
+++ b/src/compile/functions.c
@@ -256,18 +256,7 @@ Text_t compile_lambda(env_t *env, ast_t *ast) {
         set_binding(body_scope, arg->name, arg_type, Texts("_$", arg->name));
     }
 
-    type_t *ret_t = get_type(body_scope, lambda->body);
-    if (ret_t->tag == ReturnType) ret_t = Match(ret_t, ReturnType)->ret;
-
-    if (lambda->ret_type) {
-        type_t *declared = parse_type_ast(env, lambda->ret_type);
-        if (can_promote(ret_t, declared)) ret_t = declared;
-        else
-            code_err(ast, "This function was declared to return a value of type ", type_to_str(declared),
-                     ", but actually returns a value of type ", type_to_str(ret_t));
-    }
-
-    body_scope->fn_ret = ret_t;
+    body_scope->fn = ast;
 
     Table_t closed_vars = get_closed_vars(env, lambda->args, ast);
     if (Table$length(closed_vars) > 0) { // Create a typedef for the lambda's closure userdata
@@ -289,6 +278,7 @@ Text_t compile_lambda(env_t *env, ast_t *ast) {
         env->code->local_typedefs = Texts(env->code->local_typedefs, def);
     }
 
+    type_t *ret_t = get_function_return_type(env, ast);
     Text_t code = Texts("static ", compile_type(ret_t), " ", name, "(");
     for (arg_ast_t *arg = lambda->args; arg; arg = arg->next) {
         type_t *arg_type = get_arg_ast_type(env, arg);
@@ -623,7 +613,7 @@ Text_t compile_function(env_t *env, Text_t name_code, ast_t *ast, Text_t *static
     bool is_private = false;
     const char *function_name;
     arg_ast_t *args;
-    type_t *ret_t;
+    type_t *ret_t = get_function_return_type(env, ast);
     ast_t *body;
     ast_t *cache;
     bool is_inline;
@@ -632,14 +622,12 @@ Text_t compile_function(env_t *env, Text_t name_code, ast_t *ast, Text_t *static
         function_name = Match(fndef->name, Var)->name;
         is_private = function_name[0] == '_';
         args = fndef->args;
-        ret_t = fndef->ret_type ? parse_type_ast(env, fndef->ret_type) : Type(VoidType);
         body = fndef->body;
         cache = fndef->cache;
         is_inline = fndef->is_inline;
     } else {
         DeclareMatch(convertdef, ast, ConvertDef);
         args = convertdef->args;
-        ret_t = convertdef->ret_type ? parse_type_ast(env, convertdef->ret_type) : Type(VoidType);
         function_name = get_type_name(ret_t);
         if (!function_name)
             code_err(ast,
@@ -689,7 +677,7 @@ Text_t compile_function(env_t *env, Text_t name_code, ast_t *ast, Text_t *static
         set_binding(body_scope, arg->name, arg_type, Texts("_$", arg->name));
     }
 
-    body_scope->fn_ret = ret_t;
+    body_scope->fn = ast;
 
     type_t *body_type = get_type(body_scope, body);
     if (ret_t->tag == AbortType) {
diff --git a/src/compile/statements.c b/src/compile/statements.c
index a7c5214a..4e37838c 100644
--- a/src/compile/statements.c
+++ b/src/compile/statements.c
@@ -131,7 +131,7 @@ static Text_t _compile_statement(env_t *env, ast_t *ast) {
         return code;
     }
     case Return: {
-        if (!env->fn_ret) code_err(ast, "This return statement is not inside any function");
+        if (!env->fn) code_err(ast, "This return statement is not inside any function");
         ast_t *ret = Match(ast, Return)->value;
 
         Text_t code = EMPTY_TEXT;
@@ -139,22 +139,30 @@ static Text_t _compile_statement(env_t *env, ast_t *ast) {
             code = Texts(code, compile_statement(deferred->defer_env, deferred->block));
         }
 
+        type_t *ret_type = get_function_return_type(env, env->fn);
         if (ret) {
-            if (env->fn_ret->tag == VoidType || env->fn_ret->tag == AbortType)
+            if (ret_type->tag == VoidType || ret_type->tag == AbortType)
                 code_err(ast, "This function is not supposed to return any values, "
                               "according to its type signature");
 
-            env = with_enum_scope(env, env->fn_ret);
-            Text_t value = compile_to_type(env, ret, env->fn_ret);
+            env = with_enum_scope(env, ret_type);
+            if (env->fn->tag == ConvertDef) {
+                type_t *value_type = get_type(env, ret);
+                if (!type_eq(value_type, ret_type)) {
+                    code_err(ret, "This value is a ", type_to_text(value_type),
+                             " but this conversion needs an explicit ", type_to_text(ret_type));
+                }
+            }
+            Text_t value = compile_to_type(env, ret, ret_type);
             if (env->deferred) {
-                code = Texts(compile_declaration(env->fn_ret, Text("ret")), " = ", value, ";\n", code);
+                code = Texts(compile_declaration(ret_type, Text("ret")), " = ", value, ";\n", code);
                 value = Text("ret");
             }
 
             return Texts(code, "return ", value, ";");
         } else {
-            if (env->fn_ret->tag != VoidType)
-                code_err(ast, "This function expects you to return a ", type_to_str(env->fn_ret), " value");
+            if (ret_type->tag != VoidType)
+                code_err(ast, "This function expects you to return a ", type_to_text(ret_type), " value");
             return Texts(code, "return;");
         }
     }
diff --git a/src/environment.h b/src/environment.h
index 1ef9c1f9..c726508d 100644
--- a/src/environment.h
+++ b/src/environment.h
@@ -43,7 +43,7 @@ typedef struct env_s {
     Text_t id_suffix;
     Table_t *imports;
     compilation_unit_t *code;
-    type_t *fn_ret;
+    ast_t *fn;
     loop_ctx_t *loop_ctx;
     deferral_t *deferred;
     Closure_t *comprehension_action;
diff --git a/src/typecheck.c b/src/typecheck.c
index eff88d40..2d805cdb 100644
--- a/src/typecheck.c
+++ b/src/typecheck.c
@@ -336,12 +336,12 @@ void bind_statement(env_t *env, ast_t *statement) {
     case FunctionDef: {
         DeclareMatch(def, statement, FunctionDef);
         const char *name = Match(def->name, Var)->name;
-        type_t *type = get_function_def_type(env, statement);
+        type_t *type = get_function_type(env, statement);
         set_binding(env, name, type, namespace_name(env, env->namespace, Text$from_str(name)));
         break;
     }
     case ConvertDef: {
-        type_t *type = get_function_def_type(env, statement);
+        type_t *type = get_function_type(env, statement);
         type_t *ret_t = Match(type, FunctionType)->ret;
         const char *name = get_type_name(ret_t);
         if (!name)
@@ -577,10 +577,24 @@ void bind_statement(env_t *env, ast_t *statement) {
     }
 }
 
-type_t *get_function_def_type(env_t *env, ast_t *ast) {
-    arg_ast_t *arg_asts = ast->tag == FunctionDef ? Match(ast, FunctionDef)->args : Match(ast, ConvertDef)->args;
-    type_ast_t *ret_type =
-        ast->tag == FunctionDef ? Match(ast, FunctionDef)->ret_type : Match(ast, ConvertDef)->ret_type;
+type_t *get_function_type(env_t *env, ast_t *ast) {
+    arg_ast_t *arg_asts;
+    type_ast_t *ret_ast;
+    switch (ast->tag) {
+    case FunctionDef:
+        arg_asts = Match(ast, FunctionDef)->args;
+        ret_ast = Match(ast, FunctionDef)->ret_type;
+        break;
+    case ConvertDef:
+        arg_asts = Match(ast, ConvertDef)->args;
+        ret_ast = Match(ast, ConvertDef)->ret_type;
+        break;
+    case Lambda:
+        arg_asts = Match(ast, Lambda)->args;
+        ret_ast = Match(ast, Lambda)->ret_type;
+        break;
+    default: code_err(ast, "This was expected to be a function definition of some sort");
+    }
     arg_t *args = NULL;
     env_t *scope = fresh_scope(env);
     for (arg_ast_t *arg = arg_asts; arg; arg = arg->next) {
@@ -590,10 +604,38 @@ type_t *get_function_def_type(env_t *env, ast_t *ast) {
     }
     REVERSE_LIST(args);
 
-    type_t *ret = ret_type ? parse_type_ast(scope, ret_type) : Type(VoidType);
-    if (has_stack_memory(ret))
-        code_err(ast, "Functions can't return stack references because the reference may outlive its stack frame.");
-    return Type(FunctionType, .args = args, .ret = ret);
+    if (ast->tag == Lambda) {
+        ast_t *body = Match(ast, Lambda)->body;
+        type_t *ret_t = get_type(scope, body);
+        if (ret_t->tag == ReturnType) ret_t = Match(ret_t, ReturnType)->ret;
+        if (ret_t->tag == AbortType) ret_t = Type(VoidType);
+
+        if (ret_t->tag == OptionalType && !Match(ret_t, OptionalType)->type)
+            code_err(body, "This function doesn't return a specific optional type");
+
+        if (ret_ast) {
+            type_t *declared = parse_type_ast(env, ret_ast);
+            if (can_promote(ret_t, declared)) ret_t = declared;
+            else
+                code_err(ast, "This function was declared to return a value of type ", type_to_str(declared),
+                         ", but actually returns a value of type ", type_to_str(ret_t));
+        }
+
+        if (has_stack_memory(ret_t))
+            code_err(ast, "Functions can't return stack references because the reference may outlive its stack frame.");
+        return Type(ClosureType, Type(FunctionType, .args = args, .ret = ret_t));
+    } else {
+        type_t *ret_t = ret_ast ? parse_type_ast(scope, ret_ast) : Type(VoidType);
+        if (has_stack_memory(ret_t))
+            code_err(ast, "Functions can't return stack references because the reference may outlive its stack frame.");
+        return Type(FunctionType, .args = args, .ret = ret_t);
+    }
+}
+
+type_t *get_function_return_type(env_t *env, ast_t *ast) {
+    type_t *fn_t = get_function_type(env, ast);
+    if (fn_t->tag == ClosureType) fn_t = Match(fn_t, ClosureType)->fn;
+    return Match(fn_t, FunctionType)->ret;
 }
 
 type_t *get_method_type(env_t *env, ast_t *self, const char *name) {
@@ -1085,7 +1127,7 @@ type_t *get_type(env_t *env, ast_t *ast) {
     }
     case Return: {
         ast_t *val = Match(ast, Return)->value;
-        if (env->fn_ret) env = with_enum_scope(env, env->fn_ret);
+        if (env->fn) env = with_enum_scope(env, get_function_return_type(env, env->fn));
         return Type(ReturnType, .ret = (val ? get_type(env, val) : Type(VoidType)));
     }
     case Stop:
@@ -1394,36 +1436,7 @@ type_t *get_type(env_t *env, ast_t *ast) {
         return t;
     }
 
-    case Lambda: {
-        DeclareMatch(lambda, ast, Lambda);
-        arg_t *args = NULL;
-        env_t *scope = fresh_scope(env); // For now, just use closed variables in scope normally
-        for (arg_ast_t *arg = lambda->args; arg; arg = arg->next) {
-            type_t *t = get_arg_ast_type(env, arg);
-            args = new (arg_t, .name = arg->name, .alias = arg->alias, .type = t, .next = args);
-            set_binding(scope, arg->name, t, EMPTY_TEXT);
-        }
-        REVERSE_LIST(args);
-
-        type_t *ret = get_type(scope, lambda->body);
-        if (ret->tag == ReturnType) ret = Match(ret, ReturnType)->ret;
-        if (ret->tag == AbortType) ret = Type(VoidType);
-
-        if (ret->tag == OptionalType && !Match(ret, OptionalType)->type)
-            code_err(lambda->body, "This function doesn't return a specific optional type");
-
-        if (lambda->ret_type) {
-            type_t *declared = parse_type_ast(env, lambda->ret_type);
-            if (can_promote(ret, declared)) ret = declared;
-            else
-                code_err(ast, "This function was declared to return a value of type ", type_to_str(declared),
-                         ", but actually returns a value of type ", type_to_str(ret));
-        }
-
-        if (has_stack_memory(ret))
-            code_err(ast, "Functions can't return stack references because the reference may outlive its stack frame.");
-        return Type(ClosureType, Type(FunctionType, .args = args, .ret = ret));
-    }
+    case Lambda: return get_function_type(env, ast);
 
     case FunctionDef:
     case ConvertDef:
diff --git a/src/typecheck.h b/src/typecheck.h
index 8fc30333..d64bb316 100644
--- a/src/typecheck.h
+++ b/src/typecheck.h
@@ -16,7 +16,8 @@ void prebind_statement(env_t *env, ast_t *statement);
 void bind_statement(env_t *env, ast_t *statement);
 PUREFUNC type_t *get_math_type(env_t *env, ast_t *ast, type_t *lhs_t, type_t *rhs_t);
 PUREFUNC bool is_discardable(env_t *env, ast_t *ast);
-type_t *get_function_def_type(env_t *env, ast_t *ast);
+type_t *get_function_type(env_t *env, ast_t *ast);
+type_t *get_function_return_type(env_t *env, ast_t *ast);
 type_t *get_arg_type(env_t *env, arg_t *arg);
 type_t *get_arg_ast_type(env_t *env, arg_ast_t *arg);
 env_t *when_clause_scope(env_t *env, type_t *subject_t, when_clause_t *clause);