Add a first attempt at multiplexing adapters

Signed-off-by: Marcel Müller <neikos@neikos.email>

crates/plaixt/src/main.rs

@@ -75,8 +75,10 @@ async fn main() -> miette::Result<()> {
 
             let result = execute_query(
                 &schema,
-                Arc::new(trustfall_plaixt::PlaixtAdapter {
+                Arc::new(trustfall_plaixt::TrustfallMultiAdapter {
-                    records: records.clone(),
+                    plaixt: trustfall_plaixt::PlaixtAdapter {
+                        records: records.clone(),
+                    },
                 }),
                 &query,
                 BTreeMap::<Arc<str>, FieldValue>::from([("search".into(), "trust".into())]),

crates/plaixt/src/parsing.rs

@@ -194,11 +194,15 @@ impl TryFrom<&str> for DefinitionKind {
 
 #[derive(Debug)]
 pub struct Definition {
+    pub(crate) name: String,
     pub(crate) since: Timestamp,
     pub(crate) fields: HashMap<String, DefinitionKind>,
 }
 
-pub(crate) fn parse_definition(bytes: &str) -> miette::Result<Vec<Definition>> {
+pub(crate) fn parse_definition(
+    bytes: &str,
+    definition_name: String,
+) -> miette::Result<Vec<Definition>> {
     let doc: KdlDocument = bytes.parse()?;
 
     let mut defs = vec![];
@@ -309,7 +313,11 @@ pub(crate) fn parse_definition(bytes: &str) -> miette::Result<Vec<Definition>> {
                     })
                     .collect::<miette::Result<_>>()?;
 
-                defs.push(Definition { since, fields });
+                defs.push(Definition {
+                    since,
+                    fields,
+                    name: definition_name.clone(),
+                });
             }
             unknown => {
                 return Err(miette::diagnostic!(
@@ -349,9 +357,10 @@ pub(crate) async fn load_definitions(
         })
         .flat_map(|val| futures::stream::iter(val.transpose()))
         .and_then(|(name, bytes)| async move {
+            let definition_name = name.file_stem().unwrap().to_string();
             Ok((
-                name.file_stem().unwrap().to_string(),
+                definition_name.clone(),
-                parse_definition(&bytes).map_err(|e| {
+                parse_definition(&bytes, definition_name).map_err(|e| {
                     e.with_source_code(NamedSource::new(name, bytes).with_language("kdl"))
                 })?,
             ))

crates/plaixt/src/trustfall_plaixt.rs

@@ -1,4 +1,7 @@
 use std::collections::BTreeMap;
+use std::collections::HashMap;
+use std::fmt::Write;
+use std::ops::Not;
 use std::sync::Arc;
 
 use kdl::KdlValue;
@@ -9,69 +12,351 @@ use trustfall::provider::resolve_coercion_with;
 use trustfall::provider::resolve_neighbors_with;
 use trustfall::provider::resolve_property_with;
 use trustfall::provider::Adapter;
+use trustfall::provider::AsVertex;
 use trustfall::FieldValue;
 use trustfall::Schema;
 
 use crate::parsing::Definition;
 use crate::parsing::Record;
 
+const ADAPTER_SEP: &str = "__";
+
+#[derive(Debug, Default)]
+pub struct StartingVertex {
+    adapter_name: String,
+    start_vertex_name: String,
+    vertex_type: String,
+}
+
+impl StartingVertex {
+    pub fn new(adapter_name: String, start_vertex_name: String, start_vertex_type: String) -> Self {
+        Self {
+            adapter_name,
+            start_vertex_name,
+            vertex_type: start_vertex_type,
+        }
+    }
+
+    pub fn schema_name(&self) -> String {
+        format!(
+            "{}{ADAPTER_SEP}{}",
+            self.adapter_name, self.start_vertex_name
+        )
+    }
+
+    pub fn vertex_type(&self) -> &str {
+        &self.vertex_type
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct VertexType {
+    adapter_name: String,
+    vertex_name: String,
+    vertex_fields: HashMap<String, String>,
+    implements: Vec<String>,
+}
+
+impl VertexType {
+    pub fn new(
+        adapter_name: String,
+        vertex_name: String,
+        vertex_fields: HashMap<String, String>,
+        implements: Vec<String>,
+    ) -> Self {
+        Self {
+            adapter_name,
+            vertex_name,
+            vertex_fields,
+            implements,
+        }
+    }
+
+    pub fn schema_name(&self) -> String {
+        format!("{}{ADAPTER_SEP}{}", self.adapter_name, self.vertex_name)
+    }
+
+    pub fn schema_type(&self) -> String {
+        format!(
+            r#"type {name} {impls} {{ {fields} }}"#,
+            name = self.schema_name(),
+            impls = self
+                .implements
+                .is_empty()
+                .not()
+                .then(|| format!("implements {}", self.implements.join(" & ")))
+                .unwrap_or_else(String::new),
+            fields = self.vertex_fields.iter().fold(String::new(), |mut out, f| {
+                write!(out, "{}: {}, ", f.0, f.1).unwrap();
+                out
+            }),
+        )
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct DynamicSchema {
+    roots: Vec<StartingVertex>,
+    types: Vec<VertexType>,
+}
+
+impl DynamicSchema {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn add_root(&mut self, root: StartingVertex) {
+        self.roots.push(root);
+    }
+
+    pub fn add_type(&mut self, kind: VertexType) {
+        self.types.push(kind);
+    }
+}
+
 pub(crate) fn to_schema(definitions: &BTreeMap<String, Vec<Definition>>) -> Schema {
-    let custom_schemas = definitions
+    let mut schema = DynamicSchema::new();
-        .iter()
+
-        .map(|(name, def)| {
+    schema.add_root(StartingVertex::new(
-            let fields = def
+        "Plaixt".to_string(),
-                .last()
+        "RecordsAll".to_string(),
-                .unwrap()
+        "[Record!]!".to_string(),
+    ));
+
+    for definition in definitions.values().flat_map(|d| d.first()) {
+        let fields = VertexType::new(
+            "Plaixt".to_string(),
+            format!("{}Fields", definition.name),
+            definition
                 .fields
                 .iter()
-                .map(|(name, def)| format!("{name}: {}!", def.trustfall_kind()))
+                .map(|(name, val)| (name.clone(), format!("{}!", val.trustfall_kind())))
-                .collect::<Vec<_>>()
+                .collect(),
-                .join("\n");
+            vec![],
-
+        );
-            let field_type = format!("{name}Fields");
+        schema.add_type(VertexType::new(
-
+            "Plaixt".to_string(),
-            format!(
+            definition.name.clone(),
-                r#"
+            [
-
+                (String::from("at"), String::from("String!")),
-        type {field_type} {{
+                (String::from("kind"), String::from("String!")),
-            {fields}
+                (String::from("fields"), format!("{}!", fields.schema_name())),
-        }}
+            ]
-
+            .into(),
-        type {name} implements Record {{
+            vec![String::from("Record")],
-            at: String!
+        ));
-            kind: String!
+        schema.add_type(fields);
-            fields: {field_type}!
+    }
-        }}
-        "#
-            )
-        })
-        .collect::<Vec<_>>()
-        .join("");
 
     let schema = format!(
         r#"schema {{
-    query: RootSchemaQuery
+            query: RootSchemaQuery
-}}
+        }}
-{}
+        {}
+        type RootSchemaQuery {{
+            {roots}
+        }}
+        interface Record {{
+            at: String!,
+            kind: String!,
+        }}
 
-
+        {types}
-type RootSchemaQuery {{
+        "#,
-    RecordsAll: [Record!]!
-}}
-interface Record {{
-    at: String!,
-    kind: String!,
-}}
-
-{}
-"#,
         Schema::ALL_DIRECTIVE_DEFINITIONS,
-        custom_schemas
+        roots = schema.roots.iter().fold(String::new(), |mut out, r| {
+            write!(out, "{}: {}, ", r.schema_name(), r.vertex_type()).unwrap();
+            out
+        }),
+        types = schema.types.iter().fold(String::new(), |mut out, t| {
+            writeln!(out, "{}", t.schema_type()).unwrap();
+            out
+        }),
     );
     trace!(%schema, "Using schema");
     Schema::parse(schema).unwrap()
 }
 
+pub struct TrustfallMultiAdapter {
+    pub plaixt: PlaixtAdapter,
+}
+
+#[derive(Debug, Clone)]
+pub enum TrustfallMultiVertex {
+    Plaixt(PlaixtVertex),
+}
+
+impl AsVertex<PlaixtVertex> for TrustfallMultiVertex {
+    fn as_vertex(&self) -> Option<&PlaixtVertex> {
+        self.as_plaixt()
+    }
+
+    fn into_vertex(self) -> Option<PlaixtVertex> {
+        self.as_plaixt().cloned()
+    }
+}
+
+impl TrustfallMultiVertex {
+    pub fn as_plaixt(&self) -> Option<&PlaixtVertex> {
+        if let Self::Plaixt(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
+}
+
+impl<'v> Adapter<'v> for TrustfallMultiAdapter {
+    type Vertex = TrustfallMultiVertex;
+
+    fn resolve_starting_vertices(
+        &self,
+        edge_name: &Arc<str>,
+        parameters: &trustfall::provider::EdgeParameters,
+        resolve_info: &trustfall::provider::ResolveInfo,
+    ) -> trustfall::provider::VertexIterator<'v, Self::Vertex> {
+        let (adapter_name, edge_name) = edge_name.split_once(ADAPTER_SEP).unwrap();
+
+        trace!(?adapter_name, ?edge_name, "Got start vertex");
+
+        match adapter_name {
+            "Plaixt" => {
+                let iter = self.plaixt.resolve_starting_vertices(
+                    &Arc::from(edge_name),
+                    parameters,
+                    resolve_info,
+                );
+
+                Box::new(iter.map(TrustfallMultiVertex::Plaixt))
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    fn resolve_property<V>(
+        &self,
+        contexts: trustfall::provider::ContextIterator<'v, V>,
+        type_name: &Arc<str>,
+        property_name: &Arc<str>,
+        resolve_info: &trustfall::provider::ResolveInfo,
+    ) -> trustfall::provider::ContextOutcomeIterator<'v, V, FieldValue>
+    where
+        V: trustfall::provider::AsVertex<Self::Vertex> + 'v,
+    {
+        let (adapter_name, type_name) = type_name.split_once(ADAPTER_SEP).unwrap();
+
+        match adapter_name {
+            "Plaixt" => {
+                let contexts = contexts.collect::<Vec<_>>();
+
+                let properties = self.plaixt.resolve_property(
+                    Box::new(
+                        contexts
+                            .clone()
+                            .into_iter()
+                            .map(|v| v.map(&mut |v: V| v.into_vertex().unwrap())),
+                    ),
+                    &Arc::from(type_name),
+                    property_name,
+                    resolve_info,
+                );
+
+                Box::new(
+                    properties
+                        .into_iter()
+                        .zip(contexts)
+                        .map(|((_ctx, name), og_ctx)| (og_ctx, name)),
+                )
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    fn resolve_neighbors<V: trustfall::provider::AsVertex<Self::Vertex> + 'v>(
+        &self,
+        contexts: trustfall::provider::ContextIterator<'v, V>,
+        type_name: &Arc<str>,
+        edge_name: &Arc<str>,
+        parameters: &trustfall::provider::EdgeParameters,
+        resolve_info: &trustfall::provider::ResolveEdgeInfo,
+    ) -> trustfall::provider::ContextOutcomeIterator<
+        'v,
+        V,
+        trustfall::provider::VertexIterator<'v, Self::Vertex>,
+    > {
+        let (adapter_name, type_name) = type_name.split_once(ADAPTER_SEP).unwrap();
+
+        match adapter_name {
+            "Plaixt" => {
+                let contexts = contexts.collect::<Vec<_>>();
+
+                let properties = self.plaixt.resolve_neighbors(
+                    Box::new(
+                        contexts
+                            .clone()
+                            .into_iter()
+                            .map(|v| v.map(&mut |v: V| v.into_vertex().unwrap())),
+                    ),
+                    &Arc::from(type_name),
+                    edge_name,
+                    parameters,
+                    resolve_info,
+                );
+
+                Box::new(
+                    properties
+                        .into_iter()
+                        .zip(contexts)
+                        .map(|((_ctx, vals), og_ctx)| {
+                            (
+                                og_ctx,
+                                Box::new(vals.map(TrustfallMultiVertex::Plaixt)) as Box<_>,
+                            )
+                        }),
+                )
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    fn resolve_coercion<V: trustfall::provider::AsVertex<Self::Vertex> + 'v>(
+        &self,
+        contexts: trustfall::provider::ContextIterator<'v, V>,
+        type_name: &Arc<str>,
+        coerce_to_type: &Arc<str>,
+        resolve_info: &trustfall::provider::ResolveInfo,
+    ) -> trustfall::provider::ContextOutcomeIterator<'v, V, bool> {
+        trace!(?type_name, ?coerce_to_type, "Trying to coerce");
+        let (adapter_name, coerce_to_type) = coerce_to_type.split_once(ADAPTER_SEP).unwrap();
+
+        match adapter_name {
+            "Plaixt" => {
+                let contexts = contexts.collect::<Vec<_>>();
+
+                let properties = self.plaixt.resolve_coercion(
+                    Box::new(
+                        contexts
+                            .clone()
+                            .into_iter()
+                            .map(|v| v.map(&mut |v: V| v.into_vertex().unwrap())),
+                    ),
+                    type_name,
+                    &Arc::from(coerce_to_type),
+                    resolve_info,
+                );
+
+                Box::new(
+                    properties
+                        .into_iter()
+                        .zip(contexts)
+                        .map(|((_ctx, val), og_ctx)| (og_ctx, val)),
+                )
+            }
+            _ => unreachable!(),
+        }
+    }
+}
+
 pub(crate) struct PlaixtAdapter {
     pub(crate) records: Vec<Record>,
 }
@@ -119,6 +404,7 @@ impl<'a> Adapter<'a> for PlaixtAdapter {
         _parameters: &trustfall::provider::EdgeParameters,
         _resolve_info: &trustfall::provider::ResolveInfo,
     ) -> trustfall::provider::VertexIterator<'a, Self::Vertex> {
+        trace!(?edge_name, "Resolving start vertex");
         match edge_name.as_ref() {
             "RecordsAll" => Box::new(self.records.clone().into_iter().map(PlaixtVertex::Record)),
             _ => unreachable!(),

query

@@ -1,6 +1,6 @@
 {
-    RecordsAll {
+    Plaixt__RecordsAll {
-        ... on changelog {
+        ... on Plaixt__changelog {
             at @output
             kind @output
             fields {