refactor(core): simplify op types and codegeneration (#18843)

About 2% improvement on WS/HTTP benchmarks, possibly unlocking more
optimizations in the future.

---------

Co-authored-by: Matt Mastracci <matthew@mastracci.com>

1 files changed, 52 insertions, 67 deletions

diff --git a/core/ops.rs b/core/ops.rs
index cceeb5654..b7dcc2663 100644
--- a/core/ops.rs
+++ b/core/ops.rs
@@ -8,12 +8,10 @@ use crate::runtime::JsRuntimeState;
 use crate::OpDecl;
 use crate::OpsTracker;
 use anyhow::Error;
-use futures::future::maybe_done;
-use futures::future::FusedFuture;
 use futures::future::MaybeDone;
-use futures::ready;
-use futures::task::noop_waker;
 use futures::Future;
+use futures::FutureExt;
+use pin_project::pin_project;
 use serde::Serialize;
 use std::cell::RefCell;
 use std::ops::Deref;
@@ -22,91 +20,78 @@ use std::pin::Pin;
 use std::ptr::NonNull;
 use std::rc::Rc;
 use std::rc::Weak;
-use std::task::Context;
-use std::task::Poll;
 use v8::fast_api::CFunctionInfo;
 use v8::fast_api::CTypeInfo;
 
-/// Wrapper around a Future, which causes that Future to be polled immediately.
-///
-/// Background: ops are stored in a `FuturesUnordered` structure which polls
-/// them, but without the `OpCall` wrapper this doesn't happen until the next
-/// turn of the event loop, which is too late for certain ops.
-pub struct OpCall<T>(MaybeDone<Pin<Box<dyn Future<Output = T>>>>);
+pub type RealmIdx = u16;
+pub type PromiseId = i32;
+pub type OpId = u16;
 
-pub enum EagerPollResult<T> {
-  Ready(T),
-  Pending(OpCall<T>),
+#[pin_project]
+pub struct OpCall {
+  realm_idx: RealmIdx,
+  promise_id: PromiseId,
+  op_id: OpId,
+  /// Future is not necessarily Unpin, so we need to pin_project.
+  #[pin]
+  fut: MaybeDone<Pin<Box<dyn Future<Output = OpResult>>>>,
 }
 
-impl<T> OpCall<T> {
-  /// Wraps a future, and polls the inner future immediately.
-  /// This should be the default choice for ops.
-  pub fn eager(fut: impl Future<Output = T> + 'static) -> EagerPollResult<T> {
-    let boxed = Box::pin(fut) as Pin<Box<dyn Future<Output = T>>>;
-    let mut inner = maybe_done(boxed);
-    let waker = noop_waker();
-    let mut cx = Context::from_waker(&waker);
-    let mut pinned = Pin::new(&mut inner);
-    let poll = pinned.as_mut().poll(&mut cx);
-    match poll {
-      Poll::Ready(_) => EagerPollResult::Ready(pinned.take_output().unwrap()),
-      _ => EagerPollResult::Pending(Self(inner)),
-    }
-  }
-
+impl OpCall {
   /// Wraps a future; the inner future is polled the usual way (lazily).
-  pub fn lazy(fut: impl Future<Output = T> + 'static) -> Self {
-    let boxed = Box::pin(fut) as Pin<Box<dyn Future<Output = T>>>;
-    let inner = maybe_done(boxed);
-    Self(inner)
+  pub fn pending(
+    op_ctx: &OpCtx,
+    promise_id: PromiseId,
+    fut: Pin<Box<dyn Future<Output = OpResult> + 'static>>,
+  ) -> Self {
+    Self {
+      realm_idx: op_ctx.realm_idx,
+      op_id: op_ctx.id,
+      promise_id,
+      fut: MaybeDone::Future(fut),
+    }
   }
 
   /// Create a future by specifying its output. This is basically the same as
   /// `async { value }` or `futures::future::ready(value)`.
-  pub fn ready(value: T) -> Self {
-    Self(MaybeDone::Done(value))
+  pub fn ready(op_ctx: &OpCtx, promise_id: PromiseId, value: OpResult) -> Self {
+    Self {
+      realm_idx: op_ctx.realm_idx,
+      op_id: op_ctx.id,
+      promise_id,
+      fut: MaybeDone::Done(value),
+    }
   }
 }
 
-impl<T> Future for OpCall<T> {
-  type Output = T;
+impl Future for OpCall {
+  type Output = (RealmIdx, PromiseId, OpId, OpResult);
 
   fn poll(
     self: std::pin::Pin<&mut Self>,
     cx: &mut std::task::Context<'_>,
   ) -> std::task::Poll<Self::Output> {
-    // TODO(piscisaureus): safety comment
-    #[allow(clippy::undocumented_unsafe_blocks)]
-    let inner = unsafe { &mut self.get_unchecked_mut().0 };
-    let mut pinned = Pin::new(inner);
-    ready!(pinned.as_mut().poll(cx));
-    Poll::Ready(pinned.as_mut().take_output().unwrap())
-  }
-}
-
-impl<F> FusedFuture for OpCall<F>
-where
-  F: Future,
-{
-  fn is_terminated(&self) -> bool {
-    self.0.is_terminated()
+    let realm_idx = self.realm_idx;
+    let promise_id = self.promise_id;
+    let op_id = self.op_id;
+    let fut = &mut *self.project().fut;
+    match fut {
+      MaybeDone::Done(_) => {
+        // Let's avoid using take_output as it keeps our Pin::box
+        let res = std::mem::replace(fut, MaybeDone::Gone);
+        let MaybeDone::Done(res) = res
+        else {
+          unreachable!()
+        };
+        std::task::Poll::Ready(res)
+      }
+      MaybeDone::Future(f) => f.poll_unpin(cx),
+      MaybeDone::Gone => std::task::Poll::Pending,
+    }
+    .map(move |res| (realm_idx, promise_id, op_id, res))
   }
 }
 
-pub type RealmIdx = usize;
-pub type PromiseId = i32;
-pub type OpAsyncFuture = OpCall<(PromiseId, OpId, OpResult)>;
-pub type OpFn =
-  fn(&mut v8::HandleScope, v8::FunctionCallbackArguments, v8::ReturnValue);
-pub type OpId = usize;
-
-pub enum Op {
-  Sync(OpResult),
-  Async(OpAsyncFuture),
-  NotFound,
-}
-
 pub enum OpResult {
   Ok(serde_v8::SerializablePkg),
   Err(OpError),