Geometry Nodes: Endpoint Selection Nodes

The Endpoint Selection node allows for the Selection of an aribitrary
number of endpoints from each spline in a curve. The start and end
inputs are evaluated on the spline domain. The result is outputted
as a boolean field on the point domain.

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

7 files changed, 157 insertions, 3 deletions

diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index 6dcb35de3af..d33c5e9940c 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -1539,6 +1539,7 @@ int ntreeTexExecTree(struct bNodeTree *ntree,
 #define GEO_NODE_MESH_TO_CURVE 1124
 #define GEO_NODE_TRANSFER_ATTRIBUTE 1125
 #define GEO_NODE_SUBDIVISION_SURFACE 1126
+#define GEO_NODE_CURVE_ENDPOINT_SELECTION 1127
 
 /** \} */
 
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index d29762e2af4..145a40d30cc 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -5724,6 +5724,7 @@ static void registerGeometryNodes()
   register_node_type_geo_legacy_mesh_to_curve();
   register_node_type_geo_legacy_points_to_volume();
   register_node_type_geo_legacy_select_by_material();
+  register_node_type_geo_legacy_curve_endpoints();
   register_node_type_geo_legacy_curve_spline_type();
   register_node_type_geo_legacy_curve_reverse();
   register_node_type_geo_legacy_select_by_handle_type();
@@ -5752,7 +5753,7 @@ static void registerGeometryNodes()
   register_node_type_geo_bounding_box();
   register_node_type_geo_collection_info();
   register_node_type_geo_convex_hull();
-  register_node_type_geo_curve_endpoints();
+  register_node_type_geo_curve_endpoint_selection();
   register_node_type_geo_curve_fill();
   register_node_type_geo_curve_fillet();
   register_node_type_geo_curve_handle_type_selection();
diff --git a/source/blender/nodes/CMakeLists.txt b/source/blender/nodes/CMakeLists.txt
index a202d247e60..35cd1c8c5a9 100644
--- a/source/blender/nodes/CMakeLists.txt
+++ b/source/blender/nodes/CMakeLists.txt
@@ -198,6 +198,7 @@ set(SRC
   geometry/nodes/node_geo_collection_info.cc
   geometry/nodes/node_geo_common.cc
   geometry/nodes/node_geo_convex_hull.cc
+  geometry/nodes/node_geo_curve_endpoint_selection.cc
   geometry/nodes/node_geo_curve_fill.cc
   geometry/nodes/node_geo_curve_fillet.cc
   geometry/nodes/node_geo_curve_handle_type_selection.cc
diff --git a/source/blender/nodes/NOD_geometry.h b/source/blender/nodes/NOD_geometry.h
index 94fed3c3975..3ca8c0d02d7 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/source/blender/nodes/NOD_geometry.h
@@ -38,6 +38,7 @@ void register_node_type_geo_legacy_material_assign(void);
 void register_node_type_geo_legacy_mesh_to_curve(void);
 void register_node_type_geo_legacy_points_to_volume(void);
 void register_node_type_geo_legacy_select_by_material(void);
+void register_node_type_geo_legacy_curve_endpoints(void);
 void register_node_type_geo_legacy_curve_spline_type(void);
 void register_node_type_geo_legacy_curve_reverse(void);
 void register_node_type_geo_legacy_select_by_handle_type(void);
@@ -66,7 +67,7 @@ void register_node_type_geo_boolean(void);
 void register_node_type_geo_bounding_box(void);
 void register_node_type_geo_collection_info(void);
 void register_node_type_geo_convex_hull(void);
-void register_node_type_geo_curve_endpoints(void);
+void register_node_type_geo_curve_endpoint_selection(void);
 void register_node_type_geo_curve_fill(void);
 void register_node_type_geo_curve_fillet(void);
 void register_node_type_geo_curve_handle_type_selection(void);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index 827a12ff812..9a361c3a3f1 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/source/blender/nodes/NOD_static_types.h
@@ -326,6 +326,7 @@ DefNode(GeometryNode, GEO_NODE_MESH_BOOLEAN, def_geo_boolean, "MESH_BOOLEAN", Me
 DefNode(GeometryNode, GEO_NODE_BOUNDING_BOX, 0, "BOUNDING_BOX", BoundBox, "Bounding Box", "")
 DefNode(GeometryNode, GEO_NODE_COLLECTION_INFO, def_geo_collection_info, "COLLECTION_INFO", CollectionInfo, "Collection Info", "")
 DefNode(GeometryNode, GEO_NODE_CONVEX_HULL, 0, "CONVEX_HULL", ConvexHull, "Convex Hull", "")
+DefNode(GeometryNode, GEO_NODE_CURVE_ENDPOINT_SELECTION, 0, "CURVE_ENDPOINT_SELECTION", CurveEndpointSelection, "Endpoint Selection", "")
 DefNode(GeometryNode, GEO_NODE_FILL_CURVE, def_geo_curve_fill, "FILL_CURVE", FillCurve, "Fill Curve", "")
 DefNode(GeometryNode, GEO_NODE_FILLET_CURVE, def_geo_curve_fillet, "FILLET_CURVE", FilletCurve, "Fillet Curve", "")
 DefNode(GeometryNode, GEO_NODE_CURVE_HANDLE_TYPE_SELECTION, def_geo_curve_handle_type_selection, "CURVE_HANDLE_TYPE_SELECTION", CurveHandleTypeSelection, "Handle Type Selection", "")
diff --git a/source/blender/nodes/geometry/nodes/legacy/node_geo_curve_endpoints.cc b/source/blender/nodes/geometry/nodes/legacy/node_geo_curve_endpoints.cc
index b226cc2d3be..85d1392aa35 100644
--- a/source/blender/nodes/geometry/nodes/legacy/node_geo_curve_endpoints.cc
+++ b/source/blender/nodes/geometry/nodes/legacy/node_geo_curve_endpoints.cc
@@ -208,7 +208,7 @@ static void geo_node_curve_endpoints_exec(GeoNodeExecParams params)
 
 }  // namespace blender::nodes
 
-void register_node_type_geo_curve_endpoints()
+void register_node_type_geo_legacy_curve_endpoints()
 {
   static bNodeType ntype;
 
diff --git a/source/blender/nodes/geometry/nodes/node_geo_curve_endpoint_selection.cc b/source/blender/nodes/geometry/nodes/node_geo_curve_endpoint_selection.cc
new file mode 100644
index 00000000000..ee6cf055ecb
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/node_geo_curve_endpoint_selection.cc
@@ -0,0 +1,149 @@
+/*
+ * 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 "BKE_spline.hh"
+
+#include "UI_interface.h"
+#include "UI_resources.h"
+
+#include "node_geometry_util.hh"
+
+namespace blender::nodes {
+
+static void geo_node_curve_endpoint_selection_declare(NodeDeclarationBuilder &b)
+{
+  b.add_input<decl::Int>("Start Size")
+      .min(0)
+      .default_value(1)
+      .supports_field()
+      .description("The amount of points to select from the start of each spline");
+  b.add_input<decl::Int>("End Size")
+      .min(0)
+      .default_value(1)
+      .supports_field()
+      .description("The amount of points to select from the end of each spline");
+  b.add_output<decl::Bool>("Selection")
+      .field_source()
+      .description("The selection from the start and end of the splines based on the input sizes");
+}
+
+static void select_by_spline(const int start, const int end, MutableSpan<bool> r_selection)
+{
+  const int size = r_selection.size();
+  const int start_use = std::min(start, size);
+  const int end_use = std::min(end, size);
+
+  r_selection.slice(0, start_use).fill(true);
+  r_selection.slice(size - end_use, end_use).fill(true);
+}
+
+class EndpointFieldInput final : public fn::FieldInput {
+  Field<int> start_size_;
+  Field<int> end_size_;
+
+ public:
+  EndpointFieldInput(Field<int> start_size, Field<int> end_size)
+      : FieldInput(CPPType::get<bool>(), "Selection"), start_size_(start_size), end_size_(end_size)
+  {
+  }
+
+  const GVArray *get_varray_for_context(const fn::FieldContext &context,
+                                        IndexMask UNUSED(mask),
+                                        ResourceScope &scope) const final
+  {
+    if (const GeometryComponentFieldContext *geometry_context =
+            dynamic_cast<const GeometryComponentFieldContext *>(&context)) {
+
+      const GeometryComponent &component = geometry_context->geometry_component();
+      const AttributeDomain domain = geometry_context->domain();
+      if (component.type() != GEO_COMPONENT_TYPE_CURVE || domain != ATTR_DOMAIN_POINT) {
+        return nullptr;
+      }
+
+      const CurveComponent &curve_component = static_cast<const CurveComponent &>(component);
+      const CurveEval *curve = curve_component.get_for_read();
+
+      Array<int> control_point_offsets = curve->control_point_offsets();
+
+      if (curve == nullptr || control_point_offsets.last() == 0) {
+        return nullptr;
+      }
+
+      GeometryComponentFieldContext size_context{curve_component, ATTR_DOMAIN_CURVE};
+      fn::FieldEvaluator evaluator{size_context, curve->splines().size()};
+      evaluator.add(start_size_);
+      evaluator.add(end_size_);
+      evaluator.evaluate();
+      const VArray<int> &start_size = evaluator.get_evaluated<int>(0);
+      const VArray<int> &end_size = evaluator.get_evaluated<int>(1);
+
+      const int point_size = control_point_offsets.last();
+      Array<bool> selection(point_size, false);
+      int current_point = 0;
+      MutableSpan<bool> selection_span = selection.as_mutable_span();
+      for (int i : IndexRange(curve->splines().size())) {
+        const SplinePtr &spline = curve->splines()[i];
+        if (start_size[i] <= 0 && end_size[i] <= 0) {
+          selection_span.slice(current_point, spline->size()).fill(false);
+        }
+        else {
+          int start_use = std::max(start_size[i], 0);
+          int end_use = std::max(end_size[i], 0);
+          select_by_spline(
+              start_use, end_use, selection_span.slice(current_point, spline->size()));
+        }
+        current_point += spline->size();
+      }
+      return &scope.construct<fn::GVArray_For_ArrayContainer<Array<bool>>>(std::move(selection));
+    }
+    return nullptr;
+  };
+
+  uint64_t hash() const override
+  {
+    return get_default_hash_2(start_size_, end_size_);
+  }
+
+  bool is_equal_to(const fn::FieldNode &other) const override
+  {
+    if (const EndpointFieldInput *other_endpoint = dynamic_cast<const EndpointFieldInput *>(
+            &other)) {
+      return start_size_ == other_endpoint->start_size_ && end_size_ == other_endpoint->end_size_;
+    }
+    return false;
+  }
+};
+
+static void geo_node_curve_endpoint_selection_exec(GeoNodeExecParams params)
+{
+  Field<int> start_size = params.extract_input<Field<int>>("Start Size");
+  Field<int> end_size = params.extract_input<Field<int>>("End Size");
+  Field<bool> selection_field{std::make_shared<EndpointFieldInput>(start_size, end_size)};
+  params.set_output("Selection", std::move(selection_field));
+}
+}  // namespace blender::nodes
+
+void register_node_type_geo_curve_endpoint_selection()
+{
+  static bNodeType ntype;
+
+  geo_node_type_base(
+      &ntype, GEO_NODE_CURVE_ENDPOINT_SELECTION, "Endpoint Selection", NODE_CLASS_INPUT, 0);
+  ntype.declare = blender::nodes::geo_node_curve_endpoint_selection_declare;
+  ntype.geometry_node_execute = blender::nodes::geo_node_curve_endpoint_selection_exec;
+
+  nodeRegisterType(&ntype);
+}