Functions: add two more customizable multi-functions

3 files changed, 145 insertions, 25 deletions

diff --git a/source/blender/functions/FN_multi_function_builder.hh b/source/blender/functions/FN_multi_function_builder.hh
index 08d4ec672cd..9fcf31443b2 100644
--- a/source/blender/functions/FN_multi_function_builder.hh
+++ b/source/blender/functions/FN_multi_function_builder.hh
@@ -36,15 +36,15 @@ namespace fn {
  * 2. single output (SO) of type Out1
  *
  * This example creates a function that adds 10 to the incoming values:
- *  CustomFunction_SI_SO<int, int> fn("add 10", [](int value) { return value + 10; });
+ *  CustomMF_SI_SO<int, int> fn("add 10", [](int value) { return value + 10; });
  */
-template<typename In1, typename Out1> class CustomFunction_SI_SO : public MultiFunction {
+template<typename In1, typename Out1> class CustomMF_SI_SO : public MultiFunction {
  private:
   using FunctionT = std::function<void(IndexMask, VSpan<In1>, MutableSpan<Out1>)>;
   FunctionT m_function;
 
  public:
-  CustomFunction_SI_SO(StringRef name, FunctionT function) : m_function(std::move(function))
+  CustomMF_SI_SO(StringRef name, FunctionT function) : m_function(std::move(function))
   {
     MFSignatureBuilder signature = this->get_builder(name);
     signature.single_input<In1>("In1");
@@ -52,8 +52,8 @@ template<typename In1, typename Out1> class CustomFunction_SI_SO : public MultiF
   }
 
   template<typename ElementFuncT>
-  CustomFunction_SI_SO(StringRef name, ElementFuncT element_fn)
-      : CustomFunction_SI_SO(name, CustomFunction_SI_SO::create_function(element_fn))
+  CustomMF_SI_SO(StringRef name, ElementFuncT element_fn)
+      : CustomMF_SI_SO(name, CustomMF_SI_SO::create_function(element_fn))
   {
   }
 
@@ -79,13 +79,13 @@ template<typename In1, typename Out1> class CustomFunction_SI_SO : public MultiF
  * 3. single output (SO) of type Out1
  */
 template<typename In1, typename In2, typename Out1>
-class CustomFunction_SI_SI_SO : public MultiFunction {
+class CustomMF_SI_SI_SO : public MultiFunction {
  private:
   using FunctionT = std::function<void(IndexMask, VSpan<In1>, VSpan<In2>, MutableSpan<Out1>)>;
   FunctionT m_function;
 
  public:
-  CustomFunction_SI_SI_SO(StringRef name, FunctionT function) : m_function(std::move(function))
+  CustomMF_SI_SI_SO(StringRef name, FunctionT function) : m_function(std::move(function))
   {
     MFSignatureBuilder signature = this->get_builder(name);
     signature.single_input<In1>("In1");
@@ -94,8 +94,8 @@ class CustomFunction_SI_SI_SO : public MultiFunction {
   }
 
   template<typename ElementFuncT>
-  CustomFunction_SI_SI_SO(StringRef name, ElementFuncT element_fn)
-      : CustomFunction_SI_SI_SO(name, CustomFunction_SI_SI_SO::create_function(element_fn))
+  CustomMF_SI_SI_SO(StringRef name, ElementFuncT element_fn)
+      : CustomMF_SI_SI_SO(name, CustomMF_SI_SI_SO::create_function(element_fn))
   {
   }
 
@@ -109,7 +109,7 @@ class CustomFunction_SI_SI_SO : public MultiFunction {
   void call(IndexMask mask, MFParams params, MFContext UNUSED(context)) const override
   {
     VSpan<In1> in1 = params.readonly_single_input<In1>(0);
-    VSpan<In2> in2 = params.readonly_single_input<In1>(1);
+    VSpan<In2> in2 = params.readonly_single_input<In2>(1);
     MutableSpan<Out1> out1 = params.uninitialized_single_output<Out1>(2);
     m_function(mask, in1, in2, out1);
   }
@@ -117,23 +117,75 @@ class CustomFunction_SI_SI_SO : public MultiFunction {
 
 /**
  * Generates a multi-function with the following parameters:
+ * 1. single input (SI) of type In1
+ * 2. single input (SI) of type In2
+ * 3. single input (SI) of type In3
+ * 4. single output (SO) of type Out1
+ */
+template<typename In1, typename In2, typename In3, typename Out1>
+class CustomMF_SI_SI_SI_SO : public MultiFunction {
+ private:
+  using FunctionT =
+      std::function<void(IndexMask, VSpan<In1>, VSpan<In2>, VSpan<In3>, MutableSpan<Out1>)>;
+  FunctionT m_function;
+
+ public:
+  CustomMF_SI_SI_SI_SO(StringRef name, FunctionT function) : m_function(std::move(function))
+  {
+    MFSignatureBuilder signature = this->get_builder(name);
+    signature.single_input<In1>("In1");
+    signature.single_input<In2>("In2");
+    signature.single_input<In3>("In3");
+    signature.single_output<Out1>("Out1");
+  }
+
+  template<typename ElementFuncT>
+  CustomMF_SI_SI_SI_SO(StringRef name, ElementFuncT element_fn)
+      : CustomMF_SI_SI_SI_SO(name, CustomMF_SI_SI_SI_SO::create_function(element_fn))
+  {
+  }
+
+  template<typename ElementFuncT> static FunctionT create_function(ElementFuncT element_fn)
+  {
+    return [=](IndexMask mask,
+               VSpan<In1> in1,
+               VSpan<In2> in2,
+               VSpan<In3> in3,
+               MutableSpan<Out1> out1) {
+      mask.foreach_index(
+          [&](uint i) { new ((void *)&out1[i]) Out1(element_fn(in1[i], in2[i], in3[i])); });
+    };
+  }
+
+  void call(IndexMask mask, MFParams params, MFContext UNUSED(context)) const override
+  {
+    VSpan<In1> in1 = params.readonly_single_input<In1>(0);
+    VSpan<In2> in2 = params.readonly_single_input<In2>(1);
+    VSpan<In3> in3 = params.readonly_single_input<In3>(2);
+    MutableSpan<Out1> out1 = params.uninitialized_single_output<Out1>(3);
+    m_function(mask, in1, in2, in3, out1);
+  }
+};
+
+/**
+ * Generates a multi-function with the following parameters:
  * 1. single mutable (SM) of type Mut1
  */
-template<typename Mut1> class CustomFunction_SM : public MultiFunction {
+template<typename Mut1> class CustomMF_SM : public MultiFunction {
  private:
   using FunctionT = std::function<void(IndexMask, MutableSpan<Mut1>)>;
   FunctionT m_function;
 
  public:
-  CustomFunction_SM(StringRef name, FunctionT function) : m_function(std::move(function))
+  CustomMF_SM(StringRef name, FunctionT function) : m_function(std::move(function))
   {
     MFSignatureBuilder signature = this->get_builder(name);
     signature.single_mutable<Mut1>("Mut1");
   }
 
   template<typename ElementFuncT>
-  CustomFunction_SM(StringRef name, ElementFuncT element_fn)
-      : CustomFunction_SM(name, CustomFunction_SM::create_function(element_fn))
+  CustomMF_SM(StringRef name, ElementFuncT element_fn)
+      : CustomMF_SM(name, CustomMF_SM::create_function(element_fn))
   {
   }
 
@@ -151,6 +203,29 @@ template<typename Mut1> class CustomFunction_SM : public MultiFunction {
   }
 };
 
+/**
+ * Generates a multi-function that outputs a constant value.
+ */
+template<typename T> class CustomMF_Constant : public MultiFunction {
+ private:
+  T m_value;
+
+ public:
+  template<typename U> CustomMF_Constant(U &&value) : m_value(std::forward<U>(value))
+  {
+    MFSignatureBuilder signature = this->get_builder("Constant");
+    std::stringstream ss;
+    ss << m_value;
+    signature.single_output<T>(ss.str());
+  }
+
+  void call(IndexMask mask, MFParams params, MFContext UNUSED(context)) const override
+  {
+    MutableSpan<T> output = params.uninitialized_single_output<T>(0);
+    mask.foreach_index([&](uint i) { new (&output[i]) T(m_value); });
+  }
+};
+
 }  // namespace fn
 }  // namespace blender
 
diff --git a/tests/gtests/functions/FN_multi_function_network_test.cc b/tests/gtests/functions/FN_multi_function_network_test.cc
index 2b2cb62716c..a5fe3c4fb07 100644
--- a/tests/gtests/functions/FN_multi_function_network_test.cc
+++ b/tests/gtests/functions/FN_multi_function_network_test.cc
@@ -26,9 +26,8 @@ namespace fn {
 
 TEST(multi_function_network, Test1)
 {
-  CustomFunction_SI_SO<int, int> add_10_fn("add 10", [](int value) { return value + 10; });
-  CustomFunction_SI_SI_SO<int, int, int> multiply_fn("multiply",
-                                                     [](int a, int b) { return a * b; });
+  CustomMF_SI_SO<int, int> add_10_fn("add 10", [](int value) { return value + 10; });
+  CustomMF_SI_SI_SO<int, int, int> multiply_fn("multiply", [](int a, int b) { return a * b; });
 
   MFNetwork network;
 
@@ -171,7 +170,7 @@ class CreateRangeFunction : public MultiFunction {
 
 TEST(multi_function_network, Test2)
 {
-  CustomFunction_SI_SO<int, int> add_3_fn("add 3", [](int value) { return value + 3; });
+  CustomMF_SI_SO<int, int> add_3_fn("add 3", [](int value) { return value + 3; });
 
   ConcatVectorsFunction concat_vectors_fn;
   AppendFunction append_fn;
diff --git a/tests/gtests/functions/FN_multi_function_test.cc b/tests/gtests/functions/FN_multi_function_test.cc
index 12fa17b0d07..903c385ea67 100644
--- a/tests/gtests/functions/FN_multi_function_test.cc
+++ b/tests/gtests/functions/FN_multi_function_test.cc
@@ -219,10 +219,10 @@ TEST(multi_function, GenericAppendFunction)
   EXPECT_EQ(vectors_ref[3][0], 1);
 }
 
-TEST(multi_function, CustomFunction_SI_SO)
+TEST(multi_function, CustomMF_SI_SO)
 {
-  CustomFunction_SI_SO<std::string, uint> fn("strlen",
-                                             [](const std::string &str) { return str.size(); });
+  CustomMF_SI_SO<std::string, uint> fn("strlen",
+                                       [](const std::string &str) { return str.size(); });
 
   Array<std::string> strings = {"hello", "world", "test", "another test"};
   Array<uint> sizes(strings.size(), 0);
@@ -241,9 +241,9 @@ TEST(multi_function, CustomFunction_SI_SO)
   EXPECT_EQ(sizes[3], 12);
 }
 
-TEST(multi_function, CustomFunction_SI_SI_SO)
+TEST(multi_function, CustomMF_SI_SI_SO)
 {
-  CustomFunction_SI_SI_SO<int, int, int> fn("mul", [](int a, int b) { return a * b; });
+  CustomMF_SI_SI_SO<int, int, int> fn("mul", [](int a, int b) { return a * b; });
 
   Array<int> values_a = {4, 6, 8, 9};
   int value_b = 10;
@@ -264,9 +264,36 @@ TEST(multi_function, CustomFunction_SI_SI_SO)
   EXPECT_EQ(outputs[3], 90);
 }
 
-TEST(multi_function, CustomFunction_SM)
+TEST(multi_function, CustomMF_SI_SI_SI_SO)
 {
-  CustomFunction_SM<std::string> fn("AddSuffix", [](std::string &value) { value += " test"; });
+  CustomMF_SI_SI_SI_SO<int, std::string, bool, uint> fn{
+      "custom",
+      [](int a, const std::string &b, bool c) { return (uint)((uint)a + b.size() + (uint)c); }};
+
+  Array<int> values_a = {5, 7, 3, 8};
+  Array<std::string> values_b = {"hello", "world", "another", "test"};
+  Array<bool> values_c = {true, false, false, true};
+  Array<uint> outputs(values_a.size(), 0);
+
+  MFParamsBuilder params(fn, values_a.size());
+  params.add_readonly_single_input(values_a.as_span());
+  params.add_readonly_single_input(values_b.as_span());
+  params.add_readonly_single_input(values_c.as_span());
+  params.add_uninitialized_single_output(outputs.as_mutable_span());
+
+  MFContextBuilder context;
+
+  fn.call({1, 2, 3}, params, context);
+
+  EXPECT_EQ(outputs[0], 0);
+  EXPECT_EQ(outputs[1], 12);
+  EXPECT_EQ(outputs[2], 10);
+  EXPECT_EQ(outputs[3], 13);
+}
+
+TEST(multi_function, CustomMF_SM)
+{
+  CustomMF_SM<std::string> fn("AddSuffix", [](std::string &value) { value += " test"; });
 
   Array<std::string> values = {"a", "b", "c", "d", "e"};
 
@@ -284,5 +311,24 @@ TEST(multi_function, CustomFunction_SM)
   EXPECT_EQ(values[4], "e");
 }
 
+TEST(multi_function, CustomMF_Constant)
+{
+  CustomMF_Constant<int> fn{42};
+
+  Array<int> outputs(4, 0);
+
+  MFParamsBuilder params(fn, outputs.size());
+  params.add_uninitialized_single_output(outputs.as_mutable_span());
+
+  MFContextBuilder context;
+
+  fn.call({0, 2, 3}, params, context);
+
+  EXPECT_EQ(outputs[0], 42);
+  EXPECT_EQ(outputs[1], 0);
+  EXPECT_EQ(outputs[2], 42);
+  EXPECT_EQ(outputs[3], 42);
+}
+
 }  // namespace fn
 }  // namespace blender