Support for comparing on-stack words + zeropage optimizations

src/main/scala/millfork/assembly/AssemblyLine.scala

@@ -184,14 +184,40 @@ object AssemblyLine {
 
   def implied(opcode: Opcode.Value) = AssemblyLine(opcode, AddrMode.Implied, Constant.Zero)
 
+  private val opcodesForNopVariableOperation = Set(STA, SAX, STX, STY, STZ)
+  private val opcodesForZeroedVariableOperation = Set(ADC, EOR, ORA, AND, SBC, CMP, CPX, CPY)
+  private val opcodesForZeroedOrSignExtendedVariableOperation = Set(LDA, LDX, LDY)
+
   def variable(ctx: CompilationContext, opcode: Opcode.Value, variable: Variable, offset: Int = 0): List[AssemblyLine] =
-    variable match {
-      case v@MemoryVariable(_, _, VariableAllocationMethod.Zeropage) =>
-        List(AssemblyLine.zeropage(opcode, v.toAddress + offset))
-      case v@RelativeVariable(_, _, _, true) =>
-        List(AssemblyLine.zeropage(opcode, v.toAddress + offset))
-      case v:VariableInMemory => List(AssemblyLine.absolute(opcode, v.toAddress + offset))
-      case v:StackVariable=> List(AssemblyLine.implied(TSX), AssemblyLine.absoluteX(opcode, v.baseOffset + offset + ctx.extraStackOffset))
+    if (offset > variable.typ.size) {
+      if (opcodesForNopVariableOperation(opcode)) {
+        Nil
+      } else if (opcodesForZeroedVariableOperation(opcode)) {
+        if (variable.typ.isSigned) ???
+        List(AssemblyLine.immediate(opcode, 0))
+      } else if (opcodesForZeroedOrSignExtendedVariableOperation(opcode)) {
+        if (variable.typ.isSigned) {
+          val label = MlCompiler.nextLabel("sx")
+          AssemblyLine.variable(ctx, opcode, variable, variable.typ.size - 1) ++ List(
+            AssemblyLine.immediate(ORA, 0x7f),
+            AssemblyLine.relative(BMI, label),
+            AssemblyLine.immediate(LDA, 0),
+            AssemblyLine.label(label))
+        } else {
+          List(AssemblyLine.immediate(opcode, 0))
+        }
+      } else {
+        ???
+      }
+    } else {
+      variable match {
+        case v@MemoryVariable(_, _, VariableAllocationMethod.Zeropage) =>
+          List(AssemblyLine.zeropage(opcode, v.toAddress + offset))
+        case v@RelativeVariable(_, _, _, true) =>
+          List(AssemblyLine.zeropage(opcode, v.toAddress + offset))
+        case v: VariableInMemory => List(AssemblyLine.absolute(opcode, v.toAddress + offset))
+        case v: StackVariable => List(AssemblyLine.implied(TSX), AssemblyLine.absoluteX(opcode, v.baseOffset + offset + ctx.extraStackOffset))
+      }
     }
 
   def zeropage(opcode: Opcode.Value, addr: Constant) =

src/main/scala/millfork/compiler/BuiltIns.scala

@@ -10,7 +10,6 @@ import millfork.error.ErrorReporting
 
 import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
-import scala.reflect.macros.blackbox
 
 
 object ComparisonType extends Enumeration {
@@ -381,7 +380,7 @@ object BuiltIns {
 
   def compileWordComparison(ctx: CompilationContext, compType: ComparisonType.Value, lhs: Expression, rhs: Expression, branches: BranchSpec): List[AssemblyLine] = {
     val env = ctx.env
-    // TODO: comparing stack variables
+    // TODO: comparing longer variables
     val b = env.get[Type]("byte")
     val w = env.get[Type]("word")
 
@@ -390,7 +389,7 @@ object BuiltIns {
       case BranchIfTrue(label) => compType -> label
       case BranchIfFalse(label) => ComparisonType.negate(compType) -> label
     }
-    val (lh, ll, rh, rl, ram) = (lhs, env.eval(lhs), rhs, env.eval(rhs)) match {
+    val (lh, ll, rh, rl) = (lhs, env.eval(lhs), rhs, env.eval(rhs)) match {
       case (_, Some(NumericConstant(lc, _)), _, Some(NumericConstant(rc, _))) =>
         return if (effectiveComparisonType match {
           // TODO: those masks are probably wrong
@@ -423,81 +422,84 @@ object BuiltIns {
       case (_, Some(lc), rv: VariableInMemory, None) =>
         return compileWordComparison(ctx, ComparisonType.flip(compType), rhs, lhs, branches)
       case (v: VariableExpression, None, _, Some(rc)) =>
-        // TODO: stack variables
-        (env.get[VariableInMemory](v.name + ".hi").toAddress,
-          env.get[VariableInMemory](v.name + ".lo").toAddress,
-          rc.hiByte,
-          rc.loByte,
-          Immediate)
+        val lva = env.get[VariableInMemory](v.name)
+        (AssemblyLine.variable(ctx, STA, lva, 1),
+          AssemblyLine.variable(ctx, STA, lva, 0),
+          List(AssemblyLine.immediate(STA, rc.hiByte)),
+          List(AssemblyLine.immediate(STA, rc.loByte)))
       case (lv: VariableExpression, None, rv: VariableExpression, None) =>
-        // TODO: stack variables
-        (env.get[VariableInMemory](lv.name + ".hi").toAddress,
-          env.get[VariableInMemory](lv.name + ".lo").toAddress,
-          env.get[VariableInMemory](rv.name + ".hi").toAddress,
-          env.get[VariableInMemory](rv.name + ".lo").toAddress, Absolute)
+        val lva = env.get[VariableInMemory](lv.name)
+        val rva = env.get[VariableInMemory](rv.name)
+        (AssemblyLine.variable(ctx, STA, lva, 1),
+          AssemblyLine.variable(ctx, STA, lva, 0),
+          AssemblyLine.variable(ctx, STA, rva, 1),
+          AssemblyLine.variable(ctx, STA, rva, 0))
     }
     effectiveComparisonType match {
       case ComparisonType.Equal =>
         val innerLabel = MlCompiler.nextLabel("cp")
-        List(AssemblyLine.absolute(LDA, ll),
-          AssemblyLine(CMP, ram, rl),
-          AssemblyLine.relative(BNE, innerLabel),
-          AssemblyLine.absolute(LDA, lh),
-          AssemblyLine(CMP, ram, rh),
-          AssemblyLine.relative(BEQ, Label(x)),
-          AssemblyLine.label(innerLabel))
+        staTo(LDA, ll) ++
+          staTo(CMP, rl) ++
+          List(AssemblyLine.relative(BNE, innerLabel)) ++
+          staTo(LDA, lh) ++
+          staTo(CMP, rh) ++
+          List(
+            AssemblyLine.relative(BEQ, Label(x)),
+            AssemblyLine.label(innerLabel))
 
       case ComparisonType.NotEqual =>
-        List(AssemblyLine.absolute(LDA, ll),
-          AssemblyLine(CMP, ram, rl),
-          AssemblyLine.relative(BNE, Label(x)),
-          AssemblyLine.absolute(LDA, lh),
-          AssemblyLine(CMP, ram, rh),
-          AssemblyLine.relative(BNE, Label(x)))
+        staTo(LDA, ll) ++
+          staTo(CMP, rl) ++
+          List(AssemblyLine.relative(BNE, Label(x))) ++
+          staTo(LDA, lh) ++
+          staTo(CMP, rh) ++
+          List(AssemblyLine.relative(BNE, Label(x)))
 
       case ComparisonType.LessUnsigned =>
         val innerLabel = MlCompiler.nextLabel("cp")
-        List(AssemblyLine.absolute(LDA, lh),
-          AssemblyLine(CMP, ram, rh),
-          AssemblyLine.relative(BCC, Label(x)),
-          AssemblyLine.relative(BNE, innerLabel),
-          AssemblyLine.absolute(LDA, ll),
-          AssemblyLine(CMP, ram, rl),
-          AssemblyLine.relative(BCC, Label(x)),
-          AssemblyLine.label(innerLabel))
+        staTo(LDA, lh) ++
+          staTo(CMP, rh) ++
+          List(
+            AssemblyLine.relative(BCC, Label(x)),
+            AssemblyLine.relative(BNE, innerLabel)) ++
+          staTo(LDA, ll) ++
+          staTo(CMP, rl) ++
+          List(
+            AssemblyLine.relative(BCC, Label(x)),
+            AssemblyLine.label(innerLabel))
 
       case ComparisonType.LessOrEqualUnsigned =>
         val innerLabel = MlCompiler.nextLabel("cp")
-        List(AssemblyLine(LDA, ram, rh),
-          AssemblyLine.absolute(CMP, lh),
-          AssemblyLine.relative(BCC, innerLabel),
-          AssemblyLine.relative(BNE, x),
-          AssemblyLine(LDA, ram, rl),
-          AssemblyLine.absolute(CMP, ll),
-          AssemblyLine.relative(BCS, x),
-          AssemblyLine.label(innerLabel))
+        staTo(LDA, rh) ++
+          staTo(CMP, lh) ++
+          List(AssemblyLine.relative(BCC, innerLabel),
+            AssemblyLine.relative(BNE, x)) ++
+          staTo(LDA, rl) ++
+          staTo(CMP, ll) ++
+          List(AssemblyLine.relative(BCS, x),
+            AssemblyLine.label(innerLabel))
 
       case ComparisonType.GreaterUnsigned =>
         val innerLabel = MlCompiler.nextLabel("cp")
-        List(AssemblyLine(LDA, ram, rh),
-          AssemblyLine.absolute(CMP, lh),
-          AssemblyLine.relative(BCC, Label(x)),
-          AssemblyLine.relative(BNE, innerLabel),
-          AssemblyLine(LDA, ram, rl),
-          AssemblyLine.absolute(CMP, ll),
-          AssemblyLine.relative(BCC, Label(x)),
-          AssemblyLine.label(innerLabel))
+        staTo(LDA, rh) ++
+          staTo(CMP, lh) ++
+          List(AssemblyLine.relative(BCC, Label(x)),
+            AssemblyLine.relative(BNE, innerLabel)) ++
+          staTo(LDA, rl) ++
+          staTo(CMP, ll) ++
+          List(AssemblyLine.relative(BCC, Label(x)),
+            AssemblyLine.label(innerLabel))
 
       case ComparisonType.GreaterOrEqualUnsigned =>
         val innerLabel = MlCompiler.nextLabel("cp")
-        List(AssemblyLine.absolute(LDA, lh),
-          AssemblyLine(CMP, ram, rh),
-          AssemblyLine.relative(BCC, innerLabel),
-          AssemblyLine.relative(BNE, x),
-          AssemblyLine.absolute(LDA, ll),
-          AssemblyLine(CMP, ram, rl),
-          AssemblyLine.relative(BCS, x),
-          AssemblyLine.label(innerLabel))
+        staTo(LDA, lh) ++
+          staTo(CMP, rh) ++
+          List(AssemblyLine.relative(BCC, innerLabel),
+            AssemblyLine.relative(BNE, x)) ++
+          staTo(LDA, ll) ++
+          staTo(CMP, rl) ++
+          List(AssemblyLine.relative(BCS, x),
+            AssemblyLine.label(innerLabel))
 
       case _ => ???
       // TODO: signed word comparisons

src/main/scala/millfork/compiler/MfCompiler.scala

@@ -21,26 +21,26 @@ sealed trait BranchSpec {
 }
 
 case object NoBranching extends BranchSpec {
-  override def flip = this
+  override def flip: BranchSpec = this
 }
 
 case class BranchIfTrue(label: String) extends BranchSpec {
-  override def flip = BranchIfFalse(label)
+  override def flip: BranchSpec = BranchIfFalse(label)
 }
 
 case class BranchIfFalse(label: String) extends BranchSpec {
-  override def flip = BranchIfTrue(label)
+  override def flip: BranchSpec = BranchIfTrue(label)
 }
 
 object BranchSpec {
-  val None = NoBranching
+  val None: BranchSpec = NoBranching
 }
 
 //noinspection NotImplementedCode,ScalaUnusedSymbol
 object MlCompiler {
 
 
-  private var labelCounter = new AtomicLong
+  private val labelCounter = new AtomicLong
 
   def nextLabel(prefix: String): String = "." + prefix + "__" + labelCounter.incrementAndGet().formatted("%05d")
 
@@ -158,20 +158,21 @@ object MlCompiler {
         AssemblyLine(LDA, Immediate, expr.hiByte),
         AssemblyLine(LDY, Immediate, expr.loByte))
       case m: VariableInMemory =>
+        val addrMode = if(m.zeropage) ZeroPage else Absolute
         val addr = m.toAddress
         m.typ.size match {
           case 0 => Nil
           case 1 => List(
             AssemblyLine(LDA, Immediate, expr.loByte),
-            AssemblyLine(STA, Absolute, addr))
+            AssemblyLine(STA, addrMode, addr))
           case 2 => List(
             AssemblyLine(LDA, Immediate, expr.loByte),
-            AssemblyLine(STA, Absolute, addr),
+            AssemblyLine(STA, addrMode, addr),
             AssemblyLine(LDA, Immediate, expr.hiByte),
-            AssemblyLine(STA, Absolute, addr + 1))
+            AssemblyLine(STA, addrMode, addr + 1))
           case s => List.tabulate(s)(i => List(
             AssemblyLine(LDA, Immediate, expr.subbyte(i)),
-            AssemblyLine(STA, Absolute, addr + i))).flatten
+            AssemblyLine(STA, addrMode, addr + i))).flatten
         }
       case StackVariable(_, t, offset) =>
         t.size match {
@@ -281,9 +282,15 @@ object MlCompiler {
           case s if s > 1 =>
             v match {
               case mv: VariableInMemory =>
-                AssemblyLine.absolute(store, mv) ::
-                  AssemblyLine.immediate(LDA, 0) ::
-                  List.tabulate(s - 1)(i => AssemblyLine.absolute(STA, mv.toAddress + (i + 1)))
+                if (mv.zeropage) {
+                  AssemblyLine.zeropage(store, mv) ::
+                    AssemblyLine.immediate(LDA, 0) ::
+                    List.tabulate(s - 1)(i => AssemblyLine.zeropage(STA, mv.toAddress + (i + 1)))
+                } else {
+                  AssemblyLine.absolute(store, mv) ::
+                    AssemblyLine.immediate(LDA, 0) ::
+                    List.tabulate(s - 1)(i => AssemblyLine.absolute(STA, mv.toAddress + (i + 1)))
+                }
               case sv@StackVariable(_, _, offset) =>
                 AssemblyLine.implied(transferToA) ::
                   AssemblyLine.implied(TSX) ::
@@ -1131,12 +1138,11 @@ object MlCompiler {
             AssemblyLine.implied(TAX),
             AssemblyLine.implied(PLA)) // fuck this shit
         }
-      case (_, RegisterVariable(Register.YA, _)) => {
+      case (_, RegisterVariable(Register.YA, _)) =>
         // TODO: sign extension
         List(
           AssemblyLine.implied(TAY),
           AssemblyLine.implied(TXA))
-      }
       case (_, RegisterVariable(Register.AY, _)) =>
         // TODO: sign extension
         List(