<commit_before>// Copyright © 2018, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

use failure::{format_err, Error};
use mimir2::{
    adapters::secondary::elasticsearch::{
        self,
        configuration::{IndexConfiguration, IndexMappings, IndexParameters, IndexSettings},
    },
    domain::ports::secondary::remote::Remote,
};
use serde_json::json;
use std::path::PathBuf;
use structopt::StructOpt;

// #[cfg(test)]
// #[macro_use]
// extern crate approx;

#[derive(StructOpt, Debug)]
struct Args {
    /// cosmogony file
    #[structopt(short = "i", long = "input")]
    input: String,
    /// Elasticsearch parameters.
    #[structopt(
        short = "c",
        long = "connection-string",
        default_value = "http://localhost:9200/munin"
    )]
    connection_string: String,
    /// Name of the dataset.
    #[structopt(short = "d", long = "dataset", default_value = "fr")]
    dataset: String,
    #[structopt(parse(from_os_str), default_value = "./config/admin/mappings.json")]
    mappings: PathBuf,
    #[structopt(parse(from_os_str), default_value = "./config/admin/settings.json")]
    settings: PathBuf,
    /// Number of shards for the es index
    #[structopt(short = "s", long = "nb-shards")]
    nb_shards: Option<usize>,
    /// Number of replicas for the es index
    #[structopt(short = "r", long = "nb-replicas")]
    nb_replicas: Option<usize>,
    /// Languages codes, used to build i18n names and labels
    #[structopt(name = "lang", short, long)]
    langs: Vec<String>,
}

#[tokio::main]
async fn main() {
    mimirsbrunn::utils::launch_async(index_cosmogony).await;
}

async fn index_cosmogony(args: Args) -> Result<(), Error> {
    let pool = elasticsearch::remote::connection_pool_url(&args.connection_string)
        .await
        .map_err(|err| {
            format_err!(
                "could not create elasticsearch connection pool: {}",
                err.to_string()
            )
        })?;

    let client = pool
        .conn()
        .await
        .map_err(|err| format_err!("could not connect elasticsearch pool: {}", err.to_string()))?;

    let settings = tokio::fs::read_to_string(args.settings.clone())
        .await
        .map_err(|err| {
            format_err!(
                "could not read settings file from '{}': {}",
                args.settings.display(),
                err.to_string(),
            )
        })?;

    let mut settings: serde_json::Value = serde_json::from_str(&settings).map_err(|err| {
        format_err!(
            "could not deserialize settings file from '{}': {}",
            args.settings.display(),
            err.to_string(),
        )
    })?;

    if let Some(nb_shards) = args.nb_shards {
        settings["index"]["number_of_shards"] = json!(nb_shards);
    }
    if let Some(nb_replicas) = args.nb_replicas {
        settings["index"]["number_of_replicas"] = json!(nb_replicas);
    }

    let mappings = tokio::fs::read_to_string(args.mappings.clone())
        .await
        .map_err(|err| {
            format_err!(
                "could not read mappings file from '{}': {}",
                args.mappings.display(),
                err.to_string(),
            )
        })?;

    let mappings = serde_json::from_str(&mappings).map_err(|err| {
        format_err!(
            "could not deserialize mappings file from '{}': {}",
            args.mappings.display(),
            err.to_string(),
        )
    })?;

    let config = IndexConfiguration {
        name: args.dataset.clone(),
        parameters: IndexParameters {
            timeout: String::from("10s"),
            wait_for_active_shards: String::from("1"), // only the primary shard
        },
        settings: IndexSettings { value: settings },
        mappings: IndexMappings { value: mappings },
    };

    mimirsbrunn::admin::index_cosmogony(args.input, args.langs, config, client)
        .await
        .map_err(|err| format_err!("could not index cosmogony: {}", err.to_string(),))
}

#[cfg(test)]
mod tests {
    use approx::assert_relative_eq;
    use futures::TryStreamExt;

    use super::*;
    use common::document::ContainerDocument;
    use mimir2::domain::model::query::Query;
    use mimir2::domain::ports::primary::list_documents::ListDocuments;
    use mimir2::domain::ports::primary::search_documents::SearchDocuments;
    use mimir2::{adapters::secondary::elasticsearch::remote, utils::docker};
    use places::admin::Admin;
    use places::Place;

    fn elasticsearch_test_url() -> String {
        std::env::var(elasticsearch::remote::ES_TEST_KEY).expect("env var")
    }

    #[tokio::test]
    async fn should_return_an_error_when_given_an_invalid_es_url() {
        let url = String::from("http://example.com:demo");
        let args = Args {
            input: String::from("foo"),
            connection_string: url,
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/admin/mappings.json"),
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;
        assert!(res
            .unwrap_err()
            .to_string()
            .contains("could not create elasticsearch connection pool: Invalid Elasticsearch URL"));
    }

    #[tokio::test]
    async fn should_return_an_error_when_given_an_url_not_es() {
        let url = String::from("http://no-es.test");
        let args = Args {
            input: String::from("foo"),
            connection_string: url,
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/admin/mappings.json"),
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;
        assert!(res
            .unwrap_err()
            .to_string()
            .contains("could not connect elasticsearch pool: Connection Error: Elasticsearch Connection Error"));
    }

    #[tokio::test]
    async fn should_return_an_error_when_given_an_invalid_path_for_mappings() {
        let guard = docker::initialize()
            .await
            .expect("elasticsearch docker initialization");

        let args = Args {
            input: String::from("foo"),
            connection_string: elasticsearch_test_url(),
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/invalid.json"), // a file that does not exists
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;

        drop(guard);

        assert!(res
            .unwrap_err()
            .to_string()
            .contains("could not read mappings file from"));
    }

    #[tokio::test]
    async fn should_return_an_error_when_given_an_invalid_mappings() {
        let guard = docker::initialize()
            .await
            .expect("elasticsearch docker initialization");

        let args = Args {
            input: String::from("foo"),
            connection_string: elasticsearch_test_url(),
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./README.md"), // exists, but not json
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;

        drop(guard);

        assert!(res
            .unwrap_err()
            .to_string()
            .contains("could not deserialize mappings file from"));
    }

    #[tokio::test]
    async fn should_return_an_error_when_given_an_invalid_path_for_input() {
        let guard = docker::initialize()
            .await
            .expect("elasticsearch docker initialization");

        let args = Args {
            input: String::from("./invalid.jsonl.gz"),
            connection_string: elasticsearch_test_url(),
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/admin/mappings.json"),
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;

        drop(guard);

        assert!(res
            .unwrap_err()
            .to_string()
            .contains("could not index cosmogony: No such file or directory"));
    }

    #[tokio::test]
    async fn should_correctly_index_a_small_cosmogony_file() {
        let guard = docker::initialize()
            .await
            .expect("elasticsearch docker initialization");

        let args = Args {
            input: String::from("./tests/fixtures/cosmogony/bretagne.small.jsonl.gz"),
            connection_string: elasticsearch_test_url(),
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/admin/mappings.json"),
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let _res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;

        // Now we query the index we just created. Since it's a small cosmogony file with few entries,
        // we'll just list all the documents in the index, and check them.
        let pool = remote::connection_test_pool()
            .await
            .expect("Elasticsearch Connection Pool");

        let client = pool
            .conn()
            .await
            .expect("Elasticsearch Connection Established");

        let admins: Vec<Admin> = client
            .list_documents()
            .unwrap()
            .try_collect()
            .await
            .unwrap();

        drop(guard);

        assert!(admins.iter().all(|admin| admin.boundary.is_some()));
        assert!(admins.iter().all(|admin| admin.coord.is_valid()));
    }

    #[tokio::test]
    async fn should_correctly_index_cosmogony_with_langs() {
        let guard = docker::initialize()
            .await
            .expect("elasticsearch docker initialization");

        let args = Args {
            input: String::from("./tests/fixtures/cosmogony/bretagne.small.jsonl.gz"),
            connection_string: elasticsearch_test_url(),
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/admin/mappings.json"),
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec!["fr".into(), "en".into()],
        };

        let _res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;

        let pool = remote::connection_test_pool()
            .await
            .expect("Elasticsearch Connection Pool");
        let client = pool
            .conn()
            .await
            .expect("Elasticsearch Connection Established");

        // FIXME Maybe should get a Vec<Admin> rather than Vec<AdminDoc>
        let admins: Vec<Admin> = client
            .list_documents()
            .unwrap()
            .try_collect()
            .await
            .unwrap();

        drop(guard);

        let brittany = admins.iter().find(|a| a.name == "Bretagne").unwrap();
        assert_eq!(brittany.names.get("fr"), Some("Bretagne"));
        assert_eq!(brittany.names.get("en"), Some("Brittany"));
        assert_eq!(brittany.labels.get("en"), Some("Brittany"));
    }

    #[tokio::test]
    async fn should_index_cosmogony_with_correct_values() {
        let guard = docker::initialize()
            .await
            .expect("elasticsearch docker initialization");

        let args = Args {
            input: String::from("./tests/fixtures/cosmogony/bretagne.small.jsonl.gz"),
            connection_string: elasticsearch_test_url(),
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/admin/mappings.json"),
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let _res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;

        // Now we query the index we just created. Since a small cosmogony file with few entries,
        // we'll just list all the documents in the index, and check them.
        let pool = remote::connection_test_pool()
            .await
            .expect("Elasticsearch Connection Pool");
        let client = pool
            .conn()
            .await
            .expect("Elasticsearch Connection Established");

        let admins: Vec<Admin> = client
            .search_documents(
                vec![String::from(Admin::static_doc_type())], // Fixme Use ContainerDoc::static_doc_type()
                Query::QueryString(String::from("name:bretagne")),
            )
            .await
            .unwrap()
            .into_iter()
            .map(|json| serde_json::from_value::<Place>(json).unwrap())
            .map(|place| match place {
                Place::Admin(admin) => admin,
                _ => panic!("should only have admins"),
            })
            .collect();

        drop(guard);

        let brittany = admins.iter().find(|a| a.name == "Bretagne").unwrap();
        assert_eq!(brittany.id, "admin:osm:relation:102740");
        assert_eq!(brittany.zone_type, Some(cosmogony::ZoneType::State));
        assert_relative_eq!(brittany.weight, 0.002_298, epsilon = 1e-6);
        assert_eq!(
            brittany.codes,
            vec![
                ("ISO3166-2", "FR-BRE"),
                ("ref:INSEE", "53"),
                ("ref:nuts", "FRH;FRH0"),
                ("ref:nuts:1", "FRH"),
                ("ref:nuts:2", "FRH0"),
                ("wikidata", "Q12130")
            ]
            .into_iter()
            .map(|(k, v)| (k.to_string(), v.to_string()))
            .collect()
        )
    }
}
<commit_msg>Remove deprecated comment<commit_after>// Copyright © 2018, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

use failure::{format_err, Error};
use mimir2::{
    adapters::secondary::elasticsearch::{
        self,
        configuration::{IndexConfiguration, IndexMappings, IndexParameters, IndexSettings},
    },
    domain::ports::secondary::remote::Remote,
};
use serde_json::json;
use std::path::PathBuf;
use structopt::StructOpt;

// #[cfg(test)]
// #[macro_use]
// extern crate approx;

#[derive(StructOpt, Debug)]
struct Args {
    /// cosmogony file
    #[structopt(short = "i", long = "input")]
    input: String,
    /// Elasticsearch parameters.
    #[structopt(
        short = "c",
        long = "connection-string",
        default_value = "http://localhost:9200/munin"
    )]
    connection_string: String,
    /// Name of the dataset.
    #[structopt(short = "d", long = "dataset", default_value = "fr")]
    dataset: String,
    #[structopt(parse(from_os_str), default_value = "./config/admin/mappings.json")]
    mappings: PathBuf,
    #[structopt(parse(from_os_str), default_value = "./config/admin/settings.json")]
    settings: PathBuf,
    /// Number of shards for the es index
    #[structopt(short = "s", long = "nb-shards")]
    nb_shards: Option<usize>,
    /// Number of replicas for the es index
    #[structopt(short = "r", long = "nb-replicas")]
    nb_replicas: Option<usize>,
    /// Languages codes, used to build i18n names and labels
    #[structopt(name = "lang", short, long)]
    langs: Vec<String>,
}

#[tokio::main]
async fn main() {
    mimirsbrunn::utils::launch_async(index_cosmogony).await;
}

async fn index_cosmogony(args: Args) -> Result<(), Error> {
    let pool = elasticsearch::remote::connection_pool_url(&args.connection_string)
        .await
        .map_err(|err| {
            format_err!(
                "could not create elasticsearch connection pool: {}",
                err.to_string()
            )
        })?;

    let client = pool
        .conn()
        .await
        .map_err(|err| format_err!("could not connect elasticsearch pool: {}", err.to_string()))?;

    let settings = tokio::fs::read_to_string(args.settings.clone())
        .await
        .map_err(|err| {
            format_err!(
                "could not read settings file from '{}': {}",
                args.settings.display(),
                err.to_string(),
            )
        })?;

    let mut settings: serde_json::Value = serde_json::from_str(&settings).map_err(|err| {
        format_err!(
            "could not deserialize settings file from '{}': {}",
            args.settings.display(),
            err.to_string(),
        )
    })?;

    if let Some(nb_shards) = args.nb_shards {
        settings["index"]["number_of_shards"] = json!(nb_shards);
    }
    if let Some(nb_replicas) = args.nb_replicas {
        settings["index"]["number_of_replicas"] = json!(nb_replicas);
    }

    let mappings = tokio::fs::read_to_string(args.mappings.clone())
        .await
        .map_err(|err| {
            format_err!(
                "could not read mappings file from '{}': {}",
                args.mappings.display(),
                err.to_string(),
            )
        })?;

    let mappings = serde_json::from_str(&mappings).map_err(|err| {
        format_err!(
            "could not deserialize mappings file from '{}': {}",
            args.mappings.display(),
            err.to_string(),
        )
    })?;

    let config = IndexConfiguration {
        name: args.dataset.clone(),
        parameters: IndexParameters {
            timeout: String::from("10s"),
            wait_for_active_shards: String::from("1"), // only the primary shard
        },
        settings: IndexSettings { value: settings },
        mappings: IndexMappings { value: mappings },
    };

    mimirsbrunn::admin::index_cosmogony(args.input, args.langs, config, client)
        .await
        .map_err(|err| format_err!("could not index cosmogony: {}", err.to_string(),))
}

#[cfg(test)]
mod tests {
    use approx::assert_relative_eq;
    use futures::TryStreamExt;

    use super::*;
    use common::document::ContainerDocument;
    use mimir2::domain::model::query::Query;
    use mimir2::domain::ports::primary::list_documents::ListDocuments;
    use mimir2::domain::ports::primary::search_documents::SearchDocuments;
    use mimir2::{adapters::secondary::elasticsearch::remote, utils::docker};
    use places::admin::Admin;
    use places::Place;

    fn elasticsearch_test_url() -> String {
        std::env::var(elasticsearch::remote::ES_TEST_KEY).expect("env var")
    }

    #[tokio::test]
    async fn should_return_an_error_when_given_an_invalid_es_url() {
        let url = String::from("http://example.com:demo");
        let args = Args {
            input: String::from("foo"),
            connection_string: url,
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/admin/mappings.json"),
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;
        assert!(res
            .unwrap_err()
            .to_string()
            .contains("could not create elasticsearch connection pool: Invalid Elasticsearch URL"));
    }

    #[tokio::test]
    async fn should_return_an_error_when_given_an_url_not_es() {
        let url = String::from("http://no-es.test");
        let args = Args {
            input: String::from("foo"),
            connection_string: url,
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/admin/mappings.json"),
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;
        assert!(res
            .unwrap_err()
            .to_string()
            .contains("could not connect elasticsearch pool: Connection Error: Elasticsearch Connection Error"));
    }

    #[tokio::test]
    async fn should_return_an_error_when_given_an_invalid_path_for_mappings() {
        let guard = docker::initialize()
            .await
            .expect("elasticsearch docker initialization");

        let args = Args {
            input: String::from("foo"),
            connection_string: elasticsearch_test_url(),
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/invalid.json"), // a file that does not exists
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;

        drop(guard);

        assert!(res
            .unwrap_err()
            .to_string()
            .contains("could not read mappings file from"));
    }

    #[tokio::test]
    async fn should_return_an_error_when_given_an_invalid_mappings() {
        let guard = docker::initialize()
            .await
            .expect("elasticsearch docker initialization");

        let args = Args {
            input: String::from("foo"),
            connection_string: elasticsearch_test_url(),
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./README.md"), // exists, but not json
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;

        drop(guard);

        assert!(res
            .unwrap_err()
            .to_string()
            .contains("could not deserialize mappings file from"));
    }

    #[tokio::test]
    async fn should_return_an_error_when_given_an_invalid_path_for_input() {
        let guard = docker::initialize()
            .await
            .expect("elasticsearch docker initialization");

        let args = Args {
            input: String::from("./invalid.jsonl.gz"),
            connection_string: elasticsearch_test_url(),
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/admin/mappings.json"),
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;

        drop(guard);

        assert!(res
            .unwrap_err()
            .to_string()
            .contains("could not index cosmogony: No such file or directory"));
    }

    #[tokio::test]
    async fn should_correctly_index_a_small_cosmogony_file() {
        let guard = docker::initialize()
            .await
            .expect("elasticsearch docker initialization");

        let args = Args {
            input: String::from("./tests/fixtures/cosmogony/bretagne.small.jsonl.gz"),
            connection_string: elasticsearch_test_url(),
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/admin/mappings.json"),
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let _res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;

        // Now we query the index we just created. Since it's a small cosmogony file with few entries,
        // we'll just list all the documents in the index, and check them.
        let pool = remote::connection_test_pool()
            .await
            .expect("Elasticsearch Connection Pool");

        let client = pool
            .conn()
            .await
            .expect("Elasticsearch Connection Established");

        let admins: Vec<Admin> = client
            .list_documents()
            .unwrap()
            .try_collect()
            .await
            .unwrap();

        drop(guard);

        assert!(admins.iter().all(|admin| admin.boundary.is_some()));
        assert!(admins.iter().all(|admin| admin.coord.is_valid()));
    }

    #[tokio::test]
    async fn should_correctly_index_cosmogony_with_langs() {
        let guard = docker::initialize()
            .await
            .expect("elasticsearch docker initialization");

        let args = Args {
            input: String::from("./tests/fixtures/cosmogony/bretagne.small.jsonl.gz"),
            connection_string: elasticsearch_test_url(),
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/admin/mappings.json"),
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec!["fr".into(), "en".into()],
        };

        let _res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;

        let pool = remote::connection_test_pool()
            .await
            .expect("Elasticsearch Connection Pool");
        let client = pool
            .conn()
            .await
            .expect("Elasticsearch Connection Established");

        let admins: Vec<Admin> = client
            .list_documents()
            .unwrap()
            .try_collect()
            .await
            .unwrap();

        drop(guard);

        let brittany = admins.iter().find(|a| a.name == "Bretagne").unwrap();
        assert_eq!(brittany.names.get("fr"), Some("Bretagne"));
        assert_eq!(brittany.names.get("en"), Some("Brittany"));
        assert_eq!(brittany.labels.get("en"), Some("Brittany"));
    }

    #[tokio::test]
    async fn should_index_cosmogony_with_correct_values() {
        let guard = docker::initialize()
            .await
            .expect("elasticsearch docker initialization");

        let args = Args {
            input: String::from("./tests/fixtures/cosmogony/bretagne.small.jsonl.gz"),
            connection_string: elasticsearch_test_url(),
            dataset: String::from("dataset"),
            mappings: PathBuf::from("./config/admin/mappings.json"),
            settings: PathBuf::from("./config/admin/settings.json"),
            nb_shards: None,
            nb_replicas: None,
            langs: vec![],
        };

        let _res = mimirsbrunn::utils::launch_async_args(index_cosmogony, args).await;

        // Now we query the index we just created. Since a small cosmogony file with few entries,
        // we'll just list all the documents in the index, and check them.
        let pool = remote::connection_test_pool()
            .await
            .expect("Elasticsearch Connection Pool");
        let client = pool
            .conn()
            .await
            .expect("Elasticsearch Connection Established");

        let admins: Vec<Admin> = client
            .search_documents(
                vec![String::from(Admin::static_doc_type())], // Fixme Use ContainerDoc::static_doc_type()
                Query::QueryString(String::from("name:bretagne")),
            )
            .await
            .unwrap()
            .into_iter()
            .map(|json| serde_json::from_value::<Place>(json).unwrap())
            .map(|place| match place {
                Place::Admin(admin) => admin,
                _ => panic!("should only have admins"),
            })
            .collect();

        drop(guard);

        let brittany = admins.iter().find(|a| a.name == "Bretagne").unwrap();
        assert_eq!(brittany.id, "admin:osm:relation:102740");
        assert_eq!(brittany.zone_type, Some(cosmogony::ZoneType::State));
        assert_relative_eq!(brittany.weight, 0.002_298, epsilon = 1e-6);
        assert_eq!(
            brittany.codes,
            vec![
                ("ISO3166-2", "FR-BRE"),
                ("ref:INSEE", "53"),
                ("ref:nuts", "FRH;FRH0"),
                ("ref:nuts:1", "FRH"),
                ("ref:nuts:2", "FRH0"),
                ("wikidata", "Q12130")
            ]
            .into_iter()
            .map(|(k, v)| (k.to_string(), v.to_string()))
            .collect()
        )
    }
}
<|endoftext|>


<commit_before>#![allow(
    clippy::cognitive_complexity,
    clippy::inherent_to_string,
    clippy::large_enum_variant,
    clippy::new_without_default,
    clippy::toplevel_ref_arg
)]

mod gen;
mod syntax;

use gen::include;
use std::io::{self, Write};
use std::path::PathBuf;
use structopt::StructOpt;

#[derive(StructOpt, Debug)]
#[structopt(
    name = "cxxbridge",
    author,
    about = "https://github.com/dtolnay/cxx",
    usage = "\
    cxxbridge <input>.rs              Emit .cc file for bridge to stdout
    cxxbridge <input>.rs --header     Emit .h file for bridge to stdout
    cxxbridge --header                Emit rust/cxx.h header to stdout",
    help_message = "Print help information",
    version_message = "Print version information"
)]
struct Opt {
    /// Input Rust source file containing #[cxx::bridge]
    #[structopt(parse(from_os_str), required_unless = "header")]
    input: Option<PathBuf>,

    /// Emit header with declarations only
    #[structopt(long)]
    header: bool,

    /// Any additional headers to #include
    #[structopt(short, long)]
    include: Vec<String>,
}

fn write(content: impl AsRef<[u8]>) {
    let _ = io::stdout().lock().write_all(content.as_ref());
}

fn main() {
    let opt = Opt::from_args();

    let gen = gen::Opt {
        include: opt.include,
    };

    match (opt.input, opt.header) {
        (Some(input), true) => write(gen::do_generate_header(&input, gen)),
        (Some(input), false) => write(gen::do_generate_bridge(&input, gen)),
        (None, true) => write(include::HEADER),
        (None, false) => unreachable!(), // enforced by required_unless
    }
}
<commit_msg>Specify author explicitly.<commit_after>#![allow(
    clippy::cognitive_complexity,
    clippy::inherent_to_string,
    clippy::large_enum_variant,
    clippy::new_without_default,
    clippy::toplevel_ref_arg
)]

mod gen;
mod syntax;

use gen::include;
use std::io::{self, Write};
use std::path::PathBuf;
use structopt::StructOpt;

#[derive(StructOpt, Debug)]
#[structopt(
    name = "cxxbridge",
    author = "dtolnay@gmail.com",
    about = "https://github.com/dtolnay/cxx",
    usage = "\
    cxxbridge <input>.rs              Emit .cc file for bridge to stdout
    cxxbridge <input>.rs --header     Emit .h file for bridge to stdout
    cxxbridge --header                Emit rust/cxx.h header to stdout",
    help_message = "Print help information",
    version_message = "Print version information"
)]
struct Opt {
    /// Input Rust source file containing #[cxx::bridge]
    #[structopt(parse(from_os_str), required_unless = "header")]
    input: Option<PathBuf>,

    /// Emit header with declarations only
    #[structopt(long)]
    header: bool,

    /// Any additional headers to #include
    #[structopt(short, long)]
    include: Vec<String>,
}

fn write(content: impl AsRef<[u8]>) {
    let _ = io::stdout().lock().write_all(content.as_ref());
}

fn main() {
    let opt = Opt::from_args();

    let gen = gen::Opt {
        include: opt.include,
    };

    match (opt.input, opt.header) {
        (Some(input), true) => write(gen::do_generate_header(&input, gen)),
        (Some(input), false) => write(gen::do_generate_bridge(&input, gen)),
        (None, true) => write(include::HEADER),
        (None, false) => unreachable!(), // enforced by required_unless
    }
}
<|endoftext|>


<commit_before>use helpers::*;
use core::cell::Cell;
use core::cmp;

use pm;
use hil::spi_master;
use hil::spi_master::SpiCallback;
use hil::spi_master::ClockPolarity;
use hil::spi_master::ClockPhase;
use hil::spi_master::DataOrder;

// Driver for the SPI hardware (seperate from the USARTS, described in chapter 26 of the
// datasheet)

/// The registers used to interface with the hardware
#[repr(C, packed)]
struct SpiRegisters {
    cr: u32, // 0x0
    mr: u32, // 0x4
    rdr: u32, // 0x8
    tdr: u32, // 0xC
    sr: u32, // 0x10
    ier: u32, // 0x14
    idr: u32, // 0x18
    imr: u32, // 0x1C
    reserved0: [u32; 4], // 0x20, 0x24, 0x28, 0x2C
    csr0: u32, // 0x30
    csr1: u32, // 0x34
    csr2: u32, // 0x38
    csr3: u32, // 0x3C
    reserved1: [u32; 41], // 0x40 - 0xE0
    wpcr: u32, // 0xE4
    wpsr: u32, // 0xE8
    reserved2: [u32; 3], // 0xEC - 0xF4
    features: u32, // 0xF8
    version: u32, // 0xFC
}

const SPI_BASE: u32 = 0x40008000;

/// Values for selected peripherals
#[derive(Copy,Clone)]
pub enum Peripheral {
    Peripheral0,
    Peripheral1,
    Peripheral2,
    Peripheral3,
}

///
/// SPI implementation using the SPI hardware
/// Supports four peripherals. Each peripheral can have different settings.
///
/// The init, read, and write methods act on the currently selected peripheral.
/// The init method can be safely called more than once to configure different peripherals:
///
///     spi.set_active_peripheral(Peripheral::Peripheral0);
///     spi.init(/* Parameters for peripheral 0 */);
///     spi.set_active_peripheral(Peripheral::Peripheral1);
///     spi.init(/* Parameters for peripheral 1 */);
///
pub struct Spi {
    /// Registers
    regs: *mut SpiRegisters,
    /// Client
    callback: Cell<Option<&'static SpiCallback>>,
//    dmaRead:  dma::DMAChannel,
//    dmaWrite: dma::DMAChannel,
}

pub static mut SPI: Spi = Spi::new();

impl Spi {
    /// Creates a new SPI object, with peripheral 0 selected
   const fn new() -> Spi {
        Spi {
            regs: SPI_BASE as *mut SpiRegisters,
            callback: Cell::new(None),
        }
    }

    /// Sets the approximate baud rate for the active peripheral
    ///
    /// Since the only supported baud rates are 48 MHz / n where n is an integer from 1 to 255,
    /// the exact baud rate may not be available. In that case, the next lower baud rate will be
    /// selected.
    ///
    /// The lowest available baud rate is 188235 baud. If the requested rate is lower,
    /// 188235 baud will be selected.
    pub fn set_baud_rate(&'static self, rate: u32) {
        // Main clock frequency
        let mut real_rate = rate;
        let clock = 48000000;

        if real_rate < 188235 {
            real_rate = 188235;
        }
        if real_rate > clock {
            real_rate = clock;
        }

        // Divide and truncate, resulting in a n value that might be too low
        let mut scbr = clock / real_rate;
        // If the division was not exact, increase the n to get a slower baud rate
        if clock % rate != 0 {
            scbr += 1;
        }
        let mut csr = self.read_active_csr();
        let csr_mask: u32 = 0xFFFF00FF;
        // Clear, then write CSR bits
        csr &= csr_mask;
        csr |= (scbr & 0xFF) << 8;
        self.write_active_csr(csr);
    }

    /// Returns the currently active peripheral
    pub fn get_active_peripheral(&self) -> Peripheral {
        let mr = unsafe {volatile_load(&(*self.regs).mr)};
        let pcs = (mr >> 16) & 0xF;
        // Split into bits for matching
        let pcs_bits = ((pcs >> 3) & 1, (pcs >> 2) & 1, (pcs >> 1) & 1, pcs & 1);
        match pcs_bits {
            (_, _, _, 0) => Peripheral::Peripheral0,
            (_, _, 0, 1) => Peripheral::Peripheral1,
            (_, 0, 1, 1) => Peripheral::Peripheral2,
            (0, 1, 1, 1) => Peripheral::Peripheral3,
            _ => {
                // Invalid configuration
                // ???
                Peripheral::Peripheral0
            }
        }
    }

    /// Returns the value of CSR0, CSR1, CSR2, or CSR3, whichever corresponds to the active
    /// peripheral
    fn read_active_csr(&'static self) -> u32 {
        match self.get_active_peripheral() {
            Peripheral::Peripheral0 => unsafe {volatile_load(&(*self.regs).csr0)},
            Peripheral::Peripheral1 => unsafe {volatile_load(&(*self.regs).csr1)},
            Peripheral::Peripheral2 => unsafe {volatile_load(&(*self.regs).csr2)},
            Peripheral::Peripheral3 => unsafe {volatile_load(&(*self.regs).csr3)},
        }
    }
    /// Sets the value of CSR0, CSR1, CSR2, or CSR3, whichever corresponds to the active
    /// peripheral
    fn write_active_csr(&'static self, value: u32) {
        match self.get_active_peripheral() {
            Peripheral::Peripheral0 => unsafe {volatile_store(&mut (*self.regs).csr0, value)},
            Peripheral::Peripheral1 => unsafe {volatile_store(&mut (*self.regs).csr1, value)},
            Peripheral::Peripheral2 => unsafe {volatile_store(&mut (*self.regs).csr2, value)},
            Peripheral::Peripheral3 => unsafe {volatile_store(&mut (*self.regs).csr3, value)},
        };
    }
}

impl spi_master::SpiMaster for Spi {
    fn init(&self, callback: &'static SpiCallback) {
        // Enable clock
        unsafe { pm::enable_clock(pm::Clock::PBA(pm::PBAClock::SPI)); }
        // Configure GPIO pins
        // PA21, PA27, PC04 -> MISO
        // PA22, PA28, PC05 -> MOSI
        // PA23, PA29, PC06 -> SCK
        // PA24, PA30, PC03 -> CS0
        // PA31 , PC02-> CS1
        // PC00 -> CS2
        // PC01 -> CR3
        self.callback.set(Some(callback));
        self.set_rate(8000000); // Set initial baud rate to 8MHz
        self.set_clock(ClockPolarity::IdleLow);
        self.set_phase(ClockPhase::SampleLeading);
        self.set_order(DataOrder::MSBFirst);

        // Indicate the last transfer to disable slave select
        unsafe {volatile_store(&mut (*self.regs).cr, 1 << 24)};

        let mut mode = unsafe {volatile_load(&(*self.regs).mr)};
        mode |= 1; // Enable master mode
        mode |= 1 << 4; // Disable mode fault detection (open drain outputs not supported)
        unsafe {volatile_store(&mut (*self.regs).mr, mode)};
    }

    fn read_write_byte(&'static self, val: u8) -> u8 {
        let tdr = val as u32;
        unsafe {volatile_store(&mut (*self.regs).tdr, tdr)};
        // Wait for receive data register full
        while (unsafe {volatile_load(&(*self.regs).sr)} & 1) != 1 {}
        // Return read value
        unsafe {volatile_load(&(*self.regs).rdr) as u8}
    }

    fn write_byte(&'static self, out_byte: u8) {
        let tdr = out_byte as u32;
        unsafe {volatile_store(&mut (*self.regs).tdr, tdr)};
    }

    fn read_byte(&'static self) -> u8 {
        self.read_write_byte(0)
    }

    /// The write buffer has to be mutable because it's passed back to
    /// the caller, and the caller may want to be able write into it.
    fn read_write_bytes(&'static self,
                        mut read_buffer:  Option<&'static mut [u8]>,
                        write_buffer: Option<&'static mut [u8]>) -> bool {
        // If both are Some, read/write minimum of lengths
        // If only read is Some, read length and write zeroes
        // If only write is Some, write length and discard reads
        // If both are None, return false
        // TODO: Asynchronous
        if read_buffer.is_none() && write_buffer.is_none() {
            return false
        }
        let reading = read_buffer.is_some();
        let writing = write_buffer.is_some();
        let read_len = match read_buffer {
            Some(ref buf) => {buf.len()},
            None          => 0
        };
        let write_len = match write_buffer {
            Some(ref buf) => {buf.len()},
            None          => 0
        };
        let count = if reading && writing {cmp::min(read_len, write_len)}
                    else                  {cmp::max(read_len, write_len)};
        for i in 0..count {
            let mut txbyte: u8 = 0;
            match write_buffer {
                Some(ref buf) => {txbyte = buf[i];}
                None          => {}
            }
            // Write the value
            let rxbyte = self.read_write_byte(txbyte);
            match read_buffer.take() {
                Some(ref mut buf) => {buf[i] = rxbyte;}
                None          => {}
            }
        }
        self.callback.get().map(|cb| {
            cb.read_write_done(read_buffer, write_buffer);
        });
        true
    }

#[allow(unused_variables)]
    fn set_rate(&self, rate: u32) -> u32 { 0 }
    fn get_rate(&self) -> u32 { 0 }

#[allow(unused_variables)]
    fn set_order(&self, order: DataOrder) { }
    fn get_order(&self) -> DataOrder { DataOrder::LSBFirst }

#[allow(unused_variables)]
    fn set_clock(&self, polarity: ClockPolarity) { }
    fn get_clock(&self) -> ClockPolarity { ClockPolarity::IdleLow }

#[allow(unused_variables)]
    fn set_phase(&self, phase: ClockPhase) { }
    fn get_phase(&self) -> ClockPhase { ClockPhase::SampleTrailing }

    /// Sets the active peripheral
    fn set_chip_select(&self, cs: u8) {
        let peripheral_number: u32 = match cs {
            0 => 0b0000,
            1 => 0b0001,
            2 => 0b0011,
            3 => 0b0111,
            _ => 0b0000,
        };

        let mut mr = unsafe {volatile_load(&(*self.regs).mr)};
        // Clear and set MR.PCS
        let pcs_mask: u32 = 0xFFF0FFFF;
        mr &= pcs_mask;
        mr |= peripheral_number << 16;
        unsafe {volatile_store(&mut (*self.regs).mr, mr);}
    }

    fn clear_chip_select(&self) {
       unsafe {volatile_store(&mut (*self.regs).cr, 1 << 24)};
    }
}
<commit_msg>Adding DMA implmentation of SPI.<commit_after>use helpers::*;
use core::cell::Cell;
use core::cmp;

use pm;
use hil::spi_master;
use hil::spi_master::SpiCallback;
use hil::spi_master::ClockPolarity;
use hil::spi_master::ClockPhase;
use hil::spi_master::DataOrder;

// Driver for the SPI hardware (seperate from the USARTS, described in chapter 26 of the
// datasheet)

/// The registers used to interface with the hardware
#[repr(C, packed)]
struct SpiRegisters {
    cr: u32, // 0x0
    mr: u32, // 0x4
    rdr: u32, // 0x8
    tdr: u32, // 0xC
    sr: u32, // 0x10
    ier: u32, // 0x14
    idr: u32, // 0x18
    imr: u32, // 0x1C
    reserved0: [u32; 4], // 0x20, 0x24, 0x28, 0x2C
    csr0: u32, // 0x30
    csr1: u32, // 0x34
    csr2: u32, // 0x38
    csr3: u32, // 0x3C
    reserved1: [u32; 41], // 0x40 - 0xE0
    wpcr: u32, // 0xE4
    wpsr: u32, // 0xE8
    reserved2: [u32; 3], // 0xEC - 0xF4
    features: u32, // 0xF8
    version: u32, // 0xFC
}

const SPI_BASE: u32 = 0x40008000;

/// Values for selected peripherals
#[derive(Copy,Clone)]
pub enum Peripheral {
    Peripheral0,
    Peripheral1,
    Peripheral2,
    Peripheral3,
}

///
/// SPI implementation using the SPI hardware
/// Supports four peripherals. Each peripheral can have different settings.
///
/// The init, read, and write methods act on the currently selected peripheral.
/// The init method can be safely called more than once to configure different peripherals:
///
///     spi.set_active_peripheral(Peripheral::Peripheral0);
///     spi.init(/* Parameters for peripheral 0 */);
///     spi.set_active_peripheral(Peripheral::Peripheral1);
///     spi.init(/* Parameters for peripheral 1 */);
///
pub struct Spi {
    /// Registers
    regs: *mut SpiRegisters,
    /// Client
    callback: Cell<Option<&'static SpiCallback>>,
    dmaRead:  Option<&'static mut DMAChannel>,
    dmaWrite: Option<&'static mut DMAChannel>,
}

pub static mut SPI: Spi = Spi::new();

impl Spi {
    /// Creates a new SPI object, with peripheral 0 selected
   const fn new() -> Spi {
        Spi {
            regs: SPI_BASE as *mut SpiRegisters,
            callback: Cell::new(None),
            dmaRead:  None,
            dmaWrite: None,
        }
    }

    /// Sets the approximate baud rate for the active peripheral
    ///
    /// Since the only supported baud rates are 48 MHz / n where n is an integer from 1 to 255,
    /// the exact baud rate may not be available. In that case, the next lower baud rate will be
    /// selected.
    ///
    /// The lowest available baud rate is 188235 baud. If the requested rate is lower,
    /// 188235 baud will be selected.
    pub fn set_baud_rate(&'static self, rate: u32) {
        // Main clock frequency
        let mut real_rate = rate;
        let clock = 48000000;

        if real_rate < 188235 {
            real_rate = 188235;
        }
        if real_rate > clock {
            real_rate = clock;
        }

        // Divide and truncate, resulting in a n value that might be too low
        let mut scbr = clock / real_rate;
        // If the division was not exact, increase the n to get a slower baud rate
        if clock % rate != 0 {
            scbr += 1;
        }
        let mut csr = self.read_active_csr();
        let csr_mask: u32 = 0xFFFF00FF;
        // Clear, then write CSR bits
        csr &= csr_mask;
        csr |= (scbr & 0xFF) << 8;
        self.write_active_csr(csr);
    }

    /// Returns the currently active peripheral
    pub fn get_active_peripheral(&self) -> Peripheral {
        let mr = unsafe {volatile_load(&(*self.regs).mr)};
        let pcs = (mr >> 16) & 0xF;
        // Split into bits for matching
        let pcs_bits = ((pcs >> 3) & 1, (pcs >> 2) & 1, (pcs >> 1) & 1, pcs & 1);
        match pcs_bits {
            (_, _, _, 0) => Peripheral::Peripheral0,
            (_, _, 0, 1) => Peripheral::Peripheral1,
            (_, 0, 1, 1) => Peripheral::Peripheral2,
            (0, 1, 1, 1) => Peripheral::Peripheral3,
            _ => {
                // Invalid configuration
                // ???
                Peripheral::Peripheral0
            }
        }
    }

    /// Returns the value of CSR0, CSR1, CSR2, or CSR3, whichever corresponds to the active
    /// peripheral
    fn read_active_csr(&'static self) -> u32 {
        match self.get_active_peripheral() {
            Peripheral::Peripheral0 => unsafe {volatile_load(&(*self.regs).csr0)},
            Peripheral::Peripheral1 => unsafe {volatile_load(&(*self.regs).csr1)},
            Peripheral::Peripheral2 => unsafe {volatile_load(&(*self.regs).csr2)},
            Peripheral::Peripheral3 => unsafe {volatile_load(&(*self.regs).csr3)},
        }
    }
    /// Sets the value of CSR0, CSR1, CSR2, or CSR3, whichever corresponds to the active
    /// peripheral
    fn write_active_csr(&'static self, value: u32) {
        match self.get_active_peripheral() {
            Peripheral::Peripheral0 => unsafe {volatile_store(&mut (*self.regs).csr0, value)},
            Peripheral::Peripheral1 => unsafe {volatile_store(&mut (*self.regs).csr1, value)},
            Peripheral::Peripheral2 => unsafe {volatile_store(&mut (*self.regs).csr2, value)},
            Peripheral::Peripheral3 => unsafe {volatile_store(&mut (*self.regs).csr3, value)},
        };
    }
}

impl spi_master::SpiMaster for Spi {
    fn init(&self, callback: &'static SpiCallback) {
        // Enable clock
        unsafe { pm::enable_clock(pm::Clock::PBA(pm::PBAClock::SPI)); }
        // Configure GPIO pins
        // PA21, PA27, PC04 -> MISO
        // PA22, PA28, PC05 -> MOSI
        // PA23, PA29, PC06 -> SCK
        // PA24, PA30, PC03 -> CS0
        // PA31 , PC02-> CS1
        // PC00 -> CS2
        // PC01 -> CR3
        self.callback.set(Some(callback));
        self.set_rate(8000000); // Set initial baud rate to 8MHz
        self.set_clock(ClockPolarity::IdleLow);
        self.set_phase(ClockPhase::SampleLeading);
        self.set_order(DataOrder::MSBFirst);

        // Indicate the last transfer to disable slave select
        unsafe {volatile_store(&mut (*self.regs).cr, 1 << 24)};

        let mut mode = unsafe {volatile_load(&(*self.regs).mr)};
        mode |= 1; // Enable master mode
        mode |= 1 << 4; // Disable mode fault detection (open drain outputs not supported)
        unsafe {volatile_store(&mut (*self.regs).mr, mode)};
    }

    fn read_write_byte(&'static self, val: u8) -> u8 {
        let tdr = val as u32;
        unsafe {volatile_store(&mut (*self.regs).tdr, tdr)};
        // Wait for receive data register full
        while (unsafe {volatile_load(&(*self.regs).sr)} & 1) != 1 {}
        // Return read value
        unsafe {volatile_load(&(*self.regs).rdr) as u8}
    }

    fn write_byte(&'static self, out_byte: u8) {
        let tdr = out_byte as u32;
        unsafe {volatile_store(&mut (*self.regs).tdr, tdr)};
    }

    fn read_byte(&'static self) -> u8 {
        self.read_write_byte(0)
    }

    /// The write buffer has to be mutable because it's passed back to
    /// the caller, and the caller may want to be able write into it.
    fn read_write_bytes(&'static self,
                        mut read_buffer:  Option<&'static mut [u8]>,
                        write_buffer: Option<&'static mut [u8]>) -> bool {
        // If both are Some, read/write minimum of lengths
        // If only read is Some, read length and write zeroes
        // If only write is Some, write length and discard reads
        // If both are None, return false
        // TODO: Asynchronous
        if read_buffer.is_none() && write_buffer.is_none() {
            return false
        }
        let reading = read_buffer.is_some();
        let writing = write_buffer.is_some();
        let read_len = match read_buffer {
            Some(ref buf) => {buf.len()},
            None          => 0
        };
        let write_len = match write_buffer {
            Some(ref buf) => {buf.len()},
            None          => 0
        };
        let count = if reading && writing {cmp::min(read_len, write_len)}
                    else                  {cmp::max(read_len, write_len)};
        for i in 0..count {
            let mut txbyte: u8 = 0;
            match write_buffer {
                Some(ref buf) => {txbyte = buf[i];}
                None          => {}
            }
            // Write the value
            let rxbyte = self.read_write_byte(txbyte);
            match read_buffer.take() {
                Some(ref mut buf) => {buf[i] = rxbyte;}
                None          => {}
            }
        }
        self.callback.get().map(|cb| {
            cb.read_write_done(read_buffer, write_buffer);
        });
        true
    }

#[allow(unused_variables)]
    fn set_rate(&self, rate: u32) -> u32 { 0 }
    fn get_rate(&self) -> u32 { 0 }

#[allow(unused_variables)]
    fn set_order(&self, order: DataOrder) { }
    fn get_order(&self) -> DataOrder { DataOrder::LSBFirst }

#[allow(unused_variables)]
    fn set_clock(&self, polarity: ClockPolarity) { }
    fn get_clock(&self) -> ClockPolarity { ClockPolarity::IdleLow }

#[allow(unused_variables)]
    fn set_phase(&self, phase: ClockPhase) { }
    fn get_phase(&self) -> ClockPhase { ClockPhase::SampleTrailing }

    /// Sets the active peripheral
    fn set_chip_select(&self, cs: u8) {
        let peripheral_number: u32 = match cs {
            0 => 0b0000,
            1 => 0b0001,
            2 => 0b0011,
            3 => 0b0111,
            _ => 0b0000,
        };

        let mut mr = unsafe {volatile_load(&(*self.regs).mr)};
        // Clear and set MR.PCS
        let pcs_mask: u32 = 0xFFF0FFFF;
        mr &= pcs_mask;
        mr |= peripheral_number << 16;
        unsafe {volatile_store(&mut (*self.regs).mr, mr);}
    }

    fn clear_chip_select(&self) {
       unsafe {volatile_store(&mut (*self.regs).cr, 1 << 24)};
    }
}
<|endoftext|>


<commit_before>use Thingie as Base;
use nalgebra::{Pnt2, Vec2, Cross, Dot};
use super::{Intersect, Intersection};

pub struct Circle<S: Base> {
    pub pos: Pnt2<S>,
    pub radius: S
}

///Lines go from a to b
pub struct Line<S: Base> {
    pub a: Pnt2<S>,
    pub b: Pnt2<S>
}

impl<S> Line<S> where S: Base {
    pub fn new(a: Pnt2<S>, b: Pnt2<S>) -> Line<S> {
        Line {
            a: a,
            b: b
        }
    }
}

impl<S> Intersect<Line<S>, S> for Line<S> where S: Base {
    fn intersection(&self, other: &Line<S>) -> Intersection<S> {
        let r = Vec2::new(self.b.x - self.a.x, self.b.y - self.a.y);
        let s = Vec2::new(other.b.x - other.a.x, other.b.y - other.a.y);

        //cross returns a Vec1... wtf?!
        let unum = (other.a - self.a).cross(&r).x;
        let denom = r.cross(&s).x;

        if unum == S::zero() && denom == S::zero() {
            let t0 = (other.a - self.a).dot(&r) / r.dot(&r);
            let t1 = t0 + s.dot(&r) / r.dot(&r);
            //do they overlap
            // s and r go in opposite directions so s dot r < 0 it's t1 -> t0
            if s.dot(&r) < S::zero() && t1 <= S::one() && S::zero() <= t0 {
                return Intersection::Overlap(self.a + (r * t1), self.a + (r * t0))
            } else if t0 <= S::one() && S::zero() <= t1 {
                return Intersection::Overlap(self.a + (r * t0), self.a + (r * t1))
            }

            return Intersection::Outside
        }

        if denom == S::zero() {
            //lines are parallel
            return Intersection::Parallel
        }

        let u = unum / denom;
        let t = (other.a - self.a).cross(&s).x / denom;

        if t >= S::zero() && t <= S::one() && u >= S::zero() && u <= S::one() {
            return Intersection::Intersects(self.a + (r * t), None)
        }

        return Intersection::Outside
    }

    fn intersects(&self, other: &Line<S>) -> bool {
        match self.intersection(other) {
            Intersection::Overlap(_, _) => true,
            Intersection::Intersects(_, _) => true,
            _ => false
        }
    }

    fn contains(&self, other: &Line<S>) -> bool {
        false
    }
}

#[test]
fn line_intersects() {
    // no intersect
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(2.0, 8.0))
               .intersects(&Line::new(Pnt2::new(8.0, 0.0), Pnt2::new(0.0, 20.0))), false);
    // intersect
    assert_eq!(Line::new(Pnt2::new(0.0,10.0), Pnt2::new(2.0, 0.0))
               .intersects(&Line::new(Pnt2::new(10.0, 0.0), Pnt2::new(0.0, 5.0))), true);
    // parallel, vertical
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(0.0, 10.0))
               .intersects(&Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(2.0, 10.0))), false);
    // parallel, diagonal
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0))
               .intersects(&Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(7.0, 5.0))), false);
    //collinear, overlap
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0))
               .intersects(&Line::new(Pnt2::new(2.0, 2.0), Pnt2::new(7.0, 7.0))), true);
    //collinear, no overlap
    assert_eq!(Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0))
               .intersects(&Line::new(Pnt2::new(7.0, 7.0), Pnt2::new(10.0, 10.0))), false);
}

#[test]
fn line_intersection() {
    let mut l1 = Line::new(Pnt2::new(-5.0, 5.0), Pnt2::new(-1.0, 5.0));
    let mut l2 = Line::new(Pnt2::new(0.0, -5.0), Pnt2::new(0.0, 5.0));

    // no intersection
    assert_eq!(l1.intersection(&l2), Intersection::Outside);

    l1.b.x = 5.0;
    // intersection
    assert_eq!(l1.intersection(&l2), Intersection::Intersects(Pnt2 {x: 0.0, y: 5.0}, None));

    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(0.0, 10.0));
    l2 = Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(2.0, 10.0));
    // parallel, vertical
    assert_eq!(l1.intersection(&l2), Intersection::Parallel);

    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0));
    l2 = Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(7.0, 5.0));

    // parallel, diagonal
    assert_eq!(l1.intersection(&l2), Intersection::Parallel);

    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0));
    l2 = Line::new(Pnt2::new(2.0, 2.0), Pnt2::new(7.0, 7.0));

    //collinear, overlap
    assert_eq!(l1.intersection(&l2), Intersection::Overlap(l2.a.clone(), l2.b.clone()));
    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0));
    l2 = Line::new(Pnt2::new(7.0, 7.0), Pnt2::new(10.0, 10.0));

    //collinear, no overlap
    assert_eq!(l1.intersection(&l2), Intersection::Outside);

    l1.b.x = 15.0;
    l1.b.y = 15.0;
    //it is right but there's stupid floating point mess
    assert_eq!(l1.intersection(&l2), Intersection::Overlap(Pnt2 {x: 7.0, y: 7.0 }, Pnt2 {x: 10.000000000000002, y: 10.000000000000002}));
}

<commit_msg>Circle/Line intersection<commit_after>use Thingie as Base;
use nalgebra::{Pnt2, Vec2, Cross, Dot, POrd};
use super::{Intersect, Intersection};

pub struct Circle<S: Base> {
    pub pos: Pnt2<S>,
    pub radius: S
}

///Lines go from a to b
pub struct Line<S: Base> {
    pub a: Pnt2<S>,
    pub b: Pnt2<S>
}

impl<S> Line<S> where S: Base {
    pub fn new(a: Pnt2<S>, b: Pnt2<S>) -> Line<S> {
        Line {
            a: a,
            b: b
        }
    }
}

impl<S> Circle<S> where S: Base {
    pub fn new(pos: Pnt2<S>, radius: S) -> Circle<S> {
        Circle {
            pos: pos,
            radius: radius
        }
    }
}

impl<S> Intersect<Line<S>, S> for Line<S> where S: Base {
    fn intersection(&self, other: &Line<S>) -> Intersection<S> {
        let r = Vec2::new(self.b.x - self.a.x, self.b.y - self.a.y);
        let s = Vec2::new(other.b.x - other.a.x, other.b.y - other.a.y);

        //cross returns a Vec1... wtf?!
        let unum = (other.a - self.a).cross(&r).x;
        let denom = r.cross(&s).x;

        if unum == S::zero() && denom == S::zero() {
            let t0 = (other.a - self.a).dot(&r) / r.dot(&r);
            let t1 = t0 + s.dot(&r) / r.dot(&r);
            //do they overlap
            // s and r go in opposite directions so s dot r < 0 it's t1 -> t0
            if s.dot(&r) < S::zero() && t1 <= S::one() && S::zero() <= t0 {
                return Intersection::Overlap(self.a + (r * t1), self.a + (r * t0))
            } else if t0 <= S::one() && S::zero() <= t1 {
                return Intersection::Overlap(self.a + (r * t0), self.a + (r * t1))
            }

            return Intersection::Outside
        }

        if denom == S::zero() {
            //lines are parallel
            return Intersection::Parallel
        }

        let u = unum / denom;
        let t = (other.a - self.a).cross(&s).x / denom;

        if t >= S::zero() && t <= S::one() && u >= S::zero() && u <= S::one() {
            return Intersection::Intersects(self.a + (r * t), None)
        }

        return Intersection::Outside
    }

    fn intersects(&self, other: &Line<S>) -> bool {
        match self.intersection(other) {
            Intersection::Overlap(_, _) => true,
            Intersection::Intersects(_, _) => true,
            _ => false
        }
    }

    fn contains(&self, other: &Line<S>) -> bool {
        //this function should probably just return false
        match self.intersection(other) {
            Intersection::Overlap(a, b) => {
                a.x >= self.a.x && a.y >= self.a.y && b.x <= self.b.x && b.y <= self.b.y
            },
            _ => false
        }
    }
}

impl<S> Intersect<Line<S>, S> for Circle<S> where S: Base {
    fn intersection(&self, other: &Line<S>) -> Intersection<S> {
        let two = S::from(2.0).unwrap();
        let d = (other.b - other.a);

        if d.x == S::zero() && d.y == S::zero() { //line has zero length
            //I could be an asshole and let it pretend it's a tangent and return NaNs,
            //I guess I'm just that nice
            return self.intersection(&other.a)
        }

        let f = (other.a - self.pos);
        let a = d.dot(&d);
        let b = S::from(2.0).unwrap()*f.dot(&d);
        let c = f.dot(&f) - self.radius.powf(two);

        let mut dis = b.powf(two) - (S::from(4.0).unwrap() * a * c);
        if dis < S::zero() {
            return Intersection::Outside
        }
        else if dis == S::zero() {
            dis = dis.sqrt();

            let t0 = (-b - dis)/(two * a);
            let t1 = (-b + dis)/(two * a);
            println!("zero\t{}, {}", t0, t1);
            return Intersection::Intersects(other.a + (d * t0), None)
        }
        else {
            dis = dis.sqrt();

            let t0 = (-b - dis)/(two * a);
            let t1 = (-b + dis)/(two * a);

            println!(">zero\t{}, {}", t0, t1);
            if t0 < S::zero() && t1 <= S::one() { //A inside circle
                return Intersection::Intersects(other.a + (d * t1), None)
            }
            else if t1 > S::one() && t0 >= S::zero() { //B inside circle
                return Intersection::Intersects(other.a + (d * t0), None)
            }
            else if t1 > S::one() && t0 < S::zero() { //all inside
                return Intersection::Inside
            }
            else if t0 > S::zero() && t0 < S::one() && t1 > S::zero() && t1 < S::one() { //two intersections
                return Intersection::Intersects(other.a + (d * t0), Some(other.a + (d * t1)))
            }
        }
        Intersection::Outside
    }

    fn intersects(&self, other: &Line<S>) -> bool {
        match self.intersection(other) {
            Intersection::Intersects(_, _) => true,
            _ => false
        }
    }

    fn contains(&self, other: &Line<S>) -> bool {
        match self.intersection(other) {
            Intersection::Inside => true,
            _ => false
        }
    }
}

impl<S> Intersect<Pnt2<S>, S> for Circle<S> where S: Base {
    fn intersection(&self, other: &Pnt2<S>) -> Intersection<S> {
        let d = *other - self.pos;

        if d.x.abs() < self.radius && d.y.abs() < self.radius {
            return Intersection::Inside
        }
        else if d.x.abs() > self.radius && d.y.abs() > self.radius {
            return Intersection::Outside
        }
        else {
            return Intersection::Intersects(other.clone(), None)
        }
    }

    fn intersects(&self, other: &Pnt2<S>) -> bool {
        let d = *other - self.pos;
        //should probably add an epsilon-y type check here but eh
        d.x.abs() == self.radius && d.y.abs() == self.radius
    }

    fn contains(&self, other: &Pnt2<S>) -> bool {
        match self.intersection(other) {
            Intersection::Inside => true,
            _ => false
        }
    }
}

#[test]
fn line_intersection() {
    let mut l1 = Line::new(Pnt2::new(-5.0, 5.0), Pnt2::new(-1.0, 5.0));
    let mut l2 = Line::new(Pnt2::new(0.0, -5.0), Pnt2::new(0.0, 5.0));

    // no intersection
    assert_eq!(l1.intersection(&l2), Intersection::Outside);

    l1.b.x = 5.0;
    // intersection
    assert_eq!(l1.intersection(&l2), Intersection::Intersects(Pnt2 {x: 0.0, y: 5.0}, None));

    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(0.0, 10.0));
    l2 = Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(2.0, 10.0));
    // parallel, vertical
    assert_eq!(l1.intersection(&l2), Intersection::Parallel);

    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0));
    l2 = Line::new(Pnt2::new(2.0, 0.0), Pnt2::new(7.0, 5.0));

    // parallel, diagonal
    assert_eq!(l1.intersection(&l2), Intersection::Parallel);

    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0));
    l2 = Line::new(Pnt2::new(2.0, 2.0), Pnt2::new(7.0, 7.0));

    //collinear, overlap
    assert_eq!(l1.intersection(&l2), Intersection::Overlap(l2.a.clone(), l2.b.clone()));
    l1 = Line::new(Pnt2::new(0.0,0.0), Pnt2::new(5.0, 5.0));
    l2 = Line::new(Pnt2::new(7.0, 7.0), Pnt2::new(10.0, 10.0));

    //collinear, no overlap
    assert_eq!(l1.intersection(&l2), Intersection::Outside);

    l1.b.x = 15.0;
    l1.b.y = 15.0;
    //it is right but there's stupid floating point mess
    assert_eq!(l1.intersection(&l2), Intersection::Overlap(Pnt2 {x: 7.0, y: 7.0 }, Pnt2 {x: 10.000000000000002, y: 10.000000000000002}));
    assert_eq!(l1.contains(&l2), true)
}

#[test]
fn circle_line_intersection() {
    let circle = Circle::new(Pnt2::new(3.0,1.0), 6.0);
    let mut line = Line::new(Pnt2::new(-3.0, -3.0), Pnt2::new(-3.0, 3.0));

    //tangent
    assert_eq!(circle.intersection(&line), Intersection::Intersects(Pnt2 {x: -3.0, y: 1.0}, None));

    line.a = Pnt2::new(-6.0, 1.0);
    line.b = Pnt2::new(12.0, 1.0);

    //intersection A -> B
    assert_eq!(circle.intersection(&line), Intersection::Intersects(Pnt2 {x: -3.0, y: 1.0}, Some(Pnt2 {x: 9.0, y: 1.0})));

    let temp = line.a;
    line.a = line.b;
    line.b = temp;

    //intersecion B -> A
    assert_eq!(circle.intersection(&line), Intersection::Intersects(Pnt2 {x: 9.0, y: 1.0}, Some(Pnt2 {x: -3.0, y: 1.0})));

    line.b = line.a;
    line.a = circle.pos.clone();
    //intersection A inside circle
    assert_eq!(circle.intersection(&line), Intersection::Intersects(Pnt2 {x: 9.0, y: 1.0}, None));

    line.a = Pnt2::new(-6.0, 1.0);
    line.b = circle.pos;

    //intersection B inside circle
    assert_eq!(circle.intersection(&line), Intersection::Intersects(Pnt2 {x: -3.0, y: 1.0}, None));

    line.a = circle.pos + (circle.radius/2.0);
    line.b = circle.pos;

    //line completely inside
    assert_eq!(circle.intersection(&line), Intersection::Inside);

    line.a = Pnt2::new(-20.0, 20.0);
    line.b = Pnt2::new(20.0, 20.0);

    //line completely outside
    assert_eq!(circle.intersection(&line), Intersection::Outside);

    line.b = line.a.clone();
    //zero length line outside
    assert_eq!(circle.intersection(&line), Intersection::Outside);

    line.a = circle.pos.clone();
    line.b = line.a.clone();

    //zero length line inside
    assert_eq!(circle.intersection(&line), Intersection::Inside);
}

<|endoftext|>


<commit_before>use std::fmt;
use std::hash;

use num::{FromPrimitive};

/// please don't access VA.0 directly.
/// its provided so you can construct VA like:
/// ```
/// use lancelot::arch::VA;
/// let _ = VA(0x0);
/// ```
#[derive(Copy, Clone)]
pub struct VA(pub u64);

impl VA {
    /// Compute the VA given a delta RVA, like `self:va + other:rva`.
    /// Returns None on over/underflow.
    pub fn va(&self, other: RVA) -> Option<VA> {
        if other.0 < 0 {
            Some(VA(match self.0.checked_sub((-other.0) as u64) {
                None => return None,
                Some(v) => v,
            }))
        } else {
            Some(VA(match self.0.checked_add(other.0 as u64) {
                None => return None,
                Some(v) => v,
            }))
        }
    }

    /// Compute the delta RVA from `self:va - other:va`.
    /// Returns None on over/underflow.
    pub fn rva(&self, other: VA) -> Option<RVA> {
        if other.0 > self.0 {
            let v = match other.0.checked_sub(self.0) {
                None => return None,
                Some(v) => v,
            };
            if v > std::i64::MAX as u64{
                return None;
            }
            Some(RVA(-(v as i64)))
        } else {
            let v = match self.0.checked_sub(other.0) {
                None => return None,
                Some(v) => v
            };
            if v > std::i64::MAX as u64 {
                return None;
            }
            Some(RVA(v as i64))
        }
    }
}

impl fmt::Display for VA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:#x}", self.0)
    }
}

impl fmt::Debug for VA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:#x}", self.0)
    }
}

impl fmt::LowerHex for VA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:x}", self.0)
    }
}

impl fmt::UpperHex for VA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:X}", self.0)
    }
}

impl hash::Hash for VA {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        self.0.hash(state);
    }
}

impl PartialEq for VA {
    fn eq(&self, other: &VA) -> bool {
        self.0 == other.0
    }
}

impl Eq for VA {}

impl std::convert::From<usize> for VA {
    fn from(v: usize) -> VA {
        // this should only fail when we get to 128-bit architectures
        VA(v as u64)
    }
}

impl std::convert::From<u32> for VA {
    fn from(v: u32) -> VA {
        VA(v as u64)
    }
}

impl std::convert::From<u64> for VA {
    fn from(v: u64) -> VA {
        VA(v)
    }
}

impl std::convert::Into<u64> for VA {
    fn into(self) -> u64 {
        self.0
    }
}

/// please don't access RVA.0 directly.
/// /// its provided so you can construct RVA like:
/// ```
/// use lancelot::arch::RVA;
/// let _ = RVA(0x0);
/// ```
#[derive(Copy, Clone)]
pub struct RVA(pub i64);

impl fmt::Display for RVA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:#x}", self.0)
    }
}

impl fmt::Debug for RVA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:#x}", self.0)
    }
}

impl fmt::LowerHex for RVA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:x}", self.0)
    }
}

impl fmt::UpperHex for RVA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:X}", self.0)
    }
}

impl hash::Hash for RVA {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        self.0.hash(state);
    }
}

impl PartialEq for RVA {
    fn eq(&self, other: &RVA) -> bool {
        self.0 == other.0
    }
}

impl PartialOrd for RVA {
    fn partial_cmp(&self, other: &RVA) -> Option<std::cmp::Ordering> {
        Some(self.0.cmp(&other.0))
    }
}

impl Ord for RVA {
    fn cmp(&self, other: &RVA) -> std::cmp::Ordering {
        self.0.cmp(&other.0)
    }
}

impl Eq for RVA {}

impl std::convert::From<usize> for RVA {
    /// may panic on overflow
    fn from(v: usize) -> RVA {
        RVA(i64::from_usize(v).expect("usize too large for RVA"))
    }
}

impl std::convert::Into<usize> for RVA {
    /// may panic on over/underflow
    fn into(self) -> usize {
        if self.0 < 0 {
            panic!("usize underflow")
        }
        if self.0 as u64 > std::usize::MAX as u64 {
            panic!("usize overflow");
        }
        self.0 as usize
    }
}

impl std::convert::Into<u64> for RVA {
    fn into(self) -> u64 {
        self.0 as u64
    }
}

impl std::convert::Into<i64> for RVA {
    fn into(self) -> i64 {
        self.0
    }
}

impl std::convert::From<u32> for RVA {
    fn from(v: u32) -> RVA {
        RVA(v as i64)
    }
}

impl std::convert::From<i32> for RVA {
    fn from(v: i32) -> RVA {
        RVA(v as i64)
    }
}

impl std::convert::From<i64> for RVA {
    fn from(v: i64) -> RVA {
        RVA(v)
    }
}

impl std::ops::Add<RVA> for RVA {
    type Output = RVA;

    /// may panic on overflow
    fn add(self, rhs: RVA) -> RVA {
        RVA(self.0.checked_add(rhs.0).expect("rva overflow"))
    }
}

impl std::ops::Add<usize> for RVA {
    type Output = RVA;

    /// may panic on overflow
    fn add(self, rhs: usize) -> RVA {
        self + RVA::from(rhs)
    }
}

impl std::ops::Add<u8> for RVA {
    type Output = RVA;

    /// may panic on overflow
    fn add(self, rhs: u8) -> RVA {
        self + RVA::from(rhs as i64)
    }
}

impl std::ops::Add<i32> for RVA {
    type Output = RVA;

    /// may panic on overflow
    fn add(self, rhs: i32) -> RVA {
        self + RVA::from(rhs as i64)
    }
}

impl std::ops::Sub<RVA> for RVA {
    type Output = RVA;

    /// may panic on overflow
    fn sub(self, rhs: RVA) -> RVA {
        RVA(self.0.checked_sub(rhs.0).expect("rva underflow"))
    }
}

#[derive(Copy, Clone)]
pub enum Arch {
    X32,
    X64
}

impl Arch {
    pub fn get_pointer_size(self) -> u8 {
        match self {
            Arch::X32 => 4,
            Arch::X64 => 8,
        }
    }
}

impl fmt::Display for Arch {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Arch::X32 => write!(f, "x32"),
            Arch::X64 => write!(f, "x64"),
        }
    }
}

impl fmt::Debug for Arch {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self{
            Arch::X32 => write!(f, "x32"),
            Arch::X64 => write!(f, "x64"),
        }
    }
}

<commit_msg>arch: rva: consolidate into/from<commit_after>use std::fmt;
use std::hash;

use num::{FromPrimitive};

/// please don't access VA.0 directly.
/// its provided so you can construct VA like:
/// ```
/// use lancelot::arch::VA;
/// let _ = VA(0x0);
/// ```
#[derive(Copy, Clone)]
pub struct VA(pub u64);

impl VA {
    /// Compute the VA given a delta RVA, like `self:va + other:rva`.
    /// Returns None on over/underflow.
    pub fn va(&self, other: RVA) -> Option<VA> {
        if other.0 < 0 {
            Some(VA(match self.0.checked_sub((-other.0) as u64) {
                None => return None,
                Some(v) => v,
            }))
        } else {
            Some(VA(match self.0.checked_add(other.0 as u64) {
                None => return None,
                Some(v) => v,
            }))
        }
    }

    /// Compute the delta RVA from `self:va - other:va`.
    /// Returns None on over/underflow.
    pub fn rva(&self, other: VA) -> Option<RVA> {
        if other.0 > self.0 {
            let v = match other.0.checked_sub(self.0) {
                None => return None,
                Some(v) => v,
            };
            if v > std::i64::MAX as u64{
                return None;
            }
            Some(RVA(-(v as i64)))
        } else {
            let v = match self.0.checked_sub(other.0) {
                None => return None,
                Some(v) => v
            };
            if v > std::i64::MAX as u64 {
                return None;
            }
            Some(RVA(v as i64))
        }
    }
}

impl fmt::Display for VA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:#x}", self.0)
    }
}

impl fmt::Debug for VA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:#x}", self.0)
    }
}

impl fmt::LowerHex for VA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:x}", self.0)
    }
}

impl fmt::UpperHex for VA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:X}", self.0)
    }
}

impl hash::Hash for VA {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        self.0.hash(state);
    }
}

impl PartialEq for VA {
    fn eq(&self, other: &VA) -> bool {
        self.0 == other.0
    }
}

impl Eq for VA {}

impl std::convert::From<usize> for VA {
    fn from(v: usize) -> VA {
        // this should only fail when we get to 128-bit architectures
        VA(v as u64)
    }
}

impl std::convert::From<u32> for VA {
    fn from(v: u32) -> VA {
        VA(v as u64)
    }
}

impl std::convert::From<u64> for VA {
    fn from(v: u64) -> VA {
        VA(v)
    }
}

impl std::convert::Into<u64> for VA {
    fn into(self) -> u64 {
        self.0
    }
}

/// please don't access RVA.0 directly.
/// /// its provided so you can construct RVA like:
/// ```
/// use lancelot::arch::RVA;
/// let _ = RVA(0x0);
/// ```
#[derive(Copy, Clone)]
pub struct RVA(pub i64);

impl fmt::Display for RVA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:#x}", self.0)
    }
}

impl fmt::Debug for RVA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:#x}", self.0)
    }
}

impl fmt::LowerHex for RVA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:x}", self.0)
    }
}

impl fmt::UpperHex for RVA {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:X}", self.0)
    }
}

impl hash::Hash for RVA {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        self.0.hash(state);
    }
}

impl PartialEq for RVA {
    fn eq(&self, other: &RVA) -> bool {
        self.0 == other.0
    }
}

impl PartialOrd for RVA {
    fn partial_cmp(&self, other: &RVA) -> Option<std::cmp::Ordering> {
        Some(self.0.cmp(&other.0))
    }
}

impl Ord for RVA {
    fn cmp(&self, other: &RVA) -> std::cmp::Ordering {
        self.0.cmp(&other.0)
    }
}

impl Eq for RVA {}

impl std::convert::Into<usize> for RVA {
    /// may panic on over/underflow
    fn into(self) -> usize {
        if self.0 < 0 {
            panic!("usize underflow")
        }
        if self.0 as u64 > std::usize::MAX as u64 {
            panic!("usize overflow");
        }
        self.0 as usize
    }
}

impl std::convert::Into<u64> for RVA {
    fn into(self) -> u64 {
        self.0 as u64
    }
}

impl std::convert::Into<i64> for RVA {
    fn into(self) -> i64 {
        self.0
    }
}

impl std::convert::From<u32> for RVA {
    fn from(v: u32) -> RVA {
        RVA(v as i64)
    }
}

impl std::convert::From<i32> for RVA {
    fn from(v: i32) -> RVA {
        RVA(v as i64)
    }
}

impl std::convert::From<i64> for RVA {
    fn from(v: i64) -> RVA {
        RVA(v)
    }
}

impl std::convert::From<usize> for RVA {
    /// may panic on overflow
    fn from(v: usize) -> RVA {
        RVA(i64::from_usize(v).expect("usize too large for RVA"))
    }
}

impl std::ops::Add<RVA> for RVA {
    type Output = RVA;

    /// may panic on overflow
    fn add(self, rhs: RVA) -> RVA {
        RVA(self.0.checked_add(rhs.0).expect("rva overflow"))
    }
}

impl std::ops::Add<u8> for RVA {
    type Output = RVA;

    /// may panic on overflow
    fn add(self, rhs: u8) -> RVA {
        self + RVA::from(rhs as i64)
    }
}

impl std::ops::Add<u16> for RVA {
    type Output = RVA;

    /// may panic on overflow
    fn add(self, rhs: u16) -> RVA {
        self + RVA::from(rhs as i64)
    }
}

impl std::ops::Add<i32> for RVA {
    type Output = RVA;

    /// may panic on overflow
    fn add(self, rhs: i32) -> RVA {
        self + RVA::from(rhs as i64)
    }
}

impl std::ops::Add<u32> for RVA {
    type Output = RVA;

    /// may panic on overflow
    fn add(self, rhs: u32) -> RVA {
        self + RVA::from(rhs as i64)
    }
}

impl std::ops::Add<usize> for RVA {
    type Output = RVA;

    /// may panic on overflow
    fn add(self, rhs: usize) -> RVA {
        self + RVA::from(rhs)
    }
}

impl std::ops::Sub<RVA> for RVA {
    type Output = RVA;

    /// may panic on overflow
    fn sub(self, rhs: RVA) -> RVA {
        RVA(self.0.checked_sub(rhs.0).expect("rva underflow"))
    }
}

#[derive(Copy, Clone)]
pub enum Arch {
    X32,
    X64
}

impl Arch {
    pub fn get_pointer_size(self) -> u8 {
        match self {
            Arch::X32 => 4,
            Arch::X64 => 8,
        }
    }
}

impl fmt::Display for Arch {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Arch::X32 => write!(f, "x32"),
            Arch::X64 => write!(f, "x64"),
        }
    }
}

impl fmt::Debug for Arch {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self{
            Arch::X32 => write!(f, "x32"),
            Arch::X64 => write!(f, "x64"),
        }
    }
}

<|endoftext|>


<commit_before>//< Helpers that are useful for tests and doctests.
use crate::{
    arch::Arch,
    aspace::{AddressSpace, RelativeAddressSpace},
    module::{Module, Permissions, Section},
    RVA, VA,
};

/// configure a global logger at level==DEBUG.
pub fn init_logging() {
    let log_level = log::LevelFilter::Debug;
    fern::Dispatch::new()
        .format(move |out, message, record| {
            out.finish(format_args!(
                "{} [{:5}] {} {}",
                chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
                record.level(),
                if log_level == log::LevelFilter::Trace {
                    record.target()
                } else {
                    ""
                },
                message
            ))
        })
        .level(log_level)
        .chain(std::io::stderr())
        .filter(|metadata| !metadata.target().starts_with("goblin::pe"))
        .apply()
        .expect("failed to configure logging");
}

/// this is for testing, so will panic on error.
pub fn load_shellcode(arch: Arch, buf: &[u8]) -> Module {
    let mut address_space = RelativeAddressSpace::with_capacity(buf.len() as u64);
    address_space.map.writezx(0x0, buf).unwrap();

    Module {
        arch,
        sections: vec![Section {
            name:           "shellcode".to_string(),
            permissions:    Permissions::RWX,
            physical_range: std::ops::Range {
                start: 0x0,
                end:   buf.len() as RVA,
            },
            virtual_range:  std::ops::Range {
                start: 0x0,
                end:   buf.len() as RVA,
            },
        }],
        address_space: address_space.into_absolute(0x0).unwrap(),
    }
}

/// this is for testing, so will panic on error.
pub fn load_shellcode32(buf: &[u8]) -> Module {
    load_shellcode(Arch::X32, buf)
}

/// this is for testing, so will panic on error.
pub fn load_shellcode64(buf: &[u8]) -> Module {
    load_shellcode(Arch::X64, buf)
}

/// this is for testing, so will panic on error.
pub fn read_insn(module: &Module, va: VA) -> zydis::DecodedInstruction {
    use crate::analysis::dis;

    let decoder = dis::get_disassembler(module).unwrap();
    let mut insn_buf = [0u8; 16];
    module.address_space.read_into(va, &mut insn_buf).unwrap();
    decoder.decode(&insn_buf).unwrap().unwrap()
}

pub fn emu_from_shellcode64(code: &[u8]) -> crate::emu::Emulator {
    let m = load_shellcode64(code);
    let mut emu = crate::emu::Emulator::from_module(&m);
    emu.reg.rip = m.address_space.base_address; // 0x0

    emu.mem.mmap(0x5000, 0x2000, Permissions::RW).unwrap();
    emu.reg.rsp = 0x6000;
    emu.reg.rbp = 0x6000;

    emu
}

pub fn emu_from_shellcode32(code: &[u8]) -> crate::emu::Emulator {
    let m = load_shellcode32(code);
    let mut emu = crate::emu::Emulator::from_module(&m);
    emu.reg.set_eip(m.address_space.base_address as u32); // 0x0

    emu.mem.mmap(0x5000, 0x2000, Permissions::RW).unwrap();
    emu.reg.set_esp(0x6000);
    emu.reg.set_ebp(0x6000);

    emu
}

#[cfg(test)]
pub mod uc {
    use unicorn::Cpu;

    use super::load_shellcode64;
    use crate::{
        aspace::AddressSpace,
        emu::{mmu::PAGE_SIZE, reg::STATUS_MASK},
        module::{Module, Permissions},
    };

    pub struct Uc {
        emu: unicorn::CpuX86,
    }

    impl Uc {
        pub fn from_module(m: &Module) -> Uc {
            let emu = match m.arch {
                crate::arch::Arch::X32 => unicorn::CpuX86::new(unicorn::Mode::MODE_32).unwrap(),
                crate::arch::Arch::X64 => unicorn::CpuX86::new(unicorn::Mode::MODE_64).unwrap(),
            };

            for section in m.sections.iter() {
                let mut page_addr = section.virtual_range.start;

                let section_size = section.virtual_range.end - section.virtual_range.start;
                emu.mem_map(
                    section.virtual_range.start,
                    crate::util::align(section_size, PAGE_SIZE as u64) as usize,
                    unicorn::PROT_WRITE,
                )
                .unwrap();

                while page_addr < section.virtual_range.end {
                    let mut page = [0u8; PAGE_SIZE];

                    // AddressSpace currently allows non-page-aligned sizes.
                    let page_data = if page_addr + PAGE_SIZE as u64 > section.virtual_range.end {
                        &mut page[..(section.virtual_range.end - page_addr) as usize]
                    } else {
                        &mut page[..]
                    };

                    m.address_space.read_into(page_addr, page_data).unwrap();
                    emu.mem_write(page_addr, &page[..]).unwrap();
                    page_addr += PAGE_SIZE as u64;
                }

                let mut prot = unicorn::PROT_NONE;
                if section.permissions.intersects(Permissions::W) {
                    prot.insert(unicorn::PROT_WRITE);
                }
                if section.permissions.intersects(Permissions::R) {
                    prot.insert(unicorn::PROT_READ);
                }
                if section.permissions.intersects(Permissions::X) {
                    prot.insert(unicorn::PROT_EXEC);
                }

                emu.mem_protect(
                    section.virtual_range.start,
                    crate::util::align(section_size, PAGE_SIZE as u64) as usize,
                    prot,
                )
                .unwrap();
            }

            Uc { emu }
        }

        pub fn step(&mut self) -> Result<(), unicorn::Error> {
            let rip = self.emu.reg_read(unicorn::RegisterX86::RIP).unwrap();
            self.emu.emu_start(rip, u64::MAX, 0, 1)
        }

        /// panics if the given emulator does have the same state as this.
        pub fn check(&self, other: &crate::emu::Emulator) {
            use unicorn::RegisterX86::*;

            assert_eq!(self.emu.reg_read(RIP).unwrap(), other.reg.rip(), "register: rip",);

            assert_eq!(self.emu.reg_read(RSP).unwrap(), other.reg.rsp(), "register: rsp",);

            assert_eq!(self.emu.reg_read(RBP).unwrap(), other.reg.rbp(), "register: rbp",);

            assert_eq!(self.emu.reg_read(RAX).unwrap(), other.reg.rax(), "register: rax",);

            assert_eq!(self.emu.reg_read(RBX).unwrap(), other.reg.rbx(), "register: rbx",);

            assert_eq!(self.emu.reg_read(RCX).unwrap(), other.reg.rcx(), "register: rcx",);

            assert_eq!(self.emu.reg_read(RDX).unwrap(), other.reg.rdx(), "register: rdx",);

            assert_eq!(self.emu.reg_read(RSI).unwrap(), other.reg.rsi(), "register: rsi",);

            assert_eq!(self.emu.reg_read(RDI).unwrap(), other.reg.rdi(), "register: rdi",);

            assert_eq!(
                self.emu.reg_read(EFLAGS).unwrap() as u64 & STATUS_MASK,
                other.reg.rflags() & STATUS_MASK,
                "flags"
            );
        }
    }

    pub fn uc_from_shellcode64(code: &[u8]) -> Uc {
        let m = load_shellcode64(code);
        let mut uc = Uc::from_module(&m);

        uc.emu
            .reg_write(unicorn::RegisterX86::RIP, m.address_space.base_address)
            .unwrap(); // 0x0

        uc.emu.mem_map(0x5000, 0x2000, unicorn::PROT_ALL).unwrap();
        uc.emu.reg_write(unicorn::RegisterX86::RSP, 0x6000).unwrap();
        uc.emu.reg_write(unicorn::RegisterX86::RBP, 0x6000).unwrap();

        uc
    }
}
<commit_msg>emu: uc: test *sp and *bp too<commit_after>//< Helpers that are useful for tests and doctests.
use crate::{
    arch::Arch,
    aspace::{AddressSpace, RelativeAddressSpace},
    module::{Module, Permissions, Section},
    RVA, VA,
};

/// configure a global logger at level==DEBUG.
pub fn init_logging() {
    let log_level = log::LevelFilter::Debug;
    fern::Dispatch::new()
        .format(move |out, message, record| {
            out.finish(format_args!(
                "{} [{:5}] {} {}",
                chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
                record.level(),
                if log_level == log::LevelFilter::Trace {
                    record.target()
                } else {
                    ""
                },
                message
            ))
        })
        .level(log_level)
        .chain(std::io::stderr())
        .filter(|metadata| !metadata.target().starts_with("goblin::pe"))
        .apply()
        .expect("failed to configure logging");
}

/// this is for testing, so will panic on error.
pub fn load_shellcode(arch: Arch, buf: &[u8]) -> Module {
    let mut address_space = RelativeAddressSpace::with_capacity(buf.len() as u64);
    address_space.map.writezx(0x0, buf).unwrap();

    Module {
        arch,
        sections: vec![Section {
            name:           "shellcode".to_string(),
            permissions:    Permissions::RWX,
            physical_range: std::ops::Range {
                start: 0x0,
                end:   buf.len() as RVA,
            },
            virtual_range:  std::ops::Range {
                start: 0x0,
                end:   buf.len() as RVA,
            },
        }],
        address_space: address_space.into_absolute(0x0).unwrap(),
    }
}

/// this is for testing, so will panic on error.
pub fn load_shellcode32(buf: &[u8]) -> Module {
    load_shellcode(Arch::X32, buf)
}

/// this is for testing, so will panic on error.
pub fn load_shellcode64(buf: &[u8]) -> Module {
    load_shellcode(Arch::X64, buf)
}

/// this is for testing, so will panic on error.
pub fn read_insn(module: &Module, va: VA) -> zydis::DecodedInstruction {
    use crate::analysis::dis;

    let decoder = dis::get_disassembler(module).unwrap();
    let mut insn_buf = [0u8; 16];
    module.address_space.read_into(va, &mut insn_buf).unwrap();
    decoder.decode(&insn_buf).unwrap().unwrap()
}

pub fn emu_from_shellcode64(code: &[u8]) -> crate::emu::Emulator {
    let m = load_shellcode64(code);
    let mut emu = crate::emu::Emulator::from_module(&m);
    emu.reg.rip = m.address_space.base_address; // 0x0

    emu.mem.mmap(0x5000, 0x2000, Permissions::RW).unwrap();
    emu.reg.rsp = 0x6000;
    emu.reg.rbp = 0x6000;

    emu
}

pub fn emu_from_shellcode32(code: &[u8]) -> crate::emu::Emulator {
    let m = load_shellcode32(code);
    let mut emu = crate::emu::Emulator::from_module(&m);
    emu.reg.set_eip(m.address_space.base_address as u32); // 0x0

    emu.mem.mmap(0x5000, 0x2000, Permissions::RW).unwrap();
    emu.reg.set_esp(0x6000);
    emu.reg.set_ebp(0x6000);

    emu
}

#[cfg(test)]
pub mod uc {
    use byteorder::{ByteOrder, LittleEndian};
    use unicorn::Cpu;

    use super::load_shellcode64;
    use crate::{
        aspace::AddressSpace,
        emu::{mmu::PAGE_SIZE, reg::STATUS_MASK},
        module::{Module, Permissions},
    };

    pub struct Uc {
        emu: unicorn::CpuX86,
    }

    impl Uc {
        pub fn from_module(m: &Module) -> Uc {
            let emu = match m.arch {
                crate::arch::Arch::X32 => unicorn::CpuX86::new(unicorn::Mode::MODE_32).unwrap(),
                crate::arch::Arch::X64 => unicorn::CpuX86::new(unicorn::Mode::MODE_64).unwrap(),
            };

            for section in m.sections.iter() {
                let mut page_addr = section.virtual_range.start;

                let section_size = section.virtual_range.end - section.virtual_range.start;
                emu.mem_map(
                    section.virtual_range.start,
                    crate::util::align(section_size, PAGE_SIZE as u64) as usize,
                    unicorn::PROT_WRITE,
                )
                .unwrap();

                while page_addr < section.virtual_range.end {
                    let mut page = [0u8; PAGE_SIZE];

                    // AddressSpace currently allows non-page-aligned sizes.
                    let page_data = if page_addr + PAGE_SIZE as u64 > section.virtual_range.end {
                        &mut page[..(section.virtual_range.end - page_addr) as usize]
                    } else {
                        &mut page[..]
                    };

                    m.address_space.read_into(page_addr, page_data).unwrap();
                    emu.mem_write(page_addr, &page[..]).unwrap();
                    page_addr += PAGE_SIZE as u64;
                }

                let mut prot = unicorn::PROT_NONE;
                if section.permissions.intersects(Permissions::W) {
                    prot.insert(unicorn::PROT_WRITE);
                }
                if section.permissions.intersects(Permissions::R) {
                    prot.insert(unicorn::PROT_READ);
                }
                if section.permissions.intersects(Permissions::X) {
                    prot.insert(unicorn::PROT_EXEC);
                }

                emu.mem_protect(
                    section.virtual_range.start,
                    crate::util::align(section_size, PAGE_SIZE as u64) as usize,
                    prot,
                )
                .unwrap();
            }

            Uc { emu }
        }

        pub fn step(&mut self) -> Result<(), unicorn::Error> {
            let rip = self.emu.reg_read(unicorn::RegisterX86::RIP).unwrap();
            self.emu.emu_start(rip, u64::MAX, 0, 1)
        }

        pub fn mem_read_u64(&self, addr: u64) -> Result<u64, unicorn::Error> {
            let buf = self.emu.mem_read(addr, 8)?;
            Ok(LittleEndian::read_u64(&buf))
        }

        /// panics if the given emulator does have the same state as this.
        pub fn check(&self, other: &crate::emu::Emulator) {
            use unicorn::RegisterX86::*;

            assert_eq!(self.emu.reg_read(RIP).unwrap(), other.reg.rip(), "register: rip",);

            assert_eq!(self.emu.reg_read(RSP).unwrap(), other.reg.rsp(), "register: rsp",);

            assert_eq!(self.emu.reg_read(RBP).unwrap(), other.reg.rbp(), "register: rbp",);

            assert_eq!(self.emu.reg_read(RAX).unwrap(), other.reg.rax(), "register: rax",);

            assert_eq!(self.emu.reg_read(RBX).unwrap(), other.reg.rbx(), "register: rbx",);

            assert_eq!(self.emu.reg_read(RCX).unwrap(), other.reg.rcx(), "register: rcx",);

            assert_eq!(self.emu.reg_read(RDX).unwrap(), other.reg.rdx(), "register: rdx",);

            assert_eq!(self.emu.reg_read(RSI).unwrap(), other.reg.rsi(), "register: rsi",);

            assert_eq!(self.emu.reg_read(RDI).unwrap(), other.reg.rdi(), "register: rdi",);

            // we don't emulate all of the flags, just the status flags.
            assert_eq!(
                self.emu.reg_read(EFLAGS).unwrap() as u64 & STATUS_MASK,
                other.reg.rflags() & STATUS_MASK,
                "flags"
            );

            // dereferece $SP and $BP
            let rsp = other.reg.rsp();
            assert_eq!(self.mem_read_u64(rsp).unwrap(), other.mem.read_u64(rsp).unwrap(),);

            let rbp = other.reg.rbp();
            assert_eq!(self.mem_read_u64(rbp).unwrap(), other.mem.read_u64(rbp).unwrap(),);
        }
    }

    pub fn uc_from_shellcode64(code: &[u8]) -> Uc {
        let m = load_shellcode64(code);
        let mut uc = Uc::from_module(&m);

        uc.emu
            .reg_write(unicorn::RegisterX86::RIP, m.address_space.base_address)
            .unwrap(); // 0x0

        uc.emu.mem_map(0x5000, 0x2000, unicorn::PROT_ALL).unwrap();
        uc.emu.reg_write(unicorn::RegisterX86::RSP, 0x6000).unwrap();
        uc.emu.reg_write(unicorn::RegisterX86::RBP, 0x6000).unwrap();

        uc
    }
}
<|endoftext|>