<commit_before>// FIXME(https://github.com/rust-lang/rust/issues/19630) Remove this work-around
#[macro_export]
#[doc(hidden)]
macro_rules! diesel_internal_expr_conversion {
    ($e:expr) => { $e }
}

#[macro_export]
#[doc(hidden)]
macro_rules! column {
    ($($table:ident)::*, $column_name:ident -> $Type:ty) => {
        #[allow(non_camel_case_types, dead_code)]
        #[derive(Debug, Clone, Copy)]
        pub struct $column_name;

        impl $crate::expression::Expression for $column_name {
            type SqlType = $Type;
        }

        impl<DB> $crate::query_builder::QueryFragment<DB> for $column_name where
            DB: $crate::backend::Backend,
        {
            fn to_sql(&self, out: &mut DB::QueryBuilder) -> $crate::query_builder::BuildQueryResult {
                try!(out.push_identifier($($table)::*::name()));
                out.push_sql(".");
                out.push_identifier(stringify!($column_name))
            }
        }

        impl $crate::expression::SelectableExpression<$($table)::*> for $column_name {}

        impl<'a, ST, Left, Right> SelectableExpression<
            $crate::WithQuerySource<'a, Left, Right>, ST> for $column_name where
            $column_name: SelectableExpression<Left, ST>
        {
        }

        impl $crate::expression::NonAggregate for $column_name {}

        impl $crate::query_source::Column for $column_name {
            type Table = $($table)::*;

            fn name() -> &'static str {
                stringify!($column_name)
            }
        }
    }
}

/// Specifies that a table exists, and what columns it has. This will create a
/// new public module, with the same name, as the name of the table. In this
/// module, you'll find a unit struct named `table`, and a unit struct with the
/// names of each of the columns. In the definition, you can also specify an
/// additional set of columns which exist, but should not be selected by default
/// (for example, for things like full text search)
///
/// Example usage
/// -------------
///
/// ```rust
/// # #[macro_use] extern crate diesel;
/// table! {
///     users {
///         id -> Integer,
///         name -> VarChar,
///         favorite_color -> Nullable<VarChar>,
///     }
/// }
/// # fn main() {}
/// ```
///
/// More complex usage:
///
/// ```rust
/// # #[macro_use] extern crate diesel;
/// table! {
///     users (non_standard_primary_key) {
///         non_standard_primary_key -> Integer,
///         name -> VarChar,
///         favorite_color -> Nullable<VarChar>,
///     } no select {
///         complex_index_column -> Text,
///     }
/// }
/// # fn main() {}
/// ```
///
/// This module will also contain several helper types:
///
/// dsl
/// ---
///
/// This simply re-exports the table, renamed to the same name as the module,
/// and each of the columns. This is useful to glob import when you're dealing
/// primarily with one table, to allow writing `users.filter(name.eq("Sean"))`
/// instead of `users::table.filter(users::name.eq("Sean"))`.
///
/// all_columns
/// -----------
///
/// A constant will be assigned called `all_columns`, which will be a tuple of
/// all the columns that aren't in the "no select" group. This is what will be
/// selected if you don't otherwise specify a select clause. It's type will be
/// `table::AllColumns`. You can also get this value from the
/// `Table::all_columns` function.
///
/// star
/// ----
///
/// This will be the qualified "star" expression for this table (e.g.
/// `users.*`). Internally, we read columns by index, not by name, so this
/// column is not safe to read data out of, and it has had it's SQL type set to
/// `()` to prevent accidentally using it as such. It is sometimes useful for
/// count statements however. It can also be accessed through the `Table.star()`
/// method.
#[macro_export]
macro_rules! table {
    (
        $name:ident {
            $($column_name:ident -> $Type:ty,)+
        }
    ) => {
        table! {
            $name (id) {
                $($column_name -> $Type,)+
            }
        }
    };
    (
        $name:ident ($pk:ident) {
            $($column_name:ident -> $Type:ty,)+
        }
    ) => {
        table! {
            $name ($pk) {
                $($column_name -> $Type,)+
            } no select {}
        }
    };
    (
        $name:ident {
            $($column_name:ident -> $Type:ty,)+
        } no select {
            $($no_select_column_name:ident -> $no_select_type:ty,)*
        }
    ) => {
        table! {
            $name (id) {
                $($column_name -> $Type,)+
            } no select {
                $($no_select_column_name -> $no_select_type,)*
            }
        }
    };
    (
        $name:ident ($pk:ident) {
            $($column_name:ident -> $Type:ty,)+
        } no select {
            $($no_select_column_name:ident -> $no_select_type:ty,)*
        }
    ) => {
        table_body! {
            $name ($pk) {
                $($column_name -> $Type,)+
            } no select {
                $($no_select_column_name -> $no_select_type,)*
            }
        }
    }
}

#[macro_export]
#[doc(hidden)]
macro_rules! table_body {
    (
        $name:ident ($pk:ident) {
            $($column_name:ident -> $Type:ty,)+
        } no select {
            $($no_select_column_name:ident -> $no_select_type:ty,)*
        }
    ) => {
        pub mod $name {
            use $crate::{
                QuerySource,
                Table,
            };
            use $crate::query_builder::*;
            use $crate::query_builder::nodes::Identifier;
            use $crate::types::*;
            pub use self::columns::*;

            pub mod dsl {
                pub use super::columns::{$($column_name),+};
                pub use super::table as $name;
            }

            #[allow(non_upper_case_globals, dead_code)]
            pub const all_columns: ($($column_name),+) = ($($column_name),+);

            #[allow(non_camel_case_types)]
            #[derive(Clone, Copy)]
            pub struct table;

            impl table {
                #[allow(dead_code)]
                pub fn star(&self) -> star {
                    star
                }
            }

            pub type SqlType = ($($Type),+);

            impl QuerySource for table {
                type FromClause = Identifier<'static>;

                fn from_clause(&self) -> Self::FromClause {
                    Identifier(stringify!($name))
                }
            }

            impl AsQuery for table {
                type SqlType = SqlType;
                type Query = SelectStatement<SqlType, ($($column_name),+), Self>;

                fn as_query(self) -> Self::Query {
                    SelectStatement::simple(all_columns, self)
                }
            }

            impl Table for table {
                type PrimaryKey = columns::$pk;
                type AllColumns = ($($column_name),+);

                fn name() -> &'static str {
                    stringify!($name)
                }

                fn primary_key(&self) -> Self::PrimaryKey {
                    columns::$pk
                }

                fn all_columns() -> Self::AllColumns {
                    ($($column_name),+)
                }
            }

            pub mod columns {
                use super::table;
                use $crate::{Table, Column, Expression, SelectableExpression};
                use $crate::backend::Backend;
                use $crate::query_builder::{QueryBuilder, BuildQueryResult, QueryFragment};
                use $crate::types::*;

                #[allow(non_camel_case_types, dead_code)]
                #[derive(Debug, Clone, Copy)]
                pub struct star;

                impl Expression for star {
                    type SqlType = ();
                }

                impl<DB: Backend> QueryFragment<DB> for star {
                    fn to_sql(&self, out: &mut DB::QueryBuilder) -> BuildQueryResult {
                        try!(out.push_identifier(table::name()));
                        out.push_sql(".*");
                        Ok(())
                    }
                }

                impl SelectableExpression<table> for star {}

                $(column!(table, $column_name -> $Type);)+
                $(column!(table, $no_select_column_name -> $no_select_type);)*
            }
        }
    }
}

#[macro_export]
#[doc(hidden)]
macro_rules! joinable {
    ($child:ident -> $parent:ident ($source:ident)) => {
        joinable_inner!($child::table => $parent::table : ($child::$source = $parent::table));
        joinable_inner!($parent::table => $child::table : ($child::$source = $parent::table));
    }
}

#[macro_export]
#[doc(hidden)]
macro_rules! joinable_inner {
    ($left_table:path => $right_table:path : ($foreign_key:path = $parent_table:path)) => {
        impl<JoinType> $crate::JoinTo<$right_table, JoinType> for $left_table {
            type JoinClause = $crate::query_builder::nodes::Join<
                <$left_table as $crate::QuerySource>::FromClause,
                <$right_table as $crate::QuerySource>::FromClause,
                $crate::expression::helper_types::Eq<
                    $crate::expression::nullable::Nullable<$foreign_key>,
                    $crate::expression::nullable::Nullable<
                        <$parent_table as $crate::query_source::Table>::PrimaryKey>,
                >,
                JoinType,
            >;

            fn join_clause(&self, join_type: JoinType) -> Self::JoinClause {
                use $crate::QuerySource;

                $crate::query_builder::nodes::Join::new(
                    self.from_clause(),
                    $right_table.from_clause(),
                    $foreign_key.nullable().eq($parent_table.primary_key().nullable()),
                    join_type,
                )
            }
        }
    }
}

#[macro_export]
#[doc(hidden)]
macro_rules! select_column_workaround {
    ($parent:ident -> $child:ident ($($column_name:ident),+)) => {
        $(select_column_inner!($parent -> $child $column_name);)+
        select_column_inner!($parent -> $child star);
    }
}

#[macro_export]
#[doc(hidden)]
macro_rules! select_column_inner {
    ($parent:ident -> $child:ident $column_name:ident) => {
        impl $crate::expression::SelectableExpression<
            $crate::query_source::InnerJoinSource<$child::table, $parent::table>,
        > for $parent::$column_name
        {
        }

        impl $crate::expression::SelectableExpression<
            $crate::query_source::InnerJoinSource<$parent::table, $child::table>,
        > for $parent::$column_name
        {
        }

        impl $crate::expression::SelectableExpression<
            $crate::query_source::LeftOuterJoinSource<$child::table, $parent::table>,
            <<$parent::$column_name as $crate::Expression>::SqlType
                as $crate::types::IntoNullable>::Nullable,
        > for $parent::$column_name
        {
        }

        impl $crate::expression::SelectableExpression<
            $crate::query_source::LeftOuterJoinSource<$parent::table, $child::table>,
        > for $parent::$column_name
        {
        }
    }
}

#[macro_export]
#[doc(hidden)]
macro_rules! join_through {
    ($parent:ident -> $through:ident -> $child:ident) => {
        impl<JoinType: Copy> $crate::JoinTo<$child::table, JoinType> for $parent::table {
            type JoinClause = <
                <$parent::table as $crate::JoinTo<$through::table, JoinType>>::JoinClause
                as $crate::query_builder::nodes::CombinedJoin<
                    <$through::table as $crate::JoinTo<$child::table, JoinType>>::JoinClause,
                >>::Output;

            fn join_clause(&self, join_type: JoinType) -> Self::JoinClause {
                use $crate::query_builder::nodes::CombinedJoin;
                let parent_to_through = $crate::JoinTo::<$through::table, JoinType>
                    ::join_clause(&$parent::table, join_type);
                let through_to_child = $crate::JoinTo::<$child::table, JoinType>
                    ::join_clause(&$through::table, join_type);
                parent_to_through.combine_with(through_to_child)
            }
        }
    }
}

/// Takes a query QueryFragment expression as an argument and returns a string
/// of SQL with placeholders for the dynamic values.
///
/// # Example
///
/// ### Returning SQL from a count statment:
/// #
/// # ```rust
/// # // example requires setup for users table
/// # use diesel::users::dsl::*;
/// # use diesel::query_builder::QueryFragment;
/// #
/// # fn main() {
/// let sql = debug_sql!(users.count());
/// assert_eq!(sql, "SELECT COUNT(*) FROM users");
/// # }
/// # ```
#[macro_export]
macro_rules! debug_sql {
    ($query:expr) => {{
        use $crate::query_builder::QueryFragment;
        use $crate::query_builder::debug::DebugQueryBuilder;
        let mut query_builder = DebugQueryBuilder::new();
        QueryFragment::<$crate::backend::Debug>::to_sql(&$query, &mut query_builder).unwrap();
        query_builder.sql
    }};
}

/// Takes takes a query QueryFragment expression as an argument and prints out
/// the SQL with placeholders for the dynamic values.
///
/// # Example
///
/// ### Printing SQL from a count statment:
/// #
/// # ```rust
/// # // example requires setup for users table
/// # use diesel::users::dsl::*;
/// # use diesel::query_builder::QueryFragment;
/// #
/// # fn main() {
/// print_sql!(users.count());
/// # }
/// # ```
#[macro_export]
macro_rules! print_sql {
    ($query:expr) => {
        println!("{}", &debug_sql!($query));
    };
}
<commit_msg>Fix documentation for `debug_sql!` and `print_sql!`<commit_after>// FIXME(https://github.com/rust-lang/rust/issues/19630) Remove this work-around
#[macro_export]
#[doc(hidden)]
macro_rules! diesel_internal_expr_conversion {
    ($e:expr) => { $e }
}

#[macro_export]
#[doc(hidden)]
macro_rules! column {
    ($($table:ident)::*, $column_name:ident -> $Type:ty) => {
        #[allow(non_camel_case_types, dead_code)]
        #[derive(Debug, Clone, Copy)]
        pub struct $column_name;

        impl $crate::expression::Expression for $column_name {
            type SqlType = $Type;
        }

        impl<DB> $crate::query_builder::QueryFragment<DB> for $column_name where
            DB: $crate::backend::Backend,
        {
            fn to_sql(&self, out: &mut DB::QueryBuilder) -> $crate::query_builder::BuildQueryResult {
                try!(out.push_identifier($($table)::*::name()));
                out.push_sql(".");
                out.push_identifier(stringify!($column_name))
            }
        }

        impl $crate::expression::SelectableExpression<$($table)::*> for $column_name {}

        impl<'a, ST, Left, Right> SelectableExpression<
            $crate::WithQuerySource<'a, Left, Right>, ST> for $column_name where
            $column_name: SelectableExpression<Left, ST>
        {
        }

        impl $crate::expression::NonAggregate for $column_name {}

        impl $crate::query_source::Column for $column_name {
            type Table = $($table)::*;

            fn name() -> &'static str {
                stringify!($column_name)
            }
        }
    }
}

/// Specifies that a table exists, and what columns it has. This will create a
/// new public module, with the same name, as the name of the table. In this
/// module, you'll find a unit struct named `table`, and a unit struct with the
/// names of each of the columns. In the definition, you can also specify an
/// additional set of columns which exist, but should not be selected by default
/// (for example, for things like full text search)
///
/// Example usage
/// -------------
///
/// ```rust
/// # #[macro_use] extern crate diesel;
/// table! {
///     users {
///         id -> Integer,
///         name -> VarChar,
///         favorite_color -> Nullable<VarChar>,
///     }
/// }
/// # fn main() {}
/// ```
///
/// More complex usage:
///
/// ```rust
/// # #[macro_use] extern crate diesel;
/// table! {
///     users (non_standard_primary_key) {
///         non_standard_primary_key -> Integer,
///         name -> VarChar,
///         favorite_color -> Nullable<VarChar>,
///     } no select {
///         complex_index_column -> Text,
///     }
/// }
/// # fn main() {}
/// ```
///
/// This module will also contain several helper types:
///
/// dsl
/// ---
///
/// This simply re-exports the table, renamed to the same name as the module,
/// and each of the columns. This is useful to glob import when you're dealing
/// primarily with one table, to allow writing `users.filter(name.eq("Sean"))`
/// instead of `users::table.filter(users::name.eq("Sean"))`.
///
/// all_columns
/// -----------
///
/// A constant will be assigned called `all_columns`, which will be a tuple of
/// all the columns that aren't in the "no select" group. This is what will be
/// selected if you don't otherwise specify a select clause. It's type will be
/// `table::AllColumns`. You can also get this value from the
/// `Table::all_columns` function.
///
/// star
/// ----
///
/// This will be the qualified "star" expression for this table (e.g.
/// `users.*`). Internally, we read columns by index, not by name, so this
/// column is not safe to read data out of, and it has had it's SQL type set to
/// `()` to prevent accidentally using it as such. It is sometimes useful for
/// count statements however. It can also be accessed through the `Table.star()`
/// method.
#[macro_export]
macro_rules! table {
    (
        $name:ident {
            $($column_name:ident -> $Type:ty,)+
        }
    ) => {
        table! {
            $name (id) {
                $($column_name -> $Type,)+
            }
        }
    };
    (
        $name:ident ($pk:ident) {
            $($column_name:ident -> $Type:ty,)+
        }
    ) => {
        table! {
            $name ($pk) {
                $($column_name -> $Type,)+
            } no select {}
        }
    };
    (
        $name:ident {
            $($column_name:ident -> $Type:ty,)+
        } no select {
            $($no_select_column_name:ident -> $no_select_type:ty,)*
        }
    ) => {
        table! {
            $name (id) {
                $($column_name -> $Type,)+
            } no select {
                $($no_select_column_name -> $no_select_type,)*
            }
        }
    };
    (
        $name:ident ($pk:ident) {
            $($column_name:ident -> $Type:ty,)+
        } no select {
            $($no_select_column_name:ident -> $no_select_type:ty,)*
        }
    ) => {
        table_body! {
            $name ($pk) {
                $($column_name -> $Type,)+
            } no select {
                $($no_select_column_name -> $no_select_type,)*
            }
        }
    }
}

#[macro_export]
#[doc(hidden)]
macro_rules! table_body {
    (
        $name:ident ($pk:ident) {
            $($column_name:ident -> $Type:ty,)+
        } no select {
            $($no_select_column_name:ident -> $no_select_type:ty,)*
        }
    ) => {
        pub mod $name {
            use $crate::{
                QuerySource,
                Table,
            };
            use $crate::query_builder::*;
            use $crate::query_builder::nodes::Identifier;
            use $crate::types::*;
            pub use self::columns::*;

            pub mod dsl {
                pub use super::columns::{$($column_name),+};
                pub use super::table as $name;
            }

            #[allow(non_upper_case_globals, dead_code)]
            pub const all_columns: ($($column_name),+) = ($($column_name),+);

            #[allow(non_camel_case_types)]
            #[derive(Clone, Copy)]
            pub struct table;

            impl table {
                #[allow(dead_code)]
                pub fn star(&self) -> star {
                    star
                }
            }

            pub type SqlType = ($($Type),+);

            impl QuerySource for table {
                type FromClause = Identifier<'static>;

                fn from_clause(&self) -> Self::FromClause {
                    Identifier(stringify!($name))
                }
            }

            impl AsQuery for table {
                type SqlType = SqlType;
                type Query = SelectStatement<SqlType, ($($column_name),+), Self>;

                fn as_query(self) -> Self::Query {
                    SelectStatement::simple(all_columns, self)
                }
            }

            impl Table for table {
                type PrimaryKey = columns::$pk;
                type AllColumns = ($($column_name),+);

                fn name() -> &'static str {
                    stringify!($name)
                }

                fn primary_key(&self) -> Self::PrimaryKey {
                    columns::$pk
                }

                fn all_columns() -> Self::AllColumns {
                    ($($column_name),+)
                }
            }

            pub mod columns {
                use super::table;
                use $crate::{Table, Column, Expression, SelectableExpression};
                use $crate::backend::Backend;
                use $crate::query_builder::{QueryBuilder, BuildQueryResult, QueryFragment};
                use $crate::types::*;

                #[allow(non_camel_case_types, dead_code)]
                #[derive(Debug, Clone, Copy)]
                pub struct star;

                impl Expression for star {
                    type SqlType = ();
                }

                impl<DB: Backend> QueryFragment<DB> for star {
                    fn to_sql(&self, out: &mut DB::QueryBuilder) -> BuildQueryResult {
                        try!(out.push_identifier(table::name()));
                        out.push_sql(".*");
                        Ok(())
                    }
                }

                impl SelectableExpression<table> for star {}

                $(column!(table, $column_name -> $Type);)+
                $(column!(table, $no_select_column_name -> $no_select_type);)*
            }
        }
    }
}

#[macro_export]
#[doc(hidden)]
macro_rules! joinable {
    ($child:ident -> $parent:ident ($source:ident)) => {
        joinable_inner!($child::table => $parent::table : ($child::$source = $parent::table));
        joinable_inner!($parent::table => $child::table : ($child::$source = $parent::table));
    }
}

#[macro_export]
#[doc(hidden)]
macro_rules! joinable_inner {
    ($left_table:path => $right_table:path : ($foreign_key:path = $parent_table:path)) => {
        impl<JoinType> $crate::JoinTo<$right_table, JoinType> for $left_table {
            type JoinClause = $crate::query_builder::nodes::Join<
                <$left_table as $crate::QuerySource>::FromClause,
                <$right_table as $crate::QuerySource>::FromClause,
                $crate::expression::helper_types::Eq<
                    $crate::expression::nullable::Nullable<$foreign_key>,
                    $crate::expression::nullable::Nullable<
                        <$parent_table as $crate::query_source::Table>::PrimaryKey>,
                >,
                JoinType,
            >;

            fn join_clause(&self, join_type: JoinType) -> Self::JoinClause {
                use $crate::QuerySource;

                $crate::query_builder::nodes::Join::new(
                    self.from_clause(),
                    $right_table.from_clause(),
                    $foreign_key.nullable().eq($parent_table.primary_key().nullable()),
                    join_type,
                )
            }
        }
    }
}

#[macro_export]
#[doc(hidden)]
macro_rules! select_column_workaround {
    ($parent:ident -> $child:ident ($($column_name:ident),+)) => {
        $(select_column_inner!($parent -> $child $column_name);)+
        select_column_inner!($parent -> $child star);
    }
}

#[macro_export]
#[doc(hidden)]
macro_rules! select_column_inner {
    ($parent:ident -> $child:ident $column_name:ident) => {
        impl $crate::expression::SelectableExpression<
            $crate::query_source::InnerJoinSource<$child::table, $parent::table>,
        > for $parent::$column_name
        {
        }

        impl $crate::expression::SelectableExpression<
            $crate::query_source::InnerJoinSource<$parent::table, $child::table>,
        > for $parent::$column_name
        {
        }

        impl $crate::expression::SelectableExpression<
            $crate::query_source::LeftOuterJoinSource<$child::table, $parent::table>,
            <<$parent::$column_name as $crate::Expression>::SqlType
                as $crate::types::IntoNullable>::Nullable,
        > for $parent::$column_name
        {
        }

        impl $crate::expression::SelectableExpression<
            $crate::query_source::LeftOuterJoinSource<$parent::table, $child::table>,
        > for $parent::$column_name
        {
        }
    }
}

#[macro_export]
#[doc(hidden)]
macro_rules! join_through {
    ($parent:ident -> $through:ident -> $child:ident) => {
        impl<JoinType: Copy> $crate::JoinTo<$child::table, JoinType> for $parent::table {
            type JoinClause = <
                <$parent::table as $crate::JoinTo<$through::table, JoinType>>::JoinClause
                as $crate::query_builder::nodes::CombinedJoin<
                    <$through::table as $crate::JoinTo<$child::table, JoinType>>::JoinClause,
                >>::Output;

            fn join_clause(&self, join_type: JoinType) -> Self::JoinClause {
                use $crate::query_builder::nodes::CombinedJoin;
                let parent_to_through = $crate::JoinTo::<$through::table, JoinType>
                    ::join_clause(&$parent::table, join_type);
                let through_to_child = $crate::JoinTo::<$child::table, JoinType>
                    ::join_clause(&$through::table, join_type);
                parent_to_through.combine_with(through_to_child)
            }
        }
    }
}

/// Takes a query QueryFragment expression as an argument and returns a string
/// of SQL with placeholders for the dynamic values.
///
/// # Example
///
/// ### Returning SQL from a count statment:
///
/// ```rust
/// # #[macro_use] extern crate diesel;
/// # table! {
/// #     users {
/// #         id -> Timestamp,
/// #         n -> Integer,
/// #     }
/// # }
/// # // example requires setup for users table
/// # use self::users::dsl::*;
/// # use diesel::*;
/// #
/// # fn main() {
/// let sql = debug_sql!(users.count());
/// assert_eq!(sql, "SELECT COUNT(*) FROM `users`");
/// # }
/// ```
#[macro_export]
macro_rules! debug_sql {
    ($query:expr) => {{
        use $crate::query_builder::QueryFragment;
        use $crate::query_builder::debug::DebugQueryBuilder;
        let mut query_builder = DebugQueryBuilder::new();
        QueryFragment::<$crate::backend::Debug>::to_sql(&$query, &mut query_builder).unwrap();
        query_builder.sql
    }};
}

/// Takes takes a query QueryFragment expression as an argument and prints out
/// the SQL with placeholders for the dynamic values.
///
/// # Example
///
/// ### Printing SQL from a count statment:
///
/// ```rust
/// # #[macro_use] extern crate diesel;
/// # table! {
/// #     users {
/// #         id -> Timestamp,
/// #         n -> Integer,
/// #     }
/// # }
/// # // example requires setup for users table
/// # use self::users::dsl::*;
/// # use diesel::*;
/// #
/// # fn main() {
/// print_sql!(users.count());
/// # }
/// ```
#[macro_export]
macro_rules! print_sql {
    ($query:expr) => {
        println!("{}", &debug_sql!($query));
    };
}
<|endoftext|>


<commit_before>use time::{SteadyTime, Duration};
use std::collections::BTreeMap;
use std::cmp::Ordering;
use std::fmt;
use std::fmt::Debug;

#[derive(Clone)]
struct Task<Token> where Token: Clone {
    interval: Duration,
    run_offset: SteadyTime,
    token: Token,
    bond: TaskBond
}

#[derive(Clone, Debug)]
enum TaskBond {
    OneOff,
    Perpetual
}

impl<Token> fmt::Debug for Task<Token> where Token: Clone {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.bond {
            TaskBond::Perpetual => write!(f, "Task(every {}ms)", self.interval.num_milliseconds()),
            TaskBond::OneOff => write!(f, "Task(after {}ms)", self.interval.num_milliseconds())
        }
    }
}

impl<Token> Task<Token> where Token: Clone {
    fn new(interval: Duration, run_offset: SteadyTime, bond: TaskBond, token: Token) -> Task<Token> {
        assert!(interval > Duration::seconds(0)); // negative interval would make schedule go back in time!
        Task {
            interval: interval,
            run_offset: run_offset,
            bond: bond,
            token: token
        }
    }

    // TODO: should that return new task (immutable?)
    fn next_schedule(&mut self) -> SteadyTime {
        self.run_offset = self.run_offset + self.interval;
        self.run_offset // ?
    }
}

pub trait TimeSource {
    fn now(&self) -> SteadyTime;
    fn wait(&mut self, duration: Duration);
}

// TODO: should that be u64?
type TimePoint = u32;

pub enum SchedulerAction {
    None,
    Wait(Duration),
    Skip(Vec<TimePoint>),
    Yield(TimePoint)
}

pub enum Schedule<Token> {
    Missed(Vec<Token>),
    Current(Vec<Token>)
}

pub struct Scheduler<Token, TS> where TS: TimeSource, Token: Clone {
    offset: SteadyTime,
    time_point_interval: Duration,
    tasks: BTreeMap<TimePoint, Vec<Task<Token>>>,
    time_source: TS
}

impl<Token, TS> Scheduler<Token, TS> where TS: TimeSource, Token: Clone {
    pub fn new(time_point_interval: Duration, time_source: TS) -> Scheduler<Token, TS> where TS: TimeSource {
        Scheduler {
            offset: time_source.now(),
            time_point_interval: time_point_interval,
            tasks: BTreeMap::new(),
            time_source: time_source
        }
    }

    fn schedule(&mut self, mut task: Task<Token>) {
        let now = self.time_source.now();
        assert!(now > self.offset);
        let current_time_point = self.to_time_point(now - self.offset);

        //TODO: optimise
        //TODO: this should not loop - just schedule for next and let scheduler to report Missed
        let mut time_point;
        loop {
            time_point = self.to_time_point(task.next_schedule() - now);
            if time_point > current_time_point {
                break;
            }
        }

        self.tasks.entry(time_point).or_insert(Vec::new()).push(task);
    }

    pub fn after(&mut self, duration: Duration, token: Token) {
        self.schedule(Task::new(duration, SteadyTime::now(), TaskBond::OneOff, token));
    }

    pub fn every(&mut self, duration: Duration, token: Token) {
        self.schedule(Task::new(duration, SteadyTime::now(), TaskBond::Perpetual, token));
    }

    // TODO: we may miss some time points if wait waits longer
    fn next_action(&self) -> SchedulerAction {
        let now = self.time_source.now();
        let current_time_point = self.to_time_point(now - self.offset);

        match self.tasks.iter().next() {
            None => SchedulerAction::None,
            Some((&time_point, _)) => {
                match time_point.cmp(&current_time_point) {
                    Ordering::Greater => SchedulerAction::Wait((self.offset + self.to_duration(time_point)) - now),
                    //TODO: make task Eq for sorting by time
                    Ordering::Less => SchedulerAction::Skip(self.tasks.iter().take_while(|&(&time_point, &_)| time_point < current_time_point).map(|(time_point, _)| time_point.clone()).collect()),
                    Ordering::Equal => SchedulerAction::Yield(time_point)
                }
            }
        }
    }

    pub fn wait(&mut self) -> Option<Schedule<Token>> {
        match self.next_action() {
            SchedulerAction::None => None,
            SchedulerAction::Wait(duration) => {
                self.time_source.wait(duration);
                self.wait()
            },
            SchedulerAction::Skip(time_points) => {
                let tasks: Vec<Task<Token>> = time_points.iter().flat_map(|time_point|
                        self.tasks.remove(&time_point).unwrap()
                    ).collect();
                for task in tasks.iter() {
                    match task.bond {
                        TaskBond::Perpetual => self.schedule(task.clone()),
                        TaskBond::OneOff => ()
                    };
                }
                Some(Schedule::Missed(tasks.into_iter().map(|task| task.token).collect()))
            },
            SchedulerAction::Yield(time_point) => {
                let mut tasks = self.tasks.remove(&time_point).unwrap();
                for task in tasks.iter() {
                    match task.bond {
                        TaskBond::Perpetual => self.schedule(task.clone()),
                        TaskBond::OneOff => ()
                    };
                }

                tasks.sort_by(|a, b| a.run_offset.cmp(&b.run_offset));

                Some(Schedule::Current(tasks.into_iter().map(|task| task.token).collect()))
            }
        }
    }

    #[allow(dead_code)]
    pub fn time_source(&self) -> &TimeSource {
        &self.time_source
    }

    fn to_time_point(&self, duration: Duration) -> TimePoint {
        let interval = self.time_point_interval.num_microseconds().unwrap();
        let duration = duration.num_microseconds().unwrap();
        assert!(duration >= 0);

        let time_point = duration / interval;
        if duration % interval != 0 {
            time_point + 1; // ceil
        }
        time_point as TimePoint
    }

    fn to_duration(&self, time_point: TimePoint) -> Duration {
        self.time_point_interval * (time_point as i32)
    }
}

#[cfg(test)]
mod task {
    use time::{SteadyTime, Duration};
    use super::{Task, TaskBond};

    #[test]
    fn getting_next_schedule() {
        let now = SteadyTime::now();
        let interval = Duration::seconds(1);
        let mut task = Task::new(interval, now, TaskBond::OneOff, 42);
        println!("task: {:?}", task);
        assert_eq!(task.next_schedule(), now + interval);
        assert_eq!(task.next_schedule(), now + interval * 2);
        assert_eq!(task.next_schedule(), now + interval * 3);
    }
}

<commit_msg>task is now immutable<commit_after>use time::{SteadyTime, Duration};
use std::collections::BTreeMap;
use std::cmp::Ordering;
use std::fmt;
use std::fmt::Debug;

#[derive(Clone)]
struct Task<Token> where Token: Clone {
    interval: Duration,
    run_offset: SteadyTime,
    token: Token,
    bond: TaskBond
}

#[derive(Clone, Debug)]
enum TaskBond {
    OneOff,
    Perpetual
}

impl<Token> fmt::Debug for Task<Token> where Token: Clone {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.bond {
            TaskBond::Perpetual => write!(f, "Task(every {}ms)", self.interval.num_milliseconds()),
            TaskBond::OneOff => write!(f, "Task(after {}ms)", self.interval.num_milliseconds())
        }
    }
}

impl<Token> Task<Token> where Token: Clone {
    fn new(interval: Duration, run_offset: SteadyTime, bond: TaskBond, token: Token) -> Task<Token> {
        assert!(interval > Duration::seconds(0)); // negative interval would make schedule go back in time!
        Task {
            interval: interval,
            run_offset: run_offset,
            bond: bond,
            token: token
        }
    }

    fn next(self) -> Task<Token> {
        Task {
            run_offset: self.run_offset + self.interval,
            .. self
        }
    }

    fn schedule(&self) -> SteadyTime {
        self.run_offset + self.interval
    }
}

pub trait TimeSource {
    fn now(&self) -> SteadyTime;
    fn wait(&mut self, duration: Duration);
}

// TODO: should that be u64?
type TimePoint = u32;

pub enum SchedulerAction {
    None,
    Wait(Duration),
    Skip(Vec<TimePoint>),
    Yield(TimePoint)
}

pub enum Schedule<Token> {
    Missed(Vec<Token>),
    Current(Vec<Token>)
}

pub struct Scheduler<Token, TS> where TS: TimeSource, Token: Clone {
    offset: SteadyTime,
    time_point_interval: Duration,
    tasks: BTreeMap<TimePoint, Vec<Task<Token>>>,
    time_source: TS
}

impl<Token, TS> Scheduler<Token, TS> where TS: TimeSource, Token: Clone {
    pub fn new(time_point_interval: Duration, time_source: TS) -> Scheduler<Token, TS> where TS: TimeSource {
        Scheduler {
            offset: time_source.now(),
            time_point_interval: time_point_interval,
            tasks: BTreeMap::new(),
            time_source: time_source
        }
    }

    fn schedule(&mut self, task: Task<Token>) {
        let now = self.time_source.now();
        assert!(now > self.offset);

        let task = task.next();
        let time_point = self.to_time_point(task.schedule() - now);

        self.tasks.entry(time_point).or_insert(Vec::new()).push(task);
    }

    pub fn after(&mut self, duration: Duration, token: Token) {
        self.schedule(Task::new(duration, SteadyTime::now(), TaskBond::OneOff, token));
    }

    pub fn every(&mut self, duration: Duration, token: Token) {
        self.schedule(Task::new(duration, SteadyTime::now(), TaskBond::Perpetual, token));
    }

    // TODO: we may miss some time points if wait waits longer
    fn next_action(&self) -> SchedulerAction {
        let now = self.time_source.now();
        let current_time_point = self.to_time_point(now - self.offset);

        match self.tasks.iter().next() {
            None => SchedulerAction::None,
            Some((&time_point, _)) => {
                match time_point.cmp(&current_time_point) {
                    Ordering::Greater => SchedulerAction::Wait((self.offset + self.to_duration(time_point)) - now),
                    //TODO: make task Eq for sorting by time
                    Ordering::Less => SchedulerAction::Skip(self.tasks.iter().take_while(|&(&time_point, &_)| time_point < current_time_point).map(|(time_point, _)| time_point.clone()).collect()),
                    Ordering::Equal => SchedulerAction::Yield(time_point)
                }
            }
        }
    }

    pub fn wait(&mut self) -> Option<Schedule<Token>> {
        match self.next_action() {
            SchedulerAction::None => None,
            SchedulerAction::Wait(duration) => {
                self.time_source.wait(duration);
                self.wait()
            },
            SchedulerAction::Skip(time_points) => {
                let tasks: Vec<Task<Token>> = time_points.iter().flat_map(|time_point|
                        self.tasks.remove(&time_point).unwrap()
                    ).collect();
                for task in tasks.iter() {
                    match task.bond {
                        TaskBond::Perpetual => self.schedule(task.clone()),
                        TaskBond::OneOff => ()
                    };
                }
                Some(Schedule::Missed(tasks.into_iter().map(|task| task.token).collect()))
            },
            SchedulerAction::Yield(time_point) => {
                let mut tasks = self.tasks.remove(&time_point).unwrap();
                for task in tasks.iter() {
                    match task.bond {
                        TaskBond::Perpetual => self.schedule(task.clone()),
                        TaskBond::OneOff => ()
                    };
                }

                tasks.sort_by(|a, b| a.run_offset.cmp(&b.run_offset));

                Some(Schedule::Current(tasks.into_iter().map(|task| task.token).collect()))
            }
        }
    }

    #[allow(dead_code)]
    pub fn time_source(&self) -> &TimeSource {
        &self.time_source
    }

    fn to_time_point(&self, duration: Duration) -> TimePoint {
        let interval = self.time_point_interval.num_microseconds().unwrap();
        let duration = duration.num_microseconds().unwrap();
        assert!(duration >= 0);

        let time_point = duration / interval;
        if duration % interval != 0 {
            time_point + 1; // ceil
        }
        time_point as TimePoint
    }

    fn to_duration(&self, time_point: TimePoint) -> Duration {
        self.time_point_interval * (time_point as i32)
    }
}

#[cfg(test)]
mod task {
    use time::{SteadyTime, Duration};
    use super::{Task, TaskBond};

    #[test]
    fn getting_next_schedule() {
        let now = SteadyTime::now();
        let interval = Duration::seconds(1);
        let task = Task::new(interval, now, TaskBond::OneOff, 42);
        println!("task: {:?}", task);
        assert_eq!(task.schedule(), now + interval);
        assert_eq!(task.next().next().schedule(), now + interval * 3);
    }
}

<|endoftext|>


<commit_before>use std::mem;

use bytes::{Bytes, BytesMut};
use futures::{Async, Poll};
use http::header::{self, HeaderName, HeaderValue};
use http::{HeaderMap, StatusCode, Version};
use httparse;

use error::{ParseError, PayloadError};

use server::h1decoder::{EncodingDecoder, HeaderIndex};
use server::IoStream;

use super::response::ClientMessage;
use super::ClientResponse;

const MAX_BUFFER_SIZE: usize = 131_072;
const MAX_HEADERS: usize = 96;

#[derive(Default)]
pub struct HttpResponseParser {
    decoder: Option<EncodingDecoder>,
}

#[derive(Debug, Fail)]
pub enum HttpResponseParserError {
    /// Server disconnected
    #[fail(display = "Server disconnected")]
    Disconnect,
    #[fail(display = "{}", _0)]
    Error(#[cause] ParseError),
}

impl HttpResponseParser {
    pub fn parse<T>(
        &mut self, io: &mut T, buf: &mut BytesMut,
    ) -> Poll<ClientResponse, HttpResponseParserError>
    where
        T: IoStream,
    {
        loop {
            // Read some more data into the buffer for the parser.
            match io.read_available(buf) {
                Ok(Async::Ready((false, true))) => {
                    return Err(HttpResponseParserError::Disconnect)
                }
                Ok(Async::Ready(_)) => (),
                Ok(Async::NotReady) => return Ok(Async::NotReady),
                Err(err) => {
                    return Err(HttpResponseParserError::Error(err.into()))
                }
            }

            // Call HTTP response parser.
            match HttpResponseParser::parse_message(buf)
                .map_err(HttpResponseParserError::Error)?
            {
                Async::Ready((msg, decoder)) => {
                    self.decoder = decoder;
                    return Ok(Async::Ready(msg));
                }
                Async::NotReady => {
                    if buf.capacity() >= MAX_BUFFER_SIZE {
                        return Err(HttpResponseParserError::Error(ParseError::TooLarge));
                    }
                    // Parser needs more data.  Loop and read more data.
                }
            }
        }
    }

    pub fn parse_payload<T>(
        &mut self, io: &mut T, buf: &mut BytesMut,
    ) -> Poll<Option<Bytes>, PayloadError>
    where
        T: IoStream,
    {
        if self.decoder.is_some() {
            loop {
                // read payload
                let (not_ready, stream_finished) = match io.read_available(buf) {
                    Ok(Async::Ready((_, true))) => (false, true),
                    Ok(Async::Ready((_, false))) => (false, false),
                    Ok(Async::NotReady) => (true, false),
                    Err(err) => return Err(err.into()),
                };

                match self.decoder.as_mut().unwrap().decode(buf) {
                    Ok(Async::Ready(Some(b))) => return Ok(Async::Ready(Some(b))),
                    Ok(Async::Ready(None)) => {
                        self.decoder.take();
                        return Ok(Async::Ready(None));
                    }
                    Ok(Async::NotReady) => {
                        if not_ready {
                            return Ok(Async::NotReady);
                        }
                        if stream_finished {
                            return Err(PayloadError::Incomplete);
                        }
                    }
                    Err(err) => return Err(err.into()),
                }
            }
        } else {
            Ok(Async::Ready(None))
        }
    }

    fn parse_message(
        buf: &mut BytesMut,
    ) -> Poll<(ClientResponse, Option<EncodingDecoder>), ParseError> {
        // Unsafe: we read only this data only after httparse parses headers into.
        // performance bump for pipeline benchmarks.
        let mut headers: [HeaderIndex; MAX_HEADERS] = unsafe { mem::uninitialized() };

        let (len, version, status, headers_len) = {
            let mut parsed: [httparse::Header; MAX_HEADERS] =
                unsafe { mem::uninitialized() };

            let mut resp = httparse::Response::new(&mut parsed);
            match resp.parse(buf)? {
                httparse::Status::Complete(len) => {
                    let version = if resp.version.unwrap_or(1) == 1 {
                        Version::HTTP_11
                    } else {
                        Version::HTTP_10
                    };
                    HeaderIndex::record(buf, resp.headers, &mut headers);
                    let status = StatusCode::from_u16(resp.code.unwrap())
                        .map_err(|_| ParseError::Status)?;

                    (len, version, status, resp.headers.len())
                }
                httparse::Status::Partial => return Ok(Async::NotReady),
            }
        };

        let slice = buf.split_to(len).freeze();

        // convert headers
        let mut hdrs = HeaderMap::new();
        for idx in headers[..headers_len].iter() {
            if let Ok(name) = HeaderName::from_bytes(&slice[idx.name.0..idx.name.1]) {
                // Unsafe: httparse check header value for valid utf-8
                let value = unsafe {
                    HeaderValue::from_shared_unchecked(
                        slice.slice(idx.value.0, idx.value.1),
                    )
                };
                hdrs.append(name, value);
            } else {
                return Err(ParseError::Header);
            }
        }

        let decoder = if status == StatusCode::SWITCHING_PROTOCOLS {
            Some(EncodingDecoder::eof())
        } else if let Some(len) = hdrs.get(header::CONTENT_LENGTH) {
            // Content-Length
            if let Ok(s) = len.to_str() {
                if let Ok(len) = s.parse::<u64>() {
                    Some(EncodingDecoder::length(len))
                } else {
                    debug!("illegal Content-Length: {:?}", len);
                    return Err(ParseError::Header);
                }
            } else {
                debug!("illegal Content-Length: {:?}", len);
                return Err(ParseError::Header);
            }
        } else if chunked(&hdrs)? {
            // Chunked encoding
            Some(EncodingDecoder::chunked())
        } else {
            None
        };

        if let Some(decoder) = decoder {
            Ok(Async::Ready((
                ClientResponse::new(ClientMessage {
                    status,
                    version,
                    headers: hdrs,
                    cookies: None,
                }),
                Some(decoder),
            )))
        } else {
            Ok(Async::Ready((
                ClientResponse::new(ClientMessage {
                    status,
                    version,
                    headers: hdrs,
                    cookies: None,
                }),
                None,
            )))
        }
    }
}

/// Check if request has chunked transfer encoding
pub fn chunked(headers: &HeaderMap) -> Result<bool, ParseError> {
    if let Some(encodings) = headers.get(header::TRANSFER_ENCODING) {
        if let Ok(s) = encodings.to_str() {
            Ok(s.to_lowercase().contains("chunked"))
        } else {
            Err(ParseError::Header)
        }
    } else {
        Ok(false)
    }
}
<commit_msg>If buffer is empty, read more data before calling parser.<commit_after>use std::mem;

use bytes::{Bytes, BytesMut};
use futures::{Async, Poll};
use http::header::{self, HeaderName, HeaderValue};
use http::{HeaderMap, StatusCode, Version};
use httparse;

use error::{ParseError, PayloadError};

use server::h1decoder::{EncodingDecoder, HeaderIndex};
use server::IoStream;

use super::response::ClientMessage;
use super::ClientResponse;

const MAX_BUFFER_SIZE: usize = 131_072;
const MAX_HEADERS: usize = 96;

#[derive(Default)]
pub struct HttpResponseParser {
    decoder: Option<EncodingDecoder>,
}

#[derive(Debug, Fail)]
pub enum HttpResponseParserError {
    /// Server disconnected
    #[fail(display = "Server disconnected")]
    Disconnect,
    #[fail(display = "{}", _0)]
    Error(#[cause] ParseError),
}

impl HttpResponseParser {
    pub fn parse<T>(
        &mut self, io: &mut T, buf: &mut BytesMut,
    ) -> Poll<ClientResponse, HttpResponseParserError>
    where
        T: IoStream,
    {
        loop {
            // Don't call parser until we have data to parse.
            if !buf.is_empty() {
                match HttpResponseParser::parse_message(buf)
                    .map_err(HttpResponseParserError::Error)?
                {
                    Async::Ready((msg, decoder)) => {
                        self.decoder = decoder;
                        return Ok(Async::Ready(msg));
                    }
                    Async::NotReady => {
                        if buf.capacity() >= MAX_BUFFER_SIZE {
                            return Err(HttpResponseParserError::Error(ParseError::TooLarge));
                        }
                        // Parser needs more data.
                    }
                }
            }
            // Read some more data into the buffer for the parser.
            match io.read_available(buf) {
                Ok(Async::Ready((false, true))) => {
                    return Err(HttpResponseParserError::Disconnect)
                }
                Ok(Async::Ready(_)) => (),
                Ok(Async::NotReady) => return Ok(Async::NotReady),
                Err(err) => {
                    return Err(HttpResponseParserError::Error(err.into()))
                }
            }
        }
    }

    pub fn parse_payload<T>(
        &mut self, io: &mut T, buf: &mut BytesMut,
    ) -> Poll<Option<Bytes>, PayloadError>
    where
        T: IoStream,
    {
        if self.decoder.is_some() {
            loop {
                // read payload
                let (not_ready, stream_finished) = match io.read_available(buf) {
                    Ok(Async::Ready((_, true))) => (false, true),
                    Ok(Async::Ready((_, false))) => (false, false),
                    Ok(Async::NotReady) => (true, false),
                    Err(err) => return Err(err.into()),
                };

                match self.decoder.as_mut().unwrap().decode(buf) {
                    Ok(Async::Ready(Some(b))) => return Ok(Async::Ready(Some(b))),
                    Ok(Async::Ready(None)) => {
                        self.decoder.take();
                        return Ok(Async::Ready(None));
                    }
                    Ok(Async::NotReady) => {
                        if not_ready {
                            return Ok(Async::NotReady);
                        }
                        if stream_finished {
                            return Err(PayloadError::Incomplete);
                        }
                    }
                    Err(err) => return Err(err.into()),
                }
            }
        } else {
            Ok(Async::Ready(None))
        }
    }

    fn parse_message(
        buf: &mut BytesMut,
    ) -> Poll<(ClientResponse, Option<EncodingDecoder>), ParseError> {
        // Unsafe: we read only this data only after httparse parses headers into.
        // performance bump for pipeline benchmarks.
        let mut headers: [HeaderIndex; MAX_HEADERS] = unsafe { mem::uninitialized() };

        let (len, version, status, headers_len) = {
            let mut parsed: [httparse::Header; MAX_HEADERS] =
                unsafe { mem::uninitialized() };

            let mut resp = httparse::Response::new(&mut parsed);
            match resp.parse(buf)? {
                httparse::Status::Complete(len) => {
                    let version = if resp.version.unwrap_or(1) == 1 {
                        Version::HTTP_11
                    } else {
                        Version::HTTP_10
                    };
                    HeaderIndex::record(buf, resp.headers, &mut headers);
                    let status = StatusCode::from_u16(resp.code.unwrap())
                        .map_err(|_| ParseError::Status)?;

                    (len, version, status, resp.headers.len())
                }
                httparse::Status::Partial => return Ok(Async::NotReady),
            }
        };

        let slice = buf.split_to(len).freeze();

        // convert headers
        let mut hdrs = HeaderMap::new();
        for idx in headers[..headers_len].iter() {
            if let Ok(name) = HeaderName::from_bytes(&slice[idx.name.0..idx.name.1]) {
                // Unsafe: httparse check header value for valid utf-8
                let value = unsafe {
                    HeaderValue::from_shared_unchecked(
                        slice.slice(idx.value.0, idx.value.1),
                    )
                };
                hdrs.append(name, value);
            } else {
                return Err(ParseError::Header);
            }
        }

        let decoder = if status == StatusCode::SWITCHING_PROTOCOLS {
            Some(EncodingDecoder::eof())
        } else if let Some(len) = hdrs.get(header::CONTENT_LENGTH) {
            // Content-Length
            if let Ok(s) = len.to_str() {
                if let Ok(len) = s.parse::<u64>() {
                    Some(EncodingDecoder::length(len))
                } else {
                    debug!("illegal Content-Length: {:?}", len);
                    return Err(ParseError::Header);
                }
            } else {
                debug!("illegal Content-Length: {:?}", len);
                return Err(ParseError::Header);
            }
        } else if chunked(&hdrs)? {
            // Chunked encoding
            Some(EncodingDecoder::chunked())
        } else {
            None
        };

        if let Some(decoder) = decoder {
            Ok(Async::Ready((
                ClientResponse::new(ClientMessage {
                    status,
                    version,
                    headers: hdrs,
                    cookies: None,
                }),
                Some(decoder),
            )))
        } else {
            Ok(Async::Ready((
                ClientResponse::new(ClientMessage {
                    status,
                    version,
                    headers: hdrs,
                    cookies: None,
                }),
                None,
            )))
        }
    }
}

/// Check if request has chunked transfer encoding
pub fn chunked(headers: &HeaderMap) -> Result<bool, ParseError> {
    if let Some(encodings) = headers.get(header::TRANSFER_ENCODING) {
        if let Ok(s) = encodings.to_str() {
            Ok(s.to_lowercase().contains("chunked"))
        } else {
            Err(ParseError::Header)
        }
    } else {
        Ok(false)
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(feature = "cargo-clippy", allow(redundant_field_names))]

use std::cell::RefCell;
use std::fmt::Write as FmtWrite;
use std::io::{self, Write};

#[cfg(feature = "brotli")]
use brotli2::write::BrotliEncoder;
use bytes::{BufMut, BytesMut};
#[cfg(feature = "flate2")]
use flate2::write::{DeflateEncoder, GzEncoder};
#[cfg(feature = "flate2")]
use flate2::Compression;
use futures::{Async, Poll};
use http::header::{
    HeaderValue, CONNECTION, CONTENT_ENCODING, CONTENT_LENGTH, DATE, TRANSFER_ENCODING,
};
use http::{HttpTryFrom, Version};
use time::{self, Duration};
use tokio_io::AsyncWrite;

use body::{Binary, Body};
use header::ContentEncoding;
use server::output::{ContentEncoder, Output, TransferEncoding};
use server::WriterState;

use client::ClientRequest;

const AVERAGE_HEADER_SIZE: usize = 30;

bitflags! {
    struct Flags: u8 {
        const STARTED = 0b0000_0001;
        const UPGRADE = 0b0000_0010;
        const KEEPALIVE = 0b0000_0100;
        const DISCONNECTED = 0b0000_1000;
    }
}

pub(crate) struct HttpClientWriter {
    flags: Flags,
    written: u64,
    headers_size: u32,
    buffer: Output,
    buffer_capacity: usize,
}

impl HttpClientWriter {
    pub fn new() -> HttpClientWriter {
        HttpClientWriter {
            flags: Flags::empty(),
            written: 0,
            headers_size: 0,
            buffer_capacity: 0,
            buffer: Output::Buffer(BytesMut::new()),
        }
    }

    pub fn disconnected(&mut self) {
        self.buffer.take();
    }

    pub fn is_completed(&self) -> bool {
        self.buffer.is_empty()
    }

    // pub fn keepalive(&self) -> bool {
    // self.flags.contains(Flags::KEEPALIVE) &&
    // !self.flags.contains(Flags::UPGRADE) }

    fn write_to_stream<T: AsyncWrite>(
        &mut self, stream: &mut T,
    ) -> io::Result<WriterState> {
        while !self.buffer.is_empty() {
            match stream.write(self.buffer.as_ref().as_ref()) {
                Ok(0) => {
                    self.disconnected();
                    return Ok(WriterState::Done);
                }
                Ok(n) => {
                    let _ = self.buffer.split_to(n);
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    if self.buffer.len() > self.buffer_capacity {
                        return Ok(WriterState::Pause);
                    } else {
                        return Ok(WriterState::Done);
                    }
                }
                Err(err) => return Err(err),
            }
        }
        Ok(WriterState::Done)
    }
}

pub struct Writer<'a>(pub &'a mut BytesMut);

impl<'a> io::Write for Writer<'a> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.extend_from_slice(buf);
        Ok(buf.len())
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl HttpClientWriter {
    pub fn start(&mut self, msg: &mut ClientRequest) -> io::Result<()> {
        // prepare task
        self.buffer = content_encoder(self.buffer.take(), msg);
        self.flags.insert(Flags::STARTED);
        if msg.upgrade() {
            self.flags.insert(Flags::UPGRADE);
        }

        // render message
        {
            // output buffer
            let buffer = self.buffer.as_mut();

            // status line
            writeln!(
                Writer(buffer),
                "{} {} {:?}\r",
                msg.method(),
                msg.uri()
                    .path_and_query()
                    .map(|u| u.as_str())
                    .unwrap_or("/"),
                msg.version()
            ).map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;

            // write headers
            if let Body::Binary(ref bytes) = *msg.body() {
                buffer.reserve(msg.headers().len() * AVERAGE_HEADER_SIZE + bytes.len());
            } else {
                buffer.reserve(msg.headers().len() * AVERAGE_HEADER_SIZE);
            }

            for (key, value) in msg.headers() {
                let v = value.as_ref();
                let k = key.as_str().as_bytes();
                buffer.reserve(k.len() + v.len() + 4);
                buffer.put_slice(k);
                buffer.put_slice(b": ");
                buffer.put_slice(v);
                buffer.put_slice(b"\r\n");
            }

            // set date header
            if !msg.headers().contains_key(DATE) {
                buffer.extend_from_slice(b"date: ");
                set_date(buffer);
                buffer.extend_from_slice(b"\r\n\r\n");
            } else {
                buffer.extend_from_slice(b"\r\n");
            }
        }
        self.headers_size = self.buffer.len() as u32;

        if msg.body().is_binary() {
            if let Body::Binary(bytes) = msg.replace_body(Body::Empty) {
                self.written += bytes.len() as u64;
                self.buffer.write(bytes.as_ref())?;
            }
        } else {
            self.buffer_capacity = msg.write_buffer_capacity();
        }
        Ok(())
    }

    pub fn write(&mut self, payload: &[u8]) -> io::Result<WriterState> {
        self.written += payload.len() as u64;
        if !self.flags.contains(Flags::DISCONNECTED) {
            self.buffer.write(payload)?;
        }

        if self.buffer.len() > self.buffer_capacity {
            Ok(WriterState::Pause)
        } else {
            Ok(WriterState::Done)
        }
    }

    pub fn write_eof(&mut self) -> io::Result<()> {
        if self.buffer.write_eof()? {
            Ok(())
        } else {
            Err(io::Error::new(
                io::ErrorKind::Other,
                "Last payload item, but eof is not reached",
            ))
        }
    }

    #[inline]
    pub fn poll_completed<T: AsyncWrite>(
        &mut self, stream: &mut T, shutdown: bool,
    ) -> Poll<(), io::Error> {
        match self.write_to_stream(stream) {
            Ok(WriterState::Done) => {
                if shutdown {
                    stream.shutdown()
                } else {
                    Ok(Async::Ready(()))
                }
            }
            Ok(WriterState::Pause) => Ok(Async::NotReady),
            Err(err) => Err(err),
        }
    }
}

fn content_encoder(buf: BytesMut, req: &mut ClientRequest) -> Output {
    let version = req.version();
    let mut body = req.replace_body(Body::Empty);
    let mut encoding = req.content_encoding();

    let transfer = match body {
        Body::Empty => {
            req.headers_mut().remove(CONTENT_LENGTH);
            return Output::Empty(buf);
        }
        Body::Binary(ref mut bytes) => {
            #[cfg(any(feature = "flate2", feature = "brotli"))]
            {
                if encoding.is_compression() {
                    let mut tmp = BytesMut::new();
                    let mut transfer = TransferEncoding::eof(tmp);
                    let mut enc = match encoding {
                        #[cfg(feature = "flate2")]
                        ContentEncoding::Deflate => ContentEncoder::Deflate(
                            DeflateEncoder::new(transfer, Compression::default()),
                        ),
                        #[cfg(feature = "flate2")]
                        ContentEncoding::Gzip => ContentEncoder::Gzip(GzEncoder::new(
                            transfer,
                            Compression::default(),
                        )),
                        #[cfg(feature = "brotli")]
                        ContentEncoding::Br => {
                            ContentEncoder::Br(BrotliEncoder::new(transfer, 5))
                        }
                        ContentEncoding::Auto | ContentEncoding::Identity => {
                            unreachable!()
                        }
                    };
                    // TODO return error!
                    let _ = enc.write(bytes.as_ref());
                    let _ = enc.write_eof();
                    *bytes = Binary::from(enc.buf_mut().take());

                    req.headers_mut().insert(
                        CONTENT_ENCODING,
                        HeaderValue::from_static(encoding.as_str()),
                    );
                    encoding = ContentEncoding::Identity;
                }
                let mut b = BytesMut::new();
                let _ = write!(b, "{}", bytes.len());
                req.headers_mut()
                    .insert(CONTENT_LENGTH, HeaderValue::try_from(b.freeze()).unwrap());
                TransferEncoding::eof(buf)
            }
            #[cfg(not(any(feature = "flate2", feature = "brotli")))]
            {
                TransferEncoding::eof(buf)
            }
        }
        Body::Streaming(_) | Body::Actor(_) => {
            if req.upgrade() {
                if version == Version::HTTP_2 {
                    error!("Connection upgrade is forbidden for HTTP/2");
                } else {
                    req.headers_mut()
                        .insert(CONNECTION, HeaderValue::from_static("upgrade"));
                }
                if encoding != ContentEncoding::Identity {
                    encoding = ContentEncoding::Identity;
                    req.headers_mut().remove(CONTENT_ENCODING);
                }
                TransferEncoding::eof(buf)
            } else {
                streaming_encoding(buf, version, req)
            }
        }
    };

    if encoding.is_compression() {
        req.headers_mut().insert(
            CONTENT_ENCODING,
            HeaderValue::from_static(encoding.as_str()),
        );
    }

    req.replace_body(body);
    let enc = match encoding {
        #[cfg(feature = "flate2")]
        ContentEncoding::Deflate => ContentEncoder::Deflate(DeflateEncoder::new(
            transfer,
            Compression::default(),
        )),
        #[cfg(feature = "flate2")]
        ContentEncoding::Gzip => {
            ContentEncoder::Gzip(GzEncoder::new(transfer, Compression::default()))
        }
        #[cfg(feature = "brotli")]
        ContentEncoding::Br => ContentEncoder::Br(BrotliEncoder::new(transfer, 5)),
        ContentEncoding::Identity | ContentEncoding::Auto => return Output::TE(transfer),
    };
    Output::Encoder(enc)
}

fn streaming_encoding(
    buf: BytesMut, version: Version, req: &mut ClientRequest,
) -> TransferEncoding {
    if req.chunked() {
        // Enable transfer encoding
        req.headers_mut().remove(CONTENT_LENGTH);
        if version == Version::HTTP_2 {
            req.headers_mut().remove(TRANSFER_ENCODING);
            TransferEncoding::eof(buf)
        } else {
            req.headers_mut()
                .insert(TRANSFER_ENCODING, HeaderValue::from_static("chunked"));
            TransferEncoding::chunked(buf)
        }
    } else {
        // if Content-Length is specified, then use it as length hint
        let (len, chunked) = if let Some(len) = req.headers().get(CONTENT_LENGTH) {
            // Content-Length
            if let Ok(s) = len.to_str() {
                if let Ok(len) = s.parse::<u64>() {
                    (Some(len), false)
                } else {
                    error!("illegal Content-Length: {:?}", len);
                    (None, false)
                }
            } else {
                error!("illegal Content-Length: {:?}", len);
                (None, false)
            }
        } else {
            (None, true)
        };

        if !chunked {
            if let Some(len) = len {
                TransferEncoding::length(len, buf)
            } else {
                TransferEncoding::eof(buf)
            }
        } else {
            // Enable transfer encoding
            match version {
                Version::HTTP_11 => {
                    req.headers_mut()
                        .insert(TRANSFER_ENCODING, HeaderValue::from_static("chunked"));
                    TransferEncoding::chunked(buf)
                }
                _ => {
                    req.headers_mut().remove(TRANSFER_ENCODING);
                    TransferEncoding::eof(buf)
                }
            }
        }
    }
}

// "Sun, 06 Nov 1994 08:49:37 GMT".len()
pub const DATE_VALUE_LENGTH: usize = 29;

fn set_date(dst: &mut BytesMut) {
    CACHED.with(|cache| {
        let mut cache = cache.borrow_mut();
        let now = time::get_time();
        if now > cache.next_update {
            cache.update(now);
        }
        dst.extend_from_slice(cache.buffer());
    })
}

struct CachedDate {
    bytes: [u8; DATE_VALUE_LENGTH],
    next_update: time::Timespec,
}

thread_local!(static CACHED: RefCell<CachedDate> = RefCell::new(CachedDate {
    bytes: [0; DATE_VALUE_LENGTH],
    next_update: time::Timespec::new(0, 0),
}));

impl CachedDate {
    fn buffer(&self) -> &[u8] {
        &self.bytes[..]
    }

    fn update(&mut self, now: time::Timespec) {
        write!(&mut self.bytes[..], "{}", time::at_utc(now).rfc822()).unwrap();
        self.next_update = now + Duration::seconds(1);
        self.next_update.nsec = 0;
    }
}
<commit_msg>fix missing content length<commit_after>#![cfg_attr(feature = "cargo-clippy", allow(redundant_field_names))]

use std::cell::RefCell;
use std::fmt::Write as FmtWrite;
use std::io::{self, Write};

#[cfg(feature = "brotli")]
use brotli2::write::BrotliEncoder;
use bytes::{BufMut, BytesMut};
#[cfg(feature = "flate2")]
use flate2::write::{DeflateEncoder, GzEncoder};
#[cfg(feature = "flate2")]
use flate2::Compression;
use futures::{Async, Poll};
use http::header::{
    HeaderValue, CONNECTION, CONTENT_ENCODING, CONTENT_LENGTH, DATE, TRANSFER_ENCODING,
};
use http::{HttpTryFrom, Version};
use time::{self, Duration};
use tokio_io::AsyncWrite;

use body::{Binary, Body};
use header::ContentEncoding;
use server::output::{ContentEncoder, Output, TransferEncoding};
use server::WriterState;

use client::ClientRequest;

const AVERAGE_HEADER_SIZE: usize = 30;

bitflags! {
    struct Flags: u8 {
        const STARTED = 0b0000_0001;
        const UPGRADE = 0b0000_0010;
        const KEEPALIVE = 0b0000_0100;
        const DISCONNECTED = 0b0000_1000;
    }
}

pub(crate) struct HttpClientWriter {
    flags: Flags,
    written: u64,
    headers_size: u32,
    buffer: Output,
    buffer_capacity: usize,
}

impl HttpClientWriter {
    pub fn new() -> HttpClientWriter {
        HttpClientWriter {
            flags: Flags::empty(),
            written: 0,
            headers_size: 0,
            buffer_capacity: 0,
            buffer: Output::Buffer(BytesMut::new()),
        }
    }

    pub fn disconnected(&mut self) {
        self.buffer.take();
    }

    pub fn is_completed(&self) -> bool {
        self.buffer.is_empty()
    }

    // pub fn keepalive(&self) -> bool {
    // self.flags.contains(Flags::KEEPALIVE) &&
    // !self.flags.contains(Flags::UPGRADE) }

    fn write_to_stream<T: AsyncWrite>(
        &mut self, stream: &mut T,
    ) -> io::Result<WriterState> {
        while !self.buffer.is_empty() {
            match stream.write(self.buffer.as_ref().as_ref()) {
                Ok(0) => {
                    self.disconnected();
                    return Ok(WriterState::Done);
                }
                Ok(n) => {
                    let _ = self.buffer.split_to(n);
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    if self.buffer.len() > self.buffer_capacity {
                        return Ok(WriterState::Pause);
                    } else {
                        return Ok(WriterState::Done);
                    }
                }
                Err(err) => return Err(err),
            }
        }
        Ok(WriterState::Done)
    }
}

pub struct Writer<'a>(pub &'a mut BytesMut);

impl<'a> io::Write for Writer<'a> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.extend_from_slice(buf);
        Ok(buf.len())
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl HttpClientWriter {
    pub fn start(&mut self, msg: &mut ClientRequest) -> io::Result<()> {
        // prepare task
        self.buffer = content_encoder(self.buffer.take(), msg);
        self.flags.insert(Flags::STARTED);
        if msg.upgrade() {
            self.flags.insert(Flags::UPGRADE);
        }

        // render message
        {
            // output buffer
            let buffer = self.buffer.as_mut();

            // status line
            writeln!(
                Writer(buffer),
                "{} {} {:?}\r",
                msg.method(),
                msg.uri()
                    .path_and_query()
                    .map(|u| u.as_str())
                    .unwrap_or("/"),
                msg.version()
            ).map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;

            // write headers
            if let Body::Binary(ref bytes) = *msg.body() {
                buffer.reserve(msg.headers().len() * AVERAGE_HEADER_SIZE + bytes.len());
            } else {
                buffer.reserve(msg.headers().len() * AVERAGE_HEADER_SIZE);
            }

            for (key, value) in msg.headers() {
                let v = value.as_ref();
                let k = key.as_str().as_bytes();
                buffer.reserve(k.len() + v.len() + 4);
                buffer.put_slice(k);
                buffer.put_slice(b": ");
                buffer.put_slice(v);
                buffer.put_slice(b"\r\n");
            }

            // set date header
            if !msg.headers().contains_key(DATE) {
                buffer.extend_from_slice(b"date: ");
                set_date(buffer);
                buffer.extend_from_slice(b"\r\n\r\n");
            } else {
                buffer.extend_from_slice(b"\r\n");
            }
        }
        self.headers_size = self.buffer.len() as u32;

        if msg.body().is_binary() {
            if let Body::Binary(bytes) = msg.replace_body(Body::Empty) {
                self.written += bytes.len() as u64;
                self.buffer.write(bytes.as_ref())?;
            }
        } else {
            self.buffer_capacity = msg.write_buffer_capacity();
        }
        Ok(())
    }

    pub fn write(&mut self, payload: &[u8]) -> io::Result<WriterState> {
        self.written += payload.len() as u64;
        if !self.flags.contains(Flags::DISCONNECTED) {
            self.buffer.write(payload)?;
        }

        if self.buffer.len() > self.buffer_capacity {
            Ok(WriterState::Pause)
        } else {
            Ok(WriterState::Done)
        }
    }

    pub fn write_eof(&mut self) -> io::Result<()> {
        if self.buffer.write_eof()? {
            Ok(())
        } else {
            Err(io::Error::new(
                io::ErrorKind::Other,
                "Last payload item, but eof is not reached",
            ))
        }
    }

    #[inline]
    pub fn poll_completed<T: AsyncWrite>(
        &mut self, stream: &mut T, shutdown: bool,
    ) -> Poll<(), io::Error> {
        match self.write_to_stream(stream) {
            Ok(WriterState::Done) => {
                if shutdown {
                    stream.shutdown()
                } else {
                    Ok(Async::Ready(()))
                }
            }
            Ok(WriterState::Pause) => Ok(Async::NotReady),
            Err(err) => Err(err),
        }
    }
}

fn content_encoder(buf: BytesMut, req: &mut ClientRequest) -> Output {
    let version = req.version();
    let mut body = req.replace_body(Body::Empty);
    let mut encoding = req.content_encoding();

    let transfer = match body {
        Body::Empty => {
            req.headers_mut().remove(CONTENT_LENGTH);
            return Output::Empty(buf);
        }
        Body::Binary(ref mut bytes) => {
            #[cfg(any(feature = "flate2", feature = "brotli"))]
            {
                if encoding.is_compression() {
                    let mut tmp = BytesMut::new();
                    let mut transfer = TransferEncoding::eof(tmp);
                    let mut enc = match encoding {
                        #[cfg(feature = "flate2")]
                        ContentEncoding::Deflate => ContentEncoder::Deflate(
                            DeflateEncoder::new(transfer, Compression::default()),
                        ),
                        #[cfg(feature = "flate2")]
                        ContentEncoding::Gzip => ContentEncoder::Gzip(GzEncoder::new(
                            transfer,
                            Compression::default(),
                        )),
                        #[cfg(feature = "brotli")]
                        ContentEncoding::Br => {
                            ContentEncoder::Br(BrotliEncoder::new(transfer, 5))
                        }
                        ContentEncoding::Auto | ContentEncoding::Identity => {
                            unreachable!()
                        }
                    };
                    // TODO return error!
                    let _ = enc.write(bytes.as_ref());
                    let _ = enc.write_eof();
                    *bytes = Binary::from(enc.buf_mut().take());

                    req.headers_mut().insert(
                        CONTENT_ENCODING,
                        HeaderValue::from_static(encoding.as_str()),
                    );
                    encoding = ContentEncoding::Identity;
                }
                let mut b = BytesMut::new();
                let _ = write!(b, "{}", bytes.len());
                req.headers_mut()
                    .insert(CONTENT_LENGTH, HeaderValue::try_from(b.freeze()).unwrap());
                TransferEncoding::eof(buf)
            }
            #[cfg(not(any(feature = "flate2", feature = "brotli")))]
            {
                let mut b = BytesMut::new();
                let _ = write!(b, "{}", bytes.len());
                req.headers_mut()
                    .insert(CONTENT_LENGTH, HeaderValue::try_from(b.freeze()).unwrap());
                TransferEncoding::eof(buf)
            }
        }
        Body::Streaming(_) | Body::Actor(_) => {
            if req.upgrade() {
                if version == Version::HTTP_2 {
                    error!("Connection upgrade is forbidden for HTTP/2");
                } else {
                    req.headers_mut()
                        .insert(CONNECTION, HeaderValue::from_static("upgrade"));
                }
                if encoding != ContentEncoding::Identity {
                    encoding = ContentEncoding::Identity;
                    req.headers_mut().remove(CONTENT_ENCODING);
                }
                TransferEncoding::eof(buf)
            } else {
                streaming_encoding(buf, version, req)
            }
        }
    };

    if encoding.is_compression() {
        req.headers_mut().insert(
            CONTENT_ENCODING,
            HeaderValue::from_static(encoding.as_str()),
        );
    }

    req.replace_body(body);
    let enc = match encoding {
        #[cfg(feature = "flate2")]
        ContentEncoding::Deflate => ContentEncoder::Deflate(DeflateEncoder::new(
            transfer,
            Compression::default(),
        )),
        #[cfg(feature = "flate2")]
        ContentEncoding::Gzip => {
            ContentEncoder::Gzip(GzEncoder::new(transfer, Compression::default()))
        }
        #[cfg(feature = "brotli")]
        ContentEncoding::Br => ContentEncoder::Br(BrotliEncoder::new(transfer, 5)),
        ContentEncoding::Identity | ContentEncoding::Auto => return Output::TE(transfer),
    };
    Output::Encoder(enc)
}

fn streaming_encoding(
    buf: BytesMut, version: Version, req: &mut ClientRequest,
) -> TransferEncoding {
    if req.chunked() {
        // Enable transfer encoding
        req.headers_mut().remove(CONTENT_LENGTH);
        if version == Version::HTTP_2 {
            req.headers_mut().remove(TRANSFER_ENCODING);
            TransferEncoding::eof(buf)
        } else {
            req.headers_mut()
                .insert(TRANSFER_ENCODING, HeaderValue::from_static("chunked"));
            TransferEncoding::chunked(buf)
        }
    } else {
        // if Content-Length is specified, then use it as length hint
        let (len, chunked) = if let Some(len) = req.headers().get(CONTENT_LENGTH) {
            // Content-Length
            if let Ok(s) = len.to_str() {
                if let Ok(len) = s.parse::<u64>() {
                    (Some(len), false)
                } else {
                    error!("illegal Content-Length: {:?}", len);
                    (None, false)
                }
            } else {
                error!("illegal Content-Length: {:?}", len);
                (None, false)
            }
        } else {
            (None, true)
        };

        if !chunked {
            if let Some(len) = len {
                TransferEncoding::length(len, buf)
            } else {
                TransferEncoding::eof(buf)
            }
        } else {
            // Enable transfer encoding
            match version {
                Version::HTTP_11 => {
                    req.headers_mut()
                        .insert(TRANSFER_ENCODING, HeaderValue::from_static("chunked"));
                    TransferEncoding::chunked(buf)
                }
                _ => {
                    req.headers_mut().remove(TRANSFER_ENCODING);
                    TransferEncoding::eof(buf)
                }
            }
        }
    }
}

// "Sun, 06 Nov 1994 08:49:37 GMT".len()
pub const DATE_VALUE_LENGTH: usize = 29;

fn set_date(dst: &mut BytesMut) {
    CACHED.with(|cache| {
        let mut cache = cache.borrow_mut();
        let now = time::get_time();
        if now > cache.next_update {
            cache.update(now);
        }
        dst.extend_from_slice(cache.buffer());
    })
}

struct CachedDate {
    bytes: [u8; DATE_VALUE_LENGTH],
    next_update: time::Timespec,
}

thread_local!(static CACHED: RefCell<CachedDate> = RefCell::new(CachedDate {
    bytes: [0; DATE_VALUE_LENGTH],
    next_update: time::Timespec::new(0, 0),
}));

impl CachedDate {
    fn buffer(&self) -> &[u8] {
        &self.bytes[..]
    }

    fn update(&mut self, now: time::Timespec) {
        write!(&mut self.bytes[..], "{}", time::at_utc(now).rfc822()).unwrap();
        self.next_update = now + Duration::seconds(1);
        self.next_update.nsec = 0;
    }
}
<|endoftext|>


<commit_before>// (C)opyleft 2013,2014 Frank Denis

/*!
 * Bloom filter for Rust
 *
 * This is a simple but fast Bloom filter implementation, that requires only
 * 2 hash functions, generated with SipHash-2-4 using randomized keys.
 */

#![desc = "A fast Bloom filter implementation."]
#![license = "BSD"]
#![crate_name="bloomfilter"]
#![crate_type = "rlib"]
#![warn(non_camel_case_types, non_uppercase_statics, unnecessary_qualification)]

extern crate collections;

use std::cmp;
use std::f64;
use std::hash::Hash;
use std::hash::sip::SipHasher;
use std::hash::Hasher;
use std::rand;
use collections::bitv;

#[cfg(test)]
use std::rand::Rng;

/// Bloom filter structure
pub struct Bloom {
    bitmap: bitv::Bitv,
    bitmap_bits: u64,
    k_num: uint,
    sips: [SipHasher, ..2]
}

impl Bloom {
/// Create a new bloom filter structure.
/// bitmap_size is the size in bytes (not bits) that will be allocated in memory
/// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: uint, items_count: uint) -> Bloom {
        assert!(bitmap_size > 0u && items_count > 0u);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Bloom::optimal_k_num(bitmap_bits, items_count);
        let bitmap = bitv::Bitv::with_capacity(bitmap_bits as uint, false);
        let sips = [ Bloom::sip_new(), Bloom::sip_new() ];
        Bloom {
            bitmap: bitmap,
            bitmap_bits: bitmap_bits,
            k_num: k_num,
            sips: sips
        }
    }

/// Create a new bloom filter structure.
/// items_count is an estimation of the maximum number of items to store.
/// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: uint, fp_p: f64) -> Bloom {
        let bitmap_size = Bloom::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

/// Compute a recommended bitmap size for items_count items
/// and a fp_p rate of false positives.
/// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: uint, fp_p: f64) -> uint {
        assert!(items_count > 0u);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = f64::consts::LN_2;
        let log2_2 = log2 * log2;
        ((items_count as f64) * fp_p.ln() / (-8.0 * log2_2)).ceil() as uint
    }

/// Record the presence of an item.
    pub fn set<T: Hash>(& mut self, item: T) {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            self.bitmap.set(bit_offset, true);
        }
    }

/// Check if an item is present in the set.
/// There can be false positives, but no false negatives.
    pub fn check<T: Hash>(&self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                return false;
            }
        }
        true
    }

/// Record the presence of an item in the set,
/// and return the previous state of this item.
    pub fn check_and_set<T: Hash>(&mut self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        let mut found = true;
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

/// Return the number of bits in the filter
    pub fn number_of_bits(&self) -> u64 {
        self.bitmap_bits
    }

/// Return the number of hash functions used for `check` and `set`
    pub fn number_of_hash_functions(&self) -> uint {
        self.k_num
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: uint) -> uint {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * 2.0f64.ln()).ceil() as uint;
        cmp::max(k_num, 1)
    }

    fn bloom_hash<T: Hash>(&self, hashes: & mut [u64, ..2],
                  item: &T, k_i: uint) -> u64 {
        if k_i < 2 {
            let sip = self.sips[k_i];
            let hash = sip.hash(item);
            hashes[k_i] = hash;
            hash
        } else {
            hashes[0] + (((k_i as u64) * hashes[1]) % 0xffffffffffffffc5)
        }
    }

    fn sip_new() -> SipHasher {
        let mut rng = rand::task_rng();
        SipHasher::new_with_keys(rand::Rand::rand(& mut rng),
                                 rand::Rand::rand(& mut rng))
    }
}

#[test]
fn bloom_test_set() {
    let mut bloom = Bloom::new(10, 80);
    let key: &Vec<u8> = &rand::task_rng().gen_iter::<u8>().take(16u).collect();
    assert!(bloom.check(key) == false);
    bloom.set(&key);
    assert!(bloom.check(key.clone()) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut bloom = Bloom::new(10, 80);
    let key: &Vec<u8> = &rand::task_rng().gen_iter::<u8>().take(16u).collect();
    assert!(bloom.check_and_set(key) == false);
    assert!(bloom.check_and_set(key.clone()) == true);
}
<commit_msg>Update deprecated lint directives<commit_after>// (C)opyleft 2013,2014 Frank Denis

/*!
 * Bloom filter for Rust
 *
 * This is a simple but fast Bloom filter implementation, that requires only
 * 2 hash functions, generated with SipHash-2-4 using randomized keys.
 */

#![desc = "A fast Bloom filter implementation."]
#![license = "BSD"]
#![crate_name="bloomfilter"]
#![crate_type = "rlib"]
#![warn(non_camel_case_types, non_upper_case_globals, unused_qualifications)]

extern crate collections;

use std::cmp;
use std::f64;
use std::hash::Hash;
use std::hash::sip::SipHasher;
use std::hash::Hasher;
use std::rand;
use collections::bitv;

#[cfg(test)]
use std::rand::Rng;

/// Bloom filter structure
pub struct Bloom {
    bitmap: bitv::Bitv,
    bitmap_bits: u64,
    k_num: uint,
    sips: [SipHasher, ..2]
}

impl Bloom {
/// Create a new bloom filter structure.
/// bitmap_size is the size in bytes (not bits) that will be allocated in memory
/// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: uint, items_count: uint) -> Bloom {
        assert!(bitmap_size > 0u && items_count > 0u);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Bloom::optimal_k_num(bitmap_bits, items_count);
        let bitmap = bitv::Bitv::with_capacity(bitmap_bits as uint, false);
        let sips = [ Bloom::sip_new(), Bloom::sip_new() ];
        Bloom {
            bitmap: bitmap,
            bitmap_bits: bitmap_bits,
            k_num: k_num,
            sips: sips
        }
    }

/// Create a new bloom filter structure.
/// items_count is an estimation of the maximum number of items to store.
/// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: uint, fp_p: f64) -> Bloom {
        let bitmap_size = Bloom::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

/// Compute a recommended bitmap size for items_count items
/// and a fp_p rate of false positives.
/// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: uint, fp_p: f64) -> uint {
        assert!(items_count > 0u);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = f64::consts::LN_2;
        let log2_2 = log2 * log2;
        ((items_count as f64) * fp_p.ln() / (-8.0 * log2_2)).ceil() as uint
    }

/// Record the presence of an item.
    pub fn set<T: Hash>(& mut self, item: T) {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            self.bitmap.set(bit_offset, true);
        }
    }

/// Check if an item is present in the set.
/// There can be false positives, but no false negatives.
    pub fn check<T: Hash>(&self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                return false;
            }
        }
        true
    }

/// Record the presence of an item in the set,
/// and return the previous state of this item.
    pub fn check_and_set<T: Hash>(&mut self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        let mut found = true;
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

/// Return the number of bits in the filter
    pub fn number_of_bits(&self) -> u64 {
        self.bitmap_bits
    }

/// Return the number of hash functions used for `check` and `set`
    pub fn number_of_hash_functions(&self) -> uint {
        self.k_num
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: uint) -> uint {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * 2.0f64.ln()).ceil() as uint;
        cmp::max(k_num, 1)
    }

    fn bloom_hash<T: Hash>(&self, hashes: & mut [u64, ..2],
                  item: &T, k_i: uint) -> u64 {
        if k_i < 2 {
            let sip = self.sips[k_i];
            let hash = sip.hash(item);
            hashes[k_i] = hash;
            hash
        } else {
            hashes[0] + (((k_i as u64) * hashes[1]) % 0xffffffffffffffc5)
        }
    }

    fn sip_new() -> SipHasher {
        let mut rng = rand::task_rng();
        SipHasher::new_with_keys(rand::Rand::rand(& mut rng),
                                 rand::Rand::rand(& mut rng))
    }
}

#[test]
fn bloom_test_set() {
    let mut bloom = Bloom::new(10, 80);
    let key: &Vec<u8> = &rand::task_rng().gen_iter::<u8>().take(16u).collect();
    assert!(bloom.check(key) == false);
    bloom.set(&key);
    assert!(bloom.check(key.clone()) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut bloom = Bloom::new(10, 80);
    let key: &Vec<u8> = &rand::task_rng().gen_iter::<u8>().take(16u).collect();
    assert!(bloom.check_and_set(key) == false);
    assert!(bloom.check_and_set(key.clone()) == true);
}
<|endoftext|>