<commit_before>use chrono::NaiveDateTime;
use diesel::*;

use models::Crate;
use schema::*;
use views::EncodableCategory;

#[derive(Clone, Identifiable, Queryable, QueryableByName, Debug)]
#[table_name = "categories"]
pub struct Category {
    pub id: i32,
    pub category: String,
    pub slug: String,
    pub description: String,
    pub crates_cnt: i32,
    pub created_at: NaiveDateTime,
}

#[derive(Associations, Insertable, Identifiable, Debug, Clone, Copy)]
#[belongs_to(Category)]
#[belongs_to(Crate)]
#[table_name = "crates_categories"]
#[primary_key(crate_id, category_id)]
pub struct CrateCategory {
    crate_id: i32,
    category_id: i32,
}

impl Category {
    pub fn encodable(self) -> EncodableCategory {
        let Category {
            crates_cnt,
            category,
            slug,
            description,
            created_at,
            ..
        } = self;
        EncodableCategory {
            id: slug.clone(),
            slug: slug.clone(),
            description: description.clone(),
            created_at,
            crates_cnt,
            category: category.rsplit("::").collect::<Vec<_>>()[0].to_string(),
        }
    }

    pub fn update_crate<'a>(
        conn: &PgConnection,
        krate: &Crate,
        slugs: &[&'a str],
    ) -> QueryResult<Vec<&'a str>> {
        use diesel::dsl::any;

        conn.transaction(|| {
            let categories = categories::table
                .filter(categories::slug.eq(any(slugs)))
                .load::<Category>(conn)?;
            let invalid_categories = slugs
                .iter()
                .cloned()
                .filter(|s| !categories.iter().any(|c| c.slug == *s))
                .collect();
            let crate_categories = categories
                .iter()
                .map(|c| CrateCategory {
                    category_id: c.id,
                    crate_id: krate.id,
                })
                .collect::<Vec<_>>();

            delete(CrateCategory::belonging_to(krate)).execute(conn)?;
            insert_into(crates_categories::table)
                .values(&crate_categories)
                .execute(conn)?;
            Ok(invalid_categories)
        })
    }

    pub fn count_toplevel(conn: &PgConnection) -> QueryResult<i64> {
        use self::categories::dsl::*;

        categories
            .filter(category.not_like("%::%"))
            .count()
            .get_result(conn)
    }

    pub fn toplevel(
        conn: &PgConnection,
        sort: &str,
        limit: i64,
        offset: i64,
    ) -> QueryResult<Vec<Category>> {
        use diesel::dsl::*;
        use diesel::select;

        let sort_sql = match sort {
            "crates" => "ORDER BY crates_cnt DESC",
            _ => "ORDER BY category ASC",
        };

        // Collect all the top-level categories and sum up the crates_cnt of
        // the crates in all subcategories
        select(sql::<categories::SqlType>(&format!(
            "c.id, c.category, c.slug, c.description,
                sum(c2.crates_cnt)::int as crates_cnt, c.created_at
             FROM categories as c
             INNER JOIN categories c2 ON split_part(c2.slug, '::', 1) = c.slug
             WHERE split_part(c.slug, '::', 1) = c.slug
             GROUP BY c.id
             {} LIMIT {} OFFSET {}",
            sort_sql, limit, offset
        ))).load(conn)
    }

    pub fn subcategories(&self, conn: &PgConnection) -> QueryResult<Vec<Category>> {
        use diesel::sql_types::Text;

        sql_query(
            "SELECT c.id, c.category, c.slug, c.description, \
             COALESCE (( \
             SELECT sum(c2.crates_cnt)::int \
             FROM categories as c2 \
             WHERE c2.slug = c.slug \
             OR c2.slug LIKE c.slug || '::%' \
             ), 0) as crates_cnt, c.created_at \
             FROM categories as c \
             WHERE c.category ILIKE $1 || '::%' \
             AND c.category NOT ILIKE $1 || '::%::%'",
        ).bind::<Text, _>(&self.category)
            .load(conn)
    }

    /// Gathers the parent categories from the top-level Category to the direct parent of this Category.
    /// Returns categories as a Vector in order of traversal, not including this Category.
    /// The intention is to be able to have slugs or parent categories arrayed in order, to
    /// offer the frontend, for examples, slugs to create links to each parent category in turn.
    pub fn parent_categories(&self, conn: &PgConnection) -> QueryResult<Vec<Category>> {
        use diesel::expression::dsl::*;
        use diesel::types::Text;

        sql::<categories::SqlType>(
            "SELECT c.id, c.category, c.slug, c.description, \
            COALESCE(( \
            SELECT sum(c2.crates_cnt)::int from categories c2 \
            WHERE path <@ subltree(c.path, 0, 2)), 0) as crates_cnt, c.created_at \
            FROM categories c \
            WHERE c.path @> (select path from categories where slug = $1) \
            AND c.slug <> $1
            ORDER BY c.path",
        ).bind::<Text, _>(&self.slug)
            .load(conn)
    }
}

#[derive(Insertable, AsChangeset, Default, Debug)]
#[table_name = "categories"]
pub struct NewCategory<'a> {
    pub category: &'a str,
    pub slug: &'a str,
}

impl<'a> NewCategory<'a> {
    /// Inserts the category into the database, or updates an existing one.
    pub fn create_or_update(&self, conn: &PgConnection) -> QueryResult<Category> {
        use schema::categories::dsl::*;

        insert_into(categories)
            .values(self)
            .on_conflict(slug)
            .do_update()
            .set(self)
            .get_result(conn)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use diesel::connection::SimpleConnection;
    use dotenv::dotenv;
    use std::env;

    fn pg_connection() -> PgConnection {
        let _ = dotenv();
        let database_url =
            env::var("TEST_DATABASE_URL").expect("TEST_DATABASE_URL must be set to run tests");
        let conn = PgConnection::establish(&database_url).unwrap();
        // These tests deadlock if run concurrently
        conn.batch_execute("BEGIN; LOCK categories IN ACCESS EXCLUSIVE MODE")
            .unwrap();
        conn
    }

    #[test]
    fn category_toplevel_excludes_subcategories() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug) VALUES
            ('Cat 2', 'cat2'), ('Cat 1', 'cat1'), ('Cat 1::sub', 'cat1::sub')
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "", 10, 0)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec!["Cat 1".to_string(), "Cat 2".to_string()];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_toplevel_orders_by_crates_cnt_when_sort_given() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug, crates_cnt) VALUES
            ('Cat 1', 'cat1', 0), ('Cat 2', 'cat2', 2), ('Cat 3', 'cat3', 1)
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "crates", 10, 0)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec![
            "Cat 2".to_string(),
            "Cat 3".to_string(),
            "Cat 1".to_string(),
        ];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_toplevel_applies_limit_and_offset() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug) VALUES
            ('Cat 1', 'cat1'), ('Cat 2', 'cat2')
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "", 1, 0)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec!["Cat 1".to_string()];
        assert_eq!(expected, categories);

        let categories = Category::toplevel(&conn, "", 1, 1)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec!["Cat 2".to_string()];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_toplevel_includes_subcategories_in_crate_cnt() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug, crates_cnt) VALUES
            ('Cat 1', 'cat1', 1), ('Cat 1::sub', 'cat1::sub', 2),
            ('Cat 2', 'cat2', 3), ('Cat 2::Sub 1', 'cat2::sub1', 4),
            ('Cat 2::Sub 2', 'cat2::sub2', 5), ('Cat 3', 'cat3', 6)
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "crates", 10, 0)
            .unwrap()
            .into_iter()
            .map(|c| (c.category, c.crates_cnt))
            .collect::<Vec<_>>();
        let expected = vec![
            ("Cat 2".to_string(), 12),
            ("Cat 3".to_string(), 6),
            ("Cat 1".to_string(), 3),
        ];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_toplevel_applies_limit_and_offset_after_grouping() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug, crates_cnt) VALUES
            ('Cat 1', 'cat1', 1), ('Cat 1::sub', 'cat1::sub', 2),
            ('Cat 2', 'cat2', 3), ('Cat 2::Sub 1', 'cat2::sub1', 4),
            ('Cat 2::Sub 2', 'cat2::sub2', 5), ('Cat 3', 'cat3', 6)
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "crates", 2, 0)
            .unwrap()
            .into_iter()
            .map(|c| (c.category, c.crates_cnt))
            .collect::<Vec<_>>();
        let expected = vec![("Cat 2".to_string(), 12), ("Cat 3".to_string(), 6)];
        assert_eq!(expected, categories);

        let categories = Category::toplevel(&conn, "crates", 2, 1)
            .unwrap()
            .into_iter()
            .map(|c| (c.category, c.crates_cnt))
            .collect::<Vec<_>>();
        let expected = vec![("Cat 3".to_string(), 6), ("Cat 1".to_string(), 3)];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_parent_categories_includes_path_to_node_with_count() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug, crates_cnt) VALUES
            ('Cat 1', 'cat1', 1), ('Cat 1::sub1', 'cat1::sub1', 2),
            ('Cat 1::sub2', 'cat1::sub2', 2), ('Cat 1::sub1::subsub1', 'cat1::sub1::subsub1', 2),
            ('Cat 2', 'cat2', 3), ('Cat 2::Sub 1', 'cat2::sub1', 4),
            ('Cat 2::Sub 2', 'cat2::sub2', 5), ('Cat 3', 'cat3', 200)
            ",
        ).unwrap();

        let cat = categories::table
            .filter(categories::slug.eq("cat1::sub1"))
            .first::<Category>(&conn)
            .unwrap();
        let subcats = cat.subcategories(&conn).unwrap();
        let parents = cat.parent_categories(&conn).unwrap();

        assert_eq!(parents.len(), 1);
        assert_eq!(parents[0].slug, "cat1");
        assert_eq!(parents[0].crates_cnt, 7);
        assert_eq!(subcats.len(), 1);
        assert_eq!(subcats[0].slug, "cat1::sub1::subsub1");
    }
}
<commit_msg>use insert_into(…).values<commit_after>use chrono::NaiveDateTime;
use diesel::*;

use models::Crate;
use schema::*;
use views::EncodableCategory;

#[derive(Clone, Identifiable, Queryable, QueryableByName, Debug)]
#[table_name = "categories"]
pub struct Category {
    pub id: i32,
    pub category: String,
    pub slug: String,
    pub description: String,
    pub crates_cnt: i32,
    pub created_at: NaiveDateTime,
}

#[derive(Associations, Insertable, Identifiable, Debug, Clone, Copy)]
#[belongs_to(Category)]
#[belongs_to(Crate)]
#[table_name = "crates_categories"]
#[primary_key(crate_id, category_id)]
pub struct CrateCategory {
    crate_id: i32,
    category_id: i32,
}

impl Category {
    pub fn encodable(self) -> EncodableCategory {
        let Category {
            crates_cnt,
            category,
            slug,
            description,
            created_at,
            ..
        } = self;
        EncodableCategory {
            id: slug.clone(),
            slug: slug.clone(),
            description: description.clone(),
            created_at,
            crates_cnt,
            category: category.rsplit("::").collect::<Vec<_>>()[0].to_string(),
        }
    }

    pub fn update_crate<'a>(
        conn: &PgConnection,
        krate: &Crate,
        slugs: &[&'a str],
    ) -> QueryResult<Vec<&'a str>> {
        use diesel::dsl::any;

        conn.transaction(|| {
            let categories = categories::table
                .filter(categories::slug.eq(any(slugs)))
                .load::<Category>(conn)?;
            let invalid_categories = slugs
                .iter()
                .cloned()
                .filter(|s| !categories.iter().any(|c| c.slug == *s))
                .collect();
            let crate_categories = categories
                .iter()
                .map(|c| CrateCategory {
                    category_id: c.id,
                    crate_id: krate.id,
                })
                .collect::<Vec<_>>();

            delete(CrateCategory::belonging_to(krate)).execute(conn)?;
            insert_into(crates_categories::table)
                .values(&crate_categories)
                .execute(conn)?;
            Ok(invalid_categories)
        })
    }

    pub fn count_toplevel(conn: &PgConnection) -> QueryResult<i64> {
        use self::categories::dsl::*;

        categories
            .filter(category.not_like("%::%"))
            .count()
            .get_result(conn)
    }

    pub fn toplevel(
        conn: &PgConnection,
        sort: &str,
        limit: i64,
        offset: i64,
    ) -> QueryResult<Vec<Category>> {
        use diesel::dsl::*;
        use diesel::select;

        let sort_sql = match sort {
            "crates" => "ORDER BY crates_cnt DESC",
            _ => "ORDER BY category ASC",
        };

        // Collect all the top-level categories and sum up the crates_cnt of
        // the crates in all subcategories
        select(sql::<categories::SqlType>(&format!(
            "c.id, c.category, c.slug, c.description,
                sum(c2.crates_cnt)::int as crates_cnt, c.created_at
             FROM categories as c
             INNER JOIN categories c2 ON split_part(c2.slug, '::', 1) = c.slug
             WHERE split_part(c.slug, '::', 1) = c.slug
             GROUP BY c.id
             {} LIMIT {} OFFSET {}",
            sort_sql, limit, offset
        ))).load(conn)
    }

    pub fn subcategories(&self, conn: &PgConnection) -> QueryResult<Vec<Category>> {
        use diesel::sql_types::Text;

        sql_query(
            "SELECT c.id, c.category, c.slug, c.description, \
             COALESCE (( \
             SELECT sum(c2.crates_cnt)::int \
             FROM categories as c2 \
             WHERE c2.slug = c.slug \
             OR c2.slug LIKE c.slug || '::%' \
             ), 0) as crates_cnt, c.created_at \
             FROM categories as c \
             WHERE c.category ILIKE $1 || '::%' \
             AND c.category NOT ILIKE $1 || '::%::%'",
        ).bind::<Text, _>(&self.category)
            .load(conn)
    }

    /// Gathers the parent categories from the top-level Category to the direct parent of this Category.
    /// Returns categories as a Vector in order of traversal, not including this Category.
    /// The intention is to be able to have slugs or parent categories arrayed in order, to
    /// offer the frontend, for examples, slugs to create links to each parent category in turn.
    pub fn parent_categories(&self, conn: &PgConnection) -> QueryResult<Vec<Category>> {
        use diesel::expression::dsl::*;
        use diesel::types::Text;

        sql::<categories::SqlType>(
            "SELECT c.id, c.category, c.slug, c.description, \
            COALESCE(( \
            SELECT sum(c2.crates_cnt)::int from categories c2 \
            WHERE path <@ subltree(c.path, 0, 2)), 0) as crates_cnt, c.created_at \
            FROM categories c \
            WHERE c.path @> (select path from categories where slug = $1) \
            AND c.slug <> $1
            ORDER BY c.path",
        ).bind::<Text, _>(&self.slug)
            .load(conn)
    }
}

#[derive(Insertable, AsChangeset, Default, Debug)]
#[table_name = "categories"]
pub struct NewCategory<'a> {
    pub category: &'a str,
    pub slug: &'a str,
}

impl<'a> NewCategory<'a> {
    /// Inserts the category into the database, or updates an existing one.
    pub fn create_or_update(&self, conn: &PgConnection) -> QueryResult<Category> {
        use schema::categories::dsl::*;

        insert_into(categories)
            .values(self)
            .on_conflict(slug)
            .do_update()
            .set(self)
            .get_result(conn)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use diesel::connection::SimpleConnection;
    use dotenv::dotenv;
    use std::env;

    fn pg_connection() -> PgConnection {
        let _ = dotenv();
        let database_url =
            env::var("TEST_DATABASE_URL").expect("TEST_DATABASE_URL must be set to run tests");
        let conn = PgConnection::establish(&database_url).unwrap();
        // These tests deadlock if run concurrently
        conn.batch_execute("BEGIN; LOCK categories IN ACCESS EXCLUSIVE MODE")
            .unwrap();
        conn
    }

    #[test]
    fn category_toplevel_excludes_subcategories() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug) VALUES
            ('Cat 2', 'cat2'), ('Cat 1', 'cat1'), ('Cat 1::sub', 'cat1::sub')
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "", 10, 0)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec!["Cat 1".to_string(), "Cat 2".to_string()];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_toplevel_orders_by_crates_cnt_when_sort_given() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug, crates_cnt) VALUES
            ('Cat 1', 'cat1', 0), ('Cat 2', 'cat2', 2), ('Cat 3', 'cat3', 1)
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "crates", 10, 0)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec![
            "Cat 2".to_string(),
            "Cat 3".to_string(),
            "Cat 1".to_string(),
        ];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_toplevel_applies_limit_and_offset() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug) VALUES
            ('Cat 1', 'cat1'), ('Cat 2', 'cat2')
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "", 1, 0)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec!["Cat 1".to_string()];
        assert_eq!(expected, categories);

        let categories = Category::toplevel(&conn, "", 1, 1)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec!["Cat 2".to_string()];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_toplevel_includes_subcategories_in_crate_cnt() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug, crates_cnt) VALUES
            ('Cat 1', 'cat1', 1), ('Cat 1::sub', 'cat1::sub', 2),
            ('Cat 2', 'cat2', 3), ('Cat 2::Sub 1', 'cat2::sub1', 4),
            ('Cat 2::Sub 2', 'cat2::sub2', 5), ('Cat 3', 'cat3', 6)
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "crates", 10, 0)
            .unwrap()
            .into_iter()
            .map(|c| (c.category, c.crates_cnt))
            .collect::<Vec<_>>();
        let expected = vec![
            ("Cat 2".to_string(), 12),
            ("Cat 3".to_string(), 6),
            ("Cat 1".to_string(), 3),
        ];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_toplevel_applies_limit_and_offset_after_grouping() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug, crates_cnt) VALUES
            ('Cat 1', 'cat1', 1), ('Cat 1::sub', 'cat1::sub', 2),
            ('Cat 2', 'cat2', 3), ('Cat 2::Sub 1', 'cat2::sub1', 4),
            ('Cat 2::Sub 2', 'cat2::sub2', 5), ('Cat 3', 'cat3', 6)
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "crates", 2, 0)
            .unwrap()
            .into_iter()
            .map(|c| (c.category, c.crates_cnt))
            .collect::<Vec<_>>();
        let expected = vec![("Cat 2".to_string(), 12), ("Cat 3".to_string(), 6)];
        assert_eq!(expected, categories);

        let categories = Category::toplevel(&conn, "crates", 2, 1)
            .unwrap()
            .into_iter()
            .map(|c| (c.category, c.crates_cnt))
            .collect::<Vec<_>>();
        let expected = vec![("Cat 3".to_string(), 6), ("Cat 1".to_string(), 3)];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_parent_categories_includes_path_to_node_with_count() {
        use self::categories::dsl::*;
        let conn = pg_connection();
        insert_into(categories)
            .values(&vec![
                (category.eq("Cat 1"), slug.eq("cat1"), crates_cnt.eq(1)),
                (category.eq("Cat 1::sub1"), slug.eq("cat1::sub1"), crates_cnt.eq(2)),
                (category.eq("Cat 1::sub2"), slug.eq("cat1::sub2"), crates_cnt.eq(2)),
                (category.eq("Cat 1::sub1::subsub1"), slug.eq("cat1::sub1::subsub1"), crates_cnt.eq(2)),
                (category.eq("Cat 2"), slug.eq("cat2"), crates_cnt.eq(3)),
                (category.eq("Cat 2::Sub 1"), slug.eq("cat2::sub1"), crates_cnt.eq(4)),
                (category.eq("Cat 2::Sub 2"), slug.eq("cat2::sub2"), crates_cnt.eq(5)),
                (category.eq("Cat 3"), slug.eq("cat3"), crates_cnt.eq(200))
            ])
            .execute(&conn)
            .unwrap();

        let cat = categories
            .filter(slug.eq("cat1::sub1"))
            .first::<Category>(&conn)
            .unwrap();
        let subcats = cat.subcategories(&conn).unwrap();
        let parents = cat.parent_categories(&conn).unwrap();

        assert_eq!(parents.len(), 1);
        assert_eq!(parents[0].slug, "cat1");
        assert_eq!(parents[0].crates_cnt, 7);
        assert_eq!(subcats.len(), 1);
        assert_eq!(subcats[0].slug, "cat1::sub1::subsub1");
    }
}
<|endoftext|>


<commit_before>// The MIT License (MIT)

// Copyright (c) 2015 Y. T. Chung <zonyitoo@gmail.com>

//  Permission is hereby granted, free of charge, to any person obtaining a
//  copy of this software and associated documentation files (the "Software"),
//  to deal in the Software without restriction, including without limitation
//  the rights to use, copy, modify, merge, publish, distribute, sublicense,
//  and/or sell copies of the Software, and to permit persons to whom the
//  Software is furnished to do so, subject to the following conditions:
//
//  The above copyright notice and this permission notice shall be included in
//  all copies or substantial portions of the Software.
//
//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
//  OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
//  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
//  DEALINGS IN THE SOFTWARE.

use std::net::{ToSocketAddrs, SocketAddr};
use std::convert::From;
use std::io::{self, Write, BufWriter};

use hyper;
use hyper::http;
use hyper::buffer::BufReader;
use hyper::server::{Request, Response, Handler};
use hyper::header::{Connection, Headers, Expect};
use hyper::version::HttpVersion;
use hyper::net::{NetworkListener, NetworkStream};
use hyper::status::StatusCode;
use hyper::error::Error;

use net::http::conn::{HttpListener, HttpsListener, Ssl};

use scheduler::Scheduler;

/// A server can listen on a TCP socket.
///
/// Once listening, it will create a `Request`/`Response` pair for each
/// incoming connection, and hand them to the provided handler.
#[derive(Debug)]
pub struct Server<L = HttpListener> {
    listener: L,
}

impl<L: NetworkListener> Server<L> {
    /// Creates a new server with the provided handler.
    #[inline]
    pub fn new(listener: L) -> Server<L> {
        Server {
            listener: listener,
        }
    }
}

impl Server<HttpListener> {
    /// Creates a new server that will handle `HttpStream`s.
    pub fn http<To: ToSocketAddrs>(addr: To) -> hyper::Result<Server<HttpListener>> {
        HttpListener::new(addr).map(Server::new).map_err(From::from)
    }
}

impl<S: Ssl + Clone + Send> Server<HttpsListener<S>> {
    /// Creates a new server that will handle `HttpStream`s over SSL.
    ///
    /// You can use any SSL implementation, as long as implements `hyper::net::Ssl`.
    pub fn https<A: ToSocketAddrs>(addr: A, ssl: S) -> hyper::Result<Server<HttpsListener<S>>> {
        HttpsListener::new(addr, ssl).map(Server::new)
    }
}

impl<L: NetworkListener + Send + 'static> Server<L> {
    /// Binds to a socket.
    pub fn listen<H: Handler + 'static>(mut self, handler: H) -> hyper::Result<SocketAddr> {
        let socket = try!(self.listener.local_addr());

        Scheduler::spawn(move|| {
            use std::sync::Arc;
            use std::ops::Deref;

            let handler = Arc::new(handler);
            loop {
                let mut stream = self.listener.accept().unwrap();

                let handler = handler.clone();
                Scheduler::spawn(move|| Worker(handler.deref()).handle_connection(&mut stream));
            }
        });

        Ok(socket)
    }
}

struct Worker<'a, H: Handler + 'static>(&'a H);

impl<'a, H: Handler + 'static> Worker<'a, H> {

    fn handle_connection<S>(&self, mut stream: &mut S) where S: NetworkStream + Clone {
        debug!("Incoming stream");
        let addr = match stream.peer_addr() {
            Ok(addr) => addr,
            Err(e) => {
                error!("Peer Name error: {:?}", e);
                return;
            }
        };

        // FIXME: Use Type ascription
        let stream_clone: &mut NetworkStream = &mut stream.clone();
        let rdr = BufReader::new(stream_clone);
        let wrt = BufWriter::new(stream);

        self.keep_alive_loop(rdr, wrt, addr);
        debug!("keep_alive loop ending for {}", addr);
    }

    fn keep_alive_loop<W: Write>(&self, mut rdr: BufReader<&mut NetworkStream>, mut wrt: W, addr: SocketAddr) {
        let mut keep_alive = true;
        while keep_alive {
            let req = match Request::new(&mut rdr, addr) {
                Ok(req) => req,
                Err(Error::Io(ref e)) if e.kind() == io::ErrorKind::ConnectionAborted => {
                    trace!("tcp closed, cancelling keep-alive loop");
                    break;
                }
                Err(Error::Io(e)) => {
                    debug!("ioerror in keepalive loop = {:?}", e);
                    break;
                }
                Err(e) => {
                    //TODO: send a 400 response
                    error!("request error = {:?}", e);
                    break;
                }
            };

            if !self.handle_expect(&req, &mut wrt) {
                break;
            }

            keep_alive = http::should_keep_alive(req.version, &req.headers);
            let version = req.version;
            let mut res_headers = Headers::new();
            if !keep_alive {
                res_headers.set(Connection::close());
            }
            {
                let mut res = Response::new(&mut wrt, &mut res_headers);
                res.version = version;
                self.0.handle(req, res);
            }

            // if the request was keep-alive, we need to check that the server agrees
            // if it wasn't, then the server cannot force it to be true anyways
            if keep_alive {
                keep_alive = http::should_keep_alive(version, &res_headers);
            }

            debug!("keep_alive = {:?} for {}", keep_alive, addr);
        }

    }

    fn handle_expect<W: Write>(&self, req: &Request, wrt: &mut W) -> bool {
         if req.version == HttpVersion::Http11 && req.headers.get() == Some(&Expect::Continue) {
            let status = self.0.check_continue((&req.method, &req.uri, &req.headers));
            match write!(wrt, "{} {}\r\n\r\n", HttpVersion::Http11, status) {
                Ok(..) => (),
                Err(e) => {
                    error!("error writing 100-continue: {:?}", e);
                    return false;
                }
            }

            if status != StatusCode::Continue {
                debug!("non-100 status ({}) for Expect 100 request", status);
                return false;
            }
        }

        true
    }
}

macro_rules! try_return(
    ($e:expr) => {{
        match $e {
            Ok(v) => v,
            Err(e) => { println!("Error: {}", e); return; }
        }
    }}
);
<commit_msg>does not need to use deref<commit_after>// The MIT License (MIT)

// Copyright (c) 2015 Y. T. Chung <zonyitoo@gmail.com>

//  Permission is hereby granted, free of charge, to any person obtaining a
//  copy of this software and associated documentation files (the "Software"),
//  to deal in the Software without restriction, including without limitation
//  the rights to use, copy, modify, merge, publish, distribute, sublicense,
//  and/or sell copies of the Software, and to permit persons to whom the
//  Software is furnished to do so, subject to the following conditions:
//
//  The above copyright notice and this permission notice shall be included in
//  all copies or substantial portions of the Software.
//
//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
//  OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
//  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
//  DEALINGS IN THE SOFTWARE.

use std::net::{ToSocketAddrs, SocketAddr};
use std::convert::From;
use std::io::{self, Write, BufWriter};

use hyper;
use hyper::http;
use hyper::buffer::BufReader;
use hyper::server::{Request, Response, Handler};
use hyper::header::{Connection, Headers, Expect};
use hyper::version::HttpVersion;
use hyper::net::{NetworkListener, NetworkStream};
use hyper::status::StatusCode;
use hyper::error::Error;

use net::http::conn::{HttpListener, HttpsListener, Ssl};

use scheduler::Scheduler;

/// A server can listen on a TCP socket.
///
/// Once listening, it will create a `Request`/`Response` pair for each
/// incoming connection, and hand them to the provided handler.
#[derive(Debug)]
pub struct Server<L = HttpListener> {
    listener: L,
}

impl<L: NetworkListener> Server<L> {
    /// Creates a new server with the provided handler.
    #[inline]
    pub fn new(listener: L) -> Server<L> {
        Server {
            listener: listener,
        }
    }
}

impl Server<HttpListener> {
    /// Creates a new server that will handle `HttpStream`s.
    pub fn http<To: ToSocketAddrs>(addr: To) -> hyper::Result<Server<HttpListener>> {
        HttpListener::new(addr).map(Server::new).map_err(From::from)
    }
}

impl<S: Ssl + Clone + Send> Server<HttpsListener<S>> {
    /// Creates a new server that will handle `HttpStream`s over SSL.
    ///
    /// You can use any SSL implementation, as long as implements `hyper::net::Ssl`.
    pub fn https<A: ToSocketAddrs>(addr: A, ssl: S) -> hyper::Result<Server<HttpsListener<S>>> {
        HttpsListener::new(addr, ssl).map(Server::new)
    }
}

impl<L: NetworkListener + Send + 'static> Server<L> {
    /// Binds to a socket.
    pub fn listen<H: Handler + 'static>(mut self, handler: H) -> hyper::Result<SocketAddr> {
        let socket = try!(self.listener.local_addr());

        Scheduler::spawn(move|| {
            use std::sync::Arc;

            let handler = Arc::new(handler);
            loop {
                let mut stream = self.listener.accept().unwrap();

                let handler = handler.clone();
                Scheduler::spawn(move|| Worker(&*handler).handle_connection(&mut stream));
            }
        });

        Ok(socket)
    }
}

struct Worker<'a, H: Handler + 'static>(&'a H);

impl<'a, H: Handler + 'static> Worker<'a, H> {

    fn handle_connection<S>(&self, mut stream: &mut S) where S: NetworkStream + Clone {
        debug!("Incoming stream");
        let addr = match stream.peer_addr() {
            Ok(addr) => addr,
            Err(e) => {
                error!("Peer Name error: {:?}", e);
                return;
            }
        };

        // FIXME: Use Type ascription
        let stream_clone: &mut NetworkStream = &mut stream.clone();
        let rdr = BufReader::new(stream_clone);
        let wrt = BufWriter::new(stream);

        self.keep_alive_loop(rdr, wrt, addr);
        debug!("keep_alive loop ending for {}", addr);
    }

    fn keep_alive_loop<W: Write>(&self, mut rdr: BufReader<&mut NetworkStream>, mut wrt: W, addr: SocketAddr) {
        let mut keep_alive = true;
        while keep_alive {
            let req = match Request::new(&mut rdr, addr) {
                Ok(req) => req,
                Err(Error::Io(ref e)) if e.kind() == io::ErrorKind::ConnectionAborted => {
                    trace!("tcp closed, cancelling keep-alive loop");
                    break;
                }
                Err(Error::Io(e)) => {
                    debug!("ioerror in keepalive loop = {:?}", e);
                    break;
                }
                Err(e) => {
                    //TODO: send a 400 response
                    error!("request error = {:?}", e);
                    break;
                }
            };

            if !self.handle_expect(&req, &mut wrt) {
                break;
            }

            keep_alive = http::should_keep_alive(req.version, &req.headers);
            let version = req.version;
            let mut res_headers = Headers::new();
            if !keep_alive {
                res_headers.set(Connection::close());
            }
            {
                let mut res = Response::new(&mut wrt, &mut res_headers);
                res.version = version;
                self.0.handle(req, res);
            }

            // if the request was keep-alive, we need to check that the server agrees
            // if it wasn't, then the server cannot force it to be true anyways
            if keep_alive {
                keep_alive = http::should_keep_alive(version, &res_headers);
            }

            debug!("keep_alive = {:?} for {}", keep_alive, addr);
        }

    }

    fn handle_expect<W: Write>(&self, req: &Request, wrt: &mut W) -> bool {
         if req.version == HttpVersion::Http11 && req.headers.get() == Some(&Expect::Continue) {
            let status = self.0.check_continue((&req.method, &req.uri, &req.headers));
            match write!(wrt, "{} {}\r\n\r\n", HttpVersion::Http11, status) {
                Ok(..) => (),
                Err(e) => {
                    error!("error writing 100-continue: {:?}", e);
                    return false;
                }
            }

            if status != StatusCode::Continue {
                debug!("non-100 status ({}) for Expect 100 request", status);
                return false;
            }
        }

        true
    }
}

macro_rules! try_return(
    ($e:expr) => {{
        match $e {
            Ok(v) => v,
            Err(e) => { println!("Error: {}", e); return; }
        }
    }}
);
<|endoftext|>


<commit_before>use log::{debug};
use goblin::{Object};
use failure::{Error};
use byteorder::{ByteOrder, LittleEndian};

use super::super::super::arch::{RVA};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

pub struct ImportsAnalyzer {}

impl ImportsAnalyzer {
    pub fn new() -> ImportsAnalyzer {
        ImportsAnalyzer {}
    }
}

struct ImageImportDescriptor {
    original_first_thunk: RVA,
    time_date_stamp: u32,
    forwarder_chain: u32,
    name: RVA,
    first_thunk: RVA,
}

impl ImageImportDescriptor {
    fn is_empty(&self) -> bool {
        self.original_first_thunk == RVA(0x0) &&
            self.time_date_stamp == 0x0 &&
            self.forwarder_chain == 0x0 &&
            self.name == RVA(0x0) &&
            self.first_thunk == RVA(0x0)
    }
}

impl std::fmt::Debug for ImageImportDescriptor {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "IMAGE_IMPORT_DESCRIPTOR(FT: {:#x} OFT: {:#x} name: {:#x})",
               self.first_thunk,
               self.original_first_thunk,
               self.name
        )
    }
}

fn read_image_import_descriptor(ws: &Workspace, rva: RVA) -> Result<ImageImportDescriptor, Error> {
    let buf = ws.read_bytes(rva, 5 * 4)?;

    // these fields are all u32, even on 64-bit
    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    Ok(ImageImportDescriptor {
        original_first_thunk: RVA::from(entries[0] as i64),
        time_date_stamp: entries[1],
        forwarder_chain: entries[2],
        name: RVA::from(entries[3] as i64),
        first_thunk: RVA::from(entries[4] as i64),
    })
}

enum ImageThunkData {
    Function(RVA),
    Ordinal(u32),
}

fn read_image_thunk_data(ws: &Workspace, rva: RVA) -> Result<ImageThunkData, Error> {
    // see: https://reverseengineering.stackexchange.com/a/13387/17194
    let thunk = ws.read_rva(rva)?;
    let v: i64 = thunk.into();
    if v & (1 << (ws.loader.get_arch().get_pointer_size() * 8 - 1)) > 0x0 {
        // MSB is set, this is an ordinal
        Ok(ImageThunkData::Ordinal((v & 0xFFFF) as u32))
    } else {
        Ok(ImageThunkData::Function(thunk))
    }
}

struct ImageImportByName {
    hint: u16,
    name: String,
}

impl std::fmt::Debug for ImageImportByName {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "IMAGE_IMPORT_BY_NAME(hint: {:#x} name: {})",
               self.hint,
               self.name
        )
    }
}

fn read_image_import_by_name(ws: &Workspace, rva: RVA) -> Result<ImageImportByName, Error> {
    Ok(ImageImportByName {
        hint: ws.read_u16(rva)?,
        name: ws.read_utf8(rva + RVA::from(2))?,
    })
}

impl Analyzer for ImportsAnalyzer {
    fn get_name(&self) -> String {
        "PE imports analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::ImportsAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// let anal = ImportsAnalyzer::new();
    /// anal.analyze(&mut ws).unwrap();
    /// assert_eq!(ws.get_symbol(0x77448).unwrap(), "api-ms-win-core-appcompat-l1-1-1.dll!BaseReadAppCompatDataForProcess");
    /// ```
    fn analyze(&self, ws: &mut Workspace) -> Result<(), Error> {
        let import_directory = {
            let pe = match Object::parse(&ws.buf) {
                Ok(Object::PE(pe)) => pe,
                _ => panic!("can't analyze unexpected format"),
            };

            let opt_header = match pe.header.optional_header {
                Some(opt_header) => opt_header,
                _ => return Ok(()),
            };

            let import_directory = match opt_header.data_directories.get_import_table() {
                Some(import_directory) => import_directory,
                _ => return Ok(()),
            };

            RVA::from(import_directory.virtual_address as i64)
        };

        debug!("import directory: {:#x}", import_directory);

        let mut symbols: Vec<(RVA, String)> = vec![];

        let psize: usize = ws.loader.get_arch().get_pointer_size() as usize;
        for i in 0..std::usize::MAX {
            //
            //  0x0                    0x14
            //  +-------------------------+           0x0          ptrsize
            //  | IMAGE_IMPORT_DESCRIPTOR | --------> +------------------+
            //  +-------------------------+ \         | IMAGE_THUNK_DATA | -------> +----------------------+
            //  | ...                     |  |        +------------------+          | IMAGE_IMPORT_BY_NAME |
            //  +-------------------------+  |        | ...              |          +----------------------+
            //  | 00 00 00 00 00 00 000   |  |        +------------------+          | hint u16             |
            //  +-------------------------+  |        | 00 00 00 00 00   |          | name ascii           |
            //                               |        +------------------+          +----------------------+
            //                               |
            //                               +> dll-name (ascii)
            let import_descriptor_rva = import_directory + RVA::from(i * 0x14);
            let import_descriptor = read_image_import_descriptor(ws, import_descriptor_rva)?;
            if import_descriptor.is_empty() {
                break;
            }

            let dll_name = ws.read_utf8(import_descriptor.name)?;
            debug!("{:?} -> {}", import_descriptor, dll_name);

            for j in 0..std::usize::MAX {
                // the First Thunk (FT) is the pointer that will be overwritten upon load.
                // entries here may not point to the IMAGE_IMPORT_BY_NAME.
                let first_thunk = import_descriptor.first_thunk + RVA::from(j * psize);
                // the Original First Thunk (OFT) remains constant, and points to the IMAGE_IMPORT_BY_NAME.
                // FT and OFT are parallel arrays.
                let image_thunk_data_rva = import_descriptor.original_first_thunk + RVA::from(j * psize);
                match read_image_thunk_data(ws, image_thunk_data_rva)? {
                    ImageThunkData::Function(rva) => {
                        if rva == RVA(0x0) {
                            break;
                        } else {
                            let imp = read_image_import_by_name(ws, rva)?;
                            debug!("{:?}", imp);

                            symbols.push((first_thunk, format!("{}!{}", dll_name, imp.name)))
                        }
                    },
                    ImageThunkData::Ordinal(ord) => {
                        symbols.push((first_thunk, format!("{}!#{}", dll_name, ord)))
                    },
                };
            }
        }

        for (rva, name) in symbols.iter() {
            debug!("import: {:#x} -> {}", rva, name);
            ws.make_symbol(*rva, name)?;
            ws.analyze()?;
        }

        Ok(())
    }
}
<commit_msg>analysis/imports: tests working again<commit_after>use log::{debug};
use goblin::{Object};
use failure::{Error};
use byteorder::{ByteOrder, LittleEndian};

use super::super::super::arch::{RVA};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

pub struct ImportsAnalyzer {}

impl ImportsAnalyzer {
    pub fn new() -> ImportsAnalyzer {
        ImportsAnalyzer {}
    }
}

struct ImageImportDescriptor {
    original_first_thunk: RVA,
    time_date_stamp: u32,
    forwarder_chain: u32,
    name: RVA,
    first_thunk: RVA,
}

impl ImageImportDescriptor {
    fn is_empty(&self) -> bool {
        self.original_first_thunk == RVA(0x0) &&
            self.time_date_stamp == 0x0 &&
            self.forwarder_chain == 0x0 &&
            self.name == RVA(0x0) &&
            self.first_thunk == RVA(0x0)
    }
}

impl std::fmt::Debug for ImageImportDescriptor {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "IMAGE_IMPORT_DESCRIPTOR(FT: {:#x} OFT: {:#x} name: {:#x})",
               self.first_thunk,
               self.original_first_thunk,
               self.name
        )
    }
}

fn read_image_import_descriptor(ws: &Workspace, rva: RVA) -> Result<ImageImportDescriptor, Error> {
    let buf = ws.read_bytes(rva, 5 * 4)?;

    // these fields are all u32, even on 64-bit
    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    Ok(ImageImportDescriptor {
        original_first_thunk: RVA::from(entries[0] as i64),
        time_date_stamp: entries[1],
        forwarder_chain: entries[2],
        name: RVA::from(entries[3] as i64),
        first_thunk: RVA::from(entries[4] as i64),
    })
}

enum ImageThunkData {
    Function(RVA),
    Ordinal(u32),
}

fn read_image_thunk_data(ws: &Workspace, rva: RVA) -> Result<ImageThunkData, Error> {
    // see: https://reverseengineering.stackexchange.com/a/13387/17194
    let thunk = ws.read_rva(rva)?;
    let v: i64 = thunk.into();
    if v & (1 << (ws.loader.get_arch().get_pointer_size() * 8 - 1)) > 0x0 {
        // MSB is set, this is an ordinal
        Ok(ImageThunkData::Ordinal((v & 0xFFFF) as u32))
    } else {
        Ok(ImageThunkData::Function(thunk))
    }
}

struct ImageImportByName {
    hint: u16,
    name: String,
}

impl std::fmt::Debug for ImageImportByName {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "IMAGE_IMPORT_BY_NAME(hint: {:#x} name: {})",
               self.hint,
               self.name
        )
    }
}

fn read_image_import_by_name(ws: &Workspace, rva: RVA) -> Result<ImageImportByName, Error> {
    Ok(ImageImportByName {
        hint: ws.read_u16(rva)?,
        name: ws.read_utf8(rva + RVA::from(2))?,
    })
}

impl Analyzer for ImportsAnalyzer {
    fn get_name(&self) -> String {
        "PE imports analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::ImportsAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// let anal = ImportsAnalyzer::new();
    /// anal.analyze(&mut ws).unwrap();
    /// assert_eq!(ws.get_symbol(RVA(0x77448)).unwrap(), "api-ms-win-core-appcompat-l1-1-1.dll!BaseReadAppCompatDataForProcess");
    /// ```
    fn analyze(&self, ws: &mut Workspace) -> Result<(), Error> {
        let import_directory = {
            let pe = match Object::parse(&ws.buf) {
                Ok(Object::PE(pe)) => pe,
                _ => panic!("can't analyze unexpected format"),
            };

            let opt_header = match pe.header.optional_header {
                Some(opt_header) => opt_header,
                _ => return Ok(()),
            };

            let import_directory = match opt_header.data_directories.get_import_table() {
                Some(import_directory) => import_directory,
                _ => return Ok(()),
            };

            RVA::from(import_directory.virtual_address as i64)
        };

        debug!("import directory: {:#x}", import_directory);

        let mut symbols: Vec<(RVA, String)> = vec![];

        let psize: usize = ws.loader.get_arch().get_pointer_size() as usize;
        for i in 0..std::usize::MAX {
            //
            //  0x0                    0x14
            //  +-------------------------+           0x0          ptrsize
            //  | IMAGE_IMPORT_DESCRIPTOR | --------> +------------------+
            //  +-------------------------+ \         | IMAGE_THUNK_DATA | -------> +----------------------+
            //  | ...                     |  |        +------------------+          | IMAGE_IMPORT_BY_NAME |
            //  +-------------------------+  |        | ...              |          +----------------------+
            //  | 00 00 00 00 00 00 000   |  |        +------------------+          | hint u16             |
            //  +-------------------------+  |        | 00 00 00 00 00   |          | name ascii           |
            //                               |        +------------------+          +----------------------+
            //                               |
            //                               +> dll-name (ascii)
            let import_descriptor_rva = import_directory + RVA::from(i * 0x14);
            let import_descriptor = read_image_import_descriptor(ws, import_descriptor_rva)?;
            if import_descriptor.is_empty() {
                break;
            }

            let dll_name = ws.read_utf8(import_descriptor.name)?;
            debug!("{:?} -> {}", import_descriptor, dll_name);

            for j in 0..std::usize::MAX {
                // the First Thunk (FT) is the pointer that will be overwritten upon load.
                // entries here may not point to the IMAGE_IMPORT_BY_NAME.
                let first_thunk = import_descriptor.first_thunk + RVA::from(j * psize);
                // the Original First Thunk (OFT) remains constant, and points to the IMAGE_IMPORT_BY_NAME.
                // FT and OFT are parallel arrays.
                let image_thunk_data_rva = import_descriptor.original_first_thunk + RVA::from(j * psize);
                match read_image_thunk_data(ws, image_thunk_data_rva)? {
                    ImageThunkData::Function(rva) => {
                        if rva == RVA(0x0) {
                            break;
                        } else {
                            let imp = read_image_import_by_name(ws, rva)?;
                            debug!("{:?}", imp);

                            symbols.push((first_thunk, format!("{}!{}", dll_name, imp.name)))
                        }
                    },
                    ImageThunkData::Ordinal(ord) => {
                        symbols.push((first_thunk, format!("{}!#{}", dll_name, ord)))
                    },
                };
            }
        }

        for (rva, name) in symbols.iter() {
            debug!("import: {:#x} -> {}", rva, name);
            ws.make_symbol(*rva, name)?;
            ws.analyze()?;
        }

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


<commit_before>
use std::collections::BTreeMap;

use assembler::parser::{Literal, Node, InstructionField};

// MARK: - Semantic Analysis

/**
 @return () in the event the arguments are valid for the directive, or a failure reason.
 */
fn validate_directive_semantics<'a>(identifier: &'a str, arguments: &Vec<Literal<'a>>) -> Result<(), String> {
    /*
     The 'org' directive is used to set the current output origin address. The directive requires
     a single numeric literal in the range $000-$FFF.
     */
    if identifier == "org" {
        if arguments.len() != 1 {
            return Err(format!("Incorrect number of parameters ({}) for directive .org, expecting 1", arguments.len()));
        }

        match arguments[0] {
            Literal::Numeric(a) if a <= 4095 => {}
            _ => {
                return Err(format!("Directive .org requires 1 numeric literal in the range $000-$FFF."));
            }
        }

        return Ok(());
    }

    /*
     The 'db' directive is used to emit a series of bytes at the current location.
     A list of 1 or more literals is required, either numeric or string.
     A string literal is emitted without a null-terminator.
     */
    if identifier == "db" {
        if arguments.len() == 0 {
            return Err(format!("Incorrect number of parameters ({}) for directive .db, expecting 1 or more.", arguments.len()));
        }

        for argument in arguments {
            if let &Literal::Numeric(value) = argument {
                if value > 0xFF {
                    return Err(format!("All numeric parameters to .db must be 1-byte literals (${:X} > $FF)", value));
                }
            }
        }

        return Ok(());
    }

    Err(format!("Unrecognized directive .{}", identifier))
}

/**
 @return The number of bytes the directive represents within the output stream.
 */
fn size_of_directive<'a>(identifier: &'a str, arguments: &Vec<Literal<'a>>) -> usize {
    // The origin directive is not emitted.
    if identifier == "org" {
        return 0;
    }

    // The number of bytes emitted by .db is the sum of argument lengths.
    if identifier == "db" {
        return arguments.iter().fold(0, |acc, argument| {
            match *argument {
                Literal::Numeric(number) if number <= 0xFF => { acc + 1 }
                Literal::String(string) => { acc + string.len() }
                _ => {
                    panic!("Semantic analysis for directive .{} failed.", identifier);
                }
            }
        });
    }

    // The first pass at semantic analysis should have caught an invalid directive.
    panic!("Semantic analysis for directive .{} failed.", identifier);
}

// MARK: - Pass 1: Define Labels

/**
 In a 2-pass assembler, the first pass is used to resolve the address of all labels
 declared in the assembly. This involves walking the ASL, computing the size of instructions
 and directives in the output stream and using that information to establish the label-address map.
 @param syntax_list The un-filtered syntax list.
 @return A tuple of (the filtered syntax list with all labels removed,
     the map of labels to their corresponding addresses) on success or a string describing
     the reason the first pass failed.
 */
fn define_and_filter_labels<'a>(syntax_list: Vec<Node<'a>>) -> Result<(Vec<Node<'a>>, BTreeMap<&'a str, usize>), String> {
    let mut label_address_map = BTreeMap::new();
    let mut current_address = 0x000;

    // Define all labels and process '.org' directives.
    for node in syntax_list.iter() {
        match *node {
            Node::Directive { identifier, ref arguments } => {
                // Validate the directive and that the arguments match the identifier.
                try!(validate_directive_semantics(identifier, arguments));

                // Adjust address by directive size.
                current_address = current_address + size_of_directive(identifier, arguments);

                // The origin directive changes the current address.
                if let ("org", &Literal::Numeric(address)) = (identifier, &arguments[0]) {
                    current_address = address;
                }
            }

            Node::Label { identifier } => {
                // Map the label to the current address and remove from the ASL.
                if label_address_map.contains_key(&identifier) {
                    return Err(format!("Attempted re-definition of label {}.", identifier));
                }

                label_address_map.insert(identifier, current_address);
            }

            Node::Instruction { mnemonic: _, fields: _ } => {
                // All Chip8 instructions are two bytes.
                current_address = current_address + 2;
            }
        }
    }

    // Remove any label nodes from the ASL as they have now been resolved.
    let filtered_asl = syntax_list.into_iter().filter(|node| {
        match *node {
            Node::Label { identifier: _ } => false,
            Node::Directive { identifier: _, arguments: _ } => true,
            Node::Instruction { mnemonic: _, fields: _ } => true
        }
    }).collect();

    // Semantic analysis for directives and labels is complete.
    Ok((filtered_asl, label_address_map))
}

/**
 Analyze the ASL for the assembly and convert it into an output byte stream.
 */
pub fn assemble<'a>(syntax_list: Vec<Node<'a>>) -> Result<(), String> {
    let _ = try!(define_and_filter_labels(syntax_list));
    Ok(())
}

// MARK: - Tests

#[cfg(test)]
mod tests {

  use super::*;
  use super::validate_directive_semantics;

  use assembler::scanner::Scanner;
  use assembler::parser::*;

  #[test]
  fn test_validate_directive_semantics_for_org() {
    let mut params = vec![];
    assert_eq!(validate_directive_semantics("org", &params), Err("Incorrect number of parameters (0) for directive .org, expecting 1".to_string()));
    params = vec![Literal::Numeric(1), Literal::Numeric(3)];
    assert_eq!(validate_directive_semantics("org", &params), Err("Incorrect number of parameters (2) for directive .org, expecting 1".to_string()));
    params = vec![Literal::Numeric(0x1000)];
    assert_eq!(validate_directive_semantics("org", &params), Err("Directive .org requires 1 numeric literal in the range $000-$FFF.".to_string()));
    params = vec![Literal::String("TEST_STRING")];
    assert_eq!(validate_directive_semantics("org", &params), Err("Directive .org requires 1 numeric literal in the range $000-$FFF.".to_string()));
    params = vec![Literal::Numeric(0xFFF)];
    assert_eq!(validate_directive_semantics("org", &params), Ok(()));
    params = vec![Literal::Numeric(0x000)];
    assert_eq!(validate_directive_semantics("org", &params), Ok(()));
  }

  #[test]
  fn test_validate_directive_semantics_for_db() {
    let mut params = vec![];
    assert_eq!(validate_directive_semantics("db", &params), Err("Incorrect number of parameters (0) for directive .db, expecting 1 or more.".to_string()));
    params = vec![Literal::Numeric(0x100)];
    assert_eq!(validate_directive_semantics("db", &params), Err("All numeric parameters to .db must be 1-byte literals ($100 > $FF)".to_string()));
    params = vec![Literal::String("TEST_STRING")];
    assert_eq!(validate_directive_semantics("db", &params), Ok(()));
    params = vec![Literal::Numeric(0xFF)];
    assert_eq!(validate_directive_semantics("db", &params), Ok(()));
    params = vec![Literal::Numeric(0x00)];
    assert_eq!(validate_directive_semantics("db", &params), Ok(()));
    params = vec![Literal::String("TEST_STRING"), Literal::Numeric(0x00)];
    assert_eq!(validate_directive_semantics("db", &params), Ok(()));
  }

  #[test]
  fn test_validate_directive_semantics_for_invalid_directive() {
    let mut params = vec![];
    assert_eq!(validate_directive_semantics("dw", &params), Err("Unrecognized directive .dw".to_string()));
    params = vec![Literal::Numeric(0x100)];
    assert_eq!(validate_directive_semantics("dw", &params), Err("Unrecognized directive .dw".to_string()));
  }

}
<commit_msg>Additional tests for directive helper functions.<commit_after>
use std::collections::BTreeMap;

use assembler::parser::{Literal, Node, InstructionField};

// MARK: - Constants

const BYTES_PER_INSTRUCTION: usize = 2;
const DIRECTIVE_ORG: &'static str = "org";
const DIRECTIVE_DB : &'static str = "db";

// MARK: - Semantic Analysis

/**
 @return () in the event the arguments are valid for the directive, or a failure reason.
 */
fn validate_directive_semantics<'a>(identifier: &'a str, arguments: &Vec<Literal<'a>>) -> Result<(), String> {
    /*
     The 'org' directive is used to set the current output origin address. The directive requires
     a single numeric literal in the range $000-$FFF.
     */
    if identifier == DIRECTIVE_ORG {
        if arguments.len() != 1 {
            return Err(format!("Incorrect number of parameters ({}) for directive .org, expecting 1", arguments.len()));
        }

        match arguments[0] {
            Literal::Numeric(a) if a <= 4095 => {}
            _ => {
                return Err(format!("Directive .org requires 1 numeric literal in the range $000-$FFF."));
            }
        }

        return Ok(());
    }

    /*
     The 'db' directive is used to emit a series of bytes at the current location.
     A list of 1 or more literals is required, either numeric or string.
     A string literal is emitted without a null-terminator.
     */
    if identifier == DIRECTIVE_DB {
        if arguments.len() == 0 {
            return Err(format!("Incorrect number of parameters ({}) for directive .db, expecting 1 or more.", arguments.len()));
        }

        for argument in arguments {
            if let &Literal::Numeric(value) = argument {
                if value > 0xFF {
                    return Err(format!("All numeric parameters to .db must be 1-byte literals (${:X} > $FF)", value));
                }
            }
        }

        return Ok(());
    }

    Err(format!("Unrecognized directive .{}", identifier))
}

/**
 @return The number of bytes the directive represents within the output stream.
 */
fn size_of_directive<'a>(identifier: &'a str, arguments: &Vec<Literal<'a>>) -> usize {
    // The origin directive is not emitted.
    if identifier == DIRECTIVE_ORG {
        return 0;
    }

    // The number of bytes emitted by .db is the sum of argument lengths.
    if identifier == DIRECTIVE_DB {
        return arguments.iter().fold(0, |acc, argument| {
            match *argument {
                Literal::Numeric(number) if number <= 0xFF => { acc + 1 }
                Literal::String(string) => { acc + string.len() }
                _ => {
                    panic!("Semantic analysis for directive .{} failed.", identifier);
                }
            }
        });
    }

    // The first pass at semantic analysis should have caught an invalid directive.
    panic!("Semantic analysis for directive .{} failed.", identifier);
}

// MARK: - Pass 1: Define Labels

/**
 In a 2-pass assembler, the first pass is used to resolve the address of all labels
 declared in the assembly. This involves walking the ASL, computing the size of instructions
 and directives in the output stream and using that information to establish the label-address map.
 @param syntax_list The un-filtered syntax list.
 @return A tuple of (the filtered syntax list with all labels removed,
     the map of labels to their corresponding addresses) on success or a string describing
     the reason the first pass failed.
 */
fn define_and_filter_labels<'a>(syntax_list: Vec<Node<'a>>) -> Result<(Vec<Node<'a>>, BTreeMap<&'a str, usize>), String> {
    let mut label_address_map = BTreeMap::new();
    let mut current_address = 0x000;

    // Define all labels and process '.org' directives.
    for node in syntax_list.iter() {
        match *node {
            Node::Directive { identifier, ref arguments } => {
                // Validate the directive and that the arguments match the identifier.
                try!(validate_directive_semantics(identifier, arguments));

                // Adjust address by directive size.
                current_address = current_address + size_of_directive(identifier, arguments);

                // The origin directive changes the current address.
                if let (DIRECTIVE_ORG, &Literal::Numeric(address)) = (identifier, &arguments[0]) {
                    current_address = address;
                }
            }

            Node::Label { identifier } => {
                // Map the label to the current address and remove from the ASL.
                if label_address_map.contains_key(&identifier) {
                    return Err(format!("Attempted re-definition of label {}.", identifier));
                }

                label_address_map.insert(identifier, current_address);
            }

            Node::Instruction { mnemonic: _, fields: _ } => {
                // All Chip8 instructions are the same length.
                current_address = current_address + BYTES_PER_INSTRUCTION;
            }
        }
    }

    // Remove any label nodes from the ASL as they have now been resolved.
    let filtered_asl = syntax_list.into_iter().filter(|node| {
        match *node {
            Node::Label { identifier: _ } => false,
            Node::Directive { identifier: _, arguments: _ } => true,
            Node::Instruction { mnemonic: _, fields: _ } => true
        }
    }).collect();

    // Semantic analysis for directives and labels is complete.
    Ok((filtered_asl, label_address_map))
}

/**
 Analyze the ASL for the assembly and convert it into an output byte stream.
 */
pub fn assemble<'a>(syntax_list: Vec<Node<'a>>) -> Result<(), String> {
    let _ = try!(define_and_filter_labels(syntax_list));
    Ok(())
}

// MARK: - Tests

#[cfg(test)]
mod tests {

  use super::*;
  use super::{validate_directive_semantics, size_of_directive};

  use assembler::scanner::Scanner;
  use assembler::parser::*;

  #[test]
  fn test_validate_directive_semantics_for_org() {
    let mut params = vec![];
    assert_eq!(validate_directive_semantics("org", &params), Err("Incorrect number of parameters (0) for directive .org, expecting 1".to_string()));
    params = vec![Literal::Numeric(1), Literal::Numeric(3)];
    assert_eq!(validate_directive_semantics("org", &params), Err("Incorrect number of parameters (2) for directive .org, expecting 1".to_string()));
    params = vec![Literal::Numeric(0x1000)];
    assert_eq!(validate_directive_semantics("org", &params), Err("Directive .org requires 1 numeric literal in the range $000-$FFF.".to_string()));
    params = vec![Literal::String("TEST_STRING")];
    assert_eq!(validate_directive_semantics("org", &params), Err("Directive .org requires 1 numeric literal in the range $000-$FFF.".to_string()));
    params = vec![Literal::Numeric(0xFFF)];
    assert_eq!(validate_directive_semantics("org", &params), Ok(()));
    params = vec![Literal::Numeric(0x000)];
    assert_eq!(validate_directive_semantics("org", &params), Ok(()));
  }

  #[test]
  fn test_validate_directive_semantics_for_db() {
    let mut params = vec![];
    assert_eq!(validate_directive_semantics("db", &params), Err("Incorrect number of parameters (0) for directive .db, expecting 1 or more.".to_string()));
    params = vec![Literal::Numeric(0x100)];
    assert_eq!(validate_directive_semantics("db", &params), Err("All numeric parameters to .db must be 1-byte literals ($100 > $FF)".to_string()));
    params = vec![Literal::String("TEST_STRING")];
    assert_eq!(validate_directive_semantics("db", &params), Ok(()));
    params = vec![Literal::Numeric(0xFF)];
    assert_eq!(validate_directive_semantics("db", &params), Ok(()));
    params = vec![Literal::Numeric(0x00)];
    assert_eq!(validate_directive_semantics("db", &params), Ok(()));
    params = vec![Literal::String("TEST_STRING"), Literal::Numeric(0x00)];
    assert_eq!(validate_directive_semantics("db", &params), Ok(()));
  }

  #[test]
  fn test_validate_directive_semantics_for_invalid_directive() {
    let mut params = vec![];
    assert_eq!(validate_directive_semantics("dw", &params), Err("Unrecognized directive .dw".to_string()));
    params = vec![Literal::Numeric(0x100)];
    assert_eq!(validate_directive_semantics("dw", &params), Err("Unrecognized directive .dw".to_string()));
  }

  #[test]
  fn test_size_of_directive_for_org() {
    assert_eq!(size_of_directive("org", &vec![]), 0);
    assert_eq!(size_of_directive("org", &vec![Literal::Numeric(0x100)]), 0);
    assert_eq!(size_of_directive("org", &vec![Literal::Numeric(0x100), Literal::String("TEST_STRING")]), 0);
  }

  #[test]
  fn test_size_of_directive_for_db() {
    assert_eq!(size_of_directive("db", &vec![]), 0);
    assert_eq!(size_of_directive("db", &vec![Literal::Numeric(0xFF)]), 1);
    assert_eq!(size_of_directive("db", &vec![Literal::String("TEST_STRING")]), 11);
    assert_eq!(size_of_directive("db", &vec![Literal::String("TEST_STRING"), Literal::Numeric(0x00)]), 12);
  }

}
<|endoftext|>


<commit_before>#[macro_export]
macro_rules! implement_lua_push {
    ($ty:ty, $cb:expr) => {
        impl<L> $crate::Push<L> for Sound where L: $crate::AsMutLua {
            fn push_to_lua(self, lua: L) -> $crate::PushGuard<L> {
                $crate::userdata::push_userdata(self, lua, $cb)
            }
        }
    };
}

#[macro_export]
macro_rules! implement_lua_read {
    ($ty:ty) => {
        impl<'s, 'c> hlua::LuaRead<&'c mut hlua::InsideCallback> for &'s mut Sound {
            fn lua_read_at_position(lua: &'c mut hlua::InsideCallback, index: i32) -> Result<&'s mut Sound, &'c mut hlua::InsideCallback> {
                // FIXME:
                unsafe { ::std::mem::transmute($crate::userdata::read_userdata::<$ty>(lua, index)) }
            }
        }

        impl<'s, 'c> hlua::LuaRead<&'c mut hlua::InsideCallback> for &'s Sound {
            fn lua_read_at_position(lua: &'c mut hlua::InsideCallback, index: i32) -> Result<&'s Sound, &'c mut hlua::InsideCallback> {
                // FIXME:
                unsafe { ::std::mem::transmute($crate::userdata::read_userdata::<$ty>(lua, index)) }
            }
        }
    };
}
<commit_msg>Remove hard-coded type from macros<commit_after>#[macro_export]
macro_rules! implement_lua_push {
    ($ty:ty, $cb:expr) => {
        impl<L> $crate::Push<L> for $ty where L: $crate::AsMutLua {
            fn push_to_lua(self, lua: L) -> $crate::PushGuard<L> {
                $crate::userdata::push_userdata(self, lua, $cb)
            }
        }
    };
}

#[macro_export]
macro_rules! implement_lua_read {
    ($ty:ty) => {
        impl<'s, 'c> hlua::LuaRead<&'c mut hlua::InsideCallback> for &'s mut $ty {
            fn lua_read_at_position(lua: &'c mut hlua::InsideCallback, index: i32) -> Result<&'s mut $ty, &'c mut hlua::InsideCallback> {
                // FIXME:
                unsafe { ::std::mem::transmute($crate::userdata::read_userdata::<$ty>(lua, index)) }
            }
        }

        impl<'s, 'c> hlua::LuaRead<&'c mut hlua::InsideCallback> for &'s $ty {
            fn lua_read_at_position(lua: &'c mut hlua::InsideCallback, index: i32) -> Result<&'s $ty, &'c mut hlua::InsideCallback> {
                // FIXME:
                unsafe { ::std::mem::transmute($crate::userdata::read_userdata::<$ty>(lua, index)) }
            }
        }
    };
}
<|endoftext|>


<commit_before>use super::internal::*;
use super::*;
use std::cmp::min;

/// Test if a byte is the start of a UTF-8 character.
/// (extracted from `str::is_char_boundary`)
fn is_char_boundary(b: u8) -> bool {
    // This is bit magic equivalent to: b < 128 || b >= 192
    (b as i8) >= -0x40
}

/// Find the index of a character boundary near the midpoint.
fn find_char_midpoint(chars: &str) -> usize {
    let mid = chars.len() / 2;

    // We want to split near the midpoint, but we need to find an actual
    // character boundary.  So we look at the raw bytes, first scanning
    // forward from the midpoint for a boundary, then trying backward.
    let (left, right) = chars.as_bytes().split_at(mid);
    right.iter()
        .cloned()
        .position(is_char_boundary)
        .map(|i| mid + i)
        .or_else(|| left.iter().cloned().rposition(is_char_boundary))
        .unwrap_or(0)
}

/// Parallel extensions for strings.
pub trait ParallelString {
    /// Returns a parallel iterator over the characters of a string.
    fn par_chars(&self) -> ParChars;

    /// Returns a parallel iterator over substrings separated by a
    /// given character, similar to `str::split`.
    fn par_split<P: Pattern>(&self, P) -> ParSplit<P>;

    /// Returns a parallel iterator over substrings terminated by a
    /// given character, similar to `str::split_terminator`.  It's
    /// equivalent to `par_split`, except it doesn't produce an empty
    /// substring after a trailing terminator.
    fn par_split_terminator<P: Pattern>(&self, P) -> ParSplitTerminator<P>;

    /// Returns a parallel iterator over the lines of a string, ending with an
    /// optional carriage return and with a newline (`\r\n` or just `\n`).
    /// The final line ending is optional, and line endings are not included in
    /// the output strings.
    fn par_lines(&self) -> ParLines;

    /// Returns a parallel iterator over the sub-slices of a string that are
    /// separated by any amount of whitespace.
    ///
    /// As with `str::split_whitespace`, 'whitespace' is defined according to
    /// the terms of the Unicode Derived Core Property `White_Space`.
    fn par_split_whitespace(&self) -> ParSplitWhitespace;
}

impl ParallelString for str {
    fn par_chars(&self) -> ParChars {
        ParChars { chars: self }
    }

    fn par_split<P: Pattern>(&self, separator: P) -> ParSplit<P> {
        ParSplit::new(self, separator)
    }

    fn par_split_terminator<P: Pattern>(&self, terminator: P) -> ParSplitTerminator<P> {
        ParSplitTerminator::new(self, terminator)
    }

    fn par_lines(&self) -> ParLines {
        ParLines(self)
    }

    fn par_split_whitespace(&self) -> ParSplitWhitespace {
        ParSplitWhitespace(self)
    }
}

// /////////////////////////////////////////////////////////////////////////

/// Pattern-matching trait for `ParallelString`, somewhat like a mix of
/// `std::str::pattern::{Pattern, Searcher}`.
pub trait Pattern: Sized + Sync {
    fn find_in(&self, &str) -> Option<usize>;
    fn rfind_in(&self, &str) -> Option<usize>;
    fn is_suffix_of(&self, &str) -> bool;
    fn fold_with<'ch, F>(&self, &'ch str, folder: F, skip_last: bool) -> F
        where F: Folder<&'ch str>;
}

impl Pattern for char {
    fn find_in(&self, chars: &str) -> Option<usize> {
        chars.find(*self)
    }

    fn rfind_in(&self, chars: &str) -> Option<usize> {
        chars.rfind(*self)
    }

    fn is_suffix_of(&self, chars: &str) -> bool {
        chars.ends_with(*self)
    }

    fn fold_with<'ch, F>(&self, chars: &'ch str, folder: F, skip_last: bool) -> F
        where F: Folder<&'ch str>
    {
        let mut split = chars.split(*self);
        if skip_last {
            split.next_back();
        }
        folder.consume_iter(split)
    }
}

impl<FN: Sync + Fn(char) -> bool> Pattern for FN {
    fn find_in(&self, chars: &str) -> Option<usize> {
        chars.find(self)
    }

    fn rfind_in(&self, chars: &str) -> Option<usize> {
        chars.rfind(self)
    }

    fn is_suffix_of(&self, chars: &str) -> bool {
        chars.ends_with(self)
    }

    fn fold_with<'ch, F>(&self, chars: &'ch str, folder: F, skip_last: bool) -> F
        where F: Folder<&'ch str>
    {
        let mut split = chars.split(self);
        if skip_last {
            split.next_back();
        }
        folder.consume_iter(split)
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParChars<'ch> {
    chars: &'ch str,
}

impl<'ch> ParallelIterator for ParChars<'ch> {
    type Item = char;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge_unindexed(self, consumer)
    }
}

impl<'ch> UnindexedProducer for ParChars<'ch> {
    type Item = char;

    fn split(&mut self) -> Option<Self> {
        let index = find_char_midpoint(self.chars);
        if index > 0 {
            let (left, right) = self.chars.split_at(index);
            self.chars = left;
            Some(ParChars { chars: right })
        } else {
            None
        }
    }

    fn fold_with<F>(self, folder: F) -> F
        where F: Folder<Self::Item>
    {
        folder.consume_iter(self.chars.chars())
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParSplit<'ch, P: Pattern> {
    chars: &'ch str,
    separator: P,
}

pub struct ParSplitProducer<'ch, 'sep, P: Pattern + 'sep> {
    chars: &'ch str,
    separator: &'sep P,

    /// Marks the endpoint beyond which we've already found no separators.
    tail: usize,
}

impl<'ch, P: Pattern> ParSplit<'ch, P> {
    fn new(chars: &'ch str, separator: P) -> Self {
        ParSplit {
            chars: chars,
            separator: separator,
        }
    }
}

impl<'ch, 'sep, P: Pattern + 'sep> ParSplitProducer<'ch, 'sep, P> {
    fn new(split: &'sep ParSplit<'ch, P>) -> Self {
        ParSplitProducer {
            chars: split.chars,
            separator: &split.separator,
            tail: split.chars.len(),
        }
    }

    /// Common `fold_with` implementation, integrating `ParSplitTerminator`'s
    /// need to sometimes skip its final empty item.
    fn fold_with<F>(self, folder: F, skip_last: bool) -> F
        where F: Folder<<Self as UnindexedProducer>::Item>
    {
        let ParSplitProducer { chars, separator, tail } = self;

        if tail == chars.len() {
            // No tail section, so just let `str::split` handle it.
            separator.fold_with(chars, folder, skip_last)

        } else if let Some(index) = separator.rfind_in(&chars[..tail]) {
            // We found the last separator to complete the tail, so
            // end with that slice after `str::split` finds the rest.
            let (left, right) = chars.split_at(index);
            let folder = separator.fold_with(left, folder, false);
            if skip_last || folder.full() {
                folder
            } else {
                let mut right_iter = right.chars();
                right_iter.next(); // skip the separator
                folder.consume(right_iter.as_str())
            }

        } else {
            // We know there are no separators at all.  Return our whole string.
            if skip_last {
                folder
            } else {
                folder.consume(chars)
            }
        }
    }
}

impl<'ch, P: Pattern> ParallelIterator for ParSplit<'ch, P> {
    type Item = &'ch str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        let producer = ParSplitProducer::new(&self);
        bridge_unindexed(producer, consumer)
    }
}

impl<'ch, 'sep, P: Pattern + 'sep> UnindexedProducer for ParSplitProducer<'ch, 'sep, P> {
    type Item = &'ch str;

    fn split(&mut self) -> Option<Self> {
        let ParSplitProducer { chars, separator, tail } = *self;

        // First find a suitable UTF-8 boundary in the unsearched region.
        let char_index = find_char_midpoint(&chars[..tail]);

        // Look forward for the separator, and failing that look backward.
        let index = separator.find_in(&chars[char_index..tail])
            .map(|i| char_index + i)
            .or_else(|| separator.rfind_in(&chars[..char_index]));

        if let Some(index) = index {
            let (left, right) = chars.split_at(index);

            // Update `self` as the region before the separator.
            self.chars = left;
            self.tail = min(char_index, index);

            // Create the right split following the separator.
            let mut right_iter = right.chars();
            right_iter.next(); // skip the separator
            let right_chars = right_iter.as_str();
            let right_index = chars.len() - right_chars.len();

            let mut right = ParSplitProducer {
                chars: right_chars,
                separator: separator,
                tail: tail - right_index,
            };

            // If we scanned backwards to find the separator, everything in
            // the right side is exhausted, with no separators left to find.
            if index < char_index {
                right.tail = 0;
            }

            Some(right)

        } else {
            self.tail = 0;
            None
        }
    }

    fn fold_with<F>(self, folder: F) -> F
        where F: Folder<Self::Item>
    {
        self.fold_with(folder, false)
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParSplitTerminator<'ch, P: Pattern> {
    splitter: ParSplit<'ch, P>,
}

pub struct ParSplitTerminatorProducer<'ch, 'sep, P: Pattern + 'sep> {
    splitter: ParSplitProducer<'ch, 'sep, P>,
    endpoint: bool,
}

impl<'ch, P: Pattern> ParSplitTerminator<'ch, P> {
    fn new(chars: &'ch str, terminator: P) -> Self {
        ParSplitTerminator {
            splitter: ParSplit::new(chars, terminator),
        }
    }
}

impl<'ch, 'sep, P: Pattern + 'sep> ParSplitTerminatorProducer<'ch, 'sep, P> {
    fn new(split: &'sep ParSplitTerminator<'ch, P>) -> Self {
        ParSplitTerminatorProducer {
            splitter: ParSplitProducer::new(&split.splitter),
            endpoint: true,
        }
    }
}

impl<'ch, P: Pattern> ParallelIterator for ParSplitTerminator<'ch, P> {
    type Item = &'ch str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        let producer = ParSplitTerminatorProducer::new(&self);
        bridge_unindexed(producer, consumer)
    }
}

impl<'ch, 'sep, P: Pattern + 'sep> UnindexedProducer for ParSplitTerminatorProducer<'ch, 'sep, P> {
    type Item = &'ch str;

    fn split(&mut self) -> Option<Self> {
        self.splitter.split().map(|right| {
            let endpoint = self.endpoint;
            self.endpoint = false;
            ParSplitTerminatorProducer {
                splitter: right,
                endpoint: endpoint,
            }
        })
    }

    fn fold_with<F>(self, folder: F) -> F
        where F: Folder<Self::Item>
    {
        // See if we need to eat the empty trailing substring
        let skip_last = if self.endpoint {
            let chars = self.splitter.chars;
            let terminator = self.splitter.separator;
            chars.is_empty() || terminator.is_suffix_of(chars)
        } else {
            false
        };

        self.splitter.fold_with(folder, skip_last)
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParLines<'ch>(&'ch str);

impl<'ch> ParallelIterator for ParLines<'ch> {
    type Item = &'ch str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        self.0
            .par_split_terminator('\n')
            .map(|line| {
                if line.ends_with('\r') {
                    &line[..line.len() - 1]
                } else {
                    line
                }
            })
            .drive_unindexed(consumer)
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParSplitWhitespace<'ch>(&'ch str);

impl<'ch> ParallelIterator for ParSplitWhitespace<'ch> {
    type Item = &'ch str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        self.0
            .par_split(char::is_whitespace)
            .filter(|string| !string.is_empty())
            .drive_unindexed(consumer)
    }
}
<commit_msg>Close the ParallelString trait<commit_after>use super::internal::*;
use super::*;
use std::cmp::min;

mod private {
    /// If this type is pub but not publically rechable, third parties can't
    /// name it and can't implement traits using it.
    pub struct PrivateMarker;
}

/// Test if a byte is the start of a UTF-8 character.
/// (extracted from `str::is_char_boundary`)
fn is_char_boundary(b: u8) -> bool {
    // This is bit magic equivalent to: b < 128 || b >= 192
    (b as i8) >= -0x40
}

/// Find the index of a character boundary near the midpoint.
fn find_char_midpoint(chars: &str) -> usize {
    let mid = chars.len() / 2;

    // We want to split near the midpoint, but we need to find an actual
    // character boundary.  So we look at the raw bytes, first scanning
    // forward from the midpoint for a boundary, then trying backward.
    let (left, right) = chars.as_bytes().split_at(mid);
    right.iter()
        .cloned()
        .position(is_char_boundary)
        .map(|i| mid + i)
        .or_else(|| left.iter().cloned().rposition(is_char_boundary))
        .unwrap_or(0)
}

/// Parallel extensions for strings.
pub trait ParallelString {
    /// This trait is private; this method exists to make it impossible
    /// to implement outside the crate.
    #[doc(hidden)]
    fn private_parallel_string(&self) -> private::PrivateMarker;

    /// Returns a parallel iterator over the characters of a string.
    fn par_chars(&self) -> ParChars;

    /// Returns a parallel iterator over substrings separated by a
    /// given character, similar to `str::split`.
    fn par_split<P: Pattern>(&self, P) -> ParSplit<P>;

    /// Returns a parallel iterator over substrings terminated by a
    /// given character, similar to `str::split_terminator`.  It's
    /// equivalent to `par_split`, except it doesn't produce an empty
    /// substring after a trailing terminator.
    fn par_split_terminator<P: Pattern>(&self, P) -> ParSplitTerminator<P>;

    /// Returns a parallel iterator over the lines of a string, ending with an
    /// optional carriage return and with a newline (`\r\n` or just `\n`).
    /// The final line ending is optional, and line endings are not included in
    /// the output strings.
    fn par_lines(&self) -> ParLines;

    /// Returns a parallel iterator over the sub-slices of a string that are
    /// separated by any amount of whitespace.
    ///
    /// As with `str::split_whitespace`, 'whitespace' is defined according to
    /// the terms of the Unicode Derived Core Property `White_Space`.
    fn par_split_whitespace(&self) -> ParSplitWhitespace;
}

impl ParallelString for str {
    fn private_parallel_string(&self) -> private::PrivateMarker {
        private::PrivateMarker
    }

    fn par_chars(&self) -> ParChars {
        ParChars { chars: self }
    }

    fn par_split<P: Pattern>(&self, separator: P) -> ParSplit<P> {
        ParSplit::new(self, separator)
    }

    fn par_split_terminator<P: Pattern>(&self, terminator: P) -> ParSplitTerminator<P> {
        ParSplitTerminator::new(self, terminator)
    }

    fn par_lines(&self) -> ParLines {
        ParLines(self)
    }

    fn par_split_whitespace(&self) -> ParSplitWhitespace {
        ParSplitWhitespace(self)
    }
}

// /////////////////////////////////////////////////////////////////////////

/// Pattern-matching trait for `ParallelString`, somewhat like a mix of
/// `std::str::pattern::{Pattern, Searcher}`.
pub trait Pattern: Sized + Sync {
    fn find_in(&self, &str) -> Option<usize>;
    fn rfind_in(&self, &str) -> Option<usize>;
    fn is_suffix_of(&self, &str) -> bool;
    fn fold_with<'ch, F>(&self, &'ch str, folder: F, skip_last: bool) -> F
        where F: Folder<&'ch str>;
}

impl Pattern for char {
    fn find_in(&self, chars: &str) -> Option<usize> {
        chars.find(*self)
    }

    fn rfind_in(&self, chars: &str) -> Option<usize> {
        chars.rfind(*self)
    }

    fn is_suffix_of(&self, chars: &str) -> bool {
        chars.ends_with(*self)
    }

    fn fold_with<'ch, F>(&self, chars: &'ch str, folder: F, skip_last: bool) -> F
        where F: Folder<&'ch str>
    {
        let mut split = chars.split(*self);
        if skip_last {
            split.next_back();
        }
        folder.consume_iter(split)
    }
}

impl<FN: Sync + Fn(char) -> bool> Pattern for FN {
    fn find_in(&self, chars: &str) -> Option<usize> {
        chars.find(self)
    }

    fn rfind_in(&self, chars: &str) -> Option<usize> {
        chars.rfind(self)
    }

    fn is_suffix_of(&self, chars: &str) -> bool {
        chars.ends_with(self)
    }

    fn fold_with<'ch, F>(&self, chars: &'ch str, folder: F, skip_last: bool) -> F
        where F: Folder<&'ch str>
    {
        let mut split = chars.split(self);
        if skip_last {
            split.next_back();
        }
        folder.consume_iter(split)
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParChars<'ch> {
    chars: &'ch str,
}

impl<'ch> ParallelIterator for ParChars<'ch> {
    type Item = char;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge_unindexed(self, consumer)
    }
}

impl<'ch> UnindexedProducer for ParChars<'ch> {
    type Item = char;

    fn split(&mut self) -> Option<Self> {
        let index = find_char_midpoint(self.chars);
        if index > 0 {
            let (left, right) = self.chars.split_at(index);
            self.chars = left;
            Some(ParChars { chars: right })
        } else {
            None
        }
    }

    fn fold_with<F>(self, folder: F) -> F
        where F: Folder<Self::Item>
    {
        folder.consume_iter(self.chars.chars())
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParSplit<'ch, P: Pattern> {
    chars: &'ch str,
    separator: P,
}

pub struct ParSplitProducer<'ch, 'sep, P: Pattern + 'sep> {
    chars: &'ch str,
    separator: &'sep P,

    /// Marks the endpoint beyond which we've already found no separators.
    tail: usize,
}

impl<'ch, P: Pattern> ParSplit<'ch, P> {
    fn new(chars: &'ch str, separator: P) -> Self {
        ParSplit {
            chars: chars,
            separator: separator,
        }
    }
}

impl<'ch, 'sep, P: Pattern + 'sep> ParSplitProducer<'ch, 'sep, P> {
    fn new(split: &'sep ParSplit<'ch, P>) -> Self {
        ParSplitProducer {
            chars: split.chars,
            separator: &split.separator,
            tail: split.chars.len(),
        }
    }

    /// Common `fold_with` implementation, integrating `ParSplitTerminator`'s
    /// need to sometimes skip its final empty item.
    fn fold_with<F>(self, folder: F, skip_last: bool) -> F
        where F: Folder<<Self as UnindexedProducer>::Item>
    {
        let ParSplitProducer { chars, separator, tail } = self;

        if tail == chars.len() {
            // No tail section, so just let `str::split` handle it.
            separator.fold_with(chars, folder, skip_last)

        } else if let Some(index) = separator.rfind_in(&chars[..tail]) {
            // We found the last separator to complete the tail, so
            // end with that slice after `str::split` finds the rest.
            let (left, right) = chars.split_at(index);
            let folder = separator.fold_with(left, folder, false);
            if skip_last || folder.full() {
                folder
            } else {
                let mut right_iter = right.chars();
                right_iter.next(); // skip the separator
                folder.consume(right_iter.as_str())
            }

        } else {
            // We know there are no separators at all.  Return our whole string.
            if skip_last {
                folder
            } else {
                folder.consume(chars)
            }
        }
    }
}

impl<'ch, P: Pattern> ParallelIterator for ParSplit<'ch, P> {
    type Item = &'ch str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        let producer = ParSplitProducer::new(&self);
        bridge_unindexed(producer, consumer)
    }
}

impl<'ch, 'sep, P: Pattern + 'sep> UnindexedProducer for ParSplitProducer<'ch, 'sep, P> {
    type Item = &'ch str;

    fn split(&mut self) -> Option<Self> {
        let ParSplitProducer { chars, separator, tail } = *self;

        // First find a suitable UTF-8 boundary in the unsearched region.
        let char_index = find_char_midpoint(&chars[..tail]);

        // Look forward for the separator, and failing that look backward.
        let index = separator.find_in(&chars[char_index..tail])
            .map(|i| char_index + i)
            .or_else(|| separator.rfind_in(&chars[..char_index]));

        if let Some(index) = index {
            let (left, right) = chars.split_at(index);

            // Update `self` as the region before the separator.
            self.chars = left;
            self.tail = min(char_index, index);

            // Create the right split following the separator.
            let mut right_iter = right.chars();
            right_iter.next(); // skip the separator
            let right_chars = right_iter.as_str();
            let right_index = chars.len() - right_chars.len();

            let mut right = ParSplitProducer {
                chars: right_chars,
                separator: separator,
                tail: tail - right_index,
            };

            // If we scanned backwards to find the separator, everything in
            // the right side is exhausted, with no separators left to find.
            if index < char_index {
                right.tail = 0;
            }

            Some(right)

        } else {
            self.tail = 0;
            None
        }
    }

    fn fold_with<F>(self, folder: F) -> F
        where F: Folder<Self::Item>
    {
        self.fold_with(folder, false)
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParSplitTerminator<'ch, P: Pattern> {
    splitter: ParSplit<'ch, P>,
}

pub struct ParSplitTerminatorProducer<'ch, 'sep, P: Pattern + 'sep> {
    splitter: ParSplitProducer<'ch, 'sep, P>,
    endpoint: bool,
}

impl<'ch, P: Pattern> ParSplitTerminator<'ch, P> {
    fn new(chars: &'ch str, terminator: P) -> Self {
        ParSplitTerminator {
            splitter: ParSplit::new(chars, terminator),
        }
    }
}

impl<'ch, 'sep, P: Pattern + 'sep> ParSplitTerminatorProducer<'ch, 'sep, P> {
    fn new(split: &'sep ParSplitTerminator<'ch, P>) -> Self {
        ParSplitTerminatorProducer {
            splitter: ParSplitProducer::new(&split.splitter),
            endpoint: true,
        }
    }
}

impl<'ch, P: Pattern> ParallelIterator for ParSplitTerminator<'ch, P> {
    type Item = &'ch str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        let producer = ParSplitTerminatorProducer::new(&self);
        bridge_unindexed(producer, consumer)
    }
}

impl<'ch, 'sep, P: Pattern + 'sep> UnindexedProducer for ParSplitTerminatorProducer<'ch, 'sep, P> {
    type Item = &'ch str;

    fn split(&mut self) -> Option<Self> {
        self.splitter.split().map(|right| {
            let endpoint = self.endpoint;
            self.endpoint = false;
            ParSplitTerminatorProducer {
                splitter: right,
                endpoint: endpoint,
            }
        })
    }

    fn fold_with<F>(self, folder: F) -> F
        where F: Folder<Self::Item>
    {
        // See if we need to eat the empty trailing substring
        let skip_last = if self.endpoint {
            let chars = self.splitter.chars;
            let terminator = self.splitter.separator;
            chars.is_empty() || terminator.is_suffix_of(chars)
        } else {
            false
        };

        self.splitter.fold_with(folder, skip_last)
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParLines<'ch>(&'ch str);

impl<'ch> ParallelIterator for ParLines<'ch> {
    type Item = &'ch str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        self.0
            .par_split_terminator('\n')
            .map(|line| {
                if line.ends_with('\r') {
                    &line[..line.len() - 1]
                } else {
                    line
                }
            })
            .drive_unindexed(consumer)
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParSplitWhitespace<'ch>(&'ch str);

impl<'ch> ParallelIterator for ParSplitWhitespace<'ch> {
    type Item = &'ch str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        self.0
            .par_split(char::is_whitespace)
            .filter(|string| !string.is_empty())
            .drive_unindexed(consumer)
    }
}
<|endoftext|>