Multi-Objects Metaball: Selection refactor - fix pick and box selection

This is inspired/based on the code we use for armature bone selection.
Both pick selection, and box selection should be working now.

5 files changed, 148 insertions, 69 deletions

diff --git a/source/blender/draw/modes/edit_metaball_mode.c b/source/blender/draw/modes/edit_metaball_mode.c
index 9e27ed7bd17..912c30c1482 100644
--- a/source/blender/draw/modes/edit_metaball_mode.c
+++ b/source/blender/draw/modes/edit_metaball_mode.c
@@ -31,6 +31,8 @@
 #include "BKE_object.h"
 #include "BKE_mball.h"
 
+#include "ED_mball.h"
+
 /* If builtin shaders are needed */
 #include "GPU_shader.h"
 #include "GPU_select.h"
@@ -147,8 +149,6 @@ static void EDIT_METABALL_cache_populate(void *vedata, Object *ob)
 
 			const bool is_select = DRW_state_is_select();
 
-			int selection_id = 0;
-
 			float draw_scale_xform[3][4]; /* Matrix of Scale and Translation */
 			{
 				float scamat[3][3];
@@ -165,7 +165,8 @@ static void EDIT_METABALL_cache_populate(void *vedata, Object *ob)
 				copy_v3_v3(draw_scale_xform[2], scamat[2]);
 			}
 
-			for (MetaElem *ml = mb->editelems->first; ml != NULL; ml = ml->next) {
+			int selection_id = ob->select_color;
+			for (MetaElem *ml = mb->editelems->first; ml != NULL; ml = ml->next, selection_id += 0x10000) {
 				float world_pos[3];
 				mul_v3_m4v3(world_pos, ob->obmat, &ml->x);
 				draw_scale_xform[0][3] = world_pos[0];
@@ -178,8 +179,7 @@ static void EDIT_METABALL_cache_populate(void *vedata, Object *ob)
 				else color = col_radius;
 
 				if (is_select) {
-					ml->selcol1 = ++selection_id;
-					DRW_select_load_id(selection_id);
+					DRW_select_load_id(selection_id | MBALLSEL_RADIUS);
 				}
 
 				DRW_shgroup_call_dynamic_add(group, draw_scale_xform, &ml->rad, color);
@@ -188,8 +188,7 @@ static void EDIT_METABALL_cache_populate(void *vedata, Object *ob)
 				else color = col_stiffness;
 
 				if (is_select) {
-					ml->selcol2 = ++selection_id;
-					DRW_select_load_id(selection_id);
+					DRW_select_load_id(selection_id | MBALLSEL_STIFF);
 				}
 
 				DRW_shgroup_call_dynamic_add(group, draw_scale_xform, &draw_stiffness_radius, color);
diff --git a/source/blender/editors/include/ED_mball.h b/source/blender/editors/include/ED_mball.h
index 46c34d36efc..47ea3929883 100644
--- a/source/blender/editors/include/ED_mball.h
+++ b/source/blender/editors/include/ED_mball.h
@@ -51,4 +51,11 @@ void ED_mball_editmball_load(struct Object *obedit);
 /* editmball_undo.c */
 void ED_mball_undosys_type(struct UndoType *ut);
 
+
+#define MBALLSEL_STIFF    (1 << 29)
+#define MBALLSEL_RADIUS   (1 << 30)
+#define MBALLSEL_ANY      (MBALLSEL_STIFF | MBALLSEL_RADIUS)
+
+#define MBALL_NOSEL   (1u << 31u)
+
 #endif  /* __ED_MBALL_H__ */
diff --git a/source/blender/editors/metaball/mball_edit.c b/source/blender/editors/metaball/mball_edit.c
index 83e454c0d35..6d1ba2d3fe0 100644
--- a/source/blender/editors/metaball/mball_edit.c
+++ b/source/blender/editors/metaball/mball_edit.c
@@ -56,6 +56,7 @@
 #include "DEG_depsgraph.h"
 
 #include "ED_mball.h"
+#include "ED_object.h"
 #include "ED_screen.h"
 #include "ED_select_utils.h"
 #include "ED_view3d.h"
@@ -676,10 +677,7 @@ void MBALL_OT_reveal_metaelems(wmOperatorType *ot)
 bool ED_mball_select_pick(bContext *C, const int mval[2], bool extend, bool deselect, bool toggle)
 {
 	static MetaElem *startelem = NULL;
-	Object *obedit = CTX_data_edit_object(C);
 	ViewContext vc;
-	MetaBall *mb = (MetaBall *)obedit->data;
-	MetaElem *ml, *ml_act = NULL;
 	int a, hits;
 	unsigned int buffer[MAXPICKBUF];
 	rcti rect;
@@ -692,66 +690,116 @@ bool ED_mball_select_pick(bContext *C, const int mval[2], bool extend, bool dese
 	        &vc, buffer, MAXPICKBUF, &rect,
 	        VIEW3D_SELECT_PICK_NEAREST, VIEW3D_SELECT_FILTER_NOP);
 
-	/* does startelem exist? */
-	ml = mb->editelems->first;
-	while (ml) {
-		if (ml == startelem) break;
-		ml = ml->next;
-	}
-
-	if (ml == NULL) startelem = mb->editelems->first;
+	FOREACH_BASE_IN_EDIT_MODE_BEGIN (vc.view_layer, base) {
+		ED_view3d_viewcontext_init_object(&vc, base->object);
+		MetaBall *mb = (MetaBall *)base->object->data;
+		MetaElem *ml, *ml_act = NULL;
 
-	if (hits > 0) {
-		ml = startelem;
+		/* does startelem exist? */
+		ml = mb->editelems->first;
 		while (ml) {
-			for (a = 0; a < hits; a++) {
-				/* index converted for gl stuff */
-				if (ml->selcol1 == buffer[4 * a + 3]) {
-					ml->flag |= MB_SCALE_RAD;
-					ml_act = ml;
-				}
-				if (ml->selcol2 == buffer[4 * a + 3]) {
-					ml->flag &= ~MB_SCALE_RAD;
-					ml_act = ml;
-				}
-			}
-			if (ml_act) break;
-			ml = ml->next;
-			if (ml == NULL) ml = mb->editelems->first;
 			if (ml == startelem) break;
+			ml = ml->next;
 		}
 
-		/* When some metaelem was found, then it is necessary to select or
-		 * deselect it. */
-		if (ml_act) {
-			if (extend) {
-				ml_act->flag |= SELECT;
-			}
-			else if (deselect) {
-				ml_act->flag &= ~SELECT;
+		if (ml == NULL) startelem = mb->editelems->first;
+
+		if (hits > 0) {
+			int metaelem_id = 0;
+			ml = startelem;
+			while (ml) {
+				for (a = 0; a < hits; a++) {
+					int hitresult = buffer[(4 * a) + 3];
+					if (hitresult == -1) {
+						continue;
+					}
+					else if (hitresult & MBALL_NOSEL) {
+						continue;
+					}
+
+					const uint hit_object = hitresult & 0xFFFF;
+					if (vc.obedit->select_color != hit_object) {
+						continue;
+					}
+
+					if (metaelem_id != (hitresult & 0xFFFF0000 & ~(MBALLSEL_ANY))) {
+						continue;
+					}
+
+					if (hitresult & MBALLSEL_RADIUS) {
+						ml->flag |= MB_SCALE_RAD;
+						ml_act = ml;
+						break;
+					}
+
+					if (hitresult & MBALLSEL_STIFF) {
+						ml->flag &= ~MB_SCALE_RAD;
+						ml_act = ml;
+						break;
+					}
+				}
+
+				if (ml_act) break;
+				ml = ml->next;
+				if (ml == NULL) ml = mb->editelems->first;
+				if (ml == startelem) break;
+
+				metaelem_id += 0x10000;
 			}
-			else if (toggle) {
-				if (ml_act->flag & SELECT)
+
+			/* When some metaelem was found, then it is necessary to select or
+			* deselect it. */
+			if (ml_act) {
+				if (!extend && !deselect && !toggle) {
+					uint objects_len;
+					Object **objects = BKE_view_layer_array_from_objects_in_edit_mode_unique_data(vc.view_layer, &objects_len);
+					for (uint ob_index = 0; ob_index < objects_len; ob_index++) {
+						Object *ob_iter = objects[ob_index];
+
+						if (ob_iter == base->object) {
+							continue;
+						}
+
+						BKE_mball_deselect_all((MetaBall *)ob_iter->data);
+						DEG_id_tag_update(ob_iter->data, DEG_TAG_SELECT_UPDATE);
+						WM_event_add_notifier(C, NC_GEOM | ND_SELECT, ob_iter->data);
+					}
+					MEM_freeN(objects);
+				}
+
+				if (extend) {
+					ml_act->flag |= SELECT;
+				}
+				else if (deselect) {
 					ml_act->flag &= ~SELECT;
-				else
+				}
+				else if (toggle) {
+					if (ml_act->flag & SELECT)
+						ml_act->flag &= ~SELECT;
+					else
+						ml_act->flag |= SELECT;
+				}
+				else {
+					/* Deselect all existing metaelems */
+					BKE_mball_deselect_all(mb);
+
+					/* Select only metaelem clicked on */
 					ml_act->flag |= SELECT;
-			}
-			else {
-				/* Deselect all existing metaelems */
-				BKE_mball_deselect_all(mb);
+				}
 
-				/* Select only metaelem clicked on */
-				ml_act->flag |= SELECT;
-			}
+				mb->lastelem = ml_act;
 
-			mb->lastelem = ml_act;
+				DEG_id_tag_update(&mb->id, DEG_TAG_SELECT_UPDATE);
+				WM_event_add_notifier(C, NC_GEOM | ND_SELECT, mb);
 
-			DEG_id_tag_update(&mb->id, DEG_TAG_SELECT_UPDATE);
-			WM_event_add_notifier(C, NC_GEOM | ND_SELECT, mb);
+				if (vc.view_layer->basact != base) {
+					ED_object_base_activate(C, base);
+				}
 
-			return true;
+				return true;
+			}
 		}
-	}
+	} FOREACH_BASE_IN_EDIT_MODE_END;
 
 	return false;
 }
diff --git a/source/blender/editors/space_view3d/view3d_select.c b/source/blender/editors/space_view3d/view3d_select.c
index 230bbd11450..c421c96e530 100644
--- a/source/blender/editors/space_view3d/view3d_select.c
+++ b/source/blender/editors/space_view3d/view3d_select.c
@@ -2169,7 +2169,8 @@ static int do_meta_box_select(
         ViewContext *vc,
         const rcti *rect, const eSelectOp sel_op)
 {
-	MetaBall *mb = (MetaBall *)vc->obedit->data;
+	Object *ob = vc->obedit;
+	MetaBall *mb = (MetaBall *)ob->data;
 	MetaElem *ml;
 	int a;
 
@@ -2184,27 +2185,50 @@ static int do_meta_box_select(
 		BKE_mball_deselect_all(mb);
 	}
 
-	for (ml = mb->editelems->first; ml; ml = ml->next) {
-		bool is_inside_1 = false;
-		bool is_inside_2 = false;
+	int metaelem_id = 0;
+	for (ml = mb->editelems->first; ml; ml = ml->next, metaelem_id += 0x10000) {
+		bool is_inside_radius = false;
+		bool is_inside_stiff = false;
+
 		for (a = 0; a < hits; a++) {
-			if (ml->selcol1 == buffer[(4 * a) + 3]) {
-				is_inside_1 = true;
+			int hitresult = buffer[(4 * a) + 3];
+
+			if (hitresult == -1) {
+				continue;
+			}
+			else if (hitresult & MBALL_NOSEL) {
+				continue;
+			}
+
+			const uint hit_object = hitresult & 0xFFFF;
+			if (vc->obedit->select_color != hit_object) {
+				continue;
+			}
+
+			if (metaelem_id != (hitresult & 0xFFFF0000 & ~(MBALLSEL_ANY))) {
+				continue;
+			}
+
+			if (hitresult & MBALLSEL_RADIUS) {
+				is_inside_radius = true;
 				break;
 			}
-			if (ml->selcol2 == buffer[(4 * a) + 3]) {
-				is_inside_2 = true;
+
+			if (hitresult & MBALLSEL_STIFF) {
+				is_inside_stiff = true;
 				break;
 			}
 		}
-		if (is_inside_1) {
+		if (is_inside_radius) {
 			ml->flag |= MB_SCALE_RAD;
 		}
-		if (is_inside_2) {
+		if (is_inside_stiff) {
 			ml->flag &= ~MB_SCALE_RAD;
 		}
+
 		const bool is_select = (ml->flag & SELECT);
-		const bool is_inside = is_inside_1 || is_inside_2;
+		const bool is_inside = is_inside_radius || is_inside_stiff;
+
 		const int sel_op_result = ED_select_op_action_deselected(sel_op, is_select, is_inside);
 		if (sel_op_result != -1) {
 			SET_FLAG_FROM_TEST(ml->flag, sel_op_result, SELECT);
diff --git a/source/blender/makesdna/DNA_meta_types.h b/source/blender/makesdna/DNA_meta_types.h
index c312d45567f..3f0b1b302c7 100644
--- a/source/blender/makesdna/DNA_meta_types.h
+++ b/source/blender/makesdna/DNA_meta_types.h
@@ -47,7 +47,8 @@ typedef struct MetaElem {
 
 	struct BoundBox *bb;        /* Bound Box of MetaElem */
 
-	short type, flag, selcol1, selcol2;
+	short type, flag;
+	short pad[2];
 	float x, y, z;          /* Position of center of MetaElem */
 	float quat[4];          /* Rotation of MetaElem (MUST be kept normalized) */
 	float expx; /* dimension parameters, used for some types like cubes */