From d63ddde7566262514659c00fe13473188c113d63 Mon Sep 17 00:00:00 2001
From: Philipp Oppermann <dev@phil-opp.com>
Date: Fri, 27 Mar 2020 15:54:04 +0100
Subject: [PATCH] Add waker support to ScancodeStream

---
 src/task/keyboard.rs | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/task/keyboard.rs b/src/task/keyboard.rs
index 946acf35..c4af6ccd 100644
--- a/src/task/keyboard.rs
+++ b/src/task/keyboard.rs
@@ -5,9 +5,10 @@ use core::{
     task::{Context, Poll},
 };
 use crossbeam_queue::ArrayQueue;
-use futures_util::stream::Stream;
+use futures_util::{stream::Stream, task::AtomicWaker};
 
 static SCANCODE_QUEUE: OnceCell<ArrayQueue<u8>> = OnceCell::uninit();
+static WAKER: AtomicWaker = AtomicWaker::new();
 
 /// Called by the keyboard interrupt handler
 ///
@@ -16,6 +17,8 @@ pub(crate) fn add_scancode(scancode: u8) {
     if let Ok(queue) = SCANCODE_QUEUE.try_get() {
         if let Err(_) = queue.push(scancode) {
             println!("WARNING: scancode queue full; dropping keyboard input");
+        } else {
+            WAKER.wake();
         }
     } else {
         println!("WARNING: scancode queue uninitialized");
@@ -39,9 +42,21 @@ impl Stream for ScancodeStream {
     type Item = u8;
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<u8>> {
-        let queue = SCANCODE_QUEUE.try_get().expect("not initialized");
+        let queue = SCANCODE_QUEUE
+            .try_get()
+            .expect("scancode queue not initialized");
+
+        // fast path
+        if let Ok(scancode) = queue.pop() {
+            return Poll::Ready(Some(scancode));
+        }
+
+        WAKER.register(&cx.waker());
         match queue.pop() {
-            Ok(scancode) => Poll::Ready(Some(scancode)),
+            Ok(scancode) => {
+                WAKER.take();
+                Poll::Ready(Some(scancode))
+            }
             Err(crossbeam_queue::PopError) => Poll::Pending,
         }
     }