Fix and improve nonet operations

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

@@ -94,6 +94,18 @@ object AlwaysGoodOptimizations {
       HasOpcode(LSR) & HasAddrMode(Implied) & DoesntMatterWhatItDoesWith(State.C)) ~~> { (code, ctx) =>
       AssemblyLine.immediate(LDA, (ctx.get[Int](0) & 0xff) >> 1) :: Nil
     },
+
+    (Elidable & MatchA(0) &
+      HasOpcode(ASL) & HasAddrMode(Implied)) ~~> { (code, ctx) =>
+      val v = ctx.get[Int](0)
+      AssemblyLine.immediate(LDA, (v << 1) & 0xff) :: AssemblyLine.implied(if(v.&(0x80) != 0) SEC else CLC) :: Nil
+    },
+    (Elidable & MatchA(0) &
+      HasOpcode(LSR) & HasAddrMode(Implied)) ~~> { (code, ctx) =>
+      val v = ctx.get[Int](0)
+      AssemblyLine.immediate(LDA, (v & 0xff) >> 1) :: AssemblyLine.implied(if(v.&(1) != 0) SEC else CLC) :: Nil
+    },
+
     (Elidable & MatchA(0) &
       HasClear(State.C) & HasOpcode(ROL) & HasAddrMode(Implied) & DoesntMatterWhatItDoesWith(State.C)) ~~> { (code, ctx) =>
       AssemblyLine.immediate(LDA, ctx.get[Int](0) << 1) :: Nil
@@ -110,6 +122,28 @@ object AlwaysGoodOptimizations {
       HasSet(State.C) & HasOpcode(ROR) & HasAddrMode(Implied) & DoesntMatterWhatItDoesWith(State.C)) ~~> { (code, ctx) =>
       AssemblyLine.immediate(LDA, 0x80 + (ctx.get[Int](0) & 0xff) / 2) :: Nil
     },
+
+    (Elidable & MatchA(0) &
+      HasClear(State.C) & HasOpcode(ROL) & HasAddrMode(Implied)) ~~> { (code, ctx) =>
+      val v = ctx.get[Int](0)
+      AssemblyLine.immediate(LDA, v << 1) :: AssemblyLine.implied(if(v.&(0x80) != 0) SEC else CLC) :: Nil
+    },
+    (Elidable & MatchA(0) &
+      HasClear(State.C) & HasOpcode(ROR) & HasAddrMode(Implied)) ~~> { (code, ctx) =>
+      val v = ctx.get[Int](0)
+      AssemblyLine.immediate(LDA, (ctx.get[Int](0) & 0xff) >> 1) :: AssemblyLine.implied(if(v.&(1) != 0) SEC else CLC) :: Nil
+    },
+    (Elidable & MatchA(0) &
+      HasSet(State.C) & HasOpcode(ROL) & HasAddrMode(Implied)) ~~> { (code, ctx) =>
+      val v = ctx.get[Int](0)
+      AssemblyLine.immediate(LDA, (v * 2 + 1) & 0xff) :: AssemblyLine.implied(if(v.&(0x80) != 0) SEC else CLC) :: Nil
+    },
+    (Elidable & MatchA(0) &
+      HasSet(State.C) & HasOpcode(ROR) & HasAddrMode(Implied)) ~~> { (code, ctx) =>
+      val v = ctx.get[Int](0)
+      AssemblyLine.immediate(LDA, 0x80 + (v & 0xff) / 2) :: AssemblyLine.implied(if(v.&(1) != 0) SEC else CLC) :: Nil
+    },
+
     (Elidable &
       MatchA(0) & MatchParameter(1) &
       HasOpcode(ADC) & HasAddrMode(Immediate) &
@@ -535,6 +569,19 @@ object AlwaysGoodOptimizations {
       LAX, DoesntMatterWhatItDoesWith(State.N, State.Z),
       Not(ConcernsX) & Not(ConcernsA),
       SAX, DoesntMatterWhatItDoesWith(State.A, State.X)),
+
+    (Elidable & HasOpcode(ASL) & HasAddrMode(Implied)) ~
+      (Linear & Not(ConcernsA) & Not(ConcernsC) & Not(ReadsNOrZ)).* ~
+      (Elidable & HasOpcode(ROR) & HasAddrMode(Implied) & DoesntMatterWhatItDoesWith(State.C, State.N, State.Z)) ~~> (_.init.tail),
+    (Elidable & HasOpcode(LSR) & HasAddrMode(Implied)) ~
+      (Linear & Not(ConcernsA) & Not(ConcernsC) & Not(ReadsNOrZ)).* ~
+      (Elidable & HasOpcode(ROL) & HasAddrMode(Implied) & DoesntMatterWhatItDoesWith(State.C, State.N, State.Z)) ~~> (_.init.tail),
+    (Elidable & HasOpcode(ROL) & HasAddrMode(Implied)) ~
+      (Linear & Not(ConcernsA) & Not(ConcernsC) & Not(ReadsNOrZ)).* ~
+      (Elidable & HasOpcode(ROR) & HasAddrMode(Implied) & DoesntMatterWhatItDoesWith(State.N, State.Z)) ~~> (_.init.tail),
+    (Elidable & HasOpcode(ROR) & HasAddrMode(Implied)) ~
+      (Linear & Not(ConcernsA) & Not(ConcernsC) & Not(ReadsNOrZ)).* ~
+      (Elidable & HasOpcode(ROL) & HasAddrMode(Implied) & DoesntMatterWhatItDoesWith(State.N, State.Z)) ~~> (_.init.tail),
   )
 
   val PointlessStackStashing = new RuleBasedAssemblyOptimization("Pointless stack stashing",

src/main/scala/millfork/compiler/AbstractStatementPreprocessor.scala

@@ -192,6 +192,10 @@ abstract class AbstractStatementPreprocessor(ctx: CompilationContext, statements
         FunctionCallExpression("nonet", args.map(arg => optimizeExpr(arg, Map()))).pos(pos)
       case FunctionCallExpression(name, args) =>
         FunctionCallExpression(name, args.map(arg => optimizeExpr(arg, currentVarValues))).pos(pos)
+      case SumExpression(expressions, decimal) =>
+        // don't collapse additions, let the later stages deal with it
+        // expecially important when inside a nonet operation
+        SumExpression(expressions.map{case (minus, arg) => minus -> optimizeExpr(arg, currentVarValues)}, decimal)
       case _ => expr // TODO
     }
   }

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

@@ -260,6 +260,18 @@ object BuiltIns {
 
   def compileNonetOps(ctx: CompilationContext, lhs: Expression, rhs: Expression): List[AssemblyLine] = {
     val env = ctx.env
+    lhs match {
+      case FunctionCallExpression("nonet", List(lparam)) =>
+        (env.eval(lhs), env.eval(rhs)) match {
+          case (None, Some(NumericConstant(0, _))) =>
+            return MosExpressionCompiler.compile(ctx, lparam, None, NoBranching)
+          case (None, Some(NumericConstant(n, _))) if n > 0 =>
+            return MosExpressionCompiler.compile(ctx, lparam, None, NoBranching) ++
+              (AssemblyLine.implied(ROR) :: List.fill(n.toInt - 1)(AssemblyLine.implied(LSR)))
+          case _ =>
+        }
+      case _ =>
+    }
     val b = env.get[Type]("byte")
     val (ldaHi, ldaLo) = env.eval(lhs) match {
       case Some(c) =>

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

@@ -361,7 +361,7 @@ case class CompoundConstant(operator: MathOperator.Value, lhs: Constant, rhs: Co
           case _ => return this
         }
         operator match {
-          case MathOperator.Plus9 | MathOperator.DecimalPlus9 =>
+          case MathOperator.Plus9 | MathOperator.DecimalPlus9 | MathOperator.Shl9 | MathOperator.DecimalShl9 =>
             size = 2
           case MathOperator.Times | MathOperator.Shl =>
             val mask = (1 << (size * 8)) - 1

src/main/scala/millfork/env/Environment.scala

@@ -549,9 +549,9 @@ class Environment(val parent: Option[Environment], val prefix: String, val cpuFa
                 constantOperation(MathOperator.Shl9, ps)
               case List(FunctionCallExpression("<<'", ps@List(_, _))) =>
                 constantOperation(MathOperator.DecimalShl9, ps)
-              case List(SumExpression(ps@List((true,_),(true,_)), false)) =>
+              case List(SumExpression(ps@List((false,_),(false,_)), false)) =>
                 constantOperation(MathOperator.Plus9, ps.map(_._2))
-              case List(SumExpression(ps@List((true,_),(true,_)), true)) =>
+              case List(SumExpression(ps@List((false,_),(false,_)), true)) =>
                 constantOperation(MathOperator.DecimalPlus9, ps.map(_._2))
               case List(_) =>
                 None

src/test/scala/millfork/test/NonetSuite.scala

@@ -102,4 +102,45 @@ class NonetSuite extends FunSuite with Matchers {
       m.readWord(0xc006) should equal(0x8080)
     }
   }
+
+  test("Nonet shift right") {
+    EmuCrossPlatformBenchmarkRun(Cpu.Mos, Cpu.Z80)(
+      """
+        | byte output0 @$c000
+        | byte output1 @$c002
+        | byte output2 @$c004
+        | noinline void perform(byte x) {
+        |   output0 = nonet(150+160) >>>> 1
+        |   output1 = nonet(153+output0) >>>> 1
+        |   output2 = nonet(150+x) >>>> 1
+        | }
+        | void main () {
+        |   perform(200)
+        | }
+        | noinline byte three() { return 3 }
+      """.stripMargin) { m =>
+      m.readByte(0xc000) should equal(155)
+      m.readByte(0xc002) should equal(154)
+      m.readByte(0xc004) should equal(175)
+    }
+  }
+
+  test("Nonet shift right 2") {
+    EmuCrossPlatformBenchmarkRun(Cpu.Mos, Cpu.Z80)(
+      """
+        | byte output0 @$c000
+        | byte output1 @$c002
+        | noinline void perform(byte x) {
+        |   output0 = nonet(180 << 1) >>>> 1
+        |   output1 = nonet(output0 << 1) >>>> 1
+        | }
+        | void main () {
+        |   perform(200)
+        | }
+        | noinline byte three() { return 3 }
+      """.stripMargin) { m =>
+      m.readByte(0xc000) should equal(180)
+      m.readByte(0xc002) should equal(180)
+    }
+  }
 }