add common base class for field input/operationtemp-geometry-nodes-fields--fields-jacques

5 files changed, 109 insertions, 179 deletions

diff --git a/source/blender/functions/FN_field.hh b/source/blender/functions/FN_field.hh
index 4b0dee88491..9217d98d75b 100644
--- a/source/blender/functions/FN_field.hh
+++ b/source/blender/functions/FN_field.hh
@@ -41,69 +41,57 @@
 
 namespace blender::fn {
 
-class FieldInput;
-class FieldFunction;
+class FieldSource {
+ public:
+  ~FieldSource() = default;
+
+  virtual const CPPType &cpp_type_of_output_index(int output_index) const = 0;
+
+  virtual bool is_input() const
+  {
+    return false;
+  }
+};
 
 /**
  * Descibes the output of a function. Generally corresponds to the combination of an output socket
  * and link combination in a node graph.
  */
 class GField {
-  /**
-   * 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.
-   */
-  std::shared_ptr<FieldFunction> function_;
-  /**
-   * Which output of the function this field corresponds to.
-   */
-  int output_index_ = 0;
-
-  std::shared_ptr<FieldInput> input_;
+  std::shared_ptr<FieldSource> source_;
+  int source_output_index_ = 0;
 
  public:
   GField() = default;
 
-  GField(std::shared_ptr<FieldFunction> function, const int output_index)
-      : function_(std::move(function)), output_index_(output_index)
+  GField(std::shared_ptr<FieldSource> source, const int source_output_index = 0)
+      : source_(std::move(source)), source_output_index_(source_output_index)
   {
   }
 
-  GField(std::shared_ptr<FieldInput> input) : input_(std::move(input))
+  operator bool() const
   {
+    return source_ != nullptr;
   }
 
-  const fn::CPPType &cpp_type() const;
-
-  bool is_input() const
-  {
-    return input_.get() != nullptr;
-  }
-  const FieldInput &input() const
+  const fn::CPPType &cpp_type() const
   {
-    BLI_assert(!function_);
-    BLI_assert(input_);
-    return *input_;
+    return source_->cpp_type_of_output_index(source_output_index_);
   }
 
-  bool is_function() const
+  bool is_input() const
   {
-    return function_.get() != nullptr;
+    return source_->is_input();
   }
-  const FieldFunction &function() const
+
+  const FieldSource &source() const
   {
-    BLI_assert(function_ != nullptr);
-    BLI_assert(input_ == nullptr);
-    return *function_;
+    return *source_;
   }
 
-  int function_output_index() const
+  int source_output_index() const
   {
-    BLI_assert(function_ != nullptr);
-    BLI_assert(input_ == nullptr);
-    return output_index_;
+    return source_output_index_;
   }
 };
 
@@ -117,30 +105,20 @@ template<typename T> class Field : public GField {
   }
 };
 
-/**
- * An operation acting on data described by fields. Generally corresponds
- * to a node or a subset of a node in a node graph.
- */
-class FieldFunction {
-  /**
-   * The function used to calculate the field.
-   */
+class FieldOperation : public FieldSource {
   std::unique_ptr<const MultiFunction> owned_function_;
   const MultiFunction *function_;
 
-  /**
-   * References to descriptions of the results from the functions this function depends on.
-   */
   blender::Vector<GField> inputs_;
 
  public:
-  FieldFunction(std::unique_ptr<const MultiFunction> function, Vector<GField> inputs = {})
+  FieldOperation(std::unique_ptr<const MultiFunction> function, Vector<GField> inputs = {})
       : owned_function_(std::move(function)), inputs_(std::move(inputs))
   {
     function_ = owned_function_.get();
   }
 
-  FieldFunction(const MultiFunction &function, Vector<GField> inputs = {})
+  FieldOperation(const MultiFunction &function, Vector<GField> inputs = {})
       : function_(&function), inputs_(std::move(inputs))
   {
   }
@@ -155,7 +133,7 @@ class FieldFunction {
     return *function_;
   }
 
-  const CPPType &cpp_type_of_output_index(int output_index) const
+  const CPPType &cpp_type_of_output_index(int output_index) const override
   {
     int output_counter = 0;
     for (const int param_index : function_->param_indices()) {
@@ -172,7 +150,7 @@ class FieldFunction {
   }
 };
 
-class FieldInput {
+class FieldInput : public FieldSource {
  protected:
   const CPPType *type_;
   std::string debug_name_;
@@ -194,6 +172,17 @@ class FieldInput {
   {
     return *type_;
   }
+
+  const CPPType &cpp_type_of_output_index(int output_index) const override
+  {
+    BLI_assert(output_index == 0);
+    return *type_;
+  }
+
+  bool is_input() const override
+  {
+    return true;
+  }
 };
 
 /**
@@ -216,20 +205,8 @@ template<typename T> T evaluate_constant_field(const Field<T> &field)
 template<typename T> Field<T> make_constant_field(T value)
 {
   auto constant_fn = std::make_unique<fn::CustomMF_Constant<T>>(std::forward<T>(value));
-  auto field_fn = std::make_shared<FieldFunction>(std::move(constant_fn));
-  return Field<T>{GField{std::move(field_fn), 0}};
-}
-
-/* --------------------------------------------------------------------
- * GField inline methods.
- */
-
-inline const CPPType &GField::cpp_type() const
-{
-  if (this->is_function()) {
-    return function_->cpp_type_of_output_index(output_index_);
-  }
-  return input_->cpp_type();
+  auto operation = std::make_shared<FieldOperation>(std::move(constant_fn));
+  return Field<T>{GField{std::move(operation), 0}};
 }
 
 }  // namespace blender::fn
diff --git a/source/blender/functions/intern/field.cc b/source/blender/functions/intern/field.cc
index 69a31699592..87cfc60d502 100644
--- a/source/blender/functions/intern/field.cc
+++ b/source/blender/functions/intern/field.cc
@@ -20,60 +20,12 @@
 
 #include "FN_field.hh"
 
-/**
- * TODO: There might be a more obvious way to implement this, or we might end up with
- * a separate map for functions and inputs anyway, so we could just remove it.
- */
-struct InputOrFunction {
-  const void *ptr;
-
- public:
-  InputOrFunction(const blender::fn::FieldFunction &function) : ptr(&function)
-  {
-  }
-  InputOrFunction(const blender::fn::FieldInput &input) : ptr(&input)
-  {
-  }
-  InputOrFunction(const blender::fn::GField &field) /* Maybe this is too clever. */
-  {
-    if (field.is_function()) {
-      ptr = &field.function();
-    }
-    else {
-      ptr = &field.input();
-    }
-  }
-  friend bool operator==(const InputOrFunction &a, const InputOrFunction &b)
-  {
-    return a.ptr == b.ptr;
-  }
-};
-
-template<> struct blender::DefaultHash<InputOrFunction> {
-  uint64_t operator()(const InputOrFunction &value) const
-  {
-    return DefaultHash<const void *>{}(value.ptr);
-  }
-};
-
 namespace blender::fn {
 
 /**
- * TODO: This exists because it seemed helpful for the procedure creation to be able to store
- * mutable data for each input or function output. That still may be helpful in the future, but
- * currently it isn't useful.
- */
-struct FieldVariable {
-  MFVariable *mf_variable;
-  FieldVariable(MFVariable &variable) : mf_variable(&variable)
-  {
-  }
-};
-
-/**
  * A map to hold the output variables for each function output or input so they can be reused.
  */
-using VariableMap = Map<InputOrFunction, Vector<FieldVariable>>;
+using VariableMap = Map<const FieldSource *, Vector<MFVariable *>>;
 
 /**
  * A map of the computed inputs for all of a field system's inputs, to avoid creating duplicates.
@@ -81,27 +33,27 @@ using VariableMap = Map<InputOrFunction, Vector<FieldVariable>>;
  */
 using ComputedInputMap = Map<const MFVariable *, GVArrayPtr>;
 
-static FieldVariable &get_field_variable(const GField &field, VariableMap &unique_variables)
+static MFVariable &get_field_variable(const GField &field, VariableMap &unique_variables)
 {
   if (field.is_input()) {
-    const FieldInput &input = field.input();
-    return unique_variables.lookup(input).first();
+    const FieldInput &input = dynamic_cast<const FieldInput &>(field.source());
+    return *unique_variables.lookup(&input).first();
   }
-  const FieldFunction &function = field.function();
-  MutableSpan<FieldVariable> function_outputs = unique_variables.lookup(function);
-  return function_outputs[field.function_output_index()];
+  const FieldOperation &operation = dynamic_cast<const FieldOperation &>(field.source());
+  MutableSpan<MFVariable *> operation_outputs = unique_variables.lookup(&operation);
+  return *operation_outputs[field.source_output_index()];
 }
 
-static const FieldVariable &get_field_variable(const GField &field,
-                                               const VariableMap &unique_variables)
+static const MFVariable &get_field_variable(const GField &field,
+                                            const VariableMap &unique_variables)
 {
   if (field.is_input()) {
-    const FieldInput &input = field.input();
-    return unique_variables.lookup(input).first();
+    const FieldInput &input = dynamic_cast<const FieldInput &>(field.source());
+    return *unique_variables.lookup(&input).first();
   }
-  const FieldFunction &function = field.function();
-  Span<FieldVariable> function_outputs = unique_variables.lookup(function);
-  return function_outputs[field.function_output_index()];
+  const FieldOperation &operation = dynamic_cast<const FieldOperation &>(field.source());
+  Span<MFVariable *> operation_outputs = unique_variables.lookup(&operation);
+  return *operation_outputs[field.source_output_index()];
 }
 
 /**
@@ -113,20 +65,20 @@ static void add_variables_for_input(const GField &field,
                                     VariableMap &unique_variables)
 {
   fields_to_visit.pop();
-  const FieldInput &input = field.input();
+  const FieldInput &input = dynamic_cast<const FieldInput &>(field.source());
   MFVariable &variable = builder.add_input_parameter(MFDataType::ForSingle(field.cpp_type()),
                                                      input.debug_name());
-  unique_variables.add(input, {variable});
+  unique_variables.add(&input, {&variable});
 }
 
-static void add_variables_for_function(const GField &field,
-                                       Stack<const GField *> &fields_to_visit,
-                                       MFProcedureBuilder &builder,
-                                       VariableMap &unique_variables)
+static void add_variables_for_operation(const GField &field,
+                                        Stack<const GField *> &fields_to_visit,
+                                        MFProcedureBuilder &builder,
+                                        VariableMap &unique_variables)
 {
-  const FieldFunction &function = field.function();
-  for (const GField &input_field : function.inputs()) {
-    if (!unique_variables.contains(input_field)) {
+  const FieldOperation &operation = dynamic_cast<const FieldOperation &>(field.source());
+  for (const GField &input_field : operation.inputs()) {
+    if (!unique_variables.contains(&input_field.source())) {
       /* The field for this input hasn't been handled yet. Handle it now, so that we know all
        * of this field's function inputs already have variables. TODO: Verify that this is the
        * best way to do a depth first traversal. These extra lookups don't seem ideal. */
@@ -138,17 +90,17 @@ static void add_variables_for_function(const GField &field,
   fields_to_visit.pop();
 
   Vector<MFVariable *> inputs;
-  Set<FieldVariable *> unique_inputs;
-  for (const GField &input_field : function.inputs()) {
-    FieldVariable &input = get_field_variable(input_field, unique_variables);
+  Set<MFVariable *> unique_inputs;
+  for (const GField &input_field : operation.inputs()) {
+    MFVariable &input = get_field_variable(input_field, unique_variables);
     unique_inputs.add(&input);
-    inputs.append(input.mf_variable);
+    inputs.append(&input);
   }
 
-  Vector<MFVariable *> outputs = builder.add_call(function.multi_function(), inputs);
-  Vector<FieldVariable> &unique_outputs = unique_variables.lookup_or_add(function, {});
+  Vector<MFVariable *> outputs = builder.add_call(operation.multi_function(), inputs);
+  Vector<MFVariable *> &unique_outputs = unique_variables.lookup_or_add(&operation, {});
   for (MFVariable *output : outputs) {
-    unique_outputs.append(*output);
+    unique_outputs.append(output);
   }
 }
 
@@ -163,7 +115,7 @@ static void add_unique_variables(const Span<GField> fields,
 
   while (!fields_to_visit.is_empty()) {
     const GField &field = *fields_to_visit.peek();
-    if (unique_variables.contains(field)) {
+    if (unique_variables.contains(&field.source())) {
       fields_to_visit.pop();
       continue;
     }
@@ -172,7 +124,7 @@ static void add_unique_variables(const Span<GField> fields,
       add_variables_for_input(field, fields_to_visit, builder, unique_variables);
     }
     else {
-      add_variables_for_function(field, fields_to_visit, builder, unique_variables);
+      add_variables_for_operation(field, fields_to_visit, builder, unique_variables);
     }
   }
 }
@@ -188,18 +140,18 @@ static void add_destructs(const Span<GField> fields,
                           VariableMap &unique_variables)
 {
   Set<MFVariable *> destructed_variables;
-  Set<const FieldVariable *> outputs;
+  Set<MFVariable *> outputs;
   for (const GField &field : fields) {
     /* Currently input fields are handled separately in the evaluator. */
     BLI_assert(!field.is_input());
     outputs.add(&get_field_variable(field, unique_variables));
   }
 
-  for (Span<FieldVariable> variables : unique_variables.values()) {
-    for (const FieldVariable &variable : variables) {
+  for (MutableSpan<MFVariable *> variables : unique_variables.values()) {
+    for (MFVariable *variable : variables) {
       /* Don't destruct the variable if it is used as an output parameter. */
-      if (!outputs.contains(&variable)) {
-        builder.add_destruct(*variable.mf_variable);
+      if (!outputs.contains(variable)) {
+        builder.add_destruct(*variable);
       }
     }
   }
@@ -218,7 +170,7 @@ static void build_procedure(const Span<GField> fields,
   builder.add_return();
 
   for (const GField &field : fields) {
-    MFVariable &input = *get_field_variable(field, unique_variables).mf_variable;
+    MFVariable &input = get_field_variable(field, unique_variables);
     builder.add_output_parameter(input);
   }
 
@@ -246,18 +198,18 @@ static void gather_inputs(const Span<GField> fields,
   while (!fields_to_visit.is_empty()) {
     const GField &field = *fields_to_visit.pop();
     if (field.is_input()) {
-      const FieldInput &input = field.input();
-      const FieldVariable &variable = get_field_variable(field, unique_variables);
-      if (!computed_inputs.contains(variable.mf_variable)) {
+      const FieldInput &input = dynamic_cast<const FieldInput &>(field.source());
+      const MFVariable &variable = get_field_variable(field, unique_variables);
+      if (!computed_inputs.contains(&variable)) {
         GVArrayPtr data = input.get_varray_generic_context(mask);
-        computed_inputs.add_new(variable.mf_variable);
+        computed_inputs.add_new(&variable);
         params.add_readonly_single_input(*data, input.debug_name());
         r_inputs.append(std::move(data));
       }
     }
     else {
-      const FieldFunction &function = field.function();
-      for (const GField &input_field : function.inputs()) {
+      const FieldOperation &operation = dynamic_cast<const FieldOperation &>(field.source());
+      for (const GField &input_field : operation.inputs()) {
         fields_to_visit.push(&input_field);
       }
     }
@@ -309,8 +261,9 @@ void evaluate_fields(const Span<GField> fields,
   Vector<GMutableSpan> non_input_outputs{outputs};
   for (int i = fields.size() - 1; i >= 0; i--) {
     if (non_input_fields[i].is_input()) {
-      non_input_fields[i].input().get_varray_generic_context(mask)->materialize(mask,
-                                                                                outputs[i].data());
+      dynamic_cast<const FieldInput &>(non_input_fields[i].source())
+          .get_varray_generic_context(mask)
+          ->materialize(mask, outputs[i].data());
 
       non_input_fields.remove_and_reorder(i);
       non_input_outputs.remove_and_reorder(i);
diff --git a/source/blender/functions/tests/FN_field_test.cc b/source/blender/functions/tests/FN_field_test.cc
index 291fcd929b6..f11a0d690c2 100644
--- a/source/blender/functions/tests/FN_field_test.cc
+++ b/source/blender/functions/tests/FN_field_test.cc
@@ -10,9 +10,9 @@ namespace blender::fn::tests {
 
 TEST(field, ConstantFunction)
 {
-  /* TODO: Figure out how to not use another "FieldFunction(" inside of std::make_shared. */
-  GField constant_field{std::make_shared<FieldFunction>(
-                            FieldFunction(std::make_unique<CustomMF_Constant<int>>(10), {})),
+  /* TODO: Figure out how to not use another "FieldOperation(" inside of std::make_shared. */
+  GField constant_field{std::make_shared<FieldOperation>(
+                            FieldOperation(std::make_unique<CustomMF_Constant<int>>(10), {})),
                         0};
 
   Array<int> result(4);
@@ -90,8 +90,8 @@ TEST(field, InputAndFunction)
 
   std::unique_ptr<MultiFunction> add_fn = std::make_unique<CustomMF_SI_SI_SO<int, int, int>>(
       "add", [](int a, int b) { return a + b; });
-  GField output_field{std::make_shared<FieldFunction>(
-                          FieldFunction(std::move(add_fn), {index_field, index_field})),
+  GField output_field{std::make_shared<FieldOperation>(
+                          FieldOperation(std::move(add_fn), {index_field, index_field})),
                       0};
 
   Array<int> result(10);
@@ -109,14 +109,14 @@ TEST(field, TwoFunctions)
 
   std::unique_ptr<MultiFunction> add_fn = std::make_unique<CustomMF_SI_SI_SO<int, int, int>>(
       "add", [](int a, int b) { return a + b; });
-  GField add_field{std::make_shared<FieldFunction>(
-                       FieldFunction(std::move(add_fn), {index_field, index_field})),
+  GField add_field{std::make_shared<FieldOperation>(
+                       FieldOperation(std::move(add_fn), {index_field, index_field})),
                    0};
 
   std::unique_ptr<MultiFunction> add_10_fn = std::make_unique<CustomMF_SI_SO<int, int>>(
       "add_10", [](int a) { return a + 10; });
   GField result_field{
-      std::make_shared<FieldFunction>(FieldFunction(std::move(add_10_fn), {add_field})), 0};
+      std::make_shared<FieldOperation>(FieldOperation(std::move(add_10_fn), {add_field})), 0};
 
   Array<int> result(10);
   GMutableSpan result_generic(result.as_mutable_span());
@@ -162,7 +162,7 @@ TEST(field, FunctionTwoOutputs)
   GField index_field_1{std::make_shared<IndexFieldInput>()};
   GField index_field_2{std::make_shared<IndexFieldInput>()};
 
-  std::shared_ptr<FieldFunction> fn = std::make_shared<FieldFunction>(FieldFunction(
+  std::shared_ptr<FieldOperation> fn = std::make_shared<FieldOperation>(FieldOperation(
       std::make_unique<TwoOutputFunction>("SI_SI_SO_SO"), {index_field_1, index_field_2}));
 
   GField result_field_1{fn, 0};
@@ -188,7 +188,7 @@ TEST(field, TwoFunctionsTwoOutputs)
 {
   GField index_field{std::make_shared<IndexFieldInput>()};
 
-  std::shared_ptr<FieldFunction> fn = std::make_shared<FieldFunction>(FieldFunction(
+  std::shared_ptr<FieldOperation> fn = std::make_shared<FieldOperation>(FieldOperation(
       std::make_unique<TwoOutputFunction>("SI_SI_SO_SO"), {index_field, index_field}));
 
   GField result_field_1{fn, 0};
@@ -197,7 +197,7 @@ TEST(field, TwoFunctionsTwoOutputs)
   std::unique_ptr<MultiFunction> add_10_fn = std::make_unique<CustomMF_SI_SO<int, int>>(
       "add_10", [](int a) { return a + 10; });
   GField result_field_2{
-      std::make_shared<FieldFunction>(FieldFunction(std::move(add_10_fn), {intermediate_field})),
+      std::make_shared<FieldOperation>(FieldOperation(std::move(add_10_fn), {intermediate_field})),
       0};
 
   Array<int> result_1(10);
diff --git a/source/blender/modifiers/intern/MOD_nodes_evaluator.cc b/source/blender/modifiers/intern/MOD_nodes_evaluator.cc
index 0251163d9a7..4f21924619b 100644
--- a/source/blender/modifiers/intern/MOD_nodes_evaluator.cc
+++ b/source/blender/modifiers/intern/MOD_nodes_evaluator.cc
@@ -880,7 +880,7 @@ class GeometryNodesEvaluator {
       input_fields.append(std::move(*(GField *)single_value.value));
     }
 
-    auto field_fn = std::make_shared<fn::FieldFunction>(fn, std::move(input_fields));
+    auto operation = std::make_shared<fn::FieldOperation>(fn, std::move(input_fields));
 
     /* Forward outputs. */
     int output_index = 0;
@@ -892,7 +892,7 @@ class GeometryNodesEvaluator {
       OutputState &output_state = node_state.outputs[i];
       const DOutputSocket socket{node.context(), &socket_ref};
       const CPPType *cpp_type = get_socket_cpp_type(socket_ref);
-      GField &field = *allocator.construct<GField>(field_fn, output_index).release();
+      GField &field = *allocator.construct<GField>(operation, output_index).release();
       this->forward_output(socket, {cpp_type, &field});
       output_state.has_been_computed = true;
       output_index++;
@@ -1392,8 +1392,8 @@ class GeometryNodesEvaluator {
         const MultiFunction &fn = *conversions_.get_conversion_multi_function(
             MFDataType::ForSingle(from_base_type), MFDataType::ForSingle(to_base_type));
         const GField &from_field = *(const GField *)from_value;
-        auto field_fn = std::make_shared<fn::FieldFunction>(fn, Vector<GField>{from_field});
-        new (to_value) GField(std::move(field_fn), 0);
+        auto operation = std::make_shared<fn::FieldOperation>(fn, Vector<GField>{from_field});
+        new (to_value) GField(std::move(operation), 0);
         return;
       }
     }
@@ -1413,8 +1413,8 @@ class GeometryNodesEvaluator {
       const CPPType &base_type = field_cpp_type->field_type();
       auto constant_fn = std::make_unique<fn::CustomMF_GenericConstant>(base_type,
                                                                         base_type.default_value());
-      auto field_fn = std::make_shared<fn::FieldFunction>(std::move(constant_fn));
-      new (r_value) GField(std::move(field_fn), 0);
+      auto operation = std::make_shared<fn::FieldOperation>(std::move(constant_fn));
+      new (r_value) GField(std::move(operation), 0);
       return;
     }
     type.copy_construct(type.default_value(), r_value);
diff --git a/source/blender/nodes/NOD_geometry_exec.hh b/source/blender/nodes/NOD_geometry_exec.hh
index 6c49e71d7ab..6f0d12f3dba 100644
--- a/source/blender/nodes/NOD_geometry_exec.hh
+++ b/source/blender/nodes/NOD_geometry_exec.hh
@@ -40,8 +40,8 @@ using bke::ReadAttributeLookup;
 using bke::WriteAttributeLookup;
 using fn::CPPType;
 using fn::Field;
-using fn::FieldFunction;
 using fn::FieldInput;
+using fn::FieldOperation;
 using fn::GField;
 using fn::GMutablePointer;
 using fn::GMutableSpan;