From 86f02e8b8389469b6ca3f6b375d1c1e0605250ee Mon Sep 17 00:00:00 2001
From: Jeff Epler <jepler@gmail.com>
Date: Sun, 8 Sep 2024 08:01:09 -0500
Subject: [PATCH] Add wait jmppin & wait irq prev/next

---
 adafruit_pioasm.py    | 34 +++++++++++++++++++++++++++-------
 tests/test_version.py | 13 +++++++++++++
 2 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/adafruit_pioasm.py b/adafruit_pioasm.py
index 8662e04..1ddcb1b 100644
--- a/adafruit_pioasm.py
+++ b/adafruit_pioasm.py
@@ -278,16 +278,36 @@ class Program:  # pylint: disable=too-few-public-methods
                 #                instr delay p sr index
                 assembled.append(0b001_00000_0_00_00000)
                 polarity = int(instruction[1], 0)
+                source = instruction[2]
                 if not 0 <= polarity <= 1:
                     raise RuntimeError("Invalid polarity")
                 assembled[-1] |= polarity << 7
-                assembled[-1] |= WAIT_SOURCES.index(instruction[2]) << 5
-                num = int(instruction[3], 0)
-                if not 0 <= num <= 31:
-                    raise RuntimeError("Wait num out of range")
-                assembled[-1] |= num
-                if instruction[-1] == "rel":
-                    assembled[-1] |= 0x10  # Set the high bit of the irq value
+                if instruction[2] == "jmppin":
+                    require_version(1, "wait jmppin")
+                    num = 0
+                    print("wait jmppin", instruction)
+                    if len(instruction) > 3:
+                        if len(instruction) < 5 or instruction[3] != "+":
+                            raise RuntimeError("invalid wait jmppin")
+                        num = int_in_range(instruction[4], 0, 4, "wait jmppin offset")
+                    assembled[-1] |= num
+                    assembled[-1] |= 0b11 << 5  # JMPPIN wait source
+                else:
+                    assembled[-1] |= WAIT_SOURCES.index(instruction[2]) << 5
+                    num = int(instruction[3], 0)
+                    if not 0 <= num <= 31:
+                        raise RuntimeError("Wait num out of range")
+                    assembled[-1] |= num
+                    # The flag index is decoded in the same way as the IRQ
+                    # index field, decoding down from the two MSBs
+                    if instruction[-1] == "next":
+                        require_version(1, "wait irq next")
+                        assembled[-1] |= 0b11000
+                    elif instruction[-1] == "prev":
+                        require_version(1, "wait irq prev")
+                        assembled[-1] |= 0b01000
+                    elif instruction[-1] == "rel":
+                        assembled[-1] |= 0b10000
             elif instruction[0] == "in":
                 #                instr delay src count
                 assembled.append(0b010_00000_000_00000)
diff --git a/tests/test_version.py b/tests/test_version.py
index 2999caf..380d5aa 100644
--- a/tests/test_version.py
+++ b/tests/test_version.py
@@ -125,3 +125,16 @@ def test_mov_v1() -> None:
 
     assert_assembly_fails("mov pindirs, null", errtype=ValueError)
     assert_assembles_to(prefix + "mov pindirs, null", [0b101_00000_01100011])
+
+
+def test_wait_v1() -> None:
+    assert_assembly_fails("wait 0 jmppin")
+    assert_assembly_fails("wait 0 irq 5 next")
+    prefix = ".pio_version 1\n"
+    assert_assembly_fails(prefix + "wait 0 jmppin +")
+    assert_assembly_fails(prefix + "wait 0 jmppin + 7")
+    assert_assembles_to(prefix + "wait 0 jmppin + 3", [0b001_00000_0_11_00011])
+    assert_assembles_to(prefix + "wait 1 jmppin", [0b001_00000_1_11_00000])
+
+    assert_assembles_to(prefix + "wait 0 irq 5 next", [0b001_00000_0_10_11_101])
+    assert_assembles_to(prefix + "wait 1 irq 4 prev", [0b001_00000_1_10_01_100])