From e59944dba4eb60f6cadad9f666f10576ebaa3aed Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@blender.org>
Date: Mon, 18 Sep 2023 18:04:37 +0200
Subject: [PATCH] Anim: armature edit mode 'select by bone color' operator

In the armature edit mode 'select similar' operator, replace the 'select
by same bone group' operator with 'select by same bone color'.
---
 source/blender/animrig/ANIM_bonecolor.hh      |  1 +
 source/blender/animrig/intern/bonecolor.cc    |  5 ++
 .../editors/armature/armature_select.cc       | 46 +++++++++++++++++--
 3 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/source/blender/animrig/ANIM_bonecolor.hh b/source/blender/animrig/ANIM_bonecolor.hh
index fa60bbe84c5..bbb67bfce57 100644
--- a/source/blender/animrig/ANIM_bonecolor.hh
+++ b/source/blender/animrig/ANIM_bonecolor.hh
@@ -32,6 +32,7 @@ class BoneColor : public ::BoneColor {
 
   /* Support for storing in a blender::Set<BoneColor>.*/
   bool operator==(const BoneColor &other) const;
+  bool operator!=(const BoneColor &other) const;
   uint64_t hash() const;
 };
 
diff --git a/source/blender/animrig/intern/bonecolor.cc b/source/blender/animrig/intern/bonecolor.cc
index da021b177e0..15547a8df64 100644
--- a/source/blender/animrig/intern/bonecolor.cc
+++ b/source/blender/animrig/intern/bonecolor.cc
@@ -62,6 +62,11 @@ bool BoneColor::operator==(const BoneColor &other) const
   return true;
 }
 
+bool BoneColor::operator!=(const BoneColor &other) const
+{
+  return !(*this == other);
+}
+
 uint64_t BoneColor::hash() const
 {
   if (palette_index >= 0) {
diff --git a/source/blender/editors/armature/armature_select.cc b/source/blender/editors/armature/armature_select.cc
index 42d2decc4f8..7b5bbd0fcf2 100644
--- a/source/blender/editors/armature/armature_select.cc
+++ b/source/blender/editors/armature/armature_select.cc
@@ -44,6 +44,7 @@
 #include "GPU_select.h"
 
 #include "ANIM_bone_collections.h"
+#include "ANIM_bonecolor.hh"
 
 #include "armature_intern.h"
 
@@ -1597,7 +1598,7 @@ enum {
   SIMEDBONE_PREFIX,
   SIMEDBONE_SUFFIX,
   SIMEDBONE_COLLECTION,
-  SIMEDBONE_GROUP,
+  SIMEDBONE_COLOR,
   SIMEDBONE_SHAPE,
 };
 
@@ -1610,7 +1611,7 @@ static const EnumPropertyItem prop_similar_types[] = {
     {SIMEDBONE_PREFIX, "PREFIX", 0, "Prefix", ""},
     {SIMEDBONE_SUFFIX, "SUFFIX", 0, "Suffix", ""},
     {SIMEDBONE_COLLECTION, "BONE_COLLECTION", 0, "Bone Collection", ""},
-    {SIMEDBONE_GROUP, "GROUP", 0, "Group", ""},
+    {SIMEDBONE_COLOR, "COLOR", 0, "Color", ""},
     {SIMEDBONE_SHAPE, "SHAPE", 0, "Shape", ""},
     {0, nullptr, 0, nullptr, nullptr},
 };
@@ -1756,6 +1757,43 @@ static void select_similar_bone_collection(bContext *C)
   }
   MEM_freeN(objects);
 }
+static void select_similar_bone_color(bContext *C)
+{
+  const Scene *scene = CTX_data_scene(C);
+  ViewLayer *view_layer = CTX_data_view_layer(C);
+  EditBone *ebone_act = CTX_data_active_bone(C);
+
+  const blender::animrig::BoneColor &active_bone_color = ebone_act->color.wrap();
+
+  uint objects_len = 0;
+  Object **objects = BKE_view_layer_array_from_objects_in_edit_mode_unique_data(
+      scene, view_layer, CTX_wm_view3d(C), &objects_len);
+  for (uint ob_index = 0; ob_index < objects_len; ob_index++) {
+    Object *ob = objects[ob_index];
+    bArmature *arm = static_cast<bArmature *>(ob->data);
+    bool changed = false;
+
+    LISTBASE_FOREACH (EditBone *, ebone, arm->edbo) {
+      if (!EBONE_SELECTABLE(arm, ebone)) {
+        continue;
+      }
+
+      const blender::animrig::BoneColor &bone_color = ebone->color.wrap();
+      if (bone_color != active_bone_color) {
+        continue;
+      }
+
+      ED_armature_ebone_select_set(ebone, true);
+      changed = true;
+    }
+
+    if (changed) {
+      WM_event_add_notifier(C, NC_OBJECT | ND_BONE_SELECT, ob);
+      DEG_id_tag_update(&ob->id, ID_RECALC_COPY_ON_WRITE);
+    }
+  }
+  MEM_freeN(objects);
+}
 
 static void select_similar_prefix(bContext *C)
 {
@@ -1986,8 +2024,8 @@ static int armature_select_similar_exec(bContext *C, wmOperator *op)
     case SIMEDBONE_COLLECTION:
       select_similar_bone_collection(C);
       break;
-    case SIMEDBONE_GROUP:
-      select_similar_data_pchan(C, STRUCT_SIZE_AND_OFFSET(bPoseChannel, agrp_index));
+    case SIMEDBONE_COLOR:
+      select_similar_bone_color(C);
       break;
     case SIMEDBONE_SHAPE:
       select_similar_data_pchan(C, STRUCT_SIZE_AND_OFFSET(bPoseChannel, custom));