DRW: Implement Hair Weight drawing

Fixes T57931 Particle weight edit mode is not supported.

There is a bug that prevent refresh of the toolsettings on which is based
the weight / non-weight display selection (see T58086).

source/blender/draw/intern/draw_cache.c

@@ -3465,9 +3465,10 @@ GPUBatch *DRW_cache_particles_get_dots(Object *object, ParticleSystem *psys)
 GPUBatch *DRW_cache_particles_get_edit_strands(
         Object *object,
         ParticleSystem *psys,
-        struct PTCacheEdit *edit)
+        struct PTCacheEdit *edit,
+        bool use_weight)
 {
-	return DRW_particles_batch_cache_get_edit_strands(object, psys, edit);
+	return DRW_particles_batch_cache_get_edit_strands(object, psys, edit, use_weight);
 }
 
 GPUBatch *DRW_cache_particles_get_edit_inner_points(

source/blender/draw/intern/draw_cache.h

@@ -200,7 +200,7 @@ struct GPUBatch *DRW_cache_particles_get_hair(
 struct GPUBatch *DRW_cache_particles_get_dots(
         struct Object *object, struct ParticleSystem *psys);
 struct GPUBatch *DRW_cache_particles_get_edit_strands(
-        struct Object *object, struct ParticleSystem *psys, struct PTCacheEdit *edit);
+        struct Object *object, struct ParticleSystem *psys, struct PTCacheEdit *edit, bool use_weight);
 struct GPUBatch *DRW_cache_particles_get_edit_inner_points(
         struct Object *object, struct ParticleSystem *psys, struct PTCacheEdit *edit);
 struct GPUBatch *DRW_cache_particles_get_edit_tip_points(

source/blender/draw/intern/draw_cache_impl.h

@@ -212,7 +212,7 @@ struct GPUBatch *DRW_particles_batch_cache_get_hair(
 struct GPUBatch *DRW_particles_batch_cache_get_dots(
         struct Object *object, struct ParticleSystem *psys);
 struct GPUBatch *DRW_particles_batch_cache_get_edit_strands(
-        struct Object *object, struct ParticleSystem *psys, struct PTCacheEdit *edit);
+        struct Object *object, struct ParticleSystem *psys, struct PTCacheEdit *edit, bool use_weight);
 struct GPUBatch *DRW_particles_batch_cache_get_edit_inner_points(
         struct Object *object, struct ParticleSystem *psys, struct PTCacheEdit *edit);
 struct GPUBatch *DRW_particles_batch_cache_get_edit_tip_points(

source/blender/draw/intern/draw_cache_impl_particles.c

@@ -93,6 +93,7 @@ typedef struct ParticleBatchCache {
 
 	/* Settings to determine if cache is invalid. */
 	bool is_dirty;
+	bool edit_is_weight;
 } ParticleBatchCache;
 
 /* GPUBatch cache management. */
@@ -692,8 +693,27 @@ static float particle_key_select_ratio(const PTCacheEdit *edit, int strand, floa
 	}
 }
 
+static float particle_key_weight(const ParticleData *particle, int strand, float t)
+{
+	const ParticleData *part = particle + strand;
+	const HairKey *hkeys = part->hair;
+	float edit_key_seg_t = 1.0f / (part->totkey - 1);
+	if (t == 1.0) {
+		return hkeys[part->totkey - 1].weight;
+	}
+	else {
+		float interp = t / edit_key_seg_t;
+		int index = (int)interp;
+		interp -= floorf(interp); /* Time between 2 edit key */
+		float s1 = hkeys[index].weight;
+		float s2 = hkeys[index+1].weight;
+		return s1 + interp * (s2 - s1);
+	}
+}
+
 static int particle_batch_cache_fill_segments_edit(
-        const PTCacheEdit *edit,
+        const PTCacheEdit *edit, /* NULL for weight data */
+        const ParticleData *particle, /* NULL for select data */
         ParticleCacheKey **path_cache,
         const int start_index,
         const int num_path_keys,
@@ -710,8 +730,15 @@ static int particle_batch_cache_fill_segments_edit(
 			EditStrandData *seg_data = (EditStrandData *)GPU_vertbuf_raw_step(attr_step);
 			copy_v3_v3(seg_data->pos, path[j].co);
 			float strand_t = (float)(j) / path->segments;
-			float selected = particle_key_select_ratio(edit, i, strand_t);
-			seg_data->color = (uchar)(0xFF * selected);
+			if (particle) {
+				float weight = particle_key_weight(particle, i, strand_t);
+				/* NaN or unclamped become 0xFF */
+				seg_data->color = (uchar)((weight <= 1.0f) ? 0xFE * weight : 0xFF);
+			}
+			else {
+				float selected = particle_key_select_ratio(edit, i, strand_t);
+				seg_data->color = (uchar)(0xFF * selected);
+			}
 			GPU_indexbuf_add_generic_vert(elb, curr_point);
 			curr_point++;
 		}
@@ -1436,14 +1463,17 @@ GPUBatch *DRW_particles_batch_cache_get_dots(Object *object, ParticleSystem *psy
 
 static void particle_batch_cache_ensure_edit_pos_and_seg(
         PTCacheEdit *edit,
-        ParticleSystem *UNUSED(psys),
+        ParticleSystem *psys,
         ModifierData *UNUSED(md),
-        ParticleHairCache *hair_cache)
+        ParticleHairCache *hair_cache,
+        bool use_weight)
 {
 	if (hair_cache->pos != NULL && hair_cache->indices != NULL) {
 		return;
 	}
 
+	ParticleData *particle = (use_weight) ? psys->particles : NULL;
+
 	GPU_VERTBUF_DISCARD_SAFE(hair_cache->pos);
 	GPU_INDEXBUF_DISCARD_SAFE(hair_cache->indices);
 
@@ -1464,7 +1494,7 @@ static void particle_batch_cache_ensure_edit_pos_and_seg(
 
 	if (edit != NULL && edit->pathcache != NULL) {
 		particle_batch_cache_fill_segments_edit(
-		        edit, edit->pathcache,
+		        edit, particle, edit->pathcache,
 		        0, edit->totcached,
 		        &elb, &data_step);
 	}
@@ -1477,19 +1507,25 @@ static void particle_batch_cache_ensure_edit_pos_and_seg(
 GPUBatch *DRW_particles_batch_cache_get_edit_strands(
         Object *object,
         ParticleSystem *psys,
-        PTCacheEdit *edit)
+        PTCacheEdit *edit,
+        bool use_weight)
 {
 	ParticleBatchCache *cache = particle_batch_cache_get(psys);
+	if (cache->edit_is_weight != use_weight) {
+		GPU_VERTBUF_DISCARD_SAFE(cache->edit_hair.pos);
+		GPU_BATCH_DISCARD_SAFE(cache->edit_hair.hairs);
+	}
 	if (cache->edit_hair.hairs != NULL) {
 		return cache->edit_hair.hairs;
 	}
 	drw_particle_update_ptcache_edit(object, psys, edit);
 	ensure_seg_pt_count(edit, psys, &cache->edit_hair);
-	particle_batch_cache_ensure_edit_pos_and_seg(edit, psys, NULL, &cache->edit_hair);
+	particle_batch_cache_ensure_edit_pos_and_seg(edit, psys, NULL, &cache->edit_hair, use_weight);
 	cache->edit_hair.hairs = GPU_batch_create(
 	        GPU_PRIM_LINE_STRIP,
 	        cache->edit_hair.pos,
 	        cache->edit_hair.indices);
+	cache->edit_is_weight = use_weight;
 	return cache->edit_hair.hairs;
 }
 

source/blender/draw/modes/particle_mode.c

@@ -83,6 +83,7 @@ typedef struct PARTICLE_Data {
 
 static struct {
 	struct GPUShader *strands_shader;
+	struct GPUShader *strands_weight_shader;
 	struct GPUShader *points_shader;
 } e_data = {NULL}; /* Engine data */
 
@@ -104,6 +105,13 @@ static void particle_engine_init(void *UNUSED(vedata))
 		        datatoc_common_globals_lib_glsl,
 		        "");
 
+		e_data.strands_weight_shader = DRW_shader_create_with_lib(
+		        datatoc_particle_strand_vert_glsl,
+		        NULL,
+		        datatoc_particle_strand_frag_glsl,
+		        datatoc_common_globals_lib_glsl,
+		        "#define USE_WEIGHT");
+
 		e_data.points_shader = DRW_shader_create_with_lib(
 		        datatoc_particle_strand_vert_glsl,
 		        NULL,
@@ -117,6 +125,9 @@ static void particle_cache_init(void *vedata)
 {
 	PARTICLE_PassList *psl = ((PARTICLE_Data *)vedata)->psl;
 	PARTICLE_StorageList *stl = ((PARTICLE_Data *)vedata)->stl;
+	const DRWContextState *draw_ctx = DRW_context_state_get();
+	ParticleEditSettings *pset = PE_settings(draw_ctx->scene);
+	const bool use_weight = (pset->brushtype == PE_BRUSH_WEIGHT);
 
 	if (!stl->g_data) {
 		/* Alloc transient pointers */
@@ -131,12 +142,10 @@ static void particle_cache_init(void *vedata)
 	                                       DRW_STATE_WIRE |
 	                                       DRW_STATE_POINT));
 
-	stl->g_data->strands_group = DRW_shgroup_create(
-	        e_data.strands_shader, psl->psys_edit_pass);
-	stl->g_data->inner_points_group = DRW_shgroup_create(
-	        e_data.points_shader, psl->psys_edit_pass);
-	stl->g_data->tip_points_group = DRW_shgroup_create(
-	        e_data.points_shader, psl->psys_edit_pass);
+	GPUShader *strand_shader = (use_weight) ? e_data.strands_weight_shader : e_data.strands_shader;
+	stl->g_data->strands_group = DRW_shgroup_create(strand_shader, psl->psys_edit_pass);
+	stl->g_data->inner_points_group = DRW_shgroup_create(e_data.points_shader, psl->psys_edit_pass);
+	stl->g_data->tip_points_group = DRW_shgroup_create(e_data.points_shader, psl->psys_edit_pass);
 
 	DRW_shgroup_uniform_block(stl->g_data->strands_group, "globalsBlock", globals_ubo);
 	DRW_shgroup_uniform_block(stl->g_data->inner_points_group, "globalsBlock", globals_ubo);
@@ -151,9 +160,10 @@ static void particle_edit_cache_populate(void *vedata,
 	PARTICLE_StorageList *stl = ((PARTICLE_Data *)vedata)->stl;
 	const DRWContextState *draw_ctx = DRW_context_state_get();
 	ParticleEditSettings *pset = PE_settings(draw_ctx->scene);
+	const bool use_weight = (pset->brushtype == PE_BRUSH_WEIGHT);
 	{
 		struct GPUBatch *strands =
-		        DRW_cache_particles_get_edit_strands(object, psys, edit);
+		        DRW_cache_particles_get_edit_strands(object, psys, edit, use_weight);
 		DRW_shgroup_call_add(stl->g_data->strands_group, strands, NULL);
 	}
 	if (pset->selectmode == SCE_SELECT_POINT) {
@@ -229,6 +239,7 @@ static void particle_draw_scene(void *vedata)
 static void particle_engine_free(void)
 {
 	DRW_SHADER_FREE_SAFE(e_data.strands_shader);
+	DRW_SHADER_FREE_SAFE(e_data.strands_weight_shader);
 	DRW_SHADER_FREE_SAFE(e_data.points_shader);
 }
 

source/blender/draw/modes/shaders/particle_strand_vert.glsl

@@ -9,11 +9,53 @@ out vec4 finalColor;
 out vec2 radii;
 #endif
 
+vec3 weight_to_rgb(float weight)
+{
+	vec3 r_rgb;
+	float blend = ((weight / 2.0) + 0.5);
+
+	if (weight <= 0.25) {    /* blue->cyan */
+		r_rgb[0] = 0.0;
+		r_rgb[1] = blend * weight * 4.0;
+		r_rgb[2] = blend;
+	}
+	else if (weight <= 0.50) {  /* cyan->green */
+		r_rgb[0] = 0.0;
+		r_rgb[1] = blend;
+		r_rgb[2] = blend * (1.0 - ((weight - 0.25) * 4.0));
+	}
+	else if (weight <= 0.75) {  /* green->yellow */
+		r_rgb[0] = blend * ((weight - 0.50) * 4.0);
+		r_rgb[1] = blend;
+		r_rgb[2] = 0.0;
+	}
+	else if (weight <= 1.0) {  /* yellow->red */
+		r_rgb[0] = blend;
+		r_rgb[1] = blend * (1.0 - ((weight - 0.75) * 4.0));
+		r_rgb[2] = 0.0;
+	}
+	else {
+		/* exceptional value, unclamped or nan,
+		 * avoid uninitialized memory use */
+		r_rgb[0] = 1.0;
+		r_rgb[1] = 0.0;
+		r_rgb[2] = 1.0;
+	}
+
+	return r_rgb;
+}
+
+#define DECOMPRESS_RANGE 1.0039
+
 void main()
 {
 	gl_Position = ModelViewProjectionMatrix * vec4(pos, 1.0);
 
+#ifdef USE_WEIGHT
+	finalColor = vec4(weight_to_rgb(color * DECOMPRESS_RANGE), 1.0);
+#else
 	finalColor = mix(colorWire, colorEdgeSelect, color);
+#endif
 
 #ifdef USE_POINTS
 	gl_PointSize = sizeVertex;