<commit_before>//! CSV Virtual Table
//! Port of [csv](http://www.sqlite.org/cgi/src/finfo?name=ext/misc/csv.c) C extension.
extern crate csv;
use std::fs::File;
use std::os::raw::{c_char, c_int, c_void};
use std::path::Path;
use std::result;
use std::str;

use ffi;
use types::Null;
use vtab::{
    declare_vtab, dequote, escape_double_quote, parse_boolean, Context, IndexInfo, VTab,
    VTabCursor, Values,
};
use {Connection, Error, Result};

/// Register the "csv" module.
/// ```sql
/// CREATE VIRTUAL TABLE vtab USING csv(
///   filename=FILENAME -- Name of file containing CSV content
///   [, schema=SCHEMA] -- Alternative CSV schema. 'CREATE TABLE x(col1 TEXT NOT NULL, col2 INT, ...);'
///   [, header=YES|NO] -- First row of CSV defines the names of columns if "yes". Default "no".
///   [, columns=N] -- Assume the CSV file contains N columns.
///   [, delimiter=C] -- CSV delimiter. Default ','.
///   [, quote=C] -- CSV quote. Default '"'. 0 means no quote.
/// );
/// ```
pub fn load_module(conn: &Connection) -> Result<()> {
    let aux: Option<()> = None;
    conn.create_module("csv", &CSV_MODULE, aux)
}

init_module!(
    CSV_MODULE,
    CSVTab,
    (),
    CSVTabCursor,
    csv_create,
    csv_connect,
    csv_best_index,
    csv_disconnect,
    csv_disconnect,
    csv_open,
    csv_close,
    csv_filter,
    csv_next,
    csv_eof,
    csv_column,
    csv_rowid
);

/// An instance of the CSV virtual table
#[repr(C)]
struct CSVTab {
    /// Base class. Must be first
    base: ffi::sqlite3_vtab,
    /// Name of the CSV file
    filename: String,
    has_headers: bool,
    delimiter: u8,
    quote: u8,
    /// Offset to start of data
    offset_first_row: csv::Position,
}

impl CSVTab {
    fn reader(&self) -> result::Result<csv::Reader<File>, csv::Error> {
        csv::ReaderBuilder::new()
            .has_headers(self.has_headers)
            .delimiter(self.delimiter)
            .quote(self.quote)
            .from_path(&self.filename)
    }

    fn parameter(c_slice: &[u8]) -> Result<(&str, &str)> {
        let arg = try!(str::from_utf8(c_slice)).trim();
        let mut split = arg.split('=');
        if let Some(key) = split.next() {
            if let Some(value) = split.next() {
                let param = key.trim();
                let value = dequote(value);
                return Ok((param, value));
            }
        }
        Err(Error::ModuleError(format!("illegal argument: '{}'", arg)))
    }

    fn parse_byte(arg: &str) -> Option<u8> {
        if arg.len() == 1 {
            arg.bytes().next()
        } else {
            None
        }
    }
}

impl VTab for CSVTab {
    type Aux = ();
    type Cursor = CSVTabCursor;

    unsafe fn connect(db: *mut ffi::sqlite3, _aux: *mut (), args: &[&[u8]]) -> Result<CSVTab> {
        if args.len() < 4 {
            return Err(Error::ModuleError("no CSV file specified".to_owned()));
        }

        let mut vtab = CSVTab {
            base: Default::default(),
            filename: "".to_owned(),
            has_headers: false,
            delimiter: b',',
            quote: b'"',
            offset_first_row: csv::Position::new(),
        };
        let mut schema = None;
        let mut n_col = None;

        let args = &args[3..];
        for c_slice in args {
            let (param, value) = try!(CSVTab::parameter(c_slice));
            match param {
                "filename" => {
                    if !Path::new(value).exists() {
                        return Err(Error::ModuleError(format!(
                            "file '{}' does not exist",
                            value
                        )));
                    }
                    vtab.filename = value.to_owned();
                }
                "schema" => {
                    schema = Some(value.to_owned());
                }
                "columns" => {
                    if let Ok(n) = value.parse::<u16>() {
                        if n_col.is_some() {
                            return Err(Error::ModuleError(
                                "more than one 'columns' parameter".to_owned(),
                            ));
                        } else if n == 0 {
                            return Err(Error::ModuleError(
                                "must have at least one column".to_owned(),
                            ));
                        }
                        n_col = Some(n);
                    } else {
                        return Err(Error::ModuleError(format!(
                            "unrecognized argument to 'columns': {}",
                            value
                        )));
                    }
                }
                "header" => {
                    if let Some(b) = parse_boolean(value) {
                        vtab.has_headers = b;
                    } else {
                        return Err(Error::ModuleError(format!(
                            "unrecognized argument to 'header': {}",
                            value
                        )));
                    }
                }
                "delimiter" => {
                    if let Some(b) = CSVTab::parse_byte(value) {
                        vtab.delimiter = b;
                    } else {
                        return Err(Error::ModuleError(format!(
                            "unrecognized argument to 'delimiter': {}",
                            value
                        )));
                    }
                }
                "quote" => {
                    if let Some(b) = CSVTab::parse_byte(value) {
                        if b == b'0' {
                            vtab.quote = 0;
                        } else {
                            vtab.quote = b;
                        }
                    } else {
                        return Err(Error::ModuleError(format!(
                            "unrecognized argument to 'quote': {}",
                            value
                        )));
                    }
                }
                _ => {
                    return Err(Error::ModuleError(format!(
                        "unrecognized parameter '{}'",
                        param
                    )));
                }
            }
        }

        if vtab.filename.is_empty() {
            return Err(Error::ModuleError("no CSV file specified".to_owned()));
        }

        let mut cols: Vec<String> = Vec::new();
        if vtab.has_headers || (n_col.is_none() && schema.is_none()) {
            let mut reader = try!(vtab.reader());
            if vtab.has_headers {
                {
                    let headers = try!(reader.headers());
                    // headers ignored if cols is not empty
                    if n_col.is_none() && schema.is_none() {
                        cols = headers
                            .into_iter()
                            .map(|header| escape_double_quote(&header).into_owned())
                            .collect();
                    }
                }
                vtab.offset_first_row = reader.position().clone();
            } else {
                let mut record = csv::ByteRecord::new();
                if try!(reader.read_byte_record(&mut record)) {
                    for (i, _) in record.iter().enumerate() {
                        cols.push(format!("c{}", i));
                    }
                }
            }
        }

        if cols.is_empty() && schema.is_none() {
            return Err(Error::ModuleError("no column specified".to_owned()));
        }

        if schema.is_none() {
            let mut sql = String::from("CREATE TABLE x(");
            for (i, col) in cols.iter().enumerate() {
                sql.push('"');
                sql.push_str(col);
                sql.push_str("\" TEXT");
                if i == cols.len() - 1 {
                    sql.push_str(");");
                } else {
                    sql.push_str(", ");
                }
            }
            schema = Some(sql);
        }

        try!(declare_vtab(db, &schema.unwrap()));
        Ok(vtab)
    }

    // Only a forward full table scan is supported.
    fn best_index(&self, info: &mut IndexInfo) -> Result<()> {
        info.set_estimated_cost(1_000_000.);
        Ok(())
    }

    fn open(&self) -> Result<CSVTabCursor> {
        Ok(CSVTabCursor::new(try!(self.reader())))
    }
}

/// A cursor for the CSV virtual table
#[repr(C)]
struct CSVTabCursor {
    /// Base class. Must be first
    base: ffi::sqlite3_vtab_cursor,
    /// The CSV reader object
    reader: csv::Reader<File>,
    /// Current cursor position used as rowid
    row_number: usize,
    /// Values of the current row
    cols: csv::StringRecord,
    eof: bool,
}

impl CSVTabCursor {
    fn new(reader: csv::Reader<File>) -> CSVTabCursor {
        CSVTabCursor {
            base: Default::default(),
            reader,
            row_number: 0,
            cols: csv::StringRecord::new(),
            eof: false,
        }
    }
}

impl VTabCursor for CSVTabCursor {
    type Table = CSVTab;

    fn vtab(&self) -> &CSVTab {
        unsafe { &*(self.base.pVtab as *const CSVTab) }
    }

    // Only a full table scan is supported.  So `filter` simply rewinds to
    // the beginning.
    fn filter(&mut self, _idx_num: c_int, _idx_str: Option<&str>, _args: &Values) -> Result<()> {
        {
            let offset_first_row = self.vtab().offset_first_row.clone();
            try!(self.reader.seek(offset_first_row));
        }
        self.row_number = 0;
        self.next()
    }
    fn next(&mut self) -> Result<()> {
        {
            self.eof = self.reader.is_done();
            if self.eof {
                return Ok(());
            }

            self.eof = !try!(self.reader.read_record(&mut self.cols));
        }

        self.row_number += 1;
        Ok(())
    }
    fn eof(&self) -> bool {
        self.eof
    }
    fn column(&self, ctx: &mut Context, col: c_int) -> Result<()> {
        if col < 0 || col as usize >= self.cols.len() {
            return Err(Error::ModuleError(format!(
                "column index out of bounds: {}",
                col
            )));
        }
        if self.cols.is_empty() {
            ctx.set_result(&Null);
            return Ok(());
        }
        // TODO Affinity
        ctx.set_result(&self.cols[col as usize].to_owned());
        Ok(())
    }
    fn rowid(&self) -> Result<i64> {
        Ok(self.row_number as i64)
    }
}

impl From<csv::Error> for Error {
    fn from(err: csv::Error) -> Error {
        use std::error::Error as StdError;
        Error::ModuleError(String::from(err.description()))
    }
}

#[cfg(test)]
mod test {
    use vtab::csvtab;
    use {Connection, Result};

    #[test]
    fn test_csv_module() {
        let db = Connection::open_in_memory().unwrap();
        csvtab::load_module(&db).unwrap();
        db.execute_batch("CREATE VIRTUAL TABLE vtab USING csv(filename='test.csv', header=yes)")
            .unwrap();

        {
            let mut s = db.prepare("SELECT rowid, * FROM vtab").unwrap();
            {
                let headers = s.column_names();
                assert_eq!(vec!["rowid", "colA", "colB", "colC"], headers);
            }

            let ids: Result<Vec<i32>> = s.query_map(&[], |row| row.get::<i32, i32>(0))
                .unwrap()
                .collect();
            let sum = ids.unwrap().iter().fold(0, |acc, &id| acc + id);
            assert_eq!(sum, 15);
        }
        db.execute_batch("DROP TABLE vtab").unwrap();
    }

    #[test]
    fn test_csv_cursor() {
        let db = Connection::open_in_memory().unwrap();
        csvtab::load_module(&db).unwrap();
        db.execute_batch("CREATE VIRTUAL TABLE vtab USING csv(filename='test.csv', header=yes)")
            .unwrap();

        {
            let mut s = db.prepare(
                "SELECT v1.rowid, v1.* FROM vtab v1 NATURAL JOIN vtab v2 WHERE \
                 v1.rowid < v2.rowid",
            ).unwrap();

            let mut rows = s.query(&[]).unwrap();
            let row = rows.next().unwrap().unwrap();
            assert_eq!(row.get::<i32, i32>(0), 2);
        }
        db.execute_batch("DROP TABLE vtab").unwrap();
    }
}
<commit_msg>csvtab: use `columns` when specified<commit_after>//! CSV Virtual Table
//! Port of [csv](http://www.sqlite.org/cgi/src/finfo?name=ext/misc/csv.c) C extension.
extern crate csv;
use std::fs::File;
use std::os::raw::{c_char, c_int, c_void};
use std::path::Path;
use std::result;
use std::str;

use ffi;
use types::Null;
use vtab::{
    declare_vtab, dequote, escape_double_quote, parse_boolean, Context, IndexInfo, VTab,
    VTabCursor, Values,
};
use {Connection, Error, Result};

/// Register the "csv" module.
/// ```sql
/// CREATE VIRTUAL TABLE vtab USING csv(
///   filename=FILENAME -- Name of file containing CSV content
///   [, schema=SCHEMA] -- Alternative CSV schema. 'CREATE TABLE x(col1 TEXT NOT NULL, col2 INT, ...);'
///   [, header=YES|NO] -- First row of CSV defines the names of columns if "yes". Default "no".
///   [, columns=N] -- Assume the CSV file contains N columns.
///   [, delimiter=C] -- CSV delimiter. Default ','.
///   [, quote=C] -- CSV quote. Default '"'. 0 means no quote.
/// );
/// ```
pub fn load_module(conn: &Connection) -> Result<()> {
    let aux: Option<()> = None;
    conn.create_module("csv", &CSV_MODULE, aux)
}

init_module!(
    CSV_MODULE,
    CSVTab,
    (),
    CSVTabCursor,
    csv_create,
    csv_connect,
    csv_best_index,
    csv_disconnect,
    csv_disconnect,
    csv_open,
    csv_close,
    csv_filter,
    csv_next,
    csv_eof,
    csv_column,
    csv_rowid
);

/// An instance of the CSV virtual table
#[repr(C)]
struct CSVTab {
    /// Base class. Must be first
    base: ffi::sqlite3_vtab,
    /// Name of the CSV file
    filename: String,
    has_headers: bool,
    delimiter: u8,
    quote: u8,
    /// Offset to start of data
    offset_first_row: csv::Position,
}

impl CSVTab {
    fn reader(&self) -> result::Result<csv::Reader<File>, csv::Error> {
        csv::ReaderBuilder::new()
            .has_headers(self.has_headers)
            .delimiter(self.delimiter)
            .quote(self.quote)
            .from_path(&self.filename)
    }

    fn parameter(c_slice: &[u8]) -> Result<(&str, &str)> {
        let arg = try!(str::from_utf8(c_slice)).trim();
        let mut split = arg.split('=');
        if let Some(key) = split.next() {
            if let Some(value) = split.next() {
                let param = key.trim();
                let value = dequote(value);
                return Ok((param, value));
            }
        }
        Err(Error::ModuleError(format!("illegal argument: '{}'", arg)))
    }

    fn parse_byte(arg: &str) -> Option<u8> {
        if arg.len() == 1 {
            arg.bytes().next()
        } else {
            None
        }
    }
}

impl VTab for CSVTab {
    type Aux = ();
    type Cursor = CSVTabCursor;

    unsafe fn connect(db: *mut ffi::sqlite3, _aux: *mut (), args: &[&[u8]]) -> Result<CSVTab> {
        if args.len() < 4 {
            return Err(Error::ModuleError("no CSV file specified".to_owned()));
        }

        let mut vtab = CSVTab {
            base: Default::default(),
            filename: "".to_owned(),
            has_headers: false,
            delimiter: b',',
            quote: b'"',
            offset_first_row: csv::Position::new(),
        };
        let mut schema = None;
        let mut n_col = None;

        let args = &args[3..];
        for c_slice in args {
            let (param, value) = try!(CSVTab::parameter(c_slice));
            match param {
                "filename" => {
                    if !Path::new(value).exists() {
                        return Err(Error::ModuleError(format!(
                            "file '{}' does not exist",
                            value
                        )));
                    }
                    vtab.filename = value.to_owned();
                }
                "schema" => {
                    schema = Some(value.to_owned());
                }
                "columns" => {
                    if let Ok(n) = value.parse::<u16>() {
                        if n_col.is_some() {
                            return Err(Error::ModuleError(
                                "more than one 'columns' parameter".to_owned(),
                            ));
                        } else if n == 0 {
                            return Err(Error::ModuleError(
                                "must have at least one column".to_owned(),
                            ));
                        }
                        n_col = Some(n);
                    } else {
                        return Err(Error::ModuleError(format!(
                            "unrecognized argument to 'columns': {}",
                            value
                        )));
                    }
                }
                "header" => {
                    if let Some(b) = parse_boolean(value) {
                        vtab.has_headers = b;
                    } else {
                        return Err(Error::ModuleError(format!(
                            "unrecognized argument to 'header': {}",
                            value
                        )));
                    }
                }
                "delimiter" => {
                    if let Some(b) = CSVTab::parse_byte(value) {
                        vtab.delimiter = b;
                    } else {
                        return Err(Error::ModuleError(format!(
                            "unrecognized argument to 'delimiter': {}",
                            value
                        )));
                    }
                }
                "quote" => {
                    if let Some(b) = CSVTab::parse_byte(value) {
                        if b == b'0' {
                            vtab.quote = 0;
                        } else {
                            vtab.quote = b;
                        }
                    } else {
                        return Err(Error::ModuleError(format!(
                            "unrecognized argument to 'quote': {}",
                            value
                        )));
                    }
                }
                _ => {
                    return Err(Error::ModuleError(format!(
                        "unrecognized parameter '{}'",
                        param
                    )));
                }
            }
        }

        if vtab.filename.is_empty() {
            return Err(Error::ModuleError("no CSV file specified".to_owned()));
        }

        let mut cols: Vec<String> = Vec::new();
        if vtab.has_headers || (n_col.is_none() && schema.is_none()) {
            let mut reader = try!(vtab.reader());
            if vtab.has_headers {
                {
                    let headers = try!(reader.headers());
                    // headers ignored if cols is not empty
                    if n_col.is_none() && schema.is_none() {
                        cols = headers
                            .into_iter()
                            .map(|header| escape_double_quote(&header).into_owned())
                            .collect();
                    }
                }
                vtab.offset_first_row = reader.position().clone();
            } else {
                let mut record = csv::ByteRecord::new();
                if try!(reader.read_byte_record(&mut record)) {
                    for (i, _) in record.iter().enumerate() {
                        cols.push(format!("c{}", i));
                    }
                }
            }
        } else if let Some(n_col) = n_col {
            for i in 0..n_col {
                cols.push(format!("c{}", i));
            }
        }

        if cols.is_empty() && schema.is_none() {
            return Err(Error::ModuleError("no column specified".to_owned()));
        }

        if schema.is_none() {
            let mut sql = String::from("CREATE TABLE x(");
            for (i, col) in cols.iter().enumerate() {
                sql.push('"');
                sql.push_str(col);
                sql.push_str("\" TEXT");
                if i == cols.len() - 1 {
                    sql.push_str(");");
                } else {
                    sql.push_str(", ");
                }
            }
            schema = Some(sql);
        }

        try!(declare_vtab(db, &schema.unwrap()));
        Ok(vtab)
    }

    // Only a forward full table scan is supported.
    fn best_index(&self, info: &mut IndexInfo) -> Result<()> {
        info.set_estimated_cost(1_000_000.);
        Ok(())
    }

    fn open(&self) -> Result<CSVTabCursor> {
        Ok(CSVTabCursor::new(try!(self.reader())))
    }
}

/// A cursor for the CSV virtual table
#[repr(C)]
struct CSVTabCursor {
    /// Base class. Must be first
    base: ffi::sqlite3_vtab_cursor,
    /// The CSV reader object
    reader: csv::Reader<File>,
    /// Current cursor position used as rowid
    row_number: usize,
    /// Values of the current row
    cols: csv::StringRecord,
    eof: bool,
}

impl CSVTabCursor {
    fn new(reader: csv::Reader<File>) -> CSVTabCursor {
        CSVTabCursor {
            base: Default::default(),
            reader,
            row_number: 0,
            cols: csv::StringRecord::new(),
            eof: false,
        }
    }
}

impl VTabCursor for CSVTabCursor {
    type Table = CSVTab;

    fn vtab(&self) -> &CSVTab {
        unsafe { &*(self.base.pVtab as *const CSVTab) }
    }

    // Only a full table scan is supported.  So `filter` simply rewinds to
    // the beginning.
    fn filter(&mut self, _idx_num: c_int, _idx_str: Option<&str>, _args: &Values) -> Result<()> {
        {
            let offset_first_row = self.vtab().offset_first_row.clone();
            try!(self.reader.seek(offset_first_row));
        }
        self.row_number = 0;
        self.next()
    }
    fn next(&mut self) -> Result<()> {
        {
            self.eof = self.reader.is_done();
            if self.eof {
                return Ok(());
            }

            self.eof = !try!(self.reader.read_record(&mut self.cols));
        }

        self.row_number += 1;
        Ok(())
    }
    fn eof(&self) -> bool {
        self.eof
    }
    fn column(&self, ctx: &mut Context, col: c_int) -> Result<()> {
        if col < 0 || col as usize >= self.cols.len() {
            return Err(Error::ModuleError(format!(
                "column index out of bounds: {}",
                col
            )));
        }
        if self.cols.is_empty() {
            ctx.set_result(&Null);
            return Ok(());
        }
        // TODO Affinity
        ctx.set_result(&self.cols[col as usize].to_owned());
        Ok(())
    }
    fn rowid(&self) -> Result<i64> {
        Ok(self.row_number as i64)
    }
}

impl From<csv::Error> for Error {
    fn from(err: csv::Error) -> Error {
        use std::error::Error as StdError;
        Error::ModuleError(String::from(err.description()))
    }
}

#[cfg(test)]
mod test {
    use vtab::csvtab;
    use {Connection, Result};

    #[test]
    fn test_csv_module() {
        let db = Connection::open_in_memory().unwrap();
        csvtab::load_module(&db).unwrap();
        db.execute_batch("CREATE VIRTUAL TABLE vtab USING csv(filename='test.csv', header=yes)")
            .unwrap();

        {
            let mut s = db.prepare("SELECT rowid, * FROM vtab").unwrap();
            {
                let headers = s.column_names();
                assert_eq!(vec!["rowid", "colA", "colB", "colC"], headers);
            }

            let ids: Result<Vec<i32>> = s.query_map(&[], |row| row.get::<i32, i32>(0))
                .unwrap()
                .collect();
            let sum = ids.unwrap().iter().fold(0, |acc, &id| acc + id);
            assert_eq!(sum, 15);
        }
        db.execute_batch("DROP TABLE vtab").unwrap();
    }

    #[test]
    fn test_csv_cursor() {
        let db = Connection::open_in_memory().unwrap();
        csvtab::load_module(&db).unwrap();
        db.execute_batch("CREATE VIRTUAL TABLE vtab USING csv(filename='test.csv', header=yes)")
            .unwrap();

        {
            let mut s = db.prepare(
                "SELECT v1.rowid, v1.* FROM vtab v1 NATURAL JOIN vtab v2 WHERE \
                 v1.rowid < v2.rowid",
            ).unwrap();

            let mut rows = s.query(&[]).unwrap();
            let row = rows.next().unwrap().unwrap();
            assert_eq!(row.get::<i32, i32>(0), 2);
        }
        db.execute_batch("DROP TABLE vtab").unwrap();
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::prelude::*;
use std::io::BufReader;
use std::io::SeekFrom;
use std::fs::File;
use std::cmp::Ordering;
use std::mem;
use std::thread;
use utils;
use errors::Word2VecError;

pub struct WordVector {
    vocabulary: Vec<(String, Vec<f32>)>,
    vector_size: usize,
    clusters: Option<Vec<String>>
}

impl WordVector{
	pub fn load_from_binary(file_name: &str) -> Result<WordVector, Word2VecError>
	{
		let mut file = try!(File::open(file_name));
		let mut reader = BufReader::new(file);
		let mut header = String::new();
		try!(reader.read_line(&mut header));
		let mut header_info = header.split_whitespace().filter_map(|x| x.parse::<usize>().ok()).take(2).collect::<Vec<usize>>();
		if header_info.len() != 2{
			return Err(Word2VecError::WrongHeader)
		}
		let vocabulary_size = header_info[0];
		let vector_size = header_info[1];
		let mut vocabulary: Vec<(String, Vec<f32>)> = Vec::with_capacity(vocabulary_size);
		for i in 0..vocabulary_size{
			let mut word_bytes: Vec<u8> = Vec::new();
			try!(reader.read_until(b' ', &mut word_bytes));
			word_bytes.pop();
			let word = try!(String::from_utf8(word_bytes));
			let mut current_vector: Vec<f32> = Vec::with_capacity(vector_size);
			for j in 0..vector_size{
				let mut buf: [u8; 4] = [0; 4];
				reader.read(&mut buf);
				let vec = unsafe{ mem::transmute::<[u8; 4], f32>(buf) };
				current_vector.push(vec);
			}
			current_vector = utils::vector_norm(current_vector);
			vocabulary.push((word, current_vector));
			reader.seek(SeekFrom::Current(1));
		}
		Ok(WordVector{
			vocabulary: vocabulary,
			vector_size: vector_size,
			clusters: None
		})
	}

	fn get_index(&self, word: &str) -> Option<usize>
	{
		return self.vocabulary.iter().position(|x| x.0 == word);
	}

	pub fn get_vector(&self, word: &str) -> Option<&Vec<f32>>
	{
		let index = self.get_index(word);
		match index {
		    Some(val) => {
		    	return Some(&self.vocabulary[val].1)
		    },
		    None => None,
		}
	}

	pub fn cosine(&self, word: &str) -> Option<Vec<(String, f32)>>
	{
		let word_vector = self.get_vector(word);
		match word_vector {
		    Some(val) => {
		        let mut metrics: Vec<(String, f32)> = Vec::new();
		        for word in self.vocabulary.iter(){
		    	    metrics.push((word.0.clone(), utils::dot_product(&word.1, val)));
		        }
		        metrics.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));
		        metrics.remove(0);
		        metrics.truncate(20);
		        return Some(metrics)
	        },
		    None => None,
		}
	}

	pub fn analogy(&self, pos: Vec<&str>, neg: Vec<&str>) -> Option<Vec<(String, f32)>>
	{
		let mut vectors: Vec<Vec<f32>> = Vec::new();
		for word in pos.iter(){
			match self.get_vector(word){
				Some(val) => {
					vectors.push(val.clone())
				},
				None => {}
			}
		}
		for word in neg.iter(){
			match self.get_vector(word) {
			    Some(val) => vectors.push(val.iter().map(|x| -x).collect::<Vec<f32>>()),
			    None => {},
			}
		}
		let mut mean: Vec<f32> = Vec::with_capacity(self.vector_size);
		for i in 0..self.vector_size{
			let mut new_vec = Vec::new();
			for vector in vectors.iter(){
				new_vec.push(vector[i]);
			}
			mean.push(utils::mean(new_vec));
		}
		let mut metrics: Vec<(String, f32)> = Vec::new();
		for word in self.vocabulary.iter(){
		    metrics.push((word.0.clone(), utils::dot_product(&word.1, &mean)));
		}
		metrics.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));
		metrics.remove(0);
		metrics.truncate(20);
		return Some(metrics)
	}

}<commit_msg>Fix bug with pos/neg words in analogy method<commit_after>use std::io;
use std::io::prelude::*;
use std::io::BufReader;
use std::io::SeekFrom;
use std::fs::File;
use std::cmp::Ordering;
use std::mem;
use std::thread;
use std::sync::{Arc, Mutex};
use utils;
use errors::Word2VecError;

pub struct WordVector {
    vocabulary: Vec<(String, Vec<f32>)>,
    vector_size: usize,
    clusters: Option<Vec<String>>
}

impl WordVector{
	pub fn load_from_binary(file_name: &str) -> Result<WordVector, Word2VecError>
	{
		let file = try!(File::open(file_name));
		let mut reader = BufReader::new(file);
		let mut header = String::new();
		try!(reader.read_line(&mut header));
		let header_info = header.split_whitespace().filter_map(|x| x.parse::<usize>().ok()).take(2).collect::<Vec<usize>>();
		if header_info.len() != 2{
			return Err(Word2VecError::WrongHeader)
		}
		let vocabulary_size = header_info[0];
		let vector_size = header_info[1];
		let mut vocabulary: Vec<(String, Vec<f32>)> = Vec::with_capacity(vocabulary_size);
		for _ in 0..vocabulary_size{
			let mut word_bytes: Vec<u8> = Vec::new();
			try!(reader.read_until(b' ', &mut word_bytes));
			word_bytes.pop();
			let word = try!(String::from_utf8(word_bytes));
			let mut current_vector: Vec<f32> = Vec::with_capacity(vector_size);
			for _ in 0..vector_size{
				let mut buf: [u8; 4] = [0; 4];
				try!(reader.read(&mut buf));
				let vec = unsafe{ mem::transmute::<[u8; 4], f32>(buf) };
				current_vector.push(vec);
			}
			current_vector = utils::vector_norm(current_vector);
			vocabulary.push((word, current_vector));
			try!(reader.seek(SeekFrom::Current(1)));
		}
		Ok(WordVector{
			vocabulary: vocabulary,
			vector_size: vector_size,
			clusters: None
		})
	}

	fn get_index(&self, word: &str) -> Option<usize>
	{
		return self.vocabulary.iter().position(|x| x.0 == word);
	}

	pub fn get_vector(&self, word: &str) -> Option<&Vec<f32>>
	{
		let index = self.get_index(word);
		match index {
		    Some(val) => {
		    	return Some(&self.vocabulary[val].1)
		    },
		    None => None,
		}
	}

	pub fn cosine(&self, word: &str, n: usize) -> Option<Vec<(String, f32)>>
	{
		let word_vector = self.get_vector(word);
		match word_vector {
		    Some(val) => {
		        let mut metrics: Vec<(String, f32)> = Vec::with_capacity(self.vocabulary.len());
		        for word in self.vocabulary.iter(){
		    	    metrics.push((word.0.clone(), utils::dot_product(&word.1, val)));
		        }
		        metrics.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));
		        metrics.remove(0);
		        metrics.truncate(n);
		        return Some(metrics)
	        },
		    None => None,
		}
	}

	pub fn analogy(&self, pos: Vec<&str>, neg: Vec<&str>, n: usize) -> Option<Vec<(String, f32)>>
	{
		let mut vectors: Vec<Vec<f32>> = Vec::new();
		let mut exclude: Vec<String> = Vec::new();
		for word in pos.iter(){
			exclude.push(word.to_string());
			match self.get_vector(word){
				Some(val) => vectors.push(val.clone()),
				None => {}
			}
		}
		for word in neg.iter(){
			exclude.push(word.to_string());
			match self.get_vector(word) {
			    Some(val) => vectors.push(val.iter().map(|x| -x).collect::<Vec<f32>>()),
			    None => {},
			}
		}
		let mut mean: Vec<f32> = Vec::with_capacity(self.vector_size);
		for i in 0..self.vector_size{
			let mut new_vec = Vec::new();
			for vector in vectors.iter(){
				new_vec.push(vector[i]);
			}
			mean.push(utils::mean(new_vec));
		}
		let mut metrics: Vec<(String, f32)> = Vec::new();
		for word in self.vocabulary.iter(){
		    metrics.push((word.0.clone(), utils::dot_product(&word.1, &mean)));
		}
		metrics.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));
		metrics.retain(|x| exclude.contains(&x.0) == false);
		metrics.truncate(n);
		return Some(metrics)
	}

}<|endoftext|>


<commit_before>extern crate serenity;

use std::env;

use serenity::{
    model::{channel::Message, gateway::Ready},
    prelude::*
};

struct Handler;

impl EventHandler for Handler {
    fn message(&self, _: Context, msg: Message) {
        if msg.content == "!hello" {
            // The create message builder allows you to easily create embeds and messages
            // using a builder syntax.
            // This example will create a message that says "Hello, World!", with an embed that has
            // a title, description, three fields, and footer.
            if let Err(why) = msg.channel_id.send_message(|m| m
                .content("Hello, World!")
                .embed(|e| e
                    .title("This is a title")
                    .description("This is a description")
                    .fields(vec![
                        ("This is the first field", "This is a field body", true),
                        ("This is the second field", "Both of these fields are inline", true),
                    ])
                    .field("This is the third field", "This is not an inline field", false)
                    .footer(|f| f
                        .text("This is a footer")))) {
                println!("Error sending message: {:?}", why);
            }
        }
    }

    fn ready(&self, _: Context, ready: Ready) {
        println!("{} is connected!", ready.user.name);
    }
}

fn main() {
    // Configure the client with your Discord bot token in the environment.
    let token = env::var("DISCORD_TOKEN")
        .expect("Expected a token in the environment");
    let mut client = Client::new(&token, Handler).expect("Err creating client");

    if let Err(why) = client.start() {
        println!("Client error: {:?}", why);
    }
}
<commit_msg>Add colour to example 11 (#475)<commit_after>extern crate serenity;

use std::env;

use serenity::{
    model::{channel::Message, gateway::Ready},
    prelude::*
};

struct Handler;

impl EventHandler for Handler {
    fn message(&self, _: Context, msg: Message) {
        if msg.content == "!hello" {
            // The create message builder allows you to easily create embeds and messages
            // using a builder syntax.
            // This example will create a message that says "Hello, World!", with an embed that has
            // a title, description, three fields, footer and colour.
            if let Err(why) = msg.channel_id.send_message(|m| m
                .content("Hello, World!")
                .embed(|e| e
                    .title("This is a title")
                    .description("This is a description")
                    .fields(vec![
                        ("This is the first field", "This is a field body", true),
                        ("This is the second field", "Both of these fields are inline", true),
                    ])
                    .field("This is the third field", "This is not an inline field", false)
                    .footer(|f| f
                        .text("This is a footer"))
                    .colour((246, 111, 0)))) {
                println!("Error sending message: {:?}", why);
            }
        }
    }

    fn ready(&self, _: Context, ready: Ready) {
        println!("{} is connected!", ready.user.name);
    }
}

fn main() {
    // Configure the client with your Discord bot token in the environment.
    let token = env::var("DISCORD_TOKEN")
        .expect("Expected a token in the environment");
    let mut client = Client::new(&token, Handler).expect("Err creating client");

    if let Err(why) = client.start() {
        println!("Client error: {:?}", why);
    }
}
<|endoftext|>


<commit_before>#[crate_id(name="du", vers="1.0.0", author="Derek Chiang")];

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Derek Chiang <derekchiang93@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

#[feature(macro_rules)];

extern crate extra;
extern crate getopts;
extern crate sync;

use std::os;
use std::io::fs;
use std::io::FileStat;
use std::option::Option;
use std::path::Path;
use extra::time::Timespec;
use sync::{Arc, Future};

#[path = "../util.rs"]
mod util;

static NAME: &'static str = "du";
static VERSION: &'static str = "1.0.0";

struct Options {
    all: bool,
    max_depth: Option<uint>,
    total: bool,
    separate_dirs: bool,
}

fn du(path: &Path, options_arc: Arc<Options>, depth: uint) -> ~[Arc<FileStat>] {
    let mut stats = ~[];
    let mut futures = ~[];
    let options = options_arc.get();
    let mut my_stat = safe_unwrap!(path.stat());

    for f in safe_unwrap!(fs::readdir(path)).move_iter() {
        match f.is_file() {
            true => {
                let stat = safe_unwrap!(f.stat());
                my_stat.size += stat.size;
                my_stat.unstable.blocks += stat.unstable.blocks;
                if options.all {
                    stats.push(Arc::new(stat))
                }
            }
            false => {
                let oa_clone = options_arc.clone();
                futures.push(Future::spawn(proc() { du(&f, oa_clone, depth + 1) }))
            }
        }
    }

    for future in futures.mut_iter() {
        for stat_arc in future.get().move_rev_iter() {
            let stat = stat_arc.get();
            if !options.separate_dirs && stat.path.dir_path() == my_stat.path {
                my_stat.size += stat.size;
                my_stat.unstable.blocks += stat.unstable.blocks;
            }
            if options.max_depth == None || depth < options.max_depth.unwrap() {
                stats.push(stat_arc.clone());
            }
        }
    }

    stats.push(Arc::new(my_stat));

    return stats;
}

fn main() {
    let args = os::args();
    let program = args[0].clone();
    let opts = ~[
        // In task
        getopts::optflag("a", "all", " write counts for all files, not just directories"),
        // In main
        getopts::optflag("", "apparent-size", "print apparent sizes,  rather  than  disk  usage;
            although  the apparent  size is usually smaller, it may be larger due to holes
            in ('sparse') files, internal  fragmentation,  indirect  blocks, and the like"),
        // In main
        getopts::optopt("B", "block-size", "scale sizes  by  SIZE before printing them.
            E.g., '-BM' prints sizes in units of 1,048,576 bytes.  See SIZE format below.",
            "SIZE"),
        // In main
        getopts::optflag("b", "bytes", "equivalent to '--apparent-size --block-size=1'"),
        // In main
        getopts::optflag("c", "total", "produce a grand total"),
        // In task
        // getopts::optflag("D", "dereference-args", "dereference only symlinks that are listed
        //     on the command line"),
        // In main
        // getopts::optopt("", "files0-from", "summarize disk usage of the NUL-terminated file
        //                   names specified in file F;
        //                   If F is - then read names from standard input", "F"),
        // // In task
        // getopts::optflag("H", "", "equivalent to --dereference-args (-D)"),
        // In main
        getopts::optflag("h", "human-readable", "print sizes in human readable format (e.g., 1K 234M 2G)"),
        // In main
        getopts::optflag("", "si", "like -h, but use powers of 1000 not 1024"),
        // In main
        getopts::optflag("k", "", "like --block-size=1K"),
        // In task
        getopts::optflag("l", "count-links", "count sizes many times if hard linked"),
        // // In main
        getopts::optflag("m", "", "like --block-size=1M"),
        // // In task
        // getopts::optflag("L", "dereference", "dereference all symbolic links"),
        // // In task
        // getopts::optflag("P", "no-dereference", "don't follow any symbolic links (this is the default)"),
        // // In main
        getopts::optflag("0", "null", "end each output line with 0 byte rather than newline"),
        // In main
        getopts::optflag("S", "separate-dirs", "do not include size of subdirectories"),
        // In main
        getopts::optflag("s", "summarize", "display only a total for each argument"),
        // // In task
        // getopts::optflag("x", "one-file-system", "skip directories on different file systems"),
        // // In task
        // getopts::optopt("X", "exclude-from", "exclude files that match any pattern in FILE", "FILE"),
        // // In task
        // getopts::optopt("", "exclude", "exclude files that match PATTERN", "PATTERN"),
        // In main
        getopts::optopt("d", "max-depth", "print the total for a directory (or file, with --all)
            only if it is N or fewer levels below the command
            line argument;  --max-depth=0 is the same as --summarize", "N"),
        // In main
        getopts::optflagopt("", "time", "show time of the last modification of any file in the
            directory, or any of its subdirectories.  If WORD is given, show time as WORD instead of modification time:
            atime, access, use, ctime or status", "WORD"),
        // In main
        getopts::optopt("", "time-style", "show times using style STYLE:
            full-iso, long-iso, iso, +FORMAT FORMAT is interpreted like 'date'", "STYLE"),
        getopts::optflag("", "help", "display this help and exit"),
        getopts::optflag("", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(f) => {
            show_error!(1, "Invalid options\n{}", f.to_err_msg());
            return
        }
    };

    if matches.opt_present("help") {
        println!("du {} - estimate file space usage", VERSION);
        println!("");
        println!("Usage:");
        println!("  {0:s} [OPTION]... [FILE]...", program);
        println!("  {0:s} [OPTION]... --files0-from=F", program);
        println!("");
        println!("{}", getopts::usage("Summarize disk usage of each FILE, recursively for directories.", opts));
        println!("Display  values  are  in  units  of  the  first  available  SIZE from
--block-size,  and the DU_BLOCK_SIZE, BLOCK_SIZE and BLOCKSIZE environ‐
ment variables.  Otherwise, units default to  1024  bytes  (or  512  if
POSIXLY_CORRECT is set).

SIZE  is  an  integer and optional unit (example: 10M is 10*1024*1024).
Units are K, M, G, T, P, E, Z, Y (powers of 1024) or KB, MB, ...  (pow‐
ers of 1000).");
        return
    }

    let options = Options{
        all: matches.opt_present("all"),
        max_depth: match (matches.opt_present("summarize"), matches.opt_str("max-depth")) {
            (true, Some(s)) => match from_str::<uint>(s) {
                Some(_) => {
                    println!("du: warning: summarizing conflicts with --max-depth={:s}", s);
                    return
                },
                None => {
                    println!("du: invalid maximum depth '{:s}'", s);
                    return
                }
            },
            (true, None) => Some(0),
            (false, Some(s)) => match from_str::<uint>(s) {
                Some(u) => Some(u),
                None => {
                    println!("du: invalid maximum depth '{:s}'", s);
                    return
                }
            },
            (false, None) => None
        },
        total: matches.opt_present("total"),
        separate_dirs: matches.opt_present("S"),
    };

    let strs = matches.free.clone();
    let strs = match strs.is_empty() {
        true => ~[~"./"],
        false => strs
    };

    let options_arc = Arc::new(options);

    let MB = match matches.opt_present("si") {
        true => 1000 * 1000,
        false => 1024 * 1024,
    };
    let KB = match matches.opt_present("si") {
        true => 1000,
        false => 1024,
    };

    let block_size = match matches.opt_str("block-size") {
        Some(s) => {
            let mut found_number = false;
            let mut found_letter = false;
            let mut numbers = ~[];
            let mut letters = ~[];
            for c in s.chars() {
                if found_letter && c.is_digit() || !found_number && !c.is_digit() {
                    println!("du: invalid --block-size argument '{}'", s);
                    return
                } else if c.is_digit() {
                    found_number = true;
                    numbers.push(c as u8);
                } else if c.is_alphabetic() {
                    found_letter = true;
                    letters.push(c);
                }
            }
            let number = std::uint::parse_bytes(numbers, 10).unwrap();
            let multiple = match std::str::from_chars(letters).as_slice() {
                "K" => 1024, "M" => 1024 * 1024, "G" => 1024 * 1024 * 1024,
                "T" => 1024 * 1024 * 1024 * 1024, "P" => 1024 * 1024 * 1024 * 1024 * 1024,
                "E" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "Z" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "Y" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "KB" => 1000, "MB" => 1000 * 1000, "GB" => 1000 * 1000 * 1000,
                "TB" => 1000 * 1000 * 1000 * 1000, "PB" => 1000 * 1000 * 1000 * 1000 * 1000,
                "EB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                "ZB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                "YB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                _ => {
                    println!("du: invalid --block-size argument '{}'", s); return
                }
            };
            number * multiple
        },
        None => 1024
    };

    let convert_size = |size: u64| -> ~str {
        if matches.opt_present("human-readable") || matches.opt_present("si") {
            if size > MB {
                format!("{:.1f}M", (size as f64) / (MB as f64))
            } else if size > KB {
                format!("{:.1f}K", (size as f64) / (KB as f64))
            } else {
                format!("{}B", size)
            }
        } else if matches.opt_present("k") {
            format!("{}", ((size as f64) / (KB as f64)).ceil())
        } else if matches.opt_present("m") {
            format!("{}", ((size as f64) / (MB as f64)).ceil())
        } else {
            format!("{}", ((size as f64) / (block_size as f64)).ceil())
        }
    };

    let time_format_str = match matches.opt_str("time-style") {
        Some(s) => {
            match s.as_slice() {
                "full-iso" => "%Y-%m-%d %H:%M:%S.%f %z",
                "long-iso" => "%Y-%m-%d %H:%M",
                "iso" => "%Y-%m-%d",
                _ => {
                    println!("
du: invalid argument 'awdwa' for 'time style'
Valid arguments are:
- 'full-iso'
- 'long-iso'
- 'iso'
Try 'du --help' for more information.");
                    return
                }
            }
        },
        None => "%Y-%m-%d %H:%M"
    };

    let line_separator = match matches.opt_present("0") {
        true => "\0",
        false => "\n",
    };

    let mut grand_total = 0;
    for path_str in strs.iter() {
        let path = Path::new(path_str.clone());
        let iter = du(&path, options_arc.clone(), 0).move_iter();
        let (_, len) = iter.size_hint();
        let len = len.unwrap();
        for (index, stat_arc) in iter.enumerate() {
            let stat = stat_arc.get();
            let size = match matches.opt_present("apparent-size") {
                true => stat.unstable.nlink * stat.size,
                // C's stat is such that each block is assume to be 512 bytes
                // See: http://linux.die.net/man/2/stat
                false => stat.unstable.blocks * 512,
            };
            if matches.opt_present("time") {
                let time_str = {
                    let (secs, nsecs) = {
                        let time = match matches.opt_str("time") {
                            Some(s) => match s.as_slice() {
                                "accessed" => stat.accessed,
                                "created" => stat.created,
                                "modified" => stat.modified,
                                _ => {
                                    println!("du: invalid argument 'modified' for '--time'
    Valid arguments are:
      - 'accessed', 'created', 'modified'
    Try 'du --help' for more information.");
                                    return
                                }
                            },
                            None => stat.modified
                        };
                        ((time / 1000) as i64, (time % 1000 * 1000000) as i32)
                    };
                    let time_spec = Timespec::new(secs, nsecs);
                    extra::time::at(time_spec).strftime(time_format_str)
                };
                print!("{:<10} {:<30} {}", convert_size(size), time_str, stat.path.display());
            } else {
                print!("{:<10} {}", convert_size(size), stat.path.display());
            }
            print!("{}", line_separator);
            if options.total && index == (len - 1) {
                // The last element will be the total size of the the path under
                // path_str.  We add it to the grand total.
                grand_total += size;
            }
        }
    }

    if options.total {
        print!("{:<10} total", convert_size(grand_total));
        print!("{}", line_separator);
    }
}
<commit_msg>du --version was not implemented<commit_after>#[crate_id(name="du", vers="1.0.0", author="Derek Chiang")];

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Derek Chiang <derekchiang93@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

#[feature(macro_rules)];

extern crate extra;
extern crate getopts;
extern crate sync;

use std::os;
use std::io::fs;
use std::io::FileStat;
use std::option::Option;
use std::path::Path;
use extra::time::Timespec;
use sync::{Arc, Future};

#[path = "../util.rs"]
mod util;

static NAME: &'static str = "du";
static VERSION: &'static str = "1.0.0";

struct Options {
    all: bool,
    max_depth: Option<uint>,
    total: bool,
    separate_dirs: bool,
}

fn du(path: &Path, options_arc: Arc<Options>, depth: uint) -> ~[Arc<FileStat>] {
    let mut stats = ~[];
    let mut futures = ~[];
    let options = options_arc.get();
    let mut my_stat = safe_unwrap!(path.stat());

    for f in safe_unwrap!(fs::readdir(path)).move_iter() {
        match f.is_file() {
            true => {
                let stat = safe_unwrap!(f.stat());
                my_stat.size += stat.size;
                my_stat.unstable.blocks += stat.unstable.blocks;
                if options.all {
                    stats.push(Arc::new(stat))
                }
            }
            false => {
                let oa_clone = options_arc.clone();
                futures.push(Future::spawn(proc() { du(&f, oa_clone, depth + 1) }))
            }
        }
    }

    for future in futures.mut_iter() {
        for stat_arc in future.get().move_rev_iter() {
            let stat = stat_arc.get();
            if !options.separate_dirs && stat.path.dir_path() == my_stat.path {
                my_stat.size += stat.size;
                my_stat.unstable.blocks += stat.unstable.blocks;
            }
            if options.max_depth == None || depth < options.max_depth.unwrap() {
                stats.push(stat_arc.clone());
            }
        }
    }

    stats.push(Arc::new(my_stat));

    return stats;
}

fn main() {
    let args = os::args();
    let program = args[0].clone();
    let opts = ~[
        // In task
        getopts::optflag("a", "all", " write counts for all files, not just directories"),
        // In main
        getopts::optflag("", "apparent-size", "print apparent sizes,  rather  than  disk  usage;
            although  the apparent  size is usually smaller, it may be larger due to holes
            in ('sparse') files, internal  fragmentation,  indirect  blocks, and the like"),
        // In main
        getopts::optopt("B", "block-size", "scale sizes  by  SIZE before printing them.
            E.g., '-BM' prints sizes in units of 1,048,576 bytes.  See SIZE format below.",
            "SIZE"),
        // In main
        getopts::optflag("b", "bytes", "equivalent to '--apparent-size --block-size=1'"),
        // In main
        getopts::optflag("c", "total", "produce a grand total"),
        // In task
        // getopts::optflag("D", "dereference-args", "dereference only symlinks that are listed
        //     on the command line"),
        // In main
        // getopts::optopt("", "files0-from", "summarize disk usage of the NUL-terminated file
        //                   names specified in file F;
        //                   If F is - then read names from standard input", "F"),
        // // In task
        // getopts::optflag("H", "", "equivalent to --dereference-args (-D)"),
        // In main
        getopts::optflag("h", "human-readable", "print sizes in human readable format (e.g., 1K 234M 2G)"),
        // In main
        getopts::optflag("", "si", "like -h, but use powers of 1000 not 1024"),
        // In main
        getopts::optflag("k", "", "like --block-size=1K"),
        // In task
        getopts::optflag("l", "count-links", "count sizes many times if hard linked"),
        // // In main
        getopts::optflag("m", "", "like --block-size=1M"),
        // // In task
        // getopts::optflag("L", "dereference", "dereference all symbolic links"),
        // // In task
        // getopts::optflag("P", "no-dereference", "don't follow any symbolic links (this is the default)"),
        // // In main
        getopts::optflag("0", "null", "end each output line with 0 byte rather than newline"),
        // In main
        getopts::optflag("S", "separate-dirs", "do not include size of subdirectories"),
        // In main
        getopts::optflag("s", "summarize", "display only a total for each argument"),
        // // In task
        // getopts::optflag("x", "one-file-system", "skip directories on different file systems"),
        // // In task
        // getopts::optopt("X", "exclude-from", "exclude files that match any pattern in FILE", "FILE"),
        // // In task
        // getopts::optopt("", "exclude", "exclude files that match PATTERN", "PATTERN"),
        // In main
        getopts::optopt("d", "max-depth", "print the total for a directory (or file, with --all)
            only if it is N or fewer levels below the command
            line argument;  --max-depth=0 is the same as --summarize", "N"),
        // In main
        getopts::optflagopt("", "time", "show time of the last modification of any file in the
            directory, or any of its subdirectories.  If WORD is given, show time as WORD instead of modification time:
            atime, access, use, ctime or status", "WORD"),
        // In main
        getopts::optopt("", "time-style", "show times using style STYLE:
            full-iso, long-iso, iso, +FORMAT FORMAT is interpreted like 'date'", "STYLE"),
        getopts::optflag("", "help", "display this help and exit"),
        getopts::optflag("", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(f) => {
            show_error!(1, "Invalid options\n{}", f.to_err_msg());
            return
        }
    };

    if matches.opt_present("help") {
        println!("du {} - estimate file space usage", VERSION);
        println!("");
        println!("Usage:");
        println!("  {0:s} [OPTION]... [FILE]...", program);
        println!("  {0:s} [OPTION]... --files0-from=F", program);
        println!("");
        println!("{}", getopts::usage("Summarize disk usage of each FILE, recursively for directories.", opts));
        println!("Display  values  are  in  units  of  the  first  available  SIZE from
--block-size,  and the DU_BLOCK_SIZE, BLOCK_SIZE and BLOCKSIZE environ‐
ment variables.  Otherwise, units default to  1024  bytes  (or  512  if
POSIXLY_CORRECT is set).

SIZE  is  an  integer and optional unit (example: 10M is 10*1024*1024).
Units are K, M, G, T, P, E, Z, Y (powers of 1024) or KB, MB, ...  (pow‐
ers of 1000).");
        return
    } else if matches.opt_present("version") {
        println!("du version: {}", VERSION);
        return
    }

    let options = Options{
        all: matches.opt_present("all"),
        max_depth: match (matches.opt_present("summarize"), matches.opt_str("max-depth")) {
            (true, Some(s)) => match from_str::<uint>(s) {
                Some(_) => {
                    println!("du: warning: summarizing conflicts with --max-depth={:s}", s);
                    return
                },
                None => {
                    println!("du: invalid maximum depth '{:s}'", s);
                    return
                }
            },
            (true, None) => Some(0),
            (false, Some(s)) => match from_str::<uint>(s) {
                Some(u) => Some(u),
                None => {
                    println!("du: invalid maximum depth '{:s}'", s);
                    return
                }
            },
            (false, None) => None
        },
        total: matches.opt_present("total"),
        separate_dirs: matches.opt_present("S"),
    };

    let strs = matches.free.clone();
    let strs = match strs.is_empty() {
        true => ~[~"./"],
        false => strs
    };

    let options_arc = Arc::new(options);

    let MB = match matches.opt_present("si") {
        true => 1000 * 1000,
        false => 1024 * 1024,
    };
    let KB = match matches.opt_present("si") {
        true => 1000,
        false => 1024,
    };

    let block_size = match matches.opt_str("block-size") {
        Some(s) => {
            let mut found_number = false;
            let mut found_letter = false;
            let mut numbers = ~[];
            let mut letters = ~[];
            for c in s.chars() {
                if found_letter && c.is_digit() || !found_number && !c.is_digit() {
                    println!("du: invalid --block-size argument '{}'", s);
                    return
                } else if c.is_digit() {
                    found_number = true;
                    numbers.push(c as u8);
                } else if c.is_alphabetic() {
                    found_letter = true;
                    letters.push(c);
                }
            }
            let number = std::uint::parse_bytes(numbers, 10).unwrap();
            let multiple = match std::str::from_chars(letters).as_slice() {
                "K" => 1024, "M" => 1024 * 1024, "G" => 1024 * 1024 * 1024,
                "T" => 1024 * 1024 * 1024 * 1024, "P" => 1024 * 1024 * 1024 * 1024 * 1024,
                "E" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "Z" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "Y" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "KB" => 1000, "MB" => 1000 * 1000, "GB" => 1000 * 1000 * 1000,
                "TB" => 1000 * 1000 * 1000 * 1000, "PB" => 1000 * 1000 * 1000 * 1000 * 1000,
                "EB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                "ZB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                "YB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                _ => {
                    println!("du: invalid --block-size argument '{}'", s); return
                }
            };
            number * multiple
        },
        None => 1024
    };

    let convert_size = |size: u64| -> ~str {
        if matches.opt_present("human-readable") || matches.opt_present("si") {
            if size > MB {
                format!("{:.1f}M", (size as f64) / (MB as f64))
            } else if size > KB {
                format!("{:.1f}K", (size as f64) / (KB as f64))
            } else {
                format!("{}B", size)
            }
        } else if matches.opt_present("k") {
            format!("{}", ((size as f64) / (KB as f64)).ceil())
        } else if matches.opt_present("m") {
            format!("{}", ((size as f64) / (MB as f64)).ceil())
        } else {
            format!("{}", ((size as f64) / (block_size as f64)).ceil())
        }
    };

    let time_format_str = match matches.opt_str("time-style") {
        Some(s) => {
            match s.as_slice() {
                "full-iso" => "%Y-%m-%d %H:%M:%S.%f %z",
                "long-iso" => "%Y-%m-%d %H:%M",
                "iso" => "%Y-%m-%d",
                _ => {
                    println!("
du: invalid argument 'awdwa' for 'time style'
Valid arguments are:
- 'full-iso'
- 'long-iso'
- 'iso'
Try 'du --help' for more information.");
                    return
                }
            }
        },
        None => "%Y-%m-%d %H:%M"
    };

    let line_separator = match matches.opt_present("0") {
        true => "\0",
        false => "\n",
    };

    let mut grand_total = 0;
    for path_str in strs.iter() {
        let path = Path::new(path_str.clone());
        let iter = du(&path, options_arc.clone(), 0).move_iter();
        let (_, len) = iter.size_hint();
        let len = len.unwrap();
        for (index, stat_arc) in iter.enumerate() {
            let stat = stat_arc.get();
            let size = match matches.opt_present("apparent-size") {
                true => stat.unstable.nlink * stat.size,
                // C's stat is such that each block is assume to be 512 bytes
                // See: http://linux.die.net/man/2/stat
                false => stat.unstable.blocks * 512,
            };
            if matches.opt_present("time") {
                let time_str = {
                    let (secs, nsecs) = {
                        let time = match matches.opt_str("time") {
                            Some(s) => match s.as_slice() {
                                "accessed" => stat.accessed,
                                "created" => stat.created,
                                "modified" => stat.modified,
                                _ => {
                                    println!("du: invalid argument 'modified' for '--time'
    Valid arguments are:
      - 'accessed', 'created', 'modified'
    Try 'du --help' for more information.");
                                    return
                                }
                            },
                            None => stat.modified
                        };
                        ((time / 1000) as i64, (time % 1000 * 1000000) as i32)
                    };
                    let time_spec = Timespec::new(secs, nsecs);
                    extra::time::at(time_spec).strftime(time_format_str)
                };
                print!("{:<10} {:<30} {}", convert_size(size), time_str, stat.path.display());
            } else {
                print!("{:<10} {}", convert_size(size), stat.path.display());
            }
            print!("{}", line_separator);
            if options.total && index == (len - 1) {
                // The last element will be the total size of the the path under
                // path_str.  We add it to the grand total.
                grand_total += size;
            }
        }
    }

    if options.total {
        print!("{:<10} total", convert_size(grand_total));
        print!("{}", line_separator);
    }
}
<|endoftext|>


<commit_before>use std::{
    fmt,
    iter::FromIterator,
    sync::Arc,
};

use ra_syntax::{AstNode, TreeArc, SourceFile};
use ra_db::{
    ParseQuery, FileTextQuery, SourceRootId,
    salsa::{Database, debug::{DebugQueryTable, TableEntry}},
};

use crate::{
    FileId, db::RootDatabase,
    symbol_index::{SymbolIndex, LibrarySymbolsQuery},
};

pub(crate) fn syntax_tree_stats(db: &RootDatabase) -> SyntaxTreeStats {
    db.query(ParseQuery).entries::<SyntaxTreeStats>()
}

pub(crate) fn status(db: &RootDatabase) -> String {
    let files_stats = db.query(FileTextQuery).entries::<FilesStats>();
    let syntax_tree_stats = syntax_tree_stats(db);
    let symbols_stats = db
        .query(LibrarySymbolsQuery)
        .entries::<LibrarySymbolsStats>();
    let n_defs = {
        let interner: &hir::HirInterner = db.as_ref();
        interner.len()
    };
    format!(
        "{}\n{}\n{}\nn_defs {}\n\njemalloc: {}\nGC {:?} seconds ago",
        files_stats,
        symbols_stats,
        syntax_tree_stats,
        n_defs,
        MemoryStats::current(),
        db.last_gc.elapsed().as_secs(),
    )
}

#[derive(Default)]
struct FilesStats {
    total: usize,
    size: Bytes,
}

impl fmt::Display for FilesStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{} ({}) files", self.total, self.size)
    }
}

impl FromIterator<TableEntry<FileId, Arc<String>>> for FilesStats {
    fn from_iter<T>(iter: T) -> FilesStats
    where
        T: IntoIterator<Item = TableEntry<FileId, Arc<String>>>,
    {
        let mut res = FilesStats::default();
        for entry in iter {
            res.total += 1;
            res.size += entry.value.unwrap().len();
        }
        res
    }
}

#[derive(Default)]
pub(crate) struct SyntaxTreeStats {
    total: usize,
    pub(crate) retained: usize,
    retained_size: Bytes,
}

impl fmt::Display for SyntaxTreeStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(
            fmt,
            "{} trees, {} ({}) retained",
            self.total, self.retained, self.retained_size,
        )
    }
}

impl FromIterator<TableEntry<FileId, TreeArc<SourceFile>>> for SyntaxTreeStats {
    fn from_iter<T>(iter: T) -> SyntaxTreeStats
    where
        T: IntoIterator<Item = TableEntry<FileId, TreeArc<SourceFile>>>,
    {
        let mut res = SyntaxTreeStats::default();
        for entry in iter {
            res.total += 1;
            if let Some(value) = entry.value {
                res.retained += 1;
                res.retained_size += value.syntax().memory_size_of_subtree();
            }
        }
        res
    }
}

#[derive(Default)]
struct LibrarySymbolsStats {
    total: usize,
    size: Bytes,
}

impl fmt::Display for LibrarySymbolsStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{} ({}) symbols", self.total, self.size,)
    }
}

impl FromIterator<TableEntry<SourceRootId, Arc<SymbolIndex>>> for LibrarySymbolsStats {
    fn from_iter<T>(iter: T) -> LibrarySymbolsStats
    where
        T: IntoIterator<Item = TableEntry<SourceRootId, Arc<SymbolIndex>>>,
    {
        let mut res = LibrarySymbolsStats::default();
        for entry in iter {
            let value = entry.value.unwrap();
            res.total += value.len();
            res.size += value.memory_size();
        }
        res
    }
}

struct MemoryStats {
    allocated: Bytes,
    resident: Bytes,
}

impl MemoryStats {
    fn current() -> MemoryStats {
        jemalloc_ctl::epoch().unwrap();
        MemoryStats {
            allocated: Bytes(jemalloc_ctl::stats::allocated().unwrap()),
            resident: Bytes(jemalloc_ctl::stats::resident().unwrap()),
        }
    }
}

impl fmt::Display for MemoryStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(
            fmt,
            "{} allocated {} resident",
            self.allocated, self.resident,
        )
    }
}

#[derive(Default)]
struct Bytes(usize);

impl fmt::Display for Bytes {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let bytes = self.0;
        if bytes < 4096 {
            return write!(f, "{} bytes", bytes);
        }
        let kb = bytes / 1024;
        if kb < 4096 {
            return write!(f, "{}kb", kb);
        }
        let mb = kb / 1024;
        write!(f, "{}mb", mb)
    }
}

impl std::ops::AddAssign<usize> for Bytes {
    fn add_assign(&mut self, x: usize) {
        self.0 += x;
    }
}
<commit_msg>tweak words<commit_after>use std::{
    fmt,
    iter::FromIterator,
    sync::Arc,
};

use ra_syntax::{AstNode, TreeArc, SourceFile};
use ra_db::{
    ParseQuery, FileTextQuery, SourceRootId,
    salsa::{Database, debug::{DebugQueryTable, TableEntry}},
};

use crate::{
    FileId, db::RootDatabase,
    symbol_index::{SymbolIndex, LibrarySymbolsQuery},
};

pub(crate) fn syntax_tree_stats(db: &RootDatabase) -> SyntaxTreeStats {
    db.query(ParseQuery).entries::<SyntaxTreeStats>()
}

pub(crate) fn status(db: &RootDatabase) -> String {
    let files_stats = db.query(FileTextQuery).entries::<FilesStats>();
    let syntax_tree_stats = syntax_tree_stats(db);
    let symbols_stats = db
        .query(LibrarySymbolsQuery)
        .entries::<LibrarySymbolsStats>();
    let n_defs = {
        let interner: &hir::HirInterner = db.as_ref();
        interner.len()
    };
    format!(
        "{}\n{}\n{}\n{} defs\n\nmemory:\n{}\ngc {:?} seconds ago",
        files_stats,
        symbols_stats,
        syntax_tree_stats,
        n_defs,
        MemoryStats::current(),
        db.last_gc.elapsed().as_secs(),
    )
}

#[derive(Default)]
struct FilesStats {
    total: usize,
    size: Bytes,
}

impl fmt::Display for FilesStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{} ({}) files", self.total, self.size)
    }
}

impl FromIterator<TableEntry<FileId, Arc<String>>> for FilesStats {
    fn from_iter<T>(iter: T) -> FilesStats
    where
        T: IntoIterator<Item = TableEntry<FileId, Arc<String>>>,
    {
        let mut res = FilesStats::default();
        for entry in iter {
            res.total += 1;
            res.size += entry.value.unwrap().len();
        }
        res
    }
}

#[derive(Default)]
pub(crate) struct SyntaxTreeStats {
    total: usize,
    pub(crate) retained: usize,
    retained_size: Bytes,
}

impl fmt::Display for SyntaxTreeStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(
            fmt,
            "{} trees, {} ({}) retained",
            self.total, self.retained, self.retained_size,
        )
    }
}

impl FromIterator<TableEntry<FileId, TreeArc<SourceFile>>> for SyntaxTreeStats {
    fn from_iter<T>(iter: T) -> SyntaxTreeStats
    where
        T: IntoIterator<Item = TableEntry<FileId, TreeArc<SourceFile>>>,
    {
        let mut res = SyntaxTreeStats::default();
        for entry in iter {
            res.total += 1;
            if let Some(value) = entry.value {
                res.retained += 1;
                res.retained_size += value.syntax().memory_size_of_subtree();
            }
        }
        res
    }
}

#[derive(Default)]
struct LibrarySymbolsStats {
    total: usize,
    size: Bytes,
}

impl fmt::Display for LibrarySymbolsStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{} ({}) symbols", self.total, self.size,)
    }
}

impl FromIterator<TableEntry<SourceRootId, Arc<SymbolIndex>>> for LibrarySymbolsStats {
    fn from_iter<T>(iter: T) -> LibrarySymbolsStats
    where
        T: IntoIterator<Item = TableEntry<SourceRootId, Arc<SymbolIndex>>>,
    {
        let mut res = LibrarySymbolsStats::default();
        for entry in iter {
            let value = entry.value.unwrap();
            res.total += value.len();
            res.size += value.memory_size();
        }
        res
    }
}

struct MemoryStats {
    allocated: Bytes,
    resident: Bytes,
}

impl MemoryStats {
    fn current() -> MemoryStats {
        jemalloc_ctl::epoch().unwrap();
        MemoryStats {
            allocated: Bytes(jemalloc_ctl::stats::allocated().unwrap()),
            resident: Bytes(jemalloc_ctl::stats::resident().unwrap()),
        }
    }
}

impl fmt::Display for MemoryStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(
            fmt,
            "{} allocated {} resident",
            self.allocated, self.resident,
        )
    }
}

#[derive(Default)]
struct Bytes(usize);

impl fmt::Display for Bytes {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let bytes = self.0;
        if bytes < 4096 {
            return write!(f, "{} bytes", bytes);
        }
        let kb = bytes / 1024;
        if kb < 4096 {
            return write!(f, "{}kb", kb);
        }
        let mb = kb / 1024;
        write!(f, "{}mb", mb)
    }
}

impl std::ops::AddAssign<usize> for Bytes {
    fn add_assign(&mut self, x: usize) {
        self.0 += x;
    }
}
<|endoftext|>


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

use prost::Message;
use tink_core::{utils::wrap_err, TinkError};

#[test]
fn test_templates() {
    tink_prf::init();
    let template = tink_prf::hmac_sha256_prf_key_template();
    assert!(
        check_hmac_template(&template, 32, tink_proto::HashType::Sha256).is_ok(),
        "incorrect HMAC_SHA256PRFKeyTemplate"
    );
    let template = tink_prf::hmac_sha512_prf_key_template();
    assert!(
        check_hmac_template(&template, 64, tink_proto::HashType::Sha512).is_ok(),
        "incorrect HMAC_SHA512PRFKeyTemplate"
    );
    let template = tink_prf::hkdf_sha256_prf_key_template();
    assert!(
        check_hkdf_template(&template, 32, &[], tink_proto::HashType::Sha256).is_ok(),
        "incorrect HKDFSHA256PRFKeyTemplate"
    );
    let template = tink_prf::aes_cmac_prf_key_template();
    assert!(
        check_cmac_template(&template, 32).is_ok(),
        "incorrect AESCMACSPRFKeyTemplate"
    );
}

fn check_hmac_template(
    template: &tink_proto::KeyTemplate,
    key_size: u32,
    hash_type: tink_proto::HashType,
) -> Result<(), TinkError> {
    if template.type_url != tink_tests::HMAC_PRF_TYPE_URL {
        return Err("TypeUrl is incorrect".into());
    }
    if template.output_prefix_type != tink_proto::OutputPrefixType::Raw as i32 {
        return Err("Not RAW output prefix".into());
    }
    let format = tink_proto::HmacPrfKeyFormat::decode(template.value.as_ref())
        .map_err(|_| TinkError::new("unable to unmarshal serialized key format"))?;
    if format.key_size != key_size || format.params.unwrap().hash != hash_type as i32 {
        return Err("KeyFormat is incorrect".into());
    }
    let keymanager = tink_core::registry::get_key_manager(tink_tests::HMAC_PRF_TYPE_URL)
        .map_err(|e| wrap_err("Could not obtain HMAC key manager", e))?;
    assert!(
        keymanager.new_key(&template.value).is_ok(),
        "HMAC key manager cannot create key"
    );
    Ok(())
}

fn check_hkdf_template(
    template: &tink_proto::KeyTemplate,
    key_size: u32,
    salt: &[u8],
    hash_type: tink_proto::HashType,
) -> Result<(), TinkError> {
    if template.type_url != tink_tests::HKDF_PRF_TYPE_URL {
        return Err("TypeUrl is incorrect".into());
    }
    if template.output_prefix_type != tink_proto::OutputPrefixType::Raw as i32 {
        return Err("Not RAW output prefix".into());
    }

    let format = tink_proto::HkdfPrfKeyFormat::decode(template.value.as_ref())
        .map_err(|_| TinkError::new("unable to unmarshal serialized key format"))?;
    if format.key_size != key_size
        || format.params.as_ref().unwrap().hash != hash_type as i32
        || hex::encode(salt) != hex::encode(format.params.unwrap().salt)
    {
        return Err("KeyFormat is incorrect".into());
    }
    let keymanager = tink_core::registry::get_key_manager(tink_tests::HKDF_PRF_TYPE_URL)
        .map_err(|e| wrap_err("Could not obtain HKDF key manager", e))?;
    assert!(
        keymanager.new_key(&template.value).is_ok(),
        "HKDF key manager cannot create key"
    );
    Ok(())
}

fn check_cmac_template(template: &tink_proto::KeyTemplate, key_size: u32) -> Result<(), TinkError> {
    if template.type_url != tink_tests::AES_CMAC_PRF_TYPE_URL {
        return Err("TypeUrl is incorrect".into());
    }
    if template.output_prefix_type != tink_proto::OutputPrefixType::Raw as i32 {
        return Err("Not RAW output prefix".into());
    }

    let format = tink_proto::AesCmacPrfKeyFormat::decode(template.value.as_ref())
        .map_err(|_| TinkError::new("unable to unmarshal serialized key format"))?;
    if format.key_size != key_size {
        return Err("KeyFormat is incorrect".into());
    }
    let keymanager = tink_core::registry::get_key_manager(tink_tests::AES_CMAC_PRF_TYPE_URL)
        .map_err(|e| wrap_err("Could not obtain AES-CMAC key manager", e))?;
    assert!(
        keymanager.new_key(&template.value).is_ok(),
        "AES-CMAC key manager cannot create key"
    );
    Ok(())
}
<commit_msg>tink-prf: add key template tests<commit_after>// Copyright 2020 The Tink-Rust Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
////////////////////////////////////////////////////////////////////////////////

use prost::Message;
use tink_core::{utils::wrap_err, TinkError};

#[test]
fn test_key_templates() {
    tink_prf::init();
    let test_cases = vec![
        ("HMAC_PRF_SHA256", tink_prf::hmac_sha256_prf_key_template()),
        ("HMAC_PRF_SHA512", tink_prf::hmac_sha512_prf_key_template()),
        ("HKDF_PRF_SHA256", tink_prf::hkdf_sha256_prf_key_template()),
        ("AES_CMAC_PRF", tink_prf::aes_cmac_prf_key_template()),
    ];
    for (name, template) in test_cases {
        let want = tink_tests::key_template_proto("prf", name).unwrap();
        assert_eq!(want, template);

        let handle = tink_core::keyset::Handle::new(&template).unwrap();
        let primitive = tink_prf::Set::new(&handle).unwrap();

        let msg = b"This is an IF that needs to be redacted";
        let output = primitive.compute_primary_prf(msg, 16).unwrap();
        assert_eq!(output.len(), 16);
    }
}

#[test]
fn test_templates() {
    tink_prf::init();
    let template = tink_prf::hmac_sha256_prf_key_template();
    assert!(
        check_hmac_template(&template, 32, tink_proto::HashType::Sha256).is_ok(),
        "incorrect HMAC_SHA256PRFKeyTemplate"
    );
    let template = tink_prf::hmac_sha512_prf_key_template();
    assert!(
        check_hmac_template(&template, 64, tink_proto::HashType::Sha512).is_ok(),
        "incorrect HMAC_SHA512PRFKeyTemplate"
    );
    let template = tink_prf::hkdf_sha256_prf_key_template();
    assert!(
        check_hkdf_template(&template, 32, &[], tink_proto::HashType::Sha256).is_ok(),
        "incorrect HKDFSHA256PRFKeyTemplate"
    );
    let template = tink_prf::aes_cmac_prf_key_template();
    assert!(
        check_cmac_template(&template, 32).is_ok(),
        "incorrect AESCMACSPRFKeyTemplate"
    );
}

fn check_hmac_template(
    template: &tink_proto::KeyTemplate,
    key_size: u32,
    hash_type: tink_proto::HashType,
) -> Result<(), TinkError> {
    if template.type_url != tink_tests::HMAC_PRF_TYPE_URL {
        return Err("TypeUrl is incorrect".into());
    }
    if template.output_prefix_type != tink_proto::OutputPrefixType::Raw as i32 {
        return Err("Not RAW output prefix".into());
    }
    let format = tink_proto::HmacPrfKeyFormat::decode(template.value.as_ref())
        .map_err(|_| TinkError::new("unable to unmarshal serialized key format"))?;
    if format.key_size != key_size || format.params.unwrap().hash != hash_type as i32 {
        return Err("KeyFormat is incorrect".into());
    }
    let keymanager = tink_core::registry::get_key_manager(tink_tests::HMAC_PRF_TYPE_URL)
        .map_err(|e| wrap_err("Could not obtain HMAC key manager", e))?;
    assert!(
        keymanager.new_key(&template.value).is_ok(),
        "HMAC key manager cannot create key"
    );
    Ok(())
}

fn check_hkdf_template(
    template: &tink_proto::KeyTemplate,
    key_size: u32,
    salt: &[u8],
    hash_type: tink_proto::HashType,
) -> Result<(), TinkError> {
    if template.type_url != tink_tests::HKDF_PRF_TYPE_URL {
        return Err("TypeUrl is incorrect".into());
    }
    if template.output_prefix_type != tink_proto::OutputPrefixType::Raw as i32 {
        return Err("Not RAW output prefix".into());
    }

    let format = tink_proto::HkdfPrfKeyFormat::decode(template.value.as_ref())
        .map_err(|_| TinkError::new("unable to unmarshal serialized key format"))?;
    if format.key_size != key_size
        || format.params.as_ref().unwrap().hash != hash_type as i32
        || hex::encode(salt) != hex::encode(format.params.unwrap().salt)
    {
        return Err("KeyFormat is incorrect".into());
    }
    let keymanager = tink_core::registry::get_key_manager(tink_tests::HKDF_PRF_TYPE_URL)
        .map_err(|e| wrap_err("Could not obtain HKDF key manager", e))?;
    assert!(
        keymanager.new_key(&template.value).is_ok(),
        "HKDF key manager cannot create key"
    );
    Ok(())
}

fn check_cmac_template(template: &tink_proto::KeyTemplate, key_size: u32) -> Result<(), TinkError> {
    if template.type_url != tink_tests::AES_CMAC_PRF_TYPE_URL {
        return Err("TypeUrl is incorrect".into());
    }
    if template.output_prefix_type != tink_proto::OutputPrefixType::Raw as i32 {
        return Err("Not RAW output prefix".into());
    }

    let format = tink_proto::AesCmacPrfKeyFormat::decode(template.value.as_ref())
        .map_err(|_| TinkError::new("unable to unmarshal serialized key format"))?;
    if format.key_size != key_size {
        return Err("KeyFormat is incorrect".into());
    }
    let keymanager = tink_core::registry::get_key_manager(tink_tests::AES_CMAC_PRF_TYPE_URL)
        .map_err(|e| wrap_err("Could not obtain AES-CMAC key manager", e))?;
    assert!(
        keymanager.new_key(&template.value).is_ok(),
        "AES-CMAC key manager cannot create key"
    );
    Ok(())
}
<|endoftext|>


<commit_before>use std::f32;
use std::f64;

use protobuf::reflect::FieldDescriptor;
use protobuf::reflect::MessageDescriptor;
use protobuf::reflect::ReflectValueBox;
use protobuf::reflect::RuntimeFieldType;
use protobuf::reflect::RuntimeTypeBox;
use protobuf::reflect::RuntimeTypeDynamic;
use protobuf::well_known_types::Value;
use protobuf::well_known_types::value;
use protobuf::Message;

pub fn value_for_runtime_type(field_type: &dyn RuntimeTypeDynamic) -> ReflectValueBox {
    match field_type.to_box() {
        RuntimeTypeBox::U32 => ReflectValueBox::U32(11),
        RuntimeTypeBox::U64 => ReflectValueBox::U64(12),
        RuntimeTypeBox::I32 => ReflectValueBox::I32(13),
        RuntimeTypeBox::I64 => ReflectValueBox::I64(14),
        RuntimeTypeBox::F32 => ReflectValueBox::F32(15.5),
        RuntimeTypeBox::F64 => ReflectValueBox::F64(16.5),
        RuntimeTypeBox::Bool => ReflectValueBox::Bool(true),
        RuntimeTypeBox::String | RuntimeTypeBox::Chars => {
            ReflectValueBox::String("here".to_owned())
        }
        RuntimeTypeBox::VecU8 | RuntimeTypeBox::CarllercheBytes => {
            ReflectValueBox::Bytes(b"there".as_ref().to_owned())
        }
        RuntimeTypeBox::Enum(e) => ReflectValueBox::from(&e.values()[0]),
        RuntimeTypeBox::Message(m) => ReflectValueBox::Message(m.new_instance()),
    }
}

fn values_for_message_type(descriptor: &MessageDescriptor) -> Vec<Box<dyn Message>> {
    if descriptor == Value::descriptor_static() {
        // special handling because empty `Value` is not valid
        let mut value = Value::new();
        value.kind = Some(value::Kind::number_value(23.0));
        vec![Box::new(value)]
    } else {
        vec![
            // TODO: populated messages
            descriptor.new_instance(),
        ]
    }
}

pub fn values_for_runtime_type(field_type: &dyn RuntimeTypeDynamic) -> Vec<ReflectValueBox> {
    match field_type.to_box() {
        RuntimeTypeBox::U32 => vec![
            ReflectValueBox::U32(11),
            ReflectValueBox::U32(0),
            ReflectValueBox::U32(0x7fff_ffff),
            ReflectValueBox::U32(0x8000_0000),
            ReflectValueBox::U32(0xffff_ffff),
        ],
        RuntimeTypeBox::U64 => vec![
            ReflectValueBox::U64(12),
            ReflectValueBox::U64(0),
            ReflectValueBox::U64(0xffff_ffff),
            ReflectValueBox::U64(0xffff_ffff_ffff_ffff),
        ],
        RuntimeTypeBox::I32 => vec![
            ReflectValueBox::I32(13),
            ReflectValueBox::I32(0),
            ReflectValueBox::I32(-1),
            ReflectValueBox::I32(0x7fff_ffff),
            ReflectValueBox::I32(0x1000_0000),
        ],
        RuntimeTypeBox::I64 => vec![
            ReflectValueBox::I64(14),
            ReflectValueBox::I64(0),
            ReflectValueBox::I64(-1),
            ReflectValueBox::I64(0x7fff_ffff_ffff_ffff),
            ReflectValueBox::I64(0x1000_0000_0000_0000),
        ],
        RuntimeTypeBox::F32 => vec![
            ReflectValueBox::F32(15.5),
            ReflectValueBox::F32(0.0),
            ReflectValueBox::F32(-0.0),
            ReflectValueBox::F32(1. / 3.0),
            ReflectValueBox::F32(-1. / 3.0),
            ReflectValueBox::F32(f32::NAN),
            ReflectValueBox::F32(f32::INFINITY),
            ReflectValueBox::F32(-f32::INFINITY),
        ],
        RuntimeTypeBox::F64 => vec![
            ReflectValueBox::F64(16.5),
            ReflectValueBox::F64(0.0),
            ReflectValueBox::F64(-0.0),
            ReflectValueBox::F64(1.0 / 3.0),
            ReflectValueBox::F64(-1.0 / 3.0),
            ReflectValueBox::F64(f64::NAN),
            ReflectValueBox::F64(f64::INFINITY),
            ReflectValueBox::F64(-f64::INFINITY),
        ],
        RuntimeTypeBox::Bool => vec![ReflectValueBox::Bool(true), ReflectValueBox::Bool(false)],
        RuntimeTypeBox::String | RuntimeTypeBox::Chars => vec![
            ReflectValueBox::String("here".to_owned()),
            ReflectValueBox::String("".to_owned()),
            ReflectValueBox::String(" \t\n".to_owned()),
            ReflectValueBox::String("\0".to_owned()),
        ],
        RuntimeTypeBox::VecU8 | RuntimeTypeBox::CarllercheBytes => vec![
            ReflectValueBox::Bytes(b"there".as_ref().to_owned()),
            ReflectValueBox::Bytes(b"".as_ref().to_owned()),
        ],
        RuntimeTypeBox::Enum(e) => vec![
            ReflectValueBox::from(&e.values()[0]),
            ReflectValueBox::from(&e.values()[e.values().len() - 1]),
        ],
        RuntimeTypeBox::Message(m) => values_for_message_type(m)
            .into_iter()
            .map(ReflectValueBox::from)
            .collect(),
    }
}

pub fn special_values_for_field(f: &FieldDescriptor, d: &MessageDescriptor) -> Vec<Box<dyn Message>> {
    let mut r = Vec::new();
    match f.runtime_field_type() {
        RuntimeFieldType::Singular(t) => {
            for v in values_for_runtime_type(t) {
                let mut m = d.new_instance();
                f.set_singular_field(&mut *m, v);
                r.push(m);
            }
        }
        RuntimeFieldType::Repeated(t) => {
            // TODO: empty repeated
            // TODO: repeated of more than one element
            let mut m = d.new_instance();
            f.mut_repeated(&mut *m).push(value_for_runtime_type(t));
            r.push(m);
        }
        RuntimeFieldType::Map(k, v) => {
            // TODO: empty map
            // TODO: map of more than one element
            let mut m = d.new_instance();
            let k = value_for_runtime_type(k);
            let v = value_for_runtime_type(v);
            f.mut_map(&mut *m).insert(k, v);
            r.push(m);
        }
    }
    r
}

pub fn special_messages(d: &MessageDescriptor) -> Vec<Box<dyn Message>> {
    let mut r = Vec::new();
    for f in d.fields() {
        r.extend(special_values_for_field(f, d));
    }
    r
}

pub fn special_messages_typed<M: Message>() -> Vec<M> {
    let mut r = Vec::new();
    for m in special_messages(M::descriptor_static()) {
        r.push(*m.downcast_box().unwrap());
    }
    r
}
<commit_msg>Fix test compilation<commit_after>use std::f32;
use std::f64;

use protobuf::reflect::FieldDescriptor;
use protobuf::reflect::MessageDescriptor;
use protobuf::reflect::ReflectValueBox;
use protobuf::reflect::RuntimeFieldType;
use protobuf::reflect::RuntimeTypeBox;
use protobuf::reflect::RuntimeTypeDynamic;
use protobuf::well_known_types::Value;
use protobuf::well_known_types::value;
use protobuf::Message;

pub fn value_for_runtime_type(field_type: &dyn RuntimeTypeDynamic) -> ReflectValueBox {
    match field_type.to_box() {
        RuntimeTypeBox::U32 => ReflectValueBox::U32(11),
        RuntimeTypeBox::U64 => ReflectValueBox::U64(12),
        RuntimeTypeBox::I32 => ReflectValueBox::I32(13),
        RuntimeTypeBox::I64 => ReflectValueBox::I64(14),
        RuntimeTypeBox::F32 => ReflectValueBox::F32(15.5),
        RuntimeTypeBox::F64 => ReflectValueBox::F64(16.5),
        RuntimeTypeBox::Bool => ReflectValueBox::Bool(true),
        RuntimeTypeBox::String => ReflectValueBox::String("here".to_owned()),
        RuntimeTypeBox::VecU8 => ReflectValueBox::Bytes(b"there".as_ref().to_owned()),
        RuntimeTypeBox::Enum(e) => ReflectValueBox::from(&e.values()[0]),
        RuntimeTypeBox::Message(m) => ReflectValueBox::Message(m.new_instance()),
    }
}

fn values_for_message_type(descriptor: &MessageDescriptor) -> Vec<Box<dyn Message>> {
    if descriptor == Value::descriptor_static() {
        // special handling because empty `Value` is not valid
        let mut value = Value::new();
        value.kind = Some(value::Kind::number_value(23.0));
        vec![Box::new(value)]
    } else {
        vec![
            // TODO: populated messages
            descriptor.new_instance(),
        ]
    }
}

pub fn values_for_runtime_type(field_type: &dyn RuntimeTypeDynamic) -> Vec<ReflectValueBox> {
    match field_type.to_box() {
        RuntimeTypeBox::U32 => vec![
            ReflectValueBox::U32(11),
            ReflectValueBox::U32(0),
            ReflectValueBox::U32(0x7fff_ffff),
            ReflectValueBox::U32(0x8000_0000),
            ReflectValueBox::U32(0xffff_ffff),
        ],
        RuntimeTypeBox::U64 => vec![
            ReflectValueBox::U64(12),
            ReflectValueBox::U64(0),
            ReflectValueBox::U64(0xffff_ffff),
            ReflectValueBox::U64(0xffff_ffff_ffff_ffff),
        ],
        RuntimeTypeBox::I32 => vec![
            ReflectValueBox::I32(13),
            ReflectValueBox::I32(0),
            ReflectValueBox::I32(-1),
            ReflectValueBox::I32(0x7fff_ffff),
            ReflectValueBox::I32(0x1000_0000),
        ],
        RuntimeTypeBox::I64 => vec![
            ReflectValueBox::I64(14),
            ReflectValueBox::I64(0),
            ReflectValueBox::I64(-1),
            ReflectValueBox::I64(0x7fff_ffff_ffff_ffff),
            ReflectValueBox::I64(0x1000_0000_0000_0000),
        ],
        RuntimeTypeBox::F32 => vec![
            ReflectValueBox::F32(15.5),
            ReflectValueBox::F32(0.0),
            ReflectValueBox::F32(-0.0),
            ReflectValueBox::F32(1. / 3.0),
            ReflectValueBox::F32(-1. / 3.0),
            ReflectValueBox::F32(f32::NAN),
            ReflectValueBox::F32(f32::INFINITY),
            ReflectValueBox::F32(-f32::INFINITY),
        ],
        RuntimeTypeBox::F64 => vec![
            ReflectValueBox::F64(16.5),
            ReflectValueBox::F64(0.0),
            ReflectValueBox::F64(-0.0),
            ReflectValueBox::F64(1.0 / 3.0),
            ReflectValueBox::F64(-1.0 / 3.0),
            ReflectValueBox::F64(f64::NAN),
            ReflectValueBox::F64(f64::INFINITY),
            ReflectValueBox::F64(-f64::INFINITY),
        ],
        RuntimeTypeBox::Bool => vec![ReflectValueBox::Bool(true), ReflectValueBox::Bool(false)],
        RuntimeTypeBox::String => vec![
            ReflectValueBox::String("here".to_owned()),
            ReflectValueBox::String("".to_owned()),
            ReflectValueBox::String(" \t\n".to_owned()),
            ReflectValueBox::String("\0".to_owned()),
        ],
        RuntimeTypeBox::VecU8 => vec![
            ReflectValueBox::Bytes(b"there".as_ref().to_owned()),
            ReflectValueBox::Bytes(b"".as_ref().to_owned()),
        ],
        RuntimeTypeBox::Enum(e) => vec![
            ReflectValueBox::from(&e.values()[0]),
            ReflectValueBox::from(&e.values()[e.values().len() - 1]),
        ],
        RuntimeTypeBox::Message(m) => values_for_message_type(m)
            .into_iter()
            .map(ReflectValueBox::from)
            .collect(),
    }
}

pub fn special_values_for_field(f: &FieldDescriptor, d: &MessageDescriptor) -> Vec<Box<dyn Message>> {
    let mut r = Vec::new();
    match f.runtime_field_type() {
        RuntimeFieldType::Singular(t) => {
            for v in values_for_runtime_type(t) {
                let mut m = d.new_instance();
                f.set_singular_field(&mut *m, v);
                r.push(m);
            }
        }
        RuntimeFieldType::Repeated(t) => {
            // TODO: empty repeated
            // TODO: repeated of more than one element
            let mut m = d.new_instance();
            f.mut_repeated(&mut *m).push(value_for_runtime_type(t));
            r.push(m);
        }
        RuntimeFieldType::Map(k, v) => {
            // TODO: empty map
            // TODO: map of more than one element
            let mut m = d.new_instance();
            let k = value_for_runtime_type(k);
            let v = value_for_runtime_type(v);
            f.mut_map(&mut *m).insert(k, v);
            r.push(m);
        }
    }
    r
}

pub fn special_messages(d: &MessageDescriptor) -> Vec<Box<dyn Message>> {
    let mut r = Vec::new();
    for f in d.fields() {
        r.extend(special_values_for_field(f, d));
    }
    r
}

pub fn special_messages_typed<M: Message>() -> Vec<M> {
    let mut r = Vec::new();
    for m in special_messages(M::descriptor_static()) {
        r.push(*m.downcast_box().unwrap());
    }
    r
}
<|endoftext|>