<commit_before>use std::str::from_utf8;
use std::sync::Arc;
use std::io::BufWriter;
use std::mem;

use futures::Async;
use futures::future::{FutureResult, ok};
use minihttp::{Status, Version};
use minihttp::client as http;
use tokio_core::io::Io;
use rustc_serialize::Encodable;
use rustc_serialize::json::{self, as_json, Json};

use intern::SessionId;
use base64::Base64;
use proxy::{Response};
use chat::{Cid, ConnectionSender, ConnectionMessage};
use chat::cid::{serialize_cid};
use chat::error::MessageError;
use chat::message::{AuthData, Auth, Call, Meta, Args, Kwargs};
use chat::processor::{ProcessorPool, Action};
use chat::authorize::{parse_userinfo, good_status};
use chat::ConnectionMessage::{Hello, StopSocket};
use chat::CloseReason::{AuthHttp};

enum AuthState {
    Init(String, AuthData),
    Wait,
    Headers(Status),
    Done(Response),
    Void,
}

enum CallState {
    Init {
        auth: Arc<String>,
        path: String, meta: Meta,
        args: Args, kw: Kwargs
    },
    Wait(Meta),
    Headers(Meta, Status),
    Void,
}

pub struct AuthCodec {
    state: AuthState,
    chat: ProcessorPool,
    conn_id: Cid,
    sender: ConnectionSender,
}

pub struct CallCodec {
    state: CallState,
    sender: ConnectionSender,
}

impl AuthCodec {
    pub fn new(path: String, cid: Cid, req: AuthData,
        chat: ProcessorPool, tx: ConnectionSender)
        -> AuthCodec
    {
        AuthCodec {
            state: AuthState::Init(path, req),
            chat: chat,
            conn_id: cid,
            sender: tx,
        }
    }
}

impl CallCodec {
    pub fn new(auth: Arc<String>, path: String,
        meta: Meta, args: Args, kw: Kwargs, sender: ConnectionSender)
        -> CallCodec
    {
        CallCodec {
            state: CallState::Init {
                auth: auth,
                path: path,
                meta: meta,
                args: args,
                kw: kw,
            },
            sender: sender,
        }
    }
}

fn write_json_request<S: Io, E>(mut e: http::Encoder<S>, data: &E)
    -> http::EncoderDone<S>
    where E: Encodable,
{
    use std::io::Write;
    e.add_header("Content-Type", "application/json").unwrap();
    e.add_chunked().unwrap();
    e.done_headers().unwrap();
    let mut buf = BufWriter::new(e);
    write!(&mut buf, "{}", as_json(data)).unwrap();
    match buf.into_inner() {
        Ok(x) => x.done(),
        Err(_) => unreachable!(),
    }
}

impl<S: Io> http::Codec<S> for AuthCodec {
    type Future = FutureResult<http::EncoderDone<S>, http::Error>;

    fn start_write(&mut self, mut e: http::Encoder<S>) -> Self::Future {
        use self::AuthState::*;
        if let Init(p, i) = mem::replace(&mut self.state, Void)
        {
            self.state = Wait;
            e.request_line("POST", &p, Version::Http11);
            ok(write_json_request(e,
                &Auth(&serialize_cid(&self.conn_id), &i)))
        } else {
            panic!("wrong state");
        }
    }
    fn headers_received(&mut self, headers: &http::Head)
        -> Result<http::RecvMode, http::Error>
    {
        use self::AuthState::*;
        if let Wait = mem::replace(&mut self.state, Void) {
            self.state = Headers(
                headers.status().unwrap_or(Status::InternalServerError));
            // TODO(tailhook) limit and streaming
            Ok(http::RecvMode::Buffered(10_485_760))
        } else {
            panic!("wrong state");
        }
    }
    fn data_received(&mut self, data: &[u8], end: bool)
        -> Result<Async<usize>, http::Error>
    {
        use self::AuthState::*;
        // TODO(tailhook) streaming
        assert!(end);
        match mem::replace(&mut self.state, Void) {
            Headers(Status::Ok) => {
                let result = from_utf8(data)
                    .map_err(|e| debug!("Invalid utf-8 in auth data: {}", e))
                .and_then(|s| Json::from_str(s)
                    .map_err(|e| debug!("Invalid json in auth data: {}", e)))
                .and_then(|j| parse_userinfo(j)
                    .map_err(|e| debug!("Bad user info in auth data: {}", e)));
                match result {
                    Ok((sess_id, userinfo)) => {
                        let userinfo = Arc::new(userinfo);
                        debug!("Auth data received {:?}: {:?}",
                            sess_id, userinfo);
                        self.chat.send(Action::Associate {
                            conn_id: self.conn_id,
                            session_id: sess_id.clone(),
                            metadata: userinfo.clone(),
                        });
                        self.sender.send(Hello(sess_id, userinfo));
                    }
                    Err(()) => {
                        debug!("Auth error");
                        self.sender.send(StopSocket(
                            AuthHttp(Status::InternalServerError)));
                    }
                };
            }
            Headers(status) => {
                if good_status(status) {
                    self.sender.send(StopSocket(AuthHttp(status)));
                } else {
                    self.sender.send(StopSocket(
                        AuthHttp(Status::InternalServerError)));
                }
            }
            _ => unreachable!(),
        }
        Ok((Async::Ready(data.len())))
    }
}

impl<S: Io> http::Codec<S> for CallCodec {
    type Future = FutureResult<http::EncoderDone<S>, http::Error>;

    fn start_write(&mut self, mut e: http::Encoder<S>) -> Self::Future {
        use self::CallState::*;
        if
            let Init { auth, path, meta, args, kw} =
            mem::replace(&mut self.state, Void)
        {
            e.request_line("POST", &path, Version::Http11);
            // TODO(tailhook) implement authrization
            e.add_header("Authorization", &*auth).unwrap();
            let done = write_json_request(e, &Call(&meta, &args, &kw));
            self.state = Wait(meta);
            ok(done)
        } else {
            panic!("wrong state");
        }
    }
    fn headers_received(&mut self, headers: &http::Head)
        -> Result<http::RecvMode, http::Error>
    {
        use self::CallState::*;
        if let Wait(meta) = mem::replace(&mut self.state, Void) {
            self.state = Headers(meta,
                headers.status().unwrap_or(Status::InternalServerError));
            // TODO(tailhook) configure limit
            Ok(http::RecvMode::Buffered(10_485_760))
        } else {
            panic!("wrong state");
        }
    }
    fn data_received(&mut self, data: &[u8], end: bool)
        -> Result<Async<usize>, http::Error>
    {
        use self::CallState::*;
        assert!(end);
        match mem::replace(&mut self.state, Void) {
            Headers(meta, Status::Ok) => {
                match parse_response(data) {
                    Ok(x) => {
                        self.sender.send(ConnectionMessage::Result(meta, x));
                    }
                    Err(e) => {
                        self.sender.send(ConnectionMessage::Error(meta, e));
                    }
                }
            }
            Headers(meta, status) => {
                self.sender.send(ConnectionMessage::Error(meta,
                    // TODO(tailhook) should we put body here?
                    MessageError::HttpError(status, None)));
            }
            _ => unreachable!(),
        }
        Ok((Async::Ready(data.len())))
    }
}

fn parse_response(data: &[u8]) -> Result<Json, MessageError> {
    Ok(Json::from_str(from_utf8(data)?)?)
}
<commit_msg>Elimitate race condition in chat auth<commit_after>use std::str::from_utf8;
use std::sync::Arc;
use std::io::BufWriter;
use std::mem;

use futures::Async;
use futures::future::{FutureResult, ok};
use minihttp::{Status, Version};
use minihttp::client as http;
use tokio_core::io::Io;
use rustc_serialize::Encodable;
use rustc_serialize::json::{as_json, Json};

use proxy::{Response};
use chat::{Cid, ConnectionSender, ConnectionMessage};
use chat::cid::{serialize_cid};
use chat::error::MessageError;
use chat::message::{AuthData, Auth, Call, Meta, Args, Kwargs};
use chat::processor::{ProcessorPool, Action};
use chat::authorize::{parse_userinfo, good_status};
use chat::ConnectionMessage::{Hello, StopSocket};
use chat::CloseReason::{AuthHttp};

enum AuthState {
    Init(String, AuthData),
    Wait,
    Headers(Status),
    Done(Response),
    Void,
}

enum CallState {
    Init {
        auth: Arc<String>,
        path: String, meta: Meta,
        args: Args, kw: Kwargs
    },
    Wait(Meta),
    Headers(Meta, Status),
    Void,
}

pub struct AuthCodec {
    state: AuthState,
    chat: ProcessorPool,
    conn_id: Cid,
    sender: ConnectionSender,
}

pub struct CallCodec {
    state: CallState,
    sender: ConnectionSender,
}

impl AuthCodec {
    pub fn new(path: String, cid: Cid, req: AuthData,
        chat: ProcessorPool, tx: ConnectionSender)
        -> AuthCodec
    {
        AuthCodec {
            state: AuthState::Init(path, req),
            chat: chat,
            conn_id: cid,
            sender: tx,
        }
    }
}

impl CallCodec {
    pub fn new(auth: Arc<String>, path: String,
        meta: Meta, args: Args, kw: Kwargs, sender: ConnectionSender)
        -> CallCodec
    {
        CallCodec {
            state: CallState::Init {
                auth: auth,
                path: path,
                meta: meta,
                args: args,
                kw: kw,
            },
            sender: sender,
        }
    }
}

fn write_json_request<S: Io, E>(mut e: http::Encoder<S>, data: &E)
    -> http::EncoderDone<S>
    where E: Encodable,
{
    use std::io::Write;
    e.add_header("Content-Type", "application/json").unwrap();
    e.add_chunked().unwrap();
    e.done_headers().unwrap();
    let mut buf = BufWriter::new(e);
    write!(&mut buf, "{}", as_json(data)).unwrap();
    match buf.into_inner() {
        Ok(x) => x.done(),
        Err(_) => unreachable!(),
    }
}

impl<S: Io> http::Codec<S> for AuthCodec {
    type Future = FutureResult<http::EncoderDone<S>, http::Error>;

    fn start_write(&mut self, mut e: http::Encoder<S>) -> Self::Future {
        use self::AuthState::*;
        if let Init(p, i) = mem::replace(&mut self.state, Void)
        {
            self.state = Wait;
            e.request_line("POST", &p, Version::Http11);
            ok(write_json_request(e,
                &Auth(&serialize_cid(&self.conn_id), &i)))
        } else {
            panic!("wrong state");
        }
    }
    fn headers_received(&mut self, headers: &http::Head)
        -> Result<http::RecvMode, http::Error>
    {
        use self::AuthState::*;
        if let Wait = mem::replace(&mut self.state, Void) {
            self.state = Headers(
                headers.status().unwrap_or(Status::InternalServerError));
            // TODO(tailhook) limit and streaming
            Ok(http::RecvMode::Buffered(10_485_760))
        } else {
            panic!("wrong state");
        }
    }
    fn data_received(&mut self, data: &[u8], end: bool)
        -> Result<Async<usize>, http::Error>
    {
        use self::AuthState::*;
        // TODO(tailhook) streaming
        assert!(end);
        match mem::replace(&mut self.state, Void) {
            Headers(Status::Ok) => {
                let result = from_utf8(data)
                    .map_err(|e| debug!("Invalid utf-8 in auth data: {}", e))
                .and_then(|s| Json::from_str(s)
                    .map_err(|e| debug!("Invalid json in auth data: {}", e)))
                .and_then(|j| parse_userinfo(j)
                    .map_err(|e| debug!("Bad user info in auth data: {}", e)));
                match result {
                    Ok((sess_id, userinfo)) => {
                        let userinfo = Arc::new(userinfo);
                        debug!("Auth data received {:?}: {:?}",
                            sess_id, userinfo);
                        self.sender.send(Hello(sess_id.clone(),
                                               userinfo.clone()));
                        self.chat.send(Action::Associate {
                            conn_id: self.conn_id,
                            session_id: sess_id,
                            metadata: userinfo,
                        });
                    }
                    Err(()) => {
                        debug!("Auth error");
                        self.sender.send(StopSocket(
                            AuthHttp(Status::InternalServerError)));
                    }
                };
            }
            Headers(status) => {
                if good_status(status) {
                    self.sender.send(StopSocket(AuthHttp(status)));
                } else {
                    self.sender.send(StopSocket(
                        AuthHttp(Status::InternalServerError)));
                }
            }
            _ => unreachable!(),
        }
        Ok((Async::Ready(data.len())))
    }
}

impl<S: Io> http::Codec<S> for CallCodec {
    type Future = FutureResult<http::EncoderDone<S>, http::Error>;

    fn start_write(&mut self, mut e: http::Encoder<S>) -> Self::Future {
        use self::CallState::*;
        if
            let Init { auth, path, meta, args, kw} =
            mem::replace(&mut self.state, Void)
        {
            e.request_line("POST", &path, Version::Http11);
            // TODO(tailhook) implement authrization
            e.add_header("Authorization", &*auth).unwrap();
            let done = write_json_request(e, &Call(&meta, &args, &kw));
            self.state = Wait(meta);
            ok(done)
        } else {
            panic!("wrong state");
        }
    }
    fn headers_received(&mut self, headers: &http::Head)
        -> Result<http::RecvMode, http::Error>
    {
        use self::CallState::*;
        if let Wait(meta) = mem::replace(&mut self.state, Void) {
            self.state = Headers(meta,
                headers.status().unwrap_or(Status::InternalServerError));
            // TODO(tailhook) configure limit
            Ok(http::RecvMode::Buffered(10_485_760))
        } else {
            panic!("wrong state");
        }
    }
    fn data_received(&mut self, data: &[u8], end: bool)
        -> Result<Async<usize>, http::Error>
    {
        use self::CallState::*;
        assert!(end);
        match mem::replace(&mut self.state, Void) {
            Headers(meta, Status::Ok) => {
                match parse_response(data) {
                    Ok(x) => {
                        self.sender.send(ConnectionMessage::Result(meta, x));
                    }
                    Err(e) => {
                        self.sender.send(ConnectionMessage::Error(meta, e));
                    }
                }
            }
            Headers(meta, status) => {
                self.sender.send(ConnectionMessage::Error(meta,
                    // TODO(tailhook) should we put body here?
                    MessageError::HttpError(status, None)));
            }
            _ => unreachable!(),
        }
        Ok((Async::Ready(data.len())))
    }
}

fn parse_response(data: &[u8]) -> Result<Json, MessageError> {
    Ok(Json::from_str(from_utf8(data)?)?)
}
<|endoftext|>


<commit_before>//! Compile time optimized maps
#[crate_id="github.com/sfackler/rust-phf/phf"];
#[doc(html_root_url="http://www.rust-ci.org/sfackler/rust-phf/doc")];
#[crate_type="rlib"];
#[crate_type="dylib"];
#[warn(missing_doc)];

use std::iter;
use std::vec;

/// An immutable map constructed at compile time.
///
/// `PhfMap`s may be created with the `phf_map` macro:
///
/// ```rust
/// static my_map: PhfMap<int> = phf_map!(
///    "hello" => 10,
///    "world" => 11,
/// );
/// ```
///
/// # Note
///
/// The fields of this struct are public so that they may be initialized by the
/// `phf_map` macro. They are subject to change at any time and should never
/// be accessed directly.
pub struct PhfMap<T> {
    #[doc(hidden)]
    len: uint,
    #[doc(hidden)]
    k1: u64,
    #[doc(hidden)]
    k2: u64,
    #[doc(hidden)]
    disps: &'static [(uint, uint)],
    #[doc(hidden)]
    entries: &'static [Option<(&'static str, T)>],
}

static LOG_MAX_SIZE: uint = 21;

#[doc(hidden)]
pub static MAX_SIZE: uint = 1 << LOG_MAX_SIZE;

#[doc(hidden)]
#[inline]
pub fn hash(s: &str, k1: u64, k2: u64) -> (uint, uint, uint) {
    let hash = s.hash_keyed(k1, k2);
    let mask = (MAX_SIZE - 1) as u64;

    ((hash & mask) as uint,
     ((hash >> LOG_MAX_SIZE) & mask) as uint,
     ((hash >> (2 * LOG_MAX_SIZE)) & mask) as uint)
}

#[doc(hidden)]
#[inline]
pub fn displace(f1: uint, f2: uint, d1: uint, d2: uint) -> uint {
    d2 + f1 * d1 + f2
}

impl<T> Container for PhfMap<T> {
    #[inline]
    fn len(&self) -> uint {
        self.len
    }
}

impl<'a, T> Map<&'a str, T> for PhfMap<T> {
    #[inline]
    fn find<'a>(&'a self, key: & &str) -> Option<&'a T> {
        let (g, f1, f2) = hash(*key, self.k1, self.k2);
        let (d1, d2) = self.disps[g % self.disps.len()];
        match self.entries[displace(f1, f2, d1, d2) % self.entries.len()] {
            Some((s, ref value)) if s == *key => Some(value),
            _ => None
        }
    }
}

impl<T> PhfMap<T> {
    /// Returns an iterator over the key/value pairs in the map.
    ///
    /// Entries are retuned in an arbitrary order.
    #[inline]
    pub fn entries<'a>(&'a self) -> PhfMapEntries<'a, T> {
        PhfMapEntries {
            iter: self.entries.iter().filter_map(|e| {
                match *e {
                    Some((key, ref value)) => Some((key, value)),
                    None => None
                }
            })
        }
    }

    /// Returns an iterator over the keys in the map.
    ///
    /// Keys are returned in an arbitrary order.
    #[inline]
    pub fn keys<'a>(&'a self) -> PhfMapKeys<'a, T> {
        PhfMapKeys { iter: self.entries() }
    }

    /// Returns an iterator over the values in the map.
    ///
    /// Values are returned in an arbitrary order.
    #[inline]
    pub fn values<'a>(&'a self) -> PhfMapValues<'a, T> {
        PhfMapValues { iter: self.entries() }
    }
}

/// An iterator over the key/value pairs in a `PhfMap`.
pub struct PhfMapEntries<'a, T> {
    priv iter: iter::FilterMap<'a,
                               &'a Option<(&'static str, T)>,
                               (&'static str, &'a T),
                               vec::Items<'a, Option<(&'static str, T)>>>,
}

impl<'a, T> Iterator<(&'static str, &'a T)> for PhfMapEntries<'a, T> {
    #[inline]
    fn next(&mut self) -> Option<(&'static str, &'a T)> {
        self.iter.next()
    }

    #[inline]
    fn size_hint(&self) -> (uint, Option<uint>) {
        self.iter.size_hint()
    }
}

/// An iterator over the keys in a `PhfMap`.
pub struct PhfMapKeys<'a, T> {
    priv iter: PhfMapEntries<'a, T>,
}

impl<'a, T> Iterator<&'static str> for PhfMapKeys<'a, T> {
    #[inline]
    fn next(&mut self) -> Option<&'static str> {
        self.iter.next().map(|(key, _)| key)
    }

    #[inline]
    fn size_hint(&self) -> (uint, Option<uint>) {
        self.iter.size_hint()
    }
}

/// An iterator over the values in a `PhfMap`.
pub struct PhfMapValues<'a, T> {
    priv iter: PhfMapEntries<'a, T>,
}

impl<'a, T> Iterator<&'a T> for PhfMapValues<'a, T> {
    #[inline]
    fn next(&mut self) -> Option<&'a T> {
        self.iter.next().map(|(_, value)| value)
    }

    #[inline]
    fn size_hint(&self) -> (uint, Option<uint>) {
        self.iter.size_hint()
    }
}
<commit_msg>Remove #[inline] from PhfMap API<commit_after>//! Compile time optimized maps
#[crate_id="github.com/sfackler/rust-phf/phf"];
#[doc(html_root_url="http://www.rust-ci.org/sfackler/rust-phf/doc")];
#[crate_type="rlib"];
#[crate_type="dylib"];
#[warn(missing_doc)];

use std::iter;
use std::vec;

/// An immutable map constructed at compile time.
///
/// `PhfMap`s may be created with the `phf_map` macro:
///
/// ```rust
/// static my_map: PhfMap<int> = phf_map!(
///    "hello" => 10,
///    "world" => 11,
/// );
/// ```
///
/// # Note
///
/// The fields of this struct are public so that they may be initialized by the
/// `phf_map` macro. They are subject to change at any time and should never
/// be accessed directly.
pub struct PhfMap<T> {
    #[doc(hidden)]
    len: uint,
    #[doc(hidden)]
    k1: u64,
    #[doc(hidden)]
    k2: u64,
    #[doc(hidden)]
    disps: &'static [(uint, uint)],
    #[doc(hidden)]
    entries: &'static [Option<(&'static str, T)>],
}

static LOG_MAX_SIZE: uint = 21;

#[doc(hidden)]
pub static MAX_SIZE: uint = 1 << LOG_MAX_SIZE;

#[doc(hidden)]
#[inline]
pub fn hash(s: &str, k1: u64, k2: u64) -> (uint, uint, uint) {
    let hash = s.hash_keyed(k1, k2);
    let mask = (MAX_SIZE - 1) as u64;

    ((hash & mask) as uint,
     ((hash >> LOG_MAX_SIZE) & mask) as uint,
     ((hash >> (2 * LOG_MAX_SIZE)) & mask) as uint)
}

#[doc(hidden)]
#[inline]
pub fn displace(f1: uint, f2: uint, d1: uint, d2: uint) -> uint {
    d2 + f1 * d1 + f2
}

impl<T> Container for PhfMap<T> {
    fn len(&self) -> uint {
        self.len
    }
}

impl<'a, T> Map<&'a str, T> for PhfMap<T> {
    fn find<'a>(&'a self, key: & &str) -> Option<&'a T> {
        let (g, f1, f2) = hash(*key, self.k1, self.k2);
        let (d1, d2) = self.disps[g % self.disps.len()];
        match self.entries[displace(f1, f2, d1, d2) % self.entries.len()] {
            Some((s, ref value)) if s == *key => Some(value),
            _ => None
        }
    }
}

impl<T> PhfMap<T> {
    /// Returns an iterator over the key/value pairs in the map.
    ///
    /// Entries are retuned in an arbitrary order.
    #[inline]
    pub fn entries<'a>(&'a self) -> PhfMapEntries<'a, T> {
        PhfMapEntries {
            iter: self.entries.iter().filter_map(|e| {
                match *e {
                    Some((key, ref value)) => Some((key, value)),
                    None => None
                }
            })
        }
    }

    /// Returns an iterator over the keys in the map.
    ///
    /// Keys are returned in an arbitrary order.
    #[inline]
    pub fn keys<'a>(&'a self) -> PhfMapKeys<'a, T> {
        PhfMapKeys { iter: self.entries() }
    }

    /// Returns an iterator over the values in the map.
    ///
    /// Values are returned in an arbitrary order.
    #[inline]
    pub fn values<'a>(&'a self) -> PhfMapValues<'a, T> {
        PhfMapValues { iter: self.entries() }
    }
}

/// An iterator over the key/value pairs in a `PhfMap`.
pub struct PhfMapEntries<'a, T> {
    priv iter: iter::FilterMap<'a,
                               &'a Option<(&'static str, T)>,
                               (&'static str, &'a T),
                               vec::Items<'a, Option<(&'static str, T)>>>,
}

impl<'a, T> Iterator<(&'static str, &'a T)> for PhfMapEntries<'a, T> {
    fn next(&mut self) -> Option<(&'static str, &'a T)> {
        self.iter.next()
    }

    fn size_hint(&self) -> (uint, Option<uint>) {
        self.iter.size_hint()
    }
}

/// An iterator over the keys in a `PhfMap`.
pub struct PhfMapKeys<'a, T> {
    priv iter: PhfMapEntries<'a, T>,
}

impl<'a, T> Iterator<&'static str> for PhfMapKeys<'a, T> {
    fn next(&mut self) -> Option<&'static str> {
        self.iter.next().map(|(key, _)| key)
    }

    fn size_hint(&self) -> (uint, Option<uint>) {
        self.iter.size_hint()
    }
}

/// An iterator over the values in a `PhfMap`.
pub struct PhfMapValues<'a, T> {
    priv iter: PhfMapEntries<'a, T>,
}

impl<'a, T> Iterator<&'a T> for PhfMapValues<'a, T> {
    fn next(&mut self) -> Option<&'a T> {
        self.iter.next().map(|(_, value)| value)
    }

    fn size_hint(&self) -> (uint, Option<uint>) {
        self.iter.size_hint()
    }
}
<|endoftext|>


<commit_before>use std::rc::Rc;
use std::cell::RefCell;
use std::net::SocketAddr;
use std::collections::HashMap;
use rand::random;
use mio::net::TcpStream;

use network::retry::ExponentialBackoffPolicy;
use network::{AppId,Backend,ConnectionError};

pub struct BackendMap {
  pub instances:    HashMap<AppId, BackendList>,
  pub max_failures: usize,
}

impl BackendMap {
  pub fn new() -> BackendMap {
    BackendMap {
      instances:    HashMap::new(),
      max_failures: 3,
    }
  }

  pub fn add_instance(&mut self, app_id: &str, instance_address: &SocketAddr) {
    self.instances.entry(app_id.to_string()).or_insert(BackendList::new()).add_instance(instance_address);
  }

  pub fn remove_instance(&mut self, app_id: &str, instance_address: &SocketAddr) {
    if let Some(instances) = self.instances.get_mut(app_id) {
      instances.remove_instance(instance_address);
    } else {
      error!("Instance was already removed");
    }
  }

  pub fn close_backend_connection(&mut self, app_id: &str, addr: &SocketAddr) {
    if let Some(app_instances) = self.instances.get_mut(app_id) {
      if let Some(ref mut backend) = app_instances.find_instance(addr) {
        (*backend.borrow_mut()).dec_connections();
      }
    }
  }

  pub fn backend_from_app_id(&mut self, app_id: &str) -> Result<(Rc<RefCell<Backend<ExponentialBackoffPolicy>>>,TcpStream),ConnectionError> {
    if let Some(ref mut app_instances) = self.instances.get_mut(app_id) {
      if app_instances.instances.len() == 0 {
        return Err(ConnectionError::NoBackendAvailable);
      }

      for _ in 0..self.max_failures {
        if let Some(ref mut b) = app_instances.next_available_instance() {
          let ref mut backend = *b.borrow_mut();
          info!("Connecting {} -> {:?}", app_id, (backend.address, backend.active_connections, backend.failures));
          let conn = backend.try_connect();
          if backend.failures >= MAX_FAILURES_PER_BACKEND {
            error!("backend {:?} connections failed {} times, disabling it", (backend.address, backend.active_connections), backend.failures);
          }

          return conn.map(|c| (b.clone(), c));
        } else {
          error!("no more available backends for app {}", app_id);
          return Err(ConnectionError::NoBackendAvailable);
        }
      }
      Err(ConnectionError::NoBackendAvailable)
    } else {
      Err(ConnectionError::NoBackendAvailable)
    }
  }

  pub fn backend_from_sticky_session(&mut self, app_id: &str, sticky_session: u32) -> Result<(Rc<RefCell<Backend<ExponentialBackoffPolicy>>>,TcpStream),ConnectionError> {
    let sticky_conn: Option<Result<(Rc<RefCell<Backend<ExponentialBackoffPolicy>>>,TcpStream),ConnectionError>> = self.instances
        .get_mut(app_id)
        .and_then(|app_instances| app_instances.find_sticky(sticky_session))
        .map(|b| {
          let ref mut backend = *b.borrow_mut();
          let conn = backend.try_connect();
          info!("Connecting {} -> {:?} using session {}", app_id, (backend.address, backend.active_connections, backend.failures), sticky_session);
          if backend.failures >= MAX_FAILURES_PER_BACKEND {
            error!("backend {:?} connections failed {} times, disabling it", (backend.address, backend.active_connections), backend.failures);
          }

          conn.map(|c| (b.clone(), c))
        });

    if let Some(res) = sticky_conn {
      return res;
    } else {
      debug!("Couldn't find a backend corresponding to sticky_session {} for app {}", sticky_session, app_id);
      return self.backend_from_app_id(app_id);
    }
  }
}

const MAX_FAILURES_PER_BACKEND: usize = 10;

pub struct BackendList {
  pub instances: Vec<Rc<RefCell<Backend<ExponentialBackoffPolicy>>>>,
  pub next_id:   u32,
}

impl BackendList {
  pub fn new() -> BackendList {
    BackendList {
      instances: Vec::new(),
      next_id:   0,
    }
  }

  pub fn add_instance(&mut self, instance_address: &SocketAddr) {
    if self.instances.iter().find(|b| &(*b.borrow()).address == instance_address).is_none() {
      let backend = Rc::new(RefCell::new(Backend::<ExponentialBackoffPolicy>::new(*instance_address, self.next_id)));
      self.instances.push(backend);
      self.next_id += 1;
    }
  }

  pub fn remove_instance(&mut self, instance_address: &SocketAddr) {
    self.instances.retain(|backend| &(*backend.borrow()).address != instance_address);
  }

  pub fn find_instance(&mut self, instance_address: &SocketAddr) -> Option<&mut Rc<RefCell<Backend<ExponentialBackoffPolicy>>>> {
    self.instances.iter_mut().find(|backend| &(*backend.borrow()).address == instance_address)
  }

  pub fn find_sticky(&mut self, sticky_session: u32) -> Option<&mut Rc<RefCell<Backend<ExponentialBackoffPolicy>>>> {
    self.instances.iter_mut()
      .find(|b| b.borrow().id == sticky_session )
      .and_then(|b| {
        if b.borrow().can_open() {
          Some(b)
        } else {
          None
        }
      })
  }

  pub fn available_instances(&mut self) -> Vec<&mut Rc<RefCell<Backend<ExponentialBackoffPolicy>>>> {
    self.instances.iter_mut()
      .filter(|backend| (*backend.borrow()).can_open())
      .collect()
  }

  pub fn next_available_instance(&mut self) -> Option<&mut Rc<RefCell<Backend<ExponentialBackoffPolicy>>>> {
    let mut instances:Vec<&mut Rc<RefCell<Backend<ExponentialBackoffPolicy>>>> = self.available_instances();
    if instances.is_empty() {
      return None;
    }

    let rnd = random::<usize>();
    let idx = rnd % instances.len();

    Some(instances.remove(idx))
  }
}
<commit_msg>fix indentation<commit_after>use std::rc::Rc;
use std::cell::RefCell;
use std::net::SocketAddr;
use std::collections::HashMap;
use rand::random;
use mio::net::TcpStream;

use network::retry::ExponentialBackoffPolicy;
use network::{AppId,Backend,ConnectionError};

pub struct BackendMap {
  pub instances:    HashMap<AppId, BackendList>,
  pub max_failures: usize,
}

impl BackendMap {
  pub fn new() -> BackendMap {
    BackendMap {
      instances:    HashMap::new(),
      max_failures: 3,
    }
  }

  pub fn add_instance(&mut self, app_id: &str, instance_address: &SocketAddr) {
    self.instances.entry(app_id.to_string()).or_insert(BackendList::new()).add_instance(instance_address);
  }

  pub fn remove_instance(&mut self, app_id: &str, instance_address: &SocketAddr) {
    if let Some(instances) = self.instances.get_mut(app_id) {
      instances.remove_instance(instance_address);
    } else {
      error!("Instance was already removed");
    }
  }

  pub fn close_backend_connection(&mut self, app_id: &str, addr: &SocketAddr) {
    if let Some(app_instances) = self.instances.get_mut(app_id) {
      if let Some(ref mut backend) = app_instances.find_instance(addr) {
        (*backend.borrow_mut()).dec_connections();
      }
    }
  }

  pub fn backend_from_app_id(&mut self, app_id: &str) -> Result<(Rc<RefCell<Backend<ExponentialBackoffPolicy>>>,TcpStream),ConnectionError> {
    if let Some(ref mut app_instances) = self.instances.get_mut(app_id) {
      if app_instances.instances.len() == 0 {
        return Err(ConnectionError::NoBackendAvailable);
      }

      for _ in 0..self.max_failures {
        if let Some(ref mut b) = app_instances.next_available_instance() {
          let ref mut backend = *b.borrow_mut();
          info!("Connecting {} -> {:?}", app_id, (backend.address, backend.active_connections, backend.failures));
          let conn = backend.try_connect();
          if backend.failures >= MAX_FAILURES_PER_BACKEND {
            error!("backend {:?} connections failed {} times, disabling it", (backend.address, backend.active_connections), backend.failures);
          }

          return conn.map(|c| (b.clone(), c));
        } else {
          error!("no more available backends for app {}", app_id);
          return Err(ConnectionError::NoBackendAvailable);
        }
      }
      Err(ConnectionError::NoBackendAvailable)
    } else {
      Err(ConnectionError::NoBackendAvailable)
    }
  }

  pub fn backend_from_sticky_session(&mut self, app_id: &str, sticky_session: u32) -> Result<(Rc<RefCell<Backend<ExponentialBackoffPolicy>>>,TcpStream),ConnectionError> {
    let sticky_conn: Option<Result<(Rc<RefCell<Backend<ExponentialBackoffPolicy>>>,TcpStream),ConnectionError>> = self.instances
      .get_mut(app_id)
      .and_then(|app_instances| app_instances.find_sticky(sticky_session))
      .map(|b| {
        let ref mut backend = *b.borrow_mut();
        let conn = backend.try_connect();
        info!("Connecting {} -> {:?} using session {}", app_id, (backend.address, backend.active_connections, backend.failures), sticky_session);
        if backend.failures >= MAX_FAILURES_PER_BACKEND {
          error!("backend {:?} connections failed {} times, disabling it", (backend.address, backend.active_connections), backend.failures);
        }

        conn.map(|c| (b.clone(), c))
      });

    if let Some(res) = sticky_conn {
      return res;
    } else {
      debug!("Couldn't find a backend corresponding to sticky_session {} for app {}", sticky_session, app_id);
      return self.backend_from_app_id(app_id);
    }
  }
}

const MAX_FAILURES_PER_BACKEND: usize = 10;

pub struct BackendList {
  pub instances: Vec<Rc<RefCell<Backend<ExponentialBackoffPolicy>>>>,
  pub next_id:   u32,
}

impl BackendList {
  pub fn new() -> BackendList {
    BackendList {
      instances: Vec::new(),
      next_id:   0,
    }
  }

  pub fn add_instance(&mut self, instance_address: &SocketAddr) {
    if self.instances.iter().find(|b| &(*b.borrow()).address == instance_address).is_none() {
      let backend = Rc::new(RefCell::new(Backend::<ExponentialBackoffPolicy>::new(*instance_address, self.next_id)));
      self.instances.push(backend);
      self.next_id += 1;
    }
  }

  pub fn remove_instance(&mut self, instance_address: &SocketAddr) {
    self.instances.retain(|backend| &(*backend.borrow()).address != instance_address);
  }

  pub fn find_instance(&mut self, instance_address: &SocketAddr) -> Option<&mut Rc<RefCell<Backend<ExponentialBackoffPolicy>>>> {
    self.instances.iter_mut().find(|backend| &(*backend.borrow()).address == instance_address)
  }

  pub fn find_sticky(&mut self, sticky_session: u32) -> Option<&mut Rc<RefCell<Backend<ExponentialBackoffPolicy>>>> {
    self.instances.iter_mut()
      .find(|b| b.borrow().id == sticky_session )
      .and_then(|b| {
        if b.borrow().can_open() {
          Some(b)
        } else {
          None
        }
      })
  }

  pub fn available_instances(&mut self) -> Vec<&mut Rc<RefCell<Backend<ExponentialBackoffPolicy>>>> {
    self.instances.iter_mut()
      .filter(|backend| (*backend.borrow()).can_open())
      .collect()
  }

  pub fn next_available_instance(&mut self) -> Option<&mut Rc<RefCell<Backend<ExponentialBackoffPolicy>>>> {
    let mut instances:Vec<&mut Rc<RefCell<Backend<ExponentialBackoffPolicy>>>> = self.available_instances();
    if instances.is_empty() {
      return None;
    }

    let rnd = random::<usize>();
    let idx = rnd % instances.len();

    Some(instances.remove(idx))
  }
}
<|endoftext|>


<commit_before>use std::path::PathBuf;

use std::fmt::{Display, Debug, Formatter};
use std::fmt::Error as FmtError;
use std::result::Result as RResult;
use std::path::Components;

use libimagerror::into::IntoError;

use error::StoreErrorKind as SEK;
use store::Result;
use store::Store;

/// The Index into the Store
#[derive(Debug, Clone, PartialEq, Hash, Eq, PartialOrd, Ord)]
pub struct StoreId {
    base: Option<PathBuf>,
    id:   PathBuf,
}

impl StoreId {

    pub fn new(base: Option<PathBuf>, id: PathBuf) -> Result<StoreId> {
        StoreId::new_baseless(id).map(|mut sid| { sid.base = base; sid })
    }

    pub fn new_baseless(id: PathBuf) -> Result<StoreId> {
        if id.is_absolute() {
            Err(SEK::StoreIdLocalPartAbsoluteError.into_error())
        } else {
            Ok(StoreId {
                base: None,
                id: id
            })
        }
    }

    pub fn without_base(mut self) -> StoreId {
        self.base = None;
        self
    }

    pub fn with_base(mut self, base: PathBuf) -> Self {
        self.base = Some(base);
        self
    }

    pub fn storified(self, store: &Store) -> StoreId {
        StoreId {
            base: Some(store.path().clone()),
            id: self.id
        }
    }

    pub fn exists(&self) -> bool {
        let pb : PathBuf = self.clone().into();
        pb.exists()
    }

    pub fn is_file(&self) -> bool {
        true
    }

    pub fn is_dir(&self) -> bool {
        false
    }

    pub fn to_str(&self) -> Result<String> {
        if self.base.is_some() {
            let mut base = self.base.as_ref().cloned().unwrap();
            base.push(self.id.clone());
            base
        } else {
            self.id.clone()
        }
        .to_str()
        .map(String::from)
        .ok_or(SEK::StoreIdHandlingError.into_error())
    }

    /// Returns the components of the `id` part of the StoreId object.
    ///
    /// Can be used to check whether a StoreId points to an entry in a specific collection of
    /// StoreIds.
    pub fn components(&self) -> Components {
        self.id.components()
    }

    /// Check whether a StoreId points to an entry in a specific collection.
    ///
    /// A "collection" here is simply a directory. So `foo/bar/baz` is an entry which is in
    /// collection ["foo", "bar"].
    ///
    /// # Warning
    ///
    /// The collection specification _has_ to start with the module name. Otherwise this function
    /// may return false negatives.
    ///
    pub fn is_in_collection(&self, colls: &[&str]) -> bool {
        use std::path::Component;

        self.id
            .components()
            .zip(colls)
            .all(|(component, pred_coll)| match component {
                Component::Normal(ref s) => s.to_str().map(|ref s| s == pred_coll).unwrap_or(false),
                _ => false
            }) && colls.last().map(|last| !self.id.ends_with(last)).unwrap_or(false)
    }

}

impl Into<PathBuf> for StoreId {

    fn into(self) -> PathBuf {
        let mut base = self.base.unwrap_or(PathBuf::from("/"));
        base.push(self.id);
        base
    }

}

impl Display for StoreId {

    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        match self.id.to_str() {
            Some(s) => write!(fmt, "{}", s),
            None    => write!(fmt, "{}", self.id.to_string_lossy()),
        }
    }

}

/// This Trait allows you to convert various representations to a single one
/// suitable for usage in the Store
pub trait IntoStoreId {
    fn into_storeid(self) -> Result<StoreId>;
}

impl IntoStoreId for StoreId {
    fn into_storeid(self) -> Result<StoreId> {
        Ok(self)
    }
}

impl IntoStoreId for PathBuf {
    fn into_storeid(self) -> Result<StoreId> {
        StoreId::new_baseless(self)
    }
}

#[macro_export]
macro_rules! module_entry_path_mod {
    ($name:expr) => (
        #[deny(missing_docs,
                missing_copy_implementations,
                trivial_casts, trivial_numeric_casts,
                unsafe_code,
                unstable_features,
                unused_import_braces, unused_qualifications,
                unused_imports)]
        /// A helper module to create valid module entry paths
        pub mod module_path {
            use std::convert::AsRef;
            use std::path::Path;
            use std::path::PathBuf;

            use $crate::storeid::StoreId;
            use $crate::store::Result;

            /// A Struct giving you the ability to choose store entries assigned
            /// to it.
            ///
            /// It is created through a call to `new`.
            pub struct ModuleEntryPath(PathBuf);

            impl ModuleEntryPath {
                /// Path has to be a valid UTF-8 string or this will panic!
                pub fn new<P: AsRef<Path>>(pa: P) -> ModuleEntryPath {
                    let mut path = PathBuf::new();
                    path.push(format!("{}", $name));
                    path.push(pa.as_ref().clone());
                    let name = pa.as_ref().file_name().unwrap()
                        .to_str().unwrap();
                    path.set_file_name(name);
                    ModuleEntryPath(path)
                }
            }

            impl $crate::storeid::IntoStoreId for ModuleEntryPath {
                fn into_storeid(self) -> Result<$crate::storeid::StoreId> {
                    StoreId::new(None, self.0)
                }
            }
        }
    )
}

pub struct StoreIdIterator {
    iter: Box<Iterator<Item = StoreId>>,
}

impl Debug for StoreIdIterator {

    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        write!(fmt, "StoreIdIterator")
    }

}

impl StoreIdIterator {

    pub fn new(iter: Box<Iterator<Item = StoreId>>) -> StoreIdIterator {
        StoreIdIterator {
            iter: iter,
        }
    }

}

impl Iterator for StoreIdIterator {
    type Item = StoreId;

    fn next(&mut self) -> Option<StoreId> {
        self.iter.next()
    }

}

#[cfg(test)]
mod test {

    use storeid::IntoStoreId;

    module_entry_path_mod!("test");

    #[test]
    fn correct_path() {
        let p = module_path::ModuleEntryPath::new("test");

        assert_eq!(p.into_storeid().unwrap().to_str().unwrap(), "test/test");
    }

    #[test]
    fn storeid_in_collection() {
        let p = module_path::ModuleEntryPath::new("1/2/3/4/5/6/7/8/9/0").into_storeid().unwrap();

        assert!(p.is_in_collection(&["test", "1"]));
        assert!(p.is_in_collection(&["test", "1", "2"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3", "4"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3", "4", "5"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3", "4", "5", "6"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3", "4", "5", "6", "7"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3", "4", "5", "6", "7", "8"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3", "4", "5", "6", "7", "8", "9"]));

        // "0" is the filename, not a collection
        assert!(!p.is_in_collection(&["test", "1", "2", "3", "4", "5", "6", "7", "8", "9", "0"]));

        assert!(!p.is_in_collection(&["test", "0", "2", "3", "4", "5", "6", "7", "8", "9", "0"]));
        assert!(!p.is_in_collection(&["test", "1", "2", "3", "4", "5", "6", "8"]));
        assert!(!p.is_in_collection(&["test", "1", "2", "3", "leet", "5", "6", "7"]));
    }

}
<commit_msg>Change the interface of StoreId::is_in_collection()<commit_after>use std::path::PathBuf;

use std::fmt::{Display, Debug, Formatter};
use std::fmt::Error as FmtError;
use std::result::Result as RResult;
use std::path::Components;

use libimagerror::into::IntoError;

use error::StoreErrorKind as SEK;
use store::Result;
use store::Store;

/// The Index into the Store
#[derive(Debug, Clone, PartialEq, Hash, Eq, PartialOrd, Ord)]
pub struct StoreId {
    base: Option<PathBuf>,
    id:   PathBuf,
}

impl StoreId {

    pub fn new(base: Option<PathBuf>, id: PathBuf) -> Result<StoreId> {
        StoreId::new_baseless(id).map(|mut sid| { sid.base = base; sid })
    }

    pub fn new_baseless(id: PathBuf) -> Result<StoreId> {
        if id.is_absolute() {
            Err(SEK::StoreIdLocalPartAbsoluteError.into_error())
        } else {
            Ok(StoreId {
                base: None,
                id: id
            })
        }
    }

    pub fn without_base(mut self) -> StoreId {
        self.base = None;
        self
    }

    pub fn with_base(mut self, base: PathBuf) -> Self {
        self.base = Some(base);
        self
    }

    pub fn storified(self, store: &Store) -> StoreId {
        StoreId {
            base: Some(store.path().clone()),
            id: self.id
        }
    }

    pub fn exists(&self) -> bool {
        let pb : PathBuf = self.clone().into();
        pb.exists()
    }

    pub fn is_file(&self) -> bool {
        true
    }

    pub fn is_dir(&self) -> bool {
        false
    }

    pub fn to_str(&self) -> Result<String> {
        if self.base.is_some() {
            let mut base = self.base.as_ref().cloned().unwrap();
            base.push(self.id.clone());
            base
        } else {
            self.id.clone()
        }
        .to_str()
        .map(String::from)
        .ok_or(SEK::StoreIdHandlingError.into_error())
    }

    /// Returns the components of the `id` part of the StoreId object.
    ///
    /// Can be used to check whether a StoreId points to an entry in a specific collection of
    /// StoreIds.
    pub fn components(&self) -> Components {
        self.id.components()
    }

    /// Convenience function over PathBuf::starts_with().
    ///
    /// This function calls `PathBuf::starts_with()` on the _local_ part of `self`.
    pub fn is_in_collection(&self, collspec: &PathBuf) -> bool {
        self.id.starts_with(collspec)
    }

}

impl Into<PathBuf> for StoreId {

    fn into(self) -> PathBuf {
        let mut base = self.base.unwrap_or(PathBuf::from("/"));
        base.push(self.id);
        base
    }

}

impl Display for StoreId {

    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        match self.id.to_str() {
            Some(s) => write!(fmt, "{}", s),
            None    => write!(fmt, "{}", self.id.to_string_lossy()),
        }
    }

}

/// This Trait allows you to convert various representations to a single one
/// suitable for usage in the Store
pub trait IntoStoreId {
    fn into_storeid(self) -> Result<StoreId>;
}

impl IntoStoreId for StoreId {
    fn into_storeid(self) -> Result<StoreId> {
        Ok(self)
    }
}

impl IntoStoreId for PathBuf {
    fn into_storeid(self) -> Result<StoreId> {
        StoreId::new_baseless(self)
    }
}

#[macro_export]
macro_rules! module_entry_path_mod {
    ($name:expr) => (
        #[deny(missing_docs,
                missing_copy_implementations,
                trivial_casts, trivial_numeric_casts,
                unsafe_code,
                unstable_features,
                unused_import_braces, unused_qualifications,
                unused_imports)]
        /// A helper module to create valid module entry paths
        pub mod module_path {
            use std::convert::AsRef;
            use std::path::Path;
            use std::path::PathBuf;

            use $crate::storeid::StoreId;
            use $crate::store::Result;

            /// A Struct giving you the ability to choose store entries assigned
            /// to it.
            ///
            /// It is created through a call to `new`.
            pub struct ModuleEntryPath(PathBuf);

            impl ModuleEntryPath {
                /// Path has to be a valid UTF-8 string or this will panic!
                pub fn new<P: AsRef<Path>>(pa: P) -> ModuleEntryPath {
                    let mut path = PathBuf::new();
                    path.push(format!("{}", $name));
                    path.push(pa.as_ref().clone());
                    let name = pa.as_ref().file_name().unwrap()
                        .to_str().unwrap();
                    path.set_file_name(name);
                    ModuleEntryPath(path)
                }
            }

            impl $crate::storeid::IntoStoreId for ModuleEntryPath {
                fn into_storeid(self) -> Result<$crate::storeid::StoreId> {
                    StoreId::new(None, self.0)
                }
            }
        }
    )
}

pub struct StoreIdIterator {
    iter: Box<Iterator<Item = StoreId>>,
}

impl Debug for StoreIdIterator {

    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        write!(fmt, "StoreIdIterator")
    }

}

impl StoreIdIterator {

    pub fn new(iter: Box<Iterator<Item = StoreId>>) -> StoreIdIterator {
        StoreIdIterator {
            iter: iter,
        }
    }

}

impl Iterator for StoreIdIterator {
    type Item = StoreId;

    fn next(&mut self) -> Option<StoreId> {
        self.iter.next()
    }

}

#[cfg(test)]
mod test {

    use storeid::IntoStoreId;

    module_entry_path_mod!("test");

    #[test]
    fn correct_path() {
        let p = module_path::ModuleEntryPath::new("test");

        assert_eq!(p.into_storeid().unwrap().to_str().unwrap(), "test/test");
    }

    #[test]
    fn storeid_in_collection() {
        let p = module_path::ModuleEntryPath::new("1/2/3/4/5/6/7/8/9/0").into_storeid().unwrap();

        assert!(p.is_in_collection(&["test", "1"]));
        assert!(p.is_in_collection(&["test", "1", "2"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3", "4"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3", "4", "5"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3", "4", "5", "6"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3", "4", "5", "6", "7"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3", "4", "5", "6", "7", "8"]));
        assert!(p.is_in_collection(&["test", "1", "2", "3", "4", "5", "6", "7", "8", "9"]));

        // "0" is the filename, not a collection
        assert!(!p.is_in_collection(&["test", "1", "2", "3", "4", "5", "6", "7", "8", "9", "0"]));

        assert!(!p.is_in_collection(&["test", "0", "2", "3", "4", "5", "6", "7", "8", "9", "0"]));
        assert!(!p.is_in_collection(&["test", "1", "2", "3", "4", "5", "6", "8"]));
        assert!(!p.is_in_collection(&["test", "1", "2", "3", "leet", "5", "6", "7"]));
    }

}
<|endoftext|>


<commit_before>use std::{
    hash::{Hash, Hasher},
    sync::Arc,
};

use mbe::MacroRules;
use ra_db::{salsa, CrateId, FileId};
use ra_prof::profile;
use ra_syntax::{
    ast::{self, AstNode},
    Parse, SyntaxNode,
};

use crate::{ast_id_map::FileAstId, db::AstDatabase};

macro_rules! impl_intern_key {
    ($name:ident) => {
        impl salsa::InternKey for $name {
            fn from_intern_id(v: salsa::InternId) -> Self {
                $name(v)
            }
            fn as_intern_id(&self) -> salsa::InternId {
                self.0
            }
        }
    };
}

/// Input to the analyzer is a set of files, where each file is identified by
/// `FileId` and contains source code. However, another source of source code in
/// Rust are macros: each macro can be thought of as producing a "temporary
/// file". To assign an id to such a file, we use the id of the macro call that
/// produced the file. So, a `HirFileId` is either a `FileId` (source code
/// written by user), or a `MacroCallId` (source code produced by macro).
///
/// What is a `MacroCallId`? Simplifying, it's a `HirFileId` of a file
/// containing the call plus the offset of the macro call in the file. Note that
/// this is a recursive definition! However, the size_of of `HirFileId` is
/// finite (because everything bottoms out at the real `FileId`) and small
/// (`MacroCallId` uses the location interner).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum HirFileId {
    FileId(FileId),
    MacroFile(MacroFile),
}

impl From<FileId> for HirFileId {
    fn from(id: FileId) -> Self {
        HirFileId::FileId(id)
    }
}

impl From<MacroFile> for HirFileId {
    fn from(id: MacroFile) -> Self {
        HirFileId::MacroFile(id)
    }
}

impl HirFileId {
    /// For macro-expansion files, returns the file original source file the
    /// expansion originated from.
    pub fn original_file(self, db: &impl AstDatabase) -> FileId {
        match self {
            HirFileId::FileId(file_id) => file_id,
            HirFileId::MacroFile(macro_file) => {
                let loc = db.lookup_intern_macro(macro_file.macro_call_id);
                loc.ast_id.file_id().original_file(db)
            }
        }
    }

    /// Get the crate which the macro lives in, if it is a macro file.
    pub fn macro_crate(self, db: &impl AstDatabase) -> Option<CrateId> {
        match self {
            HirFileId::FileId(_) => None,
            HirFileId::MacroFile(macro_file) => {
                let loc = db.lookup_intern_macro(macro_file.macro_call_id);
                Some(loc.def.krate)
            }
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MacroFile {
    macro_call_id: MacroCallId,
    macro_file_kind: MacroFileKind,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum MacroFileKind {
    Items,
    Expr,
}

/// `MacroCallId` identifies a particular macro invocation, like
/// `println!("Hello, {}", world)`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MacroCallId(salsa::InternId);
impl_intern_key!(MacroCallId);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MacroDefId {
    pub krate: CrateId,
    pub ast_id: AstId<ast::MacroCall>,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct MacroCallLoc {
    pub def: MacroDefId,
    pub ast_id: AstId<ast::MacroCall>,
}

impl MacroCallId {
    pub fn loc(self, db: &impl AstDatabase) -> MacroCallLoc {
        db.lookup_intern_macro(self)
    }

    pub fn as_file(self, kind: MacroFileKind) -> HirFileId {
        let macro_file = MacroFile { macro_call_id: self, macro_file_kind: kind };
        macro_file.into()
    }
}

impl MacroCallLoc {
    pub fn id(self, db: &impl AstDatabase) -> MacroCallId {
        db.intern_macro(self)
    }
}

/// `AstId` points to an AST node in any file.
///
/// It is stable across reparses, and can be used as salsa key/value.
// FIXME: isn't this just a `Source<FileAstId<N>>` ?
#[derive(Debug)]
pub struct AstId<N: AstNode> {
    file_id: HirFileId,
    file_ast_id: FileAstId<N>,
}

impl<N: AstNode> Clone for AstId<N> {
    fn clone(&self) -> AstId<N> {
        *self
    }
}
impl<N: AstNode> Copy for AstId<N> {}

impl<N: AstNode> PartialEq for AstId<N> {
    fn eq(&self, other: &Self) -> bool {
        (self.file_id, self.file_ast_id) == (other.file_id, other.file_ast_id)
    }
}
impl<N: AstNode> Eq for AstId<N> {}
impl<N: AstNode> Hash for AstId<N> {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        (self.file_id, self.file_ast_id).hash(hasher);
    }
}

impl<N: AstNode> AstId<N> {
    pub fn new(file_id: HirFileId, file_ast_id: FileAstId<N>) -> AstId<N> {
        AstId { file_id, file_ast_id }
    }

    pub fn file_id(&self) -> HirFileId {
        self.file_id
    }

    pub fn to_node(&self, db: &impl AstDatabase) -> N {
        let syntax_node = db.ast_id_to_node(self.file_id, self.file_ast_id.into());
        N::cast(syntax_node).unwrap()
    }
}

pub(crate) fn macro_def_query(db: &impl AstDatabase, id: MacroDefId) -> Option<Arc<MacroRules>> {
    let macro_call = id.ast_id.to_node(db);
    let arg = macro_call.token_tree()?;
    let (tt, _) = mbe::ast_to_token_tree(&arg).or_else(|| {
        log::warn!("fail on macro_def to token tree: {:#?}", arg);
        None
    })?;
    let rules = MacroRules::parse(&tt).ok().or_else(|| {
        log::warn!("fail on macro_def parse: {:#?}", tt);
        None
    })?;
    Some(Arc::new(rules))
}

pub(crate) fn macro_arg_query(db: &impl AstDatabase, id: MacroCallId) -> Option<Arc<tt::Subtree>> {
    let loc = db.lookup_intern_macro(id);
    let macro_call = loc.ast_id.to_node(db);
    let arg = macro_call.token_tree()?;
    let (tt, _) = mbe::ast_to_token_tree(&arg)?;
    Some(Arc::new(tt))
}

pub(crate) fn macro_expand_query(
    db: &impl AstDatabase,
    id: MacroCallId,
) -> Result<Arc<tt::Subtree>, String> {
    let loc = db.lookup_intern_macro(id);
    let macro_arg = db.macro_arg(id).ok_or("Fail to args in to tt::TokenTree")?;

    let macro_rules = db.macro_def(loc.def).ok_or("Fail to find macro definition")?;
    let tt = macro_rules.expand(&macro_arg).map_err(|err| format!("{:?}", err))?;
    // Set a hard limit for the expanded tt
    let count = tt.count();
    if count > 65536 {
        return Err(format!("Total tokens count exceed limit : count = {}", count));
    }
    Ok(Arc::new(tt))
}

pub(crate) fn parse_or_expand_query(
    db: &impl AstDatabase,
    file_id: HirFileId,
) -> Option<SyntaxNode> {
    match file_id {
        HirFileId::FileId(file_id) => Some(db.parse(file_id).tree().syntax().clone()),
        HirFileId::MacroFile(macro_file) => db.parse_macro(macro_file).map(|it| it.syntax_node()),
    }
}

pub(crate) fn parse_macro_query(
    db: &impl AstDatabase,
    macro_file: MacroFile,
) -> Option<Parse<SyntaxNode>> {
    let _p = profile("parse_macro_query");
    let macro_call_id = macro_file.macro_call_id;
    let tt = db
        .macro_expand(macro_call_id)
        .map_err(|err| {
            // Note:
            // The final goal we would like to make all parse_macro success,
            // such that the following log will not call anyway.
            log::warn!("fail on macro_parse: (reason: {})", err,);
        })
        .ok()?;
    match macro_file.macro_file_kind {
        MacroFileKind::Items => mbe::token_tree_to_items(&tt).ok().map(Parse::to_syntax),
        MacroFileKind::Expr => mbe::token_tree_to_expr(&tt).ok().map(Parse::to_syntax),
    }
}
<commit_msg>make file id repr private again<commit_after>use std::{
    hash::{Hash, Hasher},
    sync::Arc,
};

use mbe::MacroRules;
use ra_db::{salsa, CrateId, FileId};
use ra_prof::profile;
use ra_syntax::{
    ast::{self, AstNode},
    Parse, SyntaxNode,
};

use crate::{ast_id_map::FileAstId, db::AstDatabase};

macro_rules! impl_intern_key {
    ($name:ident) => {
        impl salsa::InternKey for $name {
            fn from_intern_id(v: salsa::InternId) -> Self {
                $name(v)
            }
            fn as_intern_id(&self) -> salsa::InternId {
                self.0
            }
        }
    };
}

/// Input to the analyzer is a set of files, where each file is identified by
/// `FileId` and contains source code. However, another source of source code in
/// Rust are macros: each macro can be thought of as producing a "temporary
/// file". To assign an id to such a file, we use the id of the macro call that
/// produced the file. So, a `HirFileId` is either a `FileId` (source code
/// written by user), or a `MacroCallId` (source code produced by macro).
///
/// What is a `MacroCallId`? Simplifying, it's a `HirFileId` of a file
/// containing the call plus the offset of the macro call in the file. Note that
/// this is a recursive definition! However, the size_of of `HirFileId` is
/// finite (because everything bottoms out at the real `FileId`) and small
/// (`MacroCallId` uses the location interner).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct HirFileId(HirFileIdRepr);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum HirFileIdRepr {
    FileId(FileId),
    MacroFile(MacroFile),
}

impl From<FileId> for HirFileId {
    fn from(id: FileId) -> Self {
        HirFileId(HirFileIdRepr::FileId(id))
    }
}

impl From<MacroFile> for HirFileId {
    fn from(id: MacroFile) -> Self {
        HirFileId(HirFileIdRepr::MacroFile(id))
    }
}

impl HirFileId {
    /// For macro-expansion files, returns the file original source file the
    /// expansion originated from.
    pub fn original_file(self, db: &impl AstDatabase) -> FileId {
        match self.0 {
            HirFileIdRepr::FileId(file_id) => file_id,
            HirFileIdRepr::MacroFile(macro_file) => {
                let loc = db.lookup_intern_macro(macro_file.macro_call_id);
                loc.ast_id.file_id().original_file(db)
            }
        }
    }

    /// Get the crate which the macro lives in, if it is a macro file.
    pub fn macro_crate(self, db: &impl AstDatabase) -> Option<CrateId> {
        match self.0 {
            HirFileIdRepr::FileId(_) => None,
            HirFileIdRepr::MacroFile(macro_file) => {
                let loc = db.lookup_intern_macro(macro_file.macro_call_id);
                Some(loc.def.krate)
            }
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MacroFile {
    macro_call_id: MacroCallId,
    macro_file_kind: MacroFileKind,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum MacroFileKind {
    Items,
    Expr,
}

/// `MacroCallId` identifies a particular macro invocation, like
/// `println!("Hello, {}", world)`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MacroCallId(salsa::InternId);
impl_intern_key!(MacroCallId);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MacroDefId {
    pub krate: CrateId,
    pub ast_id: AstId<ast::MacroCall>,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct MacroCallLoc {
    pub def: MacroDefId,
    pub ast_id: AstId<ast::MacroCall>,
}

impl MacroCallId {
    pub fn loc(self, db: &impl AstDatabase) -> MacroCallLoc {
        db.lookup_intern_macro(self)
    }

    pub fn as_file(self, kind: MacroFileKind) -> HirFileId {
        let macro_file = MacroFile { macro_call_id: self, macro_file_kind: kind };
        macro_file.into()
    }
}

impl MacroCallLoc {
    pub fn id(self, db: &impl AstDatabase) -> MacroCallId {
        db.intern_macro(self)
    }
}

/// `AstId` points to an AST node in any file.
///
/// It is stable across reparses, and can be used as salsa key/value.
// FIXME: isn't this just a `Source<FileAstId<N>>` ?
#[derive(Debug)]
pub struct AstId<N: AstNode> {
    file_id: HirFileId,
    file_ast_id: FileAstId<N>,
}

impl<N: AstNode> Clone for AstId<N> {
    fn clone(&self) -> AstId<N> {
        *self
    }
}
impl<N: AstNode> Copy for AstId<N> {}

impl<N: AstNode> PartialEq for AstId<N> {
    fn eq(&self, other: &Self) -> bool {
        (self.file_id, self.file_ast_id) == (other.file_id, other.file_ast_id)
    }
}
impl<N: AstNode> Eq for AstId<N> {}
impl<N: AstNode> Hash for AstId<N> {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        (self.file_id, self.file_ast_id).hash(hasher);
    }
}

impl<N: AstNode> AstId<N> {
    pub fn new(file_id: HirFileId, file_ast_id: FileAstId<N>) -> AstId<N> {
        AstId { file_id, file_ast_id }
    }

    pub fn file_id(&self) -> HirFileId {
        self.file_id
    }

    pub fn to_node(&self, db: &impl AstDatabase) -> N {
        let syntax_node = db.ast_id_to_node(self.file_id, self.file_ast_id.into());
        N::cast(syntax_node).unwrap()
    }
}

pub(crate) fn macro_def_query(db: &impl AstDatabase, id: MacroDefId) -> Option<Arc<MacroRules>> {
    let macro_call = id.ast_id.to_node(db);
    let arg = macro_call.token_tree()?;
    let (tt, _) = mbe::ast_to_token_tree(&arg).or_else(|| {
        log::warn!("fail on macro_def to token tree: {:#?}", arg);
        None
    })?;
    let rules = MacroRules::parse(&tt).ok().or_else(|| {
        log::warn!("fail on macro_def parse: {:#?}", tt);
        None
    })?;
    Some(Arc::new(rules))
}

pub(crate) fn macro_arg_query(db: &impl AstDatabase, id: MacroCallId) -> Option<Arc<tt::Subtree>> {
    let loc = db.lookup_intern_macro(id);
    let macro_call = loc.ast_id.to_node(db);
    let arg = macro_call.token_tree()?;
    let (tt, _) = mbe::ast_to_token_tree(&arg)?;
    Some(Arc::new(tt))
}

pub(crate) fn macro_expand_query(
    db: &impl AstDatabase,
    id: MacroCallId,
) -> Result<Arc<tt::Subtree>, String> {
    let loc = db.lookup_intern_macro(id);
    let macro_arg = db.macro_arg(id).ok_or("Fail to args in to tt::TokenTree")?;

    let macro_rules = db.macro_def(loc.def).ok_or("Fail to find macro definition")?;
    let tt = macro_rules.expand(&macro_arg).map_err(|err| format!("{:?}", err))?;
    // Set a hard limit for the expanded tt
    let count = tt.count();
    if count > 65536 {
        return Err(format!("Total tokens count exceed limit : count = {}", count));
    }
    Ok(Arc::new(tt))
}

pub(crate) fn parse_or_expand_query(
    db: &impl AstDatabase,
    file_id: HirFileId,
) -> Option<SyntaxNode> {
    match file_id.0 {
        HirFileIdRepr::FileId(file_id) => Some(db.parse(file_id).tree().syntax().clone()),
        HirFileIdRepr::MacroFile(macro_file) => {
            db.parse_macro(macro_file).map(|it| it.syntax_node())
        }
    }
}

pub(crate) fn parse_macro_query(
    db: &impl AstDatabase,
    macro_file: MacroFile,
) -> Option<Parse<SyntaxNode>> {
    let _p = profile("parse_macro_query");
    let macro_call_id = macro_file.macro_call_id;
    let tt = db
        .macro_expand(macro_call_id)
        .map_err(|err| {
            // Note:
            // The final goal we would like to make all parse_macro success,
            // such that the following log will not call anyway.
            log::warn!("fail on macro_parse: (reason: {})", err,);
        })
        .ok()?;
    match macro_file.macro_file_kind {
        MacroFileKind::Items => mbe::token_tree_to_items(&tt).ok().map(Parse::to_syntax),
        MacroFileKind::Expr => mbe::token_tree_to_expr(&tt).ok().map(Parse::to_syntax),
    }
}
<|endoftext|>


<commit_before>extern crate getopts;
#[macro_use]
extern crate modparse_derive;
#[macro_use]
extern crate derive_builder;

#[macro_use]
extern crate torchrs;

use torchrs::autograd::{Variable, BackwardArgs};
use torchrs::tensor::Tensor;
use torchrs::optim;

use torchrs::macros::*;
use torchrs::nn;
use torchrs::nn::{ModuleStruct, ModIntf, ModDelegate, Module};
use torchrs::nn::functional as F;
use torchrs::utils::data as D;

use getopts::Options;
use std::env;

fn MNIST() -> D::Dataset<D::Batch<f32, i64>> {
    panic!("just a stub")
}

#[derive(Clone, Builder)]
struct NetArgs {
    #[builder(default="64")]
    batch_size: usize,
    #[builder(default="1000")]
    test_batch_size: usize,
    #[builder(default="10")]
    epochs: u32,
    #[builder(default="0.01")]
    lr: f32,
    #[builder(default="0.5")]
    momentum: f32,
    #[builder(default="true")]
    cuda: bool,
    #[builder(default="1")]
    seed: usize,
    #[builder(default="10")]
    log_interval: usize,
}

impl Default for NetArgs {
    fn default() -> Self {
        NetArgsBuilder::default().build().unwrap()
    }
}

fn parse_args() -> NetArgs {
    let cmd_args: Vec<String> = env::args().collect();
    let mut args = NetArgs::default();
    let mut opts = Options::new();
    /*
    	* XXX do parsey stuff
    	*/
    args
}

#[derive(ModParse)]
struct Net<T: Default> {
    delegate: nn::Module<T>,
    conv1: nn::Conv2d<T>,
    conv2: nn::Conv2d<T>,
    fc1: nn::Linear<T>,
    fc2: nn::Linear<T>,
}

impl<T: Default> Net<T> {
    pub fn new() -> Net<T> {
        Net {
                delegate: nn::Module::new(),
                conv1: nn::Conv2d::build(1, 10, (5, 5)).done(),
                conv2: nn::Conv2d::build(10, 20, (5, 5)).done(),
                fc1: nn::Linear::build(320, 50).done(),
                fc2: nn::Linear::build(50, 10).done(),
            }
            .init_module()
    }
}
impl_mod_delegate!(Net);

// The forward operations could take on one of two implementations.
// The first supporting a near verbatim version of the python
// implementation, and the second supporting a slightly more
// idiomatic to Rust method chaining.

// a) as a near verbatim implementation of the python version
impl ModIntf<f32> for Net<f32> {
    fn forward(&mut self, args: &mut Variable<f32>) -> Variable<f32> {
        let training = self.delegate.training;
        let pool_val = F::MaxPool2dArgs::default();
        let dropout_val = F::DropoutArgs::default();
        let mut x = F::relu(&F::max_pool2d(&self.conv1.f(args), (2, 2), &pool_val));
        let mut x = F::relu(&F::max_pool2d(&F::dropout2d(&self.conv2.f(&mut x), &dropout_val),
                                           (2, 2),
                                           &pool_val));
        let mut x = x.view(&[-1, 320]);
        let x = F::relu(&self.fc1.f(&mut x));
        let mut x = F::dropout(&x, &dropout_val);
        let x = self.fc2.f(&mut x);
        F::log_softmax(&x)
    }
}

fn train(model: &mut Net<f32>,
         args: &NetArgs,
         train_loader: &D::BatchLoader<f32, i64>,
         epoch: u32,
         optimizer: &mut optim::OptIntf) {
    model.train();
    for (batch_idx, (ref data, ref target)) in train_loader.iter().enumerate() {
        let (mut data, target) = if args.cuda {
            (Variable::new(data.cuda().clone()), Variable::new(target.cuda().clone()))
        } else {
            (Variable::new(data.clone()), Variable::new(target.clone()))
        };
        optimizer.zero_grad();
        let output = model.f(&mut data);
        let mut loss = F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        loss.backward(&BackwardArgs::default());
        optimizer.step();
        if batch_idx % args.log_interval == 0 {
            println!("Train Epoch: {} [{}/{} ({:.0}%)]\tLoss: {:.6}",
                     epoch,
                     batch_idx * data.data().len(),
                     train_loader.dataset.len(),
                     100. * (batch_idx as f32) / train_loader.len() as f32,
                     loss.data()[0]);
        }
    }
}

fn test(model: &mut Net<f32>, args: &NetArgs, test_loader: &D::BatchLoader<f32, i64>, epoch: u32) {
    model.eval();
    let mut test_loss = 0.;
    let mut correct = 0;
    for (ref data, ref target) in test_loader.iter() {
        let (mut data, mut target) = if args.cuda {
            (Variable::new_volatile(data.cuda().clone()), Variable::new(target.cuda().clone()))
        } else {
            (Variable::new(data.clone()), Variable::new(target.clone()))
        };
        let mut output = model.f(&mut data);
        test_loss += F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        let pred = output.data().reduce_max(1).1;
        correct += pred.tensor_eq(target.data()).cpu().sum();
    }
    test_loss /= test_loader.len() as f32;
    println!("\nTest set: Average loss: {:.4}, Accuracy: {}/{} ({:.0}%)\n",
             test_loss,
             correct,
             test_loader.dataset.len(),
             100. * (correct as f32) / (test_loader.dataset.len() as f32))

}

fn main() {
    // no data loaders yet so demonstrate with Vec placeholder
    let train_loader: D::BatchLoader<f32, i64> = D::DataLoader::new(MNIST(),
                                                                    D::DataLoaderArgs::default());
    let test_loader: D::BatchLoader<f32, i64> = D::DataLoader::new(MNIST(),
                                                                   D::DataLoaderArgs::default());
    let args = parse_args();
    let mut optimizer = optim::SGD::new();

    let mut model = Net::new();
    for epoch in 1..args.epochs + 1 {
        train(&mut model, &args, &train_loader, epoch, &mut optimizer);
        test(&mut model, &args, &test_loader, epoch);
    }
}
<commit_msg>fix for API changes<commit_after>extern crate getopts;
#[macro_use]
extern crate modparse_derive;
#[macro_use]
extern crate derive_builder;

#[macro_use]
extern crate torchrs;

use torchrs::autograd::{Variable, VariableArgs, BackwardArgs};
use torchrs::tensor::Tensor;
use torchrs::optim;

use torchrs::macros::*;
use torchrs::nn;
use torchrs::nn::{ModuleStruct, ModIntf, ModDelegate, Module};
use torchrs::nn::functional as F;
use torchrs::utils::data as D;

use getopts::Options;
use std::env;

fn MNIST() -> D::Dataset<D::Batch<f32, i64>> {
    panic!("just a stub")
}

#[derive(Clone, Builder)]
struct NetArgs {
    #[builder(default="64")]
    batch_size: usize,
    #[builder(default="1000")]
    test_batch_size: usize,
    #[builder(default="10")]
    epochs: u32,
    #[builder(default="0.01")]
    lr: f32,
    #[builder(default="0.5")]
    momentum: f32,
    #[builder(default="true")]
    cuda: bool,
    #[builder(default="1")]
    seed: usize,
    #[builder(default="10")]
    log_interval: usize,
}

impl Default for NetArgs {
    fn default() -> Self {
        NetArgsBuilder::default().build().unwrap()
    }
}

fn parse_args() -> NetArgs {
    let cmd_args: Vec<String> = env::args().collect();
    let mut args = NetArgs::default();
    let mut opts = Options::new();
    /*
    	* XXX do parsey stuff
    	*/
    args
}

#[derive(ModParse)]
struct Net<T: Default> {
    delegate: nn::Module<T>,
    conv1: nn::Conv2d<T>,
    conv2: nn::Conv2d<T>,
    fc1: nn::Linear<T>,
    fc2: nn::Linear<T>,
}

impl<T: Default> Net<T> {
    pub fn new() -> Net<T> {
        Net {
                delegate: nn::Module::new(),
                conv1: nn::Conv2d::build(1, 10, (5, 5)).done(),
                conv2: nn::Conv2d::build(10, 20, (5, 5)).done(),
                fc1: nn::Linear::build(320, 50).done(),
                fc2: nn::Linear::build(50, 10).done(),
            }
            .init_module()
    }
}
impl_mod_delegate!(Net);

// The forward operations could take on one of two implementations.
// The first supporting a near verbatim version of the python
// implementation, and the second supporting a slightly more
// idiomatic to Rust method chaining.

// a) as a near verbatim implementation of the python version
impl ModIntf<f32> for Net<f32> {
    fn forward(&mut self, args: &mut Variable<f32>) -> Variable<f32> {
        let training = self.delegate.training;
        let pool_val = F::MaxPool2dArgs::default();
        let dropout_val = F::DropoutArgs::default();
        let mut x = F::relu(&F::max_pool2d(&self.conv1.f(args), (2, 2), &pool_val));
        let mut x = F::relu(&F::max_pool2d(&F::dropout2d(&self.conv2.f(&mut x), &dropout_val),
                                           (2, 2),
                                           &pool_val));
        let mut x = x.view(&[-1, 320]);
        let x = F::relu(&self.fc1.f(&mut x));
        let mut x = F::dropout(&x, &dropout_val);
        let x = self.fc2.f(&mut x);
        F::log_softmax(&x)
    }
}

fn train(model: &mut Net<f32>,
         args: &NetArgs,
         train_loader: &D::BatchLoader<f32, i64>,
         epoch: u32,
         optimizer: &mut optim::OptIntf) {
    model.train();
    for (batch_idx, (ref data, ref target)) in train_loader.iter().enumerate() {
        let (mut data, target) = if args.cuda {
            (Variable::new(data.cuda().clone()), Variable::new(target.cuda().clone()))
        } else {
            (Variable::new(data.clone()), Variable::new(target.clone()))
        };
        optimizer.zero_grad();
        let output = model.f(&mut data);
        let mut loss = F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        loss.backward(&BackwardArgs::default());
        optimizer.step();
        if batch_idx % args.log_interval == 0 {
            println!("Train Epoch: {} [{}/{} ({:.0}%)]\tLoss: {:.6}",
                     epoch,
                     batch_idx * data.data().len(),
                     train_loader.dataset.len(),
                     100. * (batch_idx as f32) / train_loader.len() as f32,
                     loss.data()[0]);
        }
    }
}

fn test(model: &mut Net<f32>, args: &NetArgs, test_loader: &D::BatchLoader<f32, i64>, epoch: u32) {
    model.eval();
    let mut test_loss = 0.;
    let mut correct = 0;
    for (ref data, ref target) in test_loader.iter() {
        let varargs = VariableArgs {
            volatile: true,
            ..Default::default()
        };
        let (mut data, mut target) = if args.cuda {
            (Variable::new_args(data.cuda().clone(), varargs), Variable::new(target.cuda().clone()))
        } else {
            (Variable::new_args(data.clone(), varargs), Variable::new(target.clone()))
        };
        let mut output = model.f(&mut data);
        test_loss += F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        let pred = output.data().reduce_max(1).1;
        correct += pred.tensor_eq(&*target.data()).cpu().sum();
    }
    test_loss /= test_loader.len() as f32;
    println!("\nTest set: Average loss: {:.4}, Accuracy: {}/{} ({:.0}%)\n",
             test_loss,
             correct,
             test_loader.dataset.len(),
             100. * (correct as f32) / (test_loader.dataset.len() as f32))

}

fn main() {
    // no data loaders yet so demonstrate with Vec placeholder
    let train_loader: D::BatchLoader<f32, i64> = D::DataLoader::new(MNIST(),
                                                                    D::DataLoaderArgs::default());
    let test_loader: D::BatchLoader<f32, i64> = D::DataLoader::new(MNIST(),
                                                                   D::DataLoaderArgs::default());
    let args = parse_args();
    let mut optimizer = optim::SGD::new();

    let mut model = Net::new();
    for epoch in 1..args.epochs + 1 {
        train(&mut model, &args, &train_loader, epoch, &mut optimizer);
        test(&mut model, &args, &test_loader, epoch);
    }
}
<|endoftext|>


<commit_before>extern crate polydraw;

use std::cmp::{min, max};

use polydraw::{Application, Renderer, RenderFrame};

struct MouseRenderer {
   mouse_x: u32,
   mouse_y: u32,
   old: Option<(u32, u32)>,
}

impl MouseRenderer {
   fn new() -> Self {
      MouseRenderer {
         mouse_x: 0,
         mouse_y: 0,
         old: None,
      }
   }
}

impl Renderer for MouseRenderer {
   fn render(&mut self, frame: &mut RenderFrame) {
      let half = 35_u32;

      match self.old {
         Some((old_x, old_y)) => {
            let old_x_start = max(0, (old_x as i32) - half as i32) as u32;
            let old_x_end = min(frame.width, old_x + half);

            let old_y_start = max(0, (old_y as i32) - half as i32) as u32;
            let old_y_end = min(frame.height, old_y + half);

            for y in old_y_start..old_y_end {
               let row_i = 3 * y * frame.width;

               let mut i = (3 * old_x_start + row_i) as usize;
               frame.data[i] = 33;
               frame.data[i + 1] = 168;
               frame.data[i + 2] = 222;

               i = (3 * old_x_end + row_i) as usize;
               frame.data[i] = 33;
               frame.data[i + 1] = 168;
               frame.data[i + 2] = 222;
            }

            for x in old_x_start..old_x_end {
               let mut row_i = 3 * old_y_start * frame.width;

               let mut i = (3 * x + row_i) as usize;
               frame.data[i] = 33;
               frame.data[i + 1] = 168;
               frame.data[i + 2] = 222;

               row_i = 3 * old_y_end * frame.width;

               i = (3 * x + row_i) as usize;
               frame.data[i] = 33;
               frame.data[i + 1] = 168;
               frame.data[i + 2] = 222;
            }
         },
         None => {},
      }

      let x_start = max(0, (self.mouse_x as i32) - half as i32) as u32;
      let x_end = min(frame.width, self.mouse_x + half);

      let y_start = max(0, (self.mouse_y as i32) - half as i32) as u32;
      let y_end = min(frame.height, self.mouse_y + half);

      for y in y_start..y_end {
         let row_i = 3 * y * frame.width;
         for x in x_start..x_end {
            let i = (3 * x + row_i) as usize;
            frame.data[i] = 33;
            frame.data[i + 1] = 168;
            frame.data[i + 2] = 222;
         }
      }

      for y in y_start..y_end {
         let row_i = 3 * y * frame.width;

         let mut i = (3 * x_start + row_i) as usize;
         frame.data[i] = 255;
         frame.data[i + 1] = 255;
         frame.data[i + 2] = 255;

         i = (3 * x_end + row_i) as usize;
         frame.data[i] = 255;
         frame.data[i + 1] = 255;
         frame.data[i + 2] = 255;
      }

      for x in x_start..x_end {
         let mut row_i = 3 * y_start * frame.width;

         let mut i = (3 * x + row_i) as usize;
         frame.data[i] = 255;
         frame.data[i + 1] = 255;
         frame.data[i + 2] = 255;

         row_i = 3 * y_end * frame.width;

         i = (3 * x + row_i) as usize;
         frame.data[i] = 255;
         frame.data[i + 1] = 255;
         frame.data[i + 2] = 255;
      }

      self.old = Some((self.mouse_x, self.mouse_y));
   }

   fn mouse_moved(&mut self, x: u32, y: u32) {
      self.mouse_x = x;
      self.mouse_y = y;
   }
}

fn main() {
   let mut renderer = MouseRenderer::new();

   Application::new()
      .renderer(&mut renderer)
      .title("Mouse")
      .run();
}
<commit_msg>Simplifying the mouse example<commit_after>extern crate polydraw;

use std::cmp::{min, max};

use polydraw::{Application, Renderer, RenderFrame};

struct MouseRenderer {
   mouse_x: u32,
   mouse_y: u32,
}

impl MouseRenderer {
   fn new() -> Self {
      MouseRenderer {
         mouse_x: 1000000,
         mouse_y: 1000000,
      }
   }
}

impl Renderer for MouseRenderer {
   fn render(&mut self, frame: &mut RenderFrame) {
      for i in 0 as usize..(frame.width * frame.height * 3) as usize {
         frame.data[i] = 0;
      }

      let half = 35_u32;

      let x_start = max(0, (self.mouse_x as i32) - half as i32) as u32;
      let x_end = min(frame.width, self.mouse_x + half);

      let y_start = max(0, (self.mouse_y as i32) - half as i32) as u32;
      let y_end = min(frame.height, self.mouse_y + half);

      for y in y_start..y_end {
         let row_i = 3 * y * frame.width;
         for x in x_start..x_end {
            let i = (3 * x + row_i) as usize;
            frame.data[i] = 33;
            frame.data[i + 1] = 168;
            frame.data[i + 2] = 222;
         }
      }
   }

   fn mouse_moved(&mut self, x: u32, y: u32) {
      self.mouse_x = x;
      self.mouse_y = y;
   }
}

fn main() {
   let mut renderer = MouseRenderer::new();

   Application::new()
      .renderer(&mut renderer)
      .title("Mouse")
      .run();
}
<|endoftext|>


<commit_before>extern crate kailua_env;
extern crate kailua_diag;
extern crate kailua_syntax;
extern crate kailua_check;
#[macro_use] extern crate log;
extern crate env_logger;

use std::str;
use std::io;
use std::env;
use std::cell::RefCell;
use std::rc::Rc;
use std::path::Path;

use kailua_env::{Source, SourceFile};
use kailua_diag::{Report, ConsoleReport, TrackMaxKind};
use kailua_syntax::{parse_chunk, Chunk};
use kailua_check::{FsSource, FsOptions, Context, check_from_chunk};

struct LocalFsSource {
    source: Rc<RefCell<Source>>,
    report: Rc<Report>,
}

impl FsSource for LocalFsSource {
    fn chunk_from_path(&self, resolved_path: &Path) -> Result<Option<Chunk>, String> {
        match SourceFile::from_file(resolved_path) {
            Ok(file) => {
                let mut source = self.source.borrow_mut();
                let span = source.add(file);
                if let Ok(chunk) = parse_chunk(&source, span, &*self.report) {
                    Ok(Some(chunk))
                } else {
                    Err(format!("parse error"))
                }
            }
            Err(e) => {
                if e.kind() == io::ErrorKind::NotFound {
                    Ok(None)
                } else {
                    Err(e.to_string())
                }
            }
        }
    }
}

fn parse_and_check(mainpath: &Path) -> Result<(), String> {
    let source = Rc::new(RefCell::new(Source::new()));
    let report = Rc::new(TrackMaxKind::new(ConsoleReport::new(source.clone())));
    let mut context = Context::new(report.clone());

    let fssource = LocalFsSource { source: source, report: report.clone() };
    let filechunk = fssource.chunk_from_path(mainpath)?;
    let filechunk = filechunk.ok_or_else(|| format!("cannot found the main path"))?;

    let root = mainpath.parent().unwrap_or(&Path::new(".."));
    let opts = Rc::new(RefCell::new(FsOptions::new(fssource, root.to_owned())));

    check_from_chunk(&mut context, filechunk, opts)?;
    if report.can_continue() {
        Ok(())
    } else {
        Err("stopped due to prior errors".into())
    }
}

pub fn main() {
    env_logger::init().unwrap();
    for path in env::args().skip(1) {
        println!("--== {} ==--", path);
        if let Err(e) = parse_and_check(&Path::new(&path)) {
            // hide the internal error message, which will be eventually removed
            println!("stopped due to prior errors.");
            info!("error while checking {}: {}", path, e);
        } else {
            println!("done.");
        }
        println!("");
    }
}

<commit_msg>Fixed `kailua-check` binary broken from the incompatible borrow.<commit_after>extern crate kailua_env;
extern crate kailua_diag;
extern crate kailua_syntax;
extern crate kailua_check;
#[macro_use] extern crate log;
extern crate env_logger;

use std::str;
use std::io;
use std::env;
use std::cell::RefCell;
use std::rc::Rc;
use std::path::Path;

use kailua_env::{Source, SourceFile};
use kailua_diag::{Report, ConsoleReport, TrackMaxKind};
use kailua_syntax::{parse_chunk, Chunk};
use kailua_check::{FsSource, FsOptions, Context, check_from_chunk};

struct LocalFsSource {
    source: Rc<RefCell<Source>>,
    report: Rc<Report>,
}

impl FsSource for LocalFsSource {
    fn chunk_from_path(&self, resolved_path: &Path) -> Result<Option<Chunk>, String> {
        match SourceFile::from_file(resolved_path) {
            Ok(file) => {
                let span = self.source.borrow_mut().add(file);
                if let Ok(chunk) = parse_chunk(&self.source.borrow(), span, &*self.report) {
                    Ok(Some(chunk))
                } else {
                    Err(format!("parse error"))
                }
            }
            Err(e) => {
                if e.kind() == io::ErrorKind::NotFound {
                    Ok(None)
                } else {
                    Err(e.to_string())
                }
            }
        }
    }
}

fn parse_and_check(mainpath: &Path) -> Result<(), String> {
    let source = Rc::new(RefCell::new(Source::new()));
    let report = Rc::new(TrackMaxKind::new(ConsoleReport::new(source.clone())));
    let mut context = Context::new(report.clone());

    let fssource = LocalFsSource { source: source, report: report.clone() };
    let filechunk = fssource.chunk_from_path(mainpath)?;
    let filechunk = filechunk.ok_or_else(|| format!("cannot found the main path"))?;

    let root = mainpath.parent().unwrap_or(&Path::new(".."));
    let opts = Rc::new(RefCell::new(FsOptions::new(fssource, root.to_owned())));

    check_from_chunk(&mut context, filechunk, opts)?;
    if report.can_continue() {
        Ok(())
    } else {
        Err("stopped due to prior errors".into())
    }
}

pub fn main() {
    env_logger::init().unwrap();
    for path in env::args().skip(1) {
        println!("--== {} ==--", path);
        if let Err(e) = parse_and_check(&Path::new(&path)) {
            // hide the internal error message, which will be eventually removed
            println!("stopped due to prior errors.");
            info!("error while checking {}: {}", path, e);
        } else {
            println!("done.");
        }
        println!("");
    }
}

<|endoftext|>