Nodes: Add Refract and Faceforward functions to Vector Maths nodes

Cycles, Eevee, OSL, Geo, Attribute

Based on outdated refract patch D6619 by @cubic_sloth

`refract` and `faceforward` are standard functions in GLSL, OSL and Godot shader languages.
Adding these functions provides Blender shader artists access to these standard functions.

Reviewed By: brecht

Differential Revision: https://developer.blender.org/D10622

6 files changed, 126 insertions, 23 deletions

diff --git a/source/blender/nodes/NOD_math_functions.hh b/source/blender/nodes/NOD_math_functions.hh
index c662220fea7..08496e044c5 100644
--- a/source/blender/nodes/NOD_math_functions.hh
+++ b/source/blender/nodes/NOD_math_functions.hh
@@ -351,6 +351,35 @@ inline bool try_dispatch_float_math_fl3_fl3_fl3_to_fl3(const NodeVectorMathOpera
       return dispatch([](float3 a, float3 b, float3 c) {
         return float3(wrapf(a.x, b.x, c.x), wrapf(a.y, b.y, c.y), wrapf(a.z, b.z, c.z));
       });
+    case NODE_VECTOR_MATH_FACEFORWARD:
+      return dispatch([](float3 a, float3 b, float3 c) { return float3::faceforward(a, b, c); });
+    default:
+      return false;
+  }
+  return false;
+}
+
+/**
+ * This is similar to try_dispatch_float_math_fl_to_fl, just with a different callback signature.
+ */
+template<typename Callback>
+inline bool try_dispatch_float_math_fl3_fl3_fl_to_fl3(const NodeVectorMathOperation operation,
+                                                      Callback &&callback)
+{
+  const FloatMathOperationInfo *info = get_float3_math_operation_info(operation);
+  if (info == nullptr) {
+    return false;
+  }
+
+  /* This is just a utility function to keep the individual cases smaller. */
+  auto dispatch = [&](auto math_function) -> bool {
+    callback(math_function, *info);
+    return true;
+  };
+
+  switch (operation) {
+    case NODE_VECTOR_MATH_REFRACT:
+      return dispatch([](float3 a, float3 b, float c) { return float3::refract(a, b.normalized(), c); });
     default:
       return false;
   }
diff --git a/source/blender/nodes/geometry/nodes/node_geo_attribute_vector_math.cc b/source/blender/nodes/geometry/nodes/node_geo_attribute_vector_math.cc
index aa207860f32..8b6960a05a6 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_attribute_vector_math.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_attribute_vector_math.cc
@@ -40,6 +40,7 @@ static bNodeSocketTemplate geo_node_attribute_vector_math_in[] = {
     {SOCK_FLOAT, N_("B"), 0.0f, 0.0f, 0.0f, 0.0f, -FLT_MAX, FLT_MAX},
     {SOCK_STRING, N_("C")},
     {SOCK_VECTOR, N_("C"), 0.0f, 0.0f, 0.0f, 0.0f, -FLT_MAX, FLT_MAX},
+    {SOCK_FLOAT, N_("C"), 0.0f, 0.0f, 0.0f, 0.0f, -FLT_MAX, FLT_MAX},
     {SOCK_STRING, N_("Result")},
     {-1, ""},
 };
@@ -65,7 +66,8 @@ static bool operation_use_input_b(const NodeVectorMathOperation operation)
 
 static bool operation_use_input_c(const NodeVectorMathOperation operation)
 {
-  return operation == NODE_VECTOR_MATH_WRAP;
+  return ELEM(
+      operation, NODE_VECTOR_MATH_WRAP, NODE_VECTOR_MATH_REFRACT, NODE_VECTOR_MATH_FACEFORWARD);
 }
 
 static void geo_node_attribute_vector_math_layout(uiLayout *layout,
@@ -97,6 +99,14 @@ static CustomDataType operation_get_read_type_b(const NodeVectorMathOperation op
   return CD_PROP_FLOAT3;
 }
 
+static CustomDataType operation_get_read_type_c(const NodeVectorMathOperation operation)
+{
+  if (operation == NODE_VECTOR_MATH_REFRACT) {
+    return CD_PROP_FLOAT;
+  }
+  return CD_PROP_FLOAT3;
+}
+
 static void geo_node_attribute_vector_math_init(bNodeTree *UNUSED(tree), bNode *node)
 {
   NodeAttributeVectorMath *data = (NodeAttributeVectorMath *)MEM_callocN(
@@ -132,6 +142,8 @@ static CustomDataType operation_get_result_type(const NodeVectorMathOperation op
     case NODE_VECTOR_MATH_SINE:
     case NODE_VECTOR_MATH_COSINE:
     case NODE_VECTOR_MATH_TANGENT:
+    case NODE_VECTOR_MATH_REFRACT:
+    case NODE_VECTOR_MATH_FACEFORWARD:
       return CD_PROP_FLOAT3;
     case NODE_VECTOR_MATH_DOT_PRODUCT:
     case NODE_VECTOR_MATH_DISTANCE:
@@ -223,6 +235,37 @@ static void do_math_operation_fl3_fl3_fl3_to_fl3(const Float3ReadAttribute &inpu
   UNUSED_VARS_NDEBUG(success);
 }
 
+static void do_math_operation_fl3_fl3_fl_to_fl3(const Float3ReadAttribute &input_a,
+                                                const Float3ReadAttribute &input_b,
+                                                const FloatReadAttribute &input_c,
+                                                Float3WriteAttribute result,
+                                                const NodeVectorMathOperation operation)
+{
+  const int size = input_a.size();
+
+  Span<float3> span_a = input_a.get_span();
+  Span<float3> span_b = input_b.get_span();
+  Span<float> span_c = input_c.get_span();
+  MutableSpan<float3> span_result = result.get_span_for_write_only();
+
+  bool success = try_dispatch_float_math_fl3_fl3_fl_to_fl3(
+      operation, [&](auto math_function, const FloatMathOperationInfo &UNUSED(info)) {
+        for (const int i : IndexRange(size)) {
+          const float3 a = span_a[i];
+          const float3 b = span_b[i];
+          const float c = span_c[i];
+          const float3 out = math_function(a, b, c);
+          span_result[i] = out;
+        }
+      });
+
+  result.apply_span();
+
+  /* The operation is not supported by this node currently. */
+  BLI_assert(success);
+  UNUSED_VARS_NDEBUG(success);
+}
+
 static void do_math_operation_fl3_fl3_to_fl(const Float3ReadAttribute &input_a,
                                             const Float3ReadAttribute &input_b,
                                             FloatWriteAttribute result,
@@ -364,7 +407,7 @@ static void attribute_vector_math_calc(GeometryComponent &component,
   const bool use_input_b = operation_use_input_b(operation);
   const CustomDataType read_type_b = operation_get_read_type_b(operation);
   const bool use_input_c = operation_use_input_c(operation);
-  const CustomDataType read_type_c = CD_PROP_FLOAT3;
+  const CustomDataType read_type_c = operation_get_read_type_c(operation);
 
   /* The result domain is always point for now. */
   const CustomDataType result_type = operation_get_result_type(operation);
@@ -433,9 +476,14 @@ static void attribute_vector_math_calc(GeometryComponent &component,
       do_math_operation_fl3_to_fl3(*attribute_a, *attribute_result, operation);
       break;
     case NODE_VECTOR_MATH_WRAP:
+    case NODE_VECTOR_MATH_FACEFORWARD:
       do_math_operation_fl3_fl3_fl3_to_fl3(
           *attribute_a, *attribute_b, *attribute_c, *attribute_result, operation);
       break;
+    case NODE_VECTOR_MATH_REFRACT:
+      do_math_operation_fl3_fl3_fl_to_fl3(
+          *attribute_a, *attribute_b, *attribute_c, *attribute_result, operation);
+      break;
   }
   attribute_result.save();
 }
diff --git a/source/blender/nodes/intern/math_functions.cc b/source/blender/nodes/intern/math_functions.cc
index 14de2fce9b3..fb34144abf6 100644
--- a/source/blender/nodes/intern/math_functions.cc
+++ b/source/blender/nodes/intern/math_functions.cc
@@ -205,6 +205,10 @@ const FloatMathOperationInfo *get_float3_math_operation_info(const int operation
       RETURN_OPERATION_INFO("Cosine", "vector_math_cosine");
     case NODE_VECTOR_MATH_TANGENT:
       RETURN_OPERATION_INFO("Tangent", "vector_math_tangent");
+    case NODE_VECTOR_MATH_REFRACT:
+      RETURN_OPERATION_INFO("Refract", "vector_math_refract");
+    case NODE_VECTOR_MATH_FACEFORWARD:
+      RETURN_OPERATION_INFO("Faceforward", "vector_math_faceforward");
   }
 
 #undef RETURN_OPERATION_INFO
diff --git a/source/blender/nodes/intern/node_util.c b/source/blender/nodes/intern/node_util.c
index c7c3ced4e56..00db819721c 100644
--- a/source/blender/nodes/intern/node_util.c
+++ b/source/blender/nodes/intern/node_util.c
@@ -90,6 +90,13 @@ void node_sock_label(bNodeSocket *sock, const char *name)
   BLI_strncpy(sock->label, name, MAX_NAME);
 }
 
+void node_sock_label_clear(bNodeSocket *sock)
+{
+  if (sock->label[0] != '\0') {
+    sock->label[0] = '\0';
+  }
+}
+
 void node_math_update(bNodeTree *UNUSED(ntree), bNode *node)
 {
   bNodeSocket *sock1 = BLI_findlink(&node->inputs, 0);
@@ -127,15 +134,9 @@ void node_math_update(bNodeTree *UNUSED(ntree), bNode *node)
                                  NODE_MATH_SMOOTH_MIN,
                                  NODE_MATH_SMOOTH_MAX));
 
-  if (sock1->label[0] != '\0') {
-    sock1->label[0] = '\0';
-  }
-  if (sock2->label[0] != '\0') {
-    sock2->label[0] = '\0';
-  }
-  if (sock3->label[0] != '\0') {
-    sock3->label[0] = '\0';
-  }
+  node_sock_label_clear(sock1);
+  node_sock_label_clear(sock2);
+  node_sock_label_clear(sock3);
 
   switch (node->custom1) {
     case NODE_MATH_WRAP:
diff --git a/source/blender/nodes/intern/node_util.h b/source/blender/nodes/intern/node_util.h
index 1b542a9420a..9cbb21e02f7 100644
--- a/source/blender/nodes/intern/node_util.h
+++ b/source/blender/nodes/intern/node_util.h
@@ -71,6 +71,7 @@ extern void *node_initexec_curves(struct bNodeExecContext *context,
 
 /**** Updates ****/
 void node_sock_label(struct bNodeSocket *sock, const char *name);
+void node_sock_label_clear(struct bNodeSocket *sock);
 void node_math_update(struct bNodeTree *ntree, struct bNode *node);
 
 /**** Labels ****/
diff --git a/source/blender/nodes/shader/nodes/node_shader_vector_math.cc b/source/blender/nodes/shader/nodes/node_shader_vector_math.cc
index e8396c7cfc1..becc83dd694 100644
--- a/source/blender/nodes/shader/nodes/node_shader_vector_math.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_vector_math.cc
@@ -88,6 +88,10 @@ static const char *gpu_shader_get_name(int mode)
       return "vector_math_cosine";
     case NODE_VECTOR_MATH_TANGENT:
       return "vector_math_tangent";
+    case NODE_VECTOR_MATH_REFRACT:
+      return "vector_math_refract";
+    case NODE_VECTOR_MATH_FACEFORWARD:
+      return "vector_math_faceforward";
   }
 
   return nullptr;
@@ -128,8 +132,10 @@ static void node_shader_update_vector_math(bNodeTree *UNUSED(ntree), bNode *node
                                   NODE_VECTOR_MATH_ABSOLUTE,
                                   NODE_VECTOR_MATH_FRACTION,
                                   NODE_VECTOR_MATH_NORMALIZE));
-  nodeSetSocketAvailability(sockC, ELEM(node->custom1, NODE_VECTOR_MATH_WRAP));
-  nodeSetSocketAvailability(sockScale, node->custom1 == NODE_VECTOR_MATH_SCALE);
+  nodeSetSocketAvailability(
+      sockC, ELEM(node->custom1, NODE_VECTOR_MATH_WRAP, NODE_VECTOR_MATH_FACEFORWARD));
+  nodeSetSocketAvailability(sockScale,
+                            ELEM(node->custom1, NODE_VECTOR_MATH_SCALE, NODE_VECTOR_MATH_REFRACT));
   nodeSetSocketAvailability(sockVector,
                             !ELEM(node->custom1,
                                   NODE_VECTOR_MATH_LENGTH,
@@ -142,13 +148,14 @@ static void node_shader_update_vector_math(bNodeTree *UNUSED(ntree), bNode *node
                                  NODE_VECTOR_MATH_DOT_PRODUCT));
 
   /* Labels */
-  if (sockB->label[0] != '\0') {
-    sockB->label[0] = '\0';
-  }
-  if (sockC->label[0] != '\0') {
-    sockC->label[0] = '\0';
-  }
+  node_sock_label_clear(sockB);
+  node_sock_label_clear(sockC);
+  node_sock_label_clear(sockScale);
   switch (node->custom1) {
+    case NODE_VECTOR_MATH_FACEFORWARD:
+      node_sock_label(sockB, "Incident");
+      node_sock_label(sockC, "Reference");
+      break;
     case NODE_VECTOR_MATH_WRAP:
       node_sock_label(sockB, "Max");
       node_sock_label(sockC, "Min");
@@ -156,6 +163,12 @@ static void node_shader_update_vector_math(bNodeTree *UNUSED(ntree), bNode *node
     case NODE_VECTOR_MATH_SNAP:
       node_sock_label(sockB, "Increment");
       break;
+    case NODE_VECTOR_MATH_REFRACT:
+      node_sock_label(sockScale, "Ior");
+      break;
+    case NODE_VECTOR_MATH_SCALE:
+      node_sock_label(sockScale, "Scale");
+      break;
   }
 }
 
@@ -186,7 +199,6 @@ static const blender::fn::MultiFunction &get_multi_function(
           "Divide", [](float3 a, float3 b) { return float3::safe_divide(a, b); }};
       return fn;
     }
-
     case NODE_VECTOR_MATH_CROSS_PRODUCT: {
       static blender::fn::CustomMF_SI_SI_SO<float3, float3, float3> fn{
           "Cross Product", float3::cross_high_precision};
@@ -205,7 +217,6 @@ static const blender::fn::MultiFunction &get_multi_function(
       static blender::fn::CustomMF_SI_SI_SO<float3, float3, float> fn{"Dot Product", float3::dot};
       return fn;
     }
-
     case NODE_VECTOR_MATH_DISTANCE: {
       static blender::fn::CustomMF_SI_SI_SO<float3, float3, float> fn{"Distance",
                                                                       float3::distance};
@@ -226,7 +237,17 @@ static const blender::fn::MultiFunction &get_multi_function(
           "Normalize", [](float3 a) { return a.normalized(); }};
       return fn;
     }
-
+    case NODE_VECTOR_MATH_REFRACT: {
+      static blender::fn::CustomMF_SI_SI_SI_SO<float3, float3, float, float3> fn{
+          "Refract",
+          [](float3 a, float3 b, float c) { return float3::refract(a, b.normalized(), c); }};
+      return fn;
+    }
+    case NODE_VECTOR_MATH_FACEFORWARD: {
+      static blender::fn::CustomMF_SI_SI_SI_SO<float3, float3, float3, float3> fn{
+          "Faceforward", float3::faceforward};
+      return fn;
+    }
     case NODE_VECTOR_MATH_SNAP: {
       return builder.get_not_implemented_fn();
     }
@@ -263,7 +284,6 @@ static const blender::fn::MultiFunction &get_multi_function(
     case NODE_VECTOR_MATH_TANGENT: {
       return builder.get_not_implemented_fn();
     }
-
     default:
       BLI_assert(false);
       return builder.get_not_implemented_fn();