Geometry Nodes: simplify using selection when evaluating fields

We often had to use two `FieldEvaluator` instances to first evaluate
the selection and then the remaining fields. Now both can be done
with a single `FieldEvaluator`. This results in less boilerplate code in
many cases.

Performance is not affected by this change. In a separate patch we
could improve performance by reusing evaluated sub-fields that are
used by the selection and the other fields.

Differential Revision: https://developer.blender.org/D13571

2 files changed, 50 insertions, 5 deletions

diff --git a/source/blender/functions/FN_field.hh b/source/blender/functions/FN_field.hh
index cf96eff62bd..a591aaed34a 100644
--- a/source/blender/functions/FN_field.hh
+++ b/source/blender/functions/FN_field.hh
@@ -313,6 +313,9 @@ class FieldEvaluator : NonMovable, NonCopyable {
   Vector<OutputPointerInfo> output_pointer_infos_;
   bool is_evaluated_ = false;
 
+  Field<bool> selection_field_;
+  IndexMask selection_mask_;
+
  public:
   /** Takes #mask by pointer because the mask has to live longer than the evaluator. */
   FieldEvaluator(const FieldContext &context, const IndexMask *mask)
@@ -333,6 +336,18 @@ class FieldEvaluator : NonMovable, NonCopyable {
   }
 
   /**
+   * The selection field is evaluated first to determine which indices of the other fields should
+   * be evaluated. Calling this method multiple times will just replace the previously set
+   * selection field. Only the elements selected by both this selection and the selection provided
+   * in the constructor are calculated. If no selection field is set, it is assumed that all
+   * indices passed to the constructor are selected.
+   */
+  void set_selection(Field<bool> selection)
+  {
+    selection_field_ = std::move(selection);
+  }
+
+  /**
    * \param field: Field to add to the evaluator.
    * \param dst: Mutable virtual array that the evaluated result for this field is be written into.
    */
@@ -403,6 +418,8 @@ class FieldEvaluator : NonMovable, NonCopyable {
     return this->get_evaluated(field_index).typed<T>();
   }
 
+  IndexMask get_evaluated_selection_as_mask();
+
   /**
    * Retrieve the output of an evaluated boolean field and convert it to a mask, which can be used
    * to avoid calculations for unnecessary elements later on. The evaluator will own the indices in
diff --git a/source/blender/functions/intern/field.cc b/source/blender/functions/intern/field.cc
index 5fa493c8336..604e5c6d13f 100644
--- a/source/blender/functions/intern/field.cc
+++ b/source/blender/functions/intern/field.cc
@@ -624,7 +624,7 @@ FieldInput::FieldInput(const CPPType &type, std::string debug_name)
  * FieldEvaluator.
  */
 
-static Vector<int64_t> indices_from_selection(const VArray<bool> &selection)
+static Vector<int64_t> indices_from_selection(IndexMask mask, const VArray<bool> &selection)
 {
   /* If the selection is just a single value, it's best to avoid calling this
    * function when constructing an IndexMask and use an IndexRange instead. */
@@ -633,14 +633,14 @@ static Vector<int64_t> indices_from_selection(const VArray<bool> &selection)
   Vector<int64_t> indices;
   if (selection.is_span()) {
     Span<bool> span = selection.get_internal_span();
-    for (const int64_t i : span.index_range()) {
+    for (const int64_t i : mask) {
       if (span[i]) {
         indices.append(i);
       }
     }
   }
   else {
-    for (const int i : selection.index_range()) {
+    for (const int i : mask) {
       if (selection[i]) {
         indices.append(i);
       }
@@ -681,14 +681,36 @@ int FieldEvaluator::add(GField field)
   return field_index;
 }
 
+static IndexMask evaluate_selection(const Field<bool> &selection_field,
+                                    const FieldContext &context,
+                                    IndexMask full_mask,
+                                    ResourceScope &scope)
+{
+  if (selection_field) {
+    VArray<bool> selection =
+        evaluate_fields(scope, {selection_field}, full_mask, context)[0].typed<bool>();
+    if (selection.is_single()) {
+      if (selection.get_internal_single()) {
+        return full_mask;
+      }
+      return IndexRange(0);
+    }
+    return scope.add_value(indices_from_selection(full_mask, selection)).as_span();
+  }
+  return full_mask;
+}
+
 void FieldEvaluator::evaluate()
 {
   BLI_assert_msg(!is_evaluated_, "Cannot evaluate fields twice.");
+
+  selection_mask_ = evaluate_selection(selection_field_, context_, mask_, scope_);
+
   Array<GFieldRef> fields(fields_to_evaluate_.size());
   for (const int i : fields_to_evaluate_.index_range()) {
     fields[i] = fields_to_evaluate_[i];
   }
-  evaluated_varrays_ = evaluate_fields(scope_, fields, mask_, context_, dst_varrays_);
+  evaluated_varrays_ = evaluate_fields(scope_, fields, selection_mask_, context_, dst_varrays_);
   BLI_assert(fields_to_evaluate_.size() == evaluated_varrays_.size());
   for (const int i : fields_to_evaluate_.index_range()) {
     OutputPointerInfo &info = output_pointer_infos_[i];
@@ -710,7 +732,13 @@ IndexMask FieldEvaluator::get_evaluated_as_mask(const int field_index)
     return IndexRange(0);
   }
 
-  return scope_.add_value(indices_from_selection(varray)).as_span();
+  return scope_.add_value(indices_from_selection(mask_, varray)).as_span();
+}
+
+IndexMask FieldEvaluator::get_evaluated_selection_as_mask()
+{
+  BLI_assert(is_evaluated_);
+  return selection_mask_;
 }
 
 }  // namespace blender::fn