rojo/src/snapshot/patch_apply.rs

//! Defines the algorithm for applying generated patches.

use std::{
    collections::{HashMap, HashSet},
    mem::take,
};

use rbx_dom_weak::{
    types::{Ref, Variant},
    ustr, Ustr,
};

use super::{
    patch::{AppliedPatchSet, AppliedPatchUpdate, PatchSet, PatchUpdate},
    InstanceSnapshot, RojoTree,
};
use crate::{
    multimap::MultiMap, RojoRef, REF_ID_ATTRIBUTE_NAME, REF_PATH_ATTRIBUTE_PREFIX,
    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
/// tree in sync with Rojo's.
#[profiling::function]
pub fn apply_patch_set(tree: &mut RojoTree, patch_set: PatchSet) -> AppliedPatchSet {
    let mut context = PatchApplyContext::default();

    {
        profiling::scope!("removals");
        for removed_id in patch_set.removed_instances {
            apply_remove_instance(&mut context, tree, removed_id);
        }
    }

    {
        profiling::scope!("additions");
        for add_patch in patch_set.added_instances {
            apply_add_child(&mut context, tree, add_patch.parent_id, add_patch.instance);
        }
    }

    {
        profiling::scope!("updates");
        // Updates need to be applied after additions, which reduces the complexity
        // of updates significantly.
        for update_patch in patch_set.updated_instances {
            apply_update_child(&mut context, tree, update_patch);
        }
    }

    finalize_patch_application(context, tree)
}

/// All of the ephemeral state needing during application of a patch.
#[derive(Default)]
struct PatchApplyContext {
    /// A map from transient snapshot IDs (generated by snapshot middleware) to
    /// instance IDs in the actual tree. These are both the same data type so
    /// that they fit into the same `Variant::Ref` type.
    ///
    /// At this point in the patch process, IDs in instance properties have been
    /// partially translated from 'snapshot space' into 'tree space' by the
    /// patch computation process. An ID not existing in this map means either:
    ///
    /// 1. The ID is already in tree space and refers to an instance that
    ///    existed in the tree before this patch was applied.
    ///
    /// 2. The ID if in snapshot space, but points to an instance that was not
    ///    part of the snapshot that was put through the patch computation
    ///    function.
    ///
    /// #2 should not occur in well-formed projects, but is indistinguishable
    /// from #1 right now. It could happen if two model files try to reference
    /// eachother.
    snapshot_id_to_instance_id: HashMap<Ref, Ref>,

    /// Tracks all of the instances added by this patch that have refs that need
    /// to be rewritten.
    has_refs_to_rewrite: HashSet<Ref>,

    /// Legacy: Tracks ref properties using the old ID-based system (Rojo_Target_).
    /// Kept for backwards compatibility with existing projects.
    attribute_refs_to_rewrite: MultiMap<Ref, (Ustr, String)>,

    /// Tracks path-based ref properties (Rojo_Ref_). This is the preferred system
    /// that stores paths to target instances.
    path_refs_to_rewrite: MultiMap<Ref, (Ustr, String)>,

    /// Tracks instances that have Rojo ref attributes that need to be cleaned up
    /// after resolution. The attributes are only used as transport in the file
    /// format and should not persist in the live game.
    instances_needing_attr_cleanup: HashSet<Ref>,

    /// The current applied patch result, describing changes made to the tree.
    applied_patch_set: AppliedPatchSet,
}

/// Finalize this patch application, consuming the context, applying any
/// deferred property updates, and returning the finally applied patch set.
///
/// Ref properties from snapshots refer to eachother via snapshot ID. Some of
/// these properties are transformed when the patch is computed, notably the
/// instances that the patch computing method is able to pair up.
///
/// The remaining Ref properties need to be handled during patch application,
/// where we build up a map of snapshot IDs to instance IDs as they're created,
/// then apply properties all at once at the end.
#[profiling::function]
fn finalize_patch_application(context: PatchApplyContext, tree: &mut RojoTree) -> AppliedPatchSet {
    for id in context.has_refs_to_rewrite {
        // This should always succeed since instances marked as added in our
        // patch should be added without fail.
        let mut instance = tree
            .get_instance_mut(id)
            .expect("Invalid instance ID in deferred property map");

        for value in instance.properties_mut().values_mut() {
            if let Variant::Ref(referent) = value {
                if let Some(&instance_referent) = context.snapshot_id_to_instance_id.get(referent) {
                    *value = Variant::Ref(instance_referent);
                }
            }
        }
    }

    // IMPORTANT: Rojo_Target_* (legacy) must be resolved BEFORE Rojo_Ref_*
    // (path-based) so that path-based results overwrite legacy results when
    // both exist for the same property. Do not reorder these two blocks.
    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(..));
    }

    // Handle path-based references (Rojo_Ref_)
    let mut path_rewrites = Vec::new();
    for (id, map) in context.path_refs_to_rewrite {
        for (prop_name, path) in map {
            if let Some(target) = tree.resolve_ref_path(&path, id) {
                log::debug!(
                    "Resolved path reference {} -> {} (path: {})",
                    prop_name,
                    target,
                    path
                );
                path_rewrites.push((prop_name, Variant::Ref(target)));
            } else {
                log::debug!(
                    "Could not resolve path reference {} at path '{}' - instance not found in tree",
                    prop_name,
                    path
                );
            }
        }
        if !path_rewrites.is_empty() {
            let mut instance = tree
                .get_instance_mut(id)
                .expect("Invalid instance ID in deferred path ref map");
            instance.properties_mut().extend(path_rewrites.drain(..));
        }
    }

    // Clean up Rojo ref attributes - they're only used for transport in the file
    // format and should not persist in the live game
    for id in context.instances_needing_attr_cleanup {
        let mut instance = match tree.get_instance_mut(id) {
            Some(inst) => inst,
            None => continue,
        };

        let attributes = match instance.properties_mut().get_mut(&ustr("Attributes")) {
            Some(Variant::Attributes(attrs)) => attrs,
            _ => continue,
        };

        // Remove all Rojo ref-related attributes
        let keys_to_remove: Vec<String> = attributes
            .iter()
            .filter_map(|(name, _)| {
                if name == REF_ID_ATTRIBUTE_NAME
                    || name.starts_with(REF_POINTER_ATTRIBUTE_PREFIX)
                    || name.starts_with(REF_PATH_ATTRIBUTE_PREFIX)
                {
                    Some(name.to_string())
                } else {
                    None
                }
            })
            .collect();

        for key in keys_to_remove {
            attributes.remove(key);
        }
    }

    context.applied_patch_set
}

fn apply_remove_instance(context: &mut PatchApplyContext, tree: &mut RojoTree, removed_id: Ref) {
    tree.remove(removed_id);
    context.applied_patch_set.removed.push(removed_id);
}

fn apply_add_child(
    context: &mut PatchApplyContext,
    tree: &mut RojoTree,
    parent_id: Ref,
    mut snapshot: InstanceSnapshot,
) {
    let snapshot_id = snapshot.snapshot_id;
    let children = take(&mut snapshot.children);

    // If an object we're adding has a non-null referent, we'll note this
    // instance down as needing to be revisited later.
    let has_refs = snapshot.properties.values().any(|value| match value {
        Variant::Ref(value) => value.is_some(),
        _ => false,
    });

    let id = tree.insert_instance(parent_id, snapshot);
    context.applied_patch_set.added.push(id);

    if has_refs {
        context.has_refs_to_rewrite.insert(id);
    }

    if snapshot_id.is_some() {
        context.snapshot_id_to_instance_id.insert(snapshot_id, id);
    }

    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) {
    let mut applied_patch = AppliedPatchUpdate::new(patch.id);

    if let Some(metadata) = patch.changed_metadata {
        tree.update_metadata(patch.id, metadata.clone());
        applied_patch.changed_metadata = Some(metadata);
    }

    // Capture old class name before the mutable borrow so we can update
    // the script index afterward if the class changes.
    let old_class_name = patch.changed_class_name.and_then(|_| {
        tree.get_instance(patch.id)
            .map(|inst| inst.class_name().to_string())
    });

    let mut instance = match tree.get_instance_mut(patch.id) {
        Some(instance) => instance,
        None => {
            log::warn!(
                "Patch misapplication: Instance {:?}, referred to by update patch, did not exist.",
                patch.id
            );
            return;
        }
    };

    if let Some(name) = patch.changed_name {
        *instance.name_mut() = name.clone();
        applied_patch.changed_name = Some(name);
    }

    if let Some(class_name) = patch.changed_class_name {
        instance.set_class_name(class_name);
        applied_patch.changed_class_name = Some(class_name);
    }

    for (key, property_entry) in patch.changed_properties {
        match property_entry {
            // Ref values need to be potentially rewritten from snapshot IDs to
            // instance IDs if they referred to an instance that was created as
            // part of this patch.
            Some(Variant::Ref(referent)) => {
                if referent.is_none() {
                    continue;
                }

                // If our ID is not found in this map, then it either refers to
                // an existing instance NOT added by this patch, or there was an
                // error. See `PatchApplyContext::snapshot_id_to_instance_id`
                // for more info.
                let new_referent = context
                    .snapshot_id_to_instance_id
                    .get(&referent)
                    .copied()
                    .unwrap_or(referent);

                instance
                    .properties_mut()
                    .insert(key, Variant::Ref(new_referent));
            }
            Some(ref value) => {
                instance.properties_mut().insert(key, value.clone());
            }
            None => {
                instance.properties_mut().remove(&key);
            }
        }

        applied_patch.changed_properties.insert(key, property_entry);
    }

    if let (Some(old_class), Some(new_class)) = (&old_class_name, &applied_patch.changed_class_name)
    {
        tree.update_script_tracking(patch.id, old_class, new_class.as_str());
    }

    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(&ustr("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 Some(id_str) = crate::variant_as_str(attr_value, attr_name) {
                attr_id = Some(RojoRef::new(id_str.to_string()));
            }
        }
        if let Some(prop_name) = attr_name.strip_prefix(REF_POINTER_ATTRIBUTE_PREFIX) {
            context.instances_needing_attr_cleanup.insert(id);
            if let Some(id_str) = crate::variant_as_str(attr_value, attr_name) {
                context
                    .attribute_refs_to_rewrite
                    .insert(id, (ustr(prop_name), id_str.to_string()));
            }
        }
        // Handle path-based references (Rojo_Ref_)
        if let Some(prop_name) = attr_name.strip_prefix(REF_PATH_ATTRIBUTE_PREFIX) {
            context.instances_needing_attr_cleanup.insert(id);
            if let Some(path) = crate::variant_as_str(attr_value, attr_name) {
                context
                    .path_refs_to_rewrite
                    .insert(id, (ustr(prop_name), path.to_string()));
            }
        }
    }
    if let Some(specified_id) = attr_id {
        // Legacy Rojo_Id also needs cleanup
        context.instances_needing_attr_cleanup.insert(id);
        tree.set_specified_id(id, specified_id);
    }
}

#[cfg(test)]
mod test {
    use super::*;

    use std::borrow::Cow;

    use rbx_dom_weak::{types::Variant, UstrMap};

    use super::super::PatchAdd;

    #[test]
    fn add_from_empty() {
        let _ = tracing_subscriber::fmt::try_init();

        let mut tree = RojoTree::new(InstanceSnapshot::new());

        let root_id = tree.get_root_id();

        let snapshot = InstanceSnapshot {
            snapshot_id: Ref::none(),
            metadata: Default::default(),
            name: Cow::Borrowed("Foo"),
            class_name: ustr("Bar"),
            properties: UstrMap::from_iter([(ustr("Baz"), Variant::Int32(5))]),
            children: Vec::new(),
        };

        let patch_set = PatchSet {
            added_instances: vec![PatchAdd {
                parent_id: root_id,
                instance: snapshot.clone(),
            }],
            stage_ids: HashSet::new(),
            ..Default::default()
        };

        apply_patch_set(&mut tree, patch_set);

        let root_instance = tree.get_instance(root_id).unwrap();
        let child_id = root_instance.children()[0];
        let child_instance = tree.get_instance(child_id).unwrap();

        assert_eq!(child_instance.name(), &snapshot.name);
        assert_eq!(child_instance.class_name(), snapshot.class_name);
        assert_eq!(child_instance.properties(), &snapshot.properties);
        assert!(child_instance.children().is_empty());
    }

    #[test]
    fn update_existing() {
        let _ = tracing_subscriber::fmt::try_init();

        let mut tree = RojoTree::new(
            InstanceSnapshot::new()
                .class_name("OldClassName")
                .name("OldName")
                .property("Foo", 7i32)
                .property("Bar", 3i32)
                .property("Unchanged", -5i32),
        );

        let root_id = tree.get_root_id();

        let patch = PatchUpdate {
            id: root_id,
            changed_name: Some("Foo".to_owned()),
            changed_class_name: Some(ustr("NewClassName")),
            changed_properties: UstrMap::from_iter([
                // The value of Foo has changed
                (ustr("Foo"), Some(Variant::Int32(8))),
                // Bar has been deleted
                (ustr("Bar"), None),
                // Baz has been added
                (ustr("Baz"), Some(Variant::Int32(10))),
            ]),
            changed_metadata: None,
        };

        let patch_set = PatchSet {
            updated_instances: vec![patch],
            stage_ids: HashSet::new(),
            ..Default::default()
        };

        apply_patch_set(&mut tree, patch_set);

        let expected_properties = UstrMap::from_iter([
            (ustr("Foo"), Variant::Int32(8)),
            (ustr("Baz"), Variant::Int32(10)),
            (ustr("Unchanged"), Variant::Int32(-5)),
        ]);

        let root_instance = tree.get_instance(root_id).unwrap();
        assert_eq!(root_instance.name(), "Foo");
        assert_eq!(root_instance.class_name(), "NewClassName");
        assert_eq!(root_instance.properties(), &expected_properties);
    }
}