Refactor field storage again, to allow reuse of function outputs

3 files changed, 103 insertions, 112 deletions

diff --git a/source/blender/functions/FN_field.hh b/source/blender/functions/FN_field.hh
index 9ccdca4238c..ff38d40e395 100644
--- a/source/blender/functions/FN_field.hh
+++ b/source/blender/functions/FN_field.hh
@@ -20,11 +20,8 @@
  * \ingroup fn
  */
 
-#include "BLI_function_ref.hh"
-#include "BLI_map.hh"
 #include "BLI_vector.hh"
 
-#include "FN_generic_virtual_array.hh"
 #include "FN_multi_function_procedure.hh"
 #include "FN_multi_function_procedure_builder.hh"
 #include "FN_multi_function_procedure_executor.hh"
@@ -33,19 +30,65 @@ namespace blender::fn {
 
 class Field;
 
+/**
+ * An operation acting on data described by fields. Generally corresponds
+ * to a node or a subset of a node in a node graph.
+ */
+class Function {
+  /**
+   * The function used to calculate the
+   */
+  std::unique_ptr<MultiFunction> function_;
+
+  /**
+   * References to descriptions of the results from the functions this function depends on.
+   */
+  blender::Vector<Field *> inputs_;
+
+ public:
+  Function(std::unique_ptr<MultiFunction> function, Span<Field *> inputs)
+      : function_(std::move(function)), inputs_(inputs)
+  {
+  }
+
+  Span<Field *> inputs() const
+  {
+    return inputs_;
+  }
+
+  const MultiFunction &multi_function() const
+  {
+    return *function_;
+  }
+};
+
+/**
+ * Descibes the output of a function. Generally corresponds to the combination of an output socket
+ * and link combination in a node graph.
+ */
 class Field {
+  /**
+   * The type of this field's result.
+   */
   const fn::CPPType *type_;
-  // std::unique_ptr<MultiFunction> function_;
-  const MultiFunction *function_;
 
-  blender::Vector<std::shared_ptr<Field>> input_fields_;
+  /**
+   * The function that calculates this field's values. Many fields can share the same function,
+   * since a function can have many outputs, just like a node graph, where a single output can be
+   * used as multiple inputs. This avoids calling the same function many times, only using one of
+   * its results.
+   */
+  const Function *function_;
+  /**
+   * Which output of the function this field corresponds to.
+   */
+  int output_index_;
 
   std::string debug_name_ = "";
 
  public:
-  ~Field() = default;
-  Field(const fn::CPPType &type, const MultiFunction &function)
-      : type_(&type), function_(&function)
+  Field(const fn::CPPType &type, const Function &function, const int output_index)
+      : type_(&type), function_(&function), output_index_(output_index)
   {
   }
 
@@ -55,100 +98,29 @@ class Field {
     return *type_;
   }
 
-  const MultiFunction &function() const
+  const Function &function() const
   {
     BLI_assert(function_ != nullptr);
     return *function_;
   }
 
-  blender::StringRef debug_name() const
+  int function_output_index() const
   {
-    return debug_name_;
+    return output_index_;
   }
 
-  void foreach_input(blender::FunctionRef<void(const Field &input)> fn) const
-  {
-    for (const std::shared_ptr<Field> &field : input_fields_) {
-      fn(*field);
-    }
-  }
-  void foreach_input_recursive(blender::FunctionRef<void(const Field &input)> fn) const
+  blender::StringRef debug_name() const
   {
-    for (const std::shared_ptr<Field> &field : input_fields_) {
-      fn(*field);
-      field->foreach_input(fn);
-    }
+    return debug_name_;
   }
 };
 
 /**
- * A field that doesn't have any dependencies on other fields.
- *
- * TODO: It might be an elegant simplification if every single field was a multi-function field,
- * and input fields just happened to have no inputs. Then it might not need to be a virtual class,
- * since the dynamic behavior would be contained in the multifunction, which would be very nice.
- */
-// class InputField : public Field {
-//  public:
-//   InputField(const CPPType &type) : Field(type)
-//   {
-//   }
-
-//   void foreach_input(blender::FunctionRef<void(const Field &input)> UNUSED(fn)) const final
-//   {
-//   }
-//   void foreach_input_recursive(
-//       blender::FunctionRef<void(const Field &input)> UNUSED(fn)) const final
-//   {
-//   }
-
-//   virtual GVArrayPtr get_data(IndexMask mask) const = 0;
-
-//   /**
-//    * Return true when the field input is the same as another field, used as an
-//    * optimization to avoid creating multiple virtual arrays for the same input node.
-//    */
-//   virtual bool equals(const InputField &UNUSED(other))
-//   {
-//     return false;
-//   }
-// };
-
-/**
- * A field that takes inputs
- */
-// class MultiFunctionField final : public Field {
-//   blender::Vector<FieldPtr> input_fields_;
-//   const MultiFunction *function_;
-
-//  public:
-//   void foreach_input(blender::FunctionRef<void(const Field &input)> fn) const final
-//   {
-//     for (const FieldPtr &field : input_fields_) {
-//       fn(*field);
-//     }
-//   }
-//   void foreach_input_recursive(blender::FunctionRef<void(const Field &input)> fn) const final
-//   {
-//     for (const FieldPtr &field : input_fields_) {
-//       fn(*field);
-//       field->foreach_input(fn);
-//     }
-//   }
-
-//   const MultiFunction &function() const
-//   {
-//     BLI_assert(function_ != nullptr);
-//     return *function_;
-//   }
-// };
-
-/**
  * Evaluate more than one field at a time, as an optimization
  * in case they share inputs or various intermediate values.
  */
-void evaluate_fields(const blender::Span<Field> fields,
-                     const blender::MutableSpan<GMutableSpan> outputs,
-                     const blender::IndexMask mask);
+void evaluate_fields(blender::Span<Field> fields,
+                     blender::IndexMask mask,
+                     blender::MutableSpan<GMutableSpan> outputs);
 
 }  // namespace blender::fn
\ No newline at end of file
diff --git a/source/blender/functions/intern/field.cc b/source/blender/functions/intern/field.cc
index ac017822733..77331f5681c 100644
--- a/source/blender/functions/intern/field.cc
+++ b/source/blender/functions/intern/field.cc
@@ -14,49 +14,68 @@
  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  */
 
+#include "BLI_map.hh"
+#include "BLI_multi_value_map.hh"
+
 #include "FN_field.hh"
 
 namespace blender::fn {
 
-static void add_field_parameters(const Field &field,
-                                 MFProcedureBuilder &builder,
-                                 Map<const Field *, MFVariable *> &output_map)
+/* A map to hold the output variables for each function so they can be reused. */
+using OutputMap = MultiValueMap<const Function *, MFVariable *>;
+
+static MFVariable *get_field_variable(const Field &field, const OutputMap &output_map)
 {
-  /* Recursively make sure all of the inputs have entries in the variable map. */
-  field.foreach_input_recursive(
-      [&](const Field &input_field) { add_field_parameters(input_field, builder, output_map); });
+  const Function &input_field_function = field.function();
+  const Span<MFVariable *> input_function_outputs = output_map.lookup(&input_field_function);
+  return input_function_outputs[field.function_output_index()];
+}
 
-  /* Add the immediate inputs to this field. */
+/**
+ * Traverse the fields recursively. Eventually there will be a field whose function has no
+ * inputs. Start adding multi-function variables there. Those outputs are then used as inputs
+ * for the dependent functions, and the rest of the field tree is built up from there.
+ */
+static void add_field_variables(const Field &field,
+                                MFProcedureBuilder &builder,
+                                OutputMap &output_map)
+{
+  const Function &function = field.function();
+  for (const Field *input_field : function.inputs()) {
+    add_field_variables(*input_field, builder, output_map);
+  }
+
+  /* Add the immediate inputs to this field, which were added before in the recursive call.
+   * This will be skipped for functions with no inputs. */
   Vector<MFVariable *> inputs;
-  field.foreach_input([&](const Field &input_field) {
-    MFVariable *input = output_map.lookup(&input_field);
+  for (const Field *input_field : function.inputs()) {
+    MFVariable *input = get_field_variable(*input_field, output_map);
     builder.add_input_parameter(input->data_type());
     inputs.append(input);
-  });
+  }
 
-  Vector<MFVariable *> outputs = builder.add_call(field.function(), inputs);
+  Vector<MFVariable *> outputs = builder.add_call(function.multi_function(), inputs);
 
   builder.add_destruct(inputs);
 
-  /* TODO: How to support multiple outputs?! */
-  BLI_assert(outputs.size() == 1);
-  output_map.add_new(&field, outputs.first());
+  output_map.add_multiple(&function, outputs);
 }
 
 static void build_procedure(const Span<Field> fields, MFProcedure &procedure)
 {
   MFProcedureBuilder builder{procedure};
 
-  Map<const Field *, MFVariable *> output_map;
+  OutputMap output_map;
 
   for (const Field &field : fields) {
-    add_field_parameters(field, builder, output_map);
+    add_field_variables(field, builder, output_map);
   }
 
   builder.add_return();
 
   for (const Field &field : fields) {
-    builder.add_output_parameter(*output_map.lookup(&field));
+    MFVariable *input = get_field_variable(field, output_map);
+    builder.add_output_parameter(*input);
   }
 
   std::cout << procedure.to_dot();
@@ -84,8 +103,8 @@ static void evaluate_procedure(MFProcedure &procedure,
  * Evaluate more than one prodecure at a time
  */
 void evaluate_fields(const Span<Field> fields,
-                     const MutableSpan<GMutableSpan> outputs,
-                     const IndexMask mask)
+                     const IndexMask mask,
+                     const MutableSpan<GMutableSpan> outputs)
 {
   MFProcedure procedure;
   build_procedure(fields, procedure);
@@ -93,4 +112,4 @@ void evaluate_fields(const Span<Field> fields,
   evaluate_procedure(procedure, mask, outputs);
 }
 
-}  // namespace blender::fn
\ No newline at end of file
+}  // namespace blender::fn
diff --git a/source/blender/functions/tests/FN_field_test.cc b/source/blender/functions/tests/FN_field_test.cc
index 9949235699f..5178b7f5973 100644
--- a/source/blender/functions/tests/FN_field_test.cc
+++ b/source/blender/functions/tests/FN_field_test.cc
@@ -10,13 +10,13 @@ namespace blender::fn::tests {
 
 TEST(field, ConstantFieldTest)
 {
-  CustomMF_Constant<int> const_fn{10};
-  Field constant_field = Field(CPPType::get<int>(), const_fn);
+  std::unique_ptr<CustomMF_Constant<int>> const_fn = std::make_unique<CustomMF_Constant<int>>(10);
+  Function function = Function(std::move(const_fn), {});
+  Field constant_field = Field(CPPType::get<int>(), function, 0);
 
   Array<int> result(4);
   GMutableSpan result_generic(result.as_mutable_span());
-  // MutableSpan<GMutableSpan> result_span{&result_generic, 1};
-  evaluate_fields({&constant_field, 1}, {&result_generic, 1}, IndexMask(IndexRange(4)));
+  evaluate_fields({&constant_field, 1}, IndexMask(IndexRange(4)), {&result_generic, 1});
 
   ASSERT_EQ(result[0], 10);
   ASSERT_EQ(result[1], 10);