prog8/compiler/test/codegeneration/TestVariousCodeGen.kt

package prog8tests.codegeneration

import io.kotest.assertions.withClue
import io.kotest.core.spec.style.FunSpec
import io.kotest.engine.spec.tempdir
import io.kotest.matchers.ints.shouldBeGreaterThan
import io.kotest.matchers.shouldBe
import io.kotest.matchers.shouldNotBe
import io.kotest.matchers.string.shouldContain
import io.kotest.matchers.string.shouldStartWith
import io.kotest.matchers.types.instanceOf
import io.kotest.matchers.types.shouldBeInstanceOf
import prog8.code.StNodeType
import prog8.code.ast.*
import prog8.code.core.BaseDataType
import prog8.code.core.DataType
import prog8.code.target.*
import prog8tests.helpers.ErrorReporterForTests
import prog8tests.helpers.compileText
import prog8tests.helpers.shouldContainInOrder
import kotlin.io.path.readText

class TestVariousCodeGen: FunSpec({

    val outputDir = tempdir().toPath()

    test("nested scoping") {
        val text="""
main {
    sub start() {
        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
    }
}"""
        compileText(C64Target(), false, text, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("word array indexing") {
        val text="""
main {
    sub start() {
        uword[3] seed
        cx16.r0 = seed[0] + seed[1] + seed[2]
    }
}"""
        compileText(C64Target(), false, text, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("ast result from compileText") {
        val text="""
main {
    sub start() {
        uword[3] seed
        cx16.r0 = seed[0] + seed[1] + seed[2]
    }
}"""
        val result = compileText(C64Target(), false, text, outputDir, writeAssembly = true)!!
        result.compilerAst.name shouldStartWith  "on_the_fly"
        val ast = result.codegenAst!!
        ast.name shouldBe result.compilerAst.name
        ast.children.size shouldBeGreaterThan 2
        val start = ast.entrypoint()!!
        start.name shouldBe "p8s_start"
        start.children.size shouldBeGreaterThan 2
        val seed = start.children[1] as PtVariable
        seed.name shouldBe "p8v_seed"
        seed.value shouldBe null
        seed.type shouldBe DataType.arrayFor(BaseDataType.UWORD)
        val assign = start.children[2] as PtAssignment
        assign.target.identifier!!.name shouldBe "cx16.r0"
        assign.value shouldBe instanceOf<PtBinaryExpression>()

        val st = result.codegenSymboltable!!
        st.flat.size shouldBeGreaterThan 100
        st.flat["cbm.CHROUT"]?.type shouldBe StNodeType.EXTSUB
        st.lookup("cbm.CHROUT")?.type shouldBe StNodeType.EXTSUB
    }

    test("peek and poke argument types") {
        val text="""
main {
    sub start() {
        uword[3] arr
        ubyte @shared i = 42
        uword @shared ww = peekw(arr[i])
        ubyte @shared xx = peek(arr[i])
        xx = @(arr[i])
        
        @(arr[i]) = 42
        poke(arr[i], 42)
        pokew(arr[i], 4242)
    }
}"""
        compileText(C64Target(), false, text, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("assigning memory byte into arrays works") {
        val text="""
main {
    sub start() {
        uword @shared factor1
        ubyte[3] bytearray
        uword[3] wordarray
        @(factor1) = bytearray[0]
        bytearray[0] = @(factor1)
        @(factor1) = lsb(wordarray[0])
        wordarray[0] = @(factor1)
        @(5000) = bytearray[0]
        @(5000) = lsb(wordarray[0])
        bytearray[0] = @(5000)
        wordarray[0] = @(5000)
    }
}"""
        compileText(C64Target(), false, text, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("reading memory from unknown var gives proper error") {
        val text="""
main {
    sub start() {
        cx16.r0L = @(doesnotexist)
    }
}"""
        val errors = ErrorReporterForTests()
        compileText(C64Target(), false, text, outputDir, writeAssembly = true, errors = errors)
        errors.errors.size shouldBe 2
        errors.errors[0] shouldContain "invalid address type"
        errors.errors[1] shouldContain "undefined symbol: doesnotexist"
    }

    test("shifting by word value is ok") {
        val text="""
main {
    sub start() {
        ubyte @shared c = 1
        @(15000 + c<<${'$'}0003) = 42
        @(15000 + (c<<${'$'}0003)) = 42
        @(15000 + c*${'$'}0008) = 42       ; *8 becomes a shift after opt

        uword @shared qq = 15000 + c<<${'$'}0003
        qq = 15000 + (c<<${'$'}0003)
        qq = 16000 + c*${'$'}0008
    }
}"""
        compileText(C64Target(), true, text, outputDir, writeAssembly = true) shouldNotBe null
        compileText(VMTarget(), true, text, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("builtin func in float expression") {
        val src="""
%import floats
main {
    sub start() {
        float @shared fl =25.1
        fl = abs(fl)+0.5
        fl = sqrt(fl)+0.5
    }
}"""
        compileText(C64Target(), false, src, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("string vars in inlined subroutines are ok") {
        val src="""
main {
    sub start() {
        void block2.curdir()
        void block2.other()
    }
}

block2 {
    str result="zzzz"
    sub curdir() -> str {
        return result
    }

    sub other() -> str {
        return "other"
    }
}"""

        compileText(C64Target(), true, src, outputDir, writeAssembly = true) shouldNotBe null
        compileText(VMTarget(), true, src, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("array with pointers") {
        val src = """
main {
    sub start() {
        str localstr = "hello"
        ubyte[] otherarray = [1,2,3]
        uword[] words = [1111,2222,"three",&localstr,&otherarray]
        uword @shared zz = &words
        bool @shared result = 2222 in words
        zz = words[2]
        zz++
        zz = words[3]
    }
}"""
        val othertarget = Cx16Target()
        compileText(othertarget, true, src, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("case sensitive symbols") {
        val src = """
main {
    sub start() {
        ubyte bytevar = 11      ; var at 0
        ubyte byteVAR = 22      ; var at 1
        ubyte ByteVar = 33      ; var at 2
        ubyte @shared total = bytevar+byteVAR+ByteVar   ; var at 3
        goto skipLABEL
SkipLabel:
        return
skipLABEL:
        bytevar = 42
    }
}"""
        val target = Cx16Target()
        compileText(target, true, src, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("addresses from labels/subroutines") {
        val src = """
main {
    sub start() {

mylabel:
        ubyte @shared variable
        uword @shared pointer1 = &main.start
        uword @shared pointer2 = &start
        uword @shared pointer3 = &main.start.mylabel
        uword @shared pointer4 = &mylabel
        uword[] @shared ptrs = [&variable, &start, &main.start, &mylabel, &main.start.mylabel]
    }
}

"""
        compileText(Cx16Target(), true, src, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("duplicate symbols okay other block and variable") {
        val src = """
main {
    ubyte derp
    sub start() {
        derp++
        foo.bar()
    }
}

foo {
    sub bar() {
        derp.print()
    }
}

derp {
    sub print() {
        cx16.r0++
        cx16.r1++
    }
}"""

        compileText(VMTarget(), false, src, outputDir, writeAssembly = true) shouldNotBe null
        compileText(Cx16Target(), false, src, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("ambiguous symbol name variable vs block") {
        val src = """
main {
    sub start() {
        uword module
        module++
        module.test++
    }
}

module {
    ubyte @shared test
}
"""
        val errors=ErrorReporterForTests()
        compileText(VMTarget(), false, src, outputDir, writeAssembly = false, errors = errors) shouldBe null
        errors.errors.size shouldBe 1
        errors.errors[0] shouldContain "ambiguous symbol"
    }

    test("prefix expressions with typecasting") {
        val src="""
main
{
    sub start()
    {
        uword @shared uw = 54321
        ubyte @shared ub = 123
        word @shared sw = -12345
        byte @shared sb = -123

        func_uw(~ub as uword) 
        func_ub(~uw as ubyte) 
        func_uw(~sb as uword) 
        func_ub(~sw as ubyte) 
        func_w(-sb as word)   
        func_b(-sw as byte)   
    }
    
    sub func_uw(uword arg) {
    }
    sub func_w(word arg) {
    }
    sub func_ub(ubyte arg) {
    }
    sub func_b(byte arg) {
    }
}"""

        compileText(VMTarget(), false, src, outputDir, writeAssembly = true) shouldNotBe null
        compileText(Cx16Target(), false, src, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("inlining sub with 2 statements") {
        val src="""
main {
    sub start() {
        init()
    }

    sub init() {
        init_handler()
        return
    }

    sub init_handler() {
        cx16.r0++
    }
}"""
        compileText(VMTarget(), true, src, outputDir, writeAssembly = false) shouldNotBe null
        compileText(Cx16Target(), true, src, outputDir, writeAssembly = false) shouldNotBe null
    }

    test("syslib correctly available for raw outputs") {
        val text = """
%output raw
%launcher none
%address ${'$'}2000

main {
    sub start() {
        cx16.r0++
        sys.clear_carry()
    }
}
"""
        compileText(Cx16Target(), false, text, outputDir, writeAssembly = true) shouldNotBe null
        compileText(C64Target(), false, text, outputDir, writeAssembly = true) shouldNotBe null
        compileText(C128Target(), false, text, outputDir, writeAssembly = true) shouldNotBe null
        compileText(PETTarget(), false, text, outputDir, writeAssembly = true) shouldNotBe null
        compileText(VMTarget(), false, text, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("block start address must be greater than program load address") {
        val src = """
%output raw
%launcher none
%address ${'$'}2000

main $2000 {
    sub start() {
        sys.clear_carry()
    }
}

otherblock ${'$'}2013 {
    %option force_output
}

thirdblock ${'$'}2014 {
    %option force_output
}"""
        val errors = ErrorReporterForTests()
        compileText(C64Target(), false, src, outputDir, writeAssembly = false, errors = errors) shouldBe null
        errors.errors.size shouldBe 2
        errors.errors[0] shouldContain "6:1: block address must be at least program load address + 20"
        errors.errors[1] shouldContain "12:1: block address must be at least program load address + 20"
    }

    test("for loops with just 1 iteration") {
        val src="""
main {
    sub start() {
        for cx16.r0L in 100 downto 100 {
            cx16.r1++
        }
        for cx16.r0L in 100 to 100 {
            cx16.r1++
        }
        for cx16.r0 in 2222 downto 2222 {
            cx16.r1++
        }
        for cx16.r0 in 2222 to 2222 {
            cx16.r1++
        }

        for cx16.r0L in 100 downto 100 step -5 {
            cx16.r1++
        }
        for cx16.r0L in 100 to 100 step 5 {
            cx16.r1++
        }
        for cx16.r0 in 2222 downto 2222 step -5 {
            cx16.r1++
        }
        for cx16.r0 in 2222 to 2222 step 5 {
            cx16.r1++
        }
    }
}"""
        compileText(VMTarget(), true, src, outputDir, writeAssembly = true) shouldNotBe null
        compileText(VMTarget(), false, src, outputDir, writeAssembly = true) shouldNotBe null
        compileText(Cx16Target(), true, src, outputDir, writeAssembly = true) shouldNotBe null
        compileText(Cx16Target(), false, src, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("multiple status flags return values from asmsub") {
        val src="""
main {
    extsub 5000 = carryAndNegativeAndByteAndWord() -> bool @Pc, bool @Pn, ubyte @X, uword @AY

    sub start() {
        ubyte @shared x
        uword @shared w
        bool @shared flag1
        bool @shared flag2

        flag1, flag2, x, w = carryAndNegativeAndByteAndWord()
        flag1, void, void, w = carryAndNegativeAndByteAndWord()
        void, void, void, void = carryAndNegativeAndByteAndWord()
        void carryAndNegativeAndByteAndWord()
    }    
}"""
        compileText(Cx16Target(), false, src, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("extsub with float return in multi-assign") {
        val src="""
%import floats

main {
    extsub 5000 = get_float_and_byte() -> float @FAC1, ubyte @X

    sub start() {
        float @shared f
        ubyte @shared b

        f, b = get_float_and_byte()
    }
}"""
        val errors = ErrorReporterForTests()
        compileText(Cx16Target(), false, src, outputDir, writeAssembly = true, errors = errors) shouldNotBe null
        errors.errors.size shouldBe 0
    }

    test("extsub with mixed returns including float") {
        // Test with multiple return types including FP, status flags, and CPU registers
        val src="""
%import floats

main {
    extsub 5000 = get_multi() -> bool @Pc, float @FAC1, ubyte @X, uword @AY

    sub start() {
        bool @shared flag
        float @shared f
        ubyte @shared b
        uword @shared w

        flag, f, b, w = get_multi()
        void, f, void, void = get_multi()
    }
}"""
        val errors = ErrorReporterForTests()
        compileText(Cx16Target(), false, src, outputDir, writeAssembly = true, errors = errors) shouldNotBe null
        errors.errors.size shouldBe 0
    }

    test("regular sub with float return in multi-assign") {
        // Test that regular Prog8 subroutines with float returns work correctly
        // Float can be in any position now - the compiler ensures FAC1 is assigned last
        // to avoid being clobbered by other return value computations
        val src="""
%import floats

main {
    sub start() {
        bool flag
        ubyte b
        float f
        flag, b, f = multi_test(42.0)
    }

    sub multi_test(float input) -> bool, ubyte, float {
        return true, input as ubyte, 3.14
    }
}"""
        val errors = ErrorReporterForTests()
        compileText(Cx16Target(), false, src, outputDir, writeAssembly = true, errors = errors) shouldNotBe null
        errors.errors.size shouldBe 0
    }

    test("regular sub with float in middle of multi-return") {
        // Test float as the middle return value - this used to fail because
        // computing the ubyte return value clobbered FAC1
        val src="""
%import floats

main {
    sub start() {
        bool flag
        float f
        ubyte b
        flag, f, b = multi_test(42.0)
    }

    sub multi_test(float input) -> bool, float, ubyte {
        return true, 3.14, input as ubyte
    }
}"""
        val errors = ErrorReporterForTests()
        compileText(Cx16Target(), false, src, outputDir, writeAssembly = true, errors = errors) shouldNotBe null
        errors.errors.size shouldBe 0
    }

    test("regular sub with float first in multi-return") {
        // Test float as the first return value
        val src="""
%import floats

main {
    sub start() {
        float f
        bool flag
        ubyte b
        f, flag, b = multi_test(42.0)
    }

    sub multi_test(float input) -> float, bool, ubyte {
        return 3.14, true, input as ubyte
    }
}"""
        val errors = ErrorReporterForTests()
        compileText(Cx16Target(), false, src, outputDir, writeAssembly = true, errors = errors) shouldNotBe null
        errors.errors.size shouldBe 0
    }

    test("missing rts in asmsub") {
        val src="""
main {
    sub start() {
        test()
        test2()
    }

    asmsub test() {
        %asm {{
            nop
            nop
        }}
    }

    inline asmsub test2() {
        %asm {{
            nop
            nop
        }}
    }
}"""

        val errors = ErrorReporterForTests()
        compileText(C64Target(), false, src, outputDir, writeAssembly = true, errors = errors) shouldBe null
        errors.errors.size shouldBe 1
        errors.errors[0] shouldContain "asmsub seems to never return"
    }

    test("missing rts in asmsub suppressed") {
        val src="""
main {
    sub start() {
        test()
    }

    asmsub test() {
        %asm {{
            nop
            nop
            ; !notreached!
        }}
    }
}"""

        compileText(C64Target(), false, src, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("if not without else is not swapped") {
        val src="""
main {
    sub start() {
        if not thing()
            cx16.r0++
    }

    sub thing() -> bool {
        cx16.r0++
        return false
    }
}"""
        compileText(C64Target(), false, src, outputDir, writeAssembly = true) shouldNotBe null
        val result = compileText(VMTarget(), false, src, outputDir, writeAssembly = true)!!
        val st = result.codegenAst!!.entrypoint()!!.children
        st.size shouldBe 3
        val ifelse = st[1] as PtIfElse
        ifelse.hasElse() shouldBe false
        (ifelse.condition as PtPrefix).operator shouldBe "not"
    }

    test("if not with else is swapped") {
        val src="""
main {
    sub start() {
        if not thing()
            cx16.r0++
        else
            cx16.r1++
    }

    sub thing() -> bool {
        cx16.r0++
        return false
    }
}"""
        compileText(C64Target(), false, src, outputDir, writeAssembly = true) shouldNotBe null
        val result = compileText(VMTarget(), false, src, outputDir, writeAssembly = true)!!
        val st = result.codegenAst!!.entrypoint()!!.children
        st.size shouldBe 3
        val ifelse = st[1] as PtIfElse
        ifelse.hasElse() shouldBe true
        ifelse.condition shouldBe instanceOf<PtFunctionCall>()
    }

    test("bit instruction is used for testing bits 6 and 7 of a byte") {
        val text = """
main {
    sub start() {   
        if cx16.r0L & ${'$'}80 != 0
            return
        if cx16.r1L & ${'$'}80 == 0
            return
        if cx16.r2L & ${'$'}40 != 0
            return
        if cx16.r3L & ${'$'}40 == 0
            return
        cx16.r9L = if cx16.r4L & ${'$'}80 != 0  11 else 22
        cx16.r10L = if cx16.r5L & ${'$'}40 == 0  11 else 22
    } 
}"""
        val result = compileText(C64Target(), true, text, outputDir, writeAssembly = true)!!
        val assemblyFile = result.compilationOptions.outputDir.resolve(result.compilerAst.name + ".asm")
        val assembly = assemblyFile.readText()
        assembly.shouldContainInOrder(
            "bit  cx16.r0L",
            "bit  cx16.r1L",
            "bit  cx16.r2L",
            "bit  cx16.r3L",
            "bit  cx16.r4L",
            "bit  cx16.r5L"
        )

        val resultIR = compileText(VMTarget(), true, text, outputDir, writeAssembly = true)!!
        val irFile = resultIR.compilationOptions.outputDir.resolve(result.compilerAst.name + ".p8ir")
        val ir = irFile.readText()
        ir.shouldContainInOrder(
            "bit.b ${'$'}ff02",     // r0
            "bit.b ${'$'}ff04",     // r1
            "bit.b ${'$'}ff06",     // f2
            "bit.b ${'$'}ff08",     // r3
            "bit.b ${'$'}ff0a",     // r4
            "bit.b ${'$'}ff0c"      // r5
        )
    } 

    test("strings in if expression") {
        val src="""
main {
    sub start() {
       str foo = "foo"
       str bar = "bar"
       bool flag = true
       uword @shared foobar = if flag foo else bar
    }
}"""
        compileText(C64Target(), false, src, outputDir, writeAssembly = true) shouldNotBe null
        compileText(VMTarget(), false, src, outputDir, writeAssembly = true) shouldNotBe null
    }

    test("word bitshift with byte operand") {
        val src="""
main{
    sub start() {
        cx16.r0 >>= 4
        cx16.r1 <<= 4
    }
}"""

        compileText(C64Target(), false, src, outputDir) shouldNotBe null
        compileText(VMTarget(), false, src, outputDir) shouldNotBe null
        compileText(C64Target(), true, src, outputDir) shouldNotBe null
        compileText(VMTarget(), true, src, outputDir) shouldNotBe null
    }

    test("symbol prefixing inconsistency with struct types and merged blocks") {
        // see https://github.com/irmen/prog8/issues/198
        val src="""
plane {
   struct Point {
       ubyte x
       ubyte y
   }
}

mytxt {
    %option no_symbol_prefixing

    sub print_ub(ubyte value) {
        cx16.r0L++
    }
}

mytxt {
    %option merge
    sub print_pt(^^plane.Point p) {
        mytxt.print_ub(p.x)
        mytxt.print_ub(p.y)
    }
}

main {
    sub start() {
        ^^plane.Point origin = ^^plane.Point:[0,0]
        mytxt.print_pt(origin)
    }
}"""

        compileText(C64Target(), false, src, outputDir) shouldNotBe null
    }

    test("long argument to asmsub") {
        val src="""
main {

    sub start() {
        long @shared lv = 1234567
        test(999999)
        test(lv)
        test(0)
    }
    
    asmsub test(long arg @R0R1) {
        %asm {{
            rts
        }}
    }
}"""
        compileText(C64Target(), false, src, outputDir, writeAssembly = true)!!
        compileText(Cx16Target(), false, src, outputDir, writeAssembly = true)!!
    }

    test("long argument and another to asmsub") {
        val src= $$"""
main {
    sub start() {
        extcommand_print_ulhex($12345678, true)
        extcommand_print_ulhex($abcdef99, false)
    }

    asmsub extcommand_print_ulhex(long value @R0R1, bool prefix @A) clobbers(A,X,Y) {
        %asm {{
            rts
        }}
    }

}"""
        compileText(C64Target(), false, src, outputDir, writeAssembly = true)!!
        compileText(Cx16Target(), false, src, outputDir, writeAssembly = true)!!
    }

    test("returning multivalue functioncall that can be short-circuited") {
        val src= """
main {

    sub start() {
        ubyte a,b = multi()
    }

    sub multi() -> ubyte, ubyte {
        cx16.r0++
        return multi2()
    }

    sub multi2() -> ubyte, ubyte {
        cx16.r0++
        return 42, 99
    }
}"""

        compileText(Cx16Target(), false, src, outputDir) shouldNotBe null
        compileText(VMTarget(), false, src, outputDir) shouldNotBe null
    }

    test("returning multivalue functioncall that cannot be short-circuited") {
        val src= """
main {

    sub start() {
        ubyte a,b = multi()
    }

    sub multi() -> ubyte, ubyte {
        cx16.r0++
        return multi2(99)
    }

    sub multi2(ubyte x) -> ubyte, ubyte {
        x++
        return 42, 99
    }
}"""

        compileText(Cx16Target(), false, src, outputDir) shouldNotBe null
        compileText(VMTarget(), false, src, outputDir) shouldNotBe null
    }

    test("efficient 6502 array indexing code for s[index+const]") {
        val source = """
            %zeropage basicsafe
            main {
                sub rstrip(str s) {
                    s[cx16.r0L+1] = 0
                }
                sub start() {
                    rstrip("hello")
                }
            }
        """.trimIndent()

        val result = compileText(Cx16Target(), true, source, outputDir)!!

        // 1. Check SimpleAst (codegenAst)
        // We expect operand order optimization to have swapped 's' and '(cx16.r0L + 1)'
        // so that the pointer 's' is on the right, because that is the least complex operand

        val additions = mutableListOf<PtBinaryExpression>()
        walkAst(result.codegenAst!!) { node, _ ->
            if (node is PtBinaryExpression && node.operator == "+" && !node.position.file.contains("library:")) {
                additions.add(node)
            }
            true
        }

        // In the SimpleAst for this program, we expect:
        // 1. cx16.r0L + 1
        // 2. (cx16.r0L + 1) + main.rstrip.s

        additions.size shouldBe 2

        val outerAddition = additions.find { it.type.isPointer || it.type.isUnsignedWord }!!

        withClue("Pointer 's' should be on the right due to operand order optimization") {
            val right = outerAddition.right
            right.shouldBeInstanceOf<PtIdentifier>()
            right.name shouldContain "p8v_s"
        }

        // 2. Check Assembly
        val asmFile = result.compilationOptions.outputDir.resolve(result.compilerAst.name + ".asm")
        val asm = asmFile.readText()

        withClue("Should use efficient sta (zp),y instruction") {
            asm shouldContain "sta  (p8v_s),y"
        }
    }
})