Always require a full optimization context when optimizing

src/main/scala/millfork/assembly/AssemblyOptimization.scala

@@ -1,7 +1,7 @@
 package millfork.assembly
 
 import millfork.CompilationOptions
-import millfork.env.NormalFunction
+import millfork.env.{NormalFunction, ThingInMemory}
 import millfork.node.NiceFunctionProperty
 
 /**
@@ -9,12 +9,11 @@ import millfork.node.NiceFunctionProperty
   */
 case class OptimizationContext(options: CompilationOptions,
                                labelMap: Map[String, Int],
+                               zreg: Option[ThingInMemory],
                                niceFunctionProperties: Set[(NiceFunctionProperty, String)])
 
 trait AssemblyOptimization[T <: AbstractCode] {
   def name: String
 
-  def optimize(f: NormalFunction, code: List[T], context: OptimizationContext): List[T] = optimize(f, code, context.options)
-
-  def optimize(f: NormalFunction, code: List[T], options: CompilationOptions): List[T] = optimize(f, code, OptimizationContext(options, Map(), Set()))
+  def optimize(f: NormalFunction, code: List[T], context: OptimizationContext): List[T]
 }

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

@@ -30,7 +30,7 @@ object FlowInfoRequirement extends Enumeration {
 }
 
 trait AssemblyRuleSet{
-  def flatten: List[AssemblyRule]
+  def flatten: Seq[AssemblyRule]
 }
 
 class RuleBasedAssemblyOptimization(val name: String, val needsFlowInfo: FlowInfoRequirement.Value, val rules: AssemblyRuleSet*) extends AssemblyOptimization[ZLine]{
@@ -38,18 +38,18 @@ class RuleBasedAssemblyOptimization(val name: String, val needsFlowInfo: FlowInf
   private val actualRules = rules.flatMap(_.flatten)
   actualRules.foreach(_.pattern.validate(needsFlowInfo))
 
-  override def optimize(f: NormalFunction, code: List[ZLine], options: CompilationOptions): List[ZLine] = {
+  override def optimize(f: NormalFunction, code: List[ZLine], optimizationContext: OptimizationContext): List[ZLine] = {
     val effectiveCode = code.map(a => a.copy(parameter = a.parameter.quickSimplify))
-    val taggedCode = FlowAnalyzer.analyze(f, effectiveCode, options, needsFlowInfo)
-    optimizeImpl(f, taggedCode, options)
+    val taggedCode = FlowAnalyzer.analyze(f, effectiveCode, optimizationContext.options, needsFlowInfo)
+    optimizeImpl(f, taggedCode, optimizationContext)
   }
 
-  def optimizeImpl(f: NormalFunction, code: List[(FlowInfo, ZLine)], options: CompilationOptions): List[ZLine] = {
+  def optimizeImpl(f: NormalFunction, code: List[(FlowInfo, ZLine)], optimizationContext: OptimizationContext): List[ZLine] = {
     code match {
       case Nil => Nil
       case head :: tail =>
         for ((rule, index) <- actualRules.zipWithIndex) {
-          val ctx = new AssemblyMatchingContext(options)
+          val ctx = new AssemblyMatchingContext(optimizationContext.options)
           rule.pattern.matchTo(ctx, code) match {
             case Some(rest: List[(FlowInfo, ZLine)]) =>
               val matchedChunkToOptimize: List[ZLine] = code.take(code.length - rest.length).map(_._2)
@@ -65,14 +65,14 @@ class RuleBasedAssemblyOptimization(val name: String, val needsFlowInfo: FlowInf
               ErrorReporting.trace("     ↓")
               optimizedChunk.filter(_.isPrintable).foreach(l => ErrorReporting.trace(l.toString))
               if (needsFlowInfo != FlowInfoRequirement.NoRequirement) {
-                return optimizedChunk ++ optimizeImpl(f, rest, options)
+                return optimizedChunk ++ optimizeImpl(f, rest, optimizationContext)
               } else {
-                return optimize(f, optimizedChunk ++ rest.map(_._2), options)
+                return optimize(f, optimizedChunk ++ rest.map(_._2), optimizationContext)
               }
             case None => ()
           }
         }
-        head._2 :: optimizeImpl(f, tail, options)
+        head._2 :: optimizeImpl(f, tail, optimizationContext)
     }
   }
 }
@@ -216,11 +216,11 @@ object HelperCheckers {
   }
 }
 case class AssemblyRule(pattern: AssemblyPattern, result: (List[ZLine], AssemblyMatchingContext) => List[ZLine]) extends AssemblyRuleSet {
-  override def flatten: List[AssemblyRule] = List(this)
+  override def flatten: Seq[AssemblyRule] = List(this)
 }
 
-case class MultipleAssemblyRules(list: List[AssemblyRuleSet]) extends AssemblyRuleSet {
-  override def flatten: List[AssemblyRule] = list.flatMap(_.flatten)
+case class MultipleAssemblyRules(list: Seq[AssemblyRuleSet]) extends AssemblyRuleSet {
+  override def flatten: Seq[AssemblyRule] = list.flatMap(_.flatten)
 }
 
 trait AssemblyPattern {

src/main/scala/millfork/output/AbstractAssembler.scala

@@ -474,7 +474,7 @@ abstract class AbstractAssembler[T <: AbstractCode](private val program: Program
         compiler))
     unoptimizedCodeSize += unoptimized.map(_.sizeInBytes).sum
     val code = optimizations.foldLeft(unoptimized) { (c, opt) =>
-      opt.optimize(f, c, OptimizationContext(options, labelMap, niceFunctionProperties))
+      opt.optimize(f, c, OptimizationContext(options, labelMap, env.maybeGet[ThingInMemory]("__reg"), niceFunctionProperties))
     }
     performFinalOptimizationPass(f, optimizations.nonEmpty, options, code)
   }

src/main/scala/millfork/output/MosAssembler.scala

@@ -21,7 +21,7 @@ class MosAssembler(program: Program,
 
 
   override def performFinalOptimizationPass(f: NormalFunction, actuallyOptimize: Boolean, options: CompilationOptions, code: List[AssemblyLine]):List[AssemblyLine] = {
-    val optimizationContext = OptimizationContext(options, Map(), Set())
+    val optimizationContext = OptimizationContext(options, Map(), f.environment.maybeGet[ThingInMemory]("__reg"), Set())
     if (actuallyOptimize) {
       val finalCode = if (options.flag(CompilationFlag.EmitHudsonOpcodes)) HudsonOptimizations.removeLoadZero(f, code, optimizationContext) else code
       JumpShortening(f, JumpShortening(f, JumpFixing(f, finalCode, options), optimizationContext), optimizationContext)