feat: add --base flag to upload command for merging with place files

src/cli/upload.rs

@@ -1,8 +1,14 @@
-use std::path::PathBuf;
+use std::{
+    io::BufReader,
+    mem::forget,
+    path::{Path, PathBuf},
+};
 
 use anyhow::{bail, Context};
 use clap::Parser;
+use fs_err::File;
 use memofs::Vfs;
+use rbx_dom_weak::{types::Ref, Instance, WeakDom};
 use reqwest::{
     header::{ACCEPT, CONTENT_TYPE, COOKIE, USER_AGENT},
     StatusCode,
@@ -19,6 +25,13 @@ pub struct UploadCommand {
     #[clap(default_value = "")]
     pub project: PathBuf,
 
+    /// Path to a base .rbxl or .rbxlx file to merge with the project before
+    /// uploading. When provided, the Rojo project tree is merged into this
+    /// file so that the upload contains both the base content (3D assets,
+    /// terrain, lighting, etc.) and the project's scripts.
+    #[clap(long)]
+    pub base: Option<PathBuf>,
+
     /// Authenication cookie to use. If not specified, Rojo will attempt to find one from the system automatically.
     #[clap(long)]
     pub cookie: Option<String>,
@@ -44,19 +57,38 @@ impl UploadCommand {
 
         let session = ServeSession::new(vfs, project_path)?;
 
-        let tree = session.tree();
-        let inner_tree = tree.inner();
-        let root = inner_tree.root();
-
-        let encode_ids = match root.class.as_str() {
-            "DataModel" => root.children().to_vec(),
-            _ => vec![root.referent()],
-        };
-
         let mut buffer = Vec::new();
 
-        log::trace!("Encoding binary model");
+        if let Some(base_path) = &self.base {
-        rbx_binary::to_writer(&mut buffer, tree.inner(), &encode_ids)?;
+            let base_path = resolve_path(base_path);
+
+            log::trace!("Reading base file: {}", base_path.display());
+            let base_dom = read_base_file(&base_path)?;
+
+            log::trace!("Merging Rojo project into base file");
+            let merged_dom = merge_rojo_into_base(base_dom, session.tree().inner())?;
+
+            let encode_ids = merged_dom.root().children().to_vec();
+
+            log::trace!("Encoding merged binary model");
+            rbx_binary::to_writer(&mut buffer, &merged_dom, &encode_ids)?;
+        } else {
+            let tree = session.tree();
+            let inner_tree = tree.inner();
+            let root = inner_tree.root();
+
+            let encode_ids = match root.class.as_str() {
+                "DataModel" => root.children().to_vec(),
+                _ => vec![root.referent()],
+            };
+
+            log::trace!("Encoding binary model");
+            rbx_binary::to_writer(&mut buffer, inner_tree, &encode_ids)?;
+        }
+
+        // Avoid dropping ServeSession: it's potentially very expensive to
+        // drop and we're about to exit anyway.
+        forget(session);
 
         match (self.cookie, self.api_key, self.universe_id) {
             (cookie, None, universe) => {
@@ -90,6 +122,134 @@ impl UploadCommand {
     }
 }
 
+/// Reads a base .rbxl or .rbxlx file into a WeakDom.
+fn read_base_file(path: &Path) -> anyhow::Result<WeakDom> {
+    let extension = path
+        .extension()
+        .and_then(|ext| ext.to_str())
+        .context("Base file must have a .rbxl or .rbxlx extension")?;
+
+    let content = BufReader::new(
+        File::open(path)
+            .with_context(|| format!("Could not open base file at {}", path.display()))?,
+    );
+
+    match extension {
+        "rbxl" => rbx_binary::from_reader(content).with_context(|| {
+            format!(
+                "Could not deserialize binary place file at {}",
+                path.display()
+            )
+        }),
+        "rbxlx" => {
+            let config = rbx_xml::DecodeOptions::new()
+                .property_behavior(rbx_xml::DecodePropertyBehavior::ReadUnknown);
+            rbx_xml::from_reader(content, config).with_context(|| {
+                format!(
+                    "Could not deserialize XML place file at {}",
+                    path.display()
+                )
+            })
+        }
+        _ => bail!("Base file must be .rbxl or .rbxlx, got .{}", extension),
+    }
+}
+
+/// Merges the Rojo project tree into a base WeakDom.
+///
+/// For each service that exists in the Rojo tree, finds the matching service
+/// in the base DOM by name and className and recursively merges Rojo children
+/// into it. Instances in the base that Rojo doesn't manage are left untouched,
+/// preserving 3D assets, terrain, lighting, and other non-scripted content.
+fn merge_rojo_into_base(mut base: WeakDom, rojo: &WeakDom) -> anyhow::Result<WeakDom> {
+    let rojo_root = rojo.root();
+    let base_root_ref = base.root_ref();
+
+    if rojo_root.class.as_str() != "DataModel" {
+        // Non-place project (model): clone the entire Rojo tree as a child
+        let cloned_ref = rojo.clone_into_external(rojo_root.referent(), &mut base);
+        base.transfer_within(cloned_ref, base_root_ref);
+        return Ok(base);
+    }
+
+    merge_instances(&mut base, base_root_ref, rojo, rojo.root_ref());
+
+    Ok(base)
+}
+
+/// Recursively merges the children of a Rojo instance into the corresponding
+/// base instance.
+///
+/// For each Rojo child:
+///   - If a matching child (by name+class) exists in base: update its
+///     properties and recurse into its children
+///   - If no match: clone the entire Rojo subtree into base under the
+///     current parent
+///
+/// Base children that have no Rojo counterpart are never touched, preserving
+/// 3D assets and other content that Rojo doesn't manage.
+fn merge_instances(
+    base: &mut WeakDom,
+    base_parent_ref: Ref,
+    rojo: &WeakDom,
+    rojo_parent_ref: Ref,
+) {
+    for &rojo_child_ref in rojo.get_by_ref(rojo_parent_ref).unwrap().children() {
+        let rojo_child = rojo.get_by_ref(rojo_child_ref).unwrap();
+
+        let base_match = find_child_by_name_and_class(
+            base,
+            base_parent_ref,
+            &rojo_child.name,
+            rojo_child.class.as_str(),
+        );
+
+        match base_match {
+            Some(base_child_ref) => {
+                update_properties(base, base_child_ref, rojo_child);
+                merge_instances(base, base_child_ref, rojo, rojo_child_ref);
+            }
+            None => {
+                let cloned_ref = rojo.clone_into_external(rojo_child_ref, base);
+                base.transfer_within(cloned_ref, base_parent_ref);
+            }
+        }
+    }
+}
+
+/// Finds a child of `parent_ref` in the given DOM that matches both name and
+/// className.
+fn find_child_by_name_and_class(
+    dom: &WeakDom,
+    parent_ref: Ref,
+    name: &str,
+    class: &str,
+) -> Option<Ref> {
+    let parent = dom.get_by_ref(parent_ref)?;
+    for &child_ref in parent.children() {
+        let child = dom.get_by_ref(child_ref)?;
+        if child.name == name && child.class.as_str() == class {
+            return Some(child_ref);
+        }
+    }
+    None
+}
+
+/// Overlays Rojo properties onto a base instance. Properties that exist only
+/// in the base instance are preserved; properties from Rojo overwrite any
+/// existing base values.
+fn update_properties(base: &mut WeakDom, base_ref: Ref, rojo_instance: &Instance) {
+    let base_instance = base.get_by_ref_mut(base_ref).unwrap();
+
+    base_instance.name.clone_from(&rojo_instance.name);
+
+    for (prop_name, prop_value) in &rojo_instance.properties {
+        base_instance
+            .properties
+            .insert(*prop_name, prop_value.clone());
+    }
+}
+
 fn do_upload(buffer: Vec<u8>, asset_id: u64, cookie: &str) -> anyhow::Result<()> {
     let url = format!(
         "https://data.roblox.com/Data/Upload.ashx?assetid={}",