From 7f61eb3a901cc0bdbda2fd0509e97cb38e69cf66 Mon Sep 17 00:00:00 2001
From: Nolan Darilek <nolan@thewordnerd.info>
Date: Tue, 27 Jul 2021 20:26:16 -0500
Subject: [PATCH] Even more cleanup.

---
 src/pathfinding.rs | 32 +++++++++++++++-----------------
 1 file changed, 15 insertions(+), 17 deletions(-)

diff --git a/src/pathfinding.rs b/src/pathfinding.rs
index 67a8e42..e4ba62a 100644
--- a/src/pathfinding.rs
+++ b/src/pathfinding.rs
@@ -226,31 +226,28 @@ fn calculate_path(
     }
 }
 
-fn poll_tasks(
-    mut commands: Commands,
-    mut query: Query<(Entity, &mut Calculating)>,
-    destinations: Query<&Destination>,
-) {
+fn poll_tasks(mut commands: Commands, mut query: Query<(Entity, &mut Calculating)>) {
     for (entity, mut calculating) in query.iter_mut() {
         println!("Got a calculating task for {:?}", entity);
-        if destinations.get(entity).is_ok() {
-            if let Some(result) = future::block_on(future::poll_once(&mut **calculating)) {
-                println!("Got a result, setting and returning");
-                if let Some(path) = result {
-                    commands.entity(entity).insert(path);
-                } else {
-                    commands.entity(entity).insert(NoPath);
-                }
-                println!("Removing calculation task");
-                commands.entity(entity).remove::<Calculating>();
+        if let Some(result) = future::block_on(future::poll_once(&mut **calculating)) {
+            println!("Got a result, setting and returning");
+            if let Some(path) = result {
+                commands.entity(entity).insert(path);
+            } else {
+                commands.entity(entity).insert(NoPath);
             }
-        } else {
-            println!("Destination was removed, clearing task");
+            println!("Removing calculation task");
             commands.entity(entity).remove::<Calculating>();
         }
     }
 }
 
+fn remove_destination(mut commands: Commands, removed: RemovedComponents<Destination>) {
+    for entity in removed.iter() {
+        commands.entity(entity).remove::<Calculating>();
+    }
+}
+
 fn negotiate_path(
     mut commands: Commands,
     mut query: Query<(
@@ -336,6 +333,7 @@ pub struct PathfindingPlugin;
 impl Plugin for PathfindingPlugin {
     fn build(&self, app: &mut AppBuilder) {
         app.add_system(calculate_path.system())
+            .add_system_to_stage(CoreStage::PostUpdate, remove_destination.system())
             .add_system(poll_tasks.system())
             .add_system(negotiate_path.system())
             .add_system_to_stage(CoreStage::PostUpdate, remove_calculating.system());