From 0da9142009f54971d913141581b44d9b117ea29c Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Wed, 13 Dec 2023 00:53:01 +0100
Subject: [PATCH] asm: also work for asmsub that return N or Z flag (Carry
 already worked)

---
 .../cpu6502/assignment/AssignmentAsmGen.kt    | 17 ++++++-
 docs/source/todo.rst                          |  3 +-
 examples/test.p8                              | 48 +++++++++++++------
 3 files changed, 51 insertions(+), 17 deletions(-)

diff --git a/codeGenCpu6502/src/prog8/codegen/cpu6502/assignment/AssignmentAsmGen.kt b/codeGenCpu6502/src/prog8/codegen/cpu6502/assignment/AssignmentAsmGen.kt
index 3d259c405..cfa33368c 100644
--- a/codeGenCpu6502/src/prog8/codegen/cpu6502/assignment/AssignmentAsmGen.kt
+++ b/codeGenCpu6502/src/prog8/codegen/cpu6502/assignment/AssignmentAsmGen.kt
@@ -1870,7 +1870,22 @@ internal class AssignmentAsmGen(private val program: PtProgram,
 +                   lda  #1
 +""")
             }
-            else -> throw AssemblyError("can't use Z or N flags as return 'values'")
+            Statusflag.Pz -> {
+                asmgen.out("""
+                    beq  +
+                    lda  #0
+                    beq  ++
++                   lda  #1
++""")
+            }
+            Statusflag.Pn -> {
+                asmgen.out("""
+                    bmi  +
+                    lda  #0
+                    beq  ++
++                   lda  #1
++""")
+            }
         }
         assignRegisterByte(target, CpuRegister.A, false, true)
     }
diff --git a/docs/source/todo.rst b/docs/source/todo.rst
index b6a57bf39..90524cfcf 100644
--- a/docs/source/todo.rst
+++ b/docs/source/todo.rst
@@ -3,7 +3,6 @@ TODO
 ====
 
 - optimize if-else expressions whose condition returns the boolean status in a status register to use a branch opcode instead of a comparison against 0
-- fix "can't use Z or N flags as return 'values'" in 6502 codegen?
 
 - merge branch optimize-st  for some optimizations regardign SymbolTable use
 
@@ -12,6 +11,8 @@ TODO
 
 - [on branch: shortcircuit] investigate McCarthy evaluation again? this may also reduce code size perhaps for things like if a>4 or a<2 ....
 
+- IR (expericodegen): fix code for calling routines that return a boolean in a status register such as Carry flag, it has to store the flag value somewhere
+
 ...
 
 
diff --git a/examples/test.p8 b/examples/test.p8
index b0baed729..2279f48dd 100644
--- a/examples/test.p8
+++ b/examples/test.p8
@@ -1,24 +1,42 @@
 %import textio
-%import math
+%import string
 %zeropage basicsafe
 
 main {
     sub start() {
-        ubyte x,y,z = math.rnd()
+        bool @shared blerp = test(100)
 
-        txt.print_ub(x)
-        txt.nl()
-        txt.print_ub(y)
-        txt.nl()
-        txt.print_ub(z)
-        txt.nl()
+        if test(100)
+            goto skip
+        txt.print("no0")
 
-        x=y=z=math.rnd()
-        txt.print_ub(x)
-        txt.nl()
-        txt.print_ub(y)
-        txt.nl()
-        txt.print_ub(z)
-        txt.nl()
+skip:
+        if test(100) {
+            txt.print("yes1")
+            goto skip2
+        }
+        txt.print("no1")
+
+skip2:
+        if test(100)
+            txt.print("yes2")
+        else
+            txt.print("no2")
+
+
+        while test(100) {
+            cx16.r0++
+        }
+
+        do {
+            cx16.r0++
+        } until test(100)
+    }
+
+    asmsub test(ubyte value @A) -> bool @Pz {
+        %asm {{
+            lda  #0
+            rts
+        }}
     }
 }