Handle removing folders and their path-to-ID associations better

server/Cargo.lock

@@ -641,7 +641,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 [[package]]
 name = "rbx-tree"
 version = "0.1.0"
-source = "git+https://github.com/LPGhatguy/rbx-tree.git#124175066fa31e114fe55cca627dad4cae571a63"
+source = "git+https://github.com/LPGhatguy/rbx-tree.git#258b458cffbcbd01e3f1cf5bbd701cc75260f01a"
 dependencies = [
  "serde 1.0.75 (registry+https://github.com/rust-lang/crates.io-index)",
  "serde_derive 1.0.75 (registry+https://github.com/rust-lang/crates.io-index)",

server/src/message_queue.rs

@@ -24,8 +24,8 @@ pub fn get_listener_id() -> ListenerId {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "type")]
 pub enum Message {
-    InstanceChanged {
+    InstancesRemoved {
-        id: RbxId,
+        ids: Vec<RbxId>,
     },
 }
 
@@ -50,29 +50,23 @@ impl MessageQueue {
             messages.extend_from_slice(new_messages);
         }
 
-        {
+        for listener in message_listeners.values() {
-            for listener in message_listeners.values() {
+            listener.send(()).unwrap();
-                listener.send(()).unwrap();
-            }
         }
     }
 
     pub fn subscribe(&self, sender: mpsc::Sender<()>) -> ListenerId {
         let id = get_listener_id();
 
-        {
+        let mut message_listeners = self.message_listeners.lock().unwrap();
-            let mut message_listeners = self.message_listeners.lock().unwrap();
+        message_listeners.insert(id, sender);
-            message_listeners.insert(id, sender);
-        }
 
         id
     }
 
     pub fn unsubscribe(&self, id: ListenerId) {
-        {
+        let mut message_listeners = self.message_listeners.lock().unwrap();
-            let mut message_listeners = self.message_listeners.lock().unwrap();
+        message_listeners.remove(&id);
-            message_listeners.remove(&id);
-        }
     }
 
     pub fn get_message_cursor(&self) -> u32 {

server/src/rbx_session.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::HashMap,
+    collections::{HashMap, HashSet},
     path::{Path, PathBuf},
     sync::{Arc, Mutex},
     str,
@@ -9,13 +9,83 @@ use rbx_tree::{RbxTree, RbxId, RbxInstance, RbxValue};
 
 use crate::{
     project::{Project, ProjectNode, InstanceProjectNode},
-    message_queue::MessageQueue,
+    message_queue::{Message, MessageQueue},
     imfs::{Imfs, ImfsItem},
 };
 
+#[derive(Debug)]
+struct PathIdNode {
+    id: RbxId,
+    children: HashSet<PathBuf>,
+}
+
+/// A map from paths to instance IDs, with a bit of additional data that enables
+/// removing a path and all of its child paths from the tree in constant time.
+#[derive(Debug)]
+struct PathIdTree {
+    nodes: HashMap<PathBuf, PathIdNode>,
+}
+
+impl PathIdTree {
+    pub fn new() -> PathIdTree {
+        PathIdTree {
+            nodes: HashMap::new(),
+        }
+    }
+
+    pub fn insert(&mut self, path: &Path, id: RbxId) {
+        if let Some(parent_path) = path.parent() {
+            if let Some(mut parent) = self.nodes.get_mut(parent_path) {
+                parent.children.insert(path.to_path_buf());
+            }
+        }
+
+        self.nodes.insert(path.to_path_buf(), PathIdNode {
+            id,
+            children: HashSet::new(),
+        });
+    }
+
+    pub fn remove(&mut self, root_path: &Path) -> Option<RbxId> {
+        if let Some(parent_path) = root_path.parent() {
+            if let Some(mut parent) = self.nodes.get_mut(parent_path) {
+                parent.children.remove(root_path);
+            }
+        }
+
+        let mut root_node = match self.nodes.remove(root_path) {
+            Some(node) => node,
+            None => return None,
+        };
+
+        let root_id = root_node.id;
+        let mut to_visit: Vec<PathBuf> = root_node.children.drain().collect();
+
+        loop {
+            let next_path = match to_visit.pop() {
+                Some(path) => path,
+                None => break,
+            };
+
+            match self.nodes.remove(&next_path) {
+                Some(mut node) => {
+                    for child in node.children.drain() {
+                        to_visit.push(child);
+                    }
+                },
+                None => {
+                    warn!("Consistency issue; tried to remove {} but it was already removed", next_path.display());
+                },
+            }
+        }
+
+        Some(root_id)
+    }
+}
+
 pub struct RbxSession {
     tree: RbxTree,
-    paths_to_ids: HashMap<PathBuf, RbxId>,
+    path_id_tree: PathIdTree,
     ids_to_project_paths: HashMap<RbxId, String>,
     message_queue: Arc<MessageQueue>,
     imfs: Arc<Mutex<Imfs>>,
@@ -24,14 +94,14 @@ pub struct RbxSession {
 
 impl RbxSession {
     pub fn new(project: Arc<Project>, imfs: Arc<Mutex<Imfs>>, message_queue: Arc<MessageQueue>) -> RbxSession {
-        let (tree, paths_to_ids, ids_to_project_paths) = {
+        let (tree, path_id_tree, ids_to_project_paths) = {
             let temp_imfs = imfs.lock().unwrap();
             construct_initial_tree(&project, &temp_imfs)
         };
 
         RbxSession {
             tree,
-            paths_to_ids,
+            path_id_tree,
             ids_to_project_paths,
             message_queue,
             imfs,
@@ -49,10 +119,33 @@ impl RbxSession {
 
     pub fn path_removed(&mut self, path: &Path) {
         info!("Path removed: {}", path.display());
+
+        let instance_id = match self.path_id_tree.remove(path) {
+            Some(id) => id,
+            None => return,
+        };
+
+        let removed_subtree = match self.tree.remove_instance(instance_id) {
+            Some(tree) => tree,
+            None => {
+                warn!("Rojo tried to remove an instance that was half cleaned-up. This is probably a bug in Rojo.");
+                return;
+            },
+        };
+
+        let removed_ids: Vec<RbxId> = removed_subtree.iter_all_ids().collect();
+
+        self.message_queue.push_messages(&[
+            Message::InstancesRemoved {
+                ids: removed_ids,
+            },
+        ]);
     }
 
     pub fn path_renamed(&mut self, from_path: &Path, to_path: &Path) {
         info!("Path renamed from {} to {}", from_path.display(), to_path.display());
+        self.path_removed(from_path);
+        self.path_created(to_path);
     }
 
     pub fn get_tree(&self) -> &RbxTree {
@@ -64,11 +157,18 @@ impl RbxSession {
     }
 }
 
+struct ConstructContext<'a> {
+    tree: RbxTree,
+    imfs: &'a Imfs,
+    path_id_tree: PathIdTree,
+    ids_to_project_paths: HashMap<RbxId, String>,
+}
+
 fn construct_initial_tree(
     project: &Project,
     imfs: &Imfs,
-) -> (RbxTree, HashMap<PathBuf, RbxId>, HashMap<RbxId, String>) {
+) -> (RbxTree, PathIdTree, HashMap<RbxId, String>) {
-    let paths_to_ids = HashMap::new();
+    let path_id_tree = PathIdTree::new();
     let ids_to_project_paths = HashMap::new();
     let tree = RbxTree::new(RbxInstance {
         name: "this isn't supposed to be here".to_string(),
@@ -81,7 +181,7 @@ fn construct_initial_tree(
     let mut context = ConstructContext {
         tree,
         imfs,
-        paths_to_ids,
+        path_id_tree,
         ids_to_project_paths,
     };
 
@@ -93,14 +193,7 @@ fn construct_initial_tree(
         &project.tree,
     );
 
-    (context.tree, context.paths_to_ids, context.ids_to_project_paths)
+    (context.tree, context.path_id_tree, context.ids_to_project_paths)
-}
-
-struct ConstructContext<'a> {
-    tree: RbxTree,
-    imfs: &'a Imfs,
-    paths_to_ids: HashMap<PathBuf, RbxId>,
-    ids_to_project_paths: HashMap<RbxId, String>,
 }
 
 fn construct_project_node(
@@ -165,7 +258,7 @@ fn construct_sync_point_node(
             };
 
             let id = context.tree.insert_instance(instance, parent_instance_id);
-            context.paths_to_ids.insert(file.path.clone(), id);
+            context.path_id_tree.insert(&file.path, id);
 
             id
         },
@@ -177,7 +270,7 @@ fn construct_sync_point_node(
             };
 
             let id = context.tree.insert_instance(instance, parent_instance_id);
-            context.paths_to_ids.insert(directory.path.clone(), id);
+            context.path_id_tree.insert(&directory.path, id);
 
             for child_path in &directory.children {
                 let child_instance_name = child_path.file_name().unwrap().to_str().unwrap();