<commit_before>extern crate regex;

extern crate rusoto_core;
extern crate rusoto_sts;
extern crate rusoto_dynamodb;

pub mod SessionHandler {
    use std::fs::{File, OpenOptions};
    use std::io::BufReader;
    use std::io;
    use std::io::prelude::*;
    use std::io::{stdin, stdout};
    use std::env;
    use rusoto_core::{ProfileProvider, Region, default_tls_client};
    use rusoto_sts::{Sts, StsClient, GetSessionTokenRequest, GetSessionTokenResponse, GetSessionTokenError, Credentials};
    use std::collections::HashMap;

    struct AWSConfig {
        profiles: HashMap<String, AWSProfile>,
    }

    #[derive(Debug)]
    struct AWSProfile {
        source_profile: Option<String>,
        region: Option<String>,
        output: Option<String>,
        role_arn: Option<String>,
        mfa_serial: Option<String>,
        ot_session_name: Option<String>,
        ot_expiration: Option<String>,
        ot_source_profile: Option<String>,
        ot_role_arn: Option<String>,
    }

    impl AWSConfig {
        fn save(&self) -> Result<(), io::Error> {
            println!("Saving config to disk");
            let mut bluenine_config_path = env::home_dir().unwrap().display().to_string();
            bluenine_config_path.push_str("/.aws/bluenine_config");
            let mut file = OpenOptions::new()
                        .write(true)
                        .truncate(true)
                        .open(bluenine_config_path)
                        .unwrap();

            let mut prof = String::new();
            for (name, profile) in &self.profiles {
                prof.push_str(&format!("[profile {}]\n", name));

                let source_profile = &profile.source_profile;
                match source_profile {
                    &Some(ref source_profile) => {
                        prof.push_str(&format!("source_profile = {}\n", source_profile));
                    },
                    &None => {},
                };

                let region = &profile.region;
                match region {
                    &Some(ref region) => {
                        prof.push_str(&format!("region = {}\n", region));
                    },
                    &None => {},
                };

                let output = &profile.output;
                match output {
                    &Some(ref output) => {
                        prof.push_str(&format!("output = {}\n", output));
                    },
                    &None => {},
                };

                let role_arn = &profile.role_arn;
                match role_arn {
                    &Some(ref role_arn) => {
                        prof.push_str(&format!("role_arn = {}\n", role_arn));
                    },
                    &None => {},
                };

                let mfa_serial = &profile.mfa_serial;
                match mfa_serial {
                    &Some(ref mfa_serial) => {
                        prof.push_str(&format!("mfa_serial = {}\n", mfa_serial));
                    },
                    &None => {},
                };

                let ot_session_name = &profile.ot_session_name;
                match ot_session_name {
                    &Some(ref ot_session_name) => {
                        prof.push_str(&format!("ot_session_name = {}\n", ot_session_name));
                    },
                    &None => {},
                };

                let ot_expiration = &profile.ot_expiration;
                match ot_expiration {
                    &Some(ref ot_expiration) => {
                        prof.push_str(&format!("ot_expiration = {}\n", ot_expiration));
                    },
                    &None => {},
                };

                let ot_expiration = &profile.ot_expiration;
                match ot_expiration {
                    &Some(ref ot_expiration) => {
                        prof.push_str(&format!("ot_expiration = {}\n", ot_expiration));
                    },
                    &None => {},
                };

                let ot_source_profile = &profile.ot_source_profile;
                match ot_source_profile {
                    &Some(ref ot_source_profile) => {
                        prof.push_str(&format!("ot_source_profile = {}\n", ot_source_profile));
                    },
                    &None => {},
                };

                let ot_role_arn = &profile.ot_role_arn;
                match ot_role_arn {
                    &Some(ref ot_role_arn) => {
                        prof.push_str(&format!("ot_role_arn = {}\n", ot_role_arn));
                    },
                    &None => {},
                };

                prof.push_str("\n");
            }

            try!(file.write_all(prof.as_bytes()));
            Ok(())
        }

        fn get_profile(&self, name: &str) -> &AWSProfile {
            if (&self.profiles).contains_key(name) {
                &self.profiles[name]
            }
            else {
                panic!("Could not find profile {}", name);
            }
        }

        fn contains_profile(&self, name: &str) -> bool {
            if (&self.profiles).contains_key(name) {
                true
            }
            else {
                false
            }
        }
    }

    pub fn create(profile_name: &str) {
        println!("Loading config file");
        let aws_config = load_config();

        let mut session_name = profile_name.clone().to_string();
        session_name.push_str("-session");
        if aws_config.contains_profile(&session_name) {
            println!("Session {} for profile {} already exists", session_name, profile_name);
        }

        println!("Creating session for profile \"{}\"...", profile_name);
        let aws_profile = aws_config.get_profile(profile_name);

        let mut provider = ProfileProvider::new().unwrap();
        provider.set_profile(profile_name);
        // let region = &aws_profile.region;
        // match region {
        //     &Some(ref region) => region,
        //     &None => panic!("You must specify a region for profile {}", profile_name)
        // };
        // println!("Region: {:?}", region.as_ref().unwrap());
        let client = StsClient::new(default_tls_client().unwrap(), provider, Region::EuWest1);

        let mfa_serial = &aws_profile.mfa_serial;
        if mfa_serial.is_some() {
            print!("Enter AWS MFA code for profile [{}]: ", profile_name);
            stdout().flush();
            let mut token_code = String::new();
            stdin().read_line(&mut token_code)
                .ok()
                .expect("Couldn't read line");
            token_code.pop(); // Remove newline

            let request = GetSessionTokenRequest {
                duration_seconds: None,
                serial_number: Some(mfa_serial.as_ref().unwrap().to_owned()),
                token_code: Some(token_code),
            };
            let response = client.get_session_token(&request);
            match response {
                Ok(response) => {
                    match save_profile(profile_name, &aws_profile) {
                        Ok(_) => println!("Saved profile to config file."),
                        Err(err) => println!("Error saving profile config to file: {:?}", err)
                    };
                    let credentials = response.credentials.unwrap();
                    match save_credentials(profile_name, credentials) {
                        Ok(_) => println!("Saved to credentials file."),
                        Err(err) => println!("Error saving credentials to file: {:?}", err)
                    };
                },
                Err(err) => panic!("Failed to get session token for profile {}: {:?}", profile_name, err),
            };
        }
    }
    pub fn show(profile_name: &str) {
        println!("Showing config for profile {}...", profile_name);
    }

    pub fn refresh() {
        println!("Refreshing sessions...");
    }

    pub fn clean(profile_name: &str) {
        println!("Loading config file");
        let mut aws_config = load_config();
        let mut session_name = profile_name.clone().to_string();
        session_name.push_str("-session");

        if aws_config.contains_profile(&session_name) {
            println!("Cleaning session {}", session_name);
            aws_config.profiles.remove(&session_name);
            aws_config.save();
            // Remove credentials also
            remove_credentials(&session_name);
        }
        else {
            println!("No profiles to clean");
        }
    }

    fn load_config() -> AWSConfig {
        let aws_config_file = read_aws_config_file();
        let profiles = split_config_file(aws_config_file);

        // Create an instance of the AWSConfig struct
        let mut aws_config = AWSConfig {
            profiles: HashMap::new()
        };

        // Iterate over profile chunks
        for profile in profiles {
            // Create a vector of string that holds each line
            let split = profile.split("\n");
            let lines: Vec<String> = split.map(|s| s.to_string()).collect();

            // Get the profile name from the first line
            let profile_line = lines[0].to_owned();

            // Lets ignore the default for now
            if profile_line == "[default]" {
                continue;
            }

            let split2 = profile_line.split(" ");
            let mut words: Vec<String> = split2.map(|s| s.to_string()).collect();
            let mut profile_name = String::new();
            if words.len() > 1 {
                words[1].pop();
                profile_name = words[1].trim_right().to_owned();
            }

            let mut aws_profile = AWSProfile{
                source_profile: None,
                region: None,
                output: None,
                role_arn: None,
                mfa_serial: None,
                ot_expiration: None,
                ot_session_name: None,
                ot_source_profile: None,
                ot_role_arn: None,
            };

            for i in 1..lines.len() {
                let split = lines[i].split(" = ");
                let config: Vec<String> = split.map(|s| s.to_string()).collect();
                let key = &config[0];
                let value = config[1].to_owned();
                match key.as_ref() {
                    "mfa_serial" => aws_profile.mfa_serial = Some(value),
                    "output" => aws_profile.output = Some(value),
                    "region" => aws_profile.region = Some(value),
                    "role_arn" => aws_profile.role_arn = Some(value),
                    "source_profile" => aws_profile.source_profile = Some(value),
                    "ot_session_name" => aws_profile.ot_session_name = Some(value),
                    "ot_expiration" => aws_profile.ot_expiration = Some(value),
                    "ot_source_profile" => aws_profile.ot_source_profile = Some(value),
                    "ot_role_arn" => aws_profile.ot_role_arn = Some(value),
                    _ => (),
                }
            }

            if profile_name != "" {
                aws_config.profiles.insert(profile_name.clone(), aws_profile);
            }
        }

        aws_config
    }

    fn read_aws_config_file() -> String {
        let path = match env::home_dir() {
            Some(path) => path,
            None => panic!("Could not retrieve user's home directory."),
        };
        let config_file_path = format!("{}/.aws/bluenine_config", path.display());

        let f = File::open(config_file_path).expect("Could not find AWS config file.");
        let mut buf_reader = BufReader::new(f);
        let mut contents = String::new();
        buf_reader.read_to_string(&mut contents).expect("Found config file but could not read it.");

        contents
    }

    fn read_aws_credentials_file() -> String {
        let path = match env::home_dir() {
            Some(path) => path,
            None => panic!("Could not retrieve user's home directory."),
        };
        let config_file_path = format!("{}/.aws/bluenine_credentials", path.display());

        let f = File::open(config_file_path).expect("Could not find AWS credentials file.");
        let mut buf_reader = BufReader::new(f);
        let mut contents = String::new();
        buf_reader.read_to_string(&mut contents).expect("Found credentials file but could not read it.");

        contents
    }

    fn split_config_file(aws_config: String) -> Vec<String> {
        let split = aws_config.split("\n\n");
        let mut profiles: Vec<String> = split.map(|s| s.to_string()).collect();;

        profiles.pop(); // Remove last element as it is always empty
        profiles
    }

    fn save_credentials(profile_name: &str, credentials: Credentials) -> Result<(), io::Error> {
        let mut aws_credentials_path = env::home_dir().unwrap().display().to_string();
        aws_credentials_path.push_str("/.aws/bluenine_credentials");
        let mut file = OpenOptions::new()
                       .write(true)
                       .append(true)
                       .open(aws_credentials_path)
                       .unwrap();

        let mut creds = String::new();
        creds.push_str(&format!("[{}-session]\n", profile_name));
        creds.push_str(&format!("aws_access_key_id = {}\n", credentials.access_key_id));
        creds.push_str(&format!("expiration = {}\n", credentials.expiration));
        creds.push_str(&format!("aws_secret_access_key = {}\n", credentials.secret_access_key));
        creds.push_str(&format!("aws_session_token = {}\n\n", credentials.session_token));
        try!(file.write_all(creds.as_bytes()));
        Ok(())
    }

    fn remove_credentials(profile_name: &str) -> Result<(), io::Error> {
        println!("Saving credentials to disk");
        let aws_credentials_file = read_aws_credentials_file();
        let credentials = split_config_file(aws_credentials_file);
        let mut aws_credentials_path = env::home_dir().unwrap().display().to_string();
        aws_credentials_path.push_str("/.aws/bluenine_credentials");
        let mut file = OpenOptions::new()
                .write(true)
                .truncate(true)
                .open(aws_credentials_path)
                .unwrap();

        let mut creds = String::new();
        for credential in credentials {
            let split = credential.split("\n");
            let lines: Vec<String> = split.map(|s| s.to_string()).collect();

            // Get the profile name from the first line
            let profile_line = lines[0].to_owned();
            if profile_line == format!("[{}]", profile_name) {
                continue;
            }
            creds.push_str(&credential);
            creds.push_str("\n\n");
        }
        try!(file.write_all(creds.as_bytes()));
        Ok(())
    }

    fn save_profile(profile_name: &str, aws_profile: &AWSProfile) -> Result<(), io::Error> {
        let mut aws_config_path = env::home_dir().unwrap().display().to_string();
        aws_config_path.push_str("/.aws/bluenine_config");
        let mut file = OpenOptions::new()
                       .write(true)
                       .append(true)
                       .open(aws_config_path)
                       .unwrap();

        let region = &aws_profile.region;
        match region {
            &Some(ref region) => {
                let mut prof = String::new();
                prof.push_str(&format!("[profile {}-session]\n", profile_name));
                prof.push_str(&format!("region = {}\n\n", region));
                try!(file.write_all(prof.as_bytes()));
                Ok(())
            },
            &None => {
                let region = "us-east-1";
                let mut prof = String::new();
                prof.push_str(&format!("[profile {}-session]\n", profile_name));
                prof.push_str(&format!("region = {}\n\n", region));
                try!(file.write_all(prof.as_bytes()));
                Ok(())
            },
        }
    }
}
<commit_msg>Update config credentials path<commit_after>extern crate regex;

extern crate rusoto_core;
extern crate rusoto_sts;
extern crate rusoto_dynamodb;

pub mod SessionHandler {
    use std::fs::{File, OpenOptions};
    use std::io::BufReader;
    use std::io;
    use std::io::prelude::*;
    use std::io::{stdin, stdout};
    use std::env;
    use rusoto_core::{ProfileProvider, Region, default_tls_client};
    use rusoto_sts::{Sts, StsClient, GetSessionTokenRequest, GetSessionTokenResponse, GetSessionTokenError, Credentials};
    use std::collections::HashMap;

    struct AWSConfig {
        profiles: HashMap<String, AWSProfile>,
    }

    #[derive(Debug)]
    struct AWSProfile {
        source_profile: Option<String>,
        region: Option<String>,
        output: Option<String>,
        role_arn: Option<String>,
        mfa_serial: Option<String>,
        ot_session_name: Option<String>,
        ot_expiration: Option<String>,
        ot_source_profile: Option<String>,
        ot_role_arn: Option<String>,
    }

    impl AWSConfig {
        fn save(&self) -> Result<(), io::Error> {
            println!("Saving config to disk");
            let mut bluenine_config_path = env::home_dir().unwrap().display().to_string();
            bluenine_config_path.push_str("/.aws/config");
            let mut file = OpenOptions::new()
                        .write(true)
                        .truncate(true)
                        .open(bluenine_config_path)
                        .unwrap();

            let mut prof = String::new();
            for (name, profile) in &self.profiles {
                prof.push_str(&format!("[profile {}]\n", name));

                let source_profile = &profile.source_profile;
                match source_profile {
                    &Some(ref source_profile) => {
                        prof.push_str(&format!("source_profile = {}\n", source_profile));
                    },
                    &None => {},
                };

                let region = &profile.region;
                match region {
                    &Some(ref region) => {
                        prof.push_str(&format!("region = {}\n", region));
                    },
                    &None => {},
                };

                let output = &profile.output;
                match output {
                    &Some(ref output) => {
                        prof.push_str(&format!("output = {}\n", output));
                    },
                    &None => {},
                };

                let role_arn = &profile.role_arn;
                match role_arn {
                    &Some(ref role_arn) => {
                        prof.push_str(&format!("role_arn = {}\n", role_arn));
                    },
                    &None => {},
                };

                let mfa_serial = &profile.mfa_serial;
                match mfa_serial {
                    &Some(ref mfa_serial) => {
                        prof.push_str(&format!("mfa_serial = {}\n", mfa_serial));
                    },
                    &None => {},
                };

                let ot_session_name = &profile.ot_session_name;
                match ot_session_name {
                    &Some(ref ot_session_name) => {
                        prof.push_str(&format!("ot_session_name = {}\n", ot_session_name));
                    },
                    &None => {},
                };

                let ot_expiration = &profile.ot_expiration;
                match ot_expiration {
                    &Some(ref ot_expiration) => {
                        prof.push_str(&format!("ot_expiration = {}\n", ot_expiration));
                    },
                    &None => {},
                };

                let ot_expiration = &profile.ot_expiration;
                match ot_expiration {
                    &Some(ref ot_expiration) => {
                        prof.push_str(&format!("ot_expiration = {}\n", ot_expiration));
                    },
                    &None => {},
                };

                let ot_source_profile = &profile.ot_source_profile;
                match ot_source_profile {
                    &Some(ref ot_source_profile) => {
                        prof.push_str(&format!("ot_source_profile = {}\n", ot_source_profile));
                    },
                    &None => {},
                };

                let ot_role_arn = &profile.ot_role_arn;
                match ot_role_arn {
                    &Some(ref ot_role_arn) => {
                        prof.push_str(&format!("ot_role_arn = {}\n", ot_role_arn));
                    },
                    &None => {},
                };

                prof.push_str("\n");
            }

            try!(file.write_all(prof.as_bytes()));
            Ok(())
        }

        fn get_profile(&self, name: &str) -> &AWSProfile {
            if (&self.profiles).contains_key(name) {
                &self.profiles[name]
            }
            else {
                panic!("Could not find profile {}", name);
            }
        }

        fn contains_profile(&self, name: &str) -> bool {
            if (&self.profiles).contains_key(name) {
                true
            }
            else {
                false
            }
        }
    }

    pub fn create(profile_name: &str) {
        println!("Loading config file");
        let aws_config = load_config();

        let mut session_name = profile_name.clone().to_string();
        session_name.push_str("-session");
        if aws_config.contains_profile(&session_name) {
            println!("Session {} for profile {} already exists", session_name, profile_name);
        }

        println!("Creating session for profile \"{}\"...", profile_name);
        let aws_profile = aws_config.get_profile(profile_name);

        let mut provider = ProfileProvider::new().unwrap();
        provider.set_profile(profile_name);
        // let region = &aws_profile.region;
        // match region {
        //     &Some(ref region) => region,
        //     &None => panic!("You must specify a region for profile {}", profile_name)
        // };
        // println!("Region: {:?}", region.as_ref().unwrap());
        let client = StsClient::new(default_tls_client().unwrap(), provider, Region::EuWest1);

        let mfa_serial = &aws_profile.mfa_serial;
        if mfa_serial.is_some() {
            print!("Enter AWS MFA code for profile [{}]: ", profile_name);
            stdout().flush();
            let mut token_code = String::new();
            stdin().read_line(&mut token_code)
                .ok()
                .expect("Couldn't read line");
            token_code.pop(); // Remove newline

            let request = GetSessionTokenRequest {
                duration_seconds: None,
                serial_number: Some(mfa_serial.as_ref().unwrap().to_owned()),
                token_code: Some(token_code),
            };
            let response = client.get_session_token(&request);
            match response {
                Ok(response) => {
                    match save_profile(profile_name, &aws_profile) {
                        Ok(_) => println!("Saved profile to config file."),
                        Err(err) => println!("Error saving profile config to file: {:?}", err)
                    };
                    let credentials = response.credentials.unwrap();
                    match save_credentials(profile_name, credentials) {
                        Ok(_) => println!("Saved to credentials file."),
                        Err(err) => println!("Error saving credentials to file: {:?}", err)
                    };
                },
                Err(err) => panic!("Failed to get session token for profile {}: {:?}", profile_name, err),
            };
        }
    }
    pub fn show(profile_name: &str) {
        println!("Showing config for profile {}...", profile_name);
    }

    pub fn refresh() {
        println!("Refreshing sessions...");
    }

    pub fn clean(profile_name: &str) {
        println!("Loading config file");
        let mut aws_config = load_config();
        let mut session_name = profile_name.clone().to_string();
        session_name.push_str("-session");

        if aws_config.contains_profile(&session_name) {
            println!("Cleaning session {}", session_name);
            aws_config.profiles.remove(&session_name);
            aws_config.save();
            // Remove credentials also
            remove_credentials(&session_name);
        }
        else {
            println!("No profiles to clean");
        }
    }

    fn load_config() -> AWSConfig {
        let aws_config_file = read_aws_config_file();
        let profiles = split_config_file(aws_config_file);

        // Create an instance of the AWSConfig struct
        let mut aws_config = AWSConfig {
            profiles: HashMap::new()
        };

        // Iterate over profile chunks
        for profile in profiles {
            // Create a vector of string that holds each line
            let split = profile.split("\n");
            let lines: Vec<String> = split.map(|s| s.to_string()).collect();

            // Get the profile name from the first line
            let profile_line = lines[0].to_owned();

            // Lets ignore the default for now
            if profile_line == "[default]" {
                continue;
            }

            let split2 = profile_line.split(" ");
            let mut words: Vec<String> = split2.map(|s| s.to_string()).collect();
            let mut profile_name = String::new();
            if words.len() > 1 {
                words[1].pop();
                profile_name = words[1].trim_right().to_owned();
            }

            let mut aws_profile = AWSProfile{
                source_profile: None,
                region: None,
                output: None,
                role_arn: None,
                mfa_serial: None,
                ot_expiration: None,
                ot_session_name: None,
                ot_source_profile: None,
                ot_role_arn: None,
            };

            for i in 1..lines.len() {
                let split = lines[i].split(" = ");
                let config: Vec<String> = split.map(|s| s.to_string()).collect();
                let key = &config[0];
                let value = config[1].to_owned();
                match key.as_ref() {
                    "mfa_serial" => aws_profile.mfa_serial = Some(value),
                    "output" => aws_profile.output = Some(value),
                    "region" => aws_profile.region = Some(value),
                    "role_arn" => aws_profile.role_arn = Some(value),
                    "source_profile" => aws_profile.source_profile = Some(value),
                    "ot_session_name" => aws_profile.ot_session_name = Some(value),
                    "ot_expiration" => aws_profile.ot_expiration = Some(value),
                    "ot_source_profile" => aws_profile.ot_source_profile = Some(value),
                    "ot_role_arn" => aws_profile.ot_role_arn = Some(value),
                    _ => (),
                }
            }

            if profile_name != "" {
                aws_config.profiles.insert(profile_name.clone(), aws_profile);
            }
        }

        aws_config
    }

    fn read_aws_config_file() -> String {
        let path = match env::home_dir() {
            Some(path) => path,
            None => panic!("Could not retrieve user's home directory."),
        };
        let config_file_path = format!("{}/.aws/config", path.display());

        let f = File::open(config_file_path).expect("Could not find AWS config file.");
        let mut buf_reader = BufReader::new(f);
        let mut contents = String::new();
        buf_reader.read_to_string(&mut contents).expect("Found config file but could not read it.");

        contents
    }

    fn read_aws_credentials_file() -> String {
        let path = match env::home_dir() {
            Some(path) => path,
            None => panic!("Could not retrieve user's home directory."),
        };
        let config_file_path = format!("{}/.aws/credentials", path.display());

        let f = File::open(config_file_path).expect("Could not find AWS credentials file.");
        let mut buf_reader = BufReader::new(f);
        let mut contents = String::new();
        buf_reader.read_to_string(&mut contents).expect("Found credentials file but could not read it.");

        contents
    }

    fn split_config_file(aws_config: String) -> Vec<String> {
        let split = aws_config.split("\n\n");
        let mut profiles: Vec<String> = split.map(|s| s.to_string()).collect();;

        profiles.pop(); // Remove last element as it is always empty
        profiles
    }

    fn save_credentials(profile_name: &str, credentials: Credentials) -> Result<(), io::Error> {
        let mut aws_credentials_path = env::home_dir().unwrap().display().to_string();
        aws_credentials_path.push_str("/.aws/credentials");
        let mut file = OpenOptions::new()
                       .write(true)
                       .append(true)
                       .open(aws_credentials_path)
                       .unwrap();

        let mut creds = String::new();
        creds.push_str(&format!("[{}-session]\n", profile_name));
        creds.push_str(&format!("aws_access_key_id = {}\n", credentials.access_key_id));
        creds.push_str(&format!("expiration = {}\n", credentials.expiration));
        creds.push_str(&format!("aws_secret_access_key = {}\n", credentials.secret_access_key));
        creds.push_str(&format!("aws_session_token = {}\n\n", credentials.session_token));
        try!(file.write_all(creds.as_bytes()));
        Ok(())
    }

    fn remove_credentials(profile_name: &str) -> Result<(), io::Error> {
        println!("Saving credentials to disk");
        let aws_credentials_file = read_aws_credentials_file();
        let credentials = split_config_file(aws_credentials_file);
        let mut aws_credentials_path = env::home_dir().unwrap().display().to_string();
        aws_credentials_path.push_str("/.aws/credentials");
        let mut file = OpenOptions::new()
                .write(true)
                .truncate(true)
                .open(aws_credentials_path)
                .unwrap();

        let mut creds = String::new();
        for credential in credentials {
            let split = credential.split("\n");
            let lines: Vec<String> = split.map(|s| s.to_string()).collect();

            // Get the profile name from the first line
            let profile_line = lines[0].to_owned();
            if profile_line == format!("[{}]", profile_name) {
                continue;
            }
            creds.push_str(&credential);
            creds.push_str("\n\n");
        }
        try!(file.write_all(creds.as_bytes()));
        Ok(())
    }

    fn save_profile(profile_name: &str, aws_profile: &AWSProfile) -> Result<(), io::Error> {
        let mut aws_config_path = env::home_dir().unwrap().display().to_string();
        aws_config_path.push_str("/.aws/config");
        let mut file = OpenOptions::new()
                       .write(true)
                       .append(true)
                       .open(aws_config_path)
                       .unwrap();

        let region = &aws_profile.region;
        match region {
            &Some(ref region) => {
                let mut prof = String::new();
                prof.push_str(&format!("[profile {}-session]\n", profile_name));
                prof.push_str(&format!("region = {}\n\n", region));
                try!(file.write_all(prof.as_bytes()));
                Ok(())
            },
            &None => {
                let region = "us-east-1";
                let mut prof = String::new();
                prof.push_str(&format!("[profile {}-session]\n", profile_name));
                prof.push_str(&format!("region = {}\n\n", region));
                try!(file.write_all(prof.as_bytes()));
                Ok(())
            },
        }
    }
}
<|endoftext|>


<commit_before>use std::borrow::Borrow;
use std::collections::HashMap;
use std::hash::Hash;
use std::iter::Chain;
use std::collections::hash_map;

pub struct RehashingHashMap<K, V> {
    // NOTE: I tried to make an array of 2 elements, but run into borrowing problems
    hashmap1: HashMap<K, V>,
    hashmap2: HashMap<K, V>,
    is1main: bool,
    rehashing: bool,
}

impl<K, V> RehashingHashMap<K, V>
    where K: Eq + Hash + Clone
{
    pub fn new() -> RehashingHashMap<K, V> {
        RehashingHashMap {
            hashmap1: HashMap::new(),
            hashmap2: HashMap::new(),
            is1main: true,
            rehashing: false,
        }
    }

    pub fn with_capacity(capacity: usize) -> RehashingHashMap<K, V> {
        RehashingHashMap {
            hashmap1: HashMap::with_capacity(capacity),
            hashmap2: HashMap::new(),
            is1main: true,
            rehashing: false,
        }
    }

    fn get_main(&self) -> &HashMap<K, V> {
        if self.is1main { &self.hashmap1 } else { &self.hashmap2 }
    }

    fn get_mut_main(&mut self) -> &mut HashMap<K, V> {
        if self.is1main { &mut self.hashmap1 } else { &mut self.hashmap2 }
    }

    fn get_secondary(&self) -> &HashMap<K, V> {
        if self.is1main { &self.hashmap2 } else { &self.hashmap1 }
    }

    fn get_mut_secondary(&mut self) -> &mut HashMap<K, V> {
        if self.is1main { &mut self.hashmap2 } else { &mut self.hashmap1 }
    }

    pub fn rehash(&mut self) {
        if self.rehashing {
            if self.get_secondary().len() == 0 {
                self.drop_secondary();
                return;
            }
            let (mut main, mut sec) = if self.is1main {
                (&mut self.hashmap1, &mut self.hashmap2)
            } else {
                (&mut self.hashmap2, &mut self.hashmap1)
            };
            // unwrap is safe, checked len() > 0 already
            let k: K = sec.keys().take(1).next().unwrap().clone();
            // FIXME: I wish I did not have to clone they key
            // unwrap is safe, we know the key exists in the hashmap
            let val = sec.remove(&k).unwrap();
            main.insert(k, val);
        }
    }

    pub fn capacity(&self) -> usize {
        self.get_main().capacity()
    }

    pub fn reserve(&mut self, additional: usize) {
        self.rehash();
        self.get_mut_main().reserve(additional)
    }

    pub fn is_rehashing(&self) -> bool {
        if !self.rehashing {
            assert_eq!(self.get_secondary().len(), 0);
        }
        self.rehashing
    }

    pub fn shrink_to_fit(&mut self) {
        if !self.rehashing {
            self.rehashing = true;
            self.is1main = !self.is1main;
            let len = self.len();
            self.get_mut_main().reserve(len)
        }
    }

    pub fn len(&self) -> usize {
        self.get_main().len() + self.get_secondary().len()
    }

    pub fn is_empty(&self) -> bool {
        self.get_main().is_empty() && self.get_secondary().is_empty()
    }

    fn drop_secondary(&mut self) {
        self.rehashing = false;
        assert_eq!(self.get_secondary().len(), 0);
        if self.is1main {
            self.hashmap2 = HashMap::new();
        } else {
            self.hashmap1 = HashMap::new();
        }
    }

    fn assert_state(&self) {
        #![allow(dead_code)]
        if self.rehashing {
            assert!(self.get_secondary().capacity() > 0);
        } else {
            assert!(self.get_secondary().capacity() == 0);
        }
    }

    pub fn clear(&mut self) {
        self.get_mut_main().clear();
        self.drop_secondary();
    }

    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
        // while rehashing, they key can be in either hashmap1 or hashmap2
        // but we want to remove them from wherever it is and add it to main
        let mut ret = None;
        if self.rehashing || self.is1main {
            ret = self.hashmap1.remove(&k);
        }
        if ret.is_none() && (self.rehashing || !self.is1main) {
            ret = self.hashmap2.remove(&k);
        }
        self.get_mut_main().insert(k, v);
        self.rehash();
        ret
    }

    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
            where K: Borrow<Q>, Q: Hash + Eq {
        if self.rehashing {
            match self.get_main().get(k) {
                Some(ref v) => Some(v),
                None => self.get_secondary().get(k),
            }
        } else {
            self.get_main().get(k)
        }
    }

    pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
        where K: Borrow<Q>, Q: Hash + Eq {
        if self.rehashing {
            self.rehash();
            match self.get_mut_main().remove(k) {
                Some(v) => Some(v),
                None => self.get_mut_secondary().remove(k),
            }
        } else {
            self.get_mut_main().remove(k)
        }
    }

    pub fn iter(&self) -> Iter<K, V> {
        Iter { inner: self.hashmap1.iter().chain(self.hashmap2.iter()) }
    }

    pub fn keys(&self) -> Keys<K, V> {
        Keys { inner: self.hashmap1.keys().chain(self.hashmap2.keys()) }
    }

    pub fn values(&self) -> Values<K, V> {
        Values { inner: self.hashmap1.values().chain(self.hashmap2.values()) }
    }
}

#[derive(Clone)]
pub struct Iter<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Iter<'a, K, V>, hash_map::Iter<'a, K, V>>
}

impl<'a, K, V> Iterator for Iter<'a, K, V> {
    type Item = (&'a K, &'a V);

    #[inline] fn next(&mut self) -> Option<(&'a K, &'a V)> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

#[derive(Clone)]
pub struct Keys<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Keys<'a, K, V>, hash_map::Keys<'a, K, V>>
}

impl<'a, K, V> Iterator for Keys<'a, K, V> {
    type Item = &'a K;

    #[inline] fn next(&mut self) -> Option<&'a K> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

#[derive(Clone)]
pub struct Values<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Values<'a, K, V>, hash_map::Values<'a, K, V>>
}

impl<'a, K, V> Iterator for Values<'a, K, V> {
    type Item = &'a V;

    #[inline] fn next(&mut self) -> Option<&'a V> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

#[test]
fn capacity() {
    let mut hash:RehashingHashMap<u8, u8> = RehashingHashMap::with_capacity(20);
    assert!(hash.capacity() >= 20);
    hash.reserve(40);
    assert!(hash.capacity() >= 40);
}

#[test]
fn insert() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value1 = 2;
    let value2 = 3;

    assert_eq!(hash.insert(key.clone(), value1.clone()), None);
    assert_eq!(hash.insert(key.clone(), value2.clone()), Some(value1.clone()));
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert_eq!(hash.insert(key.clone(), value1.clone()), Some(value2.clone()));
    assert!(!hash.is_rehashing());
    hash.assert_state();
}

#[test]
fn insert_many_rehash_get() {
    let mut hash = RehashingHashMap::new();

    let len = 1000;

    for i in 0..len {
        hash.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2){
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    assert_eq!(hash.len(), len);

    for i in 0..len {
        assert_eq!(hash.get(&i).unwrap(), &i);
    }
    for i in len..(len * 2) {
        assert!(hash.get(&i).is_none());
    }

    for _ in 0..(len / 2 + 1){
        hash.rehash();
    }
    assert!(!hash.is_rehashing());
    hash.assert_state();

    assert_eq!(hash.len(), len);

    for i in 0..len {
        assert_eq!(hash.get(&i).unwrap(), &i);
    }
    for i in len..(len * 2) {
        assert!(hash.get(&i).is_none());
    }
}

#[test]
fn is_empty() {
    let mut hash = RehashingHashMap::new();
    assert!(hash.is_empty());

    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    assert!(!hash.is_empty());
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert!(!hash.is_empty());
    hash.rehash();
    hash.rehash();
    assert!(!hash.is_rehashing());
    assert!(!hash.is_empty());
}

#[test]
fn clear() {
    let mut hash = RehashingHashMap::with_capacity(1000);
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.clear();
    hash.assert_state();

    assert!(hash.capacity() >= 1000);
}

#[test]
fn remove0() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn remove1() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    hash.rehash();
    assert!(hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn remove2() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    hash.rehash();
    hash.rehash();
    assert!(!hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn iterator() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    for (_, i) in hash.iter() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}

#[test]
fn keys() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    for i in hash.keys() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}

#[test]
fn values() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    for i in hash.values() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}
<commit_msg>Iter, keys and values implement exact size<commit_after>use std::borrow::Borrow;
use std::collections::HashMap;
use std::hash::Hash;
use std::iter::Chain;
use std::collections::hash_map;

pub struct RehashingHashMap<K, V> {
    // NOTE: I tried to make an array of 2 elements, but run into borrowing problems
    hashmap1: HashMap<K, V>,
    hashmap2: HashMap<K, V>,
    is1main: bool,
    rehashing: bool,
}

impl<K, V> RehashingHashMap<K, V>
    where K: Eq + Hash + Clone
{
    pub fn new() -> RehashingHashMap<K, V> {
        RehashingHashMap {
            hashmap1: HashMap::new(),
            hashmap2: HashMap::new(),
            is1main: true,
            rehashing: false,
        }
    }

    pub fn with_capacity(capacity: usize) -> RehashingHashMap<K, V> {
        RehashingHashMap {
            hashmap1: HashMap::with_capacity(capacity),
            hashmap2: HashMap::new(),
            is1main: true,
            rehashing: false,
        }
    }

    fn get_main(&self) -> &HashMap<K, V> {
        if self.is1main { &self.hashmap1 } else { &self.hashmap2 }
    }

    fn get_mut_main(&mut self) -> &mut HashMap<K, V> {
        if self.is1main { &mut self.hashmap1 } else { &mut self.hashmap2 }
    }

    fn get_secondary(&self) -> &HashMap<K, V> {
        if self.is1main { &self.hashmap2 } else { &self.hashmap1 }
    }

    fn get_mut_secondary(&mut self) -> &mut HashMap<K, V> {
        if self.is1main { &mut self.hashmap2 } else { &mut self.hashmap1 }
    }

    pub fn rehash(&mut self) {
        if self.rehashing {
            if self.get_secondary().len() == 0 {
                self.drop_secondary();
                return;
            }
            let (mut main, mut sec) = if self.is1main {
                (&mut self.hashmap1, &mut self.hashmap2)
            } else {
                (&mut self.hashmap2, &mut self.hashmap1)
            };
            // unwrap is safe, checked len() > 0 already
            let k: K = sec.keys().take(1).next().unwrap().clone();
            // FIXME: I wish I did not have to clone they key
            // unwrap is safe, we know the key exists in the hashmap
            let val = sec.remove(&k).unwrap();
            main.insert(k, val);
        }
    }

    pub fn capacity(&self) -> usize {
        self.get_main().capacity()
    }

    pub fn reserve(&mut self, additional: usize) {
        self.rehash();
        self.get_mut_main().reserve(additional)
    }

    pub fn is_rehashing(&self) -> bool {
        if !self.rehashing {
            assert_eq!(self.get_secondary().len(), 0);
        }
        self.rehashing
    }

    pub fn shrink_to_fit(&mut self) {
        if !self.rehashing {
            self.rehashing = true;
            self.is1main = !self.is1main;
            let len = self.len();
            self.get_mut_main().reserve(len)
        }
    }

    pub fn len(&self) -> usize {
        self.get_main().len() + self.get_secondary().len()
    }

    pub fn is_empty(&self) -> bool {
        self.get_main().is_empty() && self.get_secondary().is_empty()
    }

    fn drop_secondary(&mut self) {
        self.rehashing = false;
        assert_eq!(self.get_secondary().len(), 0);
        if self.is1main {
            self.hashmap2 = HashMap::new();
        } else {
            self.hashmap1 = HashMap::new();
        }
    }

    fn assert_state(&self) {
        #![allow(dead_code)]
        if self.rehashing {
            assert!(self.get_secondary().capacity() > 0);
        } else {
            assert!(self.get_secondary().capacity() == 0);
        }
    }

    pub fn clear(&mut self) {
        self.get_mut_main().clear();
        self.drop_secondary();
    }

    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
        // while rehashing, they key can be in either hashmap1 or hashmap2
        // but we want to remove them from wherever it is and add it to main
        let mut ret = None;
        if self.rehashing || self.is1main {
            ret = self.hashmap1.remove(&k);
        }
        if ret.is_none() && (self.rehashing || !self.is1main) {
            ret = self.hashmap2.remove(&k);
        }
        self.get_mut_main().insert(k, v);
        self.rehash();
        ret
    }

    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
            where K: Borrow<Q>, Q: Hash + Eq {
        if self.rehashing {
            match self.get_main().get(k) {
                Some(ref v) => Some(v),
                None => self.get_secondary().get(k),
            }
        } else {
            self.get_main().get(k)
        }
    }

    pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
        where K: Borrow<Q>, Q: Hash + Eq {
        if self.rehashing {
            self.rehash();
            match self.get_mut_main().remove(k) {
                Some(v) => Some(v),
                None => self.get_mut_secondary().remove(k),
            }
        } else {
            self.get_mut_main().remove(k)
        }
    }

    pub fn iter(&self) -> Iter<K, V> {
        Iter {
            inner: self.hashmap1.iter().chain(self.hashmap2.iter()),
            len: self.hashmap1.len() + self.hashmap2.len(),
        }
    }

    pub fn keys(&self) -> Keys<K, V> {
        Keys {
            inner: self.hashmap1.keys().chain(self.hashmap2.keys()),
            len: self.hashmap1.len() + self.hashmap2.len(),
        }
    }

    pub fn values(&self) -> Values<K, V> {
        Values {
            inner: self.hashmap1.values().chain(self.hashmap2.values()),
            len: self.hashmap1.len() + self.hashmap2.len(),
        }
    }
}

#[derive(Clone)]
pub struct Iter<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Iter<'a, K, V>, hash_map::Iter<'a, K, V>>,
    len: usize,
}

impl<'a, K, V> Iterator for Iter<'a, K, V> {
    type Item = (&'a K, &'a V);

    #[inline] fn next(&mut self) -> Option<(&'a K, &'a V)> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

impl<'a, K, V> ExactSizeIterator for Iter<'a, K, V> {
    #[inline] fn len(&self) -> usize { self.len }
}

#[derive(Clone)]
pub struct Keys<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Keys<'a, K, V>, hash_map::Keys<'a, K, V>>,
    len: usize,
}

impl<'a, K, V> Iterator for Keys<'a, K, V> {
    type Item = &'a K;

    #[inline] fn next(&mut self) -> Option<&'a K> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

impl<'a, K, V> ExactSizeIterator for Keys<'a, K, V> {
    #[inline] fn len(&self) -> usize { self.len }
}

#[derive(Clone)]
pub struct Values<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Values<'a, K, V>, hash_map::Values<'a, K, V>>,
    len: usize,
}

impl<'a, K, V> Iterator for Values<'a, K, V> {
    type Item = &'a V;

    #[inline] fn next(&mut self) -> Option<&'a V> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

impl<'a, K, V> ExactSizeIterator for Values<'a, K, V> {
    #[inline] fn len(&self) -> usize { self.len }
}

#[test]
fn capacity() {
    let mut hash:RehashingHashMap<u8, u8> = RehashingHashMap::with_capacity(20);
    assert!(hash.capacity() >= 20);
    hash.reserve(40);
    assert!(hash.capacity() >= 40);
}

#[test]
fn insert() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value1 = 2;
    let value2 = 3;

    assert_eq!(hash.insert(key.clone(), value1.clone()), None);
    assert_eq!(hash.insert(key.clone(), value2.clone()), Some(value1.clone()));
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert_eq!(hash.insert(key.clone(), value1.clone()), Some(value2.clone()));
    assert!(!hash.is_rehashing());
    hash.assert_state();
}

#[test]
fn insert_many_rehash_get() {
    let mut hash = RehashingHashMap::new();

    let len = 1000;

    for i in 0..len {
        hash.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2){
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    assert_eq!(hash.len(), len);

    for i in 0..len {
        assert_eq!(hash.get(&i).unwrap(), &i);
    }
    for i in len..(len * 2) {
        assert!(hash.get(&i).is_none());
    }

    for _ in 0..(len / 2 + 1){
        hash.rehash();
    }
    assert!(!hash.is_rehashing());
    hash.assert_state();

    assert_eq!(hash.len(), len);

    for i in 0..len {
        assert_eq!(hash.get(&i).unwrap(), &i);
    }
    for i in len..(len * 2) {
        assert!(hash.get(&i).is_none());
    }
}

#[test]
fn is_empty() {
    let mut hash = RehashingHashMap::new();
    assert!(hash.is_empty());

    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    assert!(!hash.is_empty());
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert!(!hash.is_empty());
    hash.rehash();
    hash.rehash();
    assert!(!hash.is_rehashing());
    assert!(!hash.is_empty());
}

#[test]
fn clear() {
    let mut hash = RehashingHashMap::with_capacity(1000);
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.clear();
    hash.assert_state();

    assert!(hash.capacity() >= 1000);
}

#[test]
fn remove0() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn remove1() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    hash.rehash();
    assert!(hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn remove2() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    hash.rehash();
    hash.rehash();
    assert!(!hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn iterator() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    assert_eq!(hash.iter().len(), len);
    for (_, i) in hash.iter() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}

#[test]
fn keys() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    assert_eq!(hash.keys().len(), len);
    for i in hash.keys() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}

#[test]
fn values() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    assert_eq!(hash.values().len(), len);
    for i in hash.values() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this file,
 * You can obtain one at http://mozilla.org/MPL/2.0/. */

#![crate_name = "js"]
#![crate_type = "rlib"]

#![feature(globs, link_args, phase, unsafe_destructor)]

#![allow(non_upper_case_globals, non_camel_case_types, non_snake_case, improper_ctypes)]

extern crate libc;
#[phase(plugin, link)]
extern crate log;
extern crate serialize;

use libc::c_uint;
use libc::types::common::c99::uint32_t;
use jsapi::{JSBool, JSContext, JSPropertyOp, JSStrictPropertyOp, JSEnumerateOp,
            JSObject, jsid, JSResolveOp, JSConvertOp, JSFinalizeOp, JSTraceOp,
            JSProto_LIMIT, JSHandleObject, JSCheckAccessOp, JSNative, JSHasInstanceOp};
use jsapi::JS_ComputeThis;
use jsval::JSVal;

// These are just macros in jsapi.h
pub use jsapi::JS_Init as JS_NewRuntime;
pub use jsapi::JS_Finish as JS_DestroyRuntime;
/*
FIXME: Not sure where JS_Lock is
pub use jsapi::bindgen::JS_Lock as JS_LockRuntime;
pub use jsapi::bindgen::JS_Unlock as JS_UnlockRuntime;
*/

pub use jsfriendapi::JSJitInfo;

pub mod jsapi;
pub mod linkhack;
pub mod rust;
pub mod glue;
pub mod trace;
pub mod jsval;
pub mod jsfriendapi;

// FIXME: Add the remaining options
pub const JSOPTION_STRICT: uint32_t =    0b00000000000001u32;
pub const JSOPTION_WERROR: uint32_t =    0b00000000000010u32;
pub const JSOPTION_VAROBJFIX: uint32_t = 0b00000000000100u32;
pub const JSOPTION_DONT_REPORT_UNCAUGHT: uint32_t = (1u32 << 8) as u32;
pub const JSOPTION_METHODJIT: uint32_t = (1u32 << 14) as u32;
pub const JSOPTION_TYPE_INFERENCE: uint32_t = (1u32 << 18) as u32;

pub const default_heapsize: u32 = 32_u32 * 1024_u32 * 1024_u32;
pub const default_stacksize: uint = 8192u;
pub const ERR: JSBool = 0_i32;

pub const JSID_TYPE_STRING: i64 = 0;
pub const JSID_TYPE_INT: i64 = 1;
pub const JSID_TYPE_VOID: i64 = 2;
pub const JSID_TYPE_OBJECT: i64 = 4;
pub const JSID_TYPE_DEFAULT_XML_NAMESPACE: i64 = 6;
pub const JSID_TYPE_MASK: i64 = 7;

pub const JSID_VOID: jsid = JSID_TYPE_VOID as jsid;

pub const JSFUN_CONSTRUCTOR: u32 = 0x200; /* native that can be called as a ctor */

pub const JSPROP_ENUMERATE: c_uint = 0x01;
pub const JSPROP_READONLY: c_uint  = 0x02;
pub const JSPROP_PERMANENT: c_uint = 0x04;
pub const JSPROP_GETTER: c_uint = 0x10;
pub const JSPROP_SETTER: c_uint = 0x20;
pub const JSPROP_SHARED: c_uint =    0x40;
pub const JSPROP_NATIVE_ACCESSORS: c_uint = 0x08;

pub const JSCLASS_RESERVED_SLOTS_SHIFT: c_uint = 8;
pub const JSCLASS_RESERVED_SLOTS_WIDTH: c_uint = 8;
pub const JSCLASS_RESERVED_SLOTS_MASK: c_uint = ((1u << JSCLASS_RESERVED_SLOTS_WIDTH as uint) - 1) as c_uint;

pub const JSCLASS_HIGH_FLAGS_SHIFT: c_uint =
    JSCLASS_RESERVED_SLOTS_SHIFT + JSCLASS_RESERVED_SLOTS_WIDTH;
pub const JSCLASS_IS_GLOBAL: c_uint = (1<<((JSCLASS_HIGH_FLAGS_SHIFT as uint)+1));

pub const JSCLASS_GLOBAL_SLOT_COUNT: c_uint = JSProto_LIMIT * 3 + 24;

pub const JSCLASS_IS_DOMJSCLASS: u32 = 1 << 4;
pub const JSCLASS_USERBIT1: u32 = 1 << 7;

pub const JSSLOT_PROXY_PRIVATE: u32 = 1;

pub const JSRESOLVE_QUALIFIED: u32 = 0x01;
pub const JSRESOLVE_ASSIGNING: u32 = 0x02;
pub const JSRESOLVE_DETECTING: u32 = 0x04;

pub enum JSGCTraceKind {
    JSTRACE_OBJECT,
    JSTRACE_STRING,
    JSTRACE_SCRIPT
}

pub fn JSCLASS_HAS_RESERVED_SLOTS(n: c_uint) -> c_uint {
    (n & JSCLASS_RESERVED_SLOTS_MASK) << (JSCLASS_RESERVED_SLOTS_SHIFT as uint)
}

#[inline(always)]
pub unsafe fn JS_ARGV(_cx: *mut JSContext, vp: *mut JSVal) -> *mut JSVal {
    vp.offset(2)
}

pub unsafe fn JS_SET_RVAL(_cx: *mut JSContext, vp: *mut JSVal, v: JSVal) {
    *vp = v;
}

#[inline(alwyas)]
pub unsafe fn JS_THIS_OBJECT(cx: *mut JSContext, vp: *mut JSVal) -> *mut JSObject {
    let r =
        if (*(vp.offset(1))).is_primitive() {
            JS_ComputeThis(cx, vp)
        } else {
            *(vp.offset(1))
        };
    r.to_object_or_null()
}

#[inline(always)]
pub unsafe fn JS_CALLEE(_cx: *mut JSContext, vp: *mut JSVal) -> JSVal {
    *vp
}

pub type JSObjectOp = extern "C" fn(*mut JSContext, JSHandleObject) -> *mut JSObject;

pub struct Class {
    pub name: *const libc::c_char,
    pub flags: uint32_t,
    pub addProperty: JSPropertyOp,
    pub delProperty: JSPropertyOp,
    pub getProperty: JSPropertyOp,
    pub setProperty: JSStrictPropertyOp,
    pub enumerate: JSEnumerateOp,
    pub resolve: JSResolveOp,
    pub convert: JSConvertOp,
    pub finalize: JSFinalizeOp,
    pub checkAccess: JSCheckAccessOp,
    pub call: JSNative,
    pub hasInstance: JSHasInstanceOp,
    pub construct: JSNative,
    pub trace: JSTraceOp,

    pub ext: ClassExtension,
    pub ops: ObjectOps,
}

pub struct ClassExtension {
    pub equality: *const u8,
    pub outerObject: Option<JSObjectOp>,
    pub innerObject: Option<JSObjectOp>,
    pub iteratorObject: *const u8,
    pub unused: *const u8,
    pub isWrappedNative: *const u8,
}

pub struct ObjectOps {
    pub lookupGeneric: *const u8,
    pub lookupProperty: *const u8,
    pub lookupElement: *const u8,
    pub lookupSpecial: *const u8,
    pub defineGeneric: *const u8,
    pub defineProperty: *const u8,
    pub defineElement: *const u8,
    pub defineSpecial: *const u8,
    pub getGeneric: *const u8,
    pub getProperty: *const u8,
    pub getElement: *const u8,
    pub getElementIfPresent: *const u8,
    pub getSpecial: *const u8,
    pub setGeneric: *const u8,
    pub setProperty: *const u8,
    pub setElement: *const u8,
    pub setSpecial: *const u8,
    pub getGenericAttributes: *const u8,
    pub getPropertyAttributes: *const u8,
    pub getElementAttributes: *const u8,
    pub getSpecialAttributes: *const u8,
    pub setGenericAttributes: *const u8,
    pub setPropertyAttributes: *const u8,
    pub setElementAttributes: *const u8,
    pub setSpecialAttributes: *const u8,
    pub deleteProperty: *const u8,
    pub deleteElement: *const u8,
    pub deleteSpecial: *const u8,

    pub enumerate: *const u8,
    pub typeOf: *const u8,
    pub thisObject: Option<JSObjectOp>,
    pub clear: *const u8,
}
<commit_msg>Add the JSOPTION_AUTOJSAPI_OWNS_ERROR_REPORTING constant.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this file,
 * You can obtain one at http://mozilla.org/MPL/2.0/. */

#![crate_name = "js"]
#![crate_type = "rlib"]

#![feature(globs, link_args, phase, unsafe_destructor)]

#![allow(non_upper_case_globals, non_camel_case_types, non_snake_case, improper_ctypes)]

extern crate libc;
#[phase(plugin, link)]
extern crate log;
extern crate serialize;

use libc::c_uint;
use libc::types::common::c99::uint32_t;
use jsapi::{JSBool, JSContext, JSPropertyOp, JSStrictPropertyOp, JSEnumerateOp,
            JSObject, jsid, JSResolveOp, JSConvertOp, JSFinalizeOp, JSTraceOp,
            JSProto_LIMIT, JSHandleObject, JSCheckAccessOp, JSNative, JSHasInstanceOp};
use jsapi::JS_ComputeThis;
use jsval::JSVal;

// These are just macros in jsapi.h
pub use jsapi::JS_Init as JS_NewRuntime;
pub use jsapi::JS_Finish as JS_DestroyRuntime;
/*
FIXME: Not sure where JS_Lock is
pub use jsapi::bindgen::JS_Lock as JS_LockRuntime;
pub use jsapi::bindgen::JS_Unlock as JS_UnlockRuntime;
*/

pub use jsfriendapi::JSJitInfo;

pub mod jsapi;
pub mod linkhack;
pub mod rust;
pub mod glue;
pub mod trace;
pub mod jsval;
pub mod jsfriendapi;

// FIXME: Add the remaining options
pub const JSOPTION_STRICT: uint32_t =    0b00000000000001u32;
pub const JSOPTION_WERROR: uint32_t =    0b00000000000010u32;
pub const JSOPTION_VAROBJFIX: uint32_t = 0b00000000000100u32;
pub const JSOPTION_DONT_REPORT_UNCAUGHT: uint32_t = (1u32 << 8) as u32;
pub const JSOPTION_AUTOJSAPI_OWNS_ERROR_REPORTING: uint32_t = (1u32 << 10) as u32;
pub const JSOPTION_METHODJIT: uint32_t = (1u32 << 14) as u32;
pub const JSOPTION_TYPE_INFERENCE: uint32_t = (1u32 << 18) as u32;

pub const default_heapsize: u32 = 32_u32 * 1024_u32 * 1024_u32;
pub const default_stacksize: uint = 8192u;
pub const ERR: JSBool = 0_i32;

pub const JSID_TYPE_STRING: i64 = 0;
pub const JSID_TYPE_INT: i64 = 1;
pub const JSID_TYPE_VOID: i64 = 2;
pub const JSID_TYPE_OBJECT: i64 = 4;
pub const JSID_TYPE_DEFAULT_XML_NAMESPACE: i64 = 6;
pub const JSID_TYPE_MASK: i64 = 7;

pub const JSID_VOID: jsid = JSID_TYPE_VOID as jsid;

pub const JSFUN_CONSTRUCTOR: u32 = 0x200; /* native that can be called as a ctor */

pub const JSPROP_ENUMERATE: c_uint = 0x01;
pub const JSPROP_READONLY: c_uint  = 0x02;
pub const JSPROP_PERMANENT: c_uint = 0x04;
pub const JSPROP_GETTER: c_uint = 0x10;
pub const JSPROP_SETTER: c_uint = 0x20;
pub const JSPROP_SHARED: c_uint =    0x40;
pub const JSPROP_NATIVE_ACCESSORS: c_uint = 0x08;

pub const JSCLASS_RESERVED_SLOTS_SHIFT: c_uint = 8;
pub const JSCLASS_RESERVED_SLOTS_WIDTH: c_uint = 8;
pub const JSCLASS_RESERVED_SLOTS_MASK: c_uint = ((1u << JSCLASS_RESERVED_SLOTS_WIDTH as uint) - 1) as c_uint;

pub const JSCLASS_HIGH_FLAGS_SHIFT: c_uint =
    JSCLASS_RESERVED_SLOTS_SHIFT + JSCLASS_RESERVED_SLOTS_WIDTH;
pub const JSCLASS_IS_GLOBAL: c_uint = (1<<((JSCLASS_HIGH_FLAGS_SHIFT as uint)+1));

pub const JSCLASS_GLOBAL_SLOT_COUNT: c_uint = JSProto_LIMIT * 3 + 24;

pub const JSCLASS_IS_DOMJSCLASS: u32 = 1 << 4;
pub const JSCLASS_USERBIT1: u32 = 1 << 7;

pub const JSSLOT_PROXY_PRIVATE: u32 = 1;

pub const JSRESOLVE_QUALIFIED: u32 = 0x01;
pub const JSRESOLVE_ASSIGNING: u32 = 0x02;
pub const JSRESOLVE_DETECTING: u32 = 0x04;

pub enum JSGCTraceKind {
    JSTRACE_OBJECT,
    JSTRACE_STRING,
    JSTRACE_SCRIPT
}

pub fn JSCLASS_HAS_RESERVED_SLOTS(n: c_uint) -> c_uint {
    (n & JSCLASS_RESERVED_SLOTS_MASK) << (JSCLASS_RESERVED_SLOTS_SHIFT as uint)
}

#[inline(always)]
pub unsafe fn JS_ARGV(_cx: *mut JSContext, vp: *mut JSVal) -> *mut JSVal {
    vp.offset(2)
}

pub unsafe fn JS_SET_RVAL(_cx: *mut JSContext, vp: *mut JSVal, v: JSVal) {
    *vp = v;
}

#[inline(alwyas)]
pub unsafe fn JS_THIS_OBJECT(cx: *mut JSContext, vp: *mut JSVal) -> *mut JSObject {
    let r =
        if (*(vp.offset(1))).is_primitive() {
            JS_ComputeThis(cx, vp)
        } else {
            *(vp.offset(1))
        };
    r.to_object_or_null()
}

#[inline(always)]
pub unsafe fn JS_CALLEE(_cx: *mut JSContext, vp: *mut JSVal) -> JSVal {
    *vp
}

pub type JSObjectOp = extern "C" fn(*mut JSContext, JSHandleObject) -> *mut JSObject;

pub struct Class {
    pub name: *const libc::c_char,
    pub flags: uint32_t,
    pub addProperty: JSPropertyOp,
    pub delProperty: JSPropertyOp,
    pub getProperty: JSPropertyOp,
    pub setProperty: JSStrictPropertyOp,
    pub enumerate: JSEnumerateOp,
    pub resolve: JSResolveOp,
    pub convert: JSConvertOp,
    pub finalize: JSFinalizeOp,
    pub checkAccess: JSCheckAccessOp,
    pub call: JSNative,
    pub hasInstance: JSHasInstanceOp,
    pub construct: JSNative,
    pub trace: JSTraceOp,

    pub ext: ClassExtension,
    pub ops: ObjectOps,
}

pub struct ClassExtension {
    pub equality: *const u8,
    pub outerObject: Option<JSObjectOp>,
    pub innerObject: Option<JSObjectOp>,
    pub iteratorObject: *const u8,
    pub unused: *const u8,
    pub isWrappedNative: *const u8,
}

pub struct ObjectOps {
    pub lookupGeneric: *const u8,
    pub lookupProperty: *const u8,
    pub lookupElement: *const u8,
    pub lookupSpecial: *const u8,
    pub defineGeneric: *const u8,
    pub defineProperty: *const u8,
    pub defineElement: *const u8,
    pub defineSpecial: *const u8,
    pub getGeneric: *const u8,
    pub getProperty: *const u8,
    pub getElement: *const u8,
    pub getElementIfPresent: *const u8,
    pub getSpecial: *const u8,
    pub setGeneric: *const u8,
    pub setProperty: *const u8,
    pub setElement: *const u8,
    pub setSpecial: *const u8,
    pub getGenericAttributes: *const u8,
    pub getPropertyAttributes: *const u8,
    pub getElementAttributes: *const u8,
    pub getSpecialAttributes: *const u8,
    pub setGenericAttributes: *const u8,
    pub setPropertyAttributes: *const u8,
    pub setElementAttributes: *const u8,
    pub setSpecialAttributes: *const u8,
    pub deleteProperty: *const u8,
    pub deleteElement: *const u8,
    pub deleteSpecial: *const u8,

    pub enumerate: *const u8,
    pub typeOf: *const u8,
    pub thisObject: Option<JSObjectOp>,
    pub clear: *const u8,
}
<|endoftext|>


<commit_before>
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::{Relaxed, SeqCst};

pub trait AtomicCounter {
    type PrimitiveType;

    fn inc(&self);
    fn add(&self, amount: Self::PrimitiveType);
    fn reset(&self) -> Self::PrimitiveType;
}

#[derive(Debug, Default)]
pub struct RelaxedCounter(AtomicUsize);

impl RelaxedCounter {
    pub fn new(initial_count: usize) -> RelaxedCounter {
        RelaxedCounter(AtomicUsize::new(initial_count))
    }
}

impl AtomicCounter for RelaxedCounter {
    type PrimitiveType = usize;

    fn inc(&self) {
        self.add(1)
    }

    fn add(&self, amount: usize) {
        self.0.fetch_add(amount, Relaxed);
    }

    fn reset(&self) -> usize {
        self.0.swap(0, Relaxed)
    }
}

#[derive(Debug, Default)]
pub struct ConsistentCounter(AtomicUsize);

impl ConsistentCounter {
    pub fn new(initial_count: usize) -> ConsistentCounter {
        ConsistentCounter(AtomicUsize::new(initial_count))
    }
}

impl AtomicCounter for ConsistentCounter {
    type PrimitiveType = usize;

    fn inc(&self) {
        self.add(1)
    }

    fn add(&self, amount: usize) {
        self.0.fetch_add(amount, SeqCst);
    }

    fn reset(&self) -> usize {
        self.0.swap(0, SeqCst)
    }
}

#[cfg(test)]
mod tests {

    use std::thread;
    use std::sync::Arc;
    use std::ops::Deref;

    use super::*;

    const NUM_THREADS: usize = 29;
    const NUM_ITERATIONS: usize = 7_000_000;

    #[test]
    fn test_inc() {
        let mut join_handles = Vec::new();
        let counter = Arc::new(RelaxedCounter::new(0));
        println!("test_inc: Spawning {} threads, each with {} iterations...", NUM_THREADS, NUM_ITERATIONS);
        for _ in 0..NUM_THREADS {
            let counter_ref = counter.clone();
            join_handles.push(thread::spawn(move || {
                //make sure we're not going though Arc on each iteration
                let counter: &RelaxedCounter = counter_ref.deref();
                for _ in 0..NUM_ITERATIONS {
                    counter.inc();
                }
            }));
        }
        for handle in join_handles {
            handle.join().unwrap();
        }
        let count = counter.reset();
        println!("test_inc: Got count: {}", count);
        assert_eq!(NUM_THREADS * NUM_ITERATIONS, count);
    }

    #[test]
    fn test_add() {
        let mut join_handles = Vec::new();
        let counter = Arc::new(RelaxedCounter::new(0));
        println!("test_add: Spawning {} threads, each with {} iterations...", NUM_THREADS, NUM_ITERATIONS);
        let mut expected_count = 0;
        for to_add in 0..NUM_THREADS {
            let counter_ref = counter.clone();
            expected_count += to_add * NUM_ITERATIONS;
            join_handles.push(thread::spawn(move || {
                //make sure we're not going though Arc on each iteration
                let counter: &RelaxedCounter = counter_ref.deref();
                for _ in 0..NUM_ITERATIONS {
                    counter.add(to_add);
                }
            }));
        }
        for handle in join_handles {
            handle.join().unwrap();
        }
        let count = counter.reset();
        println!("test_add: Expected count: {}, got count: {}", expected_count, count);
        assert_eq!(expected_count, count);
    }

    #[test]
    fn test_add_reset() {
        let mut join_handles = Vec::new();
        let counter = Arc::new(RelaxedCounter::new(0));
        println!("test_add_reset: Spawning {} threads, each with {} iterations...", NUM_THREADS, NUM_ITERATIONS);
        let mut expected_count = 0;
        for to_add in 0..NUM_THREADS {
            expected_count += to_add * NUM_ITERATIONS;
        }

        // setup thread that `reset()`s all the time
        let counter_ref = counter.clone();
        let reset_handle = thread::spawn(move || {
            // Usually, you would check for some better termination condition.
            // I don't want to pollute my test with thread synchronization
            // operations outside of AtomicCounter, hence this approach.
            let mut total_count = 0;
            let counter: &RelaxedCounter = counter_ref.deref();
            while total_count < expected_count {
                total_count += counter.reset();
            }
            // Ok, now we got the total_count but this could just be lucky.
            // Better do some more resets to be sure... ;)
            for _ in 0..NUM_ITERATIONS {
                total_count += counter.reset();
            }
            total_count
        });

        for to_add in 0..NUM_THREADS {
            let counter_ref = counter.clone();

            join_handles.push(thread::spawn(move || {
                //make sure we're not going though Arc on each iteration
                let counter: &RelaxedCounter = counter_ref.deref();
                for _ in 0..NUM_ITERATIONS {
                    counter.add(to_add);
                }
            }));
        }
        for handle in join_handles {
            handle.join().unwrap();
        }
        let actual_count = reset_handle.join().unwrap();
        println!("test_add_reset: Expected count: {}, got count: {}", expected_count, actual_count);
        assert_eq!(expected_count, actual_count);
    }

}
<commit_msg>Crate documentation; add into_inner(self); Add tests for _both_ Counters;<commit_after>
//! # AtomicCounter
//!
//! Atomic (thread-safe) counters for Rust.
//!
//! This crate contains an [`AtomicCounter`](trait.AtomicCounter.html) trait
//! that can safely be shared across threads.
//!
//! This crate provides two implementations:
//!
//! * [`RelaxedCounter`](struct.RelaxedCounter.html) which is suitable for
//!     e.g. collecting metrics or generate IDs, but which does not provide
//!     ["Sequential Consistency"](https://doc.rust-lang.org/nomicon/atomics.html#sequentially-consistent).
//!     `RelaxedCounter` uses [`Relaxed`](https://doc.rust-lang.org/std/sync/atomic/enum.Ordering.html#variant.Relaxed)
//!     memory ordering.
//!
//! * [`ConsistentCounter`](struct.ConsistentCounter.html) which provides the
//!     same interface but is sequentially consistent. Use this counter if the
//!     order of update from multiple threads is important.
//!     `ConsistentCounter` uses [`Sequentially Consistent`](https://doc.rust-lang.org/std/sync/atomic/enum.Ordering.html#variant.SeqCst)
//!     memory ordering.
//!
//! Both implementations are lock-free. Both are a very thin layer over
//! [`AtomicUsize`](https://doc.rust-lang.org/std/sync/atomic/struct.AtomicUsize.html)
//! which is more powerful but might be harder to use correctly.
//!
//! ## Which counter to use
//!
//! * If you are just collecting metrics, the `RelaxedCounter` is probably right choice.
//!
//! * If you are generating IDs, but don't make strong assumptions (like allocating
//!     memory based on the ID count), `RelaxedCounter` is probably the right choice.
//!
//! * If you are generating multiple IDs where you maintain an ordering
//!     invariant (e.g. ID `a` is always greater than ID `b`), you need "Sequential
//!     Consistency" and thus need to use `SequentialCounter`. The same is true
//!     for all use cases where the _ordering_ of incrementing the counter is
//!     important.
//!
//! ## No updates are lost - It's just about the ordering!
//!
//! Note that in both implementations, _no count is lost_ and all operations are atomic.
//! The difference is _only_ in how the order of operations are observed by different
//! threads.
//!
//! ## Example:
//! Assume `a` is 5 and `b` is 4. You always want to maintain `a > b`.
//!
//! Thread 1 executes this code:
//!
//! ```rust,ignore
//!
//! a.inc();
//! b.inc();
//! ```
//!
//! Thread 2 gets counts:
//!
//! ```rust,ignore
//!
//! let a_local = a.get();
//! let b_local = b.get();
//! ```
//!
//! What are the values for `a_local` and `b_local`? That depends on the order
//! in which thread 1 and 2 have run:
//!
//! * `a_local` could still be 5 and `b_local` is still be 4 (e.g. if thread 2 ran before thread 1)
//! * `a_local` could be increment to 6 while `b_local` is still at 4 (e.g. if thread 1 and 2 ran in parallel)
//! * `a_local` could be increment to 6 and `b_local` be incremented to 5 (e.g. if thread 2 ran after thread 1).
//! * Additionally, if at least one counter is a `RelaxedCounter`, we cannot make
//!     assumption on the order of `a.inc()` and `b.inc()`. Thus, in this case
//!     thread 2 can also observe `a_local` to be 5 (not incremented yet) but
//!     `b_local` to be incremented to 5, _breaking the invariant_ `a > b`.
//!     Note that if thread 2 (or any other thread) `get()` the counts
//!     again, at some point they will observe both values to be incremented.
//!     No operations will be lost. It is only the _ordering_ of the operations
//!     that cannot be assumed if `Ordering` is `Relaxed`.
//!
//! So in order to maintain invariants such as `a > b` across multiple threads,
//! use `ConsistentCounter`.

use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::{Relaxed, SeqCst};

/// Provides an atomic counter trait that can be shared across threads.
pub trait AtomicCounter: Send + Sync {
    /// Underlying primitive type that is being shared atomically.
    type PrimitiveType;

    /// Atomically increments the counter by one.
    fn inc(&self);

    /// Atomically increments the counter by amount.
    fn add(&self, amount: Self::PrimitiveType);

    /// Atomically gets the current value of the counter, without modifying the counter.
    fn get(&self) -> Self::PrimitiveType;

    /// Atomically returns the current value of the counter, while resetting to count to zero.
    fn reset(&self) -> Self::PrimitiveType;

    /// Consume the atomic counter and return the primitive type.
    ///
    /// This is safe because passing self by value guarantees that no other threads are concurrently accessing the atomic data.
    fn into_inner(self) -> Self::PrimitiveType;
}

/// Implementation of [`AtomicCounter`](trait.AtomicCounter.html) that uses
/// [`Relaxed`](https://doc.rust-lang.org/std/sync/atomic/enum.Ordering.html#variant.Relaxed)
/// memory ordering.
///
/// See [crate level documentation](index.html) for more details.
///
/// Note that all operations wrap if the counter is incremented beyond usize::max_value().
#[derive(Debug, Default)]
pub struct RelaxedCounter(AtomicUsize);

impl RelaxedCounter {

    /// Creates a new counter with initial_value
    pub fn new(initial_count: usize) -> RelaxedCounter {
        RelaxedCounter(AtomicUsize::new(initial_count))
    }
}

impl AtomicCounter for RelaxedCounter {
    type PrimitiveType = usize;

    fn inc(&self) {
        self.add(1)
    }

    fn add(&self, amount: usize) {
        self.0.fetch_add(amount, Relaxed);
    }

    fn get(&self) -> usize {
        self.0.load(Relaxed)
    }

    fn reset(&self) -> usize {
        self.0.swap(0, Relaxed)
    }

    fn into_inner(self) -> usize {
        self.0.into_inner()
    }
}

/// Implementation of [`AtomicCounter`](trait.AtomicCounter.html) that uses
/// [`Sequentially Consistent`](https://doc.rust-lang.org/std/sync/atomic/enum.Ordering.html#variant.SeqCst)
/// memory ordering.
///
/// See [crate level documentation](index.html) for more details.
///
/// Note that all operations wrap if the counter is incremented beyond usize::max_value().
#[derive(Debug, Default)]
pub struct ConsistentCounter(AtomicUsize);

impl ConsistentCounter {

    /// Creates a new counter with initial_value
    pub fn new(initial_count: usize) -> ConsistentCounter {
        ConsistentCounter(AtomicUsize::new(initial_count))
    }
}

impl AtomicCounter for ConsistentCounter {
    type PrimitiveType = usize;

    fn inc(&self) {
        self.add(1)
    }

    fn add(&self, amount: usize) {
        self.0.fetch_add(amount, SeqCst);
    }

    fn get(&self) -> usize {
        self.0.load(SeqCst)
    }

    fn reset(&self) -> usize {
        self.0.swap(0, SeqCst)
    }

    fn into_inner(self) -> usize {
        self.0.into_inner()
    }
}

#[cfg(test)]
mod tests {

    use std::fmt::Debug;
    use std::thread;
    use std::sync::Arc;
    use std::ops::Deref;

    use super::*;

    const NUM_THREADS: usize = 29;
    const NUM_ITERATIONS: usize = 7_000_000;

    fn test_inc_with<Counter>(counter: Arc<Counter>)
        where Counter: AtomicCounter<PrimitiveType=usize> + 'static + Debug
    {
        let mut join_handles = Vec::new();
        println!("test_inc: Spawning {} threads, each with {} iterations...",
                 NUM_THREADS,
                 NUM_ITERATIONS);
        for _ in 0..NUM_THREADS {
            let counter_ref = counter.clone();
            join_handles.push(thread::spawn(move || {
                //make sure we're not going though Arc on each iteration
                let counter: &Counter = counter_ref.deref();
                for _ in 0..NUM_ITERATIONS {
                    counter.inc();
                }
            }));
        }
        for handle in join_handles {
            handle.join().unwrap();
        }
        let count = Arc::try_unwrap(counter).unwrap().into_inner();
        println!("test_inc: Got count: {}", count);
        assert_eq!(NUM_THREADS * NUM_ITERATIONS, count);
    }

    #[test]
    fn test_inc_relaxed() {
        test_inc_with(Arc::new(RelaxedCounter::new(0)));
    }

    #[test]
    fn test_inc_consistent() {
        test_inc_with(Arc::new(ConsistentCounter::new(0)));
    }

    fn test_add_with<Counter>(counter: Arc<Counter>)
        where Counter: AtomicCounter<PrimitiveType=usize> + 'static + Debug
    {
        let mut join_handles = Vec::new();
        println!("test_add: Spawning {} threads, each with {} iterations...",
                 NUM_THREADS,
                 NUM_ITERATIONS);
        let mut expected_count = 0;
        for to_add in 0..NUM_THREADS {
            let counter_ref = counter.clone();
            expected_count += to_add * NUM_ITERATIONS;
            join_handles.push(thread::spawn(move || {
                //make sure we're not going though Arc on each iteration
                let counter: &Counter = counter_ref.deref();
                for _ in 0..NUM_ITERATIONS {
                    counter.add(to_add);
                }
            }));
        }
        for handle in join_handles {
            handle.join().unwrap();
        }
        let count = Arc::try_unwrap(counter).unwrap().into_inner();
        println!("test_add: Expected count: {}, got count: {}",
                 expected_count,
                 count);
        assert_eq!(expected_count, count);
    }

    #[test]
    fn test_add_relaxed() {
        test_add_with(Arc::new(RelaxedCounter::new(0)));
    }

    #[test]
    fn test_add_consistent() {
        test_add_with(Arc::new(ConsistentCounter::new(0)));
    }

    fn test_reset_with<Counter>(counter: Arc<Counter>)
        where Counter: AtomicCounter<PrimitiveType=usize> + 'static + Debug
    {
        let mut join_handles = Vec::new();
        println!("test_add_reset: Spawning {} threads, each with {} iterations...",
                 NUM_THREADS,
                 NUM_ITERATIONS);
        let mut expected_count = 0;
        for to_add in 0..NUM_THREADS {
            expected_count += to_add * NUM_ITERATIONS;
        }

        // setup thread that `reset()`s all the time
        let counter_ref = counter.clone();
        let reset_handle = thread::spawn(move || {
            // Usually, you would check for some better termination condition.
            // I don't want to pollute my test with thread synchronization
            // operations outside of AtomicCounter, hence this approach.
            let mut total_count = 0;
            let counter: &Counter = counter_ref.deref();
            while total_count < expected_count {
                total_count += counter.reset();
            }
            // Ok, now we got the total_count but this could just be lucky.
            // Better do some more resets to be sure... ;)
            for _ in 0..NUM_ITERATIONS {
                total_count += counter.reset();
            }
            total_count
        });

        for to_add in 0..NUM_THREADS {
            let counter_ref = counter.clone();

            join_handles.push(thread::spawn(move || {
                //make sure we're not going though Arc on each iteration
                let counter: &Counter = counter_ref.deref();
                for _ in 0..NUM_ITERATIONS {
                    counter.add(to_add);
                }
            }));
        }
        for handle in join_handles {
            handle.join().unwrap();
        }
        let actual_count = reset_handle.join().unwrap();
        println!("test_add_reset: Expected count: {}, got count: {}",
                 expected_count,
                 actual_count);
        assert_eq!(expected_count, actual_count);
    }

    #[test]
    fn test_reset_consistent() {
        test_reset_with(Arc::new(ConsistentCounter::new(0)));
    }

    #[test]
    fn test_reset_relaxed() {
        test_reset_with(Arc::new(RelaxedCounter::new(0)));
    }

}
<|endoftext|>


<commit_before>#![feature(io)]
#![feature(path)]
#![feature(fs)]
#![feature(collections)]
#![feature(core)]

use std::fs::{File};
use std::string::String;
use std::path::AsPath;
use std::error::FromError;
use std::io::{Read, BufReader, Write};
use std::vec::Vec;
use std::num::ParseIntError;
use std::mem;

#[derive(Debug)]
pub enum Error {
    IO,
    Parse,
    DataType,
    FieldType,
    Malformed,
    NotImplemented
}

impl FromError<std::io::Error> for Error {
    fn from_error(_: std::io::Error) -> Error {
        Error::IO
    }
}

impl FromError<ParseIntError> for Error {
    fn from_error(_: ParseIntError) -> Error {
        Error::Parse
    }
}

#[derive(Debug)]
pub enum DataType {
    XDRFloat,
    FloatLE,
}

impl DataType {
    fn from_str(s: &str) -> Result<DataType, Error> {
        match s {
            "float_le" => Ok(DataType::FloatLE),
//            "xdr_float" => Ok(DataType::XDRFloat),
            _ => Err(Error::DataType)
        }
    }
}

#[derive(Debug)]
pub enum FieldType {
    Uniform
}

impl FieldType {
    fn from_str(s: &str) -> Result<FieldType, Error> {
        match s {
            "uniform" => Ok(FieldType::Uniform),
            _ => Err(Error::FieldType)
        }
    }
}

pub struct AVSFile<'a> {
    pub ndim: usize,
    pub sizes: Vec<usize>,
    pub data_type: DataType,
    pub field_type: FieldType,
    reader: Box<Read + 'a>
}

impl<'a> AVSFile<'a> {
    pub fn write<W: Write, T>(
                writer: &mut W, dims: &[usize], data: &[T])
                    -> Result<(), Error> {
        // header
        let ndim = dims.len();
        try!(writer.write_fmt(format_args!("# AVS FLD file (written by avsfldrs github.com/greyhill/avsfldrs)\n")));
        try!(writer.write_fmt(format_args!("ndim={}\n", ndim)));
        try!(writer.write_fmt(format_args!("veclen=1\n")));
        try!(writer.write_fmt(format_args!("nspace={}\n", ndim)));
        try!(writer.write_fmt(format_args!("field=uniform\n")));
        try!(writer.write_fmt(format_args!("data=float_le\n"))); // TODO
        for (id, size) in dims.iter().enumerate() {
            try!(writer.write_fmt(format_args!("dim{}={}\n", id+1, size)));
        }
        try!(writer.write_fmt(format_args!("{}{}", 12 as char, 12 as char)));
        let b: &[u8] = unsafe {
            std::slice::from_raw_parts(data.as_ptr() as *const u8,
                                       data.len()*mem::size_of::<T>())
        };
        try!(writer.write_all(b));
        Ok(())
    }

    pub fn read<T>(self: &mut Self) -> Result<Vec<T>, Error> {
        let size = self.sizes.iter().fold(1 as usize, |l, r| l * *r);
        let mut buf_u8 = Vec::<u8>::with_capacity(mem::size_of::<T>()*size);
        try!(self.reader.read_to_end(&mut buf_u8));
        let buf: Vec<T> = unsafe {
            let ptr = buf_u8.as_mut_ptr();
            let cap = buf_u8.capacity();
            mem::forget(buf_u8);
            Vec::<T>::from_raw_parts(
                mem::transmute(ptr),
                size,
                cap / mem::size_of::<T>())
        };
        Ok(buf)
    }

    pub fn open<P: AsPath>(path: &P) -> Result<AVSFile, Error> {
        let mut reader = BufReader::new(try!(File::open(path)));

        let mut ndim: Option<usize> = None;
        let mut sizes = Vec::<Option<usize>>::new();
        let mut data_type: Option<DataType> = None;
        let mut field_type: Option<FieldType> = None;
        let mut external: Option<String> = None;

        let mut line = String::new();
        let mut last_char: u8 = 0;
        loop {
            let mut new_char_buf: [u8;1] = [ 0u8 ];
            try!(reader.read(&mut new_char_buf));

            // break on two chr 14s
            let new_char = new_char_buf[0];
            if (new_char, last_char) == (12u8, 12u8) {
                break;
            }
            last_char = new_char;

            line.push(new_char as char);

            // new line; process the line and discard
            if new_char == 10 {
                let tokens: Vec<&str> = line.split('=')
                    .map(|s| s.trim()).collect();
                match &tokens[..] {
                    ["ndim", s] => {
                        let nd = try!(s.parse::<usize>());
                        ndim = Some(nd);
                        for _ in 0..nd {
                            sizes.push(None);
                        }
                    },
                    ["dim1", s] => sizes[0] = Some(try!(s.parse::<usize>())),
                    ["dim2", s] => sizes[1] = Some(try!(s.parse::<usize>())),
                    ["dim3", s] => sizes[2] = Some(try!(s.parse::<usize>())),
                    ["dim4", s] => sizes[3] = Some(try!(s.parse::<usize>())),
                    ["dim5", s] => sizes[4] = Some(try!(s.parse::<usize>())),
                    ["dim6", s] => sizes[5] = Some(try!(s.parse::<usize>())),
                    ["dim7", s] => sizes[6] = Some(try!(s.parse::<usize>())),
                    ["data", s] =>
                        data_type = Some(try!(DataType::from_str(s))),
                    ["field", s] =>
                        field_type = Some(try!(FieldType::from_str(s))),
                    ["variable 1 file", s] =>
                        external = Some(String::from_str(s)),
                    _ => {}
                }
            }
            // hack?  code smell?  need borrow in previous block to expire
            if new_char == 10 {
                line.clear();
            }
        }

        match external {
            None => {
                let mut tr = AVSFile {
                    ndim: try!(ndim.ok_or(Error::Malformed)),
                    sizes: Vec::<usize>::new(),
                    data_type: try!(data_type.ok_or(Error::Malformed)),
                    field_type: try!(field_type.ok_or(Error::Malformed)),
                    reader: Box::new(reader),
                };
                for idx in 0..ndim.unwrap() {
                    tr.sizes.push(
                        try!(sizes[idx].ok_or(Error::Malformed)));
                }
                Ok(tr)
            },
            Some(path) => {
                let new_reader = BufReader::new(try!(File::open(&path)));
                let mut tr = AVSFile {
                    ndim: try!(ndim.ok_or(Error::Malformed)),
                    sizes: Vec::<usize>::new(),
                    data_type: try!(data_type.ok_or(Error::Malformed)),
                    field_type: try!(field_type.ok_or(Error::Malformed)),
                    reader: Box::new(new_reader),
                };
                for idx in 0..ndim.unwrap() {
                    tr.sizes.push(
                        try!(sizes[idx].ok_or(Error::Malformed)));
                }
                Ok(tr)
            },
        }
    }
}

<commit_msg>updated to more recent rust<commit_after>#![feature(collections)]
#![feature(slice_patterns)]

use std::fs::{File};
use std::string::String;
use std::path::Path;
use std::error::FromError;
use std::io::{Read, BufReader, Write};
use std::vec::Vec;
use std::num::ParseIntError;
use std::mem;

#[derive(Debug)]
pub enum Error {
    IO,
    Parse,
    DataType,
    FieldType,
    Malformed,
    NotImplemented
}

impl FromError<std::io::Error> for Error {
    fn from_error(_: std::io::Error) -> Error {
        Error::IO
    }
}

impl FromError<ParseIntError> for Error {
    fn from_error(_: ParseIntError) -> Error {
        Error::Parse
    }
}

#[derive(Debug)]
pub enum DataType {
    XDRFloat,
    FloatLE,
}

impl DataType {
    fn from_str(s: &str) -> Result<DataType, Error> {
        match s {
            "float_le" => Ok(DataType::FloatLE),
//            "xdr_float" => Ok(DataType::XDRFloat),
            _ => Err(Error::DataType)
        }
    }
}

#[derive(Debug)]
pub enum FieldType {
    Uniform
}

impl FieldType {
    fn from_str(s: &str) -> Result<FieldType, Error> {
        match s {
            "uniform" => Ok(FieldType::Uniform),
            _ => Err(Error::FieldType)
        }
    }
}

pub struct AVSFile {
    pub ndim: usize,
    pub sizes: Vec<usize>,
    pub data_type: DataType,
    pub field_type: FieldType,
    reader: Box<Read>
}

impl AVSFile {
    pub fn write<W: Write, T>(
                writer: &mut W, dims: &[usize], data: &[T])
                    -> Result<(), Error> {
        // header
        let ndim = dims.len();
        try!(writer.write_fmt(format_args!("# AVS FLD file (written by avsfldrs github.com/greyhill/avsfldrs)\n")));
        try!(writer.write_fmt(format_args!("ndim={}\n", ndim)));
        try!(writer.write_fmt(format_args!("veclen=1\n")));
        try!(writer.write_fmt(format_args!("nspace={}\n", ndim)));
        try!(writer.write_fmt(format_args!("field=uniform\n")));
        try!(writer.write_fmt(format_args!("data=float_le\n"))); // TODO
        for (id, size) in dims.iter().enumerate() {
            try!(writer.write_fmt(format_args!("dim{}={}\n", id+1, size)));
        }
        try!(writer.write_fmt(format_args!("{}{}", 12 as char, 12 as char)));
        let b: &[u8] = unsafe {
            std::slice::from_raw_parts(data.as_ptr() as *const u8,
                                       data.len()*mem::size_of::<T>())
        };
        try!(writer.write_all(b));
        Ok(())
    }

    pub fn read<T>(self: &mut Self) -> Result<Vec<T>, Error> {
        let size = self.sizes.iter().fold(1 as usize, |l, r| l * *r);
        let mut buf_u8 = Vec::<u8>::with_capacity(mem::size_of::<T>()*size);
        try!(self.reader.read_to_end(&mut buf_u8));
        let buf: Vec<T> = unsafe {
            let ptr = buf_u8.as_mut_ptr();
            let cap = buf_u8.capacity();
            mem::forget(buf_u8);
            Vec::<T>::from_raw_parts(
                mem::transmute(ptr),
                size,
                cap / mem::size_of::<T>())
        };
        Ok(buf)
    }

    pub fn open(path: &Path) -> Result<AVSFile, Error> {
        let mut reader = BufReader::new(try!(File::open(path)));

        let mut ndim: Option<usize> = None;
        let mut sizes = Vec::<Option<usize>>::new();
        let mut data_type: Option<DataType> = None;
        let mut field_type: Option<FieldType> = None;
        let mut external: Option<String> = None;

        let mut line = String::new();
        let mut last_char: u8 = 0;
        loop {
            let mut new_char_buf: [u8;1] = [ 0u8 ];
            try!(reader.read(&mut new_char_buf));

            // break on two chr 14s
            let new_char = new_char_buf[0];
            if (new_char, last_char) == (12u8, 12u8) {
                break;
            }
            last_char = new_char;

            line.push(new_char as char);

            // new line; process the line and discard
            if new_char == 10 {
                let tokens: Vec<&str> = line.split('=')
                    .map(|s| s.trim()).collect();
                match &tokens[..] {
                    ["ndim", s] => {
                        let nd = try!(s.parse::<usize>());
                        ndim = Some(nd);
                        for _ in 0..nd {
                            sizes.push(None);
                        }
                    },
                    ["dim1", s] => sizes[0] = Some(try!(s.parse::<usize>())),
                    ["dim2", s] => sizes[1] = Some(try!(s.parse::<usize>())),
                    ["dim3", s] => sizes[2] = Some(try!(s.parse::<usize>())),
                    ["dim4", s] => sizes[3] = Some(try!(s.parse::<usize>())),
                    ["dim5", s] => sizes[4] = Some(try!(s.parse::<usize>())),
                    ["dim6", s] => sizes[5] = Some(try!(s.parse::<usize>())),
                    ["dim7", s] => sizes[6] = Some(try!(s.parse::<usize>())),
                    ["data", s] =>
                        data_type = Some(try!(DataType::from_str(s))),
                    ["field", s] =>
                        field_type = Some(try!(FieldType::from_str(s))),
                    ["variable 1 file", s] =>
                        external = Some(String::from_str(s)),
                    _ => {}
                }
            }
            // hack?  code smell?  need borrow in previous block to expire
            if new_char == 10 {
                line.clear();
            }
        }

        match external {
            None => {
                let mut tr = AVSFile {
                    ndim: try!(ndim.ok_or(Error::Malformed)),
                    sizes: Vec::<usize>::new(),
                    data_type: try!(data_type.ok_or(Error::Malformed)),
                    field_type: try!(field_type.ok_or(Error::Malformed)),
                    reader: Box::new(reader),
                };
                for idx in 0..ndim.unwrap() {
                    tr.sizes.push(
                        try!(sizes[idx].ok_or(Error::Malformed)));
                }
                Ok(tr)
            },
            Some(path) => {
                let new_reader = BufReader::new(try!(File::open(&path)));
                let mut tr = AVSFile {
                    ndim: try!(ndim.ok_or(Error::Malformed)),
                    sizes: Vec::<usize>::new(),
                    data_type: try!(data_type.ok_or(Error::Malformed)),
                    field_type: try!(field_type.ok_or(Error::Malformed)),
                    reader: Box::new(new_reader),
                };
                for idx in 0..ndim.unwrap() {
                    tr.sizes.push(
                        try!(sizes[idx].ok_or(Error::Malformed)));
                }
                Ok(tr)
            },
        }
    }
}

<|endoftext|>


<commit_before>goto: file://|DOC_PATH|/test_docs/type.SomeType.html
assert-all: (".top-block .docblock p", {"font-weight": "400"})
<commit_msg>Remove erroneously created file<commit_after><|endoftext|>