From 63dd131d812be4b8d4894fc20ca9968e25996b07 Mon Sep 17 00:00:00 2001
From: Joel Challis <git@zvecr.com>
Date: Thu, 14 Mar 2024 10:45:12 +0000
Subject: [PATCH] Refactor vusb to protocol use pre/post task (#14944)

---
 quantum/main.c                      | 35 +++++++++++++++--------------
 tmk_core/protocol/chibios/chibios.c | 13 -----------
 tmk_core/protocol/lufa/lufa.c       |  6 +----
 tmk_core/protocol/vusb/protocol.c   | 32 ++++++--------------------
 tmk_core/protocol/vusb/vusb.c       | 10 ++++++++-
 5 files changed, 35 insertions(+), 61 deletions(-)

diff --git a/quantum/main.c b/quantum/main.c
index 3b101c522c..3159c55850 100644
--- a/quantum/main.c
+++ b/quantum/main.c
@@ -25,22 +25,9 @@ void protocol_pre_task(void);
 void protocol_post_task(void);
 
 // Bodge as refactoring this area sucks....
-void protocol_init(void) __attribute__((weak));
-void protocol_init(void) {
-    protocol_pre_init();
-
-    keyboard_init();
-
-    protocol_post_init();
-}
-
-void protocol_task(void) __attribute__((weak));
-void protocol_task(void) {
-    protocol_pre_task();
-
+void protocol_keyboard_task(void) __attribute__((weak));
+void protocol_keyboard_task(void) {
     keyboard_task();
-
-    protocol_post_task();
 }
 
 /** \brief Main
@@ -53,11 +40,25 @@ int main(void) {
     protocol_setup();
     keyboard_setup();
 
-    protocol_init();
+    protocol_pre_init();
+    keyboard_init();
+    protocol_post_init();
 
     /* Main loop */
     while (true) {
-        protocol_task();
+        protocol_pre_task();
+        protocol_keyboard_task();
+        protocol_post_task();
+
+#ifdef RAW_ENABLE
+        void raw_hid_task(void);
+        raw_hid_task();
+#endif
+
+#ifdef CONSOLE_ENABLE
+        void console_task(void);
+        console_task();
+#endif
 
 #ifdef QUANTUM_PAINTER_ENABLE
         // Run Quantum Painter task
diff --git a/tmk_core/protocol/chibios/chibios.c b/tmk_core/protocol/chibios/chibios.c
index 76a37ae538..360e6b4b04 100644
--- a/tmk_core/protocol/chibios/chibios.c
+++ b/tmk_core/protocol/chibios/chibios.c
@@ -70,13 +70,6 @@ host_driver_t chibios_driver = {keyboard_leds, send_keyboard, send_nkro, send_mo
 void virtser_task(void);
 #endif
 
-#ifdef RAW_ENABLE
-void raw_hid_task(void);
-#endif
-
-#ifdef CONSOLE_ENABLE
-void console_task(void);
-#endif
 #ifdef MIDI_ENABLE
 void midi_ep_task(void);
 #endif
@@ -209,17 +202,11 @@ void protocol_pre_task(void) {
 }
 
 void protocol_post_task(void) {
-#ifdef CONSOLE_ENABLE
-    console_task();
-#endif
 #ifdef MIDI_ENABLE
     midi_ep_task();
 #endif
 #ifdef VIRTSER_ENABLE
     virtser_task();
-#endif
-#ifdef RAW_ENABLE
-    raw_hid_task();
 #endif
     usb_idle_task();
 }
diff --git a/tmk_core/protocol/lufa/lufa.c b/tmk_core/protocol/lufa/lufa.c
index 22cc0db8ce..d6f0c69b6b 100644
--- a/tmk_core/protocol/lufa/lufa.c
+++ b/tmk_core/protocol/lufa/lufa.c
@@ -151,7 +151,7 @@ __attribute__((weak)) void raw_hid_receive(uint8_t *data, uint8_t length) {
  *
  * FIXME: Needs doc
  */
-static void raw_hid_task(void) {
+void raw_hid_task(void) {
     // Create a temporary buffer to hold the read in data from the host
     uint8_t data[RAW_EPSIZE];
     bool    data_read = false;
@@ -865,10 +865,6 @@ void protocol_post_task(void) {
     CDC_Device_USBTask(&cdc_device);
 #endif
 
-#ifdef RAW_ENABLE
-    raw_hid_task();
-#endif
-
 #if !defined(INTERRUPT_CONTROL_ENDPOINT)
     USB_USBTask();
 #endif
diff --git a/tmk_core/protocol/vusb/protocol.c b/tmk_core/protocol/vusb/protocol.c
index 6178d48ef2..41ccf451fd 100644
--- a/tmk_core/protocol/vusb/protocol.c
+++ b/tmk_core/protocol/vusb/protocol.c
@@ -31,14 +31,6 @@
 #    include "sleep_led.h"
 #endif
 
-#ifdef CONSOLE_ENABLE
-void console_task(void);
-#endif
-
-#ifdef RAW_ENABLE
-void raw_hid_task(void);
-#endif
-
 /* This is from main.c of USBaspLoader */
 static void initForUsbConnectivity(void) {
     uint8_t i = 0;
@@ -136,7 +128,7 @@ static inline bool should_do_suspend(void) {
     return vusb_suspended;
 }
 
-void protocol_task(void) {
+void protocol_pre_task(void) {
 #if !defined(NO_USB_STARTUP_CHECK)
     if (should_do_suspend()) {
         dprintln("suspending keyboard");
@@ -159,7 +151,9 @@ void protocol_task(void) {
         vusb_wakeup();
     }
 #endif
+}
 
+void protocol_keyboard_task(void) {
     usbPoll();
 
     // TODO: configuration process is inconsistent. it sometime fails.
@@ -167,20 +161,8 @@ void protocol_task(void) {
     if (usbConfiguration && usbInterruptIsReady()) {
         keyboard_task();
     }
-
-#ifdef RAW_ENABLE
-    usbPoll();
-
-    if (usbConfiguration && usbInterruptIsReady4()) {
-        raw_hid_task();
-    }
-#endif
-
-#ifdef CONSOLE_ENABLE
-    usbPoll();
-
-    if (usbConfiguration && usbInterruptIsReady3()) {
-        console_task();
-    }
-#endif
+}
+
+void protocol_post_task(void) {
+    // do nothing
 }
diff --git a/tmk_core/protocol/vusb/vusb.c b/tmk_core/protocol/vusb/vusb.c
index cfeeed3712..c8ab494253 100644
--- a/tmk_core/protocol/vusb/vusb.c
+++ b/tmk_core/protocol/vusb/vusb.c
@@ -162,6 +162,12 @@ __attribute__((weak)) void raw_hid_receive(uint8_t *data, uint8_t length) {
 }
 
 void raw_hid_task(void) {
+    usbPoll();
+
+    if (!usbConfiguration || !usbInterruptIsReady4()) {
+        return;
+    }
+
     if (raw_output_received_bytes == RAW_BUFFER_SIZE) {
         raw_hid_receive(raw_output_buffer, RAW_BUFFER_SIZE);
         raw_output_received_bytes = 0;
@@ -182,7 +188,9 @@ int8_t sendchar(uint8_t c) {
 }
 
 void console_task(void) {
-    if (!usbConfiguration) {
+    usbPoll();
+
+    if (!usbConfiguration || !usbInterruptIsReady3()) {
         return;
     }