<commit_before>//! Pragma helpers

use std::ops::Deref;

use crate::error::Error;
use crate::ffi;
use crate::types::{ToSql, ToSqlOutput, ValueRef};
use crate::{Connection, DatabaseName, Result, Row, NO_PARAMS};

pub struct Sql {
    buf: String,
}

impl Sql {
    pub fn new() -> Sql {
        Sql { buf: String::new() }
    }

    pub fn push_pragma(
        &mut self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
    ) -> Result<()> {
        self.push_keyword("PRAGMA")?;
        self.push_space();
        if let Some(schema_name) = schema_name {
            self.push_schema_name(schema_name);
            self.push_dot();
        }
        self.push_keyword(pragma_name)
    }

    pub fn push_keyword(&mut self, keyword: &str) -> Result<()> {
        if !keyword.is_empty() && is_identifier(keyword) {
            self.buf.push_str(keyword);
            Ok(())
        } else {
            Err(Error::SqliteFailure(
                ffi::Error::new(ffi::SQLITE_MISUSE),
                Some(format!("Invalid keyword \"{}\"", keyword)),
            ))
        }
    }

    pub fn push_schema_name(&mut self, schema_name: DatabaseName<'_>) {
        match schema_name {
            DatabaseName::Main => self.buf.push_str("main"),
            DatabaseName::Temp => self.buf.push_str("temp"),
            DatabaseName::Attached(s) => self.push_identifier(s),
        };
    }

    pub fn push_identifier(&mut self, s: &str) {
        if is_identifier(s) {
            self.buf.push_str(s);
        } else {
            self.wrap_and_escape(s, '"');
        }
    }

    pub fn push_value(&mut self, value: &dyn ToSql) -> Result<()> {
        let value = value.to_sql()?;
        let value = match value {
            ToSqlOutput::Borrowed(v) => v,
            ToSqlOutput::Owned(ref v) => ValueRef::from(v),
            #[cfg(feature = "blob")]
            ToSqlOutput::ZeroBlob(_) => {
                return Err(Error::SqliteFailure(
                    ffi::Error::new(ffi::SQLITE_MISUSE),
                    Some(format!("Unsupported value \"{:?}\"", value)),
                ));
            }
            #[cfg(feature = "array")]
            ToSqlOutput::Array(_) => {
                return Err(Error::SqliteFailure(
                    ffi::Error::new(ffi::SQLITE_MISUSE),
                    Some(format!("Unsupported value \"{:?}\"", value)),
                ));
            }
        };
        match value {
            ValueRef::Integer(i) => {
                self.push_int(i);
            }
            ValueRef::Real(r) => {
                self.push_real(r);
            }
            ValueRef::Text(s) => {
                self.push_string_literal(s);
            }
            _ => {
                return Err(Error::SqliteFailure(
                    ffi::Error::new(ffi::SQLITE_MISUSE),
                    Some(format!("Unsupported value \"{:?}\"", value)),
                ));
            }
        };
        Ok(())
    }

    pub fn push_string_literal(&mut self, s: &str) {
        self.wrap_and_escape(s, '\'');
    }

    pub fn push_int(&mut self, i: i64) {
        self.buf.push_str(&i.to_string());
    }

    pub fn push_real(&mut self, f: f64) {
        self.buf.push_str(&f.to_string());
    }

    pub fn push_space(&mut self) {
        self.buf.push(' ');
    }

    pub fn push_dot(&mut self) {
        self.buf.push('.');
    }

    pub fn push_equal_sign(&mut self) {
        self.buf.push('=');
    }

    pub fn open_brace(&mut self) {
        self.buf.push('(');
    }

    pub fn close_brace(&mut self) {
        self.buf.push(')');
    }

    pub fn as_str(&self) -> &str {
        &self.buf
    }

    fn wrap_and_escape(&mut self, s: &str, quote: char) {
        self.buf.push(quote);
        let chars = s.chars();
        for ch in chars {
            // escape `quote` by doubling it
            if ch == quote {
                self.buf.push(ch);
            }
            self.buf.push(ch)
        }
        self.buf.push(quote);
    }
}

impl Deref for Sql {
    type Target = str;

    fn deref(&self) -> &str {
        self.as_str()
    }
}

impl Connection {
    /// Query the current value of `pragma_name`.
    ///
    /// Some pragmas will return multiple rows/values which cannot be retrieved
    /// with this method.
    ///
    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
    /// `SELECT user_version FROM pragma_user_version;`
    pub fn pragma_query_value<T, F>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        f: F,
    ) -> Result<T>
    where
        F: FnOnce(Row<'_, '_>) -> Result<T>,
    {
        let mut query = Sql::new();
        query.push_pragma(schema_name, pragma_name)?;
        let mut stmt = self.prepare(&query)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        if let Some(result_row) = rows.next() {
            let row = result_row?;
            f(row)
        } else {
            Err(Error::QueryReturnedNoRows)
        }
    }

    /// Query the current rows/values of `pragma_name`.
    ///
    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
    /// `SELECT * FROM pragma_collation_list;`
    pub fn pragma_query<F>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        mut f: F,
    ) -> Result<()>
    where
        F: FnMut(&Row<'_, '_>) -> Result<()>,
    {
        let mut query = Sql::new();
        query.push_pragma(schema_name, pragma_name)?;
        let mut stmt = self.prepare(&query)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        while let Some(result_row) = rows.next() {
            let row = result_row?;
            f(&row)?;
        }
        Ok(())
    }

    /// Query the current value(s) of `pragma_name` associated to
    /// `pragma_value`.
    ///
    /// This method can be used with query-only pragmas which need an argument
    /// (e.g. `table_info('one_tbl')`) or pragmas which returns value(s)
    /// (e.g. `integrity_check`).
    ///
    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
    /// `SELECT * FROM pragma_table_info(?);`
    pub fn pragma<F>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        pragma_value: &dyn ToSql,
        mut f: F,
    ) -> Result<()>
    where
        F: FnMut(&Row<'_, '_>) -> Result<()>,
    {
        let mut sql = Sql::new();
        sql.push_pragma(schema_name, pragma_name)?;
        // The argument is may be either in parentheses
        // or it may be separated from the pragma name by an equal sign.
        // The two syntaxes yield identical results.
        sql.open_brace();
        sql.push_value(pragma_value)?;
        sql.close_brace();
        let mut stmt = self.prepare(&sql)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        while let Some(result_row) = rows.next() {
            let row = result_row?;
            f(&row)?;
        }
        Ok(())
    }

    /// Set a new value to `pragma_name`.
    ///
    /// Some pragmas will return the updated value which cannot be retrieved
    /// with this method.
    pub fn pragma_update(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        pragma_value: &dyn ToSql,
    ) -> Result<()> {
        let mut sql = Sql::new();
        sql.push_pragma(schema_name, pragma_name)?;
        // The argument is may be either in parentheses
        // or it may be separated from the pragma name by an equal sign.
        // The two syntaxes yield identical results.
        sql.push_equal_sign();
        sql.push_value(pragma_value)?;
        self.execute_batch(&sql)
    }

    /// Set a new value to `pragma_name` and return the updated value.
    ///
    /// Only few pragmas automatically return the updated value.
    pub fn pragma_update_and_check<F, T>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        pragma_value: &dyn ToSql,
        f: F,
    ) -> Result<T>
    where
        F: FnOnce(Row<'_, '_>) -> Result<T>,
    {
        let mut sql = Sql::new();
        sql.push_pragma(schema_name, pragma_name)?;
        // The argument is may be either in parentheses
        // or it may be separated from the pragma name by an equal sign.
        // The two syntaxes yield identical results.
        sql.push_equal_sign();
        sql.push_value(pragma_value)?;
        let mut stmt = self.prepare(&sql)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        if let Some(result_row) = rows.next() {
            let row = result_row?;
            f(row)
        } else {
            Err(Error::QueryReturnedNoRows)
        }
    }
}

fn is_identifier(s: &str) -> bool {
    let chars = s.char_indices();
    for (i, ch) in chars {
        if i == 0 {
            if !is_identifier_start(ch) {
                return false;
            }
        } else if !is_identifier_continue(ch) {
            return false;
        }
    }
    true
}

fn is_identifier_start(c: char) -> bool {
    (c >= 'A' && c <= 'Z') || c == '_' || (c >= 'a' && c <= 'z') || c > '\x7F'
}

fn is_identifier_continue(c: char) -> bool {
    c == '$'
        || (c >= '0' && c <= '9')
        || (c >= 'A' && c <= 'Z')
        || c == '_'
        || (c >= 'a' && c <= 'z')
        || c > '\x7F'
}

#[cfg(test)]
mod test {
    use super::Sql;
    use crate::pragma;
    use crate::{Connection, DatabaseName};

    #[test]
    fn pragma_query_value() {
        let db = Connection::open_in_memory().unwrap();
        let user_version: i32 = db
            .pragma_query_value(None, "user_version", |row| row.get_checked(0))
            .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    #[cfg(feature = "bundled")]
    fn pragma_func_query_value() {
        use crate::NO_PARAMS;

        let db = Connection::open_in_memory().unwrap();
        let user_version: i32 = db
            .query_row(
                "SELECT user_version FROM pragma_user_version",
                NO_PARAMS,
                |row| row.get(0),
            )
            .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    fn pragma_query_no_schema() {
        let db = Connection::open_in_memory().unwrap();
        let mut user_version = -1;
        db.pragma_query(None, "user_version", |row| {
            user_version = row.get_checked(0)?;
            Ok(())
        })
        .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    fn pragma_query_with_schema() {
        let db = Connection::open_in_memory().unwrap();
        let mut user_version = -1;
        db.pragma_query(Some(DatabaseName::Main), "user_version", |row| {
            user_version = row.get_checked(0)?;
            Ok(())
        })
        .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    fn pragma() {
        let db = Connection::open_in_memory().unwrap();
        let mut columns = Vec::new();
        db.pragma(None, "table_info", &"sqlite_master", |row| {
            let column: String = row.get_checked(1)?;
            columns.push(column);
            Ok(())
        })
        .unwrap();
        assert_eq!(5, columns.len());
    }

    #[test]
    #[cfg(feature = "bundled")]
    fn pragma_func() {
        let db = Connection::open_in_memory().unwrap();
        let mut table_info = db.prepare("SELECT * FROM pragma_table_info(?)").unwrap();
        let mut columns = Vec::new();
        let mut rows = table_info.query(&["sqlite_master"]).unwrap();

        while let Some(row) = rows.next() {
            let row = row.unwrap();
            let column: String = row.get(1);
            columns.push(column);
        }
        assert_eq!(5, columns.len());
    }

    #[test]
    fn pragma_update() {
        let db = Connection::open_in_memory().unwrap();
        db.pragma_update(None, "user_version", &1).unwrap();
    }

    #[test]
    fn pragma_update_and_check() {
        let db = Connection::open_in_memory().unwrap();
        let journal_mode: String = db
            .pragma_update_and_check(None, "journal_mode", &"OFF", |row| row.get_checked(0))
            .unwrap();
        assert_eq!("off", &journal_mode);
    }

    #[test]
    fn is_identifier() {
        assert!(pragma::is_identifier("full"));
        assert!(pragma::is_identifier("r2d2"));
        assert!(!pragma::is_identifier("sp ce"));
        assert!(!pragma::is_identifier("semi;colon"));
    }

    #[test]
    fn double_quote() {
        let mut sql = Sql::new();
        sql.push_schema_name(DatabaseName::Attached(r#"schema";--"#));
        assert_eq!(r#""schema"";--""#, sql.as_str());
    }

    #[test]
    fn wrap_and_escape() {
        let mut sql = Sql::new();
        sql.push_string_literal("value'; --");
        assert_eq!("'value''; --'", sql.as_str());
    }
}
<commit_msg>Fix typos<commit_after>//! Pragma helpers

use std::ops::Deref;

use crate::error::Error;
use crate::ffi;
use crate::types::{ToSql, ToSqlOutput, ValueRef};
use crate::{Connection, DatabaseName, Result, Row, NO_PARAMS};

pub struct Sql {
    buf: String,
}

impl Sql {
    pub fn new() -> Sql {
        Sql { buf: String::new() }
    }

    pub fn push_pragma(
        &mut self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
    ) -> Result<()> {
        self.push_keyword("PRAGMA")?;
        self.push_space();
        if let Some(schema_name) = schema_name {
            self.push_schema_name(schema_name);
            self.push_dot();
        }
        self.push_keyword(pragma_name)
    }

    pub fn push_keyword(&mut self, keyword: &str) -> Result<()> {
        if !keyword.is_empty() && is_identifier(keyword) {
            self.buf.push_str(keyword);
            Ok(())
        } else {
            Err(Error::SqliteFailure(
                ffi::Error::new(ffi::SQLITE_MISUSE),
                Some(format!("Invalid keyword \"{}\"", keyword)),
            ))
        }
    }

    pub fn push_schema_name(&mut self, schema_name: DatabaseName<'_>) {
        match schema_name {
            DatabaseName::Main => self.buf.push_str("main"),
            DatabaseName::Temp => self.buf.push_str("temp"),
            DatabaseName::Attached(s) => self.push_identifier(s),
        };
    }

    pub fn push_identifier(&mut self, s: &str) {
        if is_identifier(s) {
            self.buf.push_str(s);
        } else {
            self.wrap_and_escape(s, '"');
        }
    }

    pub fn push_value(&mut self, value: &dyn ToSql) -> Result<()> {
        let value = value.to_sql()?;
        let value = match value {
            ToSqlOutput::Borrowed(v) => v,
            ToSqlOutput::Owned(ref v) => ValueRef::from(v),
            #[cfg(feature = "blob")]
            ToSqlOutput::ZeroBlob(_) => {
                return Err(Error::SqliteFailure(
                    ffi::Error::new(ffi::SQLITE_MISUSE),
                    Some(format!("Unsupported value \"{:?}\"", value)),
                ));
            }
            #[cfg(feature = "array")]
            ToSqlOutput::Array(_) => {
                return Err(Error::SqliteFailure(
                    ffi::Error::new(ffi::SQLITE_MISUSE),
                    Some(format!("Unsupported value \"{:?}\"", value)),
                ));
            }
        };
        match value {
            ValueRef::Integer(i) => {
                self.push_int(i);
            }
            ValueRef::Real(r) => {
                self.push_real(r);
            }
            ValueRef::Text(s) => {
                self.push_string_literal(s);
            }
            _ => {
                return Err(Error::SqliteFailure(
                    ffi::Error::new(ffi::SQLITE_MISUSE),
                    Some(format!("Unsupported value \"{:?}\"", value)),
                ));
            }
        };
        Ok(())
    }

    pub fn push_string_literal(&mut self, s: &str) {
        self.wrap_and_escape(s, '\'');
    }

    pub fn push_int(&mut self, i: i64) {
        self.buf.push_str(&i.to_string());
    }

    pub fn push_real(&mut self, f: f64) {
        self.buf.push_str(&f.to_string());
    }

    pub fn push_space(&mut self) {
        self.buf.push(' ');
    }

    pub fn push_dot(&mut self) {
        self.buf.push('.');
    }

    pub fn push_equal_sign(&mut self) {
        self.buf.push('=');
    }

    pub fn open_brace(&mut self) {
        self.buf.push('(');
    }

    pub fn close_brace(&mut self) {
        self.buf.push(')');
    }

    pub fn as_str(&self) -> &str {
        &self.buf
    }

    fn wrap_and_escape(&mut self, s: &str, quote: char) {
        self.buf.push(quote);
        let chars = s.chars();
        for ch in chars {
            // escape `quote` by doubling it
            if ch == quote {
                self.buf.push(ch);
            }
            self.buf.push(ch)
        }
        self.buf.push(quote);
    }
}

impl Deref for Sql {
    type Target = str;

    fn deref(&self) -> &str {
        self.as_str()
    }
}

impl Connection {
    /// Query the current value of `pragma_name`.
    ///
    /// Some pragmas will return multiple rows/values which cannot be retrieved
    /// with this method.
    ///
    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
    /// `SELECT user_version FROM pragma_user_version;`
    pub fn pragma_query_value<T, F>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        f: F,
    ) -> Result<T>
    where
        F: FnOnce(Row<'_, '_>) -> Result<T>,
    {
        let mut query = Sql::new();
        query.push_pragma(schema_name, pragma_name)?;
        let mut stmt = self.prepare(&query)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        if let Some(result_row) = rows.next() {
            let row = result_row?;
            f(row)
        } else {
            Err(Error::QueryReturnedNoRows)
        }
    }

    /// Query the current rows/values of `pragma_name`.
    ///
    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
    /// `SELECT * FROM pragma_collation_list;`
    pub fn pragma_query<F>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        mut f: F,
    ) -> Result<()>
    where
        F: FnMut(&Row<'_, '_>) -> Result<()>,
    {
        let mut query = Sql::new();
        query.push_pragma(schema_name, pragma_name)?;
        let mut stmt = self.prepare(&query)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        while let Some(result_row) = rows.next() {
            let row = result_row?;
            f(&row)?;
        }
        Ok(())
    }

    /// Query the current value(s) of `pragma_name` associated to
    /// `pragma_value`.
    ///
    /// This method can be used with query-only pragmas which need an argument
    /// (e.g. `table_info('one_tbl')`) or pragmas which returns value(s)
    /// (e.g. `integrity_check`).
    ///
    /// Prefer [PRAGMA function](https://sqlite.org/pragma.html#pragfunc) introduced in SQLite 3.20:
    /// `SELECT * FROM pragma_table_info(?);`
    pub fn pragma<F>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        pragma_value: &dyn ToSql,
        mut f: F,
    ) -> Result<()>
    where
        F: FnMut(&Row<'_, '_>) -> Result<()>,
    {
        let mut sql = Sql::new();
        sql.push_pragma(schema_name, pragma_name)?;
        // The argument may be either in parentheses
        // or it may be separated from the pragma name by an equal sign.
        // The two syntaxes yield identical results.
        sql.open_brace();
        sql.push_value(pragma_value)?;
        sql.close_brace();
        let mut stmt = self.prepare(&sql)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        while let Some(result_row) = rows.next() {
            let row = result_row?;
            f(&row)?;
        }
        Ok(())
    }

    /// Set a new value to `pragma_name`.
    ///
    /// Some pragmas will return the updated value which cannot be retrieved
    /// with this method.
    pub fn pragma_update(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        pragma_value: &dyn ToSql,
    ) -> Result<()> {
        let mut sql = Sql::new();
        sql.push_pragma(schema_name, pragma_name)?;
        // The argument may be either in parentheses
        // or it may be separated from the pragma name by an equal sign.
        // The two syntaxes yield identical results.
        sql.push_equal_sign();
        sql.push_value(pragma_value)?;
        self.execute_batch(&sql)
    }

    /// Set a new value to `pragma_name` and return the updated value.
    ///
    /// Only few pragmas automatically return the updated value.
    pub fn pragma_update_and_check<F, T>(
        &self,
        schema_name: Option<DatabaseName<'_>>,
        pragma_name: &str,
        pragma_value: &dyn ToSql,
        f: F,
    ) -> Result<T>
    where
        F: FnOnce(Row<'_, '_>) -> Result<T>,
    {
        let mut sql = Sql::new();
        sql.push_pragma(schema_name, pragma_name)?;
        // The argument may be either in parentheses
        // or it may be separated from the pragma name by an equal sign.
        // The two syntaxes yield identical results.
        sql.push_equal_sign();
        sql.push_value(pragma_value)?;
        let mut stmt = self.prepare(&sql)?;
        let mut rows = stmt.query(NO_PARAMS)?;
        if let Some(result_row) = rows.next() {
            let row = result_row?;
            f(row)
        } else {
            Err(Error::QueryReturnedNoRows)
        }
    }
}

fn is_identifier(s: &str) -> bool {
    let chars = s.char_indices();
    for (i, ch) in chars {
        if i == 0 {
            if !is_identifier_start(ch) {
                return false;
            }
        } else if !is_identifier_continue(ch) {
            return false;
        }
    }
    true
}

fn is_identifier_start(c: char) -> bool {
    (c >= 'A' && c <= 'Z') || c == '_' || (c >= 'a' && c <= 'z') || c > '\x7F'
}

fn is_identifier_continue(c: char) -> bool {
    c == '$'
        || (c >= '0' && c <= '9')
        || (c >= 'A' && c <= 'Z')
        || c == '_'
        || (c >= 'a' && c <= 'z')
        || c > '\x7F'
}

#[cfg(test)]
mod test {
    use super::Sql;
    use crate::pragma;
    use crate::{Connection, DatabaseName};

    #[test]
    fn pragma_query_value() {
        let db = Connection::open_in_memory().unwrap();
        let user_version: i32 = db
            .pragma_query_value(None, "user_version", |row| row.get_checked(0))
            .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    #[cfg(feature = "bundled")]
    fn pragma_func_query_value() {
        use crate::NO_PARAMS;

        let db = Connection::open_in_memory().unwrap();
        let user_version: i32 = db
            .query_row(
                "SELECT user_version FROM pragma_user_version",
                NO_PARAMS,
                |row| row.get(0),
            )
            .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    fn pragma_query_no_schema() {
        let db = Connection::open_in_memory().unwrap();
        let mut user_version = -1;
        db.pragma_query(None, "user_version", |row| {
            user_version = row.get_checked(0)?;
            Ok(())
        })
        .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    fn pragma_query_with_schema() {
        let db = Connection::open_in_memory().unwrap();
        let mut user_version = -1;
        db.pragma_query(Some(DatabaseName::Main), "user_version", |row| {
            user_version = row.get_checked(0)?;
            Ok(())
        })
        .unwrap();
        assert_eq!(0, user_version);
    }

    #[test]
    fn pragma() {
        let db = Connection::open_in_memory().unwrap();
        let mut columns = Vec::new();
        db.pragma(None, "table_info", &"sqlite_master", |row| {
            let column: String = row.get_checked(1)?;
            columns.push(column);
            Ok(())
        })
        .unwrap();
        assert_eq!(5, columns.len());
    }

    #[test]
    #[cfg(feature = "bundled")]
    fn pragma_func() {
        let db = Connection::open_in_memory().unwrap();
        let mut table_info = db.prepare("SELECT * FROM pragma_table_info(?)").unwrap();
        let mut columns = Vec::new();
        let mut rows = table_info.query(&["sqlite_master"]).unwrap();

        while let Some(row) = rows.next() {
            let row = row.unwrap();
            let column: String = row.get(1);
            columns.push(column);
        }
        assert_eq!(5, columns.len());
    }

    #[test]
    fn pragma_update() {
        let db = Connection::open_in_memory().unwrap();
        db.pragma_update(None, "user_version", &1).unwrap();
    }

    #[test]
    fn pragma_update_and_check() {
        let db = Connection::open_in_memory().unwrap();
        let journal_mode: String = db
            .pragma_update_and_check(None, "journal_mode", &"OFF", |row| row.get_checked(0))
            .unwrap();
        assert_eq!("off", &journal_mode);
    }

    #[test]
    fn is_identifier() {
        assert!(pragma::is_identifier("full"));
        assert!(pragma::is_identifier("r2d2"));
        assert!(!pragma::is_identifier("sp ce"));
        assert!(!pragma::is_identifier("semi;colon"));
    }

    #[test]
    fn double_quote() {
        let mut sql = Sql::new();
        sql.push_schema_name(DatabaseName::Attached(r#"schema";--"#));
        assert_eq!(r#""schema"";--""#, sql.as_str());
    }

    #[test]
    fn wrap_and_escape() {
        let mut sql = Sql::new();
        sql.push_string_literal("value'; --");
        assert_eq!("'value''; --'", sql.as_str());
    }
}
<|endoftext|>


<commit_before>use ammonia::Ammonia;
use comrak;
use htmlescape::encode_minimal;

use util::CargoResult;

/// Context for markdown to HTML rendering.
#[allow(missing_debug_implementations)]
pub struct MarkdownRenderer<'a> {
    html_sanitizer: Ammonia<'a>,
}

impl<'a> MarkdownRenderer<'a> {
    /// Creates a new renderer instance.
    pub fn new() -> MarkdownRenderer<'a> {
        let tags = [
            "a",
            "b",
            "blockquote",
            "br",
            "code",
            "dd",
            "del",
            "dl",
            "dt",
            "em",
            "h1",
            "h2",
            "h3",
            "hr",
            "i",
            "img",
            "input",
            "kbd",
            "li",
            "ol",
            "p",
            "pre",
            "s",
            "strike",
            "strong",
            "sub",
            "sup",
            "table",
            "tbody",
            "td",
            "th",
            "thead",
            "tr",
            "ul",
            "hr",
            "span",
        ].iter()
            .cloned()
            .collect();
        let tag_attributes = [
            ("a", ["href", "target"].iter().cloned().collect()),
            ("code", ["class"].iter().cloned().collect()),
            (
                "img",
                ["width", "height", "src", "alt", "align"]
                    .iter()
                    .cloned()
                    .collect(),
            ),
            (
                "input",
                ["checked", "disabled", "type"].iter().cloned().collect(),
            ),
        ].iter()
            .cloned()
            .collect();
        let allowed_classes = [
            (
                "code",
                [
                    "language-bash",
                    "language-clike",
                    "language-glsl",
                    "language-go",
                    "language-ini",
                    "language-javascript",
                    "language-json",
                    "language-markup",
                    "language-protobuf",
                    "language-ruby",
                    "language-rust",
                    "language-scss",
                    "language-sql",
                    "yaml",
                ].iter()
                    .cloned()
                    .collect(),
            ),
        ].iter()
            .cloned()
            .collect();
        let html_sanitizer = Ammonia {
            link_rel: Some("nofollow noopener noreferrer"),
            keep_cleaned_elements: true,
            tags: tags,
            tag_attributes: tag_attributes,
            allowed_classes: allowed_classes,
            ..Ammonia::default()
        };
        MarkdownRenderer {
            html_sanitizer: html_sanitizer,
        }
    }

    /// Renders the given markdown to HTML using the current settings.
    pub fn to_html(&self, text: &str) -> CargoResult<String> {
        let options = comrak::ComrakOptions {
            ext_autolink: true,
            ext_strikethrough: true,
            ext_table: true,
            ext_tagfilter: true,
            ext_tasklist: true,
            ..comrak::ComrakOptions::default()
        };
        let rendered = comrak::markdown_to_html(text, &options);
        Ok(self.html_sanitizer.clean(&rendered))
    }
}

impl<'a> Default for MarkdownRenderer<'a> {
    fn default() -> Self {
        Self::new()
    }
}

/// Renders a markdown text to sanitized HTML.
///
/// The returned text should not contain any harmful HTML tag or attribute (such as iframe,
/// onclick, onmouseover, etc.).
///
/// # Examples
///
/// ```
/// use render::markdown_to_html;
///
/// let text = "[Rust](https://rust-lang.org/) is an awesome *systems programming* language!";
/// let rendered = markdown_to_html(text)?;
/// ```
fn markdown_to_html(text: &str) -> CargoResult<String> {
    let renderer = MarkdownRenderer::new();
    renderer.to_html(text)
}

static MARKDOWN_EXTENSIONS: [&'static str; 7] = [
    ".md",
    ".markdown",
    ".mdown",
    ".mdwn",
    ".mkd",
    ".mkdn",
    ".mkdown",
];

pub fn readme_to_html(text: &str, filename: &str) -> CargoResult<String> {
    for e in &MARKDOWN_EXTENSIONS {
        if filename.to_lowercase().ends_with(e) {
            return markdown_to_html(text);
        }
    }

    Ok(encode_minimal(text).replace("\n", "
\n"))
}

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

    #[test]
    fn empty_text() {
        let text = "";
        let result = markdown_to_html(text).unwrap();
        assert_eq!(result, "");
    }

    #[test]
    fn text_with_script_tag() {
        let text = "foo_readme\n\n";
        let result = markdown_to_html(text).unwrap();
        assert_eq!(
            result,
            "
foo_readme
\n<script>alert(\'Hello World\')</script>\n"
        );
    }

    #[test]
    fn text_with_iframe_tag() {
        let text = "foo_readme\n\nalert('Hello World')";
        let result = markdown_to_html(text).unwrap();
        assert_eq!(
            result,
            "
foo_readme
\n<iframe>alert(\'Hello World\')</iframe>\n"
        );
    }

    #[test]
    fn text_with_unknown_tag() {
        let text = "foo_readme\n\n<unknown>alert('Hello World')</unknown>";
        let result = markdown_to_html(text).unwrap();
        assert_eq!(result, "
foo_readme
\n
alert(\'Hello World\')
\n");
    }

    #[test]
    fn text_with_inline_javascript() {
        let text = r#"foo_readme\n\nCrate page"#;
        let result = markdown_to_html(text).unwrap();
        assert_eq!(
            result,
            "
foo_readme\\n\\nCrate page
\n"
        );
    }

    // See https://github.com/kivikakk/comrak/issues/37. This panic happened
    // in comrak 0.1.8 but was fixed in 0.1.9.
    #[test]
    fn text_with_fancy_single_quotes() {
        let text = r#"wb’"#;
        let result = markdown_to_html(text).unwrap();
        assert_eq!(result, "
wb’
\n");
    }

    #[test]
    fn code_block_with_syntax_highlighting() {
        let code_block = r#"```rust \
                            println!("Hello World"); \
                           ```"#;
        let result = markdown_to_html(code_block).unwrap();
        assert!(result.contains(""));
    }

    #[test]
    fn text_with_forbidden_class_attribute() {
        let text = "
Hello World!
";
        let result = markdown_to_html(text).unwrap();
        assert_eq!(result, "
Hello World!
\n");
    }

    #[test]
    fn readme_to_html_renders_markdown() {
        for f in &["readme.md", "README.MARKDOWN", "whatever.mkd"] {
            assert_eq!(
                readme_to_html("*lobster*", f).unwrap(),
                "
lobster
\n"
            );
        }
    }

    #[test]
    fn readme_to_html_renders_other_things() {
        for f in &["readme", "readem.org", "blah.adoc"] {
            assert_eq!(
                readme_to_html("\n\nis my friend\n", f).unwrap(),
                "<script>lobster</script>
\n
\nis my friend
\n"
            );
        }
    }
}
<commit_msg>Cleanup based on @kureuil's comments<commit_after>use ammonia::Ammonia;
use comrak;
use htmlescape::encode_minimal;

use util::CargoResult;

/// Context for markdown to HTML rendering.
#[allow(missing_debug_implementations)]
struct MarkdownRenderer<'a> {
    html_sanitizer: Ammonia<'a>,
}

impl<'a> MarkdownRenderer<'a> {
    /// Creates a new renderer instance.
    fn new() -> MarkdownRenderer<'a> {
        let tags = [
            "a",
            "b",
            "blockquote",
            "br",
            "code",
            "dd",
            "del",
            "dl",
            "dt",
            "em",
            "h1",
            "h2",
            "h3",
            "hr",
            "i",
            "img",
            "input",
            "kbd",
            "li",
            "ol",
            "p",
            "pre",
            "s",
            "strike",
            "strong",
            "sub",
            "sup",
            "table",
            "tbody",
            "td",
            "th",
            "thead",
            "tr",
            "ul",
            "hr",
            "span",
        ].iter()
            .cloned()
            .collect();
        let tag_attributes = [
            ("a", ["href", "target"].iter().cloned().collect()),
            ("code", ["class"].iter().cloned().collect()),
            (
                "img",
                ["width", "height", "src", "alt", "align"]
                    .iter()
                    .cloned()
                    .collect(),
            ),
            (
                "input",
                ["checked", "disabled", "type"].iter().cloned().collect(),
            ),
        ].iter()
            .cloned()
            .collect();
        let allowed_classes = [
            (
                "code",
                [
                    "language-bash",
                    "language-clike",
                    "language-glsl",
                    "language-go",
                    "language-ini",
                    "language-javascript",
                    "language-json",
                    "language-markup",
                    "language-protobuf",
                    "language-ruby",
                    "language-rust",
                    "language-scss",
                    "language-sql",
                    "yaml",
                ].iter()
                    .cloned()
                    .collect(),
            ),
        ].iter()
            .cloned()
            .collect();
        let html_sanitizer = Ammonia {
            link_rel: Some("nofollow noopener noreferrer"),
            keep_cleaned_elements: true,
            tags: tags,
            tag_attributes: tag_attributes,
            allowed_classes: allowed_classes,
            ..Ammonia::default()
        };
        MarkdownRenderer {
            html_sanitizer: html_sanitizer,
        }
    }

    /// Renders the given markdown to HTML using the current settings.
    fn to_html(&self, text: &str) -> CargoResult<String> {
        let options = comrak::ComrakOptions {
            ext_autolink: true,
            ext_strikethrough: true,
            ext_table: true,
            ext_tagfilter: true,
            ext_tasklist: true,
            ..comrak::ComrakOptions::default()
        };
        let rendered = comrak::markdown_to_html(text, &options);
        Ok(self.html_sanitizer.clean(&rendered))
    }
}

impl<'a> Default for MarkdownRenderer<'a> {
    fn default() -> Self {
        Self::new()
    }
}

/// Renders Markdown text to sanitized HTML.
fn markdown_to_html(text: &str) -> CargoResult<String> {
    let renderer = MarkdownRenderer::new();
    renderer.to_html(text)
}

/// Any readme with a filename ending in one of these extensions will be rendered as Markdown.
static MARKDOWN_EXTENSIONS: [&'static str; 7] = [
    ".md",
    ".markdown",
    ".mdown",
    ".mdwn",
    ".mkd",
    ".mkdn",
    ".mkdown",
];

/// Renders a readme to sanitized HTML.  An appropriate rendering method is chosen depending
/// on the extension of the supplied filename.
///
/// The returned text should not contain any harmful HTML tag or attribute (such as iframe,
/// onclick, onmouseover, etc.).
///
/// # Examples
///
/// ```
/// use render::render_to_html;
///
/// let text = "[Rust](https://rust-lang.org/) is an awesome *systems programming* language!";
/// let rendered = readme_to_html(text, "README.md")?;
/// ```
pub fn readme_to_html(text: &str, filename: &str) -> CargoResult<String> {
    for e in &MARKDOWN_EXTENSIONS {
        if filename.to_lowercase().ends_with(e) {
            return markdown_to_html(text);
        }
    }

    Ok(encode_minimal(text).replace("\n", "
\n"))
}

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

    #[test]
    fn empty_text() {
        let text = "";
        let result = markdown_to_html(text).unwrap();
        assert_eq!(result, "");
    }

    #[test]
    fn text_with_script_tag() {
        let text = "foo_readme\n\n";
        let result = markdown_to_html(text).unwrap();
        assert_eq!(
            result,
            "
foo_readme
\n<script>alert(\'Hello World\')</script>\n"
        );
    }

    #[test]
    fn text_with_iframe_tag() {
        let text = "foo_readme\n\nalert('Hello World')";
        let result = markdown_to_html(text).unwrap();
        assert_eq!(
            result,
            "
foo_readme
\n<iframe>alert(\'Hello World\')</iframe>\n"
        );
    }

    #[test]
    fn text_with_unknown_tag() {
        let text = "foo_readme\n\n<unknown>alert('Hello World')</unknown>";
        let result = markdown_to_html(text).unwrap();
        assert_eq!(result, "
foo_readme
\n
alert(\'Hello World\')
\n");
    }

    #[test]
    fn text_with_inline_javascript() {
        let text = r#"foo_readme\n\nCrate page"#;
        let result = markdown_to_html(text).unwrap();
        assert_eq!(
            result,
            "
foo_readme\\n\\nCrate page
\n"
        );
    }

    // See https://github.com/kivikakk/comrak/issues/37. This panic happened
    // in comrak 0.1.8 but was fixed in 0.1.9.
    #[test]
    fn text_with_fancy_single_quotes() {
        let text = r#"wb’"#;
        let result = markdown_to_html(text).unwrap();
        assert_eq!(result, "
wb’
\n");
    }

    #[test]
    fn code_block_with_syntax_highlighting() {
        let code_block = r#"```rust \
                            println!("Hello World"); \
                           ```"#;
        let result = markdown_to_html(code_block).unwrap();
        assert!(result.contains(""));
    }

    #[test]
    fn text_with_forbidden_class_attribute() {
        let text = "
Hello World!
";
        let result = markdown_to_html(text).unwrap();
        assert_eq!(result, "
Hello World!
\n");
    }

    #[test]
    fn readme_to_html_renders_markdown() {
        for f in &["readme.md", "README.MARKDOWN", "whatever.mkd"] {
            assert_eq!(
                readme_to_html("*lobster*", f).unwrap(),
                "
lobster
\n"
            );
        }
    }

    #[test]
    fn readme_to_html_renders_other_things() {
        for f in &["readme", "readem.org", "blah.adoc"] {
            assert_eq!(
                readme_to_html("\n\nis my friend\n", f).unwrap(),
                "<script>lobster</script>
\n
\nis my friend
\n"
            );
        }
    }
}
<|endoftext|>


<commit_before>use std::result:: { Result };
use std::collections:: { HashMap };
use std::default::Default;
use lcov_parser:: {
    LCOVParser, LCOVRecord, LineData, FunctionData as FunctionDataRecord,
    BranchData as BranchDataRecord,
    FunctionName, ParseError, FromFile
};
use branch:: { BranchUnit };
use test:: { Test };
use test_sum:: { TestSum };
use file:: { File, CheckSum, FunctionData };

/// Read the trace file of LCOV
///
/// # Examples
///
/// ```
/// use lcov_merge::parse_file;
///
/// let files = parse_file("tests/fixtures/fixture1.info").unwrap();
/// assert_eq!(files.keys().count(), 3);
///
/// let fixture = files.get("/home/vagrant/shared/lcov-merge/tests/fixtures/fixture.c").unwrap();
///
/// assert_eq!(fixture.sum().get_line_count(&4), Some(&1));
/// assert_eq!(fixture.sum().get_line_count(&6), Some(&1));
/// assert_eq!(fixture.sum().get_line_count(&7), Some(&1));
/// assert_eq!(fixture.sum().get_line_count(&8), Some(&1));
/// assert_eq!(fixture.sum().get_line_count(&1), None);
///
/// assert_eq!(fixture.get_test(&"example".to_string()).unwrap().get_line_count(&4), Some(&1));
/// ```
pub fn parse_file(file: &str) -> Result<HashMap<String, File>, ParseError> {
    let mut parse = ReportParser::new();
    parse.parse(file)
}

struct ReportParser {
    test_name: Option<String>,
    source_name: Option<String>,
    tests: HashMap<String, Test>,
    sum: TestSum,
    checksum: CheckSum,
    func: FunctionData,
    files: HashMap<String, File>
}

impl ReportParser {
    fn new() -> Self {
        ReportParser {
            test_name: None,
            source_name: None,
            tests: HashMap::new(),
            sum: TestSum::new(),
            checksum: HashMap::new(),
            func: HashMap::new(),
            files: HashMap::new()
        }
    }
    fn parse(&mut self, file: &str) -> Result<HashMap<String, File>, ParseError> {
        let mut parser = try!(LCOVParser::from_file(file));

        loop {
            let result = try!(parser.next());

            if result.is_none() {
                break;
            }
            let record = result.unwrap();

            match record {
                LCOVRecord::TestName(ref name) => self.on_test_name(name),
                LCOVRecord::SourceFile(ref name) => self.on_source_file(name),
                LCOVRecord::Data(ref data) => self.on_data(data),
                LCOVRecord::FunctionName(ref func_name) => self.on_func_name(func_name),
                LCOVRecord::FunctionData(ref func_data) => self.on_func_data(func_data),
                LCOVRecord::BranchData(ref branch_data) => self.on_branch_data(branch_data),
                LCOVRecord::EndOfRecord => self.on_end_of_record(),
                _ => { continue; }
            };
        }
        Ok(self.files.clone())
    }
    fn on_test_name(&mut self, test_name: &Option<String>) {
        self.test_name = test_name.clone().or( Some(String::new()) );
    }
    fn on_source_file(&mut self, source_name: &String) {
        self.source_name = Some(source_name.clone());

        let current_test_name = self.test_name.as_ref().unwrap();

        if !self.tests.contains_key(current_test_name) {
            self.tests.insert(current_test_name.to_string(), Test::default());
        }
    }
    fn on_data(&mut self, data: &LineData) {
        self.sum.add_line_count(&data.line, &data.count);

        if self.test_name.is_some() {
            let test_name = self.test_name.clone().unwrap();
            let mut test = self.tests.get_mut(&test_name).unwrap();
            test.add_line_count(&data.line, &data.count);
        }

        if data.checksum.is_none() {
            return;
        }

        if !self.checksum.contains_key(&data.line) {
            let checksum_value = data.checksum.clone().unwrap();
            self.checksum.insert(data.line.clone(), checksum_value);
            return;
        }

        let checksum_value = data.checksum.clone().unwrap();
        let current_checksum = self.checksum.get(&data.line).unwrap();
        if current_checksum != &checksum_value {
            println!("{} {}", current_checksum, checksum_value);
        }
    }
    fn on_func_name(&mut self, func_name: &FunctionName) {
        let _ = self.func.entry(func_name.name.clone())
            .or_insert(func_name.line.clone());
    }
    fn on_func_data(&mut self, func_data: &FunctionDataRecord) {
        self.sum.add_func_count(&func_data.name, &func_data.count);

        if self.test_name.is_none() {
            return;
        }

        let test_name = self.test_name.clone().unwrap();
        let mut test = self.tests.get_mut(&test_name).unwrap();

        test.add_func_count(&func_data.name, &func_data.count);
    }
    fn on_branch_data(&mut self, branch_data: &BranchDataRecord) {
        let ref branch_unit = BranchUnit::new(branch_data.block.clone(), branch_data.branch.clone());

        self.sum.add_branch_count(
            &branch_data.line,
            branch_unit,
            &branch_data.taken
        );

        if self.test_name.is_some() {
            let ref test_name = self.test_name.clone().unwrap();
            let mut test = self.tests.get_mut(test_name).unwrap();
            test.add_branch_count(
                &branch_data.line,
                branch_unit,
                &branch_data.taken
            );
        }
    }
    fn on_end_of_record(&mut self) {
        let source_name = self.source_name.clone().unwrap();

        self.files.insert(source_name, File::new(
            self.sum.clone(),
            self.tests.clone(),
            self.checksum.clone(),
            self.func.clone()
        ));
    }
}
<commit_msg>replace to += operation<commit_after>use std::result:: { Result };
use std::collections:: { HashMap };
use std::default::Default;
use lcov_parser:: {
    LCOVParser, LCOVRecord, LineData, FunctionData as FunctionDataRecord,
    BranchData as BranchDataRecord,
    FunctionName, ParseError, FromFile
};
use branch:: { BranchUnit };
use test:: { Test };
use test_sum:: { TestSum };
use file:: { File, CheckSum, FunctionData };

/// Read the trace file of LCOV
///
/// # Examples
///
/// ```
/// use lcov_merge::parse_file;
///
/// let files = parse_file("tests/fixtures/fixture1.info").unwrap();
/// assert_eq!(files.keys().count(), 3);
///
/// let fixture = files.get("/home/vagrant/shared/lcov-merge/tests/fixtures/fixture.c").unwrap();
///
/// assert_eq!(fixture.sum().get_line_count(&4), Some(&1));
/// assert_eq!(fixture.sum().get_line_count(&6), Some(&1));
/// assert_eq!(fixture.sum().get_line_count(&7), Some(&1));
/// assert_eq!(fixture.sum().get_line_count(&8), Some(&1));
/// assert_eq!(fixture.sum().get_line_count(&1), None);
///
/// assert_eq!(fixture.get_test(&"example".to_string()).unwrap().get_line_count(&4), Some(&1));
/// ```
pub fn parse_file(file: &str) -> Result<HashMap<String, File>, ParseError> {
    let mut parse = ReportParser::new();
    parse.parse(file)
}

struct ReportParser {
    test_name: Option<String>,
    source_name: Option<String>,
    tests: HashMap<String, Test>,
    sum: TestSum,
    checksum: CheckSum,
    func: FunctionData,
    files: HashMap<String, File>
}

impl ReportParser {
    fn new() -> Self {
        ReportParser {
            test_name: None,
            source_name: None,
            tests: HashMap::new(),
            sum: TestSum::new(),
            checksum: HashMap::new(),
            func: HashMap::new(),
            files: HashMap::new()
        }
    }
    fn parse(&mut self, file: &str) -> Result<HashMap<String, File>, ParseError> {
        let mut parser = try!(LCOVParser::from_file(file));

        loop {
            let result = try!(parser.next());

            if result.is_none() {
                break;
            }
            let record = result.unwrap();

            match record {
                LCOVRecord::TestName(ref name) => self.on_test_name(name),
                LCOVRecord::SourceFile(ref name) => self.on_source_file(name),
                LCOVRecord::Data(ref data) => self.on_data(data),
                LCOVRecord::FunctionName(ref func_name) => self.on_func_name(func_name),
                LCOVRecord::FunctionData(ref func_data) => self.on_func_data(func_data),
                LCOVRecord::BranchData(ref branch_data) => self.on_branch_data(branch_data),
                LCOVRecord::EndOfRecord => self.on_end_of_record(),
                _ => { continue; }
            };
        }
        Ok(self.files.clone())
    }
    fn on_test_name(&mut self, test_name: &Option<String>) {
        self.test_name = test_name.clone().or( Some(String::new()) );
    }
    fn on_source_file(&mut self, source_name: &String) {
        self.source_name = Some(source_name.clone());

        let current_test_name = self.test_name.as_ref().unwrap();

        if !self.tests.contains_key(current_test_name) {
            self.tests.insert(current_test_name.to_string(), Test::default());
        }
    }
    fn on_data(&mut self, data: &LineData) {
        self.sum += data;

        if self.test_name.is_some() {
            let test_name = self.test_name.clone().unwrap();
            let mut test = self.tests.get_mut(&test_name).unwrap();
            test.add_line_count(&data.line, &data.count);
        }

        if data.checksum.is_none() {
            return;
        }

        if !self.checksum.contains_key(&data.line) {
            let checksum_value = data.checksum.clone().unwrap();
            self.checksum.insert(data.line.clone(), checksum_value);
            return;
        }

        let checksum_value = data.checksum.clone().unwrap();
        let current_checksum = self.checksum.get(&data.line).unwrap();
        if current_checksum != &checksum_value {
            println!("{} {}", current_checksum, checksum_value);
        }
    }
    fn on_func_name(&mut self, func_name: &FunctionName) {
        let _ = self.func.entry(func_name.name.clone())
            .or_insert(func_name.line.clone());
    }
    fn on_func_data(&mut self, func_data: &FunctionDataRecord) {
        self.sum += func_data;

        if self.test_name.is_none() {
            return;
        }

        let test_name = self.test_name.clone().unwrap();
        let mut test = self.tests.get_mut(&test_name).unwrap();

        test.add_func_count(&func_data.name, &func_data.count);
    }
    fn on_branch_data(&mut self, branch_data: &BranchDataRecord) {
        let ref branch_unit = BranchUnit::new(branch_data.block.clone(), branch_data.branch.clone());

        self.sum += branch_data;

        if self.test_name.is_some() {
            let ref test_name = self.test_name.clone().unwrap();
            let mut test = self.tests.get_mut(test_name).unwrap();
            test.add_branch_count(
                &branch_data.line,
                branch_unit,
                &branch_data.taken
            );
        }
    }
    fn on_end_of_record(&mut self) {
        let source_name = self.source_name.clone().unwrap();

        self.files.insert(source_name, File::new(
            self.sum.clone(),
            self.tests.clone(),
            self.checksum.clone(),
            self.func.clone()
        ));
    }
}
<|endoftext|>


<commit_before>use regex::Regex;
use term;

pub fn print_matches(pattern: &str, subject: &str) {
    // Attempt to compile the given regex pattern.
    let regex = match Regex::new(&pattern) {
        Ok(result) => { result },
        Err(_) => {
            println!("Invalid regular expression pattern '{}'.", pattern);
            return;
        }
    };

    // Grab a terminal object to write formatted output with.
    let mut term = term::stdout().unwrap();

    // Loop over each match in the subject and pretty-print the match as well as
    // any trailing, non-matching text preceding the match.
    let mut last_index = 0;
    for captures in regex.captures_iter(subject) {
        let positions = captures.pos(0).unwrap();

        print!("{}", &subject[last_index .. positions.0]);
        term.bg(term::color::BLUE).unwrap();
        print!("{}", captures.at(0).unwrap());
        term.reset().unwrap();

        last_index = positions.1;
    }

    // Print trailing unmatched text if there is any.
    if last_index < subject.len() {
        print!("{}", &subject[last_index ..]);
    }
}
<commit_msg>Separate print functions<commit_after>use regex::{Captures, Regex};
use term::{self, Terminal};

/// Finds all matches for the given pattern on the given subject and prints out
/// the subject with matches highlighted.
pub fn print_matches(pattern: &str, subject: &str) {
    // Attempt to compile the given regex pattern.
    let regex = match Regex::new(&pattern) {
        Ok(result) => { result },
        Err(_) => {
            println!("Invalid regular expression pattern '{}'.", pattern);
            return;
        }
    };

    // Loop over each match in the subject and pretty-print the match as well as
    // any trailing, non-matching text preceding the match.
    let mut last_index = 0;
    for captures in regex.captures_iter(subject) {
        let positions = captures.pos(0).unwrap();

        print!("{}", &subject[last_index .. positions.0]);
        print_captures(captures);
        last_index = positions.1;
    }

    // Print trailing unmatched text if there is any.
    if last_index < subject.len() {
        print!("{}", &subject[last_index ..]);
    }
}

/// Prints out a capture group using color formatting.
fn print_captures(captures: Captures) {
    // Grab a terminal object to write formatted output with.
    let mut terminal = term::stdout().unwrap();

    terminal.bg(term::color::BLUE).unwrap();
    print!("{}", captures.at(0).unwrap());
    terminal.reset().unwrap();
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::collections::hash_map::{Occupied, Vacant};
use iron::{Request, Response, Handler, IronResult, Error, IronError, Set};
use iron::{status, method};
use iron::response::modifiers::Status;
use iron::typemap::Assoc;
use recognizer::Router as Recognizer;
use recognizer::{Match, Params};

/// `Router` provides an interface for creating complex routes as middleware
/// for the Iron framework.
pub struct Router {
    // The routers, specialized by method.
    routers: HashMap<method::Method, Recognizer<Box<Handler + Send + Sync>>>,
    error: Option<Box<Handler + Send + Sync>>
}

#[deriving(Show)]
/// The error thrown by router if there is no matching route.
pub struct NoRoute;

impl Error for NoRoute {
    fn name(&self) -> &'static str { "No Route" }
}

impl Router {
    /// `new` constructs a new, blank `Router`.
    pub fn new() -> Router { Router { routers: HashMap::new(), error: None } }
    /// Add a new route to a `Router`, matching both a method and glob pattern.
    ///
    /// `route` supports glob patterns: `*` for a single wildcard segment and
    /// `:param` for matching storing that segment of the request url in the `Params`
    /// object, which is stored in the request `extensions`.
    ///
    /// For instance, to route `Get` requests on any route matching
    /// `/users/:userid/:friend` and store `userid` and `friend` in
    /// the exposed Params object:
    ///
    /// ```ignore
    /// router.route(http::method::Get, "/users/:userid/:friendid", controller);
    /// ```
    ///
    /// The controller provided to route can be any `Handler`, which allows
    /// extreme flexibility when handling routes. For instance, you could provide
    /// a `Chain`, a `Handler`, which contains an authorization middleware and
    /// a controller function, so that you can confirm that the request is
    /// authorized for this route before handling it.
    pub fn route<H: Handler, S: Str>(&mut self, method: method::Method, glob: S, handler: H) -> &mut Router {
        match self.routers.entry(method) {
            Vacant(entry)   => entry.set(Recognizer::new()),
            Occupied(entry) => entry.into_mut()
        }.add(glob.as_slice(), box handler as Box<Handler + Send + Sync>);
        self
    }

    /// Like route, but specialized to the `Get` method.
    pub fn get<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Get, glob, handler)
    }

    /// Like route, but specialized to the `Post` method.
    pub fn post<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Post, glob, handler)
    }

    /// Like route, but specialized to the `Put` method.
    pub fn put<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Put, glob, handler)
    }

    /// Like route, but specialized to the `Delete` method.
    pub fn delete<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Delete, glob, handler)
    }

    /// Like route, but specialized to the `Head` method.
    pub fn head<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Head, glob, handler)
    }

    /// Like route, but specialized to the `Patch` method.
    pub fn patch<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Patch, glob, handler)
    }

    /// Like route, but specialized to the `Options` method.
    pub fn options<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Options, glob, handler)
    }

    /// Add a Handler to be used for this Router's `catch` method.
    pub fn error<H: Handler>(&mut self, handler: H) -> &mut Router {
        self.error = Some(box handler as Box<Handler + Send + Sync>);
        self
    }

    fn recognize<'a>(&'a self, method: &method::Method, path: &str)
                     -> Option<Match<&'a Box<Handler + Send + Sync>>> {
        self.routers.get(method).and_then(|router| router.recognize(path).ok())
    }
}

impl Assoc<Params> for Router {}

impl Handler for Router {
    fn call(&self, req: &mut Request) -> IronResult<Response> {
        let matched = match self.recognize(&req.method, req.url.path.connect("/").as_slice()) {
            Some(matched) => matched,
            // No match.
            None => return Err(box NoRoute as IronError)
        };

        req.extensions.insert::<Router, Params>(matched.params);
        matched.handler.call(req)
    }

    fn catch(&self, req: &mut Request, err: IronError) -> (Response, IronResult<()>) {
        match self.error {
            Some(ref error_handler) => error_handler.catch(req, err),
            // Error that is not caught by anything!
            None => (Response::new().set(Status(status::InternalServerError)), Err(err))
        }
    }
}
<commit_msg>(doc) Remove use http from the docs.<commit_after>use std::collections::HashMap;
use std::collections::hash_map::{Occupied, Vacant};
use iron::{Request, Response, Handler, IronResult, Error, IronError, Set};
use iron::{status, method};
use iron::response::modifiers::Status;
use iron::typemap::Assoc;
use recognizer::Router as Recognizer;
use recognizer::{Match, Params};

/// `Router` provides an interface for creating complex routes as middleware
/// for the Iron framework.
pub struct Router {
    // The routers, specialized by method.
    routers: HashMap<method::Method, Recognizer<Box<Handler + Send + Sync>>>,
    error: Option<Box<Handler + Send + Sync>>
}

#[deriving(Show)]
/// The error thrown by router if there is no matching route.
pub struct NoRoute;

impl Error for NoRoute {
    fn name(&self) -> &'static str { "No Route" }
}

impl Router {
    /// `new` constructs a new, blank `Router`.
    pub fn new() -> Router { Router { routers: HashMap::new(), error: None } }
    /// Add a new route to a `Router`, matching both a method and glob pattern.
    ///
    /// `route` supports glob patterns: `*` for a single wildcard segment and
    /// `:param` for matching storing that segment of the request url in the `Params`
    /// object, which is stored in the request `extensions`.
    ///
    /// For instance, to route `Get` requests on any route matching
    /// `/users/:userid/:friend` and store `userid` and `friend` in
    /// the exposed Params object:
    ///
    /// ```ignore
    /// router.route(iron::method::Method::Get, "/users/:userid/:friendid", controller);
    /// ```
    ///
    /// The controller provided to route can be any `Handler`, which allows
    /// extreme flexibility when handling routes. For instance, you could provide
    /// a `Chain`, a `Handler`, which contains an authorization middleware and
    /// a controller function, so that you can confirm that the request is
    /// authorized for this route before handling it.
    pub fn route<H: Handler, S: Str>(&mut self, method: method::Method, glob: S, handler: H) -> &mut Router {
        match self.routers.entry(method) {
            Vacant(entry)   => entry.set(Recognizer::new()),
            Occupied(entry) => entry.into_mut()
        }.add(glob.as_slice(), box handler as Box<Handler + Send + Sync>);
        self
    }

    /// Like route, but specialized to the `Get` method.
    pub fn get<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Get, glob, handler)
    }

    /// Like route, but specialized to the `Post` method.
    pub fn post<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Post, glob, handler)
    }

    /// Like route, but specialized to the `Put` method.
    pub fn put<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Put, glob, handler)
    }

    /// Like route, but specialized to the `Delete` method.
    pub fn delete<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Delete, glob, handler)
    }

    /// Like route, but specialized to the `Head` method.
    pub fn head<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Head, glob, handler)
    }

    /// Like route, but specialized to the `Patch` method.
    pub fn patch<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Patch, glob, handler)
    }

    /// Like route, but specialized to the `Options` method.
    pub fn options<H: Handler, S: Str>(&mut self, glob: S, handler: H) -> &mut Router {
        self.route(method::Options, glob, handler)
    }

    /// Add a Handler to be used for this Router's `catch` method.
    pub fn error<H: Handler>(&mut self, handler: H) -> &mut Router {
        self.error = Some(box handler as Box<Handler + Send + Sync>);
        self
    }

    fn recognize<'a>(&'a self, method: &method::Method, path: &str)
                     -> Option<Match<&'a Box<Handler + Send + Sync>>> {
        self.routers.get(method).and_then(|router| router.recognize(path).ok())
    }
}

impl Assoc<Params> for Router {}

impl Handler for Router {
    fn call(&self, req: &mut Request) -> IronResult<Response> {
        let matched = match self.recognize(&req.method, req.url.path.connect("/").as_slice()) {
            Some(matched) => matched,
            // No match.
            None => return Err(box NoRoute as IronError)
        };

        req.extensions.insert::<Router, Params>(matched.params);
        matched.handler.call(req)
    }

    fn catch(&self, req: &mut Request, err: IronError) -> (Response, IronResult<()>) {
        match self.error {
            Some(ref error_handler) => error_handler.catch(req, err),
            // Error that is not caught by anything!
            None => (Response::new().set(Status(status::InternalServerError)), Err(err))
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::error::Error;
use std::fmt;
use std::sync::Arc;

use iron::{Request, Response, Handler, IronResult, IronError};
use iron::{status, method, headers};
use iron::typemap::Key;
use iron::modifiers::Redirect;

use recognizer::Router as Recognizer;
use recognizer::{Match, Params};

pub struct RouterInner {
    // The routers, specialized by method.
    pub routers: Recognizer<HashMap<method::Method, Arc<Handler>>>,
    // Routes that accept any method.
    pub wildcard: Recognizer<Arc<Handler>>,
    // Used in URL generation.
    pub route_ids: HashMap<String, String>
}

/// `Router` provides an interface for creating complex routes as middleware
/// for the Iron framework.
pub struct Router {
    inner: Arc<RouterInner>
}

impl Router {
    /// Construct a new, empty `Router`.
    ///
    /// ```
    /// # use router::Router;
    /// let router = Router::new();
    /// ```
    pub fn new() -> Router {
        Router {
            inner: Arc::new(RouterInner {
                routers: Recognizer::new(),
                wildcard: Recognizer::new(),
                route_ids: HashMap::new()
            })
        }
    }

    fn mut_inner(&mut self) -> &mut RouterInner {
        Arc::get_mut(&mut self.inner).expect("Cannot modify router at this point.")
    }

    /// Add a new route to a `Router`, matching both a method and glob pattern.
    ///
    /// `route` supports glob patterns: `*` for a single wildcard segment and
    /// `:param` for matching storing that segment of the request url in the `Params`
    /// object, which is stored in the request `extensions`.
    ///
    /// For instance, to route `Get` requests on any route matching
    /// `/users/:userid/:friend` and store `userid` and `friend` in
    /// the exposed Params object:
    ///
    /// ```ignore
    /// let mut router = Router::new();
    /// router.route(method::Get, "/users/:userid/:friendid", controller, "user_friend");
    /// ```
    ///
    /// `route_id` is a unique name for your route, and is used when generating an URL with
    /// `url_for`.
    ///
    /// The controller provided to route can be any `Handler`, which allows
    /// extreme flexibility when handling routes. For instance, you could provide
    /// a `Chain`, a `Handler`, which contains an authorization middleware and
    /// a controller function, so that you can confirm that the request is
    /// authorized for this route before handling it.
    pub fn route<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, method: method::Method, glob: S, handler: H, route_id: I) -> &mut Router {

        let mut hash: HashMap<method::Method, Arc<Handler>>;

        if let Some(s) = self.mut_inner().routers.recognize(glob.as_ref()).ok() {
            hash = s.handler.clone();
        } else {
            hash = HashMap::new();
        }

        hash.insert(method, Arc::new(handler));
        self.mut_inner().routers.add(glob.as_ref(), hash);
        self.route_id(route_id.as_ref(), glob.as_ref());
        self
    }

    fn route_id(&mut self, id: &str, glob: &str) {
        let mut mut_inner = self.mut_inner();
        let ref mut route_ids = mut_inner.route_ids;

        match route_ids.get(id) {
            Some(other_glob) if glob != other_glob => panic!("Duplicate route_id: {}", id),
            _ => ()
        };

        route_ids.insert(id.to_owned(), glob.to_owned());
    }

    /// Like route, but specialized to the `Get` method.
    pub fn get<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Get, glob, handler, route_id)
    }

    /// Like route, but specialized to the `Post` method.
    pub fn post<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Post, glob, handler, route_id)
    }

    /// Like route, but specialized to the `Put` method.
    pub fn put<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Put, glob, handler, route_id)
    }

    /// Like route, but specialized to the `Delete` method.
    pub fn delete<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Delete, glob, handler, route_id)
    }

    /// Like route, but specialized to the `Head` method.
    pub fn head<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Head, glob, handler, route_id)
    }

    /// Like route, but specialized to the `Patch` method.
    pub fn patch<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Patch, glob, handler, route_id)
    }

    /// Like route, but specialized to the `Options` method.
    pub fn options<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Options, glob, handler, route_id)
    }

    /// Route will match any method, including gibberish.
    /// In case of ambiguity, handlers specific to methods will be preferred.
    pub fn any<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.mut_inner().wildcard.add(glob.as_ref(), Arc::new(handler));
        self.route_id(route_id.as_ref(), glob.as_ref());
        self
    }

    fn recognize(&self, method: &method::Method, path: &str)
                      -> Result<Match<&Arc<Handler>>, RouterError> {
        let mut found_path = false;
        if let Ok(s) = self.inner.routers.recognize(path) {
            found_path = true;
            if let Some(h) = s.handler.get(method) {
                return Ok(Match::new(h, s.params));
            }
        };

        if let Ok(x) = self.inner.wildcard.recognize(path) {
            return Ok(x)
        };

        Err(if found_path { RouterError::MethodNotAllowed } else { RouterError::NotFound })
    }

    fn handle_options(&self, path: &str) -> Response {
        static METHODS: &'static [method::Method] =
            &[method::Get, method::Post, method::Put,
              method::Delete, method::Head, method::Patch];

        // Get all the available methods and return them.
        let mut options = vec![];

        for method in METHODS.iter() {
            if let Ok(s) = self.inner.routers.recognize(path) {
                if let Some(_) = s.handler.get(method) {
                    options.push(method.clone());
                }
            }
        }
        // If GET is there, HEAD is also there.
        if options.contains(&method::Get) && !options.contains(&method::Head) {
            options.push(method::Head);
        }

        let mut res = Response::with(status::Ok);
        res.headers.set(headers::Allow(options));
        res
    }

    // Tests for a match by adding or removing a trailing slash.
    fn redirect_slash(&self, req: &Request) -> Option<IronError> {
        let mut url = req.url.clone();
        let mut path = url.path().join("/");

        if let Some(last_char) = path.chars().last() {
            {
                let mut path_segments = url.as_mut().path_segments_mut().unwrap();
                if last_char == '/' {
                    // We didn't recognize anything without a trailing slash; try again with one appended.
                    path.pop();
                    path_segments.pop();
                } else {
                    // We didn't recognize anything with a trailing slash; try again without it.
                    path.push('/');
                    path_segments.push("");
                }
            }
        }

        self.recognize(&req.method, &path).ok().and(
            Some(IronError::new(RouterError::TrailingSlash,
                                (status::MovedPermanently, Redirect(url))))
        )
    }

    fn handle_method(&self, req: &mut Request, path: &str) -> IronResult<Response> {
        match self.recognize(&req.method, path) {
            Ok(matched) => {
                req.extensions.insert::<Router>(matched.params);
                req.extensions.insert::<RouterInner>(self.inner.clone());
                matched.handler.handle(req)
            },
            Err(RouterError::MethodNotAllowed) => {
                Err(IronError::new(RouterError::MethodNotAllowed, status::MethodNotAllowed))
            },
            Err(_) => {
                if let Some(err) = self.redirect_slash(req) {
                    return Err(err)
                };

                match req.method {
                    method::Options => Ok(self.handle_options(path)),
                    // For HEAD, fall back to GET. Hyper ensures no response body is written.
                    method::Head => {
                        req.method = method::Get;
                        self.handle_method(req, path)
                    },
                    _ => Err(IronError::new(RouterError::NotFound, status::NotFound))
                }
            }
        }
    }
}

impl Key for Router { type Value = Params; }

impl Key for RouterInner { type Value = Arc<RouterInner>; }

impl Handler for Router {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        let path = req.url.path().join("/");
        self.handle_method(req, &path)
    }
}

/// A set of errors that can occur parsing HTTP streams.
#[derive(Debug, PartialEq)]
pub enum RouterError {
    /// The error thrown by router if there is no matching method in existing route.
    MethodNotAllowed,
    /// The error thrown by router if there is no matching route.
    NotFound,
    /// The error thrown by router if a request was redirected by adding or removing a trailing slash.
    TrailingSlash
}

impl fmt::Display for RouterError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(&self.description())
    }
}

impl Error for RouterError {
    fn description(&self) -> &str {
        match *self {
            RouterError::MethodNotAllowed => "Method Not Allowed",
            RouterError::NotFound => "No matching route found.",
            RouterError::TrailingSlash => "The request had a trailing slash."
        }
    }
}

/// The error thrown by router if there is no matching route,
/// it is always accompanied by a NotFound response.
#[derive(Debug, PartialEq, Eq)]
pub struct NoRoute;

impl fmt::Display for NoRoute {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("No matching route found.")
    }
}

impl Error for NoRoute {
    fn description(&self) -> &str { "No Route" }
}

/// The error thrown by router if a request was redirected
/// by adding or removing a trailing slash.
#[derive(Debug, PartialEq, Eq)]
pub struct TrailingSlash;

impl fmt::Display for TrailingSlash {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("The request had a trailing slash.")
    }
}

impl Error for TrailingSlash {
    fn description(&self) -> &str { "Trailing Slash" }
}

#[cfg(test)]
mod test {
    use super::{Router, RouterError};
    use iron::{headers, method, status, Request, Response};

    #[test]
    fn test_handle_options_post() {
        let mut router = Router::new();
        router.post("/", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        let resp = router.handle_options("/");
        let headers = resp.headers.get::<headers::Allow>().unwrap();
        let expected = headers::Allow(vec![method::Method::Post]);
        assert_eq!(&expected, headers);
    }

    #[test]
    fn test_handle_options_get_head() {
        let mut router = Router::new();
        router.get("/", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        let resp = router.handle_options("/");
        let headers = resp.headers.get::<headers::Allow>().unwrap();
        let expected = headers::Allow(vec![method::Method::Get, method::Method::Head]);
        assert_eq!(&expected, headers);
    }

    #[test]
    fn test_not_allowed_method() {
        let mut router = Router::new();

        router.post("/post", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");

        router.get("/post/", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "another_route");

        match router.recognize(&method::Get, "/post") {
            Ok(_) => {
                panic!();
            },
            Err(e) => {
                assert_eq!(RouterError::MethodNotAllowed, e);
            }
        }
    }

    #[test]
    fn test_handle_any_ok() {
        let mut router = Router::new();
        router.post("/post", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        router.any("/post", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        router.put("/post", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        router.any("/get", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "any");

        assert!(router.recognize(&method::Get, "/post").is_ok());
        assert!(router.recognize(&method::Get, "/get").is_ok());
    }

    #[test]
    fn test_request() {
        let mut router = Router::new();
        router.post("/post", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        router.get("/post", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");

        assert!(router.recognize(&method::Post, "/post").is_ok());
        assert!(router.recognize(&method::Get, "/post").is_ok());
        assert!(router.recognize(&method::Put, "/post").is_err());
        assert!(router.recognize(&method::Get, "/post/").is_err());
    }

    #[test]
    fn test_not_found() {
        let mut router = Router::new();
        router.put("/put", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        match router.recognize(&method::Patch, "/patch") {
            Ok(_) => {
                panic!();
            },
            Err(e) => {
                assert_eq!(RouterError::NotFound, e);
            }
        }
    }

    #[test]
    #[should_panic]
    fn test_same_route_id() {
        let mut router = Router::new();
        router.put("/put", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "my_route_id");
        router.get("/get", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "my_route_id");
    }

    #[test]
    fn test_wildcard_regression() {
        let mut router = Router::new();
        router.options("*", |_: &mut Request| Ok(Response::with((status::Ok, ""))), "id1");
        router.put("/upload/*filename", |_: &mut Request| Ok(Response::with((status::Ok, ""))), "id2");
        assert!(router.recognize(&method::Options, "/foo").is_ok());
        assert!(router.recognize(&method::Put, "/foo").is_err());
        assert!(router.recognize(&method::Put, "/upload/foo").is_ok());
    }
}
<commit_msg>Revert "Simplify logic"<commit_after>use std::collections::HashMap;
use std::error::Error;
use std::fmt;
use std::sync::Arc;

use iron::{Request, Response, Handler, IronResult, IronError};
use iron::{status, method, headers};
use iron::typemap::Key;
use iron::modifiers::Redirect;

use recognizer::Router as Recognizer;
use recognizer::{Match, Params};

pub struct RouterInner {
    // The routers, specialized by method.
    pub routers: Recognizer<HashMap<method::Method, Arc<Handler>>>,
    // Routes that accept any method.
    pub wildcard: Recognizer<Arc<Handler>>,
    // Used in URL generation.
    pub route_ids: HashMap<String, String>
}

/// `Router` provides an interface for creating complex routes as middleware
/// for the Iron framework.
pub struct Router {
    inner: Arc<RouterInner>
}

impl Router {
    /// Construct a new, empty `Router`.
    ///
    /// ```
    /// # use router::Router;
    /// let router = Router::new();
    /// ```
    pub fn new() -> Router {
        Router {
            inner: Arc::new(RouterInner {
                routers: Recognizer::new(),
                wildcard: Recognizer::new(),
                route_ids: HashMap::new()
            })
        }
    }

    fn mut_inner(&mut self) -> &mut RouterInner {
        Arc::get_mut(&mut self.inner).expect("Cannot modify router at this point.")
    }

    /// Add a new route to a `Router`, matching both a method and glob pattern.
    ///
    /// `route` supports glob patterns: `*` for a single wildcard segment and
    /// `:param` for matching storing that segment of the request url in the `Params`
    /// object, which is stored in the request `extensions`.
    ///
    /// For instance, to route `Get` requests on any route matching
    /// `/users/:userid/:friend` and store `userid` and `friend` in
    /// the exposed Params object:
    ///
    /// ```ignore
    /// let mut router = Router::new();
    /// router.route(method::Get, "/users/:userid/:friendid", controller, "user_friend");
    /// ```
    ///
    /// `route_id` is a unique name for your route, and is used when generating an URL with
    /// `url_for`.
    ///
    /// The controller provided to route can be any `Handler`, which allows
    /// extreme flexibility when handling routes. For instance, you could provide
    /// a `Chain`, a `Handler`, which contains an authorization middleware and
    /// a controller function, so that you can confirm that the request is
    /// authorized for this route before handling it.
    pub fn route<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, method: method::Method, glob: S, handler: H, route_id: I) -> &mut Router {

        let mut hash: HashMap<method::Method, Arc<Handler>>;

        if let Some(s) = self.mut_inner().routers.recognize(glob.as_ref()).ok() {
            hash = s.handler.clone();
        } else {
            hash = HashMap::new();
        }

        hash.insert(method, Arc::new(handler));
        self.mut_inner().routers.add(glob.as_ref(), hash);
        self.route_id(route_id.as_ref(), glob.as_ref());
        self
    }

    fn route_id(&mut self, id: &str, glob: &str) {
        let mut mut_inner = self.mut_inner();
        let ref mut route_ids = mut_inner.route_ids;

        match route_ids.get(id) {
            Some(other_glob) if glob != other_glob => panic!("Duplicate route_id: {}", id),
            _ => ()
        };

        route_ids.insert(id.to_owned(), glob.to_owned());
    }

    /// Like route, but specialized to the `Get` method.
    pub fn get<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Get, glob, handler, route_id)
    }

    /// Like route, but specialized to the `Post` method.
    pub fn post<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Post, glob, handler, route_id)
    }

    /// Like route, but specialized to the `Put` method.
    pub fn put<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Put, glob, handler, route_id)
    }

    /// Like route, but specialized to the `Delete` method.
    pub fn delete<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Delete, glob, handler, route_id)
    }

    /// Like route, but specialized to the `Head` method.
    pub fn head<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Head, glob, handler, route_id)
    }

    /// Like route, but specialized to the `Patch` method.
    pub fn patch<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Patch, glob, handler, route_id)
    }

    /// Like route, but specialized to the `Options` method.
    pub fn options<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.route(method::Options, glob, handler, route_id)
    }

    /// Route will match any method, including gibberish.
    /// In case of ambiguity, handlers specific to methods will be preferred.
    pub fn any<S: AsRef<str>, H: Handler, I: AsRef<str>>(&mut self, glob: S, handler: H, route_id: I) -> &mut Router {
        self.mut_inner().wildcard.add(glob.as_ref(), Arc::new(handler));
        self.route_id(route_id.as_ref(), glob.as_ref());
        self
    }

    fn recognize(&self, method: &method::Method, path: &str)
                      -> Result<Match<&Arc<Handler>>, RouterError> {
        match self.inner.routers.recognize(path)
            .map(|s|
                match s.handler.get(method) {
                    Some(h) => Ok(Match::new(h, s.params)),
                    None => self.inner.wildcard.recognize(path).ok().map_or(Err(RouterError::MethodNotAllowed), |s| Ok(s))
                }
        ).map_err(|_| self.inner.wildcard.recognize(path).ok().map_or(Err(RouterError::NotFound), |s| Ok(s))) {
            Ok(s) => s,
            Err(e) => e
        }
    }

    fn handle_options(&self, path: &str) -> Response {
        static METHODS: &'static [method::Method] =
            &[method::Get, method::Post, method::Put,
              method::Delete, method::Head, method::Patch];

        // Get all the available methods and return them.
        let mut options = vec![];

        for method in METHODS.iter() {
            if let Ok(s) = self.inner.routers.recognize(path) {
                if let Some(_) = s.handler.get(method) {
                    options.push(method.clone());
                }
            }
        }
        // If GET is there, HEAD is also there.
        if options.contains(&method::Get) && !options.contains(&method::Head) {
            options.push(method::Head);
        }

        let mut res = Response::with(status::Ok);
        res.headers.set(headers::Allow(options));
        res
    }

    // Tests for a match by adding or removing a trailing slash.
    fn redirect_slash(&self, req: &Request) -> Option<IronError> {
        let mut url = req.url.clone();
        let mut path = url.path().join("/");

        if let Some(last_char) = path.chars().last() {
            {
                let mut path_segments = url.as_mut().path_segments_mut().unwrap();
                if last_char == '/' {
                    // We didn't recognize anything without a trailing slash; try again with one appended.
                    path.pop();
                    path_segments.pop();
                } else {
                    // We didn't recognize anything with a trailing slash; try again without it.
                    path.push('/');
                    path_segments.push("");
                }
            }
        }

        self.recognize(&req.method, &path).ok().and(
            Some(IronError::new(RouterError::TrailingSlash,
                                (status::MovedPermanently, Redirect(url))))
        )
    }

    fn handle_method(&self, req: &mut Request, path: &str) -> IronResult<Response> {
        match self.recognize(&req.method, path) {
            Ok(matched) => {
                req.extensions.insert::<Router>(matched.params);
                req.extensions.insert::<RouterInner>(self.inner.clone());
                matched.handler.handle(req)
            },
            Err(RouterError::MethodNotAllowed) => {
                Err(IronError::new(RouterError::MethodNotAllowed, status::MethodNotAllowed))
            },
            Err(_) => {
                if let Some(err) = self.redirect_slash(req) {
                    return Err(err)
                };

                match req.method {
                    method::Options => Ok(self.handle_options(path)),
                    // For HEAD, fall back to GET. Hyper ensures no response body is written.
                    method::Head => {
                        req.method = method::Get;
                        self.handle_method(req, path)
                    },
                    _ => Err(IronError::new(RouterError::NotFound, status::NotFound))
                }
            }
        }
    }
}

impl Key for Router { type Value = Params; }

impl Key for RouterInner { type Value = Arc<RouterInner>; }

impl Handler for Router {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        let path = req.url.path().join("/");
        self.handle_method(req, &path)
    }
}

/// A set of errors that can occur parsing HTTP streams.
#[derive(Debug, PartialEq)]
pub enum RouterError {
    /// The error thrown by router if there is no matching method in existing route.
    MethodNotAllowed,
    /// The error thrown by router if there is no matching route.
    NotFound,
    /// The error thrown by router if a request was redirected by adding or removing a trailing slash.
    TrailingSlash
}

impl fmt::Display for RouterError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(&self.description())
    }
}

impl Error for RouterError {
    fn description(&self) -> &str {
        match *self {
            RouterError::MethodNotAllowed => "Method Not Allowed",
            RouterError::NotFound => "No matching route found.",
            RouterError::TrailingSlash => "The request had a trailing slash."
        }
    }
}

/// The error thrown by router if there is no matching route,
/// it is always accompanied by a NotFound response.
#[derive(Debug, PartialEq, Eq)]
pub struct NoRoute;

impl fmt::Display for NoRoute {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("No matching route found.")
    }
}

impl Error for NoRoute {
    fn description(&self) -> &str { "No Route" }
}

/// The error thrown by router if a request was redirected
/// by adding or removing a trailing slash.
#[derive(Debug, PartialEq, Eq)]
pub struct TrailingSlash;

impl fmt::Display for TrailingSlash {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("The request had a trailing slash.")
    }
}

impl Error for TrailingSlash {
    fn description(&self) -> &str { "Trailing Slash" }
}

#[cfg(test)]
mod test {
    use super::{Router, RouterError};
    use iron::{headers, method, status, Request, Response};

    #[test]
    fn test_handle_options_post() {
        let mut router = Router::new();
        router.post("/", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        let resp = router.handle_options("/");
        let headers = resp.headers.get::<headers::Allow>().unwrap();
        let expected = headers::Allow(vec![method::Method::Post]);
        assert_eq!(&expected, headers);
    }

    #[test]
    fn test_handle_options_get_head() {
        let mut router = Router::new();
        router.get("/", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        let resp = router.handle_options("/");
        let headers = resp.headers.get::<headers::Allow>().unwrap();
        let expected = headers::Allow(vec![method::Method::Get, method::Method::Head]);
        assert_eq!(&expected, headers);
    }

    #[test]
    fn test_not_allowed_method() {
        let mut router = Router::new();

        router.post("/post", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");

        router.get("/post/", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "another_route");

        match router.recognize(&method::Get, "/post") {
            Ok(_) => {
                panic!();
            },
            Err(e) => {
                assert_eq!(RouterError::MethodNotAllowed, e);
            }
        }
    }

    #[test]
    fn test_handle_any_ok() {
        let mut router = Router::new();
        router.post("/post", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        router.any("/post", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        router.put("/post", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        router.any("/get", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "any");

        assert!(router.recognize(&method::Get, "/post").is_ok());
        assert!(router.recognize(&method::Get, "/get").is_ok());
    }

    #[test]
    fn test_request() {
        let mut router = Router::new();
        router.post("/post", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        router.get("/post", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");

        assert!(router.recognize(&method::Post, "/post").is_ok());
        assert!(router.recognize(&method::Get, "/post").is_ok());
        assert!(router.recognize(&method::Put, "/post").is_err());
        assert!(router.recognize(&method::Get, "/post/").is_err());
    }

    #[test]
    fn test_not_found() {
        let mut router = Router::new();
        router.put("/put", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "");
        match router.recognize(&method::Patch, "/patch") {
            Ok(_) => {
                panic!();
            },
            Err(e) => {
                assert_eq!(RouterError::NotFound, e);
            }
        }
    }

    #[test]
    #[should_panic]
    fn test_same_route_id() {
        let mut router = Router::new();
        router.put("/put", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "my_route_id");
        router.get("/get", |_: &mut Request| {
            Ok(Response::with((status::Ok, "")))
        }, "my_route_id");
    }

    #[test]
    fn test_wildcard_regression() {
        let mut router = Router::new();
        router.options("*", |_: &mut Request| Ok(Response::with((status::Ok, ""))), "id1");
        router.put("/upload/*filename", |_: &mut Request| Ok(Response::with((status::Ok, ""))), "id2");
        assert!(router.recognize(&method::Options, "/foo").is_ok());
        assert!(router.recognize(&method::Put, "/foo").is_err());
        assert!(router.recognize(&method::Put, "/upload/foo").is_ok());
    }
}
<|endoftext|>