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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

//! An attribute that would create a projection struct covering all the fields.
//!
//! ## Examples
//!
//! ```rust
//! use pin_project::unsafe_project;
//! use std::marker::Unpin;
//! use std::pin::Pin;
//!
//! #[unsafe_project]
//! struct Foo<T, U> {
//!     #[pin]
//!     future: T,
//!     field: U,
//! }
//!
//! impl<T, U> Foo<T, U> {
//!     fn baz(mut self: Pin<&mut Self>) {
//!         let this = self.project();
//!         let _: Pin<&mut T> = this.future; // Pinned reference to the field
//!         let _: &mut U = this.field; // Normal reference to the field
//!     }
//! }
//!
//! impl<T: Unpin, U> Unpin for Foo<T, U> {} // Conditional Unpin impl
//! ```
//!
//! ## Rust Version
//!
//! The current version of pin-project requires Rust nightly 2018-12-26 or later.
//!

#![crate_type = "proc-macro"]
#![recursion_limit = "256"]
#![doc(html_root_url = "https://docs.rs/pin-project/0.1.1")]

extern crate proc_macro;

use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::{quote, ToTokens};
use syn::{Attribute, Field, Fields, FieldsNamed, Ident, ItemStruct};

/// An attribute that would create a projection struct covering all the fields.
#[proc_macro_attribute]
pub fn unsafe_project(args: TokenStream, input: TokenStream) -> TokenStream {
    if !args.is_empty() {
        return compile_err("`unsafe_project` do not requires arguments");
    }

    let mut item: ItemStruct = match syn::parse(input) {
        Err(_) => return compile_err("`unsafe_project` may only be used on structs"),
        Ok(i) => i,
    };

    let fields = match &mut item.fields {
        Fields::Named(FieldsNamed { named, .. }) if !named.is_empty() => named,
        Fields::Named(_) => return err("zero fields"),
        Fields::Unnamed(_) => return err("unnamed fields"),
        Fields::Unit => return err("with units"),
    };

    let mut proj_fields = Vec::with_capacity(fields.len());
    let mut proj_init = Vec::with_capacity(fields.len());
    let pin = quote!(core::pin::Pin);

    fields.iter_mut().for_each(
        |Field {
             attrs, ident, ty, ..
         }| {
            match find_remove(attrs, "pin") {
                Some(_) => {
                    proj_fields.push(quote!(#ident: #pin<&'__a mut #ty>));
                    proj_init
                        .push(quote!(#ident: unsafe { #pin::new_unchecked(&mut this.#ident) }));
                }
                None => {
                    proj_fields.push(quote!(#ident: &'__a mut #ty));
                    proj_init.push(quote!(#ident: &mut this.#ident));
                }
            }
        },
    );

    let proj_ident = Ident::new(&format!("__{}Projection", item.ident), Span::call_site());
    let proj_generics = {
        let generics = item.generics.params.iter();
        quote!(<'__a, #(#generics),*>)
    };
    let proj_item = quote! {
        struct #proj_ident #proj_generics {
            #(#proj_fields,)*
        }
    };

    let ident = &item.ident;
    let (impl_generics, ty_generics, where_clause) = item.generics.split_for_impl();
    let proj_impl = quote! {
        impl #impl_generics #ident #ty_generics #where_clause {
            fn project<'__a>(self: #pin<&'__a mut Self>) -> #proj_ident #proj_generics {
                let this = unsafe { #pin::get_unchecked_mut(self) };
                #proj_ident { #(#proj_init,)* }
            }
        }
    };

    let mut item = item.into_token_stream();
    item.extend(proj_item);
    item.extend(proj_impl);
    TokenStream::from(item)
}

#[inline(never)]
fn compile_err(msg: &str) -> TokenStream {
    TokenStream::from(quote!(compile_error!(#msg);))
}

#[inline(never)]
fn err(msg: &str) -> TokenStream {
    compile_err(&format!("cannot be implemented for structs with {}", msg))
}

fn find_remove(attrs: &mut Vec<Attribute>, ident: &str) -> Option<Attribute> {
    fn remove<T>(v: &mut Vec<T>, index: usize) -> T {
        match v.len() {
            1 => v.pop().unwrap(),
            2 => v.swap_remove(index),
            _ => v.remove(index),
        }
    }

    attrs
        .iter()
        .position(|attr| attr.path.is_ident(ident) && attr.tts.is_empty())
        .map(|i| remove(attrs, i))
}