trait-gen 2.0.7

// Copyright (c) 2025 Redglyph (@gmail.com). All Rights Reserved.
//
// Code substitution code. Independent of proc_macro.

use std::fmt::{Debug, Formatter};
use proc_macro_error2::abort;
use quote::{quote, ToTokens};
use syn::{Generics, GenericParam, TypePath, Path, PathArguments, Expr, Lit, LitStr, ExprLit, Macro, parse_str, Attribute, PathSegment, Type, MetaList, TypeParam};
use syn::spanned::Spanned;
use syn::token::PathSep;
use syn::visit_mut::VisitMut;
use crate::args::{CondParams, TraitGen};
use crate::lib_macros::{VERBOSE, VERBOSE_TF};
use crate::utils::{path_prefix_len, pathname, replace_str};

//==============================================================================
// Attributes that may be inside the content scanned by trait-gen, and which need
// to be processed (the other attributes, either standard or from 3rd-party crates
// are attached normally on the code generated by trait-gen).

// Embedded trait-gen attributes (trait-gen will check for types / paths that must
// be changed).
// Note: when the feature "type_gen" is disabled, we avoid matching "type_gen" by
//       making both constants equal (those constants are used in a match statement).
const TRAIT_GEN: &str = "trait_gen";
const TYPE_GEN: &str = if cfg!(feature = "no_type_gen") { TRAIT_GEN } else { "type_gen" };

// Attributes for conditional implementation.
// Note: when the feature "type_gen" is disabled, we avoid matching "type_gen_if" by
//       making both constants equal (those constants are used in a match statement).
const TRAIT_GEN_IF: &str = "trait_gen_if";
const TYPE_GEN_IF: &str = if cfg!(feature = "no_type_gen") { TRAIT_GEN_IF } else { "type_gen_if" };

//==============================================================================
// Substitution types

#[derive(Clone, PartialEq)]
/// Substitution item, either a Path (`super::Type`) or a Type (`&mut Type`)
pub enum SubstType {
    Path(Path),
    Type(Type)
}

impl ToTokens for SubstType {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        match self {
            SubstType::Path(path) => path.to_tokens(tokens),
            SubstType::Type(ty) => ty.to_tokens(tokens)
        }
    }
}

impl Debug for SubstType {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            SubstType::Path(p) => write!(f, "Path({})", pathname(p)),
            SubstType::Type(t) => write!(f, "Type({})", pathname(t)),
        }
    }
}

/// Converts a list of `Type` to a list of `SubstType`, making them all `Path` if possible,
/// otherwise all `Type`.
pub(crate) fn to_subst_types(mut types: Vec<Type>) -> (bool, Vec<SubstType>) {
    let mut visitor = TurboFish;
    for ty in types.iter_mut() {
        visitor.visit_type_mut(ty);
    }
    let is_path = types.iter().all(|ty| matches!(ty, Type::Path(_)));
    let subst_types = types.into_iter()
        .map(|ty|
            if is_path {
                if let Type::Path(p) = ty {
                    SubstType::Path(p.path)
                } else {
                    panic!("this should match Type::Path: {:?}", ty)
                }
            } else {
                SubstType::Type(ty)
            })
        .collect::<Vec<_>>();
    (is_path, subst_types)
}

/// This type is only used to implement the VisitMut trait.
struct TurboFish;

/// Adds the turbofish double-colon whenever possible to avoid post-substitution problems.
///
/// The replaced part may be an expression requiring it, or a type that doesn't require the
/// double-colon but accepts it. Handling both cases would be complicated so we always include it.
impl VisitMut for TurboFish {
    fn visit_path_mut(&mut self, node: &mut Path) {
        if VERBOSE_TF {
            println!("TURBOF: path '{}'", pathname(node));
        }
        for segment in &mut node.segments {
            if let PathArguments::AngleBracketed(generic_args) = &mut segment.arguments {
                generic_args.colon2_token = Some(PathSep::default());
            }
        }
    }
}

#[derive(Clone)]
/// Attribute substitution data used to replace the generic argument in `generic_arg` with the
/// types in `types`. The first item in `types` is the one currently being replaced in the
/// `VisitMut` implementation.
pub(crate) struct Subst<'a> {
    /// generic argument to replace
    pub generic_arg: Path,
    /// types that replace the generic argument
    pub types: Vec<SubstType>,
    /// Path substitution items if true, Type items if false
    pub is_path: bool,
    /// Context stack, cannot substitue paths when last is false (can substitute if empty)
    pub can_subst_path: Vec<bool>,
    pub type_helper: Option<&'a Vec<SubstType>>
}

impl<'a> Subst<'a> {
    pub fn from_trait_gen(attribute: TraitGen, generic_arg: Path) -> Self {
        let (is_path, types) = to_subst_types(attribute.types);
        Subst {
            generic_arg,
            types,
            is_path,
            can_subst_path: vec![],
            type_helper: None
        }
    }

    fn can_subst_path(&self) -> bool {
        *self.can_subst_path.last().unwrap_or(&true)
    }

    /// Gives a helpful list of type names that might be used in a substitution list.
    fn get_example_types(&self) -> String {
        // This is called for error messages, which happen only during the first visit_mut pass over
        // the inner attributes: we know that Subst still has all the types in `self.new_types`.
        let mut examples = self.type_helper.unwrap_or(&Vec::new()).iter().map(pathname).take(3).collect::<Vec<_>>();
        while examples.len() < 3 {
            examples.push(format!("Type{}", examples.len() + 1));
        }
        examples.join(", ")
    }
}

impl Debug for Subst<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "PathTypes {{ generic_arg: {}, types: [{}], is_path: {}, enabled: {}, type_helper: {} }}",
               pathname(&self.generic_arg),
               self.types.iter().map(pathname).collect::<Vec<_>>().join(", "),
               self.is_path,
               self.can_subst_path.iter().map(|e| e.to_string()).collect::<Vec<_>>().join(", "),
               self.type_helper.unwrap_or(&Vec::new()).iter().map(pathname).collect::<Vec<_>>().join(", "),
        )
    }
}

//==============================================================================
// Main substitution code

impl VisitMut for Subst<'_> {
    fn visit_attribute_mut(&mut self, node: &mut Attribute) {
        // Takes the last segment in case there's a specific path to the attribute. This means we don't support other attributes
        // with the same name inside the outer attribute, but it shouldn't be a problem as long as they could be renamed in the `use`
        // declaration (in the unlikely event that another attribute shared the same name).
        if let Some(PathSegment { ident, .. }) = node.path().segments.last() {
            let ident = ident.to_string();
            match ident.as_str() {
                #[allow(unreachable_patterns)]
                // conditional pseudo-attribute (TYPE_GEN_IF == TRAIT_GEN_IF when type_gen is disabled)
                TRAIT_GEN_IF | TYPE_GEN_IF => {
                    // checks that the syntax is fine and performs the type substitutions if required
                    match node.parse_args::<CondParams>() {
                        Ok(cond) => {
                            let mut output = proc_macro2::TokenStream::new();
                            if VERBOSE { println!("- {} -> {}", pathname(&self.generic_arg), pathname(self.types.first().unwrap())); }
                            let mut g = cond.generic_arg;
                            self.visit_type_mut(&mut g);
                            if cond.is_negated {
                                output.extend(quote!(!#g in));
                            } else {
                                output.extend(quote!(#g in));
                            }
                            let mut first = true;
                            for mut ty in cond.types {
                                // checks if substitutions must be made in that argument:
                                self.visit_type_mut(&mut ty);
                                if !first {
                                    output.extend(quote!(, ));
                                }
                                output.extend(quote!(#ty));
                                first = false;
                            }
                            let original = pathname(&node);
                            if let syn::Meta::List(MetaList { ref mut tokens, .. }) = node.meta {
                                if VERBOSE { println!("  {original} => {}", pathname(&output)); }
                                *tokens = output;
                                return;
                            } else {
                                // shouldn't happen
                                abort!(node.meta.span(), "invalid {} attribute format", ident;
                                    help = "The expected format is: #[{}({} -> {})]", ident, pathname(&self.generic_arg), self.get_example_types());
                            };
                        }
                        Err(err) => {
                            // gives a personalized hint
                            abort!(err.span(), err;
                                help = "The expected format is: #[{}({} in {})]", ident, pathname(&self.generic_arg), self.get_example_types());
                        }
                    };
                }
                #[allow(unreachable_patterns)]
                // embedded trait-gen attributes
                TRAIT_GEN | TYPE_GEN => {
                    // Perform substitutions in the arguments of the inner attribute if necessary.
                    // #[trait_gen(U -> i32, u32)]     // <== self
                    // #[trait_gen(T -> &U, &mut U)]   // <== change 'U' to 'i32' and 'u32'
                    // impl Neg for T { /* .... */ }
                    match node.parse_args::<TraitGen>() {
                        Ok(mut types) => {
                            let mut output = proc_macro2::TokenStream::new();
                            // It would be nice to give a warning if the format of the attributes were different,
                            // once there is a way to generate custom warnings (an error for that seems too harsh).
                            let g = types.args;
                            output.extend(quote!(#g -> ));
                            let mut first = true;
                            for ty in &mut types.types {
                                // checks if substitutions must be made in that argument:
                                self.visit_type_mut(ty);
                                if !first {
                                    output.extend(quote!(, ));
                                }
                                output.extend(quote!(#ty));
                                first = false;
                            }
                            if let syn::Meta::List(MetaList { ref mut tokens, .. }) = node.meta {
                                *tokens = output;
                                return;
                            } else {
                                // shouldn't happen
                                abort!(node.meta.span(), "invalid {} attribute format", ident;
                                    help = "The expected format is: #[{}({} -> {})]", ident, pathname(&self.generic_arg), self.get_example_types());
                            };
                        }
                        Err(err) => {
                            // gives a personalized hint
                            abort!(err.span(), err;
                                help = "The expected format is: #[{}({} -> {})]", ident, pathname(&self.generic_arg), self.get_example_types());
                        }
                    };

                }
                _ => ()
            }
        }
        syn::visit_mut::visit_attribute_mut(self, node);
    }

    fn visit_expr_mut(&mut self, node: &mut Expr) {
        let mut enabled = self.can_subst_path();
        match node {
            // allows substitutions for the nodes below, and until a new Expr is met:
            Expr::Call(_) => enabled = true,
            Expr::Cast(_) => enabled = true,
            Expr::Struct(_) => enabled = true,

            // 'ExprPath' is the node checking for authorization through ExprPath.path,
            // so the current 'enabled' is preserved: (see also visit_path_mut())
            Expr::Path(_) => { /* don't change */ }

            // all other expressions in general must disable substitution:
            _ => enabled = false,
        };
        self.can_subst_path.push(enabled);
        syn::visit_mut::visit_expr_mut(self, node);
        self.can_subst_path.pop();
    }

    fn visit_expr_lit_mut(&mut self, node: &mut ExprLit) {
        if let Lit::Str(lit) = &node.lit {
            // substitutes "${T}" in expression literals (not used in macros, see visit_macro_mut)
            if let Some(ts_str) = replace_str(
                &lit.to_token_stream().to_string(),
                &format!("${{{}}}", pathname(&self.generic_arg)),
                &pathname(self.types.first().unwrap()))
            {
                let new_lit: LitStr = parse_str(&ts_str).unwrap_or_else(|_| panic!("parsing LitStr failed: {}", ts_str));
                node.lit = Lit::Str(new_lit);
            } else {
                syn::visit_mut::visit_expr_lit_mut(self, node);
            }
        }
    }

    fn visit_generics_mut(&mut self, i: &mut Generics) {
        if let Some(segment) = self.generic_arg.segments.first() {
            let current_ident = &segment.ident;
            for t in i.params.iter() {
                match &t {
                    GenericParam::Type(TypeParam { ident, .. }) => {
                        if ident == current_ident {
                            abort!(ident.span(),
                                "Type '{}' is reserved for the substitution.", current_ident.to_string();
                                help = "Use another identifier for this local generic type."
                            );

                            // replace the 'abort!' above with this once it is stable:
                            //
                            // t.span().unwrap()
                            //     .error(format!("Type '{}' is reserved for the substitution.", self.current_type.to_string()))
                            //     .help("Use another identifier for this local generic type.")
                            //     .emit();
                        }
                    }
                    _ => {}
                }
            }
        }
        syn::visit_mut::visit_generics_mut(self, i);
    }

    fn visit_macro_mut(&mut self, node: &mut Macro) {
        // substitutes "${T}" in macros
        if let Some(ts_str) = replace_str(
            &node.tokens.to_string(),
            &format!("${{{}}}", pathname(&self.generic_arg)),
            &pathname(self.types.first().unwrap()))
        {
            let new_ts: proc_macro2::TokenStream = ts_str.parse().unwrap_or_else(|_| panic!("parsing Macro failed: {}", ts_str));
            node.tokens = new_ts;
        } else {
            syn::visit_mut::visit_macro_mut(self, node);
        }
    }

    fn visit_path_mut(&mut self, path: &mut Path) {
        let path_name = pathname(path);
        let path_length = path.segments.len();
        if let Some(length) = path_prefix_len(&self.generic_arg, path) {
            // If U is both a constant and the generic argument, in an expression (so when
            // self.substitution_enabled() == false), we must distinguish two cases:
            // - U::MAX must be replaced (length < path_length)
            // - U or U.add(1) must stay
            if length < path_length || self.can_subst_path() {
                if VERBOSE { print!("[path] path: {} length = {}", path_name, length); }
                match self.types.first().unwrap() {
                    SubstType::Path(p) => {
                        let mut new_seg = p.segments.clone();
                        for seg in path.segments.iter().skip(length) {
                            new_seg.push(seg.clone());
                        }
                        // check if orphan arguments:
                        //     #[trait_gen(g::T -> mod::Name, ...) { ... g::T<'_> ... }
                        //     path     = g :: T   <'_>    len = 2, subst enabled
                        //     new_path = mod :: Name        len = 2
                        //  => new_seg  = mod :: Name<'_>
                        let nth_new_seg = new_seg.last_mut().unwrap();
                        let nth_seg = path.segments.iter().nth(length - 1).unwrap();
                        if nth_new_seg.arguments.is_empty() && !nth_seg.arguments.is_empty() {
                            nth_new_seg.arguments = nth_seg.arguments.clone();
                        }
                        path.segments = new_seg;
                        if VERBOSE { println!(" -> {}", pathname(path)); }
                    }
                    SubstType::Type(ty) => {
                        if VERBOSE { println!(" -> Path '{}' cannot be substituted by type '{}'", path_name, pathname(ty)); }
                        // note: emit-warning is unstable...
                        abort!(ty.span(), "Path '{}' cannot be substituted by type '{}'", path_name, pathname(ty));
                    }
                }
            } else {
                if VERBOSE { println!("disabled path: {}", path_name); }
                syn::visit_mut::visit_path_mut(self, path);
            }
        } else {
            if VERBOSE { println!("path: {} mismatch", path_name); }
            syn::visit_mut::visit_path_mut(self, path);
        }
    }

    fn visit_type_mut(&mut self, node: &mut Type) {
        if !self.is_path {
            match node {
                Type::Path(TypePath { path, .. }) => {
                    let path_name = pathname(path);
                    let path_length = path.segments.len();
                    if let Some(length) = path_prefix_len(&self.generic_arg, path) {
                        if /*length < path_length ||*/ self.can_subst_path() {
                            assert!(length == path_length, "length={length}, path_length={path_length}");
                            // must implement length < path_length if such a case can be found, but it's been elusive
                            if VERBOSE {
                                print!("[type] type path: {} length = {length}, path length = {path_length} {} -> ",
                                       path_name, if self.can_subst_path() { ", can_subst" } else { "" });
                            }
                            *node = if let SubstType::Type(ty) = self.types.first().unwrap() {
                                if VERBOSE { println!("{}", pathname(ty)); }
                                ty.clone()
                            } else {
                                panic!("found path item instead of type in SubstType")
                            };
                        }
                    } else {
                        syn::visit_mut::visit_type_mut(self, node);
                    }
                }
                _ => syn::visit_mut::visit_type_mut(self, node),
            }
        } else {
            syn::visit_mut::visit_type_mut(self, node);
        }
    }

    fn visit_type_path_mut(&mut self, typepath: &mut TypePath) {
        self.can_subst_path.push(true);
        let TypePath { path, .. } = typepath;
        if VERBOSE { println!("typepath: {}", pathname(path)); }
        syn::visit_mut::visit_type_path_mut(self, typepath);
        self.can_subst_path.pop();
    }
}