From fdf8e9f78c63c5c921ecd18df86a5cce02ccaa76 Mon Sep 17 00:00:00 2001
From: tonasz <27835465+Tonasz@users.noreply.github.com>
Date: Sun, 16 Jan 2022 19:36:24 +0100
Subject: [PATCH] Add RP2040 PIO UART implementation for split keyboards

---
 docs/split_keyboards.md  |  1 +
 kmk/handlers/pio_uart.py | 90 ++++++++++++++++++++++++++++++++++++++++
 kmk/modules/split.py     | 25 ++++++++---
 3 files changed, 110 insertions(+), 6 deletions(-)
 create mode 100644 kmk/handlers/pio_uart.py

diff --git a/docs/split_keyboards.md b/docs/split_keyboards.md
index 5d3f969..17016cc 100644
--- a/docs/split_keyboards.md
+++ b/docs/split_keyboards.md
@@ -41,6 +41,7 @@ split = Split(
     data_pin2=None,  # Second uart pin to allow 2 way communication
     target_left=True,  # Assumes that left will be the one on USB. Set to folse if it will be the right
     uart_flip=True,  # Reverses the RX and TX pins if both are provided
+    use_pio=False,  # Use RP2040 PIO implementation of UART. Required if you want to use other pins than RX/TX
 )
 
 ```
diff --git a/kmk/handlers/pio_uart.py b/kmk/handlers/pio_uart.py
new file mode 100644
index 0000000..b912813
--- /dev/null
+++ b/kmk/handlers/pio_uart.py
@@ -0,0 +1,90 @@
+# Original source of these examples: https://github.com/adafruit/Adafruit_CircuitPython_PIOASM/tree/main/examples
+# SPDX-FileCopyrightText: 2021 Jeff Epler, written for Adafruit Industries
+# SPDX-License-Identifier: MIT
+
+import rp2pio
+import adafruit_pioasm
+
+tx_code = adafruit_pioasm.assemble(
+    """
+.program uart_tx
+.side_set 1 opt
+; An 8n1 UART transmit program.
+; OUT pin 0 and side-set pin 0 are both mapped to UART TX pin.
+  pull side 1 [7] ; Assert stop bit, or stall with line in idle state
+  set x, 7 side 0 [7] ; Preload bit counter, assert start bit for 8 clocks
+bitloop: ; This loop will run 8 times (8n1 UART)
+  out pins, 1 ; Shift 1 bit from OSR to the first OUT pin
+  jmp x-- bitloop [6] ; Each loop iteration is 8 cycles.
+"""
+)
+
+rx_code = adafruit_pioasm.assemble(
+    """
+.program uart_rx_mini
+
+; Minimum viable 8n1 UART receiver. Wait for the start bit, then sample 8 bits
+; with the correct timing.
+; IN pin 0 is mapped to the GPIO used as UART RX.
+; Autopush must be enabled, with a threshold of 8.
+
+    wait 0 pin 0        ; Wait for start bit
+    set x, 7 [10]       ; Preload bit counter, delay until eye of first data bit
+bitloop:                ; Loop 8 times
+    in pins, 1          ; Sample data
+    jmp x-- bitloop [6] ; Each iteration is 8 cycles
+
+"""
+)
+
+
+class PIO_UART:
+    def __init__(self, *, tx, rx, baudrate=9600):
+        if tx:
+            self.tx_pio = rp2pio.StateMachine(
+                tx_code,
+                first_out_pin=tx,
+                first_sideset_pin=tx,
+                frequency=8 * baudrate,
+                initial_sideset_pin_state=1,
+                initial_sideset_pin_direction=1,
+                initial_out_pin_state=1,
+                initial_out_pin_direction=1,
+                sideset_enable=True,
+            )
+        if rx:
+            self.rx_pio = rp2pio.StateMachine(
+                rx_code,
+                first_in_pin=rx,
+                frequency=8 * baudrate,
+                auto_push=True,
+                push_threshold=8,
+            )
+
+    @property
+    def timeout(self):
+        return 0
+
+    @property
+    def baudrate(self):
+        return self.tx_pio.frequency // 8
+
+    @baudrate.setter
+    def baudrate(self, frequency):
+        self.tx_pio.frequency = frequency * 8
+        self.rx_pio.frequency = frequency * 8
+
+    def write(self, buf):
+        return self.tx_pio.write(buf)
+
+    @property
+    def in_waiting(self):
+        return self.rx_pio.in_waiting
+
+    def read(self, n):
+        b = bytearray(n)
+        n = self.rx_pio.readinto(b)
+        return b[:n]
+
+    def readinto(self, buf):  # pylint: disable=unused-argument
+        return self.rx_pio.readinto(buf)  # pylint: disable=undefined-variable
diff --git a/kmk/modules/split.py b/kmk/modules/split.py
index c273395..5409397 100644
--- a/kmk/modules/split.py
+++ b/kmk/modules/split.py
@@ -36,6 +36,7 @@ class Split(Module):
         data_pin2=None,
         target_left=True,
         uart_flip=True,
+        use_pio=False,
         debug_enabled=False,
     ):
         self._is_target = True
@@ -49,6 +50,7 @@ class Split(Module):
         self.data_pin2 = data_pin2
         self.target_left = target_left
         self.uart_flip = uart_flip
+        self._use_pio = use_pio
         self._uart = None
         self._uart_interval = uart_interval
         self._debug_enabled = debug_enabled
@@ -75,6 +77,11 @@ class Split(Module):
             self._advertising = False
             self._psave_enable = False
 
+        if self._use_pio:
+            from kmk.handlers.pio_uart import PIO_UART
+
+            self.PIO_UART = PIO_UART
+
     def during_bootup(self, keyboard):
         # Set up name for target side detection and BLE advertisment
         name = str(getmount('/').label)
@@ -116,13 +123,19 @@ class Split(Module):
 
         if self.split_type == SplitType.UART and self.data_pin is not None:
             if self._is_target or not self.uart_flip:
-                self._uart = busio.UART(
-                    tx=self.data_pin2, rx=self.data_pin, timeout=self._uart_interval
-                )
+                if self._use_pio:
+                    self._uart = self.PIO_UART(tx=self.data_pin2, rx=self.data_pin)
+                else:
+                    self._uart = busio.UART(
+                        tx=self.data_pin2, rx=self.data_pin, timeout=self._uart_interval
+                    )
             else:
-                self._uart = busio.UART(
-                    tx=self.data_pin, rx=self.data_pin2, timeout=self._uart_interval
-                )
+                if self._use_pio:
+                    self._uart = self.PIO_UART(tx=self.data_pin, rx=self.data_pin2)
+                else:
+                    self._uart = busio.UART(
+                        tx=self.data_pin, rx=self.data_pin2, timeout=self._uart_interval
+                    )
 
         # Attempt to sanely guess a coord_mapping if one is not provided.
         if not keyboard.coord_mapping: