fix(jit): union support (#2452)

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

src/core/ir/resolver_context_like.rs

@@ -96,7 +96,10 @@ impl SelectionField {
         field: &crate::core::jit::Field<Nested<ConstValue>, ConstValue>,
     ) -> SelectionField {
         let name = field.name.clone();
-        let selection_set = field.nested_iter().map(Self::from_jit_field).collect();
+        let selection_set = field
+            .nested_iter(field.type_of.name())
+            .map(Self::from_jit_field)
+            .collect();
         let args = field
             .args
             .iter()

src/core/jit/builder.rs

@@ -129,7 +129,7 @@ impl Builder {
     fn iter(
         &self,
         selection: &SelectionSet,
-        type_of: &str,
+        type_condition: &str,
         exts: Option<Flat>,
         fragments: &HashMap<&str, &FragmentDefinition>,
     ) -> Vec<Field<Flat, Value>> {
@@ -170,7 +170,7 @@ impl Builder {
                         .map(|(k, v)| (k.node.as_str().to_string(), v.node.to_owned()))
                         .collect::<HashMap<_, _>>();
 
-                    if let Some(field_def) = self.index.get_field(type_of, field_name) {
+                    if let Some(field_def) = self.index.get_field(type_condition, field_name) {
                         let mut args = Vec::with_capacity(request_args.len());
                         if let QueryField::Field((_, schema_args)) = field_def {
                             for (arg_name, arg_value) in schema_args {
@@ -220,6 +220,7 @@ impl Builder {
                             ir,
                             is_scalar: self.index.type_is_scalar(type_of.name()),
                             type_of,
+                            type_condition: type_condition.to_string(),
                             skip,
                             include,
                             args,
@@ -246,7 +247,20 @@ impl Builder {
                         ));
                     }
                 }
-                _ => {}
+                Selection::InlineFragment(Positioned { node: fragment, .. }) => {
+                    let type_of = fragment
+                        .type_condition
+                        .as_ref()
+                        .map(|cond| cond.node.on.node.as_str())
+                        .unwrap_or(type_condition);
+
+                    fields.extend(self.iter(
+                        &fragment.selection_set.node,
+                        type_of,
+                        exts.clone(),
+                        fragments,
+                    ));
+                }
             }
         }
 
@@ -429,9 +443,15 @@ mod tests {
               phone
             }
 
+            fragment PostPII on Post {
+              title
+              body
+            }
+
             query {
               user(id:1) {
                 ...UserPII
+                ...PostPII
               }
             }
         "#,

src/core/jit/common/jp.rs

@@ -7,6 +7,7 @@ use crate::core::blueprint::Blueprint;
 use crate::core::config::{Config, ConfigModule};
 use crate::core::jit;
 use crate::core::jit::builder::Builder;
+use crate::core::jit::exec::TypedValue;
 use crate::core::jit::store::{Data, Store};
 use crate::core::jit::synth::Synth;
 use crate::core::jit::{OperationPlan, Variables};
@@ -27,9 +28,11 @@ struct TestData<Value> {
     users: Vec<Value>,
 }
 
+type Entry<Value> = Data<Result<TypedValue<Value>, Positioned<jit::Error>>>;
+
 struct ProcessedTestData<Value> {
     posts: Value,
-    users: HashMap<usize, Data<Result<Value, Positioned<jit::Error>>>>,
+    users: HashMap<usize, Entry<Value>>,
 }
 
 impl<'a, Value: JsonLike<'a> + Deserialize<'a> + Clone + 'a> TestData<Value> {
@@ -74,6 +77,7 @@ impl<'a, Value: JsonLike<'a> + Deserialize<'a> + Clone + 'a> TestData<Value> {
                     Value::null()
                 }
             })
+            .map(TypedValue::new)
             .map(Ok)
             .map(Data::Single)
             .enumerate()
@@ -128,7 +132,7 @@ impl<
             .to_owned();
 
         let store = [
-            (posts_id, Data::Single(Ok(posts))),
+            (posts_id, Data::Single(Ok(TypedValue::new(posts)))),
             (users_id, Data::Multiple(users)),
         ]
         .into_iter()

src/core/jit/error.rs

@@ -30,6 +30,8 @@ pub enum ValidationError {
     // with async_graphql error message for this case
     #[error(r#"internal: invalid value for scalar "{type_of}", expected "FieldValue::Value""#)]
     ScalarInvalid { type_of: String, path: String },
+    #[error("TypeName shape doesn't satisfy the processed object")]
+    TypeNameMismatch,
 }
 
 #[derive(Debug, Clone, Error)]
@@ -58,11 +60,9 @@ impl ErrorExtensions for Error {
 impl Error {
     pub fn path(&self) -> Vec<PathSegment> {
         match self {
-            Error::Validation(error) => match error {
-                ValidationError::ScalarInvalid { type_of: _, path } => {
-                    vec![PathSegment::Field(path.clone())]
-                }
-            },
+            Error::Validation(ValidationError::ScalarInvalid { type_of: _, path }) => {
+                vec![PathSegment::Field(path.clone())]
+            }
             _ => Vec::new(),
         }
     }

src/core/jit/exec.rs

@@ -9,11 +9,12 @@ use futures_util::future::join_all;
 use super::context::Context;
 use super::{DataPath, OperationPlan, Request, Response, Store};
 use crate::core::ir::model::IR;
+use crate::core::ir::TypeName;
 use crate::core::jit;
 use crate::core::jit::synth::Synth;
 use crate::core::json::{JsonLike, JsonObjectLike};
 
-type SharedStore<Output, Error> = Arc<Mutex<Store<Result<Output, Positioned<Error>>>>>;
+type SharedStore<Output, Error> = Arc<Mutex<Store<Result<TypedValue<Output>, Positioned<Error>>>>>;
 
 ///
 /// Default GraphQL executor that takes in a GraphQL Request and produces a
@@ -37,7 +38,7 @@ where
     pub async fn store(
         &self,
         request: Request<Input>,
-    ) -> Store<Result<Output, Positioned<jit::Error>>> {
+    ) -> Store<Result<TypedValue<Output>, Positioned<jit::Error>>> {
         let store = Arc::new(Mutex::new(Store::new()));
         let mut ctx = ExecutorInner::new(request, store.clone(), self.plan.to_owned(), &self.exec);
         ctx.init().await;
@@ -104,16 +105,23 @@ where
         &'b self,
         ctx: &'b Context<'b, Input, Output>,
         data_path: &DataPath,
-        value: &'b Output,
+        result: TypedValueRef<'b, Output>,
     ) -> Result<(), Error> {
         let field = ctx.field();
+        let TypedValueRef { value, type_name } = result;
         // Array
         // Check if the field expects a list
         if field.type_of.is_list() {
             // Check if the value is an array
             if let Some(array) = value.as_array() {
-                join_all(field.nested_iter().map(|field| {
-                    join_all(array.iter().enumerate().map(|(index, value)| {
+                join_all(array.iter().enumerate().map(|(index, value)| {
+                    let type_name = match &type_name {
+                        Some(TypeName::Single(type_name)) => type_name, /* TODO: should throw */
+                        // ValidationError
+                        Some(TypeName::Vec(v)) => &v[index],
+                        None => field.type_of.name(),
+                    };
+                    join_all(field.nested_iter(type_name).map(|field| {
                         let ctx = ctx.with_value_and_field(value, field);
                         let data_path = data_path.clone().with_index(index);
                         async move { self.execute(&ctx, data_path).await }
@@ -128,7 +136,13 @@ where
         // TODO: Validate if the value is an Object
         // Has to be an Object, we don't do anything while executing if its a Scalar
         else {
-            join_all(field.nested_iter().map(|child| {
+            let type_name = match &type_name {
+                Some(TypeName::Single(type_name)) => type_name,
+                Some(TypeName::Vec(_)) => panic!("TypeName type mismatch"), /* TODO: should throw ValidationError */
+                None => field.type_of.name(),
+            };
+
+            join_all(field.nested_iter(type_name).map(|child| {
                 let ctx = ctx.with_value_and_field(value, child);
                 let data_path = data_path.clone();
                 async move { self.execute(&ctx, data_path).await }
@@ -149,8 +163,8 @@ where
         if let Some(ir) = &field.ir {
             let result = self.ir_exec.execute(ir, ctx).await;
 
-            if let Ok(ref value) = result {
-                self.iter_field(ctx, &data_path, value).await?;
+            if let Ok(ref result) = result {
+                self.iter_field(ctx, &data_path, result.as_ref()).await?;
             }
 
             let mut store = self.store.lock().unwrap();
@@ -175,15 +189,50 @@ where
                 // here without doing the "fix"
                 .unwrap_or(&default_obj);
 
-            self.iter_field(ctx, &data_path, value).await?;
+            let result = TypedValueRef { value, type_name: None };
+
+            self.iter_field(ctx, &data_path, result).await?;
         }
 
         Ok(())
     }
 }
 
+#[derive(Clone)]
+pub struct TypedValue<V> {
+    pub value: V,
+    pub type_name: Option<TypeName>,
+}
+
+pub struct TypedValueRef<'a, V> {
+    pub value: &'a V,
+    pub type_name: Option<&'a TypeName>,
+}
+
+impl<V> TypedValue<V> {
+    pub fn new(value: V) -> Self {
+        Self { value, type_name: None }
+    }
+
+    pub fn as_ref(&self) -> TypedValueRef<'_, V> {
+        TypedValueRef { value: &self.value, type_name: self.type_name.as_ref() }
+    }
+}
+
+impl<'a, V> TypedValueRef<'a, V> {
+    pub fn new(value: &'a V) -> Self {
+        Self { value, type_name: None }
+    }
+
+    pub fn map<'out, U>(&self, map: impl FnOnce(&V) -> &'out U) -> TypedValueRef<'out, U>
+    where
+        'a: 'out,
+    {
+        TypedValueRef { value: map(self.value), type_name: self.type_name }
+    }
+}
+
 /// Executor for IR
-#[async_trait::async_trait]
 pub trait IRExecutor {
     type Input;
     type Output;
@@ -192,5 +241,5 @@ pub trait IRExecutor {
         &'a self,
         ir: &'a IR,
         ctx: &'a Context<'a, Self::Input, Self::Output>,
-    ) -> Result<Self::Output, Self::Error>;
+    ) -> Result<TypedValue<Self::Output>, Self::Error>;
 }

src/core/jit/exec_const.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use async_graphql_value::ConstValue;
 
 use super::context::Context;
-use super::exec::{Executor, IRExecutor};
+use super::exec::{Executor, IRExecutor, TypedValue};
 use super::{Error, OperationPlan, Request, Response, Result};
 use crate::core::app_context::AppContext;
 use crate::core::http::RequestContext;
@@ -47,7 +47,6 @@ impl<'a> ConstValueExec<'a> {
     }
 }
 
-#[async_trait::async_trait]
 impl<'ctx> IRExecutor for ConstValueExec<'ctx> {
     type Input = ConstValue;
     type Output = ConstValue;
@@ -57,9 +56,13 @@ impl<'ctx> IRExecutor for ConstValueExec<'ctx> {
         &'a self,
         ir: &'a IR,
         ctx: &'a Context<'a, Self::Input, Self::Output>,
-    ) -> Result<Self::Output> {
+    ) -> Result<TypedValue<Self::Output>> {
         let req_context = &self.req_context;
-        let mut ctx = EvalContext::new(req_context, ctx);
-        Ok(ir.eval(&mut ctx).await?)
+        let mut eval_ctx = EvalContext::new(req_context, ctx);
+
+        Ok(ir
+            .eval(&mut eval_ctx)
+            .await
+            .map(|value| TypedValue { value, type_name: eval_ctx.type_name.take() })?)
     }
 }