Add UV attribute creation to cone node

3 files changed, 100 insertions, 13 deletions

diff --git a/source/blender/nodes/geometry/nodes/node_geo_mesh_primitive_cone.cc b/source/blender/nodes/geometry/nodes/node_geo_mesh_primitive_cone.cc
index 756c419c8a6..1fb5144b981 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_mesh_primitive_cone.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_mesh_primitive_cone.cc
@@ -188,6 +188,94 @@ static int face_total(const GeometryNodeMeshCircleFillType fill_type,
   return face_total;
 }
 
+static void calculate_uvs(Mesh *mesh,
+                          const bool use_top,
+                          const bool use_bottom,
+                          const int verts_num,
+                          const GeometryNodeMeshCircleFillType fill_type)
+{
+  MeshComponent mesh_component;
+  mesh_component.replace(mesh, GeometryOwnershipType::Editable);
+  OutputAttributePtr uv_attribute = mesh_component.attribute_try_get_for_output(
+      "uv", ATTR_DOMAIN_CORNER, CD_PROP_FLOAT2, nullptr);
+  MutableSpan<float2> uvs = uv_attribute->get_span_for_write_only<float2>();
+
+  Array<float2> circle(verts_num);
+  float angle = 0.0f;
+  const float angle_delta = 2.0f * M_PI / static_cast<float>(verts_num);
+  for (const int i : IndexRange(verts_num)) {
+    circle[i].x = std::cos(angle) * 0.225f + 0.25f;
+    circle[i].y = std::sin(angle) * 0.225f + 0.25f;
+    angle += angle_delta;
+  }
+
+  int loop_index = 0;
+  if (use_top) {
+    if (fill_type == GEO_NODE_MESH_CIRCLE_FILL_NGON) {
+      for (const int i : IndexRange(verts_num)) {
+        uvs[loop_index++] = circle[i];
+      }
+    }
+    else if (fill_type == GEO_NODE_MESH_CIRCLE_FILL_TRIANGLE_FAN) {
+      for (const int i : IndexRange(verts_num)) {
+        uvs[loop_index++] = circle[i];
+        uvs[loop_index++] = circle[(i + 1) % verts_num];
+        uvs[loop_index++] = float2(0.25f, 0.25f);
+      }
+    }
+  }
+
+  /* Create side corners and faces. */
+  if (use_top && use_bottom) {
+    const float bottom = (fill_type == GEO_NODE_MESH_CIRCLE_FILL_NONE) ? 0.0f : 0.5f;
+    /* Quads connect the top and bottom. */
+    for (const int i : IndexRange(verts_num)) {
+      const float vert = static_cast<float>(i);
+      uvs[loop_index++] = float2(vert / verts_num, bottom);
+      uvs[loop_index++] = float2((vert + 1.0f) / verts_num, bottom);
+      uvs[loop_index++] = float2((vert + 1.0f) / verts_num, 1.0f);
+      uvs[loop_index++] = float2(vert / verts_num, 1.0f);
+    }
+  }
+  else {
+    /* Triangles connect the top and bottom section. */
+    if (use_top) {
+      for (const int i : IndexRange(verts_num)) {
+        uvs[loop_index++] = circle[i] + float2(0.5f, 0.0f);
+        uvs[loop_index++] = float2(0.75f, 0.25f);
+        uvs[loop_index++] = circle[(i + 1) % verts_num] + float2(0.5f, 0.0f);
+      }
+    }
+    else {
+      BLI_assert(use_bottom);
+      for (const int i : IndexRange(verts_num)) {
+        uvs[loop_index++] = circle[i];
+        uvs[loop_index++] = circle[(i + 1) % verts_num];
+        uvs[loop_index++] = float2(0.25f, 0.25f);
+      }
+    }
+  }
+
+  /* Create bottom corners and faces. */
+  if (use_bottom) {
+    if (fill_type == GEO_NODE_MESH_CIRCLE_FILL_NGON) {
+      for (const int i : IndexRange(verts_num)) {
+        /* Go backwards because of reversed face normal. */
+        uvs[loop_index++] = circle[verts_num - 1 - i] + float2(0.5f, 0.0f);
+      }
+    }
+    else if (fill_type == GEO_NODE_MESH_CIRCLE_FILL_TRIANGLE_FAN) {
+      for (const int i : IndexRange(verts_num)) {
+        uvs[loop_index++] = circle[i] + float2(0.5f, 0.0f);
+        uvs[loop_index++] = float2(0.75f, 0.25f);
+        uvs[loop_index++] = circle[(i + 1) % verts_num] + float2(0.5f, 0.0f);
+      }
+    }
+  }
+
+  uv_attribute.apply_span_and_save();
+}
+
 Mesh *create_cylinder_or_cone_mesh(const float radius_top,
                                    const float radius_bottom,
                                    const float depth,
@@ -441,6 +529,10 @@ Mesh *create_cylinder_or_cone_mesh(const float radius_top,
 
   BKE_mesh_calc_normals(mesh);
 
+  calculate_uvs(mesh, use_top, use_bottom, verts_num, fill_type);
+
+  BLI_assert(BKE_mesh_is_valid(mesh));
+
   return mesh;
 }
 
@@ -467,8 +559,6 @@ static void geo_node_mesh_primitive_cone_exec(GeoNodeExecParams params)
   Mesh *mesh = create_cylinder_or_cone_mesh(
       radius_top, radius_bottom, depth, verts_num, fill_type);
 
-  BLI_assert(BKE_mesh_is_valid(mesh));
-
   /* Transform the mesh so that the base of the cone is at the origin. */
   transform_mesh(mesh, location + float3(0.0f, 0.0f, depth), rotation, float3(1));
 
diff --git a/source/blender/nodes/geometry/nodes/node_geo_mesh_primitive_cylinder.cc b/source/blender/nodes/geometry/nodes/node_geo_mesh_primitive_cylinder.cc
index 21fa6745d61..4b8b6a35058 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_mesh_primitive_cylinder.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_mesh_primitive_cylinder.cc
@@ -83,8 +83,6 @@ static void geo_node_mesh_primitive_cylinder_exec(GeoNodeExecParams params)
 
   transform_mesh(mesh, location, rotation, float3(1));
 
-  BLI_assert(BKE_mesh_is_valid(mesh));
-
   params.set_output("Geometry", GeometrySet::create_with_mesh(mesh));
 }
 
diff --git a/source/blender/nodes/geometry/nodes/node_geo_mesh_primitive_uv_sphere.cc b/source/blender/nodes/geometry/nodes/node_geo_mesh_primitive_uv_sphere.cc
index 33af67f8d60..1fd9ab21a40 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_mesh_primitive_uv_sphere.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_mesh_primitive_uv_sphere.cc
@@ -217,10 +217,10 @@ static void calculate_sphere_faces(MutableSpan<MLoop> loops,
   }
 }
 
-static void calculate_uv_attribute(MeshComponent &mesh_component,
-                                   const float segments,
-                                   const float rings)
+static void calculate_uvs(Mesh *mesh, const float segments, const float rings)
 {
+  MeshComponent mesh_component;
+  mesh_component.replace(mesh, GeometryOwnershipType::Editable);
   OutputAttributePtr uv_attribute = mesh_component.attribute_try_get_for_output(
       "uv", ATTR_DOMAIN_CORNER, CD_PROP_FLOAT2, nullptr);
   MutableSpan<float2> uvs = uv_attribute->get_span_for_write_only<float2>();
@@ -252,6 +252,8 @@ static void calculate_uv_attribute(MeshComponent &mesh_component,
     uvs[loop_index++] = float2((segment + 1.0f) / segments, 1.0f - dy);
     uvs[loop_index++] = float2(segment / segments, 1.0f - dy);
   }
+
+  uv_attribute.apply_span_and_save();
 }
 
 static Mesh *create_uv_sphere_mesh(const float radius, const int segments, const int rings)
@@ -272,6 +274,8 @@ static Mesh *create_uv_sphere_mesh(const float radius, const int segments, const
 
   calculate_sphere_faces(loops, polys, segments, rings);
 
+  calculate_uvs(mesh, segments, rings);
+
   BLI_assert(BKE_mesh_is_valid(mesh));
 
   return mesh;
@@ -294,12 +298,7 @@ static void geo_node_mesh_primitive_uv_sphere_exec(GeoNodeExecParams params)
 
   transform_mesh(mesh, location, rotation, float3(1));
 
-  GeometrySet result_set = GeometrySet::create_with_mesh(mesh);
-
-  calculate_uv_attribute(
-      result_set.get_component_for_write<MeshComponent>(), segments_num, rings_num);
-
-  params.set_output("Geometry", result_set);
+  params.set_output("Geometry", GeometrySet::create_with_mesh(mesh));
 }
 
 }  // namespace blender::nodes