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

4 files changed, 120 insertions, 26 deletions

diff --git a/source/blender/editors/space_node/drawnode.c b/source/blender/editors/space_node/drawnode.c
index b16ccb9bce4..7d0ff2331d0 100644
--- a/source/blender/editors/space_node/drawnode.c
+++ b/source/blender/editors/space_node/drawnode.c
@@ -3336,6 +3336,7 @@ static void node_geometry_buts_align_rotation_to_vector(uiLayout *layout,
                                                         PointerRNA *ptr)
 {
   uiItemR(layout, ptr, "axis", DEFAULT_FLAGS | UI_ITEM_R_EXPAND, NULL, ICON_NONE);
+  uiItemR(layout, ptr, "pivot_axis", DEFAULT_FLAGS, IFACE_("Pivot"), ICON_NONE);
   uiLayout *col = uiLayoutColumn(layout, false);
   uiItemR(col, ptr, "input_type_factor", DEFAULT_FLAGS, IFACE_("Factor"), ICON_NONE);
   uiItemR(col, ptr, "input_type_vector", DEFAULT_FLAGS, IFACE_("Vector"), ICON_NONE);
diff --git a/source/blender/makesdna/DNA_node_types.h b/source/blender/makesdna/DNA_node_types.h
index edab43c752e..73b6a1b66c6 100644
--- a/source/blender/makesdna/DNA_node_types.h
+++ b/source/blender/makesdna/DNA_node_types.h
@@ -1148,12 +1148,12 @@ typedef struct NodeGeometryRotatePoints {
 typedef struct NodeGeometryAlignRotationToVector {
   /* GeometryNodeAlignRotationToVectorAxis */
   uint8_t axis;
+  /* GeometryNodeAlignRotationToVectorPivotAxis */
+  uint8_t pivot_axis;
 
   /* GeometryNodeAttributeInputMode */
   uint8_t input_type_factor;
   uint8_t input_type_vector;
-
-  char _pad[5];
 } NodeGeometryAlignRotationToVector;
 
 typedef struct NodeGeometryPointScale {
@@ -1660,6 +1660,13 @@ typedef enum GeometryNodeAlignRotationToVectorAxis {
   GEO_NODE_ALIGN_ROTATION_TO_VECTOR_AXIS_Z = 2,
 } GeometryNodeAlignRotationToVectorAxis;
 
+typedef enum GeometryNodeAlignRotationToVectorPivotAxis {
+  GEO_NODE_ALIGN_ROTATION_TO_VECTOR_PIVOT_AXIS_AUTO = 0,
+  GEO_NODE_ALIGN_ROTATION_TO_VECTOR_PIVOT_AXIS_X = 1,
+  GEO_NODE_ALIGN_ROTATION_TO_VECTOR_PIVOT_AXIS_Y = 2,
+  GEO_NODE_ALIGN_ROTATION_TO_VECTOR_PIVOT_AXIS_Z = 3,
+} GeometryNodeAlignRotationToVectorPivotAxis;
+
 typedef enum GeometryNodeTransformSpace {
   GEO_NODE_TRANSFORM_SPACE_ORIGINAL = 0,
   GEO_NODE_TRANSFORM_SPACE_RELATIVE = 1,
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index 7f66d5ff70f..77334f10cd0 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/source/blender/makesrna/intern/rna_nodetree.c
@@ -8793,6 +8793,30 @@ static void def_geo_align_rotation_to_vector(StructRNA *srna)
       {0, NULL, 0, NULL, NULL},
   };
 
+  static const EnumPropertyItem pivot_axis_items[] = {
+      {GEO_NODE_ALIGN_ROTATION_TO_VECTOR_PIVOT_AXIS_AUTO,
+       "AUTO",
+       ICON_NONE,
+       "Auto",
+       "Automatically detect the best rotation axis to rotate towards the vector"},
+      {GEO_NODE_ALIGN_ROTATION_TO_VECTOR_PIVOT_AXIS_X,
+       "X",
+       ICON_NONE,
+       "X",
+       "Rotate around the local X axis"},
+      {GEO_NODE_ALIGN_ROTATION_TO_VECTOR_PIVOT_AXIS_Y,
+       "Y",
+       ICON_NONE,
+       "Y",
+       "Rotate around the local Y axis"},
+      {GEO_NODE_ALIGN_ROTATION_TO_VECTOR_PIVOT_AXIS_Z,
+       "Z",
+       ICON_NONE,
+       "Z",
+       "Rotate around the local Z axis"},
+      {0, NULL, 0, NULL, NULL},
+  };
+
   PropertyRNA *prop;
 
   RNA_def_struct_sdna_from(srna, "NodeGeometryAlignRotationToVector", "storage");
@@ -8802,6 +8826,11 @@ static void def_geo_align_rotation_to_vector(StructRNA *srna)
   RNA_def_property_ui_text(prop, "Axis", "Axis to align to the vector");
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_update");
 
+  prop = RNA_def_property(srna, "pivot_axis", PROP_ENUM, PROP_NONE);
+  RNA_def_property_enum_items(prop, pivot_axis_items);
+  RNA_def_property_ui_text(prop, "Pivot Axis", "Axis to rotate around");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_update");
+
   prop = RNA_def_property(srna, "input_type_factor", PROP_ENUM, PROP_NONE);
   RNA_def_property_enum_items(prop, rna_node_geometry_attribute_input_type_items_float);
   RNA_def_property_ui_text(prop, "Input Type Factor", "");
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();
 }