From 3f5877dbccbdd72a5e5badcc25db50c858f65087 Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Mon, 23 May 2022 21:24:36 +0200
Subject: [PATCH] vm: fix array iteration

---
 .../src/prog8/codegen/virtual/CodeGen.kt      | 32 ++++++++++++-------
 examples/primes.p8                            | 12 +++++++
 examples/test.p8                              | 13 ++++----
 3 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/codeGenVirtual/src/prog8/codegen/virtual/CodeGen.kt b/codeGenVirtual/src/prog8/codegen/virtual/CodeGen.kt
index d714186f1..1b01bb940 100644
--- a/codeGenVirtual/src/prog8/codegen/virtual/CodeGen.kt
+++ b/codeGenVirtual/src/prog8/codegen/virtual/CodeGen.kt
@@ -221,17 +221,27 @@ class CodeGen(internal val program: PtProgram,
                     val elementDt = ArrayToElementTypes.getValue(iterable.type)
                     val elementSize = program.memsizer.memorySize(elementDt)
                     val lengthBytes = iterableVar.length!! * elementSize
-                    val lengthReg = vmRegisters.nextFree()
-                    code += VmCodeInstruction(Opcode.LOAD, VmDataType.BYTE, reg1=indexReg, value=0)
-                    code += VmCodeInstruction(Opcode.LOAD, VmDataType.BYTE, reg1=lengthReg, value=lengthBytes)
-                    code += VmCodeLabel(loopLabel)
-                    code += VmCodeInstruction(Opcode.BEQ, VmDataType.BYTE, reg1=indexReg, reg2=lengthReg, labelSymbol = endLabel)
-                    code += VmCodeInstruction(Opcode.LOADX, vmType(elementDt), reg1=tmpReg, reg2=indexReg, value=arrayAddress)
-                    code += VmCodeInstruction(Opcode.STOREM, vmType(elementDt), reg1=tmpReg, value = loopvarAddress)
-                    code += translateNode(forLoop.statements)
-                    code += addConstReg(VmDataType.BYTE, indexReg, elementSize)
-                    code += VmCodeInstruction(Opcode.JUMP, labelSymbol = loopLabel)
-                    code += VmCodeLabel(endLabel)
+                    if(lengthBytes<256) {
+                        val lengthReg = vmRegisters.nextFree()
+                        code += VmCodeInstruction(Opcode.LOAD, VmDataType.BYTE, reg1=indexReg, value=0)
+                        code += VmCodeInstruction(Opcode.LOAD, VmDataType.BYTE, reg1=lengthReg, value=lengthBytes)
+                        code += VmCodeLabel(loopLabel)
+                        code += VmCodeInstruction(Opcode.LOADX, vmType(elementDt), reg1=tmpReg, reg2=indexReg, value=arrayAddress)
+                        code += VmCodeInstruction(Opcode.STOREM, vmType(elementDt), reg1=tmpReg, value = loopvarAddress)
+                        code += translateNode(forLoop.statements)
+                        code += addConstReg(VmDataType.BYTE, indexReg, elementSize)
+                        code += VmCodeInstruction(Opcode.BNE, VmDataType.BYTE, reg1=indexReg, reg2=lengthReg, labelSymbol = loopLabel)
+                    } else if(lengthBytes==256) {
+                        code += VmCodeInstruction(Opcode.LOAD, VmDataType.BYTE, reg1=indexReg, value=0)
+                        code += VmCodeLabel(loopLabel)
+                        code += VmCodeInstruction(Opcode.LOADX, vmType(elementDt), reg1=tmpReg, reg2=indexReg, value=arrayAddress)
+                        code += VmCodeInstruction(Opcode.STOREM, vmType(elementDt), reg1=tmpReg, value = loopvarAddress)
+                        code += translateNode(forLoop.statements)
+                        code += addConstReg(VmDataType.BYTE, indexReg, elementSize)
+                        code += VmCodeInstruction(Opcode.BNZ, VmDataType.BYTE, reg1=indexReg, labelSymbol = loopLabel)
+                    } else {
+                        throw AssemblyError("iterator length should never exceed 256")
+                    }
                 }
             }
             else -> throw AssemblyError("weird for iterable")
diff --git a/examples/primes.p8 b/examples/primes.p8
index 7c1619d10..143249f83 100644
--- a/examples/primes.p8
+++ b/examples/primes.p8
@@ -40,6 +40,10 @@ main {
                 return 0        ; we wrapped; no more primes available in the sieve
         }
 
+        txt.print("candidate: ")
+        txt.print_ub(candidate_prime)
+        txt.nl()
+
         ; found next one, mark the multiples and return it.
         sieve[candidate_prime] = true
         uword multiple = candidate_prime
@@ -49,6 +53,14 @@ main {
             sieve[lsb(multiple)] = true
             multiple += candidate_prime
         }
+
+        ubyte xx
+        for xx in sieve {
+            txt.print_ub(xx)
+            txt.spc()
+        }
+        txt.nl()
+
         return candidate_prime
     }
 }
diff --git a/examples/test.p8 b/examples/test.p8
index b678ccdc8..72bf33af3 100644
--- a/examples/test.p8
+++ b/examples/test.p8
@@ -31,16 +31,17 @@ main {
     sub start() {
         ; mcCarthy()
 
-        ubyte[256] sieve
+        ubyte[20] sieve
+        uword count=0
         ubyte xx
-        for xx in 0 to 255 {
-            sieve[xx] = false
-        }
-        for xx in 0 to 255 {
-            txt.print_ub(sieve[xx])
+        for xx in sieve {
+            txt.print_ub(xx)
             txt.spc()
+            count++
         }
         txt.nl()
+        txt.print_uw(count)
+        txt.nl()
 
 
 ;        ; a "pixelshader":