<commit_before>#![feature(drain)]
#![feature(plugin)]
#![plugin(phf_macros)]

extern crate phf;
extern crate itertools;
extern crate unicode_segmentation;
extern crate unicode_normalization;
extern crate rustc_serialize;

mod utils;
mod suffix;
mod nickname;
mod title;
mod surname;
mod namecase;
mod namepart;
mod parse;
mod comparison;
mod serialization;

#[cfg(feature = "name_eq_hash")]
mod eq_hash;

use std::borrow::Cow;
use itertools::Itertools;
use utils::is_mixed_case;

pub struct Name {
    words: Vec<String>,
    surname_index: usize,
    suffix_index: usize,
    initials: String,
    word_indices_in_initials: Vec<(usize,usize)>,
}

impl Name {

    /// Parses a string represent a single person's full name into a canonical
    /// representation.
    ///
    /// # Examples
    /// ```
    /// use human_name::Name;
    ///
    /// let name = Name::parse("Jane Doe").unwrap();
    /// assert_eq!("Doe", name.surname());
    /// assert_eq!(Some("Jane"), name.given_name());
    ///
    /// let name = Name::parse("Doe, J").unwrap();
    /// assert_eq!("Doe", name.surname());
    /// assert_eq!(None, name.given_name());
    /// assert_eq!('J', name.first_initial());
    ///
    /// let name = Name::parse("Dr. Juan Alberto T. Velasquez y Garcia III").unwrap();
    /// assert_eq!("Velasquez y Garcia", name.surname());
    /// assert_eq!(Some("Juan"), name.given_name());
    /// assert_eq!(Some("AT"), name.middle_initials());
    /// assert_eq!(Some("III"), name.suffix());
    /// ```
    ///
    /// # Supported formats
    ///
    /// Supports a variety of formats, including prefix and postfix titles,
    /// parenthesized nicknames, initials with and without periods, and sort
    /// order ("Doe, Jane"). Makes use of heuristics based on case when
    /// applicable (e.g., "AL Doe" is parsed as "A. L. Doe", while "Al Doe" is
    /// parsed as a given name and surname), as well as _small_ sets of known
    /// particles, conjunctions, titles, etc.
    ///
    /// # Limitations
    ///
    /// Errs on the side of producing parse output rather than giving up, so
    /// this function is _not_ suitable as a way of guessing whether a given
    /// string actually represents a name.
    ///
    /// However, success requires at least an apparent surname and first initial.
    /// Single-word names cannot be parsed (you may or may not wish to assume
    /// they are given names).
    ///
    /// Does not preserve titles (other than generational suffixes such as "III")
    /// or nicknames. Does not handle plural forms specially: "Mr. & Mrs. John
    /// Doe" will be parsed as "John Doe", and "Jane Doe, et al" will be parsed
    /// as "Jane Doe".
    ///
    /// Works best on Latin names - i.e., data from North or South America or
    /// Europe. Does not understand surname-first formats without commas: "Kim
    /// Il-sung" will be parsed as having the first name "Kim".
    ///
    /// Handles non-Latin unicode strings, but without any particular intelligence.
    /// Attempts at least to fail nicely, such that either `parse` returns `None`,
    /// or calling `display_full()` on the parsed result returns the input,
    /// plus or minus whitespace.
    ///
    /// Of course, [there is no perfect algorithm](http://www.kalzumeus.com/2010/06/17/falsehoods-programmers-believe-about-names/)
    /// for canonicalizing names. The goal here is to do the best we can without
    /// large statistical models.
    pub fn parse(name: &str) -> Option<Name> {
        if name.len() >= 1000 || !name.chars().any(char::is_alphabetic) {
            return None;
        }

        let mixed_case = is_mixed_case(name);
        let name = nickname::strip_nickname(name);

        let result = parse::parse(&*name, mixed_case);
        if result.is_none() {
            return None;
        }

        let (words, surname_index, suffix_index) = result.unwrap();

        let mut names: Vec<String> = Vec::with_capacity(words.len());
        let mut initials = String::with_capacity(surname_index);
        let mut surname_index_in_names = surname_index;
        let mut suffix_index_in_names = suffix_index;
        let mut word_indices_in_initials: Vec<(usize,usize)> = Vec::with_capacity(surname_index);

        for (i, word) in words.into_iter().enumerate() {
            if word.is_initials() && i < surname_index {
                initials.extend(word.namecased.chars()
                    .filter(|c| c.is_alphabetic())
                    .flat_map(|c| c.to_uppercase()));

                surname_index_in_names -= 1;
                suffix_index_in_names -= 1;
            } else if i < surname_index {
                let prior_len = initials.len();

                initials.extend(word.namecased.split('-')
                    .filter_map(|w| w.chars().find(|c| c.is_alphabetic())));

                names.push(word.namecased.into_owned());
                word_indices_in_initials.push((prior_len, initials.len()));
            } else if i < suffix_index {
                names.push(word.namecased.into_owned());
            } else {
                names.push(suffix::namecase(&word));
            }
        }

        names.shrink_to_fit();
        word_indices_in_initials.shrink_to_fit();

        Some(Name {
            words: names,
            surname_index: surname_index_in_names,
            suffix_index: suffix_index_in_names,
            initials: initials,
            word_indices_in_initials: word_indices_in_initials,
        })
    }

    /// First initial (always present)
    pub fn first_initial(&self) -> char {
        self.initials.chars().nth(0).unwrap()
    }

    /// Given name as a string, if present
    pub fn given_name(&self) -> Option<&str> {
        if self.surname_index > 0 {
            Some(&*self.words[0])
        } else {
            None
        }
    }

    /// Does this person use a middle name in place of their given name (e.g., T. Boone Pickens)?
    pub fn goes_by_middle_name(&self) -> bool {
        !self.word_indices_in_initials.is_empty() && self.word_indices_in_initials[0].0 > 0
    }

    /// First and middle initials as a string (always present)
    pub fn initials(&self) -> &str {
        &self.initials
    }

    /// Middle names as an array of words, if present
    pub fn middle_names(&self) -> Option<&[String]> {
        if self.surname_index > 1 {
            Some(&self.words[1..self.surname_index])
        } else {
            None
        }
    }

    /// Middle names as a string, if present
    pub fn middle_name(&self) -> Option<Cow<str>> {
        match self.middle_names() {
            Some(words) => {
                if words.len() == 1 {
                    Some(Cow::Borrowed(&*words[0]))
                } else {
                    Some(Cow::Owned(words.join(" ")))
                }
            }
            None => None,
        }
    }

    /// Middle initials as a string, if present
    pub fn middle_initials(&self) -> Option<&str> {
        match self.initials().char_indices().skip(1).nth(0) {
            Some((i, _)) => Some(&self.initials[i..]),
            None => None,
        }
    }

    /// Surname as a slice of words (always present)
    pub fn surnames(&self) -> &[String] {
        &self.words[self.surname_index..self.suffix_index]
    }

    /// Surname as a string (always present)
    pub fn surname(&self) -> Cow<str> {
        if self.surnames().len() > 1 {
            Cow::Owned(self.surnames().join(" "))
        } else {
            Cow::Borrowed(&*self.surnames()[0])
        }
    }

    /// Generational suffix, if present
    pub fn suffix(&self) -> Option<&str> {
        if self.words.len() > self.suffix_index {
            Some(&*self.words[self.suffix_index])
        } else {
            None
        }
    }

    /// Given name and surname, if given name is known, otherwise first initial
    /// and surname.
    ///
    /// ```
    /// use human_name::Name;
    ///
    /// let name = Name::parse("JOHN ALLEN Q MACDONALD JR").unwrap();
    /// assert_eq!("John MacDonald", name.display_short());
    /// ```
    pub fn display_short(&self) -> String {
        match self.given_name() {
            Some(ref name) => {
                format!("{} {}", name, self.surname())
            }
            None => {
                format!("{}. {}", self.first_initial(), self.surname())
            }
        }
    }

    /// The full name, or as much of it as was preserved from the input,
    /// including given name, middle names, surname and suffix.
    ///
    /// ```
    /// use human_name::Name;
    ///
    /// let name = Name::parse("JOHN ALLEN Q MACDONALD JR").unwrap();
    /// assert_eq!("John Allen Q. MacDonald, Jr.", name.display_full());
    /// ```
    pub fn display_full(&self) -> String {
        let mut results: Vec<Cow<str>> = Vec::with_capacity(self.words.len());

        let mut initials = self.initials.chars().enumerate();
        let mut known_names = self.words[0..self.surname_index].iter();
        let mut known_name_indices = self.word_indices_in_initials.iter().peekable();

        loop {
            match initials.next() {
                Some((i,initial)) => {
                    let mut next_name = None;
                    if let Some(&&(j,k)) = known_name_indices.peek() {
                        if j == i {
                            known_name_indices.next();
                            next_name = known_names.next();

                            // Handle case of hyphenated name for which we have 2+ initials
                            let mut num_initials_for_name = k-j;
                            while num_initials_for_name > 1 {
                                initials.next();
                                num_initials_for_name -= 1;
                            }
                        }
                    }

                    if let Some(name) = next_name {
                        results.push(Cow::Borrowed(name));
                    } else {
                        results.push(Cow::Owned(format!("{}.", initial)));
                    }
                },
                None => {
                    break;
                }
            }
        }

        for surname in self.surnames() {
            results.push(Cow::Borrowed(&surname));
        }

        let mut result = results.join(" ");

        if let Some(suffix) = self.suffix() {
            result = result + &format!(", {}", suffix);
        }

        result
    }
}
<commit_msg>Factor out looping through given/middle names or initials<commit_after>#![feature(drain)]
#![feature(plugin)]
#![plugin(phf_macros)]

extern crate phf;
extern crate itertools;
extern crate unicode_segmentation;
extern crate unicode_normalization;
extern crate rustc_serialize;

mod utils;
mod suffix;
mod nickname;
mod title;
mod surname;
mod namecase;
mod namepart;
mod parse;
mod comparison;
mod serialization;

#[cfg(feature = "name_eq_hash")]
mod eq_hash;

use std::borrow::Cow;
use itertools::Itertools;
use utils::is_mixed_case;

pub struct Name {
    words: Vec<String>,
    surname_index: usize,
    suffix_index: usize,
    initials: String,
    word_indices_in_initials: Vec<(usize,usize)>,
}

enum NameWordOrInitial<'a> {
    Word(&'a str),
    Initial(char),
}

impl Name {

    /// Parses a string represent a single person's full name into a canonical
    /// representation.
    ///
    /// # Examples
    /// ```
    /// use human_name::Name;
    ///
    /// let name = Name::parse("Jane Doe").unwrap();
    /// assert_eq!("Doe", name.surname());
    /// assert_eq!(Some("Jane"), name.given_name());
    ///
    /// let name = Name::parse("Doe, J").unwrap();
    /// assert_eq!("Doe", name.surname());
    /// assert_eq!(None, name.given_name());
    /// assert_eq!('J', name.first_initial());
    ///
    /// let name = Name::parse("Dr. Juan Alberto T. Velasquez y Garcia III").unwrap();
    /// assert_eq!("Velasquez y Garcia", name.surname());
    /// assert_eq!(Some("Juan"), name.given_name());
    /// assert_eq!(Some("AT"), name.middle_initials());
    /// assert_eq!(Some("III"), name.suffix());
    /// ```
    ///
    /// # Supported formats
    ///
    /// Supports a variety of formats, including prefix and postfix titles,
    /// parenthesized nicknames, initials with and without periods, and sort
    /// order ("Doe, Jane"). Makes use of heuristics based on case when
    /// applicable (e.g., "AL Doe" is parsed as "A. L. Doe", while "Al Doe" is
    /// parsed as a given name and surname), as well as _small_ sets of known
    /// particles, conjunctions, titles, etc.
    ///
    /// # Limitations
    ///
    /// Errs on the side of producing parse output rather than giving up, so
    /// this function is _not_ suitable as a way of guessing whether a given
    /// string actually represents a name.
    ///
    /// However, success requires at least an apparent surname and first initial.
    /// Single-word names cannot be parsed (you may or may not wish to assume
    /// they are given names).
    ///
    /// Does not preserve titles (other than generational suffixes such as "III")
    /// or nicknames. Does not handle plural forms specially: "Mr. & Mrs. John
    /// Doe" will be parsed as "John Doe", and "Jane Doe, et al" will be parsed
    /// as "Jane Doe".
    ///
    /// Works best on Latin names - i.e., data from North or South America or
    /// Europe. Does not understand surname-first formats without commas: "Kim
    /// Il-sung" will be parsed as having the first name "Kim".
    ///
    /// Handles non-Latin unicode strings, but without any particular intelligence.
    /// Attempts at least to fail nicely, such that either `parse` returns `None`,
    /// or calling `display_full()` on the parsed result returns the input,
    /// plus or minus whitespace.
    ///
    /// Of course, [there is no perfect algorithm](http://www.kalzumeus.com/2010/06/17/falsehoods-programmers-believe-about-names/)
    /// for canonicalizing names. The goal here is to do the best we can without
    /// large statistical models.
    pub fn parse(name: &str) -> Option<Name> {
        if name.len() >= 1000 || !name.chars().any(char::is_alphabetic) {
            return None;
        }

        let mixed_case = is_mixed_case(name);
        let name = nickname::strip_nickname(name);

        let result = parse::parse(&*name, mixed_case);
        if result.is_none() {
            return None;
        }

        let (words, surname_index, suffix_index) = result.unwrap();

        let mut names: Vec<String> = Vec::with_capacity(words.len());
        let mut initials = String::with_capacity(surname_index);
        let mut surname_index_in_names = surname_index;
        let mut suffix_index_in_names = suffix_index;
        let mut word_indices_in_initials: Vec<(usize,usize)> = Vec::with_capacity(surname_index);

        for (i, word) in words.into_iter().enumerate() {
            if word.is_initials() && i < surname_index {
                initials.extend(word.namecased.chars()
                    .filter(|c| c.is_alphabetic())
                    .flat_map(|c| c.to_uppercase()));

                surname_index_in_names -= 1;
                suffix_index_in_names -= 1;
            } else if i < surname_index {
                let prior_len = initials.len();

                initials.extend(word.namecased.split('-')
                    .filter_map(|w| w.chars().find(|c| c.is_alphabetic())));

                names.push(word.namecased.into_owned());
                word_indices_in_initials.push((prior_len, initials.len()));
            } else if i < suffix_index {
                names.push(word.namecased.into_owned());
            } else {
                names.push(suffix::namecase(&word));
            }
        }

        names.shrink_to_fit();
        word_indices_in_initials.shrink_to_fit();

        Some(Name {
            words: names,
            surname_index: surname_index_in_names,
            suffix_index: suffix_index_in_names,
            initials: initials,
            word_indices_in_initials: word_indices_in_initials,
        })
    }

    /// First initial (always present)
    pub fn first_initial(&self) -> char {
        self.initials.chars().nth(0).unwrap()
    }

    /// Given name as a string, if present
    pub fn given_name(&self) -> Option<&str> {
        if self.surname_index > 0 {
            Some(&*self.words[0])
        } else {
            None
        }
    }

    /// Does this person use a middle name in place of their given name (e.g., T. Boone Pickens)?
    pub fn goes_by_middle_name(&self) -> bool {
        !self.word_indices_in_initials.is_empty() && self.word_indices_in_initials[0].0 > 0
    }

    /// First and middle initials as a string (always present)
    pub fn initials(&self) -> &str {
        &self.initials
    }

    /// Middle names as an array of words, if present
    pub fn middle_names(&self) -> Option<&[String]> {
        if self.surname_index > 1 {
            Some(&self.words[1..self.surname_index])
        } else {
            None
        }
    }

    /// Middle names as a string, if present
    pub fn middle_name(&self) -> Option<Cow<str>> {
        match self.middle_names() {
            Some(words) => {
                if words.len() == 1 {
                    Some(Cow::Borrowed(&*words[0]))
                } else {
                    Some(Cow::Owned(words.join(" ")))
                }
            }
            None => None,
        }
    }

    /// Middle initials as a string, if present
    pub fn middle_initials(&self) -> Option<&str> {
        match self.initials().char_indices().skip(1).nth(0) {
            Some((i, _)) => Some(&self.initials[i..]),
            None => None,
        }
    }

    /// Surname as a slice of words (always present)
    pub fn surnames(&self) -> &[String] {
        &self.words[self.surname_index..self.suffix_index]
    }

    /// Surname as a string (always present)
    pub fn surname(&self) -> Cow<str> {
        if self.surnames().len() > 1 {
            Cow::Owned(self.surnames().join(" "))
        } else {
            Cow::Borrowed(&*self.surnames()[0])
        }
    }

    /// Generational suffix, if present
    pub fn suffix(&self) -> Option<&str> {
        if self.words.len() > self.suffix_index {
            Some(&*self.words[self.suffix_index])
        } else {
            None
        }
    }

    fn with_each_given_name_or_initial<F: FnMut(NameWordOrInitial, usize)>(&self, cb: &mut F) {
        let mut initials = self.initials.chars().enumerate();
        let mut known_names = self.words[0..self.surname_index].iter();
        let mut known_name_indices = self.word_indices_in_initials.iter().peekable();

        loop {
            match initials.next() {
                Some((i,initial)) => {
                    let mut next_name = None;
                    if let Some(&&(j,k)) = known_name_indices.peek() {
                        if j == i {
                            known_name_indices.next();
                            next_name = known_names.next();

                            // Handle case of hyphenated name for which we have 2+ initials
                            let mut num_initials_for_name = k-j;
                            while num_initials_for_name > 1 {
                                initials.next();
                                num_initials_for_name -= 1;
                            }
                        }
                    }

                    if let Some(name) = next_name {
                        cb(NameWordOrInitial::Word(name), i);
                    } else {
                        cb(NameWordOrInitial::Initial(initial), i);
                    }
                },
                None => {
                    break;
                }
            }
        }
    }

    /// Given name and surname, if given name is known, otherwise first initial
    /// and surname.
    ///
    /// ```
    /// use human_name::Name;
    ///
    /// let name = Name::parse("JOHN ALLEN Q DE LA MACDONALD JR").unwrap();
    /// assert_eq!("John de la MacDonald", name.display_short());
    /// ```
    pub fn display_short(&self) -> String {
        match self.given_name() {
            Some(ref name) => {
                format!("{} {}", name, self.surname())
            }
            None => {
                format!("{}. {}", self.first_initial(), self.surname())
            }
        }
    }

    /// The full name, or as much of it as was preserved from the input,
    /// including given name, middle names, surname and suffix.
    ///
    /// ```
    /// use human_name::Name;
    ///
    /// let name = Name::parse("JOHN ALLEN Q DE LA MACDONALD JR").unwrap();
    /// assert_eq!("John Allen Q. de la MacDonald, Jr.", name.display_full());
    /// ```
    pub fn display_full(&self) -> String {
        // This will be correct assuming only ASCII and only words, no initials,
        // no suffix, otherwise it'll be too short, which is ok
        let min_len = self.words.iter().fold(self.words.len()-1, |sum,ref word| sum + word.len());

        let mut result = String::with_capacity(min_len);

        self.with_each_given_name_or_initial( &mut |part, _|
            match part {
                NameWordOrInitial::Word(name) => {
                    result.push_str(name);
                    result.push(' ');
                },
                NameWordOrInitial::Initial(initial) => {
                    result.push(initial);
                    result.push_str(". ");
                },
            }
        );

        let surnames = self.surnames();
        if surnames.len() > 1 {
            for word in surnames[0..surnames.len() - 1].iter() {
                result.push_str(word);
                result.push(' ');
            }
        }
        result.push_str(&surnames[surnames.len() - 1]);

        if let Some(suffix) = self.suffix() {
            result.push_str(", ");
            result.push_str(suffix);
        }

        result
    }
}
<|endoftext|>


<commit_before>mod alias_dist;

use alias_dist::AliasDistribution;
use std::collections::HashMap;
use std::hash::Hash;

#[derive(Debug)]
pub struct MarkovNode<T: Hash + Eq + Clone> {
    pub value: T,
    dist: AliasDistribution<T>,
}

impl<T: Hash + Eq + Clone> MarkovNode<T> {
    pub fn new(value: T, counts: HashMap<T, usize>) -> MarkovNode<T> {
        MarkovNode {
            value: value.clone(),
            dist: AliasDistribution::new(counts),
        }
    }

    pub fn next(&self) -> &T {
        self.dist.choice().unwrap()
    }

    pub fn entropy(&self) -> f64 {
        self.dist.entropy
    }
}

#[derive(Debug)]
pub struct PassphraseMarkovChain {
    nodes: HashMap<String, MarkovNode<String>>,
    starting_dist: AliasDistribution<String>,
}

impl PassphraseMarkovChain {
    pub fn new<U: Clone>(ngrams: U) -> Result<PassphraseMarkovChain, &'static str>
            where U: Iterator<Item=String> {
        let mut transition_map = HashMap::new();
        let mut starting_counts = HashMap::new();
        let mut ngrams_copy = ngrams.clone().cycle();
        ngrams_copy.next();
        for (a, b) in ngrams.zip(ngrams_copy) {
            if b.starts_with(" ") {
                let count = starting_counts.entry(b.clone()).or_insert(0);
                *count += 1
            }
            let transitions = transition_map.entry(a).or_insert(HashMap::new());
            let count = transitions.entry(b).or_insert(0);
            *count += 1;
        };

        let mut total_entropy: f64 = 0.0;
        let mut nodes = HashMap::new();
        for (ngram, transition_counts) in transition_map.into_iter() {
            let node = MarkovNode::new(ngram.clone(), transition_counts);
            total_entropy += node.entropy();
            nodes.insert(ngram.clone(), node);
        };
        let starting_dist = AliasDistribution::new(starting_counts);

        if total_entropy == 0.0 {
            return Err("No entropy found in input.");
        } else if starting_dist.entropy == 0.0 {
            return Err("No start of word entropy found in input.");
        };

        Ok(PassphraseMarkovChain {
            nodes: nodes,
            starting_dist: starting_dist,
        })
    }

    pub fn get_node(&self) -> &MarkovNode<String> {
        self.nodes.get(self.starting_dist.choice().unwrap()).unwrap()
    }

    pub fn passphrase(&self, min_entropy: f64) -> (String, f64) {
        let mut node = self.get_node();
        let mut entropy = self.starting_dist.entropy;
        let mut nodes = Vec::new();
        loop {
            nodes.push(node);
            entropy += node.entropy();
            if entropy >= min_entropy && node.value.ends_with(" ") { break };
            node = self.nodes.get(node.next()).unwrap();
        }
        let tail = nodes.last().unwrap().value.chars().skip(1);
        let chars = nodes.iter().map(|n| n.value.chars().next().unwrap()).chain(tail);
        let passphrase = chars.collect::<String>().trim().to_string();

        (passphrase, entropy)
    }
}
<commit_msg>Another minor style change<commit_after>mod alias_dist;

use alias_dist::AliasDistribution;
use std::collections::HashMap;
use std::hash::Hash;

#[derive(Debug)]
pub struct MarkovNode<T: Hash + Eq + Clone> {
    pub value: T,
    dist: AliasDistribution<T>,
}

impl<T: Hash + Eq + Clone> MarkovNode<T> {
    pub fn new(value: T, counts: HashMap<T, usize>) -> MarkovNode<T> {
        MarkovNode {
            value: value.clone(),
            dist: AliasDistribution::new(counts),
        }
    }

    pub fn next(&self) -> &T {
        self.dist.choice().unwrap()
    }

    pub fn entropy(&self) -> f64 {
        self.dist.entropy
    }
}

#[derive(Debug)]
pub struct PassphraseMarkovChain {
    nodes: HashMap<String, MarkovNode<String>>,
    starting_dist: AliasDistribution<String>,
}

impl PassphraseMarkovChain {
    pub fn new<U: Iterator<Item=String> + Clone>(ngrams: U)
            -> Result<PassphraseMarkovChain, &'static str> {
        let mut transition_map = HashMap::new();
        let mut starting_counts = HashMap::new();
        let mut ngrams_copy = ngrams.clone().cycle();
        ngrams_copy.next();
        for (a, b) in ngrams.zip(ngrams_copy) {
            if b.starts_with(" ") {
                let count = starting_counts.entry(b.clone()).or_insert(0);
                *count += 1
            }
            let transitions = transition_map.entry(a).or_insert(HashMap::new());
            let count = transitions.entry(b).or_insert(0);
            *count += 1;
        };

        let mut total_entropy: f64 = 0.0;
        let mut nodes = HashMap::new();
        for (ngram, transition_counts) in transition_map.into_iter() {
            let node = MarkovNode::new(ngram.clone(), transition_counts);
            total_entropy += node.entropy();
            nodes.insert(ngram.clone(), node);
        };
        let starting_dist = AliasDistribution::new(starting_counts);

        if total_entropy == 0.0 {
            return Err("No entropy found in input.");
        } else if starting_dist.entropy == 0.0 {
            return Err("No start of word entropy found in input.");
        };

        Ok(PassphraseMarkovChain {
            nodes: nodes,
            starting_dist: starting_dist,
        })
    }

    pub fn get_node(&self) -> &MarkovNode<String> {
        self.nodes.get(self.starting_dist.choice().unwrap()).unwrap()
    }

    pub fn passphrase(&self, min_entropy: f64) -> (String, f64) {
        let mut node = self.get_node();
        let mut entropy = self.starting_dist.entropy;
        let mut nodes = Vec::new();
        loop {
            nodes.push(node);
            entropy += node.entropy();
            if entropy >= min_entropy && node.value.ends_with(" ") { break };
            node = self.nodes.get(node.next()).unwrap();
        }
        let tail = nodes.last().unwrap().value.chars().skip(1);
        let chars = nodes.iter().map(|n| n.value.chars().next().unwrap()).chain(tail);
        let passphrase = chars.collect::<String>().trim().to_string();

        (passphrase, entropy)
    }
}
<|endoftext|>


<commit_before>///! # multiaddr
///!
///! Implementation of [multiaddr](https://github.com/jbenet/multiaddr)
///! in Rust.
extern crate byteorder;
extern crate cid;
extern crate integer_encoding;

mod protocol;
mod errors;

pub use errors::{Result, Error};
pub use protocol::{Protocol, ProtocolArgSize, AddrComponent};

use std::fmt;
use std::iter::FromIterator;
use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6, IpAddr, Ipv4Addr, Ipv6Addr};
use std::str::FromStr;

/// Representation of a Multiaddr.
#[derive(PartialEq, Eq, Clone, Hash)]
pub struct Multiaddr {
    bytes: Vec<u8>,
}

impl fmt::Debug for Multiaddr {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.to_string().fmt(f)
    }
}

impl fmt::Display for Multiaddr {
    /// Convert a Multiaddr to a string
    ///
    /// # Examples
    ///
    /// ```
    /// use multiaddr::Multiaddr;
    ///
    /// let address: Multiaddr = "/ip4/127.0.0.1/udt".parse().unwrap();
    /// assert_eq!(address.to_string(), "/ip4/127.0.0.1/udt");
    /// ```
    ///
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for s in self.iter() {
            s.to_string().fmt(f)?;
        }
        Ok(())
    }
}

impl Multiaddr {
    /// Returns the raw bytes representation of the multiaddr.
    #[inline]
    pub fn into_bytes(self) -> Vec<u8> {
        self.bytes
    }

    /// Return a copy to disallow changing the bytes directly
    pub fn to_bytes(&self) -> Vec<u8> {
        self.bytes.to_owned()
    }

    /// Produces a `Multiaddr` from its bytes representation.
    pub fn from_bytes(bytes: Vec<u8>) -> Result<Multiaddr> {
        {
            let mut ptr = &bytes[..];
            while !ptr.is_empty() {
                let (_, new_ptr) = AddrComponent::from_bytes(ptr)?;
                ptr = new_ptr;
            }
        }

        Ok(Multiaddr { bytes })
    }

    /// Extracts a slice containing the entire underlying vector.
    pub fn as_slice(&self) -> &[u8] {
        &self.bytes
    }

    /// Return a list of protocols
    ///
    /// # Examples
    ///
    /// A single protocol
    ///
    /// ```
    /// use multiaddr::{Multiaddr, Protocol};
    ///
    /// let address: Multiaddr = "/ip4/127.0.0.1".parse().unwrap();
    /// assert_eq!(address.protocol(), vec![Protocol::IP4]);
    /// ```
    ///
    #[inline]
    #[deprecated(note = "Use `self.iter().map(|addr| addr.protocol_id())` instead")]
    pub fn protocol(&self) -> Vec<Protocol> {
        self.iter().map(|addr| addr.protocol_id()).collect()
    }

    /// Wrap a given Multiaddr and return the combination.
    ///
    /// # Examples
    ///
    /// ```
    /// use multiaddr::Multiaddr;
    ///
    /// let address: Multiaddr = "/ip4/127.0.0.1".parse().unwrap();
    /// let nested = address.encapsulate("/udt").unwrap();
    /// assert_eq!(nested, "/ip4/127.0.0.1/udt".parse().unwrap());
    /// ```
    ///
    pub fn encapsulate<T: ToMultiaddr>(&self, input: T) -> Result<Multiaddr> {
        let new = input.to_multiaddr()?;
        let mut bytes = self.bytes.clone();

        bytes.extend(new.to_bytes());

        Ok(Multiaddr { bytes: bytes })
    }

    /// Adds an already-parsed address component to the end of this multiaddr.
    ///
    /// # Examples
    ///
    /// ```
    /// use multiaddr::{Multiaddr, AddrComponent};
    ///
    /// let mut address: Multiaddr = "/ip4/127.0.0.1".parse().unwrap();
    /// address.append(AddrComponent::TCP(10000));
    /// assert_eq!(address, "/ip4/127.0.0.1/tcp/10000".parse().unwrap());
    /// ```
    ///
    #[inline]
    pub fn append(&mut self, component: AddrComponent) {
        component.write_bytes(&mut self.bytes).expect("writing to a Vec never fails")
    }

    /// Remove the outermost address.
    ///
    /// # Examples
    ///
    /// ```
    /// use multiaddr::{Multiaddr, ToMultiaddr};
    ///
    /// let address: Multiaddr = "/ip4/127.0.0.1/udt/sctp/5678".parse().unwrap();
    /// let unwrapped = address.decapsulate("/udt").unwrap();
    /// assert_eq!(unwrapped, "/ip4/127.0.0.1".parse().unwrap());
    ///
    /// assert_eq!(
    ///     address.decapsulate("/udt").unwrap(),
    ///     "/ip4/127.0.0.1".to_multiaddr().unwrap()
    /// );
    /// ```
    ///
    /// Returns the original if the passed in address is not found
    ///
    /// ```
    /// use multiaddr::ToMultiaddr;
    ///
    /// let address = "/ip4/127.0.0.1/udt/sctp/5678".to_multiaddr().unwrap();
    /// let unwrapped = address.decapsulate("/ip4/127.0.1.1").unwrap();
    /// assert_eq!(unwrapped, address);
    /// ```
    ///
    pub fn decapsulate<T: ToMultiaddr>(&self, input: T) -> Result<Multiaddr> {
        let input = input.to_multiaddr()?.to_bytes();

        let bytes_len = self.bytes.len();
        let input_length = input.len();

        let mut input_pos = 0;
        let mut matches = false;

        for (i, _) in self.bytes.iter().enumerate() {
            let next = i + input_length;

            if next > bytes_len {
                continue;
            }

            if &self.bytes[i..next] == input.as_slice() {
                matches = true;
                input_pos = i;
                break;
            }
        }

        if !matches {
            return Ok(Multiaddr { bytes: self.bytes.clone() });
        }

        let mut bytes = self.bytes.clone();
        bytes.truncate(input_pos);

        Ok(Multiaddr { bytes: bytes })
    }

    /// Returns the components of this multiaddress.
    ///
    /// ```
    /// use std::net::Ipv4Addr;
    /// use multiaddr::AddrComponent;
    /// use multiaddr::Multiaddr;
    ///
    /// let address: Multiaddr = "/ip4/127.0.0.1/udt/sctp/5678".parse().unwrap();
    ///
    /// let components = address.iter().collect::<Vec<_>>();
    /// assert_eq!(components[0], AddrComponent::IP4(Ipv4Addr::new(127, 0, 0, 1)));
    /// assert_eq!(components[1], AddrComponent::UDT);
    /// assert_eq!(components[2], AddrComponent::SCTP(5678));
    /// ```
    ///
    #[inline]
    pub fn iter(&self) -> Iter {
        Iter(&self.bytes)
    }

    /// Pops the last `AddrComponent` of this multiaddr, or `None` if the multiaddr is empty.
    /// ```
    /// use multiaddr::AddrComponent;
    /// use multiaddr::Multiaddr;
    ///
    /// let address: Multiaddr = "/ip4/127.0.0.1/udt/sctp/5678".parse().unwrap();
    ///
    /// assert_eq!(address.pop().unwrap(), AddrComponent::SCTP(5678));
    /// assert_eq!(address.pop().unwrap(), AddrComponent::UDT);
    /// ```
    ///
    pub fn pop(&mut self) -> Option<AddrComponent> {
        // Note: could be more optimized
        let mut list = self.iter().collect::<Vec<_>>();
        let last_elem = list.pop();
        *self = list.into_iter().collect();
        last_elem
    }
}

impl From<AddrComponent> for Multiaddr {
    fn from(addr: AddrComponent) -> Multiaddr {
        let mut out = Vec::new();
        addr.write_bytes(&mut out).expect("writing to a Vec never fails");
        Multiaddr { bytes: out }
    }
}

impl<'a> IntoIterator for &'a Multiaddr {
    type Item = AddrComponent;
    type IntoIter = Iter<'a>;

    #[inline]
    fn into_iter(self) -> Iter<'a> {
        Iter(&self.bytes)
    }
}

impl FromIterator<AddrComponent> for Multiaddr {
    fn from_iter<T>(iter: T) -> Self
        where T: IntoIterator<Item = AddrComponent>
    {
        let mut bytes = Vec::new();
        for cmp in iter {
            cmp.write_bytes(&mut bytes).expect("writing to a Vec never fails");
        }
        Multiaddr { bytes: bytes }
    }
}

impl FromStr for Multiaddr {
    type Err = Error;

    #[inline]
    fn from_str(input: &str) -> Result<Self> {
        let mut bytes = Vec::new();

        let mut parts = input.split('/');
        // A multiaddr must start with `/`
        if !parts.next().ok_or(Error::InvalidMultiaddr)?.is_empty() {
            return Err(Error::InvalidMultiaddr);
        }

        while let Some(part) = parts.next() {
            let protocol: Protocol = part.parse()?;
            let addr_component = match protocol.size() {
                ProtocolArgSize::Fixed { bytes: 0 } => {
                    protocol.parse_data("")?     // TODO: bad design
                },
                _ => {
                    let data = parts.next().ok_or(Error::MissingAddress)?;
                    protocol.parse_data(data)?
                },
            };

            addr_component.write_bytes(&mut bytes).expect("writing to a Vec never fails");
        }

        Ok(Multiaddr { bytes: bytes })
    }
}

/// Iterator for the address components in a multiaddr.
pub struct Iter<'a>(&'a [u8]);

impl<'a> Iterator for Iter<'a> {
    type Item = AddrComponent;

    fn next(&mut self) -> Option<AddrComponent> {
        if self.0.is_empty() {
            return None;
        }

        let (component, next_data) = AddrComponent::from_bytes(self.0)
                                                        .expect("multiaddr is known to be valid");
        self.0 = next_data;
        Some(component)
    }
}

/// A trait for objects which can be converted to a
/// Multiaddr.
///
/// This trait is implemented by default for
///
/// * `SocketAddr`, `SocketAddrV4` and `SocketAddrV6`, assuming that the
///   the given port is a tcp port.
///
/// * `Ipv4Addr`, `Ipv6Addr`
///
/// * `String` and `&str`, requiring the default string format for a Multiaddr.
///
pub trait ToMultiaddr {
    /// Converts this object to a Multiaddr
    ///
    /// # Errors
    ///
    /// Any errors encountered during parsing will be returned
    /// as an `Err`.
    fn to_multiaddr(&self) -> Result<Multiaddr>;
}

impl ToMultiaddr for SocketAddr {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        match *self {
            SocketAddr::V4(ref a) => (*a).to_multiaddr(),
            SocketAddr::V6(ref a) => (*a).to_multiaddr(),
        }
    }
}

impl ToMultiaddr for SocketAddrV4 {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        format!("/ip4/{}/tcp/{}", self.ip(), self.port()).parse()
    }
}

impl ToMultiaddr for SocketAddrV6 {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        // TODO: Should how should we handle `flowinfo` and `scope_id`?
        format!("/ip6/{}/tcp/{}", self.ip(), self.port()).parse()
    }
}

impl ToMultiaddr for IpAddr {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        match *self {
            IpAddr::V4(ref a) => (*a).to_multiaddr(),
            IpAddr::V6(ref a) => (*a).to_multiaddr(),
        }
    }
}

impl ToMultiaddr for Ipv4Addr {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        format!("/ip4/{}", &self).parse()
    }
}

impl ToMultiaddr for Ipv6Addr {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        format!("/ip6/{}", &self).parse()
    }
}

impl ToMultiaddr for String {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        self.parse()
    }
}

impl<'a> ToMultiaddr for &'a str {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        self.parse()
    }
}

impl ToMultiaddr for Multiaddr {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        Ok(self.clone())
    }
}
<commit_msg>fix: doc test in beta and stable<commit_after>///! # multiaddr
///!
///! Implementation of [multiaddr](https://github.com/jbenet/multiaddr)
///! in Rust.
extern crate byteorder;
extern crate cid;
extern crate integer_encoding;

mod protocol;
mod errors;

pub use errors::{Result, Error};
pub use protocol::{Protocol, ProtocolArgSize, AddrComponent};

use std::fmt;
use std::iter::FromIterator;
use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6, IpAddr, Ipv4Addr, Ipv6Addr};
use std::str::FromStr;

/// Representation of a Multiaddr.
#[derive(PartialEq, Eq, Clone, Hash)]
pub struct Multiaddr {
    bytes: Vec<u8>,
}

impl fmt::Debug for Multiaddr {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.to_string().fmt(f)
    }
}

impl fmt::Display for Multiaddr {
    /// Convert a Multiaddr to a string
    ///
    /// # Examples
    ///
    /// ```
    /// use multiaddr::Multiaddr;
    ///
    /// let address: Multiaddr = "/ip4/127.0.0.1/udt".parse().unwrap();
    /// assert_eq!(address.to_string(), "/ip4/127.0.0.1/udt");
    /// ```
    ///
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for s in self.iter() {
            s.to_string().fmt(f)?;
        }
        Ok(())
    }
}

impl Multiaddr {
    /// Returns the raw bytes representation of the multiaddr.
    #[inline]
    pub fn into_bytes(self) -> Vec<u8> {
        self.bytes
    }

    /// Return a copy to disallow changing the bytes directly
    pub fn to_bytes(&self) -> Vec<u8> {
        self.bytes.to_owned()
    }

    /// Produces a `Multiaddr` from its bytes representation.
    pub fn from_bytes(bytes: Vec<u8>) -> Result<Multiaddr> {
        {
            let mut ptr = &bytes[..];
            while !ptr.is_empty() {
                let (_, new_ptr) = AddrComponent::from_bytes(ptr)?;
                ptr = new_ptr;
            }
        }

        Ok(Multiaddr { bytes })
    }

    /// Extracts a slice containing the entire underlying vector.
    pub fn as_slice(&self) -> &[u8] {
        &self.bytes
    }

    /// Return a list of protocols
    ///
    /// # Examples
    ///
    /// A single protocol
    ///
    /// ```
    /// use multiaddr::{Multiaddr, Protocol};
    ///
    /// let address: Multiaddr = "/ip4/127.0.0.1".parse().unwrap();
    /// assert_eq!(address.protocol(), vec![Protocol::IP4]);
    /// ```
    ///
    #[inline]
    #[deprecated(note = "Use `self.iter().map(|addr| addr.protocol_id())` instead")]
    pub fn protocol(&self) -> Vec<Protocol> {
        self.iter().map(|addr| addr.protocol_id()).collect()
    }

    /// Wrap a given Multiaddr and return the combination.
    ///
    /// # Examples
    ///
    /// ```
    /// use multiaddr::Multiaddr;
    ///
    /// let address: Multiaddr = "/ip4/127.0.0.1".parse().unwrap();
    /// let nested = address.encapsulate("/udt").unwrap();
    /// assert_eq!(nested, "/ip4/127.0.0.1/udt".parse().unwrap());
    /// ```
    ///
    pub fn encapsulate<T: ToMultiaddr>(&self, input: T) -> Result<Multiaddr> {
        let new = input.to_multiaddr()?;
        let mut bytes = self.bytes.clone();

        bytes.extend(new.to_bytes());

        Ok(Multiaddr { bytes: bytes })
    }

    /// Adds an already-parsed address component to the end of this multiaddr.
    ///
    /// # Examples
    ///
    /// ```
    /// use multiaddr::{Multiaddr, AddrComponent};
    ///
    /// let mut address: Multiaddr = "/ip4/127.0.0.1".parse().unwrap();
    /// address.append(AddrComponent::TCP(10000));
    /// assert_eq!(address, "/ip4/127.0.0.1/tcp/10000".parse().unwrap());
    /// ```
    ///
    #[inline]
    pub fn append(&mut self, component: AddrComponent) {
        component.write_bytes(&mut self.bytes).expect(
            "writing to a Vec never fails",
        )
    }

    /// Remove the outermost address.
    ///
    /// # Examples
    ///
    /// ```
    /// use multiaddr::{Multiaddr, ToMultiaddr};
    ///
    /// let address: Multiaddr = "/ip4/127.0.0.1/udt/sctp/5678".parse().unwrap();
    /// let unwrapped = address.decapsulate("/udt").unwrap();
    /// assert_eq!(unwrapped, "/ip4/127.0.0.1".parse().unwrap());
    ///
    /// assert_eq!(
    ///     address.decapsulate("/udt").unwrap(),
    ///     "/ip4/127.0.0.1".to_multiaddr().unwrap()
    /// );
    /// ```
    ///
    /// Returns the original if the passed in address is not found
    ///
    /// ```
    /// use multiaddr::ToMultiaddr;
    ///
    /// let address = "/ip4/127.0.0.1/udt/sctp/5678".to_multiaddr().unwrap();
    /// let unwrapped = address.decapsulate("/ip4/127.0.1.1").unwrap();
    /// assert_eq!(unwrapped, address);
    /// ```
    ///
    pub fn decapsulate<T: ToMultiaddr>(&self, input: T) -> Result<Multiaddr> {
        let input = input.to_multiaddr()?.to_bytes();

        let bytes_len = self.bytes.len();
        let input_length = input.len();

        let mut input_pos = 0;
        let mut matches = false;

        for (i, _) in self.bytes.iter().enumerate() {
            let next = i + input_length;

            if next > bytes_len {
                continue;
            }

            if &self.bytes[i..next] == input.as_slice() {
                matches = true;
                input_pos = i;
                break;
            }
        }

        if !matches {
            return Ok(Multiaddr { bytes: self.bytes.clone() });
        }

        let mut bytes = self.bytes.clone();
        bytes.truncate(input_pos);

        Ok(Multiaddr { bytes: bytes })
    }

    /// Returns the components of this multiaddress.
    ///
    /// ```
    /// use std::net::Ipv4Addr;
    /// use multiaddr::AddrComponent;
    /// use multiaddr::Multiaddr;
    ///
    /// let address: Multiaddr = "/ip4/127.0.0.1/udt/sctp/5678".parse().unwrap();
    ///
    /// let components = address.iter().collect::<Vec<_>>();
    /// assert_eq!(components[0], AddrComponent::IP4(Ipv4Addr::new(127, 0, 0, 1)));
    /// assert_eq!(components[1], AddrComponent::UDT);
    /// assert_eq!(components[2], AddrComponent::SCTP(5678));
    /// ```
    ///
    #[inline]
    pub fn iter(&self) -> Iter {
        Iter(&self.bytes)
    }

    /// Pops the last `AddrComponent` of this multiaddr, or `None` if the multiaddr is empty.
    /// ```
    /// use multiaddr::AddrComponent;
    /// use multiaddr::Multiaddr;
    ///
    /// let mut address: Multiaddr = "/ip4/127.0.0.1/udt/sctp/5678".parse().unwrap();
    ///
    /// assert_eq!(address.pop().unwrap(), AddrComponent::SCTP(5678));
    /// assert_eq!(address.pop().unwrap(), AddrComponent::UDT);
    /// ```
    ///
    pub fn pop(&mut self) -> Option<AddrComponent> {
        // Note: could be more optimized
        let mut list = self.iter().collect::<Vec<_>>();
        let last_elem = list.pop();
        *self = list.into_iter().collect();
        last_elem
    }
}

impl From<AddrComponent> for Multiaddr {
    fn from(addr: AddrComponent) -> Multiaddr {
        let mut out = Vec::new();
        addr.write_bytes(&mut out).expect(
            "writing to a Vec never fails",
        );
        Multiaddr { bytes: out }
    }
}

impl<'a> IntoIterator for &'a Multiaddr {
    type Item = AddrComponent;
    type IntoIter = Iter<'a>;

    #[inline]
    fn into_iter(self) -> Iter<'a> {
        Iter(&self.bytes)
    }
}

impl FromIterator<AddrComponent> for Multiaddr {
    fn from_iter<T>(iter: T) -> Self
    where
        T: IntoIterator<Item = AddrComponent>,
    {
        let mut bytes = Vec::new();
        for cmp in iter {
            cmp.write_bytes(&mut bytes).expect(
                "writing to a Vec never fails",
            );
        }
        Multiaddr { bytes: bytes }
    }
}

impl FromStr for Multiaddr {
    type Err = Error;

    #[inline]
    fn from_str(input: &str) -> Result<Self> {
        let mut bytes = Vec::new();

        let mut parts = input.split('/');
        // A multiaddr must start with `/`
        if !parts.next().ok_or(Error::InvalidMultiaddr)?.is_empty() {
            return Err(Error::InvalidMultiaddr);
        }

        while let Some(part) = parts.next() {
            let protocol: Protocol = part.parse()?;
            let addr_component = match protocol.size() {
                ProtocolArgSize::Fixed { bytes: 0 } => {
                    protocol.parse_data("")? // TODO: bad design
                }
                _ => {
                    let data = parts.next().ok_or(Error::MissingAddress)?;
                    protocol.parse_data(data)?
                }
            };

            addr_component.write_bytes(&mut bytes).expect(
                "writing to a Vec never fails",
            );
        }

        Ok(Multiaddr { bytes: bytes })
    }
}

/// Iterator for the address components in a multiaddr.
pub struct Iter<'a>(&'a [u8]);

impl<'a> Iterator for Iter<'a> {
    type Item = AddrComponent;

    fn next(&mut self) -> Option<AddrComponent> {
        if self.0.is_empty() {
            return None;
        }

        let (component, next_data) =
            AddrComponent::from_bytes(self.0).expect("multiaddr is known to be valid");
        self.0 = next_data;
        Some(component)
    }
}

/// A trait for objects which can be converted to a
/// Multiaddr.
///
/// This trait is implemented by default for
///
/// * `SocketAddr`, `SocketAddrV4` and `SocketAddrV6`, assuming that the
///   the given port is a tcp port.
///
/// * `Ipv4Addr`, `Ipv6Addr`
///
/// * `String` and `&str`, requiring the default string format for a Multiaddr.
///
pub trait ToMultiaddr {
    /// Converts this object to a Multiaddr
    ///
    /// # Errors
    ///
    /// Any errors encountered during parsing will be returned
    /// as an `Err`.
    fn to_multiaddr(&self) -> Result<Multiaddr>;
}

impl ToMultiaddr for SocketAddr {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        match *self {
            SocketAddr::V4(ref a) => (*a).to_multiaddr(),
            SocketAddr::V6(ref a) => (*a).to_multiaddr(),
        }
    }
}

impl ToMultiaddr for SocketAddrV4 {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        format!("/ip4/{}/tcp/{}", self.ip(), self.port()).parse()
    }
}

impl ToMultiaddr for SocketAddrV6 {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        // TODO: Should how should we handle `flowinfo` and `scope_id`?
        format!("/ip6/{}/tcp/{}", self.ip(), self.port()).parse()
    }
}

impl ToMultiaddr for IpAddr {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        match *self {
            IpAddr::V4(ref a) => (*a).to_multiaddr(),
            IpAddr::V6(ref a) => (*a).to_multiaddr(),
        }
    }
}

impl ToMultiaddr for Ipv4Addr {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        format!("/ip4/{}", &self).parse()
    }
}

impl ToMultiaddr for Ipv6Addr {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        format!("/ip6/{}", &self).parse()
    }
}

impl ToMultiaddr for String {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        self.parse()
    }
}

impl<'a> ToMultiaddr for &'a str {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        self.parse()
    }
}

impl ToMultiaddr for Multiaddr {
    fn to_multiaddr(&self) -> Result<Multiaddr> {
        Ok(self.clone())
    }
}
<|endoftext|>


<commit_before>#![crate_name = "jsonway"]
#![comment = "JSON builder for Rust"]
#![license = "MIT"]
#![crate_type = "rlib"]
// #![deny(warnings)]
// #![deny(bad_style)]
#![feature(macro_rules, phase, tuple_indexing)]

extern crate serialize;

use std::collections::TreeMap;
use std::to_string::ToString;
use serialize::json;
use serialize::json::{JsonList, Json, JsonObject, ToJson};

pub mod mutable_json;

trait Builder {

}

struct ListBuilder {
    list: JsonList,
    null: bool
}

impl ListBuilder {

    pub fn new() -> ListBuilder {
        ListBuilder { list: vec![], null: false }
    }

    pub fn from_json(list: Json) -> Option<ListBuilder> {
        match list {
            json::List(list) => Some(ListBuilder { list: list, null: false }),
            _ => None
        }
    }

    pub fn build(builder: |&mut ListBuilder|) -> ListBuilder {
        let mut bldr = ListBuilder::new();
        builder(&mut bldr);

        bldr
    }

    pub fn move_to_json(self) -> Json {
        if self.null {
            json::Null
        } else {
            json::List(self.list)
        }
    }
}

impl<T: ToJson> ListBuilder {
    pub fn push(&mut self, value: T) {
        self.list.push(value.to_json());
    }

    pub fn list(&mut self, builder: |&mut ListBuilder|) {
        self.push(ListBuilder::build(builder).move_to_json());
    }

    pub fn object(&mut self, builder: |&mut ObjectBuilder|) {
        self.push(ObjectBuilder::build(builder).move_to_json());
    }

    pub fn null(&mut self) {
        self.null = true;
    }
}

impl ToJson for ListBuilder {
    fn to_json(&self) -> Json {
         if self.null { json::Null } else { self.list.to_json() }
    }
}

struct ObjectBuilder {
    object: JsonObject,
    null: bool
}

impl ObjectBuilder {
    pub fn new() -> ObjectBuilder {
        ObjectBuilder { object: TreeMap::new(), null: false }
    }

    pub fn from_json(object: Json) -> Option<ObjectBuilder> {
        match object {
            json::Object(object) => Some(ObjectBuilder { object: object, null: false }),
            _ => None
        }
    }

    pub fn build(builder: |&mut ObjectBuilder|) -> ObjectBuilder {
        let mut bldr = ObjectBuilder::new();
        builder(&mut bldr);

        bldr
    }

    pub fn move_to_json(self) -> Json {
        if self.null {
            json::Null
        } else {
            json::Object(self.object)
        }
    }

    pub fn null(&mut self) {
        self.null = true;
    }
}

impl<V: ToJson, N: ToString> ObjectBuilder {
    pub fn set(&mut self, name: N, value: V) {
        self.object.insert(name.to_string(), value.to_json());
    }
}

impl<N: ToString> ObjectBuilder {
    pub fn list(&mut self, name: N, builder: |&mut ListBuilder|) {
        self.set(name, ListBuilder::build(builder).move_to_json());
    }

    pub fn object(&mut self, name: N, builder: |&mut ObjectBuilder|) {
        self.set(name, ObjectBuilder::build(builder).move_to_json());
    }
}

impl ToJson for ObjectBuilder {
    fn to_json(&self) -> Json {
        if self.null { json::Null } else { self.object.to_json() }
    }
}

#[test]
fn test() {
    let b = ListBuilder::new();
}<commit_msg>feat: add more common methods<commit_after>#![crate_name = "jsonway"]
#![comment = "JSON builder for Rust"]
#![license = "MIT"]
#![crate_type = "rlib"]
// #![deny(warnings)]
// #![deny(bad_style)]
#![feature(macro_rules, phase, tuple_indexing)]

extern crate serialize;

use std::collections::TreeMap;
use std::to_string::ToString;
use serialize::json;
use serialize::json::{JsonList, Json, JsonObject, ToJson};

pub use mutable_json::MutableJson;

pub mod mutable_json;

pub struct ListBuilder {
    list: JsonList,
    null: bool
}

impl ListBuilder {

    pub fn new() -> ListBuilder {
        ListBuilder { list: vec![], null: false }
    }

    pub fn from_json(list: Json) -> Option<ListBuilder> {
        match list {
            json::List(list) => Some(ListBuilder { list: list, null: false }),
            _ => None
        }
    }

    pub fn build(builder: |&mut ListBuilder|) -> ListBuilder {
        let mut bldr = ListBuilder::new();
        builder(&mut bldr);

        bldr
    }

    pub fn move_to_json(self) -> Json {
        if self.null {
            json::Null
        } else {
            json::List(self.list)
        }
    }

    pub fn push_json(&mut self, value: Json) {
        self.list.push(value);
    }

    pub fn list(&mut self, builder: |&mut ListBuilder|) {
        self.push(ListBuilder::build(builder).move_to_json());
    }

    pub fn object(&mut self, builder: |&mut ObjectBuilder|) {
        self.push(ObjectBuilder::build(builder).move_to_json());
    }

    pub fn null(&mut self) {
        self.null = true;
    }
}

impl<T: ToJson> ListBuilder {
    pub fn push(&mut self, value: T) {
        self.push_json(value.to_json());
    }
}

impl<A, T: Iterator<A>> ListBuilder {
    pub fn objects(&mut self, iter: &mut T, func: |A, &mut ObjectBuilder|) {
        let mut stop = false;
        while stop {
            let a = iter.next();
            if a.is_some() {
                let mut bldr = ObjectBuilder::new();
                func(a.unwrap(), &mut bldr);
                self.push(bldr.move_to_json())
            } else {
                stop = true;
            }
        }
    }

    pub fn lists(&mut self, iter: &mut T, func: |A, &mut ListBuilder|) {
        let mut stop = false;
        while stop {
            let a = iter.next();
            if a.is_some() {
                let mut bldr = ListBuilder::new();
                func(a.unwrap(), &mut bldr);
                self.push(bldr.move_to_json())
            } else {
                stop = true;
            }
        }
    }

    pub fn map(&mut self, iter: &mut T, func: |A| -> Json) {
        let mut stop = false;
        while stop {
            let a = iter.next();
            if a.is_some() {
                self.push(func(a.unwrap()))
            } else {
                stop = true;
            }
        }
    }
}

impl ToJson for ListBuilder {
    fn to_json(&self) -> Json {
         if self.null { json::Null } else { self.list.to_json() }
    }
}

pub struct ObjectBuilder {
    object: JsonObject,
    null: bool
}

impl ObjectBuilder {
    pub fn new() -> ObjectBuilder {
        ObjectBuilder { object: TreeMap::new(), null: false }
    }

    pub fn from_json(object: Json) -> Option<ObjectBuilder> {
        match object {
            json::Object(object) => Some(ObjectBuilder { object: object, null: false }),
            _ => None
        }
    }

    pub fn build(builder: |&mut ObjectBuilder|) -> ObjectBuilder {
        let mut bldr = ObjectBuilder::new();
        builder(&mut bldr);

        bldr
    }

    pub fn move_to_json(self) -> Json {
        if self.null {
            json::Null
        } else {
            json::Object(self.object)
        }
    }

    pub fn null(&mut self) {
        self.null = true;
    }

    pub fn index(&mut self) -> bool {
        true
    }
}

impl<V: ToJson, N: ToString> ObjectBuilder {
    pub fn set(&mut self, name: N, value: V) {
        self.set_json(name.to_string(), value.to_json());
    }

    pub fn call(&mut self, name: N, value: V) {
        self.set(name, value);
    }
}

impl<N: ToString> ObjectBuilder {
    pub fn set_json(&mut self, name: N, value: Json) {
        self.object.insert(name.to_string(), value);
    }

    pub fn list(&mut self, name: N, builder: |&mut ListBuilder|) {
        self.set(name, ListBuilder::build(builder).move_to_json());
    }

    pub fn object(&mut self, name: N, builder: |&mut ObjectBuilder|) {
        self.set(name, ObjectBuilder::build(builder).move_to_json());
    }
}

impl ToJson for ObjectBuilder {
    fn to_json(&self) -> Json {
        if self.null { json::Null } else { self.object.to_json() }
    }
}

#[test]
fn simple_object() {
    ObjectBuilder::build(|json| {
        json.set("first_name", "Luke".to_string());
        json.set("last_name", "Skywalker".to_string());
    });
}

#[test]
fn simple_list() {
    ListBuilder::build(|json| {
        json.push("Luke".to_string());
        json.push("Skywalker".to_string());
    });
}<|endoftext|>


<commit_before>#![allow(dead_code)]

//! A wrapper around [cfitsio].
//!
//! This wrapper attempts to add some safety to the [cfitsio] library, and allows `.fits` files to
//! be read from rust.
//!
//! The main entry point to the code is `FitsFile::open` which creates a new file, and
//! `FitsFile::create` which creates a new file.
//!
//! [cfitsio]: https://heasarc.gsfc.nasa.gov/docs/software/fitsio/fitsio.html

pub mod raw;

extern crate libc;

use raw::*;
use libc::{c_int, c_long, c_char, c_void};
use std::ptr;
use std::ffi;
use std::result;

/// Error type
#[derive(Debug)]
pub struct FitsError {
    status: i32,
    message: String,
}

pub type Result<T> = result::Result<T, FitsError>;

/// Internal function to get the fits error description from a status code
fn status_to_string(status: c_int) -> Option<String> {
    match status {
        0 => None,
        status => {
            let mut buffer: Vec<c_char> = vec![0; 31];
            unsafe {
                ffgerr(status, buffer.as_mut_ptr());
            }
            let result_str = String::from_utf8(buffer.iter()
                    .map(|&x| x as u8)
                    .filter(|&x| x != 0)
                    .collect())
                .unwrap();
            Some(result_str)
        }
    }
}

/// General type defining what kind of HDU we're talking about
#[derive(Eq, PartialEq, Debug)]
pub enum HduType {
    ImageHDU,
    AsciiTableHDU,
    BinTableHDU,
}

/// Wrapper around a C `fitsfile` pointer.
///
/// This handles [opening][FitsFile::open], [creating][FitsFile::create] and automatically
/// closing (through the `Drop` trait) the file.
///
/// All subsequent file access is controled through this object.
///
/// [FitsFile::open]: #method.open
/// [FitsFile::create]: #method.create
pub struct FitsFile {
    fptr: *mut fitsfile,
    status: c_int,
    pub filename: String,
}

impl FitsFile {
    /// Open a fits file for reading
    ///
    /// * `filename` - Filename to pass to `cfitsio`. Can conform to the
    /// [Extended Filename Syntax][extended-filename-syntax].
    ///
    /// [extended-filename-syntax]:
    ///     https://heasarc.gsfc.nasa.gov/docs/software/fitsio/c/c_user/node82.html
    ///
    /// Examples
    ///
    /// ```
    /// # use fitsio::FitsFile;
    /// # fn main() {
    ///     match FitsFile::open("testdata/full_example.fits") {
    ///         Ok(f) => { },
    ///         Err(e) => panic!("{:?}", e),
    ///     }
    /// # }
    pub fn open(filename: &str) -> Result<Self> {
        let mut fptr = ptr::null_mut();
        let mut status = 0;
        let c_filename = ffi::CString::new(filename).unwrap();

        unsafe {
            ffopen(&mut fptr as *mut *mut fitsfile,
                   c_filename.as_ptr(),
                   0,
                   &mut status);
        }

        return match status {
            0 => {
                Ok(FitsFile {
                    fptr: fptr,
                    status: status,
                    filename: filename.to_string(),
                })
            }
            status => {
                Err(FitsError {
                    status: status,
                    message: status_to_string(status).unwrap(),
                })
            }
        };

    }

    /// Create a new fits file
    ///
    /// An empty primary header is added so when the `drop` method is called, the file is valid.
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate fitsio;
    /// # extern crate tempdir;
    /// # use fitsio::FitsFile;
    /// # fn main() {
    /// # let tdir = tempdir::TempDir::new("fitsio-").unwrap();
    /// # let tdir_path = tdir.path();
    /// # let filename = tdir_path.join("test.fits");
    /// # let filename = filename.to_str().unwrap();
    /// match FitsFile::create(filename) {
    ///     Ok(f) => println!("Fits file created ok"),
    ///     Err(e) => panic!("Error: {:?}", e),
    /// }
    /// # }
    /// ```
    pub fn create(path: &str) -> Result<Self> {
        let mut fptr = ptr::null_mut();
        let mut status = 0;
        let c_filename = ffi::CString::new(path).unwrap();

        unsafe {
            ffinit(&mut fptr as *mut *mut fitsfile,
                   c_filename.as_ptr(),
                   &mut status);
        }

        return match status {
            0 => {
                Ok(FitsFile {
                    fptr: fptr,
                    status: status,
                    filename: path.to_string(),
                })
            }
            status => {
                Err(FitsError {
                    status: status,
                    message: status_to_string(status).unwrap(),
                })
            }
        };
    }

    /// Function to check that the status code is ok.
    ///
    /// If the value of `self.status` is not 0 then exit the current process as an error has
    /// occurred.
    fn check(&self) {
        match self.status {
            0 => {}
            status => {
                panic!("Status code {} encountered, msg: {}",
                       status,
                       status_to_string(status).unwrap())
            }
        }
    }

    /// Returns the current HDU number, 0-indexed
    ///
    /// The FITS standard is 1-indexed (thanks Fortran), so this function
    /// converts the indexing to be 0-based so the primary HDU is at
    /// index 0.
    ///
    /// # Examples
    ///
    /// ```
    /// # use fitsio::FitsFile;
    /// let f = FitsFile::open("testdata/full_example.fits").unwrap();
    /// assert_eq!(f.current_hdu_number(), 0);
    /// ```
    pub fn current_hdu_number(&self) -> u32 {
        let mut hdu_num = 0;
        unsafe {
            ffghdn(self.fptr, &mut hdu_num);
        }
        self.check();
        assert!(hdu_num >= 1);
        (hdu_num - 1) as u32
    }

    /// Change the current HDU
    ///
    /// Examples
    ///
    /// ```
    /// # use fitsio::FitsFile;
    /// # fn main() {
    /// # let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
    /// assert_eq!(f.current_hdu_number(), 0);
    /// f.change_hdu(1);
    /// assert_eq!(f.current_hdu_number(), 1);
    /// # }
    /// ```
    pub fn change_hdu(&mut self, hdu_num: u32) {
        let mut _hdu_type = 0;
        unsafe {
            ffmahd(self.fptr,
                   (hdu_num + 1) as i32,
                   &mut _hdu_type,
                   &mut self.status);
        }
        self.check();
    }

    /// Get which type of HDU the current HDU is
    ///
    /// Results in one of the `HduType` options.
    ///
    /// # Examples
    ///
    /// ```
    /// # use fitsio::{FitsFile, HduType};
    /// # fn main() {
    /// let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
    /// // Primary HDUs are always image hdus
    /// assert_eq!(f.get_hdu_type(), HduType::ImageHDU);
    /// # }
    /// ```
    pub fn get_hdu_type(&mut self) -> HduType {
        let mut hdu_type = 3;
        unsafe {
            ffghdt(self.fptr, &mut hdu_type, &mut self.status);
        }
        self.check();

        match hdu_type {
            0 => HduType::ImageHDU,
            1 => HduType::AsciiTableHDU,
            2 => HduType::BinTableHDU,
            _ => panic!("Unknown hdu type: {}", hdu_type),
        }
    }

    /// Return a `FitsHDU` object for the specified HDU
    ///
    /// # Examples
    ///
    /// ```
    /// # use fitsio::{FitsFile, HduType};
    /// # fn main() {
    /// # let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
    /// let primary_hdu = f.get_hdu(0);
    /// assert_eq!(primary_hdu.hdu_type, HduType::ImageHDU);
    /// # }
    /// ```
    pub fn get_hdu(&mut self, index: usize) -> FitsHDU {
        self.change_hdu(index as u32);
        let hdu_type = self.get_hdu_type();

        let image_shape = if hdu_type == HduType::ImageHDU {
            let mut naxis = vec![0, 0];
            unsafe {
                ffgisz(self.fptr, 2, naxis.as_mut_ptr(), &mut self.status);
            }
            println!("{:?}", naxis);
            (naxis[0] as usize, naxis[1] as usize)
        } else {
            (0, 0)
        };

        FitsHDU {
            fitsfile: self,
            hdu_type: hdu_type,
            image_shape: image_shape,
        }
    }
}

impl Drop for FitsFile {
    fn drop(&mut self) {
        unsafe {
            ffclos(self.fptr, &mut self.status);
        }
    }
}

/// Wrapper around a FITS HDU
///
/// This struct is the main interface around reading and writing the file contents.
pub struct FitsHDU<'a> {
    fitsfile: &'a FitsFile,
    pub hdu_type: HduType,
    image_shape: (usize, usize),
}

impl<'a> FitsHDU<'a> {
    /// Read a header key as a string
    ///
    /// The user is responsible for converting the value type from a string to whatever type the
    /// header key contains.
    ///
    /// # Examples
    ///
    /// ```
    /// # use fitsio::FitsFile;
    /// # fn main() {
    /// let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
    /// let mut primary_hdu = f.get_hdu(0);
    /// // Image is 2-dimensional
    /// let naxis = primary_hdu.get_key("NAXIS").parse::<i32>().unwrap();
    /// assert_eq!(naxis, 2);
    /// # }
    /// ```
    pub fn get_key(&mut self, key: &str) -> String {
        let fptr = &self.fitsfile.fptr;
        let mut value: Vec<c_char> = vec![0; MAX_VALUE_LENGTH];
        let keyname = ffi::CString::new(key).unwrap();
        let mut status = 0;

        unsafe {
            ffgkys(*fptr,
                   keyname.as_ptr(),
                   value.as_mut_ptr(),
                   ptr::null_mut(),
                   &mut status);
        }

        match status {
            0 => {
                let value: Vec<u8> = value.iter()
                    .map(|&x| x as u8)
                    .filter(|&x| x != 0)
                    .collect();
                return String::from_utf8(value).unwrap();
            }
            _ => {
                panic!("Invalid status code: {}, msg: {}",
                       status,
                       status_to_string(status).unwrap())
            }
        }
    }

    fn read_all_i32(&mut self, buffer: &mut Vec<i32>) {
        let npix = self.image_shape.0 * self.image_shape.1;
        buffer.resize(npix, 0);
        let mut status = 0;

        unsafe {
            ffgpv(self.fitsfile.fptr,
                  31, // TINT
                  1,
                  npix as i64,
                  ptr::null_mut() as *mut c_void,
                  buffer.as_mut_ptr() as *mut c_void,
                  ptr::null_mut(),
                  &mut status);
        }
        match status {
            0 => {}
            status => panic!("Bad status value: {}", status),
        }
        self.fitsfile.check();
    }
}

#[cfg(test)]
mod test {
    extern crate tempdir;
    use super::*;
    use super::status_to_string;

    #[test]
    fn returning_error_messages() {
        assert_eq!(status_to_string(105).unwrap(),
                   "couldn't create the named file");
    }

    #[test]
    fn opening_an_existing_file() {
        match FitsFile::open("testdata/full_example.fits") {
            Ok(f) => assert_eq!(f.status, 0),
            Err(e) => panic!("{:?}", e),
        }
    }

    #[test]
    fn creating_a_new_file() {
        let tdir = tempdir::TempDir::new("fitsio-").unwrap();
        let tdir_path = tdir.path();
        let filename = tdir_path.join("test.fits");
        assert!(!filename.exists());

        match FitsFile::create(filename.to_str().unwrap()) {
            Ok(f) => assert!(filename.exists()),
            Err(e) => panic!("Error: {:?}", e),
        }
    }

    #[test]
    fn filename_is_stored() {
        let f = FitsFile::open("testdata/full_example.fits").unwrap();
        assert_eq!(f.filename, "testdata/full_example.fits");
    }

    #[test]
    fn change_hdu() {
        let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
        f.change_hdu(1);
        assert_eq!(f.current_hdu_number(), 1u32);
    }

    #[test]
    fn getting_current_hdu_number() {
        let f = FitsFile::open("testdata/full_example.fits").unwrap();
        assert_eq!(f.current_hdu_number(), 0u32);
    }

    #[test]
    fn getting_hdu_object() {
        let mut f = FitsFile::open("testdata/full_example.fits").unwrap();

        // TODO: get rid of these scopes
        //
        // They're there because `get_hdu` mutably borrows the `self` object and
        // the result is still in scope. We need to expire the `FitsHDU` objects
        // and hence the scopes.
        {
            let primary_hdu = f.get_hdu(0);
            assert_eq!(primary_hdu.hdu_type, HduType::ImageHDU);
        }

        {
            let table_hdu = f.get_hdu(1);
            assert_eq!(table_hdu.hdu_type, HduType::BinTableHDU);
        }
    }

    #[test]
    fn reading_in_image_data() {
        let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
        let mut primary_hdu = f.get_hdu(0);
        let mut data = Vec::new();
        primary_hdu.read_all_i32(&mut data);
        assert_eq!(data[0], 108);
        assert_eq!(data[1], 176);
    }

    #[test]
    fn get_image_dimensions() {
        let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
        let primary_hdu = f.get_hdu(0);
        assert_eq!(primary_hdu.image_shape, (100, 100));
    }
}
<commit_msg>get_key returns error object<commit_after>#![allow(dead_code)]

//! A wrapper around [cfitsio].
//!
//! This wrapper attempts to add some safety to the [cfitsio] library, and allows `.fits` files to
//! be read from rust.
//!
//! The main entry point to the code is `FitsFile::open` which creates a new file, and
//! `FitsFile::create` which creates a new file.
//!
//! [cfitsio]: https://heasarc.gsfc.nasa.gov/docs/software/fitsio/fitsio.html

pub mod raw;

extern crate libc;

use raw::*;
use libc::{c_int, c_long, c_char, c_void};
use std::ptr;
use std::ffi;
use std::result;

/// Error type
#[derive(Debug)]
pub struct FitsError {
    status: i32,
    message: String,
}

pub type Result<T> = result::Result<T, FitsError>;

/// Internal function to get the fits error description from a status code
fn status_to_string(status: c_int) -> Option<String> {
    match status {
        0 => None,
        status => {
            let mut buffer: Vec<c_char> = vec![0; 31];
            unsafe {
                ffgerr(status, buffer.as_mut_ptr());
            }
            let result_str = String::from_utf8(buffer.iter()
                    .map(|&x| x as u8)
                    .filter(|&x| x != 0)
                    .collect())
                .unwrap();
            Some(result_str)
        }
    }
}

/// General type defining what kind of HDU we're talking about
#[derive(Eq, PartialEq, Debug)]
pub enum HduType {
    ImageHDU,
    AsciiTableHDU,
    BinTableHDU,
}

/// Wrapper around a C `fitsfile` pointer.
///
/// This handles [opening][FitsFile::open], [creating][FitsFile::create] and automatically
/// closing (through the `Drop` trait) the file.
///
/// All subsequent file access is controled through this object.
///
/// [FitsFile::open]: #method.open
/// [FitsFile::create]: #method.create
pub struct FitsFile {
    fptr: *mut fitsfile,
    status: c_int,
    pub filename: String,
}

impl FitsFile {
    /// Open a fits file for reading
    ///
    /// * `filename` - Filename to pass to `cfitsio`. Can conform to the
    /// [Extended Filename Syntax][extended-filename-syntax].
    ///
    /// [extended-filename-syntax]:
    ///     https://heasarc.gsfc.nasa.gov/docs/software/fitsio/c/c_user/node82.html
    ///
    /// Examples
    ///
    /// ```
    /// # use fitsio::FitsFile;
    /// # fn main() {
    ///     match FitsFile::open("testdata/full_example.fits") {
    ///         Ok(f) => { },
    ///         Err(e) => panic!("{:?}", e),
    ///     }
    /// # }
    pub fn open(filename: &str) -> Result<Self> {
        let mut fptr = ptr::null_mut();
        let mut status = 0;
        let c_filename = ffi::CString::new(filename).unwrap();

        unsafe {
            ffopen(&mut fptr as *mut *mut fitsfile,
                   c_filename.as_ptr(),
                   0,
                   &mut status);
        }

        return match status {
            0 => {
                Ok(FitsFile {
                    fptr: fptr,
                    status: status,
                    filename: filename.to_string(),
                })
            }
            status => {
                Err(FitsError {
                    status: status,
                    message: status_to_string(status).unwrap(),
                })
            }
        };

    }

    /// Create a new fits file
    ///
    /// An empty primary header is added so when the `drop` method is called, the file is valid.
    ///
    /// # Examples
    ///
    /// ```
    /// # extern crate fitsio;
    /// # extern crate tempdir;
    /// # use fitsio::FitsFile;
    /// # fn main() {
    /// # let tdir = tempdir::TempDir::new("fitsio-").unwrap();
    /// # let tdir_path = tdir.path();
    /// # let filename = tdir_path.join("test.fits");
    /// # let filename = filename.to_str().unwrap();
    /// match FitsFile::create(filename) {
    ///     Ok(f) => println!("Fits file created ok"),
    ///     Err(e) => panic!("Error: {:?}", e),
    /// }
    /// # }
    /// ```
    pub fn create(path: &str) -> Result<Self> {
        let mut fptr = ptr::null_mut();
        let mut status = 0;
        let c_filename = ffi::CString::new(path).unwrap();

        unsafe {
            ffinit(&mut fptr as *mut *mut fitsfile,
                   c_filename.as_ptr(),
                   &mut status);
        }

        return match status {
            0 => {
                Ok(FitsFile {
                    fptr: fptr,
                    status: status,
                    filename: path.to_string(),
                })
            }
            status => {
                Err(FitsError {
                    status: status,
                    message: status_to_string(status).unwrap(),
                })
            }
        };
    }

    /// Function to check that the status code is ok.
    ///
    /// If the value of `self.status` is not 0 then exit the current process as an error has
    /// occurred.
    fn check(&self) {
        match self.status {
            0 => {}
            status => {
                panic!("Status code {} encountered, msg: {}",
                       status,
                       status_to_string(status).unwrap())
            }
        }
    }

    /// Returns the current HDU number, 0-indexed
    ///
    /// The FITS standard is 1-indexed (thanks Fortran), so this function
    /// converts the indexing to be 0-based so the primary HDU is at
    /// index 0.
    ///
    /// # Examples
    ///
    /// ```
    /// # use fitsio::FitsFile;
    /// let f = FitsFile::open("testdata/full_example.fits").unwrap();
    /// assert_eq!(f.current_hdu_number(), 0);
    /// ```
    pub fn current_hdu_number(&self) -> u32 {
        let mut hdu_num = 0;
        unsafe {
            ffghdn(self.fptr, &mut hdu_num);
        }
        self.check();
        assert!(hdu_num >= 1);
        (hdu_num - 1) as u32
    }

    /// Change the current HDU
    ///
    /// Examples
    ///
    /// ```
    /// # use fitsio::FitsFile;
    /// # fn main() {
    /// # let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
    /// assert_eq!(f.current_hdu_number(), 0);
    /// f.change_hdu(1);
    /// assert_eq!(f.current_hdu_number(), 1);
    /// # }
    /// ```
    pub fn change_hdu(&mut self, hdu_num: u32) {
        let mut _hdu_type = 0;
        unsafe {
            ffmahd(self.fptr,
                   (hdu_num + 1) as i32,
                   &mut _hdu_type,
                   &mut self.status);
        }
        self.check();
    }

    /// Get which type of HDU the current HDU is
    ///
    /// Results in one of the `HduType` options.
    ///
    /// # Examples
    ///
    /// ```
    /// # use fitsio::{FitsFile, HduType};
    /// # fn main() {
    /// let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
    /// // Primary HDUs are always image hdus
    /// assert_eq!(f.get_hdu_type(), HduType::ImageHDU);
    /// # }
    /// ```
    pub fn get_hdu_type(&mut self) -> HduType {
        let mut hdu_type = 3;
        unsafe {
            ffghdt(self.fptr, &mut hdu_type, &mut self.status);
        }
        self.check();

        match hdu_type {
            0 => HduType::ImageHDU,
            1 => HduType::AsciiTableHDU,
            2 => HduType::BinTableHDU,
            _ => panic!("Unknown hdu type: {}", hdu_type),
        }
    }

    /// Return a `FitsHDU` object for the specified HDU
    ///
    /// # Examples
    ///
    /// ```
    /// # use fitsio::{FitsFile, HduType};
    /// # fn main() {
    /// # let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
    /// let primary_hdu = f.get_hdu(0);
    /// assert_eq!(primary_hdu.hdu_type, HduType::ImageHDU);
    /// # }
    /// ```
    pub fn get_hdu(&mut self, index: usize) -> FitsHDU {
        self.change_hdu(index as u32);
        let hdu_type = self.get_hdu_type();

        let image_shape = if hdu_type == HduType::ImageHDU {
            let mut naxis = vec![0, 0];
            unsafe {
                ffgisz(self.fptr, 2, naxis.as_mut_ptr(), &mut self.status);
            }
            println!("{:?}", naxis);
            (naxis[0] as usize, naxis[1] as usize)
        } else {
            (0, 0)
        };

        FitsHDU {
            fitsfile: self,
            hdu_type: hdu_type,
            image_shape: image_shape,
        }
    }
}

impl Drop for FitsFile {
    fn drop(&mut self) {
        unsafe {
            ffclos(self.fptr, &mut self.status);
        }
    }
}

/// Wrapper around a FITS HDU
///
/// This struct is the main interface around reading and writing the file contents.
pub struct FitsHDU<'a> {
    fitsfile: &'a FitsFile,
    pub hdu_type: HduType,
    image_shape: (usize, usize),
}

impl<'a> FitsHDU<'a> {
    /// Read a header key as a string
    ///
    /// The user is responsible for converting the value type from a string to whatever type the
    /// header key contains.
    ///
    /// # Examples
    ///
    /// ```
    /// # use fitsio::FitsFile;
    /// # fn main() {
    /// let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
    /// let mut primary_hdu = f.get_hdu(0);
    /// // Image is 2-dimensional
    /// let naxis = primary_hdu.get_key("NAXIS").unwrap().parse::<i32>().unwrap();
    /// assert_eq!(naxis, 2);
    /// # }
    /// ```
    pub fn get_key(&mut self, key: &str) -> Result<String> {
        let fptr = &self.fitsfile.fptr;
        let mut value: Vec<c_char> = vec![0; MAX_VALUE_LENGTH];
        let keyname = ffi::CString::new(key).unwrap();
        let mut status = 0;

        unsafe {
            ffgkys(*fptr,
                   keyname.as_ptr(),
                   value.as_mut_ptr(),
                   ptr::null_mut(),
                   &mut status);
        }

        match status {
            0 => {
                let value: Vec<u8> = value.iter()
                    .map(|&x| x as u8)
                    .filter(|&x| x != 0)
                    .collect();
                Ok(String::from_utf8(value).unwrap())
            }
            status => {
                Err(FitsError {
                    status: status,
                    message: status_to_string(status).unwrap(),
                })
            }
        }
    }

    fn read_all_i32(&mut self, buffer: &mut Vec<i32>) {
        let npix = self.image_shape.0 * self.image_shape.1;
        buffer.resize(npix, 0);
        let mut status = 0;

        unsafe {
            ffgpv(self.fitsfile.fptr,
                  31, // TINT
                  1,
                  npix as i64,
                  ptr::null_mut() as *mut c_void,
                  buffer.as_mut_ptr() as *mut c_void,
                  ptr::null_mut(),
                  &mut status);
        }
        match status {
            0 => {}
            status => panic!("Bad status value: {}", status),
        }
        self.fitsfile.check();
    }
}

#[cfg(test)]
mod test {
    extern crate tempdir;
    use super::*;
    use super::status_to_string;

    #[test]
    fn returning_error_messages() {
        assert_eq!(status_to_string(105).unwrap(),
                   "couldn't create the named file");
    }

    #[test]
    fn opening_an_existing_file() {
        match FitsFile::open("testdata/full_example.fits") {
            Ok(f) => assert_eq!(f.status, 0),
            Err(e) => panic!("{:?}", e),
        }
    }

    #[test]
    fn creating_a_new_file() {
        let tdir = tempdir::TempDir::new("fitsio-").unwrap();
        let tdir_path = tdir.path();
        let filename = tdir_path.join("test.fits");
        assert!(!filename.exists());

        match FitsFile::create(filename.to_str().unwrap()) {
            Ok(f) => assert!(filename.exists()),
            Err(e) => panic!("Error: {:?}", e),
        }
    }

    #[test]
    fn filename_is_stored() {
        let f = FitsFile::open("testdata/full_example.fits").unwrap();
        assert_eq!(f.filename, "testdata/full_example.fits");
    }

    #[test]
    fn change_hdu() {
        let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
        f.change_hdu(1);
        assert_eq!(f.current_hdu_number(), 1u32);
    }

    #[test]
    fn getting_current_hdu_number() {
        let f = FitsFile::open("testdata/full_example.fits").unwrap();
        assert_eq!(f.current_hdu_number(), 0u32);
    }

    #[test]
    fn getting_hdu_object() {
        let mut f = FitsFile::open("testdata/full_example.fits").unwrap();

        // TODO: get rid of these scopes
        //
        // They're there because `get_hdu` mutably borrows the `self` object and
        // the result is still in scope. We need to expire the `FitsHDU` objects
        // and hence the scopes.
        {
            let primary_hdu = f.get_hdu(0);
            assert_eq!(primary_hdu.hdu_type, HduType::ImageHDU);
        }

        {
            let table_hdu = f.get_hdu(1);
            assert_eq!(table_hdu.hdu_type, HduType::BinTableHDU);
        }
    }

    #[test]
    fn reading_in_image_data() {
        let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
        let mut primary_hdu = f.get_hdu(0);
        let mut data = Vec::new();
        primary_hdu.read_all_i32(&mut data);
        assert_eq!(data[0], 108);
        assert_eq!(data[1], 176);
    }

    #[test]
    fn get_image_dimensions() {
        let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
        let primary_hdu = f.get_hdu(0);
        assert_eq!(primary_hdu.image_shape, (100, 100));
    }

    #[test]
    fn get_key_returns_error_for_missing_key() {
        let mut f = FitsFile::open("testdata/full_example.fits").unwrap();
        let mut primary_hdu = f.get_hdu(0);

        match primary_hdu.get_key("THISKEYDOESNOTEXIST") {
            Err(e) => assert_eq!(e.status, 202),
            Ok(f) => panic!("No error thrown"),
        }
    }
}
<|endoftext|>


<commit_before>#![no_std]

#[macro_export]
macro_rules! simple_bitfield_fields {
   ($t:ty,) => {};
   ($t:ty, _, $setter:ident: $msb:expr, $lsb:expr, $($rest:tt)*) => {
       pub fn $setter(&mut self, value: $t) {
           use $crate::BitRange;
           self.set_bit_range($msb, $lsb, value);
       }
       simple_bitfield_fields!{$t, $($rest)*}
   };
   ($t:ty, _, $setter:ident: $msb:expr, $lsb:expr; $count:expr, $($rest:tt)*) => {
       #[allow(unknown_lints)]
       #[allow(eq_op)]
       pub fn $setter(&mut self, index: usize, value: $t) {
           use $crate::BitRange;
           debug_assert!(index < $count);
           let width = $msb - $lsb + 1;
           let lsb = $lsb + index*width;
           let msb = lsb + width - 1;
           self.set_bit_range(msb, lsb, value);
       }
       simple_bitfield_fields!{$t, $($rest)*}
   };
   ($t:ty, $getter:ident, _: $msb:expr, $lsb:expr, $($rest:tt)*) => {
       pub fn $getter(&self) -> $t {
           use $crate::BitRange;
           self.bit_range($msb, $lsb)
       }
       simple_bitfield_fields!{$t, $($rest)*}
   };
   ($t:ty, $getter:ident, _: $msb:expr, $lsb:expr; $count:expr, $($rest:tt)*) => {
       #[allow(unknown_lints)]
       #[allow(eq_op)]
       pub fn $getter(&self, index: usize) -> $t {
           use $crate::BitRange;
           debug_assert!(index < $count);
           let width = $msb - $lsb + 1;
           let lsb = $lsb + index*width;
           let msb = lsb + width - 1;
           self.bit_range(msb, lsb)
       }
       simple_bitfield_fields!{$t, $($rest)*}
   };
   ($t:ty, $getter:ident, $setter:ident: $msb:expr, $lsb:expr, $($rest:tt)*) => {
       simple_bitfield_fields!{$t, $getter, _: $msb, $lsb, }
       simple_bitfield_fields!{$t, _, $setter: $msb, $lsb, }
       simple_bitfield_fields!{$t, $($rest)*}
   };
   ($t:ty, $getter:ident, $setter:ident: $msb:expr, $lsb:expr; $count:expr, $($rest:tt)*) => {
         simple_bitfield_fields!{$t, $getter, _: $msb, $lsb; $count, }
         simple_bitfield_fields!{$t, _, $setter: $msb, $lsb; $count, }
         simple_bitfield_fields!{$t, $($rest)*}
   };
}

#[macro_export]
macro_rules! simple_bitfield {
    ($name:ident, [$t:ty]; $($rest:tt)*) => {
        pub struct $name<T>(pub T);

        impl<T: AsMut<[$t]> + AsRef<[$t]>> $crate::BitRange<u64> for $name<T> {
            fn bit_range(&self, msb: usize, lsb: usize) -> u64 {
                let bit_len = $crate::size_of::<$t>()*8;
                let mut value = 0;
                for i in (lsb..msb+1).rev() {
                    value <<= 1;
                    value |= ((self.0.as_ref()[i/bit_len] >> (i%bit_len)) & 1) as u64;
                }
                value
            }

            fn set_bit_range(&mut self, msb: usize, lsb: usize, value: u64) {
                let bit_len = $crate::size_of::<$t>()*8;
                let mut value = value;
                for i in lsb..msb+1 {
                    self.0.as_mut()[i/bit_len] &= !(1 << (i%bit_len));
                    self.0.as_mut()[i/bit_len] |= (value & 1) as $t << (i%bit_len);
                    value >>= 1;
                }
            }
        }

        impl<T: AsMut<[$t]> + AsRef<[$t]>> $name<T> {
            simple_bitfield_fields!{u64, $($rest)*}
        }
    };
    ($name:ident, $t:ty; $($rest:tt)*) => {
        #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
        #[repr(C)]
        pub struct $name(pub $t);

        impl $crate::BitRange<$t> for $name {
            fn bit_range(&self, msb: usize, lsb: usize) -> $t {
                self.0.bit_range(msb, lsb)
            }
            fn set_bit_range(&mut self, msb: usize, lsb: usize, value: $t) {
                self.0.set_bit_range(msb, lsb, value);
            }
        }

        impl $name {
            simple_bitfield_fields!{$t, $($rest)*}
         }
    };
}

#[doc(hidden)]
pub fn size_of<T>() -> usize {
    core::mem::size_of::<T>()
}

pub trait BitRange<T> {
    fn bit_range(&self, msb: usize, lsb: usize) -> T;
    fn set_bit_range(&mut self, msb: usize, lsb: usize, value: T);
}

macro_rules! impl_bitrange_for_u {
    ($t:ty) => {
        impl BitRange<$t> for $t {
            #[inline]
            fn bit_range(&self, msb: usize, lsb: usize) -> $t {
                let bit_len = size_of::<$t>()*8;
                (*self << (bit_len - msb - 1)) >> (bit_len - msb - 1 + lsb)
            }

            #[inline]
            fn set_bit_range(&mut self, msb: usize, lsb: usize, value: $t) {
                let bit_len = size_of::<$t>()*8;
                let mask: $t = !(0 as $t)
                    << (bit_len - msb - 1)
                    >> (bit_len - msb - 1 + lsb)
                        << (lsb);
                *self &= !mask;
                *self |= (value << lsb) & mask;
            }
        }
    }
}

impl_bitrange_for_u!{u8}
impl_bitrange_for_u!{u16}
impl_bitrange_for_u!{u32}
impl_bitrange_for_u!{u64}
<commit_msg>Add doc comments for BitRange<commit_after>#![no_std]

#[macro_export]
macro_rules! simple_bitfield_fields {
   ($t:ty,) => {};
   ($t:ty, _, $setter:ident: $msb:expr, $lsb:expr, $($rest:tt)*) => {
       pub fn $setter(&mut self, value: $t) {
           use $crate::BitRange;
           self.set_bit_range($msb, $lsb, value);
       }
       simple_bitfield_fields!{$t, $($rest)*}
   };
   ($t:ty, _, $setter:ident: $msb:expr, $lsb:expr; $count:expr, $($rest:tt)*) => {
       #[allow(unknown_lints)]
       #[allow(eq_op)]
       pub fn $setter(&mut self, index: usize, value: $t) {
           use $crate::BitRange;
           debug_assert!(index < $count);
           let width = $msb - $lsb + 1;
           let lsb = $lsb + index*width;
           let msb = lsb + width - 1;
           self.set_bit_range(msb, lsb, value);
       }
       simple_bitfield_fields!{$t, $($rest)*}
   };
   ($t:ty, $getter:ident, _: $msb:expr, $lsb:expr, $($rest:tt)*) => {
       pub fn $getter(&self) -> $t {
           use $crate::BitRange;
           self.bit_range($msb, $lsb)
       }
       simple_bitfield_fields!{$t, $($rest)*}
   };
   ($t:ty, $getter:ident, _: $msb:expr, $lsb:expr; $count:expr, $($rest:tt)*) => {
       #[allow(unknown_lints)]
       #[allow(eq_op)]
       pub fn $getter(&self, index: usize) -> $t {
           use $crate::BitRange;
           debug_assert!(index < $count);
           let width = $msb - $lsb + 1;
           let lsb = $lsb + index*width;
           let msb = lsb + width - 1;
           self.bit_range(msb, lsb)
       }
       simple_bitfield_fields!{$t, $($rest)*}
   };
   ($t:ty, $getter:ident, $setter:ident: $msb:expr, $lsb:expr, $($rest:tt)*) => {
       simple_bitfield_fields!{$t, $getter, _: $msb, $lsb, }
       simple_bitfield_fields!{$t, _, $setter: $msb, $lsb, }
       simple_bitfield_fields!{$t, $($rest)*}
   };
   ($t:ty, $getter:ident, $setter:ident: $msb:expr, $lsb:expr; $count:expr, $($rest:tt)*) => {
         simple_bitfield_fields!{$t, $getter, _: $msb, $lsb; $count, }
         simple_bitfield_fields!{$t, _, $setter: $msb, $lsb; $count, }
         simple_bitfield_fields!{$t, $($rest)*}
   };
}

#[macro_export]
macro_rules! simple_bitfield {
    ($name:ident, [$t:ty]; $($rest:tt)*) => {
        pub struct $name<T>(pub T);

        impl<T: AsMut<[$t]> + AsRef<[$t]>> $crate::BitRange<u64> for $name<T> {
            fn bit_range(&self, msb: usize, lsb: usize) -> u64 {
                let bit_len = $crate::size_of::<$t>()*8;
                let mut value = 0;
                for i in (lsb..msb+1).rev() {
                    value <<= 1;
                    value |= ((self.0.as_ref()[i/bit_len] >> (i%bit_len)) & 1) as u64;
                }
                value
            }

            fn set_bit_range(&mut self, msb: usize, lsb: usize, value: u64) {
                let bit_len = $crate::size_of::<$t>()*8;
                let mut value = value;
                for i in lsb..msb+1 {
                    self.0.as_mut()[i/bit_len] &= !(1 << (i%bit_len));
                    self.0.as_mut()[i/bit_len] |= (value & 1) as $t << (i%bit_len);
                    value >>= 1;
                }
            }
        }

        impl<T: AsMut<[$t]> + AsRef<[$t]>> $name<T> {
            simple_bitfield_fields!{u64, $($rest)*}
        }
    };
    ($name:ident, $t:ty; $($rest:tt)*) => {
        #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
        #[repr(C)]
        pub struct $name(pub $t);

        impl $crate::BitRange<$t> for $name {
            fn bit_range(&self, msb: usize, lsb: usize) -> $t {
                self.0.bit_range(msb, lsb)
            }
            fn set_bit_range(&mut self, msb: usize, lsb: usize, value: $t) {
                self.0.set_bit_range(msb, lsb, value);
            }
        }

        impl $name {
            simple_bitfield_fields!{$t, $($rest)*}
         }
    };
}

#[doc(hidden)]
pub fn size_of<T>() -> usize {
    core::mem::size_of::<T>()
}

/// A trait to get or set ranges of bits.
pub trait BitRange<T> {
    /// Get a range of bits.
    fn bit_range(&self, msb: usize, lsb: usize) -> T;
    /// Set a range of bits.
    fn set_bit_range(&mut self, msb: usize, lsb: usize, value: T);
}

macro_rules! impl_bitrange_for_u {
    ($t:ty) => {
        impl BitRange<$t> for $t {
            #[inline]
            fn bit_range(&self, msb: usize, lsb: usize) -> $t {
                let bit_len = size_of::<$t>()*8;
                (*self << (bit_len - msb - 1)) >> (bit_len - msb - 1 + lsb)
            }

            #[inline]
            fn set_bit_range(&mut self, msb: usize, lsb: usize, value: $t) {
                let bit_len = size_of::<$t>()*8;
                let mask: $t = !(0 as $t)
                    << (bit_len - msb - 1)
                    >> (bit_len - msb - 1 + lsb)
                        << (lsb);
                *self &= !mask;
                *self |= (value << lsb) & mask;
            }
        }
    }
}

impl_bitrange_for_u!{u8}
impl_bitrange_for_u!{u16}
impl_bitrange_for_u!{u32}
impl_bitrange_for_u!{u64}
<|endoftext|>