obi1kenobi · obi1kenobi · Jun 11, 2024 · Jun 11, 2024
diff --git a/trustfall_core/src/frontend/mod.rs b/trustfall_core/src/frontend/mod.rs
@@ -1309,7 +1309,7 @@ where
             };
 
             let prior_output_by_that_name =
-                fold_specific_outputs.insert(final_output_name.clone(), fold_specific_field.kind);
+                fold_specific_outputs.insert(final_output_name.clone(), field_ref.clone());
             if let Some(prior_output_kind) = prior_output_by_that_name {
                 errors.push(FrontendError::MultipleOutputsWithSameName(DuplicatedNamesConflict {
                     duplicates: btreemap! {
@@ -1324,9 +1324,7 @@ where
         for tag_directive in &transform_group.tag {
             let tag_name = tag_directive.name.as_ref().map(|x| x.as_ref());
             if let Some(tag_name) = tag_name {
-                let field = FieldRef::FoldSpecificField(fold_specific_field.clone());
-
-                if let Err(e) = tags.register_tag(tag_name, field, component_path) {
+                if let Err(e) = tags.register_tag(tag_name, field_ref.clone(), component_path) {
                     errors.push(FrontendError::MultipleTagsWithSameName(tag_name.to_string()));
                 }
             } else {

diff --git a/trustfall_core/src/interpreter/execution.rs b/trustfall_core/src/interpreter/execution.rs
@@ -500,33 +500,42 @@ fn compute_fold<'query, AdapterT: Adapter<'query> + 'query>(
     //
     // For example, if `@filter(op: ">", value: ["$ten"])` is our only filter on the count
     // of the fold, we can stop computing the rest of the fold after seeing we have 11 elements.
-    let min_fold_size =
-        if let Some(min_fold_size) = get_min_fold_count_limit(carrier, fold.as_ref()) {
-            let no_outputs_in_fold = fold.component.outputs.is_empty();
-            let has_output_on_fold_count =
-                fold.fold_specific_outputs.values().any(|x| *x == FoldSpecificFieldKind::Count);
-            let has_tag_on_fold_count = parent_component.vertices.values().any(|vertex| {
-                vertex.filters.iter().any(|filter| {
-                    let Some(Argument::Tag(FieldRef::FoldSpecificField(tagged_fold_count))) =
-                        filter.right()
-                    else {
-                        return false;
-                    };
+    let min_fold_size = if let Some(min_fold_size) =
+        get_min_fold_count_limit(carrier, fold.as_ref())
+    {
+        let no_outputs_in_fold = fold.component.outputs.is_empty();
+        let has_output_on_fold_count = fold.fold_specific_outputs.values().any(|x| {
+            x.refers_to_fold_specific_field()
+                .is_some_and(|f| f.kind == FoldSpecificFieldKind::Count)
+        });
+        let has_tag_on_fold_count = parent_component.vertices.values().any(|vertex| {
+            // TODO: If we allow referencing the tagged values located inside a fold in filters
+            //       outside that fold, then this logic will be buggy (not conservative enough).
+            //       It currently assumes that use of tags on `@fold @transform(op: "count")`
+            //       can only happen within the direct parent component of that fold.
+            vertex.filters.iter().any(|filter| {
+                let Some(Argument::Tag(field_ref)) = filter.right() else {
+                    return false;
+                };
 
-                    tagged_fold_count.fold_root_vid == fold.to_vid
-                        && tagged_fold_count.fold_eid == fold.eid
-                        && tagged_fold_count.kind == FoldSpecificFieldKind::Count
-                })
-            });
+                let Some(fold_specific_field) = field_ref.refers_to_fold_specific_field() else {
+                    return false;
+                };
 
-            if no_outputs_in_fold && !has_output_on_fold_count && !has_tag_on_fold_count {
-                Some(min_fold_size)
-            } else {
-                None
-            }
+                fold_specific_field.fold_root_vid == fold.to_vid
+                    && fold_specific_field.fold_eid == fold.eid
+                    && fold_specific_field.kind == FoldSpecificFieldKind::Count
+            })
+        });
+
+        if no_outputs_in_fold && !has_output_on_fold_count && !has_tag_on_fold_count {
+            Some(min_fold_size)
         } else {
             None
-        };
+        }
+    } else {
+        None
+    };
 
     let moved_fold = fold.clone();
     let folded_iterator = edge_iterator.filter_map(move |(mut context, neighbors)| {
@@ -599,14 +608,19 @@ mismatch on whether the fold below {expanding_from_vid:?} was inside an `@option
         );
 
         // Add any fold-specific field outputs to the context's folded values.
-        for (output_name, fold_specific_field) in &fold.fold_specific_outputs {
+        for (output_name, field_ref) in &fold.fold_specific_outputs {
             // If the @fold is inside an @optional that doesn't exist,
             // its outputs should be `null` rather than empty lists (the usual for empty folds).
             // Transformed outputs should also be `null` rather than their usual transformed defaults.
-            let value = fold_elements.as_ref().map(|elements| match fold_specific_field {
-                FoldSpecificFieldKind::Count => {
-                    ValueOrVec::Value(FieldValue::Uint64(elements.len() as u64))
-                }
+            let value = fold_elements.as_ref().map(|elements| match field_ref {
+                FieldRef::FoldSpecificField(field) => match field.kind {
+                    FoldSpecificFieldKind::Count => {
+                        ValueOrVec::Value(FieldValue::Uint64(elements.len() as u64))
+                    }
+                },
+                FieldRef::ContextField(_) => unreachable!(
+                    "found ContextField inside a fold's fold-specific outputs: {fold:#?}"
+                ),
             });
             ctx.folded_values
                 .insert_or_error((fold_eid, output_name.clone()), value)

diff --git a/trustfall_core/src/ir/mod.rs b/trustfall_core/src/ir/mod.rs
@@ -210,8 +210,12 @@ pub struct IRFold {
     #[serde(default, skip_serializing_if = "Vec::is_empty")]
     pub imported_tags: Vec<FieldRef>,
 
+    /// Outputs from this fold that are derived from fold-specific fields.
+    ///
+    /// All [`FieldRef`] values in the map are guaranteed to have
+    /// `[FieldRef].refers_to_fold_specific_field().is_some() == true`.
     #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
-    pub fold_specific_outputs: BTreeMap<Arc<str>, FoldSpecificFieldKind>,
+    pub fold_specific_outputs: BTreeMap<Arc<str>, FieldRef>,
 
     #[serde(default, skip_serializing_if = "Vec::is_empty")]
     pub post_filters: Vec<Operation<FoldSpecificFieldKind, Argument>>,
@@ -326,6 +330,13 @@ impl FieldRef {
             FieldRef::FoldSpecificField(f) => f.fold_root_vid,
         }
     }
+
+    pub fn refers_to_fold_specific_field(&self) -> Option<&FoldSpecificField> {
+        match self {
+            FieldRef::ContextField(_) => None,
+            FieldRef::FoldSpecificField(fold_specific) => Some(fold_specific),
+        }
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]

diff --git a/trustfall_core/test_data/tests/valid_queries/fold_count_filter_with_inner_filter.ir.ron b/trustfall_core/test_data/tests/valid_queries/fold_count_filter_with_inner_filter.ir.ron
@@ -49,7 +49,11 @@ Ok(TestIRQuery(
             },
           ),
           fold_specific_outputs: {
-            "primeFactorcount": Count,
+            "primeFactorcount": FoldSpecificField(FoldSpecificField(
+              fold_eid: Eid(1),
+              fold_root_vid: Vid(2),
+              kind: Count,
+            )),
           },
           post_filters: [
             Equals(Count, Variable(VariableRef(

diff --git a/trustfall_core/test_data/tests/valid_queries/fold_count_filter_with_inner_filter.trace.ron b/trustfall_core/test_data/tests/valid_queries/fold_count_filter_with_inner_filter.trace.ron
@@ -473,7 +473,11 @@ TestInterpreterOutputTrace(
               },
             ),
             fold_specific_outputs: {
-              "primeFactorcount": Count,
+              "primeFactorcount": FoldSpecificField(FoldSpecificField(
+                fold_eid: Eid(1),
+                fold_root_vid: Vid(2),
+                kind: Count,
+              )),
             },
             post_filters: [
               Equals(Count, Variable(VariableRef(

diff --git a/trustfall_core/test_data/tests/valid_queries/fold_count_prefixed_output_name.ir.ron b/trustfall_core/test_data/tests/valid_queries/fold_count_prefixed_output_name.ir.ron
@@ -33,7 +33,11 @@ Ok(TestIRQuery(
             },
           ),
           fold_specific_outputs: {
-            "factor_count": Count,
+            "factor_count": FoldSpecificField(FoldSpecificField(
+              fold_eid: Eid(1),
+              fold_root_vid: Vid(2),
+              kind: Count,
+            )),
           },
         ),
       },

diff --git a/trustfall_core/test_data/tests/valid_queries/fold_count_prefixed_output_name.trace.ron b/trustfall_core/test_data/tests/valid_queries/fold_count_prefixed_output_name.trace.ron
@@ -171,7 +171,11 @@ TestInterpreterOutputTrace(
               },
             ),
             fold_specific_outputs: {
-              "factor_count": Count,
+              "factor_count": FoldSpecificField(FoldSpecificField(
+                fold_eid: Eid(1),
+                fold_root_vid: Vid(2),
+                kind: Count,
+              )),
             },
           ),
         },

diff --git a/trustfall_core/test_data/tests/valid_queries/fold_count_tag_on_nonexistent_optional.ir.ron b/trustfall_core/test_data/tests/valid_queries/fold_count_tag_on_nonexistent_optional.ir.ron
@@ -74,7 +74,11 @@ Ok(TestIRQuery(
             },
           ),
           fold_specific_outputs: {
-            "count": Count,
+            "count": FoldSpecificField(FoldSpecificField(
+              fold_eid: Eid(3),
+              fold_root_vid: Vid(4),
+              kind: Count,
+            )),
           },
         ),
       },

diff --git a/...tfall_core/test_data/tests/valid_queries/fold_count_tag_on_nonexistent_optional.trace.ron b/...tfall_core/test_data/tests/valid_queries/fold_count_tag_on_nonexistent_optional.trace.ron
@@ -521,7 +521,11 @@ TestInterpreterOutputTrace(
               },
             ),
             fold_specific_outputs: {
-              "count": Count,
+              "count": FoldSpecificField(FoldSpecificField(
+                fold_eid: Eid(3),
+                fold_root_vid: Vid(4),
+                kind: Count,
+              )),
             },
           ),
         },

diff --git a/trustfall_core/test_data/tests/valid_queries/fold_count_tag_then_filter_on_fold.ir.ron b/trustfall_core/test_data/tests/valid_queries/fold_count_tag_then_filter_on_fold.ir.ron
@@ -41,7 +41,11 @@ Ok(TestIRQuery(
             },
           ),
           fold_specific_outputs: {
-            "predecessorcount": Count,
+            "predecessorcount": FoldSpecificField(FoldSpecificField(
+              fold_eid: Eid(2),
+              fold_root_vid: Vid(3),
+              kind: Count,
+            )),
           },
           post_filters: [
             Equals(Count, Tag(FoldSpecificField(FoldSpecificField(

diff --git a/trustfall_core/test_data/tests/valid_queries/fold_count_tag_then_filter_on_fold.trace.ron b/trustfall_core/test_data/tests/valid_queries/fold_count_tag_then_filter_on_fold.trace.ron
@@ -196,7 +196,11 @@ TestInterpreterOutputTrace(
               },
             ),
             fold_specific_outputs: {
-              "predecessorcount": Count,
+              "predecessorcount": FoldSpecificField(FoldSpecificField(
+                fold_eid: Eid(2),
+                fold_root_vid: Vid(3),
+                kind: Count,
+              )),
             },
             post_filters: [
               Equals(Count, Tag(FoldSpecificField(FoldSpecificField(

diff --git a/...all_core/test_data/tests/valid_queries/fold_with_no_inner_output_with_count_output.ir.ron b/...all_core/test_data/tests/valid_queries/fold_with_no_inner_output_with_count_output.ir.ron
@@ -33,7 +33,11 @@ Ok(TestIRQuery(
             },
           ),
           fold_specific_outputs: {
-            "primeFactorcount": Count,
+            "primeFactorcount": FoldSpecificField(FoldSpecificField(
+              fold_eid: Eid(1),
+              fold_root_vid: Vid(2),
+              kind: Count,
+            )),
           },
           post_filters: [
             GreaterThanOrEqual(Count, Variable(VariableRef(

diff --git a/..._core/test_data/tests/valid_queries/fold_with_no_inner_output_with_count_output.trace.ron b/..._core/test_data/tests/valid_queries/fold_with_no_inner_output_with_count_output.trace.ron
@@ -304,7 +304,11 @@ TestInterpreterOutputTrace(
               },
             ),
             fold_specific_outputs: {
-              "primeFactorcount": Count,
+              "primeFactorcount": FoldSpecificField(FoldSpecificField(
+                fold_eid: Eid(1),
+                fold_root_vid: Vid(2),
+                kind: Count,
+              )),
             },
             post_filters: [
               GreaterThanOrEqual(Count, Variable(VariableRef(

diff --git a/trustfall_core/test_data/tests/valid_queries/multiple_fold_count_outputs.ir.ron b/trustfall_core/test_data/tests/valid_queries/multiple_fold_count_outputs.ir.ron
@@ -40,7 +40,11 @@ Ok(TestIRQuery(
             },
           ),
           fold_specific_outputs: {
-            "primeFactorcount": Count,
+            "primeFactorcount": FoldSpecificField(FoldSpecificField(
+              fold_eid: Eid(1),
+              fold_root_vid: Vid(2),
+              kind: Count,
+            )),
           },
         ),
         Eid(2): IRFold(
@@ -70,7 +74,11 @@ Ok(TestIRQuery(
             },
           ),
           fold_specific_outputs: {
-            "multiplecount": Count,
+            "multiplecount": FoldSpecificField(FoldSpecificField(
+              fold_eid: Eid(2),
+              fold_root_vid: Vid(3),
+              kind: Count,
+            )),
           },
         ),
       },

diff --git a/trustfall_core/test_data/tests/valid_queries/multiple_fold_count_outputs.trace.ron b/trustfall_core/test_data/tests/valid_queries/multiple_fold_count_outputs.trace.ron
@@ -722,7 +722,11 @@ TestInterpreterOutputTrace(
               },
             ),
             fold_specific_outputs: {
-              "primeFactorcount": Count,
+              "primeFactorcount": FoldSpecificField(FoldSpecificField(
+                fold_eid: Eid(1),
+                fold_root_vid: Vid(2),
+                kind: Count,
+              )),
             },
           ),
           Eid(2): IRFold(
@@ -752,7 +756,11 @@ TestInterpreterOutputTrace(
               },
             ),
             fold_specific_outputs: {
-              "multiplecount": Count,
+              "multiplecount": FoldSpecificField(FoldSpecificField(
+                fold_eid: Eid(2),
+                fold_root_vid: Vid(3),
+                kind: Count,
+              )),
             },
           ),
         },

diff --git a/trustfall_core/test_data/tests/valid_queries/nested_fold_count_output.ir.ron b/trustfall_core/test_data/tests/valid_queries/nested_fold_count_output.ir.ron
@@ -59,7 +59,11 @@ Ok(TestIRQuery(
                   },
                 ),
                 fold_specific_outputs: {
-                  "multiplecount": Count,
+                  "multiplecount": FoldSpecificField(FoldSpecificField(
+                    fold_eid: Eid(2),
+                    fold_root_vid: Vid(3),
+                    kind: Count,
+                  )),
                 },
               ),
             },
@@ -72,7 +76,11 @@ Ok(TestIRQuery(
             },
           ),
           fold_specific_outputs: {
-            "primeFactorcount": Count,
+            "primeFactorcount": FoldSpecificField(FoldSpecificField(
+              fold_eid: Eid(1),
+              fold_root_vid: Vid(2),
+              kind: Count,
+            )),
           },
         ),
       },

diff --git a/trustfall_core/test_data/tests/valid_queries/nested_fold_count_output.trace.ron b/trustfall_core/test_data/tests/valid_queries/nested_fold_count_output.trace.ron
@@ -980,7 +980,11 @@ TestInterpreterOutputTrace(
                     },
                   ),
                   fold_specific_outputs: {
-                    "multiplecount": Count,
+                    "multiplecount": FoldSpecificField(FoldSpecificField(
+                      fold_eid: Eid(2),
+                      fold_root_vid: Vid(3),
+                      kind: Count,
+                    )),
                   },
                 ),
               },
@@ -993,7 +997,11 @@ TestInterpreterOutputTrace(
               },
             ),
             fold_specific_outputs: {
-              "primeFactorcount": Count,
+              "primeFactorcount": FoldSpecificField(FoldSpecificField(
+                fold_eid: Eid(1),
+                fold_root_vid: Vid(2),
+                kind: Count,
+              )),
             },
           ),
         },

diff --git a/trustfall_core/test_data/tests/valid_queries/nonexistent_optional_with_nested_fold.ir.ron b/trustfall_core/test_data/tests/valid_queries/nonexistent_optional_with_nested_fold.ir.ron
@@ -46,7 +46,11 @@ Ok(TestIRQuery(
             },
           ),
           fold_specific_outputs: {
-            "successor_counts": Count,
+            "successor_counts": FoldSpecificField(FoldSpecificField(
+              fold_eid: Eid(2),
+              fold_root_vid: Vid(3),
+              kind: Count,
+            )),
           },
         ),
       },

diff --git a/trustfall_core/test_data/tests/valid_queries/nonexistent_optional_with_nested_fold.trace.ron b/trustfall_core/test_data/tests/valid_queries/nonexistent_optional_with_nested_fold.trace.ron
@@ -299,7 +299,11 @@ TestInterpreterOutputTrace(
               },
             ),
             fold_specific_outputs: {
-              "successor_counts": Count,
+              "successor_counts": FoldSpecificField(FoldSpecificField(
+                fold_eid: Eid(2),
+                fold_root_vid: Vid(3),
+                kind: Count,
+              )),
             },
           ),
         },