From 63d4bf2317b44156fff5512597b6b1026ef75101 Mon Sep 17 00:00:00 2001
From: Karol Stasiak <karol.m.stasiak@gmail.com>
Date: Tue, 11 Feb 2020 01:10:49 +0100
Subject: [PATCH] 6502 Fix an optimization bug with the carry flag (fixes #46)

---
 .../mos/opt/AlwaysGoodOptimizations.scala     | 15 +++++--
 src/test/scala/millfork/test/MacroSuite.scala | 39 ++++++++++++++++++-
 .../millfork/test/emu/EmuOptimizedRun.scala   | 15 +++++++
 3 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/src/main/scala/millfork/assembly/mos/opt/AlwaysGoodOptimizations.scala b/src/main/scala/millfork/assembly/mos/opt/AlwaysGoodOptimizations.scala
index 4b1d2371..120f11e8 100644
--- a/src/main/scala/millfork/assembly/mos/opt/AlwaysGoodOptimizations.scala
+++ b/src/main/scala/millfork/assembly/mos/opt/AlwaysGoodOptimizations.scala
@@ -10,6 +10,7 @@ import millfork.assembly.mos.AddrMode._
 import millfork.env._
 import millfork.error.FatalErrorReporting
 import millfork.node.LiteralExpression
+import millfork.node.MosNiceFunctionProperty.DoesntChangeA
 
 /**
   * These optimizations should not remove opportunities for more complex optimizations to trigger.
@@ -1863,9 +1864,17 @@ object AlwaysGoodOptimizations {
         ctx.get[List[AssemblyLine]](5) ++
         ctx.get[List[AssemblyLine]](6) ++ (
           if (elseLabelUseCount == 1) Nil
-          else ctx.get[List[AssemblyLine]](10) ++ List(
-            AssemblyLine.immediate(LDA, if (firstSet) (if (zeroIfSet) nonZero else 0) else (if (zeroIfSet) 0 else nonZero))
-          ) ++ ctx.get[List[AssemblyLine]](5) ++ ctx.get[List[AssemblyLine]](6)
+          else {
+            val storeAndSecondReturn = ctx.get[List[AssemblyLine]](5) ++ ctx.get[List[AssemblyLine]](6)
+            val set = Set(RTS, JSR, RTI, RTL, BSR)
+            if (storeAndSecondReturn.exists(p => set(p.opcode))) {
+              ctx.get[List[AssemblyLine]](10) ++ List(
+                AssemblyLine.immediate(LDA, if (firstSet) (if (zeroIfSet) nonZero else 0) else (if (zeroIfSet) 0 else nonZero))
+              ) ++ ctx.get[List[AssemblyLine]](5) ++ ctx.get[List[AssemblyLine]](6)
+            } else {
+              ctx.get[List[AssemblyLine]](10)
+            }
+          }
         )
     }
   }
diff --git a/src/test/scala/millfork/test/MacroSuite.scala b/src/test/scala/millfork/test/MacroSuite.scala
index d1781d22..3247dd1b 100644
--- a/src/test/scala/millfork/test/MacroSuite.scala
+++ b/src/test/scala/millfork/test/MacroSuite.scala
@@ -1,7 +1,8 @@
 package millfork.test
 
 import millfork.Cpu
-import millfork.test.emu.{EmuBenchmarkRun, EmuCrossPlatformBenchmarkRun, EmuUnoptimizedCrossPlatformRun, ShouldNotCompile}
+import millfork.output.MemoryBank
+import millfork.test.emu.{EmuBenchmarkRun, EmuCrossPlatformBenchmarkRun, EmuOptimizedAccordingToLevelRun, EmuOptimizedRun, EmuUnoptimizedCrossPlatformRun, ShouldNotCompile}
 import org.scalatest.{FunSuite, Matchers}
 
 /**
@@ -178,4 +179,40 @@ class MacroSuite extends FunSuite with Matchers {
       m.readByte(0xc000) should equal(1)
     }
   }
+
+  test("Issue #46") {
+    val src =
+      """
+        |bool output
+        |
+        |byte input
+        |byte y
+        |
+        |void main() {
+        |    bool result
+        |    input = 13
+        |    y = 12
+        |    test(result, input)
+        |
+        |    if (result) {
+        |        output = true
+        |    }
+        |}
+        |
+        |macro void test(bool returnVal, byte x) {
+        |    returnVal = x >= y
+        |    returnVal = returnVal && x <= y + 8
+        |}
+        |""".stripMargin
+
+
+    def assertAtLeastOneTrueInZeroPage(m: MemoryBank): Unit = {
+      val countOfTrues = 0.to(0xff).count(a => m.readByte(a) == 1)
+      countOfTrues should be > (0)
+    }
+
+    assertAtLeastOneTrueInZeroPage(EmuOptimizedAccordingToLevelRun(1)(src))
+    assertAtLeastOneTrueInZeroPage(EmuOptimizedAccordingToLevelRun(2)(src))
+    assertAtLeastOneTrueInZeroPage(EmuOptimizedAccordingToLevelRun(3)(src))
+  }
 }
diff --git a/src/test/scala/millfork/test/emu/EmuOptimizedRun.scala b/src/test/scala/millfork/test/emu/EmuOptimizedRun.scala
index e8085deb..61b0baf4 100644
--- a/src/test/scala/millfork/test/emu/EmuOptimizedRun.scala
+++ b/src/test/scala/millfork/test/emu/EmuOptimizedRun.scala
@@ -25,6 +25,21 @@ object EmuOptimizedRun extends EmuRun(
     OptimizationPresets.Good ++
     OptimizationPresets.Good)
 
+case class EmuOptimizedAccordingToLevelRun(optLevel: Int) extends EmuRun(
+  Cpu.Ricoh,
+  OptimizationPresets.NodeOpt,
+  optLevel match {
+    case 0 => Nil
+    case 1 => OptimizationPresets.QuickPreset
+    case _ =>
+      val goodExtras = ZeropageRegisterOptimizations.All
+      val extras = Nil
+      val goodCycle = List.fill(optLevel - 2)(OptimizationPresets.Good ++ goodExtras).flatten
+      val mainCycle = List.fill(optLevel - 1)(OptimizationPresets.AssOpt ++ extras).flatten
+      goodCycle ++ mainCycle ++ goodCycle
+  }
+)
+
 object EmuSizeOptimizedRun extends EmuRun(
   Cpu.StrictMos,
   OptimizationPresets.NodeOpt,