<commit_before>extern crate env_logger;
extern crate futures;
extern crate thrussh;
extern crate thrussh_keys;
extern crate tokio;
use anyhow::Result;
use futures::Future;
use std::io::{Read, Write};
use std::sync::Arc;
use thrussh::*;
// use thrussh::server::{Auth, Session};
use thrussh_keys::*;

struct Client {}

impl client::Handler for Client {
    type Error = thrussh::Error;
    type FutureUnit = futures::future::Ready<Result<(Self, client::Session), Self::Error>>;
    type FutureBool = futures::future::Ready<Result<(Self, bool), Self::Error>>;

    fn finished_bool(self, b: bool) -> Self::FutureBool {
        futures::future::ready(Ok((self, b)))
    }
    fn finished(self, session: client::Session) -> Self::FutureUnit {
        println!("FINISHED");
        futures::future::ready(Ok((self, session)))
    }
    fn check_server_key(self, server_public_key: &key::PublicKey) -> Self::FutureBool {
        println!(
            "check_server_key: {:?}",
            server_public_key.public_key_base64()
        );
        // TODO: compare against preshared key?
        self.finished_bool(true)
    }
    /*
    // FIXME: this here makes the session hang
    fn channel_open_confirmation(
        self,
        channel: ChannelId,
        max_packet_size: u32,
        window_size: u32,
        session: client::Session,
    ) -> Self::FutureUnit {
        println!("channel_open_confirmation: {:?}", channel);
        self.finished(session)
    }
    fn data(self, channel: ChannelId, data: &[u8], session: client::Session) -> Self::FutureUnit {
        println!(
            "data on channel {:?}: {:?}",
            channel,
            std::str::from_utf8(data)
        );
        self.finished(session)
    }
    */
}

// from https://nest.pijul.com/pijul/thrussh/discussions/20
pub struct Session {
    session: client::Handle<Client>,
}

impl Session {
    async fn connect(
        key_file: &str,
        user: impl Into<String>,
        addr: impl std::net::ToSocketAddrs,
    ) -> Result<Self> {
        let key_pair = thrussh_keys::load_secret_key(key_file, None)?;
        // TODO: import openssl for RSA key support
        /*
        let key_pair = key::KeyPair::RSA {
            key: openssl::rsa::Rsa::private_key_from_pem(pem)?,
            hash: key::SignatureHash::SHA2_512,
        };
        */
        let config = client::Config::default();
        let config = Arc::new(config);
        let sh = Client {};
        let mut agent = agent::client::AgentClient::connect_env().await?;
        agent.add_identity(&key_pair, &[]).await?;
        let mut identities = agent.request_identities().await?;
        let mut session = client::connect(config, addr, sh).await?;
        let pubkey = identities.pop().unwrap();
        let (_, auth_res) = session.authenticate_future(user, pubkey, agent).await;
        let _auth_res = auth_res?;
        Ok(Self { session })
    }

    async fn call(&mut self, command: &str) -> Result<CommandResult> {
        let mut channel = self.session.channel_open_session().await?;
        channel.exec(true, command).await?;
        let mut output = Vec::new();
        let mut code = None;
        while let Some(msg) = channel.wait().await {
            match msg {
                thrussh::ChannelMsg::Data { ref data } => {
                    output.write_all(&data).unwrap();
                }
                thrussh::ChannelMsg::ExitStatus { exit_status } => {
                    code = Some(exit_status);
                }
                _ => {}
            }
        }
        Ok(CommandResult { output, code })
    }

    async fn close(&mut self) -> Result<()> {
        self.session
            .disconnect(Disconnect::ByApplication, "", "English")
            .await?;
        Ok(())
    }
}

struct CommandResult {
    output: Vec<u8>,
    code: Option<u32>,
}

impl CommandResult {
    fn output(&self) -> String {
        String::from_utf8_lossy(&self.output).into()
    }

    fn success(&self) -> bool {
        self.code == Some(0)
    }
}

#[tokio::main]
pub async fn ssh() -> Result<()> {
    let host = "localhost:22";
    let key_file = "/home/dama/.ssh/id_ed25519";

    let user: String;
    match std::env::var("USER") {
        Ok(val) => user = val,
        Err(e) => {
            user = "root".to_string();
            println!("No USER set({}); going with {}", e, user);
        }
    }

    let mut ssh = Session::connect(key_file, user, host).await?;
    let r = ssh.call("whoami").await?;
    assert!(r.success());
    println!("Who am I, anyway? {:?}", r.output());
    ssh.close().await?;
    Ok(())
}
<commit_msg>Rewrite in Rust: paste server example from docs<commit_after>extern crate env_logger;
extern crate futures;
extern crate thrussh;
extern crate thrussh_keys;
extern crate tokio;
use anyhow::Result;
use futures::Future;
use std::io::{Read, Write};
use std::sync::Arc;
use thrussh::*;
// use thrussh::server::{Auth, Session};
use thrussh_keys::*;

struct Client {}

impl client::Handler for Client {
    type Error = thrussh::Error;
    type FutureUnit = futures::future::Ready<Result<(Self, client::Session), Self::Error>>;
    type FutureBool = futures::future::Ready<Result<(Self, bool), Self::Error>>;

    fn finished_bool(self, b: bool) -> Self::FutureBool {
        futures::future::ready(Ok((self, b)))
    }
    fn finished(self, session: client::Session) -> Self::FutureUnit {
        println!("FINISHED");
        futures::future::ready(Ok((self, session)))
    }
    fn check_server_key(self, server_public_key: &key::PublicKey) -> Self::FutureBool {
        println!(
            "check_server_key: {:?}",
            server_public_key.public_key_base64()
        );
        // TODO: compare against preshared key?
        self.finished_bool(true)
    }
    /*
    // FIXME: this here makes the session hang
    fn channel_open_confirmation(
        self,
        channel: ChannelId,
        max_packet_size: u32,
        window_size: u32,
        session: client::Session,
    ) -> Self::FutureUnit {
        println!("channel_open_confirmation: {:?}", channel);
        self.finished(session)
    }
    fn data(self, channel: ChannelId, data: &[u8], session: client::Session) -> Self::FutureUnit {
        println!(
            "data on channel {:?}: {:?}",
            channel,
            std::str::from_utf8(data)
        );
        self.finished(session)
    }
    */
}

// from https://nest.pijul.com/pijul/thrussh/discussions/20
pub struct Session {
    session: client::Handle<Client>,
}

impl Session {
    async fn connect(
        key_file: &str,
        user: impl Into<String>,
        addr: impl std::net::ToSocketAddrs,
    ) -> Result<Self> {
        let key_pair = thrussh_keys::load_secret_key(key_file, None)?;
        // TODO: import openssl for RSA key support
        /*
        let key_pair = key::KeyPair::RSA {
            key: openssl::rsa::Rsa::private_key_from_pem(pem)?,
            hash: key::SignatureHash::SHA2_512,
        };
        */
        let config = client::Config::default();
        let config = Arc::new(config);
        let sh = Client {};
        let mut agent = agent::client::AgentClient::connect_env().await?;
        agent.add_identity(&key_pair, &[]).await?;
        let mut identities = agent.request_identities().await?;
        let mut session = client::connect(config, addr, sh).await?;
        let pubkey = identities.pop().unwrap();
        let (_, auth_res) = session.authenticate_future(user, pubkey, agent).await;
        let _auth_res = auth_res?;
        Ok(Self { session })
    }

    async fn call(&mut self, command: &str) -> Result<CommandResult> {
        let mut channel = self.session.channel_open_session().await?;
        channel.exec(true, command).await?;
        let mut output = Vec::new();
        let mut code = None;
        while let Some(msg) = channel.wait().await {
            match msg {
                thrussh::ChannelMsg::Data { ref data } => {
                    output.write_all(&data).unwrap();
                }
                thrussh::ChannelMsg::ExitStatus { exit_status } => {
                    code = Some(exit_status);
                }
                _ => {}
            }
        }
        Ok(CommandResult { output, code })
    }

    async fn close(&mut self) -> Result<()> {
        self.session
            .disconnect(Disconnect::ByApplication, "", "English")
            .await?;
        Ok(())
    }
}

struct CommandResult {
    output: Vec<u8>,
    code: Option<u32>,
}

impl CommandResult {
    fn output(&self) -> String {
        String::from_utf8_lossy(&self.output).into()
    }

    fn success(&self) -> bool {
        self.code == Some(0)
    }
}

#[tokio::main]
pub async fn ssh() -> Result<()> {
    let host = "localhost:22";
    let key_file = "/home/dama/.ssh/id_ed25519";

    let user: String;
    match std::env::var("USER") {
        Ok(val) => user = val,
        Err(e) => {
            user = "root".to_string();
            println!("No USER set({}); going with {}", e, user);
        }
    }

    let mut ssh = Session::connect(key_file, user, host).await?;
    let r = ssh.call("whoami").await?;
    assert!(r.success());
    println!("Who am I, anyway? {:?}", r.output());
    ssh.close().await?;
    Ok(())
}

/*
#[derive(Clone)]
struct Server {
    client_pubkey: Arc<thrussh_keys::key::PublicKey>,
    clients: Arc<Mutex<HashMap<(usize, ChannelId), thrussh::server::Handle>>>,
    id: usize,
}

impl server::Server for Server {
    type Handler = Self;
    fn new(&mut self, _: Option<std::net::SocketAddr>) -> Self {
        let s = self.clone();
        self.id += 1;
        s
    }
}

impl server::Handler for Server {
    type Error = anyhow::Error;
    type FutureAuth = futures::future::Ready<Result<(Self, server::Auth), anyhow::Error>>;
    type FutureUnit = futures::future::Ready<Result<(Self, Session), anyhow::Error>>;
    type FutureBool = futures::future::Ready<Result<(Self, Session, bool), anyhow::Error>>;

    fn finished_auth(mut self, auth: Auth) -> Self::FutureAuth {
        futures::future::ready(Ok((self, auth)))
    }
    fn finished_bool(self, b: bool, s: Session) -> Self::FutureBool {
        futures::future::ready(Ok((self, s, b)))
    }
    fn finished(self, s: Session) -> Self::FutureUnit {
        futures::future::ready(Ok((self, s)))
    }
    fn channel_open_session(self, channel: ChannelId, session: Session) -> Self::FutureUnit {
        {
            let mut clients = self.clients.lock().unwrap();
            clients.insert((self.id, channel), session.handle());
        }
        self.finished(session)
    }
    fn auth_publickey(self, _: &str, _: &key::PublicKey) -> Self::FutureAuth {
        self.finished_auth(server::Auth::Accept)
    }
    fn data(self, channel: ChannelId, data: &[u8], mut session: Session) -> Self::FutureUnit {
        {
            let mut clients = self.clients.lock().unwrap();
            for ((id, channel), ref mut s) in clients.iter_mut() {
                if *id != self.id {
                    s.data(*channel, CryptoVec::from_slice(data));
                }
            }
        }
        session.data(channel, CryptoVec::from_slice(data));
        self.finished(session)
    }
}

#[tokio::main]
async fn main() {
    let client_key = thrussh_keys::key::KeyPair::generate_ed25519().unwrap();
    let client_pubkey = Arc::new(client_key.clone_public_key());
    let mut config = thrussh::server::Config::default();
    config.connection_timeout = Some(std::time::Duration::from_secs(3));
    config.auth_rejection_time = std::time::Duration::from_secs(3);
    config.keys.push(thrussh_keys::key::KeyPair::generate_ed25519().unwrap());
    let config = Arc::new(config);
    let sh = Server{
        client_pubkey,
        clients: Arc::new(Mutex::new(HashMap::new())),
        id: 0
    };
    tokio::time::timeout(
       std::time::Duration::from_secs(1),
       thrussh::server::run(config, "0.0.0.0:2222", sh)
    ).await.unwrap_or(Ok(()));
}
*/
<|endoftext|>


<commit_before>/// Components are types that associate an entity with a set of properties
/// for each type a list of all its instances is kept so they can be iterated
/// over efficiently by systems.
///
/// Whenever a new component is created it is added to a list.
///
/// component! { physics, body: physics.RigidBody, physics_id: physics.ID }
///

use std::collections::{ HashSet, HashMap };
use std::any::{ Any, TypeId };
use shared_mutex::{ SharedMutex, SharedMutexReadGuard, SharedMutexWriteGuard };

pub mod example;

#[derive(Clone)]
pub struct Component {
	pub name: TypeId,
	pub get_component_list: fn () -> Box<Any + 'static>
}

lazy_static! {
	pub static ref COMPONENTS: SharedMutex<HashMap<TypeId, Component>> = SharedMutex::new(HashMap::new());
}

pub fn register(component : Component) {
	let mut components = COMPONENTS.write().expect("COMPONENTS lock corrupted");
	components.insert(component.name, component);
}

pub fn get_components_lock<'mutex, T>(id : TypeId) -> SharedMutexReadGuard<'mutex, T> where T : Any {
	let components = COMPONENTS.read().expect("COMPONENTS lock corrupted");
	let component = components.get(&id).expect("Unknown component type requested");
	let component_list_any = (component.get_component_list)();
	let component_list : &'static SharedMutex<T> = *component_list_any.downcast_ref().expect("Invalid type for components list");
	return component_list.read().expect("Component LIST lock corrupted");
}

#[macro_export]
macro_rules! component {
	( $component_name:ident , $( $name:ident : $field:ty ),* ) => (
		pub mod $component_name {
			use shared_mutex::{ SharedMutex, SharedMutexWriteGuard };
			use entity_rust::entities::{ ComponentList, EntityID };

			#[derive(Default)]
			pub struct Component {
				pub $($name : $field),*
			}

			lazy_static! {
				pub static ref LIST: SharedMutex<ComponentList<Component>> = SharedMutex::new(Vec::new());
			}

			pub fn create(mut list: SharedMutexWriteGuard<ComponentList<Component>>, entity: EntityID, $($name : $field),*) {
				let c = Component { $($name : $name),* };
				list.push((entity,c));
			}
		}
	)
}
<commit_msg>add getter for write lock<commit_after>/// Components are types that associate an entity with a set of properties
/// for each type a list of all its instances is kept so they can be iterated
/// over efficiently by systems.
///
/// Whenever a new component is created it is added to a list.
///
/// component! { physics, body: physics.RigidBody, physics_id: physics.ID }
///

use std::collections::{ HashSet, HashMap };
use std::any::{ Any, TypeId };
use shared_mutex::{ SharedMutex, SharedMutexReadGuard, SharedMutexWriteGuard };

pub mod example;

#[derive(Clone)]
pub struct Component {
	pub name: TypeId,
	pub get_component_list: fn () -> Box<Any + 'static>
}

lazy_static! {
	pub static ref COMPONENTS: SharedMutex<HashMap<TypeId, Component>> = SharedMutex::new(HashMap::new());
}

pub fn register(component : Component) {
	let mut components = COMPONENTS.write().expect("COMPONENTS lock corrupted");
	components.insert(component.name, component);
}

pub fn get_components_lock<'mutex, T>(id : TypeId) -> SharedMutexReadGuard<'mutex, T> where T : Any {
	let components = COMPONENTS.read().expect("COMPONENTS lock corrupted");
	let component = components.get(&id).expect("Unknown component type requested");
	let component_list_any = (component.get_component_list)();
	let component_list : &'static SharedMutex<T> = *component_list_any.downcast_ref().expect("Invalid type for components list");
	return component_list.read().expect("Component LIST lock corrupted");
}

pub fn get_components_lock_mut<'mutex, T>(id : TypeId) -> SharedMutexWriteGuard<'mutex, T> where T : Any {
	let components = COMPONENTS.read().expect("COMPONENTS lock corrupted");
	let component = components.get(&id).expect("Unknown component type requested");
	let component_list_any = (component.get_component_list)();
	let component_list : &'static SharedMutex<T> = *component_list_any.downcast_ref().expect("Invalid type for components list");
	return component_list.write().expect("Component LIST lock corrupted");
}

#[macro_export]
macro_rules! component {
	( $component_name:ident , $( $name:ident : $field:ty ),* ) => (
		pub mod $component_name {
			use shared_mutex::{ SharedMutex, SharedMutexWriteGuard };
			use entity_rust::entities::{ ComponentList, EntityID };

			#[derive(Default)]
			pub struct Component {
				pub $($name : $field),*
			}

			lazy_static! {
				pub static ref LIST: SharedMutex<ComponentList<Component>> = SharedMutex::new(Vec::new());
			}

			pub fn create(mut list: SharedMutexWriteGuard<ComponentList<Component>>, entity: EntityID, $($name : $field),*) {
				let c = Component { $($name : $name),* };
				list.push((entity,c));
			}
		}
	)
}
<|endoftext|>


<commit_before>use {Error, Result, Statement};
use types::ToSql;

impl<'conn> Statement<'conn> {
    /// Execute an INSERT and return the ROWID.
    ///
    /// # Failure
    /// Will return `Err` if no row is inserted or many rows are inserted.
    pub fn insert(&mut self, params: &[&ToSql]) -> Result<i64> {
        // Some non-insertion queries could still return 1 change (an UPDATE, for example), so
        // to guard against that we can check that the connection's last_insert_rowid() changes
        // after we execute the statement.
        let prev_rowid = self.conn.last_insert_rowid();
        let changes = try!(self.execute(params));
        let new_rowid = self.conn.last_insert_rowid();
        match changes {
            1 if prev_rowid != new_rowid => Ok(new_rowid),
            1 if prev_rowid == new_rowid => Err(Error::StatementFailedToInsertRow),
            _ => Err(Error::StatementChangedRows(changes))
        }
    }

    /// Return `true` if a query in the SQL statement it executes returns one or more rows
    /// and `false` if the SQL returns an empty set.
    pub fn exists(&mut self, params: &[&ToSql]) -> Result<bool> {
        self.reset_if_needed();
        let mut rows = try!(self.query(params));
        match rows.next() {
            Some(_) => Ok(true),
            None => Ok(false),
        }
    }
}

#[cfg(test)]
mod test {
    use {Connection, Error};

    #[test]
    fn test_insert() {
        let db = Connection::open_in_memory().unwrap();
        db.execute_batch("CREATE TABLE foo(x INTEGER UNIQUE)").unwrap();
        let mut stmt = db.prepare("INSERT OR IGNORE INTO foo (x) VALUES (?)").unwrap();
        assert_eq!(stmt.insert(&[&1i32]).unwrap(), 1);
        assert_eq!(stmt.insert(&[&2i32]).unwrap(), 2);
        match stmt.insert(&[&1i32]).unwrap_err() {
            Error::StatementChangedRows(0) => (),
            err => panic!("Unexpected error {}", err),
        }
        let mut multi = db.prepare("INSERT INTO foo (x) SELECT 3 UNION ALL SELECT 4").unwrap();
        match multi.insert(&[]).unwrap_err() {
            Error::StatementChangedRows(2) => (),
            err => panic!("Unexpected error {}", err),
        }
    }

    #[test]
    fn test_insert_failures() {
        let db = Connection::open_in_memory().unwrap();
        db.execute_batch("CREATE TABLE foo(x INTEGER UNIQUE)").unwrap();
        let mut insert = db.prepare("INSERT INTO foo (x) VALUES (?)").unwrap();
        let mut update = db.prepare("UPDATE foo SET x = ?").unwrap();

        assert_eq!(insert.insert(&[&1i32]).unwrap(), 1);

        match update.insert(&[&2i32]) {
            Err(Error::StatementFailedToInsertRow) => (),
            r => panic!("Unexpected result {:?}", r),
        }
    }

    #[test]
    fn test_exists() {
        let db = Connection::open_in_memory().unwrap();
        let sql = "BEGIN;
                   CREATE TABLE foo(x INTEGER);
                   INSERT INTO foo VALUES(1);
                   INSERT INTO foo VALUES(2);
                   END;";
        db.execute_batch(sql).unwrap();
        let mut stmt = db.prepare("SELECT 1 FROM foo WHERE x = ?").unwrap();
        assert!(stmt.exists(&[&1i32]).unwrap());
        assert!(stmt.exists(&[&2i32]).unwrap());
        assert!(!stmt.exists(&[&0i32]).unwrap());
    }
}
<commit_msg>rustfmt<commit_after>use {Error, Result, Statement};
use types::ToSql;

impl<'conn> Statement<'conn> {
    /// Execute an INSERT and return the ROWID.
    ///
    /// # Failure
    /// Will return `Err` if no row is inserted or many rows are inserted.
    pub fn insert(&mut self, params: &[&ToSql]) -> Result<i64> {
        // Some non-insertion queries could still return 1 change (an UPDATE, for example), so
        // to guard against that we can check that the connection's last_insert_rowid() changes
        // after we execute the statement.
        let prev_rowid = self.conn.last_insert_rowid();
        let changes = try!(self.execute(params));
        let new_rowid = self.conn.last_insert_rowid();
        match changes {
            1 if prev_rowid != new_rowid => Ok(new_rowid),
            1 if prev_rowid == new_rowid => Err(Error::StatementFailedToInsertRow),
            _ => Err(Error::StatementChangedRows(changes)),
        }
    }

    /// Return `true` if a query in the SQL statement it executes returns one or more rows
    /// and `false` if the SQL returns an empty set.
    pub fn exists(&mut self, params: &[&ToSql]) -> Result<bool> {
        self.reset_if_needed();
        let mut rows = try!(self.query(params));
        match rows.next() {
            Some(_) => Ok(true),
            None => Ok(false),
        }
    }
}

#[cfg(test)]
mod test {
    use {Connection, Error};

    #[test]
    fn test_insert() {
        let db = Connection::open_in_memory().unwrap();
        db.execute_batch("CREATE TABLE foo(x INTEGER UNIQUE)").unwrap();
        let mut stmt = db.prepare("INSERT OR IGNORE INTO foo (x) VALUES (?)").unwrap();
        assert_eq!(stmt.insert(&[&1i32]).unwrap(), 1);
        assert_eq!(stmt.insert(&[&2i32]).unwrap(), 2);
        match stmt.insert(&[&1i32]).unwrap_err() {
            Error::StatementChangedRows(0) => (),
            err => panic!("Unexpected error {}", err),
        }
        let mut multi = db.prepare("INSERT INTO foo (x) SELECT 3 UNION ALL SELECT 4").unwrap();
        match multi.insert(&[]).unwrap_err() {
            Error::StatementChangedRows(2) => (),
            err => panic!("Unexpected error {}", err),
        }
    }

    #[test]
    fn test_insert_failures() {
        let db = Connection::open_in_memory().unwrap();
        db.execute_batch("CREATE TABLE foo(x INTEGER UNIQUE)").unwrap();
        let mut insert = db.prepare("INSERT INTO foo (x) VALUES (?)").unwrap();
        let mut update = db.prepare("UPDATE foo SET x = ?").unwrap();

        assert_eq!(insert.insert(&[&1i32]).unwrap(), 1);

        match update.insert(&[&2i32]) {
            Err(Error::StatementFailedToInsertRow) => (),
            r => panic!("Unexpected result {:?}", r),
        }
    }

    #[test]
    fn test_exists() {
        let db = Connection::open_in_memory().unwrap();
        let sql = "BEGIN;
                   CREATE TABLE foo(x INTEGER);
                   INSERT INTO foo VALUES(1);
                   INSERT INTO foo VALUES(2);
                   END;";
        db.execute_batch(sql).unwrap();
        let mut stmt = db.prepare("SELECT 1 FROM foo WHERE x = ?").unwrap();
        assert!(stmt.exists(&[&1i32]).unwrap());
        assert!(stmt.exists(&[&2i32]).unwrap());
        assert!(!stmt.exists(&[&0i32]).unwrap());
    }
}
<|endoftext|>


<commit_before>use jwk::*;
use result::*;
use openssl::crypto::rsa::RSA;
use openssl::nid::Nid;
use openssl::bn::BigNum;
use rust_crypto::digest::Digest;
use rust_crypto::sha2::{Sha256, Sha384, Sha512};
use rustc_serialize::base64::*;
use algorithm::*;
use signer::*;
use header::*;
use validation::*;
use json::*;
use key_type::*;
use bignum::*;

pub fn bn_to_b64(bn: BigNum) -> JwtResult<String> {
    use Part;

    let component = BigNumComponent(bn);
    component.to_base64().map_err(From::from)
}

pub trait RsaParameters {
    fn dp(&self) -> JwtResult<BigNum>;
    fn dq(&self) -> JwtResult<BigNum>;
    fn qi(&self) -> JwtResult<BigNum>;
}

impl RsaParameters for RSA {
    fn dp(&self) -> JwtResult<BigNum> {
        unsafe {
            BigNum::new_from_ffi((*self.as_ptr()).dmp1)
        }.map_err(From::from)
    }

    fn dq(&self) -> JwtResult<BigNum> {
        unsafe {
            BigNum::new_from_ffi((*self.as_ptr()).dmq1)
        }.map_err(From::from)
    }

    fn qi(&self) -> JwtResult<BigNum> {
        unsafe {
            BigNum::new_from_ffi((*self.as_ptr()).iqmp)
        }.map_err(From::from)
    }
}

pub trait RsaJwk {
    fn convert_private_pem_to_jwk(self) -> JwtResult<Jwk>;

    fn convert_public_pem_to_jwk(self) -> JwtResult<Jwk>;
}

impl RsaJwk for RSA {
    fn convert_private_pem_to_jwk(self) -> JwtResult<Jwk> {
        let mut vs = ValidationState::new();

        if !self.has_n() {
            vs.reject("n", ValidationError::MissingRequiredValue("n".to_owned()));
        }

        if !self.has_e() {
            vs.reject("e", ValidationError::MissingRequiredValue("e".to_owned()));
        }

        if vs.valid {
            let mut jwk = Jwk::new(KeyType::RSA);

            let n = try!(self.n());
            jwk.set_value("n", &try!(bn_to_b64(n)));

            let e = try!(self.e());
            jwk.set_value("e", &try!(bn_to_b64(e)));

            let d = try!(self.d());
            jwk.set_value("d", &try!(bn_to_b64(d)));

            let p = try!(self.p());
            jwk.set_value("p", &try!(bn_to_b64(p)));

            let q = try!(self.q());
            jwk.set_value("q", &try!(bn_to_b64(q)));

            let dp = try!(self.dp());
            jwk.set_value("dp", &try!(bn_to_b64(dp)));

            let dq = try!(self.dq());
            jwk.set_value("dq", &try!(bn_to_b64(dq)));

            let qi = try!(self.qi());
            jwk.set_value("qi", &try!(bn_to_b64(qi)));

            Ok(jwk)
        } else {
            Err(From::from(ValidationError::ValidationError(vs)))
        }
    }

    fn convert_public_pem_to_jwk(self) -> JwtResult<Jwk> {
        use Part;

        let mut vs = ValidationState::new();

        if !self.has_n() {
            vs.reject("n", ValidationError::MissingRequiredValue("n".to_owned()));
        }

        if !self.has_e() {
            vs.reject("e", ValidationError::MissingRequiredValue("e".to_owned()));
        }

        if vs.valid {
            let mut jwk = Jwk::new(KeyType::RSA);

            let n = try!(self.n());
            let n_bn = BigNumComponent(n);
            let n_b64 = try!(n_bn.to_base64());
            jwk.set_value("n", &n_b64);

            let e = try!(self.e());
            let e_bn = BigNumComponent(e);
            let e_b64 = try!(e_bn.to_base64());
            jwk.set_value("e", &e_b64);

            Ok(jwk)
        } else {
            Err(From::from(ValidationError::ValidationError(vs)))
        }
    }
}

pub trait RsaKey {
    fn public_key(&self) -> JwtResult<RSA>;

    fn private_key(&self) -> JwtResult<RSA>;

    fn is_private_key(&self) -> bool;
}

impl RsaKey for Jwk {
    fn public_key(&self) -> JwtResult<RSA> {
        let n = try!(self.get_bignum_param("n"));
        let e = try!(self.get_bignum_param("e"));

        RSA::from_public_components(n, e).map_err(JwtError::from)
    }

    fn private_key(&self) -> JwtResult<RSA> {
        let n = try!(self.get_bignum_param("n"));
        let e = try!(self.get_bignum_param("e"));
        let d = try!(self.get_bignum_param("d"));
        let p = try!(self.get_bignum_param("p"));
        let q = try!(self.get_bignum_param("q"));
        let dp = try!(self.get_bignum_param("dp"));
        let dq = try!(self.get_bignum_param("dq"));
        let qi = try!(self.get_bignum_param("qi"));

        RSA::from_private_components(n, e, d, p, q, dp, dq, qi).map_err(JwtError::from)
    }

    fn is_private_key(&self) -> bool {
        self.params.contains_key("d")
    }
}

#[derive(Clone, Debug)]
pub struct RsaSigner {
    private_key: Jwk,
    algorithms: Vec<Algorithm>
}

impl RsaSigner {
    pub fn new(private_key: Jwk) -> RsaSigner {
        RsaSigner {
            algorithms: compatible_algorithms(),
            private_key: private_key,
        }
    }

    fn sign_with_digest<D: Digest>(mut d: D, nid: Nid, private_key: &Jwk, input: &[u8]) -> JwtResult<Vec<u8>> {
        let mut hash = vec![0;d.output_bytes()];

        d.input(&input);
        d.result(&mut hash);

        let rsa = try!(private_key.private_key());

        rsa.sign(nid, &hash).map_err(From::from)
    }
}

impl Signer for RsaSigner {
    fn sign(&self, header: &Header, signing_input: &[u8]) -> JwtResult<String> {
        let result = try!(match header.alg {
            Algorithm::RS256 => Self::sign_with_digest(Sha256::new(), Nid::SHA256, &self.private_key, signing_input),
            Algorithm::RS384 => Self::sign_with_digest(Sha384::new(), Nid::SHA256, &self.private_key, signing_input),
            Algorithm::RS512 => Self::sign_with_digest(Sha512::new(), Nid::SHA256, &self.private_key, signing_input),
            _ => Err(JwtError::InvalidAlgorithm("algorithm is not a supported mac: ".to_owned(), header.alg))
        });

        Ok(result.to_base64(URL_SAFE))
    }

    fn has_algorithm(&self, a: Algorithm) -> bool {
        self.algorithms.contains(&a)
    }
}

pub fn compatible_algorithms() -> Vec<Algorithm> {
    vec![Algorithm::RS256, Algorithm::RS384, Algorithm::RS512]
}

#[cfg(test)]
mod test {
    use std::fs::File;
    use std::io::Write;
    use serde_json;
    use serde;
    use rust_crypto::sha2::Sha256;
    use rust_crypto::digest::Digest;
    use openssl;

    use jwk::*;
    use json::*;
    use super::*;

    pub fn load_json<T: serde::Deserialize>(file_name: &str) -> T {
        let mut f = File::open(file_name).unwrap();

        f.read_json().unwrap()
    }

    #[test]
    pub fn test_rsa_sign_and_verify() {
        let key = load_json::<Jwk>("samples/jws/a2 - rs256/4-rsa-key.json");

        let public_key = key.public_key().unwrap();
        let private_key = key.private_key().unwrap();

        let message = load_json::<Vec<u8>>("samples/jws/a2 - rs256/3-signing-input-octets.json");

        let mut sha = Sha256::new();
        sha.input(&message);
        let mut hash = vec![0;32];
        sha.result(&mut hash);

        println!("digest: {:?}", hash);

        let sig = private_key.sign(openssl::nid::Nid::SHA256, &hash).unwrap();

        println!("{:?}", sig);

        let verified = public_key.verify(openssl::nid::Nid::SHA256, &hash, &sig).unwrap();

        assert!(verified);
    }

    #[test]
    pub fn test_write_to_pem() {
        let key = load_json::<Jwk>("samples/jws/a2 - rs256/4-rsa-key.json");

        let public_key = key.public_key().unwrap();
        let private_key = key.private_key().unwrap();

        let mut buffer = File::create("key.pem").unwrap();
        private_key.private_key_to_pem(&mut buffer).unwrap();
        buffer.flush().unwrap();

        let mut buffer = File::create("pubkey.pem").unwrap();
        public_key.public_key_to_pem(&mut buffer).unwrap();
        buffer.flush().unwrap();
    }
}<commit_msg>swap trait names<commit_after>use jwk::*;
use result::*;
use openssl::crypto::rsa::RSA;
use openssl::nid::Nid;
use openssl::bn::BigNum;
use rust_crypto::digest::Digest;
use rust_crypto::sha2::{Sha256, Sha384, Sha512};
use rustc_serialize::base64::*;
use algorithm::*;
use signer::*;
use header::*;
use validation::*;
use json::*;
use key_type::*;
use bignum::*;

pub fn bn_to_b64(bn: BigNum) -> JwtResult<String> {
    use Part;

    let component = BigNumComponent(bn);
    component.to_base64().map_err(From::from)
}

pub trait RsaParameters {
    fn dp(&self) -> JwtResult<BigNum>;
    fn dq(&self) -> JwtResult<BigNum>;
    fn qi(&self) -> JwtResult<BigNum>;
}

impl RsaParameters for RSA {
    fn dp(&self) -> JwtResult<BigNum> {
        unsafe {
            BigNum::new_from_ffi((*self.as_ptr()).dmp1)
        }.map_err(From::from)
    }

    fn dq(&self) -> JwtResult<BigNum> {
        unsafe {
            BigNum::new_from_ffi((*self.as_ptr()).dmq1)
        }.map_err(From::from)
    }

    fn qi(&self) -> JwtResult<BigNum> {
        unsafe {
            BigNum::new_from_ffi((*self.as_ptr()).iqmp)
        }.map_err(From::from)
    }
}

pub trait RsaKey {
    fn convert_private_pem_to_jwk(self) -> JwtResult<Jwk>;

    fn convert_public_pem_to_jwk(self) -> JwtResult<Jwk>;
}

impl RsaKey for RSA {
    fn convert_private_pem_to_jwk(self) -> JwtResult<Jwk> {
        let mut vs = ValidationState::new();

        if !self.has_n() {
            vs.reject("n", ValidationError::MissingRequiredValue("n".to_owned()));
        }

        if !self.has_e() {
            vs.reject("e", ValidationError::MissingRequiredValue("e".to_owned()));
        }

        if vs.valid {
            let mut jwk = Jwk::new(KeyType::RSA);

            let n = try!(self.n());
            jwk.set_value("n", &try!(bn_to_b64(n)));

            let e = try!(self.e());
            jwk.set_value("e", &try!(bn_to_b64(e)));

            let d = try!(self.d());
            jwk.set_value("d", &try!(bn_to_b64(d)));

            let p = try!(self.p());
            jwk.set_value("p", &try!(bn_to_b64(p)));

            let q = try!(self.q());
            jwk.set_value("q", &try!(bn_to_b64(q)));

            let dp = try!(self.dp());
            jwk.set_value("dp", &try!(bn_to_b64(dp)));

            let dq = try!(self.dq());
            jwk.set_value("dq", &try!(bn_to_b64(dq)));

            let qi = try!(self.qi());
            jwk.set_value("qi", &try!(bn_to_b64(qi)));

            Ok(jwk)
        } else {
            Err(From::from(ValidationError::ValidationError(vs)))
        }
    }

    fn convert_public_pem_to_jwk(self) -> JwtResult<Jwk> {
        use Part;

        let mut vs = ValidationState::new();

        if !self.has_n() {
            vs.reject("n", ValidationError::MissingRequiredValue("n".to_owned()));
        }

        if !self.has_e() {
            vs.reject("e", ValidationError::MissingRequiredValue("e".to_owned()));
        }

        if vs.valid {
            let mut jwk = Jwk::new(KeyType::RSA);

            let n = try!(self.n());
            let n_bn = BigNumComponent(n);
            let n_b64 = try!(n_bn.to_base64());
            jwk.set_value("n", &n_b64);

            let e = try!(self.e());
            let e_bn = BigNumComponent(e);
            let e_b64 = try!(e_bn.to_base64());
            jwk.set_value("e", &e_b64);

            Ok(jwk)
        } else {
            Err(From::from(ValidationError::ValidationError(vs)))
        }
    }
}

pub trait RsaJwk {
    fn public_key(&self) -> JwtResult<RSA>;

    fn private_key(&self) -> JwtResult<RSA>;

    fn is_private_key(&self) -> bool;
}

impl RsaJwk for Jwk {
    fn public_key(&self) -> JwtResult<RSA> {
        let n = try!(self.get_bignum_param("n"));
        let e = try!(self.get_bignum_param("e"));

        RSA::from_public_components(n, e).map_err(JwtError::from)
    }

    fn private_key(&self) -> JwtResult<RSA> {
        let n = try!(self.get_bignum_param("n"));
        let e = try!(self.get_bignum_param("e"));
        let d = try!(self.get_bignum_param("d"));
        let p = try!(self.get_bignum_param("p"));
        let q = try!(self.get_bignum_param("q"));
        let dp = try!(self.get_bignum_param("dp"));
        let dq = try!(self.get_bignum_param("dq"));
        let qi = try!(self.get_bignum_param("qi"));

        RSA::from_private_components(n, e, d, p, q, dp, dq, qi).map_err(JwtError::from)
    }

    fn is_private_key(&self) -> bool {
        self.params.contains_key("d")
    }
}

#[derive(Clone, Debug)]
pub struct RsaSigner {
    private_key: Jwk,
    algorithms: Vec<Algorithm>
}

impl RsaSigner {
    pub fn new(private_key: Jwk) -> RsaSigner {
        RsaSigner {
            algorithms: compatible_algorithms(),
            private_key: private_key,
        }
    }

    fn sign_with_digest<D: Digest>(mut d: D, nid: Nid, private_key: &Jwk, input: &[u8]) -> JwtResult<Vec<u8>> {
        let mut hash = vec![0;d.output_bytes()];

        d.input(&input);
        d.result(&mut hash);

        let rsa = try!(private_key.private_key());

        rsa.sign(nid, &hash).map_err(From::from)
    }
}

impl Signer for RsaSigner {
    fn sign(&self, header: &Header, signing_input: &[u8]) -> JwtResult<String> {
        let result = try!(match header.alg {
            Algorithm::RS256 => Self::sign_with_digest(Sha256::new(), Nid::SHA256, &self.private_key, signing_input),
            Algorithm::RS384 => Self::sign_with_digest(Sha384::new(), Nid::SHA256, &self.private_key, signing_input),
            Algorithm::RS512 => Self::sign_with_digest(Sha512::new(), Nid::SHA256, &self.private_key, signing_input),
            _ => Err(JwtError::InvalidAlgorithm("algorithm is not a supported mac: ".to_owned(), header.alg))
        });

        Ok(result.to_base64(URL_SAFE))
    }

    fn has_algorithm(&self, a: Algorithm) -> bool {
        self.algorithms.contains(&a)
    }
}

pub fn compatible_algorithms() -> Vec<Algorithm> {
    vec![Algorithm::RS256, Algorithm::RS384, Algorithm::RS512]
}

#[cfg(test)]
mod test {
    use std::fs::File;
    use std::io::Write;
    use serde;
    use rust_crypto::sha2::Sha256;
    use rust_crypto::digest::Digest;
    use openssl;

    use jwk::*;
    use json::*;
    use super::*;

    pub fn load_json<T: serde::Deserialize>(file_name: &str) -> T {
        let mut f = File::open(file_name).unwrap();

        f.read_json().unwrap()
    }

    #[test]
    pub fn test_rsa_sign_and_verify() {
        let key = load_json::<Jwk>("samples/jws/a2 - rs256/4-rsa-key.json");

        let public_key = key.public_key().unwrap();
        let private_key = key.private_key().unwrap();

        let message = load_json::<Vec<u8>>("samples/jws/a2 - rs256/3-signing-input-octets.json");

        let mut sha = Sha256::new();
        sha.input(&message);
        let mut hash = vec![0;32];
        sha.result(&mut hash);

        println!("digest: {:?}", hash);

        let sig = private_key.sign(openssl::nid::Nid::SHA256, &hash).unwrap();

        println!("{:?}", sig);

        let verified = public_key.verify(openssl::nid::Nid::SHA256, &hash, &sig).unwrap();

        assert!(verified);
    }

    #[test]
    pub fn test_write_to_pem() {
        let key = load_json::<Jwk>("samples/jws/a2 - rs256/4-rsa-key.json");

        let public_key = key.public_key().unwrap();
        let private_key = key.private_key().unwrap();

        let mut buffer = File::create("key.pem").unwrap();
        private_key.private_key_to_pem(&mut buffer).unwrap();
        buffer.flush().unwrap();

        let mut buffer = File::create("pubkey.pem").unwrap();
        public_key.public_key_to_pem(&mut buffer).unwrap();
        buffer.flush().unwrap();
    }
}<|endoftext|>


<commit_before>use crate::actually_private::Private;
use alloc::borrow::Cow;
use alloc::string::String;
use core::fmt::{self, Debug, Display};
use core::marker::{PhantomData, PhantomPinned};
use core::mem::MaybeUninit;
use core::pin::Pin;
use core::slice;
use core::str::{self, Utf8Error};

extern "C" {
    #[link_name = "cxxbridge1$cxx_string$init"]
    fn string_init(this: &mut MaybeUninit<CxxString>, ptr: *const u8, len: usize);
    #[link_name = "cxxbridge1$cxx_string$destroy"]
    fn string_destroy(this: &mut MaybeUninit<CxxString>);
    #[link_name = "cxxbridge1$cxx_string$data"]
    fn string_data(this: &CxxString) -> *const u8;
    #[link_name = "cxxbridge1$cxx_string$length"]
    fn string_length(this: &CxxString) -> usize;
    #[link_name = "cxxbridge1$cxx_string$push"]
    fn string_push(this: Pin<&mut CxxString>, ptr: *const u8, len: usize);
}

/// Binding to C++ `std::string`.
///
/// # Invariants
///
/// As an invariant of this API and the static analysis of the cxx::bridge
/// macro, in Rust code we can never obtain a `CxxString` by value. C++'s string
/// requires a move constructor and may hold internal pointers, which is not
/// compatible with Rust's move behavior. Instead in Rust code we will only ever
/// look at a CxxString through a reference or smart pointer, as in `&CxxString`
/// or `UniquePtr<CxxString>`.
#[repr(C)]
pub struct CxxString {
    _private: [u8; 0],
    _pinned: PhantomData<PhantomPinned>,
}

/// Construct a C++ std::string on the Rust stack.
///
/// # Syntax
///
/// In statement position:
///
/// ```
/// # use cxx::let_cxx_string;
/// # let expression = "";
/// let_cxx_string!(var = expression);
/// ```
///
/// The `expression` may have any type that implements `AsRef<[u8]>`. Commonly
/// it will be a string literal, but for example `&[u8]` and `String` would work
/// as well.
///
/// The macro expands to something resembling `let $var: Pin<&mut CxxString> =
/// /*???*/;`. The resulting [`Pin`] can be deref'd to `&CxxString` as needed.
///
/// # Example
///
/// ```
/// use cxx::{let_cxx_string, CxxString};
///
/// fn f(s: &CxxString) {/* ... */}
///
/// fn main() {
///     let_cxx_string!(s = "example");
///     f(&s);
/// }
/// ```
#[macro_export]
macro_rules! let_cxx_string {
    ($var:ident = $value:expr $(,)?) => {
        let mut cxx_stack_string = $crate::private::StackString::new();
        #[allow(unused_mut, unused_unsafe)]
        let mut $var = match $value {
            let_cxx_string => unsafe { cxx_stack_string.init(let_cxx_string) },
        };
    };
}

impl CxxString {
    /// `CxxString` is not constructible via `new`. Instead, use the
    /// [`let_cxx_string!`] macro.
    pub fn new<T: Private>() -> Self {
        unreachable!()
    }

    /// Returns the length of the string in bytes.
    ///
    /// Matches the behavior of C++ [std::string::size][size].
    ///
    /// [size]: https://en.cppreference.com/w/cpp/string/basic_string/size
    pub fn len(&self) -> usize {
        unsafe { string_length(self) }
    }

    /// Returns true if `self` has a length of zero bytes.
    ///
    /// Matches the behavior of C++ [std::string::empty][empty].
    ///
    /// [empty]: https://en.cppreference.com/w/cpp/string/basic_string/empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns a byte slice of this string's contents.
    pub fn as_bytes(&self) -> &[u8] {
        let data = self.as_ptr();
        let len = self.len();
        unsafe { slice::from_raw_parts(data, len) }
    }

    /// Produces a pointer to the first character of the string.
    ///
    /// Matches the behavior of C++ [std::string::data][data].
    ///
    /// Note that the return type may look like `const char *` but is not a
    /// `const char *` in the typical C sense, as C++ strings may contain
    /// internal null bytes. As such, the returned pointer only makes sense as a
    /// string in combination with the length returned by [`len()`][len].
    ///
    /// [data]: https://en.cppreference.com/w/cpp/string/basic_string/data
    /// [len]: #method.len
    pub fn as_ptr(&self) -> *const u8 {
        unsafe { string_data(self) }
    }

    /// Validates that the C++ string contains UTF-8 data and produces a view of
    /// it as a Rust &str, otherwise an error.
    pub fn to_str(&self) -> Result<&str, Utf8Error> {
        str::from_utf8(self.as_bytes())
    }

    /// If the contents of the C++ string are valid UTF-8, this function returns
    /// a view as a Cow::Borrowed &str. Otherwise replaces any invalid UTF-8
    /// sequences with the U+FFFD [replacement character] and returns a
    /// Cow::Owned String.
    ///
    /// [replacement character]: https://doc.rust-lang.org/std/char/constant.REPLACEMENT_CHARACTER.html
    pub fn to_string_lossy(&self) -> Cow<str> {
        String::from_utf8_lossy(self.as_bytes())
    }

    /// Appends a given string slice onto the end of this C++ string.
    pub fn push_str(self: Pin<&mut Self>, s: &str) {
        self.push_bytes(s.as_bytes());
    }

    /// Appends arbitrary bytes onto the end of this C++ string.
    pub fn push_bytes(self: Pin<&mut Self>, bytes: &[u8]) {
        unsafe { string_push(self, bytes.as_ptr(), bytes.len()) }
    }
}

impl Display for CxxString {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        Display::fmt(self.to_string_lossy().as_ref(), f)
    }
}

impl Debug for CxxString {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        Debug::fmt(self.to_string_lossy().as_ref(), f)
    }
}

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

impl PartialEq<CxxString> for str {
    fn eq(&self, other: &CxxString) -> bool {
        self.as_bytes() == other.as_bytes()
    }
}

impl PartialEq<str> for CxxString {
    fn eq(&self, other: &str) -> bool {
        self.as_bytes() == other.as_bytes()
    }
}

impl Eq for CxxString {}

impl PartialOrd for CxxString {
    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
        self.as_bytes().partial_cmp(other.as_bytes())
    }
}

impl PartialOrd<str> for CxxString {
    fn partial_cmp(&self, other: &str) -> Option<core::cmp::Ordering> {
        self.as_bytes().partial_cmp(other.as_bytes())
    }
}

impl PartialOrd<CxxString> for str {
    fn partial_cmp(&self, other: &CxxString) -> Option<core::cmp::Ordering> {
        self.as_bytes().partial_cmp(other.as_bytes())
    }
}

impl Ord for CxxString {
    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
        self.as_bytes().cmp(other.as_bytes())
    }
}

impl core::hash::Hash for CxxString {
    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
        self.as_bytes().hash(state)
    }
}

#[doc(hidden)]
#[repr(C)]
pub struct StackString {
    // Static assertions in cxx.cc validate that this is large enough and
    // aligned enough.
    space: MaybeUninit<[usize; 8]>,
}

impl StackString {
    pub fn new() -> Self {
        StackString {
            space: MaybeUninit::uninit(),
        }
    }

    pub unsafe fn init(&mut self, value: impl AsRef<[u8]>) -> Pin<&mut CxxString> {
        let value = value.as_ref();
        let this = &mut *self.space.as_mut_ptr().cast::<MaybeUninit<CxxString>>();
        string_init(this, value.as_ptr(), value.len());
        Pin::new_unchecked(&mut *this.as_mut_ptr())
    }
}

impl Drop for StackString {
    fn drop(&mut self) {
        unsafe {
            let this = &mut *self.space.as_mut_ptr().cast::<MaybeUninit<CxxString>>();
            string_destroy(this);
        }
    }
}
<commit_msg>Touch up PR 743<commit_after>use crate::actually_private::Private;
use alloc::borrow::Cow;
use alloc::string::String;
use core::cmp::Ordering;
use core::fmt::{self, Debug, Display};
use core::hash::{Hash, Hasher};
use core::marker::{PhantomData, PhantomPinned};
use core::mem::MaybeUninit;
use core::pin::Pin;
use core::slice;
use core::str::{self, Utf8Error};

extern "C" {
    #[link_name = "cxxbridge1$cxx_string$init"]
    fn string_init(this: &mut MaybeUninit<CxxString>, ptr: *const u8, len: usize);
    #[link_name = "cxxbridge1$cxx_string$destroy"]
    fn string_destroy(this: &mut MaybeUninit<CxxString>);
    #[link_name = "cxxbridge1$cxx_string$data"]
    fn string_data(this: &CxxString) -> *const u8;
    #[link_name = "cxxbridge1$cxx_string$length"]
    fn string_length(this: &CxxString) -> usize;
    #[link_name = "cxxbridge1$cxx_string$push"]
    fn string_push(this: Pin<&mut CxxString>, ptr: *const u8, len: usize);
}

/// Binding to C++ `std::string`.
///
/// # Invariants
///
/// As an invariant of this API and the static analysis of the cxx::bridge
/// macro, in Rust code we can never obtain a `CxxString` by value. C++'s string
/// requires a move constructor and may hold internal pointers, which is not
/// compatible with Rust's move behavior. Instead in Rust code we will only ever
/// look at a CxxString through a reference or smart pointer, as in `&CxxString`
/// or `UniquePtr<CxxString>`.
#[repr(C)]
pub struct CxxString {
    _private: [u8; 0],
    _pinned: PhantomData<PhantomPinned>,
}

/// Construct a C++ std::string on the Rust stack.
///
/// # Syntax
///
/// In statement position:
///
/// ```
/// # use cxx::let_cxx_string;
/// # let expression = "";
/// let_cxx_string!(var = expression);
/// ```
///
/// The `expression` may have any type that implements `AsRef<[u8]>`. Commonly
/// it will be a string literal, but for example `&[u8]` and `String` would work
/// as well.
///
/// The macro expands to something resembling `let $var: Pin<&mut CxxString> =
/// /*???*/;`. The resulting [`Pin`] can be deref'd to `&CxxString` as needed.
///
/// # Example
///
/// ```
/// use cxx::{let_cxx_string, CxxString};
///
/// fn f(s: &CxxString) {/* ... */}
///
/// fn main() {
///     let_cxx_string!(s = "example");
///     f(&s);
/// }
/// ```
#[macro_export]
macro_rules! let_cxx_string {
    ($var:ident = $value:expr $(,)?) => {
        let mut cxx_stack_string = $crate::private::StackString::new();
        #[allow(unused_mut, unused_unsafe)]
        let mut $var = match $value {
            let_cxx_string => unsafe { cxx_stack_string.init(let_cxx_string) },
        };
    };
}

impl CxxString {
    /// `CxxString` is not constructible via `new`. Instead, use the
    /// [`let_cxx_string!`] macro.
    pub fn new<T: Private>() -> Self {
        unreachable!()
    }

    /// Returns the length of the string in bytes.
    ///
    /// Matches the behavior of C++ [std::string::size][size].
    ///
    /// [size]: https://en.cppreference.com/w/cpp/string/basic_string/size
    pub fn len(&self) -> usize {
        unsafe { string_length(self) }
    }

    /// Returns true if `self` has a length of zero bytes.
    ///
    /// Matches the behavior of C++ [std::string::empty][empty].
    ///
    /// [empty]: https://en.cppreference.com/w/cpp/string/basic_string/empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns a byte slice of this string's contents.
    pub fn as_bytes(&self) -> &[u8] {
        let data = self.as_ptr();
        let len = self.len();
        unsafe { slice::from_raw_parts(data, len) }
    }

    /// Produces a pointer to the first character of the string.
    ///
    /// Matches the behavior of C++ [std::string::data][data].
    ///
    /// Note that the return type may look like `const char *` but is not a
    /// `const char *` in the typical C sense, as C++ strings may contain
    /// internal null bytes. As such, the returned pointer only makes sense as a
    /// string in combination with the length returned by [`len()`][len].
    ///
    /// [data]: https://en.cppreference.com/w/cpp/string/basic_string/data
    /// [len]: #method.len
    pub fn as_ptr(&self) -> *const u8 {
        unsafe { string_data(self) }
    }

    /// Validates that the C++ string contains UTF-8 data and produces a view of
    /// it as a Rust &str, otherwise an error.
    pub fn to_str(&self) -> Result<&str, Utf8Error> {
        str::from_utf8(self.as_bytes())
    }

    /// If the contents of the C++ string are valid UTF-8, this function returns
    /// a view as a Cow::Borrowed &str. Otherwise replaces any invalid UTF-8
    /// sequences with the U+FFFD [replacement character] and returns a
    /// Cow::Owned String.
    ///
    /// [replacement character]: https://doc.rust-lang.org/std/char/constant.REPLACEMENT_CHARACTER.html
    pub fn to_string_lossy(&self) -> Cow<str> {
        String::from_utf8_lossy(self.as_bytes())
    }

    /// Appends a given string slice onto the end of this C++ string.
    pub fn push_str(self: Pin<&mut Self>, s: &str) {
        self.push_bytes(s.as_bytes());
    }

    /// Appends arbitrary bytes onto the end of this C++ string.
    pub fn push_bytes(self: Pin<&mut Self>, bytes: &[u8]) {
        unsafe { string_push(self, bytes.as_ptr(), bytes.len()) }
    }
}

impl Display for CxxString {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        Display::fmt(self.to_string_lossy().as_ref(), f)
    }
}

impl Debug for CxxString {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        Debug::fmt(self.to_string_lossy().as_ref(), f)
    }
}

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

impl PartialEq<CxxString> for str {
    fn eq(&self, other: &CxxString) -> bool {
        self.as_bytes() == other.as_bytes()
    }
}

impl PartialEq<str> for CxxString {
    fn eq(&self, other: &str) -> bool {
        self.as_bytes() == other.as_bytes()
    }
}

impl Eq for CxxString {}

impl PartialOrd for CxxString {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.as_bytes().partial_cmp(other.as_bytes())
    }
}

impl PartialOrd<str> for CxxString {
    fn partial_cmp(&self, other: &str) -> Option<Ordering> {
        self.as_bytes().partial_cmp(other.as_bytes())
    }
}

impl PartialOrd<CxxString> for str {
    fn partial_cmp(&self, other: &CxxString) -> Option<Ordering> {
        self.as_bytes().partial_cmp(other.as_bytes())
    }
}

impl Ord for CxxString {
    fn cmp(&self, other: &Self) -> Ordering {
        self.as_bytes().cmp(other.as_bytes())
    }
}

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

#[doc(hidden)]
#[repr(C)]
pub struct StackString {
    // Static assertions in cxx.cc validate that this is large enough and
    // aligned enough.
    space: MaybeUninit<[usize; 8]>,
}

impl StackString {
    pub fn new() -> Self {
        StackString {
            space: MaybeUninit::uninit(),
        }
    }

    pub unsafe fn init(&mut self, value: impl AsRef<[u8]>) -> Pin<&mut CxxString> {
        let value = value.as_ref();
        let this = &mut *self.space.as_mut_ptr().cast::<MaybeUninit<CxxString>>();
        string_init(this, value.as_ptr(), value.len());
        Pin::new_unchecked(&mut *this.as_mut_ptr())
    }
}

impl Drop for StackString {
    fn drop(&mut self) {
        unsafe {
            let this = &mut *self.space.as_mut_ptr().cast::<MaybeUninit<CxxString>>();
            string_destroy(this);
        }
    }
}
<|endoftext|>


<commit_before>use data::timeinfo::*;
use data::pixel::*;
use data::audiosample::*;

use alloc::heap::allocate;
use bytes::Bytes;
use std::ptr::write_bytes;

error_chain! {
    errors {
        InvalidIndex
    }
}

pub struct VideoFrame {
    pub width: usize,
    pub height: usize,
    pub format: Box<Formaton>,
}

impl VideoFrame {
    fn size(&self, align: usize) -> usize {
        let mut size = 0;
        for &component in self.format.into_iter() {
            if let Some(c) = component {
                size += c.get_data_size(self.width, self.height, align);
            }
        }
        size
    }
}

pub struct AudioFrame {
    pub samples: usize,
    pub rate: usize,
    pub map: ChannelMap,
    pub format: Box<Soniton>,
}

impl AudioFrame {
    fn size(&self, align: usize) -> usize {
        unimplemented!()
    }
}

pub enum FrameKind {
    Video(VideoFrame),
    Audio(AudioFrame),
}

use self::FrameKind::*;

pub trait FrameBuffer {
    fn as_slice<'a>(&'a self, idx: usize) -> Result<&'a [u8]>;
    fn as_mut_slice<'a>(&'a mut self, idx: usize) -> Result<&'a mut [u8]>;
    fn linesize(&self, idx: usize) -> usize;
    fn count(&self) -> usize;
}

pub struct Frame {
    kind: FrameKind,
    buf: Box<FrameBuffer>,
    t: TimeInfo,
}

const ALIGNMENT: usize = 32;

struct Plane {
    buf: Bytes,
    linesize: usize,
}

struct DefaultFrameBuffer {
    buf: Bytes,
    planes: Vec<Plane>,
}

impl DefaultFrameBuffer {
    pub fn new(kind: &FrameKind) -> DefaultFrameBuffer {
        match kind {
            &Video(ref video) => {
                let size = video.size(ALIGNMENT);
                let data = unsafe { allocate(size, ALIGNMENT) };
                // unsafe { write_bytes(data, 0, size) };
                let mut buf = Bytes::from(unsafe { Vec::from_raw_parts(data, size, size) });
                let mut buffer = DefaultFrameBuffer {
                    buf: buf,
                    planes: Vec::new(),
                };
                for &component in video.format.into_iter() {
                    if let Some(c) = component {
                        let planesize = c.get_data_size(video.width, video.height, ALIGNMENT);
                        let linesize = c.get_linesize(video.width, ALIGNMENT);
                        buffer.planes.push(Plane {
                            buf: buffer.buf.split_to(planesize),
                            linesize: linesize,
                        });
                    }
                }
                buffer
            }
            _ => unimplemented!(),
        }
    }
}
<commit_msg>frame: Suppress some warnings from the wip code<commit_after>#![allow(dead_code, unused_variables)]
use data::timeinfo::*;
use data::pixel::*;
use data::audiosample::*;

use alloc::heap::allocate;
use bytes::Bytes;
// use std::ptr::write_bytes;

error_chain! {
    errors {
        InvalidIndex
    }
}

pub struct VideoFrame {
    pub width: usize,
    pub height: usize,
    pub format: Box<Formaton>,
}

impl VideoFrame {
    fn size(&self, align: usize) -> usize {
        let mut size = 0;
        for &component in self.format.into_iter() {
            if let Some(c) = component {
                size += c.get_data_size(self.width, self.height, align);
            }
        }
        size
    }
}

pub struct AudioFrame {
    pub samples: usize,
    pub rate: usize,
    pub map: ChannelMap,
    pub format: Box<Soniton>,
}

impl AudioFrame {
    fn size(&self, align: usize) -> usize {
        unimplemented!()
    }
}

pub enum FrameKind {
    Video(VideoFrame),
    Audio(AudioFrame),
}

use self::FrameKind::*;

pub trait FrameBuffer {
    fn as_slice<'a>(&'a self, idx: usize) -> Result<&'a [u8]>;
    fn as_mut_slice<'a>(&'a mut self, idx: usize) -> Result<&'a mut [u8]>;
    fn linesize(&self, idx: usize) -> usize;
    fn count(&self) -> usize;
}

pub struct Frame {
    kind: FrameKind,
    buf: Box<FrameBuffer>,
    t: TimeInfo,
}

const ALIGNMENT: usize = 32;

struct Plane {
    buf: Bytes,
    linesize: usize,
}

struct DefaultFrameBuffer {
    buf: Bytes,
    planes: Vec<Plane>,
}

impl DefaultFrameBuffer {
    pub fn new(kind: &FrameKind) -> DefaultFrameBuffer {
        match kind {
            &Video(ref video) => {
                let size = video.size(ALIGNMENT);
                let data = unsafe { allocate(size, ALIGNMENT) };
                // unsafe { write_bytes(data, 0, size) };
                let buf = Bytes::from(unsafe { Vec::from_raw_parts(data, size, size) });
                let mut buffer = DefaultFrameBuffer {
                    buf: buf,
                    planes: Vec::new(),
                };
                for &component in video.format.into_iter() {
                    if let Some(c) = component {
                        let planesize = c.get_data_size(video.width, video.height, ALIGNMENT);
                        let linesize = c.get_linesize(video.width, ALIGNMENT);
                        buffer.planes.push(Plane {
                            buf: buffer.buf.split_to(planesize),
                            linesize: linesize,
                        });
                    }
                }
                buffer
            }
            _ => unimplemented!(),
        }
    }
}
<|endoftext|>


<commit_before>use super::book::{Book, BookAccounting};
use std::collections::HashMap;
use std::fmt;
use std::cmp::Ordering;
use ::error::PoloError;
use super::timeseries::Timeseries;
use super::tradestats::{TradeStats, TimeStats};
use time;

#[derive(Clone, Debug, PartialEq)]
pub struct Record {
  pub rate: f64,
  pub amount: f64
}

#[derive(Clone, Debug, PartialEq)]
pub struct BookStats {
  pub min_sell: f64,
  pub max_buy: f64,
  pub sum_sell: f64,
  pub sum_buy: f64,
  pub vec_buy: Vec<Record>,
  pub vec_sell: Vec<Record>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct BookWithStats {
  book: Book,
  pub stats: BookStats,
  pub trade_series_1s: Timeseries<TradeStats>,
  pub trade_stats_1m: TradeStats,
}

// BookStats operations

impl BookStats {
  pub fn new(book: &Book) -> BookStats {
    let (sum_buy, mut vec_buy) = hash_to_vec(&book.buy);
    vec_buy.sort_unstable_by(|rec1, rec2| f64cmp(&rec2.rate, &rec1.rate));
    let (sum_sell, mut vec_sell) = hash_to_vec(&book.sell);
    vec_sell.sort_unstable_by(|rec1, rec2| f64cmp(&rec1.rate, &rec2.rate));

    BookStats {
      min_sell: vec_sell.first().map_or(0.0, |rec| rec.rate),
      max_buy: vec_buy.first().map_or(0.0, |rec| rec.rate),
      sum_sell,
      sum_buy,
      vec_buy,
      vec_sell,
    }
  }

  pub fn update_sell(&mut self, rate: f64, amount: f64, prev_amount: Option<f64>) {
    let idx_r = self.vec_sell.binary_search_by(|rec| f64cmp(&rec.rate, &rate));
    let stat_cmp = self.min_sell > rate;
    update_sorted_vec(idx_r, &mut self.vec_sell, &mut self.min_sell, rate, amount, stat_cmp);
    self.sum_sell = self.sum_sell + amount - prev_amount.unwrap_or(0.0);
  }

  pub fn update_buy(&mut self, rate: f64, amount: f64, prev_amount: Option<f64>) {
    let idx_r = self.vec_buy.binary_search_by(|rec| f64cmp(&rate, &rec.rate));
    let stat_cmp = self.max_buy < rate;
    update_sorted_vec(idx_r, &mut self.vec_buy, &mut self.max_buy, rate, amount, stat_cmp);
    self.sum_buy = self.sum_buy + amount - prev_amount.unwrap_or(0.0);
  }
}

impl fmt::Display for BookStats {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "(BUY min {} sum {} | SELL max {} sum {})", self.max_buy, self.sum_buy, self.min_sell, self.sum_sell)
  }
}

// BookWithStats operations

impl BookWithStats {
  pub fn new(book: Book) -> BookWithStats {
    BookWithStats {
      stats: BookStats::new(&book),
      trade_series_1s: Timeseries::new(),
      trade_stats_1m: TradeStats::default(),
      book
    }
  }
}

impl fmt::Display for BookWithStats {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "  ** STATS ")?;
    self.stats.fmt(f)?;
    if let Some(last_sec) = self.trade_series_1s.data.front() {
      write!(f, "\n  >> T1s: ")?;
      last_sec.fmt(f)?;
    }
    write!(f, "\n  >> T1m: ")?;
    self.trade_stats_1m.fmt(f)
  }
}

impl BookAccounting for BookWithStats {
  fn update_sell(&mut self, rate: String, amount: f64) -> Option<f64> {
    let rate_f64 = rate.parse::<f64>();
    let prev_amount = self.book.update_sell(rate, amount);
    if let Ok(rate_f64) = rate_f64 {
      self.stats.update_sell(rate_f64, amount, prev_amount);
    };
    prev_amount
  }

  fn update_buy(&mut self, rate: String, amount: f64) -> Option<f64> {
    let rate_f64 = rate.parse::<f64>();
    let prev_amount = self.book.update_buy(rate, amount);
    if let Ok(rate_f64) = rate_f64 {
      self.stats.update_buy(rate_f64, amount, prev_amount);
    };
    prev_amount
  }

  fn new_deal(&mut self, id: u64, rate: String, amount: f64) -> Result<f64, PoloError> {
    let rate_f64 = self.book.new_deal(id, rate, amount)?;
    Ok(rate_f64)
  }

  fn book_ref(&self) -> &Book {
    &self.book
  }
}

impl TimeStats for BookWithStats {
  fn update_stats_1s(&mut self) {
    let timestamp = time::get_time();
    // cleanup deals time series
    self.book.deals.drain_until(time::Timespec { sec: timestamp.sec - 600, nsec: 0 });
    // update stats with last second of data
    let last_second = self.book.deals.vec_after(time::Timespec { sec: timestamp.sec - 1, nsec: timestamp.nsec });
    let stats = TradeStats::new(last_second);
    self.trade_stats_1m = self.trade_stats_1m + &stats;
    self.trade_series_1s.add(stats);
    if self.trade_series_1s.data.len() > 60 {
      let stats_1m_ago = &self.trade_series_1s.data[self.trade_series_1s.data.len() - 60];
      self.trade_stats_1m = self.trade_stats_1m - stats_1m_ago;
    }
  }
}

/**
 ** Library functions
 **/

fn f64cmp(f1: &f64, f2: &f64) -> Ordering {
  f1.partial_cmp(f2).unwrap_or(Ordering::Equal)
}

fn hash_to_vec(hash: &HashMap<String, f64>) -> (f64, Vec<Record>) {
  let mut sum = 0.0;
  let vec = hash.iter().filter_map(|(rate_s, amount)| {
    rate_s.parse::<f64>()
      .and_then(|rate| {
        sum += *amount;
        Ok(Record { rate, amount: *amount })
      }).ok()
  }).collect();
  (sum, vec)
}

fn update_sorted_vec(idx_r: Result<usize, usize>, vec: &mut Vec<Record>, stat: &mut f64, rate: f64, amount: f64, stat_cmp: bool) {
  if amount == 0.0 {
    if let Ok(idx) = idx_r {
      vec.remove(idx);
    };
    if *stat == rate || *stat == 0.0 {
      *stat = vec.first().map_or(*stat, |rec| rec.rate);
    }
  } else if amount > 0.0  {
    match idx_r {
      Ok(idx) => vec[idx].amount = amount,
      Err(idx) => vec.insert(idx, Record { rate, amount })
    };
    if stat_cmp || *stat == 0.0 {
      *stat = rate
    }
  };
}

/**
 ** TESTS TESTS TESTS
 **/

#[cfg(test)]
mod tests {
  use super::BookWithStats;
  use ::data::book::Book;
  use json;
  use std::convert::TryFrom;

  #[test]
  fn stats_init() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let book_stats = BookWithStats::new(book).stats;
    assert_eq!(book_stats.min_sell, 0.13161901);
    assert_eq!(book_stats.max_buy, 0.13109621);
  }

  #[test]
  fn stats_init_wrong_order() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13361901": 0.23709568, "0.13164313": "0.17328089"}, {"0.12909621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let book_stats = BookWithStats::new(book).stats;
    assert_eq!(book_stats.min_sell, 0.13164313);
    assert_eq!(book_stats.max_buy, 0.13069621);
  }

  #[test]
  fn stats_update_sell_zero() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let mut book_stats = BookWithStats::new(book).stats;
    book_stats.update_sell(0.1, 0.0, None);
    assert_eq!(book_stats.min_sell, 0.13161901);
  }

  #[test]
  fn stats_update_sell_shift_zero() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let mut book_stats = BookWithStats::new(book).stats;
    book_stats.update_sell(0.13161901, 0.0, Some(0.23709568));
    assert_eq!(book_stats.min_sell, 0.13164313);
  }

  #[test]
  fn stats_update_sell() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let mut book_stats = BookWithStats::new(book).stats;
    book_stats.update_sell(0.1, 1.0, None);
    assert_eq!(book_stats.min_sell, 0.1);
  }

  #[test]
  fn stats_update_buy_zero() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let mut book_stats = BookWithStats::new(book).stats;
    book_stats.update_buy(100.0, 0.0, None);
    assert_eq!(book_stats.max_buy, 0.13109621);
  }

  #[test]
  fn stats_update_buy_shift_zero() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let mut book_stats = BookWithStats::new(book).stats;
    book_stats.update_buy(0.13109621, 0.0, Some(0.2331));
    assert_eq!(book_stats.max_buy, 0.13069621);
  }

  #[test]
  fn stats_update_buy() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let mut book_stats = BookWithStats::new(book).stats;
    book_stats.update_buy(100.0, 1.0, None);
    assert_eq!(book_stats.max_buy, 100.0);
  }

}<commit_msg>add debug print<commit_after>use super::book::{Book, BookAccounting};
use std::collections::HashMap;
use std::fmt;
use std::cmp::Ordering;
use ::error::PoloError;
use super::timeseries::Timeseries;
use super::tradestats::{TradeStats, TimeStats};
use time;

#[derive(Clone, Debug, PartialEq)]
pub struct Record {
  pub rate: f64,
  pub amount: f64
}

#[derive(Clone, Debug, PartialEq)]
pub struct BookStats {
  pub min_sell: f64,
  pub max_buy: f64,
  pub sum_sell: f64,
  pub sum_buy: f64,
  pub vec_buy: Vec<Record>,
  pub vec_sell: Vec<Record>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct BookWithStats {
  book: Book,
  pub stats: BookStats,
  pub trade_series_1s: Timeseries<TradeStats>,
  pub trade_stats_1m: TradeStats,
}

// BookStats operations

impl BookStats {
  pub fn new(book: &Book) -> BookStats {
    let (sum_buy, mut vec_buy) = hash_to_vec(&book.buy);
    vec_buy.sort_unstable_by(|rec1, rec2| f64cmp(&rec2.rate, &rec1.rate));
    let (sum_sell, mut vec_sell) = hash_to_vec(&book.sell);
    vec_sell.sort_unstable_by(|rec1, rec2| f64cmp(&rec1.rate, &rec2.rate));

    BookStats {
      min_sell: vec_sell.first().map_or(0.0, |rec| rec.rate),
      max_buy: vec_buy.first().map_or(0.0, |rec| rec.rate),
      sum_sell,
      sum_buy,
      vec_buy,
      vec_sell,
    }
  }

  pub fn update_sell(&mut self, rate: f64, amount: f64, prev_amount: Option<f64>) {
    let idx_r = self.vec_sell.binary_search_by(|rec| f64cmp(&rec.rate, &rate));
    let stat_cmp = self.min_sell > rate;
    update_sorted_vec(idx_r, &mut self.vec_sell, &mut self.min_sell, rate, amount, stat_cmp);
    self.sum_sell = self.sum_sell + amount - prev_amount.unwrap_or(0.0);
  }

  pub fn update_buy(&mut self, rate: f64, amount: f64, prev_amount: Option<f64>) {
    let idx_r = self.vec_buy.binary_search_by(|rec| f64cmp(&rate, &rec.rate));
    let stat_cmp = self.max_buy < rate;
    update_sorted_vec(idx_r, &mut self.vec_buy, &mut self.max_buy, rate, amount, stat_cmp);
    self.sum_buy = self.sum_buy + amount - prev_amount.unwrap_or(0.0);
  }
}

impl fmt::Display for BookStats {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "(BUY min {} sum {} | SELL max {} sum {})", self.max_buy, self.sum_buy, self.min_sell, self.sum_sell)
  }
}

// BookWithStats operations

impl BookWithStats {
  pub fn new(book: Book) -> BookWithStats {
    BookWithStats {
      stats: BookStats::new(&book),
      trade_series_1s: Timeseries::new(),
      trade_stats_1m: TradeStats::default(),
      book
    }
  }
}

impl fmt::Display for BookWithStats {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "  ** STATS ")?;
    self.stats.fmt(f)?;
    if let Some(last_sec) = self.trade_series_1s.data.front() {
      write!(f, "\n  >> T1s: ")?;
      last_sec.fmt(f)?;
    }
    write!(f, "\n  >> T1m: ")?;
    self.trade_stats_1m.fmt(f)
  }
}

impl BookAccounting for BookWithStats {
  fn update_sell(&mut self, rate: String, amount: f64) -> Option<f64> {
    let rate_f64 = rate.parse::<f64>();
    let prev_amount = self.book.update_sell(rate, amount);
    if let Ok(rate_f64) = rate_f64 {
      self.stats.update_sell(rate_f64, amount, prev_amount);
    };
    prev_amount
  }

  fn update_buy(&mut self, rate: String, amount: f64) -> Option<f64> {
    let rate_f64 = rate.parse::<f64>();
    let prev_amount = self.book.update_buy(rate, amount);
    if let Ok(rate_f64) = rate_f64 {
      self.stats.update_buy(rate_f64, amount, prev_amount);
    };
    prev_amount
  }

  fn new_deal(&mut self, id: u64, rate: String, amount: f64) -> Result<f64, PoloError> {
    let rate_f64 = self.book.new_deal(id, rate, amount)?;
    Ok(rate_f64)
  }

  fn book_ref(&self) -> &Book {
    &self.book
  }
}

impl TimeStats for BookWithStats {
  fn update_stats_1s(&mut self) {
    let timestamp = time::get_time();
    // cleanup deals time series
    self.book.deals.drain_until(time::Timespec { sec: timestamp.sec - 600, nsec: 0 });
    // update stats with last second of data
    let last_second = self.book.deals.vec_after(time::Timespec { sec: timestamp.sec - 1, nsec: timestamp.nsec });
    let stats = TradeStats::new(last_second);
    self.trade_stats_1m = self.trade_stats_1m + &stats;
    self.trade_series_1s.add(stats);
    if self.trade_series_1s.data.len() > 60 {
      let stats_1m_ago = &self.trade_series_1s.data[self.trade_series_1s.data.len() - 60];
      println!("stats_1m_ago {:?}", stats_1m_ago);
      self.trade_stats_1m = self.trade_stats_1m - stats_1m_ago;
    }
  }
}

/**
 ** Library functions
 **/

fn f64cmp(f1: &f64, f2: &f64) -> Ordering {
  f1.partial_cmp(f2).unwrap_or(Ordering::Equal)
}

fn hash_to_vec(hash: &HashMap<String, f64>) -> (f64, Vec<Record>) {
  let mut sum = 0.0;
  let vec = hash.iter().filter_map(|(rate_s, amount)| {
    rate_s.parse::<f64>()
      .and_then(|rate| {
        sum += *amount;
        Ok(Record { rate, amount: *amount })
      }).ok()
  }).collect();
  (sum, vec)
}

fn update_sorted_vec(idx_r: Result<usize, usize>, vec: &mut Vec<Record>, stat: &mut f64, rate: f64, amount: f64, stat_cmp: bool) {
  if amount == 0.0 {
    if let Ok(idx) = idx_r {
      vec.remove(idx);
    };
    if *stat == rate || *stat == 0.0 {
      *stat = vec.first().map_or(*stat, |rec| rec.rate);
    }
  } else if amount > 0.0  {
    match idx_r {
      Ok(idx) => vec[idx].amount = amount,
      Err(idx) => vec.insert(idx, Record { rate, amount })
    };
    if stat_cmp || *stat == 0.0 {
      *stat = rate
    }
  };
}

/**
 ** TESTS TESTS TESTS
 **/

#[cfg(test)]
mod tests {
  use super::BookWithStats;
  use ::data::book::Book;
  use json;
  use std::convert::TryFrom;

  #[test]
  fn stats_init() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let book_stats = BookWithStats::new(book).stats;
    assert_eq!(book_stats.min_sell, 0.13161901);
    assert_eq!(book_stats.max_buy, 0.13109621);
  }

  #[test]
  fn stats_init_wrong_order() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13361901": 0.23709568, "0.13164313": "0.17328089"}, {"0.12909621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let book_stats = BookWithStats::new(book).stats;
    assert_eq!(book_stats.min_sell, 0.13164313);
    assert_eq!(book_stats.max_buy, 0.13069621);
  }

  #[test]
  fn stats_update_sell_zero() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let mut book_stats = BookWithStats::new(book).stats;
    book_stats.update_sell(0.1, 0.0, None);
    assert_eq!(book_stats.min_sell, 0.13161901);
  }

  #[test]
  fn stats_update_sell_shift_zero() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let mut book_stats = BookWithStats::new(book).stats;
    book_stats.update_sell(0.13161901, 0.0, Some(0.23709568));
    assert_eq!(book_stats.min_sell, 0.13164313);
  }

  #[test]
  fn stats_update_sell() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let mut book_stats = BookWithStats::new(book).stats;
    book_stats.update_sell(0.1, 1.0, None);
    assert_eq!(book_stats.min_sell, 0.1);
  }

  #[test]
  fn stats_update_buy_zero() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let mut book_stats = BookWithStats::new(book).stats;
    book_stats.update_buy(100.0, 0.0, None);
    assert_eq!(book_stats.max_buy, 0.13109621);
  }

  #[test]
  fn stats_update_buy_shift_zero() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let mut book_stats = BookWithStats::new(book).stats;
    book_stats.update_buy(0.13109621, 0.0, Some(0.2331));
    assert_eq!(book_stats.max_buy, 0.13069621);
  }

  #[test]
  fn stats_update_buy() {
    let book_init = r#"{"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13109621": 0.2331, "0.13069621": 0.2331}]}"#;
    let book = Book::try_from(&json::parse(book_init).unwrap()).unwrap();
    let mut book_stats = BookWithStats::new(book).stats;
    book_stats.update_buy(100.0, 1.0, None);
    assert_eq!(book_stats.max_buy, 100.0);
  }

}<|endoftext|>


<commit_before>use sized_enum::SizedEnum;
use enum_primitive::FromPrimitive;
use map::RoomId;
use game_state::{GameState, GameMode};

impl GameState {
  fn print_help(&self) {
    println!(
      "Debug mode commands:\n\n\

       quit            - exit debug mode\n\
       goto <room id>  - teleport to a room\n\
       rooms           - list rooms + their contents\n"
    );
  }

  fn goto_room(&mut self, args: Vec<&str>) {
    if let Ok(number) = args[0].parse::<usize>() {
      if let Some(room_id) = RoomId::from_usize(number) {
        self.curr_room = room_id;
        println!("Teleported to {:?}.", room_id);
        return;
      }
    }
    println!("Unknown room id.");
  }

  fn list_rooms(&self) {
    for i in 0..RoomId::size() {
      RoomId::from_usize(i).map(|room_id| {
        let room = self.map.room(room_id);
        println!("Room {} - {:?}", i, room_id);
        if let Some(contents) = room.contents {
          println!("  Contains {:?}", contents);
        }
      });
    }
  }

  pub fn tick_debug_mode(&mut self) {
    self.ask("debug> ", |state, input| {
      if input == "q" || input == "quit" {
        state.set_mode(GameMode::Primary);
      } else if input == "h" || input == "?" || input == "help" {
        state.print_help();
      } else if input == "rooms" {
        state.list_rooms();
      } else if input.starts_with("goto ") {
        state.goto_room(input.split_whitespace().skip(1).collect());
      } else if input.len() > 0 {
        println!("Unrecognized command. Type ? for help.");
      }
    });
  }
}
<commit_msg>Add 'version' debug command<commit_after>use sized_enum::SizedEnum;
use enum_primitive::FromPrimitive;
use map::RoomId;
use game_state::{GameState, GameMode};

impl GameState {
  fn print_help(&self) {
    println!(
      "Debug mode commands:\n\n\

       quit            - exit debug mode\n\
       goto <room id>  - teleport to a room\n\
       rooms           - list rooms + their contents\n\
       version         - show version number\n"
    );
  }

  fn goto_room(&mut self, args: Vec<&str>) {
    if let Ok(number) = args[0].parse::<usize>() {
      if let Some(room_id) = RoomId::from_usize(number) {
        self.curr_room = room_id;
        println!("Teleported to {:?}.", room_id);
        return;
      }
    }
    println!("Unknown room id.");
  }

  fn list_rooms(&self) {
    for i in 0..RoomId::size() {
      RoomId::from_usize(i).map(|room_id| {
        let room = self.map.room(room_id);
        println!("Room {} - {:?}", i, room_id);
        if let Some(contents) = room.contents {
          println!("  Contains {:?}", contents);
        }
      });
    }
  }

  pub fn tick_debug_mode(&mut self) {
    self.ask("debug> ", |state, input| {
      if input == "q" || input == "quit" {
        state.set_mode(GameMode::Primary);
      } else if input == "h" || input == "?" || input == "help" {
        state.print_help();
      } else if input == "rooms" {
        state.list_rooms();
      } else if input.starts_with("goto ") {
        state.goto_room(input.split_whitespace().skip(1).collect());
      } else if input == "version" {
        println!("{} {}", env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"));
      } else if input.len() > 0 {
        println!("Unrecognized command. Type ? for help.");
      }
    });
  }
}
<|endoftext|>


<commit_before>fn main() {

	struct GoalExplore {
		name: ~str
	}

	struct GoalRest {
		name: ~str
	}

	trait Goal {
		fn new() -> Self;
		fn process(&self) { }
	}

	impl Goal for GoalExplore {
		fn new() -> GoalExplore {
			GoalExplore { name: ~"GoalExplore" }
		}

		fn process(&self) {
			println(fmt!("Exploring"));
		}
	}

	impl Goal for GoalRest {
		fn new() -> GoalRest {
			GoalRest { name: ~"GoalRest"}
		}

		fn process(&self) {
			println(fmt!("Resting"))
		}
	}

	enum Tree<T> {
		Leaf(T),
		Node(~[Tree<T>]),
	}

	//fn foo<G: Goal>(goal: &G) { .. }

	impl<T> Tree<T> {
		fn push(&mut self, node: Tree<T>) {
			match *self {
				Node(ref mut children) => children.push(node),
				Leaf(_) => fail!(~"Cannot add node to leaf node")
			}
		}

		/*fn to_str(&mut self) {
			match *self {
				Node(ref mut children) => for x in children.iter() { x.to_str() },
				Leaf(v) => ()//println!("{:?}", v)
			}
		}*/
	}

	let t: @GoalExplore = @Goal::new();
	t.process();
	let rest: @GoalRest = @Goal::new();
	rest.process();

	let mut tree: Tree<@Goal> = Node(~[]);

	tree.push(Leaf(t as @Goal));

	println!("{:?}", tree);
}
/*
enum Tree<T> {
	Leaf(T),
	Node { children: ~[Tree<T>] }
}

let mut my_tree: Tree<uint> = Node { children: ~[] };

match my_tree {
	Node { children: ref mut children } => children.push(Leaf(1u)),
	_ => ()
}

my_tree  */<commit_msg>Working data structure for composite pattern<commit_after>fn main() {

	struct GoalExplore {
		name: ~str
	}

	struct GoalRest {
		name: ~str
	}

	trait Goal {
		fn new() -> Self;
		fn process(&self) { }
		fn to_str(&self) -> ~str;
	}

	impl Goal for GoalExplore {
		fn new() -> GoalExplore {
			GoalExplore { name: ~"GoalExplore" }
		}

		fn process(&self) {
			println(fmt!("Exploring"));
		}

		fn to_str(&self) -> ~str {
			self.name.clone()
		}
	}

	impl Goal for GoalRest {
		fn new() -> GoalRest {
			GoalRest { name: ~"GoalRest"}
		}

		fn process(&self) {
			println(fmt!("Resting"))
		}

		fn to_str(&self) -> ~str {
			self.name.clone()
		}
	}

	enum Tree {
		Leaf(@Goal),
		Node(~[Tree]),
	}

	//fn foo<G: Goal>(goal: &G) { .. }

	impl Tree {
		fn push(&mut self, node: Tree) {
			match *self {
				Node(ref mut children) => children.push(node),
				Leaf(_) => fail!(~"Cannot add node to leaf node")
			}
		}

		fn process(&mut self) {
			match *self {
				Node(ref mut children) => for x in children.mut_iter() { x.process() },
				Leaf(l) => println(fmt!("Leaf: %s", l.to_str()))
			}
		}
	}

	let mut t: @mut GoalExplore = @mut Goal::new();
	t.process();
	let mut rest: @mut GoalRest = @mut Goal::new();
	rest.process();

	let mut tree: Tree = Node(~[]);

	tree.push(Leaf(t as @Goal));

	println!("{:?}", tree);
}<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::fs;
use std::path::{Path, PathBuf};

use glob::Pattern;
use shellexpand;
use walkdir::{DirEntry, WalkDir, WalkDirIterator};

use app::{Cache, Config, InitialStr};
use vcs;
use util;

use super::repository::Repository;

// inner representation of cache format.
#[derive(Default, Serialize, Deserialize)]
struct CacheData {
  repositories: Vec<Repository>,
}

#[derive(Serialize, Deserialize)]
struct ConfigData {
  root: Option<String>,
  includes: Option<Vec<String>>,
  excludes: Option<Vec<String>>,
}

impl ConfigData {
  pub fn includes(&self) -> &[String] {
    self.includes
        .as_ref()
        .map(Vec::as_slice)
        .unwrap_or(&[])
  }

  #[allow(dead_code)]
  pub fn excludes(&self) -> &[String] {
    self.excludes
        .as_ref()
        .map(Vec::as_slice)
        .unwrap_or(&[])
  }
}

impl InitialStr for ConfigData {
  #[inline]
  fn initial_str() -> &'static str {
    include_str!("config.toml")
  }
}

pub struct Workspace<'a> {
  cache: Cache<CacheData>,
  config: Config<ConfigData>,
  root: Option<&'a Path>,
}

impl<'a> Workspace<'a> {
  pub fn new(root: Option<&'a Path>) -> ::Result<Workspace<'a>> {
    let cache = Cache::load()?;
    let config = Config::load()?;
    Ok(Workspace {
         cache: cache,
         config: config,
         root: root,
       })
  }

  /// Returns root directory of the workspace.
  pub fn root_dir(&self) -> Option<Cow<Path>> {
    self.root
        .map(Into::into)
        .or_else(|| {
                   self.config
                       .get()
                       .root
                       .as_ref()
                       .and_then(|root| util::make_path_buf(root).ok())
                       .map(Into::into)
                 })
  }

  pub fn base_dirs(&self) -> Vec<PathBuf> {
    self.config
        .get()
        .includes()
        .into_iter()
        .filter_map(|root| util::make_path_buf(&root).ok())
        .collect()
  }

  pub fn exclude_patterns(&self) -> Vec<Pattern> {
    self.config
        .get()
        .excludes()
        .into_iter()
        .filter_map(|ex| {
                      shellexpand::full(&ex)
                        .ok()
                        .map(|ex| ex.replace(r"\", "/"))
                        .and_then(|ex| Pattern::new(&ex).ok())
                    })
        .collect()
  }

  pub fn add_repository(&mut self, repo: Repository) {
    let ref mut repos = self.cache.get_mut().repositories;
    if let Some(mut r) = repos.iter_mut().find(|r| r.is_same_local(&repo)) {
      *r = repo;
      return;
    }
    repos.push(repo);
  }

  /// Returns a list of managed repositories.
  /// Note that this method returns None if cache has not created yet.
  pub fn repositories(&self) -> Option<&[Repository]> {
    self.cache
        .get_opt()
        .map(|cache| cache.repositories.as_slice())
  }

  /// Scan repositories and update state.
  pub fn scan_repositories(&mut self, verbose: bool, prune: bool, depth: Option<usize>) -> ::Result<()> {
    let mut repos = Vec::new();
    for repo in self.collect_base_dirs(depth) {
      if verbose {
        println!("Found at {}", repo.path_string());
      }
      repos.push(repo);
    }

    let outside_repos = self.collect_outsides();
    for repo in outside_repos {
      if prune {
        println!("Dropped: {}", repo.path_string());
      } else {
        if verbose {
          println!("Found at outside: {}", repo.path_string());
        }
        repos.push(repo);
      }
    }

    self.cache.get_mut().repositories = repos;
    Ok(())
  }

  /// Save current state of workspace to cache file.
  pub fn save_cache(&self) -> ::Result<()> {
    self.cache.dump()?;
    Ok(())
  }

  /// Collect repositories located at inside of base directories
  fn collect_base_dirs(&self, depth: Option<usize>) -> Vec<Repository> {
    self.base_dirs()
        .into_iter()
        .flat_map(|root| collect_repositories_from(root, depth, self.exclude_patterns()))
        .filter_map(|path| Repository::from_path(path).ok())
        .collect()
  }

  /// Collect managed repositories located at outside of base directories
  fn collect_outsides(&self) -> Vec<Repository> {
    let cache = match self.cache.get_opt() {
      Some(cache) => cache,
      None => return Vec::new(),
    };

    let mut repos = Vec::with_capacity(cache.repositories.len());
    for repo in cache.repositories.clone() {
      let under_management = self.base_dirs()
                                 .into_iter()
                                 .any(|root| repo.is_contained(root));
      if !under_management && repo.is_vcs() {
        repos.push(repo);
      }
    }
    repos
  }
}

fn collect_repositories_from<P>(root: P, depth: Option<usize>, excludes: Vec<Pattern>) -> Vec<PathBuf>
  where P: AsRef<Path>
{
  let filter = |entry: &DirEntry| {
    entry.path() == root.as_ref() ||
    entry.path()
         .parent()
         .map(|path| vcs::detect_from_path(&path).is_none())
         .unwrap_or(true)
  };

  let mut walkdir = WalkDir::new(root.as_ref()).follow_links(true);
  if let Some(depth) = depth {
    walkdir = walkdir.max_depth(depth);
  }
  walkdir.into_iter()
         .filter_entry(filter)
         .filter_map(|e| e.ok())
         .filter_map(|e| fs::canonicalize(e.path()).ok())
         .filter(move |path| {
                   let path = path.to_str().unwrap().trim_left_matches(r"\\?\");
                   excludes.iter().all(|ex| !ex.matches(path))
                 })
         .filter(|ref path| vcs::detect_from_path(path).is_some())
         .collect()
}
<commit_msg>exclude pattern at filter_entry()<commit_after>use std::borrow::Cow;
use std::fs;
use std::path::{Path, PathBuf};

use glob::Pattern;
use shellexpand;
use walkdir::{DirEntry, WalkDir, WalkDirIterator};

use app::{Cache, Config, InitialStr};
use vcs;
use util;

use super::repository::Repository;

// inner representation of cache format.
#[derive(Default, Serialize, Deserialize)]
struct CacheData {
  repositories: Vec<Repository>,
}

#[derive(Serialize, Deserialize)]
struct ConfigData {
  root: Option<String>,
  includes: Option<Vec<String>>,
  excludes: Option<Vec<String>>,
}

impl ConfigData {
  pub fn includes(&self) -> &[String] {
    self.includes
        .as_ref()
        .map(Vec::as_slice)
        .unwrap_or(&[])
  }

  #[allow(dead_code)]
  pub fn excludes(&self) -> &[String] {
    self.excludes
        .as_ref()
        .map(Vec::as_slice)
        .unwrap_or(&[])
  }
}

impl InitialStr for ConfigData {
  #[inline]
  fn initial_str() -> &'static str {
    include_str!("config.toml")
  }
}

pub struct Workspace<'a> {
  cache: Cache<CacheData>,
  config: Config<ConfigData>,
  root: Option<&'a Path>,
}

impl<'a> Workspace<'a> {
  pub fn new(root: Option<&'a Path>) -> ::Result<Workspace<'a>> {
    let cache = Cache::load()?;
    let config = Config::load()?;
    Ok(Workspace {
         cache: cache,
         config: config,
         root: root,
       })
  }

  /// Returns root directory of the workspace.
  pub fn root_dir(&self) -> Option<Cow<Path>> {
    self.root
        .map(Into::into)
        .or_else(|| {
                   self.config
                       .get()
                       .root
                       .as_ref()
                       .and_then(|root| util::make_path_buf(root).ok())
                       .map(Into::into)
                 })
  }

  pub fn base_dirs(&self) -> Vec<PathBuf> {
    self.config
        .get()
        .includes()
        .into_iter()
        .filter_map(|root| util::make_path_buf(&root).ok())
        .collect()
  }

  pub fn exclude_patterns(&self) -> Vec<Pattern> {
    self.config
        .get()
        .excludes()
        .into_iter()
        .filter_map(|ex| {
                      shellexpand::full(&ex)
                        .ok()
                        .map(|ex| ex.replace(r"\", "/"))
                        .and_then(|ex| Pattern::new(&ex).ok())
                    })
        .collect()
  }

  pub fn add_repository(&mut self, repo: Repository) {
    let ref mut repos = self.cache.get_mut().repositories;
    if let Some(mut r) = repos.iter_mut().find(|r| r.is_same_local(&repo)) {
      *r = repo;
      return;
    }
    repos.push(repo);
  }

  /// Returns a list of managed repositories.
  /// Note that this method returns None if cache has not created yet.
  pub fn repositories(&self) -> Option<&[Repository]> {
    self.cache
        .get_opt()
        .map(|cache| cache.repositories.as_slice())
  }

  /// Scan repositories and update state.
  pub fn scan_repositories(&mut self, verbose: bool, prune: bool, depth: Option<usize>) -> ::Result<()> {
    let mut repos = Vec::new();
    for repo in self.collect_base_dirs(depth) {
      if verbose {
        println!("Found at {}", repo.path_string());
      }
      repos.push(repo);
    }

    let outside_repos = self.collect_outsides();
    for repo in outside_repos {
      if prune {
        println!("Dropped: {}", repo.path_string());
      } else {
        if verbose {
          println!("Found at outside: {}", repo.path_string());
        }
        repos.push(repo);
      }
    }

    self.cache.get_mut().repositories = repos;
    Ok(())
  }

  /// Save current state of workspace to cache file.
  pub fn save_cache(&self) -> ::Result<()> {
    self.cache.dump()?;
    Ok(())
  }

  /// Collect repositories located at inside of base directories
  fn collect_base_dirs(&self, depth: Option<usize>) -> Vec<Repository> {
    self.base_dirs()
        .into_iter()
        .flat_map(|root| collect_repositories_from(root, depth, self.exclude_patterns()))
        .filter_map(|path| Repository::from_path(path).ok())
        .collect()
  }

  /// Collect managed repositories located at outside of base directories
  fn collect_outsides(&self) -> Vec<Repository> {
    let cache = match self.cache.get_opt() {
      Some(cache) => cache,
      None => return Vec::new(),
    };

    let mut repos = Vec::with_capacity(cache.repositories.len());
    for repo in cache.repositories.clone() {
      let under_management = self.base_dirs()
                                 .into_iter()
                                 .any(|root| repo.is_contained(root));
      if !under_management && repo.is_vcs() {
        repos.push(repo);
      }
    }
    repos
  }
}

fn collect_repositories_from<P>(root: P, depth: Option<usize>, excludes: Vec<Pattern>) -> Vec<PathBuf>
  where P: AsRef<Path>
{
  let filter = {
    let root = root.as_ref();
    move |entry: &DirEntry| {
      if entry.path() == root {
        return true;
      }
      !entry.path()
            .parent()
            .map(|path| vcs::detect_from_path(&path).is_some())
            .unwrap_or(false) &&
      entry.path()
           .canonicalize()
           .ok()
           .map(|path| {
                  let path = path.to_str().unwrap().trim_left_matches(r"\\?\");
                  excludes.iter().all(|ex| !ex.matches(path))
                })
           .unwrap_or(false)
    }
  };

  let mut walkdir = WalkDir::new(root.as_ref()).follow_links(true);
  if let Some(depth) = depth {
    walkdir = walkdir.max_depth(depth);
  }

  walkdir.into_iter()
         .filter_entry(filter)
         .filter_map(|entry| {
                       entry.ok()
                            .and_then(|entry| fs::canonicalize(entry.path()).ok())
                     })
         .filter(|ref path| vcs::detect_from_path(path).is_some())
         .collect()
}
<|endoftext|>