fix(jit): implement skip and include (tailcallhq#2357)

Co-authored-by: Tushar Mathur <[email protected]>

src/core/jit/builder.rs

@@ -1,16 +1,49 @@
 use std::collections::HashMap;
+use std::ops::Deref;
 
 use async_graphql::parser::types::{
-    DocumentOperations, ExecutableDocument, FragmentDefinition, OperationType, Selection,
-    SelectionSet,
+    Directive, DocumentOperations, ExecutableDocument, FragmentDefinition, OperationType,
+    Selection, SelectionSet,
 };
 use async_graphql::Positioned;
+use async_graphql_value::Value;
 
 use super::model::*;
 use crate::core::blueprint::{Blueprint, Index, QueryField};
 use crate::core::counter::{Count, Counter};
 use crate::core::merge_right::MergeRight;
 
+#[derive(PartialEq)]
+enum Condition {
+    True,
+    False,
+    Variable(Variable),
+}
+
+struct Conditions {
+    skip: Option<Condition>,
+    include: Option<Condition>,
+}
+
+impl Conditions {
+    /// Checks if the field should be skipped always
+    fn is_const_skip(&self) -> bool {
+        matches!(self.skip, Some(Condition::True)) ^ matches!(self.include, Some(Condition::True))
+    }
+
+    fn into_variable_tuple(self) -> (Option<Variable>, Option<Variable>) {
+        let comp = |condition| match condition? {
+            Condition::Variable(var) => Some(var),
+            _ => None,
+        };
+
+        let include = comp(self.include);
+        let skip = comp(self.skip);
+
+        (include, skip)
+    }
+}
+
 pub struct Builder {
     pub index: Index,
     pub arg_id: Counter<usize>,
@@ -29,7 +62,48 @@ impl Builder {
         }
     }
 
+    #[inline(always)]
+    fn include(
+        &self,
+        directives: &[Positioned<async_graphql::parser::types::Directive>],
+    ) -> Conditions {
+        fn get_condition(dir: &Directive) -> Option<Condition> {
+            let arg = dir.get_argument("if").map(|pos| &pos.node);
+            let is_include = dir.name.node.as_str() == "include";
+            match arg {
+                None => None,
+                Some(value) => match value {
+                    Value::Boolean(bool) => {
+                        let condition = if is_include ^ bool {
+                            Condition::True
+                        } else {
+                            Condition::False
+                        };
+                        Some(condition)
+                    }
+                    Value::Variable(var) => {
+                        Some(Condition::Variable(Variable::new(var.deref().to_owned())))
+                    }
+                    _ => None,
+                },
+            }
+        }
+        Conditions {
+            skip: directives
+                .iter()
+                .find(|d| d.node.name.node.as_str() == "skip")
+                .map(|d| &d.node)
+                .and_then(get_condition),
+            include: directives
+                .iter()
+                .find(|d| d.node.name.node.as_str() == "include")
+                .map(|d| &d.node)
+                .and_then(get_condition),
+        }
+    }
+
     #[allow(clippy::too_many_arguments)]
+    #[inline(always)]
     fn iter(
         &self,
         selection: &SelectionSet,
@@ -41,6 +115,16 @@ impl Builder {
         for selection in &selection.items {
             match &selection.node {
                 Selection::Field(Positioned { node: gql_field, .. }) => {
+                    let conditions = self.include(&gql_field.directives);
+
+                    // if include is always false xor skip is always true,
+                    // then we can skip the field from the plan
+                    if conditions.is_const_skip() {
+                        continue;
+                    }
+
+                    let (include, skip) = conditions.into_variable_tuple();
+
                     let field_name = gql_field.name.node.as_str();
                     let field_args = gql_field
                         .arguments
@@ -91,6 +175,8 @@ impl Builder {
                             name,
                             ir,
                             type_of,
+                            skip,
+                            include,
                             args,
                             extensions: exts.clone(),
                         });
@@ -116,6 +202,7 @@ impl Builder {
         fields
     }
 
+    #[inline(always)]
     fn get_type(&self, ty: OperationType) -> Option<&str> {
         match ty {
             OperationType::Query => Some(self.index.get_query()),
@@ -124,6 +211,7 @@ impl Builder {
         }
     }
 
+    #[inline(always)]
     pub fn build(&self) -> Result<ExecutionPlan, String> {
         let mut fields = Vec::new();
         let mut fragments: HashMap<&str, &FragmentDefinition> = HashMap::new();
@@ -176,7 +264,6 @@ mod tests {
         let config = Config::from_sdl(CONFIG).to_result().unwrap();
         let blueprint = Blueprint::try_from(&config.into()).unwrap();
         let document = async_graphql::parser::parse_query(query).unwrap();
-
         Builder::new(&blueprint, document).build().unwrap()
     }
 

src/core/jit/common/json_placeholder.rs

@@ -7,6 +7,7 @@ use crate::core::config::{Config, ConfigModule};
 use crate::core::jit::builder::Builder;
 use crate::core::jit::store::{Data, Store};
 use crate::core::jit::synth::Synth;
+use crate::core::jit::Variables;
 use crate::core::json::JsonLike;
 use crate::core::valid::Validator;
 
@@ -82,6 +83,7 @@ impl JsonPlaceholder {
             store
         });
 
-        Synth::new(plan, store)
+        let vars = Variables::new();
+        Synth::new(plan, store, vars)
     }
 }

src/core/jit/exec.rs

@@ -20,10 +20,10 @@ pub struct Executor<Synth, IRExec> {
     exec: IRExec,
 }
 
-impl<Input, Output, Error, Synth, Exec> Executor<Synth, Exec>
+impl<Input: Clone, Output, Error, Synth, Exec> Executor<Synth, Exec>
 where
     Output: JsonLike<Json = Output> + Debug,
-    Synth: Synthesizer<Value = Result<Output, Error>>,
+    Synth: Synthesizer<Value = Result<Output, Error>, Variable = Input>,
     Exec: IRExecutor<Input = Input, Output = Output, Error = Error>,
 {
     pub fn new(plan: ExecutionPlan, synth: Synth, exec: Exec) -> Self {
@@ -40,8 +40,9 @@ where
     }
 
     pub async fn execute(self, request: Request<Input>) -> Response<Output, Error> {
+        let vars = request.variables.clone();
         let store = self.execute_inner(request).await;
-        Response::new(self.synth.synthesize(store))
+        Response::new(self.synth.synthesize(store, vars))
     }
 }
 

src/core/jit/model.rs

@@ -1,7 +1,37 @@
+use std::collections::HashMap;
 use std::fmt::{Debug, Formatter};
 
+use serde::Deserialize;
+
 use crate::core::ir::model::IR;
 
+#[derive(Debug, Deserialize, Clone)]
+pub struct Variables<Value>(HashMap<String, Value>);
+
+impl<Value> Default for Variables<Value> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<Value> Variables<Value> {
+    pub fn new() -> Self {
+        Self(HashMap::new())
+    }
+    pub fn get(&self, key: &str) -> Option<&Value> {
+        self.0.get(key)
+    }
+    pub fn insert(&mut self, key: String, value: Value) {
+        self.0.insert(key, value);
+    }
+}
+
+impl<V> FromIterator<(String, V)> for Variables<V> {
+    fn from_iter<T: IntoIterator<Item = (String, V)>>(iter: T) -> Self {
+        Self(iter.into_iter().collect())
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct Arg {
     pub id: ArgId,
@@ -50,9 +80,40 @@ pub struct Field<Extensions> {
     pub name: String,
     pub ir: Option<IR>,
     pub type_of: crate::core::blueprint::Type,
+    pub skip: Option<Variable>,
+    pub include: Option<Variable>,
     pub args: Vec<Arg>,
     pub extensions: Option<Extensions>,
 }
+#[derive(Clone, Debug, PartialEq)]
+pub struct Variable(String);
+
+impl Variable {
+    pub fn new(name: String) -> Self {
+        Variable(name)
+    }
+}
+
+impl<A> Field<A> {
+    #[inline(always)]
+    pub fn skip(&self, variables: &Variables<async_graphql_value::ConstValue>) -> bool {
+        let eval = |variable_option: Option<&Variable>,
+                    variables: &Variables<async_graphql_value::ConstValue>,
+                    default: bool| {
+            match variable_option {
+                Some(Variable(name)) => variables.get(name).map_or(default, |value| match value {
+                    async_graphql_value::ConstValue::Boolean(b) => *b,
+                    _ => default,
+                }),
+                None => default,
+            }
+        };
+        let skip = eval(self.skip.as_ref(), variables, false);
+        let include = eval(self.include.as_ref(), variables, true);
+
+        skip == include
+    }
+}
 
 const EMPTY_VEC: &Vec<Field<Nested>> = &Vec::new();
 impl Field<Nested> {
@@ -90,6 +151,8 @@ impl Field<Flat> {
             name: self.name,
             ir: self.ir,
             type_of: self.type_of,
+            skip: self.skip,
+            include: self.include,
             args: self.args,
             extensions,
         }

src/core/jit/request.rs

@@ -5,15 +5,16 @@ use serde::Deserialize;
 
 use super::{Builder, Error, ExecutionPlan, Result};
 use crate::core::blueprint::Blueprint;
+use crate::core::jit::model::Variables;
 
-#[derive(Debug, Deserialize, Setters)]
+#[derive(Debug, Deserialize, Setters, Clone)]
 pub struct Request<Value> {
     #[serde(default)]
     pub query: String,
     #[serde(default, rename = "operationName")]
     pub operation_name: Option<String>,
     #[serde(default)]
-    pub variables: HashMap<String, Value>,
+    pub variables: Variables<Value>,
     #[serde(default)]
     pub extensions: HashMap<String, Value>,
 }
@@ -25,9 +26,9 @@ impl From<async_graphql::Request> for Request<async_graphql_value::ConstValue> {
             operation_name: value.operation_name,
             variables: match value.variables.into_value() {
                 async_graphql_value::ConstValue::Object(val) => {
-                    HashMap::from_iter(val.into_iter().map(|(k, v)| (k.to_string(), v)))
+                    Variables::from_iter(val.into_iter().map(|(name, val)| (name.to_string(), val)))
                 }
-                _ => HashMap::new(),
+                _ => Variables::default(),
             },
             extensions: value.extensions,
         }
@@ -47,7 +48,7 @@ impl<A> Request<A> {
         Self {
             query: query.to_string(),
             operation_name: None,
-            variables: HashMap::new(),
+            variables: Variables::default(),
             extensions: HashMap::new(),
         }
     }

src/core/jit/synth/mod.rs

@@ -4,8 +4,14 @@ mod synth_const;
 // pub use synth_borrow::SynthBorrow;
 pub use synth_const::{Synth, SynthConst};
 
-use super::Store;
+use super::{Store, Variables};
+use crate::core::json::JsonLike;
 pub trait Synthesizer {
     type Value;
-    fn synthesize(self, store: Store<Self::Value>) -> Self::Value;
+    type Variable: JsonLike;
+    fn synthesize(
+        self,
+        store: Store<Self::Value>,
+        variables: Variables<Self::Variable>,
+    ) -> Self::Value;
 }

src/core/jit/synth/synth_borrow.rs

@@ -131,6 +131,7 @@ mod tests {
     use crate::core::jit::model::FieldId;
     use crate::core::jit::store::{Data, Store};
     use crate::core::jit::synth::SynthBorrow;
+    use crate::core::jit::Variables;
     use crate::core::valid::Validator;
 
     const POSTS: &str = r#"
@@ -207,7 +208,11 @@ mod tests {
         let config = Config::from_sdl(CONFIG).to_result().unwrap();
         let config = ConfigModule::from(config);
 
-        let builder = Builder::new(&Blueprint::try_from(&config).unwrap(), doc);
+        let builder = Builder::new(
+            &Blueprint::try_from(&config).unwrap(),
+            doc,
+            Variables::new(),
+        );
         let plan = builder.build().unwrap();
 
         let store = store