6502: Fix optimization bug

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

@@ -4,7 +4,7 @@ import millfork.assembly._
 import millfork.assembly.mos._
 import millfork.assembly.opt.FlowCache
 import millfork.env._
-import millfork.node.MosRegister
+import millfork.node.{MosRegister, NiceFunctionProperty}
 
 /**
   * @author Karol Stasiak
@@ -143,10 +143,12 @@ object ReverseFlowAnalyzer {
     m = Important, w = Important,
     r0 = Important, r1 = Important, r2 = Important, r3 = Important)
 
+  def analyze(f: NormalFunction, code: List[AssemblyLine], optimizationContext: OptimizationContext): List[CpuImportance] =
+    analyze(code, optimizationContext.niceFunctionProperties)
+
   //noinspection RedundantNewCaseClass
-  def analyze(f: NormalFunction, code: List[AssemblyLine], optimizationContext: OptimizationContext): List[CpuImportance] = {
+  def analyze(code: List[AssemblyLine], niceFunctionProperties: Set[(NiceFunctionProperty, String)]): List[CpuImportance] = {
     cache.get(code).foreach(return _)
-    val niceFunctionProperties = optimizationContext.niceFunctionProperties
     val importanceArray = Array.fill[CpuImportance](code.length)(new CpuImportance())
     val codeArray = code.toArray
 

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

@@ -90,7 +90,7 @@ object TwoVariablesToIndexRegistersOptimization extends AssemblyOptimization[Ass
     }.map(_.name).toSet
 
     val variablesWithLifetimes = localVariables.map(v =>
-      v.name -> VariableLifetime.apply(v.name, code, expandToIncludeIndexing = true)
+      v.name -> VariableLifetime.apply(v.name, code, expandToIncludeIndexing = true, expandToIncludeUsesOfLoadedIndices = Some(optimizationContext.niceFunctionProperties))
     ).toMap
 
     val removeVariablesForReal = !options.flag(CompilationFlag.InternalCurrentlyOptimizingForMeasurement)
@@ -190,12 +190,25 @@ object TwoVariablesToIndexRegistersOptimization extends AssemblyOptimization[Ass
       case (AssemblyLine0(LDX, Absolute | ZeroPage, MemoryAddressConstant(th)), _) :: xs if th.name == vy =>
         canBeInlined(vx, vy, vy, loadedY, xs).map(_ + CyclesAndBytes(bytes = 2, cycles = 2))
 
-      case (AssemblyLine0(LDX, _, _), _) :: xs if "--" == vx =>
+      case (AssemblyLine0(LDX, _, constant), _) :: xs if "--" == vx && !constant.refersTo(vy) =>
         canBeInlined(vx, vy, "-", loadedY, xs)
 
-      case (AssemblyLine0(LDY, _, _), _) :: xs if "--" == vy =>
+      case (AssemblyLine0(LDY, _, constant), _) :: xs if "--" == vy && !constant.refersTo(vx) =>
+        println(s"$constant doesn't refer to $vx")
         canBeInlined(vx, vy, loadedX, "-", xs)
 
+      case (l@AssemblyLine0(STY, ZeroPage | Absolute, _), _) :: xs if loadedY == vx =>
+        canBeInlined(vx, vy, loadedX, loadedY, xs)
+
+      case (l@AssemblyLine0(STX, ZeroPage | Absolute, _), _) :: xs if loadedX == vy =>
+        canBeInlined(vx, vy, loadedX, loadedY, xs)
+
+      case (l@AssemblyLine0(_,_, _), _) :: _ if l.concernsX && loadedX == vy =>
+        fail(71)
+
+      case (l@AssemblyLine0(_,_, _), _) :: _ if l.concernsY && loadedY == vx =>
+        fail(72)
+
       case (AssemblyLine0(_, AbsoluteX, _), _) :: xs if loadedX == vx || vx == "--" && loadedX == "-" =>
           canBeInlined(vx, vy, loadedX, loadedY, xs)
 
@@ -362,6 +375,12 @@ object TwoVariablesToIndexRegistersOptimization extends AssemblyOptimization[Ass
       case (l@AssemblyLine0(_, AbsoluteY, _), _) :: xs if loadedY == vx =>
         tailcall(inlineVars(vx, vy, loadedX, loadedY, xs)).map(l.copy(addrMode = AbsoluteX) :: _)
 
+      case (l@AssemblyLine0(STY, ZeroPage | Absolute, _), _) :: xs if loadedY == vx =>
+        tailcall(inlineVars(vx, vy, loadedX, loadedY, xs)).map(l.copy(opcode = STX) :: _)
+
+      case (l@AssemblyLine0(STX, ZeroPage | Absolute, _), _) :: xs if loadedX == vy =>
+        tailcall(inlineVars(vx, vy, loadedX, loadedY, xs)).map(l.copy(opcode = STY) :: _)
+
       case (x, _) :: xs => inlineVars(vx, vy, loadedX, loadedY, xs).map(x ::  _)
 
       case Nil => done(Nil)

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

@@ -1,8 +1,9 @@
 package millfork.assembly.mos.opt
 
-import millfork.assembly.mos.{AssemblyLine, OpcodeClasses}
+import millfork.assembly.mos.{AssemblyLine, AssemblyLine0, OpcodeClasses, State}
 import millfork.env._
 import millfork.error.ConsoleLogger
+import millfork.node.NiceFunctionProperty
 
 /**
   * @author Karol Stasiak
@@ -10,7 +11,7 @@ import millfork.error.ConsoleLogger
 object VariableLifetime {
 
   // This only works for non-stack variables.
-  def apply(variableName: String, code: List[AssemblyLine], expandToIncludeIndexing: Boolean = false): Range = {
+  def apply(variableName: String, code: List[AssemblyLine], expandToIncludeIndexing: Boolean = false, expandToIncludeUsesOfLoadedIndices: Option[Set[(NiceFunctionProperty, String)]] = None): Range = {
     val flags = code.map(_.parameter match {
       case MemoryAddressConstant(MemoryVariable(n, _, _)) if n == variableName => true
       case CompoundConstant(MathOperator.Plus, MemoryAddressConstant(MemoryVariable(n, _, _)), NumericConstant(_, 1)) if n == variableName => true
@@ -61,6 +62,59 @@ object VariableLifetime {
       }
     }
 
+    if (expandToIncludeUsesOfLoadedIndices.isDefined) {
+      // if the range ends with something like `LDY variableName`, then include the lifetime of that value in the register
+      // does not handle the Z register
+      import millfork.assembly.mos.Opcode._
+      val flow = ReverseFlowAnalyzer.analyze(code, expandToIncludeUsesOfLoadedIndices.get)
+      val maskX = Array.fill(code.length)(false)
+      val maskY = Array.fill(code.length)(false)
+      def mark(start: Int, mask: Array[Boolean], state: CpuImportance => Importance, readsReg: AssemblyLine => Boolean): Unit = {
+        var i = start
+        if (mask(i)) return
+        while (true) {
+          val line = code(i)
+          println(line)
+          if (state(flow(i)) != Important && !readsReg(line)) return
+          println("masking...")
+          mask(i) = true
+          val op = code(i).opcode
+          line match {
+            case AssemblyLine0(_, _, MemoryAddressConstant(Label(l1))) if OpcodeClasses.AllDirectJumps(op) =>
+              for (j <- labelMap.getOrElse(l1, Set())) {
+                mark(j, mask, state, readsReg)
+              }
+            case _ =>
+          }
+          if (op == RTS || op == RTI || op == BRA || op == JMP || op == BRL || op == KIL) return
+          if (state(flow(i)) != Important) return
+          i += 1
+          if (i >= code.length) return
+        }
+      }
+      @inline
+      def markX(i: Int): Unit = mark(i, maskX, _.x, _.reads(State.X))
+      @inline
+      def markY(i: Int): Unit = mark(i, maskY, _.y, _.reads(State.Y))
+      code.indices.foreach { i =>
+        val l = code(i)
+        l match {
+          case AssemblyLine0(LDX | LDX_W, _, MemoryAddressConstant(th)) if th.name == variableName => markX(i)
+          case AssemblyLine0(LDY | LDY_W, _, MemoryAddressConstant(th)) if th.name == variableName => markY(i)
+          case _ =>
+        }
+      }
+      @inline
+      def spread(arr: Array[Boolean]): Unit = {
+        val first = arr.indexOf(true)
+        if (first >= 0) min = min min first
+        val last = arr.lastIndexOf(true)
+        if (last >= 0) max = max max (last + 1)
+      }
+      spread(maskX)
+      spread(maskY)
+    }
+
 //    val log = new ConsoleLogger
 //    log.verbosity = 3
 //    log.trace("Lifetime for " + variableName)

src/test/scala/millfork/test/AbiSuite.scala

@@ -0,0 +1,30 @@
+package millfork.test
+
+import millfork.Cpu
+import millfork.test.emu.{EmuBenchmarkRun, EmuCrossPlatformBenchmarkRun, EmuOptimizedCmosRun, EmuOptimizedHudsonRun, EmuOptimizedRun, EmuUndocumentedRun, EmuUnoptimizedCrossPlatformRun, EmuUnoptimizedHudsonRun, EmuUnoptimizedRun, ShouldNotCompile}
+import org.scalatest.{AppendedClues, FunSuite, Matchers}
+
+/**
+  * @author Karol Stasiak
+  */
+class AbiSuite extends FunSuite with Matchers with AppendedClues {
+
+  test("Test parameter storage #1") {
+    EmuBenchmarkRun(
+      """
+        | byte output @$c000
+        | void main() {
+        |   f(42)
+        | }
+        |
+        | noinline void f(byte parameter) {
+        |   asm {
+        |     ldy parameter
+        |     sty output
+        |   }
+        | }
+        |""".stripMargin) { m =>
+      m.readByte(0xc000) should equal(42)
+    }
+  }
+}