From afa871abcf6795b4b074453561b70e26d07f1371 Mon Sep 17 00:00:00 2001
From: Karol Stasiak <karol.m.stasiak@gmail.com>
Date: Wed, 26 Jun 2019 02:11:53 +0200
Subject: [PATCH] Z80: Fix stack variable optimizations when addresses of stack
 variables are taken

---
 .../ByteVariableToRegisterOptimization.scala  |  6 +++++
 .../assembly/z80/opt/CompactStackFrame.scala  | 17 +++++++++++++-
 .../scala/millfork/test/StackVarSuite.scala   | 22 ++++++++++++++++++-
 3 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/main/scala/millfork/assembly/z80/opt/ByteVariableToRegisterOptimization.scala b/src/main/scala/millfork/assembly/z80/opt/ByteVariableToRegisterOptimization.scala
index e20279f7..05cf03ee 100644
--- a/src/main/scala/millfork/assembly/z80/opt/ByteVariableToRegisterOptimization.scala
+++ b/src/main/scala/millfork/assembly/z80/opt/ByteVariableToRegisterOptimization.scala
@@ -241,6 +241,12 @@ object ByteVariableToRegisterOptimization extends AssemblyOptimization[ZLine] {
       def unapply(c: Int): Option[Int] = if ("IY+" + c == vname) Some(c) else None
     }
     code match {
+      case (_, ZLine0(LD_16, TwoRegisters(_, SP), _)) :: _ if vname.contains("+") => None
+      case (_, ZLine0(LD_16, TwoRegisters(SP, _), _)) :: _ if vname.contains("+") => None
+      case (_, ZLine0(ADD_16 | SBC_16, TwoRegisters(_, SP), _)) :: _ if vname.contains("+") => None
+      case (_, ZLine0(LD_HLSP | LD_DESP, _, _)) :: _ if vname.contains("+") => None
+      case (_, ZLine0(EX_SP, _, _)) :: _ if vname.contains("+") => None
+
       case (_, ZLine0(LD, TwoRegisters(A, MEM_ABS_8), ThisVar(_))) :: xs =>
         canBeInlined(vname, synced, target, addressInHl, addressInBc, addressInDe, xs).map(add(CyclesAndBytes(9, 2)))
       case (_, ZLine0(LD, TwoRegisters(MEM_ABS_8, A), ThisVar(_))) :: xs =>
diff --git a/src/main/scala/millfork/assembly/z80/opt/CompactStackFrame.scala b/src/main/scala/millfork/assembly/z80/opt/CompactStackFrame.scala
index e5778283..ae61019d 100644
--- a/src/main/scala/millfork/assembly/z80/opt/CompactStackFrame.scala
+++ b/src/main/scala/millfork/assembly/z80/opt/CompactStackFrame.scala
@@ -58,12 +58,27 @@ object CompactStackFrame extends AssemblyOptimization[ZLine] {
 
 
   def findUsedOffsets(code: List[ZLine], Mem: ZRegister.Value): Set[Int] = {
-    code.flatMap {
+    val used = code.flatMap {
       case ZLine0(_, OneRegisterOffset(Mem, offset), _) => Some(offset)
       case ZLine0(_, TwoRegistersOffset(_, Mem, offset), _) => Some(offset)
       case ZLine0(_, TwoRegistersOffset(Mem, _, offset), _) => Some(offset)
       case _ => None
     }.toSet
+    import ZOpcode._
+    import ZRegister._
+    if (used.isEmpty){
+      used
+    } else if (code.exists {
+      case ZLine0(LD_16, TwoRegisters(_, SP), _) => true
+      case ZLine0(LD_16, TwoRegisters(SP, _), _) => true
+      case ZLine0(ADD_16 | SBC_16, TwoRegisters(_, SP), _) => true
+      case ZLine0(LD_HLSP | LD_DESP, _, _) => true
+      case ZLine0(EX_SP, _, _) => true
+      case _ => false
+    }){
+      (0 to used.max).toSet
+    } else used
+
   }
 
   def optimizeContinue(code: List[ZLine], Index: ZRegister.Value, sourceSize: Int, targetSize: Int, mapping: Map[Int, Int]): Option[List[ZLine]] = {
diff --git a/src/test/scala/millfork/test/StackVarSuite.scala b/src/test/scala/millfork/test/StackVarSuite.scala
index f0019023..20f48adc 100644
--- a/src/test/scala/millfork/test/StackVarSuite.scala
+++ b/src/test/scala/millfork/test/StackVarSuite.scala
@@ -1,7 +1,7 @@
 package millfork.test
 
 import millfork.Cpu
-import millfork.test.emu.{EmuCrossPlatformBenchmarkRun, EmuSoftwareStackBenchmarkRun, EmuUnoptimizedCrossPlatformRun}
+import millfork.test.emu.{EmuCrossPlatformBenchmarkRun, EmuOptimizedZ80Run, EmuSoftwareStackBenchmarkRun, EmuUnoptimizedCrossPlatformRun}
 import org.scalatest.{FunSuite, Matchers}
 
 /**
@@ -390,5 +390,25 @@ class StackVarSuite extends FunSuite with Matchers {
     }
   }
 
+  test("Pointers to stack variables 3") {
+    val m = EmuOptimizedZ80Run(
+      """
+        | noinline byte f() = 6
+        | noinline void g(byte x) {}
+        | void main () {
+        |   stack byte b
+        |   stack byte a
+        |   stack byte c
+        |      b = b
+        |      c = c
+        |   a = f()
+        |   a += 5
+        |   a += a.addr.lo
+        |   g(a)
+        |   if f() == 0 { main() }
+        | }
+      """.stripMargin)
+  }
+
 
 }
\ No newline at end of file