<commit_before>//! Functions for computing [integral images](https://en.wikipedia.org/wiki/Summed_area_table)
//! and running sums of rows and columns.

extern crate image;

use image::{
    Luma,
    GenericImage,
    ImageBuffer
};

use definitions::{
    VecBuffer
};

/// Computes the integral image of an 8bpp grayscale image.
/// I is the integral image of an image F if I(x, y) is the
/// sum of F(x', y') for x' <= x, y' <= y. i.e. each pixel
/// in the integral image contains the sum of the pixel intensities
/// of all input pixels that are above it and to its left.
/// The integral image has the helpful property that it lets us
/// compute the sum of pixel intensities from any rectangular region
/// in the input image in constant time.
/// Specifically, given a rectangle in F with clockwise corners
/// A, B, C, D, with A at the upper left, the total pixel intensity
/// of this rectangle is I(C) - I(B) - I(D) + I(A).
// TODO: Support more formats.
// TODO: This is extremely slow. Fix that!
pub fn integral_image<I>(image: &I) -> VecBuffer<Luma<u32>>
    where I: GenericImage<Pixel=Luma<u8>> {
    padded_integral_image(image, 0, 0)
}

/// Computes the integral image of the result of padding image with
/// its boundary pixels for x_padding columns on either side and y_padding
/// rows at its top and bottom. Returned image has width image.width() + 2 * x_padding
/// and height image.height() + 2 * y_padding.
pub fn padded_integral_image<I>(image: &I, x_padding: u32, y_padding: u32) -> VecBuffer<Luma<u32>>
    where I: GenericImage<Pixel=Luma<u8>>
{
    let (in_width, in_height) = image.dimensions();
    let out_width = in_width + 2 * x_padding;
    let out_height = in_height + 2 * y_padding;

    let mut out: ImageBuffer<Luma<u32>, Vec<u32>> = ImageBuffer::new(out_width, out_height);

    if in_width == 0 || in_height == 0 {
        return out;
    }

    let mut top_sum = 0u32;
    for x in 0..out_width {
        let x_in = adjust_for_padding(x, x_padding, in_width);
        unsafe {
            top_sum += image.unsafe_get_pixel(x_in, 0)[0] as u32;
            out.unsafe_put_pixel(x, 0, Luma([top_sum]));
        }
    }

    for y in 1..out_height {
        let y_in = adjust_for_padding(y, y_padding, in_height);
        let mut sum = 0;

        for x in 0..out_width {
            let x_in = adjust_for_padding(x, x_padding, in_width);

            unsafe {
                sum += image.unsafe_get_pixel(x_in, y_in)[0] as u32;
                let above = out.unsafe_get_pixel(x, y - 1)[0];
                out.unsafe_put_pixel(x, y, Luma([above + sum]))
            }
        }
    }

    out
}

#[inline(always)]
fn adjust_for_padding(out_index: u32, padding: u32, upper: u32) -> u32 {
    if out_index < padding {
        return 0;
    } else if out_index >= upper + padding {
        return upper - 1;
    }
    return out_index - padding;
}

/// Computes the running sum of one row of image, padded
/// at the beginning and end. The padding is by continuity.
/// Takes a reference to buffer so that this can be reused
/// for all rows in an image.
// TODO: faster, more formats
pub fn row_running_sum<I>(image: &I, row: u32, buffer: &mut [u32], padding: u32)
    where I: GenericImage<Pixel=Luma<u8>> {

    let (width, height) = image.dimensions();
    assert!(buffer.len() >= (width + 2 * padding) as usize,
        format!("Buffer length {} is less than {} + 2 * {}", buffer.len(), width, padding));
    assert!(row < height, format!("row out of bound: {} >= {}", row, height));

    unsafe {
        let mut sum = 0;
        for x in 0..padding {
            sum += image.unsafe_get_pixel(0, row)[0] as u32;
            *buffer.get_unchecked_mut(x as usize) = sum;
        }

        for x in 0..width {
            sum += image.unsafe_get_pixel(x, row)[0] as u32;
            *buffer.get_unchecked_mut((x + padding) as usize) = sum;
        }

        for x in 0..padding {
            sum += image.unsafe_get_pixel(width - 1, row)[0] as u32;
            *buffer.get_unchecked_mut((x + width + padding) as usize) = sum;
        }
    }
}

/// Computes the running sum of one column of image, padded
/// at the top and bottom. The padding is by continuity.
/// Takes a reference to buffer so that this can be reused
/// for all columns in an image.
// TODO: faster, more formats
pub fn column_running_sum<I>(image: &I, column: u32, buffer: &mut [u32], padding: u32)
    where I: GenericImage<Pixel=Luma<u8>> {

    let (width, height) = image.dimensions();
    assert!(buffer.len() >= (height + 2 * padding) as usize,
        format!("Buffer length {} is less than {} + 2 * {}", buffer.len(), height, padding));
    assert!(column < width, format!("column out of bound: {} >= {}", column, width));

    unsafe {
        let mut sum = 0;
        for y in 0..padding {
            sum += image.unsafe_get_pixel(column, 0)[0] as u32;
            *buffer.get_unchecked_mut(y as usize) = sum;
        }

        for y in 0..height {
            sum += image.unsafe_get_pixel(column, y)[0] as u32;
            *buffer.get_unchecked_mut((y + padding) as usize) = sum;
        }

        for y in 0..padding {
            sum += image.unsafe_get_pixel(column, height - 1)[0] as u32;
            *buffer.get_unchecked_mut((y + height + padding) as usize) = sum;
        }
    }
}

#[cfg(test)]
mod test {

    use super::{
        column_running_sum,
        integral_image,
        padded_integral_image,
        row_running_sum
    };
    use utils::{
        gray_bench_image,
        GrayTestImage
    };
    use image::{
        GenericImage,
        GrayImage,
        ImageBuffer,
        Luma
    };
    use quickcheck::{
        quickcheck
    };
    use definitions::{
        VecBuffer
    };
    use test;

    #[test]
    fn test_integral_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 2, vec![
            1, 2, 3,
            4, 5, 6]).unwrap();

        let expected: ImageBuffer<Luma<u32>, Vec<u32>>
            = ImageBuffer::from_raw(3, 2, vec![
            1,  3,  6,
            5, 12, 21]).unwrap();

        assert_pixels_eq!(integral_image(&image), expected);
    }

    #[bench]
    fn bench_integral_image(b: &mut test::Bencher) {
        let image = gray_bench_image(500, 500);
        b.iter(|| {
            let integral = integral_image(&image);
            test::black_box(integral);
            });
    }

    #[test]
    fn test_padded_integral_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 2, vec![
            1, 2, 3,
            4, 5, 6]).unwrap();

        let expected: ImageBuffer<Luma<u32>, Vec<u32>>
            = ImageBuffer::from_raw(5, 4, vec![
              1,  2,   4,   7,  10,
              2,  4,   8,  14,  20,
              6, 12,  21,  33,  45,
             10, 20,  34,  52,  70]).unwrap();

        assert_pixels_eq!(padded_integral_image(&image, 1, 1), expected);
    }

    #[test]
    fn test_row_running_sum() {
        let image: GrayImage = ImageBuffer::from_raw(3, 2, vec![
            1, 2, 3,
            4, 5, 6]).unwrap();

        let expected = [1, 2, 4, 7, 10];

        let mut buffer = [0; 5];
        row_running_sum(&image, 0, &mut buffer, 1);

        assert_eq!(buffer, expected);
    }

    #[bench]
    fn bench_row_running_sum(b: &mut test::Bencher) {
        let image = gray_bench_image(1000, 1);
        let mut buffer = [0; 1010];
        b.iter(|| {
            row_running_sum(&image, 0, &mut buffer, 5);
            });
    }

    #[test]
    fn test_column_running_sum() {
        let image: GrayImage = ImageBuffer::from_raw(2, 3, vec![
            1, 4,
            2, 5,
            3, 6]).unwrap();

        let expected = [1, 2, 4, 7, 10];

        let mut buffer = [0; 5];
        column_running_sum(&image, 0, &mut buffer, 1);

        assert_eq!(buffer, expected);
    }

    #[bench]
    fn bench_column_running_sum(b: &mut test::Bencher) {
        let image = gray_bench_image(100, 1000);
        let mut buffer = [0; 1010];
        b.iter(|| {
            column_running_sum(&image, 0, &mut buffer, 5);
            });
    }

    /// Simple implementation of integral_image to validate faster versions against.
    fn integral_image_ref<I>(image: &I) -> VecBuffer<Luma<u32>>
        where I: GenericImage<Pixel=Luma<u8>>
    {
        let (width, height) = image.dimensions();
        let mut out = ImageBuffer::new(width, height);

        for y in 0..height {
            for x in 0..width {
                let mut sum = 0u32;

                for iy in 0..(y + 1) {
                    for ix in 0..(x + 1) {
                        sum += image.get_pixel(ix, iy)[0] as u32;
                    }
                }

                out.put_pixel(x, y, Luma([sum]));
            }
        }

        out
    }

    #[test]
    fn test_integral_image_matches_reference_implementation() {
        fn prop(image: GrayTestImage) -> bool {
            let expected = integral_image_ref(&image.0);
            let actual = integral_image(&image.0);
            actual.pixels().eq(expected.pixels())
        }
        quickcheck(prop as fn(GrayTestImage) -> bool);
    }
}
<commit_msg>Fix error message<commit_after>//! Functions for computing [integral images](https://en.wikipedia.org/wiki/Summed_area_table)
//! and running sums of rows and columns.

extern crate image;

use image::{
    Luma,
    GenericImage,
    ImageBuffer
};

use definitions::{
    VecBuffer
};

/// Computes the integral image of an 8bpp grayscale image.
/// I is the integral image of an image F if I(x, y) is the
/// sum of F(x', y') for x' <= x, y' <= y. i.e. each pixel
/// in the integral image contains the sum of the pixel intensities
/// of all input pixels that are above it and to its left.
/// The integral image has the helpful property that it lets us
/// compute the sum of pixel intensities from any rectangular region
/// in the input image in constant time.
/// Specifically, given a rectangle in F with clockwise corners
/// A, B, C, D, with A at the upper left, the total pixel intensity
/// of this rectangle is I(C) - I(B) - I(D) + I(A).
// TODO: Support more formats.
// TODO: This is extremely slow. Fix that!
pub fn integral_image<I>(image: &I) -> VecBuffer<Luma<u32>>
    where I: GenericImage<Pixel=Luma<u8>> {
    padded_integral_image(image, 0, 0)
}

/// Computes the integral image of the result of padding image with
/// its boundary pixels for x_padding columns on either side and y_padding
/// rows at its top and bottom. Returned image has width image.width() + 2 * x_padding
/// and height image.height() + 2 * y_padding.
pub fn padded_integral_image<I>(image: &I, x_padding: u32, y_padding: u32) -> VecBuffer<Luma<u32>>
    where I: GenericImage<Pixel=Luma<u8>>
{
    let (in_width, in_height) = image.dimensions();
    let out_width = in_width + 2 * x_padding;
    let out_height = in_height + 2 * y_padding;

    let mut out: ImageBuffer<Luma<u32>, Vec<u32>> = ImageBuffer::new(out_width, out_height);

    if in_width == 0 || in_height == 0 {
        return out;
    }

    let mut top_sum = 0u32;
    for x in 0..out_width {
        let x_in = adjust_for_padding(x, x_padding, in_width);
        unsafe {
            top_sum += image.unsafe_get_pixel(x_in, 0)[0] as u32;
            out.unsafe_put_pixel(x, 0, Luma([top_sum]));
        }
    }

    for y in 1..out_height {
        let y_in = adjust_for_padding(y, y_padding, in_height);
        let mut sum = 0;

        for x in 0..out_width {
            let x_in = adjust_for_padding(x, x_padding, in_width);

            unsafe {
                sum += image.unsafe_get_pixel(x_in, y_in)[0] as u32;
                let above = out.unsafe_get_pixel(x, y - 1)[0];
                out.unsafe_put_pixel(x, y, Luma([above + sum]))
            }
        }
    }

    out
}

#[inline(always)]
fn adjust_for_padding(out_index: u32, padding: u32, upper: u32) -> u32 {
    if out_index < padding {
        return 0;
    } else if out_index >= upper + padding {
        return upper - 1;
    }
    return out_index - padding;
}

/// Computes the running sum of one row of image, padded
/// at the beginning and end. The padding is by continuity.
/// Takes a reference to buffer so that this can be reused
/// for all rows in an image.
// TODO: faster, more formats
pub fn row_running_sum<I>(image: &I, row: u32, buffer: &mut [u32], padding: u32)
    where I: GenericImage<Pixel=Luma<u8>> {

    let (width, height) = image.dimensions();
    assert!(buffer.len() >= (width + 2 * padding) as usize,
        format!("Buffer length {} is less than {} + 2 * {}", buffer.len(), width, padding));
    assert!(row < height, format!("row out of bounds: {} >= {}", row, height));

    unsafe {
        let mut sum = 0;
        for x in 0..padding {
            sum += image.unsafe_get_pixel(0, row)[0] as u32;
            *buffer.get_unchecked_mut(x as usize) = sum;
        }

        for x in 0..width {
            sum += image.unsafe_get_pixel(x, row)[0] as u32;
            *buffer.get_unchecked_mut((x + padding) as usize) = sum;
        }

        for x in 0..padding {
            sum += image.unsafe_get_pixel(width - 1, row)[0] as u32;
            *buffer.get_unchecked_mut((x + width + padding) as usize) = sum;
        }
    }
}

/// Computes the running sum of one column of image, padded
/// at the top and bottom. The padding is by continuity.
/// Takes a reference to buffer so that this can be reused
/// for all columns in an image.
// TODO: faster, more formats
pub fn column_running_sum<I>(image: &I, column: u32, buffer: &mut [u32], padding: u32)
    where I: GenericImage<Pixel=Luma<u8>> {

    let (width, height) = image.dimensions();
    assert!(buffer.len() >= (height + 2 * padding) as usize,
        format!("Buffer length {} is less than {} + 2 * {}", buffer.len(), height, padding));
    assert!(column < width, format!("column out of bounds: {} >= {}", column, width));

    unsafe {
        let mut sum = 0;
        for y in 0..padding {
            sum += image.unsafe_get_pixel(column, 0)[0] as u32;
            *buffer.get_unchecked_mut(y as usize) = sum;
        }

        for y in 0..height {
            sum += image.unsafe_get_pixel(column, y)[0] as u32;
            *buffer.get_unchecked_mut((y + padding) as usize) = sum;
        }

        for y in 0..padding {
            sum += image.unsafe_get_pixel(column, height - 1)[0] as u32;
            *buffer.get_unchecked_mut((y + height + padding) as usize) = sum;
        }
    }
}

#[cfg(test)]
mod test {

    use super::{
        column_running_sum,
        integral_image,
        padded_integral_image,
        row_running_sum
    };
    use utils::{
        gray_bench_image,
        GrayTestImage
    };
    use image::{
        GenericImage,
        GrayImage,
        ImageBuffer,
        Luma
    };
    use quickcheck::{
        quickcheck
    };
    use definitions::{
        VecBuffer
    };
    use test;

    #[test]
    fn test_integral_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 2, vec![
            1, 2, 3,
            4, 5, 6]).unwrap();

        let expected: ImageBuffer<Luma<u32>, Vec<u32>>
            = ImageBuffer::from_raw(3, 2, vec![
            1,  3,  6,
            5, 12, 21]).unwrap();

        assert_pixels_eq!(integral_image(&image), expected);
    }

    #[bench]
    fn bench_integral_image(b: &mut test::Bencher) {
        let image = gray_bench_image(500, 500);
        b.iter(|| {
            let integral = integral_image(&image);
            test::black_box(integral);
            });
    }

    #[test]
    fn test_padded_integral_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 2, vec![
            1, 2, 3,
            4, 5, 6]).unwrap();

        let expected: ImageBuffer<Luma<u32>, Vec<u32>>
            = ImageBuffer::from_raw(5, 4, vec![
              1,  2,   4,   7,  10,
              2,  4,   8,  14,  20,
              6, 12,  21,  33,  45,
             10, 20,  34,  52,  70]).unwrap();

        assert_pixels_eq!(padded_integral_image(&image, 1, 1), expected);
    }

    #[test]
    fn test_row_running_sum() {
        let image: GrayImage = ImageBuffer::from_raw(3, 2, vec![
            1, 2, 3,
            4, 5, 6]).unwrap();

        let expected = [1, 2, 4, 7, 10];

        let mut buffer = [0; 5];
        row_running_sum(&image, 0, &mut buffer, 1);

        assert_eq!(buffer, expected);
    }

    #[bench]
    fn bench_row_running_sum(b: &mut test::Bencher) {
        let image = gray_bench_image(1000, 1);
        let mut buffer = [0; 1010];
        b.iter(|| {
            row_running_sum(&image, 0, &mut buffer, 5);
            });
    }

    #[test]
    fn test_column_running_sum() {
        let image: GrayImage = ImageBuffer::from_raw(2, 3, vec![
            1, 4,
            2, 5,
            3, 6]).unwrap();

        let expected = [1, 2, 4, 7, 10];

        let mut buffer = [0; 5];
        column_running_sum(&image, 0, &mut buffer, 1);

        assert_eq!(buffer, expected);
    }

    #[bench]
    fn bench_column_running_sum(b: &mut test::Bencher) {
        let image = gray_bench_image(100, 1000);
        let mut buffer = [0; 1010];
        b.iter(|| {
            column_running_sum(&image, 0, &mut buffer, 5);
            });
    }

    /// Simple implementation of integral_image to validate faster versions against.
    fn integral_image_ref<I>(image: &I) -> VecBuffer<Luma<u32>>
        where I: GenericImage<Pixel=Luma<u8>>
    {
        let (width, height) = image.dimensions();
        let mut out = ImageBuffer::new(width, height);

        for y in 0..height {
            for x in 0..width {
                let mut sum = 0u32;

                for iy in 0..(y + 1) {
                    for ix in 0..(x + 1) {
                        sum += image.get_pixel(ix, iy)[0] as u32;
                    }
                }

                out.put_pixel(x, y, Luma([sum]));
            }
        }

        out
    }

    #[test]
    fn test_integral_image_matches_reference_implementation() {
        fn prop(image: GrayTestImage) -> bool {
            let expected = integral_image_ref(&image.0);
            let actual = integral_image(&image.0);
            actual.pixels().eq(expected.pixels())
        }
        quickcheck(prop as fn(GrayTestImage) -> bool);
    }
}
<|endoftext|>


<commit_before>use regex::Regex;

use config::JiraConfig;
use github::{Commit, CommitLike, PullRequest, PushCommit};
use jira;
use jira::Transition;

fn get_jira_keys(strings: Vec<String>) -> Vec<String> {
    let re = Regex::new(r"\b([A-Z]+-[0-9]+)\b").unwrap();

    let mut all_keys = vec![];
    for s in strings {
        all_keys.extend(re.captures_iter(&s).map(|c| c[1].to_string()));
    }

    all_keys.sort();
    all_keys.dedup();

    all_keys
}

fn get_fixed_jira_keys<T: CommitLike>(commits: &Vec<T>) -> Vec<String> {
    // Fix [ABC-123][OTHER-567], [YEAH-999]
    let re = Regex::new(r"(?i)(?:Fix):?\s*(?-i)((\[?([A-Z]+-[0-9]+)(?:\]|\b)[\s,]*)+)").unwrap();

    // first extract jiras with fix markers
    let mut all_refs = vec![];
    for c in commits {
        all_refs.extend(re.captures_iter(c.message()).map(|c| c[1].to_string()));
    }

    get_jira_keys(all_refs)
}

fn get_referenced_jira_keys<T: CommitLike>(commits: &Vec<T>) -> Vec<String> {
    let fixed = get_fixed_jira_keys(commits);

    let mut refd = get_jira_keys(commits.iter().map(|c| c.message().to_string()).collect());

    // only return ones not marked as fixed
    refd.retain(|s| fixed.iter().position(|s2| s == s2).is_none());
    refd
}

pub fn submit_for_review(pr: &PullRequest, commits: &Vec<Commit>, jira: &jira::api::Session, config: &JiraConfig) {
    for key in get_fixed_jira_keys(commits) {
        // add comment
        if let Err(e) = jira.comment_issue(&key, &format!("Review submitted for branch {}: {}", pr.base.ref_name, pr.html_url)) {
            error!("Error commenting on key [{}]: {}", key, e);
            continue; // give up on transitioning if we can't comment.
        }

        // try to transition to in-progress
        try_transition(&key, &config.progress_states(), jira);
        // try transition to pending-review
        try_transition(&key, &config.review_states(), jira);
    }

    for key in get_referenced_jira_keys(commits) {
        // add comment
        if let Err(e) = jira.comment_issue(&key, &format!("Referenced by review submitted for branch {}: {}", pr.base.ref_name, pr.html_url)) {
            error!("Error commenting on key [{}]: {}", key, e);
            continue; // give up on transitioning if we can't comment.
        }

        // try to transition to in-progress
        try_transition(&key, &config.progress_states(), jira);
    }
}

pub fn resolve_issue(branch: &str, version: Option<&str>, commits: &Vec<PushCommit>, jira: &jira::api::Session, config: &JiraConfig) {
    for commit in commits {
        let desc = format!("[{}|{}]\n{{quote}}{}{{quote}}", Commit::short_hash(&commit), commit.html_url(), Commit::title(&commit));

        let version_desc = match version {
            None => String::new(),
            Some(v) => format!("\nIncluded in version {}", v),
        };

        let fix_msg = format!("Merged into branch {}: {}{}", branch, desc, version_desc);
        let ref_msg = format!("Referenced by commit merged into branch {}: {}{}", branch, desc, version_desc);

        for key in get_fixed_jira_keys(&vec![commit]) {
            if let Err(e) = jira.comment_issue(&key, &fix_msg) {
                error!("Error commenting on key [{}]: {}", key, e);
            }

            let to = config.resolved_states();
            match find_transition(&key, &to, jira) {
                Ok(Some(transition)) => {
                    let mut req = transition.new_request();

                    if let Some(ref fields) = transition.fields {
                        if let Some(ref resolution) = fields.resolution {
                            for res in &resolution.allowed_values {
                                for resolution in config.fixed_resolutions() {
                                    if res.name == resolution {
                                        req.set_resolution(&res);
                                        break;
                                    }
                                }
                                if req.fields.is_some() {
                                    break;
                                }
                            }
                            if req.fields.is_none() {
                                error!("Could not find fixed resolution in allowed values: [{:?}]!", resolution.allowed_values);
                            }
                        }
                    }

                    if let Err(e) = jira.transition_issue(&key, &req) {
                        error!("Error transitioning JIRA issue [{}] to one of [{:?}]: {}", key, to, e);
                    } else {
                        info!("Transitioned [{}] to one of [{:?}]", key, to);
                    }
                },
                Ok(None) => info!("JIRA [{}] cannot be transitioned to  any of [{:?}]", key, to),
                Err(e) => error!("{}", e),
            };
        }

        // add comment only to referenced jiras
        for key in get_referenced_jira_keys(&vec![commit]) {
            if let Err(e) = jira.comment_issue(&key, &ref_msg) {
                error!("Error commenting on key [{}]: {}", key, e);
            }
        }
    }
}

fn try_transition(key: &str, to: &Vec<String>, jira: &jira::api::Session) {
    match find_transition(&key, to, jira) {
        Ok(Some(transition)) => {
            let req = transition.new_request();
            if let Err(e) = jira.transition_issue(&key, &req) {
                error!("Error transitioning JIRA issue [{}] to one of [{:?}]: {}", key, to, e);
            } else {
                info!("Transitioned [{}] to one of [{:?}]", key, to);
            }
        },
        Ok(None) => info!("JIRA [{}] cannot be transitioned to any of [{:?}]", key, to),
        Err(e) => error!("{}", e),
    };
}

fn find_transition(key: &str, to: &Vec<String>, jira: &jira::api::Session) -> Result<Option<Transition>, String> {
    let transitions = match jira.get_transitions(&key) {
        Ok(t) => t,
        Err(e) => return Err(format!("Error looking up JIRA transitions for key [{}]: {}", key, e)),
    };

    Ok(pick_transition(to, &transitions))
}

fn pick_transition(to: &Vec<String>, choices: &Vec<Transition>) -> Option<Transition> {
    for t in choices {
        for name in to {
            if &t.name == name || &t.to.name == name {
                return Some(t.clone())
            }
        }
    }

    None
}

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

    #[test]
    pub fn test_get_jira_keys() {
        let mut commit = Commit::new();
        assert_eq!(Vec::<String>::new(), get_fixed_jira_keys(&vec![commit.clone()]));
        assert_eq!(Vec::<String>::new(), get_referenced_jira_keys(&vec![commit.clone()]));

        commit.commit.message = "Fix [KEY-1][KEY-2] Some thing that also fixed [KEY-3], [KEY-4]".into();
        assert_eq!(vec!["KEY-1", "KEY-2", "KEY-3", "KEY-4"], get_fixed_jira_keys(&vec![commit.clone()]));
        assert_eq!(Vec::<String>::new(), get_referenced_jira_keys(&vec![commit.clone()]));

        commit.commit.message += "\n\nOh, I forgot it also fixes [KEY-4], and also mentions [KEY-4], [KEY-5] but not [lowercase-99]";
        assert_eq!(vec!["KEY-1", "KEY-2", "KEY-3", "KEY-4"], get_fixed_jira_keys(&vec![commit.clone()]));
        assert_eq!(vec!["KEY-5"], get_referenced_jira_keys(&vec![commit.clone()]));
    }

    #[test]
    pub fn test_get_jira_keys_alt_format() {
        let mut commit = Commit::new();
        commit.commit.message = "KEY-1, KEY-2:Some thing that also fixed\n\nAlso [KEY-3], OTHER-5".into();
        assert_eq!(Vec::<String>::new(), get_fixed_jira_keys(&vec![commit.clone()]));
        assert_eq!(vec!["KEY-1", "KEY-2", "KEY-3", "OTHER-5"], get_referenced_jira_keys(&vec![commit.clone()]));
    }

    #[test]
    pub fn test_pick_transition() {
        let t1 = Transition {
            id: "1".into(),
            name: "t1".into(),
            to: TransitionTo {
                id: "10".into(),
                name: "inside-t1".into(),
            },
            fields: None,
        };
        let t2 = Transition {
            id: "2".into(),
            name: "t2".into(),
            to: TransitionTo {
                id: "20".into(),
                name: "inside-t2".into(),
            },
            fields: None,
        };
        assert_eq!(Some(t1.clone()), pick_transition(&vec!["t1".into()], &vec![t1.clone(), t2.clone()]));
        assert_eq!(Some(t1.clone()), pick_transition(&vec!["inside-t1".into(), "t2".into()], &vec![t1.clone(), t2.clone()]));
        assert_eq!(Some(t2.clone()), pick_transition(&vec!["inside-t2".into()], &vec![t1.clone(), t2.clone()]));
        assert_eq!(None, pick_transition(&vec!["something-else".into()], &vec![t1.clone(), t2.clone()]));
    }

}
<commit_msg>Whoops. fix tests (#51)<commit_after>use regex::Regex;

use config::JiraConfig;
use github::{Commit, CommitLike, PullRequest, PushCommit};
use jira;
use jira::Transition;

fn get_jira_keys(strings: Vec<String>) -> Vec<String> {
    let re = Regex::new(r"\b([A-Z]+-[0-9]+)\b").unwrap();

    let mut all_keys = vec![];
    for s in strings {
        all_keys.extend(re.captures_iter(&s).map(|c| c[1].to_string()));
    }

    all_keys.sort();
    all_keys.dedup();

    all_keys
}

fn get_fixed_jira_keys<T: CommitLike>(commits: &Vec<T>) -> Vec<String> {
    // Fix [ABC-123][OTHER-567], [YEAH-999]
    let re = Regex::new(r"(?i)(?:Fix):?\s*(?-i)((\[?([A-Z]+-[0-9]+)(?:\]|\b)[\s,]*)+)").unwrap();

    // first extract jiras with fix markers
    let mut all_refs = vec![];
    for c in commits {
        all_refs.extend(re.captures_iter(c.message()).map(|c| c[1].to_string()));
    }

    get_jira_keys(all_refs)
}

fn get_referenced_jira_keys<T: CommitLike>(commits: &Vec<T>) -> Vec<String> {
    let fixed = get_fixed_jira_keys(commits);

    let mut refd = get_jira_keys(commits.iter().map(|c| c.message().to_string()).collect());

    // only return ones not marked as fixed
    refd.retain(|s| fixed.iter().position(|s2| s == s2).is_none());
    refd
}

pub fn submit_for_review(pr: &PullRequest, commits: &Vec<Commit>, jira: &jira::api::Session, config: &JiraConfig) {
    for key in get_fixed_jira_keys(commits) {
        // add comment
        if let Err(e) = jira.comment_issue(&key, &format!("Review submitted for branch {}: {}", pr.base.ref_name, pr.html_url)) {
            error!("Error commenting on key [{}]: {}", key, e);
            continue; // give up on transitioning if we can't comment.
        }

        // try to transition to in-progress
        try_transition(&key, &config.progress_states(), jira);
        // try transition to pending-review
        try_transition(&key, &config.review_states(), jira);
    }

    for key in get_referenced_jira_keys(commits) {
        // add comment
        if let Err(e) = jira.comment_issue(&key, &format!("Referenced by review submitted for branch {}: {}", pr.base.ref_name, pr.html_url)) {
            error!("Error commenting on key [{}]: {}", key, e);
            continue; // give up on transitioning if we can't comment.
        }

        // try to transition to in-progress
        try_transition(&key, &config.progress_states(), jira);
    }
}

pub fn resolve_issue(branch: &str, version: Option<&str>, commits: &Vec<PushCommit>, jira: &jira::api::Session, config: &JiraConfig) {
    for commit in commits {
        let desc = format!("[{}|{}]\n{{quote}}{}{{quote}}", Commit::short_hash(&commit), commit.html_url(), Commit::title(&commit));

        let version_desc = match version {
            None => String::new(),
            Some(v) => format!("\nIncluded in version {}", v),
        };

        let fix_msg = format!("Merged into branch {}: {}{}", branch, desc, version_desc);
        let ref_msg = format!("Referenced by commit merged into branch {}: {}{}", branch, desc, version_desc);

        for key in get_fixed_jira_keys(&vec![commit]) {
            if let Err(e) = jira.comment_issue(&key, &fix_msg) {
                error!("Error commenting on key [{}]: {}", key, e);
            }

            let to = config.resolved_states();
            match find_transition(&key, &to, jira) {
                Ok(Some(transition)) => {
                    let mut req = transition.new_request();

                    if let Some(ref fields) = transition.fields {
                        if let Some(ref resolution) = fields.resolution {
                            for res in &resolution.allowed_values {
                                for resolution in config.fixed_resolutions() {
                                    if res.name == resolution {
                                        req.set_resolution(&res);
                                        break;
                                    }
                                }
                                if req.fields.is_some() {
                                    break;
                                }
                            }
                            if req.fields.is_none() {
                                error!("Could not find fixed resolution in allowed values: [{:?}]!", resolution.allowed_values);
                            }
                        }
                    }

                    if let Err(e) = jira.transition_issue(&key, &req) {
                        error!("Error transitioning JIRA issue [{}] to one of [{:?}]: {}", key, to, e);
                    } else {
                        info!("Transitioned [{}] to one of [{:?}]", key, to);
                    }
                },
                Ok(None) => info!("JIRA [{}] cannot be transitioned to  any of [{:?}]", key, to),
                Err(e) => error!("{}", e),
            };
        }

        // add comment only to referenced jiras
        for key in get_referenced_jira_keys(&vec![commit]) {
            if let Err(e) = jira.comment_issue(&key, &ref_msg) {
                error!("Error commenting on key [{}]: {}", key, e);
            }
        }
    }
}

fn try_transition(key: &str, to: &Vec<String>, jira: &jira::api::Session) {
    match find_transition(&key, to, jira) {
        Ok(Some(transition)) => {
            let req = transition.new_request();
            if let Err(e) = jira.transition_issue(&key, &req) {
                error!("Error transitioning JIRA issue [{}] to one of [{:?}]: {}", key, to, e);
            } else {
                info!("Transitioned [{}] to one of [{:?}]", key, to);
            }
        },
        Ok(None) => info!("JIRA [{}] cannot be transitioned to any of [{:?}]", key, to),
        Err(e) => error!("{}", e),
    };
}

fn find_transition(key: &str, to: &Vec<String>, jira: &jira::api::Session) -> Result<Option<Transition>, String> {
    let transitions = match jira.get_transitions(&key) {
        Ok(t) => t,
        Err(e) => return Err(format!("Error looking up JIRA transitions for key [{}]: {}", key, e)),
    };

    Ok(pick_transition(to, &transitions))
}

fn pick_transition(to: &Vec<String>, choices: &Vec<Transition>) -> Option<Transition> {
    for t in choices {
        for name in to {
            if &t.name == name || &t.to.name == name {
                return Some(t.clone())
            }
        }
    }

    None
}

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

    #[test]
    pub fn test_get_jira_keys() {
        let mut commit = Commit::new();
        assert_eq!(Vec::<String>::new(), get_fixed_jira_keys(&vec![commit.clone()]));
        assert_eq!(Vec::<String>::new(), get_referenced_jira_keys(&vec![commit.clone()]));

        commit.commit.message = "Fix [KEY-1][KEY-2], [KEY-3] Some thing that also fixed [KEY-4]".into();
        assert_eq!(vec!["KEY-1", "KEY-2", "KEY-3"], get_fixed_jira_keys(&vec![commit.clone()]));
        assert_eq!(vec!["KEY-4"], get_referenced_jira_keys(&vec![commit.clone()]));

        commit.commit.message += "\n\nFix: [KEY-4], and also mentions [KEY-4], [KEY-5] but not [lowercase-99]";
        assert_eq!(vec!["KEY-1", "KEY-2", "KEY-3", "KEY-4"], get_fixed_jira_keys(&vec![commit.clone()]));
        assert_eq!(vec!["KEY-5"], get_referenced_jira_keys(&vec![commit.clone()]));
    }

    #[test]
    pub fn test_get_jira_keys_alt_format() {
        let mut commit = Commit::new();
        commit.commit.message = "KEY-1, KEY-2:Some thing that also fixed\n\nAlso [KEY-3], OTHER-5".into();
        assert_eq!(Vec::<String>::new(), get_fixed_jira_keys(&vec![commit.clone()]));
        assert_eq!(vec!["KEY-1", "KEY-2", "KEY-3", "OTHER-5"], get_referenced_jira_keys(&vec![commit.clone()]));
    }

    #[test]
    pub fn test_pick_transition() {
        let t1 = Transition {
            id: "1".into(),
            name: "t1".into(),
            to: TransitionTo {
                id: "10".into(),
                name: "inside-t1".into(),
            },
            fields: None,
        };
        let t2 = Transition {
            id: "2".into(),
            name: "t2".into(),
            to: TransitionTo {
                id: "20".into(),
                name: "inside-t2".into(),
            },
            fields: None,
        };
        assert_eq!(Some(t1.clone()), pick_transition(&vec!["t1".into()], &vec![t1.clone(), t2.clone()]));
        assert_eq!(Some(t1.clone()), pick_transition(&vec!["inside-t1".into(), "t2".into()], &vec![t1.clone(), t2.clone()]));
        assert_eq!(Some(t2.clone()), pick_transition(&vec!["inside-t2".into()], &vec![t1.clone(), t2.clone()]));
        assert_eq!(None, pick_transition(&vec!["something-else".into()], &vec![t1.clone(), t2.clone()]));
    }

}
<|endoftext|>


<commit_before>use std::io::net::ip::{SocketAddr, Ipv4Addr};
use std::io::net::get_host_addresses;
use std::io::{Buffer, Reader, Writer, IoResult, BufferedStream, standard_error};
use std::io;
use std::collections::TreeMap;
use url::Url;

#[cfg(test)]
use test::Bencher;
#[cfg(test)]
use serialize::json::ToJson;

use nonce::Nonce;
use message::{WSMessage, WSHeader, WS_MASK, WS_LEN, WS_LEN16, WS_LEN64};
use stream::NetworkStream;

pub struct WebSocket<S = NetworkStream> {
    stream: Option<BufferedStream<S>>,
    connected: bool,
    pub remote_addr: Option<SocketAddr>,
    pub url: Url,
    use_ssl: bool,
}

impl WebSocket {
    pub fn new(url: Url) -> IoResult<WebSocket> {
        let addr = match try!(url.domain()
            .map(|h| get_host_addresses(h)
                 .map(|v| v.into_iter().find(|&a| {
                     match a {
                         Ipv4Addr(..) => true,
                         _ => false
                     }
                 })))
            .unwrap_or(Err(standard_error(io::InvalidInput)))) {
                Some(a) => a,
                None => return Err(standard_error(io::FileNotFound))
            };

        let use_ssl = url.scheme.as_slice() == "wss";

        let port = match url.port() {
            Some(p) => p,
            None if use_ssl => 443,
            _ => 80
        };

        Ok(WebSocket {
            stream: None,
            connected: false,
            remote_addr: Some(SocketAddr{ ip: addr, port: port }),
            url: url,
            use_ssl: use_ssl,
        })
    }

    #[allow(unused_variable)]
    fn try_connect(&mut self) -> IoResult<()> {
        self.stream = Some(BufferedStream::new(try!(self.remote_addr.map(|ref a| NetworkStream::connect(format!("{}", a.ip).as_slice(), a.port, self.use_ssl))
                               .unwrap_or_else(|| Err(standard_error(io::InvalidInput))))));
        Ok(())
    }

    fn write_request(&mut self, nonce: &str) -> IoResult<()> {
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::NotConnected)) };

        try!(s.write(format!("GET {} HTTP/1.1\r\n", self.url.serialize_path().unwrap_or("/".to_string())).as_bytes()));
        try!(s.write(format!("Host: {}\r\n", self.url.host().unwrap()).as_bytes()));
        try!(s.write(format!("Origin: {}\r\n", self.url.serialize_no_fragment()).as_bytes()));
        try!(s.write(format!("Sec-WebSocket-Key: {}\r\n", nonce).as_bytes()));

        try!(s.write(b"Upgrade: websocket\r\n"));
        try!(s.write(b"Connection: Upgrade\r\n"));
        try!(s.write(b"Sec-WebSocket-Protocol: char, superchat\r\n"));
        try!(s.write(b"Sec-WebSocket-Version: 13\r\n"));
        try!(s.write(b"\r\n"));

        s.flush()
    }

    fn read_response(&mut self, nonce: &str) -> IoResult<()> {
        let spaces: &[_] = &[' ', '\t', '\r', '\n'];
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::NotConnected)) };
        let status = try!(s.read_line()).as_slice().splitn(2, ' ').nth(1).and_then(|s| from_str::<uint>(s));

        match status {
            Some(101) => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        let headers = s.lines().map(|r| r.unwrap_or("\r\n".to_string())) .take_while(|l| l.as_slice() != "\r\n")
            .map(|s| s.as_slice().splitn(1, ':').map(|s| s.trim_chars(spaces).to_string()).collect::<Vec<String>>())
            .map(|p| (p[0].to_string(), p[1].to_string()))
            .collect::<TreeMap<String, String>>();

        try!(s.flush());

        let response = headers.find(&"Sec-WebSocket-Accept".to_string());
        match response {
            Some(r) if nonce == r.as_slice() => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        Ok(())
    }

    pub fn connect(&mut self) -> IoResult<()> {
        let mut nonce = Nonce::new();

        try!(self.try_connect());
        try!(self.write_request(nonce.as_slice()));

        nonce = nonce.encode();
        try!(self.read_response(nonce.as_slice()));

        self.connected = true;

        Ok(())
    }

    fn read_header(&mut self) -> IoResult<WSHeader> {
        // XXX: this is a bug, WSHeader should accept u16
        Ok(WSHeader::from_bits_truncate(try!(self.read_be_u16())))
    }

    fn read_length(&mut self, header: &WSHeader) -> IoResult<uint> {
        let wslen = header & WS_LEN;
        if wslen == WS_LEN16 { self.read_be_u16().map(|v| v as uint) }
        else if wslen == WS_LEN64 { self.read_be_u64().map(|v| v as uint) }
        else { Ok(wslen.bits() as uint) }
    }

    pub fn read_message(&mut self) -> IoResult<WSMessage> {
        let header = try!(self.read_header());
        let len = try!(self.read_length(&header));

        let data = if header.contains(WS_MASK) {
            WebSocket::unmask_data(try!(self.read_exact(len)), try!(self.read_be_u32()))
        } else {
            try!(self.read_exact(len))
        };

        Ok(WSMessage { header: header, data: data })
    }

    fn unmask_data(data: Vec<u8>, mask: u32) -> Vec<u8> {
        data.iter().enumerate().map(|(i, b)| b ^ (mask >> ((i % 4) << 3) & 0xff) as u8).collect::<Vec<u8>>()
    }

    // TODO: send_message(&mut self, &WSMessage) -> IoResult<()>

    pub fn iter(&mut self) -> WSMessages {
        WSMessages { sock: self }
    }
}

impl Reader for WebSocket {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
        match self.stream {
            Some(ref mut s) => s.read(buf),
            None => Err(standard_error(io::NotConnected))
        }
    }
}

impl Writer for WebSocket {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.write(buf),
            None => Err(standard_error(io::NotConnected))
        }
    }

    fn flush(&mut self) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.flush(),
            None => Err(standard_error(io::NotConnected))
        }
    }
}

impl Buffer for WebSocket {
    fn fill_buf<'a>(&'a mut self) -> IoResult<&'a [u8]> {
        match self.stream {
            Some(ref mut s) => s.fill_buf(),
            None => Err(standard_error(io::NotConnected))
        }
    }

    fn consume(&mut self, amt: uint) {
        match self.stream {
            Some(ref mut s) => s.consume(amt),
            None => ()
        }
    }
}

pub struct WSMessages<'a> {
    sock: &'a mut WebSocket
}

pub struct WSDefragMessages<'a> {
    underlying: &'a mut WSMessages<'a>,
    buffer: WSMessage
}

impl<'a> WSMessages<'a> {
    pub fn defrag(&'a mut self) -> WSDefragMessages<'a> {
        WSDefragMessages{ underlying: self, buffer: WSMessage{ header: WSHeader.empty(), data: Vec::new() } }
    }
}

impl<'a> Iterator<WSMessage> for WSMessages<'a> {
    fn next(&mut self) -> Option<WSMessage> {
        self.sock.read_message().ok()
    }
}

impl<'a> WSDefragMessages<'a> {
    fn popbuf(&mut self) -> Option<WSMessage> {
        if self.buffer.data.is_empty() {
            None
        } else {
            let buf = WSMessage{ header: WSHeader.empty(), data: Vec::new() };
            mem::swap(self.buffer, &mut buf);
            Some(buf)
        }
    }

    fn swapbuf(&mut self, msg: &mut WSMessage) -> WSMessage {
        mem::swap(self.buffer, msg);
        return msg;
    }
}

impl<'a> Iterator<WSMessage> for WSDefragMessages<'a> {
    fn next(&mut self) -> Option<WSMessage> {
        loop {
            match self.underlying.next() {
                None => return self.popbuf(),
                Some(msg) => if msg.header.contains(WS_FIN) {
                    if msg.header & WS_OPCODE == WS_OPCONT {
                        self.buffer.push(msg);
                        return self.popbuf();
                    } else {
                        return Some(msg);
                    }

                } else {
                    if msg.header & WS_OPCODE == WS_OPCONT {
                        self.buffer.push(msg);
                    } else {
                        return self.swapbuf(&mut msg);
                    }
                }
            }
        }
    }
}

#[bench]
#[allow(dead_code)]
fn test_connect(b: &mut Bencher) {
    let url = Url::parse("wss://stream.pushbullet.com/websocket/").unwrap();
    let mut ws = WebSocket::new(url).unwrap();

    match ws.connect() {
        Err(e) => panic!("error: {}", e),
        _ => ()
    }
    let msg = ws.read_message().unwrap();
    println!("received: {} {} {}", msg, msg.to_string(), msg.to_json());
    for msg in ws.iter() {
        println!("{}", msg.to_string());
    }
}
<commit_msg>TreeMap -> BTreeMap<commit_after>use std::io::net::ip::{SocketAddr, Ipv4Addr};
use std::io::net::get_host_addresses;
use std::io::{Buffer, Reader, Writer, IoResult, BufferedStream, standard_error};
use std::io;
use std::collections::BTreeMap;
use url::Url;

#[cfg(test)]
use test::Bencher;
#[cfg(test)]
use serialize::json::ToJson;

use nonce::Nonce;
use message::{WSMessage, WSHeader, WS_MASK, WS_LEN, WS_LEN16, WS_LEN64};
use stream::NetworkStream;

pub struct WebSocket<S = NetworkStream> {
    stream: Option<BufferedStream<S>>,
    connected: bool,
    pub remote_addr: Option<SocketAddr>,
    pub url: Url,
    use_ssl: bool,
}

impl WebSocket {
    pub fn new(url: Url) -> IoResult<WebSocket> {
        let addr = match try!(url.domain()
            .map(|h| get_host_addresses(h)
                 .map(|v| v.into_iter().find(|&a| {
                     match a {
                         Ipv4Addr(..) => true,
                         _ => false
                     }
                 })))
            .unwrap_or(Err(standard_error(io::InvalidInput)))) {
                Some(a) => a,
                None => return Err(standard_error(io::FileNotFound))
            };

        let use_ssl = url.scheme.as_slice() == "wss";

        let port = match url.port() {
            Some(p) => p,
            None if use_ssl => 443,
            _ => 80
        };

        Ok(WebSocket {
            stream: None,
            connected: false,
            remote_addr: Some(SocketAddr{ ip: addr, port: port }),
            url: url,
            use_ssl: use_ssl,
        })
    }

    #[allow(unused_variable)]
    fn try_connect(&mut self) -> IoResult<()> {
        self.stream = Some(BufferedStream::new(try!(self.remote_addr.map(|ref a| NetworkStream::connect(format!("{}", a.ip).as_slice(), a.port, self.use_ssl))
                               .unwrap_or_else(|| Err(standard_error(io::InvalidInput))))));
        Ok(())
    }

    fn write_request(&mut self, nonce: &str) -> IoResult<()> {
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::NotConnected)) };

        try!(s.write(format!("GET {} HTTP/1.1\r\n", self.url.serialize_path().unwrap_or("/".to_string())).as_bytes()));
        try!(s.write(format!("Host: {}\r\n", self.url.host().unwrap()).as_bytes()));
        try!(s.write(format!("Origin: {}\r\n", self.url.serialize_no_fragment()).as_bytes()));
        try!(s.write(format!("Sec-WebSocket-Key: {}\r\n", nonce).as_bytes()));

        try!(s.write(b"Upgrade: websocket\r\n"));
        try!(s.write(b"Connection: Upgrade\r\n"));
        try!(s.write(b"Sec-WebSocket-Protocol: char, superchat\r\n"));
        try!(s.write(b"Sec-WebSocket-Version: 13\r\n"));
        try!(s.write(b"\r\n"));

        s.flush()
    }

    fn read_response(&mut self, nonce: &str) -> IoResult<()> {
        let spaces: &[_] = &[' ', '\t', '\r', '\n'];
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::NotConnected)) };
        let status = try!(s.read_line()).as_slice().splitn(2, ' ').nth(1).and_then(|s| from_str::<uint>(s));

        match status {
            Some(101) => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        let headers = s.lines().map(|r| r.unwrap_or("\r\n".to_string())) .take_while(|l| l.as_slice() != "\r\n")
            .map(|s| s.as_slice().splitn(1, ':').map(|s| s.trim_chars(spaces).to_string()).collect::<Vec<String>>())
            .map(|p| (p[0].to_string(), p[1].to_string()))
            .collect::<BTreeMap<String, String>>();

        try!(s.flush());

        let response = headers.find(&"Sec-WebSocket-Accept".to_string());
        match response {
            Some(r) if nonce == r.as_slice() => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        Ok(())
    }

    pub fn connect(&mut self) -> IoResult<()> {
        let mut nonce = Nonce::new();

        try!(self.try_connect());
        try!(self.write_request(nonce.as_slice()));

        nonce = nonce.encode();
        try!(self.read_response(nonce.as_slice()));

        self.connected = true;

        Ok(())
    }

    fn read_header(&mut self) -> IoResult<WSHeader> {
        // XXX: this is a bug, WSHeader should accept u16
        Ok(WSHeader::from_bits_truncate(try!(self.read_be_u16())))
    }

    fn read_length(&mut self, header: &WSHeader) -> IoResult<uint> {
        let wslen = header & WS_LEN;
        if wslen == WS_LEN16 { self.read_be_u16().map(|v| v as uint) }
        else if wslen == WS_LEN64 { self.read_be_u64().map(|v| v as uint) }
        else { Ok(wslen.bits() as uint) }
    }

    pub fn read_message(&mut self) -> IoResult<WSMessage> {
        let header = try!(self.read_header());
        let len = try!(self.read_length(&header));

        let data = if header.contains(WS_MASK) {
            WebSocket::unmask_data(try!(self.read_exact(len)), try!(self.read_be_u32()))
        } else {
            try!(self.read_exact(len))
        };

        Ok(WSMessage { header: header, data: data })
    }

    fn unmask_data(data: Vec<u8>, mask: u32) -> Vec<u8> {
        data.iter().enumerate().map(|(i, b)| b ^ (mask >> ((i % 4) << 3) & 0xff) as u8).collect::<Vec<u8>>()
    }

    // TODO: send_message(&mut self, &WSMessage) -> IoResult<()>

    pub fn iter(&mut self) -> WSMessages {
        WSMessages { sock: self }
    }
}

impl Reader for WebSocket {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
        match self.stream {
            Some(ref mut s) => s.read(buf),
            None => Err(standard_error(io::NotConnected))
        }
    }
}

impl Writer for WebSocket {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.write(buf),
            None => Err(standard_error(io::NotConnected))
        }
    }

    fn flush(&mut self) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.flush(),
            None => Err(standard_error(io::NotConnected))
        }
    }
}

impl Buffer for WebSocket {
    fn fill_buf<'a>(&'a mut self) -> IoResult<&'a [u8]> {
        match self.stream {
            Some(ref mut s) => s.fill_buf(),
            None => Err(standard_error(io::NotConnected))
        }
    }

    fn consume(&mut self, amt: uint) {
        match self.stream {
            Some(ref mut s) => s.consume(amt),
            None => ()
        }
    }
}

pub struct WSMessages<'a> {
    sock: &'a mut WebSocket
}

pub struct WSDefragMessages<'a> {
    underlying: &'a mut WSMessages<'a>,
    buffer: WSMessage
}

impl<'a> WSMessages<'a> {
    pub fn defrag(&'a mut self) -> WSDefragMessages<'a> {
        WSDefragMessages{ underlying: self, buffer: WSMessage{ header: WSHeader.empty(), data: Vec::new() } }
    }
}

impl<'a> Iterator<WSMessage> for WSMessages<'a> {
    fn next(&mut self) -> Option<WSMessage> {
        self.sock.read_message().ok()
    }
}

impl<'a> WSDefragMessages<'a> {
    fn popbuf(&mut self) -> Option<WSMessage> {
        if self.buffer.data.is_empty() {
            None
        } else {
            let buf = WSMessage{ header: WSHeader.empty(), data: Vec::new() };
            mem::swap(self.buffer, &mut buf);
            Some(buf)
        }
    }

    fn swapbuf(&mut self, msg: &mut WSMessage) -> WSMessage {
        mem::swap(self.buffer, msg);
        return msg;
    }
}

impl<'a> Iterator<WSMessage> for WSDefragMessages<'a> {
    fn next(&mut self) -> Option<WSMessage> {
        loop {
            match self.underlying.next() {
                None => return self.popbuf(),
                Some(msg) => if msg.header.contains(WS_FIN) {
                    if msg.header & WS_OPCODE == WS_OPCONT {
                        self.buffer.push(msg);
                        return self.popbuf();
                    } else {
                        return Some(msg);
                    }

                } else {
                    if msg.header & WS_OPCODE == WS_OPCONT {
                        self.buffer.push(msg);
                    } else {
                        return self.swapbuf(&mut msg);
                    }
                }
            }
        }
    }
}

#[bench]
#[allow(dead_code)]
fn test_connect(b: &mut Bencher) {
    let url = Url::parse("wss://stream.pushbullet.com/websocket/").unwrap();
    let mut ws = WebSocket::new(url).unwrap();

    match ws.connect() {
        Err(e) => panic!("error: {}", e),
        _ => ()
    }
    let msg = ws.read_message().unwrap();
    println!("received: {} {} {}", msg, msg.to_string(), msg.to_json());
    for msg in ws.iter() {
        println!("{}", msg.to_string());
    }
}
<|endoftext|>


<commit_before>extern crate orbclient;

use std::env;
use std::fs::{self, File};
use std::process::Command;
use std::thread;

use orbclient::{BmpFile, Color, EventOption, Window};

use package::Package;

pub mod package;

fn get_packages() -> Vec<Package> {
    let mut packages: Vec<Package> = Vec::new();

    //TODO: Use a directory walk
    for entry_result in fs::read_dir("/apps/").unwrap() {
        let entry = entry_result.unwrap();
        if entry.file_type().unwrap().is_dir() {
            packages.push(Package::from_path(&("/apps/".to_string() + entry.file_name().to_str().unwrap())));
        }
    }

    packages
}

fn draw(window: &mut Window, packages: &Vec<Package>, shutdown: &BmpFile, mouse_x: i32, mouse_y: i32){
    window.set(Color::rgba(0, 0, 0, 0));
    let w = window.width();
    window.rect(0, 16, w, 32, Color::rgba(0, 0, 0, 128));

    let mut x = 0;
    for package in packages.iter() {
        if package.icon.has_data() {
            let y = window.height() as isize - package.icon.height() as isize;

            if mouse_y >= y as i32 && mouse_x >= x && mouse_x < x + package.icon.width() as i32 {
                window.rect(x as i32, y as i32,
                                  package.icon.width() as u32, package.icon.height() as u32,
                                  Color::rgba(128, 128, 128, 128));

                window.rect(x as i32, y as i32 - 16,
                    package.name.len() as u32 * 8, 16,
                    Color::rgba(0, 0, 0, 128));

                let mut c_x = x;
                for c in package.name.chars() {
                    window.char(c_x as i32, y as i32 - 16, c, Color::rgb(255, 255, 255));
                    c_x += 8;
                }
            }

            window.image(x as i32, y as i32,
                        package.icon.width() as u32,
                        package.icon.height() as u32,
                        &package.icon);
            x = x + package.icon.width() as i32;
        }
    }

    if shutdown.has_data() {
        x = window.width() as i32 - shutdown.width() as i32;
        let y = window.height() as isize - shutdown.height() as isize;

        if mouse_y >= y as i32 && mouse_x >= x &&
           mouse_x < x + shutdown.width() as i32 {
            window.rect(x as i32, y as i32,
                              shutdown.width() as u32, shutdown.height() as u32,
                              Color::rgba(128, 128, 128, 128));
        }

        window.image(x as i32, y as i32,
                        shutdown.width() as u32, shutdown.height() as u32,
                        &shutdown);
        x = x + shutdown.width() as i32;
    }

    window.sync();
}

fn draw_chooser(window: &mut Window, packages: &Vec<Package>, mouse_x: i32, mouse_y: i32){
    let w = window.width();

    window.set(Color::rgb(255, 255, 255));

    let mut y = 0;
    for package in packages.iter() {
        if mouse_y >= y as i32 && mouse_y < y + 32 {
            window.rect(0, y, w, 32, Color::rgb(128, 128, 128));
        }

        if package.icon.has_data() {
            window.image(0, y, package.icon.width() as u32, package.icon.height() as u32, &package.icon);
        }

        let mut c_x = 40;
        for c in package.name.chars() {
            window.char(c_x as i32, y + 8, c, Color::rgb(0, 0, 0));
            c_x += 8;
        }

        y += 32;
    }

    window.sync();
}

fn main() {
    let paths = env::args().skip(1);
    if paths.len() > 0 {
        for ref path in paths {
            let mut packages = get_packages();

            packages.retain(|package| -> bool {
                for accept in package.accepts.iter() {
                    if (accept.starts_with('*') && path.ends_with(&accept[1 ..])) ||
                       (accept.ends_with('*') && path.starts_with(&accept[.. accept.len() - 1])) {
                        return true;
                    }
                }
                false
            });

            if packages.len() > 1 {
                for package in packages.iter() {
                    println!("{:?}: {}", package.binary, package.icon.has_data());
                }

                let mut window = Window::new(-1, -1, 400, packages.len() as u32 * 32, path).unwrap();

                draw_chooser(&mut window, &packages, -1, -1);
                'choosing: loop {
                    for event in window.events() {
                        match event.to_option() {
                            EventOption::Mouse(mouse_event) => {
                                draw_chooser(&mut window, &packages, mouse_event.x, mouse_event.y);

                                if mouse_event.left_button {
                                    let mut y = 0;
                                    for package in packages.iter() {
                                        if package.icon.has_data() {
                                            if mouse_event.y >= y && mouse_event.y < y + 32 {
                                                if let Err(err) = Command::new(&package.binary).arg(path).spawn() {
                                                    println!("{}: Failed to launch: {}", package.binary, err);
                                                }
                                                break 'choosing;
                                            }
                                            y += 32;
                                        }
                                    }
                                }
                            },
                            EventOption::Quit(_) => break 'choosing,
                            _ => ()
                        }
                    }

                    thread::yield_now();
                }
            } else if let Some(package) = packages.get(0) {
                if let Err(err) = Command::new(&package.binary).arg(&path).spawn() {
                    println!("launcher: failed to launch '{}': {}", package.binary, err);
                }
            } else {
                println!("launcher: no application found for '{}'", path);
            }
        }
    } else {
        let packages = get_packages();

        let shutdown = BmpFile::from_path("/ui/actions/system-shutdown.bmp");
        if ! shutdown.has_data() {
            println!("launcher: failed to read shutdown icon");
        }

        let mut window = Window::new(0, 600 - 48, 800, 48, "").unwrap();

        draw(&mut window, &packages, &shutdown, -1, -1);
        'running: loop {
            for event in window.events() {
                match event.to_option() {
                    EventOption::Mouse(mouse_event) => {
                        draw(&mut window, &packages, &shutdown, mouse_event.x, mouse_event.y);

                        if mouse_event.left_button {
                            let mut x = 0;
                            for package in packages.iter() {
                                if package.icon.has_data() {
                                    let y = window.height() as i32 - package.icon.height() as i32;
                                    if mouse_event.y >= y && mouse_event.x >= x &&
                                       mouse_event.x < x + package.icon.width() as i32 {
                                        if let Err(err) = Command::new(&package.binary).spawn() {
                                            println!("{}: Failed to launch: {}", package.binary, err);
                                        }
                                    }
                                    x = x + package.icon.width() as i32;
                                }
                            }

                            if shutdown.has_data() {
                                x = window.width() as i32 - shutdown.width() as i32;
                                let y = window.height() as i32 - shutdown.height() as i32;
                                if mouse_event.y >= y && mouse_event.x >= x &&
                                   mouse_event.x < x + shutdown.width() as i32 {
                                       File::create("acpi:off");
                                       break 'running;
                                }
                            }
                        }
                    },
                    EventOption::Quit(_) => break 'running,
                    _ => ()
                }
            }

            thread::yield_now();
        }
    }
}
<commit_msg>Remove todo<commit_after>extern crate orbclient;

use std::env;
use std::fs::{self, File};
use std::process::Command;
use std::thread;

use orbclient::{BmpFile, Color, EventOption, Window};

use package::Package;

pub mod package;

fn get_packages() -> Vec<Package> {
    let mut packages: Vec<Package> = Vec::new();

    for entry_result in fs::read_dir("/apps/").unwrap() {
        let entry = entry_result.unwrap();
        if entry.file_type().unwrap().is_dir() {
            packages.push(Package::from_path(&("/apps/".to_string() + entry.file_name().to_str().unwrap())));
        }
    }

    packages
}

fn draw(window: &mut Window, packages: &Vec<Package>, shutdown: &BmpFile, mouse_x: i32, mouse_y: i32){
    window.set(Color::rgba(0, 0, 0, 0));
    let w = window.width();
    window.rect(0, 16, w, 32, Color::rgba(0, 0, 0, 128));

    let mut x = 0;
    for package in packages.iter() {
        if package.icon.has_data() {
            let y = window.height() as isize - package.icon.height() as isize;

            if mouse_y >= y as i32 && mouse_x >= x && mouse_x < x + package.icon.width() as i32 {
                window.rect(x as i32, y as i32,
                                  package.icon.width() as u32, package.icon.height() as u32,
                                  Color::rgba(128, 128, 128, 128));

                window.rect(x as i32, y as i32 - 16,
                    package.name.len() as u32 * 8, 16,
                    Color::rgba(0, 0, 0, 128));

                let mut c_x = x;
                for c in package.name.chars() {
                    window.char(c_x as i32, y as i32 - 16, c, Color::rgb(255, 255, 255));
                    c_x += 8;
                }
            }

            window.image(x as i32, y as i32,
                        package.icon.width() as u32,
                        package.icon.height() as u32,
                        &package.icon);
            x = x + package.icon.width() as i32;
        }
    }

    if shutdown.has_data() {
        x = window.width() as i32 - shutdown.width() as i32;
        let y = window.height() as isize - shutdown.height() as isize;

        if mouse_y >= y as i32 && mouse_x >= x &&
           mouse_x < x + shutdown.width() as i32 {
            window.rect(x as i32, y as i32,
                              shutdown.width() as u32, shutdown.height() as u32,
                              Color::rgba(128, 128, 128, 128));
        }

        window.image(x as i32, y as i32,
                        shutdown.width() as u32, shutdown.height() as u32,
                        &shutdown);
        x = x + shutdown.width() as i32;
    }

    window.sync();
}

fn draw_chooser(window: &mut Window, packages: &Vec<Package>, mouse_x: i32, mouse_y: i32){
    let w = window.width();

    window.set(Color::rgb(255, 255, 255));

    let mut y = 0;
    for package in packages.iter() {
        if mouse_y >= y as i32 && mouse_y < y + 32 {
            window.rect(0, y, w, 32, Color::rgb(128, 128, 128));
        }

        if package.icon.has_data() {
            window.image(0, y, package.icon.width() as u32, package.icon.height() as u32, &package.icon);
        }

        let mut c_x = 40;
        for c in package.name.chars() {
            window.char(c_x as i32, y + 8, c, Color::rgb(0, 0, 0));
            c_x += 8;
        }

        y += 32;
    }

    window.sync();
}

fn main() {
    let paths = env::args().skip(1);
    if paths.len() > 0 {
        for ref path in paths {
            let mut packages = get_packages();

            packages.retain(|package| -> bool {
                for accept in package.accepts.iter() {
                    if (accept.starts_with('*') && path.ends_with(&accept[1 ..])) ||
                       (accept.ends_with('*') && path.starts_with(&accept[.. accept.len() - 1])) {
                        return true;
                    }
                }
                false
            });

            if packages.len() > 1 {
                for package in packages.iter() {
                    println!("{:?}: {}", package.binary, package.icon.has_data());
                }

                let mut window = Window::new(-1, -1, 400, packages.len() as u32 * 32, path).unwrap();

                draw_chooser(&mut window, &packages, -1, -1);
                'choosing: loop {
                    for event in window.events() {
                        match event.to_option() {
                            EventOption::Mouse(mouse_event) => {
                                draw_chooser(&mut window, &packages, mouse_event.x, mouse_event.y);

                                if mouse_event.left_button {
                                    let mut y = 0;
                                    for package in packages.iter() {
                                        if package.icon.has_data() {
                                            if mouse_event.y >= y && mouse_event.y < y + 32 {
                                                if let Err(err) = Command::new(&package.binary).arg(path).spawn() {
                                                    println!("{}: Failed to launch: {}", package.binary, err);
                                                }
                                                break 'choosing;
                                            }
                                            y += 32;
                                        }
                                    }
                                }
                            },
                            EventOption::Quit(_) => break 'choosing,
                            _ => ()
                        }
                    }

                    thread::yield_now();
                }
            } else if let Some(package) = packages.get(0) {
                if let Err(err) = Command::new(&package.binary).arg(&path).spawn() {
                    println!("launcher: failed to launch '{}': {}", package.binary, err);
                }
            } else {
                println!("launcher: no application found for '{}'", path);
            }
        }
    } else {
        let packages = get_packages();

        let shutdown = BmpFile::from_path("/ui/actions/system-shutdown.bmp");
        if ! shutdown.has_data() {
            println!("launcher: failed to read shutdown icon");
        }

        let mut window = Window::new(0, 600 - 48, 800, 48, "").unwrap();

        draw(&mut window, &packages, &shutdown, -1, -1);
        'running: loop {
            for event in window.events() {
                match event.to_option() {
                    EventOption::Mouse(mouse_event) => {
                        draw(&mut window, &packages, &shutdown, mouse_event.x, mouse_event.y);

                        if mouse_event.left_button {
                            let mut x = 0;
                            for package in packages.iter() {
                                if package.icon.has_data() {
                                    let y = window.height() as i32 - package.icon.height() as i32;
                                    if mouse_event.y >= y && mouse_event.x >= x &&
                                       mouse_event.x < x + package.icon.width() as i32 {
                                        if let Err(err) = Command::new(&package.binary).spawn() {
                                            println!("{}: Failed to launch: {}", package.binary, err);
                                        }
                                    }
                                    x = x + package.icon.width() as i32;
                                }
                            }

                            if shutdown.has_data() {
                                x = window.width() as i32 - shutdown.width() as i32;
                                let y = window.height() as i32 - shutdown.height() as i32;
                                if mouse_event.y >= y && mouse_event.x >= x &&
                                   mouse_event.x < x + shutdown.width() as i32 {
                                       File::create("acpi:off");
                                       break 'running;
                                }
                            }
                        }
                    },
                    EventOption::Quit(_) => break 'running,
                    _ => ()
                }
            }

            thread::yield_now();
        }
    }
}
<|endoftext|>


<commit_before>//! Uses inter proc analysis to fill in call summary information.

use std::collections::HashSet;

use petgraph::graph::NodeIndex;

use analysis::interproc::transfer::InterProcAnalysis;
use frontend::radeco_containers::{RadecoModule, RadecoFunction, CGInfo};
use middle::ssa::ssa_traits::{SSA, NodeType};
use middle::ssa::cfg_traits::CFG;
use middle::ir::MOpcode;

#[derive(Clone, Debug, Default)]
pub struct CallSummary { }

impl InterProcAnalysis for CallSummary {
    fn new() -> CallSummary {
        Default::default()
    }

    // Compute fn arguments, modifides and returns lists.
    // TODO: Add support for memory args and modifides.
    fn transfer(&mut self, rmod: &mut RadecoModule, fn_ref: u64) {
        {
            // fn arguments.
            let mut args = HashSet::new();
            let mut modifides = HashSet::new();
            let mut returns = HashSet::new();
            let rfn = rmod.functions.get_mut(&fn_ref);
            if rfn.is_none() {
                // Nothing to analyze. This is not a function defined inside the loaded binary. So
                // it must be an import. Use the calling convention information for analysis and
                // return.
                // XXX: For now, we hardcode some analysis information. This will be loaded from
                // r2/Source once the information is avaliable.
                return;
            }
            let rfn = rfn.unwrap();
            {
                let locals = rfn.bindings().iter()
                    .enumerate()
                    .filter(|&(_, l)| l.btype.is_local())
                    .map(|(i, _)| i)
                    .collect::<HashSet<_>>();
                let ssa = rfn.ssa();
                // Get register state at the start block.
                let rs = registers_in_err!(ssa, entry_node_err!(ssa),
                    ssa.invalid_value().unwrap());
                // For every register in the starting block.
                for (i, reg) in ssa.operands_of(rs).iter().enumerate() {
                    // Get the uses of the register
                    let uses = ssa.uses_of(*reg);
                    if !uses.is_empty() {
                        let insert = uses.iter().any(|x| {
                            if let Ok(ref data) = ssa.node_data(*x) {
                                match data.nt {
                                    NodeType::Op(MOpcode::OpLoad) |
                                    NodeType::Op(MOpcode::OpStore) => {
                                        // If the store is to a local variable, then it is still an
                                        // argument. Otherwise, it is part of preservation code and
                                        // must not be considered an argument to the function.
                                        locals.contains(&i)
                                    },
                                    _ => true,
                                }
                            } else {
                                false
                            }
                        });
                        if insert {
                            args.insert(i);
                        }
                    }
                }
                let rs = registers_in_err!(ssa, exit_node_err!(ssa),
                    ssa.invalid_value().unwrap());
                for (i, r) in ssa.operands_of(rs).iter().enumerate() {
                    let (insert_r, insert_m) = if let Ok(ref data) = ssa.node_data(*r) {
                        match data.nt {
                            NodeType::Comment(_) => (false, false),
                            NodeType::Op(MOpcode::OpLoad) => {
                                (locals.contains(&i), true)
                            }
                            _ => (true, true),
                        }
                    } else {
                        (false, false)
                    };

                    if insert_m {
                        modifides.insert(i);
                    }
                    if insert_r {
                        returns.insert(i);
                    }
                }
            }

            rfn.set_returns(&returns.into_iter().map(From::from).collect::<Vec<_>>());
            rfn.set_modifides(&modifides.into_iter().map(From::from).collect::<Vec<_>>());
            rfn.set_args(&args.into_iter().map(From::from).collect::<Vec<_>>());
        }
    }

    // Iterate through all the call-sites in a function(`fn_ref`) and pull in changes from the
    // callees. i.e. Replace args_list by args_list in the callee and replace modifides by the
    // values the callee actually modifies.
    fn propagate(&mut self, rmod: &mut RadecoModule, fn_ref: u64) {
        if let Some(rfn) = rmod.functions.get(&fn_ref) {
            for context in rfn.call_sites(rmod.callgraph()) {
                let callee = rmod.callgraph().callees(context.csite_node).next()
                    .map(|x| x.0).unwrap_or_else(|| {
                    radeco_err!("Call site cannot have callee as `None`");
                    0
                });
                let args = if let Some(callee) = rmod.functions.get(&callee) {
                    callee.args().clone()
                } else {
                    // XXX
                    // No information is inferable. Load data from r2. For now, we use fake data.
                    //unimplemented!()
                    Vec::new()
                };
            }
        }
    }
}

impl CallSummary {
    fn detect_arguments() -> HashSet<NodeIndex> {
        unimplemented!()
    }

    fn detect_modifides() -> HashSet<NodeIndex> {
        unimplemented!()
    }

    fn detect_locals() -> HashSet<NodeIndex> {
        unimplemented!()
    }
}
<commit_msg>Fix compile errors<commit_after>//! Uses inter proc analysis to fill in call summary information.

use std::collections::HashSet;

use petgraph::graph::NodeIndex;

use analysis::interproc::transfer::InterProcAnalysis;
use frontend::radeco_containers::{RadecoModule, RadecoFunction, CGInfo};
use middle::ssa::ssa_traits::{SSA, NodeType};
use middle::ssa::cfg_traits::CFG;
use middle::ir::MOpcode;

#[derive(Clone, Debug, Default)]
pub struct CallSummary { }

impl InterProcAnalysis for CallSummary {
    fn new() -> CallSummary {
        Default::default()
    }

    // Compute fn arguments, modifides and returns lists.
    // TODO: Add support for memory args and modifides.
    fn transfer(&mut self, rmod: &mut RadecoModule, fn_ref: u64) {
        {
            // fn arguments.
            let mut args = HashSet::new();
            let mut modifides = HashSet::new();
            let mut returns = HashSet::new();
            let rfn = rmod.functions.get_mut(&fn_ref);
            if rfn.is_none() {
                // Nothing to analyze. This is not a function defined inside the loaded binary. So
                // it must be an import. Use the calling convention information for analysis and
                // return.
                // XXX: For now, we hardcode some analysis information. This will be loaded from
                // r2/Source once the information is avaliable.
                return;
            }
            let rfn = rfn.unwrap();
            {
                let locals = rfn.bindings().iter()
                    .enumerate()
                    .filter(|&(_, l)| l.btype.is_local())
                    .map(|(i, _)| i)
                    .collect::<HashSet<_>>();
                let ssa = rfn.ssa();
                // Get register state at the start block.
                let rs = registers_in_err!(ssa, entry_node_err!(ssa),
                    ssa.invalid_value().unwrap());
                // For every register in the starting block.
                for (i, reg) in ssa.operands_of(rs).iter().enumerate() {
                    // Get the uses of the register
                    let uses = ssa.uses_of(*reg);
                    if !uses.is_empty() {
                        let insert = uses.iter().any(|x| {
                            if let Ok(ref data) = ssa.node_data(*x) {
                                match data.nt {
                                    NodeType::Op(MOpcode::OpLoad) |
                                    NodeType::Op(MOpcode::OpStore) => {
                                        // If the store is to a local variable, then it is still an
                                        // argument. Otherwise, it is part of preservation code and
                                        // must not be considered an argument to the function.
                                        locals.contains(&i)
                                    },
                                    _ => true,
                                }
                            } else {
                                false
                            }
                        });
                        if insert {
                            args.insert(i);
                        }
                    }
                }
                let rs = registers_in_err!(ssa, exit_node_err!(ssa),
                    ssa.invalid_value().unwrap());
                for (i, r) in ssa.operands_of(rs).iter().enumerate() {
                    let (insert_r, insert_m) = if let Ok(ref data) = ssa.node_data(*r) {
                        match data.nt {
                            NodeType::Comment(_) => (false, false),
                            NodeType::Op(MOpcode::OpLoad) => {
                                (locals.contains(&i), true)
                            }
                            _ => (true, true),
                        }
                    } else {
                        (false, false)
                    };

                    if insert_m {
                        modifides.insert(i);
                    }
                    if insert_r {
                        returns.insert(i);
                    }
                }
            }

            rfn.set_returns(&returns.into_iter().map(From::from).collect::<Vec<_>>());
            rfn.set_modifides(&modifides.into_iter().map(From::from).collect::<Vec<_>>());
            rfn.set_args(&args.into_iter().map(From::from).collect::<Vec<_>>());
        }
    }

    // Iterate through all the call-sites in a function(`fn_ref`) and pull in changes from the
    // callees. i.e. Replace args_list by args_list in the callee and replace modifides by the
    // values the callee actually modifies.
    fn propagate(&mut self, rmod: &mut RadecoModule, fn_ref: u64) {
        if let Some(rfn) = rmod.functions.get(&fn_ref) {
            for context in rfn.call_sites(&rmod.callgraph) {
                let callee = rmod.callgraph.callees(context.csite_node).next()
                    .map(|x| x.0).unwrap_or_else(|| {
                    radeco_err!("Call site cannot have callee as `None`");
                    0
                });
                let args = if let Some(callee) = rmod.functions.get(&callee) {
                    callee.args().clone()
                } else {
                    // XXX
                    // No information is inferable. Load data from r2. For now, we use fake data.
                    //unimplemented!()
                    Vec::new()
                };
            }
        }
    }
}

impl CallSummary {
    fn detect_arguments() -> HashSet<NodeIndex> {
        unimplemented!()
    }

    fn detect_modifides() -> HashSet<NodeIndex> {
        unimplemented!()
    }

    fn detect_locals() -> HashSet<NodeIndex> {
        unimplemented!()
    }
}
<|endoftext|>


<commit_before>#[test]
fn it_works() {
}
<commit_msg>Remove superfluous test<commit_after><|endoftext|>


<commit_before>mod entry;
mod table;
use ::PAGE_SIZE;
use ::Frame;

use self::table::{Table, Level4};
use core::ptr::Unique;

pub use self::entry::*;
use ::FrameAllocator;

pub type PhysicalAddress = usize;
pub type VirtualAddress = usize;

const ENTRY_COUNT: usize = 512;

pub struct Page {
   number: usize,
}

pub struct ActivePageTable {
    p4: Unique<Table<Level4>>,
}

impl ActivePageTable {
    pub unsafe fn new() -> ActivePageTable {
        ActivePageTable {
            p4: Unique::new_unchecked(table::P4),
        }
    }

    fn p4(&self) -> &Table<Level4> {
        unsafe { self.p4.as_ref() }
    }

    fn p4_mut(&mut self) -> &mut Table<Level4> {
        unsafe { self.p4.as_mut() }
    }

    pub fn translate(&self, virtual_address: VirtualAddress)
    -> Option<PhysicalAddress>
    {
        let offset = virtual_address % PAGE_SIZE;
        self.translate_page(Page::containing_address(virtual_address))
            .map(|frame| frame.number * PAGE_SIZE + offset)
    }

    fn translate_page(&self, page: Page) -> Option<Frame> {
        let p3 = self.p4().next_table(page.p4_index());

        let huge_page = || {
            p3.and_then(|p3| {
                let p3_entry = &p3[page.p3_index()];
                // 1GiB page?
                if let Some(start_frame) = p3_entry.pointed_frame() {
                    if p3_entry.flags().contains(EntryFlags::HUGE_PAGE) {
                        // address must be 1GiB aligned
                        assert!(start_frame.number % (ENTRY_COUNT * ENTRY_COUNT) == 0);
                        return Some(Frame {
                            number: start_frame.number + page.p2_index() *
                                    ENTRY_COUNT + page.p1_index(),
                        });
                    }
                }
                if let Some(p2) = p3.next_table(page.p3_index()) {
                    let p2_entry = &p2[page.p2_index()];
                    // 2MiB page?
                    if let Some(start_frame) = p2_entry.pointed_frame() {
                        if p2_entry.flags().contains(EntryFlags::HUGE_PAGE) {
                            // address must be 2MiB aligned
                            assert!(start_frame.number % ENTRY_COUNT == 0);
                            return Some(Frame {
                                number: start_frame.number + page.p1_index()
                            });
                        }
                    }
                }
                None
            })
        };

        p3.and_then(|p3| p3.next_table(page.p3_index()))
        .and_then(|p2| p2.next_table(page.p2_index()))
        .and_then(|p1| p1[page.p1_index()].pointed_frame())
        .or_else(huge_page)
    }

    pub fn map_to<A>(&mut self, page: Page, frame: Frame, flags: EntryFlags, allocator: &mut A)
        where A: FrameAllocator
    {
        let mut p3 = self.p4_mut().next_table_create(page.p4_index(), allocator);
        let mut p2 = p3.next_table_create(page.p3_index(), allocator);
        let mut p1 = p2.next_table_create(page.p2_index(), allocator);

        assert!(p1[page.p1_index()].is_unused());
        p1[page.p1_index()].set(frame, flags | EntryFlags::PRESENT);
    }

    pub fn map<A>(&mut self, page: Page, flags: EntryFlags, allocator: &mut A)
        where A: FrameAllocator
    {
        let frame = allocator.allocate_frame().expect("out of memory");
        self.map_to(page, frame, flags, allocator)
    }

    pub fn identity_map<A>(&mut self, frame: Frame, flags: EntryFlags, allocator: &mut A)
        where A: FrameAllocator
    {
        let page = Page::containing_address(frame.start_address());
        self.map_to(page, frame, flags, allocator)
    }

}

impl Page {
    pub fn containing_address(address: VirtualAddress) -> Page {
        assert!(address < 0x0000_8000_0000_0000 ||
                address >= 0xffff_8000_0000_0000,
                "invalid address: 0x{:x}", address);
        Page { number: address / PAGE_SIZE }
    }

    fn start_address(&self) -> usize {
        self.number * PAGE_SIZE
    }

    fn p4_index(&self) -> usize {
        (self.number >> 27) & 0o777
    }
    fn p3_index(&self) -> usize {
        (self.number >> 18) & 0o777
    }
    fn p2_index(&self) -> usize {
        (self.number >> 9) & 0o777
    }
    fn p1_index(&self) -> usize {
        (self.number >> 0) & 0o777
    }
}
<commit_msg>paging<commit_after>mod entry;
mod table;
use ::PAGE_SIZE;
use ::Frame;

use self::table::{Table, Level4};
use core::ptr::Unique;

pub use self::entry::*;
use ::FrameAllocator;

pub type PhysicalAddress = usize;
pub type VirtualAddress = usize;

const ENTRY_COUNT: usize = 512;

pub struct Page {
   number: usize,
}

pub struct ActivePageTable {
    p4: Unique<Table<Level4>>,
}

impl ActivePageTable {
    pub unsafe fn new() -> ActivePageTable {
        ActivePageTable {
            p4: Unique::new_unchecked(table::P4),
        }
    }

    fn p4(&self) -> &Table<Level4> {
        unsafe { self.p4.as_ref() }
    }

    fn p4_mut(&mut self) -> &mut Table<Level4> {
        unsafe { self.p4.as_mut() }
    }

    pub fn translate(&self, virtual_address: VirtualAddress)
    -> Option<PhysicalAddress>
    {
        let offset = virtual_address % PAGE_SIZE;
        self.translate_page(Page::containing_address(virtual_address))
            .map(|frame| frame.number * PAGE_SIZE + offset)
    }

    fn translate_page(&self, page: Page) -> Option<Frame> {
        let p3 = self.p4().next_table(page.p4_index());

        let huge_page = || {
            p3.and_then(|p3| {
                let p3_entry = &p3[page.p3_index()];
                // 1GiB page?
                if let Some(start_frame) = p3_entry.pointed_frame() {
                    if p3_entry.flags().contains(EntryFlags::HUGE_PAGE) {
                        // address must be 1GiB aligned
                        assert!(start_frame.number % (ENTRY_COUNT * ENTRY_COUNT) == 0);
                        return Some(Frame {
                            number: start_frame.number + page.p2_index() *
                                    ENTRY_COUNT + page.p1_index(),
                        });
                    }
                }
                if let Some(p2) = p3.next_table(page.p3_index()) {
                    let p2_entry = &p2[page.p2_index()];
                    // 2MiB page?
                    if let Some(start_frame) = p2_entry.pointed_frame() {
                        if p2_entry.flags().contains(EntryFlags::HUGE_PAGE) {
                            // address must be 2MiB aligned
                            assert!(start_frame.number % ENTRY_COUNT == 0);
                            return Some(Frame {
                                number: start_frame.number + page.p1_index()
                            });
                        }
                    }
                }
                None
            })
        };

        p3.and_then(|p3| p3.next_table(page.p3_index()))
        .and_then(|p2| p2.next_table(page.p2_index()))
        .and_then(|p1| p1[page.p1_index()].pointed_frame())
        .or_else(huge_page)
    }

    pub fn map_to<A>(&mut self, page: Page, frame: Frame, flags: EntryFlags, allocator: &mut A)
        where A: FrameAllocator
    {
        let mut p3 = self.p4_mut().next_table_create(page.p4_index(), allocator);
        let mut p2 = p3.next_table_create(page.p3_index(), allocator);
        let mut p1 = p2.next_table_create(page.p2_index(), allocator);

        assert!(p1[page.p1_index()].is_unused());
        p1[page.p1_index()].set(frame, flags | EntryFlags::PRESENT);
    }

    pub fn map<A>(&mut self, page: Page, flags: EntryFlags, allocator: &mut A)
        where A: FrameAllocator
    {
        let frame = allocator.allocate_frame().expect("out of memory");
        self.map_to(page, frame, flags, allocator)
    }

    pub fn identity_map<A>(&mut self, frame: Frame, flags: EntryFlags, allocator: &mut A)
        where A: FrameAllocator
    {
        let page = Page::containing_address(frame.start_address());
        self.map_to(page, frame, flags, allocator)
    }

    fn unmap<A>(&mut self, page: Page, allocator: &mut A)
        where A: FrameAllocator
    {
        assert!(self.translate(page.start_address()).is_some());

        let p1 = self.p4_mut()
                     .next_table_mut(page.p4_index())
                     .and_then(|p3| p3.next_table_mut(page.p3_index()))
                     .and_then(|p2| p2.next_table_mut(page.p2_index()))
                     .expect("mapping code does not support huge pages");
        let frame = p1[page.p1_index()].pointed_frame().unwrap();
        p1[page.p1_index()].set_unused();
        // TODO free p(1,2,3) table if empty
        allocator.deallocate_frame(frame);
    }

}

impl Page {
    pub fn containing_address(address: VirtualAddress) -> Page {
        assert!(address < 0x0000_8000_0000_0000 ||
                address >= 0xffff_8000_0000_0000,
                "invalid address: 0x{:x}", address);
        Page { number: address / PAGE_SIZE }
    }

    fn start_address(&self) -> usize {
        self.number * PAGE_SIZE
    }

    fn p4_index(&self) -> usize {
        (self.number >> 27) & 0o777
    }
    fn p3_index(&self) -> usize {
        (self.number >> 18) & 0o777
    }
    fn p2_index(&self) -> usize {
        (self.number >> 9) & 0o777
    }
    fn p1_index(&self) -> usize {
        (self.number >> 0) & 0o777
    }
}
<|endoftext|>