Refactor: Unify vertex and sculpt colors into new

          color attribute system.

This commit removes sculpt colors from experimental
status and unifies it with vertex colors. It
introduces the concept of "color attributes", which
are any attributes that represents colors.  Color
attributes can be represented with byte or floating-point
numbers and can be stored in either vertices or
face corners.

Color attributes share a common namespace
(so you can no longer have a floating-point
sculpt color attribute and a byte vertex color
attribute with the same name).

Note: this commit does not include vertex paint mode,
      which is a separate patch, see:
      https://developer.blender.org/D14179

Differential Revision: https://developer.blender.org/D12587
Ref D12587

4 files changed, 406 insertions, 109 deletions

diff --git a/source/blender/draw/intern/draw_cache_extract.h b/source/blender/draw/intern/draw_cache_extract.h
index aaec47b4dce..4567e470146 100644
--- a/source/blender/draw/intern/draw_cache_extract.h
+++ b/source/blender/draw/intern/draw_cache_extract.h
@@ -308,6 +308,9 @@ typedef struct MeshBatchCache {
   float tot_area, tot_uv_area;
 
   bool no_loose_wire;
+
+  eV3DShadingColorType color_type;
+  bool pbvh_is_drawing;
 } MeshBatchCache;
 
 #define MBC_EDITUV \
diff --git a/source/blender/draw/intern/draw_cache_impl_mesh.c b/source/blender/draw/intern/draw_cache_impl_mesh.c
index 79a080cfccd..e6f34d3dd0d 100644
--- a/source/blender/draw/intern/draw_cache_impl_mesh.c
+++ b/source/blender/draw/intern/draw_cache_impl_mesh.c
@@ -454,29 +454,24 @@ static void mesh_cd_calc_active_mask_uv_layer(const Object *object,
   }
 }
 
-static void mesh_cd_calc_active_vcol_layer(const Object *object,
-                                           const Mesh *me,
-                                           DRW_MeshAttributes *attrs_used)
-{
-  const Mesh *me_final = editmesh_final_or_this(object, me);
-  const CustomData *cd_vdata = mesh_cd_vdata_get_from_mesh(me_final);
-
-  int layer = CustomData_get_active_layer(cd_vdata, CD_PROP_COLOR);
-  if (layer != -1) {
-    drw_mesh_attributes_add_request(attrs_used, CD_PROP_COLOR, layer, ATTR_DOMAIN_POINT);
-  }
-}
-
 static void mesh_cd_calc_active_mloopcol_layer(const Object *object,
                                                const Mesh *me,
                                                DRW_MeshCDMask *cd_used)
 {
   const Mesh *me_final = editmesh_final_or_this(object, me);
+  Mesh me_query = {0};
+
+  const CustomData *cd_vdata = mesh_cd_vdata_get_from_mesh(me_final);
   const CustomData *cd_ldata = mesh_cd_ldata_get_from_mesh(me_final);
 
-  int layer = CustomData_get_active_layer(cd_ldata, CD_MLOOPCOL);
-  if (layer != -1) {
-    cd_used->vcol |= (1 << layer);
+  BKE_id_attribute_copy_domains_temp(ID_ME, cd_vdata, NULL, cd_ldata, NULL, NULL, &me_query.id);
+
+  CustomDataLayer *layer = BKE_id_attributes_active_color_get(&me_query.id);
+  int layer_i = BKE_id_attribute_to_index(
+      &me_query.id, layer, ATTR_DOMAIN_MASK_COLOR, CD_MASK_COLOR_ALL);
+
+  if (layer_i != -1) {
+    cd_used->vcol |= (1UL << (uint)layer_i);
   }
 }
 
@@ -510,6 +505,51 @@ static bool custom_data_match_attribute(const CustomData *custom_data,
   return false;
 }
 
+static uint mesh_cd_calc_gpu_layers_vcol_used(const Mesh *me_query,
+                                              const CustomData *cd_vdata,
+                                              const CustomData *cd_ldata,
+                                              const char name[])
+{
+  CustomDataLayer *layer = NULL;
+  AttributeDomain domain;
+
+  if (name[0]) {
+    int layer_i = 0;
+
+    domain = ATTR_DOMAIN_POINT;
+    layer_i = CustomData_get_named_layer_index(cd_vdata, CD_PROP_COLOR, name);
+    layer_i = layer_i == -1 ? CustomData_get_named_layer_index(cd_vdata, CD_MLOOPCOL, name) :
+                              layer_i;
+
+    if (layer_i == -1) {
+      domain = ATTR_DOMAIN_CORNER;
+      layer_i = layer_i == -1 ? CustomData_get_named_layer_index(cd_ldata, CD_PROP_COLOR, name) :
+                                layer_i;
+      layer_i = layer_i == -1 ? CustomData_get_named_layer_index(cd_ldata, CD_MLOOPCOL, name) :
+                                layer_i;
+    }
+
+    /* Note: this is not the same as the layer_i below. */
+    if (layer_i != -1) {
+      layer = (domain == ATTR_DOMAIN_POINT ? cd_vdata : cd_ldata)->layers + layer_i;
+    }
+  }
+  else {
+    layer = BKE_id_attributes_render_color_get(&me_query->id);
+  }
+
+  if (!layer) {
+    return -1;
+  }
+
+  /* Note: this is the logical index into the color attribute list,
+   * not the customdata index. */
+  int vcol_i = BKE_id_attribute_to_index(
+      (ID *)me_query, layer, ATTR_DOMAIN_MASK_COLOR, CD_MASK_COLOR_ALL);
+
+  return vcol_i;
+}
+
 static DRW_MeshCDMask mesh_cd_calc_used_gpu_layers(const Object *object,
                                                    const Mesh *me,
                                                    struct GPUMaterial **gpumat_array,
@@ -522,6 +562,13 @@ static DRW_MeshCDMask mesh_cd_calc_used_gpu_layers(const Object *object,
   const CustomData *cd_vdata = mesh_cd_vdata_get_from_mesh(me_final);
   const CustomData *cd_edata = mesh_cd_edata_get_from_mesh(me_final);
 
+  /* Create a mesh with final customdata domains
+   * we can query with attribute API. */
+  Mesh me_query = {0};
+
+  BKE_id_attribute_copy_domains_temp(
+      ID_ME, cd_vdata, cd_edata, cd_ldata, cd_pdata, NULL, &me_query.id);
+
   /* See: DM_vertex_attributes_from_gpu for similar logic */
   DRW_MeshCDMask cd_used;
   mesh_cd_layers_type_clear(&cd_used);
@@ -547,8 +594,35 @@ static DRW_MeshCDMask mesh_cd_calc_used_gpu_layers(const Object *object,
             type = CD_MTFACE;
 
             if (layer == -1) {
+              layer = CustomData_get_named_layer(cd_vdata, CD_PROP_COLOR, name);
+              if (layer != -1) {
+                type = CD_PROP_COLOR;
+                domain = ATTR_DOMAIN_POINT;
+              }
+            }
+
+            if (layer == -1) {
+              layer = CustomData_get_named_layer(cd_ldata, CD_PROP_COLOR, name);
+              if (layer != -1) {
+                type = CD_PROP_COLOR;
+                domain = ATTR_DOMAIN_CORNER;
+              }
+            }
+
+            if (layer == -1) {
+              layer = CustomData_get_named_layer(cd_vdata, CD_MLOOPCOL, name);
+              if (layer != -1) {
+                type = CD_MLOOPCOL;
+                domain = ATTR_DOMAIN_POINT;
+              }
+            }
+
+            if (layer == -1) {
               layer = CustomData_get_named_layer(cd_ldata, CD_MLOOPCOL, name);
-              type = CD_MCOL;
+              if (layer != -1) {
+                type = CD_MLOOPCOL;
+                domain = ATTR_DOMAIN_CORNER;
+              }
             }
 
 #if 0 /* Tangents are always from UV's - this will never happen. */
@@ -619,29 +693,33 @@ static DRW_MeshCDMask mesh_cd_calc_used_gpu_layers(const Object *object,
             }
             break;
           }
-          case CD_MCOL: {
-            /* Vertex Color Data */
-            if (layer == -1) {
-              layer = (name[0] != '\0') ? CustomData_get_named_layer(cd_ldata, CD_MLOOPCOL, name) :
-                                          CustomData_get_render_layer(cd_ldata, CD_MLOOPCOL);
-            }
-            if (layer != -1) {
-              cd_used.vcol |= (1 << layer);
-            }
 
-            break;
-          }
           case CD_ORCO: {
             cd_used.orco = 1;
             break;
           }
+
+          /* Note: attr->type will always be CD_PROP_COLOR even for
+           * CD_MLOOPCOL layers, see node_shader_gpu_vertex_color in
+           * node_shader_vertex_color.cc.
+           */
+          case CD_MCOL:
+          case CD_MLOOPCOL:
+          case CD_PROP_COLOR: {
+            int vcol_bit = mesh_cd_calc_gpu_layers_vcol_used(&me_query, cd_vdata, cd_ldata, name);
+
+            if (vcol_bit != -1) {
+              cd_used.vcol |= 1UL << (uint)vcol_bit;
+            }
+
+            break;
+          }
+          case CD_PROP_FLOAT3:
           case CD_PROP_BOOL:
           case CD_PROP_INT8:
           case CD_PROP_INT32:
           case CD_PROP_FLOAT:
-          case CD_PROP_FLOAT2:
-          case CD_PROP_FLOAT3:
-          case CD_PROP_COLOR: {
+          case CD_PROP_FLOAT2: {
             if (layer != -1 && domain != ATTR_DOMAIN_NUM) {
               drw_mesh_attributes_add_request(attributes, type, layer, domain);
             }
@@ -799,6 +877,12 @@ static bool mesh_batch_cache_valid(Object *object, Mesh *me)
     return false;
   }
 
+  if (object->sculpt && object->sculpt->pbvh) {
+    if (cache->pbvh_is_drawing != BKE_pbvh_is_drawing(object->sculpt->pbvh)) {
+      return false;
+    }
+  }
+
   if (cache->is_editmode != (me->edit_mesh != NULL)) {
     return false;
   }
@@ -827,6 +911,10 @@ static void mesh_batch_cache_init(Object *object, Mesh *me)
 
   cache->is_editmode = me->edit_mesh != NULL;
 
+  if (object->sculpt && object->sculpt->pbvh) {
+    cache->pbvh_is_drawing = BKE_pbvh_is_drawing(object->sculpt->pbvh);
+  }
+
   if (cache->is_editmode == false) {
     // cache->edge_len = mesh_render_edges_len_get(me);
     // cache->tri_len = mesh_render_looptri_len_get(me);
@@ -1120,14 +1208,28 @@ static void texpaint_request_active_vcol(MeshBatchCache *cache, Object *object,
 
 static void sculpt_request_active_vcol(MeshBatchCache *cache, Object *object, Mesh *me)
 {
-  DRW_MeshAttributes attrs_needed;
-  drw_mesh_attributes_clear(&attrs_needed);
-  mesh_cd_calc_active_vcol_layer(object, me, &attrs_needed);
+  const Mesh *me_final = editmesh_final_or_this(object, me);
+  const CustomData *cd_vdata = mesh_cd_vdata_get_from_mesh(me_final);
+  const CustomData *cd_ldata = mesh_cd_ldata_get_from_mesh(me_final);
+
+  Mesh me_query = {0};
+  BKE_id_attribute_copy_domains_temp(ID_ME, cd_vdata, NULL, cd_ldata, NULL, NULL, &me_query.id);
 
-  BLI_assert(attrs_needed.num_requests != 0 &&
-             "No MPropCol layer available in Sculpt, but batches requested anyway!");
+  CustomDataLayer *active = BKE_id_attributes_active_color_get(&me_query.id);
+  CustomDataLayer *render = BKE_id_attributes_render_color_get(&me_query.id);
 
-  drw_mesh_attributes_merge(&cache->attr_needed, &attrs_needed, me->runtime.render_mutex);
+  int active_i = BKE_id_attribute_to_index(
+      &me_query.id, active, ATTR_DOMAIN_MASK_COLOR, CD_MASK_COLOR_ALL);
+  int render_i = BKE_id_attribute_to_index(
+      &me_query.id, render, ATTR_DOMAIN_MASK_COLOR, CD_MASK_COLOR_ALL);
+
+  if (active_i >= 0) {
+    cache->cd_used.vcol |= 1UL << (uint)active_i;
+  }
+
+  if (render_i >= 0) {
+    cache->cd_used.vcol |= 1UL << (uint)render_i;
+  }
 }
 
 GPUBatch *DRW_mesh_batch_cache_get_all_verts(Mesh *me)
@@ -1148,6 +1250,7 @@ GPUBatch *DRW_mesh_batch_cache_get_surface(Mesh *me)
 {
   MeshBatchCache *cache = mesh_batch_cache_get(me);
   mesh_batch_cache_request_surface_batches(cache);
+
   return cache->batch.surface;
 }
 
diff --git a/source/blender/draw/intern/mesh_extractors/extract_mesh_vbo_attributes.cc b/source/blender/draw/intern/mesh_extractors/extract_mesh_vbo_attributes.cc
index 4f4aa764fbc..f9e58709c6e 100644
--- a/source/blender/draw/intern/mesh_extractors/extract_mesh_vbo_attributes.cc
+++ b/source/blender/draw/intern/mesh_extractors/extract_mesh_vbo_attributes.cc
@@ -159,17 +159,6 @@ static void init_vbo_for_attribute(const MeshRenderData *mr,
   GPU_vertformat_deinterleave(&format);
   GPU_vertformat_attr_add(&format, attr_name, comp_type, comp_size, fetch_mode);
 
-  /* Ensure Sculpt Vertex Colors are properly aliased. */
-  if (request.cd_type == CD_PROP_COLOR && request.domain == ATTR_DOMAIN_POINT) {
-    CustomData *cd_vdata = get_custom_data_for_domain(mr, ATTR_DOMAIN_POINT);
-    if (request.layer_index == CustomData_get_render_layer(cd_vdata, CD_PROP_COLOR)) {
-      GPU_vertformat_alias_add(&format, "c");
-    }
-    if (request.layer_index == CustomData_get_active_layer(cd_vdata, CD_PROP_COLOR)) {
-      GPU_vertformat_alias_add(&format, "ac");
-    }
-  }
-
   if (build_on_device) {
     GPU_vertbuf_init_build_on_device(vbo, &format, len);
   }
diff --git a/source/blender/draw/intern/mesh_extractors/extract_mesh_vbo_vcol.cc b/source/blender/draw/intern/mesh_extractors/extract_mesh_vbo_vcol.cc
index 7a8f4a9a17e..7d159eb3df2 100644
--- a/source/blender/draw/intern/mesh_extractors/extract_mesh_vbo_vcol.cc
+++ b/source/blender/draw/intern/mesh_extractors/extract_mesh_vbo_vcol.cc
@@ -7,11 +7,63 @@
 
 #include "MEM_guardedalloc.h"
 
+#include "BKE_attribute.h"
 #include "BLI_string.h"
+#include "BLI_vector.hh"
 
 #include "draw_subdivision.h"
 #include "extract_mesh.h"
 
+struct VColRef {
+  const CustomDataLayer *layer;
+  AttributeDomain domain;
+};
+
+/** Get all vcol layers as AttributeRefs.
+ *
+ * \param vcol_layers: bitmask to filter vcol layers by, each bit
+ *                     corresponds to the integer position of the attribute
+ *                     within the global color attribute list.
+ */
+static blender::Vector<VColRef> get_vcol_refs(const CustomData *cd_vdata,
+                                              const CustomData *cd_ldata,
+                                              const uint vcol_layers)
+{
+  blender::Vector<VColRef> refs;
+  uint layeri = 0;
+
+  auto buildList = [&](const CustomData *cdata, AttributeDomain domain) {
+    for (int i = 0; i < cdata->totlayer; i++) {
+      const CustomDataLayer *layer = cdata->layers + i;
+
+      if (!(CD_TYPE_AS_MASK(layer->type) & CD_MASK_COLOR_ALL)) {
+        continue;
+      }
+
+      if (layer->flag & CD_FLAG_TEMPORARY) {
+        continue;
+      }
+
+      if (!(vcol_layers & (1UL << layeri))) {
+        layeri++;
+        continue;
+      }
+
+      VColRef ref;
+      ref.domain = domain;
+      ref.layer = layer;
+
+      refs.append(ref);
+      layeri++;
+    }
+  };
+
+  buildList(cd_vdata, ATTR_DOMAIN_POINT);
+  buildList(cd_ldata, ATTR_DOMAIN_CORNER);
+
+  return refs;
+}
+
 namespace blender::draw {
 
 /* ---------------------------------------------------------------------- */
@@ -21,34 +73,44 @@ namespace blender::draw {
 /* Initialize the common vertex format for vcol for coarse and subdivided meshes. */
 static void init_vcol_format(GPUVertFormat *format,
                              const MeshBatchCache *cache,
-                             CustomData *cd_ldata)
+                             CustomData *cd_vdata,
+                             CustomData *cd_ldata,
+                             CustomDataLayer *active,
+                             CustomDataLayer *render)
 {
   GPU_vertformat_deinterleave(format);
 
   const uint32_t vcol_layers = cache->cd_used.vcol;
 
-  for (int i = 0; i < MAX_MCOL; i++) {
-    if (vcol_layers & (1 << i)) {
-      char attr_name[32], attr_safe_name[GPU_MAX_SAFE_ATTR_NAME];
-      const char *layer_name = CustomData_get_layer_name(cd_ldata, CD_MLOOPCOL, i);
-      GPU_vertformat_safe_attr_name(layer_name, attr_safe_name, GPU_MAX_SAFE_ATTR_NAME);
+  blender::Vector<VColRef> refs = get_vcol_refs(cd_vdata, cd_ldata, vcol_layers);
 
-      BLI_snprintf(attr_name, sizeof(attr_name), "c%s", attr_safe_name);
-      GPU_vertformat_attr_add(format, attr_name, GPU_COMP_U16, 4, GPU_FETCH_INT_TO_FLOAT_UNIT);
+  for (const VColRef &ref : refs) {
+    char attr_name[32], attr_safe_name[GPU_MAX_SAFE_ATTR_NAME];
 
-      if (i == CustomData_get_render_layer(cd_ldata, CD_MLOOPCOL)) {
-        GPU_vertformat_alias_add(format, "c");
-      }
-      if (i == CustomData_get_active_layer(cd_ldata, CD_MLOOPCOL)) {
-        GPU_vertformat_alias_add(format, "ac");
-      }
+    GPU_vertformat_safe_attr_name(ref.layer->name, attr_safe_name, GPU_MAX_SAFE_ATTR_NAME);
 
-      /* Gather number of auto layers. */
-      /* We only do `vcols` that are not overridden by `uvs`. */
-      if (CustomData_get_named_layer_index(cd_ldata, CD_MLOOPUV, layer_name) == -1) {
-        BLI_snprintf(attr_name, sizeof(attr_name), "a%s", attr_safe_name);
-        GPU_vertformat_alias_add(format, attr_name);
-      }
+    /* VCol layer name. */
+    BLI_snprintf(attr_name, sizeof(attr_name), "c%s", attr_safe_name);
+    GPU_vertformat_attr_add(format, attr_name, GPU_COMP_U16, 4, GPU_FETCH_INT_TO_FLOAT_UNIT);
+
+    /* Active layer name. */
+    if (ref.layer == active) {
+      GPU_vertformat_alias_add(format, "ac");
+    }
+
+    /* Active render layer name. */
+    if (ref.layer == render) {
+      GPU_vertformat_alias_add(format, "c");
+    }
+
+    /* Gather number of auto layers. */
+    /* We only do `vcols` that are not overridden by `uvs`. */
+    bool bad = ref.domain == ATTR_DOMAIN_CORNER;
+    bad = bad && CustomData_get_named_layer_index(cd_ldata, CD_MLOOPUV, ref.layer->name) != -1;
+
+    if (!bad) {
+      BLI_snprintf(attr_name, sizeof(attr_name), "a%s", attr_safe_name);
+      GPU_vertformat_alias_add(format, attr_name);
     }
   }
 }
@@ -77,41 +139,101 @@ static void extract_vcol_init(const MeshRenderData *mr,
 {
   GPUVertBuf *vbo = static_cast<GPUVertBuf *>(buf);
   GPUVertFormat format = {0};
+
+  CustomData *cd_vdata = (mr->extract_type == MR_EXTRACT_BMESH) ? &mr->bm->vdata : &mr->me->vdata;
   CustomData *cd_ldata = (mr->extract_type == MR_EXTRACT_BMESH) ? &mr->bm->ldata : &mr->me->ldata;
+
+  Mesh me_query = {0};
+
+  BKE_id_attribute_copy_domains_temp(
+      ID_ME, cd_vdata, nullptr, cd_ldata, nullptr, nullptr, &me_query.id);
+
+  CustomDataLayer *active_color = BKE_id_attributes_active_color_get(&me_query.id);
+  CustomDataLayer *render_color = BKE_id_attributes_render_color_get(&me_query.id);
+
   const uint32_t vcol_layers = cache->cd_used.vcol;
-  init_vcol_format(&format, cache, cd_ldata);
+  init_vcol_format(&format, cache, cd_vdata, cd_ldata, active_color, render_color);
 
   GPU_vertbuf_init_with_format(vbo, &format);
   GPU_vertbuf_data_alloc(vbo, mr->loop_len);
 
   gpuMeshVcol *vcol_data = (gpuMeshVcol *)GPU_vertbuf_get_data(vbo);
 
-  for (int i = 0; i < MAX_MCOL; i++) {
-    if (vcol_layers & (1 << i)) {
-      if (mr->extract_type == MR_EXTRACT_BMESH) {
-        int cd_ofs = CustomData_get_n_offset(cd_ldata, CD_MLOOPCOL, i);
-        BMIter f_iter;
-        BMFace *efa;
-        BM_ITER_MESH (efa, &f_iter, mr->bm, BM_FACES_OF_MESH) {
-          BMLoop *l_iter, *l_first;
-          l_iter = l_first = BM_FACE_FIRST_LOOP(efa);
-          do {
-            const MLoopCol *mloopcol = (const MLoopCol *)BM_ELEM_CD_GET_VOID_P(l_iter, cd_ofs);
+  blender::Vector<VColRef> refs = get_vcol_refs(cd_vdata, cd_ldata, vcol_layers);
+
+  for (const VColRef &ref : refs) {
+    CustomData *cdata = ref.domain == ATTR_DOMAIN_POINT ? cd_vdata : cd_ldata;
+
+    if (mr->extract_type == MR_EXTRACT_BMESH) {
+      int cd_ofs = ref.layer->offset;
+
+      if (cd_ofs == -1) {
+        vcol_data += ref.domain == ATTR_DOMAIN_POINT ? mr->bm->totvert : mr->bm->totloop;
+        continue;
+      }
+
+      BMIter iter;
+      const bool is_byte = ref.layer->type == CD_MLOOPCOL;
+      const bool is_point = ref.domain == ATTR_DOMAIN_POINT;
+
+      BMFace *f;
+      BM_ITER_MESH (f, &iter, mr->bm, BM_FACES_OF_MESH) {
+        BMLoop *l_iter = f->l_first;
+        do {
+          BMElem *elem = is_point ? reinterpret_cast<BMElem *>(l_iter->v) :
+                                    reinterpret_cast<BMElem *>(l_iter);
+          if (is_byte) {
+            const MLoopCol *mloopcol = (const MLoopCol *)BM_ELEM_CD_GET_VOID_P(elem, cd_ofs);
             vcol_data->r = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mloopcol->r]);
             vcol_data->g = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mloopcol->g]);
             vcol_data->b = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mloopcol->b]);
             vcol_data->a = unit_float_to_ushort_clamp(mloopcol->a * (1.0f / 255.0f));
             vcol_data++;
-          } while ((l_iter = l_iter->next) != l_first);
-        }
+          }
+          else {
+            const MPropCol *mpcol = (const MPropCol *)BM_ELEM_CD_GET_VOID_P(elem, cd_ofs);
+            vcol_data->r = unit_float_to_ushort_clamp(mpcol->color[0]);
+            vcol_data->g = unit_float_to_ushort_clamp(mpcol->color[1]);
+            vcol_data->b = unit_float_to_ushort_clamp(mpcol->color[2]);
+            vcol_data->a = unit_float_to_ushort_clamp(mpcol->color[3]);
+            vcol_data++;
+          }
+        } while ((l_iter = l_iter->next) != f->l_first);
+      }
+    }
+    else {
+      int totloop = mr->loop_len;
+      int idx = CustomData_get_named_layer_index(cdata, ref.layer->type, ref.layer->name);
+
+      MLoopCol *mcol = nullptr;
+      MPropCol *pcol = nullptr;
+      const MLoop *mloop = mr->mloop;
+
+      if (ref.layer->type == CD_PROP_COLOR) {
+        pcol = static_cast<MPropCol *>(cdata->layers[idx].data);
       }
       else {
-        const MLoopCol *mloopcol = (MLoopCol *)CustomData_get_layer_n(cd_ldata, CD_MLOOPCOL, i);
-        for (int ml_index = 0; ml_index < mr->loop_len; ml_index++, mloopcol++, vcol_data++) {
-          vcol_data->r = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mloopcol->r]);
-          vcol_data->g = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mloopcol->g]);
-          vcol_data->b = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mloopcol->b]);
-          vcol_data->a = unit_float_to_ushort_clamp(mloopcol->a * (1.0f / 255.0f));
+        mcol = static_cast<MLoopCol *>(cdata->layers[idx].data);
+      }
+
+      const bool is_corner = ref.domain == ATTR_DOMAIN_CORNER;
+
+      for (int i = 0; i < totloop; i++, mloop++) {
+        const int v_i = is_corner ? i : mloop->v;
+
+        if (mcol) {
+          vcol_data->r = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mcol[v_i].r]);
+          vcol_data->g = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mcol[v_i].g]);
+          vcol_data->b = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mcol[v_i].b]);
+          vcol_data->a = unit_float_to_ushort_clamp(mcol[v_i].a * (1.0f / 255.0f));
+          vcol_data++;
+        }
+        else if (pcol) {
+          vcol_data->r = unit_float_to_ushort_clamp(pcol[v_i].color[0]);
+          vcol_data->g = unit_float_to_ushort_clamp(pcol[v_i].color[1]);
+          vcol_data->b = unit_float_to_ushort_clamp(pcol[v_i].color[2]);
+          vcol_data->a = unit_float_to_ushort_clamp(pcol[v_i].color[3]);
+          vcol_data++;
         }
       }
     }
@@ -119,7 +241,7 @@ static void extract_vcol_init(const MeshRenderData *mr,
 }
 
 static void extract_vcol_init_subdiv(const DRWSubdivCache *subdiv_cache,
-                                     const MeshRenderData *UNUSED(mr),
+                                     const MeshRenderData *mr,
                                      struct MeshBatchCache *cache,
                                      void *buffer,
                                      void *UNUSED(data))
@@ -127,8 +249,23 @@ static void extract_vcol_init_subdiv(const DRWSubdivCache *subdiv_cache,
   GPUVertBuf *dst_buffer = static_cast<GPUVertBuf *>(buffer);
   Mesh *coarse_mesh = subdiv_cache->mesh;
 
+  bool extract_bmesh = mr->extract_type == MR_EXTRACT_BMESH;
+
+  const CustomData *cd_vdata = extract_bmesh ? &coarse_mesh->edit_mesh->bm->vdata :
+                                               &coarse_mesh->vdata;
+  const CustomData *cd_ldata = extract_bmesh ? &coarse_mesh->edit_mesh->bm->ldata :
+                                               &coarse_mesh->ldata;
+
+  Mesh me_query = *coarse_mesh;
+  BKE_id_attribute_copy_domains_temp(
+      ID_ME, cd_vdata, nullptr, cd_ldata, nullptr, nullptr, &me_query.id);
+
+  CustomDataLayer *active_color = BKE_id_attributes_active_color_get(&me_query.id);
+  CustomDataLayer *render_color = BKE_id_attributes_render_color_get(&me_query.id);
+
   GPUVertFormat format = {0};
-  init_vcol_format(&format, cache, &coarse_mesh->ldata);
+  init_vcol_format(
+      &format, cache, &coarse_mesh->vdata, &coarse_mesh->ldata, active_color, render_color);
 
   GPU_vertbuf_init_build_on_device(dst_buffer, &format, subdiv_cache->num_subdiv_loops);
 
@@ -140,32 +277,97 @@ static void extract_vcol_init_subdiv(const DRWSubdivCache *subdiv_cache,
 
   gpuMeshVcol *mesh_vcol = (gpuMeshVcol *)GPU_vertbuf_get_data(src_data);
 
-  const CustomData *cd_ldata = &coarse_mesh->ldata;
-
   const uint vcol_layers = cache->cd_used.vcol;
 
+  blender::Vector<VColRef> refs = get_vcol_refs(cd_vdata, cd_ldata, vcol_layers);
+
+  gpuMeshVcol *vcol = mesh_vcol;
+
   /* Index of the vertex color layer in the compact buffer. Used vertex color layers are stored in
    * a single buffer. */
   int pack_layer_index = 0;
-  for (int i = 0; i < MAX_MTFACE; i++) {
-    if (vcol_layers & (1 << i)) {
-      /* Include stride in offset, we use a stride of 2 since colors are packed into 2 uints. */
-      const int dst_offset = (int)subdiv_cache->num_subdiv_loops * 2 * pack_layer_index++;
-      const MLoopCol *mloopcol = (MLoopCol *)CustomData_get_layer_n(cd_ldata, CD_MLOOPCOL, i);
-
-      gpuMeshVcol *vcol = mesh_vcol;
-
-      for (int ml_index = 0; ml_index < coarse_mesh->totloop; ml_index++, vcol++, mloopcol++) {
-        vcol->r = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mloopcol->r]);
-        vcol->g = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mloopcol->g]);
-        vcol->b = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mloopcol->b]);
-        vcol->a = unit_float_to_ushort_clamp(mloopcol->a * (1.0f / 255.0f));
-      }
+  for (const VColRef &ref : refs) {
+    /* Include stride in offset, we use a stride of 2 since colors are packed into 2 uints. */
+    const int dst_offset = (int)subdiv_cache->num_subdiv_loops * 2 * pack_layer_index++;
+
+    const CustomData *cdata = ref.domain == ATTR_DOMAIN_POINT ? cd_vdata : cd_ldata;
+    const MLoop *ml = coarse_mesh->mloop;
+
+    int layer_i = CustomData_get_named_layer_index(cdata, ref.layer->type, ref.layer->name);
 
-      /* Ensure data is uploaded properly. */
-      GPU_vertbuf_tag_dirty(src_data);
-      draw_subdiv_interp_custom_data(subdiv_cache, src_data, dst_buffer, 4, dst_offset, true);
+    if (layer_i == -1) {
+      printf("%s: missing color layer %s\n", __func__, ref.layer->name);
+      vcol += coarse_mesh->totloop;
+      continue;
+    }
+
+    MLoopCol *mcol = nullptr;
+    MPropCol *pcol = nullptr;
+
+    if (ref.layer->type == CD_PROP_COLOR) {
+      pcol = static_cast<MPropCol *>(cdata->layers[layer_i].data);
+    }
+    else {
+      mcol = static_cast<MLoopCol *>(cdata->layers[layer_i].data);
     }
+
+    const bool is_vert = ref.domain == ATTR_DOMAIN_POINT;
+
+    if (extract_bmesh) {
+      BMesh *bm = coarse_mesh->edit_mesh->bm;
+      BMIter iter;
+      BMFace *f;
+      int cd_ofs = cdata->layers[layer_i].offset;
+      const bool is_byte = ref.layer->type == CD_MLOOPCOL;
+
+      BM_ITER_MESH (f, &iter, bm, BM_FACES_OF_MESH) {
+        BMLoop *l_iter = f->l_first;
+
+        do {
+          BMElem *elem = is_vert ? reinterpret_cast<BMElem *>(l_iter->v) :
+                                   reinterpret_cast<BMElem *>(l_iter);
+
+          if (is_byte) {
+            MLoopCol *mcol2 = static_cast<MLoopCol *>(BM_ELEM_CD_GET_VOID_P(elem, cd_ofs));
+
+            vcol->r = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mcol2->r]);
+            vcol->g = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mcol2->g]);
+            vcol->b = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mcol2->b]);
+            vcol->a = unit_float_to_ushort_clamp(mcol2->a * (1.0f / 255.0f));
+          }
+          else {
+            MPropCol *pcol2 = static_cast<MPropCol *>(BM_ELEM_CD_GET_VOID_P(elem, cd_ofs));
+
+            vcol->r = unit_float_to_ushort_clamp(pcol2->color[0]);
+            vcol->g = unit_float_to_ushort_clamp(pcol2->color[1]);
+            vcol->b = unit_float_to_ushort_clamp(pcol2->color[2]);
+            vcol->a = unit_float_to_ushort_clamp(pcol2->color[3]);
+          }
+        } while ((l_iter = l_iter->next) != f->l_first);
+      }
+    }
+    else {
+      for (int ml_index = 0; ml_index < coarse_mesh->totloop; ml_index++, vcol++, ml++) {
+        int idx = is_vert ? ml->v : ml_index;
+
+        if (mcol) {
+          vcol->r = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mcol[idx].r]);
+          vcol->g = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mcol[idx].g]);
+          vcol->b = unit_float_to_ushort_clamp(BLI_color_from_srgb_table[mcol[idx].b]);
+          vcol->a = unit_float_to_ushort_clamp(mcol[idx].a * (1.0f / 255.0f));
+        }
+        else if (pcol) {
+          vcol->r = unit_float_to_ushort_clamp(pcol[idx].color[0]);
+          vcol->g = unit_float_to_ushort_clamp(pcol[idx].color[1]);
+          vcol->b = unit_float_to_ushort_clamp(pcol[idx].color[2]);
+          vcol->a = unit_float_to_ushort_clamp(pcol[idx].color[3]);
+        }
+      }
+    }
+
+    /* Ensure data is uploaded properly. */
+    GPU_vertbuf_tag_dirty(src_data);
+    draw_subdiv_interp_custom_data(subdiv_cache, src_data, dst_buffer, 4, dst_offset, true);
   }
 
   GPU_vertbuf_discard(src_data);