<commit_before>//! Details of the `ruma_api` procedural macro.

use std::convert::{TryFrom, TryInto as _};

use proc_macro2::{Span, TokenStream};
use proc_macro_crate::crate_name;
use quote::{quote, ToTokens};
use syn::{
    braced,
    parse::{Parse, ParseStream},
    spanned::Spanned,
    Field, FieldValue, Ident, Token, Type,
};

pub(crate) mod attribute;
pub(crate) mod metadata;
pub(crate) mod request;
pub(crate) mod response;

use self::{metadata::Metadata, request::Request, response::Response};
use crate::util;

/// Removes `serde` attributes from struct fields.
pub fn strip_serde_attrs(field: &Field) -> Field {
    let mut field = field.clone();
    field.attrs.retain(|attr| !attr.path.is_ident("serde"));
    field
}

pub fn import_ruma_api() -> TokenStream {
    if let Ok(name) = crate_name("ruma-api") {
        let import = Ident::new(&name, Span::call_site());
        quote! { ::#import }
    } else {
        quote! { ::ruma_api }
    }
}

/// The result of processing the `ruma_api` macro, ready for output back to source code.
pub struct Api {
    /// The `metadata` section of the macro.
    metadata: Metadata,
    /// The `request` section of the macro.
    request: Request,
    /// The `response` section of the macro.
    response: Response,
    /// The `error` section of the macro.
    error: Type,
}

impl TryFrom<RawApi> for Api {
    type Error = syn::Error;

    fn try_from(raw_api: RawApi) -> syn::Result<Self> {
        let res = Self {
            metadata: raw_api.metadata.try_into()?,
            request: raw_api.request.try_into()?,
            response: raw_api.response.try_into()?,
            error: raw_api
                .error
                .map_or(syn::parse_str::<Type>("ruma_api::error::Void").unwrap(), |err| err.ty),
        };

        let newtype_body_field = res.request.newtype_body_field();
        if res.metadata.method == "GET"
            && (res.request.has_body_fields() || newtype_body_field.is_some())
        {
            let mut combined_error: Option<syn::Error> = None;
            let mut add_error = |field| {
                let error = syn::Error::new_spanned(field, "GET endpoints can't have body fields");
                if let Some(combined_error_ref) = &mut combined_error {
                    combined_error_ref.combine(error);
                } else {
                    combined_error = Some(error);
                }
            };

            for field in res.request.body_fields() {
                add_error(field);
            }

            if let Some(field) = newtype_body_field {
                add_error(field);
            }

            Err(combined_error.unwrap())
        } else {
            Ok(res)
        }
    }
}

impl ToTokens for Api {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        // Guarantee `ruma_api` is available and named something we can refer to.
        let ruma_api_import = import_ruma_api();

        let description = &self.metadata.description;
        let method = &self.metadata.method;
        // We don't (currently) use this literal as a literal in the generated code. Instead we just
        // put it into doc comments, for which the span information is irrelevant. So we can work
        // with only the literal's value from here on.
        let name = &self.metadata.name.value();
        let path = &self.metadata.path;
        let rate_limited = &self.metadata.rate_limited;
        let requires_authentication = &self.metadata.requires_authentication;

        let request_type = &self.request;
        let response_type = &self.response;

        let request_try_from_type = if self.request.contains_lifetimes() {
            quote!(IncomingRequest)
        } else {
            quote!(Request)
        };

        let extract_request_path = if self.request.has_path_fields() {
            quote! {
                let path_segments: ::std::vec::Vec<&::std::primitive::str> =
                    request.uri().path()[1..].split('/').collect();
            }
        } else {
            TokenStream::new()
        };

        let (request_path_string, parse_request_path) =
            util::request_path_string_and_parse(&self.request, &self.metadata);

        let request_query_string = util::build_query_string(&self.request);

        let extract_request_query = util::extract_request_query(&self.request);

        let parse_request_query = if let Some(field) = self.request.query_map_field() {
            let field_name = field.ident.as_ref().expect("expected field to have an identifier");

            quote! {
                #field_name: request_query,
            }
        } else {
            self.request.request_init_query_fields()
        };

        let mut header_kvs = self.request.append_header_kvs();
        if requires_authentication.value {
            header_kvs.push(quote! {
                #ruma_api_import::exports::http::header::AUTHORIZATION,
                #ruma_api_import::exports::http::header::HeaderValue::from_str(
                    &::std::format!(
                        "Bearer {}",
                        access_token.ok_or_else(
                            #ruma_api_import::error::IntoHttpError::needs_authentication
                        )?
                    )
                )?
            });
        }

        let extract_request_headers = if self.request.has_header_fields() {
            quote! {
                let headers = request.headers();
            }
        } else {
            TokenStream::new()
        };

        let extract_request_body = if self.request.has_body_fields()
            || self.request.newtype_body_field().is_some()
        {
            let body_lifetimes = if self.request.has_body_lifetimes() {
                // duplicate the anonymous lifetime as many times as needed
                let lifetimes =
                    std::iter::repeat(quote! { '_ }).take(self.request.body_lifetime_count());
                quote! { < #( #lifetimes, )* >}
            } else {
                TokenStream::new()
            };
            quote! {
                let request_body: <RequestBody #body_lifetimes as #ruma_api_import::Outgoing>::Incoming =
                    #ruma_api_import::try_deserialize!(
                        request,
                        #ruma_api_import::exports::serde_json::from_slice(request.body().as_slice())
                    );
            }
        } else {
            TokenStream::new()
        };

        let parse_request_headers = if self.request.has_header_fields() {
            self.request.parse_headers_from_request()
        } else {
            TokenStream::new()
        };

        let request_body = util::build_request_body(&self.request);

        let parse_request_body = util::parse_request_body(&self.request);

        let extract_response_headers = if self.response.has_header_fields() {
            quote! {
                let mut headers = response.headers().clone();
            }
        } else {
            TokenStream::new()
        };

        let typed_response_body_decl = if self.response.has_body_fields()
            || self.response.newtype_body_field().is_some()
        {
            quote! {
                let response_body: <ResponseBody as #ruma_api_import::Outgoing>::Incoming =
                    #ruma_api_import::try_deserialize!(
                        response,
                        #ruma_api_import::exports::serde_json::from_slice(response.body().as_slice()),
                    );
            }
        } else {
            TokenStream::new()
        };

        let response_init_fields = self.response.init_fields();

        let serialize_response_headers = self.response.apply_header_fields();

        let body = self.response.to_body();

        let request_doc = format!(
            "Data for a request to the `{}` API endpoint.\n\n{}",
            name,
            description.value()
        );
        let response_doc = format!("Data in the response from the `{}` API endpoint.", name);

        let error = &self.error;

        let request_lifetimes = self.request.combine_lifetimes();

        let non_auth_endpoint_impl = if requires_authentication.value {
            TokenStream::new()
        } else {
            quote! {
                impl #request_lifetimes #ruma_api_import::NonAuthEndpoint
                    for Request #request_lifetimes
                {}
            }
        };

        let api = quote! {

            #[doc = #request_doc]
            #request_type

            impl ::std::convert::TryFrom<#ruma_api_import::exports::http::Request<Vec<u8>>> for #request_try_from_type {
                type Error = #ruma_api_import::error::FromHttpRequestError;

                #[allow(unused_variables)]
                fn try_from(
                    request: #ruma_api_import::exports::http::Request<Vec<u8>>
                ) -> ::std::result::Result<Self, Self::Error> {
                    #extract_request_path
                    #extract_request_query
                    #extract_request_headers
                    #extract_request_body

                    Ok(Self {
                        #parse_request_path
                        #parse_request_query
                        #parse_request_headers
                        #parse_request_body
                    })
                }
            }

            #[doc = #response_doc]
            #response_type

            impl ::std::convert::TryFrom<Response> for #ruma_api_import::exports::http::Response<Vec<u8>> {
                type Error = #ruma_api_import::error::IntoHttpError;

                #[allow(unused_variables)]
                fn try_from(response: Response) -> ::std::result::Result<Self, Self::Error> {
                    let response = #ruma_api_import::exports::http::Response::builder()
                        .header(#ruma_api_import::exports::http::header::CONTENT_TYPE, "application/json")
                        #serialize_response_headers
                        .body(#body)
                        // Since we require header names to come from the `http::header` module,
                        // this cannot fail.
                        .unwrap();
                    Ok(response)
                }
            }

            impl ::std::convert::TryFrom<#ruma_api_import::exports::http::Response<Vec<u8>>> for Response {
                type Error = #ruma_api_import::error::FromHttpResponseError<#error>;

                #[allow(unused_variables)]
                fn try_from(
                    response: #ruma_api_import::exports::http::Response<Vec<u8>>,
                ) -> ::std::result::Result<Self, Self::Error> {
                    if response.status().as_u16() < 400 {
                        #extract_response_headers

                        #typed_response_body_decl

                        Ok(Self {
                            #response_init_fields
                        })
                    } else {
                        match <#error as #ruma_api_import::EndpointError>::try_from_response(response) {
                            Ok(err) => Err(#ruma_api_import::error::ServerError::Known(err).into()),
                            Err(response_err) => {
                                Err(#ruma_api_import::error::ServerError::Unknown(response_err).into())
                            }
                        }
                    }
                }
            }

            impl #request_lifetimes #ruma_api_import::Endpoint for Request #request_lifetimes {
                type Response = Response;
                type ResponseError = #error;
                type IncomingRequest = <Self as #ruma_api_import::Outgoing>::Incoming;
                type IncomingResponse = <Response as #ruma_api_import::Outgoing>::Incoming;

                /// Metadata for the `#name` endpoint.
                const METADATA: #ruma_api_import::Metadata = #ruma_api_import::Metadata {
                    description: #description,
                    method: #ruma_api_import::exports::http::Method::#method,
                    name: #name,
                    path: #path,
                    rate_limited: #rate_limited,
                    requires_authentication: #requires_authentication,
                };

                #[allow(unused_mut, unused_variables)]
                fn try_into_http_request(
                    self,
                    base_url: &::std::primitive::str,
                    access_token: ::std::option::Option<&str>,
                ) -> ::std::result::Result<
                    #ruma_api_import::exports::http::Request<Vec<u8>>,
                    #ruma_api_import::error::IntoHttpError,
                > {
                    let metadata = Request::METADATA;

                    let http_request = #ruma_api_import::exports::http::Request::builder()
                        .method(#ruma_api_import::exports::http::Method::#method)
                        .uri(::std::format!(
                            "{}{}{}",
                            base_url.strip_suffix("/").unwrap_or(base_url),
                            #request_path_string,
                            #request_query_string,
                        ))
                        #( .header(#header_kvs) )*
                        .body(#request_body)?;

                    Ok(http_request)
                }
            }

            #non_auth_endpoint_impl
        };

        api.to_tokens(tokens);
    }
}

/// Custom keyword macros for syn.
mod kw {
    use syn::custom_keyword;

    custom_keyword!(metadata);
    custom_keyword!(request);
    custom_keyword!(response);
    custom_keyword!(error);
}

/// The entire `ruma_api!` macro structure directly as it appears in the source code..
pub struct RawApi {
    /// The `metadata` section of the macro.
    pub metadata: RawMetadata,
    /// The `request` section of the macro.
    pub request: RawRequest,
    /// The `response` section of the macro.
    pub response: RawResponse,
    /// The `error` section of the macro.
    pub error: Option<RawErrorType>,
}

impl Parse for RawApi {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        Ok(Self {
            metadata: input.parse()?,
            request: input.parse()?,
            response: input.parse()?,
            error: input.parse().ok(),
        })
    }
}

pub struct RawMetadata {
    pub metadata_kw: kw::metadata,
    pub field_values: Vec<FieldValue>,
}

impl Parse for RawMetadata {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let metadata_kw = input.parse::<kw::metadata>()?;
        input.parse::<Token![:]>()?;
        let field_values;
        braced!(field_values in input);

        Ok(Self {
            metadata_kw,
            field_values: field_values
                .parse_terminated::<FieldValue, Token![,]>(FieldValue::parse)?
                .into_iter()
                .collect(),
        })
    }
}

pub struct RawRequest {
    pub request_kw: kw::request,
    pub fields: Vec<Field>,
}

impl Parse for RawRequest {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let request_kw = input.parse::<kw::request>()?;
        input.parse::<Token![:]>()?;
        let fields;
        braced!(fields in input);

        Ok(Self {
            request_kw,
            fields: fields
                .parse_terminated::<Field, Token![,]>(Field::parse_named)?
                .into_iter()
                .collect(),
        })
    }
}

pub struct RawResponse {
    pub response_kw: kw::response,
    pub fields: Vec<Field>,
}

impl Parse for RawResponse {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let response_kw = input.parse::<kw::response>()?;
        input.parse::<Token![:]>()?;
        let fields;
        braced!(fields in input);

        Ok(Self {
            response_kw,
            fields: fields
                .parse_terminated::<Field, Token![,]>(Field::parse_named)?
                .into_iter()
                .map(|f| {
                    if util::has_lifetime(&f.ty) {
                        Err(syn::Error::new(
                            f.ident.span(),
                            "Lifetimes on Response fields cannot be supported until GAT are stable",
                        ))
                    } else {
                        Ok(f)
                    }
                })
                .collect::<Result<Vec<_>, _>>()?,
        })
    }
}

pub struct RawErrorType {
    pub error_kw: kw::error,
    pub ty: Type,
}

impl Parse for RawErrorType {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let error_kw = input.parse::<kw::error>()?;
        input.parse::<Token![:]>()?;
        let ty = input.parse()?;

        Ok(Self { error_kw, ty })
    }
}
<commit_msg>Allow the main ruma crate to re-export the ruma-api-macro things<commit_after>//! Details of the `ruma_api` procedural macro.

use std::convert::{TryFrom, TryInto as _};

use proc_macro2::{Span, TokenStream};
use proc_macro_crate::crate_name;
use quote::{quote, ToTokens};
use syn::{
    braced,
    parse::{Parse, ParseStream},
    spanned::Spanned,
    Field, FieldValue, Ident, Token, Type,
};

pub(crate) mod attribute;
pub(crate) mod metadata;
pub(crate) mod request;
pub(crate) mod response;

use self::{metadata::Metadata, request::Request, response::Response};
use crate::util;

/// Removes `serde` attributes from struct fields.
pub fn strip_serde_attrs(field: &Field) -> Field {
    let mut field = field.clone();
    field.attrs.retain(|attr| !attr.path.is_ident("serde"));
    field
}

pub fn import_ruma_api() -> TokenStream {
    if let Ok(possibly_renamed) = crate_name("ruma-api") {
        let import = Ident::new(&possibly_renamed, Span::call_site());
        quote! { ::#import }
    } else if let Ok(possibly_renamed) = crate_name("ruma") {
        let import = Ident::new(&possibly_renamed, Span::call_site());
        quote! { ::#import::api }
    } else {
        quote! { ::ruma_api }
    }
}

/// The result of processing the `ruma_api` macro, ready for output back to source code.
pub struct Api {
    /// The `metadata` section of the macro.
    metadata: Metadata,
    /// The `request` section of the macro.
    request: Request,
    /// The `response` section of the macro.
    response: Response,
    /// The `error` section of the macro.
    error: Type,
}

impl TryFrom<RawApi> for Api {
    type Error = syn::Error;

    fn try_from(raw_api: RawApi) -> syn::Result<Self> {
        let res = Self {
            metadata: raw_api.metadata.try_into()?,
            request: raw_api.request.try_into()?,
            response: raw_api.response.try_into()?,
            error: raw_api
                .error
                .map_or(syn::parse_str::<Type>("ruma_api::error::Void").unwrap(), |err| err.ty),
        };

        let newtype_body_field = res.request.newtype_body_field();
        if res.metadata.method == "GET"
            && (res.request.has_body_fields() || newtype_body_field.is_some())
        {
            let mut combined_error: Option<syn::Error> = None;
            let mut add_error = |field| {
                let error = syn::Error::new_spanned(field, "GET endpoints can't have body fields");
                if let Some(combined_error_ref) = &mut combined_error {
                    combined_error_ref.combine(error);
                } else {
                    combined_error = Some(error);
                }
            };

            for field in res.request.body_fields() {
                add_error(field);
            }

            if let Some(field) = newtype_body_field {
                add_error(field);
            }

            Err(combined_error.unwrap())
        } else {
            Ok(res)
        }
    }
}

impl ToTokens for Api {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        // Guarantee `ruma_api` is available and named something we can refer to.
        let ruma_api_import = import_ruma_api();

        let description = &self.metadata.description;
        let method = &self.metadata.method;
        // We don't (currently) use this literal as a literal in the generated code. Instead we just
        // put it into doc comments, for which the span information is irrelevant. So we can work
        // with only the literal's value from here on.
        let name = &self.metadata.name.value();
        let path = &self.metadata.path;
        let rate_limited = &self.metadata.rate_limited;
        let requires_authentication = &self.metadata.requires_authentication;

        let request_type = &self.request;
        let response_type = &self.response;

        let request_try_from_type = if self.request.contains_lifetimes() {
            quote!(IncomingRequest)
        } else {
            quote!(Request)
        };

        let extract_request_path = if self.request.has_path_fields() {
            quote! {
                let path_segments: ::std::vec::Vec<&::std::primitive::str> =
                    request.uri().path()[1..].split('/').collect();
            }
        } else {
            TokenStream::new()
        };

        let (request_path_string, parse_request_path) =
            util::request_path_string_and_parse(&self.request, &self.metadata);

        let request_query_string = util::build_query_string(&self.request);

        let extract_request_query = util::extract_request_query(&self.request);

        let parse_request_query = if let Some(field) = self.request.query_map_field() {
            let field_name = field.ident.as_ref().expect("expected field to have an identifier");

            quote! {
                #field_name: request_query,
            }
        } else {
            self.request.request_init_query_fields()
        };

        let mut header_kvs = self.request.append_header_kvs();
        if requires_authentication.value {
            header_kvs.push(quote! {
                #ruma_api_import::exports::http::header::AUTHORIZATION,
                #ruma_api_import::exports::http::header::HeaderValue::from_str(
                    &::std::format!(
                        "Bearer {}",
                        access_token.ok_or_else(
                            #ruma_api_import::error::IntoHttpError::needs_authentication
                        )?
                    )
                )?
            });
        }

        let extract_request_headers = if self.request.has_header_fields() {
            quote! {
                let headers = request.headers();
            }
        } else {
            TokenStream::new()
        };

        let extract_request_body = if self.request.has_body_fields()
            || self.request.newtype_body_field().is_some()
        {
            let body_lifetimes = if self.request.has_body_lifetimes() {
                // duplicate the anonymous lifetime as many times as needed
                let lifetimes =
                    std::iter::repeat(quote! { '_ }).take(self.request.body_lifetime_count());
                quote! { < #( #lifetimes, )* >}
            } else {
                TokenStream::new()
            };
            quote! {
                let request_body: <RequestBody #body_lifetimes as #ruma_api_import::Outgoing>::Incoming =
                    #ruma_api_import::try_deserialize!(
                        request,
                        #ruma_api_import::exports::serde_json::from_slice(request.body().as_slice())
                    );
            }
        } else {
            TokenStream::new()
        };

        let parse_request_headers = if self.request.has_header_fields() {
            self.request.parse_headers_from_request()
        } else {
            TokenStream::new()
        };

        let request_body = util::build_request_body(&self.request);

        let parse_request_body = util::parse_request_body(&self.request);

        let extract_response_headers = if self.response.has_header_fields() {
            quote! {
                let mut headers = response.headers().clone();
            }
        } else {
            TokenStream::new()
        };

        let typed_response_body_decl = if self.response.has_body_fields()
            || self.response.newtype_body_field().is_some()
        {
            quote! {
                let response_body: <ResponseBody as #ruma_api_import::Outgoing>::Incoming =
                    #ruma_api_import::try_deserialize!(
                        response,
                        #ruma_api_import::exports::serde_json::from_slice(response.body().as_slice()),
                    );
            }
        } else {
            TokenStream::new()
        };

        let response_init_fields = self.response.init_fields();

        let serialize_response_headers = self.response.apply_header_fields();

        let body = self.response.to_body();

        let request_doc = format!(
            "Data for a request to the `{}` API endpoint.\n\n{}",
            name,
            description.value()
        );
        let response_doc = format!("Data in the response from the `{}` API endpoint.", name);

        let error = &self.error;

        let request_lifetimes = self.request.combine_lifetimes();

        let non_auth_endpoint_impl = if requires_authentication.value {
            TokenStream::new()
        } else {
            quote! {
                impl #request_lifetimes #ruma_api_import::NonAuthEndpoint
                    for Request #request_lifetimes
                {}
            }
        };

        let api = quote! {

            #[doc = #request_doc]
            #request_type

            impl ::std::convert::TryFrom<#ruma_api_import::exports::http::Request<Vec<u8>>> for #request_try_from_type {
                type Error = #ruma_api_import::error::FromHttpRequestError;

                #[allow(unused_variables)]
                fn try_from(
                    request: #ruma_api_import::exports::http::Request<Vec<u8>>
                ) -> ::std::result::Result<Self, Self::Error> {
                    #extract_request_path
                    #extract_request_query
                    #extract_request_headers
                    #extract_request_body

                    Ok(Self {
                        #parse_request_path
                        #parse_request_query
                        #parse_request_headers
                        #parse_request_body
                    })
                }
            }

            #[doc = #response_doc]
            #response_type

            impl ::std::convert::TryFrom<Response> for #ruma_api_import::exports::http::Response<Vec<u8>> {
                type Error = #ruma_api_import::error::IntoHttpError;

                #[allow(unused_variables)]
                fn try_from(response: Response) -> ::std::result::Result<Self, Self::Error> {
                    let response = #ruma_api_import::exports::http::Response::builder()
                        .header(#ruma_api_import::exports::http::header::CONTENT_TYPE, "application/json")
                        #serialize_response_headers
                        .body(#body)
                        // Since we require header names to come from the `http::header` module,
                        // this cannot fail.
                        .unwrap();
                    Ok(response)
                }
            }

            impl ::std::convert::TryFrom<#ruma_api_import::exports::http::Response<Vec<u8>>> for Response {
                type Error = #ruma_api_import::error::FromHttpResponseError<#error>;

                #[allow(unused_variables)]
                fn try_from(
                    response: #ruma_api_import::exports::http::Response<Vec<u8>>,
                ) -> ::std::result::Result<Self, Self::Error> {
                    if response.status().as_u16() < 400 {
                        #extract_response_headers

                        #typed_response_body_decl

                        Ok(Self {
                            #response_init_fields
                        })
                    } else {
                        match <#error as #ruma_api_import::EndpointError>::try_from_response(response) {
                            Ok(err) => Err(#ruma_api_import::error::ServerError::Known(err).into()),
                            Err(response_err) => {
                                Err(#ruma_api_import::error::ServerError::Unknown(response_err).into())
                            }
                        }
                    }
                }
            }

            impl #request_lifetimes #ruma_api_import::Endpoint for Request #request_lifetimes {
                type Response = Response;
                type ResponseError = #error;
                type IncomingRequest = <Self as #ruma_api_import::Outgoing>::Incoming;
                type IncomingResponse = <Response as #ruma_api_import::Outgoing>::Incoming;

                /// Metadata for the `#name` endpoint.
                const METADATA: #ruma_api_import::Metadata = #ruma_api_import::Metadata {
                    description: #description,
                    method: #ruma_api_import::exports::http::Method::#method,
                    name: #name,
                    path: #path,
                    rate_limited: #rate_limited,
                    requires_authentication: #requires_authentication,
                };

                #[allow(unused_mut, unused_variables)]
                fn try_into_http_request(
                    self,
                    base_url: &::std::primitive::str,
                    access_token: ::std::option::Option<&str>,
                ) -> ::std::result::Result<
                    #ruma_api_import::exports::http::Request<Vec<u8>>,
                    #ruma_api_import::error::IntoHttpError,
                > {
                    let metadata = Request::METADATA;

                    let http_request = #ruma_api_import::exports::http::Request::builder()
                        .method(#ruma_api_import::exports::http::Method::#method)
                        .uri(::std::format!(
                            "{}{}{}",
                            base_url.strip_suffix("/").unwrap_or(base_url),
                            #request_path_string,
                            #request_query_string,
                        ))
                        #( .header(#header_kvs) )*
                        .body(#request_body)?;

                    Ok(http_request)
                }
            }

            #non_auth_endpoint_impl
        };

        api.to_tokens(tokens);
    }
}

/// Custom keyword macros for syn.
mod kw {
    use syn::custom_keyword;

    custom_keyword!(metadata);
    custom_keyword!(request);
    custom_keyword!(response);
    custom_keyword!(error);
}

/// The entire `ruma_api!` macro structure directly as it appears in the source code..
pub struct RawApi {
    /// The `metadata` section of the macro.
    pub metadata: RawMetadata,
    /// The `request` section of the macro.
    pub request: RawRequest,
    /// The `response` section of the macro.
    pub response: RawResponse,
    /// The `error` section of the macro.
    pub error: Option<RawErrorType>,
}

impl Parse for RawApi {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        Ok(Self {
            metadata: input.parse()?,
            request: input.parse()?,
            response: input.parse()?,
            error: input.parse().ok(),
        })
    }
}

pub struct RawMetadata {
    pub metadata_kw: kw::metadata,
    pub field_values: Vec<FieldValue>,
}

impl Parse for RawMetadata {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let metadata_kw = input.parse::<kw::metadata>()?;
        input.parse::<Token![:]>()?;
        let field_values;
        braced!(field_values in input);

        Ok(Self {
            metadata_kw,
            field_values: field_values
                .parse_terminated::<FieldValue, Token![,]>(FieldValue::parse)?
                .into_iter()
                .collect(),
        })
    }
}

pub struct RawRequest {
    pub request_kw: kw::request,
    pub fields: Vec<Field>,
}

impl Parse for RawRequest {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let request_kw = input.parse::<kw::request>()?;
        input.parse::<Token![:]>()?;
        let fields;
        braced!(fields in input);

        Ok(Self {
            request_kw,
            fields: fields
                .parse_terminated::<Field, Token![,]>(Field::parse_named)?
                .into_iter()
                .collect(),
        })
    }
}

pub struct RawResponse {
    pub response_kw: kw::response,
    pub fields: Vec<Field>,
}

impl Parse for RawResponse {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let response_kw = input.parse::<kw::response>()?;
        input.parse::<Token![:]>()?;
        let fields;
        braced!(fields in input);

        Ok(Self {
            response_kw,
            fields: fields
                .parse_terminated::<Field, Token![,]>(Field::parse_named)?
                .into_iter()
                .map(|f| {
                    if util::has_lifetime(&f.ty) {
                        Err(syn::Error::new(
                            f.ident.span(),
                            "Lifetimes on Response fields cannot be supported until GAT are stable",
                        ))
                    } else {
                        Ok(f)
                    }
                })
                .collect::<Result<Vec<_>, _>>()?,
        })
    }
}

pub struct RawErrorType {
    pub error_kw: kw::error,
    pub ty: Type,
}

impl Parse for RawErrorType {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let error_kw = input.parse::<kw::error>()?;
        input.parse::<Token![:]>()?;
        let ty = input.parse()?;

        Ok(Self { error_kw, ty })
    }
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

fn main() -> Result<(), Box<dyn std::error::Error>> {
    tonic_build::compile_protos("../../format/Flight.proto")?;
    Ok(())
}
<commit_msg>ARROW-7794: [Rust] [Flight] Remove hard-coded relative path to Flight.proto<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use std::env;

fn main() -> Result<(), String> {
    // The current working directory can vary depending on how the project is being
    // built or released so we build an absolute path to the proto file
    let mut path = env::current_dir().map_err(|e| format!("{:?}", e))?;
    loop {
        // build path to format/Flight.proto
        let mut flight_proto = path.clone();
        flight_proto.push("format");
        flight_proto.push("Flight.proto");

        if flight_proto.exists() {
            tonic_build::compile_protos(flight_proto).map_err(|e| format!("{:?}", e))?;
            return Ok(());
        }

        if !path.pop() {
            // reached root of file system
            break;
        }
    }

    Err("Failed to locate format/Flight.proto in any parent directory".to_string())
}
<|endoftext|>


<commit_before>//! This example shows how to configure fern to output really nicely colored logs
//!  - when the log level is error, the whole line is red
//!  - when the log level is warn, the whole line is yellow
//!  - when the log level is info, the level name is green and the rest of the line is white
//!  - when the log level is debug, the whole line is white
//!  - when the log level is trace, the whole line is gray ("bright black")

extern crate chrono;
extern crate fern;
#[macro_use]
extern crate log;

use fern::colors::{Color, ColoredLevelConfig};

// ===================== Logging Set Up =====================
fn main() {
    set_up_logging();
    // let's simulate some logging
    info!("starting simulation!");
    for i in 0..26 {
        trace!("loading: {}%, very verbose debbuging information", 4*i);
        if 5 == i {
            debug!("this is taking so long... boooring!");
        } else if 10 == i {
            debug!("still alive! yay!");
        } else if 13 == i {
            info!("halfway there!");
        } else if 16 == i {
            debug!("*scratches nose*");
            warn!("nose is itching, continuing anyways");
        } else if 20 == i {
            debug!("uh oh");
            warn!(">nose itching intensifies");
            error!("HATCHOOO!");
            debug!("encountered minor problem, trying to recover");
            info!("gesundheit");
            debug!("recovered from minor problem, continuing");
        } else if 25 == i {
            info!("successfully loaded nothing");
            info!("have a good time!");
        }
    }
}

// ===================== Logging Set Up =====================
fn set_up_logging() {
    // configure colors for the whole line
    let colors_line = ColoredLevelConfig::new()
        .error(Color::Red)
        .warn(Color::Yellow)
        // we actually don't need to specify the color for debug and info, they are white by default
        .info(Color::White)
        .debug(Color::White)
        // depending on the terminals color scheme, this is the same as the background color
        .trace(Color::BrightBlack);

    // configure colors for the name of the level.
    // since almost all of them are the some as the color for the whole line, we just clone
    // `colors_line` and overwrite our changes
    let colors_level = colors_line.clone()
        .info(Color::Green);
    // here we set up our fern Dispatch
    fern::Dispatch::new()
        .format(move |out, message, record| {
            out.finish(format_args!(
                "{color_line}[{date}][{target}][{level}{color_line}] {message}\x1B[0m",
                color_line = format_args!("\x1B[{}m", colors_line.get_color(&record.level()).to_fg_str()),
                date = chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
                target = record.target(),
                level = colors_level.color(record.level()),
                message = message,
            ));
        })
        // set the default log level
        .level(log::LevelFilter::Trace)
        // set module (actually, it's target) specific log levels
        .level_for("more_color", log::LevelFilter::Trace)
        .level_for("silly", log::LevelFilter::Trace)
        // output to stdout
        .chain(std::io::stdout())
        .apply().unwrap();

    debug!("finished setting up logging! yay!");
}
<commit_msg>Fix minor mistakes<commit_after>//! This example shows how to configure fern to output really nicely colored logs
//!  - when the log level is error, the whole line is red
//!  - when the log level is warn, the whole line is yellow
//!  - when the log level is info, the level name is green and the rest of the line is white
//!  - when the log level is debug, the whole line is white
//!  - when the log level is trace, the whole line is gray ("bright black")

extern crate chrono;
extern crate fern;
#[macro_use]
extern crate log;

use fern::colors::{Color, ColoredLevelConfig};

fn main() {
    set_up_logging();
    // let's simulate some logging
    info!("starting simulation!");
    for i in 0..26 {
        trace!("loading: {}%, very verbose debbuging information", 4*i);
        if 5 == i {
            debug!("this is taking so long... boooring!");
        } else if 10 == i {
            debug!("still alive! yay!");
        } else if 13 == i {
            info!("halfway there!");
        } else if 16 == i {
            debug!("*scratches nose*");
            warn!("nose is itching, continuing anyways");
        } else if 20 == i {
            debug!("uh oh");
            warn!(">nose itching intensifies");
            error!("HATCHOOO!");
            debug!("encountered minor problem, trying to recover");
            info!("gesundheit");
            debug!("recovered from minor problem, continuing");
        } else if 25 == i {
            info!("successfully loaded nothing");
            info!("have a good time!");
        }
    }
}

// ===================== Logging Set Up =====================
fn set_up_logging() {
    // configure colors for the whole line
    let colors_line = ColoredLevelConfig::new()
        .error(Color::Red)
        .warn(Color::Yellow)
        // we actually don't need to specify the color for debug and info, they are white by default
        .info(Color::White)
        .debug(Color::White)
        // depending on the terminals color scheme, this is the same as the background color
        .trace(Color::BrightBlack);

    // configure colors for the name of the level.
    // since almost all of them are the some as the color for the whole line, we just clone
    // `colors_line` and overwrite our changes
    let colors_level = colors_line.clone()
        .info(Color::Green);
    // here we set up our fern Dispatch
    fern::Dispatch::new()
        .format(move |out, message, record| {
            out.finish(format_args!(
                "{color_line}[{date}][{target}][{level}{color_line}] {message}\x1B[0m",
                color_line = format_args!("\x1B[{}m", colors_line.get_color(&record.level()).to_fg_str()),
                date = chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
                target = record.target(),
                level = colors_level.color(record.level()),
                message = message,
            ));
        })
        // set the default log level
        .level(log::LevelFilter::Warn)
        // set module (actually, it's target) specific log levels
        .level_for("pretty_colored", log::LevelFilter::Trace)
        // output to stdout
        .chain(std::io::stdout())
        .apply().unwrap();

    debug!("finished setting up logging! yay!");
}
<|endoftext|>


<commit_before>//! Procedural macro for the `#[async]` attribute.
//!
//! This crate is an implementation of the `#[async]` attribute as a procedural
//! macro. This is nightly-only for now as it's using the unstable features of
//! procedural macros. Furthermore it's generating code that's using a new
//! keyword, `yield`, and a new construct, generators, both of which are also
//! unstable.
//!
//! Currently this crate depends on `syn` and `quote` to do all the heavy
//! lifting, this is just a very small shim around creating a closure/future out
//! of a generator.

#![feature(proc_macro)]
#![recursion_limit = "128"]

extern crate proc_macro;
extern crate proc_macro2;
#[macro_use]
extern crate quote;
extern crate syn;

use proc_macro::{TokenStream, TokenTree, Delimiter, TokenNode};
use syn::*;
use syn::fold::Folder;

#[proc_macro_attribute]
pub fn async(attribute: TokenStream, function: TokenStream) -> TokenStream {
    // Handle arguments to the #[async] attribute, if any
    let attribute = attribute.to_string();
    let boxed = if attribute == "( boxed )" {
        true
    } else if attribute == "" {
        false
    } else {
        panic!("the #[async] attribute currently only takes `boxed` as an arg");
    };

    // Parse our item, expecting a function. This function may be an actual
    // top-level function or it could be a method (typically dictated by the
    // arguments). We then extract everything we'd like to use.
    let Item { attrs, node } = syn::parse(function)
        .expect("failed to parse tokens as a function");
    let ItemFn {
        ident,
        vis,
        unsafety,
        constness,
        abi,
        block,
        decl,
        ..
    } = match node {
        ItemKind::Fn(item) => item,
        _ => panic!("#[async] can only be applied to functions"),
    };
    let FnDecl { inputs, output, variadic, generics, .. } = { *decl };
    let where_clause = &generics.where_clause;
    assert!(!variadic, "variadic functions cannot be async");
    let output = match output {
        FunctionRetTy::Ty(t, _) => t,
        FunctionRetTy::Default => {
            TyTup {
                tys: Default::default(),
                lone_comma: Default::default(),
                paren_token: Default::default(),
            }.into()
        }
    };

    // We've got to get a bit creative with our handling of arguments. For a
    // number of reasons we translate this:
    //
    //      fn foo(ref a: u32) -> Result<u32, u32> {
    //          // ...
    //      }
    //
    // into roughly:
    //
    //      fn foo(__arg_0: u32) -> impl Future<...> {
    //          gen(move || {
    //              let ref a = __arg0;
    //
    //              // ...
    //          })
    //      }
    //
    // The intention here is to ensure that all local function variables get
    // moved into the generator we're creating, and they're also all then bound
    // appropriately according to their patterns and whatnot.
    //
    // We notably skip everything related to `self` which typically doesn't have
    // many patterns with it and just gets captured naturally.
    let mut inputs_no_patterns = Vec::new();
    let mut patterns = Vec::new();
    let mut temp_bindings = Vec::new();
    for (i, input) in inputs.into_iter().enumerate() {
        let input = input.into_item();

        // `self: Box<Self>` will get captured naturally
        let mut is_input_no_pattern = false;
        if let FnArg::Captured(ref arg) = input {
            if let Pat::Ident(PatIdent { ref ident, ..}) = arg.pat {
                if ident == "self" {
                    is_input_no_pattern = true;
                }
            }
        }
        if is_input_no_pattern {
            inputs_no_patterns.push(input);
            continue
        }

        match input {
            // `a: B`
            FnArg::Captured(ArgCaptured { pat, ty, .. }) => {
                patterns.push(pat);
                let ident = Ident::from(format!("__arg_{}", i));
                temp_bindings.push(ident.clone());
                let pat = PatIdent {
                    mode: BindingMode::ByValue(Mutability::Immutable),
                    ident: ident,
                    at_token: None,
                    subpat: None,
                };
                inputs_no_patterns.push(ArgCaptured {
                    pat: pat.into(),
                    ty: ty,
                    colon_token: Default::default(),
                }.into());
            }

            // Other `self`-related arguments get captured naturally
            _ => {
                inputs_no_patterns.push(input);
            }
        }
    }

    // This is the point where we handle
    //
    //      #[async]
    //      for x in y {
    //      }
    //
    // Basically just take all those expression and expand them.
    let block = ExpandAsyncFor.fold_block(*block);

    // TODO: can we lift the restriction that `futures` must be at the root of
    //       the crate?
    let return_ty = if boxed {
        quote! {
            Box<::futures::Future<
                Item = <#output as ::futures::__rt::MyTry>::MyOk,
                Error = <#output as ::futures::__rt::MyTry>::MyError,
            >>
        }
    } else {
        // Dunno why this is buggy, hits weird typecheck errors in tests
        //
        // quote! {
        //     impl ::futures::Future<
        //         Item = <#output as ::futures::__rt::MyTry>::MyOk,
        //         Error = <#output as ::futures::__rt::MyTry>::MyError,
        //     >
        // }

        quote! { impl ::futures::__rt::MyFuture<#output> }
    };

    let maybe_boxed = if boxed {
        quote! { Box::new }
    } else {
        quote! { }
    };

    let output = quote! {
        #(#attrs)*
        #vis #unsafety #abi #constness
        fn #ident #generics(#(#inputs_no_patterns),*) -> #return_ty
            #where_clause
        {
            #maybe_boxed (::futures::__rt::gen(move || {
                let __e: #output = {
                    #( let #patterns = #temp_bindings; )*
                    #block
                };

                // Ensure that this closure is a generator, even if it doesn't
                // have any `yield` statements.
                #[allow(unreachable_code)]
                {
                    return __e;
                    loop { yield }
                }
            }))
        }
    };

    // println!("{}", output);
    output.into()
}

#[proc_macro]
pub fn async_block(input: TokenStream) -> TokenStream {
    let input = TokenStream::from(TokenTree {
        kind: TokenNode::Group(Delimiter::Brace, input),
        span: Default::default(),
    });
    let expr = syn::parse(input)
        .expect("failed to parse tokens as an expression");
    let expr = ExpandAsyncFor.fold_expr(expr);

    (quote! {
        ::futures::__rt::gen(move || { #expr })
    }).into()
}

struct ExpandAsyncFor;

impl Folder for ExpandAsyncFor {
    fn fold_expr(&mut self, expr: Expr) -> Expr {
        let expr = fold::noop_fold_expr(self, expr);
        if expr.attrs.len() != 1 {
            return expr
        }
        // TODO: more validation here
        if expr.attrs[0].path.segments.get(0).item().ident != "async" {
            return expr
        }
        let all = match expr.node {
            ExprKind::ForLoop(item) => item,
            _ => panic!("only for expressions can have #[async]"),
        };
        let ExprForLoop { pat, expr, body, label, colon_token, .. } = all;

        // Basically just expand to a `poll` loop
        let tokens = quote! {{
            let mut __stream = #expr;
            #label
            #colon_token
            loop {
                let #pat = {
                    extern crate futures_await;
                    let r = futures_await::Stream::poll(&mut __stream)?;
                    match r {
                        futures_await::Async::Ready(e) => {
                            match e {
                                futures_await::__rt::Some(e) => e,
                                futures_await::__rt::None => break,
                            }
                        }
                        futures_await::Async::NotReady => {
                            yield;
                            continue
                        }
                    }
                };

                #body
            }
        }};
        syn::parse(tokens.into()).unwrap()
    }

    // Don't recurse into items
    fn fold_item(&mut self, item: Item) -> Item {
        item
    }
}
<commit_msg>Remove an unused crate<commit_after>//! Procedural macro for the `#[async]` attribute.
//!
//! This crate is an implementation of the `#[async]` attribute as a procedural
//! macro. This is nightly-only for now as it's using the unstable features of
//! procedural macros. Furthermore it's generating code that's using a new
//! keyword, `yield`, and a new construct, generators, both of which are also
//! unstable.
//!
//! Currently this crate depends on `syn` and `quote` to do all the heavy
//! lifting, this is just a very small shim around creating a closure/future out
//! of a generator.

#![feature(proc_macro)]
#![recursion_limit = "128"]

extern crate proc_macro;
#[macro_use]
extern crate quote;
extern crate syn;

use proc_macro::{TokenStream, TokenTree, Delimiter, TokenNode};
use syn::*;
use syn::fold::Folder;

#[proc_macro_attribute]
pub fn async(attribute: TokenStream, function: TokenStream) -> TokenStream {
    // Handle arguments to the #[async] attribute, if any
    let attribute = attribute.to_string();
    let boxed = if attribute == "( boxed )" {
        true
    } else if attribute == "" {
        false
    } else {
        panic!("the #[async] attribute currently only takes `boxed` as an arg");
    };

    // Parse our item, expecting a function. This function may be an actual
    // top-level function or it could be a method (typically dictated by the
    // arguments). We then extract everything we'd like to use.
    let Item { attrs, node } = syn::parse(function)
        .expect("failed to parse tokens as a function");
    let ItemFn {
        ident,
        vis,
        unsafety,
        constness,
        abi,
        block,
        decl,
        ..
    } = match node {
        ItemKind::Fn(item) => item,
        _ => panic!("#[async] can only be applied to functions"),
    };
    let FnDecl { inputs, output, variadic, generics, .. } = { *decl };
    let where_clause = &generics.where_clause;
    assert!(!variadic, "variadic functions cannot be async");
    let output = match output {
        FunctionRetTy::Ty(t, _) => t,
        FunctionRetTy::Default => {
            TyTup {
                tys: Default::default(),
                lone_comma: Default::default(),
                paren_token: Default::default(),
            }.into()
        }
    };

    // We've got to get a bit creative with our handling of arguments. For a
    // number of reasons we translate this:
    //
    //      fn foo(ref a: u32) -> Result<u32, u32> {
    //          // ...
    //      }
    //
    // into roughly:
    //
    //      fn foo(__arg_0: u32) -> impl Future<...> {
    //          gen(move || {
    //              let ref a = __arg0;
    //
    //              // ...
    //          })
    //      }
    //
    // The intention here is to ensure that all local function variables get
    // moved into the generator we're creating, and they're also all then bound
    // appropriately according to their patterns and whatnot.
    //
    // We notably skip everything related to `self` which typically doesn't have
    // many patterns with it and just gets captured naturally.
    let mut inputs_no_patterns = Vec::new();
    let mut patterns = Vec::new();
    let mut temp_bindings = Vec::new();
    for (i, input) in inputs.into_iter().enumerate() {
        let input = input.into_item();

        // `self: Box<Self>` will get captured naturally
        let mut is_input_no_pattern = false;
        if let FnArg::Captured(ref arg) = input {
            if let Pat::Ident(PatIdent { ref ident, ..}) = arg.pat {
                if ident == "self" {
                    is_input_no_pattern = true;
                }
            }
        }
        if is_input_no_pattern {
            inputs_no_patterns.push(input);
            continue
        }

        match input {
            // `a: B`
            FnArg::Captured(ArgCaptured { pat, ty, .. }) => {
                patterns.push(pat);
                let ident = Ident::from(format!("__arg_{}", i));
                temp_bindings.push(ident.clone());
                let pat = PatIdent {
                    mode: BindingMode::ByValue(Mutability::Immutable),
                    ident: ident,
                    at_token: None,
                    subpat: None,
                };
                inputs_no_patterns.push(ArgCaptured {
                    pat: pat.into(),
                    ty: ty,
                    colon_token: Default::default(),
                }.into());
            }

            // Other `self`-related arguments get captured naturally
            _ => {
                inputs_no_patterns.push(input);
            }
        }
    }

    // This is the point where we handle
    //
    //      #[async]
    //      for x in y {
    //      }
    //
    // Basically just take all those expression and expand them.
    let block = ExpandAsyncFor.fold_block(*block);

    // TODO: can we lift the restriction that `futures` must be at the root of
    //       the crate?
    let return_ty = if boxed {
        quote! {
            Box<::futures::Future<
                Item = <#output as ::futures::__rt::MyTry>::MyOk,
                Error = <#output as ::futures::__rt::MyTry>::MyError,
            >>
        }
    } else {
        // Dunno why this is buggy, hits weird typecheck errors in tests
        //
        // quote! {
        //     impl ::futures::Future<
        //         Item = <#output as ::futures::__rt::MyTry>::MyOk,
        //         Error = <#output as ::futures::__rt::MyTry>::MyError,
        //     >
        // }

        quote! { impl ::futures::__rt::MyFuture<#output> }
    };

    let maybe_boxed = if boxed {
        quote! { Box::new }
    } else {
        quote! { }
    };

    let output = quote! {
        #(#attrs)*
        #vis #unsafety #abi #constness
        fn #ident #generics(#(#inputs_no_patterns),*) -> #return_ty
            #where_clause
        {
            #maybe_boxed (::futures::__rt::gen(move || {
                let __e: #output = {
                    #( let #patterns = #temp_bindings; )*
                    #block
                };

                // Ensure that this closure is a generator, even if it doesn't
                // have any `yield` statements.
                #[allow(unreachable_code)]
                {
                    return __e;
                    loop { yield }
                }
            }))
        }
    };

    // println!("{}", output);
    output.into()
}

#[proc_macro]
pub fn async_block(input: TokenStream) -> TokenStream {
    let input = TokenStream::from(TokenTree {
        kind: TokenNode::Group(Delimiter::Brace, input),
        span: Default::default(),
    });
    let expr = syn::parse(input)
        .expect("failed to parse tokens as an expression");
    let expr = ExpandAsyncFor.fold_expr(expr);

    (quote! {
        ::futures::__rt::gen(move || { #expr })
    }).into()
}

struct ExpandAsyncFor;

impl Folder for ExpandAsyncFor {
    fn fold_expr(&mut self, expr: Expr) -> Expr {
        let expr = fold::noop_fold_expr(self, expr);
        if expr.attrs.len() != 1 {
            return expr
        }
        // TODO: more validation here
        if expr.attrs[0].path.segments.get(0).item().ident != "async" {
            return expr
        }
        let all = match expr.node {
            ExprKind::ForLoop(item) => item,
            _ => panic!("only for expressions can have #[async]"),
        };
        let ExprForLoop { pat, expr, body, label, colon_token, .. } = all;

        // Basically just expand to a `poll` loop
        let tokens = quote! {{
            let mut __stream = #expr;
            #label
            #colon_token
            loop {
                let #pat = {
                    extern crate futures_await;
                    let r = futures_await::Stream::poll(&mut __stream)?;
                    match r {
                        futures_await::Async::Ready(e) => {
                            match e {
                                futures_await::__rt::Some(e) => e,
                                futures_await::__rt::None => break,
                            }
                        }
                        futures_await::Async::NotReady => {
                            yield;
                            continue
                        }
                    }
                };

                #body
            }
        }};
        syn::parse(tokens.into()).unwrap()
    }

    // Don't recurse into items
    fn fold_item(&mut self, item: Item) -> Item {
        item
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::collections::HashMap;
use std::collections::HashSet;
use std::collections::VecDeque;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering;

use crate::datatype::AbstractDomain;
use crate::graph::Graph;
use crate::graph::SuccessorNodes;
use crate::wpo::WeakPartialOrdering;
use crate::wpo::WpoIdx;

// Unlike the C++ version, we don't treat this as a base interface
// for all iterators, instead, it's only a trait for concrete type
// to analyze nodes and edges. And this can be composited in the
// fixpoint iterator type. We will revisit this part in the following
// diff for analyzer interface.
pub trait FixpointIteratorTransformer<G: Graph, D: AbstractDomain> {
    /// The *current_state* could be updated in place.
    fn analyze_node(&mut self, n: G::NodeId, current_state: &mut D);

    fn analyze_edge(&mut self, e: G::EdgeId, exit_state_at_src: &D) -> D;
}

pub struct MonotonicFixpointIteratorContext<G: Graph, D: AbstractDomain> {
    init_value: D,
    local_iterations: HashMap<G::NodeId, u32>,
    global_iterations: HashMap<G::NodeId, u32>,
}

impl<G, D> MonotonicFixpointIteratorContext<G, D>
where
    G: Graph,
    D: AbstractDomain,
{
    pub fn get_local_iterations_for(&self, n: G::NodeId) -> u32 {
        *self.local_iterations.get(&n).unwrap_or(&0)
    }

    pub fn get_global_iterations_for(&self, n: G::NodeId) -> u32 {
        *self.global_iterations.get(&n).unwrap_or(&0)
    }

    pub fn get_init_value(&self) -> &D {
        &self.init_value
    }

    fn increase_iteration_count(n: G::NodeId, table: &mut HashMap<G::NodeId, u32>) {
        *table.entry(n).or_default() += 1;
    }

    pub fn increase_iteration_count_for(&mut self, n: G::NodeId) {
        Self::increase_iteration_count(n, &mut self.local_iterations);
        Self::increase_iteration_count(n, &mut self.global_iterations);
    }

    pub fn reset_local_iteration_count_for(&mut self, n: G::NodeId) {
        *self.local_iterations.entry(n).or_default() = 0;
    }

    pub fn new(init_value: D) -> Self {
        Self {
            init_value,
            local_iterations: Default::default(),
            global_iterations: Default::default(),
        }
    }

    pub fn with_nodes(mut self, nodes: &HashSet<G::NodeId>) -> Self {
        for &node in nodes {
            *self.global_iterations.entry(node).or_default() = 0;
            *self.local_iterations.entry(node).or_default() = 0;
        }
        self
    }
}

pub struct MonotonicFixpointIterator<
    'g,
    G: Graph,
    D: AbstractDomain,
    T: FixpointIteratorTransformer<G, D>,
> {
    graph: &'g G,
    entry_states: HashMap<G::NodeId, D>,
    exit_states: HashMap<G::NodeId, D>,
    transformer: T,
    wpo: WeakPartialOrdering<G::NodeId>,
}

impl<'g, G, D, T> MonotonicFixpointIterator<'g, G, D, T>
where
    G: Graph,
    D: AbstractDomain,
    T: FixpointIteratorTransformer<G, D>,
{
    pub fn new<SN>(g: &'g G, cfg_size_hint: usize, transformer: T, successors_nodes: &SN) -> Self
    where
        SN: SuccessorNodes<NodeId = G::NodeId>,
    {
        let wpo = WeakPartialOrdering::new(g.entry(), g.size(), successors_nodes);
        Self {
            graph: g,
            entry_states: HashMap::with_capacity(cfg_size_hint),
            exit_states: HashMap::with_capacity(cfg_size_hint),
            transformer,
            wpo,
        }
    }

    pub fn run(&mut self, init_value: D) {
        self.clear();

        let mut context = MonotonicFixpointIteratorContext::new(init_value);
        let wpo_counter: Vec<AtomicU32> =
            (0..self.wpo.size()).map(|_| Default::default()).collect();

        let mut worklist = VecDeque::new();
        let entry_idx = self.wpo.get_entry();
        worklist.push_front(entry_idx);
        assert_eq!(self.wpo.get_num_preds(entry_idx), 0);

        let mut process_node = |wpo_idx: WpoIdx, worklist: &mut VecDeque<WpoIdx>| {
            assert_eq!(
                wpo_counter[wpo_idx as usize].load(Ordering::Relaxed),
                self.wpo.get_num_preds(wpo_idx)
            );

            wpo_counter[wpo_idx as usize].store(0, Ordering::Relaxed);

            if !self.wpo.is_exit(wpo_idx) {
                self.analyze_vertex(&context, self.wpo.get_node(wpo_idx));

                for &succ_idx in self.wpo.get_successors(wpo_idx) {
                    let old_counter =
                        wpo_counter[succ_idx as usize].fetch_add(1, Ordering::Relaxed);
                    if old_counter + 1 == self.wpo.get_num_preds(succ_idx) {
                        worklist.push_back(succ_idx);
                    }
                }

                return;
            }

            let head_idx = self.wpo.get_head_of_exit(wpo_idx);
            let head = self.wpo.get_node(head_idx);
            let current_state = self.entry_states.entry(head).or_insert_with(D::bottom);
            let mut new_state = D::bottom();
            Self::compute_entry_state(
                self.graph,
                &self.exit_states,
                &mut self.transformer,
                &context,
                head,
                &mut new_state,
            );

            if new_state.leq(current_state) {
                context.reset_local_iteration_count_for(head);
                *current_state = new_state;

                for &succ_idx in self.wpo.get_successors(wpo_idx) {
                    let old_counter =
                        wpo_counter[succ_idx as usize].fetch_add(1, Ordering::Relaxed);
                    if old_counter + 1 == self.wpo.get_num_preds(succ_idx) {
                        worklist.push_back(succ_idx);
                    }
                }
            } else {
                Self::extrapolate(&context, head, current_state, new_state);
                context.increase_iteration_count_for(head);
                for (&component_idx, &num) in self.wpo.get_num_outer_preds(wpo_idx) {
                    if component_idx != entry_idx {
                        let old_counter =
                            wpo_counter[component_idx as usize].fetch_add(num, Ordering::Relaxed);
                        if old_counter + num == self.wpo.get_num_preds(component_idx) {
                            worklist.push_back(component_idx);
                        }
                    }
                }

                if head_idx == entry_idx {
                    worklist.push_back(head_idx);
                }
            }
        };

        while let Some(idx) = worklist.pop_front() {
            process_node(idx, &mut worklist);
        }

        for counter in wpo_counter {
            assert_eq!(counter.load(Ordering::Relaxed), 0);
        }
    }

    /// Default strategy for applying widening operator (apply
    /// join at the first iteration and then widening in all
    /// rest iterations).
    pub fn extrapolate(
        context: &MonotonicFixpointIteratorContext<G, D>,
        n: G::NodeId,
        current_state: &mut D,
        new_state: D,
    ) {
        if 0 == context.get_global_iterations_for(n) {
            // TODO: we need to revisit this design, should we use
            // move or clone of domain?
            current_state.join_with(new_state);
        } else {
            current_state.widen_with(new_state);
        }
    }

    fn get_state_at_or_bottom(states: &HashMap<G::NodeId, D>, n: G::NodeId) -> D {
        states.get(&n).cloned().unwrap_or_else(D::bottom)
    }

    pub fn get_entry_state_at(&self, n: G::NodeId) -> D {
        Self::get_state_at_or_bottom(&self.entry_states, n)
    }

    pub fn get_exit_state_at(&self, n: G::NodeId) -> D {
        Self::get_state_at_or_bottom(&self.exit_states, n)
    }

    pub fn clear(&mut self) {
        self.entry_states.clear();
        self.entry_states.shrink_to_fit();
        self.exit_states.clear();
        self.exit_states.shrink_to_fit();
    }

    pub fn set_all_to_bottom(&mut self, all_nodes: &HashSet<G::NodeId>) {
        for &node in all_nodes {
            self.entry_states
                .entry(node)
                .and_modify(|s| *s = D::bottom())
                .or_insert_with(D::bottom);
            self.exit_states
                .entry(node)
                .and_modify(|s| *s = D::bottom())
                .or_insert_with(D::bottom);
        }
    }

    pub fn compute_entry_state(
        graph: &'g G,
        exit_states: &HashMap<G::NodeId, D>,
        transformer: &mut T,
        context: &MonotonicFixpointIteratorContext<G, D>,
        n: G::NodeId,
        entry_state: &mut D,
    ) {
        if n == graph.entry() {
            entry_state.join_with(context.get_init_value().clone());
        }

        for e in graph.predecessors(n) {
            let d = Self::get_state_at_or_bottom(exit_states, graph.source(e));
            entry_state.join_with(transformer.analyze_edge(e, &d));
        }
    }

    pub fn analyze_vertex(
        &mut self,
        context: &MonotonicFixpointIteratorContext<G, D>,
        n: G::NodeId,
    ) {
        let entry_state = self.entry_states.entry(n).or_insert_with(D::bottom);
        Self::compute_entry_state(
            self.graph,
            &self.exit_states,
            &mut self.transformer,
            context,
            n,
            entry_state,
        );
        let exit_state = self
            .exit_states
            .entry(n)
            .and_modify(|s| *s = entry_state.clone())
            .or_insert_with(|| entry_state.clone());
        self.transformer.analyze_node(n, exit_state);
    }
}
<commit_msg>Change the if condition to assertion in fixpoint iterator in run function<commit_after>/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::collections::HashMap;
use std::collections::HashSet;
use std::collections::VecDeque;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering;

use crate::datatype::AbstractDomain;
use crate::graph::Graph;
use crate::graph::SuccessorNodes;
use crate::wpo::WeakPartialOrdering;
use crate::wpo::WpoIdx;

// Unlike the C++ version, we don't treat this as a base interface
// for all iterators, instead, it's only a trait for concrete type
// to analyze nodes and edges. And this can be composited in the
// fixpoint iterator type. We will revisit this part in the following
// diff for analyzer interface.
pub trait FixpointIteratorTransformer<G: Graph, D: AbstractDomain> {
    /// The *current_state* could be updated in place.
    fn analyze_node(&mut self, n: G::NodeId, current_state: &mut D);

    fn analyze_edge(&mut self, e: G::EdgeId, exit_state_at_src: &D) -> D;
}

pub struct MonotonicFixpointIteratorContext<G: Graph, D: AbstractDomain> {
    init_value: D,
    local_iterations: HashMap<G::NodeId, u32>,
    global_iterations: HashMap<G::NodeId, u32>,
}

impl<G, D> MonotonicFixpointIteratorContext<G, D>
where
    G: Graph,
    D: AbstractDomain,
{
    pub fn get_local_iterations_for(&self, n: G::NodeId) -> u32 {
        *self.local_iterations.get(&n).unwrap_or(&0)
    }

    pub fn get_global_iterations_for(&self, n: G::NodeId) -> u32 {
        *self.global_iterations.get(&n).unwrap_or(&0)
    }

    pub fn get_init_value(&self) -> &D {
        &self.init_value
    }

    fn increase_iteration_count(n: G::NodeId, table: &mut HashMap<G::NodeId, u32>) {
        *table.entry(n).or_default() += 1;
    }

    pub fn increase_iteration_count_for(&mut self, n: G::NodeId) {
        Self::increase_iteration_count(n, &mut self.local_iterations);
        Self::increase_iteration_count(n, &mut self.global_iterations);
    }

    pub fn reset_local_iteration_count_for(&mut self, n: G::NodeId) {
        *self.local_iterations.entry(n).or_default() = 0;
    }

    pub fn new(init_value: D) -> Self {
        Self {
            init_value,
            local_iterations: Default::default(),
            global_iterations: Default::default(),
        }
    }

    pub fn with_nodes(mut self, nodes: &HashSet<G::NodeId>) -> Self {
        for &node in nodes {
            *self.global_iterations.entry(node).or_default() = 0;
            *self.local_iterations.entry(node).or_default() = 0;
        }
        self
    }
}

pub struct MonotonicFixpointIterator<
    'g,
    G: Graph,
    D: AbstractDomain,
    T: FixpointIteratorTransformer<G, D>,
> {
    graph: &'g G,
    entry_states: HashMap<G::NodeId, D>,
    exit_states: HashMap<G::NodeId, D>,
    transformer: T,
    wpo: WeakPartialOrdering<G::NodeId>,
}

impl<'g, G, D, T> MonotonicFixpointIterator<'g, G, D, T>
where
    G: Graph,
    D: AbstractDomain,
    T: FixpointIteratorTransformer<G, D>,
{
    pub fn new<SN>(g: &'g G, cfg_size_hint: usize, transformer: T, successors_nodes: &SN) -> Self
    where
        SN: SuccessorNodes<NodeId = G::NodeId>,
    {
        let wpo = WeakPartialOrdering::new(g.entry(), g.size(), successors_nodes);
        Self {
            graph: g,
            entry_states: HashMap::with_capacity(cfg_size_hint),
            exit_states: HashMap::with_capacity(cfg_size_hint),
            transformer,
            wpo,
        }
    }

    pub fn run(&mut self, init_value: D) {
        self.clear();

        let mut context = MonotonicFixpointIteratorContext::new(init_value);
        let wpo_counter: Vec<AtomicU32> =
            (0..self.wpo.size()).map(|_| Default::default()).collect();

        let mut worklist = VecDeque::new();
        let entry_idx = self.wpo.get_entry();
        worklist.push_front(entry_idx);
        assert_eq!(self.wpo.get_num_preds(entry_idx), 0);

        let mut process_node = |wpo_idx: WpoIdx, worklist: &mut VecDeque<WpoIdx>| {
            assert_eq!(
                wpo_counter[wpo_idx as usize].load(Ordering::Relaxed),
                self.wpo.get_num_preds(wpo_idx)
            );

            wpo_counter[wpo_idx as usize].store(0, Ordering::Relaxed);

            if !self.wpo.is_exit(wpo_idx) {
                self.analyze_vertex(&context, self.wpo.get_node(wpo_idx));

                for &succ_idx in self.wpo.get_successors(wpo_idx) {
                    let old_counter =
                        wpo_counter[succ_idx as usize].fetch_add(1, Ordering::Relaxed);
                    if old_counter + 1 == self.wpo.get_num_preds(succ_idx) {
                        worklist.push_back(succ_idx);
                    }
                }

                return;
            }

            let head_idx = self.wpo.get_head_of_exit(wpo_idx);
            let head = self.wpo.get_node(head_idx);
            let current_state = self.entry_states.entry(head).or_insert_with(D::bottom);
            let mut new_state = D::bottom();
            Self::compute_entry_state(
                self.graph,
                &self.exit_states,
                &mut self.transformer,
                &context,
                head,
                &mut new_state,
            );

            if new_state.leq(current_state) {
                context.reset_local_iteration_count_for(head);
                *current_state = new_state;

                for &succ_idx in self.wpo.get_successors(wpo_idx) {
                    let old_counter =
                        wpo_counter[succ_idx as usize].fetch_add(1, Ordering::Relaxed);
                    if old_counter + 1 == self.wpo.get_num_preds(succ_idx) {
                        worklist.push_back(succ_idx);
                    }
                }
            } else {
                Self::extrapolate(&context, head, current_state, new_state);
                context.increase_iteration_count_for(head);
                for (&component_idx, &num) in self.wpo.get_num_outer_preds(wpo_idx) {
                    assert!(component_idx != entry_idx);
                    let old_counter =
                        wpo_counter[component_idx as usize].fetch_add(num, Ordering::Relaxed);
                    if old_counter + num == self.wpo.get_num_preds(component_idx) {
                        worklist.push_back(component_idx);
                    }
                }

                if head_idx == entry_idx {
                    worklist.push_back(head_idx);
                }
            }
        };

        while let Some(idx) = worklist.pop_front() {
            process_node(idx, &mut worklist);
        }

        for counter in wpo_counter {
            assert_eq!(counter.load(Ordering::Relaxed), 0);
        }
    }

    /// Default strategy for applying widening operator (apply
    /// join at the first iteration and then widening in all
    /// rest iterations).
    pub fn extrapolate(
        context: &MonotonicFixpointIteratorContext<G, D>,
        n: G::NodeId,
        current_state: &mut D,
        new_state: D,
    ) {
        if 0 == context.get_global_iterations_for(n) {
            // TODO: we need to revisit this design, should we use
            // move or clone of domain?
            current_state.join_with(new_state);
        } else {
            current_state.widen_with(new_state);
        }
    }

    fn get_state_at_or_bottom(states: &HashMap<G::NodeId, D>, n: G::NodeId) -> D {
        states.get(&n).cloned().unwrap_or_else(D::bottom)
    }

    pub fn get_entry_state_at(&self, n: G::NodeId) -> D {
        Self::get_state_at_or_bottom(&self.entry_states, n)
    }

    pub fn get_exit_state_at(&self, n: G::NodeId) -> D {
        Self::get_state_at_or_bottom(&self.exit_states, n)
    }

    pub fn clear(&mut self) {
        self.entry_states.clear();
        self.entry_states.shrink_to_fit();
        self.exit_states.clear();
        self.exit_states.shrink_to_fit();
    }

    pub fn set_all_to_bottom(&mut self, all_nodes: &HashSet<G::NodeId>) {
        for &node in all_nodes {
            self.entry_states
                .entry(node)
                .and_modify(|s| *s = D::bottom())
                .or_insert_with(D::bottom);
            self.exit_states
                .entry(node)
                .and_modify(|s| *s = D::bottom())
                .or_insert_with(D::bottom);
        }
    }

    pub fn compute_entry_state(
        graph: &'g G,
        exit_states: &HashMap<G::NodeId, D>,
        transformer: &mut T,
        context: &MonotonicFixpointIteratorContext<G, D>,
        n: G::NodeId,
        entry_state: &mut D,
    ) {
        if n == graph.entry() {
            entry_state.join_with(context.get_init_value().clone());
        }

        for e in graph.predecessors(n) {
            let d = Self::get_state_at_or_bottom(exit_states, graph.source(e));
            entry_state.join_with(transformer.analyze_edge(e, &d));
        }
    }

    pub fn analyze_vertex(
        &mut self,
        context: &MonotonicFixpointIteratorContext<G, D>,
        n: G::NodeId,
    ) {
        let entry_state = self.entry_states.entry(n).or_insert_with(D::bottom);
        Self::compute_entry_state(
            self.graph,
            &self.exit_states,
            &mut self.transformer,
            context,
            n,
            entry_state,
        );
        let exit_state = self
            .exit_states
            .entry(n)
            .and_modify(|s| *s = entry_state.clone())
            .or_insert_with(|| entry_state.clone());
        self.transformer.analyze_node(n, exit_state);
    }
}
<|endoftext|>


<commit_before>//
// base64.rs - base64 implementation
//
// The Base64 Alphabet
//
//  Value Encoding  Value Encoding  Value Encoding  Value Encoding
//   0 A            17 R            34 i            51 z
//   1 B            18 S            35 j            52 0
//   2 C            19 T            36 k            53 1
//   3 D            20 U            37 l            54 2
//   4 E            21 V            38 m            55 3
//   5 F            22 W            39 n            56 4
//   6 G            23 X            40 o            57 5
//   7 H            24 Y            41 p            58 6
//   8 I            25 Z            42 q            59 7
//   9 J            26 a            43 r            60 8
//  10 K            27 b            44 s            61 9
//  11 L            28 c            45 t            62 +
//  12 M            29 d            46 u            63 /
//  13 N            30 e            47 v
//  14 O            31 f            48 w         (pad) =
//  15 P            32 g            49 x
//  16 Q            33 h            50 y
//

use std;

import std::map;

export base64, mk_base64;

iface base64 {
    fn encode(src: [u8]) -> [u8];
    fn decode(src: [u8]) -> [u8];
    fn urlsafe_encode(src: [u8]) -> [u8];
    fn urlsafe_decode(src: [u8]) -> [u8];
}

fn mk_base64() -> base64 {
    type _base64 = {padd: u8, table: [u8]};
    impl of base64 for _base64 {
        fn encode(src: [u8]) -> [u8] {
            let srclen = vec::len(src);
            let targ = if srclen % 3u == 0u {
                vec::init_elt_mut((srclen / 3u) * 4u, 0u8)
            } else {
                vec::init_elt_mut((srclen / 3u + 1u) * 4u, 0u8)
            };
            let input = vec::init_elt_mut(3u, 0u8);
            let output = vec::init_elt_mut(4u, 0u8);
            let curr = 0u, src_curr = 0u;
            let table = self.table;

            while srclen > 2u {
                input[0] = src[src_curr];
                input[1] = src[src_curr + 1u];
                input[2] = src[src_curr + 2u];
                srclen -= 3u; src_curr += 3u;

                output[0] = input[0] >> 2u8;
                output[1] = ((input[0] &  3u8) << 4u8) | (input[1] >> 4u8);
                output[2] = ((input[1] & 15u8) << 2u8) | (input[2] >> 6u8);
                output[3] = input[2] & 63u8;

                targ[curr] = table[output[0]]; curr += 1u;
                targ[curr] = table[output[1]]; curr += 1u;
                targ[curr] = table[output[2]]; curr += 1u;
                targ[curr] = table[output[3]]; curr += 1u;
            }

            if srclen != 0u {
                input[0] = 0u8; input[1] = 0u8; input[2] = 0u8;

                alt srclen {
                  1u { input[0] = src[src_curr]; }
                  2u { input[0] = src[src_curr];
                       input[1] = src[src_curr + 1u]; }
                  _  { }
                }

                output[0] = input[0] >> 2u8;
                output[1] = ((input[0] &  3u8) << 4u8) | (input[1] >> 4u8);
                output[2] = ((input[1] & 15u8) << 2u8) | (input[2] >> 6u8);

                targ[curr] = table[output[0]]; curr += 1u;
                targ[curr] = table[output[1]]; curr += 1u;
                targ[curr] = if srclen == 1u {
                    self.padd
                } else {
                    table[output[2]]
                };
                curr += 1u; targ[curr] = self.padd; curr += 1u;
            }

            vec::from_mut(targ)
        }
        fn decode(src: [u8]) -> [u8] {
            let src = src;
            let srclen = vec::len(src);
            let input  = vec::init_elt_mut(4u, 0u8);
            let output = vec::init_elt_mut(4u, 0u8);
            let targ = if src[vec::len(src) - 2u] == self.padd {
                vec::init_elt_mut(srclen / 4u * 3u - 2u, 0u8)
            } else if src[vec::len(src) - 1u] == self.padd {
                vec::init_elt_mut(srclen / 4u * 3u - 1u, 0u8)
            } else {
                vec::init_elt_mut(srclen / 4u * 3u, 0u8)
            };
            let c = -1;
            let curr = 0u, src_curr = 0u;
            let table = self.table;

            while srclen > 4u {
                input[0] = src[src_curr];
                input[1] = src[src_curr + 1u];
                input[2] = src[src_curr + 2u];
                input[3] = src[src_curr + 3u];
                srclen -= 4u; src_curr += 4u;

                c = idx(table, input[0]);
                if c == -1 { fail "malformed base64 string"; }
                output[0] = c as u8;

                c = idx(table, input[1]);
                if c == -1 { fail "malformed base64 string"; }
                output[1] = c as u8;

                c = idx(table, input[2]);
                if c == -1 { fail "malformed base64 string"; }
                output[2] = c as u8;

                c = idx(table, input[3]);
                if c == -1 { fail "malformed base64 string"; }
                output[3] = c as u8;

                targ[curr] = (output[0] << 2u8) | (output[1] >> 4u8);
                curr += 1u;

                targ[curr] =
                    ((output[1] & 15u8) << 4u8) | (output[2] >> 2u8);
                curr += 1u;

                targ[curr] = ((output[2] &  3u8) << 6u8) | output[3];
                curr += 1u;
            }

            if srclen == 4u {
                input[0] = src[src_curr];
                input[1] = src[src_curr + 1u];
                input[2] = src[src_curr + 2u];
                input[3] = src[src_curr + 3u];

                c = idx(table, input[0]);
                if c == -1 { fail "malformed base64 string"; }
                output[0] = c as u8;

                c = idx(table, input[1]);
                if c == -1 { fail "malformed base64 string"; }
                output[1] = c as u8;

                targ[curr] = (output[0] << 2u8) | (output[1] >> 4u8);
                curr += 1u;

                if input[2] == self.padd { ret vec::from_mut(targ); }

                c = idx(table, input[2]);
                if c == -1 { fail "malformed base64 string"; }
                output[2] = c as u8;

                targ[curr] = ((output[1] & 15u8) << 4u8) | (output[2] >> 2u8);
                curr += 1u;

                if input[3] == self.padd { ret vec::from_mut(targ); }

                c = idx(table, input[3]);
                if c == -1 { fail "malformed base64 string"; }
                output[3] = c as u8;

                targ[curr] = ((output[2] & 3u8) << 6u8) | output[3];
            } else {
                fail "malformed base64 string";
            }

            vec::from_mut(targ)
        }
        fn urlsafe_encode(src: [u8]) -> [u8] {
            let res = self.encode(src);
            vec::map(res) { |i|
                if i == 43u8 { 45u8 }
                else if i == 47u8 { 95u8 }
                else { i }
            }
        }
        fn urlsafe_decode(src: [u8]) -> [u8] {
            let src0 = vec::map(src) { |i|
                if i == 45u8 { 43u8 }
                else if i == 95u8 { 47u8 }
                else { i }
            };
            self.decode(src0)
        }
    }

    let table = vec::init_elt_mut(64u, 0u8), i = 0u8;
    u8::range(65u8, 91u8)  { |j| table[i] = j; i += 1u8; }
    u8::range(97u8, 123u8) { |j| table[i] = j; i += 1u8; }
    u8::range(48u8, 58u8)  { |j| table[i] = j; i += 1u8; }
    table[i] = 43u8; i += 1u8;
    table[i] = 47u8; i += 1u8;

    {padd: 61u8, table: vec::from_mut(table)} as base64
}

fn idx(elems: [const u8], x: u8) -> int {
    let i = 0;
    for e in elems {
        if e == x { ret i; }
        i += 1;
    }
    ret -1;
}

#[cfg(test)]
mod tests {
    import core::str::{from_bytes, bytes};
    enum mode {
        t_decode,
        t_encode,
        t_urlsafe_encode,
        t_urlsafe_decode
    }
    fn setup(t: mode) -> map::hashmap<str, str> {
        let m = map::new_str_hash::<str>();
        alt t {
          t_decode | t_urlsafe_decode {
            m.insert("cGxlYXN1cmUu", "pleasure.");
            m.insert("bGVhc3VyZS4=",  "leasure.");
            m.insert("ZWFzdXJlLg==",   "easure.");
            m.insert("YXN1cmUu",        "asure.");
            m.insert("c3VyZS4=",         "sure.");
          }
          t_encode | t_urlsafe_encode {
            m.insert("pleasure.", "cGxlYXN1cmUu");
            m.insert("leasure.",  "bGVhc3VyZS4=");
            m.insert("easure.",   "ZWFzdXJlLg==");
            m.insert("asure.",        "YXN1cmUu");
            m.insert("sure.",         "c3VyZS4=");
          }
        }
        m
    }
    fn do_test(t: mode) {
        let b64 = mk_base64();
        let m = setup(t);
        m.keys { |k|
            let expected = m.get(k);
            let actual = alt t {
              t_decode { from_bytes(b64.decode(bytes(k))) }
              t_encode { from_bytes(b64.encode(bytes(k))) }
              t_urlsafe_decode { from_bytes(b64.urlsafe_decode(bytes(k))) }
              t_urlsafe_encode { from_bytes(b64.urlsafe_encode(bytes(k))) }
            };
            #debug("expected: %s", expected);
            #debug("actual:   %s", actual);
            assert expected == actual;
        }
    }
    #[test]
    fn decode() { do_test(t_decode); }
    #[test]
    fn encode() { do_test(t_encode); }
    #[test]
    fn urlsafe_decode() { do_test(t_urlsafe_decode); }
    #[test]
    fn urlsafe_encode() { do_test(t_urlsafe_encode); }
}
<commit_msg>changes<commit_after>//
// base64.rs - base64 implementation
//
// The Base64 Alphabet
//
//  Value Encoding  Value Encoding  Value Encoding  Value Encoding
//   0 A            17 R            34 i            51 z
//   1 B            18 S            35 j            52 0
//   2 C            19 T            36 k            53 1
//   3 D            20 U            37 l            54 2
//   4 E            21 V            38 m            55 3
//   5 F            22 W            39 n            56 4
//   6 G            23 X            40 o            57 5
//   7 H            24 Y            41 p            58 6
//   8 I            25 Z            42 q            59 7
//   9 J            26 a            43 r            60 8
//  10 K            27 b            44 s            61 9
//  11 L            28 c            45 t            62 +
//  12 M            29 d            46 u            63 /
//  13 N            30 e            47 v
//  14 O            31 f            48 w         (pad) =
//  15 P            32 g            49 x
//  16 Q            33 h            50 y
//

use std;

import std::map;

export base64, mk_base64;

iface base64 {
    fn encode(src: [u8]) -> [u8];
    fn decode(src: [u8]) -> [u8];
    fn urlsafe_encode(src: [u8]) -> [u8];
    fn urlsafe_decode(src: [u8]) -> [u8];
}

fn mk_base64() -> base64 {
    type _base64 = {padd: u8, table: [u8]};
    impl of base64 for _base64 {
        fn encode(src: [u8]) -> [u8] {
            let srclen = vec::len(src);
            let targ = if srclen % 3u == 0u {
                vec::init_elt_mut((srclen / 3u) * 4u, 0u8)
            } else {
                vec::init_elt_mut((srclen / 3u + 1u) * 4u, 0u8)
            };
            let input = vec::init_elt_mut(3u, 0u8);
            let output = vec::init_elt_mut(4u, 0u8);
            let curr = 0u, src_curr = 0u;
            let table = self.table;

            while srclen > 2u {
                input[0] = src[src_curr];
                input[1] = src[src_curr + 1u];
                input[2] = src[src_curr + 2u];
                srclen -= 3u; src_curr += 3u;

                output[0] = input[0] >> 2u8;
                output[1] = ((input[0] &  3u8) << 4u8) | (input[1] >> 4u8);
                output[2] = ((input[1] & 15u8) << 2u8) | (input[2] >> 6u8);
                output[3] = input[2] & 63u8;

                targ[curr] = table[output[0]]; curr += 1u;
                targ[curr] = table[output[1]]; curr += 1u;
                targ[curr] = table[output[2]]; curr += 1u;
                targ[curr] = table[output[3]]; curr += 1u;
            }

            if srclen != 0u {
                input[0] = 0u8; input[1] = 0u8; input[2] = 0u8;

                alt srclen {
                  1u { input[0] = src[src_curr]; }
                  2u { input[0] = src[src_curr];
                       input[1] = src[src_curr + 1u]; }
                  _  { }
                }

                output[0] = input[0] >> 2u8;
                output[1] = ((input[0] &  3u8) << 4u8) | (input[1] >> 4u8);
                output[2] = ((input[1] & 15u8) << 2u8) | (input[2] >> 6u8);

                targ[curr] = table[output[0]]; curr += 1u;
                targ[curr] = table[output[1]]; curr += 1u;
                targ[curr] = if srclen == 1u {
                    self.padd
                } else {
                    table[output[2]]
                };
                curr += 1u; targ[curr] = self.padd; curr += 1u;
            }

            vec::from_mut(targ)
        }
        fn decode(src: [u8]) -> [u8] {
            let src = src;
            let srclen = vec::len(src);
            let input  = vec::init_elt_mut(4u, 0u8);
            let output = vec::init_elt_mut(4u, 0u8);
            let targ = if src[srclen - 2u] == self.padd {
                vec::init_elt_mut(srclen / 4u * 3u - 2u, 0u8)
            } else if src[srclen - 1u] == self.padd {
                vec::init_elt_mut(srclen / 4u * 3u - 1u, 0u8)
            } else {
                vec::init_elt_mut(srclen / 4u * 3u, 0u8)
            };
            let c = -1;
            let curr = 0u, src_curr = 0u;
            let table = self.table;

            while srclen > 4u {
                input[0] = src[src_curr];
                input[1] = src[src_curr + 1u];
                input[2] = src[src_curr + 2u];
                input[3] = src[src_curr + 3u];
                srclen -= 4u; src_curr += 4u;

                c = idx(table, input[0]);
                if c == -1 { fail "malformed base64 string"; }
                output[0] = c as u8;

                c = idx(table, input[1]);
                if c == -1 { fail "malformed base64 string"; }
                output[1] = c as u8;

                c = idx(table, input[2]);
                if c == -1 { fail "malformed base64 string"; }
                output[2] = c as u8;

                c = idx(table, input[3]);
                if c == -1 { fail "malformed base64 string"; }
                output[3] = c as u8;

                targ[curr] = (output[0] << 2u8) | (output[1] >> 4u8);
                curr += 1u;

                targ[curr] =
                    ((output[1] & 15u8) << 4u8) | (output[2] >> 2u8);
                curr += 1u;

                targ[curr] = ((output[2] &  3u8) << 6u8) | output[3];
                curr += 1u;
            }

            if srclen == 4u {
                input[0] = src[src_curr];
                input[1] = src[src_curr + 1u];
                input[2] = src[src_curr + 2u];
                input[3] = src[src_curr + 3u];

                c = idx(table, input[0]);
                if c == -1 { fail "malformed base64 string"; }
                output[0] = c as u8;

                c = idx(table, input[1]);
                if c == -1 { fail "malformed base64 string"; }
                output[1] = c as u8;

                targ[curr] = (output[0] << 2u8) | (output[1] >> 4u8);
                curr += 1u;

                if input[2] == self.padd { ret vec::from_mut(targ); }

                c = idx(table, input[2]);
                if c == -1 { fail "malformed base64 string"; }
                output[2] = c as u8;

                targ[curr] = ((output[1] & 15u8) << 4u8) | (output[2] >> 2u8);
                curr += 1u;

                if input[3] == self.padd { ret vec::from_mut(targ); }

                c = idx(table, input[3]);
                if c == -1 { fail "malformed base64 string"; }
                output[3] = c as u8;

                targ[curr] = ((output[2] & 3u8) << 6u8) | output[3];
            } else {
                fail "malformed base64 string";
            }

            vec::from_mut(targ)
        }
        fn urlsafe_encode(src: [u8]) -> [u8] {
            let res = self.encode(src);
            vec::map(res) { |i|
                if i == 43u8 { 45u8 }
                else if i == 47u8 { 95u8 }
                else { i }
            }
        }
        fn urlsafe_decode(src: [u8]) -> [u8] {
            let src0 = vec::map(src) { |i|
                if i == 45u8 { 43u8 }
                else if i == 95u8 { 47u8 }
                else { i }
            };
            self.decode(src0)
        }
    }

    let table = vec::init_elt_mut(64u, 0u8), i = 0u8;
    u8::range(65u8, 91u8)  { |j| table[i] = j; i += 1u8; }
    u8::range(97u8, 123u8) { |j| table[i] = j; i += 1u8; }
    u8::range(48u8, 58u8)  { |j| table[i] = j; i += 1u8; }
    table[i] = 43u8; i += 1u8;
    table[i] = 47u8; i += 1u8;

    {padd: 61u8, table: vec::from_mut(table)} as base64
}

fn idx(elems: [const u8], x: u8) -> int {
    let i = 0;
    for e in elems {
        if e == x { ret i; }
        i += 1;
    }
    ret -1;
}

#[cfg(test)]
mod tests {
    import core::str::{from_bytes, bytes};
    enum mode {
        t_decode,
        t_encode,
        t_urlsafe_encode,
        t_urlsafe_decode
    }
    fn setup(t: mode) -> map::hashmap<str, str> {
        let m = map::new_str_hash::<str>();
        alt t {
          t_decode | t_urlsafe_decode {
            m.insert("cGxlYXN1cmUu", "pleasure.");
            m.insert("bGVhc3VyZS4=",  "leasure.");
            m.insert("ZWFzdXJlLg==",   "easure.");
            m.insert("YXN1cmUu",        "asure.");
            m.insert("c3VyZS4=",         "sure.");
          }
          t_encode | t_urlsafe_encode {
            m.insert("pleasure.", "cGxlYXN1cmUu");
            m.insert("leasure.",  "bGVhc3VyZS4=");
            m.insert("easure.",   "ZWFzdXJlLg==");
            m.insert("asure.",        "YXN1cmUu");
            m.insert("sure.",         "c3VyZS4=");
          }
        }
        m
    }
    fn do_test(t: mode) {
        let b64 = mk_base64();
        let m = setup(t);
        m.keys { |k|
            let expected = m.get(k);
            let actual = alt t {
              t_decode { from_bytes(b64.decode(bytes(k))) }
              t_encode { from_bytes(b64.encode(bytes(k))) }
              t_urlsafe_decode { from_bytes(b64.urlsafe_decode(bytes(k))) }
              t_urlsafe_encode { from_bytes(b64.urlsafe_encode(bytes(k))) }
            };
            #debug("expected: %s", expected);
            #debug("actual:   %s", actual);
            assert expected == actual;
        }
    }
    #[test]
    fn decode() { do_test(t_decode); }
    #[test]
    fn encode() { do_test(t_encode); }
    #[test]
    fn urlsafe_decode() { do_test(t_urlsafe_decode); }
    #[test]
    fn urlsafe_encode() { do_test(t_urlsafe_encode); }
}
<|endoftext|>