<commit_before>use crate::query_builder::upsert::into_conflict_clause::IntoConflictValueClause;
use crate::query_builder::upsert::on_conflict_actions::*;
use crate::query_builder::upsert::on_conflict_clause::*;
use crate::query_builder::upsert::on_conflict_target::*;
use crate::query_builder::{AsChangeset, InsertStatement, UndecoratedInsertRecord};
use crate::query_source::QuerySource;

impl<T, U, Op, Ret> InsertStatement<T, U, Op, Ret>
where
    U: UndecoratedInsertRecord<T> + IntoConflictValueClause,
{
    /// Adds `ON CONFLICT DO NOTHING` to the insert statement, without
    /// specifying any columns or constraints to restrict the conflict to.
    ///
    /// # Examples
    ///
    /// ### Single Record
    ///
    /// ```rust
    /// # include!("on_conflict_docs_setup.rs");
    /// #
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// #     let conn = establish_connection();
    /// #     #[cfg(feature = "postgres")]
    /// #     conn.execute("TRUNCATE TABLE users").unwrap();
    /// #     #[cfg(feature = "sqlite")]
    /// #     conn.execute("DELETE FROM users").unwrap();
    /// let user = User { id: 1, name: "Sean", };
    ///
    /// let inserted_row_count = diesel::insert_into(users)
    ///     .values(&user)
    ///     .on_conflict_do_nothing()
    ///     .execute(&conn);
    /// assert_eq!(Ok(1), inserted_row_count);
    ///
    /// let inserted_row_count = diesel::insert_into(users)
    ///     .values(&user)
    ///     .on_conflict_do_nothing()
    ///     .execute(&conn);
    /// assert_eq!(Ok(0), inserted_row_count);
    /// # }
    /// ```
    ///
    /// ### Vec of Records
    ///
    /// ```rust
    /// # include!("on_conflict_docs_setup.rs");
    /// #
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// #     let conn = establish_connection();
    /// #     #[cfg(feature = "postgres")]
    /// #     conn.execute("TRUNCATE TABLE users").unwrap();
    /// # #[cfg(feature = "postgres")]
    /// let user = User { id: 1, name: "Sean", };
    ///
    /// # #[cfg(feature = "postgres")]
    /// let inserted_row_count = diesel::insert_into(users)
    ///     .values(&vec![user, user])
    ///     .on_conflict_do_nothing()
    ///     .execute(&conn);
    /// # #[cfg(feature = "postgres")]
    /// assert_eq!(Ok(1), inserted_row_count);
    /// # }
    /// ```
    pub fn on_conflict_do_nothing(
        self,
    ) -> InsertStatement<T, OnConflictValues<U::ValueClause, NoConflictTarget, DoNothing>, Op, Ret>
    {
        self.replace_values(|values| OnConflictValues::do_nothing(values.into_value_clause()))
    }

    /// Adds an `ON CONFLICT` to the insert statement, if a conflict occurs
    /// for the given unique constraint.
    ///
    /// `Target` can be one of:
    ///
    /// - A column
    /// - A tuple of columns
    /// - [`on_constraint("constraint_name")`][`on_constraint`]
    ///
    /// # Examples
    ///
    /// ### Specifying a column as the target
    ///
    /// ```rust
    /// # include!("on_conflict_docs_setup.rs");
    /// #
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// #     let conn = establish_connection();
    /// #     #[cfg(feature = "postgres")]
    /// #     conn.execute("TRUNCATE TABLE users").unwrap();
    /// #     #[cfg(feature = "sqlite")]
    /// #     conn.execute("DELETE FROM users").unwrap();
    /// conn.execute("CREATE UNIQUE INDEX users_name ON users (name)").unwrap();
    /// let user = User { id: 1, name: "Sean", };
    /// let same_name_different_id = User { id: 2, name: "Sean" };
    /// let same_id_different_name = User { id: 1, name: "Pascal" };
    ///
    /// assert_eq!(Ok(1), diesel::insert_into(users).values(&user).execute(&conn));
    ///
    /// let inserted_row_count = diesel::insert_into(users)
    ///     .values(&same_name_different_id)
    ///     .on_conflict(name)
    ///     .do_nothing()
    ///     .execute(&conn);
    /// assert_eq!(Ok(0), inserted_row_count);
    ///
    /// let pk_conflict_result = diesel::insert_into(users)
    ///     .values(&same_id_different_name)
    ///     .on_conflict(name)
    ///     .do_nothing()
    ///     .execute(&conn);
    /// assert!(pk_conflict_result.is_err());
    /// # }
    /// ```
    ///
    /// ### Specifying multiple columns as the target
    ///
    /// ```rust
    /// # include!("../doctest_setup.rs");
    /// #
    /// # table! {
    /// #     users {
    /// #         id -> Integer,
    /// #         name -> VarChar,
    /// #         hair_color -> VarChar,
    /// #     }
    /// # }
    /// #
    /// # #[derive(Clone, Copy, Insertable)]
    /// # #[table_name="users"]
    /// # struct User<'a> {
    /// #     id: i32,
    /// #     name: &'a str,
    /// #     hair_color: &'a str,
    /// # }
    /// #
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// use diesel::upsert::*;
    ///
    /// #     let conn = establish_connection();
    /// #     conn.execute("DROP TABLE users").unwrap();
    /// #     conn.execute("CREATE TABLE users (id SERIAL PRIMARY KEY, name TEXT, hair_color TEXT)").unwrap();
    /// conn.execute("CREATE UNIQUE INDEX users_name_hair_color ON users (name, hair_color)").unwrap();
    /// let user = User { id: 1, name: "Sean", hair_color: "black" };
    /// let same_name_different_hair_color = User { id: 2, name: "Sean", hair_color: "brown" };
    /// let same_name_same_hair_color = User { id: 3, name: "Sean", hair_color: "black" };
    ///
    /// assert_eq!(Ok(1), diesel::insert_into(users).values(&user).execute(&conn));
    ///
    /// let inserted_row_count = diesel::insert_into(users)
    ///     .values(&same_name_different_hair_color)
    ///     .on_conflict((name, hair_color))
    ///     .do_nothing()
    ///     .execute(&conn);
    /// assert_eq!(Ok(1), inserted_row_count);
    ///
    /// let inserted_row_count = diesel::insert_into(users)
    ///     .values(&same_name_same_hair_color)
    ///     .on_conflict((name, hair_color))
    ///     .do_nothing()
    ///     .execute(&conn);
    /// assert_eq!(Ok(0), inserted_row_count);
    /// # }
    /// ```
    ///
    /// See the documentation for [`on_constraint`] and [`do_update`] for
    /// more examples.
    ///
    /// [`on_constraint`]: ../upsert/fn.on_constraint.html
    /// [`do_update`]: ../upsert/struct.IncompleteOnConflict.html#method.do_update
    pub fn on_conflict<Target>(
        self,
        target: Target,
    ) -> IncompleteOnConflict<InsertStatement<T, U::ValueClause, Op, Ret>, ConflictTarget<Target>>
    where
        ConflictTarget<Target>: OnConflictTarget<T>,
    {
        IncompleteOnConflict {
            stmt: self.replace_values(IntoConflictValueClause::into_value_clause),
            target: ConflictTarget(target),
        }
    }
}

/// A partially constructed `ON CONFLICT` clause.
#[derive(Debug, Clone, Copy)]
pub struct IncompleteOnConflict<Stmt, Target> {
    stmt: Stmt,
    target: Target,
}

impl<T, U, Op, Ret, Target> IncompleteOnConflict<InsertStatement<T, U, Op, Ret>, Target> {
    /// Creates a query with `ON CONFLICT (target) DO NOTHING`
    ///
    /// If you want to do nothing when *any* constraint conflicts, use
    /// [`on_conflict_do_nothing`] instead. See [`on_conflict`] for usage
    /// examples.
    ///
    /// [`on_conflict_do_nothing`]: ../../query_builder/struct.InsertStatement.html#method.on_conflict_do_nothing
    /// [`on_conflict`]: ../../query_builder/struct.InsertStatement.html#method.on_conflict
    pub fn do_nothing(self) -> InsertStatement<T, OnConflictValues<U, Target, DoNothing>, Op, Ret> {
        let target = self.target;
        self.stmt
            .replace_values(|values| OnConflictValues::new(values, target, DoNothing))
    }
}

impl<Stmt, Target> IncompleteOnConflict<Stmt, Target> {
    /// Used to create a query in the form `ON CONFLICT (...) DO UPDATE ...`
    ///
    /// Call `.set` on the result of this function with the changes you want to
    /// apply. The argument to `set` can be anything that implements `AsChangeset`
    /// (e.g. anything you could pass to `set` on a normal update statement).
    ///
    /// Note: When inserting more than one row at a time, this query can still fail
    /// if the rows being inserted conflict with each other.
    ///
    /// # Examples
    ///
    /// ## Set specific value on conflict
    ///
    /// ```rust
    /// # include!("on_conflict_docs_setup.rs");
    /// #
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// #     let conn = establish_connection();
    /// #     #[cfg(feature = "postgres")]
    /// #     conn.execute("TRUNCATE TABLE users").unwrap();
    /// #     #[cfg(feature = "sqlite")]
    /// #     conn.execute("DELETE FROM users").unwrap();
    /// let user = User { id: 1, name: "Pascal" };
    /// let user2 = User { id: 1, name: "Sean" };
    ///
    /// assert_eq!(Ok(1), diesel::insert_into(users).values(&user).execute(&conn));
    ///
    /// let insert_count = diesel::insert_into(users)
    ///     .values(&user2)
    ///     .on_conflict(id)
    ///     .do_update()
    ///     .set(name.eq("I DONT KNOW ANYMORE"))
    ///     .execute(&conn);
    /// assert_eq!(Ok(1), insert_count);
    ///
    /// let users_in_db = users.load(&conn);
    /// assert_eq!(Ok(vec![(1, "I DONT KNOW ANYMORE".to_string())]), users_in_db);
    /// # }
    /// ```
    ///
    /// ## Set `AsChangeset` struct on conflict
    ///
    /// ```rust
    /// # include!("on_conflict_docs_setup.rs");
    /// #
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// #     let conn = establish_connection();
    /// #     #[cfg(feature = "postgres")]
    /// #     conn.execute("TRUNCATE TABLE users").unwrap();
    /// #     #[cfg(feature = "sqlite")]
    /// #     conn.execute("DELETE FROM users").unwrap();
    /// let user = User { id: 1, name: "Pascal" };
    /// let user2 = User { id: 1, name: "Sean" };
    ///
    /// assert_eq!(Ok(1), diesel::insert_into(users).values(&user).execute(&conn));
    ///
    /// let insert_count = diesel::insert_into(users)
    ///     .values(&user2)
    ///     .on_conflict(id)
    ///     .do_update()
    ///     .set(&user2)
    ///     .execute(&conn);
    /// assert_eq!(Ok(1), insert_count);
    ///
    /// let users_in_db = users.load(&conn);
    /// assert_eq!(Ok(vec![(1, "Sean".to_string())]), users_in_db);
    /// # }
    /// ```
    ///
    /// ## Use `excluded` to get the rejected value
    ///
    /// ```rust
    /// # include!("on_conflict_docs_setup.rs");
    /// #
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// use diesel::upsert::excluded;
    ///
    /// #     let conn = establish_connection();
    /// #     #[cfg(feature = "postgres")]
    /// #     conn.execute("TRUNCATE TABLE users").unwrap();
    /// let user = User { id: 1, name: "Pascal" };
    /// let user2 = User { id: 1, name: "Sean" };
    /// let user3 = User { id: 2, name: "Tess" };
    ///
    /// # #[cfg(feature = "postgres")]
    /// assert_eq!(Ok(1), diesel::insert_into(users).values(&user).execute(&conn));
    ///
    /// #[cfg(feature = "postgres")]
    /// let insert_count = diesel::insert_into(users)
    ///     .values(&vec![user2, user3])
    ///     .on_conflict(id)
    ///     .do_update()
    ///     .set(name.eq(excluded(name)))
    ///     .execute(&conn);
    /// # #[cfg(feature = "postgres")]
    /// assert_eq!(Ok(2), insert_count);
    ///
    /// # #[cfg(feature = "postgres")]
    /// let users_in_db = users.load(&conn);
    /// # #[cfg(feature = "postgres")]
    /// assert_eq!(Ok(vec![(1, "Sean".to_string()), (2, "Tess".to_string())]), users_in_db);
    /// # }
    /// ```
    pub fn do_update(self) -> IncompleteDoUpdate<Stmt, Target> {
        IncompleteDoUpdate {
            stmt: self.stmt,
            target: self.target,
        }
    }
}

/// A partially constructed `ON CONFLICT DO UPDATE` clause.
#[derive(Debug, Clone, Copy)]
pub struct IncompleteDoUpdate<Stmt, Target> {
    stmt: Stmt,
    target: Target,
}

impl<T, U, Op, Ret, Target> IncompleteDoUpdate<InsertStatement<T, U, Op, Ret>, Target> {
    /// See [`do_update`] for usage examples.
    ///
    /// [`do_update`]: struct.IncompleteOnConflict.html#method.do_update
    pub fn set<Changes>(
        self,
        changes: Changes,
    ) -> InsertStatement<T, OnConflictValues<U, Target, DoUpdate<Changes::Changeset>>, Op, Ret>
    where
        T: QuerySource,
        Changes: AsChangeset<Target = T>,
    {
        let target = self.target;
        self.stmt.replace_values(|values| {
            OnConflictValues::new(values, target, DoUpdate::new(changes.as_changeset()))
        })
    }
}
<commit_msg>Fix doc tests<commit_after>use crate::query_builder::upsert::into_conflict_clause::IntoConflictValueClause;
use crate::query_builder::upsert::on_conflict_actions::*;
use crate::query_builder::upsert::on_conflict_clause::*;
use crate::query_builder::upsert::on_conflict_target::*;
use crate::query_builder::{AsChangeset, InsertStatement, UndecoratedInsertRecord};
use crate::query_source::QuerySource;

impl<T, U, Op, Ret> InsertStatement<T, U, Op, Ret>
where
    U: UndecoratedInsertRecord<T> + IntoConflictValueClause,
{
    /// Adds `ON CONFLICT DO NOTHING` to the insert statement, without
    /// specifying any columns or constraints to restrict the conflict to.
    ///
    /// # Examples
    ///
    /// ### Single Record
    ///
    /// ```rust
    /// # include!("on_conflict_docs_setup.rs");
    /// #
    /// # #[cfg(any(feature = "sqlite", feature = "postgres"))]
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// #     let conn = establish_connection();
    /// #     #[cfg(feature = "postgres")]
    /// #     conn.execute("TRUNCATE TABLE users").unwrap();
    /// #     #[cfg(feature = "sqlite")]
    /// #     conn.execute("DELETE FROM users").unwrap();
    /// let user = User { id: 1, name: "Sean", };
    ///
    /// let inserted_row_count = diesel::insert_into(users)
    ///     .values(&user)
    ///     .on_conflict_do_nothing()
    ///     .execute(&conn);
    /// assert_eq!(Ok(1), inserted_row_count);
    ///
    /// let inserted_row_count = diesel::insert_into(users)
    ///     .values(&user)
    ///     .on_conflict_do_nothing()
    ///     .execute(&conn);
    /// assert_eq!(Ok(0), inserted_row_count);
    /// # }
    /// # #[cfg(feature = "mysql")]
    /// # fn main() {}
    /// ```
    ///
    /// ### Vec of Records
    ///
    /// ```rust
    /// # include!("on_conflict_docs_setup.rs");
    /// #
    /// # #[cfg(any(feature = "sqlite", feature = "postgres"))]
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// #     let conn = establish_connection();
    /// #     #[cfg(feature = "postgres")]
    /// #     conn.execute("TRUNCATE TABLE users").unwrap();
    /// # #[cfg(feature = "postgres")]
    /// let user = User { id: 1, name: "Sean", };
    ///
    /// # #[cfg(feature = "postgres")]
    /// let inserted_row_count = diesel::insert_into(users)
    ///     .values(&vec![user, user])
    ///     .on_conflict_do_nothing()
    ///     .execute(&conn);
    /// # #[cfg(feature = "postgres")]
    /// assert_eq!(Ok(1), inserted_row_count);
    /// # }
    /// # #[cfg(feature = "mysql")]
    /// # fn main() {}
    /// ```
    pub fn on_conflict_do_nothing(
        self,
    ) -> InsertStatement<T, OnConflictValues<U::ValueClause, NoConflictTarget, DoNothing>, Op, Ret>
    {
        self.replace_values(|values| OnConflictValues::do_nothing(values.into_value_clause()))
    }

    /// Adds an `ON CONFLICT` to the insert statement, if a conflict occurs
    /// for the given unique constraint.
    ///
    /// `Target` can be one of:
    ///
    /// - A column
    /// - A tuple of columns
    /// - [`on_constraint("constraint_name")`][`on_constraint`]
    ///
    /// # Examples
    ///
    /// ### Specifying a column as the target
    ///
    /// ```rust
    /// # include!("on_conflict_docs_setup.rs");
    /// #
    /// # #[cfg(any(feature = "sqlite", feature = "postgres"))]
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// #     let conn = establish_connection();
    /// #     #[cfg(feature = "postgres")]
    /// #     conn.execute("TRUNCATE TABLE users").unwrap();
    /// #     #[cfg(feature = "sqlite")]
    /// #     conn.execute("DELETE FROM users").unwrap();
    /// conn.execute("CREATE UNIQUE INDEX users_name ON users (name)").unwrap();
    /// let user = User { id: 1, name: "Sean", };
    /// let same_name_different_id = User { id: 2, name: "Sean" };
    /// let same_id_different_name = User { id: 1, name: "Pascal" };
    ///
    /// assert_eq!(Ok(1), diesel::insert_into(users).values(&user).execute(&conn));
    ///
    /// let inserted_row_count = diesel::insert_into(users)
    ///     .values(&same_name_different_id)
    ///     .on_conflict(name)
    ///     .do_nothing()
    ///     .execute(&conn);
    /// assert_eq!(Ok(0), inserted_row_count);
    ///
    /// let pk_conflict_result = diesel::insert_into(users)
    ///     .values(&same_id_different_name)
    ///     .on_conflict(name)
    ///     .do_nothing()
    ///     .execute(&conn);
    /// assert!(pk_conflict_result.is_err());
    /// # }
    /// # #[cfg(feature = "mysql")]
    /// # fn main() {}
    /// ```
    ///
    /// ### Specifying multiple columns as the target
    ///
    /// ```rust
    /// # include!("../doctest_setup.rs");
    /// #
    /// # table! {
    /// #     users {
    /// #         id -> Integer,
    /// #         name -> VarChar,
    /// #         hair_color -> VarChar,
    /// #     }
    /// # }
    /// #
    /// # #[derive(Clone, Copy, Insertable)]
    /// # #[table_name="users"]
    /// # struct User<'a> {
    /// #     id: i32,
    /// #     name: &'a str,
    /// #     hair_color: &'a str,
    /// # }
    /// #
    /// # #[cfg(any(feature = "sqlite", feature = "postgres"))]
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// use diesel::upsert::*;
    ///
    /// #     let conn = establish_connection();
    /// #     conn.execute("DROP TABLE users").unwrap();
    /// #     conn.execute("CREATE TABLE users (id SERIAL PRIMARY KEY, name TEXT, hair_color TEXT)").unwrap();
    /// conn.execute("CREATE UNIQUE INDEX users_name_hair_color ON users (name, hair_color)").unwrap();
    /// let user = User { id: 1, name: "Sean", hair_color: "black" };
    /// let same_name_different_hair_color = User { id: 2, name: "Sean", hair_color: "brown" };
    /// let same_name_same_hair_color = User { id: 3, name: "Sean", hair_color: "black" };
    ///
    /// assert_eq!(Ok(1), diesel::insert_into(users).values(&user).execute(&conn));
    ///
    /// let inserted_row_count = diesel::insert_into(users)
    ///     .values(&same_name_different_hair_color)
    ///     .on_conflict((name, hair_color))
    ///     .do_nothing()
    ///     .execute(&conn);
    /// assert_eq!(Ok(1), inserted_row_count);
    ///
    /// let inserted_row_count = diesel::insert_into(users)
    ///     .values(&same_name_same_hair_color)
    ///     .on_conflict((name, hair_color))
    ///     .do_nothing()
    ///     .execute(&conn);
    /// assert_eq!(Ok(0), inserted_row_count);
    /// # }
    /// # #[cfg(feature = "mysql")]
    /// # fn main() {}
    /// ```
    ///
    /// See the documentation for [`on_constraint`] and [`do_update`] for
    /// more examples.
    ///
    /// [`on_constraint`]: ../upsert/fn.on_constraint.html
    /// [`do_update`]: ../upsert/struct.IncompleteOnConflict.html#method.do_update
    pub fn on_conflict<Target>(
        self,
        target: Target,
    ) -> IncompleteOnConflict<InsertStatement<T, U::ValueClause, Op, Ret>, ConflictTarget<Target>>
    where
        ConflictTarget<Target>: OnConflictTarget<T>,
    {
        IncompleteOnConflict {
            stmt: self.replace_values(IntoConflictValueClause::into_value_clause),
            target: ConflictTarget(target),
        }
    }
}

/// A partially constructed `ON CONFLICT` clause.
#[derive(Debug, Clone, Copy)]
pub struct IncompleteOnConflict<Stmt, Target> {
    stmt: Stmt,
    target: Target,
}

impl<T, U, Op, Ret, Target> IncompleteOnConflict<InsertStatement<T, U, Op, Ret>, Target> {
    /// Creates a query with `ON CONFLICT (target) DO NOTHING`
    ///
    /// If you want to do nothing when *any* constraint conflicts, use
    /// [`on_conflict_do_nothing`] instead. See [`on_conflict`] for usage
    /// examples.
    ///
    /// [`on_conflict_do_nothing`]: ../../query_builder/struct.InsertStatement.html#method.on_conflict_do_nothing
    /// [`on_conflict`]: ../../query_builder/struct.InsertStatement.html#method.on_conflict
    pub fn do_nothing(self) -> InsertStatement<T, OnConflictValues<U, Target, DoNothing>, Op, Ret> {
        let target = self.target;
        self.stmt
            .replace_values(|values| OnConflictValues::new(values, target, DoNothing))
    }
}

impl<Stmt, Target> IncompleteOnConflict<Stmt, Target> {
    /// Used to create a query in the form `ON CONFLICT (...) DO UPDATE ...`
    ///
    /// Call `.set` on the result of this function with the changes you want to
    /// apply. The argument to `set` can be anything that implements `AsChangeset`
    /// (e.g. anything you could pass to `set` on a normal update statement).
    ///
    /// Note: When inserting more than one row at a time, this query can still fail
    /// if the rows being inserted conflict with each other.
    ///
    /// # Examples
    ///
    /// ## Set specific value on conflict
    ///
    /// ```rust
    /// # include!("on_conflict_docs_setup.rs");
    /// #
    /// # #[cfg(any(feature = "sqlite", feature = "postgres"))]
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// #     let conn = establish_connection();
    /// #     #[cfg(feature = "postgres")]
    /// #     conn.execute("TRUNCATE TABLE users").unwrap();
    /// #     #[cfg(feature = "sqlite")]
    /// #     conn.execute("DELETE FROM users").unwrap();
    /// let user = User { id: 1, name: "Pascal" };
    /// let user2 = User { id: 1, name: "Sean" };
    ///
    /// assert_eq!(Ok(1), diesel::insert_into(users).values(&user).execute(&conn));
    ///
    /// let insert_count = diesel::insert_into(users)
    ///     .values(&user2)
    ///     .on_conflict(id)
    ///     .do_update()
    ///     .set(name.eq("I DONT KNOW ANYMORE"))
    ///     .execute(&conn);
    /// assert_eq!(Ok(1), insert_count);
    ///
    /// let users_in_db = users.load(&conn);
    /// assert_eq!(Ok(vec![(1, "I DONT KNOW ANYMORE".to_string())]), users_in_db);
    /// # }
    /// # #[cfg(feature = "mysql")]
    /// # fn main() {}
    /// ```
    ///
    /// ## Set `AsChangeset` struct on conflict
    ///
    /// ```rust
    /// # include!("on_conflict_docs_setup.rs");
    /// #
    /// # #[cfg(any(feature = "sqlite", feature = "postgres"))]
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// #     let conn = establish_connection();
    /// #     #[cfg(feature = "postgres")]
    /// #     conn.execute("TRUNCATE TABLE users").unwrap();
    /// #     #[cfg(feature = "sqlite")]
    /// #     conn.execute("DELETE FROM users").unwrap();
    /// let user = User { id: 1, name: "Pascal" };
    /// let user2 = User { id: 1, name: "Sean" };
    ///
    /// assert_eq!(Ok(1), diesel::insert_into(users).values(&user).execute(&conn));
    ///
    /// let insert_count = diesel::insert_into(users)
    ///     .values(&user2)
    ///     .on_conflict(id)
    ///     .do_update()
    ///     .set(&user2)
    ///     .execute(&conn);
    /// assert_eq!(Ok(1), insert_count);
    ///
    /// let users_in_db = users.load(&conn);
    /// assert_eq!(Ok(vec![(1, "Sean".to_string())]), users_in_db);
    /// # }
    /// # #[cfg(feature = "mysql")]
    /// # fn main() {}
    /// ```
    ///
    /// ## Use `excluded` to get the rejected value
    ///
    /// ```rust
    /// # include!("on_conflict_docs_setup.rs");
    /// #
    /// # #[cfg(any(feature = "sqlite", feature = "postgres"))]
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// use diesel::upsert::excluded;
    ///
    /// #     let conn = establish_connection();
    /// #     #[cfg(feature = "postgres")]
    /// #     conn.execute("TRUNCATE TABLE users").unwrap();
    /// let user = User { id: 1, name: "Pascal" };
    /// let user2 = User { id: 1, name: "Sean" };
    /// let user3 = User { id: 2, name: "Tess" };
    ///
    /// # #[cfg(feature = "postgres")]
    /// assert_eq!(Ok(1), diesel::insert_into(users).values(&user).execute(&conn));
    ///
    /// #[cfg(feature = "postgres")]
    /// let insert_count = diesel::insert_into(users)
    ///     .values(&vec![user2, user3])
    ///     .on_conflict(id)
    ///     .do_update()
    ///     .set(name.eq(excluded(name)))
    ///     .execute(&conn);
    /// # #[cfg(feature = "postgres")]
    /// assert_eq!(Ok(2), insert_count);
    ///
    /// # #[cfg(feature = "postgres")]
    /// let users_in_db = users.load(&conn);
    /// # #[cfg(feature = "postgres")]
    /// assert_eq!(Ok(vec![(1, "Sean".to_string()), (2, "Tess".to_string())]), users_in_db);
    /// # }
    /// # #[cfg(feature = "mysql")]
    /// # fn main() {}
    /// ```
    pub fn do_update(self) -> IncompleteDoUpdate<Stmt, Target> {
        IncompleteDoUpdate {
            stmt: self.stmt,
            target: self.target,
        }
    }
}

/// A partially constructed `ON CONFLICT DO UPDATE` clause.
#[derive(Debug, Clone, Copy)]
pub struct IncompleteDoUpdate<Stmt, Target> {
    stmt: Stmt,
    target: Target,
}

impl<T, U, Op, Ret, Target> IncompleteDoUpdate<InsertStatement<T, U, Op, Ret>, Target> {
    /// See [`do_update`] for usage examples.
    ///
    /// [`do_update`]: struct.IncompleteOnConflict.html#method.do_update
    pub fn set<Changes>(
        self,
        changes: Changes,
    ) -> InsertStatement<T, OnConflictValues<U, Target, DoUpdate<Changes::Changeset>>, Op, Ret>
    where
        T: QuerySource,
        Changes: AsChangeset<Target = T>,
    {
        let target = self.target;
        self.stmt.replace_values(|values| {
            OnConflictValues::new(values, target, DoUpdate::new(changes.as_changeset()))
        })
    }
}
<|endoftext|>


<commit_before>//! Implements `cargo miri setup` via xargo

use std::env;
use std::ffi::OsStr;
use std::fs::{self};
use std::io::BufRead;
use std::ops::Not;
use std::path::{Path, PathBuf};
use std::process::{self, Command};

use crate::{util::*, version::*};

fn xargo_version() -> Option<(u32, u32, u32)> {
    let out = xargo_check().arg("--version").output().ok()?;
    if !out.status.success() {
        return None;
    }
    // Parse output. The first line looks like "xargo 0.3.12 (b004f1c 2018-12-13)".
    let line = out
        .stderr
        .lines()
        .next()
        .expect("malformed `xargo --version` output: not at least one line")
        .expect("malformed `xargo --version` output: error reading first line");
    let (name, version) = {
        let mut split = line.split(' ');
        (
            split.next().expect("malformed `xargo --version` output: empty"),
            split.next().expect("malformed `xargo --version` output: not at least two words"),
        )
    };
    if name != "xargo" {
        // This is some fork of xargo
        return None;
    }
    let mut version_pieces = version.split('.');
    let major = version_pieces
        .next()
        .expect("malformed `xargo --version` output: not a major version piece")
        .parse()
        .expect("malformed `xargo --version` output: major version is not an integer");
    let minor = version_pieces
        .next()
        .expect("malformed `xargo --version` output: not a minor version piece")
        .parse()
        .expect("malformed `xargo --version` output: minor version is not an integer");
    let patch = version_pieces
        .next()
        .expect("malformed `xargo --version` output: not a patch version piece")
        .parse()
        .expect("malformed `xargo --version` output: patch version is not an integer");
    if version_pieces.next().is_some() {
        panic!("malformed `xargo --version` output: more than three pieces in version");
    }
    Some((major, minor, patch))
}

/// Performs the setup required to make `cargo miri` work: Getting a custom-built libstd. Then sets
/// `MIRI_SYSROOT`. Skipped if `MIRI_SYSROOT` is already set, in which case we expect the user has
/// done all this already.
pub fn setup(subcommand: &MiriCommand, host: &str, target: &str) {
    let only_setup = matches!(subcommand, MiriCommand::Setup);
    let ask_user = !only_setup;
    let print_sysroot = only_setup && has_arg_flag("--print-sysroot"); // whether we just print the sysroot path
    if std::env::var_os("MIRI_SYSROOT").is_some() {
        if only_setup {
            println!("WARNING: MIRI_SYSROOT already set, not doing anything.")
        }
        return;
    }

    // First, we need xargo.
    if xargo_version().map_or(true, |v| v < XARGO_MIN_VERSION) {
        if std::env::var_os("XARGO_CHECK").is_some() {
            // The user manually gave us a xargo binary; don't do anything automatically.
            show_error!("xargo is too old; please upgrade to the latest version")
        }
        let mut cmd = cargo();
        cmd.args(&["install", "xargo"]);
        ask_to_run(cmd, ask_user, "install a recent enough xargo");
    }

    // Determine where the rust sources are located.  The env vars manually setting the source
    // (`MIRI_LIB_SRC`, `XARGO_RUST_SRC`) trump auto-detection.
    let rust_src_env_var =
        std::env::var_os("MIRI_LIB_SRC").or_else(|| std::env::var_os("XARGO_RUST_SRC"));
    let rust_src = match rust_src_env_var {
        Some(path) => {
            let path = PathBuf::from(path);
            // Make path absolute if possible.
            path.canonicalize().unwrap_or(path)
        }
        None => {
            // Check for `rust-src` rustup component.
            let output = miri_for_host()
                .args(&["--print", "sysroot"])
                .output()
                .expect("failed to determine sysroot");
            if !output.status.success() {
                show_error!(
                    "Failed to determine sysroot; Miri said:\n{}",
                    String::from_utf8_lossy(&output.stderr).trim_end()
                );
            }
            let sysroot = std::str::from_utf8(&output.stdout).unwrap();
            let sysroot = Path::new(sysroot.trim_end_matches('\n'));
            // Check for `$SYSROOT/lib/rustlib/src/rust/library`; test if that contains `std/Cargo.toml`.
            let rustup_src =
                sysroot.join("lib").join("rustlib").join("src").join("rust").join("library");
            if !rustup_src.join("std").join("Cargo.toml").exists() {
                // Ask the user to install the `rust-src` component, and use that.
                let mut cmd = Command::new("rustup");
                cmd.args(&["component", "add", "rust-src"]);
                ask_to_run(
                    cmd,
                    ask_user,
                    "install the `rust-src` component for the selected toolchain",
                );
            }
            rustup_src
        }
    };
    if !rust_src.exists() {
        show_error!("given Rust source directory `{}` does not exist.", rust_src.display());
    }
    if rust_src.file_name().and_then(OsStr::to_str) != Some("library") {
        show_error!(
            "given Rust source directory `{}` does not seem to be the `library` subdirectory of \
             a Rust source checkout.",
            rust_src.display()
        );
    }

    // Next, we need our own libstd. Prepare a xargo project for that purpose.
    // We will do this work in whatever is a good cache dir for this platform.
    let dirs = directories::ProjectDirs::from("org", "rust-lang", "miri").unwrap();
    let dir = dirs.cache_dir();
    if !dir.exists() {
        fs::create_dir_all(&dir).unwrap();
    }
    // The interesting bit: Xargo.toml (only needs content if we actually need std)
    let xargo_toml = if std::env::var_os("MIRI_NO_STD").is_some() {
        ""
    } else {
        r#"
[dependencies.std]
default_features = false
# We support unwinding, so enable that panic runtime.
features = ["panic_unwind", "backtrace"]

[dependencies.test]
"#
    };
    write_to_file(&dir.join("Xargo.toml"), xargo_toml);
    // The boring bits: a dummy project for xargo.
    // FIXME: With xargo-check, can we avoid doing this?
    write_to_file(
        &dir.join("Cargo.toml"),
        r#"
[package]
name = "miri-xargo"
description = "A dummy project for building libstd with xargo."
version = "0.0.0"

[lib]
path = "lib.rs"
"#,
    );
    write_to_file(&dir.join("lib.rs"), "#![no_std]");

    // Figure out where xargo will build its stuff.
    // Unfortunately, it puts things into a different directory when the
    // architecture matches the host.
    let sysroot = if target == host { dir.join("HOST") } else { PathBuf::from(dir) };
    // Make sure all target-level Miri invocations know their sysroot.
    std::env::set_var("MIRI_SYSROOT", &sysroot);

    // Now invoke xargo.
    let mut command = xargo_check();
    command.arg("check").arg("-q");
    command.current_dir(&dir);
    command.env("XARGO_HOME", &dir);
    command.env("XARGO_RUST_SRC", &rust_src);
    // We always need to set a target so rustc bootstrap can tell apart host from target crates.
    command.arg("--target").arg(target);
    // Use Miri as rustc to build a libstd compatible with us (and use the right flags).
    // However, when we are running in bootstrap, we cannot just overwrite `RUSTC`,
    // because we still need bootstrap to distinguish between host and target crates.
    // In that case we overwrite `RUSTC_REAL` instead which determines the rustc used
    // for target crates.
    // We set ourselves (`cargo-miri`) instead of Miri directly to be able to patch the flags
    // for `libpanic_abort` (usually this is done by bootstrap but we have to do it ourselves).
    // The `MIRI_CALLED_FROM_XARGO` will mean we dispatch to `phase_setup_rustc`.
    let cargo_miri_path = std::env::current_exe().expect("current executable path invalid");
    if env::var_os("RUSTC_STAGE").is_some() {
        assert!(env::var_os("RUSTC").is_some());
        command.env("RUSTC_REAL", &cargo_miri_path);
    } else {
        command.env("RUSTC", &cargo_miri_path);
    }
    command.env("MIRI_CALLED_FROM_XARGO", "1");
    // Make sure there are no other wrappers getting in our way
    // (Cc https://github.com/rust-lang/miri/issues/1421, https://github.com/rust-lang/miri/issues/2429).
    // Looks like setting `RUSTC_WRAPPER` to the empty string overwrites `build.rustc-wrapper` set via `config.toml`.
    command.env("RUSTC_WRAPPER", "");
    // Disable debug assertions in the standard library -- Miri is already slow enough. But keep the
    // overflow checks, they are cheap. This completely overwrites flags the user might have set,
    // which is consistent with normal `cargo build` that does not apply `RUSTFLAGS` to the sysroot
    // either.
    command.env("RUSTFLAGS", "-Cdebug-assertions=off -Coverflow-checks=on");
    // Manage the output the user sees.
    if only_setup {
        // We want to be explicit.
        eprintln!("Preparing a sysroot for Miri (target: {target})...");
        if print_sysroot {
            // Be extra sure there is no noise on stdout.
            command.stdout(process::Stdio::null());
        }
    } else {
        // We want to be quiet, but still let the user know that something is happening.
        eprint!("Preparing a sysroot for Miri (target: {target})... ");
        command.stdout(process::Stdio::null());
        command.stderr(process::Stdio::null());
    }

    // Finally run it!
    if command.status().expect("failed to run xargo").success().not() {
        if only_setup {
            show_error!("failed to run xargo, see error details above")
        } else {
            show_error!("failed to run xargo; run `cargo miri setup` to see the error details")
        }
    }

    // Figure out what to print.
    if only_setup {
        eprintln!("A sysroot for Miri is now available in `{}`.", sysroot.display());
    } else {
        eprintln!("done");
    }
    if print_sysroot {
        // Print just the sysroot and nothing else to stdout; this way we do not need any escaping.
        println!("{}", sysroot.display());
    }
}
<commit_msg>Clippy after rustup<commit_after>//! Implements `cargo miri setup` via xargo

use std::env;
use std::ffi::OsStr;
use std::fs::{self};
use std::io::BufRead;
use std::ops::Not;
use std::path::{Path, PathBuf};
use std::process::{self, Command};

use crate::{util::*, version::*};

fn xargo_version() -> Option<(u32, u32, u32)> {
    let out = xargo_check().arg("--version").output().ok()?;
    if !out.status.success() {
        return None;
    }
    // Parse output. The first line looks like "xargo 0.3.12 (b004f1c 2018-12-13)".
    let line = out
        .stderr
        .lines()
        .next()
        .expect("malformed `xargo --version` output: not at least one line")
        .expect("malformed `xargo --version` output: error reading first line");
    let (name, version) = {
        let mut split = line.split(' ');
        (
            split.next().expect("malformed `xargo --version` output: empty"),
            split.next().expect("malformed `xargo --version` output: not at least two words"),
        )
    };
    if name != "xargo" {
        // This is some fork of xargo
        return None;
    }
    let mut version_pieces = version.split('.');
    let major = version_pieces
        .next()
        .expect("malformed `xargo --version` output: not a major version piece")
        .parse()
        .expect("malformed `xargo --version` output: major version is not an integer");
    let minor = version_pieces
        .next()
        .expect("malformed `xargo --version` output: not a minor version piece")
        .parse()
        .expect("malformed `xargo --version` output: minor version is not an integer");
    let patch = version_pieces
        .next()
        .expect("malformed `xargo --version` output: not a patch version piece")
        .parse()
        .expect("malformed `xargo --version` output: patch version is not an integer");
    if version_pieces.next().is_some() {
        panic!("malformed `xargo --version` output: more than three pieces in version");
    }
    Some((major, minor, patch))
}

/// Performs the setup required to make `cargo miri` work: Getting a custom-built libstd. Then sets
/// `MIRI_SYSROOT`. Skipped if `MIRI_SYSROOT` is already set, in which case we expect the user has
/// done all this already.
pub fn setup(subcommand: &MiriCommand, host: &str, target: &str) {
    let only_setup = matches!(subcommand, MiriCommand::Setup);
    let ask_user = !only_setup;
    let print_sysroot = only_setup && has_arg_flag("--print-sysroot"); // whether we just print the sysroot path
    if std::env::var_os("MIRI_SYSROOT").is_some() {
        if only_setup {
            println!("WARNING: MIRI_SYSROOT already set, not doing anything.")
        }
        return;
    }

    // First, we need xargo.
    if xargo_version().map_or(true, |v| v < XARGO_MIN_VERSION) {
        if std::env::var_os("XARGO_CHECK").is_some() {
            // The user manually gave us a xargo binary; don't do anything automatically.
            show_error!("xargo is too old; please upgrade to the latest version")
        }
        let mut cmd = cargo();
        cmd.args(["install", "xargo"]);
        ask_to_run(cmd, ask_user, "install a recent enough xargo");
    }

    // Determine where the rust sources are located.  The env vars manually setting the source
    // (`MIRI_LIB_SRC`, `XARGO_RUST_SRC`) trump auto-detection.
    let rust_src_env_var =
        std::env::var_os("MIRI_LIB_SRC").or_else(|| std::env::var_os("XARGO_RUST_SRC"));
    let rust_src = match rust_src_env_var {
        Some(path) => {
            let path = PathBuf::from(path);
            // Make path absolute if possible.
            path.canonicalize().unwrap_or(path)
        }
        None => {
            // Check for `rust-src` rustup component.
            let output = miri_for_host()
                .args(["--print", "sysroot"])
                .output()
                .expect("failed to determine sysroot");
            if !output.status.success() {
                show_error!(
                    "Failed to determine sysroot; Miri said:\n{}",
                    String::from_utf8_lossy(&output.stderr).trim_end()
                );
            }
            let sysroot = std::str::from_utf8(&output.stdout).unwrap();
            let sysroot = Path::new(sysroot.trim_end_matches('\n'));
            // Check for `$SYSROOT/lib/rustlib/src/rust/library`; test if that contains `std/Cargo.toml`.
            let rustup_src =
                sysroot.join("lib").join("rustlib").join("src").join("rust").join("library");
            if !rustup_src.join("std").join("Cargo.toml").exists() {
                // Ask the user to install the `rust-src` component, and use that.
                let mut cmd = Command::new("rustup");
                cmd.args(["component", "add", "rust-src"]);
                ask_to_run(
                    cmd,
                    ask_user,
                    "install the `rust-src` component for the selected toolchain",
                );
            }
            rustup_src
        }
    };
    if !rust_src.exists() {
        show_error!("given Rust source directory `{}` does not exist.", rust_src.display());
    }
    if rust_src.file_name().and_then(OsStr::to_str) != Some("library") {
        show_error!(
            "given Rust source directory `{}` does not seem to be the `library` subdirectory of \
             a Rust source checkout.",
            rust_src.display()
        );
    }

    // Next, we need our own libstd. Prepare a xargo project for that purpose.
    // We will do this work in whatever is a good cache dir for this platform.
    let dirs = directories::ProjectDirs::from("org", "rust-lang", "miri").unwrap();
    let dir = dirs.cache_dir();
    if !dir.exists() {
        fs::create_dir_all(dir).unwrap();
    }
    // The interesting bit: Xargo.toml (only needs content if we actually need std)
    let xargo_toml = if std::env::var_os("MIRI_NO_STD").is_some() {
        ""
    } else {
        r#"
[dependencies.std]
default_features = false
# We support unwinding, so enable that panic runtime.
features = ["panic_unwind", "backtrace"]

[dependencies.test]
"#
    };
    write_to_file(&dir.join("Xargo.toml"), xargo_toml);
    // The boring bits: a dummy project for xargo.
    // FIXME: With xargo-check, can we avoid doing this?
    write_to_file(
        &dir.join("Cargo.toml"),
        r#"
[package]
name = "miri-xargo"
description = "A dummy project for building libstd with xargo."
version = "0.0.0"

[lib]
path = "lib.rs"
"#,
    );
    write_to_file(&dir.join("lib.rs"), "#![no_std]");

    // Figure out where xargo will build its stuff.
    // Unfortunately, it puts things into a different directory when the
    // architecture matches the host.
    let sysroot = if target == host { dir.join("HOST") } else { PathBuf::from(dir) };
    // Make sure all target-level Miri invocations know their sysroot.
    std::env::set_var("MIRI_SYSROOT", &sysroot);

    // Now invoke xargo.
    let mut command = xargo_check();
    command.arg("check").arg("-q");
    command.current_dir(dir);
    command.env("XARGO_HOME", dir);
    command.env("XARGO_RUST_SRC", &rust_src);
    // We always need to set a target so rustc bootstrap can tell apart host from target crates.
    command.arg("--target").arg(target);
    // Use Miri as rustc to build a libstd compatible with us (and use the right flags).
    // However, when we are running in bootstrap, we cannot just overwrite `RUSTC`,
    // because we still need bootstrap to distinguish between host and target crates.
    // In that case we overwrite `RUSTC_REAL` instead which determines the rustc used
    // for target crates.
    // We set ourselves (`cargo-miri`) instead of Miri directly to be able to patch the flags
    // for `libpanic_abort` (usually this is done by bootstrap but we have to do it ourselves).
    // The `MIRI_CALLED_FROM_XARGO` will mean we dispatch to `phase_setup_rustc`.
    let cargo_miri_path = std::env::current_exe().expect("current executable path invalid");
    if env::var_os("RUSTC_STAGE").is_some() {
        assert!(env::var_os("RUSTC").is_some());
        command.env("RUSTC_REAL", &cargo_miri_path);
    } else {
        command.env("RUSTC", &cargo_miri_path);
    }
    command.env("MIRI_CALLED_FROM_XARGO", "1");
    // Make sure there are no other wrappers getting in our way
    // (Cc https://github.com/rust-lang/miri/issues/1421, https://github.com/rust-lang/miri/issues/2429).
    // Looks like setting `RUSTC_WRAPPER` to the empty string overwrites `build.rustc-wrapper` set via `config.toml`.
    command.env("RUSTC_WRAPPER", "");
    // Disable debug assertions in the standard library -- Miri is already slow enough. But keep the
    // overflow checks, they are cheap. This completely overwrites flags the user might have set,
    // which is consistent with normal `cargo build` that does not apply `RUSTFLAGS` to the sysroot
    // either.
    command.env("RUSTFLAGS", "-Cdebug-assertions=off -Coverflow-checks=on");
    // Manage the output the user sees.
    if only_setup {
        // We want to be explicit.
        eprintln!("Preparing a sysroot for Miri (target: {target})...");
        if print_sysroot {
            // Be extra sure there is no noise on stdout.
            command.stdout(process::Stdio::null());
        }
    } else {
        // We want to be quiet, but still let the user know that something is happening.
        eprint!("Preparing a sysroot for Miri (target: {target})... ");
        command.stdout(process::Stdio::null());
        command.stderr(process::Stdio::null());
    }

    // Finally run it!
    if command.status().expect("failed to run xargo").success().not() {
        if only_setup {
            show_error!("failed to run xargo, see error details above")
        } else {
            show_error!("failed to run xargo; run `cargo miri setup` to see the error details")
        }
    }

    // Figure out what to print.
    if only_setup {
        eprintln!("A sysroot for Miri is now available in `{}`.", sysroot.display());
    } else {
        eprintln!("done");
    }
    if print_sysroot {
        // Print just the sysroot and nothing else to stdout; this way we do not need any escaping.
        println!("{}", sysroot.display());
    }
}
<|endoftext|>


<commit_before>use console::{Console, ConsoleTextKind};
use memmap::{Mmap, Protection};
use pipeline_matcher::PathMatch;
use std::fs;
use std::io;
use std::io::{Write, Error};
use std::process;
use std::sync::mpsc::{Receiver, Sender};
use time;
use util::{catch, decode_error, PipelineInfo};
use tempfile::NamedTempFile;

// ---------------------------------------------------------------------------------------------------------------------
// PipelineReplacer
// ---------------------------------------------------------------------------------------------------------------------

pub trait PipelineReplacer {
    fn replace( &mut self, replacement: &[u8], rx: Receiver<PipelineInfo<PathMatch>>, tx: Sender<PipelineInfo<()>> );
}

// ---------------------------------------------------------------------------------------------------------------------
// SimplePipelineReplacer
// ---------------------------------------------------------------------------------------------------------------------

pub struct SimplePipelineReplacer {
    pub is_color      : bool,
    pub is_interactive: bool,
    pub print_file    : bool,
    pub print_column  : bool,
    pub infos         : Vec<String>,
    pub errors        : Vec<String>,
    console           : Console,
    all_replace       : bool,
    time_beg          : u64,
    time_end          : u64,
    time_bsy          : u64,
}

impl SimplePipelineReplacer {
    pub fn new() -> Self {
        SimplePipelineReplacer {
            is_color      : true,
            is_interactive: true,
            print_file    : true,
            print_column  : true,
            infos         : Vec::new(),
            errors        : Vec::new(),
            console       : Console::new(),
            all_replace   : false,
            time_beg      : 0,
            time_end      : 0,
            time_bsy      : 0,
        }
    }

    fn replace_match( &mut self, replacement: &[u8], pm: PathMatch ) {
        if pm.matches.is_empty() { return; }
        self.console.is_color = self.is_color;

        let result = catch::<_, (), Error> ( || {
            let mut tmpfile = try!( NamedTempFile::new_in( pm.path.parent().unwrap_or( &pm.path ) )  );

            {
                let mmap = try!( Mmap::open_path( &pm.path, Protection::Read ) );
                let src  = unsafe { mmap.as_slice() };

                let mut i       = 0;
                let mut pos     = 0;
                let mut column  = 0;
                let mut last_lf = 0;
                for m in &pm.matches {
                    try!( tmpfile.write_all( &src[i..m.beg] ) );

                    let mut do_replace = true;
                    if self.is_interactive & !self.all_replace {
                        if self.print_file {
                            self.console.write( ConsoleTextKind::Filename, pm.path.to_str().unwrap() );
                            self.console.write( ConsoleTextKind::Other, ": " );
                        }
                        if self.print_column {
                            while pos < m.beg {
                                if src[pos] == 0x0a {
                                    column += 1;
                                    last_lf = pos;
                                }
                                pos += 1;
                            }
                            self.console.write( ConsoleTextKind::Other, &format!( "{}:{}:", column + 1, m.beg - last_lf ) );
                        }

                        self.console.write_match_line( src, m );

                        loop {
                            self.console.write( ConsoleTextKind::Other, "Replace keyword? ( Yes[Y], No[N], All[A], Quit[Q] ):" );
                            let mut buf = String::new();
                            io::stdin().read_line( &mut buf ).unwrap();
                            match buf.trim().to_lowercase().as_ref() {
                                "y"    => { do_replace = true ; break },
                                "yes"  => { do_replace = true ; break },
                                "n"    => { do_replace = false; break },
                                "no"   => { do_replace = false; break },
                                "a"    => { self.all_replace = true ; break },
                                "all"  => { self.all_replace = true ; break },
                                "q"    => { let _ = tmpfile.close(); process::exit( 0 ) },
                                "quit" => { let _ = tmpfile.close(); process::exit( 0 ) },
                                _      => continue,
                            }
                        }
                    }

                    if do_replace {
                        try!( tmpfile.write_all( &replacement ) );
                    } else {
                        try!( tmpfile.write_all( &src[m.beg..m.end] ) );
                    }
                    i = m.end;
                }

                if i < src.len() {
                    try!( tmpfile.write_all( &src[i..src.len()] ) );
                }
                try!( tmpfile.flush() );
            }

            let metadata = try!( fs::metadata( &pm.path ) );

            try!( fs::rename( tmpfile.path(), &pm.path ) );
            try!( fs::set_permissions( &pm.path, metadata.permissions() ) );

            Ok(())
        } );
        match result {
            Ok ( _ ) => (),
            Err( e ) => self.console.write( ConsoleTextKind::Error, &format!( "Error: {} @ {:?}\n", decode_error( e.kind() ), pm.path ) ),
        }
    }
}

impl PipelineReplacer for SimplePipelineReplacer {
    fn replace( &mut self, replacement: &[u8], rx: Receiver<PipelineInfo<PathMatch>>, tx: Sender<PipelineInfo<()>> ) {
        loop {
            match rx.recv() {
                Ok( PipelineInfo::Ok( pm ) ) => {
                    let beg = time::precise_time_ns();

                    self.replace_match( replacement, pm );
                    let _ = tx.send( PipelineInfo::Ok( () ) );

                    let end = time::precise_time_ns();
                    self.time_bsy += end - beg;
                },

                Ok( PipelineInfo::Beg( x ) ) => {
                    self.infos  = Vec::new();
                    self.errors = Vec::new();

                    self.time_beg = time::precise_time_ns();
                    let _ = tx.send( PipelineInfo::Beg( x ) );
                },

                Ok( PipelineInfo::End( x ) ) => {
                    for i in &self.infos  { let _ = tx.send( PipelineInfo::Info( i.clone() ) ); }
                    for e in &self.errors { let _ = tx.send( PipelineInfo::Err ( e.clone() ) ); }

                    self.time_end = time::precise_time_ns();
                    let _ = tx.send( PipelineInfo::Time( self.time_bsy, self.time_end - self.time_beg ) );
                    let _ = tx.send( PipelineInfo::End( x ) );
                    break;
                },

                Ok ( PipelineInfo::Info( e      ) ) => { let _ = tx.send( PipelineInfo::Info( e      ) ); },
                Ok ( PipelineInfo::Err ( e      ) ) => { let _ = tx.send( PipelineInfo::Err ( e      ) ); },
                Ok ( PipelineInfo::Time( t0, t1 ) ) => { let _ = tx.send( PipelineInfo::Time( t0, t1 ) ); },
                Err( _                            ) => break,
            }
        }
    }
}

// ---------------------------------------------------------------------------------------------------------------------
// Test
// ---------------------------------------------------------------------------------------------------------------------

//#[cfg(test)]
//mod tests {
//    use super::*;
//    use matcher::QuickSearchMatcher;
//    use std::path::PathBuf;
//    use std::thread;
//    use std::sync::mpsc;
//    use util::PipelineInfo;
//
//    #[test]
//    fn test_simple_pipeline_printer() {
//    }
//}

<commit_msg>Set permissions first.<commit_after>use console::{Console, ConsoleTextKind};
use memmap::{Mmap, Protection};
use pipeline_matcher::PathMatch;
use std::fs;
use std::io;
use std::io::{Write, Error};
use std::process;
use std::sync::mpsc::{Receiver, Sender};
use time;
use util::{catch, decode_error, PipelineInfo};
use tempfile::NamedTempFile;

// ---------------------------------------------------------------------------------------------------------------------
// PipelineReplacer
// ---------------------------------------------------------------------------------------------------------------------

pub trait PipelineReplacer {
    fn replace( &mut self, replacement: &[u8], rx: Receiver<PipelineInfo<PathMatch>>, tx: Sender<PipelineInfo<()>> );
}

// ---------------------------------------------------------------------------------------------------------------------
// SimplePipelineReplacer
// ---------------------------------------------------------------------------------------------------------------------

pub struct SimplePipelineReplacer {
    pub is_color      : bool,
    pub is_interactive: bool,
    pub print_file    : bool,
    pub print_column  : bool,
    pub infos         : Vec<String>,
    pub errors        : Vec<String>,
    console           : Console,
    all_replace       : bool,
    time_beg          : u64,
    time_end          : u64,
    time_bsy          : u64,
}

impl SimplePipelineReplacer {
    pub fn new() -> Self {
        SimplePipelineReplacer {
            is_color      : true,
            is_interactive: true,
            print_file    : true,
            print_column  : true,
            infos         : Vec::new(),
            errors        : Vec::new(),
            console       : Console::new(),
            all_replace   : false,
            time_beg      : 0,
            time_end      : 0,
            time_bsy      : 0,
        }
    }

    fn replace_match( &mut self, replacement: &[u8], pm: PathMatch ) {
        if pm.matches.is_empty() { return; }
        self.console.is_color = self.is_color;

        let result = catch::<_, (), Error> ( || {
            let mut tmpfile = try!( NamedTempFile::new_in( pm.path.parent().unwrap_or( &pm.path ) )  );

            {
                let mmap = try!( Mmap::open_path( &pm.path, Protection::Read ) );
                let src  = unsafe { mmap.as_slice() };

                let mut i       = 0;
                let mut pos     = 0;
                let mut column  = 0;
                let mut last_lf = 0;
                for m in &pm.matches {
                    try!( tmpfile.write_all( &src[i..m.beg] ) );

                    let mut do_replace = true;
                    if self.is_interactive & !self.all_replace {
                        if self.print_file {
                            self.console.write( ConsoleTextKind::Filename, pm.path.to_str().unwrap() );
                            self.console.write( ConsoleTextKind::Other, ": " );
                        }
                        if self.print_column {
                            while pos < m.beg {
                                if src[pos] == 0x0a {
                                    column += 1;
                                    last_lf = pos;
                                }
                                pos += 1;
                            }
                            self.console.write( ConsoleTextKind::Other, &format!( "{}:{}:", column + 1, m.beg - last_lf ) );
                        }

                        self.console.write_match_line( src, m );

                        loop {
                            self.console.write( ConsoleTextKind::Other, "Replace keyword? ( Yes[Y], No[N], All[A], Quit[Q] ):" );
                            let mut buf = String::new();
                            io::stdin().read_line( &mut buf ).unwrap();
                            match buf.trim().to_lowercase().as_ref() {
                                "y"    => { do_replace = true ; break },
                                "yes"  => { do_replace = true ; break },
                                "n"    => { do_replace = false; break },
                                "no"   => { do_replace = false; break },
                                "a"    => { self.all_replace = true ; break },
                                "all"  => { self.all_replace = true ; break },
                                "q"    => { let _ = tmpfile.close(); process::exit( 0 ) },
                                "quit" => { let _ = tmpfile.close(); process::exit( 0 ) },
                                _      => continue,
                            }
                        }
                    }

                    if do_replace {
                        try!( tmpfile.write_all( &replacement ) );
                    } else {
                        try!( tmpfile.write_all( &src[m.beg..m.end] ) );
                    }
                    i = m.end;
                }

                if i < src.len() {
                    try!( tmpfile.write_all( &src[i..src.len()] ) );
                }
                try!( tmpfile.flush() );
            }

            let metadata = try!( fs::metadata( &pm.path ) );

            try!( fs::set_permissions( tmpfile.path(), metadata.permissions() ) );
            try!( fs::rename( tmpfile.path(), &pm.path ) );

            Ok(())
        } );
        match result {
            Ok ( _ ) => (),
            Err( e ) => self.console.write( ConsoleTextKind::Error, &format!( "Error: {} @ {:?}\n", decode_error( e.kind() ), pm.path ) ),
        }
    }
}

impl PipelineReplacer for SimplePipelineReplacer {
    fn replace( &mut self, replacement: &[u8], rx: Receiver<PipelineInfo<PathMatch>>, tx: Sender<PipelineInfo<()>> ) {
        loop {
            match rx.recv() {
                Ok( PipelineInfo::Ok( pm ) ) => {
                    let beg = time::precise_time_ns();

                    self.replace_match( replacement, pm );
                    let _ = tx.send( PipelineInfo::Ok( () ) );

                    let end = time::precise_time_ns();
                    self.time_bsy += end - beg;
                },

                Ok( PipelineInfo::Beg( x ) ) => {
                    self.infos  = Vec::new();
                    self.errors = Vec::new();

                    self.time_beg = time::precise_time_ns();
                    let _ = tx.send( PipelineInfo::Beg( x ) );
                },

                Ok( PipelineInfo::End( x ) ) => {
                    for i in &self.infos  { let _ = tx.send( PipelineInfo::Info( i.clone() ) ); }
                    for e in &self.errors { let _ = tx.send( PipelineInfo::Err ( e.clone() ) ); }

                    self.time_end = time::precise_time_ns();
                    let _ = tx.send( PipelineInfo::Time( self.time_bsy, self.time_end - self.time_beg ) );
                    let _ = tx.send( PipelineInfo::End( x ) );
                    break;
                },

                Ok ( PipelineInfo::Info( e      ) ) => { let _ = tx.send( PipelineInfo::Info( e      ) ); },
                Ok ( PipelineInfo::Err ( e      ) ) => { let _ = tx.send( PipelineInfo::Err ( e      ) ); },
                Ok ( PipelineInfo::Time( t0, t1 ) ) => { let _ = tx.send( PipelineInfo::Time( t0, t1 ) ); },
                Err( _                            ) => break,
            }
        }
    }
}

// ---------------------------------------------------------------------------------------------------------------------
// Test
// ---------------------------------------------------------------------------------------------------------------------

//#[cfg(test)]
//mod tests {
//    use super::*;
//    use matcher::QuickSearchMatcher;
//    use std::path::PathBuf;
//    use std::thread;
//    use std::sync::mpsc;
//    use util::PipelineInfo;
//
//    #[test]
//    fn test_simple_pipeline_printer() {
//    }
//}

<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

trait Serializable<'self, T> { //~ ERROR lifetimes cannot use keyword names
    fn serialize(val : &'self T) -> Vec<u8>;
    //~^ ERROR lifetimes cannot use keyword names
    fn deserialize(repr : &[u8]) -> &'self T;
    //~^ ERROR lifetimes cannot use keyword names
}

impl<'self> Serializable<str> for &'self str { //~ ERROR lifetimes cannot use keyword names
    //~^ ERROR lifetimes cannot use keyword names
    //~| ERROR missing lifetime specifier
    fn serialize(val : &'self str) -> Vec<u8> {
    //~^ ERROR lifetimes cannot use keyword names
        vec![1]
    }
    fn deserialize(repr: &[u8]) -> &'self str {
    //~^ ERROR lifetimes cannot use keyword names
        "hi"
    }
}

fn main() {
    println!("hello");
    let x = "foo".to_string();
    let y = x;
    println!("{}", y);
}
<commit_msg>remove unnecessary changes<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

trait Serializable<'self, T> { //~ ERROR lifetimes cannot use keyword names
    fn serialize(val : &'self T) -> Vec<u8>; //~ ERROR lifetimes cannot use keyword names
    fn deserialize(repr : &[u8]) -> &'self T; //~ ERROR lifetimes cannot use keyword names
}

impl<'self> Serializable<str> for &'self str { //~ ERROR lifetimes cannot use keyword names
    //~^ ERROR lifetimes cannot use keyword names
    //~| ERROR missing lifetime specifier
    fn serialize(val : &'self str) -> Vec<u8> { //~ ERROR lifetimes cannot use keyword names
        vec![1]
    }
    fn deserialize(repr: &[u8]) -> &'self str { //~ ERROR lifetimes cannot use keyword names
        "hi"
    }
}

fn main() {
    println!("hello");
    let x = "foo".to_string();
    let y = x;
    println!("{}", y);
}
<|endoftext|>


<commit_before>#![allow(dead_code)]
#![allow(unused_must_use)]

use std::collections::HashMap;

use mio::{Token, Handler, EventLoop, EventSet, PollOpt};
use mio::tcp::{TcpListener, TcpStream};

use raftserver::{SERVER_TOKEN, FIRST_CUSTOM_TOKEN, DEFAULT_BASE_TICK_MS, INVALID_TOKEN};
use raftserver::{Msg, Sender, Result, ConnData, TimerMsg};
use raftserver::conn::Conn;
use raftserver::handler::ServerHandler;

pub struct Server<T: ServerHandler> {
    pub listener: TcpListener,
    pub conns: HashMap<Token, Conn>,
    pub token_counter: usize,
    pub sender: Sender,

    peers: HashMap<String, Token>,
    handler: T,
}

impl<T: ServerHandler> Server<T> {
    pub fn new(h: T, l: TcpListener, sender: Sender) -> Server<T> {
        Server {
            handler: h,
            listener: l,
            sender: sender,
            conns: HashMap::new(),
            peers: HashMap::new(),
            token_counter: FIRST_CUSTOM_TOKEN.as_usize(),
        }
    }

    pub fn register_tick(&mut self, event_loop: &mut EventLoop<Server<T>>) -> Result<()> {
        let token = Msg::Tick;
        // must ok, maybe check error later.
        event_loop.timeout_ms(token, DEFAULT_BASE_TICK_MS);
        Ok(())
    }

    fn register_timer(&mut self,
                      event_loop: &mut EventLoop<Server<T>>,
                      delay: u64,
                      msg: TimerMsg) {
        let token = Msg::Timer {
            delay: delay,
            msg: msg,
        };
        event_loop.timeout_ms(token, delay);
    }

    fn handle_error(&mut self, event_loop: &mut EventLoop<Server<T>>, token: Token) {
        let conn = self.conns.remove(&token);
        match conn {
            Some(conn) => {
                // if connected to remote peer, remove this too.
                if let Some(addr) = conn.peer_addr {
                    self.peers.remove(&addr);
                }

                event_loop.deregister(&conn.sock);
            }
            None => {
                warn!("missing connection for token {}", token.as_usize());
            }
        }

    }

    fn add_new_conn(&mut self,
                    event_loop: &mut EventLoop<Server<T>>,
                    sock: TcpStream,
                    peer_addr: Option<String>)
                    -> Result<(Token)> {
        let new_token = Token(self.token_counter);
        self.token_counter += 1;

        try!(sock.set_nodelay(true));

        let conn = Conn::new(sock, new_token, peer_addr);

        self.conns.insert(new_token, conn);

        try!(event_loop.register(&self.conns[&new_token].sock,
                                 new_token,
                                 EventSet::readable(),
                                 PollOpt::edge()));
        Ok(new_token)
    }

    fn handle_readeable(&mut self, event_loop: &mut EventLoop<Server<T>>, token: Token) {
        match token {
            SERVER_TOKEN => {
                loop {
                    // For edge trigger, we must accept all connections until None.
                    let sock = match self.listener.accept() {
                        Err(e) => {
                            error!("accept error: {:?}", e);
                            return;
                        }
                        Ok(None) => {
                            debug!("no connection, accept later.");
                            return;
                        }
                        Ok(Some((sock, addr))) => {
                            debug!("accept conn {}", addr);
                            sock
                        }
                    };

                    self.add_new_conn(event_loop, sock, None)
                        .map_err(|e| {
                            error!("register conn err {:?}", e);
                        });
                }
            }
            token => {
                let msgs;
                match self.conns.get_mut(&token) {
                    None => {
                        warn!("missing conn for token {:?}", token);
                        return;
                    }
                    Some(conn) => msgs = conn.readable(event_loop),
                }

                msgs.and_then(|msgs| {
                        if msgs.len() == 0 {
                            return Ok(msgs);
                        }

                        self.handler.handle_read_data(&self.sender, token, msgs)
                    })
                    .and_then(|res| {
                        if res.len() == 0 {
                            return Ok(());
                        }

                        // append to write buffer here, no need using sender to notify.
                        if let Some(conn) = self.conns.get_mut(&token) {
                            for data in res {
                                conn.append_write_buf(data);
                            }
                            // write buffer here directly?
                            try!(conn.reregister(event_loop));
                        }
                        Ok(())
                    })
                    .map_err(|e| warn!("handle read conn err {:?}", e));
            }

        }
    }

    fn handle_writeable(&mut self, event_loop: &mut EventLoop<Server<T>>, token: Token) {
        match self.conns.get_mut(&token) {
            None => {
                warn!("missing conn for token {:?}", token);
                return;
            }
            Some(conn) => conn.writeable(event_loop),
        };
    }

    fn handle_writedata(&mut self,
                        event_loop: &mut EventLoop<Server<T>>,
                        token: Token,
                        data: ConnData) {
        if let Some(conn) = self.conns.get_mut(&token) {
            conn.append_write_buf(data);
            // write buffer here directly?
            conn.reregister(event_loop);
        }
    }

    fn handle_tick(&mut self, event_loop: &mut EventLoop<Server<T>>) {
        self.handler
            .handle_tick(&self.sender)
            .map_err(|e| warn!("handle tick err {:?}", e));

        self.register_tick(event_loop);
    }

    fn handle_timer(&mut self, _: &mut EventLoop<Server<T>>, msg: TimerMsg) {
        self.handler
            .handle_timer(&self.sender, msg)
            .map_err(|e| warn!("handle timer err {:?}", e));
    }

    fn connect_peer(&mut self, event_loop: &mut EventLoop<Server<T>>, addr: &str) -> Result<Token> {
        let peer_addr = try!(addr.parse());
        let sock = try!(TcpStream::connect(&peer_addr));
        let token = try!(self.add_new_conn(event_loop, sock, Some(addr.to_string())));
        self.peers.insert(addr.to_string(), token);
        Ok(token)
    }

    fn handle_sendpeer(&mut self,
                       event_loop: &mut EventLoop<Server<T>>,
                       addr: String,
                       data: ConnData) {
        // check the corresponding token for peer address.
        let mut token = self.peers.get(&addr).map_or(INVALID_TOKEN, |t| *t);

        if token == INVALID_TOKEN {
            match self.connect_peer(event_loop, &addr) {
                Err(e) => {
                    error!("connect {:?} err {:?}", addr, e);
                    return;
                }
                Ok(new_token) => token = new_token,
            }
        }

        self.handle_writedata(event_loop, token, data);
    }
}

impl<T: ServerHandler> Handler for Server<T> {
    type Timeout = Msg;
    type Message = Msg;

    fn ready(&mut self, event_loop: &mut EventLoop<Server<T>>, token: Token, events: EventSet) {
        if events.is_hup() | events.is_error() {
            self.handle_error(event_loop, token);
            return;
        }

        if events.is_readable() {
            self.handle_readeable(event_loop, token);
        }

        if events.is_writable() {
            self.handle_writeable(event_loop, token);
        }
    }

    fn notify(&mut self, event_loop: &mut EventLoop<Server<T>>, msg: Msg) {
        match msg {
            Msg::Quit => event_loop.shutdown(),
            Msg::WriteData{token, data} => self.handle_writedata(event_loop, token, data),
            Msg::Timer{delay, msg} => self.register_timer(event_loop, delay, msg),
            Msg::SendPeer{addr, data} => self.handle_sendpeer(event_loop, addr, data),
            _ => panic!("unexpected msg"),
        }
    }

    fn timeout(&mut self, event_loop: &mut EventLoop<Server<T>>, msg: Msg) {
        match msg {
            Msg::Tick => self.handle_tick(event_loop),
            Msg::Timer{msg, ..} => self.handle_timer(event_loop, msg),
            _ => panic!("unexpected msg"),
        }
    }
}
<commit_msg>raftserver: remove conn when read/write error.<commit_after>#![allow(dead_code)]
#![allow(unused_must_use)]

use std::collections::HashMap;

use mio::{Token, Handler, EventLoop, EventSet, PollOpt};
use mio::tcp::{TcpListener, TcpStream};

use raftserver::{SERVER_TOKEN, FIRST_CUSTOM_TOKEN, DEFAULT_BASE_TICK_MS, INVALID_TOKEN};
use raftserver::{Msg, Sender, Result, ConnData, TimerMsg};
use raftserver::conn::Conn;
use raftserver::handler::ServerHandler;

pub struct Server<T: ServerHandler> {
    pub listener: TcpListener,
    pub conns: HashMap<Token, Conn>,
    pub token_counter: usize,
    pub sender: Sender,

    peers: HashMap<String, Token>,
    handler: T,
}

impl<T: ServerHandler> Server<T> {
    pub fn new(h: T, l: TcpListener, sender: Sender) -> Server<T> {
        Server {
            handler: h,
            listener: l,
            sender: sender,
            conns: HashMap::new(),
            peers: HashMap::new(),
            token_counter: FIRST_CUSTOM_TOKEN.as_usize(),
        }
    }

    pub fn register_tick(&mut self, event_loop: &mut EventLoop<Server<T>>) -> Result<()> {
        let token = Msg::Tick;
        // must ok, maybe check error later.
        event_loop.timeout_ms(token, DEFAULT_BASE_TICK_MS);
        Ok(())
    }

    fn register_timer(&mut self,
                      event_loop: &mut EventLoop<Server<T>>,
                      delay: u64,
                      msg: TimerMsg) {
        let token = Msg::Timer {
            delay: delay,
            msg: msg,
        };
        event_loop.timeout_ms(token, delay);
    }

    fn remove_conn(&mut self, event_loop: &mut EventLoop<Server<T>>, token: Token) {
        let conn = self.conns.remove(&token);
        match conn {
            Some(conn) => {
                // if connected to remote peer, remove this too.
                if let Some(addr) = conn.peer_addr {
                    self.peers.remove(&addr);
                }

                event_loop.deregister(&conn.sock);
            }
            None => {
                warn!("missing connection for token {}", token.as_usize());
            }
        }
    }

    fn add_new_conn(&mut self,
                    event_loop: &mut EventLoop<Server<T>>,
                    sock: TcpStream,
                    peer_addr: Option<String>)
                    -> Result<(Token)> {
        let new_token = Token(self.token_counter);
        self.token_counter += 1;

        try!(sock.set_nodelay(true));

        let conn = Conn::new(sock, new_token, peer_addr);

        self.conns.insert(new_token, conn);

        try!(event_loop.register(&self.conns[&new_token].sock,
                                 new_token,
                                 EventSet::readable(),
                                 PollOpt::edge())
                       .or_else(|e| {
                           // register error, remove from conn map.
                           self.conns.remove(&new_token);
                           Err(e)
                       }));
        Ok(new_token)
    }

    fn handle_readeable(&mut self, event_loop: &mut EventLoop<Server<T>>, token: Token) {
        match token {
            SERVER_TOKEN => {
                loop {
                    // For edge trigger, we must accept all connections until None.
                    let sock = match self.listener.accept() {
                        Err(e) => {
                            error!("accept error: {:?}", e);
                            return;
                        }
                        Ok(None) => {
                            debug!("no connection, accept later.");
                            return;
                        }
                        Ok(Some((sock, addr))) => {
                            debug!("accept conn {}", addr);
                            sock
                        }
                    };

                    self.add_new_conn(event_loop, sock, None)
                        .map_err(|e| {
                            error!("register conn err {:?}", e);
                        });
                }
            }
            token => {
                let msgs = match self.conns.get_mut(&token) {
                    None => {
                        warn!("missing conn for token {:?}", token);
                        return;
                    }
                    Some(conn) => conn.readable(event_loop),
                };

                msgs.and_then(|msgs| {
                        if msgs.len() == 0 {
                            return Ok(msgs);
                        }

                        self.handler.handle_read_data(&self.sender, token, msgs)
                    })
                    .and_then(|res| {
                        if res.len() == 0 {
                            return Ok(());
                        }

                        // append to write buffer here, no need using sender to notify.
                        if let Some(conn) = self.conns.get_mut(&token) {
                            for data in res {
                                conn.append_write_buf(data);
                            }
                            // write buffer here directly?
                            try!(conn.reregister(event_loop));
                        }
                        Ok(())
                    })
                    .map_err(|e| {
                        warn!("handle read conn err {:?}, remove", e);
                        self.remove_conn(event_loop, token);
                    });
            }

        }
    }

    fn handle_writeable(&mut self, event_loop: &mut EventLoop<Server<T>>, token: Token) {
        let res = match self.conns.get_mut(&token) {
            None => {
                warn!("missing conn for token {:?}", token);
                return;
            }
            Some(conn) => conn.writeable(event_loop),
        };

        res.map_err(|e| {
            warn!("handle write conn err {:?}, remove", e);
            self.remove_conn(event_loop, token);
        });
    }

    fn handle_writedata(&mut self,
                        event_loop: &mut EventLoop<Server<T>>,
                        token: Token,
                        data: ConnData) {
        if let Some(conn) = self.conns.get_mut(&token) {
            conn.append_write_buf(data);
            // write buffer here directly?
            conn.reregister(event_loop);
        }
    }

    fn handle_tick(&mut self, event_loop: &mut EventLoop<Server<T>>) {
        self.handler
            .handle_tick(&self.sender)
            .map_err(|e| warn!("handle tick err {:?}", e));

        self.register_tick(event_loop);
    }

    fn handle_timer(&mut self, _: &mut EventLoop<Server<T>>, msg: TimerMsg) {
        self.handler
            .handle_timer(&self.sender, msg)
            .map_err(|e| warn!("handle timer err {:?}", e));
    }

    fn connect_peer(&mut self, event_loop: &mut EventLoop<Server<T>>, addr: &str) -> Result<Token> {
        let peer_addr = try!(addr.parse());
        let sock = try!(TcpStream::connect(&peer_addr));
        let token = try!(self.add_new_conn(event_loop, sock, Some(addr.to_string())));
        self.peers.insert(addr.to_string(), token);
        Ok(token)
    }

    fn handle_sendpeer(&mut self,
                       event_loop: &mut EventLoop<Server<T>>,
                       addr: String,
                       data: ConnData) {
        // check the corresponding token for peer address.
        let mut token = self.peers.get(&addr).map_or(INVALID_TOKEN, |t| *t);

        if token == INVALID_TOKEN {
            match self.connect_peer(event_loop, &addr) {
                Err(e) => {
                    error!("connect {:?} err {:?}", addr, e);
                    return;
                }
                Ok(new_token) => token = new_token,
            }
        }

        self.handle_writedata(event_loop, token, data);
    }
}

impl<T: ServerHandler> Handler for Server<T> {
    type Timeout = Msg;
    type Message = Msg;

    fn ready(&mut self, event_loop: &mut EventLoop<Server<T>>, token: Token, events: EventSet) {
        if events.is_hup() || events.is_error() {
            self.remove_conn(event_loop, token);
            return;
        }

        if events.is_readable() {
            self.handle_readeable(event_loop, token);
        }

        if events.is_writable() {
            self.handle_writeable(event_loop, token);
        }
    }

    fn notify(&mut self, event_loop: &mut EventLoop<Server<T>>, msg: Msg) {
        match msg {
            Msg::Quit => event_loop.shutdown(),
            Msg::WriteData{token, data} => self.handle_writedata(event_loop, token, data),
            Msg::Timer{delay, msg} => self.register_timer(event_loop, delay, msg),
            Msg::SendPeer{addr, data} => self.handle_sendpeer(event_loop, addr, data),
            _ => panic!("unexpected msg"),
        }
    }

    fn timeout(&mut self, event_loop: &mut EventLoop<Server<T>>, msg: Msg) {
        match msg {
            Msg::Tick => self.handle_tick(event_loop),
            Msg::Timer{msg, ..} => self.handle_timer(event_loop, msg),
            _ => panic!("unexpected msg"),
        }
    }
}
<|endoftext|>


<commit_before>import std::map::hashmap;
import syntax::ast;
import ast::{ty, pat};
import syntax::codemap::{span};
import syntax::visit;
import syntax::print;

fn indent<R>(op: fn() -> R) -> R {
    // Use in conjunction with the log post-processor like `src/etc/indenter`
    // to make debug output more readable.
    #debug[">>"];
    let r <- op();
    #debug["<< (Result = %?)", r];
    ret r;
}

resource _indenter(_i: ()) {
    #debug["<<"];
}

fn indenter() -> _indenter {
    #debug[">>"];
    _indenter(())
}

type flag = hashmap<str, ()>;

fn field_expr(f: ast::field) -> @ast::expr { ret f.node.expr; }

fn field_exprs(fields: [ast::field]) -> [@ast::expr] {
    let mut es = [];
    for fields.each {|f| es += [f.node.expr]; }
    ret es;
}

fn has_nonlocal_exits(b: ast::blk) -> bool {
    let has_exits = @mut false;
    fn visit_expr(flag: @mut bool, e: @ast::expr) {
        alt e.node {
          ast::expr_break { *flag = true; }
          ast::expr_cont { *flag = true; }
          _ { }
        }
    }
    let v =
        visit::mk_simple_visitor(@{visit_expr: bind visit_expr(has_exits, _)
                                      with *visit::default_simple_visitor()});
    visit::visit_block(b, (), v);
    ret *has_exits;
}

/* FIXME: copy/paste, yuck */
fn may_break(b: ast::blk) -> bool {
    let has_exits = @mut false;
    fn visit_expr(flag: @mut bool, e: @ast::expr) {
        alt e.node {
          ast::expr_break { *flag = true; }
          _ { }
        }
    }
    let v =
        visit::mk_simple_visitor(@{visit_expr: bind visit_expr(has_exits, _)
                                      with *visit::default_simple_visitor()});
    visit::visit_block(b, (), v);
    ret *has_exits;
}

fn local_rhs_span(l: @ast::local, def: span) -> span {
    alt l.node.init { some(i) { ret i.expr.span; } _ { ret def; } }
}

fn is_main_name(path: syntax::ast_map::path) -> bool {
    // FIXME: path should be a constrained type, so we know
    // the call to last doesn't fail
    vec::last(path) == syntax::ast_map::path_name(@"main")
}

//
// Local Variables:
// mode: rust
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// End:
//
<commit_msg>Refactor; annotate a FIXME<commit_after>import std::map::hashmap;
import syntax::ast;
import ast::{ty, pat};
import syntax::codemap::{span};
import syntax::visit;
import syntax::print;

fn indent<R>(op: fn() -> R) -> R {
    // Use in conjunction with the log post-processor like `src/etc/indenter`
    // to make debug output more readable.
    #debug[">>"];
    let r <- op();
    #debug["<< (Result = %?)", r];
    ret r;
}

resource _indenter(_i: ()) {
    #debug["<<"];
}

fn indenter() -> _indenter {
    #debug[">>"];
    _indenter(())
}

type flag = hashmap<str, ()>;

fn field_expr(f: ast::field) -> @ast::expr { ret f.node.expr; }

fn field_exprs(fields: [ast::field]) -> [@ast::expr] {
    let mut es = [];
    for fields.each {|f| es += [f.node.expr]; }
    ret es;
}

// Takes a predicate p, returns true iff p is true for any subexpressions
// of b
fn block_expr_query(b: ast::blk, p: fn@(ast::expr_) -> bool) -> bool {
    let rs = @mut false;
    let visit_expr = {|flag: @mut bool, e: @ast::expr| *flag |= p(e.node)};
    let v =
        visit::mk_simple_visitor(@{visit_expr: bind visit_expr(rs, _)
                                      with *visit::default_simple_visitor()});
    visit::visit_block(b, (), v);
    ret *rs;
}

fn has_nonlocal_exits(b: ast::blk) -> bool {
    block_expr_query(b) {|e| alt e {
      ast::expr_break | ast::expr_cont { true }
      _ { false }}}
}

fn may_break(b: ast::blk) -> bool {
    block_expr_query(b) {|e| alt e {
      ast::expr_break { true }
      _ { false }}}
}

fn local_rhs_span(l: @ast::local, def: span) -> span {
    alt l.node.init { some(i) { ret i.expr.span; } _ { ret def; } }
}

fn is_main_name(path: syntax::ast_map::path) -> bool {
    // FIXME: path should be a constrained type, so we know
    // the call to last doesn't fail (#34)
    vec::last(path) == syntax::ast_map::path_name(@"main")
}

//
// Local Variables:
// mode: rust
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// End:
//
<|endoftext|>


<commit_before>use std::io::BufferedReader;
use std::io::net::ip::SocketAddr;
use std::io::net::tcp::TcpStream;
use std::option::Option;
use uuid::{Uuid, UuidVersion, Version4Random};
use super::super::events::*;
use super::parsers::{read_rpc, as_network_msg};
use super::types::*;

// Each peer should have one of these, and they should be consistent across
// nodes.
pub struct NetPeer {
    pub conf: ~NetPeerConfig,
    // If we have an open connection to this peer, then this will be Some(...).
    pub stream: Option<TcpStream>,
    // When we iterate over the NetPeers, we can check if we have a reply from
    // an AppendEntriesReq/etc. on this receiver.
    // In this case, this means when we send an AppendEntriesReq up to Schooner,
    // we need to assign this as the receiving end of the AppendEntriesReq's
    // to_leader channel.
    rpc_send: ~Sender<RaftRpc>,
}

impl NetPeer {
    /*
     * sender: the port we can use to send received commands back to the server.
     * should be a clone of a single Sender attached to a single Receiver, basically.
     * since each command will contain a port we can use for replies, we don't
     * need two ports in the arguments.
     */
    pub fn new(config: ~NetPeerConfig, sender: ~Sender<RaftRpc>) -> NetPeer {
        NetPeer {
            conf: config,
            stream: None,
            rpc_send: sender,
        }
    }

    /*
     * Try connecting to the peer from here. Note that this is only one way
     * we can establish a peer connection - the other way is if the peer tries
     * to connect to *us*.
     */
    pub fn try_spawn(&mut self) -> bool {
        if self.stream.is_some() {
            return false;
        }
        match TcpStream::connect_timeout(self.conf.address, 10000) {
            Ok(mut stream) => {
                stream.write_uint(self.conf.id);
                stream.write_line("");
                self.stream = Some(stream.clone());
                self.listen();
                true
            }
            err => {
                false
            }
        }
    }

    fn listen(&mut self) {
        if self.stream.is_none() {
            return;
        }
        let stream = self.stream.clone().unwrap();
        let sender = self.rpc_send.clone();
        spawn(proc() {
            loop {
                let either_rpc = read_rpc(stream.clone());
                match either_rpc {
                    Ok(rpc) => {
                        sender.send(rpc);
                    }
                    Err(e) => {
                        debug!("Dropped peer: {}", e);
                        break;
                    }
                }
            }
        });
        self.stream = None;
    }

    /*
     * If the node chose to connect to us, then we got a connection on our listening
     * address and need to give the stream to us here.
     *
     * Returns: True if we successfully connected, false if we thought we already had
     * an open connection to this peer (this is an invalid state; we should probably
     * crash or handle it somehow).
     */
    pub fn add_connection(&mut self, stream: &TcpStream) -> bool {
        if self.stream.is_some() {
            return false;
        }
        self.stream = Some(stream.clone());
        true
    }

    /*
     * Used by the leader to send commands to followers, and by candidates, etc.
     */
    pub fn send(&mut self, cmd: RaftRpc) -> Option<Receiver<RaftRpc>> {
        if self.stream.is_none() {
            return None;
        }
        let (rpc_send, rpc_recv) = channel();
        let stream = self.stream.clone().unwrap();
        spawn(proc() {
            //stream.write(as_network_msg(cmd));
            // TODO: replies.
            /*
            reply = // wait for a reply on the TCP connection
            // Probably we should break the channel if the TCPstream dies,
            // so the leader will know we didn't get a reply.
            rpc_send.send(reply);
            */
         });
        Some(rpc_recv)
    }
}

// TODO: Get the old parsing code out of the Git history and work it into
// this configuration.

#[cfg(test)]
mod test {
    use std::io::{TcpStream, BufferedReader, IoResult, IoError, InvalidInput};
    use std::io::net::ip::{SocketAddr, Ipv4Addr};
    use std::io::{Acceptor, Listener, TcpListener, TcpStream};
    use std::io::net::tcp::TcpAcceptor;

    use super::super::super::events::*;
    use uuid::{Uuid, UuidVersion, Version4Random};
    use super::NetPeer;
    use super::super::types::*;
    use super::super::parsers::{as_network_msg};

    /*
     * Can we parse content length fields?
     */
    #[test]
    fn test_spawn() {
        let pc = ~NetPeerConfig {
            id: 1,
            address: SocketAddr {
                ip: Ipv4Addr(127, 0, 0, 1),
                port: 8844,
            },
            client_addr: SocketAddr {
                ip: Ipv4Addr(127, 0, 0, 1),
                port: 8840,
            },
        };
        let (send1, recv1) = channel();
        let (send2, recv2) = channel();
        let mut peer1 = NetPeer::new(pc.clone(), ~send1);
        let mut peer2 = NetPeer::new(pc, ~send2);
        let listen_addr = SocketAddr {
            ip: Ipv4Addr(127, 0, 0, 1),
            port: 8844,
        };
        let listener: TcpListener = TcpListener::bind(listen_addr).unwrap();
        let mut acceptor: TcpAcceptor = listener.listen().unwrap();
        // Spawn two peers
        peer1.try_spawn();
        peer2.try_spawn();
        let mut count = 0;
        // Send each peer the vote
        let vote = RpcVRQ(VoteReq {
            term: 0,
            candidate_id: 0,
            last_log_index: 0,
            last_log_term: 0,
            uuid: Uuid::new(Version4Random).unwrap(),
        });
        for mut stream in acceptor.incoming() {
            let vote_bytes = as_network_msg(vote.clone());
            stream.write(vote_bytes);
            count += 1;
            if count > 1 {
                break;
            }
        }
        let mut replies = 0;
        // We should get the votes back out on the port that we were waiting on
        assert!(recv1.recv() == vote);
        assert!(recv2.recv() == vote);
        drop(acceptor);
    }
}
<commit_msg>rm peer.rs; unused<commit_after><|endoftext|>


<commit_before>pub struct IndexStats {
    /// Total number of active documents in the index
    pub total_docs: u64,
}

pub struct FieldStats {
    pub sum_total_term_freq: u64,
}

pub struct FieldTermStats {
    /// Total number of docs that contain this term in this field
    pub total_docs: u64,
}

#[derive(Debug, PartialEq)]
pub enum SimilarityModel {
    TF_IDF,
    BM25{k1: f64, b: f64},
}

/// idf(term_docs, total_docs) = log((total_docs + 1.0) / (term_docs + 1.0)) + 1.0
fn idf(term_docs: u64, total_docs: u64) -> f64 {
    let term_docs = term_docs as f64;
    let total_docs = total_docs as f64;
    ((total_docs + 1.0) / (term_docs + 1.0)).log(10.0) + 1.0
}

impl SimilarityModel {
    pub fn score(&self, term_frequency: u32, doc_length: u32, index_stats: &IndexStats, field_stats: &FieldStats, field_term_stats: &FieldTermStats) -> f64 {
        match *self {
            SimilarityModel::TF_IDF => {
                let tf = (term_frequency as f64).sqrt();
                let idf = idf(field_term_stats.total_docs, index_stats.total_docs);

                tf * idf
            }
            SimilarityModel::BM25{k1, b} => {
                let tf = (term_frequency as f64).sqrt();
                let idf = idf(field_term_stats.total_docs, index_stats.total_docs);

                let norm_value = 1.0f64 /* field boost */ / tf;
                let average_length = field_stats.sum_total_term_freq as f64 / index_stats.total_docs as f64;

                idf * (k1 + 1.0) * (tf / (tf + (k1 * ((1.0 - b) + b * norm_value / average_length))))
            }
        }
    }
}
<commit_msg>Tweaks to BM25<commit_after>pub struct IndexStats {
    /// Total number of active documents in the index
    pub total_docs: u64,
}

pub struct FieldStats {
    pub sum_total_term_freq: u64,
}

pub struct FieldTermStats {
    /// Total number of docs that contain this term in this field
    pub total_docs: u64,
}

#[derive(Debug, PartialEq)]
pub enum SimilarityModel {
    TF_IDF,
    BM25{k1: f64, b: f64},
}

/// idf(term_docs, total_docs) = log((total_docs + 1.0) / (term_docs + 1.0)) + 1.0
fn idf(term_docs: u64, total_docs: u64) -> f64 {
    let term_docs = term_docs as f64;
    let total_docs = total_docs as f64;
    ((total_docs + 1.0) / (term_docs + 1.0)).log(10.0) + 1.0
}

impl SimilarityModel {
    pub fn score(&self, term_frequency: u32, doc_length: u32, index_stats: &IndexStats, field_stats: &FieldStats, field_term_stats: &FieldTermStats) -> f64 {
        match *self {
            SimilarityModel::TF_IDF => {
                let tf = (term_frequency as f64).sqrt();
                let idf = idf(field_term_stats.total_docs, index_stats.total_docs);

                tf * idf
            }
            SimilarityModel::BM25{k1, b} => {
                let tf = (term_frequency as f64).sqrt();
                let idf = idf(field_term_stats.total_docs, index_stats.total_docs);
                let average_length = field_stats.sum_total_term_freq as f64 / index_stats.total_docs as f64;

                idf * (k1 + 1.0) * (tf / (tf + (k1 * ((1.0 - b) + b * (doc_length as f64) / average_length))))
            }
        }
    }
}
<|endoftext|>