From 95ec6e4dd373f58eebf3119e470ead4156c724b4 Mon Sep 17 00:00:00 2001
From: Jacques Lucke <jacques@blender.org>
Date: Fri, 24 Sep 2021 16:02:59 +0200
Subject: Geometry Nodes: make index field more reusable

Some inputs will be the index field implicitly, so we want this
class to be available outside of `node_geo_input_index.cc`.
---
 source/blender/functions/FN_field.hh     |  9 +++++++++
 source/blender/functions/intern/field.cc | 15 +++++++++++++++
 2 files changed, 24 insertions(+)

(limited to 'source/blender/functions')

diff --git a/source/blender/functions/FN_field.hh b/source/blender/functions/FN_field.hh
index d4375b625ce..3ce0993da59 100644
--- a/source/blender/functions/FN_field.hh
+++ b/source/blender/functions/FN_field.hh
@@ -484,4 +484,13 @@ template<typename T> Field<T> make_constant_field(T value)
 
 GField make_field_constant_if_possible(GField field);
 
+class IndexFieldInput final : public FieldInput {
+ public:
+  IndexFieldInput();
+
+  const GVArray *get_varray_for_context(const FieldContext &context,
+                                        IndexMask mask,
+                                        ResourceScope &scope) const final;
+};
+
 }  // namespace blender::fn
diff --git a/source/blender/functions/intern/field.cc b/source/blender/functions/intern/field.cc
index 599e4d4595a..fbd35c2c377 100644
--- a/source/blender/functions/intern/field.cc
+++ b/source/blender/functions/intern/field.cc
@@ -526,6 +526,21 @@ const GVArray *FieldContext::get_varray_for_input(const FieldInput &field_input,
   return field_input.get_varray_for_context(*this, mask, scope);
 }
 
+IndexFieldInput::IndexFieldInput() : FieldInput(CPPType::get<int>(), "Index")
+{
+}
+
+const GVArray *IndexFieldInput::get_varray_for_context(const fn::FieldContext &UNUSED(context),
+                                                       IndexMask mask,
+                                                       ResourceScope &scope) const
+{
+  /* TODO: Investigate a similar method to IndexRange::as_span() */
+  auto index_func = [](int i) { return i; };
+  return &scope.construct<
+      fn::GVArray_For_EmbeddedVArray<int, VArray_For_Func<int, decltype(index_func)>>>(
+      mask.min_array_size(), mask.min_array_size(), index_func);
+}
+
 /* --------------------------------------------------------------------
  * FieldOperation.
  */
-- 
cgit v1.2.3