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 5e8c4fb9edf..be7d7de18f6 100644 --- a/source/blender/nodes/geometry/nodes/node_geo_set_position.cc +++ b/source/blender/nodes/geometry/nodes/node_geo_set_position.cc @@ -2,11 +2,11 @@ * * SPDX-License-Identifier: GPL-2.0-or-later */ -#include "BLI_task.hh" - #include "BKE_curves.hh" #include "BKE_grease_pencil.hh" #include "BKE_instances.hh" +#include "BKE_mesh.hh" +#include "BKE_pointcloud.hh" #include "node_geometry_util.hh" @@ -21,214 +21,137 @@ static void node_declare(NodeDeclarationBuilder &b) b.add_output("Geometry").propagate_all(); } -constexpr GrainSize grain_size{10000}; - -static bool check_positions_are_original(const AttributeAccessor &attributes, - const VArray &in_positions) +static const auto &get_add_fn() { - const bke::AttributeReader positions_read_only = attributes.lookup("position"); - if (positions_read_only.varray.is_span() && in_positions.is_span()) { - return positions_read_only.varray.get_internal_span().data() == - in_positions.get_internal_span().data(); - } - return false; + static const auto fn = mf::build::SI2_SO( + "Add", + [](const float3 a, const float3 b) { return a + b; }, + mf::build::exec_presets::AllSpanOrSingle()); + return fn; } -static void write_offset_positions(const bool positions_are_original, - const IndexMask &selection, - const VArray &in_positions, - const VArray &in_offsets, - VMutableArray &out_positions) +static const auto &get_sub_fn() { - if (positions_are_original) { - if (const std::optional offset = in_offsets.get_if_single()) { - if (math::is_zero(*offset)) { - return; - } - } - } - - MutableVArraySpan out_positions_span = out_positions; - if (positions_are_original) { - devirtualize_varray(in_offsets, [&](const auto in_offsets) { - selection.foreach_index_optimized( - grain_size, [&](const int i) { out_positions_span[i] += in_offsets[i]; }); - }); - } - else { - devirtualize_varray2( - in_positions, in_offsets, [&](const auto in_positions, const auto in_offsets) { - selection.foreach_index_optimized(grain_size, [&](const int i) { - out_positions_span[i] = in_positions[i] + in_offsets[i]; - }); - }); - } - out_positions_span.save(); + static const auto fn = mf::build::SI2_SO( + "Add", + [](const float3 a, const float3 b) { return a - b; }, + mf::build::exec_presets::AllSpanOrSingle()); + return fn; } -static void set_computed_position_and_offset(GeometryComponent &component, - const VArray &in_positions, - const VArray &in_offsets, - const IndexMask &selection, - MutableAttributeAccessor attributes) +static void set_points_position(bke::MutableAttributeAccessor attributes, + const fn::FieldContext &field_context, + const Field &selection_field, + const Field &position_field) { - /* Optimize the case when `in_positions` references the original positions array. */ - switch (component.type()) { - case GeometryComponent::Type::Curve: { - if (attributes.contains("handle_right") && attributes.contains("handle_left")) { - CurveComponent &curve_component = static_cast(component); - Curves &curves_id = *curve_component.get_for_write(); - bke::CurvesGeometry &curves = curves_id.geometry.wrap(); - SpanAttributeWriter handle_right_attribute = - attributes.lookup_or_add_for_write_span("handle_right", AttrDomain::Point); - SpanAttributeWriter handle_left_attribute = - attributes.lookup_or_add_for_write_span("handle_left", AttrDomain::Point); - - AttributeWriter positions = attributes.lookup_for_write("position"); - MutableVArraySpan out_positions_span = positions.varray; - devirtualize_varray2( - in_positions, in_offsets, [&](const auto in_positions, const auto in_offsets) { - selection.foreach_index_optimized(grain_size, [&](const int i) { - const float3 new_position = in_positions[i] + in_offsets[i]; - const float3 delta = new_position - out_positions_span[i]; - handle_right_attribute.span[i] += delta; - handle_left_attribute.span[i] += delta; - out_positions_span[i] = new_position; - }); - }); - - out_positions_span.save(); - positions.finish(); - handle_right_attribute.finish(); - handle_left_attribute.finish(); - - /* Automatic Bezier handles must be recalculated based on the new positions. */ - curves.calculate_bezier_auto_handles(); - break; - } - AttributeWriter positions = attributes.lookup_for_write("position"); - write_offset_positions(check_positions_are_original(attributes, in_positions), - selection, - in_positions, - in_offsets, - positions.varray); - positions.finish(); - break; - } - case GeometryComponent::Type::Instance: { - /* Special case for "position" which is no longer an attribute on instances. */ - auto &instances_component = reinterpret_cast(component); - bke::Instances &instances = *instances_component.get_for_write(); - VMutableArray positions = bke::instance_position_varray_for_write(instances); - write_offset_positions(false, selection, in_positions, in_offsets, positions); - break; - } - default: { - AttributeWriter positions = attributes.lookup_for_write("position"); - write_offset_positions(check_positions_are_original(attributes, in_positions), - selection, - in_positions, - in_offsets, - positions.varray); - positions.finish(); - break; - } - } + bke::try_capture_field_on_geometry(attributes, + field_context, + "position", + bke::AttrDomain::Point, + selection_field, + position_field); } -static void set_position_in_grease_pencil(GreasePencilComponent &grease_pencil_component, +static void set_curves_position(bke::CurvesGeometry &curves, + const fn::FieldContext &field_context, + const Field &selection_field, + const Field &position_field) +{ + MutableAttributeAccessor attributes = curves.attributes_for_write(); + if (attributes.contains("handle_right") && attributes.contains("handle_left")) { + fn::Field delta(fn::FieldOperation::Create( + get_sub_fn(), {position_field, bke::AttributeFieldInput::Create("position")})); + for (const StringRef name : {"handle_left", "handle_right"}) { + bke::try_capture_field_on_geometry( + attributes, + field_context, + name, + bke::AttrDomain::Point, + selection_field, + Field(fn::FieldOperation::Create( + get_add_fn(), {bke::AttributeFieldInput::Create(name), delta}))); + } + } + set_points_position(attributes, field_context, selection_field, position_field); + curves.calculate_bezier_auto_handles(); +} + +static void set_position_in_grease_pencil(GreasePencil &grease_pencil, const Field &selection_field, - const Field &position_field, - const Field &offset_field) + const Field &position_field) { using namespace blender::bke::greasepencil; - GreasePencil &grease_pencil = *grease_pencil_component.get_for_write(); - /* Set position for each layer. */ for (const int layer_index : grease_pencil.layers().index_range()) { Drawing *drawing = bke::greasepencil::get_eval_grease_pencil_layer_drawing_for_write( grease_pencil, layer_index); if (drawing == nullptr || drawing->strokes().points_num() == 0) { continue; } - bke::GreasePencilLayerFieldContext field_context( - grease_pencil, AttrDomain::Point, layer_index); - fn::FieldEvaluator evaluator{field_context, drawing->strokes().points_num()}; - evaluator.set_selection(selection_field); - evaluator.add(position_field); - evaluator.add(offset_field); - evaluator.evaluate(); - - const IndexMask selection = evaluator.get_evaluated_selection_as_mask(); - if (selection.is_empty()) { - continue; - } - - MutableAttributeAccessor attributes = drawing->strokes_for_write().attributes_for_write(); - const VArray positions_input = evaluator.get_evaluated(0); - const VArray offsets_input = evaluator.get_evaluated(1); - set_computed_position_and_offset( - grease_pencil_component, positions_input, offsets_input, selection, attributes); + set_curves_position( + drawing->strokes_for_write(), + bke::GreasePencilLayerFieldContext(grease_pencil, bke::AttrDomain::Point, layer_index), + selection_field, + position_field); drawing->tag_positions_changed(); } } -static void set_position_in_component(GeometrySet &geometry, - GeometryComponent::Type component_type, - const Field &selection_field, - const Field &position_field, - const Field &offset_field) +static void set_instances_position(bke::Instances &instances, + const Field &selection_field, + const Field &position_field) { - const GeometryComponent &component = *geometry.get_component(component_type); - const AttrDomain domain = component.type() == GeometryComponent::Type::Instance ? - AttrDomain::Instance : - AttrDomain::Point; - const int domain_size = component.attribute_domain_size(domain); - if (domain_size == 0) { - return; - } - - bke::GeometryFieldContext field_context{component, domain}; - fn::FieldEvaluator evaluator{field_context, domain_size}; + const bke::InstancesFieldContext context(instances); + fn::FieldEvaluator evaluator(context, instances.instances_num()); evaluator.set_selection(selection_field); - evaluator.add(position_field); - evaluator.add(offset_field); + + /* Use a temporary array for the output to avoid potentially reading from freed memory if + * retrieving the transforms has to make a mutable copy (then we can't depend on the user count + * of the original read-only data). */ + Array result(instances.instances_num()); + evaluator.add_with_destination(position_field, result.as_mutable_span()); evaluator.evaluate(); - const IndexMask selection = evaluator.get_evaluated_selection_as_mask(); - if (selection.is_empty()) { - return; - } - - GeometryComponent &mutable_component = geometry.get_component_for_write(component_type); - MutableAttributeAccessor attributes = *mutable_component.attributes_for_write(); - const VArray positions_input = evaluator.get_evaluated(0); - const VArray offsets_input = evaluator.get_evaluated(1); - set_computed_position_and_offset( - mutable_component, positions_input, offsets_input, selection, attributes); + MutableSpan transforms = instances.transforms_for_write(); + threading::parallel_for(transforms.index_range(), 2048, [&](const IndexRange range) { + for (const int i : range) { + transforms[i].location() = result[i]; + } + }); } static void node_geo_exec(GeoNodeExecParams params) { GeometrySet geometry = params.extract_input("Geometry"); - Field selection_field = params.extract_input>("Selection"); - Field offset_field = params.extract_input>("Offset"); - Field position_field = params.extract_input>("Position"); + const Field selection_field = params.extract_input>("Selection"); + const fn::Field position_field( + fn::FieldOperation::Create(get_add_fn(), + {params.extract_input>("Position"), + params.extract_input>("Offset")})); - if (geometry.has_grease_pencil()) { - set_position_in_grease_pencil(geometry.get_component_for_write(), - selection_field, - position_field, - offset_field); + if (Mesh *mesh = geometry.get_mesh_for_write()) { + set_points_position(mesh->attributes_for_write(), + bke::MeshFieldContext(*mesh, bke::AttrDomain::Point), + selection_field, + position_field); } - - for (const GeometryComponent::Type type : {GeometryComponent::Type::Mesh, - GeometryComponent::Type::PointCloud, - GeometryComponent::Type::Curve, - GeometryComponent::Type::Instance}) - { - if (geometry.has(type)) { - set_position_in_component(geometry, type, selection_field, position_field, offset_field); - } + if (PointCloud *point_cloud = geometry.get_pointcloud_for_write()) { + set_points_position(point_cloud->attributes_for_write(), + bke::PointCloudFieldContext(*point_cloud), + selection_field, + position_field); + } + if (Curves *curves_id = geometry.get_curves_for_write()) { + bke::CurvesGeometry &curves = curves_id->geometry.wrap(); + set_curves_position(curves, + bke::CurvesFieldContext(curves, bke::AttrDomain::Point), + selection_field, + position_field); + } + if (GreasePencil *grease_pencil = geometry.get_grease_pencil_for_write()) { + set_position_in_grease_pencil(*grease_pencil, selection_field, position_field); + } + if (bke::Instances *instances = geometry.get_instances_for_write()) { + set_instances_position(*instances, selection_field, position_field); } params.set_output("Geometry", std::move(geometry));