Weight Paint: add a pie menu for locking and unlocking vertex groups.

Provide different options for locking and unlocking vertex groups
using bone selection, accessible via a pie menu triggered via the
'K' hotkey. To implement a variety of operations, extend the old
operator with a new option to mask it by bone selection. If the
X Mirror option is enabled, selection is automatically mirrored.

This follows D6533 as the next step in improving accessibility of
vertex group locking during weight painting.

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

4 files changed, 199 insertions, 8 deletions

diff --git a/release/scripts/presets/keyconfig/keymap_data/blender_default.py b/release/scripts/presets/keyconfig/keymap_data/blender_default.py
index 3f5b7917f9c..cd531119155 100644
--- a/release/scripts/presets/keyconfig/keymap_data/blender_default.py
+++ b/release/scripts/presets/keyconfig/keymap_data/blender_default.py
@@ -3869,6 +3869,7 @@ def km_weight_paint(params):
          {"properties": [("data_path", 'weight_paint_object.data.use_paint_mask_vertex')]}),
         ("wm.context_toggle", {"type": 'S', "value": 'PRESS', "shift": True},
          {"properties": [("data_path", 'tool_settings.weight_paint.brush.use_smooth_stroke')]}),
+        op_menu_pie("VIEW3D_MT_wpaint_vgroup_lock_pie", {"type" : 'K', "value": 'PRESS'}),
         *_template_items_context_panel("VIEW3D_PT_paint_weight_context_menu", params.context_menu_event),
     ])
 
diff --git a/release/scripts/startup/bl_ui/properties_data_mesh.py b/release/scripts/startup/bl_ui/properties_data_mesh.py
index 3edce6b3b52..d6aa986613d 100644
--- a/release/scripts/startup/bl_ui/properties_data_mesh.py
+++ b/release/scripts/startup/bl_ui/properties_data_mesh.py
@@ -55,9 +55,12 @@ class MESH_MT_vertex_group_context_menu(Menu):
         layout.operator("object.vertex_group_remove", text="Delete All Unlocked Groups").all_unlocked = True
         layout.operator("object.vertex_group_remove", text="Delete All Groups").all = True
         layout.separator()
-        layout.operator("object.vertex_group_lock", icon='LOCKED', text="Lock All").action = 'LOCK'
-        layout.operator("object.vertex_group_lock", icon='UNLOCKED', text="UnLock All").action = 'UNLOCK'
-        layout.operator("object.vertex_group_lock", text="Lock Invert All").action = 'INVERT'
+        props = layout.operator("object.vertex_group_lock", icon='LOCKED', text="Lock All")
+        props.action, props.mask = 'LOCK', 'ALL'
+        props = layout.operator("object.vertex_group_lock", icon='UNLOCKED', text="UnLock All")
+        props.action, props.mask = 'UNLOCK', 'ALL'
+        props = layout.operator("object.vertex_group_lock", text="Lock Invert All")
+        props.action, props.mask = 'INVERT', 'ALL'
 
 
 class MESH_MT_shape_key_context_menu(Menu):
diff --git a/release/scripts/startup/bl_ui/space_view3d.py b/release/scripts/startup/bl_ui/space_view3d.py
index 537a4a4761a..cffa52c8052 100644
--- a/release/scripts/startup/bl_ui/space_view3d.py
+++ b/release/scripts/startup/bl_ui/space_view3d.py
@@ -5020,6 +5020,39 @@ class VIEW3D_MT_sculpt_mask_edit_pie(Menu):
         op.auto_iteration_count = False
 
 
+class VIEW3D_MT_wpaint_vgroup_lock_pie(Menu):
+    bl_label = "Vertex Group Locks"
+
+    def draw(self, _context):
+        layout = self.layout
+        pie = layout.menu_pie()
+
+        # 1: Left
+        op = pie.operator("object.vertex_group_lock", icon='LOCKED', text="Lock All")
+        op.action, op.mask = 'LOCK', 'ALL'
+        # 2: Right
+        op = pie.operator("object.vertex_group_lock", icon='UNLOCKED', text="Unlock All")
+        op.action, op.mask = 'UNLOCK', 'ALL'
+        # 3: Down
+        op = pie.operator("object.vertex_group_lock", icon='UNLOCKED', text="Unlock Selected")
+        op.action, op.mask = 'UNLOCK', 'SELECTED'
+        # 4: Up
+        op = pie.operator("object.vertex_group_lock", icon='LOCKED', text="Lock Selected")
+        op.action, op.mask = 'LOCK', 'SELECTED'
+        # 5: Up/Left
+        op = pie.operator("object.vertex_group_lock", icon='LOCKED', text="Lock Unselected")
+        op.action, op.mask = 'LOCK', 'UNSELECTED'
+        # 6: Up/Right
+        op = pie.operator("object.vertex_group_lock", text="Lock Only Selected")
+        op.action, op.mask = 'LOCK', 'INVERT_UNSELECTED'
+        # 7: Down/Left
+        op = pie.operator("object.vertex_group_lock", text="Lock Only Unselected")
+        op.action, op.mask = 'UNLOCK', 'INVERT_UNSELECTED'
+        # 8: Down/Right
+        op = pie.operator("object.vertex_group_lock", text="Invert Locks")
+        op.action, op.mask = 'INVERT', 'ALL'
+
+
 # ********** Panel **********
 
 
@@ -7090,6 +7123,7 @@ classes = (
     VIEW3D_MT_orientations_pie,
     VIEW3D_MT_proportional_editing_falloff_pie,
     VIEW3D_MT_sculpt_mask_edit_pie,
+    VIEW3D_MT_wpaint_vgroup_lock_pie,
     VIEW3D_PT_active_tool,
     VIEW3D_PT_active_tool_duplicate,
     VIEW3D_PT_view3d_properties,
diff --git a/source/blender/editors/object/object_vgroup.c b/source/blender/editors/object/object_vgroup.c
index 9ccbc7a1a0a..17cc266b9c8 100644
--- a/source/blender/editors/object/object_vgroup.c
+++ b/source/blender/editors/object/object_vgroup.c
@@ -47,6 +47,7 @@
 #include "BKE_context.h"
 #include "BKE_customdata.h"
 #include "BKE_deform.h"
+#include "BKE_mesh.h"
 #include "BKE_mesh_mapping.h"
 #include "BKE_mesh_runtime.h"
 #include "BKE_editmesh.h"
@@ -1687,13 +1688,78 @@ static const EnumPropertyItem vgroup_lock_actions[] = {
     {0, NULL, 0, NULL, NULL},
 };
 
-static void vgroup_lock_all(Object *ob, int action)
+enum {
+  VGROUP_MASK_ALL,
+  VGROUP_MASK_SELECTED,
+  VGROUP_MASK_UNSELECTED,
+  VGROUP_MASK_INVERT_UNSELECTED,
+};
+
+static const EnumPropertyItem vgroup_lock_mask[] = {
+    {VGROUP_MASK_ALL, "ALL", 0, "All", "Apply action to all vertex groups"},
+    {VGROUP_MASK_SELECTED, "SELECTED", 0, "Selected", "Apply to selected vertex groups"},
+    {VGROUP_MASK_UNSELECTED, "UNSELECTED", 0, "Unselected", "Apply to unselected vertex groups"},
+    {VGROUP_MASK_INVERT_UNSELECTED,
+     "INVERT_UNSELECTED",
+     0,
+     "Invert Unselected",
+     "Apply the opposite of Lock/Unlock to unselected vertex groups"},
+    {0, NULL, 0, NULL, NULL},
+};
+
+static bool *vgroup_selected_get(Object *ob)
+{
+  int sel_count = 0, defbase_tot = BLI_listbase_count(&ob->defbase);
+  bool *mask;
+
+  if (ob->mode & OB_MODE_WEIGHT_PAINT) {
+    mask = BKE_object_defgroup_selected_get(ob, defbase_tot, &sel_count);
+
+    /* Mirror the selection if X Mirror is enabled. */
+    Mesh *me = BKE_mesh_from_object(ob);
+
+    if (me && (me->editflag & ME_EDIT_MIRROR_X) != 0) {
+      BKE_object_defgroup_mirror_selection(ob, defbase_tot, mask, mask, &sel_count);
+    }
+  }
+  else {
+    mask = MEM_callocN(defbase_tot * sizeof(bool), __func__);
+  }
+
+  if (sel_count == 0 && ob->actdef >= 1 && ob->actdef <= defbase_tot) {
+    mask[ob->actdef - 1] = true;
+  }
+
+  return mask;
+}
+
+static void vgroup_lock_all(Object *ob, int action, int mask)
 {
   bDeformGroup *dg;
+  bool *selected = NULL;
+  int i;
+
+  if (mask != VGROUP_MASK_ALL) {
+    selected = vgroup_selected_get(ob);
+  }
 
   if (action == VGROUP_TOGGLE) {
     action = VGROUP_LOCK;
-    for (dg = ob->defbase.first; dg; dg = dg->next) {
+
+    for (dg = ob->defbase.first, i = 0; dg; dg = dg->next, i++) {
+      switch (mask) {
+        case VGROUP_MASK_INVERT_UNSELECTED:
+        case VGROUP_MASK_SELECTED:
+          if (!selected[i])
+            continue;
+          break;
+        case VGROUP_MASK_UNSELECTED:
+          if (selected[i])
+            continue;
+          break;
+        default:;
+      }
+
       if (dg->flag & DG_LOCK_WEIGHT) {
         action = VGROUP_UNLOCK;
         break;
@@ -1701,7 +1767,19 @@ static void vgroup_lock_all(Object *ob, int action)
     }
   }
 
-  for (dg = ob->defbase.first; dg; dg = dg->next) {
+  for (dg = ob->defbase.first, i = 0; dg; dg = dg->next, i++) {
+    switch (mask) {
+      case VGROUP_MASK_SELECTED:
+        if (!selected[i])
+          continue;
+        break;
+      case VGROUP_MASK_UNSELECTED:
+        if (selected[i])
+          continue;
+        break;
+      default:;
+    }
+
     switch (action) {
       case VGROUP_LOCK:
         dg->flag |= DG_LOCK_WEIGHT;
@@ -1713,6 +1791,14 @@ static void vgroup_lock_all(Object *ob, int action)
         dg->flag ^= DG_LOCK_WEIGHT;
         break;
     }
+
+    if (mask == VGROUP_MASK_INVERT_UNSELECTED && !selected[i]) {
+      dg->flag ^= DG_LOCK_WEIGHT;
+    }
+  }
+
+  if (selected) {
+    MEM_freeN(selected);
   }
 }
 
@@ -3176,24 +3262,84 @@ static int vertex_group_lock_exec(bContext *C, wmOperator *op)
   Object *ob = CTX_data_active_object(C);
 
   int action = RNA_enum_get(op->ptr, "action");
+  int mask = RNA_enum_get(op->ptr, "mask");
 
-  vgroup_lock_all(ob, action);
+  vgroup_lock_all(ob, action, mask);
 
   WM_event_add_notifier(C, NC_OBJECT | ND_DRAW, ob);
 
   return OPERATOR_FINISHED;
 }
 
+static char *vertex_group_lock_description(struct bContext *UNUSED(C),
+                                           struct wmOperatorType *UNUSED(op),
+                                           struct PointerRNA *params)
+{
+  int action = RNA_enum_get(params, "action");
+  int mask = RNA_enum_get(params, "mask");
+
+  const char *action_str, *target_str;
+
+  switch (action) {
+    case VGROUP_LOCK:
+      action_str = "Lock";
+      break;
+    case VGROUP_UNLOCK:
+      action_str = "Unlock";
+      break;
+    case VGROUP_TOGGLE:
+      action_str = "Toggle locks of";
+      break;
+    case VGROUP_INVERT:
+      action_str = "Invert locks of";
+      break;
+    default:
+      return NULL;
+  }
+
+  switch (mask) {
+    case VGROUP_MASK_ALL:
+      target_str = "all";
+      break;
+    case VGROUP_MASK_SELECTED:
+      target_str = "selected";
+      break;
+    case VGROUP_MASK_UNSELECTED:
+      target_str = "unselected";
+      break;
+    case VGROUP_MASK_INVERT_UNSELECTED:
+      switch (action) {
+        case VGROUP_INVERT:
+          target_str = "selected";
+          break;
+        case VGROUP_LOCK:
+          target_str = "selected and unlock unselected";
+          break;
+        case VGROUP_UNLOCK:
+          target_str = "selected and lock unselected";
+          break;
+        default:
+          target_str = "all and invert unselected";
+      }
+      break;
+    default:
+      return NULL;
+  }
+
+  return BLI_sprintfN("%s %s vertex groups of the active object", action_str, target_str);
+}
+
 void OBJECT_OT_vertex_group_lock(wmOperatorType *ot)
 {
   /* identifiers */
   ot->name = "Change the Lock On Vertex Groups";
   ot->idname = "OBJECT_OT_vertex_group_lock";
-  ot->description = "Change the lock state of all vertex groups of active object";
+  ot->description = "Change the lock state of all or some vertex groups of active object";
 
   /* api callbacks */
   ot->poll = vertex_group_poll;
   ot->exec = vertex_group_lock_exec;
+  ot->get_description = vertex_group_lock_description;
 
   /* flags */
   ot->flag = OPTYPE_REGISTER | OPTYPE_UNDO;
@@ -3204,6 +3350,13 @@ void OBJECT_OT_vertex_group_lock(wmOperatorType *ot)
                VGROUP_TOGGLE,
                "Action",
                "Lock action to execute on vertex groups");
+
+  RNA_def_enum(ot->srna,
+               "mask",
+               vgroup_lock_mask,
+               VGROUP_MASK_ALL,
+               "Mask",
+               "Apply the action based on vertex group selection");
 }
 
 static int vertex_group_invert_exec(bContext *C, wmOperator *op)