Fix T98960: Baking to active color attribute uses wrong layer

Baking assumed that color attributes could only have two configurations:
float color data type on vertices, or byte color type on face corners.
In reality the options can be combined to make four total options.
This commit handles the four cases explicitly with a somewhat
more scaleable approach (though this should really be C++ code).

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

1 files changed, 62 insertions, 32 deletions

diff --git a/source/blender/editors/object/object_bake_api.c b/source/blender/editors/object/object_bake_api.c
index a7379d7e492..82f31536be8 100644
--- a/source/blender/editors/object/object_bake_api.c
+++ b/source/blender/editors/object/object_bake_api.c
@@ -21,6 +21,8 @@
 #include "BLI_path_util.h"
 #include "BLI_string.h"
 
+#include "BKE_attribute.h"
+#include "BKE_callbacks.h"
 #include "BKE_context.h"
 #include "BKE_global.h"
 #include "BKE_image.h"
@@ -446,10 +448,7 @@ static bool bake_object_check(ViewLayer *view_layer,
   }
 
   if (target == R_BAKE_TARGET_VERTEX_COLORS) {
-    MPropCol *mcol = CustomData_get_layer(&me->vdata, CD_PROP_COLOR);
-    const bool mcol_valid = (mcol != NULL);
-    MLoopCol *mloopcol = CustomData_get_layer(&me->ldata, CD_PROP_BYTE_COLOR);
-    if (mloopcol == NULL && !mcol_valid) {
+    if (BKE_id_attributes_active_color_get(&me->id) == NULL) {
       BKE_reportf(reports,
                   RPT_ERROR,
                   "No vertex colors layer found in the object \"%s\"",
@@ -943,10 +942,8 @@ static bool bake_targets_init_vertex_colors(BakeTargets *targets, Object *ob, Re
   }
 
   Mesh *me = ob->data;
-  MPropCol *mcol = CustomData_get_layer(&me->vdata, CD_PROP_COLOR);
-  const bool mcol_valid = (mcol != NULL);
-  MLoopCol *mloopcol = CustomData_get_layer(&me->ldata, CD_PROP_BYTE_COLOR);
-  if (mloopcol == NULL && !mcol_valid) {
+
+  if (BKE_id_attributes_active_color_get(&me->id) == NULL) {
     BKE_report(reports, RPT_ERROR, "No vertex colors layer found to bake to");
     return false;
   }
@@ -996,10 +993,10 @@ static int find_original_loop(const Mesh *me_orig,
   return ORIGINDEX_NONE;
 }
 
-static void bake_targets_populate_pixels_vertex_colors(BakeTargets *targets,
-                                                       Object *ob,
-                                                       Mesh *me_eval,
-                                                       BakePixel *pixel_array)
+static void bake_targets_populate_pixels_color_attributes(BakeTargets *targets,
+                                                          Object *ob,
+                                                          Mesh *me_eval,
+                                                          BakePixel *pixel_array)
 {
   Mesh *me = ob->data;
   const int pixels_num = targets->pixels_num;
@@ -1094,19 +1091,42 @@ static void bake_result_add_to_rgba(float rgba[4], const float *result, const in
   }
 }
 
+static void convert_float_color_to_byte_color(const MPropCol *float_colors,
+                                              const int num,
+                                              const bool is_noncolor,
+                                              MLoopCol *byte_colors)
+{
+  if (is_noncolor) {
+    for (int i = 0; i < num; i++) {
+      unit_float_to_uchar_clamp_v4(&byte_colors->r, float_colors[i].color);
+    }
+  }
+  else {
+    for (int i = 0; i < num; i++) {
+      linearrgb_to_srgb_uchar4(&byte_colors[i].r, float_colors[i].color);
+    }
+  }
+}
+
 static bool bake_targets_output_vertex_colors(BakeTargets *targets, Object *ob)
 {
   Mesh *me = ob->data;
-  MPropCol *mcol = CustomData_get_layer(&me->vdata, CD_PROP_COLOR);
-  const bool mcol_valid = (mcol != NULL);
-  MLoopCol *mloopcol = CustomData_get_layer(&me->ldata, CD_PROP_BYTE_COLOR);
+  CustomDataLayer *active_color_layer = BKE_id_attributes_active_color_get(&me->id);
+  BLI_assert(active_color_layer != NULL);
+  const AttributeDomain domain = BKE_id_attribute_domain(&me->id, active_color_layer);
+
   const int channels_num = targets->channels_num;
+  const bool is_noncolor = targets->is_noncolor;
   const float *result = targets->result;
 
-  if (mcol_valid) {
+  if (domain == ATTR_DOMAIN_POINT) {
     const int totvert = me->totvert;
     const int totloop = me->totloop;
 
+    MPropCol *mcol = active_color_layer->type == CD_PROP_COLOR ?
+                         active_color_layer->data :
+                         MEM_malloc_arrayN(totvert, sizeof(MPropCol), __func__);
+
     /* Accumulate float vertex colors in scene linear color space. */
     int *num_loops_for_vertex = MEM_callocN(sizeof(int) * me->totvert, "num_loops_for_vertex");
     memset(mcol, 0, sizeof(MPropCol) * me->totvert);
@@ -1125,25 +1145,35 @@ static bool bake_targets_output_vertex_colors(BakeTargets *targets, Object *ob)
       }
     }
 
+    if (mcol != active_color_layer->data) {
+      convert_float_color_to_byte_color(mcol, totvert, is_noncolor, active_color_layer->data);
+      MEM_freeN(mcol);
+    }
+
     MEM_SAFE_FREE(num_loops_for_vertex);
   }
-  else {
-    /* Byte loop colors in sRGB colors space. */
-    MLoop *mloop = me->mloop;
-    const int totloop = me->totloop;
-    const bool is_noncolor = targets->is_noncolor;
-
-    for (int i = 0; i < totloop; i++, mloop++, mloopcol++) {
-      float rgba[4];
-      zero_v4(rgba);
-      bake_result_add_to_rgba(rgba, &result[i * channels_num], channels_num);
-
-      if (is_noncolor) {
-        unit_float_to_uchar_clamp_v4(&mloopcol->r, rgba);
+  else if (domain == ATTR_DOMAIN_CORNER) {
+    switch (active_color_layer->type) {
+      case CD_PROP_COLOR: {
+        MPropCol *colors = active_color_layer->data;
+        for (int i = 0; i < me->totloop; i++) {
+          zero_v4(colors[i].color);
+          bake_result_add_to_rgba(colors[i].color, &result[i * channels_num], channels_num);
+        }
+        break;
       }
-      else {
-        linearrgb_to_srgb_uchar4(&mloopcol->r, rgba);
+      case CD_PROP_BYTE_COLOR: {
+        MLoopCol *colors = active_color_layer->data;
+        for (int i = 0; i < me->totloop; i++) {
+          MPropCol color;
+          zero_v4(color.color);
+          bake_result_add_to_rgba(color.color, &result[i * channels_num], channels_num);
+          convert_float_color_to_byte_color(&color, 1, is_noncolor, &colors[i]);
+        }
+        break;
       }
+      default:
+        BLI_assert_unreachable();
     }
   }
 
@@ -1197,7 +1227,7 @@ static void bake_targets_populate_pixels(const BakeAPIRender *bkr,
                                          BakePixel *pixel_array)
 {
   if (bkr->target == R_BAKE_TARGET_VERTEX_COLORS) {
-    bake_targets_populate_pixels_vertex_colors(targets, ob, me_eval, pixel_array);
+    bake_targets_populate_pixels_color_attributes(targets, ob, me_eval, pixel_array);
   }
   else {
     RE_bake_pixels_populate(me_eval, pixel_array, targets->pixels_num, targets, bkr->uv_layer);