Implement execution for @transform applied to filters' left-hand operand. (#626)

obi1kenobi · web-flow · commit 56521dcce880 · 2024-06-15T13:43:05.000-04:00
diff --git a/trustfall_core/src/interpreter/execution.rs b/trustfall_core/src/interpreter/execution.rs
@@ -8,15 +8,17 @@ use crate::{
     ir::{
         Argument, ContextField, EdgeParameters, Eid, FieldRef, FieldValue, FoldSpecificFieldKind,
         IREdge, IRFold, IRQueryComponent, IRVertex, IndexedQuery, LocalField, Operation,
-        OperationSubject, Recursive, Vid,
+        OperationSubject, Recursive, TransformBase, Vid,
     },
     util::BTreeMapTryInsertExt,
 };
 
 use super::{
-    error::QueryArgumentsError, filtering::apply_filter, Adapter, AsVertex, ContextIterator,
-    ContextOutcomeIterator, DataContext, InterpretedQuery, ResolveEdgeInfo, ResolveInfo,
-    TaggedValue, ValueOrVec, VertexIterator,
+    error::QueryArgumentsError,
+    filtering::apply_filter,
+    transformation::{apply_transforms, push_transform_argument_tag_values_onto_stack},
+    Adapter, AsVertex, ContextIterator, ContextOutcomeIterator, DataContext, InterpretedQuery,
+    ResolveEdgeInfo, ResolveInfo, TaggedValue, ValueOrVec, VertexIterator,
 };
 
 #[derive(Debug, Clone)]
@@ -621,25 +623,15 @@ fn compute_fold<'query, AdapterT: Adapter<'query> + 'query>(
     let mut post_filtered_iterator: ContextIterator<'query, AdapterT::Vertex> =
         Box::new(folded_iterator);
     for post_fold_filter in fold.post_filters.iter() {
-        let left = post_fold_filter.left();
-        match left {
-            OperationSubject::FoldSpecificField(fold_specific_field) => {
-                let remapped_operation = post_fold_filter.map(|_| fold_specific_field.kind, |x| x);
-                post_filtered_iterator = apply_fold_specific_filter(
-                    adapter.as_ref(),
-                    carrier,
-                    parent_component,
-                    fold.as_ref(),
-                    expanding_from.vid,
-                    &remapped_operation,
-                    post_filtered_iterator,
-                );
-            }
-            OperationSubject::TransformedField(_) => todo!(),
-            OperationSubject::LocalField(_) => {
-                unreachable!("unexpectedly found a fold post-filtering step that references a LocalField: {fold:#?}");
-            }
-        }
+        post_filtered_iterator = apply_fold_specific_filter(
+            adapter.as_ref(),
+            carrier,
+            parent_component,
+            fold.as_ref(),
+            expanding_from.vid,
+            post_fold_filter,
+            post_filtered_iterator,
+        );
     }
 
     // Compute the outputs from this fold.
@@ -809,7 +801,63 @@ fn apply_filter_with_non_folded_field_subject<'query, AdapterT: Adapter<'query>>
             filter.map_left(|_| field),
             iterator,
         ),
-        OperationSubject::TransformedField(_) => todo!(),
+        OperationSubject::TransformedField(transformed) => {
+            let prepped_iterator = push_transform_argument_tag_values_onto_stack(
+                adapter,
+                carrier,
+                component,
+                current_vid,
+                &transformed.value.transforms,
+                iterator,
+            );
+
+            let query_variables =
+                Arc::clone(&carrier.query.as_ref().expect("query was not returned").arguments);
+            let transform_data = Arc::clone(&transformed.value);
+
+            match &transformed.value.base {
+                TransformBase::ContextField(field) => {
+                    assert_eq!(current_vid, field.vertex_id, "filter left-hand side was a transformed field from a different vertex: {current_vid:?} {filter:?}");
+                    let local_field = LocalField {
+                        field_name: field.field_name.clone(),
+                        field_type: field.field_type.clone(),
+                    };
+
+                    let filter_input_iterator = Box::new(
+                        compute_local_field_with_separate_value(
+                            adapter,
+                            carrier,
+                            component,
+                            current_vid,
+                            &local_field,
+                            prepped_iterator,
+                        )
+                        .map(move |(mut ctx, mut value)| {
+                            value = apply_transforms(
+                                &transform_data,
+                                &query_variables,
+                                &mut ctx.values,
+                                value,
+                            );
+                            ctx.values.push(value);
+                            ctx
+                        }),
+                    );
+
+                    apply_filter(
+                        adapter,
+                        carrier,
+                        component,
+                        current_vid,
+                        &filter.map(|_| (), |r| r),
+                        filter_input_iterator,
+                    )
+                }
+                TransformBase::FoldSpecificField(..) => unreachable!(
+                    "illegal filter over fold-specific field passed to this function: {filter:?}"
+                ),
+            }
+        }
         OperationSubject::FoldSpecificField(..) => unreachable!(
             "illegal filter over fold-specific field passed to this function: {filter:?}"
         ),
@@ -844,27 +892,70 @@ fn apply_fold_specific_filter<'query, AdapterT: Adapter<'query>>(
     component: &IRQueryComponent,
     fold: &IRFold,
     current_vid: Vid,
-    filter: &Operation<FoldSpecificFieldKind, &Argument>,
+    filter: &Operation<OperationSubject, Argument>,
     iterator: ContextIterator<'query, AdapterT::Vertex>,
 ) -> ContextIterator<'query, AdapterT::Vertex> {
-    let fold_specific_field = filter.left();
-    let field_iterator = Box::new(compute_fold_specific_field_with_separate_value(fold.eid, fold_specific_field, iterator).map(|(mut ctx, tagged_value)| {
-        let value = match tagged_value {
-            TaggedValue::Some(value) => value,
-            TaggedValue::NonexistentOptional => {
-                unreachable!("while applying fold-specific filter, the @fold turned out to not exist: {ctx:?}")
+    let left = filter.left();
+    let (fold_specific_field, transform_data) = match left {
+        OperationSubject::FoldSpecificField(field) => (field, None),
+        OperationSubject::TransformedField(transformed) => match &transformed.value.base {
+            TransformBase::FoldSpecificField(field) => (field, Some(&transformed.value)),
+            TransformBase::ContextField(_) => {
+                unreachable!("post-fold filter does not refer to a fold-specific field: {left:?}")
             }
-        };
-        ctx.values.push(value);
-        ctx
-    }));
+        },
+        OperationSubject::LocalField(_) => {
+            unreachable!("post-fold filter does not refer to a fold-specific field: {left:?}")
+        }
+    };
+
+    let field_iterator: ContextIterator<'query, AdapterT::Vertex> = if let Some(transform_data) =
+        transform_data
+    {
+        let prepped_iterator = push_transform_argument_tag_values_onto_stack(
+            adapter,
+            carrier,
+            component,
+            current_vid,
+            &transform_data.transforms,
+            iterator,
+        );
+
+        let query_variables =
+            Arc::clone(&carrier.query.as_ref().expect("query was not returned").arguments);
+        let transform_data = Arc::clone(transform_data);
+        Box::new(compute_fold_specific_field_with_separate_value(fold.eid, &fold_specific_field.kind, prepped_iterator).map(move |(mut ctx, tagged_value)| {
+            let mut value = match tagged_value {
+                TaggedValue::Some(value) => value,
+                TaggedValue::NonexistentOptional => {
+                    unreachable!("while applying fold-specific filter, the @fold turned out to not exist: {ctx:?}")
+                }
+            };
+
+            value = apply_transforms(&transform_data, &query_variables, &mut ctx.values, value);
+
+            ctx.values.push(value);
+            ctx
+        }))
+    } else {
+        Box::new(compute_fold_specific_field_with_separate_value(fold.eid, &fold_specific_field.kind, iterator).map(|(mut ctx, tagged_value)| {
+            let value = match tagged_value {
+                TaggedValue::Some(value) => value,
+                TaggedValue::NonexistentOptional => {
+                    unreachable!("while applying fold-specific filter, the @fold turned out to not exist: {ctx:?}")
+                }
+            };
+            ctx.values.push(value);
+            ctx
+        }))
+    };
 
     apply_filter(
         adapter,
         carrier,
         component,
         current_vid,
-        &filter.map(|_| (), |r| *r),
+        &filter.map(|_| (), |r| r),
         field_iterator,
     )
 }
diff --git a/trustfall_core/src/interpreter/filtering.rs b/trustfall_core/src/interpreter/filtering.rs
@@ -2,14 +2,11 @@ use std::{fmt::Debug, mem};
 
 use regex::Regex;
 
-use crate::ir::{Argument, FieldRef, FieldValue, IRQueryComponent, LocalField, Operation, Vid};
+use crate::ir::{Argument, FieldValue, IRQueryComponent, Operation, Vid};
 
 use super::{
-    execution::{
-        compute_context_field_with_separate_value, compute_fold_specific_field_with_separate_value,
-        compute_local_field_with_separate_value, QueryCarrier,
-    },
-    Adapter, ContextIterator, ContextOutcomeIterator, TaggedValue,
+    execution::QueryCarrier, tags::compute_tag_with_separate_value, Adapter, ContextIterator,
+    ContextOutcomeIterator, TaggedValue,
 };
 
 #[inline(always)]
@@ -281,56 +278,15 @@ pub(super) fn apply_filter<'query, AdapterT: Adapter<'query>>(
             let right_value = query_arguments[var.variable_name.as_ref()].to_owned();
             apply_filter_with_static_argument_value(filter, right_value, iterator)
         }
-        Some(Argument::Tag(FieldRef::ContextField(context_field))) => {
-            // TODO: Benchmark if it would be faster to duplicate the filtering code to special-case
-            //       the situation when the tag is always known to exist, so we don't have to unwrap
-            //       a TaggedValue enum, because we know it would be TaggedValue::Some.
-            let argument_value_iterator = if context_field.vertex_id == current_vid {
-                // This tag is from the vertex we're currently filtering. That means the field
-                // whose value we want to get is actually local, so there's no need to compute it
-                // using the more expensive approach we use for non-local fields.
-                let local_equivalent_field = LocalField {
-                    field_name: context_field.field_name.clone(),
-                    field_type: context_field.field_type.clone(),
-                };
-                Box::new(
-                    compute_local_field_with_separate_value(
-                        adapter,
-                        carrier,
-                        component,
-                        current_vid,
-                        &local_equivalent_field,
-                        iterator,
-                    )
-                    .map(|(ctx, value)| (ctx, TaggedValue::Some(value))),
-                )
-            } else {
-                compute_context_field_with_separate_value(
-                    adapter,
-                    carrier,
-                    component,
-                    context_field,
-                    iterator,
-                )
-            };
-            apply_filter_with_tagged_argument_value(filter, argument_value_iterator)
-        }
-        Some(Argument::Tag(field_ref @ FieldRef::FoldSpecificField(fold_field))) => {
-            let argument_value_iterator = if component.folds.contains_key(&fold_field.fold_eid) {
-                compute_fold_specific_field_with_separate_value(
-                    fold_field.fold_eid,
-                    &fold_field.kind,
-                    iterator,
-                )
-            } else {
-                // This value represents an imported tag value from an outer component.
-                // Grab its value from the context itself.
-                let cloned_ref = field_ref.clone();
-                Box::new(iterator.map(move |ctx| {
-                    let right_value = ctx.imported_tags[&cloned_ref].clone();
-                    (ctx, right_value)
-                }))
-            };
+        Some(Argument::Tag(field_ref)) => {
+            let argument_value_iterator = compute_tag_with_separate_value(
+                adapter,
+                carrier,
+                component,
+                current_vid,
+                field_ref,
+                iterator,
+            );
             apply_filter_with_tagged_argument_value(filter, argument_value_iterator)
         }
         Some(Argument::Tag(FieldRef::TransformedField(_))) => {
diff --git a/trustfall_core/src/interpreter/mod.rs b/trustfall_core/src/interpreter/mod.rs
@@ -17,7 +17,9 @@ mod filtering;
 pub mod helpers;
 mod hints;
 pub mod replay;
+mod tags;
 pub mod trace;
+mod transformation;
 
 pub use hints::{
     CandidateValue, DynamicallyResolvedValue, EdgeInfo, NeighborInfo, QueryInfo, Range,
diff --git a/trustfall_core/src/interpreter/tags.rs b/trustfall_core/src/interpreter/tags.rs
@@ -0,0 +1,71 @@
+use crate::ir::{FieldRef, IRQueryComponent, LocalField, Vid};
+
+use super::{
+    execution::{
+        compute_context_field_with_separate_value, compute_fold_specific_field_with_separate_value,
+        compute_local_field_with_separate_value, QueryCarrier,
+    },
+    Adapter, ContextIterator, DataContext, TaggedValue,
+};
+
+pub(super) fn compute_tag_with_separate_value<'query, AdapterT: Adapter<'query>>(
+    adapter: &AdapterT,
+    carrier: &mut QueryCarrier,
+    component: &IRQueryComponent,
+    current_vid: Vid,
+    field_ref: &FieldRef,
+    iterator: ContextIterator<'query, AdapterT::Vertex>,
+) -> Box<dyn Iterator<Item = (DataContext<AdapterT::Vertex>, TaggedValue)> + 'query> {
+    match field_ref {
+        FieldRef::ContextField(context_field) => {
+            // TODO: Benchmark if it would be faster to duplicate the code to special-case
+            //       the situation when the tag is always known to exist, so we don't have to unwrap
+            //       a TaggedValue enum, because we know it would be TaggedValue::Some.
+            if context_field.vertex_id == current_vid {
+                // This tag is from the vertex we're currently evaluating. That means the field
+                // whose value we want to get is actually local, so there's no need to compute it
+                // using the more expensive approach we use for non-local fields.
+                let local_equivalent_field = LocalField {
+                    field_name: context_field.field_name.clone(),
+                    field_type: context_field.field_type.clone(),
+                };
+                Box::new(
+                    compute_local_field_with_separate_value(
+                        adapter,
+                        carrier,
+                        component,
+                        current_vid,
+                        &local_equivalent_field,
+                        iterator,
+                    )
+                    .map(|(ctx, value)| (ctx, TaggedValue::Some(value))),
+                )
+            } else {
+                compute_context_field_with_separate_value(
+                    adapter,
+                    carrier,
+                    component,
+                    context_field,
+                    iterator,
+                )
+            }
+        }
+        FieldRef::FoldSpecificField(fold_field) => {
+            if component.folds.contains_key(&fold_field.fold_eid) {
+                compute_fold_specific_field_with_separate_value(
+                    fold_field.fold_eid,
+                    &fold_field.kind,
+                    iterator,
+                )
+            } else {
+                // This value represents an imported tag value from an outer component.
+                // Grab its value from the context itself.
+                let cloned_ref = field_ref.clone();
+                Box::new(iterator.map(move |ctx| {
+                    let right_value = ctx.imported_tags[&cloned_ref].clone();
+                    (ctx, right_value)
+                }))
+            }
+        }
+    }
+}
diff --git a/trustfall_core/src/interpreter/transformation.rs b/trustfall_core/src/interpreter/transformation.rs
diff --git a/trustfall_core/src/ir/mod.rs b/trustfall_core/src/ir/mod.rs
