From c6d293cc6bc3fe582eed2bb2a54d4add4f42a286 Mon Sep 17 00:00:00 2001
From: Quaternions <krakow20@gmail.com>
Date: Sat, 18 Nov 2023 15:48:30 -0800
Subject: [PATCH] write "Ref" attribute based on internal roblox part id

---
 src/main.rs | 62 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)

diff --git a/src/main.rs b/src/main.rs
index 3343008..aa058f7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -17,6 +17,7 @@ enum Commands {
     ConvertTextures,
     DownloadMeshes(PathBufList),
     Extract(PathBufList),
+    WriteAttributes,
     Interactive,
     Replace,
     Scan,
@@ -56,6 +57,16 @@ fn recursive_collect_superclass(objects: &mut std::vec::Vec<rbx_dom_weak::types:
         }
     }
 }
+fn recursive_collect_regex(objects: &mut std::vec::Vec<rbx_dom_weak::types::Ref>,dom: &rbx_dom_weak::WeakDom, instance: &rbx_dom_weak::Instance, regex: &lazy_regex::Lazy<lazy_regex::Regex>){
+    for &referent in instance.children() {
+        if let Some(c) = dom.get_by_ref(referent) {
+            if regex.captures(c.name.as_str()).is_some(){
+                objects.push(c.referent());//copy ref
+            }
+            recursive_collect_regex(objects,dom,c,regex);
+        }
+    }
+}
 fn get_full_name(dom:&rbx_dom_weak::WeakDom,instance:&rbx_dom_weak::Instance) -> String{
     let mut full_name=instance.name.clone();
     let mut pref=instance.parent();
@@ -90,6 +101,11 @@ fn get_script_refs(dom:&rbx_dom_weak::WeakDom) -> Vec<rbx_dom_weak::types::Ref>{
     recursive_collect_superclass(&mut scripts, dom, dom.root(),"LuaSourceContainer");
     scripts
 }
+fn get_button_refs(dom:&rbx_dom_weak::WeakDom) -> Vec<rbx_dom_weak::types::Ref>{
+    let mut buttons = std::vec::Vec::new();
+    recursive_collect_regex(&mut buttons, dom, dom.root(),lazy_regex::regex!(r"Button(\d+)$"));
+    buttons
+}
 fn get_texture_refs(dom:&rbx_dom_weak::WeakDom) -> Vec<rbx_dom_weak::types::Ref>{
     let mut objects = std::vec::Vec::new();
     recursive_collect_superclass(&mut objects, dom, dom.root(),"Decal");
@@ -961,6 +977,51 @@ fn unzip_all()->AResult<()>{
     Ok(())
 }
 
+fn write_attributes() -> AResult<()>{
+    for entry in std::fs::read_dir("maps/unprocessed")? {
+        let file_thing=entry?;
+        println!("processing map={:?}",file_thing.file_name());
+        let mut input = std::io::BufReader::new(std::fs::File::open(file_thing.path())?);
+        let mut dom = get_dom(&mut input)?;
+
+        let button_refs = get_button_refs(&dom);
+
+        for &button_ref in &button_refs {
+            if let Some(button)=dom.get_by_ref_mut(button_ref){
+                match button.properties.get_mut("Attributes"){
+                    Some(rbx_dom_weak::types::Variant::Attributes(attributes))=>{
+                        println!("Appending Ref={} to existing attributes for {}",button_ref,button.name);
+                        attributes.insert("Ref".to_string(),rbx_dom_weak::types::Variant::String(button_ref.to_string()));
+                    },
+                    None=>{
+                        println!("Creating new attributes with Ref={} for {}",button_ref,button.name);
+                        let mut attributes=rbx_dom_weak::types::Attributes::new();
+                        attributes.insert("Ref".to_string(),rbx_dom_weak::types::Variant::String(button_ref.to_string()));
+                        button.properties.insert("Attributes".to_string(),rbx_dom_weak::types::Variant::Attributes(attributes));
+                    }
+                    _=>unreachable!("Fetching attributes did not return attributes."),
+                }
+            }
+        }
+        let mut dest={
+            let mut dest=std::path::PathBuf::from("maps/attributes");
+            dest.push(file_thing.file_name());
+            let output = std::io::BufWriter::new(std::fs::File::create(dest)?);
+            //write workspace:GetChildren()[1]
+            let workspace_children=dom.root().children();
+            if workspace_children.len()!=1{
+                return Err(anyhow::Error::msg("there can only be one model"));
+            }
+            rbx_binary::to_writer(output, &dom, &[workspace_children[0]])?;
+            //move original to processed folder
+            std::path::PathBuf::from("maps/unaltered")
+        };
+        dest.push(file_thing.file_name());
+        std::fs::rename(file_thing.path(), dest)?;
+    }
+    Ok(())
+}
+
 fn main() -> AResult<()> {
     let cli = Cli::parse();
     match cli.command {
@@ -969,6 +1030,7 @@ fn main() -> AResult<()> {
         Commands::ConvertTextures=>convert_textures(),
         Commands::DownloadMeshes(pathlist)=>download_meshes(pathlist.paths),
         Commands::Extract(pathlist)=>extract(pathlist.paths),
+        Commands::WriteAttributes=>write_attributes(),
         Commands::Interactive=>interactive(),
         Commands::Replace=>replace(),
         Commands::Scan=>scan(),