From 91acbfae2f21eeaa8cae08e0c70ba337ab6cc3fa Mon Sep 17 00:00:00 2001
From: Karol Stasiak <karol.m.stasiak@gmail.com>
Date: Sat, 20 Jan 2018 18:51:49 +0100
Subject: [PATCH] Using short jumps whenever possible

---
 .../assembly/opt/JumpShortening.scala         | 78 +++++++++++++++++++
 .../scala/millfork/output/Assembler.scala     |  5 +-
 2 files changed, 81 insertions(+), 2 deletions(-)
 create mode 100644 src/main/scala/millfork/assembly/opt/JumpShortening.scala

diff --git a/src/main/scala/millfork/assembly/opt/JumpShortening.scala b/src/main/scala/millfork/assembly/opt/JumpShortening.scala
new file mode 100644
index 00000000..c095c8f0
--- /dev/null
+++ b/src/main/scala/millfork/assembly/opt/JumpShortening.scala
@@ -0,0 +1,78 @@
+package millfork.assembly.opt
+
+import millfork.{CompilationFlag, CompilationOptions}
+import millfork.assembly.{AddrMode, AssemblyLine, Opcode}
+import millfork.assembly.Opcode._
+import millfork.env.{Label, MemoryAddressConstant, NormalFunction}
+import millfork.error.ErrorReporting
+
+/**
+  * @author Karol Stasiak
+  */
+object JumpShortening {
+
+  def validShortJump(thisOffset: Int, labelOffset: Int): Boolean = {
+    val distance = labelOffset - (thisOffset + 2)
+    distance.toByte == distance
+  }
+
+  def apply(f: NormalFunction, code: List[AssemblyLine], options: CompilationOptions): List[AssemblyLine] = {
+    val offsets = new Array[Int](code.length)
+    var o = 0
+    code.zipWithIndex.foreach{
+      case (line, ix) =>
+        offsets(ix) = o
+        o += line.sizeInBytes
+    }
+    val labelOffsets = code.zipWithIndex.flatMap {
+      case (AssemblyLine(LABEL, _, MemoryAddressConstant(Label(label)), _), ix) => Some(label -> offsets(ix))
+      case _ => None
+    }.toMap
+    val cmos = options.flags(CompilationFlag.EmitCmosOpcodes)
+    if (cmos) {
+      code.zipWithIndex.map {
+        case (line@AssemblyLine(JMP, AddrMode.Absolute, MemoryAddressConstant(Label(label)), true), ix) =>
+          labelOffsets.get(label).fold(line) { labelOffset =>
+            val thisOffset = offsets(ix)
+            if (validShortJump(thisOffset, labelOffset)) {
+              val result = line.copy(opcode = BRA, addrMode = AddrMode.Relative)
+              ErrorReporting.debug("Changing branch from long to short")
+              ErrorReporting.trace(line.toString)
+              ErrorReporting.trace("     ↓")
+              ErrorReporting.trace(result.toString)
+              result
+            } else line
+          }
+          case (line, _) => line
+      }
+    } else {
+      FlowAnalyzer.analyze(f, code, options, FlowInfoRequirement.ForwardFlow).zipWithIndex.map {
+        case ((info, line@AssemblyLine(JMP, AddrMode.Absolute, MemoryAddressConstant(Label(label)), _)), ix) =>
+          labelOffsets.get(label).fold(line) { labelOffset =>
+            val thisOffset = offsets(ix)
+            if (validShortJump(thisOffset, labelOffset)) {
+              val bra =
+                if (info.statusBefore.z.contains(true)) BEQ
+                else if (info.statusBefore.z.contains(false)) BNE
+                else if (info.statusBefore.n.contains(true)) BMI
+                else if (info.statusBefore.n.contains(false)) BPL
+                else if (info.statusBefore.v.contains(true)) BVS
+                else if (info.statusBefore.v.contains(false)) BVC
+                else if (info.statusBefore.c.contains(true)) BCS
+                else if (info.statusBefore.c.contains(false)) BCC
+                else JMP
+              if (bra != JMP) {
+                val result = line.copy(opcode = bra, addrMode = AddrMode.Relative)
+                ErrorReporting.debug("Changing branch from long to short")
+                ErrorReporting.trace(line.toString)
+                ErrorReporting.trace("     ↓")
+                ErrorReporting.trace(result.toString)
+                result
+              } else line
+            } else line
+          }
+        case ((_, line), _) => line
+      }
+    }
+  }
+}
diff --git a/src/main/scala/millfork/output/Assembler.scala b/src/main/scala/millfork/output/Assembler.scala
index a1dc37ba..f94b1db4 100644
--- a/src/main/scala/millfork/output/Assembler.scala
+++ b/src/main/scala/millfork/output/Assembler.scala
@@ -1,6 +1,6 @@
 package millfork.output
 
-import millfork.assembly.opt.AssemblyOptimization
+import millfork.assembly.opt.{AssemblyOptimization, JumpShortening}
 import millfork.assembly.{AddrMode, AssemblyLine, Opcode}
 import millfork.compiler.{CompilationContext, MfCompiler}
 import millfork.env._
@@ -305,7 +305,8 @@ class Assembler(private val program: Program, private val rootEnv: Environment)
     val code = optimizations.foldLeft(unoptimized) { (c, opt) =>
       opt.optimize(f, c, options)
     }
-    code
+    if (optimizations.nonEmpty) JumpShortening(f, JumpShortening(f, code, options), options)
+    else code
   }
 
   private def outputFunction(code: List[AssemblyLine], startFrom: Int, assOut: mutable.ArrayBuffer[String], options: CompilationOptions): Int = {