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
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

extern crate proc_macro;
extern crate syn;

#[macro_use]
extern crate quote;

use {
    proc_macro::TokenStream,
    syn::{Attribute, Data, DeriveInput, Fields, Ident, Type},
};

#[proc_macro_derive(SocketWrapper, attributes(sink, stream))]
pub fn socket_derive(input: TokenStream) -> TokenStream {
    let input: DeriveInput = syn::parse(input).unwrap();

    if !only_has_inner_socket(&input.data) {
        panic!("Expected to derive for struct with inner: Socket");
    }

    let name = input.ident;

    let from_parts = quote! {
        impl From<(zmq::Socket, EventedFile)> for #name {
            fn from(tup: (zmq::Socket, EventedFile)) -> Self {
                #name {
                    inner: tup.into()
                }
            }
        }
    };

    let as_socket = quote! {
        impl ::prelude::AsSocket for #name {
            fn socket(self) -> Socket {
                self.inner
            }
        }
    };

    let stream = if has_attr(&input.attrs, "stream") {
        quote!{
            impl ::prelude::StreamSocket for #name {}
        }
    } else {
        quote!{}
    };

    let sink = if has_attr(&input.attrs, "sink") {
        quote!{
            impl ::prelude::SinkSocket for #name {}
        }
    } else {
        quote!{}
    };

    let full = quote! {
        #from_parts
        #as_socket
        #stream
        #sink
    };

    full.into()
}

fn has_attr(attrs: &[Attribute], name: &str) -> bool {
    attrs.iter().any(|attr| {
        attr.path
            .segments
            .last()
            .map(|seg| seg.into_value())
            .map(|seg| seg.ident == Ident::new(name, seg.ident.span()))
            .unwrap_or(false)
    })
}

fn only_has_inner_socket(input: &Data) -> bool {
    let data_struct = match *input {
        Data::Struct(ref data_struct) => data_struct,
        _ => return false, // TODO: Make this work for enums with sockets in each varient?
    };

    let fields_named = match data_struct.fields {
        Fields::Named(ref fields_named) => fields_named,
        _ => return false, // TODO: Allow other kinds of structs?
    };

    if fields_named.named.len() != 1 {
        return false;
    }

    let field = fields_named.named.first().unwrap().into_value();

    let found = field
        .ident
        .as_ref()
        .map(|id| *id == Ident::new("inner", id.span()))
        .unwrap_or(false);

    if !found {
        return false;
    }

    let type_path = match field.ty {
        Type::Path(ref type_path) => type_path,
        _ => return false,
    };

    type_path
        .path
        .segments
        .last()
        .map(|ps| ps.into_value())
        .map(|ps| ps.ident == Ident::new("Socket", ps.ident.span()))
        .unwrap_or(false)
}