<commit_before>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![deny(unused_imports)]
#![deny(unused_variables)]

extern crate png;
extern crate test;
extern crate regex;
extern crate url;

use std::ascii::AsciiExt;
use std::io;
use std::io::{File, Reader, Command};
use std::io::process::ExitStatus;
use std::io::fs::PathExtensions;
use std::os;
use std::path::Path;
use test::{AutoColor, DynTestName, DynTestFn, TestDesc, TestOpts, TestDescAndFn};
use test::run_tests_console;
use regex::Regex;
use url::Url;

bitflags!(
    flags RenderMode: u32 {
        const CPU_RENDERING  = 0x00000001,
        const GPU_RENDERING  = 0x00000010,
        const LINUX_TARGET   = 0x00000100,
        const MACOS_TARGET   = 0x00001000,
        const ANDROID_TARGET = 0x00010000
    }
)

fn main() {
    let args = os::args();
    let mut parts = args.tail().split(|e| "--" == e.as_slice());

    let harness_args = parts.next().unwrap();  // .split() is never empty
    let servo_args = parts.next().unwrap_or(&[]);

    let (render_mode_string, base_path, testname) = match harness_args {
        [] | [_] => panic!("USAGE: cpu|gpu base_path [testname regex]"),
        [ref render_mode_string, ref base_path] => (render_mode_string, base_path, None),
        [ref render_mode_string, ref base_path, ref testname, ..] => (render_mode_string, base_path, Some(Regex::new(testname.as_slice()).unwrap())),
    };

    let mut render_mode = match render_mode_string.as_slice() {
        "cpu" => CPU_RENDERING,
        "gpu" => GPU_RENDERING,
        _ => panic!("First argument must specify cpu or gpu as rendering mode")
    };
    if cfg!(target_os = "linux") {
        render_mode.insert(LINUX_TARGET);
    }
    if cfg!(target_os = "macos") {
        render_mode.insert(MACOS_TARGET);
    }
    if cfg!(target_os = "android") {
        render_mode.insert(ANDROID_TARGET);
    }

    let mut all_tests = vec!();
    println!("Scanning {} for manifests\n", base_path);

    for file in io::fs::walk_dir(&Path::new(base_path.as_slice())).unwrap() {
        let maybe_extension = file.extension_str();
        match maybe_extension {
            Some(extension) => {
                if extension.to_ascii_lower().as_slice() == "list" && file.is_file() {
                    let tests = parse_lists(&file, servo_args, render_mode, all_tests.len());
                    println!("\t{} [{} tests]", file.display(), tests.len());
                    all_tests.extend(tests.into_iter());
                }
            }
            _ => {}
        }
    }

    let test_opts = TestOpts {
        filter: testname,
        run_ignored: false,
        logfile: None,
        run_tests: true,
        run_benchmarks: false,
        ratchet_noise_percent: None,
        ratchet_metrics: None,
        save_metrics: None,
        test_shard: None,
        nocapture: false,
        color: AutoColor
    };

    match run_tests_console(&test_opts, all_tests) {
        Ok(false) => os::set_exit_status(1), // tests failed
        Err(_) => os::set_exit_status(2),    // I/O-related failure
        _ => (),
    }
}

#[deriving(PartialEq)]
enum ReftestKind {
    Same,
    Different,
}

struct Reftest {
    name: String,
    kind: ReftestKind,
    files: [Path, ..2],
    id: uint,
    servo_args: Vec<String>,
    render_mode: RenderMode,
    is_flaky: bool,
    experimental: bool,
    fragment_identifier: Option<String>,
}

struct TestLine<'a> {
    conditions: &'a str,
    kind: &'a str,
    file_left: &'a str,
    file_right: &'a str,
}

fn parse_lists(file: &Path, servo_args: &[String], render_mode: RenderMode, id_offset: uint) -> Vec<TestDescAndFn> {
    let mut tests = Vec::new();
    let contents = File::open_mode(file, io::Open, io::Read)
                       .and_then(|mut f| f.read_to_string())
                       .ok().expect("Could not read file");

    for line in contents.as_slice().lines() {
        // ignore comments or empty lines
        if line.starts_with("#") || line.is_empty() {
            continue;
        }

        let parts: Vec<&str> = line.split(' ').filter(|p| !p.is_empty()).collect();

        let test_line = match parts.len() {
            3 => TestLine {
                conditions: "",
                kind: parts[0],
                file_left: parts[1],
                file_right: parts[2],
            },
            4 => TestLine {
                conditions: parts[0],
                kind: parts[1],
                file_left: parts[2],
                file_right: parts[3],
            },
            _ => panic!("reftest line: '{:s}' doesn't match '[CONDITIONS] KIND LEFT RIGHT'", line),
        };

        let kind = match test_line.kind {
            "==" => Same,
            "!=" => Different,
            part => panic!("reftest line: '{:s}' has invalid kind '{:s}'", line, part)
        };

        // If we're running this directly, file.dir_path() might be relative.
        // (see issue #3521)
        let base = match file.dir_path().is_relative() {
            true  => os::getcwd().join(file.dir_path()),
            false => file.dir_path()
        };

        let file_left =  base.join(test_line.file_left);
        let file_right = base.join(test_line.file_right);

        let mut conditions_list = test_line.conditions.split(',');
        let mut flakiness = RenderMode::empty();
        let mut experimental = false;
        let mut fragment_identifier = None;
        for condition in conditions_list {
            match condition {
                "flaky_cpu" => flakiness.insert(CPU_RENDERING),
                "flaky_gpu" => flakiness.insert(GPU_RENDERING),
                "flaky_linux" => flakiness.insert(LINUX_TARGET),
                "flaky_macos" => flakiness.insert(MACOS_TARGET),
                "experimental" => experimental = true,
                _ => (),
            }
            if condition.starts_with("fragment=") {
                fragment_identifier = Some(condition.slice_from("fragment=".len()).to_string());
            }
        }

        let reftest = Reftest {
            name: format!("{} {} {}", test_line.file_left, test_line.kind, test_line.file_right),
            kind: kind,
            files: [file_left, file_right],
            id: id_offset + tests.len(),
            render_mode: render_mode,
            servo_args: servo_args.iter().map(|x| x.clone()).collect(),
            is_flaky: render_mode.intersects(flakiness),
            experimental: experimental,
            fragment_identifier: fragment_identifier,
        };

        tests.push(make_test(reftest));
    }
    tests
}

fn make_test(reftest: Reftest) -> TestDescAndFn {
    let name = reftest.name.clone();
    TestDescAndFn {
        desc: TestDesc {
            name: DynTestName(name),
            ignore: false,
            should_fail: false,
        },
        testfn: DynTestFn(proc() {
            check_reftest(reftest);
        }),
    }
}

fn capture(reftest: &Reftest, side: uint) -> (u32, u32, Vec<u8>) {
    let png_filename = format!("/tmp/servo-reftest-{:06u}-{:u}.png", reftest.id, side);
    let mut command = Command::new("target/servo");
    command
        .args(reftest.servo_args.as_slice())
        // Allows pixel perfect rendering of Ahem font for reftests.
        .arg("-Z")
        .arg("disable-text-aa")
        .args(["-f", "-o"])
        .arg(png_filename.as_slice())
        .arg({
            let mut url = Url::from_file_path(&reftest.files[side]).unwrap();
            url.fragment = reftest.fragment_identifier.clone();
            url.to_string()
        });
    // CPU rendering is the default
    if reftest.render_mode.contains(CPU_RENDERING) {
        command.arg("-c");
    }
    if reftest.render_mode.contains(GPU_RENDERING) {
        command.arg("-g");
    }
    if reftest.experimental {
        command.arg("--experimental");
    }
    let retval = match command.status() {
        Ok(status) => status,
        Err(e) => panic!("failed to execute process: {}", e),
    };
    assert_eq!(retval, ExitStatus(0));

    let image = png::load_png(&from_str::<Path>(png_filename.as_slice()).unwrap()).unwrap();
    let rgba8_bytes = match image.pixels {
        png::RGBA8(pixels) => pixels,
        _ => panic!(),
    };
    (image.width, image.height, rgba8_bytes)
}

fn check_reftest(reftest: Reftest) {
    let (left_width, left_height, left_bytes) = capture(&reftest, 0);
    let (right_width, right_height, right_bytes) = capture(&reftest, 1);

    assert_eq!(left_width, right_width);
    assert_eq!(left_height, right_height);

    let left_all_white = left_bytes.iter().all(|&p| p == 255);
    let right_all_white = right_bytes.iter().all(|&p| p == 255);

    if left_all_white && right_all_white {
        panic!("Both renderings are empty")
    }

    let pixels = left_bytes.iter().zip(right_bytes.iter()).map(|(&a, &b)| {
        if a as i8 - b as i8 == 0 {
            // White for correct
            0xFF
        } else {
            // "1100" in the RGBA channel with an error for an incorrect value
            // This results in some number of C0 and FFs, which is much more
            // readable (and distinguishable) than the previous difference-wise
            // scaling but does not require reconstructing the actual RGBA pixel.
            0xC0
        }
    }).collect::<Vec<u8>>();

    if pixels.iter().any(|&a| a < 255) {
        let output_str = format!("/tmp/servo-reftest-{:06u}-diff.png", reftest.id);
        let output = from_str::<Path>(output_str.as_slice()).unwrap();

        let mut img = png::Image {
            width: left_width,
            height: left_height,
            pixels: png::RGBA8(pixels),
        };
        let res = png::store_png(&mut img, &output);
        assert!(res.is_ok());

        match (reftest.kind, reftest.is_flaky) {
            (Same, true) => println!("flaky test - rendering difference: {}", output_str),
            (Same, false) => panic!("rendering difference: {}", output_str),
            (Different, _) => {}   // Result was different and that's what was expected
        }
    } else {
        assert!(reftest.is_flaky || reftest.kind == Same);
    }
}
<commit_msg>auto merge of #4155 : mttr/servo/reftest_unknown_servo_args, r=jdm<commit_after>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![deny(unused_imports)]
#![deny(unused_variables)]

extern crate png;
extern crate test;
extern crate regex;
extern crate url;

use std::ascii::AsciiExt;
use std::io;
use std::io::{File, Reader, Command, IoResult};
use std::io::process::ExitStatus;
use std::io::fs::PathExtensions;
use std::os;
use std::path::Path;
use test::{AutoColor, DynTestName, DynTestFn, TestDesc, TestOpts, TestDescAndFn};
use test::run_tests_console;
use regex::Regex;
use url::Url;

bitflags!(
    flags RenderMode: u32 {
        const CPU_RENDERING  = 0x00000001,
        const GPU_RENDERING  = 0x00000010,
        const LINUX_TARGET   = 0x00000100,
        const MACOS_TARGET   = 0x00001000,
        const ANDROID_TARGET = 0x00010000
    }
)

fn main() {
    let args = os::args();
    let mut parts = args.tail().split(|e| "--" == e.as_slice());

    let harness_args = parts.next().unwrap();  // .split() is never empty
    let servo_args = parts.next().unwrap_or(&[]);

    let (render_mode_string, base_path, testname) = match harness_args {
        [] | [_] => panic!("USAGE: cpu|gpu base_path [testname regex]"),
        [ref render_mode_string, ref base_path] => (render_mode_string, base_path, None),
        [ref render_mode_string, ref base_path, ref testname, ..] => (render_mode_string, base_path, Some(Regex::new(testname.as_slice()).unwrap())),
    };

    let mut render_mode = match render_mode_string.as_slice() {
        "cpu" => CPU_RENDERING,
        "gpu" => GPU_RENDERING,
        _ => panic!("First argument must specify cpu or gpu as rendering mode")
    };
    if cfg!(target_os = "linux") {
        render_mode.insert(LINUX_TARGET);
    }
    if cfg!(target_os = "macos") {
        render_mode.insert(MACOS_TARGET);
    }
    if cfg!(target_os = "android") {
        render_mode.insert(ANDROID_TARGET);
    }

    let mut all_tests = vec!();
    println!("Scanning {} for manifests\n", base_path);

    for file in io::fs::walk_dir(&Path::new(base_path.as_slice())).unwrap() {
        let maybe_extension = file.extension_str();
        match maybe_extension {
            Some(extension) => {
                if extension.to_ascii_lower().as_slice() == "list" && file.is_file() {
                    let tests = parse_lists(&file, servo_args, render_mode, all_tests.len());
                    println!("\t{} [{} tests]", file.display(), tests.len());
                    all_tests.extend(tests.into_iter());
                }
            }
            _ => {}
        }
    }

    let test_opts = TestOpts {
        filter: testname,
        run_ignored: false,
        logfile: None,
        run_tests: true,
        run_benchmarks: false,
        ratchet_noise_percent: None,
        ratchet_metrics: None,
        save_metrics: None,
        test_shard: None,
        nocapture: false,
        color: AutoColor
    };

    match run(test_opts,
              all_tests,
              servo_args.iter().map(|x| x.clone()).collect()) {
        Ok(false) => os::set_exit_status(1), // tests failed
        Err(_) => os::set_exit_status(2),    // I/O-related failure
        _ => (),
    }
}

fn run(test_opts: TestOpts, all_tests: Vec<TestDescAndFn>,
       servo_args: Vec<String>) -> IoResult<bool> {
    // Verify that we're passing in valid servo arguments. Otherwise, servo
    // will exit before we've run any tests, and it will appear to us as if
    // all the tests are failing.
    let mut command = match Command::new("target/servo").args(servo_args.as_slice()).spawn() {
        Ok(p) => p,
        Err(e) => panic!("failed to execute process: {}", e),
    };
    let stderr = command.stderr.as_mut().unwrap().read_to_string().unwrap();

    if stderr.as_slice().contains("Unrecognized") {
        println!("Servo: {}", stderr.as_slice());
        return Ok(false);
    }

    run_tests_console(&test_opts, all_tests)
}

#[deriving(PartialEq)]
enum ReftestKind {
    Same,
    Different,
}

struct Reftest {
    name: String,
    kind: ReftestKind,
    files: [Path, ..2],
    id: uint,
    servo_args: Vec<String>,
    render_mode: RenderMode,
    is_flaky: bool,
    experimental: bool,
    fragment_identifier: Option<String>,
}

struct TestLine<'a> {
    conditions: &'a str,
    kind: &'a str,
    file_left: &'a str,
    file_right: &'a str,
}

fn parse_lists(file: &Path, servo_args: &[String], render_mode: RenderMode, id_offset: uint) -> Vec<TestDescAndFn> {
    let mut tests = Vec::new();
    let contents = File::open_mode(file, io::Open, io::Read)
                       .and_then(|mut f| f.read_to_string())
                       .ok().expect("Could not read file");

    for line in contents.as_slice().lines() {
        // ignore comments or empty lines
        if line.starts_with("#") || line.is_empty() {
            continue;
        }

        let parts: Vec<&str> = line.split(' ').filter(|p| !p.is_empty()).collect();

        let test_line = match parts.len() {
            3 => TestLine {
                conditions: "",
                kind: parts[0],
                file_left: parts[1],
                file_right: parts[2],
            },
            4 => TestLine {
                conditions: parts[0],
                kind: parts[1],
                file_left: parts[2],
                file_right: parts[3],
            },
            _ => panic!("reftest line: '{:s}' doesn't match '[CONDITIONS] KIND LEFT RIGHT'", line),
        };

        let kind = match test_line.kind {
            "==" => Same,
            "!=" => Different,
            part => panic!("reftest line: '{:s}' has invalid kind '{:s}'", line, part)
        };

        // If we're running this directly, file.dir_path() might be relative.
        // (see issue #3521)
        let base = match file.dir_path().is_relative() {
            true  => os::getcwd().join(file.dir_path()),
            false => file.dir_path()
        };

        let file_left =  base.join(test_line.file_left);
        let file_right = base.join(test_line.file_right);

        let mut conditions_list = test_line.conditions.split(',');
        let mut flakiness = RenderMode::empty();
        let mut experimental = false;
        let mut fragment_identifier = None;
        for condition in conditions_list {
            match condition {
                "flaky_cpu" => flakiness.insert(CPU_RENDERING),
                "flaky_gpu" => flakiness.insert(GPU_RENDERING),
                "flaky_linux" => flakiness.insert(LINUX_TARGET),
                "flaky_macos" => flakiness.insert(MACOS_TARGET),
                "experimental" => experimental = true,
                _ => (),
            }
            if condition.starts_with("fragment=") {
                fragment_identifier = Some(condition.slice_from("fragment=".len()).to_string());
            }
        }

        let reftest = Reftest {
            name: format!("{} {} {}", test_line.file_left, test_line.kind, test_line.file_right),
            kind: kind,
            files: [file_left, file_right],
            id: id_offset + tests.len(),
            render_mode: render_mode,
            servo_args: servo_args.iter().map(|x| x.clone()).collect(),
            is_flaky: render_mode.intersects(flakiness),
            experimental: experimental,
            fragment_identifier: fragment_identifier,
        };

        tests.push(make_test(reftest));
    }
    tests
}

fn make_test(reftest: Reftest) -> TestDescAndFn {
    let name = reftest.name.clone();
    TestDescAndFn {
        desc: TestDesc {
            name: DynTestName(name),
            ignore: false,
            should_fail: false,
        },
        testfn: DynTestFn(proc() {
            check_reftest(reftest);
        }),
    }
}

fn capture(reftest: &Reftest, side: uint) -> (u32, u32, Vec<u8>) {
    let png_filename = format!("/tmp/servo-reftest-{:06u}-{:u}.png", reftest.id, side);
    let mut command = Command::new("target/servo");
    command
        .args(reftest.servo_args.as_slice())
        // Allows pixel perfect rendering of Ahem font for reftests.
        .arg("-Z")
        .arg("disable-text-aa")
        .args(["-f", "-o"])
        .arg(png_filename.as_slice())
        .arg({
            let mut url = Url::from_file_path(&reftest.files[side]).unwrap();
            url.fragment = reftest.fragment_identifier.clone();
            url.to_string()
        });
    // CPU rendering is the default
    if reftest.render_mode.contains(CPU_RENDERING) {
        command.arg("-c");
    }
    if reftest.render_mode.contains(GPU_RENDERING) {
        command.arg("-g");
    }
    if reftest.experimental {
        command.arg("--experimental");
    }
    let retval = match command.status() {
        Ok(status) => status,
        Err(e) => panic!("failed to execute process: {}", e),
    };
    assert_eq!(retval, ExitStatus(0));

    let image = png::load_png(&from_str::<Path>(png_filename.as_slice()).unwrap()).unwrap();
    let rgba8_bytes = match image.pixels {
        png::RGBA8(pixels) => pixels,
        _ => panic!(),
    };
    (image.width, image.height, rgba8_bytes)
}

fn check_reftest(reftest: Reftest) {
    let (left_width, left_height, left_bytes) = capture(&reftest, 0);
    let (right_width, right_height, right_bytes) = capture(&reftest, 1);

    assert_eq!(left_width, right_width);
    assert_eq!(left_height, right_height);

    let left_all_white = left_bytes.iter().all(|&p| p == 255);
    let right_all_white = right_bytes.iter().all(|&p| p == 255);

    if left_all_white && right_all_white {
        panic!("Both renderings are empty")
    }

    let pixels = left_bytes.iter().zip(right_bytes.iter()).map(|(&a, &b)| {
        if a as i8 - b as i8 == 0 {
            // White for correct
            0xFF
        } else {
            // "1100" in the RGBA channel with an error for an incorrect value
            // This results in some number of C0 and FFs, which is much more
            // readable (and distinguishable) than the previous difference-wise
            // scaling but does not require reconstructing the actual RGBA pixel.
            0xC0
        }
    }).collect::<Vec<u8>>();

    if pixels.iter().any(|&a| a < 255) {
        let output_str = format!("/tmp/servo-reftest-{:06u}-diff.png", reftest.id);
        let output = from_str::<Path>(output_str.as_slice()).unwrap();

        let mut img = png::Image {
            width: left_width,
            height: left_height,
            pixels: png::RGBA8(pixels),
        };
        let res = png::store_png(&mut img, &output);
        assert!(res.is_ok());

        match (reftest.kind, reftest.is_flaky) {
            (Same, true) => println!("flaky test - rendering difference: {}", output_str),
            (Same, false) => panic!("rendering difference: {}", output_str),
            (Different, _) => {}   // Result was different and that's what was expected
        }
    } else {
        assert!(reftest.is_flaky || reftest.kind == Same);
    }
}
<|endoftext|>


<commit_before>// Copyright 2016, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

use std::slice;
use std::io::{Read, Write, Error};
use std::process;
use std::thread;

use libc::{c_void, c_uint};

use ffi;
use ffi::enums::Status;
use error::IoError;
use ImageSurface;

struct CallbackGuard;

impl Drop for CallbackGuard {
    fn drop(&mut self) {
        if thread::panicking() {
            process::exit(101);
        }
    }
}

struct ReadEnv<'a, R: 'a + Read> {
    reader: &'a mut R,
    error: Option<Error>,
}

unsafe extern "C" fn read_func<R: Read>(closure: *mut c_void, data: *mut u8, len: c_uint) -> Status {
    let _cbguard = CallbackGuard;
    let read_env: &mut ReadEnv<R> = &mut *(closure as *mut ReadEnv<R>);
    let buffer = slice::from_raw_parts_mut(data, len as usize);
    match read_env.reader.read_exact(buffer) {
        Ok(()) => Status::Success,
        Err(error) => {
            read_env.error = Some(error);
            Status::ReadError
        },
    }
}

struct WriteEnv<'a, W: 'a + Write> {
    writer: &'a mut W,
    error: Option<Error>,
}

unsafe extern "C" fn write_func<W: Write>(closure: *mut c_void, data: *mut u8, len: c_uint) -> Status {
    let _cbguard = CallbackGuard;
    let write_env: &mut WriteEnv<W> = &mut *(closure as *mut WriteEnv<W>);
    let buffer = slice::from_raw_parts(data, len as usize);
    match write_env.writer.write_all(buffer) {
        Ok(()) => Status::Success,
        Err(error) => {
            write_env.error = Some(error);
            Status::WriteError
        },
    }
}

impl ImageSurface {
    pub fn create_from_png<R: Read>(stream: &mut R) -> Result<ImageSurface, IoError> {
        let mut env = ReadEnv{ reader: stream, error: None };
        unsafe {
            let raw_surface = ffi::cairo_image_surface_create_from_png_stream(
                Some(read_func::<R>),
                &mut env as *mut ReadEnv<R> as *mut c_void);

            let surface = ImageSurface::from_raw_full(raw_surface)?;

            match env.error {
                None => Ok(surface),
                Some(err) => Err(IoError::Io(err)),
            }
        }
    }

    pub fn write_to_png<W: Write>(&self, stream: &mut W) -> Result<(), IoError> {
        let mut env = WriteEnv{ writer: stream, error: None };
        let status = unsafe { ffi::cairo_surface_write_to_png_stream(self.to_raw_none(),
            Some(write_func::<W>), &mut env as *mut WriteEnv<W> as *mut c_void) };
        match env.error {
            None => match status {
                Status::Success => Ok(()),
                st => Err(IoError::Cairo(st)),
            },
            Some(err) => Err(IoError::Io(err)),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::io::ErrorKind;
    use super::*;
    use ffi::enums::Format;

    struct IoErrorReader;

    // A reader that always returns an error
    impl Read for IoErrorReader {
        fn read(&mut self, _: &mut [u8]) -> Result<usize, Error> {
            Err(Error::new(ErrorKind::Other, "yikes!"))
        }
    }

    #[test]
    fn valid_png_reads_correctly() {
        // A 1x1 PNG, RGB, no alpha, with a single pixel with (42, 42, 42) values
        let png_data: Vec<u8> = vec![
            0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a,  0x00, 0x00, 0x00, 0x0d, 0x49, 0x48, 0x44, 0x52,
            0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,  0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53,
            0xde, 0x00, 0x00, 0x00, 0x0c, 0x49, 0x44, 0x41,  0x54, 0x08, 0xd7, 0x63, 0xd0, 0xd2, 0xd2, 0x02,
            0x00, 0x01, 0x00, 0x00, 0x7f, 0x09, 0xa9, 0x5a,  0x4d, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4e,
            0x44, 0xae, 0x42, 0x60, 0x82
        ];

        let r = ImageSurface::create_from_png(&mut &png_data[..]);
        assert!(r.is_ok());

        let mut surface = r.unwrap();
        assert!(surface.get_width() == 1);
        assert!(surface.get_height() == 1);
        assert!(surface.get_format() == Format::Rgb24);

        let data = surface.get_data().unwrap();
        assert!(data.len() >= 3);

        let slice = &data[0..3];
        assert!(slice[0] == 42);
        assert!(slice[1] == 42);
        assert!(slice[2] == 42);
    }

    #[test]
    fn invalid_png_yields_error() {
        let png_data: Vec<u8> = vec![
            //      v--- this byte is modified
            0x89, 0x40, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a,  0x00, 0x00, 0x00, 0x0d, 0x49, 0x48, 0x44, 0x52,
            0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,  0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53,
            0xde, 0x00, 0x00, 0x00, 0x0c, 0x49, 0x44, 0x41,  0x54, 0x08, 0xd7, 0x63, 0xd0, 0xd2, 0xd2, 0x02,
            0x00, 0x01, 0x00, 0x00, 0x7f, 0x09, 0xa9, 0x5a,  0x4d, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4e,
            0x44, 0xae, 0x42, 0x60, 0x82
        ];

        match ImageSurface::create_from_png(&mut &png_data[..]) {
            Err(IoError::Cairo(_)) => (),
            _ => unreachable!()
        }
    }

    #[test]
    fn io_error_yields_cairo_read_error() {
        let mut r = IoErrorReader;

        match ImageSurface::create_from_png(&mut r) {
            Err(IoError::Cairo(Status::ReadError)) => (),
            _ => unreachable!()
        }
    }
}
<commit_msg>Ignore failing tests on osx<commit_after>// Copyright 2016, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

use std::slice;
use std::io::{Read, Write, Error};
use std::process;
use std::thread;

use libc::{c_void, c_uint};

use ffi;
use ffi::enums::Status;
use error::IoError;
use ImageSurface;

struct CallbackGuard;

impl Drop for CallbackGuard {
    fn drop(&mut self) {
        if thread::panicking() {
            process::exit(101);
        }
    }
}

struct ReadEnv<'a, R: 'a + Read> {
    reader: &'a mut R,
    error: Option<Error>,
}

unsafe extern "C" fn read_func<R: Read>(closure: *mut c_void, data: *mut u8, len: c_uint) -> Status {
    let _cbguard = CallbackGuard;
    let read_env: &mut ReadEnv<R> = &mut *(closure as *mut ReadEnv<R>);
    let buffer = slice::from_raw_parts_mut(data, len as usize);
    match read_env.reader.read_exact(buffer) {
        Ok(()) => Status::Success,
        Err(error) => {
            read_env.error = Some(error);
            Status::ReadError
        },
    }
}

struct WriteEnv<'a, W: 'a + Write> {
    writer: &'a mut W,
    error: Option<Error>,
}

unsafe extern "C" fn write_func<W: Write>(closure: *mut c_void, data: *mut u8, len: c_uint) -> Status {
    let _cbguard = CallbackGuard;
    let write_env: &mut WriteEnv<W> = &mut *(closure as *mut WriteEnv<W>);
    let buffer = slice::from_raw_parts(data, len as usize);
    match write_env.writer.write_all(buffer) {
        Ok(()) => Status::Success,
        Err(error) => {
            write_env.error = Some(error);
            Status::WriteError
        },
    }
}

impl ImageSurface {
    pub fn create_from_png<R: Read>(stream: &mut R) -> Result<ImageSurface, IoError> {
        let mut env = ReadEnv{ reader: stream, error: None };
        unsafe {
            let raw_surface = ffi::cairo_image_surface_create_from_png_stream(
                Some(read_func::<R>),
                &mut env as *mut ReadEnv<R> as *mut c_void);

            let surface = ImageSurface::from_raw_full(raw_surface)?;

            match env.error {
                None => Ok(surface),
                Some(err) => Err(IoError::Io(err)),
            }
        }
    }

    pub fn write_to_png<W: Write>(&self, stream: &mut W) -> Result<(), IoError> {
        let mut env = WriteEnv{ writer: stream, error: None };
        let status = unsafe { ffi::cairo_surface_write_to_png_stream(self.to_raw_none(),
            Some(write_func::<W>), &mut env as *mut WriteEnv<W> as *mut c_void) };
        match env.error {
            None => match status {
                Status::Success => Ok(()),
                st => Err(IoError::Cairo(st)),
            },
            Some(err) => Err(IoError::Io(err)),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::io::ErrorKind;
    use super::*;
    use ffi::enums::Format;

    struct IoErrorReader;

    // A reader that always returns an error
    impl Read for IoErrorReader {
        fn read(&mut self, _: &mut [u8]) -> Result<usize, Error> {
            Err(Error::new(ErrorKind::Other, "yikes!"))
        }
    }

    #[test]
    fn valid_png_reads_correctly() {
        // A 1x1 PNG, RGB, no alpha, with a single pixel with (42, 42, 42) values
        let png_data: Vec<u8> = vec![
            0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a,  0x00, 0x00, 0x00, 0x0d, 0x49, 0x48, 0x44, 0x52,
            0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,  0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53,
            0xde, 0x00, 0x00, 0x00, 0x0c, 0x49, 0x44, 0x41,  0x54, 0x08, 0xd7, 0x63, 0xd0, 0xd2, 0xd2, 0x02,
            0x00, 0x01, 0x00, 0x00, 0x7f, 0x09, 0xa9, 0x5a,  0x4d, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4e,
            0x44, 0xae, 0x42, 0x60, 0x82
        ];

        let r = ImageSurface::create_from_png(&mut &png_data[..]);
        assert!(r.is_ok());

        let mut surface = r.unwrap();
        assert!(surface.get_width() == 1);
        assert!(surface.get_height() == 1);
        assert!(surface.get_format() == Format::Rgb24);

        let data = surface.get_data().unwrap();
        assert!(data.len() >= 3);

        let slice = &data[0..3];
        assert!(slice[0] == 42);
        assert!(slice[1] == 42);
        assert!(slice[2] == 42);
    }

    #[cfg(not(target_os = "macos"))]
    #[test]
    fn invalid_png_yields_error() {
        let png_data: Vec<u8> = vec![
            //      v--- this byte is modified
            0x89, 0x40, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a,  0x00, 0x00, 0x00, 0x0d, 0x49, 0x48, 0x44, 0x52,
            0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,  0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53,
            0xde, 0x00, 0x00, 0x00, 0x0c, 0x49, 0x44, 0x41,  0x54, 0x08, 0xd7, 0x63, 0xd0, 0xd2, 0xd2, 0x02,
            0x00, 0x01, 0x00, 0x00, 0x7f, 0x09, 0xa9, 0x5a,  0x4d, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4e,
            0x44, 0xae, 0x42, 0x60, 0x82
        ];

        match ImageSurface::create_from_png(&mut &png_data[..]) {
            Err(IoError::Cairo(_)) => (),
            _ => unreachable!()
        }
    }

    #[cfg(not(target_os = "macos"))]
    #[test]
    fn io_error_yields_cairo_read_error() {
        let mut r = IoErrorReader;

        match ImageSurface::create_from_png(&mut r) {
            Err(IoError::Cairo(Status::ReadError)) => (),
            _ => unreachable!()
        }
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]
use std::{
    env,
    fs::File,
    io::prelude::*,
    path::Path,
    process::{Command, Stdio},
    sync::{
        atomic::{AtomicUsize, Ordering},
        RwLock,
    },
};

use tempdir::TempDir;

use gpgme::{self, Context, PassphraseRequest, PinentryMode};

#[macro_export]
macro_rules! fail_if_err {
    ($e:expr) => {
        match $e {
            Ok(v) => v,
            Err(err) => panic!("Operation failed: {}", err),
        }
    };
}

macro_rules! count {
    () => {0usize};
    ($_head:tt $($tail:tt)*) => {1usize + count!($($tail)*)};
}

macro_rules! test_case {
    (@impl $name:ident($tester:ident) $body:block,) => {
        #[test]
        fn $name() {
            let $tester = TEST_CASE.new_test();
            $body
        }
    };
    (@impl $name:ident($tester:ident) $body:block,
     $($rest_name:ident($rest_tester:ident) $rest_body:block,)+) => {
        test_case!(@impl $name($tester) $body,);
        test_case!(@impl $($rest_name($rest_tester) $rest_body,)+);
    };
    ($($name:ident($tester:ident) $body:block,)+) => {
        static TEST_CASE: ::once_cell::sync::Lazy<$crate::support::TestCase>
            = ::once_cell::sync::Lazy::new(|| $crate::support::TestCase::new(count!($($name)+)));
        test_case!(@impl $($name($tester) $body,)+);
    };
    ($($name:ident($tester:ident) $body:block),+) => {
        test_case!($($name($tester) $body,)+);
    };
}

pub fn passphrase_cb(_req: PassphraseRequest<'_>, out: &mut dyn Write) -> gpgme::Result<()> {
    out.write_all(b"abc")?;
    Ok(())
}

fn import_key(key: &[u8]) {
    let gpg = env::var_os("GPG").unwrap_or("gpg".into());
    let mut child = Command::new(&gpg)
        .arg("--no-permission-warning")
        .arg("--passphrase")
        .arg("abc")
        .arg("--import")
        .stdin(Stdio::piped())
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()
        .unwrap();
    child.stdin.as_mut().unwrap().write_all(key).unwrap();
    assert!(child.wait().unwrap().success());
}

fn setup_agent(dir: &Path) {
    env::set_var("GNUPGHOME", dir);
    env::set_var("GPG_AGENT_INFO", "");
    let pinentry = {
        let mut path = env::current_exe().unwrap();
        path.pop();
        if path.ends_with("deps") {
            path.pop();
        }
        path.push("pinentry");
        path.set_extension(env::consts::EXE_EXTENSION);
        path
    };
    if !pinentry.exists() {
        panic!("Unable to find pinentry program");
    }

    let agent_conf = dir.join("gpg-agent.conf");
    let mut agent_conf = File::create(agent_conf).unwrap();
    agent_conf
        .write_all(b"ignore-invalid-option allow-loopback-pinentry\n")
        .unwrap();
    agent_conf.write_all(b"allow-loopback-pinentry\n").unwrap();
    agent_conf
        .write_all(b"ignore-invalid-option pinentry-mode\n")
        .unwrap();
    agent_conf.write_all(b"pinentry-mode loopback\n").unwrap();
    agent_conf.write_all(b"pinentry-program ").unwrap();
    agent_conf
        .write_all(pinentry.to_str().unwrap().as_ref())
        .unwrap();
    agent_conf.write_all(b"\n").unwrap();
}

pub struct TestCase {
    count: AtomicUsize,
    homedir: RwLock<Option<TempDir>>,
}

impl TestCase {
    pub fn new(count: usize) -> TestCase {
        let dir = TempDir::new(".test-gpgme").unwrap();
        setup_agent(dir.path());
        import_key(include_bytes!("./data/pubdemo.asc"));
        import_key(include_bytes!("./data/secdemo.asc"));
        TestCase {
            count: AtomicUsize::new(count),
            homedir: RwLock::new(Some(dir)),
        }
    }

    pub fn new_test(&self) -> Test<'_> {
        Test { parent: self }
    }

    pub fn kill_agent(&self) {
        let socket = {
            let homedir = self.homedir.read().unwrap();
            homedir.as_ref().unwrap().path().join("S.gpg-agent")
        };
        let mut child = match Command::new("gpg-connect-agent")
            .arg("-S")
            .arg(socket)
            .stdin(Stdio::piped())
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .spawn()
        {
            Ok(child) => child,
            Err(err) => {
                println!("Unable to kill agent: {}", err);
                return;
            }
        };
        if let Some(ref mut stdin) = child.stdin {
            let _ = stdin.write_all(b"KILLAGENT\nBYE\n");
            let _ = stdin.flush();
        }
        if let Err(err) = child.wait() {
            println!("Unable to kill agent: {}", err);
        }
    }

    fn drop(&self) {
        if self.count.fetch_sub(1, Ordering::SeqCst) == 1 {
            self.kill_agent();
            self.homedir.write().unwrap().take();
        }
    }
}

pub struct Test<'a> {
    parent: &'a TestCase,
}

impl<'a> Drop for Test<'a> {
    fn drop(&mut self) {
        self.parent.drop();
    }
}

impl<'a> Test<'a> {
    pub fn create_context(&self) -> Context {
        let mut ctx = fail_if_err!(Context::from_protocol(gpgme::Protocol::OpenPgp));
        let _ = ctx.set_pinentry_mode(PinentryMode::Loopback);
        ctx
    }
}
<commit_msg>Fix keysign test for recent versions of gpg<commit_after>#![allow(dead_code)]
use std::{
    env, fs,
    io::prelude::*,
    path::Path,
    process::{Command, Stdio},
    sync::{
        atomic::{AtomicUsize, Ordering},
        RwLock,
    },
};

use tempdir::TempDir;

use gpgme::{self, Context, PassphraseRequest, PinentryMode};

#[macro_export]
macro_rules! fail_if_err {
    ($e:expr) => {
        match $e {
            Ok(v) => v,
            Err(err) => panic!("Operation failed: {}", err),
        }
    };
}

macro_rules! count {
    () => {0usize};
    ($_head:tt $($tail:tt)*) => {1usize + count!($($tail)*)};
}

macro_rules! test_case {
    (@impl $name:ident($tester:ident) $body:block,) => {
        #[test]
        fn $name() {
            let $tester = TEST_CASE.new_test();
            $body
        }
    };
    (@impl $name:ident($tester:ident) $body:block,
     $($rest_name:ident($rest_tester:ident) $rest_body:block,)+) => {
        test_case!(@impl $name($tester) $body,);
        test_case!(@impl $($rest_name($rest_tester) $rest_body,)+);
    };
    ($($name:ident($tester:ident) $body:block,)+) => {
        static TEST_CASE: ::once_cell::sync::Lazy<$crate::support::TestCase>
            = ::once_cell::sync::Lazy::new(|| $crate::support::TestCase::new(count!($($name)+)));
        test_case!(@impl $($name($tester) $body,)+);
    };
    ($($name:ident($tester:ident) $body:block),+) => {
        test_case!($($name($tester) $body,)+);
    };
}

pub fn passphrase_cb(_req: PassphraseRequest<'_>, out: &mut dyn Write) -> gpgme::Result<()> {
    out.write_all(b"abc")?;
    Ok(())
}

fn import_key(key: &[u8]) {
    let gpg = env::var_os("GPG").unwrap_or("gpg".into());
    let mut child = Command::new(&gpg)
        .arg("--no-permission-warning")
        .arg("--passphrase")
        .arg("abc")
        .arg("--import")
        .stdin(Stdio::piped())
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()
        .unwrap();
    child.stdin.as_mut().unwrap().write_all(key).unwrap();
    assert!(child.wait().unwrap().success());
}

fn setup_agent(dir: &Path) {
    env::set_var("GNUPGHOME", dir);
    env::set_var("GPG_AGENT_INFO", "");
    let pinentry = {
        let mut path = env::current_exe().unwrap();
        path.pop();
        if path.ends_with("deps") {
            path.pop();
        }
        path.push("pinentry");
        path.set_extension(env::consts::EXE_EXTENSION);
        path
    };
    if !pinentry.exists() {
        panic!("Unable to find pinentry program");
    }

    let conf = dir.join("gpg.conf");
    fs::write(
        conf,
        "ignore-invalid-option allow-weak-key-signatures\n\
         allow-weak-key-signatures\n",
    )
    .unwrap();

    let agent_conf = dir.join("gpg-agent.conf");
    fs::write(
        agent_conf,
        format!(
            "ignore-invalid-option allow-loopback-pinentry\n\
             allow-loopback-pinentry\n\
             ignore-invalid-option pinentry-mode\n\
             pinentry-mode loopback\n\
             pinentry-program {}\n",
            pinentry.to_str().unwrap()
        ),
    )
    .unwrap();
}

pub struct TestCase {
    count: AtomicUsize,
    homedir: RwLock<Option<TempDir>>,
}

impl TestCase {
    pub fn new(count: usize) -> TestCase {
        let dir = TempDir::new(".test-gpgme").unwrap();
        setup_agent(dir.path());
        import_key(include_bytes!("./data/pubdemo.asc"));
        import_key(include_bytes!("./data/secdemo.asc"));
        TestCase {
            count: AtomicUsize::new(count),
            homedir: RwLock::new(Some(dir)),
        }
    }

    pub fn new_test(&self) -> Test<'_> {
        Test { parent: self }
    }

    pub fn kill_agent(&self) {
        let socket = {
            let homedir = self.homedir.read().unwrap();
            homedir.as_ref().unwrap().path().join("S.gpg-agent")
        };
        let mut child = match Command::new("gpg-connect-agent")
            .arg("-S")
            .arg(socket)
            .stdin(Stdio::piped())
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .spawn()
        {
            Ok(child) => child,
            Err(err) => {
                println!("Unable to kill agent: {}", err);
                return;
            }
        };
        if let Some(ref mut stdin) = child.stdin {
            let _ = stdin.write_all(b"KILLAGENT\nBYE\n");
            let _ = stdin.flush();
        }
        if let Err(err) = child.wait() {
            println!("Unable to kill agent: {}", err);
        }
    }

    fn drop(&self) {
        if self.count.fetch_sub(1, Ordering::SeqCst) == 1 {
            self.kill_agent();
            self.homedir.write().unwrap().take();
        }
    }
}

pub struct Test<'a> {
    parent: &'a TestCase,
}

impl Drop for Test<'_> {
    fn drop(&mut self) {
        self.parent.drop();
    }
}

impl Test<'_> {
    pub fn create_context(&self) -> Context {
        let mut ctx = fail_if_err!(Context::from_protocol(gpgme::Protocol::OpenPgp));
        let _ = ctx.set_pinentry_mode(PinentryMode::Loopback);
        ctx
    }
}
<|endoftext|>


<commit_before>use distribution;
use random;

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    k: usize,
    p: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// It should hold that `p[i] >= 0`, `p[i] <= 1`, and `sum(p) == 1`.
    #[inline]
    pub fn new(p: &[f64]) -> Categorical {
        should!(is_probability_vector(p), {
            const EPSILON: f64 = 1e-12;
            let mut in_unit = true;
            let mut sum = 0.0;
            for &p in p.iter() {
                if p < 0.0 || p > 1.0 {
                    in_unit = false;
                    break;
                }
                sum += p;
            }
            in_unit && (sum - 1.0).abs() <= EPSILON
        });
        Categorical { k: p.len(), p: p.to_vec() }
    }

    /// Return the number of categories.
    #[inline(always)]
    pub fn k(&self) -> usize { self.k }

    /// Return the event probabilities.
    #[inline(always)]
    pub fn p(&self) -> &[f64] { &self.p }
}

impl distribution::Distribution for Categorical {
    type Value = usize;

    fn cdf(&self, x: f64) -> f64 {
        if x < 0.0 {
            return 0.0;
        }
        let x = x as usize;
        if x >= self.k - 1 {
            1.0
        } else {
            self.p.iter().take(x + 1).fold(0.0, |a, b| a + b)
        }
    }
}

impl distribution::Discrete for Categorical {
    #[inline]
    fn pmf(&self, x: usize) -> f64 {
        should!(x < self.k);
        self.p[x]
    }
}

impl distribution::Entropy for Categorical {
    fn entropy(&self) -> f64 {
        -self.p.iter().fold(0.0, |sum, p| sum + p * p.ln())
    }
}

impl distribution::Inverse for Categorical {
    fn inv_cdf(&self, p: f64) -> usize {
        should!(0.0 <= p && p <= 1.0);
        let mut sum = 0.0;
        self.p.iter().position(|&pi| {
            sum += pi;
            pi > 0.0 && sum >= p
        }).unwrap_or_else(|| {
            self.p.iter().rposition(|&pi| pi > 0.0).unwrap()
        })
    }
}

impl distribution::Kurtosis for Categorical {
    fn kurtosis(&self) -> f64 {
        use distribution::{Mean, Variance};
        let (mean, variance) = (self.mean(), self.variance());
        let kurt = self.p.iter().enumerate().fold(0.0, |sum, (i, p)| {
            sum + (i as f64 - mean).powi(4) * p
        });
        kurt / variance.powi(2) - 3.0
    }
}

impl distribution::Mean for Categorical {
    fn mean(&self) -> f64 {
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + i as f64 * p)
    }
}

impl distribution::Median for Categorical {
    fn median(&self) -> f64 {
        if self.p[0] > 0.5 {
            return 0.0;
        } else if self.p[0] == 0.5 {
            return 0.5;
        }
        let mut sum = 0.0;
        for i in 0..self.k {
            sum += self.p[i];
            if sum == 0.5 {
                return (2 * i - 1) as f64 / 2.0;
            } else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }
}

impl distribution::Modes for Categorical {
    fn modes(&self) -> Vec<usize> {
        let mut modes = Vec::new();
        let mut max = 0.0;
        for (i, &p) in self.p.iter().enumerate() {
            if p == max {
                modes.push(i);
            }
            if p > max {
                max = p;
                modes = vec![i];
            }
        }
        modes
    }
}

impl distribution::Sample for Categorical {
    #[inline]
    fn sample<S>(&self, source: &mut S) -> usize where S: random::Source {
        use distribution::Inverse;
        self.inv_cdf(source.read::<f64>())
    }
}

impl distribution::Skewness for Categorical {
    fn skewness(&self) -> f64 {
        use distribution::{Mean, Variance};
        let (mean, variance) = (self.mean(), self.variance());
        let skew = self.p.iter().enumerate().fold(0.0, |sum, (i, p)| {
            sum + (i as f64 - mean).powi(3) * p
        });
        skew / (variance * variance.sqrt())
    }
}

impl distribution::Variance for Categorical {
    fn variance(&self) -> f64 {
        use distribution::Mean;
        let mean = self.mean();
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| {
            sum + (i as f64 - mean).powi(2) * p
        })
    }
}

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

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1.0 / $k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    fn cdf() {
        let d = new!([0.0, 0.75, 0.25, 0.0]);
        let p = vec![0.0, 0.0, 0.75, 1.0, 1.0];

        let x = (-1..4).map(|x| d.cdf(x as f64)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let x = (-1..4).map(|x| d.cdf(x as f64 + 0.5)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let d = new!(equal 3);
        let p = vec![0.0, 1.0 / 3.0, 2.0 / 3.0, 1.0];

        let x = (-1..3).map(|x| d.cdf(x as f64)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let x = (-1..3).map(|x| d.cdf(x as f64 + 0.5)).collect::<Vec<_>>();
        assert_eq!(&x, &p);
    }

    #[test]
    fn pmf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d = new!(p);
        assert_eq!(&(0..4).map(|x| d.pmf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d = new!(equal 3);
        assert_eq!(&(0..3).map(|x| d.pmf(x)).collect::<Vec<_>>(), &vec![1.0 / 3.0; 3])
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn inv_cdf() {
        let d = new!([0.0, 0.75, 0.25, 0.0]);
        let p = vec![0.0, 0.75, 0.7500001, 1.0];
        assert_eq!(&p.iter().map(|&p| d.inv_cdf(p)).collect::<Vec<_>>(), &vec![1, 1, 2, 2]);

        let d = new!(equal 3);
        let p = vec![0.0, 0.5, 0.75, 1.0];
        assert_eq!(&p.iter().map(|&p| d.inv_cdf(p)).collect::<Vec<_>>(), &vec![0, 1, 2, 2]);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.0);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).mean(), 1.5);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).modes(), vec![1, 2]);
    }

    #[test]
    fn sample() {
        let mut source = random::default();
        let sum = Independent(&new!([0.0, 0.5, 0.5]), &mut source).take(100).fold(0, |a, b| a + b);
        assert!(100 <= sum && sum <= 200);
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).skewness(), 0.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn variance() {
        assert_eq!(new!(equal 3).variance(), 2.0 / 3.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).variance(), 11.0 / 12.0);
    }
}
<commit_msg>categorical: cache the cumulative sum<commit_after>use distribution;
use random;

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    k: usize,
    p: Vec<f64>,
    cumsum: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// It should hold that `p[i] >= 0`, `p[i] <= 1`, and `sum(p) == 1`.
    pub fn new(p: &[f64]) -> Categorical {
        should!(is_probability_vector(p), {
            const EPSILON: f64 = 1e-12;
            p.iter().all(|&p| p >= 0.0 && p <= 1.0) &&
                (p.iter().fold(0.0, |sum, &p| sum + p) - 1.0).abs() < EPSILON
        });

        let k = p.len();
        let mut cumsum = p.to_vec();
        for i in 1..(k - 1) {
            cumsum[i] += cumsum[i - 1];
        }
        cumsum[k - 1] = 1.0;
        Categorical { k: k, p: p.to_vec(), cumsum: cumsum }
    }

    /// Return the number of categories.
    #[inline(always)]
    pub fn k(&self) -> usize { self.k }

    /// Return the event probabilities.
    #[inline(always)]
    pub fn p(&self) -> &[f64] { &self.p }
}

impl distribution::Distribution for Categorical {
    type Value = usize;

    fn cdf(&self, x: f64) -> f64 {
        if x < 0.0 {
            return 0.0;
        }
        let x = x as usize;
        if x >= self.k {
            return 1.0;
        }
        self.cumsum[x]
    }
}

impl distribution::Discrete for Categorical {
    #[inline]
    fn pmf(&self, x: usize) -> f64 {
        should!(x < self.k);
        self.p[x]
    }
}

impl distribution::Entropy for Categorical {
    fn entropy(&self) -> f64 {
        -self.p.iter().fold(0.0, |sum, p| sum + p * p.ln())
    }
}

impl distribution::Inverse for Categorical {
    fn inv_cdf(&self, p: f64) -> usize {
        should!(0.0 <= p && p <= 1.0);
        self.cumsum.iter().position(|&sum| sum > 0.0 && sum >= p).unwrap_or_else(|| {
            self.p.iter().rposition(|&p| p > 0.0).unwrap()
        })
    }
}

impl distribution::Kurtosis for Categorical {
    fn kurtosis(&self) -> f64 {
        use distribution::{Mean, Variance};
        let (mean, variance) = (self.mean(), self.variance());
        let kurt = self.p.iter().enumerate().fold(0.0, |sum, (i, p)| {
            sum + (i as f64 - mean).powi(4) * p
        });
        kurt / variance.powi(2) - 3.0
    }
}

impl distribution::Mean for Categorical {
    fn mean(&self) -> f64 {
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + i as f64 * p)
    }
}

impl distribution::Median for Categorical {
    fn median(&self) -> f64 {
        if self.p[0] > 0.5 {
            return 0.0;
        }
        if self.p[0] == 0.5 {
            return 0.5;
        }
        for (i, &sum) in self.cumsum.iter().enumerate() {
            if sum == 0.5 {
                return (2 * i - 1) as f64 / 2.0;
            } else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }
}

impl distribution::Modes for Categorical {
    fn modes(&self) -> Vec<usize> {
        let mut modes = Vec::new();
        let mut max = 0.0;
        for (i, &p) in self.p.iter().enumerate() {
            if p == max {
                modes.push(i);
            }
            if p > max {
                max = p;
                modes = vec![i];
            }
        }
        modes
    }
}

impl distribution::Sample for Categorical {
    #[inline]
    fn sample<S>(&self, source: &mut S) -> usize where S: random::Source {
        use distribution::Inverse;
        self.inv_cdf(source.read::<f64>())
    }
}

impl distribution::Skewness for Categorical {
    fn skewness(&self) -> f64 {
        use distribution::{Mean, Variance};
        let (mean, variance) = (self.mean(), self.variance());
        let skew = self.p.iter().enumerate().fold(0.0, |sum, (i, p)| {
            sum + (i as f64 - mean).powi(3) * p
        });
        skew / (variance * variance.sqrt())
    }
}

impl distribution::Variance for Categorical {
    fn variance(&self) -> f64 {
        use distribution::Mean;
        let mean = self.mean();
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| {
            sum + (i as f64 - mean).powi(2) * p
        })
    }
}

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

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1.0 / $k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    fn cdf() {
        let d = new!([0.0, 0.75, 0.25, 0.0]);
        let p = vec![0.0, 0.0, 0.75, 1.0, 1.0];

        let x = (-1..4).map(|x| d.cdf(x as f64)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let x = (-1..4).map(|x| d.cdf(x as f64 + 0.5)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let d = new!(equal 3);
        let p = vec![0.0, 1.0 / 3.0, 2.0 / 3.0, 1.0];

        let x = (-1..3).map(|x| d.cdf(x as f64)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let x = (-1..3).map(|x| d.cdf(x as f64 + 0.5)).collect::<Vec<_>>();
        assert_eq!(&x, &p);
    }

    #[test]
    fn pmf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d = new!(p);
        assert_eq!(&(0..4).map(|x| d.pmf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d = new!(equal 3);
        assert_eq!(&(0..3).map(|x| d.pmf(x)).collect::<Vec<_>>(), &vec![1.0 / 3.0; 3])
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn inv_cdf() {
        let d = new!([0.0, 0.75, 0.25, 0.0]);
        let p = vec![0.0, 0.75, 0.7500001, 1.0];
        assert_eq!(&p.iter().map(|&p| d.inv_cdf(p)).collect::<Vec<_>>(), &vec![1, 1, 2, 2]);

        let d = new!(equal 3);
        let p = vec![0.0, 0.5, 0.75, 1.0];
        assert_eq!(&p.iter().map(|&p| d.inv_cdf(p)).collect::<Vec<_>>(), &vec![0, 1, 2, 2]);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.0);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).mean(), 1.5);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).modes(), vec![1, 2]);
    }

    #[test]
    fn sample() {
        let mut source = random::default();
        let sum = Independent(&new!([0.0, 0.5, 0.5]), &mut source).take(100).fold(0, |a, b| a + b);
        assert!(100 <= sum && sum <= 200);
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).skewness(), 0.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn variance() {
        assert_eq!(new!(equal 3).variance(), 2.0 / 3.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).variance(), 11.0 / 12.0);
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

/// Represents an identifier name in program source code.
pub type Identifier = usize;

#[derive(Debug)]
pub struct NameTable<'a> {
    identifier_names: HashMap<Identifier, &'a str>,
    identifier_ids: HashMap<&'a str, Identifier>,
    max_id: usize,
}

impl<'a> NameTable<'a> {
    pub fn new() -> NameTable<'a> {
        NameTable {
            identifier_names: HashMap::new(),
            identifier_ids: HashMap::new(),
            max_id: 0,
        }
    }

    pub fn new_id(&mut self, name: &'a str) -> Identifier {
        if let Some(id) = self.get_id(name) {
            return id;
        }
        let id = self.max_id;
        self.max_id += 1;
        self.identifier_names.insert(id, name);
        self.identifier_ids.entry(name).or_insert(id);
        id
    }
    /// Creates a new identifier which cannot be looked up by name, since it doesn't have one.
    pub fn new_anon(&mut self) -> Identifier {
        let id = self.max_id;
        self.max_id += 1;
        self.identifier_names.insert(id, "*anon*");
        id
    }
    pub fn get_id(&'a self, name: &'a str) -> Option<Identifier> {
        self.identifier_ids.get(name).map(|x| *x)
    }
    pub fn get_name(&'a self, id: Identifier) -> Option<&'a str> {
        self.identifier_names.get(&id).map(|x| *x)
    }
}
<commit_msg>Add the ability to check if an identifier is anon<commit_after>use std::collections::HashMap;

/// Represents an identifier name in program source code.
pub type Identifier = usize;

#[derive(Debug)]
pub struct NameTable<'a> {
    identifier_names: HashMap<Identifier, &'a str>,
    identifier_ids: HashMap<&'a str, Identifier>,
    max_id: usize,
}

const ANON_NAME: &'static str = "*anon*";

impl<'a> NameTable<'a> {
    pub fn new() -> NameTable<'a> {
        NameTable {
            identifier_names: HashMap::new(),
            identifier_ids: HashMap::new(),
            max_id: 0,
        }
    }

    pub fn new_id(&mut self, name: &'a str) -> Identifier {
        if let Some(id) = self.get_id(name) {
            return id;
        }
        let id = self.max_id;
        self.max_id += 1;
        self.identifier_names.insert(id, name);
        self.identifier_ids.entry(name).or_insert(id);
        id
    }
    /// Creates a new identifier which cannot be looked up by name.
    pub fn new_anon(&mut self) -> Identifier {
        let id = self.max_id;
        self.max_id += 1;
        self.identifier_names.insert(id, ANON_NAME);
        id
    }
    pub fn is_anon(&self, id: Identifier) -> Option<bool> {
        self.identifier_names.get(&id).map(|&x| x == ANON_NAME)
    }
    pub fn get_id(&'a self, name: &'a str) -> Option<Identifier> {
        self.identifier_ids.get(name).map(|x| *x)
    }
    pub fn get_name(&'a self, id: Identifier) -> Option<&'a str> {
        self.identifier_names.get(&id).map(|x| *x)
    }
}
<|endoftext|>


<commit_before>use std::borrow::Borrow;
use std::cmp::{max, min};
use std::io::Write;
use std::io;
use std::mem;

use rustbox::{RustBox, Key};
use tui::style::Style;
use tui::widget::{Widget, WidgetRet};

static LINEBREAK : char = '\\';

pub struct MsgArea {
    lines  : Vec<Line>,

    width  : i32,
    height : i32,

    /// Total number of lines added to the widget. Note that this is not the
    /// same thing as lines.len(): This simply returns number of lines added to
    /// the widget, rather than number of lines shown by the widget.
    total_lines : i32,

    /// Vertical scroll
    scroll : i32,

    // TODO: logging
}

#[derive(Debug)]
struct Line {
    /// A line. INVARIANT: Not longer than the width of the widget, and shown as
    /// single line in the TUI.
    msg : String,

    /// Is this continuation of previous line?
    continuation : bool,

    /// Index of the line in the buffer. Continuations have same indexes as the
    /// original message.
    line_idx : i32,

    style : Style,
}

impl MsgArea {
    pub fn new(width : i32, height : i32) -> MsgArea {
        MsgArea {
            lines: Vec::new(),
            width: width,
            height: height,
            total_lines: 0,
            scroll: 0,
        }
    }

    ////////////////////////////////////////////////////////////////////////////
    // Resizing

    fn resize_(&mut self, width : i32, height : i32) {
        // either scroll to bottom ...
        let scroll_to_bottom       = self.need_to_scroll();
        // ... or make sure the first visible line is still visible
        let first_visible_line_idx =
            if self.lines.is_empty() { 0 } else { self.lines[self.scroll as usize].line_idx };

        // Update the size
        self.width = width;
        self.height = height;

        // Combine/re-split lines
        self.resplit_lines();

        // Update the scroll
        if (self.lines.len() as i32) < self.height {
            self.scroll = 0;
        } else if scroll_to_bottom {
            self.scroll = (self.lines.len() as i32) - self.height;
        } else {
            match self.find_line_idx(first_visible_line_idx) {
                None => {/* TODO: Log this somewhere, this is a bug */},
                Some(idx) => { self.scroll = idx as i32; }
            }
        }
    }

    /// Combine continuations with original lines, add lines again in the
    /// original order. Should be called after updating width to have any
    /// effect.
    fn resplit_lines(&mut self) {
        // We could probably modify in-place using two indexes, but whatever.

        let old_lines : Vec<Line> = {
            let total_lines = self.lines.len();
            mem::replace(&mut self.lines, Vec::with_capacity(total_lines))
        };

        writeln!(&mut io::stderr(), "old lines: {:?}", old_lines).unwrap();

        let mut line_idx = 0;
        while line_idx < old_lines.len() {
            // How many bytes combined string needs?
            let mut total_len = old_lines[line_idx].msg.len();
            {
                let mut cont_idx  = line_idx + 1;
                while cont_idx < old_lines.len() && old_lines[cont_idx].continuation {
                    total_len += old_lines[cont_idx].msg.len();
                    cont_idx  += 1;
                }
            }

            // Combine all lines
            let mut new_line = String::with_capacity(total_len);
            new_line.push_str(old_lines[line_idx].msg.borrow());
            let mut cont_idx = line_idx + 1;
            while cont_idx < old_lines.len() && old_lines[cont_idx].continuation {
                new_line.push_str(old_lines[cont_idx].msg.borrow());
                cont_idx += 1;
            }

            // Finally add the combined line to the fresh buffer
            self.add_msg_(new_line.chars().collect::<Vec<char>>().borrow(),
                          old_lines[line_idx].style);

            line_idx = cont_idx;
        }

        writeln!(&mut io::stderr(), "new lines: {:?}", self.lines).unwrap();
    }

    /// Find index of given line in 'self.lines'.
    fn find_line_idx(&self, line_idx : i32) -> Option<usize> {
        for (line_idx_in_vec, line) in self.lines.iter().enumerate() {
            if line.line_idx == line_idx as i32 {
                return Some(line_idx_in_vec)
            }
        }
        None
    }

    ////////////////////////////////////////////////////////////////////////////

    #[inline]
    pub fn add_msg_str(&mut self, msg_str : &str, style : Style) {
        writeln!(&mut io::stderr(), "adding msg: {:?}", msg_str).unwrap();
        self.add_msg_str_(msg_str, style)
    }

    #[inline]
    pub fn add_msg(&mut self, msg : &[char], style : Style) {
        self.add_msg_(msg, style)
    }

    fn draw_(&self, rustbox : &RustBox, pos_x : i32, pos_y : i32) {
        let mut row = self.height - 1;
        let mut line_idx = min(self.scroll + self.height, self.lines.len() as i32) - 1;
        while line_idx >= 0 && row >= 0 {
            let line = &self.lines[line_idx as usize];

            if line.continuation {
                rustbox.print_char(pos_x as usize, (pos_y + row) as usize,
                                   line.style.style, line.style.fg, line.style.bg, LINEBREAK);
            }

            let pos_x = if line.continuation { pos_x + 1 } else { pos_x };

            rustbox.print(pos_x as usize, (pos_y + row) as usize,
                          line.style.style, line.style.fg, line.style.bg, &line.msg);

            row -= 1;
            line_idx -= 1;
        }
    }

    ////////////////////////////////////////////////////////////////////////////
    // Adding new messages

    fn add_msg_str_(&mut self, msg : &str, style : Style) {
        // Take the fast path when number of bytes (which gives the max number
        // of characters possibly be in the string) is smaller than the width.
        if msg.len() <= self.width as usize {
            if self.need_to_scroll() {
                self.scroll += 1;
            }

            writeln!(&mut io::stderr(), "msg_area adding msg: {}", msg).unwrap();
            self.lines.push(Line {
                msg: msg.to_owned(),
                continuation: false,
                line_idx: self.total_lines,
                style: style,
            });

            self.total_lines += 1;
        } else {
            // Need to split the lines, taking the slow path that uses an
            // intermediate vector.
            self.add_msg_(&msg.chars().collect::<Vec<char>>().borrow(), style);
        }
    }

    fn add_msg_(&mut self, mut msg : &[char], style : Style) {
        let mut lines : Vec<Line> = Vec::with_capacity(1);
        while msg.len() != 0 {
            let first_line = lines.len() == 0;
            let split = if first_line { self.width } else { self.width - 1 };

            let (line, rest) = msg.split_at(min(msg.len(), split as usize));
            msg = rest;

            lines.push(Line {
                msg: line.iter().cloned().collect(),
                continuation: !first_line,
                line_idx: self.total_lines,
                style: style,
            });
        }

        // We need to check the scroll before adding the messages, as the check
        // uses current scroll and number of messages.
        if self.need_to_scroll() {
            self.scroll += lines.len() as i32;
        }

        self.lines.append(&mut lines);

        self.total_lines += 1;
    }

    ////////////////////////////////////////////////////////////////////////////

    pub fn scroll_up(&mut self) {
        if self.scroll > 0 {
            self.scroll -= 1;
        }
    }

    pub fn scroll_down(&mut self) {
        if (self.lines.len() as i32) > self.scroll + self.height {
            self.scroll += 1;
        }
    }

    pub fn page_up(&mut self) {
        self.scroll = max(0, self.scroll - 10);
    }

    pub fn page_down(&mut self) {
        self.scroll = min(self.scroll + 10, (self.lines.len() as i32) - self.height);
    }

    /// Do we need to scroll when adding a new message?
    #[inline]
    fn need_to_scroll(&self) -> bool {
        self.scroll + self.height == self.lines.len() as i32
    }
}

impl Widget for MsgArea {
    fn draw(&self, rustbox : &RustBox, pos_x : i32, pos_y : i32) {
        self.draw_(rustbox, pos_x, pos_y)
    }

    fn keypressed(&mut self, _ : Key) -> WidgetRet {
        WidgetRet::KeyIgnored
    }

    fn resize(&mut self, width : i32, height : i32) {
        self.resize_(width, height)
    }
}
<commit_msg>remove some debug prints<commit_after>use std::borrow::Borrow;
use std::cmp::{max, min};
use std::io::Write;
use std::io;
use std::mem;

use rustbox::{RustBox, Key};
use tui::style::Style;
use tui::widget::{Widget, WidgetRet};

static LINEBREAK : char = '\\';

pub struct MsgArea {
    lines  : Vec<Line>,

    width  : i32,
    height : i32,

    /// Total number of lines added to the widget. Note that this is not the
    /// same thing as lines.len(): This simply returns number of lines added to
    /// the widget, rather than number of lines shown by the widget.
    total_lines : i32,

    /// Vertical scroll
    scroll : i32,

    // TODO: logging
}

#[derive(Debug)]
struct Line {
    /// A line. INVARIANT: Not longer than the width of the widget, and shown as
    /// single line in the TUI.
    msg : String,

    /// Is this continuation of previous line?
    continuation : bool,

    /// Index of the line in the buffer. Continuations have same indexes as the
    /// original message.
    line_idx : i32,

    style : Style,
}

impl MsgArea {
    pub fn new(width : i32, height : i32) -> MsgArea {
        MsgArea {
            lines: Vec::new(),
            width: width,
            height: height,
            total_lines: 0,
            scroll: 0,
        }
    }

    ////////////////////////////////////////////////////////////////////////////
    // Resizing

    fn resize_(&mut self, width : i32, height : i32) {
        // either scroll to bottom ...
        let scroll_to_bottom       = self.need_to_scroll();
        // ... or make sure the first visible line is still visible
        let first_visible_line_idx =
            if self.lines.is_empty() { 0 } else { self.lines[self.scroll as usize].line_idx };

        // Update the size
        self.width = width;
        self.height = height;

        // Combine/re-split lines
        self.resplit_lines();

        // Update the scroll
        if (self.lines.len() as i32) < self.height {
            self.scroll = 0;
        } else if scroll_to_bottom {
            self.scroll = (self.lines.len() as i32) - self.height;
        } else {
            match self.find_line_idx(first_visible_line_idx) {
                None => {/* TODO: Log this somewhere, this is a bug */},
                Some(idx) => { self.scroll = idx as i32; }
            }
        }
    }

    /// Combine continuations with original lines, add lines again in the
    /// original order. Should be called after updating width to have any
    /// effect.
    fn resplit_lines(&mut self) {
        // We could probably modify in-place using two indexes, but whatever.

        let old_lines : Vec<Line> = {
            let total_lines = self.lines.len();
            mem::replace(&mut self.lines, Vec::with_capacity(total_lines))
        };

        let mut line_idx = 0;
        while line_idx < old_lines.len() {
            // How many bytes combined string needs?
            let mut total_len = old_lines[line_idx].msg.len();
            {
                let mut cont_idx  = line_idx + 1;
                while cont_idx < old_lines.len() && old_lines[cont_idx].continuation {
                    total_len += old_lines[cont_idx].msg.len();
                    cont_idx  += 1;
                }
            }

            // Combine all lines
            let mut new_line = String::with_capacity(total_len);
            new_line.push_str(old_lines[line_idx].msg.borrow());
            let mut cont_idx = line_idx + 1;
            while cont_idx < old_lines.len() && old_lines[cont_idx].continuation {
                new_line.push_str(old_lines[cont_idx].msg.borrow());
                cont_idx += 1;
            }

            // Finally add the combined line to the fresh buffer
            self.add_msg_(new_line.chars().collect::<Vec<char>>().borrow(),
                          old_lines[line_idx].style);

            line_idx = cont_idx;
        }
    }

    /// Find index of given line in 'self.lines'.
    fn find_line_idx(&self, line_idx : i32) -> Option<usize> {
        for (line_idx_in_vec, line) in self.lines.iter().enumerate() {
            if line.line_idx == line_idx as i32 {
                return Some(line_idx_in_vec)
            }
        }
        None
    }

    ////////////////////////////////////////////////////////////////////////////

    #[inline]
    pub fn add_msg_str(&mut self, msg_str : &str, style : Style) {
        self.add_msg_str_(msg_str, style)
    }

    #[inline]
    pub fn add_msg(&mut self, msg : &[char], style : Style) {
        self.add_msg_(msg, style)
    }

    fn draw_(&self, rustbox : &RustBox, pos_x : i32, pos_y : i32) {
        let mut row = self.height - 1;
        let mut line_idx = min(self.scroll + self.height, self.lines.len() as i32) - 1;
        while line_idx >= 0 && row >= 0 {
            let line = &self.lines[line_idx as usize];

            if line.continuation {
                rustbox.print_char(pos_x as usize, (pos_y + row) as usize,
                                   line.style.style, line.style.fg, line.style.bg, LINEBREAK);
            }

            let pos_x = if line.continuation { pos_x + 1 } else { pos_x };

            rustbox.print(pos_x as usize, (pos_y + row) as usize,
                          line.style.style, line.style.fg, line.style.bg, &line.msg);

            row -= 1;
            line_idx -= 1;
        }
    }

    ////////////////////////////////////////////////////////////////////////////
    // Adding new messages

    fn add_msg_str_(&mut self, msg : &str, style : Style) {
        // Take the fast path when number of bytes (which gives the max number
        // of characters possibly be in the string) is smaller than the width.
        if msg.len() <= self.width as usize {
            if self.need_to_scroll() {
                self.scroll += 1;
            }

            self.lines.push(Line {
                msg: msg.to_owned(),
                continuation: false,
                line_idx: self.total_lines,
                style: style,
            });

            self.total_lines += 1;
        } else {
            // Need to split the lines, taking the slow path that uses an
            // intermediate vector.
            self.add_msg_(&msg.chars().collect::<Vec<char>>().borrow(), style);
        }
    }

    fn add_msg_(&mut self, mut msg : &[char], style : Style) {
        let mut lines : Vec<Line> = Vec::with_capacity(1);
        while msg.len() != 0 {
            let first_line = lines.len() == 0;
            let split = if first_line { self.width } else { self.width - 1 };

            let (line, rest) = msg.split_at(min(msg.len(), split as usize));
            msg = rest;

            lines.push(Line {
                msg: line.iter().cloned().collect(),
                continuation: !first_line,
                line_idx: self.total_lines,
                style: style,
            });
        }

        // We need to check the scroll before adding the messages, as the check
        // uses current scroll and number of messages.
        if self.need_to_scroll() {
            self.scroll += lines.len() as i32;
        }

        self.lines.append(&mut lines);

        self.total_lines += 1;
    }

    ////////////////////////////////////////////////////////////////////////////

    pub fn scroll_up(&mut self) {
        if self.scroll > 0 {
            self.scroll -= 1;
        }
    }

    pub fn scroll_down(&mut self) {
        if (self.lines.len() as i32) > self.scroll + self.height {
            self.scroll += 1;
        }
    }

    pub fn page_up(&mut self) {
        self.scroll = max(0, self.scroll - 10);
    }

    pub fn page_down(&mut self) {
        self.scroll = min(self.scroll + 10, (self.lines.len() as i32) - self.height);
    }

    /// Do we need to scroll when adding a new message?
    #[inline]
    fn need_to_scroll(&self) -> bool {
        self.scroll + self.height == self.lines.len() as i32
    }
}

impl Widget for MsgArea {
    fn draw(&self, rustbox : &RustBox, pos_x : i32, pos_y : i32) {
        self.draw_(rustbox, pos_x, pos_y)
    }

    fn keypressed(&mut self, _ : Key) -> WidgetRet {
        WidgetRet::KeyIgnored
    }

    fn resize(&mut self, width : i32, height : i32) {
        self.resize_(width, height)
    }
}
<|endoftext|>