6502: Faster page-aligned loops

src/main/scala/millfork/compiler/AbstractStatementCompiler.scala

@@ -15,8 +15,7 @@ abstract class AbstractStatementCompiler[T <: AbstractCode] {
     getStatementPreprocessor(ctx, statements)().flatMap(s => compile(ctx, s))
   }
 
-  def getStatementPreprocessor(ctx: CompilationContext, statements: List[ExecutableStatement]) =
+  def getStatementPreprocessor(ctx: CompilationContext, statements: List[ExecutableStatement]): AbstractStatementPreprocessor
-    new AbstractStatementPreprocessor(ctx, statements)
 
   def compile(ctx: CompilationContext, statement: ExecutableStatement): List[T]
 

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

@@ -12,13 +12,13 @@ import scala.collection.mutable.ListBuffer
 /**
   * @author Karol Stasiak
   */
-class AbstractStatementPreprocessor(ctx: CompilationContext, statements: List[ExecutableStatement]) {
+abstract class AbstractStatementPreprocessor(ctx: CompilationContext, statements: List[ExecutableStatement]) {
   type VV = Map[String, Constant]
-  private val optimize = true // TODO
+  protected val optimize = true // TODO
-  private val env: Environment = ctx.env
+  protected val env: Environment = ctx.env
-  private val localPrefix = ctx.function.name + "$"
+  protected val localPrefix = ctx.function.name + "$"
-  private val usedIdentifiers = if (optimize) statements.flatMap(_.getAllExpressions).flatMap(_.getAllIdentifiers) else Set()
+  protected val usedIdentifiers = if (optimize) statements.flatMap(_.getAllExpressions).flatMap(_.getAllIdentifiers) else Set()
-  private val trackableVars: Set[String] = if (optimize) {
+  protected val trackableVars: Set[String] = if (optimize) {
     env.getAllLocalVariables.map(_.name.stripPrefix(localPrefix))
       .filterNot(_.contains("."))
       .filterNot(_.contains("$"))
@@ -30,12 +30,12 @@ class AbstractStatementPreprocessor(ctx: CompilationContext, statements: List[Ex
   if (ErrorReporting.traceEnabled && trackableVars.nonEmpty) {
     ErrorReporting.trace("Tracking local variables: " + trackableVars.mkString(", "))
   }
-  private val reentrantVars: Set[String] = trackableVars.filter(v => env.get[Variable](v) match {
+  protected val reentrantVars: Set[String] = trackableVars.filter(v => env.get[Variable](v) match {
     case _: StackVariable => true
     case UninitializedMemoryVariable(_, _, VariableAllocationMethod.Auto, _) => ctx.options.flag(CompilationFlag.DangerousOptimizations)
     case _ => false
   })
-  private val nonreentrantVars: Set[String] = trackableVars -- reentrantVars
+  protected val nonreentrantVars: Set[String] = trackableVars -- reentrantVars
 
 
   def apply(): List[ExecutableStatement] = {
@@ -53,6 +53,8 @@ class AbstractStatementPreprocessor(ctx: CompilationContext, statements: List[Ex
     result.toList -> cv
   }
 
+  def maybeOptimizeForStatement(f: ForStatement): Option[(ExecutableStatement, VV)]
+
   def optimizeStmt(stmt: ExecutableStatement, currentVarValues: VV): (ExecutableStatement, VV) = {
     var cv = currentVarValues
     val pos = stmt.position
@@ -95,11 +97,15 @@ class AbstractStatementPreprocessor(ctx: CompilationContext, statements: List[Ex
         val (b, _) = optimizeStmts(body, Map())
         val (i, _) = optimizeStmts(inc, Map())
         DoWhileStatement(b, i, c, labels).pos(pos) -> Map()
-      case ForStatement(v, st, en, dir, body) =>
+      case f@ForStatement(v, st, en, dir, body) =>
-        val s = optimizeExpr(st, Map())
+        maybeOptimizeForStatement(f) match {
-        val e = optimizeExpr(en, Map())
+          case Some(x) => x
-        val (b, _) = optimizeStmts(body, Map())
+          case None =>
-        ForStatement(v, s, e, dir, b).pos(pos) -> Map()
+            val s = optimizeExpr(st, Map())
+            val e = optimizeExpr(en, Map())
+            val (b, _) = optimizeStmts(body, Map())
+            ForStatement(v, s, e, dir, b).pos(pos) -> Map()
+        }
       case _ => stmt -> Map()
     }
   }

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

@@ -297,4 +297,7 @@ object MosStatementCompiler extends AbstractStatementCompiler[AssemblyLine] {
   }
 
   override def nextLabel(prefix: String): String = MosCompiler.nextLabel(prefix)
+
+  override def getStatementPreprocessor(ctx: CompilationContext, statements: List[ExecutableStatement]): AbstractStatementPreprocessor =
+    new MosStatementPreprocessor(ctx, statements)
 }

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

@@ -0,0 +1,79 @@
+package millfork.compiler.mos
+
+import millfork.compiler.{AbstractStatementPreprocessor, CompilationContext}
+import millfork.env.{NumericConstant, Variable}
+import millfork.error.ErrorReporting
+import millfork.node._
+
+/**
+  * @author Karol Stasiak
+  */
+class MosStatementPreprocessor(ctx: CompilationContext, statements: List[ExecutableStatement]) extends AbstractStatementPreprocessor(ctx, statements) {
+
+  def maybeOptimizeForStatement(f: ForStatement): Option[(ExecutableStatement, VV)] = {
+    if (optimize && !f.variable.contains(".") && env.get[Variable](f.variable).typ.size == 2) {
+      (env.eval(f.start), env.eval(f.end)) match {
+        case (Some(NumericConstant(s, _)), Some(NumericConstant(e, _))) if (s & 0xffff) == s && (e & 0xffff) == e =>
+          f.direction match {
+            case ForDirection.Until | ForDirection.ParallelUntil =>
+              if (s.&(0xff) == 0 && e.&(0xff) == 0) {
+                ErrorReporting.debug(s"Loop across whole memory pages", f.position)
+                Some(ForStatement(
+                  f.variable + ".hi",
+                  FunctionCallExpression("hi", List(f.start)).pos(f.start.position),
+                  FunctionCallExpression("hi", List(f.end)).pos(f.end.position),
+                  f.direction,
+                  List(
+                    ForStatement(
+                      f.variable + ".lo",
+                      LiteralExpression(0, 1).pos(f.start.position),
+                      LiteralExpression(256, 2).pos(f.end.position),
+                      f.direction,
+                      f.body
+                    ).pos(f.position))
+                ).pos(f.position) -> Map())
+              } else None
+            case ForDirection.To | ForDirection.ParallelTo =>
+              if (s.&(0xff) == 0 && e.&(0xff) == 0xff) {
+                ErrorReporting.debug(s"Loop across whole memory pages", f.position)
+                Some(ForStatement(
+                  f.variable + ".hi",
+                  FunctionCallExpression("hi", List(f.start)).pos(f.start.position),
+                  FunctionCallExpression("hi", List(f.end)).pos(f.end.position),
+                  f.direction,
+                  List(
+                    ForStatement(
+                      f.variable + ".lo",
+                      LiteralExpression(0, 1).pos(f.start.position),
+                      LiteralExpression(255, 1).pos(f.end.position),
+                      f.direction,
+                      f.body
+                    ).pos(f.position))
+                ).pos(f.position) -> Map())
+              } else None
+            case ForDirection.DownTo | ForDirection.ParallelTo =>
+              if (s.&(0xff) == 0xff && e.&(0xff) == 0) {
+                ErrorReporting.debug(s"Loop across whole memory pages", f.position)
+                Some(ForStatement(
+                  f.variable + ".hi",
+                  FunctionCallExpression("hi", List(f.start)).pos(f.start.position),
+                  FunctionCallExpression("hi", List(f.end)).pos(f.end.position),
+                  f.direction,
+                  List(
+                    ForStatement(
+                      f.variable + ".lo",
+                      LiteralExpression(255, 1).pos(f.start.position),
+                      LiteralExpression(0, 1).pos(f.end.position),
+                      f.direction,
+                      f.body
+                    ).pos(f.position))
+                ).pos(f.position) -> Map())
+              } else None
+            case _ => None
+          }
+        case _ => None
+      }
+    } else None
+  }
+
+}

src/main/scala/millfork/compiler/z80/Z80StatementCompiler.scala

@@ -2,7 +2,7 @@ package millfork.compiler.z80
 
 import millfork.assembly.BranchingOpcodeMapping
 import millfork.assembly.z80._
-import millfork.compiler.{AbstractExpressionCompiler, AbstractStatementCompiler, BranchSpec, CompilationContext}
+import millfork.compiler._
 import millfork.env._
 import millfork.node._
 import millfork.assembly.z80.ZOpcode._
@@ -215,4 +215,7 @@ object Z80StatementCompiler extends AbstractStatementCompiler[ZLine] {
 
   override def compileExpressionForBranching(ctx: CompilationContext, expr: Expression, branching: BranchSpec): List[ZLine] =
     Z80ExpressionCompiler.compile(ctx, expr, ZExpressionTarget.NOTHING, branching)
+
+  override def getStatementPreprocessor(ctx: CompilationContext, statements: List[ExecutableStatement]): AbstractStatementPreprocessor =
+    new Z80StatementPreprocessor(ctx, statements)
 }

src/main/scala/millfork/compiler/z80/Z80StatementPreprocessor.scala

@@ -0,0 +1,12 @@
+package millfork.compiler.z80
+
+import millfork.compiler.{AbstractStatementPreprocessor, CompilationContext}
+import millfork.node.{ExecutableStatement, ForStatement}
+
+/**
+  * @author Karol Stasiak
+  */
+class Z80StatementPreprocessor(ctx: CompilationContext, statements: List[ExecutableStatement]) extends AbstractStatementPreprocessor(ctx, statements) {
+
+  def maybeOptimizeForStatement(f: ForStatement): Option[(ExecutableStatement, VV)] = None
+}

src/test/scala/millfork/test/ForLoopSuite.scala

@@ -174,6 +174,24 @@ class ForLoopSuite extends FunSuite with Matchers {
     }
   }
 
+  test("Screen fill") {
+    EmuCrossPlatformBenchmarkRun(Cpu.Mos, Cpu.Z80)(
+      """
+        | array output[$400]@$c000
+        | void main () {
+        |   pointer p
+        |   for p,$c000,paralleluntil,$c400{
+        |     p[0] = 34
+        |   }
+        | }
+        | void _panic(){while(true){}}
+      """.stripMargin){ m=>
+      for(i <- 0xc000 to 0xc3ff) {
+        m.readByte(i) should equal (34)
+      }
+    }
+  }
+
   test("Various bulk operations") {
     EmuCrossPlatformBenchmarkRun(Cpu.Mos, Cpu.Z80)(
       """