Few optimization improvements for 6502

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

@@ -2,6 +2,7 @@ package millfork.assembly.mos.opt
 
 import java.util.concurrent.atomic.AtomicInteger
 
+import millfork.CompilationFlag
 import millfork.assembly.mos.Opcode._
 import millfork.assembly.mos.OpcodeClasses._
 import millfork.assembly.mos.{AddrMode, opt, _}
@@ -132,6 +133,12 @@ object AlwaysGoodOptimizations {
       HasClear(State.D) & HasSet(State.C) & DoesntMatterWhatItDoesWith(State.C, State.V)) ~~> { (code, ctx) =>
       AssemblyLine.immediate(LDA, (ctx.get[Constant](1) + (ctx.get[Int](0) + 1)).quickSimplify.loByte) :: Nil
     },
+    (Elidable &
+      MatchA(0) & MatchParameter(1) &
+      HasOpcode(ADC) & HasAddrMode(Immediate) &
+      HasSet(State.D) & HasClear(State.C) & DoesntMatterWhatItDoesWith(State.C, State.V)) ~~> { (code, ctx) =>
+      AssemblyLine.immediate(LDA, CompoundConstant(MathOperator.DecimalPlus, ctx.get[Constant](1), NumericConstant(ctx.get[Int](0), 1)).quickSimplify.loByte) :: Nil
+    },
     (Elidable &
       MatchA(0) & MatchParameter(1) &
       HasOpcode(SBC) & HasAddrMode(Immediate) &
@@ -2333,8 +2340,8 @@ object AlwaysGoodOptimizations {
         ctx.addObject(3, ZeroPage)
         true
       }) ~
-      (Linear & DoesNotConcernMemoryAt(3,4) & DoesNotConcernMemoryAt(3,5)).* ~
+      (Linear & DoesNotConcernMemoryAt(3, 4) & DoesNotConcernMemoryAt(3, 5)).* ~
-      (HasOpcode(STA) & MatchA(1) & HasAddrModeIn(Absolute, ZeroPage) & MatchParameter(5)) ~
+      (HasOpcode(STA) & HasAddrModeIn(Absolute, ZeroPage) & MatchA(1) & MatchParameter(5)) ~
       Where(ctx => {
         val lo = ctx.get[Int](0) & 0xff
         val hi = ctx.get[Int](1) & 0xff
@@ -2342,7 +2349,10 @@ object AlwaysGoodOptimizations {
         true
       }) ~
       (Linear & DoesNotConcernMemoryAt(3,4) & DoesNotConcernMemoryAt(3,5)).* ~
-      (Elidable & MatchParameter(6) & HasAddrModeIn(IndexedZ, IndexedY)) ~~> { (code, ctx) =>
+      (Elidable & MatchParameter(4) & HasAddrModeIn(IndexedZ, IndexedY) & MatchAddrMode(9)) ~
+      Where(ctx => {
+        ctx.get[AddrMode.Value](9) == IndexedY || !ctx.compilationOptions.flag(CompilationFlag.Emit65CE02Opcodes)
+      }) ~~> { (code, ctx) =>
       val addr = ctx.get[Int](2)
       val last = code.last
       code.init :+ last.copy(parameter = NumericConstant(addr, 2), addrMode = if (last.addrMode == IndexedZ) Absolute else AbsoluteY)
@@ -2363,8 +2373,11 @@ object AlwaysGoodOptimizations {
         ctx.addObject(2, hi * 256 + lo)
         true
       }) ~
-      (Linear & DoesNotConcernMemoryAt(3,4) & DoesNotConcernMemoryAt(3,5)).* ~
+      (Linear & DoesNotConcernMemoryAt(3, 4) & DoesNotConcernMemoryAt(3, 5)).* ~
-      (Elidable & MatchParameter(6) & HasAddrModeIn(IndexedZ, IndexedY)) ~~> { (code, ctx) =>
+      (Elidable & MatchParameter(4) & HasAddrModeIn(IndexedZ, IndexedY) & MatchAddrMode(9)) ~
+      Where(ctx => {
+        ctx.get[AddrMode.Value](9) == IndexedY || !ctx.compilationOptions.flag(CompilationFlag.Emit65CE02Opcodes)
+      }) ~~> { (code, ctx) =>
       val addr = ctx.get[Int](2)
       val last = code.last
       code.init :+ last.copy(parameter = NumericConstant(addr, 2), addrMode = if (last.addrMode == IndexedZ) Absolute else AbsoluteY)

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

@@ -1,8 +1,7 @@
 package millfork.assembly.mos.opt
 
-import millfork.CompilationOptions
 import millfork.assembly.mos.{AssemblyLine, OpcodeClasses}
-import millfork.assembly.AssemblyOptimization
+import millfork.assembly.{AssemblyOptimization, OptimizationContext}
 import millfork.env.NormalFunction
 import millfork.error.ErrorReporting
 
@@ -33,7 +32,7 @@ class ChangeIndexRegisterOptimization(preferX2Y: Boolean) extends AssemblyOptimi
 
   override def name = "Changing index registers"
 
-  override def optimize(f: NormalFunction, code: List[AssemblyLine], options: CompilationOptions): List[AssemblyLine] = {
+  override def optimize(f: NormalFunction, code: List[AssemblyLine], optimizationContext: OptimizationContext): List[AssemblyLine] = {
     val usesX = code.exists(l =>
       OpcodeClasses.ReadsXAlways(l.opcode) ||
         OpcodeClasses.ReadsYAlways(l.opcode) ||

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

@@ -29,9 +29,15 @@ object FlowInfoRequirement extends Enumeration {
   }
 }
 
-class RuleBasedAssemblyOptimization(val name: String, val needsFlowInfo: FlowInfoRequirement.Value, val rules: AssemblyRule*) extends AssemblyOptimization[AssemblyLine] {
+trait AssemblyRuleSet{
+  def flatten: Seq[AssemblyRule]
+}
+
+class RuleBasedAssemblyOptimization(val name: String, val needsFlowInfo: FlowInfoRequirement.Value, val rules: AssemblyRuleSet*) extends AssemblyOptimization[AssemblyLine] {
+
+  private val actualRules = rules.flatMap(_.flatten)
+  actualRules.foreach(_.pattern.validate(needsFlowInfo))
 
-  rules.foreach(_.pattern.validate(needsFlowInfo))
 
   override def optimize(f: NormalFunction, code: List[AssemblyLine], optimizationContext: OptimizationContext): List[AssemblyLine] = {
     val effectiveCode = code.map(a => a.copy(parameter = a.parameter.quickSimplify))
@@ -43,8 +49,8 @@ class RuleBasedAssemblyOptimization(val name: String, val needsFlowInfo: FlowInf
     code match {
       case Nil => Nil
       case head :: tail =>
-        for ((rule, index) <- rules.zipWithIndex) {
+        for ((rule, index) <- actualRules.zipWithIndex) {
-          val ctx = new AssemblyMatchingContext(optimizationContext.options, optimizationContext.labelMap, optimizationContext.niceFunctionProperties)
+          val ctx = new AssemblyMatchingContext(optimizationContext.options, optimizationContext.labelMap, optimizationContext.zreg, optimizationContext.niceFunctionProperties)
           rule.pattern.matchTo(ctx, code) match {
             case Some(rest: List[(FlowInfo, AssemblyLine)]) =>
               val matchedChunkToOptimize: List[AssemblyLine] = code.take(code.length - rest.length).map(_._2)
@@ -74,6 +80,7 @@ class RuleBasedAssemblyOptimization(val name: String, val needsFlowInfo: FlowInf
 
 class AssemblyMatchingContext(val compilationOptions: CompilationOptions,
                               val labelMap: Map[String, Int],
+                              val zeropageRegister: Option[ThingInMemory],
                               val niceFunctionProperties: Set[(NiceFunctionProperty, String)]) {
 
   def functionChangesA(name: String): Boolean = !niceFunctionProperties(MosNiceFunctionProperty.DoesntChangeA -> name)
@@ -189,10 +196,18 @@ class AssemblyMatchingContext(val compilationOptions: CompilationOptions,
     jumps.isEmpty
   }
 
+  def zreg(i: Int): Constant = {
+    MemoryAddressConstant(zeropageRegister.get) + i
+  }
+
 }
 
-case class AssemblyRule(pattern: AssemblyPattern, result: (List[AssemblyLine], AssemblyMatchingContext) => List[AssemblyLine]) {
+case class AssemblyRule(pattern: AssemblyPattern, result: (List[AssemblyLine], AssemblyMatchingContext) => List[AssemblyLine]) extends AssemblyRuleSet {
+  override def flatten: Seq[AssemblyRule] = List(this)
+}
 
+case class MultipleAssemblyRules(list: Seq[AssemblyRuleSet]) extends AssemblyRuleSet {
+  override def flatten: Seq[AssemblyRule] = list.flatMap(_.flatten)
 }
 
 trait AssemblyPattern {

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

@@ -57,7 +57,7 @@ object SuperOptimizer extends AssemblyOptimization[AssemblyLine] {
     )
     val optionsForMeasurements = options.copy(commandLineFlags = options.commandLineFlags + (CompilationFlag.InternalCurrentlyOptimizingForMeasurement -> true))
     val optimizationContextForMeasurements = optimizationContext.copy(options = optionsForMeasurements)
-    val quicklyCleanedCode = quickScrub.foldLeft(code)((c, o) => o.optimize(m, c, optionsForMeasurements))
+    val quicklyCleanedCode = quickScrub.foldLeft(code)((c, o) => o.optimize(m, c, optimizationContextForMeasurements))
     seenSoFar += viewCode(quicklyCleanedCode)
     queue.enqueue(quickScrub.reverse -> quicklyCleanedCode)
 
@@ -91,7 +91,7 @@ object SuperOptimizer extends AssemblyOptimization[AssemblyLine] {
     ErrorReporting.debug(s"Visited ${leaves.size} leaves")
     ErrorReporting.debug(s"${code.map(_.sizeInBytes).sum} B -> ${result._2.map(_.sizeInBytes).sum} B:   ${result._1.reverse.map(_.name).mkString(" -> ")}")
     result._1.reverse.foldLeft(code){(c, opt) =>
-      val n = opt.optimize(m, c, options)
+      val n = opt.optimize(m, c, optimizationContext)
 //      println(c.mkString("","",""))
 //      println(n.mkString("","",""))
       n

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

@@ -1,9 +1,8 @@
 package millfork.assembly.mos.opt
 
-import millfork.CompilationOptions
 import millfork.assembly.mos.AssemblyLine
 import millfork.assembly.mos.Opcode._
-import millfork.assembly.AssemblyOptimization
+import millfork.assembly.{AssemblyOptimization, OptimizationContext}
 import millfork.env._
 import millfork.error.ErrorReporting
 
@@ -12,7 +11,7 @@ import millfork.error.ErrorReporting
   */
 object UnusedLabelRemoval extends AssemblyOptimization[AssemblyLine] {
 
-  override def optimize(f: NormalFunction, code: List[AssemblyLine], options: CompilationOptions): List[AssemblyLine] = {
+  override def optimize(f: NormalFunction, code: List[AssemblyLine], optimizationContext: OptimizationContext): List[AssemblyLine] = {
     val usedLabels = code.flatMap {
       case AssemblyLine(LABEL, _, _, _) => None
       case AssemblyLine(_, _, MemoryAddressConstant(Label(l)), _) => Some(l)

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

@@ -3,7 +3,7 @@ package millfork.assembly.mos.opt
 import millfork.assembly.mos.Opcode._
 import millfork.assembly.mos.AddrMode._
 import millfork.assembly.AssemblyOptimization
-import millfork.assembly.mos.{AssemblyLine, State}
+import millfork.assembly.mos.{AssemblyLine, Opcode, State}
 import millfork.env.{CompoundConstant, Constant, MathOperator}
 
 /**
@@ -84,10 +84,30 @@ object ZeropageRegisterOptimizations {
 
   )
 
+  val StashInRegInsteadOfStack = new RuleBasedAssemblyOptimization("Stashing in zeropage register instead of stack",
+    needsFlowInfo = FlowInfoRequirement.BothFlows,
+    MultipleAssemblyRules((0 to 1).map{ zregIndex =>
+      (Elidable & HasOpcode(PHA) & DoesntMatterWhatItDoesWithReg(zregIndex)) ~
+        (Linear & Not(ConcernsS) & Not(RefersToOrUses("__reg", zregIndex))).*.capture(21) ~
+        (Elidable & HasOpcode(TSX)) ~
+        HasOpcodeIn(CLC, SED, CLD, SEC).*.capture(22) ~
+        (Elidable & MatchOpcode(1) & HasAddrMode(AbsoluteX) & HasParameterWhere(p => p.isProvably(0x101))).* ~
+        (Elidable & HasOpcode(INX)) ~
+        (Elidable & HasOpcode(TXS) & DoesntMatterWhatItDoesWith(State.Z, State.N)) ~~> { (code, ctx) =>
+        AssemblyLine.zeropage(STA, ctx.zreg(zregIndex)) :: (
+          ctx.get[List[AssemblyLine]](21) ++
+            ctx.get[List[AssemblyLine]](22) ++ List(
+            AssemblyLine.zeropage(ctx.get[Opcode.Value](1), ctx.zreg(zregIndex)),
+            AssemblyLine.implied(TSX)))
+      }
+    })
+  )
+
   val All: List[AssemblyOptimization[AssemblyLine]] = List(
     ConstantMultiplication,
     DeadRegStore,
     DeadRegStoreFromFlow,
+    StashInRegInsteadOfStack,
   )
 
 }