├── lua51-lifter ├── .gitignore ├── Cargo.toml └── src │ └── main.rs ├── luau-lifter ├── .gitignore ├── src │ ├── main.rs │ ├── deserializer │ │ ├── list.rs │ │ ├── bytecode.rs │ │ ├── mod.rs │ │ ├── chunk.rs │ │ ├── constant.rs │ │ └── function.rs │ ├── instruction.rs │ └── lib.rs └── Cargo.toml ├── cfg ├── .gitignore ├── src │ ├── ssa.rs │ ├── lib.rs │ ├── block.rs │ ├── pattern.rs │ ├── dot.rs │ ├── ssa │ │ ├── destruct │ │ │ └── liveness.rs │ │ ├── upvalues.rs │ │ └── param_dependency_graph.rs │ └── function.rs └── Cargo.toml ├── playermodule.bin ├── .gitignore ├── lua51-deserializer ├── src │ ├── lib.rs │ ├── instruction │ │ ├── position.rs │ │ ├── argument.rs │ │ ├── layout.rs │ │ └── operation_code.rs │ ├── local.rs │ ├── chunk │ │ ├── mod.rs │ │ └── header.rs │ ├── value.rs │ └── function.rs └── Cargo.toml ├── README.md ├── script.lua ├── ast ├── src │ ├── break.rs │ ├── vararg.rs │ ├── continue.rs │ ├── close.rs │ ├── side_effects.rs │ ├── name_gen.rs │ ├── paste.txt │ ├── global.rs │ ├── goto.rs │ ├── return.rs │ ├── while.rs │ ├── repeat.rs │ ├── index.rs │ ├── if.rs │ ├── replace_locals.rs │ ├── assign.rs │ ├── closure.rs │ ├── literal.rs │ ├── set_list.rs │ ├── local.rs │ ├── table.rs │ ├── call.rs │ ├── traverse.rs │ ├── name_locals.rs │ ├── type_system.rs │ ├── local_declarations.rs │ ├── lib.rs │ └── binary.rs └── Cargo.toml ├── decomp-server ├── Cargo.toml └── src │ └── main.rs ├── luau-worker ├── Cargo.toml ├── wrangler.toml └── src │ └── lib.rs ├── restructure ├── Cargo.toml └── src │ ├── jump.rs │ └── conditional.rs ├── Cargo.toml └── LICENSE.txt /lua51-lifter/.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | -------------------------------------------------------------------------------- /luau-lifter/.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | -------------------------------------------------------------------------------- /cfg/.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | Cargo.lock 3 | -------------------------------------------------------------------------------- /playermodule.bin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xeonise/medal2/HEAD/playermodule.bin -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | target/ 2 | .idea 3 | .vscode 4 | Cargo.lock 5 | flamegraph.svg 6 | result.lua -------------------------------------------------------------------------------- /lua51-deserializer/src/lib.rs: -------------------------------------------------------------------------------- 1 | pub use function::Function; 2 | pub use instruction::{argument, Instruction}; 3 | pub use value::Value; 4 | 5 | pub mod chunk; 6 | pub mod function; 7 | pub mod instruction; 8 | pub mod local; 9 | pub mod value; 10 | -------------------------------------------------------------------------------- /cfg/src/ssa.rs: -------------------------------------------------------------------------------- 1 | pub mod construct; 2 | mod destruct; 3 | pub mod inline; 4 | mod param_dependency_graph; 5 | pub mod structuring; 6 | pub mod upvalues; 7 | //pub mod dataflow; 8 | 9 | pub use construct::construct; 10 | pub use destruct::Destructor; 11 | -------------------------------------------------------------------------------- /cfg/src/lib.rs: -------------------------------------------------------------------------------- 1 | #![feature(box_patterns)] 2 | #![feature(box_into_inner)] 3 | #![feature(let_chains)] 4 | #![feature(if_let_guard)] 5 | #![feature(iter_order_by)] 6 | 7 | pub mod block; 8 | pub mod dot; 9 | pub mod function; 10 | pub mod pattern; 11 | pub mod ssa; 12 | -------------------------------------------------------------------------------- /lua51-deserializer/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "lua51-deserializer" 3 | version = "0.1.0" 4 | edition.workspace = true 5 | authors.workspace = true 6 | 7 | [dependencies] 8 | nom = "7.1.1" 9 | num-traits = "0.2.15" 10 | num-derive = "0.3.3" 11 | either = "1.8.0" 12 | enum-as-inner = "0.5.1" 13 | strum_macros = "0.24.3" -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | Medal's LuaU decompiler 2 | 3 | All credits to this project goes to in honor and memory of: 4 | Jujhar Singh (KowalskiFX) 5 | Mathias Pedersen (Costomality) 6 | 7 | While details of how they passed and our relationship with them are completely irrelevant its better if their legacy 8 | does not go in vain. 9 | 10 | Keep the Singh and Pedersen family in you guys prayers. 11 | We love you both. 12 | -------------------------------------------------------------------------------- /script.lua: -------------------------------------------------------------------------------- 1 | function medaldecom(script) 2 | httpresponse = request({ 3 | Url = "http://127.0.0.1:9002", 4 | Body = base64.encode(getscriptbytecode(script)), -- The base64 encoded bytecode 5 | Method = "POST", 6 | Headers = { 7 | ["Content-Type"] = "text/plain" 8 | }, 9 | }) 10 | return httpresponse.Body 11 | end 12 | getgenv().decompile = medaldecom 13 | -------------------------------------------------------------------------------- /ast/src/break.rs: -------------------------------------------------------------------------------- 1 | use std::fmt; 2 | 3 | use crate::{has_side_effects, LocalRw, Traverse}; 4 | 5 | #[derive(Debug, PartialEq, Eq, Clone)] 6 | pub struct Break {} 7 | 8 | has_side_effects!(Break); 9 | 10 | impl LocalRw for Break {} 11 | 12 | impl Traverse for Break {} 13 | 14 | impl fmt::Display for Break { 15 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 16 | write!(f, "break") 17 | } 18 | } 19 | -------------------------------------------------------------------------------- /ast/src/vararg.rs: -------------------------------------------------------------------------------- 1 | use std::fmt; 2 | 3 | use crate::{LocalRw, SideEffects, Traverse}; 4 | 5 | #[derive(Debug, Clone, PartialEq, Eq)] 6 | pub struct VarArg; 7 | 8 | impl LocalRw for VarArg {} 9 | 10 | impl SideEffects for VarArg {} 11 | 12 | impl Traverse for VarArg {} 13 | 14 | impl fmt::Display for VarArg { 15 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 16 | write!(f, "...") 17 | } 18 | } 19 | -------------------------------------------------------------------------------- /ast/src/continue.rs: -------------------------------------------------------------------------------- 1 | use std::fmt; 2 | 3 | use crate::{has_side_effects, LocalRw, Traverse}; 4 | 5 | #[derive(Debug, PartialEq, Eq, Clone)] 6 | pub struct Continue {} 7 | 8 | has_side_effects!(Continue); 9 | 10 | impl LocalRw for Continue {} 11 | 12 | impl Traverse for Continue {} 13 | 14 | impl fmt::Display for Continue { 15 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 16 | write!(f, "continue") 17 | } 18 | } 19 | -------------------------------------------------------------------------------- /luau-lifter/src/main.rs: -------------------------------------------------------------------------------- 1 | fn main() { 2 | let file_name = std::env::args().nth(1).expect("expected exactly one file"); 3 | let key = std::env::args() 4 | .nth(2) 5 | .or_else(|| None) 6 | .map(|s| if s == "-e" { 203 } else { panic!() }) 7 | .unwrap_or(1); 8 | let bytecode = std::fs::read(file_name).expect("failed to read file"); 9 | println!("{}", luau_lifter::decompile_bytecode(&bytecode, key)); 10 | } 11 | -------------------------------------------------------------------------------- /decomp-server/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "decomp-server" 3 | version = "0.0.1" 4 | edition = "2021" 5 | 6 | [dependencies] 7 | tokio = { version = "1", features = ["full"] } 8 | hyper = { version = "0.14", features = ["full"] } 9 | base64 = "0.13" 10 | luau-lifter = { path = "../luau-lifter" } 11 | serde = { version = "1.0", features = ["derive"] } 12 | serde_json = "1.0" 13 | tokio-tungstenite = "*" 14 | futures-util = "*" 15 | log = "*" 16 | env_logger = "*" 17 | -------------------------------------------------------------------------------- /luau-worker/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "luau-worker" 3 | version = "0.1.0" 4 | edition.workspace = true 5 | authors.workspace = true 6 | 7 | [package.metadata.wasm-pack.profile.release] 8 | wasm-opt = false 9 | 10 | [dependencies] 11 | console_error_panic_hook = "0.1.7" 12 | worker = "0.3.2" 13 | futures-util = "0.3.30" 14 | luau-lifter = { path = "../luau-lifter" } 15 | base64 = "0.22.1" 16 | chrono = "0.4.38" 17 | serde_json = "1.0.117" 18 | serde = "1.0.202" 19 | 20 | [lib] 21 | crate-type = ["cdylib"] 22 | -------------------------------------------------------------------------------- /restructure/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "restructure" 3 | version = "0.1.0" 4 | edition.workspace = true 5 | authors.workspace = true 6 | 7 | [dependencies] 8 | petgraph = { git = "https://github.com/jujhar16/petgraph.git", branch="ensure_len_resize_with" } 9 | ast = { path="../ast" } 10 | rustc-hash = "1.1.0" 11 | derive_more = "0.99.17" 12 | enum-as-inner = "0.5.1" 13 | itertools = "0.10.5" 14 | array_tool = "1.0.3" 15 | tuple = "0.5.1" 16 | cfg = { path = "../cfg" } 17 | triomphe = "0.1.8" 18 | parking_lot = "0.12.1" -------------------------------------------------------------------------------- /ast/src/close.rs: -------------------------------------------------------------------------------- 1 | use itertools::Itertools; 2 | 3 | use crate::{LocalRw, RcLocal, SideEffects, Traverse}; 4 | 5 | #[derive(Debug, Clone, PartialEq, Eq)] 6 | pub struct Close { 7 | pub locals: Vec, 8 | } 9 | 10 | impl std::fmt::Display for Close { 11 | fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { 12 | write!(f, "__close_uv({})", self.locals.iter().join(", ")) 13 | } 14 | } 15 | 16 | impl LocalRw for Close {} 17 | impl SideEffects for Close {} 18 | impl Traverse for Close {} 19 | -------------------------------------------------------------------------------- /ast/src/side_effects.rs: -------------------------------------------------------------------------------- 1 | use enum_dispatch::enum_dispatch; 2 | 3 | #[enum_dispatch] 4 | pub trait SideEffects { 5 | fn has_side_effects(&self) -> bool { 6 | false 7 | } 8 | } 9 | 10 | macro_rules! has_side_effects { 11 | ($($name:ty),*) => { 12 | $( 13 | impl $crate::SideEffects for $name { 14 | fn has_side_effects(&self) -> bool { 15 | true 16 | } 17 | } 18 | )* 19 | }; 20 | } 21 | 22 | pub(crate) use has_side_effects; 23 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [workspace] 2 | members = [ 3 | "cfg", 4 | "ast", 5 | "lua51-lifter", 6 | "lua51-deserializer", 7 | "luau-lifter", 8 | "restructure", 9 | "luau-worker", 10 | "decomp-server" 11 | ] 12 | 13 | [workspace.package] 14 | edition = "2021" 15 | authors = [ 16 | "Mathias Pedersen ", 17 | "Jujhar Singh ", 18 | "Amit Bashan ", 19 | ] 20 | 21 | [profile.release] 22 | debug = true 23 | lto = true 24 | # strip = true 25 | # panic = "abort" 26 | -------------------------------------------------------------------------------- /ast/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ast" 3 | version = "0.1.0" 4 | edition.workspace = true 5 | authors.workspace = true 6 | 7 | [dependencies] 8 | derive_more = "0.99.17" 9 | enum-as-inner = "0.5.1" 10 | itertools = "0.10.5" 11 | enum_dispatch = "0.3.8" 12 | by_address = "1.1.0" 13 | rustc-hash = "1.1.0" 14 | indexmap = "1.9.1" 15 | petgraph = { git = "https://github.com/jujhar16/petgraph.git", branch = "ensure_len_resize_with" } 16 | array_tool = "1.0.3" 17 | itoa = "1.0.4" 18 | ryu = "1.0.11" 19 | nohash-hasher = "0.2.0" 20 | triomphe = "0.1.8" 21 | parking_lot = "0.12.1" -------------------------------------------------------------------------------- /cfg/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "cfg" 3 | version = "0.1.0" 4 | edition.workspace = true 5 | authors.workspace = true 6 | 7 | [dependencies] 8 | thiserror = "1.0.37" 9 | enum_dispatch = "0.3.8" 10 | enum-as-inner = "0.5.1" 11 | petgraph = { git = "https://github.com/jujhar16/petgraph.git", branch = "ensure_len_resize_with" } 12 | indexmap = "1.9.1" 13 | ast = { path = "../ast" } 14 | dot = { version = "0.1.4" } 15 | rustc-hash = "1.1.0" 16 | itertools = "0.10.5" 17 | contracts = "0.6.3" 18 | hashlink = "0.8.1" 19 | array_tool = "1.0.3" 20 | rangemap = "1.0.3" 21 | tuple = "0.5.1" 22 | -------------------------------------------------------------------------------- /luau-worker/wrangler.toml: -------------------------------------------------------------------------------- 1 | name = "luau-worker" 2 | type = "rust" 3 | workers_dev = true 4 | compatibility_date = "2022-01-20" 5 | 6 | [vars] 7 | WORKERS_RS_VERSION = "0.0.9" 8 | 9 | [build] 10 | command = "cargo install -q worker-build && worker-build --release" # required 11 | 12 | [build.upload] 13 | dir = "build/worker" 14 | format = "modules" 15 | main = "./shim.mjs" 16 | 17 | [[build.upload.rules]] 18 | globs = ["**/*.wasm"] 19 | type = "CompiledWasm" 20 | 21 | # read more about configuring your Worker via wrangler.toml at: 22 | # https://developers.cloudflare.com/workers/cli-wrangler/configuration -------------------------------------------------------------------------------- /luau-lifter/src/deserializer/list.rs: -------------------------------------------------------------------------------- 1 | use nom::{multi::count, IResult}; 2 | use nom_leb128::leb128_usize; 3 | 4 | pub(crate) fn parse_list<'a, T>( 5 | input: &'a [u8], 6 | parser: impl Fn(&'a [u8]) -> IResult<&'a [u8], T>, 7 | ) -> IResult<&'a [u8], Vec> { 8 | let (input, length) = leb128_usize(input)?; 9 | let (input, items) = count(parser, length)(input)?; 10 | Ok((input, items)) 11 | } 12 | 13 | pub(crate) fn parse_list_len<'a, T>( 14 | input: &'a [u8], 15 | parser: impl Fn(&'a [u8]) -> IResult<&'a [u8], T>, 16 | length: usize, 17 | ) -> IResult<&'a [u8], Vec> { 18 | let (input, items) = count(parser, length)(input)?; 19 | Ok((input, items)) 20 | } 21 | -------------------------------------------------------------------------------- /ast/src/name_gen.rs: -------------------------------------------------------------------------------- 1 | use super::{Literal, RValue}; 2 | 3 | pub trait NameGenerator { 4 | fn generate_name(&self, rvalue: &RValue, identifier: usize) -> Option; 5 | } 6 | 7 | pub struct DefaultNameGenerator {} 8 | 9 | impl NameGenerator for DefaultNameGenerator { 10 | fn generate_name(&self, rvalue: &RValue, identifier: usize) -> Option { 11 | let hint = match rvalue { 12 | RValue::Global(global) => Some(global.to_string()), 13 | RValue::Index(index) => match &*index.right { 14 | RValue::Literal(Literal::String(string)) => Some(string.to_string()), 15 | _ => None, 16 | }, 17 | _ => None, 18 | }; 19 | hint.map(|hint| format!("v_{}_{}", identifier, hint)) 20 | } 21 | } 22 | -------------------------------------------------------------------------------- /lua51-deserializer/src/instruction/position.rs: -------------------------------------------------------------------------------- 1 | use nom::{multi::count, number::complete::le_u32, IResult}; 2 | 3 | #[derive(Debug)] 4 | pub struct Position { 5 | pub instruction: usize, 6 | pub source: u32, 7 | } 8 | 9 | impl Position { 10 | pub fn parse(input: &[u8]) -> IResult<&[u8], Vec> { 11 | let (input, positions_length) = le_u32(input)?; 12 | let (input, source_positions) = count(le_u32, positions_length as usize)(input)?; 13 | 14 | Ok(( 15 | input, 16 | source_positions 17 | .iter() 18 | .enumerate() 19 | .map(|(instruction, &source)| Self { 20 | instruction, 21 | source, 22 | }) 23 | .collect(), 24 | )) 25 | } 26 | } 27 | -------------------------------------------------------------------------------- /lua51-deserializer/src/instruction/argument.rs: -------------------------------------------------------------------------------- 1 | use either::Either; 2 | 3 | #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] 4 | pub struct Register(pub u8); 5 | 6 | impl From for Register { 7 | fn from(value: u8) -> Self { 8 | Self(value) 9 | } 10 | } 11 | 12 | #[derive(Debug, Copy, Clone)] 13 | pub struct Constant(pub u32); 14 | 15 | #[derive(Debug, Copy, Clone)] 16 | pub struct RegisterOrConstant(pub Either); 17 | 18 | impl From for RegisterOrConstant { 19 | fn from(value: u32) -> Self { 20 | Self(if value > 255 { 21 | Either::Right(Constant(value - 256)) 22 | } else { 23 | Either::Left(Register(value as u8)) 24 | }) 25 | } 26 | } 27 | 28 | #[derive(Debug, Clone)] 29 | pub struct Upvalue(pub u8); 30 | 31 | #[derive(Debug, Clone)] 32 | pub struct Function(pub u32); 33 | -------------------------------------------------------------------------------- /lua51-deserializer/src/local.rs: -------------------------------------------------------------------------------- 1 | use std::ops::Range; 2 | 3 | use nom::{multi::count, number::complete::le_u32, IResult}; 4 | 5 | use crate::value::parse_string; 6 | 7 | #[derive(Debug)] 8 | pub struct Local<'a> { 9 | pub name: &'a [u8], 10 | pub range: Range, 11 | } 12 | 13 | impl<'a> Local<'a> { 14 | pub fn parse_list(input: &'a [u8]) -> IResult<&'a [u8], Vec> { 15 | let (input, length) = le_u32(input)?; 16 | 17 | count(Self::parse, length as usize)(input) 18 | } 19 | 20 | fn parse(input: &'a [u8]) -> IResult<&'a [u8], Self> { 21 | let (input, name) = parse_string(input)?; 22 | let (input, start) = le_u32(input)?; 23 | let (input, end) = le_u32(input)?; 24 | 25 | Ok(( 26 | input, 27 | Self { 28 | name: &name[..name.len() - 1], 29 | range: (start..end), 30 | }, 31 | )) 32 | } 33 | } 34 | -------------------------------------------------------------------------------- /luau-lifter/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "luau-lifter" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | num_enum = "0.5.6" 10 | nom = "7.1.0" 11 | nom-leb128 = "0.2.0" 12 | clap = { version = "4.0.26", features = ["derive"] } 13 | anyhow = { version = "1.0.53", features = ["backtrace"] } 14 | cfg = { path = "../cfg" } 15 | ast = { path = "../ast" } 16 | rustc-hash = "1.1.0" 17 | dhat = "0.3.1" 18 | either = "1.6.1" 19 | petgraph = { git = "https://github.com/jujhar16/petgraph.git", branch = "ensure_len_resize_with" } 20 | restructure = { path = "../restructure" } 21 | lazy_static = "1.4.0" 22 | itertools = "0.10.5" 23 | indexmap = "1.9.1" 24 | by_address = "1.1.0" 25 | rayon = "1.5.3" 26 | triomphe = "0.1.8" 27 | parking_lot = "0.12.1" 28 | walkdir = "2.3.2" 29 | 30 | [features] 31 | dhat-heap = [] 32 | panic-handled = [] -------------------------------------------------------------------------------- /lua51-lifter/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "lua51-lifter" 3 | version = "0.1.0" 4 | edition.workspace = true 5 | authors.workspace = true 6 | 7 | [dependencies] 8 | num_enum = "0.5.7" 9 | nom = "7.1.1" 10 | clap = { version = "4.0.10", features = ["derive"] } 11 | anyhow = { version = "1.0.65", features = ["backtrace"] } 12 | cfg = { path = "../cfg" } 13 | lua51-deserializer = { path = "../lua51-deserializer" } 14 | # graph = { path = "../graph", features = ["dot"] } 15 | petgraph = { git = "https://github.com/jujhar16/petgraph.git", branch="ensure_len_resize_with" } 16 | indexmap = "1.9.1" 17 | ast = { path = "../ast" } 18 | dhat = "0.3.1" 19 | rustc-hash = "1.1.0" 20 | either = "1.8.0" 21 | restructure = { path = "../restructure" } 22 | enum-as-inner = "0.5.1" 23 | itertools = "0.10.5" 24 | by_address = "1.1.0" 25 | rayon = "1.5.3" 26 | triomphe = "0.1.8" 27 | parking_lot = "0.12.1" 28 | 29 | [features] 30 | dhat-heap = [] 31 | panic-handled = [] -------------------------------------------------------------------------------- /ast/src/paste.txt: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | def search_files_for_string(search_string, root_dir='.'): 4 | # Traverse the directory 5 | for dirpath, _, filenames in os.walk(root_dir): 6 | for filename in filenames: 7 | # Construct the full file path 8 | file_path = os.path.join(dirpath, filename) 9 | 10 | try: 11 | # Open the file and read its content 12 | with open(file_path, 'r', encoding='utf-8') as file: 13 | # Check if the search string is in the file 14 | content = file.read() 15 | if search_string in content: 16 | print(f"Found '{search_string}' in: {file_path}") 17 | except (UnicodeDecodeError, PermissionError): 18 | # Skip files that can't be opened (binary files, permission issues) 19 | continue 20 | 21 | # Call the function to search for 'v_' 22 | search_files_for_string('v_') 23 | -------------------------------------------------------------------------------- /luau-lifter/src/deserializer/bytecode.rs: -------------------------------------------------------------------------------- 1 | use nom::{bytes::complete::take, number::complete::le_u8, IResult}; 2 | 3 | use super::chunk::Chunk; 4 | 5 | #[derive(Debug)] 6 | pub enum Bytecode { 7 | Error(String), 8 | Chunk(Chunk), 9 | } 10 | 11 | impl Bytecode { 12 | pub fn parse(input: &[u8], encode_key: u8) -> IResult<&[u8], Bytecode> { 13 | let (input, status_code) = le_u8(input)?; 14 | match status_code { 15 | 0 => { 16 | let (input, error_msg) = take(input.len())(input)?; 17 | Ok(( 18 | input, 19 | Bytecode::Error(String::from_utf8_lossy(error_msg).to_string()), 20 | )) 21 | } 22 | 4..=6 => { 23 | let (input, chunk) = Chunk::parse(input, encode_key, status_code)?; 24 | Ok((input, Bytecode::Chunk(chunk))) 25 | } 26 | _ => panic!("Unsupported bytecode version: {}", status_code), 27 | } 28 | } 29 | } 30 | -------------------------------------------------------------------------------- /luau-lifter/src/deserializer/mod.rs: -------------------------------------------------------------------------------- 1 | use nom::{bytes::complete::take, IResult}; 2 | use nom_leb128::leb128_usize; 3 | 4 | pub mod bytecode; 5 | pub mod chunk; 6 | pub mod constant; 7 | pub mod function; 8 | mod list; 9 | 10 | fn parse_string(input: &[u8]) -> IResult<&[u8], Vec> { 11 | let (input, length) = leb128_usize(input)?; 12 | let (input, bytes) = take(length)(input)?; 13 | Ok((input, bytes.to_owned())) 14 | } 15 | 16 | pub fn deserialize(bytecode: &[u8], encode_key: u8) -> Result { 17 | match bytecode::Bytecode::parse(bytecode, encode_key) { 18 | Ok((_, deserialized_bytecode)) => Ok(deserialized_bytecode), 19 | Err(err) => Err(err.to_string()), 20 | } 21 | } 22 | 23 | /*#[test] 24 | fn main() -> anyhow::Result<()> { 25 | let compiler = Compiler::new() 26 | .set_debug_level(1).set_optimization_level(2); 27 | let bytecode = compiler.compile("asd = test"); 28 | println!("{:#?}", bytecode); 29 | let deserialized = deserialize(&bytecode); 30 | println!("{:#?}", deserialized); 31 | Ok(()) 32 | }*/ 33 | -------------------------------------------------------------------------------- /ast/src/global.rs: -------------------------------------------------------------------------------- 1 | use derive_more::From; 2 | use std::fmt; 3 | 4 | use crate::{formatter::Formatter, LocalRw, SideEffects, Traverse}; 5 | 6 | #[derive(Debug, From, PartialEq, Eq, PartialOrd, Clone)] 7 | pub struct Global(pub Vec); 8 | 9 | impl Global { 10 | pub fn new(name: Vec) -> Self { 11 | Self(name) 12 | } 13 | } 14 | 15 | impl LocalRw for Global {} 16 | 17 | impl SideEffects for Global { 18 | fn has_side_effects(&self) -> bool { 19 | true 20 | } 21 | } 22 | 23 | impl Traverse for Global {} 24 | 25 | impl<'a> From<&'a str> for Global { 26 | fn from(name: &'a str) -> Self { 27 | Self::new(name.into()) 28 | } 29 | } 30 | 31 | impl fmt::Display for Global { 32 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 33 | if Formatter::::is_valid_name(&self.0) { 34 | write!(f, "{}", std::str::from_utf8(&self.0).unwrap()) 35 | } else { 36 | write!( 37 | f, 38 | "__FENV[\"{}\"]", 39 | Formatter::::escape_string(&self.0) 40 | ) 41 | } 42 | } 43 | } 44 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Jujhar Singh, Mathias Pedersen 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /lua51-deserializer/src/chunk/mod.rs: -------------------------------------------------------------------------------- 1 | use std::mem; 2 | 3 | use nom::IResult; 4 | 5 | pub use header::Header; 6 | 7 | use crate::{ 8 | chunk::header::{Endianness, Format}, 9 | function::Function, 10 | }; 11 | 12 | pub mod header; 13 | 14 | #[derive(Debug)] 15 | pub struct Chunk<'a> { 16 | pub function: Function<'a>, 17 | } 18 | 19 | impl<'a> Chunk<'a> { 20 | pub fn parse(input: &'a [u8]) -> IResult<&[u8], Self> { 21 | let (input, header) = Header::parse(input)?; 22 | // TODO: pass header to Function::parse 23 | assert_eq!(header.version_number, 0x51); 24 | assert_eq!(header.format, Format::Official); 25 | assert_eq!(header.endianness, Endianness::Little); 26 | assert_eq!(header.int_width as usize, mem::size_of::()); 27 | assert_eq!(header.size_t_width as usize, mem::size_of::()); 28 | assert_eq!(header.instr_width as usize, mem::size_of::()); 29 | assert_eq!(header.number_width as usize, mem::size_of::()); 30 | assert!(!header.number_is_integral); 31 | let (input, function) = Function::parse(input)?; 32 | 33 | Ok((input, Self { function })) 34 | } 35 | } 36 | -------------------------------------------------------------------------------- /ast/src/goto.rs: -------------------------------------------------------------------------------- 1 | use std::fmt; 2 | 3 | use crate::{has_side_effects, LocalRw, SideEffects, Traverse}; 4 | 5 | // TODO: Rc 6 | #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] 7 | pub struct Label(pub String); 8 | 9 | impl SideEffects for Label {} 10 | 11 | impl From<&str> for Label { 12 | fn from(str: &str) -> Self { 13 | Label(str.into()) 14 | } 15 | } 16 | 17 | impl From for Label { 18 | fn from(str: String) -> Self { 19 | Label(str) 20 | } 21 | } 22 | 23 | impl LocalRw for Label {} 24 | 25 | impl Traverse for Label {} 26 | 27 | impl fmt::Display for Label { 28 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 29 | write!(f, "::{}::", self.0) 30 | } 31 | } 32 | 33 | #[derive(Debug, PartialEq, Eq, Clone)] 34 | pub struct Goto(pub Label); 35 | 36 | impl Traverse for Goto {} 37 | 38 | has_side_effects!(Goto); 39 | 40 | impl Goto { 41 | pub fn new(label: Label) -> Self { 42 | Self(label) 43 | } 44 | } 45 | 46 | impl LocalRw for Goto {} 47 | 48 | impl fmt::Display for Goto { 49 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 50 | write!(f, "goto {}", self.0 .0) 51 | } 52 | } 53 | -------------------------------------------------------------------------------- /cfg/src/block.rs: -------------------------------------------------------------------------------- 1 | use std::fmt; 2 | 3 | use ast::{RValue, RcLocal}; 4 | 5 | #[derive(Debug, Clone, PartialEq, Eq, Default)] 6 | pub enum BranchType { 7 | #[default] 8 | Unconditional, 9 | Then, 10 | Else, 11 | } 12 | 13 | #[derive(Debug, Clone, Default)] 14 | pub struct BlockEdge { 15 | pub branch_type: BranchType, 16 | // TODO: why is this not a hash map? 17 | pub arguments: Vec<(RcLocal, RValue)>, 18 | } 19 | 20 | impl BlockEdge { 21 | pub fn new(branch_type: BranchType) -> Self { 22 | Self { 23 | branch_type, 24 | arguments: Vec::new(), 25 | } 26 | } 27 | } 28 | 29 | impl fmt::Display for BlockEdge { 30 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 31 | match self.branch_type { 32 | BranchType::Unconditional => write!(f, "u"), 33 | BranchType::Then => write!(f, "t"), 34 | BranchType::Else => write!(f, "e"), 35 | }?; 36 | if !self.arguments.is_empty() { 37 | for (i, (local, new_local)) in self.arguments.iter().enumerate() { 38 | write!(f, "{} -> {}", local, new_local)?; 39 | if i + 1 != self.arguments.len() { 40 | writeln!(f)?; 41 | } 42 | } 43 | } 44 | Ok(()) 45 | } 46 | } 47 | -------------------------------------------------------------------------------- /ast/src/return.rs: -------------------------------------------------------------------------------- 1 | use std::fmt; 2 | 3 | use crate::{formatter::Formatter, has_side_effects, LocalRw, RcLocal, Traverse}; 4 | 5 | use super::RValue; 6 | 7 | #[derive(Debug, PartialEq, Clone, Default)] 8 | pub struct Return { 9 | pub values: Vec, 10 | } 11 | 12 | has_side_effects!(Return); 13 | 14 | impl Return { 15 | pub fn new(values: Vec) -> Self { 16 | Self { values } 17 | } 18 | } 19 | 20 | impl Traverse for Return { 21 | fn rvalues_mut(&mut self) -> Vec<&mut RValue> { 22 | self.values.iter_mut().collect() 23 | } 24 | 25 | fn rvalues(&self) -> Vec<&RValue> { 26 | self.values.iter().collect() 27 | } 28 | } 29 | 30 | impl LocalRw for Return { 31 | fn values_read(&self) -> Vec<&RcLocal> { 32 | self.values.iter().flat_map(|r| r.values_read()).collect() 33 | } 34 | 35 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 36 | self.values 37 | .iter_mut() 38 | .flat_map(|r| r.values_read_mut()) 39 | .collect() 40 | } 41 | } 42 | 43 | impl fmt::Display for Return { 44 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 45 | Formatter { 46 | indentation_level: 0, 47 | indentation_mode: Default::default(), 48 | output: f, 49 | } 50 | .format_return(self) 51 | } 52 | } 53 | -------------------------------------------------------------------------------- /luau-lifter/src/deserializer/chunk.rs: -------------------------------------------------------------------------------- 1 | use super::{function::Function, list::parse_list, parse_string}; 2 | use nom::character::complete::char; 3 | use nom::multi::many_till; 4 | use nom::number::complete::le_u8; 5 | use nom::IResult; 6 | use nom_leb128::leb128_usize; 7 | 8 | #[derive(Debug)] 9 | pub struct Chunk { 10 | pub string_table: Vec>, 11 | pub functions: Vec, 12 | pub main: usize, 13 | } 14 | 15 | impl Chunk { 16 | pub(crate) fn parse(input: &[u8], encode_key: u8, version: u8) -> IResult<&[u8], Self> { 17 | let (input, types_version) = if version >= 4 { 18 | le_u8(input)? 19 | } else { 20 | (input, 0) 21 | }; 22 | if types_version > 3 { 23 | panic!("unsupported types version"); 24 | } 25 | let (input, string_table) = parse_list(input, parse_string)?; 26 | let input = if types_version == 3 { 27 | many_till(leb128_usize, char('\0'))(input)?.0 28 | } else { 29 | input 30 | }; 31 | let (input, functions) = parse_list(input, |i| Function::parse(i, encode_key))?; 32 | let (input, main) = leb128_usize(input)?; 33 | 34 | Ok(( 35 | input, 36 | Self { 37 | string_table, 38 | functions, 39 | main, 40 | }, 41 | )) 42 | } 43 | } 44 | -------------------------------------------------------------------------------- /ast/src/while.rs: -------------------------------------------------------------------------------- 1 | use parking_lot::Mutex; 2 | use triomphe::Arc; 3 | 4 | use crate::{formatter::Formatter, has_side_effects, Block, LocalRw, RValue, RcLocal, Traverse}; 5 | use std::fmt; 6 | 7 | #[derive(Debug, Clone)] 8 | pub struct While { 9 | pub condition: RValue, 10 | pub block: Arc>, 11 | } 12 | 13 | impl PartialEq for While { 14 | fn eq(&self, _other: &Self) -> bool { 15 | // TODO: compare block 16 | false 17 | } 18 | } 19 | 20 | has_side_effects!(While); 21 | 22 | impl While { 23 | pub fn new(condition: RValue, block: Block) -> Self { 24 | Self { 25 | condition, 26 | block: Arc::new(block.into()), 27 | } 28 | } 29 | } 30 | 31 | impl Traverse for While { 32 | fn rvalues_mut(&mut self) -> Vec<&mut RValue> { 33 | vec![&mut self.condition] 34 | } 35 | 36 | fn rvalues(&self) -> Vec<&RValue> { 37 | vec![&self.condition] 38 | } 39 | } 40 | 41 | impl LocalRw for While { 42 | fn values_read(&self) -> Vec<&RcLocal> { 43 | self.condition.values_read() 44 | } 45 | 46 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 47 | self.condition.values_read_mut() 48 | } 49 | } 50 | 51 | impl fmt::Display for While { 52 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 53 | Formatter { 54 | indentation_level: 0, 55 | indentation_mode: Default::default(), 56 | output: f, 57 | } 58 | .format_while(self) 59 | } 60 | } 61 | -------------------------------------------------------------------------------- /ast/src/repeat.rs: -------------------------------------------------------------------------------- 1 | use parking_lot::Mutex; 2 | use triomphe::Arc; 3 | 4 | use crate::{formatter::Formatter, has_side_effects, Block, LocalRw, RValue, RcLocal, Traverse}; 5 | use std::fmt; 6 | 7 | // TODO: move condition after block 8 | #[derive(Debug, Clone)] 9 | pub struct Repeat { 10 | pub condition: RValue, 11 | pub block: Arc>, 12 | } 13 | 14 | impl PartialEq for Repeat { 15 | fn eq(&self, _other: &Self) -> bool { 16 | // TODO: compare block 17 | false 18 | } 19 | } 20 | 21 | has_side_effects!(Repeat); 22 | 23 | impl Repeat { 24 | pub fn new(condition: RValue, block: Block) -> Self { 25 | Self { 26 | condition, 27 | block: Arc::new(block.into()), 28 | } 29 | } 30 | } 31 | 32 | impl Traverse for Repeat { 33 | fn rvalues_mut(&mut self) -> Vec<&mut RValue> { 34 | vec![&mut self.condition] 35 | } 36 | 37 | fn rvalues(&self) -> Vec<&RValue> { 38 | vec![&self.condition] 39 | } 40 | } 41 | 42 | impl LocalRw for Repeat { 43 | fn values_read(&self) -> Vec<&RcLocal> { 44 | self.condition.values_read() 45 | } 46 | 47 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 48 | self.condition.values_read_mut() 49 | } 50 | } 51 | 52 | impl fmt::Display for Repeat { 53 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 54 | Formatter { 55 | indentation_level: 0, 56 | indentation_mode: Default::default(), 57 | output: f, 58 | } 59 | .format_repeat(self) 60 | } 61 | } 62 | -------------------------------------------------------------------------------- /ast/src/index.rs: -------------------------------------------------------------------------------- 1 | use crate::{formatter::Formatter, has_side_effects, LocalRw, RcLocal, Traverse}; 2 | 3 | use super::RValue; 4 | use std::fmt; 5 | 6 | #[derive(Debug, Clone, PartialEq)] 7 | pub struct Index { 8 | pub left: Box, 9 | pub right: Box, 10 | } 11 | 12 | // this should be the same as MethodCall 13 | has_side_effects!(Index); 14 | 15 | impl Index { 16 | pub fn new(left: RValue, right: RValue) -> Self { 17 | Self { 18 | left: Box::new(left), 19 | right: Box::new(right), 20 | } 21 | } 22 | } 23 | 24 | impl LocalRw for Index { 25 | fn values_read(&self) -> Vec<&RcLocal> { 26 | self.left 27 | .values_read() 28 | .into_iter() 29 | .chain(self.right.values_read().into_iter()) 30 | .collect() 31 | } 32 | 33 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 34 | self.left 35 | .values_read_mut() 36 | .into_iter() 37 | .chain(self.right.values_read_mut().into_iter()) 38 | .collect() 39 | } 40 | } 41 | 42 | impl Traverse for Index { 43 | fn rvalues_mut(&mut self) -> Vec<&mut RValue> { 44 | vec![&mut self.left, &mut self.right] 45 | } 46 | 47 | fn rvalues(&self) -> Vec<&RValue> { 48 | vec![&self.left, &self.right] 49 | } 50 | } 51 | 52 | impl fmt::Display for Index { 53 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 54 | Formatter { 55 | indentation_level: 0, 56 | indentation_mode: Default::default(), 57 | output: f, 58 | } 59 | .format_index(self) 60 | } 61 | } 62 | -------------------------------------------------------------------------------- /decomp-server/src/main.rs: -------------------------------------------------------------------------------- 1 | use futures_util::StreamExt; 2 | use hyper::service::{make_service_fn, service_fn}; 3 | use hyper::{Body, Request, Response, Server}; 4 | use log::*; 5 | use serde::Deserialize; 6 | use std::convert::Infallible; 7 | use std::net::SocketAddr; 8 | use tokio::net::TcpListener; 9 | use tokio::sync::mpsc; 10 | use base64::decode; 11 | use luau_lifter::decompile_bytecode; 12 | 13 | async fn decompiler(base64_b: String) -> Result { 14 | match decode(base64_b.as_bytes()) { 15 | Ok(bytecode) => { 16 | let decompiled = decompile_bytecode(&bytecode, 203); 17 | Ok(format!("{}", decompiled)) 18 | } 19 | Err(e) => Err(format!("Failed to decode base64 bytecode: {}", e)), 20 | } 21 | } 22 | 23 | async fn handle_request(req: Request) -> Result, Infallible> { 24 | let bytes = hyper::body::to_bytes(req.into_body()).await.unwrap(); 25 | let base64_bytecode = String::from_utf8(bytes.to_vec()).unwrap(); 26 | 27 | match decompiler(base64_bytecode).await { 28 | Ok(decompiled_code) => Ok(Response::new(Body::from(decompiled_code))), 29 | Err(e) => Ok(Response::new(Body::from(e))), 30 | } 31 | } 32 | 33 | #[tokio::main] 34 | async fn main() { 35 | env_logger::init(); 36 | 37 | let addr = SocketAddr::from(([127, 0, 0, 1], 9002)); 38 | 39 | let make_svc = make_service_fn(|_conn| { 40 | async { 41 | Ok::<_, Infallible>(service_fn(handle_request)) 42 | } 43 | }); 44 | 45 | let server = Server::bind(&addr).serve(make_svc); 46 | 47 | println!("Listening on http://{}", addr); 48 | 49 | if let Err(e) = server.await { 50 | eprintln!("Server error: {}", e); 51 | } 52 | } 53 | -------------------------------------------------------------------------------- /ast/src/if.rs: -------------------------------------------------------------------------------- 1 | use parking_lot::Mutex; 2 | use triomphe::Arc; 3 | 4 | use crate::{formatter::Formatter, LocalRw, RcLocal, SideEffects, Traverse}; 5 | 6 | use super::{Block, RValue}; 7 | 8 | use std::fmt; 9 | 10 | #[derive(Debug, Clone)] 11 | pub struct If { 12 | pub condition: RValue, 13 | pub then_block: Arc>, 14 | pub else_block: Arc>, 15 | } 16 | 17 | impl PartialEq for If { 18 | fn eq(&self, _other: &Self) -> bool { 19 | // TODO: compare block 20 | false 21 | } 22 | } 23 | 24 | impl If { 25 | pub fn new(condition: RValue, then_block: Block, else_block: Block) -> Self { 26 | Self { 27 | condition, 28 | then_block: Arc::new(then_block.into()), 29 | else_block: Arc::new(else_block.into()), 30 | } 31 | } 32 | } 33 | 34 | impl Traverse for If { 35 | fn rvalues_mut(&mut self) -> Vec<&mut RValue> { 36 | vec![&mut self.condition] 37 | } 38 | 39 | fn rvalues(&self) -> Vec<&RValue> { 40 | vec![&self.condition] 41 | } 42 | } 43 | 44 | impl SideEffects for If { 45 | // TODO: side effects for blocks 46 | fn has_side_effects(&self) -> bool { 47 | true 48 | } 49 | } 50 | 51 | impl LocalRw for If { 52 | fn values_read(&self) -> Vec<&RcLocal> { 53 | self.condition.values_read() 54 | } 55 | 56 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 57 | self.condition.values_read_mut() 58 | } 59 | } 60 | 61 | impl fmt::Display for If { 62 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 63 | Formatter { 64 | indentation_level: 0, 65 | indentation_mode: Default::default(), 66 | output: f, 67 | } 68 | .format_if(self) 69 | } 70 | } 71 | -------------------------------------------------------------------------------- /ast/src/replace_locals.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | 3 | use itertools::Either; 4 | 5 | use crate::{Block, LocalRw, RValue, RcLocal, Statement, Traverse}; 6 | 7 | pub fn replace_locals( 8 | block: &mut Block, 9 | map: &HashMap, 10 | ) { 11 | for statement in &mut block.0 { 12 | for local in statement.values_read_mut() { 13 | if let Some(new_local) = map.get(local) { 14 | *local = new_local.clone(); 15 | } 16 | } 17 | for local in statement.values_written_mut() { 18 | if let Some(new_local) = map.get(local) { 19 | *local = new_local.clone(); 20 | } 21 | } 22 | // TODO: traverse_values 23 | statement.post_traverse_values(&mut |value| -> Option<()> { 24 | if let Either::Right(RValue::Closure(closure)) = value { 25 | replace_locals(&mut closure.function.lock().body, map) 26 | }; 27 | None 28 | }); 29 | match statement { 30 | Statement::If(r#if) => { 31 | replace_locals(&mut r#if.then_block.lock(), map); 32 | replace_locals(&mut r#if.else_block.lock(), map); 33 | } 34 | Statement::While(r#while) => { 35 | replace_locals(&mut r#while.block.lock(), map); 36 | } 37 | Statement::Repeat(repeat) => { 38 | replace_locals(&mut repeat.block.lock(), map); 39 | } 40 | Statement::NumericFor(numeric_for) => { 41 | replace_locals(&mut numeric_for.block.lock(), map); 42 | } 43 | Statement::GenericFor(generic_for) => { 44 | replace_locals(&mut generic_for.block.lock(), map); 45 | } 46 | _ => {} 47 | } 48 | } 49 | } 50 | -------------------------------------------------------------------------------- /lua51-deserializer/src/instruction/layout.rs: -------------------------------------------------------------------------------- 1 | use nom::{ 2 | error::{Error, ErrorKind, ParseError}, 3 | number::complete::le_u32, 4 | Err, IResult, 5 | }; 6 | use num_traits::FromPrimitive; 7 | use strum_macros::EnumDiscriminants; 8 | 9 | use super::OperationCode; 10 | 11 | #[derive(Debug, EnumDiscriminants)] 12 | pub enum Layout { 13 | BC { a: u8, b: u16, c: u16 }, 14 | // b extended 15 | BX { a: u8, b_x: u32 }, 16 | // b signed, extended 17 | BSx { a: u8, b_sx: i32 }, 18 | } 19 | 20 | impl Layout { 21 | pub fn parse(input: &[u8], operation_code: u8) -> IResult<&[u8], Self> { 22 | let (input, instruction) = le_u32(input)?; 23 | 24 | match OperationCode::from_u8(operation_code).map(|o: OperationCode| o.instruction_layout()) 25 | { 26 | Some(LayoutDiscriminants::BC) => { 27 | let a = ((instruction >> 6) & 0xFF) as u8; 28 | let c = ((instruction >> 14) & 0x1FF) as u16; 29 | let b = ((instruction >> 23) & 0x1FF) as u16; 30 | 31 | Ok((input, Self::BC { a, b, c })) 32 | } 33 | Some(LayoutDiscriminants::BX) => { 34 | let a = ((instruction >> 6) & 0xFF) as u8; 35 | let b_x = (instruction >> 14) & 0x3FFFF; 36 | 37 | Ok((input, Self::BX { a, b_x })) 38 | } 39 | Some(LayoutDiscriminants::BSx) => { 40 | let a = ((instruction >> 6) & 0xFF) as u8; 41 | let b_x = (instruction >> 14) & 0x3FFFF; 42 | // subtract maximum 18 bit signed int 43 | let b_sx = b_x as i32 - (((1 << 18) - 1) >> 1); 44 | 45 | Ok((input, Self::BSx { a, b_sx })) 46 | } 47 | _ => Err(Err::Failure(Error::from_error_kind( 48 | input, 49 | ErrorKind::Switch, 50 | ))), 51 | } 52 | } 53 | } 54 | -------------------------------------------------------------------------------- /lua51-deserializer/src/value.rs: -------------------------------------------------------------------------------- 1 | use enum_as_inner::EnumAsInner; 2 | use nom::{ 3 | bytes::complete::take, 4 | error::{Error, ErrorKind, ParseError}, 5 | multi::count, 6 | number::complete::{le_f64, le_u32, le_u8}, 7 | Err, IResult, 8 | }; 9 | 10 | #[derive(Debug, EnumAsInner)] 11 | pub enum Value<'a> { 12 | Nil, 13 | Boolean(bool), 14 | Number(f64), 15 | String(&'a [u8]), 16 | } 17 | 18 | impl<'a> Value<'a> { 19 | pub fn parse(input: &'a [u8]) -> IResult<&'a [u8], Self> { 20 | let (input, kind) = le_u8(input)?; 21 | 22 | match kind { 23 | 0 => Ok((input, Self::Nil)), 24 | 1 => { 25 | let (input, value) = le_u8(input)?; 26 | 27 | Ok((input, Self::Boolean(value != 0))) 28 | } 29 | 3 => { 30 | let (input, value) = le_f64(input)?; 31 | 32 | Ok((input, Self::Number(value))) 33 | } 34 | 4 => { 35 | let (input, value) = parse_string(input)?; 36 | 37 | // TODO: lua bytecode actually allows the string to be completely empty 38 | // it sets the type to string but gc to NULL 39 | // this probably causes some weird behavior 40 | assert!(!value.is_empty()); 41 | 42 | // exclude null terminator 43 | Ok((input, Self::String(&value[..value.len() - 1]))) 44 | } 45 | _ => Err(Err::Failure(Error::from_error_kind( 46 | input, 47 | ErrorKind::Switch, 48 | ))), 49 | } 50 | } 51 | } 52 | 53 | pub fn parse_string(input: &[u8]) -> IResult<&[u8], &[u8]> { 54 | let (input, string_length) = le_u32(input)?; 55 | take(string_length as usize)(input) 56 | } 57 | 58 | pub fn parse_strings(input: &[u8]) -> IResult<&[u8], Vec<&[u8]>> { 59 | let (input, string_count) = le_u32(input)?; 60 | let (input, strings) = count(parse_string, string_count as usize)(input)?; 61 | 62 | Ok((input, strings)) 63 | } 64 | -------------------------------------------------------------------------------- /ast/src/assign.rs: -------------------------------------------------------------------------------- 1 | use std::fmt; 2 | 3 | use crate::{formatter::Formatter, RcLocal, SideEffects, Traverse}; 4 | 5 | use super::{LValue, LocalRw, RValue}; 6 | 7 | #[derive(Debug, Clone, PartialEq)] 8 | pub struct Assign { 9 | pub left: Vec, 10 | pub right: Vec, 11 | pub prefix: bool, 12 | pub parallel: bool, 13 | } 14 | 15 | impl Assign { 16 | pub fn new(left: Vec, right: Vec) -> Self { 17 | Self { 18 | left, 19 | right, 20 | prefix: false, 21 | parallel: false, 22 | } 23 | } 24 | } 25 | 26 | impl Traverse for Assign { 27 | fn lvalues_mut(&mut self) -> Vec<&mut LValue> { 28 | self.left.iter_mut().collect() 29 | } 30 | 31 | fn rvalues_mut(&mut self) -> Vec<&mut RValue> { 32 | self.right.iter_mut().collect() 33 | } 34 | 35 | fn rvalues(&self) -> Vec<&RValue> { 36 | self.right.iter().collect() 37 | } 38 | } 39 | 40 | impl SideEffects for Assign { 41 | fn has_side_effects(&self) -> bool { 42 | self.right.iter().any(|r| r.has_side_effects()) 43 | || self.left.iter().any(|l| l.has_side_effects()) 44 | } 45 | } 46 | 47 | impl LocalRw for Assign { 48 | fn values_read(&self) -> Vec<&RcLocal> { 49 | self.left 50 | .iter() 51 | .flat_map(|l| l.values_read()) 52 | .chain(self.right.iter().flat_map(|r| r.values_read())) 53 | .collect() 54 | } 55 | 56 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 57 | self.left 58 | .iter_mut() 59 | .flat_map(|l| l.values_read_mut()) 60 | .chain(self.right.iter_mut().flat_map(|r| r.values_read_mut())) 61 | .collect() 62 | } 63 | 64 | fn values_written(&self) -> Vec<&RcLocal> { 65 | self.left.iter().flat_map(|l| l.values_written()).collect() 66 | } 67 | 68 | fn values_written_mut(&mut self) -> Vec<&mut RcLocal> { 69 | self.left 70 | .iter_mut() 71 | .flat_map(|l| l.values_written_mut()) 72 | .collect() 73 | } 74 | } 75 | 76 | impl fmt::Display for Assign { 77 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 78 | Formatter { 79 | indentation_level: 0, 80 | indentation_mode: Default::default(), 81 | output: f, 82 | } 83 | .format_assign(self) 84 | } 85 | } 86 | -------------------------------------------------------------------------------- /luau-lifter/src/deserializer/constant.rs: -------------------------------------------------------------------------------- 1 | use super::list::parse_list; 2 | use nom::{ 3 | number::complete::{le_f32, le_f64, le_u32, le_u8}, 4 | IResult, 5 | }; 6 | use nom_leb128::leb128_usize; 7 | 8 | const CONSTANT_NIL: u8 = 0; 9 | const CONSTANT_BOOLEAN: u8 = 1; 10 | const CONSTANT_NUMBER: u8 = 2; 11 | const CONSTANT_STRING: u8 = 3; 12 | const CONSTANT_IMPORT: u8 = 4; 13 | const CONSTANT_TABLE: u8 = 5; 14 | const CONSTANT_CLOSURE: u8 = 6; 15 | const CONSTANT_VECTOR: u8 = 7; 16 | 17 | #[derive(Debug)] 18 | pub enum Constant { 19 | Nil, 20 | Boolean(bool), 21 | Number(f64), 22 | String(usize), 23 | Import(usize), 24 | Table(Vec), 25 | Closure(usize), 26 | Vector(f32, f32, f32, f32), 27 | } 28 | 29 | impl Constant { 30 | pub(crate) fn parse(input: &[u8]) -> IResult<&[u8], Self> { 31 | let (input, tag) = le_u8(input)?; 32 | match tag { 33 | CONSTANT_NIL => Ok((input, Constant::Nil)), 34 | CONSTANT_BOOLEAN => { 35 | let (input, value) = le_u8(input)?; 36 | Ok((input, Constant::Boolean(value != 0u8))) 37 | } 38 | CONSTANT_NUMBER => { 39 | let (input, value) = le_f64(input)?; 40 | Ok((input, Constant::Number(value))) 41 | } 42 | CONSTANT_STRING => { 43 | let (input, string_index) = leb128_usize(input)?; 44 | Ok((input, Constant::String(string_index))) 45 | } 46 | CONSTANT_IMPORT => { 47 | let (input, import_index) = le_u32(input)?; 48 | Ok((input, Constant::Import(import_index as usize))) 49 | } 50 | CONSTANT_TABLE => { 51 | let (input, keys) = parse_list(input, leb128_usize)?; 52 | Ok((input, Constant::Table(keys))) 53 | } 54 | CONSTANT_CLOSURE => { 55 | let (input, f_id) = leb128_usize(input)?; 56 | Ok((input, Constant::Closure(f_id))) 57 | } 58 | CONSTANT_VECTOR => { 59 | let (input, x) = le_f32(input)?; 60 | let (input, y) = le_f32(input)?; 61 | let (input, z) = le_f32(input)?; 62 | let (input, w) = le_f32(input)?; 63 | Ok((input, Constant::Vector(x, y, z, w))) 64 | } 65 | _ => panic!("{}", tag), 66 | } 67 | } 68 | } 69 | -------------------------------------------------------------------------------- /ast/src/closure.rs: -------------------------------------------------------------------------------- 1 | use std::fmt; 2 | 3 | use by_address::ByAddress; 4 | use parking_lot::Mutex; 5 | use triomphe::Arc; 6 | 7 | use crate::{ 8 | formatter::Formatter, 9 | type_system::{Infer, TypeSystem}, 10 | Block, Literal, LocalRw, RcLocal, Reduce, SideEffects, Traverse, Type, 11 | }; 12 | 13 | #[derive(Debug, PartialEq, Clone)] 14 | pub enum Upvalue { 15 | Copy(RcLocal), 16 | Ref(RcLocal), 17 | } 18 | 19 | #[derive(Default, Debug, PartialEq, Clone)] 20 | pub struct Function { 21 | pub name: Option, 22 | pub parameters: Vec, 23 | pub is_variadic: bool, 24 | pub body: Block, 25 | } 26 | 27 | #[derive(Debug, PartialEq, Clone)] 28 | pub struct Closure { 29 | pub function: ByAddress>>, 30 | pub upvalues: Vec, 31 | } 32 | 33 | impl Reduce for Closure { 34 | fn reduce(self) -> crate::RValue { 35 | self.into() 36 | } 37 | 38 | fn reduce_condition(self) -> crate::RValue { 39 | Literal::Boolean(true).into() 40 | } 41 | } 42 | 43 | impl Infer for Closure { 44 | fn infer<'a: 'b, 'b>(&'a mut self, _system: &mut TypeSystem<'b>) -> Type { 45 | todo!() 46 | // let return_values = system.analyze_block(&mut self.body); 47 | // let parameters = self 48 | // .parameters 49 | // .iter_mut() 50 | // .map(|l| l.infer(system)) 51 | // .collect_vec(); 52 | 53 | // Type::Function(parameters, return_values) 54 | } 55 | } 56 | 57 | impl fmt::Display for Closure { 58 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 59 | Formatter { 60 | indentation_level: 0, 61 | indentation_mode: Default::default(), 62 | output: f, 63 | } 64 | .format_closure(self) 65 | } 66 | } 67 | 68 | impl LocalRw for Closure { 69 | fn values_read(&self) -> Vec<&RcLocal> { 70 | self.upvalues 71 | .iter() 72 | .map(|u| match u { 73 | Upvalue::Copy(l) | Upvalue::Ref(l) => l, 74 | }) 75 | .collect() 76 | } 77 | 78 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 79 | self.upvalues 80 | .iter_mut() 81 | .map(|u| match u { 82 | Upvalue::Copy(l) | Upvalue::Ref(l) => l, 83 | }) 84 | .collect() 85 | } 86 | } 87 | 88 | impl SideEffects for Closure {} 89 | 90 | impl Traverse for Closure {} 91 | -------------------------------------------------------------------------------- /lua51-deserializer/src/chunk/header.rs: -------------------------------------------------------------------------------- 1 | use nom::{ 2 | bytes::complete::tag, 3 | error::{Error, ErrorKind, ParseError}, 4 | number::complete::le_u8, 5 | Err, IResult, 6 | }; 7 | 8 | #[derive(Debug, PartialEq, Eq)] 9 | pub enum Endianness { 10 | Big, 11 | Little, 12 | } 13 | 14 | #[derive(Debug, PartialEq, Eq)] 15 | pub enum Format { 16 | Official, 17 | } 18 | 19 | #[derive(Debug)] 20 | pub struct Header { 21 | pub(crate) version_number: u8, 22 | pub(crate) format: Format, 23 | pub(crate) endianness: Endianness, 24 | pub(crate) int_width: u8, 25 | pub(crate) size_t_width: u8, 26 | pub(crate) instr_width: u8, 27 | pub(crate) number_width: u8, 28 | pub(crate) number_is_integral: bool, 29 | } 30 | 31 | impl Header { 32 | pub fn parse(input: &[u8]) -> IResult<&[u8], Self> { 33 | let (input, _) = tag("\x1BLua")(input)?; 34 | let (input, version_number) = le_u8(input)?; 35 | let (input, format) = match le_u8(input)? { 36 | (input, 0) => Ok((input, Format::Official)), 37 | _ => Err(Err::Failure(Error::from_error_kind( 38 | input, 39 | ErrorKind::Switch, 40 | ))), 41 | }?; 42 | // TODO: try_into instead 43 | let (input, endianness) = match le_u8(input)? { 44 | (input, 0) => Ok((input, Endianness::Big)), 45 | (input, 1) => Ok((input, Endianness::Little)), 46 | _ => Err(Err::Failure(Error::from_error_kind( 47 | input, 48 | ErrorKind::Switch, 49 | ))), 50 | }?; 51 | let (input, int_width) = le_u8(input)?; 52 | let (input, size_t_width) = le_u8(input)?; 53 | let (input, instr_width) = le_u8(input)?; 54 | let (input, number_width) = le_u8(input)?; 55 | let (input, number_is_integral) = match le_u8(input)? { 56 | (input, 0) => Ok((input, false)), 57 | (input, 1) => Ok((input, true)), 58 | _ => Err(Err::Failure(Error::from_error_kind( 59 | input, 60 | ErrorKind::Switch, 61 | ))), 62 | }?; 63 | 64 | Ok(( 65 | input, 66 | Self { 67 | version_number, 68 | format, 69 | endianness, 70 | int_width, 71 | size_t_width, 72 | instr_width, 73 | number_width, 74 | number_is_integral, 75 | }, 76 | )) 77 | } 78 | } 79 | -------------------------------------------------------------------------------- /lua51-deserializer/src/function.rs: -------------------------------------------------------------------------------- 1 | use nom::{ 2 | combinator::opt, 3 | multi::count, 4 | number::complete::{le_u32, le_u8}, 5 | IResult, 6 | }; 7 | 8 | use crate::{ 9 | instruction::{position::Position, Instruction}, 10 | local::Local, 11 | value::{self, Value}, 12 | }; 13 | 14 | #[derive(Debug)] 15 | pub struct Function<'a> { 16 | pub name: &'a [u8], 17 | pub line_defined: u32, 18 | pub last_line_defined: u32, 19 | pub number_of_upvalues: u8, 20 | pub vararg_flag: u8, 21 | pub maximum_stack_size: u8, 22 | pub code: Vec, 23 | pub constants: Vec>, 24 | pub closures: Vec>, 25 | pub positions: Vec, 26 | pub locals: Vec>, 27 | pub upvalues: Vec<&'a [u8]>, 28 | pub number_of_parameters: u8, 29 | } 30 | 31 | impl<'a> Function<'a> { 32 | pub fn parse(input: &'a [u8]) -> IResult<&'a [u8], Self> { 33 | let (input, name) = value::parse_string(input)?; 34 | let (input, line_defined) = le_u32(input)?; 35 | let (input, last_line_defined) = le_u32(input)?; 36 | let (input, number_of_upvalues) = le_u8(input)?; 37 | let (input, number_of_parameters) = le_u8(input)?; 38 | let (input, vararg_flag) = le_u8(input)?; 39 | let (input, maximum_stack_size) = le_u8(input)?; 40 | let (input, code_length) = le_u32(input)?; 41 | let (input, code) = count(Instruction::parse, code_length as usize)(input)?; 42 | let (input, constants_length) = le_u32(input)?; 43 | let (input, constants) = count(Value::parse, constants_length as usize)(input)?; 44 | let (input, closures_length) = le_u32(input)?; 45 | let (input, closures) = count(Self::parse, closures_length as usize)(input)?; 46 | let (input, positions) = opt(Position::parse)(input)?; 47 | let (input, locals) = opt(Local::parse_list)(input)?; 48 | let (input, upvalues) = opt(value::parse_strings)(input)?; 49 | 50 | Ok(( 51 | input, 52 | Self { 53 | name, 54 | line_defined, 55 | last_line_defined, 56 | number_of_upvalues, 57 | vararg_flag, 58 | maximum_stack_size, 59 | code, 60 | constants, 61 | closures, 62 | positions: positions.unwrap_or_default(), 63 | locals: locals.unwrap_or_default(), 64 | upvalues: upvalues.unwrap_or_default(), 65 | number_of_parameters, 66 | }, 67 | )) 68 | } 69 | } 70 | -------------------------------------------------------------------------------- /ast/src/literal.rs: -------------------------------------------------------------------------------- 1 | use derive_more::From; 2 | use enum_as_inner::EnumAsInner; 3 | use std::fmt; 4 | 5 | use crate::{ 6 | formatter::Formatter, type_system::Infer, LocalRw, Reduce, SideEffects, Traverse, Type, 7 | TypeSystem, 8 | }; 9 | 10 | #[derive(Debug, From, Clone, PartialEq, PartialOrd, EnumAsInner)] 11 | pub enum Literal { 12 | Nil, 13 | Boolean(bool), 14 | Number(f64), 15 | String(Vec), 16 | Vector(f32, f32, f32), 17 | } 18 | 19 | impl Reduce for Literal { 20 | fn reduce(self) -> crate::RValue { 21 | self.into() 22 | } 23 | 24 | fn reduce_condition(self) -> crate::RValue { 25 | Literal::Boolean(match self { 26 | Literal::Boolean(false) | Literal::Nil => false, 27 | Literal::Boolean(true) 28 | | Literal::Number(_) 29 | | Literal::String(_) 30 | | Literal::Vector(..) => true, 31 | }) 32 | .into() 33 | } 34 | } 35 | 36 | impl Infer for Literal { 37 | fn infer<'a: 'b, 'b>(&'a mut self, _: &mut TypeSystem<'b>) -> Type { 38 | match self { 39 | Literal::Nil => Type::Nil, 40 | Literal::Boolean(_) => Type::Boolean, 41 | Literal::Number(_) => Type::Number, 42 | Literal::String(_) => Type::String, 43 | Literal::Vector(..) => Type::Vector, 44 | } 45 | } 46 | } 47 | 48 | impl From<&str> for Literal { 49 | fn from(value: &str) -> Self { 50 | Self::String(value.into()) 51 | } 52 | } 53 | 54 | impl LocalRw for Literal {} 55 | 56 | impl SideEffects for Literal {} 57 | 58 | impl Traverse for Literal {} 59 | 60 | impl fmt::Display for Literal { 61 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 62 | match self { 63 | Literal::Nil => write!(f, "nil"), 64 | Literal::Boolean(value) => write!(f, "{}", value), 65 | &Literal::Number(value) => { 66 | // TODO: this is a bit messy, just use `buffer.format` here and format_finite 67 | // in formatter.rs 68 | debug_assert!(value.is_finite()); 69 | // TODO: fork ryu to remove ".0" 70 | let mut buffer = ryu::Buffer::new(); 71 | let printed = buffer.format_finite(value); 72 | write!(f, "{}", printed.strip_suffix(".0").unwrap_or(printed)) 73 | } 74 | Literal::String(value) => { 75 | write!( 76 | f, 77 | "\"{}\"", 78 | Formatter::::escape_string(value) 79 | ) 80 | } 81 | Literal::Vector(x, y, z) => write!(f, "Vector3.new({}, {}, {})", x, y, z), 82 | } 83 | } 84 | } 85 | -------------------------------------------------------------------------------- /ast/src/set_list.rs: -------------------------------------------------------------------------------- 1 | use crate::{formatter, LocalRw, RValue, RcLocal, SideEffects, Traverse}; 2 | 3 | #[derive(Debug, Clone, PartialEq)] 4 | pub struct SetList { 5 | pub object_local: RcLocal, 6 | pub index: usize, 7 | pub values: Vec, 8 | pub tail: Option, 9 | } 10 | 11 | impl SetList { 12 | pub fn new( 13 | object_local: RcLocal, 14 | index: usize, 15 | values: Vec, 16 | tail: Option, 17 | ) -> Self { 18 | Self { 19 | object_local, 20 | index, 21 | values, 22 | tail, 23 | } 24 | } 25 | } 26 | 27 | impl LocalRw for SetList { 28 | fn values_read(&self) -> Vec<&RcLocal> { 29 | let tail_locals = self 30 | .tail 31 | .as_ref() 32 | .map(|t| t.values_read()) 33 | .unwrap_or_default(); 34 | std::iter::once(&self.object_local) 35 | .chain(self.values.iter().flat_map(|rvalue| rvalue.values_read())) 36 | .chain(tail_locals) 37 | .collect() 38 | } 39 | 40 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 41 | let tail_locals = self 42 | .tail 43 | .as_mut() 44 | .map(|t| t.values_read_mut()) 45 | .unwrap_or_default(); 46 | std::iter::once(&mut self.object_local) 47 | .chain( 48 | self.values 49 | .iter_mut() 50 | .flat_map(|rvalue| rvalue.values_read_mut()), 51 | ) 52 | .chain(tail_locals) 53 | .collect() 54 | } 55 | } 56 | 57 | impl SideEffects for SetList { 58 | fn has_side_effects(&self) -> bool { 59 | self.values 60 | .iter() 61 | .chain(self.tail.as_ref()) 62 | .any(|rvalue| rvalue.has_side_effects()) 63 | } 64 | } 65 | 66 | impl Traverse for SetList { 67 | fn rvalues(&self) -> Vec<&RValue> { 68 | self.values.iter().chain(self.tail.as_ref()).collect() 69 | } 70 | 71 | fn rvalues_mut(&mut self) -> Vec<&mut RValue> { 72 | self.values.iter_mut().chain(self.tail.as_mut()).collect() 73 | } 74 | } 75 | 76 | impl std::fmt::Display for SetList { 77 | fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { 78 | write!( 79 | f, 80 | "__set_list({}, {}, {{{}}})", 81 | self.object_local, 82 | self.index, 83 | // TODO: bad 84 | formatter::format_arg_list( 85 | &self 86 | .values 87 | .iter() 88 | .chain(self.tail.as_ref()) 89 | .cloned() 90 | .collect::>() 91 | ) 92 | ) 93 | } 94 | } 95 | -------------------------------------------------------------------------------- /luau-worker/src/lib.rs: -------------------------------------------------------------------------------- 1 | use futures_util::StreamExt; 2 | extern crate console_error_panic_hook; 3 | 4 | use base64::prelude::*; 5 | use luau_lifter::decompile_bytecode; 6 | use serde::{Deserialize, Serialize}; 7 | use worker::*; 8 | 9 | const AUTH_SECRET: &str = "ymjKH2O3BbO3bDSsKmpo3ek3vHxIWYLQfj0"; 10 | 11 | #[derive(Deserialize)] 12 | struct DecompileMessage { 13 | id: String, 14 | encoded_bytecode: String, 15 | } 16 | 17 | #[derive(Serialize)] 18 | struct DecompileResponse { 19 | id: String, 20 | decompilation: String, 21 | } 22 | 23 | #[event(fetch, respond_with_errors)] 24 | pub async fn main(req: Request, env: Env, _ctx: worker::Context) -> Result { 25 | console_error_panic_hook::set_once(); 26 | 27 | let router = Router::new(); 28 | router 29 | .get_async("/decompile_ws", |req, _ctx| async move { 30 | let license = req 31 | .headers() 32 | .get("Authorization") 33 | .unwrap_or_default() 34 | .expect("authorization header is required"); 35 | 36 | if license != AUTH_SECRET { 37 | return Response::error("invalid license", 403); 38 | } 39 | 40 | let pair = WebSocketPair::new()?; 41 | let server = pair.server; 42 | server.accept()?; 43 | 44 | wasm_bindgen_futures::spawn_local(async move { 45 | let mut event_stream = server.events().expect("could not open stream"); 46 | while let Some(event) = event_stream.next().await { 47 | if let WebsocketEvent::Message(msg) = 48 | event.expect("received error in websocket") 49 | { 50 | let msg = msg 51 | .json::() 52 | .expect("malformed decompile message"); 53 | let bytecode = BASE64_STANDARD 54 | .decode(msg.encoded_bytecode) 55 | .expect("bytecode must be base64 encoded"); 56 | let resp = DecompileResponse { 57 | id: msg.id, 58 | decompilation: decompile_bytecode(&bytecode, 1), 59 | }; 60 | server 61 | .send_with_str(serde_json::to_string(&resp).unwrap()) 62 | .unwrap(); 63 | } 64 | } 65 | }); 66 | 67 | Response::from_websocket(pair.client) 68 | }) 69 | .post_async("/decompile", |mut req, _ctx| async move { 70 | let license = req 71 | .headers() 72 | .get("Authorization") 73 | .unwrap_or_default() 74 | .expect("authorization header is required"); 75 | 76 | if license != AUTH_SECRET { 77 | return Response::error("invalid license", 403); 78 | } 79 | 80 | let encoded_bytecode = req.bytes().await?; 81 | match BASE64_STANDARD.decode(encoded_bytecode) { 82 | Ok(bytecode) => Response::ok(decompile_bytecode(&bytecode, 203)), 83 | Err(_) => Response::error("invalid bytecode", 400), 84 | } 85 | }) 86 | .run(req, env) 87 | .await 88 | } 89 | -------------------------------------------------------------------------------- /ast/src/local.rs: -------------------------------------------------------------------------------- 1 | use crate::{type_system::Infer, SideEffects, Traverse, Type, TypeSystem}; 2 | use by_address::ByAddress; 3 | use derive_more::From; 4 | use enum_dispatch::enum_dispatch; 5 | use nohash_hasher::NoHashHasher; 6 | use parking_lot::Mutex; 7 | use std::{ 8 | fmt::{self, Display}, 9 | hash::{Hash, Hasher}, 10 | }; 11 | use triomphe::Arc; 12 | 13 | #[derive(Debug, Default, From, Clone, PartialEq, PartialOrd, Ord, Eq, Hash)] 14 | pub struct Local(pub Option); 15 | 16 | impl Local { 17 | pub fn new(name: Option) -> Self { 18 | Self(name) 19 | } 20 | } 21 | 22 | impl fmt::Display for Local { 23 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 24 | match &self.0 { 25 | Some(name) => write!(f, "{}", name), 26 | None => write!(f, "UNNAMED_LOCAL"), 27 | } 28 | } 29 | } 30 | 31 | #[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] 32 | pub struct RcLocal(pub ByAddress>>); 33 | 34 | impl Infer for RcLocal { 35 | fn infer<'a: 'b, 'b>(&'a mut self, system: &mut TypeSystem<'b>) -> Type { 36 | system.type_of(self).clone() 37 | } 38 | } 39 | 40 | impl Display for RcLocal { 41 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 42 | match &self.0 .0.lock().0 { 43 | Some(name) => write!(f, "{}", name), 44 | None => { 45 | let mut hasher = NoHashHasher::::default(); 46 | self.hash(&mut hasher); 47 | write!(f, "UNNAMED_{}", hasher.finish()) 48 | } 49 | } 50 | } 51 | } 52 | 53 | impl SideEffects for RcLocal {} 54 | 55 | impl Traverse for RcLocal {} 56 | 57 | impl RcLocal { 58 | pub fn new(local: Local) -> Self { 59 | Self(ByAddress(Arc::new(Mutex::new(local)))) 60 | } 61 | } 62 | 63 | impl LocalRw for RcLocal { 64 | fn values_read(&self) -> Vec<&RcLocal> { 65 | vec![self] 66 | } 67 | 68 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 69 | vec![self] 70 | } 71 | } 72 | 73 | #[enum_dispatch] 74 | pub trait LocalRw { 75 | fn values_read(&self) -> Vec<&RcLocal> { 76 | Vec::new() 77 | } 78 | 79 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 80 | Vec::new() 81 | } 82 | 83 | fn values_written(&self) -> Vec<&RcLocal> { 84 | Vec::new() 85 | } 86 | 87 | fn values_written_mut(&mut self) -> Vec<&mut RcLocal> { 88 | Vec::new() 89 | } 90 | 91 | fn values(&self) -> Vec<&RcLocal> { 92 | self.values_read() 93 | .into_iter() 94 | .chain(self.values_written()) 95 | .collect() 96 | } 97 | 98 | fn replace_values_read(&mut self, old: &RcLocal, new: &RcLocal) { 99 | for value in self.values_read_mut() { 100 | if value == old { 101 | *value = new.clone(); 102 | } 103 | } 104 | } 105 | 106 | fn replace_values_written(&mut self, old: &RcLocal, new: &RcLocal) { 107 | for value in self.values_written_mut() { 108 | if value == old { 109 | *value = new.clone(); 110 | } 111 | } 112 | } 113 | 114 | fn replace_values(&mut self, old: &RcLocal, new: &RcLocal) { 115 | self.replace_values_read(old, new); 116 | self.replace_values_written(old, new); 117 | } 118 | } 119 | -------------------------------------------------------------------------------- /cfg/src/pattern.rs: -------------------------------------------------------------------------------- 1 | use std::fmt; 2 | 3 | use petgraph::{ 4 | stable_graph::{NodeIndex, StableDiGraph}, 5 | visit::{Dfs, EdgeRef, Walker}, 6 | Direction, 7 | }; 8 | use rustc_hash::{FxHashMap, FxHashSet}; 9 | 10 | use crate::{block::BlockEdge, function::Function}; 11 | 12 | pub trait NodeChecker { 13 | fn check(&self, node: NodeIndex) -> bool; 14 | } 15 | 16 | #[derive(Debug)] 17 | pub struct PatternNode { 18 | pub allow_external_neighbors: bool, 19 | } 20 | 21 | impl PatternNode { 22 | pub fn new(allow_external_neighbors: bool) -> Self { 23 | Self { 24 | allow_external_neighbors, 25 | } 26 | } 27 | } 28 | 29 | impl fmt::Display for PatternNode { 30 | fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { 31 | Ok(()) 32 | } 33 | } 34 | 35 | pub type PatternEdge = BlockEdge; 36 | pub type PatternGraph = StableDiGraph; 37 | 38 | #[derive(Debug)] 39 | pub struct Pattern { 40 | root: NodeIndex, 41 | graph: PatternGraph, 42 | } 43 | 44 | impl Pattern { 45 | fn new(root: NodeIndex, graph: PatternGraph) -> Self { 46 | // make sure all nodes in pattern are connected 47 | assert!(Dfs::new(&graph, root).iter(&graph).count() == graph.node_count()); 48 | Self { root, graph } 49 | } 50 | } 51 | 52 | #[derive(Debug)] 53 | pub struct PatternChecker<'a> { 54 | pattern: &'a Pattern, 55 | function: &'a Function, 56 | visited: FxHashSet, 57 | mapping: FxHashMap, 58 | } 59 | 60 | impl<'a> PatternChecker<'a> { 61 | fn check_successors(&self, _node: NodeIndex) {} 62 | 63 | fn check_pattern_rec(&self, pattern_node: NodeIndex, function_node: NodeIndex) -> bool { 64 | let _function_successors = self 65 | .function 66 | .successor_blocks(function_node) 67 | .collect::>(); 68 | 69 | let _pattern_successors = self 70 | .pattern 71 | .graph 72 | .neighbors_directed(pattern_node, Direction::Outgoing) 73 | .collect::>(); 74 | 75 | for function_edge in self.function.edges(pattern_node) { 76 | let function_node = function_edge.target(); 77 | if let Some(&pattern_successor_node) = self.mapping.get(&function_node) { 78 | if let Some(_pattern_edge) = self 79 | .pattern 80 | .graph 81 | .find_edge(pattern_node, pattern_successor_node) 82 | { 83 | todo!(); 84 | } else { 85 | return false; 86 | } 87 | } else { 88 | todo!(); 89 | } 90 | } 91 | 92 | true 93 | } 94 | 95 | fn check_pattern(mut self, root: NodeIndex) -> bool { 96 | self.mapping.insert(self.pattern.root, root); 97 | self.check_pattern_rec(self.pattern.root, root) 98 | } 99 | } 100 | 101 | impl Pattern { 102 | fn node_matches(&self, _function: &Function, _node: NodeIndex) -> bool { 103 | let _mapping = FxHashMap::::default(); 104 | 105 | false 106 | } 107 | 108 | pub fn match_on(&self, _function: &Function, _node: NodeIndex) -> Option { 109 | None 110 | } 111 | } 112 | 113 | #[derive(Debug, Default)] 114 | pub struct Match { 115 | // pattern node to function node 116 | pub mapping: FxHashMap, 117 | } 118 | -------------------------------------------------------------------------------- /ast/src/table.rs: -------------------------------------------------------------------------------- 1 | use crate::{ 2 | formatter::Formatter, Literal, LocalRw, RValue, RcLocal, Reduce, SideEffects, Traverse, 3 | }; 4 | 5 | use std::{fmt, iter}; 6 | 7 | #[derive(Debug, Clone, PartialEq, Default)] 8 | pub struct Table(pub Vec<(Option, RValue)>); 9 | 10 | impl Reduce for Table { 11 | fn reduce(self) -> RValue { 12 | self.into() 13 | } 14 | 15 | fn reduce_condition(self) -> RValue { 16 | if self.has_side_effects() { 17 | // TODO: remove all members w/o side effects 18 | self.into() 19 | } else { 20 | Literal::Boolean(true).into() 21 | } 22 | } 23 | } 24 | 25 | /*impl Infer for Table { 26 | fn infer<'a: 'b, 'b>(&'a mut self, system: &mut TypeSystem<'b>) -> Type { 27 | let elements: BTreeSet<_> = self 28 | .0 29 | .iter_mut() 30 | .map(|(f, v)| (f.clone(), v.infer(system))) 31 | .collect(); 32 | let elements: BTreeSet<_> = elements 33 | .iter() 34 | .filter(|(f, t)| { 35 | f.is_some() || !elements.iter().any(|(_, x)| t != x && t.is_subtype_of(x)) 36 | }) 37 | .cloned() 38 | .collect(); 39 | let (elements, fields): (BTreeSet<_>, BTreeMap<_, _>) = 40 | elements.into_iter().partition_map(|(f, t)| match f { 41 | None => Either::Left(t), 42 | Some(f) => Either::Right((f, t)), 43 | }); 44 | 45 | Type::Table { 46 | indexer: Box::new(( 47 | Type::Any, 48 | if elements.len() > 1 { 49 | Type::Union(elements) 50 | } else { 51 | elements.into_iter().next().unwrap_or(Type::Any) 52 | }, 53 | )), 54 | fields, 55 | } 56 | } 57 | }*/ 58 | 59 | impl LocalRw for Table { 60 | fn values_read(&self) -> Vec<&RcLocal> { 61 | self.0 62 | .iter() 63 | .flat_map(|(k, v)| k.iter().chain(iter::once(v))) 64 | .flat_map(|v| v.values_read()) 65 | .collect() 66 | } 67 | 68 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 69 | self.0 70 | .iter_mut() 71 | .flat_map(|(k, v)| k.iter_mut().chain(iter::once(v))) 72 | .flat_map(|v| v.values_read_mut()) 73 | .collect() 74 | } 75 | } 76 | 77 | impl Traverse for Table { 78 | fn rvalues_mut(&mut self) -> Vec<&mut RValue> { 79 | self.0 80 | .iter_mut() 81 | .flat_map(|(k, v)| k.iter_mut().chain(iter::once(v))) 82 | .collect() 83 | } 84 | 85 | fn rvalues(&self) -> Vec<&RValue> { 86 | self.0 87 | .iter() 88 | .flat_map(|(k, v)| k.iter().chain(iter::once(v))) 89 | .collect() 90 | } 91 | } 92 | 93 | impl SideEffects for Table { 94 | fn has_side_effects(&self) -> bool { 95 | self.0 96 | .iter() 97 | .flat_map(|(k, v)| k.iter().chain(iter::once(v))) 98 | .any(|r| r.has_side_effects()) 99 | } 100 | } 101 | 102 | /*impl fmt::Display for Table { 103 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 104 | write!( 105 | f, 106 | "{{{}}}", 107 | self.0 108 | .iter() 109 | .map(|(key, value)| match key { 110 | Some(key) => format!("{} = {}", key, value), 111 | None => value.to_string(), 112 | }) 113 | .join(", ") 114 | ) 115 | } 116 | }*/ 117 | 118 | impl fmt::Display for Table { 119 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 120 | Formatter { 121 | indentation_level: 0, 122 | indentation_mode: Default::default(), 123 | output: f, 124 | } 125 | .format_table(self) 126 | } 127 | } 128 | -------------------------------------------------------------------------------- /ast/src/call.rs: -------------------------------------------------------------------------------- 1 | use std::fmt; 2 | 3 | use crate::{formatter::Formatter, has_side_effects, LocalRw, RcLocal, Traverse}; 4 | 5 | use super::RValue; 6 | 7 | #[derive(Debug, Clone, PartialEq)] 8 | pub struct Call { 9 | pub value: Box, 10 | pub arguments: Vec, 11 | } 12 | 13 | impl Call { 14 | pub fn new(value: RValue, arguments: Vec) -> Self { 15 | Self { 16 | value: Box::new(value), 17 | arguments, 18 | } 19 | } 20 | } 21 | 22 | // call can error 23 | has_side_effects!(Call); 24 | // impl SideEffects for Call { 25 | // fn has_side_effects(&self) -> bool { 26 | // matches!(self.value, box RValue::Local(_)) 27 | // || self.value.has_side_effects() 28 | // || self.arguments.iter().any(|arg| arg.has_side_effects()) 29 | // } 30 | // } 31 | 32 | impl Traverse for Call { 33 | fn rvalues_mut(&mut self) -> Vec<&mut RValue> { 34 | std::iter::once(self.value.as_mut()) 35 | .chain(self.arguments.iter_mut()) 36 | .collect() 37 | } 38 | 39 | fn rvalues(&self) -> Vec<&RValue> { 40 | std::iter::once(self.value.as_ref()) 41 | .chain(self.arguments.iter()) 42 | .collect() 43 | } 44 | } 45 | 46 | impl LocalRw for Call { 47 | fn values_read(&self) -> Vec<&RcLocal> { 48 | self.value 49 | .values_read() 50 | .into_iter() 51 | .chain(self.arguments.iter().flat_map(|r| r.values_read())) 52 | .collect() 53 | } 54 | 55 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 56 | self.value 57 | .values_read_mut() 58 | .into_iter() 59 | .chain(self.arguments.iter_mut().flat_map(|r| r.values_read_mut())) 60 | .collect() 61 | } 62 | } 63 | 64 | impl fmt::Display for Call { 65 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 66 | Formatter { 67 | indentation_level: 0, 68 | indentation_mode: Default::default(), 69 | output: f, 70 | } 71 | .format_call(self) 72 | } 73 | } 74 | 75 | #[derive(Debug, Clone, PartialEq)] 76 | pub struct MethodCall { 77 | // TODO: STYLE: rename to object? 78 | pub value: Box, 79 | pub method: String, 80 | pub arguments: Vec, 81 | } 82 | 83 | impl MethodCall { 84 | pub fn new(value: RValue, method: String, arguments: Vec) -> Self { 85 | Self { 86 | value: Box::new(value), 87 | method, 88 | arguments, 89 | } 90 | } 91 | } 92 | 93 | // this should reflect Index 94 | has_side_effects!(MethodCall); 95 | 96 | impl Traverse for MethodCall { 97 | fn rvalues_mut(&mut self) -> Vec<&mut RValue> { 98 | std::iter::once(self.value.as_mut()) 99 | .chain(self.arguments.iter_mut()) 100 | .collect() 101 | } 102 | 103 | fn rvalues(&self) -> Vec<&RValue> { 104 | std::iter::once(self.value.as_ref()) 105 | .chain(self.arguments.iter()) 106 | .collect() 107 | } 108 | } 109 | 110 | impl LocalRw for MethodCall { 111 | fn values_read(&self) -> Vec<&RcLocal> { 112 | self.value 113 | .values_read() 114 | .into_iter() 115 | .chain(self.arguments.iter().flat_map(|r| r.values_read())) 116 | .collect() 117 | } 118 | 119 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 120 | self.value 121 | .values_read_mut() 122 | .into_iter() 123 | .chain(self.arguments.iter_mut().flat_map(|r| r.values_read_mut())) 124 | .collect() 125 | } 126 | } 127 | 128 | impl fmt::Display for MethodCall { 129 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 130 | Formatter { 131 | indentation_level: 0, 132 | indentation_mode: Default::default(), 133 | output: f, 134 | } 135 | .format_method_call(self) 136 | } 137 | } 138 | -------------------------------------------------------------------------------- /ast/src/traverse.rs: -------------------------------------------------------------------------------- 1 | use crate::{LValue, RValue}; 2 | use enum_dispatch::enum_dispatch; 3 | use itertools::Either; 4 | 5 | pub enum PreOrPost { 6 | Pre, 7 | Post, 8 | } 9 | 10 | #[enum_dispatch] 11 | pub trait Traverse { 12 | fn lvalues_mut(&mut self) -> Vec<&mut LValue> { 13 | Vec::new() 14 | } 15 | 16 | fn rvalues_mut(&mut self) -> Vec<&mut RValue> { 17 | Vec::new() 18 | } 19 | 20 | fn rvalues(&self) -> Vec<&RValue> { 21 | Vec::new() 22 | } 23 | 24 | // fn traverse_lvalues( 25 | // &mut self, 26 | // lvalue_callback: &impl Fn(&mut LValue), 27 | // rvalue_callback: &impl Fn(&mut RValue), 28 | // ) { 29 | // self.rvalues_mut().into_iter().for_each(rvalue_callback); 30 | // self.lvalues_mut().into_iter().for_each(lvalue_callback); 31 | // self.lvalues_mut().into_iter().for_each(|lvalue| { 32 | // lvalue.traverse_lvalues(lvalue_callback, rvalue_callback); 33 | // }); 34 | // } 35 | 36 | fn traverse_rvalues(&mut self, callback: &mut F) 37 | where 38 | F: FnMut(&mut RValue), 39 | { 40 | for lvalue in self.lvalues_mut() { 41 | lvalue.traverse_rvalues(callback); 42 | } 43 | for rvalue in self.rvalues_mut() { 44 | callback(rvalue); 45 | rvalue.traverse_rvalues(callback); 46 | } 47 | } 48 | 49 | fn post_traverse_rvalues(&mut self, callback: &mut F) -> Option 50 | where 51 | F: FnMut(&mut RValue) -> Option, 52 | { 53 | for lvalue in self.lvalues_mut() { 54 | if let Some(res) = lvalue.post_traverse_rvalues(callback) { 55 | return Some(res); 56 | } 57 | } 58 | for rvalue in self.rvalues_mut() { 59 | if let Some(res) = rvalue.post_traverse_rvalues(callback) { 60 | return Some(res); 61 | } 62 | if let Some(res) = callback(rvalue) { 63 | return Some(res); 64 | } 65 | } 66 | 67 | None 68 | } 69 | 70 | fn post_traverse_values(&mut self, callback: &mut F) -> Option 71 | where 72 | // TODO: REFACTOR: use an enum called Value instead of Either 73 | F: FnMut(Either<&mut LValue, &mut RValue>) -> Option, 74 | { 75 | for lvalue in self.lvalues_mut() { 76 | if let Some(res) = lvalue.post_traverse_values(callback) { 77 | return Some(res); 78 | } 79 | if let Some(res) = callback(Either::Left(lvalue)) { 80 | return Some(res); 81 | } 82 | } 83 | for rvalue in self.rvalues_mut() { 84 | if let Some(res) = rvalue.post_traverse_values(callback) { 85 | return Some(res); 86 | } 87 | if let Some(res) = callback(Either::Right(rvalue)) { 88 | return Some(res); 89 | } 90 | } 91 | 92 | None 93 | } 94 | 95 | fn traverse_values(&mut self, callback: &mut F) -> Option 96 | where 97 | // TODO: REFACTOR: use an enum called Value instead of Either 98 | F: FnMut(PreOrPost, Either<&mut LValue, &mut RValue>) -> Option, 99 | { 100 | for lvalue in self.lvalues_mut() { 101 | if let Some(res) = callback(PreOrPost::Pre, Either::Left(lvalue)) { 102 | return Some(res); 103 | } 104 | if let Some(res) = lvalue.traverse_values(callback) { 105 | return Some(res); 106 | } 107 | if let Some(res) = callback(PreOrPost::Post, Either::Left(lvalue)) { 108 | return Some(res); 109 | } 110 | } 111 | for rvalue in self.rvalues_mut() { 112 | if let Some(res) = callback(PreOrPost::Pre, Either::Right(rvalue)) { 113 | return Some(res); 114 | } 115 | if let Some(res) = rvalue.traverse_values(callback) { 116 | return Some(res); 117 | } 118 | if let Some(res) = callback(PreOrPost::Post, Either::Right(rvalue)) { 119 | return Some(res); 120 | } 121 | } 122 | 123 | None 124 | } 125 | } 126 | -------------------------------------------------------------------------------- /lua51-deserializer/src/instruction/operation_code.rs: -------------------------------------------------------------------------------- 1 | use crate::instruction::layout::LayoutDiscriminants; 2 | use nom::{ 3 | error::{Error, ErrorKind, ParseError}, 4 | number::complete::le_u8, 5 | Err, IResult, 6 | }; 7 | use num_derive::{FromPrimitive, ToPrimitive}; 8 | use num_traits::FromPrimitive; 9 | 10 | #[derive(Debug, FromPrimitive, ToPrimitive)] 11 | pub enum OperationCode { 12 | Move = 0, 13 | LoadConstant, 14 | LoadBoolean, 15 | LoadNil, 16 | GetUpvalue, 17 | GetGlobal, 18 | GetIndex, 19 | SetGlobal, 20 | SetUpvalue, 21 | SetIndex, 22 | NewTable, 23 | PrepMethodCall, 24 | Add, 25 | Subtract, 26 | Multiply, 27 | Divide, 28 | Modulo, 29 | Power, 30 | Minus, 31 | Not, 32 | Length, 33 | Concatenate, 34 | Jump, 35 | Equal, 36 | LessThan, 37 | LessThanOrEqual, 38 | Test, 39 | TestSet, 40 | Call, 41 | TailCall, 42 | Return, 43 | IterateNumericForLoop, 44 | InitNumericForLoop, 45 | IterateGenericForLoop, 46 | SetList, 47 | Close, 48 | Closure, 49 | VarArg, 50 | } 51 | 52 | impl OperationCode { 53 | pub fn parse(input: &[u8]) -> IResult<&[u8], Self> { 54 | let (input, operation_code) = le_u8(input)?; 55 | let operation_code = operation_code & 0x3F; 56 | 57 | Ok(( 58 | input, 59 | match FromPrimitive::from_u8(operation_code) { 60 | None => { 61 | return Err(Err::Failure(Error::from_error_kind( 62 | input, 63 | ErrorKind::Switch, 64 | ))) 65 | } 66 | Some(operation_code) => operation_code, 67 | }, 68 | )) 69 | } 70 | 71 | pub fn instruction_layout(&self) -> LayoutDiscriminants { 72 | /* 73 | 0 = BC 74 | 1 = BX 75 | 2 = BSx 76 | */ 77 | 78 | match self { 79 | Self::Move => LayoutDiscriminants::BC, 80 | Self::LoadConstant => LayoutDiscriminants::BX, 81 | Self::LoadBoolean => LayoutDiscriminants::BC, 82 | Self::LoadNil => LayoutDiscriminants::BC, 83 | Self::GetUpvalue => LayoutDiscriminants::BC, 84 | Self::GetGlobal => LayoutDiscriminants::BX, 85 | Self::GetIndex => LayoutDiscriminants::BC, 86 | Self::SetGlobal => LayoutDiscriminants::BX, 87 | Self::SetUpvalue => LayoutDiscriminants::BC, 88 | Self::SetIndex => LayoutDiscriminants::BC, 89 | Self::NewTable => LayoutDiscriminants::BC, 90 | Self::PrepMethodCall => LayoutDiscriminants::BC, 91 | Self::Add => LayoutDiscriminants::BC, 92 | Self::Subtract => LayoutDiscriminants::BC, 93 | Self::Multiply => LayoutDiscriminants::BC, 94 | Self::Divide => LayoutDiscriminants::BC, 95 | Self::Modulo => LayoutDiscriminants::BC, 96 | Self::Power => LayoutDiscriminants::BC, 97 | Self::Minus => LayoutDiscriminants::BC, 98 | Self::Not => LayoutDiscriminants::BC, 99 | Self::Length => LayoutDiscriminants::BC, 100 | Self::Concatenate => LayoutDiscriminants::BC, 101 | Self::Jump => LayoutDiscriminants::BSx, 102 | Self::Equal => LayoutDiscriminants::BC, 103 | Self::LessThan => LayoutDiscriminants::BC, 104 | Self::LessThanOrEqual => LayoutDiscriminants::BC, 105 | Self::Test => LayoutDiscriminants::BC, 106 | Self::TestSet => LayoutDiscriminants::BC, 107 | Self::Call => LayoutDiscriminants::BC, 108 | Self::TailCall => LayoutDiscriminants::BC, 109 | Self::Return => LayoutDiscriminants::BC, 110 | Self::IterateNumericForLoop => LayoutDiscriminants::BSx, 111 | Self::InitNumericForLoop => LayoutDiscriminants::BSx, 112 | Self::IterateGenericForLoop => LayoutDiscriminants::BC, 113 | Self::SetList => LayoutDiscriminants::BC, 114 | Self::Close => LayoutDiscriminants::BC, 115 | Self::Closure => LayoutDiscriminants::BX, 116 | Self::VarArg => LayoutDiscriminants::BC, 117 | } 118 | } 119 | } 120 | -------------------------------------------------------------------------------- /cfg/src/dot.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | borrow::{Borrow, Cow}, 3 | cell::RefCell, 4 | io::Write, 5 | }; 6 | 7 | use ast::LocalRw; 8 | use dot::{GraphWalk, LabelText, Labeller}; 9 | 10 | use itertools::Itertools; 11 | use petgraph::{ 12 | stable_graph::{EdgeIndex, NodeIndex}, 13 | visit::{Bfs, Walker}, 14 | }; 15 | 16 | use crate::function::Function; 17 | 18 | fn arguments(args: &Vec<(ast::RcLocal, ast::RValue)>) -> String { 19 | let mut s = String::new(); 20 | for (i, (local, new_local)) in args.iter().enumerate() { 21 | use std::fmt::Write; 22 | write!(s, "{} -> {}", local, new_local).unwrap(); 23 | if i + 1 != args.len() { 24 | s.push('\n'); 25 | } 26 | } 27 | s 28 | } 29 | 30 | struct FunctionLabeller<'a> { 31 | function: &'a Function, 32 | counter: RefCell, 33 | } 34 | 35 | impl<'a> Labeller<'a, NodeIndex, EdgeIndex> for FunctionLabeller<'a> { 36 | fn graph_id(&'a self) -> dot::Id<'a> { 37 | dot::Id::new("cfg").unwrap() 38 | } 39 | 40 | fn node_label<'b>(&'b self, n: &NodeIndex) -> dot::LabelText<'b> { 41 | let block = self.function.block(*n).unwrap(); 42 | let prefix = if self.function.entry() == &Some(*n) { 43 | "entry" 44 | } else { 45 | "" 46 | }; 47 | dot::LabelText::LabelStr( 48 | block 49 | .iter() 50 | .map(|s| { 51 | for local in s.values() { 52 | let name = &mut local.0 .0.lock().0; 53 | if name.is_none() { 54 | // TODO: ugly 55 | *name = Some(format!("v{}", self.counter.borrow())); 56 | *self.counter.borrow_mut() += 1; 57 | } 58 | } 59 | s 60 | }) 61 | .join("\n") 62 | .into(), 63 | ) 64 | .prefix_line(dot::LabelText::LabelStr( 65 | format!("{} {}", n.index(), prefix).into(), 66 | )) 67 | } 68 | 69 | fn edge_label<'b>(&'b self, e: &EdgeIndex) -> dot::LabelText<'b> { 70 | let edge = self.function.graph().edge_weight(*e).unwrap(); 71 | match edge.branch_type { 72 | crate::block::BranchType::Unconditional => { 73 | dot::LabelText::LabelStr(arguments(&edge.arguments).into()) 74 | } 75 | crate::block::BranchType::Then => { 76 | let arguments = arguments(&edge.arguments); 77 | if !arguments.is_empty() { 78 | dot::LabelText::LabelStr(format!("t\n{}", arguments).into()) 79 | } else { 80 | dot::LabelText::LabelStr("t".into()) 81 | } 82 | } 83 | crate::block::BranchType::Else => { 84 | let arguments = arguments(&edge.arguments); 85 | if !arguments.is_empty() { 86 | dot::LabelText::LabelStr(format!("e\n{}", arguments).into()) 87 | } else { 88 | dot::LabelText::LabelStr("e".into()) 89 | } 90 | } 91 | } 92 | } 93 | 94 | fn node_id(&'a self, n: &NodeIndex) -> dot::Id<'a> { 95 | dot::Id::new(format!("N{}", n.index())).unwrap() 96 | } 97 | 98 | fn node_shape(&'a self, _n: &NodeIndex) -> Option> { 99 | Some(LabelText::LabelStr("rect".into())) 100 | } 101 | } 102 | 103 | impl<'a> GraphWalk<'a, NodeIndex, EdgeIndex> for FunctionLabeller<'a> { 104 | fn nodes(&'a self) -> dot::Nodes<'a, NodeIndex> { 105 | Cow::Owned( 106 | Bfs::new(self.function.graph(), self.function.entry().unwrap()) 107 | .iter(self.function.graph()) 108 | .collect::>(), 109 | ) 110 | } 111 | 112 | fn edges(&'a self) -> dot::Edges<'a, EdgeIndex> { 113 | Cow::Owned(self.function.graph().edge_indices().collect()) 114 | } 115 | 116 | fn source(&self, e: &EdgeIndex) -> NodeIndex { 117 | self.function.graph().edge_endpoints(*e).unwrap().0 118 | } 119 | 120 | fn target(&self, e: &EdgeIndex) -> NodeIndex { 121 | self.function.graph().edge_endpoints(*e).unwrap().1 122 | } 123 | } 124 | 125 | pub fn render_to(function: &Function, output: &mut W) -> std::io::Result<()> { 126 | dot::render( 127 | &FunctionLabeller { 128 | function, 129 | counter: RefCell::new(1), 130 | }, 131 | output, 132 | ) 133 | } 134 | -------------------------------------------------------------------------------- /luau-lifter/src/instruction.rs: -------------------------------------------------------------------------------- 1 | use std::convert::TryFrom; 2 | 3 | use crate::op_code::OpCode; 4 | 5 | /* 6 | 7 | registers in () are not used 8 | f prefix means no registers are used but its parsed as said type 9 | 10 | LOP_NOP, f abc 11 | LOP_BREAK, f abc 12 | LOP_LOADNIL, a (bc) 13 | LOP_LOADB, abc 14 | LOP_LOADN, ad 15 | LOP_LOADK, ad 16 | LOP_MOVE, ab (c) 17 | LOP_GETGLOBAL, a (b) c aux 18 | LOP_SETGLOBAL, a (b) c aux 19 | LOP_GETUPVAL, ab (c) 20 | LOP_SETUPVAL, ab (c) 21 | LOP_CLOSEUPVALS, a (bc) 22 | LOP_GETIMPORT, ad aux 23 | LOP_GETTABLE, abc 24 | LOP_SETTABLE, abc 25 | LOP_GETTABLEKS, abc aux 26 | LOP_SETTABLEKS, abc aux 27 | LOP_GETTABLEN, abc 28 | LOP_SETTABLEN, abc 29 | LOP_NEWCLOSURE, ad 30 | LOP_NAMECALL, abc aux 31 | LOP_CALL, abc 32 | LOP_RETURN, ab (c) 33 | LOP_JUMP, (a) d 34 | LOP_JUMPBACK, (a) d 35 | LOP_JUMPIF, ad 36 | LOP_JUMPIFNOT, ad 37 | LOP_JUMPIFEQ, ad aux 38 | LOP_JUMPIFLE, ad aux 39 | LOP_JUMPIFLT, ad aux 40 | LOP_JUMPIFNOTEQ, ad aux 41 | LOP_JUMPIFNOTLE, ad aux 42 | LOP_JUMPIFNOTLT, ad aux 43 | LOP_ADD, abc 44 | LOP_SUB, abc 45 | LOP_MUL, abc 46 | LOP_DIV, abc 47 | LOP_MOD, abc 48 | LOP_POW, abc 49 | LOP_ADDK, abc 50 | LOP_SUBK, abc 51 | LOP_MULK, abc 52 | LOP_DIVK, abc 53 | LOP_MODK, abc 54 | LOP_POWK, abc 55 | LOP_AND, abc 56 | LOP_OR, abc 57 | LOP_ANDK, abc 58 | LOP_ORK, abc 59 | LOP_CONCAT, abc 60 | LOP_NOT, ab (c) 61 | LOP_MINUS, ab (c) 62 | LOP_LENGTH, ab (c) 63 | LOP_NEWTABLE, ab (c) aux 64 | LOP_DUPTABLE, ad 65 | LOP_SETLIST, abc aux 66 | LOP_FORNPREP, ad 67 | LOP_FORNLOOP, ad 68 | LOP_FORGLOOP, ad aux 69 | LOP_FORGPREP_INEXT, ad 70 | LOP_FORGLOOP_INEXT, ad 71 | LOP_FORGPREP_NEXT, ad 72 | LOP_FORGLOOP_NEXT, ad 73 | LOP_GETVARARGS, ab (c) 74 | LOP_DUPCLOSURE, ad 75 | LOP_PREPVARARGS, a (bc) 76 | LOP_LOADKX, a (bc) aux 77 | LOP_JUMPX, e 78 | LOP_FASTCALL, a (b) c 79 | LOP_COVERAGE, e 80 | LOP_CAPTURE, ab (c) 81 | LOP_JUMPIFEQK, ad aux 82 | LOP_JUMPIFNOTEQK, ad aux 83 | LOP_FASTCALL1, abc 84 | LOP_FASTCALL2, abc aux 85 | LOP_FASTCALL2K, abc aux 86 | LOP_FORGPREP, ad 87 | 88 | LOP_IDIV, abc 89 | 90 | store aud mh 91 | 92 | */ 93 | 94 | #[derive(Debug, Clone, Copy)] 95 | pub enum Instruction { 96 | BC { 97 | op_code: OpCode, 98 | a: u8, 99 | b: u8, 100 | c: u8, 101 | aux: u32, 102 | }, 103 | AD { 104 | op_code: OpCode, 105 | a: u8, 106 | d: i16, 107 | aux: u32, 108 | }, 109 | E { 110 | op_code: OpCode, 111 | e: i32, 112 | }, 113 | } 114 | 115 | impl Instruction { 116 | pub fn parse(insn: u32, encode_key: u8) -> Result { 117 | let op_code = (insn & 0xFF) as u8; 118 | let op_code = op_code.wrapping_mul(encode_key); 119 | match op_code { 120 | 0 121 | | 1 122 | | 2 123 | | 3 124 | | 6..=11 125 | | 13..=18 126 | | 20..=22 127 | | 33..=53 128 | | 55 129 | | 60 130 | | 63 131 | | 65 132 | | 66 133 | | 68 134 | | 70 135 | | 71..=75 136 | | 81 137 | | 82 => { 138 | let (a, b, c) = Self::parse_abc(insn); 139 | 140 | Ok(Self::BC { 141 | op_code: OpCode::try_from(op_code).unwrap(), 142 | a, 143 | b, 144 | c, 145 | aux: 0, 146 | }) 147 | } 148 | 4 | 5 | 12 | 19 | 23..=32 | 54 | 56..=59 | 61 | 62 | 64 | 76..=80 => { 149 | let (a, d) = Self::parse_ad(insn); 150 | 151 | Ok(Self::AD { 152 | op_code: OpCode::try_from(op_code).unwrap(), 153 | a, 154 | d, 155 | aux: 0, 156 | }) 157 | } 158 | 67 | 69 => { 159 | let e = Self::parse_e(insn); 160 | 161 | Ok(Self::E { 162 | op_code: OpCode::try_from(op_code).unwrap(), 163 | e, 164 | }) 165 | } 166 | 97 => Ok(Self::BC { 167 | op_code: OpCode::try_from(0).unwrap(), 168 | a: 0, 169 | b: 0, 170 | c: 0, 171 | aux: 0, 172 | }), 173 | _ => unreachable!("{}", op_code), 174 | } 175 | } 176 | 177 | fn parse_abc(insn: u32) -> (u8, u8, u8) { 178 | let a = ((insn >> 8) & 0xFF) as u8; 179 | let b = ((insn >> 16) & 0xFF) as u8; 180 | let c = ((insn >> 24) & 0xFF) as u8; 181 | 182 | (a, b, c) 183 | } 184 | 185 | fn parse_ad(insn: u32) -> (u8, i16) { 186 | let a = ((insn >> 8) & 0xFF) as u8; 187 | let d = ((insn >> 16) & 0xFFFF) as i16; 188 | 189 | (a, d) 190 | } 191 | 192 | fn parse_e(insn: u32) -> i32 { 193 | (insn as i32) >> 8 194 | } 195 | } 196 | -------------------------------------------------------------------------------- /ast/src/name_locals.rs: -------------------------------------------------------------------------------- 1 | use rustc_hash::FxHashSet; 2 | use triomphe::Arc; 3 | 4 | use crate::{Block, RValue, RcLocal, Statement, Traverse, Upvalue}; 5 | 6 | struct Namer { 7 | rename: bool, 8 | counter: usize, 9 | upvalues: FxHashSet, 10 | } 11 | 12 | impl Namer { 13 | fn name_local(&mut self, prefix: &str, local: &RcLocal) { 14 | let mut lock = local.0 .0.lock(); 15 | if self.rename || lock.0.is_none() { 16 | // TODO: hacky and slow 17 | if Arc::count(&local.0 .0) == 1 { 18 | lock.0 = Some("_".to_string()); 19 | } else { 20 | let prefix = prefix.to_string() 21 | + if self.upvalues.contains(local) { 22 | "_u_" 23 | } else { 24 | "" 25 | }; 26 | lock.0 = Some(format!("{}{}", prefix, self.counter)); 27 | self.counter += 1; 28 | } 29 | } 30 | } 31 | 32 | fn name_locals(&mut self, block: &mut Block) { 33 | for statement in &mut block.0 { 34 | // TODO: traverse_rvalues 35 | statement.post_traverse_values(&mut |value| -> Option<()> { 36 | if let itertools::Either::Right(RValue::Closure(closure)) = value { 37 | let mut function = closure.function.lock(); 38 | for param in &function.parameters { 39 | self.name_local("p", param); 40 | } 41 | self.name_locals(&mut function.body); 42 | }; 43 | None 44 | }); 45 | match statement { 46 | Statement::Assign(assign) if assign.prefix => { 47 | for lvalue in &assign.left { 48 | self.name_local("v", lvalue.as_local().unwrap()); 49 | } 50 | } 51 | Statement::If(r#if) => { 52 | self.name_locals(&mut r#if.then_block.lock()); 53 | self.name_locals(&mut r#if.else_block.lock()); 54 | } 55 | Statement::While(r#while) => { 56 | self.name_locals(&mut r#while.block.lock()); 57 | } 58 | Statement::Repeat(repeat) => { 59 | self.name_locals(&mut repeat.block.lock()); 60 | } 61 | Statement::NumericFor(numeric_for) => { 62 | self.name_local("v", &numeric_for.counter); 63 | self.name_locals(&mut numeric_for.block.lock()); 64 | } 65 | Statement::GenericFor(generic_for) => { 66 | for res_local in &generic_for.res_locals { 67 | self.name_local("v", res_local); 68 | } 69 | self.name_locals(&mut generic_for.block.lock()); 70 | } 71 | _ => {} 72 | } 73 | } 74 | } 75 | 76 | // TODO: does this need to be mut? 77 | fn find_upvalues(&mut self, block: &mut Block) { 78 | for statement in &mut block.0 { 79 | // TODO: traverse_values 80 | // TODO: doesnt need to be mut 81 | statement.post_traverse_values(&mut |value| -> Option<()> { 82 | if let itertools::Either::Right(RValue::Closure(closure)) = value { 83 | self.upvalues.extend( 84 | closure 85 | .upvalues 86 | .iter() 87 | .map(|u| match u { 88 | Upvalue::Copy(l) | Upvalue::Ref(l) => l, 89 | }) 90 | .cloned(), 91 | ); 92 | self.find_upvalues(&mut closure.function.lock().body); 93 | }; 94 | None 95 | }); 96 | match statement { 97 | Statement::If(r#if) => { 98 | self.find_upvalues(&mut r#if.then_block.lock()); 99 | self.find_upvalues(&mut r#if.else_block.lock()); 100 | } 101 | Statement::While(r#while) => { 102 | self.find_upvalues(&mut r#while.block.lock()); 103 | } 104 | Statement::Repeat(repeat) => { 105 | self.find_upvalues(&mut repeat.block.lock()); 106 | } 107 | Statement::NumericFor(numeric_for) => { 108 | self.find_upvalues(&mut numeric_for.block.lock()); 109 | } 110 | Statement::GenericFor(generic_for) => { 111 | self.find_upvalues(&mut generic_for.block.lock()); 112 | } 113 | _ => {} 114 | } 115 | } 116 | } 117 | } 118 | 119 | pub fn name_locals(block: &mut Block, rename: bool) { 120 | let mut namer = Namer { 121 | rename, 122 | counter: 1, 123 | upvalues: FxHashSet::default(), 124 | }; 125 | namer.find_upvalues(block); 126 | namer.name_locals(block); 127 | } 128 | -------------------------------------------------------------------------------- /cfg/src/ssa/destruct/liveness.rs: -------------------------------------------------------------------------------- 1 | use ast::{LocalRw, RcLocal}; 2 | use rustc_hash::{FxHashMap, FxHashSet}; 3 | 4 | use petgraph::{stable_graph::NodeIndex, visit::Walker}; 5 | 6 | use crate::function::Function; 7 | 8 | #[derive(Debug, Default)] 9 | struct BlockLiveness<'a> { 10 | // the locals that are used in this block 11 | uses: FxHashSet<&'a RcLocal>, 12 | // the locals that are defined in this block 13 | defs: FxHashSet<&'a RcLocal>, 14 | // the locals that are used by arguments passed from this block to its successor 15 | arg_out_uses: FxHashSet<&'a RcLocal>, 16 | // the locals that are defined by the parameters passed to this block by its predecessor 17 | params: FxHashSet<&'a RcLocal>, 18 | live_sets: LiveSets, 19 | } 20 | 21 | #[derive(Debug, Default)] 22 | pub struct LiveSets { 23 | // the set LiveIn(B) = params(B) ⋃ ( [uses(B) ⋃ live_out(B)] ∖ defs(B)) 24 | pub live_in: FxHashSet, 25 | // the set LiveOut(B) = ( ⋃_{S ∊ successor(B)} [live_in(S)∖params(S)] ) ⋃ arg_out_uses(B) 26 | pub live_out: FxHashSet, 27 | } 28 | 29 | #[derive(Debug)] 30 | pub struct Liveness<'a> { 31 | block_liveness: FxHashMap>, 32 | } 33 | 34 | impl<'a> Liveness<'a> { 35 | fn explore_all_paths( 36 | liveness: &mut Liveness, 37 | function: &'a Function, 38 | node: NodeIndex, 39 | variable: &'a RcLocal, 40 | ) { 41 | let mut stack = vec![node]; 42 | while let Some(node) = stack.pop() { 43 | let block_liveness = liveness.block_liveness.get_mut(&node).unwrap(); 44 | if block_liveness.defs.contains(variable) 45 | // block already visited 46 | || block_liveness.live_sets.live_in.contains(variable) 47 | { 48 | continue; 49 | } 50 | block_liveness.live_sets.live_in.insert(variable.clone()); 51 | if block_liveness.params.contains(variable) { 52 | continue; 53 | } 54 | for pred in function.predecessor_blocks(node) { 55 | liveness 56 | .block_liveness 57 | .get_mut(&pred) 58 | .unwrap() 59 | .live_sets 60 | .live_out 61 | .insert(variable.clone()); 62 | stack.push(pred); 63 | } 64 | } 65 | } 66 | 67 | pub fn calculate(function: &'a Function) -> FxHashMap { 68 | let mut liveness = Liveness { 69 | block_liveness: FxHashMap::with_capacity_and_hasher( 70 | function.graph().node_count(), 71 | Default::default(), 72 | ), 73 | }; 74 | for (node, block) in function.blocks() { 75 | let block_liveness = liveness.block_liveness.entry(node).or_default(); 76 | for instruction in block.iter() { 77 | block_liveness 78 | .uses 79 | .extend(instruction.values_read().into_iter()); 80 | block_liveness 81 | .defs 82 | .extend(instruction.values_written().into_iter()); 83 | } 84 | for (pred, edge) in function.edges_to_block(node) { 85 | liveness 86 | .block_liveness 87 | .get_mut(&node) 88 | .unwrap() 89 | .params 90 | .extend(edge.arguments.iter().map(|(k, _)| k)); 91 | let block_liveness = liveness.block_liveness.entry(pred).or_default(); 92 | block_liveness 93 | .arg_out_uses 94 | .extend(edge.arguments.iter().flat_map(|(_, v)| v.values_read())); 95 | } 96 | } 97 | for node in function.graph().node_indices() { 98 | let block_liveness = liveness.block_liveness.get_mut(&node).unwrap(); 99 | block_liveness.live_sets.live_in.reserve( 100 | block_liveness.params.len() 101 | + block_liveness 102 | .uses 103 | .len() 104 | .saturating_sub(block_liveness.defs.len()), 105 | ); 106 | let arg_out_uses = std::mem::take(&mut block_liveness.arg_out_uses); 107 | block_liveness 108 | .live_sets 109 | .live_out 110 | .reserve(arg_out_uses.len()); 111 | for variable in arg_out_uses { 112 | let block_liveness = liveness.block_liveness.get_mut(&node).unwrap(); 113 | block_liveness.live_sets.live_out.insert(variable.clone()); 114 | Self::explore_all_paths(&mut liveness, function, node, variable); 115 | } 116 | let block_liveness = liveness.block_liveness.get_mut(&node).unwrap(); 117 | for variable in block_liveness.uses.clone() { 118 | Self::explore_all_paths(&mut liveness, function, node, variable); 119 | } 120 | } 121 | liveness 122 | .block_liveness 123 | .into_iter() 124 | .map(|(n, l)| (n, l.live_sets)) 125 | .collect() 126 | } 127 | } 128 | -------------------------------------------------------------------------------- /cfg/src/ssa/upvalues.rs: -------------------------------------------------------------------------------- 1 | use petgraph::stable_graph::NodeIndex; 2 | use rangemap::RangeInclusiveMap; 3 | use rustc_hash::{FxHashMap, FxHashSet}; 4 | 5 | use crate::function::Function; 6 | 7 | #[derive(Debug)] 8 | pub(crate) struct UpvaluesOpen { 9 | pub open: FxHashMap< 10 | NodeIndex, 11 | FxHashMap>>, 12 | >, 13 | old_locals: FxHashMap, 14 | } 15 | 16 | impl UpvaluesOpen { 17 | pub fn new(function: &Function, old_locals: FxHashMap) -> Self { 18 | let mut this = Self { 19 | open: Default::default(), 20 | old_locals, 21 | }; 22 | let entry = function.entry().unwrap(); 23 | let mut stack = vec![entry]; 24 | let mut visited = FxHashSet::default(); 25 | while let Some(node) = stack.pop() { 26 | visited.insert(node); 27 | let block = function.block(node).unwrap(); 28 | let block_opened = this.open.entry(node).or_default(); 29 | for (stat_index, statement) in block.iter().enumerate() { 30 | // TODO: use traverse rvalues instead 31 | // this is because the lifter isnt guaranteed to be lifting bytecode 32 | // it could be lifting lua source code for deobfuscation purposes 33 | if let ast::Statement::Assign(assign) = statement { 34 | for opened in assign 35 | .right 36 | .iter() 37 | .filter_map(|r| r.as_closure()) 38 | .flat_map(|c| c.upvalues.iter()) 39 | .filter_map(|u| match u { 40 | ast::Upvalue::Copy(_) => None, 41 | ast::Upvalue::Ref(l) => Some(l), 42 | }) 43 | .map(|l| this.old_locals[l].clone()) 44 | { 45 | let open_ranges = block_opened.entry(opened).or_default(); 46 | let mut open_locations = Vec::new(); 47 | if let Some((_prev_range, prev_locations)) = 48 | open_ranges.get_key_value(&stat_index) 49 | { 50 | // TODO: this assert fails in Luau with the below code, 51 | // but i dont know why. it appears to work fine with the 52 | // assert commented out, but we should double check it. 53 | /* 54 | local u = a 55 | 56 | if u then 57 | print'hi' 58 | end 59 | 60 | function f() 61 | return u 62 | end 63 | */ 64 | // assert!(prev_range.contains(&(block.len() - 1))); 65 | open_locations.extend(prev_locations); 66 | } 67 | open_locations.push((node, stat_index)); 68 | open_ranges.insert(stat_index..=block.len() - 1, open_locations); 69 | } 70 | } else if let ast::Statement::Close(close) = statement { 71 | for closed in &close.locals { 72 | if let Some(open_ranges) = block_opened.get_mut(closed) { 73 | open_ranges.remove(stat_index..=block.len() - 1); 74 | } 75 | } 76 | } 77 | } 78 | for successor in function.successor_blocks(node) { 79 | // TODO: is there any case where successor is visited but has open stuff 80 | // that wasnt already discovered? 81 | // maybe possible with multiple opens 82 | if !visited.contains(&successor) { 83 | let successor_block = function.block(successor).unwrap(); 84 | let open_at_end = this.open[&node] 85 | .iter() 86 | .filter_map(|(l, m)| { 87 | Some((l.clone(), m.get(&(block.len().saturating_sub(1)))?.clone())) 88 | }) 89 | .collect::>(); 90 | let successor_open = this.open.entry(successor).or_default(); 91 | for (open, mut locations) in open_at_end { 92 | let open_ranges = successor_open.entry(open).or_default(); 93 | // TODO: sorta ugly doing a saturating subtraction, use uninclusive ranges instead? 94 | let range = 0..=successor_block.len().saturating_sub(1); 95 | if let Some((prev_range, prev_locations)) = open_ranges.get_key_value(&0) { 96 | assert_eq!(prev_range, &range); 97 | locations.extend(prev_locations); 98 | } 99 | open_ranges.insert(range, locations); 100 | } 101 | 102 | stack.push(successor); 103 | } 104 | } 105 | } 106 | this 107 | } 108 | } 109 | -------------------------------------------------------------------------------- /cfg/src/function.rs: -------------------------------------------------------------------------------- 1 | use ast::{LocalRw, RcLocal}; 2 | use contracts::requires; 3 | 4 | use petgraph::{ 5 | stable_graph::{EdgeReference, Neighbors, NodeIndex, StableDiGraph}, 6 | visit::{EdgeRef, IntoEdgesDirected}, 7 | Direction, 8 | }; 9 | 10 | use crate::block::{BlockEdge, BranchType}; 11 | 12 | #[derive(Debug, Clone, Default)] 13 | pub struct Function { 14 | pub id: usize, 15 | pub name: Option, 16 | pub parameters: Vec, 17 | pub is_variadic: bool, 18 | graph: StableDiGraph, 19 | entry: Option, 20 | } 21 | 22 | impl Function { 23 | pub fn new(id: usize) -> Self { 24 | Self { 25 | id, 26 | name: None, 27 | parameters: Vec::new(), 28 | is_variadic: false, 29 | graph: StableDiGraph::new(), 30 | entry: None, 31 | } 32 | } 33 | 34 | pub fn name_mut(&mut self) -> &mut Option { 35 | &mut self.name 36 | } 37 | 38 | pub fn entry(&self) -> &Option { 39 | &self.entry 40 | } 41 | 42 | #[requires(self.has_block(new_entry))] 43 | pub fn set_entry(&mut self, new_entry: NodeIndex) { 44 | self.entry = Some(new_entry); 45 | } 46 | 47 | pub fn graph(&self) -> &StableDiGraph { 48 | &self.graph 49 | } 50 | 51 | pub fn graph_mut(&mut self) -> &mut StableDiGraph { 52 | &mut self.graph 53 | } 54 | 55 | pub fn has_block(&self, block: NodeIndex) -> bool { 56 | self.graph.contains_node(block) 57 | } 58 | 59 | pub fn block(&self, block: NodeIndex) -> Option<&ast::Block> { 60 | self.graph.node_weight(block) 61 | } 62 | 63 | pub fn block_mut(&mut self, block: NodeIndex) -> Option<&mut ast::Block> { 64 | self.graph.node_weight_mut(block) 65 | } 66 | 67 | pub fn blocks(&self) -> impl Iterator { 68 | self.graph 69 | .node_indices() 70 | .map(|i| (i, self.graph.node_weight(i).unwrap())) 71 | } 72 | 73 | pub fn blocks_mut(&mut self) -> impl Iterator { 74 | self.graph.node_weights_mut() 75 | } 76 | 77 | pub fn successor_blocks(&self, block: NodeIndex) -> Neighbors { 78 | self.graph.neighbors_directed(block, Direction::Outgoing) 79 | } 80 | 81 | pub fn predecessor_blocks(&self, block: NodeIndex) -> Neighbors { 82 | self.graph.neighbors_directed(block, Direction::Incoming) 83 | } 84 | 85 | pub fn edges_to_block(&self, node: NodeIndex) -> impl Iterator { 86 | let mut edges = self.predecessor_blocks(node).detach(); 87 | std::iter::from_fn(move || edges.next_edge(&self.graph)).filter_map(move |e| { 88 | let (source, target) = self.graph.edge_endpoints(e).unwrap(); 89 | if target == node { 90 | Some((source, self.graph.edge_weight(e).unwrap())) 91 | } else { 92 | None 93 | } 94 | }) 95 | } 96 | 97 | pub fn edges(&self, node: NodeIndex) -> impl Iterator> { 98 | self.graph.edges_directed(node, Direction::Outgoing) 99 | } 100 | 101 | pub fn remove_edges(&mut self, node: NodeIndex) -> Vec<(NodeIndex, BlockEdge)> { 102 | let mut edges = Vec::new(); 103 | for (target, edge) in self 104 | .edges(node) 105 | .map(|e| (e.target(), e.id())) 106 | .collect::>() 107 | { 108 | edges.push((target, self.graph.remove_edge(edge).unwrap())); 109 | } 110 | edges 111 | } 112 | 113 | // returns previous edges 114 | pub fn set_edges( 115 | &mut self, 116 | node: NodeIndex, 117 | new_edges: Vec<(NodeIndex, BlockEdge)>, 118 | ) -> Vec<(NodeIndex, BlockEdge)> { 119 | let prev_edges = self.remove_edges(node); 120 | for (target, edge) in new_edges { 121 | self.graph.add_edge(node, target, edge); 122 | } 123 | prev_edges 124 | } 125 | 126 | pub fn conditional_edges( 127 | &self, 128 | node: NodeIndex, 129 | ) -> Option<(EdgeReference, EdgeReference)> { 130 | let edges = self 131 | .graph 132 | .edges_directed(node, Direction::Outgoing) 133 | .collect::>(); 134 | if let [e0, e1] = edges[..] { 135 | let mut res = (e0, e1); 136 | if res.1.weight().branch_type == BranchType::Then { 137 | std::mem::swap(&mut res.0, &mut res.1); 138 | } 139 | assert!(res.0.weight().branch_type == BranchType::Then); 140 | assert!(res.1.weight().branch_type == BranchType::Else); 141 | Some(res) 142 | } else { 143 | None 144 | } 145 | } 146 | 147 | pub fn unconditional_edge(&self, node: NodeIndex) -> Option> { 148 | let edges = self 149 | .graph 150 | .edges_directed(node, Direction::Outgoing) 151 | .collect::>(); 152 | if let [e] = edges[..] { 153 | Some(e) 154 | } else { 155 | None 156 | } 157 | } 158 | 159 | // TODO: disable_contracts for production builds 160 | #[requires(self.has_block(node))] 161 | pub fn values_read(&self, node: NodeIndex) -> impl Iterator { 162 | self.block(node) 163 | .unwrap() 164 | .0 165 | .iter() 166 | .flat_map(|s| s.values_read()) 167 | .chain(self.edges(node).flat_map(|e| { 168 | e.weight() 169 | .arguments 170 | .iter() 171 | .flat_map(|(_, a)| a.values_read()) 172 | })) 173 | } 174 | 175 | pub fn new_block(&mut self) -> NodeIndex { 176 | self.graph.add_node(ast::Block::default()) 177 | } 178 | 179 | pub fn remove_block(&mut self, block: NodeIndex) -> Option { 180 | self.graph.remove_node(block) 181 | } 182 | } 183 | -------------------------------------------------------------------------------- /restructure/src/jump.rs: -------------------------------------------------------------------------------- 1 | use ast::SideEffects; 2 | use cfg::block::{BlockEdge, BranchType}; 3 | use itertools::Itertools; 4 | use petgraph::{ 5 | algo::dominators::Dominators, 6 | stable_graph::NodeIndex, 7 | visit::{EdgeRef, IntoEdgeReferences}, 8 | Direction, 9 | }; 10 | 11 | impl super::GraphStructurer { 12 | // TODO: STYLE: better name 13 | // TODO: this is the same as in structuring.rs but w/o block params 14 | // maybe we can use the same function? 15 | pub(crate) fn try_remove_unnecessary_condition(&mut self, node: NodeIndex) -> bool { 16 | let block = self.function.block(node).unwrap(); 17 | if !block.is_empty() 18 | && block.last().unwrap().as_if().is_some() 19 | && let Some((then_edge, else_edge)) = self.function.conditional_edges(node) 20 | && then_edge.target() == else_edge.target() 21 | { 22 | let target = then_edge.target(); 23 | let cond = self 24 | .function 25 | .block_mut(node) 26 | .unwrap() 27 | .pop() 28 | .unwrap() 29 | .into_if() 30 | .unwrap() 31 | .condition; 32 | 33 | let new_stat = match cond { 34 | ast::RValue::Call(call) => Some(call.into()), 35 | ast::RValue::MethodCall(method_call) => Some(method_call.into()), 36 | cond if cond.has_side_effects() => Some( 37 | ast::Assign { 38 | left: vec![ast::RcLocal::default().into()], 39 | right: vec![cond], 40 | prefix: true, 41 | parallel: false, 42 | } 43 | .into(), 44 | ), 45 | _ => None, 46 | }; 47 | self.function.block_mut(node).unwrap().extend(new_stat); 48 | self.function.set_edges( 49 | node, 50 | vec![(target, BlockEdge::new(BranchType::Unconditional))], 51 | ); 52 | true 53 | } else { 54 | false 55 | } 56 | } 57 | 58 | pub(crate) fn match_jump(&mut self, node: NodeIndex, target: Option) -> bool { 59 | if let Some(target) = target { 60 | if node == target { 61 | return false; 62 | } 63 | if !self.is_for_next(node) { 64 | assert!(self.function.unconditional_edge(node).is_some()); 65 | if Self::block_is_no_op(self.function.block(node).unwrap()) 66 | && self.function.entry() != &Some(node) 67 | && !self.is_loop_header(node) 68 | { 69 | for (source, edge) in self 70 | .function 71 | .graph() 72 | .edges_directed(node, Direction::Incoming) 73 | .map(|e| (e.source(), e.id())) 74 | .collect::>() 75 | { 76 | let edge = self.function.graph_mut().remove_edge(edge).unwrap(); 77 | self.function.graph_mut().add_edge(source, target, edge); 78 | self.try_remove_unnecessary_condition(source); 79 | } 80 | self.function.remove_block(node); 81 | true 82 | } else if self.function.predecessor_blocks(target).count() == 1 83 | && !self.function.edges_to_block(node).any(|(t, _)| t == target) 84 | && !self 85 | .function 86 | .edges_to_block(target) 87 | .any(|(t, _)| t == target) 88 | { 89 | if self.function.entry() != &Some(target) 90 | && !self.is_loop_header(target) 91 | && !self.is_for_next(target) 92 | { 93 | let edges = self.function.remove_edges(target); 94 | let block = self.function.remove_block(target).unwrap(); 95 | self.function.block_mut(node).unwrap().extend(block.0); 96 | self.function.set_edges(node, edges); 97 | true 98 | } else if self.function.entry() != &Some(node) && !self.is_loop_header(node) { 99 | // TODO: test 100 | for (source, edge) in self 101 | .function 102 | .graph() 103 | .edges_directed(node, Direction::Incoming) 104 | .map(|e| (e.source(), e.id())) 105 | .collect::>() 106 | { 107 | let edge = self.function.graph_mut().remove_edge(edge).unwrap(); 108 | self.function.graph_mut().add_edge(source, target, edge); 109 | self.try_remove_unnecessary_condition(source); 110 | } 111 | let mut block = self.function.remove_block(node).unwrap(); 112 | block.extend(std::mem::take(self.function.block_mut(target).unwrap()).0); 113 | *self.function.block_mut(target).unwrap() = block; 114 | true 115 | } else { 116 | false 117 | } 118 | } else { 119 | false 120 | } 121 | } else { 122 | false 123 | } 124 | } 125 | // node is terminating 126 | // TODO: block_is_no_op returns true for blocks with comments, do we wanna remove the block if it has comments? 127 | else if Self::block_is_no_op(self.function.block(node).unwrap()) 128 | && self.function.entry() != &Some(node) 129 | && !self.is_loop_header(node) 130 | && !self.is_for_next(node) 131 | { 132 | let mut invalid = false; 133 | for pred in self.function.predecessor_blocks(node).collect_vec() { 134 | if self.function.successor_blocks(pred).collect_vec().len() != 1 { 135 | invalid = true; 136 | break; 137 | } 138 | } 139 | if !invalid { 140 | for edge in self 141 | .function 142 | .graph() 143 | .edges_directed(node, Direction::Incoming) 144 | .map(|e| e.id()) 145 | .collect::>() 146 | { 147 | assert_eq!( 148 | self.function 149 | .graph_mut() 150 | .remove_edge(edge) 151 | .unwrap() 152 | .branch_type, 153 | BranchType::Unconditional 154 | ); 155 | } 156 | self.function.remove_block(node); 157 | true 158 | } else { 159 | false 160 | } 161 | } else { 162 | false 163 | } 164 | } 165 | } 166 | -------------------------------------------------------------------------------- /luau-lifter/src/deserializer/function.rs: -------------------------------------------------------------------------------- 1 | use core::num; 2 | 3 | use nom::{ 4 | complete::take, 5 | number::complete::{le_u32, le_u8}, 6 | IResult, 7 | }; 8 | use nom_leb128::leb128_usize; 9 | 10 | use super::{ 11 | constant::Constant, 12 | list::{parse_list, parse_list_len}, 13 | }; 14 | 15 | use crate::{instruction::*, op_code::OpCode}; 16 | 17 | #[derive(Debug)] 18 | pub struct Function { 19 | pub max_stack_size: u8, 20 | pub num_parameters: u8, 21 | pub num_upvalues: u8, 22 | pub is_vararg: bool, 23 | //pub instructions: Vec, 24 | pub instructions: Vec, 25 | pub constants: Vec, 26 | pub functions: Vec, 27 | pub line_defined: usize, 28 | pub function_name: usize, 29 | pub line_gap_log2: Option, 30 | pub line_info_delta: Option>, 31 | pub abs_line_info_delta: Option>, 32 | } 33 | 34 | impl Function { 35 | fn parse_instructions(vec: &Vec, encode_key: u8) -> Vec { 36 | let mut v: Vec = Vec::new(); 37 | let mut pc = 0; 38 | 39 | loop { 40 | let ins = Instruction::parse(vec[pc], encode_key).unwrap(); 41 | let op = match ins { 42 | Instruction::BC { op_code, .. } => op_code, 43 | Instruction::AD { op_code, .. } => op_code, 44 | Instruction::E { op_code, .. } => op_code, 45 | }; 46 | 47 | // handle ops with aux values 48 | match op { 49 | OpCode::LOP_GETGLOBAL 50 | | OpCode::LOP_SETGLOBAL 51 | | OpCode::LOP_GETIMPORT 52 | | OpCode::LOP_GETTABLEKS 53 | | OpCode::LOP_SETTABLEKS 54 | | OpCode::LOP_NAMECALL 55 | | OpCode::LOP_JUMPIFEQ 56 | | OpCode::LOP_JUMPIFLE 57 | | OpCode::LOP_JUMPIFLT 58 | | OpCode::LOP_JUMPIFNOTEQ 59 | | OpCode::LOP_JUMPIFNOTLE 60 | | OpCode::LOP_JUMPIFNOTLT 61 | | OpCode::LOP_NEWTABLE 62 | | OpCode::LOP_SETLIST 63 | | OpCode::LOP_FORGLOOP 64 | | OpCode::LOP_LOADKX 65 | | OpCode::LOP_FASTCALL2 66 | | OpCode::LOP_FASTCALL2K 67 | | OpCode::LOP_FASTCALL3 68 | | OpCode::LOP_JUMPXEQKNIL 69 | | OpCode::LOP_JUMPXEQKB 70 | | OpCode::LOP_JUMPXEQKN 71 | | OpCode::LOP_JUMPXEQKS => { 72 | let aux = vec[pc + 1]; 73 | pc += 2; 74 | match ins { 75 | Instruction::BC { 76 | op_code, a, b, c, .. 77 | } => { 78 | v.push(Instruction::BC { 79 | op_code, 80 | a, 81 | b, 82 | c, 83 | aux, 84 | }); 85 | } 86 | Instruction::AD { op_code, a, d, .. } => { 87 | v.push(Instruction::AD { op_code, a, d, aux }); 88 | } 89 | _ => unreachable!(), 90 | } 91 | v.push(Instruction::BC { 92 | op_code: OpCode::LOP_NOP, 93 | a: 0, 94 | b: 0, 95 | c: 0, 96 | aux: 0, 97 | }); 98 | } 99 | _ => { 100 | v.push(ins); 101 | pc += 1; 102 | } 103 | } 104 | 105 | if pc == vec.len() { 106 | break; 107 | } 108 | } 109 | 110 | v 111 | } 112 | 113 | pub(crate) fn parse(input: &[u8], encode_key: u8) -> IResult<&[u8], Self> { 114 | let (input, max_stack_size) = le_u8(input)?; 115 | let (input, num_parameters) = le_u8(input)?; 116 | let (input, num_upvalues) = le_u8(input)?; 117 | let (input, is_vararg) = le_u8(input)?; 118 | 119 | let (input, flags) = le_u8(input)?; 120 | let (input, _) = parse_list(input, le_u8)?; 121 | 122 | let (input, u32_instructions) = parse_list(input, le_u32)?; 123 | //let (input, instructions) = parse_list(input, Function::parse_instrution)?; 124 | let instructions = Self::parse_instructions(&u32_instructions, encode_key); 125 | let (input, constants) = parse_list(input, Constant::parse)?; 126 | let (input, functions) = parse_list(input, leb128_usize)?; 127 | let (input, line_defined) = leb128_usize(input)?; 128 | let (input, function_name) = leb128_usize(input)?; 129 | let (input, has_line_info) = le_u8(input)?; 130 | let (input, line_gap_log2) = match has_line_info { 131 | 0 => (input, None), 132 | _ => { 133 | let (input, line_gap_log2) = le_u8(input)?; 134 | (input, Some(line_gap_log2)) 135 | } 136 | }; 137 | let (input, line_info_delta) = match has_line_info { 138 | 0 => (input, None), 139 | _ => { 140 | let (input, line_info_delta) = 141 | parse_list_len(input, le_u8, u32_instructions.len())?; 142 | (input, Some(line_info_delta)) 143 | } 144 | }; 145 | let (input, abs_line_info_delta) = match has_line_info { 146 | 0 => (input, None), 147 | _ => { 148 | let (input, abs_line_info_delta) = parse_list_len( 149 | input, 150 | le_u32, 151 | ((u32_instructions.len() - 1) >> line_gap_log2.unwrap()) + 1, 152 | )?; 153 | (input, Some(abs_line_info_delta)) 154 | } 155 | }; 156 | let input = match le_u8(input)? { 157 | (input, 0) => input, 158 | (input, _) => { 159 | panic!("we have debug info"); 160 | let (mut input, num_locvars) = leb128_usize(input)?; 161 | for _ in 0..num_locvars { 162 | (input, _) = leb128_usize(input)?; 163 | (input, _) = leb128_usize(input)?; 164 | (input, _) = leb128_usize(input)?; 165 | (input, _) = le_u8(input)?; 166 | } 167 | let (mut input, num_upvalues) = leb128_usize(input)?; 168 | for _ in 0..num_upvalues { 169 | (input, _) = leb128_usize(input)?; 170 | } 171 | input 172 | } 173 | }; 174 | Ok(( 175 | input, 176 | Self { 177 | max_stack_size, 178 | num_parameters, 179 | num_upvalues, 180 | is_vararg: is_vararg != 0u8, 181 | instructions, 182 | constants, 183 | functions, 184 | line_defined, 185 | function_name, 186 | line_gap_log2, 187 | line_info_delta, 188 | abs_line_info_delta, 189 | }, 190 | )) 191 | } 192 | } 193 | -------------------------------------------------------------------------------- /ast/src/type_system.rs: -------------------------------------------------------------------------------- 1 | use crate::{Block, RcLocal}; 2 | use itertools::Itertools; 3 | use std::{ 4 | borrow::Cow, 5 | collections::{BTreeMap, BTreeSet}, 6 | fmt::{Display, Formatter}, 7 | }; 8 | 9 | #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)] 10 | pub enum Type { 11 | Any, 12 | Nil, 13 | Boolean, 14 | Number, 15 | String, 16 | Table { 17 | indexer: Box<(Type, Type)>, 18 | fields: BTreeMap, 19 | }, 20 | Function(Vec, Vec), 21 | Optional(Box), 22 | Union(BTreeSet), 23 | Intersection(BTreeSet), 24 | VarArg, 25 | Vector, 26 | } 27 | 28 | impl Type { 29 | pub fn is_subtype_of(&self, t: &Self) -> bool { 30 | match t { 31 | Self::Any => true, 32 | Self::Table { 33 | box indexer, 34 | fields, 35 | } => { 36 | let t_fields = fields; 37 | let (indexer_type, element_type) = indexer; 38 | 39 | match self { 40 | Self::Table { 41 | box indexer, 42 | fields, 43 | } if indexer.0.is_subtype_of(indexer_type) 44 | && indexer.1.is_subtype_of(element_type) => 45 | { 46 | t_fields 47 | .keys() 48 | .all(|k| fields[k].is_subtype_of(&t_fields[k])) 49 | } 50 | _ => false, 51 | } 52 | } 53 | Self::Union(union) => match self { 54 | Self::Union(u) => union.iter().all(|t| u.contains(t)), 55 | _ => union.contains(self), 56 | }, 57 | _ => false, 58 | } 59 | } 60 | 61 | pub fn precedence(&self) -> usize { 62 | match self { 63 | Self::Any => 0, 64 | Self::Nil => 0, 65 | Self::Boolean => 0, 66 | Self::Number => 0, 67 | Self::String => 0, 68 | Self::Table { .. } => 1, 69 | Self::Function(_, _) => 1, 70 | Self::Optional(_) => 0, 71 | Self::Union(_) => 2, 72 | Self::Intersection(_) => 2, 73 | Self::VarArg => 0, 74 | Self::Vector => 0, 75 | } 76 | } 77 | } 78 | 79 | impl Display for Type { 80 | fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { 81 | write!( 82 | f, 83 | "{}", 84 | match self { 85 | Type::Any => Cow::Borrowed("any"), 86 | Type::Nil => Cow::Borrowed("nil"), 87 | Type::Boolean => Cow::Borrowed("boolean"), 88 | Type::Number => Cow::Borrowed("number"), 89 | Type::String => Cow::Borrowed("string"), 90 | Type::Table { indexer, fields } => { 91 | let (indexer_type, element_type) = indexer.as_ref(); 92 | 93 | Cow::Owned(format!( 94 | "{{{}{}{}}}", 95 | if indexer_type == &Type::Number && fields.is_empty() { 96 | element_type.to_string() 97 | } else { 98 | format!("[{}]: {}", indexer_type, element_type) 99 | }, 100 | (!fields.is_empty()).then_some(", ").unwrap_or_default(), 101 | fields 102 | .iter() 103 | .map(|(field, r#type)| { format!("{}: {}", field, r#type) }) 104 | .join(", ") 105 | )) 106 | } 107 | Type::Function(domain, codomain) => Cow::Owned(format!( 108 | "({}) -> {}", 109 | domain.iter().join(", "), 110 | if (codomain.len() == 1 && self.precedence() >= codomain[0].precedence()) 111 | || codomain.len() > 1 112 | { 113 | format!("({})", codomain.iter().join(", ")) 114 | } else { 115 | codomain.iter().join(", ") 116 | } 117 | )), 118 | Type::Optional(r#type) => Cow::Owned(format!("{}?", r#type)), 119 | Type::Union(types) => { 120 | Cow::Owned(types.iter().join(" | ")) 121 | } 122 | Type::Intersection(types) => { 123 | Cow::Owned(types.iter().join(" & ")) 124 | } 125 | Type::VarArg => Cow::Borrowed("..."), 126 | Type::Vector => Cow::Borrowed("vector"), 127 | } 128 | ) 129 | } 130 | } 131 | 132 | pub struct TypeSystem<'a> { 133 | // TODO: use hash map? 134 | annotations: BTreeMap<&'a RcLocal, &'a mut Type>, 135 | } 136 | 137 | impl<'a> TypeSystem<'a> { 138 | pub fn analyze(block: &'a mut Block) { 139 | let mut system = Self { 140 | annotations: BTreeMap::new(), 141 | }; 142 | 143 | system.analyze_block(block); 144 | } 145 | 146 | pub fn analyze_block(&mut self, _block: &'a mut Block) -> Vec { 147 | todo!() 148 | // let mut return_values = Vec::new(); 149 | 150 | // for statement in &mut block.0 { 151 | // match statement { 152 | // /*Statement::Assign(assign) => { 153 | // for ((lvalue, annotation), rvalue) in 154 | // assign.left.iter_mut().zip(assign.right.iter_mut()) 155 | // { 156 | // let r#type = rvalue.infer(self); 157 | 158 | // if let LValue::Local(local) = lvalue { 159 | // if let Some(annotation) = self.annotations.get_mut(local) { 160 | // if let Type::Union(types) = annotation { 161 | // types.insert(r#type); 162 | // } else { 163 | // let mut types = BTreeSet::new(); 164 | 165 | // types.insert(annotation.clone()); 166 | // types.insert(r#type); 167 | 168 | // **annotation = Type::Union(types); 169 | // } 170 | // } else { 171 | // *annotation = Some(r#type); 172 | // self.annotations.insert(local, annotation.as_mut().unwrap()); 173 | // } 174 | // } 175 | // } 176 | // }*/ 177 | // Statement::If(r#if) => { 178 | // self.analyze_block(&mut r#if.then_block); 179 | // self.analyze_block(&mut r#if.else_block); 180 | // } 181 | // // TODO: repeat and for loops 182 | // Statement::While(r#while) => { 183 | // self.analyze_block(&mut r#while.block.lock()); 184 | // } 185 | // Statement::Return(r#return) => { 186 | // return_values.extend(r#return.values.iter_mut().map(|v| v.infer(self))); 187 | // } 188 | // // TODO: numericfor, genericfor 189 | // _ => {} 190 | // } 191 | // } 192 | 193 | // return_values 194 | } 195 | 196 | pub fn type_of(&self, local: &RcLocal) -> &Type { 197 | self.annotations 198 | .get(local) 199 | .map(|t| t as &Type) 200 | .unwrap_or(&Type::Any) 201 | } 202 | } 203 | 204 | pub trait Infer { 205 | fn infer<'a: 'b, 'b>(&'a mut self, system: &mut TypeSystem<'b>) -> Type; 206 | } 207 | -------------------------------------------------------------------------------- /ast/src/local_declarations.rs: -------------------------------------------------------------------------------- 1 | use std::collections::BTreeMap; 2 | 3 | use array_tool::vec::Intersect; 4 | use by_address::ByAddress; 5 | use indexmap::{IndexMap, IndexSet}; 6 | use itertools::Itertools; 7 | use parking_lot::Mutex; 8 | use petgraph::{ 9 | algo::dominators::simple_fast, 10 | prelude::{DiGraph, NodeIndex}, 11 | Direction, 12 | }; 13 | use rustc_hash::{FxHashMap, FxHashSet}; 14 | use triomphe::Arc; 15 | 16 | use crate::{Assign, Block, LocalRw, RcLocal, Statement}; 17 | 18 | #[derive(Default)] 19 | pub struct LocalDeclarer { 20 | block_to_node: FxHashMap>>, NodeIndex>, 21 | graph: DiGraph<(Option>>, usize), ()>, 22 | local_usages: IndexMap>, 23 | declarations: FxHashMap>>, BTreeMap>>, 24 | } 25 | 26 | impl LocalDeclarer { 27 | fn visit(&mut self, block: Arc>, stat_index: usize) -> NodeIndex { 28 | let node = self.graph.add_node((Some(block.clone()), stat_index)); 29 | self.block_to_node.insert(block.clone().into(), node); 30 | for (stat_index, stat) in block.lock().iter().enumerate() { 31 | // for loops already declare their own locals :) 32 | if !matches!(stat, Statement::GenericFor(_) | Statement::NumericFor(_)) { 33 | // we only visit locals written because locals are guaranteed to be written 34 | // before they are read. 35 | // TODO: move to seperate function and visit breadth-first? 36 | for local in stat.values_written() { 37 | self.local_usages 38 | .entry(local.clone()) 39 | .or_default() 40 | .entry(node) 41 | .or_insert(stat_index); 42 | } 43 | } 44 | match stat { 45 | Statement::If(r#if) => { 46 | let if_node = self.graph.add_node((None, stat_index)); 47 | self.graph.add_edge(node, if_node, ()); 48 | let then_node = self.visit(r#if.then_block.clone(), stat_index); 49 | self.graph.add_edge(if_node, then_node, ()); 50 | let else_node = self.visit(r#if.else_block.clone(), stat_index); 51 | self.graph.add_edge(if_node, else_node, ()); 52 | } 53 | Statement::While(r#while) => { 54 | let child = self.visit(r#while.block.clone(), stat_index); 55 | self.graph.add_edge(node, child, ()); 56 | } 57 | Statement::Repeat(repeat) => { 58 | let child = self.visit(r#repeat.block.clone(), stat_index); 59 | self.graph.add_edge(node, child, ()); 60 | } 61 | Statement::NumericFor(numeric_for) => { 62 | let child = self.visit(r#numeric_for.block.clone(), stat_index); 63 | self.graph.add_edge(node, child, ()); 64 | } 65 | Statement::GenericFor(generic_for) => { 66 | let child = self.visit(r#generic_for.block.clone(), stat_index); 67 | self.graph.add_edge(node, child, ()); 68 | } 69 | _ => {} 70 | } 71 | } 72 | node 73 | } 74 | 75 | pub fn declare_locals( 76 | mut self, 77 | root_block: Arc>, 78 | locals_to_ignore: &FxHashSet, 79 | ) { 80 | let root_node = self.visit(root_block, 0); 81 | let dominators = simple_fast(&self.graph, root_node); 82 | for (local, usages) in self.local_usages { 83 | if locals_to_ignore.contains(&local) { 84 | continue; 85 | } 86 | let (mut node, mut first_stat_index) = if usages.len() == 1 { 87 | usages.into_iter().next().unwrap() 88 | } else { 89 | let node_dominators = usages 90 | .keys() 91 | .map(|&n| dominators.dominators(n).unwrap().collect_vec()) 92 | .collect_vec(); 93 | let mut dom_iter = node_dominators.iter().cloned(); 94 | let mut common_dominators = dom_iter.next().unwrap(); 95 | for node_dominators in dom_iter { 96 | common_dominators = common_dominators.intersect(node_dominators); 97 | } 98 | let common_dominator = common_dominators[0]; 99 | if let Some((_, first_stat_index)) = 100 | usages.into_iter().find(|&(n, _)| n == common_dominator) 101 | { 102 | (common_dominator, first_stat_index) 103 | } else { 104 | // find the left-most dominated node 105 | let mut first_stat_index = None; 106 | for child in self 107 | .graph 108 | .neighbors_directed(common_dominator, Direction::Outgoing) 109 | { 110 | for node_dominators in &node_dominators { 111 | if node_dominators.contains(&child) { 112 | first_stat_index = Some(self.graph.node_weight(child).unwrap().1); 113 | } 114 | } 115 | } 116 | (common_dominator, first_stat_index.unwrap()) 117 | } 118 | }; 119 | while let (block, parent_stat_index) = self.graph.node_weight(node).unwrap() 120 | && block.is_none() 121 | { 122 | let parent = self 123 | .graph 124 | .neighbors_directed(node, Direction::Incoming) 125 | .exactly_one() 126 | .unwrap(); 127 | (node, first_stat_index) = (parent, *parent_stat_index); 128 | } 129 | let block = self 130 | .graph 131 | .node_weight(node) 132 | .unwrap() 133 | .0 134 | .as_ref() 135 | .unwrap() 136 | .clone(); 137 | self.declarations 138 | .entry(block.into()) 139 | .or_default() 140 | .entry(first_stat_index) 141 | .or_default() 142 | .insert(local); 143 | } 144 | 145 | for (ByAddress(block), declarations) in self.declarations { 146 | let mut block = block.lock(); 147 | for (stat_index, mut locals) in declarations.into_iter().rev() { 148 | match &mut block[stat_index] { 149 | Statement::Assign(assign) 150 | if assign 151 | .left 152 | .iter() 153 | .all(|l| l.as_local().is_some_and(|l| locals.contains(l))) => 154 | { 155 | locals.retain(|l| { 156 | !assign 157 | .left 158 | .iter() 159 | .map(|l| l.as_local().unwrap()) 160 | .contains(l) 161 | }); 162 | assign.prefix = true; 163 | } 164 | _ => {} 165 | } 166 | if !locals.is_empty() { 167 | let mut declaration = 168 | Assign::new(locals.into_iter().map(|l| l.into()).collect_vec(), vec![]); 169 | declaration.prefix = true; 170 | block.insert(stat_index, declaration.into()); 171 | } 172 | } 173 | } 174 | } 175 | } 176 | -------------------------------------------------------------------------------- /lua51-lifter/src/main.rs: -------------------------------------------------------------------------------- 1 | #![feature(box_patterns)] 2 | #![feature(let_chains)] 3 | 4 | use ast::{ 5 | local_declarations::LocalDeclarer, name_locals::name_locals, replace_locals::replace_locals, 6 | Traverse, 7 | }; 8 | use by_address::ByAddress; 9 | use cfg::ssa::{ 10 | self, 11 | structuring::{structure_conditionals, structure_jumps, structure_method_calls}, 12 | }; 13 | use indexmap::IndexMap; 14 | use lifter::Lifter; 15 | use parking_lot::Mutex; 16 | use petgraph::algo::dominators::simple_fast; 17 | use rayon::iter::ParallelIterator; 18 | use rayon::prelude::IntoParallelIterator; 19 | use rustc_hash::FxHashMap; 20 | use std::{ 21 | fs::File, 22 | io::{Read, Write}, 23 | path::Path, 24 | time::Instant, 25 | }; 26 | use triomphe::Arc; 27 | 28 | use clap::Parser; 29 | 30 | use lua51_deserializer::chunk::Chunk; 31 | 32 | mod lifter; 33 | 34 | #[cfg(feature = "dhat-heap")] 35 | #[global_allocator] 36 | static ALLOC: dhat::Alloc = dhat::Alloc; 37 | 38 | #[derive(Parser, Debug)] 39 | #[clap(about, version, author)] 40 | struct Args { 41 | #[clap(short, long)] 42 | file: String, 43 | } 44 | 45 | fn main() -> anyhow::Result<()> { 46 | #[cfg(feature = "dhat-heap")] 47 | let _profiler = dhat::Profiler::new_heap(); 48 | 49 | let args = Args::parse(); 50 | let path = Path::new(&args.file); 51 | let mut input = File::open(path)?; 52 | let mut buffer = vec![0; input.metadata()?.len() as usize]; 53 | input.read_exact(&mut buffer)?; 54 | 55 | let start = Instant::now(); 56 | let chunk = Chunk::parse(&buffer).unwrap().1; 57 | let mut lifted = Vec::new(); 58 | let (function, upvalues) = Lifter::lift(&chunk.function, &mut lifted); 59 | lifted.push((Arc::>::default(), function, upvalues)); 60 | lifted.reverse(); 61 | 62 | let (main, ..) = lifted.first().unwrap().clone(); 63 | let mut upvalues = lifted 64 | .into_iter() 65 | .map(|(ast_function, mut function, upvalues_in)| { 66 | let (local_count, local_groups, upvalue_in_groups, upvalue_passed_groups) = 67 | cfg::ssa::construct(&mut function, &upvalues_in); 68 | let upvalue_to_group = upvalue_in_groups 69 | .into_iter() 70 | .chain( 71 | upvalue_passed_groups 72 | .into_iter() 73 | .map(|m| (ast::RcLocal::default(), m)), 74 | ) 75 | .flat_map(|(i, g)| g.into_iter().map(move |u| (u, i.clone()))) 76 | .collect::>(); 77 | // TODO: do we even need this? 78 | let local_to_group = local_groups 79 | .into_iter() 80 | .enumerate() 81 | .flat_map(|(i, g)| g.into_iter().map(move |l| (l, i))) 82 | .collect::>(); 83 | // TODO: REFACTOR: some way to write a macro that states 84 | // if cfg::ssa::inline results in change then structure_jumps, structure_compound_conditionals, 85 | // structure_for_loops and remove_unnecessary_params must run again. 86 | // if structure_compound_conditionals results in change then dominators and post dominators 87 | // must be recalculated. 88 | // etc. 89 | // the macro could also maybe generate an optimal ordering? 90 | let mut changed = true; 91 | while changed { 92 | changed = false; 93 | 94 | let dominators = simple_fast(function.graph(), function.entry().unwrap()); 95 | changed |= structure_jumps(&mut function, &dominators); 96 | 97 | ssa::inline::inline(&mut function, &local_to_group, &upvalue_to_group); 98 | 99 | if structure_conditionals(&mut function) 100 | // || { 101 | // let post_dominators = post_dominators(function.graph_mut()); 102 | // structure_for_loops(&mut function, &dominators, &post_dominators) 103 | // } 104 | || structure_method_calls(&mut function) 105 | { 106 | changed = true; 107 | } 108 | let mut local_map = FxHashMap::default(); 109 | // TODO: loop until returns false? 110 | if ssa::construct::remove_unnecessary_params(&mut function, &mut local_map) { 111 | changed = true; 112 | } 113 | ssa::construct::apply_local_map(&mut function, local_map); 114 | } 115 | ssa::Destructor::new( 116 | &mut function, 117 | upvalue_to_group, 118 | upvalues_in.iter().cloned().collect(), 119 | local_count, 120 | ) 121 | .destruct(); 122 | 123 | let params = std::mem::take(&mut function.parameters); 124 | let is_variadic = function.is_variadic; 125 | let block = Arc::new(restructure::lift(function).into()); 126 | LocalDeclarer::default().declare_locals( 127 | // TODO: why does block.clone() not work? 128 | Arc::clone(&block), 129 | &upvalues_in.iter().chain(params.iter()).cloned().collect(), 130 | ); 131 | 132 | { 133 | let mut ast_function = ast_function.lock(); 134 | ast_function.body = Arc::try_unwrap(block).unwrap().into_inner(); 135 | ast_function.parameters = params; 136 | ast_function.is_variadic = is_variadic; 137 | } 138 | (ByAddress(ast_function), upvalues_in) 139 | }) 140 | .collect::>(); 141 | 142 | let main = ByAddress(main); 143 | upvalues.remove(&main); 144 | let mut body = Arc::try_unwrap(main.0).unwrap().into_inner().body; 145 | link_upvalues(&mut body, &mut upvalues); 146 | name_locals(&mut body, true); 147 | let res = body.to_string(); 148 | let duration = start.elapsed(); 149 | 150 | // TODO: use BufWriter? 151 | let mut out = File::create(path.with_extension("dec.51.lua").file_name().unwrap())?; 152 | writeln!(out, "-- decompiled by Sentinel (took {:?})", duration)?; 153 | writeln!(out, "{}", res)?; 154 | 155 | Ok(()) 156 | } 157 | 158 | fn link_upvalues( 159 | body: &mut ast::Block, 160 | upvalues: &mut FxHashMap>>, Vec>, 161 | ) { 162 | for stat in &mut body.0 { 163 | stat.traverse_rvalues(&mut |rvalue| { 164 | if let ast::RValue::Closure(closure) = rvalue { 165 | let old_upvalues = upvalues.remove(&closure.function).unwrap(); 166 | let mut function = closure.function.lock(); 167 | // TODO: inefficient, try constructing a map of all up -> new up first 168 | // and then call replace_locals on main body 169 | let mut local_map = 170 | FxHashMap::with_capacity_and_hasher(old_upvalues.len(), Default::default()); 171 | for (old, new) in 172 | old_upvalues 173 | .iter() 174 | .zip(closure.upvalues.iter().map(|u| match u { 175 | ast::Upvalue::Copy(l) | ast::Upvalue::Ref(l) => l, 176 | })) 177 | { 178 | // println!("{} -> {}", old, new); 179 | local_map.insert(old.clone(), new.clone()); 180 | } 181 | link_upvalues(&mut function.body, upvalues); 182 | replace_locals(&mut function.body, &local_map); 183 | } 184 | }); 185 | match stat { 186 | ast::Statement::If(r#if) => { 187 | link_upvalues(&mut r#if.then_block.lock(), upvalues); 188 | link_upvalues(&mut r#if.else_block.lock(), upvalues); 189 | } 190 | ast::Statement::While(r#while) => { 191 | link_upvalues(&mut r#while.block.lock(), upvalues); 192 | } 193 | ast::Statement::Repeat(repeat) => { 194 | link_upvalues(&mut repeat.block.lock(), upvalues); 195 | } 196 | ast::Statement::NumericFor(numeric_for) => { 197 | link_upvalues(&mut numeric_for.block.lock(), upvalues); 198 | } 199 | ast::Statement::GenericFor(generic_for) => { 200 | link_upvalues(&mut generic_for.block.lock(), upvalues); 201 | } 202 | _ => {} 203 | } 204 | } 205 | } 206 | -------------------------------------------------------------------------------- /luau-lifter/src/lib.rs: -------------------------------------------------------------------------------- 1 | mod deserializer; 2 | mod instruction; 3 | mod lifter; 4 | mod op_code; 5 | 6 | use ast::{ 7 | local_declarations::LocalDeclarer, name_locals::name_locals, replace_locals::replace_locals, 8 | Traverse, 9 | }; 10 | 11 | use by_address::ByAddress; 12 | use cfg::{ 13 | function::Function, 14 | ssa::{ 15 | self, 16 | structuring::{structure_conditionals, structure_jumps}, 17 | }, 18 | }; 19 | use indexmap::IndexMap; 20 | 21 | use lifter::Lifter; 22 | 23 | //use cfg_ir::{dot, function::Function, ssa}; 24 | use clap::Parser; 25 | use parking_lot::Mutex; 26 | use petgraph::algo::dominators::simple_fast; 27 | use rayon::prelude::*; 28 | 29 | use anyhow::anyhow; 30 | use rustc_hash::FxHashMap; 31 | use triomphe::Arc; 32 | use walkdir::WalkDir; 33 | 34 | use std::{ 35 | fs::File, 36 | io::{Read, Write}, 37 | path::Path, 38 | time::Instant, 39 | }; 40 | 41 | use deserializer::bytecode::Bytecode; 42 | 43 | #[cfg(feature = "dhat-heap")] 44 | #[global_allocator] 45 | static ALLOC: dhat::Alloc = dhat::Alloc; 46 | 47 | #[derive(Parser, Debug)] 48 | #[clap(about, version, author)] 49 | struct Args { 50 | paths: Vec, 51 | /// Number of threads to use (0 = automatic) 52 | #[clap(short, long, default_value_t = 0)] 53 | threads: usize, 54 | /// op = op * key % 256 55 | /// For Roblox client bytecode, use 203 56 | #[clap(short, long, default_value_t = 1)] 57 | key: u8, 58 | #[clap(short, long)] 59 | recursive: bool, 60 | #[clap(short, long)] 61 | verbose: bool, 62 | } 63 | 64 | pub fn decompile_bytecode(bytecode: &[u8], encode_key: u8) -> String { 65 | let chunk = deserializer::deserialize(bytecode, encode_key).unwrap(); 66 | match chunk { 67 | Bytecode::Error(msg) => msg, 68 | Bytecode::Chunk(chunk) => { 69 | let mut lifted = Vec::new(); 70 | let mut stack = vec![(Arc::>::default(), chunk.main)]; 71 | while let Some((ast_func, func_id)) = stack.pop() { 72 | let (function, upvalues, child_functions) = 73 | Lifter::lift(&chunk.functions, &chunk.string_table, func_id); 74 | lifted.push((ast_func, function, upvalues)); 75 | stack.extend(child_functions.into_iter().map(|(a, f)| (a.0, f))); 76 | } 77 | 78 | let (main, ..) = lifted.first().unwrap().clone(); 79 | let mut upvalues = lifted 80 | .into_iter() 81 | .map(|(ast_function, function, upvalues_in)| { 82 | use std::{backtrace::Backtrace, cell::RefCell, fmt::Write, panic}; 83 | 84 | thread_local! { 85 | static BACKTRACE: RefCell> = const { RefCell::new(None) }; 86 | } 87 | 88 | let function_id = function.id; 89 | let mut args = std::panic::AssertUnwindSafe(Some(( 90 | ast_function.clone(), 91 | function, 92 | upvalues_in, 93 | ))); 94 | 95 | let prev_hook = panic::take_hook(); 96 | panic::set_hook(Box::new(|_| { 97 | let trace = Backtrace::capture(); 98 | BACKTRACE.with(move |b| b.borrow_mut().replace(trace)); 99 | })); 100 | let result = panic::catch_unwind(move || { 101 | let (ast_function, function, upvalues_in) = args.take().unwrap(); 102 | decompile_function(ast_function, function, upvalues_in) 103 | }); 104 | panic::set_hook(prev_hook); 105 | 106 | match result { 107 | Ok(r) => r, 108 | Err(e) => { 109 | let panic_information = match e.downcast::() { 110 | Ok(v) => *v, 111 | Err(e) => match e.downcast::<&str>() { 112 | Ok(v) => v.to_string(), 113 | _ => "Unknown Source of Error".to_owned(), 114 | }, 115 | }; 116 | 117 | let mut message = String::new(); 118 | writeln!(message, "failed to decompile").unwrap(); 119 | // writeln!(message, "function {} panicked at '{}'", function_id, panic_information).unwrap(); 120 | // if let Some(backtrace) = BACKTRACE.with(|b| b.borrow_mut().take()) { 121 | // write!(message, "stack backtrace:\n{}", backtrace).unwrap(); 122 | // } 123 | 124 | ast_function.lock().body.extend( 125 | message 126 | .trim_end() 127 | .split('\n') 128 | .map(|s| ast::Comment::new(s.to_string()).into()), 129 | ); 130 | (ByAddress(ast_function), Vec::new()) 131 | } 132 | } 133 | }) 134 | .collect::>(); 135 | 136 | let main = ByAddress(main); 137 | upvalues.remove(&main); 138 | let mut body = Arc::try_unwrap(main.0).unwrap().into_inner().body; 139 | link_upvalues(&mut body, &mut upvalues); 140 | name_locals(&mut body, true); 141 | body.to_string() 142 | } 143 | } 144 | } 145 | 146 | fn decompile_function( 147 | ast_function: Arc>, 148 | mut function: Function, 149 | upvalues_in: Vec, 150 | ) -> (ByAddress>>, Vec) { 151 | let (local_count, local_groups, upvalue_in_groups, upvalue_passed_groups) = 152 | cfg::ssa::construct(&mut function, &upvalues_in); 153 | let upvalue_to_group = upvalue_in_groups 154 | .into_iter() 155 | .chain( 156 | upvalue_passed_groups 157 | .into_iter() 158 | .map(|m| (ast::RcLocal::default(), m)), 159 | ) 160 | .flat_map(|(i, g)| g.into_iter().map(move |u| (u, i.clone()))) 161 | .collect::>(); 162 | // TODO: do we even need this? 163 | let local_to_group = local_groups 164 | .into_iter() 165 | .enumerate() 166 | .flat_map(|(i, g)| g.into_iter().map(move |l| (l, i))) 167 | .collect::>(); 168 | // TODO: REFACTOR: some way to write a macro that states 169 | // if cfg::ssa::inline results in change then structure_jumps, structure_compound_conditionals, 170 | // structure_for_loops and remove_unnecessary_params must run again. 171 | // if structure_compound_conditionals results in change then dominators and post dominators 172 | // must be recalculated. 173 | // etc. 174 | // the macro could also maybe generate an optimal ordering? 175 | let mut changed = true; 176 | while changed { 177 | changed = false; 178 | 179 | let dominators = simple_fast(function.graph(), function.entry().unwrap()); 180 | changed |= structure_jumps(&mut function, &dominators); 181 | 182 | ssa::inline::inline(&mut function, &local_to_group, &upvalue_to_group); 183 | 184 | if structure_conditionals(&mut function) 185 | // || { 186 | // let post_dominators = post_dominators(function.graph_mut()); 187 | // structure_for_loops(&mut function, &dominators, &post_dominators) 188 | // } 189 | // we can't structure method calls like this because of __namecall 190 | // || structure_method_calls(&mut function) 191 | { 192 | changed = true; 193 | } 194 | let mut local_map = FxHashMap::default(); 195 | // TODO: loop until returns false? 196 | if ssa::construct::remove_unnecessary_params(&mut function, &mut local_map) { 197 | changed = true; 198 | } 199 | ssa::construct::apply_local_map(&mut function, local_map); 200 | } 201 | // cfg::dot::render_to(&function, &mut std::io::stdout()).unwrap(); 202 | ssa::Destructor::new( 203 | &mut function, 204 | upvalue_to_group, 205 | upvalues_in.iter().cloned().collect(), 206 | local_count, 207 | ) 208 | .destruct(); 209 | 210 | let params = std::mem::take(&mut function.parameters); 211 | let is_variadic = function.is_variadic; 212 | let block = Arc::new(restructure::lift(function).into()); 213 | LocalDeclarer::default().declare_locals( 214 | // TODO: why does block.clone() not work? 215 | Arc::clone(&block), 216 | &upvalues_in.iter().chain(params.iter()).cloned().collect(), 217 | ); 218 | 219 | { 220 | let mut ast_function = ast_function.lock(); 221 | ast_function.body = Arc::try_unwrap(block).unwrap().into_inner(); 222 | ast_function.parameters = params; 223 | ast_function.is_variadic = is_variadic; 224 | } 225 | (ByAddress(ast_function), upvalues_in) 226 | } 227 | 228 | fn link_upvalues( 229 | body: &mut ast::Block, 230 | upvalues: &mut FxHashMap>>, Vec>, 231 | ) { 232 | for stat in &mut body.0 { 233 | stat.traverse_rvalues(&mut |rvalue| { 234 | if let ast::RValue::Closure(closure) = rvalue { 235 | let old_upvalues = &upvalues[&closure.function]; 236 | let mut function = closure.function.lock(); 237 | // TODO: inefficient, try constructing a map of all up -> new up first 238 | // and then call replace_locals on main body 239 | let mut local_map = 240 | FxHashMap::with_capacity_and_hasher(old_upvalues.len(), Default::default()); 241 | for (old, new) in 242 | old_upvalues 243 | .iter() 244 | .zip(closure.upvalues.iter().map(|u| match u { 245 | ast::Upvalue::Copy(l) | ast::Upvalue::Ref(l) => l, 246 | })) 247 | { 248 | // println!("{} -> {}", old, new); 249 | local_map.insert(old.clone(), new.clone()); 250 | } 251 | link_upvalues(&mut function.body, upvalues); 252 | replace_locals(&mut function.body, &local_map); 253 | } 254 | }); 255 | match stat { 256 | ast::Statement::If(r#if) => { 257 | link_upvalues(&mut r#if.then_block.lock(), upvalues); 258 | link_upvalues(&mut r#if.else_block.lock(), upvalues); 259 | } 260 | ast::Statement::While(r#while) => { 261 | link_upvalues(&mut r#while.block.lock(), upvalues); 262 | } 263 | ast::Statement::Repeat(repeat) => { 264 | link_upvalues(&mut repeat.block.lock(), upvalues); 265 | } 266 | ast::Statement::NumericFor(numeric_for) => { 267 | link_upvalues(&mut numeric_for.block.lock(), upvalues); 268 | } 269 | ast::Statement::GenericFor(generic_for) => { 270 | link_upvalues(&mut generic_for.block.lock(), upvalues); 271 | } 272 | _ => {} 273 | } 274 | } 275 | } 276 | -------------------------------------------------------------------------------- /cfg/src/ssa/param_dependency_graph.rs: -------------------------------------------------------------------------------- 1 | use ast::{LocalRw, RcLocal}; 2 | use indexmap::IndexSet; 3 | use petgraph::{ 4 | prelude::DiGraph, 5 | stable_graph::NodeIndex, 6 | visit::{DfsPostOrder, Walker}, 7 | }; 8 | use rustc_hash::{FxHashMap, FxHashSet}; 9 | 10 | use crate::function::Function; 11 | 12 | // https://github.com/fkie-cad/dewolf/blob/7afe5b46e79a7b56e9904e63f29d54bd8f7302d9/decompiler/pipeline/ssa/phi_dependency_graph.py 13 | #[derive(Debug)] 14 | pub struct ParamDependencyGraph { 15 | // TODO: does this need to be a stable graph? 16 | pub graph: DiGraph, 17 | pub local_to_node: FxHashMap, 18 | } 19 | 20 | impl ParamDependencyGraph { 21 | pub fn new(function: &mut Function, node: NodeIndex) -> Self { 22 | let mut this = Self { 23 | graph: DiGraph::new(), 24 | local_to_node: FxHashMap::default(), 25 | }; 26 | 27 | let edges = function.edges_to_block(node); 28 | for (edge_index, edge) in edges.enumerate() { 29 | if edge_index == 0 { 30 | for (param, _) in &edge.1.arguments { 31 | this.add_node(param.clone()); 32 | } 33 | } 34 | // TODO: support non-local block arguments 35 | for (param, arg) in &edge.1.arguments { 36 | for read in arg.values_read() { 37 | if let Some(&defining_param_node) = this.local_to_node.get(read) { 38 | let param_node = this.local_to_node[param]; 39 | if param_node != defining_param_node { 40 | this.graph.add_edge(param_node, defining_param_node, ()); 41 | } 42 | } 43 | } 44 | } 45 | } 46 | 47 | this 48 | } 49 | 50 | pub fn remove_node(&mut self, node: NodeIndex) -> Option { 51 | if let Some(local) = self.graph.remove_node(node) { 52 | self.local_to_node.remove(&local); 53 | Some(local) 54 | } else { 55 | None 56 | } 57 | } 58 | 59 | pub fn add_node(&mut self, local: RcLocal) -> NodeIndex { 60 | let node = self.graph.add_node(local.clone()); 61 | self.local_to_node.insert(local, node); 62 | node 63 | } 64 | 65 | // This function computes a directed feedback vertex (directed fvs) set of a given graph. 66 | // Since this problem is NP-hard, we only compute an approximate solution. 67 | pub fn compute_directed_fvs(&self) -> IndexSet { 68 | let mut directed_fvs = IndexSet::new(); 69 | let mut dfs_post_order = DfsPostOrder::empty(&self.graph); 70 | dfs_post_order.stack.extend(self.graph.node_indices()); 71 | let mut topological_order = dfs_post_order.iter(&self.graph).collect::>(); 72 | topological_order.reverse(); 73 | 74 | let mut smaller_order = FxHashSet::default(); 75 | for node in topological_order { 76 | if self 77 | .graph 78 | .neighbors(node) 79 | .any(|n| smaller_order.contains(&n)) 80 | { 81 | directed_fvs.insert(node); 82 | } else { 83 | smaller_order.insert(node); 84 | } 85 | } 86 | 87 | directed_fvs 88 | } 89 | } 90 | 91 | // TODO: fix 92 | // #[cfg(test)] 93 | // mod tests { 94 | // use petgraph::dot::Dot; 95 | 96 | // use crate::block::{BasicBlockEdge, Terminator}; 97 | 98 | // use super::*; 99 | 100 | // #[test] 101 | // fn test_dependency_graph() { 102 | // let mut function = Function::default(); 103 | 104 | // let local_y1 = ast::RcLocal::new(ast::Local::new("y1".to_string().into())); 105 | // let local_y2 = ast::RcLocal::new(ast::Local::new("y2".to_string().into())); 106 | // let local_z1 = ast::RcLocal::new(ast::Local::new("z1".to_string().into())); 107 | // let local_z2 = ast::RcLocal::new(ast::Local::new("z2".to_string().into())); 108 | 109 | // let entry_node = function.new_block(); 110 | // let block1_node = function.new_block(); 111 | // let block2_node = function.new_block(); 112 | // function.set_entry(entry_node); 113 | 114 | // let arguments = vec![(local_y2.clone(), local_y1), (local_z2.clone(), local_z1)]; 115 | // function.set_block_terminator( 116 | // entry_node, 117 | // Some(Terminator::Jump(BasicBlockEdge { 118 | // node: block1_node, 119 | // arguments, 120 | // })), 121 | // ); 122 | 123 | // function.set_block_terminator( 124 | // block1_node, 125 | // Some(Terminator::Jump(BasicBlockEdge { 126 | // node: block2_node, 127 | // arguments: Vec::new(), 128 | // })), 129 | // ); 130 | 131 | // let arguments = vec![ 132 | // (local_y2.clone(), local_z2.clone()), 133 | // (local_z2.clone(), local_y2.clone()), 134 | // ]; 135 | // function.set_block_terminator( 136 | // block2_node, 137 | // Some(Terminator::Jump(BasicBlockEdge { 138 | // node: block1_node, 139 | // arguments, 140 | // })), 141 | // ); 142 | 143 | // let dependency_graph = ParamDependencyGraph::new(&mut function, block1_node); 144 | // println!("{:?}", Dot::new(&dependency_graph.graph)); 145 | 146 | // assert!( 147 | // dependency_graph 148 | // .graph 149 | // .neighbors(dependency_graph.local_to_node[&local_z2]) 150 | // .next() 151 | // .unwrap() 152 | // == dependency_graph.local_to_node[&local_y2] 153 | // ); 154 | // assert!( 155 | // dependency_graph 156 | // .graph 157 | // .neighbors(dependency_graph.local_to_node[&local_y2]) 158 | // .next() 159 | // .unwrap() 160 | // == dependency_graph.local_to_node[&local_z2] 161 | // ); 162 | // } 163 | 164 | // #[test] 165 | // fn test_directed_fvs() { 166 | // let mut function = Function::default(); 167 | 168 | // let local_y1 = ast::RcLocal::new(ast::Local::new("y1".to_string().into())); 169 | // let local_y2 = ast::RcLocal::new(ast::Local::new("y2".to_string().into())); 170 | // let local_z1 = ast::RcLocal::new(ast::Local::new("z1".to_string().into())); 171 | // let local_z2 = ast::RcLocal::new(ast::Local::new("z2".to_string().into())); 172 | 173 | // let entry_node = function.new_block(); 174 | // let block1_node = function.new_block(); 175 | // let block2_node = function.new_block(); 176 | // function.set_entry(entry_node); 177 | 178 | // let arguments = vec![(local_y2.clone(), local_y1), (local_z2.clone(), local_z1)]; 179 | // function.set_block_terminator( 180 | // entry_node, 181 | // Some(Terminator::Jump(BasicBlockEdge { 182 | // node: block1_node, 183 | // arguments, 184 | // })), 185 | // ); 186 | 187 | // function.set_block_terminator( 188 | // block1_node, 189 | // Some(Terminator::Jump(BasicBlockEdge { 190 | // node: block2_node, 191 | // arguments: Vec::new(), 192 | // })), 193 | // ); 194 | 195 | // let arguments = vec![ 196 | // (local_y2.clone(), local_z2.clone()), 197 | // (local_z2.clone(), local_y2.clone()), 198 | // ]; 199 | // function.set_block_terminator( 200 | // block2_node, 201 | // Some(Terminator::Jump(BasicBlockEdge { 202 | // node: block1_node, 203 | // arguments, 204 | // })), 205 | // ); 206 | 207 | // let dependency_graph = ParamDependencyGraph::new(&mut function, block1_node); 208 | // println!("{:?}", Dot::new(&dependency_graph.graph)); 209 | 210 | // let directed_fvs = dependency_graph.compute_directed_fvs(); 211 | // println!("{:?}", directed_fvs); 212 | // assert!( 213 | // directed_fvs.contains(&dependency_graph.local_to_node[&local_y2]) 214 | // || directed_fvs.contains(&dependency_graph.local_to_node[&local_z2]) 215 | // ); 216 | // } 217 | 218 | // #[test] 219 | // fn test_multiple_directed_fvs() { 220 | // let mut function = Function::default(); 221 | 222 | // let local_y1 = ast::RcLocal::new(ast::Local::new("y1".to_string().into())); 223 | // let local_y2 = ast::RcLocal::new(ast::Local::new("y2".to_string().into())); 224 | // let local_z1 = ast::RcLocal::new(ast::Local::new("z1".to_string().into())); 225 | // let local_z2 = ast::RcLocal::new(ast::Local::new("z2".to_string().into())); 226 | // let local_a1 = ast::RcLocal::new(ast::Local::new("a1".to_string().into())); 227 | // let local_a2 = ast::RcLocal::new(ast::Local::new("a2".to_string().into())); 228 | // let local_b1 = ast::RcLocal::new(ast::Local::new("b1".to_string().into())); 229 | // let local_b2 = ast::RcLocal::new(ast::Local::new("b2".to_string().into())); 230 | 231 | // let entry_node = function.new_block(); 232 | // let block1_node = function.new_block(); 233 | // let block2_node = function.new_block(); 234 | // function.set_entry(entry_node); 235 | 236 | // let arguments = vec![ 237 | // (local_y2.clone(), local_y1), 238 | // (local_z2.clone(), local_z1), 239 | // (local_a2.clone(), local_a1), 240 | // (local_b2.clone(), local_b1), 241 | // ]; 242 | 243 | // function.set_block_terminator( 244 | // entry_node, 245 | // Some(Terminator::Jump(BasicBlockEdge { 246 | // node: block1_node, 247 | // arguments, 248 | // })), 249 | // ); 250 | 251 | // function.set_block_terminator( 252 | // block1_node, 253 | // Some(Terminator::Jump(BasicBlockEdge { 254 | // node: block2_node, 255 | // arguments: Vec::new(), 256 | // })), 257 | // ); 258 | 259 | // let arguments = vec![ 260 | // (local_y2.clone(), local_z2.clone()), 261 | // (local_z2.clone(), local_y2.clone()), 262 | // (local_a2.clone(), local_b2.clone()), 263 | // (local_b2.clone(), local_a2.clone()), 264 | // ]; 265 | // function.set_block_terminator( 266 | // block2_node, 267 | // Some(Terminator::Jump(BasicBlockEdge { 268 | // node: block1_node, 269 | // arguments, 270 | // })), 271 | // ); 272 | 273 | // let dependency_graph = ParamDependencyGraph::new(&mut function, block1_node); 274 | // println!("{:?}", Dot::new(&dependency_graph.graph)); 275 | 276 | // let directed_fvs = dependency_graph.compute_directed_fvs(); 277 | // println!("{:?}", directed_fvs); 278 | // assert!( 279 | // directed_fvs.contains(&dependency_graph.local_to_node[&local_y2]) 280 | // || directed_fvs.contains(&dependency_graph.local_to_node[&local_z2]) 281 | // ); 282 | // assert!( 283 | // directed_fvs.contains(&dependency_graph.local_to_node[&local_a2]) 284 | // || directed_fvs.contains(&dependency_graph.local_to_node[&local_b2]) 285 | // ); 286 | // } 287 | // } 288 | -------------------------------------------------------------------------------- /ast/src/lib.rs: -------------------------------------------------------------------------------- 1 | #![feature(box_patterns)] 2 | #![feature(let_chains)] 3 | 4 | use derive_more::From; 5 | use enum_as_inner::EnumAsInner; 6 | use enum_dispatch::enum_dispatch; 7 | use formatter::Formatter; 8 | use itertools::Either; 9 | 10 | use std::{ 11 | fmt, 12 | ops::{Deref, DerefMut}, 13 | }; 14 | 15 | mod assign; 16 | mod binary; 17 | mod r#break; 18 | mod call; 19 | mod close; 20 | mod closure; 21 | mod r#continue; 22 | mod r#for; 23 | pub mod formatter; 24 | mod global; 25 | mod goto; 26 | mod r#if; 27 | mod index; 28 | mod literal; 29 | mod local; 30 | //mod name_gen; 31 | pub mod local_declarations; 32 | pub mod name_locals; 33 | mod repeat; 34 | pub mod replace_locals; 35 | mod r#return; 36 | mod set_list; 37 | mod side_effects; 38 | mod table; 39 | mod traverse; 40 | pub mod type_system; 41 | mod unary; 42 | mod vararg; 43 | mod r#while; 44 | 45 | pub use assign::*; 46 | pub use binary::*; 47 | pub use call::*; 48 | pub use close::*; 49 | pub use closure::*; 50 | pub use global::*; 51 | pub use goto::*; 52 | pub use index::*; 53 | pub use literal::*; 54 | pub use local::*; 55 | pub use r#break::*; 56 | pub use r#continue::*; 57 | pub use r#for::*; 58 | pub use r#if::*; 59 | pub use r#return::*; 60 | pub use r#while::*; 61 | pub use repeat::*; 62 | pub use set_list::*; 63 | pub use side_effects::*; 64 | pub use table::*; 65 | pub use traverse::*; 66 | use type_system::{Type, TypeSystem}; 67 | pub use unary::*; 68 | pub use vararg::*; 69 | 70 | pub trait Reduce { 71 | fn reduce(self) -> RValue; 72 | fn reduce_condition(self) -> RValue; 73 | } 74 | 75 | #[enum_dispatch(LocalRw, SideEffects, Traverse)] 76 | #[derive(Debug, Clone, PartialEq, EnumAsInner)] 77 | pub enum Select { 78 | VarArg(VarArg), 79 | Call(Call), 80 | MethodCall(MethodCall), 81 | } 82 | 83 | impl fmt::Display for Select { 84 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 85 | match self { 86 | Select::VarArg(var_arg) => write!(f, "{}", var_arg), 87 | Select::Call(call) => write!(f, "{}", call), 88 | Select::MethodCall(method_call) => write!(f, "{}", method_call), 89 | } 90 | } 91 | } 92 | 93 | #[enum_dispatch(LocalRw, SideEffects, Traverse)] 94 | #[derive(Debug, Clone, PartialEq, EnumAsInner)] 95 | pub enum RValue { 96 | Local(RcLocal), 97 | Global(Global), 98 | Call(Call), 99 | MethodCall(MethodCall), 100 | VarArg(VarArg), 101 | Table(Table), 102 | Literal(Literal), 103 | Index(Index), 104 | Unary(Unary), 105 | Binary(Binary), 106 | Closure(Closure), 107 | Select(Select), 108 | } 109 | 110 | impl type_system::Infer for RValue { 111 | fn infer<'a: 'b, 'b>(&'a mut self, system: &mut TypeSystem<'b>) -> Type { 112 | match self { 113 | RValue::Local(local) => local.infer(system), 114 | RValue::Global(_) => Type::Any, 115 | RValue::Call(_) => Type::Any, 116 | //RValue::Table(table) => table.infer(system), 117 | RValue::Literal(literal) => literal.infer(system), 118 | RValue::Index(_) => Type::Any, 119 | RValue::Unary(_) => Type::Any, 120 | RValue::Binary(_) => Type::Any, 121 | RValue::Closure(closure) => closure.infer(system), 122 | _ => Type::VarArg, 123 | } 124 | } 125 | } 126 | 127 | impl<'a: 'b, 'b> Reduce for RValue { 128 | fn reduce(self) -> RValue { 129 | match self { 130 | Self::Unary(unary) => unary.reduce(), 131 | Self::Binary(binary) => binary.reduce(), 132 | Self::Literal(literal) => literal.reduce(), 133 | Self::Table(table) => table.reduce(), 134 | Self::Closure(closure) => closure.reduce(), 135 | other => other, 136 | } 137 | } 138 | 139 | fn reduce_condition(self) -> RValue { 140 | match self { 141 | Self::Unary(unary) => unary.reduce_condition(), 142 | Self::Binary(binary) => binary.reduce_condition(), 143 | Self::Literal(literal) => literal.reduce_condition(), 144 | Self::Table(table) => table.reduce_condition(), 145 | Self::Closure(closure) => closure.reduce_condition(), 146 | other => other, 147 | } 148 | } 149 | } 150 | 151 | impl RValue { 152 | pub fn precedence(&self) -> usize { 153 | match self { 154 | Self::Binary(binary) => binary.precedence(), 155 | Self::Unary(unary) => unary.precedence(), 156 | RValue::Literal(Literal::Number(n)) if n.is_finite() && n.is_sign_negative() => { 157 | return 7; 158 | } 159 | _ => 9, 160 | } 161 | } 162 | 163 | pub fn into_lvalue(self) -> Option { 164 | match self { 165 | Self::Local(local) => Some(LValue::Local(local)), 166 | Self::Global(global) => Some(LValue::Global(global)), 167 | Self::Index(index) => Some(LValue::Index(index)), 168 | _ => None, 169 | } 170 | } 171 | } 172 | 173 | impl fmt::Display for RValue { 174 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 175 | match self { 176 | RValue::Local(local) => write!(f, "{}", local), 177 | RValue::Global(global) => write!(f, "{}", global), 178 | RValue::Literal(literal) => write!(f, "{}", literal), 179 | RValue::Call(call) => write!(f, "{}", call), 180 | RValue::MethodCall(method_call) => write!(f, "{}", method_call), 181 | RValue::VarArg(var_arg) => write!(f, "{}", var_arg), 182 | RValue::Table(table) => write!(f, "{}", table), 183 | RValue::Index(index) => write!(f, "{}", index), 184 | RValue::Unary(unary) => write!(f, "{}", unary), 185 | RValue::Binary(binary) => write!(f, "{}", binary), 186 | RValue::Closure(closure) => write!(f, "{}", closure), 187 | RValue::Select(select) => write!(f, "{}", select), 188 | } 189 | } 190 | } 191 | 192 | #[enum_dispatch(SideEffects, Traverse)] 193 | #[derive(Debug, Clone, PartialEq, EnumAsInner)] 194 | pub enum LValue { 195 | Local(RcLocal), 196 | Global(Global), 197 | Index(Index), 198 | } 199 | 200 | impl LocalRw for LValue { 201 | fn values_read(&self) -> Vec<&RcLocal> { 202 | match self { 203 | LValue::Local(_) => Vec::new(), 204 | LValue::Global(global) => global.values_read(), 205 | LValue::Index(index) => index.values_read(), 206 | } 207 | } 208 | 209 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 210 | match self { 211 | LValue::Local(_) => Vec::new(), 212 | LValue::Global(global) => global.values_read_mut(), 213 | LValue::Index(index) => index.values_read_mut(), 214 | } 215 | } 216 | 217 | fn values_written(&self) -> Vec<&RcLocal> { 218 | match self { 219 | LValue::Local(local) => vec![local], 220 | LValue::Global(global) => global.values_written(), 221 | LValue::Index(index) => index.values_written(), 222 | } 223 | } 224 | 225 | fn values_written_mut(&mut self) -> Vec<&mut RcLocal> { 226 | match self { 227 | LValue::Local(local) => vec![local], 228 | LValue::Global(global) => global.values_written_mut(), 229 | LValue::Index(index) => index.values_written_mut(), 230 | } 231 | } 232 | } 233 | 234 | impl fmt::Display for LValue { 235 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 236 | match self { 237 | LValue::Local(local) => write!(f, "{}", local), 238 | LValue::Global(global) => write!(f, "{}", global), 239 | LValue::Index(index) => write!(f, "{}", index), 240 | } 241 | } 242 | } 243 | 244 | #[derive(Debug, PartialEq, Eq, Clone)] 245 | pub struct Comment { 246 | pub text: String, 247 | } 248 | 249 | impl Comment { 250 | pub fn new(text: String) -> Self { 251 | Self { text } 252 | } 253 | } 254 | 255 | impl Traverse for Comment {} 256 | 257 | impl SideEffects for Comment {} 258 | 259 | impl LocalRw for Comment {} 260 | 261 | #[enum_dispatch(LocalRw, SideEffects, Traverse)] 262 | #[derive(Debug, Clone, PartialEq, EnumAsInner)] 263 | pub enum Statement { 264 | Empty(Empty), 265 | Call(Call), 266 | MethodCall(MethodCall), 267 | Assign(Assign), 268 | If(If), 269 | Goto(Goto), 270 | Label(Label), 271 | While(While), 272 | Repeat(Repeat), 273 | NumForInit(NumForInit), 274 | NumForNext(NumForNext), 275 | NumericFor(NumericFor), 276 | GenericForInit(GenericForInit), 277 | GenericForNext(GenericForNext), 278 | GenericFor(GenericFor), 279 | Return(Return), 280 | Continue(Continue), 281 | Break(Break), 282 | Close(Close), 283 | SetList(SetList), 284 | Comment(Comment), 285 | } 286 | 287 | #[derive(Debug, PartialEq, Eq, Clone)] 288 | pub struct Empty {} 289 | 290 | impl SideEffects for Empty {} 291 | 292 | impl LocalRw for Empty {} 293 | 294 | impl Traverse for Empty {} 295 | 296 | impl fmt::Display for Empty { 297 | fn fmt(&self, _: &mut fmt::Formatter) -> fmt::Result { 298 | Ok(()) 299 | } 300 | } 301 | 302 | impl fmt::Display for Comment { 303 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 304 | write!(f, "-- {}", self.text) 305 | } 306 | } 307 | 308 | impl fmt::Display for Statement { 309 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 310 | match self { 311 | // TODO: STYLE: order in same order as `Statement` enum 312 | Statement::Call(call) => write!(f, "{}", call), 313 | Statement::MethodCall(method_call) => write!(f, "{}", method_call), 314 | Statement::Assign(assign) => write!(f, "{}", assign), 315 | // TODO: STYLE: replace all `if_` with `r#if`, etc 316 | Statement::If(if_) => write!(f, "{}", if_), 317 | Statement::Goto(goto) => write!(f, "{}", goto), 318 | Statement::Label(label) => write!(f, "{}", label), 319 | Statement::While(while_) => write!(f, "{}", while_), 320 | Statement::Repeat(repeat) => write!(f, "{}", repeat), 321 | Statement::NumForInit(num_for_init) => write!(f, "{}", num_for_init), 322 | Statement::NumForNext(num_for_next) => write!(f, "{}", num_for_next), 323 | Statement::NumericFor(numeric_for) => write!(f, "{}", numeric_for), 324 | Statement::GenericForInit(generic_for_init) => write!(f, "{}", generic_for_init), 325 | Statement::GenericForNext(generic_for_next) => write!(f, "{}", generic_for_next), 326 | Statement::GenericFor(generic_for) => write!(f, "{}", generic_for), 327 | Statement::Return(return_) => write!(f, "{}", return_), 328 | Statement::Continue(continue_) => write!(f, "{}", continue_), 329 | Statement::Break(break_) => write!(f, "{}", break_), 330 | Statement::Comment(comment) => write!(f, "{}", comment), 331 | Statement::SetList(setlist) => write!(f, "{}", setlist), 332 | Statement::Close(close) => write!(f, "{}", close), 333 | Statement::Empty(empty) => write!(f, "{}", empty), 334 | } 335 | } 336 | } 337 | 338 | #[derive(Debug, PartialEq, Clone, Default, From)] 339 | pub struct Block(pub Vec); 340 | 341 | // rust-analyzer doesnt like derive_more :/ 342 | impl Deref for Block { 343 | type Target = Vec; 344 | 345 | fn deref(&self) -> &Self::Target { 346 | &self.0 347 | } 348 | } 349 | 350 | impl DerefMut for Block { 351 | fn deref_mut(&mut self) -> &mut Self::Target { 352 | &mut self.0 353 | } 354 | } 355 | 356 | impl fmt::Display for Block { 357 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 358 | Formatter::format(self, f, Default::default()) 359 | } 360 | } 361 | -------------------------------------------------------------------------------- /ast/src/binary.rs: -------------------------------------------------------------------------------- 1 | use std::fmt; 2 | 3 | use crate::{Literal, LocalRw, RValue, RcLocal, Reduce, SideEffects, Traverse}; 4 | 5 | use super::{Unary, UnaryOperation}; 6 | 7 | #[derive(Debug, PartialEq, Eq, PartialOrd, Copy, Clone)] 8 | pub enum BinaryOperation { 9 | Add, 10 | Sub, 11 | Mul, 12 | Div, 13 | Mod, 14 | Pow, 15 | Concat, 16 | Equal, 17 | NotEqual, 18 | LessThanOrEqual, 19 | GreaterThanOrEqual, 20 | LessThan, 21 | GreaterThan, 22 | And, 23 | Or, 24 | IDiv, 25 | } 26 | 27 | impl BinaryOperation { 28 | pub fn is_comparator(&self) -> bool { 29 | matches!( 30 | self, 31 | BinaryOperation::Equal 32 | | BinaryOperation::NotEqual 33 | | BinaryOperation::LessThanOrEqual 34 | | BinaryOperation::GreaterThanOrEqual 35 | | BinaryOperation::LessThan 36 | | BinaryOperation::GreaterThan 37 | ) 38 | } 39 | } 40 | 41 | impl fmt::Display for BinaryOperation { 42 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 43 | write!( 44 | f, 45 | "{}", 46 | match self { 47 | BinaryOperation::Add => "+", 48 | BinaryOperation::Sub => "-", 49 | BinaryOperation::Mul => "*", 50 | BinaryOperation::Div => "/", 51 | BinaryOperation::Mod => "%", 52 | BinaryOperation::Pow => "^", 53 | BinaryOperation::Concat => "..", 54 | BinaryOperation::Equal => "==", 55 | BinaryOperation::NotEqual => "~=", 56 | BinaryOperation::LessThanOrEqual => "<=", 57 | BinaryOperation::GreaterThanOrEqual => ">=", 58 | BinaryOperation::LessThan => "<", 59 | BinaryOperation::GreaterThan => ">", 60 | BinaryOperation::And => "and", 61 | BinaryOperation::Or => "or", 62 | BinaryOperation::IDiv => "//", 63 | } 64 | ) 65 | } 66 | } 67 | 68 | #[derive(Debug, Clone, PartialEq)] 69 | pub struct Binary { 70 | pub left: Box, 71 | pub right: Box, 72 | pub operation: BinaryOperation, 73 | } 74 | 75 | impl Traverse for Binary { 76 | fn rvalues_mut(&mut self) -> Vec<&mut RValue> { 77 | vec![&mut self.left, &mut self.right] 78 | } 79 | 80 | fn rvalues(&self) -> Vec<&RValue> { 81 | vec![&self.left, &self.right] 82 | } 83 | } 84 | 85 | impl SideEffects for Binary { 86 | fn has_side_effects(&self) -> bool { 87 | // TODO: do this properly 88 | match self.operation { 89 | BinaryOperation::And | BinaryOperation::Or => { 90 | self.left.has_side_effects() || self.right.has_side_effects() 91 | } 92 | _ => true, 93 | } 94 | } 95 | } 96 | 97 | impl<'a: 'b, 'b> Reduce for Binary { 98 | fn reduce(self) -> RValue { 99 | // TODO: true == true, true == false, etc. 100 | // really anything without side effects should be true if l == r 101 | match (self.left.reduce(), self.right.reduce(), self.operation) { 102 | ( 103 | RValue::Unary(Unary { 104 | operation: UnaryOperation::Not, 105 | value: left, 106 | }), 107 | RValue::Unary(Unary { 108 | operation: UnaryOperation::Not, 109 | value: right, 110 | }), 111 | BinaryOperation::And | BinaryOperation::Or, 112 | ) => Unary { 113 | value: Box::new( 114 | Binary { 115 | left, 116 | right, 117 | operation: if self.operation == BinaryOperation::And { 118 | BinaryOperation::Or 119 | } else { 120 | BinaryOperation::And 121 | }, 122 | } 123 | .into(), 124 | ), 125 | operation: UnaryOperation::Not, 126 | } 127 | .into(), 128 | ( 129 | RValue::Literal(Literal::Boolean(left)), 130 | RValue::Literal(Literal::Boolean(right)), 131 | BinaryOperation::And | BinaryOperation::Or, 132 | ) => Literal::Boolean(if self.operation == BinaryOperation::And { 133 | left && right 134 | } else { 135 | left || right 136 | }) 137 | .into(), 138 | ( 139 | RValue::Literal(Literal::Boolean(left)), 140 | right, 141 | BinaryOperation::And | BinaryOperation::Or, 142 | ) => match self.operation { 143 | BinaryOperation::And if !left => RValue::Literal(Literal::Boolean(false)), 144 | BinaryOperation::And => right.reduce(), 145 | BinaryOperation::Or if left => RValue::Literal(Literal::Boolean(true)), 146 | BinaryOperation::Or => right.reduce(), 147 | _ => unreachable!(), 148 | }, 149 | (left, right, BinaryOperation::And) 150 | if !left.has_side_effects() && !right.has_side_effects() && left == right => 151 | { 152 | left 153 | } 154 | ( 155 | RValue::Binary(Binary { 156 | left: 157 | box value @ RValue::Unary(Unary { 158 | operation: UnaryOperation::Not, 159 | .. 160 | }), 161 | right: box RValue::Literal(Literal::Boolean(true)), 162 | operation: BinaryOperation::And, 163 | }), 164 | RValue::Literal(Literal::Boolean(false)), 165 | BinaryOperation::Or, 166 | ) => value, 167 | (left, right, BinaryOperation::Or) if left == right => left, 168 | // TODO: concat numbers 169 | ( 170 | RValue::Literal(Literal::String(left)), 171 | RValue::Literal(Literal::String(right)), 172 | BinaryOperation::Concat, 173 | ) => RValue::Literal(Literal::String( 174 | left.into_iter().chain(right.into_iter()).collect(), 175 | )), 176 | (left, right, operation) => Self { 177 | left: Box::new(left), 178 | right: Box::new(right), 179 | operation, 180 | } 181 | .into(), 182 | } 183 | } 184 | 185 | fn reduce_condition(self) -> RValue { 186 | let (left, right) = if matches!(self.operation, BinaryOperation::And | BinaryOperation::Or) 187 | { 188 | (self.left.reduce_condition(), self.right.reduce_condition()) 189 | } else { 190 | (self.left.reduce(), self.right.reduce()) 191 | }; 192 | match (left, right, self.operation) { 193 | ( 194 | RValue::Unary(Unary { 195 | operation: UnaryOperation::Not, 196 | value: left, 197 | }), 198 | RValue::Unary(Unary { 199 | operation: UnaryOperation::Not, 200 | value: right, 201 | }), 202 | BinaryOperation::And | BinaryOperation::Or, 203 | ) => Unary { 204 | value: Box::new( 205 | Binary { 206 | left, 207 | right, 208 | operation: if self.operation == BinaryOperation::And { 209 | BinaryOperation::Or 210 | } else { 211 | BinaryOperation::And 212 | }, 213 | } 214 | .into(), 215 | ), 216 | operation: UnaryOperation::Not, 217 | } 218 | .into(), 219 | ( 220 | RValue::Literal(Literal::Boolean(left)), 221 | RValue::Literal(Literal::Boolean(right)), 222 | BinaryOperation::And | BinaryOperation::Or, 223 | ) => Literal::Boolean(if self.operation == BinaryOperation::And { 224 | left && right 225 | } else { 226 | left || right 227 | }) 228 | .into(), 229 | ( 230 | RValue::Literal(Literal::Boolean(left)), 231 | right, 232 | BinaryOperation::And | BinaryOperation::Or, 233 | ) => match self.operation { 234 | BinaryOperation::And if !left => RValue::Literal(Literal::Boolean(false)), 235 | BinaryOperation::And => right.reduce(), 236 | BinaryOperation::Or if left => RValue::Literal(Literal::Boolean(true)), 237 | BinaryOperation::Or => right.reduce(), 238 | _ => unreachable!(), 239 | }, 240 | ( 241 | left, 242 | RValue::Literal(Literal::Boolean(right)), 243 | BinaryOperation::And | BinaryOperation::Or, 244 | ) => match self.operation { 245 | BinaryOperation::And if !right => RValue::Literal(Literal::Boolean(false)), 246 | BinaryOperation::And => left.reduce(), 247 | BinaryOperation::Or if right => RValue::Literal(Literal::Boolean(true)), 248 | BinaryOperation::Or => left.reduce(), 249 | _ => unreachable!(), 250 | }, 251 | // TODO: concat numbers 252 | ( 253 | RValue::Literal(Literal::String(left)), 254 | RValue::Literal(Literal::String(right)), 255 | BinaryOperation::Concat, 256 | ) => RValue::Literal(Literal::String( 257 | left.into_iter().chain(right.into_iter()).collect(), 258 | )), 259 | (left, right, operation) => Self { 260 | left: Box::new(left), 261 | right: Box::new(right), 262 | operation, 263 | } 264 | .into(), 265 | } 266 | } 267 | } 268 | 269 | impl Binary { 270 | pub fn new(left: RValue, right: RValue, operation: BinaryOperation) -> Self { 271 | Self { 272 | left: Box::new(left), 273 | right: Box::new(right), 274 | operation, 275 | } 276 | } 277 | 278 | pub fn precedence(&self) -> usize { 279 | match self.operation { 280 | BinaryOperation::Pow => 8, 281 | BinaryOperation::Mul 282 | | BinaryOperation::Div 283 | | BinaryOperation::Mod 284 | | BinaryOperation::IDiv => 6, 285 | BinaryOperation::Add | BinaryOperation::Sub => 5, 286 | BinaryOperation::Concat => 4, 287 | BinaryOperation::LessThan 288 | | BinaryOperation::GreaterThan 289 | | BinaryOperation::LessThanOrEqual 290 | | BinaryOperation::GreaterThanOrEqual 291 | | BinaryOperation::Equal 292 | | BinaryOperation::NotEqual => 3, 293 | BinaryOperation::And => 2, 294 | BinaryOperation::Or => 1, 295 | } 296 | } 297 | 298 | pub fn right_associative(&self) -> bool { 299 | matches!( 300 | self.operation, 301 | BinaryOperation::Pow | BinaryOperation::Concat 302 | ) 303 | } 304 | 305 | pub fn left_group(&self) -> bool { 306 | self.precedence() > self.left.precedence() 307 | || (self.precedence() == self.left.precedence() && self.right_associative()) 308 | } 309 | 310 | pub fn right_group(&self) -> bool { 311 | self.precedence() > self.right.precedence() 312 | || (self.precedence() == self.right.precedence() && !self.right_associative()) 313 | } 314 | } 315 | 316 | impl LocalRw for Binary { 317 | fn values_read(&self) -> Vec<&RcLocal> { 318 | self.left 319 | .values_read() 320 | .into_iter() 321 | .chain(self.right.values_read().into_iter()) 322 | .collect() 323 | } 324 | 325 | fn values_read_mut(&mut self) -> Vec<&mut RcLocal> { 326 | self.left 327 | .values_read_mut() 328 | .into_iter() 329 | .chain(self.right.values_read_mut().into_iter()) 330 | .collect() 331 | } 332 | } 333 | 334 | impl fmt::Display for Binary { 335 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 336 | let parentheses = |group: bool, rvalue: &RValue| { 337 | if group { 338 | format!("({})", rvalue) 339 | } else { 340 | format!("{}", rvalue) 341 | } 342 | }; 343 | 344 | write!( 345 | f, 346 | "{} {} {}", 347 | parentheses(self.left_group(), self.left.as_ref()), 348 | self.operation, 349 | parentheses(self.right_group(), self.right.as_ref()), 350 | ) 351 | } 352 | } 353 | -------------------------------------------------------------------------------- /restructure/src/conditional.rs: -------------------------------------------------------------------------------- 1 | use ast::Reduce; 2 | use cfg::block::{BlockEdge, BranchType}; 3 | use itertools::Itertools; 4 | use parking_lot::Mutex; 5 | use petgraph::visit::EdgeRef; 6 | use triomphe::Arc; 7 | use tuple::Map; 8 | 9 | use crate::GraphStructurer; 10 | use petgraph::{algo::dominators::Dominators, stable_graph::NodeIndex}; 11 | 12 | impl GraphStructurer { 13 | fn simplify_if(if_stat: &mut ast::If) { 14 | if let Some(unary) = if_stat.condition.as_unary() { 15 | if unary.operation == ast::UnaryOperation::Not { 16 | if_stat.condition = *unary.value.clone(); 17 | std::mem::swap(&mut if_stat.then_block, &mut if_stat.else_block); 18 | } 19 | } 20 | } 21 | 22 | fn expand_if(if_stat: &mut ast::If) -> Option { 23 | let mut then_block = if_stat.then_block.lock(); 24 | let mut else_block = if_stat.else_block.lock(); 25 | let then_return = then_block.last().and_then(|x| x.as_return()); 26 | let else_return = else_block.last().and_then(|x| x.as_return()); 27 | if let Some(then_return) = then_return 28 | && let Some(else_return) = else_return 29 | { 30 | if then_return.values.is_empty() && else_return.values.is_empty() { 31 | then_block.pop(); 32 | else_block.pop(); 33 | None 34 | } else if !then_return.values.is_empty() && else_return.values.is_empty() { 35 | Some(std::mem::take(&mut else_block)) 36 | } else if then_return.values.is_empty() && !else_return.values.is_empty() { 37 | let then_block = std::mem::replace::( 38 | &mut then_block, 39 | std::mem::take(&mut else_block), 40 | ); 41 | // TODO: unnecessary clone (also other cases) 42 | if_stat.condition = 43 | ast::Unary::new(if_stat.condition.clone(), ast::UnaryOperation::Not) 44 | .reduce_condition(); 45 | Some(then_block) 46 | } else { 47 | match then_block.len().cmp(&else_block.len()) { 48 | std::cmp::Ordering::Less => Some(std::mem::take(&mut else_block)), 49 | std::cmp::Ordering::Greater => { 50 | let then_block = std::mem::replace::( 51 | &mut then_block, 52 | std::mem::take(&mut else_block), 53 | ); 54 | if_stat.condition = 55 | ast::Unary::new(if_stat.condition.clone(), ast::UnaryOperation::Not) 56 | .reduce_condition(); 57 | Some(then_block) 58 | } 59 | // TODO: `Some(std::mem::take(&mut if_stat.else_block))`? 60 | std::cmp::Ordering::Equal => None, 61 | } 62 | } 63 | } else { 64 | None 65 | } 66 | } 67 | 68 | // a -> b -> d + a -> c -> d 69 | // results in a -> d 70 | fn match_diamond_conditional( 71 | &mut self, 72 | entry: NodeIndex, 73 | then_node: NodeIndex, 74 | else_node: NodeIndex, 75 | ) -> bool { 76 | let mut then_successors = self.function.successor_blocks(then_node).collect_vec(); 77 | let mut else_successors = self.function.successor_blocks(else_node).collect_vec(); 78 | 79 | if then_successors.len() > 1 || else_successors.len() > 1 { 80 | if self.is_loop_header(entry) { 81 | // TODO: ugly 82 | let t = if let Some(index) = then_successors.iter().position(|n| *n == entry) { 83 | then_successors.swap_remove(index); 84 | true 85 | } else { 86 | false 87 | }; 88 | let e = if let Some(index) = else_successors.iter().position(|n| *n == entry) { 89 | else_successors.swap_remove(index); 90 | true 91 | } else { 92 | false 93 | }; 94 | if (!t && then_successors.len() != 1) || (!e && else_successors.len() != 1) { 95 | return false; 96 | } 97 | 98 | if then_successors != else_successors { 99 | return false; 100 | } 101 | 102 | let mut refine = |n| { 103 | let (then_target, else_target) = self 104 | .function 105 | .conditional_edges(n) 106 | .unwrap() 107 | .map(|e| e.target()); 108 | let block = self.function.block_mut(n).unwrap(); 109 | if let Some(if_stat) = block.last_mut().unwrap().as_if_mut() { 110 | if then_target == entry { 111 | if_stat.then_block = 112 | Arc::new(Mutex::new(vec![ast::Continue {}.into()].into())); 113 | true 114 | } else if else_target == entry { 115 | if_stat.else_block = 116 | Arc::new(Mutex::new(vec![ast::Continue {}.into()].into())); 117 | true 118 | } else { 119 | false 120 | } 121 | } else { 122 | false 123 | } 124 | }; 125 | 126 | let then_changed = if t { refine(then_node) } else { false }; 127 | let else_changed = if e { refine(else_node) } else { false }; 128 | if !then_changed && !else_changed { 129 | return false; 130 | } 131 | if t && e { 132 | assert!(then_changed && else_changed) 133 | } 134 | } else { 135 | return false; 136 | } 137 | } else if then_successors != else_successors { 138 | return false; 139 | } 140 | 141 | if self.function.predecessor_blocks(then_node).count() != 1 142 | || self.function.predecessor_blocks(else_node).count() != 1 143 | { 144 | return false; 145 | } 146 | 147 | let then_block = self.function.remove_block(then_node).unwrap(); 148 | let else_block = self.function.remove_block(else_node).unwrap(); 149 | 150 | let block = self.function.block_mut(entry).unwrap(); 151 | // TODO: STYLE: rename to r#if? 152 | let if_stat = block.last_mut().unwrap().as_if_mut().unwrap(); 153 | if_stat.then_block = Arc::new(then_block.into()); 154 | if_stat.else_block = Arc::new(else_block.into()); 155 | Self::simplify_if(if_stat); 156 | 157 | let after = Self::expand_if(if_stat); 158 | if if_stat.then_block.lock().is_empty() { 159 | // TODO: unnecessary clone 160 | if_stat.condition = 161 | ast::Unary::new(if_stat.condition.clone(), ast::UnaryOperation::Not) 162 | .reduce_condition(); 163 | std::mem::swap(&mut if_stat.then_block, &mut if_stat.else_block); 164 | } 165 | if let Some(after) = after { 166 | block.extend(after.0); 167 | } 168 | 169 | let exit = then_successors.first().cloned(); 170 | if let Some(exit) = exit { 171 | self.function.set_edges( 172 | entry, 173 | vec![(exit, BlockEdge::new(BranchType::Unconditional))], 174 | ); 175 | } else { 176 | self.function.remove_edges(entry); 177 | } 178 | self.match_jump(entry, exit); 179 | 180 | true 181 | } 182 | 183 | // a -> b -> c + a -> c 184 | // results in a -> c 185 | fn match_triangle_conditional( 186 | &mut self, 187 | entry: NodeIndex, 188 | then_node: NodeIndex, 189 | else_node: NodeIndex, 190 | ) -> bool { 191 | let mut _match_triangle_conditional = |then_node, else_node, inverted| { 192 | let then_successors = self.function.successor_blocks(then_node).collect_vec(); 193 | 194 | if then_successors.len() > 1 { 195 | return false; 196 | } 197 | 198 | if self.function.predecessor_blocks(then_node).count() != 1 { 199 | return false; 200 | } 201 | 202 | if !then_successors.is_empty() && then_successors[0] != else_node { 203 | return false; 204 | } 205 | 206 | let then_block = self.function.remove_block(then_node).unwrap(); 207 | 208 | let block = self.function.block_mut(entry).unwrap(); 209 | let if_stat = block.last_mut().unwrap().as_if_mut().unwrap(); 210 | if_stat.then_block = Arc::new(then_block.into()); 211 | 212 | if inverted { 213 | if_stat.condition = 214 | ast::Unary::new(if_stat.condition.clone(), ast::UnaryOperation::Not) 215 | .reduce_condition() 216 | } 217 | 218 | //Self::simplify_if(if_stat); 219 | 220 | self.function.set_edges( 221 | entry, 222 | vec![(else_node, BlockEdge::new(BranchType::Unconditional))], 223 | ); 224 | 225 | self.match_jump(entry, Some(else_node)); 226 | 227 | true 228 | }; 229 | 230 | _match_triangle_conditional(then_node, else_node, false) 231 | || _match_triangle_conditional(else_node, then_node, true) 232 | } 233 | 234 | // a -> b a -> c 235 | pub(crate) fn refine_virtual_edge_jump( 236 | &mut self, 237 | post_dom: &Dominators, 238 | entry: NodeIndex, 239 | node: NodeIndex, 240 | header: NodeIndex, 241 | next: Option, 242 | ) -> bool { 243 | if node == header { 244 | // TODO: only check back edges? 245 | if !self 246 | .function 247 | .predecessor_blocks(header) 248 | .filter(|&n| n != entry) 249 | .any(|n| { 250 | post_dom 251 | .dominators(entry) 252 | .is_some_and(|mut p| p.contains(&n)) 253 | }) 254 | { 255 | return false; 256 | } 257 | let block = &mut self.function.block_mut(entry).unwrap(); 258 | block.push(ast::Continue {}.into()); 259 | } else if Some(node) == next { 260 | let block = &mut self.function.block_mut(entry).unwrap(); 261 | block.push(ast::Break {}.into()); 262 | } 263 | self.function.set_edges(entry, vec![]); 264 | true 265 | } 266 | 267 | pub(crate) fn refine_virtual_edge_conditional( 268 | &mut self, 269 | post_dom: &Dominators, 270 | entry: NodeIndex, 271 | then_node: NodeIndex, 272 | else_node: NodeIndex, 273 | header: NodeIndex, 274 | next: Option, 275 | ) -> bool { 276 | let then_main_cont = self 277 | .function 278 | .predecessor_blocks(header) 279 | .filter(|&n| n != entry) 280 | .any(|n| { 281 | post_dom 282 | .dominators(then_node) 283 | .is_some_and(|mut p| p.contains(&n)) 284 | }); 285 | 286 | let else_main_cont = self 287 | .function 288 | .predecessor_blocks(header) 289 | .filter(|&n| n != entry) 290 | .any(|n| { 291 | post_dom 292 | .dominators(else_node) 293 | .is_some_and(|mut p| p.contains(&n)) 294 | }); 295 | 296 | let mut changed = false; 297 | let header_successors = self.function.successor_blocks(header).collect_vec(); 298 | let block = self.function.block_mut(entry).unwrap(); 299 | if let Some(if_stat) = block.last_mut().unwrap().as_if_mut() { 300 | if then_node == header && !header_successors.contains(&entry) && then_main_cont { 301 | if_stat.then_block = Arc::new(Mutex::new(vec![ast::Continue {}.into()].into())); 302 | changed = true; 303 | } else if Some(then_node) == next { 304 | if_stat.then_block = Arc::new(Mutex::new(vec![ast::Break {}.into()].into())); 305 | changed = true; 306 | } 307 | if else_node == header && !header_successors.contains(&entry) && else_main_cont { 308 | if_stat.else_block = Arc::new(Mutex::new(vec![ast::Continue {}.into()].into())); 309 | changed = true; 310 | } else if Some(else_node) == next { 311 | if_stat.else_block = Arc::new(Mutex::new(vec![ast::Break {}.into()].into())); 312 | changed = true; 313 | } 314 | if !if_stat.then_block.lock().is_empty() && if_stat.else_block.lock().is_empty() { 315 | self.function.set_edges( 316 | entry, 317 | vec![(else_node, BlockEdge::new(BranchType::Unconditional))], 318 | ); 319 | changed = true; 320 | } else if if_stat.then_block.lock().is_empty() && !if_stat.else_block.lock().is_empty() 321 | { 322 | if_stat.condition = 323 | ast::Unary::new(if_stat.condition.clone(), ast::UnaryOperation::Not) 324 | .reduce_condition(); 325 | std::mem::swap(&mut if_stat.then_block, &mut if_stat.else_block); 326 | self.function.set_edges( 327 | entry, 328 | vec![(then_node, BlockEdge::new(BranchType::Unconditional))], 329 | ); 330 | changed = true; 331 | } else if !if_stat.then_block.lock().is_empty() && !if_stat.else_block.lock().is_empty() 332 | { 333 | self.function.remove_edges(entry); 334 | changed = true; 335 | } 336 | } 337 | changed 338 | } 339 | 340 | pub(crate) fn match_conditional( 341 | &mut self, 342 | entry: NodeIndex, 343 | then_node: NodeIndex, 344 | else_node: NodeIndex, 345 | ) -> bool { 346 | let block = self.function.block_mut(entry).unwrap(); 347 | if block.last_mut().unwrap().as_if_mut().is_none() { 348 | // for loops 349 | return false; 350 | } 351 | 352 | self.match_diamond_conditional(entry, then_node, else_node) 353 | || self.match_triangle_conditional(entry, then_node, else_node) 354 | } 355 | } 356 | --------------------------------------------------------------------------------