cubecl-macros 0.12.0-pre.1

use proc_macro2::{Span, TokenStream};
use quote::{format_ident, quote, quote_spanned};
use syn::{Token, Type, spanned::Spanned};

use crate::{
    expression::Expression,
    parse::kernel::expand_kernel_ty,
    paths::{frontend_type, prelude_type},
    scope::Context,
    statement::{DefineKind, Statement},
};

/// Names the value of the variable `var` as `name` when its type implements `CubeDebug`. A value
/// of any other type, for example a comptime enum, keeps no name, so it needs no `CubeDebug`.
pub fn name_debug_var(name: &str, var: &TokenStream) -> TokenStream {
    let debug_var = frontend_type("DebugVar");
    let named = frontend_type("DebugVarNamed");
    let unnamed = frontend_type("DebugVarUnnamed");
    quote! {{
        use #named as _;
        use #unnamed as _;
        (&#debug_var(&#var)).debug_var(scope, #name);
    }}
}

impl Statement {
    pub fn to_tokens(&self, context: &mut Context) -> TokenStream {
        match self {
            Statement::Local { variable, init, .. } => {
                let name = &variable.name;
                let is_mut = variable.is_mut_owned || init.as_deref().is_some_and(is_mut_owned);
                let mutable = variable.is_mut_owned.then(|| quote![mut]);
                let is_const = init.as_ref().is_some_and(|it| it.is_const());
                let init = if is_mut {
                    if let Some(as_const) =
                        init.as_ref().and_then(|it| it.as_const_primitive(context))
                    {
                        let expand = frontend_type("NativeExpand");
                        Some(quote_spanned![as_const.span()=> #expand::from_lit(scope, #as_const)])
                    } else if let Some(as_const) = init.as_ref().and_then(|it| it.as_const(context))
                    {
                        Some(quote_spanned![as_const.span()=> #as_const.clone()])
                    } else {
                        init.as_ref().map(|it| it.to_tokens(context))
                    }
                } else {
                    init.as_ref().map(|init| {
                        init.as_const(context)
                            .unwrap_or_else(|| init.to_tokens(context))
                    })
                };
                let ty = variable.ty.as_ref().map(|ty| {
                    let ty = expand_kernel_ty(ty.clone(), is_const)
                        .unwrap_or_else(|err| Type::Verbatim(err.into_compile_error()));
                    quote_spanned! {
                        ty.span()=> :#ty
                    }
                });

                let init = match (is_mut, init) {
                    (true, Some(init)) => {
                        let into_mut = frontend_type("IntoMut");
                        let init_ty =
                            quote_spanned![init.span()=> #into_mut::into_mut(#init, scope)];
                        Some(quote! {{
                            #init_ty
                        }})
                    }
                    (_, init) => init,
                };

                if let Some(mut init) = init {
                    if is_mut || !is_const {
                        let name_str = name.to_string();
                        let init_var = if context.debug_symbols {
                            let var = quote![__cube_var];
                            let name_var = name_debug_var(&name_str, &var);

                            quote![
                                let #var #ty = #init;
                                #name_var
                                #var
                            ]
                        } else {
                            quote![#init]
                        };
                        init = quote! {{
                            #init_var
                        }};
                    }

                    quote![let #mutable #name #ty = #init;]
                } else {
                    let uninit_local = prelude_type("uninit_local");

                    quote![let mut #name #ty = #uninit_local(scope);]
                }
            }
            Statement::Define { name, kind, init } => {
                let value = init
                    .as_const(context)
                    .unwrap_or_else(|| init.to_tokens(context));
                let define_func = match kind {
                    DefineKind::Type => prelude_type("define_scalar"),
                    DefineKind::Size => prelude_type("define_size"),
                };
                // A definition, of a type or a size, lasts to the end of its block.
                let definition = format_ident!("__{}_definition", name.to_string().to_lowercase());
                let register = match kind {
                    DefineKind::Size => {
                        quote![let #definition = scope.define_size::<#name>(#value);]
                    }
                    DefineKind::Type => {
                        quote![let #definition = scope.define_type::<#name>(#value);]
                    }
                };
                quote! {
                    #define_func!(#name);
                    #register
                }
            }
            Statement::Expression {
                expression,
                terminated,
                ..
            } => {
                let terminator = terminated.then(|| Token![;](Span::call_site()));
                if let Some(as_const) = expression.as_const(context) {
                    quote![#as_const #terminator]
                } else {
                    let expression = expression.to_tokens(context);
                    quote![#expression #terminator]
                }
            }
            Statement::Verbatim { tokens } => tokens.clone(),
        }
    }
}

fn is_mut_owned(init: &Expression) -> bool {
    match init {
        Expression::Variable(var) => var.is_mut_owned,
        Expression::FieldAccess { base, .. } => is_mut_owned(base),
        _ => false,
    }
}