From c968bacb01f1ce7dee24f21505f1dbfc4b3edc69 Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Wed, 8 Mar 2023 23:29:57 +0100
Subject: [PATCH] fix pokew() crash with certain address expressions

---
 .../codegen/cpu6502/BuiltinFunctionsAsmGen.kt  | 11 ++++++-----
 .../test/codegeneration/TestVariousCodeGen.kt  | 18 ++++++++++++++++++
 docs/source/todo.rst                           |  2 --
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/codeGenCpu6502/src/prog8/codegen/cpu6502/BuiltinFunctionsAsmGen.kt b/codeGenCpu6502/src/prog8/codegen/cpu6502/BuiltinFunctionsAsmGen.kt
index 903cbbc94..0e04387fd 100644
--- a/codeGenCpu6502/src/prog8/codegen/cpu6502/BuiltinFunctionsAsmGen.kt
+++ b/codeGenCpu6502/src/prog8/codegen/cpu6502/BuiltinFunctionsAsmGen.kt
@@ -680,12 +680,13 @@ internal class BuiltinFunctionsAsmGen(private val program: PtProgram,
                     }
                 }
             }
-            else -> throw AssemblyError("wrong pokew arg type")
+            else -> {
+                asmgen.assignExpressionToVariable(fcall.args[0], "P8ZP_SCRATCH_W1", DataType.UWORD, null)
+                asmgen.assignExpressionToRegister(fcall.args[1], RegisterOrPair.AY)
+                asmgen.out("  jsr  prog8_lib.func_pokew")
+                return
+            }
         }
-
-        asmgen.assignExpressionToVariable(fcall.args[0], "P8ZP_SCRATCH_W1", DataType.UWORD, null)
-        asmgen.assignExpressionToRegister(fcall.args[1], RegisterOrPair.AY)
-        asmgen.out("  jsr  prog8_lib.func_pokew")
     }
 
     private fun funcPeekW(fcall: PtBuiltinFunctionCall, resultToStack: Boolean, resultRegister: RegisterOrPair?) {
diff --git a/compiler/test/codegeneration/TestVariousCodeGen.kt b/compiler/test/codegeneration/TestVariousCodeGen.kt
index 5d2c99ff7..80ef7752a 100644
--- a/compiler/test/codegeneration/TestVariousCodeGen.kt
+++ b/compiler/test/codegeneration/TestVariousCodeGen.kt
@@ -91,4 +91,22 @@ main {
         assign.target.identifier!!.name shouldBe "cx16.r0"
         assign.value shouldBe instanceOf<PtArrayIndexer>()
     }
+
+    test("peek and poke argument types") {
+        val text="""
+main {
+    sub start() {
+        uword[3] arr
+        ubyte i = 42
+        uword ww = peekw(arr[i])
+        ubyte xx = peek(arr[i])
+        xx = @(arr[i])
+        
+        @(arr[i]) = 42
+        poke(arr[i], 42)
+        pokew(arr[i], 4242)
+    }
+}"""
+        compileText(C64Target(), false, text, writeAssembly = true) shouldNotBe null
+    }
 })
\ No newline at end of file
diff --git a/docs/source/todo.rst b/docs/source/todo.rst
index e5f30fe17..2a936a30a 100644
--- a/docs/source/todo.rst
+++ b/docs/source/todo.rst
@@ -3,8 +3,6 @@ TODO
 
 For next minor release
 ^^^^^^^^^^^^^^^^^^^^^^
-- check that peekw() and peek() and @(...) accept a[i] where a=uword
-- fix pokew() crash with args like: a[i] (where a=uword) ,  z % twoIMinusOne
 - fix Github issue with X register https://github.com/irmen/prog8/issues/94
 - fix Github issue with array problems https://github.com/irmen/prog8/issues/99
 - fix IR/VM: animals.p8 example somehow doesn't print the animal name correctly in the first question, and exits after 1 animal instead of looping