Do not allocate removed variables; also, some weird bugfix

src/main/scala/millfork/assembly/opt/VariableToRegisterOptimization.scala

@@ -132,7 +132,7 @@ object VariableToRegisterOptimization extends AssemblyOptimization {
 
     val variablesWithLifetimes = localVariables.map(v =>
       v.name -> VariableLifetime.apply(v.name, code)
-    )
+    ).toMap
 
     val costFunction: CyclesAndBytes => Int = if (options.flag(CompilationFlag.OptimizeForSpeed)) _.cycles else _.bytes
     val importances = ReverseFlowAnalyzer.analyze(f, code)
@@ -240,10 +240,6 @@ object VariableToRegisterOptimization extends AssemblyOptimization {
     }
 
     if (bestXs.nonEmpty || bestYs.nonEmpty || bestZs.nonEmpty || bestAs.nonEmpty) {
-      bestXs.foreach(v => f.environment.removeVariable(v._1))
-      bestYs.foreach(v => f.environment.removeVariable(v._1))
-      bestZs.foreach(v => f.environment.removeVariable(v._1))
-      bestAs.foreach(v => f.environment.removeVariable(v._1))
       val output = ListBuffer[AssemblyLine]()
       var i = 0
       while (i < code.length) {
@@ -256,6 +252,9 @@ object VariableToRegisterOptimization extends AssemblyOptimization {
             reportOptimizedBlock(oldCode, newCode)
             output ++= newCode
             i = range.end
+            if (contains(range, variablesWithLifetimes(v))) {
+              f.environment.removeVariable(v)
+            }
             done = true
         }
         if (!done) {
@@ -267,6 +266,9 @@ object VariableToRegisterOptimization extends AssemblyOptimization {
               reportOptimizedBlock(oldCode, newCode)
               output ++= newCode
               i = range.end
+              if (contains(range, variablesWithLifetimes(v))) {
+                f.environment.removeVariable(v)
+              }
               done = true
           }
         }
@@ -279,6 +281,9 @@ object VariableToRegisterOptimization extends AssemblyOptimization {
               reportOptimizedBlock(oldCode, newCode)
               output ++= newCode
               i = range.end
+              if (contains(range, variablesWithLifetimes(v))) {
+                f.environment.removeVariable(v)
+              }
               done = true
           }
         }
@@ -291,6 +296,9 @@ object VariableToRegisterOptimization extends AssemblyOptimization {
               reportOptimizedBlock(oldCode, newCode)
               output ++= newCode
               i = range.end
+              if (contains(range, variablesWithLifetimes(v))) {
+                f.environment.removeVariable(v)
+              }
               done = true
           }
         }
@@ -305,6 +313,10 @@ object VariableToRegisterOptimization extends AssemblyOptimization {
     }
   }
 
+  def contains(outer: Range, inner: Range): Boolean = {
+    outer.contains(inner.start) && outer.contains(inner.end - 1)
+  }
+
   // TODO: STA has different flag behaviour than TAX, keep it in mind!
   def canBeInlined(xCandidate: Option[String], yCandidate: Option[String], zCandidate: Option[String], features: Features, lines: List[(AssemblyLine, CpuImportance)]): Option[CyclesAndBytes] = {
     val vx = xCandidate.getOrElse("-")
@@ -628,6 +640,15 @@ object VariableToRegisterOptimization extends AssemblyOptimization {
           None
         }
 
+      case (AssemblyLine(LDA, Absolute | ZeroPage, MemoryAddressConstant(th), true), imp) :: xs
+        if th.name == candidate =>
+        // removing LDA saves 3 bytes
+        if (imp.z == Unimportant && imp.n == Unimportant) {
+          canBeInlinedToAccumulator(options, start = false, synced = true, candidate, xs).map(_ + CyclesAndBytes(bytes = 3, cycles = 4))
+        } else {
+          canBeInlinedToAccumulator(options, start = false, synced = true, candidate, xs).map(_ + CyclesAndBytes(bytes = 1, cycles = 2))
+        }
+
       case (AssemblyLine(LDA, _, _, elidable),_) :: (AssemblyLine(op, Absolute | ZeroPage, MemoryAddressConstant(th), elidable2),_) :: xs
         if opcodesCommutative(op) =>
         if (th.name == candidate) {
@@ -642,15 +663,6 @@ object VariableToRegisterOptimization extends AssemblyOptimization {
           else None
         } else canBeInlinedToAccumulator(options, start = false, synced = synced, candidate, xs)
 
-      case (AssemblyLine(LDA, Absolute | ZeroPage, MemoryAddressConstant(th), true), imp) :: xs
-        if th.name == candidate =>
-        // removing LDA saves 3 bytes
-        if (imp.z == Unimportant && imp.n == Unimportant) {
-          canBeInlinedToAccumulator(options, start = false, synced = true, candidate, xs).map(_ + CyclesAndBytes(bytes = 3, cycles = 4))
-        } else {
-          canBeInlinedToAccumulator(options, start = false, synced = true, candidate, xs).map(_ + CyclesAndBytes(bytes = 1, cycles = 2))
-        }
-
       case (AssemblyLine(LDX | LDY | LAX, Absolute | ZeroPage, MemoryAddressConstant(th), elidable),_) :: xs
         if th.name == candidate =>
         // converting a load into a transfer saves 2 bytes
@@ -694,6 +706,12 @@ object VariableToRegisterOptimization extends AssemblyOptimization {
     }
   }
 
+  def isNot(v: String, param: Constant): Boolean = param match {
+    case MemoryAddressConstant(th) => th.name != v
+    case CompoundConstant(_, MemoryAddressConstant(th), _) =>  th.name != v
+    case _ => true
+  }
+
   def inlineVars(xCandidate: Option[String], yCandidate: Option[String], zCandidate: Option[String], aCandidate: Option[String], features: Features, lines: List[(AssemblyLine, CpuImportance)]): List[AssemblyLine] = {
     val vx = xCandidate.getOrElse("-")
     val vy = yCandidate.getOrElse("-")
@@ -841,17 +859,17 @@ object VariableToRegisterOptimization extends AssemblyOptimization {
         AssemblyLine(LDZ, am, param) :: AssemblyLine.implied(TZA) :: inlineVars(xCandidate, yCandidate, zCandidate, aCandidate, features, xs)
 
       case (AssemblyLine(LDA, Absolute | ZeroPage, MemoryAddressConstant(th), _), _) :: (AssemblyLine(CMP, am, param, true), _) :: xs
-        if th.name == vx && CpxyzAddrModes(am) =>
+        if th.name == vx && CpxyzAddrModes(am) && isNot(vx, param) =>
         // ditto
         AssemblyLine.implied(TXA) :: AssemblyLine(CPX, am, param) :: inlineVars(xCandidate, yCandidate, zCandidate, aCandidate, features, xs)
 
       case (AssemblyLine(LDA, Absolute | ZeroPage, MemoryAddressConstant(th), _), _) :: (AssemblyLine(CMP, am, param, true), _) :: xs
-        if th.name == vy && CpxyzAddrModes(am) =>
+        if th.name == vy && CpxyzAddrModes(am) && isNot(vx, param) =>
         // ditto
         AssemblyLine.implied(TYA) :: AssemblyLine(CPY, am, param) :: inlineVars(xCandidate, yCandidate, zCandidate, aCandidate, features, xs)
 
       case (AssemblyLine(LDA, Absolute | ZeroPage, MemoryAddressConstant(th), _), _) :: (AssemblyLine(CMP, am, param, true), _) :: xs
-        if th.name == vy && CpxyzAddrModes(am) =>
+        if th.name == vy && CpxyzAddrModes(am) && isNot(vx, param) =>
         // ditto
         AssemblyLine.implied(TZA) :: AssemblyLine(CPZ, am, param) :: inlineVars(xCandidate, yCandidate, zCandidate, aCandidate, features, xs)
 

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

@@ -20,10 +20,9 @@ class Environment(val parent: Option[Environment], val prefix: String) {
 
 
   private var baseStackOffset = 0x101
-  private val relVarId = new AtomicLong
 
   def genRelativeVariable(constant: Constant, typ: Type, zeropage: Boolean): RelativeVariable = {
-    val variable = RelativeVariable(".rv__" + relVarId.incrementAndGet().formatted("%06d"), constant, typ, zeropage = zeropage, declaredBank = None /*TODO*/)
+    val variable = RelativeVariable(".rv__" + Environment.relVarId.incrementAndGet().formatted("%06d"), constant, typ, zeropage = zeropage, declaredBank = None /*TODO*/)
     addThing(variable, None)
     variable
   }
@@ -149,6 +148,7 @@ class Environment(val parent: Option[Environment], val prefix: String) {
   }
 
   val things: mutable.Map[String, Thing] = mutable.Map()
+  val removedThings: mutable.Set[String] = mutable.Set()
 
   private def addThing(t: Thing, position: Option[Position]): Unit = {
     assertNotDefined(t.name, position)
@@ -156,8 +156,24 @@ class Environment(val parent: Option[Environment], val prefix: String) {
   }
 
   def removeVariable(str: String): Unit = {
+    ErrorReporting.trace("Removing variable: " + str)
+    removeVariableImpl(str)
+  }
+
+  private def removeVariableImpl(str: String): Unit = {
+    removedThings += str
+    removedThings += str + ".addr"
+    removedThings += str + ".addr.lo"
+    removedThings += str + ".addr.hi"
     things -= str
     things -= str + ".addr"
+    things -= str + ".addr.lo"
+    things -= str + ".addr.hi"
+    things -= str.stripPrefix(prefix)
+    things -= str.stripPrefix(prefix) + ".addr"
+    things -= str.stripPrefix(prefix) + ".addr.lo"
+    things -= str.stripPrefix(prefix) + ".addr.hi"
+    parent.foreach(_ removeVariableImpl str)
   }
 
   def get[T <: Thing : Manifest](name: String, position: Option[Position] = None): T = {
@@ -877,4 +893,5 @@ class Environment(val parent: Option[Environment], val prefix: String) {
 
 object Environment {
   val predefinedFunctions = Set("not", "hi", "lo", "nonet")
+  val relVarId = new AtomicLong
 }

src/main/scala/millfork/output/Assembler.scala

@@ -197,6 +197,7 @@ class Assembler(private val program: Program, private val rootEnv: Environment,
             compiledFunctions(f) = code
             optimizedCodeSize += code.map(_.sizeInBytes).sum
         }
+        function.environment.removedThings.foreach(env.removeVariable)
       }
     }