asset-tool/rox_compiler/src/decompile.rs

use std::path::PathBuf;
use rbx_dom_weak::types::Ref;
use crate::common::{sanitize,Style,PropertiesOverride};

#[derive(PartialEq)]
enum Class{
	Folder,
	ModuleScript,
	LocalScript,
	Script,
	Model,
}

struct TreeNode{
	name:String,
	referent:Ref,
	parent:Ref,
	class:Class,
	children:Vec<Ref>,
}
impl TreeNode{
	fn new(name:String,referent:Ref,parent:Ref,class:Class)->Self{
		Self{
			name,
			referent,
			parent,
			class,
			children:Vec::new(),
		}
	}
}

enum TrimStackInstruction{
	Referent(Ref),
	IncrementScript,
	DecrementScript,
}

enum WriteStackInstruction<'a>{
	Node(&'a TreeNode,u32),//(Node,NameTally)
	PushFolder(String),
	PopFolder,
	Destroy(Ref),
}

#[derive(Debug)]
pub enum WriteError{
	ClassNotScript(String),
	IO(std::io::Error),
	EncodeError(rbx_xml::EncodeError),
}
impl std::fmt::Display for WriteError{
	fn fmt(&self,f:&mut std::fmt::Formatter<'_>)->std::fmt::Result{
		write!(f,"{self:?}")
	}
}
impl std::error::Error for WriteError{}

fn write_item(dom:&rbx_dom_weak::WeakDom,mut file:PathBuf,node:&TreeNode,node_name_override:String,mut properties:PropertiesOverride,style:Style,write_models:bool,write_scripts:bool)->Result<(),WriteError>{
	file.push(sanitize(node_name_override.as_str()).as_ref());
	match node.class{
		Class::Folder=>(),
		Class::ModuleScript|Class::LocalScript|Class::Script=>{
			if !write_scripts{
				return Ok(())
			}

			//set extension
			match style{
				Style::Rox=>assert!(file.set_extension("lua"),"could not set extension"),
				Style::RoxRojo|Style::Rojo=>{
					match properties.class.as_deref(){
						Some("LocalScript")=>{
							file.set_extension("client.lua");
							properties.class=None;
						},
						Some("Script")=>{
							file.set_extension("server.lua");
							properties.class=None;
						},
						// Some("ModuleScript")=>{
						// 	file.set_extension("module");
						// 	properties.class=None;
						// },
						None=>assert!(file.set_extension("lua"),"could not set extension"),
						Some(other)=>Err(WriteError::ClassNotScript(other.to_owned()))?,
					}
				}
			}

			if let Some(item)=dom.get_by_ref(node.referent){
				//TODO: delete disabled scripts
				if let Some(rbx_dom_weak::types::Variant::String(source))=item.properties.get("Source"){
					if properties.is_some(){
						//rox style
						let source=properties.to_string()+source.as_str();
						std::fs::write(file,source).map_err(WriteError::IO)?;
					}else{
						std::fs::write(file,source).map_err(WriteError::IO)?;
					}
				}
			}
		},
		Class::Model=>{
			if !write_models{
				return Ok(())
			}
			assert!(file.set_extension("rbxmx"));
			let output=std::io::BufWriter::new(std::fs::File::create(file).map_err(WriteError::IO)?);
			rbx_xml::to_writer_default(output,dom,&[node.referent]).map_err(WriteError::EncodeError)?;
		},
	}
	Ok(())
}

pub struct WriteConfig{
	pub style:Style,
	pub output_folder:PathBuf,
	pub write_template:bool,
	pub write_models:bool,
	pub write_scripts:bool,
}

pub struct DecompiledContext{
	dom:rbx_dom_weak::WeakDom,
	tree_refs:std::collections::HashMap<rbx_dom_weak::types::Ref,TreeNode>,
}

impl DecompiledContext{
	/// Will panic on circular tree structure but otherwise infallible
	pub fn from_dom(dom:rbx_dom_weak::WeakDom)->Self{
		let mut tree_refs=std::collections::HashMap::new();
		tree_refs.insert(dom.root_ref(),TreeNode::new(
			"src".to_owned(),
			dom.root_ref(),
			Ref::none(),
			Class::Folder
		));

		//run rules
		let mut stack=vec![dom.root()];
		while let Some(item)=stack.pop(){
			let class=match item.class.as_str(){
				"ModuleScript"=>Class::ModuleScript,
				"LocalScript"=>Class::LocalScript,
				"Script"=>Class::Script,
				"Model"=>Class::Model,
				_=>Class::Folder,
			};
			let skip=class==Class::Model;
			if let Some(parent_node)=tree_refs.get_mut(&item.parent()){
				let referent=item.referent();
				let node=TreeNode::new(item.name.clone(),referent,parent_node.referent,class);
				parent_node.children.push(referent);
				tree_refs.insert(referent,node);
			}
			//look no further, turn this node and all its children into a model
			if skip{
				continue;
			}
			for &referent in item.children(){
				if let Some(c)=dom.get_by_ref(referent){
					stack.push(c);
				}
			}
		}

		//trim empty folders
		let mut script_count=0;
		let mut stack:Vec<TrimStackInstruction>=tree_refs.get(&dom.root_ref()).unwrap().children
		.iter().map(|&c|TrimStackInstruction::Referent(c)).collect();
		while let Some(instruction)=stack.pop(){
			match instruction{
				TrimStackInstruction::IncrementScript=>script_count+=1,
				TrimStackInstruction::DecrementScript=>script_count-=1,
				TrimStackInstruction::Referent(referent)=>{
					let mut delete=None;
					if let Some(node)=tree_refs.get_mut(&referent){
						if node.class==Class::Folder&&script_count!=0{
							node.class=Class::Model
						}
						if node.class==Class::Folder&&node.children.is_empty(){
							delete=Some(node.parent);
						}else{
							//how the hell do I do this better without recursion
							let is_script=matches!(
								node.class,
								Class::ModuleScript|Class::LocalScript|Class::Script
							);
							//stack is popped from back
							if is_script{
								stack.push(TrimStackInstruction::DecrementScript);
							}
							for &child_referent in &node.children{
								stack.push(TrimStackInstruction::Referent(child_referent));
							}
							if is_script{
								stack.push(TrimStackInstruction::IncrementScript);
							}
						}
					}
					//trim referent
					if let Some(parent_ref)=delete{
						let parent_node=tree_refs.get_mut(&parent_ref)
						.expect("parent_ref does not exist in tree_refs");
						parent_node.children.remove(
							parent_node.children.iter()
							.position(|&r|r==referent)
							.expect("parent.children does not contain referent")
						);
						tree_refs.remove(&referent);
					}
				},
			}
		}

		Self{
			dom,
			tree_refs,
		}
	}
	pub async fn write_files(mut self,config:WriteConfig)->Result<(),WriteError>{
		let mut write_queue=Vec::new();
		let mut destroy_queue=Vec::new();

		let mut name_tally=std::collections::HashMap::<String,u32>::new();
		let mut folder=config.output_folder.clone();
		let mut stack=vec![WriteStackInstruction::Node(self.tree_refs.get(&self.dom.root_ref()).unwrap(),0)];
		while let Some(instruction)=stack.pop(){
			match instruction{
				WriteStackInstruction::PushFolder(component)=>folder.push(component),
				WriteStackInstruction::PopFolder=>assert!(folder.pop(),"weirdness"),
				WriteStackInstruction::Destroy(referent)=>destroy_queue.push(referent),
				WriteStackInstruction::Node(node,name_count)=>{
					//track properties that must be overriden to compile folder structure back into a place file
					let mut properties=PropertiesOverride::default();
					let has_children=node.children.is_empty();
					match node.class{
						Class::Folder=>(),
						Class::ModuleScript=>(),//.lua files are ModuleScript by default
						Class::LocalScript=>properties.class=Some("LocalScript".to_owned()),
						Class::Script=>properties.class=Some("Script".to_owned()),
						Class::Model=>(),
					}
					let name_override=if 0<name_count{
						properties.name=Some(node.name.clone());
						format!("{}_{}",node.name,name_count)
					}else{
						node.name.clone()
					};

					if has_children{
						//push temp subfolder
						let mut subfolder=folder.clone();
						subfolder.push(sanitize(name_override.as_str()).as_ref());
						//make folder
						tokio::fs::create_dir(subfolder.clone()).await.map_err(WriteError::IO)?;

						let name_final=match config.style{
							Style::Rox
							|Style::RoxRojo=>name_override.clone(),
							Style::Rojo=>"init".to_owned(),
						};

						//write item in subfolder
						write_queue.push((subfolder,node,name_final,properties,config.style));
					}else{
						//write item
						write_queue.push((folder.clone(),node,name_override.clone(),properties,config.style));
					}
					//queue item to be deleted from dom after child objects are handled (stack is popped from the back)
					match node.class{
						Class::Folder=>(),
						_=>stack.push(WriteStackInstruction::Destroy(node.referent)),
					}
					if has_children{
						stack.push(WriteStackInstruction::PopFolder);
						name_tally.clear();
						for referent in &node.children{
							if let Some(c)=self.tree_refs.get(referent){
								let v=name_tally.entry(c.name.clone()).and_modify(|v|*v+=1).or_default();
								stack.push(WriteStackInstruction::Node(c,*v));
							}
						}
						stack.push(WriteStackInstruction::PushFolder(sanitize(name_override.as_str()).to_string()));
					}
				},
			}
		}

		//run the async
		{
			let dom=&self.dom;
			let write_models=config.write_models;
			let write_scripts=config.write_scripts;
			let results:Vec<Result<(),WriteError>>=rayon::iter::ParallelIterator::collect(rayon::iter::ParallelIterator::map(rayon::iter::IntoParallelIterator::into_par_iter(write_queue),|(write_path,node,node_name_override,properties,style)|{
				write_item(dom,write_path,node,node_name_override,properties,style,write_models,write_scripts)
			}));
			for result in results{
				result?;
			}
		}

		//run the destroy
		for destroy_ref in destroy_queue{
			self.dom.destroy(destroy_ref);
		}

		//write what remains in template.rbxlx
		if config.write_template{
			let mut file=config.output_folder.clone();
			file.push("template");
			assert!(file.set_extension("rbxlx"));
			let output=std::io::BufWriter::new(std::fs::File::create(file).map_err(WriteError::IO)?);
			rbx_xml::to_writer_default(output,&self.dom,self.dom.root().children()).map_err(WriteError::EncodeError)?;
		}

		Ok(())
	}
}