From 7426587c3866cc9ef583b1cbaa6bb7ba5032eb2e Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Mon, 23 Mar 2020 19:49:11 +0100
Subject: [PATCH] fix: add proper return statement type cast if needed, now
 also for non constant values

---
 .../src/prog8/ast/processing/AstChecker.kt    |  2 +-
 .../ast/processing/AstIdentifiersChecker.kt   | 13 -----------
 .../prog8/ast/processing/TypecastsAdder.kt    | 22 +++++++++++++++++++
 3 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/compiler/src/prog8/ast/processing/AstChecker.kt b/compiler/src/prog8/ast/processing/AstChecker.kt
index 97cdaf901..c0c83400a 100644
--- a/compiler/src/prog8/ast/processing/AstChecker.kt
+++ b/compiler/src/prog8/ast/processing/AstChecker.kt
@@ -95,7 +95,7 @@ internal class AstChecker(private val program: Program,
                 errors.err("return value type mismatch", returnStmt.value!!.position)
             } else {
                 if (expectedReturnValues[0] != valueDt.typeOrElse(DataType.STRUCT))
-                    errors.err("type $valueDt of return value doesn't match subroutine's return type", returnStmt.value!!.position)
+                    errors.err("type $valueDt of return value doesn't match subroutine's return type ${expectedReturnValues[0]}", returnStmt.value!!.position)
             }
         }
         super.visit(returnStmt)
diff --git a/compiler/src/prog8/ast/processing/AstIdentifiersChecker.kt b/compiler/src/prog8/ast/processing/AstIdentifiersChecker.kt
index e8b1284ca..2dfdad08f 100644
--- a/compiler/src/prog8/ast/processing/AstIdentifiersChecker.kt
+++ b/compiler/src/prog8/ast/processing/AstIdentifiersChecker.kt
@@ -217,19 +217,6 @@ internal class AstIdentifiersChecker(private val program: Program,
         return super.visit(assignTarget)
     }
 
-    override fun visit(returnStmt: Return): Statement {
-        if(returnStmt.value!=null) {
-            // possibly adjust any literal values returned, into the desired returning data type
-            val subroutine = returnStmt.definingSubroutine()!!
-            if(subroutine.returntypes.size!=1)
-                return returnStmt  // mismatch in number of return values, error will be printed later.
-            val lval = returnStmt.value as? NumericLiteralValue
-            returnStmt.value = lval?.cast(subroutine.returntypes.single()) ?: returnStmt.value!!
-        }
-        return super.visit(returnStmt)
-    }
-
-
     override fun visit(arrayLiteral: ArrayLiteralValue): Expression {
         val array = super.visit(arrayLiteral)
         if(array is ArrayLiteralValue) {
diff --git a/compiler/src/prog8/ast/processing/TypecastsAdder.kt b/compiler/src/prog8/ast/processing/TypecastsAdder.kt
index 93587226b..59a7d1a63 100644
--- a/compiler/src/prog8/ast/processing/TypecastsAdder.kt
+++ b/compiler/src/prog8/ast/processing/TypecastsAdder.kt
@@ -186,4 +186,26 @@ class TypecastsAdder(val program: Program, val errors: ErrorReporter) : AstWalke
         }
         return emptyList()
     }
+
+    override fun after(returnStmt: Return, parent: Node): Iterable<IAstModification> {
+        // add a typecast to the return type if it doesn't match the subroutine's signature
+        val returnValue = returnStmt.value
+        if(returnValue!=null) {
+            val subroutine = returnStmt.definingSubroutine()!!
+            if(subroutine.returntypes.size==1) {
+                val subReturnType = subroutine.returntypes.first()
+                if (returnValue.inferType(program).istype(subReturnType))
+                    return emptyList()
+                if (returnValue is NumericLiteralValue) {
+                    returnStmt.value = returnValue.cast(subroutine.returntypes.single())
+                } else {
+                    return listOf(IAstModification.ReplaceNode(
+                            returnValue,
+                            TypecastExpression(returnValue, subReturnType, true, returnValue.position),
+                            returnStmt))
+                }
+            }
+        }
+        return emptyList()
+    }
 }