allow simple binary expressions as array indexing too, but not more

compiler/src/prog8/ast/AstToplevel.kt

@@ -37,6 +37,18 @@ interface Node {
         throw FatalAstException("scope missing from $this")
     }
 
+    fun definingBlock(): Block {
+        if(this is Block)
+            return this
+        return findParentNode<Block>(this)!!
+    }
+
+    fun containingStatement(): Statement {
+        if(this is Statement)
+            return this
+        return findParentNode<Statement>(this)!!
+    }
+
     fun replaceChildNode(node: Node, replacement: Node)
 }
 
@@ -47,6 +59,7 @@ interface IFunctionCall {
 
 
 class AsmGenInfo {
+    var usedAutoArrayIndexer = false
     var usedRegsaveA = false
     var usedRegsaveX = false
     var usedRegsaveY = false

compiler/src/prog8/ast/processing/AstChecker.kt

@@ -1054,10 +1054,8 @@ internal class AstChecker(private val program: Program,
             errors.err("array indexing is limited to byte size 0..255", arrayIndexedExpression.position)
 
         // check type of array indexer
-        if(arrayIndexedExpression.indexer.indexVar==null && arrayIndexedExpression.indexer.indexNum==null) {
+        if(arrayIndexedExpression.indexer.indexVar==null && arrayIndexedExpression.indexer.indexNum==null)
-            errors.err("array indexing can only be done with a number or a variable, not an arbitrary expression. Use a temp var?", arrayIndexedExpression.indexer.position)
+            errors.err("array indexing can only be done with a number or a variable, or a simple binary expression, but not an arbitrary expression. Use a temp var?", arrayIndexedExpression.indexer.position)
-            // TODO we can probably deal with simple binary expressions such as  array[i+1] or lsb(w)/msb(w)  or array[i*2]  to do this automatically...
-        }
 
         super.visit(arrayIndexedExpression)
     }

compiler/src/prog8/ast/processing/StatementReorderer.kt

@@ -81,6 +81,34 @@ internal class StatementReorderer(val program: Program, val errors: ErrorReporte
             arrayIndexedExpression.indexer.indexVar = expr
             arrayIndexedExpression.indexer.origExpression = null
         }
+        else if(expr is BinaryExpression) {
+            if((expr.left is NumericLiteralValue || expr.left is IdentifierReference) &&
+                    (expr.right is NumericLiteralValue || expr.right is IdentifierReference)) {
+                val modifications = mutableListOf<IAstModification>()
+                val indexerVarName = "prog8_autovar_index"
+                val block = expr.definingBlock()
+                if(!block.asmGenInfo.usedAutoArrayIndexer) {
+                    // create the indexer var at block level scope
+                    val vardecl = VarDecl(VarDeclType.VAR, DataType.UBYTE, ZeropageWish.PREFER_ZEROPAGE, null, indexerVarName, null, null, false, true, expr.position)
+                    modifications.add(IAstModification.InsertFirst(vardecl, block))
+                    block.asmGenInfo.usedAutoArrayIndexer = true
+                }
+
+                // assign the indexing expression to the helper variable, and replace the indexer with just the variable
+                val target = AssignTarget(IdentifierReference(listOf(indexerVarName), expr.position), null, null, expr.position)
+                val assign = Assignment(target, expr, expr.position)
+                val beforeWhat = arrayIndexedExpression.containingStatement()
+                modifications.add(IAstModification.InsertBefore(beforeWhat, assign, beforeWhat.parent))
+                modifications.add(IAstModification.SetExpression( {
+                    arrayIndexedExpression.indexer.indexVar = it as IdentifierReference
+                    arrayIndexedExpression.indexer.indexNum = null
+                    arrayIndexedExpression.indexer.origExpression = null
+                }, target.identifier!!.copy(), arrayIndexedExpression.indexer))
+
+                return modifications
+            }
+        }
+
         return noModifications
     }
 

compiler/src/prog8/ast/statements/AstStatements.kt

@@ -29,12 +29,6 @@ sealed class Statement : Node {
             scope.add(name)
         return scope.joinToString(".")
     }
-
-    fun definingBlock(): Block {
-        if(this is Block)
-            return this
-        return findParentNode<Block>(this)!!
-    }
 }
 
 

compiler/test/UnitTests.kt

@@ -188,9 +188,9 @@ class TestC64Zeropage {
         val zp1 = C64Zeropage(CompilationOptions(OutputType.RAW, LauncherType.NONE, ZeropageType.BASICSAFE, emptyList(), true, false))
         assertEquals(18, zp1.available())
         val zp2 = C64Zeropage(CompilationOptions(OutputType.RAW, LauncherType.NONE, ZeropageType.FLOATSAFE, emptyList(), false, false))
-        assertEquals(91, zp2.available())
+        assertEquals(89, zp2.available())
         val zp3 = C64Zeropage(CompilationOptions(OutputType.RAW, LauncherType.NONE, ZeropageType.KERNALSAFE, emptyList(), false, false))
-        assertEquals(127, zp3.available())
+        assertEquals(125, zp3.available())
         val zp4 = C64Zeropage(CompilationOptions(OutputType.RAW, LauncherType.NONE, ZeropageType.FULL, emptyList(), false, false))
         assertEquals(238, zp4.available())
     }

docs/source/programming.rst

@@ -257,7 +257,9 @@ Array types are also supported. They can be made of bytes, words or floats, stri
 .. note::
     To avoid slow and complex assembly code generation, Prog8 currently enforces some limits on
     what you can index the array with.  *It is not possible to use an arbitrary expression/calculation as an index value*.
-    You can use a numerical constant value or a single variable as an index value.
+    You can use a numerical constant value or a single variable as an index value, and simple expressions such as
+    "i+1"  or "i*2".   If you need more complex indexing expressions, the compiler refuses the statement and
+    will suggest to use a temporary indexer variable instead.
 
 
 You can split an array initializer list over several lines if you want.

examples/sprites.p8

@@ -41,9 +41,9 @@ main {
         ubyte @zp i
         for i in 0 to 7 {
             c64.SPRPTR[i] = $0a00 / 64
-            ubyte twoi = i*2        ; TODO is index for array
+            ubyte twoi = i*2
             c64.SPXY[twoi] = 50+25*i
-            twoi++                  ; TODO is index for array
+            twoi++
             c64.SPXY[twoi] = rnd()
         }
 
@@ -61,8 +61,7 @@ irq {
         ; float up & wobble horizontally
         ubyte @zp i
         for i in 0 to 14 step 2 {
-            ubyte ipp=i+1           ; TODO is index for array
+            c64.SPXY[i+1]--
-            c64.SPXY[ipp]--
             ubyte @zp r = rnd()
             if r>200
                 c64.SPXY[i]++

examples/tehtriz.p8

@@ -218,9 +218,8 @@ waitkey:
                 if linepos and blocklogic.isLineFull(linepos)
                     blocklogic.collapse(linepos)
             lines += num_lines
-            uword[] scores = [10, 25, 50, 100]      ; can never clear more than 4 lines
+            uword[] scores = [10, 25, 50, 100]      ; can never clear more than 4 lines at once
-            ubyte scorei = num_lines-1      ; TODO is index for array
+            score += scores[num_lines-1]
-            score += scores[scorei]
             speedlevel = 1+lsb(lines/10)
             drawScore()
         }

examples/test.p8

@@ -7,6 +7,25 @@ main {
 
     sub start() {
 
+        ubyte[] array = [1,2,3,4]
+
+        const ubyte ic = 2
+        ubyte ib = 2
+
+        ib = array[ic+1]
+        ib = array[ib]
+        ib = array[ib+1]
+        ib = array[ib+2]
+        ib = array[ib-1]
+        ib = array[ib-2]
+        ib = array[ib*2]
+        ib = array[2*ib]
+        ib = array[ib*3]
+        ib = array[3*ib]
+        ib = array[ib*4]
+        ib = array[4*ib]
+;        ib = array[lsb(ib)]
+;        ib = array[msb(ib)]
         testX()
     }
 

examples/textelite.p8

@@ -506,33 +506,29 @@ galaxy {
     sub make_current_planet_name() -> str {
         ubyte ni = 0
         str name = "         "    ; max 8
-        ubyte pn_pair1_p1 = pn_pair1+1      ; TODO is index for array
-        ubyte pn_pair2_p1 = pn_pair2+1      ; TODO is index for array
-        ubyte pn_pair3_p1 = pn_pair3+1      ; TODO is index for array
-        ubyte pn_pair4_p1 = pn_pair4+1      ; TODO is index for array
 
         if pn_pairs[pn_pair1] != '.' {
             name[ni] = pn_pairs[pn_pair1]
             ni++
         }
-        if pn_pairs[pn_pair1_p1] != '.' {
+        if pn_pairs[pn_pair1+1] != '.' {
-            name[ni] = pn_pairs[pn_pair1_p1]
+            name[ni] = pn_pairs[pn_pair1+1]
             ni++
         }
         if pn_pairs[pn_pair2] != '.' {
             name[ni] = pn_pairs[pn_pair2]
             ni++
         }
-        if pn_pairs[pn_pair2_p1] != '.' {
+        if pn_pairs[pn_pair2+1] != '.' {
-            name[ni] = pn_pairs[pn_pair2_p1]
+            name[ni] = pn_pairs[pn_pair2+1]
             ni++
         }
         if pn_pairs[pn_pair3] != '.' {
             name[ni] = pn_pairs[pn_pair3]
             ni++
         }
-        if pn_pairs[pn_pair3_p1] != '.' {
+        if pn_pairs[pn_pair3+1] != '.' {
-            name[ni] = pn_pairs[pn_pair3_p1]
+            name[ni] = pn_pairs[pn_pair3+1]
             ni++
         }
 
@@ -541,8 +537,8 @@ galaxy {
                 name[ni] = pn_pairs[pn_pair4]
                 ni++
             }
-            if pn_pairs[pn_pair4_p1] != '.' {
+            if pn_pairs[pn_pair4+1] != '.' {
-                name[ni] = pn_pairs[pn_pair4_p1]
+                name[ni] = pn_pairs[pn_pair4+1]
                 ni++
             }
         }
@@ -710,7 +706,7 @@ planet {
                 name[nx] = pairs0[x]
                 nx++
             }
-            x++     ; TODO is index for array
+            x++
             if pairs0[x] != '.' {
                 name[nx] = pairs0[x]
                 nx++