proc-macro-assertions 0.1.5

use std::{
    collections::{BTreeMap, BTreeSet, HashMap},
    iter::{FlatMap, Repeat, Zip},
};

use crate::{
    raw_assert::RawAssert,
    token_store::{AddTokens, TokenLevel, TokenStore},
};

use super::{context::Context, generatable, maybe_borrowed::MaybeBorrowed};
use better_any::{Tid, TidAble};

type Asserts<'a, AssertT> = BTreeSet<MaybeBorrowed<'a, AssertT>>;

type Generics<'a, AssertT> =
    HashMap<Option<MaybeBorrowed<'a, syn::Generics>>, Asserts<'a, AssertT>>;

type Templates<'a, Generatable, AssertT> =
    BTreeMap<MaybeBorrowed<'a, Generatable>, Generics<'a, AssertT>>;

/// A wrapper around a [`Generatable`](crate::Generatable) which uniquely stores templates and asserts for
/// each Type implementing [`Generatable`](crate::Generatable)
#[derive(Tid)]
#[allow(clippy::module_name_repetitions)]
pub struct GenerationWrapper<'a, Generatable, Assert>
where
    Generatable: generatable::Generatable<'a, Assert> + Ord + Eq,
    Assert: 'a + TidAble<'a>,
{
    templates: Templates<'a, Generatable, Assert>,
}

/// Trait to add functionality to flatten values of an iterator over a [`std::iter::BTreeMap`] or [`std::iter::HashMap`]
/// which are themselves iteratable into an iterator of tuples, repeating keys for each value returned by the inner iterator.
trait FlattenValues {
    /// The keys of the tuple
    type Keys;
    /// The values of the tuple (themselves implementing `IntoIterator`)
    type Values;
    /// The iterator returned
    type Iterator: Iterator<Item = (Self::Keys, Self::Values)>;

    /// Flatten values into tuples of key and single values.
    fn flatten_values(self) -> Self::Iterator;
}

impl<'a, T, K, V> FlattenValues for T
where
    K: 'a,
    V: 'a,
    T: IntoIterator<Item = (&'a K, &'a V)>,
    &'a V: IntoIterator,
{
    type Keys = &'a K;

    type Values = <&'a V as IntoIterator>::Item;

    type Iterator = FlatMap<
        <T as IntoIterator>::IntoIter,
        Zip<Repeat<&'a K>, <&'a V as IntoIterator>::IntoIter>,
        fn((&'a K, &'a V)) -> Zip<Repeat<&'a K>, <&'a V as IntoIterator>::IntoIter>,
    >;

    fn flatten_values(self) -> Self::Iterator {
        self.into_iter()
            .flat_map(|(x, y)| std::iter::repeat(x).zip(y))
    }
}

impl<'a, Generatable, A> GenerationWrapper<'a, Generatable, A>
where
    Generatable: generatable::Generatable<'a, A> + Ord + Eq,
    A: Ord + Eq + TidAble<'a>,
{
    pub fn add_assert(&mut self, assert: RawAssert<'a, Generatable, A>) {
        self.templates
            .entry(assert.template)
            .or_default()
            .entry(assert.generics)
            .or_default()
            .insert(assert.assert);
    }

    pub fn do_generation(&self, context: &mut Context<'_>, token_store: &mut TokenStore) {
        if Generatable::EMITS_NON_CONSTANT_CODE {
            context.set_require_nonconstant_code();
        }

        let generatable_data = Generatable::generatable(context)
            .seperate_tokens_into(TokenLevel::Generatable, token_store);

        for (template, asserts) in &self.templates {
            let template_data = template
                .template(context, &generatable_data)
                .seperate_tokens_into(TokenLevel::Template, token_store);

            for (generics, assert) in asserts.iter().flatten_values() {
                if let Some(tokens) =
                    template.assert(context, (&generatable_data, &template_data), &**assert)
                {
                    token_store.add_tokens(TokenLevel::Assert(generics), tokens);
                }
            }
        }
    }
}

impl<'a, Generatable, A> Default for GenerationWrapper<'a, Generatable, A>
where
    Generatable: generatable::Generatable<'a, A> + Ord + Eq,
    A: 'a + TidAble<'a>,
{
    fn default() -> Self {
        Self {
            templates: BTreeMap::new(),
        }
    }
}