derive/src/visit.rs

use std::collections::HashSet;

use proc_macro::{self, TokenStream};
use proc_macro2::{Span, TokenStream as TokenStream2};
use quote::quote;
use syn::{
  parse::Parse, parse_macro_input, parse_quote, Attribute, Data, DataEnum, DeriveInput, Field, Fields,
  GenericParam, Generics, Ident, Member, Token, Type, Visibility,
};

pub fn derive_visit_children(input: TokenStream) -> TokenStream {
  let DeriveInput {
    ident,
    data,
    generics,
    attrs,
    ..
  } = parse_macro_input!(input);

  let options: Vec<VisitOptions> = attrs
    .iter()
    .filter_map(|attr| {
      if attr.path.is_ident("visit") {
        let opts: VisitOptions = attr.parse_args().unwrap();
        Some(opts)
      } else {
        None
      }
    })
    .collect();

  let visit_types = if let Some(attr) = attrs.iter().find(|attr| attr.path.is_ident("visit_types")) {
    let types: VisitTypes = attr.parse_args().unwrap();
    let types = types.types;
    Some(quote! { crate::visit_types!(#(#types)|*) })
  } else {
    None
  };

  if options.is_empty() {
    derive(&ident, &data, &generics, None, visit_types)
  } else {
    options
      .into_iter()
      .map(|options| derive(&ident, &data, &generics, Some(options), visit_types.clone()))
      .collect()
  }
}

fn derive(
  ident: &Ident,
  data: &Data,
  generics: &Generics,
  options: Option<VisitOptions>,
  visit_types: Option<TokenStream2>,
) -> TokenStream {
  let mut impl_generics = generics.clone();
  let mut type_defs = quote! {};
  let generics = if let Some(VisitOptions {
    generic: Some(generic), ..
  }) = &options
  {
    let mappings = generics
      .type_params()
      .zip(generic.type_params())
      .map(|(a, b)| quote! { type #a = #b; });
    type_defs = quote! { #(#mappings)* };
    impl_generics.params.clear();
    generic
  } else {
    &generics
  };

  if impl_generics.lifetimes().next().is_none() {
    impl_generics.params.insert(0, parse_quote! { 'i })
  }

  let lifetime = impl_generics.lifetimes().next().unwrap().clone();
  let t = impl_generics.type_params().find(|g| &g.ident.to_string() == "R");
  let v = quote! { __V };
  let t = if let Some(t) = t {
    GenericParam::Type(t.ident.clone().into())
  } else {
    let t: GenericParam = parse_quote! { __T };
    impl_generics
      .params
      .push(parse_quote! { #t: crate::visitor::Visit<#lifetime, __T, #v> });
    t
  };

  impl_generics
    .params
    .push(parse_quote! { #v: ?Sized + crate::visitor::Visitor<#lifetime, #t> });

  for ty in generics.type_params() {
    let name = &ty.ident;
    impl_generics.make_where_clause().predicates.push(parse_quote! {
      #name: Visit<#lifetime, #t, #v>
    })
  }

  let mut seen_types = HashSet::new();
  let mut child_types = Vec::new();
  let mut visit = Vec::new();
  match data {
    Data::Struct(s) => {
      for (
        index,
        Field {
          vis, ty, ident, attrs, ..
        },
      ) in s.fields.iter().enumerate()
      {
        if attrs.iter().any(|attr| attr.path.is_ident("skip_visit")) {
          continue;
        }

        if matches!(ty, Type::Reference(_)) || !matches!(vis, Visibility::Public(..)) {
          continue;
        }

        if visit_types.is_none() && !seen_types.contains(ty) && !skip_type(attrs) {
          seen_types.insert(ty.clone());
          child_types.push(quote! {
            <#ty as Visit<#lifetime, #t, #v>>::CHILD_TYPES.bits()
          });
        }

        let name = ident
          .as_ref()
          .map_or_else(|| Member::Unnamed(index.into()), |ident| Member::Named(ident.clone()));
        visit.push(quote! { self.#name.visit(visitor)?; })
      }
    }
    Data::Enum(DataEnum { variants, .. }) => {
      let variants = variants
        .iter()
        .map(|variant| {
          let name = &variant.ident;
          let mut field_names = Vec::new();
          let mut visit_fields = Vec::new();
          for (index, Field { ty, ident, attrs, .. }) in variant.fields.iter().enumerate() {
            let name = ident.as_ref().map_or_else(
              || Ident::new(&format!("_{}", index), Span::call_site()),
              |ident| ident.clone(),
            );
            field_names.push(name.clone());

            if matches!(ty, Type::Reference(_)) {
              continue;
            }

            if visit_types.is_none() && !seen_types.contains(ty) && !skip_type(attrs) && !skip_type(&variant.attrs)
            {
              seen_types.insert(ty.clone());
              child_types.push(quote! {
                <#ty as Visit<#lifetime, #t, #v>>::CHILD_TYPES.bits()
              });
            }

            visit_fields.push(quote! { #name.visit(visitor)?; })
          }

          match variant.fields {
            Fields::Unnamed(_) => {
              quote! {
                Self::#name(#(#field_names),*) => {
                  #(#visit_fields)*
                }
              }
            }
            Fields::Named(_) => {
              quote! {
                Self::#name { #(#field_names),* } => {
                  #(#visit_fields)*
                }
              }
            }
            Fields::Unit => quote! {},
          }
        })
        .collect::<proc_macro2::TokenStream>();

      visit.push(quote! {
        match self {
          #variants
          _ => {}
        }
      })
    }
    _ => {}
  }

  if visit_types.is_none() && child_types.is_empty() {
    child_types.push(quote! { crate::visitor::VisitTypes::empty().bits() });
  }

  let (_, ty_generics, _) = generics.split_for_impl();
  let (impl_generics, _, where_clause) = impl_generics.split_for_impl();

  let self_visit = if let Some(VisitOptions {
    visit: Some(visit),
    kind: Some(kind),
    ..
  }) = &options
  {
    child_types.push(quote! { crate::visitor::VisitTypes::#kind.bits() });

    quote! {
      fn visit(&mut self, visitor: &mut #v) -> Result<(), #v::Error> {
        if visitor.visit_types().contains(crate::visitor::VisitTypes::#kind) {
          visitor.#visit(self)
        } else {
          self.visit_children(visitor)
        }
      }
    }
  } else {
    quote! {}
  };

  let child_types = visit_types.unwrap_or_else(|| {
    quote! {
      {
        #type_defs
        crate::visitor::VisitTypes::from_bits_retain(#(#child_types)|*)
      }
    }
  });

  let output = quote! {
    impl #impl_generics Visit<#lifetime, #t, #v> for #ident #ty_generics #where_clause {
      const CHILD_TYPES: crate::visitor::VisitTypes = #child_types;

      #self_visit

      fn visit_children(&mut self, visitor: &mut #v) -> Result<(), #v::Error> {
        if !<Self as Visit<#lifetime, #t, #v>>::CHILD_TYPES.intersects(visitor.visit_types()) {
          return Ok(())
        }

        #(#visit)*

        Ok(())
      }
    }
  };

  output.into()
}

fn skip_type(attrs: &Vec<Attribute>) -> bool {
  attrs.iter().any(|attr| attr.path.is_ident("skip_type"))
}

struct VisitOptions {
  visit: Option<Ident>,
  kind: Option<Ident>,
  generic: Option<Generics>,
}

impl Parse for VisitOptions {
  fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
    let (visit, kind, comma) = if input.peek(Ident) {
      let visit: Ident = input.parse()?;
      let _: Token![,] = input.parse()?;
      let kind: Ident = input.parse()?;
      let comma: Result<Token![,], _> = input.parse();
      (Some(visit), Some(kind), comma.is_ok())
    } else {
      (None, None, true)
    };
    let generic: Option<Generics> = if comma { Some(input.parse()?) } else { None };
    Ok(Self { visit, kind, generic })
  }
}

struct VisitTypes {
  types: Vec<Ident>,
}

impl Parse for VisitTypes {
  fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
    let first: Ident = input.parse()?;
    let mut types = vec![first];
    while input.parse::<Token![|]>().is_ok() {
      let id: Ident = input.parse()?;
      types.push(id);
    }
    Ok(Self { types })
  }
}