fix sub return register

codeGenCpu6502/src/prog8/codegen/cpu6502/AsmGen.kt

@@ -835,7 +835,7 @@ $repeatLabel    lda  $counterVar
     private fun translate(ret: PtReturn, withRts: Boolean=true) {
         ret.value?.let { returnvalue ->
             val sub = ret.definingSub()!!
-            val returnReg = RegisterOrStatusflag(RegisterOrPair.A, null) //  TODO("what is returnReg of the sub?")
+            val returnReg = sub.returnRegister()!!
             when (sub.returntype) {
                 in NumericDatatypes -> {
                     assignExpressionToRegister(returnvalue, returnReg.registerOrPair!!)

codeGenCpu6502/src/prog8/codegen/cpu6502/ExpressionsAsmGen.kt

@@ -44,7 +44,7 @@ internal class ExpressionsAsmGen(private val program: PtProgram,
     private fun translateFunctionCallResultOntoStack(call: PtFunctionCall) {
         // only for use in nested expression evaluation
 
-        val sub = call.targetSubroutine(program)!!
+        val sub = call.targetSubroutine(program)
         asmgen.saveXbeforeCall(call)
         asmgen.translateFunctionCall(call, true)
         if(sub.regXasResult()) {

codeGenCpu6502/src/prog8/codegen/cpu6502/Extensions.kt

@@ -66,6 +66,7 @@ internal fun PtIdentifier.targetStatement(program: PtProgram): PtNode {
 }
 
 internal fun PtProgram.lookup(name: String): PtNode {
+    // TODO should be cached?
     fun searchLocalSymbol(node: PtNode, namePart: String): PtNode? {
         when(node) {
             is PtProgram -> {
@@ -131,8 +132,8 @@ internal fun PtAsmSub.shouldKeepA(): KeepAresult {
     return KeepAresult(false, saveAonReturn)
 }
 
-internal fun PtFunctionCall.targetSubroutine(program: PtProgram): IPtSubroutine? =
-    this.targetStatement(program) as? IPtSubroutine
+internal fun PtFunctionCall.targetSubroutine(program: PtProgram): IPtSubroutine =
+    this.targetStatement(program) as IPtSubroutine
 
 internal fun PtFunctionCall.targetStatement(program: PtProgram): PtNode {
     return if(name in BuiltinFunctions)
@@ -185,4 +186,14 @@ internal fun PtExpression.clone(): PtExpression {
         is PtString -> return withClonedChildrenFrom(this, PtString(value, encoding, position))
         is PtTypeCast -> return withClonedChildrenFrom(this, PtTypeCast(type, position))
     }
+}
+
+internal fun PtSub.returnRegister(): RegisterOrStatusflag? {
+    return when(returntype) {
+        in ByteDatatypes -> RegisterOrStatusflag(RegisterOrPair.A, null)
+        in WordDatatypes -> RegisterOrStatusflag(RegisterOrPair.AY, null)
+        DataType.FLOAT -> RegisterOrStatusflag(RegisterOrPair.FAC1, null)
+        null -> null
+        else -> RegisterOrStatusflag(RegisterOrPair.AY, null)
+    }
 }

codeGenCpu6502/src/prog8/codegen/cpu6502/FunctionCallAsmGen.kt

@@ -18,7 +18,7 @@ internal class FunctionCallAsmGen(private val program: PtProgram, private val as
     }
 
     internal fun saveXbeforeCall(stmt: PtFunctionCall) {
-        val sub = stmt.targetSubroutine(program) ?: throw AssemblyError("undefined subroutine ${stmt.name}")
+        val sub = stmt.targetSubroutine(program)
         if(sub.shouldSaveX()) {
             if(sub is PtAsmSub) {
                 val regSaveOnStack = sub.address == null       // rom-routines don't require registers to be saved on stack, normal subroutines do because they can contain nested calls
@@ -32,7 +32,7 @@ internal class FunctionCallAsmGen(private val program: PtProgram, private val as
     }
 
     internal fun restoreXafterCall(stmt: PtFunctionCall) {
-        val sub = stmt.targetSubroutine(program) ?: throw AssemblyError("undefined subroutine ${stmt.name}")
+        val sub = stmt.targetSubroutine(program)
         if(sub.shouldSaveX()) {
             if(sub is PtAsmSub) {
                 val regSaveOnStack = sub.address == null       // rom-routines don't require registers to be saved on stack, normal subroutines do because they can contain nested calls
@@ -55,7 +55,7 @@ internal class FunctionCallAsmGen(private val program: PtProgram, private val as
         // NOTE: does NOT output code to save/restore the X register for this call! Every caller should deal with this in their own way!!
         //       (you can use subroutine.shouldSaveX() and saveX()/restoreX() routines as a help for this)
 
-        val sub: IPtSubroutine = call.targetSubroutine(program) ?: throw AssemblyError("undefined subroutine ${call.name}")
+        val sub: IPtSubroutine = call.targetSubroutine(program)
         val subAsmName = asmgen.asmSymbolName(call.name)
 
         if(sub is PtAsmSub) {

codeGenCpu6502/src/prog8/codegen/cpu6502/assignment/AsmAssignment.kt

@@ -158,7 +158,7 @@ internal class AsmAssignSource(val kind: SourceStorageKind,
                     AsmAssignSource(SourceStorageKind.EXPRESSION, program, asmgen, value.type, expression = value)
                 }
                 is PtFunctionCall -> {
-                    val sub = value.targetSubroutine(program)!!
+                    val sub = value.targetSubroutine(program)
                     val returnType = sub.returnsWhatWhere().firstOrNull { rr -> rr.second.registerOrPair != null || rr.second.statusflag!=null }?.first
                             ?: throw AssemblyError("can't translate zero return values in assignment")
 

codeGenCpu6502/src/prog8/codegen/cpu6502/assignment/AssignmentAsmGen.kt

@@ -178,7 +178,7 @@ internal class AssignmentAsmGen(private val program: PtProgram,
             is PtMemoryByte -> throw AssemblyError("source kind should have been memory")
             is PtTypeCast -> assignTypeCastedValue(assign.target, value.type, value.value, value)
             is PtFunctionCall -> {
-                val sub = value.targetSubroutine(program)!!
+                val sub = value.targetSubroutine(program)
                 asmgen.saveXbeforeCall(value)
                 asmgen.translateFunctionCall(value, true)
                 val returnValue = sub.returnsWhatWhere().singleOrNull() { it.second.registerOrPair!=null } ?: sub.returnsWhatWhere().single() { it.second.statusflag!=null }

examples/test.p8

@@ -1,22 +1,23 @@
-%import textio
-%option no_sysinit
-%zeropage basicsafe
-
 main {
 
     sub start() {
-        rsavex()
-
-        ubyte @shared ub
-        main.normalfoo.arg=99
-        void normalfoo(42)
-somelabel:
-        ub++
-        txt.print_ub(ub)
-    }
-
-    sub normalfoo(ubyte arg) -> ubyte {
-        arg++
-        return 42
+        testscope.duplicate()
+        cx16.r0L = testscope.duplicate2()
+    }
+}
+
+testscope {
+
+    sub sub1() {
+        ubyte @shared duplicate
+        ubyte @shared duplicate2
+    }
+
+    sub duplicate() {
+        ; do nothing
+    }
+
+    sub duplicate2() -> ubyte {
+        return cx16.r0L
     }
 }