Geometry Nodes: Add a selection input to the store named attribute node

This patch is a response to T101313.

Adds a selection to the Store Named Attribute node.

If the attribute does not exist unselected parts
are filled with zero values. Otherwise, only the
selected parts are filled.

Differential Revision: https://developer.blender.org/D16237
This commit is contained in:
Charlie Jolly 2023-01-07 14:32:49 +00:00
parent 8a3a1a0c14
commit d397ecae32
Notes: blender-bot 2023-05-31 04:43:10 +02:00
Referenced by commit 2fb829a2de, Fix T103837: crash opening a specific file
Referenced by issue #103837, Regression : crash opening file with uv from GN primitive
3 changed files with 56 additions and 10 deletions

View File

@ -315,6 +315,12 @@ bool try_capture_field_on_geometry(GeometryComponent &component,
const eAttrDomain domain,
const fn::GField &field);
bool try_capture_field_on_geometry(GeometryComponent &component,
const AttributeIDRef &attribute_id,
const eAttrDomain domain,
const fn::Field<bool> &selection,
const fn::GField &field);
/**
* Try to find the geometry domain that the field should be evaluated on. If it is not obvious
* which domain is correct, none is returned.

View File

@ -1,5 +1,7 @@
/* SPDX-License-Identifier: GPL-2.0-or-later */
#include "BLI_array_utils.hh"
#include "BKE_attribute.hh"
#include "BKE_curves.hh"
#include "BKE_geometry_fields.hh"
@ -408,6 +410,7 @@ bool NormalFieldInput::is_equal_to(const fn::FieldNode &other) const
bool try_capture_field_on_geometry(GeometryComponent &component,
const AttributeIDRef &attribute_id,
const eAttrDomain domain,
const fn::Field<bool> &selection,
const fn::GField &field)
{
MutableAttributeAccessor attributes = *component.attributes_for_write();
@ -423,19 +426,44 @@ bool try_capture_field_on_geometry(GeometryComponent &component,
const IndexMask mask{IndexMask(domain_size)};
const bke::AttributeValidator validator = attributes.lookup_validator(attribute_id);
/* Could avoid allocating a new buffer if:
* - We are writing to an attribute that exists already with the correct domain and type.
* - The field does not depend on that attribute (we can't easily check for that yet). */
void *buffer = MEM_mallocN(type.size() * domain_size, __func__);
const std::optional<AttributeMetaData> meta_data = attributes.lookup_meta_data(attribute_id);
const bool attribute_exists = meta_data && meta_data->domain == domain &&
meta_data->data_type == data_type;
/* We are writing to an attribute that exists already with the correct domain and type. */
if (attribute_exists) {
if (GSpanAttributeWriter dst_attribute = attributes.lookup_for_write_span(attribute_id)) {
bke::GeometryFieldContext field_context{component, domain};
const IndexMask mask{IndexMask(domain_size)};
fn::FieldEvaluator evaluator{field_context, &mask};
evaluator.add(validator.validate_field_if_necessary(field));
evaluator.set_selection(selection);
evaluator.evaluate();
const IndexMask selection = evaluator.get_evaluated_selection_as_mask();
array_utils::copy(evaluator.get_evaluated(0), selection, dst_attribute.span);
dst_attribute.finish();
return true;
}
}
/* Could avoid allocating a new buffer if:
* - The field does not depend on that attribute (we can't easily check for that yet). */
void *buffer = MEM_mallocN_aligned(type.size() * domain_size, type.alignment(), __func__);
if (selection.node().depends_on_input() || !fn::evaluate_constant_field(selection)) {
/* If every element might not be selected, the buffer must be initialized. */
type.value_initialize_n(buffer, domain_size);
}
fn::FieldEvaluator evaluator{field_context, &mask};
evaluator.add_with_destination(validator.validate_field_if_necessary(field),
GMutableSpan{type, buffer, domain_size});
evaluator.set_selection(selection);
evaluator.evaluate();
const std::optional<AttributeMetaData> meta_data = attributes.lookup_meta_data(attribute_id);
if (meta_data && meta_data->domain == domain && meta_data->data_type == data_type) {
if (attribute_exists) {
if (GAttributeWriter attribute = attributes.lookup_for_write(attribute_id)) {
attribute.varray.set_all(buffer);
attribute.finish();
@ -457,6 +485,15 @@ bool try_capture_field_on_geometry(GeometryComponent &component,
return false;
}
bool try_capture_field_on_geometry(GeometryComponent &component,
const AttributeIDRef &attribute_id,
const eAttrDomain domain,
const fn::GField &field)
{
const fn::Field<bool> selection = fn::make_constant_field<bool>(true);
return try_capture_field_on_geometry(component, attribute_id, domain, selection, field);
}
std::optional<eAttrDomain> try_detect_field_domain(const GeometryComponent &component,
const fn::GField &field)
{

View File

@ -20,6 +20,7 @@ NODE_STORAGE_FUNCS(NodeGeometryStoreNamedAttribute)
static void node_declare(NodeDeclarationBuilder &b)
{
b.add_input<decl::Geometry>(N_("Geometry"));
b.add_input<decl::Bool>(N_("Selection")).default_value(true).hide_value().field_on_all();
b.add_input<decl::String>(N_("Name")).is_attribute_name();
b.add_input<decl::Vector>(N_("Value"), "Value_Vector").field_on_all();
b.add_input<decl::Float>(N_("Value"), "Value_Float").field_on_all();
@ -52,7 +53,7 @@ static void node_update(bNodeTree *ntree, bNode *node)
const eCustomDataType data_type = eCustomDataType(storage.data_type);
bNodeSocket *socket_geometry = static_cast<bNodeSocket *>(node->inputs.first);
bNodeSocket *socket_name = socket_geometry->next;
bNodeSocket *socket_name = socket_geometry->next->next;
bNodeSocket *socket_vector = socket_name->next;
bNodeSocket *socket_float = socket_vector->next;
bNodeSocket *socket_color4f = socket_float->next;
@ -108,6 +109,8 @@ static void node_geo_exec(GeoNodeExecParams params)
const eCustomDataType data_type = eCustomDataType(storage.data_type);
const eAttrDomain domain = eAttrDomain(storage.domain);
const Field<bool> selection = params.extract_input<Field<bool>>("Selection");
GField field;
switch (data_type) {
case CD_PROP_FLOAT:
@ -147,7 +150,7 @@ static void node_geo_exec(GeoNodeExecParams params)
if (geometry_set.has_instances()) {
GeometryComponent &component = geometry_set.get_component_for_write(
GEO_COMPONENT_TYPE_INSTANCES);
if (!bke::try_capture_field_on_geometry(component, name, domain, field)) {
if (!bke::try_capture_field_on_geometry(component, name, domain, selection, field)) {
if (component.attribute_domain_size(domain) != 0) {
failure.store(true);
}
@ -160,7 +163,7 @@ static void node_geo_exec(GeoNodeExecParams params)
{GEO_COMPONENT_TYPE_MESH, GEO_COMPONENT_TYPE_POINT_CLOUD, GEO_COMPONENT_TYPE_CURVE}) {
if (geometry_set.has(type)) {
GeometryComponent &component = geometry_set.get_component_for_write(type);
if (!bke::try_capture_field_on_geometry(component, name, domain, field)) {
if (!bke::try_capture_field_on_geometry(component, name, domain, selection, field)) {
if (component.attribute_domain_size(domain) != 0) {
failure.store(true);
}