add min_fold_count_limit optimization

trustfall_core/src/interpreter/execution.rs

@@ -317,10 +317,53 @@ fn get_max_fold_count_limit(carrier: &mut QueryCarrier, fold: &IRFold) -> Option
     result
 }
 
+/// If this IRFold has a filter on the folded element count, and that filter imposes
+/// a min size that can be statically determined, return that min size so it can
+/// be used for further optimizations. Otherwise, return None.
+fn get_min_fold_count_limit(carrier: &mut QueryCarrier, fold: &IRFold) -> Option<usize> {
+    let mut result: Option<usize> = None;
+
+    let query_arguments = &carrier.query.as_ref().expect("query was not returned").arguments;
+    for post_fold_filter in fold.post_filters.iter() {
+        let next_limit = match post_fold_filter {
+            Operation::GreaterThanOrEqual(
+                FoldSpecificFieldKind::Count,
+                Argument::Variable(var_ref),
+            ) => {
+                let variable_value =
+                    usize_from_field_value(&query_arguments[&var_ref.variable_name])
+                        .expect("for field value to be coercible to usize");
+                Some(variable_value)
+            }
+            Operation::GreaterThan(FoldSpecificFieldKind::Count, Argument::Variable(var_ref)) => {
+                let variable_value =
+                    usize_from_field_value(&query_arguments[&var_ref.variable_name])
+                        .expect("for field value to be coercible to usize");
+                Some(variable_value.saturating_add(1))
+            }
+            // If we don't know how many elements of the fold are required
+            // to satisfy this filter, fail early.
+            _ => return None,
+        };
+
+        match (result, next_limit) {
+            (None, _) => result = next_limit,
+            (Some(l), Some(r)) if l < r => result = next_limit,
+            _ => {}
+        }
+    }
+
+    result
+}
+
 fn collect_fold_elements<'query, Vertex: Clone + Debug + 'query>(
     mut iterator: ContextIterator<'query, Vertex>,
     max_fold_count_limit: &Option<usize>,
+    min_fold_count_limit: &Option<usize>,
 ) -> Option<Vec<DataContext<Vertex>>> {
+    // If we must collect the fold up to our upperbound of `max_fold_count_limit`,
+    // then we won't use our lowerbound of `min_fold_count_limit`, as by definition
+    // the upperbound will be larger than the lowerbound.
     if let Some(max_fold_count_limit) = max_fold_count_limit {
         // If this fold has more than `max_fold_count_limit` elements,
         // it will get filtered out by a post-fold filter.
@@ -350,9 +393,16 @@ fn collect_fold_elements<'query, Vertex: Clone + Debug + 'query>(
 
         Some(fold_elements)
     } else {
-        // We weren't able to find any early-termination condition for materializing the fold,
-        // so materialize the whole thing and return it.
-        Some(iterator.collect())
+        let collected = match min_fold_count_limit {
+            // Some queries may be able to be optimized by only partially expanding the fold,
+            // just enough to check any filters that may be applied to the fold count.
+            Some(min_fold_count_limit) => iterator.take(*min_fold_count_limit).collect(),
+            // We weren't able to find any early-termination condition for materializing the fold,
+            // so materialize the whole thing and return it.
+            _ => iterator.collect(),
+        };
+
+        Some(collected)
     }
 }
 
@@ -443,6 +493,41 @@ fn compute_fold<'query, AdapterT: Adapter<'query> + 'query>(
     let fold_component = fold.component.clone();
     let fold_eid = fold.eid;
     let max_fold_size = get_max_fold_count_limit(carrier, fold.as_ref());
+
+    // Queries that do not observe the fold count nor any fold contents may be able to
+    // be optimized by only partially expanding the fold, just enough to check any filters
+    // that may be applied to the fold count.
+    //
+    // 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;
+                    };
+
+                    tagged_fold_count.fold_root_vid == fold.to_vid
+                        && tagged_fold_count.fold_eid == fold.eid
+                        && tagged_fold_count.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)| {
         let imported_tags = context.imported_tags.clone();
@@ -466,7 +551,7 @@ fn compute_fold<'query, AdapterT: Adapter<'query> + 'query>(
         let fold_elements = if fold_exists {
             // N.B.: Note the `?` at the end here!
             //       This lets us early-discard folds that failed a post-processing filter.
-            Some(collect_fold_elements(computed_iterator, &max_fold_size)?)
+            Some(collect_fold_elements(computed_iterator, &max_fold_size, &min_fold_size)?)
         } else {
             None
         };

trustfall_core/test_data/tests/valid_queries/fold_with_no_outputs.trace.ron

@@ -118,16 +118,6 @@ TestInterpreterOutputTrace(
       ),
       Opid(15): TraceOp(
         opid: Opid(15),
-        parent_opid: Some(Opid(12)),
-        content: YieldFrom(ResolveNeighborsInner(2, Prime(PrimeNumber(5)))),
-      ),
-      Opid(16): TraceOp(
-        opid: Opid(16),
-        parent_opid: Some(Opid(12)),
-        content: OutputIteratorExhausted,
-      ),
-      Opid(17): TraceOp(
-        opid: Opid(17),
         parent_opid: Some(Opid(4)),
         content: YieldInto(SerializableContext(
           active_vertex: Some(Composite(CompositeNumber(30, [
@@ -156,18 +146,12 @@ TestInterpreterOutputTrace(
                   Vid(2): Some(Prime(PrimeNumber(3))),
                 },
               ),
-              SerializableContext(
-                active_vertex: Some(Prime(PrimeNumber(5))),
-                vertices: {
-                  Vid(2): Some(Prime(PrimeNumber(5))),
-                },
-              ),
             ]),
           },
         )),
       ),
-      Opid(18): TraceOp(
-        opid: Opid(18),
+      Opid(16): TraceOp(
+        opid: Opid(16),
         parent_opid: Some(Opid(4)),
         content: YieldFrom(ResolveProperty(SerializableContext(
           active_vertex: Some(Composite(CompositeNumber(30, [
@@ -196,70 +180,64 @@ TestInterpreterOutputTrace(
                   Vid(2): Some(Prime(PrimeNumber(3))),
                 },
               ),
-              SerializableContext(
-                active_vertex: Some(Prime(PrimeNumber(5))),
-                vertices: {
-                  Vid(2): Some(Prime(PrimeNumber(5))),
-                },
-              ),
             ]),
           },
         ), Int64(30))),
       ),
-      Opid(19): TraceOp(
-        opid: Opid(19),
+      Opid(17): TraceOp(
+        opid: Opid(17),
         parent_opid: None,
         content: ProduceQueryResult({
           "value": Int64(30),
         }),
       ),
+      Opid(18): TraceOp(
+        opid: Opid(18),
+        parent_opid: Some(Opid(4)),
+        content: AdvanceInputIterator,
+      ),
+      Opid(19): TraceOp(
+        opid: Opid(19),
+        parent_opid: Some(Opid(3)),
+        content: AdvanceInputIterator,
+      ),
       Opid(20): TraceOp(
         opid: Opid(20),
-        parent_opid: Some(Opid(4)),
+        parent_opid: Some(Opid(2)),
         content: AdvanceInputIterator,
       ),
       Opid(21): TraceOp(
         opid: Opid(21),
-        parent_opid: Some(Opid(3)),
-        content: AdvanceInputIterator,
+        parent_opid: Some(Opid(1)),
+        content: OutputIteratorExhausted,
       ),
       Opid(22): TraceOp(
         opid: Opid(22),
         parent_opid: Some(Opid(2)),
-        content: AdvanceInputIterator,
+        content: InputIteratorExhausted,
       ),
       Opid(23): TraceOp(
         opid: Opid(23),
-        parent_opid: Some(Opid(1)),
+        parent_opid: Some(Opid(2)),
         content: OutputIteratorExhausted,
       ),
       Opid(24): TraceOp(
         opid: Opid(24),
-        parent_opid: Some(Opid(2)),
+        parent_opid: Some(Opid(3)),
         content: InputIteratorExhausted,
       ),
       Opid(25): TraceOp(
         opid: Opid(25),
-        parent_opid: Some(Opid(2)),
+        parent_opid: Some(Opid(3)),
         content: OutputIteratorExhausted,
       ),
       Opid(26): TraceOp(
         opid: Opid(26),
-        parent_opid: Some(Opid(3)),
+        parent_opid: Some(Opid(4)),
         content: InputIteratorExhausted,
       ),
       Opid(27): TraceOp(
         opid: Opid(27),
-        parent_opid: Some(Opid(3)),
-        content: OutputIteratorExhausted,
-      ),
-      Opid(28): TraceOp(
-        opid: Opid(28),
-        parent_opid: Some(Opid(4)),
-        content: InputIteratorExhausted,
-      ),
-      Opid(29): TraceOp(
-        opid: Opid(29),
         parent_opid: Some(Opid(4)),
         content: OutputIteratorExhausted,
       ),

trustfall_core/test_data/tests/valid_queries/fold_with_no_outputs_elided_braces.trace.ron

@@ -118,16 +118,6 @@ TestInterpreterOutputTrace(
       ),
       Opid(15): TraceOp(
         opid: Opid(15),
-        parent_opid: Some(Opid(12)),
-        content: YieldFrom(ResolveNeighborsInner(2, Prime(PrimeNumber(5)))),
-      ),
-      Opid(16): TraceOp(
-        opid: Opid(16),
-        parent_opid: Some(Opid(12)),
-        content: OutputIteratorExhausted,
-      ),
-      Opid(17): TraceOp(
-        opid: Opid(17),
         parent_opid: Some(Opid(4)),
         content: YieldInto(SerializableContext(
           active_vertex: Some(Composite(CompositeNumber(30, [
@@ -156,18 +146,12 @@ TestInterpreterOutputTrace(
                   Vid(2): Some(Prime(PrimeNumber(3))),
                 },
               ),
-              SerializableContext(
-                active_vertex: Some(Prime(PrimeNumber(5))),
-                vertices: {
-                  Vid(2): Some(Prime(PrimeNumber(5))),
-                },
-              ),
             ]),
           },
         )),
       ),
-      Opid(18): TraceOp(
-        opid: Opid(18),
+      Opid(16): TraceOp(
+        opid: Opid(16),
         parent_opid: Some(Opid(4)),
         content: YieldFrom(ResolveProperty(SerializableContext(
           active_vertex: Some(Composite(CompositeNumber(30, [
@@ -196,70 +180,64 @@ TestInterpreterOutputTrace(
                   Vid(2): Some(Prime(PrimeNumber(3))),
                 },
               ),
-              SerializableContext(
-                active_vertex: Some(Prime(PrimeNumber(5))),
-                vertices: {
-                  Vid(2): Some(Prime(PrimeNumber(5))),
-                },
-              ),
             ]),
           },
         ), Int64(30))),
       ),
-      Opid(19): TraceOp(
-        opid: Opid(19),
+      Opid(17): TraceOp(
+        opid: Opid(17),
         parent_opid: None,
         content: ProduceQueryResult({
           "value": Int64(30),
         }),
       ),
+      Opid(18): TraceOp(
+        opid: Opid(18),
+        parent_opid: Some(Opid(4)),
+        content: AdvanceInputIterator,
+      ),
+      Opid(19): TraceOp(
+        opid: Opid(19),
+        parent_opid: Some(Opid(3)),
+        content: AdvanceInputIterator,
+      ),
       Opid(20): TraceOp(
         opid: Opid(20),
-        parent_opid: Some(Opid(4)),
+        parent_opid: Some(Opid(2)),
         content: AdvanceInputIterator,
       ),
       Opid(21): TraceOp(
         opid: Opid(21),
-        parent_opid: Some(Opid(3)),
-        content: AdvanceInputIterator,
+        parent_opid: Some(Opid(1)),
+        content: OutputIteratorExhausted,
       ),
       Opid(22): TraceOp(
         opid: Opid(22),
         parent_opid: Some(Opid(2)),
-        content: AdvanceInputIterator,
+        content: InputIteratorExhausted,
       ),
       Opid(23): TraceOp(
         opid: Opid(23),
-        parent_opid: Some(Opid(1)),
+        parent_opid: Some(Opid(2)),
         content: OutputIteratorExhausted,
       ),
       Opid(24): TraceOp(
         opid: Opid(24),
-        parent_opid: Some(Opid(2)),
+        parent_opid: Some(Opid(3)),
         content: InputIteratorExhausted,
       ),
       Opid(25): TraceOp(
         opid: Opid(25),
-        parent_opid: Some(Opid(2)),
+        parent_opid: Some(Opid(3)),
         content: OutputIteratorExhausted,
       ),
       Opid(26): TraceOp(
         opid: Opid(26),
-        parent_opid: Some(Opid(3)),
+        parent_opid: Some(Opid(4)),
         content: InputIteratorExhausted,
       ),
       Opid(27): TraceOp(
         opid: Opid(27),
-        parent_opid: Some(Opid(3)),
-        content: OutputIteratorExhausted,
-      ),
-      Opid(28): TraceOp(
-        opid: Opid(28),
-        parent_opid: Some(Opid(4)),
-        content: InputIteratorExhausted,
-      ),
-      Opid(29): TraceOp(
-        opid: Opid(29),
         parent_opid: Some(Opid(4)),
         content: OutputIteratorExhausted,
       ),