Geometry Nodes: support fixed pivot axis in Align Rotation to Vector node

When the pivot axis is not set to auto, the node will try to align the rotation
to vector as best as possible, given the selected rotation axis.

Ref T85211.

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

1 files changed, 81 insertions, 24 deletions

diff --git a/source/blender/nodes/geometry/nodes/node_geo_align_rotation_to_vector.cc b/source/blender/nodes/geometry/nodes/node_geo_align_rotation_to_vector.cc
index eac77b25bd6..bdec3599bfa 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_align_rotation_to_vector.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_align_rotation_to_vector.cc
@@ -34,30 +34,12 @@ static bNodeSocketTemplate geo_node_align_rotation_to_vector_out[] = {
 
 namespace blender::nodes {
 
-static void align_rotations_on_component(GeometryComponent &component,
-                                         const GeoNodeExecParams &params)
+static void align_rotations_auto_pivot(const Float3ReadAttribute &vectors,
+                                       const FloatReadAttribute &factors,
+                                       const float3 local_main_axis,
+                                       MutableSpan<float3> rotations)
 {
-  const bNode &node = params.node();
-  const NodeGeometryAlignRotationToVector &storage = *(const NodeGeometryAlignRotationToVector *)
-                                                          node.storage;
-
-  OutputAttributePtr rotation_attribute = component.attribute_try_get_for_output(
-      "rotation", ATTR_DOMAIN_POINT, CD_PROP_FLOAT3);
-  if (!rotation_attribute) {
-    return;
-  }
-  MutableSpan<float3> rotations = rotation_attribute->get_span<float3>();
-
-  FloatReadAttribute factors = params.get_input_attribute<float>(
-      "Factor", component, ATTR_DOMAIN_POINT, 1.0f);
-  Float3ReadAttribute vectors = params.get_input_attribute<float3>(
-      "Vector", component, ATTR_DOMAIN_POINT, {0, 0, 1});
-
-  float3 main_axis{0, 0, 0};
-  main_axis[storage.axis] = 1;
-
-  const int domain_size = component.attribute_domain_size(ATTR_DOMAIN_POINT);
-  for (const int i : IndexRange(domain_size)) {
+  for (const int i : IndexRange(vectors.size())) {
     const float3 vector = vectors[i];
     if (is_zero_v3(vector)) {
       continue;
@@ -66,7 +48,7 @@ static void align_rotations_on_component(GeometryComponent &component,
     float old_rotation[3][3];
     eul_to_mat3(old_rotation, rotations[i]);
     float3 old_axis;
-    mul_v3_m3v3(old_axis, old_rotation, main_axis);
+    mul_v3_m3v3(old_axis, old_rotation, local_main_axis);
 
     const float3 new_axis = vector.normalized();
     const float3 rotation_axis = float3::cross_high_precision(old_axis, new_axis);
@@ -84,6 +66,81 @@ static void align_rotations_on_component(GeometryComponent &component,
 
     rotations[i] = new_rotation;
   }
+}
+
+static void align_rotations_fixed_pivot(const Float3ReadAttribute &vectors,
+                                        const FloatReadAttribute &factors,
+                                        const float3 local_main_axis,
+                                        const float3 local_pivot_axis,
+                                        MutableSpan<float3> rotations)
+{
+  if (local_main_axis == local_pivot_axis) {
+    /* Can't compute any meaningful rotation angle in this case. */
+    return;
+  }
+
+  for (const int i : IndexRange(vectors.size())) {
+    const float3 vector = vectors[i];
+    if (is_zero_v3(vector)) {
+      continue;
+    }
+
+    float old_rotation[3][3];
+    eul_to_mat3(old_rotation, rotations[i]);
+    float3 old_axis;
+    mul_v3_m3v3(old_axis, old_rotation, local_main_axis);
+    float3 pivot_axis;
+    mul_v3_m3v3(pivot_axis, old_rotation, local_pivot_axis);
+
+    float full_angle = angle_signed_on_axis_v3v3_v3(vector, old_axis, pivot_axis);
+    if (full_angle > M_PI) {
+      /* Make sure the point is rotated as little as possible. */
+      full_angle -= 2.0f * M_PI;
+    }
+    const float angle = factors[i] * full_angle;
+
+    float rotation[3][3];
+    axis_angle_to_mat3(rotation, pivot_axis, angle);
+
+    float new_rotation_matrix[3][3];
+    mul_m3_m3m3(new_rotation_matrix, rotation, old_rotation);
+
+    float3 new_rotation;
+    mat3_to_eul(new_rotation, new_rotation_matrix);
+
+    rotations[i] = new_rotation;
+  }
+}
+
+static void align_rotations_on_component(GeometryComponent &component,
+                                         const GeoNodeExecParams &params)
+{
+  const bNode &node = params.node();
+  const NodeGeometryAlignRotationToVector &storage = *(const NodeGeometryAlignRotationToVector *)
+                                                          node.storage;
+
+  OutputAttributePtr rotation_attribute = component.attribute_try_get_for_output(
+      "rotation", ATTR_DOMAIN_POINT, CD_PROP_FLOAT3);
+  if (!rotation_attribute) {
+    return;
+  }
+  MutableSpan<float3> rotations = rotation_attribute->get_span<float3>();
+
+  FloatReadAttribute factors = params.get_input_attribute<float>(
+      "Factor", component, ATTR_DOMAIN_POINT, 1.0f);
+  Float3ReadAttribute vectors = params.get_input_attribute<float3>(
+      "Vector", component, ATTR_DOMAIN_POINT, {0, 0, 1});
+
+  float3 local_main_axis{0, 0, 0};
+  local_main_axis[storage.axis] = 1;
+  if (storage.pivot_axis == GEO_NODE_ALIGN_ROTATION_TO_VECTOR_PIVOT_AXIS_AUTO) {
+    align_rotations_auto_pivot(vectors, factors, local_main_axis, rotations);
+  }
+  else {
+    float3 local_pivot_axis{0, 0, 0};
+    local_pivot_axis[storage.pivot_axis - 1] = 1;
+    align_rotations_fixed_pivot(vectors, factors, local_main_axis, local_pivot_axis, rotations);
+  }
 
   rotation_attribute.apply_span_and_save();
 }