From 8b09941cef47b0929dbeccb5e80a52ca524918a7 Mon Sep 17 00:00:00 2001
From: Karol Stasiak <karol.m.stasiak@gmail.com>
Date: Wed, 8 Aug 2018 13:45:38 +0200
Subject: [PATCH] Z80: Jump shortening

---
 .../assembly/z80/opt/JumpShortening.scala     | 58 +++++++++++++++++++
 .../scala/millfork/output/Z80Assembler.scala  |  7 ++-
 2 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 src/main/scala/millfork/assembly/z80/opt/JumpShortening.scala

diff --git a/src/main/scala/millfork/assembly/z80/opt/JumpShortening.scala b/src/main/scala/millfork/assembly/z80/opt/JumpShortening.scala
new file mode 100644
index 00000000..56dc48d0
--- /dev/null
+++ b/src/main/scala/millfork/assembly/z80/opt/JumpShortening.scala
@@ -0,0 +1,58 @@
+package millfork.assembly.z80.opt
+
+import millfork.{CompilationFlag, CompilationOptions}
+import millfork.assembly.OptimizationContext
+import millfork.assembly.z80.ZOpcode._
+import millfork.assembly.z80._
+import millfork.env.{Label, MemoryAddressConstant, NormalFunction}
+
+/**
+  * @author Karol Stasiak
+  */
+object JumpShortening {
+
+  def validShortJump(thisOffset: Int, labelOffset: Int): Boolean = {
+    val distance = labelOffset - (thisOffset + 2)
+    // TODO: I don't trust the instruction size calculations
+    // distance.toByte == distance
+    distance.abs < 100
+  }
+
+  def apply(f: NormalFunction, code: List[ZLine], options: CompilationOptions): List[ZLine] = {
+    if (!options.flags(CompilationFlag.EmitExtended80Opcodes)) return code
+    val shouldOptimizeAlways = options.flag(CompilationFlag.EmitSharpOpcodes) || options.flag(CompilationFlag.OptimizeForSize)
+    val bePessimistic = options.flag(CompilationFlag.OptimizeForSpeed)
+    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 (ZLine(LABEL, _, MemoryAddressConstant(Label(label)), _), ix) => Some(label -> offsets(ix))
+      case _ => None
+    }.toMap
+    code.zipWithIndex.map {
+      case (line@ZLine(JP, NoRegisters | IfFlagSet(ZFlag.Z | ZFlag.C) | IfFlagClear(ZFlag.Z | ZFlag.C), MemoryAddressConstant(Label(label)), true), ix) =>
+        labelOffsets.get(label).fold(line) { labelOffset =>
+          val thisOffset = offsets(ix)
+          val willBeTaken  = line.registers == NoRegisters ||
+            labelOffset < thisOffset ||  // jumping back means a loopty-loop TODO: better heuristics
+            bePessimistic // assume jump forward will happen if compiling with -Of
+          val willBeBetter = shouldOptimizeAlways || !willBeTaken
+          if (willBeBetter && validShortJump(thisOffset, labelOffset)) {
+            val result = line.copy(opcode = JR)
+            options.log.debug("Changing branch from long to short")
+            if (options.log.traceEnabled) {
+              options.log.trace(line.toString)
+              options.log.trace("     ↓")
+              options.log.trace(result.toString)
+            }
+            result
+          } else line
+        }
+      case (line, _) => line
+    }
+  }
+}
diff --git a/src/main/scala/millfork/output/Z80Assembler.scala b/src/main/scala/millfork/output/Z80Assembler.scala
index 46ba876f..6115aa75 100644
--- a/src/main/scala/millfork/output/Z80Assembler.scala
+++ b/src/main/scala/millfork/output/Z80Assembler.scala
@@ -2,6 +2,7 @@ package millfork.output
 
 import millfork.{CompilationFlag, CompilationOptions, Cpu, Platform}
 import millfork.assembly.z80._
+import millfork.assembly.z80.opt.JumpShortening
 import millfork.compiler.z80.Z80Compiler
 import millfork.env._
 import millfork.error.ConsoleLogger
@@ -15,7 +16,11 @@ import scala.collection.mutable
 class Z80Assembler(program: Program,
                    rootEnv: Environment,
                    platform: Platform) extends AbstractAssembler[ZLine](program, rootEnv, platform, Z80InliningCalculator, Z80Compiler) {
-  override def performFinalOptimizationPass(f: NormalFunction, actuallyOptimize: Boolean, options: CompilationOptions, code: List[ZLine]): List[ZLine] = code
+  override def performFinalOptimizationPass(f: NormalFunction, actuallyOptimize: Boolean, options: CompilationOptions, code: List[ZLine]): List[ZLine] = {
+    if (actuallyOptimize) {
+      JumpShortening(f, code, options)
+    } else code
+  }
 
   private def internalRegisterIndex(reg: ZRegister.Value): Int = reg match {
     case ZRegister.B => 0