If a variable is used wholly within a loop body and initialized conditionally, do not remove the last store to it

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

@@ -17,6 +17,10 @@ object EmptyMemoryStoreRemoval extends AssemblyOptimization[AssemblyLine] {
   override def name = "Removing pointless stores to automatic variables"
 
   private val storeAddrModes = Set(Absolute, ZeroPage, AbsoluteX, AbsoluteY, ZeroPageX, ZeroPageY)
+  private val directStorageOpcodes = Set(
+    STA, STX, STY, SAX, STZ, SHX, SHY, AHX,
+    STA_W, STX_W, STY_W, STZ_W
+  )
 
   override def optimize(f: NormalFunction, code: List[AssemblyLine], optimizationContext: OptimizationContext): List[AssemblyLine] = {
     val paramVariables = f.params match {
@@ -55,20 +59,26 @@ object EmptyMemoryStoreRemoval extends AssemblyOptimization[AssemblyLine] {
 
     for(v <- localVariables) {
       val lifetime = VariableLifetime.apply(v.name, code)
+      val firstaccess = lifetime.head
       val lastaccess = lifetime.last
-      if (lastaccess >= 0) {
+      if (firstaccess >= 0 && lastaccess >= 0) {
+        val firstVariableAccess = code(firstaccess)
         val lastVariableAccess = code(lastaccess)
-        if (lastVariableAccess.elidable && storeAddrModes(lastVariableAccess.addrMode) && (lastVariableAccess.opcode match {
+        if (lastVariableAccess.elidable &&
-          case STA | STX | STY | SAX | STZ | SHX | SHY | AHX =>
+          storeAddrModes(lastVariableAccess.addrMode) &&
+          storeAddrModes(firstVariableAccess.addrMode) &&
+          directStorageOpcodes(firstVariableAccess.opcode) &&
+          (lastVariableAccess.opcode match {
+          case op if directStorageOpcodes(op) =>
             true
           case TSB | TRB =>
             if (importances eq null) importances = ReverseFlowAnalyzer.analyze(f, code, optimizationContext)
             importances(lastaccess).z != Important
-          case INC | DEC =>
+          case INC | DEC | INC_W | DEC_W =>
             if (importances eq null) importances = ReverseFlowAnalyzer.analyze(f, code, optimizationContext)
             importances(lastaccess).z != Important &&
               importances(lastaccess).n != Important
-          case ASL | LSR | ROL | ROR | DCP =>
+          case ASL | LSR | ROL | ROR | DCP | ASL_W | LSR_W | ROL_W | ROR_W =>
             if (importances eq null) importances = ReverseFlowAnalyzer.analyze(f, code, optimizationContext)
             importances(lastaccess).z != Important &&
               importances(lastaccess).n != Important &&
@@ -78,7 +88,7 @@ object EmptyMemoryStoreRemoval extends AssemblyOptimization[AssemblyLine] {
             importances(lastaccess).z != Important &&
               importances(lastaccess).n != Important &&
               importances(lastaccess).a != Important
-          case DCP | SLO | SRE | RLA =>
+          case SLO | SRE | RLA =>
             if (importances eq null) importances = ReverseFlowAnalyzer.analyze(f, code, optimizationContext)
             importances(lastaccess).z != Important &&
               importances(lastaccess).n != Important &&

src/main/scala/millfork/assembly/z80/opt/EmptyMemoryStoreRemoval.scala

@@ -25,11 +25,17 @@ object EmptyMemoryStoreRemoval extends AssemblyOptimization[ZLine] {
     val badVariables = mutable.Set[String]()
 
     for((v, lifetime) <- vs.variablesWithLifetimes if lifetime.nonEmpty) {
+      val firstaccess = lifetime.head
       val lastaccess = lifetime.last
-      if (lastaccess >= 0) {
+      if (firstaccess >= 0 && lastaccess >= 0) {
+        val firstVariableAccess = code(firstaccess)
         val lastVariableAccess = code(lastaccess)
         import millfork.assembly.z80.ZOpcode._
-        if (lastVariableAccess match {
+        if ((firstVariableAccess match {
+          case ZLine(LD, TwoRegisters(MEM_HL, _), _, Elidability.Elidable, _) => true
+          case ZLine(LD | LD_16, TwoRegisters(MEM_ABS_8 | MEM_ABS_16, _), _, Elidability.Elidable, _) => true
+          case _ => false
+        }) && (lastVariableAccess match {
           case ZLine(LD, TwoRegisters(MEM_HL, _), _, Elidability.Elidable, _) => true
           case ZLine(LD | LD_16, TwoRegisters(MEM_ABS_8 | MEM_ABS_16, _), _, Elidability.Elidable, _) => true
           case ZLine(INC | DEC, OneRegister(MEM_HL), _, Elidability.Elidable, _) =>
@@ -39,7 +45,7 @@ object EmptyMemoryStoreRemoval extends AssemblyOptimization[ZLine] {
             val importances = vs.codeWithFlow(lastaccess)._1.importanceAfter
             Seq(importances.sf, importances.zf, importances.cf).forall(_ == Unimportant)
           case _ => false
-        }) {
+        })) {
           badVariables += v.name
           toRemove += lastaccess
         }

src/test/scala/millfork/test/SecondAssemblyOptimizationSuite.scala

@@ -103,4 +103,45 @@ class SecondAssemblyOptimizationSuite extends FunSuite with Matchers {
       m.readByte(0xc001) should equal(4)
     }
   }
+
+  test("Conditional variable initialization") {
+    EmuCrossPlatformBenchmarkRun(Cpu.Mos, Cpu.Z80)(
+      """
+        | array output [16] @$c000
+        | void main () {
+        |    byte entropy
+        |    byte noise
+        |    byte i
+        |    entropy = 0
+        |    for i,0,until,output.length {
+        |      if entropy == 0 {
+        |        entropy = 8
+        |        noise = rand()
+        |      }
+        |      output[i] = noise & 1
+        |      noise >>= 1
+        |      entropy -= 1
+        |    }
+        | }
+        | noinline byte rand() { return $42 }
+      """.stripMargin) { m =>
+      m.readByte(0xc000) should equal(0)
+      m.readByte(0xc001) should equal(1)
+      m.readByte(0xc002) should equal(0)
+      m.readByte(0xc003) should equal(0)
+      m.readByte(0xc004) should equal(0)
+      m.readByte(0xc005) should equal(0)
+      m.readByte(0xc006) should equal(1)
+      m.readByte(0xc007) should equal(0)
+      m.readByte(0xc008) should equal(0)
+      m.readByte(0xc009) should equal(1)
+      m.readByte(0xc00a) should equal(0)
+      m.readByte(0xc00b) should equal(0)
+      m.readByte(0xc00c) should equal(0)
+      m.readByte(0xc00d) should equal(0)
+      m.readByte(0xc00e) should equal(1)
+      m.readByte(0xc00f) should equal(0)
+
+    }
+  }
 }