Geometry Nodes: Sort Elements node

Implements the design in #109983

Pull Request: https://projects.blender.org/blender/blender/pulls/114194
This commit is contained in:
Iliya Katushenock
2024-01-12 14:30:34 +01:00
committed by Hans Goudey
parent 4ac0267567
commit 37b2c12cfa
18 changed files with 740 additions and 24 deletions

View File

@@ -215,6 +215,7 @@ class NODE_MT_geometry_node_GEO_GEOMETRY_OPERATIONS(Menu):
node_add_menu.add_node_type(layout, "GeometryNodeDeleteGeometry")
node_add_menu.add_node_type(layout, "GeometryNodeDuplicateElements")
node_add_menu.add_node_type(layout, "GeometryNodeMergeByDistance")
node_add_menu.add_node_type(layout, "GeometryNodeSortElements")
node_add_menu.add_node_type(layout, "GeometryNodeTransform")
layout.separator()
node_add_menu.add_node_type(layout, "GeometryNodeSeparateComponents")

View File

@@ -447,6 +447,7 @@ class MeshComponent : public GeometryComponent {
public:
MeshComponent();
MeshComponent(Mesh *mesh, GeometryOwnershipType ownership = GeometryOwnershipType::Owned);
~MeshComponent();
GeometryComponentPtr copy() const override;
@@ -501,6 +502,8 @@ class PointCloudComponent : public GeometryComponent {
public:
PointCloudComponent();
PointCloudComponent(PointCloud *pointcloud,
GeometryOwnershipType ownership = GeometryOwnershipType::Owned);
~PointCloudComponent();
GeometryComponentPtr copy() const override;
@@ -561,6 +564,7 @@ class CurveComponent : public GeometryComponent {
public:
CurveComponent();
CurveComponent(Curves *curve, GeometryOwnershipType ownership = GeometryOwnershipType::Owned);
~CurveComponent();
GeometryComponentPtr copy() const override;
@@ -602,6 +606,8 @@ class InstancesComponent : public GeometryComponent {
public:
InstancesComponent();
InstancesComponent(Instances *instances,
GeometryOwnershipType ownership = GeometryOwnershipType::Owned);
~InstancesComponent();
GeometryComponentPtr copy() const override;

View File

@@ -1327,6 +1327,7 @@ void BKE_nodetree_remove_layer_n(struct bNodeTree *ntree, struct Scene *scene, i
#define GEO_NODE_BAKE 2120
#define GEO_NODE_GET_NAMED_GRID 2121
#define GEO_NODE_STORE_NAMED_GRID 2122
#define GEO_NODE_SORT_ELEMENTS 2123
/** \} */

View File

@@ -936,29 +936,6 @@ void gather_attributes(const AttributeAccessor src_attributes,
});
}
static bool indices_are_range(const Span<int> indices, const IndexRange range)
{
if (indices.size() != range.size()) {
return false;
}
return threading::parallel_reduce(
range,
4096,
true,
[&](const IndexRange range, const bool init) {
if (!init) {
return false;
}
for (const int i : range) {
if (indices[i] != i) {
return false;
}
}
return true;
},
std::logical_and());
}
void gather_attributes(const AttributeAccessor src_attributes,
const AttrDomain domain,
const AnonymousAttributePropagationInfo &propagation_info,
@@ -966,7 +943,7 @@ void gather_attributes(const AttributeAccessor src_attributes,
const Span<int> indices,
MutableAttributeAccessor dst_attributes)
{
if (indices_are_range(indices, IndexRange(src_attributes.domain_size(domain)))) {
if (array_utils::indices_are_range(indices, IndexRange(src_attributes.domain_size(domain)))) {
copy_attributes(src_attributes, domain, propagation_info, skip, dst_attributes);
}
else {

View File

@@ -27,6 +27,11 @@ namespace blender::bke {
CurveComponent::CurveComponent() : GeometryComponent(Type::Curve) {}
CurveComponent::CurveComponent(Curves *curve, GeometryOwnershipType ownership)
: GeometryComponent(Type::Curve), curves_(curve), ownership_(ownership)
{
}
CurveComponent::~CurveComponent()
{
this->clear();

View File

@@ -32,6 +32,11 @@ namespace blender::bke {
InstancesComponent::InstancesComponent() : GeometryComponent(Type::Instance) {}
InstancesComponent::InstancesComponent(Instances *instances, GeometryOwnershipType ownership)
: GeometryComponent(Type::Instance), instances_(instances), ownership_(ownership)
{
}
InstancesComponent::~InstancesComponent()
{
this->clear();

View File

@@ -28,6 +28,11 @@ namespace blender::bke {
MeshComponent::MeshComponent() : GeometryComponent(Type::Mesh) {}
MeshComponent::MeshComponent(Mesh *mesh, GeometryOwnershipType ownership)
: GeometryComponent(Type::Mesh), mesh_(mesh), ownership_(ownership)
{
}
MeshComponent::~MeshComponent()
{
this->clear();

View File

@@ -18,6 +18,11 @@ namespace blender::bke {
PointCloudComponent::PointCloudComponent() : GeometryComponent(Type::PointCloud) {}
PointCloudComponent::PointCloudComponent(PointCloud *pointcloud, GeometryOwnershipType ownership)
: GeometryComponent(Type::PointCloud), pointcloud_(pointcloud), ownership_(ownership)
{
}
PointCloudComponent::~PointCloudComponent()
{
this->clear();

View File

@@ -83,6 +83,19 @@ inline void scatter(const Span<T> src,
});
}
template<typename T>
inline void scatter(const Span<T> src,
const IndexMask &indices,
MutableSpan<T> dst,
const int64_t grain_size = 4096)
{
BLI_assert(indices.size() == src.size());
BLI_assert(indices.min_array_size() <= dst.size());
indices.foreach_index_optimized<int64_t>(
GrainSize(grain_size),
[&](const int64_t index, const int64_t pos) { dst[index] = src[pos]; });
}
/**
* Fill the destination span by gathering indexed values from the `src` array.
*/
@@ -281,4 +294,6 @@ bool indexed_data_equal(const Span<T> all_values, const Span<int> indices, const
return false;
}
bool indices_are_range(Span<int> indices, IndexRange range);
} // namespace blender::array_utils

View File

@@ -155,6 +155,8 @@ void copy_group_sizes(OffsetIndices<int> offsets, const IndexMask &mask, Mutable
/** Gather the number of indices in each indexed group to sizes. */
void gather_group_sizes(OffsetIndices<int> offsets, const IndexMask &mask, MutableSpan<int> sizes);
void gather_group_sizes(OffsetIndices<int> offsets, Span<int> indices, MutableSpan<int> sizes);
/** Build new offsets that contains only the groups chosen by \a selection. */
OffsetIndices<int> gather_selected_offsets(OffsetIndices<int> src_offsets,
const IndexMask &selection,

View File

@@ -2,6 +2,8 @@
*
* SPDX-License-Identifier: GPL-2.0-or-later */
#include <functional>
#include "BLI_array_utils.hh"
#include "BLI_threads.h"
@@ -196,4 +198,22 @@ int64_t count_booleans(const VArray<bool> &varray)
return count_booleans(varray, IndexMask(varray.size()));
}
bool indices_are_range(Span<int> indices, IndexRange range)
{
if (indices.size() != range.size()) {
return false;
}
return threading::parallel_reduce(
range.index_range(),
4096,
true,
[&](const IndexRange part, const bool is_range) {
const Span<int> local_indices = indices.slice(part);
const IndexRange local_range = range.slice(part);
return is_range &&
std::equal(local_indices.begin(), local_indices.end(), local_range.begin());
},
std::logical_and<bool>());
}
} // namespace blender::array_utils

View File

@@ -48,6 +48,17 @@ void gather_group_sizes(const OffsetIndices<int> offsets,
});
}
void gather_group_sizes(const OffsetIndices<int> offsets,
const Span<int> indices,
MutableSpan<int> sizes)
{
threading::parallel_for(indices.index_range(), 4096, [&](const IndexRange range) {
for (const int i : range) {
sizes[i] = offsets[indices[i]].size();
}
});
}
OffsetIndices<int> gather_selected_offsets(const OffsetIndices<int> src_offsets,
const IndexMask &selection,
const int start_offset,

View File

@@ -34,6 +34,7 @@ set(SRC
intern/points_to_volume.cc
intern/randomize.cc
intern/realize_instances.cc
intern/reorder.cc
intern/resample_curves.cc
intern/reverse_uv_sampler.cc
intern/set_curve_type.cc
@@ -61,6 +62,7 @@ set(SRC
GEO_points_to_volume.hh
GEO_randomize.hh
GEO_realize_instances.hh
GEO_reorder.hh
GEO_resample_curves.hh
GEO_reverse_uv_sampler.hh
GEO_set_curve_type.hh

View File

@@ -0,0 +1,42 @@
/* SPDX-FileCopyrightText: 2024 Blender Authors
*
* SPDX-License-Identifier: GPL-2.0-or-later */
#pragma once
#include "BLI_multi_value_map.hh"
#include "BLI_span.hh"
#include "BKE_anonymous_attribute_id.hh"
#include "BKE_attribute.hh"
#include "BKE_geometry_set.hh"
namespace blender::geometry {
const MultiValueMap<bke::GeometryComponent::Type, bke::AttrDomain> &
components_supported_reordering();
Mesh *reorder_mesh(const Mesh &src_mesh,
Span<int> old_by_new_map,
bke::AttrDomain domain,
const bke::AnonymousAttributePropagationInfo &propagation_info);
PointCloud *reorder_points(const PointCloud &src_pointcloud,
Span<int> old_by_new_map,
const bke::AnonymousAttributePropagationInfo &propagation_info);
Curves *reorder_curves(const Curves &src_curves,
Span<int> old_by_new_map,
const bke::AnonymousAttributePropagationInfo &propagation_info);
bke::Instances *reorder_instaces(const bke::Instances &src_instances,
Span<int> old_by_new_map,
const bke::AnonymousAttributePropagationInfo &propagation_info);
bke::GeometryComponentPtr reordered_component(
const bke::GeometryComponent &src_component,
Span<int> old_by_new_map,
bke::AttrDomain domain,
const bke::AnonymousAttributePropagationInfo &propagation_info);
}; // namespace blender::geometry

View File

@@ -0,0 +1,327 @@
/* SPDX-FileCopyrightText: 2024 Blender Authors
*
* SPDX-License-Identifier: GPL-2.0-or-later */
#include "BKE_anonymous_attribute_id.hh"
#include "BKE_attribute.hh"
#include "BKE_attribute_math.hh"
#include "BKE_curves.hh"
#include "BKE_geometry_set.hh"
#include "BKE_instances.hh"
#include "BKE_mesh.hh"
#include "BKE_pointcloud.h"
#include "BLI_array.hh"
#include "BLI_array_utils.hh"
#include "BLI_multi_value_map.hh"
#include "DNA_curves_types.h"
#include "DNA_mesh_types.h"
#include "DNA_pointcloud_types.h"
#include "GEO_reorder.hh"
namespace blender::geometry {
const MultiValueMap<bke::GeometryComponent::Type, bke::AttrDomain> &
components_supported_reordering()
{
using namespace bke;
const static MultiValueMap<GeometryComponent::Type, AttrDomain> supported_types_and_domains =
[]() {
MultiValueMap<GeometryComponent::Type, AttrDomain> supported_types_and_domains;
supported_types_and_domains.add_multiple(
GeometryComponent::Type::Mesh,
{AttrDomain::Point, AttrDomain::Edge, AttrDomain::Face});
supported_types_and_domains.add(GeometryComponent::Type::Curve, AttrDomain::Curve);
supported_types_and_domains.add(GeometryComponent::Type::PointCloud, AttrDomain::Point);
supported_types_and_domains.add(GeometryComponent::Type::Instance, AttrDomain::Instance);
return supported_types_and_domains;
}();
return supported_types_and_domains;
}
static void reorder_attributes_group_to_group(const bke::AttributeAccessor src_attributes,
const bke::AttrDomain domain,
const OffsetIndices<int> src_offsets,
const OffsetIndices<int> dst_offsets,
const Span<int> old_by_new_map,
bke::MutableAttributeAccessor dst_attributes)
{
src_attributes.for_all(
[&](const bke::AttributeIDRef &id, const bke::AttributeMetaData meta_data) {
if (meta_data.domain != domain) {
return true;
}
const GVArray src = *src_attributes.lookup(id, domain);
bke::GSpanAttributeWriter dst = dst_attributes.lookup_or_add_for_write_only_span(
id, domain, meta_data.data_type);
if (!dst) {
return true;
}
threading::parallel_for(old_by_new_map.index_range(), 1024, [&](const IndexRange range) {
for (const int new_i : range) {
const int old_i = old_by_new_map[new_i];
array_utils::copy(src.slice(src_offsets[old_i]), dst.span.slice(dst_offsets[new_i]));
}
});
dst.finish();
return true;
});
}
static Array<int> invert_permutation(const Span<int> permutation)
{
Array<int> data(permutation.size());
threading::parallel_for(permutation.index_range(), 2048, [&](const IndexRange range) {
for (const int64_t i : range) {
data[permutation[i]] = i;
}
});
return data;
}
static void reorder_mesh_verts_exec(const Mesh &src_mesh,
const Span<int> old_by_new_map,
Mesh &dst_mesh)
{
bke::gather_attributes(src_mesh.attributes(),
bke::AttrDomain::Point,
{},
{},
old_by_new_map,
dst_mesh.attributes_for_write());
const Array<int> new_by_old_map = invert_permutation(old_by_new_map);
array_utils::gather(new_by_old_map.as_span(),
dst_mesh.edges().cast<int>(),
dst_mesh.edges_for_write().cast<int>());
array_utils::gather(
new_by_old_map.as_span(), dst_mesh.corner_verts(), dst_mesh.corner_verts_for_write());
}
static void reorder_mesh_edges_exec(const Mesh &src_mesh,
const Span<int> old_by_new_map,
Mesh &dst_mesh)
{
bke::gather_attributes(src_mesh.attributes(),
bke::AttrDomain::Edge,
{},
{},
old_by_new_map,
dst_mesh.attributes_for_write());
const Array<int> new_by_old_map = invert_permutation(old_by_new_map);
array_utils::gather(
new_by_old_map.as_span(), dst_mesh.corner_edges(), dst_mesh.corner_edges_for_write());
}
static void reorder_mesh_faces_exec(const Mesh &src_mesh,
const Span<int> old_by_new_map,
Mesh &dst_mesh)
{
bke::gather_attributes(src_mesh.attributes(),
bke::AttrDomain::Face,
{},
{},
old_by_new_map,
dst_mesh.attributes_for_write());
const Span<int> old_offsets = src_mesh.face_offsets();
MutableSpan<int> new_offsets = dst_mesh.face_offsets_for_write();
offset_indices::gather_group_sizes(old_offsets, old_by_new_map, new_offsets);
offset_indices::accumulate_counts_to_offsets(new_offsets);
reorder_attributes_group_to_group(src_mesh.attributes(),
bke::AttrDomain::Corner,
old_offsets,
new_offsets.as_span(),
old_by_new_map,
dst_mesh.attributes_for_write());
}
static void reorder_mesh_exec(const Mesh &src_mesh,
const Span<int> old_by_new_map,
const bke::AttrDomain domain,
Mesh &dst_mesh)
{
switch (domain) {
case bke::AttrDomain::Point:
reorder_mesh_verts_exec(src_mesh, old_by_new_map, dst_mesh);
break;
case bke::AttrDomain::Edge:
reorder_mesh_edges_exec(src_mesh, old_by_new_map, dst_mesh);
break;
case bke::AttrDomain::Face:
reorder_mesh_faces_exec(src_mesh, old_by_new_map, dst_mesh);
break;
default:
break;
}
dst_mesh.tag_positions_changed();
dst_mesh.tag_topology_changed();
}
static void reorder_points_exec(const PointCloud &src_pointcloud,
const Span<int> old_by_new_map,
PointCloud &dst_pointcloud)
{
bke::gather_attributes(src_pointcloud.attributes(),
bke::AttrDomain::Point,
{},
{},
old_by_new_map,
dst_pointcloud.attributes_for_write());
dst_pointcloud.tag_positions_changed();
dst_pointcloud.tag_radii_changed();
}
static void reorder_curves_exec(const bke::CurvesGeometry &src_curves,
const Span<int> old_by_new_map,
bke::CurvesGeometry &dst_curves)
{
bke::gather_attributes(src_curves.attributes(),
bke::AttrDomain::Curve,
{},
{},
old_by_new_map,
dst_curves.attributes_for_write());
const Span<int> old_offsets = src_curves.offsets();
MutableSpan<int> new_offsets = dst_curves.offsets_for_write();
offset_indices::gather_group_sizes(old_offsets, old_by_new_map, new_offsets);
offset_indices::accumulate_counts_to_offsets(new_offsets);
reorder_attributes_group_to_group(src_curves.attributes(),
bke::AttrDomain::Point,
old_offsets,
new_offsets.as_span(),
old_by_new_map,
dst_curves.attributes_for_write());
dst_curves.tag_topology_changed();
}
static void reorder_instaces_exec(const bke::Instances &src_instances,
const Span<int> old_by_new_map,
bke::Instances &dst_instances)
{
bke::gather_attributes(src_instances.attributes(),
bke::AttrDomain::Instance,
{},
{},
old_by_new_map,
dst_instances.attributes_for_write());
const Span<int> old_reference_handles = src_instances.reference_handles();
MutableSpan<int> new_reference_handles = dst_instances.reference_handles();
array_utils::gather(old_reference_handles, old_by_new_map, new_reference_handles);
const Span<float4x4> old_transforms = src_instances.transforms();
MutableSpan<float4x4> new_transforms = dst_instances.transforms();
array_utils::gather(old_transforms, old_by_new_map, new_transforms);
}
static void clean_unused_attributes(const bke::AnonymousAttributePropagationInfo &propagation_info,
bke::MutableAttributeAccessor attributes)
{
Vector<std::string> unused_ids;
attributes.for_all(
[&](const bke::AttributeIDRef &id, const bke::AttributeMetaData /*meta_data*/) {
if (!id.is_anonymous()) {
return true;
}
if (propagation_info.propagate(id.anonymous_id())) {
return true;
}
unused_ids.append(id.name());
return true;
});
for (const std::string &unused_id : unused_ids) {
attributes.remove(unused_id);
}
}
Mesh *reorder_mesh(const Mesh &src_mesh,
Span<int> old_by_new_map,
bke::AttrDomain domain,
const bke::AnonymousAttributePropagationInfo &propagation_info)
{
Mesh *dst_mesh = BKE_mesh_copy_for_eval(&src_mesh);
clean_unused_attributes(propagation_info, dst_mesh->attributes_for_write());
reorder_mesh_exec(src_mesh, old_by_new_map, domain, *dst_mesh);
return dst_mesh;
}
PointCloud *reorder_points(const PointCloud &src_pointcloud,
Span<int> old_by_new_map,
const bke::AnonymousAttributePropagationInfo &propagation_info)
{
PointCloud *dst_pointcloud = BKE_pointcloud_copy_for_eval(&src_pointcloud);
clean_unused_attributes(propagation_info, dst_pointcloud->attributes_for_write());
reorder_points_exec(src_pointcloud, old_by_new_map, *dst_pointcloud);
return dst_pointcloud;
}
Curves *reorder_curves(const Curves &src_curves,
Span<int> old_by_new_map,
const bke::AnonymousAttributePropagationInfo &propagation_info)
{
const bke::CurvesGeometry src_curve_geometry = src_curves.geometry.wrap();
Curves *dst_curves = BKE_curves_copy_for_eval(&src_curves);
bke::CurvesGeometry &dst_curve_geometry = dst_curves->geometry.wrap();
clean_unused_attributes(propagation_info, dst_curve_geometry.attributes_for_write());
reorder_curves_exec(src_curve_geometry, old_by_new_map, dst_curve_geometry);
return dst_curves;
}
bke::Instances *reorder_instaces(const bke::Instances &src_instances,
Span<int> old_by_new_map,
const bke::AnonymousAttributePropagationInfo &propagation_info)
{
bke::Instances *dst_instances = new bke::Instances(src_instances);
clean_unused_attributes(propagation_info, dst_instances->attributes_for_write());
reorder_instaces_exec(src_instances, old_by_new_map, *dst_instances);
return dst_instances;
}
bke::GeometryComponentPtr reordered_component(
const bke::GeometryComponent &src_component,
const Span<int> old_by_new_map,
const bke::AttrDomain domain,
const bke::AnonymousAttributePropagationInfo &propagation_info)
{
BLI_assert(!src_component.is_empty());
if (const bke::MeshComponent *src_mesh_component = dynamic_cast<const bke::MeshComponent *>(
&src_component))
{
Mesh *result_mesh = reorder_mesh(
*src_mesh_component->get(), old_by_new_map, domain, propagation_info);
return bke::GeometryComponentPtr(new bke::MeshComponent(result_mesh));
}
else if (const bke::PointCloudComponent *src_points_component =
dynamic_cast<const bke::PointCloudComponent *>(&src_component))
{
PointCloud *result_point_cloud = reorder_points(
*src_points_component->get(), old_by_new_map, propagation_info);
return bke::GeometryComponentPtr(new bke::PointCloudComponent(result_point_cloud));
}
else if (const bke::CurveComponent *src_curves_component =
dynamic_cast<const bke::CurveComponent *>(&src_component))
{
Curves *result_curves = reorder_curves(
*src_curves_component->get(), old_by_new_map, propagation_info);
return bke::GeometryComponentPtr(new bke::CurveComponent(result_curves));
}
else if (const bke::InstancesComponent *src_instances_component =
dynamic_cast<const bke::InstancesComponent *>(&src_component))
{
bke::Instances *result_instances = reorder_instaces(
*src_instances_component->get(), old_by_new_map, propagation_info);
return bke::GeometryComponentPtr(new bke::InstancesComponent(result_instances));
}
BLI_assert_unreachable();
return {};
}
} // namespace blender::geometry

View File

@@ -438,6 +438,7 @@ DefNode(GeometryNode, GEO_NODE_SET_SPLINE_CYCLIC, 0, "SET_SPLINE_CYCLIC", SetSpl
DefNode(GeometryNode, GEO_NODE_SET_SPLINE_RESOLUTION, 0, "SET_SPLINE_RESOLUTION", SetSplineResolution, "Set Spline Resolution", "Control how many evaluated points should be generated on every curve segment")
DefNode(GeometryNode, GEO_NODE_SIMULATION_INPUT, def_geo_simulation_input, "SIMULATION_INPUT", SimulationInput, "Simulation Input", "Input data for the simulation zone")
DefNode(GeometryNode, GEO_NODE_SIMULATION_OUTPUT, def_geo_simulation_output, "SIMULATION_OUTPUT", SimulationOutput, "Simulation Output", "Output data from the simulation zone")
DefNode(GeometryNode, GEO_NODE_SORT_ELEMENTS, 0, "SORT_ELEMENTS", SortElements, "Sort Elements", "Rearrange geometry elements, changing their indices")
DefNode(GeometryNode, GEO_NODE_SPLIT_TO_INSTANCES, 0, "Split to Instances", SplitToInstances, "Split to Instances", "Create separate geometries containing the elements from the same group")
DefNode(GeometryNode, GEO_NODE_SPLIT_EDGES, 0, "SPLIT_EDGES", SplitEdges, "Split Edges", "Duplicate mesh edges and break connections with the surrounding faces")
DefNode(GeometryNode, GEO_NODE_STORE_NAMED_ATTRIBUTE, 0, "STORE_NAMED_ATTRIBUTE", StoreNamedAttribute, "Store Named Attribute", "Store the result of a field on a geometry as an attribute with the specified name")

View File

@@ -179,6 +179,7 @@ set(SRC
nodes/node_geo_set_spline_resolution.cc
nodes/node_geo_simulation_input.cc
nodes/node_geo_simulation_output.cc
nodes/node_geo_sort_elements.cc
nodes/node_geo_split_to_instances.cc
nodes/node_geo_store_named_attribute.cc
nodes/node_geo_store_named_grid.cc

View File

@@ -0,0 +1,290 @@
/* SPDX-FileCopyrightText: 2024 Blender Authors
*
* SPDX-License-Identifier: GPL-2.0-or-later */
#include <atomic>
#include "BKE_attribute.hh"
#include "BKE_mesh.hh"
#include "BLI_array_utils.hh"
#include "BLI_index_mask.hh"
#include "BLI_sort.hh"
#include "BLI_task.hh"
#include "DNA_mesh_types.h"
#include "GEO_reorder.hh"
#include "NOD_rna_define.hh"
#include "RNA_enum_types.hh"
#include "UI_interface.hh"
#include "UI_resources.hh"
#include "node_geometry_util.hh"
namespace blender::nodes::node_geo_sort_elements_cc {
static void node_declare(NodeDeclarationBuilder &b)
{
b.add_input<decl::Geometry>("Geometry");
b.add_input<decl::Bool>("Selection").default_value(true).field_on_all().hide_value();
b.add_input<decl::Int>("Group ID").field_on_all().hide_value();
b.add_input<decl::Float>("Sort Weight").field_on_all().hide_value();
b.add_output<decl::Geometry>("Geometry").propagate_all();
}
static void node_layout(uiLayout *layout, bContext * /*C*/, PointerRNA *ptr)
{
uiItemR(layout, ptr, "domain", UI_ITEM_NONE, "", ICON_NONE);
}
static void node_init(bNodeTree * /*tree*/, bNode *node)
{
node->custom1 = int(bke::AttrDomain::Point);
}
static void grouped_sort(const OffsetIndices<int> offsets,
const Span<float> weights,
MutableSpan<int> indices)
{
const auto comparator = [&](const int index_a, const int index_b) {
const float weight_a = weights[index_a];
const float weight_b = weights[index_b];
if (UNLIKELY(weight_a == weight_b)) {
/* Approach to make it stable. */
return index_a < index_b;
}
return weight_a < weight_b;
};
threading::parallel_for(offsets.index_range(), 250, [&](const IndexRange range) {
for (const int group_index : range) {
MutableSpan<int> group = indices.slice(offsets[group_index]);
parallel_sort(group.begin(), group.end(), comparator);
}
});
}
static void find_points_by_group_index(const Span<int> indices,
MutableSpan<int> r_offsets,
MutableSpan<int> r_indices)
{
offset_indices::build_reverse_offsets(indices, r_offsets);
Array<int> counts(r_offsets.size(), 0);
for (const int64_t index : indices.index_range()) {
const int curve_index = indices[index];
r_indices[r_offsets[curve_index] + counts[curve_index]] = int(index);
counts[curve_index]++;
}
}
template<typename T, typename Func>
static void parallel_transform(MutableSpan<T> values, const int64_t grain_size, const Func &func)
{
threading::parallel_for(values.index_range(), grain_size, [&](const IndexRange range) {
MutableSpan<T> values_range = values.slice(range);
std::transform(values_range.begin(), values_range.end(), values_range.begin(), func);
});
}
static Array<int> invert_permutation(const Span<int> permutation)
{
Array<int> data(permutation.size());
threading::parallel_for(permutation.index_range(), 2048, [&](const IndexRange range) {
for (const int64_t i : range) {
data[permutation[i]] = i;
}
});
return data;
}
static int identifiers_to_indices(MutableSpan<int> r_identifiers_to_indices)
{
const VectorSet<int> deduplicated_identifiers(r_identifiers_to_indices);
parallel_transform(r_identifiers_to_indices, 2048, [&](const int identifier) {
return deduplicated_identifiers.index_of(identifier);
});
Array<int> indices(deduplicated_identifiers.size());
array_utils::fill_index_range<int>(indices);
parallel_sort(indices.begin(), indices.end(), [&](const int index_a, const int index_b) {
return deduplicated_identifiers[index_a] < deduplicated_identifiers[index_b];
});
Array<int> permutation = invert_permutation(indices);
parallel_transform(
r_identifiers_to_indices, 4096, [&](const int index) { return permutation[index]; });
return deduplicated_identifiers.size();
}
static std::optional<Array<int>> sorted_indices(const bke::GeometryComponent &component,
const Field<bool> selection_field,
const Field<int> group_id_field,
const Field<float> weight_field,
const bke::AttrDomain domain)
{
const int domain_size = component.attribute_domain_size(domain);
if (domain_size == 0) {
return std::nullopt;
}
const bke::GeometryFieldContext context(component, domain);
FieldEvaluator evaluator(context, domain_size);
evaluator.set_selection(selection_field);
evaluator.add(group_id_field);
evaluator.add(weight_field);
evaluator.evaluate();
const IndexMask mask = evaluator.get_evaluated_selection_as_mask();
const VArray<int> group_id = evaluator.get_evaluated<int>(0);
const VArray<float> weight = evaluator.get_evaluated<float>(1);
if (group_id.is_single() && weight.is_single()) {
return std::nullopt;
}
if (mask.is_empty()) {
return std::nullopt;
}
Array<int> gathered_indices(mask.size());
if (group_id.is_single()) {
mask.to_indices<int>(gathered_indices);
Array<float> weight_span(domain_size);
array_utils::copy(weight, mask, weight_span.as_mutable_span());
grouped_sort(Span({0, int(mask.size())}), weight_span, gathered_indices);
}
else {
Array<int> gathered_group_id(mask.size());
array_utils::gather(group_id, mask, gathered_group_id.as_mutable_span());
const int total_groups = identifiers_to_indices(gathered_group_id);
Array<int> offsets_to_sort(total_groups + 1, 0);
find_points_by_group_index(gathered_group_id, offsets_to_sort, gathered_indices);
if (!weight.is_single()) {
Array<float> weight_span(mask.size());
array_utils::gather(weight, mask, weight_span.as_mutable_span());
grouped_sort(offsets_to_sort.as_span(), weight_span, gathered_indices);
}
parallel_transform<int>(gathered_indices, 2048, [&](const int pos) { return mask[pos]; });
}
if (array_utils::indices_are_range(gathered_indices, IndexRange(domain_size))) {
return std::nullopt;
}
if (mask.size() == domain_size) {
return gathered_indices;
}
IndexMaskMemory memory;
const IndexMask unselected = mask.complement(IndexRange(domain_size), memory);
Array<int> indices(domain_size);
array_utils::scatter<int>(gathered_indices, mask, indices);
unselected.foreach_index_optimized<int>(GrainSize(2048),
[&](const int index) { indices[index] = index; });
if (array_utils::indices_are_range(indices, indices.index_range())) {
return std::nullopt;
}
return indices;
}
static void node_geo_exec(GeoNodeExecParams params)
{
GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
const Field<bool> selection_field = params.extract_input<Field<bool>>("Selection");
const Field<int> group_id_field = params.extract_input<Field<int>>("Group ID");
const Field<float> weight_field = params.extract_input<Field<float>>("Sort Weight");
const bke::AttrDomain domain = bke::AttrDomain(params.node().custom1);
const bke::AnonymousAttributePropagationInfo propagation_info =
params.get_output_propagation_info("Geometry");
std::atomic<bool> has_reorder = false;
std::atomic<bool> has_unsupported = false;
geometry_set.modify_geometry_sets([&](GeometrySet &geometry_set) {
for (const auto [type, domains] : geometry::components_supported_reordering().items()) {
const bke::GeometryComponent *src_component = geometry_set.get_component(type);
if (src_component == nullptr || src_component->is_empty()) {
continue;
}
if (!domains.contains(domain)) {
has_unsupported = true;
continue;
}
has_reorder = true;
const std::optional<Array<int>> indices = sorted_indices(
*src_component, selection_field, group_id_field, weight_field, domain);
if (!indices.has_value()) {
continue;
}
bke::GeometryComponentPtr dst_component = geometry::reordered_component(
*src_component, *indices, domain, propagation_info);
geometry_set.remove(type);
geometry_set.add(*dst_component.get());
}
});
if (has_unsupported && !has_reorder) {
params.error_message_add(NodeWarningType::Info,
TIP_("Domain and geometry type combination is unsupported"));
}
params.set_output("Geometry", std::move(geometry_set));
}
template<typename T>
static Vector<EnumPropertyItem> items_value_in(const Span<T> values,
const EnumPropertyItem *src_items)
{
Vector<EnumPropertyItem> items;
for (const EnumPropertyItem *item = src_items; item->identifier != nullptr; item++) {
if (values.contains(T(item->value))) {
items.append(*item);
}
}
items.append({0, nullptr, 0, nullptr, nullptr});
return items;
}
static void node_rna(StructRNA *srna)
{
static const Vector<EnumPropertyItem> supported_items = items_value_in<bke::AttrDomain>(
{bke::AttrDomain::Point,
bke::AttrDomain::Edge,
bke::AttrDomain::Face,
bke::AttrDomain::Curve,
bke::AttrDomain::Instance},
rna_enum_attribute_domain_items);
RNA_def_node_enum(srna,
"domain",
"Domain",
"",
supported_items.data(),
NOD_inline_enum_accessors(custom1),
int(bke::AttrDomain::Point));
}
static void node_register()
{
static bNodeType ntype;
geo_node_type_base(&ntype, GEO_NODE_SORT_ELEMENTS, "Sort Elements", NODE_CLASS_GEOMETRY);
ntype.declare = node_declare;
ntype.initfunc = node_init;
ntype.geometry_node_execute = node_geo_exec;
ntype.draw_buttons = node_layout;
nodeRegisterType(&ntype);
node_rna(ntype.rna_ext.srna);
}
NOD_REGISTER_NODE(node_register)
} // namespace blender::nodes::node_geo_sort_elements_cc