<commit_before>/*!
Library for reading/writing Labeled Tab-Separated Values

# Example

~~~~~~~~~~~~~~~~~~~~~~
extern mod ltsv;
use ltsv::LTSVWriter;
use ltsv::LTSVReader;

fn main() {
    let infile = io::file_reader(&Path("path/fo/file.tlsv")).get();
    for infile.read_ltsv().each |record| {
        for record.each |&(label, value)| {
            io::println(fmt!("%s: %s", *label, *value));
        }
    }
}
~~~~~~~~~~~~~~~~~~~~~~
*/
#[link(name = "ltsv",
       vers = "0.2",
       uuid = "E0EA0251-E165-4612-919F-38E89ACECBE9",
       url  = "https://github.com/tychosci/rust-ltsv/")];

#[comment = "Library for reading/writing Labeled Tab-Separated Values"];
#[license = "MIT license"];
#[crate_type = "lib"];

use core::container::Map;
use core::hashmap::linear::LinearMap;
use core::io::WriterUtil;

pub type Record = LinearMap<~str, ~str>;

enum ParseType {
    FieldLabel,
    FieldValue,
    Field,
    Record,
    Ltsv
}

#[deriving(Eq)]
enum ParseDelimiter {
    EOF, TAB, NL, MISC
}

enum ParseResult<T> {
    ParseError(~str),
    ParseOk(ParseType, ParseDelimiter, T)
}

pub trait LTSVWriter {
    fn write_ltsv(&self, ltsv: &[Record]);
    fn write_ltsv_record(&self, record: &Record);
}

pub trait LTSVReader {
    fn read_ltsv(&self) -> ~[Record];
    fn each_ltsv_record(&self, f: &fn(&Record) -> bool);
    fn each_ltsv_field(&self, f: &fn(&(~str, ~str)) -> bool);
}

impl<T: io::Writer> LTSVWriter for T {
    fn write_ltsv(&self, ltsv: &[Record]) {
        for ltsv.each |record| {
            self.write_ltsv_record(record);
            self.write_char('\n');
        }
    }
    fn write_ltsv_record(&self, record: &Record) {
        let mut is_first = true;
        for record.each |&(k, v)| {
            if !is_first { self.write_char('\t'); }
            self.write_str(fmt!("%s:%s", *k, *v));
            if is_first { is_first = false; }
        }
    }
}

impl<T: io::Reader> LTSVReader for T {
    fn read_ltsv(&self) -> ~[Record] {
        let mut parser = LTSVParser::new(self);
        match parser.parse_ltsv() {
            ParseError(reason) => fail!(reason),
            ParseOk(_, _, records) => records
        }
    }
    fn each_ltsv_record(&self, f: &fn(&Record) -> bool) {
        let mut parser = LTSVParser::new(self);
        while !parser.eof() {
            match parser.parse_record() {
                ParseError(reason) => fail!(reason),
                ParseOk(_, _, record) => if !f(&record) { break; }
            }
        }
    }
    fn each_ltsv_field(&self, f: &fn(&(~str, ~str)) -> bool) {
        let mut parser = LTSVParser::new(self);
        while !parser.eof() {
            match parser.parse_field() {
                ParseError(reason) => fail!(reason),
                ParseOk(_, _, field) => if !f(&field) { break; }
            }
        }
    }
}

struct LTSVParser<'self, T> {
    priv rd: &'self T,
    priv cur: @mut int
}

pub impl<'self, T: io::Reader> LTSVParser<'self, T> {
    fn new(rd: &'self T) -> LTSVParser<'self, T> {
        let cur = @mut rd.read_byte();
        LTSVParser { rd: rd, cur: cur }
    }

    fn eof(&self) -> bool { *self.cur == -1 }

    fn bump(&self) {
        if !self.eof() {
            *self.cur = self.rd.read_byte();
        }
    }

    fn parse_ltsv(&self) -> ParseResult<~[Record]> {
        let mut records = ~[];
        loop {
            match self.parse_record() {
                ParseError(reason) => {
                    return ParseError(reason);
                }
                ParseOk(_, EOF, record) => {
                    records.push(record);
                    break;
                }
                ParseOk(_, _, record) => {
                    records.push(record);
                }
            }
        }
        ParseOk(Ltsv, EOF, records)
    }

    fn parse_record(&self) -> ParseResult<Record> {
        let mut record = LinearMap::new();
        loop {
            match self.parse_field() {
                ParseError(reason) => {
                    return ParseError(reason);
                }
                ParseOk(_, TAB, (label, value)) => {
                    record.insert(label, value);
                }
                ParseOk(_, delim, (label, value)) => {
                    record.insert(label, value);
                    return ParseOk(Record, delim, record);
                }
            }
        }
    }

    fn parse_field(&self) -> ParseResult<(~str, ~str)> {
        self.skip_whitespaces();
        let label = match self.parse_field_label() {
            ParseError(reason) => return ParseError(reason),
            ParseOk(_, _, label) => { self.bump(); label }
        };
        match self.parse_field_value() {
            ParseError(reason) => {
                ParseError(reason)
            }
            ParseOk(_, delim, value) => {
                self.bump();
                // avoid skipping whitespaces in the middle of parsing record.
                if delim != TAB { self.skip_whitespaces(); }
                // re-check EOF
                let delim = if self.eof() { EOF } else { delim };
                ParseOk(Field, delim, (label, value))
            }
        }
    }

    priv fn parse_field_label(&self) -> ParseResult<~str> {
        let mut bytes = ~[];
        loop {
            match *self.cur {
                0x30..0x39 | 0x41..0x5a | 0x61..0x7a | 0x5f |
                0x2e | 0x2d => bytes.push(*self.cur as u8),
                0x3a if bytes.len() == 0 => return ParseError(~"label is empty"),
                0x3a => return ParseOk(FieldLabel, MISC, str::from_bytes(bytes)),
                -1   => return ParseError(~"EOF while parsing field label"),
                _    => return ParseError(~"invalid byte detected")
            }
            self.bump();
        }
    }

    priv fn parse_field_value(&self) -> ParseResult<~str> {
        let mut bytes = ~[];
        loop {
            match *self.cur {
                0x01..0x08 | 0x0b | 0x0c |
                0x0e..0xff => bytes.push(*self.cur as u8),
                0x0d => return self.consume_forward_LF(str::from_bytes(bytes)),
                0x0a => return ParseOk(FieldValue, NL, str::from_bytes(bytes)),
                0x09 => return ParseOk(FieldValue, TAB, str::from_bytes(bytes)),
                -1   => return ParseOk(FieldValue, EOF, str::from_bytes(bytes)),
                _    => return ParseError(~"invalid byte detected")
            }
            self.bump();
        }
    }

    priv fn consume_forward_LF(&self, rv: ~str) -> ParseResult<~str> {
        self.bump();
        if *self.cur != 0x0a {
            ParseError(~"CR detected, but not provided with LF")
        } else {
            ParseOk(FieldValue, NL, rv)
        }
    }

    priv fn skip_whitespaces(&self) {
        while char::is_whitespace(*self.cur as char) {
            self.bump();
        }
    }
}

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

    fn mk_record_string() -> ~str {
        do io::with_str_writer |wr| {
            // genzairyou
            wr.write_str(fmt!("%s:%s\t", "tofu", "豆"));
            wr.write_str(fmt!("%s:%s\t", "kamaboko", "魚"));
            wr.write_str(fmt!("%s:%s\n", "sukonbu", "海藻"));
            // konomi
            wr.write_str(fmt!("%s:%s\t", "tofu", "好き"));
            wr.write_str(fmt!("%s:%s\t", "kamaboko", "普通"));
            wr.write_str(fmt!("%s:%s\n", "sukonbu", "苦手"));
        }
    }

    #[test]
    fn test_parse_simple() {
        let records = io::with_str_reader(~"a:1\tb:2", |rd| rd.read_ltsv());
        assert_eq!(records.len(), 1);
    }

    #[test]
    fn test_parse_full() {
        let s = mk_record_string();
        let records = io::with_str_reader(s, |rd| rd.read_ltsv());
        assert_eq!(records.len(), 2);
    }

    #[test]
    fn test_parse_ltsv_trailing_nl_and_write() {
        let s = mk_record_string();
        let records_1 = io::with_str_reader(s, |rd| rd.read_ltsv());
        let s2 = io::with_str_writer(|wr| wr.write_ltsv(records_1));
        let records_2 = io::with_str_reader(s2, |rd| rd.read_ltsv());
        assert_eq!(records_1, records_2);
    }

    #[test]
    fn test_each_read_each_record() {
        let s = mk_record_string();
        let ks = [~"tofu", ~"kamaboko", ~"sukonbu"];
        let vs = [~"豆", ~"魚", ~"海藻", ~"好き", ~"普通", ~"苦手"];
        do io::with_str_reader(s) |rd| {
            for rd.each_ltsv_record |record| {
                for record.each |&(k, v)| {
                    assert!(ks.contains(k));
                    assert!(vs.contains(v));
                }
            }
        }
    }
}
<commit_msg>Cleanup<commit_after>/*!
Library for reading/writing Labeled Tab-Separated Values

# Example

~~~~~~~~~~~~~~~~~~~~~~
extern mod ltsv;
use ltsv::LTSVWriter;
use ltsv::LTSVReader;

fn main() {
    let infile = io::file_reader(&Path("path/fo/file.tlsv")).get();
    for infile.read_ltsv().each |record| {
        for record.each |&(label, value)| {
            io::println(fmt!("%s: %s", *label, *value));
        }
    }
}
~~~~~~~~~~~~~~~~~~~~~~
*/
#[link(name = "ltsv",
       vers = "0.2",
       uuid = "E0EA0251-E165-4612-919F-38E89ACECBE9",
       url  = "https://github.com/tychosci/rust-ltsv/")];

#[comment = "Library for reading/writing Labeled Tab-Separated Values"];
#[license = "MIT license"];
#[crate_type = "lib"];

use core::container::Map;
use core::hashmap::linear::LinearMap;
use core::io::WriterUtil;

pub type Record = LinearMap<~str, ~str>;

enum ParseType {
    FieldLabel,
    FieldValue,
    Field,
    Record,
    Ltsv
}

#[deriving(Eq)]
enum ParseDelimiter {
    EOF, TAB, NL, MISC
}

enum ParseResult<T> {
    ParseError(~str),
    ParseOk(ParseType, ParseDelimiter, T)
}

pub trait LTSVWriter {
    fn write_ltsv(&self, ltsv: &[Record]);
    fn write_ltsv_record(&self, record: &Record);
}

pub trait LTSVReader {
    fn read_ltsv(&self) -> ~[Record];
    fn each_ltsv_record(&self, f: &fn(&Record) -> bool);
    fn each_ltsv_field(&self, f: &fn(&(~str, ~str)) -> bool);
}

impl<T: io::Writer> LTSVWriter for T {
    fn write_ltsv(&self, ltsv: &[Record]) {
        for ltsv.each |record| {
            self.write_ltsv_record(record);
            self.write_char('\n');
        }
    }
    fn write_ltsv_record(&self, record: &Record) {
        let mut is_first = true;
        for record.each |&(k, v)| {
            if !is_first { self.write_char('\t'); }
            self.write_str(fmt!("%s:%s", *k, *v));
            if is_first { is_first = false; }
        }
    }
}

impl<T: io::Reader> LTSVReader for T {
    fn read_ltsv(&self) -> ~[Record] {
        let mut parser = LTSVParser::new(self);
        match parser.parse_ltsv() {
            ParseError(reason) => fail!(reason),
            ParseOk(_, _, records) => records
        }
    }
    fn each_ltsv_record(&self, f: &fn(&Record) -> bool) {
        let mut parser = LTSVParser::new(self);
        while !parser.eof() {
            match parser.parse_record() {
                ParseError(reason) => fail!(reason),
                ParseOk(_, _, record) => if !f(&record) { break; }
            }
        }
    }
    fn each_ltsv_field(&self, f: &fn(&(~str, ~str)) -> bool) {
        let mut parser = LTSVParser::new(self);
        while !parser.eof() {
            match parser.parse_field() {
                ParseError(reason) => fail!(reason),
                ParseOk(_, _, field) => if !f(&field) { break; }
            }
        }
    }
}

struct LTSVParser<'self, T> {
    priv rd: &'self T,
    priv cur: @mut int
}

pub impl<'self, T: io::Reader> LTSVParser<'self, T> {
    fn new(rd: &'self T) -> LTSVParser<'self, T> {
        let cur = @mut rd.read_byte();
        LTSVParser { rd: rd, cur: cur }
    }

    fn eof(&self) -> bool {
        *self.cur == -1
    }

    fn parse_ltsv(&self) -> ParseResult<~[Record]> {
        let mut records = ~[];
        loop {
            match self.parse_record() {
                ParseError(reason) => {
                    return ParseError(reason);
                }
                ParseOk(_, EOF, record) => {
                    records.push(record);
                    break;
                }
                ParseOk(_, _, record) => {
                    records.push(record);
                }
            }
        }
        ParseOk(Ltsv, EOF, records)
    }

    fn parse_record(&self) -> ParseResult<Record> {
        let mut record = LinearMap::new();
        loop {
            match self.parse_field() {
                ParseError(reason) => {
                    return ParseError(reason);
                }
                ParseOk(_, TAB, (label, value)) => {
                    record.insert(label, value);
                }
                ParseOk(_, delim, (label, value)) => {
                    record.insert(label, value);
                    return ParseOk(Record, delim, record);
                }
            }
        }
    }

    fn parse_field(&self) -> ParseResult<(~str, ~str)> {
        self.skip_whitespaces();
        let label = match self.parse_field_label() {
            ParseError(reason) => return ParseError(reason),
            ParseOk(_, _, label) => { self.bump(); label }
        };
        match self.parse_field_value() {
            ParseError(reason) => {
                ParseError(reason)
            }
            ParseOk(_, delim, value) => {
                self.bump();
                // avoid skipping whitespaces in the middle of parsing record.
                if delim != TAB { self.skip_whitespaces(); }
                // re-check EOF
                let delim = if self.eof() { EOF } else { delim };
                ParseOk(Field, delim, (label, value))
            }
        }
    }

    priv fn parse_field_label(&self) -> ParseResult<~str> {
        let mut bytes = ~[];
        loop {
            match *self.cur {
                0x30..0x39 | 0x41..0x5a | 0x61..0x7a | 0x5f |
                0x2e | 0x2d => bytes.push(*self.cur as u8),
                0x3a if bytes.len() == 0 => return ParseError(~"label is empty"),
                0x3a => return ParseOk(FieldLabel, MISC, str::from_bytes(bytes)),
                -1   => return ParseError(~"EOF while parsing field label"),
                _    => return ParseError(~"invalid byte detected")
            }
            self.bump();
        }
    }

    priv fn parse_field_value(&self) -> ParseResult<~str> {
        let mut bytes = ~[];
        loop {
            match *self.cur {
                0x01..0x08 | 0x0b | 0x0c |
                0x0e..0xff => bytes.push(*self.cur as u8),
                0x0d => return self.consume_forward_LF(str::from_bytes(bytes)),
                0x0a => return ParseOk(FieldValue, NL, str::from_bytes(bytes)),
                0x09 => return ParseOk(FieldValue, TAB, str::from_bytes(bytes)),
                -1   => return ParseOk(FieldValue, EOF, str::from_bytes(bytes)),
                _    => return ParseError(~"invalid byte detected")
            }
            self.bump();
        }
    }

    priv fn consume_forward_LF(&self, rv: ~str) -> ParseResult<~str> {
        self.bump();
        if *self.cur != 0x0a {
            ParseError(~"CR detected, but not provided with LF")
        } else {
            ParseOk(FieldValue, NL, rv)
        }
    }

    priv fn bump(&self) {
        if !self.eof() {
            *self.cur = self.rd.read_byte();
        }
    }

    priv fn skip_whitespaces(&self) {
        while char::is_whitespace(*self.cur as char) {
            self.bump();
        }
    }
}

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

    fn mk_record_string() -> ~str {
        do io::with_str_writer |wr| {
            // genzairyou
            wr.write_str(fmt!("%s:%s\t", "tofu", "豆"));
            wr.write_str(fmt!("%s:%s\t", "kamaboko", "魚"));
            wr.write_str(fmt!("%s:%s\n", "sukonbu", "海藻"));
            // konomi
            wr.write_str(fmt!("%s:%s\t", "tofu", "好き"));
            wr.write_str(fmt!("%s:%s\t", "kamaboko", "普通"));
            wr.write_str(fmt!("%s:%s\n", "sukonbu", "苦手"));
        }
    }

    #[test]
    fn test_parse_simple() {
        let records = io::with_str_reader(~"a:1\tb:2", |rd| rd.read_ltsv());
        assert_eq!(records.len(), 1);
    }

    #[test]
    fn test_parse_full() {
        let s = mk_record_string();
        let records = io::with_str_reader(s, |rd| rd.read_ltsv());
        assert_eq!(records.len(), 2);
    }

    #[test]
    fn test_parse_ltsv_trailing_nl_and_write() {
        let s = mk_record_string();
        let records_1 = io::with_str_reader(s, |rd| rd.read_ltsv());
        let s2 = io::with_str_writer(|wr| wr.write_ltsv(records_1));
        let records_2 = io::with_str_reader(s2, |rd| rd.read_ltsv());
        assert_eq!(records_1, records_2);
    }

    #[test]
    fn test_each_read_each_record() {
        let s = mk_record_string();
        let ks = [~"tofu", ~"kamaboko", ~"sukonbu"];
        let vs = [~"豆", ~"魚", ~"海藻", ~"好き", ~"普通", ~"苦手"];
        do io::with_str_reader(s) |rd| {
            for rd.each_ltsv_record |record| {
                for record.each |&(k, v)| {
                    assert!(ks.contains(k));
                    assert!(vs.contains(v));
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// Originally from https://github.com/rust-lang/crates.io/blob/master/src/s3/lib.rs
//#![deny(warnings)]

use std::ascii::AsciiExt;
use std::fmt;
use std::io::prelude::*;
use std::io;

use simples3::credential::*;
use crypto::digest::Digest;
use crypto::hmac::Hmac;
use crypto::mac::Mac;
use crypto::sha1::Sha1;
use hyper::{self,header};
use rustc_serialize::base64::{ToBase64, STANDARD};
use time;

#[derive(Debug, Copy, Clone)]
#[allow(dead_code)]
/// Whether or not to use SSL.
pub enum Ssl {
    /// Use SSL.
    Yes,
    /// Do not use SSL.
    No,
}

fn base_url(bucket_name: &str, ssl: Ssl, region: Option<&str>) -> String {
    format!("{}://{}.s3{}.amazonaws.com/",
            match ssl {
                Ssl::Yes => "https",
                Ssl::No => "http",
            },
            bucket_name,
            match region {
                Some(ref r) => format!("-{}", r),
                None => String::new(),
            })
}

fn hmac<D: Digest>(d: D, key: &[u8], data: &[u8]) -> Vec<u8> {
    let mut hmac = Hmac::new(d, key);
    hmac.input(data);
    hmac.result().code().iter().map(|b| *b).collect::<Vec<u8>>()
}

fn signature(string_to_sign: &str, signing_key: &str) -> String {
    hmac(Sha1::new(), signing_key.as_bytes(), string_to_sign.as_bytes()).to_base64(STANDARD)
}

/// An S3 bucket.
pub struct Bucket {
    name: String,
    base_url: String,
    client: hyper::Client,
}

impl fmt::Display for Bucket {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Bucket(name={}, base_url={})", self.name, self.base_url)
    }
}

/// Errors from making S3 requests.
#[derive(Debug)]
pub enum S3Error {
    HyperError(hyper::Error),
    IOError(io::Error),
    BadHTTPStatus(hyper::status::StatusCode),
}

impl Bucket {
    pub fn new(name: &str, ssl: Ssl) -> Bucket {
        let base_url = base_url(&name, ssl, None);
        Bucket {
            name: name.to_owned(),
            base_url: base_url,
            client: hyper::Client::new(),
        }
    }

    pub fn get(&self, key: &str) -> Result<Vec<u8>, S3Error> {
        let url = format!("{}{}", self.base_url, key);
        debug!("GET {}", url);
        match self.client.get(&url).send() {
            Err(e) => Err(S3Error::HyperError(e)),
            Ok(mut res) => {
                if res.status.class() == hyper::status::StatusClass::Success {
                    let mut body = vec!();
                    res.read_to_end(&mut body)
                        .or_else(|e| Err(S3Error::IOError(e)))
                        .map(|_| body)
                } else {
                    Err(S3Error::BadHTTPStatus(res.status))
                }
            }
        }
    }

    pub fn put(&self, key: &str, content: &[u8],
               creds: &AwsCredentials)
               -> Result<(), S3Error> {
        let content_type = "application/octet-stream";
        let date = time::now().rfc822z().to_string();
        let mut headers = header::Headers::new();
        let mut canonical_headers = String::new();
        let token = creds.token().as_ref().map(|s| s.as_str());
        // Keep the list of header values sorted!
        for (header, maybe_value) in vec![
            ("x-amz-acl", Some("public-read")),
            ("x-amz-security-token", token),
            ("x-amz-storage-class", Some("REDUCED_REDUNDANCY")),
            ] {
            if let Some(ref value) = maybe_value {
                headers.set_raw(header, vec!(value.as_bytes().to_vec()));
                canonical_headers.push_str(format!("{}:{}\n", header.to_ascii_lowercase(), value).as_ref());
            }
        }
        let auth = self.auth("PUT", &date, key, "", &canonical_headers, content_type, creds);
        headers.set_raw("Date", vec!(date.into_bytes()));
        headers.set(header::ContentType(content_type.parse().unwrap()));
        headers.set(header::CacheControl(vec![
            // Two weeks
            header::CacheDirective::MaxAge(1296000)
                ]));
        headers.set_raw("Authorization", vec!(auth.into_bytes()));
        let url = format!("{}{}", self.base_url, key);
        debug!("PUT {}", url);

        match self.client.put(&url)
            .body(content)
            .headers(headers)
            .send() {
                Err(e) => {
                    trace!("PUT failed with error: {:?}", e);
                    Err(S3Error::HyperError(e))
                }
                Ok(res) => {
                    if res.status.class() == hyper::status::StatusClass::Success {
                        trace!("PUT succeeded");
                        Ok(())
                    } else {
                        trace!("PUT failed with HTTP status: {}", res.status);
                        Err(S3Error::BadHTTPStatus(res.status))
                    }
                }
            }
    }

    // http://docs.aws.amazon.com/AmazonS3/latest/dev/RESTAuthentication.html
    fn auth(&self, verb: &str, date: &str, path: &str,
            md5: &str, headers: &str, content_type: &str, creds: &AwsCredentials) -> String {
        let string = format!("{verb}\n{md5}\n{ty}\n{date}\n{headers}{resource}",
                             verb = verb,
                             md5 = md5,
                             ty = content_type,
                             date = date,
                             headers = headers,
                             resource = format!("/{}/{}", self.name, path));
        let signature = signature(&string, creds.aws_secret_access_key());
        format!("AWS {}:{}", creds.aws_access_key_id(), signature)
    }
}
<commit_msg>stop setting public-read ACL, rely on bucket having it set in policy<commit_after>// Originally from https://github.com/rust-lang/crates.io/blob/master/src/s3/lib.rs
//#![deny(warnings)]

use std::ascii::AsciiExt;
use std::fmt;
use std::io::prelude::*;
use std::io;

use simples3::credential::*;
use crypto::digest::Digest;
use crypto::hmac::Hmac;
use crypto::mac::Mac;
use crypto::sha1::Sha1;
use hyper::{self,header};
use rustc_serialize::base64::{ToBase64, STANDARD};
use time;

#[derive(Debug, Copy, Clone)]
#[allow(dead_code)]
/// Whether or not to use SSL.
pub enum Ssl {
    /// Use SSL.
    Yes,
    /// Do not use SSL.
    No,
}

fn base_url(bucket_name: &str, ssl: Ssl, region: Option<&str>) -> String {
    format!("{}://{}.s3{}.amazonaws.com/",
            match ssl {
                Ssl::Yes => "https",
                Ssl::No => "http",
            },
            bucket_name,
            match region {
                Some(ref r) => format!("-{}", r),
                None => String::new(),
            })
}

fn hmac<D: Digest>(d: D, key: &[u8], data: &[u8]) -> Vec<u8> {
    let mut hmac = Hmac::new(d, key);
    hmac.input(data);
    hmac.result().code().iter().map(|b| *b).collect::<Vec<u8>>()
}

fn signature(string_to_sign: &str, signing_key: &str) -> String {
    hmac(Sha1::new(), signing_key.as_bytes(), string_to_sign.as_bytes()).to_base64(STANDARD)
}

/// An S3 bucket.
pub struct Bucket {
    name: String,
    base_url: String,
    client: hyper::Client,
}

impl fmt::Display for Bucket {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Bucket(name={}, base_url={})", self.name, self.base_url)
    }
}

/// Errors from making S3 requests.
#[derive(Debug)]
pub enum S3Error {
    HyperError(hyper::Error),
    IOError(io::Error),
    BadHTTPStatus(hyper::status::StatusCode),
}

impl Bucket {
    pub fn new(name: &str, ssl: Ssl) -> Bucket {
        let base_url = base_url(&name, ssl, None);
        Bucket {
            name: name.to_owned(),
            base_url: base_url,
            client: hyper::Client::new(),
        }
    }

    pub fn get(&self, key: &str) -> Result<Vec<u8>, S3Error> {
        let url = format!("{}{}", self.base_url, key);
        debug!("GET {}", url);
        match self.client.get(&url).send() {
            Err(e) => Err(S3Error::HyperError(e)),
            Ok(mut res) => {
                if res.status.class() == hyper::status::StatusClass::Success {
                    let mut body = vec!();
                    res.read_to_end(&mut body)
                        .or_else(|e| Err(S3Error::IOError(e)))
                        .map(|_| body)
                } else {
                    Err(S3Error::BadHTTPStatus(res.status))
                }
            }
        }
    }

    pub fn put(&self, key: &str, content: &[u8],
               creds: &AwsCredentials)
               -> Result<(), S3Error> {
        let content_type = "application/octet-stream";
        let date = time::now().rfc822z().to_string();
        let mut headers = header::Headers::new();
        let mut canonical_headers = String::new();
        let token = creds.token().as_ref().map(|s| s.as_str());
        // Keep the list of header values sorted!
        for (header, maybe_value) in vec![
            ("x-amz-security-token", token),
            ("x-amz-storage-class", Some("REDUCED_REDUNDANCY")),
            ] {
            if let Some(ref value) = maybe_value {
                headers.set_raw(header, vec!(value.as_bytes().to_vec()));
                canonical_headers.push_str(format!("{}:{}\n", header.to_ascii_lowercase(), value).as_ref());
            }
        }
        let auth = self.auth("PUT", &date, key, "", &canonical_headers, content_type, creds);
        headers.set_raw("Date", vec!(date.into_bytes()));
        headers.set(header::ContentType(content_type.parse().unwrap()));
        headers.set(header::CacheControl(vec![
            // Two weeks
            header::CacheDirective::MaxAge(1296000)
                ]));
        headers.set_raw("Authorization", vec!(auth.into_bytes()));
        let url = format!("{}{}", self.base_url, key);
        debug!("PUT {}", url);

        match self.client.put(&url)
            .body(content)
            .headers(headers)
            .send() {
                Err(e) => {
                    trace!("PUT failed with error: {:?}", e);
                    Err(S3Error::HyperError(e))
                }
                Ok(res) => {
                    if res.status.class() == hyper::status::StatusClass::Success {
                        trace!("PUT succeeded");
                        Ok(())
                    } else {
                        trace!("PUT failed with HTTP status: {}", res.status);
                        Err(S3Error::BadHTTPStatus(res.status))
                    }
                }
            }
    }

    // http://docs.aws.amazon.com/AmazonS3/latest/dev/RESTAuthentication.html
    fn auth(&self, verb: &str, date: &str, path: &str,
            md5: &str, headers: &str, content_type: &str, creds: &AwsCredentials) -> String {
        let string = format!("{verb}\n{md5}\n{ty}\n{date}\n{headers}{resource}",
                             verb = verb,
                             md5 = md5,
                             ty = content_type,
                             date = date,
                             headers = headers,
                             resource = format!("/{}/{}", self.name, path));
        let signature = signature(&string, creds.aws_secret_access_key());
        format!("AWS {}:{}", creds.aws_access_key_id(), signature)
    }
}
<|endoftext|>


<commit_before>/// Strongly borrowed from TypedArena
/// Differs in growth strategy (non-doubling)
/// And only for variable-sized strings

use std::borrow::BorrowFrom;
use std::cell::{Cell,RefCell};
use std::cmp::max;
use std::collections::DList;
use std::collections::hash_map::Entry::{Occupied,Vacant};
use std::collections::hash_map::HashMap;
use std::collections::hash_state::DefaultState;
use std::default::Default;
use std::ops::Deref;
use std::raw::Slice;
use std::rt::heap::{allocate, deallocate};
use std::{fmt,hash,mem,ptr,str};
use xxhash::XXHasher;

struct Chunk {
    start: *mut u8,
    capacity: usize,
}

impl Chunk {
    fn new(capacity: usize) -> Chunk {
        Chunk {
            start: unsafe { allocate(capacity, mem::min_align_of::<u8>()) },
            capacity: capacity,
        }
    }

    // Returns a pointer to the beginning of the allocated space.
    #[inline]
    fn start(&self) -> *const u8 {
        self.start as *const u8
    }

    // Returns a pointer to the end of the allocated space.
    #[inline]
    fn end(&self) -> *const u8 {
        unsafe { self.start().offset(self.capacity as isize) }
    }
}

impl Drop for Chunk {
    fn drop(&mut self) {
        unsafe {
            deallocate(self.start, self.capacity, mem::min_align_of::<u8>());
        }
    }
}

#[derive(Copy)]
pub struct InternedString {
    slice: Slice<u8>,
}

impl InternedString {
    fn from_parts(data: *const u8, len: usize) -> InternedString {
        InternedString {
            slice: Slice {
                data: data,
                len: len,
            },
        }
    }

    pub fn from_str(s: &str) -> InternedString {
        let bytes: &[u8] = s.as_bytes();
        let slice: Slice<u8> = unsafe { mem::transmute(bytes) };
        InternedString {
            slice: slice,
        }
    }

    pub fn as_slice<'s>(&self) -> &'s str {
        unsafe {
            let bytes: &[u8] = mem::transmute(self.slice);
            str::from_utf8_unchecked(bytes)
        }
    }
}

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

impl PartialEq for InternedString {
    fn eq(&self, other: &InternedString) -> bool {
        self.as_slice().eq(other.as_slice())
    }
}

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

impl Eq for InternedString {}

impl<S> hash::Hash<S> for InternedString
    where S: hash::Writer + hash::Hasher
{
    fn hash(&self, state: &mut S) {
        self.as_slice().hash(state)
    }
}

impl BorrowFrom<InternedString> for str {
    fn borrow_from(owned: &InternedString) -> &str {
        owned.as_slice()
    }
}

impl Deref for InternedString {
    type Target = str;

    fn deref(&self) -> &str {
        self.as_slice()
    }
}

pub struct StringPool {
    start: Cell<*mut u8>,
    end: Cell<*const u8>,
    chunks: RefCell<DList<Chunk>>,
    index: RefCell<HashMap<InternedString, InternedString, DefaultState<XXHasher>>>,
}

static CAPACITY: usize = 10240;

impl StringPool {
    pub fn new() -> StringPool {
        StringPool{
            start: Cell::new(ptr::null_mut()),
            end: Cell::new(ptr::null()),
            chunks: RefCell::new(DList::new()),
            index: RefCell::new(Default::default()),
        }
    }

    pub fn intern<'s>(&'s self, s: &str) -> &'s str {
        if s == "" { return ""; }

        let search_string = InternedString::from_str(s);

        let interned_str = match self.index.borrow_mut().entry(search_string) {
            Occupied(entry) => *entry.get(),
            Vacant(entry) => *entry.insert(self.do_intern(s)),
        };

        // The lifetime is really matched to us
        unsafe { mem::transmute(interned_str) }
    }

    fn do_intern(&self, s: &str) -> InternedString {
        self.ensure_capacity(s.len());
        self.store(s)
    }

    fn ensure_capacity(&self, str_len: usize) {
        let remaining = self.end.get() as usize - self.start.get() as usize;
        if remaining < str_len {
            self.allocate_another(max(CAPACITY, str_len))
        }
    }

    fn allocate_another(&self, capacity: usize) {
        let chunk = Chunk::new(capacity);
        self.start.set(chunk.start() as *mut u8);
        self.end.set(chunk.end());
        self.chunks.borrow_mut().push_front(chunk);
    }

    fn store(&self, s: &str) -> InternedString {
        let str_start = s.as_bytes().as_ptr();
        let str_len = s.len();

        unsafe {
            ptr::copy_nonoverlapping_memory(self.start.get(), str_start, str_len);

            // Rebuild the string from our buffer
            let interned_str = InternedString::from_parts(self.start.get() as *const u8, str_len);

            // Increase current pointer
            self.start.set(self.start.get().offset(str_len as isize));

            interned_str
        }
    }
}

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

    #[test]
    fn keeps_the_same_string() {
        let s = StringPool::new();

        let interned = s.intern("hello");

        assert_eq!(interned, "hello");
    }

    #[test]
    fn does_not_reuse_the_pointer_of_the_input() {
        let s = StringPool::new();
        let input = "hello";

        let interned = s.intern(input);

        assert!(
            input.as_bytes().as_ptr() !=
            interned.as_bytes().as_ptr()
        );
    }

    #[test]
    fn reuses_the_pointer_for_repeated_input() {
        let s = StringPool::new();

        let interned1 = s.intern("world");
        let interned2 = s.intern("world");

        assert_eq!(
            interned1.as_bytes().as_ptr(),
            interned2.as_bytes().as_ptr()
        );
    }

    // #[test]
    // #[compile_failure]
    // fn string_cannot_outlive_the_pool() {
    //     let z = {
    //         let s = StringPool::new();
    //         s.intern("hello")
    //     };
    // }

    #[test]
    fn ignores_the_lifetime_of_the_input_string() {
        let s = StringPool::new();

        let interned = {
            let allocated_string = String::from_str("green");
            s.intern(&allocated_string)
        };

        // allocated_string is gone now, but we should be able to
        // access the result value until the storage goes away.

        assert_eq!(interned, "green");
    }

    #[test]
    fn can_be_dropped_immediately() {
        StringPool::new();
    }

    fn return_populated_storage() -> (StringPool, *const u8) {
        let s = StringPool::new();
        let ptr = {
            let interned = s.intern("hello");
            interned.as_ptr()
        };
        (s, ptr)
    }

    #[test]
    fn can_return_storage_populated_with_values() {
        let (s, ptr_val) = return_populated_storage();
        let interned = s.intern("hello");
        assert_eq!(interned.as_ptr(), ptr_val);
    }
}

#[cfg(test)]
mod bench {
    use test::Bencher;

    use super::StringPool;

    #[bench]
    fn single_string(b: &mut Bencher) {
        let s = StringPool::new();
        b.iter(|| s.intern("hello"));
        b.bytes = "hello".len() as u64;
    }

    #[bench]
    fn many_unique_string(b: &mut Bencher) {
        let s = StringPool::new();

        let strings: Vec<String> =
            range(0, 1000)
            .map(|i| format!("str{}str", i))
            .collect();
        b.iter(|| {
            for ss in strings.iter() {
                s.intern(ss);
            }
        });
        b.bytes = strings.iter().fold(0, |a, s| a + s.len()) as u64;
    }

    #[bench]
    fn many_repeated_string(b: &mut Bencher) {
        let s = StringPool::new();

        let strings: Vec<String> =
            range(0, 1000)
            .map(|i| format!("str{}str", i % 100))
            .collect();
        b.iter(|| {
            for ss in strings.iter() {
                s.intern(ss);
            }
        });
        b.bytes = strings.iter().fold(0, |a, s| a + s.len()) as u64;
    }
}
<commit_msg>BorrowFrom -> Borrow<commit_after>/// Strongly borrowed from TypedArena
/// Differs in growth strategy (non-doubling)
/// And only for variable-sized strings

use std::borrow::Borrow;
use std::cell::{Cell,RefCell};
use std::cmp::max;
use std::collections::DList;
use std::collections::hash_map::Entry::{Occupied,Vacant};
use std::collections::hash_map::HashMap;
use std::collections::hash_state::DefaultState;
use std::default::Default;
use std::ops::Deref;
use std::raw::Slice;
use std::rt::heap::{allocate, deallocate};
use std::{fmt,hash,mem,ptr,str};
use xxhash::XXHasher;

struct Chunk {
    start: *mut u8,
    capacity: usize,
}

impl Chunk {
    fn new(capacity: usize) -> Chunk {
        Chunk {
            start: unsafe { allocate(capacity, mem::min_align_of::<u8>()) },
            capacity: capacity,
        }
    }

    // Returns a pointer to the beginning of the allocated space.
    #[inline]
    fn start(&self) -> *const u8 {
        self.start as *const u8
    }

    // Returns a pointer to the end of the allocated space.
    #[inline]
    fn end(&self) -> *const u8 {
        unsafe { self.start().offset(self.capacity as isize) }
    }
}

impl Drop for Chunk {
    fn drop(&mut self) {
        unsafe {
            deallocate(self.start, self.capacity, mem::min_align_of::<u8>());
        }
    }
}

#[derive(Copy)]
pub struct InternedString {
    slice: Slice<u8>,
}

impl InternedString {
    fn from_parts(data: *const u8, len: usize) -> InternedString {
        InternedString {
            slice: Slice {
                data: data,
                len: len,
            },
        }
    }

    pub fn from_str(s: &str) -> InternedString {
        let bytes: &[u8] = s.as_bytes();
        let slice: Slice<u8> = unsafe { mem::transmute(bytes) };
        InternedString {
            slice: slice,
        }
    }

    pub fn as_slice<'s>(&self) -> &'s str {
        unsafe {
            let bytes: &[u8] = mem::transmute(self.slice);
            str::from_utf8_unchecked(bytes)
        }
    }
}

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

impl PartialEq for InternedString {
    fn eq(&self, other: &InternedString) -> bool {
        self.as_slice().eq(other.as_slice())
    }
}

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

impl Eq for InternedString {}

impl<S> hash::Hash<S> for InternedString
    where S: hash::Writer + hash::Hasher
{
    fn hash(&self, state: &mut S) {
        self.as_slice().hash(state)
    }
}

impl Borrow<str> for InternedString {
    fn borrow(&self) -> &str {
        self.as_slice()
    }
}

impl Deref for InternedString {
    type Target = str;

    fn deref(&self) -> &str {
        self.as_slice()
    }
}

pub struct StringPool {
    start: Cell<*mut u8>,
    end: Cell<*const u8>,
    chunks: RefCell<DList<Chunk>>,
    index: RefCell<HashMap<InternedString, InternedString, DefaultState<XXHasher>>>,
}

static CAPACITY: usize = 10240;

impl StringPool {
    pub fn new() -> StringPool {
        StringPool{
            start: Cell::new(ptr::null_mut()),
            end: Cell::new(ptr::null()),
            chunks: RefCell::new(DList::new()),
            index: RefCell::new(Default::default()),
        }
    }

    pub fn intern<'s>(&'s self, s: &str) -> &'s str {
        if s == "" { return ""; }

        let search_string = InternedString::from_str(s);

        let interned_str = match self.index.borrow_mut().entry(search_string) {
            Occupied(entry) => *entry.get(),
            Vacant(entry) => *entry.insert(self.do_intern(s)),
        };

        // The lifetime is really matched to us
        unsafe { mem::transmute(interned_str) }
    }

    fn do_intern(&self, s: &str) -> InternedString {
        self.ensure_capacity(s.len());
        self.store(s)
    }

    fn ensure_capacity(&self, str_len: usize) {
        let remaining = self.end.get() as usize - self.start.get() as usize;
        if remaining < str_len {
            self.allocate_another(max(CAPACITY, str_len))
        }
    }

    fn allocate_another(&self, capacity: usize) {
        let chunk = Chunk::new(capacity);
        self.start.set(chunk.start() as *mut u8);
        self.end.set(chunk.end());
        self.chunks.borrow_mut().push_front(chunk);
    }

    fn store(&self, s: &str) -> InternedString {
        let str_start = s.as_bytes().as_ptr();
        let str_len = s.len();

        unsafe {
            ptr::copy_nonoverlapping_memory(self.start.get(), str_start, str_len);

            // Rebuild the string from our buffer
            let interned_str = InternedString::from_parts(self.start.get() as *const u8, str_len);

            // Increase current pointer
            self.start.set(self.start.get().offset(str_len as isize));

            interned_str
        }
    }
}

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

    #[test]
    fn keeps_the_same_string() {
        let s = StringPool::new();

        let interned = s.intern("hello");

        assert_eq!(interned, "hello");
    }

    #[test]
    fn does_not_reuse_the_pointer_of_the_input() {
        let s = StringPool::new();
        let input = "hello";

        let interned = s.intern(input);

        assert!(
            input.as_bytes().as_ptr() !=
            interned.as_bytes().as_ptr()
        );
    }

    #[test]
    fn reuses_the_pointer_for_repeated_input() {
        let s = StringPool::new();

        let interned1 = s.intern("world");
        let interned2 = s.intern("world");

        assert_eq!(
            interned1.as_bytes().as_ptr(),
            interned2.as_bytes().as_ptr()
        );
    }

    // #[test]
    // #[compile_failure]
    // fn string_cannot_outlive_the_pool() {
    //     let z = {
    //         let s = StringPool::new();
    //         s.intern("hello")
    //     };
    // }

    #[test]
    fn ignores_the_lifetime_of_the_input_string() {
        let s = StringPool::new();

        let interned = {
            let allocated_string = String::from_str("green");
            s.intern(&allocated_string)
        };

        // allocated_string is gone now, but we should be able to
        // access the result value until the storage goes away.

        assert_eq!(interned, "green");
    }

    #[test]
    fn can_be_dropped_immediately() {
        StringPool::new();
    }

    fn return_populated_storage() -> (StringPool, *const u8) {
        let s = StringPool::new();
        let ptr = {
            let interned = s.intern("hello");
            interned.as_ptr()
        };
        (s, ptr)
    }

    #[test]
    fn can_return_storage_populated_with_values() {
        let (s, ptr_val) = return_populated_storage();
        let interned = s.intern("hello");
        assert_eq!(interned.as_ptr(), ptr_val);
    }
}

#[cfg(test)]
mod bench {
    use test::Bencher;

    use super::StringPool;

    #[bench]
    fn single_string(b: &mut Bencher) {
        let s = StringPool::new();
        b.iter(|| s.intern("hello"));
        b.bytes = "hello".len() as u64;
    }

    #[bench]
    fn many_unique_string(b: &mut Bencher) {
        let s = StringPool::new();

        let strings: Vec<String> =
            range(0, 1000)
            .map(|i| format!("str{}str", i))
            .collect();
        b.iter(|| {
            for ss in strings.iter() {
                s.intern(ss);
            }
        });
        b.bytes = strings.iter().fold(0, |a, s| a + s.len()) as u64;
    }

    #[bench]
    fn many_repeated_string(b: &mut Bencher) {
        let s = StringPool::new();

        let strings: Vec<String> =
            range(0, 1000)
            .map(|i| format!("str{}str", i % 100))
            .collect();
        b.iter(|| {
            for ss in strings.iter() {
                s.intern(ss);
            }
        });
        b.bytes = strings.iter().fold(0, |a, s| a + s.len()) as u64;
    }
}
<|endoftext|>


<commit_before>use grammar::*;
use grammar::Statement::*;
use grammar::Expression::*;
use grammar::ClassMember::*;
use grammar::OperatorType::*;

pub struct Settings {
    pub transform_block_scope: bool,
    pub transform_arrow: bool,
    pub transform_object: bool,
    pub transform_exponentation: bool,
}

#[inline(always)]
fn bind_this(function: Expression) -> Expression {
    Expression::call(Expression::member(function, "bind"), vec![ThisExpression])
}

impl Settings {
    pub fn target_es5() -> Settings {
        let mut settings = Settings::target_es2015();

        settings.transform_block_scope = true;
        settings.transform_arrow = true;
        settings.transform_object = true;

        settings
    }

    pub fn target_es2015() -> Settings {
        let mut settings = Settings::no_transform();

        settings.transform_exponentation = true;

        settings
    }

    pub fn no_transform() -> Settings {
        Settings {
            transform_block_scope: false,
            transform_arrow: false,
            transform_object: false,
            transform_exponentation: false,
        }
    }
}

/// The `Transformable` trait provides an interface for instances of grammar
/// to alter the AST, either by mutating self, or by returning a new node.
///
/// NOTE: Returning `None` means no changes are necessary!
trait Transformable {
    fn transform(&mut self, _: &Settings) {}

    fn contains_this(&self) -> bool {
        false
    }
}

impl Transformable for Parameter {}

impl Transformable for Expression {
    fn transform(&mut self, settings: &Settings) {
        *self = match *self {
            ArrowFunctionExpression {
                ref mut params,
                ref mut body,
            } => {
                params.transform(settings);
                body.transform(settings);

                // transformation flag check
                if !settings.transform_arrow {
                    return;
                }

                let body = match **body {
                    BlockStatement { ref mut body } => body.split_off(0),
                    ExpressionStatement(ref expr)   => vec![
                        ReturnStatement {
                            value: Some(expr.clone())
                        }
                    ],
                    ref statement => {
                        panic!("Invalid arrow function body {:#?}", statement);
                    }
                };

                let bind = body.contains_this();

                let function = FunctionExpression {
                    name: None,
                    params: params.split_off(0),
                    body: body,
                };

                if bind {
                    bind_this(function)
                } else {
                    function
                }
            },

            ArrayExpression(ref mut items) => {
                items.transform(settings);
                return;
            },

            ObjectExpression(ref mut members) => {
                members.transform(settings);

                // transformation flag check
                if !settings.transform_object {
                    return;
                }

                let (mut computed, literal): (Vec<ObjectMember>, Vec<ObjectMember>)
                = members.drain(..).partition(|member| {
                    match *member {
                        ObjectMember::Computed { .. } => true,
                        _                             => false,
                    }
                });

                if computed.is_empty() {
                    *members = literal;
                    return;
                }

                let mut body = Vec::new();

                body.push(VariableDeclarationStatement {
                    kind: VariableDeclarationKind::Var,
                    declarators: vec![
                        VariableDeclarator {
                            name: "___".to_string(),
                            value: Some(ObjectExpression(literal)),
                        }
                    ]
                });

                for member in computed.drain(..) {
                    if let ObjectMember::Computed { key, value } = member {
                        body.push(ExpressionStatement(BinaryExpression {
                            left: Box::new(MemberExpression {
                                object: Box::new(Expression::ident("___")),
                                property: Box::new(
                                    MemberKey::Computed(key)
                                )
                            }),
                            operator: Assign,
                            right: Box::new(value),
                        }));
                    }
                }

                body.push(ReturnStatement {
                    value: Some(Expression::ident("___"))
                });

                Expression::call(FunctionExpression {
                    name: None,
                    params: Vec::new(),
                    body: body,
                }, Vec::new())
            },

            CallExpression {
                ref mut callee,
                ref mut arguments,
            } => {
                callee.transform(settings);
                arguments.transform(settings);
                return;
            },

            BinaryExpression {
                ref mut left,
                ref mut operator,
                ref mut right,
            } => {
                left.transform(settings);
                right.transform(settings);

                if !settings.transform_exponentation {
                    return;
                }

                match *operator {
                    Exponent => Expression::call(
                        Expression::member(Expression::ident("Math"), "pow"),
                        vec![*left.clone(), *right.clone()]
                    ),

                    ExponentAssign => {
                        *operator = Assign;
                        *right = Box::new(Expression::call(
                            Expression::member(Expression::ident("Math"), "pow"),
                            vec![*left.clone(), *right.clone()]
                        ));
                        return;
                    },

                    _ => return,
                }
            }

            _ => return,
        }
    }

    fn contains_this(&self) -> bool {
        match *self {
            ThisExpression => true,

            ArrayExpression(ref items) => items.contains_this(),

            ObjectExpression(ref members) => members.contains_this(),

            MemberExpression {
                ref object,
                ..
            } => object.contains_this(),

            CallExpression {
                ref callee,
                ref arguments,
            } => callee.contains_this() || arguments.contains_this(),

            BinaryExpression {
                ref left,
                ref right,
                ..
            } => left.contains_this() || right.contains_this(),

            _ => false,
        }
    }
}

impl Transformable for ObjectMember {
    fn transform(&mut self, settings: &Settings) {
        *self = match *self {

            ObjectMember::Shorthand {
                ref key,
            } => {
                // transformation flag check
                if !settings.transform_object {
                    return;
                }

                ObjectMember::Literal {
                    key: key.clone(),
                    value: IdentifierExpression(key.clone()),
                }
            },

            ObjectMember::Literal {
                ref mut value,
                ..
            } => {
                value.transform(settings);
                return;
            },

            ObjectMember::Computed {
                ref mut key,
                ref mut value,
            } => {
                key.transform(settings);
                value.transform(settings);
                return;
            },

            ObjectMember::Method {
                ref mut name,
                ref mut params,
                ref mut body,
            } => {
                body.transform(settings);
                params.transform(settings);

                // transformation flag check
                if !settings.transform_object {
                    return;
                }

                ObjectMember::Literal {
                    key: name.clone(),
                    value: FunctionExpression {
                        name: Some(name.clone()),
                        params: params.split_off(0),
                        body: body.split_off(0),
                    }
                }
            },

            ObjectMember::ComputedMethod {
                ref mut name,
                ref mut params,
                ref mut body,
            } => {
                body.transform(settings);
                params.transform(settings);

                // transformation flag check
                if !settings.transform_object {
                    return;
                }

                ObjectMember::Computed {
                    key: name.clone(),
                    value: FunctionExpression {
                        name: None,
                        params: params.split_off(0),
                        body: body.split_off(0),
                    }
                }
            },
        }
    }

    fn contains_this(&self) -> bool {
        match *self {
            ObjectMember::Literal {
                ref value,
                ..
            } => value.contains_this(),

            ObjectMember::Computed {
                ref key,
                ref value,
            } => key.contains_this() || value.contains_this(),

            _ => false,
        }
    }
}

impl Transformable for ClassMember {
    fn transform(&mut self, settings: &Settings) {
        match *self {
            Constructor {
                ref mut params,
                ref mut body,
            } => {
                params.transform(settings);
                body.transform(settings);
            },

            Method {
                ref mut params,
                ref mut body,
                ..
            } => {
                params.transform(settings);
                body.transform(settings);
            },

            Property {
                ref mut value,
                ..
            } => {
                value.transform(settings);
            }
        }
    }
}

impl Transformable for VariableDeclarator {
    fn transform(&mut self, settings: &Settings) {
        match self.value {
            Some(ref mut expression) => expression.transform(settings),
            _                        => {},
        }
    }

    fn contains_this(&self) -> bool {
        match self.value {
            Some(ref expression) => expression.contains_this(),
            _                    => false,
        }
    }
}

impl Transformable for Statement {
    fn transform(&mut self, settings: &Settings) {
        *self = match *self {
            BlockStatement {
                ref mut body,
            } => {
                body.transform(settings);
                return;
            },

            LabeledStatement {
                ref mut body,
                ..
            } => {
                body.transform(settings);
                return;
            },

            VariableDeclarationStatement {
                ref mut kind,
                ref mut declarators,
            } => {
                declarators.transform(settings);

                // transformation flag check
                if !settings.transform_block_scope {
                    return;
                }

                *kind = VariableDeclarationKind::Var;
                return;
            },

            ExpressionStatement(ref mut expression) => {
                expression.transform(settings);
                return;
            },

            IfStatement {
                ref mut test,
                ref mut consequent,
                ref mut alternate,
                ..
            } => {
                test.transform(settings);
                consequent.transform(settings);
                if let Some(ref mut alternate) = *alternate {
                    alternate.transform(settings);
                }
                return;
            },

            ClassStatement {
                ref mut body,
                ..
            } => {
                body.transform(settings);
                return;
            }

            _ => return,
        }
    }

    fn contains_this(&self) -> bool {
        match *self {
            VariableDeclarationStatement {
                ref declarators,
                ..
            } => declarators.contains_this(),

            ExpressionStatement(ref expression) => expression.contains_this(),

            ReturnStatement {
                value: Some(ref expression)
            } => expression.contains_this(),

            _ => false,
        }
    }
}

impl<T: Transformable> Transformable for Vec<T> {
    fn transform(&mut self, settings: &Settings) {
        for item in self.iter_mut() {
            item.transform(settings);
        }
    }

    fn contains_this(&self) -> bool {
        for item in self {
            if item.contains_this() {
                return true;
            }
        }
        return false;
    }
}

pub fn transform(program: &mut Program, settings: Settings) {
    program.body.transform(&settings);
}
<commit_msg>Starting class transformers<commit_after>use grammar::*;
use grammar::Statement::*;
use grammar::Expression::*;
use grammar::ClassMember::*;
use grammar::OperatorType::*;

pub struct Settings {
    pub transform_block_scope: bool,
    pub transform_arrow: bool,
    pub transform_object: bool,
    pub transform_exponentation: bool,
    pub transform_class_properties: bool,
}

#[inline(always)]
fn bind_this(function: Expression) -> Expression {
    Expression::call(Expression::member(function, "bind"), vec![ThisExpression])
}

impl Settings {
    pub fn target_es5() -> Settings {
        let mut settings = Settings::target_es2015();

        settings.transform_block_scope = true;
        settings.transform_arrow = true;
        settings.transform_object = true;

        settings
    }

    pub fn target_es2015() -> Settings {
        let mut settings = Settings::no_transform();

        settings.transform_exponentation = true;
        settings.transform_class_properties = true;

        settings
    }

    pub fn no_transform() -> Settings {
        Settings {
            transform_block_scope: false,
            transform_arrow: false,
            transform_object: false,
            transform_exponentation: false,
            transform_class_properties: false,
        }
    }
}

/// The `Transformable` trait provides an interface for instances of grammar
/// to alter the AST, either by mutating self, or by returning a new node.
///
/// NOTE: Returning `None` means no changes are necessary!
trait Transformable {
    fn transform(&mut self, _: &Settings) {}

    fn contains_this(&self) -> bool {
        false
    }
}

impl Transformable for Parameter {}

impl Transformable for Expression {
    fn transform(&mut self, settings: &Settings) {
        *self = match *self {
            ArrowFunctionExpression {
                ref mut params,
                ref mut body,
            } => {
                params.transform(settings);
                body.transform(settings);

                // transformation flag check
                if !settings.transform_arrow {
                    return;
                }

                let body = match **body {
                    BlockStatement { ref mut body } => body.split_off(0),
                    ExpressionStatement(ref expr)   => vec![
                        ReturnStatement {
                            value: Some(expr.clone())
                        }
                    ],
                    ref statement => {
                        panic!("Invalid arrow function body {:#?}", statement);
                    }
                };

                let bind = body.contains_this();

                let function = FunctionExpression {
                    name: None,
                    params: params.split_off(0),
                    body: body,
                };

                if bind {
                    bind_this(function)
                } else {
                    function
                }
            },

            ArrayExpression(ref mut items) => {
                items.transform(settings);
                return;
            },

            ObjectExpression(ref mut members) => {
                members.transform(settings);

                // transformation flag check
                if !settings.transform_object {
                    return;
                }

                let (mut computed, literal): (Vec<ObjectMember>, Vec<ObjectMember>)
                = members.drain(..).partition(|member| {
                    match *member {
                        ObjectMember::Computed { .. } => true,
                        _                             => false,
                    }
                });

                if computed.is_empty() {
                    *members = literal;
                    return;
                }

                let mut body = Vec::new();

                body.push(VariableDeclarationStatement {
                    kind: VariableDeclarationKind::Var,
                    declarators: vec![
                        VariableDeclarator {
                            name: "___".to_string(),
                            value: Some(ObjectExpression(literal)),
                        }
                    ]
                });

                for member in computed.drain(..) {
                    if let ObjectMember::Computed { key, value } = member {
                        body.push(ExpressionStatement(BinaryExpression {
                            left: Box::new(MemberExpression {
                                object: Box::new(Expression::ident("___")),
                                property: Box::new(
                                    MemberKey::Computed(key)
                                )
                            }),
                            operator: Assign,
                            right: Box::new(value),
                        }));
                    }
                }

                body.push(ReturnStatement {
                    value: Some(Expression::ident("___"))
                });

                Expression::call(FunctionExpression {
                    name: None,
                    params: Vec::new(),
                    body: body,
                }, Vec::new())
            },

            CallExpression {
                ref mut callee,
                ref mut arguments,
            } => {
                callee.transform(settings);
                arguments.transform(settings);
                return;
            },

            BinaryExpression {
                ref mut left,
                ref mut operator,
                ref mut right,
            } => {
                left.transform(settings);
                right.transform(settings);

                if !settings.transform_exponentation {
                    return;
                }

                match *operator {
                    Exponent => Expression::call(
                        Expression::member(Expression::ident("Math"), "pow"),
                        vec![*left.clone(), *right.clone()]
                    ),

                    ExponentAssign => {
                        *operator = Assign;
                        *right = Box::new(Expression::call(
                            Expression::member(Expression::ident("Math"), "pow"),
                            vec![*left.clone(), *right.clone()]
                        ));
                        return;
                    },

                    _ => return,
                }
            }

            _ => return,
        }
    }

    fn contains_this(&self) -> bool {
        match *self {
            ThisExpression => true,

            ArrayExpression(ref items) => items.contains_this(),

            ObjectExpression(ref members) => members.contains_this(),

            MemberExpression {
                ref object,
                ..
            } => object.contains_this(),

            CallExpression {
                ref callee,
                ref arguments,
            } => callee.contains_this() || arguments.contains_this(),

            BinaryExpression {
                ref left,
                ref right,
                ..
            } => left.contains_this() || right.contains_this(),

            _ => false,
        }
    }
}

impl Transformable for ObjectMember {
    fn transform(&mut self, settings: &Settings) {
        *self = match *self {

            ObjectMember::Shorthand {
                ref key,
            } => {
                // transformation flag check
                if !settings.transform_object {
                    return;
                }

                ObjectMember::Literal {
                    key: key.clone(),
                    value: IdentifierExpression(key.clone()),
                }
            },

            ObjectMember::Literal {
                ref mut value,
                ..
            } => {
                value.transform(settings);
                return;
            },

            ObjectMember::Computed {
                ref mut key,
                ref mut value,
            } => {
                key.transform(settings);
                value.transform(settings);
                return;
            },

            ObjectMember::Method {
                ref mut name,
                ref mut params,
                ref mut body,
            } => {
                body.transform(settings);
                params.transform(settings);

                // transformation flag check
                if !settings.transform_object {
                    return;
                }

                ObjectMember::Literal {
                    key: name.clone(),
                    value: FunctionExpression {
                        name: Some(name.clone()),
                        params: params.split_off(0),
                        body: body.split_off(0),
                    }
                }
            },

            ObjectMember::ComputedMethod {
                ref mut name,
                ref mut params,
                ref mut body,
            } => {
                body.transform(settings);
                params.transform(settings);

                // transformation flag check
                if !settings.transform_object {
                    return;
                }

                ObjectMember::Computed {
                    key: name.clone(),
                    value: FunctionExpression {
                        name: None,
                        params: params.split_off(0),
                        body: body.split_off(0),
                    }
                }
            },
        }
    }

    fn contains_this(&self) -> bool {
        match *self {
            ObjectMember::Literal {
                ref value,
                ..
            } => value.contains_this(),

            ObjectMember::Computed {
                ref key,
                ref value,
            } => key.contains_this() || value.contains_this(),

            _ => false,
        }
    }
}

impl Transformable for ClassMember {
    fn transform(&mut self, settings: &Settings) {
        match *self {
            Constructor {
                ref mut params,
                ref mut body,
            } => {
                params.transform(settings);
                body.transform(settings);
            },

            Method {
                ref mut params,
                ref mut body,
                ..
            } => {
                params.transform(settings);
                body.transform(settings);
            },

            Property {
                ref mut value,
                ..
            } => {
                value.transform(settings);
            }
        }
    }
}

impl Transformable for VariableDeclarator {
    fn transform(&mut self, settings: &Settings) {
        match self.value {
            Some(ref mut expression) => expression.transform(settings),
            _                        => {},
        }
    }

    fn contains_this(&self) -> bool {
        match self.value {
            Some(ref expression) => expression.contains_this(),
            _                    => false,
        }
    }
}

impl Transformable for Statement {
    fn transform(&mut self, settings: &Settings) {
        *self = match *self {
            BlockStatement {
                ref mut body,
            } => {
                body.transform(settings);
                return;
            },

            LabeledStatement {
                ref mut body,
                ..
            } => {
                body.transform(settings);
                return;
            },

            VariableDeclarationStatement {
                ref mut kind,
                ref mut declarators,
            } => {
                declarators.transform(settings);

                // transformation flag check
                if !settings.transform_block_scope {
                    return;
                }

                *kind = VariableDeclarationKind::Var;
                return;
            },

            ExpressionStatement(ref mut expression) => {
                expression.transform(settings);
                return;
            },

            IfStatement {
                ref mut test,
                ref mut consequent,
                ref mut alternate,
                ..
            } => {
                test.transform(settings);
                consequent.transform(settings);
                if let Some(ref mut alternate) = *alternate {
                    alternate.transform(settings);
                }
                return;
            },

            ClassStatement {
                ref mut body,
                ..
            } => {
                body.transform(settings);

                if !settings.transform_class_properties {
                    return;
                }

                let (mut props, methods): (Vec<ClassMember>, Vec<ClassMember>)
                = body.drain(..).partition(|member| {
                    match *member {
                        ClassMember::Property { .. } => true,
                        _                            => false,
                    }
                });

                *body = methods;

                if !props.is_empty() {
                    // copy the props into the class constructor
                }

                return;
            }

            _ => return,
        }
    }

    fn contains_this(&self) -> bool {
        match *self {
            VariableDeclarationStatement {
                ref declarators,
                ..
            } => declarators.contains_this(),

            ExpressionStatement(ref expression) => expression.contains_this(),

            ReturnStatement {
                value: Some(ref expression)
            } => expression.contains_this(),

            _ => false,
        }
    }
}

impl<T: Transformable> Transformable for Vec<T> {
    fn transform(&mut self, settings: &Settings) {
        for item in self.iter_mut() {
            item.transform(settings);
        }
    }

    fn contains_this(&self) -> bool {
        for item in self {
            if item.contains_this() {
                return true;
            }
        }
        return false;
    }
}

pub fn transform(program: &mut Program, settings: Settings) {
    program.body.transform(&settings);
}
<|endoftext|>


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

//! The macros crate contains all useful needed macros.

/// Get the count of macro's arguments.
///
/// # Examples
///
/// ```
/// # #[macro_use] extern crate tikv;
/// # fn main() {
/// assert_eq!(count_args!(), 0);
/// assert_eq!(count_args!(1), 1);
/// assert_eq!(count_args!(1, 2), 2);
/// assert_eq!(count_args!(1, 2, 3), 3);
/// # }
/// ```
///
/// [rfc#88](https://github.com/rust-lang/rfcs/pull/88) proposes to use $# to count the number
/// of args, but it has not been implemented yet.
#[macro_export]
macro_rules! count_args {
    () => { 0 };
    ($head:expr $(, $tail:expr)*) => { 1 + count_args!($($tail),*) };
}

/// Initial a `HashMap` with specify key-value pairs.
///
/// # Examples
///
/// ```
/// # #[macro_use] extern crate tikv;
/// # fn main() {
/// // empty map
/// let m: tikv::util::collections::HashMap<u8, u8> = map!();
/// assert!(m.is_empty());
///
/// // one initial kv pairs.
/// let m = map!("key" => "value");
/// assert_eq!(m.len(), 1);
/// assert_eq!(m["key"], "value");
///
/// // initialize with multiple kv pairs.
/// let m = map!("key1" => "value1", "key2" => "value2");
/// assert_eq!(m.len(), 2);
/// assert_eq!(m["key1"], "value1");
/// assert_eq!(m["key2"], "value2");
/// # }
/// ```
///
/// This macro may be removed once
/// [official implementation](https://github.com/rust-lang/rfcs/issues/542) is provided.
#[macro_export]
macro_rules! map {
    () => {
        {
            $crate::util::collections::HashMap::default()
        }
    };
    ( $( $k:expr => $v:expr ),+ ) => {
        {
            let mut temp_map =
                $crate::util::collections::HashMap::with_capacity_and_hasher(
                    count_args!($(($k, $v)),+),
                    Default::default()
                );
            $(
                temp_map.insert($k, $v);
            )+
            temp_map
        }
    };
}

/// box try will box error first, and then do the same thing as try!.
#[macro_export]
macro_rules! box_try {
    ($expr:expr) => ({
        match $expr {
            Ok(r) => r,
            Err(e) => return Err(box_err!(e)),
        }
    })
}

/// A shortcut to box an error.
#[macro_export]
macro_rules! box_err {
    ($e:expr) => ({
        use std::error::Error;
        let e: Box<Error + Sync + Send> = ($e).into();
        e.into()
    });
    ($f:tt, $($arg:expr),+) => ({
        box_err!(format!($f, $($arg),+))
    });
}

#[allow(doc_markdown)]
/// Recover from panicable closure.
///
/// Please note that this macro assume the closure is able to be forced to implement `UnwindSafe`.
/// Also see https://doc.rust-lang.org/std/panic/struct.AssertUnwindSafe.html
// Maybe we should define a recover macro too.
#[macro_export]
macro_rules! recover_safe {
    ($e:expr) => ({
        use std::panic::{AssertUnwindSafe, catch_unwind};
        use $crate::util::panic_hook;
        panic_hook::mute();
        let res = catch_unwind(AssertUnwindSafe($e));
        panic_hook::unmute();
        res
    })
}

/// Log slow operations with warn!.
macro_rules! slow_log {
    ($t:expr, $($arg:tt)*) => {{
        if $t.is_slow() {
            warn!("{} [takes {:?}]", format_args!($($arg)*), $t.elapsed());
        }
    }}
}

/// make a thread name with additional tag inheriting from current thread.
#[macro_export]
macro_rules! thd_name {
    ($name:expr) => ({
        $crate::util::get_tag_from_thread_name().map(|tag| {
            format!("{}::{}", $name, tag)
        }).unwrap_or_else(|| $name.to_owned())
    });
}

/// Simulating go's defer.
///
/// Please note that, different from go, this defer is bound to scope.
#[macro_export]
macro_rules! defer {
    ($t:expr) => (
        let __ctx = $crate::util::DeferContext::new(|| $t);
    );
}

/// `wait_op!` waits for async operation. It returns `Option<Res>`
/// after the expression get executed.
#[macro_export]
macro_rules! wait_op {
    ($expr:expr) => {
        wait_op!(IMPL $expr, None)
    };
    ($expr:expr, $timeout:expr) => {
        wait_op!(IMPL $expr, Some($timeout))
    };
    (IMPL $expr:expr, $timeout:expr) => {
        {
            use std::sync::mpsc;
            let (tx, rx) = mpsc::channel();
            let cb = box move |res| {
                 // we don't care error actually.
                let _ = tx.send(res);
            };
            $expr(cb);
            match $timeout {
                None => rx.recv().ok(),
                Some(timeout) => rx.recv_timeout(timeout).ok()
            }
        }
    }
}

/// `try_opt` check `Result<Option<T>>`, return early when met `Err` or `Ok(None)`.
#[macro_export]
macro_rules! try_opt {
    ($expr:expr) => ({
        match $expr {
            Err(e) => return Err(e.into()),
            Ok(None) => return Ok(None),
            Ok(Some(res)) => res,
        }
    });
}
<commit_msg>util/macros: update count_args doc (#2417)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

//! The macros crate contains all useful needed macros.

/// Get the count of macro's arguments.
///
/// # Examples
///
/// ```
/// # #[macro_use] extern crate tikv;
/// # fn main() {
/// assert_eq!(count_args!(), 0);
/// assert_eq!(count_args!(1), 1);
/// assert_eq!(count_args!(1, 2), 2);
/// assert_eq!(count_args!(1, 2, 3), 3);
/// # }
/// ```
#[macro_export]
macro_rules! count_args {
    () => { 0 };
    ($head:expr $(, $tail:expr)*) => { 1 + count_args!($($tail),*) };
}

/// Initial a `HashMap` with specify key-value pairs.
///
/// # Examples
///
/// ```
/// # #[macro_use] extern crate tikv;
/// # fn main() {
/// // empty map
/// let m: tikv::util::collections::HashMap<u8, u8> = map!();
/// assert!(m.is_empty());
///
/// // one initial kv pairs.
/// let m = map!("key" => "value");
/// assert_eq!(m.len(), 1);
/// assert_eq!(m["key"], "value");
///
/// // initialize with multiple kv pairs.
/// let m = map!("key1" => "value1", "key2" => "value2");
/// assert_eq!(m.len(), 2);
/// assert_eq!(m["key1"], "value1");
/// assert_eq!(m["key2"], "value2");
/// # }
/// ```
///
/// This macro may be removed once
/// [official implementation](https://github.com/rust-lang/rfcs/issues/542) is provided.
#[macro_export]
macro_rules! map {
    () => {
        {
            $crate::util::collections::HashMap::default()
        }
    };
    ( $( $k:expr => $v:expr ),+ ) => {
        {
            let mut temp_map =
                $crate::util::collections::HashMap::with_capacity_and_hasher(
                    count_args!($(($k, $v)),+),
                    Default::default()
                );
            $(
                temp_map.insert($k, $v);
            )+
            temp_map
        }
    };
}

/// box try will box error first, and then do the same thing as try!.
#[macro_export]
macro_rules! box_try {
    ($expr:expr) => ({
        match $expr {
            Ok(r) => r,
            Err(e) => return Err(box_err!(e)),
        }
    })
}

/// A shortcut to box an error.
#[macro_export]
macro_rules! box_err {
    ($e:expr) => ({
        use std::error::Error;
        let e: Box<Error + Sync + Send> = ($e).into();
        e.into()
    });
    ($f:tt, $($arg:expr),+) => ({
        box_err!(format!($f, $($arg),+))
    });
}

#[allow(doc_markdown)]
/// Recover from panicable closure.
///
/// Please note that this macro assume the closure is able to be forced to implement `UnwindSafe`.
/// Also see https://doc.rust-lang.org/std/panic/struct.AssertUnwindSafe.html
// Maybe we should define a recover macro too.
#[macro_export]
macro_rules! recover_safe {
    ($e:expr) => ({
        use std::panic::{AssertUnwindSafe, catch_unwind};
        use $crate::util::panic_hook;
        panic_hook::mute();
        let res = catch_unwind(AssertUnwindSafe($e));
        panic_hook::unmute();
        res
    })
}

/// Log slow operations with warn!.
macro_rules! slow_log {
    ($t:expr, $($arg:tt)*) => {{
        if $t.is_slow() {
            warn!("{} [takes {:?}]", format_args!($($arg)*), $t.elapsed());
        }
    }}
}

/// make a thread name with additional tag inheriting from current thread.
#[macro_export]
macro_rules! thd_name {
    ($name:expr) => ({
        $crate::util::get_tag_from_thread_name().map(|tag| {
            format!("{}::{}", $name, tag)
        }).unwrap_or_else(|| $name.to_owned())
    });
}

/// Simulating go's defer.
///
/// Please note that, different from go, this defer is bound to scope.
#[macro_export]
macro_rules! defer {
    ($t:expr) => (
        let __ctx = $crate::util::DeferContext::new(|| $t);
    );
}

/// `wait_op!` waits for async operation. It returns `Option<Res>`
/// after the expression get executed.
#[macro_export]
macro_rules! wait_op {
    ($expr:expr) => {
        wait_op!(IMPL $expr, None)
    };
    ($expr:expr, $timeout:expr) => {
        wait_op!(IMPL $expr, Some($timeout))
    };
    (IMPL $expr:expr, $timeout:expr) => {
        {
            use std::sync::mpsc;
            let (tx, rx) = mpsc::channel();
            let cb = box move |res| {
                 // we don't care error actually.
                let _ = tx.send(res);
            };
            $expr(cb);
            match $timeout {
                None => rx.recv().ok(),
                Some(timeout) => rx.recv_timeout(timeout).ok()
            }
        }
    }
}

/// `try_opt` check `Result<Option<T>>`, return early when met `Err` or `Ok(None)`.
#[macro_export]
macro_rules! try_opt {
    ($expr:expr) => ({
        match $expr {
            Err(e) => return Err(e.into()),
            Ok(None) => return Ok(None),
            Ok(Some(res)) => res,
        }
    });
}
<|endoftext|>


<commit_before>use bls::get_withdrawal_credentials;
use deposit_contract::{encode_eth1_tx_data, DEPOSIT_GAS};
use futures::compat::Future01CompatExt;
use hex;
use ssz::{Decode, Encode};
use ssz_derive::{Decode, Encode};
use std::fs;
use std::fs::File;
use std::io::prelude::*;
use std::os::unix::fs::PermissionsExt;
use std::path::PathBuf;
use types::{
    test_utils::generate_deterministic_keypair, ChainSpec, DepositData, Hash256, Keypair,
    PublicKey, SecretKey, Signature,
};
use web3::{
    types::{Address, TransactionRequest, U256},
    Transport, Web3,
};

const VOTING_KEY_PREFIX: &str = "voting";
const WITHDRAWAL_KEY_PREFIX: &str = "withdrawal";
const ETH1_DEPOSIT_DATA_FILE: &str = "eth1_deposit_data.rlp";

/// Returns the filename of a keypair file.
fn keypair_file(prefix: &str) -> String {
    format!("{}_keypair", prefix)
}

/// Returns the name of the folder to be generated for a validator with the given voting key.
fn dir_name(voting_pubkey: &PublicKey) -> String {
    format!("0x{}", hex::encode(voting_pubkey.as_ssz_bytes()))
}

/// Represents the files/objects for each dedicated lighthouse validator directory.
///
/// Generally lives in `~/.lighthouse/validators/`.
#[derive(Debug, PartialEq, Clone)]
pub struct ValidatorDirectory {
    pub directory: PathBuf,
    pub voting_keypair: Option<Keypair>,
    pub withdrawal_keypair: Option<Keypair>,
    pub deposit_data: Option<Vec<u8>>,
}

impl ValidatorDirectory {
    /// Attempts to load a validator from the given directory, requiring only components necessary
    /// for signing messages.
    pub fn load_for_signing(directory: PathBuf) -> Result<Self, String> {
        if !directory.exists() {
            return Err(format!(
                "Validator directory does not exist: {:?}",
                directory
            ));
        }

        Ok(Self {
            voting_keypair: Some(
                load_keypair(directory.clone(), VOTING_KEY_PREFIX)
                    .map_err(|e| format!("Unable to get voting keypair: {}", e))?,
            ),
            withdrawal_keypair: load_keypair(directory.clone(), WITHDRAWAL_KEY_PREFIX).ok(),
            deposit_data: load_eth1_deposit_data(directory.clone()).ok(),
            directory,
        })
    }
}

/// Load a `Keypair` from a file.
fn load_keypair(base_path: PathBuf, file_prefix: &str) -> Result<Keypair, String> {
    let path = base_path.join(keypair_file(file_prefix));

    if !path.exists() {
        return Err(format!("Keypair file does not exist: {:?}", path));
    }

    let mut bytes = vec![];

    File::open(&path)
        .map_err(|e| format!("Unable to open keypair file: {}", e))?
        .read_to_end(&mut bytes)
        .map_err(|e| format!("Unable to read keypair file: {}", e))?;

    SszEncodableKeypair::from_ssz_bytes(&bytes)
        .map(Into::into)
        .map_err(|e| format!("Unable to decode keypair: {:?}", e))
}

/// Load eth1_deposit_data from file.
fn load_eth1_deposit_data(base_path: PathBuf) -> Result<Vec<u8>, String> {
    let path = base_path.join(ETH1_DEPOSIT_DATA_FILE);

    if !path.exists() {
        return Err(format!("Eth1 deposit data file does not exist: {:?}", path));
    }

    let mut bytes = vec![];

    File::open(&path)
        .map_err(|e| format!("Unable to open eth1 deposit data file: {}", e))?
        .read_to_end(&mut bytes)
        .map_err(|e| format!("Unable to read eth1 deposit data file: {}", e))?;

    let string = String::from_utf8_lossy(&bytes);
    if string.starts_with("0x") {
        hex::decode(&string[2..])
            .map_err(|e| format!("Unable to decode eth1 data file as hex: {}", e))
    } else {
        Err(format!("String did not start with 0x: {}", string))
    }
}

/// A helper struct to allow SSZ enc/dec for a `Keypair`.
#[derive(Encode, Decode)]
struct SszEncodableKeypair {
    pk: PublicKey,
    sk: SecretKey,
}

impl Into<Keypair> for SszEncodableKeypair {
    fn into(self) -> Keypair {
        Keypair {
            sk: self.sk,
            pk: self.pk,
        }
    }
}

impl From<Keypair> for SszEncodableKeypair {
    fn from(kp: Keypair) -> Self {
        Self {
            sk: kp.sk,
            pk: kp.pk,
        }
    }
}

/// Builds a `ValidatorDirectory`, both in-memory and on-disk.
#[derive(Default)]
pub struct ValidatorDirectoryBuilder {
    directory: Option<PathBuf>,
    voting_keypair: Option<Keypair>,
    withdrawal_keypair: Option<Keypair>,
    amount: Option<u64>,
    deposit_data: Option<Vec<u8>>,
    spec: Option<ChainSpec>,
}

impl ValidatorDirectoryBuilder {
    /// Set the specification for this validator.
    pub fn spec(mut self, spec: ChainSpec) -> Self {
        self.spec = Some(spec);
        self
    }

    /// Use the `MAX_EFFECTIVE_BALANCE` as this validator's deposit.
    pub fn full_deposit_amount(mut self) -> Result<Self, String> {
        let spec = self
            .spec
            .as_ref()
            .ok_or_else(|| "full_deposit_amount requires a spec")?;
        self.amount = Some(spec.max_effective_balance);
        Ok(self)
    }

    /// Use a validator deposit of `gwei`.
    pub fn custom_deposit_amount(mut self, gwei: u64) -> Self {
        self.amount = Some(gwei);
        self
    }

    /// Generate keypairs using `Keypair::random()`.
    pub fn thread_random_keypairs(mut self) -> Self {
        self.voting_keypair = Some(Keypair::random());
        self.withdrawal_keypair = Some(Keypair::random());
        self
    }

    /// Generate insecure, deterministic keypairs.
    ///
    ///
    /// ## Warning
    /// Only for use in testing. Do not store value in these keys.
    pub fn insecure_keypairs(mut self, index: usize) -> Self {
        let keypair = generate_deterministic_keypair(index);
        self.voting_keypair = Some(keypair.clone());
        self.withdrawal_keypair = Some(keypair);
        self
    }

    /// Creates a validator directory in the given `base_path` (e.g., `~/.lighthouse/validators/`).
    pub fn create_directory(mut self, base_path: PathBuf) -> Result<Self, String> {
        let voting_keypair = self
            .voting_keypair
            .as_ref()
            .ok_or_else(|| "directory requires a voting_keypair")?;

        let directory = base_path.join(dir_name(&voting_keypair.pk));

        if directory.exists() {
            return Err(format!(
                "Validator directory already exists: {:?}",
                directory
            ));
        }

        fs::create_dir_all(&directory)
            .map_err(|e| format!("Unable to create validator directory: {}", e))?;

        self.directory = Some(directory);

        Ok(self)
    }

    /// Write the validators keypairs to disk.
    pub fn write_keypair_files(self) -> Result<Self, String> {
        let voting_keypair = self
            .voting_keypair
            .clone()
            .ok_or_else(|| "write_keypair_files requires a voting_keypair")?;
        let withdrawal_keypair = self
            .withdrawal_keypair
            .clone()
            .ok_or_else(|| "write_keypair_files requires a withdrawal_keypair")?;

        self.save_keypair(voting_keypair, VOTING_KEY_PREFIX)?;
        self.save_keypair(withdrawal_keypair, WITHDRAWAL_KEY_PREFIX)?;
        Ok(self)
    }

    fn save_keypair(&self, keypair: Keypair, file_prefix: &str) -> Result<(), String> {
        let path = self
            .directory
            .as_ref()
            .map(|directory| directory.join(keypair_file(file_prefix)))
            .ok_or_else(|| "save_keypair requires a directory")?;

        if path.exists() {
            return Err(format!("Keypair file already exists at: {:?}", path));
        }

        let mut file = File::create(&path).map_err(|e| format!("Unable to create file: {}", e))?;

        // Ensure file has correct permissions.
        let mut perm = file
            .metadata()
            .map_err(|e| format!("Unable to get file metadata: {}", e))?
            .permissions();
        perm.set_mode((libc::S_IWUSR | libc::S_IRUSR) as u32);
        file.set_permissions(perm)
            .map_err(|e| format!("Unable to set file permissions: {}", e))?;

        file.write_all(&SszEncodableKeypair::from(keypair).as_ssz_bytes())
            .map_err(|e| format!("Unable to write keypair to file: {}", e))?;

        Ok(())
    }

    fn get_deposit_data(&self) -> Result<(Vec<u8>, u64), String> {
        let voting_keypair = self
            .voting_keypair
            .as_ref()
            .ok_or_else(|| "write_eth1_data_file requires a voting_keypair")?;
        let withdrawal_keypair = self
            .withdrawal_keypair
            .as_ref()
            .ok_or_else(|| "write_eth1_data_file requires a withdrawal_keypair")?;
        let amount = self
            .amount
            .ok_or_else(|| "write_eth1_data_file requires an amount")?;
        let spec = self.spec.as_ref().ok_or_else(|| "build requires a spec")?;

        let withdrawal_credentials = Hash256::from_slice(&get_withdrawal_credentials(
            &withdrawal_keypair.pk,
            spec.bls_withdrawal_prefix_byte,
        ));

        let mut deposit_data = DepositData {
            pubkey: voting_keypair.pk.clone().into(),
            withdrawal_credentials,
            amount,
            signature: Signature::empty_signature().into(),
        };

        deposit_data.signature = deposit_data.create_signature(&voting_keypair.sk, &spec);

        let deposit_data = encode_eth1_tx_data(&deposit_data)
            .map_err(|e| format!("Unable to encode eth1 deposit tx data: {:?}", e))?;

        Ok((deposit_data, amount))
    }

    pub fn write_eth1_data_file(mut self) -> Result<Self, String> {
        let path = self
            .directory
            .as_ref()
            .map(|directory| directory.join(ETH1_DEPOSIT_DATA_FILE))
            .ok_or_else(|| "write_eth1_data_file requires a directory")?;

        let (deposit_data, _) = self.get_deposit_data()?;

        if path.exists() {
            return Err(format!("Eth1 data file already exists at: {:?}", path));
        }

        File::create(&path)
            .map_err(|e| format!("Unable to create file: {}", e))?
            .write_all(&format!("0x{}", hex::encode(&deposit_data)).as_bytes())
            .map_err(|e| format!("Unable to write eth1 data file: {}", e))?;

        self.deposit_data = Some(deposit_data);

        Ok(self)
    }

    pub async fn submit_eth1_deposit<T: Transport>(
        self,
        web3: Web3<T>,
        from: Address,
        deposit_contract: Address,
    ) -> Result<(Self, Hash256), String> {
        let (deposit_data, deposit_amount) = self.get_deposit_data()?;
        web3.eth()
            .send_transaction(TransactionRequest {
                from,
                to: Some(deposit_contract),
                gas: Some(DEPOSIT_GAS.into()),
                gas_price: None,
                value: Some(from_gwei(deposit_amount)),
                data: Some(deposit_data.into()),
                nonce: None,
                condition: None,
            })
            .compat()
            .await
            .map_err(|e| format!("Failed to send transaction: {:?}", e))
            .map(|tx| (self, tx))
    }

    pub fn build(self) -> Result<ValidatorDirectory, String> {
        let directory = self.directory.ok_or_else(|| "build requires a directory")?;

        Ok(ValidatorDirectory {
            directory,
            voting_keypair: self.voting_keypair,
            withdrawal_keypair: self.withdrawal_keypair,
            deposit_data: self.deposit_data,
        })
    }
}

/// Converts gwei to wei.
fn from_gwei(gwei: u64) -> U256 {
    U256::from(gwei) * U256::exp10(9)
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempdir::TempDir;
    use types::{EthSpec, MinimalEthSpec};

    type E = MinimalEthSpec;

    #[test]
    fn random_keypairs_round_trip() {
        let spec = E::default_spec();
        let temp_dir = TempDir::new("acc_manager").expect("should create test dir");

        let created_dir = ValidatorDirectoryBuilder::default()
            .spec(spec)
            .full_deposit_amount()
            .expect("should set full deposit amount")
            .thread_random_keypairs()
            .create_directory(temp_dir.path().into())
            .expect("should create directory")
            .write_keypair_files()
            .expect("should write keypair files")
            .write_eth1_data_file()
            .expect("should write eth1 data file")
            .build()
            .expect("should build dir");

        let loaded_dir = ValidatorDirectory::load_for_signing(created_dir.directory.clone())
            .expect("should load directory");

        assert_eq!(
            created_dir, loaded_dir,
            "the directory created should match the one loaded"
        );
    }

    #[test]
    fn deterministic_keypairs_round_trip() {
        let spec = E::default_spec();
        let temp_dir = TempDir::new("acc_manager").expect("should create test dir");
        let index = 42;

        let created_dir = ValidatorDirectoryBuilder::default()
            .spec(spec)
            .full_deposit_amount()
            .expect("should set full deposit amount")
            .insecure_keypairs(index)
            .create_directory(temp_dir.path().into())
            .expect("should create directory")
            .write_keypair_files()
            .expect("should write keypair files")
            .write_eth1_data_file()
            .expect("should write eth1 data file")
            .build()
            .expect("should build dir");

        assert!(
            created_dir.directory.exists(),
            "should have created directory"
        );

        let mut parent = created_dir.directory.clone();
        parent.pop();
        assert_eq!(
            parent,
            PathBuf::from(temp_dir.path()),
            "should have created directory ontop of base dir"
        );

        let expected_keypair = generate_deterministic_keypair(index);
        assert_eq!(
            created_dir.voting_keypair,
            Some(expected_keypair.clone()),
            "voting keypair should be as expected"
        );
        assert_eq!(
            created_dir.withdrawal_keypair,
            Some(expected_keypair),
            "withdrawal keypair should be as expected"
        );
        assert!(
            !created_dir
                .deposit_data
                .clone()
                .expect("should have data")
                .is_empty(),
            "should have some deposit data"
        );

        let loaded_dir = ValidatorDirectory::load_for_signing(created_dir.directory.clone())
            .expect("should load directory");

        assert_eq!(
            created_dir, loaded_dir,
            "the directory created should match the one loaded"
        );
    }
}
<commit_msg>Removing validator_client/src/validator_directory.rs as logic has been moved to account_manager/src/validator/deposit.rs (#1266)<commit_after><|endoftext|>