Allocating arbitrary variables in the zeropage and using zeropage addressing when appropriate

src/main/scala/millfork/Platform.scala

@@ -23,6 +23,7 @@ class Platform(
                 val outputPackager: OutputPackager,
                 val codeAllocators: Map[String, UpwardByteAllocator],
                 val variableAllocators: Map[String, VariableAllocator],
+                val freeZpPointers: List[Int],
                 val fileExtension: String,
                 val generateBbcMicroInfFile: Boolean,
                 val bankNumbers: Map[String, Int],
@@ -171,6 +172,7 @@ object Platform {
     new Platform(cpu, flagOverrides, startingModules, outputPackager,
       codeAllocators.toMap,
       variableAllocators.toMap,
+      freePointers,
       if (fileExtension == "" || fileExtension.startsWith(".")) fileExtension else "." + fileExtension,
       generateBbcMicroInfFile,
       bankNumbers,

src/main/scala/millfork/assembly/AssemblyOptimization.scala

@@ -9,5 +9,7 @@ import millfork.env.NormalFunction
 trait AssemblyOptimization[T <: AbstractCode] {
   def name: String
 
-  def optimize(f: NormalFunction, code: List[T], options: CompilationOptions): List[T]
+  def optimize(f: NormalFunction, code: List[T], options: CompilationOptions, labelMap: Map[String, Int]): List[T] = optimize(f, code, options)
+
+  def optimize(f: NormalFunction, code: List[T], options: CompilationOptions): List[T] = optimize(f, code, options, Map())
 }

src/main/scala/millfork/assembly/mos/opt/RuleBasedAssemblyOptimization.scala

@@ -32,18 +32,18 @@ class RuleBasedAssemblyOptimization(val name: String, val needsFlowInfo: FlowInf
 
   rules.foreach(_.pattern.validate(needsFlowInfo))
 
-  override def optimize(f: NormalFunction, code: List[AssemblyLine], options: CompilationOptions): List[AssemblyLine] = {
+  override def optimize(f: NormalFunction, code: List[AssemblyLine], options: CompilationOptions, labelMap: Map[String, Int]): List[AssemblyLine] = {
     val effectiveCode = code.map(a => a.copy(parameter = a.parameter.quickSimplify))
     val taggedCode = FlowAnalyzer.analyze(f, effectiveCode, options, needsFlowInfo)
-    optimizeImpl(f, taggedCode, options)
+    optimizeImpl(f, taggedCode, options, labelMap)
   }
 
-  def optimizeImpl(f: NormalFunction, code: List[(FlowInfo, AssemblyLine)], options: CompilationOptions): List[AssemblyLine] = {
+  def optimizeImpl(f: NormalFunction, code: List[(FlowInfo, AssemblyLine)], options: CompilationOptions, labelMap: Map[String, Int]): List[AssemblyLine] = {
     code match {
       case Nil => Nil
       case head :: tail =>
         for ((rule, index) <- rules.zipWithIndex) {
-          val ctx = new AssemblyMatchingContext(options)
+          val ctx = new AssemblyMatchingContext(options, labelMap)
           rule.pattern.matchTo(ctx, code) match {
             case Some(rest: List[(FlowInfo, AssemblyLine)]) =>
               val matchedChunkToOptimize: List[AssemblyLine] = code.take(code.length - rest.length).map(_._2)
@@ -59,19 +59,19 @@ class RuleBasedAssemblyOptimization(val name: String, val needsFlowInfo: FlowInf
               ErrorReporting.trace("     ↓")
               optimizedChunk.filter(_.isPrintable).foreach(l => ErrorReporting.trace(l.toString))
               if (needsFlowInfo != FlowInfoRequirement.NoRequirement) {
-                return optimizedChunk ++ optimizeImpl(f, rest, options)
+                return optimizedChunk ++ optimizeImpl(f, rest, options, labelMap)
               } else {
                 return optimize(f, optimizedChunk ++ rest.map(_._2), options)
               }
             case None => ()
           }
         }
-        head._2 :: optimizeImpl(f, tail, options)
+        head._2 :: optimizeImpl(f, tail, options, labelMap)
     }
   }
 }
 
-class AssemblyMatchingContext(val compilationOptions: CompilationOptions) {
+class AssemblyMatchingContext(val compilationOptions: CompilationOptions, val labelMap: Map[String, Int]) {
   private val map = mutable.Map[Int, Any]()
 
   override def toString: String = map.mkString(", ")
@@ -1020,3 +1020,29 @@ case class MatchElidableCopyOf(i: Int, firstLinePattern: AssemblyLinePattern, la
     Some(after)
   }
 }
+
+case object IsZeroPage extends AssemblyLinePattern {
+  override def matchLineTo(ctx: AssemblyMatchingContext, flowInfo: FlowInfo, line: AssemblyLine): Boolean = {
+    import Opcode._
+    line match {
+      case AssemblyLine(_, AddrMode.ZeroPage, _, _) => true
+      case l@AssemblyLine(LDA | STA | CMP |
+                          LDX | STX | CPX |
+                          LDY | STY | CPY |
+                          LDZ | STZ | CPZ |
+                          BIT |
+                          ADC | SBC | AND | ORA | EOR |
+                          INC | DEC | ROL | ROR | ASL | LSR |
+                          ISC | DCP | LAX | SAX | RLA | RRA | SLO | SRE, AddrMode.Absolute, p, true) =>
+        p match {
+          case NumericConstant(n, _) => n <= 255
+          case MemoryAddressConstant(th) => ctx.labelMap.getOrElse(th.name, 0x800) < 0x100
+          case CompoundConstant(MathOperator.Plus,
+          MemoryAddressConstant(th),
+          NumericConstant(n, _)) => ctx.labelMap.getOrElse(th.name, 0x800) + n < 0x100
+          case _ => false
+        }
+      case _ => false
+    }
+  }
+}

src/main/scala/millfork/assembly/mos/opt/SuperOptimizer.scala

@@ -13,7 +13,7 @@ import scala.collection.mutable
   */
 object SuperOptimizer extends AssemblyOptimization[AssemblyLine] {
 
-  def optimize(m: NormalFunction, code: List[AssemblyLine], options: CompilationOptions): List[AssemblyLine] = {
+  override def optimize(m: NormalFunction, code: List[AssemblyLine], options: CompilationOptions): List[AssemblyLine] = {
     val oldVerbosity = ErrorReporting.verbosity
     ErrorReporting.verbosity = -1
     var allOptimizers = OptimizationPresets.Good ++ LaterOptimizations.All

src/main/scala/millfork/env/Environment.scala

@@ -8,7 +8,7 @@ import millfork.assembly.mos.Opcode
 import millfork.assembly.z80.{IfFlagClear, IfFlagSet, ZFlag}
 import millfork.error.ErrorReporting
 import millfork.node._
-import millfork.output.{CompiledMemory, VariableAllocator}
+import millfork.output.{AllocationLocation, CompiledMemory, VariableAllocator}
 
 import scala.collection.mutable
 
@@ -89,7 +89,14 @@ class Environment(val parent: Option[Environment], val prefix: String) {
     }
   }
 
-  def allocateVariables(nf: Option[NormalFunction], mem: CompiledMemory, callGraph: CallGraph, allocators: Map[String, VariableAllocator], options: CompilationOptions, onEachVariable: (String, Int) => Unit): Unit = {
+  def allocateVariables(nf: Option[NormalFunction],
+                        mem: CompiledMemory,
+                        callGraph: CallGraph,
+                        allocators: Map[String, VariableAllocator],
+                        options: CompilationOptions,
+                        onEachVariable: (String, Int) => Unit,
+                        pass: Int,
+                        forZpOnly: Boolean): Unit = {
     val b = get[Type]("byte")
     val p = get[Type]("pointer")
     val params = nf.fold(List[String]()) { f =>
@@ -100,8 +107,13 @@ class Environment(val parent: Option[Environment], val prefix: String) {
           Nil
       }
     }.toSet
+    def passForAlloc(m: VariableAllocationMethod.Value) = m match {
+      case VariableAllocationMethod.Zeropage => 1
+      case VariableAllocationMethod.Register => 2
+      case _ => 3
+    }
     val toAdd = things.values.flatMap {
-      case m: UninitializedMemory if nf.isDefined == isLocalVariableName(m.name) && !m.name.endsWith(".addr") && maybeGet[Thing](m.name + ".array").isEmpty =>
+      case m: UninitializedMemory if passForAlloc(m.alloc) == pass && nf.isDefined == isLocalVariableName(m.name) && !m.name.endsWith(".addr") && maybeGet[Thing](m.name + ".array").isEmpty =>
         val vertex = if (options.flag(CompilationFlag.VariableOverlap)) {
           nf.fold[VariableVertex](GlobalVertex) { f =>
             if (m.alloc == VariableAllocationMethod.Static) {
@@ -118,38 +130,44 @@ class Environment(val parent: Option[Environment], val prefix: String) {
           case VariableAllocationMethod.None =>
             Nil
           case VariableAllocationMethod.Zeropage =>
-            m.sizeInBytes match {
+            if (forZpOnly) {
-              case 2 =>
               val addr =
-                  allocators(bank).allocatePointer(mem.banks(bank), callGraph, vertex)
+                allocators(bank).allocateBytes(mem.banks(bank), callGraph, vertex, options, m.sizeInBytes, initialized = false, writeable = true, location = AllocationLocation.Zeropage)
               onEachVariable(m.name, addr)
               List(
                 ConstantThing(m.name.stripPrefix(prefix) + "`", NumericConstant(addr, 2), p)
               )
-            }
+            } else Nil
           case VariableAllocationMethod.Auto | VariableAllocationMethod.Register | VariableAllocationMethod.Static =>
             if (m.alloc == VariableAllocationMethod.Register) {
               ErrorReporting.warn(s"Failed to inline variable `${m.name}` into a register", options, None)
             }
-            m.sizeInBytes match {
+            if (m.sizeInBytes == 0) Nil else {
-              case 0 => Nil
+              val graveName = m.name.stripPrefix(prefix) + "`"
-              case 2 =>
+              if (forZpOnly) {
-                val addr =
+                if (bank == "default") {
-                  allocators(bank).allocateBytes(mem.banks(bank), callGraph, vertex, options, 2, initialized = false, writeable = true)
+                  allocators(bank).tryAllocateZeropageBytes(mem.banks(bank), callGraph, vertex, options, m.sizeInBytes) match {
+                    case None => Nil
+                    case Some(addr) =>
                       onEachVariable(m.name, addr)
                       List(
                         ConstantThing(m.name.stripPrefix(prefix) + "`", NumericConstant(addr, 2), p)
                       )
-              case count =>
+                  }
-                val addr = allocators(bank).allocateBytes(mem.banks(bank), callGraph, vertex, options, count, initialized = false, writeable = true)
+                } else Nil
+              } else if (things.contains(graveName)) {
+                Nil
+              } else {
+                val addr = allocators(bank).allocateBytes(mem.banks(bank), callGraph, vertex, options, m.sizeInBytes, initialized = false, writeable = true, location = AllocationLocation.Either)
                 onEachVariable(m.name, addr)
                 List(
-                  ConstantThing(m.name.stripPrefix(prefix) + "`", NumericConstant(addr, 2), p)
+                  ConstantThing(graveName, NumericConstant(addr, 2), p)
                 )
               }
             }
+        }
       case f: NormalFunction =>
-        f.environment.allocateVariables(Some(f), mem, callGraph, allocators, options, onEachVariable)
+        f.environment.allocateVariables(Some(f), mem, callGraph, allocators, options, onEachVariable, pass, forZpOnly)
         Nil
       case _ => Nil
     }.toList

src/main/scala/millfork/output/AbstractAssembler.scala

@@ -170,6 +170,13 @@ abstract class AbstractAssembler[T <: AbstractCode](private val program: Program
 
   def assemble(callGraph: CallGraph, optimizations: Seq[AssemblyOptimization[T]], options: CompilationOptions): AssemblerOutput = {
     val platform = options.platform
+    val variableAllocators = platform.variableAllocators
+    val zpOccupied = mem.banks("default").occupied
+    (0 until 0x100).foreach(i => zpOccupied(i) = true)
+    platform.freeZpPointers.foreach { i =>
+      zpOccupied(i) = false
+      zpOccupied(i + 1) = false
+    }
 
     val assembly = mutable.ArrayBuffer[String]()
 
@@ -186,12 +193,16 @@ abstract class AbstractAssembler[T <: AbstractCode](private val program: Program
     val potentiallyInlineable: Map[String, Int] = inliningResult.potentiallyInlineableFunctions
     var nonInlineableFunctions: Set[String] = inliningResult.nonInlineableFunctions
 
+    env.allocateVariables(None, mem, callGraph, variableAllocators, options, labelMap.put, 1, forZpOnly = true)
+    env.allocateVariables(None, mem, callGraph, variableAllocators, options, labelMap.put, 2, forZpOnly = true)
+    env.allocateVariables(None, mem, callGraph, variableAllocators, options, labelMap.put, 3, forZpOnly = true)
+
     var inlinedFunctions = Map[String, List[T]]()
     val compiledFunctions = mutable.Map[String, List[T]]()
     val recommendedCompilationOrder = callGraph.recommendedCompilationOrder
     recommendedCompilationOrder.foreach { f =>
       env.maybeGet[NormalFunction](f).foreach { function =>
-        val code = compileFunction(function, optimizations, options, inlinedFunctions)
+        val code = compileFunction(function, optimizations, options, inlinedFunctions, labelMap.toMap)
         val strippedCodeForInlining = for {
           limit <- potentiallyInlineable.get(f)
           if code.map(_.sizeInBytes).sum <= limit
@@ -269,7 +280,7 @@ abstract class AbstractAssembler[T <: AbstractCode](private val program: Program
         val code = compiledFunctions(f.name)
         if (code.nonEmpty) {
           val size = code.map(_.sizeInBytes).sum
-          val index = codeAllocators(bank).allocateBytes(mem.banks(bank), options, size, initialized = true, writeable = false)
+          val index = codeAllocators(bank).allocateBytes(mem.banks(bank), options, size, initialized = true, writeable = false, location = AllocationLocation.High)
           labelMap(f.name) = index
           justAfterCode += bank -> outputFunction(bank, code, index, assembly, options)
         }
@@ -282,7 +293,7 @@ abstract class AbstractAssembler[T <: AbstractCode](private val program: Program
         val code = compiledFunctions(f.name)
         if (code.nonEmpty) {
           val size = code.map(_.sizeInBytes).sum
-          val index = codeAllocators(bank).allocateBytes(bank0, options, size, initialized = true, writeable = false)
+          val index = codeAllocators(bank).allocateBytes(bank0, options, size, initialized = true, writeable = false, location = AllocationLocation.High)
           labelMap(f.name) = index
           justAfterCode += bank -> outputFunction(bank, code, index, assembly, options)
         }
@@ -297,7 +308,7 @@ abstract class AbstractAssembler[T <: AbstractCode](private val program: Program
             val bank = m.bank(options)
             if (bank != "default") ???
             val bank0 = mem.banks(bank)
-            var index = codeAllocators(bank).allocateBytes(bank0, options, typ.size + 1, initialized = true, writeable = false)
+            var index = codeAllocators(bank).allocateBytes(bank0, options, typ.size + 1, initialized = true, writeable = false, location = AllocationLocation.High)
             labelMap(name) = index + 1
             val altName = m.name.stripPrefix(env.prefix) + "`"
             val thing = if (name.endsWith(".addr")) env.get[ThingInMemory](name.stripSuffix(".addr")) else env.get[ThingInMemory](name + ".array")
@@ -318,7 +329,7 @@ abstract class AbstractAssembler[T <: AbstractCode](private val program: Program
           }
         case _ => ()
       }
-      val index = codeAllocators("default").allocateBytes(mem.banks("default"), options, 1, initialized = true, writeable = false)
+      val index = codeAllocators("default").allocateBytes(mem.banks("default"), options, 1, initialized = true, writeable = false, location = AllocationLocation.High)
       writeByte("default", index, 2.toByte) // BIT abs
       assembly.append("* = $" + index.toHexString)
       assembly.append("    !byte 2 ;; end of LUnix relocatable segment")
@@ -328,7 +339,7 @@ abstract class AbstractAssembler[T <: AbstractCode](private val program: Program
       case thing@InitializedArray(name, None, items, _) =>
         val bank = thing.bank(options)
         val bank0 = mem.banks(bank)
-        var index = codeAllocators(bank).allocateBytes(bank0, options, items.size, initialized = true, writeable = true)
+        var index = codeAllocators(bank).allocateBytes(bank0, options, items.size, initialized = true, writeable = true, location = AllocationLocation.High)
         labelMap(name) = index
         assembly.append("* = $" + index.toHexString)
         assembly.append(name)
@@ -350,7 +361,7 @@ abstract class AbstractAssembler[T <: AbstractCode](private val program: Program
       case m@InitializedMemoryVariable(name, None, typ, value, _)  =>
         val bank = m.bank(options)
         val bank0 = mem.banks(bank)
-        var index = codeAllocators(bank).allocateBytes(bank0, options, typ.size, initialized = true, writeable = true)
+        var index = codeAllocators(bank).allocateBytes(bank0, options, typ.size, initialized = true, writeable = true, location = AllocationLocation.High)
         labelMap(name) = index
         val altName = m.name.stripPrefix(env.prefix) + "`"
         env.things += altName -> ConstantThing(altName, NumericConstant(index, 2), env.get[Type]("pointer"))
@@ -376,15 +387,23 @@ abstract class AbstractAssembler[T <: AbstractCode](private val program: Program
         val bank0 = mem.banks(bank)
         for(i <- 0 until size) bank0.occupied(addr + i) = true
     }
-    val variableAllocators = platform.variableAllocators
     variableAllocators.foreach { case (b, a) => a.notifyAboutEndOfCode(justAfterCode(b)) }
-    env.allocateVariables(None, mem, callGraph, variableAllocators, options, labelMap.put)
+    env.allocateVariables(None, mem, callGraph, variableAllocators, options, labelMap.put, 2, forZpOnly = false)
+    env.allocateVariables(None, mem, callGraph, variableAllocators, options, labelMap.put, 3, forZpOnly = false)
 
-    val zeropageOccupation = mem.banks("default").occupied.slice(variableAllocators("default").pointers.head, variableAllocators("default").pointers.last + 2)
+    if (platform.freeZpPointers.nonEmpty) {
+      val zpUsageOffset = platform.freeZpPointers.min
+      val zeropageOccupation = zpOccupied.slice(zpUsageOffset, platform.freeZpPointers.max + 2)
       labelMap += "__zeropage_usage" -> (zeropageOccupation.lastIndexOf(true) - zeropageOccupation.indexOf(true) + 1)
-    labelMap += "__zeropage_first" -> (zeropageOccupation.indexOf(true) max 0)
+      labelMap += "__zeropage_first" -> (zpUsageOffset + (zeropageOccupation.indexOf(true) max 0))
-    labelMap += "__zeropage_last" -> (zeropageOccupation.lastIndexOf(true) max 0)
+      labelMap += "__zeropage_last" -> (zpUsageOffset + (zeropageOccupation.lastIndexOf(true) max 0))
-    labelMap += "__zeropage_end" -> (zeropageOccupation.lastIndexOf(true) + 1)
+      labelMap += "__zeropage_end" -> (zpUsageOffset + zeropageOccupation.lastIndexOf(true) + 1)
+    } else {
+      labelMap += "__zeropage_usage" -> 0
+      labelMap += "__zeropage_first" -> 3
+      labelMap += "__zeropage_last" -> 2
+      labelMap += "__zeropage_end" -> 3
+    }
 
     env = rootEnv.allThings
 
@@ -421,16 +440,18 @@ abstract class AbstractAssembler[T <: AbstractCode](private val program: Program
     AssemblerOutput(code, assembly.toArray, labelMap.toList)
   }
 
-  private def compileFunction(f: NormalFunction, optimizations: Seq[AssemblyOptimization[T]], options: CompilationOptions, inlinedFunctions: Map[String, List[T]]): List[T] = {
+  def injectLabels(labelMap: Map[String, Int], code: List[T]): List[T]
+
+  private def compileFunction(f: NormalFunction, optimizations: Seq[AssemblyOptimization[T]], options: CompilationOptions, inlinedFunctions: Map[String, List[T]], labelMap: Map[String, Int]): List[T] = {
     ErrorReporting.debug("Compiling: " + f.name, f.position)
     val unoptimized: List[T] =
-      inliningCalculator.inline(
+      injectLabels(labelMap, inliningCalculator.inline(
         compiler.compile(CompilationContext(env = f.environment, function = f, extraStackOffset = 0, options = options)),
         inlinedFunctions,
-        compiler)
+        compiler))
     unoptimizedCodeSize += unoptimized.map(_.sizeInBytes).sum
     val code = optimizations.foldLeft(unoptimized) { (c, opt) =>
-      opt.optimize(f, c, options)
+      opt.optimize(f, c, options, labelMap)
     }
     performFinalOptimizationPass(f, optimizations.nonEmpty, options, code)
   }

src/main/scala/millfork/output/MosAssembler.scala

@@ -63,6 +63,30 @@ class MosAssembler(program: Program,
         index + 4
     }
   }
+
+  override def injectLabels(labelMap: Map[String, Int], code: List[AssemblyLine]): List[AssemblyLine] = {
+    import Opcode._
+    code.map {
+      case l@AssemblyLine(LDA | STA | CMP |
+                          LDX | STX | CPX |
+                          LDY | STY | CPY |
+                          LDZ | STZ | CPZ |
+                          BIT |
+                          ADC | SBC | AND | ORA | EOR |
+                          INC | DEC | ROL | ROR | ASL | LSR |
+                          ISC | DCP | LAX | SAX | RLA | RRA | SLO | SRE, AddrMode.Absolute, p, true) =>
+        p match {
+          case NumericConstant(n, _) => if (n <= 255) l.copy(addrMode = AddrMode.ZeroPage) else l
+          case MemoryAddressConstant(th) => if (labelMap.getOrElse(th.name, 0x800) < 0x100) l.copy(addrMode = AddrMode.ZeroPage) else l
+          case CompoundConstant(MathOperator.Plus,
+          MemoryAddressConstant(th),
+          NumericConstant(n, _)) => if (labelMap.getOrElse(th.name, 0x800) + n < 0x100) l.copy(addrMode = AddrMode.ZeroPage) else l
+          case _ => l
+        }
+
+      case l => l
+    }
+  }
 }
 
 

src/main/scala/millfork/output/VariableAllocator.scala

@@ -11,9 +11,20 @@ import scala.collection.mutable
   * @author Karol Stasiak
   */
 
+object AllocationLocation extends Enumeration {
+  val Zeropage, High, Either = Value
+  def matches(addr:Int, location: AllocationLocation.Value): Boolean = location match {
+    case Zeropage => addr < 0x100
+    case High => addr >= 0x100
+    case Either => true
+  }
+}
+
 sealed trait ByteAllocator {
-  protected def startAt: Int
+  def startAt: Int
-  protected def endBefore: Int
+  def endBefore: Int
+
+  def preferredOrder: Option[List[Int]]
 
   def notifyAboutEndOfCode(org: Int): Unit
 
@@ -21,7 +32,7 @@ sealed trait ByteAllocator {
     var lastFree = startAt
     var counter = 0
     val occupied = mem.occupied
-    for(i <- startAt until endBefore) {
+    for(i <- preferredOrder.getOrElse(startAt until endBefore)) {
       if (occupied(i) || counter == 0 && count == 2 && i.&(0xff) == 0xff && options.flags(CompilationFlag.PreventJmpIndirectBug)) {
         counter = 0
       } else {
@@ -37,72 +48,99 @@ sealed trait ByteAllocator {
         }
       }
     }
-    ErrorReporting.fatal("Out of high memory")
+    -1
   }
 }
 
 class UpwardByteAllocator(val startAt: Int, val endBefore: Int) extends ByteAllocator {
   def notifyAboutEndOfCode(org: Int): Unit = ()
+  override def preferredOrder: Option[List[Int]] = None
+}
+
+class ZeropageAllocator(val freeZpPointers: List[Int]) extends ByteAllocator {
+
+  def notifyAboutEndOfCode(org: Int): Unit = ()
+  override def preferredOrder: Option[List[Int]] = if (freeZpPointers.isEmpty) None else Some(freeZpPointers.flatMap(i => Seq(i,i+1)))
+
+  override def startAt: Int = if (freeZpPointers.isEmpty) 2 else freeZpPointers.min
+
+  override def endBefore: Int = if (freeZpPointers.isEmpty) 2 else freeZpPointers.max + 2
 }
 
 class AfterCodeByteAllocator(val endBefore: Int) extends ByteAllocator {
   var startAt = 0x200
   def notifyAboutEndOfCode(org: Int): Unit = startAt = org
+
+  override def preferredOrder: Option[List[Int]] = None
 }
 
-class VariableAllocator(val pointers: List[Int], private val bytes: ByteAllocator) {
+class VariableAllocator(pointers: List[Int], private val bytes: ByteAllocator) {
+
+  val zeropage: ByteAllocator = new ZeropageAllocator(pointers)
 
-  private val pointerMap = mutable.Map[Int, Set[VariableVertex]]()
   private val variableMap = mutable.Map[Int, mutable.Map[Int, Set[VariableVertex]]]()
 
-  def allocatePointer(mem: MemoryBank, callGraph: CallGraph, p: VariableVertex): Int = {
+  def allocateBytes(mem: MemoryBank, callGraph: CallGraph, p: VariableVertex, options: CompilationOptions, count: Int, initialized: Boolean, writeable: Boolean, location: AllocationLocation.Value): Int = {
-    // TODO: search for free zeropage locations
-    pointerMap.foreach { case (addr, alreadyThere) =>
-      if (alreadyThere.forall(q => callGraph.canOverlap(p, q))) {
-        pointerMap(addr) += p
-        return addr
-      }
-    }
-    @tailrec
-    def pickFreePointer(ps: List[Int]): Int =
-      ps match {
-        case Nil =>
-          ErrorReporting.trace(pointerMap.mkString(", "))
-          ErrorReporting.fatal("Out of zero-page memory")
-        case next :: rest =>
-          if (mem.occupied(next) || mem.occupied(next + 1)) {
-            pickFreePointer(rest)
-          } else {
-            mem.readable(next) = true
-            mem.readable(next + 1) = true
-            mem.occupied(next) = true
-            mem.occupied(next + 1) = true
-            mem.writeable(next) = true
-            mem.writeable(next + 1) = true
-            pointerMap(next) = Set(p)
-            next
-          }
-      }
-    pickFreePointer(pointers)
-  }
-
-  def allocateBytes(mem: MemoryBank, callGraph: CallGraph, p: VariableVertex, options: CompilationOptions, count: Int, initialized: Boolean, writeable: Boolean): Int = {
     if (!variableMap.contains(count)) {
       variableMap(count) = mutable.Map()
     }
     variableMap(count).foreach { case (a, alreadyThere) =>
-      if (alreadyThere.forall(q => callGraph.canOverlap(p, q))) {
+      if (AllocationLocation.matches(a, location) && alreadyThere.forall(q => callGraph.canOverlap(p, q))) {
         variableMap(count)(a) += p
         return a
       }
     }
-    val addr = allocateBytes(mem, options, count, initialized, writeable)
+    val addr = allocateBytes(mem, options, count, initialized, writeable, location)
     variableMap(count)(addr) = Set(p)
     addr
   }
 
-  def allocateBytes(mem: MemoryBank, options: CompilationOptions, count: Int, initialized: Boolean, writeable: Boolean): Int = {
+  def tryAllocateZeropageBytes(mem: MemoryBank, callGraph: CallGraph, p: VariableVertex, options: CompilationOptions, count: Int): Option[Int]={
-    val addr = bytes.findFreeBytes(mem, count, options)
+    if (!variableMap.contains(count)) {
+      variableMap(count) = mutable.Map()
+    }
+    variableMap(count).foreach { case (a, alreadyThere) =>
+      if (a < 0x100 && alreadyThere.forall(q => callGraph.canOverlap(p, q))) {
+        variableMap(count)(a) += p
+        return Some(a)
+      }
+    }
+    val addr = zeropage.findFreeBytes(mem, count, options)
+    if (addr < 0) None else {
+      markBytes(mem, addr, count, initialized = false, writeable = true)
+      Some(addr)
+    }
+  }
+
+  def allocateBytes(mem: MemoryBank, options: CompilationOptions, count: Int, initialized: Boolean, writeable: Boolean, location: AllocationLocation.Value): Int = {
+    val addr = location match {
+      case AllocationLocation.Zeropage =>
+        val a = zeropage.findFreeBytes(mem, count, options)
+        if (a < 0) {
+          ErrorReporting.fatal("Out of zeropage memory")
+        }
+        a
+      case AllocationLocation.High =>
+        val a = bytes.findFreeBytes(mem, count, options)
+        if (a < 0) {
+          ErrorReporting.fatal("Out of high memory")
+        }
+        a
+      case AllocationLocation.Either =>
+        var a = zeropage.findFreeBytes(mem, count, options)
+        if (a < 0) {
+          a = bytes.findFreeBytes(mem, count, options)
+          if (a < 0) {
+            ErrorReporting.fatal("Out of high memory")
+          }
+        }
+        a
+    }
+    markBytes(mem, addr, count, initialized, writeable)
+    addr
+  }
+
+  private def markBytes(mem: MemoryBank, addr: Int, count: Int, initialized: Boolean, writeable: Boolean): Unit = {
     ErrorReporting.trace(s"allocating $count bytes at $$${addr.toHexString}")
     (addr until (addr + count)).foreach { i =>
       if (mem.occupied(i)) ErrorReporting.fatal("Overlapping objects")
@@ -111,7 +149,6 @@ class VariableAllocator(val pointers: List[Int], private val bytes: ByteAllocato
       mem.initialized(i) = initialized
       mem.writeable(i) = writeable
     }
-    addr
   }
 
   def notifyAboutEndOfCode(org: Int): Unit = bytes.notifyAboutEndOfCode(org)

src/main/scala/millfork/output/Z80Assembler.scala

@@ -18,6 +18,8 @@ class Z80Assembler(program: Program,
     // TODO
     index
   }
+
+  override def injectLabels(labelMap: Map[String, Int], code: List[ZLine]): List[ZLine] = code // TODO
 }
 object Z80Assembler {
 

src/test/scala/millfork/test/emu/EmuPlatform.scala

@@ -7,13 +7,16 @@ import millfork.{Cpu, OutputStyle, Platform}
   * @author Karol Stasiak
   */
 object EmuPlatform {
+  private val pointers: List[Int] = (0 until 256 by 2).toList
+
   def get(cpu: Cpu.Value) = new Platform(
     cpu,
     Map(),
     Nil,
     CurrentBankFragmentOutput(0, 0xffff),
     Map("default" -> new UpwardByteAllocator(0x200, 0xb000)),
-    Map("default" -> new VariableAllocator((0 until 256 by 2).toList, new AfterCodeByteAllocator(0xff00))),
+    Map("default" -> new VariableAllocator(pointers, new AfterCodeByteAllocator(0xff00))),
+    pointers,
     ".bin",
     false,
     Map("default" -> 0),