variable zp allocation now only done in the allocator

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

@@ -1,6 +1,5 @@
 package prog8.codegen.cpu6502
 
-import com.github.michaelbull.result.fold
 import prog8.ast.IFunctionCall
 import prog8.ast.Program
 import prog8.ast.antlr.escape
@@ -164,9 +163,9 @@ internal class ProgramGen(
         asmgen.outputSourceLine(block)
 
         val vardecls = block.statements.filterIsInstance<VarDecl>()
-        zeropagevars2asm(vardecls, block)
+        zeropagevars2asm(vardecls)
         memdefs2asmVars(vardecls, block)
-        memdefs2asmAsmsubs(block.statements.filterIsInstance<Subroutine>(), block)
+        memdefs2asmAsmsubs(block.statements.filterIsInstance<Subroutine>())
         vardecls2asm(vardecls, block)
 
         vardecls.forEach {
@@ -226,9 +225,9 @@ internal class ProgramGen(
             // regular subroutine
             asmgen.out("${sub.name}\t.proc")
             val vardecls = sub.statements.filterIsInstance<VarDecl>()
-            zeropagevars2asm(vardecls, null)
+            zeropagevars2asm(vardecls)
             memdefs2asmVars(vardecls, null)
-            memdefs2asmAsmsubs(sub.statements.filterIsInstance<Subroutine>(), null)
+            memdefs2asmAsmsubs(sub.statements.filterIsInstance<Subroutine>())
 
             // the main.start subroutine is the program's entrypoint and should perform some initialization logic
             if(sub.name=="start" && sub.definingBlock.name=="main")
@@ -349,40 +348,25 @@ internal class ProgramGen(
                     clc""")
     }
 
-    private fun zeropagevars2asm(vardecls: List<VarDecl>, inBlock: Block?) {
+    private fun zeropagevars2asm(vardecls: List<VarDecl>) {
+        // TODO get list of variables directly from allocator or zeropage
         val zp = zeropage
         asmgen.out("; vars allocated on zeropage")
-        val variables = vardecls.filter { it.type==VarDeclType.VAR }
-        val blockname = inBlock?.name
-        for(variable in variables) {
-            if(blockname=="prog8_lib" && variable.name.startsWith("P8ZP_SCRATCH_"))
-                continue       // the "hooks" to the temp vars are not generated as new variables
-            val scopedName = variable.scopedName
-            val zpAlloc = zp.variables[scopedName]
-            if (zpAlloc == null) {
-                // TODO NO LONGER ALLOCATE HERE, IT'S ALL BEEN DONE IN THE VARIABLEALLOCATOR ALREADY
-                // This var is not on the ZP yet. Attempt to move it there if it's an integer type
-                if(variable.zeropage != ZeropageWish.NOT_IN_ZEROPAGE &&
-                    variable.datatype in IntegerDatatypes
-                    && options.zeropage != ZeropageType.DONTUSE) {
-                    val result = zp.allocate(scopedName, variable.datatype, null, null, null, errors)
-                    errors.report()
-                    result.fold(
-                        success = { (address, _) -> asmgen.out("${variable.name} = $address\t; zp ${variable.datatype}") },
-                        failure = { /* leave it as it is, not on zeropage. */ }
-                    )
+        vardecls
+            .filter { it.type==VarDeclType.VAR }
+            .forEach { variable ->
+                val scopedName = variable.scopedName
+                val zpAlloc = zp.variables[scopedName]
+                if (zpAlloc != null) {
+                    val lenspec = when(zpAlloc.dt) {
+                        DataType.FLOAT,
+                        DataType.STR,
+                        in ArrayDatatypes -> " ${zpAlloc.size} bytes"
+                        else -> ""
+                    }
+                    asmgen.out("${variable.name} = ${zpAlloc.address}\t; zp ${variable.datatype} $lenspec")
                 }
-            } else {
-                // Var has been placed in ZP, just output the address
-                val lenspec = when(zpAlloc.dt) {
-                    DataType.FLOAT,
-                    DataType.STR,
-                    in ArrayDatatypes -> " ${zpAlloc.size} bytes"
-                    else -> ""
-                }
-                asmgen.out("${variable.name} = ${zpAlloc.address}\t; zp ${variable.datatype} $lenspec")
             }
-        }
     }
 
     private fun vardecl2asm(decl: VarDecl, nameOverride: String?=null) {
@@ -496,7 +480,7 @@ internal class ProgramGen(
             }
     }
 
-    private fun memdefs2asmAsmsubs(subroutines: List<Subroutine>, inBlock: Block?) {
+    private fun memdefs2asmAsmsubs(subroutines: List<Subroutine>) {
         subroutines
             .filter { it.isAsmSubroutine }
             .forEach { sub->

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

@@ -3,6 +3,7 @@ package prog8.codegen.cpu6502
 import com.github.michaelbull.result.onFailure
 import prog8.ast.base.ArrayDatatypes
 import prog8.ast.base.DataType
+import prog8.ast.base.IntegerDatatypes
 import prog8.ast.expressions.StringLiteralValue
 import prog8.ast.statements.Subroutine
 import prog8.ast.statements.VarDecl
@@ -37,44 +38,6 @@ internal class VariableAllocator(val vars: IVariablesAndConsts, val errors: IErr
             .filterNot { it.second.last() == "tbl" }   // TODO HACK -- NOT REALLY NECESSARY, BUT OLD CallGraph DIDN'T CONTAIN IT EITHER
             .filterNot { it.second.last() in setOf("retval_interm_w",  "retval_interm_b", "retval_interm_w2", "retval_interm_b2") }   // TODO HACK TO REMOVE THESE UNUSED VARS
 
-        val varsRequiringZp = allVariables
-            .filter { it.first.zeropage == ZeropageWish.REQUIRE_ZEROPAGE }
-        val varsPreferringZp = allVariables
-            .filter { it.first.zeropage == ZeropageWish.PREFER_ZEROPAGE }
-            .sortedBy { options.compTarget.memorySize(it.first.datatype) }      // allocate the smallest DT first
-        val varsDontCare = allVariables
-            .filter { it.first.zeropage == ZeropageWish.DONTCARE }
-            .sortedBy { options.compTarget.memorySize(it.first.datatype) }      // allocate the smallest DT first
-
-/*
-        // OLD CODE CHECKING:
-        if(true) {
-            val allVariablesFoundInCallgraph = callGraphForCheck.allIdentifiers.asSequence()
-                .map { it.value }
-                .filterIsInstance<VarDecl>()
-                .filter { it.type == VarDeclType.VAR && it.origin != VarDeclOrigin.SUBROUTINEPARAM }
-                .map { it.name to it.position }
-                .toSet()
-            val newAllVars = (vars.blockVars.flatMap { it.value }
-                .map { it.name to it.position } + vars.subroutineVars.flatMap { it.value }
-                .map { it.name to it.position }).toSet()
-            val extraVarsInCallgraph = allVariablesFoundInCallgraph - newAllVars
-            val extraVarsInNew = newAllVars - allVariablesFoundInCallgraph
-
-            if (extraVarsInCallgraph.any() || extraVarsInNew.any()) {
-                println("EXTRA VARS IN CALLGRAPH: ${extraVarsInCallgraph.size}")
-                extraVarsInCallgraph.forEach {
-                    println("  $it")
-                }
-                println("EXTRA VARS IN VARIABLESOBJ: ${extraVarsInNew.size}")
-                extraVarsInNew.forEach {
-                    println("  $it")
-                }
-                //TODO("fix differences")
-            }
-        }
-*/
-
         val zeropage = options.compTarget.machine.zeropage
 
         fun numArrayElements(vardecl: VarDecl): Int? = when(vardecl.datatype) {
@@ -87,6 +50,10 @@ internal class VariableAllocator(val vars: IVariablesAndConsts, val errors: IErr
             else -> null
         }
 
+        val varsRequiringZp = allVariables.filter { it.first.zeropage == ZeropageWish.REQUIRE_ZEROPAGE }
+        val varsPreferringZp = allVariables.filter { it.first.zeropage == ZeropageWish.PREFER_ZEROPAGE }
+        val varsDontCare = allVariables.filter { it.first.zeropage == ZeropageWish.DONTCARE }
+
         varsRequiringZp.forEach { (vardecl, scopedname) ->
             val numElements = numArrayElements(vardecl)
             val result = zeropage.allocate(scopedname, vardecl.datatype, numElements, vardecl.value, vardecl.position, errors)
@@ -100,14 +67,18 @@ internal class VariableAllocator(val vars: IVariablesAndConsts, val errors: IErr
                 //  no need to check for allocation error, if there is one, just allocate in normal system ram.
             }
 
-            // TODO enable zp allocation here and remove it from other places
-//            if(errors.noErrors()) {
-//                varsDontCare.forEach { (vardecl, scopedname) ->
-//                    val numElements = numArrayElements(vardecl)
-//                    zeropage.allocate(scopedname, vardecl.datatype, numElements, vardecl.value, vardecl.position, errors)
-//                    //  no need to check for allocation error, if there is one, just allocate in normal system ram.
-//                }
-//            }
+            // try to allocate any other interger variables into the zeropage until it is full.
+            // TODO some form of intelligent priorization? most often used variables first? loopcounter vars first? ...?
+            if(errors.noErrors()) {
+                for ((vardecl, scopedname) in varsDontCare) {
+                    if(vardecl.datatype in IntegerDatatypes) {
+                        val numElements = numArrayElements(vardecl)
+                        zeropage.allocate(scopedname, vardecl.datatype, numElements, vardecl.value, vardecl.position, errors)
+                        if(zeropage.free.isEmpty())
+                            break
+                    }
+                }
+            }
         }
     }
 

docs/source/todo.rst

@@ -6,6 +6,9 @@ For next release
 - programGen: don't generate variables from the VarDecl nodes, use allocator/zeropage tables
 - (newvaralloc) UnusedCodeRemover after(decl: VarDecl): fix that vars defined in a library can also safely be removed if unused. Currently this breaks programs such as textelite (due to diskio.save().end_address ?)
 - remove hacks in VariableAllocator
+- make it so that subroutine parameters as variables can again be allocated in ZP, if there's still space
+- wormfood became a lot larger??? why???  (and chess a little bit larger, but usually program size is down)
+- check: are prog8_lib.P8ZP_SCRATCH emitted again from prog8_lib.p8 source, or do we only have the correct hardcoded redefines.
 
 
 Need help with

examples/test.p8

@@ -17,6 +17,8 @@ main {
     str @requirezp zpname = "irmenzp"
 
     sub start() {
+        prog8_lib.P8ZP_SCRATCH_B1 = 1
+
         txt.print_uw(nullwords[1])
         txt.nl()
         txt.print_ub(nullbytes[1])