diff --git a/compiler/res/prog8lib/diskio.p8 b/compiler/res/prog8lib/diskio.p8
index a8b7c2ac2..ad70aee24 100644
--- a/compiler/res/prog8lib/diskio.p8
+++ b/compiler/res/prog8lib/diskio.p8
@@ -3,11 +3,8 @@
 
 ; Note: this code is compatible with C64 and CX16.
 
-; TODO directory() BROKEN ON C64, SEEMS TO WORK ON CX16
-
 diskio {
 
-
     sub directory(ubyte drivenumber) -> ubyte {
         ; -- Prints the directory contents of disk drive 8-11 to the screen. Returns success.
 
@@ -34,7 +31,7 @@ diskio {
             ubyte @zp char
             repeat {
                 char = c64.CHRIN()
-                if char==0                  ; TODO doesn't work???
+                if char==0
                     break
                 txt.chrout(char)
             }
diff --git a/compiler/src/prog8/compiler/target/c64/codegen/AsmGen.kt b/compiler/src/prog8/compiler/target/c64/codegen/AsmGen.kt
index d42ea016f..b0f3ea42f 100644
--- a/compiler/src/prog8/compiler/target/c64/codegen/AsmGen.kt
+++ b/compiler/src/prog8/compiler/target/c64/codegen/AsmGen.kt
@@ -548,18 +548,27 @@ internal class AsmGen(private val program: Program,
     private  fun fixNameSymbols(name: String) = name.replace("<", "prog8_").replace(">", "")     // take care of the autogenerated invalid (anon) label names
 
     internal fun saveRegisterLocal(register: CpuRegister, scope: Subroutine) {
-        when (register) {
-            CpuRegister.A -> {
-                out("  sta  _prog8_regsaveA")
-                scope.asmGenInfo.usedRegsaveA = true
+        if (CompilationTarget.instance.machine.cpu == CpuType.CPU65c02) {
+            // just use the cpu's stack for all registers, shorter code
+            when (register) {
+                CpuRegister.A -> out("  pha")
+                CpuRegister.X -> out("  phx")
+                CpuRegister.Y -> out("  phy")
             }
-            CpuRegister.X -> {
-                out("  stx  _prog8_regsaveX")
-                scope.asmGenInfo.usedRegsaveX = true
-            }
-            CpuRegister.Y -> {
-                out("  sty  _prog8_regsaveY")
-                scope.asmGenInfo.usedRegsaveY = true
+        } else {
+            when (register) {
+                CpuRegister.A -> {
+                    // just use the stack, only for A
+                    out("  pha")
+                }
+                CpuRegister.X -> {
+                    out("  stx  _prog8_regsaveX")
+                    scope.asmGenInfo.usedRegsaveX = true
+                }
+                CpuRegister.Y -> {
+                    out("  sty  _prog8_regsaveY")
+                    scope.asmGenInfo.usedRegsaveY = true
+                }
             }
         }
     }
@@ -589,10 +598,20 @@ internal class AsmGen(private val program: Program,
     }
 
     internal fun restoreRegisterLocal(register: CpuRegister) {
-        when (register) {
-            CpuRegister.A -> out("  lda  _prog8_regsaveA")
-            CpuRegister.X -> out("  ldx  _prog8_regsaveX")
-            CpuRegister.Y -> out("  ldy  _prog8_regsaveY")
+        if (CompilationTarget.instance.machine.cpu == CpuType.CPU65c02) {
+            when (register) {
+                // this just used the stack, for all registers. Shorter code.
+                CpuRegister.A -> out("  pla")
+                CpuRegister.X -> out("  plx")
+                CpuRegister.Y -> out("  ply")
+            }
+
+        } else {
+            when (register) {
+                CpuRegister.A -> out("  pla")   // this just used the stack but only for A
+                CpuRegister.X -> out("  ldx  _prog8_regsaveX")
+                CpuRegister.Y -> out("  ldy  _prog8_regsaveY")
+            }
         }
     }
 
diff --git a/examples/cx16/imageviewer/imageviewer.p8 b/examples/cx16/imageviewer/imageviewer.p8
index 51902b0c2..066b48777 100644
--- a/examples/cx16/imageviewer/imageviewer.p8
+++ b/examples/cx16/imageviewer/imageviewer.p8
@@ -8,7 +8,6 @@
 %import bmp_module
 ;; %import ci_module
 
-; TODO WHY is this 400 bytes larger than a few hours ago?
 
 main {
     sub start() {