6502: Fix and optimize sign extension

src/main/scala/millfork/OptimizationPresets.scala

@@ -179,6 +179,7 @@ object OptimizationPresets {
     AlwaysGoodOptimizations.SimplifiableStackOperation,
     LaterOptimizations.UseBit,
     LaterOptimizations.ReplaceableLoad,
+    LaterOptimizations.BranchlessSignExtension,
   )
 
   val Good: List[AssemblyOptimization[AssemblyLine]] = List[AssemblyOptimization[AssemblyLine]](
@@ -261,6 +262,7 @@ object OptimizationPresets {
     UseAccumulatorInsteadOfYRegister,
     VariableToRegisterOptimization,
     TwoVariablesToIndexRegistersOptimization,
+    LaterOptimizations.BranchlessSignExtension,
   )
 
   val QuickPreset: List[AssemblyOptimization[AssemblyLine]] = List[AssemblyOptimization[AssemblyLine]](

src/main/scala/millfork/assembly/mos/opt/LaterOptimizations.scala

@@ -559,7 +559,34 @@ object LaterOptimizations {
 
   )
 
+  val BranchlessSignExtension = new RuleBasedAssemblyOptimization("Branchless sign extension",
+    needsFlowInfo = FlowInfoRequirement.BothFlows,
+    (NotFixed & HasOpcode(LDA) & HasAddrModeIn(Absolute, ZeroPage, AbsoluteX, IndexedX, IndexedY, AbsoluteY, ZeroPageX) & HasClear(State.D)) ~
+      (Elidable & HasOpcode(ORA) & HasImmediate(0x7F)) ~
+      (Elidable & HasOpcode(BMI) & MatchParameter(10)) ~
+      (Elidable & HasOpcode(LDA) & HasImmediate(0)) ~
+      (Elidable & HasOpcode(LABEL) & IsNotALabelUsedManyTimes & MatchParameter(10) & DoesntMatterWhatItDoesWith(State.V)) ~~> { code =>
+      List(AssemblyLine.immediate(LDA, 0x7F), code.head.copy(opcode = CMP), AssemblyLine.immediate(SBC, 0x7F))
+    },
+
+    (NotFixed & HasOpcode(LDA) & HasAddrModeIn(Absolute, ZeroPage, AbsoluteX, IndexedX, IndexedY, AbsoluteY, ZeroPageX) & HasClear(State.D)) ~
+      (Elidable & HasOpcode(AND) & HasImmediate(0x80)) ~
+      (Elidable & HasOpcode(BPL) & MatchParameter(10)) ~
+      (Elidable & HasOpcode(LDA) & HasImmediate(0xff)) ~
+      (Elidable & HasOpcode(LABEL) & IsNotALabelUsedManyTimes & MatchParameter(10) & DoesntMatterWhatItDoesWith(State.V)) ~~> { code =>
+      List(AssemblyLine.immediate(LDA, 0x7F), code.head.copy(opcode = CMP), AssemblyLine.immediate(SBC, 0x7F))
+    },
+
+  )
+
+//  AssemblyLine.immediate(ORA, 0x7F),
+//  AssemblyLine.relative(BMI, label),
+//  AssemblyLine.immediate(LDA, 0),
+//  AssemblyLine.label(label))
+
+  // use the lists in OptimizationPresets to actually add these to the normal pipeline
   val All = List(
+    BranchlessSignExtension,
     CommutativeInPlaceModifications,
     DontUseIndexRegisters,
     DoubleLoadToDifferentRegisters,

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

@@ -61,9 +61,9 @@ object MosStatementCompiler extends AbstractStatementCompiler[AssemblyLine] {
             AssemblyLine.zeropage(STA, reg,i).copy(elidability = Elidability.Volatile))
         }.toList
       } else Nil)
-    val someRegisterA = Some(b, RegisterVariable(MosRegister.A, b))
+//    val someRegisterA = Some(b, RegisterVariable(MosRegister.A, b))
-    val someRegisterAX = Some(w, RegisterVariable(MosRegister.AX, w))
+//    val someRegisterAX = Some(w, RegisterVariable(MosRegister.AX, w))
-    val someRegisterYA = Some(w, RegisterVariable(MosRegister.YA, w))
+//    val someRegisterYA = Some(w, RegisterVariable(MosRegister.YA, w))
     lazy val returnInstructions = if (m.interrupt) {
       if (ctx.options.flag(CompilationFlag.EmitNative65816Opcodes)) {
         if (zpRegisterSize > 0) {
@@ -241,6 +241,7 @@ object MosStatementCompiler extends AbstractStatementCompiler[AssemblyLine] {
       case s : ReturnDispatchStatement =>
         MosReturnDispatch.compile(ctx, s) -> Nil
       case ReturnStatement(Some(e)) =>
+        val exprType = AbstractExpressionCompiler.getExpressionType(ctx, e)
         (m.returnType match {
           case _: BooleanType =>
             m.returnType.size match {
@@ -248,9 +249,9 @@ object MosStatementCompiler extends AbstractStatementCompiler[AssemblyLine] {
                 ctx.log.error("Cannot return anything from a void function", statement.position)
                 stackPointerFixBeforeReturn(ctx) ++ returnInstructions
               case 1 =>
-                MosExpressionCompiler.compile(ctx, e, someRegisterA, NoBranching) ++ stackPointerFixBeforeReturn(ctx, preserveA = true) ++ returnInstructions
+                MosExpressionCompiler.compile(ctx, e, Some(exprType, RegisterVariable(MosRegister.A, b)), NoBranching) ++ stackPointerFixBeforeReturn(ctx, preserveA = true) ++ returnInstructions
               case 2 =>
-                MosExpressionCompiler.compile(ctx, e, someRegisterAX, NoBranching) ++ stackPointerFixBeforeReturn(ctx, preserveA = true, preserveX = true) ++ returnInstructions
+                MosExpressionCompiler.compile(ctx, e, Some(exprType, RegisterVariable(MosRegister.AX, w)), NoBranching) ++ stackPointerFixBeforeReturn(ctx, preserveA = true, preserveX = true) ++ returnInstructions
               case _ =>
                 // TODO: is this case ever used?
                 MosExpressionCompiler.compileAssignment(ctx, e, VariableExpression(ctx.function.name + "`return")) ++
@@ -267,14 +268,14 @@ object MosStatementCompiler extends AbstractStatementCompiler[AssemblyLine] {
                 ctx.log.error("Cannot return anything from a void function", statement.position)
                 stackPointerFixBeforeReturn(ctx) ++ List(AssemblyLine.discardAF(), AssemblyLine.discardXF(), AssemblyLine.discardYF()) ++ returnInstructions
               case 1 =>
-                MosExpressionCompiler.compile(ctx, e, someRegisterA, NoBranching) ++ stackPointerFixBeforeReturn(ctx, preserveA = true) ++ List(AssemblyLine.discardXF(), AssemblyLine.discardYF()) ++ returnInstructions
+                MosExpressionCompiler.compile(ctx, e, Some(exprType, RegisterVariable(MosRegister.A, w)), NoBranching) ++ stackPointerFixBeforeReturn(ctx, preserveA = true) ++ List(AssemblyLine.discardXF(), AssemblyLine.discardYF()) ++ returnInstructions
               case 2 =>
                 // TODO: ???
                 val stackPointerFix = stackPointerFixBeforeReturn(ctx, preserveA = true, preserveY = true)
                 if (stackPointerFix.isEmpty) {
-                  MosExpressionCompiler.compile(ctx, e, someRegisterAX, NoBranching) ++ List(AssemblyLine.discardYF()) ++ returnInstructions
+                  MosExpressionCompiler.compile(ctx, e, Some(exprType, RegisterVariable(MosRegister.AX, w)), NoBranching) ++ List(AssemblyLine.discardYF()) ++ returnInstructions
                 } else {
-                  MosExpressionCompiler.compile(ctx, e, someRegisterYA, NoBranching) ++
+                  MosExpressionCompiler.compile(ctx, e, Some(exprType, RegisterVariable(MosRegister.YA, w)), NoBranching) ++
                     stackPointerFix ++
                     List(AssemblyLine.implied(TAX), AssemblyLine.implied(TYA), AssemblyLine.discardYF()) ++
                     returnInstructions

src/test/scala/millfork/test/SignExtensionSuite.scala

@@ -14,11 +14,11 @@ class SignExtensionSuite extends FunSuite with Matchers {
         | word output @$c000
         | void main () {
         |   sbyte b
-        |   b = -1
+        |   b = -2
         |   output = b
         | }
       """.stripMargin){m =>
-      m.readWord(0xc000) should equal(0xffff)
+      m.readWord(0xc000) should equal(0xfffe)
     }
   }
   test("Sbyte to Word 2") {
@@ -28,9 +28,9 @@ class SignExtensionSuite extends FunSuite with Matchers {
         |   output = b()
         | }
         | sbyte b() {
-        |   return -1
+        |   return -2
         | }
-      """.stripMargin){m => m.readWord(0xc000) should equal(0xffff)}
+      """.stripMargin){m => m.readWord(0xc000) should equal(0xfffe)}
   }
   test("Sbyte to Long") {
     EmuCrossPlatformBenchmarkRun(Cpu.Mos, Cpu.Z80, Cpu.Intel8080, Cpu.Sharp, Cpu.Intel8086)("""
@@ -75,12 +75,12 @@ class SignExtensionSuite extends FunSuite with Matchers {
         | word output @$c000
         | void main () {
         |   sbyte b
-        |   b = -1
+        |   b = -2
         |   memory_barrier()
         |   output = byte(b)
         | }
       """.stripMargin){m =>
-      m.readWord(0xc000) should equal(0x00ff)
+      m.readWord(0xc000) should equal(0x00fe)
     }
   }
 
@@ -89,12 +89,12 @@ class SignExtensionSuite extends FunSuite with Matchers {
         | word output @$c000
         | void main () {
         |   sbyte b
-        |   b = -1
+        |   b = -2
         |   memory_barrier()
         |   output = word(byte(b))
         | }
       """.stripMargin){m =>
-      m.readWord(0xc000) should equal(0x00ff)
+      m.readWord(0xc000) should equal(0x00fe)
     }
   }
 
@@ -102,11 +102,11 @@ class SignExtensionSuite extends FunSuite with Matchers {
     EmuUnoptimizedCrossPlatformRun(Cpu.Mos, Cpu.Z80, Cpu.Motorola6809)("""
         | word output @$c000
         | void main () {
-        |   output = f($ff)
+        |   output = f($f4)
         | }
         | noinline word f(byte x) = sbyte(x)
       """.stripMargin){m =>
-      m.readWord(0xc000) should equal(0xffff)
+      m.readWord(0xc000) should equal(0xfff4)
     }
   }
 
@@ -130,11 +130,11 @@ class SignExtensionSuite extends FunSuite with Matchers {
         |word output @$c000
         |word output2 @$c002
         |void main() {
-        |  output = -1
+        |  output = -3
         |  output2 = -30000
         |}
         |""".stripMargin) { m =>
-      m.readWord(0xc000) should equal(0xffff)
+      m.readWord(0xc000) should equal(0xfffd)
       m.readWord(0xc002).toShort.toInt should equal(-30000)
     }
   }
@@ -146,12 +146,12 @@ class SignExtensionSuite extends FunSuite with Matchers {
         |
         |void main() {
         |  static volatile signed16 tmp
-        |  tmp = -1
+        |  tmp = -3
         |  memory_barrier()
         |  output = tmp
         |}
         |""".stripMargin){ m =>
-      m.readLong(0xc000) should equal(-1)
+      m.readLong(0xc000) should equal(-3)
     }
   }
 
@@ -162,13 +162,33 @@ class SignExtensionSuite extends FunSuite with Matchers {
         |
         |void main() {
         |  static volatile signed16 tmp
-        |  tmp = -1
+        |  tmp = -3
         |  output = 0
         |  memory_barrier()
         |  output += tmp
         |}
         |""".stripMargin){ m =>
-      m.readLong(0xc000) should equal(-1)
+      m.readLong(0xc000) should equal(-3)
+    }
+  }
+
+  test("Trivial implicit sbyte to word extension") {
+    EmuCrossPlatformBenchmarkRun(Cpu.Mos, Cpu.Z80, Cpu.Motorola6809)(
+      """
+        |word output @$c000
+        |
+        |sbyte c
+        |
+        |noinline word f() {
+        |  return c
+        |}
+        |
+        |void main() {
+        |  c = -2
+        |  output = f()
+        |}
+        |""".stripMargin){ m =>
+      m.readWord(0xc000) should equal(0xFFFE)
     }
   }