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();
 }