1 files changed, 29 insertions, 42 deletions

diff --git a/src/compile/statements.c b/src/compile/statements.c
index a7c5214a..bde9ae36 100644
--- a/src/compile/statements.c
+++ b/src/compile/statements.c
@@ -25,6 +25,14 @@ Text_t with_source_info(env_t *env, ast_t *ast, Text_t code) {
     return Texts("\n#line ", line, "\n", code);
 }
 
+static Text_t compile_simple_update_assignment(env_t *env, ast_t *ast, const char *op) {
+    binary_operands_t update = BINARY_OPERANDS(ast);
+    type_t *lhs_t = get_type(env, update.lhs);
+    if (is_idempotent(update.lhs) && (lhs_t->tag == IntType || lhs_t->tag == NumType || lhs_t->tag == ByteType))
+        return Texts(compile_lvalue(env, update.lhs), " ", op, "= ", compile_to_type(env, update.rhs, lhs_t), ";");
+    return compile_update_assignment(env, ast);
+}
+
 static Text_t _compile_statement(env_t *env, ast_t *ast) {
     switch (ast->tag) {
     case When: return compile_when_statement(env, ast);
@@ -47,41 +55,12 @@ static Text_t _compile_statement(env_t *env, ast_t *ast) {
         }
     }
     case Assign: return compile_assignment_statement(env, ast);
-    case PlusUpdate: {
-        DeclareMatch(update, ast, PlusUpdate);
-        type_t *lhs_t = get_type(env, update->lhs);
-        if (is_idempotent(update->lhs) && (lhs_t->tag == IntType || lhs_t->tag == NumType || lhs_t->tag == ByteType))
-            return Texts(compile_lvalue(env, update->lhs), " += ", compile_to_type(env, update->rhs, lhs_t), ";");
-        return compile_update_assignment(env, ast);
-    }
-    case MinusUpdate: {
-        DeclareMatch(update, ast, MinusUpdate);
-        type_t *lhs_t = get_type(env, update->lhs);
-        if (is_idempotent(update->lhs) && (lhs_t->tag == IntType || lhs_t->tag == NumType || lhs_t->tag == ByteType))
-            return Texts(compile_lvalue(env, update->lhs), " -= ", compile_to_type(env, update->rhs, lhs_t), ";");
-        return compile_update_assignment(env, ast);
-    }
-    case MultiplyUpdate: {
-        DeclareMatch(update, ast, MultiplyUpdate);
-        type_t *lhs_t = get_type(env, update->lhs);
-        if (is_idempotent(update->lhs) && (lhs_t->tag == IntType || lhs_t->tag == NumType || lhs_t->tag == ByteType))
-            return Texts(compile_lvalue(env, update->lhs), " *= ", compile_to_type(env, update->rhs, lhs_t), ";");
-        return compile_update_assignment(env, ast);
-    }
-    case DivideUpdate: {
-        DeclareMatch(update, ast, DivideUpdate);
-        type_t *lhs_t = get_type(env, update->lhs);
-        if (is_idempotent(update->lhs) && (lhs_t->tag == IntType || lhs_t->tag == NumType || lhs_t->tag == ByteType))
-            return Texts(compile_lvalue(env, update->lhs), " /= ", compile_to_type(env, update->rhs, lhs_t), ";");
-        return compile_update_assignment(env, ast);
-    }
-    case ModUpdate: {
-        DeclareMatch(update, ast, ModUpdate);
-        type_t *lhs_t = get_type(env, update->lhs);
-        if (is_idempotent(update->lhs) && (lhs_t->tag == IntType || lhs_t->tag == NumType || lhs_t->tag == ByteType))
-            return Texts(compile_lvalue(env, update->lhs), " %= ", compile_to_type(env, update->rhs, lhs_t), ";");
-        return compile_update_assignment(env, ast);
-    }
+    case PlusUpdate: return compile_simple_update_assignment(env, ast, "+");
+    case MinusUpdate: return compile_simple_update_assignment(env, ast, "-");
+    case MultiplyUpdate: return compile_simple_update_assignment(env, ast, "*");
+    case DivideUpdate: return compile_simple_update_assignment(env, ast, "/");
+    case ModUpdate: return compile_simple_update_assignment(env, ast, "%");
+
     case PowerUpdate:
     case Mod1Update:
     case ConcatUpdate:
@@ -131,7 +110,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 +118,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;");
         }
     }