Optimize mouse UI picking on BoxContainer, FlowContainer and GridCont…

…ainer

scene/gui/box_container.cpp

@@ -199,6 +199,9 @@ void BoxContainer::_resort() {
 		delta = -1;
 	}
 
+	cached_children_pos.clear();
+	cached_children_idx.clear();
+
 	for (int i = start; i != end; i += delta) {
 		Control *c = Object::cast_to<Control>(get_child(i));
 		if (!c || !c->is_visible_in_tree()) {
@@ -237,6 +240,8 @@ void BoxContainer::_resort() {
 		}
 
 		fit_child_in_rect(c, rect);
+		cached_children_pos.push_back(vertical ? rect.position.y : rect.position.x);
+		cached_children_idx.push_back(i);
 
 		ofs = to;
 		idx++;
@@ -288,6 +293,34 @@ Size2 BoxContainer::get_minimum_size() const {
 	return minimum;
 }
 
+Vector<CanvasItem *> BoxContainer::get_children_at_pos(const Point2 &p_pos) const {
+	if(cached_children_pos.is_empty()) {
+		return Vector<CanvasItem *>();
+	}
+
+	int coord = vertical ? p_pos.y : p_pos.x;
+	// Custom binary search; we are searching for the closest value to coord, not an exact one.
+	int bin_start = 0;
+	int bin_end = cached_children_pos.size();
+	while (bin_start < bin_end) {
+		int mid = (bin_start + bin_end) / 2;
+		if (coord < cached_children_pos[mid]) {
+			bin_end = mid;
+		} else {
+			bin_start = mid + 1;
+		}
+	}
+
+	int idx = cached_children_idx[MAX(bin_start - 1, 0)];
+
+	Vector<CanvasItem *> children;
+	CanvasItem *child = Object::cast_to<CanvasItem>(get_child(idx));
+	if (child) {
+		children.push_back(child);
+	}
+	return children;
+}
+
 void BoxContainer::_update_theme_item_cache() {
 	Container::_update_theme_item_cache();
 

scene/gui/box_container.h

@@ -51,6 +51,9 @@ class BoxContainer : public Container {
 		int separation = 0;
 	} theme_cache;
 
+	Vector<int> cached_children_pos;
+	Vector<int> cached_children_idx;
+
 	void _resort();
 
 protected:
@@ -72,6 +75,7 @@ class BoxContainer : public Container {
 	bool is_vertical() const;
 
 	virtual Size2 get_minimum_size() const override;
+	virtual Vector<CanvasItem *> get_children_at_pos(const Point2 &p_pos) const override;
 
 	virtual Vector<int> get_allowed_size_flags_horizontal() const override;
 	virtual Vector<int> get_allowed_size_flags_vertical() const override;

scene/gui/container.cpp

@@ -81,6 +81,17 @@ void Container::remove_child_notify(Node *p_child) {
 	queue_sort();
 }
 
+Vector<CanvasItem *> Container::get_children_at_pos(const Point2 &p_pos) const {
+	Vector<CanvasItem *> children;
+	for (int i = get_child_count() - 1; i >= 0; i--) {
+		CanvasItem *child = Object::cast_to<CanvasItem>(get_child(i));
+		if (child && child->is_visible()) {
+			children.push_back(child);
+		}
+	}
+	return children;
+}
+
 void Container::_sort_children() {
 	if (!is_inside_tree()) {
 		return;

scene/gui/container.h

@@ -32,6 +32,7 @@
 #define CONTAINER_H
 
 #include "scene/gui/control.h"
+#include "core/templates/vector.h"
 
 class Container : public Control {
 	GDCLASS(Container, Control);
@@ -59,6 +60,7 @@ class Container : public Control {
 	};
 
 	void fit_child_in_rect(Control *p_child, const Rect2 &p_rect);
+	virtual Vector<CanvasItem*> get_children_at_pos(const Point2 &p_pos) const;
 
 	virtual Vector<int> get_allowed_size_flags_horizontal() const;
 	virtual Vector<int> get_allowed_size_flags_vertical() const;

scene/gui/flow_container.cpp

@@ -57,6 +57,9 @@ void FlowContainer::_resort() {
 	int current_container_size = vertical ? get_rect().size.y : get_rect().size.x;
 	int children_in_current_line = 0;
 
+	cached_children_pos.clear();
+	cached_children_idx.clear();
+
 	// First pass for line wrapping and minimum size calculation.
 	for (int i = 0; i < get_child_count(); i++) {
 		Control *child = Object::cast_to<Control>(get_child(i));
@@ -128,6 +131,9 @@ void FlowContainer::_resort() {
 	ofs.x = 0;
 	ofs.y = 0;
 
+	cached_children_pos.push_back(0); // First child
+	cached_children_idx.push_back(0);
+
 	for (int i = 0; i < get_child_count(); i++) {
 		Control *child = Object::cast_to<Control>(get_child(i));
 		if (!child || !child->is_visible()) {
@@ -145,10 +151,13 @@ void FlowContainer::_resort() {
 			if (vertical) {
 				ofs.x += line_data.min_line_height + theme_cache.h_separation;
 				ofs.y = 0;
+				cached_children_pos.push_back(ofs.x);
 			} else {
 				ofs.x = 0;
 				ofs.y += line_data.min_line_height + theme_cache.v_separation;
+				cached_children_pos.push_back(ofs.y);
 			}
+			cached_children_idx.push_back(i);
 			line_data = lines_data[current_line_idx];
 		}
 
@@ -245,6 +254,42 @@ Size2 FlowContainer::get_minimum_size() const {
 	return minimum;
 }
 
+Vector<CanvasItem *> FlowContainer::get_children_at_pos(const Point2 &p_pos) const
+{
+	if (cached_children_pos.is_empty()) {
+		return Vector<CanvasItem *>();
+	}
+
+	int coord = vertical ? p_pos.x : p_pos.y;
+	// Custom binary search; we are searching for the closest value to coord, not an exact one.
+	int bin_start = 0;
+	int bin_end = cached_children_pos.size();
+	while(bin_start < bin_end) {
+		int mid = (bin_start + bin_end) / 2;
+		if(coord < cached_children_pos[mid]) {
+			bin_end = mid;
+		} else {
+			bin_start = mid + 1;
+		}
+	}
+
+	int start_idx = cached_children_idx[MAX(bin_start - 1, 0)];
+	int end_idx = cached_children_idx[MIN(bin_end, cached_children_idx.size() - 1)];
+
+	if(start_idx == end_idx) {
+		end_idx = get_child_count();
+	}
+
+	Vector<CanvasItem *> children;
+	for(end_idx--; end_idx >= start_idx; end_idx--) {
+		CanvasItem *child = Object::cast_to<CanvasItem>(get_child(end_idx));
+		if (child && child->is_visible()) {
+			children.push_back(child);
+		}
+	}
+	return children;
+}
+
 Vector<int> FlowContainer::get_allowed_size_flags_horizontal() const {
 	Vector<int> flags;
 	flags.append(SIZE_FILL);

scene/gui/flow_container.h

@@ -32,6 +32,7 @@
 #define FLOW_CONTAINER_H
 
 #include "scene/gui/container.h"
+#include "core/templates/vector.h"
 
 class FlowContainer : public Container {
 	GDCLASS(FlowContainer, Container);
@@ -57,6 +58,10 @@ class FlowContainer : public Container {
 
 	void _resort();
 
+	Vector<int> cached_children_pos;
+	Vector<int> cached_children_idx;
+	Vector<Node *> cached_visible_children;
+
 protected:
 	bool is_fixed = false;
 
@@ -76,6 +81,7 @@ class FlowContainer : public Container {
 	bool is_vertical() const;
 
 	virtual Size2 get_minimum_size() const override;
+	virtual Vector<CanvasItem *> get_children_at_pos(const Point2& p_pos) const override;
 
 	virtual Vector<int> get_allowed_size_flags_horizontal() const override;
 	virtual Vector<int> get_allowed_size_flags_vertical() const override;

scene/gui/grid_container.cpp

@@ -190,6 +190,11 @@ void GridContainer::_notification(int p_what) {
 			}
 
 			valid_controls_index = 0;
+			cached_children_pos.clear();
+			cached_children_idx.clear();
+
+			cached_children_pos.push_back(0);
+			cached_children_idx.push_back(0);
 			for (int i = 0; i < get_child_count(); i++) {
 				Control *c = Object::cast_to<Control>(get_child(i));
 				if (!c || !c->is_visible_in_tree()) {
@@ -212,6 +217,8 @@ void GridContainer::_notification(int p_what) {
 							// Apply the remaining pixel of the previous row.
 							row_ofs++;
 						}
+						cached_children_pos.push_back(row_ofs);
+						cached_children_idx.push_back(i);
 					}
 				}
 
@@ -320,4 +327,38 @@ Size2 GridContainer::get_minimum_size() const {
 	return ms;
 }
 
+Vector<CanvasItem *> GridContainer::get_children_at_pos(const Point2 &p_pos) const {
+	if (cached_children_pos.is_empty()) {
+		return Vector<CanvasItem *>();
+	}
+
+	// Custom binary search; we are searching for the closest value to coord, not an exact one.
+	int bin_start = 0;
+	int bin_end = cached_children_pos.size();
+	while (bin_start < bin_end) {
+		int mid = (bin_start + bin_end) / 2;
+		if (p_pos.y < cached_children_pos[mid]) {
+			bin_end = mid;
+		} else {
+			bin_start = mid + 1;
+		}
+	}
+
+	int start_idx = cached_children_idx[MAX(bin_start - 1, 0)];
+	int end_idx = cached_children_idx[MIN(bin_end, cached_children_idx.size() - 1)];
+
+	if (start_idx == end_idx) {
+		end_idx = get_child_count();
+	}
+
+	Vector<CanvasItem *> children;
+	for (end_idx--; end_idx >= start_idx; end_idx--) {
+		CanvasItem *child = Object::cast_to<CanvasItem>(get_child(end_idx));
+		if (child && child->is_visible()) {
+			children.push_back(child);
+		}
+	}
+	return children;
+}
+
 GridContainer::GridContainer() {}

scene/gui/grid_container.h

@@ -43,6 +43,9 @@ class GridContainer : public Container {
 		int v_separation = 0;
 	} theme_cache;
 
+	Vector<int> cached_children_pos;
+	Vector<int> cached_children_idx;
+
 protected:
 	virtual void _update_theme_item_cache() override;
 
@@ -53,6 +56,7 @@ class GridContainer : public Container {
 	void set_columns(int p_columns);
 	int get_columns() const;
 	virtual Size2 get_minimum_size() const override;
+	virtual Vector<CanvasItem *> get_children_at_pos(const Point2 &p_pos) const override;
 
 	GridContainer();
 };

scene/main/viewport.cpp

@@ -1681,8 +1681,26 @@ Control *Viewport::_gui_find_control_at_pos(CanvasItem *p_node, const Point2 &p_
 	}
 
 	Control *c = Object::cast_to<Control>(p_node);
+	Point2 point = matrix.affine_inverse().xform(p_global);
+
+	if (!c || !c->is_clipping_contents() || c->has_point(point)) {
+		// If node is a Container, it can optimize the search at the given position.
+		Container *cnt = Object::cast_to<Container>(p_node);
+		if (cnt && cnt->has_point(point)) {
+			Vector<CanvasItem *> children = cnt->get_children_at_pos(point);
+			for (CanvasItem *ci : children) {
+				if (ci->is_set_as_top_level()) {
+					continue;
+				}
+
+				Control *ret = _gui_find_control_at_pos(ci, p_global, matrix);
+				if (ret) {
+					return ret;
+				}
+			}
+		}
 
-	if (!c || !c->is_clipping_contents() || c->has_point(matrix.affine_inverse().xform(p_global))) {
+		// Fallback; iterate over all children
 		for (int i = p_node->get_child_count() - 1; i >= 0; i--) {
 			CanvasItem *ci = Object::cast_to<CanvasItem>(p_node->get_child(i));
 			if (!ci || ci->is_set_as_top_level()) {