fix(jit): remove excessive plan & field cloning (#3016)

Co-authored-by: Sandipsinh Rathod <[email protected]>
Co-authored-by: Sandipsinh Dilipsinh Rathod <[email protected]>
Co-authored-by: Tushar Mathur <[email protected]>
Co-authored-by: laststylebender <[email protected]>

Cargo.toml

@@ -270,6 +270,10 @@ debug = false
 incremental = false
 overflow-checks = false
 
+[profile.benchmark]
+inherits = "release"
+debug = true
+
 [profile.release.package.tailcall-cloudflare]
 strip = true
 codegen-units = 1

benches/bench_synth.rs

@@ -1,5 +1,5 @@
 use criterion::Criterion;
-use tailcall::core::jit::common::JP;
+use tailcall::core::jit::fixtures::JP;
 
 pub fn bench_synth_nested(c: &mut Criterion) {
     c.bench_function("synth_nested", |b| {

src/core/jit/context.rs

@@ -8,21 +8,21 @@ use super::error::*;
 use super::{Field, OperationPlan, Positioned};
 use crate::core::ir::{ResolverContextLike, SelectionField};
 
-#[derive(Debug, Clone)]
-pub struct RequestContext<Input> {
-    plan: OperationPlan<Input>,
+#[derive(Debug)]
+pub struct RequestContext<'a, Input> {
+    plan: &'a OperationPlan<Input>,
     errors: Arc<Mutex<Vec<Positioned<Error>>>>,
 }
 
-impl<Input> RequestContext<Input> {
-    pub fn new(plan: OperationPlan<Input>) -> Self {
+impl<'a, Input> RequestContext<'a, Input> {
+    pub fn new(plan: &'a OperationPlan<Input>) -> Self {
         Self { plan, errors: Arc::new(Mutex::new(vec![])) }
     }
     pub fn add_error(&self, new_error: Positioned<Error>) {
         self.errors().push(new_error);
     }
     pub fn plan(&self) -> &OperationPlan<Input> {
-        &self.plan
+        self.plan
     }
     pub fn errors(&self) -> MutexGuard<Vec<Positioned<Error>>> {
         self.errors.lock().unwrap()
@@ -37,7 +37,7 @@ pub struct Context<'a, Input, Output> {
     // TODO: remove the args, since they're already present inside the fields and add support for
     // default values.
     field: &'a Field<Input>,
-    request: &'a RequestContext<Input>,
+    request: &'a RequestContext<'a, Input>,
 }
 impl<'a, Input: Clone, Output> Context<'a, Input, Output> {
     pub fn new(field: &'a Field<Input>, request: &'a RequestContext<Input>) -> Self {
@@ -136,7 +136,7 @@ mod test {
     fn test_field() {
         let plan = setup("query {posts {id title}}").unwrap();
         let field = &plan.selection;
-        let env = RequestContext::new(plan.clone());
+        let env = RequestContext::new(&plan);
         let ctx = Context::<ConstValue, ConstValue>::new(&field[0], &env);
         let expected = <Context<_, _> as ResolverContextLike>::field(&ctx).unwrap();
         insta::assert_debug_snapshot!(expected);
@@ -145,7 +145,7 @@ mod test {
     #[test]
     fn test_is_query() {
         let plan = setup("query {posts {id title}}").unwrap();
-        let env = RequestContext::new(plan.clone());
+        let env = RequestContext::new(&plan);
         let ctx = Context::new(&plan.selection[0], &env);
         assert!(ctx.is_query());
     }

src/core/jit/error.rs

@@ -51,6 +51,8 @@ pub enum Error {
     Validation(#[from] ValidationError),
     #[error("{0}")]
     ServerError(async_graphql::ServerError),
+    #[error("Unexpected error")]
+    Unknown,
 }
 
 impl ErrorExtensions for Error {
@@ -61,6 +63,7 @@ impl ErrorExtensions for Error {
             Error::IR(error) => error.extend(),
             Error::Validation(error) => error.extend(),
             Error::ServerError(error) => error.extend(),
+            Error::Unknown => async_graphql::Error::new(self.to_string()),
         }
     }
 }

src/core/jit/exec.rs

@@ -18,18 +18,18 @@ type SharedStore<Output, Error> = Arc<Mutex<Store<Result<Output, Positioned<Erro
 ///
 /// Default GraphQL executor that takes in a GraphQL Request and produces a
 /// GraphQL Response
-pub struct Executor<IRExec, Input> {
-    ctx: RequestContext<Input>,
+pub struct Executor<'a, IRExec, Input> {
+    ctx: RequestContext<'a, Input>,
     exec: IRExec,
 }
 
-impl<Input, Output, Exec> Executor<Exec, Input>
+impl<'a, Input, Output, Exec> Executor<'a, Exec, Input>
 where
     Output: for<'b> JsonLike<'b> + Debug + Clone,
     Input: Clone + Debug,
     Exec: IRExecutor<Input = Input, Output = Output, Error = jit::Error>,
 {
-    pub fn new(plan: OperationPlan<Input>, exec: Exec) -> Self {
+    pub fn new(plan: &'a OperationPlan<Input>, exec: Exec) -> Self {
         Self { exec, ctx: RequestContext::new(plan) }
     }
 
@@ -42,7 +42,7 @@ where
         store
     }
 
-    pub async fn execute(self, synth: Synth<Output>) -> Response<Output, jit::Error> {
+    pub async fn execute(self, synth: Synth<'a, Output>) -> Response<Output, jit::Error> {
         let mut response = Response::new(synth.synthesize());
         response.add_errors(self.ctx.errors().clone());
         response
@@ -53,7 +53,7 @@ where
 struct ExecutorInner<'a, Input, Output, Error, Exec> {
     store: SharedStore<Output, Error>,
     ir_exec: &'a Exec,
-    request: &'a RequestContext<Input>,
+    request: &'a RequestContext<'a, Input>,
 }
 
 impl<'a, Input, Output, Error, Exec> ExecutorInner<'a, Input, Output, Error, Exec>

src/core/jit/exec_const.rs

@@ -45,10 +45,17 @@ impl ConstValueExecutor {
         let is_const = self.plan.is_const;
 
         // Attempt to skip unnecessary fields
-        let plan = transform::Skip::new(variables)
-            .transform(self.plan.clone())
+        let Ok(plan) = transform::Skip::new(variables)
+            .transform(self.plan)
             .to_result()
-            .unwrap_or(self.plan);
+        else {
+            // this shouldn't actually ever happen
+            return Response {
+                data: None,
+                errors: vec![Positioned::new(Error::Unknown, Pos { line: 0, column: 0 })],
+                extensions: Default::default(),
+            };
+        };
 
         // Attempt to replace variables in the plan with the actual values
         // TODO: operation from [ExecutableDocument] could contain definitions for
@@ -71,16 +78,16 @@ impl ConstValueExecutor {
             }
         };
 
-        // TODO: drop the clones in plan
-        let exec = ConstValueExec::new(plan.clone(), req_ctx);
+        let exec = ConstValueExec::new(&plan, req_ctx);
+        // PERF: remove this particular clone?
         let vars = request.variables.clone();
-        let exe = Executor::new(plan.clone(), exec);
+        let exe = Executor::new(&plan, exec);
         let store = exe.store().await;
-        let synth = Synth::new(plan, store, vars);
+        let synth = Synth::new(&plan, store, vars);
         let response = exe.execute(synth).await;
 
         // Cache the response if we know the output is always the same
-        if is_const && self.response.is_none() {
+        if is_const {
             self.response = Some(response.clone());
         }
 
@@ -89,12 +96,12 @@ impl ConstValueExecutor {
 }
 
 struct ConstValueExec<'a> {
-    plan: OperationPlan<ConstValue>,
+    plan: &'a OperationPlan<ConstValue>,
     req_context: &'a RequestContext,
 }
 
 impl<'a> ConstValueExec<'a> {
-    pub fn new(plan: OperationPlan<ConstValue>, req_context: &'a RequestContext) -> Self {
+    pub fn new(plan: &'a OperationPlan<ConstValue>, req_context: &'a RequestContext) -> Self {
         Self { req_context, plan }
     }
 

src/core/jit/common/jp.rs → src/core/jit/fixtures/jp.rs

@@ -117,11 +117,11 @@ impl<'a, Value: Deserialize<'a> + Clone + 'a + JsonLike<'a> + std::fmt::Debug> J
 
     pub fn synth(&'a self) -> Synth<Value> {
         let ProcessedTestData { posts, users } = self.test_data.to_processed();
-        let plan = self.plan.clone();
         let vars = self.vars.clone();
 
-        let posts_id = plan.find_field_path(&["posts"]).unwrap().id.to_owned();
-        let users_id = plan
+        let posts_id = self.plan.find_field_path(&["posts"]).unwrap().id.to_owned();
+        let users_id = self
+            .plan
             .find_field_path(&["posts", "user"])
             .unwrap()
             .id
@@ -134,6 +134,6 @@ impl<'a, Value: Deserialize<'a> + Clone + 'a + JsonLike<'a> + std::fmt::Debug> J
                 store
             });
 
-        Synth::new(plan, store, vars)
+        Synth::new(&self.plan, store, vars)
     }
 }

src/core/jit/mod.rs

@@ -14,7 +14,7 @@ mod response;
 
 // NOTE: Only used in tests and benchmarks
 mod builder;
-pub mod common;
+pub mod fixtures;
 mod graphql_executor;
 
 // Public Exports

src/core/jit/request.rs

@@ -45,12 +45,14 @@ impl Request<ConstValue> {
         let builder = Builder::new(blueprint, doc);
         let plan = builder.build(self.operation_name.as_deref())?;
 
-        Ok(transform::CheckConst::new()
+        transform::CheckConst::new()
             .pipe(transform::CheckDedupe::new())
             .transform(plan)
             .to_result()
-            // NOTE: Unwrapping because these transformations fail with () error
-            .unwrap())
+            // both transformers are infallible right now
+            // but we can't just unwrap this in stable rust
+            // so convert to the Unknown error
+            .map_err(|_| super::Error::Unknown)
     }
 }
 

src/core/jit/synth/synth.rs

@@ -8,24 +8,24 @@ use crate::core::scalar;
 
 type ValueStore<Value> = Store<Result<Value, Positioned<Error>>>;
 
-pub struct Synth<Value> {
-    plan: OperationPlan<Value>,
+pub struct Synth<'a, Value> {
+    plan: &'a OperationPlan<Value>,
     store: ValueStore<Value>,
     variables: Variables<Value>,
 }
 
-impl<Value> Synth<Value> {
+impl<'a, Value> Synth<'a, Value> {
     #[inline(always)]
     pub fn new(
-        plan: OperationPlan<Value>,
+        plan: &'a OperationPlan<Value>,
         store: ValueStore<Value>,
         variables: Variables<Value>,
     ) -> Self {
         Self { plan, store, variables }
     }
 }
 
-impl<'a, Value> Synth<Value>
+impl<'a, Value> Synth<'a, Value>
 where
     Value: JsonLike<'a> + Clone + std::fmt::Debug,
 {
@@ -233,13 +233,15 @@ mod tests {
     use async_graphql_value::ConstValue;
     use serde::{Deserialize, Serialize};
 
+    use super::ValueStore;
     use crate::core::blueprint::Blueprint;
     use crate::core::config::{Config, ConfigModule};
     use crate::core::jit::builder::Builder;
-    use crate::core::jit::common::JP;
+    use crate::core::jit::fixtures::JP;
     use crate::core::jit::model::{FieldId, Variables};
     use crate::core::jit::store::Store;
     use crate::core::jit::synth::Synth;
+    use crate::core::jit::OperationPlan;
     use crate::core::json::JsonLike;
     use crate::core::valid::Validator;
 
@@ -308,7 +310,10 @@ mod tests {
 
     const CONFIG: &str = include_str!("../fixtures/jsonplaceholder-mutation.graphql");
 
-    fn make_store<'a, Value>(query: &str, store: Vec<(FieldId, TestData)>) -> Synth<Value>
+    fn make_store<'a, Value>(
+        query: &str,
+        store: Vec<(FieldId, TestData)>,
+    ) -> (OperationPlan<Value>, ValueStore<Value>, Variables<Value>)
     where
         Value: Deserialize<'a> + JsonLike<'a> + Serialize + Clone + std::fmt::Debug,
     {
@@ -343,27 +348,21 @@ mod tests {
             });
         let vars = Variables::new();
 
-        super::Synth::new(plan, store, vars)
+        (plan, store, vars)
     }
 
-    struct Synths<'a> {
-        synth_const: Synth<async_graphql::Value>,
-        synth_borrow: Synth<serde_json_borrow::Value<'a>>,
-    }
-
-    impl<'a> Synths<'a> {
-        fn init(query: &str, store: Vec<(FieldId, TestData)>) -> Self {
-            let synth_const = make_store::<ConstValue>(query, store.clone());
-            let synth_borrow = make_store::<serde_json_borrow::Value>(query, store.clone());
-            Self { synth_const, synth_borrow }
-        }
-        fn assert(self) {
-            let val_const = self.synth_const.synthesize().unwrap();
-            let val_const = serde_json::to_string_pretty(&val_const).unwrap();
-            let val_borrow = self.synth_borrow.synthesize().unwrap();
-            let val_borrow = serde_json::to_string_pretty(&val_borrow).unwrap();
-            assert_eq!(val_const, val_borrow);
-        }
+    fn assert_synths(query: &str, store: Vec<(FieldId, TestData)>) {
+        let (plan, value_store, vars) = make_store::<ConstValue>(query, store.clone());
+        let synth_const = Synth::new(&plan, value_store, vars);
+        let (plan, value_store, vars) =
+            make_store::<serde_json_borrow::Value>(query, store.clone());
+        let synth_borrow = Synth::new(&plan, value_store, vars);
+
+        let val_const = synth_const.synthesize().unwrap();
+        let val_const = serde_json::to_string_pretty(&val_const).unwrap();
+        let val_borrow = synth_borrow.synthesize().unwrap();
+        let val_borrow = serde_json::to_string_pretty(&val_borrow).unwrap();
+        assert_eq!(val_const, val_borrow);
     }
 
     #[test]
@@ -375,8 +374,7 @@ mod tests {
             }
         "#;
 
-        let synths = Synths::init(query, store);
-        synths.assert();
+        assert_synths(query, store);
     }
 
     #[test]
@@ -388,8 +386,7 @@ mod tests {
                 }
             "#;
 
-        let synths = Synths::init(query, store);
-        synths.assert();
+        assert_synths(query, store);
     }
 
     #[test]
@@ -403,8 +400,7 @@ mod tests {
                     posts { id title user { id name } }
                 }
             "#;
-        let synths = Synths::init(query, store);
-        synths.assert();
+        assert_synths(query, store);
     }
 
     #[test]
@@ -420,8 +416,7 @@ mod tests {
                     users { id name }
                 }
             "#;
-        let synths = Synths::init(query, store);
-        synths.assert();
+        assert_synths(query, store);
     }
 
     #[test]

src/core/jit/transform/check_const.rs

@@ -1,3 +1,4 @@
+use std::convert::Infallible;
 use std::marker::PhantomData;
 
 use crate::core::ir::model::IR;
@@ -31,7 +32,7 @@ pub fn is_const(ir: &IR) -> bool {
 
 impl<A> Transform for CheckConst<A> {
     type Value = OperationPlan<A>;
-    type Error = ();
+    type Error = Infallible;
 
     fn transform(&self, mut plan: Self::Value) -> Valid<Self::Value, Self::Error> {
         let is_const = plan.iter_dfs().all(|field| match field.ir {