1 files changed, 9 insertions, 5 deletions

diff --git a/tmk_core/protocol/ps2_mouse.c b/tmk_core/protocol/ps2_mouse.c
index 5415453a0..525aeb45a 100644
--- a/tmk_core/protocol/ps2_mouse.c
+++ b/tmk_core/protocol/ps2_mouse.c
@@ -16,9 +16,13 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
 
 #include <stdbool.h>
-#include <avr/io.h>
-#include <util/delay.h>
+
+#if defined(__AVR__)
+#    include <avr/io.h>
+#endif
+
 #include "ps2_mouse.h"
+#include "wait.h"
 #include "host.h"
 #include "timer.h"
 #include "print.h"
@@ -42,7 +46,7 @@ static inline void ps2_mouse_scroll_button_task(report_mouse_t *mouse_report);
 void ps2_mouse_init(void) {
     ps2_host_init();
 
-    _delay_ms(PS2_MOUSE_INIT_DELAY);  // wait for powering up
+    wait_ms(PS2_MOUSE_INIT_DELAY);  // wait for powering up
 
     PS2_MOUSE_SEND(PS2_MOUSE_RESET, "ps2_mouse_init: sending reset");
 
@@ -210,7 +214,7 @@ static inline void ps2_mouse_enable_scrolling(void) {
     PS2_MOUSE_SEND(PS2_MOUSE_SET_SAMPLE_RATE, "Set sample rate");
     PS2_MOUSE_SEND(80, "80");
     PS2_MOUSE_SEND(PS2_MOUSE_GET_DEVICE_ID, "Finished enabling scroll wheel");
-    _delay_ms(20);
+    wait_ms(20);
 }
 
 #define PRESS_SCROLL_BUTTONS mouse_report->buttons |= (PS2_MOUSE_SCROLL_BTN_MASK)
@@ -252,7 +256,7 @@ static inline void ps2_mouse_scroll_button_task(report_mouse_t *mouse_report) {
         if (scroll_state == SCROLL_BTN && timer_elapsed(scroll_button_time) < PS2_MOUSE_SCROLL_BTN_SEND) {
             PRESS_SCROLL_BUTTONS;
             host_mouse_send(mouse_report);
-            _delay_ms(100);
+            wait_ms(100);
             RELEASE_SCROLL_BUTTONS;
         }
 #endif