From 146e5986362d3196e89d04a57e637dd7f7d9ccf1 Mon Sep 17 00:00:00 2001
From: Karol Stasiak <karol.m.stasiak@gmail.com>
Date: Mon, 2 Apr 2018 21:57:53 +0200
Subject: [PATCH] Optimizing CMP/CPX/CPY/CPZ #0

---
 src/main/scala/millfork/OptimizationPresets.scala   |  2 ++
 .../assembly/opt/AlwaysGoodOptimizations.scala      |  8 ++++++++
 .../scala/millfork/assembly/opt/CpuStatus.scala     | 13 +++++++++++++
 .../opt/RuleBasedAssemblyOptimization.scala         |  8 ++++++++
 4 files changed, 31 insertions(+)

diff --git a/src/main/scala/millfork/OptimizationPresets.scala b/src/main/scala/millfork/OptimizationPresets.scala
index 5d5cc034..be9c9e3a 100644
--- a/src/main/scala/millfork/OptimizationPresets.scala
+++ b/src/main/scala/millfork/OptimizationPresets.scala
@@ -118,6 +118,7 @@ object OptimizationPresets {
     AlwaysGoodOptimizations.PointlessMathFromFlow,
     AlwaysGoodOptimizations.PointlessMathFromFlow,
     AlwaysGoodOptimizations.PointlessMathFromFlow,
+    AlwaysGoodOptimizations.OptimizeZeroComparisons,
     AlwaysGoodOptimizations.SimplifiableCondition,
     AlwaysGoodOptimizations.IncrementingIndexRegistersAfterTransfer,
     AlwaysGoodOptimizations.MathOperationOnTwoIdenticalMemoryOperands,
@@ -160,6 +161,7 @@ object OptimizationPresets {
     AlwaysGoodOptimizations.NonetAddition,
     AlwaysGoodOptimizations.NonetBitOp,
     AlwaysGoodOptimizations.OperationsAroundShifting,
+    AlwaysGoodOptimizations.OptimizeZeroComparisons,
     AlwaysGoodOptimizations.PoinlessFlagChange,
     AlwaysGoodOptimizations.PointlessLoadAfterLoadOrStore,
     AlwaysGoodOptimizations.PoinlessLoadBeforeAnotherLoad,
diff --git a/src/main/scala/millfork/assembly/opt/AlwaysGoodOptimizations.scala b/src/main/scala/millfork/assembly/opt/AlwaysGoodOptimizations.scala
index 7a453742..838421ea 100644
--- a/src/main/scala/millfork/assembly/opt/AlwaysGoodOptimizations.scala
+++ b/src/main/scala/millfork/assembly/opt/AlwaysGoodOptimizations.scala
@@ -1600,6 +1600,14 @@ object AlwaysGoodOptimizations {
     },
   )
 
+  val OptimizeZeroComparisons = new RuleBasedAssemblyOptimization("Optimizing zero comparisons",
+    needsFlowInfo = FlowInfoRequirement.BothFlows,
+    (Elidable & HasSourceOfNZ(State.A) & HasOpcode(CMP) & HasImmediate(0) & DoesntMatterWhatItDoesWith(State.C)) ~~> (_.init),
+    (Elidable & HasSourceOfNZ(State.X) & HasOpcode(CPX) & HasImmediate(0) & DoesntMatterWhatItDoesWith(State.C)) ~~> (_.init),
+    (Elidable & HasSourceOfNZ(State.Y) & HasOpcode(CPY) & HasImmediate(0) & DoesntMatterWhatItDoesWith(State.C)) ~~> (_.init),
+    (Elidable & HasSourceOfNZ(State.IZ) & HasOpcode(CPZ) & HasImmediate(0) & DoesntMatterWhatItDoesWith(State.C)) ~~> (_.init),
+  )
+
   private def remapZ2N(line: AssemblyLine) = line.opcode match {
     case BNE => line.copy(opcode = BMI)
     case BEQ => line.copy(opcode = BPL)
diff --git a/src/main/scala/millfork/assembly/opt/CpuStatus.scala b/src/main/scala/millfork/assembly/opt/CpuStatus.scala
index 4e7cb4d3..e704bda3 100644
--- a/src/main/scala/millfork/assembly/opt/CpuStatus.scala
+++ b/src/main/scala/millfork/assembly/opt/CpuStatus.scala
@@ -42,6 +42,11 @@ sealed trait Status[+T] {
     case SingleStatus(x) => f(x)
     case _ => AnyStatus
   }
+
+  def exists(predicate: T => Boolean): Boolean = this match {
+    case AnyStatus | UnknownStatus => false
+    case SingleStatus(x) => predicate(x)
+  }
 }
 
 object SourceOfNZ {
@@ -57,6 +62,14 @@ object SourceOfNZ {
 }
 
 case class SourceOfNZ(a: Boolean = false, aw: Boolean = false, x: Boolean = false, y: Boolean = false, iz: Boolean = false) {
+  def matches(state: State.Value): Boolean = state match {
+    case State.A => a
+    case State.X => x
+    case State.Y => y
+    case State.IZ => iz
+    case _ => throw new IllegalArgumentException
+  }
+
   override def toString: String = {
     val builder = new StringBuilder
     if (a) builder += 'A'
diff --git a/src/main/scala/millfork/assembly/opt/RuleBasedAssemblyOptimization.scala b/src/main/scala/millfork/assembly/opt/RuleBasedAssemblyOptimization.scala
index 9443f5fe..d939bc33 100644
--- a/src/main/scala/millfork/assembly/opt/RuleBasedAssemblyOptimization.scala
+++ b/src/main/scala/millfork/assembly/opt/RuleBasedAssemblyOptimization.scala
@@ -556,6 +556,14 @@ case class HasSet(state: State.Value) extends AssemblyLinePattern {
     flowInfo.hasSet(state)
 }
 
+case class HasSourceOfNZ(state: State.Value) extends AssemblyLinePattern {
+  override def validate(needsFlowInfo: FlowInfoRequirement.Value): Unit =
+    FlowInfoRequirement.assertForward(needsFlowInfo)
+
+  override def matchLineTo(ctx: AssemblyMatchingContext, flowInfo: FlowInfo, line: AssemblyLine): Boolean =
+    flowInfo.statusBefore.src.exists(s => s.matches(state))
+}
+
 object HasAccu8 extends AssemblyLinePattern {
   override def validate(needsFlowInfo: FlowInfoRequirement.Value): Unit =
     FlowInfoRequirement.assertForward(needsFlowInfo)