<commit_before>use matching::{MatchResult, match_request};

use pact_matching::models::{Pact, Interaction, Request, OptionalBody};
use pact_matching::models::matchingrules::*;
use pact_matching::models::generators::*;
use pact_matching::models::parse_query_string;

use std::collections::HashMap;
use std::sync::Arc;
use log::{log, error, warn, info, debug};
use hyper::{Body, Response, Server, Error};
use hyper::http::response::{Builder as ResponseBuilder};
use hyper::http::header::{HeaderName, HeaderValue};
use hyper::service::service_fn;
use futures::future;
use futures::future::Future;
use futures::stream::Stream;
use serde_json::json;

enum InteractionError {
    RequestHeaderEncodingError,
    RequestBodyError,
    ResponseHeaderEncodingError,
    ResponseBodyError
}

fn extract_path(uri: &hyper::Uri) -> String {
    uri.path_and_query()
        .map(|path_and_query| path_and_query.path())
        .unwrap_or("/")
        .into()
}

fn extract_query_string(uri: &hyper::Uri) -> Option<HashMap<String, Vec<String>>> {
    uri.path_and_query()
        .and_then(|path_and_query| path_and_query.query())
        .and_then(|query| parse_query_string(&query.into()))
}

fn extract_headers(headers: &hyper::HeaderMap) -> Result<Option<HashMap<String, String>>, InteractionError> {
    if headers.len() > 0 {
        let result: Result<HashMap<String, String>, InteractionError> = headers.keys()
            .map(|name| -> Result<(String, String), InteractionError> {
                let values = headers.get_all(name);
                let mut iter = values.iter();

                let first_value = iter.next().unwrap();

                if iter.next().is_some() {
                    warn!("Multiple headers associated with '{}', but only the first is used", name);
                }

                Ok((
                    name.as_str().into(),
                    first_value.to_str()
                        .map_err(|_| InteractionError::RequestHeaderEncodingError)?
                        .into()
                    )
                )
            })
            .collect();

        result.map(|map| Some(map))
    } else {
        Ok(None)
    }
}

pub fn extract_body(chunk: hyper::Chunk) -> OptionalBody {
    let bytes = chunk.into_bytes();
    if bytes.len() > 0 {
        OptionalBody::Present(bytes.to_vec())
    } else {
        OptionalBody::Empty
    }
}

fn hyper_request_to_pact_request(req: hyper::Request<Body>) -> impl Future<Item = Request, Error = InteractionError> {
    let method = req.method().to_string();
    let path = extract_path(req.uri());
    let query = extract_query_string(req.uri());
    let headers = extract_headers(req.headers());

    future::done(headers)
        .and_then(move |headers| {
            req.into_body()
                .concat2()
                .map_err(|_| InteractionError::RequestBodyError)
                .map(|body_chunk| (headers, body_chunk))
        })
        .and_then(|(headers, body_chunk)|
            Ok(Request {
                method: method,
                path: path,
                query: query,
                headers: headers,
                body: extract_body(body_chunk),
                matching_rules: MatchingRules::default(),
                generators: Generators::default()
            })
        )
}

fn set_hyper_headers(builder: &mut ResponseBuilder, headers: &Option<HashMap<String, String>>) -> Result<(), InteractionError> {
    let hyper_headers = builder.headers_mut().unwrap();
    match headers {
        Some(header_map) => {
            for (k, v) in header_map {
                // FIXME?: Headers are not sent in "raw" mode.
                // Names are converted to lower case and values are parsed.
                hyper_headers.insert(
                    HeaderName::from_bytes(k.as_bytes())
                        .map_err(|err| {
                            error!("Invalid header name '{}' ({})", k, err);
                            InteractionError::ResponseHeaderEncodingError
                        })?,
                    v.parse::<HeaderValue>()
                        .map_err(|err| {
                            error!("Invalid header value '{}': '{}' ({})", k, v, err);
                            InteractionError::ResponseHeaderEncodingError
                        })?
                );
            }
        },
        _ => {}
    }
    Ok(())
}

fn error_body(request: &Request, error: &String) -> String {
    let body = json!({ "error" : format!("{} : {:?}", error, request) });
    body.to_string()
}

fn match_result_to_hyper_response(request: &Request, match_result: MatchResult) -> Result<Response<Body>, InteractionError> {
    match match_result {
        MatchResult::RequestMatch(ref interaction) => {
            let response = pact_matching::generate_response(&interaction.response);
            info!("Request matched, sending response {:?}", response);
            info!("     body: '{}'\n\n", interaction.response.body.str_value());
            info!("     body: '{}'\n\n", interaction.response.body.str_value());

            let mut builder = Response::builder();

            builder.status(response.status);
            builder.header(hyper::header::ACCESS_CONTROL_ALLOW_ORIGIN, "*");
            set_hyper_headers(&mut builder, &response.headers)?;

            builder.body(match response.body {
                OptionalBody::Present(ref s) => Body::from(s.clone()),
                _ => Body::empty()
            })
                .map_err(|_| InteractionError::ResponseBodyError)
        },
        _ => {
            Response::builder()
                .status(500)
                .header(hyper::header::ACCESS_CONTROL_ALLOW_ORIGIN, "*")
                .header(hyper::header::CONTENT_TYPE, "application/json; charset=utf-8")
                .header("X-Pact", match_result.match_key())
                .body(Body::from(error_body(&request, &match_result.match_key())))
                .map_err(|_| InteractionError::ResponseBodyError)
        }
    }
}

fn handle_request(
    req: hyper::Request<Body>,
    pact: Arc<Pact>,
    matches_tx: std::sync::mpsc::Sender<MatchResult>
) -> impl Future<Item = Response<Body>, Error = InteractionError> {
    debug!("Creating pact request from hyper request");

    hyper_request_to_pact_request(req)
        .and_then(move |request| {
            info!("Received request {:?}", request);
            let match_result = match_request(&request, &pact.interactions);

            matches_tx.send(match_result.clone()).unwrap();

            match_result_to_hyper_response(&request, match_result)
        })
}

// TODO: Should instead use some form of X-Pact headers
fn handle_mock_request_error(result: Result<Response<Body>, InteractionError>) -> Result<Response<Body>, Error> {
    match result {
        Ok(response) => Ok(response),
        Err(error) => {
            let response = match error {
                InteractionError::RequestHeaderEncodingError => Response::builder()
                    .status(400)
                    .body(Body::from("Found an invalid header encoding")),
                InteractionError::RequestBodyError => Response::builder()
                    .status(500)
                    .body(Body::from("Could not process request body")),
                InteractionError::ResponseBodyError => Response::builder()
                    .status(500)
                    .body(Body::from("Could not process response body")),
                InteractionError::ResponseHeaderEncodingError => Response::builder()
                    .status(500)
                    .body(Body::from("Could not set response header"))
            };
            Ok(response.unwrap())
        }
    }
}

pub fn create_and_bind(
    pact: Pact,
    port: u16,
    shutdown: impl Future<Item = (), Error = ()>,
    matches_tx: std::sync::mpsc::Sender<MatchResult>
) -> Result<(impl Future<Item = (), Error = ()>, std::net::SocketAddr), hyper::Error> {
    let pact = Arc::new(pact);
    let addr = ([0, 0, 0, 0], port).into();

    let server = Server::try_bind(&addr)?
        .serve(move || {
            let pact = pact.clone();
            let matches_tx = matches_tx.clone();

            service_fn(move |req| {
                handle_request(req, pact.clone(), matches_tx.clone())
                    .then(handle_mock_request_error)
            })
        });

    let socket_addr = server.local_addr();

    let prepared_server = server
        .with_graceful_shutdown(shutdown)
        .map_err(move |err| {
            eprintln!("server error: {}", err);
        });

    Ok((prepared_server, socket_addr))
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::runtime::current_thread::Runtime;

    #[test]
    fn can_fetch_results_on_current_thread() {
        let mut runtime = Runtime::new().unwrap();

        let (shutdown_tx, shutdown_rx) = futures::sync::oneshot::channel();
        let (matches_tx, matches_rx) = std::sync::mpsc::channel();

        let (future, _) = create_and_bind(Pact::default(), 0, shutdown_rx.map_err(|_| ()), matches_tx).unwrap();

        runtime.spawn(future);
        shutdown_tx.send(()).unwrap();

        // Channel is still open and has no items
        assert!(matches_rx.try_recv().is_err());

        // Server has shut down, now flush the server future from runtime
        runtime.run().unwrap();

        // Match channel should now be closed and can read 0 matches.
        let all_matches: Vec<MatchResult> = matches_rx.iter().collect();
        assert_eq!(all_matches, vec![]);
    }
}
<commit_msg>Remove unused import<commit_after>use matching::{MatchResult, match_request};

use pact_matching::models::{Pact, Request, OptionalBody};
use pact_matching::models::matchingrules::*;
use pact_matching::models::generators::*;
use pact_matching::models::parse_query_string;

use std::collections::HashMap;
use std::sync::Arc;
use log::{log, error, warn, info, debug};
use hyper::{Body, Response, Server, Error};
use hyper::http::response::{Builder as ResponseBuilder};
use hyper::http::header::{HeaderName, HeaderValue};
use hyper::service::service_fn;
use futures::future;
use futures::future::Future;
use futures::stream::Stream;
use serde_json::json;

enum InteractionError {
    RequestHeaderEncodingError,
    RequestBodyError,
    ResponseHeaderEncodingError,
    ResponseBodyError
}

fn extract_path(uri: &hyper::Uri) -> String {
    uri.path_and_query()
        .map(|path_and_query| path_and_query.path())
        .unwrap_or("/")
        .into()
}

fn extract_query_string(uri: &hyper::Uri) -> Option<HashMap<String, Vec<String>>> {
    uri.path_and_query()
        .and_then(|path_and_query| path_and_query.query())
        .and_then(|query| parse_query_string(&query.into()))
}

fn extract_headers(headers: &hyper::HeaderMap) -> Result<Option<HashMap<String, String>>, InteractionError> {
    if headers.len() > 0 {
        let result: Result<HashMap<String, String>, InteractionError> = headers.keys()
            .map(|name| -> Result<(String, String), InteractionError> {
                let values = headers.get_all(name);
                let mut iter = values.iter();

                let first_value = iter.next().unwrap();

                if iter.next().is_some() {
                    warn!("Multiple headers associated with '{}', but only the first is used", name);
                }

                Ok((
                    name.as_str().into(),
                    first_value.to_str()
                        .map_err(|_| InteractionError::RequestHeaderEncodingError)?
                        .into()
                    )
                )
            })
            .collect();

        result.map(|map| Some(map))
    } else {
        Ok(None)
    }
}

pub fn extract_body(chunk: hyper::Chunk) -> OptionalBody {
    let bytes = chunk.into_bytes();
    if bytes.len() > 0 {
        OptionalBody::Present(bytes.to_vec())
    } else {
        OptionalBody::Empty
    }
}

fn hyper_request_to_pact_request(req: hyper::Request<Body>) -> impl Future<Item = Request, Error = InteractionError> {
    let method = req.method().to_string();
    let path = extract_path(req.uri());
    let query = extract_query_string(req.uri());
    let headers = extract_headers(req.headers());

    future::done(headers)
        .and_then(move |headers| {
            req.into_body()
                .concat2()
                .map_err(|_| InteractionError::RequestBodyError)
                .map(|body_chunk| (headers, body_chunk))
        })
        .and_then(|(headers, body_chunk)|
            Ok(Request {
                method: method,
                path: path,
                query: query,
                headers: headers,
                body: extract_body(body_chunk),
                matching_rules: MatchingRules::default(),
                generators: Generators::default()
            })
        )
}

fn set_hyper_headers(builder: &mut ResponseBuilder, headers: &Option<HashMap<String, String>>) -> Result<(), InteractionError> {
    let hyper_headers = builder.headers_mut().unwrap();
    match headers {
        Some(header_map) => {
            for (k, v) in header_map {
                // FIXME?: Headers are not sent in "raw" mode.
                // Names are converted to lower case and values are parsed.
                hyper_headers.insert(
                    HeaderName::from_bytes(k.as_bytes())
                        .map_err(|err| {
                            error!("Invalid header name '{}' ({})", k, err);
                            InteractionError::ResponseHeaderEncodingError
                        })?,
                    v.parse::<HeaderValue>()
                        .map_err(|err| {
                            error!("Invalid header value '{}': '{}' ({})", k, v, err);
                            InteractionError::ResponseHeaderEncodingError
                        })?
                );
            }
        },
        _ => {}
    }
    Ok(())
}

fn error_body(request: &Request, error: &String) -> String {
    let body = json!({ "error" : format!("{} : {:?}", error, request) });
    body.to_string()
}

fn match_result_to_hyper_response(request: &Request, match_result: MatchResult) -> Result<Response<Body>, InteractionError> {
    match match_result {
        MatchResult::RequestMatch(ref interaction) => {
            let response = pact_matching::generate_response(&interaction.response);
            info!("Request matched, sending response {:?}", response);
            info!("     body: '{}'\n\n", interaction.response.body.str_value());
            info!("     body: '{}'\n\n", interaction.response.body.str_value());

            let mut builder = Response::builder();

            builder.status(response.status);
            builder.header(hyper::header::ACCESS_CONTROL_ALLOW_ORIGIN, "*");
            set_hyper_headers(&mut builder, &response.headers)?;

            builder.body(match response.body {
                OptionalBody::Present(ref s) => Body::from(s.clone()),
                _ => Body::empty()
            })
                .map_err(|_| InteractionError::ResponseBodyError)
        },
        _ => {
            Response::builder()
                .status(500)
                .header(hyper::header::ACCESS_CONTROL_ALLOW_ORIGIN, "*")
                .header(hyper::header::CONTENT_TYPE, "application/json; charset=utf-8")
                .header("X-Pact", match_result.match_key())
                .body(Body::from(error_body(&request, &match_result.match_key())))
                .map_err(|_| InteractionError::ResponseBodyError)
        }
    }
}

fn handle_request(
    req: hyper::Request<Body>,
    pact: Arc<Pact>,
    matches_tx: std::sync::mpsc::Sender<MatchResult>
) -> impl Future<Item = Response<Body>, Error = InteractionError> {
    debug!("Creating pact request from hyper request");

    hyper_request_to_pact_request(req)
        .and_then(move |request| {
            info!("Received request {:?}", request);
            let match_result = match_request(&request, &pact.interactions);

            matches_tx.send(match_result.clone()).unwrap();

            match_result_to_hyper_response(&request, match_result)
        })
}

// TODO: Should instead use some form of X-Pact headers
fn handle_mock_request_error(result: Result<Response<Body>, InteractionError>) -> Result<Response<Body>, Error> {
    match result {
        Ok(response) => Ok(response),
        Err(error) => {
            let response = match error {
                InteractionError::RequestHeaderEncodingError => Response::builder()
                    .status(400)
                    .body(Body::from("Found an invalid header encoding")),
                InteractionError::RequestBodyError => Response::builder()
                    .status(500)
                    .body(Body::from("Could not process request body")),
                InteractionError::ResponseBodyError => Response::builder()
                    .status(500)
                    .body(Body::from("Could not process response body")),
                InteractionError::ResponseHeaderEncodingError => Response::builder()
                    .status(500)
                    .body(Body::from("Could not set response header"))
            };
            Ok(response.unwrap())
        }
    }
}

pub fn create_and_bind(
    pact: Pact,
    port: u16,
    shutdown: impl Future<Item = (), Error = ()>,
    matches_tx: std::sync::mpsc::Sender<MatchResult>
) -> Result<(impl Future<Item = (), Error = ()>, std::net::SocketAddr), hyper::Error> {
    let pact = Arc::new(pact);
    let addr = ([0, 0, 0, 0], port).into();

    let server = Server::try_bind(&addr)?
        .serve(move || {
            let pact = pact.clone();
            let matches_tx = matches_tx.clone();

            service_fn(move |req| {
                handle_request(req, pact.clone(), matches_tx.clone())
                    .then(handle_mock_request_error)
            })
        });

    let socket_addr = server.local_addr();

    let prepared_server = server
        .with_graceful_shutdown(shutdown)
        .map_err(move |err| {
            eprintln!("server error: {}", err);
        });

    Ok((prepared_server, socket_addr))
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::runtime::current_thread::Runtime;

    #[test]
    fn can_fetch_results_on_current_thread() {
        let mut runtime = Runtime::new().unwrap();

        let (shutdown_tx, shutdown_rx) = futures::sync::oneshot::channel();
        let (matches_tx, matches_rx) = std::sync::mpsc::channel();

        let (future, _) = create_and_bind(Pact::default(), 0, shutdown_rx.map_err(|_| ()), matches_tx).unwrap();

        runtime.spawn(future);
        shutdown_tx.send(()).unwrap();

        // Channel is still open and has no items
        assert!(matches_rx.try_recv().is_err());

        // Server has shut down, now flush the server future from runtime
        runtime.run().unwrap();

        // Match channel should now be closed and can read 0 matches.
        let all_matches: Vec<MatchResult> = matches_rx.iter().collect();
        assert_eq!(all_matches, vec![]);
    }
}
<|endoftext|>


<commit_before>extern crate robot_name as robot;
extern crate itertools;

// These are the expected signatures:
//
// impl Robot {
// pub fn new() -> Robot { ... }
// pub fn name<'a>(&'a self) -> &'a str { ... }
// pub fn reset_name(&mut self) { ... }
// }
//

fn assert_name_matches_pattern(n: &str) {
    assert!(n.len() == 5, "name is exactly 5 characters long");
    assert!(n[0..2].chars().all(|c| c >= 'A' && c <= 'Z'),
            "name starts with 2 uppercase letters");
    assert!(n[2..].chars().all(|c| c >= '0' && c <= '9'),
            "name ends with 3 numbers");
}

fn assert_name_is_persistent(r: &robot::Robot) {
    // The type system already proves this, but why not.
    let n1 = r.name();
    let n2 = r.name();
    let n3 = r.name();
    assert_eq!(n1, n2);
    assert_eq!(n2, n3);
}

#[test]
fn test_name_should_match_expected_pattern() {
    let r = robot::Robot::new();
    assert_name_matches_pattern(r.name());
}

#[test]
fn test_name_is_persistent() {
    assert_name_is_persistent(&robot::Robot::new());
}

#[test]
fn test_different_robots_have_different_names() {
    let r1 = robot::Robot::new();
    let r2 = robot::Robot::new();
    assert!(r1.name() != r2.name(), "Robot names should be different");
}

#[test]
fn test_new_name_should_match_expected_pattern() {
    let mut r = robot::Robot::new();
    assert_name_matches_pattern(r.name());
    r.reset_name();
    assert_name_matches_pattern(r.name());
}

#[test]
fn test_new_name_is_persistent() {
    let mut r = robot::Robot::new();
    r.reset_name();
    assert_name_is_persistent(&r);
}

#[test]
fn test_new_name_is_different_from_old_name() {
    let mut r = robot::Robot::new();
    let n1 = r.name().to_string();
    r.reset_name();
    let n2 = r.name().to_string();
    assert!(n1 != n2, "Robot name should change when reset");
}

#[test]
fn test_when_registry_is_full_with_all_generated_names() {
    let v = vec!['B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q',
                 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z'];

    // clean the registry first
    robot::ROBOTSREGISTRY.lock().unwrap().clear();

    use itertools::Itertools;
    v.clone()
        .into_iter()
        .cartesian_product(v.into_iter())
        .cartesian_product((100..1000).into_iter())
        .map(|((a, b), c)| format!("{}{}{}", a, b, c))
        .fold(&robot::ROBOTSREGISTRY, |ref acc, val| {
            acc.lock().unwrap().insert(val);
            acc
        });

    // adding one more robot should cause the robot name to be already present in the registry
    // after some iterations
    let r = robot::Robot::new();
    assert!(robot::ROBOTSREGISTRY.lock().unwrap().contains(r.name()));
}
<commit_msg>Update test naming for Robot Name<commit_after>extern crate robot_name as robot;
extern crate itertools;

// These are the expected signatures:
//
// impl Robot {
// pub fn new() -> Robot { ... }
// pub fn name<'a>(&'a self) -> &'a str { ... }
// pub fn reset_name(&mut self) { ... }
// }
//

fn assert_name_matches_pattern(n: &str) {
    assert!(n.len() == 5, "name is exactly 5 characters long");
    assert!(n[0..2].chars().all(|c| c >= 'A' && c <= 'Z'),
            "name starts with 2 uppercase letters");
    assert!(n[2..].chars().all(|c| c >= '0' && c <= '9'),
            "name ends with 3 numbers");
}

fn assert_name_is_persistent(r: &robot::Robot) {
    // The type system already proves this, but why not.
    let n1 = r.name();
    let n2 = r.name();
    let n3 = r.name();
    assert_eq!(n1, n2);
    assert_eq!(n2, n3);
}

#[test]
fn test_name_should_match_expected_pattern() {
    let r = robot::Robot::new();
    assert_name_matches_pattern(r.name());
}

#[test]
fn test_name_is_persistent() {
    assert_name_is_persistent(&robot::Robot::new());
}

#[test]
fn test_different_robots_have_different_names() {
    let r1 = robot::Robot::new();
    let r2 = robot::Robot::new();
    assert!(r1.name() != r2.name(), "Robot names should be different");
}

#[test]
fn test_new_name_should_match_expected_pattern() {
    let mut r = robot::Robot::new();
    assert_name_matches_pattern(r.name());
    r.reset_name();
    assert_name_matches_pattern(r.name());
}

#[test]
fn test_new_name_is_persistent() {
    let mut r = robot::Robot::new();
    r.reset_name();
    assert_name_is_persistent(&r);
}

#[test]
fn test_new_name_is_different_from_old_name() {
    let mut r = robot::Robot::new();
    let n1 = r.name().to_string();
    r.reset_name();
    let n2 = r.name().to_string();
    assert!(n1 != n2, "Robot name should change when reset");
}

#[test]
fn test_when_robot_name_already_in_registry() {
    let v = vec!['B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q',
                 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z'];

    // clean the registry first
    robot::ROBOTSREGISTRY.lock().unwrap().clear();

    use itertools::Itertools;
    v.clone()
        .into_iter()
        .cartesian_product(v.into_iter())
        .cartesian_product((100..1000).into_iter())
        .map(|((a, b), c)| format!("{}{}{}", a, b, c))
        .fold(&robot::ROBOTSREGISTRY, |ref acc, val| {
            acc.lock().unwrap().insert(val);
            acc
        });

    // adding one more robot should cause the robot name to be already present in the registry
    // after some iterations
    let r = robot::Robot::new();
    assert!(robot::ROBOTSREGISTRY.lock().unwrap().contains(r.name()));
}
<|endoftext|>


<commit_before>use linked_hash_map::LinkedHashMap;
use std::collections::HashSet;

use crate::options::{ImportStyle, Options, StringTransform};
use crate::shape::{self, Shape};
use crate::to_singular::to_singular;
use crate::util::{kebab_case, lower_camel_case, snake_case, type_case};

#[derive(PartialEq, PartialOrd, Ord, Eq, Hash, Clone, Copy)]
enum Import {
    Any,
    Optional,
    BaseModel,
    Field,
}

impl Import {
    fn pair(&self) -> (&'static str, &'static str) {
        match self {
            Import::Any => ("typing", "Any"),
            Import::Optional => ("typing", "Optional"),
            Import::BaseModel => ("pydantic", "BaseModel"),
            Import::Field => ("pydantic", "Field"),
        }
    }
    fn module(&self) -> &'static str {
        self.pair().0
    }
    fn identifier(&self) -> &'static str {
        self.pair().1
    }
    fn qualified(&self) -> String {
        let (module, identifier) = self.pair();
        format!("{}.{}", module, identifier)
    }
}

struct Ctxt {
    options: Options,
    type_names: HashSet<String>,
    imports: HashSet<Import>,
    created_classes: Vec<(Shape, Ident)>,
}

pub type Ident = String;
pub type Code = String;

pub fn python_types(name: &str, shape: &Shape, options: Options) -> Code {
    let mut ctxt = Ctxt {
        options,
        type_names: HashSet::new(),
        imports: HashSet::new(),
        created_classes: Vec::new(),
    };

    if !matches!(shape, Shape::Struct { .. }) {
        // reserve the requested name
        ctxt.type_names.insert(name.to_string());
    }

    let (ident, code) = type_from_shape(&mut ctxt, name, shape);
    let mut code = code.unwrap_or_default();

    if !ctxt.imports.is_empty() {
        let mut imports: Vec<_> = ctxt.imports.drain().collect();
        imports.sort();
        let mut import_code = String::new();
        match ctxt.options.import_style {
            ImportStyle::AssumeExisting => {}
            ImportStyle::AddImports => {
                for import in imports {
                    let (module, identifier) = import.pair();
                    import_code += &format!("from {} import {}\n", module, identifier);
                }
            }
            ImportStyle::QualifiedPaths => {
                let mut seen = HashSet::new();
                for import in imports {
                    let module = import.module();
                    if seen.insert(module) {
                        import_code += &format!("import {}\n", module);
                    }
                }
            }
        }
        if !import_code.is_empty(){
            import_code += "\n\n";
            code = import_code + &code;
        }
    }

    if ident != name {
        code += &format!("\n\n{} = {}", name, ident);
    }
    code
}

fn type_from_shape(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    use crate::shape::Shape::*;
    match shape {
        Null | Any | Bottom => (import(ctxt, Import::Any), None),
        Bool => ("bool".into(), None),
        StringT => ("str".into(), None),
        Integer => ("int".into(), None),
        Floating => ("float".into(), None),
        Tuple(shapes, _n) => {
            let folded = shape::fold_shapes(shapes.clone());
            if folded == Any && shapes.iter().any(|s| s != &Any) {
                generate_tuple_type(ctxt, path, shapes)
            } else {
                generate_vec_type(ctxt, path, &folded)
            }
        }
        VecT { elem_type: e } => generate_vec_type(ctxt, path, e),
        Struct { fields } => generate_data_class(ctxt, path, fields, shape),
        MapT { val_type: v } => generate_map_type(ctxt, path, v),
        Opaque(t) => (t.clone(), None),
        Optional(e) => {
            let (inner, defs) = type_from_shape(ctxt, path, e);
            if ctxt.options.use_default_for_missing_fields {
                (inner, defs)
            } else {
                let optional = import(ctxt, Import::Optional);
                (format!("{}[{}]", optional, inner), defs)
            }
        }
    }
}

fn generate_vec_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("list[{}]", inner), defs)
}

fn generate_map_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("dict[str, {}]", inner), defs)
}

fn generate_tuple_type(ctxt: &mut Ctxt, path: &str, shapes: &[Shape]) -> (Ident, Option<Code>) {
    let mut types = Vec::new();
    let mut defs = Vec::new();

    for shape in shapes {
        let (typ, def) = type_from_shape(ctxt, path, shape);
        types.push(typ);
        if let Some(code) = def {
            if !code.is_empty() {
                defs.push(code)
            }
        }
    }

    (
        format!("tuple[{}]", types.join(", ")),
        Some(defs.join("\n\n")),
    )
}

fn field_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "field", snake_case)
}

fn type_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "GeneratedType", type_case)
}

// https://docs.python.org/3/reference/lexical_analysis.html#keywords
#[rustfmt::skip]
const PYTHON_KEYWORDS: &[&str] = &[
    "False", "None", "True",
    "and", "as", "assert", "async", "await", "break", "class", "continue",
    "def", "del", "elif", "else", "except", "finally", "for", "from", "global",
    "if", "import", "in", "is", "lambda", "nonlocal", "not", "or", "pass",
    "raise", "return", "try", "while", "with", "yield",
];

fn type_or_field_name(
    name: &str,
    used_names: &HashSet<String>,
    default_name: &str,
    case_fn: fn(&str) -> String,
) -> Ident {
    let name = name.trim();
    let mut output_name = case_fn(name);
    if PYTHON_KEYWORDS.contains(&&*output_name) {
        output_name.push_str("_field");
    }
    if output_name.is_empty() {
        output_name.push_str(default_name);
    }
    if let Some(c) = output_name.chars().next() {
        if c.is_ascii_digit() {
            output_name = String::from("n") + &output_name;
        }
    }
    if !used_names.contains(&output_name) {
        return output_name;
    }
    for n in 2.. {
        let temp = format!("{}{}", output_name, n);
        if !used_names.contains(&temp) {
            return temp;
        }
    }
    unreachable!()
}

fn import(ctxt: &mut Ctxt, import: Import) -> String {
    ctxt.imports.insert(import);
    match ctxt.options.import_style {
        ImportStyle::QualifiedPaths => import.qualified(),
        _ => import.identifier().into(),
    }
}

fn generate_data_class(
    ctxt: &mut Ctxt,
    path: &str,
    field_shapes: &LinkedHashMap<String, Shape>,
    containing_shape: &Shape,
) -> (Ident, Option<Code>) {
    for (created_for_shape, ident) in ctxt.created_classes.iter() {
        if created_for_shape.is_acceptable_substitution_for(containing_shape) {
            return (ident.into(), None);
        }
    }

    let type_name = type_name(path, &ctxt.type_names);
    ctxt.type_names.insert(type_name.clone());
    ctxt.created_classes
        .push((containing_shape.clone(), type_name.clone()));

    let mut field_names = HashSet::new();
    let mut defs = Vec::new();

    let fields: Vec<Code> = field_shapes
        .iter()
        .map(|(name, typ)| {
            let field_name = field_name(name, &field_names);
            field_names.insert(field_name.clone());

            let (field_type, child_defs) = type_from_shape(ctxt, name, typ);

            if let Some(code) = child_defs {
                if !code.is_empty() {
                    defs.push(code);
                }
            }

            let mut field_code = String::new();
            let transformed = apply_transform(ctxt, &field_name, name);
            if transformed != field_name {
                field_code += &format!(" = {}(alias = \"{}\")", import(ctxt, Import::Field), transformed)
            }

            format!("    {}: {}{}", field_name, field_type, field_code)
        })
        .collect();

    let mut code = String::new();

    code += &format!(
        "class {}({}):\n",
        type_name,
        import(ctxt, Import::BaseModel)
    );

    if fields.is_empty() {
        code += "    pass\n";
    } else {
        code += &fields.join("\n");
        code += "\n";
    }

    if !defs.is_empty() {
        let mut d = defs.join("\n\n");
        d += "\n\n";
        d += &code;
        code = d;
    }

    (type_name, Some(code))
}

fn apply_transform(ctxt: &Ctxt, field_name: &str, name: &str) -> String {
    match ctxt.options.property_name_format {
        Some(StringTransform::LowerCase) => field_name.to_ascii_lowercase(),
        Some(StringTransform::PascalCase) => type_case(field_name),
        Some(StringTransform::SnakeCase) => snake_case(field_name),
        Some(StringTransform::KebabCase) => kebab_case(field_name),
        Some(StringTransform::UpperCase) => field_name.to_ascii_uppercase(),
        Some(StringTransform::CamelCase) => lower_camel_case(field_name),
        Some(StringTransform::ScreamingSnakeCase) => snake_case(field_name).to_ascii_uppercase(),
        Some(StringTransform::ScreamingKebabCase) => kebab_case(field_name).to_ascii_uppercase(),
        None => name.to_string(),
    }
}

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

    #[test]
    fn field_names_test() {
        fn field_name_test(from: &str, to: &str) {
            assert_eq!(
                field_name(from, &HashSet::new()),
                to.to_string(),
                r#"From "{}" to "{}""#,
                from,
                to
            );
        }

        field_name_test("valid", "valid");
        field_name_test("1", "n1");
        field_name_test("+1", "n1");
        field_name_test("", "field");
        field_name_test("def", "def_field");
    }
}
<commit_msg>shared: Remove spaces around alias= keyword argument.<commit_after>use linked_hash_map::LinkedHashMap;
use std::collections::HashSet;

use crate::options::{ImportStyle, Options, StringTransform};
use crate::shape::{self, Shape};
use crate::to_singular::to_singular;
use crate::util::{kebab_case, lower_camel_case, snake_case, type_case};

#[derive(PartialEq, PartialOrd, Ord, Eq, Hash, Clone, Copy)]
enum Import {
    Any,
    Optional,
    BaseModel,
    Field,
}

impl Import {
    fn pair(&self) -> (&'static str, &'static str) {
        match self {
            Import::Any => ("typing", "Any"),
            Import::Optional => ("typing", "Optional"),
            Import::BaseModel => ("pydantic", "BaseModel"),
            Import::Field => ("pydantic", "Field"),
        }
    }
    fn module(&self) -> &'static str {
        self.pair().0
    }
    fn identifier(&self) -> &'static str {
        self.pair().1
    }
    fn qualified(&self) -> String {
        let (module, identifier) = self.pair();
        format!("{}.{}", module, identifier)
    }
}

struct Ctxt {
    options: Options,
    type_names: HashSet<String>,
    imports: HashSet<Import>,
    created_classes: Vec<(Shape, Ident)>,
}

pub type Ident = String;
pub type Code = String;

pub fn python_types(name: &str, shape: &Shape, options: Options) -> Code {
    let mut ctxt = Ctxt {
        options,
        type_names: HashSet::new(),
        imports: HashSet::new(),
        created_classes: Vec::new(),
    };

    if !matches!(shape, Shape::Struct { .. }) {
        // reserve the requested name
        ctxt.type_names.insert(name.to_string());
    }

    let (ident, code) = type_from_shape(&mut ctxt, name, shape);
    let mut code = code.unwrap_or_default();

    if !ctxt.imports.is_empty() {
        let mut imports: Vec<_> = ctxt.imports.drain().collect();
        imports.sort();
        let mut import_code = String::new();
        match ctxt.options.import_style {
            ImportStyle::AssumeExisting => {}
            ImportStyle::AddImports => {
                for import in imports {
                    let (module, identifier) = import.pair();
                    import_code += &format!("from {} import {}\n", module, identifier);
                }
            }
            ImportStyle::QualifiedPaths => {
                let mut seen = HashSet::new();
                for import in imports {
                    let module = import.module();
                    if seen.insert(module) {
                        import_code += &format!("import {}\n", module);
                    }
                }
            }
        }
        if !import_code.is_empty(){
            import_code += "\n\n";
            code = import_code + &code;
        }
    }

    if ident != name {
        code += &format!("\n\n{} = {}", name, ident);
    }
    code
}

fn type_from_shape(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    use crate::shape::Shape::*;
    match shape {
        Null | Any | Bottom => (import(ctxt, Import::Any), None),
        Bool => ("bool".into(), None),
        StringT => ("str".into(), None),
        Integer => ("int".into(), None),
        Floating => ("float".into(), None),
        Tuple(shapes, _n) => {
            let folded = shape::fold_shapes(shapes.clone());
            if folded == Any && shapes.iter().any(|s| s != &Any) {
                generate_tuple_type(ctxt, path, shapes)
            } else {
                generate_vec_type(ctxt, path, &folded)
            }
        }
        VecT { elem_type: e } => generate_vec_type(ctxt, path, e),
        Struct { fields } => generate_data_class(ctxt, path, fields, shape),
        MapT { val_type: v } => generate_map_type(ctxt, path, v),
        Opaque(t) => (t.clone(), None),
        Optional(e) => {
            let (inner, defs) = type_from_shape(ctxt, path, e);
            if ctxt.options.use_default_for_missing_fields {
                (inner, defs)
            } else {
                let optional = import(ctxt, Import::Optional);
                (format!("{}[{}]", optional, inner), defs)
            }
        }
    }
}

fn generate_vec_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("list[{}]", inner), defs)
}

fn generate_map_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("dict[str, {}]", inner), defs)
}

fn generate_tuple_type(ctxt: &mut Ctxt, path: &str, shapes: &[Shape]) -> (Ident, Option<Code>) {
    let mut types = Vec::new();
    let mut defs = Vec::new();

    for shape in shapes {
        let (typ, def) = type_from_shape(ctxt, path, shape);
        types.push(typ);
        if let Some(code) = def {
            if !code.is_empty() {
                defs.push(code)
            }
        }
    }

    (
        format!("tuple[{}]", types.join(", ")),
        Some(defs.join("\n\n")),
    )
}

fn field_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "field", snake_case)
}

fn type_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "GeneratedType", type_case)
}

// https://docs.python.org/3/reference/lexical_analysis.html#keywords
#[rustfmt::skip]
const PYTHON_KEYWORDS: &[&str] = &[
    "False", "None", "True",
    "and", "as", "assert", "async", "await", "break", "class", "continue",
    "def", "del", "elif", "else", "except", "finally", "for", "from", "global",
    "if", "import", "in", "is", "lambda", "nonlocal", "not", "or", "pass",
    "raise", "return", "try", "while", "with", "yield",
];

fn type_or_field_name(
    name: &str,
    used_names: &HashSet<String>,
    default_name: &str,
    case_fn: fn(&str) -> String,
) -> Ident {
    let name = name.trim();
    let mut output_name = case_fn(name);
    if PYTHON_KEYWORDS.contains(&&*output_name) {
        output_name.push_str("_field");
    }
    if output_name.is_empty() {
        output_name.push_str(default_name);
    }
    if let Some(c) = output_name.chars().next() {
        if c.is_ascii_digit() {
            output_name = String::from("n") + &output_name;
        }
    }
    if !used_names.contains(&output_name) {
        return output_name;
    }
    for n in 2.. {
        let temp = format!("{}{}", output_name, n);
        if !used_names.contains(&temp) {
            return temp;
        }
    }
    unreachable!()
}

fn import(ctxt: &mut Ctxt, import: Import) -> String {
    ctxt.imports.insert(import);
    match ctxt.options.import_style {
        ImportStyle::QualifiedPaths => import.qualified(),
        _ => import.identifier().into(),
    }
}

fn generate_data_class(
    ctxt: &mut Ctxt,
    path: &str,
    field_shapes: &LinkedHashMap<String, Shape>,
    containing_shape: &Shape,
) -> (Ident, Option<Code>) {
    for (created_for_shape, ident) in ctxt.created_classes.iter() {
        if created_for_shape.is_acceptable_substitution_for(containing_shape) {
            return (ident.into(), None);
        }
    }

    let type_name = type_name(path, &ctxt.type_names);
    ctxt.type_names.insert(type_name.clone());
    ctxt.created_classes
        .push((containing_shape.clone(), type_name.clone()));

    let mut field_names = HashSet::new();
    let mut defs = Vec::new();

    let fields: Vec<Code> = field_shapes
        .iter()
        .map(|(name, typ)| {
            let field_name = field_name(name, &field_names);
            field_names.insert(field_name.clone());

            let (field_type, child_defs) = type_from_shape(ctxt, name, typ);

            if let Some(code) = child_defs {
                if !code.is_empty() {
                    defs.push(code);
                }
            }

            let mut field_code = String::new();
            let transformed = apply_transform(ctxt, &field_name, name);
            if transformed != field_name {
                field_code += &format!(" = {}(alias=\"{}\")", import(ctxt, Import::Field), transformed)
            }

            format!("    {}: {}{}", field_name, field_type, field_code)
        })
        .collect();

    let mut code = String::new();

    code += &format!(
        "class {}({}):\n",
        type_name,
        import(ctxt, Import::BaseModel)
    );

    if fields.is_empty() {
        code += "    pass\n";
    } else {
        code += &fields.join("\n");
        code += "\n";
    }

    if !defs.is_empty() {
        let mut d = defs.join("\n\n");
        d += "\n\n";
        d += &code;
        code = d;
    }

    (type_name, Some(code))
}

fn apply_transform(ctxt: &Ctxt, field_name: &str, name: &str) -> String {
    match ctxt.options.property_name_format {
        Some(StringTransform::LowerCase) => field_name.to_ascii_lowercase(),
        Some(StringTransform::PascalCase) => type_case(field_name),
        Some(StringTransform::SnakeCase) => snake_case(field_name),
        Some(StringTransform::KebabCase) => kebab_case(field_name),
        Some(StringTransform::UpperCase) => field_name.to_ascii_uppercase(),
        Some(StringTransform::CamelCase) => lower_camel_case(field_name),
        Some(StringTransform::ScreamingSnakeCase) => snake_case(field_name).to_ascii_uppercase(),
        Some(StringTransform::ScreamingKebabCase) => kebab_case(field_name).to_ascii_uppercase(),
        None => name.to_string(),
    }
}

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

    #[test]
    fn field_names_test() {
        fn field_name_test(from: &str, to: &str) {
            assert_eq!(
                field_name(from, &HashSet::new()),
                to.to_string(),
                r#"From "{}" to "{}""#,
                from,
                to
            );
        }

        field_name_test("valid", "valid");
        field_name_test("1", "n1");
        field_name_test("+1", "n1");
        field_name_test("", "field");
        field_name_test("def", "def_field");
    }
}
<|endoftext|>


<commit_before>#![cfg(feature = "serde")]

extern crate serde;

use EnumMap;
use Internal;

use self::serde::ser::{Serialize, Serializer, SerializeMap};
use self::serde::de::{self, Deserialize, Deserializer, Error, MapAccess};

use core::fmt;
use core::marker::PhantomData;

/// Requires crate feature `"serde"`
impl<K: Internal<V> + Serialize, V: Serialize> Serialize for EnumMap<K, V> {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let mut map = serializer.serialize_map(Some(self.len()))?;
        for (key, value) in self {
            map.serialize_entry(&key, value)?;
        }
        map.end()
    }
}

/// Requires crate feature `"serde"`
impl<'de, K, V> Deserialize<'de> for EnumMap<K, V>
where
    K: Internal<V>
        + Internal<Option<V>>
        + Deserialize<'de>,
    V: Deserialize<'de>,
{
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        deserializer.deserialize_map(Visitor(PhantomData))
    }
}

struct Visitor<K, V>(PhantomData<(K, V)>);

impl<'de, K, V> de::Visitor<'de> for Visitor<K, V>
where
    K: Internal<V>
        + Internal<Option<V>>
        + Deserialize<'de>,
    V: Deserialize<'de>,
{
    type Value = EnumMap<K, V>;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        write!(formatter, "a map")
    }

    fn visit_map<M: MapAccess<'de>>(self, mut access: M) -> Result<Self::Value, M::Error> {
        let mut entries = enum_map! { _ => None };
        while let Some((key, value)) = access.next_entry()? {
            entries[key] = Some(value);
        }
        for value in entries.values() {
            value.ok_or_else(|| M::Error::custom("key not specified"))?;
        }
        Ok(enum_map! { key => entries[key].take().unwrap() })
    }
}
<commit_msg>Borrow options in visit_map instead of moving them<commit_after>#![cfg(feature = "serde")]

extern crate serde;

use EnumMap;
use Internal;

use self::serde::ser::{Serialize, Serializer, SerializeMap};
use self::serde::de::{self, Deserialize, Deserializer, Error, MapAccess};

use core::fmt;
use core::marker::PhantomData;

/// Requires crate feature `"serde"`
impl<K: Internal<V> + Serialize, V: Serialize> Serialize for EnumMap<K, V> {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let mut map = serializer.serialize_map(Some(self.len()))?;
        for (key, value) in self {
            map.serialize_entry(&key, value)?;
        }
        map.end()
    }
}

/// Requires crate feature `"serde"`
impl<'de, K, V> Deserialize<'de> for EnumMap<K, V>
where
    K: Internal<V>
        + Internal<Option<V>>
        + Deserialize<'de>,
    V: Deserialize<'de>,
{
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        deserializer.deserialize_map(Visitor(PhantomData))
    }
}

struct Visitor<K, V>(PhantomData<(K, V)>);

impl<'de, K, V> de::Visitor<'de> for Visitor<K, V>
where
    K: Internal<V>
        + Internal<Option<V>>
        + Deserialize<'de>,
    V: Deserialize<'de>,
{
    type Value = EnumMap<K, V>;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        write!(formatter, "a map")
    }

    fn visit_map<M: MapAccess<'de>>(self, mut access: M) -> Result<Self::Value, M::Error> {
        let mut entries = enum_map! { _ => None };
        while let Some((key, value)) = access.next_entry()? {
            entries[key] = Some(value);
        }
        for value in entries.values() {
            value.as_ref().ok_or_else(|| M::Error::custom("key not specified"))?;
        }
        Ok(enum_map! { key => entries[key].take().unwrap() })
    }
}
<|endoftext|>


<commit_before>use std::error::Error;
use std::collections::{HashMap, BTreeMap, btree_map};
use std::str;

use itertools::Itertools;
use rust_htslib::bcf;

pub struct VarIndex {
    inner: HashMap<Vec<u8>, BTreeMap<u32, Vec<Variant>>>,
    max_dist: u32
}

impl VarIndex {
    pub fn new(mut reader: bcf::Reader, max_dist: u32) -> Result<Self, Box<Error>> {
        let mut inner = HashMap::new();
        let mut i = 0;
        let mut rec = bcf::Record::new();
        loop {
            if let Err(e) = reader.read(&mut rec) {
                if e.is_eof() {
                    break;
                }
                return Err(Box::new(e));
            }
            if let Some(rid) = rec.rid() {
                let chrom = reader.header().rid2name(rid);
                let recs = inner.entry(chrom.to_owned()).or_insert(BTreeMap::new());
                recs.entry(rec.pos()).or_insert_with(|| Vec::new()).push(Variant::new(&mut rec, &mut i)?);
                //recs.insert(rec.pos(), Variant::new(&mut rec, &mut i)?);
            } else {
                // skip records without rid
                let alt_count = rec.alleles().len() as u32 - 1;
                i += alt_count;
            }
        }

        Ok(VarIndex {
            inner: inner,
            max_dist: max_dist
        })
    }

    pub fn range(&self, chrom: &[u8], pos: u32) -> Option<btree_map::Range<u32, Vec<Variant>>> {
        self.inner.get(chrom).map(
            |recs| recs.range(pos.saturating_sub(self.max_dist)..pos + self.max_dist)
        )
    }
}

pub fn match_variants(
    matchbcf: &str, max_dist: u32, max_len_diff: u32
) -> Result<(), Box<Error>> {
    let mut inbcf = bcf::Reader::from_stdin()?;
    let mut header = bcf::Header::with_template(inbcf.header());

    header.push_record(
        format!("##INFO=<ID=MATCHING,Number=A,Type=Integer,\
        Description=\"For each alternative allele, -1 if it does not match a variant in another VCF/BCF. \
        If it matches a variant, an id i>=0 points to the i-th variant in the VCF/BCF (counting each \
        alternative allele separately). For indels, matching is fuzzy: distance of centres <= {}, difference of \
        lengths <= {}\">", max_dist, max_len_diff).as_bytes()
    );
    let mut outbcf = bcf::Writer::from_path(&"-", &header, false, false)?;
    let index = VarIndex::new(bcf::Reader::from_path(matchbcf)?, max_dist)?;

    let mut rec = bcf::Record::new();
    let mut i = 0;
    loop {
        if let Err(e) = inbcf.read(&mut rec) {
            if e.is_eof() {
                break;
            }
            return Err(Box::new(e));
        }
        outbcf.translate(&mut rec);

        if let Some(rid) = rec.rid() {
            let chrom = inbcf.header().rid2name(rid);
            let pos = rec.pos();

            let var = Variant::new(&mut rec, &mut i)?;
            let matching = var.alleles.iter().map(|a| {
                if let Some(range) = index.range(chrom, pos) {
                    for v in range.map(|(_, idx_vars)| idx_vars).flatten() {
                        if let Some(id) = var.matches(v, a, max_dist, max_len_diff) {
                            return id as i32;
                        }
                    }
                }
                -1
            }).collect_vec();

            rec.push_info_integer(b"MATCHING", &matching)?;
            outbcf.write(&rec)?;
        } else {
            outbcf.write(&rec)?;
        }

        if (i) % 1000 == 0 {
            info!("{} variants written.", i);
        }
    }
    info!("{} variants written.", i);

    Ok(())
}

#[derive(Debug)]
pub struct Variant {
    id: u32,
    rid: u32,
    pos: u32,
    alleles: Vec<VariantType>
}

impl Variant {
    pub fn new(rec: &mut bcf::Record, id: &mut u32) -> Result<Self, Box<Error>> {
        let pos = rec.pos();

        let svlen = if let Ok(Some(svlen)) = rec.info(b"SVLEN").integer() {
            Some(svlen[0].abs() as u32)
        } else { None };
        let svtype = if let Ok(Some(svtype)) = rec.info(b"SVTYPE").string() {
            Some(svtype[0].to_owned())
        } else { None };
        let end = if let Ok(Some(end)) = rec.info(b"END").integer() {
            Some(end[0] as u32)
        } else { None };
        let inslen = if let Ok(Some(inslen)) = rec.info(b"INSLEN").integer() {
            Some(inslen[0] as u32)
        } else { None };
        let alleles = rec.alleles();
        let refallele = alleles[0];

        let _alleles: Vec<VariantType> = if let Some(svtype) = svtype {
            vec![
                if svtype == b"INS" {
                    match (svlen, inslen) {
                        (Some(svlen), _)     => VariantType::Insertion(svlen),
                        (None, Some(inslen)) => VariantType::Insertion(inslen),
                        _ => {
                            warn!("Unsupported variant INS without SVLEN or INSLEN");
                            VariantType::Unsupported
                        }
                    }
                } else if svtype == b"DEL" {
                    let svlen = match(svlen, end) {
                        (Some(svlen), _)  => svlen,
                        (None, Some(end)) => end - 1 - pos,
                        _ => {
                            return Err(Box::new(MatchError::MissingTag("SVLEN or END".to_owned())));
                        }
                    };
                    VariantType::Deletion(svlen)
                } else {
                    warn!("Unsupported variant {}", try!(str::from_utf8(&svtype)));
                    VariantType::Unsupported
                }
            ]
        } else {
            let mut _alleles = Vec::with_capacity(alleles.len() - 1);
            for a in &alleles[1..] {
                _alleles.push(
                    if a.len() < refallele.len() {
                        VariantType::Deletion((refallele.len() - a.len()) as u32)
                    } else if a.len() > refallele.len() {
                        VariantType::Insertion((a.len() - refallele.len()) as u32)
                    } else if a.len() == 1 {
                        VariantType::SNV(a[0])
                    } else {
                        warn!("Unsupported variant {} -> {}", try!(str::from_utf8(refallele)), try!(str::from_utf8(a)));
                        VariantType::Unsupported
                    }
                );
            }
            _alleles
        };
        let var = Variant {
            id: *id,
            rid: rec.rid().unwrap(),
            pos: pos,
            alleles: _alleles
        };
        *id += alleles.len() as u32 - 1;
        Ok(var)
    }

    pub fn centerpoint(&self, allele: &VariantType) -> u32 {
        match allele {
            &VariantType::SNV(_) => self.pos,
            &VariantType::Insertion(_) => self.pos,
            &VariantType::Deletion(len) => (self.pos as f64 + len as f64 / 2.0) as u32,
            &VariantType::Unsupported => panic!("Unsupported variant.")
        }
    }

    pub fn matches(&self, other: &Variant, allele: &VariantType, max_dist: u32, max_len_diff: u32) -> Option<u32> {
        if allele.is_unsupported() {
            return None;
        }
        for (j, b) in other.alleles.iter().enumerate() {
            if b.is_unsupported() {
                continue;
            }
            let dist = (self.centerpoint(allele) as i32 - other.centerpoint(b) as i32).abs() as u32;
            match (allele, b) {
                (&VariantType::SNV(a), &VariantType::SNV(b)) => {
                    if a == b && dist == 0 {
                        return Some(other.id(j));
                    }
                },
                (&VariantType::Insertion(l1), &VariantType::Insertion(l2)) |
                (&VariantType::Deletion(l1), &VariantType::Deletion(l2)) => {
                    if (l1 as i32 - l2 as i32).abs() as u32 <= max_len_diff && dist <= max_dist {
                        return Some(other.id(j));
                    }
                },
                // TODO: for now, ignore complex variants
                _ => continue
            }
        }
        None
    }

    pub fn id(&self, allele: usize) -> u32 {
        self.id + allele as u32
    }
}

#[derive(Debug)]
pub enum VariantType {
    SNV(u8),
    Insertion(u32),
    Deletion(u32),
    Unsupported
}

impl VariantType {
    pub fn is_unsupported(&self) -> bool {
        match self {
            &VariantType::Unsupported => true,
            _ => false
        }
    }
}

quick_error! {
    #[derive(Debug)]
    pub enum MatchError {
        MissingTag(tag: String) {
            description("missing tag")
            display("missing tag {}", tag)
        }
    }
}
<commit_msg>fix Itertools::flatten regression in nightly. (#6)<commit_after>use std::error::Error;
use std::collections::{HashMap, BTreeMap, btree_map};
use std::str;

use itertools::Itertools;
use rust_htslib::bcf;

pub struct VarIndex {
    inner: HashMap<Vec<u8>, BTreeMap<u32, Vec<Variant>>>,
    max_dist: u32
}

impl VarIndex {
    pub fn new(mut reader: bcf::Reader, max_dist: u32) -> Result<Self, Box<Error>> {
        let mut inner = HashMap::new();
        let mut i = 0;
        let mut rec = bcf::Record::new();
        loop {
            if let Err(e) = reader.read(&mut rec) {
                if e.is_eof() {
                    break;
                }
                return Err(Box::new(e));
            }
            if let Some(rid) = rec.rid() {
                let chrom = reader.header().rid2name(rid);
                let recs = inner.entry(chrom.to_owned()).or_insert(BTreeMap::new());
                recs.entry(rec.pos()).or_insert_with(|| Vec::new()).push(Variant::new(&mut rec, &mut i)?);
                //recs.insert(rec.pos(), Variant::new(&mut rec, &mut i)?);
            } else {
                // skip records without rid
                let alt_count = rec.alleles().len() as u32 - 1;
                i += alt_count;
            }
        }

        Ok(VarIndex {
            inner: inner,
            max_dist: max_dist
        })
    }

    pub fn range(&self, chrom: &[u8], pos: u32) -> Option<btree_map::Range<u32, Vec<Variant>>> {
        self.inner.get(chrom).map(
            |recs| recs.range(pos.saturating_sub(self.max_dist)..pos + self.max_dist)
        )
    }
}

pub fn match_variants(
    matchbcf: &str, max_dist: u32, max_len_diff: u32
) -> Result<(), Box<Error>> {
    let mut inbcf = bcf::Reader::from_stdin()?;
    let mut header = bcf::Header::with_template(inbcf.header());

    header.push_record(
        format!("##INFO=<ID=MATCHING,Number=A,Type=Integer,\
        Description=\"For each alternative allele, -1 if it does not match a variant in another VCF/BCF. \
        If it matches a variant, an id i>=0 points to the i-th variant in the VCF/BCF (counting each \
        alternative allele separately). For indels, matching is fuzzy: distance of centres <= {}, difference of \
        lengths <= {}\">", max_dist, max_len_diff).as_bytes()
    );
    let mut outbcf = bcf::Writer::from_path(&"-", &header, false, false)?;
    let index = VarIndex::new(bcf::Reader::from_path(matchbcf)?, max_dist)?;

    let mut rec = bcf::Record::new();
    let mut i = 0;
    loop {
        if let Err(e) = inbcf.read(&mut rec) {
            if e.is_eof() {
                break;
            }
            return Err(Box::new(e));
        }
        outbcf.translate(&mut rec);

        if let Some(rid) = rec.rid() {
            let chrom = inbcf.header().rid2name(rid);
            let pos = rec.pos();

            let var = Variant::new(&mut rec, &mut i)?;
            let matching = var.alleles.iter().map(|a| {
                if let Some(range) = index.range(chrom, pos) {
                    for v in Itertools::flatten(range.map(|(_, idx_vars)| idx_vars)) {
                        if let Some(id) = var.matches(v, a, max_dist, max_len_diff) {
                            return id as i32;
                        }
                    }
                }
                -1
            }).collect_vec();

            rec.push_info_integer(b"MATCHING", &matching)?;
            outbcf.write(&rec)?;
        } else {
            outbcf.write(&rec)?;
        }

        if (i) % 1000 == 0 {
            info!("{} variants written.", i);
        }
    }
    info!("{} variants written.", i);

    Ok(())
}

#[derive(Debug)]
pub struct Variant {
    id: u32,
    rid: u32,
    pos: u32,
    alleles: Vec<VariantType>
}

impl Variant {
    pub fn new(rec: &mut bcf::Record, id: &mut u32) -> Result<Self, Box<Error>> {
        let pos = rec.pos();

        let svlen = if let Ok(Some(svlen)) = rec.info(b"SVLEN").integer() {
            Some(svlen[0].abs() as u32)
        } else { None };
        let svtype = if let Ok(Some(svtype)) = rec.info(b"SVTYPE").string() {
            Some(svtype[0].to_owned())
        } else { None };
        let end = if let Ok(Some(end)) = rec.info(b"END").integer() {
            Some(end[0] as u32)
        } else { None };
        let inslen = if let Ok(Some(inslen)) = rec.info(b"INSLEN").integer() {
            Some(inslen[0] as u32)
        } else { None };
        let alleles = rec.alleles();
        let refallele = alleles[0];

        let _alleles: Vec<VariantType> = if let Some(svtype) = svtype {
            vec![
                if svtype == b"INS" {
                    match (svlen, inslen) {
                        (Some(svlen), _)     => VariantType::Insertion(svlen),
                        (None, Some(inslen)) => VariantType::Insertion(inslen),
                        _ => {
                            warn!("Unsupported variant INS without SVLEN or INSLEN");
                            VariantType::Unsupported
                        }
                    }
                } else if svtype == b"DEL" {
                    let svlen = match(svlen, end) {
                        (Some(svlen), _)  => svlen,
                        (None, Some(end)) => end - 1 - pos,
                        _ => {
                            return Err(Box::new(MatchError::MissingTag("SVLEN or END".to_owned())));
                        }
                    };
                    VariantType::Deletion(svlen)
                } else {
                    warn!("Unsupported variant {}", try!(str::from_utf8(&svtype)));
                    VariantType::Unsupported
                }
            ]
        } else {
            let mut _alleles = Vec::with_capacity(alleles.len() - 1);
            for a in &alleles[1..] {
                _alleles.push(
                    if a.len() < refallele.len() {
                        VariantType::Deletion((refallele.len() - a.len()) as u32)
                    } else if a.len() > refallele.len() {
                        VariantType::Insertion((a.len() - refallele.len()) as u32)
                    } else if a.len() == 1 {
                        VariantType::SNV(a[0])
                    } else {
                        warn!("Unsupported variant {} -> {}", try!(str::from_utf8(refallele)), try!(str::from_utf8(a)));
                        VariantType::Unsupported
                    }
                );
            }
            _alleles
        };
        let var = Variant {
            id: *id,
            rid: rec.rid().unwrap(),
            pos: pos,
            alleles: _alleles
        };
        *id += alleles.len() as u32 - 1;
        Ok(var)
    }

    pub fn centerpoint(&self, allele: &VariantType) -> u32 {
        match allele {
            &VariantType::SNV(_) => self.pos,
            &VariantType::Insertion(_) => self.pos,
            &VariantType::Deletion(len) => (self.pos as f64 + len as f64 / 2.0) as u32,
            &VariantType::Unsupported => panic!("Unsupported variant.")
        }
    }

    pub fn matches(&self, other: &Variant, allele: &VariantType, max_dist: u32, max_len_diff: u32) -> Option<u32> {
        if allele.is_unsupported() {
            return None;
        }
        for (j, b) in other.alleles.iter().enumerate() {
            if b.is_unsupported() {
                continue;
            }
            let dist = (self.centerpoint(allele) as i32 - other.centerpoint(b) as i32).abs() as u32;
            match (allele, b) {
                (&VariantType::SNV(a), &VariantType::SNV(b)) => {
                    if a == b && dist == 0 {
                        return Some(other.id(j));
                    }
                },
                (&VariantType::Insertion(l1), &VariantType::Insertion(l2)) |
                (&VariantType::Deletion(l1), &VariantType::Deletion(l2)) => {
                    if (l1 as i32 - l2 as i32).abs() as u32 <= max_len_diff && dist <= max_dist {
                        return Some(other.id(j));
                    }
                },
                // TODO: for now, ignore complex variants
                _ => continue
            }
        }
        None
    }

    pub fn id(&self, allele: usize) -> u32 {
        self.id + allele as u32
    }
}

#[derive(Debug)]
pub enum VariantType {
    SNV(u8),
    Insertion(u32),
    Deletion(u32),
    Unsupported
}

impl VariantType {
    pub fn is_unsupported(&self) -> bool {
        match self {
            &VariantType::Unsupported => true,
            _ => false
        }
    }
}

quick_error! {
    #[derive(Debug)]
    pub enum MatchError {
        MissingTag(tag: String) {
            description("missing tag")
            display("missing tag {}", tag)
        }
    }
}
<|endoftext|>


<commit_before>extern crate rand;
extern crate evo;
extern crate petgraph;
extern crate graph_annealing;
extern crate clap;
extern crate crossbeam;
extern crate simple_parallel;
extern crate num_cpus;
extern crate pcg;
extern crate graph_sgf;
extern crate triadic_census;

use std::f32;
use std::fmt::Debug;
use std::str::FromStr;
use clap::{App, Arg};
use rand::{Rng, SeedableRng};
use rand::distributions::{IndependentSample, Range};
use rand::os::OsRng;
use pcg::PcgRng;

use evo::Probability;
use evo::nsga2::{self, FitnessEval, MultiObjective3};
use graph_annealing::repr::edge_ops_genome::{EdgeOpsGenome, Toolbox};
use graph_annealing::helper::draw_graph;
use graph_annealing::goal::Goal;
use simple_parallel::Pool;

use petgraph::{Directed, Graph};
use graph_sgf::{PetgraphReader, Unweighted};
use triadic_census::OptDenseDigraph;

use std::io::BufReader;
use std::fs::File;

// maximum connectivity
const MAX_DEGREE: u32 = 20;

#[derive(Debug, Copy, Clone)]
enum FitnessFunction {
    Null,
    ConnectedComponents,
    StronglyConnectedComponents,
    NeighborMatching,
    TriadicDistance,
}

fn apply_fitness_function<N: Clone + Default, E: Clone + Default>(fitfun: FitnessFunction,
                                                                  goal: &Goal<N, E>,
                                                                  g: &OptDenseDigraph<(), ()>)
                                                                  -> f32 {
    match fitfun {
        FitnessFunction::Null => {
            0.0
        }
        FitnessFunction::ConnectedComponents => {
            goal.connected_components_distance(g) as f32
        }
        FitnessFunction::StronglyConnectedComponents => {
            goal.strongly_connected_components_distance(g) as f32
        }
        FitnessFunction::NeighborMatching => {
            goal.neighbor_matching_score(g) as f32
        }
        FitnessFunction::TriadicDistance => {
            goal.triadic_distance(g) as f32
        }
    }
}

struct MyEval<N, E> {
    goal: Goal<N, E>,
    pool: Pool,
    fitness_functions: (FitnessFunction, FitnessFunction, FitnessFunction),
}

fn fitness<N: Clone + Default, E: Clone + Default>(fitness_functions: (FitnessFunction,
                                                                       FitnessFunction,
                                                                       FitnessFunction),
                                                   goal: &Goal<N, E>,
                                                   ind: &EdgeOpsGenome)
                                                   -> MultiObjective3<f32> {
    let g = ind.to_graph(MAX_DEGREE);
    MultiObjective3::from((apply_fitness_function(fitness_functions.0, goal, &g),
                           apply_fitness_function(fitness_functions.1, goal, &g),
                           apply_fitness_function(fitness_functions.2, goal, &g)))

}

impl<N:Clone+Sync+Default,E:Clone+Sync+Default> FitnessEval<EdgeOpsGenome, MultiObjective3<f32>> for MyEval<N,E> {
    fn fitness(&mut self, pop: &[EdgeOpsGenome]) -> Vec<MultiObjective3<f32>> {
        let pool = &mut self.pool;
        let goal = &self.goal;

        let fitness_functions = self.fitness_functions;

        crossbeam::scope(|scope| {
            pool.map(scope, pop, |ind| fitness(fitness_functions, goal, ind))
                .collect()
        })
    }
}

#[derive(Debug)]
struct Stat<T: Debug> {
    min: T,
    max: T,
    avg: T,
}

#[allow(non_snake_case)]
fn main() {
    let ncpus = num_cpus::get();
    println!("Using {} CPUs", ncpus);
    let matches = App::new("nsga2_edge")
                      .arg(Arg::with_name("NGEN")
                               .long("ngen")
                               .help("Number of generations to run")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("OPS")
                               .long("ops")
                               .help("Edge operation and weight specification, e.g. \
                                      Dup:1,Split:3,Parent:2")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("GRAPH")
                               .long("graph")
                               .help("SGF graph file")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("MU")
                               .long("mu")
                               .help("Population size")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("LAMBDA")
                               .long("lambda")
                               .help("Size of offspring population")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("K")
                               .long("k")
                               .help("Tournament selection (default: 2)")
                               .takes_value(true))
                      .arg(Arg::with_name("SEED")
                               .long("seed")
                               .help("Seed value for Rng")
                               .takes_value(true))
                      .arg(Arg::with_name("PMUT")
                               .long("pmut")
                               .help("Probability for element mutation")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("ILEN")
                               .long("ilen")
                               .help("Initial genome length (random range)")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("OBJECTIVES")
                               .long("objectives")
                               .help("Specify 3 objective functions comma separated (null, cc, \
                                      scc, nm, td), e.g. cc,nm,td")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("THRESHOLD")
                               .long("threshold")
                               .help("Abort if for all fitness[i] <= value[i] (default: 0,0,0)")
                               .takes_value(true)
                               .required(false))
                      .get_matches();

    // number of generations
    let NGEN: usize = FromStr::from_str(matches.value_of("NGEN").unwrap()).unwrap();
    println!("NGEN: {}", NGEN);

    // size of population
    let MU: usize = FromStr::from_str(matches.value_of("MU").unwrap()).unwrap();
    println!("MU: {}", MU);

    // size of offspring population
    let LAMBDA: usize = FromStr::from_str(matches.value_of("LAMBDA").unwrap()).unwrap();
    println!("LAMBDA: {}", LAMBDA);

    // tournament selection
    let K: usize = FromStr::from_str(matches.value_of("K").unwrap_or("2")).unwrap();
    assert!(K > 0);
    println!("K: {}", K);

    let seed: Vec<u64>;
    if let Some(seed_str) = matches.value_of("SEED") {
        seed = seed_str.split(",").map(|s| FromStr::from_str(s).unwrap()).collect();
    } else {
        println!("Use OsRng to generate seed..");
        let mut rng = OsRng::new().unwrap();
        seed = (0..2).map(|_| rng.next_u64()).collect();
    }
    println!("SEED: {:?}", seed);

    let PMUT: f32 = FromStr::from_str(matches.value_of("PMUT").unwrap()).unwrap();
    println!("PMUT: {:?}", PMUT);

    let ilen_str = matches.value_of("ILEN").unwrap();
    let ilen: Vec<usize> = ilen_str.split(",").map(|s| FromStr::from_str(s).unwrap()).collect();
    assert!(ilen.len() == 1 || ilen.len() == 2);

    let ilen_from = ilen[0];
    let ilen_to = if ilen.len() == 1 {
        ilen[0]
    } else {
        ilen[1]
    };
    assert!(ilen_from <= ilen_to);

    // read objective functions
    let mut objectives_arr: Vec<FitnessFunction> = Vec::new();
    for s in matches.value_of("OBJECTIVES").unwrap().split(",") {
        objectives_arr.push(match s {
            "null" => {
                FitnessFunction::Null
            }
            "cc" => {
                FitnessFunction::ConnectedComponents
            }
            "scc" => {
                FitnessFunction::StronglyConnectedComponents
            }
            "nm" => {
                FitnessFunction::NeighborMatching
            }
            "td" => {
                FitnessFunction::TriadicDistance
            }
            _ => {
                panic!("Invalid objective function");
            }
        });
    }

    while objectives_arr.len() < 3 {
        objectives_arr.push(FitnessFunction::Null);
    }

    if objectives_arr.len() > 3 {
        panic!("Max 3 objectives allowed");
    }

    println!("objectives={:?}", objectives_arr);

    // read objective functions
    let mut threshold_arr: Vec<f32> = Vec::new();
    for s in matches.value_of("THRESHOLD").unwrap_or("").split(",") {
        let value: f32 = FromStr::from_str(s).unwrap();
        threshold_arr.push(value);
    }

    while threshold_arr.len() < 3 {
        threshold_arr.push(0.0);
    }

    if threshold_arr.len() > 3 {
        panic!("Max 3 threshold values allowed");
    }

    // read graph
    let graph_file = matches.value_of("GRAPH").unwrap();
    println!("Using graph file: {}", graph_file);
    let graph = {
        let mut f = BufReader::new(File::open(graph_file).unwrap());
        let graph: Graph<Unweighted, Unweighted, Directed> = PetgraphReader::from_sgf(&mut f);
        graph
    };

    // Parse weighted operation choice from command line
    let ops = Toolbox::parse_weighted_op_choice_list(matches.value_of("OPS").unwrap()).unwrap();
    println!("edge ops: {:?}", ops);
    let mut toolbox = Toolbox::new(Probability::new(PMUT), &ops[..]);

    let mut evaluator = MyEval {
        goal: Goal::new(OptDenseDigraph::from(graph)),
        pool: Pool::new(ncpus),
        fitness_functions: (objectives_arr[0], objectives_arr[1], objectives_arr[2]),
    };

    assert!(seed.len() == 2);
    let mut rng: PcgRng = SeedableRng::from_seed([seed[0], seed[1]]);

    // create initial random population
    let initial_population: Vec<EdgeOpsGenome> = (0..MU)
                                                     .map(|_| {
                                                         let len = if ilen_from == ilen_to {
                                                             ilen_from
                                                         } else {
                                                             Range::new(ilen_from, ilen_to)
                                                                 .ind_sample(&mut rng)
                                                         };
                                                         toolbox.random_genome(&mut rng, len)
                                                     })
                                                     .collect();

    // evaluate fitness
    let fitness: Vec<_> = evaluator.fitness(&initial_population[..]);
    assert!(fitness.len() == initial_population.len());

    let mut pop = initial_population;
    let mut fit = fitness;

    for i in 0..NGEN {
        println!("Iteration: {}", i);
        let (new_pop, new_fit) = nsga2::iterate(&mut rng,
                                                pop,
                                                fit,
                                                &mut evaluator,
                                                MU,
                                                LAMBDA,
                                                K,
                                                &mut toolbox);
        pop = new_pop;
        fit = new_fit;

        let stats: Vec<Stat<f32>> = [0, 1, 2]
                                        .iter()
                                        .map(|&i| {
                                            let min = fit.iter().fold(f32::INFINITY, |acc, f| {
                                                let x = f.objectives[i];
                                                if x < acc {
                                                    x
                                                } else {
                                                    acc
                                                }
                                            });
                                            let max = fit.iter().fold(-f32::INFINITY, |acc, f| {
                                                let x = f.objectives[i];
                                                if x > acc {
                                                    x
                                                } else {
                                                    acc
                                                }
                                            });
                                            let sum = fit.iter()
                                                         .fold(0.0, |acc, f| acc + f.objectives[i]);
                                            Stat {
                                                min: min,
                                                max: max,
                                                avg: sum / fit.len() as f32,
                                            }
                                        })
                                        .collect();
        println!("stats: {:?}", stats);
        let mut found_optimum = false;
        for f in fit.iter() {
            if f.objectives[0] <= threshold_arr[0] && f.objectives[1] <= threshold_arr[1] &&
               f.objectives[2] <= threshold_arr[2] {
                found_optimum = true;
                break;
            }
        }
        if found_optimum {
            println!("Found premature optimum in Iteration {}", i);
            break;
        }
    }
    println!("===========================================================");

    // finally evaluate rank and crowding distance (using select()).
    let rank_dist = nsga2::select(&fit[..], MU);
    assert!(rank_dist.len() == MU);

    let mut j = 0;

    for rd in rank_dist.iter() {
        if rd.rank == 0 {
            println!("-------------------------------------------");
            println!("rd: {:?}", rd);
            println!("fitness: {:?}", fit[rd.idx]);
            println!("genome: {:?}", pop[rd.idx]);

            // if fit[rd.idx].objectives[0] < 1.0 {
            draw_graph((&pop[rd.idx].to_graph(MAX_DEGREE)).ref_graph(),
                       // XXX: name
                       &format!("nsga2edgeops_g{}_f{}_i{}.svg",
                                NGEN,
                                fit[rd.idx].objectives[1] as usize,
                                j));
            j += 1;
            // }
        }
    }
}
<commit_msg>Parameter --max-degree<commit_after>extern crate rand;
extern crate evo;
extern crate petgraph;
extern crate graph_annealing;
extern crate clap;
extern crate crossbeam;
extern crate simple_parallel;
extern crate num_cpus;
extern crate pcg;
extern crate graph_sgf;
extern crate triadic_census;

use std::f32;
use std::fmt::Debug;
use std::str::FromStr;
use clap::{App, Arg};
use rand::{Rng, SeedableRng};
use rand::distributions::{IndependentSample, Range};
use rand::os::OsRng;
use pcg::PcgRng;

use evo::Probability;
use evo::nsga2::{self, FitnessEval, MultiObjective3};
use graph_annealing::repr::edge_ops_genome::{EdgeOpsGenome, Toolbox};
use graph_annealing::helper::draw_graph;
use graph_annealing::goal::Goal;
use simple_parallel::Pool;

use petgraph::{Directed, Graph};
use graph_sgf::{PetgraphReader, Unweighted};
use triadic_census::OptDenseDigraph;

use std::io::BufReader;
use std::fs::File;

#[derive(Debug, Copy, Clone)]
enum FitnessFunction {
    Null,
    ConnectedComponents,
    StronglyConnectedComponents,
    NeighborMatching,
    TriadicDistance,
}

fn apply_fitness_function<N: Clone + Default, E: Clone + Default>(fitfun: FitnessFunction,
                                                                  goal: &Goal<N, E>,
                                                                  g: &OptDenseDigraph<(), ()>)
                                                                  -> f32 {
    match fitfun {
        FitnessFunction::Null => {
            0.0
        }
        FitnessFunction::ConnectedComponents => {
            goal.connected_components_distance(g) as f32
        }
        FitnessFunction::StronglyConnectedComponents => {
            goal.strongly_connected_components_distance(g) as f32
        }
        FitnessFunction::NeighborMatching => {
            goal.neighbor_matching_score(g) as f32
        }
        FitnessFunction::TriadicDistance => {
            goal.triadic_distance(g) as f32
        }
    }
}

struct MyEval<N, E> {
    goal: Goal<N, E>,
    pool: Pool,
    fitness_functions: (FitnessFunction, FitnessFunction, FitnessFunction),
    max_degree: u32,
}

#[inline]
fn fitness<N: Clone + Default, E: Clone + Default>(fitness_functions: (FitnessFunction,
                                                                       FitnessFunction,
                                                                       FitnessFunction),
                                                   max_degree: u32,
                                                   goal: &Goal<N, E>,
                                                   ind: &EdgeOpsGenome)
                                                   -> MultiObjective3<f32> {
    let g = ind.to_graph(max_degree);
    MultiObjective3::from((apply_fitness_function(fitness_functions.0, goal, &g),
                           apply_fitness_function(fitness_functions.1, goal, &g),
                           apply_fitness_function(fitness_functions.2, goal, &g)))

}

impl<N:Clone+Sync+Default,E:Clone+Sync+Default> FitnessEval<EdgeOpsGenome, MultiObjective3<f32>> for MyEval<N,E> {
    fn fitness(&mut self, pop: &[EdgeOpsGenome]) -> Vec<MultiObjective3<f32>> {
        let pool = &mut self.pool;
        let goal = &self.goal;

        let fitness_functions = self.fitness_functions;
        let max_degree = self.max_degree;

        crossbeam::scope(|scope| {
            pool.map(scope,
                     pop,
                     |ind| fitness(fitness_functions, max_degree, goal, ind))
                .collect()
        })
    }
}

#[derive(Debug)]
struct Stat<T: Debug> {
    min: T,
    max: T,
    avg: T,
}

#[allow(non_snake_case)]
fn main() {
    let ncpus = num_cpus::get();
    println!("Using {} CPUs", ncpus);
    let matches = App::new("nsga2_edge")
                      .arg(Arg::with_name("NGEN")
                               .long("ngen")
                               .help("Number of generations to run")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("OPS")
                               .long("ops")
                               .help("Edge operation and weight specification, e.g. \
                                      Dup:1,Split:3,Parent:2")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("GRAPH")
                               .long("graph")
                               .help("SGF graph file")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("MU")
                               .long("mu")
                               .help("Population size")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("LAMBDA")
                               .long("lambda")
                               .help("Size of offspring population")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("K")
                               .long("k")
                               .help("Tournament selection (default: 2)")
                               .takes_value(true))
                      .arg(Arg::with_name("SEED")
                               .long("seed")
                               .help("Seed value for Rng")
                               .takes_value(true))
                      .arg(Arg::with_name("PMUT")
                               .long("pmut")
                               .help("Probability for element mutation")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("ILEN")
                               .long("ilen")
                               .help("Initial genome length (random range)")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("OBJECTIVES")
                               .long("objectives")
                               .help("Specify 3 objective functions comma separated (null, cc, \
                                      scc, nm, td), e.g. cc,nm,td")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("THRESHOLD")
                               .long("threshold")
                               .help("Abort if for all fitness[i] <= value[i] (default: 0,0,0)")
                               .takes_value(true)
                               .required(false))
                      .arg(Arg::with_name("MAX_DEGREE")
                               .long("max-degree")
                               .help("Maximum degree (default: 20)")
                               .takes_value(true)
                               .required(false))
                      .get_matches();

    // number of generations
    let NGEN: usize = FromStr::from_str(matches.value_of("NGEN").unwrap()).unwrap();
    println!("NGEN: {}", NGEN);

    // size of population
    let MU: usize = FromStr::from_str(matches.value_of("MU").unwrap()).unwrap();
    println!("MU: {}", MU);

    // size of offspring population
    let LAMBDA: usize = FromStr::from_str(matches.value_of("LAMBDA").unwrap()).unwrap();
    println!("LAMBDA: {}", LAMBDA);

    // tournament selection
    let K: usize = FromStr::from_str(matches.value_of("K").unwrap_or("2")).unwrap();
    assert!(K > 0);
    println!("K: {}", K);

    let seed: Vec<u64>;
    if let Some(seed_str) = matches.value_of("SEED") {
        seed = seed_str.split(",").map(|s| FromStr::from_str(s).unwrap()).collect();
    } else {
        println!("Use OsRng to generate seed..");
        let mut rng = OsRng::new().unwrap();
        seed = (0..2).map(|_| rng.next_u64()).collect();
    }
    println!("SEED: {:?}", seed);

    let PMUT: f32 = FromStr::from_str(matches.value_of("PMUT").unwrap()).unwrap();
    println!("PMUT: {:?}", PMUT);

    let ilen_str = matches.value_of("ILEN").unwrap();
    let ilen: Vec<usize> = ilen_str.split(",").map(|s| FromStr::from_str(s).unwrap()).collect();
    assert!(ilen.len() == 1 || ilen.len() == 2);

    let ilen_from = ilen[0];
    let ilen_to = if ilen.len() == 1 {
        ilen[0]
    } else {
        ilen[1]
    };
    assert!(ilen_from <= ilen_to);

    // read objective functions
    let mut objectives_arr: Vec<FitnessFunction> = Vec::new();
    for s in matches.value_of("OBJECTIVES").unwrap().split(",") {
        objectives_arr.push(match s {
            "null" => {
                FitnessFunction::Null
            }
            "cc" => {
                FitnessFunction::ConnectedComponents
            }
            "scc" => {
                FitnessFunction::StronglyConnectedComponents
            }
            "nm" => {
                FitnessFunction::NeighborMatching
            }
            "td" => {
                FitnessFunction::TriadicDistance
            }
            _ => {
                panic!("Invalid objective function");
            }
        });
    }

    while objectives_arr.len() < 3 {
        objectives_arr.push(FitnessFunction::Null);
    }

    if objectives_arr.len() > 3 {
        panic!("Max 3 objectives allowed");
    }

    println!("objectives={:?}", objectives_arr);

    // read objective functions
    let mut threshold_arr: Vec<f32> = Vec::new();
    for s in matches.value_of("THRESHOLD").unwrap_or("").split(",") {
        let value: f32 = FromStr::from_str(s).unwrap();
        threshold_arr.push(value);
    }

    while threshold_arr.len() < 3 {
        threshold_arr.push(0.0);
    }

    if threshold_arr.len() > 3 {
        panic!("Max 3 threshold values allowed");
    }

    // read graph
    let graph_file = matches.value_of("GRAPH").unwrap();
    println!("Using graph file: {}", graph_file);
    let graph = {
        let mut f = BufReader::new(File::open(graph_file).unwrap());
        let graph: Graph<Unweighted, Unweighted, Directed> = PetgraphReader::from_sgf(&mut f);
        graph
    };

    let MAX_DEGREE: u32 = FromStr::from_str(matches.value_of("MAX_DEGREE").unwrap_or("20"))
                              .unwrap();

    // Parse weighted operation choice from command line
    let ops = Toolbox::parse_weighted_op_choice_list(matches.value_of("OPS").unwrap()).unwrap();
    println!("edge ops: {:?}", ops);
    let mut toolbox = Toolbox::new(Probability::new(PMUT), &ops[..]);

    let mut evaluator = MyEval {
        goal: Goal::new(OptDenseDigraph::from(graph)),
        pool: Pool::new(ncpus),
        max_degree: MAX_DEGREE,
        fitness_functions: (objectives_arr[0], objectives_arr[1], objectives_arr[2]),
    };

    assert!(seed.len() == 2);
    let mut rng: PcgRng = SeedableRng::from_seed([seed[0], seed[1]]);

    // create initial random population
    let initial_population: Vec<EdgeOpsGenome> = (0..MU)
                                                     .map(|_| {
                                                         let len = if ilen_from == ilen_to {
                                                             ilen_from
                                                         } else {
                                                             Range::new(ilen_from, ilen_to)
                                                                 .ind_sample(&mut rng)
                                                         };
                                                         toolbox.random_genome(&mut rng, len)
                                                     })
                                                     .collect();

    // evaluate fitness
    let fitness: Vec<_> = evaluator.fitness(&initial_population[..]);
    assert!(fitness.len() == initial_population.len());

    let mut pop = initial_population;
    let mut fit = fitness;

    for i in 0..NGEN {
        println!("Iteration: {}", i);
        let (new_pop, new_fit) = nsga2::iterate(&mut rng,
                                                pop,
                                                fit,
                                                &mut evaluator,
                                                MU,
                                                LAMBDA,
                                                K,
                                                &mut toolbox);
        pop = new_pop;
        fit = new_fit;

        let stats: Vec<Stat<f32>> = [0, 1, 2]
                                        .iter()
                                        .map(|&i| {
                                            let min = fit.iter().fold(f32::INFINITY, |acc, f| {
                                                let x = f.objectives[i];
                                                if x < acc {
                                                    x
                                                } else {
                                                    acc
                                                }
                                            });
                                            let max = fit.iter().fold(-f32::INFINITY, |acc, f| {
                                                let x = f.objectives[i];
                                                if x > acc {
                                                    x
                                                } else {
                                                    acc
                                                }
                                            });
                                            let sum = fit.iter()
                                                         .fold(0.0, |acc, f| acc + f.objectives[i]);
                                            Stat {
                                                min: min,
                                                max: max,
                                                avg: sum / fit.len() as f32,
                                            }
                                        })
                                        .collect();
        println!("stats: {:?}", stats);
        let mut found_optimum = false;
        for f in fit.iter() {
            if f.objectives[0] <= threshold_arr[0] && f.objectives[1] <= threshold_arr[1] &&
               f.objectives[2] <= threshold_arr[2] {
                found_optimum = true;
                break;
            }
        }
        if found_optimum {
            println!("Found premature optimum in Iteration {}", i);
            break;
        }
    }
    println!("===========================================================");

    // finally evaluate rank and crowding distance (using select()).
    let rank_dist = nsga2::select(&fit[..], MU);
    assert!(rank_dist.len() == MU);

    let mut j = 0;

    for rd in rank_dist.iter() {
        if rd.rank == 0 {
            println!("-------------------------------------------");
            println!("rd: {:?}", rd);
            println!("fitness: {:?}", fit[rd.idx]);
            println!("genome: {:?}", pop[rd.idx]);

            // if fit[rd.idx].objectives[0] < 1.0 {
            draw_graph((&pop[rd.idx].to_graph(evaluator.max_degree)).ref_graph(),
                       // XXX: name
                       &format!("nsga2edgeops_g{}_f{}_i{}.svg",
                                NGEN,
                                fit[rd.idx].objectives[1] as usize,
                                j));
            j += 1;
            // }
        }
    }
}
<|endoftext|>


<commit_before>mod example {
	use std::sync::{ Mutex };

	lazy_static! { static ref List: Mutex<Vec<Component>> = Mutex::new(vec![]); }

	/// Example component
	struct Component {
		x: f64,
		y: f64
	}

	impl Component {
		fn create(x: f64, y: f64) {
			let c = Component { x: x, y: y };
			let mut list = List.lock().unwrap();
			list.push(c);
		}
	}
}<commit_msg>document strategy<commit_after>mod example {
	use std::sync::{ Mutex };

	/// The coolest way to implement this would be to have two lists references.
	/// One mutable, one is immutable.
	///
	/// The event system will have a trigger function that takes an
	/// argument that decides if the event will be triggered this tick
	/// or the next tick. The default should be next tick to prevent
	/// inadvertent infinite loop bugs.
	///
	/// The event system can then schedule all events handlers that
	/// do not mutate components to be run in parallel.
	/// All handlers that do mutate the component list can be run
	/// sequentially after that. Perhaps we can even schedule
	/// those with non-overlapping mutable component lists to
	/// run in parallel.
	///
	/// So the next step at least is to make a mutable list iterator
	/// that we can use and have all events be executed sequentially.
	///
	/// An easy way could be that the event loop is responsible for taking
	/// the lock and all the functions here take a reference to the unlocked
	/// component list given to them by the event system.

	lazy_static! { static ref List: Mutex<Vec<(usize, Component)>> = Mutex::new(vec![]); }

	/// Example component
	struct Component {
		x: f64,
		y: f64
	}

	fn create(entity: usize, x: f64, y: f64) {
		let c = Component { x: x, y: y };
		let mut list = List.lock().unwrap();
		list.push((entity,c));
	}

	fn iter_mut() -> {

	}
}
<|endoftext|>