1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

use proc_macro2::{Ident, TokenStream};
use quote::quote;
use syn::{Attribute, LifetimeParam, Meta, TypeParam};

enum Variance {
    Covariant,
    Contravariant,
    Invariant,
}

pub trait HasVarianceAttribute {
    fn attrs(&mut self) -> &mut Vec<Attribute>;
}

impl HasVarianceAttribute for TypeParam {
    fn attrs(&mut self) -> &mut Vec<Attribute> {
        &mut self.attrs
    }
}

impl HasVarianceAttribute for LifetimeParam {
    fn attrs(&mut self) -> &mut Vec<Attribute> {
        &mut self.attrs
    }
}

pub fn apply(
    param: &mut dyn HasVarianceAttribute,
    base: TokenStream,
    type_param: &Ident,
) -> TokenStream {
    let mut variance = Variance::Covariant;

    let attrs = param.attrs();
    *attrs = attrs
        .drain(..)
        .filter(|attr| {
            if let Meta::Path(attr_path) = &attr.meta {
                if attr_path.is_ident("contra") {
                    variance = Variance::Contravariant;
                    return false;
                } else if attr_path.is_ident("invariant") {
                    variance = Variance::Invariant;
                    return false;
                }
            }
            true
        })
        .collect();

    let phantom = quote!(self::#type_param<#base>);
    match variance {
        Variance::Covariant => base,
        Variance::Contravariant => quote!(fn(#phantom)),
        Variance::Invariant => quote!(fn(#phantom) -> #phantom),
    }
}