rojo/src/snapshot/patch_compute.rs

//! Defines the algorithm for computing a roughly-minimal patch set given an
//! existing instance tree and an instance snapshot.

use std::{collections::HashMap, mem::take, sync::Arc};

use rbx_dom_weak::{
    types::{Ref, Variant},
    ustr, HashMapExt as _, UstrMap, UstrSet,
};

use crate::{
    variant_eq::variant_eq, RojoRef, REF_ID_ATTRIBUTE_NAME, REF_PATH_ATTRIBUTE_PREFIX,
    REF_POINTER_ATTRIBUTE_PREFIX,
};

use super::{
    patch::{PatchAdd, PatchSet, PatchUpdate},
    InstanceSnapshot, InstanceWithMeta, InstigatingSource, PathIgnoreRule, RojoTree,
};

#[profiling::function]
pub fn compute_patch_set(snapshot: Option<InstanceSnapshot>, tree: &RojoTree, id: Ref) -> PatchSet {
    use super::matching::MatchingSession;

    let mut patch_set = PatchSet::new();

    if let Some(snapshot) = snapshot {
        let mut context = ComputePatchContext::default();
        let session = MatchingSession::new();

        compute_patch_set_internal(&mut context, snapshot, tree, id, &mut patch_set, &session);

        // Rewrite Ref properties to refer to instance IDs instead of snapshot IDs
        // for all of the IDs that we know about so far.
        rewrite_refs_in_updates(&context, &mut patch_set.updated_instances);
        rewrite_refs_in_additions(&context, &mut patch_set.added_instances);
    } else if id != tree.get_root_id() {
        patch_set.removed_instances.push(id);
    }

    patch_set
}

#[derive(Default)]
struct ComputePatchContext {
    snapshot_id_to_instance_id: HashMap<Ref, Ref>,
    /// Maps `Rojo_Id` attribute values (from snapshots being patched) to the
    /// tree Ref of the instance that carries them.  Populated top-down so that
    /// parent IDs are available when children try to resolve `Rojo_Target_*`.
    snapshot_specified_ids: HashMap<String, Ref>,
}

fn rewrite_refs_in_updates(context: &ComputePatchContext, updates: &mut [PatchUpdate]) {
    for update in updates {
        for property_value in update.changed_properties.values_mut() {
            if let Some(Variant::Ref(referent)) = property_value {
                if let Some(&instance_ref) = context.snapshot_id_to_instance_id.get(referent) {
                    *property_value = Some(Variant::Ref(instance_ref));
                }
            }
        }
    }
}

fn rewrite_refs_in_additions(context: &ComputePatchContext, additions: &mut [PatchAdd]) {
    for addition in additions {
        rewrite_refs_in_snapshot(context, &mut addition.instance);
    }
}

fn rewrite_refs_in_snapshot(context: &ComputePatchContext, snapshot: &mut InstanceSnapshot) {
    for property_value in snapshot.properties.values_mut() {
        if let Variant::Ref(referent) = property_value {
            if let Some(&instance_referent) = context.snapshot_id_to_instance_id.get(referent) {
                *property_value = Variant::Ref(instance_referent);
            }
        }
    }

    for child in &mut snapshot.children {
        rewrite_refs_in_snapshot(context, child);
    }
}

fn compute_patch_set_internal(
    context: &mut ComputePatchContext,
    mut snapshot: InstanceSnapshot,
    tree: &RojoTree,
    id: Ref,
    patch_set: &mut PatchSet,
    session: &super::matching::MatchingSession,
) {
    if snapshot.snapshot_id.is_some() {
        context
            .snapshot_id_to_instance_id
            .insert(snapshot.snapshot_id, id);
    }

    if let Some(Variant::Attributes(attrs)) = snapshot.properties.get(&ustr("Attributes")) {
        if let Some(id_value) = attrs.get(REF_ID_ATTRIBUTE_NAME) {
            if let Some(id_str) = crate::variant_as_str(id_value, REF_ID_ATTRIBUTE_NAME) {
                context
                    .snapshot_specified_ids
                    .insert(id_str.to_string(), id);
            }
        }
    }

    let instance = tree
        .get_instance(id)
        .expect("Instance did not exist in tree");

    let glob_ignore_rules = if snapshot.metadata.glob_ignored_children {
        Some(Arc::clone(&snapshot.metadata.context.path_ignore_rules))
    } else {
        None
    };

    compute_property_patches(
        &context.snapshot_specified_ids,
        &mut snapshot,
        &instance,
        patch_set,
        tree,
    );
    compute_children_patches(
        context,
        &mut snapshot,
        tree,
        id,
        patch_set,
        session,
        glob_ignore_rules.as_ref().map(|v| v.as_slice()),
    );
}

fn compute_property_patches(
    snapshot_specified_ids: &HashMap<String, Ref>,
    snapshot: &mut InstanceSnapshot,
    instance: &InstanceWithMeta,
    patch_set: &mut PatchSet,
    tree: &RojoTree,
) {
    let mut visited_properties = UstrSet::default();
    let mut changed_properties = UstrMap::new();

    let attribute_ref_properties =
        compute_ref_properties(snapshot, tree, instance.id(), snapshot_specified_ids);

    // Strip Rojo ref transport attributes from the snapshot before comparison.
    // finalize_patch_application strips these from the tree after resolving refs,
    // so the snapshot must be stripped too to prevent spurious diffs on duplicate
    // VFS events (common on macOS). The refs were already extracted above.
    if let Some(Variant::Attributes(attrs)) = snapshot.properties.get_mut(&ustr("Attributes")) {
        let keys_to_remove: Vec<String> = attrs
            .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 {
            attrs.remove(key);
        }
    }

    let changed_name = if snapshot.name == instance.name() {
        None
    } else {
        Some(take(&mut snapshot.name).into_owned())
    };

    let changed_class_name = if snapshot.class_name == instance.class_name() {
        None
    } else {
        Some(take(&mut snapshot.class_name))
    };

    let changed_metadata = if &snapshot.metadata == instance.metadata() {
        None
    } else {
        Some(take(&mut snapshot.metadata))
    };

    for (name, snapshot_value) in take(&mut snapshot.properties) {
        visited_properties.insert(name);

        match instance.properties().get(&name) {
            Some(instance_value) => {
                if !variant_eq(&snapshot_value, instance_value) {
                    changed_properties.insert(name, Some(snapshot_value));
                }
            }
            None => {
                changed_properties.insert(name, Some(snapshot_value));
            }
        }
    }

    for name in instance.properties().keys() {
        if visited_properties.contains(name) {
            continue;
        }

        changed_properties.insert(*name, None);
    }

    for (name, ref_value) in attribute_ref_properties {
        match (&ref_value, instance.properties().get(&name)) {
            (Some(referent), Some(instance_value)) => {
                if referent != instance_value {
                    changed_properties.insert(name, ref_value);
                } else {
                    changed_properties.remove(&name);
                }
            }
            (Some(_), None) | (None, Some(_)) => {
                changed_properties.insert(name, ref_value);
            }
            (None, None) => {
                changed_properties.remove(&name);
            }
        }
    }

    // !!!!!!!!!! UGLY HACK !!!!!!!!!!
    //
    // See RojoTree::insert_instance. Adjust that code also if you are touching this.
    let actual_class = changed_class_name.unwrap_or(instance.class_name());
    match actual_class.as_str() {
        "Model" | "Actor" | "Tool" | "HopperBin" | "Flag" | "WorldModel" | "Workspace"
        | "Status" => {
            let migration_prop = ustr("NeedsPivotMigration");
            // We want to just ignore this if it's being removed by a patch.
            // Normally this would not matter because serving != building but
            // if we start syncing models using SerializationService
            // (or GetObjects) it will affect how Studio deserializes things.
            if !instance.properties().contains_key(&migration_prop) {
                changed_properties.insert(migration_prop, Some(Variant::Bool(false)));
            }
            match changed_properties.get(&migration_prop) {
                Some(Some(Variant::Bool(_))) => {}
                Some(None) => {
                    changed_properties.remove(&migration_prop);
                }
                _ => {
                    changed_properties.insert(migration_prop, Some(Variant::Bool(false)));
                }
            }
        }
        _ => {}
    };

    if changed_properties.is_empty()
        && changed_name.is_none()
        && changed_class_name.is_none()
        && changed_metadata.is_none()
    {
        return;
    }

    patch_set.updated_instances.push(PatchUpdate {
        id: instance.id(),
        changed_name,
        changed_class_name,
        changed_properties,
        changed_metadata,
    });
}

fn compute_children_patches(
    context: &mut ComputePatchContext,
    snapshot: &mut InstanceSnapshot,
    tree: &RojoTree,
    id: Ref,
    patch_set: &mut PatchSet,
    session: &super::matching::MatchingSession,
    glob_ignore_rules: Option<&[PathIgnoreRule]>,
) {
    use super::matching::match_forward;

    let instance = tree
        .get_instance(id)
        .expect("Instance did not exist in tree");

    let instance_children = instance.children();
    let snapshot_children = take(&mut snapshot.children);

    if instance_children.is_empty() {
        for snapshot_child in snapshot_children {
            patch_set.added_instances.push(PatchAdd {
                parent_id: id,
                instance: snapshot_child,
            });
        }
        return;
    }

    let instance_children = instance_children.to_vec();
    let match_result = match_forward(snapshot_children, &instance_children, tree, session);

    // Matched pairs: recursively compute patches.
    for (snapshot_child, tree_child_id) in match_result.matched {
        compute_patch_set_internal(
            context,
            snapshot_child,
            tree,
            tree_child_id,
            patch_set,
            session,
        );
    }

    // Unmatched snapshots: new instances to be added.
    for snapshot_child in match_result.unmatched_snapshot {
        patch_set.added_instances.push(PatchAdd {
            parent_id: id,
            instance: snapshot_child,
        });
    }

    // Unmatched tree children: instances to be removed. When the parent has
    // glob-ignored children, each child is checked individually -- only
    // children whose InstigatingSource path matches a glob rule are kept.
    // Standard $ignoreUnknownInstances (scripts-only mode) still allows
    // server-side removals; the plugin's shouldDeleteChild handles the
    // selective keep logic for non-scripts.
    for tree_child_id in match_result.unmatched_tree {
        if let Some(rules) = glob_ignore_rules {
            let is_glob_ignored = tree
                .get_instance(tree_child_id)
                .and_then(|inst| inst.metadata().instigating_source.as_ref())
                .map(|source| match source {
                    InstigatingSource::Path(path) => !rules.iter().all(|rule| rule.passes(path)),
                    _ => false,
                })
                .unwrap_or(false);
            if is_glob_ignored {
                continue;
            }
        }
        patch_set.removed_instances.push(tree_child_id);
    }
}

fn compute_ref_properties(
    snapshot: &InstanceSnapshot,
    tree: &RojoTree,
    source_ref: Ref,
    snapshot_specified_ids: &HashMap<String, Ref>,
) -> UstrMap<Option<Variant>> {
    let mut map = UstrMap::new();
    let attributes = match snapshot.properties.get(&ustr("Attributes")) {
        Some(Variant::Attributes(attrs)) => attrs,
        _ => return map,
    };

    for (attr_name, attr_value) in attributes.iter() {
        // Handle legacy ID-based refs (Rojo_Target_) -- lower priority.
        // BTreeMap iterates alphabetically, so Rojo_Ref_* ('R') is visited
        // BEFORE Rojo_Target_* ('T'). Without the contains_key guard,
        // Rojo_Target_* would silently overwrite Rojo_Ref_*. The guard
        // ensures path-based refs always win when both exist.
        if let Some(prop_name) = attr_name.strip_prefix(REF_POINTER_ATTRIBUTE_PREFIX) {
            let key = ustr(prop_name);
            // Skip if Rojo_Ref_* already set a value for this property
            if map.contains_key(&key) {
                continue;
            }
            let Some(id_str) = crate::variant_as_str(attr_value, attr_name) else {
                continue;
            };
            let rojo_ref = RojoRef::new(id_str.to_string());
            if let Some(target_id) = tree.get_specified_id(&rojo_ref) {
                map.insert(key, Some(Variant::Ref(target_id)));
            } else if let Some(&target_id) = snapshot_specified_ids.get(id_str) {
                map.insert(key, Some(Variant::Ref(target_id)));
            } else {
                map.insert(key, None);
            }
            continue;
        }

        // Handle path-based refs (Rojo_Ref_) -- higher priority.
        // Always inserted unconditionally; if a Rojo_Target_* for the same
        // property was already visited (shouldn't happen due to BTreeMap
        // ordering, but defensive), this overwrites it.
        if let Some(prop_name) = attr_name.strip_prefix(REF_PATH_ATTRIBUTE_PREFIX) {
            let Some(path) = crate::variant_as_str(attr_value, attr_name) else {
                continue;
            };
            if let Some(target_id) = tree.resolve_ref_path(path, source_ref) {
                map.insert(ustr(prop_name), Some(Variant::Ref(target_id)));
            } else {
                // Path not found yet - will be resolved during patch apply
                map.insert(ustr(prop_name), None);
            }
            continue;
        }
    }

    map
}

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

    use std::borrow::Cow;

    /// This test makes sure that rewriting refs in instance update patches to
    /// instances that already exists works. We should be able to correlate the
    /// snapshot ID and instance ID during patch computation and replace the
    /// value before returning from compute_patch_set.
    #[test]
    fn rewrite_ref_existing_instance_update() {
        let tree = RojoTree::new(InstanceSnapshot::new().name("foo").class_name("foo"));

        let root_id = tree.get_root_id();

        // This snapshot should be identical to the existing tree except for the
        // addition of a prop named Self, which is a self-referential Ref.
        let snapshot_id = Ref::new();
        let snapshot = InstanceSnapshot {
            snapshot_id,
            properties: UstrMap::from_iter([(ustr("Self"), Variant::Ref(snapshot_id))]),

            metadata: Default::default(),
            name: Cow::Borrowed("foo"),
            class_name: ustr("foo"),
            children: Vec::new(),
        };

        let patch_set = compute_patch_set(Some(snapshot), &tree, root_id);

        let expected_patch_set = PatchSet {
            updated_instances: vec![PatchUpdate {
                id: root_id,
                changed_name: None,
                changed_class_name: None,
                changed_properties: UstrMap::from_iter([(
                    ustr("Self"),
                    Some(Variant::Ref(root_id)),
                )]),
                changed_metadata: None,
            }],
            added_instances: Vec::new(),
            removed_instances: Vec::new(),
            stage_ids: std::collections::HashSet::new(),
            stage_paths: Vec::new(),
        };

        assert_eq!(patch_set, expected_patch_set);
    }

    #[test]
    fn glob_ignored_children_suppresses_glob_matched_removals() {
        use super::super::patch_apply::apply_patch_set;
        use crate::glob::Glob;
        use std::path::PathBuf;

        let mut tree = RojoTree::new(InstanceSnapshot::new().name("root").class_name("Folder"));
        let root_id = tree.get_root_id();

        let add_patch = PatchSet {
            added_instances: vec![
                PatchAdd {
                    parent_id: root_id,
                    instance: InstanceSnapshot::new()
                        .name("A")
                        .class_name("Folder")
                        .metadata(
                            super::super::InstanceMetadata::new()
                                .instigating_source(PathBuf::from("/project/src/A")),
                        ),
                },
                PatchAdd {
                    parent_id: root_id,
                    instance: InstanceSnapshot::new()
                        .name("B")
                        .class_name("Folder")
                        .metadata(
                            super::super::InstanceMetadata::new()
                                .instigating_source(PathBuf::from("/project/src/generated/B")),
                        ),
                },
            ],
            updated_instances: Vec::new(),
            removed_instances: Vec::new(),
            stage_ids: std::collections::HashSet::new(),
            stage_paths: Vec::new(),
        };
        apply_patch_set(&mut tree, add_patch);

        let mut context = super::super::InstanceContext::new();
        context.add_path_ignore_rules(vec![PathIgnoreRule {
            base_path: PathBuf::from("/project"),
            glob: Glob::new("src/generated/**").unwrap(),
        }]);

        // Snapshot has only A. B is unmatched but its path matches the
        // glob rule, so it should NOT be removed.
        let snapshot = InstanceSnapshot {
            snapshot_id: Ref::none(),
            metadata: super::super::InstanceMetadata::new()
                .ignore_unknown_instances(true)
                .glob_ignored_children(true)
                .context(&context),
            properties: UstrMap::new(),
            name: Cow::Borrowed("root"),
            class_name: ustr("Folder"),
            children: vec![InstanceSnapshot::new().name("A").class_name("Folder")],
        };

        let patch_set = compute_patch_set(Some(snapshot), &tree, root_id);

        assert!(
            patch_set.removed_instances.is_empty(),
            "Glob-matched child should NOT be removed"
        );
    }

    #[test]
    fn glob_ignored_children_per_child_matching() {
        use super::super::patch_apply::apply_patch_set;
        use crate::glob::Glob;
        use std::path::PathBuf;

        let mut tree = RojoTree::new(InstanceSnapshot::new().name("root").class_name("Folder"));
        let root_id = tree.get_root_id();

        // A: will be in the snapshot (matched).
        // B: deleted from filesystem, path does NOT match glob → removed.
        // C: glob-ignored, path matches glob → NOT removed.
        let add_patch = PatchSet {
            added_instances: vec![
                PatchAdd {
                    parent_id: root_id,
                    instance: InstanceSnapshot::new()
                        .name("A")
                        .class_name("Folder")
                        .metadata(
                            super::super::InstanceMetadata::new()
                                .instigating_source(PathBuf::from("/project/src/A")),
                        ),
                },
                PatchAdd {
                    parent_id: root_id,
                    instance: InstanceSnapshot::new()
                        .name("B")
                        .class_name("Folder")
                        .metadata(
                            super::super::InstanceMetadata::new()
                                .instigating_source(PathBuf::from("/project/src/B")),
                        ),
                },
                PatchAdd {
                    parent_id: root_id,
                    instance: InstanceSnapshot::new()
                        .name("C")
                        .class_name("Folder")
                        .metadata(
                            super::super::InstanceMetadata::new()
                                .instigating_source(PathBuf::from("/project/src/generated/C")),
                        ),
                },
            ],
            updated_instances: Vec::new(),
            removed_instances: Vec::new(),
            stage_ids: std::collections::HashSet::new(),
            stage_paths: Vec::new(),
        };
        apply_patch_set(&mut tree, add_patch);

        let mut context = super::super::InstanceContext::new();
        context.add_path_ignore_rules(vec![PathIgnoreRule {
            base_path: PathBuf::from("/project"),
            glob: Glob::new("src/generated/**").unwrap(),
        }]);

        let snapshot = InstanceSnapshot {
            snapshot_id: Ref::none(),
            metadata: super::super::InstanceMetadata::new()
                .ignore_unknown_instances(true)
                .glob_ignored_children(true)
                .context(&context),
            properties: UstrMap::new(),
            name: Cow::Borrowed("root"),
            class_name: ustr("Folder"),
            children: vec![InstanceSnapshot::new().name("A").class_name("Folder")],
        };

        let patch_set = compute_patch_set(Some(snapshot), &tree, root_id);

        assert_eq!(
            patch_set.removed_instances.len(),
            1,
            "Only genuinely deleted child B should be removed, not glob-ignored child C"
        );

        let removed_inst = tree
            .get_instance(patch_set.removed_instances[0])
            .expect("Removed instance should exist in tree");
        assert_eq!(
            removed_inst.name(),
            "B",
            "The removed instance should be B (genuine deletion)"
        );
    }

    #[test]
    fn ignore_unknown_instances_alone_still_removes() {
        use super::super::patch_apply::apply_patch_set;

        let mut tree = RojoTree::new(InstanceSnapshot::new().name("root").class_name("Folder"));
        let root_id = tree.get_root_id();

        let add_patch = PatchSet {
            added_instances: vec![
                PatchAdd {
                    parent_id: root_id,
                    instance: InstanceSnapshot::new().name("A").class_name("Folder"),
                },
                PatchAdd {
                    parent_id: root_id,
                    instance: InstanceSnapshot::new().name("B").class_name("Folder"),
                },
            ],
            updated_instances: Vec::new(),
            removed_instances: Vec::new(),
            stage_ids: std::collections::HashSet::new(),
            stage_paths: Vec::new(),
        };
        apply_patch_set(&mut tree, add_patch);

        // ignore_unknown_instances = true (scripts-only mode) but
        // glob_ignored_children = false. Server still sends removals;
        // the plugin's shouldDeleteChild handles the selective keep.
        let snapshot = InstanceSnapshot {
            snapshot_id: Ref::none(),
            metadata: super::super::InstanceMetadata::new().ignore_unknown_instances(true),
            properties: UstrMap::new(),
            name: Cow::Borrowed("root"),
            class_name: ustr("Folder"),
            children: vec![InstanceSnapshot::new().name("A").class_name("Folder")],
        };

        let patch_set = compute_patch_set(Some(snapshot), &tree, root_id);

        assert_eq!(
            patch_set.removed_instances.len(),
            1,
            "ignore_unknown_instances alone should NOT suppress server-side removals"
        );
    }

    #[test]
    fn default_removes_unmatched_children() {
        use super::super::patch_apply::apply_patch_set;

        let mut tree = RojoTree::new(InstanceSnapshot::new().name("root").class_name("Folder"));
        let root_id = tree.get_root_id();

        let add_patch = PatchSet {
            added_instances: vec![
                PatchAdd {
                    parent_id: root_id,
                    instance: InstanceSnapshot::new().name("A").class_name("Folder"),
                },
                PatchAdd {
                    parent_id: root_id,
                    instance: InstanceSnapshot::new().name("B").class_name("Folder"),
                },
            ],
            updated_instances: Vec::new(),
            removed_instances: Vec::new(),
            stage_ids: std::collections::HashSet::new(),
            stage_paths: Vec::new(),
        };
        apply_patch_set(&mut tree, add_patch);

        let snapshot = InstanceSnapshot {
            snapshot_id: Ref::none(),
            metadata: Default::default(),
            properties: UstrMap::new(),
            name: Cow::Borrowed("root"),
            class_name: ustr("Folder"),
            children: vec![InstanceSnapshot::new().name("A").class_name("Folder")],
        };

        let patch_set = compute_patch_set(Some(snapshot), &tree, root_id);

        assert_eq!(
            patch_set.removed_instances.len(),
            1,
            "Unmatched tree children should be removed with default metadata"
        );
    }

    /// The same as rewrite_ref_existing_instance_update, except that the
    /// property is added in a new instance instead of modifying an existing
    /// one.
    #[test]
    fn rewrite_ref_existing_instance_addition() {
        let tree = RojoTree::new(InstanceSnapshot::new().name("foo").class_name("foo"));

        let root_id = tree.get_root_id();

        // This patch describes the existing instance with a new child added.
        let snapshot_id = Ref::new();
        let snapshot = InstanceSnapshot {
            snapshot_id,
            children: vec![InstanceSnapshot {
                properties: UstrMap::from_iter([(ustr("Self"), Variant::Ref(snapshot_id))]),

                snapshot_id: Ref::none(),
                metadata: Default::default(),
                name: Cow::Borrowed("child"),
                class_name: ustr("child"),
                children: Vec::new(),
            }],

            metadata: Default::default(),
            properties: UstrMap::new(),
            name: Cow::Borrowed("foo"),
            class_name: ustr("foo"),
        };

        let patch_set = compute_patch_set(Some(snapshot), &tree, root_id);

        let expected_patch_set = PatchSet {
            added_instances: vec![PatchAdd {
                parent_id: root_id,
                instance: InstanceSnapshot {
                    snapshot_id: Ref::none(),
                    metadata: Default::default(),
                    properties: UstrMap::from_iter([(ustr("Self"), Variant::Ref(root_id))]),
                    name: Cow::Borrowed("child"),
                    class_name: ustr("child"),
                    children: Vec::new(),
                },
            }],
            updated_instances: Vec::new(),
            removed_instances: Vec::new(),
            stage_ids: std::collections::HashSet::new(),
            stage_paths: Vec::new(),
        };

        assert_eq!(patch_set, expected_patch_set);
    }
}