Support setting referent properties via attributes (#843)

Co-authored-by: Kenneth Loeffler <kenloef@gmail.com>

src/snapshot/tree.rs

@@ -8,7 +8,7 @@ use rbx_dom_weak::{
     Instance, InstanceBuilder, WeakDom,
 };
 
-use crate::multimap::MultiMap;
+use crate::{multimap::MultiMap, RojoRef};
 
 use super::{InstanceMetadata, InstanceSnapshot};
 
@@ -33,6 +33,12 @@ pub struct RojoTree {
     /// appearing multiple times in the same Rojo project. This is sometimes
     /// called "path aliasing" in various Rojo documentation.
     path_to_ids: MultiMap<PathBuf, Ref>,
+
+    /// A map of specified RojoRefs to underlying Refs they represent.
+    /// This field is a MultiMap to allow for the possibility of the user specifying
+    /// the same RojoRef for multiple different instances. An entry containing
+    /// multiple elements is an error condition that should be raised to the user.
+    specified_id_to_refs: MultiMap<RojoRef, Ref>,
 }
 
 impl RojoTree {
@@ -45,6 +51,7 @@ impl RojoTree {
             inner: WeakDom::new(root_builder),
             metadata_map: HashMap::new(),
             path_to_ids: MultiMap::new(),
+            specified_id_to_refs: MultiMap::new(),
         };
 
         let root_ref = tree.inner.root_ref();
@@ -137,6 +144,20 @@ impl RojoTree {
                         self.path_to_ids.insert(new_path.clone(), id);
                     }
                 }
+                if existing_metadata.specified_id != metadata.specified_id {
+                    // We need to uphold the invariant that each ID can only map
+                    // to one referent.
+                    if let Some(new) = &metadata.specified_id {
+                        if self.specified_id_to_refs.get(new).len() > 0 {
+                            log::error!("Duplicate user-specified referent '{new}'");
+                        }
+
+                        self.specified_id_to_refs.insert(new.clone(), id);
+                    }
+                    if let Some(old) = &existing_metadata.specified_id {
+                        self.specified_id_to_refs.remove(old, id);
+                    }
+                }
 
                 entry.insert(metadata);
             }
@@ -161,11 +182,37 @@ impl RojoTree {
         self.metadata_map.get(&id)
     }
 
+    /// Get the backing Ref of the given RojoRef. If the RojoRef maps to exactly
+    /// one Ref, this method returns Some. Otherwise, it returns None.
+    pub fn get_specified_id(&self, specified: &RojoRef) -> Option<Ref> {
+        match self.specified_id_to_refs.get(specified)[..] {
+            [referent] => Some(referent),
+            _ => None,
+        }
+    }
+
+    pub fn set_specified_id(&mut self, id: Ref, specified: RojoRef) {
+        if let Some(metadata) = self.metadata_map.get_mut(&id) {
+            if let Some(old) = metadata.specified_id.replace(specified.clone()) {
+                self.specified_id_to_refs.remove(&old, id);
+            }
+        }
+        self.specified_id_to_refs.insert(specified, id);
+    }
+
     fn insert_metadata(&mut self, id: Ref, metadata: InstanceMetadata) {
         for path in &metadata.relevant_paths {
             self.path_to_ids.insert(path.clone(), id);
         }
 
+        if let Some(specified_id) = &metadata.specified_id {
+            if self.specified_id_to_refs.get(specified_id).len() > 0 {
+                log::error!("Duplicate user-specified referent '{specified_id}'");
+            }
+
+            self.set_specified_id(id, specified_id.clone());
+        }
+
         self.metadata_map.insert(id, metadata);
     }
 
@@ -174,6 +221,10 @@ impl RojoTree {
     fn remove_metadata(&mut self, id: Ref) {
         let metadata = self.metadata_map.remove(&id).unwrap();
 
+        if let Some(specified) = metadata.specified_id {
+            self.specified_id_to_refs.remove(&specified, id);
+        }
+
         for path in &metadata.relevant_paths {
             self.path_to_ids.remove(path, id);
         }
@@ -297,3 +348,30 @@ impl InstanceWithMetaMut<'_> {
         self.metadata
     }
 }
+
+#[cfg(test)]
+mod test {
+    use crate::{
+        snapshot::{InstanceMetadata, InstanceSnapshot},
+        RojoRef,
+    };
+
+    use super::RojoTree;
+
+    #[test]
+    fn swap_duped_specified_ids() {
+        let custom_ref = RojoRef::new("MyCoolRef".into());
+        let snapshot = InstanceSnapshot::new()
+            .metadata(InstanceMetadata::new().specified_id(Some(custom_ref.clone())));
+        let mut tree = RojoTree::new(InstanceSnapshot::new());
+
+        let original = tree.insert_instance(tree.get_root_id(), snapshot.clone());
+        assert_eq!(tree.get_specified_id(&custom_ref.clone()), Some(original));
+
+        let duped = tree.insert_instance(tree.get_root_id(), snapshot.clone());
+        assert_eq!(tree.get_specified_id(&custom_ref.clone()), None);
+
+        tree.remove(original);
+        assert_eq!(tree.get_specified_id(&custom_ref.clone()), Some(duped));
+    }
+}