allow to reserve space in OAHashMap explicitly and also in AStar.

* also handle overflow occurring in _get_probe_length

3 files changed, 67 insertions, 9 deletions

diff --git a/core/math/a_star.cpp b/core/math/a_star.cpp
index aea42a1edf..60b7326c29 100644
--- a/core/math/a_star.cpp
+++ b/core/math/a_star.cpp
@@ -243,6 +243,20 @@ void AStar::clear() {
 	points.clear();
 }
 
+int AStar::get_point_count() const {
+	return points.get_num_elements();
+}
+
+int AStar::get_point_capacity() const {
+	return points.get_capacity();
+}
+
+void AStar::reserve_space(int p_num_nodes) {
+	ERR_FAIL_COND_MSG(p_num_nodes <= 0, "New capacity must be greater than 0, was: " + itos(p_num_nodes) + ".");
+	ERR_FAIL_COND_MSG((uint32_t)p_num_nodes < points.get_capacity(), "New capacity must be greater than current capacity: " + itos(points.get_capacity()) + ", new was: " + itos(p_num_nodes) + ".");
+	points.reserve(p_num_nodes);
+}
+
 int AStar::get_closest_point(const Vector3 &p_point) const {
 
 	int closest_id = -1;
@@ -521,6 +535,9 @@ void AStar::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("disconnect_points", "id", "to_id"), &AStar::disconnect_points);
 	ClassDB::bind_method(D_METHOD("are_points_connected", "id", "to_id"), &AStar::are_points_connected);
 
+	ClassDB::bind_method(D_METHOD("get_point_count"), &AStar::get_point_count);
+	ClassDB::bind_method(D_METHOD("get_point_capacity"), &AStar::get_point_capacity);
+	ClassDB::bind_method(D_METHOD("reserve_space", "num_nodes"), &AStar::reserve_space);
 	ClassDB::bind_method(D_METHOD("clear"), &AStar::clear);
 
 	ClassDB::bind_method(D_METHOD("get_closest_point", "to_position"), &AStar::get_closest_point);
@@ -605,10 +622,22 @@ bool AStar2D::are_points_connected(int p_id, int p_with_id) const {
 	return astar.are_points_connected(p_id, p_with_id);
 }
 
+int AStar2D::get_point_count() const {
+	return astar.get_point_count();
+}
+
+int AStar2D::get_point_capacity() const {
+	return astar.get_point_capacity();
+}
+
 void AStar2D::clear() {
 	astar.clear();
 }
 
+void AStar2D::reserve_space(int p_num_nodes) {
+	astar.reserve_space(p_num_nodes);
+}
+
 int AStar2D::get_closest_point(const Vector2 &p_point) const {
 	return astar.get_closest_point(Vector3(p_point.x, p_point.y, 0));
 }
@@ -659,6 +688,9 @@ void AStar2D::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("disconnect_points", "id", "to_id"), &AStar2D::disconnect_points);
 	ClassDB::bind_method(D_METHOD("are_points_connected", "id", "to_id"), &AStar2D::are_points_connected);
 
+	ClassDB::bind_method(D_METHOD("get_point_count"), &AStar2D::get_point_count);
+	ClassDB::bind_method(D_METHOD("get_point_capacity"), &AStar2D::get_point_capacity);
+	ClassDB::bind_method(D_METHOD("reserve_space", "num_nodes"), &AStar2D::reserve_space);
 	ClassDB::bind_method(D_METHOD("clear"), &AStar2D::clear);
 
 	ClassDB::bind_method(D_METHOD("get_closest_point", "to_position"), &AStar2D::get_closest_point);
diff --git a/core/math/a_star.h b/core/math/a_star.h
index 53aaaa1f6c..ec2a06f07f 100644
--- a/core/math/a_star.h
+++ b/core/math/a_star.h
@@ -46,6 +46,10 @@ class AStar : public Reference {
 
 	struct Point {
 
+		Point() :
+				neighbours(4u),
+				unlinked_neighbours(4u) {}
+
 		int id;
 		Vector3 pos;
 		real_t weight_scale;
@@ -132,6 +136,9 @@ public:
 	void disconnect_points(int p_id, int p_with_id);
 	bool are_points_connected(int p_id, int p_with_id) const;
 
+	int get_point_count() const;
+	int get_point_capacity() const;
+	void reserve_space(int p_num_nodes);
 	void clear();
 
 	int get_closest_point(const Vector3 &p_point) const;
@@ -171,6 +178,9 @@ public:
 	void disconnect_points(int p_id, int p_with_id);
 	bool are_points_connected(int p_id, int p_with_id) const;
 
+	int get_point_count() const;
+	int get_point_capacity() const;
+	void reserve_space(int p_num_nodes);
 	void clear();
 
 	int get_closest_point(const Vector2 &p_point) const;
diff --git a/core/oa_hash_map.h b/core/oa_hash_map.h
index 83621bec14..5ea6d8b0d4 100644
--- a/core/oa_hash_map.h
+++ b/core/oa_hash_map.h
@@ -78,8 +78,7 @@ private:
 		p_hash = p_hash & ~DELETED_HASH_BIT; // we don't care if it was deleted or not
 
 		uint32_t original_pos = p_hash % capacity;
-
-		return p_pos - original_pos;
+		return (p_pos - original_pos) % capacity;
 	}
 
 	_FORCE_INLINE_ void _construct(uint32_t p_pos, uint32_t p_hash, const TKey &p_key, const TValue &p_value) {
@@ -152,17 +151,16 @@ private:
 		}
 	}
 
-	void _resize_and_rehash() {
+	void _resize_and_rehash(uint32_t p_new_capacity) {
+
+		uint32_t old_capacity = capacity;
+		capacity = p_new_capacity;
 
 		TKey *old_keys = keys;
 		TValue *old_values = values;
 		uint32_t *old_hashes = hashes;
 
-		uint32_t old_capacity = capacity;
-
-		capacity = old_capacity * 2;
 		num_elements = 0;
-
 		keys = memnew_arr(TKey, capacity);
 		values = memnew_arr(TValue, capacity);
 		hashes = memnew_arr(uint32_t, capacity);
@@ -187,6 +185,10 @@ private:
 		memdelete_arr(old_hashes);
 	}
 
+	void _resize_and_rehash() {
+		_resize_and_rehash(capacity * 2);
+	}
+
 public:
 	_FORCE_INLINE_ uint32_t get_capacity() const { return capacity; }
 	_FORCE_INLINE_ uint32_t get_num_elements() const { return num_elements; }
@@ -199,11 +201,15 @@ public:
 
 		for (uint32_t i = 0; i < capacity; i++) {
 
+			if (hashes[i] == EMPTY_HASH) {
+				continue;
+			}
+
 			if (hashes[i] & DELETED_HASH_BIT) {
 				continue;
 			}
 
-			hashes[i] |= DELETED_HASH_BIT;
+			hashes[i] = EMPTY_HASH;
 			values[i].~TValue();
 			keys[i].~TKey();
 		}
@@ -272,6 +278,16 @@ public:
 		num_elements--;
 	}
 
+	/**
+	 * reserves space for a number of elements, useful to avoid many resizes and rehashes
+	 *  if adding a known (possibly large) number of elements at once, must be larger than old
+	 *  capacity.
+	 **/
+	void reserve(uint32_t p_new_capacity) {
+		ERR_FAIL_COND(p_new_capacity < capacity);
+		_resize_and_rehash(p_new_capacity);
+	}
+
 	struct Iterator {
 		bool valid;
 
@@ -336,7 +352,7 @@ public:
 		hashes = memnew_arr(uint32_t, p_initial_capacity);
 
 		for (uint32_t i = 0; i < p_initial_capacity; i++) {
-			hashes[i] = 0;
+			hashes[i] = EMPTY_HASH;
 		}
 	}