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() 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 } })