refactor: replace manual future enum with auto_enum

Signed-off-by: MrCroxx <mrcroxx@outlook.com>

foyer-storage/Cargo.toml

@@ -18,6 +18,7 @@ anyhow = "1.0"
 # TODO(MrCroxx): use `array_chunks` after `#![feature(array_chunks)]` is stable.
 array-util = "1"
 async-channel = "2"
+auto_enums = { version = "0.8", features = ["futures03"] }
 bincode = "1"
 bitflags = "2.3.1"
 bytes = "1"

foyer-storage/src/engine.rs

@@ -12,15 +12,10 @@
 //  See the License for the specific language governing permissions and
 //  limitations under the License.
 
-use std::{
-    fmt::Debug,
-    marker::PhantomData,
-    pin::Pin,
-    sync::Arc,
-    task::{Context, Poll},
-};
+use std::{fmt::Debug, marker::PhantomData, sync::Arc};
 
 use ahash::RandomState;
+use auto_enums::auto_enum;
 use foyer_common::code::{HashBuilder, StorageKey, StorageValue};
 use foyer_memory::CacheEntry;
 use futures::Future;
@@ -96,45 +91,6 @@ where
     }
 }
 
-enum StoreFuture<F1, F2, F3, F4> {
-    Noop(F1),
-    Large(F2),
-    Small(F3),
-    Combined(F4),
-}
-
-impl<F1, F2, F3, F4> StoreFuture<F1, F2, F3, F4> {
-    pub fn as_pin_mut(self: Pin<&mut Self>) -> StoreFuture<Pin<&mut F1>, Pin<&mut F2>, Pin<&mut F3>, Pin<&mut F4>> {
-        unsafe {
-            match *Pin::get_unchecked_mut(self) {
-                StoreFuture::Noop(ref mut inner) => StoreFuture::Noop(Pin::new_unchecked(inner)),
-                StoreFuture::Large(ref mut inner) => StoreFuture::Large(Pin::new_unchecked(inner)),
-                StoreFuture::Small(ref mut inner) => StoreFuture::Small(Pin::new_unchecked(inner)),
-                StoreFuture::Combined(ref mut inner) => StoreFuture::Combined(Pin::new_unchecked(inner)),
-            }
-        }
-    }
-}
-
-impl<F1, F2, F3, F4> Future for StoreFuture<F1, F2, F3, F4>
-where
-    F1: Future,
-    F2: Future<Output = F1::Output>,
-    F3: Future<Output = F1::Output>,
-    F4: Future<Output = F1::Output>,
-{
-    type Output = F1::Output;
-
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        match self.as_pin_mut() {
-            StoreFuture::Noop(future) => future.poll(cx),
-            StoreFuture::Large(future) => future.poll(cx),
-            StoreFuture::Small(future) => future.poll(cx),
-            StoreFuture::Combined(future) => future.poll(cx),
-        }
-    }
-}
-
 #[expect(clippy::type_complexity)]
 pub enum EngineConfig<K, V, S = RandomState>
 where
@@ -251,12 +207,13 @@ where
         }
     }
 
+    #[auto_enum(Future)]
     fn load(&self, hash: u64) -> impl Future<Output = Result<Option<(Self::Key, Self::Value)>>> + Send + 'static {
         match self {
-            Engine::Noop(storage) => StoreFuture::Noop(storage.load(hash)),
-            Engine::Large(storage) => StoreFuture::Large(storage.load(hash)),
-            Engine::Small(storage) => StoreFuture::Small(storage.load(hash)),
-            Engine::Combined(storage) => StoreFuture::Combined(storage.load(hash)),
+            Engine::Noop(storage) => storage.load(hash),
+            Engine::Large(storage) => storage.load(hash),
+            Engine::Small(storage) => storage.load(hash),
+            Engine::Combined(storage) => storage.load(hash),
         }
     }
 
@@ -296,12 +253,13 @@ where
         }
     }
 
+    #[auto_enum(Future)]
     fn wait(&self) -> impl Future<Output = ()> + Send + 'static {
         match self {
-            Engine::Noop(storage) => StoreFuture::Noop(storage.wait()),
-            Engine::Large(storage) => StoreFuture::Large(storage.wait()),
-            Engine::Small(storage) => StoreFuture::Small(storage.wait()),
-            Engine::Combined(storage) => StoreFuture::Combined(storage.wait()),
+            Engine::Noop(storage) => storage.wait(),
+            Engine::Large(storage) => storage.wait(),
+            Engine::Small(storage) => storage.wait(),
+            Engine::Combined(storage) => storage.wait(),
         }
     }
 }

foyer-storage/src/storage/either.rs

@@ -12,13 +12,9 @@
 //  See the License for the specific language governing permissions and
 //  limitations under the License.
 
-use std::{
-    fmt::Debug,
-    pin::Pin,
-    sync::Arc,
-    task::{Context, Poll},
-};
+use std::{fmt::Debug, sync::Arc};
 
+use auto_enums::auto_enum;
 use foyer_common::code::{HashBuilder, StorageKey, StorageValue};
 use foyer_memory::CacheEntry;
 use futures::{
@@ -28,41 +24,6 @@ use futures::{
 
 use crate::{error::Result, storage::Storage, DeviceStats};
 
-enum OrderFuture<F1, F2, F3> {
-    LeftFirst(F1),
-    RightFirst(F2),
-    Parallel(F3),
-}
-
-impl<F1, F2, F3> OrderFuture<F1, F2, F3> {
-    pub fn as_pin_mut(self: Pin<&mut Self>) -> OrderFuture<Pin<&mut F1>, Pin<&mut F2>, Pin<&mut F3>> {
-        unsafe {
-            match *Pin::get_unchecked_mut(self) {
-                OrderFuture::LeftFirst(ref mut inner) => OrderFuture::LeftFirst(Pin::new_unchecked(inner)),
-                OrderFuture::RightFirst(ref mut inner) => OrderFuture::RightFirst(Pin::new_unchecked(inner)),
-                OrderFuture::Parallel(ref mut inner) => OrderFuture::Parallel(Pin::new_unchecked(inner)),
-            }
-        }
-    }
-}
-
-impl<F1, F2, F3> Future for OrderFuture<F1, F2, F3>
-where
-    F1: Future,
-    F2: Future<Output = F1::Output>,
-    F3: Future<Output = F1::Output>,
-{
-    type Output = F1::Output;
-
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        match self.as_pin_mut() {
-            OrderFuture::LeftFirst(future) => future.poll(cx),
-            OrderFuture::RightFirst(future) => future.poll(cx),
-            OrderFuture::Parallel(future) => future.poll(cx),
-        }
-    }
-}
-
 /// Order of ops.
 #[derive(Debug, Clone, Copy)]
 pub enum Order {
@@ -226,24 +187,25 @@ where
         }
     }
 
+    #[auto_enum(Future)]
     fn load(&self, hash: u64) -> impl Future<Output = Result<Option<(Self::Key, Self::Value)>>> + Send + 'static {
         let fleft = self.left.load(hash);
         let fright = self.right.load(hash);
         match self.load_order {
             // FIXME(MrCroxx): false-positive on hash collision.
-            Order::LeftFirst => OrderFuture::LeftFirst(fleft.then(|res| match res {
+            Order::LeftFirst => fleft.then(|res| match res {
                 Ok(Some(kv)) => ready(Ok(Some(kv))).left_future(),
                 Err(e) => ready(Err(e)).left_future(),
                 Ok(None) => fright.right_future(),
-            })),
+            }),
             // FIXME(MrCroxx): false-positive on hash collision.
-            Order::RightFirst => OrderFuture::RightFirst(fright.then(|res| match res {
+            Order::RightFirst => fright.then(|res| match res {
                 Ok(Some(kv)) => ready(Ok(Some(kv))).left_future(),
                 Err(e) => ready(Err(e)).left_future(),
                 Ok(None) => fleft.right_future(),
-            })),
+            }),
             Order::Parallel => {
-                OrderFuture::Parallel(async move {
+                async move {
                     pin_mut!(fleft);
                     pin_mut!(fright);
                     // Returns a 4-way `Either` by nesting `Either` in `Either`.
@@ -261,7 +223,7 @@ where
                             },
                         })
                         .await
-                })
+                }
             }
         }
     }