Add utility methods for field destinations and selection fields

2 files changed, 86 insertions, 8 deletions

diff --git a/source/blender/functions/FN_field.hh b/source/blender/functions/FN_field.hh
index 2262bfdb641..81c18b2dc54 100644
--- a/source/blender/functions/FN_field.hh
+++ b/source/blender/functions/FN_field.hh
@@ -142,7 +142,8 @@ class GField : public GFieldBase<std::shared_ptr<FieldNode>> {
   }
 };
 
-/** Same as #GField but is cheaper to copy/move around, because it does not contain a
+/**
+ * Same as #GField but is cheaper to copy/move around, because it does not contain a
  * #std::shared_ptr.
  */
 class GFieldRef : public GFieldBase<const FieldNode *> {
@@ -287,6 +288,8 @@ void evaluate_fields_to_spans(Span<GFieldRef> fields_to_evaluate,
                               const FieldContext &context,
                               Span<GMutableSpan> out_spans);
 
+Vector<int64_t> indices_from_selection(const VArray<bool> &selection);
+
 template<typename T> T evaluate_constant_field(const Field<T> &field)
 {
   T value;
@@ -326,6 +329,9 @@ class FieldEvaluator : NonMovable, NonCopyable {
       : context_(context), mask_(*mask)
   {
   }
+  FieldEvaluator(const FieldContext &context, const int64_t size) : context_(context), mask_(size)
+  {
+  }
 
   /**
    * \param field: Field to add to the evaluator.
@@ -347,6 +353,34 @@ class FieldEvaluator : NonMovable, NonCopyable {
     return this->add_with_destination(GField(std::move(field)), generic_dst_hint);
   }
 
+  /**
+   * \param field: Field to add to the evaluator.
+   * \param dst: Mutable span that the evaluated result for this field is be written into.
+   * \note: When the output may only be used as a single value, the version of this function with
+   * a virtual array result array should be used.
+   */
+  int add_with_destination(GField field, GMutableSpan dst)
+  {
+    const int field_index = fields_to_evaluate_.append_and_get_index(std::move(field));
+    dst_hints_.append(&scope_.construct<GVMutableArray_For_GMutableSpan>(__func__, dst));
+    output_pointer_infos_.append({});
+    return field_index;
+  }
+
+  /**
+   * \param field: Field to add to the evaluator.
+   * \param dst: Mutable span that the evaluated result for this field is be written into.
+   * \note: When the output may only be used as a single value, the version of this function with
+   * a virtual array result array should be used.
+   */
+  template<typename T> int add_with_destination(Field<T> field, MutableSpan<T> dst)
+  {
+    const int field_index = fields_to_evaluate_.append_and_get_index(std::move(field));
+    dst_hints_.append(&scope_.construct<GVMutableArray_For_MutableSpan<T>>(__func__, dst));
+    output_pointer_infos_.append({});
+    return field_index;
+  }
+
   int add(GField field, const GVArray **varray_ptr)
   {
     const int field_index = fields_to_evaluate_.append_and_get_index(std::move(field));
@@ -419,6 +453,26 @@ class FieldEvaluator : NonMovable, NonCopyable {
     GVArray_Typed<T> &typed_varray = scope_.construct<GVArray_Typed<T>>(__func__, varray);
     return *typed_varray;
   }
+
+  /**
+   * 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
+   * some cases, so it must live at least as long as the returned mask.
+   */
+  IndexMask get_evaluated_as_mask(const int field_index)
+  {
+    const GVArray &varray = this->get_evaluated(field_index);
+    GVArray_Typed<bool> typed_varray{varray};
+
+    if (typed_varray->is_single()) {
+      if (typed_varray->get_internal_single()) {
+        return IndexRange(typed_varray.size());
+      }
+      return IndexRange(0);
+    }
+
+    return scope_.add_value(indices_from_selection(*typed_varray), __func__).as_span();
+  }
 };
 
 }  // namespace blender::fn
diff --git a/source/blender/functions/intern/field.cc b/source/blender/functions/intern/field.cc
index f10a154dcd6..97c3b0b368e 100644
--- a/source/blender/functions/intern/field.cc
+++ b/source/blender/functions/intern/field.cc
@@ -32,8 +32,8 @@ struct FieldTreeInfo {
    */
   MultiValueMap<GFieldRef, GFieldRef> field_users;
   /**
-   * The same field input may exist in the field tree as as separate nodes due to the way the tree
-   * is constructed. This set contains every input only once.
+   * The same field input may exist in the field tree as as separate nodes due to the way
+   * the tree is constructed. This set contains every different input only once.
    */
   VectorSet<std::reference_wrapper<const FieldInput>> deduplicated_field_inputs;
 };
@@ -205,8 +205,8 @@ static void build_multi_function_procedure_for_fields(MFProcedure &procedure,
   for (const GFieldRef &field : output_fields) {
     MFVariable *variable = variable_by_field.lookup(field);
     if (!already_output_variables.add(variable)) {
-      /* One variable can be output at most once. To output the same value twice, we have to make a
-       * copy first. */
+      /* One variable can be output at most once. To output the same value twice, we have to make
+       * a copy first. */
       const MultiFunction &copy_fn = scope.construct<CustomMF_GenericCopy>(
           __func__, "copy", variable->data_type());
       variable = builder.add_call<1>(copy_fn, {variable})[0];
@@ -258,7 +258,7 @@ struct PartiallyInitializedArray : NonCopyable, NonMovable {
  *   instead of into newly created ones. That allows making the computed data live longer than
  *   #scope and is more efficient when the data will be written into those virtual arrays
  *   later anyway.
- * \return The computed virtual arrays for each provided field. If #dst_hints were passed, the
+ * \return The computed virtual arrays for each provided field. If #dst_hints is passed, the
  *   provided virtual arrays are returned.
  */
 Vector<const GVArray *> evaluate_fields(ResourceScope &scope,
@@ -284,7 +284,7 @@ Vector<const GVArray *> evaluate_fields(ResourceScope &scope,
   Vector<const GVArray *> field_context_inputs = get_field_context_inputs(
       scope, mask, context, field_tree_info.deduplicated_field_inputs);
 
-  /* Finish fields that output a context varray directly. For those we don't have to do any further
+  /* Finish fields that output an input varray directly. For those we don't have to do any further
    * processing. */
   for (const int out_index : fields_to_evaluate.index_range()) {
     const GFieldRef &field = fields_to_evaluate[out_index];
@@ -330,7 +330,7 @@ Vector<const GVArray *> evaluate_fields(ResourceScope &scope,
     build_multi_function_procedure_for_fields(
         procedure, scope, field_tree_info, varying_fields_to_evaluate);
     MFProcedureExecutor procedure_executor{"Procedure", procedure};
-    MFParamsBuilder mf_params{procedure_executor, mask.min_array_size()};
+    MFParamsBuilder mf_params{procedure_executor, array_size};
     MFContextBuilder mf_context;
 
     /* Provide inputs to the procedure executor. */
@@ -480,4 +480,28 @@ const GVArray *FieldContext::try_get_varray_for_context(const FieldInput &field_
   return field_input.try_get_varray_for_context(*this, mask, scope);
 }
 
+Vector<int64_t> indices_from_selection(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. */
+  BLI_assert(!selection.is_single());
+  Vector<int64_t> indices;
+  if (selection.is_span()) {
+    Span<bool> span = selection.get_internal_span();
+    for (const int64_t i : span.index_range()) {
+      if (span[i]) {
+        indices.append(i);
+      }
+    }
+  }
+  else {
+    for (const int i : selection.index_range()) {
+      if (selection[i]) {
+        indices.append(i);
+      }
+    }
+  }
+  return indices;
+}
+
 }  // namespace blender::fn