Improve support for directed graphs in AStar

1 files changed, 33 insertions, 20 deletions

diff --git a/core/math/a_star.cpp b/core/math/a_star.cpp
index 60b7326c29..482d7d8cd5 100644
--- a/core/math/a_star.cpp
+++ b/core/math/a_star.cpp
@@ -164,23 +164,21 @@ void AStar::connect_points(int p_id, int p_with_id, bool bidirectional) {
 	}
 
 	Segment s(p_id, p_with_id);
-	if (s.from == p_id) {
-		s.from_point = a;
-		s.to_point = b;
-	} else {
-		s.from_point = b;
-		s.to_point = a;
-	}
-
+	s.from_point = a;
+	s.to_point = b;
 	segments.insert(s);
+
+	if (bidirectional) {
+		SWAP(s.from, s.to);
+		SWAP(s.from_point, s.to_point);
+		segments.insert(s);
+	}
 }
 
-void AStar::disconnect_points(int p_id, int p_with_id) {
+void AStar::disconnect_points(int p_id, int p_with_id, bool bidirectional) {
 
 	Segment s(p_id, p_with_id);
-	ERR_FAIL_COND(!segments.has(s));
-
-	segments.erase(s);
+	Segment t(p_with_id, p_id);
 
 	Point *a;
 	bool a_exists = points.lookup(p_id, a);
@@ -190,10 +188,24 @@ void AStar::disconnect_points(int p_id, int p_with_id) {
 	bool b_exists = points.lookup(p_with_id, b);
 	CRASH_COND(!b_exists);
 
-	a->neighbours.remove(b->id);
-	a->unlinked_neighbours.remove(b->id);
-	b->neighbours.remove(a->id);
-	b->unlinked_neighbours.remove(a->id);
+	bool warned = false;
+
+	if (segments.has(s)) {
+		segments.erase(s);
+		a->neighbours.remove(b->id);
+		b->unlinked_neighbours.remove(a->id);
+	} else {
+		warned = true;
+		WARN_PRINT("The edge to be removed does not exist.");
+	}
+
+	if (bidirectional && segments.has(t)) {
+		segments.erase(t);
+		b->neighbours.remove(a->id);
+		a->unlinked_neighbours.remove(b->id);
+	} else if (bidirectional && !warned) {
+		WARN_PRINT("The reverse edge to be removed does not exist.");
+	}
 }
 
 bool AStar::has_point(int p_id) const {
@@ -227,10 +239,11 @@ PoolVector<int> AStar::get_point_connections(int p_id) {
 	return point_list;
 }
 
-bool AStar::are_points_connected(int p_id, int p_with_id) const {
+bool AStar::are_points_connected(int p_id, int p_with_id, bool bidirectional) const {
 
 	Segment s(p_id, p_with_id);
-	return segments.has(s);
+	Segment t(p_with_id, p_id);
+	return segments.has(s) || (bidirectional && segments.has(t));
 }
 
 void AStar::clear() {
@@ -532,8 +545,8 @@ void AStar::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("is_point_disabled", "id"), &AStar::is_point_disabled);
 
 	ClassDB::bind_method(D_METHOD("connect_points", "id", "to_id", "bidirectional"), &AStar::connect_points, DEFVAL(true));
-	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("disconnect_points", "id", "to_id", "bidirectional"), &AStar::disconnect_points, DEFVAL(true));
+	ClassDB::bind_method(D_METHOD("are_points_connected", "id", "to_id", "bidirectional"), &AStar::are_points_connected, DEFVAL(true));
 
 	ClassDB::bind_method(D_METHOD("get_point_count"), &AStar::get_point_count);
 	ClassDB::bind_method(D_METHOD("get_point_capacity"), &AStar::get_point_capacity);