<commit_before>use proc_macro2::*;
use quote::ToTokens;
use syn::parse::{self, Parse, ParseStream};
use syn::punctuated::Punctuated;

use meta::*;
use util::*;

// Extremely curious why this triggers on a nearly branchless function
#[allow(clippy::cognitive_complexity)]
// for loop comes from `quote!`
#[allow(clippy::for_loop_over_option)]
// https://github.com/rust-lang/rust-clippy/issues/3768
#[allow(clippy::useless_let_if_seq)]
pub(crate) fn expand(input: SqlFunctionDecl) -> Result<TokenStream, Diagnostic> {
    let SqlFunctionDecl {
        mut attributes,
        fn_token,
        fn_name,
        mut generics,
        args,
        return_type,
    } = input;

    let sql_name = MetaItem::with_name(&attributes, "sql_name")
        .map(|m| m.str_value())
        .unwrap_or_else(|| Ok(fn_name.to_string()))?;
    let is_aggregate = MetaItem::with_name(&attributes, "aggregate").is_some();

    attributes.retain(|attr| {
        attr.parse_meta()
            .map(|m| !m.path().is_ident("sql_name") && !m.path().is_ident("aggregate"))
            .unwrap_or(true)
    });

    let args = &args;
    let (ref arg_name, ref arg_type): (Vec<_>, Vec<_>) = args
        .iter()
        .map(|StrictFnArg { name, ty, .. }| (name, ty))
        .unzip();
    let arg_struct_assign = args.iter().map(
        |StrictFnArg {
             name, colon_token, ..
         }| {
            let name2 = name.clone();
            quote!(#name #colon_token #name2.as_expression())
        },
    );

    let type_args = &generics
        .type_params()
        .map(|type_param| type_param.ident.clone())
        .collect::<Vec<_>>();
    let type_args2 = &type_args.clone();
    for StrictFnArg { name, .. } in args {
        generics.params.push(parse_quote!(#name));
    }

    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
    // Even if we force an empty where clause, it still won't print the where
    // token with no bounds.
    let where_clause = where_clause
        .map(|w| quote!(#w))
        .unwrap_or_else(|| quote!(where));

    let mut generics_with_internal = generics.clone();
    generics_with_internal
        .params
        .push(parse_quote!(__DieselInternal));
    let (impl_generics_internal, _, _) = generics_with_internal.split_for_impl();

    let sql_type;
    let numeric_derive;

    if arg_name.is_empty() {
        sql_type = None;
        // FIXME: We can always derive once trivial bounds are stable
        numeric_derive = None;
    } else {
        sql_type = Some(quote!((#(#arg_name),*)));
        numeric_derive = Some(quote!(#[derive(diesel::sql_types::DieselNumericOps)]));
    }

    let args_iter = args.iter();
    let mut tokens = quote! {
        use diesel::{self, QueryResult};
        use diesel::expression::{AsExpression, Expression, SelectableExpression, AppearsOnTable};
        use diesel::query_builder::{QueryFragment, AstPass};
        use diesel::sql_types::*;
        use super::*;

        #[derive(Debug, Clone, Copy, diesel::query_builder::QueryId)]
        #numeric_derive
        pub struct #fn_name #ty_generics {
            #(pub(in super) #args_iter,)*
            #(pub(in super) #type_args: ::std::marker::PhantomData<#type_args>,)*
        }

        pub type HelperType #ty_generics = #fn_name <
            #(#type_args,)*
            #(<#arg_name as AsExpression<#arg_type>>::Expression,)*
        >;

        impl #impl_generics Expression for #fn_name #ty_generics
        #where_clause
            #(#sql_type: Expression,)*
        {
            type SqlType = #return_type;
        }

        // __DieselInternal is what we call QS normally
        impl #impl_generics_internal SelectableExpression<__DieselInternal>
            for #fn_name #ty_generics
        #where_clause
            #(#arg_name: SelectableExpression<__DieselInternal>,)*
            Self: AppearsOnTable<__DieselInternal>,
        {
        }

        // __DieselInternal is what we call QS normally
        impl #impl_generics_internal AppearsOnTable<__DieselInternal>
            for #fn_name #ty_generics
        #where_clause
            #(#arg_name: AppearsOnTable<__DieselInternal>,)*
            Self: Expression,
        {
        }

        // __DieselInternal is what we call DB normally
        impl #impl_generics_internal QueryFragment<__DieselInternal>
            for #fn_name #ty_generics
        where
            __DieselInternal: diesel::backend::Backend,
            for<'a> (#(&'a #arg_name),*): QueryFragment<__DieselInternal>,
        {
            fn walk_ast(&self, mut out: AstPass<__DieselInternal>) -> QueryResult<()> {
                out.push_sql(concat!(#sql_name, "("));
                (#(&self.#arg_name,)*).walk_ast(out.reborrow())?;
                out.push_sql(")");
                Ok(())
            }
        }
    };

    if is_aggregate {
        tokens = quote! {
            #tokens

            impl #impl_generics_internal ValidGrouping<__DieselInternal>
                for #fn_name #ty_generics
            #where_clause
                #(#arg_name: diesel::expression::NonAggregate,)*
                Self: Expression,
            {
                type IsAggregate = diesel::expression::is_aggregate::Yes;
            }
        };
        if cfg!(feature = "sqlite") && type_args.is_empty() {
            tokens = quote! {
                #tokens

                use diesel::sqlite::{Sqlite, SqliteConnection};
                use diesel::serialize::ToSql;
                use diesel::deserialize::{FromSqlRow, StaticallySizedRow};
                use diesel::sqlite::SqliteAggregateFunction;
                use diesel::sql_types::IntoNullable;
            };

            match arg_name.len() {
                x if x > 1 => {
                    tokens = quote! {
                        #tokens

                        #[allow(dead_code)]
                        /// Registers an implementation for this aggregate function on the given connection
                        ///
                        /// This function must be called for every `SqliteConnection` before
                        /// this SQL function can be used on SQLite. The implementation must be
                        /// deterministic (returns the same result given the same arguments).
                        pub fn register_impl<A, #(#arg_name,)*>(
                            conn: &SqliteConnection
                        ) -> QueryResult<()>
                            where
                            A: SqliteAggregateFunction<(#(#arg_name,)*)> + Send + 'static,
                            A::Output: ToSql<#return_type, Sqlite>,
                            (#(#arg_name,)*): FromSqlRow<(#(#arg_type,)*), Sqlite> +
                                StaticallySizedRow<(#(#arg_type,)*), Sqlite>,
                        {
                            conn.register_aggregate_function::<(#(#arg_type,)*), #return_type, _, _, A>(#sql_name)
                        }
                    };
                }
                x if x == 1 => {
                    let arg_name = arg_name[0];
                    let arg_type = arg_type[0];

                    tokens = quote! {
                        #tokens

                        #[allow(dead_code)]
                        /// Registers an implementation for this aggregate function on the given connection
                        ///
                        /// This function must be called for every `SqliteConnection` before
                        /// this SQL function can be used on SQLite. The implementation must be
                        /// deterministic (returns the same result given the same arguments).
                        pub fn register_impl<A, #arg_name>(
                            conn: &SqliteConnection
                        ) -> QueryResult<()>
                            where
                            A: SqliteAggregateFunction<#arg_name> + Send + 'static,
                            A::Output: ToSql<#return_type, Sqlite>,
                            #arg_name: FromSqlRow<#arg_type, Sqlite> +
                                StaticallySizedRow<#arg_type, Sqlite>,
                            {
                                conn.register_aggregate_function::<#arg_type, #return_type, _, _, A>(#sql_name)
                            }
                    };
                }
                _ => (),
            }
        }
    } else {
        tokens = quote! {
            #tokens

            #[derive(ValidGrouping)]
            pub struct __Derived<#(#arg_name,)*>(#(#arg_name,)*);

            impl #impl_generics_internal ValidGrouping<__DieselInternal>
                for #fn_name #ty_generics
            where
                __Derived<#(#arg_name,)*>: ValidGrouping<__DieselInternal>,
            {
                type IsAggregate = <__Derived<#(#arg_name,)*> as ValidGrouping<__DieselInternal>>::IsAggregate;
            }
        };

        if cfg!(feature = "sqlite") && type_args.is_empty() {
            tokens = quote! {
                #tokens

                use diesel::sqlite::{Sqlite, SqliteConnection};
                use diesel::serialize::ToSql;
                use diesel::deserialize::{FromSqlRow, StaticallySizedRow};

                #[allow(dead_code)]
                /// Registers an implementation for this function on the given connection
                ///
                /// This function must be called for every `SqliteConnection` before
                /// this SQL function can be used on SQLite. The implementation must be
                /// deterministic (returns the same result given the same arguments). If
                /// the function is nondeterministic, call
                /// `register_nondeterministic_impl` instead.
                pub fn register_impl<F, Ret, #(#arg_name,)*>(
                    conn: &SqliteConnection,
                    f: F,
                ) -> QueryResult<()>
                where
                    F: Fn(#(#arg_name,)*) -> Ret + Send + 'static,
                    (#(#arg_name,)*): FromSqlRow<(#(#arg_type,)*), Sqlite> +
                        StaticallySizedRow<(#(#arg_type,)*), Sqlite>,
                    Ret: ToSql<#return_type, Sqlite>,
                {
                    conn.register_sql_function::<(#(#arg_type,)*), #return_type, _, _, _>(
                        #sql_name,
                        true,
                        move |(#(#arg_name,)*)| f(#(#arg_name,)*),
                    )
                }

                #[allow(dead_code)]
                /// Registers an implementation for this function on the given connection
                ///
                /// This function must be called for every `SqliteConnection` before
                /// this SQL function can be used on SQLite.
                /// `register_nondeterministic_impl` should only be used if your
                /// function can return different results with the same arguments (e.g.
                /// `random`). If your function is deterministic, you should call
                /// `register_impl` instead.
                pub fn register_nondeterministic_impl<F, Ret, #(#arg_name,)*>(
                    conn: &SqliteConnection,
                    mut f: F,
                ) -> QueryResult<()>
                where
                    F: FnMut(#(#arg_name,)*) -> Ret + Send + 'static,
                    (#(#arg_name,)*): FromSqlRow<(#(#arg_type,)*), Sqlite> +
                        StaticallySizedRow<(#(#arg_type,)*), Sqlite>,
                    Ret: ToSql<#return_type, Sqlite>,
                {
                    conn.register_sql_function::<(#(#arg_type,)*), #return_type, _, _, _>(
                        #sql_name,
                        false,
                        move |(#(#arg_name,)*)| f(#(#arg_name,)*),
                    )
                }
            };
        }
    }

    let args_iter = args.iter();
    tokens = quote! {
        #(#attributes)*
        #[allow(non_camel_case_types)]
        pub #fn_token #fn_name #impl_generics (#(#args_iter,)*)
            -> #fn_name::HelperType #ty_generics
        #where_clause
            #(#arg_name: ::diesel::expression::AsExpression<#arg_type>,)*
        {
            #fn_name::#fn_name {
                #(#arg_struct_assign,)*
                #(#type_args: ::std::marker::PhantomData,)*
            }
        }

        #[doc(hidden)]
        #[allow(non_camel_case_types, non_snake_case, unused_imports)]
        pub(crate) mod #fn_name {
            #tokens
        }
    };

    Ok(tokens)
}

pub(crate) struct SqlFunctionDecl {
    attributes: Vec<syn::Attribute>,
    fn_token: Token![fn],
    fn_name: syn::Ident,
    generics: syn::Generics,
    args: Punctuated<StrictFnArg, Token![,]>,
    return_type: syn::Type,
}

impl Parse for SqlFunctionDecl {
    fn parse(input: ParseStream) -> parse::Result<Self> {
        let attributes = syn::Attribute::parse_outer(input)?;
        let fn_token: Token![fn] = input.parse()?;
        let fn_name = syn::Ident::parse(input)?;
        let generics = syn::Generics::parse(input)?;
        let args;
        let _paren = parenthesized!(args in input);
        let args = args.parse_terminated::<_, Token![,]>(StrictFnArg::parse)?;
        let return_type = if Option::<Token![->]>::parse(input)?.is_some() {
            syn::Type::parse(input)?
        } else {
            parse_quote!(diesel::expression::expression_types::NotSelectable)
        };
        let _semi = Option::<Token![;]>::parse(input)?;

        Ok(Self {
            attributes,
            fn_token,
            fn_name,
            generics,
            args,
            return_type,
        })
    }
}

/// Essentially the same as syn::ArgCaptured, but only allowing ident patterns
struct StrictFnArg {
    name: syn::Ident,
    colon_token: Token![:],
    ty: syn::Type,
}

impl Parse for StrictFnArg {
    fn parse(input: ParseStream) -> parse::Result<Self> {
        let name = input.parse()?;
        let colon_token = input.parse()?;
        let ty = input.parse()?;
        Ok(Self {
            name,
            colon_token,
            ty,
        })
    }
}

impl ToTokens for StrictFnArg {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        self.name.to_tokens(tokens);
        self.colon_token.to_tokens(tokens);
        self.name.to_tokens(tokens);
    }
}
<commit_msg>Post rebase fixes<commit_after>use proc_macro2::*;
use quote::ToTokens;
use syn::parse::{self, Parse, ParseStream};
use syn::punctuated::Punctuated;

use meta::*;
use util::*;

// Extremely curious why this triggers on a nearly branchless function
#[allow(clippy::cognitive_complexity)]
// for loop comes from `quote!`
#[allow(clippy::for_loop_over_option)]
// https://github.com/rust-lang/rust-clippy/issues/3768
#[allow(clippy::useless_let_if_seq)]
pub(crate) fn expand(input: SqlFunctionDecl) -> Result<TokenStream, Diagnostic> {
    let SqlFunctionDecl {
        mut attributes,
        fn_token,
        fn_name,
        mut generics,
        args,
        return_type,
    } = input;

    let sql_name = MetaItem::with_name(&attributes, "sql_name")
        .map(|m| m.str_value())
        .unwrap_or_else(|| Ok(fn_name.to_string()))?;
    let is_aggregate = MetaItem::with_name(&attributes, "aggregate").is_some();

    attributes.retain(|attr| {
        attr.parse_meta()
            .map(|m| !m.path().is_ident("sql_name") && !m.path().is_ident("aggregate"))
            .unwrap_or(true)
    });

    let args = &args;
    let (ref arg_name, ref arg_type): (Vec<_>, Vec<_>) = args
        .iter()
        .map(|StrictFnArg { name, ty, .. }| (name, ty))
        .unzip();
    let arg_struct_assign = args.iter().map(
        |StrictFnArg {
             name, colon_token, ..
         }| {
            let name2 = name.clone();
            quote!(#name #colon_token #name2.as_expression())
        },
    );

    let type_args = &generics
        .type_params()
        .map(|type_param| type_param.ident.clone())
        .collect::<Vec<_>>();

    for StrictFnArg { name, .. } in args {
        generics.params.push(parse_quote!(#name));
    }

    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
    // Even if we force an empty where clause, it still won't print the where
    // token with no bounds.
    let where_clause = where_clause
        .map(|w| quote!(#w))
        .unwrap_or_else(|| quote!(where));

    let mut generics_with_internal = generics.clone();
    generics_with_internal
        .params
        .push(parse_quote!(__DieselInternal));
    let (impl_generics_internal, _, _) = generics_with_internal.split_for_impl();

    let sql_type;
    let numeric_derive;

    if arg_name.is_empty() {
        sql_type = None;
        // FIXME: We can always derive once trivial bounds are stable
        numeric_derive = None;
    } else {
        sql_type = Some(quote!((#(#arg_name),*): Expression,));
        numeric_derive = Some(quote!(#[derive(diesel::sql_types::DieselNumericOps)]));
    }

    let args_iter = args.iter();
    let mut tokens = quote! {
        use diesel::{self, QueryResult};
        use diesel::expression::{AsExpression, Expression, SelectableExpression, AppearsOnTable, ValidGrouping};
        use diesel::query_builder::{QueryFragment, AstPass};
        use diesel::sql_types::*;
        use super::*;

        #[derive(Debug, Clone, Copy, diesel::query_builder::QueryId)]
        #numeric_derive
        pub struct #fn_name #ty_generics {
            #(pub(in super) #args_iter,)*
            #(pub(in super) #type_args: ::std::marker::PhantomData<#type_args>,)*
        }

        pub type HelperType #ty_generics = #fn_name <
            #(#type_args,)*
            #(<#arg_name as AsExpression<#arg_type>>::Expression,)*
        >;

        impl #impl_generics Expression for #fn_name #ty_generics
        #where_clause
            #sql_type
        {
            type SqlType = #return_type;
        }

        // __DieselInternal is what we call QS normally
        impl #impl_generics_internal SelectableExpression<__DieselInternal>
            for #fn_name #ty_generics
        #where_clause
            #(#arg_name: SelectableExpression<__DieselInternal>,)*
            Self: AppearsOnTable<__DieselInternal>,
        {
        }

        // __DieselInternal is what we call QS normally
        impl #impl_generics_internal AppearsOnTable<__DieselInternal>
            for #fn_name #ty_generics
        #where_clause
            #(#arg_name: AppearsOnTable<__DieselInternal>,)*
            Self: Expression,
        {
        }

        // __DieselInternal is what we call DB normally
        impl #impl_generics_internal QueryFragment<__DieselInternal>
            for #fn_name #ty_generics
        where
            __DieselInternal: diesel::backend::Backend,
            for<'a> (#(&'a #arg_name),*): QueryFragment<__DieselInternal>,
        {
            fn walk_ast(&self, mut out: AstPass<__DieselInternal>) -> QueryResult<()> {
                out.push_sql(concat!(#sql_name, "("));
                (#(&self.#arg_name,)*).walk_ast(out.reborrow())?;
                out.push_sql(")");
                Ok(())
            }
        }
    };

    if is_aggregate {
        tokens = quote! {
            #tokens

            impl #impl_generics_internal ValidGrouping<__DieselInternal>
                for #fn_name #ty_generics
            {
                type IsAggregate = diesel::expression::is_aggregate::Yes;
            }
        };
        if cfg!(feature = "sqlite") && type_args.is_empty() {
            tokens = quote! {
                #tokens

                use diesel::sqlite::{Sqlite, SqliteConnection};
                use diesel::serialize::ToSql;
                use diesel::deserialize::{FromSqlRow, StaticallySizedRow};
                use diesel::sqlite::SqliteAggregateFunction;
                use diesel::sql_types::IntoNullable;
            };

            match arg_name.len() {
                x if x > 1 => {
                    tokens = quote! {
                        #tokens

                        #[allow(dead_code)]
                        /// Registers an implementation for this aggregate function on the given connection
                        ///
                        /// This function must be called for every `SqliteConnection` before
                        /// this SQL function can be used on SQLite. The implementation must be
                        /// deterministic (returns the same result given the same arguments).
                        pub fn register_impl<A, #(#arg_name,)*>(
                            conn: &SqliteConnection
                        ) -> QueryResult<()>
                            where
                            A: SqliteAggregateFunction<(#(#arg_name,)*)> + Send + 'static,
                            A::Output: ToSql<#return_type, Sqlite>,
                            (#(#arg_name,)*): FromSqlRow<(#(#arg_type,)*), Sqlite> +
                                StaticallySizedRow<(#(#arg_type,)*), Sqlite>,
                        {
                            conn.register_aggregate_function::<(#(#arg_type,)*), #return_type, _, _, A>(#sql_name)
                        }
                    };
                }
                x if x == 1 => {
                    let arg_name = arg_name[0];
                    let arg_type = arg_type[0];

                    tokens = quote! {
                        #tokens

                        #[allow(dead_code)]
                        /// Registers an implementation for this aggregate function on the given connection
                        ///
                        /// This function must be called for every `SqliteConnection` before
                        /// this SQL function can be used on SQLite. The implementation must be
                        /// deterministic (returns the same result given the same arguments).
                        pub fn register_impl<A, #arg_name>(
                            conn: &SqliteConnection
                        ) -> QueryResult<()>
                            where
                            A: SqliteAggregateFunction<#arg_name> + Send + 'static,
                            A::Output: ToSql<#return_type, Sqlite>,
                            #arg_name: FromSqlRow<#arg_type, Sqlite> +
                                StaticallySizedRow<#arg_type, Sqlite>,
                            {
                                conn.register_aggregate_function::<#arg_type, #return_type, _, _, A>(#sql_name)
                            }
                    };
                }
                _ => (),
            }
        }
    } else {
        tokens = quote! {
            #tokens

            #[derive(ValidGrouping)]
            pub struct __Derived<#(#arg_name,)*>(#(#arg_name,)*);

            impl #impl_generics_internal ValidGrouping<__DieselInternal>
                for #fn_name #ty_generics
            where
                __Derived<#(#arg_name,)*>: ValidGrouping<__DieselInternal>,
            {
                type IsAggregate = <__Derived<#(#arg_name,)*> as ValidGrouping<__DieselInternal>>::IsAggregate;
            }
        };

        if cfg!(feature = "sqlite") && type_args.is_empty() && !arg_name.is_empty() {
            tokens = quote! {
                #tokens

                use diesel::sqlite::{Sqlite, SqliteConnection};
                use diesel::serialize::ToSql;
                use diesel::deserialize::{FromSqlRow, StaticallySizedRow};

                #[allow(dead_code)]
                /// Registers an implementation for this function on the given connection
                ///
                /// This function must be called for every `SqliteConnection` before
                /// this SQL function can be used on SQLite. The implementation must be
                /// deterministic (returns the same result given the same arguments). If
                /// the function is nondeterministic, call
                /// `register_nondeterministic_impl` instead.
                pub fn register_impl<F, Ret, #(#arg_name,)*>(
                    conn: &SqliteConnection,
                    f: F,
                ) -> QueryResult<()>
                where
                    F: Fn(#(#arg_name,)*) -> Ret + Send + 'static,
                    (#(#arg_name,)*): FromSqlRow<(#(#arg_type,)*), Sqlite> +
                        StaticallySizedRow<(#(#arg_type,)*), Sqlite>,
                    Ret: ToSql<#return_type, Sqlite>,
                {
                    conn.register_sql_function::<(#(#arg_type,)*), #return_type, _, _, _>(
                        #sql_name,
                        true,
                        move |(#(#arg_name,)*)| f(#(#arg_name,)*),
                    )
                }

                #[allow(dead_code)]
                /// Registers an implementation for this function on the given connection
                ///
                /// This function must be called for every `SqliteConnection` before
                /// this SQL function can be used on SQLite.
                /// `register_nondeterministic_impl` should only be used if your
                /// function can return different results with the same arguments (e.g.
                /// `random`). If your function is deterministic, you should call
                /// `register_impl` instead.
                pub fn register_nondeterministic_impl<F, Ret, #(#arg_name,)*>(
                    conn: &SqliteConnection,
                    mut f: F,
                ) -> QueryResult<()>
                where
                    F: FnMut(#(#arg_name,)*) -> Ret + Send + 'static,
                    (#(#arg_name,)*): FromSqlRow<(#(#arg_type,)*), Sqlite> +
                        StaticallySizedRow<(#(#arg_type,)*), Sqlite>,
                    Ret: ToSql<#return_type, Sqlite>,
                {
                    conn.register_sql_function::<(#(#arg_type,)*), #return_type, _, _, _>(
                        #sql_name,
                        false,
                        move |(#(#arg_name,)*)| f(#(#arg_name,)*),
                    )
                }
            };
        }
    }

    let args_iter = args.iter();
    tokens = quote! {
        #(#attributes)*
        #[allow(non_camel_case_types)]
        pub #fn_token #fn_name #impl_generics (#(#args_iter,)*)
            -> #fn_name::HelperType #ty_generics
        #where_clause
            #(#arg_name: ::diesel::expression::AsExpression<#arg_type>,)*
        {
            #fn_name::#fn_name {
                #(#arg_struct_assign,)*
                #(#type_args: ::std::marker::PhantomData,)*
            }
        }

        #[doc(hidden)]
        #[allow(non_camel_case_types, non_snake_case, unused_imports)]
        pub(crate) mod #fn_name {
            #tokens
        }
    };

    Ok(tokens)
}

pub(crate) struct SqlFunctionDecl {
    attributes: Vec<syn::Attribute>,
    fn_token: Token![fn],
    fn_name: syn::Ident,
    generics: syn::Generics,
    args: Punctuated<StrictFnArg, Token![,]>,
    return_type: syn::Type,
}

impl Parse for SqlFunctionDecl {
    fn parse(input: ParseStream) -> parse::Result<Self> {
        let attributes = syn::Attribute::parse_outer(input)?;
        let fn_token: Token![fn] = input.parse()?;
        let fn_name = syn::Ident::parse(input)?;
        let generics = syn::Generics::parse(input)?;
        let args;
        let _paren = parenthesized!(args in input);
        let args = args.parse_terminated::<_, Token![,]>(StrictFnArg::parse)?;
        let return_type = if Option::<Token![->]>::parse(input)?.is_some() {
            syn::Type::parse(input)?
        } else {
            parse_quote!(diesel::expression::expression_types::NotSelectable)
        };
        let _semi = Option::<Token![;]>::parse(input)?;

        Ok(Self {
            attributes,
            fn_token,
            fn_name,
            generics,
            args,
            return_type,
        })
    }
}

/// Essentially the same as syn::ArgCaptured, but only allowing ident patterns
struct StrictFnArg {
    name: syn::Ident,
    colon_token: Token![:],
    ty: syn::Type,
}

impl Parse for StrictFnArg {
    fn parse(input: ParseStream) -> parse::Result<Self> {
        let name = input.parse()?;
        let colon_token = input.parse()?;
        let ty = input.parse()?;
        Ok(Self {
            name,
            colon_token,
            ty,
        })
    }
}

impl ToTokens for StrictFnArg {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        self.name.to_tokens(tokens);
        self.colon_token.to_tokens(tokens);
        self.name.to_tokens(tokens);
    }
}
<|endoftext|>


<commit_before>use std::cmp;
use AttackResult::*;
use EffectClass::*;

#[derive(Debug, Copy, Clone)]
struct Character {
    hp: i32,
    mana: i32,
    damage: i32,
    armor: i32
}

#[derive(Debug, Clone)]
enum AttackResult {
    Win,
    Lose,
    Continue(State)
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum EffectClass {
    Shield,
    Poison,
    Recharge
}

#[derive(Debug, Copy, Clone)]
struct Effect {
    turns: i32,
    intensity: i32,
    class: EffectClass
}

#[derive(Debug, Copy, Clone)]
struct Spell {
    cost: i32,
    damage: i32,
    healing: i32,
    effect: Option<Effect>
}

#[derive(Debug, Clone)]
struct State {
    player: Character,
    boss: Character,
    effects: Vec<Effect>
}

fn process_effects(state: &mut State) {
    let mut effects = Vec::new();
    let mut armor = 0;

    for mut effect in state.effects.iter().filter(|x| x.turns > 0).cloned() {
        effect.turns -= 1;

        if effect.turns > 0 {
            effects.push(effect);
        }

        match effect.class {
            Shield => armor = effect.intensity,
            Poison => state.boss.hp -= effect.intensity,
            Recharge => state.player.mana += effect.intensity
        }
    }

    state.player.armor = armor;
    state.effects = effects;
}

fn cast_spell(state: &State, spell: Spell) -> AttackResult {
    let mut result = state.clone();
    process_effects(&mut result);

    // Cast spell

    result.player.mana -= spell.cost;
    result.player.hp += spell.healing;
    result.boss.hp -= spell.damage;

    if result.player.mana <= 0 {
        return Lose;
    }

    // Deal with effects

    if let Some(effect) = spell.effect {
        if result.effects.iter().all(|x| x.class != effect.class) {
            result.effects.push(effect);
        } else {
            return Lose;
        }
    }

    // Counter attack

    process_effects(&mut result);

    if result.boss.hp > 0 {
        let damage = cmp::max(result.boss.damage - result.player.armor, 1);
        result.player.hp = cmp::max(result.player.hp - damage, 0);

        if result.player.hp <= 0 {
            return Lose;
        }
    } else {
        return Win;
    }

    Continue(result)
}

fn cheapest_strategy(spells: &Vec<Spell>, state: &State) -> Option<Vec<Spell>> {
    spells.iter().filter_map(|&spell| match cast_spell(state, spell) {
        Continue(next_state) => {
            Some(match cheapest_strategy(spells, &next_state) {
                Some(mut strategy) => {
                    strategy.push(spell);
                    strategy
                },
                None => return None
            })
        },
        Win => Some(vec![spell]),
        Lose => None
    }).min_by_key(|strategy| {
        strategy.into_iter().map(|spell| spell.cost).sum::<i32>()
    })
}

fn main() {
    let spells = vec![
        Spell {
            cost: 53,
            damage: 4,
            healing: 0,
            effect: None
        },
        Spell {
            cost: 73,
            damage: 2,
            healing: 2,
            effect: None
        },
        Spell {
            cost: 113,
            damage: 0,
            healing: 0,
            effect: Some(Effect {turns: 6, intensity: 7, class: Shield})
        },
        Spell {
            cost: 173,
            damage: 0,
            healing: 0,
            effect: Some(Effect {turns: 6, intensity: 3, class: Poison})
        },
        Spell {
            cost: 229,
            damage: 0,
            healing: 0,
            effect: Some(Effect {turns: 5, intensity: 101, class: Recharge})
        }
    ];

    let boss = Character {
        hp: 51,
        mana: 0,
        damage: 9,
        armor: 0
    };

    let player = Character {
        hp: 50,
        mana: 500,
        damage: 0,
        armor: 0
    };

    let start = State {player, boss, effects: Vec::new()};
    let mut strategy = cheapest_strategy(&spells, &start).unwrap();

    strategy.reverse();

    println!("{:#?}", strategy);
}
<commit_msg>BFS implementation<commit_after>use std::cmp;
use std::collections::VecDeque;
use AttackResult::*;
use EffectClass::*;

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
struct Character {
    hp: i32,
    mana: i32,
    damage: i32,
    armor: i32
}

#[derive(Debug, Clone)]
enum AttackResult {
    Win,
    Lose,
    Continue(State)
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum EffectClass {
    Shield,
    Poison,
    Recharge
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
struct Effect {
    turns: i32,
    intensity: i32,
    class: EffectClass
}

#[derive(Debug, Copy, Clone)]
struct Spell {
    cost: i32,
    damage: i32,
    healing: i32,
    effect: Option<Effect>
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct State {
    player: Character,
    boss: Character,
    effects: Vec<Effect>
}

fn process_effects(state: &mut State) {
    let mut effects = Vec::new();
    let mut armor = 0;

    for mut effect in state.effects.iter().filter(|x| x.turns > 0).cloned() {
        effect.turns -= 1;

        if effect.turns > 0 {
            effects.push(effect);
        }

        match effect.class {
            Shield => armor = effect.intensity,
            Poison => state.boss.hp -= effect.intensity,
            Recharge => state.player.mana += effect.intensity
        }
    }

    state.player.armor = armor;
    state.effects = effects;
}

fn cast_spell(state: &State, spell: Spell) -> AttackResult {
    let mut result = state.clone();
    process_effects(&mut result);

    // Cast spell

    result.player.mana -= spell.cost;
    result.player.hp += spell.healing;
    result.boss.hp -= spell.damage;

    if result.player.mana <= 0 {
        return Lose;
    }

    // Deal with effects

    if let Some(effect) = spell.effect {
        if result.effects.iter().all(|x| x.class != effect.class) {
            result.effects.push(effect);
        } else {
            return Lose;
        }
    }

    // Counter attack

    process_effects(&mut result);

    if result.boss.hp > 0 {
        let damage = cmp::max(result.boss.damage - result.player.armor, 1);
        result.player.hp = cmp::max(result.player.hp - damage, 0);

        if result.player.hp <= 0 {
            return Lose;
        }
    } else {
        return Win;
    }

    Continue(result)
}

fn cheapest_strategy(spells: &Vec<Spell>, start: &State) -> Option<Vec<Spell>> {
    let mut queue: VecDeque<(Vec<Spell>, State)> = VecDeque::new();
    let mut min = None;

    queue.push_back((vec![], start.clone()));

    while queue.len() > 0 {
        let (substrategy, state) = queue.pop_front().unwrap();
        let cost = substrategy.iter().map(|spell| spell.cost).sum::<i32>();

        for &spell in spells {
            match cast_spell(&state, spell) {
                Lose => continue,
                attack_result => {
                    if let Some((min_cost, _)) = min {
                        if cost + spell.cost >= min_cost {
                            continue;
                        }
                    }

                    let mut strategy = substrategy.clone();
                    strategy.push(spell);

                    if let Continue(next_state) = attack_result {
                        queue.push_back((strategy, next_state));
                    } else {
                        min = Some((cost + spell.cost, strategy));
                    }
                }
            }
        }
    }

    match min {
        None => None,
        Some((_, mut strategy)) => {
            strategy.reverse();
            Some(strategy)
        }
    }
}

fn main() {
    let spells = vec![
        Spell {
            cost: 53,
            damage: 4,
            healing: 0,
            effect: None
        },
        Spell {
            cost: 73,
            damage: 2,
            healing: 2,
            effect: None
        },
        Spell {
            cost: 113,
            damage: 0,
            healing: 0,
            effect: Some(Effect {turns: 6, intensity: 7, class: Shield})
        },
        Spell {
            cost: 173,
            damage: 0,
            healing: 0,
            effect: Some(Effect {turns: 6, intensity: 3, class: Poison})
        },
        Spell {
            cost: 229,
            damage: 0,
            healing: 0,
            effect: Some(Effect {turns: 5, intensity: 101, class: Recharge})
        }
    ];

    let boss = Character {
        hp: 51,
        mana: 0,
        damage: 9,
        armor: 0
    };

    let player = Character {
        hp: 50,
        mana: 500,
        damage: 0,
        armor: 0
    };

    let start = State {player, boss, effects: Vec::new()};
    let strategy = cheapest_strategy(&spells, &start).unwrap();
    let cost = strategy.iter().map(|spell| spell.cost).sum::<i32>();

    println!("Part 1: {}", cost);
}
<|endoftext|>


<commit_before>use crate::internal::consts::{self, MAX_REGULAR_STREAM_ID, NO_STREAM};
use crate::internal::{self, Version};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::io::{self, Read, Write};
use uuid::Uuid;

// ========================================================================= //

macro_rules! malformed {
    ($e:expr) => { invalid_data!("Malformed directory entry ({})", $e) };
    ($fmt:expr, $($arg:tt)+) => {
        invalid_data!("Malformed directory entry ({})",
                      format!($fmt, $($arg)+))
    };
}

// ========================================================================= //

#[derive(Clone)]
pub struct DirEntry {
    pub name: String,
    pub obj_type: u8,
    pub color: u8,
    pub left_sibling: u32,
    pub right_sibling: u32,
    pub child: u32,
    pub clsid: Uuid,
    pub state_bits: u32,
    pub creation_time: u64,
    pub modified_time: u64,
    pub start_sector: u32,
    pub stream_len: u64,
}

impl DirEntry {
    pub fn unallocated() -> DirEntry {
        DirEntry {
            name: String::new(),
            obj_type: consts::OBJ_TYPE_UNALLOCATED,
            color: 0,
            left_sibling: NO_STREAM,
            right_sibling: NO_STREAM,
            child: NO_STREAM,
            clsid: Uuid::nil(),
            state_bits: 0,
            creation_time: 0,
            modified_time: 0,
            start_sector: 0,
            stream_len: 0,
        }
    }

    pub fn read_clsid<R: Read>(reader: &mut R) -> io::Result<Uuid> {
        let d1 = reader.read_u32::<LittleEndian>()?;
        let d2 = reader.read_u16::<LittleEndian>()?;
        let d3 = reader.read_u16::<LittleEndian>()?;
        let mut d4 = [0u8; 8];
        reader.read_exact(&mut d4)?;
        Ok(Uuid::from_fields(d1, d2, d3, &d4).unwrap())
    }

    pub fn write_clsid<W: Write>(
        writer: &mut W,
        clsid: &Uuid,
    ) -> io::Result<()> {
        let (d1, d2, d3, d4) = clsid.as_fields();
        writer.write_u32::<LittleEndian>(d1)?;
        writer.write_u16::<LittleEndian>(d2)?;
        writer.write_u16::<LittleEndian>(d3)?;
        writer.write_all(d4)?;
        Ok(())
    }

    pub fn read_from<R: Read>(
        reader: &mut R,
        version: Version,
    ) -> io::Result<DirEntry> {
        let name: String = {
            let mut name_chars: Vec<u16> = Vec::with_capacity(32);
            for _ in 0..32 {
                name_chars.push(reader.read_u16::<LittleEndian>()?);
            }
            let name_len_bytes = reader.read_u16::<LittleEndian>()?;
            if name_len_bytes > 64 {
                malformed!("name length too large: {}", name_len_bytes);
            } else if name_len_bytes % 2 != 0 {
                malformed!("odd name length: {}", name_len_bytes);
            }
            let name_len_chars = if name_len_bytes > 0 {
                (name_len_bytes / 2 - 1) as usize
            } else {
                0
            };
            debug_assert!(name_len_chars < name_chars.len());
            if name_chars[name_len_chars] != 0 {
                malformed!("name not null-terminated");
            }
            match String::from_utf16(&name_chars[0..name_len_chars]) {
                Ok(name) => name,
                Err(_) => malformed!("name not valid UTF-16"),
            }
        };
        internal::path::validate_name(&name)?;
        let obj_type = reader.read_u8()?;
        if obj_type != consts::OBJ_TYPE_UNALLOCATED
            && obj_type != consts::OBJ_TYPE_STORAGE
            && obj_type != consts::OBJ_TYPE_STREAM
            && obj_type != consts::OBJ_TYPE_ROOT
        {
            malformed!("invalid object type: {}", obj_type);
        }
        let color = reader.read_u8()?;
        if color != consts::COLOR_RED && color != consts::COLOR_BLACK {
            malformed!("invalid color: {}", color);
        }
        let left_sibling = reader.read_u32::<LittleEndian>()?;
        if left_sibling != NO_STREAM && left_sibling > MAX_REGULAR_STREAM_ID {
            malformed!("invalid left sibling: {}", left_sibling);
        }
        let right_sibling = reader.read_u32::<LittleEndian>()?;
        if right_sibling != NO_STREAM && right_sibling > MAX_REGULAR_STREAM_ID
        {
            malformed!("invalid right sibling: {}", right_sibling);
        }
        let child = reader.read_u32::<LittleEndian>()?;
        if child != NO_STREAM {
            if obj_type == consts::OBJ_TYPE_STREAM {
                malformed!("non-empty stream child: {}", child);
            } else if child > MAX_REGULAR_STREAM_ID {
                malformed!("invalid child: {}", child);
            }
        }
        let clsid = DirEntry::read_clsid(reader)?;
        if obj_type == consts::OBJ_TYPE_STREAM && !clsid.is_nil() {
            malformed!("non-null stream CLSID: {:?}", clsid);
        }
        let state_bits = reader.read_u32::<LittleEndian>()?;
        let creation_time = reader.read_u64::<LittleEndian>()?;
        let modified_time = reader.read_u64::<LittleEndian>()?;
        let start_sector = reader.read_u32::<LittleEndian>()?;

        // Spec say this is suppose to be zero for DirEntries
        // but some cfb implementations set start_sector to FREE_SECTOR instead
        if obj_type == consts::OBJ_TYPE_STORAGE
            && !(start_sector == 0 || start_sector == consts::FREE_SECTOR)
        {
            malformed!("non-zero storage start sector: {}", start_sector);
        }
        let stream_len =
            reader.read_u64::<LittleEndian>()? & version.stream_len_mask();
        if obj_type == consts::OBJ_TYPE_STORAGE && stream_len != 0 {
            malformed!("non-zero storage stream length: {}", stream_len);
        }
        Ok(DirEntry {
            name,
            obj_type,
            color,
            left_sibling,
            right_sibling,
            child,
            clsid,
            state_bits,
            creation_time,
            modified_time,
            start_sector,
            stream_len,
        })
    }

    pub fn write_to<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        debug_assert!(internal::path::validate_name(&self.name).is_ok());
        let name_utf16: Vec<u16> = self.name.encode_utf16().collect();
        debug_assert!(name_utf16.len() < 32);
        for &chr in name_utf16.iter() {
            writer.write_u16::<LittleEndian>(chr)?;
        }
        for _ in name_utf16.len()..32 {
            writer.write_u16::<LittleEndian>(0)?;
        }
        writer.write_u16::<LittleEndian>((name_utf16.len() as u16 + 1) * 2)?;
        writer.write_u8(self.obj_type)?;
        writer.write_u8(self.color)?;
        writer.write_u32::<LittleEndian>(self.left_sibling)?;
        writer.write_u32::<LittleEndian>(self.right_sibling)?;
        writer.write_u32::<LittleEndian>(self.child)?;
        DirEntry::write_clsid(writer, &self.clsid)?;
        writer.write_u32::<LittleEndian>(self.state_bits)?;
        writer.write_u64::<LittleEndian>(self.creation_time)?;
        writer.write_u64::<LittleEndian>(self.modified_time)?;
        writer.write_u32::<LittleEndian>(self.start_sector)?;
        writer.write_u64::<LittleEndian>(self.stream_len)?;
        Ok(())
    }
}

// ========================================================================= //

#[cfg(test)]
mod tests {
    use super::DirEntry;
    use crate::internal::consts;
    use crate::internal::Version;
    use uuid::Uuid;

    #[test]
    fn parse_valid_storage_entry() {
        let input: [u8; consts::DIR_ENTRY_LEN] = [
            // Name:
            70, 0, 111, 0, 111, 0, 98, 0, 97, 0, 114, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 14, 0, // name length
            1, // obj type
            1, // color,
            12, 0, 0, 0, // left sibling
            34, 0, 0, 0, // right sibling
            56, 0, 0, 0, // child
            0xe0, 0x85, 0x9f, 0xf2, 0xf9, 0x4f, 0x68, 0x10, // CLSID
            0xab, 0x91, 0x08, 0x00, 0x2b, 0x27, 0xb3, 0xd9, // CLSID
            239, 190, 173, 222, // state bits
            0, 0, 0, 0, 0, 0, 0, 0, // created
            0, 0, 0, 0, 0, 0, 0, 0, // modified
            0, 0, 0, 0, // start sector
            0, 0, 0, 0, 0, 0, 0, 0, // stream length
        ];
        let dir_entry =
            DirEntry::read_from(&mut (&input as &[u8]), Version::V4).unwrap();
        assert_eq!(&dir_entry.name, "Foobar");
        assert_eq!(dir_entry.obj_type, consts::OBJ_TYPE_STORAGE);
        assert_eq!(dir_entry.color, consts::COLOR_BLACK);
        assert_eq!(dir_entry.left_sibling, 12);
        assert_eq!(dir_entry.right_sibling, 34);
        assert_eq!(dir_entry.child, 56);
        assert_eq!(
            dir_entry.clsid,
            Uuid::parse_str("F29F85E0-4FF9-1068-AB91-08002B27B3D9").unwrap()
        );
        assert_eq!(dir_entry.state_bits, 0xdeadbeef);
        assert_eq!(dir_entry.creation_time, 0);
        assert_eq!(dir_entry.modified_time, 0);
        assert_eq!(dir_entry.start_sector, 0);
        assert_eq!(dir_entry.stream_len, 0);
    }

    #[test]
    #[should_panic(expected = "Malformed directory entry \
                               (invalid object type: 3)")]
    fn invalid_object_type() {
        let input: [u8; consts::DIR_ENTRY_LEN] = [
            // Name:
            70, 0, 111, 0, 111, 0, 98, 0, 97, 0, 114, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 14, 0, // name length
            3, // obj type
            1, // color,
            12, 0, 0, 0, // left sibling
            34, 0, 0, 0, // right sibling
            56, 0, 0, 0, // child
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // CLSID
            239, 190, 173, 222, // state bits
            0, 0, 0, 0, 0, 0, 0, 0, // created
            0, 0, 0, 0, 0, 0, 0, 0, // modified
            0, 0, 0, 0, // start sector
            0, 0, 0, 0, 0, 0, 0, 0, // stream length
        ];
        DirEntry::read_from(&mut (&input as &[u8]), Version::V4).unwrap();
    }

    #[test]
    #[should_panic(expected = "Malformed directory entry (invalid color: 2)")]
    fn invalid_color() {
        let input: [u8; consts::DIR_ENTRY_LEN] = [
            // Name:
            70, 0, 111, 0, 111, 0, 98, 0, 97, 0, 114, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 14, 0, // name length
            1, // obj type
            2, // color,
            12, 0, 0, 0, // left sibling
            34, 0, 0, 0, // right sibling
            56, 0, 0, 0, // child
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // CLSID
            239, 190, 173, 222, // state bits
            0, 0, 0, 0, 0, 0, 0, 0, // created
            0, 0, 0, 0, 0, 0, 0, 0, // modified
            0, 0, 0, 0, // start sector
            0, 0, 0, 0, 0, 0, 0, 0, // stream length
        ];
        DirEntry::read_from(&mut (&input as &[u8]), Version::V4).unwrap();
    }
}

// ========================================================================= //
<commit_msg>Loosen test on directory entries to parse supposedly valid CFB files from Altium Designer.<commit_after>use crate::internal::consts::{self, MAX_REGULAR_STREAM_ID, NO_STREAM};
use crate::internal::{self, Version};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::io::{self, Read, Write};
use uuid::Uuid;

// ========================================================================= //

macro_rules! malformed {
    ($e:expr) => { invalid_data!("Malformed directory entry ({})", $e) };
    ($fmt:expr, $($arg:tt)+) => {
        invalid_data!("Malformed directory entry ({})",
                      format!($fmt, $($arg)+))
    };
}

// ========================================================================= //

#[derive(Clone)]
pub struct DirEntry {
    pub name: String,
    pub obj_type: u8,
    pub color: u8,
    pub left_sibling: u32,
    pub right_sibling: u32,
    pub child: u32,
    pub clsid: Uuid,
    pub state_bits: u32,
    pub creation_time: u64,
    pub modified_time: u64,
    pub start_sector: u32,
    pub stream_len: u64,
}

impl DirEntry {
    pub fn unallocated() -> DirEntry {
        DirEntry {
            name: String::new(),
            obj_type: consts::OBJ_TYPE_UNALLOCATED,
            color: 0,
            left_sibling: NO_STREAM,
            right_sibling: NO_STREAM,
            child: NO_STREAM,
            clsid: Uuid::nil(),
            state_bits: 0,
            creation_time: 0,
            modified_time: 0,
            start_sector: 0,
            stream_len: 0,
        }
    }

    pub fn read_clsid<R: Read>(reader: &mut R) -> io::Result<Uuid> {
        let d1 = reader.read_u32::<LittleEndian>()?;
        let d2 = reader.read_u16::<LittleEndian>()?;
        let d3 = reader.read_u16::<LittleEndian>()?;
        let mut d4 = [0u8; 8];
        reader.read_exact(&mut d4)?;
        Ok(Uuid::from_fields(d1, d2, d3, &d4).unwrap())
    }

    pub fn write_clsid<W: Write>(
        writer: &mut W,
        clsid: &Uuid,
    ) -> io::Result<()> {
        let (d1, d2, d3, d4) = clsid.as_fields();
        writer.write_u32::<LittleEndian>(d1)?;
        writer.write_u16::<LittleEndian>(d2)?;
        writer.write_u16::<LittleEndian>(d3)?;
        writer.write_all(d4)?;
        Ok(())
    }

    pub fn read_from<R: Read>(
        reader: &mut R,
        version: Version,
    ) -> io::Result<DirEntry> {
        let name: String = {
            let mut name_chars: Vec<u16> = Vec::with_capacity(32);
            for _ in 0..32 {
                name_chars.push(reader.read_u16::<LittleEndian>()?);
            }
            let name_len_bytes = reader.read_u16::<LittleEndian>()?;
            if name_len_bytes > 64 {
                malformed!("name length too large: {}", name_len_bytes);
            } else if name_len_bytes % 2 != 0 {
                malformed!("odd name length: {}", name_len_bytes);
            }
            let name_len_chars = if name_len_bytes > 0 {
                (name_len_bytes / 2 - 1) as usize
            } else {
                0
            };
            debug_assert!(name_len_chars < name_chars.len());
            if name_chars[name_len_chars] != 0 {
                malformed!("name not null-terminated");
            }
            match String::from_utf16(&name_chars[0..name_len_chars]) {
                Ok(name) => name,
                Err(_) => malformed!("name not valid UTF-16"),
            }
        };
        internal::path::validate_name(&name)?;
        let obj_type = reader.read_u8()?;
        if obj_type != consts::OBJ_TYPE_UNALLOCATED
            && obj_type != consts::OBJ_TYPE_STORAGE
            && obj_type != consts::OBJ_TYPE_STREAM
            && obj_type != consts::OBJ_TYPE_ROOT
        {
            malformed!("invalid object type: {}", obj_type);
        }
        let color = reader.read_u8()?;
        if color != consts::COLOR_RED && color != consts::COLOR_BLACK {
            malformed!("invalid color: {}", color);
        }
        let left_sibling = reader.read_u32::<LittleEndian>()?;
        if left_sibling != NO_STREAM && left_sibling > MAX_REGULAR_STREAM_ID {
            malformed!("invalid left sibling: {}", left_sibling);
        }
        let right_sibling = reader.read_u32::<LittleEndian>()?;
        if right_sibling != NO_STREAM && right_sibling > MAX_REGULAR_STREAM_ID
        {
            malformed!("invalid right sibling: {}", right_sibling);
        }
        let child = reader.read_u32::<LittleEndian>()?;
        if child != NO_STREAM {
            if obj_type == consts::OBJ_TYPE_STREAM {
                malformed!("non-empty stream child: {}", child);
            } else if child > MAX_REGULAR_STREAM_ID {
                malformed!("invalid child: {}", child);
            }
        }
        let clsid = DirEntry::read_clsid(reader)?;
        if obj_type == consts::OBJ_TYPE_STREAM && !clsid.is_nil() {
            malformed!("non-null stream CLSID: {:?}", clsid);
        }
        let state_bits = reader.read_u32::<LittleEndian>()?;
        let creation_time = reader.read_u64::<LittleEndian>()?;
        let modified_time = reader.read_u64::<LittleEndian>()?;
        let start_sector = reader.read_u32::<LittleEndian>()?;

        // Spec say this is suppose to be zero for DirEntries
        // but some cfb implementations set start_sector to FREE_SECTOR instead
        if obj_type == consts::OBJ_TYPE_STORAGE
            && !(start_sector == 0 || start_sector == consts::FREE_SECTOR
            || start_sector == (consts::FREE_SECTOR - 1))
        {
            malformed!("non-zero storage start sector: {:x} vs {:x}", start_sector, consts::FREE_SECTOR);
        }
        let stream_len =
            reader.read_u64::<LittleEndian>()? & version.stream_len_mask();
        if obj_type == consts::OBJ_TYPE_STORAGE && stream_len != 0 {
            malformed!("non-zero storage stream length: {}", stream_len);
        }
        Ok(DirEntry {
            name,
            obj_type,
            color,
            left_sibling,
            right_sibling,
            child,
            clsid,
            state_bits,
            creation_time,
            modified_time,
            start_sector,
            stream_len,
        })
    }

    pub fn write_to<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        debug_assert!(internal::path::validate_name(&self.name).is_ok());
        let name_utf16: Vec<u16> = self.name.encode_utf16().collect();
        debug_assert!(name_utf16.len() < 32);
        for &chr in name_utf16.iter() {
            writer.write_u16::<LittleEndian>(chr)?;
        }
        for _ in name_utf16.len()..32 {
            writer.write_u16::<LittleEndian>(0)?;
        }
        writer.write_u16::<LittleEndian>((name_utf16.len() as u16 + 1) * 2)?;
        writer.write_u8(self.obj_type)?;
        writer.write_u8(self.color)?;
        writer.write_u32::<LittleEndian>(self.left_sibling)?;
        writer.write_u32::<LittleEndian>(self.right_sibling)?;
        writer.write_u32::<LittleEndian>(self.child)?;
        DirEntry::write_clsid(writer, &self.clsid)?;
        writer.write_u32::<LittleEndian>(self.state_bits)?;
        writer.write_u64::<LittleEndian>(self.creation_time)?;
        writer.write_u64::<LittleEndian>(self.modified_time)?;
        writer.write_u32::<LittleEndian>(self.start_sector)?;
        writer.write_u64::<LittleEndian>(self.stream_len)?;
        Ok(())
    }
}

// ========================================================================= //

#[cfg(test)]
mod tests {
    use super::DirEntry;
    use crate::internal::consts;
    use crate::internal::Version;
    use uuid::Uuid;

    #[test]
    fn parse_valid_storage_entry() {
        let input: [u8; consts::DIR_ENTRY_LEN] = [
            // Name:
            70, 0, 111, 0, 111, 0, 98, 0, 97, 0, 114, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 14, 0, // name length
            1, // obj type
            1, // color,
            12, 0, 0, 0, // left sibling
            34, 0, 0, 0, // right sibling
            56, 0, 0, 0, // child
            0xe0, 0x85, 0x9f, 0xf2, 0xf9, 0x4f, 0x68, 0x10, // CLSID
            0xab, 0x91, 0x08, 0x00, 0x2b, 0x27, 0xb3, 0xd9, // CLSID
            239, 190, 173, 222, // state bits
            0, 0, 0, 0, 0, 0, 0, 0, // created
            0, 0, 0, 0, 0, 0, 0, 0, // modified
            0, 0, 0, 0, // start sector
            0, 0, 0, 0, 0, 0, 0, 0, // stream length
        ];
        let dir_entry =
            DirEntry::read_from(&mut (&input as &[u8]), Version::V4).unwrap();
        assert_eq!(&dir_entry.name, "Foobar");
        assert_eq!(dir_entry.obj_type, consts::OBJ_TYPE_STORAGE);
        assert_eq!(dir_entry.color, consts::COLOR_BLACK);
        assert_eq!(dir_entry.left_sibling, 12);
        assert_eq!(dir_entry.right_sibling, 34);
        assert_eq!(dir_entry.child, 56);
        assert_eq!(
            dir_entry.clsid,
            Uuid::parse_str("F29F85E0-4FF9-1068-AB91-08002B27B3D9").unwrap()
        );
        assert_eq!(dir_entry.state_bits, 0xdeadbeef);
        assert_eq!(dir_entry.creation_time, 0);
        assert_eq!(dir_entry.modified_time, 0);
        assert_eq!(dir_entry.start_sector, 0);
        assert_eq!(dir_entry.stream_len, 0);
    }

    #[test]
    #[should_panic(expected = "Malformed directory entry \
                               (invalid object type: 3)")]
    fn invalid_object_type() {
        let input: [u8; consts::DIR_ENTRY_LEN] = [
            // Name:
            70, 0, 111, 0, 111, 0, 98, 0, 97, 0, 114, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 14, 0, // name length
            3, // obj type
            1, // color,
            12, 0, 0, 0, // left sibling
            34, 0, 0, 0, // right sibling
            56, 0, 0, 0, // child
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // CLSID
            239, 190, 173, 222, // state bits
            0, 0, 0, 0, 0, 0, 0, 0, // created
            0, 0, 0, 0, 0, 0, 0, 0, // modified
            0, 0, 0, 0, // start sector
            0, 0, 0, 0, 0, 0, 0, 0, // stream length
        ];
        DirEntry::read_from(&mut (&input as &[u8]), Version::V4).unwrap();
    }

    #[test]
    #[should_panic(expected = "Malformed directory entry (invalid color: 2)")]
    fn invalid_color() {
        let input: [u8; consts::DIR_ENTRY_LEN] = [
            // Name:
            70, 0, 111, 0, 111, 0, 98, 0, 97, 0, 114, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 14, 0, // name length
            1, // obj type
            2, // color,
            12, 0, 0, 0, // left sibling
            34, 0, 0, 0, // right sibling
            56, 0, 0, 0, // child
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // CLSID
            239, 190, 173, 222, // state bits
            0, 0, 0, 0, 0, 0, 0, 0, // created
            0, 0, 0, 0, 0, 0, 0, 0, // modified
            0, 0, 0, 0, // start sector
            0, 0, 0, 0, 0, 0, 0, 0, // stream length
        ];
        DirEntry::read_from(&mut (&input as &[u8]), Version::V4).unwrap();
    }
}

// ========================================================================= //
<|endoftext|>


<commit_before>use ffi;
use libc;

use std::io::Cursor;
use std::io::Read;
use std::io::Error as IoError;
use std::mem;
use std::ptr;

use AsMutLua;

use LuaRead;
use LuaError;
use Push;
use PushGuard;
use PushOne;

pub struct LuaCode<'a>(&'a str);

impl<'lua, 'c, L> Push<L> for LuaCode<'c> where L: AsMutLua<'lua> {
    type Err = LuaError;

    #[inline]
    fn push_to_lua(self, lua: L) -> Result<PushGuard<L>, (LuaError, L)> {
        LuaCodeFromReader(Cursor::new(self.0.as_bytes())).push_to_lua(lua)
    }
}

impl<'lua, 'c, L> PushOne<L> for LuaCode<'c> where L: AsMutLua<'lua> {
}

pub struct LuaCodeFromReader<R>(R);

impl<'lua, L, R> Push<L> for LuaCodeFromReader<R>
    where L: AsMutLua<'lua>,
          R: Read
{
    type Err = LuaError;

    #[inline]
    fn push_to_lua(self, mut lua: L) -> Result<PushGuard<L>, (LuaError, L)> {
        struct ReadData<R> {
            reader: R,
            buffer: [u8; 128],
            triggered_error: Option<IoError>,
        }

        extern "C" fn reader<R>(_: *mut ffi::lua_State,
                                data_raw: *mut libc::c_void,
                                size: *mut libc::size_t)
                                -> *const libc::c_char
            where R: Read
        {
            let data: &mut ReadData<R> = unsafe { mem::transmute(data_raw) };

            if data.triggered_error.is_some() {
                unsafe { (*size) = 0 }
                return data.buffer.as_ptr() as *const libc::c_char;
            }

            match data.reader.read(&mut data.buffer) {
                Ok(len) => unsafe { (*size) = len as libc::size_t },
                Err(e) => {
                    unsafe { (*size) = 0 }
                    data.triggered_error = Some(e)
                }
            };

            data.buffer.as_ptr() as *const libc::c_char
        }

        let readdata = ReadData {
            reader: self.0,
            buffer: unsafe { ::std::mem::uninitialized() },
            triggered_error: None,
        };

        let (load_return_value, pushed_value) = unsafe {
            let code = ffi::lua_load(lua.as_mut_lua().0,
                                     reader::<R>,
                                     mem::transmute(&readdata),
                                     b"chunk\0".as_ptr() as *const _,
                                     ptr::null());
            let raw_lua = lua.as_lua();
            (code,
             PushGuard {
                 lua: lua,
                 size: 1,
                 raw_lua: raw_lua,
             })
        };

        if readdata.triggered_error.is_some() {
            let error = readdata.triggered_error.unwrap();
            return Err((LuaError::ReadError(error), pushed_value.into_inner()));
        }

        if load_return_value == 0 {
            return Ok(pushed_value);
        }

        let error_msg: String = LuaRead::lua_read(&pushed_value)
            .ok()
            .expect("can't find error message at the top of the Lua stack");

        if load_return_value == ffi::LUA_ERRMEM {
            panic!("LUA_ERRMEM");
        }

        if load_return_value == ffi::LUA_ERRSYNTAX {
            return Err((LuaError::SyntaxError(error_msg), pushed_value.into_inner()));
        }

        panic!("Unknown error while calling lua_load");
    }
}

impl<'lua, L, R> PushOne<L> for LuaCodeFromReader<R>
    where L: AsMutLua<'lua>,
          R: Read
{
}

///
pub struct LuaFunction<L> {
    variable: L,
}

impl<'lua, L> LuaFunction<L>
    where L: AsMutLua<'lua>
{
    /// Calls the `LuaFunction`.
    #[inline]
    pub fn call<'a, V>(&'a mut self) -> Result<V, LuaError>
        where V: LuaRead<PushGuard<&'a mut L>>
    {
        // calling pcall pops the parameters and pushes output
        let (pcall_return_value, pushed_value) = unsafe {
            // lua_pcall pops the function, so we have to make a copy of it
            ffi::lua_pushvalue(self.variable.as_mut_lua().0, -1);
            let pcall_return_value = ffi::lua_pcall(self.variable.as_mut_lua().0, 0, 1, 0);     // TODO: arguments
            let raw_lua = self.variable.as_lua();
            (pcall_return_value,
             PushGuard {
                 lua: &mut self.variable,
                 size: 1,
                 raw_lua: raw_lua,
             })
        };

        // if pcall succeeded, returning
        if pcall_return_value == 0 {
            return match LuaRead::lua_read(pushed_value) {
                Err(_) => Err(LuaError::WrongType),
                Ok(x) => Ok(x),
            };
        }

        // an error occured during execution
        if pcall_return_value == ffi::LUA_ERRMEM {
            panic!("lua_pcall returned LUA_ERRMEM");
        }

        if pcall_return_value == ffi::LUA_ERRRUN {
            let error_msg: String = LuaRead::lua_read(pushed_value)
                .ok()
                .expect("can't find error message at the top of the Lua stack");
            return Err(LuaError::ExecutionError(error_msg));
        }

        panic!("Unknown error code returned by lua_pcall: {}", pcall_return_value)
    }

    /// Builds a new `LuaFunction` from the code of a reader.
    #[inline]
    pub fn load_from_reader<R>(lua: L, code: R) -> Result<LuaFunction<PushGuard<L>>, LuaError>
        where R: Read
    {
        match LuaCodeFromReader(code).push_to_lua(lua) {
            Ok(pushed) => Ok(LuaFunction { variable: pushed }),
            Err((err, _)) => Err(err),
        }
    }

    /// Builds a new `LuaFunction` from a raw string.
    #[inline]
    pub fn load(lua: L, code: &str) -> Result<LuaFunction<PushGuard<L>>, LuaError> {
        let reader = Cursor::new(code.as_bytes());
        LuaFunction::load_from_reader(lua, reader)
    }
}

// TODO: return Result<Ret, ExecutionError> instead
// impl<'a, 'lua, Ret: CopyRead> ::std::ops::FnMut<(), Ret> for LuaFunction<'a,'lua> {
// fn call_mut(&mut self, _: ()) -> Ret {
// self.call().unwrap()
// }
// }

impl<'lua, L> LuaRead<L> for LuaFunction<L>
    where L: AsMutLua<'lua>
{
    #[inline]
    fn lua_read_at_position(mut lua: L, index: i32) -> Result<LuaFunction<L>, L> {
        assert!(index == -1);   // FIXME:
        if unsafe { ffi::lua_isfunction(lua.as_mut_lua().0, -1) } {
            Ok(LuaFunction { variable: lua })
        } else {
            Err(lua)
        }
    }
}
<commit_msg>Cleanup functions_read<commit_after>use ffi;
use libc;

use std::io::Cursor;
use std::io::Read;
use std::io::Error as IoError;
use std::mem;
use std::ptr;

use AsMutLua;

use LuaRead;
use LuaError;
use Push;
use PushGuard;
use PushOne;

/// Wrapper around a `&str`. When pushed, the content will be parsed as Lua code and turned into a
/// function.
///
/// Since pushing this value can fail in case of a parsing error, you must use the `checked_set`
/// method instead of `set`.
///
/// > **Note**: This struct is a wrapper around `LuaCodeFromReader`. There's no advantage in using
/// > it except that it is more convenient. More advanced usages (such as returning a Lua function
/// > from a Rust function) can be done with `LuaCodeFromReader`.
///
/// # Example
///
/// ```
/// let mut lua = hlua::Lua::new();
/// lua.checked_set("hello", hlua::LuaCode("return 5")).unwrap();
///
/// let r: i32 = lua.execute("return hello();").unwrap();
/// assert_eq!(r, 5);
/// ```
pub struct LuaCode<'a>(pub &'a str);

impl<'lua, 'c, L> Push<L> for LuaCode<'c>
    where L: AsMutLua<'lua>
{
    type Err = LuaError;

    #[inline]
    fn push_to_lua(self, lua: L) -> Result<PushGuard<L>, (LuaError, L)> {
        LuaCodeFromReader(Cursor::new(self.0.as_bytes())).push_to_lua(lua)
    }
}

impl<'lua, 'c, L> PushOne<L> for LuaCode<'c> where L: AsMutLua<'lua> {}

/// Wrapper around a `Read` object. When pushed, the content will be parsed as Lua code and turned
/// into a function.
///
/// Since pushing this value can fail in case of a reading error or a parsing error, you must use
/// the `checked_set` method instead of `set`.
///
/// # Example: returning a Lua function from a Rust function
///
/// ```
/// use std::io::Cursor;
///
/// let mut lua = hlua::Lua::new();
///
/// lua.set("call_rust", hlua::function0(|| -> hlua::LuaCodeFromReader<Cursor<String>> {
///     let lua_code = "return 18;";
///     return hlua::LuaCodeFromReader(Cursor::new(lua_code.to_owned()));
/// }));
///
/// let r: i32 = lua.execute("local lua_func = call_rust(); return lua_func();").unwrap();
/// assert_eq!(r, 18);
/// ```
pub struct LuaCodeFromReader<R>(pub R);

impl<'lua, L, R> Push<L> for LuaCodeFromReader<R>
    where L: AsMutLua<'lua>,
          R: Read
{
    type Err = LuaError;

    #[inline]
    fn push_to_lua(self, mut lua: L) -> Result<PushGuard<L>, (LuaError, L)> {
        unsafe {
            struct ReadData<R> {
                reader: R,
                buffer: [u8; 128],
                triggered_error: Option<IoError>,
            }

            let mut read_data = ReadData {
                reader: self.0,
                buffer: mem::uninitialized(),
                triggered_error: None,
            };

            extern "C" fn reader<R>(_: *mut ffi::lua_State,
                                    data: *mut libc::c_void,
                                    size: *mut libc::size_t)
                                    -> *const libc::c_char
                where R: Read
            {
                unsafe {
                    let data: *mut ReadData<R> = data as *mut _;
                    let data: &mut ReadData<R> = &mut *data;

                    if data.triggered_error.is_some() {
                        (*size) = 0;
                        return data.buffer.as_ptr() as *const libc::c_char;
                    }

                    match data.reader.read(&mut data.buffer) {
                        Ok(len) => (*size) = len as libc::size_t,
                        Err(e) => {
                            (*size) = 0;
                            data.triggered_error = Some(e);
                        }
                    };

                    data.buffer.as_ptr() as *const libc::c_char
                }
            }

            let (load_return_value, pushed_value) = {
                let code = ffi::lua_load(lua.as_mut_lua().0,
                                         reader::<R>,
                                         &mut read_data as *mut ReadData<_> as *mut libc::c_void,
                                         b"chunk\0".as_ptr() as *const _,
                                         ptr::null());
                let raw_lua = lua.as_lua();
                (code,
                 PushGuard {
                     lua: lua,
                     size: 1,
                     raw_lua: raw_lua,
                 })
            };

            if read_data.triggered_error.is_some() {
                let error = read_data.triggered_error.unwrap();
                return Err((LuaError::ReadError(error), pushed_value.into_inner()));
            }

            if load_return_value == 0 {
                return Ok(pushed_value);
            }

            let error_msg: String = LuaRead::lua_read(&pushed_value)
                .ok()
                .expect("can't find error message at the top of the Lua stack");

            if load_return_value == ffi::LUA_ERRMEM {
                panic!("LUA_ERRMEM");
            }

            if load_return_value == ffi::LUA_ERRSYNTAX {
                return Err((LuaError::SyntaxError(error_msg), pushed_value.into_inner()));
            }

            panic!("Unknown error while calling lua_load");
        }
    }
}

impl<'lua, L, R> PushOne<L> for LuaCodeFromReader<R>
    where L: AsMutLua<'lua>,
          R: Read
{
}

/// Handle to a function in the Lua context.
///
/// Just like you can read variables as integers and strings, you can also read Lua functions by
/// requesting a `LuaFunction` object. Once you have a `LuaFunction` you can call it with `call()`.
///
/// > **Note**: Passing parameters when calling the function is not yet implemented.
///
/// # Example
///
/// ```
/// let mut lua = hlua::Lua::new();
/// lua.execute::<()>("function foo() return 12 end").unwrap();
///
/// let mut foo: hlua::LuaFunction<_> = lua.get("foo").unwrap();
/// let result: i32 = foo.call().unwrap();
/// assert_eq!(result, 12);
/// ```
// TODO: example for how to get a LuaFunction as a parameter of a Rust function
pub struct LuaFunction<L> {
    variable: L,
}

impl<'lua, L> LuaFunction<L>
    where L: AsMutLua<'lua>
{
    /// Calls the function.
    ///
    /// Returns an error if there is an error while executing the Lua code (eg. a function call
    /// returns an error), or if the requested return type doesn't match the actual return type.
    ///
    /// > **Note**: Passing parameters when calling is not yet implemented.
    #[inline]
    pub fn call<'a, V>(&'a mut self) -> Result<V, LuaError>
        where V: LuaRead<PushGuard<&'a mut L>>
    {
        // calling pcall pops the parameters and pushes output
        let (pcall_return_value, pushed_value) = unsafe {
            // lua_pcall pops the function, so we have to make a copy of it
            ffi::lua_pushvalue(self.variable.as_mut_lua().0, -1);
            let pcall_return_value = ffi::lua_pcall(self.variable.as_mut_lua().0, 0, 1, 0);     // TODO: arguments
            let raw_lua = self.variable.as_lua();
            (pcall_return_value,
             PushGuard {
                 lua: &mut self.variable,
                 size: 1,
                 raw_lua: raw_lua,
             })
        };

        // if pcall succeeded, returning
        if pcall_return_value == 0 {
            return match LuaRead::lua_read(pushed_value) {
                Err(_) => Err(LuaError::WrongType),
                Ok(x) => Ok(x),
            };
        }

        // an error occured during execution
        if pcall_return_value == ffi::LUA_ERRMEM {
            panic!("lua_pcall returned LUA_ERRMEM");
        }

        if pcall_return_value == ffi::LUA_ERRRUN {
            let error_msg: String = LuaRead::lua_read(pushed_value)
                .ok()
                .expect("can't find error message at the top of the Lua stack");
            return Err(LuaError::ExecutionError(error_msg));
        }

        panic!("Unknown error code returned by lua_pcall: {}",
               pcall_return_value)
    }

    /// Builds a new `LuaFunction` from the code of a reader.
    ///
    /// Returns an error if reading from the `Read` object fails or if there is a syntax error in
    /// the code.
    ///
    /// # Example
    ///
    /// ```
    /// use std::io::Cursor;
    ///
    /// let mut lua = hlua::Lua::new();
    ///
    /// let mut f = hlua::LuaFunction::load_from_reader(&mut lua, Cursor::new("return 8")).unwrap();
    /// let ret: i32 = f.call().unwrap();
    /// assert_eq!(ret, 8);
    /// ```
    #[inline]
    pub fn load_from_reader<R>(lua: L, code: R) -> Result<LuaFunction<PushGuard<L>>, LuaError>
        where R: Read
    {
        match LuaCodeFromReader(code).push_to_lua(lua) {
            Ok(pushed) => Ok(LuaFunction { variable: pushed }),
            Err((err, _)) => Err(err),
        }
    }

    /// Builds a new `LuaFunction` from a raw string.
    ///
    /// > **Note**: This is just a wrapper around `load_from_reader`. There is no advantage in
    /// > using `load` except that it is more convenient.
    // TODO: remove this function? it's only a thin wrapper and it's for a very niche situation
    #[inline]
    pub fn load(lua: L, code: &str) -> Result<LuaFunction<PushGuard<L>>, LuaError> {
        let reader = Cursor::new(code.as_bytes());
        LuaFunction::load_from_reader(lua, reader)
    }
}

// TODO: return Result<Ret, ExecutionError> instead
// impl<'a, 'lua, Ret: CopyRead> ::std::ops::FnMut<(), Ret> for LuaFunction<'a,'lua> {
// fn call_mut(&mut self, _: ()) -> Ret {
// self.call().unwrap()
// }
// }

impl<'lua, L> LuaRead<L> for LuaFunction<L>
    where L: AsMutLua<'lua>
{
    #[inline]
    fn lua_read_at_position(mut lua: L, index: i32) -> Result<LuaFunction<L>, L> {
        assert!(index == -1);   // FIXME:
        if unsafe { ffi::lua_isfunction(lua.as_mut_lua().0, -1) } {
            Ok(LuaFunction { variable: lua })
        } else {
            Err(lua)
        }
    }
}
<|endoftext|>


<commit_before>use anyhow::{Context, Error};
use next_swc::{custom_before_pass, TransformOptions};
use once_cell::sync::Lazy;
use std::sync::Arc;
use swc::{config::JsMinifyOptions, config::ParseOptions, try_with_handler, Compiler};
use swc_common::{comments::Comments, errors::ColorConfig, FileName, FilePathMapping, SourceMap};
use swc_ecmascript::transforms::pass::noop;
use wasm_bindgen::prelude::*;

fn convert_err(err: Error) -> JsValue {
    format!("{:?}", err).into()
}

#[wasm_bindgen(js_name = "minifySync")]
pub fn minify_sync(s: &str, opts: JsValue) -> Result<JsValue, JsValue> {
    console_error_panic_hook::set_once();

    let c = compiler();

    try_with_handler(
        c.cm.clone(),
        swc::HandlerOpts {
            color: ColorConfig::Never,
            skip_filename: false,
        },
        |handler| {
            let opts: JsMinifyOptions = opts.into_serde().context("failed to parse options")?;

            let fm = c.cm.new_source_file(FileName::Anon, s.into());
            let program = c
                .minify(fm, handler, &opts)
                .context("failed to minify file")?;

            JsValue::from_serde(&program).context("failed to serialize json")
        },
    )
    .map_err(convert_err)
}

#[wasm_bindgen(js_name = "transformSync")]
pub fn transform_sync(s: &str, opts: JsValue) -> Result<JsValue, JsValue> {
    console_error_panic_hook::set_once();

    let c = compiler();

    try_with_handler(
        c.cm.clone(),
        swc::HandlerOpts {
            color: ColorConfig::Never,
            skip_filename: false,
        },
        |handler| {
            let opts: TransformOptions = opts.into_serde().context("failed to parse options")?;

            let fm = c.cm.new_source_file(
                if opts.swc.filename.is_empty() {
                    FileName::Anon
                } else {
                    FileName::Real(opts.swc.filename.clone().into())
                },
                s.into(),
            );
            let cm = c.cm.clone();
            let file = fm.clone();
            let out = c
                .process_js_with_custom_pass(
                    fm,
                    None,
                    handler,
                    &opts.swc,
                    |_, comments| {
                        custom_before_pass(cm, file, &opts, comments.clone(), Default::default())
                    },
                    |_, _| noop(),
                )
                .context("failed to process js file")?;

            JsValue::from_serde(&out).context("failed to serialize json")
        },
    )
    .map_err(convert_err)
}

#[wasm_bindgen(js_name = "parseSync")]
pub fn parse_sync(s: &str, opts: JsValue) -> Result<JsValue, JsValue> {
    console_error_panic_hook::set_once();

    let c = compiler();

    try_with_handler(
        c.cm.clone(),
        swc::HandlerOpts {
            ..Default::default()
        },
        |handler| {
            c.run(|| {
                let opts: ParseOptions = opts.into_serde().context("failed to parse options")?;

                let fm = c.cm.new_source_file(FileName::Anon, s.into());

                let cmts = c.comments().clone();
                let comments = if opts.comments {
                    Some(&cmts as &dyn Comments)
                } else {
                    None
                };

                let program = c
                    .parse_js(
                        fm,
                        handler,
                        opts.target,
                        opts.syntax,
                        opts.is_module,
                        comments,
                    )
                    .context("failed to parse code")?;

                let s = serde_json::to_string(&program).unwrap();
                Ok(JsValue::from_str(&s))
            })
        },
    )
    .map_err(convert_err)
}

/// Get global sourcemap
fn compiler() -> Arc<Compiler> {
    static C: Lazy<Arc<Compiler>> = Lazy::new(|| {
        let cm = Arc::new(SourceMap::new(FilePathMapping::empty()));

        Arc::new(Compiler::new(cm))
    });

    C.clone()
}
<commit_msg>next-swc/wasm: Make span start at 0 (#36532)<commit_after>use anyhow::{Context, Error};
use next_swc::{custom_before_pass, TransformOptions};
use once_cell::sync::Lazy;
use std::sync::Arc;
use swc::{config::JsMinifyOptions, config::ParseOptions, try_with_handler, Compiler};
use swc_common::{comments::Comments, errors::ColorConfig, FileName, FilePathMapping, SourceMap};
use swc_ecmascript::transforms::pass::noop;
use wasm_bindgen::prelude::*;

fn convert_err(err: Error) -> JsValue {
    format!("{:?}", err).into()
}

#[wasm_bindgen(js_name = "minifySync")]
pub fn minify_sync(s: &str, opts: JsValue) -> Result<JsValue, JsValue> {
    console_error_panic_hook::set_once();

    let c = compiler();

    try_with_handler(
        c.cm.clone(),
        swc::HandlerOpts {
            color: ColorConfig::Never,
            skip_filename: false,
        },
        |handler| {
            let opts: JsMinifyOptions = opts.into_serde().context("failed to parse options")?;

            let fm = c.cm.new_source_file(FileName::Anon, s.into());
            let program = c
                .minify(fm, handler, &opts)
                .context("failed to minify file")?;

            JsValue::from_serde(&program).context("failed to serialize json")
        },
    )
    .map_err(convert_err)
}

#[wasm_bindgen(js_name = "transformSync")]
pub fn transform_sync(s: &str, opts: JsValue) -> Result<JsValue, JsValue> {
    console_error_panic_hook::set_once();

    let c = compiler();

    try_with_handler(
        c.cm.clone(),
        swc::HandlerOpts {
            color: ColorConfig::Never,
            skip_filename: false,
        },
        |handler| {
            let opts: TransformOptions = opts.into_serde().context("failed to parse options")?;

            let fm = c.cm.new_source_file(
                if opts.swc.filename.is_empty() {
                    FileName::Anon
                } else {
                    FileName::Real(opts.swc.filename.clone().into())
                },
                s.into(),
            );
            let cm = c.cm.clone();
            let file = fm.clone();
            let out = c
                .process_js_with_custom_pass(
                    fm,
                    None,
                    handler,
                    &opts.swc,
                    |_, comments| {
                        custom_before_pass(cm, file, &opts, comments.clone(), Default::default())
                    },
                    |_, _| noop(),
                )
                .context("failed to process js file")?;

            JsValue::from_serde(&out).context("failed to serialize json")
        },
    )
    .map_err(convert_err)
}

#[wasm_bindgen(js_name = "parseSync")]
pub fn parse_sync(s: &str, opts: JsValue) -> Result<JsValue, JsValue> {
    console_error_panic_hook::set_once();

    let c = swc::Compiler::new(Arc::new(SourceMap::new(FilePathMapping::empty())));

    try_with_handler(
        c.cm.clone(),
        swc::HandlerOpts {
            ..Default::default()
        },
        |handler| {
            c.run(|| {
                let opts: ParseOptions = opts.into_serde().context("failed to parse options")?;

                let fm = c.cm.new_source_file(FileName::Anon, s.into());

                let cmts = c.comments().clone();
                let comments = if opts.comments {
                    Some(&cmts as &dyn Comments)
                } else {
                    None
                };

                let program = c
                    .parse_js(
                        fm,
                        handler,
                        opts.target,
                        opts.syntax,
                        opts.is_module,
                        comments,
                    )
                    .context("failed to parse code")?;

                let s = serde_json::to_string(&program).unwrap();
                Ok(JsValue::from_str(&s))
            })
        },
    )
    .map_err(convert_err)
}

/// Get global sourcemap
fn compiler() -> Arc<Compiler> {
    static C: Lazy<Arc<Compiler>> = Lazy::new(|| {
        let cm = Arc::new(SourceMap::new(FilePathMapping::empty()));

        Arc::new(Compiler::new(cm))
    });

    C.clone()
}
<|endoftext|>


<commit_before>use rustc_serialize::json;
use hyper;
use hyper::Client;
use url::form_urlencoded;
use crossbeam;

use std::io::Read;

use text_view::view_as_text;
use text_view::insert_annotation;
use token::Token;
use token::Data;
use error::{Error, Result, Source};

/// Represents a grammar error from language tool
///
/// Note: lots of fields are missing
#[derive(RustcDecodable, RustcEncodable, Debug)]
pub struct GrammarError {
    pub message: String,
    pub offset: usize,
    pub length: usize,
    pub short_message: Option<String>,
    pub issue_type: Option<String>,
}

/// Contains a list of matches to errors
///
/// Corresponds to the JSON that LanguageTool-server sends back
///
/// Note: lots of fields are missing
#[derive(RustcDecodable, RustcEncodable, Debug)]
pub struct GrammarCheck {
    pub matches: Vec<GrammarError>,
}

/// GrammarChecker
pub struct GrammarChecker {
    lang: String,
    port: usize,
}

impl GrammarChecker {
    /// Initialize the grammarchecker
    pub fn new<S: Into<String>>(port: usize, lang: S) -> Result<GrammarChecker> {
        let checker = GrammarChecker {
            lang: lang.into(),
            port: port,
        };

        let res = try!(Client::new()
            .get(&format!("http://localhost:{}/v2/languages", port))
            .send()
            .map_err(|e| {
                Error::grammar_check(Source::empty(),
                                     lformat!("could not connect to language tool server: {error}",
                                              error = e))
            }));
        if res.status != hyper::Ok {
            return Err(Error::grammar_check(Source::empty(),
                                            lformat!("server didn't respond with a OK status \
                                                      code")));
        }
        Ok(checker)
    }

    /// Send a query to LanguageTools server and get back a list of errors
    pub fn check(&self, text: &str) -> Result<GrammarCheck> {
        let query: String = form_urlencoded::Serializer::new(String::new())
            .append_pair("language", &self.lang)
            .append_pair("text", text)
            .finish();

        let client = Client::new();

        let mut res = try!(client.post(&format!("http://localhost:{}/v2/check", self.port))
            .body(&query)
            .send()
            .map_err(|e| {
                Error::grammar_check(Source::empty(),
                                     lformat!("could not send request to server: {error}",
                                              error = e))
            }));

        if res.status != hyper::Ok {
            return Err(Error::grammar_check(Source::empty(),
                                            lformat!("server didn't respond with a OK status \
                                                      code")));
        }

        let mut s = String::new();
        try!(res.read_to_string(&mut s)
            .map_err(|e| {
                Error::grammar_check(Source::empty(),
                                     lformat!("could not read response: {error}", error = e))
            }));
        let reponse: GrammarCheck = try!(json::decode(&s).map_err(|e| {
            Error::default(Source::empty(),
                           lformat!("could not decode JSON: {error}", error = e))
        }));
        Ok(reponse)
    }
}

impl GrammarChecker {
    /// Check the grammar in a vector of tokens.
    ///
    /// This modifies the AST
    pub fn check_chapter(&self, tokens: &mut Vec<Token>) -> Result<()> {
        let n_threads = 4; // seems to give best results but quite random
        let len = tokens.len();
        let subvecs = tokens.chunks_mut(len / n_threads);

        let mut result = Ok(());
        crossbeam::scope(|scope| {
            let mut handles = vec![];
            for subvec in subvecs {
                for mut token in subvec.iter_mut() {
                    match *token {
                        Token::Paragraph(ref mut v) |
                        Token::Header(_, ref mut v) |
                        Token::BlockQuote(ref mut v) |
                        Token::List(ref mut v) |
                        Token::OrderedList(_, ref mut v) => {
                            handles.push(scope.spawn(move || {
                                let check = try!(self.check(&view_as_text(v)));
                                for error in check.matches {
                                    insert_annotation(v,
                                                      &Data::GrammarError(error.message.clone()),
                                                      error.offset,
                                                      error.length);
                                }
                                Ok(())
                            }))
                        }
                        _ => (),
                    }
                }
            }
            for handle in handles.into_iter() {
                if let Err(err) = handle.join() {
                    result = Err(err);
                    break;
                }
            }
        });
        result

        // crossbeam::scope(|scope| {
        //     for subvec in subvecs {
        //         scope.spawn(|| Self::check_grammar(lang, subvec)
        //     }
        // });

        //        Ok(())
    }
}
<commit_msg>Use ? in grammar_check.rs<commit_after>use rustc_serialize::json;
use hyper;
use hyper::Client;
use url::form_urlencoded;
use crossbeam;

use std::io::Read;

use text_view::view_as_text;
use text_view::insert_annotation;
use token::Token;
use token::Data;
use error::{Error, Result, Source};

/// Represents a grammar error from language tool
///
/// Note: lots of fields are missing
#[derive(RustcDecodable, RustcEncodable, Debug)]
pub struct GrammarError {
    pub message: String,
    pub offset: usize,
    pub length: usize,
    pub short_message: Option<String>,
    pub issue_type: Option<String>,
}

/// Contains a list of matches to errors
///
/// Corresponds to the JSON that LanguageTool-server sends back
///
/// Note: lots of fields are missing
#[derive(RustcDecodable, RustcEncodable, Debug)]
pub struct GrammarCheck {
    pub matches: Vec<GrammarError>,
}

/// GrammarChecker
pub struct GrammarChecker {
    lang: String,
    port: usize,
}

impl GrammarChecker {
    /// Initialize the grammarchecker
    pub fn new<S: Into<String>>(port: usize, lang: S) -> Result<GrammarChecker> {
        let checker = GrammarChecker {
            lang: lang.into(),
            port: port,
        };

        let res = Client::new()
            .get(&format!("http://localhost:{}/v2/languages", port))
            .send()
            .map_err(|e| {
                Error::grammar_check(Source::empty(),
                                     lformat!("could not connect to language tool server: {error}",
                                              error = e))
            })?;
        if res.status != hyper::Ok {
            return Err(Error::grammar_check(Source::empty(),
                                            lformat!("server didn't respond with a OK status \
                                                      code")));
        }
        Ok(checker)
    }

    /// Send a query to LanguageTools server and get back a list of errors
    pub fn check(&self, text: &str) -> Result<GrammarCheck> {
        let query: String = form_urlencoded::Serializer::new(String::new())
            .append_pair("language", &self.lang)
            .append_pair("text", text)
            .finish();

        let client = Client::new();

        let mut res = client.post(&format!("http://localhost:{}/v2/check", self.port))
            .body(&query)
            .send()
            .map_err(|e| {
                Error::grammar_check(Source::empty(),
                                     lformat!("could not send request to server: {error}",
                                              error = e))
            })?;

        if res.status != hyper::Ok {
            return Err(Error::grammar_check(Source::empty(),
                                            lformat!("server didn't respond with a OK status \
                                                      code")));
        }

        let mut s = String::new();
        res.read_to_string(&mut s)
            .map_err(|e| {
                Error::grammar_check(Source::empty(),
                                     lformat!("could not read response: {error}", error = e))
            })?;
        let reponse: GrammarCheck = json::decode(&s)
            .map_err(|e| {
                Error::default(Source::empty(),
                               lformat!("could not decode JSON: {error}", error = e))
            })?;
        Ok(reponse)
    }
}

impl GrammarChecker {
    /// Check the grammar in a vector of tokens.
    ///
    /// This modifies the AST
    pub fn check_chapter(&self, tokens: &mut Vec<Token>) -> Result<()> {
        let n_threads = 4; // seems to give best results but quite random
        let len = tokens.len();
        let subvecs = tokens.chunks_mut(len / n_threads);

        let mut result = Ok(());
        crossbeam::scope(|scope| {
            let mut handles = vec![];
            for subvec in subvecs {
                for mut token in subvec.iter_mut() {
                    match *token {
                        Token::Paragraph(ref mut v) |
                        Token::Header(_, ref mut v) |
                        Token::BlockQuote(ref mut v) |
                        Token::List(ref mut v) |
                        Token::OrderedList(_, ref mut v) => {
                            handles.push(scope.spawn(move || {
                                let check = self.check(&view_as_text(v))?;
                                for error in check.matches {
                                    insert_annotation(v,
                                                      &Data::GrammarError(error.message.clone()),
                                                      error.offset,
                                                      error.length);
                                }
                                Ok(())
                            }))
                        }
                        _ => (),
                    }
                }
            }
            for handle in handles.into_iter() {
                if let Err(err) = handle.join() {
                    result = Err(err);
                    break;
                }
            }
        });
        result

        // crossbeam::scope(|scope| {
        //     for subvec in subvecs {
        //         scope.spawn(|| Self::check_grammar(lang, subvec)
        //     }
        // });

        //        Ok(())
    }
}
<|endoftext|>