Weight Paint Tools: Add "Subset" option to "Normalize All"

This option is needed when vertex groups are used for both armature
deformation and the other purpose such as influence of mesh modifier.

Thanks to Campbell for code review!

3 files changed, 97 insertions, 50 deletions

diff --git a/source/blender/blenkernel/BKE_deform.h b/source/blender/blenkernel/BKE_deform.h
index a43a243106d..435cad17e57 100644
--- a/source/blender/blenkernel/BKE_deform.h
+++ b/source/blender/blenkernel/BKE_deform.h
@@ -78,8 +78,12 @@ void defvert_flip_merged(struct MDeformVert *dvert, const int *flip_map, const i
 void defvert_normalize(struct MDeformVert *dvert);
 void defvert_normalize_subset(struct MDeformVert *dvert,
                               const bool *vgroup_subset, const int vgroup_tot);
-void defvert_normalize_lock_single(struct MDeformVert *dvert, const int def_nr_lock);
-void defvert_normalize_lock_map(struct MDeformVert *dvert, const bool *lock_flags, const int defbase_tot);
+void defvert_normalize_lock_single(struct MDeformVert *dvert,
+                                   const bool *vgroup_subset, const int vgroup_tot,
+                                   const int def_nr_lock);
+void defvert_normalize_lock_map(struct MDeformVert *dvert,
+                                const bool *vgroup_subset, const int vgroup_tot,
+                                const bool *lock_flags, const int defbase_tot);
 
 /* utility function, note that MAX_VGROUP_NAME chars is the maximum string length since its only
  * used with defgroups currently */
diff --git a/source/blender/blenkernel/intern/deform.c b/source/blender/blenkernel/intern/deform.c
index 882085aa5db..a183872552d 100644
--- a/source/blender/blenkernel/intern/deform.c
+++ b/source/blender/blenkernel/intern/deform.c
@@ -217,24 +217,35 @@ void defvert_remap(MDeformVert *dvert, int *map, const int map_len)
 void defvert_normalize_subset(MDeformVert *dvert,
                               const bool *vgroup_subset, const int vgroup_tot)
 {
-	MDeformWeight *dw;
-	unsigned int i;
-	float tot_weight = 0.0f;
-
-	for (i = dvert->totweight, dw = dvert->dw; i != 0; i--, dw++) {
+	if (dvert->totweight == 0) {
+		/* nothing */
+	}
+	else if (dvert->totweight == 1) {
+		MDeformWeight *dw = dvert->dw;
 		if ((dw->def_nr < vgroup_tot) && vgroup_subset[dw->def_nr]) {
-			tot_weight += dw->weight;
+			dw->weight = 1.0f;
 		}
 	}
+	else {
+		MDeformWeight *dw;
+		unsigned int i;
+		float tot_weight = 0.0f;
 
-	if (tot_weight > 0.0f) {
-		float scalar = 1.0f / tot_weight;
 		for (i = dvert->totweight, dw = dvert->dw; i != 0; i--, dw++) {
 			if ((dw->def_nr < vgroup_tot) && vgroup_subset[dw->def_nr]) {
-				dw->weight *= scalar;
-				
-				/* in case of division errors with very low weights */
-				CLAMP(dw->weight, 0.0f, 1.0f);
+				tot_weight += dw->weight;
+			}
+		}
+
+		if (tot_weight > 0.0f) {
+			float scalar = 1.0f / tot_weight;
+			for (i = dvert->totweight, dw = dvert->dw; i != 0; i--, dw++) {
+				if ((dw->def_nr < vgroup_tot) && vgroup_subset[dw->def_nr]) {
+					dw->weight *= scalar;
+
+					/* in case of division errors with very low weights */
+					CLAMP(dw->weight, 0.0f, 1.0f);
+				}
 			}
 		}
 	}
@@ -242,7 +253,7 @@ void defvert_normalize_subset(MDeformVert *dvert,
 
 void defvert_normalize(MDeformVert *dvert)
 {
-	if (dvert->totweight <= 0) {
+	if (dvert->totweight == 0) {
 		/* nothing */
 	}
 	else if (dvert->totweight == 1) {
@@ -269,14 +280,20 @@ void defvert_normalize(MDeformVert *dvert)
 	}
 }
 
-void defvert_normalize_lock_single(MDeformVert *dvert, const int def_nr_lock)
+/* Same as defvert_normalize() if the locked vgroup is not a member of the subset */
+void defvert_normalize_lock_single(MDeformVert *dvert,
+                                   const bool *vgroup_subset, const int vgroup_tot,
+                                   const int def_nr_lock)
 {
-	if (dvert->totweight <= 0) {
+	if (dvert->totweight == 0) {
 		/* nothing */
 	}
 	else if (dvert->totweight == 1) {
-		if (def_nr_lock != 0) {
-			dvert->dw[0].weight = 1.0f;
+		MDeformWeight *dw = dvert->dw;
+		if ((dw->def_nr < vgroup_tot) && vgroup_subset[dw->def_nr]) {
+			if (def_nr_lock != 0) {
+				dw->weight = 1.0f;
+			}
 		}
 	}
 	else {
@@ -287,13 +304,15 @@ void defvert_normalize_lock_single(MDeformVert *dvert, const int def_nr_lock)
 		float lock_iweight = 1.0f;
 
 		for (i = dvert->totweight, dw = dvert->dw; i != 0; i--, dw++) {
-			if (dw->def_nr != def_nr_lock) {
-				tot_weight += dw->weight;
-			}
-			else {
-				dw_lock = dw;
-				lock_iweight = (1.0f - dw_lock->weight);
-				CLAMP(lock_iweight, 0.0f, 1.0f);
+			if ((dw->def_nr < vgroup_tot) && vgroup_subset[dw->def_nr]) {
+				if (dw->def_nr != def_nr_lock) {
+					tot_weight += dw->weight;
+				}
+				else {
+					dw_lock = dw;
+					lock_iweight = (1.0f - dw_lock->weight);
+					CLAMP(lock_iweight, 0.0f, 1.0f);
+				}
 			}
 		}
 
@@ -302,25 +321,33 @@ void defvert_normalize_lock_single(MDeformVert *dvert, const int def_nr_lock)
 
 			float scalar = (1.0f / tot_weight) * lock_iweight;
 			for (i = dvert->totweight, dw = dvert->dw; i != 0; i--, dw++) {
-				if (dw != dw_lock) {
-					dw->weight *= scalar;
+				if ((dw->def_nr < vgroup_tot) && vgroup_subset[dw->def_nr]) {
+					if (dw != dw_lock) {
+						dw->weight *= scalar;
 
-					/* in case of division errors with very low weights */
-					CLAMP(dw->weight, 0.0f, 1.0f);
+						/* in case of division errors with very low weights */
+						CLAMP(dw->weight, 0.0f, 1.0f);
+					}
 				}
 			}
 		}
 	}
 }
 
-void defvert_normalize_lock_map(MDeformVert *dvert, const bool *lock_flags, const int defbase_tot)
+/* Same as defvert_normalize() if no locked vgroup is a member of the subset */
+void defvert_normalize_lock_map(MDeformVert *dvert,
+				const bool *vgroup_subset, const int vgroup_tot,
+				const bool *lock_flags, const int defbase_tot)
 {
-	if (dvert->totweight <= 0) {
+	if (dvert->totweight == 0) {
 		/* nothing */
 	}
 	else if (dvert->totweight == 1) {
-		if (LIKELY(defbase_tot >= 1) && lock_flags[0]) {
-			dvert->dw[0].weight = 1.0f;
+		MDeformWeight *dw = dvert->dw;
+		if ((dw->def_nr < vgroup_tot) && vgroup_subset[dw->def_nr]) {
+			if (LIKELY(defbase_tot >= 1) && lock_flags[0]) {
+				dw->weight = 1.0f;
+			}
 		}
 	}
 	else {
@@ -330,12 +357,14 @@ void defvert_normalize_lock_map(MDeformVert *dvert, const bool *lock_flags, cons
 		float lock_iweight = 0.0f;
 
 		for (i = dvert->totweight, dw = dvert->dw; i != 0; i--, dw++) {
-			if ((dw->def_nr < defbase_tot) && (lock_flags[dw->def_nr] == FALSE)) {
-				tot_weight += dw->weight;
-			}
-			else {
-				/* invert after */
-				lock_iweight += dw->weight;
+			if ((dw->def_nr < vgroup_tot) && vgroup_subset[dw->def_nr]) {
+				if ((dw->def_nr < defbase_tot) && (lock_flags[dw->def_nr] == FALSE)) {
+					tot_weight += dw->weight;
+				}
+				else {
+					/* invert after */
+					lock_iweight += dw->weight;
+				}
 			}
 		}
 
@@ -346,11 +375,13 @@ void defvert_normalize_lock_map(MDeformVert *dvert, const bool *lock_flags, cons
 
 			float scalar = (1.0f / tot_weight) * lock_iweight;
 			for (i = dvert->totweight, dw = dvert->dw; i != 0; i--, dw++) {
-				if ((dw->def_nr < defbase_tot) && (lock_flags[dw->def_nr] == FALSE)) {
-					dw->weight *= scalar;
+				if ((dw->def_nr < vgroup_tot) && vgroup_subset[dw->def_nr]) {
+					if ((dw->def_nr < defbase_tot) && (lock_flags[dw->def_nr] == FALSE)) {
+						dw->weight *= scalar;
 
-					/* in case of division errors with very low weights */
-					CLAMP(dw->weight, 0.0f, 1.0f);
+						/* in case of division errors with very low weights */
+						CLAMP(dw->weight, 0.0f, 1.0f);
+					}
 				}
 			}
 		}
diff --git a/source/blender/editors/object/object_vgroup.c b/source/blender/editors/object/object_vgroup.c
index b0d884568ca..d6c365e9247 100644
--- a/source/blender/editors/object/object_vgroup.c
+++ b/source/blender/editors/object/object_vgroup.c
@@ -2001,7 +2001,11 @@ static void vgroup_levels_subset(Object *ob, const bool *vgroup_validmap, const
 	}
 }
 
-static void vgroup_normalize_all(Object *ob, const bool lock_active)
+static void vgroup_normalize_all(Object *ob,
+                                 const bool *vgroup_validmap,
+                                 const int vgroup_tot,
+                                 const int subset_count,
+                                 const bool lock_active)
 {
 	MDeformVert *dv, **dvert_array = NULL;
 	int i, dvert_tot = 0;
@@ -2009,7 +2013,7 @@ static void vgroup_normalize_all(Object *ob, const bool lock_active)
 
 	const int use_vert_sel = vertex_group_use_vert_sel(ob);
 
-	if (lock_active && !BLI_findlink(&ob->defbase, def_nr)) {
+	if ((lock_active && !BLI_findlink(&ob->defbase, def_nr)) || subset_count == 0) {
 		return;
 	}
 
@@ -2030,13 +2034,15 @@ static void vgroup_normalize_all(Object *ob, const bool lock_active)
 			/* in case its not selected */
 			if ((dv = dvert_array[i])) {
 				if (lock_flags) {
-					defvert_normalize_lock_map(dv, lock_flags, defbase_tot);
+					defvert_normalize_lock_map(dv, vgroup_validmap, vgroup_tot,
+					                           lock_flags, defbase_tot);
 				}
 				else if (lock_active) {
-					defvert_normalize_lock_single(dv, def_nr);
+					defvert_normalize_lock_single(dv, vgroup_validmap, vgroup_tot,
+					                              def_nr);
 				}
 				else {
-					defvert_normalize(dv);
+					defvert_normalize_subset(dv, vgroup_validmap, vgroup_tot);
 				}
 			}
 		}
@@ -3474,8 +3480,13 @@ static int vertex_group_normalize_all_exec(bContext *C, wmOperator *op)
 {
 	Object *ob = ED_object_context(C);
 	bool lock_active = RNA_boolean_get(op->ptr, "lock_active");
+	eVGroupSelect subset_type  = RNA_enum_get(op->ptr, "group_select_mode");
+
+	int subset_count, vgroup_tot;
 
-	vgroup_normalize_all(ob, lock_active);
+	const bool *vgroup_validmap = ED_vgroup_subset_from_select_type(ob, subset_type, &vgroup_tot, &subset_count);
+	vgroup_normalize_all(ob, vgroup_validmap, vgroup_tot, subset_count, lock_active);
+	MEM_freeN((void *)vgroup_validmap);
 
 	DAG_id_tag_update(&ob->id, OB_RECALC_DATA);
 	WM_event_add_notifier(C, NC_OBJECT | ND_DRAW, ob);
@@ -3499,6 +3510,7 @@ void OBJECT_OT_vertex_group_normalize_all(wmOperatorType *ot)
 	/* flags */
 	ot->flag = OPTYPE_REGISTER | OPTYPE_UNDO;
 
+	vgroup_operator_subset_select_props(ot, false);
 	RNA_def_boolean(ot->srna, "lock_active", true, "Lock Active",
 	                "Keep the values of the active group while normalizing others");
 }