2 files changed, 36 insertions, 3 deletions

diff --git a/platforms/suspend.c b/platforms/suspend.c
index fea23cbd02..4756796ea4 100644
--- a/platforms/suspend.c
+++ b/platforms/suspend.c
@@ -4,6 +4,9 @@
 #include "suspend.h"
 #include "matrix.h"
 
+extern matrix_row_t matrix_previous[MATRIX_ROWS];
+static matrix_row_t wakeup_matrix[MATRIX_ROWS];
+
 // TODO: Move to more correct location
 __attribute__((weak)) void matrix_power_up(void) {}
 __attribute__((weak)) void matrix_power_down(void) {}
@@ -44,8 +47,34 @@ bool suspend_wakeup_condition(void) {
     matrix_power_up();
     matrix_scan();
     matrix_power_down();
-    for (uint8_t r = 0; r < MATRIX_ROWS; r++) {
-        if (matrix_get_row(r)) return true;
+
+    bool wakeup = false;
+    for (uint8_t row = 0; row < MATRIX_ROWS; row++) {
+        wakeup_matrix[row] = matrix_get_row(row);
+        wakeup |= wakeup_matrix[row] != 0;
+    }
+
+    return wakeup;
+}
+
+void update_matrix_state_after_wakeup(void) {
+    matrix_power_up();
+    matrix_scan();
+    matrix_power_down();
+
+    for (uint8_t row = 0; row < MATRIX_ROWS; row++) {
+        const matrix_row_t current_row = matrix_get_row(row);
+        wakeup_matrix[row] |= current_row & ~matrix_previous[row];
+        matrix_previous[row] |= current_row;
+    }
+}
+
+bool keypress_is_wakeup_key(uint8_t row, uint8_t col) {
+    return (wakeup_matrix[row] & ((matrix_row_t)1 << col));
+}
+
+void wakeup_matrix_handle_key_event(uint8_t row, uint8_t col, bool pressed) {
+    if (!pressed) {
+        wakeup_matrix[row] &= ~((matrix_row_t)1 << col);
     }
-    return false;
 }
diff --git a/platforms/suspend.h b/platforms/suspend.h
index e4f7f39ddb..3a7f51f9f0 100644
--- a/platforms/suspend.h
+++ b/platforms/suspend.h
@@ -14,6 +14,10 @@ void suspend_power_down_user(void);
 void suspend_power_down_kb(void);
 void suspend_power_down_quantum(void);
 
+bool keypress_is_wakeup_key(uint8_t row, uint8_t col);
+void update_matrix_state_after_wakeup(void);
+void wakeup_matrix_handle_key_event(uint8_t row, uint8_t col, bool pressed);
+
 #ifndef USB_SUSPEND_WAKEUP_DELAY
 #    define USB_SUSPEND_WAKEUP_DELAY 0
 #endif