proc-macro-assertions 0.1.5

use std::{collections::HashMap, hash::Hash};

use proc_macro2::{TokenStream, TokenTree};
use quote::quote;

use crate::{context::Context, maybe_borrowed::MaybeBorrowed};

/// Indicates which Stage (see [`Generatable`](crate::generatable::Generatable)) the tokens are from.
/// For the assert stage it also includes the [`Generics`](syn::Generics).
#[derive(Debug, PartialEq, Eq, Hash)]
pub enum TokenLevel<'a> {
    Generatable,
    Template,
    Assert(&'a Option<MaybeBorrowed<'a, syn::Generics>>),
}

/// A token store to store tokens which were generated by [`Generatable`](crate::generatable::Generatable) types.
/// Can be turned into one continous [`TokenStream`](proc_macro_2::TokenStream) using [`TokenStore::into_tokens`].
#[derive(Debug, Default)]
pub struct TokenStore {
    generatable: TokenStream,
    template: TokenStream,
    assert: TokenStream,
    generics: HashMap<syn::Generics, TokenStream>,
}

/// A trait representing some kind of store where the tokens `I` from [`Generatable`](crate::generatable::Generatable) can
/// be added into.
pub trait AddTokens<I> {
    /// Add some tokens to the store, with some [`TokenLevel`].
    fn add_tokens<T>(&mut self, level: TokenLevel, tokens: T)
    where
        T: IntoIterator<Item = I>;
}

fn hash_map_get_or_default_clone_key<'a, K, V>(map: &'a mut HashMap<K, V>, key: &K) -> &'a mut V
where
    K: Eq + Hash + Clone,
    V: Default,
{
    if !map.contains_key(key) {
        map.insert(key.clone(), V::default());
    }

    map.get_mut(key).unwrap()
}

impl AddTokens<TokenStream> for TokenStore {
    fn add_tokens<T>(&mut self, level: TokenLevel, tokens: T)
    where
        T: IntoIterator<Item = TokenStream>,
    {
        match level {
            TokenLevel::Generatable => self.generatable.extend(tokens),
            TokenLevel::Template => self.template.extend(tokens),
            TokenLevel::Assert(&None) => self.assert.extend(tokens),
            TokenLevel::Assert(Some(t)) => {
                hash_map_get_or_default_clone_key(&mut self.generics, t).extend(tokens);
            }
        }
    }
}
impl AddTokens<TokenTree> for TokenStore {
    fn add_tokens<T>(&mut self, level: TokenLevel, tokens: T)
    where
        T: IntoIterator<Item = TokenTree>,
    {
        match level {
            TokenLevel::Generatable => self.generatable.extend(tokens),
            TokenLevel::Template => self.template.extend(tokens),
            TokenLevel::Assert(&None) => self.assert.extend(tokens),
            TokenLevel::Assert(Some(t)) => {
                hash_map_get_or_default_clone_key(&mut self.generics, t).extend(tokens);
            }
        }
    }
}

impl TokenStore {
    /// Create a new token store
    pub fn new() -> Self {
        Self::default()
    }

    /// Turn the token store into one continous token stream.
    pub fn into_tokens(
        self,
        Context {
            ident_generator, ..
        }: &mut Context,
    ) -> TokenStream {
        let Self {
            generatable,
            template,
            assert,
            generics,
        } = self;

        let generics: TokenStream = generics
            .into_iter()
            .map(|(generics, tokens)| {
                let generics_ident = ident_generator.prefixed("generics");
                let (impl_generics, _type_generics, where_clause) = generics.split_for_impl();

                quote! {
                    fn #generics_ident #impl_generics () #where_clause {
                        #tokens
                    }
                }
            })
            .collect();

        quote! {
            #generatable
            #template
            #assert
            #generics
        }
    }
}