From 70818cc3d24aa2780807d6c71246d84fc7acbe2e Mon Sep 17 00:00:00 2001
From: Karol Stasiak <karol.m.stasiak@gmail.com>
Date: Mon, 11 Jun 2018 00:04:24 +0200
Subject: [PATCH] Fixing invalid short jumps

---
 docs/abi/generated-labels.md                  |  2 +
 .../millfork/assembly/opt/JumpFixing.scala    | 79 +++++++++++++++++++
 .../scala/millfork/output/Assembler.scala     |  6 +-
 3 files changed, 84 insertions(+), 3 deletions(-)
 create mode 100644 src/main/scala/millfork/assembly/opt/JumpFixing.scala

diff --git a/docs/abi/generated-labels.md b/docs/abi/generated-labels.md
index 3907ddc0..d9f49458 100644
--- a/docs/abi/generated-labels.md
+++ b/docs/abi/generated-labels.md
@@ -38,6 +38,8 @@ where `11111` is a sequential number and `xx` is the type:
 
 * `is` – optimized addition of carry using undocumented instructions
 
+* `lj` – extra labels generated when converting invalid short jumps to long jumps
+
 * `no` – nonet to word extension caused by the `nonet` operator
 
 * `od` – end of  a `do-while` statement
diff --git a/src/main/scala/millfork/assembly/opt/JumpFixing.scala b/src/main/scala/millfork/assembly/opt/JumpFixing.scala
new file mode 100644
index 00000000..55fc4393
--- /dev/null
+++ b/src/main/scala/millfork/assembly/opt/JumpFixing.scala
@@ -0,0 +1,79 @@
+package millfork.assembly.opt
+
+import java.util.concurrent.atomic.AtomicInteger
+
+import millfork.assembly.Opcode._
+import millfork.assembly.{AddrMode, AssemblyLine, Opcode}
+import millfork.compiler.MfCompiler
+import millfork.env.{Label, MemoryAddressConstant, NormalFunction}
+import millfork.error.ErrorReporting
+import millfork.{CompilationFlag, CompilationOptions}
+
+/**
+  * @author Karol Stasiak
+  */
+object JumpFixing {
+  val counter = new AtomicInteger(80000)
+
+  def generateNextLabel() = f".lj${counter.getAndIncrement()}%05d"
+
+  def invalidShortJump(thisOffset: Int, labelOffset: Int): Boolean = {
+    val distance = labelOffset - (thisOffset + 2)
+    distance.toByte != distance
+  }
+
+  private def negate(opcode: Opcode.Value) = opcode match {
+    case BEQ => BNE
+    case BNE => BEQ
+    case BCC => BCS
+    case BCS => BCC
+    case BVC => BVS
+    case BVS => BVC
+    case BMI => BPL
+    case BPL => BMI
+  }
+
+  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
+    var changed = false
+    val result = code.zipWithIndex.flatMap {
+      case (line@AssemblyLine(op, AddrMode.Relative, MemoryAddressConstant(Label(label)), true), ix) =>
+        labelOffsets.get(label).fold(List(line)) { labelOffset =>
+          val thisOffset = offsets(ix)
+          if (invalidShortJump(thisOffset, labelOffset)) {
+            changed = true
+            val long: List[AssemblyLine] = op match {
+              case BRA => List(line.copy(opcode = JMP, addrMode = AddrMode.Absolute))
+              case BSR => List(line.copy(opcode = JSR, addrMode = AddrMode.Absolute))
+              case _ =>
+                val label = generateNextLabel()
+                List(
+                  AssemblyLine.relative(negate(op), label),
+                  line.copy(opcode = JMP, addrMode = AddrMode.Absolute),
+                  AssemblyLine.label(label)
+                )
+            }
+            ErrorReporting.debug("Changing branch from short to long")
+            ErrorReporting.trace(line.toString)
+            ErrorReporting.trace("     ↓")
+            long.foreach(l => ErrorReporting.trace(l.toString))
+            long
+          } else List(line)
+        }
+      case (line, _) => List(line)
+    }
+    if (changed) apply(f, result, options) else result
+  }
+
+
+}
diff --git a/src/main/scala/millfork/output/Assembler.scala b/src/main/scala/millfork/output/Assembler.scala
index 1b59cc87..87a9c77d 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, HudsonOptimizations, JumpShortening}
+import millfork.assembly.opt.{AssemblyOptimization, HudsonOptimizations, JumpFixing, JumpShortening}
 import millfork.assembly.{AddrMode, AssemblyLine, Opcode}
 import millfork.compiler.{CompilationContext, MfCompiler}
 import millfork.env._
@@ -446,9 +446,9 @@ class Assembler(private val program: Program, private val rootEnv: Environment,
     }
     if (optimizations.nonEmpty) {
       val finalCode = if (options.flag(CompilationFlag.EmitHudsonOpcodes)) HudsonOptimizations.removeLoadZero(code) else code
-      JumpShortening(f, JumpShortening(f, finalCode, options), options)
+      JumpShortening(f, JumpShortening(f, JumpFixing(f, finalCode, options), options), options)
     }
-    else code
+    else JumpFixing(f, code, options)
   }
 
   private def outputFunction(bank: String, code: List[AssemblyLine], startFrom: Int, assOut: mutable.ArrayBuffer[String], options: CompilationOptions): Int = {