improvements

3 files changed, 162 insertions, 62 deletions

diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index 3fcc0db017f..c140bbc49e4 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -96,6 +96,12 @@
 
 #define NODE_DEFAULT_MAX_WIDTH 700
 
+using blender::Array;
+using blender::Span;
+using blender::Stack;
+using blender::Vector;
+using namespace blender::nodes::node_tree_ref_types;
+
 /* Fallback types for undefined tree, nodes, sockets */
 static bNodeTreeType NodeTreeTypeUndefined;
 bNodeType NodeTypeUndefined;
@@ -3950,6 +3956,9 @@ NodeDeclarationHandle *nodeDeclarationEnsure(bNodeTree *UNUSED(ntree), bNode *no
   if (node->typeinfo->declare == nullptr) {
     return nullptr;
   }
+  if (node->declaration != nullptr) {
+    return node->declaration;
+  }
 
   node->declaration = new blender::nodes::NodeDeclaration();
   blender::nodes::NodeDeclarationBuilder builder{*node->declaration};
@@ -4469,103 +4478,187 @@ void ntreeUpdateAllUsers(Main *main, ID *id)
   }
 }
 
-static void update_socket_shapes_for_fields(bNodeTree &btree)
+static bool is_field_socket_type(eNodeSocketDatatype type)
 {
-  using namespace blender;
-  using namespace blender::nodes;
-  if (btree.type != NTREE_GEOMETRY) {
-    return;
+  return ELEM(type, SOCK_FLOAT, SOCK_INT, SOCK_BOOLEAN, SOCK_VECTOR, SOCK_RGBA);
+}
+
+static bool sockets_have_links(blender::Span<const SocketRef *> sockets)
+{
+  for (const SocketRef *socket : sockets) {
+    if (!socket->directly_linked_links().is_empty()) {
+      return true;
+    }
   }
+  return false;
+}
 
-  NodeTreeRef tree{&btree};
-  Vector<const NodeRef *> input_nodes;
-  Vector<const NodeRef *> output_nodes;
+static Vector<const NodeRef *> toposort_nodes(const NodeTreeRef &tree, bool left_to_right = true)
+{
+  Vector<const NodeRef *> toposort;
+  toposort.reserve(tree.nodes().size());
+  Array<bool> node_is_pushed_by_id(tree.nodes().size(), false);
+  Stack<const NodeRef *> nodes_to_check;
 
   for (const NodeRef *node : tree.nodes()) {
-    if (node->inputs().is_empty()) {
-      input_nodes.append(node);
+    if (!sockets_have_links(node->inputs_or_outputs(!left_to_right))) {
+      node_is_pushed_by_id[node->id()] = true;
+      nodes_to_check.push(node);
     }
-    if (node->outputs().is_empty()) {
-      output_nodes.append(node);
+  }
+
+  while (!nodes_to_check.is_empty()) {
+    const NodeRef *node = nodes_to_check.pop();
+    toposort.append(node);
+
+    for (const SocketRef *input_socket : node->inputs_or_outputs(left_to_right)) {
+      for (const SocketRef *linked_socket : input_socket->directly_linked_sockets()) {
+        const NodeRef &linked_node = linked_socket->node();
+        const int linked_node_id = linked_node.id();
+        if (!node_is_pushed_by_id[linked_node_id]) {
+          node_is_pushed_by_id[linked_node_id] = true;
+          nodes_to_check.push(&linked_node);
+        }
+      }
     }
   }
 
-  Array<int> field_level_by_node_id(tree.nodes().size(), 0);
-  Array<int> field_level_by_socket_id(tree.sockets().size(), 0);
-  Array<bool> node_is_enqueued_by_id(tree.nodes().size(), false);
-  std::queue<const NodeRef *> nodes_to_check;
+  toposort.as_mutable_span().reverse();
+  return toposort;
+}
 
-  for (const NodeRef *node : input_nodes) {
-    nodes_to_check.push(node);
-    node_is_enqueued_by_id[node->id()] = true;
+static void update_socket_shapes_for_fields(bNodeTree &btree)
+{
+  using namespace blender;
+  using namespace blender::nodes;
+  if (btree.type != NTREE_GEOMETRY) {
+    return;
   }
 
-  while (!nodes_to_check.empty()) {
-    const NodeRef &node = *nodes_to_check.front();
-    nodes_to_check.pop();
-    NodeDeclaration *node_decl = nodeDeclarationEnsure(&btree, node.bnode());
+  NodeTreeRef tree{&btree};
+  Vector<const NodeRef *> toposort_left_to_right = toposort_nodes(tree, true);
+  Vector<const NodeRef *> toposort_right_to_left = toposort_nodes(tree, false);
+
+  struct SocketFieldState {
+    bool requires_single = false;
+    bool is_field = false;
+    bool is_field_source = false;
+  };
+
+  Array<SocketFieldState> field_state_by_socket_id(tree.sockets().size());
+
+  auto check_if_node_is_fixed = [](const NodeRef &node) {
+    const StringRef node_idname = node.idname();
+    return node_idname.startswith("GeometryNode");
+  };
 
-    int max_node_field_level = 0;
-    StringRef node_idname = node.idname();
-    if (node_idname.startswith("ShaderNode") || node_idname.startswith("FunctionNode")) {
-      max_node_field_level = 1;
+  auto has_field_dependency = [&](const OutputSocketRef &output_socket,
+                                  const InputSocketRef &input_socket) {
+    if (!is_field_socket_type((eNodeSocketDatatype)output_socket.typeinfo()->type)) {
+      return false;
     }
+    if (!is_field_socket_type((eNodeSocketDatatype)input_socket.typeinfo()->type)) {
+      return false;
+    }
+    return true;
+  };
+
+  for (const NodeRef *node : toposort_left_to_right) {
+    NodeDeclaration *node_decl = nodeDeclarationEnsure(&btree, node->bnode());
+
+    const bool node_is_fixed = check_if_node_is_fixed(*node);
 
-    int node_field_level = 0;
-    for (const SocketRef *input_socket : node.inputs()) {
-      const int input_index = input_socket->index();
-      int normal_input_field_level = 0;
+    auto input_supports_field = [&](const InputSocketRef &socket) {
+      if (!node_is_fixed) {
+        return true;
+      }
       if (node_decl != nullptr) {
-        const SocketDeclaration &socket_decl = *node_decl->inputs()[input_index];
-        normal_input_field_level = socket_decl.is_field();
+        const SocketDeclaration &socket_decl = *node_decl->inputs()[socket.index()];
+        return socket_decl.is_field();
       }
-
-      int input_field_level = normal_input_field_level;
-      for (const SocketRef *origin_socket : input_socket->directly_linked_sockets()) {
-        const int origin_field_level = field_level_by_socket_id[origin_socket->id()];
-        input_field_level = std::max(input_field_level, origin_field_level);
+      return false;
+    };
+
+    for (const InputSocketRef *input_socket : node->inputs()) {
+      bool is_field = false;
+      if (input_supports_field(*input_socket)) {
+        for (const SocketRef *origin_socket : input_socket->directly_linked_sockets()) {
+          is_field |= field_state_by_socket_id[origin_socket->id()].is_field;
+        }
       }
-
-      const int field_level_increase = input_field_level - normal_input_field_level;
-      node_field_level = std::min(max_node_field_level,
-                                  std::max(node_field_level, field_level_increase));
-
-      field_level_by_socket_id[input_socket->id()] = input_field_level;
+      field_state_by_socket_id[input_socket->id()].is_field = is_field;
     }
 
-    for (const SocketRef *output_socket : node.outputs()) {
+    for (const OutputSocketRef *output_socket : node->outputs()) {
       const int output_index = output_socket->index();
-      int normal_output_field_level = 0;
+
+      bool output_is_field = false;
+
       if (node_decl != nullptr) {
         const SocketDeclaration &socket_decl = *node_decl->outputs()[output_index];
-        normal_output_field_level = socket_decl.is_field();
+        output_is_field = socket_decl.is_field();
+        field_state_by_socket_id[output_index].is_field_source = true;
       }
 
-      const int output_field_level = normal_output_field_level + node_field_level;
-      field_level_by_socket_id[output_socket->id()] = output_field_level;
-
-      for (const SocketRef *target_socket : output_socket->directly_linked_sockets()) {
-        const NodeRef &target_node = target_socket->node();
-        if (!node_is_enqueued_by_id[target_node.id()]) {
-          nodes_to_check.push(&target_node);
-          node_is_enqueued_by_id[target_node.id()] = true;
+      if (!output_is_field) {
+        for (const InputSocketRef *input_socket : node->inputs()) {
+          if (has_field_dependency(*output_socket, *input_socket)) {
+            output_is_field |= field_state_by_socket_id[input_socket->id()].is_field;
+          }
         }
       }
+      field_state_by_socket_id[output_socket->id()].is_field = output_is_field;
     }
+  }
+
+  for (const NodeRef *node : toposort_right_to_left) {
+    NodeDeclaration *node_decl = nodeDeclarationEnsure(&btree, node->bnode());
 
-    std::cout << node.name() << ": " << node_field_level << "\n";
+    const bool node_is_fixed = check_if_node_is_fixed(*node);
+
+    for (const OutputSocketRef *output_socket : node->outputs()) {
+      bool requires_single = false;
+      if (node_decl != nullptr && node_is_fixed) {
+        const SocketDeclaration &socket_decl = *node_decl->outputs()[output_socket->index()];
+        requires_single = !socket_decl.is_field();
+      }
+      for (const InputSocketRef *target_socket : output_socket->directly_linked_sockets()) {
+        requires_single |= field_state_by_socket_id[target_socket->id()].requires_single;
+      }
+      field_state_by_socket_id[output_socket->id()].requires_single = requires_single;
+    }
+    for (const InputSocketRef *input_socket : node->inputs()) {
+      bool requires_single = false;
+      if (node_decl != nullptr && node_is_fixed) {
+        const SocketDeclaration &socket_decl = *node_decl->inputs()[input_socket->index()];
+        requires_single = !socket_decl.is_field();
+      }
+      if (!requires_single) {
+        for (const OutputSocketRef *output_socket : node->outputs()) {
+          if (has_field_dependency(*output_socket, *input_socket)) {
+            requires_single |= field_state_by_socket_id[output_socket->id()].requires_single;
+          }
+        }
+      }
+      field_state_by_socket_id[input_socket->id()].requires_single = requires_single;
+    }
   }
-  std::cout << "\n";
 
   for (const SocketRef *socket : tree.sockets()) {
     bNodeSocket *bsocket = socket->bsocket();
-    const int field_level = field_level_by_socket_id[socket->id()];
-    if (field_level == 0) {
-      bsocket->display_shape = SOCK_DISPLAY_SHAPE_CIRCLE;
+    SocketFieldState &state = field_state_by_socket_id[socket->id()];
+    if (state.is_field_source) {
+      bsocket->display_shape = SOCK_DISPLAY_SHAPE_DIAMOND;
+    }
+    else if (state.requires_single) {
+      bsocket->display_shape = SOCK_DISPLAY_SHAPE_CIRCLE_DOT;
     }
-    else if (field_level == 1) {
+    else if (state.is_field) {
       bsocket->display_shape = SOCK_DISPLAY_SHAPE_DIAMOND;
     }
+    else {
+      bsocket->display_shape = SOCK_DISPLAY_SHAPE_CIRCLE;
+    }
   }
 }
 
diff --git a/source/blender/nodes/NOD_node_tree_ref.hh b/source/blender/nodes/NOD_node_tree_ref.hh
index 4f2565cbbaf..363e47ccabc 100644
--- a/source/blender/nodes/NOD_node_tree_ref.hh
+++ b/source/blender/nodes/NOD_node_tree_ref.hh
@@ -181,6 +181,7 @@ class NodeRef : NonCopyable, NonMovable {
 
   Span<const InputSocketRef *> inputs() const;
   Span<const OutputSocketRef *> outputs() const;
+  Span<const SocketRef *> inputs_or_outputs(bool get_inputs) const;
   Span<const InternalLinkRef *> internal_links() const;
 
   const InputSocketRef &input(int index) const;
@@ -496,6 +497,12 @@ inline Span<const OutputSocketRef *> NodeRef::outputs() const
   return outputs_;
 }
 
+inline Span<const SocketRef *> NodeRef::inputs_or_outputs(bool get_inputs) const
+{
+  return get_inputs ? inputs_.as_span().cast<const SocketRef *>() :
+                      outputs_.as_span().cast<const SocketRef *>();
+}
+
 inline Span<const InternalLinkRef *> NodeRef::internal_links() const
 {
   return internal_links_;
diff --git a/source/blender/nodes/geometry/nodes/node_geo_set_position.cc b/source/blender/nodes/geometry/nodes/node_geo_set_position.cc
index a1bc061f86e..0bd11c75ff8 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_set_position.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_set_position.cc
@@ -23,7 +23,7 @@ namespace blender::nodes {
 static void geo_node_set_position_declare(NodeDeclarationBuilder &b)
 {
   b.add_input<decl::Geometry>("Geometry");
-  b.add_input<decl::Vector>("Position");
+  b.add_input<decl::Vector>("Position").is_field();
   b.add_input<decl::Bool>("Selection").default_value(true).hide_value().is_field();
   b.add_output<decl::Geometry>("Geometry");
 }