<commit_before>#![allow(unstable)]

extern crate time;

use std::io::{File, SeekStyle, FileMode, FileAccess, USER_RWX, BufferedReader,
    Lines};
use std::io::fs::{unlink, PathExtensions, mkdir};
use std::os::{homedir, args};
use std::time::Duration;
use time::{now_utc, Tm, empty_tm, strptime};

fn main() {
    match args().get(1) {
        None => panic!("No command given"),
        Some(command) => {
            let mut time_clock = TimeClock::new();
            match &command[] {
                "in" => time_clock.punch_in(),
                "out" => time_clock.punch_out(),
                "status" => time_clock.status(),
                "report" => time_clock.report_daily_hours(),
                _ => panic!("unknown command")
            }
        }
    }
}

struct TimeClock {
    now: Tm,
    timesheet: File,
    currently_working: bool,
    state_path: Path
}

impl TimeClock {
    fn new() -> TimeClock {
        let now = now_utc();
        let home = homedir().unwrap();
        let base_dir = home.join(Path::new(".punch"));
        let timesheet_path = base_dir.join("timesheet");
        let working_state_path = base_dir.join("state");
        if !base_dir.exists() {
            mkdir(&base_dir, USER_RWX).unwrap();
        }
        let timesheet = File::open_mode(&timesheet_path,
                                            FileMode::Append,
                                            FileAccess::Write).unwrap();
        TimeClock {
            timesheet: timesheet,
            currently_working: working_state_path.exists(),
            state_path: working_state_path,
            now: now
        }
    }

    // commands

    fn punch_in(&mut self) {
        if self.currently_working {
            panic!("You're already working");
        }
        self.timesheet.seek(0, SeekStyle::SeekEnd).unwrap();
        writeln!(&mut self.timesheet, "in: {}", self.now.rfc822()).unwrap();
        self.set_current_working_state(true);
    }

    fn punch_out(&mut self) {
        if !self.currently_working {
            panic!("Can't punch out if you're not working!");
        }
        self.timesheet.seek(0, SeekStyle::SeekEnd).unwrap();
        writeln!(&mut self.timesheet, "out: {}", self.now.rfc822()).unwrap();
        self.set_current_working_state(false);
    }

    fn status(&self) {
        if self.currently_working {
            println!("You're punched in")
        } else {
            println!("You're punched out")
        }
    }

    fn report_daily_hours(&mut self) {
        self.timesheet.seek(0, SeekStyle::SeekSet).unwrap();
        let mut buf =
            BufferedReader::new(File::open(self.timesheet.path()).unwrap());
        let mut current_day = empty_tm();
        let mut time_worked_today = Duration::zero();

        let mut intervals = IntervalIter::from_lines(buf.lines());
        for (start, end) in intervals {
            if !same_day(&start, &current_day) {
                if !time_worked_today.is_zero() {
                    print_time_worked(&time_worked_today, &current_day);
                }
                current_day = start;
                time_worked_today = Duration::zero();
            }
            time_worked_today =
                time_worked_today + (end.to_timespec() - start.to_timespec());
        }

        if !time_worked_today.is_zero() {
            print_time_worked(&time_worked_today, &current_day);
        }
    }

    // aux. methods

    fn set_current_working_state(&mut self, currently_working: bool) {
        self.currently_working = currently_working;
        if currently_working {
            File::create(&self.state_path).unwrap();
        } else {
            unlink(&self.state_path).unwrap();
        }
    }
}

struct IntervalIter<'a> {
    lines: Lines<'a, BufferedReader<File>>
}

impl <'a> IntervalIter<'a> {
    fn from_lines(lines: Lines<'a, BufferedReader<File>>) -> IntervalIter<'a> {
        IntervalIter {lines: lines}
    }
}

impl <'a> Iterator for IntervalIter<'a> {
    type Item = (Tm, Tm);
    fn next(&mut self) -> Option<(Tm, Tm)> {
        match (self.lines.next(), self.lines.next()) {
            (None, None) => None,
            (Some(Ok(start_line)), o_end_line) => {
                if !start_line.starts_with("in: ") {
                    panic!("Bad data in timesheet!");
                }
                let start = parse_time(&start_line[4..]);
                let end = match o_end_line {
                    None => now_utc(),
                    Some(Ok(end_line)) => {
                        if !end_line.starts_with("out: ") {
                            panic!("Bad data in timesheet!");
                        }
                        parse_time(&end_line[5..])
                    },
                    Some(Err(_)) => panic!("Weird IO error")
                };
                Some((start, end))
            },
            _ => panic!("Weird IO error")
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (0, None)
    }
}

fn parse_time(s: &str) -> Tm {
    strptime(s.slice_to(s.len() - 1), "%a, %d %b %Y %T %Z").unwrap()
}

fn same_day(t1: &Tm, t2: &Tm) -> bool {
    t1.tm_year == t2.tm_year &&
    t1.tm_mon == t2.tm_mon &&
    t1.tm_mday == t2.tm_mday
}

fn print_time_worked(t: &Duration, day: &Tm) {
    println!("{}: {}:{}",
        day.strftime("%a, %d %b %Y").unwrap(),
        t.num_hours() ,
        t.num_minutes() % 60
    );
}
<commit_msg>Get rid of most possible panics in favour of proper error handling<commit_after>#![allow(unstable)]

extern crate time;

use std::io::{File, SeekStyle, FileMode, FileAccess, USER_RWX, BufferedReader,
    Lines};
use std::io::fs::{unlink, PathExtensions, mkdir};
use std::os::{homedir, args, set_exit_status};
use std::fmt;
use std::time::Duration;
use time::{now_utc, Tm, empty_tm, strptime};

fn main() {
    let result = match args().get(1) {
        None => Err(PunchClockError::NoCommandGiven),
        Some(command) => {
            let mut time_clock = TimeClock::new();
            match &command[] {
                "in" => time_clock.punch_in(),
                "out" => time_clock.punch_out(),
                "status" => time_clock.status(),
                "report" => time_clock.report_daily_hours(),
                _ => Err(PunchClockError::UnknownCommand)
            }
        }
    };

    if let Err(e) = result {
        println!("Error: {}", e);
        set_exit_status(1);
    }
}

enum PunchClockError {
    NoCommandGiven,
    UnknownCommand,
    AlreadyPunchedIn,
    AlreadyPunchedOut,
    CorruptedTimeSheet,
    IoError(std::io::IoError),
}

impl fmt::String for PunchClockError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        use PunchClockError::*;
        fmt.write_str(
            match *self {
                NoCommandGiven => "No command given",
                UnknownCommand => "Unknown command",
                AlreadyPunchedIn => "You are already punched in",
                AlreadyPunchedOut => "You're not currently punched in",
                CorruptedTimeSheet => "Bad data in timesheet",
                IoError(_) => "IO error"
            }
        )
    }
}

type PunchClockResult<T> = Result<T, PunchClockError>;

struct TimeClock {
    now: Tm,
    timesheet: File,
    currently_working: bool,
    state_path: Path
}

impl TimeClock {
    fn new() -> TimeClock {
        let now = now_utc();
        let home = homedir().unwrap();
        let base_dir = home.join(Path::new(".punch"));
        let timesheet_path = base_dir.join("timesheet");
        let working_state_path = base_dir.join("state");
        if !base_dir.exists() {
            mkdir(&base_dir, USER_RWX).unwrap();
        }
        let timesheet = File::open_mode(&timesheet_path,
                                            FileMode::Append,
                                            FileAccess::Write).unwrap();
        TimeClock {
            timesheet: timesheet,
            currently_working: working_state_path.exists(),
            state_path: working_state_path,
            now: now
        }
    }

    // commands

    fn punch_in(&mut self) -> PunchClockResult<()> {
        if self.currently_working {
            return Err(PunchClockError::AlreadyPunchedIn);
        }
        self.timesheet.seek(0, SeekStyle::SeekEnd).unwrap();
        writeln!(&mut self.timesheet, "in: {}", self.now.rfc822()).unwrap();
        self.set_current_working_state(true);
        Ok(())
    }

    fn punch_out(&mut self) -> PunchClockResult<()> {
        if !self.currently_working {
            return Err(PunchClockError::AlreadyPunchedOut);
        }
        self.timesheet.seek(0, SeekStyle::SeekEnd).unwrap();
        writeln!(&mut self.timesheet, "out: {}", self.now.rfc822()).unwrap();
        self.set_current_working_state(false);
        Ok(())
    }

    fn status(&self) -> PunchClockResult<()> {
        if self.currently_working {
            println!("You're punched in");
        } else {
            println!("You're punched out");
        }
        Ok(())
    }

    fn report_daily_hours(&mut self) -> PunchClockResult<()> {
        self.timesheet.seek(0, SeekStyle::SeekSet).unwrap();
        let mut buf =
            BufferedReader::new(File::open(self.timesheet.path()).unwrap());
        let mut current_day = empty_tm();
        let mut time_worked_today = Duration::zero();

        let mut intervals = IntervalIter::from_lines(buf.lines());
        for interval in intervals {
            let (start, end) = try!(interval);
            if !same_day(&start, &current_day) {
                if !time_worked_today.is_zero() {
                    print_time_worked(&time_worked_today, &current_day);
                }
                current_day = start;
                time_worked_today = Duration::zero();
            }
            time_worked_today =
                time_worked_today + (end.to_timespec() - start.to_timespec());
        }

        if !time_worked_today.is_zero() {
            print_time_worked(&time_worked_today, &current_day);
        }
        Ok(())
    }

    // aux. methods

    fn set_current_working_state(&mut self, currently_working: bool) {
        self.currently_working = currently_working;
        if currently_working {
            File::create(&self.state_path).unwrap();
        } else {
            unlink(&self.state_path).unwrap();
        }
    }
}

struct IntervalIter<'a> {
    lines: Lines<'a, BufferedReader<File>>
}

impl <'a> IntervalIter<'a> {
    fn from_lines(lines: Lines<'a, BufferedReader<File>>) -> IntervalIter<'a> {
        IntervalIter {lines: lines}
    }
}

impl <'a> Iterator for IntervalIter<'a> {
    type Item = PunchClockResult<(Tm, Tm)>;
    fn next(&mut self) -> Option<PunchClockResult<(Tm, Tm)>> {

        // helper function to make error handling a bit nicer
        fn inner_unwrap<T>(x: Option<std::io::IoResult<T>>)
                -> PunchClockResult<Option<T>> {
            match x {
                None => Ok(None),
                Some(Ok(inner)) => Ok(Some(inner)),
                Some(Err(e)) => Err(PunchClockError::IoError(e))
            }
        }

        let line_1 = match inner_unwrap(self.lines.next()) {
            Ok(l) => l,
            Err(e) => return Some(Err(e))
        };
        let line_2 = match inner_unwrap(self.lines.next()) {
            Ok(l) => l,
            Err(e) => return Some(Err(e))
        };

        match (line_1, line_2) {
            (None, None) => None,
            (Some(start_line), o_end_line) => {
                if !start_line.starts_with("in: ") {
                    return Some(Err(PunchClockError::CorruptedTimeSheet));
                }
                let start = parse_time(&start_line[4..]);
                let end = match o_end_line {
                    None => now_utc(),
                    Some(end_line) => {
                        if !end_line.starts_with("out: ") {
                            return Some(Err(PunchClockError::CorruptedTimeSheet));
                        }
                        parse_time(&end_line[5..])
                    },
                };
                Some(Ok((start, end)))
            },
            _ => unreachable!() // (None, Some(l)) should not happen
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (0, None)
    }
}

fn parse_time(s: &str) -> Tm {
    strptime(s.slice_to(s.len() - 1), "%a, %d %b %Y %T %Z").unwrap()
}

fn same_day(t1: &Tm, t2: &Tm) -> bool {
    t1.tm_year == t2.tm_year &&
    t1.tm_mon == t2.tm_mon &&
    t1.tm_mday == t2.tm_mday
}

fn print_time_worked(t: &Duration, day: &Tm) {
    println!("{}: {}:{}",
        day.strftime("%a, %d %b %Y").unwrap(),
        t.num_hours() ,
        t.num_minutes() % 60
    );
}
<|endoftext|>


<commit_before>extern crate rand;
extern crate time;
extern crate docopt;
extern crate itertools;
extern crate rustc_serialize;
// piston stuff
extern crate piston;
extern crate graphics;
extern crate glutin_window;
extern crate opengl_graphics;

mod piston_app;
mod console_app;
mod world;
mod parser;
mod util;

use docopt::Docopt;
use world::World;

static USAGE: &'static str = "
Usage:
    golr [options]

Options:
    --help                                      Show this message
    -h <height>, --height <height>              World height, points                         [default: 25]
    -w <width>, --width <width>                 World width, points                          [default: 80]
    -p <period>, --period <period>              World generational change period, ms         [default: 350]
    -i <inputfile>, --inputfile <inputfile>     Populate world from file
    -r <render>, --render <render>              Choose render engine [0-console | 1-piston]  [default: 0]
    --cell-size <cell_size>                     Choose cell size (Only for visual renderers) [default: 2]
    --gl-version <gl_version>                   OpenGl version [32, 30]                      [default: 32]
";

#[derive(RustcDecodable, Debug)]
struct CliArgs {
    flag_width: isize,
    flag_height: isize,
    flag_period: u64,
    flag_inputfile: Option<String>,
    flag_render: isize,
    flag_cell_size: u32,
    flag_gl_version: u32,
}

fn main() {
    let args = Docopt::new(USAGE)
        .and_then(|d| d.decode::<CliArgs>())
        .unwrap_or_else(|e| e.exit());

    let seed = args.flag_inputfile.as_ref()
                   .and_then(|fif| util::read_file_to_string(fif)
                                  .and_then(|d: String| Ok(Some(parser::parse_plaintext(d))))
                                  .unwrap_or(None));

    let mut world = World::new(args.flag_width, args.flag_height).seed(seed);

    match args.flag_render {
        0 => console_app::start_console_app(&mut world, args.flag_period),
        1 => piston_app::start_piston_app(world, args.flag_period, args.flag_cell_size, args.flag_gl_version),
        _ => panic!("Unknown render engine")
    };
}
<commit_msg>Remove Result wrapping from seed<commit_after>extern crate rand;
extern crate time;
extern crate docopt;
extern crate itertools;
extern crate rustc_serialize;
// piston stuff
extern crate piston;
extern crate graphics;
extern crate glutin_window;
extern crate opengl_graphics;

mod piston_app;
mod console_app;
mod world;
mod parser;
mod util;

use docopt::Docopt;
use world::World;

static USAGE: &'static str = "
Usage:
    golr [options]

Options:
    --help                                      Show this message
    -h <height>, --height <height>              World height, points                         [default: 25]
    -w <width>, --width <width>                 World width, points                          [default: 80]
    -p <period>, --period <period>              World generational change period, ms         [default: 350]
    -i <inputfile>, --inputfile <inputfile>     Populate world from file
    -r <render>, --render <render>              Choose render engine [0-console | 1-piston]  [default: 0]
    --cell-size <cell_size>                     Choose cell size (Only for visual renderers) [default: 2]
    --gl-version <gl_version>                   OpenGl version [32, 30]                      [default: 32]
";

#[derive(RustcDecodable, Debug)]
struct CliArgs {
    flag_width: isize,
    flag_height: isize,
    flag_period: u64,
    flag_inputfile: Option<String>,
    flag_render: isize,
    flag_cell_size: u32,
    flag_gl_version: u32,
}

fn main() {
    let args = Docopt::new(USAGE)
        .and_then(|d| d.decode::<CliArgs>())
        .unwrap_or_else(|e| e.exit());

    let seed = args.flag_inputfile.as_ref()
                   .and_then(|fif| util::read_file_to_string(fif).ok()
                                   .and_then(|d| Some(parser::parse_plaintext(d))));

    let mut world = World::new(args.flag_width, args.flag_height).seed(seed);

    match args.flag_render {
        0 => console_app::start_console_app(&mut world, args.flag_period),
        1 => piston_app::start_piston_app(world, args.flag_period, args.flag_cell_size, args.flag_gl_version),
        _ => panic!("Unknown render engine")
    };
}
<|endoftext|>


<commit_before>use std::iter::Peekable;
use std::str::Chars;

use {Error, Result};

/// The data attribute of a path.
///
/// http://www.w3.org/TR/SVG/paths.html#PathData
pub struct Data {
    commands: Vec<Command>,
}

/// A command used to draw a path.
#[derive(Debug)]
pub enum Command {
    /// [Establish][1] a new current point.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataMovetoCommands
    MoveTo(Positioning, Vec<f64>),

    /// [End][1] the current subpath.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataClosePathCommand
    ClosePath,

    /// [Draw][1] straight lines.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataLinetoCommands
    LineTo(Positioning, Vec<f64>),

    /// [Draw][1] horizontal lines.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataLinetoCommands
    HorizontalLineTo(Positioning, Vec<f64>),

    /// [Draw][1] vertical lines.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataLinetoCommands
    VerticalLineTo(Positioning, Vec<f64>),

    /// [Draw][1] a cubic Bézier curve.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataCubicBezierCommands
    CurveTo(Positioning, Vec<f64>),

    /// [Draw][1] a cubic Bézier curve assuming the first control point to be
    /// the reflection of the second control point on the previous command
    /// relative to the current point.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataCubicBezierCommands
    SmoothCurveTo(Positioning, Vec<f64>),

    /// [Draw][1] a quadratic Bézier curve.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataQuadraticBezierCommands
    QuadraticBezierCurveTo(Positioning, Vec<f64>),

    /// [Draw][1] a quadratic Bézier curve assuming the control point to be the
    /// reflection of the control point on the previous command relative to the
    /// current point.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataQuadraticBezierCommands
    SmoothQuadraticBezierCurveTo(Positioning, Vec<f64>),

    /// [Draw][1] an elliptical arc.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataEllipticalArcCommands
    EllipticalArc(Positioning, Vec<f64>),
}

/// An attribute of a command indicating whether the coordinates of the command
/// are absolute or relative to the current position.
#[derive(Debug)]
pub enum Positioning {
    Absolute,
    Relative,
}

impl Data {
    /// Parse the data attribute of a path.
    #[inline]
    pub fn parse(line: &str) -> Result<Data> {
        DataParser::new(line).process()
    }

    #[inline]
    pub fn iter(&self) -> ::std::slice::Iter<Command> {
        self.commands.iter()
    }
}

struct DataParser<'s> {
    line: &'s str,
    position: usize,
    cursor: Peekable<Chars<'s>>,
}

macro_rules! ok(
    ($result:expr) => (
        match $result {
            Ok(result) => result,
            Err(error) => return Err(error),
        }
    )
);

macro_rules! raise(
    ($parser:expr, $($arg:tt)*) => (
        return Err(Error {
            line: 1,
            column: $parser.position,
            message: format!($($arg)*),
        })
    );
);

impl<'s> DataParser<'s> {
    fn new(line: &'s str) -> DataParser<'s> {
        DataParser {
            line: line,
            position: 0,
            cursor: line.chars().peekable(),
        }
    }

    fn process(&mut self) -> Result<Data> {
        let mut commands = Vec::new();

        loop {
            self.skip_whitespace();

            match ok!(self.read_command()) {
                Some(command) => commands.push(command),
                _ => break,
            }
        }

        Ok(Data {
            commands: commands,
        })
    }

    fn read_command(&mut self) -> Result<Option<Command>> {
        use self::Command::*;
        use self::Positioning::*;

        let name = match self.next() {
            Some(name) => match name {
                'A'...'Z' | 'a'...'z' => name,
                _ => raise!(self, "expected a path command"),
            },
            _ => return Ok(None),
        };

        self.skip_whitespace();

        let params = ok!(self.read_parameters());

        Ok(Some(match name {
            'M' => MoveTo(Absolute, params),
            'm' => MoveTo(Relative, params),

            'Z' | 'z' => ClosePath,

            'L' => LineTo(Absolute, params),
            'l' => LineTo(Relative, params),

            'H' => HorizontalLineTo(Absolute, params),
            'h' => HorizontalLineTo(Relative, params),

            'V' => VerticalLineTo(Absolute, params),
            'v' => VerticalLineTo(Relative, params),

            'C' => CurveTo(Absolute, params),
            'c' => CurveTo(Relative, params),

            'S' => SmoothCurveTo(Absolute, params),
            's' => SmoothCurveTo(Relative, params),

            'Q' => QuadraticBezierCurveTo(Absolute, params),
            'q' => QuadraticBezierCurveTo(Relative, params),

            'T' => SmoothQuadraticBezierCurveTo(Absolute, params),
            't' => SmoothQuadraticBezierCurveTo(Relative, params),

            'A' => EllipticalArc(Absolute, params),
            'a' => EllipticalArc(Relative, params),

            _ => raise!(self, "found an unknown path command '{}'", name),
        }))
    }

    fn read_parameters(&mut self) -> Result<Vec<f64>> {
        let mut params = Vec::new();

        loop {
            match ok!(self.read_number()) {
                Some(number) => params.push(number),
                _ => break,
            }

            self.skip_whitespace();
            self.skip(",");
        }

        Ok(params)
    }

    fn read_number(&mut self) -> Result<Option<f64>> {
        self.skip_whitespace();

        let start = self.position;

        self.skip("-");
        self.skip_digits();
        self.skip(".");
        self.skip_digits();

        let end = self.position;

        if start == end {
            return Ok(None)
        }

        let number = &self.line[start..end];

        match number.parse() {
            Ok(number) => Ok(Some(number)),
            Err(_) => raise!(self, "failed to parse a number '{}'", number),
        }
    }

    #[inline]
    fn skip_whitespace(&mut self) {
        self.skip(" \t")
    }

    #[inline]
    fn skip_digits(&mut self) {
        self.skip("0123456789")
    }

    fn skip(&mut self, chars: &str) {
        loop {
            match self.peek() {
                Some(c) => {
                    if chars.contains_char(c) {
                        self.next();
                    } else {
                        break;
                    }
                },
                _ => break,
            }
        }
    }

    #[inline]
    fn next(&mut self) -> Option<char> {
        match self.cursor.next() {
            Some(c) => {
                self.position += 1;
                Some(c)
            },
            _ => None,
        }
    }

    #[inline]
    fn peek(&mut self) -> Option<char> {
        self.cursor.peek().and_then(|&c| Some(c))
    }
}

#[cfg(test)]
mod tests {
    use super::{Data, DataParser};
    use super::Command::*;
    use super::Positioning::*;

    #[test]
    fn data_parser() {
        let data = Data::parse("M1,2 l3,4").ok().unwrap();

        assert_eq!(data.commands.len(), 2);

        match data.commands[0] {
            MoveTo(Absolute, ref params) => assert_eq!(*params, vec![1.0, 2.0]),
            _ => assert!(false),
        }

        match data.commands[1] {
            LineTo(Relative, ref params) => assert_eq!(*params, vec![3.0, 4.0]),
            _ => assert!(false),
        }
    }

    #[test]
    fn data_parser_read_command() {
        macro_rules! run(
            ($line:expr) => ({
                let mut parser = DataParser::new($line);
                parser.read_command().ok().unwrap().unwrap()
            });
        );

        macro_rules! test(
            ($line:expr, $command:ident, $positioning:ident, $params:expr) => (
                match run!($line) {
                    $command($positioning, params) => assert_eq!(params, $params),
                    _ => assert!(false),
                }
            );
            ($line:expr, $command:ident) => (
                match run!($line) {
                    $command => {},
                    _ => assert!(false),
                }
            );
        );

        test!("M4,2", MoveTo, Absolute, vec![4.0, 2.0]);
        test!("m4,2", MoveTo, Relative, vec![4.0, 2.0]);

        test!("Z", ClosePath);
        test!("z", ClosePath);

        test!("L7, 8  9", LineTo, Absolute, vec![7.0, 8.0, 9.0]);
        test!("l 7,8 9", LineTo, Relative, vec![7.0, 8.0, 9.0]);

        test!("H\t6,9", HorizontalLineTo, Absolute, vec![6.0, 9.0]);
        test!("h6,  \t9", HorizontalLineTo, Relative, vec![6.0, 9.0]);

        test!("V2.1,-3", VerticalLineTo, Absolute, vec![2.1, -3.0]);
        test!("v2.1 -3", VerticalLineTo, Relative, vec![2.1, -3.0]);

        test!("C0,1 0,2", CurveTo, Absolute, vec![0.0, 1.0, 0.0, 2.0]);
        test!("c0 ,1 0,  2", CurveTo, Relative, vec![0.0, 1.0, 0.0, 2.0]);

        test!("S42,0", SmoothCurveTo, Absolute, vec![42.0, 0.0]);
        test!("s \t 42,0", SmoothCurveTo, Relative, vec![42.0, 0.0]);

        test!("Q90.5 0", QuadraticBezierCurveTo, Absolute, vec![90.5, 0.0]);
        test!("q90.5, 0", QuadraticBezierCurveTo, Relative, vec![90.5, 0.0]);

        test!("T-1", SmoothQuadraticBezierCurveTo, Absolute, vec![-1.0]);
        test!("t -1", SmoothQuadraticBezierCurveTo, Relative, vec![-1.0]);

        test!("A2.6,0 -7", EllipticalArc, Absolute, vec![2.6, 0.0, -7.0]);
        test!("a 2.6 ,0 -7", EllipticalArc, Relative, vec![2.6, 0.0, -7.0]);
    }

    #[test]
    fn data_parser_read_parameters() {
        let mut parser = DataParser::new("1,2 3,4 5 6.7");
        let params = parser.read_parameters().ok().unwrap();
        assert_eq!(params, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.7]);
    }

    #[test]
    fn data_parser_read_number() {
        let lines = vec!["-1", "3", "3.14"];
        let numbers = vec![-1.0, 3.0, 3.14];

        for (line, &number) in lines.iter().zip(numbers.iter()) {
            let mut parser = DataParser::new(line);
            assert_eq!(parser.read_number().ok().unwrap().unwrap(), number);
        }
    }

    #[test]
    fn data_parser_skip_whitespace() {
        let mut parser = DataParser::new(" \t    m ");
        parser.skip_whitespace();
        assert_eq!(parser.position, 6);
    }
}
<commit_msg>Add the possibility to have new lines in path data<commit_after>use std::iter::Peekable;
use std::str::Chars;

use {Error, Result};

/// The data attribute of a path.
///
/// http://www.w3.org/TR/SVG/paths.html#PathData
pub struct Data {
    commands: Vec<Command>,
}

/// A command used to draw a path.
#[derive(Debug)]
pub enum Command {
    /// [Establish][1] a new current point.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataMovetoCommands
    MoveTo(Positioning, Vec<f64>),

    /// [End][1] the current subpath.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataClosePathCommand
    ClosePath,

    /// [Draw][1] straight lines.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataLinetoCommands
    LineTo(Positioning, Vec<f64>),

    /// [Draw][1] horizontal lines.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataLinetoCommands
    HorizontalLineTo(Positioning, Vec<f64>),

    /// [Draw][1] vertical lines.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataLinetoCommands
    VerticalLineTo(Positioning, Vec<f64>),

    /// [Draw][1] a cubic Bézier curve.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataCubicBezierCommands
    CurveTo(Positioning, Vec<f64>),

    /// [Draw][1] a cubic Bézier curve assuming the first control point to be
    /// the reflection of the second control point on the previous command
    /// relative to the current point.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataCubicBezierCommands
    SmoothCurveTo(Positioning, Vec<f64>),

    /// [Draw][1] a quadratic Bézier curve.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataQuadraticBezierCommands
    QuadraticBezierCurveTo(Positioning, Vec<f64>),

    /// [Draw][1] a quadratic Bézier curve assuming the control point to be the
    /// reflection of the control point on the previous command relative to the
    /// current point.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataQuadraticBezierCommands
    SmoothQuadraticBezierCurveTo(Positioning, Vec<f64>),

    /// [Draw][1] an elliptical arc.
    ///
    /// [1]: http://www.w3.org/TR/SVG/paths.html#PathDataEllipticalArcCommands
    EllipticalArc(Positioning, Vec<f64>),
}

/// An attribute of a command indicating whether the coordinates of the command
/// are absolute or relative to the current position.
#[derive(Debug)]
pub enum Positioning {
    Absolute,
    Relative,
}

impl Data {
    /// Parse the data attribute of a path.
    #[inline]
    pub fn parse(text: &str) -> Result<Data> {
        DataParser::new(text).process()
    }

    #[inline]
    pub fn iter(&self) -> ::std::slice::Iter<Command> {
        self.commands.iter()
    }
}

struct DataParser<'s> {
    text: &'s str,

    line: usize,
    column: usize,
    offset: usize,

    cursor: Peekable<Chars<'s>>,
}

macro_rules! ok(
    ($result:expr) => (
        match $result {
            Ok(result) => result,
            Err(error) => return Err(error),
        }
    )
);

macro_rules! raise(
    ($parser:expr, $($arg:tt)*) => (
        return Err(Error {
            line: $parser.line,
            column: $parser.column,
            message: format!($($arg)*),
        })
    );
);

impl<'s> DataParser<'s> {
    fn new(text: &'s str) -> DataParser<'s> {
        DataParser {
            text: text,

            line: 1,
            column: 0,
            offset: 0,

            cursor: text.chars().peekable(),
        }
    }

    fn process(&mut self) -> Result<Data> {
        let mut commands = Vec::new();

        loop {
            self.skip_whitespace();

            match ok!(self.read_command()) {
                Some(command) => commands.push(command),
                _ => break,
            }
        }

        Ok(Data {
            commands: commands,
        })
    }

    fn read_command(&mut self) -> Result<Option<Command>> {
        use self::Command::*;
        use self::Positioning::*;

        let name = match self.next() {
            Some(name) => match name {
                'A'...'Z' | 'a'...'z' => name,
                _ => raise!(self, "expected a path command"),
            },
            _ => return Ok(None),
        };

        self.skip_whitespace();

        let params = ok!(self.read_parameters());

        Ok(Some(match name {
            'M' => MoveTo(Absolute, params),
            'm' => MoveTo(Relative, params),

            'Z' | 'z' => ClosePath,

            'L' => LineTo(Absolute, params),
            'l' => LineTo(Relative, params),

            'H' => HorizontalLineTo(Absolute, params),
            'h' => HorizontalLineTo(Relative, params),

            'V' => VerticalLineTo(Absolute, params),
            'v' => VerticalLineTo(Relative, params),

            'C' => CurveTo(Absolute, params),
            'c' => CurveTo(Relative, params),

            'S' => SmoothCurveTo(Absolute, params),
            's' => SmoothCurveTo(Relative, params),

            'Q' => QuadraticBezierCurveTo(Absolute, params),
            'q' => QuadraticBezierCurveTo(Relative, params),

            'T' => SmoothQuadraticBezierCurveTo(Absolute, params),
            't' => SmoothQuadraticBezierCurveTo(Relative, params),

            'A' => EllipticalArc(Absolute, params),
            'a' => EllipticalArc(Relative, params),

            _ => raise!(self, "found an unknown path command '{}'", name),
        }))
    }

    fn read_parameters(&mut self) -> Result<Vec<f64>> {
        let mut params = Vec::new();

        loop {
            match ok!(self.read_number()) {
                Some(number) => params.push(number),
                _ => break,
            }

            self.skip_whitespace();
            self.skip(",");
        }

        Ok(params)
    }

    fn read_number(&mut self) -> Result<Option<f64>> {
        self.skip_whitespace();

        let start = self.offset;

        self.skip("-");
        self.skip_digits();
        self.skip(".");
        self.skip_digits();

        let end = self.offset;

        if start == end {
            return Ok(None)
        }

        let number = &self.text[start..end];

        match number.parse() {
            Ok(number) => Ok(Some(number)),
            Err(_) => raise!(self, "failed to parse a number '{}'", number),
        }
    }

    #[inline]
    fn skip_whitespace(&mut self) {
        self.skip(" \t\n")
    }

    #[inline]
    fn skip_digits(&mut self) {
        self.skip("0123456789")
    }

    fn skip(&mut self, chars: &str) {
        loop {
            match self.peek() {
                Some(c) => {
                    if chars.contains_char(c) {
                        self.next();
                    } else {
                        break;
                    }
                },
                _ => break,
            }
        }
    }

    #[inline]
    fn next(&mut self) -> Option<char> {
        match self.cursor.next() {
            Some(c) => {
                if c == '\n' {
                    self.line += 1;
                    self.column = 0;
                }
                self.column += 1;
                self.offset += 1;
                Some(c)
            }
            _ => None,
        }
    }

    #[inline]
    fn peek(&mut self) -> Option<char> {
        self.cursor.peek().and_then(|&c| Some(c))
    }
}

#[cfg(test)]
mod tests {
    use super::{Data, DataParser};
    use super::Command::*;
    use super::Positioning::*;

    #[test]
    fn data_parser() {
        let data = Data::parse("M1,2 l3,4").ok().unwrap();

        assert_eq!(data.commands.len(), 2);

        match data.commands[0] {
            MoveTo(Absolute, ref params) => assert_eq!(*params, vec![1.0, 2.0]),
            _ => assert!(false),
        }

        match data.commands[1] {
            LineTo(Relative, ref params) => assert_eq!(*params, vec![3.0, 4.0]),
            _ => assert!(false),
        }
    }

    #[test]
    fn data_parser_read_command() {
        macro_rules! run(
            ($text:expr) => ({
                let mut parser = DataParser::new($text);
                parser.read_command().ok().unwrap().unwrap()
            });
        );

        macro_rules! test(
            ($text:expr, $command:ident, $positioning:ident, $params:expr) => (
                match run!($text) {
                    $command($positioning, params) => assert_eq!(params, $params),
                    _ => assert!(false),
                }
            );
            ($text:expr, $command:ident) => (
                match run!($text) {
                    $command => {},
                    _ => assert!(false),
                }
            );
        );

        test!("M4,2", MoveTo, Absolute, vec![4.0, 2.0]);
        test!("m4,\n2", MoveTo, Relative, vec![4.0, 2.0]);

        test!("Z", ClosePath);
        test!("z", ClosePath);

        test!("L7, 8  9", LineTo, Absolute, vec![7.0, 8.0, 9.0]);
        test!("l 7,8 \n9", LineTo, Relative, vec![7.0, 8.0, 9.0]);

        test!("H\t6,9", HorizontalLineTo, Absolute, vec![6.0, 9.0]);
        test!("h6,  \t9", HorizontalLineTo, Relative, vec![6.0, 9.0]);

        test!("V2.1,-3", VerticalLineTo, Absolute, vec![2.1, -3.0]);
        test!("v\n2.1 -3", VerticalLineTo, Relative, vec![2.1, -3.0]);

        test!("C0,1 0,2", CurveTo, Absolute, vec![0.0, 1.0, 0.0, 2.0]);
        test!("c0 ,1 0,  2", CurveTo, Relative, vec![0.0, 1.0, 0.0, 2.0]);

        test!("S42,0", SmoothCurveTo, Absolute, vec![42.0, 0.0]);
        test!("s \t 42,0", SmoothCurveTo, Relative, vec![42.0, 0.0]);

        test!("Q90.5 0", QuadraticBezierCurveTo, Absolute, vec![90.5, 0.0]);
        test!("q90.5\n, 0", QuadraticBezierCurveTo, Relative, vec![90.5, 0.0]);

        test!("T-1", SmoothQuadraticBezierCurveTo, Absolute, vec![-1.0]);
        test!("t -1", SmoothQuadraticBezierCurveTo, Relative, vec![-1.0]);

        test!("A2.6,0 -7", EllipticalArc, Absolute, vec![2.6, 0.0, -7.0]);
        test!("a 2.6 ,0 -7", EllipticalArc, Relative, vec![2.6, 0.0, -7.0]);
    }

    #[test]
    fn data_parser_read_parameters() {
        let mut parser = DataParser::new("1,2 3,4 5 6.7");
        let params = parser.read_parameters().ok().unwrap();
        assert_eq!(params, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.7]);
    }

    #[test]
    fn data_parser_read_number() {
        let texts = vec!["-1", "3", "3.14"];
        let numbers = vec![-1.0, 3.0, 3.14];

        for (text, &number) in texts.iter().zip(numbers.iter()) {
            let mut parser = DataParser::new(text);
            assert_eq!(parser.read_number().ok().unwrap().unwrap(), number);
        }
    }

    #[test]
    fn data_parser_skip_whitespace() {
        let mut parser = DataParser::new(" \t  \n\n  \tm ");
        parser.skip_whitespace();
        assert_eq!(parser.line, 3);
        assert_eq!(parser.column, 4);
        assert_eq!(parser.offset, 9);
    }
}
<|endoftext|>


<commit_before>use std::rc::Rc;
use std::cell::RefCell;
use std::io;

use byteorder::{BigEndian, WriteBytesExt};

use mio;

use EventLoop;
use protocol::Protocol as Protocol;
use transport::Connection as Connection;

enum PipeStateIdx {
	Handshake,
	Ready
}

pub struct Pipe {
	state: PipeStateIdx,
	handshake_state: HandshakePipeState,
	ready_state: ReadyPipeState
}

impl Pipe {

	pub fn new(id: usize, protocol: Rc<Box<Protocol>>, connection: Box<Connection>) -> Pipe {
		let conn_ref = Rc::new(RefCell::new(connection));

		Pipe {
			state: PipeStateIdx::Handshake,
			handshake_state: HandshakePipeState::new(id, protocol.clone(), conn_ref.clone()),
			ready_state: ReadyPipeState::new(id, protocol.clone(), conn_ref.clone())
		}
	}

	fn get_state<'a>(&'a mut self) -> &'a mut PipeState {
		match self.state {
			PipeStateIdx::Handshake => &mut self.handshake_state,
			PipeStateIdx::Ready => &mut self.ready_state
		}
	}

	fn set_state(&mut self, state: PipeStateIdx, event_loop: &mut EventLoop) -> io::Result<()> {
		self.state = state;
		self.get_state().enter(event_loop)
	}

	pub fn init(&mut self, event_loop: &mut EventLoop) -> io::Result<()> {
		self.set_state(PipeStateIdx::Handshake, event_loop)
	}

	pub fn readable(&mut self, event_loop: &mut EventLoop, hint: mio::ReadHint)	-> io::Result<()> {
		if let Some(new_state) = try!(self.get_state().readable(event_loop)) {
			try!(self.set_state(new_state, event_loop));
		}

		Ok(())
	}

	pub fn writable(&mut self, event_loop: &mut EventLoop) -> io::Result<()> {
		if let Some(new_state) = try!(self.get_state().writable(event_loop)) {
			try!(self.set_state(new_state, event_loop));
		}

		Ok(())
	}

}

trait PipeState {
	fn enter(&mut self, event_loop: &mut EventLoop) -> io::Result<()>;
	fn readable(&mut self, event_loop: &mut EventLoop) -> io::Result<Option<PipeStateIdx>>;
	fn writable(&mut self, event_loop: &mut EventLoop) -> io::Result<Option<PipeStateIdx>>;
}

struct HandshakePipeState {
	id: usize,
	protocol: Rc<Box<Protocol>>,
	connection: Rc<RefCell<Box<Connection>>>,
	sent: bool,
	received: bool
}

impl HandshakePipeState {
	fn new(id: usize, protocol: Rc<Box<Protocol>>, connection: Rc<RefCell<Box<Connection>>>) -> HandshakePipeState {
		HandshakePipeState {
			id: id,
			protocol: protocol,
			connection: connection,
			sent: false,
			received: false }
	}

	fn check_received_handshake(&self, handshake: &[u8; 8]) -> io::Result<()> {
		let mut expected_handshake = vec!(0, 83, 80, 0);
		let protocol_id = self.protocol.peer_id();
		try!(expected_handshake.write_u16::<BigEndian>(protocol_id));
		try!(expected_handshake.write_u16::<BigEndian>(0));
		let mut both = handshake.iter().zip(expected_handshake.iter());

		if both.all(|(l,r)| l == r) {
			Ok(())
		} else {
			Err(io::Error::new(io::ErrorKind::InvalidData, "received wrong handshake"))
		}
	}

	fn check_sent_handshake(&self, written: Option<usize>) -> io::Result<()> {
		match written {
			Some(8) => Ok(()),
			_ => Err(io::Error::new(io::ErrorKind::InvalidData, "failed to send handshake"))
		}
	}

	fn register(&self, event_loop: &mut EventLoop, interest: mio::Interest, poll_opt: mio::PollOpt) -> io::Result<()> {
		let token = mio::Token(self.id);
		let connection = self.connection.borrow_mut();
		let fd = connection.as_evented();

		event_loop.register_opt(fd, token, interest, poll_opt)
	}

	fn reregister(&self, event_loop: &mut EventLoop, interest: mio::Interest, poll_opt: mio::PollOpt) -> io::Result<()> {
		let token = mio::Token(self.id);
		let connection = self.connection.borrow_mut();
		let fd = connection.as_evented();

		event_loop.reregister(fd, token, interest, poll_opt)
	}
}

impl PipeState for HandshakePipeState {
	fn enter(&mut self, event_loop: &mut EventLoop) -> io::Result<()> {
		debug!("Enter handshake state");
		self.sent = false;
		self.received = false;

		let interest = mio::Interest::readable() | mio::Interest::writable() | mio::Interest::hup();
		let poll_opt = mio::PollOpt::oneshot();

		try!(self.register(event_loop, interest, poll_opt));

		Ok(())
	}

	fn readable(&mut self, event_loop: &mut EventLoop) -> io::Result<Option<PipeStateIdx>> {
		if self.received {
			return Ok(None)
		}

		let mut handshake = [0u8; 8];
		let mut connection = self.connection.borrow_mut();
		try!(
			connection.try_read(&mut handshake).
			and_then(|_| self.check_received_handshake(&handshake)));
		drop(connection);

		debug!("handshake received !");
		self.received = true;

		if self.sent {
			Ok(Some(PipeStateIdx::Ready))
		} else {
			let interest = mio::Interest::hup() | mio::Interest::writable();
			let poll_opt = mio::PollOpt::oneshot();

			try!(self.reregister(event_loop, interest, poll_opt));
			Ok(None)
		}
	}

	fn writable(&mut self, event_loop: &mut EventLoop) -> io::Result<Option<PipeStateIdx>> {
		if self.sent {
			return Ok(None)
		}
		// handshake is Zero, 'S', 'P', Version, Proto, Rsvd
		let mut handshake = vec!(0, 83, 80, 0);
		let protocol_id = self.protocol.id();
		try!(handshake.write_u16::<BigEndian>(protocol_id));
		try!(handshake.write_u16::<BigEndian>(0));

		let mut connection = self.connection.borrow_mut();
		try!(
			connection.try_write(&handshake).
			and_then(|w| self.check_sent_handshake(w)));
		drop(connection);

		debug!("handshake sent !");
		self.sent = true;

		if self.received {
			Ok(Some(PipeStateIdx::Ready))
		} else {
			let interest = mio::Interest::hup() | mio::Interest::readable();
			let poll_opt = mio::PollOpt::oneshot();
			try!(self.reregister(event_loop, interest, poll_opt));
			Ok(None)
		}
	}
}

struct ReadyPipeState {
	id: usize,
	protocol: Rc<Box<Protocol>>,
	connection: Rc<RefCell<Box<Connection>>>,
	sent: bool
}

impl ReadyPipeState {
	fn new(id: usize, protocol: Rc<Box<Protocol>>, connection: Rc<RefCell<Box<Connection>>>) -> ReadyPipeState {
		ReadyPipeState {
			id: id,
			protocol: protocol,
			connection: connection,
			sent: false }
	}
}

impl PipeState for ReadyPipeState {
	fn enter(&mut self, event_loop: &mut EventLoop) -> io::Result<()> {
		debug!("Enter ready state");
		let token = mio::Token(self.id);
		let connection = self.connection.borrow_mut();
		let fd = connection.as_evented();
		let interest = mio::Interest::hup() | mio::Interest::writable();
		let poll_opt = mio::PollOpt::edge();
		try!(event_loop.reregister(fd, token, interest, poll_opt));
		Ok(())
	}

	fn readable(&mut self, _: &mut EventLoop) -> io::Result<Option<PipeStateIdx>> {
		Ok(None)
	}

	fn writable(&mut self, _: &mut EventLoop) -> io::Result<Option<PipeStateIdx>> {
		if self.sent {
			Ok(None)
		} else {
			let mut connection = self.connection.borrow_mut();
			let header = [0, 0, 0, 0, 0, 0, 0, 3];
			let body = [65, 86, 67];

			try!(connection.try_write(&header));
			try!(connection.try_write(&body));
			self.sent = true;
			debug!("dummy msg sent !");

			Ok(None)
		}
	}
}<commit_msg>Moves read and write handshare to dedicated functions<commit_after>use std::rc::Rc;
use std::cell::RefCell;
use std::io;

use byteorder::{BigEndian, WriteBytesExt};

use mio;

use EventLoop;
use protocol::Protocol as Protocol;
use transport::Connection as Connection;

enum PipeStateIdx {
	Handshake,
	Ready
}

pub struct Pipe {
	state: PipeStateIdx,
	handshake_state: HandshakePipeState,
	ready_state: ReadyPipeState
}

impl Pipe {

	pub fn new(id: usize, protocol: Rc<Box<Protocol>>, connection: Box<Connection>) -> Pipe {
		let conn_ref = Rc::new(RefCell::new(connection));

		Pipe {
			state: PipeStateIdx::Handshake,
			handshake_state: HandshakePipeState::new(id, protocol.clone(), conn_ref.clone()),
			ready_state: ReadyPipeState::new(id, protocol.clone(), conn_ref.clone())
		}
	}

	fn get_state<'a>(&'a mut self) -> &'a mut PipeState {
		match self.state {
			PipeStateIdx::Handshake => &mut self.handshake_state,
			PipeStateIdx::Ready => &mut self.ready_state
		}
	}

	fn set_state(&mut self, state: PipeStateIdx, event_loop: &mut EventLoop) -> io::Result<()> {
		self.state = state;
		self.get_state().enter(event_loop)
	}

	pub fn init(&mut self, event_loop: &mut EventLoop) -> io::Result<()> {
		self.set_state(PipeStateIdx::Handshake, event_loop)
	}

	pub fn readable(&mut self, event_loop: &mut EventLoop, hint: mio::ReadHint)	-> io::Result<()> {
		if let Some(new_state) = try!(self.get_state().readable(event_loop)) {
			try!(self.set_state(new_state, event_loop));
		}

		Ok(())
	}

	pub fn writable(&mut self, event_loop: &mut EventLoop) -> io::Result<()> {
		if let Some(new_state) = try!(self.get_state().writable(event_loop)) {
			try!(self.set_state(new_state, event_loop));
		}

		Ok(())
	}

}

trait PipeState {
	fn enter(&mut self, event_loop: &mut EventLoop) -> io::Result<()>;
	fn readable(&mut self, event_loop: &mut EventLoop) -> io::Result<Option<PipeStateIdx>>;
	fn writable(&mut self, event_loop: &mut EventLoop) -> io::Result<Option<PipeStateIdx>>;
}

struct HandshakePipeState {
	id: usize,
	protocol: Rc<Box<Protocol>>,
	connection: Rc<RefCell<Box<Connection>>>,
	sent: bool,
	received: bool
}

impl HandshakePipeState {

	fn new(id: usize, protocol: Rc<Box<Protocol>>, connection: Rc<RefCell<Box<Connection>>>) -> HandshakePipeState {
		HandshakePipeState {
			id: id,
			protocol: protocol,
			connection: connection,
			sent: false,
			received: false }
	}

	fn check_received_handshake(&self, handshake: &[u8; 8]) -> io::Result<()> {
		let mut expected_handshake = vec!(0, 83, 80, 0);
		let protocol_id = self.protocol.peer_id();
		try!(expected_handshake.write_u16::<BigEndian>(protocol_id));
		try!(expected_handshake.write_u16::<BigEndian>(0));
		let mut both = handshake.iter().zip(expected_handshake.iter());

		if both.all(|(l,r)| l == r) {
			Ok(())
		} else {
			Err(io::Error::new(io::ErrorKind::InvalidData, "received wrong handshake"))
		}
	}

	fn check_sent_handshake(&self, written: Option<usize>) -> io::Result<()> {
		match written {
			Some(8) => Ok(()),
			_ => Err(io::Error::new(io::ErrorKind::InvalidData, "failed to send handshake"))
		}
	}

	fn read_handshake(&mut self) -> io::Result<()> {
		let mut handshake = [0u8; 8];
		let mut connection = self.connection.borrow_mut();
		try!(
			connection.try_read(&mut handshake).
			and_then(|_| self.check_received_handshake(&handshake)));
		debug!("handshake received !");
		self.received = true;
		Ok(())
	}

	fn write_handshake(&mut self) -> io::Result<()> {
		// handshake is Zero, 'S', 'P', Version, Proto, Rsvd
		let mut handshake = vec!(0, 83, 80, 0);
		let protocol_id = self.protocol.id();
		try!(handshake.write_u16::<BigEndian>(protocol_id));
		try!(handshake.write_u16::<BigEndian>(0));
		let mut connection = self.connection.borrow_mut();
		try!(
			connection.try_write(&handshake).
			and_then(|w| self.check_sent_handshake(w)));
		debug!("handshake sent !");
		self.sent = true;
		Ok(())
	}

	fn register(&self, event_loop: &mut EventLoop, interest: mio::Interest, poll_opt: mio::PollOpt) -> io::Result<()> {
		let token = mio::Token(self.id);
		let connection = self.connection.borrow();
		let fd = connection.as_evented();

		event_loop.register_opt(fd, token, interest, poll_opt)
	}

	fn reregister(&self, event_loop: &mut EventLoop, interest: mio::Interest, poll_opt: mio::PollOpt) -> io::Result<()> {
		let token = mio::Token(self.id);
		let connection = self.connection.borrow();
		let fd = connection.as_evented();

		event_loop.reregister(fd, token, interest, poll_opt)
	}

}

impl PipeState for HandshakePipeState {

	fn enter(&mut self, event_loop: &mut EventLoop) -> io::Result<()> {
		debug!("Enter handshake state");
		self.sent = false;
		self.received = false;

		let interest = mio::Interest::readable() | mio::Interest::writable() | mio::Interest::hup();
		let poll_opt = mio::PollOpt::oneshot();

		try!(self.register(event_loop, interest, poll_opt));

		Ok(())
	}

	fn readable(&mut self, event_loop: &mut EventLoop) -> io::Result<Option<PipeStateIdx>> {
		if self.received {
			return Ok(None)
		}

		try!(self.read_handshake());

		if self.sent {
			Ok(Some(PipeStateIdx::Ready))
		} else {
			let interest = mio::Interest::hup() | mio::Interest::writable();
			let poll_opt = mio::PollOpt::oneshot();

			try!(self.reregister(event_loop, interest, poll_opt));
			Ok(None)
		}
	}

	fn writable(&mut self, event_loop: &mut EventLoop) -> io::Result<Option<PipeStateIdx>> {
		if self.sent {
			return Ok(None)
		}

		try!(self.write_handshake());

		if self.received {
			Ok(Some(PipeStateIdx::Ready))
		} else {
			let interest = mio::Interest::hup() | mio::Interest::readable();
			let poll_opt = mio::PollOpt::oneshot();
			try!(self.reregister(event_loop, interest, poll_opt));
			Ok(None)
		}
	}

}

struct ReadyPipeState {
	id: usize,
	protocol: Rc<Box<Protocol>>,
	connection: Rc<RefCell<Box<Connection>>>,
	sent: bool
}

impl ReadyPipeState {
	fn new(id: usize, protocol: Rc<Box<Protocol>>, connection: Rc<RefCell<Box<Connection>>>) -> ReadyPipeState {
		ReadyPipeState {
			id: id,
			protocol: protocol,
			connection: connection,
			sent: false }
	}
}

impl PipeState for ReadyPipeState {
	fn enter(&mut self, event_loop: &mut EventLoop) -> io::Result<()> {
		debug!("Enter ready state");
		let token = mio::Token(self.id);
		let connection = self.connection.borrow_mut();
		let fd = connection.as_evented();
		let interest = mio::Interest::hup() | mio::Interest::writable();
		let poll_opt = mio::PollOpt::edge();
		try!(event_loop.reregister(fd, token, interest, poll_opt));
		Ok(())
	}

	fn readable(&mut self, _: &mut EventLoop) -> io::Result<Option<PipeStateIdx>> {
		Ok(None)
	}

	fn writable(&mut self, _: &mut EventLoop) -> io::Result<Option<PipeStateIdx>> {
		if self.sent {
			Ok(None)
		} else {
			let mut connection = self.connection.borrow_mut();
			let header = [0, 0, 0, 0, 0, 0, 0, 3];
			let body = [65, 86, 67];

			try!(connection.try_write(&header));
			try!(connection.try_write(&body));
			self.sent = true;
			debug!("dummy msg sent !");

			Ok(None)
		}
	}
}<|endoftext|>


<commit_before>use std::io::fs;
use std::path::posix::Path;
use std::os::homedir;
#[cfg(test)] use test_helper::{cleanup_dir,random_name};

/// Create the muxed directory and return the path if creation is successful.
pub fn create_muxed_dir(name: &String) -> Path {
    let path = &Path::new(format!("{}/.{}", homedir_string(), name));
    match fs::mkdir(path, ::std::io::UserRWX) {
        Ok(()) => (),
        Err(_e) => println!("Failed to create project {}", path.filename()),
    }

    path.clone()
}

/// Return a boolean if the ~/.muxed/ dir exists.
pub fn muxed_dir_exists(name: &String) -> bool {
    let path = &Path::new(format!("{}/.{}", homedir_string(), name));
    path.exists()
}

/// Return the users current homedir as a string.
pub fn homedir_string() -> String {
    let home_unwrap = homedir().unwrap();
    format!("{}", home_unwrap.display())
}

#[test]
fn muxed_dir_exists_returns_false() {
    assert!(!muxed_dir_exists(&random_name()));
}

#[test]
fn muxed_dir_exists_returns_true() {
    let dir = random_name();
    create_muxed_dir(&dir);
    assert!(muxed_dir_exists(&dir));

    let muxed_path  = &Path::new(format!("{}/.{}/", homedir_string(), dir));
    cleanup_dir(muxed_path);
}

#[test]
fn creates_muxed_dir() {
    let dir = random_name();
    let muxed_path  = &Path::new(format!("{}/.{}/", homedir_string(), dir));
    create_muxed_dir(&dir);
    assert!(muxed_path.exists());
    cleanup_dir(muxed_path);
}
<commit_msg>Use a Path static variable and cange it in the test config. This will allow me to do large scale testing outside of the default environment.<commit_after>use std::io::fs;
use std::path::posix::Path;
use std::os::homedir;
#[cfg(test)] use test_helper::{cleanup_dir,random_name};

static MUXED_NAME_STR: &'static str = "muxed";
#[cfg(test)] static MUXED_PATH: Path = Path::new(format!("/tmp/.{}/", &MUXED_NAME_STR.to_string()));
//#[cfg(not_test)] static MUXED_DIR: Path = Path::new(format!("{}/.{}/", root::homedir_string(), &DEFAULT_MUXED_DIR.to_string()));

/// Create the muxed directory and return the path if creation is successful.
pub fn create_muxed_dir(name: &String) -> Path {
    let path = &Path::new(format!("{}/.{}", homedir_string(), name));
    match fs::mkdir(path, ::std::io::UserRWX) {
        Ok(()) => (),
        Err(_e) => println!("Failed to create project {}", path.filename()),
    }

    path.clone()
}

/// Return a boolean if the ~/.muxed/ dir exists.
pub fn muxed_dir_exists(name: &String) -> bool {
    let path = &Path::new(format!("{}/.{}", homedir_string(), name));
    path.exists()
}

/// Return the users current homedir as a string.
pub fn homedir_string() -> String {
    let home_unwrap = homedir().unwrap();
    format!("{}", home_unwrap.display())
}

#[test]
fn muxed_dir_exists_returns_false() {
    assert!(!muxed_dir_exists(&random_name()));
}

#[test]
fn muxed_dir_exists_returns_true() {
    let dir = random_name();
    create_muxed_dir(&dir);
    assert!(muxed_dir_exists(&dir));

    let muxed_path  = &Path::new(format!("{}/.{}/", homedir_string(), dir));
    cleanup_dir(muxed_path);
}

#[test]
fn creates_muxed_dir() {
    let dir = random_name();
    let muxed_path  = &Path::new(format!("{}/.{}/", homedir_string(), dir));
    create_muxed_dir(&dir);
    assert!(muxed_path.exists());
    cleanup_dir(muxed_path);
}
<|endoftext|>


<commit_before>use crate::direction::Direction;

pub struct Ship {
    state: ShipState,
    position: Vec<[u8; 2]>,
    dir: Direction,
}

impl Ship {
    /// Creates a new `Ship` with the given position.
    pub fn new(pos: Vec<[u8; 2]>) -> Ship {
        let dir = Direction::from_positions(&pos[1], &pos[0]).unwrap();

        Ship {
            state: ShipState::Placement,
            position: pos,
            dir: dir,
        }
    }

    /// Returns the ship's position.
    pub fn pos(&self) -> &[[u8; 2]] {
        &self.position
    }

    /// Sets the ship's position.
    ///
    /// # Errors
    ///
    /// Returns an error if `pos` does not form a vertical or horizontal line.
    pub fn set_pos(&mut self, pos: Vec<[u8; 2]>) -> Result<(), &'static str> {
        let dir = Direction::from_positions(&pos[1], &pos[0])?;

        self.position = pos;
        self.dir = dir;

        Ok(())
    }

    /// Returns the ship's direction.
    pub fn dir(&self) -> Direction {
        self.dir
    }

    /// Returns the ship's length.
    pub fn len(&self) -> usize {
        self.position.len()
    }

    /// Returns whether the ship is in the placement state.
    pub fn is_placement(&self) -> bool {
        self.state == ShipState::Placement
    }

    /// Returns whether the ship is active.
    pub fn is_active(&self) -> bool {
        self.state == ShipState::Active
    }

    /// Sets the ship as active.
    ///
    /// # Errors
    ///
    /// Returns an error if the ship's state is not `ShipState::Placement`.
    pub fn set_active(&mut self) -> Result<(), &'static str> {
        if self.state != ShipState::Placement {
            Err("tried to set a ship as active that wasn't in placement state")
        } else {
            self.state = ShipState::Active;

            Ok(())
        }
    }

    /// Returns whether the ship has sunk.
    pub fn is_sunk(&self) -> bool {
        self.state == ShipState::Sunk
    }

    /// Sets the ship as having been sunk.
    ///
    /// # Errors
    ///
    /// Returns an error if the ship's state is not `ShipState::Active`.
    pub fn set_sunk(&mut self) -> Result<(), &'static str> {
        if self.state != ShipState::Active {
            Err("tried to sink a ship that was not active")
        } else {
            self.state = ShipState::Sunk;

            Ok(())
        }
    }
}

#[derive(PartialEq)]
enum ShipState {
    Placement,
    Active,
    Sunk,
}

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

    #[test]
    fn pos() {
        let pos = vec![[0, 0], [0, 1]];
        let ship = Ship::new(pos.clone());
        assert_eq!(ship.pos(), pos.as_slice());
    }

    #[test]
    fn set_pos() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert!(ship.set_pos(vec![[1, 0], [0, 0]]).is_ok());
        assert_eq!(ship.dir(), Direction::East);

        assert!(ship.set_pos(vec![[0, 1], [0, 0]]).is_ok());
        assert_eq!(ship.dir(), Direction::South);

        assert!(ship.set_pos(vec![[0, 0], [1, 0]]).is_ok());
        assert_eq!(ship.dir(), Direction::West);

        assert!(ship.set_pos(vec![[0, 0], [0, 1]]).is_ok());
        assert_eq!(ship.dir(), Direction::North);

        assert!(ship.set_pos(vec![[0, 0], [0, 0]]).is_err());
        assert!(ship.set_pos(vec![[0, 0], [0, 2]]).is_err());
        assert!(ship.set_pos(vec![]).is_err());
    }

    #[test]
    fn dir() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert_eq!(ship.dir(), Direction::North);

        ship = Ship::new(vec![[0, 1], [0, 0]]);
        assert_eq!(ship.dir(), Direction::South);

        ship = Ship::new(vec![[0, 0], [1, 0]]);
        assert_eq!(ship.dir(), Direction::West);

        ship = Ship::new(vec![[1, 0], [0, 0]]);
        assert_eq!(ship.dir(), Direction::East);
    }

    #[test]
    fn len() {
        let pos = vec![[0, 0], [0, 1]];
        let ship = Ship::new(pos.clone());
        assert_eq!(ship.len(), pos.len());
    }

    #[test]
    fn is_placement() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert!(ship.is_placement());
        assert!(ship.set_active().is_ok());
        assert!(!ship.is_placement());
        assert!(ship.set_sunk().is_ok());
        assert!(!ship.is_placement());
    }

    #[test]
    fn is_active() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert!(!ship.is_active());
        assert!(ship.set_active().is_ok());
        assert!(ship.is_active());
        assert!(ship.set_sunk().is_ok());
        assert!(!ship.is_active());
    }

    #[test]
    fn set_active() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert!(ship.set_active().is_ok());
        assert!(ship.set_active().is_err());
        assert!(ship.set_sunk().is_ok());
        assert!(ship.set_active().is_err());
    }

    #[test]
    fn is_sunk() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert!(!ship.is_sunk());
        assert!(ship.set_active().is_ok());
        assert!(!ship.is_sunk());
        assert!(ship.set_sunk().is_ok());
        assert!(ship.is_sunk());
    }

    #[test]
    fn set_sunk() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert!(ship.set_sunk().is_err());
        assert!(ship.set_active().is_ok());
        assert!(ship.set_sunk().is_ok());
        assert!(ship.set_sunk().is_err());
    }
}
<commit_msg>Ship: fix issues with set_pos()<commit_after>use crate::direction::Direction;

pub struct Ship {
    state: ShipState,
    position: Vec<[u8; 2]>,
    dir: Direction,
}

impl Ship {
    /// Creates a new `Ship` with the given position.
    pub fn new(pos: Vec<[u8; 2]>) -> Ship {
        let dir = Direction::from_positions(&pos[1], &pos[0]).unwrap();

        Ship {
            state: ShipState::Placement,
            position: pos,
            dir: dir,
        }
    }

    /// Returns the ship's position.
    pub fn pos(&self) -> &[[u8; 2]] {
        &self.position
    }

    /// Sets the ship's position.
    ///
    /// # Errors
    ///
    /// Returns an error if `pos` does not form a vertical or horizontal line.
    pub fn set_pos(&mut self, pos: Vec<[u8; 2]>) -> Result<(), &'static str> {
        if pos.is_empty() {
            Err("tried to set an empty position to a ship")
        } else if pos.len() == 1 {
            self.position = pos;

            Ok(())
        } else {
            let mut valid = true;
            let dir = Direction::from_positions(&pos[1], &pos[0])?;

            for i in 0..pos.len() - 1 {
                let x_diff = (pos[i][0] as i16 - pos[i + 1][0] as i16).abs() as u8;
                let y_diff = (pos[i][1] as i16 - pos[i + 1][1] as i16).abs() as u8;

                if x_diff + y_diff != 1 {
                    valid = false;
                }

                if valid {
                    let next_dir = Direction::from_positions(&pos[i + 1], &pos[i])?;

                    if next_dir != dir {
                        valid = false;
                    }
                }
            }

            if !valid {
                Err("ship position does not form a continuous line")
            } else {
                let dir = Direction::from_positions(&pos[1], &pos[0])?;

                self.position = pos;
                self.dir = dir;

                Ok(())
            }
        }
    }

    /// Returns the ship's direction.
    pub fn dir(&self) -> Direction {
        self.dir
    }

    /// Returns the ship's length.
    pub fn len(&self) -> usize {
        self.position.len()
    }

    /// Returns whether the ship is in the placement state.
    pub fn is_placement(&self) -> bool {
        self.state == ShipState::Placement
    }

    /// Returns whether the ship is active.
    pub fn is_active(&self) -> bool {
        self.state == ShipState::Active
    }

    /// Sets the ship as active.
    ///
    /// # Errors
    ///
    /// Returns an error if the ship's state is not `ShipState::Placement`.
    pub fn set_active(&mut self) -> Result<(), &'static str> {
        if self.state != ShipState::Placement {
            Err("tried to set a ship as active that wasn't in placement state")
        } else {
            self.state = ShipState::Active;

            Ok(())
        }
    }

    /// Returns whether the ship has sunk.
    pub fn is_sunk(&self) -> bool {
        self.state == ShipState::Sunk
    }

    /// Sets the ship as having been sunk.
    ///
    /// # Errors
    ///
    /// Returns an error if the ship's state is not `ShipState::Active`.
    pub fn set_sunk(&mut self) -> Result<(), &'static str> {
        if self.state != ShipState::Active {
            Err("tried to sink a ship that was not active")
        } else {
            self.state = ShipState::Sunk;

            Ok(())
        }
    }
}

#[derive(PartialEq)]
enum ShipState {
    Placement,
    Active,
    Sunk,
}

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

    #[test]
    fn pos() {
        let pos = vec![[0, 0], [0, 1]];
        let ship = Ship::new(pos.clone());
        assert_eq!(ship.pos(), pos.as_slice());
    }

    #[test]
    fn set_pos() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert!(ship.set_pos(vec![[1, 0], [0, 0]]).is_ok());
        assert_eq!(ship.dir(), Direction::East);

        assert!(ship.set_pos(vec![[0, 1], [0, 0]]).is_ok());
        assert_eq!(ship.dir(), Direction::South);

        assert!(ship.set_pos(vec![[0, 0], [1, 0]]).is_ok());
        assert_eq!(ship.dir(), Direction::West);

        assert!(ship.set_pos(vec![[0, 0], [0, 1]]).is_ok());
        assert_eq!(ship.dir(), Direction::North);

        assert!(ship.set_pos(vec![[0, 0], [0, 0]]).is_err());
        assert!(ship.set_pos(vec![[0, 0], [0, 2]]).is_err());
        assert!(ship.set_pos(vec![]).is_err());
    }

    #[test]
    fn dir() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert_eq!(ship.dir(), Direction::North);

        ship = Ship::new(vec![[0, 1], [0, 0]]);
        assert_eq!(ship.dir(), Direction::South);

        ship = Ship::new(vec![[0, 0], [1, 0]]);
        assert_eq!(ship.dir(), Direction::West);

        ship = Ship::new(vec![[1, 0], [0, 0]]);
        assert_eq!(ship.dir(), Direction::East);
    }

    #[test]
    fn len() {
        let pos = vec![[0, 0], [0, 1]];
        let ship = Ship::new(pos.clone());
        assert_eq!(ship.len(), pos.len());
    }

    #[test]
    fn is_placement() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert!(ship.is_placement());
        assert!(ship.set_active().is_ok());
        assert!(!ship.is_placement());
        assert!(ship.set_sunk().is_ok());
        assert!(!ship.is_placement());
    }

    #[test]
    fn is_active() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert!(!ship.is_active());
        assert!(ship.set_active().is_ok());
        assert!(ship.is_active());
        assert!(ship.set_sunk().is_ok());
        assert!(!ship.is_active());
    }

    #[test]
    fn set_active() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert!(ship.set_active().is_ok());
        assert!(ship.set_active().is_err());
        assert!(ship.set_sunk().is_ok());
        assert!(ship.set_active().is_err());
    }

    #[test]
    fn is_sunk() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert!(!ship.is_sunk());
        assert!(ship.set_active().is_ok());
        assert!(!ship.is_sunk());
        assert!(ship.set_sunk().is_ok());
        assert!(ship.is_sunk());
    }

    #[test]
    fn set_sunk() {
        let mut ship = Ship::new(vec![[0, 0], [0, 1]]);
        assert!(ship.set_sunk().is_err());
        assert!(ship.set_active().is_ok());
        assert!(ship.set_sunk().is_ok());
        assert!(ship.set_sunk().is_err());
    }
}
<|endoftext|>


<commit_before>use std::io;

#[cfg(not(target_os = "redox"))]
use libc::c_ushort;

#[cfg(not(target_os = "redox"))]
#[repr(C)]
struct TermSize {
    row: c_ushort,
    col: c_ushort,
    _x: c_ushort,
    _y: c_ushort,
}

// Since attributes on non-item statements is not stable yet, we use a function.
#[cfg(not(target_os = "redox"))]
#[cfg(target_pointer_width = "64")]
fn tiocgwinsz() -> u64 {
    use termios::TIOCGWINSZ;
    TIOCGWINSZ as u64
}
#[cfg(not(target_os = "redox"))]
#[cfg(target_pointer_width = "32")]
fn tiocgwinsz() -> u32 {
    use termios::TIOCGWINSZ;
    TIOCGWINSZ as u32
}

/// Get the size of the terminal.
#[cfg(not(target_os = "redox"))]
pub fn terminal_size() -> io::Result<(u16, u16)> {
    use libc::ioctl;
    use libc::STDOUT_FILENO;

    use std::mem;
    unsafe {
        let mut size: TermSize = mem::zeroed();

        if ioctl(STDOUT_FILENO, tiocgwinsz(), &mut size as *mut _) == 0 {
            Ok((size.col as u16, size.row as u16))
        } else {
            Err(io::Error::new(io::ErrorKind::Other, "Unable to get the terminal size."))
        }
    }
}

/// Get the size of the terminal.
#[cfg(target_os = "redox")]
pub fn terminal_size() -> io::Result<(u16, u16)> {
    use std::env;

    let width = try!(env::var("COLUMNS").map_err(|x| io::Error::new(io::ErrorKind::NotFound, x))).parse().unwrap_or(0);
    let height = try!(env::var("LINES").map_err(|x| io::Error::new(io::ErrorKind::NotFound, x))).parse().unwrap_or(0);

    Ok((width, height))
}

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

    #[test]
    fn test_size() {
        assert!(terminal_size().is_ok());
    }
}
<commit_msg>bug/lib.rs Allow musl builds (#58)<commit_after>use std::io;

#[cfg(not(target_os = "redox"))]
use libc::c_ushort;

#[cfg(not(target_os = "redox"))]
#[repr(C)]
struct TermSize {
    row: c_ushort,
    col: c_ushort,
    _x: c_ushort,
    _y: c_ushort,
}

// Since attributes on non-item statements is not stable yet, we use a function.
#[cfg(not(target_os = "redox"))]
#[cfg(target_pointer_width = "64")]
#[cfg(not(target_env = "musl"))]
fn tiocgwinsz() -> u64 {
    use termios::TIOCGWINSZ;
    TIOCGWINSZ as u64
}
#[cfg(not(target_os = "redox"))]
#[cfg(target_pointer_width = "32")]
#[cfg(not(target_env = "musl"))]
fn tiocgwinsz() -> u32 {
    use termios::TIOCGWINSZ;
    TIOCGWINSZ as u32
}

#[cfg(target_env = "musl")]
fn tiocgwinsz() -> i32 {
    use termios::TIOCGWINSZ;
    TIOCGWINSZ as i32
}

/// Get the size of the terminal.
#[cfg(not(target_os = "redox"))]
pub fn terminal_size() -> io::Result<(u16, u16)> {
    use libc::ioctl;
    use libc::STDOUT_FILENO;

    use std::mem;
    unsafe {
        let mut size: TermSize = mem::zeroed();

        if ioctl(STDOUT_FILENO, tiocgwinsz(), &mut size as *mut _) == 0 {
            Ok((size.col as u16, size.row as u16))
        } else {
            Err(io::Error::new(io::ErrorKind::Other, "Unable to get the terminal size."))
        }
    }
}

/// Get the size of the terminal.
#[cfg(target_os = "redox")]
pub fn terminal_size() -> io::Result<(u16, u16)> {
    use std::env;

    let width = try!(env::var("COLUMNS").map_err(|x| io::Error::new(io::ErrorKind::NotFound, x)))
        .parse()
        .unwrap_or(0);
    let height = try!(env::var("LINES").map_err(|x| io::Error::new(io::ErrorKind::NotFound, x)))
        .parse()
        .unwrap_or(0);

    Ok((width, height))
}

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

    #[test]
    fn test_size() {
        assert!(terminal_size().is_ok());
    }
}
<|endoftext|>


<commit_before>// TODO: Break up Ring Trait in WriteableRing and ReadableRing
pub trait Ring {
    fn input_buffer(&mut self) -> &mut [u8];
    fn output_buffer(&mut self) -> &mut [u8];

    fn get_in_cons(&self) -> usize;
    fn get_out_cons(&self) -> usize;
    fn get_in_prod(&self) -> usize;
    fn get_out_prod(&self) -> usize;

    fn set_in_cons(&mut self, usize);
    fn set_out_cons(&mut self, usize);
    fn set_in_prod(&mut self, usize);
    fn set_out_prod(&mut self, usize);

    fn input_mask(&mut self) -> usize {
        let len = self.input_buffer().len();
        return len - 1;
    }

    fn output_mask(&mut self) -> usize {
        let len = self.output_buffer().len();
        return len - 1;
    }

    fn write(&mut self) {
    }
    fn read(&mut self) {
    }
}
<commit_msg>Implemented ring buffer trait, still some work to do though<commit_after>use ::xen::arch::mem::*;
use ::xen::event_channels::send;

// TODO: Break up Ring Trait in WriteableRing and ReadableRing
pub trait Ring {
    fn input_buffer(&mut self) -> &mut [u8];
    fn output_buffer(&mut self) -> &mut [u8];

    fn get_in_cons(&self) -> usize;
    fn get_out_cons(&self) -> usize;
    fn get_in_prod(&self) -> usize;
    fn get_out_prod(&self) -> usize;

    fn set_in_cons(&mut self, usize);
    fn set_out_cons(&mut self, usize);
    fn set_in_prod(&mut self, usize);
    fn set_out_prod(&mut self, usize);

    fn input_mask(&mut self) -> usize {
        let len = self.input_buffer().len();
        return len - 1;
    }

    fn output_mask(&mut self) -> usize {
        let len = self.output_buffer().len();
        return len - 1;
    }

    fn write(&mut self, buf : &[u8]) {
        let mut sent = 0usize;

        let cons = self.get_out_cons();
        let mut prod = self.get_out_prod();

        mb();

        {
            let output_mask = { self.output_mask() };
            let output_len = { self.output_buffer().len() };
            let mut output = self.output_buffer();

            while (sent < buf.len()) && ((prod - cons) < output_len) {
                let idx = prod & output_mask; // mask the index to make sure we don't overflow
                prod = prod + 1;
                output[idx] = buf[sent];
                sent = sent + 1;
            }
        }

        wmb();

        self.set_out_prod(prod);

    }

    fn write_notify(&mut self, buf : &[u8], evtchn : u32) {
        self.write(buf);
        unsafe { send(evtchn); }
    }

    fn read(&mut self) {
        // TODO
    }
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! The repartition operator maps N input partitions to M output partitions based on a
//! partitioning scheme.

use std::any::Any;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use crate::error::{DataFusionError, Result};
use crate::physical_plan::{ExecutionPlan, Partitioning};
use arrow::datatypes::SchemaRef;
use arrow::error::Result as ArrowResult;
use arrow::record_batch::RecordBatch;

use super::{RecordBatchStream, SendableRecordBatchStream};
use async_trait::async_trait;

use crossbeam::channel::{unbounded, Receiver, Sender};
use futures::stream::Stream;
use futures::StreamExt;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;

type MaybeBatch = Option<ArrowResult<RecordBatch>>;

/// The repartition operator maps N input partitions to M output partitions based on a
/// partitioning scheme. No guarantees are made about the order of the resulting partitions.
#[derive(Debug)]
pub struct RepartitionExec {
    /// Input execution plan
    input: Arc<dyn ExecutionPlan>,
    /// Partitioning scheme to use
    partitioning: Partitioning,
    /// Channels for sending batches from input partitions to output partitions
    /// there is one entry in this Vec for each output partition
    channels: Arc<Mutex<Vec<(Sender<MaybeBatch>, Receiver<MaybeBatch>)>>>,
}

impl RepartitionExec {
    /// Input execution plan
    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
        &self.input
    }

    /// Partitioning scheme to use
    pub fn partitioning(&self) -> &Partitioning {
        &self.partitioning
    }
}

#[async_trait]
impl ExecutionPlan for RepartitionExec {
    /// Return a reference to Any that can be used for downcasting
    fn as_any(&self) -> &dyn Any {
        self
    }

    /// Get the schema for this execution plan
    fn schema(&self) -> SchemaRef {
        self.input.schema()
    }

    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
        vec![self.input.clone()]
    }

    fn with_new_children(
        &self,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        match children.len() {
            1 => Ok(Arc::new(RepartitionExec::try_new(
                children[0].clone(),
                self.partitioning.clone(),
            )?)),
            _ => Err(DataFusionError::Internal(
                "RepartitionExec wrong number of children".to_string(),
            )),
        }
    }

    fn output_partitioning(&self) -> Partitioning {
        self.partitioning.clone()
    }

    async fn execute(&self, partition: usize) -> Result<SendableRecordBatchStream> {
        // lock mutexes
        let mut channels = self.channels.lock().await;

        let num_input_partitions = self.input.output_partitioning().partition_count();
        let num_output_partitions = self.partitioning.partition_count();

        // if this is the first partition to be invoked then we need to set up initial state
        if channels.is_empty() {
            // create one channel per *output* partition
            for _ in 0..num_output_partitions {
                // Note that this operator uses unbounded channels to avoid deadlocks because
                // the output partitions can be read in any order and this could cause input
                // partitions to be blocked when sending data to output receivers that are not
                // being read yet. This may cause high memory usage if the next operator is
                // reading output partitions in order rather than concurrently. One workaround
                // for this would be to add spill-to-disk capabilities.
                let (sender, receiver) = unbounded::<Option<ArrowResult<RecordBatch>>>();
                channels.push((sender, receiver));
            }
            // launch one async task per *input* partition
            for i in 0..num_input_partitions {
                let input = self.input.clone();
                let mut channels = channels.clone();
                let partitioning = self.partitioning.clone();
                let join_handle: JoinHandle<Result<()>> = tokio::spawn(async move {
                    let mut stream = input.execute(i).await?;
                    let mut counter = 0;
                    while let Some(result) = stream.next().await {
                        match partitioning {
                            Partitioning::RoundRobinBatch(_) => {
                                let output_partition = counter % num_output_partitions;
                                let tx = &mut channels[output_partition].0;
                                tx.send(Some(result)).map_err(|e| {
                                    DataFusionError::Execution(e.to_string())
                                })?;
                            }
                            other => {
                                // this should be unreachable as long as the validation logic
                                // in the constructor is kept up-to-date
                                return Err(DataFusionError::NotImplemented(format!(
                                    "Unsupported repartitioning scheme {:?}",
                                    other
                                )));
                            }
                        }
                        counter += 1;
                    }

                    // notify each output partition that this input partition has no more data
                    for channel in channels.iter_mut().take(num_output_partitions) {
                        let tx = &mut channel.0;
                        tx.send(None)
                            .map_err(|e| DataFusionError::Execution(e.to_string()))?;
                    }
                    Ok(())
                });
                join_handle
                    .await
                    .map(|_| ())
                    .map_err(|e| DataFusionError::Execution(e.to_string()))?;
            }
        }

        // now return stream for the specified *output* partition which will
        // read from the channel
        Ok(Box::pin(RepartitionStream {
            num_input_partitions,
            num_input_partitions_processed: 0,
            schema: self.input.schema(),
            input: channels[partition].1.clone(),
        }))
    }
}

impl RepartitionExec {
    /// Create a new RepartitionExec
    pub fn try_new(
        input: Arc<dyn ExecutionPlan>,
        partitioning: Partitioning,
    ) -> Result<Self> {
        match &partitioning {
            Partitioning::RoundRobinBatch(_) => Ok(RepartitionExec {
                input,
                partitioning,
                channels: Arc::new(Mutex::new(vec![])),
            }),
            other => Err(DataFusionError::NotImplemented(format!(
                "Partitioning scheme not supported yet: {:?}",
                other
            ))),
        }
    }
}

struct RepartitionStream {
    /// Number of input partitions that will be sending batches to this output channel
    num_input_partitions: usize,
    /// Number of input partitions that have finished sending batches to this output channel
    num_input_partitions_processed: usize,
    /// Schema
    schema: SchemaRef,
    /// channel containing the repartitioned batches
    input: Receiver<Option<ArrowResult<RecordBatch>>>,
}

impl Stream for RepartitionStream {
    type Item = ArrowResult<RecordBatch>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        match self.input.recv() {
            Ok(Some(batch)) => Poll::Ready(Some(batch)),
            // End of results from one input partition
            Ok(None) => {
                self.num_input_partitions_processed += 1;
                if self.num_input_partitions == self.num_input_partitions_processed {
                    // all input partitions have finished sending batches
                    Poll::Ready(None)
                } else {
                    // other partitions still have data to send
                    self.poll_next(cx)
                }
            }
            // RecvError means receiver has exited and closed the channel
            Err(_) => Poll::Ready(None),
        }
    }
}

impl RecordBatchStream for RepartitionStream {
    /// Get the schema
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::physical_plan::memory::MemoryExec;
    use arrow::array::UInt32Array;
    use arrow::datatypes::{DataType, Field, Schema};
    use arrow::record_batch::RecordBatch;

    #[tokio::test(flavor = "multi_thread")]
    async fn one_to_many_round_robin() -> Result<()> {
        // define input partitions
        let schema = test_schema();
        let partition = create_vec_batches(&schema, 50);
        let partitions = vec![partition];

        // repartition from 1 input to 4 output
        let output_partitions =
            repartition(&schema, partitions, Partitioning::RoundRobinBatch(4)).await?;

        assert_eq!(4, output_partitions.len());
        assert_eq!(13, output_partitions[0].len());
        assert_eq!(13, output_partitions[1].len());
        assert_eq!(12, output_partitions[2].len());
        assert_eq!(12, output_partitions[3].len());

        Ok(())
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn many_to_one_round_robin() -> Result<()> {
        // define input partitions
        let schema = test_schema();
        let partition = create_vec_batches(&schema, 50);
        let partitions = vec![partition.clone(), partition.clone(), partition.clone()];

        // repartition from 3 input to 1 output
        let output_partitions =
            repartition(&schema, partitions, Partitioning::RoundRobinBatch(1)).await?;

        assert_eq!(1, output_partitions.len());
        assert_eq!(150, output_partitions[0].len());

        Ok(())
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn many_to_many_round_robin() -> Result<()> {
        // define input partitions
        let schema = test_schema();
        let partition = create_vec_batches(&schema, 50);
        let partitions = vec![partition.clone(), partition.clone(), partition.clone()];

        // repartition from 3 input to 5 output
        let output_partitions =
            repartition(&schema, partitions, Partitioning::RoundRobinBatch(5)).await?;

        assert_eq!(5, output_partitions.len());
        assert_eq!(30, output_partitions[0].len());
        assert_eq!(30, output_partitions[1].len());
        assert_eq!(30, output_partitions[2].len());
        assert_eq!(30, output_partitions[3].len());
        assert_eq!(30, output_partitions[4].len());

        Ok(())
    }

    fn test_schema() -> Arc<Schema> {
        Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]))
    }

    fn create_vec_batches(schema: &Arc<Schema>, n: usize) -> Vec<RecordBatch> {
        let batch = create_batch(schema);
        let mut vec = Vec::with_capacity(n);
        for _ in 0..n {
            vec.push(batch.clone());
        }
        vec
    }

    fn create_batch(schema: &Arc<Schema>) -> RecordBatch {
        RecordBatch::try_new(
            schema.clone(),
            vec![Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]))],
        )
        .unwrap()
    }

    async fn repartition(
        schema: &SchemaRef,
        input_partitions: Vec<Vec<RecordBatch>>,
        partitioning: Partitioning,
    ) -> Result<Vec<Vec<RecordBatch>>> {
        // create physical plan
        let exec = MemoryExec::try_new(&input_partitions, schema.clone(), None)?;
        let exec = RepartitionExec::try_new(Arc::new(exec), partitioning)?;

        // execute and collect results
        let mut output_partitions = vec![];
        for i in 0..exec.partitioning.partition_count() {
            // execute this *output* partition and collect all batches
            let mut stream = exec.execute(i).await?;
            let mut batches = vec![];
            while let Some(result) = stream.next().await {
                batches.push(result?);
            }
            output_partitions.push(batches);
        }
        Ok(output_partitions)
    }
}
<commit_msg>ARROW-11791: [Rust][DataFusion] Fix RepartitionExec Blocking<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! The repartition operator maps N input partitions to M output partitions based on a
//! partitioning scheme.

use std::any::Any;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use crate::error::{DataFusionError, Result};
use crate::physical_plan::{ExecutionPlan, Partitioning};
use arrow::datatypes::SchemaRef;
use arrow::error::Result as ArrowResult;
use arrow::record_batch::RecordBatch;

use super::{RecordBatchStream, SendableRecordBatchStream};
use async_trait::async_trait;

use crossbeam::channel::{unbounded, Receiver, Sender};
use futures::stream::Stream;
use futures::StreamExt;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;

type MaybeBatch = Option<ArrowResult<RecordBatch>>;

/// The repartition operator maps N input partitions to M output partitions based on a
/// partitioning scheme. No guarantees are made about the order of the resulting partitions.
#[derive(Debug)]
pub struct RepartitionExec {
    /// Input execution plan
    input: Arc<dyn ExecutionPlan>,
    /// Partitioning scheme to use
    partitioning: Partitioning,
    /// Channels for sending batches from input partitions to output partitions
    /// there is one entry in this Vec for each output partition
    channels: Arc<Mutex<Vec<(Sender<MaybeBatch>, Receiver<MaybeBatch>)>>>,
}

impl RepartitionExec {
    /// Input execution plan
    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
        &self.input
    }

    /// Partitioning scheme to use
    pub fn partitioning(&self) -> &Partitioning {
        &self.partitioning
    }
}

#[async_trait]
impl ExecutionPlan for RepartitionExec {
    /// Return a reference to Any that can be used for downcasting
    fn as_any(&self) -> &dyn Any {
        self
    }

    /// Get the schema for this execution plan
    fn schema(&self) -> SchemaRef {
        self.input.schema()
    }

    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
        vec![self.input.clone()]
    }

    fn with_new_children(
        &self,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        match children.len() {
            1 => Ok(Arc::new(RepartitionExec::try_new(
                children[0].clone(),
                self.partitioning.clone(),
            )?)),
            _ => Err(DataFusionError::Internal(
                "RepartitionExec wrong number of children".to_string(),
            )),
        }
    }

    fn output_partitioning(&self) -> Partitioning {
        self.partitioning.clone()
    }

    async fn execute(&self, partition: usize) -> Result<SendableRecordBatchStream> {
        // lock mutexes
        let mut channels = self.channels.lock().await;

        let num_input_partitions = self.input.output_partitioning().partition_count();
        let num_output_partitions = self.partitioning.partition_count();

        // if this is the first partition to be invoked then we need to set up initial state
        if channels.is_empty() {
            // create one channel per *output* partition
            for _ in 0..num_output_partitions {
                // Note that this operator uses unbounded channels to avoid deadlocks because
                // the output partitions can be read in any order and this could cause input
                // partitions to be blocked when sending data to output receivers that are not
                // being read yet. This may cause high memory usage if the next operator is
                // reading output partitions in order rather than concurrently. One workaround
                // for this would be to add spill-to-disk capabilities.
                let (sender, receiver) = unbounded::<Option<ArrowResult<RecordBatch>>>();
                channels.push((sender, receiver));
            }
            // launch one async task per *input* partition
            for i in 0..num_input_partitions {
                let input = self.input.clone();
                let mut channels = channels.clone();
                let partitioning = self.partitioning.clone();
                let _: JoinHandle<Result<()>> = tokio::spawn(async move {
                    let mut stream = input.execute(i).await?;
                    let mut counter = 0;
                    while let Some(result) = stream.next().await {
                        match partitioning {
                            Partitioning::RoundRobinBatch(_) => {
                                let output_partition = counter % num_output_partitions;
                                let tx = &mut channels[output_partition].0;
                                tx.send(Some(result)).map_err(|e| {
                                    DataFusionError::Execution(e.to_string())
                                })?;
                            }
                            other => {
                                // this should be unreachable as long as the validation logic
                                // in the constructor is kept up-to-date
                                return Err(DataFusionError::NotImplemented(format!(
                                    "Unsupported repartitioning scheme {:?}",
                                    other
                                )));
                            }
                        }
                        counter += 1;
                    }

                    // notify each output partition that this input partition has no more data
                    for channel in channels.iter_mut().take(num_output_partitions) {
                        let tx = &mut channel.0;
                        tx.send(None)
                            .map_err(|e| DataFusionError::Execution(e.to_string()))?;
                    }
                    Ok(())
                });
                tokio::task::yield_now().await;
            }
        }

        // now return stream for the specified *output* partition which will
        // read from the channel
        Ok(Box::pin(RepartitionStream {
            num_input_partitions,
            num_input_partitions_processed: 0,
            schema: self.input.schema(),
            input: channels[partition].1.clone(),
        }))
    }
}

impl RepartitionExec {
    /// Create a new RepartitionExec
    pub fn try_new(
        input: Arc<dyn ExecutionPlan>,
        partitioning: Partitioning,
    ) -> Result<Self> {
        match &partitioning {
            Partitioning::RoundRobinBatch(_) => Ok(RepartitionExec {
                input,
                partitioning,
                channels: Arc::new(Mutex::new(vec![])),
            }),
            other => Err(DataFusionError::NotImplemented(format!(
                "Partitioning scheme not supported yet: {:?}",
                other
            ))),
        }
    }
}

struct RepartitionStream {
    /// Number of input partitions that will be sending batches to this output channel
    num_input_partitions: usize,
    /// Number of input partitions that have finished sending batches to this output channel
    num_input_partitions_processed: usize,
    /// Schema
    schema: SchemaRef,
    /// channel containing the repartitioned batches
    input: Receiver<Option<ArrowResult<RecordBatch>>>,
}

impl Stream for RepartitionStream {
    type Item = ArrowResult<RecordBatch>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        match self.input.recv() {
            Ok(Some(batch)) => Poll::Ready(Some(batch)),
            // End of results from one input partition
            Ok(None) => {
                self.num_input_partitions_processed += 1;
                if self.num_input_partitions == self.num_input_partitions_processed {
                    // all input partitions have finished sending batches
                    Poll::Ready(None)
                } else {
                    // other partitions still have data to send
                    self.poll_next(cx)
                }
            }
            // RecvError means receiver has exited and closed the channel
            Err(_) => Poll::Ready(None),
        }
    }
}

impl RecordBatchStream for RepartitionStream {
    /// Get the schema
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::physical_plan::memory::MemoryExec;
    use arrow::array::UInt32Array;
    use arrow::datatypes::{DataType, Field, Schema};
    use arrow::record_batch::RecordBatch;

    #[tokio::test(flavor = "multi_thread")]
    async fn one_to_many_round_robin() -> Result<()> {
        // define input partitions
        let schema = test_schema();
        let partition = create_vec_batches(&schema, 50);
        let partitions = vec![partition];

        // repartition from 1 input to 4 output
        let output_partitions =
            repartition(&schema, partitions, Partitioning::RoundRobinBatch(4)).await?;

        assert_eq!(4, output_partitions.len());
        assert_eq!(13, output_partitions[0].len());
        assert_eq!(13, output_partitions[1].len());
        assert_eq!(12, output_partitions[2].len());
        assert_eq!(12, output_partitions[3].len());

        Ok(())
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn many_to_one_round_robin() -> Result<()> {
        // define input partitions
        let schema = test_schema();
        let partition = create_vec_batches(&schema, 50);
        let partitions = vec![partition.clone(), partition.clone(), partition.clone()];

        // repartition from 3 input to 1 output
        let output_partitions =
            repartition(&schema, partitions, Partitioning::RoundRobinBatch(1)).await?;

        assert_eq!(1, output_partitions.len());
        assert_eq!(150, output_partitions[0].len());

        Ok(())
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn many_to_many_round_robin() -> Result<()> {
        // define input partitions
        let schema = test_schema();
        let partition = create_vec_batches(&schema, 50);
        let partitions = vec![partition.clone(), partition.clone(), partition.clone()];

        // repartition from 3 input to 5 output
        let output_partitions =
            repartition(&schema, partitions, Partitioning::RoundRobinBatch(5)).await?;

        assert_eq!(5, output_partitions.len());
        assert_eq!(30, output_partitions[0].len());
        assert_eq!(30, output_partitions[1].len());
        assert_eq!(30, output_partitions[2].len());
        assert_eq!(30, output_partitions[3].len());
        assert_eq!(30, output_partitions[4].len());

        Ok(())
    }

    fn test_schema() -> Arc<Schema> {
        Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]))
    }

    fn create_vec_batches(schema: &Arc<Schema>, n: usize) -> Vec<RecordBatch> {
        let batch = create_batch(schema);
        let mut vec = Vec::with_capacity(n);
        for _ in 0..n {
            vec.push(batch.clone());
        }
        vec
    }

    fn create_batch(schema: &Arc<Schema>) -> RecordBatch {
        RecordBatch::try_new(
            schema.clone(),
            vec![Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]))],
        )
        .unwrap()
    }

    async fn repartition(
        schema: &SchemaRef,
        input_partitions: Vec<Vec<RecordBatch>>,
        partitioning: Partitioning,
    ) -> Result<Vec<Vec<RecordBatch>>> {
        // create physical plan
        let exec = MemoryExec::try_new(&input_partitions, schema.clone(), None)?;
        let exec = RepartitionExec::try_new(Arc::new(exec), partitioning)?;

        // execute and collect results
        let mut output_partitions = vec![];
        for i in 0..exec.partitioning.partition_count() {
            // execute this *output* partition and collect all batches
            let mut stream = exec.execute(i).await?;
            let mut batches = vec![];
            while let Some(result) = stream.next().await {
                batches.push(result?);
            }
            output_partitions.push(batches);
        }
        Ok(output_partitions)
    }
}
<|endoftext|>