From aa091ef5820cd3b8c688b4fe624c7512776c1041 Mon Sep 17 00:00:00 2001
From: Putta Khunchalee <ultimaweapon@outlook.com>
Date: Sun, 15 Dec 2024 01:04:50 +0700
Subject: [PATCH] Implements custom async executor for winit event loop

---
 gui/src/rt/mod.rs   | 80 ++++++++++++++++++++++++++++++++++-----------
 gui/src/rt/task.rs  | 34 +++++++++++++++++++
 gui/src/rt/waker.rs | 26 +++++++++++++++
 3 files changed, 121 insertions(+), 19 deletions(-)
 create mode 100644 gui/src/rt/task.rs
 create mode 100644 gui/src/rt/waker.rs

diff --git a/gui/src/rt/mod.rs b/gui/src/rt/mod.rs
index 95251877a..2c7d771a2 100644
--- a/gui/src/rt/mod.rs
+++ b/gui/src/rt/mod.rs
@@ -1,35 +1,38 @@
 pub use self::context::*;
 
 use self::event::WindowEvent;
-use futures::executor::LocalPool;
-use futures::task::LocalSpawnExt;
+use self::task::TaskList;
+use self::waker::Waker;
 use std::future::Future;
+use std::sync::Arc;
 use thiserror::Error;
 use winit::application::ApplicationHandler;
 use winit::error::EventLoopError;
-use winit::event_loop::{ActiveEventLoop, EventLoop};
+use winit::event_loop::{ActiveEventLoop, EventLoop, EventLoopProxy};
 use winit::window::WindowId;
 
 mod context;
 mod event;
+mod task;
+mod waker;
 
 pub fn block_on(main: impl Future<Output = ()> + 'static) -> Result<(), RuntimeError> {
     // Setup winit event loop.
     let mut el = EventLoop::<Event>::with_user_event();
     let el = el.build().map_err(RuntimeError::CreateEventLoop)?;
-    let executor = LocalPool::new();
-
-    executor
-        .spawner()
-        .spawn_local(async move {
-            main.await;
-            RuntimeContext::with(|cx| cx.event_loop().exit());
-        })
-        .unwrap();
+    let main = async move {
+        main.await;
+        RuntimeContext::with(|cx| cx.event_loop().exit());
+    };
 
     // Run event loop.
+    let mut tasks = TaskList::default();
+    let main: Box<dyn Future<Output = ()>> = Box::new(main);
+    let main = tasks.insert(Box::into_pin(main));
     let mut rt = Runtime {
-        executor,
+        el: el.create_proxy(),
+        tasks,
+        main,
         on_close: WindowEvent::default(),
     };
 
@@ -38,18 +41,55 @@ pub fn block_on(main: impl Future<Output = ()> + 'static) -> Result<(), RuntimeE
 
 /// Implementation of [`ApplicationHandler`] to drive [`Future`].
 struct Runtime {
-    executor: LocalPool,
+    el: EventLoopProxy<Event>,
+    tasks: TaskList,
+    main: u64,
     on_close: WindowEvent<()>,
 }
 
-impl ApplicationHandler<Event> for Runtime {
-    fn resumed(&mut self, event_loop: &ActiveEventLoop) {
+impl Runtime {
+    pub fn dispatch(&mut self, el: &ActiveEventLoop, task: u64) -> bool {
+        // Get target task.
+        let mut task = match self.tasks.get(task) {
+            Some(v) => v,
+            None => {
+                // It is possible for the waker to wake the same task multiple times. In this case
+                // the previous wake may complete the task.
+                return false;
+            }
+        };
+
+        // Poll the task.
+        let waker = Arc::new(Waker::new(self.el.clone(), *task.key()));
         let mut cx = RuntimeContext {
-            el: event_loop,
+            el,
             on_close: &mut self.on_close,
         };
 
-        cx.run(|| self.executor.run_until_stalled());
+        cx.run(|| {
+            let waker = std::task::Waker::from(waker);
+            let mut cx = std::task::Context::from_waker(&waker);
+
+            if task.get_mut().as_mut().poll(&mut cx).is_ready() {
+                drop(task.remove());
+            }
+        });
+
+        true
+    }
+}
+
+impl ApplicationHandler<Event> for Runtime {
+    fn resumed(&mut self, event_loop: &ActiveEventLoop) {
+        assert!(self.dispatch(event_loop, self.main));
+    }
+
+    fn user_event(&mut self, event_loop: &ActiveEventLoop, event: Event) {
+        match event {
+            Event::TaskReady(task) => {
+                self.dispatch(event_loop, task);
+            }
+        }
     }
 
     fn window_event(
@@ -69,7 +109,9 @@ impl ApplicationHandler<Event> for Runtime {
 }
 
 /// Event to wakeup winit event loop.
-enum Event {}
+enum Event {
+    TaskReady(u64),
+}
 
 /// Represents an error when [`block_on()`] fails.
 #[derive(Debug, Error)]
diff --git a/gui/src/rt/task.rs b/gui/src/rt/task.rs
new file mode 100644
index 000000000..65565ffaa
--- /dev/null
+++ b/gui/src/rt/task.rs
@@ -0,0 +1,34 @@
+use std::collections::hash_map::OccupiedEntry;
+use std::collections::HashMap;
+use std::future::Future;
+use std::pin::Pin;
+
+/// List of pending tasks.
+#[derive(Default)]
+pub struct TaskList {
+    list: HashMap<u64, Pin<Box<dyn Future<Output = ()>>>>,
+    next: u64,
+}
+
+impl TaskList {
+    pub fn insert(&mut self, f: Pin<Box<dyn Future<Output = ()>>>) -> u64 {
+        let id = self.next;
+
+        assert!(self.list.insert(id, f).is_none());
+        self.next = self.next.checked_add(1).unwrap();
+
+        id
+    }
+
+    pub fn get(
+        &mut self,
+        id: u64,
+    ) -> Option<OccupiedEntry<u64, Pin<Box<dyn Future<Output = ()>>>>> {
+        use std::collections::hash_map::Entry;
+
+        match self.list.entry(id) {
+            Entry::Occupied(e) => Some(e),
+            Entry::Vacant(_) => None,
+        }
+    }
+}
diff --git a/gui/src/rt/waker.rs b/gui/src/rt/waker.rs
new file mode 100644
index 000000000..ca96bd8b3
--- /dev/null
+++ b/gui/src/rt/waker.rs
@@ -0,0 +1,26 @@
+use super::Event;
+use std::sync::Arc;
+use std::task::Wake;
+use winit::event_loop::EventLoopProxy;
+
+/// Implementation of [`Wake`].
+pub struct Waker {
+    el: EventLoopProxy<Event>,
+    task: u64,
+}
+
+impl Waker {
+    pub fn new(el: EventLoopProxy<Event>, task: u64) -> Self {
+        Self { el, task }
+    }
+}
+
+impl Wake for Waker {
+    fn wake(self: Arc<Self>) {
+        self.wake_by_ref();
+    }
+
+    fn wake_by_ref(self: &Arc<Self>) {
+        drop(self.el.send_event(Event::TaskReady(self.task)));
+    }
+}