prog8/compiler/test/codegeneration/Test2DArrays.kt

package prog8tests.codegeneration

import io.kotest.assertions.fail
import io.kotest.core.spec.style.FunSpec
import io.kotest.engine.spec.tempdir
import io.kotest.matchers.shouldBe
import io.kotest.matchers.shouldNotBe
import io.kotest.matchers.string.shouldContain
import prog8.code.ast.PtVariable
import prog8.code.ast.walkAst
import prog8.code.target.C64Target
import prog8.code.target.VMTarget
import prog8tests.helpers.ErrorReporterForTests
import prog8tests.helpers.compileText
import prog8tests.helpers.shouldContainInOrder
import kotlin.io.path.readText

class Test2DArrays: FunSpec({
    val outputDir = tempdir().toPath()

    test("2D array declaration and indexing compiles on multiple targets") {
        // Covers: constant indexing, variable indexing, expressions, address-of, function args, optimized
        val sources = listOf(
            // Constant indexing
            """
            main {
                sub start() {
                    ubyte[3][4] m
                    m[0][0] = 1
                    m[2][3] = 12
                    ubyte @shared v = m[0][0]
                }
            }""",
            // Variable indexing
            """
            main {
                sub start() {
                    ubyte[3][4] m
                    ubyte @shared r = 1
                    ubyte @shared c = 2
                    m[r][c] = 42
                    ubyte @shared v = m[r][c]
                }
            }""",
            // Used in expressions
            """
            main {
                sub start() {
                    ubyte[2][3] m
                    m[0][0] = 10
                    m[1][1] = 15
                    ubyte @shared s = m[0][0] + m[1][1]
                }
            }""",
            // As function argument
            """
            main {
                sub start() {
                    ubyte[2][3] m
                    m[1][1] = 42
                    process(m[0][0])
                }
                sub process(ubyte v) { }
            }""",
            // Assignment from expression
            """
            main {
                sub start() {
                    ubyte[2][3] src
                    ubyte[2][3] dst
                    src[1][2] = 99
                    dst[1][2] = src[1][2]
                }
            }""",
            // In for loop
            """
            main {
                ubyte[4][4] m
                sub start() {
                    ubyte @shared r = 0
                    ubyte @shared c = 0
                    for r in 0 to 3 {
                        for c in 0 to 3 {
                            m[r][c] = r * c
                        }
                    }
                }
            }""",
            // Address-of operation
            """
            main {
                sub start() {
                    ubyte[3][4] m
                    pointer @shared addr = &m
                }
            }"""
        )
        for(src in sources) {
            val result = compileText(VMTarget(), false, src, outputDir) ?: fail("Compilation failed for source:\n$src")
            val resultc64 = compileText(C64Target(), false, src, outputDir) ?: fail("Compilation failed for source:\n$src")
        }
        // Also verify optimized compilation
        compileText(VMTarget(), true, sources[0], outputDir) shouldNotBe null
    }

    test("2D array with different element types compiles") {
        val types = listOf(
            "byte[2][3] m; m[0][0] = -10; byte v = m[0][0]",
            "word[2][3] m; m[0][0] = -1000; word v = m[0][0]",
            "uword[2][3] m; m[0][0] = 1000; uword v = m[0][0]"
        )
        for(body in types) {
            val src = """
            main {
                sub start() {
                    $body
                }
            }"""
            compileText(VMTarget(), false, src, outputDir) shouldNotBe null
            compileText(C64Target(), false, src, outputDir) shouldNotBe null
        }
    }

    test("2D array desugaring and constant folding produce correct flat index") {
        // matrix[row][col] should desugar to matrix[row * numCols + col]
        // For ubyte[3][4]: matrix[1][2] -> index 1*4+2=6, matrix[2][3] -> index 2*4+3=11
        
        // Verify assembly contains correct indexing pattern
        val src1 = """
        main {
            sub start() {
                ubyte[3][4] m
                m[1][2] = 42
            }
        }"""
        val result1 = compileText(C64Target(), false, src1, outputDir, writeAssembly = true)!!
        val asmFile1 = result1.compilationOptions.outputDir.resolve(result1.compilerAst.name + ".asm")
        val asm1 = asmFile1.readText()
        asm1.shouldContainInOrder("m")  // Should reference the array
        
        // Verify IR after constant folding
        val src2 = """
        main {
            ubyte[3][4] m
            sub start() {
                m[2][3] = 100
            }
        }"""
        val result2 = compileText(VMTarget(), true, src2, outputDir, writeAssembly = true)!!
        val virtfile = result2.compilationOptions.outputDir.resolve(result2.compilerAst.name + ".p8ir")
        val ir = virtfile.readText()
        ir.shouldContainInOrder("m")  // Should contain the array reference
    }

    test("2D array dimensions with const identifiers flatten to correct total size in Simple AST") {
        // Mirrors examples/test.p8: const dims must be folded during processAst(),
        // Simple AST should only contain the consolidated 1D sizes.
        val src = """
            main {
                const WIDTH=5
                const HEIGHT=8

                sub start() {
                    ubyte[WIDTH] @shared tokens
                    bool[2][4] @shared matrix1
                    bool[2][HEIGHT] @shared matrix2
                    bool[WIDTH][HEIGHT] @shared matrix3
                }
            }"""
        val result = compileText(VMTarget(), false, src, outputDir) ?: fail("Compilation failed")
        val expected = mapOf(
            "tokens" to 5u,
            "matrix1" to 8u,
            "matrix2" to 16u,
            "matrix3" to 40u
        )
        // 1. Check Simple AST (PtVariable.arraySize is the flattened rows*cols total)
        val ptVars = mutableMapOf<String, UInt?>()
        walkAst(result.codegenAst!!) { node, _ ->
            if(node is PtVariable)
                ptVars[node.name.substringAfterLast('.').removePrefix("p8v_")] = node.arraySize
            true
        }
        for((name, size) in expected)
            ptVars[name] shouldBe size
        // 2. Check Simple AST symbol table lengths agree
        for((name, size) in expected) {
            val stVar = result.codegenSymboltable!!.allVariables.first { it.name.endsWith(name) }
            stVar.length shouldBe size
        }
    }

    test("invalid 2D array usage produces compile errors") {
        // Chained indexing on 1D array
        val src1 = """
        main {
            sub start() {
                ubyte[12] arr
                ubyte x = arr[0][1]
            }
        }"""
        val errors1 = ErrorReporterForTests()
        compileText(VMTarget(), false, src1, outputDir, errors = errors1) shouldBe null
        errors1.errors.isNotEmpty() shouldBe true
        errors1.errors.any { it.contains("chained indexing requires the variable to be declared as a 2D array") } shouldBe true

        // Chained indexing on @shared uword pointer variable
        val srcShared = """
        main {
            sub start() {
                uword @shared ptr = 2222
                ubyte x = ptr[55][22]
            }
        }"""
        val errorsShared = ErrorReporterForTests()
        compileText(VMTarget(), false, srcShared, outputDir, errors = errorsShared) shouldBe null
        errorsShared.errors.size shouldBe 1
        errorsShared.errors[0] shouldContain "chained indexing requires the variable to be declared as a 2D array"

        // Chained indexing on regular uword pointer variable
        val srcUword = """
        main {
            sub start() {
                uword ptr = 3333
                ubyte x = ptr[55][22]
            }
        }"""
        val errorsUword = ErrorReporterForTests()
        compileText(VMTarget(), false, srcUword, outputDir, errors = errorsUword) shouldBe null
        errorsUword.errors.size shouldBe 1
        errorsUword.errors[0] shouldContain "chained indexing requires the variable to be declared as a 2D array"

        // 3D array declaration (grammar limit)
        val src2 = """
        main {
            sub start() {
                ubyte[3][4][5] cube
            }
        }"""
        compileText(VMTarget(), false, src2, outputDir) shouldBe null

        // 2D array as subroutine parameter
        val src3 = """
        main {
            sub process(ubyte[3][4] m) { }
            sub start() { }
        }"""
        compileText(VMTarget(), false, src3, outputDir) shouldBe null
    }
})