Z80: Return dispatch

src/main/scala/millfork/compiler/AbstractReturnDispatch.scala

@@ -0,0 +1,165 @@
+package millfork.compiler
+
+import millfork.CompilationFlag
+import millfork.assembly.AbstractCode
+import millfork.env._
+import millfork.error.ErrorReporting
+import millfork.node._
+
+import scala.collection.mutable
+
+/**
+  * @author Karol Stasiak
+  */
+abstract class AbstractReturnDispatch[T <: AbstractCode] {
+
+  def compile(ctx: CompilationContext, stmt: ReturnDispatchStatement): List[T] = {
+    if (stmt.branches.isEmpty) {
+      ErrorReporting.error("At least one branch is required", stmt.position)
+      return Nil
+    }
+
+    def toConstant(e: Expression) = {
+      ctx.env.eval(e).getOrElse {
+        ErrorReporting.error("Non-constant parameter for dispatch branch", e.position)
+        Constant.Zero
+      }
+    }
+
+    def toInt(e: Expression): Int = {
+      ctx.env.eval(e) match {
+        case Some(NumericConstant(i, _)) =>
+          if (i < 0 || i > 255) ErrorReporting.error("Branch labels have to be in the 0-255 range", e.position)
+          i.toInt & 0xff
+        case _ =>
+          ErrorReporting.error("Branch labels have to early resolvable constants", e.position)
+          0
+      }
+    }
+
+    val indexerType = AbstractExpressionCompiler.getExpressionType(ctx, stmt.indexer)
+    if (indexerType.size != 1) {
+      ErrorReporting.error("Return dispatch index expression type has to be a byte", stmt.indexer.position)
+    }
+    if (indexerType.isSigned) {
+      ErrorReporting.warn("Return dispatch index expression type will be automatically casted to unsigned", ctx.options, stmt.indexer.position)
+    }
+    stmt.params.foreach {
+      case e@VariableExpression(name) =>
+        if (ctx.env.get[Variable](name).typ.size != 1) {
+          ErrorReporting.error("Dispatch parameters should be bytes", e.position)
+        }
+      case _ => ()
+    }
+
+    val returnType = ctx.function.returnType
+    val map: mutable.Map[Int, (Option[ThingInMemory], List[Expression])] = mutable.Map()
+    var min = Option.empty[Int]
+    var max = Option.empty[Int]
+    var default = Option.empty[(Option[ThingInMemory], List[Expression])]
+    stmt.branches.foreach { branch =>
+      val function: String = ctx.env.evalForAsm(branch.function) match {
+        case Some(MemoryAddressConstant(f: FunctionInMemory)) =>
+          if (f.returnType.name != returnType.name) {
+            ErrorReporting.warn(s"Dispatching to a function of different return type: dispatcher return type: ${returnType.name}, dispatchee return type: ${f.returnType.name}", ctx.options, branch.function.position)
+          }
+          f.name
+        case _ =>
+          ErrorReporting.error("Undefined function or Non-constant function address for dispatch branch", branch.function.position)
+          ""
+      }
+      branch.params.foreach(toConstant)
+      val params = branch.params
+      if (params.length > stmt.params.length) {
+        ErrorReporting.error("Too many parameters for dispatch branch", branch.params.head.position)
+      }
+      branch.label match {
+        case DefaultReturnDispatchLabel(start, end) =>
+          if (default.isDefined) {
+            ErrorReporting.error(s"Duplicate default dispatch label", branch.position)
+          }
+          min = start.map(toInt)
+          max = end.map(toInt)
+          default = Some(Some(ctx.env.get[FunctionInMemory](function)) -> params)
+        case StandardReturnDispatchLabel(labels) =>
+          labels.foreach { label =>
+            val i = toInt(label)
+            if (map.contains(i)) {
+              ErrorReporting.error(s"Duplicate dispatch label: $label = $i", label.position)
+            }
+            map(i) = Some(ctx.env.get[FunctionInMemory](function)) -> params
+          }
+      }
+    }
+    val nonDefaultMin = map.keys.reduceOption(_ min _)
+    val nonDefaultMax = map.keys.reduceOption(_ max _)
+    val defaultMin = min.orElse(nonDefaultMin).getOrElse(0)
+    val defaultMax = max.orElse(nonDefaultMax).getOrElse {
+      ErrorReporting.error("Undefined maximum label for dispatch", stmt.position)
+      defaultMin
+    }
+    val actualMin = defaultMin min nonDefaultMin.getOrElse(defaultMin)
+    val actualMax = defaultMax max nonDefaultMax.getOrElse(defaultMax)
+    val zeroes = None -> List[Expression]()
+    for (i <- actualMin to actualMax) {
+      if (!map.contains(i)) map(i) = default.getOrElse {
+        // TODO: warning?
+        zeroes
+      }
+    }
+
+    val compactParams = ctx.options.flag(CompilationFlag.CompactReturnDispatchParams)
+    val paramMins = stmt.params.indices.map { paramIndex =>
+      if (compactParams) map.filter(_._2._2.length > paramIndex).keys.reduceOption(_ min _).getOrElse(0)
+      else actualMin
+    }
+    val paramMaxes = stmt.params.indices.map { paramIndex =>
+      if (compactParams) map.filter(_._2._2.length > paramIndex).keys.reduceOption(_ max _).getOrElse(0)
+      else actualMax
+    }
+
+    val env = ctx.env.root
+    val b = env.get[VariableType]("byte")
+    val label = nextLabel("di")
+    val paramArrays = stmt.params.indices.map { ix =>
+      val a = InitializedArray(label + "$" + ix + ".array", None, (paramMins(ix) to paramMaxes(ix)).map { key =>
+        map(key)._2.lift(ix).getOrElse(LiteralExpression(0, 1))
+      }.toList,
+        ctx.function.declaredBank, b, b)
+      env.registerUnnamedArray(a)
+      a
+    }
+    compileImpl(ctx, stmt, label, actualMin, actualMax, paramArrays, paramMins, map)
+  }
+
+  def nextLabel(prefix: String): String
+
+  def compileImpl(ctx: CompilationContext,
+                           stmt: ReturnDispatchStatement,
+                           label: String,
+                           actualMin: Int,
+                           actualMax: Int,
+                           paramArrays: IndexedSeq[InitializedArray],
+                           paramMins: IndexedSeq[Int],
+                           map: mutable.Map[Int, (Option[ThingInMemory], List[Expression])]): List[T]
+
+  protected def zeroOr(function: Option[ThingInMemory])(F: ThingInMemory => Constant): Expression =
+    function.fold[Expression](LiteralExpression(0, 1))(F andThen ConstantArrayElementExpression)
+
+  protected def lobyte0(function: Option[ThingInMemory]): Expression = {
+    zeroOr(function)(f => MemoryAddressConstant(f).loByte)
+  }
+
+  protected def hibyte0(function: Option[ThingInMemory]): Expression = {
+    zeroOr(function)(f => MemoryAddressConstant(f).hiByte)
+  }
+
+  protected def lobyte1(function: Option[ThingInMemory]): Expression = {
+    zeroOr(function)(f => MemoryAddressConstant(f).-(1).loByte)
+  }
+
+  protected def hibyte1(function: Option[ThingInMemory]): Expression = {
+    zeroOr(function)(f => MemoryAddressConstant(f).-(1).hiByte)
+  }
+
+}

src/main/scala/millfork/compiler/mos/MosReturnDispatch.scala

@@ -2,8 +2,8 @@ package millfork.compiler.mos
 
 import millfork.CompilationFlag
 import millfork.assembly.mos.Opcode._
-import millfork.assembly.mos._
-import millfork.compiler.{BranchSpec, CompilationContext}
+import millfork.assembly.mos.{AddrMode, _}
+import millfork.compiler.{AbstractReturnDispatch, BranchSpec, CompilationContext}
 import millfork.env._
 import millfork.error.ErrorReporting
 import millfork.node._
@@ -13,126 +13,18 @@ import scala.collection.mutable
 /**
   * @author Karol Stasiak
   */
-object MosReturnDispatch {
+object MosReturnDispatch extends AbstractReturnDispatch[AssemblyLine] {
 
-  def compile(ctx: CompilationContext, stmt: ReturnDispatchStatement): List[AssemblyLine] = {
-    if (stmt.branches.isEmpty) {
-      ErrorReporting.error("At least one branch is required", stmt.position)
-      return Nil
-    }
-
-    def toConstant(e: Expression) = {
-      ctx.env.eval(e).getOrElse {
-        ErrorReporting.error("Non-constant parameter for dispatch branch", e.position)
-        Constant.Zero
-      }
-    }
-
-    def toInt(e: Expression): Int = {
-      ctx.env.eval(e) match {
-        case Some(NumericConstant(i, _)) =>
-          if (i < 0 || i > 255) ErrorReporting.error("Branch labels have to be in the 0-255 range", e.position)
-          i.toInt & 0xff
-        case _ =>
-          ErrorReporting.error("Branch labels have to early resolvable constants", e.position)
-          0
-      }
-    }
-
-    val indexerType = MosExpressionCompiler.getExpressionType(ctx, stmt.indexer)
-    if (indexerType.size != 1) {
-      ErrorReporting.error("Return dispatch index expression type has to be a byte", stmt.indexer.position)
-    }
-    if (indexerType.isSigned) {
-      ErrorReporting.warn("Return dispatch index expression type will be automatically casted to unsigned", ctx.options, stmt.indexer.position)
-    }
-    stmt.params.foreach{
-      case e@VariableExpression(name) =>
-        if (ctx.env.get[Variable](name).typ.size != 1) {
-          ErrorReporting.error("Dispatch parameters should be bytes", e.position)
-        }
-      case _ => ()
-    }
-
-    var env = ctx.env
-    val returnType = ctx.function.returnType
-    val map: mutable.Map[Int, (Option[ThingInMemory], List[Expression])] = mutable.Map()
-    var min = Option.empty[Int]
-    var max = Option.empty[Int]
-    var default = Option.empty[(Option[ThingInMemory], List[Expression])]
-    stmt.branches.foreach { branch =>
-      val function: String = ctx.env.evalForAsm(branch.function) match {
-        case Some(MemoryAddressConstant(f: FunctionInMemory)) =>
-          if (f.returnType.name != returnType.name) {
-            ErrorReporting.warn(s"Dispatching to a function of different return type: dispatcher return type: ${returnType.name}, dispatchee return type: ${f.returnType.name}", ctx.options, branch.function.position)
-          }
-          f.name
-        case _ =>
-          ErrorReporting.error("Undefined function or Non-constant function address for dispatch branch", branch.function.position)
-          ""
-      }
-      branch.params.foreach(toConstant)
-      val params = branch.params
-      if (params.length > stmt.params.length) {
-        ErrorReporting.error("Too many parameters for dispatch branch", branch.params.head.position)
-      }
-      branch.label match {
-        case DefaultReturnDispatchLabel(start, end) =>
-          if (default.isDefined) {
-            ErrorReporting.error(s"Duplicate default dispatch label", branch.position)
-          }
-          min = start.map(toInt)
-          max = end.map(toInt)
-          default = Some(Some(env.get[FunctionInMemory](function)) -> params)
-        case StandardReturnDispatchLabel(labels) =>
-          labels.foreach { label =>
-            val i = toInt(label)
-            if (map.contains(i)) {
-              ErrorReporting.error(s"Duplicate dispatch label: $label = $i", label.position)
-            }
-            map(i) = Some(env.get[FunctionInMemory](function)) -> params
-          }
-      }
-    }
-    val nonDefaultMin = map.keys.reduceOption(_ min _)
-    val nonDefaultMax = map.keys.reduceOption(_ max _)
-    val defaultMin = min.orElse(nonDefaultMin).getOrElse(0)
-    val defaultMax = max.orElse(nonDefaultMax).getOrElse {
-      ErrorReporting.error("Undefined maximum label for dispatch", stmt.position)
-      defaultMin
-    }
-    val actualMin = defaultMin min nonDefaultMin.getOrElse(defaultMin)
-    val actualMax = defaultMax max nonDefaultMax.getOrElse(defaultMax)
-    val zeroes = None -> List[Expression]()
-    for (i <- actualMin to actualMax) {
-      if (!map.contains(i)) map(i) = default.getOrElse {
-        // TODO: warning?
-        zeroes
-      }
-    }
-
-    val compactParams = ctx.options.flag(CompilationFlag.CompactReturnDispatchParams)
-    val paramMins = stmt.params.indices.map { paramIndex =>
-      if (compactParams) map.filter(_._2._2.length > paramIndex).keys.reduceOption(_ min _).getOrElse(0)
-      else actualMin
-    }
-    val paramMaxes = stmt.params.indices.map { paramIndex =>
-      if (compactParams) map.filter(_._2._2.length > paramIndex).keys.reduceOption(_ max _).getOrElse(0)
-      else actualMax
-    }
-
-    val b = ctx.env.get[VariableType]("byte")
-
-    while (env.parent.isDefined) env = env.parent.get
-    val label = MosCompiler.nextLabel("di")
-    val paramArrays = stmt.params.indices.map { ix =>
-      val a = InitializedArray(label + "$" + ix + ".array", None, (paramMins(ix) to paramMaxes(ix)).map { key =>
-        map(key)._2.lift(ix).getOrElse(LiteralExpression(0, 1))
-      }.toList,
-        ctx.function.declaredBank, b, b)
-      env.registerUnnamedArray(a)
-      a
-    }
+  override def compileImpl(ctx: CompilationContext,
+                  stmt: ReturnDispatchStatement,
+                  label: String,
+                  actualMin: Int,
+                  actualMax: Int,
+                  paramArrays: IndexedSeq[InitializedArray],
+                  paramMins: IndexedSeq[Int],
+                  map: mutable.Map[Int, (Option[ThingInMemory], List[Expression])]): List[AssemblyLine] = {
+    val env = ctx.env.root
+    val b = env.get[VariableType]("byte")
 
     val useJmpaix = ctx.options.flag(CompilationFlag.EmitCmosOpcodes) && !ctx.options.flag(CompilationFlag.LUnixRelocatableCode) && (actualMax - actualMin) <= 127
 
@@ -188,22 +80,5 @@ object MosReturnDispatch {
     }
   }
 
-  private def zeroOr(function: Option[ThingInMemory])(F: ThingInMemory => Constant): Expression =
-    function.fold[Expression](LiteralExpression(0, 1))(F andThen ConstantArrayElementExpression)
-
-  private def lobyte0(function: Option[ThingInMemory]): Expression = {
-    zeroOr(function)(f => MemoryAddressConstant(f).loByte)
-  }
-
-  private def hibyte0(function: Option[ThingInMemory]): Expression = {
-    zeroOr(function)(f => MemoryAddressConstant(f).hiByte)
-  }
-
-  private def lobyte1(function: Option[ThingInMemory]): Expression = {
-    zeroOr(function)(f => MemoryAddressConstant(f).-(1).loByte)
-  }
-
-  private def hibyte1(function: Option[ThingInMemory]): Expression = {
-    zeroOr(function)(f => MemoryAddressConstant(f).-(1).hiByte)
-  }
+  override def nextLabel(prefix: String): String = MosCompiler.nextLabel("di")
 }

src/main/scala/millfork/compiler/z80/Z80ReturnDispatch.scala

@@ -0,0 +1,74 @@
+package millfork.compiler.z80
+
+import millfork.assembly.z80.{ZLine, ZOpcode}
+import millfork.compiler.{AbstractReturnDispatch, CompilationContext}
+import millfork.env.{Constant, InitializedArray, ThingInMemory, VariableType}
+import millfork.error.ErrorReporting
+import millfork.node.{Expression, ReturnDispatchStatement, ZRegister}
+
+import scala.collection.mutable
+
+/**
+  * @author Karol Stasiak
+  */
+object Z80ReturnDispatch extends AbstractReturnDispatch[ZLine] {
+  override def nextLabel(prefix: String): String = Z80Compiler.nextLabel("di")
+
+  override def compileImpl(ctx: CompilationContext,
+                           stmt: ReturnDispatchStatement,
+                           label: String,
+                           actualMin: Int,
+                           actualMax: Int,
+                           paramArrays: IndexedSeq[InitializedArray],
+                           paramMins: IndexedSeq[Int],
+                           map: mutable.Map[Int, (Option[ThingInMemory], List[Expression])]): List[ZLine] = {
+
+    import ZRegister._
+    import ZOpcode._
+    val env = ctx.env.root
+    val b = env.get[VariableType]("byte")
+    val loadIndex = Z80ExpressionCompiler.compileToHL(ctx, stmt.indexer)
+    val ctxForStoringParams = ctx.neverCheckArrayBounds
+    val pair = stmt.params.zipWithIndex.foldLeft(Constant.Zero -> List[List[ZLine]]()) { (p1, p2) =>
+      (p1, p2) match {
+        case ((offset, reversedResult), (paramVar, paramIndex)) =>
+          val storeParam =
+            Z80ExpressionCompiler.stashBCIfChanged(
+              Z80ExpressionCompiler.stashHLIfChanged(
+                Z80ExpressionCompiler.storeA(ctxForStoringParams, paramVar, signedSource = false)))
+          if (storeParam.exists(l => l.changesRegister(A))) {
+            ErrorReporting.error("Invalid/too complex target parameter variable", paramVar.position)
+            storeParam.foreach(l => ErrorReporting.debug(l.toString))
+          }
+          val nextArray = (paramArrays(paramIndex).toAddress - paramMins(paramIndex)).quickSimplify
+          nextArray -> ((
+            ZLine.ldImm16(BC, (nextArray - offset).quickSimplify) ::
+              ZLine.registers(ADD_16, HL, BC) ::
+              ZLine.ld8(A, MEM_HL) ::
+              storeParam
+            ) :: reversedResult)
+      }
+    }
+    val copyParams = pair._2.reverse.flatten
+    val offsetAfterParams = pair._1
+    val jumpTableLo = InitializedArray(label + "$jl.array", None, (actualMin to actualMax).map(i => lobyte0(map(i)._1)).toList, ctx.function.declaredBank, b, b)
+    val jumpTableHi = InitializedArray(label + "$jh.array", None, (actualMin to actualMax).map(i => hibyte0(map(i)._1)).toList, ctx.function.declaredBank, b, b)
+    env.registerUnnamedArray(jumpTableLo)
+    env.registerUnnamedArray(jumpTableHi)
+
+    val moveOffsetToLo = (jumpTableLo.toAddress - offsetAfterParams - actualMin).quickSimplify
+    val moveOffsetToHi = (jumpTableHi.toAddress - jumpTableLo.toAddress).quickSimplify
+    val loadAddressToDE = List(
+      ZLine.ldImm16(BC, moveOffsetToLo),
+      ZLine.registers(ADD_16, HL, BC),
+      ZLine.ld8(E, MEM_HL),
+      ZLine.ldImm16(BC, moveOffsetToHi.quickSimplify),
+      ZLine.registers(ADD_16, HL, BC),
+      ZLine.ld8(D, MEM_HL))
+
+    loadIndex ++ copyParams ++ loadAddressToDE ++ List(
+      ZLine.ld8(H, D),
+      ZLine.ld8(L, E),
+      ZLine.register(JP, HL))
+  }
+}

src/main/scala/millfork/compiler/z80/Z80StatementCompiler.scala

@@ -79,6 +79,8 @@ object Z80StatementCompiler extends AbstractStatementCompiler[ZLine] {
         compileWhileStatement(ctx, s)
       case s: DoWhileStatement =>
         compileDoWhileStatement(ctx, s)
+      case s: ReturnDispatchStatement =>
+        Z80ReturnDispatch.compile(ctx, s)
 
       case f@ForStatement(_, _, _, _, List(Assignment(target: IndexedExpression, source: IndexedExpression))) =>
         Z80BulkMemoryOperations.compileMemcpy(ctx, target, source, f)

src/main/scala/millfork/env/Constant.scala

@@ -227,6 +227,10 @@ case class CompoundConstant(operator: MathOperator.Value, lhs: Constant, rhs: Co
         } else {
           CompoundConstant(MathOperator.Plus, a, rr - ll).quickSimplify
         }
+      case (_, CompoundConstant(MathOperator.Minus, a, b)) if operator == MathOperator.Minus =>
+        ((l + b) - a).quickSimplify
+      case (_, CompoundConstant(MathOperator.Plus, a, b)) if operator == MathOperator.Minus =>
+        ((l - a) - b).quickSimplify
       case (CompoundConstant(MathOperator.Shl, SubbyteConstant(c1, 1), NumericConstant(8, _)), SubbyteConstant(c2, 0)) if operator == MathOperator.Or && c1 == c2 => c1
       case (NumericConstant(lv, ls), NumericConstant(rv, rs)) =>
         var size = ls max rs

src/main/scala/millfork/output/Z80Assembler.scala

@@ -42,6 +42,7 @@ class Z80Assembler(program: Program,
 
   override def emitInstruction(bank: String, options: CompilationOptions, index: Int, instr: ZLine): Int = {
     import millfork.assembly.z80.ZOpcode._
+    import ZRegister._
     import Z80Assembler._
     instr match {
       case ZLine(LABEL, NoRegisters, MemoryAddressConstant(Label(labelName)), _) =>
@@ -289,6 +290,17 @@ class Z80Assembler(program: Program,
         writeByte(bank, index, 0xfa)
         writeWord(bank, index + 1, param)
         index + 3
+      case ZLine(JP, OneRegister(HL), _, _) =>
+        writeByte(bank, index, 0xe9)
+        index + 1
+      case ZLine(JP, OneRegister(IX), _, _) =>
+        writeByte(bank, index, 0xdd)
+        writeByte(bank, index, 0xe9)
+        index + 2
+      case ZLine(JP, OneRegister(IY), _, _) =>
+        writeByte(bank, index, 0xfd)
+        writeByte(bank, index, 0xe9)
+        index + 2
 
       case ZLine(RET, IfFlagClear(ZFlag.Z), _, _) =>
         writeByte(bank, index, 0xc0)

src/main/scala/millfork/output/Z80InliningCalculator.scala

@@ -31,6 +31,7 @@ object Z80InliningCalculator extends AbstractInliningCalculator[ZLine] {
       case ZLine(op, _, MemoryAddressConstant(Label(l)), _) if jumpingRelatedOpcodes(op) =>
         !l.startsWith(".")
       case ZLine(CALL, _, MemoryAddressConstant(th: ExternFunction), _) => false
+      case ZLine(JP, OneRegister(_), _, _) => false
       case ZLine(CALL, _, MemoryAddressConstant(th: NormalFunction), _) =>
         !functionsAlreadyKnownToBeNonInlineable(th.name)
       case ZLine(op, _, _, _) if jumpingRelatedOpcodes(op) || badOpcodes(op) => true

src/test/scala/millfork/test/ReturnDispatchSuite.scala

@@ -1,6 +1,7 @@
 package millfork.test
 
-import millfork.test.emu.{EmuBenchmarkRun, EmuCmosBenchmarkRun, EmuUnoptimizedRun}
+import millfork.Cpu
+import millfork.test.emu.EmuCrossPlatformBenchmarkRun
 import org.scalatest.{FunSuite, Matchers}
 
 /**
@@ -9,7 +10,7 @@ import org.scalatest.{FunSuite, Matchers}
 class ReturnDispatchSuite extends FunSuite with Matchers {
 
   test("Trivial test") {
-    EmuCmosBenchmarkRun(
+    EmuCrossPlatformBenchmarkRun(Cpu.Mos, Cpu.Cmos, Cpu.Z80)(
       """
         | byte output @$c000
         | void main () {
@@ -27,7 +28,7 @@ class ReturnDispatchSuite extends FunSuite with Matchers {
     }
   }
   test("Parameter test") {
-    EmuCmosBenchmarkRun(
+    EmuCrossPlatformBenchmarkRun(Cpu.Mos, Cpu.Cmos, Cpu.Z80)(
       """
         | array output [200] @$c000
         | sbyte param