From 270099c58b11a7857de5abb456c33dd9ce94af3a Mon Sep 17 00:00:00 2001
From: "Aode (lion)" <asonix@asonix.dog>
Date: Mon, 7 Mar 2022 19:56:29 -0600
Subject: [PATCH] Select

---
 .gitignore |  2 ++
 Cargo.toml |  8 +++++
 src/lib.rs | 89 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+)
 create mode 100644 .gitignore
 create mode 100644 Cargo.toml
 create mode 100644 src/lib.rs

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..96ef6c0
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+/target
+Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
new file mode 100644
index 0000000..e651f6c
--- /dev/null
+++ b/Cargo.toml
@@ -0,0 +1,8 @@
+[package]
+name = "select"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..98d7a68
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,89 @@
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+    task::{Context, Poll, Wake, Waker},
+};
+
+pub enum Either<L, R> {
+    Left(L),
+    Right(R),
+}
+
+pub fn join<L, R>(left: L, right: R) -> impl Future<Output = Either<L::Output, R::Output>> + Unpin
+where
+    L: Future + Unpin,
+    R: Future + Unpin,
+{
+    Select {
+        left: SelectState {
+            woken: Arc::new(AtomicBool::new(true)),
+            future: left,
+        },
+        right: SelectState {
+            woken: Arc::new(AtomicBool::new(true)),
+            future: right,
+        },
+    }
+}
+
+struct SelectWaker {
+    woken: Arc<AtomicBool>,
+    parent: Waker,
+}
+
+struct SelectState<T> {
+    woken: Arc<AtomicBool>,
+    future: T,
+}
+
+struct Select<L, R> {
+    left: SelectState<L>,
+    right: SelectState<R>,
+}
+
+impl<T: Future> SelectState<T> {
+    fn check_for_poll(&self) -> bool {
+        self.woken.swap(false, Ordering::AcqRel)
+    }
+}
+
+impl Wake for SelectWaker {
+    fn wake(self: Arc<Self>) {
+        self.wake_by_ref();
+    }
+
+    fn wake_by_ref(self: &Arc<Self>) {
+        self.woken.store(true, Ordering::Release);
+        self.parent.wake_by_ref();
+    }
+}
+
+impl<L, R> Future for Select<L, R>
+where
+    L: Future + Unpin,
+    R: Future + Unpin,
+{
+    type Output = Either<L::Output, R::Output>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.get_mut();
+
+        if this.left.check_for_poll() {
+            if let Poll::Ready(out) = Pin::new(&mut this.left.future).poll(cx) {
+                return Poll::Ready(Either::Left(out));
+            }
+        }
+
+        if this.right.check_for_poll() {
+            if let Poll::Ready(out) = Pin::new(&mut this.right.future).poll(cx) {
+                return Poll::Ready(Either::Right(out));
+            }
+        }
+
+        Poll::Pending
+    }
+}