Geometry Nodes: Domain Size Node

The Domain Size node has a single geometry input and a selection for
the component type. Based on the component chosen, outputs containing
single values for the related domains are shown.

Mesh:
  - Point Count
  - Edge Count
  - Face Count
  - Face Corner Count
Curve:
  - Point Count
  - Spline Count
Point Cloud:
  - Point Count
Instances:
  - Instance Count

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

9 files changed, 185 insertions, 25 deletions

diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index 5ed1b512fd5..48f0aaeeb8d 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -1554,6 +1554,7 @@ int ntreeTexExecTree(struct bNodeTree *ntree,
 #define GEO_NODE_VOLUME_TO_MESH 1133
 #define GEO_NODE_INPUT_ID 1134
 #define GEO_NODE_SET_ID 1135
+#define GEO_NODE_ATTRIBUTE_DOMAIN_SIZE 1136
 
 /** \} */
 
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index 7d5ba4189a3..b414c0d37d0 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -5853,6 +5853,7 @@ static void registerGeometryNodes()
   register_node_type_geo_attribute_compare();
   register_node_type_geo_attribute_convert();
   register_node_type_geo_attribute_curve_map();
+  register_node_type_geo_attribute_domain_size();
   register_node_type_geo_attribute_fill();
   register_node_type_geo_attribute_map_range();
   register_node_type_geo_attribute_math();
diff --git a/source/blender/makesrna/RNA_enum_items.h b/source/blender/makesrna/RNA_enum_items.h
index fb18802483d..78dd977af20 100644
--- a/source/blender/makesrna/RNA_enum_items.h
+++ b/source/blender/makesrna/RNA_enum_items.h
@@ -210,6 +210,7 @@ DEF_ENUM(rna_enum_attribute_type_with_auto_items)
 DEF_ENUM(rna_enum_attribute_domain_items)
 DEF_ENUM(rna_enum_attribute_domain_without_corner_items)
 DEF_ENUM(rna_enum_attribute_domain_with_auto_items)
+DEF_ENUM(rna_enum_geometry_component_type_items)
 
 DEF_ENUM(rna_enum_volume_grid_data_type_items)
 
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index aef3823338b..2db517d6b58 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/source/blender/makesrna/intern/rna_nodetree.c
@@ -39,6 +39,7 @@
 #include "BKE_animsys.h"
 #include "BKE_attribute.h"
 #include "BKE_cryptomatte.h"
+#include "BKE_geometry_set.h"
 #include "BKE_image.h"
 #include "BKE_node.h"
 #include "BKE_texture.h"
@@ -9166,6 +9167,16 @@ static void def_geo_boolean(StructRNA *srna)
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_socket_update");
 }
 
+static void def_geo_attribute_domain_size(StructRNA *srna)
+{
+  PropertyRNA *prop = RNA_def_property(srna, "component", PROP_ENUM, PROP_NONE);
+  RNA_def_property_enum_sdna(prop, NULL, "custom1");
+  RNA_def_property_enum_items(prop, rna_enum_geometry_component_type_items);
+  RNA_def_property_enum_default(prop, GEO_COMPONENT_TYPE_MESH);
+  RNA_def_property_ui_text(prop, "Component", "");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_socket_update");
+}
+
 static void def_geo_curve_primitive_bezier_segment(StructRNA *srna)
 {
   static const EnumPropertyItem mode_items[] = {
diff --git a/source/blender/makesrna/intern/rna_space.c b/source/blender/makesrna/intern/rna_space.c
index 8f04524c5c1..bf7a912e9e6 100644
--- a/source/blender/makesrna/intern/rna_space.c
+++ b/source/blender/makesrna/intern/rna_space.c
@@ -73,6 +73,30 @@
 
 #include "RNA_enum_types.h"
 
+const EnumPropertyItem rna_enum_geometry_component_type_items[] = {
+    {GEO_COMPONENT_TYPE_MESH,
+     "MESH",
+     ICON_MESH_DATA,
+     "Mesh",
+     "Mesh component containing point, corner, edge and face data"},
+    {GEO_COMPONENT_TYPE_POINT_CLOUD,
+     "POINTCLOUD",
+     ICON_POINTCLOUD_DATA,
+     "Point Cloud",
+     "Point cloud component containing only point data"},
+    {GEO_COMPONENT_TYPE_CURVE,
+     "CURVE",
+     ICON_CURVE_DATA,
+     "Curve",
+     "Curve component containing spline and control point data"},
+    {GEO_COMPONENT_TYPE_INSTANCES,
+     "INSTANCES",
+     ICON_EMPTY_AXIS,
+     "Instances",
+     "Instances of objects or collections"},
+    {0, NULL, 0, NULL, NULL},
+};
+
 const EnumPropertyItem rna_enum_space_type_items[] = {
     /* empty must be here for python, is skipped for UI */
     {SPACE_EMPTY, "EMPTY", ICON_NONE, "Empty", ""},
@@ -7832,30 +7856,6 @@ static void rna_def_space_spreadsheet(BlenderRNA *brna)
   PropertyRNA *prop;
   StructRNA *srna;
 
-  static const EnumPropertyItem geometry_component_type_items[] = {
-      {GEO_COMPONENT_TYPE_MESH,
-       "MESH",
-       ICON_MESH_DATA,
-       "Mesh",
-       "Mesh component containing point, corner, edge and face data"},
-      {GEO_COMPONENT_TYPE_POINT_CLOUD,
-       "POINTCLOUD",
-       ICON_POINTCLOUD_DATA,
-       "Point Cloud",
-       "Point cloud component containing only point data"},
-      {GEO_COMPONENT_TYPE_CURVE,
-       "CURVE",
-       ICON_CURVE_DATA,
-       "Curve",
-       "Curve component containing spline and control point data"},
-      {GEO_COMPONENT_TYPE_INSTANCES,
-       "INSTANCES",
-       ICON_EMPTY_AXIS,
-       "Instances",
-       "Instances of objects or collections"},
-      {0, NULL, 0, NULL, NULL},
-  };
-
   static const EnumPropertyItem object_eval_state_items[] = {
       {SPREADSHEET_OBJECT_EVAL_STATE_EVALUATED,
        "EVALUATED",
@@ -7914,7 +7914,7 @@ static void rna_def_space_spreadsheet(BlenderRNA *brna)
   RNA_def_property_update(prop, NC_SPACE | ND_SPACE_SPREADSHEET, NULL);
 
   prop = RNA_def_property(srna, "geometry_component_type", PROP_ENUM, PROP_NONE);
-  RNA_def_property_enum_items(prop, geometry_component_type_items);
+  RNA_def_property_enum_items(prop, rna_enum_geometry_component_type_items);
   RNA_def_property_ui_text(
       prop, "Geometry Component", "Part of the geometry to display data from");
   RNA_def_property_update(prop,
diff --git a/source/blender/nodes/NOD_geometry.h b/source/blender/nodes/NOD_geometry.h
index ea3458af065..44cfd89423d 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/source/blender/nodes/NOD_geometry.h
@@ -57,6 +57,7 @@ void register_node_type_geo_attribute_combine_xyz(void);
 void register_node_type_geo_attribute_compare(void);
 void register_node_type_geo_attribute_convert(void);
 void register_node_type_geo_attribute_curve_map(void);
+void register_node_type_geo_attribute_domain_size(void);
 void register_node_type_geo_attribute_fill(void);
 void register_node_type_geo_attribute_map_range(void);
 void register_node_type_geo_attribute_math(void);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index 20ad4d359f1..9f8b99cdc97 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/source/blender/nodes/NOD_static_types.h
@@ -325,6 +325,7 @@ DefNode(GeometryNode, GEO_NODE_LEGACY_VOLUME_TO_MESH, def_geo_volume_to_mesh, "L
 
 DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_REMOVE, 0, "ATTRIBUTE_REMOVE", AttributeRemove, "Attribute Remove", "")
 DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_STATISTIC, def_geo_attribute_statistic, "ATTRIBUTE_STATISTIC", AttributeStatistic, "Attribute Statistic", "")
+DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_DOMAIN_SIZE, def_geo_attribute_domain_size, "ATTRIBUTE_DOMAIN_SIZE", AttributeDomainSize, "Domain Size", "")
 DefNode(GeometryNode, GEO_NODE_BOUNDING_BOX, 0, "BOUNDING_BOX", BoundBox, "Bounding Box", "")
 DefNode(GeometryNode, GEO_NODE_CAPTURE_ATTRIBUTE, def_geo_attribute_capture, "CAPTURE_ATTRIBUTE", CaptureAttribute, "Capture Attribute", "")
 DefNode(GeometryNode, GEO_NODE_COLLECTION_INFO, def_geo_collection_info, "COLLECTION_INFO", CollectionInfo, "Collection Info", "")
diff --git a/source/blender/nodes/geometry/CMakeLists.txt b/source/blender/nodes/geometry/CMakeLists.txt
index 576c5249e98..7f30614dbbc 100644
--- a/source/blender/nodes/geometry/CMakeLists.txt
+++ b/source/blender/nodes/geometry/CMakeLists.txt
@@ -81,6 +81,7 @@ set(SRC
   nodes/legacy/node_geo_legacy_volume_to_mesh.cc
 
   nodes/node_geo_attribute_capture.cc
+  nodes/node_geo_attribute_domain_size.cc
   nodes/node_geo_attribute_remove.cc
   nodes/node_geo_attribute_statistic.cc
   nodes/node_geo_boolean.cc
diff --git a/source/blender/nodes/geometry/nodes/node_geo_attribute_domain_size.cc b/source/blender/nodes/geometry/nodes/node_geo_attribute_domain_size.cc
new file mode 100644
index 00000000000..082ea87853e
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/node_geo_attribute_domain_size.cc
@@ -0,0 +1,143 @@
+/*
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ */
+
+#include "UI_interface.h"
+#include "UI_resources.h"
+
+#include "node_geometry_util.hh"
+
+namespace blender::nodes::node_geo_attribute_domain_size_cc {
+
+static void node_declare(NodeDeclarationBuilder &b)
+{
+  b.add_input<decl::Geometry>("Geometry");
+  b.add_output<decl::Int>("Point Count");
+  b.add_output<decl::Int>("Edge Count");
+  b.add_output<decl::Int>("Face Count");
+  b.add_output<decl::Int>("Face Corner Count");
+  b.add_output<decl::Int>("Spline Count");
+  b.add_output<decl::Int>("Instance Count");
+}
+
+static void node_layout(uiLayout *layout, bContext *UNUSED(C), PointerRNA *ptr)
+{
+  uiItemR(layout, ptr, "component", 0, "", ICON_NONE);
+}
+
+static void node_init(bNodeTree *UNUSED(tree), bNode *node)
+{
+  node->custom1 = GEO_COMPONENT_TYPE_MESH;
+}
+
+static void node_update(bNodeTree *ntree, bNode *node)
+{
+  bNodeSocket *point_socket = (bNodeSocket *)node->outputs.first;
+  bNodeSocket *edge_socket = point_socket->next;
+  bNodeSocket *face_socket = edge_socket->next;
+  bNodeSocket *face_corner_socket = face_socket->next;
+  bNodeSocket *spline_socket = face_corner_socket->next;
+  bNodeSocket *instances_socket = spline_socket->next;
+
+  nodeSetSocketAvailability(ntree,
+                            point_socket,
+                            ELEM(node->custom1,
+                                 GEO_COMPONENT_TYPE_MESH,
+                                 GEO_COMPONENT_TYPE_CURVE,
+                                 GEO_COMPONENT_TYPE_POINT_CLOUD));
+  nodeSetSocketAvailability(ntree, edge_socket, node->custom1 == GEO_COMPONENT_TYPE_MESH);
+  nodeSetSocketAvailability(ntree, face_socket, node->custom1 == GEO_COMPONENT_TYPE_MESH);
+  nodeSetSocketAvailability(ntree, face_corner_socket, node->custom1 == GEO_COMPONENT_TYPE_MESH);
+  nodeSetSocketAvailability(ntree, spline_socket, node->custom1 == GEO_COMPONENT_TYPE_CURVE);
+  nodeSetSocketAvailability(
+      ntree, instances_socket, node->custom1 == GEO_COMPONENT_TYPE_INSTANCES);
+}
+
+static void node_geo_exec(GeoNodeExecParams params)
+{
+  GeometryComponentType component = (GeometryComponentType)params.node().custom1;
+  GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
+
+  switch (component) {
+    case GEO_COMPONENT_TYPE_MESH: {
+      if (geometry_set.has_mesh()) {
+        const MeshComponent *component = geometry_set.get_component_for_read<MeshComponent>();
+        params.set_output("Point Count", component->attribute_domain_size(ATTR_DOMAIN_POINT));
+        params.set_output("Edge Count", component->attribute_domain_size(ATTR_DOMAIN_EDGE));
+        params.set_output("Face Count", component->attribute_domain_size(ATTR_DOMAIN_FACE));
+        params.set_output("Face Corner Count",
+                          component->attribute_domain_size(ATTR_DOMAIN_CORNER));
+      }
+      else {
+        params.set_default_remaining_outputs();
+      }
+      break;
+    }
+    case GEO_COMPONENT_TYPE_CURVE: {
+      if (geometry_set.has_curve()) {
+        const CurveComponent *component = geometry_set.get_component_for_read<CurveComponent>();
+        params.set_output("Point Count", component->attribute_domain_size(ATTR_DOMAIN_POINT));
+        params.set_output("Spline Count", component->attribute_domain_size(ATTR_DOMAIN_CURVE));
+      }
+      else {
+        params.set_default_remaining_outputs();
+      }
+      break;
+    }
+    case GEO_COMPONENT_TYPE_POINT_CLOUD: {
+      if (geometry_set.has_pointcloud()) {
+        const PointCloudComponent *component =
+            geometry_set.get_component_for_read<PointCloudComponent>();
+        params.set_output("Point Count", component->attribute_domain_size(ATTR_DOMAIN_POINT));
+      }
+      else {
+        params.set_default_remaining_outputs();
+      }
+      break;
+    }
+    case GEO_COMPONENT_TYPE_INSTANCES: {
+      if (geometry_set.has_instances()) {
+        const InstancesComponent *component =
+            geometry_set.get_component_for_read<InstancesComponent>();
+        params.set_output("Instance Count",
+                          component->attribute_domain_size(ATTR_DOMAIN_INSTANCE));
+      }
+      else {
+        params.set_default_remaining_outputs();
+      }
+      break;
+    }
+    default:
+      BLI_assert_unreachable();
+  }
+}
+
+}  // namespace blender::nodes::node_geo_attribute_domain_size_cc
+
+void register_node_type_geo_attribute_domain_size()
+{
+  namespace file_ns = blender::nodes::node_geo_attribute_domain_size_cc;
+
+  static bNodeType ntype;
+  geo_node_type_base(
+      &ntype, GEO_NODE_ATTRIBUTE_DOMAIN_SIZE, "Domain Size", NODE_CLASS_ATTRIBUTE, 0);
+  ntype.geometry_node_execute = file_ns::node_geo_exec;
+  ntype.declare = file_ns::node_declare;
+  ntype.draw_buttons = file_ns::node_layout;
+  node_type_init(&ntype, file_ns::node_init);
+  ntype.updatefunc = file_ns::node_update;
+
+  nodeRegisterType(&ntype);
+}