Cleanup: Deduplicate sampling at index in geometry nodes

The "Evaluate at Index" and "Sample Index" nodes are exactly the same
once they retrieve the values to copy and the indices (apart from the
clamping option anyway). This also allows devirtualizing the index input
in the evaluate at index node like the sample index node.
This commit is contained in:
Hans Goudey 2023-05-03 11:46:18 -04:00
parent 09c1a93c21
commit 7f040099e3
3 changed files with 57 additions and 54 deletions

View File

@ -149,4 +149,9 @@ void simulation_state_to_values(const Span<NodeSimulationItem> node_simulation_i
const bNode &sim_output_node,
Span<void *> r_output_values);
void copy_with_checked_indices(const GVArray &src,
const VArray<int> &indices,
IndexMask mask,
GMutableSpan dst);
} // namespace blender::nodes

View File

@ -44,26 +44,9 @@ GVArray FieldAtIndexInput::get_varray_for_context(const bke::GeometryFieldContex
index_evaluator.evaluate();
const VArray<int> indices = index_evaluator.get_evaluated<int>(0);
GVArray output_array;
attribute_math::convert_to_static_type(*type_, [&](auto dummy) {
using T = decltype(dummy);
Array<T> dst_array(mask.min_array_size());
VArray<T> src_values = values.typed<T>();
threading::parallel_for(mask.index_range(), 1024, [&](const IndexRange range) {
for (const int i : mask.slice(range)) {
const int index = indices[i];
if (src_values.index_range().contains(index)) {
dst_array[i] = src_values[index];
}
else {
dst_array[i] = {};
}
}
});
output_array = VArray<T>::ForContainer(std::move(dst_array));
});
return output_array;
GArray<> dst_array(values.type(), mask.min_array_size());
copy_with_checked_indices(values, indices, mask, dst_array);
return GVArray::ForGArray(std::move(dst_array));
}
} // namespace blender::nodes
@ -176,13 +159,13 @@ static void node_geo_exec(GeoNodeExecParams params)
const eAttrDomain domain = eAttrDomain(node.custom1);
const eCustomDataType data_type = eCustomDataType(node.custom2);
Field<int> index_field = params.extract_input<Field<int>>("Index");
attribute_math::convert_to_static_type(data_type, [&](auto dummy) {
using T = decltype(dummy);
static const std::string identifier = "Value_" + identifier_suffix(data_type);
Field<T> value_field = params.extract_input<Field<T>>(identifier);
Field<T> output_field{std::make_shared<FieldAtIndexInput>(
std::move(index_field), std::move(value_field), domain)};
Field<T> output_field{
std::make_shared<FieldAtIndexInput>(params.extract_input<Field<int>>("Index"),
params.extract_input<Field<T>>(identifier),
domain)};
params.set_output(identifier, std::move(output_field));
});
}

View File

@ -11,6 +11,43 @@
#include "node_geometry_util.hh"
namespace blender::nodes {
template<typename T>
void copy_with_checked_indices(const VArray<T> &src,
const VArray<int> &indices,
const IndexMask mask,
MutableSpan<T> dst)
{
const IndexRange src_range = src.index_range();
devirtualize_varray2(src, indices, [&](const auto src, const auto indices) {
threading::parallel_for(mask.index_range(), 4096, [&](IndexRange range) {
for (const int i : mask.slice(range)) {
const int index = indices[i];
if (src_range.contains(index)) {
dst[i] = src[index];
}
else {
dst[i] = {};
}
}
});
});
}
void copy_with_checked_indices(const GVArray &src,
const VArray<int> &indices,
const IndexMask mask,
GMutableSpan dst)
{
attribute_math::convert_to_static_type(src.type(), [&](auto dummy) {
using T = decltype(dummy);
copy_with_checked_indices(src.typed<T>(), indices, mask, dst.typed<T>());
});
}
} // namespace blender::nodes
namespace blender::nodes::node_geo_sample_index_cc {
NODE_STORAGE_FUNCS(NodeGeometrySampleIndex);
@ -132,28 +169,6 @@ static const GeometryComponent *find_source_component(const GeometrySet &geometr
return nullptr;
}
template<typename T>
void copy_with_indices(const VArray<T> &src,
const VArray<int> &indices,
const IndexMask mask,
MutableSpan<T> dst)
{
const IndexRange src_range = src.index_range();
devirtualize_varray2(src, indices, [&](const auto src, const auto indices) {
threading::parallel_for(mask.index_range(), 4096, [&](IndexRange range) {
for (const int i : mask.slice(range)) {
const int index = indices[i];
if (src_range.contains(index)) {
dst[i] = src[index];
}
else {
dst[i] = {};
}
}
});
});
}
template<typename T>
void copy_with_clamped_indices(const VArray<T> &src,
const VArray<int> &indices,
@ -233,15 +248,15 @@ class SampleIndexFunction : public mf::MultiFunction {
return;
}
attribute_math::convert_to_static_type(type, [&](auto dummy) {
using T = decltype(dummy);
if (clamp_) {
if (clamp_) {
attribute_math::convert_to_static_type(type, [&](auto dummy) {
using T = decltype(dummy);
copy_with_clamped_indices(src_data_->typed<T>(), indices, mask, dst.typed<T>());
}
else {
copy_with_indices(src_data_->typed<T>(), indices, mask, dst.typed<T>());
}
});
});
}
else {
copy_with_checked_indices(*src_data_, indices, mask, dst);
}
}
};