<commit_before>//! Simple module to store files in database.
//!
//! cratesfyi is generating more than 5 million files, they are small and mostly html files.
//! They are using so many inodes and it is better to store them in database instead of
//! filesystem. This module is adding files into database and retrieving them.

use std::path::{PathBuf, Path};
use postgres::Connection;
use rustc_serialize::json::{Json, ToJson};
use std::cmp;
use std::fs;
use std::io::Read;
use error::Result;
use failure::err_msg;
use rusoto_s3::{S3, PutObjectRequest, GetObjectRequest, S3Client};
use rusoto_core::region::Region;
use rusoto_credential::DefaultCredentialsProvider;

const MAX_CONCURRENT_UPLOADS: usize = 1000;

fn get_file_list_from_dir<P: AsRef<Path>>(path: P,
                                          files: &mut Vec<PathBuf>)
                                          -> Result<()> {
    let path = path.as_ref();

    for file in try!(path.read_dir()) {
        let file = try!(file);

        if try!(file.file_type()).is_file() {
            files.push(file.path());
        } else if try!(file.file_type()).is_dir() {
            try!(get_file_list_from_dir(file.path(), files));
        }
    }

    Ok(())
}

pub fn get_file_list<P: AsRef<Path>>(path: P) -> Result<Vec<PathBuf>> {
    let path = path.as_ref();
    let mut files = Vec::new();

    if !path.exists() {
        return Err(err_msg("File not found"));
    } else if path.is_file() {
        files.push(PathBuf::from(path.file_name().unwrap()));
    } else if path.is_dir() {
        try!(get_file_list_from_dir(path, &mut files));
        for file_path in &mut files {
            // We want the paths in this list to not be {path}/bar.txt but just bar.txt
            *file_path = PathBuf::from(file_path.strip_prefix(path).unwrap());
        }
    }

    Ok(files)
}

pub struct Blob {
    pub path: String,
    pub mime: String,
    pub date_updated: time::Timespec,
    pub content: Vec<u8>,
}

pub fn get_path(conn: &Connection, path: &str) -> Option<Blob> {
    if let Some(client) = s3_client() {
        let res = client.get_object(GetObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.into(),
            ..Default::default()
        }).sync();

        let res = match res {
            Ok(r) => r,
            Err(_) => {
                return None;
            }
        };

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::new();
        b.read_to_end(&mut content).unwrap();

        let last_modified = res.last_modified.unwrap();
        let last_modified = time::strptime(&last_modified, "%a, %d %b %Y %H:%M:%S %Z")
            .unwrap_or_else(|e| panic!("failed to parse {:?} as timespec: {:?}", last_modified, e))
            .to_timespec();

        Some(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated: last_modified,
            content,
        })
    } else {
        let rows = conn.query("SELECT path, mime, date_updated, content
                            FROM files
                            WHERE path = $1", &[&path]).unwrap();

        if rows.len() == 0 {
            None
        } else {
            let row = rows.get(0);

            Some(Blob {
                path: row.get(0),
                mime: row.get(1),
                date_updated: row.get(2),
                content: row.get(3),
            })
        }
    }
}

fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() && std::env::var_os("FORCE_S3").is_none() {
        return None;
    }
    let creds = match DefaultCredentialsProvider::new() {
        Ok(creds) => creds,
        Err(err) => {
            warn!("failed to retrieve AWS credentials: {}", err);
            return None;
        }
    };
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        creds,
        std::env::var("S3_ENDPOINT").ok().map(|e| Region::Custom {
            name: "us-west-1".to_owned(),
            endpoint: e,
        }).unwrap_or(Region::UsWest1),
    ))
}

/// Adds files into database and returns list of files with their mime type in Json
pub fn add_path_into_database<P: AsRef<Path>>(conn: &Connection,
                                              prefix: &str,
                                              path: P)
                                              -> Result<Json> {
    use magic::{Cookie, flags};
    let cookie = try!(Cookie::open(flags::MIME_TYPE));
    try!(cookie.load::<&str>(&[]));

    let trans = try!(conn.transaction());

    use std::collections::HashMap;
    let mut file_paths_and_mimes: HashMap<PathBuf, String> = HashMap::new();

    let mut rt = ::tokio::runtime::Runtime::new().unwrap();

    let mut to_upload = try!(get_file_list(&path));
    let mut batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
    let mut currently_uploading: Vec<_> = to_upload.drain(..batch_size).collect();
    let mut attempts = 0;

    while !to_upload.is_empty() || !currently_uploading.is_empty() {

        let mut futures = Vec::new();
        let client = s3_client();

        for file_path in &currently_uploading {
            let path = Path::new(path.as_ref()).join(&file_path);
            // Some files have insufficient permissions (like .lock file created by cargo in
            // documentation directory). We are skipping this files.
            let mut file = match fs::File::open(path) {
                Ok(f) => f,
                Err(_) => continue,
            };
            let mut content: Vec<u8> = Vec::new();
            try!(file.read_to_end(&mut content));
            let bucket_path = Path::new(prefix).join(&file_path)
                .into_os_string().into_string().unwrap();

            let mime = {
                let mime = try!(cookie.buffer(&content));
                // css's are causing some problem in browsers
                // magic will return text/plain for css file types
                // convert them to text/css
                // do the same for javascript files
                if mime == "text/plain" {
                    let e = file_path.extension().unwrap_or_default();
                    if e == "css" {
                        "text/css".to_owned()
                    } else if e == "js" {
                        "application/javascript".to_owned()
                    } else {
                        mime.to_owned()
                    }
                } else {
                    mime.to_owned()
                }
            };

            if let Some(client) = &client {
                futures.push(client.put_object(PutObjectRequest {
                    bucket: "rust-docs-rs".into(),
                    key: bucket_path.clone(),
                    body: Some(content.clone().into()),
                    content_type: Some(mime.clone()),
                    ..Default::default()
                }));
            } else {
                // If AWS credentials are configured, don't insert/update the database
                // check if file already exists in database
                let rows = try!(conn.query("SELECT COUNT(*) FROM files WHERE path = $1", &[&bucket_path]));

                if rows.get(0).get::<usize, i64>(0) == 0 {
                    try!(trans.query("INSERT INTO files (path, mime, content) VALUES ($1, $2, $3)",
                                    &[&bucket_path, &mime, &content]));
                } else {
                    try!(trans.query("UPDATE files SET mime = $2, content = $3, date_updated = NOW() \
                                    WHERE path = $1",
                                    &[&bucket_path, &mime, &content]));
                }
            }

            file_paths_and_mimes.insert(file_path.clone(), mime.clone());
        }

        if !futures.is_empty() {
            attempts += 1;

            if rt.block_on(::futures::future::join_all(futures)).is_ok() {
                // this batch was successful, start another batch if there are still more files
                crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(batch_size as i64);
                batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
                currently_uploading = to_upload.drain(..batch_size).collect();
                attempts = 0;
            } else {
                // if any futures error, leave `currently_uploading` in tact so that we can retry the batch
                if attempts > 2 {
                    panic!("failed to upload 3 times, exiting");
                }
            }
        } else {
            crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(batch_size as i64);
            batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
            currently_uploading = to_upload.drain(..batch_size).collect();
        }
    }

    try!(trans.commit());

    let file_list_with_mimes: Vec<(String, PathBuf)> = file_paths_and_mimes
        .into_iter()
        .map(|(file_path, mime)| (mime, file_path))
        .collect();
    file_list_to_json(file_list_with_mimes)
}

fn file_list_to_json(file_list: Vec<(String, PathBuf)>) -> Result<Json> {

    let mut file_list_json: Vec<Json> = Vec::new();

    for file in file_list {
        let mut v: Vec<String> = Vec::new();
        v.push(file.0.clone());
        v.push(file.1.into_os_string().into_string().unwrap());
        file_list_json.push(v.to_json());
    }

    Ok(file_list_json.to_json())
}

pub fn move_to_s3(conn: &Connection, n: usize) -> Result<usize> {
    let trans = try!(conn.transaction());
    let client = s3_client().expect("configured s3");

    let rows = try!(trans.query(
            &format!("SELECT path, mime, content FROM files WHERE content != E'in-s3' LIMIT {}", n),
            &[]));
    let count = rows.len();

    let mut rt = ::tokio::runtime::Runtime::new().unwrap();
    let mut futures = Vec::new();
    for row in &rows {
        let path: String = row.get(0);
        let mime: String = row.get(1);
        let content: Vec<u8> = row.get(2);
        let path_1 = path.clone();
        futures.push(client.put_object(PutObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.clone(),
            body: Some(content.into()),
            content_type: Some(mime),
            ..Default::default()
        }).map(move |_| {
            path_1
        }).map_err(move |e| {
            panic!("failed to upload to {}: {:?}", path, e)
        }));
    }

    use ::futures::future::Future;
    match rt.block_on(::futures::future::join_all(futures)) {
        Ok(paths) => {
            let statement = trans.prepare("DELETE FROM files WHERE path = $1").unwrap();
            for path in paths {
                statement.execute(&[&path]).unwrap();
            }
        }
        Err(e) => {
            panic!("results err: {:?}", e);
        }
    }

    try!(trans.commit());

    Ok(count)
}

#[cfg(test)]
mod test {
    extern crate env_logger;
    use std::env;
    use super::get_file_list;

    #[test]
    fn test_get_file_list() {
        let _ = env_logger::try_init();

        let files = get_file_list(env::current_dir().unwrap());
        assert!(files.is_ok());
        assert!(files.unwrap().len() > 0);

        let files = get_file_list(env::current_dir().unwrap().join("Cargo.toml")).unwrap();
        assert_eq!(files[0], std::path::Path::new("Cargo.toml"));
    }
}
<commit_msg>increment upload metric for each successful file<commit_after>//! Simple module to store files in database.
//!
//! cratesfyi is generating more than 5 million files, they are small and mostly html files.
//! They are using so many inodes and it is better to store them in database instead of
//! filesystem. This module is adding files into database and retrieving them.

use std::path::{PathBuf, Path};
use postgres::Connection;
use rustc_serialize::json::{Json, ToJson};
use std::cmp;
use std::fs;
use std::io::Read;
use error::Result;
use failure::err_msg;
use rusoto_s3::{S3, PutObjectRequest, GetObjectRequest, S3Client};
use rusoto_core::region::Region;
use rusoto_credential::DefaultCredentialsProvider;

const MAX_CONCURRENT_UPLOADS: usize = 1000;

fn get_file_list_from_dir<P: AsRef<Path>>(path: P,
                                          files: &mut Vec<PathBuf>)
                                          -> Result<()> {
    let path = path.as_ref();

    for file in try!(path.read_dir()) {
        let file = try!(file);

        if try!(file.file_type()).is_file() {
            files.push(file.path());
        } else if try!(file.file_type()).is_dir() {
            try!(get_file_list_from_dir(file.path(), files));
        }
    }

    Ok(())
}

pub fn get_file_list<P: AsRef<Path>>(path: P) -> Result<Vec<PathBuf>> {
    let path = path.as_ref();
    let mut files = Vec::new();

    if !path.exists() {
        return Err(err_msg("File not found"));
    } else if path.is_file() {
        files.push(PathBuf::from(path.file_name().unwrap()));
    } else if path.is_dir() {
        try!(get_file_list_from_dir(path, &mut files));
        for file_path in &mut files {
            // We want the paths in this list to not be {path}/bar.txt but just bar.txt
            *file_path = PathBuf::from(file_path.strip_prefix(path).unwrap());
        }
    }

    Ok(files)
}

pub struct Blob {
    pub path: String,
    pub mime: String,
    pub date_updated: time::Timespec,
    pub content: Vec<u8>,
}

pub fn get_path(conn: &Connection, path: &str) -> Option<Blob> {
    if let Some(client) = s3_client() {
        let res = client.get_object(GetObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.into(),
            ..Default::default()
        }).sync();

        let res = match res {
            Ok(r) => r,
            Err(_) => {
                return None;
            }
        };

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::new();
        b.read_to_end(&mut content).unwrap();

        let last_modified = res.last_modified.unwrap();
        let last_modified = time::strptime(&last_modified, "%a, %d %b %Y %H:%M:%S %Z")
            .unwrap_or_else(|e| panic!("failed to parse {:?} as timespec: {:?}", last_modified, e))
            .to_timespec();

        Some(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated: last_modified,
            content,
        })
    } else {
        let rows = conn.query("SELECT path, mime, date_updated, content
                            FROM files
                            WHERE path = $1", &[&path]).unwrap();

        if rows.len() == 0 {
            None
        } else {
            let row = rows.get(0);

            Some(Blob {
                path: row.get(0),
                mime: row.get(1),
                date_updated: row.get(2),
                content: row.get(3),
            })
        }
    }
}

fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() && std::env::var_os("FORCE_S3").is_none() {
        return None;
    }
    let creds = match DefaultCredentialsProvider::new() {
        Ok(creds) => creds,
        Err(err) => {
            warn!("failed to retrieve AWS credentials: {}", err);
            return None;
        }
    };
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        creds,
        std::env::var("S3_ENDPOINT").ok().map(|e| Region::Custom {
            name: "us-west-1".to_owned(),
            endpoint: e,
        }).unwrap_or(Region::UsWest1),
    ))
}

/// Adds files into database and returns list of files with their mime type in Json
pub fn add_path_into_database<P: AsRef<Path>>(conn: &Connection,
                                              prefix: &str,
                                              path: P)
                                              -> Result<Json> {
    use magic::{Cookie, flags};
    let cookie = try!(Cookie::open(flags::MIME_TYPE));
    try!(cookie.load::<&str>(&[]));

    let trans = try!(conn.transaction());

    use std::collections::HashMap;
    let mut file_paths_and_mimes: HashMap<PathBuf, String> = HashMap::new();

    use futures::future::Future;

    let mut rt = ::tokio::runtime::Runtime::new().unwrap();

    let mut to_upload = try!(get_file_list(&path));
    let mut batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
    let mut currently_uploading: Vec<_> = to_upload.drain(..batch_size).collect();
    let mut attempts = 0;

    while !to_upload.is_empty() || !currently_uploading.is_empty() {

        let mut futures = Vec::new();
        let client = s3_client();

        for file_path in &currently_uploading {
            let path = Path::new(path.as_ref()).join(&file_path);
            // Some files have insufficient permissions (like .lock file created by cargo in
            // documentation directory). We are skipping this files.
            let mut file = match fs::File::open(path) {
                Ok(f) => f,
                Err(_) => continue,
            };
            let mut content: Vec<u8> = Vec::new();
            try!(file.read_to_end(&mut content));
            let bucket_path = Path::new(prefix).join(&file_path)
                .into_os_string().into_string().unwrap();

            let mime = {
                let mime = try!(cookie.buffer(&content));
                // css's are causing some problem in browsers
                // magic will return text/plain for css file types
                // convert them to text/css
                // do the same for javascript files
                if mime == "text/plain" {
                    let e = file_path.extension().unwrap_or_default();
                    if e == "css" {
                        "text/css".to_owned()
                    } else if e == "js" {
                        "application/javascript".to_owned()
                    } else {
                        mime.to_owned()
                    }
                } else {
                    mime.to_owned()
                }
            };

            if let Some(client) = &client {
                futures.push(client.put_object(PutObjectRequest {
                    bucket: "rust-docs-rs".into(),
                    key: bucket_path.clone(),
                    body: Some(content.clone().into()),
                    content_type: Some(mime.clone()),
                    ..Default::default()
                }).inspect(|_| {
                    crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(1);
                }));
            } else {
                // If AWS credentials are configured, don't insert/update the database
                // check if file already exists in database
                let rows = try!(conn.query("SELECT COUNT(*) FROM files WHERE path = $1", &[&bucket_path]));

                if rows.get(0).get::<usize, i64>(0) == 0 {
                    try!(trans.query("INSERT INTO files (path, mime, content) VALUES ($1, $2, $3)",
                                    &[&bucket_path, &mime, &content]));
                } else {
                    try!(trans.query("UPDATE files SET mime = $2, content = $3, date_updated = NOW() \
                                    WHERE path = $1",
                                    &[&bucket_path, &mime, &content]));
                }
            }

            file_paths_and_mimes.insert(file_path.clone(), mime.clone());
        }

        if !futures.is_empty() {
            attempts += 1;

            if rt.block_on(::futures::future::join_all(futures)).is_ok() {
                // this batch was successful, start another batch if there are still more files
                batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
                currently_uploading = to_upload.drain(..batch_size).collect();
                attempts = 0;
            } else {
                // if any futures error, leave `currently_uploading` in tact so that we can retry the batch
                if attempts > 2 {
                    panic!("failed to upload 3 times, exiting");
                }
            }
        } else {
            batch_size = cmp::min(to_upload.len(), MAX_CONCURRENT_UPLOADS);
            currently_uploading = to_upload.drain(..batch_size).collect();
        }
    }

    try!(trans.commit());

    let file_list_with_mimes: Vec<(String, PathBuf)> = file_paths_and_mimes
        .into_iter()
        .map(|(file_path, mime)| (mime, file_path))
        .collect();
    file_list_to_json(file_list_with_mimes)
}

fn file_list_to_json(file_list: Vec<(String, PathBuf)>) -> Result<Json> {

    let mut file_list_json: Vec<Json> = Vec::new();

    for file in file_list {
        let mut v: Vec<String> = Vec::new();
        v.push(file.0.clone());
        v.push(file.1.into_os_string().into_string().unwrap());
        file_list_json.push(v.to_json());
    }

    Ok(file_list_json.to_json())
}

pub fn move_to_s3(conn: &Connection, n: usize) -> Result<usize> {
    let trans = try!(conn.transaction());
    let client = s3_client().expect("configured s3");

    let rows = try!(trans.query(
            &format!("SELECT path, mime, content FROM files WHERE content != E'in-s3' LIMIT {}", n),
            &[]));
    let count = rows.len();

    let mut rt = ::tokio::runtime::Runtime::new().unwrap();
    let mut futures = Vec::new();
    for row in &rows {
        let path: String = row.get(0);
        let mime: String = row.get(1);
        let content: Vec<u8> = row.get(2);
        let path_1 = path.clone();
        futures.push(client.put_object(PutObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.clone(),
            body: Some(content.into()),
            content_type: Some(mime),
            ..Default::default()
        }).map(move |_| {
            path_1
        }).map_err(move |e| {
            panic!("failed to upload to {}: {:?}", path, e)
        }));
    }

    use ::futures::future::Future;
    match rt.block_on(::futures::future::join_all(futures)) {
        Ok(paths) => {
            let statement = trans.prepare("DELETE FROM files WHERE path = $1").unwrap();
            for path in paths {
                statement.execute(&[&path]).unwrap();
            }
        }
        Err(e) => {
            panic!("results err: {:?}", e);
        }
    }

    try!(trans.commit());

    Ok(count)
}

#[cfg(test)]
mod test {
    extern crate env_logger;
    use std::env;
    use super::get_file_list;

    #[test]
    fn test_get_file_list() {
        let _ = env_logger::try_init();

        let files = get_file_list(env::current_dir().unwrap());
        assert!(files.is_ok());
        assert!(files.unwrap().len() > 0);

        let files = get_file_list(env::current_dir().unwrap().join("Cargo.toml")).unwrap();
        assert_eq!(files[0], std::path::Path::new("Cargo.toml"));
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 Jonathan Eyolfson

use std::ptr;

use raw;

pub struct Display {
    ptr: *mut raw::wl_display
}

impl Display {
    pub fn new() -> Display {
        unsafe {
            let ptr = raw::wl_display_connect(ptr::null());
            assert!(!ptr.is_null());
            Display { ptr: ptr }
        }
    }
    pub unsafe fn to_ptr(&mut self) -> *mut raw::wl_display { self.ptr }
    pub fn dispatch(&mut self) {
        unsafe {
            let ret = raw::wl_display_dispatch(self.ptr);
            assert!(ret >= 0);
        }
    }
    pub fn roundtrip(&mut self) {
        unsafe {
            let ret = raw::wl_display_roundtrip(self.ptr);
            assert!(ret >= 0);
        }
    }
    pub fn flush(&mut self) {
        unsafe {
            let ret = raw::wl_display_flush(self.ptr);
            assert!(ret >= 0);
        }
    }
}

impl Drop for Display {
    fn drop(&mut self) {
        unsafe {
            raw::wl_display_disconnect(self.ptr);
        }
    }
}
<commit_msg>Display wrapper style changed and added more functions<commit_after>// Copyright 2014 Jonathan Eyolfson

use std::ptr;

use raw;

pub struct Display {
    ptr: *mut raw::wl_display
}

impl Display {
    pub fn connect_to_env_or_default() -> Display {
        unsafe {
            let ptr = raw::wl_display_connect(ptr::null());
            assert!(!ptr.is_null(), "wl_display_connect failed");
            Display { ptr: ptr }
        }
    }
    pub fn connect(name: &str) -> Display {
        unsafe {
            let ptr = raw::wl_display_connect(name.as_ptr() as *const i8);
            assert!(!ptr.is_null(), "wl_display_connect failed");
            Display { ptr: ptr }
        }
    }
    pub fn roundtrip(&mut self) -> i32 {
        unsafe {
            let r = raw::wl_display_roundtrip(self.ptr);
            assert!(r != -1, "wl_display_roundtrip failed");
            r
        }
    }
    pub fn read_events(&mut self) {
        unsafe {
            let r = raw::wl_display_read_events(self.ptr);
            assert!(r != -1, "wl_display_read_events failed");
        }
    }
    pub fn prepare_read(&mut self) {
        unsafe {
            let r = raw::wl_display_prepare_read(self.ptr);
            assert!(r != -1, "wl_display_prepare_read failed");
        }
    }
    pub fn cancel_read(&mut self) {
        unsafe {
            raw::wl_display_cancel_read(self.ptr);
        }
    }
    pub fn dispatch(&mut self) -> i32 {
        unsafe {
            let r = raw::wl_display_dispatch(self.ptr);
            assert!(r != -1);
            r
        }
    }
    pub fn flush(&mut self) -> i32 {
        unsafe {
            let r = raw::wl_display_flush(self.ptr);
            assert!(r != -1);
            r
        }
    }
    pub unsafe fn to_ptr(&mut self) -> *mut raw::wl_display { self.ptr }
}

impl Drop for Display {
    fn drop(&mut self) {
        unsafe {
            raw::wl_display_disconnect(self.ptr);
        }
    }
}
<|endoftext|>


<commit_before>use std::char;
use std::cmp::{max, min};
use std::collections::VecDeque;
use num::{Complex, Float};
use rustty;
use rustty::{Attr, Color, Terminal, Cell, CellAccessor, HasSize};
use rustty::ui::{Alignable, Widget, VerticalAlign};
use itertools::{Itertools, EitherOrBoth};

pub struct Canvas {
    term: Terminal,
    spectrum: Widget,
    waterfall: Widget,
    history: VecDeque<Vec<f32>>,
}

impl Canvas {
    pub fn new() -> Result<Self, rustty::Error> {
        let term = try!(Terminal::new());
        let (cols, rows) = term.size();

        let spectrum_height = rows / 2;
        let waterfall_height = if rows % 2 != 0 { rows / 2 } else { rows / 2 + 1 };

        let mut spectrum = Widget::new(cols, spectrum_height);
        spectrum.valign(&term, VerticalAlign::Top, 0);

        let mut waterfall = Widget::new(cols, waterfall_height);
        waterfall.valign(&term, VerticalAlign::Bottom, 0);

        Ok(Canvas {
            term: term,
            spectrum: spectrum,
            waterfall: waterfall,
            history: VecDeque::with_capacity(waterfall_height * 2),
        })
    }

    /// Adds a spectrum to the history and draws it on the waterfall
    /// and the spectrum view.
    pub fn add_spectrum(&mut self, spec: Vec<Complex<f32>>) {
        draw_spectrum(&mut self.spectrum, &spec);

        let normalized = normalize_spectrum(&spec, 50.0);
        let averaged = normalized.chunks(2).map(|v| (v[0] + v[1]) / 2.0).collect();
        self.history.push_front(averaged);

        draw_waterfall(&mut self.waterfall, &self.history);

        self.spectrum.draw_into(&mut self.term);
        self.waterfall.draw_into(&mut self.term);
        self.term.swap_buffers().unwrap();
    }

    pub fn get_term(&mut self) -> &mut Terminal {
        &mut self.term
    }
}

fn draw_waterfall<T: CellAccessor + HasSize>(canvas: &mut T, spectra: &VecDeque<Vec<f32>>) {
    let (cols, rows) = canvas.size();
    for (row, mut specs) in (0..rows).zip(&spectra.iter().chunks_lazy(2)) {
        let upper_heights = specs.next().into_iter().flat_map(|x| x);
        let lower_heights = specs.next().into_iter().flat_map(|x| x);
        for (c, heights) in (0..cols).zip(upper_heights.zip_longest(lower_heights)) {
            let (u, l) = match heights {
                EitherOrBoth::Both(&upper, &lower) => (upper, lower),
                EitherOrBoth::Left(&upper) => (upper, 0.0),
                EitherOrBoth::Right(&lower) => (0.0, lower),
            };
            *canvas.get_mut(c, row).unwrap() = spectrum_heights_to_waterfall_cell(u, l);
        }
    }
}

fn spectrum_heights_to_waterfall_cell(upper: f32, lower: f32) -> Cell {
    Cell::new('▀',
              Color::Byte(color_mapping(upper)),
              Color::Byte(color_mapping(lower)),
              Attr::Default)
}

/// Assumes `f` is between 0 and 1. Anything outside of this range
/// will be clamped.
fn color_mapping(f: f32) -> u8 {
    let mapping = [16, 17, 18, 19, 21, 27, 33, 39, 45, 51,
                   50, 49, 48, 47, 46, 82, 118, 154, 190, 226];
    let idx = (f * (mapping.len() as f32)) as i32;
    if idx < 0 {
        mapping[0]
    } else if idx >= mapping.len() as i32 {
        mapping[mapping.len() - 1]
    } else {
        mapping[idx as usize]
    }
    //let lower = 232.0;
    //let upper = 255.0;
    //let mapped = f * (upper - lower) + lower;
    //if mapped < lower {
        //lower as u8
    //} else if mapped > upper {
        //upper as u8
    //} else {
        //mapped as u8
    //}
}

fn normalize_spectrum(spec: &[Complex<f32>], max_db: f32) -> Vec<f32> {
    // FFT shift
    let (first_half, last_half) = spec.split_at((spec.len() + 1) / 2);
    let shifted_spec = last_half.iter().chain(first_half.iter());

    // normalize and take the log
    shifted_spec.map(Complex::norm)
                .map(Float::log10)
                .map(|x| 10.0 * x)
                .map(|x| x / max_db)
                .collect()
}

// indexing is from the top of the cell
fn pixel_nums_to_braille(p1: Option<u8>, p2: Option<u8>) -> char {
    let pixel_map = [[0x01, 0x08],
                     [0x02, 0x10],
                     [0x04, 0x20],
                     [0x40, 0x80]];

    let mut c = 0;
    if let Some(p) = p1 {
        for i in p..4 {
            c |= pixel_map[i as usize][0];
        }
    }

    if let Some(p) = p2 {
        for i in p..4 {
            c |= pixel_map[i as usize][1];
        }
    }

    char::from_u32((0x2800 + c) as u32).unwrap()
}

fn fft_shift<T: Clone>(spec: &mut [T]) {
    let spec_copy = spec.to_owned();

    let (first_half, last_half) = spec_copy.split_at((spec_copy.len() + 1) / 2);
    let shifted_spec = last_half.iter().chain(first_half.iter());
    for (x, y) in spec.iter_mut().zip(shifted_spec) {
        *x = y.clone();
    }
}

fn spectrum_to_bin_heights(spec: &[Complex<f32>], dest: &mut [f32]) {
    //TODO should be plotting in log scale

    // subsample
    let mut last_idx = -1isize;
    for (i, x) in spec.iter().map(Complex::norm).enumerate() {
        if (i * dest.len() / spec.len()) as isize > last_idx {
            last_idx += 1;
            dest[last_idx as usize] = x;
        }
    }

    //TODO unnecessary allocation
    fft_shift(dest);
}

fn char_to_cell(c: char) -> Cell {
    let mut cell = Cell::with_char(c);
    cell.set_attrs(Attr::Bold);
    cell
}

fn draw_pixel_pair<T>(canvas: &mut T, col_idx: usize, p1: usize, p2: usize)
    where T: CellAccessor + HasSize
{
    let (_, rows) = canvas.size();
    let max_pixel_height = 4 * rows;

    // clamp heights
    let p1 = if p1 >= max_pixel_height { max_pixel_height - 1} else { p1 };
    let p2 = if p2 >= max_pixel_height { max_pixel_height - 1} else { p2 };

    // Reverse it, since the terminal indexing is from the top
    let p1 = max_pixel_height - p1 - 1;
    let p2 = max_pixel_height - p2 - 1;

    // cell indices
    let c1 = p1 / 4;
    let c2 = p2 / 4;

    // Fill in full height cells.
    let full_cell_char = pixel_nums_to_braille(Some(0), Some(0));
    for row_idx in max(c1, c2)..rows {
        *canvas.get_mut(col_idx, row_idx).unwrap() = char_to_cell(full_cell_char);
    }

    let left_fill_cell_char = pixel_nums_to_braille(Some(0), None);
    for row_idx in min(c1, c2)..c2 {
        *canvas.get_mut(col_idx, row_idx).unwrap() = char_to_cell(left_fill_cell_char);
    }

    let right_fill_cell_char = pixel_nums_to_braille(None, Some(0));
    for row_idx in min(c1, c2)..c1 {
        *canvas.get_mut(col_idx, row_idx).unwrap() = char_to_cell(right_fill_cell_char);
    }

    // Now fill in partial height cells.
    if c1 == c2 {
        // top pixels are in the same cell
        *canvas.get_mut(col_idx, c1).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some((p1 % 4) as u8), Some((p2 % 4) as u8)));
    } else if c1 > c2 {
        // right pixel is in a higher cell.
        *canvas.get_mut(col_idx, c1).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some((p1 % 4) as u8), Some(0)));
        *canvas.get_mut(col_idx, c2).unwrap() = char_to_cell(
            pixel_nums_to_braille(None, Some((p2 % 4) as u8)));
    } else {
        // left pixel is in a higher cell.
        *canvas.get_mut(col_idx, c1).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some((p1 % 4) as u8), None));
        *canvas.get_mut(col_idx, c2).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some(0), Some((p2 % 4) as u8)));
    }
}

fn draw_spectrum<T: CellAccessor + HasSize>(canvas: &mut T, spec: &[Complex<f32>]) {
    canvas.clear(Cell::default());
    let (num_cols, num_rows) = canvas.size();
    let pixel_height = num_rows * 4;
    let pixel_width = num_cols * 2;
    // TODO what should this value be?
    let max_height = 500.0;

    let mut bins = vec![0.0; pixel_width];
    spectrum_to_bin_heights(spec, &mut bins[..]);

    for col_idx in 0..num_cols {
        // height in float between 0 and 1.
        let h1 = bins[col_idx * 2] / max_height;
        let h2 = bins[col_idx * 2 + 1] / max_height;

        // The "pixel" height of each point.
        let p1 = (h1 * pixel_height as f32).floor() as usize;
        let p2 = (h2 * pixel_height as f32).floor() as usize;

        draw_pixel_pair(canvas, col_idx, p1, p2);
    }
}

#[cfg(test)]
mod tests {
    use super::{pixel_nums_to_braille, fft_shift, draw_pixel_pair};
    use rustty::Terminal;

    #[test]
    fn test_fft_shift_dc() {
        let len = 9;
        let mut spec = vec![0; len];
        spec[0] = 1;
        fft_shift(&mut spec[..]);
        assert_eq!(spec[len / 2], 1);
    }

    #[test]
    fn test_fft_shift_even() {
        let mut before: Vec<usize> = (0..10).collect();
        let after = vec![5, 6, 7, 8, 9, 0, 1, 2, 3, 4];
        fft_shift(&mut before[..]);
        assert_eq!(before, after);
    }

    #[test]
    fn test_pixel_nums() {
        assert_eq!(pixel_nums_to_braille(Some(0), Some(0)), '⣿');
        assert_eq!(pixel_nums_to_braille(Some(1), Some(2)), '⣦');
        assert_eq!(pixel_nums_to_braille(None, Some(3)), '⢀');
        assert_eq!(pixel_nums_to_braille(Some(2), None), '⡄');
        assert_eq!(pixel_nums_to_braille(None, None), '⠀');
    }

    #[test]
    fn test_draw_pixel_pair() {
        let mut term = Terminal::new().unwrap();

        // Test drawing with the same top cell
        draw_pixel_pair(&mut term, 0, 4, 6);
        assert_eq!(term[(0, term.rows() - 3)].ch(), ' ');
        assert_eq!(term[(0, term.rows() - 2)].ch(), '⣰');
        assert_eq!(term[(0, term.rows() - 1)].ch(), '⣿');
        term.clear().unwrap();

        // Test drawing with the top pixel in each column being in
        // different cells
        draw_pixel_pair(&mut term, 0, 4, 8);
        assert_eq!(term[(0, term.rows() - 4)].ch(), ' ');
        assert_eq!(term[(0, term.rows() - 3)].ch(), '⢀');
        assert_eq!(term[(0, term.rows() - 2)].ch(), '⣸');
        assert_eq!(term[(0, term.rows() - 1)].ch(), '⣿');
        term.clear().unwrap();

        draw_pixel_pair(&mut term, 1, 13, 2);
        assert_eq!(term[(1, term.rows() - 5)].ch(), ' ');
        assert_eq!(term[(1, term.rows() - 4)].ch(), '⡄');
        assert_eq!(term[(1, term.rows() - 3)].ch(), '⡇');
        assert_eq!(term[(1, term.rows() - 2)].ch(), '⡇');
        assert_eq!(term[(1, term.rows() - 1)].ch(), '⣷');
        term.clear().unwrap();
    }
}
<commit_msg>Added resizing to windows.<commit_after>use std::char;
use std::cmp::{max, min};
use std::collections::VecDeque;
use num::{Complex, Float};
use rustty;
use rustty::{Attr, Color, Terminal, Cell, CellAccessor, HasSize};
use rustty::ui::{Alignable, Widget, VerticalAlign, HorizontalAlign};
use itertools::{Itertools, EitherOrBoth};

pub struct Canvas {
    term: Terminal,
    spectrum: Widget,
    waterfall: Widget,
    history: VecDeque<Vec<f32>>,
}

impl Canvas {
    pub fn new() -> Result<Self, rustty::Error> {
        let term = try!(Terminal::new());

        let mut canvas = Canvas {
            term: term,
            spectrum: Widget::new(0, 0),
            waterfall: Widget::new(0, 0),
            history: VecDeque::new(),
        };

        canvas.resize();

        Ok(canvas)
    }

    fn resize(&mut self) {
        let (cols, rows) = self.term.size();
        let spectrum_height = rows / 2;
        let waterfall_height = if rows % 2 == 0 { rows / 2 } else { rows / 2 + 1 };

        self.spectrum = Widget::new(cols, spectrum_height);
        self.spectrum.align(&self.term, HorizontalAlign::Middle, VerticalAlign::Top, 0);

        self.waterfall = Widget::new(cols, waterfall_height);
        self.waterfall.align(&self.term, HorizontalAlign::Middle, VerticalAlign::Bottom, 0);

        self.history.reserve(waterfall_height * 2);
    }

    fn check_and_resize(&mut self) {
        let (cols, rows) = self.term.size();
        let (spectrum_cols, spectrum_rows) = self.spectrum.size();
        let (waterfall_cols, waterfall_rows) = self.waterfall.size();
        // if the terminal size has changed...
        if cols != spectrum_cols || cols != waterfall_cols ||
            rows != (spectrum_rows + waterfall_rows) {
            self.resize();
        }
    }

    /// Adds a spectrum to the history and draws it on the waterfall
    /// and the spectrum view.
    pub fn add_spectrum(&mut self, spec: Vec<Complex<f32>>) {
        draw_spectrum(&mut self.spectrum, &spec);

        let normalized = normalize_spectrum(&spec, 50.0);
        let averaged = normalized.chunks(2).map(|v| (v[0] + v[1]) / 2.0).collect();
        self.history.push_front(averaged);

        draw_waterfall(&mut self.waterfall, &self.history);

        self.spectrum.draw_into(&mut self.term);
        self.waterfall.draw_into(&mut self.term);
        self.term.swap_buffers().unwrap();

        self.check_and_resize();
    }

    pub fn get_term(&mut self) -> &mut Terminal {
        &mut self.term
    }
}

fn draw_waterfall<T: CellAccessor + HasSize>(canvas: &mut T, spectra: &VecDeque<Vec<f32>>) {
    let (cols, rows) = canvas.size();
    for (row, mut specs) in (0..rows).zip(&spectra.iter().chunks_lazy(2)) {
        let upper_heights = specs.next().into_iter().flat_map(|x| x);
        let lower_heights = specs.next().into_iter().flat_map(|x| x);
        for (c, heights) in (0..cols).zip(upper_heights.zip_longest(lower_heights)) {
            let (u, l) = match heights {
                EitherOrBoth::Both(&upper, &lower) => (upper, lower),
                EitherOrBoth::Left(&upper) => (upper, 0.0),
                EitherOrBoth::Right(&lower) => (0.0, lower),
            };
            *canvas.get_mut(c, row).unwrap() = spectrum_heights_to_waterfall_cell(u, l);
        }
    }
}

fn spectrum_heights_to_waterfall_cell(upper: f32, lower: f32) -> Cell {
    Cell::new('▀',
              Color::Byte(color_mapping(upper)),
              Color::Byte(color_mapping(lower)),
              Attr::Default)
}

/// Assumes `f` is between 0 and 1. Anything outside of this range
/// will be clamped.
fn color_mapping(f: f32) -> u8 {
    let mapping = [16, 17, 18, 19, 21, 27, 33, 39, 45, 51,
                   50, 49, 48, 47, 46, 82, 118, 154, 190, 226];
    let idx = (f * (mapping.len() as f32)) as i32;
    if idx < 0 {
        mapping[0]
    } else if idx >= mapping.len() as i32 {
        mapping[mapping.len() - 1]
    } else {
        mapping[idx as usize]
    }
    //let lower = 232.0;
    //let upper = 255.0;
    //let mapped = f * (upper - lower) + lower;
    //if mapped < lower {
        //lower as u8
    //} else if mapped > upper {
        //upper as u8
    //} else {
        //mapped as u8
    //}
}

fn normalize_spectrum(spec: &[Complex<f32>], max_db: f32) -> Vec<f32> {
    // FFT shift
    let (first_half, last_half) = spec.split_at((spec.len() + 1) / 2);
    let shifted_spec = last_half.iter().chain(first_half.iter());

    // normalize and take the log
    shifted_spec.map(Complex::norm)
                .map(Float::log10)
                .map(|x| 10.0 * x)
                .map(|x| x / max_db)
                .collect()
}

// indexing is from the top of the cell
fn pixel_nums_to_braille(p1: Option<u8>, p2: Option<u8>) -> char {
    let pixel_map = [[0x01, 0x08],
                     [0x02, 0x10],
                     [0x04, 0x20],
                     [0x40, 0x80]];

    let mut c = 0;
    if let Some(p) = p1 {
        for i in p..4 {
            c |= pixel_map[i as usize][0];
        }
    }

    if let Some(p) = p2 {
        for i in p..4 {
            c |= pixel_map[i as usize][1];
        }
    }

    char::from_u32((0x2800 + c) as u32).unwrap()
}

fn fft_shift<T: Clone>(spec: &mut [T]) {
    let spec_copy = spec.to_owned();

    let (first_half, last_half) = spec_copy.split_at((spec_copy.len() + 1) / 2);
    let shifted_spec = last_half.iter().chain(first_half.iter());
    for (x, y) in spec.iter_mut().zip(shifted_spec) {
        *x = y.clone();
    }
}

fn spectrum_to_bin_heights(spec: &[Complex<f32>], dest: &mut [f32]) {
    //TODO should be plotting in log scale

    // subsample
    let mut last_idx = -1isize;
    for (i, x) in spec.iter().map(Complex::norm).enumerate() {
        if (i * dest.len() / spec.len()) as isize > last_idx {
            last_idx += 1;
            dest[last_idx as usize] = x;
        }
    }

    //TODO unnecessary allocation
    fft_shift(dest);
}

fn char_to_cell(c: char) -> Cell {
    let mut cell = Cell::with_char(c);
    cell.set_attrs(Attr::Bold);
    cell
}

fn draw_pixel_pair<T>(canvas: &mut T, col_idx: usize, p1: usize, p2: usize)
    where T: CellAccessor + HasSize
{
    let (_, rows) = canvas.size();
    let max_pixel_height = 4 * rows;

    // clamp heights
    let p1 = if p1 >= max_pixel_height { max_pixel_height - 1} else { p1 };
    let p2 = if p2 >= max_pixel_height { max_pixel_height - 1} else { p2 };

    // Reverse it, since the terminal indexing is from the top
    let p1 = max_pixel_height - p1 - 1;
    let p2 = max_pixel_height - p2 - 1;

    // cell indices
    let c1 = p1 / 4;
    let c2 = p2 / 4;

    // Fill in full height cells.
    let full_cell_char = pixel_nums_to_braille(Some(0), Some(0));
    for row_idx in max(c1, c2)..rows {
        *canvas.get_mut(col_idx, row_idx).unwrap() = char_to_cell(full_cell_char);
    }

    let left_fill_cell_char = pixel_nums_to_braille(Some(0), None);
    for row_idx in min(c1, c2)..c2 {
        *canvas.get_mut(col_idx, row_idx).unwrap() = char_to_cell(left_fill_cell_char);
    }

    let right_fill_cell_char = pixel_nums_to_braille(None, Some(0));
    for row_idx in min(c1, c2)..c1 {
        *canvas.get_mut(col_idx, row_idx).unwrap() = char_to_cell(right_fill_cell_char);
    }

    // Now fill in partial height cells.
    if c1 == c2 {
        // top pixels are in the same cell
        *canvas.get_mut(col_idx, c1).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some((p1 % 4) as u8), Some((p2 % 4) as u8)));
    } else if c1 > c2 {
        // right pixel is in a higher cell.
        *canvas.get_mut(col_idx, c1).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some((p1 % 4) as u8), Some(0)));
        *canvas.get_mut(col_idx, c2).unwrap() = char_to_cell(
            pixel_nums_to_braille(None, Some((p2 % 4) as u8)));
    } else {
        // left pixel is in a higher cell.
        *canvas.get_mut(col_idx, c1).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some((p1 % 4) as u8), None));
        *canvas.get_mut(col_idx, c2).unwrap() = char_to_cell(
            pixel_nums_to_braille(Some(0), Some((p2 % 4) as u8)));
    }
}

fn draw_spectrum<T: CellAccessor + HasSize>(canvas: &mut T, spec: &[Complex<f32>]) {
    canvas.clear(Cell::default());
    let (num_cols, num_rows) = canvas.size();
    let pixel_height = num_rows * 4;
    let pixel_width = num_cols * 2;
    // TODO what should this value be?
    let max_height = 500.0;

    let mut bins = vec![0.0; pixel_width];
    spectrum_to_bin_heights(spec, &mut bins[..]);

    for col_idx in 0..num_cols {
        // height in float between 0 and 1.
        let h1 = bins[col_idx * 2] / max_height;
        let h2 = bins[col_idx * 2 + 1] / max_height;

        // The "pixel" height of each point.
        let p1 = (h1 * pixel_height as f32).floor() as usize;
        let p2 = (h2 * pixel_height as f32).floor() as usize;

        draw_pixel_pair(canvas, col_idx, p1, p2);
    }
}

#[cfg(test)]
mod tests {
    use super::{pixel_nums_to_braille, fft_shift, draw_pixel_pair};
    use rustty::Terminal;

    #[test]
    fn test_fft_shift_dc() {
        let len = 9;
        let mut spec = vec![0; len];
        spec[0] = 1;
        fft_shift(&mut spec[..]);
        assert_eq!(spec[len / 2], 1);
    }

    #[test]
    fn test_fft_shift_even() {
        let mut before: Vec<usize> = (0..10).collect();
        let after = vec![5, 6, 7, 8, 9, 0, 1, 2, 3, 4];
        fft_shift(&mut before[..]);
        assert_eq!(before, after);
    }

    #[test]
    fn test_pixel_nums() {
        assert_eq!(pixel_nums_to_braille(Some(0), Some(0)), '⣿');
        assert_eq!(pixel_nums_to_braille(Some(1), Some(2)), '⣦');
        assert_eq!(pixel_nums_to_braille(None, Some(3)), '⢀');
        assert_eq!(pixel_nums_to_braille(Some(2), None), '⡄');
        assert_eq!(pixel_nums_to_braille(None, None), '⠀');
    }

    #[test]
    fn test_draw_pixel_pair() {
        let mut term = Terminal::new().unwrap();

        // Test drawing with the same top cell
        draw_pixel_pair(&mut term, 0, 4, 6);
        assert_eq!(term[(0, term.rows() - 3)].ch(), ' ');
        assert_eq!(term[(0, term.rows() - 2)].ch(), '⣰');
        assert_eq!(term[(0, term.rows() - 1)].ch(), '⣿');
        term.clear().unwrap();

        // Test drawing with the top pixel in each column being in
        // different cells
        draw_pixel_pair(&mut term, 0, 4, 8);
        assert_eq!(term[(0, term.rows() - 4)].ch(), ' ');
        assert_eq!(term[(0, term.rows() - 3)].ch(), '⢀');
        assert_eq!(term[(0, term.rows() - 2)].ch(), '⣸');
        assert_eq!(term[(0, term.rows() - 1)].ch(), '⣿');
        term.clear().unwrap();

        draw_pixel_pair(&mut term, 1, 13, 2);
        assert_eq!(term[(1, term.rows() - 5)].ch(), ' ');
        assert_eq!(term[(1, term.rows() - 4)].ch(), '⡄');
        assert_eq!(term[(1, term.rows() - 3)].ch(), '⡇');
        assert_eq!(term[(1, term.rows() - 2)].ch(), '⡇');
        assert_eq!(term[(1, term.rows() - 1)].ch(), '⣷');
        term.clear().unwrap();
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::prelude::*;
use std::collections::HashMap;

pub fn main() {
    let value: HashMap<char, usize> = (b'A'..b'Z' + 1)
                                          .enumerate()
                                          .map(|(i, c)| (c as char, i + 1))
                                          .collect();

    println!("Pipe the puzzle input to stdin");
    let stdin = io::stdin();
    for line in stdin.lock().lines() {
        let line = line.unwrap().replace("\"", "");
        let mut names: Vec<&str> = line.split(",").collect();
        names.sort();
        println!("{}",
                 names.iter()
                      .enumerate()
                      .map(|(i, name)| {
                          (i + 1) *
                          name.chars()
                              .map(|c| value.get(&c).unwrap())
                              .fold(0, |acc, i| acc + i)
                      })
                      .fold(0, |acc, i| acc + i));
    }
}
<commit_msg>use char instead of 1-length String<commit_after>use std::io;
use std::io::prelude::*;
use std::collections::HashMap;

pub fn main() {
    let value: HashMap<char, usize> = (b'A'..b'Z' + 1)
                                          .enumerate()
                                          .map(|(i, c)| (c as char, i + 1))
                                          .collect();

    println!("Pipe the puzzle input to stdin");
    let stdin = io::stdin();
    for line in stdin.lock().lines() {
        let line = line.unwrap().replace("\"", "");
        let mut names: Vec<&str> = line.split(',').collect();
        names.sort();
        println!("{}",
                 names.iter()
                      .enumerate()
                      .map(|(i, name)| {
                          (i + 1) *
                          name.chars()
                              .map(|c| value.get(&c).unwrap())
                              .fold(0, |acc, i| acc + i)
                      })
                      .fold(0, |acc, i| acc + i));
    }
}
<|endoftext|>


<commit_before>use std::char;

pub fn main() {
    let primes = (2..).filter(|&n| (2..(n as f32).sqrt() as u32 + 1).all(|i| n % i != 0));
    let mut max = 0;
    for p in primes.take_while(|&p| p <= 7654321) {
        let s = p.to_string();
        if (1..s.len() + 1)
               .all(|n| s.chars().any(|p| p == char::from_digit(n as u32, 10).unwrap())) {
            if p > max {
                max = p;
            }
        }
    }
    println!("{}", max);
}
<commit_msg>euler 41: solve without mutable variables<commit_after>use std::char;

pub fn main() {
    let primes = (2..).filter(|&n| (2..(n as f32).sqrt() as u32 + 1).all(|i| n % i != 0));
    let pandigital = |number: &u32| {
        let s = number.to_string();
        (1..s.len() + 1).all(|n| s.chars().any(|p| p == char::from_digit(n as u32, 10).unwrap()))
    };
    let ans = primes.take_while(|&p| p <= 7654321).filter(pandigital).last().unwrap();
    println!("{}", ans);
}
<|endoftext|>


<commit_before>//! Unix-specific utilities for working with the filesystem.

use std::ffi::{OsStr, OsString};
use std::path::Path;

use log::warn;

/// An extended attribute of a file.
///
/// On Linux, extended attributes can be obtained using `getfattr` and set with
/// `setfattr` utilities. On macOS, they can be manipulated through `getxattr`
/// and `setxattr` utilities.
///
/// See the [Wikipedia] article for more details.
///
/// [Wikipedia]: https://en.wikipedia.org/wiki/Extended_file_attributes
pub struct ExtAttr {
    /// A name of the extended attribute.
    pub key: OsString,
    /// A value of the extended attribute.
    pub value: Option<Vec<u8>>,
}

/// Returns an iterator over extended attributes of the specified file.
///
/// # Errors
///
/// The function will fail if a list of extended attributes of the file cannot
/// be obtained (e.g. when the file doesn't exist). However, all errors when
/// that can occur when inspecting values for particular attribute are logged
/// and forgotten.
///
/// # Examples
///
/// ```no_run
/// let attrs = rrg::fs::unix::ext_attrs(&"/tmp/foo")
///     .expect("failed to collect extended attributes");
///
/// for attr in attrs {
///     let key = attr.key.to_string_lossy();
///     match attr.value {
///         Some(value) => {
///             let value = std::str::from_utf8(&value)
///                 .expect("failed to convert attribute value");
///
///             println!("{}: {}", key, value);
///         },
///         None => println!("{}", key),
///     }
/// }
/// ```
pub fn ext_attrs<'p, P>(path: &'p P) -> std::io::Result<ExtAttrs<'p>>
where
    P: AsRef<Path>,
{
    let iter = xattr::list(&path)?;

    Ok(ExtAttrs {
        path: path.as_ref(),
        iter: iter,
    })
}

/// Iterator over extended attributes of a file.
///
/// Note that this iterator always returns an attribute. All errors that can
/// occur when obtaining values for particular attributes are swallowed.
///
/// The iterator can be constructed with the [`ext_attrs`] function.
///
/// [`ext_attrs`]: fn.ext_attrs.html
pub struct ExtAttrs<'p> {
    path: &'p Path,
    iter: xattr::XAttrs,
}

impl<'p> Iterator for ExtAttrs<'p> {

    type Item = ExtAttr;

    fn next(&mut self) -> Option<ExtAttr> {
        for key in &mut self.iter {
            let value = match ext_attr_value(self.path, &key) {
                Ok(value) => value,
                Err(()) => continue,
            };

            return Some(ExtAttr {
                key: key,
                value: value,
            });
        }

        None
    }
}

/// Collects value of an extended attribute at the given `key`.
///
/// This is a tiny wrapper around `xattr::get`, but logs and forgets the error
/// (if occurs).
fn ext_attr_value<P>(path: P, key: &OsStr) -> Result<Option<Vec<u8>>, ()>
where
    P: AsRef<Path>,
{
    xattr::get(&path, key).map_err(|error| {
        warn! {
            "failed to collect {key:?} of '{path}': {cause}",
            key = key,
            path = path.as_ref().display(),
            cause = error,
        };
    })
}
<commit_msg>Change nomenclature used for attribute names.<commit_after>//! Unix-specific utilities for working with the filesystem.

use std::ffi::{OsStr, OsString};
use std::path::Path;

use log::warn;

/// An extended attribute of a file.
///
/// On Linux, extended attributes can be obtained using `getfattr` and set with
/// `setfattr` utilities. On macOS, they can be manipulated through `getxattr`
/// and `setxattr` utilities.
///
/// See the [Wikipedia] article for more details.
///
/// [Wikipedia]: https://en.wikipedia.org/wiki/Extended_file_attributes
pub struct ExtAttr {
    /// A name of the extended attribute.
    pub name: OsString,
    /// A value of the extended attribute.
    pub value: Option<Vec<u8>>,
}

/// Returns an iterator over extended attributes of the specified file.
///
/// # Errors
///
/// The function will fail if a list of extended attributes of the file cannot
/// be obtained (e.g. when the file doesn't exist). However, all errors when
/// that can occur when inspecting values for particular attribute are logged
/// and forgotten.
///
/// # Examples
///
/// ```no_run
/// let attrs = rrg::fs::unix::ext_attrs(&"/tmp/foo")
///     .expect("failed to collect extended attributes");
///
/// for attr in attrs {
///     let name = attr.name.to_string_lossy();
///     match attr.value {
///         Some(value) => {
///             let value = std::str::from_utf8(&value)
///                 .expect("failed to convert attribute value");
///
///             println!("{}: {}", name, value);
///         },
///         None => println!("{}", name),
///     }
/// }
/// ```
pub fn ext_attrs<'p, P>(path: &'p P) -> std::io::Result<ExtAttrs<'p>>
where
    P: AsRef<Path>,
{
    let iter = xattr::list(&path)?;

    Ok(ExtAttrs {
        path: path.as_ref(),
        iter: iter,
    })
}

/// Iterator over extended attributes of a file.
///
/// Note that this iterator always returns an attribute. All errors that can
/// occur when obtaining values for particular attributes are swallowed.
///
/// The iterator can be constructed with the [`ext_attrs`] function.
///
/// [`ext_attrs`]: fn.ext_attrs.html
pub struct ExtAttrs<'p> {
    path: &'p Path,
    iter: xattr::XAttrs,
}

impl<'p> Iterator for ExtAttrs<'p> {

    type Item = ExtAttr;

    fn next(&mut self) -> Option<ExtAttr> {
        for name in &mut self.iter {
            let value = match ext_attr_value(self.path, &name) {
                Ok(value) => value,
                Err(()) => continue,
            };

            return Some(ExtAttr {
                name: name,
                value: value,
            });
        }

        None
    }
}

/// Collects value of an extended attribute with the specified name.
///
/// This is a tiny wrapper around `xattr::get`, but logs and forgets the error
/// (if occurs).
fn ext_attr_value<P>(path: P, name: &OsStr) -> Result<Option<Vec<u8>>, ()>
where
    P: AsRef<Path>,
{
    xattr::get(&path, name).map_err(|error| {
        warn! {
            "failed to collect attribute '{:?}' of '{path}': {cause}",
            name = name,
            path = path.as_ref().display(),
            cause = error,
        };
    })
}
<|endoftext|>


<commit_before>//! Unix-specific utilities for working with the filesystem.

use std::ffi::{OsStr, OsString};
use std::path::Path;

use log::warn;

/// An extended attribute of a file.
///
/// On Linux, extended attributes can be obtained using `getfattr` and set with
/// `setfattr` utilities. On macOS, they can be manipulated through `getxattr`
/// and `setxattr` utilities.
///
/// See the [Wikipedia] article for more details.
///
/// [Wikipedia]: https://en.wikipedia.org/wiki/Extended_file_attributes
pub struct ExtAttr {
    /// A name of the extended attribute.
    pub name: OsString,
    /// A value of the extended attribute.
    pub value: Option<OsString>,
}

/// Returns an iterator over extended attributes of the specified file.
///
/// # Errors
///
/// The function will fail if a list of extended attributes of the file cannot
/// be obtained (e.g. when the file doesn't exist). However, all errors when
/// that can occur when inspecting values for particular attribute are logged
/// and forgotten.
///
/// # Examples
///
/// ```no_run
/// let attrs = rrg::fs::unix::ext_attrs(&"/tmp/foo")
///     .expect("failed to collect extended attributes");
///
/// for attr in attrs {
///     let name = attr.name.to_string_lossy();
///     match attr.value {
///         Some(value) => println!("{}: {}", name, value.to_string_lossy()),
///         None => println!("{}", name),
///     }
/// }
/// ```
pub fn ext_attrs<'p, P>(path: &'p P) -> std::io::Result<ExtAttrs<'p>>
where
    P: AsRef<Path>,
{
    let iter = xattr::list(&path)?;

    Ok(ExtAttrs {
        path: path.as_ref(),
        iter: iter,
    })
}

/// Iterator over extended attributes of a file.
///
/// Note that this iterator always returns an attribute. All errors that can
/// occur when obtaining values for particular attributes are swallowed.
///
/// The iterator can be constructed with the [`ext_attrs`] function.
///
/// [`ext_attrs`]: fn.ext_attrs.html
pub struct ExtAttrs<'p> {
    path: &'p Path,
    iter: xattr::XAttrs,
}

impl<'p> Iterator for ExtAttrs<'p> {

    type Item = ExtAttr;

    fn next(&mut self) -> Option<ExtAttr> {
        for name in &mut self.iter {
            let value = match ext_attr_value(self.path, &name) {
                Ok(value) => value,
                Err(()) => continue,
            };

            return Some(ExtAttr {
                name: name,
                value: value,
            });
        }

        None
    }
}

/// Collects value of an extended attribute with the specified name.
///
/// This is a tiny wrapper around `xattr::get`, but logs and forgets the error
/// (if occurs).
fn ext_attr_value<P>(path: P, name: &OsStr) -> Result<Option<OsString>, ()>
where
    P: AsRef<Path>,
{
    match xattr::get(&path, name) {
        Ok(value) => Ok(value.map(std::os::unix::ffi::OsStringExt::from_vec)),
        Err(error) => Err(warn! {
            "failed to collect attribute '{:?}' of '{path}': {cause}",
            name = name,
            path = path.as_ref().display(),
            cause = error,
        }),
    }
}
<commit_msg>Add missing license header.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Unix-specific utilities for working with the filesystem.

use std::ffi::{OsStr, OsString};
use std::path::Path;

use log::warn;

/// An extended attribute of a file.
///
/// On Linux, extended attributes can be obtained using `getfattr` and set with
/// `setfattr` utilities. On macOS, they can be manipulated through `getxattr`
/// and `setxattr` utilities.
///
/// See the [Wikipedia] article for more details.
///
/// [Wikipedia]: https://en.wikipedia.org/wiki/Extended_file_attributes
pub struct ExtAttr {
    /// A name of the extended attribute.
    pub name: OsString,
    /// A value of the extended attribute.
    pub value: Option<OsString>,
}

/// Returns an iterator over extended attributes of the specified file.
///
/// # Errors
///
/// The function will fail if a list of extended attributes of the file cannot
/// be obtained (e.g. when the file doesn't exist). However, all errors when
/// that can occur when inspecting values for particular attribute are logged
/// and forgotten.
///
/// # Examples
///
/// ```no_run
/// let attrs = rrg::fs::unix::ext_attrs(&"/tmp/foo")
///     .expect("failed to collect extended attributes");
///
/// for attr in attrs {
///     let name = attr.name.to_string_lossy();
///     match attr.value {
///         Some(value) => println!("{}: {}", name, value.to_string_lossy()),
///         None => println!("{}", name),
///     }
/// }
/// ```
pub fn ext_attrs<'p, P>(path: &'p P) -> std::io::Result<ExtAttrs<'p>>
where
    P: AsRef<Path>,
{
    let iter = xattr::list(&path)?;

    Ok(ExtAttrs {
        path: path.as_ref(),
        iter: iter,
    })
}

/// Iterator over extended attributes of a file.
///
/// Note that this iterator always returns an attribute. All errors that can
/// occur when obtaining values for particular attributes are swallowed.
///
/// The iterator can be constructed with the [`ext_attrs`] function.
///
/// [`ext_attrs`]: fn.ext_attrs.html
pub struct ExtAttrs<'p> {
    path: &'p Path,
    iter: xattr::XAttrs,
}

impl<'p> Iterator for ExtAttrs<'p> {

    type Item = ExtAttr;

    fn next(&mut self) -> Option<ExtAttr> {
        for name in &mut self.iter {
            let value = match ext_attr_value(self.path, &name) {
                Ok(value) => value,
                Err(()) => continue,
            };

            return Some(ExtAttr {
                name: name,
                value: value,
            });
        }

        None
    }
}

/// Collects value of an extended attribute with the specified name.
///
/// This is a tiny wrapper around `xattr::get`, but logs and forgets the error
/// (if occurs).
fn ext_attr_value<P>(path: P, name: &OsStr) -> Result<Option<OsString>, ()>
where
    P: AsRef<Path>,
{
    match xattr::get(&path, name) {
        Ok(value) => Ok(value.map(std::os::unix::ffi::OsStringExt::from_vec)),
        Err(error) => Err(warn! {
            "failed to collect attribute '{:?}' of '{path}': {cause}",
            name = name,
            path = path.as_ref().display(),
            cause = error,
        }),
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::ops::Add;
use std::ops::Range;
use std::path::Path;
use std::process::Command;

use git;

#[derive(Debug, Clone)]
pub struct Entry<'a> { name: &'a str, time: &'a str, comment: &'a str, commit: &'a str, }

impl<'a> fmt::Display for Entry<'a> {
    fn fmt(& self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Entry: {:?} {:?} {:?}: {:?}", self.commit, self.time, self.name, self.comment)
    }
}

impl<'a> Entry<'a> {
    pub fn new<S:?Sized + AsRef<str>>(commit: &'a S, name: &'a S, time: &'a S, comment: &'a S) -> Entry<'a> {
        Entry { commit : commit.as_ref(), name: name.as_ref(), time: time.as_ref(), comment: comment.as_ref() }
    }

    pub fn name(&self) -> &str { self.name }
    pub fn time(&self) -> &str { self.time }
    pub fn comment(&self) -> &str { self.comment }
    pub fn commit(&self) -> &str { self.commit }
}

#[derive(Debug, Clone)]
pub struct History<'a> {
    cwd: &'a Path,
    entries_count: Option<i32>,
    page_size: i32,
}

impl<'a> History<'a> {
    const GIT_ONE_LINE_DETAILS : [&'static str;2] = ["log", "--pretty=format:%h|%an|%cr|%s"];

    pub fn new(size: i32, cwd: &Path) -> History {
        History { page_size: size, cwd : cwd, entries_count: Self::__count__(cwd) }
    }

    pub fn entries_count(&self) -> Option<i32> { self.entries_count }

    pub fn page_size(&self) -> i32 { self.page_size }

    pub fn page_count(&self) -> Option<i32> {
        self.entries_count()
            .map(|e| e + self.page_size())
            .map(|e| e / self.page_size())
    }

    fn __count__(cwd: &Path) -> Option<i32> {
        let args = Self::GIT_ONE_LINE_DETAILS.to_vec();

        git::git_pipe(Command::new("wc").arg("-l"),
                      cwd, args.as_ref())
            .map(|s| s.trim().parse::<i32>().unwrap_or(0))
            .ok()
    }

    pub fn get_page(&self, page: i32) -> Option<Vec<Entry>> {
        self.get_page_data(page)
            .map(|strings| strings.into_iter()
                      .filter_map(|ref line| self.make_entry(line))
                      .inspect(|ref e| debug!("Inspect 1: {}", e))
                      .collect::<Vec<_>>()
                      .into_iter().inspect(|ref e| debug!("Inspect 2: {}", e))
                      .collect::<Vec<_>>()
                )
    }

    fn make_entry(&self, s: &'a String) -> Option<Entry> {
        let parts: Vec<_> = s.trim().split('|').collect();
        debug!("make_entry: {:?} -> {:?}", s, parts);

        if let [commit, name, time, comment] = &parts[..] {
            Some(Entry::new(commit, name, time, comment))
        } else {
            None
        }
    }

    fn get_page_data(&self, page: i32) -> Option<Vec<String>> {
        let mut args = Self::GIT_ONE_LINE_DETAILS
                        .iter()
                        .map(|e| String::from(*e))
                        .collect::<Vec<String>>();

        let arg1 = format!("--skip={}", self.page_size() * page);
        let arg2 = format!("--max-count={}", self.page_size());

        args.push(arg1);
        args.push(arg2);

        git::git_exec(self.cwd, args.as_ref())
            .ok()
            .map(|s| s.trim().lines()
                 .map(|e| e.trim().to_owned())
                 .filter(|e| !e.is_empty())
                 .collect::<Vec<String>>())
    }
}
<commit_msg>fix: Correct history records (memory corruption)<commit_after>use std::fmt;
use std::ops::Add;
use std::ops::Range;
use std::path::Path;
use std::process::Command;

use git;

#[derive(Debug, Clone)]
pub struct Entry { name: String, time: String, comment: String, commit: String, }

impl fmt::Display for Entry {
    fn fmt(& self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Entry: {:?} {:?} {:?}: {:?}", self.commit, self.time, self.name, self.comment)
    }
}

impl Entry {
    pub fn new(commit: &str, name: &str, time: &str, comment: &str) -> Entry {
        Entry { commit : commit.to_owned(), name: name.to_owned(), time: time.to_owned(), comment: comment.to_owned() }
    }

    pub fn name(&self) -> &str { &*self.name }
    pub fn time(&self) -> &str { &*self.time }
    pub fn comment(&self) -> &str { &*self.comment }
    pub fn commit(&self) -> &str { &*self.commit }
}

#[derive(Debug, Clone)]
pub struct History<'a> {
    cwd: &'a Path,
    entries_count: Option<i32>,
    page_size: i32,
}

impl<'a> History<'a> {
    const GIT_ONE_LINE_DETAILS : [&'static str;2] = ["log", "--pretty=format:%h|%an|%cr|%s"];

    pub fn new(size: i32, cwd: &Path) -> History {
        History { page_size: size, cwd : cwd, entries_count: Self::__count__(cwd) }
    }

    pub fn entries_count(&self) -> Option<i32> { self.entries_count }

    pub fn page_size(&self) -> i32 { self.page_size }

    pub fn page_count(&self) -> Option<i32> {
        self.entries_count()
            .map(|e| e + self.page_size())
            .map(|e| e / self.page_size())
    }

    fn __count__(cwd: &Path) -> Option<i32> {
        let args = Self::GIT_ONE_LINE_DETAILS.to_vec();

        git::git_pipe(Command::new("wc").arg("-l"),
                      cwd, args.as_ref())
            .map(|s| s.trim().parse::<i32>().unwrap_or(0))
            .ok()
    }

    pub fn get_page(&self, page: i32) -> Option<Vec<Entry>> {
        self.get_page_data(page)
            .map(|strings| strings.into_iter()
                    .filter_map(|line| Self::make_entry(line))
                    .collect::<Vec<_>>())
    }

    fn make_entry(line: String) -> Option<Entry> {
        let parts: Vec<_> = line.trim().split('|')
            .collect();
        debug!("make_entry: {:?} -> {:?}", line, parts);

        if let [commit, name, time, comment] = &parts[..] {
            Some(Entry::new(commit, name, time, comment))
        } else {
            None
        }
    }

    fn get_page_data(&self, page: i32) -> Option<Vec<String>> {
        let mut args = Self::GIT_ONE_LINE_DETAILS
                        .iter()
                        .map(|e| String::from(*e))
                        .collect::<Vec<String>>();

        let arg1 = format!("--skip={}", self.page_size() * page);
        let arg2 = format!("--max-count={}", self.page_size());

        args.push(arg1);
        args.push(arg2);

        git::git_exec(self.cwd, args.as_ref())
            .ok()
            .map(|s| s.trim().lines()
                 .map(|e| e.trim().to_owned())
                 .filter(|e| !e.is_empty())
                 .collect::<Vec<String>>())
    }
}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer, channels::Channels, connection_status::ConnectionState, frames::Frames,
    internal_rpc::InternalRPC, thread::ThreadHandle, waker::Waker, Configuration, ConnectionStatus,
    Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker as MioWaker};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    connection_io_loop_handle: ThreadHandle,
    waker: Waker,
    socket: T,
    status: Status,
    poll: Poll,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        connection_io_loop_handle: ThreadHandle,
        waker: Waker,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        waker.set_waker(MioWaker::new(poll.registry(), WAKER)?);
        let frame_size = std::cmp::max(8192, configuration.frame_max() as usize);
        let mut inner = Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            connection_io_loop_handle,
            waker,
            socket,
            status: Status::Initial,
            poll,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status == Status::Initial
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serialized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
        );
        waker.wake()
    }

    fn cancel_serialized_frames(&self, error: Error) -> Result<()> {
        for (_, resolver) in self.serialized_frames.iter() {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                // Due to a bug in epoll/mio, it doesn't seem like we can trust this, it's sometimes missing when it should be there
                /*
                if event.is_readable() {
                    self.can_read = true;
                }
                */
                self.can_read = true;
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        let res = self.do_run();
        self.internal_rpc.poll(&self.channels).and(res)
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            self.internal_rpc.poll(&self.channels)?;
            if self.connection_status.closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.internal_rpc.poll(&self.channels)?;
            if self.should_heartbeat() {
                self.channels.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection_status.errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.channels.set_connection_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            }
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer.read_from(&mut self.socket).map(|sz| {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            self.send_continue()
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.channels.set_connection_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.channels.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.channels.set_connection_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<commit_msg>IoLoop: flush socket after writes<commit_after>use crate::{
    buffer::Buffer, channels::Channels, connection_status::ConnectionState, frames::Frames,
    internal_rpc::InternalRPC, thread::ThreadHandle, waker::Waker, Configuration, ConnectionStatus,
    Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker as MioWaker};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    connection_io_loop_handle: ThreadHandle,
    waker: Waker,
    socket: T,
    status: Status,
    poll: Poll,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        connection_io_loop_handle: ThreadHandle,
        waker: Waker,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        waker.set_waker(MioWaker::new(poll.registry(), WAKER)?);
        let frame_size = std::cmp::max(8192, configuration.frame_max() as usize);
        let mut inner = Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            connection_io_loop_handle,
            waker,
            socket,
            status: Status::Initial,
            poll,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status == Status::Initial
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serialized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
        );
        waker.wake()
    }

    fn cancel_serialized_frames(&self, error: Error) -> Result<()> {
        for (_, resolver) in self.serialized_frames.iter() {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                // Due to a bug in epoll/mio, it doesn't seem like we can trust this, it's sometimes missing when it should be there
                /*
                if event.is_readable() {
                    self.can_read = true;
                }
                */
                self.can_read = true;
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        let res = self.do_run();
        self.internal_rpc.poll(&self.channels).and(res)
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            self.internal_rpc.poll(&self.channels)?;
            if self.connection_status.closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.internal_rpc.poll(&self.channels)?;
            if self.should_heartbeat() {
                self.channels.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection_status.errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.channels.set_connection_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            } else {
                self.socket.flush()?;
            }
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer.read_from(&mut self.socket).map(|sz| {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            self.send_continue()
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.channels.set_connection_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.channels.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.channels.set_connection_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<|endoftext|>