Support setting referent properties via attributes (#843)

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

src/snapshot/metadata.rs

@@ -12,6 +12,7 @@ use crate::{
     path_serializer,
     project::ProjectNode,
     snapshot_middleware::{emit_legacy_scripts_default, Middleware},
+    RojoRef,
 };
 
 /// Rojo-specific metadata that can be associated with an instance or a snapshot
@@ -58,6 +59,9 @@ pub struct InstanceMetadata {
     /// that instance's instigating source is snapshotted directly, the same
     /// context will be passed into it.
     pub context: InstanceContext,
+
+    /// Indicates the ID used for Ref properties pointing to this Instance.
+    pub specified_id: Option<RojoRef>,
 }
 
 impl InstanceMetadata {
@@ -67,6 +71,7 @@ impl InstanceMetadata {
             instigating_source: None,
             relevant_paths: Vec::new(),
             context: InstanceContext::default(),
+            specified_id: None,
         }
     }
 
@@ -97,6 +102,13 @@ impl InstanceMetadata {
             ..self
         }
     }
+
+    pub fn specified_id(self, id: Option<RojoRef>) -> Self {
+        Self {
+            specified_id: id,
+            ..self
+        }
+    }
 }
 
 impl Default for InstanceMetadata {

src/snapshot/patch_apply.rs

@@ -11,6 +11,7 @@ use super::{
     patch::{AppliedPatchSet, AppliedPatchUpdate, PatchSet, PatchUpdate},
     InstanceSnapshot, RojoTree,
 };
+use crate::{multimap::MultiMap, RojoRef, REF_ID_ATTRIBUTE_NAME, REF_POINTER_ATTRIBUTE_PREFIX};
 
 /// Consumes the input `PatchSet`, applying all of its prescribed changes to the
 /// tree and returns an `AppliedPatchSet`, which can be used to keep another
@@ -72,6 +73,11 @@ struct PatchApplyContext {
     /// to be rewritten.
     has_refs_to_rewrite: HashSet<Ref>,
 
+    /// Tracks all ref properties that were specified using attributes. This has
+    /// to be handled after everything else is done just like normal referent
+    /// properties.
+    attribute_refs_to_rewrite: MultiMap<Ref, (String, String)>,
+
     /// The current applied patch result, describing changes made to the tree.
     applied_patch_set: AppliedPatchSet,
 }
@@ -104,6 +110,22 @@ fn finalize_patch_application(context: PatchApplyContext, tree: &mut RojoTree) -
         }
     }
 
+    // This is to get around the fact that `RojoTre::get_specified_id` borrows
+    // the tree as immutable, but we need to hold a mutable reference to it.
+    // Not exactly elegant, but it does the job.
+    let mut real_rewrites = Vec::new();
+    for (id, map) in context.attribute_refs_to_rewrite {
+        for (prop_name, prop_value) in map {
+            if let Some(target) = tree.get_specified_id(&RojoRef::new(prop_value)) {
+                real_rewrites.push((prop_name, Variant::Ref(target)))
+            }
+        }
+        let mut instance = tree
+            .get_instance_mut(id)
+            .expect("Invalid instance ID in deferred attribute ref map");
+        instance.properties_mut().extend(real_rewrites.drain(..));
+    }
+
     context.applied_patch_set
 }
 
@@ -142,6 +164,8 @@ fn apply_add_child(
     for child in children {
         apply_add_child(context, tree, id, child);
     }
+
+    defer_ref_properties(tree, id, context);
 }
 
 fn apply_update_child(context: &mut PatchApplyContext, tree: &mut RojoTree, patch: PatchUpdate) {
@@ -208,9 +232,72 @@ fn apply_update_child(context: &mut PatchApplyContext, tree: &mut RojoTree, patc
         applied_patch.changed_properties.insert(key, property_entry);
     }
 
+    defer_ref_properties(tree, patch.id, context);
+
     context.applied_patch_set.updated.push(applied_patch)
 }
 
+/// Calculates manually-specified Ref properties and marks them in the provided
+/// `PatchApplyContext` to be rewritten at the end of the patch application
+/// process.
+///
+/// Currently, this only uses attributes but it can easily handle rewriting
+/// referents in other ways too!
+fn defer_ref_properties(tree: &mut RojoTree, id: Ref, context: &mut PatchApplyContext) {
+    let instance = tree
+        .get_instance(id)
+        .expect("Instances should exist when calculating deferred refs");
+    let attributes = match instance.properties().get("Attributes") {
+        Some(Variant::Attributes(attrs)) => attrs,
+        _ => return,
+    };
+
+    let mut attr_id = None;
+    for (attr_name, attr_value) in attributes.iter() {
+        if attr_name == REF_ID_ATTRIBUTE_NAME {
+            if let Variant::String(specified_id) = attr_value {
+                attr_id = Some(RojoRef::new(specified_id.clone()));
+            } else if let Variant::BinaryString(specified_id) = attr_value {
+                if let Ok(str) = std::str::from_utf8(specified_id.as_ref()) {
+                    attr_id = Some(RojoRef::new(str.to_string()))
+                } else {
+                    log::error!("Specified IDs must be valid UTF-8 strings.")
+                }
+            } else {
+                log::warn!(
+                    "Attribute {attr_name} is of type {:?} when it was \
+                    expected to be a String",
+                    attr_value.ty()
+                )
+            }
+        }
+        if let Some(prop_name) = attr_name.strip_prefix(REF_POINTER_ATTRIBUTE_PREFIX) {
+            if let Variant::String(prop_value) = attr_value {
+                context
+                    .attribute_refs_to_rewrite
+                    .insert(id, (prop_name.to_owned(), prop_value.clone()));
+            } else if let Variant::BinaryString(prop_value) = attr_value {
+                if let Ok(str) = std::str::from_utf8(prop_value.as_ref()) {
+                    context
+                        .attribute_refs_to_rewrite
+                        .insert(id, (prop_name.to_owned(), str.to_string()));
+                } else {
+                    log::error!("IDs specified by referent property attributes must be valid UTF-8 strings.")
+                }
+            } else {
+                log::warn!(
+                    "Attribute {attr_name} is of type {:?} when it was \
+                    expected to be a String",
+                    attr_value.ty()
+                )
+            }
+        }
+    }
+    if let Some(specified_id) = attr_id {
+        tree.set_specified_id(id, specified_id);
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

src/snapshot/tests/snapshots/librojo__snapshot__tests__apply__add_property.snap

@@ -13,5 +13,6 @@ metadata:
   relevant_paths: []
   context:
     emit_legacy_scripts: true
+  specified_id: ~
 children: []
 

src/snapshot/tests/snapshots/librojo__snapshot__tests__apply__remove_property_after_patch.snap

@@ -11,5 +11,6 @@ metadata:
   relevant_paths: []
   context:
     emit_legacy_scripts: true
+  specified_id: ~
 children: []
 

src/snapshot/tests/snapshots/librojo__snapshot__tests__apply__remove_property_initial.snap

@@ -13,5 +13,6 @@ metadata:
   relevant_paths: []
   context:
     emit_legacy_scripts: true
+  specified_id: ~
 children: []
 

src/snapshot/tests/snapshots/librojo__snapshot__tests__apply__set_name_and_class_name.snap

@@ -11,5 +11,6 @@ metadata:
   relevant_paths: []
   context:
     emit_legacy_scripts: true
+  specified_id: ~
 children: []
 

src/snapshot/tests/snapshots/librojo__snapshot__tests__compute__add_child.snap

@@ -12,6 +12,7 @@ added_instances:
         relevant_paths: []
         context:
           emit_legacy_scripts: true
+        specified_id: ~
       name: New
       class_name: Folder
       properties: {}

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));
+    }
+}