<commit_before>// Copyright (c) 2016 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::Arc;
use std::path::Path;
use std::time::Duration;

use hab_core::util::sys;
use hyper::client::Client as HyperClient;
use hyper::client::RequestBuilder;
use hyper::client::pool::{Config, Pool};
use hyper::header::UserAgent;
use hyper::http::h1::Http11Protocol;
use hyper::net::{HttpsConnector, Openssl};
use openssl::ssl::{SslContext, SslMethod, SSL_OP_NO_SSLV2, SSL_OP_NO_SSLV3, SSL_OP_NO_COMPRESSION};
use url::Url;

use error::Result;
use net::ProxyHttpsConnector;
use proxy::{ProxyInfo, proxy_unless_domain_exempted};
use ssl;

header! { (ProxyAuthorization, "Proxy-Authorization") => [String] }

/// A generic wrapper around a Hyper HTTP client intended for API-like usage.
///
/// When an `ApiClient` is created, it has a constant URL base which is assumed to be some API
/// endpoint. This allows the underlying Hyper client to load and use any relevant HTTP proxy
/// support and to provide reasonable User-Agent HTTP headers, etc.
#[derive(Debug)]
pub struct ApiClient {
    /// The base URL for the client.
    endpoint: Url,
    /// An instance of a `hyper::Client` which is configured with an SSL context and optionally
    /// using an HTTP proxy.
    inner: HyperClient,
    /// Proxy information, if a proxy is being used.
    proxy: Option<ProxyInfo>,
    /// The URL scheme of the endpoint.
    target_scheme: String,
    /// The `User-Agent` header string to use for HTTP calls.
    user_agent_header: UserAgent,
}

impl ApiClient {
    /// Creates and returns a new `ApiClient` instance.
    ///
    /// # Errors
    ///
    /// * If the underlying Hyper client cannot be created
    /// * If a suitable SSL context cannot be established
    /// * If an HTTP proxy cannot be correctly setup
    /// * If a `User-Agent` HTTP header string cannot be constructed
    pub fn new(endpoint: &Url,
               product: &str,
               version: &str,
               fs_root_path: Option<&Path>)
               -> Result<Self> {
        Ok(ApiClient {
            endpoint: endpoint.clone(),
            inner: try!(new_hyper_client(Some(endpoint), fs_root_path)),
            proxy: try!(proxy_unless_domain_exempted(Some(endpoint))),
            target_scheme: endpoint.scheme().to_string(),
            user_agent_header: try!(user_agent(product, version)),
        })
    }

    /// Builds an HTTP GET request for a given path.
    pub fn get(&self, path: &str) -> RequestBuilder {
        self.get_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP GET request for a given path with the ability to customize the target URL.
    pub fn get_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("GET {} with {:?}", &url, &self);
        self.add_headers(self.inner.get(url))
    }

    /// Builds an HTTP HEAD request for a given path.
    pub fn head(&self, path: &str) -> RequestBuilder {
        self.head_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP HEAD request for a given path with the ability to customize the target URL.
    pub fn head_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("HEAD {} with {:?}", &url, &self);
        self.add_headers(self.inner.head(url))
    }

    /// Builds an HTTP PATCH request for a given path.
    pub fn patch(&self, path: &str) -> RequestBuilder {
        self.patch_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP PATCH request for a given path with the ability to customize the target URL.
    pub fn patch_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("PATH {} with {:?}", &url, &self);
        self.add_headers(self.inner.patch(url))
    }

    /// Builds an HTTP POST request for a given path.
    pub fn post(&self, path: &str) -> RequestBuilder {
        self.post_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP POST request for a given path with the ability to customize the target URL.
    pub fn post_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("POST {} with {:?}", &url, &self);
        self.add_headers(self.inner.post(url))
    }

    /// Builds an HTTP PUT request for a given path.
    pub fn put(&self, path: &str) -> RequestBuilder {
        self.put_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP PUT request for a given path with the ability to customize the target URL.
    pub fn put_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("PUT {} with {:?}", &url, &self);
        self.add_headers(self.inner.put(url))
    }

    /// Builds an HTTP DELETE request for a given path.
    pub fn delete(&self, path: &str) -> RequestBuilder {
        self.delete_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP DELETE request for a given path with the ability to customize the target URL.
    pub fn delete_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("DELETE {} with {:?}", &url, &self);
        self.add_headers(self.inner.delete(url))
    }

    fn add_headers<'a>(&'a self, rb: RequestBuilder<'a>) -> RequestBuilder {
        let mut rb = rb.header(self.user_agent_header.clone());
        // If the target URL is an `"http"` scheme and we're using a proxy server, then add the
        // proxy authorization header if appropriate. Note that for `"https"` targets, the proxy
        // server will be operating in TCP tuneling mode and will be authenticated on connection to
        // the proxy server which is why we should not add an additional header in this latter
        // case.
        if self.target_scheme == "http" {
            if let Some(ref info) = self.proxy {
                if let Some(header_value) = info.authorization_header_value() {
                    rb = rb.header(ProxyAuthorization(header_value));
                }
            }
        }
        rb
    }

    fn url_for(&self, path: &str) -> Url {
        let mut url = self.endpoint.clone();
        if path.is_empty() {
            return url;
        }
        if url.path().ends_with("/") {
            url.set_path(&format!("{}{}", self.endpoint.path(), path));
        } else {
            url.set_path(&format!("{}/{}", self.endpoint.path(), path));
        }
        url
    }
}

/// Builds a new hyper HTTP client with appropriate SSL configuration and HTTP/HTTPS proxy support.
///
/// ## Linux Platforms
///
/// We need a set of root certificates when connected to SSL/TLS web endpoints and this usually
/// boild down to using an all-in-one certificate file (such as a `cacert.pem` file) or a directory
/// of files which are certificates. The strategy to location or use a reasonable set of
/// certificates is as follows:
///
/// 1. If the `SSL_CERT_FILE` environment variable is set, then use its value for the certificates.
///    Interally this is triggering default OpenSSL behavior for this environment variable.
/// 2. If the `SSL_CERT_DIR` environment variable is set, then use its value for the directory
///    containing certificates. Like the `SSL_CERT_FILE` case above, this triggers default OpenSSL
///    behavior for this environment variable.
/// 3. If the `core/cacerts` Habitat package is installed locally, then use the latest release's
///    `cacert.pem` file.
/// 4. If none of the conditions above are met, then a `cacert.pem` will be written in an SSL cache
///    directory (by default `/hab/cache/ssl` for a root user and `$HOME/.hab/cache/ssl` for a
///    non-root user) and this will be used. The contents of this file will be inlined in this
///    crate at build time as a fallback insurance policy, meaning that if the a program using this
///    code is operating in a minimal environment which may not contain system certificates, it can
///    still operate. Once a `core/cacerts` Habitat package is present, the logic would fall back
///    to preferring the package version to the cached/inline file version.
///
/// ## Mac Platforms
///
/// The Mac platoform uses a Security Framework to store and find root certificates and the hyper
/// library will default to using this on the Mac. Therefore the behavior on the Mac remains
/// unchanged and will use the system's certificates.
///
fn new_hyper_client(for_domain: Option<&Url>, fs_root_path: Option<&Path>) -> Result<HyperClient> {
    let ctx = try!(ssl_ctx(fs_root_path));
    let ssl_client = Openssl { context: Arc::new(ctx) };
    match try!(proxy_unless_domain_exempted(for_domain)) {
        Some(proxy) => {
            debug!("Using proxy {}:{}...", proxy.host(), proxy.port());
            let connector = try!(ProxyHttpsConnector::new(proxy, ssl_client));
            let pool = Pool::with_connector(Config::default(), connector);
            let mut client = HyperClient::with_protocol(Http11Protocol::with_connector(pool));
            let timeout = Some(Duration::from_secs(5));
            client.set_read_timeout(timeout);
            client.set_write_timeout(timeout);
            Ok(client)
        }
        None => {
            let connector = HttpsConnector::new(ssl_client);
            let pool = Pool::with_connector(Config::default(), connector);
            let mut client = HyperClient::with_protocol(Http11Protocol::with_connector(pool));
            let timeout = Some(Duration::from_secs(5));
            client.set_read_timeout(timeout);
            client.set_write_timeout(timeout);
            Ok(client)
        }
    }
}

/// Returns an HTTP User-Agent string type for use by Hyper when making HTTP requests.
///
/// The general form for Habitat-related clients are of the following form:
///
/// ```text
/// <PRODUCT>/<VERSION> (<TARGET>; <KERNEL_RELEASE>)
/// ```
///
/// where:
///
/// * `<PRODUCT>`: is the provided product name
/// * `<VERSION>`: is the provided version string which may also include a release number
/// * `<TARGET>`: is the machine architecture and the kernel seperated by a dash in lower case
/// * `<KERNEL_RELEASE>`: is the kernel release string from `uname`
///
/// For example:
///
/// ```text
/// hab/0.6.0/20160606153031 (x86_64-darwin; 14.5.0)
/// ```
///
/// # Errors
///
/// * If system information cannot be obtained via `uname`
fn user_agent(product: &str, version: &str) -> Result<UserAgent> {
    let uname = try!(sys::uname());
    let ua = format!("{}/{} ({}-{}; {})",
                     product.trim(),
                     version.trim(),
                     uname.machine.trim().to_lowercase(),
                     uname.sys_name.trim().to_lowercase(),
                     uname.release.trim().to_lowercase());
    debug!("User-Agent: {}", &ua);
    Ok(UserAgent(ua))
}

fn ssl_ctx(fs_root_path: Option<&Path>) -> Result<SslContext> {
    // The spirit of this implementation is directly from Hyper's default OpensslClient function:
    // https://github.com/hyperium/hyper/blob/v0.9.5/src/net.rs#L653-L661
    let mut ctx = try!(SslContext::new(SslMethod::Sslv23));
    try!(ssl::set_ca(&mut ctx, fs_root_path));
    ctx.set_options(SSL_OP_NO_SSLV2 | SSL_OP_NO_SSLV3 | SSL_OP_NO_COMPRESSION);
    try!(ctx.set_cipher_list("ALL!EXPORT!EXPORT40!EXPORT56!aNULL!LOW!RC4@STRENGTH"));
    Ok(ctx)
}
<commit_msg>Pull out the timeout<commit_after>// Copyright (c) 2016 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::Arc;
use std::path::Path;
use std::time::Duration;

use hab_core::util::sys;
use hyper::client::Client as HyperClient;
use hyper::client::RequestBuilder;
use hyper::client::pool::{Config, Pool};
use hyper::header::UserAgent;
use hyper::http::h1::Http11Protocol;
use hyper::net::{HttpsConnector, Openssl};
use openssl::ssl::{SslContext, SslMethod, SSL_OP_NO_SSLV2, SSL_OP_NO_SSLV3, SSL_OP_NO_COMPRESSION};
use url::Url;

use error::Result;
use net::ProxyHttpsConnector;
use proxy::{ProxyInfo, proxy_unless_domain_exempted};
use ssl;

header! { (ProxyAuthorization, "Proxy-Authorization") => [String] }

/// A generic wrapper around a Hyper HTTP client intended for API-like usage.
///
/// When an `ApiClient` is created, it has a constant URL base which is assumed to be some API
/// endpoint. This allows the underlying Hyper client to load and use any relevant HTTP proxy
/// support and to provide reasonable User-Agent HTTP headers, etc.
#[derive(Debug)]
pub struct ApiClient {
    /// The base URL for the client.
    endpoint: Url,
    /// An instance of a `hyper::Client` which is configured with an SSL context and optionally
    /// using an HTTP proxy.
    inner: HyperClient,
    /// Proxy information, if a proxy is being used.
    proxy: Option<ProxyInfo>,
    /// The URL scheme of the endpoint.
    target_scheme: String,
    /// The `User-Agent` header string to use for HTTP calls.
    user_agent_header: UserAgent,
}

impl ApiClient {
    /// Creates and returns a new `ApiClient` instance.
    ///
    /// # Errors
    ///
    /// * If the underlying Hyper client cannot be created
    /// * If a suitable SSL context cannot be established
    /// * If an HTTP proxy cannot be correctly setup
    /// * If a `User-Agent` HTTP header string cannot be constructed
    pub fn new(endpoint: &Url,
               product: &str,
               version: &str,
               fs_root_path: Option<&Path>)
               -> Result<Self> {
        Ok(ApiClient {
            endpoint: endpoint.clone(),
            inner: try!(new_hyper_client(Some(endpoint), fs_root_path)),
            proxy: try!(proxy_unless_domain_exempted(Some(endpoint))),
            target_scheme: endpoint.scheme().to_string(),
            user_agent_header: try!(user_agent(product, version)),
        })
    }

    /// Builds an HTTP GET request for a given path.
    pub fn get(&self, path: &str) -> RequestBuilder {
        self.get_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP GET request for a given path with the ability to customize the target URL.
    pub fn get_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("GET {} with {:?}", &url, &self);
        self.add_headers(self.inner.get(url))
    }

    /// Builds an HTTP HEAD request for a given path.
    pub fn head(&self, path: &str) -> RequestBuilder {
        self.head_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP HEAD request for a given path with the ability to customize the target URL.
    pub fn head_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("HEAD {} with {:?}", &url, &self);
        self.add_headers(self.inner.head(url))
    }

    /// Builds an HTTP PATCH request for a given path.
    pub fn patch(&self, path: &str) -> RequestBuilder {
        self.patch_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP PATCH request for a given path with the ability to customize the target URL.
    pub fn patch_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("PATH {} with {:?}", &url, &self);
        self.add_headers(self.inner.patch(url))
    }

    /// Builds an HTTP POST request for a given path.
    pub fn post(&self, path: &str) -> RequestBuilder {
        self.post_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP POST request for a given path with the ability to customize the target URL.
    pub fn post_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("POST {} with {:?}", &url, &self);
        self.add_headers(self.inner.post(url))
    }

    /// Builds an HTTP PUT request for a given path.
    pub fn put(&self, path: &str) -> RequestBuilder {
        self.put_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP PUT request for a given path with the ability to customize the target URL.
    pub fn put_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("PUT {} with {:?}", &url, &self);
        self.add_headers(self.inner.put(url))
    }

    /// Builds an HTTP DELETE request for a given path.
    pub fn delete(&self, path: &str) -> RequestBuilder {
        self.delete_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP DELETE request for a given path with the ability to customize the target URL.
    pub fn delete_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("DELETE {} with {:?}", &url, &self);
        self.add_headers(self.inner.delete(url))
    }

    fn add_headers<'a>(&'a self, rb: RequestBuilder<'a>) -> RequestBuilder {
        let mut rb = rb.header(self.user_agent_header.clone());
        // If the target URL is an `"http"` scheme and we're using a proxy server, then add the
        // proxy authorization header if appropriate. Note that for `"https"` targets, the proxy
        // server will be operating in TCP tuneling mode and will be authenticated on connection to
        // the proxy server which is why we should not add an additional header in this latter
        // case.
        if self.target_scheme == "http" {
            if let Some(ref info) = self.proxy {
                if let Some(header_value) = info.authorization_header_value() {
                    rb = rb.header(ProxyAuthorization(header_value));
                }
            }
        }
        rb
    }

    fn url_for(&self, path: &str) -> Url {
        let mut url = self.endpoint.clone();
        if path.is_empty() {
            return url;
        }
        if url.path().ends_with("/") {
            url.set_path(&format!("{}{}", self.endpoint.path(), path));
        } else {
            url.set_path(&format!("{}/{}", self.endpoint.path(), path));
        }
        url
    }
}

/// Builds a new hyper HTTP client with appropriate SSL configuration and HTTP/HTTPS proxy support.
///
/// ## Linux Platforms
///
/// We need a set of root certificates when connected to SSL/TLS web endpoints and this usually
/// boild down to using an all-in-one certificate file (such as a `cacert.pem` file) or a directory
/// of files which are certificates. The strategy to location or use a reasonable set of
/// certificates is as follows:
///
/// 1. If the `SSL_CERT_FILE` environment variable is set, then use its value for the certificates.
///    Interally this is triggering default OpenSSL behavior for this environment variable.
/// 2. If the `SSL_CERT_DIR` environment variable is set, then use its value for the directory
///    containing certificates. Like the `SSL_CERT_FILE` case above, this triggers default OpenSSL
///    behavior for this environment variable.
/// 3. If the `core/cacerts` Habitat package is installed locally, then use the latest release's
///    `cacert.pem` file.
/// 4. If none of the conditions above are met, then a `cacert.pem` will be written in an SSL cache
///    directory (by default `/hab/cache/ssl` for a root user and `$HOME/.hab/cache/ssl` for a
///    non-root user) and this will be used. The contents of this file will be inlined in this
///    crate at build time as a fallback insurance policy, meaning that if the a program using this
///    code is operating in a minimal environment which may not contain system certificates, it can
///    still operate. Once a `core/cacerts` Habitat package is present, the logic would fall back
///    to preferring the package version to the cached/inline file version.
///
/// ## Mac Platforms
///
/// The Mac platoform uses a Security Framework to store and find root certificates and the hyper
/// library will default to using this on the Mac. Therefore the behavior on the Mac remains
/// unchanged and will use the system's certificates.
///
fn new_hyper_client(for_domain: Option<&Url>, fs_root_path: Option<&Path>) -> Result<HyperClient> {
    let ctx = try!(ssl_ctx(fs_root_path));
    let ssl_client = Openssl { context: Arc::new(ctx) };
    let timeout = Some(Duration::from_secs(5));

    match try!(proxy_unless_domain_exempted(for_domain)) {
        Some(proxy) => {
            debug!("Using proxy {}:{}...", proxy.host(), proxy.port());
            let connector = try!(ProxyHttpsConnector::new(proxy, ssl_client));
            let pool = Pool::with_connector(Config::default(), connector);
            let mut client = HyperClient::with_protocol(Http11Protocol::with_connector(pool));
            client.set_read_timeout(timeout);
            client.set_write_timeout(timeout);
            Ok(client)
        }
        None => {
            let connector = HttpsConnector::new(ssl_client);
            let pool = Pool::with_connector(Config::default(), connector);
            let mut client = HyperClient::with_protocol(Http11Protocol::with_connector(pool));
            client.set_read_timeout(timeout);
            client.set_write_timeout(timeout);
            Ok(client)
        }
    }
}

/// Returns an HTTP User-Agent string type for use by Hyper when making HTTP requests.
///
/// The general form for Habitat-related clients are of the following form:
///
/// ```text
/// <PRODUCT>/<VERSION> (<TARGET>; <KERNEL_RELEASE>)
/// ```
///
/// where:
///
/// * `<PRODUCT>`: is the provided product name
/// * `<VERSION>`: is the provided version string which may also include a release number
/// * `<TARGET>`: is the machine architecture and the kernel seperated by a dash in lower case
/// * `<KERNEL_RELEASE>`: is the kernel release string from `uname`
///
/// For example:
///
/// ```text
/// hab/0.6.0/20160606153031 (x86_64-darwin; 14.5.0)
/// ```
///
/// # Errors
///
/// * If system information cannot be obtained via `uname`
fn user_agent(product: &str, version: &str) -> Result<UserAgent> {
    let uname = try!(sys::uname());
    let ua = format!("{}/{} ({}-{}; {})",
                     product.trim(),
                     version.trim(),
                     uname.machine.trim().to_lowercase(),
                     uname.sys_name.trim().to_lowercase(),
                     uname.release.trim().to_lowercase());
    debug!("User-Agent: {}", &ua);
    Ok(UserAgent(ua))
}

fn ssl_ctx(fs_root_path: Option<&Path>) -> Result<SslContext> {
    // The spirit of this implementation is directly from Hyper's default OpensslClient function:
    // https://github.com/hyperium/hyper/blob/v0.9.5/src/net.rs#L653-L661
    let mut ctx = try!(SslContext::new(SslMethod::Sslv23));
    try!(ssl::set_ca(&mut ctx, fs_root_path));
    ctx.set_options(SSL_OP_NO_SSLV2 | SSL_OP_NO_SSLV3 | SSL_OP_NO_COMPRESSION);
    try!(ctx.set_cipher_list("ALL!EXPORT!EXPORT40!EXPORT56!aNULL!LOW!RC4@STRENGTH"));
    Ok(ctx)
}
<|endoftext|>


<commit_before>use bit_set::BitSet;
use byteorder::{ByteOrder, BigEndian};
use std::cmp::min;
use std::sync::{Arc, Condvar, Mutex};
use std::sync::mpsc::{self, TryRecvError};
use std::thread;
use std::fs;
use std::io::{self, Read, Write, Seek, SeekFrom};
use std::path::PathBuf;
use tempfile::TempFile;

use util::{FileId, IgnoreExt, ZeroFile, mkdir_existing};
use session::Session;
use stream::StreamEvent;

const CHUNK_SIZE : usize = 0x20000;

pub enum AudioFile {
    Direct(fs::File),
    Loading(AudioFileLoading)
}

pub struct AudioFileLoading {
    read_file: TempFile,

    position: u64,
    seek: mpsc::Sender<u64>,

    shared: Arc<AudioFileShared>,
}

struct AudioFileShared {
    file_id: FileId,
    size: usize,
    chunk_count: usize,
    cond: Condvar,
    bitmap: Mutex<BitSet>,
}

impl AudioFileLoading {
    fn new(session: &Session, file_id: FileId) -> AudioFileLoading {
        let mut files_iter = TempFile::shared(2).unwrap().into_iter();
        let read_file = files_iter.next().unwrap();
        let mut write_file = files_iter.next().unwrap();

        let size = session.stream(file_id, 0, 1).into_iter()
            .filter_map(|event| {
                match event {
                    StreamEvent::Header(id, ref data) if id == 0x3 => {
                        Some(BigEndian::read_u32(data) as usize * 4)
                    }
                    _ => None
                }
            }).next().unwrap();

        let chunk_count = (size + CHUNK_SIZE / 2) / CHUNK_SIZE;

        let shared = Arc::new(AudioFileShared {
            file_id: file_id,
            size: size,
            chunk_count: chunk_count,
            cond: Condvar::new(),
            bitmap: Mutex::new(BitSet::with_capacity(chunk_count)),
        });

        io::copy(&mut ZeroFile::new(size as u64), &mut write_file).unwrap();

        let (seek_tx, seek_rx) = mpsc::channel();

        let _shared = shared.clone();
        let _session = session.clone();

        thread::spawn(move || {
            AudioFileLoading::fetch(&_session, _shared, write_file, seek_rx)
        });

        AudioFileLoading {
            read_file: read_file,

            position: 0,
            seek: seek_tx,

            shared: shared
        }
    }

    fn fetch(session: &Session, shared: Arc<AudioFileShared>,
             mut write_file: TempFile, seek_rx: mpsc::Receiver<u64>) {
        let mut index = 0;

        loop {
            match seek_rx.try_recv() {
                Ok(position) => {
                    index = position as usize / CHUNK_SIZE;
                }
                Err(TryRecvError::Disconnected) => break,
                Err(TryRecvError::Empty) => (),
            }

            let bitmap = shared.bitmap.lock().unwrap();
            if bitmap.len() >= shared.chunk_count {
                drop(bitmap);
                AudioFileLoading::store(session, &shared, &mut write_file);
                break;
            }

            while bitmap.contains(&index) {
                index = (index + 1) % shared.chunk_count;
            }
            drop(bitmap);

            AudioFileLoading::fetch_chunk(session, &shared, &mut write_file, index);
        }
    }

    fn fetch_chunk(session: &Session, shared: &Arc<AudioFileShared>,
                   write_file: &mut TempFile, index: usize) {

        let rx = session.stream(shared.file_id,
                     (index * CHUNK_SIZE / 4) as u32,
                     (CHUNK_SIZE / 4) as u32);

        println!("Chunk {}", index);

        write_file.seek(SeekFrom::Start((index * CHUNK_SIZE) as u64)).unwrap();

        let mut size = 0usize;
        for event in rx.iter() {
            match event {
                StreamEvent::Header(..) => (),
                StreamEvent::Data(data) => {
                    write_file.write_all(&data).unwrap();

                    size += data.len();
                    if size >= CHUNK_SIZE {
                        break
                    }
                }
            }
        }

        let mut bitmap = shared.bitmap.lock().unwrap();
        bitmap.insert(index as usize);

        shared.cond.notify_all();
    }

    fn store(session: &Session, shared: &AudioFileShared, write_file: &mut TempFile) {
        write_file.seek(SeekFrom::Start(0)).unwrap();

        mkdir_existing(&AudioFileManager::cache_dir(session, shared.file_id)).unwrap();

        let mut f = fs::File::create(AudioFileManager::cache_path(session, shared.file_id)).unwrap();
        io::copy(write_file, &mut f).unwrap();
    }
}

impl Read for AudioFileLoading {
    fn read(&mut self, output: &mut [u8]) -> io::Result<usize> {
        let index = self.position as usize / CHUNK_SIZE;
        let offset = self.position as usize % CHUNK_SIZE;
        let len = min(output.len(), CHUNK_SIZE-offset);

        let mut bitmap = self.shared.bitmap.lock().unwrap();
        while !bitmap.contains(&index) {
            bitmap = self.shared.cond.wait(bitmap).unwrap();
        }
        drop(bitmap);

        let read_len = try!(self.read_file.read(&mut output[..len]));

        self.position += read_len as u64;

        Ok(read_len)
    }
}

impl Seek for AudioFileLoading {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        self.position = try!(self.read_file.seek(pos));

        /*
         * Notify the fetch thread to get the correct block
         * This can fail if fetch thread has completed, in which case the
         * block is ready. Just ignore the error.
         */
        self.seek.send(self.position).ignore();
        Ok(self.position as u64)
    }
}

impl Read for AudioFile {
    fn read(&mut self, output: &mut [u8]) -> io::Result<usize> {
        match *self {
            AudioFile::Direct(ref mut file) => file.read(output),
            AudioFile::Loading(ref mut loading) => loading.read(output),
        }
    }
}

impl Seek for AudioFile {
    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
        match *self {
            AudioFile::Direct(ref mut file) => file.seek(pos),
            AudioFile::Loading(ref mut loading) => loading.seek(pos),
        }
    }
}

pub struct AudioFileManager;
impl AudioFileManager {
    pub fn new() -> AudioFileManager {
        AudioFileManager
    }

    pub fn cache_dir(session: &Session, file_id: FileId) -> PathBuf {
        let name = file_id.to_base16();
        session.0.config.cache_location.join(&name[0..2])
    }

    pub fn cache_path(session: &Session, file_id: FileId) -> PathBuf {
        let name = file_id.to_base16();
        AudioFileManager::cache_dir(session, file_id).join(&name[2..])
    }

    pub fn request (&mut self, session: &Session, file_id: FileId) -> AudioFile {
        match fs::File::open(AudioFileManager::cache_path(session, file_id)) {
            Ok(f) => AudioFile::Direct(f),
            Err(..) => AudioFile::Loading(AudioFileLoading::new(session, file_id))
        }
    }
}

<commit_msg>Fix formula for the number of chunks.<commit_after>use bit_set::BitSet;
use byteorder::{ByteOrder, BigEndian};
use std::cmp::min;
use std::sync::{Arc, Condvar, Mutex};
use std::sync::mpsc::{self, TryRecvError};
use std::thread;
use std::fs;
use std::io::{self, Read, Write, Seek, SeekFrom};
use std::path::PathBuf;
use tempfile::TempFile;

use util::{FileId, IgnoreExt, ZeroFile, mkdir_existing};
use session::Session;
use stream::StreamEvent;

const CHUNK_SIZE : usize = 0x20000;

pub enum AudioFile {
    Direct(fs::File),
    Loading(AudioFileLoading)
}

pub struct AudioFileLoading {
    read_file: TempFile,

    position: u64,
    seek: mpsc::Sender<u64>,

    shared: Arc<AudioFileShared>,
}

struct AudioFileShared {
    file_id: FileId,
    size: usize,
    chunk_count: usize,
    cond: Condvar,
    bitmap: Mutex<BitSet>,
}

impl AudioFileLoading {
    fn new(session: &Session, file_id: FileId) -> AudioFileLoading {
        let mut files_iter = TempFile::shared(2).unwrap().into_iter();
        let read_file = files_iter.next().unwrap();
        let mut write_file = files_iter.next().unwrap();

        let size = session.stream(file_id, 0, 1).into_iter()
            .filter_map(|event| {
                match event {
                    StreamEvent::Header(id, ref data) if id == 0x3 => {
                        Some(BigEndian::read_u32(data) as usize * 4)
                    }
                    _ => None
                }
            }).next().unwrap();

        let chunk_count = (size + CHUNK_SIZE - 1) / CHUNK_SIZE;

        let shared = Arc::new(AudioFileShared {
            file_id: file_id,
            size: size,
            chunk_count: chunk_count,
            cond: Condvar::new(),
            bitmap: Mutex::new(BitSet::with_capacity(chunk_count)),
        });

        io::copy(&mut ZeroFile::new(size as u64), &mut write_file).unwrap();

        let (seek_tx, seek_rx) = mpsc::channel();

        let _shared = shared.clone();
        let _session = session.clone();

        thread::spawn(move || {
            AudioFileLoading::fetch(&_session, _shared, write_file, seek_rx)
        });

        AudioFileLoading {
            read_file: read_file,

            position: 0,
            seek: seek_tx,

            shared: shared
        }
    }

    fn fetch(session: &Session, shared: Arc<AudioFileShared>,
             mut write_file: TempFile, seek_rx: mpsc::Receiver<u64>) {
        let mut index = 0;

        loop {
            match seek_rx.try_recv() {
                Ok(position) => {
                    index = position as usize / CHUNK_SIZE;
                }
                Err(TryRecvError::Disconnected) => break,
                Err(TryRecvError::Empty) => (),
            }

            let bitmap = shared.bitmap.lock().unwrap();
            if bitmap.len() >= shared.chunk_count {
                drop(bitmap);
                AudioFileLoading::store(session, &shared, &mut write_file);
                break;
            }

            while bitmap.contains(&index) {
                index = (index + 1) % shared.chunk_count;
            }
            drop(bitmap);

            AudioFileLoading::fetch_chunk(session, &shared, &mut write_file, index);
        }
    }

    fn fetch_chunk(session: &Session, shared: &Arc<AudioFileShared>,
                   write_file: &mut TempFile, index: usize) {

        let rx = session.stream(shared.file_id,
                     (index * CHUNK_SIZE / 4) as u32,
                     (CHUNK_SIZE / 4) as u32);

        println!("Chunk {}", index);

        write_file.seek(SeekFrom::Start((index * CHUNK_SIZE) as u64)).unwrap();

        let mut size = 0usize;
        for event in rx.iter() {
            match event {
                StreamEvent::Header(..) => (),
                StreamEvent::Data(data) => {
                    write_file.write_all(&data).unwrap();

                    size += data.len();
                    if size >= CHUNK_SIZE {
                        break
                    }
                }
            }
        }

        let mut bitmap = shared.bitmap.lock().unwrap();
        bitmap.insert(index as usize);

        shared.cond.notify_all();
    }

    fn store(session: &Session, shared: &AudioFileShared, write_file: &mut TempFile) {
        write_file.seek(SeekFrom::Start(0)).unwrap();

        mkdir_existing(&AudioFileManager::cache_dir(session, shared.file_id)).unwrap();

        let mut f = fs::File::create(AudioFileManager::cache_path(session, shared.file_id)).unwrap();
        io::copy(write_file, &mut f).unwrap();
    }
}

impl Read for AudioFileLoading {
    fn read(&mut self, output: &mut [u8]) -> io::Result<usize> {
        let index = self.position as usize / CHUNK_SIZE;
        let offset = self.position as usize % CHUNK_SIZE;
        let len = min(output.len(), CHUNK_SIZE-offset);

        let mut bitmap = self.shared.bitmap.lock().unwrap();
        while !bitmap.contains(&index) {
            bitmap = self.shared.cond.wait(bitmap).unwrap();
        }
        drop(bitmap);

        let read_len = try!(self.read_file.read(&mut output[..len]));

        self.position += read_len as u64;

        Ok(read_len)
    }
}

impl Seek for AudioFileLoading {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        self.position = try!(self.read_file.seek(pos));

        /*
         * Notify the fetch thread to get the correct block
         * This can fail if fetch thread has completed, in which case the
         * block is ready. Just ignore the error.
         */
        self.seek.send(self.position).ignore();
        Ok(self.position as u64)
    }
}

impl Read for AudioFile {
    fn read(&mut self, output: &mut [u8]) -> io::Result<usize> {
        match *self {
            AudioFile::Direct(ref mut file) => file.read(output),
            AudioFile::Loading(ref mut loading) => loading.read(output),
        }
    }
}

impl Seek for AudioFile {
    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
        match *self {
            AudioFile::Direct(ref mut file) => file.seek(pos),
            AudioFile::Loading(ref mut loading) => loading.seek(pos),
        }
    }
}

pub struct AudioFileManager;
impl AudioFileManager {
    pub fn new() -> AudioFileManager {
        AudioFileManager
    }

    pub fn cache_dir(session: &Session, file_id: FileId) -> PathBuf {
        let name = file_id.to_base16();
        session.0.config.cache_location.join(&name[0..2])
    }

    pub fn cache_path(session: &Session, file_id: FileId) -> PathBuf {
        let name = file_id.to_base16();
        AudioFileManager::cache_dir(session, file_id).join(&name[2..])
    }

    pub fn request (&mut self, session: &Session, file_id: FileId) -> AudioFile {
        match fs::File::open(AudioFileManager::cache_path(session, file_id)) {
            Ok(f) => AudioFile::Direct(f),
            Err(..) => AudioFile::Loading(AudioFileLoading::new(session, file_id))
        }
    }
}

<|endoftext|>


<commit_before>// The MIT License (MIT)

// Copyright (c) 2014 Y. T. CHUNG <zonyitoo@gmail.com>

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

use std::io::net::ip::Port;
use std::io::net::tcp::TcpStream;
use std::io::IoResult;
use std::collections::TreeMap;

use crc32::Crc32;

use proto::{Proto, Operation, MultiOperation, ServerOperation, NoReplyOperation, Error, mod};

struct Server {
    pub proto: Box<Proto + Send>,
}

impl Server {
    fn connect(addr: &str, port: Port, protocol: proto::ProtoType) -> IoResult<Server> {
        Ok(Server {
            proto: match protocol {
                proto::Binary => {
                    let stream = try!(TcpStream::connect(addr, port));
                    box proto::BinaryProto::new(stream) as Box<Proto + Send>
                }
            }
        })
    }
}

/// Memcached client
///
/// ```no_run
/// use memcached::client::Client;
/// use memcached::proto::{Operation, Binary};
///
/// let mut client = Client::connect([("127.0.0.1", 11211, Binary, 1)]);
/// client.set(b"Foo", b"Bar", 0xdeadbeef, 2).unwrap();
/// let (value, flags) = client.get(b"Foo").unwrap();
///
/// assert_eq!(value.as_slice(), b"Bar");
/// assert_eq!(flags, 0xdeadbeef);
/// ```
pub struct Client {
    servers: Vec<Server>,
    key_hasher: Crc32,
    bucket: Vec<uint>,
}

impl Client {
    pub fn connect(svrs: &[(&str, Port, proto::ProtoType, uint)]) -> Client {
        let mut servers = Vec::new();
        let mut bucket = Vec::new();
        for &(addr, port, p, weight) in svrs.iter() {
            servers.push(Server::connect(addr, port, p).unwrap_or_else(|err| {
                panic!("Cannot connect server {}:{}: {}", addr, port, err);
            }));

            for _ in range(0, weight) {
                bucket.push(servers.len() - 1);
            }
        }
        Client {
            servers: servers,
            key_hasher: Crc32::new(),
            bucket: bucket,
        }
    }

    fn find_server_by_key<'a>(&'a mut self, key: &[u8]) -> &'a mut Server {
        let hash = (self.key_hasher.crc(key) >> 16) & 0x7fff;
        let idx = (hash as uint) % self.bucket.len();
        &mut self.servers[self.bucket[idx]]
    }

    pub fn flush_all(&mut self) -> Result<(), Error> {
        for s in self.servers.iter_mut() {
            try!(s.proto.flush(0));
        }
        Ok(())
    }
}

impl Operation for Client {
    fn set(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.set(key, value, flags, expiration)
    }

    fn add(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.add(key, value, flags, expiration)
    }

    fn delete(&mut self, key: &[u8]) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.delete(key)
    }

    fn replace(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.replace(key, value, flags, expiration)
    }

    fn get(&mut self, key: &[u8]) -> Result<(Vec<u8>, u32), Error> {
        let server = self.find_server_by_key(key);
        server.proto.get(key)
    }

    fn getk(&mut self, key: &[u8]) -> Result<(Vec<u8>, Vec<u8>, u32), Error> {
       let server = self.find_server_by_key(key);
       server.proto.getk(key)
    }

    fn increment(&mut self, key: &[u8], amount: u64, initial: u64, expiration: u32) -> Result<u64, Error> {
        let server = self.find_server_by_key(key);
        server.proto.increment(key, amount, initial, expiration)
    }

    fn decrement(&mut self, key: &[u8], amount: u64, initial: u64, expiration: u32) -> Result<u64, Error> {
        let server = self.find_server_by_key(key);
        server.proto.increment(key, amount, initial, expiration)
    }

    fn append(&mut self, key: &[u8], value: &[u8]) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.append(key, value)
    }

    fn prepend(&mut self, key: &[u8], value: &[u8]) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.prepend(key, value)
    }

    fn touch(&mut self, key: &[u8], expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.touch(key, expiration)
    }
}

impl MultiOperation for Client {
    fn set_multi(&mut self, kv: TreeMap<&[u8], (&[u8], u32, u32)>) -> Result<Vec<Result<(), Error>>, Error> {
        unimplemented!();
    }

    fn delete_multi(&mut self, keys: &[&[u8]]) -> Result<(), Error> {
        unimplemented!();
    }

    fn get_multi(&mut self, keys: &[&[u8]]) -> Result<TreeMap<Vec<u8>, (Vec<u8>, u32)>, Error> {
        unimplemented!();
    }
}

impl NoReplyOperation for Client {
    fn set_noreply(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.set_noreply(key, value, flags, expiration)
    }

    fn add_noreply(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.add_noreply(key, value, flags, expiration)
    }

    fn delete_noreply(&mut self, key: &[u8]) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.delete_noreply(key)
    }

    fn replace_noreply(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.replace_noreply(key, value, flags, expiration)
    }

    fn increment_noreply(&mut self, key: &[u8], amount: u64, initial: u64, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.increment_noreply(key, amount, initial, expiration)
    }

    fn decrement_noreply(&mut self, key: &[u8], amount: u64, initial: u64, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.decrement_noreply(key, amount, initial, expiration)
    }

    fn append_noreply(&mut self, key: &[u8], value: &[u8]) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.append_noreply(key, value)
    }

    fn prepend_noreply(&mut self, key: &[u8], value: &[u8]) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.prepend_noreply(key, value)
    }
}
<commit_msg>add cas op in client<commit_after>// The MIT License (MIT)

// Copyright (c) 2014 Y. T. CHUNG <zonyitoo@gmail.com>

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

use std::io::net::ip::Port;
use std::io::net::tcp::TcpStream;
use std::io::IoResult;
use std::collections::TreeMap;

use crc32::Crc32;

use proto::{Proto, Operation, MultiOperation, ServerOperation, NoReplyOperation, CasOperation, Error, mod};

struct Server {
    pub proto: Box<Proto + Send>,
}

impl Server {
    fn connect(addr: &str, port: Port, protocol: proto::ProtoType) -> IoResult<Server> {
        Ok(Server {
            proto: match protocol {
                proto::Binary => {
                    let stream = try!(TcpStream::connect(addr, port));
                    box proto::BinaryProto::new(stream) as Box<Proto + Send>
                }
            }
        })
    }
}

/// Memcached client
///
/// ```no_run
/// use memcached::client::Client;
/// use memcached::proto::{Operation, Binary};
///
/// let mut client = Client::connect([("127.0.0.1", 11211, Binary, 1)]);
/// client.set(b"Foo", b"Bar", 0xdeadbeef, 2).unwrap();
/// let (value, flags) = client.get(b"Foo").unwrap();
///
/// assert_eq!(value.as_slice(), b"Bar");
/// assert_eq!(flags, 0xdeadbeef);
/// ```
pub struct Client {
    servers: Vec<Server>,
    key_hasher: Crc32,
    bucket: Vec<uint>,
}

impl Client {
    pub fn connect(svrs: &[(&str, Port, proto::ProtoType, uint)]) -> Client {
        let mut servers = Vec::new();
        let mut bucket = Vec::new();
        for &(addr, port, p, weight) in svrs.iter() {
            servers.push(Server::connect(addr, port, p).unwrap_or_else(|err| {
                panic!("Cannot connect server {}:{}: {}", addr, port, err);
            }));

            for _ in range(0, weight) {
                bucket.push(servers.len() - 1);
            }
        }
        Client {
            servers: servers,
            key_hasher: Crc32::new(),
            bucket: bucket,
        }
    }

    fn find_server_by_key<'a>(&'a mut self, key: &[u8]) -> &'a mut Server {
        let hash = (self.key_hasher.crc(key) >> 16) & 0x7fff;
        let idx = (hash as uint) % self.bucket.len();
        &mut self.servers[self.bucket[idx]]
    }

    pub fn flush_all(&mut self) -> Result<(), Error> {
        for s in self.servers.iter_mut() {
            try!(s.proto.flush(0));
        }
        Ok(())
    }
}

impl Operation for Client {
    fn set(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.set(key, value, flags, expiration)
    }

    fn add(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.add(key, value, flags, expiration)
    }

    fn delete(&mut self, key: &[u8]) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.delete(key)
    }

    fn replace(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.replace(key, value, flags, expiration)
    }

    fn get(&mut self, key: &[u8]) -> Result<(Vec<u8>, u32), Error> {
        let server = self.find_server_by_key(key);
        server.proto.get(key)
    }

    fn getk(&mut self, key: &[u8]) -> Result<(Vec<u8>, Vec<u8>, u32), Error> {
       let server = self.find_server_by_key(key);
       server.proto.getk(key)
    }

    fn increment(&mut self, key: &[u8], amount: u64, initial: u64, expiration: u32) -> Result<u64, Error> {
        let server = self.find_server_by_key(key);
        server.proto.increment(key, amount, initial, expiration)
    }

    fn decrement(&mut self, key: &[u8], amount: u64, initial: u64, expiration: u32) -> Result<u64, Error> {
        let server = self.find_server_by_key(key);
        server.proto.increment(key, amount, initial, expiration)
    }

    fn append(&mut self, key: &[u8], value: &[u8]) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.append(key, value)
    }

    fn prepend(&mut self, key: &[u8], value: &[u8]) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.prepend(key, value)
    }

    fn touch(&mut self, key: &[u8], expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.touch(key, expiration)
    }
}

impl MultiOperation for Client {
    fn set_multi(&mut self, kv: TreeMap<&[u8], (&[u8], u32, u32)>) -> Result<Vec<Result<(), Error>>, Error> {
        unimplemented!();
    }

    fn delete_multi(&mut self, keys: &[&[u8]]) -> Result<(), Error> {
        unimplemented!();
    }

    fn get_multi(&mut self, keys: &[&[u8]]) -> Result<TreeMap<Vec<u8>, (Vec<u8>, u32)>, Error> {
        unimplemented!();
    }
}

impl NoReplyOperation for Client {
    fn set_noreply(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.set_noreply(key, value, flags, expiration)
    }

    fn add_noreply(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.add_noreply(key, value, flags, expiration)
    }

    fn delete_noreply(&mut self, key: &[u8]) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.delete_noreply(key)
    }

    fn replace_noreply(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.replace_noreply(key, value, flags, expiration)
    }

    fn increment_noreply(&mut self, key: &[u8], amount: u64, initial: u64, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.increment_noreply(key, amount, initial, expiration)
    }

    fn decrement_noreply(&mut self, key: &[u8], amount: u64, initial: u64, expiration: u32) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.decrement_noreply(key, amount, initial, expiration)
    }

    fn append_noreply(&mut self, key: &[u8], value: &[u8]) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.append_noreply(key, value)
    }

    fn prepend_noreply(&mut self, key: &[u8], value: &[u8]) -> Result<(), Error> {
        let server = self.find_server_by_key(key);
        server.proto.prepend_noreply(key, value)
    }
}

impl CasOperation for Client {
    fn set_cas(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32, cas: u64) -> Result<u64, Error> {
        let server = self.find_server_by_key(key);
        server.proto.set_cas(key, value, flags, expiration, cas)
    }

    fn add_cas(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32) -> Result<u64, Error> {
        let server = self.find_server_by_key(key);
        server.proto.add_cas(key, value, flags, expiration)
    }

    fn replace_cas(&mut self, key: &[u8], value: &[u8], flags: u32, expiration: u32, cas: u64) -> Result<u64, Error> {
        let server = self.find_server_by_key(key);
        server.proto.replace_cas(key, value, flags, expiration, cas)
    }

    fn get_cas(&mut self, key: &[u8]) -> Result<(Vec<u8>, u32, u64), Error> {
        let server = self.find_server_by_key(key);
        server.proto.get_cas(key)
    }

    fn getk_cas(&mut self, key: &[u8]) -> Result<(Vec<u8>, Vec<u8>, u32, u64), Error> {
        let server = self.find_server_by_key(key);
        server.proto.getk_cas(key)
    }

    fn increment_cas(&mut self, key: &[u8], amount: u64, initial: u64, expiration: u32, cas: u64)
            -> Result<(u64, u64), Error> {
        let server = self.find_server_by_key(key);
        server.proto.increment_cas(key, amount, initial, expiration, cas)
    }

    fn decrement_cas(&mut self, key: &[u8], amount: u64, initial: u64, expiration: u32, cas: u64)
            -> Result<(u64, u64), Error> {
        let server = self.find_server_by_key(key);
        server.proto.decrement_cas(key, amount, initial, expiration, cas)
    }

    fn append_cas(&mut self, key: &[u8], value: &[u8], cas: u64) -> Result<u64, Error> {
        let server = self.find_server_by_key(key);
        server.proto.append_cas(key, value, cas)
    }

    fn prepend_cas(&mut self, key: &[u8], value: &[u8], cas: u64) -> Result<u64, Error> {
        let server = self.find_server_by_key(key);
        server.proto.prepend_cas(key, value, cas)
    }

    fn touch_cas(&mut self, key: &[u8], expiration: u32, cas: u64) -> Result<u64, Error> {
        let server = self.find_server_by_key(key);
        server.proto.touch_cas(key, expiration, cas)
    }
}
<|endoftext|>


<commit_before>pub mod error;

use rustc_serialize::json;
use std::env;
use std::path::PathBuf;
use std::fs::File;
use std::io::Error as StdIoError;
use std::io::ErrorKind;

use std::io::prelude::*; // Required for read_to_string use later.

use self::error::ConfigError;

use cli::arg_types::AccountType;

#[derive(Debug, RustcEncodable, RustcDecodable)]
pub struct Config {
    pub auth_token: String,
    pub current: String,
    pub savings: String,
    pub business: String,
}

impl Config {
    pub fn new<S: Into<String>>(auth_token: S, current: S, savings: S, business: S) -> Config {
        Config {
            auth_token: auth_token.into(),
            current: current.into(),
            savings: savings.into(),
            business: business.into(),
        }
    }

    pub fn new_with_auth_token_only<S: Into<String>>(auth_token: S) -> Config {
        Config::new(auth_token.into(),
                    "".to_string(),
                    "".to_string(),
                    "".to_string())
    }

    pub fn get_account_id(&self, account: &AccountType) -> String {
        let default_account_id = self.current.to_owned();
        match *account {
            AccountType::Current => self.current.to_owned(),
            AccountType::Savings => self.savings.to_owned(),
            AccountType::Business => self.business.to_owned(),
            _ => default_account_id,
        }
    }

    pub fn get_account_alias_for_id<'a>(&self, account_id: &str) -> &'a str {
        if *account_id == self.current {
            "(current)"
        } else if *account_id == self.savings {
            "(savings)"
        } else if *account_id == self.business {
            "(business)"
        } else {
            ""
        }
    }
}

pub fn get_config_path() -> PathBuf {
    let fallback_config_path = PathBuf::from("./.tellerrc");
    let append_config_file = |mut p: PathBuf| {
        p.push(".tellerrc");
        p
    };
    env::home_dir().map_or(fallback_config_path, append_config_file)
}

pub fn get_config_file(config_path: &PathBuf) -> Option<File> {
    let config_path_str = config_path.to_str().unwrap_or("[error: config_path#to_str fails]");
    info!("Checking whether config file within {} exists",
          config_path_str);
    let config_file = File::open(&config_path);
    match config_file {
        Err(ref e) if ErrorKind::NotFound == e.kind() => {
            debug!("No config file found");
            None
        }
        Err(_) => panic!("Unable to read config!"),
        Ok(config_file) => Some(config_file),
    }
}

pub fn get_config_file_to_write(config_path: &PathBuf) -> Result<File, StdIoError> {
    let config_file = File::create(&config_path);
    match config_file {
        Err(ref e) if ErrorKind::PermissionDenied == e.kind() => {
            panic!("Permission to read config denied")
        }
        _ => config_file,
    }
}

pub fn get_config() -> Option<Config> {
    let config_file_path = get_config_path();
    match get_config_file(&config_file_path) {
        None => None,
        Some(mut config_file) => {
            match read_config(&mut config_file) {
                Ok(config) => Some(config),
                Err(e) => {
                    panic!("ERROR: attempting to read file {}: {}",
                           config_file_path.display(),
                           e)
                }
            }
        }
    }
}

pub fn read_config(config_file: &mut File) -> Result<Config, ConfigError> {
    let mut content_str = String::new();
    try!(config_file.read_to_string(&mut content_str));

    debug!("config read as: {}", content_str);

    let config: Config = try!(json::decode(&content_str));

    Ok(config)
}

pub fn write_config(config_file: &mut File, config: &Config) -> Result<(), ConfigError> {
    let content_str = try!(json::encode(&config));

    try!(config_file.write_all(content_str.as_bytes()));

    Ok(())
}
<commit_msg>tests(config): starting config tests<commit_after>pub mod error;

use rustc_serialize::json;
use std::env;
use std::path::PathBuf;
use std::fs::File;
use std::io::Error as StdIoError;
use std::io::ErrorKind;

use std::io::prelude::*; // Required for read_to_string use later.

use self::error::ConfigError;

use cli::arg_types::AccountType;

#[derive(Debug, RustcEncodable, RustcDecodable)]
pub struct Config {
    pub auth_token: String,
    pub current: String,
    pub savings: String,
    pub business: String,
}

impl Config {
    pub fn new<S: Into<String>>(auth_token: S, current: S, savings: S, business: S) -> Config {
        Config {
            auth_token: auth_token.into(),
            current: current.into(),
            savings: savings.into(),
            business: business.into(),
        }
    }

    pub fn new_with_auth_token_only<S: Into<String>>(auth_token: S) -> Config {
        Config::new(auth_token.into(),
                    "".to_string(),
                    "".to_string(),
                    "".to_string())
    }

    pub fn get_account_id(&self, account: &AccountType) -> String {
        let default_account_id = self.current.to_owned();
        match *account {
            AccountType::Current => self.current.to_owned(),
            AccountType::Savings => self.savings.to_owned(),
            AccountType::Business => self.business.to_owned(),
            _ => default_account_id,
        }
    }

    pub fn get_account_alias_for_id<'a>(&self, account_id: &str) -> &'a str {
        if *account_id == self.current {
            "(current)"
        } else if *account_id == self.savings {
            "(savings)"
        } else if *account_id == self.business {
            "(business)"
        } else {
            ""
        }
    }
}

pub fn get_config_path() -> PathBuf {
    let fallback_config_path = PathBuf::from("./.tellerrc");
    let append_config_file = |mut p: PathBuf| {
        p.push(".tellerrc");
        p
    };
    env::home_dir().map_or(fallback_config_path, append_config_file)
}

pub fn get_config_file(config_path: &PathBuf) -> Option<File> {
    let config_path_str = config_path.to_str().unwrap_or("[error: config_path#to_str fails]");
    info!("Checking whether config file within {} exists",
          config_path_str);
    let config_file = File::open(&config_path);
    match config_file {
        Err(ref e) if ErrorKind::NotFound == e.kind() => {
            debug!("No config file found");
            None
        }
        Err(_) => panic!("Unable to read config!"),
        Ok(config_file) => Some(config_file),
    }
}

pub fn get_config_file_to_write(config_path: &PathBuf) -> Result<File, StdIoError> {
    let config_file = File::create(&config_path);
    match config_file {
        Err(ref e) if ErrorKind::PermissionDenied == e.kind() => {
            panic!("Permission to read config denied")
        }
        _ => config_file,
    }
}

pub fn get_config() -> Option<Config> {
    let config_file_path = get_config_path();
    match get_config_file(&config_file_path) {
        None => None,
        Some(mut config_file) => {
            match read_config(&mut config_file) {
                Ok(config) => Some(config),
                Err(e) => {
                    panic!("ERROR: attempting to read file {}: {}",
                           config_file_path.display(),
                           e)
                }
            }
        }
    }
}

pub fn read_config(config_file: &mut File) -> Result<Config, ConfigError> {
    let mut content_str = String::new();
    try!(config_file.read_to_string(&mut content_str));

    debug!("config read as: {}", content_str);

    let config: Config = try!(json::decode(&content_str));

    Ok(config)
}

pub fn write_config(config_file: &mut File, config: &Config) -> Result<(), ConfigError> {
    let content_str = try!(json::encode(&config));

    try!(config_file.write_all(content_str.as_bytes()));

    Ok(())
}

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

    #[test]
    fn can_instantiate_config() {
        let expected_auth_token = "fake-auth-token";
        let expected_current = "current-id";
        let expected_savings = "savings-id";
        let expected_business = "business-id";

        let config = Config::new(expected_auth_token, expected_current, expected_savings, expected_business);

        assert_eq!(expected_auth_token, config.auth_token);
        assert_eq!(expected_current, config.current);
        assert_eq!(expected_savings, config.savings);
        assert_eq!(expected_business, config.business);
    }

}
<|endoftext|>


<commit_before>
use crate::core::{Edge, EdgeWeighted, Directedness, trait_aliases::{
	Id,
}, Directed};
use std::iter::Iterator;

#[macro_use]
mod macros {
	#[macro_export]
	macro_rules! edges_between {
		($e:expr, $v1:expr, $v2:expr) => {
			{
				// Filter out any edge that is not connected to both vertices
				let relevant = $e.filter(move|edge|{
					(edge.source() == $v1 && edge.sink() == $v2) ||
						(edge.source() == $v2 && edge.sink() == $v1)
				});

				// Return the result
				Box::new(relevant)
			}
		}
	}
	#[macro_export]
	macro_rules! edges_incident_on {
		($e:expr, $v:expr, $i:ident) => {
			Box::new($e.into_iter().filter(move|e| e.$i() == $v))
		};
		($e:expr, $v:expr) => {
			Box::new($e.into_iter().filter(move|edge| (edge.source() == $v) || (edge.sink() == $v)))
		}
	}
}

///
/// The basic graph interface. This is the main trait for all types of graphs.
///
/// For all graphs, vertices are identified by their [value](#associatedType.Vertex), and must be unique
/// in the graph. Edges are identified by the two vertices they are incident on, the direction,
/// and their [Id](#associatedType.EdgeId); I.e. the triple `(v1,v2,e)` is an edge from `v1` to `v2`
/// with an edge id `e`. All edges are directed but need not be unique. Therefore, multiple edges with
/// the same source (`v1`), sink (`v2`), and id (`e`) may be present in a graph at the same time.
/// Edges with the same source, sink, and id are identical and must be interchangeable. E.g. if any
/// one of two or more identical edges is to be removed, then any one of them may be removed.
///
/// TODO: When Rust supports 'impl trait' consider having some of the iterators be clone too
/// (those that don't include mutable references). This allows cloning the iterators mid-iteration,
/// which can be useful when comparing each item to the ones after it.
///
pub trait Graph
{
	///
	/// Type of the graph's vertex value.
	///
	type Vertex: Id;
	type VertexWeight;
	type EdgeWeight;
	type Directedness: Directedness;

	///
	/// Returns copies of all current vertex values in the graph.
	///
	fn all_vertices_weighted<'a>(&'a self)
		 -> Box<dyn 'a + Iterator<Item=(Self::Vertex, &'a Self::VertexWeight)>>;
	///
	/// Returns copies of all current edges in the graph.
	///
	fn all_edges<'a>(&'a self) -> Box<dyn 'a + Iterator<Item=
		(Self::Vertex, Self::Vertex, &'a Self::EdgeWeight)>>;

// Optional methods

	fn all_vertices<'a>(&'a self) -> Box<dyn 'a + Iterator<Item=Self::Vertex>>
	{
		Box::new(self.all_vertices_weighted().map(|(v, _)| v))
	}
	fn vertex_weight(&self, v: Self::Vertex) -> Option<&Self::VertexWeight>
	{
		self.all_vertices_weighted().find(|&(candidate,_)| candidate == v).map(|(_,w)| w)
	}
	fn contains_vertex(&self, v: Self::Vertex) -> bool
	{
		self.vertex_weight(v).is_some()
	}
	fn all_vertex_weights<'a>(&'a self) -> Box<dyn 'a + Iterator<Item=&'a Self::VertexWeight>>
	{
		Box::new(self.all_vertices_weighted().map(|(_,w)| w))
	}

	///
	/// Returns all edges that are incident on both of the given vertices, regardless of direction.
	///
	/// I.e. all edges where e == (v1,v2,_) or e == (v2,v1,_)
	///
	fn edges_between<'a>(&'a self, v1: Self::Vertex, v2: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a Self::EdgeWeight)>>
	{
		edges_between!(self.all_edges(), v1, v2)
	}
	///
	/// Returns all edges that are sourced in the given vertex.
	///
	/// I.e. all edges where `e == (v,_,_)`
	///
	/// Only available for directed graphs
	///
	fn edges_sourced_in<'a>(&'a self, v: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a Self::EdgeWeight)>>
		where Self: Graph<Directedness=Directed>
	{
		edges_incident_on!(self.all_edges(), v, source)
	}
	///
	/// Returns all edges that are sinked in the given vertex.
	///
	/// I.e. all edges where `e == (_,v,_)`
	///
	fn edges_sinked_in<'a>(&'a self, v: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a Self::EdgeWeight)>>
		where Self: Graph<Directedness=Directed>
	{
		edges_incident_on!(self.all_edges(), v, sink)
	}
	fn edges_incident_on<'a>(&'a self, v: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a Self::EdgeWeight)>>
	{
		edges_incident_on!(self.all_edges(),v)
	}
}

pub trait GraphMut: Graph
{
	fn all_vertices_weighted_mut<'a>(&'a mut self)
									 -> Box<dyn 'a + Iterator<Item=(Self::Vertex, &'a mut Self::VertexWeight)>>;

	fn all_edges_mut<'a>(&'a mut self) -> Box<dyn 'a + Iterator<Item=
	(Self::Vertex, Self::Vertex, &'a mut Self::EdgeWeight)>>;

// Optional methods

	fn vertex_weight_mut(&mut self, v: Self::Vertex) -> Option<&mut Self::VertexWeight>
	{
		self.all_vertices_weighted_mut().find(|&(candidate,_)| candidate == v).map(|(_,w)| w)
	}

	fn edges_between_mut<'a>(&'a mut self, v1: Self::Vertex, v2: Self::Vertex)
		 -> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a mut Self::EdgeWeight)>>
	{
		edges_between!(self.all_edges_mut(), v1, v2)
	}
	fn edges_sourced_in_mut<'a>(&'a mut self, v: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a mut Self::EdgeWeight)>>
		where Self: Graph<Directedness=Directed>
	{
		edges_incident_on!(self.all_edges_mut(), v, source)
	}
	fn edges_sinked_in_mut<'a>(&'a mut self, v: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a mut Self::EdgeWeight)>>
		where Self: Graph<Directedness=Directed>
	{
		edges_incident_on!(self.all_edges_mut(), v, sink)
	}
	fn edges_incident_on_mut<'a>(&'a mut self, v: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a mut Self::EdgeWeight)>>
	{
		edges_incident_on!(self.all_edges_mut(),v)
	}

}

///
/// A graph where the vertex ids are assigned automatically.
///
pub trait AddVertex: GraphMut
{
	///
	/// Adds a new vertex with the given weight to the graph.
	/// Returns the id of the new vertex.
	///
	fn new_vertex_weighted(&mut self, w: Self::VertexWeight) -> Result<Self::Vertex,()>;
	///
	/// Removes the given vertex from the graph, returning its weight.
	/// If the vertex still has edges incident on it, they are also removed,
	/// dropping their weights.
	///
	fn remove_vertex(&mut self, v: Self::Vertex) -> Result<Self::VertexWeight,()>;

// Optional methods
	///
	/// Adds a new vertex to the graph.
	/// Returns the id of the new vertex.
	/// The weight of the vertex is the default.
	///
	fn new_vertex(&mut self) -> Result<Self::Vertex,()>
		where Self::VertexWeight: Default
	{
		self.new_vertex_weighted(Self::VertexWeight::default())
	}

}

pub trait AddEdge: GraphMut
{
	///
	/// Adds a copy of the given edge to the graph
	///
	fn add_edge_weighted<E>(&mut self, e: E) -> Result<(),()>
		where E: EdgeWeighted<Self::Vertex, Self::EdgeWeight>;
	///
	/// Removes an edge that matches the given predicate closure.
	/// If no edge is found to match and successfully removed, returns error
	/// but otherwise doesn't change the graph.
	///
	fn remove_edge_where<F>(&mut self, f: F) -> Result<(Self::Vertex, Self::Vertex, Self::EdgeWeight), ()>
		where F: Fn((Self::Vertex, Self::Vertex, &Self::EdgeWeight)) -> bool;

// Optional methods

	///
	/// Adds the given edge to the graph, regardless of whether there are existing,
	/// identical edges in the graph.
	/// The vertices the new edge is incident on must exist in the graph and the id must be valid.
	///
	/// ###Returns
	/// - `Ok` if the edge is valid and was added to the graph.
	/// - `Err` if the edge is invalid or the graph was otherwise unable to store it.
	///
	/// ###`Ok` properties:
	///
	/// - Only the given edge is added to the graph.
	/// - Existing edges are unchanged.
	/// - No vertices are introduced or removed.
	///
	/// ###`Err` properties:
	///
	/// - The graph is unchanged.
	///
	fn add_edge<E>(&mut self, e: E) -> Result<(),()>
		where
			E: Edge<Self::Vertex>,
			Self::EdgeWeight: Default,
	{
		self.add_edge_weighted((e.source(), e.sink(), Self::EdgeWeight::default()))
	}
	///
	/// Removes the given edge from the graph if it exists.
	///
	/// ###Returns
	/// - `Ok` if the edge was present before the call and was removed.
	/// - `Err` if the edge was not found in the graph or it was otherwise unable to remove it.
	///
	/// ###`Ok` properties:
	///
	/// - One edge identical to the given edge is removed.
	/// - No new edges are introduced.
	/// - No edges are changed.
	/// - No new vertices are introduced or removed.
	///
	/// ###`Err` properties:
	///
	/// - The graph is unchanged.
	///
	fn remove_edge<E>(&mut self, e: E) -> Result<Self::EdgeWeight,()>
		where E: Edge<Self::Vertex>
	{
		self.remove_edge_where_weight(e, |_| true)
	}
	fn remove_edge_where_weight<E,F>(&mut self, e: E, f: F) -> Result<Self::EdgeWeight,()>
		where
			E: Edge<Self::Vertex>,
			F: Fn(&Self::EdgeWeight) -> bool,
	{
		self.remove_edge_where(|(so,si, w)|
			(so == e.source()) && (si == e.sink()) && f(w))
			.map(|removed_edge| removed_edge.2)
	}

	///
	/// Tries to removes all edges that match the given predicate.
	/// Stops either when no more edges match, or an edge that matched couldn't be removed.
	///
	/// Returns the removed edges regardless.
	///
	fn remove_edge_where_all<F>(&mut self, f: F) -> Vec<(Self::Vertex, Self::Vertex, Self::EdgeWeight)>
		where F: Fn((Self::Vertex, Self::Vertex, &Self::EdgeWeight)) -> bool
	{
		let mut result = Vec::new();
		while let Ok(e) = self.remove_edge_where(|e| f(e)) {
			result.push(e);
		}
		result
	}
}

///
/// Graph that at all times has a finite set of vertices and edges.
///
/// TODO: redesign. What about edges? even through there are less than usize::MAX vertices
/// there can easily be be more than usize::MAX edges. Maybe this should be 2 traits.
///
pub trait ExactGraph: Graph
{

	///
	/// Returns the number of vertices in the graph.
	///
	fn vertex_count(&self) -> usize {
		self.all_vertices().count()
	}

	///
	/// Returns the number of edges in the graph.
	///
	fn edge_count(&self) -> usize {
		self.all_edges().count()
	}
}
<commit_msg>Added the methods 'constrain_directedness(_mut)' to Graph(Mut).<commit_after>
use crate::core::{Edge, EdgeWeighted, Directedness, trait_aliases::{
	Id,
}, Directed, Constrainer, BaseGraph};
use std::iter::Iterator;
use crate::core::constraint::{DirectedGraph, UndirectedGraph};

#[macro_use]
mod macros {
	#[macro_export]
	macro_rules! edges_between {
		($e:expr, $v1:expr, $v2:expr) => {
			{
				// Filter out any edge that is not connected to both vertices
				let relevant = $e.filter(move|edge|{
					(edge.source() == $v1 && edge.sink() == $v2) ||
						(edge.source() == $v2 && edge.sink() == $v1)
				});

				// Return the result
				Box::new(relevant)
			}
		}
	}
	#[macro_export]
	macro_rules! edges_incident_on {
		($e:expr, $v:expr, $i:ident) => {
			Box::new($e.into_iter().filter(move|e| e.$i() == $v))
		};
		($e:expr, $v:expr) => {
			Box::new($e.into_iter().filter(move|edge| (edge.source() == $v) || (edge.sink() == $v)))
		}
	}
}

///
/// The basic graph interface. This is the main trait for all types of graphs.
///
/// For all graphs, vertices are identified by their [value](#associatedType.Vertex), and must be unique
/// in the graph. Edges are identified by the two vertices they are incident on, the direction,
/// and their [Id](#associatedType.EdgeId); I.e. the triple `(v1,v2,e)` is an edge from `v1` to `v2`
/// with an edge id `e`. All edges are directed but need not be unique. Therefore, multiple edges with
/// the same source (`v1`), sink (`v2`), and id (`e`) may be present in a graph at the same time.
/// Edges with the same source, sink, and id are identical and must be interchangeable. E.g. if any
/// one of two or more identical edges is to be removed, then any one of them may be removed.
///
/// TODO: When Rust supports 'impl trait' consider having some of the iterators be clone too
/// (those that don't include mutable references). This allows cloning the iterators mid-iteration,
/// which can be useful when comparing each item to the ones after it.
///
pub trait Graph
{
	///
	/// Type of the graph's vertex value.
	///
	type Vertex: Id;
	type VertexWeight;
	type EdgeWeight;
	type Directedness: Directedness;

	///
	/// Returns copies of all current vertex values in the graph.
	///
	fn all_vertices_weighted<'a>(&'a self)
		 -> Box<dyn 'a + Iterator<Item=(Self::Vertex, &'a Self::VertexWeight)>>;
	///
	/// Returns copies of all current edges in the graph.
	///
	fn all_edges<'a>(&'a self) -> Box<dyn 'a + Iterator<Item=
		(Self::Vertex, Self::Vertex, &'a Self::EdgeWeight)>>;

// Optional methods

	fn all_vertices<'a>(&'a self) -> Box<dyn 'a + Iterator<Item=Self::Vertex>>
	{
		Box::new(self.all_vertices_weighted().map(|(v, _)| v))
	}
	fn vertex_weight(&self, v: Self::Vertex) -> Option<&Self::VertexWeight>
	{
		self.all_vertices_weighted().find(|&(candidate,_)| candidate == v).map(|(_,w)| w)
	}
	fn contains_vertex(&self, v: Self::Vertex) -> bool
	{
		self.vertex_weight(v).is_some()
	}
	fn all_vertex_weights<'a>(&'a self) -> Box<dyn 'a + Iterator<Item=&'a Self::VertexWeight>>
	{
		Box::new(self.all_vertices_weighted().map(|(_,w)| w))
	}

	///
	/// Returns all edges that are incident on both of the given vertices, regardless of direction.
	///
	/// I.e. all edges where e == (v1,v2,_) or e == (v2,v1,_)
	///
	fn edges_between<'a>(&'a self, v1: Self::Vertex, v2: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a Self::EdgeWeight)>>
	{
		edges_between!(self.all_edges(), v1, v2)
	}
	///
	/// Returns all edges that are sourced in the given vertex.
	///
	/// I.e. all edges where `e == (v,_,_)`
	///
	/// Only available for directed graphs
	///
	fn edges_sourced_in<'a>(&'a self, v: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a Self::EdgeWeight)>>
		where Self: Graph<Directedness=Directed>
	{
		edges_incident_on!(self.all_edges(), v, source)
	}
	///
	/// Returns all edges that are sinked in the given vertex.
	///
	/// I.e. all edges where `e == (_,v,_)`
	///
	fn edges_sinked_in<'a>(&'a self, v: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a Self::EdgeWeight)>>
		where Self: Graph<Directedness=Directed>
	{
		edges_incident_on!(self.all_edges(), v, sink)
	}
	fn edges_incident_on<'a>(&'a self, v: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a Self::EdgeWeight)>>
	{
		edges_incident_on!(self.all_edges(),v)
	}

	fn constrain_directedness(&self) -> DirectednessVariants<&Self>
		where Self: BaseGraph
	{
		if let Ok(g) = self.constrain() {
			DirectednessVariants::Directed(g)
		} else {
			DirectednessVariants::Undirected(UndirectedGraph::unchecked(self))
		}
	}
}

pub trait GraphMut: Graph
{
	fn all_vertices_weighted_mut<'a>(&'a mut self)
									 -> Box<dyn 'a + Iterator<Item=(Self::Vertex, &'a mut Self::VertexWeight)>>;

	fn all_edges_mut<'a>(&'a mut self) -> Box<dyn 'a + Iterator<Item=
	(Self::Vertex, Self::Vertex, &'a mut Self::EdgeWeight)>>;

// Optional methods

	fn vertex_weight_mut(&mut self, v: Self::Vertex) -> Option<&mut Self::VertexWeight>
	{
		self.all_vertices_weighted_mut().find(|&(candidate,_)| candidate == v).map(|(_,w)| w)
	}

	fn edges_between_mut<'a>(&'a mut self, v1: Self::Vertex, v2: Self::Vertex)
		 -> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a mut Self::EdgeWeight)>>
	{
		edges_between!(self.all_edges_mut(), v1, v2)
	}
	fn edges_sourced_in_mut<'a>(&'a mut self, v: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a mut Self::EdgeWeight)>>
		where Self: Graph<Directedness=Directed>
	{
		edges_incident_on!(self.all_edges_mut(), v, source)
	}
	fn edges_sinked_in_mut<'a>(&'a mut self, v: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a mut Self::EdgeWeight)>>
		where Self: Graph<Directedness=Directed>
	{
		edges_incident_on!(self.all_edges_mut(), v, sink)
	}
	fn edges_incident_on_mut<'a>(&'a mut self, v: Self::Vertex)
		-> Box<dyn 'a + Iterator<Item=(Self::Vertex, Self::Vertex, &'a mut Self::EdgeWeight)>>
	{
		edges_incident_on!(self.all_edges_mut(),v)
	}

	fn constrain_directedness_mut(&mut self) -> DirectednessVariants<& mut Self>
		where Self: BaseGraph
	{
		unimplemented!("I suspect the below implementation to work, but needs review.");
		/*  We use this unsafe block to allow us to use 'self' in the 'else' branch below.
			This is safe because 'self' is not used by 'g' if the 'else' branch is taken and
			therefore the reference is not shared with anyone.
			However, the compiler cannot see this and there doesn't seem to a fix in the pipeline:
			https://github.com/rust-lang/rust/issues/53528

			If/when this issue is fixed, this code can be updated to remove the unsafe block.
		*/
		let self_2: &mut Self = unsafe {  (self as *mut Self).as_mut().unwrap() } ;

		if let Ok(g) = self.constrain() {
			DirectednessVariants::Directed(g)
		} else {
			DirectednessVariants::Undirected(UndirectedGraph::unchecked(self_2))
		}
	}
}

///
/// A graph where the vertex ids are assigned automatically.
///
pub trait AddVertex: GraphMut
{
	///
	/// Adds a new vertex with the given weight to the graph.
	/// Returns the id of the new vertex.
	///
	fn new_vertex_weighted(&mut self, w: Self::VertexWeight) -> Result<Self::Vertex,()>;
	///
	/// Removes the given vertex from the graph, returning its weight.
	/// If the vertex still has edges incident on it, they are also removed,
	/// dropping their weights.
	///
	fn remove_vertex(&mut self, v: Self::Vertex) -> Result<Self::VertexWeight,()>;

// Optional methods
	///
	/// Adds a new vertex to the graph.
	/// Returns the id of the new vertex.
	/// The weight of the vertex is the default.
	///
	fn new_vertex(&mut self) -> Result<Self::Vertex,()>
		where Self::VertexWeight: Default
	{
		self.new_vertex_weighted(Self::VertexWeight::default())
	}

}

pub trait AddEdge: GraphMut
{
	///
	/// Adds a copy of the given edge to the graph
	///
	fn add_edge_weighted<E>(&mut self, e: E) -> Result<(),()>
		where E: EdgeWeighted<Self::Vertex, Self::EdgeWeight>;
	///
	/// Removes an edge that matches the given predicate closure.
	/// If no edge is found to match and successfully removed, returns error
	/// but otherwise doesn't change the graph.
	///
	fn remove_edge_where<F>(&mut self, f: F) -> Result<(Self::Vertex, Self::Vertex, Self::EdgeWeight), ()>
		where F: Fn((Self::Vertex, Self::Vertex, &Self::EdgeWeight)) -> bool;

// Optional methods

	///
	/// Adds the given edge to the graph, regardless of whether there are existing,
	/// identical edges in the graph.
	/// The vertices the new edge is incident on must exist in the graph and the id must be valid.
	///
	/// ###Returns
	/// - `Ok` if the edge is valid and was added to the graph.
	/// - `Err` if the edge is invalid or the graph was otherwise unable to store it.
	///
	/// ###`Ok` properties:
	///
	/// - Only the given edge is added to the graph.
	/// - Existing edges are unchanged.
	/// - No vertices are introduced or removed.
	///
	/// ###`Err` properties:
	///
	/// - The graph is unchanged.
	///
	fn add_edge<E>(&mut self, e: E) -> Result<(),()>
		where
			E: Edge<Self::Vertex>,
			Self::EdgeWeight: Default,
	{
		self.add_edge_weighted((e.source(), e.sink(), Self::EdgeWeight::default()))
	}
	///
	/// Removes the given edge from the graph if it exists.
	///
	/// ###Returns
	/// - `Ok` if the edge was present before the call and was removed.
	/// - `Err` if the edge was not found in the graph or it was otherwise unable to remove it.
	///
	/// ###`Ok` properties:
	///
	/// - One edge identical to the given edge is removed.
	/// - No new edges are introduced.
	/// - No edges are changed.
	/// - No new vertices are introduced or removed.
	///
	/// ###`Err` properties:
	///
	/// - The graph is unchanged.
	///
	fn remove_edge<E>(&mut self, e: E) -> Result<Self::EdgeWeight,()>
		where E: Edge<Self::Vertex>
	{
		self.remove_edge_where_weight(e, |_| true)
	}
	fn remove_edge_where_weight<E,F>(&mut self, e: E, f: F) -> Result<Self::EdgeWeight,()>
		where
			E: Edge<Self::Vertex>,
			F: Fn(&Self::EdgeWeight) -> bool,
	{
		self.remove_edge_where(|(so,si, w)|
			(so == e.source()) && (si == e.sink()) && f(w))
			.map(|removed_edge| removed_edge.2)
	}

	///
	/// Tries to removes all edges that match the given predicate.
	/// Stops either when no more edges match, or an edge that matched couldn't be removed.
	///
	/// Returns the removed edges regardless.
	///
	fn remove_edge_where_all<F>(&mut self, f: F) -> Vec<(Self::Vertex, Self::Vertex, Self::EdgeWeight)>
		where F: Fn((Self::Vertex, Self::Vertex, &Self::EdgeWeight)) -> bool
	{
		let mut result = Vec::new();
		while let Ok(e) = self.remove_edge_where(|e| f(e)) {
			result.push(e);
		}
		result
	}
}

///
/// Graph that at all times has a finite set of vertices and edges.
///
/// TODO: redesign. What about edges? even through there are less than usize::MAX vertices
/// there can easily be be more than usize::MAX edges. Maybe this should be 2 traits.
///
pub trait ExactGraph: Graph
{

	///
	/// Returns the number of vertices in the graph.
	///
	fn vertex_count(&self) -> usize {
		self.all_vertices().count()
	}

	///
	/// Returns the number of edges in the graph.
	///
	fn edge_count(&self) -> usize {
		self.all_edges().count()
	}
}

pub enum DirectednessVariants<C: Constrainer>
{
	Directed(DirectedGraph<C>),
	Undirected(UndirectedGraph<C>),
}<|endoftext|>