From fa59f74bbfde084f9d5c9d2207a8f78a2d705ed9 Mon Sep 17 00:00:00 2001
From: Andy Leiserson <aleiserson@mozilla.com>
Date: Fri, 8 Mar 2024 14:21:15 -0800
Subject: [PATCH] MaybeFuture

---
 ipa-core/src/helpers/futures.rs | 44 +++++++++++++++++++++++++++++++++
 ipa-core/src/helpers/mod.rs     |  2 ++
 2 files changed, 46 insertions(+)
 create mode 100644 ipa-core/src/helpers/futures.rs

diff --git a/ipa-core/src/helpers/futures.rs b/ipa-core/src/helpers/futures.rs
new file mode 100644
index 000000000..70bcf4ca3
--- /dev/null
+++ b/ipa-core/src/helpers/futures.rs
@@ -0,0 +1,44 @@
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use pin_project::pin_project;
+
+#[pin_project(project = MaybeFutureProj)]
+pub enum MaybeFuture<Fut: Future> {
+    Future(#[pin] Fut),
+    Value(Option<Fut::Output>),
+}
+
+impl<Fut: Future> Future for MaybeFuture<Fut> {
+    type Output = Fut::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match self.project() {
+            MaybeFutureProj::Future(fut) => fut.poll(cx),
+            MaybeFutureProj::Value(val) => Poll::Ready(val.take().unwrap()),
+        }
+    }
+}
+
+impl<Fut: Future> MaybeFuture<Fut> {
+    pub fn future(fut: Fut) -> Self {
+        MaybeFuture::Future(fut)
+    }
+
+    pub fn value(val: Fut::Output) -> Self {
+        MaybeFuture::Value(Some(val))
+    }
+}
+
+impl<Fut: Future<Output = Result<(), E>>, E> MaybeFuture<Fut> {
+    pub fn future_or_ok<F: FnOnce() -> Fut>(condition: bool, f: F) -> Self {
+        if condition {
+            MaybeFuture::Future(f())
+        } else {
+            MaybeFuture::Value(Some(Ok(())))
+        }
+    }
+}
diff --git a/ipa-core/src/helpers/mod.rs b/ipa-core/src/helpers/mod.rs
index b2b15f305..4544c2795 100644
--- a/ipa-core/src/helpers/mod.rs
+++ b/ipa-core/src/helpers/mod.rs
@@ -8,6 +8,7 @@ use generic_array::GenericArray;
 
 mod buffers;
 mod error;
+mod futures;
 mod gateway;
 pub(crate) mod prss_protocol;
 mod transport;
@@ -18,6 +19,7 @@ use std::ops::{Index, IndexMut};
 #[cfg(test)]
 pub use buffers::OrderingSender;
 pub use error::{Error, Result};
+pub use futures::MaybeFuture;
 
 #[cfg(feature = "stall-detection")]
 mod gateway_exports {