fix(sdk): reduce union/either variant if required in apply syntax

dev/tree-view.ts

@@ -21,7 +21,12 @@ export function treeView(tg: TypeGraphDS, rootIdx = 0, depth = 4) {
     const indent = "    ".repeat(path.edges.length);
     const edge = cyan(`${path.edges[path.edges.length - 1] ?? "[root]"}`);
     const idxStr = green(`${idx}`);
-    console.log(`${indent}${edge} → ${idxStr} ${type.type}:${type.title}`);
+    const injection = type.injection
+      ? ` (injection ${type.injection.source})`
+      : "";
+    console.log(
+      `${indent}${edge} → ${idxStr} ${type.type}:${type.title}${injection}`,
+    );
     return path.edges.length < depth;
   }, { allowCircular: true });
 }

typegate/tests/typecheck/apply.py

@@ -3,15 +3,40 @@
 
 simple_tpe = t.struct(
     {
-        "one": t.string(),
+        "one": t.union([t.string(), t.integer()]).optional(),
         "two": t.struct(
             {
                 "apply": t.integer(),
                 "user": t.integer(),
                 "set": t.integer().optional(),
                 "context": t.string().optional(),
             }
-        ),
+        ).optional(),
+        "branching": t.union(
+            [
+                # *.a.b
+                t.struct({"a": t.struct({"b": t.string()})}, name="V1"),
+                # *.a.b.c
+                #      .d
+                t.struct(
+                    {
+                        "a": t.struct(
+                            {
+                                "b": t.union(
+                                    [
+                                        t.struct({"c": t.string()}, name="A"),
+                                        t.struct(
+                                            {"c": t.string(), "d": t.string()}, name="B"
+                                        ),
+                                    ]
+                                )
+                            }
+                        )
+                    },
+                    name="V2",
+                ),
+            ]
+        ).optional(),
     }
 )
 
@@ -61,6 +86,9 @@ def test_apply_python(g: Graph):
             }
         )
         .with_policy(public),
+        testBranching=identity_simple.apply(
+            {"branching": {"a": {"b": {"c": "nested"}}}}
+        ).with_policy(public),
         selfReferingType=identity_self_ref.apply(
             {
                 "a": g.inherit(),  # A1

typegate/tests/typecheck/apply_syntax_test.ts

@@ -217,6 +217,40 @@ Meta.test("python(sdk): apply", async (t) => {
     },
   );
 
+  await t.should(
+    "work with nested union/either",
+    async () => {
+      await gql`
+        query {
+            testBranching {
+              branching {
+                ... on V1 { a { b } }
+                ... on V2 { a { 
+                    b {
+                      ... on A { c }
+                      ... on B { c }
+                    }
+                }
+              }
+            }
+          }
+        }
+      `
+        .expectData({
+          testBranching: {
+            branching: {
+              a: {
+                b: {
+                  c: "nested",
+                },
+              },
+            },
+          },
+        })
+        .on(e);
+    },
+  );
+
   await t.should(
     "work with self-refering type",
     async () => {

typegraph/core/src/global_store.rs

@@ -164,6 +164,83 @@ impl Store {
         })
     }
 
+    pub fn get_reduced_branching(supertype_id: TypeId, path: &[String]) -> Result<(Type, TypeId)> {
+        let supertype = supertype_id.as_type()?;
+        let map_reduce = |variants: Vec<u32>| match path.len() {
+            0 => Ok((supertype.clone(), supertype_id)), // terminal node
+            _ => {
+                let mut compatible = vec![];
+                let mut failures = vec![];
+                let chunk = path.first().unwrap();
+                for (i, variant) in variants.iter().enumerate() {
+                    let variant: TypeId = variant.into();
+                    let unwrapped_variant = variant.resolve_wrapper()?.as_type()?;
+                    match unwrapped_variant {
+                        Type::Struct(t) => {
+                            for (prop_name, prop_id) in t.iter_props() {
+                                if prop_name.eq(chunk) {
+                                    // variant is compatible with the path
+                                    // try expanding it, if it fails, just skip
+                                    match Store::get_type_by_path(prop_id, &path[1..]) {
+                                        Ok((_, solution)) => compatible.push(solution),
+                                        Err(e) => failures.push(format!(
+                                            "[v{i} → {prop_name}]: {}",
+                                            e.stack.first().unwrap().clone()
+                                        )),
+                                    }
+                                }
+                            }
+                        }
+                        Type::Either(..) | Type::Union(..) => {
+                            // get_type_by_path => get_reduced_branching
+                            match Store::get_type_by_path(variant, &path[1..]) {
+                                Ok((_, solution)) => compatible.push(solution),
+                                Err(e) => failures
+                                    .push(format!("[v{i}]: {}", e.stack.first().unwrap().clone())),
+                            }
+                        }
+                        _ => {} // skip
+                    }
+                }
+
+                if compatible.is_empty() {
+                    return Err(format!(
+                        "unable to expand variant with **.{}\nDetails:\n{}",
+                        path.join("."),
+                        failures.join("\n")
+                    )
+                    .into());
+                }
+
+                let first = compatible.first().unwrap().to_owned();
+                let ret_id = match &supertype {
+                    Type::Union(..) => first,
+                    Type::Either(..) => {
+                        if compatible.len() > 1 {
+                            return Err(format!(
+                                    "encountered an either node with more than one compatible variant to the path {:?}",
+                                    path.join("."),
+                                ).into(),
+                            );
+                        }
+                        first
+                    }
+                    _ => {
+                        return Err("invalid state: either or union expected as supertype".into());
+                    }
+                };
+
+                Ok((ret_id.as_type()?, ret_id))
+            }
+        };
+
+        match &supertype {
+            Type::Either(t) => map_reduce(t.data.variants.clone()),
+            Type::Union(t) => map_reduce(t.data.variants.clone()),
+            _ => Store::get_type_by_path(supertype_id, path), // no branching, trivial case
+        }
+    }
+
     pub fn get_type_by_path(struct_id: TypeId, path: &[String]) -> Result<(Type, TypeId)> {
         let mut ret = (struct_id.as_type()?, struct_id);
 
@@ -190,6 +267,10 @@ impl Store {
                         }
                     };
                 }
+                Type::Union(..) | Type::Either(..) => {
+                    ret = Store::get_reduced_branching(unwrapped_id, &path[pos..])?;
+                    break;
+                }
                 _ => return Err(errors::expect_object_at_path(&curr_path)),
             }
         }

typegraph/core/src/utils/mod.rs

@@ -16,7 +16,7 @@ use crate::wit::runtimes::MaterializerDenoFunc;
 use crate::wit::utils::Auth as WitAuth;
 use crate::{t, Lib};
 
-mod apply;
+pub mod apply;
 
 fn find_missing_props(
     supertype_id: TypeId,
@@ -109,7 +109,8 @@ impl crate::wit::utils::Guest for crate::Lib {
 
             if item.node.is_leaf() {
                 let path_infos = item.node.path_infos.clone();
-                let apply_value = path_infos.value;
+                let apply_value = path_infos.value.clone();
+                // let id = Store::get_reduced_branching(supertype_id.into(), &path_infos.path)?.1;
                 let id = Store::get_type_by_path(supertype_id.into(), &path_infos.path)?.1;
 
                 if apply_value.inherit && apply_value.payload.is_none() {

typegraph/core/wit/typegraph.wit

@@ -468,6 +468,7 @@ interface utils {
     // ]
     record apply-value {
         inherit: bool,
+        // json string
         payload: option<string>
     }