luby/crates/luaify/src/transform.rs

use crate::utils::{LuaType, expr_ident, pat_ident, syn_error, wrap_expr_block};
use quote::format_ident;
use std::mem;
use syn::{spanned::*, visit_mut::*, *};

pub fn transform(expr: &mut Expr) -> Result<()> {
    let mut visitor = Visitor::new();
    visitor.visit_expr_mut(expr);
    visitor.result
}

pub fn transform_chunk(block: &mut Block) -> Result<()> {
    let mut visitor = Visitor::new();
    visitor.visit_block_mut(block);
    visitor.result
}

#[derive(Debug)]
struct Visitor {
    result: Result<()>,
}

impl Visitor {
    fn new() -> Self {
        Self { result: Ok(()) }
    }
}

impl VisitMut for Visitor {
    fn visit_expr_closure_mut(&mut self, clo: &mut ExprClosure) {
        match self.transform_expr_closure(clo) {
            res @ Err(_) => self.result = res,
            _ => visit_expr_closure_mut(self, clo),
        }
    }

    fn visit_item_fn_mut(&mut self, func: &mut ItemFn) {
        match self.transform_function(func) {
            res @ Err(_) => self.result = res,
            _ => visit_item_fn_mut(self, func),
        }
    }

    fn visit_expr_mut(&mut self, expr: &mut Expr) {
        match self.transform_expr(expr) {
            res @ Err(_) => self.result = res,
            _ => visit_expr_mut(self, expr),
        }
    }

    fn visit_expr_unary_mut(&mut self, un: &mut ExprUnary) {
        match self.transform_unary(un) {
            res @ Err(_) => self.result = res,
            _ => visit_expr_unary_mut(self, un),
        }
    }

    fn visit_expr_match_mut(&mut self, mat: &mut ExprMatch) {
        match self.transform_match(mat) {
            res @ Err(_) => self.result = res,
            _ => visit_expr_match_mut(self, mat),
        }
    }
}

impl Visitor {
    fn transform_expr_closure(&mut self, clo: &mut ExprClosure) -> Result<()> {
        //
        // transforms a closure expression with input type annotations by removing the annotations
        // and inserting `as` casts at the start.
        //
        // before:
        //    |a: string, b: number| { ... }
        // after:
        //    |a, b| { a as string; b as number; ... }
        //
        let mut checks: Vec<Stmt> = vec![];
        for input in clo.inputs.iter_mut() {
            match input {
                Pat::Ident(_) => {}
                Pat::Type(typed) => {
                    let ident = pat_ident(&typed.pat)?;
                    let ty = mem::replace(&mut typed.ty, parse_quote!(_));
                    match (&*ty).try_into()? {
                        LuaType::Any => {}
                        _ => checks.push(parse_quote! { #ident as #ty; }),
                    }
                }
                _ => {}
            }
        }

        if !checks.is_empty() {
            let mut body = wrap_expr_block(&clo.body);
            body.stmts.splice(..0, checks);
            clo.body = Box::new(parse_quote! { #body });
        }

        Ok(())
    }

    fn transform_function(&mut self, func: &mut ItemFn) -> Result<()> {
        //
        // transforms a function item with input type annotations by removing the annotations
        // and inserting `as` casts at the start.
        //
        // before:
        //    fn my_func(self: table, a: string) { ... }
        // after:
        //    fn my_func(self: _, a: _) { self as table; a as string; ... }
        //
        let mut checks: Vec<Stmt> = vec![];
        for input in func.sig.inputs.iter_mut() {
            if let Some((ident, ty)) = match input {
                FnArg::Receiver(recv) if recv.colon_token.is_some() => {
                    let ty = mem::replace(&mut recv.ty, parse_quote!(_));
                    recv.colon_token = None;
                    Some((Ident::new("self", recv.self_token.span()), ty))
                }
                FnArg::Typed(typed) => {
                    let ident = pat_ident(&typed.pat)?;
                    let ty = mem::replace(&mut typed.ty, parse_quote!(_));
                    Some((ident, ty))
                }
                _ => None,
            } {
                match (&*ty).try_into()? {
                    LuaType::Any => {}
                    _ => checks.push(parse_quote! { #ident as #ty; }),
                }
            };
        }

        func.block.stmts.splice(..0, checks);
        Ok(())
    }

    fn transform_expr(&mut self, expr: &mut Expr) -> Result<()> {
        self.transform_expr_cast(expr)?;
        Ok(())
    }

    fn transform_expr_cast(&mut self, expr: &mut Expr) -> Result<()> {
        //
        // transforms an `as` cast expression into a block expression containing a runtime
        // lua type check.
        //
        // before:
        //    var as string
        // after:
        //    { if type(var) != "string" { error(...) } }
        //
        if let Expr::Cast(cast) = expr {
            let arg = (*cast.expr).clone();
            let mut prelude: Option<Stmt> = None;
            let ty: LuaType = (&*cast.ty).try_into()?;
            let ty_str = format!("{ty}");
            let (ident, msg) = match expr_ident(&arg) {
                Ok(ident) => (ident.clone(), format!("{ty} expected in '{ident}', got ")),
                Err(_) => {
                    let ident = Ident::new("_", arg.span());
                    prelude = Some(parse_quote! { let #ident = #arg; });
                    (ident, format!("{ty} expected, got "))
                }
            };

            let tmp = format_ident!("__{ident}");
            let span = cast.span();
            *expr = match ty {
                LuaType::Any => parse_quote_spanned!(span => {}),
                LuaType::Nil => parse_quote_spanned!(span => {
                    #prelude
                    assert(#ident == (), concat!(#msg, r#type(#ident)));
                }),
                LuaType::Number => parse_quote_spanned!(span => {
                    #prelude
                    let #tmp = #ident;
                    #ident = tonumber(#ident);
                    assert(#ident != (), concat!(#msg, r#type(#tmp)));
                }),
                LuaType::Integer => parse_quote_spanned!(span => {
                    #prelude
                    let #tmp = #ident;
                    #ident = tonumber(#ident);
                    assert(#ident != () && math::floor(#ident) == #ident, concat!(#msg, r#type(#tmp)));
                }),
                LuaType::String => parse_quote_spanned!(span => {
                    #prelude
                    if r#type(#ident) == "number" {
                        #ident = tostring(#ident);
                    } else {
                        assert(r#type(#ident) == "string", concat!(#msg, r#type(#ident)));
                    }
                }),
                _ => parse_quote_spanned!(span => {
                    #prelude
                    assert(r#type(#ident) == #ty_str, concat!(#msg, r#type(#ident)));
                }),
            }
        }

        Ok(())
    }

    fn transform_unary(&mut self, un: &mut ExprUnary) -> Result<()> {
        //
        // separates a nested negation unary operator with parentheses, because double hyphen
        // `--` indicates a comment in lua.
        //
        // before:
        //    --a
        // after:
        //    -(-a)
        //
        if let UnOp::Neg(_) = un.op
            && let Expr::Unary(ref inner) = *un.expr
            && let UnOp::Neg(_) = inner.op
        {
            un.expr = Box::new(parse_quote!((#inner)));
        }

        Ok(())
    }

    fn transform_match(&mut self, mat: &mut ExprMatch) -> Result<()> {
        // TODO:
        syn_error!(mat, "match-to-if transformation not yet implemented");
    }
}