<commit_before>/// Types related to the construction of an `INSERT` statement.

use std::any::*;

use backend::Backend;
use expression::{Expression, NonAggregate, SelectableExpression};
use expression::operators::Eq;
use insertable::*;
#[cfg(feature = "mysql")]
use mysql::Mysql;
use query_builder::*;
#[cfg(feature = "sqlite")]
use query_dsl::ExecuteDsl;
use query_source::{Column, Table};
use result::QueryResult;
#[cfg(feature = "sqlite")]
use sqlite::{Sqlite, SqliteConnection};
use super::returning_clause::*;

/// The structure returned by [`insert_into`](../../fn.insert_into.html).
///
/// The provided methods [`values()`] and [`default_values()`] will insert
/// data into the targeted table.
#[derive(Debug, Clone, Copy)]
pub struct IncompleteInsertStatement<T, Op> {
    target: T,
    operator: Op,
}

impl<T, Op> IncompleteInsertStatement<T, Op> {
    pub(crate) fn new(target: T, operator: Op) -> Self {
        IncompleteInsertStatement { target, operator }
    }

    /// Inserts `DEFAULT VALUES` into the targeted table.
    ///
    /// ```rust
    /// # #[macro_use] extern crate diesel;
    /// # include!("../doctest_setup.rs");
    /// #
    /// # table! {
    /// #     users (name) {
    /// #         name -> Text,
    /// #         hair_color -> Text,
    /// #     }
    /// # }
    /// #
    /// # fn main() {
    /// #     use diesel::insert_into;
    /// #     use users::dsl::*;
    /// #     let connection = connection_no_data();
    /// connection.execute("CREATE TABLE users (
    ///     name VARCHAR(255) NOT NULL DEFAULT 'Sean',
    ///     hair_color VARCHAR(255) NOT NULL DEFAULT 'Green'
    /// )").unwrap();
    ///
    /// insert_into(users)
    ///     .default_values()
    ///     .execute(&connection)
    ///     .unwrap();
    ///
    /// let expected_data = vec![
    ///     ("Sean".to_string(), "Green".to_string()),
    /// ];
    /// assert_eq!(Ok(expected_data), users.load(&connection));
    /// # }
    /// ```
    pub fn default_values(self) -> InsertStatement<T, DefaultValues, Op> {
        static STATIC_DEFAULT_VALUES: &DefaultValues = &DefaultValues;
        self.values(STATIC_DEFAULT_VALUES)
    }

    /// Inserts the given values into the table passed to `insert_into`.
    ///
    /// See the documentation of [`insert_into`](../fn.insert_into.html) for
    /// usage examples.
    pub fn values<U>(self, records: U) -> InsertStatement<T, U::Values, Op>
    where
        U: Insertable<T>,
    {
        InsertStatement::new(
            self.target,
            records.values(),
            self.operator,
            NoReturningClause,
        )
    }
}

#[derive(Debug, Copy, Clone)]
pub struct InsertStatement<T, U, Op = Insert, Ret = NoReturningClause> {
    operator: Op,
    target: T,
    records: U,
    returning: Ret,
}

impl<T, U, Op, Ret> InsertStatement<T, U, Op, Ret> {
    fn new(target: T, records: U, operator: Op, returning: Ret) -> Self {
        InsertStatement {
            operator: operator,
            target: target,
            records: records,
            returning: returning,
        }
    }

    #[cfg(feature = "postgres")]
    pub(crate) fn replace_values<F, V>(self, f: F) -> InsertStatement<T, V, Op, Ret>
    where
        F: FnOnce(U) -> V,
    {
        InsertStatement::new(self.target, f(self.records), self.operator, self.returning)
    }
}

impl<T, U, Op, Ret, DB> QueryFragment<DB> for InsertStatement<T, U, Op, Ret>
where
    DB: Backend,
    T: Table,
    T::FromClause: QueryFragment<DB>,
    U: InsertValues<T, DB> + CanInsertInSingleQuery<DB>,
    Op: QueryFragment<DB>,
    Ret: QueryFragment<DB>,
{
    fn walk_ast(&self, mut out: AstPass<DB>) -> QueryResult<()> {
        out.unsafe_to_cache_prepared();

        if self.records.rows_to_insert() == 0 {
            out.push_sql("SELECT 1 FROM ");
            self.target.from_clause().walk_ast(out.reborrow())?;
            out.push_sql(" WHERE 1=0");
            return Ok(());
        }

        self.operator.walk_ast(out.reborrow())?;
        out.push_sql(" INTO ");
        self.target.from_clause().walk_ast(out.reborrow())?;
        if self.records.is_noop() {
            out.push_sql(" DEFAULT VALUES");
        } else {
            out.push_sql(" (");
            if let Some(builder) = out.reborrow().query_builder() {
                self.records.column_names(builder)?;
            }
            out.push_sql(") VALUES ");
            if self.records.requires_parenthesis() {
                out.push_sql("(");
            }
            self.records.walk_ast(out.reborrow())?;
            if self.records.requires_parenthesis() {
                out.push_sql(")");
            }
        }
        self.returning.walk_ast(out.reborrow())?;
        Ok(())
    }
}

#[cfg(feature = "sqlite")]
impl<'a, T, U, Op> ExecuteDsl<SqliteConnection> for InsertStatement<T, &'a [U], Op>
where
    &'a U: Insertable<T>,
    InsertStatement<T, <&'a U as Insertable<T>>::Values, Op>: QueryFragment<Sqlite>,
    T: Copy,
    Op: Copy,
{
    fn execute(self, conn: &SqliteConnection) -> QueryResult<usize> {
        use connection::Connection;
        conn.transaction(|| {
            let mut result = 0;
            for record in self.records {
                result += InsertStatement::new(
                    self.target,
                    record.values(),
                    self.operator,
                    self.returning,
                ).execute(conn)?;
            }
            Ok(result)
        })
    }
}

impl_query_id!(noop: InsertStatement<T, U, Op, Ret>);

impl<T, U, Op> AsQuery for InsertStatement<T, U, Op, NoReturningClause>
where
    T: Table,
    InsertStatement<T, U, Op, ReturningClause<T::AllColumns>>: Query,
{
    type SqlType = <Self::Query as Query>::SqlType;
    type Query = InsertStatement<T, U, Op, ReturningClause<T::AllColumns>>;

    fn as_query(self) -> Self::Query {
        self.returning(T::all_columns())
    }
}

impl<T, U, Op, Ret> Query for InsertStatement<T, U, Op, ReturningClause<Ret>>
where
    Ret: Expression + SelectableExpression<T> + NonAggregate,
{
    type SqlType = Ret::SqlType;
}

impl<T, U, Op> InsertStatement<T, U, Op> {
    /// Specify what expression is returned after execution of the `insert`.
    /// # Examples
    ///
    /// ### Inserting records:
    ///
    /// ```rust
    /// # #[macro_use] extern crate diesel;
    /// # include!("../doctest_setup.rs");
    /// #
    /// # table! {
    /// #     users {
    /// #         id -> Integer,
    /// #         name -> VarChar,
    /// #     }
    /// # }
    /// #
    /// # #[cfg(feature = "postgres")]
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// #     let connection = establish_connection();
    /// let new_users = vec![
    ///     NewUser { name: "Timmy".to_string(), },
    ///     NewUser { name: "Jimmy".to_string(), },
    /// ];
    ///
    /// let inserted_names = diesel::insert_into(users)
    ///     .values(&new_users)
    ///     .returning(name)
    ///     .get_results(&connection);
    /// assert_eq!(Ok(vec!["Timmy".to_string(), "Jimmy".to_string()]), inserted_names);
    /// # }
    /// # #[cfg(not(feature = "postgres"))]
    /// # fn main() {}
    /// ```
    pub fn returning<E>(self, returns: E) -> InsertStatement<T, U, Op, ReturningClause<E>>
    where
        InsertStatement<T, U, Op, ReturningClause<E>>: Query,
    {
        InsertStatement::new(
            self.target,
            self.records,
            self.operator,
            ReturningClause(returns),
        )
    }
}

#[derive(Debug, Copy, Clone)]
#[doc(hidden)]
pub struct Insert;

impl<DB: Backend> QueryFragment<DB> for Insert {
    fn walk_ast(&self, mut out: AstPass<DB>) -> QueryResult<()> {
        out.push_sql("INSERT");
        Ok(())
    }
}

impl_query_id!(Insert);

#[derive(Debug, Copy, Clone)]
#[doc(hidden)]
pub struct Replace;

#[cfg(feature = "sqlite")]
impl QueryFragment<Sqlite> for Replace {
    fn walk_ast(&self, mut out: AstPass<Sqlite>) -> QueryResult<()> {
        out.push_sql("REPLACE");
        Ok(())
    }
}

#[cfg(feature = "mysql")]
impl QueryFragment<Mysql> for Replace {
    fn walk_ast(&self, mut out: AstPass<Mysql>) -> QueryResult<()> {
        out.push_sql("REPLACE");
        Ok(())
    }
}

impl_query_id!(Replace);

/// Marker trait to indicate that no additional operations have been added
/// to a record for insert. Used to prevent things like
/// `insert(&vec![user.on_conflict_do_nothing(), user2.on_conflict_do_nothing()])`
/// from compiling.
pub trait UndecoratedInsertRecord<Table> {}

impl<'a, T, Tab> UndecoratedInsertRecord<Tab> for &'a T
where
    T: ?Sized + UndecoratedInsertRecord<Tab>,
{
}

impl<T, U> UndecoratedInsertRecord<T::Table> for ColumnInsertValue<T, U>
where
    T: Column,
{
}

impl<T, Table> UndecoratedInsertRecord<Table> for [T]
where
    T: UndecoratedInsertRecord<Table>,
{
}

impl<T, Table> UndecoratedInsertRecord<Table> for Vec<T>
where
    [T]: UndecoratedInsertRecord<Table>,
{
}

impl<Lhs, Rhs> UndecoratedInsertRecord<Lhs::Table> for Eq<Lhs, Rhs>
where
    Lhs: Column,
{
}

impl<Lhs, Rhs, Tab> UndecoratedInsertRecord<Tab> for Option<Eq<Lhs, Rhs>>
where
    Eq<Lhs, Rhs>: UndecoratedInsertRecord<Tab>,
{
}

#[derive(Debug, Clone, Copy)]
#[doc(hidden)]
pub struct DefaultValues;

impl<DB: Backend> CanInsertInSingleQuery<DB> for DefaultValues {
    fn rows_to_insert(&self) -> usize {
        1
    }
}

impl<'a, Tab> Insertable<Tab> for &'a DefaultValues {
    type Values = DefaultValues;

    fn values(self) -> Self::Values {
        *self
    }
}

impl<Tab, DB> InsertValues<Tab, DB> for DefaultValues
where
    Tab: Table,
    DB: Backend + Any,
{
    fn column_names(&self, _: &mut DB::QueryBuilder) -> QueryResult<()> {
        Ok(())
    }

    fn walk_ast(&self, _: AstPass<DB>) -> QueryResult<()> {
        Ok(())
    }

    #[cfg(not(feature = "mysql"))]
    fn is_noop(&self) -> bool {
        true
    }

    #[cfg(feature = "mysql")]
    fn is_noop(&self) -> bool {
        // The syntax for this on MySQL is
        // INSERT INTO table () VALUES ()
        //
        // This is hacky, but it's the easiest way to get this done without a
        // deeper restructuring of this code.
        TypeId::of::<DB>() != TypeId::of::<::mysql::Mysql>()
    }
}
<commit_msg>Make module documentation document the module [ci skip]<commit_after>//! Types related to the construction of an `INSERT` statement.

use std::any::*;

use backend::Backend;
use expression::{Expression, NonAggregate, SelectableExpression};
use expression::operators::Eq;
use insertable::*;
#[cfg(feature = "mysql")]
use mysql::Mysql;
use query_builder::*;
#[cfg(feature = "sqlite")]
use query_dsl::ExecuteDsl;
use query_source::{Column, Table};
use result::QueryResult;
#[cfg(feature = "sqlite")]
use sqlite::{Sqlite, SqliteConnection};
use super::returning_clause::*;

/// The structure returned by [`insert_into`](../../fn.insert_into.html).
///
/// The provided methods [`values()`] and [`default_values()`] will insert
/// data into the targeted table.
#[derive(Debug, Clone, Copy)]
pub struct IncompleteInsertStatement<T, Op> {
    target: T,
    operator: Op,
}

impl<T, Op> IncompleteInsertStatement<T, Op> {
    pub(crate) fn new(target: T, operator: Op) -> Self {
        IncompleteInsertStatement { target, operator }
    }

    /// Inserts `DEFAULT VALUES` into the targeted table.
    ///
    /// ```rust
    /// # #[macro_use] extern crate diesel;
    /// # include!("../doctest_setup.rs");
    /// #
    /// # table! {
    /// #     users (name) {
    /// #         name -> Text,
    /// #         hair_color -> Text,
    /// #     }
    /// # }
    /// #
    /// # fn main() {
    /// #     use diesel::insert_into;
    /// #     use users::dsl::*;
    /// #     let connection = connection_no_data();
    /// connection.execute("CREATE TABLE users (
    ///     name VARCHAR(255) NOT NULL DEFAULT 'Sean',
    ///     hair_color VARCHAR(255) NOT NULL DEFAULT 'Green'
    /// )").unwrap();
    ///
    /// insert_into(users)
    ///     .default_values()
    ///     .execute(&connection)
    ///     .unwrap();
    ///
    /// let expected_data = vec![
    ///     ("Sean".to_string(), "Green".to_string()),
    /// ];
    /// assert_eq!(Ok(expected_data), users.load(&connection));
    /// # }
    /// ```
    pub fn default_values(self) -> InsertStatement<T, DefaultValues, Op> {
        static STATIC_DEFAULT_VALUES: &DefaultValues = &DefaultValues;
        self.values(STATIC_DEFAULT_VALUES)
    }

    /// Inserts the given values into the table passed to `insert_into`.
    ///
    /// See the documentation of [`insert_into`](../fn.insert_into.html) for
    /// usage examples.
    pub fn values<U>(self, records: U) -> InsertStatement<T, U::Values, Op>
    where
        U: Insertable<T>,
    {
        InsertStatement::new(
            self.target,
            records.values(),
            self.operator,
            NoReturningClause,
        )
    }
}

#[derive(Debug, Copy, Clone)]
pub struct InsertStatement<T, U, Op = Insert, Ret = NoReturningClause> {
    operator: Op,
    target: T,
    records: U,
    returning: Ret,
}

impl<T, U, Op, Ret> InsertStatement<T, U, Op, Ret> {
    fn new(target: T, records: U, operator: Op, returning: Ret) -> Self {
        InsertStatement {
            operator: operator,
            target: target,
            records: records,
            returning: returning,
        }
    }

    #[cfg(feature = "postgres")]
    pub(crate) fn replace_values<F, V>(self, f: F) -> InsertStatement<T, V, Op, Ret>
    where
        F: FnOnce(U) -> V,
    {
        InsertStatement::new(self.target, f(self.records), self.operator, self.returning)
    }
}

impl<T, U, Op, Ret, DB> QueryFragment<DB> for InsertStatement<T, U, Op, Ret>
where
    DB: Backend,
    T: Table,
    T::FromClause: QueryFragment<DB>,
    U: InsertValues<T, DB> + CanInsertInSingleQuery<DB>,
    Op: QueryFragment<DB>,
    Ret: QueryFragment<DB>,
{
    fn walk_ast(&self, mut out: AstPass<DB>) -> QueryResult<()> {
        out.unsafe_to_cache_prepared();

        if self.records.rows_to_insert() == 0 {
            out.push_sql("SELECT 1 FROM ");
            self.target.from_clause().walk_ast(out.reborrow())?;
            out.push_sql(" WHERE 1=0");
            return Ok(());
        }

        self.operator.walk_ast(out.reborrow())?;
        out.push_sql(" INTO ");
        self.target.from_clause().walk_ast(out.reborrow())?;
        if self.records.is_noop() {
            out.push_sql(" DEFAULT VALUES");
        } else {
            out.push_sql(" (");
            if let Some(builder) = out.reborrow().query_builder() {
                self.records.column_names(builder)?;
            }
            out.push_sql(") VALUES ");
            if self.records.requires_parenthesis() {
                out.push_sql("(");
            }
            self.records.walk_ast(out.reborrow())?;
            if self.records.requires_parenthesis() {
                out.push_sql(")");
            }
        }
        self.returning.walk_ast(out.reborrow())?;
        Ok(())
    }
}

#[cfg(feature = "sqlite")]
impl<'a, T, U, Op> ExecuteDsl<SqliteConnection> for InsertStatement<T, &'a [U], Op>
where
    &'a U: Insertable<T>,
    InsertStatement<T, <&'a U as Insertable<T>>::Values, Op>: QueryFragment<Sqlite>,
    T: Copy,
    Op: Copy,
{
    fn execute(self, conn: &SqliteConnection) -> QueryResult<usize> {
        use connection::Connection;
        conn.transaction(|| {
            let mut result = 0;
            for record in self.records {
                result += InsertStatement::new(
                    self.target,
                    record.values(),
                    self.operator,
                    self.returning,
                ).execute(conn)?;
            }
            Ok(result)
        })
    }
}

impl_query_id!(noop: InsertStatement<T, U, Op, Ret>);

impl<T, U, Op> AsQuery for InsertStatement<T, U, Op, NoReturningClause>
where
    T: Table,
    InsertStatement<T, U, Op, ReturningClause<T::AllColumns>>: Query,
{
    type SqlType = <Self::Query as Query>::SqlType;
    type Query = InsertStatement<T, U, Op, ReturningClause<T::AllColumns>>;

    fn as_query(self) -> Self::Query {
        self.returning(T::all_columns())
    }
}

impl<T, U, Op, Ret> Query for InsertStatement<T, U, Op, ReturningClause<Ret>>
where
    Ret: Expression + SelectableExpression<T> + NonAggregate,
{
    type SqlType = Ret::SqlType;
}

impl<T, U, Op> InsertStatement<T, U, Op> {
    /// Specify what expression is returned after execution of the `insert`.
    /// # Examples
    ///
    /// ### Inserting records:
    ///
    /// ```rust
    /// # #[macro_use] extern crate diesel;
    /// # include!("../doctest_setup.rs");
    /// #
    /// # table! {
    /// #     users {
    /// #         id -> Integer,
    /// #         name -> VarChar,
    /// #     }
    /// # }
    /// #
    /// # #[cfg(feature = "postgres")]
    /// # fn main() {
    /// #     use self::users::dsl::*;
    /// #     let connection = establish_connection();
    /// let new_users = vec![
    ///     NewUser { name: "Timmy".to_string(), },
    ///     NewUser { name: "Jimmy".to_string(), },
    /// ];
    ///
    /// let inserted_names = diesel::insert_into(users)
    ///     .values(&new_users)
    ///     .returning(name)
    ///     .get_results(&connection);
    /// assert_eq!(Ok(vec!["Timmy".to_string(), "Jimmy".to_string()]), inserted_names);
    /// # }
    /// # #[cfg(not(feature = "postgres"))]
    /// # fn main() {}
    /// ```
    pub fn returning<E>(self, returns: E) -> InsertStatement<T, U, Op, ReturningClause<E>>
    where
        InsertStatement<T, U, Op, ReturningClause<E>>: Query,
    {
        InsertStatement::new(
            self.target,
            self.records,
            self.operator,
            ReturningClause(returns),
        )
    }
}

#[derive(Debug, Copy, Clone)]
#[doc(hidden)]
pub struct Insert;

impl<DB: Backend> QueryFragment<DB> for Insert {
    fn walk_ast(&self, mut out: AstPass<DB>) -> QueryResult<()> {
        out.push_sql("INSERT");
        Ok(())
    }
}

impl_query_id!(Insert);

#[derive(Debug, Copy, Clone)]
#[doc(hidden)]
pub struct Replace;

#[cfg(feature = "sqlite")]
impl QueryFragment<Sqlite> for Replace {
    fn walk_ast(&self, mut out: AstPass<Sqlite>) -> QueryResult<()> {
        out.push_sql("REPLACE");
        Ok(())
    }
}

#[cfg(feature = "mysql")]
impl QueryFragment<Mysql> for Replace {
    fn walk_ast(&self, mut out: AstPass<Mysql>) -> QueryResult<()> {
        out.push_sql("REPLACE");
        Ok(())
    }
}

impl_query_id!(Replace);

/// Marker trait to indicate that no additional operations have been added
/// to a record for insert. Used to prevent things like
/// `insert(&vec![user.on_conflict_do_nothing(), user2.on_conflict_do_nothing()])`
/// from compiling.
pub trait UndecoratedInsertRecord<Table> {}

impl<'a, T, Tab> UndecoratedInsertRecord<Tab> for &'a T
where
    T: ?Sized + UndecoratedInsertRecord<Tab>,
{
}

impl<T, U> UndecoratedInsertRecord<T::Table> for ColumnInsertValue<T, U>
where
    T: Column,
{
}

impl<T, Table> UndecoratedInsertRecord<Table> for [T]
where
    T: UndecoratedInsertRecord<Table>,
{
}

impl<T, Table> UndecoratedInsertRecord<Table> for Vec<T>
where
    [T]: UndecoratedInsertRecord<Table>,
{
}

impl<Lhs, Rhs> UndecoratedInsertRecord<Lhs::Table> for Eq<Lhs, Rhs>
where
    Lhs: Column,
{
}

impl<Lhs, Rhs, Tab> UndecoratedInsertRecord<Tab> for Option<Eq<Lhs, Rhs>>
where
    Eq<Lhs, Rhs>: UndecoratedInsertRecord<Tab>,
{
}

#[derive(Debug, Clone, Copy)]
#[doc(hidden)]
pub struct DefaultValues;

impl<DB: Backend> CanInsertInSingleQuery<DB> for DefaultValues {
    fn rows_to_insert(&self) -> usize {
        1
    }
}

impl<'a, Tab> Insertable<Tab> for &'a DefaultValues {
    type Values = DefaultValues;

    fn values(self) -> Self::Values {
        *self
    }
}

impl<Tab, DB> InsertValues<Tab, DB> for DefaultValues
where
    Tab: Table,
    DB: Backend + Any,
{
    fn column_names(&self, _: &mut DB::QueryBuilder) -> QueryResult<()> {
        Ok(())
    }

    fn walk_ast(&self, _: AstPass<DB>) -> QueryResult<()> {
        Ok(())
    }

    #[cfg(not(feature = "mysql"))]
    fn is_noop(&self) -> bool {
        true
    }

    #[cfg(feature = "mysql")]
    fn is_noop(&self) -> bool {
        // The syntax for this on MySQL is
        // INSERT INTO table () VALUES ()
        //
        // This is hacky, but it's the easiest way to get this done without a
        // deeper restructuring of this code.
        TypeId::of::<DB>() != TypeId::of::<::mysql::Mysql>()
    }
}
<|endoftext|>


<commit_before>//! Data consumers
//!
//! The goal of data producers is to parse data depending on the previous result.
//! It can be used to selectively seek in a file.
//!
//! ## Example
//!
//! This consumer will take 4 samples from the input, print them, then stop
//!
//! ```rust
//!  use nom::{MemProducer,Consumer,ConsumerState};
//!  use std::str;
//!
//!  struct TestPrintConsumer {
//!    counter: usize
//!  }
//!
//!  impl TestPrintConsumer {
//!    fn new() -> TestPrintConsumer {
//!      TestPrintConsumer { counter: 0 }
//!    }
//!  }
//!
//!  // Return ConsumerState::Await if it needs more data, or ConsumerDone when it ends
//!  impl Consumer for TestPrintConsumer {
//!    fn consume(&mut self, input: &[u8]) -> ConsumerState {
//!      println!("{} -> {}", self.counter, str::from_utf8(input).unwrap());
//!      self.counter = self.counter + 1;
//!      if self.counter <=4 {
//!        ConsumerState::Await
//!      } else {
//!        ConsumerState::ConsumerDone
//!      }
//!    }
//!  }
//!
//!  // It can consume data directly from a producer
//!  let mut p = MemProducer::new("abcdefghijklmnopqrstuvwx".as_bytes(), 4);
//!  let mut c = TestPrintConsumer::new();
//!  c.run(&mut p);
//! ```

use self::ConsumerState::*;
use producer::Producer;
use producer::ProducerState::*;
use internal::Err;

/// Holds the current state of the consumer
///
/// * Await if more data is needed
///
/// * ConsumerDone if the consumer does not need anymore data to be parsed
///
/// * ConsumerError when something went wrong
#[derive(Debug,PartialEq,Eq,Copy)]
pub enum ConsumerState {
  Await,
  //Incomplete,
  ConsumerDone,
  ConsumerError(Err)
}

/// Implement the consume method, taking a byte array as input and returning a consumer state
///
/// The run function takes care of continuing or not
pub trait Consumer {
  fn consume(&mut self, input: &[u8]) -> ConsumerState;
  fn run(&mut self, producer: &mut Producer) {
    let mut acc: Vec<u8> = Vec::new();
    //let mut v2: Vec<u8>  = Vec::new();
    loop {
      let state = producer.produce();
      match state {
        Data(v) => {
          println!("got data");
          acc.push_all(v)
        },
        Eof([])  => {
          println!("eof empty");
          break;
        }
        Eof(v) => {
          println!("eof with {} bytes", v.len());
          acc.push_all(v)
        }
        _ => {break;}
      }
      //v2.push_all(acc.as_slice());
      //match consumer.consume(v2.as_slice()) {
      match self.consume(&acc[]) {
        ConsumerError(e) => {
          println!("consumer error, stopping: {}", e);
        },
        ConsumerDone => {
          println!("data, done");
          acc.clear();
          //acc.push_all(i);
          break;
        },
        Await => {
          println!("await");
          acc.clear();
          //acc.push_all(i);
        }
      }
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use super::ConsumerState::*;
  use producer::MemProducer;
  use std::str;

  struct TestPrintConsumer {
    counter: usize
  }

  impl TestPrintConsumer {
    fn new() -> TestPrintConsumer {
      TestPrintConsumer { counter: 0 }
    }
  }

  impl Consumer for TestPrintConsumer {
    fn consume(&mut self, input: &[u8]) -> ConsumerState {
      println!("{} -> {}", self.counter, str::from_utf8(input).unwrap());
      self.counter = self.counter + 1;
      if self.counter <=4 {
        Await
      } else {
        ConsumerDone
      }
    }
  }

  #[test]
  fn pull() {
    let mut p = MemProducer::new("abcdefghijklmnopqrstuvwx".as_bytes(), 4);
    let mut c = TestPrintConsumer::new();
    c.run(&mut p);
  }
}
<commit_msg>Add an end method to consumers<commit_after>//! Data consumers
//!
//! The goal of data producers is to parse data depending on the previous result.
//! It can be used to selectively seek in a file.
//!
//! ## Example
//!
//! This consumer will take 4 samples from the input, print them, then stop
//!
//! ```rust
//!  use nom::{MemProducer,Consumer,ConsumerState};
//!  use std::str;
//!
//!  struct TestPrintConsumer {
//!    counter: usize
//!  }
//!
//!  impl TestPrintConsumer {
//!    fn new() -> TestPrintConsumer {
//!      TestPrintConsumer { counter: 0 }
//!    }
//!  }
//!
//!  // Return ConsumerState::Await if it needs more data, or ConsumerDone when it ends
//!  impl Consumer for TestPrintConsumer {
//!    fn consume(&mut self, input: &[u8]) -> ConsumerState {
//!      println!("{} -> {}", self.counter, str::from_utf8(input).unwrap());
//!      self.counter = self.counter + 1;
//!      if self.counter <=4 {
//!        ConsumerState::Await
//!      } else {
//!        ConsumerState::ConsumerDone
//!      }
//!    }
//!
//!    fn end(&mut self) {
//!      println!("finished");
//!    }
//!  }
//!
//!  // It can consume data directly from a producer
//!  let mut p = MemProducer::new("abcdefghijklmnopqrstuvwx".as_bytes(), 4);
//!  let mut c = TestPrintConsumer::new();
//!  c.run(&mut p);
//! ```

use self::ConsumerState::*;
use producer::Producer;
use producer::ProducerState::*;
use internal::Err;

/// Holds the current state of the consumer
///
/// * Await if more data is needed
///
/// * ConsumerDone if the consumer does not need anymore data to be parsed
///
/// * ConsumerError when something went wrong
#[derive(Debug,PartialEq,Eq,Copy)]
pub enum ConsumerState {
  Await,
  //Incomplete,
  ConsumerDone,
  ConsumerError(Err)
}

/// Implement the consume method, taking a byte array as input and returning a consumer state
///
/// The run function takes care of continuing or not
pub trait Consumer {
  fn consume(&mut self, input: &[u8]) -> ConsumerState;
  fn end(&mut self);

  fn run(&mut self, producer: &mut Producer) {
    let mut acc: Vec<u8> = Vec::new();
    //let mut v2: Vec<u8>  = Vec::new();
    loop {
      let state   = producer.produce();
      let mut eof = false;
      match state {
        Data(v) => {
          println!("got data");
          acc.push_all(v);
        },
        Eof([])  => {
          println!("eof empty");
          eof = true;
          break;
        }
        Eof(v) => {
          println!("eof with {} bytes", v.len());
          eof = true;
          acc.push_all(v);
        }
        _ => {break;}
      }
      //v2.push_all(acc.as_slice());
      //match consumer.consume(v2.as_slice()) {
      match self.consume(&acc[]) {
        ConsumerError(e) => {
          println!("consumer error, stopping: {}", e);
        },
        ConsumerDone => {
          println!("data, done");
          acc.clear();
          self.end();
          //acc.push_all(i);
          break;
        },
        Await => {
          println!("await");
          acc.clear();
          //acc.push_all(i);
        }
      }
    if eof { self.end(); }
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use super::ConsumerState::*;
  use producer::MemProducer;
  use std::str;

  struct TestPrintConsumer {
    counter: usize,
    ended:   bool
  }

  impl TestPrintConsumer {
    fn new() -> TestPrintConsumer {
      TestPrintConsumer { counter: 0, ended: false }
    }
  }

  impl Consumer for TestPrintConsumer {
    fn consume(&mut self, input: &[u8]) -> ConsumerState {
      println!("{} -> {}", self.counter, str::from_utf8(input).unwrap());
      self.counter = self.counter + 1;
      if self.counter <= 4 {
        Await
      } else {
        ConsumerDone
      }
    }

    fn end(&mut self) {
      assert_eq!(self.counter, 5);
      self.ended = true;
    }
  }

  #[test]
  fn pull() {
    let mut p = MemProducer::new("abcdefghijklmnopqrstuvwx".as_bytes(), 4);
    let mut c = TestPrintConsumer::new();
    c.run(&mut p);

    assert!(c.ended);
  }
}
<|endoftext|>


<commit_before>//! Data consumers
//!
//! The goal of data producers is to parse data depending on the previous result.
//! It can be used to selectively seek in a file.
//!
//! ## Example
//!
//! This consumer will take 4 samples from the input, print them, then stop
//!
//! ```rust
//!  use nom::{IResult,Needed,MemProducer,Consumer,ConsumerState};
//!  use std::str;
//!
//!  struct TestPrintConsumer {
//!    counter: usize
//!  }
//!
//!  impl TestPrintConsumer {
//!    fn new() -> TestPrintConsumer {
//!      TestPrintConsumer { counter: 0 }
//!    }
//!  }
//!
//!  fn take4(i:&[u8]) -> IResult<&[u8], &[u8]>{
//!    if i.len() < 4 {
//!      IResult::Incomplete(Needed::Size(4))
//!    } else {
//!      IResult::Done(&i[4..],&i[0..4])
//!    }
//!  }
//!
//!  // Return ConsumerState::Await if it needs more data, or ConsumerDone when it ends
//!  impl Consumer for TestPrintConsumer {
//!    fn consume(&mut self, input: &[u8]) -> ConsumerState {
//!      match take4(input) {
//!        IResult::Error(a)      => ConsumerState::ConsumerError(0),
//!        IResult::Incomplete(a) => ConsumerState::Await(0, 4),
//!        IResult::Done(i, o)    => {
//!          println!("{} -> {}", self.counter, str::from_utf8(o).unwrap());
//!          self.counter = self.counter + 1;
//!          if self.counter <= 4 {
//!            ConsumerState::Await(4, 4)
//!          } else {
//!            ConsumerState::ConsumerDone
//!          }
//!        }
//!      }
//!    }
//!
//!    fn end(&mut self) {
//!      println!("finished");
//!    }
//!  }
//!
//!  // It can consume data directly from a producer
//!  let mut p = MemProducer::new(b"abcdefghijklmnopqrstuvwx", 4);
//!  let mut c = TestPrintConsumer::new();
//!  c.run(&mut p);
//! ```

use self::ConsumerState::*;
use producer::Producer;
use producer::ProducerState::*;
use std::io::SeekFrom;

/// Holds the current state of the consumer
///
/// * Await if more data is needed
///
/// * ConsumerDone if the consumer does not need anymore data to be parsed
///
/// * ConsumerError when something went wrong
#[derive(Debug,PartialEq,Eq,Clone,Copy)]
pub enum ConsumerState {
  Await(
    usize,    // consumed
    usize     // needed buffer size
  ),
  Seek(
    usize,    // consumed
    SeekFrom, // new position
    usize     // needed buffer size
  ),
  Incomplete,
  ConsumerDone,
  ConsumerError(u32)
}

/// Implement the consume method, taking a byte array as input and returning a consumer state
///
/// The run function takes care of continuing or not
pub trait Consumer {
  fn consume(&mut self, input: &[u8]) -> ConsumerState;
  fn end(&mut self);

  fn run(&mut self, producer: &mut Producer) {
    let mut acc: Vec<u8>       = Vec::new();
    let mut position           = 0;
    let mut should_seek        = false;
    let mut consumed:usize     = 0;
    let mut needed:usize       = 0;
    let mut seek_from:SeekFrom = SeekFrom::Current(0);
    let mut eof = false;
    let mut end = false;

    loop {
      //self.getDataFromProducer(producer, seek_from, needed, acc);
      if !should_seek && acc.len() - consumed >= needed {
        //println!("buffer is large enough, skipping");
        let mut tmp = Vec::new();
        //println!("before:\n{}", acc.to_hex(16));
        //println!("after:\n{}", (&acc[consumed..acc.len()]).to_hex(16));
        tmp.extend(acc[consumed..acc.len()].iter().cloned());
        acc.clear();
        acc = tmp;
      } else {
        if should_seek {
          let _ = producer.seek(seek_from);
          //println!("seeking: {:?}", pos);
          should_seek = false;
          acc.clear();
        } else {
          let mut tmp = Vec::new();
          tmp.extend(acc[consumed..acc.len()].iter().cloned());
          acc.clear();
          acc = tmp;
        }

        loop {
          let state   = producer.produce();
          match state {
            Data(v) => {
              //println!("got data: {} bytes", v.len());
              acc.extend(v.iter().cloned());
              position = position + v.len();
            },
            Eof(v) => {
              if v.is_empty() {
                //println!("eof empty");
                //eof = true;
                self.end();
                return
              } else {
                //println!("eof with {} bytes", v.len());
                eof = true;
                acc.extend(v.iter().cloned());
                position = position + v.len();
                break;
              }
            }
            ProducerError(_) => {break;}
            Continue => { continue; }
          }
          //println!("acc size: {}", acc.len());
          if acc.len() >= needed { break; }
        }
      }

      //println!("full:\n{}", acc.to_hex(8));
      //println!("truncated:\n{}", (&acc[0..needed]).to_hex(16));
      match self.consume(&acc[0..needed]) {
        ConsumerError(_) => {
        //ConsumerError(e) => {
          //println!("consumer error, stopping: {}", e);
        },
        ConsumerDone => {
          //println!("data, done");
          end = true;
        },
        Seek(consumed_bytes, sf, needed_bytes) => {
          //println!("Seek: consumed {} bytes, got {:?} and asked {} bytes", consumed_bytes, sf, needed_bytes);
          seek_from = match sf {
            SeekFrom::Current(i) => SeekFrom::Current(i - (acc.len() - needed) as i64),
            a => a
          };
          should_seek = true;
          consumed = consumed_bytes;
          needed   = needed_bytes;
        },
        Await(consumed_bytes, needed_bytes) => {
          //println!("consumed: {} bytes | needed: {} bytes", consumed_bytes, needed_bytes);
          consumed = consumed_bytes;
          needed   = needed_bytes;
        },
        Incomplete => {
          //println!("incomplete");
        }
      }
      if eof {
        self.end();
        break;
      }
      if end {
        self.end();
        break;
      }
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use super::ConsumerState::*;
  use producer::MemProducer;
  use internal::{Needed,IResult};
  use std::str;

#[macro_export]
macro_rules! take(
  ($name:ident $count:expr) => (
    fn $name(i:&[u8]) -> IResult<&[u8], &[u8]>{
      if i.len() < $count {
        IResult::Incomplete(Needed::Size($count))
      } else {
        IResult::Done(&i[$count..],&i[0..$count])
      }
    }
  )
);

  struct TestPrintConsumer {
    counter: usize,
    ended:   bool
  }

  impl TestPrintConsumer {
    fn new() -> TestPrintConsumer {
      TestPrintConsumer { counter: 0, ended: false }
    }
  }

  take!(take4 4);

  impl Consumer for TestPrintConsumer {
    fn consume(&mut self, input: &[u8]) -> ConsumerState {
      match take4(input) {
        IResult::Error(_)      => ConsumerError(0),
        IResult::Incomplete(_) => Await(0, 4),
        IResult::Done(_, o)    => {
          println!("{} -> {}", self.counter, str::from_utf8(o).unwrap());
          self.counter = self.counter + 1;
          if self.counter <= 4 {
            Await(4, 4)
          } else {
            ConsumerDone
          }
        }
      }
    }

    fn end(&mut self) {
      assert_eq!(self.counter, 5);
      self.ended = true;
    }
  }

  #[test]
  fn pull() {
    let mut p = MemProducer::new(&b"abcdefghijklmnopqrstuvwx"[..], 4);
    let mut c = TestPrintConsumer::new();
    c.run(&mut p);

    assert!(c.ended);
  }
}
<commit_msg>a bit more explanations for  consumer states<commit_after>//! Data consumers
//!
//! The goal of data producers is to parse data depending on the previous result.
//! It can be used to selectively seek in a file.
//!
//! ## Example
//!
//! This consumer will take 4 samples from the input, print them, then stop
//!
//! ```rust
//!  use nom::{IResult,Needed,MemProducer,Consumer,ConsumerState};
//!  use std::str;
//!
//!  struct TestPrintConsumer {
//!    counter: usize
//!  }
//!
//!  impl TestPrintConsumer {
//!    fn new() -> TestPrintConsumer {
//!      TestPrintConsumer { counter: 0 }
//!    }
//!  }
//!
//!  fn take4(i:&[u8]) -> IResult<&[u8], &[u8]>{
//!    if i.len() < 4 {
//!      IResult::Incomplete(Needed::Size(4))
//!    } else {
//!      IResult::Done(&i[4..],&i[0..4])
//!    }
//!  }
//!
//!  // Return ConsumerState::Await if it needs more data, or ConsumerDone when it ends
//!  impl Consumer for TestPrintConsumer {
//!    fn consume(&mut self, input: &[u8]) -> ConsumerState {
//!      match take4(input) {
//!        IResult::Error(a)      => ConsumerState::ConsumerError(0),
//!        IResult::Incomplete(a) => ConsumerState::Await(0, 4),
//!        IResult::Done(i, o)    => {
//!          println!("{} -> {}", self.counter, str::from_utf8(o).unwrap());
//!          self.counter = self.counter + 1;
//!          if self.counter <= 4 {
//!            ConsumerState::Await(4, 4)
//!          } else {
//!            ConsumerState::ConsumerDone
//!          }
//!        }
//!      }
//!    }
//!
//!    fn end(&mut self) {
//!      println!("finished");
//!    }
//!  }
//!
//!  // It can consume data directly from a producer
//!  let mut p = MemProducer::new(b"abcdefghijklmnopqrstuvwx", 4);
//!  let mut c = TestPrintConsumer::new();
//!  c.run(&mut p);
//! ```

use self::ConsumerState::*;
use producer::Producer;
use producer::ProducerState::*;
use std::io::SeekFrom;

/// Holds the current state of the consumer
///
/// * Await(consumed, needed input data size) if more data is needed
///
/// * Seek(consumed, new position, needed input data size) if the consumer must move back or forth
///
/// * Incomplete if there is not enough data but not enough information for Await or Seek
///
/// * ConsumerDone if the consumer does not need anymore data to be parsed
///
/// * ConsumerError(error code) when something went wrong
#[derive(Debug,PartialEq,Eq,Clone,Copy)]
pub enum ConsumerState {
  Await(
    usize,    // consumed
    usize     // needed buffer size
  ),
  Seek(
    usize,    // consumed
    SeekFrom, // new position
    usize     // needed buffer size
  ),
  Incomplete,
  ConsumerDone,
  ConsumerError(u32)
}

/// Implement the consume method, taking a byte array as input and returning a consumer state
///
/// The run function takes care of continuing or not
pub trait Consumer {
  fn consume(&mut self, input: &[u8]) -> ConsumerState;
  fn end(&mut self);

  fn run(&mut self, producer: &mut Producer) {
    let mut acc: Vec<u8>       = Vec::new();
    let mut position           = 0;
    let mut should_seek        = false;
    let mut consumed:usize     = 0;
    let mut needed:usize       = 0;
    let mut seek_from:SeekFrom = SeekFrom::Current(0);
    let mut eof = false;
    let mut end = false;

    loop {
      //self.getDataFromProducer(producer, seek_from, needed, acc);
      if !should_seek && acc.len() - consumed >= needed {
        //println!("buffer is large enough, skipping");
        let mut tmp = Vec::new();
        //println!("before:\n{}", acc.to_hex(16));
        //println!("after:\n{}", (&acc[consumed..acc.len()]).to_hex(16));
        tmp.extend(acc[consumed..acc.len()].iter().cloned());
        acc.clear();
        acc = tmp;
      } else {
        if should_seek {
          let _ = producer.seek(seek_from);
          //println!("seeking: {:?}", pos);
          should_seek = false;
          acc.clear();
        } else {
          let mut tmp = Vec::new();
          tmp.extend(acc[consumed..acc.len()].iter().cloned());
          acc.clear();
          acc = tmp;
        }

        loop {
          let state   = producer.produce();
          match state {
            Data(v) => {
              //println!("got data: {} bytes", v.len());
              acc.extend(v.iter().cloned());
              position = position + v.len();
            },
            Eof(v) => {
              if v.is_empty() {
                //println!("eof empty");
                //eof = true;
                self.end();
                return
              } else {
                //println!("eof with {} bytes", v.len());
                eof = true;
                acc.extend(v.iter().cloned());
                position = position + v.len();
                break;
              }
            }
            ProducerError(_) => {break;}
            Continue => { continue; }
          }
          //println!("acc size: {}", acc.len());
          if acc.len() >= needed { break; }
        }
      }

      //println!("full:\n{}", acc.to_hex(8));
      //println!("truncated:\n{}", (&acc[0..needed]).to_hex(16));
      match self.consume(&acc[0..needed]) {
        ConsumerError(_) => {
        //ConsumerError(e) => {
          //println!("consumer error, stopping: {}", e);
        },
        ConsumerDone => {
          //println!("data, done");
          end = true;
        },
        Seek(consumed_bytes, sf, needed_bytes) => {
          //println!("Seek: consumed {} bytes, got {:?} and asked {} bytes", consumed_bytes, sf, needed_bytes);
          seek_from = match sf {
            SeekFrom::Current(i) => SeekFrom::Current(i - (acc.len() - needed) as i64),
            a => a
          };
          should_seek = true;
          consumed = consumed_bytes;
          needed   = needed_bytes;
        },
        Await(consumed_bytes, needed_bytes) => {
          //println!("consumed: {} bytes | needed: {} bytes", consumed_bytes, needed_bytes);
          consumed = consumed_bytes;
          needed   = needed_bytes;
        },
        Incomplete => {
          //println!("incomplete");
        }
      }
      if eof {
        self.end();
        break;
      }
      if end {
        self.end();
        break;
      }
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use super::ConsumerState::*;
  use producer::MemProducer;
  use internal::{Needed,IResult};
  use std::str;

#[macro_export]
macro_rules! take(
  ($name:ident $count:expr) => (
    fn $name(i:&[u8]) -> IResult<&[u8], &[u8]>{
      if i.len() < $count {
        IResult::Incomplete(Needed::Size($count))
      } else {
        IResult::Done(&i[$count..],&i[0..$count])
      }
    }
  )
);

  struct TestPrintConsumer {
    counter: usize,
    ended:   bool
  }

  impl TestPrintConsumer {
    fn new() -> TestPrintConsumer {
      TestPrintConsumer { counter: 0, ended: false }
    }
  }

  take!(take4 4);

  impl Consumer for TestPrintConsumer {
    fn consume(&mut self, input: &[u8]) -> ConsumerState {
      match take4(input) {
        IResult::Error(_)      => ConsumerError(0),
        IResult::Incomplete(_) => Await(0, 4),
        IResult::Done(_, o)    => {
          println!("{} -> {}", self.counter, str::from_utf8(o).unwrap());
          self.counter = self.counter + 1;
          if self.counter <= 4 {
            Await(4, 4)
          } else {
            ConsumerDone
          }
        }
      }
    }

    fn end(&mut self) {
      assert_eq!(self.counter, 5);
      self.ended = true;
    }
  }

  #[test]
  fn pull() {
    let mut p = MemProducer::new(&b"abcdefghijklmnopqrstuvwx"[..], 4);
    let mut c = TestPrintConsumer::new();
    c.run(&mut p);

    assert!(c.ended);
  }
}
<|endoftext|>


<commit_before>use std::{convert, io};
use std::fs::File;
use std::io::{Seek, Write};

use hash::{calculate_sha1, Sha1};
use metainfo::Metainfo;
use rand;
use rand::Rng;

pub const BLOCK_SIZE: u32 = 16384;

pub struct Download {
    pub our_peer_id: String,
    pub metainfo:    Metainfo,
    pieces:          Vec<Piece>,
    file:            File,
}

impl Download {
    pub fn new(our_peer_id: String, metainfo: Metainfo) -> Result<Download, Error> {
        let file_length = metainfo.info.length;
        let piece_length = metainfo.info.piece_length;
        let num_pieces = metainfo.info.num_pieces;

        // create pieces
        let mut pieces = vec![];
        for i in 0..num_pieces {
            let len = if i < (num_pieces - 1) {
                piece_length
            } else {
                (file_length - (piece_length as u64 * (num_pieces as u64 - 1))) as u32
            };
            pieces.push(Piece::new(i, len, piece_length, metainfo.info.pieces[i as usize].clone()));
        }

        // create file
        let file = try!(File::create(&metainfo.info.name));
        try!(file.set_len(metainfo.info.length));

        Ok(Download {
            our_peer_id: our_peer_id,
            metainfo:    metainfo,
            pieces:      pieces,
            file:        file,
        })
    }

    pub fn store(&mut self, piece_index: u32, block_index: u32, data: Vec<u8>) -> Result<bool, Error> {
        {
            let piece = &mut self.pieces[piece_index as usize];
            try!(piece.store(&mut self.file, block_index, data));
        }
        Ok(self.is_complete())
    }

    pub fn next_block_to_request(&self, peer_has_pieces: &[bool]) -> Option<(u32, u32, u32)> {
        match self.get_random_incomplete_piece(peer_has_pieces) {
            Some(piece) => match piece.next_block_to_request() {
                Some(block) => Some((piece.index, block.index, block.length)),
                None => None
            },
            None => None
        }
    }

    fn is_complete(&self) -> bool {
        for piece in self.pieces.iter() {
            if !piece.is_complete {
                return false
            }
        }
        true
    }

    fn get_random_incomplete_piece(&self, peer_has_pieces: &[bool]) -> Option<&Piece> {
        let mut random_num_generator = rand::thread_rng();

        let incomplete_pieces: Vec<&Piece> = self.pieces.iter().filter(|x| !x.is_complete && peer_has_pieces[x.index as usize]).collect();
        if incomplete_pieces.len() > 0 {
            let random_incomplete_piece_index = random_num_generator.gen_range(0, incomplete_pieces.len());
            let random_incomplete_piece = incomplete_pieces[random_incomplete_piece_index];
            Some(random_incomplete_piece)
        } else {
            None
        }
    }
}

struct Piece {
    index:        u32,
    piece_length: u32,
    hash:         Sha1,
    blocks:       Vec<Block>,
    is_complete:  bool,
}

impl Piece {
    fn new(index: u32, length: u32, piece_length: u32, hash: Sha1) -> Piece {
        // create blocks
        let mut blocks = vec![];
        let num_blocks = (length as f64 / BLOCK_SIZE as f64).ceil() as u32;
        for i in 0..num_blocks {
            let len = if i < (num_blocks - 1) {
                BLOCK_SIZE
            } else {
                length - (BLOCK_SIZE * (num_blocks - 1))
            };
            blocks.push(Block::new(i, len));
        }

        Piece {
            index:        index,
            piece_length: piece_length,
            hash:         hash,
            blocks:       blocks,
            is_complete:  false
        }
    }

    fn store(&mut self, file: &mut File, block_index: u32, data: Vec<u8>) -> Result<(), Error> {
        {
            // store data in the appropriate block
            let block = &mut self.blocks[block_index as usize];
            block.data = Some(data);
        }

        if self.have_all_blocks() {
            // concatenate data from blocks together
            let mut data = vec![];
            for block in self.blocks.iter() {
                data.extend(block.data.clone().unwrap());
            }

            // validate that piece data matches SHA1 hash
            if self.hash == calculate_sha1(&data) {
                println!("Piece {} is complete and correct, writing to the file.", self.index);
                let offset = self.index as u64 * self.piece_length as u64;
                try!(file.seek(io::SeekFrom::Start(offset)));
                try!(file.write_all(&data));
                self.clear_block_data();
                self.is_complete = true;
            } else {
                println!("Piece is corrupt, deleting downloaded piece data!");
                self.clear_block_data();
            }
        }
        Ok(())
    }

    fn next_block_to_request(&self) -> Option<&Block> {
        if self.is_complete {
            return None
        }

        for block in self.blocks.iter() {
            if block.data.is_none() {
                return Some(block)
            }
        }
        None
    }

    fn have_all_blocks(&self) -> bool {
        for block in self.blocks.iter() {
            if block.data.is_none() {
                return false
            }
        }
        true
    }

    fn clear_block_data(&mut self) {
        for block in self.blocks.iter_mut() {
            block.data = None;
        }
    }
}

struct Block {
    index:  u32,
    length: u32,
    data:   Option<Vec<u8>>,
}

impl Block {
    fn new(index: u32, length: u32) -> Block {
        Block {
            index:  index,
            length: length,
            data:   None,
        }
    }
}

#[derive(Debug)]
pub enum Error {
    IoError(io::Error),
}

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<commit_msg>Randomized Block selection for Peers to Download<commit_after>use std::{convert, io};
use std::fs::File;
use std::io::{Seek, Write};

use hash::{calculate_sha1, Sha1};
use metainfo::Metainfo;
use rand;
use rand::Rng;

pub const BLOCK_SIZE: u32 = 16384;

pub struct Download {
    pub our_peer_id: String,
    pub metainfo:    Metainfo,
    pieces:          Vec<Piece>,
    file:            File,
}

impl Download {
    pub fn new(our_peer_id: String, metainfo: Metainfo) -> Result<Download, Error> {
        let file_length = metainfo.info.length;
        let piece_length = metainfo.info.piece_length;
        let num_pieces = metainfo.info.num_pieces;

        // create pieces
        let mut pieces = vec![];
        for i in 0..num_pieces {
            let len = if i < (num_pieces - 1) {
                piece_length
            } else {
                (file_length - (piece_length as u64 * (num_pieces as u64 - 1))) as u32
            };
            pieces.push(Piece::new(i, len, piece_length, metainfo.info.pieces[i as usize].clone()));
        }

        // create file
        let file = try!(File::create(&metainfo.info.name));
        try!(file.set_len(metainfo.info.length));

        Ok(Download {
            our_peer_id: our_peer_id,
            metainfo:    metainfo,
            pieces:      pieces,
            file:        file,
        })
    }

    pub fn store(&mut self, piece_index: u32, block_index: u32, data: Vec<u8>) -> Result<bool, Error> {
        {
            let piece = &mut self.pieces[piece_index as usize];
            try!(piece.store(&mut self.file, block_index, data));
        }
        Ok(self.is_complete())
    }

    pub fn next_block_to_request(&self, peer_has_pieces: &[bool]) -> Option<(u32, u32, u32)> {
        match self.get_random_incomplete_piece(peer_has_pieces) {
            Some(piece) => match piece.get_random_incomplete_to_request() {
                Some(block) => Some((piece.index, block.index, block.length)),
                None => None
            },
            None => None
        }
    }

    fn is_complete(&self) -> bool {
        for piece in self.pieces.iter() {
            if !piece.is_complete {
                return false
            }
        }
        true
    }

    fn get_random_incomplete_piece(&self, peer_has_pieces: &[bool]) -> Option<&Piece> {
        let mut random_num_generator = rand::thread_rng();

        let incomplete_pieces: Vec<&Piece> = self.pieces.iter().filter(|x| !x.is_complete && peer_has_pieces[x.index as usize]).collect();
        if incomplete_pieces.len() > 0 {
            let random_incomplete_piece_index = random_num_generator.gen_range(0, incomplete_pieces.len());
            let random_incomplete_piece = incomplete_pieces[random_incomplete_piece_index];
            Some(random_incomplete_piece)
        } else {
            None
        }
    }
}

struct Piece {
    index:        u32,
    piece_length: u32,
    hash:         Sha1,
    blocks:       Vec<Block>,
    is_complete:  bool,
}

impl Piece {
    fn new(index: u32, length: u32, piece_length: u32, hash: Sha1) -> Piece {
        // create blocks
        let mut blocks = vec![];
        let num_blocks = (length as f64 / BLOCK_SIZE as f64).ceil() as u32;
        for i in 0..num_blocks {
            let len = if i < (num_blocks - 1) {
                BLOCK_SIZE
            } else {
                length - (BLOCK_SIZE * (num_blocks - 1))
            };
            blocks.push(Block::new(i, len));
        }

        Piece {
            index:        index,
            piece_length: piece_length,
            hash:         hash,
            blocks:       blocks,
            is_complete:  false
        }
    }

    fn store(&mut self, file: &mut File, block_index: u32, data: Vec<u8>) -> Result<(), Error> {
        {
            // store data in the appropriate block
            let block = &mut self.blocks[block_index as usize];
            block.data = Some(data);
        }

        if self.have_all_blocks() {
            // concatenate data from blocks together
            let mut data = vec![];
            for block in self.blocks.iter() {
                data.extend(block.data.clone().unwrap());
            }

            // validate that piece data matches SHA1 hash
            if self.hash == calculate_sha1(&data) {
                println!("Piece {} is complete and correct, writing to the file.", self.index);
                let offset = self.index as u64 * self.piece_length as u64;
                try!(file.seek(io::SeekFrom::Start(offset)));
                try!(file.write_all(&data));
                self.clear_block_data();
                self.is_complete = true;
            } else {
                println!("Piece is corrupt, deleting downloaded piece data!");
                self.clear_block_data();
            }
        }
        Ok(())
    }

    fn get_random_incomplete_to_request(&self) -> Option<&Block> {
        let mut random_num_generator = rand::thread_rng();

        if self.is_complete {
            return None
        }

        let empty_blocks: Vec<&Block> = self.blocks.iter().filter(|x| x.data.is_none()).collect();
        if empty_blocks.len() > 0 {
            let random_empty_block_index = random_num_generator.gen_range(0, empty_blocks.len());
            let random_empty_block = empty_blocks[random_empty_block_index];
            Some(random_empty_block)
        } else {
            None
        }
    }

    fn have_all_blocks(&self) -> bool {
        for block in self.blocks.iter() {
            if block.data.is_none() {
                return false
            }
        }
        true
    }

    fn clear_block_data(&mut self) {
        for block in self.blocks.iter_mut() {
            block.data = None;
        }
    }
}

struct Block {
    index:  u32,
    length: u32,
    data:   Option<Vec<u8>>,
}

impl Block {
    fn new(index: u32, length: u32) -> Block {
        Block {
            index:  index,
            length: length,
            data:   None,
        }
    }
}

#[derive(Debug)]
pub enum Error {
    IoError(io::Error),
}

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<|endoftext|>


<commit_before>extern crate handlebars;
extern crate itertools;
extern crate serde;
#[macro_use] extern crate serde_derive;
extern crate serde_json;

use handlebars::Handlebars;
use itertools::Itertools;
use serde_json::Value;
use std::collections::BTreeMap;
use std::ascii::AsciiExt;

trait Codegen {
    fn codegen(&self, handlebars: &Handlebars) -> String;
}

#[derive(Debug, Deserialize)]
pub struct AMQProtocolDefinition {
    pub name:          String,
    #[serde(rename="major-version")]
    pub major_version: u8,
    #[serde(rename="minor-version")]
    pub minor_version: u8,
    pub revision:      u8,
    pub port:          u32,
    pub copyright:     Vec<String>,
    pub domains:       Vec<AMQPDomain>,
    pub constants:     Vec<AMQPConstant>,
    pub classes:       Vec<AMQPClass>,
}

impl AMQProtocolDefinition {
    pub fn load() -> AMQProtocolDefinition {
        let specs = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/specs/amqp-rabbitmq-0.9.1.json"));

        serde_json::from_str::<AMQProtocolDefinition>(specs).expect("Failed to parse AMQP specs file")
    }

    pub fn codegen(&self, templates: &AMQPTemplates) -> String {
        let handlebars = register_templates(templates);
        let mut data   = BTreeMap::new();

        data.insert("name".to_string(),          self.name.clone());
        data.insert("major_version".to_string(), format!("{}", self.major_version));
        data.insert("minor_version".to_string(), format!("{}", self.minor_version));
        data.insert("revision".to_string(),      format!("{}", self.revision));
        data.insert("port".to_string(),          format!("{}", self.port));
        data.insert("copyright".to_string(),     self.copyright.iter().join(""));
        data.insert("domains".to_string(),       self.domains.iter().map(|domain| domain.codegen(&handlebars)).join("\n"));
        data.insert("constants".to_string(),     self.constants.iter().map(|constant| constant.codegen(&handlebars)).join("\n"));
        data.insert("classes".to_string(),       self.classes.iter().map(|klass| klass.codegen(&handlebars)).join("\n"));

        handlebars.render("main", &data).expect("Failed to render main template")
    }
}

fn register_templates(templates: &AMQPTemplates) -> Handlebars {
    let mut handlebars = Handlebars::new();

    handlebars.register_escape_fn(handlebars::no_escape);

    handlebars.register_template_string("main",     &templates.main).expect("Failed to register main template");
    handlebars.register_template_string("domain",   &templates.domain).expect("Failed to register domain template");
    handlebars.register_template_string("constant", &templates.constant).expect("Failed to register constant template");
    handlebars.register_template_string("class",    &templates.klass).expect("Failed to register class template");
    handlebars.register_template_string("method",   &templates.method).expect("Failed to register method template");
    handlebars.register_template_string("argument", &templates.argument).expect("Failed to register argument template");
    handlebars.register_template_string("property", &templates.property).expect("Failed to register property template");

    handlebars
}

pub struct AMQPTemplates {
    pub main:     String,
    pub domain:   String,
    pub constant: String,
    pub klass:    String,
    pub method:   String,
    pub argument: String,
    pub property: String,
}

#[derive(Debug, Deserialize)]
pub enum AMQPType {
    #[serde(rename="bit")]
    Bit,
    #[serde(rename="octet")]
    Octet,
    #[serde(rename="short")]
    Short,
    #[serde(rename="long")]
    Long,
    #[serde(rename="longlong")]
    LongLong,
    #[serde(rename="shortstr")]
    ShortStr,
    #[serde(rename="longstr")]
    LongStr,
    #[serde(rename="table")]
    Table,
    #[serde(rename="timestamp")]
    Timestamp,
}

impl AMQPType {
    fn to_string(&self) -> String {
        match *self {
            AMQPType::Bit       => "bit",
            AMQPType::Octet     => "octet",
            AMQPType::Short     => "short",
            AMQPType::Long      => "long",
            AMQPType::LongLong  => "longlong",
            AMQPType::ShortStr  => "shortstr",
            AMQPType::LongStr   => "longstr",
            AMQPType::Table     => "table",
            AMQPType::Timestamp => "timestamp",
        }.to_string()
    }

    fn to_rust_type(&self) -> String {
        match *self {
            AMQPType::Bit       => "bool",
            AMQPType::Octet     => "u8",
            AMQPType::Short     => "u16",
            AMQPType::Long      => "u32",
            AMQPType::LongLong  => "u64",
            AMQPType::ShortStr  => "String",
            AMQPType::LongStr   => "String",
            AMQPType::Table     => "String", /* FIXME: add a custom type */
            AMQPType::Timestamp => "u64",
        }.to_string()
    }
}

#[derive(Debug, Deserialize)]
pub struct AMQPDomain(String, AMQPType);

impl Codegen for AMQPDomain {
    fn codegen(&self, handlebars: &Handlebars) -> String {
        let mut data = BTreeMap::new();

        data.insert("name".to_string(), self.0.clone());
        data.insert("type".to_string(), self.1.to_string());

        handlebars.render("domain", &data).expect("Failed to render domain template")
    }
}

#[derive(Debug, Deserialize)]
pub struct AMQPConstant {
    pub name:  String,
    pub value: u16,
    #[serde(rename="class")]
    pub klass: Option<String>,
}

impl Codegen for AMQPConstant {
    fn codegen(&self, handlebars: &Handlebars) -> String {
        let mut data = BTreeMap::new();

        data.insert("name".to_string(),  self.name.clone());
        data.insert("value".to_string(), format!("{}", self.value));
        if let Some(ref klass) = self.klass {
            data.insert("class".to_string(), klass.clone());
        }

        handlebars.render("constant", &data).expect("Failed to render constant template")
    }
}

#[derive(Debug, Deserialize)]
pub struct AMQPClass {
    pub id:         u8,
    pub methods:    Vec<AMQPMethod>,
    pub name:       String,
    pub properties: Option<Vec<AMQPProperty>>,
}

impl Codegen for AMQPClass {
    fn codegen(&self, handlebars: &Handlebars) -> String {
        let mut data = BTreeMap::new();

        data.insert("id".to_string(),      format!("{}", self.id));
        data.insert("methods".to_string(), self.methods.iter().map(|method| method.codegen(&handlebars)).join("\n"));
        data.insert("name".to_string(),    self.name.clone());
        if let Some(ref properties) = self.properties {
            data.insert("properties".to_string(), properties.iter().map(|prop| prop.codegen(&handlebars)).join("\n"));
        }

        handlebars.render("class", &data).expect("Failed to render class template")
    }
}

#[derive(Debug, Deserialize)]
pub struct AMQPMethod {
    pub id:          u8,
    pub arguments:   Vec<AMQPArgument>,
    pub name:        String,
    pub synchronous: Option<bool>,
}

impl Codegen for AMQPMethod {
    fn codegen(&self, handlebars: &Handlebars) -> String {
        let mut data = BTreeMap::new();

        data.insert("id".to_string(),          format!("{}", self.id));
        data.insert("arguments".to_string(),   self.arguments.iter().map(|arg| arg.codegen(&handlebars)).join("\n"));
        data.insert("name".to_string(),        self.name.clone());
        data.insert("synchronous".to_string(), format!("{}", self.synchronous.unwrap_or(false)));

        handlebars.render("method", &data).expect("Failed to render method template")
    }
}

#[derive(Debug, Deserialize)]
pub struct AMQPArgument {
    #[serde(rename="type")]
    pub amqp_type:     Option<AMQPType>,
    pub name:          String,
    #[serde(rename="default-value")]
    pub default_value: Option<Value>,
    pub domain:        Option<String>,
}

impl Codegen for AMQPArgument {
    fn codegen(&self, handlebars: &Handlebars) -> String {
        let mut data = BTreeMap::new();

        if let Some(ref amqp_type) = self.amqp_type {
            data.insert("type".to_string(), amqp_type.to_string());
        }
        if let Some(ref default_value) = self.default_value {
            data.insert("default_value".to_string(), default_value.to_string());
        }
        data.insert("name".to_string(),   self.name.clone());
        if let Some(ref domain) = self.domain {
            data.insert("domain".to_string(), domain.clone());
        }

        handlebars.render("argument", &data).expect("Failed to render domain template")
    }
}

#[derive(Debug, Deserialize)]
pub struct AMQPProperty {
    #[serde(rename="type")]
    pub amqp_type: AMQPType,
    pub name:      String,
}

impl AMQPProperty {
    pub fn struct_name(&self) -> String {
        let mut new_word: bool = true;
        self.name.chars().fold("".to_string(), |mut result, ch| {
            if ch == '-' || ch == '_' || ch == ' ' {
                new_word = true;
                result
            } else {
                result.push(if new_word { ch.to_ascii_uppercase() } else { ch.to_ascii_lowercase() });
                new_word = false;
                result
            }
        })
    }
}

impl Codegen for AMQPProperty {
    fn codegen(&self, handlebars: &Handlebars) -> String {
        let mut data = BTreeMap::new();

        data.insert("type".to_string(),        self.amqp_type.to_string());
        data.insert("rust_type".to_string(),   self.amqp_type.to_rust_type());
        data.insert("name".to_string(),        self.name.clone());
        data.insert("struct_name".to_string(), self.struct_name());

        handlebars.render("property", &data).expect("Failed to render domain template")
    }
}

#[cfg(test)]
mod test {
    use super::*;

    fn specs() -> AMQProtocolDefinition {
        AMQProtocolDefinition {
            name:          "AMQP".to_string(),
            major_version: 0,
            minor_version: 9,
            revision:      1,
            port:          5672,
            copyright:     vec!["Copyright 1".to_string(), "Copyright 2".to_string()],
            domains:       vec![AMQPDomain("domain1".to_string(), AMQPType::Octet)],
            constants:     vec![
                AMQPConstant {
                    name:  "constant1".to_string(),
                    value: 128,
                    klass: Some("class1".to_string()),
                }
            ],
            classes:       vec![
                AMQPClass {
                    id:         42,
                    methods:    vec![
                        AMQPMethod {
                            id:          64,
                            arguments:   vec![
                                AMQPArgument {
                                    amqp_type:     Some(AMQPType::LongStr),
                                    name:          "argument1".to_string(),
                                    default_value: Some(From::from("value1")),
                                    domain:        Some("domain1".to_string()),
                                }
                            ],
                            name:        "method1".to_string(),
                            synchronous: Some(true),
                        }
                    ],
                    name:       "class1".to_string(),
                    properties: Some(vec![
                        AMQPProperty {
                            amqp_type: AMQPType::LongStr,
                            name:      "property1".to_string(),
                        }
                    ]),
                }
            ],
        }
    }

    fn templates() -> AMQPTemplates {
        AMQPTemplates {
            main:     r#"
{{name}} - {{major_version}}.{{minor_version}}.{{revision}}
{{copyright}}
port {{port}}
{{domains}}
{{constants}}
{{classes}}
"#.to_string(),
            domain:   "{{name}}: {{type}}".to_string(),
            constant: "{{name}}({{class}}) = {{value}}".to_string(),
            klass:    r#"
{{id}} - {{name}}
{{properties}}
{{methods}}
"#.to_string(),
            method:   r#"
{{id}} - {{name}}
synchronous: {{synchronous}}
{{arguments}}
"#.to_string(),
            argument: "{{name}}({{domain}}): {{type}} = {{default_value}}".to_string(),
            property: "{{name}}: {{type}}".to_string(),
        }
    }

    #[test]
    fn main_template() {
        assert_eq!(specs().codegen(&templates()), r#"
AMQP - 0.9.1
Copyright 1
Copyright 2
port 5672
domain1: octet
constant1(class1) = 128

42 - class1
property1: longstr

64 - method1
synchronous: true
argument1(domain1): longstr = "value1"

"#);
    }
}
<commit_msg>codegen: inject more data to method template too<commit_after>extern crate handlebars;
extern crate itertools;
extern crate serde;
#[macro_use] extern crate serde_derive;
extern crate serde_json;

use handlebars::Handlebars;
use itertools::Itertools;
use serde_json::Value;
use std::collections::BTreeMap;
use std::ascii::AsciiExt;

trait Codegen {
    fn codegen(&self, handlebars: &Handlebars) -> String;
}

#[derive(Debug, Deserialize)]
pub struct AMQProtocolDefinition {
    pub name:          String,
    #[serde(rename="major-version")]
    pub major_version: u8,
    #[serde(rename="minor-version")]
    pub minor_version: u8,
    pub revision:      u8,
    pub port:          u32,
    pub copyright:     Vec<String>,
    pub domains:       Vec<AMQPDomain>,
    pub constants:     Vec<AMQPConstant>,
    pub classes:       Vec<AMQPClass>,
}

impl AMQProtocolDefinition {
    pub fn load() -> AMQProtocolDefinition {
        let specs = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/specs/amqp-rabbitmq-0.9.1.json"));

        serde_json::from_str::<AMQProtocolDefinition>(specs).expect("Failed to parse AMQP specs file")
    }

    pub fn codegen(&self, templates: &AMQPTemplates) -> String {
        let handlebars = register_templates(templates);
        let mut data   = BTreeMap::new();

        data.insert("name".to_string(),          self.name.clone());
        data.insert("major_version".to_string(), format!("{}", self.major_version));
        data.insert("minor_version".to_string(), format!("{}", self.minor_version));
        data.insert("revision".to_string(),      format!("{}", self.revision));
        data.insert("port".to_string(),          format!("{}", self.port));
        data.insert("copyright".to_string(),     self.copyright.iter().join(""));
        data.insert("domains".to_string(),       self.domains.iter().map(|domain| domain.codegen(&handlebars)).join("\n"));
        data.insert("constants".to_string(),     self.constants.iter().map(|constant| constant.codegen(&handlebars)).join("\n"));
        data.insert("classes".to_string(),       self.classes.iter().map(|klass| klass.codegen(&handlebars)).join("\n"));

        handlebars.render("main", &data).expect("Failed to render main template")
    }
}

fn register_templates(templates: &AMQPTemplates) -> Handlebars {
    let mut handlebars = Handlebars::new();

    handlebars.register_escape_fn(handlebars::no_escape);

    handlebars.register_template_string("main",     &templates.main).expect("Failed to register main template");
    handlebars.register_template_string("domain",   &templates.domain).expect("Failed to register domain template");
    handlebars.register_template_string("constant", &templates.constant).expect("Failed to register constant template");
    handlebars.register_template_string("class",    &templates.klass).expect("Failed to register class template");
    handlebars.register_template_string("method",   &templates.method).expect("Failed to register method template");
    handlebars.register_template_string("argument", &templates.argument).expect("Failed to register argument template");
    handlebars.register_template_string("property", &templates.property).expect("Failed to register property template");

    handlebars
}

pub struct AMQPTemplates {
    pub main:     String,
    pub domain:   String,
    pub constant: String,
    pub klass:    String,
    pub method:   String,
    pub argument: String,
    pub property: String,
}

#[derive(Debug, Deserialize)]
pub enum AMQPType {
    #[serde(rename="bit")]
    Bit,
    #[serde(rename="octet")]
    Octet,
    #[serde(rename="short")]
    Short,
    #[serde(rename="long")]
    Long,
    #[serde(rename="longlong")]
    LongLong,
    #[serde(rename="shortstr")]
    ShortStr,
    #[serde(rename="longstr")]
    LongStr,
    #[serde(rename="table")]
    Table,
    #[serde(rename="timestamp")]
    Timestamp,
}

impl AMQPType {
    fn to_string(&self) -> String {
        match *self {
            AMQPType::Bit       => "bit",
            AMQPType::Octet     => "octet",
            AMQPType::Short     => "short",
            AMQPType::Long      => "long",
            AMQPType::LongLong  => "longlong",
            AMQPType::ShortStr  => "shortstr",
            AMQPType::LongStr   => "longstr",
            AMQPType::Table     => "table",
            AMQPType::Timestamp => "timestamp",
        }.to_string()
    }

    fn to_rust_type(&self) -> String {
        match *self {
            AMQPType::Bit       => "bool",
            AMQPType::Octet     => "u8",
            AMQPType::Short     => "u16",
            AMQPType::Long      => "u32",
            AMQPType::LongLong  => "u64",
            AMQPType::ShortStr  => "String",
            AMQPType::LongStr   => "String",
            AMQPType::Table     => "String", /* FIXME: add a custom type */
            AMQPType::Timestamp => "u64",
        }.to_string()
    }
}

pub fn camel_name(name: &str) -> String {
    let mut new_word: bool = true;
    name.chars().fold("".to_string(), |mut result, ch| {
        if ch == '-' || ch == '_' || ch == ' ' {
            new_word = true;
            result
        } else {
            result.push(if new_word { ch.to_ascii_uppercase() } else { ch.to_ascii_lowercase() });
            new_word = false;
            result
        }
    })
}

#[derive(Debug, Deserialize)]
pub struct AMQPDomain(String, AMQPType);

impl Codegen for AMQPDomain {
    fn codegen(&self, handlebars: &Handlebars) -> String {
        let mut data = BTreeMap::new();

        data.insert("name".to_string(), self.0.clone());
        data.insert("type".to_string(), self.1.to_string());

        handlebars.render("domain", &data).expect("Failed to render domain template")
    }
}

#[derive(Debug, Deserialize)]
pub struct AMQPConstant {
    pub name:  String,
    pub value: u16,
    #[serde(rename="class")]
    pub klass: Option<String>,
}

impl Codegen for AMQPConstant {
    fn codegen(&self, handlebars: &Handlebars) -> String {
        let mut data = BTreeMap::new();

        data.insert("name".to_string(),  self.name.clone());
        data.insert("value".to_string(), format!("{}", self.value));
        if let Some(ref klass) = self.klass {
            data.insert("class".to_string(), klass.clone());
        }

        handlebars.render("constant", &data).expect("Failed to render constant template")
    }
}

#[derive(Debug, Deserialize)]
pub struct AMQPClass {
    pub id:         u8,
    pub methods:    Vec<AMQPMethod>,
    pub name:       String,
    pub properties: Option<Vec<AMQPProperty>>,
}

impl Codegen for AMQPClass {
    fn codegen(&self, handlebars: &Handlebars) -> String {
        let mut data = BTreeMap::new();

        data.insert("id".to_string(),      format!("{}", self.id));
        data.insert("methods".to_string(), self.methods.iter().map(|method| method.codegen(&handlebars)).join("\n"));
        data.insert("name".to_string(),    self.name.clone());
        if let Some(ref properties) = self.properties {
            data.insert("properties".to_string(), properties.iter().map(|prop| prop.codegen(&handlebars)).join("\n"));
        }

        handlebars.render("class", &data).expect("Failed to render class template")
    }
}

#[derive(Debug, Deserialize)]
pub struct AMQPMethod {
    pub id:          u8,
    pub arguments:   Vec<AMQPArgument>,
    pub name:        String,
    pub synchronous: Option<bool>,
}

impl Codegen for AMQPMethod {
    fn codegen(&self, handlebars: &Handlebars) -> String {
        let mut data = BTreeMap::new();

        data.insert("id".to_string(),          format!("{}", self.id));
        data.insert("arguments".to_string(),   self.arguments.iter().map(|arg| arg.codegen(&handlebars)).join("\n"));
        data.insert("name".to_string(),        self.name.clone());
        data.insert("synchronous".to_string(), format!("{}", self.synchronous.unwrap_or(false)));
        data.insert("camel_name".to_string(),  camel_name(&self.name));

        handlebars.render("method", &data).expect("Failed to render method template")
    }
}

#[derive(Debug, Deserialize)]
pub struct AMQPArgument {
    #[serde(rename="type")]
    pub amqp_type:     Option<AMQPType>,
    pub name:          String,
    #[serde(rename="default-value")]
    pub default_value: Option<Value>,
    pub domain:        Option<String>,
}

impl Codegen for AMQPArgument {
    fn codegen(&self, handlebars: &Handlebars) -> String {
        let mut data = BTreeMap::new();

        if let Some(ref amqp_type) = self.amqp_type {
            data.insert("type".to_string(), amqp_type.to_string());
        }
        if let Some(ref default_value) = self.default_value {
            data.insert("default_value".to_string(), default_value.to_string());
        }
        data.insert("name".to_string(),   self.name.clone());
        if let Some(ref domain) = self.domain {
            data.insert("domain".to_string(), domain.clone());
        }

        handlebars.render("argument", &data).expect("Failed to render domain template")
    }
}

#[derive(Debug, Deserialize)]
pub struct AMQPProperty {
    #[serde(rename="type")]
    pub amqp_type: AMQPType,
    pub name:      String,
}

impl Codegen for AMQPProperty {
    fn codegen(&self, handlebars: &Handlebars) -> String {
        let mut data = BTreeMap::new();

        data.insert("type".to_string(),       self.amqp_type.to_string());
        data.insert("rust_type".to_string(),  self.amqp_type.to_rust_type());
        data.insert("name".to_string(),       self.name.clone());
        data.insert("camel_name".to_string(), camel_name(&self.name));

        handlebars.render("property", &data).expect("Failed to render domain template")
    }
}

#[cfg(test)]
mod test {
    use super::*;

    fn specs() -> AMQProtocolDefinition {
        AMQProtocolDefinition {
            name:          "AMQP".to_string(),
            major_version: 0,
            minor_version: 9,
            revision:      1,
            port:          5672,
            copyright:     vec!["Copyright 1".to_string(), "Copyright 2".to_string()],
            domains:       vec![AMQPDomain("domain1".to_string(), AMQPType::Octet)],
            constants:     vec![
                AMQPConstant {
                    name:  "constant1".to_string(),
                    value: 128,
                    klass: Some("class1".to_string()),
                }
            ],
            classes:       vec![
                AMQPClass {
                    id:         42,
                    methods:    vec![
                        AMQPMethod {
                            id:          64,
                            arguments:   vec![
                                AMQPArgument {
                                    amqp_type:     Some(AMQPType::LongStr),
                                    name:          "argument1".to_string(),
                                    default_value: Some(From::from("value1")),
                                    domain:        Some("domain1".to_string()),
                                }
                            ],
                            name:        "method1".to_string(),
                            synchronous: Some(true),
                        }
                    ],
                    name:       "class1".to_string(),
                    properties: Some(vec![
                        AMQPProperty {
                            amqp_type: AMQPType::LongStr,
                            name:      "property1".to_string(),
                        }
                    ]),
                }
            ],
        }
    }

    fn templates() -> AMQPTemplates {
        AMQPTemplates {
            main:     r#"
{{name}} - {{major_version}}.{{minor_version}}.{{revision}}
{{copyright}}
port {{port}}
{{domains}}
{{constants}}
{{classes}}
"#.to_string(),
            domain:   "{{name}}: {{type}}".to_string(),
            constant: "{{name}}({{class}}) = {{value}}".to_string(),
            klass:    r#"
{{id}} - {{name}}
{{properties}}
{{methods}}
"#.to_string(),
            method:   r#"
{{id}} - {{name}}
synchronous: {{synchronous}}
{{arguments}}
"#.to_string(),
            argument: "{{name}}({{domain}}): {{type}} = {{default_value}}".to_string(),
            property: "{{name}}: {{type}}".to_string(),
        }
    }

    #[test]
    fn main_template() {
        assert_eq!(specs().codegen(&templates()), r#"
AMQP - 0.9.1
Copyright 1
Copyright 2
port 5672
domain1: octet
constant1(class1) = 128

42 - class1
property1: longstr

64 - method1
synchronous: true
argument1(domain1): longstr = "value1"

"#);
    }
}
<|endoftext|>


<commit_before>//! # A strongly restrictive color type.
use std::fmt;
use palette::pixel::Srgb;
use palette::Rgb;

/// It is guaranteed that `T` can be and will always be only either `u8` or `f32`.
/// It can also be guaranteed that `Display` is not and will never be implemented
/// for `Color<f32>`.
/// It can also be guaranteed that arithmetics are not and will never be implemented
/// for `Color<u8>`.
///
/// ```rust,ignore
/// assert_fail_to_compile!(Color(0.0, 0.0, 0.0).to_string());
/// assert_fail_to_compile!(println!("{}", Color(0.0, 0.0, 0.0)));
/// assert_fail_to_compile!(Color(0, 0, 0) + Color(1, 1, 1));
/// ```
#[derive(Debug, PartialEq, Clone, Copy)]
pub struct Color<T>(pub T, pub T, pub T);

/// `Color888` can only do IO (convert to and from strings), and some simple equality tests.
/// Color is sRGB encoded.
pub type Color888 = Color<u8>;
/// `Colorf32` can only do arithmetics and convert to/from `Color888`.
/// Color is Linear RGB encoded.
pub type Colorf32 = Color<f32>;

impl fmt::Display for Color888 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let Color(r, g, b) = *self;
        write!(f, "{:02x}{:02x}{:02x}", r, g, b)
    }
}

impl Eq for Color888 {}

impl From<Color888> for Colorf32 {
    /// Convert `Color888` (sRGB) to `Colorf32` (Linear RGB).
    fn from(color: Color888) -> Self {
        let Color(r, g, b) = color;
        let rgb: Rgb<f32> = Srgb::new_u8(r, g, b).into();
        let Rgb { red: r, green: g, blue: b } = rgb;
        Color(r, g, b)
    }
}

impl Colorf32 {
    /// Convert `Colorf32` (Linear RGB) to `Color888` (sRGB).
    pub fn clamp_to_888(self) -> Color888 {
        /// This is used to "reduce" conversion error, so that
        /// `lossless_conversion` test can pass.
        const DELTA: f32 = 1e-6;

        let Color(r, g, b) = self;
        let srgb: Srgb = Rgb::new(r + DELTA, g + DELTA, b + DELTA).into();
        let Srgb { red: r, green: g, blue: b, .. } = srgb;
        let Color(r, g, b) = Color(r * 255.0, g * 255.0, b * 255.0).clamp_with(0.0, 255.0);
        Color(r as u8, g as u8, b as u8)
    }

    /// Clamp to [`min`, `max`].
    pub fn clamp_with(self, min: f32, max: f32) -> Self {
        let Color(r, g, b) = self;
        let r = r.max(min).min(max);
        let g = g.max(min).min(max);
        let b = b.max(min).min(max);
        Color(r, g, b)
    }

    /// Clamp to [0, 1].
    pub fn clamp(self) -> Self {
        self.clamp_with(0.0, 1.0)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// When this test fails, you may want to increase the `DELTA` constant
    /// in `Colorf32::clamp_to_888()`.
    #[test]
    fn lossless_conversion() {
        for i in 0..256 {
            let a = i as u8;
            let color = Color(a, a, a);
            let color2 = Colorf32::from(color).clamp_to_888();
            assert_eq!(color, color2);
        }
    }

    #[test]
    fn clampped_down_conversion() {
        let max = Color(255, 255, 255);
        for i in 256..512 {
            let a = i as f32 / 255_f32;
            let color = Color(a, a, a).clamp_to_888();
            assert_eq!(color, max);
        }
    }

    #[test]
    fn clampped_up_conversion() {
        let min = Color(0, 0, 0);
        for i in -256..0 {
            let a = i as f32 / 255_f32;
            let color = Color(a, a, a).clamp_to_888();
            assert_eq!(color, min);
        }
    }

    #[test]
    fn display_color888() {
        assert_eq!(Color(0, 0, 0).to_string(), "000000");
        assert_eq!(Color(255, 255, 255).to_string(), "ffffff");
        assert_eq!(Color(1, 254, 10).to_string(), "01fe0a");
    }
}

<commit_msg>color: implement mixing; clean up<commit_after>//! # A strongly restrictive color type.
use std::fmt;
use palette::pixel::Srgb;
use palette::{Rgb, Lch, Mix};

/// It is guaranteed that `T` can be and will always be only either `u8` or `f32`.
/// It can also be guaranteed that `Display` is not and will never be implemented
/// for `Color<f32>`.
/// It can also be guaranteed that arithmetics are not and will never be implemented
/// for `Color<u8>`.
///
/// ```rust,ignore
/// assert_fail_to_compile!(Color(0.0, 0.0, 0.0).to_string());
/// assert_fail_to_compile!(println!("{}", Color(0.0, 0.0, 0.0)));
/// assert_fail_to_compile!(Color(0, 0, 0) + Color(1, 1, 1));
/// ```
#[derive(Debug, PartialEq, Clone, Copy)]
pub struct Color<T>(pub T, pub T, pub T);

/// `Color888` can only do IO (convert to and from strings), and some simple equality tests.
/// Color is sRGB encoded.
pub type Color888 = Color<u8>;
/// `Colorf32` can only do arithmetics and convert to/from `Color888`.
/// Color is Linear RGB encoded.
pub type Colorf32 = Color<f32>;

impl fmt::Display for Color888 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let Color(r, g, b) = *self;
        write!(f, "{:02x}{:02x}{:02x}", r, g, b)
    }
}

impl Eq for Color888 {}

impl From<Color888> for Colorf32 {
    /// Convert `Color888` (sRGB) to `Colorf32` (Linear RGB).
    fn from(color: Color888) -> Self {
        let Color(r, g, b) = color;
        let rgb = Rgb::<f32>::from(Srgb::new_u8(r, g, b));
        let Rgb { red: r, green: g, blue: b } = rgb;
        Color(r, g, b)
    }
}

impl Colorf32 {
    /// Convert `Colorf32` (Linear RGB) to `Color888` (sRGB).
    pub fn clamp_to_888(self) -> Color888 {
        /// This is used to "reduce" conversion error, so that
        /// `lossless_conversion` test can pass.
        const DELTA: f32 = 1e-6;

        let Color(r, g, b) = self;
        let srgb = Srgb::from(Rgb::new(r + DELTA, g + DELTA, b + DELTA));
        let Srgb { red: r, green: g, blue: b, .. } = srgb;
        let Color(r, g, b) = Color(r * 255.0, g * 255.0, b * 255.0).clamp_with(0.0, 255.0);
        Color(r as u8, g as u8, b as u8)
    }

    /// Clamp to [`min`, `max`].
    pub fn clamp_with(self, min: f32, max: f32) -> Self {
        let Color(r, g, b) = self;
        let r = r.max(min).min(max);
        let g = g.max(min).min(max);
        let b = b.max(min).min(max);
        Color(r, g, b)
    }

    /// Clamp to [0, 1].
    pub fn clamp(self) -> Self {
        self.clamp_with(0.0, 1.0)
    }

    /// Linear interpolate between colors.
    /// This is done in the CIE LCHab color space.
    pub fn mix(self, other: Self, factor: f32) -> Self {
        let Color(r1, g1, b1) = self;
        let Color(r2, g2, b2) = other;
        let lch1 = Lch::from(Rgb::new(r1, g1, b1));
        let lch2 = Lch::from(Rgb::new(r2, g2, b2));
        let lch = lch1.mix(&lch2, factor);
        let Rgb { red: r, green: g, blue: b } = Rgb::from(lch);
        Color(r, g, b)
    }

    /// Linear interpolate between colors.
    /// This is done in the Linear RGB color space.
    pub fn mix_fast(self, other: Self, factor: f32) -> Self {
        let Color(r1, g1, b1) = self;
        let Color(r2, g2, b2) = other;
        let r = (1.0 - factor) * r1 + factor * r2;
        let g = (1.0 - factor) * g1 + factor * g2;
        let b = (1.0 - factor) * b1 + factor * b2;
        Color(r, g, b)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// When this test fails, you may want to increase the `DELTA` constant
    /// in `Colorf32::clamp_to_888()`.
    #[test]
    fn lossless_conversion() {
        for i in 0..256 {
            let a = i as u8;
            let color = Color(a, a, a);
            let color2 = Colorf32::from(color).clamp_to_888();
            assert_eq!(color, color2);
        }
    }

    #[test]
    fn clampped_down_conversion() {
        let max = Color(255, 255, 255);
        for i in 256..512 {
            let a = i as f32 / 255_f32;
            let color = Color(a, a, a).clamp_to_888();
            assert_eq!(color, max);
        }
    }

    #[test]
    fn clampped_up_conversion() {
        let min = Color(0, 0, 0);
        for i in -256..0 {
            let a = i as f32 / 255_f32;
            let color = Color(a, a, a).clamp_to_888();
            assert_eq!(color, min);
        }
    }

    #[test]
    fn display_color888() {
        assert_eq!(Color(0, 0, 0).to_string(), "000000");
        assert_eq!(Color(255, 255, 255).to_string(), "ffffff");
        assert_eq!(Color(1, 254, 10).to_string(), "01fe0a");
    }

    // TODO: automate the checking of this test.
    #[test]
    fn mixing() {
        let c1 = Color(1.0, 0.0, 0.0);
        let c2 = Color(0.0, 0.0, 1.0);
        for i in 0..256 {
            let a = (i as f32) / 255.0;
            let mix1 = c1.mix(c2, a).clamp_to_888();
            let mix2 = c1.mix_fast(c2, a).clamp_to_888();
            println!("{} {}", mix1, mix2);
        }
    }
}

<|endoftext|>


<commit_before>#[macro_use]
mod events;
pub mod data;
pub mod gfx;

use sdl2::render::Renderer;

struct_events! {
  keyboard: {
    key_escape: Escape,
    key_up: Up,
    key_down: Down,
    key_left: Left,
    key_right: Right,
    key_space: Space
  },
  else: {
    quit: Quit { .. }
  }
}

pub struct Game<'window> {
  pub events: Events,
  pub renderer: Renderer<'window>,
}

impl<'window> Game<'window> {
  fn new(events: Events, renderer: Renderer<'window>) -> Game<'window> {
    Game {
      events: events,
      renderer: renderer,
    }
  }

  pub fn output_size(&self) -> (f64, f64) {
    let (w, h) = self.renderer.output_size().unwrap();
    (w as f64, h as f64)
  }
}

fn cartesian_to_isometric(pointX: f64, pointY: f64) -> (f64, f64) {
  let x = pointX - pointY;
  let y = (pointX + pointY) / 2.0;
  (x,y)
}

pub enum ViewAction {
  None,
  Quit,
  ChangeView(Box<View>),
}

pub trait View {
    fn render(&mut self, context: &mut Game, elapsed: f64) -> ViewAction;
  }

  pub fn spawn<F>(title: &str, init: F)
  where F: Fn(&mut Game) -> Box<View> {
    let sdl_context = ::sdl2::init().unwrap();
    let video = sdl_context.video().unwrap();
    let mut timer = sdl_context.timer().unwrap();
    let _image_context = ::sdl2_image::init(::sdl2_image::INIT_PNG).unwrap();

    let window = video.window(title, 1280, 720)
    .position_centered().opengl().resizable()
    .build().unwrap();

    let mut context = Game::new(
      Events::new(sdl_context.event_pump().unwrap()),
      window.renderer()
      .accelerated()
      .build().unwrap());

    let mut current_view = init(&mut context);

    let interval = 1_000 / 60;
    let mut before = timer.ticks();
    let mut last_second = timer.ticks();
    let mut fps = 0u16;

    loop {
      let now = timer.ticks();
      let dt = now - before;
      let elapsed = dt as f64 / 1_000.0;

      if dt < interval {
        timer.delay(interval - dt);
        continue;
      }

      before = now;
      fps += 1;

      if now - last_second > 1_000 {
        println!("FPS: {}", fps);
        last_second = now;
        fps = 0;
      }

      context.events.pump(&mut context.renderer);

      match current_view.render(&mut context, elapsed) {
        ViewAction::None =>
        context.renderer.present(),

        ViewAction::Quit =>
        break,

        ViewAction::ChangeView(new_view) =>
        current_view = new_view,
      }
    }
  }
<commit_msg>Fix key bindings<commit_after>#[macro_use]
mod events;
pub mod data;
pub mod gfx;

use sdl2::render::Renderer;

struct_events! {
  keyboard: {
    key_escape: Escape,
    key_up: W,
    key_down: S,
    key_left: A,
    key_right: D
  },
  else: {
    quit: Quit { .. }
  }
}

pub struct Game<'window> {
  pub events: Events,
  pub renderer: Renderer<'window>,
}

impl<'window> Game<'window> {
  fn new(events: Events, renderer: Renderer<'window>) -> Game<'window> {
    Game {
      events: events,
      renderer: renderer,
    }
  }

  pub fn output_size(&self) -> (f64, f64) {
    let (w, h) = self.renderer.output_size().unwrap();
    (w as f64, h as f64)
  }
}

fn cartesian_to_isometric(pointX: f64, pointY: f64) -> (f64, f64) {
  let x = pointX - pointY;
  let y = (pointX + pointY) / 2.0;
  (x,y)
}

pub enum ViewAction {
  None,
  Quit,
  ChangeView(Box<View>),
}

pub trait View {
    fn render(&mut self, context: &mut Game, elapsed: f64) -> ViewAction;
  }

  pub fn spawn<F>(title: &str, init: F)
  where F: Fn(&mut Game) -> Box<View> {
    let sdl_context = ::sdl2::init().unwrap();
    let video = sdl_context.video().unwrap();
    let mut timer = sdl_context.timer().unwrap();
    let _image_context = ::sdl2_image::init(::sdl2_image::INIT_PNG).unwrap();

    let window = video.window(title, 1280, 720)
    .position_centered().opengl().resizable()
    .build().unwrap();

    let mut context = Game::new(
      Events::new(sdl_context.event_pump().unwrap()),
      window.renderer()
      .accelerated()
      .build().unwrap());

    let mut current_view = init(&mut context);

    let interval = 1_000 / 60;
    let mut before = timer.ticks();
    let mut last_second = timer.ticks();
    let mut fps = 0u16;

    loop {
      let now = timer.ticks();
      let dt = now - before;
      let elapsed = dt as f64 / 1_000.0;

      if dt < interval {
        timer.delay(interval - dt);
        continue;
      }

      before = now;
      fps += 1;

      if now - last_second > 1_000 {
        println!("FPS: {}", fps);
        last_second = now;
        fps = 0;
      }

      context.events.pump(&mut context.renderer);

      match current_view.render(&mut context, elapsed) {
        ViewAction::None =>
        context.renderer.present(),

        ViewAction::Quit =>
        break,

        ViewAction::ChangeView(new_view) =>
        current_view = new_view,
      }
    }
  }
<|endoftext|>