obi1kenobi · obi1kenobi · Jun 11, 2024 · Jun 11, 2024
diff --git a/trustfall_core/src/frontend/mod.rs b/trustfall_core/src/frontend/mod.rs
@@ -731,10 +731,7 @@ where
     // TODO: fixme, temporary hack to avoid changing the IRQueryComponent struct
     let hacked_outputs = component_outputs
         .into_iter()
-        .filter_map(|(k, v)| match v {
-            FieldRef::ContextField(c) => Some((k, c)),
-            FieldRef::FoldSpecificField(_) => None,
-        })
+        .filter(|(k, v)| !matches!(&v, FieldRef::FoldSpecificField(..)))
         .collect();
 
     Ok(IRQueryComponent {

diff --git a/trustfall_core/src/interpreter/execution.rs b/trustfall_core/src/interpreter/execution.rs
@@ -193,29 +193,37 @@ fn construct_outputs<'query, AdapterT: Adapter<'query>>(
     let mut output_iterator = iterator;
 
     for output_name in output_names.iter() {
-        let context_field = &root_component.outputs[output_name];
-        let vertex_id = context_field.vertex_id;
+        let field_ref = &root_component.outputs[output_name];
 
-        let moved_iterator = Box::new(output_iterator.map(move |context| {
-            let new_vertex = context.vertices[&vertex_id].clone();
-            context.move_to_vertex(new_vertex)
-        }));
+        match field_ref {
+            FieldRef::ContextField(context_field) => {
+                let vertex_id = context_field.vertex_id;
 
-        let resolve_info = ResolveInfo::new(query, vertex_id, true);
+                let moved_iterator = Box::new(output_iterator.map(move |context| {
+                    let new_vertex = context.vertices[&vertex_id].clone();
+                    context.move_to_vertex(new_vertex)
+                }));
 
-        let type_name = &root_component.vertices[&vertex_id].type_name;
-        let field_data_iterator = adapter.resolve_property(
-            moved_iterator,
-            type_name,
-            &context_field.field_name,
-            &resolve_info,
-        );
-        query = resolve_info.into_inner();
+                let resolve_info = ResolveInfo::new(query, vertex_id, true);
 
-        output_iterator = Box::new(field_data_iterator.map(|(mut context, value)| {
-            context.values.push(value);
-            context
-        }));
+                let type_name = &root_component.vertices[&vertex_id].type_name;
+                let field_data_iterator = adapter.resolve_property(
+                    moved_iterator,
+                    type_name,
+                    &context_field.field_name,
+                    &resolve_info,
+                );
+                query = resolve_info.into_inner();
+
+                output_iterator = Box::new(field_data_iterator.map(|(mut context, value)| {
+                    context.values.push(value);
+                    context
+                }));
+            }
+            FieldRef::FoldSpecificField(_) => {
+                unreachable!("found fold-specific field in component outputs: {root_component:#?}")
+            }
+        }
     }
     let expected_output_names: BTreeSet<_> = query.indexed_query.outputs.keys().cloned().collect();
     carrier.query = Some(query);
@@ -648,27 +656,36 @@ mismatch on whether the fold below {expanding_from_vid:?} was inside an `@option
                 // This is a slimmed-down version of computing a context field:
                 // - it does not restore the prior active vertex after getting each value
                 // - it already knows that the context field is guaranteed to exist
-                let context_field = &fold.component.outputs[output_name.as_ref()];
-                let vertex_id = context_field.vertex_id;
-                let moved_iterator = Box::new(output_iterator.map(move |context| {
-                    let new_vertex = context.vertices[&vertex_id].clone();
-                    context.move_to_vertex(new_vertex)
-                }));
-
-                let query = cloned_carrier.query.take().expect("query was not returned");
-                let resolve_info = ResolveInfo::new(query, vertex_id, true);
-                let field_data_iterator = cloned_adapter.resolve_property(
-                    moved_iterator,
-                    &fold.component.vertices[&vertex_id].type_name,
-                    &context_field.field_name,
-                    &resolve_info,
-                );
-                cloned_carrier.query = Some(resolve_info.into_inner());
-
-                output_iterator = Box::new(field_data_iterator.map(|(mut context, value)| {
-                    context.values.push(value);
-                    context
-                }));
+                let field_ref = &fold.component.outputs[output_name.as_ref()];
+
+                match field_ref {
+                    FieldRef::ContextField(context_field) => {
+                        let vertex_id = context_field.vertex_id;
+                        let moved_iterator = Box::new(output_iterator.map(move |context| {
+                            let new_vertex = context.vertices[&vertex_id].clone();
+                            context.move_to_vertex(new_vertex)
+                        }));
+
+                        let query = cloned_carrier.query.take().expect("query was not returned");
+                        let resolve_info = ResolveInfo::new(query, vertex_id, true);
+                        let field_data_iterator = cloned_adapter.resolve_property(
+                            moved_iterator,
+                            &fold.component.vertices[&vertex_id].type_name,
+                            &context_field.field_name,
+                            &resolve_info,
+                        );
+                        cloned_carrier.query = Some(resolve_info.into_inner());
+
+                        output_iterator =
+                            Box::new(field_data_iterator.map(|(mut context, value)| {
+                                context.values.push(value);
+                                context
+                            }));
+                    }
+                    FieldRef::FoldSpecificField(_) => {
+                        unreachable!("found fold-specific field in component outputs: {fold:#?}")
+                    }
+                }
             }
 
             for mut folded_context in output_iterator {

diff --git a/trustfall_core/src/interpreter/helpers/correctness.rs b/trustfall_core/src/interpreter/helpers/correctness.rs
@@ -3,8 +3,8 @@ use std::{collections::BTreeMap, fmt::Debug, num::NonZeroUsize, sync::Arc};
 use crate::{
     interpreter::{Adapter, DataContext, InterpretedQuery, ResolveEdgeInfo, ResolveInfo},
     ir::{
-        ContextField, EdgeParameters, Eid, FieldValue, IREdge, IRQuery, IRQueryComponent, IRVertex,
-        TransparentValue, Type, Vid,
+        ContextField, EdgeParameters, Eid, FieldRef, FieldValue, IREdge, IRQuery, IRQueryComponent,
+        IRVertex, TransparentValue, Type, Vid,
     },
     schema::{Schema, SchemaAdapter},
     TryIntoStruct,
@@ -166,11 +166,11 @@ fn make_resolve_info_for_property_check(
             edges: Default::default(),
             folds: Default::default(),
             outputs: btreemap! {
-                property_name.clone() => ContextField {
+                property_name.clone() => FieldRef::ContextField(ContextField {
                     vertex_id: vid,
                     field_name: property_name.clone(),
                     field_type: Type::parse(property_type).expect("not a valid type"),
-                }
+                })
             },
         }),
         variables: Default::default(),
@@ -324,11 +324,11 @@ fn make_resolve_edge_info_for_edge_check(
             },
             folds: Default::default(),
             outputs: btreemap! {
-                property_name.clone() => ContextField {
+                property_name.clone() => FieldRef::ContextField(ContextField {
                     vertex_id: vid,
                     field_name: property_name,
                     field_type: Type::parse("String!").expect("not a valid type"),
-                }
+                })
             },
         }),
         variables: Default::default(),
@@ -493,11 +493,11 @@ fn make_resolve_info_for_type_coercion(
             edges: Default::default(),
             folds: Default::default(),
             outputs: btreemap! {
-                typename_property.clone() => ContextField {
+                typename_property.clone() => FieldRef::ContextField(ContextField {
                     vertex_id: vid,
                     field_name: typename_property.clone(),
                     field_type: Type::parse("String!").expect("not a valid type"),
-                }
+                })
             },
         }),
         variables: Default::default(),

diff --git a/trustfall_core/src/interpreter/hints/vertex_info.rs b/trustfall_core/src/interpreter/hints/vertex_info.rs
@@ -165,8 +165,8 @@ impl<T: InternalVertexInfo + super::sealed::__Sealed> VertexInfo for T {
         let properties = current_component
             .outputs
             .values()
-            .filter(|c| c.vertex_id == current_vertex.vid)
-            .map(|c| RequiredProperty::new(c.field_name.clone()));
+            .filter(|c| c.defined_at() == current_vertex.vid)
+            .map(|c| RequiredProperty::new(c.field_name_arc()));
 
         let properties = properties.chain(
             current_vertex

diff --git a/trustfall_core/src/ir/indexed.rs b/trustfall_core/src/ir/indexed.rs
@@ -154,7 +154,7 @@ fn add_data_from_component(
     }
 
     for (output_name, field) in component.outputs.iter() {
-        let output_vid = field.vertex_id;
+        let output_vid = field.defined_at();
 
         // the output must be from a vertex in this component
         let output_component =
@@ -166,7 +166,7 @@ fn add_data_from_component(
         let output_name = output_name.clone();
         let output_type = get_output_type(
             output_vid,
-            &field.field_type,
+            field.field_type(),
             &component_optional_vertices,
             are_folds_optional,
         );

diff --git a/trustfall_core/src/ir/mod.rs b/trustfall_core/src/ir/mod.rs
@@ -118,7 +118,7 @@ pub struct IRQueryComponent {
     pub folds: BTreeMap<Eid, Arc<IRFold>>,
 
     #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
-    pub outputs: BTreeMap<Arc<str>, ContextField>,
+    pub outputs: BTreeMap<Arc<str>, FieldRef>,
 }
 
 /// Intermediate representation of a query
@@ -319,6 +319,13 @@ impl FieldRef {
         }
     }
 
+    pub fn field_name_arc(&self) -> Arc<str> {
+        match self {
+            FieldRef::ContextField(c) => c.field_name.clone(),
+            FieldRef::FoldSpecificField(f) => f.kind.field_name().into(),
+        }
+    }
+
     /// The vertex ID at which this reference is considered defined.
     pub fn defined_at(&self) -> Vid {
         match self {

diff --git a/trustfall_core/test_data/tests/execution_errors/both_missing_and_unused.ir.ron b/trustfall_core/test_data/tests/execution_errors/both_missing_and_unused.ir.ron
@@ -26,11 +26,11 @@ Ok(TestIRQuery(
         ),
       },
       outputs: {
-        "value": ContextField(
+        "value": ContextField(ContextField(
           vertex_id: Vid(1),
           field_name: "value",
           field_type: "Int",
-        ),
+        )),
       },
     ),
     variables: {

diff --git a/trustfall_core/test_data/tests/execution_errors/invalid_argument_types.ir.ron b/trustfall_core/test_data/tests/execution_errors/invalid_argument_types.ir.ron
@@ -33,11 +33,11 @@ Ok(TestIRQuery(
         ),
       },
       outputs: {
-        "value": ContextField(
+        "value": ContextField(ContextField(
           vertex_id: Vid(1),
           field_name: "value",
           field_type: "Int",
-        ),
+        )),
       },
     ),
     variables: {

diff --git a/trustfall_core/test_data/tests/execution_errors/invalid_null_argument_values.ir.ron b/trustfall_core/test_data/tests/execution_errors/invalid_null_argument_values.ir.ron
@@ -33,11 +33,11 @@ Ok(TestIRQuery(
         ),
       },
       outputs: {
-        "value": ContextField(
+        "value": ContextField(ContextField(
           vertex_id: Vid(1),
           field_name: "value",
           field_type: "Int",
-        ),
+        )),
       },
     ),
     variables: {

diff --git a/trustfall_core/test_data/tests/execution_errors/missing_argument.ir.ron b/trustfall_core/test_data/tests/execution_errors/missing_argument.ir.ron
@@ -26,11 +26,11 @@ Ok(TestIRQuery(
         ),
       },
       outputs: {
-        "value": ContextField(
+        "value": ContextField(ContextField(
           vertex_id: Vid(1),
           field_name: "value",
           field_type: "Int",
-        ),
+        )),
       },
     ),
     variables: {

diff --git a/trustfall_core/test_data/tests/execution_errors/missing_unused_and_invalid_arguments.ir.ron b/trustfall_core/test_data/tests/execution_errors/missing_unused_and_invalid_arguments.ir.ron
@@ -40,11 +40,11 @@ Ok(TestIRQuery(
         ),
       },
       outputs: {
-        "value": ContextField(
+        "value": ContextField(ContextField(
           vertex_id: Vid(1),
           field_name: "value",
           field_type: "Int",
-        ),
+        )),
       },
     ),
     variables: {

diff --git a/...e/test_data/tests/execution_errors/type_narrowing_prevents_list_with_null_argument.ir.ron b/...e/test_data/tests/execution_errors/type_narrowing_prevents_list_with_null_argument.ir.ron
@@ -27,16 +27,16 @@ Ok(TestIRQuery(
         ),
       },
       outputs: {
-        "intNonNullList": ContextField(
+        "intNonNullList": ContextField(ContextField(
           vertex_id: Vid(1),
           field_name: "intNonNullList",
           field_type: "[Int]!",
-        ),
-        "nonNullIntList": ContextField(
+        )),
+        "nonNullIntList": ContextField(ContextField(
           vertex_id: Vid(1),
           field_name: "nonNullIntList",
           field_type: "[Int!]",
-        ),
+        )),
       },
     ),
     variables: {

diff --git a/trustfall_core/test_data/tests/execution_errors/type_narrowing_prevents_null_argument.ir.ron b/trustfall_core/test_data/tests/execution_errors/type_narrowing_prevents_null_argument.ir.ron
@@ -27,16 +27,16 @@ Ok(TestIRQuery(
         ),
       },
       outputs: {
-        "intNonNullList": ContextField(
+        "intNonNullList": ContextField(ContextField(
           vertex_id: Vid(1),
           field_name: "intNonNullList",
           field_type: "[Int]!",
-        ),
-        "nonNullIntList": ContextField(
+        )),
+        "nonNullIntList": ContextField(ContextField(
           vertex_id: Vid(1),
           field_name: "nonNullIntList",
           field_type: "[Int!]",
-        ),
+        )),
       },
     ),
     variables: {

diff --git a/trustfall_core/test_data/tests/execution_errors/unused_argument.ir.ron b/trustfall_core/test_data/tests/execution_errors/unused_argument.ir.ron
@@ -17,16 +17,16 @@ Ok(TestIRQuery(
         ),
       },
       outputs: {
-        "name": ContextField(
+        "name": ContextField(ContextField(
           vertex_id: Vid(1),
           field_name: "name",
           field_type: "String",
-        ),
-        "value": ContextField(
+        )),
+        "value": ContextField(ContextField(
           vertex_id: Vid(1),
           field_name: "value",
           field_type: "Int",
-        ),
+        )),
       },
     ),
   ),

diff --git a/trustfall_core/test_data/tests/valid_queries/alias_and_prefix_output_names.ir.ron b/trustfall_core/test_data/tests/valid_queries/alias_and_prefix_output_names.ir.ron
@@ -23,16 +23,16 @@ Ok(TestIRQuery(
         ),
       },
       outputs: {
-        "succ_is_named": ContextField(
+        "succ_is_named": ContextField(ContextField(
           vertex_id: Vid(2),
           field_name: "name",
           field_type: "String",
-        ),
-        "zero": ContextField(
+        )),
+        "zero": ContextField(ContextField(
           vertex_id: Vid(1),
           field_name: "value",
           field_type: "Int",
-        ),
+        )),
       },
     ),
   ),