From 1bc1ab3539a64a41933067c24fa70c6ef6a692ff Mon Sep 17 00:00:00 2001
From: Karol Stasiak <karol.m.stasiak@gmail.com>
Date: Fri, 27 Jan 2023 17:00:29 +0100
Subject: [PATCH] Fix some codegen bugs on 6502

---
 src/main/scala/millfork/compiler/mos/BuiltIns.scala   |  8 +++++++-
 .../millfork/compiler/mos/MosExpressionCompiler.scala | 11 ++++++++---
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/main/scala/millfork/compiler/mos/BuiltIns.scala b/src/main/scala/millfork/compiler/mos/BuiltIns.scala
index 8f056e9b..548bf5b9 100644
--- a/src/main/scala/millfork/compiler/mos/BuiltIns.scala
+++ b/src/main/scala/millfork/compiler/mos/BuiltIns.scala
@@ -74,11 +74,17 @@ object BuiltIns {
         case FunctionCallExpression(name, List(param)) if env.maybeGet[Type](name).isDefined =>
           return simpleOperation(opcode, ctx, param, indexChoice, preserveA, commutative, decimal)
         case _: FunctionCallExpression | _: SumExpression if commutative =>
+          val code = MosExpressionCompiler.compileToA(ctx, source)
+          if (ctx.options.flags(CompilationFlag.IdentityPage)
+            && !code.exists(_.concernsX)
+            && AssemblyLine.treatment(code, State.X) == Treatment.Unchanged) {
+            return List(AssemblyLine.implied(TAX)) ++ code ++ wrapInSedCldIfNeeded(decimal, List(AssemblyLine.absoluteX(opcode, env.identityPage)))
+          }
           // TODO: is it ok?
           if (ctx.options.zpRegisterSize >= 1) {
             val reg = ctx.env.get[ThingInMemory]("__reg")
             return List(AssemblyLine.implied(PHA)) ++
-              MosExpressionCompiler.compileToA(ctx, source) ++
+              MosExpressionCompiler.fixTsx(code) ++
               List(AssemblyLine.zeropage(STA, reg), AssemblyLine.implied(PLA)) ++
               wrapInSedCldIfNeeded(decimal, List(AssemblyLine.zeropage(opcode, reg)))
           } else if (ctx.options.flag(CompilationFlag.EmitEmulation65816Opcodes)) {
diff --git a/src/main/scala/millfork/compiler/mos/MosExpressionCompiler.scala b/src/main/scala/millfork/compiler/mos/MosExpressionCompiler.scala
index c3bf8d44..7442d39d 100644
--- a/src/main/scala/millfork/compiler/mos/MosExpressionCompiler.scala
+++ b/src/main/scala/millfork/compiler/mos/MosExpressionCompiler.scala
@@ -72,6 +72,7 @@ object MosExpressionCompiler extends AbstractExpressionCompiler[AssemblyLine] {
           case 2 => if (ctx.options.flag(CompilationFlag.EmitNative65816Opcodes)) {
             AssemblyLine.accu16 :: AssemblyLine(LDA_W, WordImmediate, expr) :: (AssemblyLine.variable(ctx, STA_W, m) :+ AssemblyLine.accu8)
           } else AssemblyLine.tsx(ctx) ++ List(
+            // TODO: ???
             AssemblyLine.implied(TSX),
             AssemblyLine.immediate(LDA, expr.loByte),
             AssemblyLine.dataStackX(ctx, STA, offset),
@@ -155,10 +156,10 @@ object MosExpressionCompiler extends AbstractExpressionCompiler[AssemblyLine] {
     }
 
     val cmos = ctx.options.flag(CompilationFlag.EmitCmosOpcodes)
-    if (register == MosRegister.AX && !code.exists(_.concernsX)) {
+    if (register == MosRegister.AX && !code.exists(_.concernsX) && code.forall(_.treatment(State.X).==(Treatment.Unchanged))) {
       return preserveRegisterIfNeeded(ctx, MosRegister.A, code)
     }
-    if (register == MosRegister.AY && !code.exists(_.concernsY)) {
+    if (register == MosRegister.AY && !code.exists(_.concernsY) && code.forall(_.treatment(State.Y).==(Treatment.Unchanged))) {
       return preserveRegisterIfNeeded(ctx, MosRegister.A, code)
     }
     if (states.exists(state => AssemblyLine.treatment(code, state) != Treatment.Unchanged)) {
@@ -672,7 +673,11 @@ object MosExpressionCompiler extends AbstractExpressionCompiler[AssemblyLine] {
         }
         case RegisterVariable(MosRegister.Y, _) => actualOffset match {
           case 0 => List(tsx)
-          case 1 => List(tsx, AssemblyLine.implied(TXA), AssemblyLine.implied(TAY), AssemblyLine.implied(INY))
+          case 1 =>
+            if (ctx.options.flag(CompilationFlag.IdentityPage))
+              List(tsx, AssemblyLine.absoluteX(LDY, ctx.env.identityPage), AssemblyLine.implied(INY))
+            else
+              List(tsx, AssemblyLine.implied(TXA), AssemblyLine.implied(TAY), AssemblyLine.implied(INY))
           case _ => List(tsx, AssemblyLine.implied(TXA), AssemblyLine.implied(CLC), AssemblyLine.immediate(ADC, actualOffset), AssemblyLine.implied(TAY))
         }
         case _ =>