<commit_before>use crate::constants::*;

#[derive(Clone, Debug, PartialEq)]
pub enum Style {
    Old,
    OldInverse,
    New,
    NewInverse,
    Error,
}

#[derive(Clone, Debug, PartialEq)]
pub struct StyledToken {
    token: String,
    style: Style,
}

impl StyledToken {
    pub fn new(token: String, style: Style) -> StyledToken {
        return StyledToken { token, style };
    }

    pub fn is_whitespace(&self) -> bool {
        let mut chars_iterator = self.token.chars();
        let first_char = chars_iterator.next().unwrap();
        if chars_iterator.next().is_some() {
            // Multiple chars found in this token, but whitespace will only be
            // one per token.
            return false;
        }

        return first_char.is_whitespace();
    }

    pub fn is_word(&self) -> bool {
        let mut chars_iterator = self.token.chars();
        let first_char = chars_iterator.next().unwrap();
        let second_char = chars_iterator.next();
        if second_char.is_some() {
            // We consist of multiple characters, that means we are a word
            return true;
        }

        // If we get here, it means our token consists of one character only. If
        // that single character is alphanumeric, we are a word, otherwise not.
        return first_char.is_alphanumeric();
    }
}

pub struct TokenCollector {
    line_prefix: StyledToken,
    tokens: Vec<StyledToken>,
    rendered: bool,
}

impl Style {
    #[must_use]
    pub fn is_inverse(&self) -> bool {
        match self {
            Style::OldInverse | Style::NewInverse | Style::Error => {
                return true;
            }
            _ => {
                return false;
            }
        }
    }

    pub fn inverted(&self) -> Style {
        return match self {
            Style::Old => Style::OldInverse,
            Style::New => Style::NewInverse,
            Style::OldInverse => Style::OldInverse,
            Style::NewInverse => Style::NewInverse,
            Style::Error => Style::Error,
        };
    }

    #[must_use]
    pub fn color<'a>(&self) -> &'a str {
        match self {
            Style::Old => {
                return OLD;
            }
            Style::OldInverse => {
                return OLD;
            }
            Style::New => {
                return NEW;
            }
            Style::NewInverse => {
                return NEW;
            }
            Style::Error => {
                return ERROR;
            }
        }
    }
}

impl TokenCollector {
    #[must_use]
    pub fn create(line_prefix: StyledToken) -> Self {
        return TokenCollector {
            line_prefix,
            tokens: Vec::new(),
            rendered: false,
        };
    }

    pub fn push(&mut self, token: StyledToken) {
        self.tokens.push(token);
    }

    #[must_use]
    fn render_row(&self, row: &mut [StyledToken]) -> String {
        let mut rendered = String::new();

        if self.line_prefix.style == Style::New {
            highlight_trailing_whitespace(row);
            highlight_nonleading_tab(row);
        }
        highlight_space_between_words(row);

        // Set inverse from prefix
        let mut is_inverse = self.line_prefix.style.is_inverse();
        if is_inverse {
            rendered.push_str(INVERSE_VIDEO);
        }

        // Set line color from prefix
        let mut color = self.line_prefix.style.color();
        rendered.push_str(self.line_prefix.style.color());

        // Render prefix
        rendered.push_str(&self.line_prefix.token);

        for token in row {
            if token.style.is_inverse() && !is_inverse {
                rendered.push_str(INVERSE_VIDEO);
            }
            if is_inverse && !token.style.is_inverse() {
                rendered.push_str(NOT_INVERSE_VIDEO);
            }
            is_inverse = token.style.is_inverse();

            if token.style.color() != color {
                rendered.push_str(token.style.color());
                color = token.style.color();
            }

            rendered.push_str(&token.token);
        }

        rendered.push_str(NORMAL);

        return rendered;
    }

    /// Render all the tokens into a (most of the time multiline) string
    #[must_use]
    pub fn render(&mut self) -> String {
        assert!(!self.rendered);
        let mut current_row: Vec<StyledToken> = Vec::new();
        let mut rendered = String::new();

        let mut tokens = std::mem::take(&mut self.tokens);
        censor_multi_line_highlights(&mut tokens);

        for token in tokens {
            if token.token == "\n" {
                let rendered_row = &self.render_row(&mut current_row);
                rendered.push_str(rendered_row);
                rendered.push('\n');
                current_row.clear();
                continue;
            }

            current_row.push(token);
        }

        if !current_row.is_empty() {
            let rendered_row = &self.render_row(&mut current_row);
            rendered.push_str(rendered_row);
        }

        self.rendered = true;
        return rendered;
    }
}

fn censor_multi_line_highlights(_: &mut [StyledToken]) {
    // FIXME: Code missing here
}

fn highlight_trailing_whitespace(row: &mut [StyledToken]) {
    for token in row.iter_mut().rev() {
        if !token.is_whitespace() {
            return;
        }

        token.style = Style::Error;
    }
}

fn highlight_nonleading_tab(row: &mut [StyledToken]) {
    let mut token_iter = row.iter_mut();

    // Skip leading TABs
    loop {
        let next = token_iter.next();
        if next.is_none() {
            // Done!
            return;
        }

        let token = next.unwrap();
        if token.token != "\t" {
            // Not a TAB, this means we're out of skipping the leading TABs
            break;
        }
    }

    // Scan the rest of the line for non-leading TABs
    for token in token_iter {
        if token.token != "\t" {
            // Not a TAB, never mind
            continue;
        }

        // Non-leading TAB, mark it!
        token.style = Style::Error;
    }
}

/// Highlight single space between two highlighted words
fn highlight_space_between_words(row: &mut [StyledToken]) {
    enum FoundState {
        Nothing,
        HighlightedWord,
        WordSpace,
    }

    let mut found_state = FoundState::Nothing;
    let mut previous_token: Option<&mut StyledToken> = None;
    for token in row.iter_mut() {
        match found_state {
            FoundState::Nothing => {
                if token.style.is_inverse() && token.is_word() {
                    // Found "Monkey"
                    found_state = FoundState::HighlightedWord;
                }
            }

            FoundState::HighlightedWord => {
                if token.is_whitespace() {
                    // Found "Monkey " (note trailing space)
                    found_state = FoundState::WordSpace;
                } else if token.style.is_inverse() && token.is_word() {
                    found_state = FoundState::HighlightedWord;
                } else {
                    found_state = FoundState::Nothing;
                }
            }

            FoundState::WordSpace => {
                if token.style.is_inverse() && token.is_word() {
                    // Found "Monkey Dance"
                    if let Some(_previous_token) = previous_token {
                        _previous_token.style = _previous_token.style.inverted();
                    }

                    found_state = FoundState::HighlightedWord;
                } else {
                    found_state = FoundState::Nothing;
                }
            }
        }

        previous_token = Some(token);
    }
}

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

    #[cfg(test)]
    use pretty_assertions::assert_eq;

    #[test]
    fn test_basic() {
        let mut test_me = TokenCollector::create(StyledToken {
            token: "+".to_string(),
            style: Style::New,
        });

        test_me.push(StyledToken {
            token: "hej".to_string(),
            style: Style::New,
        });

        test_me.push(StyledToken {
            token: "\n".to_string(),
            style: Style::New,
        });

        let rendered = test_me.render();
        assert_eq!(rendered, format!("{}+hej{}\n", NEW, NORMAL));
    }

    #[test]
    fn test_add_trailing_whitespace() {
        // Just a whitespace
        let mut row = [StyledToken::new(" ".to_string(), Style::New)];
        highlight_trailing_whitespace(&mut row);
        assert_eq!(row, [StyledToken::new(" ".to_string(), Style::Error)]);

        // Trailing whitespace
        let mut row = [
            StyledToken::new("x".to_string(), Style::New),
            StyledToken::new(" ".to_string(), Style::New),
        ];
        highlight_trailing_whitespace(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("x".to_string(), Style::New),
                StyledToken::new(" ".to_string(), Style::Error),
            ]
        );

        // Leading whitespace
        let mut row = [
            StyledToken::new(" ".to_string(), Style::New),
            StyledToken::new("x".to_string(), Style::New),
        ];
        highlight_trailing_whitespace(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new(" ".to_string(), Style::New),
                StyledToken::new("x".to_string(), Style::New),
            ]
        );
    }

    #[test]
    fn test_removed_trailing_whitespace() {
        // It shouldn't be highlighted, just added ones should
        let mut test_me = TokenCollector::create(StyledToken::new("-".to_string(), Style::Old));
        test_me.push(StyledToken::new(" ".to_string(), Style::Old));
        let actual = test_me.render();

        assert_eq!(actual, format!("{}- {}", OLD, NORMAL));
    }

    #[test]
    fn test_add_nonleading_tab() {
        // Trailing TAB
        let mut row = [
            StyledToken::new("x".to_string(), Style::New),
            StyledToken::new("\t".to_string(), Style::New),
        ];
        highlight_nonleading_tab(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("x".to_string(), Style::New),
                StyledToken::new("\t".to_string(), Style::Error),
            ]
        );

        // Middle TAB
        let mut row = [
            StyledToken::new("x".to_string(), Style::New),
            StyledToken::new("\t".to_string(), Style::New),
            StyledToken::new("y".to_string(), Style::New),
        ];
        highlight_nonleading_tab(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("x".to_string(), Style::New),
                StyledToken::new("\t".to_string(), Style::Error),
                StyledToken::new("y".to_string(), Style::New),
            ]
        );

        // Leading TAB (don't highlight)
        let mut row = [
            StyledToken::new("\t".to_string(), Style::New),
            StyledToken::new("x".to_string(), Style::New),
        ];
        highlight_nonleading_tab(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("\t".to_string(), Style::New),
                StyledToken::new("x".to_string(), Style::New),
            ]
        );

        // Single TAB (don't highlight because it is leading)
        let mut row = [StyledToken::new("\t".to_string(), Style::New)];
        highlight_nonleading_tab(&mut row);
        assert_eq!(row, [StyledToken::new("\t".to_string(), Style::New),]);
    }

    #[test]
    fn test_removed_nonleading_tab() {
        // It shouldn't be highlighted, just added ones should
        let mut test_me = TokenCollector::create(StyledToken::new("-".to_string(), Style::Old));
        test_me.push(StyledToken::new("x".to_string(), Style::Old));
        test_me.push(StyledToken::new("\t".to_string(), Style::Old));
        let actual = test_me.render();

        assert_eq!(actual, format!("{}-x\t{}", OLD, NORMAL));
    }

    #[test]
    fn test_highlight_space_between_words() {
        let mut row = [
            StyledToken::new("Monkey".to_string(), Style::NewInverse),
            StyledToken::new(" ".to_string(), Style::New),
            StyledToken::new("Dance".to_string(), Style::NewInverse),
        ];

        highlight_space_between_words(&mut row);

        assert_eq!(
            row,
            [
                StyledToken::new("Monkey".to_string(), Style::NewInverse),
                StyledToken::new(" ".to_string(), Style::NewInverse),
                StyledToken::new("Dance".to_string(), Style::NewInverse),
            ]
        );
    }

    /// "_" is an unhighlighted token, "." is a highlighted token, "n" is an
    /// unhighlighted newline and "N" is a highlighted newline.
    fn test_censor_multiline_highlighting(input: &str, expected_output: &str) {
        let mut row = Vec::new();
        for c in input.chars() {
            let mut token = "x".to_string();
            if c == 'n' || c == 'N' {
                token = '\n'.to_string();
            }

            let mut style = Style::New;
            if c == '.' || c == 'N' {
                style = Style::NewInverse;
            }

            row.push(StyledToken { token, style });
        }

        censor_multi_line_highlights(&mut row);

        let mut actual_output = String::new();
        for token in row {
            if token.token != "\n" && !token.style.is_inverse() {
                actual_output.push('_');
            } else if token.token != "\n" && token.style.is_inverse() {
                actual_output.push('.');
            } else if token.token == "\n" && !token.style.is_inverse() {
                actual_output.push('n');
            } else if token.token == "\n" && token.style.is_inverse() {
                actual_output.push('N');
            } else {
                panic!("How did we get here?");
            }
        }

        assert_eq!(actual_output, expected_output);
    }

    #[test]
    fn test_censor_multiline_highlights_start_at_beginning() {
        test_censor_multiline_highlighting("...n___", "___n___");
    }

    #[test]
    fn test_censor_multiline_highlights_at_end() {
        test_censor_multiline_highlighting("___n...", "___n___");
    }

    #[test]
    fn test_censor_multiline_highlights_midline() {
        test_censor_multiline_highlighting("___n...n___", "___n___n___");
    }

    #[test]
    fn test_censor_multiline_highlights_partial_lines() {
        // Highlighted line parts should not be censored
        test_censor_multiline_highlighting("_.._", "_.._");
    }
}
<commit_msg>More tests<commit_after>use crate::constants::*;

#[derive(Clone, Debug, PartialEq)]
pub enum Style {
    Old,
    OldInverse,
    New,
    NewInverse,
    Error,
}

#[derive(Clone, Debug, PartialEq)]
pub struct StyledToken {
    token: String,
    style: Style,
}

impl StyledToken {
    pub fn new(token: String, style: Style) -> StyledToken {
        return StyledToken { token, style };
    }

    pub fn is_whitespace(&self) -> bool {
        let mut chars_iterator = self.token.chars();
        let first_char = chars_iterator.next().unwrap();
        if chars_iterator.next().is_some() {
            // Multiple chars found in this token, but whitespace will only be
            // one per token.
            return false;
        }

        return first_char.is_whitespace();
    }

    pub fn is_word(&self) -> bool {
        let mut chars_iterator = self.token.chars();
        let first_char = chars_iterator.next().unwrap();
        let second_char = chars_iterator.next();
        if second_char.is_some() {
            // We consist of multiple characters, that means we are a word
            return true;
        }

        // If we get here, it means our token consists of one character only. If
        // that single character is alphanumeric, we are a word, otherwise not.
        return first_char.is_alphanumeric();
    }
}

pub struct TokenCollector {
    line_prefix: StyledToken,
    tokens: Vec<StyledToken>,
    rendered: bool,
}

impl Style {
    #[must_use]
    pub fn is_inverse(&self) -> bool {
        match self {
            Style::OldInverse | Style::NewInverse | Style::Error => {
                return true;
            }
            _ => {
                return false;
            }
        }
    }

    pub fn inverted(&self) -> Style {
        return match self {
            Style::Old => Style::OldInverse,
            Style::New => Style::NewInverse,
            Style::OldInverse => Style::OldInverse,
            Style::NewInverse => Style::NewInverse,
            Style::Error => Style::Error,
        };
    }

    #[must_use]
    pub fn color<'a>(&self) -> &'a str {
        match self {
            Style::Old => {
                return OLD;
            }
            Style::OldInverse => {
                return OLD;
            }
            Style::New => {
                return NEW;
            }
            Style::NewInverse => {
                return NEW;
            }
            Style::Error => {
                return ERROR;
            }
        }
    }
}

impl TokenCollector {
    #[must_use]
    pub fn create(line_prefix: StyledToken) -> Self {
        return TokenCollector {
            line_prefix,
            tokens: Vec::new(),
            rendered: false,
        };
    }

    pub fn push(&mut self, token: StyledToken) {
        self.tokens.push(token);
    }

    #[must_use]
    fn render_row(&self, row: &mut [StyledToken]) -> String {
        let mut rendered = String::new();

        if self.line_prefix.style == Style::New {
            highlight_trailing_whitespace(row);
            highlight_nonleading_tab(row);
        }
        highlight_space_between_words(row);

        // Set inverse from prefix
        let mut is_inverse = self.line_prefix.style.is_inverse();
        if is_inverse {
            rendered.push_str(INVERSE_VIDEO);
        }

        // Set line color from prefix
        let mut color = self.line_prefix.style.color();
        rendered.push_str(self.line_prefix.style.color());

        // Render prefix
        rendered.push_str(&self.line_prefix.token);

        for token in row {
            if token.style.is_inverse() && !is_inverse {
                rendered.push_str(INVERSE_VIDEO);
            }
            if is_inverse && !token.style.is_inverse() {
                rendered.push_str(NOT_INVERSE_VIDEO);
            }
            is_inverse = token.style.is_inverse();

            if token.style.color() != color {
                rendered.push_str(token.style.color());
                color = token.style.color();
            }

            rendered.push_str(&token.token);
        }

        rendered.push_str(NORMAL);

        return rendered;
    }

    /// Render all the tokens into a (most of the time multiline) string
    #[must_use]
    pub fn render(&mut self) -> String {
        assert!(!self.rendered);
        let mut current_row: Vec<StyledToken> = Vec::new();
        let mut rendered = String::new();

        let mut tokens = std::mem::take(&mut self.tokens);

        // FIXME: Do highlight_space_between_words() first?
        censor_multi_line_highlights(&mut tokens);

        for token in tokens {
            if token.token == "\n" {
                let rendered_row = &self.render_row(&mut current_row);
                rendered.push_str(rendered_row);
                rendered.push('\n');
                current_row.clear();
                continue;
            }

            current_row.push(token);
        }

        if !current_row.is_empty() {
            let rendered_row = &self.render_row(&mut current_row);
            rendered.push_str(rendered_row);
        }

        self.rendered = true;
        return rendered;
    }
}

fn censor_multi_line_highlights(_: &mut [StyledToken]) {
    // FIXME: Code missing here
}

fn highlight_trailing_whitespace(row: &mut [StyledToken]) {
    for token in row.iter_mut().rev() {
        if !token.is_whitespace() {
            return;
        }

        token.style = Style::Error;
    }
}

fn highlight_nonleading_tab(row: &mut [StyledToken]) {
    let mut token_iter = row.iter_mut();

    // Skip leading TABs
    loop {
        let next = token_iter.next();
        if next.is_none() {
            // Done!
            return;
        }

        let token = next.unwrap();
        if token.token != "\t" {
            // Not a TAB, this means we're out of skipping the leading TABs
            break;
        }
    }

    // Scan the rest of the line for non-leading TABs
    for token in token_iter {
        if token.token != "\t" {
            // Not a TAB, never mind
            continue;
        }

        // Non-leading TAB, mark it!
        token.style = Style::Error;
    }
}

/// Highlight single space between two highlighted words
fn highlight_space_between_words(row: &mut [StyledToken]) {
    enum FoundState {
        Nothing,
        HighlightedWord,
        WordSpace,
    }

    let mut found_state = FoundState::Nothing;
    let mut previous_token: Option<&mut StyledToken> = None;
    for token in row.iter_mut() {
        match found_state {
            FoundState::Nothing => {
                if token.style.is_inverse() && token.is_word() {
                    // Found "Monkey"
                    found_state = FoundState::HighlightedWord;
                }
            }

            FoundState::HighlightedWord => {
                if token.is_whitespace() {
                    // Found "Monkey " (note trailing space)
                    found_state = FoundState::WordSpace;
                } else if token.style.is_inverse() && token.is_word() {
                    found_state = FoundState::HighlightedWord;
                } else {
                    found_state = FoundState::Nothing;
                }
            }

            FoundState::WordSpace => {
                if token.style.is_inverse() && token.is_word() {
                    // Found "Monkey Dance"
                    if let Some(_previous_token) = previous_token {
                        _previous_token.style = _previous_token.style.inverted();
                    }

                    found_state = FoundState::HighlightedWord;
                } else {
                    found_state = FoundState::Nothing;
                }
            }
        }

        previous_token = Some(token);
    }
}

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

    #[cfg(test)]
    use pretty_assertions::assert_eq;

    #[test]
    fn test_basic() {
        let mut test_me = TokenCollector::create(StyledToken {
            token: "+".to_string(),
            style: Style::New,
        });

        test_me.push(StyledToken {
            token: "hej".to_string(),
            style: Style::New,
        });

        test_me.push(StyledToken {
            token: "\n".to_string(),
            style: Style::New,
        });

        let rendered = test_me.render();
        assert_eq!(rendered, format!("{}+hej{}\n", NEW, NORMAL));
    }

    #[test]
    fn test_add_trailing_whitespace() {
        // Just a whitespace
        let mut row = [StyledToken::new(" ".to_string(), Style::New)];
        highlight_trailing_whitespace(&mut row);
        assert_eq!(row, [StyledToken::new(" ".to_string(), Style::Error)]);

        // Trailing whitespace
        let mut row = [
            StyledToken::new("x".to_string(), Style::New),
            StyledToken::new(" ".to_string(), Style::New),
        ];
        highlight_trailing_whitespace(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("x".to_string(), Style::New),
                StyledToken::new(" ".to_string(), Style::Error),
            ]
        );

        // Leading whitespace
        let mut row = [
            StyledToken::new(" ".to_string(), Style::New),
            StyledToken::new("x".to_string(), Style::New),
        ];
        highlight_trailing_whitespace(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new(" ".to_string(), Style::New),
                StyledToken::new("x".to_string(), Style::New),
            ]
        );
    }

    #[test]
    fn test_removed_trailing_whitespace() {
        // It shouldn't be highlighted, just added ones should
        let mut test_me = TokenCollector::create(StyledToken::new("-".to_string(), Style::Old));
        test_me.push(StyledToken::new(" ".to_string(), Style::Old));
        let actual = test_me.render();

        assert_eq!(actual, format!("{}- {}", OLD, NORMAL));
    }

    #[test]
    fn test_add_nonleading_tab() {
        // Trailing TAB
        let mut row = [
            StyledToken::new("x".to_string(), Style::New),
            StyledToken::new("\t".to_string(), Style::New),
        ];
        highlight_nonleading_tab(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("x".to_string(), Style::New),
                StyledToken::new("\t".to_string(), Style::Error),
            ]
        );

        // Middle TAB
        let mut row = [
            StyledToken::new("x".to_string(), Style::New),
            StyledToken::new("\t".to_string(), Style::New),
            StyledToken::new("y".to_string(), Style::New),
        ];
        highlight_nonleading_tab(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("x".to_string(), Style::New),
                StyledToken::new("\t".to_string(), Style::Error),
                StyledToken::new("y".to_string(), Style::New),
            ]
        );

        // Leading TAB (don't highlight)
        let mut row = [
            StyledToken::new("\t".to_string(), Style::New),
            StyledToken::new("x".to_string(), Style::New),
        ];
        highlight_nonleading_tab(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("\t".to_string(), Style::New),
                StyledToken::new("x".to_string(), Style::New),
            ]
        );

        // Single TAB (don't highlight because it is leading)
        let mut row = [StyledToken::new("\t".to_string(), Style::New)];
        highlight_nonleading_tab(&mut row);
        assert_eq!(row, [StyledToken::new("\t".to_string(), Style::New),]);
    }

    #[test]
    fn test_removed_nonleading_tab() {
        // It shouldn't be highlighted, just added ones should
        let mut test_me = TokenCollector::create(StyledToken::new("-".to_string(), Style::Old));
        test_me.push(StyledToken::new("x".to_string(), Style::Old));
        test_me.push(StyledToken::new("\t".to_string(), Style::Old));
        let actual = test_me.render();

        assert_eq!(actual, format!("{}-x\t{}", OLD, NORMAL));
    }

    #[test]
    fn test_highlight_space_between_words() {
        let mut row = [
            StyledToken::new("Monkey".to_string(), Style::NewInverse),
            StyledToken::new(" ".to_string(), Style::New),
            StyledToken::new("Dance".to_string(), Style::NewInverse),
        ];

        highlight_space_between_words(&mut row);

        assert_eq!(
            row,
            [
                StyledToken::new("Monkey".to_string(), Style::NewInverse),
                StyledToken::new(" ".to_string(), Style::NewInverse),
                StyledToken::new("Dance".to_string(), Style::NewInverse),
            ]
        );
    }

    /// "_" is an unhighlighted token, "." is a highlighted token, "n" is an
    /// unhighlighted newline and "N" is a highlighted newline.
    fn test_censor_multiline_highlighting(input: &str, expected_output: &str) {
        let mut row = Vec::new();
        for c in input.chars() {
            let mut token = "x".to_string();
            if c == 'n' || c == 'N' {
                token = '\n'.to_string();
            }

            let mut style = Style::New;
            if c == '.' || c == 'N' {
                style = Style::NewInverse;
            }

            row.push(StyledToken { token, style });
        }

        censor_multi_line_highlights(&mut row);

        let mut actual_output = String::new();
        for token in row {
            if token.token != "\n" && !token.style.is_inverse() {
                actual_output.push('_');
            } else if token.token != "\n" && token.style.is_inverse() {
                actual_output.push('.');
            } else if token.token == "\n" && !token.style.is_inverse() {
                actual_output.push('n');
            } else if token.token == "\n" && token.style.is_inverse() {
                actual_output.push('N');
            } else {
                panic!("How did we get here?");
            }
        }

        assert_eq!(actual_output, expected_output);
    }

    #[test]
    fn test_censor_multiline_highlights_start_at_beginning() {
        test_censor_multiline_highlighting("...n___", "___n___");
    }

    #[test]
    fn test_censor_multiline_highlights_at_end() {
        test_censor_multiline_highlighting("___n...", "___n___");
    }

    #[test]
    fn test_censor_multiline_highlights_midline() {
        test_censor_multiline_highlighting("___n...n___", "___n___n___");
    }

    #[test]
    fn test_censor_multiline_highlights_edgecases() {
        test_censor_multiline_highlighting("n.n", "n_n");
    }

    #[test]
    fn test_censor_multiline_highlights_dont_highlight() {
        // Highlighted line parts should not be censored
        test_censor_multiline_highlighting("_.._n.__.", "_.._n.__.");
    }
}
<|endoftext|>


<commit_before>extern crate libc;

use array::Array;
use dim4::Dim4;
use defines::{AfError, RandomEngineType};
use error::HANDLE_ERROR;
use self::libc::{uint8_t, c_int, c_uint};
use util::{DimT, HasAfEnum, MutAfArray, MutRandEngine, RandEngine, Uintl};

#[allow(dead_code)]
extern {
    fn af_set_seed(seed: Uintl) -> c_int;
    fn af_get_seed(seed: *mut Uintl) -> c_int;

    fn af_randu(out: MutAfArray, ndims: c_uint, dims: *const DimT, afdtype: uint8_t) -> c_int;
    fn af_randn(out: MutAfArray, ndims: c_uint, dims: *const DimT, afdtype: uint8_t) -> c_int;

    fn af_create_random_engine(engine: MutRandEngine, rtype: uint8_t, seed: Uintl) -> c_int;
    fn af_retain_random_engine(engine: MutRandEngine, inputEngine: RandEngine) -> c_int;
    fn af_random_engine_set_type(engine: MutRandEngine, rtpye: uint8_t) -> c_int;
    fn af_random_engine_get_type(rtype: *mut uint8_t, engine: RandEngine) -> c_int;
    fn af_random_engine_set_seed(engine: MutRandEngine, seed: Uintl) -> c_int;
    fn af_random_engine_get_seed(seed: *mut Uintl, engine: RandEngine) -> c_int;
    fn af_release_random_engine(engine: RandEngine) -> c_int;

    fn af_get_default_random_engine(engine: MutRandEngine) -> c_int;
    fn af_set_default_random_engine_type(rtype: uint8_t) -> c_int;

    fn af_random_uniform(out: MutAfArray, ndims: c_uint, dims: *const DimT,
                         aftype: uint8_t, engine: RandEngine) -> c_int;
    fn af_random_normal(out: MutAfArray, ndims: c_uint, dims: *const DimT,
                        aftype: uint8_t, engine: RandEngine) -> c_int;
}

/// Set seed for random number generation
pub fn set_seed(seed: u64) {
    unsafe {
        let err_val = af_set_seed(seed as Uintl);
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Get the seed of random number generator
#[allow(unused_mut)]
pub fn get_seed() -> u64 {
    unsafe {
        let mut temp: u64 = 0;
        let err_val = af_get_seed(&mut temp as *mut Uintl);
        HANDLE_ERROR(AfError::from(err_val));
        temp
    }
}

macro_rules! data_gen_def {
    ($doc_str: expr, $fn_name:ident, $ffi_name: ident) => (
        #[doc=$doc_str]
        ///
        ///# Parameters
        ///
        /// - `dims` is the output dimensions
        ///
        ///# Return Values
        ///
        /// An Array with random values.
        #[allow(unused_mut)]
        pub fn $fn_name<T: HasAfEnum>(dims: Dim4) -> Array {
            unsafe {
                let aftype = T::get_af_dtype();
                let mut temp: i64 = 0;
                let err_val = $ffi_name(&mut temp as MutAfArray,
                                        dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                                        aftype as uint8_t);
                HANDLE_ERROR(AfError::from(err_val));
                Array::from(temp)
            }
        }
    )
}

data_gen_def!("Create random numbers from uniform distribution", randu, af_randu);
data_gen_def!("Create random numbers from normal distribution", randn, af_randn);

/// Random number generator engine
///
/// This is a wrapper for ArrayFire's native random number generator engine.
pub struct RandomEngine {
    handle: i64,
}

/// Used for creating RandomEngine object from native resource id
impl From<i64> for RandomEngine {
    fn from(t: i64) -> RandomEngine {
        RandomEngine {handle: t}
    }
}

impl RandomEngine {
    /// Create a new random engine object
    ///
    /// # Parameters
    ///
    /// - `rengine` can be value of [RandomEngineType](./enum.RandomEngineType.html) enum.
    /// - `seed` is the initial seed value
    ///
    /// # Return Values
    ///
    /// A object of type RandomEngine
    pub fn new(rengine: RandomEngineType, seed: Option<u64>) -> RandomEngine {
        unsafe {
            let mut temp: i64 = 0;
            let err_val = af_create_random_engine(&mut temp as MutRandEngine, rengine as uint8_t,
                                                  match seed {Some(s) => s, None => 0} as Uintl);
            HANDLE_ERROR(AfError::from(err_val));
            RandomEngine::from(temp)
        }
    }

    /// Get random engine type
    pub fn get_type(&self) -> RandomEngineType {
        unsafe {
            let mut temp: u8 = 0;
            let err_val = af_random_engine_get_type(&mut temp as *mut uint8_t,
                                                    self.handle as RandEngine);
            HANDLE_ERROR(AfError::from(err_val));
            RandomEngineType::from(temp as i32)
        }
    }

    /// Get random engine type
    pub fn set_type(&mut self, engine_type: RandomEngineType) {
        unsafe {
            let err_val = af_random_engine_set_type(&mut self.handle as MutRandEngine,
                                                    engine_type as uint8_t);
            HANDLE_ERROR(AfError::from(err_val));
        }
    }

    /// Set seed for random engine
    pub fn set_seed(&mut self, seed: u64) {
        unsafe {
            let err_val = af_random_engine_set_seed(&mut self.handle as MutRandEngine,
                                                    seed as Uintl);
            HANDLE_ERROR(AfError::from(err_val));
        }
    }

    /// Get seed of the random engine
    pub fn get_seed(&self) -> u64 {
        unsafe {
            let mut seed: u64 = 0;
            let err_val = af_random_engine_get_seed(&mut seed as *mut Uintl, self.handle as RandEngine);
            HANDLE_ERROR(AfError::from(err_val));
            seed
        }
    }

    /// Returns the native FFI handle for Rust object `RandomEngine`
    pub fn get(&self) -> i64 {
        self.handle
    }
}

/// Increment reference count of RandomEngine's native resource
impl Clone for RandomEngine {
    fn clone(&self) -> RandomEngine {
        unsafe {
            let mut temp: i64 = 0;
            let err_val = af_retain_random_engine(&mut temp as MutRandEngine, self.handle as RandEngine);
            HANDLE_ERROR(AfError::from(err_val));
            RandomEngine::from(temp)
        }
    }
}

/// Free RandomEngine's native resource
impl Drop for RandomEngine {
    fn drop(&mut self) {
        unsafe {
            let err_val = af_release_random_engine(self.handle as RandEngine);
            HANDLE_ERROR(AfError::from(err_val));
        }
    }
}

/// Get default random engine
pub fn get_default_random_engine() -> RandomEngine {
    unsafe {
        let mut temp : i64 = 0;
        let mut err_val = af_get_default_random_engine(&mut temp as MutRandEngine);
        HANDLE_ERROR(AfError::from(err_val));
        let mut handle : i64 = 0;
        err_val = af_retain_random_engine(&mut handle as MutRandEngine, temp as RandEngine);
        HANDLE_ERROR(AfError::from(err_val));
        RandomEngine::from(handle)
    }
}

/// Set the random engine type for default random number generator
///
/// # Parameters
///
/// - `rtype` can take one of the values of enum [RandomEngineType](./enum.RandomEngineType.html)
pub fn set_default_random_engine_type(rtype: RandomEngineType) {
    unsafe {
        let err_val = af_set_default_random_engine_type(rtype as uint8_t);
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Generate array of uniform numbers using a random engine
///
/// # Parameters
///
/// - `dims` is output array dimensions
/// - `engine` is an object of type [RandomEngine](./struct.RandomEngine.html)
///
/// # Return Values
///
/// An Array with uniform numbers generated using random engine
pub fn random_uniform<T: HasAfEnum>(dims: Dim4, engine: RandomEngine) -> Array {
    unsafe {
        let aftype = T::get_af_dtype();
        let mut temp : i64 = 0;
        let err_val = af_random_uniform(&mut temp as MutAfArray,
                                        dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                                        aftype as uint8_t, engine.get() as RandEngine);
        HANDLE_ERROR(AfError::from(err_val));
        Array::from(temp)
    }
}

/// Generate array of normal numbers using a random engine
///
/// # Parameters
///
/// - `dims` is output array dimensions
/// - `engine` is an object of type [RandomEngine](./struct.RandomEngine.html)
///
/// # Return Values
///
/// An Array with normal numbers generated using random engine
pub fn random_normal<T: HasAfEnum>(dims: Dim4, engine: RandomEngine) -> Array {
    unsafe {
        let aftype = T::get_af_dtype();
        let mut temp : i64 = 0;
        let err_val = af_random_normal(&mut temp as MutAfArray,
                                       dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                                       aftype as uint8_t, engine.get() as RandEngine);
        HANDLE_ERROR(AfError::from(err_val));
        Array::from(temp)
    }
}
<commit_msg>Take a shared reference instead of ownership in random::random_uniform() and random::random_normal() (fix #162)<commit_after>extern crate libc;

use array::Array;
use dim4::Dim4;
use defines::{AfError, RandomEngineType};
use error::HANDLE_ERROR;
use self::libc::{uint8_t, c_int, c_uint};
use util::{DimT, HasAfEnum, MutAfArray, MutRandEngine, RandEngine, Uintl};

#[allow(dead_code)]
extern {
    fn af_set_seed(seed: Uintl) -> c_int;
    fn af_get_seed(seed: *mut Uintl) -> c_int;

    fn af_randu(out: MutAfArray, ndims: c_uint, dims: *const DimT, afdtype: uint8_t) -> c_int;
    fn af_randn(out: MutAfArray, ndims: c_uint, dims: *const DimT, afdtype: uint8_t) -> c_int;

    fn af_create_random_engine(engine: MutRandEngine, rtype: uint8_t, seed: Uintl) -> c_int;
    fn af_retain_random_engine(engine: MutRandEngine, inputEngine: RandEngine) -> c_int;
    fn af_random_engine_set_type(engine: MutRandEngine, rtpye: uint8_t) -> c_int;
    fn af_random_engine_get_type(rtype: *mut uint8_t, engine: RandEngine) -> c_int;
    fn af_random_engine_set_seed(engine: MutRandEngine, seed: Uintl) -> c_int;
    fn af_random_engine_get_seed(seed: *mut Uintl, engine: RandEngine) -> c_int;
    fn af_release_random_engine(engine: RandEngine) -> c_int;

    fn af_get_default_random_engine(engine: MutRandEngine) -> c_int;
    fn af_set_default_random_engine_type(rtype: uint8_t) -> c_int;

    fn af_random_uniform(out: MutAfArray, ndims: c_uint, dims: *const DimT,
                         aftype: uint8_t, engine: RandEngine) -> c_int;
    fn af_random_normal(out: MutAfArray, ndims: c_uint, dims: *const DimT,
                        aftype: uint8_t, engine: RandEngine) -> c_int;
}

/// Set seed for random number generation
pub fn set_seed(seed: u64) {
    unsafe {
        let err_val = af_set_seed(seed as Uintl);
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Get the seed of random number generator
#[allow(unused_mut)]
pub fn get_seed() -> u64 {
    unsafe {
        let mut temp: u64 = 0;
        let err_val = af_get_seed(&mut temp as *mut Uintl);
        HANDLE_ERROR(AfError::from(err_val));
        temp
    }
}

macro_rules! data_gen_def {
    ($doc_str: expr, $fn_name:ident, $ffi_name: ident) => (
        #[doc=$doc_str]
        ///
        ///# Parameters
        ///
        /// - `dims` is the output dimensions
        ///
        ///# Return Values
        ///
        /// An Array with random values.
        #[allow(unused_mut)]
        pub fn $fn_name<T: HasAfEnum>(dims: Dim4) -> Array {
            unsafe {
                let aftype = T::get_af_dtype();
                let mut temp: i64 = 0;
                let err_val = $ffi_name(&mut temp as MutAfArray,
                                        dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                                        aftype as uint8_t);
                HANDLE_ERROR(AfError::from(err_val));
                Array::from(temp)
            }
        }
    )
}

data_gen_def!("Create random numbers from uniform distribution", randu, af_randu);
data_gen_def!("Create random numbers from normal distribution", randn, af_randn);

/// Random number generator engine
///
/// This is a wrapper for ArrayFire's native random number generator engine.
pub struct RandomEngine {
    handle: i64,
}

/// Used for creating RandomEngine object from native resource id
impl From<i64> for RandomEngine {
    fn from(t: i64) -> RandomEngine {
        RandomEngine {handle: t}
    }
}

impl RandomEngine {
    /// Create a new random engine object
    ///
    /// # Parameters
    ///
    /// - `rengine` can be value of [RandomEngineType](./enum.RandomEngineType.html) enum.
    /// - `seed` is the initial seed value
    ///
    /// # Return Values
    ///
    /// A object of type RandomEngine
    pub fn new(rengine: RandomEngineType, seed: Option<u64>) -> RandomEngine {
        unsafe {
            let mut temp: i64 = 0;
            let err_val = af_create_random_engine(&mut temp as MutRandEngine, rengine as uint8_t,
                                                  match seed {Some(s) => s, None => 0} as Uintl);
            HANDLE_ERROR(AfError::from(err_val));
            RandomEngine::from(temp)
        }
    }

    /// Get random engine type
    pub fn get_type(&self) -> RandomEngineType {
        unsafe {
            let mut temp: u8 = 0;
            let err_val = af_random_engine_get_type(&mut temp as *mut uint8_t,
                                                    self.handle as RandEngine);
            HANDLE_ERROR(AfError::from(err_val));
            RandomEngineType::from(temp as i32)
        }
    }

    /// Get random engine type
    pub fn set_type(&mut self, engine_type: RandomEngineType) {
        unsafe {
            let err_val = af_random_engine_set_type(&mut self.handle as MutRandEngine,
                                                    engine_type as uint8_t);
            HANDLE_ERROR(AfError::from(err_val));
        }
    }

    /// Set seed for random engine
    pub fn set_seed(&mut self, seed: u64) {
        unsafe {
            let err_val = af_random_engine_set_seed(&mut self.handle as MutRandEngine,
                                                    seed as Uintl);
            HANDLE_ERROR(AfError::from(err_val));
        }
    }

    /// Get seed of the random engine
    pub fn get_seed(&self) -> u64 {
        unsafe {
            let mut seed: u64 = 0;
            let err_val = af_random_engine_get_seed(&mut seed as *mut Uintl, self.handle as RandEngine);
            HANDLE_ERROR(AfError::from(err_val));
            seed
        }
    }

    /// Returns the native FFI handle for Rust object `RandomEngine`
    pub fn get(&self) -> i64 {
        self.handle
    }
}

/// Increment reference count of RandomEngine's native resource
impl Clone for RandomEngine {
    fn clone(&self) -> RandomEngine {
        unsafe {
            let mut temp: i64 = 0;
            let err_val = af_retain_random_engine(&mut temp as MutRandEngine, self.handle as RandEngine);
            HANDLE_ERROR(AfError::from(err_val));
            RandomEngine::from(temp)
        }
    }
}

/// Free RandomEngine's native resource
impl Drop for RandomEngine {
    fn drop(&mut self) {
        unsafe {
            let err_val = af_release_random_engine(self.handle as RandEngine);
            HANDLE_ERROR(AfError::from(err_val));
        }
    }
}

/// Get default random engine
pub fn get_default_random_engine() -> RandomEngine {
    unsafe {
        let mut temp : i64 = 0;
        let mut err_val = af_get_default_random_engine(&mut temp as MutRandEngine);
        HANDLE_ERROR(AfError::from(err_val));
        let mut handle : i64 = 0;
        err_val = af_retain_random_engine(&mut handle as MutRandEngine, temp as RandEngine);
        HANDLE_ERROR(AfError::from(err_val));
        RandomEngine::from(handle)
    }
}

/// Set the random engine type for default random number generator
///
/// # Parameters
///
/// - `rtype` can take one of the values of enum [RandomEngineType](./enum.RandomEngineType.html)
pub fn set_default_random_engine_type(rtype: RandomEngineType) {
    unsafe {
        let err_val = af_set_default_random_engine_type(rtype as uint8_t);
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Generate array of uniform numbers using a random engine
///
/// # Parameters
///
/// - `dims` is output array dimensions
/// - `engine` is an object of type [RandomEngine](./struct.RandomEngine.html)
///
/// # Return Values
///
/// An Array with uniform numbers generated using random engine
pub fn random_uniform<T: HasAfEnum>(dims: Dim4, engine: &RandomEngine) -> Array {
    unsafe {
        let aftype = T::get_af_dtype();
        let mut temp : i64 = 0;
        let err_val = af_random_uniform(&mut temp as MutAfArray,
                                        dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                                        aftype as uint8_t, engine.get() as RandEngine);
        HANDLE_ERROR(AfError::from(err_val));
        Array::from(temp)
    }
}

/// Generate array of normal numbers using a random engine
///
/// # Parameters
///
/// - `dims` is output array dimensions
/// - `engine` is an object of type [RandomEngine](./struct.RandomEngine.html)
///
/// # Return Values
///
/// An Array with normal numbers generated using random engine
pub fn random_normal<T: HasAfEnum>(dims: Dim4, engine: &RandomEngine) -> Array {
    unsafe {
        let aftype = T::get_af_dtype();
        let mut temp : i64 = 0;
        let err_val = af_random_normal(&mut temp as MutAfArray,
                                       dims.ndims() as c_uint, dims.get().as_ptr() as *const DimT,
                                       aftype as uint8_t, engine.get() as RandEngine);
        HANDLE_ERROR(AfError::from(err_val));
        Array::from(temp)
    }
}
<|endoftext|>


<commit_before>// This file is part of Grust, GObject introspection bindings for Rust
//
// Copyright (C) 2013, 2014  Mikhail Zabaluev <mikhail.zabaluev@gmail.com>
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

use ffi;
use types::{gchar,gpointer};
use util::escape_bytestring;

use libc;
use std::error::Error;
use std::fmt;
use std::marker;
use std::mem;
use std::ops::Deref;
use std::slice;
use std::str;

const NUL: u8 = 0;

pub struct OwnedGStr {
    ptr: *const gchar,
}

pub unsafe fn parse_as_bytes<'a, T: ?Sized>(raw: *const gchar,
                                            life_anchor: &'a T)
                                           -> &'a [u8]
{
    assert!(!raw.is_null());
    let r = mem::copy_lifetime(life_anchor, &(raw as *const u8));
    slice::from_raw_buf(r, libc::strlen(raw) as usize)
}

#[inline]
pub unsafe fn parse_as_utf8<'a, T: ?Sized>(raw: *const gchar,
                                           life_anchor: &'a T)
                                          -> Result<&'a str, str::Utf8Error>
{
    str::from_utf8(parse_as_bytes(raw, life_anchor))
}

impl OwnedGStr {

    pub unsafe fn from_raw(ptr: *mut gchar) -> OwnedGStr {
        assert!(!ptr.is_null());
        OwnedGStr { ptr: ptr as *const gchar }
    }

    pub fn parse_as_bytes<'a>(&'a self) -> &'a [u8] {
        unsafe {
            let r = mem::copy_lifetime(self, &(self.ptr as *const u8));
            slice::from_raw_buf(r, libc::strlen(self.ptr) as usize)
        }
    }

    #[inline]
    pub fn parse_as_utf8<'a>(&'a self) -> Result<&'a str, str::Utf8Error> {
        str::from_utf8(self.parse_as_bytes())
    }

    #[inline]
    pub unsafe fn parse_as_utf8_unchecked<'a>(&'a self) -> &'a str {
        str::from_utf8_unchecked(self.parse_as_bytes())
    }
}

impl Deref for OwnedGStr {

    type Target = GStr;

    fn deref(&self) -> &GStr {
        unsafe { g_str_from_ptr_internal(self.ptr as *const u8, self) }
    }
}

impl Drop for OwnedGStr {
    fn drop(&mut self) {
        unsafe { ffi::g_free(self.ptr as gpointer) }
    }
}

impl Clone for OwnedGStr {
    fn clone(&self) -> OwnedGStr {
        unsafe {
            OwnedGStr::from_raw(ffi::g_strdup(self.ptr))
        }
    }
}

impl PartialEq for OwnedGStr {
    fn eq(&self, other: &OwnedGStr) -> bool {
        unsafe { libc::strcmp(self.ptr, other.ptr) == 0 }
    }
}

impl Eq for OwnedGStr { }

#[derive(Copy)]
pub struct NulError {
    pub position: usize
}

impl Error for NulError {

    fn description(&self) -> &str {
        "invalid data for C string: contains a NUL byte"
    }

    fn detail(&self) -> Option<String> {
        Some(format!("NUL at position {}", self.position))
    }
}

impl fmt::Show for NulError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "invalid data for C string: NUL at position {}",
               self.position)
    }
}

pub struct IntoGStrError {
    cause: NulError,
    bytes: Vec<u8>
}

impl IntoGStrError {

    pub fn nul_error(&self) -> &NulError {
        &self.cause
    }

    pub fn into_bytes(self) -> Vec<u8> {
        self.bytes
    }
}

impl Error for IntoGStrError {

    fn description(&self) -> &str {
        self.cause.description()
    }

    fn detail(&self) -> Option<String> {
        self.cause.detail()
    }
}

impl fmt::Show for IntoGStrError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self.cause)
    }
}

#[repr(C)]
pub struct GStr {
    head: gchar,
    marker: marker::NoCopy
}

#[repr(C)]
pub struct Utf8 {
    gstr: GStr
}

impl GStr {
    #[inline]
    pub fn as_ptr(&self) -> *const gchar {
        &self.head as *const gchar
    }

    pub fn from_static_bytes(bytes: &'static [u8]) -> &'static GStr {
        assert!(bytes.last() == Some(&NUL),
                "static byte string is not null-terminated: \"{}\"",
                escape_bytestring(bytes));
        unsafe { g_str_from_ptr_internal(bytes.as_ptr(), bytes) }
    }

    pub unsafe fn from_raw<'a, T: ?Sized>(ptr: *const gchar,
                                          life_anchor: &'a T)
                                         -> &'a GStr
    {
        assert!(!ptr.is_null());
        g_str_from_ptr_internal(ptr as *const u8, life_anchor)
    }
}

impl Utf8 {
    #[inline]
    pub fn as_ptr(&self) -> *const gchar {
        self.gstr.as_ptr()
    }

    #[inline]
    pub fn as_g_str(&self) -> &GStr {
        &self.gstr
    }

    pub fn from_static_str(s: &'static str) -> &'static Utf8 {
        assert!(s.ends_with("\0"),
                "static string is not null-terminated: \"{}\"", s);
        unsafe { utf8_from_ptr_internal(s.as_ptr(), s) }
    }

    pub unsafe fn from_raw<'a, T: ?Sized>(ptr: *const gchar,
                                          life_anchor: &'a T)
                                         -> &'a Utf8
    {
        assert!(!ptr.is_null());
        utf8_from_ptr_internal(ptr as *const u8, life_anchor)
    }
}

unsafe fn g_str_from_ptr_internal<'a, T: ?Sized>(ptr: *const u8,
                                                 life_anchor: &'a T)
                                                -> &'a GStr
{
    mem::copy_lifetime(life_anchor, &*(ptr as *const GStr))
}

unsafe fn utf8_from_ptr_internal<'a, T: ?Sized>(ptr: *const u8,
                                                life_anchor: &'a T)
                                               -> &'a Utf8
{
    mem::copy_lifetime(life_anchor, &*(ptr as *const Utf8))
}

fn vec_into_g_str_buf(mut v: Vec<u8>) -> GStrBuf {
    v.push(NUL);
    GStrBuf { data: v }
}

pub struct GStrBuf {
    data: Vec<u8>
}

impl Deref for GStrBuf {

    type Target = GStr;

    fn deref(&self) -> &GStr {
        unsafe { g_str_from_ptr_internal(self.data.as_ptr(), self) }
    }
}

impl GStrBuf {

    #[inline]
    pub fn from_str(s: &str) -> Result<GStrBuf, NulError> {
        GStrBuf::from_bytes(s.as_bytes())
    }

    pub fn from_bytes(bytes: &[u8]) -> Result<GStrBuf, NulError> {
        if let Some(pos) = bytes.position_elem(&NUL) {
            return Err(NulError { position: pos });
        }
        Ok(vec_into_g_str_buf(bytes.to_vec()))
    }

    pub fn from_vec(vec: Vec<u8>) -> Result<GStrBuf, IntoGStrError> {
        if let Some(pos) = vec.position_elem(&NUL) {
            return Err(IntoGStrError {
                cause: NulError { position: pos },
                bytes: vec
            });
        }
        Ok(vec_into_g_str_buf(vec))
    }
}

pub struct Utf8Buf {
    inner: GStrBuf
}

impl Deref for Utf8Buf {

    type Target = Utf8;

    fn deref(&self) -> &Utf8 {
        unsafe { utf8_from_ptr_internal(self.inner.data.as_ptr(), self) }
    }
}

fn utf8_wrap_g_str_result<E>(res: Result<GStrBuf, E>) -> Result<Utf8Buf, E> {
    res.map(|buf| {
        Utf8Buf { inner: buf }
    })
}

impl Utf8Buf {

    pub fn from_str(s: &str) -> Result<Utf8Buf, NulError> {
        utf8_wrap_g_str_result(GStrBuf::from_bytes(s.as_bytes()))
    }

    pub fn from_string(s: String) -> Result<Utf8Buf, IntoGStrError> {
        utf8_wrap_g_str_result(GStrBuf::from_vec(s.into_bytes()))
    }
}
<commit_msg>gstr: Move parse_as_* methods to GStr, parse_as_str method for Utf8<commit_after>// This file is part of Grust, GObject introspection bindings for Rust
//
// Copyright (C) 2013, 2014  Mikhail Zabaluev <mikhail.zabaluev@gmail.com>
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

use ffi;
use types::{gchar,gpointer};
use util::escape_bytestring;

use libc;
use std::error::Error;
use std::fmt;
use std::marker;
use std::mem;
use std::ops::Deref;
use std::slice;
use std::str;

const NUL: u8 = 0;

pub struct OwnedGStr {
    ptr: *const gchar,
}

pub unsafe fn parse_as_bytes<'a, T: ?Sized>(raw: *const gchar,
                                            life_anchor: &'a T)
                                           -> &'a [u8]
{
    assert!(!raw.is_null());
    let r = mem::copy_lifetime(life_anchor, &(raw as *const u8));
    slice::from_raw_buf(r, libc::strlen(raw) as usize)
}

#[inline]
pub unsafe fn parse_as_utf8<'a, T: ?Sized>(raw: *const gchar,
                                           life_anchor: &'a T)
                                          -> Result<&'a str, str::Utf8Error>
{
    str::from_utf8(parse_as_bytes(raw, life_anchor))
}

impl OwnedGStr {

    pub unsafe fn from_raw(ptr: *mut gchar) -> OwnedGStr {
        assert!(!ptr.is_null());
        OwnedGStr { ptr: ptr as *const gchar }
    }
}

impl Deref for OwnedGStr {

    type Target = GStr;

    fn deref(&self) -> &GStr {
        unsafe { g_str_from_ptr_internal(self.ptr as *const u8, self) }
    }
}

impl Drop for OwnedGStr {
    fn drop(&mut self) {
        unsafe { ffi::g_free(self.ptr as gpointer) }
    }
}

impl Clone for OwnedGStr {
    fn clone(&self) -> OwnedGStr {
        unsafe {
            OwnedGStr::from_raw(ffi::g_strdup(self.ptr))
        }
    }
}

impl PartialEq for OwnedGStr {
    fn eq(&self, other: &OwnedGStr) -> bool {
        unsafe { libc::strcmp(self.ptr, other.ptr) == 0 }
    }
}

impl Eq for OwnedGStr { }

#[derive(Copy)]
pub struct NulError {
    pub position: usize
}

impl Error for NulError {

    fn description(&self) -> &str {
        "invalid data for C string: contains a NUL byte"
    }

    fn detail(&self) -> Option<String> {
        Some(format!("NUL at position {}", self.position))
    }
}

impl fmt::Show for NulError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "invalid data for C string: NUL at position {}",
               self.position)
    }
}

pub struct IntoGStrError {
    cause: NulError,
    bytes: Vec<u8>
}

impl IntoGStrError {

    pub fn nul_error(&self) -> &NulError {
        &self.cause
    }

    pub fn into_bytes(self) -> Vec<u8> {
        self.bytes
    }
}

impl Error for IntoGStrError {

    fn description(&self) -> &str {
        self.cause.description()
    }

    fn detail(&self) -> Option<String> {
        self.cause.detail()
    }
}

impl fmt::Show for IntoGStrError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self.cause)
    }
}

#[repr(C)]
pub struct GStr {
    head: gchar,
    marker: marker::NoCopy
}

#[repr(C)]
pub struct Utf8 {
    gstr: GStr
}

impl GStr {
    #[inline]
    pub fn as_ptr(&self) -> *const gchar {
        &self.head as *const gchar
    }

    pub fn parse_as_bytes(&self) -> &[u8] {
        unsafe {
            let r = mem::copy_lifetime(self, &(self.as_ptr() as *const u8));
            slice::from_raw_buf(r, libc::strlen(self.as_ptr()) as usize)
        }
    }

    #[inline]
    pub fn parse_as_utf8(&self) -> Result<&str, str::Utf8Error> {
        str::from_utf8(self.parse_as_bytes())
    }

    #[inline]
    pub unsafe fn parse_as_utf8_unchecked(&self) -> &str {
        str::from_utf8_unchecked(self.parse_as_bytes())
    }

    pub fn from_static_bytes(bytes: &'static [u8]) -> &'static GStr {
        assert!(bytes.last() == Some(&NUL),
                "static byte string is not null-terminated: \"{}\"",
                escape_bytestring(bytes));
        unsafe { g_str_from_ptr_internal(bytes.as_ptr(), bytes) }
    }

    pub unsafe fn from_raw<'a, T: ?Sized>(ptr: *const gchar,
                                          life_anchor: &'a T)
                                         -> &'a GStr
    {
        assert!(!ptr.is_null());
        g_str_from_ptr_internal(ptr as *const u8, life_anchor)
    }
}

impl Utf8 {

    #[inline]
    pub fn as_ptr(&self) -> *const gchar {
        self.gstr.as_ptr()
    }

    #[inline]
    pub fn as_g_str(&self) -> &GStr {
        &self.gstr
    }

    #[inline]
    pub fn parse_as_str(&self) -> &str {
        unsafe { str::from_utf8_unchecked(self.gstr.parse_as_bytes()) }
    }

    pub fn from_static_str(s: &'static str) -> &'static Utf8 {
        assert!(s.ends_with("\0"),
                "static string is not null-terminated: \"{}\"", s);
        unsafe { utf8_from_ptr_internal(s.as_ptr(), s) }
    }

    pub unsafe fn from_raw<'a, T: ?Sized>(ptr: *const gchar,
                                          life_anchor: &'a T)
                                         -> &'a Utf8
    {
        assert!(!ptr.is_null());
        utf8_from_ptr_internal(ptr as *const u8, life_anchor)
    }
}

unsafe fn g_str_from_ptr_internal<'a, T: ?Sized>(ptr: *const u8,
                                                 life_anchor: &'a T)
                                                -> &'a GStr
{
    mem::copy_lifetime(life_anchor, &*(ptr as *const GStr))
}

unsafe fn utf8_from_ptr_internal<'a, T: ?Sized>(ptr: *const u8,
                                                life_anchor: &'a T)
                                               -> &'a Utf8
{
    mem::copy_lifetime(life_anchor, &*(ptr as *const Utf8))
}

fn vec_into_g_str_buf(mut v: Vec<u8>) -> GStrBuf {
    v.push(NUL);
    GStrBuf { data: v }
}

pub struct GStrBuf {
    data: Vec<u8>
}

impl Deref for GStrBuf {

    type Target = GStr;

    fn deref(&self) -> &GStr {
        unsafe { g_str_from_ptr_internal(self.data.as_ptr(), self) }
    }
}

impl GStrBuf {

    #[inline]
    pub fn from_str(s: &str) -> Result<GStrBuf, NulError> {
        GStrBuf::from_bytes(s.as_bytes())
    }

    pub fn from_bytes(bytes: &[u8]) -> Result<GStrBuf, NulError> {
        if let Some(pos) = bytes.position_elem(&NUL) {
            return Err(NulError { position: pos });
        }
        Ok(vec_into_g_str_buf(bytes.to_vec()))
    }

    pub fn from_vec(vec: Vec<u8>) -> Result<GStrBuf, IntoGStrError> {
        if let Some(pos) = vec.position_elem(&NUL) {
            return Err(IntoGStrError {
                cause: NulError { position: pos },
                bytes: vec
            });
        }
        Ok(vec_into_g_str_buf(vec))
    }
}

pub struct Utf8Buf {
    inner: GStrBuf
}

impl Deref for Utf8Buf {

    type Target = Utf8;

    fn deref(&self) -> &Utf8 {
        unsafe { utf8_from_ptr_internal(self.inner.data.as_ptr(), self) }
    }
}

fn utf8_wrap_g_str_result<E>(res: Result<GStrBuf, E>) -> Result<Utf8Buf, E> {
    res.map(|buf| {
        Utf8Buf { inner: buf }
    })
}

impl Utf8Buf {

    pub fn from_str(s: &str) -> Result<Utf8Buf, NulError> {
        utf8_wrap_g_str_result(GStrBuf::from_bytes(s.as_bytes()))
    }

    pub fn from_string(s: String) -> Result<Utf8Buf, IntoGStrError> {
        utf8_wrap_g_str_result(GStrBuf::from_vec(s.into_bytes()))
    }
}
<|endoftext|>


<commit_before>extern crate num_cpus;
extern crate itertools;
extern crate rand;
extern crate crossbeam;

use self::rand::{thread_rng, Rng};
use self::itertools::Itertools;
use self::crossbeam::{scope, ScopedJoinHandle};

use std::ops::Range;
use std::fmt::{Debug, Formatter, Result as FmtResult};
use std::sync::{Mutex, RwLock, MutexGuard};
use std::sync::mpsc::{Sender, Receiver, channel};
use std::thread::spawn;
use std::collections::BTreeSet;

use task::{TaskGenerator, Task};
use candidate::{WorkingCandidate, Candidate};
use context::Context;
use scaling::{ScalingFunction, proportionate};
use result::{Result as AbcResult, Error as AbcError};

/// Manages the parameters of the ABC algorithm.
pub struct HiveBuilder<Ctx: Context> {
    workers: usize,
    observers: usize,
    retries: usize,
    context: Ctx,
    threads: usize,
    scale: Box<ScalingFunction>,
}

impl<Ctx: Context> HiveBuilder<Ctx> {
    /// Creates a new hive.
    ///
    /// * `context` - Factory-like state that can be used while generating solutions.
    /// * `workers` - Number of working solution candidates to maintain at a time.
    pub fn new(context: Ctx, workers: usize) -> HiveBuilder<Ctx> {
        if workers == 0 {
            panic!("HiveBuilder must have at least one worker.");
        }

        HiveBuilder {
            workers: workers,
            observers: workers,
            retries: workers,

            context: context,
            threads: num_cpus::get(),
            scale: proportionate(),
        }
    }

    /// Sets the number of "bees" that will pick a candidate to work on at random.
    ///
    /// This defaults to the number of workers.
    pub fn set_observers(mut self, observers: usize) -> HiveBuilder<Ctx> {
        self.observers = observers;
        self
    }

    /// Sets the number of times a candidate can go unimproved before being reinitialized.
    ///
    /// This defaults to the number of workers.
    pub fn set_retries(mut self, retries: usize) -> HiveBuilder<Ctx> {
        self.retries = retries;
        self
    }

    /// Sets the number of worker threads to use while running.
    pub fn set_threads(mut self, threads: usize) -> HiveBuilder<Ctx> {
        self.threads = threads;
        self
    }

    /// Sets the scaling function for observers to use.
    pub fn set_scaling(mut self, scale: Box<ScalingFunction>) -> HiveBuilder<Ctx> {
        self.scale = scale;
        self
    }

    /// Activates the `HiveBuilder` to create a runnable object.
    pub fn build(self) -> AbcResult<Hive<Ctx>> {
        Hive::new(self)
    }

    fn new_candidate(&self) -> Candidate<Ctx::Solution> {
        let solution = self.context.make();
        let fitness = self.context.evaluate_fitness(&solution);
        Candidate::new(solution, fitness)
    }
}

/// Runs the ABC algorithm, maintaining any necessary state.
pub struct Hive<Ctx: Context> {
    hive: HiveBuilder<Ctx>,

    working: Vec<RwLock<WorkingCandidate<Ctx::Solution>>>,
    best: Mutex<Candidate<Ctx::Solution>>,
    scouting: RwLock<BTreeSet<usize>>,

    tasks: Mutex<Option<TaskGenerator>>,
    sender: Option<Mutex<Sender<Candidate<Ctx::Solution>>>>,
}

impl<Ctx: Context> Hive<Ctx> {
    fn new(hive: HiveBuilder<Ctx>) -> AbcResult<Hive<Ctx>> {
        // Start by populating the field with an initial set of solution candidates.

        // Feed the worker threads a total of N items, each signifying that
        // we need another candidate.
        let tokens: Mutex<Range<usize>> = Mutex::new(0..hive.workers);

        let candidates = Mutex::new(Vec::with_capacity(hive.workers));
        let mut handles = Vec::<ScopedJoinHandle<AbcResult<()>>>::with_capacity(hive.threads);

        try!(crossbeam::scope(|scope| {
            for _ in 0..hive.threads {
                handles.push(scope.spawn(|| {
                    while let Some(_) = {
                        let mut guard = tokens.lock().unwrap();
                        guard.next()
                    } {
                        let candidate = hive.new_candidate();
                        try!(candidates.lock()).push(candidate);
                    }
                    Ok(())
                }));
            }

            // Gather and return `Ok` iff all of the workers finished
            // successfully, otherwise abort the construction.
            handles.drain(..)
                   .fold(Ok(()), |result, handle| result.and(handle.join()))
        }));

        // We don't need the mutex anymore, since we're no longer populating
        // the candidate set from multiple threads.
        let mut candidates = try!(candidates.into_inner());

        // Find the current best candidate, since we want to cache the best
        // at any given moment.
        let best = {
            let best_candidate = candidates.iter()
                                           .fold1(|best, next| {
                                               if next.fitness > best.fitness {
                                                   next
                                               } else {
                                                   best
                                               }
                                           })
                                           .unwrap();
            Mutex::new(best_candidate.clone())
        };

        // Wrap the candidates in a structure that will let the eventual
        // thread swarm work on them.
        let working = candidates.drain(..)
                                .map(|c| RwLock::new(WorkingCandidate::new(c, hive.retries)))
                                .collect::<Vec<RwLock<WorkingCandidate<Ctx::Solution>>>>();

        Ok(Hive {
            hive: hive,
            working: working,
            best: best,
            scouting: RwLock::new(BTreeSet::new()),
            tasks: Mutex::new(None),
            sender: None,
        })
    }

    /// Clone a snapshot of the current set of working candidates.
    ///
    /// The goal of this function is to hold a guard for each solution for as
    /// little time as possible, so we can get out of the way of the other
    /// threads. To this end, we clone the solutions, so that the thread can do
    /// its work on a snapshot.
    fn current_working(&self) -> AbcResult<Vec<Candidate<Ctx::Solution>>> {
        let mut current_working = Vec::with_capacity(self.working.len());
        for candidate_mutex in &self.working {
            let read_guard = try!(candidate_mutex.read());
            current_working.push(read_guard.candidate.clone())
        }
        Ok(current_working)
    }

    /// Returns a guard for the current best solution found by the hive.
    ///
    /// If the hive is running, you should drop the guard returned by this
    /// function as soon as convenient, since the logic of the hive can block
    /// on the availability of the associated mutex. If you plan on performing
    /// expensive computations, you should `drop` the guard as soon as
    /// possible, or acquire and clone it within a small block.
    pub fn get(&self) -> AbcResult<MutexGuard<Candidate<Ctx::Solution>>> {
        self.best.lock().map_err(AbcError::from)
    }

    /// Perform greedy selection between a new candidate and the current best.
    fn consider_improvement(&self, candidate: &Candidate<Ctx::Solution>) -> AbcResult<()> {
        let mut best_guard = try!(self.best.lock());
        if candidate.fitness > best_guard.fitness {
            *best_guard = candidate.clone();
            if let Some(mutex) = self.sender.as_ref() {
                // We're streaming, so we need to post the improved candidate.
                let sender_guard = try!(mutex.lock());
                // If this errors, the receiver was dropped, so we're done.
                if let Err(_) = sender_guard.send(candidate.clone()) {
                    try!(self.stop());
                }
            }
        }
        Ok(())
    }

    fn work_on(&self, current_working: &[Candidate<Ctx::Solution>], n: usize) -> AbcResult<()> {
        let variant_solution = self.hive.context.explore(current_working, n);
        let variant_fitness = self.hive.context.evaluate_fitness(&variant_solution);
        let variant = Candidate::new(variant_solution, variant_fitness);
        let mut write_guard = try!(self.working[n].write());
        if variant.fitness > write_guard.candidate.fitness {
            *write_guard = WorkingCandidate::new(variant, self.hive.retries);
            try!(self.consider_improvement(&write_guard.candidate));
        } else {
            write_guard.deplete();
            // Scouting has been folded into the working process
            if write_guard.expired() {
                let mut scouting_guard = try!(self.scouting.write());
                scouting_guard.insert(n);
                drop(scouting_guard);
                drop(write_guard);

                let candidate = self.hive.new_candidate();
                let mut write_guard = try!(self.working[n].write());
                *write_guard = WorkingCandidate::new(candidate, self.hive.retries);
                try!(self.consider_improvement(&write_guard.candidate));
                drop(write_guard);

                let mut scouting_guard = try!(self.scouting.write());
                scouting_guard.remove(&n);
            }
        }
        Ok(())
    }

    fn choose(&self, current_working: &[Candidate<Ctx::Solution>]) -> AbcResult<usize> {
        let fitnesses = (self.hive.scale)(current_working.iter()
                                                         .map(|candidate| candidate.fitness)
                                                         .collect::<Vec<f64>>());

        // Avoid observing candidates that are being scouted.
        let scouting_guard = try!(self.scouting.read());
        let running_totals = fitnesses.iter()
                                      .enumerate()
                                      .filter(|&(ref i, _)| !scouting_guard.contains(i))
                                      .scan(0f64, |total, (i, fitness)| {
                                          *total += *fitness;
                                          Some((i, *total))
                                      })
                                      .collect::<Vec<(usize, f64)>>();
        drop(scouting_guard);

        // Multiplying the choice point is equivalent to, and more efficient than, normalizing
        // all of the scaled fitnesses and having a choice point in [0,1)
        match running_totals.last() {
            Some(&(_, total_fitness)) => {
                let choice_point = thread_rng().next_f64() * total_fitness;
                for &(i, total) in running_totals.iter() {
                    if total > choice_point {
                        return Ok(i);
                    }
                }
                unreachable!();
            }

            // If we are currently scouting all of the solutions, pick one at random.
            None => Ok(thread_rng().gen_range::<usize>(0, fitnesses.len())),
        }
    }

    fn execute(&self, task: &Task) -> AbcResult<()> {
        let current_working = try!(self.current_working());
        let index = match *task {
            Task::Worker(n) => {
                // If the worker's candidate is in the middle of being replaced, just skip it.
                let scouting_guard = try!(self.scouting.read());
                if scouting_guard.contains(&n) {
                    return Ok(());
                }
                n
            }
            Task::Observer(_) => try!(self.choose(&current_working)),
        };
        self.work_on(&current_working, index)
    }

    fn run(&self, tasks: TaskGenerator) -> AbcResult<()> {
        let mut guard = try!(self.tasks.lock());
        *guard = Some(tasks);
        drop(guard);

        let mut handles: Vec<ScopedJoinHandle<AbcResult<()>>> = Vec::new();

        scope(|scope| {
            for _ in 0..self.hive.threads {
                handles.push(scope.spawn(|| {
                    loop {
                        let mut guard = try!(self.tasks.lock());
                        let task = guard.as_mut().and_then(|gen| gen.next());
                        drop(guard);

                        match task {
                            Some(t) => try!(self.execute(&t)),
                            None => return Ok(()),
                        };
                    }
                }));
            }

            // Returns `Ok(())` only if all threads join cleanly, and the task
            // cycle is successfully cleared away.
            //
            // We avoid `try!` because we want all of the following logic to
            // execute unconditionally.
            handles.drain(..)
                   .fold(Ok(()), |result, handle| result.and(handle.join()))
                   .and(self.tasks
                            .lock()
                            .map(|mut tasks_guard| *tasks_guard = None)
                            .map_err(AbcError::from))
        })
    }

    /// Runs for a fixed number of rounds, then return the best solution found.
    ///
    /// If one of the worker threads panics while working, this will return
    /// `Err(abc::Error)`. Otherwise, it will return `Ok` with a `Candidate`.
    pub fn run_for_rounds(&self, rounds: usize) -> AbcResult<Candidate<Ctx::Solution>> {
        let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers).max_rounds(rounds);
        try!(self.run(tasks));
        self.get().map(|guard| guard.clone())
    }

    /// Run indefinitely.
    ///
    /// If one of the worker threads panics while working, this will return
    /// `Err(abc::Error)`. Otherwise, it will return `Ok(())`.
    pub fn run_forever(&self) -> AbcResult<()> {
        let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers);
        self.run(tasks)
    }

    /// Stops a running hive.
    ///
    /// If a worker thread has panicked, this returns `Err(abc::Error)`.
    pub fn stop(&self) -> AbcResult<()> {
        let mut tasks_guard = try!(self.tasks.lock());
        Ok(tasks_guard.as_mut().map_or((), |t| t.stop()))
    }

    /// Each new best candidate will be sent to `sender`.
    ///
    /// This is kept in a separate function so that the hive can be borrowed
    /// while running.
    pub fn set_sender(&mut self, sender: Sender<Candidate<Ctx::Solution>>) -> AbcResult<()> {
        let best_guard = try!(self.best.lock());
        sender.send(best_guard.clone()).unwrap_or(());
        self.sender = Some(Mutex::new(sender));
        Ok(())
    }

    /// Returns the current round of a running hive.
    ///
    /// If a worker thread has panicked and poisoned the task generator lock,
    /// `get_round` will return `Err(abc::Error)`. If the hive has not been
    /// run, `get_round` will return `Ok(None)`.
    ///
    /// If the hive is running, this will return `Ok(Some(n))`. `n` will start
    /// at 0, and increment each time every task in the round has been claimed
    /// (though not necessarily completed) by a worker thread.
    pub fn get_round(&self) -> AbcResult<Option<usize>> {
        let tasks_guard = try!(self.tasks.lock());
        Ok(tasks_guard.as_ref().map(|tasks| tasks.round))
    }

    /// Get a reference to the hive's context.
    pub fn context(&self) -> &Ctx {
        &self.hive.context
    }
}

impl<Ctx: Context + 'static> Hive<Ctx> {
    /// Runs indefinitely in the background, providing a stream of results.
    ///
    /// This method consumes the hive, which will run until the `HiveBuilder`
    /// object is dropped. It returns an `mpsc::Receiver`, which receives a
    /// `Candidate` each time the hive improves on its best solution.
    pub fn stream(mut self) -> Receiver<Candidate<Ctx::Solution>> {
        let (sender, receiver) = channel();
        spawn(move || {
            self.set_sender(sender).expect("Failed to set background sender");
            let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers);
            self.run(tasks)
        });
        receiver
    }
}

impl<Ctx: Context> Debug for Hive<Ctx>
    where Ctx::Solution: Debug
{
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        for mutex in (&self.working).iter() {
            let working = mutex.read().unwrap();
            try!(write!(f, "..{:?}..\n", working.candidate));
        }
        let best_candidate = self.get().unwrap();
        write!(f, ">>{:?}<<", *best_candidate)
    }
}

impl<Ctx: Context> Drop for Hive<Ctx> {
    fn drop(&mut self) {
        self.stop().unwrap_or(())
    }
}
<commit_msg>Change set_sender to return () (maintaining API)<commit_after>extern crate num_cpus;
extern crate itertools;
extern crate rand;
extern crate crossbeam;

use self::rand::{thread_rng, Rng};
use self::itertools::Itertools;
use self::crossbeam::{scope, ScopedJoinHandle};

use std::ops::Range;
use std::fmt::{Debug, Formatter, Result as FmtResult};
use std::sync::{Mutex, RwLock, MutexGuard};
use std::sync::mpsc::{Sender, Receiver, channel};
use std::thread::spawn;
use std::collections::BTreeSet;

use task::{TaskGenerator, Task};
use candidate::{WorkingCandidate, Candidate};
use context::Context;
use scaling::{ScalingFunction, proportionate};
use result::{Result as AbcResult, Error as AbcError};

/// Manages the parameters of the ABC algorithm.
pub struct HiveBuilder<Ctx: Context> {
    workers: usize,
    observers: usize,
    retries: usize,
    context: Ctx,
    threads: usize,
    scale: Box<ScalingFunction>,
}

impl<Ctx: Context> HiveBuilder<Ctx> {
    /// Creates a new hive.
    ///
    /// * `context` - Factory-like state that can be used while generating solutions.
    /// * `workers` - Number of working solution candidates to maintain at a time.
    pub fn new(context: Ctx, workers: usize) -> HiveBuilder<Ctx> {
        if workers == 0 {
            panic!("HiveBuilder must have at least one worker.");
        }

        HiveBuilder {
            workers: workers,
            observers: workers,
            retries: workers,

            context: context,
            threads: num_cpus::get(),
            scale: proportionate(),
        }
    }

    /// Sets the number of "bees" that will pick a candidate to work on at random.
    ///
    /// This defaults to the number of workers.
    pub fn set_observers(mut self, observers: usize) -> HiveBuilder<Ctx> {
        self.observers = observers;
        self
    }

    /// Sets the number of times a candidate can go unimproved before being reinitialized.
    ///
    /// This defaults to the number of workers.
    pub fn set_retries(mut self, retries: usize) -> HiveBuilder<Ctx> {
        self.retries = retries;
        self
    }

    /// Sets the number of worker threads to use while running.
    pub fn set_threads(mut self, threads: usize) -> HiveBuilder<Ctx> {
        self.threads = threads;
        self
    }

    /// Sets the scaling function for observers to use.
    pub fn set_scaling(mut self, scale: Box<ScalingFunction>) -> HiveBuilder<Ctx> {
        self.scale = scale;
        self
    }

    /// Activates the `HiveBuilder` to create a runnable object.
    pub fn build(self) -> AbcResult<Hive<Ctx>> {
        Hive::new(self)
    }

    fn new_candidate(&self) -> Candidate<Ctx::Solution> {
        let solution = self.context.make();
        let fitness = self.context.evaluate_fitness(&solution);
        Candidate::new(solution, fitness)
    }
}

/// Runs the ABC algorithm, maintaining any necessary state.
pub struct Hive<Ctx: Context> {
    hive: HiveBuilder<Ctx>,

    working: Vec<RwLock<WorkingCandidate<Ctx::Solution>>>,
    best: Mutex<Candidate<Ctx::Solution>>,
    scouting: RwLock<BTreeSet<usize>>,

    tasks: Mutex<Option<TaskGenerator>>,
    sender: Option<Mutex<Sender<Candidate<Ctx::Solution>>>>,
}

impl<Ctx: Context> Hive<Ctx> {
    fn new(hive: HiveBuilder<Ctx>) -> AbcResult<Hive<Ctx>> {
        // Start by populating the field with an initial set of solution candidates.

        // Feed the worker threads a total of N items, each signifying that
        // we need another candidate.
        let tokens: Mutex<Range<usize>> = Mutex::new(0..hive.workers);

        let candidates = Mutex::new(Vec::with_capacity(hive.workers));
        let mut handles = Vec::<ScopedJoinHandle<AbcResult<()>>>::with_capacity(hive.threads);

        try!(crossbeam::scope(|scope| {
            for _ in 0..hive.threads {
                handles.push(scope.spawn(|| {
                    while let Some(_) = {
                        let mut guard = tokens.lock().unwrap();
                        guard.next()
                    } {
                        let candidate = hive.new_candidate();
                        try!(candidates.lock()).push(candidate);
                    }
                    Ok(())
                }));
            }

            // Gather and return `Ok` iff all of the workers finished
            // successfully, otherwise abort the construction.
            handles.drain(..)
                   .fold(Ok(()), |result, handle| result.and(handle.join()))
        }));

        // We don't need the mutex anymore, since we're no longer populating
        // the candidate set from multiple threads.
        let mut candidates = try!(candidates.into_inner());

        // Find the current best candidate, since we want to cache the best
        // at any given moment.
        let best = {
            let best_candidate = candidates.iter()
                                           .fold1(|best, next| {
                                               if next.fitness > best.fitness {
                                                   next
                                               } else {
                                                   best
                                               }
                                           })
                                           .unwrap();
            Mutex::new(best_candidate.clone())
        };

        // Wrap the candidates in a structure that will let the eventual
        // thread swarm work on them.
        let working = candidates.drain(..)
                                .map(|c| RwLock::new(WorkingCandidate::new(c, hive.retries)))
                                .collect::<Vec<RwLock<WorkingCandidate<Ctx::Solution>>>>();

        Ok(Hive {
            hive: hive,
            working: working,
            best: best,
            scouting: RwLock::new(BTreeSet::new()),
            tasks: Mutex::new(None),
            sender: None,
        })
    }

    /// Clone a snapshot of the current set of working candidates.
    ///
    /// The goal of this function is to hold a guard for each solution for as
    /// little time as possible, so we can get out of the way of the other
    /// threads. To this end, we clone the solutions, so that the thread can do
    /// its work on a snapshot.
    fn current_working(&self) -> AbcResult<Vec<Candidate<Ctx::Solution>>> {
        let mut current_working = Vec::with_capacity(self.working.len());
        for candidate_mutex in &self.working {
            let read_guard = try!(candidate_mutex.read());
            current_working.push(read_guard.candidate.clone())
        }
        Ok(current_working)
    }

    /// Returns a guard for the current best solution found by the hive.
    ///
    /// If the hive is running, you should drop the guard returned by this
    /// function as soon as convenient, since the logic of the hive can block
    /// on the availability of the associated mutex. If you plan on performing
    /// expensive computations, you should `drop` the guard as soon as
    /// possible, or acquire and clone it within a small block.
    pub fn get(&self) -> AbcResult<MutexGuard<Candidate<Ctx::Solution>>> {
        self.best.lock().map_err(AbcError::from)
    }

    /// Perform greedy selection between a new candidate and the current best.
    fn consider_improvement(&self, candidate: &Candidate<Ctx::Solution>) -> AbcResult<()> {
        let mut best_guard = try!(self.best.lock());
        if candidate.fitness > best_guard.fitness {
            *best_guard = candidate.clone();
            if let Some(mutex) = self.sender.as_ref() {
                // We're streaming, so we need to post the improved candidate.
                let sender_guard = try!(mutex.lock());
                // If this errors, the receiver was dropped, so we're done.
                if let Err(_) = sender_guard.send(candidate.clone()) {
                    try!(self.stop());
                }
            }
        }
        Ok(())
    }

    fn work_on(&self, current_working: &[Candidate<Ctx::Solution>], n: usize) -> AbcResult<()> {
        let variant_solution = self.hive.context.explore(current_working, n);
        let variant_fitness = self.hive.context.evaluate_fitness(&variant_solution);
        let variant = Candidate::new(variant_solution, variant_fitness);
        let mut write_guard = try!(self.working[n].write());
        if variant.fitness > write_guard.candidate.fitness {
            *write_guard = WorkingCandidate::new(variant, self.hive.retries);
            try!(self.consider_improvement(&write_guard.candidate));
        } else {
            write_guard.deplete();
            // Scouting has been folded into the working process
            if write_guard.expired() {
                let mut scouting_guard = try!(self.scouting.write());
                scouting_guard.insert(n);
                drop(scouting_guard);
                drop(write_guard);

                let candidate = self.hive.new_candidate();
                let mut write_guard = try!(self.working[n].write());
                *write_guard = WorkingCandidate::new(candidate, self.hive.retries);
                try!(self.consider_improvement(&write_guard.candidate));
                drop(write_guard);

                let mut scouting_guard = try!(self.scouting.write());
                scouting_guard.remove(&n);
            }
        }
        Ok(())
    }

    fn choose(&self, current_working: &[Candidate<Ctx::Solution>]) -> AbcResult<usize> {
        let fitnesses = (self.hive.scale)(current_working.iter()
                                                         .map(|candidate| candidate.fitness)
                                                         .collect::<Vec<f64>>());

        // Avoid observing candidates that are being scouted.
        let scouting_guard = try!(self.scouting.read());
        let running_totals = fitnesses.iter()
                                      .enumerate()
                                      .filter(|&(ref i, _)| !scouting_guard.contains(i))
                                      .scan(0f64, |total, (i, fitness)| {
                                          *total += *fitness;
                                          Some((i, *total))
                                      })
                                      .collect::<Vec<(usize, f64)>>();
        drop(scouting_guard);

        // Multiplying the choice point is equivalent to, and more efficient than, normalizing
        // all of the scaled fitnesses and having a choice point in [0,1)
        match running_totals.last() {
            Some(&(_, total_fitness)) => {
                let choice_point = thread_rng().next_f64() * total_fitness;
                for &(i, total) in running_totals.iter() {
                    if total > choice_point {
                        return Ok(i);
                    }
                }
                unreachable!();
            }

            // If we are currently scouting all of the solutions, pick one at random.
            None => Ok(thread_rng().gen_range::<usize>(0, fitnesses.len())),
        }
    }

    fn execute(&self, task: &Task) -> AbcResult<()> {
        let current_working = try!(self.current_working());
        let index = match *task {
            Task::Worker(n) => {
                // If the worker's candidate is in the middle of being replaced, just skip it.
                let scouting_guard = try!(self.scouting.read());
                if scouting_guard.contains(&n) {
                    return Ok(());
                }
                n
            }
            Task::Observer(_) => try!(self.choose(&current_working)),
        };
        self.work_on(&current_working, index)
    }

    fn run(&self, tasks: TaskGenerator) -> AbcResult<()> {
        let mut guard = try!(self.tasks.lock());
        *guard = Some(tasks);
        drop(guard);

        let mut handles: Vec<ScopedJoinHandle<AbcResult<()>>> = Vec::new();

        scope(|scope| {
            for _ in 0..self.hive.threads {
                handles.push(scope.spawn(|| {
                    loop {
                        let mut guard = try!(self.tasks.lock());
                        let task = guard.as_mut().and_then(|gen| gen.next());
                        drop(guard);

                        match task {
                            Some(t) => try!(self.execute(&t)),
                            None => return Ok(()),
                        };
                    }
                }));
            }

            // Returns `Ok(())` only if all threads join cleanly, and the task
            // cycle is successfully cleared away.
            //
            // We avoid `try!` because we want all of the following logic to
            // execute unconditionally.
            handles.drain(..)
                   .fold(Ok(()), |result, handle| result.and(handle.join()))
                   .and(self.tasks
                            .lock()
                            .map(|mut tasks_guard| *tasks_guard = None)
                            .map_err(AbcError::from))
        })
    }

    /// Runs for a fixed number of rounds, then return the best solution found.
    ///
    /// If one of the worker threads panics while working, this will return
    /// `Err(abc::Error)`. Otherwise, it will return `Ok` with a `Candidate`.
    pub fn run_for_rounds(&self, rounds: usize) -> AbcResult<Candidate<Ctx::Solution>> {
        let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers).max_rounds(rounds);
        try!(self.run(tasks));
        self.get().map(|guard| guard.clone())
    }

    /// Run indefinitely.
    ///
    /// If one of the worker threads panics while working, this will return
    /// `Err(abc::Error)`. Otherwise, it will return `Ok(())`.
    pub fn run_forever(&self) -> AbcResult<()> {
        let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers);
        self.run(tasks)
    }

    /// Stops a running hive.
    ///
    /// If a worker thread has panicked, this returns `Err(abc::Error)`.
    pub fn stop(&self) -> AbcResult<()> {
        let mut tasks_guard = try!(self.tasks.lock());
        Ok(tasks_guard.as_mut().map_or((), |t| t.stop()))
    }

    /// Each new best candidate will be sent to `sender`.
    ///
    /// This is kept in a separate function so that the hive can be borrowed
    /// while running.
    pub fn set_sender(&mut self, sender: Sender<Candidate<Ctx::Solution>>) {
        if let Ok(best_guard) = self.best.lock() {
            sender.send(best_guard.clone()).unwrap_or(());
        }
        self.sender = Some(Mutex::new(sender));
    }

    /// Returns the current round of a running hive.
    ///
    /// If a worker thread has panicked and poisoned the task generator lock,
    /// `get_round` will return `Err(abc::Error)`. If the hive has not been
    /// run, `get_round` will return `Ok(None)`.
    ///
    /// If the hive is running, this will return `Ok(Some(n))`. `n` will start
    /// at 0, and increment each time every task in the round has been claimed
    /// (though not necessarily completed) by a worker thread.
    pub fn get_round(&self) -> AbcResult<Option<usize>> {
        let tasks_guard = try!(self.tasks.lock());
        Ok(tasks_guard.as_ref().map(|tasks| tasks.round))
    }

    /// Get a reference to the hive's context.
    pub fn context(&self) -> &Ctx {
        &self.hive.context
    }
}

impl<Ctx: Context + 'static> Hive<Ctx> {
    /// Runs indefinitely in the background, providing a stream of results.
    ///
    /// This method consumes the hive, which will run until the `HiveBuilder`
    /// object is dropped. It returns an `mpsc::Receiver`, which receives a
    /// `Candidate` each time the hive improves on its best solution.
    pub fn stream(mut self) -> Receiver<Candidate<Ctx::Solution>> {
        let (sender, receiver) = channel();
        spawn(move || {
            self.set_sender(sender);
            let tasks = TaskGenerator::new(self.hive.workers, self.hive.observers);
            self.run(tasks)
        });
        receiver
    }
}

impl<Ctx: Context> Debug for Hive<Ctx>
    where Ctx::Solution: Debug
{
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        for mutex in (&self.working).iter() {
            let working = mutex.read().unwrap();
            try!(write!(f, "..{:?}..\n", working.candidate));
        }
        let best_candidate = self.get().unwrap();
        write!(f, ">>{:?}<<", *best_candidate)
    }
}

impl<Ctx: Context> Drop for Hive<Ctx> {
    fn drop(&mut self) {
        self.stop().unwrap_or(())
    }
}
<|endoftext|>