<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// This file is imported into every module by default.

/* Reexported core operators */

pub use either::{Either, Left, Right};
pub use kinds::{Const, Copy, Owned, Durable};
pub use ops::{Add, Sub, Mul, Div, Modulo, Neg, Not};
pub use ops::{BitAnd, BitOr, BitXor};
pub use ops::{Drop};
pub use ops::{Shl, Shr, Index};
pub use option::{Option, Some, None};
pub use result::{Result, Ok, Err};

/* Reexported types and traits */

pub use clone::Clone;
pub use cmp::{Eq, Ord, TotalOrd, Ordering, Less, Equal, Greater};
pub use container::{Container, Mutable, Map, Set};
pub use hash::Hash;
pub use iter::{BaseIter, ReverseIter, ExtendedIter, EqIter, CopyableIter};
pub use iter::{CopyableOrderedIter, CopyableNonstrictIter, Times};
pub use num::NumCast;
pub use path::GenericPath;
pub use path::Path;
pub use path::PosixPath;
pub use path::WindowsPath;
pub use ptr::Ptr;
pub use str::{StrSlice, Trimmable, OwnedStr};
pub use to_bytes::IterBytes;
pub use to_str::ToStr;
pub use tuple::{CopyableTuple, ImmutableTuple, ExtendedTupleOps};
pub use vec::{CopyableVector, ImmutableVector};
pub use vec::{ImmutableEqVector, ImmutableCopyableVector};
pub use vec::{OwnedVector, OwnedCopyableVector};

/* Reexported modules */

pub use at_vec;
pub use bool;
pub use cast;
pub use char;
pub use cmp;
pub use either;
pub use f32;
pub use f64;
pub use float;
pub use i16;
pub use i32;
pub use i64;
pub use i8;
pub use int;
pub use io;
pub use iter;
pub use libc;
pub use num;
pub use ops;
pub use option;
pub use os;
pub use path;
pub use comm;
pub use unstable;
pub use ptr;
pub use rand;
pub use result;
pub use str;
pub use sys;
pub use task;
pub use to_str;
pub use u16;
pub use u32;
pub use u64;
pub use u8;
pub use uint;
pub use vec;
<commit_msg>core: Add spawn, stream and friends to prelude. #5299<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// This file is imported into every module by default.

/* Reexported core operators */

pub use either::{Either, Left, Right};
pub use kinds::{Const, Copy, Owned, Durable};
pub use ops::{Add, Sub, Mul, Div, Modulo, Neg, Not};
pub use ops::{BitAnd, BitOr, BitXor};
pub use ops::{Drop};
pub use ops::{Shl, Shr, Index};
pub use option::{Option, Some, None};
pub use result::{Result, Ok, Err};

/* Reexported types and traits */

pub use clone::Clone;
pub use cmp::{Eq, Ord, TotalOrd, Ordering, Less, Equal, Greater};
pub use container::{Container, Mutable, Map, Set};
pub use hash::Hash;
pub use iter::{BaseIter, ReverseIter, ExtendedIter, EqIter, CopyableIter};
pub use iter::{CopyableOrderedIter, CopyableNonstrictIter, Times};
pub use num::NumCast;
pub use path::GenericPath;
pub use path::Path;
pub use path::PosixPath;
pub use path::WindowsPath;
pub use ptr::Ptr;
pub use str::{StrSlice, Trimmable, OwnedStr};
pub use to_bytes::IterBytes;
pub use to_str::ToStr;
pub use tuple::{CopyableTuple, ImmutableTuple, ExtendedTupleOps};
pub use vec::{CopyableVector, ImmutableVector};
pub use vec::{ImmutableEqVector, ImmutableCopyableVector};
pub use vec::{OwnedVector, OwnedCopyableVector};

/* Reexported runtime types */
pub use comm::{stream, Port, Chan, GenericChan, GenericSmartChan, GenericPort, Peekable};
pub use task::spawn;

/* Reexported modules */

pub use at_vec;
pub use bool;
pub use cast;
pub use char;
pub use cmp;
pub use either;
pub use f32;
pub use f64;
pub use float;
pub use i16;
pub use i32;
pub use i64;
pub use i8;
pub use int;
pub use io;
pub use iter;
pub use libc;
pub use num;
pub use ops;
pub use option;
pub use os;
pub use path;
pub use comm;
pub use unstable;
pub use ptr;
pub use rand;
pub use result;
pub use str;
pub use sys;
pub use task;
pub use to_str;
pub use u16;
pub use u32;
pub use u64;
pub use u8;
pub use uint;
pub use vec;
<|endoftext|>


<commit_before>/*
 * An example redis server accepting GET and SET requests
 * listening on 127.0.0.1:8000.
 *
 * You can test it with "redis-cli -p 8000".
 *
 * Copyright (c) 2014 by Michael Neumann
 *
 */

extern crate redis = "redis#0.1";
extern crate sync = "sync#0.10-pre";
extern crate collections;

use std::io::net::ip::SocketAddr;
use std::io::net::tcp::{TcpListener,TcpStream};
use std::io::{Listener,Acceptor,Writer};
use std::io::BufferedStream;
use std::task;
use collections::hashmap::HashMap;
use sync::RWArc;

fn handle_connection(conn: TcpStream, shared_ht: RWArc<HashMap<~[u8],~[u8]>>) {
    debug!("Got connection");

    let mut io = BufferedStream::new(conn);

    loop {
        match redis::parse(&mut io).unwrap() {
            redis::List(ref lst) => {
                match lst.get(0) {
                    Some(&redis::Data(ref command)) => {
                        if command.as_slice() == bytes!("GET") {
                            match (lst.len(), lst.get(1)) {
                                (2, Some(&redis::Data(ref key))) => {
                                    debug!("GET: {:s}", std::str::from_utf8(key.as_slice()).unwrap());
                                    let mut cwr = redis::CommandWriter::new();
                                    shared_ht.read(|ht| {
                                        match ht.find(key) {
                                            Some(val) => {
                                                cwr.args(1);
                                                cwr.arg_bin(*val);
                                            }
                                            None => {
                                                cwr.nil();
                                            }
                                        }
                                    });
                                    cwr.with_buf(|bytes| {io.write(bytes); io.flush()});
                                    continue;
                                }
                                _ => { /* fallthrough: error */ }
                            }
                        }
                        else if command.as_slice() == bytes!("SET") {
                            match (lst.len(), lst.get(1), lst.get(2)) {
                                (3, Some(&redis::Data(ref key)), Some(&redis::Data(ref val))) => {
                                    debug!("SET: {:s} {:?}", std::str::from_utf8(key.as_slice()).unwrap(), val);
                                    shared_ht.write(|ht| ht.insert(key.clone(), val.clone()));
                                    let mut cwr = redis::CommandWriter::new();
                                    cwr.status("OK");
                                    cwr.with_buf(|bytes| {io.write(bytes); io.flush()});
                                    continue;
                                }
                                _ => { /* fallthrough: error */ }
                            }
                        }
                        else {
                            /* fallthrough: error */
                        }
                    }
                    _ => { /* fallthrough: error */ }
                }
            }
            _ => { /* fallthrough: error */ }
        }

        /* error */
        let mut cwr = redis::CommandWriter::new();
        cwr.error("Invalid Command");
        cwr.with_buf(|bytes| {io.write(bytes); io.flush()});
    }
}

fn main() {
    let addr: SocketAddr = from_str("127.0.0.1:8000").unwrap();
    let shared_ht = RWArc::new(HashMap::new());

    match TcpListener::bind(addr) {
        Ok(listener) => {
            match listener.listen() {
                Ok(ref mut acceptor) => {
                    loop {
                        match acceptor.accept() {
                            Ok(conn) => {
                                let ht = shared_ht.clone();
                                task::spawn(proc() {
                                    handle_connection(conn, ht)
                                });
                            }
                            Err(_) => {}
                        }
                    }
                }
                Err(_) => {}
            }
        }
        Err(_) => {}
    }
}
<commit_msg>Macro debug! moved into extern create log<commit_after>#[feature(phase)];

/*
 * An example redis server accepting GET and SET requests
 * listening on 127.0.0.1:8000.
 *
 * You can test it with "redis-cli -p 8000".
 *
 * Copyright (c) 2014 by Michael Neumann
 *
 */

extern crate redis = "redis#0.1";
extern crate sync = "sync#0.10-pre";
extern crate collections;
#[phase(syntax, link)] extern crate log;

use std::io::net::ip::SocketAddr;
use std::io::net::tcp::{TcpListener,TcpStream};
use std::io::{Listener,Acceptor,Writer};
use std::io::BufferedStream;
use std::task;
use collections::hashmap::HashMap;
use sync::RWArc;

fn handle_connection(conn: TcpStream, shared_ht: RWArc<HashMap<~[u8],~[u8]>>) {
    debug!("Got connection");

    let mut io = BufferedStream::new(conn);

    loop {
        match redis::parse(&mut io).unwrap() {
            redis::List(ref lst) => {
                match lst.get(0) {
                    Some(&redis::Data(ref command)) => {
                        if command.as_slice() == bytes!("GET") {
                            match (lst.len(), lst.get(1)) {
                                (2, Some(&redis::Data(ref key))) => {
                                    debug!("GET: {:s}", std::str::from_utf8(key.as_slice()).unwrap());
                                    let mut cwr = redis::CommandWriter::new();
                                    shared_ht.read(|ht| {
                                        match ht.find(key) {
                                            Some(val) => {
                                                cwr.args(1);
                                                cwr.arg_bin(*val);
                                            }
                                            None => {
                                                cwr.nil();
                                            }
                                        }
                                    });
                                    cwr.with_buf(|bytes| {io.write(bytes); io.flush()});
                                    continue;
                                }
                                _ => { /* fallthrough: error */ }
                            }
                        }
                        else if command.as_slice() == bytes!("SET") {
                            match (lst.len(), lst.get(1), lst.get(2)) {
                                (3, Some(&redis::Data(ref key)), Some(&redis::Data(ref val))) => {
                                    debug!("SET: {:s} {:?}", std::str::from_utf8(key.as_slice()).unwrap(), val);
                                    shared_ht.write(|ht| ht.insert(key.clone(), val.clone()));
                                    let mut cwr = redis::CommandWriter::new();
                                    cwr.status("OK");
                                    cwr.with_buf(|bytes| {io.write(bytes); io.flush()});
                                    continue;
                                }
                                _ => { /* fallthrough: error */ }
                            }
                        }
                        else {
                            /* fallthrough: error */
                        }
                    }
                    _ => { /* fallthrough: error */ }
                }
            }
            _ => { /* fallthrough: error */ }
        }

        /* error */
        let mut cwr = redis::CommandWriter::new();
        cwr.error("Invalid Command");
        cwr.with_buf(|bytes| {io.write(bytes); io.flush()});
    }
}

fn main() {
    let addr: SocketAddr = from_str("127.0.0.1:8000").unwrap();
    let shared_ht = RWArc::new(HashMap::new());

    match TcpListener::bind(addr) {
        Ok(listener) => {
            match listener.listen() {
                Ok(ref mut acceptor) => {
                    loop {
                        match acceptor.accept() {
                            Ok(conn) => {
                                let ht = shared_ht.clone();
                                task::spawn(proc() {
                                    handle_connection(conn, ht)
                                });
                            }
                            Err(_) => {}
                        }
                    }
                }
                Err(_) => {}
            }
        }
        Err(_) => {}
    }
}
<|endoftext|>


<commit_before>use std::io::{self, Read, Write};
use handy_async::sync_io::{ReadExt, WriteExt};

use {Result, Method, Attribute, ErrorKind};
use types::{U12, TransactionId};
use attribute::RawAttribute;
use message::{Request, Indication, SuccessResponse, ErrorResponse, Response};
use constants::MAGIC_COOKIE;

#[derive(Debug, Clone)]
pub struct RawMessage {
    class: Class,
    method: U12,
    transaction_id: TransactionId,
    attributes: Vec<RawAttribute>,
}
impl RawMessage {
    pub fn new(class: Class,
               method: U12,
               transaction_id: TransactionId,
               attributes: Vec<RawAttribute>)
               -> Self {
        RawMessage {
            class: class,
            method: method,
            transaction_id: transaction_id,
            attributes: attributes,
        }
    }
    pub fn class(&self) -> Class {
        self.class
    }
    pub fn method(&self) -> U12 {
        self.method
    }
    pub fn transaction_id(&self) -> &TransactionId {
        &self.transaction_id
    }
    pub fn attributes(&self) -> &[RawAttribute] {
        &self.attributes
    }
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::new();
        self.write_to(&mut buf).expect("must succeed");
        buf
    }
    pub fn write_to<W: Write>(&self, writer: &mut W) -> Result<()> {
        let mut temp_writer = io::Cursor::new(vec![0; 20]);
        let message_type = Type {
            class: self.class,
            method: self.method,
        };
        track_try!(temp_writer.write_u16be(message_type.as_u16()));
        track_try!(temp_writer.write_u16be(0)); // dummy length
        track_try!(temp_writer.write_u32be(MAGIC_COOKIE));
        track_try!(temp_writer.write_all(&self.transaction_id));
        for attr in self.attributes.iter() {
            track_try!(attr.write_to(&mut temp_writer));
        }
        let attrs_len = temp_writer.get_ref().len() - 20;
        track_assert!(attrs_len < 0x10000,
                      ErrorKind::Failed,
                      "Too large message length: actual={}, limit=0xFFFF",
                      attrs_len);
        temp_writer.set_position(2);
        track_try!(temp_writer.write_u16be(attrs_len as u16));

        let buf = temp_writer.into_inner();
        track_try!(writer.write_all(&buf));
        Ok(())
    }
    pub fn read_from<R: Read>(reader: &mut R) -> Result<Self> {
        let message_type = track_try!(reader.read_u16be());
        let message_type = track_try!(Type::from_u16(message_type));
        let message_len = track_try!(reader.read_u16be());
        track_assert!(message_len % 4 == 0,
                      ErrorKind::NotStunMessage,
                      "Unexpected message length: {} % 4 != 0",
                      message_len);
        let magic_cookie = track_try!(reader.read_u32be());
        track_assert!(magic_cookie == MAGIC_COOKIE,
                      ErrorKind::NotStunMessage,
                      "Unexpected magic cookie: actual={}, expected={}",
                      magic_cookie,
                      MAGIC_COOKIE);;

        let mut transaction_id: [u8; 12] = [0; 12];
        track_try!(reader.read_exact(&mut transaction_id));

        let mut attrs = Vec::new();
        let mut reader = reader.take(message_len as u64);
        while reader.limit() > 0 {
            let attr = track_try!(RawAttribute::read_from(&mut reader));
            attrs.push(attr);
        }
        Ok(RawMessage::new(message_type.class,
                           message_type.method,
                           transaction_id,
                           attrs))
    }

    pub fn try_into_request<M, A>(self) -> Result<Request<M, A>>
        where M: Method,
              A: Attribute
    {
        panic!()
    }
    pub fn try_into_indication<M, A>(self) -> Result<Indication<M, A>>
        where M: Method,
              A: Attribute
    {
        panic!()
    }
    pub fn try_into_response<M, A>(self) -> Result<Response<M, A>>
        where M: Method,
              A: Attribute
    {
        panic!()
    }
    pub fn try_from_request<M, A>(from: Request<M, A>) -> Result<Self>
        where M: Method,
              A: Attribute
    {
        track_err!(Self::try_from(Class::Request,
                                  from.method(),
                                  *from.transaction_id(),
                                  from.attributes()))
    }
    pub fn try_from_indication<M, A>(from: Indication<M, A>) -> Result<Self>
        where M: Method,
              A: Attribute
    {
        track_err!(Self::try_from(Class::Indication,
                                  from.method(),
                                  *from.transaction_id(),
                                  from.attributes()))
    }
    pub fn try_from_success_response<M, A>(from: SuccessResponse<M, A>) -> Result<Self>
        where M: Method,
              A: Attribute
    {
        track_err!(Self::try_from(Class::SuccessResponse,
                                  from.method(),
                                  *from.transaction_id(),
                                  from.attributes()))
    }
    pub fn try_from_error_response<M, A>(from: ErrorResponse<M, A>) -> Result<Self>
        where M: Method,
              A: Attribute
    {
        track_err!(Self::try_from(Class::ErrorResponse,
                                  from.method(),
                                  *from.transaction_id(),
                                  from.attributes()))
    }

    fn try_from<M, A>(class: Class,
                      method: &M,
                      transaction_id: TransactionId,
                      attributes: &[A])
                      -> Result<Self>
        where M: Method,
              A: Attribute
    {
        let mut m = RawMessage::new(class, method.as_u12(), transaction_id, Vec::new());
        for attr in attributes {
            let raw_attr = track_try!(attr.encode(&m));
            m.attributes.push(raw_attr);
        }
        Ok(m)
    }
}

#[derive(Debug, Clone, Copy, PartialOrd, Ord, PartialEq, Eq, Hash)]
pub enum Class {
    Request = 0b00,
    Indication = 0b01,
    SuccessResponse = 0b10,
    ErrorResponse = 0b11,
}
impl Class {
    pub fn from_u8(value: u8) -> Option<Self> {
        match value {
            0b00 => Some(Class::Request),
            0b01 => Some(Class::Indication),
            0b10 => Some(Class::SuccessResponse),
            0b11 => Some(Class::ErrorResponse),
            _ => None,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct Type {
    pub class: Class,
    pub method: U12,
}
impl Type {
    /// TODO:
    ///
    /// ```text
    /// 0                 1
    /// 2  3  4 5 6 7 8 9 0 1 2 3 4 5
    ///
    /// +--+--+-+-+-+-+-+-+-+-+-+-+-+-+
    /// |M |M |M|M|M|C|M|M|M|C|M|M|M|M|
    /// |11|10|9|8|7|1|6|5|4|0|3|2|1|0|
    /// +--+--+-+-+-+-+-+-+-+-+-+-+-+-+
    ///
    /// Figure 3: Format of STUN Message Type Field
    /// ```
    pub fn as_u16(&self) -> u16 {
        let class = self.class as u16;
        let method = self.method.as_u12().as_u16();
        ((method & 0b0000_0000_1111) << 0) | ((class & 0b01) << 4) |
        ((method & 0b0000_0111_0000) << 5) | ((class & 0b10) << 7) |
        ((method & 0b1111_1000_0000) << 9)
    }

    pub fn from_u16(value: u16) -> Result<Self> {
        track_assert!(value >> 14 == 0,
                      ErrorKind::NotStunMessage,
                      "First two-bits of STUN message must be 0");
        let class = ((value >> 4) & 0b01) | ((value >> 7) & 0b10);
        let class = Class::from_u8(class as u8).unwrap();
        let method = (value & 0b0000_0000_1111) | ((value >> 1) & 0b0000_0111_0000) |
                     ((value >> 2) & 0b1111_1000_0000);
        let method = U12::from_u16(method).unwrap();
        Ok(Type {
            class: class,
            method: method,
        })
    }
}
<commit_msg>Implement `RawMessage::try_into_XXX` methods<commit_after>use std::io::{self, Read, Write};
use handy_async::sync_io::{ReadExt, WriteExt};

use {Result, Method, Attribute, ErrorKind};
use types::{U12, TransactionId};
use attribute::RawAttribute;
use message::{Request, Indication, SuccessResponse, ErrorResponse, Response};
use constants::MAGIC_COOKIE;

#[derive(Debug, Clone)]
pub struct RawMessage {
    class: Class,
    method: U12,
    transaction_id: TransactionId,
    attributes: Vec<RawAttribute>,
}
impl RawMessage {
    pub fn new(class: Class,
               method: U12,
               transaction_id: TransactionId,
               attributes: Vec<RawAttribute>)
               -> Self {
        RawMessage {
            class: class,
            method: method,
            transaction_id: transaction_id,
            attributes: attributes,
        }
    }
    pub fn class(&self) -> Class {
        self.class
    }
    pub fn method(&self) -> U12 {
        self.method
    }
    pub fn transaction_id(&self) -> &TransactionId {
        &self.transaction_id
    }
    pub fn attributes(&self) -> &[RawAttribute] {
        &self.attributes
    }
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::new();
        self.write_to(&mut buf).expect("must succeed");
        buf
    }
    pub fn write_to<W: Write>(&self, writer: &mut W) -> Result<()> {
        let mut temp_writer = io::Cursor::new(vec![0; 20]);
        let message_type = Type {
            class: self.class,
            method: self.method,
        };
        track_try!(temp_writer.write_u16be(message_type.as_u16()));
        track_try!(temp_writer.write_u16be(0)); // dummy length
        track_try!(temp_writer.write_u32be(MAGIC_COOKIE));
        track_try!(temp_writer.write_all(&self.transaction_id));
        for attr in self.attributes.iter() {
            track_try!(attr.write_to(&mut temp_writer));
        }
        let attrs_len = temp_writer.get_ref().len() - 20;
        track_assert!(attrs_len < 0x10000,
                      ErrorKind::Failed,
                      "Too large message length: actual={}, limit=0xFFFF",
                      attrs_len);
        temp_writer.set_position(2);
        track_try!(temp_writer.write_u16be(attrs_len as u16));

        let buf = temp_writer.into_inner();
        track_try!(writer.write_all(&buf));
        Ok(())
    }
    pub fn read_from<R: Read>(reader: &mut R) -> Result<Self> {
        let message_type = track_try!(reader.read_u16be());
        let message_type = track_try!(Type::from_u16(message_type));
        let message_len = track_try!(reader.read_u16be());
        track_assert!(message_len % 4 == 0,
                      ErrorKind::NotStunMessage,
                      "Unexpected message length: {} % 4 != 0",
                      message_len);
        let magic_cookie = track_try!(reader.read_u32be());
        track_assert!(magic_cookie == MAGIC_COOKIE,
                      ErrorKind::NotStunMessage,
                      "Unexpected magic cookie: actual={}, expected={}",
                      magic_cookie,
                      MAGIC_COOKIE);;

        let mut transaction_id: [u8; 12] = [0; 12];
        track_try!(reader.read_exact(&mut transaction_id));

        let mut attrs = Vec::new();
        let mut reader = reader.take(message_len as u64);
        while reader.limit() > 0 {
            let attr = track_try!(RawAttribute::read_from(&mut reader));
            attrs.push(attr);
        }
        Ok(RawMessage::new(message_type.class,
                           message_type.method,
                           transaction_id,
                           attrs))
    }

    pub fn try_into_request<M, A>(self) -> Result<Request<M, A>>
        where M: Method,
              A: Attribute
    {
        track_assert_eq!(self.class, Class::Request, ErrorKind::Other);
        let method = track_try!(M::from_u12(self.method).ok_or(ErrorKind::Other));
        let mut attrs = Vec::new();
        for a in self.attributes.iter() {
            let a = track_try!(A::decode(&a, &self));
            attrs.push(a);
        }
        Ok(Request {
            method: method,
            transaction_id: self.transaction_id,
            attributes: attrs,
        })
    }
    pub fn try_into_indication<M, A>(self) -> Result<Indication<M, A>>
        where M: Method,
              A: Attribute
    {
        track_assert_eq!(self.class, Class::Indication, ErrorKind::Other);
        let method = track_try!(M::from_u12(self.method).ok_or(ErrorKind::Other));
        let mut attrs = Vec::new();
        for a in self.attributes.iter() {
            let a = track_try!(A::decode(&a, &self));
            attrs.push(a);
        }
        Ok(Indication {
            method: method,
            transaction_id: self.transaction_id,
            attributes: attrs,
        })
    }
    pub fn try_into_response<M, A>(self) -> Result<Response<M, A>>
        where M: Method,
              A: Attribute
    {
        track_assert!(self.class == Class::SuccessResponse || self.class == Class::ErrorResponse,
                      ErrorKind::Other);
        let method = track_try!(M::from_u12(self.method).ok_or(ErrorKind::Other));
        let mut attrs = Vec::new();
        for a in self.attributes.iter() {
            let a = track_try!(A::decode(&a, &self));
            attrs.push(a);
        }
        if self.class == Class::SuccessResponse {
            Ok(Ok(SuccessResponse {
                method: method,
                transaction_id: self.transaction_id,
                attributes: attrs,
            }))
        } else {
            panic!()
            // TODO: handle error code attribute
            // Ok(Err(ErrorResponse {
            //     method: method,
            //     transaction_id: self.transaction_id,
            //     attributes: Vec::new(),
            // }))
        }
    }
    pub fn try_from_request<M, A>(from: Request<M, A>) -> Result<Self>
        where M: Method,
              A: Attribute
    {
        track_err!(Self::try_from(Class::Request,
                                  from.method(),
                                  *from.transaction_id(),
                                  from.attributes()))
    }
    pub fn try_from_indication<M, A>(from: Indication<M, A>) -> Result<Self>
        where M: Method,
              A: Attribute
    {
        track_err!(Self::try_from(Class::Indication,
                                  from.method(),
                                  *from.transaction_id(),
                                  from.attributes()))
    }
    pub fn try_from_success_response<M, A>(from: SuccessResponse<M, A>) -> Result<Self>
        where M: Method,
              A: Attribute
    {
        track_err!(Self::try_from(Class::SuccessResponse,
                                  from.method(),
                                  *from.transaction_id(),
                                  from.attributes()))
    }
    pub fn try_from_error_response<M, A>(from: ErrorResponse<M, A>) -> Result<Self>
        where M: Method,
              A: Attribute
    {
        track_err!(Self::try_from(Class::ErrorResponse,
                                  from.method(),
                                  *from.transaction_id(),
                                  from.attributes()))
    }

    fn try_from<M, A>(class: Class,
                      method: &M,
                      transaction_id: TransactionId,
                      attributes: &[A])
                      -> Result<Self>
        where M: Method,
              A: Attribute
    {
        let mut m = RawMessage::new(class, method.as_u12(), transaction_id, Vec::new());
        for attr in attributes {
            let raw_attr = track_try!(attr.encode(&m));
            m.attributes.push(raw_attr);
        }
        Ok(m)
    }
}

#[derive(Debug, Clone, Copy, PartialOrd, Ord, PartialEq, Eq, Hash)]
pub enum Class {
    Request = 0b00,
    Indication = 0b01,
    SuccessResponse = 0b10,
    ErrorResponse = 0b11,
}
impl Class {
    pub fn from_u8(value: u8) -> Option<Self> {
        match value {
            0b00 => Some(Class::Request),
            0b01 => Some(Class::Indication),
            0b10 => Some(Class::SuccessResponse),
            0b11 => Some(Class::ErrorResponse),
            _ => None,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct Type {
    pub class: Class,
    pub method: U12,
}
impl Type {
    /// TODO:
    ///
    /// ```text
    /// 0                 1
    /// 2  3  4 5 6 7 8 9 0 1 2 3 4 5
    ///
    /// +--+--+-+-+-+-+-+-+-+-+-+-+-+-+
    /// |M |M |M|M|M|C|M|M|M|C|M|M|M|M|
    /// |11|10|9|8|7|1|6|5|4|0|3|2|1|0|
    /// +--+--+-+-+-+-+-+-+-+-+-+-+-+-+
    ///
    /// Figure 3: Format of STUN Message Type Field
    /// ```
    pub fn as_u16(&self) -> u16 {
        let class = self.class as u16;
        let method = self.method.as_u12().as_u16();
        ((method & 0b0000_0000_1111) << 0) | ((class & 0b01) << 4) |
        ((method & 0b0000_0111_0000) << 5) | ((class & 0b10) << 7) |
        ((method & 0b1111_1000_0000) << 9)
    }

    pub fn from_u16(value: u16) -> Result<Self> {
        track_assert!(value >> 14 == 0,
                      ErrorKind::NotStunMessage,
                      "First two-bits of STUN message must be 0");
        let class = ((value >> 4) & 0b01) | ((value >> 7) & 0b10);
        let class = Class::from_u8(class as u8).unwrap();
        let method = (value & 0b0000_0000_1111) | ((value >> 1) & 0b0000_0111_0000) |
                     ((value >> 2) & 0b1111_1000_0000);
        let method = U12::from_u16(method).unwrap();
        Ok(Type {
            class: class,
            method: method,
        })
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]

use std::mem;
use std::ptr;
use libc::{c_int, c_ulong};

use runtime::Class;
use {Id, Message};

#[allow(improper_ctypes)]
#[link(name = "Foundation", kind = "framework")]
extern {
    static _NSConcreteStackBlock: Class;
}

pub type BlockInvoke<A, R> =
    unsafe extern fn(*mut Block<A, R>, ...) -> R;

pub type ConcreteBlockInvoke<A, R, C> =
    unsafe extern fn(*mut ConcreteBlock<A, R, C>, ...) -> R;

pub trait BlockArguments {
    fn call_block<R>(self, block: &Block<Self, R>) -> R;

    fn invoke_for_concrete_block<R, C: Clone>() -> ConcreteBlockInvoke<Self, R, C>;
}

impl BlockArguments for () {
    fn call_block<R>(self, block: &Block<(), R>) -> R {
        let invoke: unsafe extern fn(*mut Block<(), R>) -> R = unsafe {
            mem::transmute(block.invoke)
        };
        let block_ptr = block as *const _ as *mut _;
        unsafe {
            invoke(block_ptr)
        }
    }

    fn invoke_for_concrete_block<R, C: Clone>() -> ConcreteBlockInvoke<(), R, C> {
        unsafe extern fn concrete_block_invoke_args0<R, C: Clone>(
                block_ptr: *mut ConcreteBlock<(), R, C>) -> R {
            let block = &*block_ptr;
            (block.rust_invoke)(&block.context, ())
        }

        unsafe {
            mem::transmute(concrete_block_invoke_args0::<R, C>)
        }
    }
}

#[repr(C)]
pub struct Block<A: BlockArguments, R> {
    isa: *const Class,
    flags: c_int,
    _reserved: c_int,
    invoke: BlockInvoke<A, R>,
}

impl<A: BlockArguments, R> Block<A, R> {
    pub fn copy(&self) -> Id<Block<A, R>> {
        unsafe {
            let block = msg_send![self copy] as *mut Block<A, R>;
            Id::from_retained_ptr(block)
        }
    }

    pub fn call(&self, args: A) -> R {
        args.call_block(self)
    }
}

impl<A: BlockArguments, R> Message for Block<A, R> { }

#[repr(C)]
struct ConcreteBlock<A: BlockArguments, R, C: Clone> {
    base: Block<A, R>,
    descriptor: *const BlockDescriptor<ConcreteBlock<A, R, C>>,
    rust_invoke: fn (&C, A) -> R,
    context: C,
}

impl<A: BlockArguments, R, C: Clone> ConcreteBlock<A, R, C> {
    pub fn new(invoke: fn (&C, A) -> R, context: C) -> ConcreteBlock<A, R, C> {
        let extern_invoke: ConcreteBlockInvoke<A, R, C> =
            BlockArguments::invoke_for_concrete_block();
        ConcreteBlock {
            base: Block {
                isa: &_NSConcreteStackBlock,
                // 1 << 25 = BLOCK_HAS_COPY_DISPOSE
                flags: 1 << 25,
                _reserved: 0,
                invoke: unsafe { mem::transmute(extern_invoke) },
            },
            // TODO: don't leak memory here
            descriptor: unsafe {
                mem::transmute(box BlockDescriptor::<A, R, C>::new())
            },
            rust_invoke: invoke,
            context: context,
        }
    }
}

impl<A: BlockArguments, R, C: Clone> Clone for ConcreteBlock<A, R, C> {
    fn clone(&self) -> ConcreteBlock<A, R, C> {
        ConcreteBlock::new(self.rust_invoke, self.context.clone())
    }
}

unsafe extern fn block_context_dispose<A: BlockArguments, R, C: Clone>(
        block: &mut ConcreteBlock<A, R, C>) {
    let mut context = mem::uninitialized();
    ptr::copy_nonoverlapping_memory(&mut context, &block.context, 1);
    drop(context);
}

unsafe extern fn block_context_copy<A: BlockArguments, R, C: Clone>(
        dst: &mut ConcreteBlock<A, R, C>, src: &ConcreteBlock<A, R, C>) {
    dst.context = src.context.clone();
}

#[repr(C)]
struct BlockDescriptor<B> {
    _reserved: c_ulong,
    block_size: c_ulong,
    copy_helper: unsafe extern fn(&mut B, &B),
    dispose_helper: unsafe extern fn(&mut B),
}

impl<A: BlockArguments, R, C: Clone> BlockDescriptor<ConcreteBlock<A, R, C>> {
    pub fn new() -> BlockDescriptor<ConcreteBlock<A, R, C>> {
        BlockDescriptor {
            _reserved: 0,
            block_size: mem::size_of::<ConcreteBlock<A, R, C>>() as c_ulong,
            copy_helper: block_context_copy::<A, R, C>,
            dispose_helper: block_context_dispose::<A, R, C>,
        }
    }
}

#[cfg(test)]
mod tests {
    use std::mem;
    use objc_test_utils::invoke_int_block;
    use super::ConcreteBlock;

    #[test]
    fn test_create_block() {
        fn block_get_int(context: &int, _args: ()) -> int {
            *context
        }

        let result = unsafe {
            let block = ConcreteBlock::new(block_get_int, 13i);
            invoke_int_block(mem::transmute(&block))
        };
        assert!(result == 13);
    }
}
<commit_msg>Implemented BlockArguments up to 12-tuple with a macro.<commit_after>#![allow(dead_code)]

use std::mem;
use std::ptr;
use libc::{c_int, c_ulong};

use runtime::Class;
use {Id, Message};

#[allow(improper_ctypes)]
#[link(name = "Foundation", kind = "framework")]
extern {
    static _NSConcreteStackBlock: Class;
}

pub type BlockInvoke<A, R> =
    unsafe extern fn(*mut Block<A, R>, ...) -> R;

pub type ConcreteBlockInvoke<A, R, C> =
    unsafe extern fn(*mut ConcreteBlock<A, R, C>, ...) -> R;

pub trait BlockArguments {
    fn call_block<R>(self, block: &Block<Self, R>) -> R;

    fn invoke_for_concrete_block<R, C: Clone>() -> ConcreteBlockInvoke<Self, R, C>;
}

macro_rules! block_args_impl(
    ($f:ident $(, $a:ident : $t:ident)*) => (
        impl<$($t),*> BlockArguments for ($($t,)*) {
            fn call_block<R>(self, block: &Block<($($t,)*), R>) -> R {
                let invoke: unsafe extern fn(*mut Block<($($t,)*), R> $(, $t)*) -> R = unsafe {
                    mem::transmute(block.invoke)
                };
                let ($($a,)*) = self;
                let block_ptr = block as *const _ as *mut _;
                unsafe {
                    invoke(block_ptr $(, $a)*)
                }
            }

            fn invoke_for_concrete_block<R, X: Clone>() ->
                    ConcreteBlockInvoke<($($t,)*), R, X> {
                unsafe extern fn $f<R, X: Clone $(, $t)*>(
                        block_ptr: *mut ConcreteBlock<($($t,)*), R, X>
                        $(, $a: $t)*) -> R {
                    let args = ($($a,)*);
                    let block = &*block_ptr;
                    (block.rust_invoke)(&block.context, args)
                }

                unsafe {
                    mem::transmute($f::<R, X $(, $t)*>)
                }
            }
        }
    );
);

block_args_impl!(concrete_block_invoke_args0);
block_args_impl!(concrete_block_invoke_args1, a: A);
block_args_impl!(concrete_block_invoke_args2, a: A, b: B);
block_args_impl!(concrete_block_invoke_args3, a: A, b: B, c: C);
block_args_impl!(concrete_block_invoke_args4, a: A, b: B, c: C, d: D);
block_args_impl!(concrete_block_invoke_args5, a: A, b: B, c: C, d: D, e: E);
block_args_impl!(concrete_block_invoke_args6, a: A, b: B, c: C, d: D, e: E, f: F);
block_args_impl!(concrete_block_invoke_args7, a: A, b: B, c: C, d: D, e: E, f: F, g: G);
block_args_impl!(concrete_block_invoke_args8, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H);
block_args_impl!(concrete_block_invoke_args9, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I);
block_args_impl!(concrete_block_invoke_args10, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J);
block_args_impl!(concrete_block_invoke_args11, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K);
block_args_impl!(concrete_block_invoke_args12, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K, l: L);

#[repr(C)]
pub struct Block<A: BlockArguments, R> {
    isa: *const Class,
    flags: c_int,
    _reserved: c_int,
    invoke: BlockInvoke<A, R>,
}

impl<A: BlockArguments, R> Block<A, R> {
    pub fn copy(&self) -> Id<Block<A, R>> {
        unsafe {
            let block = msg_send![self copy] as *mut Block<A, R>;
            Id::from_retained_ptr(block)
        }
    }

    pub fn call(&self, args: A) -> R {
        args.call_block(self)
    }
}

impl<A: BlockArguments, R> Message for Block<A, R> { }

#[repr(C)]
struct ConcreteBlock<A: BlockArguments, R, C: Clone> {
    base: Block<A, R>,
    descriptor: *const BlockDescriptor<ConcreteBlock<A, R, C>>,
    rust_invoke: fn (&C, A) -> R,
    context: C,
}

impl<A: BlockArguments, R, C: Clone> ConcreteBlock<A, R, C> {
    pub fn new(invoke: fn (&C, A) -> R, context: C) -> ConcreteBlock<A, R, C> {
        let extern_invoke: ConcreteBlockInvoke<A, R, C> =
            BlockArguments::invoke_for_concrete_block();
        ConcreteBlock {
            base: Block {
                isa: &_NSConcreteStackBlock,
                // 1 << 25 = BLOCK_HAS_COPY_DISPOSE
                flags: 1 << 25,
                _reserved: 0,
                invoke: unsafe { mem::transmute(extern_invoke) },
            },
            // TODO: don't leak memory here
            descriptor: unsafe {
                mem::transmute(box BlockDescriptor::<A, R, C>::new())
            },
            rust_invoke: invoke,
            context: context,
        }
    }
}

impl<A: BlockArguments, R, C: Clone> Clone for ConcreteBlock<A, R, C> {
    fn clone(&self) -> ConcreteBlock<A, R, C> {
        ConcreteBlock::new(self.rust_invoke, self.context.clone())
    }
}

unsafe extern fn block_context_dispose<A: BlockArguments, R, C: Clone>(
        block: &mut ConcreteBlock<A, R, C>) {
    let mut context = mem::uninitialized();
    ptr::copy_nonoverlapping_memory(&mut context, &block.context, 1);
    drop(context);
}

unsafe extern fn block_context_copy<A: BlockArguments, R, C: Clone>(
        dst: &mut ConcreteBlock<A, R, C>, src: &ConcreteBlock<A, R, C>) {
    dst.context = src.context.clone();
}

#[repr(C)]
struct BlockDescriptor<B> {
    _reserved: c_ulong,
    block_size: c_ulong,
    copy_helper: unsafe extern fn(&mut B, &B),
    dispose_helper: unsafe extern fn(&mut B),
}

impl<A: BlockArguments, R, C: Clone> BlockDescriptor<ConcreteBlock<A, R, C>> {
    pub fn new() -> BlockDescriptor<ConcreteBlock<A, R, C>> {
        BlockDescriptor {
            _reserved: 0,
            block_size: mem::size_of::<ConcreteBlock<A, R, C>>() as c_ulong,
            copy_helper: block_context_copy::<A, R, C>,
            dispose_helper: block_context_dispose::<A, R, C>,
        }
    }
}

#[cfg(test)]
mod tests {
    use std::mem;
    use objc_test_utils::invoke_int_block;
    use super::ConcreteBlock;

    #[test]
    fn test_create_block() {
        fn block_get_int(context: &int, _args: ()) -> int {
            *context
        }

        let result = unsafe {
            let block = ConcreteBlock::new(block_get_int, 13i);
            invoke_int_block(mem::transmute(&block))
        };
        assert!(result == 13);
    }
}
<|endoftext|>


<commit_before>//! A Rust interface for Objective-C blocks.
//!
//! For more information on the specifics of the block implementation, see
//! Clang's documentation: http://clang.llvm.org/docs/Block-ABI-Apple.html

use std::mem;
use std::ptr;
use libc::{c_int, c_ulong};

use runtime::Class;
use {Id, Message};

#[allow(improper_ctypes)]
#[link(name = "Foundation", kind = "framework")]
extern {
    static _NSConcreteStackBlock: Class;
}

/// An invoke function for a `Block`; this is the raw C function called by the
/// Objective-C runtime.
type BlockInvoke<A, R> = unsafe extern fn(*mut Block<A, R>, ...) -> R;

/// An invoke function for a `ConcreteBlock`; this is the raw C function called
/// by the Objective-C runtime.
pub type ConcreteBlockInvoke<A, R, C> =
    unsafe extern fn(*mut ConcreteBlock<A, R, C>, ...) -> R;

/// Types that may be used as the arguments to an Objective-C block.
pub trait BlockArguments {
    /// Calls the given `Block` with self as the arguments.
    fn call_block<R>(self, block: &Block<Self, R>) -> R;

    /// Returns an invoke function for a `ConcreteBlock` that takes this type
    /// of arguments.
    fn invoke_for_concrete_block<R, C: Clone>() -> ConcreteBlockInvoke<Self, R, C>;
}

macro_rules! block_args_impl(
    ($f:ident $(, $a:ident : $t:ident)*) => (
        impl<$($t),*> BlockArguments for ($($t,)*) {
            fn call_block<R>(self, block: &Block<($($t,)*), R>) -> R {
                let invoke: unsafe extern fn(*mut Block<($($t,)*), R> $(, $t)*) -> R = unsafe {
                    mem::transmute(block.invoke)
                };
                let ($($a,)*) = self;
                let block_ptr = block as *const _ as *mut _;
                unsafe {
                    invoke(block_ptr $(, $a)*)
                }
            }

            fn invoke_for_concrete_block<R, X: Clone>() ->
                    ConcreteBlockInvoke<($($t,)*), R, X> {
                unsafe extern fn $f<R, X: Clone $(, $t)*>(
                        block_ptr: *mut ConcreteBlock<($($t,)*), R, X>
                        $(, $a: $t)*) -> R {
                    let args = ($($a,)*);
                    let block = &*block_ptr;
                    (block.rust_invoke)(&block.context, args)
                }

                unsafe {
                    mem::transmute($f::<R, X $(, $t)*>)
                }
            }
        }
    );
);

block_args_impl!(concrete_block_invoke_args0);
block_args_impl!(concrete_block_invoke_args1, a: A);
block_args_impl!(concrete_block_invoke_args2, a: A, b: B);
block_args_impl!(concrete_block_invoke_args3, a: A, b: B, c: C);
block_args_impl!(concrete_block_invoke_args4, a: A, b: B, c: C, d: D);
block_args_impl!(concrete_block_invoke_args5, a: A, b: B, c: C, d: D, e: E);
block_args_impl!(concrete_block_invoke_args6, a: A, b: B, c: C, d: D, e: E, f: F);
block_args_impl!(concrete_block_invoke_args7, a: A, b: B, c: C, d: D, e: E, f: F, g: G);
block_args_impl!(concrete_block_invoke_args8, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H);
block_args_impl!(concrete_block_invoke_args9, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I);
block_args_impl!(concrete_block_invoke_args10, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J);
block_args_impl!(concrete_block_invoke_args11, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K);
block_args_impl!(concrete_block_invoke_args12, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K, l: L);

/// An Objective-C block that takes arguments of `A` when called and
/// returns a value of `R`.
#[repr(C)]
pub struct Block<A: BlockArguments, R> {
    isa: *const Class,
    flags: c_int,
    _reserved: c_int,
    invoke: BlockInvoke<A, R>,
}

impl<A: BlockArguments, R> Block<A, R> {
    /// Copy self onto the heap.
    pub fn copy(&self) -> Id<Block<A, R>> {
        unsafe {
            let block = msg_send![self copy] as *mut Block<A, R>;
            Id::from_retained_ptr(block)
        }
    }

    /// Call self with the given arguments.
    pub fn call(&self, args: A) -> R {
        args.call_block(self)
    }
}

impl<A: BlockArguments, R> Message for Block<A, R> { }

/// An Objective-C block whose size is known at compile time and may be
/// constructed on the stack.
#[repr(C)]
pub struct ConcreteBlock<A: BlockArguments, R, C> {
    base: Block<A, R>,
    descriptor: Box<BlockDescriptor<ConcreteBlock<A, R, C>>>,
    rust_invoke: fn (&C, A) -> R,
    context: C,
}

impl<A: BlockArguments, R, C: Clone> ConcreteBlock<A, R, C> {
    /// Constructs a `ConcreteBlock` with the given invoke function and context.
    /// When the block is called, it will return the value that results from
    /// calling the invoke function with a reference to its context.
    pub fn new(invoke: fn (&C, A) -> R, context: C) -> ConcreteBlock<A, R, C> {
        let extern_invoke: ConcreteBlockInvoke<A, R, C> =
            BlockArguments::invoke_for_concrete_block();
        ConcreteBlock {
            base: Block {
                isa: &_NSConcreteStackBlock,
                // 1 << 25 = BLOCK_HAS_COPY_DISPOSE
                flags: 1 << 25,
                _reserved: 0,
                invoke: unsafe { mem::transmute(extern_invoke) },
            },
            descriptor: box BlockDescriptor::<A, R, C>::new(),
            rust_invoke: invoke,
            context: context,
        }
    }
}

impl<A: BlockArguments, R, C: Clone> Clone for ConcreteBlock<A, R, C> {
    fn clone(&self) -> ConcreteBlock<A, R, C> {
        ConcreteBlock::new(self.rust_invoke, self.context.clone())
    }
}

impl<A: BlockArguments, R, C> Deref<Block<A, R>> for ConcreteBlock<A, R, C> {
    fn deref(&self) -> &Block<A, R> {
        &self.base
    }
}

unsafe extern fn block_context_dispose<A: BlockArguments, R, C>(
        block: &mut ConcreteBlock<A, R, C>) {
    // Read the block onto the stack and let it drop
    ptr::read(block);
}

unsafe extern fn block_context_copy<A: BlockArguments, R, C: Clone>(
        dst: &mut ConcreteBlock<A, R, C>, src: &ConcreteBlock<A, R, C>) {
    // Doesn't seem like the descriptor is supposed to change in this function,
    // but our descriptor isn't static (that's hard), so we just clone the box.
    ptr::write(&mut dst.descriptor, src.descriptor.clone());
    // The src block actually gets memmoved to the destination beforehand,
    // but we'll set the function pointer, too, to be safe.
    ptr::write(&mut dst.rust_invoke, src.rust_invoke);
    ptr::write(&mut dst.context, src.context.clone());
}

#[repr(C)]
struct BlockDescriptor<B> {
    _reserved: c_ulong,
    block_size: c_ulong,
    copy_helper: unsafe extern fn(&mut B, &B),
    dispose_helper: unsafe extern fn(&mut B),
}

impl<A: BlockArguments, R, C: Clone> BlockDescriptor<ConcreteBlock<A, R, C>> {
    fn new() -> BlockDescriptor<ConcreteBlock<A, R, C>> {
        BlockDescriptor {
            _reserved: 0,
            block_size: mem::size_of::<ConcreteBlock<A, R, C>>() as c_ulong,
            copy_helper: block_context_copy::<A, R, C>,
            dispose_helper: block_context_dispose::<A, R, C>,
        }
    }
}

impl<B> Copy for BlockDescriptor<B> { }

impl<B> Clone for BlockDescriptor<B> {
    fn clone(&self) -> BlockDescriptor<B> {
        *self
    }
}

#[cfg(test)]
mod tests {
    use Id;
    use objc_test_utils;
    use super::{Block, ConcreteBlock};

    fn get_int_block() -> &'static Block<(), int> {
        unsafe {
            &*(objc_test_utils::get_int_block() as *const _)
        }
    }

    fn get_int_block_with(i: int) -> Id<Block<(), int>> {
        unsafe {
            let ptr = objc_test_utils::get_int_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn get_add_block() -> &'static Block<(int,), int> {
        unsafe {
            &*(objc_test_utils::get_add_block() as *const _)
        }
    }

    fn get_add_block_with(i: int) -> Id<Block<(int,), int>> {
        unsafe {
            let ptr = objc_test_utils::get_add_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn invoke_int_block(block: &Block<(), int>) -> int {
        let ptr = block as *const _;
        unsafe {
            objc_test_utils::invoke_int_block(ptr as *const _)
        }
    }

    fn invoke_add_block(block: &Block<(int,), int>, a: int) -> int {
        let ptr = block as *const _;
        unsafe {
            objc_test_utils::invoke_add_block(ptr as *const _, a)
        }
    }

    #[test]
    fn test_call_block() {
        let block = get_int_block();
        assert!(block.call(()) == 7);

        let block = get_int_block_with(13);
        assert!(block.call(()) == 13);
    }

    #[test]
    fn test_call_block_args() {
        let block = get_add_block();
        assert!(block.call((2,)) == 9);

        let block = get_add_block_with(13);
        assert!(block.call((2,)) == 15);
    }

    #[test]
    fn test_create_block() {
        fn block_get_int(context: &int, _args: ()) -> int {
            *context
        }

        let block = ConcreteBlock::new(block_get_int, 13);
        let result = invoke_int_block(&*block);
        assert!(result == 13);
    }

    #[test]
    fn test_create_block_args() {
        fn block_add_int(context: &int, (a,): (int,)) -> int {
            a + *context
        }

        let block = ConcreteBlock::new(block_add_int, 5);
        let result = invoke_add_block(&*block, 6);
        assert!(result == 11);
    }

    #[test]
    fn test_concrete_block_clone() {
        fn block_get_string_len(context: &String, _args: ()) -> uint {
            context.len()
        }

        let s = "Hello!".into_string();
        let expected_len = s.len();
        let block = ConcreteBlock::new(block_get_string_len, s);
        assert!(block.call(()) == expected_len);

        let cloned = block.clone();
        assert!(cloned.call(()) == expected_len);

        drop(block);
        assert!(cloned.call(()) == expected_len);
    }

    #[test]
    fn test_concrete_block_copy() {
        fn block_get_string_len(context: &String, _args: ()) -> uint {
            context.len()
        }

        let s = "Hello!".into_string();
        let expected_len = s.len();
        let block = ConcreteBlock::new(block_get_string_len, s);
        assert!(block.call(()) == expected_len);

        let copied = block.copy();
        assert!(copied.call(()) == expected_len);

        drop(block);
        assert!(copied.call(()) == expected_len);
    }
}
<commit_msg>Move self in copy method and only impl it for ConcreteBlocks.<commit_after>//! A Rust interface for Objective-C blocks.
//!
//! For more information on the specifics of the block implementation, see
//! Clang's documentation: http://clang.llvm.org/docs/Block-ABI-Apple.html

use std::mem;
use std::ptr;
use libc::{c_int, c_ulong};

use runtime::Class;
use {Id, Message};

#[allow(improper_ctypes)]
#[link(name = "Foundation", kind = "framework")]
extern {
    static _NSConcreteStackBlock: Class;
}

/// An invoke function for a `Block`; this is the raw C function called by the
/// Objective-C runtime.
type BlockInvoke<A, R> = unsafe extern fn(*mut Block<A, R>, ...) -> R;

/// An invoke function for a `ConcreteBlock`; this is the raw C function called
/// by the Objective-C runtime.
pub type ConcreteBlockInvoke<A, R, C> =
    unsafe extern fn(*mut ConcreteBlock<A, R, C>, ...) -> R;

/// Types that may be used as the arguments to an Objective-C block.
pub trait BlockArguments {
    /// Calls the given `Block` with self as the arguments.
    fn call_block<R>(self, block: &Block<Self, R>) -> R;

    /// Returns an invoke function for a `ConcreteBlock` that takes this type
    /// of arguments.
    fn invoke_for_concrete_block<R, C: Clone>() -> ConcreteBlockInvoke<Self, R, C>;
}

macro_rules! block_args_impl(
    ($f:ident $(, $a:ident : $t:ident)*) => (
        impl<$($t),*> BlockArguments for ($($t,)*) {
            fn call_block<R>(self, block: &Block<($($t,)*), R>) -> R {
                let invoke: unsafe extern fn(*mut Block<($($t,)*), R> $(, $t)*) -> R = unsafe {
                    mem::transmute(block.invoke)
                };
                let ($($a,)*) = self;
                let block_ptr = block as *const _ as *mut _;
                unsafe {
                    invoke(block_ptr $(, $a)*)
                }
            }

            fn invoke_for_concrete_block<R, X: Clone>() ->
                    ConcreteBlockInvoke<($($t,)*), R, X> {
                unsafe extern fn $f<R, X: Clone $(, $t)*>(
                        block_ptr: *mut ConcreteBlock<($($t,)*), R, X>
                        $(, $a: $t)*) -> R {
                    let args = ($($a,)*);
                    let block = &*block_ptr;
                    (block.rust_invoke)(&block.context, args)
                }

                unsafe {
                    mem::transmute($f::<R, X $(, $t)*>)
                }
            }
        }
    );
);

block_args_impl!(concrete_block_invoke_args0);
block_args_impl!(concrete_block_invoke_args1, a: A);
block_args_impl!(concrete_block_invoke_args2, a: A, b: B);
block_args_impl!(concrete_block_invoke_args3, a: A, b: B, c: C);
block_args_impl!(concrete_block_invoke_args4, a: A, b: B, c: C, d: D);
block_args_impl!(concrete_block_invoke_args5, a: A, b: B, c: C, d: D, e: E);
block_args_impl!(concrete_block_invoke_args6, a: A, b: B, c: C, d: D, e: E, f: F);
block_args_impl!(concrete_block_invoke_args7, a: A, b: B, c: C, d: D, e: E, f: F, g: G);
block_args_impl!(concrete_block_invoke_args8, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H);
block_args_impl!(concrete_block_invoke_args9, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I);
block_args_impl!(concrete_block_invoke_args10, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J);
block_args_impl!(concrete_block_invoke_args11, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K);
block_args_impl!(concrete_block_invoke_args12, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K, l: L);

/// An Objective-C block that takes arguments of `A` when called and
/// returns a value of `R`.
#[repr(C)]
pub struct Block<A: BlockArguments, R> {
    isa: *const Class,
    flags: c_int,
    _reserved: c_int,
    invoke: BlockInvoke<A, R>,
}

impl<A: BlockArguments, R> Block<A, R> {
    /// Call self with the given arguments.
    pub fn call(&self, args: A) -> R {
        args.call_block(self)
    }
}

impl<A: BlockArguments, R> Message for Block<A, R> { }

/// An Objective-C block whose size is known at compile time and may be
/// constructed on the stack.
#[repr(C)]
pub struct ConcreteBlock<A: BlockArguments, R, C> {
    base: Block<A, R>,
    descriptor: Box<BlockDescriptor<ConcreteBlock<A, R, C>>>,
    rust_invoke: fn (&C, A) -> R,
    context: C,
}

impl<A: BlockArguments, R, C: Clone> ConcreteBlock<A, R, C> {
    /// Constructs a `ConcreteBlock` with the given invoke function and context.
    /// When the block is called, it will return the value that results from
    /// calling the invoke function with a reference to its context.
    pub fn new(invoke: fn (&C, A) -> R, context: C) -> ConcreteBlock<A, R, C> {
        let extern_invoke: ConcreteBlockInvoke<A, R, C> =
            BlockArguments::invoke_for_concrete_block();
        ConcreteBlock {
            base: Block {
                isa: &_NSConcreteStackBlock,
                // 1 << 25 = BLOCK_HAS_COPY_DISPOSE
                flags: 1 << 25,
                _reserved: 0,
                invoke: unsafe { mem::transmute(extern_invoke) },
            },
            descriptor: box BlockDescriptor::<A, R, C>::new(),
            rust_invoke: invoke,
            context: context,
        }
    }
}

impl<A: BlockArguments, R, C> ConcreteBlock<A, R, C> {
    /// Copy self onto the heap.
    pub fn copy(self) -> Id<Block<A, R>> {
        unsafe {
            let block = msg_send![&self.base copy] as *mut Block<A, R>;
            Id::from_retained_ptr(block)
        }
    }
}

impl<A: BlockArguments, R, C: Clone> Clone for ConcreteBlock<A, R, C> {
    fn clone(&self) -> ConcreteBlock<A, R, C> {
        ConcreteBlock::new(self.rust_invoke, self.context.clone())
    }
}

impl<A: BlockArguments, R, C> Deref<Block<A, R>> for ConcreteBlock<A, R, C> {
    fn deref(&self) -> &Block<A, R> {
        &self.base
    }
}

unsafe extern fn block_context_dispose<A: BlockArguments, R, C>(
        block: &mut ConcreteBlock<A, R, C>) {
    // Read the block onto the stack and let it drop
    ptr::read(block);
}

unsafe extern fn block_context_copy<A: BlockArguments, R, C: Clone>(
        dst: &mut ConcreteBlock<A, R, C>, src: &ConcreteBlock<A, R, C>) {
    // Doesn't seem like the descriptor is supposed to change in this function,
    // but our descriptor isn't static (that's hard), so we just clone the box.
    ptr::write(&mut dst.descriptor, src.descriptor.clone());
    // The src block actually gets memmoved to the destination beforehand,
    // but we'll set the function pointer, too, to be safe.
    ptr::write(&mut dst.rust_invoke, src.rust_invoke);
    ptr::write(&mut dst.context, src.context.clone());
}

#[repr(C)]
struct BlockDescriptor<B> {
    _reserved: c_ulong,
    block_size: c_ulong,
    copy_helper: unsafe extern fn(&mut B, &B),
    dispose_helper: unsafe extern fn(&mut B),
}

impl<A: BlockArguments, R, C: Clone> BlockDescriptor<ConcreteBlock<A, R, C>> {
    fn new() -> BlockDescriptor<ConcreteBlock<A, R, C>> {
        BlockDescriptor {
            _reserved: 0,
            block_size: mem::size_of::<ConcreteBlock<A, R, C>>() as c_ulong,
            copy_helper: block_context_copy::<A, R, C>,
            dispose_helper: block_context_dispose::<A, R, C>,
        }
    }
}

impl<B> Copy for BlockDescriptor<B> { }

impl<B> Clone for BlockDescriptor<B> {
    fn clone(&self) -> BlockDescriptor<B> {
        *self
    }
}

#[cfg(test)]
mod tests {
    use Id;
    use objc_test_utils;
    use super::{Block, ConcreteBlock};

    fn get_int_block() -> &'static Block<(), int> {
        unsafe {
            &*(objc_test_utils::get_int_block() as *const _)
        }
    }

    fn get_int_block_with(i: int) -> Id<Block<(), int>> {
        unsafe {
            let ptr = objc_test_utils::get_int_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn get_add_block() -> &'static Block<(int,), int> {
        unsafe {
            &*(objc_test_utils::get_add_block() as *const _)
        }
    }

    fn get_add_block_with(i: int) -> Id<Block<(int,), int>> {
        unsafe {
            let ptr = objc_test_utils::get_add_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn invoke_int_block(block: &Block<(), int>) -> int {
        let ptr = block as *const _;
        unsafe {
            objc_test_utils::invoke_int_block(ptr as *const _)
        }
    }

    fn invoke_add_block(block: &Block<(int,), int>, a: int) -> int {
        let ptr = block as *const _;
        unsafe {
            objc_test_utils::invoke_add_block(ptr as *const _, a)
        }
    }

    #[test]
    fn test_call_block() {
        let block = get_int_block();
        assert!(block.call(()) == 7);

        let block = get_int_block_with(13);
        assert!(block.call(()) == 13);
    }

    #[test]
    fn test_call_block_args() {
        let block = get_add_block();
        assert!(block.call((2,)) == 9);

        let block = get_add_block_with(13);
        assert!(block.call((2,)) == 15);
    }

    #[test]
    fn test_create_block() {
        fn block_get_int(context: &int, _args: ()) -> int {
            *context
        }

        let block = ConcreteBlock::new(block_get_int, 13);
        let result = invoke_int_block(&*block);
        assert!(result == 13);
    }

    #[test]
    fn test_create_block_args() {
        fn block_add_int(context: &int, (a,): (int,)) -> int {
            a + *context
        }

        let block = ConcreteBlock::new(block_add_int, 5);
        let result = invoke_add_block(&*block, 6);
        assert!(result == 11);
    }

    #[test]
    fn test_concrete_block_clone() {
        fn block_get_string_len(context: &String, _args: ()) -> uint {
            context.len()
        }

        let s = "Hello!".into_string();
        let expected_len = s.len();
        let block = ConcreteBlock::new(block_get_string_len, s);
        assert!(block.call(()) == expected_len);

        let cloned = block.clone();
        assert!(cloned.call(()) == expected_len);

        drop(block);
        assert!(cloned.call(()) == expected_len);
    }

    #[test]
    fn test_concrete_block_copy() {
        fn block_get_string_len(context: &String, _args: ()) -> uint {
            context.len()
        }

        let s = "Hello!".into_string();
        let expected_len = s.len();
        let block = ConcreteBlock::new(block_get_string_len, s);
        assert!(block.call(()) == expected_len);

        let copied = block.copy();
        assert!(copied.call(()) == expected_len);
    }
}
<|endoftext|>


<commit_before>use std::convert::From;
use std::error::Error;
use std::fmt;
use std::num::ParseIntError;

use ansi_term::{ANSIString, Color, Style};

use util::Shell;

/// Representation of config options that all modules have
#[derive(Clone)]
pub struct ModuleOptions {
    /// String to display to the left of the content
    pub padding_left: String,
    /// String to display to the right of the content
    pub padding_right: String,
    /// String to print out after the content and right padding
    pub separator: String,
    /// Background color, foreground color, etc.
    pub style: ModuleStyle,
}

/// Representation of how to style a module
#[derive(Clone)]
pub struct ModuleStyle {
    /// Color behind the text
    pub background: Option<Color>,
    /// Color of the text
    pub foreground: Option<Color>,
    /// Optional combination of attributes like bold, italicized, etc.
    pub text_properties: Option<Style>,
}

/// Formats a string with the given `ModuleOptions` for a specific
/// `Shell`.
///
/// # Parameters
///
/// - `s` - the contents of the module to be formatted
/// - `options` - the background, foreground, padding, etc. to apply
/// - `next_bg` - the background color, if any, of the next visible module
/// - `shell` - the type of shell to format the string for
pub fn format_for_module(s: &str,
                         options: ModuleOptions,
                         next_bg: Option<Color>,
                         shell: Shell)
                         -> ANSIString<'static> {

    let style = style_from_modulestyle(&options.style);

    // Each shell keeps track of the number of characters that make up
    // the prompt. The ANSI escape-sequences that color the text will
    // be accidentally included in this length *unless* we prefix and
    // suffix them with these shell-specific escape-sequences. We
    // don't want the shell to mistakenly think there's fewer
    // characters remaining on the current line than there actually
    // are.
    let (len_esc_prefix, len_esc_suffix) = match shell {
        Shell::Bash => ("\\[", "\\]"),
        Shell::Zsh => ("%{", "%}"),
    };

    unimplemented!();
}

/// Converts a `ModuleStyle` into an `ansi_term::Style`.
fn style_from_modulestyle(s: &ModuleStyle) -> Style {
    let mut style = s.text_properties.unwrap_or_default();
    if let Some(bg) = s.background {
        style = style.on(bg);
    }
    if let Some(fg) = s.foreground {
        style = style.fg(fg);
    }
    style
}

// We use methods of the form "try_*_from_*" because the user should
// just be able to specify (in their config) what color/style they
// want, and we have to be able to parse what they say into a
// meaningful type.

/// Attempts to convert a string into an `ansi_term::Color`.
///
/// Returns a `ConvertError::NoSuchMatch` if the provided string
/// doesn't match any of the colors defined in crate `ansi_term`.
///
/// # Examples
/// ```
/// assert_eq!(try_color_from_str("black"), Ok(Color::Black));
/// assert_eq!(try_color_from_str("green"), Ok(Color::Green));
/// assert_eq!(try_color_from_str("turquoise"), Err(ConvertError::NoSuchMatch));
/// ```
pub fn try_color_from_str(s: &str) -> Result<Color, ConvertError> {
    match s {
        "black" => Ok(Color::Black),
        "red" => Ok(Color::Red),
        "green" => Ok(Color::Green),
        "yellow" => Ok(Color::Yellow),
        "blue" => Ok(Color::Blue),
        "purple" => Ok(Color::Purple),
        "cyan" => Ok(Color::Cyan),
        "white" => Ok(Color::White),
        _ => Err(ConvertError::NoSuchMatch),
    }
}

/// Attempts to convert a string into an `ansi_term::Color::RGB`.
///
/// Returns a `ConvertError::InvalidForm` if the provided string is
/// not a sequence of 3 `u8`s, separated by commas.
///
/// # Examples
/// ```
/// assert_eq!(try_rgb_from_str("(14, 76, 1)"), Ok(Color::RGB(14, 76, 1)));
/// assert_eq!(try_rgb_from_str("0, 100, 0"), Ok(Color::RGB(0, 100, 0)));
/// assert_eq!(try_rgb_from_str("1000, b, c, -1"), Err(ConvertError::InvalidForm));
/// ```
pub fn try_rgb_from_str(s: &str) -> Result<Color, ConvertError> {
    // Strip out non-integer characters
    let cleaned: Vec<String> = s.split(',')
        .map(|i| i.replace(|j| j == '(' || j == ')' || j == ' ', ""))
        .collect();

    // We want to immediately return an Error if a conversion fails
    let mut ints: Vec<u8> = Vec::new();
    for i in &cleaned {
        ints.push(i.parse::<u8>()?);
    }

    // Note: somewhat-malformed input, like "()()(0, 0, 0)" can still
    // get through, but we can still make "sense" of it so it's good
    // to go.
    if ints.len() != 3 {
        Err(ConvertError::InvalidForm)
    } else {
        Ok(Color::RGB(ints[0], ints[1], ints[2]))
    }
}

// NOTE: We do not need a try_color_from_u8 OR a try_rgb_from_vec.
// The implementation would just look like:
// pub fn try_color_from_u8(i: u8) -> Color { Color::u8(i) }
// It would always succeed due to Rust's type system.

#[derive(Debug, PartialEq)]
/// Error type for when a conversion due to parsing fails
pub enum ConvertError {
    /// Input doesn't correspond to a valid result
    NoSuchMatch,
    /// Input is malformed and cannot be parsed
    InvalidForm,
}

impl fmt::Display for ConvertError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl Error for ConvertError {
    fn description(&self) -> &str {
        match *self {
            ConvertError::NoSuchMatch => "no match was found for the provided input",
            ConvertError::InvalidForm => "provided input was malformed and could not be parsed",
        }
    }
}

// So that we can use try!() and ? to return early if we encounter a
// ParseIntError
impl From<ParseIntError> for ConvertError {
    fn from(_: ParseIntError) -> ConvertError {
        ConvertError::InvalidForm
    }
}

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

    #[test]
    fn test_style_from_modulestyle() {
        const CONTENT: &'static str = "Hello";

        struct Test {
            style: ModuleStyle,
            expected: String,
        }

        let tests = [// Normal text, no bg, custom fg
                     Test {
                         style: ModuleStyle {
                             background: None,
                             foreground: Some(Color::White),
                             text_properties: None,
                         },
                         expected: format!("\x1B[37m{}\x1B[0m", CONTENT),
                     },
                     // Bold text, no fg, custom bg
                     Test {
                         style: ModuleStyle {
                             background: Some(Color::Blue),
                             foreground: None,
                             text_properties: Some(Style::default().bold()),
                         },
                         expected: format!("\x1B[1;44m{}\x1B[0m", CONTENT),
                     },
                     // Underlined text, custom bg and fg
                     Test {
                         style: ModuleStyle {
                             background: Some(Color::Blue),
                             foreground: Some(Color::White),
                             text_properties: Some(Style::default().underline()),
                         },
                         expected: format!("\x1B[4;44;37m{}\x1B[0m", CONTENT),
                     },
                     // Normal text, no bg nor fg
                     Test {
                         style: ModuleStyle {
                             background: None,
                             foreground: None,
                             text_properties: None,
                         },
                         expected: format!("{}", CONTENT),
                     }];

        for test in &tests {
            let result = format!("{}", style_from_modulestyle(&test.style).paint(CONTENT));
            assert_eq!(result, test.expected);
        }
    }

    #[test]
    fn test_color_try_from_str() {
        assert_eq!(try_color_from_str("blue"), Ok(Color::Blue));
        assert_eq!(try_color_from_str("teal"), Err(ConvertError::NoSuchMatch));
    }

    #[test]
    fn test_rgb_try_from_str() {
        // Valid inputs
        assert_eq!(try_rgb_from_str("(0, 0, 0)"), Ok(Color::RGB(0, 0, 0)));
        assert_eq!(try_rgb_from_str("(255, 255, 255)"),
                   Ok(Color::RGB(255, 255, 255)));

        // Questionable inputs (should still work)
        assert_eq!(try_rgb_from_str("(0, 0, 0))"), Ok(Color::RGB(0, 0, 0)));

        // Improperly formed ("too many" inputs)
        assert_eq!(try_rgb_from_str("(0, 0, 0,)"),
                   Err(ConvertError::InvalidForm));

        // Too few inputs
        assert_eq!(try_rgb_from_str("(0, 0)"), Err(ConvertError::InvalidForm));

        // Inputs aren't u8's
        assert_eq!(try_rgb_from_str("(1000, 0, 0)"),
                   Err(ConvertError::InvalidForm));
        assert_eq!(try_rgb_from_str("(0, 0, -1)"),
                   Err(ConvertError::InvalidForm));
    }
}
<commit_msg>feat: implement format_for_module<commit_after>use std::convert::From;
use std::error::Error;
use std::fmt;
use std::num::ParseIntError;

use ansi_term::{ANSIString, Color, Style};

use util::Shell;

/// Representation of config options that all modules have
#[derive(Clone)]
pub struct ModuleOptions<'a> {
    /// String to display to the left of the content
    pub padding_left: &'a str,
    /// String to display to the right of the content
    pub padding_right: &'a str,
    /// String to print out after the content and right padding
    pub separator: &'a str,
    /// Background color, foreground color, etc.
    pub style: ModuleStyle,
}

/// Representation of how to style a module
#[derive(Clone)]
pub struct ModuleStyle {
    /// Color behind the text
    pub background: Option<Color>,
    /// Color of the text
    pub foreground: Option<Color>,
    /// Combination of attributes like bold, italicized, etc.
    pub text_properties: Option<Style>,
}

/// Formats a string with the given `ModuleOptions` for a specific
/// `Shell`.
///
/// # Parameters
///
/// - `s` - the contents of the module to be formatted
/// - `options` - the background, foreground, padding, etc. to apply
/// - `next_bg` - the background color, if any, of the next visible module
/// - `shell` - the type of shell to format the string for
pub fn format_for_module(s: &str,
                         options: &ModuleOptions,
                         next_bg: Option<Color>,
                         shell: Shell)
                         -> ANSIString<'static> {

    let style = style_from_modulestyle(&options.style);

    // Each shell keeps track of the number of characters that make up
    // the prompt. The ANSI escape-sequences that color the text will
    // be accidentally included in this length *unless* we prefix and
    // suffix them with these shell-specific escape-sequences. We
    // don't want the shell to mistakenly think there's fewer
    // characters remaining on the current line than there actually
    // are.
    let (len_esc_prefix, len_esc_suffix) = match shell {
        Shell::Bash => ("\\[", "\\]"),
        Shell::Zsh => ("%{", "%}"),
    };

    // Every time there is a color escape-sequence, it must be
    // surrounded by the length escape-codes. We also include the
    // padding before and after the content.
    let content = format!("{}{}{}{}{}{}{}{}{}",
                    len_esc_prefix,
                    style.prefix(),
                    len_esc_suffix,
                    options.padding_left,
                    s,
                    options.padding_right,
                    len_esc_prefix,
                    style.suffix(),
                    len_esc_suffix,
    );

    // We must format the separator differently depending on whether
    // there exists a visible module after this one or not.
    let separator_style = ModuleStyle {
        foreground: options.style.background,
        background: next_bg,
        text_properties: options.style.text_properties,
    };
    let separator_style = style_from_modulestyle(&separator_style);
    let separator = format!("{}{}{}{}{}{}{}",
                            len_esc_prefix,
                            separator_style.prefix(),
                            len_esc_suffix,
                            options.separator,
                            len_esc_prefix,
                            separator_style.suffix(),
                            len_esc_suffix);

    ANSIString::from(format!("{}{}", content, separator))
}

/// Converts a `ModuleStyle` into an `ansi_term::Style`.
fn style_from_modulestyle(s: &ModuleStyle) -> Style {
    let mut style = s.text_properties.unwrap_or_default();
    if let Some(bg) = s.background {
        style = style.on(bg);
    }
    if let Some(fg) = s.foreground {
        style = style.fg(fg);
    }
    style
}

// We use methods of the form "try_*_from_*" because the user should
// just be able to specify (in their config) what color/style they
// want, and we have to be able to parse what they say into a
// meaningful type.

/// Attempts to convert a string into an `ansi_term::Color`.
///
/// Returns a `ConvertError::NoSuchMatch` if the provided string
/// doesn't match any of the colors defined in crate `ansi_term`.
///
/// # Examples
/// ```
/// assert_eq!(try_color_from_str("black"), Ok(Color::Black));
/// assert_eq!(try_color_from_str("green"), Ok(Color::Green));
/// assert_eq!(try_color_from_str("turquoise"), Err(ConvertError::NoSuchMatch));
/// ```
pub fn try_color_from_str(s: &str) -> Result<Color, ConvertError> {
    match s.to_lowercase().as_ref() {
        "black" => Ok(Color::Black),
        "red" => Ok(Color::Red),
        "green" => Ok(Color::Green),
        "yellow" => Ok(Color::Yellow),
        "blue" => Ok(Color::Blue),
        "purple" => Ok(Color::Purple),
        "cyan" => Ok(Color::Cyan),
        "white" => Ok(Color::White),
        _ => Err(ConvertError::NoSuchMatch),
    }
}

/// Attempts to convert a string into an `ansi_term::Color::RGB`.
///
/// Returns a `ConvertError::InvalidForm` if the provided string is
/// not a sequence of 3 `u8`s, separated by commas.
///
/// # Examples
/// ```
/// assert_eq!(try_rgb_from_str("(14, 76, 1)"), Ok(Color::RGB(14, 76, 1)));
/// assert_eq!(try_rgb_from_str("0, 100, 0"), Ok(Color::RGB(0, 100, 0)));
/// assert_eq!(try_rgb_from_str("1000, b, c, -1"), Err(ConvertError::InvalidForm));
/// ```
pub fn try_rgb_from_str(s: &str) -> Result<Color, ConvertError> {
    // Strip out non-integer characters
    let cleaned: Vec<String> =
        s.split(',').map(|i| i.replace(|j| j == '(' || j == ')' || j == ' ', "")).collect();

    // We want to immediately return an Error if a conversion fails
    let mut ints: Vec<u8> = Vec::new();
    for i in &cleaned {
        ints.push(i.parse::<u8>()?);
    }

    // Note: somewhat-malformed input, like "()()(0, 0, 0)" can still
    // get through, but we can still make "sense" of it so it's good
    // to go.
    if ints.len() != 3 {
        Err(ConvertError::InvalidForm)
    } else {
        Ok(Color::RGB(ints[0], ints[1], ints[2]))
    }
}

// NOTE: We do not need a try_color_from_u8 OR a try_rgb_from_vec.
// The implementation would just look like:
// pub fn try_color_from_u8(i: u8) -> Color { Color::u8(i) }
// It would always succeed due to Rust's type system.

#[derive(Debug, PartialEq)]
/// Error type for when a conversion due to parsing fails
pub enum ConvertError {
    /// Input doesn't correspond to a valid result
    NoSuchMatch,
    /// Input is malformed and cannot be parsed
    InvalidForm,
}

impl fmt::Display for ConvertError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl Error for ConvertError {
    fn description(&self) -> &str {
        match *self {
            ConvertError::NoSuchMatch => "no match was found for the provided input",
            ConvertError::InvalidForm => "provided input was malformed and could not be parsed",
        }
    }
}

// So that we can use try!() and ? to return early if we encounter a
// ParseIntError
impl From<ParseIntError> for ConvertError {
    fn from(_: ParseIntError) -> ConvertError {
        ConvertError::InvalidForm
    }
}

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

    #[test]
    fn test_format_for_module() {
        const CONTENT: &'static str = "Hello";
        const PADDING: &'static str = " ";
        const SEPARATOR: &'static str = ">";

        let options = ModuleOptions {
            padding_left: PADDING,
            padding_right: PADDING,
            separator: SEPARATOR,
            style: ModuleStyle {
                background: Some(Color::Blue),
                foreground: Some(Color::White),
                text_properties: Some(Style::default().bold()),
            },
        };

        let formatted_string = format_for_module(CONTENT, &options, None, Shell::Bash);
        assert_eq!(format!("\\[\x1B[1;44;37m\\]{}{}{}\\[\x1B[0m\\]\\[\x1B[1;34m\\]{}\\[\x1B[0m\\]",
                           PADDING,
                           CONTENT,
                           PADDING,
                           SEPARATOR),
                   format!("{}", formatted_string))
    }

    #[test]
    fn test_style_from_modulestyle() {
        const CONTENT: &'static str = "Hello";

        struct Test {
            style: ModuleStyle,
            expected: String,
        }

        let tests = [// Normal text, no bg, custom fg
                     Test {
                         style: ModuleStyle {
                             background: None,
                             foreground: Some(Color::White),
                             text_properties: None,
                         },
                         expected: format!("\x1B[37m{}\x1B[0m", CONTENT),
                     },
                     // Bold text, no fg, custom bg
                     Test {
                         style: ModuleStyle {
                             background: Some(Color::Blue),
                             foreground: None,
                             text_properties: Some(Style::default().bold()),
                         },
                         expected: format!("\x1B[1;44m{}\x1B[0m", CONTENT),
                     },
                     // Underlined text, custom bg and fg
                     Test {
                         style: ModuleStyle {
                             background: Some(Color::Blue),
                             foreground: Some(Color::White),
                             text_properties: Some(Style::default().underline()),
                         },
                         expected: format!("\x1B[4;44;37m{}\x1B[0m", CONTENT),
                     },
                     // Normal text, no bg nor fg
                     Test {
                         style: ModuleStyle {
                             background: None,
                             foreground: None,
                             text_properties: None,
                         },
                         expected: format!("{}", CONTENT),
                     }];

        for test in &tests {
            let result = format!("{}", style_from_modulestyle(&test.style).paint(CONTENT));
            assert_eq!(result, test.expected);
        }
    }

    #[test]
    fn test_color_try_from_str() {
        assert_eq!(try_color_from_str("blue"), Ok(Color::Blue));
        assert_eq!(try_color_from_str("teal"), Err(ConvertError::NoSuchMatch));
    }

    #[test]
    fn test_rgb_try_from_str() {
        // Valid inputs
        assert_eq!(try_rgb_from_str("(0, 0, 0)"), Ok(Color::RGB(0, 0, 0)));
        assert_eq!(try_rgb_from_str("(255, 255, 255)"),
                   Ok(Color::RGB(255, 255, 255)));

        // Questionable inputs (should still work)
        assert_eq!(try_rgb_from_str("(0, 0, 0))"), Ok(Color::RGB(0, 0, 0)));

        // Improperly formed ("too many" inputs)
        assert_eq!(try_rgb_from_str("(0, 0, 0,)"),
                   Err(ConvertError::InvalidForm));

        // Too few inputs
        assert_eq!(try_rgb_from_str("(0, 0)"), Err(ConvertError::InvalidForm));

        // Inputs aren't u8's
        assert_eq!(try_rgb_from_str("(1000, 0, 0)"),
                   Err(ConvertError::InvalidForm));
        assert_eq!(try_rgb_from_str("(0, 0, -1)"),
                   Err(ConvertError::InvalidForm));
    }
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::collections::{HashMap, VecDeque};
use std::vec::Vec;
use autograd::variable::*;
use tensor::*;
use ::*;

thread_local! {
    pub static FUNC_TABLE: RefCell<VecDeque<FuncImpl>> = RefCell::new(VecDeque::new());
}
pub type FuncId = i32;
pub enum RootKind {
    RootVar(VarKind),
    RootFunc(Function),
}
impl RootKind {
    pub fn requires_grad(&self) -> bool {
        match *self {
            RootKind::RootVar(v) => v.requires_grad(),
            RootKind::RootFunc(f) => f.requires_grad(),
        }
    }
}

struct FuncStub {
    delegate: Function,
}
impl_func_delegate!(FuncStub);
impl FuncIntf for FuncStub {
    fn forward(&mut self, input: &TensorKindList) -> TensorKindList {
        unreachable!()
    }
    fn backward(&mut self, input: &TensorKindList) -> TensorKindList {
        unreachable!()
    }
}
impl FuncStub {
    fn new() -> Self {
        FuncStub { delegate: Function::default() }
    }
}

pub struct FuncImpl {
    previous_functions: Vec<(RootKind, i32)>,
    saved_variables: Vec<VarId>,
    needs_input_grad: Vec<VarId>,
    non_differentiable: Vec<TensorId>,
    output_ids: HashMap<VarId, usize>,
    to_save: Vec<TensorId>,
    requires_grad: bool,
    owner: RcMut<FuncIntf>,
}
impl Default for FuncImpl {
    fn default() -> Self {
        FuncImpl {
            previous_functions: Vec::new(),
            saved_variables: Vec::new(),
            needs_input_grad: Vec::new(),
            non_differentiable: Vec::new(),
            output_ids: HashMap::new(),
            to_save: Vec::new(),
            requires_grad: false,
            owner: RcMutNew(FuncStub::new()),
        }
    }
}

impl FuncImpl {
    fn _call_hooks(&self, input: &TensorKindList, output: &TensorKindList) {
        unimplemented!();
    }
    fn init(&mut self, intf: RcMut<FuncIntf>) {
        self.owner = intf;
    }
}

#[derive(Clone)]
pub struct Function {
    id: FuncId,
}

impl Default for Function {
    fn default() -> Self {
        Function { id: -1 }
    }
}

pub struct FIWrap<T> {
    value: RcMut<T>,
}
impl<T: 'static + FuncIntf> FIWrap<T> {
    pub fn new(arg: T) -> Self {
        let t = FIWrap { value: RcMutNew(arg) };
        t.value.borrow_mut().delegate().init(t.value.clone());
        t
    }
    pub fn f(&mut self, mut input_: &mut VarKindList) -> VarKindList {
        self.value.borrow_mut().f(input_)
    }
}

impl Function {
    pub fn new() -> Self {
        use std::usize;
        let mut id = usize::MAX;
        FUNC_TABLE.with(|f| {
                            let mut table = f.borrow_mut();
                            id = table.len();
                            table.push_back(FuncImpl::default());
                        });
        Function { id: id as i32 }
    }
    pub fn init(&self, intf: RcMut<FuncIntf>) {
        //FUNC_INTF_TABLE.with(|m| m.borrow_mut().insert(self.id, intf));
        self.access().init(intf)
    }
    pub fn previous_functions(&self) -> &Vec<(RootKind, i32)> {
        &self.access().previous_functions
    }
    pub fn requires_grad(&self) -> bool {
        self.access().requires_grad
    }
    fn access(&self) -> &mut FuncImpl {
        let vecp = FUNC_TABLE.with(|f| f.as_ptr());
        let vec = unsafe { &mut *vecp };
        &mut vec[self.id as usize]
    }
    pub fn from(id: FuncId) -> Self {
        Function { id: id }
    }
    pub fn saved_tensors(&mut self) -> TensorKindList {
        // XXX see if we can't avoid the clone
        self.access()
            .saved_variables
            .iter()
            .map(|v| VarKind::from(*v).data())
            .collect()
    }
    pub fn save_for_backward<T>(&mut self, input: &RefTensorList<T>) {
        self.access().to_save = input.iter().map(|t| t.id).collect();
    }
    pub fn _do_backward(&mut self,
                        grad_output: &TensorKindList,
                        retain_variables: bool)
                        -> TensorKindList {
        let inner = self.access();
        if inner.saved_variables.is_empty() {
            panic!("Trying to backward through the graph second \
                    time, but the buffers have already been freed. Please \
                    specify retain_variables=True when calling backward for \
                    the first time.");
        };
        let grad_input = inner.owner.borrow_mut().backward(grad_output);
        inner._call_hooks(&grad_input, grad_output);
        if !retain_variables {
            inner.saved_variables.clear();
        };
        grad_input.into()
    }
}

pub trait FuncDelegate {
    fn delegate(&mut self) -> &mut Function;
}

pub trait FuncIntf: FuncDelegate {
    fn forward(&mut self, input: &TensorKindList) -> TensorKindList;
    fn backward(&mut self, input: &TensorKindList) -> TensorKindList;
    fn f(&mut self, mut input_: &mut VarKindList) -> VarKindList {
        let is_volatile = input_.iter().any(|v| v.is_volatile());
        {
            // do start graph stuff with f
            let f = self.delegate();
            let mut inner = f.access();
            if !is_volatile {
                inner.previous_functions = input_
                    .iter()
                    .enumerate()
                    .map(|(i, v)| if let Some(grad_fn) = v.grad_fn() {
                             (RootKind::RootFunc(grad_fn), v.varid())
                         } else {
                             (RootKind::RootVar(v.clone()), v.varid())
                         })
                    .collect();
                inner.needs_input_grad = input_
                    .iter()
                    .filter_map(|v| if v.requires_grad() {
                                    Some(v.varid())
                                } else {
                                    None
                                })
                    .collect();
                inner.requires_grad = inner.needs_input_grad.len() != 0;
            }
        }
        let v;
        {
            let input_tensors = input_.iter_mut().map(|v| v.data()).collect();
            v = self.forward(&input_tensors);
        }
        let f = self.delegate();
        let mut inner = f.access();
        let output = if is_volatile {
            let args = VariableArgs::build().volatile(true).done();
            v.into_iter().map(|t| VarKind::new_args(t, &args)).collect()
        } else {
            let args = VariableArgs::build()
                .creator(Some(f.clone()))
                .requires_grad(inner.requires_grad)
                .done();
            let mut output: VarKindList =
                v.into_iter().map(|t| VarKind::new_args(t, &args)).collect();
            for (i, v) in output.iter().enumerate() {
                inner.output_ids.insert(v.varid(), i);
            }
            if !inner.to_save.is_empty() {
                /* if a tensor was modified in place replace the old variable with the new one */
                let mut t2v = HashMap::new();
                for ref mut var in input_.iter_mut() {
                    t2v.insert(var.tid(), var.varid());
                }
                for ref mut var in &mut output.iter_mut() {
                    t2v.insert(var.tid(), var.varid());
                }
                for t in inner.to_save.iter() {
                    inner.saved_variables.push(t2v[t]);
                }
                inner.to_save.clear();
            };
            if !inner.non_differentiable.is_empty() {
                for ref mut var in &mut output {
                    if inner.non_differentiable.contains(&var.tid()) {
                        var.requires_nograd()
                    }
                }
                inner.non_differentiable.clear();
            };
            output
        };
        output
    }
}
<commit_msg>fix compile post update<commit_after>use std::cell::RefCell;
use std::collections::{HashMap, VecDeque};
use std::vec::Vec;
use autograd::variable::*;
use tensor::*;
use ::*;

thread_local! {
    pub static FUNC_TABLE: RefCell<VecDeque<FuncImpl>> = RefCell::new(VecDeque::new());
}
pub type FuncId = i32;
pub enum RootKind {
    RootVar(VarKind),
    RootFunc(Function),
}
impl RootKind {
    pub fn requires_grad(&self) -> bool {
        match *self {
            RootKind::RootVar(ref v) => v.requires_grad(),
            RootKind::RootFunc(ref f) => f.requires_grad(),
        }
    }
}

struct FuncStub {
    delegate: Function,
}
impl_func_delegate!(FuncStub);
impl FuncIntf for FuncStub {
    fn forward(&mut self, input: &TensorKindList) -> TensorKindList {
        unreachable!()
    }
    fn backward(&mut self, input: &TensorKindList) -> TensorKindList {
        unreachable!()
    }
}
impl FuncStub {
    fn new() -> Self {
        FuncStub { delegate: Function::default() }
    }
}

pub struct FuncImpl {
    previous_functions: Vec<(RootKind, i32)>,
    saved_variables: Vec<VarId>,
    needs_input_grad: Vec<VarId>,
    non_differentiable: Vec<TensorId>,
    output_ids: HashMap<VarId, usize>,
    to_save: Vec<TensorId>,
    requires_grad: bool,
    owner: RcMut<FuncIntf>,
}
impl Default for FuncImpl {
    fn default() -> Self {
        FuncImpl {
            previous_functions: Vec::new(),
            saved_variables: Vec::new(),
            needs_input_grad: Vec::new(),
            non_differentiable: Vec::new(),
            output_ids: HashMap::new(),
            to_save: Vec::new(),
            requires_grad: false,
            owner: RcMutNew(FuncStub::new()),
        }
    }
}

impl FuncImpl {
    fn _call_hooks(&self, input: &TensorKindList, output: &TensorKindList) {
        unimplemented!();
    }
    fn init(&mut self, intf: RcMut<FuncIntf>) {
        self.owner = intf;
    }
}

#[derive(Clone)]
pub struct Function {
    id: FuncId,
}

impl Default for Function {
    fn default() -> Self {
        Function { id: -1 }
    }
}

pub struct FIWrap<T> {
    value: RcMut<T>,
}
impl<T: 'static + FuncIntf> FIWrap<T> {
    pub fn new(arg: T) -> Self {
        let t = FIWrap { value: RcMutNew(arg) };
        t.value.borrow_mut().delegate().init(t.value.clone());
        t
    }
    pub fn f(&mut self, mut input_: &mut VarKindList) -> VarKindList {
        self.value.borrow_mut().f(input_)
    }
}

impl Function {
    pub fn new() -> Self {
        use std::usize;
        let mut id = usize::MAX;
        FUNC_TABLE.with(|f| {
                            let mut table = f.borrow_mut();
                            id = table.len();
                            table.push_back(FuncImpl::default());
                        });
        Function { id: id as i32 }
    }
    pub fn init(&self, intf: RcMut<FuncIntf>) {
        //FUNC_INTF_TABLE.with(|m| m.borrow_mut().insert(self.id, intf));
        self.access().init(intf)
    }
    pub fn previous_functions(&self) -> &Vec<(RootKind, i32)> {
        &self.access().previous_functions
    }
    pub fn requires_grad(&self) -> bool {
        self.access().requires_grad
    }
    fn access(&self) -> &mut FuncImpl {
        let vecp = FUNC_TABLE.with(|f| f.as_ptr());
        let vec = unsafe { &mut *vecp };
        &mut vec[self.id as usize]
    }
    pub fn from(id: FuncId) -> Self {
        Function { id: id }
    }
    pub fn saved_tensors(&mut self) -> TensorKindList {
        // XXX see if we can't avoid the clone
        self.access()
            .saved_variables
            .iter()
            .map(|v| VarKind::from(*v).data())
            .collect()
    }
    pub fn save_for_backward<T>(&mut self, input: &RefTensorList<T>) {
        self.access().to_save = input.iter().map(|t| t.id).collect();
    }
    pub fn _do_backward(&mut self,
                        grad_output: &TensorKindList,
                        retain_variables: bool)
                        -> TensorKindList {
        let inner = self.access();
        if inner.saved_variables.is_empty() {
            panic!("Trying to backward through the graph second \
                    time, but the buffers have already been freed. Please \
                    specify retain_variables=True when calling backward for \
                    the first time.");
        };
        let grad_input = inner.owner.borrow_mut().backward(grad_output);
        inner._call_hooks(&grad_input, grad_output);
        if !retain_variables {
            inner.saved_variables.clear();
        };
        grad_input.into()
    }
}

pub trait FuncDelegate {
    fn delegate(&mut self) -> &mut Function;
}

pub trait FuncIntf: FuncDelegate {
    fn forward(&mut self, input: &TensorKindList) -> TensorKindList;
    fn backward(&mut self, input: &TensorKindList) -> TensorKindList;
    fn f(&mut self, mut input_: &mut VarKindList) -> VarKindList {
        let is_volatile = input_.iter().any(|v| v.is_volatile());
        {
            // do start graph stuff with f
            let f = self.delegate();
            let mut inner = f.access();
            if !is_volatile {
                inner.previous_functions = input_
                    .iter()
                    .enumerate()
                    .map(|(i, v)| if let Some(grad_fn) = v.grad_fn() {
                             (RootKind::RootFunc(grad_fn), v.varid())
                         } else {
                             (RootKind::RootVar(v.clone()), v.varid())
                         })
                    .collect();
                inner.needs_input_grad = input_
                    .iter()
                    .filter_map(|v| if v.requires_grad() {
                                    Some(v.varid())
                                } else {
                                    None
                                })
                    .collect();
                inner.requires_grad = inner.needs_input_grad.len() != 0;
            }
        }
        let v;
        {
            let input_tensors = input_.iter_mut().map(|v| v.data()).collect();
            v = self.forward(&input_tensors);
        }
        let f = self.delegate();
        let mut inner = f.access();
        let output = if is_volatile {
            let args = VariableArgs::build().volatile(true).done();
            v.into_iter().map(|t| VarKind::new_args(t, &args)).collect()
        } else {
            let args = VariableArgs::build()
                .creator(Some(f.clone()))
                .requires_grad(inner.requires_grad)
                .done();
            let mut output: VarKindList =
                v.into_iter().map(|t| VarKind::new_args(t, &args)).collect();
            for (i, v) in output.iter().enumerate() {
                inner.output_ids.insert(v.varid(), i);
            }
            if !inner.to_save.is_empty() {
                /* if a tensor was modified in place replace the old variable with the new one */
                let mut t2v = HashMap::new();
                for ref mut var in input_.iter_mut() {
                    t2v.insert(var.tid(), var.varid());
                }
                for ref mut var in &mut output.iter_mut() {
                    t2v.insert(var.tid(), var.varid());
                }
                for t in inner.to_save.iter() {
                    inner.saved_variables.push(t2v[t]);
                }
                inner.to_save.clear();
            };
            if !inner.non_differentiable.is_empty() {
                for ref mut var in &mut output {
                    if inner.non_differentiable.contains(&var.tid()) {
                        var.requires_nograd()
                    }
                }
                inner.non_differentiable.clear();
            };
            output
        };
        output
    }
}
<|endoftext|>