Skip to main content
AST transformation with full-moon allows you to programmatically modify Lua code while preserving formatting. This is essential for building refactoring tools, code migration utilities, and automated code modernization systems.

​
Understanding VisitorMut

The VisitorMut trait is similar to Visitor, but instead of just inspecting nodes, it returns modified versions. Each visit_* method receives a node by value and returns a (potentially modified) node. 
impl VisitorMut for MyTransformer {
    fn visit_local_assignment(&mut self, node: ast::LocalAssignment) -> ast::LocalAssignment {
        // Transform the node and return the modified version
        node.with_names(modified_names)
    }
}
VisitorMut consumes the original AST and returns a new one. Make sure you don’t need the original before transforming.

​
Example: Function Call Renamer

Rename function calls throughout a codebase (useful for API migrations). 
use full_moon::{
    ast::{self, Ast},
    parse,
    tokenizer::{Token, TokenType, TokenReference},
    visitors::VisitorMut,
};
use std::collections::HashMap;

struct FunctionRenamer {
    renames: HashMap<String, String>,
}

impl FunctionRenamer {
    fn new() -> Self {
        let mut renames = HashMap::new();
        // Define old -> new function name mappings
        renames.insert("oldPrint".to_string(), "print".to_string());
        renames.insert("deprecatedFunction".to_string(), "newFunction".to_string());
        Self { renames }
    }
}

impl VisitorMut for FunctionRenamer {
    fn visit_function_call(&mut self, call: ast::FunctionCall) -> ast::FunctionCall {
        // Check if the prefix is a simple name we want to rename
        let call = match call.prefix() {
            ast::Prefix::Name(name) => {
                let func_name = name.token().to_string();
                
                if let Some(new_name) = self.renames.get(&func_name) {
                    // Create a new token with the new name
                    let new_token = TokenReference::new(
                        vec![], // leading trivia
                        Token::new(TokenType::Identifier {
                            identifier: new_name.clone().into(),
                        }),
                        vec![], // trailing trivia  
                    );
                    
                    call.with_prefix(ast::Prefix::Name(new_token))
                } else {
                    call
                }
            }
            _ => call,
        };
        
        // Continue visiting the rest of the call
        call.visit_mut(self)
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let code = r#"
        oldPrint("Hello")
        deprecatedFunction(42)
        normalFunction()  -- This won't be changed
        
        local x = oldPrint("test")
    "#;

    println!("Original code:\n{}", code);
    
    let ast = parse(code)?;
    let mut renamer = FunctionRenamer::new();
    let transformed = renamer.visit_ast(ast);
    
    println!("\nTransformed code:\n{}", transformed);
    
    Ok(())
}
Output: 
print("Hello")
newFunction(42)
normalFunction()  -- This won't be changed

local x = print("test")

​
Example: Add Type Annotations (Luau)

Automatically add type annotations to function parameters based on heuristics. 
use full_moon::{
    ast::{self, Ast, punctuated::Punctuated},
    parse,
    visitors::VisitorMut,
};

#[cfg(feature = "luau")]
struct TypeAnnotator;

#[cfg(feature = "luau")]
impl TypeAnnotator {
    fn infer_type_from_default(&self, expr: &ast::Expression) -> Option<String> {
        match expr {
            ast::Expression::Number(_) => Some("number".to_string()),
            ast::Expression::String(_) => Some("string".to_string()),
            ast::Expression::Symbol(token) => {
                match token.token().to_string().as_str() {
                    "true" | "false" => Some("boolean".to_string()),
                    "nil" => Some("nil".to_string()),
                    _ => None,
                }
            }
            ast::Expression::TableConstructor(_) => Some("table".to_string()),
            ast::Expression::Function(_) => Some("function".to_string()),
            _ => None,
        }
    }
}

#[cfg(feature = "luau")]
impl VisitorMut for TypeAnnotator {
    fn visit_function_body(&mut self, body: ast::FunctionBody) -> ast::FunctionBody {
        // Check if parameters already have type specifiers
        if body.type_specifiers().any(|spec| spec.is_some()) {
            return body; // Already annotated
        }

        // Try to infer types from parameter names
        let type_specifiers: Vec<Option<ast::luau::TypeSpecifier>> = body
            .parameters()
            .iter()
            .map(|param| {
                if let ast::Parameter::Name(name) = param {
                    let param_name = name.token().to_string();
                    
                    // Simple heuristics based on name
                    let type_str = if param_name.contains("count") || param_name.contains("index") {
                        Some("number")
                    } else if param_name.contains("name") || param_name.contains("message") {
                        Some("string")
                    } else if param_name.starts_with("is_") || param_name.starts_with("has_") {
                        Some("boolean")
                    } else {
                        None
                    };

                    type_str.map(|t| {
                        // Create a basic type specifier
                        // Note: This is simplified; actual implementation would
                        // need to construct proper TypeSpecifier AST nodes
                        unimplemented!("Type specifier construction")
                    })
                } else {
                    None
                }
            })
            .collect();

        body.with_type_specifiers(type_specifiers)
    }
}

​
Example: Variable Name Refactoring

Refactor variable names from camelCase to snake_case throughout the code. 
use full_moon::{
    ast::{self, Ast, punctuated::Punctuated},
    parse,
    tokenizer::{Token, TokenType, TokenReference},
    visitors::VisitorMut,
};
use std::collections::HashMap;

struct SnakeCaseConverter {
    // Track original -> converted names
    name_map: HashMap<String, String>,
}

impl SnakeCaseConverter {
    fn new() -> Self {
        Self {
            name_map: HashMap::new(),
        }
    }

    fn to_snake_case(&mut self, name: &str) -> String {
        // Check if we've already converted this name
        if let Some(converted) = self.name_map.get(name) {
            return converted.clone();
        }

        // Convert camelCase to snake_case
        let mut result = String::new();
        let mut chars = name.chars().peekable();
        
        while let Some(c) = chars.next() {
            if c.is_uppercase() && !result.is_empty() {
                result.push('_');
                result.push(c.to_lowercase().next().unwrap());
            } else {
                result.push(c);
            }
        }
        
        self.name_map.insert(name.to_string(), result.clone());
        result
    }

    fn convert_token(&mut self, token: &TokenReference) -> TokenReference {
        let old_name = token.token().to_string();
        let new_name = self.to_snake_case(&old_name);
        
        if old_name == new_name {
            return token.clone();
        }
        
        // Preserve trivia (whitespace and comments)
        let leading: Vec<_> = token.leading_trivia().cloned().collect();
        let trailing: Vec<_> = token.trailing_trivia().cloned().collect();
        
        TokenReference::new(
            leading,
            Token::new(TokenType::Identifier {
                identifier: new_name.into(),
            }),
            trailing,
        )
    }
}

impl VisitorMut for SnakeCaseConverter {
    fn visit_local_assignment(&mut self, local: ast::LocalAssignment) -> ast::LocalAssignment {
        // Convert variable names in declarations
        let names: Punctuated<TokenReference> = local
            .names()
            .pairs()
            .map(|pair| {
                let name = self.convert_token(pair.value());
                pair.to_owned().map(|_| name)
            })
            .collect();
        
        local.with_names(names).visit_mut(self)
    }

    fn visit_var(&mut self, var: ast::Var) -> ast::Var {
        // Convert variable names in usage
        match var {
            ast::Var::Name(name) => {
                ast::Var::Name(self.convert_token(&name))
            }
            _ => var.visit_mut(self),
        }
    }

    fn visit_function_declaration(&mut self, func: ast::FunctionDeclaration) -> ast::FunctionDeclaration {
        // Convert function names
        let name = func.name();
        let names: Punctuated<TokenReference> = name
            .names()
            .pairs()
            .map(|pair| {
                let name = self.convert_token(pair.value());
                pair.to_owned().map(|_| name)
            })
            .collect();
        
        let new_name = name.clone().with_names(names);
        func.with_name(new_name).visit_mut(self)
    }

    fn visit_parameter(&mut self, param: ast::Parameter) -> ast::Parameter {
        // Convert parameter names
        match param {
            ast::Parameter::Name(name) => {
                ast::Parameter::Name(self.convert_token(&name))
            }
            ast::Parameter::Ellipsis(token) => ast::Parameter::Ellipsis(token),
        }
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let code = r#"
        local myVariable = 10
        local anotherName = 20
        
        function calculateTotal(firstNumber, secondNumber)
            local totalSum = firstNumber + secondNumber
            return totalSum
        end
        
        print(myVariable)
    "#;

    println!("Original (camelCase):\n{}", code);
    
    let ast = parse(code)?;
    let mut converter = SnakeCaseConverter::new();
    let transformed = converter.visit_ast(ast);
    
    println!("\nTransformed (snake_case):\n{}", transformed);
    println!("\nName mappings:");
    for (old, new) in converter.name_map {
        println!("  {} -> {}", old, new);
    }
    
    Ok(())
}
Output: 
local my_variable = 10
local another_name = 20

function calculate_total(first_number, second_number)
    local total_sum = first_number + second_number
    return total_sum
end

print(my_variable)

Name mappings:
  myVariable -> my_variable
  anotherName -> another_name
  calculateTotal -> calculate_total
  firstNumber -> first_number
  secondNumber -> second_number
  totalSum -> total_sum

​
Example: Dead Code Elimination

Remove unreachable code after return statements. 
use full_moon::{
    ast::{self, Ast},
    parse,
    visitors::VisitorMut,
};

struct DeadCodeEliminator;

impl VisitorMut for DeadCodeEliminator {
    fn visit_block(&mut self, block: ast::Block) -> ast::Block {
        let mut statements = Vec::new();
        let mut found_return = false;
        
        // Keep statements until we hit a return
        for (stmt, semi) in block.stmts_with_semicolon() {
            if found_return {
                // Skip statements after return
                continue;
            }
            
            statements.push((stmt.clone(), semi.cloned()));
            
            // Check if this statement contains a return
            // (simplified - a real implementation would check nested blocks)
            if matches!(stmt, ast::Stmt::FunctionCall(_)) {
                // Could be a return in disguise, be conservative
            }
        }
        
        // Check last statement
        let last_stmt = if !found_return {
            block.last_stmt_with_semicolon().map(|(stmt, semi)| {
                if matches!(stmt, ast::LastStmt::Return(_)) {
                    found_return = true;
                }
                (stmt.clone(), semi.cloned())
            })
        } else {
            None
        };
        
        block
            .with_stmts(statements)
            .with_last_stmt(last_stmt)
            .visit_mut(self)
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let code = r#"
        function example(x)
            if x > 10 then
                return true
                print("This won't execute")  -- Dead code
                local y = 5  -- Dead code
            end
            return false
        end
    "#;

    println!("Before optimization:\n{}", code);
    
    let ast = parse(code)?;
    let mut eliminator = DeadCodeEliminator;
    let optimized = eliminator.visit_ast(ast);
    
    println!("\nAfter optimization:\n{}", optimized);
    
    Ok(())
}

​
Combining Transformations

Chain multiple transformers for complex refactoring: 
use full_moon::{ast::Ast, parse};

fn refactor_codebase(code: &str) -> Result<String, Box<dyn std::error::Error>> {
    let ast = parse(code)?;
    
    // Apply multiple transformations in sequence
    let ast = FunctionRenamer::new().visit_ast(ast);
    let ast = SnakeCaseConverter::new().visit_ast(ast);
    let ast = DeadCodeEliminator.visit_ast(ast);
    let ast = IndentationFormatter::new(4).visit_ast(ast);
    
    Ok(ast.to_string())
}

​
Key Concepts

1

Return modified nodes

Each visit_* method must return a node of the same type, either the original or a modified version.
2

Chain transformations

Call .visit_mut(self) at the end of your visitor method to continue traversing child nodes.
3

Use with_* methods

AST nodes have with_* builder methods to create modified copies (e.g., with_name(), with_body()).
4

Preserve tokens

Use TokenReference::new() to create new tokens while preserving trivia (whitespace/comments).

​
Best Practices

  • Test thoroughly: Parse, transform, and parse again to ensure valid Lua output
  • Handle all cases: Consider edge cases like nested structures, method calls, and table indexing
  • Preserve formatting: Copy trivia from original tokens when creating replacements
  • Be conservative: When unsure if a transformation is safe, skip it and log a warning
  • Validate scope: Track variable scopes to avoid renaming conflicts
  • Update positions: Call ast.update_positions() after transformations if you need accurate positions
Scope safety: When renaming variables or functions, be careful about shadowing and scope. A variable named x in one scope might be different from x in another scope.

​
Advanced: Custom AST Node Construction

Sometimes you need to build AST nodes from scratch: 
use full_moon::{
    ast::{self, punctuated::Punctuated},
    tokenizer::{Token, TokenType, TokenReference},
};

// Create a new local assignment: local newVar = 42
fn create_local_assignment() -> ast::LocalAssignment {
    let name = TokenReference::new(
        vec![], // no leading trivia
        Token::new(TokenType::Identifier {
            identifier: "newVar".into(),
        }),
        vec![], // no trailing trivia
    );
    
    let mut names = Punctuated::new();
    names.push(ast::punctuated::Pair::End(name));
    
    let value = ast::Expression::Number(
        TokenReference::basic_symbol("42")
    );
    
    let mut values = Punctuated::new();
    values.push(ast::punctuated::Pair::End(value));
    
    ast::LocalAssignment::new(names)
        .with_equal_token(Some(TokenReference::basic_symbol(" = ")))
        .with_expressions(values)
}

​
Next Steps

Static Analysis

Analyze code patterns with the Visitor trait

Code Formatter

Build formatters that preserve original style

Build docs developers (and LLMs) love

Get started for freeTalk to us