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

use {Error, Result};

/// The outline of a shape.
///
/// http://www.w3.org/TR/SVG/paths.html
pub struct Path;

/// The data attribute of a path.
///
/// http://www.w3.org/TR/SVG/paths.html#PathData
pub struct Data {
    /// A series of path commands.
    pub 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 {
    #[inline]
    pub fn parse(line: &str) -> Result<Data> {
        DataParser::new(line).process()
    }
}

struct DataParser<'s> {
    line: &'s str,
    position: usize,
    cursor: Peekable<CharIndices<'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.char_indices().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> {
        match self.cursor.peek() {
            Some(&(_, c)) => Some(c),
            _ => None,
        }
    }
}

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

    #[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 a test for path::Data::parse<commit_after>use std::iter::Peekable;
use std::str::CharIndices;

use {Error, Result};

/// The outline of a shape.
///
/// http://www.w3.org/TR/SVG/paths.html
pub struct Path;

/// The data attribute of a path.
///
/// http://www.w3.org/TR/SVG/paths.html#PathData
pub struct Data {
    /// A series of path commands.
    pub 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 {
    #[inline]
    pub fn parse(line: &str) -> Result<Data> {
        DataParser::new(line).process()
    }
}

struct DataParser<'s> {
    line: &'s str,
    position: usize,
    cursor: Peekable<CharIndices<'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.char_indices().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> {
        match self.cursor.peek() {
            Some(&(_, c)) => Some(c),
            _ => None,
        }
    }
}

#[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);
    }
}
<|endoftext|>


<commit_before>//! track handles parsing of GPX-spec tracks.

use errors::*;
use std::io::Read;
use xml::reader::XmlEvent;

use parser::tracksegment;
use parser::string;
use parser::Context;

use Track;

/// consume consumes a GPX track from the `reader` until it ends.
pub fn consume<R: Read>(context: &mut Context<R>) -> Result<Track> {
    let mut track: Track = Default::default();

    while let Some(event) = context.reader.next() {
        match event.chain_err(|| "error while parsing XML")? {
            XmlEvent::StartElement { name, .. } => match name.local_name.as_ref() {
                "trk" => {}
                "name" => track.name = Some(string::consume(context)?),
                "cmt" => track.comment = Some(string::consume(context)?),
                "desc" => track.description = Some(string::consume(context)?),
                "src" => track.source = Some(string::consume(context)?),
                "type" => track._type = Some(string::consume(context)?),
                "trkseg" => track.segments.push(tracksegment::consume(context)?),
                child => Err(Error::from(ErrorKind::InvalidChildElement(
                    String::from(child),
                    "track",
                )))?,
            },

            XmlEvent::EndElement { .. } => {
                return Ok(track);
            }

            _ => {}
        }
    }

    return Err("no end tag for track".into());
}

#[cfg(test)]
mod tests {
    use std::io::BufReader;
    use xml::reader::EventReader;

    use GpxVersion;
    use parser::Context;
    use super::consume;

    #[test]
    fn consume_full_track() {
        let track = consume!(
            "
            <trk>
                <name>track name</name>
                <cmt>track comment</cmt>
                <desc>track description</desc>
                <src>track source</src>
                <type>track type</type>
            </trk>
            ",
            GpxVersion::Gpx11
        );

        assert!(track.is_ok());

        let track = track.unwrap();

        assert_eq!(track.name.unwrap(), "track name");
        assert_eq!(track.comment.unwrap(), "track comment");
        assert_eq!(track.description.unwrap(), "track description");
        assert_eq!(track.source.unwrap(), "track source");
        assert_eq!(track._type.unwrap(), "track type");
    }

    #[test]
    fn consume_empty() {
        let track = consume!("<trk></trk>", GpxVersion::Gpx11);

        assert!(track.is_ok());
    }
}
<commit_msg>Rewrite track::consume cleaner and more strict - Do not accept arbitrary amounts of opening trk tags, we require exactly   one at the beginning and do not allow further ones - Require the closing tag to match the opening tag - Use reader.peek() instead of reader.next() to not consume the opening   tag for deeper parser functions - Use new custom errors InvalidClosingTag and MissingClosingTag<commit_after>//! track handles parsing of GPX-spec tracks.

use errors::*;
use std::io::Read;
use xml::reader::XmlEvent;

use parser::tracksegment;
use parser::string;
use parser::Context;
use parser::verify_starting_tag;

use Track;

/// consume consumes a GPX track from the `reader` until it ends.
pub fn consume<R: Read>(context: &mut Context<R>) -> Result<Track> {
    let mut track: Track = Default::default();
    verify_starting_tag(context, "trk")?;

    loop {
        let next_event = {
            if let Some(next) = context.reader.peek() {
                next.clone()
            } else {
                break;
            }
        };

        match next_event.chain_err(|| Error::from("error while parsing track event"))? {
            XmlEvent::StartElement { ref name, .. } => match name.local_name.as_ref() {
                "name" => {
                    track.name = Some(string::consume(context)?);
                }
                "cmt" => {
                    track.comment = Some(string::consume(context)?);
                }
                "desc" => {
                    track.description = Some(string::consume(context)?);
                }
                "src" => {
                    track.source = Some(string::consume(context)?);
                }
                "type" => {
                    track._type = Some(string::consume(context)?);
                }
                "trkseg" => {
                    track.segments.push(tracksegment::consume(context)?);
                }
                child => {
                    bail!(ErrorKind::InvalidChildElement(String::from(child), "track"));
                }
            },
            XmlEvent::EndElement { ref name } => {
                ensure!(
                    name.local_name == "trk",
                    ErrorKind::InvalidClosingTag(name.local_name.clone(), "track")
                );
                context.reader.next(); //consume the end tag
                return Ok(track);
            }
            _ => {
                context.reader.next(); //consume and ignore this event
            }
        }
    }

    bail!(ErrorKind::MissingClosingTag("track"));
}

#[cfg(test)]
mod tests {
    use std::io::BufReader;
    use xml::reader::EventReader;

    use GpxVersion;
    use parser::Context;
    use super::consume;

    #[test]
    fn consume_full_track() {
        let track = consume!(
            "
            <trk>
                <name>track name</name>
                <cmt>track comment</cmt>
                <desc>track description</desc>
                <src>track source</src>
                <type>track type</type>
            </trk>
            ",
            GpxVersion::Gpx11
        );

        assert!(track.is_ok());

        let track = track.unwrap();

        assert_eq!(track.name.unwrap(), "track name");
        assert_eq!(track.comment.unwrap(), "track comment");
        assert_eq!(track.description.unwrap(), "track description");
        assert_eq!(track.source.unwrap(), "track source");
        assert_eq!(track._type.unwrap(), "track type");
    }

    #[test]
    fn consume_empty() {
        let track = consume!("<trk></trk>", GpxVersion::Gpx11);
        assert!(track.is_ok());
    }
}
<|endoftext|>


<commit_before>use byteorder::{LittleEndian, ReadBytesExt};
use std::fs::File;
use std::io::{Read, Result, Seek, SeekFrom};
use std::path::Path;
use std::str::from_utf8;

pub enum Flags {
    BackgroundFlag = 0x00000001,
    WeaponFlag = 0x00000002,
    TransparentFlag = 0x00000004,
    TranslucentFlag = 0x00000008,
    WallFlag = 0x00000010,
    DamagingFlag = 0x00000020,
    ImpassableFlag = 0x00000040,
    WetFlag = 0x00000080,
    UnknownFlag = 0x00000100,
    SurfaceFlag = 0x00000200,
    BridgeFlag = 0x00000400,
    StackableFlag = 0x00000800,
    WindowFlag = 0x00001000,
    NoShootFlag = 0x00002000,
    PrefixAFlag = 0x00004000,
    PrefixAnFlag = 0x00008000,
    InternalFlag = 0x00010000,
    FoliageFlag = 0x00020000,
    PartialHueFlag = 0x00040000,
    Unknown1Flag = 0x00080000,
    MapFlag = 0x00100000,
    ContainerFlag = 0x00200000,
    WearableFlag = 0x00400000,
    LightSourceFlag = 0x00800000,
    AnimatedFlag = 0x01000000,
    NoDiagonalFlag = 0x02000000,
    Unknown2Flag = 0x04000000,
    ArmorFlag = 0x08000000,
    RoofFlag = 0x10000000,
    DoorFlag = 0x20000000,
    StairBackFlag = 0x40000000,
    StairRightFlag = 0x80000000,
}

// Tile data is odd, as we have [(unknown, (LAND_TILE_DATA) *32) * 512]
static GROUP_HEADER_SIZE: u32 = 4;
static MAP_TILE_SIZE: u32 = 26;
static STATIC_TILE_SIZE: u32 = 37;
static STATIC_OFFSET: u32 = 428032;

#[derive(Clone, Debug)]
pub struct MapTileData {
    pub flags: u32,
    pub texture_id: u16,
    pub name: String,
}

#[derive(Clone, Debug)]
pub struct StaticTileData {
    pub flags: u32,
    pub weight: u8,
    pub quality_layer_light_id: u8,
    pub quantity_weapon_class_armor_class: u8,
    pub anim_id: u16,
    pub hue: u8,
    pub height_capacity: u8,
    pub name: String,
}

pub struct TileDataReader<T: Read + Seek> {
    data_reader: T,
}

impl TileDataReader<File> {
    pub fn new(mul_path: &Path) -> Result<TileDataReader<File>> {
        let data_reader = File::open(mul_path)?;

        Ok(TileDataReader {
            data_reader: data_reader,
        })
    }
}

impl<T: Read + Seek> TileDataReader<T> {
    pub fn read_map_tile_data(&mut self, idx: u32) -> Result<MapTileData> {
        let offset = self.calculate_map_tile_offset(idx);
        self.data_reader.seek(SeekFrom::Start(offset))?;
        let flags = self.data_reader.read_u32::<LittleEndian>()?;
        let texture_id = self.data_reader.read_u16::<LittleEndian>()?;

        let mut raw_name = vec![];
        while match raw_name.last() {
            Some(0) => false,
            _ => true,
        } {
            raw_name.push(self.data_reader.read_u8()?);
        }

        Ok(MapTileData {
            flags: flags,
            texture_id: texture_id,
            name: String::from(from_utf8(&raw_name).unwrap_or("ERROR")),
        })
    }

    fn calculate_map_tile_offset(&self, idx: u32) -> u64 {
        //For every 32, we have to add an unknown header
        let group_header_jumps = ((idx / 32) + 1) * GROUP_HEADER_SIZE;
        ((idx * MAP_TILE_SIZE) + group_header_jumps) as u64
    }

    pub fn read_static_tile_data(&mut self, idx: u32) -> Result<StaticTileData> {
        let offset = self.calculate_static_tile_offset(idx);
        self.data_reader.seek(SeekFrom::Start(offset))?;

        let flags = self.data_reader.read_u32::<LittleEndian>()?;
        let weight = self.data_reader.read_u8()?;
        let quality = self.data_reader.read_u8()?;
        let _unknown = self.data_reader.read_u16::<LittleEndian>()?;
        let _unknown1 = self.data_reader.read_u8()?;
        let quantity = self.data_reader.read_u8()?;
        let anim_id = self.data_reader.read_u16::<LittleEndian>()?;
        let _unknown2 = self.data_reader.read_u8()?;
        let hue = self.data_reader.read_u8()?;
        let _unknown3 = self.data_reader.read_u16::<LittleEndian>()?;
        let height = self.data_reader.read_u8()?;

        let mut raw_name = vec![];
        while match raw_name.last() {
            Some(0) => false,
            _ => true,
        } {
            raw_name.push(self.data_reader.read_u8()?);
        }

        Ok(StaticTileData {
            flags: flags,
            weight: weight,
            quality_layer_light_id: quality,
            quantity_weapon_class_armor_class: quantity,
            anim_id: anim_id,
            hue: hue,
            height_capacity: height,
            name: String::from(from_utf8(&raw_name).unwrap_or("ERROR")),
        })
    }

    fn calculate_static_tile_offset(&self, idx: u32) -> u64 {
        //For every 32, we have to add an unknown header
        let group_header_jumps = ((idx / 32) + 1) * GROUP_HEADER_SIZE;
        ((idx * STATIC_TILE_SIZE) + group_header_jumps + STATIC_OFFSET) as u64
    }
}
<commit_msg>Don't read in the null terminator<commit_after>use byteorder::{LittleEndian, ReadBytesExt};
use std::fs::File;
use std::io::{Read, Result, Seek, SeekFrom};
use std::path::Path;
use std::str::from_utf8;

pub enum Flags {
    BackgroundFlag = 0x00000001,
    WeaponFlag = 0x00000002,
    TransparentFlag = 0x00000004,
    TranslucentFlag = 0x00000008,
    WallFlag = 0x00000010,
    DamagingFlag = 0x00000020,
    ImpassableFlag = 0x00000040,
    WetFlag = 0x00000080,
    UnknownFlag = 0x00000100,
    SurfaceFlag = 0x00000200,
    BridgeFlag = 0x00000400,
    StackableFlag = 0x00000800,
    WindowFlag = 0x00001000,
    NoShootFlag = 0x00002000,
    PrefixAFlag = 0x00004000,
    PrefixAnFlag = 0x00008000,
    InternalFlag = 0x00010000,
    FoliageFlag = 0x00020000,
    PartialHueFlag = 0x00040000,
    Unknown1Flag = 0x00080000,
    MapFlag = 0x00100000,
    ContainerFlag = 0x00200000,
    WearableFlag = 0x00400000,
    LightSourceFlag = 0x00800000,
    AnimatedFlag = 0x01000000,
    NoDiagonalFlag = 0x02000000,
    Unknown2Flag = 0x04000000,
    ArmorFlag = 0x08000000,
    RoofFlag = 0x10000000,
    DoorFlag = 0x20000000,
    StairBackFlag = 0x40000000,
    StairRightFlag = 0x80000000,
}

// Tile data is odd, as we have [(unknown, (LAND_TILE_DATA) *32) * 512]
static GROUP_HEADER_SIZE: u32 = 4;
static MAP_TILE_SIZE: u32 = 26;
static STATIC_TILE_SIZE: u32 = 37;
static STATIC_OFFSET: u32 = 428032;

#[derive(Clone, Debug)]
pub struct MapTileData {
    pub flags: u32,
    pub texture_id: u16,
    pub name: String,
}

#[derive(Clone, Debug)]
pub struct StaticTileData {
    pub flags: u32,
    pub weight: u8,
    pub quality_layer_light_id: u8,
    pub quantity_weapon_class_armor_class: u8,
    pub anim_id: u16,
    pub hue: u8,
    pub height_capacity: u8,
    pub name: String,
}

pub struct TileDataReader<T: Read + Seek> {
    data_reader: T,
}

impl TileDataReader<File> {
    pub fn new(mul_path: &Path) -> Result<TileDataReader<File>> {
        let data_reader = File::open(mul_path)?;

        Ok(TileDataReader {
            data_reader: data_reader,
        })
    }
}

impl<T: Read + Seek> TileDataReader<T> {
    pub fn read_map_tile_data(&mut self, idx: u32) -> Result<MapTileData> {
        let offset = self.calculate_map_tile_offset(idx);
        self.data_reader.seek(SeekFrom::Start(offset))?;
        let flags = self.data_reader.read_u32::<LittleEndian>()?;
        let texture_id = self.data_reader.read_u16::<LittleEndian>()?;

        let mut raw_name = vec![];
        loop {
            match self.data_reader.read_u8()? {
                0 => break,
                x => raw_name.push(x)
            }
        }

        Ok(MapTileData {
            flags: flags,
            texture_id: texture_id,
            name: String::from(from_utf8(&raw_name).unwrap_or("ERROR")),
        })
    }

    fn calculate_map_tile_offset(&self, idx: u32) -> u64 {
        //For every 32, we have to add an unknown header
        let group_header_jumps = ((idx / 32) + 1) * GROUP_HEADER_SIZE;
        ((idx * MAP_TILE_SIZE) + group_header_jumps) as u64
    }

    pub fn read_static_tile_data(&mut self, idx: u32) -> Result<StaticTileData> {
        let offset = self.calculate_static_tile_offset(idx);
        self.data_reader.seek(SeekFrom::Start(offset))?;

        let flags = self.data_reader.read_u32::<LittleEndian>()?;
        let weight = self.data_reader.read_u8()?;
        let quality = self.data_reader.read_u8()?;
        let _unknown = self.data_reader.read_u16::<LittleEndian>()?;
        let _unknown1 = self.data_reader.read_u8()?;
        let quantity = self.data_reader.read_u8()?;
        let anim_id = self.data_reader.read_u16::<LittleEndian>()?;
        let _unknown2 = self.data_reader.read_u8()?;
        let hue = self.data_reader.read_u8()?;
        let _unknown3 = self.data_reader.read_u16::<LittleEndian>()?;
        let height = self.data_reader.read_u8()?;

        let mut raw_name = vec![];
        loop {
            match self.data_reader.read_u8()? {
                0 => break,
                x => raw_name.push(x)
            }
        }

        Ok(StaticTileData {
            flags: flags,
            weight: weight,
            quality_layer_light_id: quality,
            quantity_weapon_class_armor_class: quantity,
            anim_id: anim_id,
            hue: hue,
            height_capacity: height,
            name: String::from(from_utf8(&raw_name).unwrap_or("ERROR")),
        })
    }

    fn calculate_static_tile_offset(&self, idx: u32) -> u64 {
        //For every 32, we have to add an unknown header
        let group_header_jumps = ((idx / 32) + 1) * GROUP_HEADER_SIZE;
        ((idx * STATIC_TILE_SIZE) + group_header_jumps + STATIC_OFFSET) as u64
    }
}
<|endoftext|>


<commit_before>use libc::system;
use std::ffi::CString;

static TMUX_NAME: &'static str = "tmux";
static NEW_SESSION: &'static str = "new -s";

pub fn open(session_name: String) {
    let system_call = CString::new(format!("{} {} {}", TMUX_NAME, NEW_SESSION, session_name)).unwrap();
    unsafe { system(system_call.as_ptr()); };
}
<commit_msg>Send all methods through a common call interface.<commit_after>use libc::system;
use std::ffi::CString;

static TMUX_NAME: &'static str = "tmux";

fn call(command: String) {
    let system_call = CString::new(format!("{} {}", TMUX_NAME, command)).unwrap();
    unsafe { system(system_call.as_ptr()); };
}

pub fn open(session_name: String) {
    let open_command = format!("new -s {} -n {}", session_name, "hello");
    call(open_command);
}

static SELECT_LAYOUT: &'static str = "select-layout";
pub fn select_layout(window: String, layout: String) {
    let command = format!("{} -t {} {}", SELECT_LAYOUT, window, layout);
    call(command);
}
<|endoftext|>


<commit_before>#![feature(macro_rules)]

extern crate rustbox;

use std::collections::dlist::DList;
use std::comm::{Receiver, Sender};
use std::io::{File, BufferedReader};
use std::str::from_utf8;

pub struct Editor {
    pub sender: Sender<rustbox::Event>,
    pub events: Receiver<rustbox::Event>,
    pub buffers: Vec<Buf>
}

pub struct Buf {
    pub lines: DList<Line>,
    pub active: bool,
    pub num_lines: int,
}

#[deriving(Clone)]
pub struct Line {
    pub data: String,
}

pub enum Response {
    Continue,
    Quit,
}

impl Buf {
    pub fn new() -> Buf {
        Buf {
            lines: DList::new(),
            active: false,
            num_lines: 0,
        }
    }

    pub fn new_from_file(filename: String) -> Buf {
        let path = Path::new(filename.to_string());

        let mut new_buffer = Buf::new();
        let mut file = BufferedReader::new(File::open(&path));
        let lines: Vec<String> = file.lines().map(|x| x.unwrap()).collect();

        for line in lines.iter() {
            new_buffer.lines.push(Line{data: line.clone()})
        }

        new_buffer
    }
}

impl Line {
    pub fn new() -> Line {
        Line {
            data: String::new(),
        }
    }
}

impl Editor {
    pub fn new(filenames: &[String]) -> Editor {
        let mut buffers = Vec::new();

        for filename in filenames.iter() {
            let mut b = Buf::new_from_file(filename.clone());
            b.active = true;
            buffers.push(b);
        }

        let (send, recv) = channel();
        Editor {
            sender: send,
            events: recv,
            buffers: buffers,
        }
    }

    pub fn handle_key_event(&self, ch: u32) -> Response {
        match std::char::from_u32(ch) {
            Some('q') => Quit,
            _ => Continue,
        }
    }

    pub fn draw(&mut self) {
        for buffer in self.buffers.iter() {
            for (index, line) in buffer.lines.iter().enumerate() {
                rustbox::print(1, index, rustbox::Bold, rustbox::White, rustbox::Black, line.data.clone());
            }
        }
    }

    pub fn start(&mut self) -> bool {
        loop {
            self.draw();
            rustbox::present();
            let status = match self.events.recv() {
                rustbox::KeyEvent(_, _, ch) => {
                    let status = self.handle_key_event(ch);
                    match status {
                        Quit => { false },
                        _ => { true },
                    }
                },
                _ => { true }
            };
            if status == false {
                return false;
            }
        }
    }

}

<commit_msg>Draw initial cursor position<commit_after>#![feature(macro_rules)]

extern crate rustbox;

use std::collections::dlist::DList;
use std::comm::{Receiver, Sender};
use std::io::{File, BufferedReader};
use std::str::from_utf8;

pub struct Editor {
    pub sender: Sender<rustbox::Event>,
    pub events: Receiver<rustbox::Event>,
    pub buffers: Vec<Buf>
}

pub struct Buf {
    pub lines: DList<Line>,
    pub active: bool,
    pub num_lines: int,
}

#[deriving(Clone)]
pub struct Line {
    pub data: String,
}

pub enum Response {
    Continue,
    Quit,
}

impl Buf {
    pub fn new() -> Buf {
        Buf {
            lines: DList::new(),
            active: false,
            num_lines: 0,
        }
    }

    pub fn new_from_file(filename: String) -> Buf {
        let path = Path::new(filename.to_string());

        let mut new_buffer = Buf::new();
        let mut file = BufferedReader::new(File::open(&path));
        let lines: Vec<String> = file.lines().map(|x| x.unwrap()).collect();

        for line in lines.iter() {
            new_buffer.lines.push(Line{data: line.clone()})
        }

        new_buffer
    }
}

impl Line {
    pub fn new() -> Line {
        Line {
            data: String::new(),
        }
    }
}

impl Editor {
    pub fn new(filenames: &[String]) -> Editor {
        let mut buffers = Vec::new();

        for filename in filenames.iter() {
            let mut b = Buf::new_from_file(filename.clone());
            b.active = true;
            buffers.push(b);
        }

        let (send, recv) = channel();
        Editor {
            sender: send,
            events: recv,
            buffers: buffers,
        }
    }

    pub fn handle_key_event(&self, ch: u32) -> Response {
        match std::char::from_u32(ch) {
            Some('q') => Quit,
            _ => Continue,
        }
    }

    pub fn draw(&mut self) {
        for buffer in self.buffers.iter() {
            if buffer.active {
                for (index, line) in buffer.lines.iter().enumerate() {
                    rustbox::print(1, index, rustbox::Bold, rustbox::White, rustbox::Black, line.data.clone());
                }
            }
        }
        rustbox::set_cursor(0, 1);
    }

    pub fn start(&mut self) -> bool {
        loop {
            self.draw();
            rustbox::present();
            let status = match self.events.recv() {
                rustbox::KeyEvent(_, _, ch) => {
                    let status = self.handle_key_event(ch);
                    match status {
                        Quit => { false },
                        _ => { true },
                    }
                },
                _ => { true }
            };
            if status == false {
                return false;
            }
        }
    }

}

<|endoftext|>


<commit_before>#![recursion_limit="128"]
#![cfg_attr(feature = "cargo-clippy", allow(cyclomatic_complexity, trivial_regex))]
use std::str;
use std::vec::Vec;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::io::stdout;
use std::ops::Range;

#[macro_use]
extern crate failure;

extern crate bam2bedgraph;
use bam2bedgraph::*;
use bam2bedgraph::error::*;
use bam2bedgraph::indexed_annotation::*;

extern crate rust_htslib;
use rust_htslib::bam::Read;
use rust_htslib::bam::IndexedReader;

extern crate structopt;
#[macro_use]
extern crate structopt_derive;
use structopt::StructOpt;

extern crate futures;
use futures::Future;
extern crate futures_cpupool;
use futures_cpupool::CpuPool;

extern crate num_cpus;

extern crate ordered_float;
use ordered_float::OrderedFloat;

#[derive(StructOpt, Debug)]
#[structopt(name = "intronrpkm", about = "Analyze RPKM values in intronic space")]
struct Options {
    // input files
    #[structopt(long="gff", help = "A genome annotation file in gff3 format", name="ANNOT_GFF_FILE")]
    annotfile_gff: Option<String>,
    #[structopt(long="gtf", help = "A genome annotation file in gtf format", name="ANNOT_GTF_FILE")]
    annotfile_gtf: Option<String>,
    #[structopt(long="bam1", short="1", help = "The set of stranded .bam files to analyze where read1 indicates strand", name="BAMFILE1")]
    bam1: Vec<String>,
    #[structopt(long="bam2", short="2", help = "The set of stranded .bam files to analyze where read2 indicates strand", name="BAMFILE2")]
    bam2: Vec<String>,
    #[structopt(long="bam", short="u", help = "The set of unstranded .bam files to analyze", name="BAMFILE")]
    bam: Vec<String>,
    #[structopt(long="chrmap", help = "Optional tab-delimited chr name mapping file", name="CHRMAP_FILE")]
    chrmap_file: Option<String>,
    #[structopt(long="vizchrmap", help = "Optional tab-delimited chr name mapping file for bigwig/bigbed exports", name="VIZCHRMAP_FILE")]
    vizchrmap_file: Option<String>,
    #[structopt(long="sizes", help = "Optional chr sizes file", name="SIZES_FILE")]
    sizes_file: Option<String>,
    // output file
    #[structopt(long="out", short="o", help = "Output file", name="OUT_FILE", default_value="-")]
    outfile: String,
    // feature types filter
    #[structopt(long="exon_type", help = "The exon type(s) to search for", name="EXON_TYPE")]
    exon_type: Vec<String>,
    #[structopt(long="transcript_type", help = "The transcript type(s) to search for", name="TRANSCRIPT_TYPE")]
    transcript_type: Vec<String>,
    #[structopt(long="gene_type", help = "The gene type(s) to search for", name="GENE_TYPE")]
    gene_type: Vec<String>,
    // flags
    #[structopt(long="cpu_threads", short="t", help = "How many threads to use for processing", default_value="0")]
    cpu_threads: usize,
}

#[derive(Ord, Eq, PartialOrd, PartialEq)]
struct Row {
    seqname: String,
    strand: String,
    start: u64,
    end: u64,
    cov: OrderedFloat<f64>,
    rpkm: OrderedFloat<f64>,
}

fn write_exon_cov(
    options: &Options,
    annot: &Arc<IndexedAnnotation>,
    total_reads: u64,
    bamfiles: &Vec<String>,
    bamstrand: &Vec<Option<bool>>)
    -> Result<()>
{
    let mut tidmaps = HashMap::<String,HashMap<String,u32>>::new();
    for bamfile in bamfiles {
        let bam = IndexedReader::from_path(bamfile)?;
        // build the tid map for this bam file
        let mut tidmap = HashMap::<String,u32>::new();
        {   let header = bam.header();
            for target_name in header.target_names() {
                let tid = header.tid(target_name).r()?;
                let target_name = String::from(std::str::from_utf8(target_name)?);
                let chr = annot.chrmap.get(&target_name).unwrap_or(&target_name);
                tidmap.insert(chr.clone(), tid);
            }
        }
        tidmaps.insert(bamfile.clone(), tidmap);
    }
    let tidmaps = Arc::new(tidmaps);

    let mut unmerged_exons = HashMap::<(String,String),Vec<Range<u64>>>::new();

    for (gene_row, gene) in annot.rows.iter().enumerate() {
        if options.gene_type.is_empty() || options.gene_type.contains(&gene.feature_type) {
            if let Some(feature_rows) = annot.row2children.get(&gene_row) {
                // get the transcript rows for this gene
                'transcript_row:
                for transcript_row in feature_rows {
                    let transcript = &annot.rows[*transcript_row];
                    // make sure this is a transcript type
                    // make sure transcript coordinates are consistent with the gene coordinates
                    if (options.transcript_type.is_empty() ||
                            options.transcript_type.contains(&transcript.feature_type)) &&
                        transcript.seqname == gene.seqname &&
                        transcript.strand == gene.strand
                    {
                        if let Some(exon_rows) = annot.row2children.get(transcript_row) {
                            // first get exon start/stop -> transcript associations
                            // and splice start/stop -> transcript associations
                            for exon_row in exon_rows {
                                let exon = &annot.rows[*exon_row];
                                if options.exon_type.is_empty() || options.exon_type.contains(&exon.feature_type) {
                                    let exon = &annot.rows[*exon_row];
                                    unmerged_exons.entry((exon.seqname.clone(),exon.strand.clone())).
                                        or_insert_with(Vec::new).push(exon.start-1..exon.end);
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    let mut merged_exons = HashMap::<(String,String),Vec<Range<u64>>>::new();
    let keys = unmerged_exons.keys().map(|k| k.clone()).collect::<Vec<_>>();
    for key in &keys {
        if let Some(ref mut unmerged) = unmerged_exons.get_mut(&key) {
            unmerged.sort_by_key(|u| u.start);
            let merged = merged_exons.entry((*key).clone()).or_insert_with(Vec::new);
            for exon in unmerged.iter() {
                let mut extend = false;
                if let Some(ref mut last) = merged.last_mut() {
                    last.end = exon.end;
                    extend = true;
                }
                if !extend {
                    merged.push(exon.clone());
                }
            }
        }
    }

    let num_cpus = num_cpus::get();
    let mut pair_futures = Vec::new();
    let pool = Arc::new(CpuPool::new(if options.cpu_threads==0 {num_cpus} else {options.cpu_threads}));
    for key in &keys {
        let exons = merged_exons[&key.clone()].clone();
        for exon in exons {
            let seqname = key.0.clone();
            let strand = key.1.clone();
            let chr = seqname.clone();
            let strand_is_plus = strand == "+";
            let bamfiles = bamfiles.clone();
            let bamstrand = bamstrand.clone();
            let tidmaps = tidmaps.clone();

            let pair_future = pool.spawn_fn(move ||->Result<Row> {
                let mut exon_cov = 0f64;
                let mut exon_reads = HashSet::<String>::new();
                for (i,bamfile) in bamfiles.iter().enumerate() {
                    let read1strand = bamstrand[i];
                    let tidmap = &tidmaps[bamfile];
                    if let Some(tid) = tidmap.get(&chr) {
                        let mut bam = IndexedReader::from_path(bamfile)?;
                        bam.fetch(*tid, exon.start as u32, exon.end as u32)?;
                        for read in bam.records() {
                            let read = read?;
                            // make sure the read's strand matches
                            if let Some(is_read1strand) = read1strand {
                                if !(((is_read1strand == read.is_first_in_template()) == !read.is_reverse()) == strand_is_plus) {
                                    continue;
                                }
                            }
                            let read_name = String::from(str::from_utf8(read.qname())?);
                            exon_reads.insert(read_name);
                            //eprintln!("Looking at read: {}", read_name);
                            let exons = cigar2exons(&read.cigar(), read.pos() as u64)?;
                            for e in exons {
                                if e.start < exon.end && exon.start < e.end {
                                    exon_cov +=
                                        (std::cmp::min(e.end, exon.end) -
                                        std::cmp::max(e.start, exon.start)) as f64;
                                }
                            }
                        }
                    }
                }
                let exon_length = exon.end-exon.start;
                let rpkm = (1e10f64 * exon_reads.len() as f64) / (total_reads as f64 * exon_length as f64);
                Ok(Row{
                    seqname: seqname,
                    strand: strand,
                    start: exon.start,
                    end: exon.end,
                    cov: OrderedFloat((exon_cov / exon_length as f64)),
                    rpkm: OrderedFloat(rpkm),
                })
            });
            pair_futures.push(pair_future);
        }
    }

    let mut rows = Vec::new();
    for future in pair_futures {
        match future.wait() {
            Ok(row) => {
                rows.push(row)
            }
            Err(ref e) => {
                eprintln!("Got Err in write_exon_cov: {:?}", e);
            }
        }
    }
    rows.sort();

    let mut output: BufWriter<Box<Write>> = BufWriter::new(
        if options.outfile == "-" { Box::new(stdout()) }
            else { Box::new(File::create(&options.outfile)?) });

    output.write_fmt(format_args!("{}\t{}\t{}\t{}\t{}\t{}\n",
                                  "seqname",
                                  "strand",
                                  "start",
                                  "end",
                                  "cov",
                                  "rpkm",
    ))?;
    for row in rows {
        output.write_fmt(format_args!("{}\t{}\t{}\t{}\t{}\t{}\n",
            row.seqname,
            row.strand,
            row.start,
            row.end,
            row.cov,
            row.rpkm,
        ))?;
    }
    Ok(())
}

fn run() -> Result<()> {
    let mut options = Options::from_args();
    // set defaults for feature types
    options.gene_type =
        (if options.gene_type.is_empty() { vec!["gene".to_string()] }
         else { options.gene_type.clone() }).into_iter().collect();
    //options.transcript_type =
    //    (if options.transcript_type.is_empty() { vec!["transcript".to_string()] }
    //     else { options.transcript_type.clone() }).into_iter().collect();
    options.exon_type =
        (if options.exon_type.is_empty() { vec!["exon".to_string()] }
         else { options.exon_type.clone() }).into_iter().collect();
    // set debug options if --debug flag is set

    // organize bamfiles by read strand type
    let mut bamfiles = Vec::<String>::new();
    bamfiles.append(&mut options.bam1.clone());
    bamfiles.append(&mut options.bam2.clone());
    bamfiles.append(&mut options.bam.clone());
    let mut bamstrand = Vec::<Option<bool>>::new();
    bamstrand.append(&mut options.bam1.iter().map(|_| Some(true)).collect());
    bamstrand.append(&mut options.bam2.iter().map(|_| Some(false)).collect());
    bamstrand.append(&mut options.bam.iter().map(|_| None).collect());
    // get the chromosome names and sizes from the first bam file
    if bamfiles.is_empty() {
        Options::clap().print_help()?;
        eprintln!("\n\nNo bam files were passed in!");
        std::process::exit(1)
    }
    let transcript_type = String::from("transcript");
    let mut annot = if let Some(annotfile_gff) = options.annotfile_gff.clone() {
        eprintln!("Reading annotation file {:?}", &annotfile_gff);
        IndexedAnnotation::from_gff(
            &annotfile_gff,
            &options.chrmap_file,
            &options.vizchrmap_file)?
    } else if let Some(annotfile_gtf) = options.annotfile_gtf.clone() {
        eprintln!("Reading annotation file {:?}", &annotfile_gtf);
        IndexedAnnotation::from_gtf(&annotfile_gtf,
            options.gene_type.get(0).r()?,
            options.transcript_type.get(0).unwrap_or(&transcript_type),
            &options.chrmap_file,
            &options.vizchrmap_file)?
    } else {
        Options::clap().print_help()?;
        eprintln!("\n\nNo annotation file was given!");
        std::process::exit(1);
    };
    eprintln!("Getting refseq lengths from bam file {:?}", &bamfiles[0]);
    let refs = match options.sizes_file.clone() {
        Some(sizes_file) => read_sizes_file(&sizes_file, &annot.chrmap)?,
        None => get_bam_refs(&bamfiles[0], &annot.chrmap)?,
    };
    annot.refs = refs;

    // get the total bam reads
    eprintln!("Running samtools idxstats to get total bam read counts");
    let total_reads = get_bam_total_reads(&bamfiles)?;
    eprintln!("Found {} total reads", total_reads);
    write_exon_cov(&options, &Arc::new(annot), total_reads, &bamfiles, &bamstrand)?;
    Ok(())
}

fn main() {
    // enable stack traces
    std::env::set_var("RUST_BACKTRACE", "full");

    if let Err(ref e) = run() {
        eprintln!("error: {}", e);
        eprintln!("backtrace: {:?}", e.backtrace());
        ::std::process::exit(1);
    }
}
<commit_msg>fix merging<commit_after>#![recursion_limit="128"]
#![cfg_attr(feature = "cargo-clippy", allow(cyclomatic_complexity, trivial_regex))]
use std::str;
use std::vec::Vec;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::io::stdout;
use std::ops::Range;

#[macro_use]
extern crate failure;

extern crate bam2bedgraph;
use bam2bedgraph::*;
use bam2bedgraph::error::*;
use bam2bedgraph::indexed_annotation::*;

extern crate rust_htslib;
use rust_htslib::bam::Read;
use rust_htslib::bam::IndexedReader;

extern crate structopt;
#[macro_use]
extern crate structopt_derive;
use structopt::StructOpt;

extern crate futures;
use futures::Future;
extern crate futures_cpupool;
use futures_cpupool::CpuPool;

extern crate num_cpus;

extern crate ordered_float;
use ordered_float::OrderedFloat;

#[derive(StructOpt, Debug)]
#[structopt(name = "intronrpkm", about = "Analyze RPKM values in intronic space")]
struct Options {
    // input files
    #[structopt(long="gff", help = "A genome annotation file in gff3 format", name="ANNOT_GFF_FILE")]
    annotfile_gff: Option<String>,
    #[structopt(long="gtf", help = "A genome annotation file in gtf format", name="ANNOT_GTF_FILE")]
    annotfile_gtf: Option<String>,
    #[structopt(long="bam1", short="1", help = "The set of stranded .bam files to analyze where read1 indicates strand", name="BAMFILE1")]
    bam1: Vec<String>,
    #[structopt(long="bam2", short="2", help = "The set of stranded .bam files to analyze where read2 indicates strand", name="BAMFILE2")]
    bam2: Vec<String>,
    #[structopt(long="bam", short="u", help = "The set of unstranded .bam files to analyze", name="BAMFILE")]
    bam: Vec<String>,
    #[structopt(long="chrmap", help = "Optional tab-delimited chr name mapping file", name="CHRMAP_FILE")]
    chrmap_file: Option<String>,
    #[structopt(long="vizchrmap", help = "Optional tab-delimited chr name mapping file for bigwig/bigbed exports", name="VIZCHRMAP_FILE")]
    vizchrmap_file: Option<String>,
    #[structopt(long="sizes", help = "Optional chr sizes file", name="SIZES_FILE")]
    sizes_file: Option<String>,
    // output file
    #[structopt(long="out", short="o", help = "Output file", name="OUT_FILE", default_value="-")]
    outfile: String,
    // feature types filter
    #[structopt(long="exon_type", help = "The exon type(s) to search for", name="EXON_TYPE")]
    exon_type: Vec<String>,
    #[structopt(long="transcript_type", help = "The transcript type(s) to search for", name="TRANSCRIPT_TYPE")]
    transcript_type: Vec<String>,
    #[structopt(long="gene_type", help = "The gene type(s) to search for", name="GENE_TYPE")]
    gene_type: Vec<String>,
    // flags
    #[structopt(long="cpu_threads", short="t", help = "How many threads to use for processing", default_value="0")]
    cpu_threads: usize,
}

#[derive(Ord, Eq, PartialOrd, PartialEq)]
struct Row {
    seqname: String,
    strand: String,
    start: u64,
    end: u64,
    cov: OrderedFloat<f64>,
    rpkm: OrderedFloat<f64>,
}

fn write_exon_cov(
    options: &Options,
    annot: &Arc<IndexedAnnotation>,
    total_reads: u64,
    bamfiles: &Vec<String>,
    bamstrand: &Vec<Option<bool>>)
    -> Result<()>
{
    let mut tidmaps = HashMap::<String,HashMap<String,u32>>::new();
    for bamfile in bamfiles {
        let bam = IndexedReader::from_path(bamfile)?;
        // build the tid map for this bam file
        let mut tidmap = HashMap::<String,u32>::new();
        {   let header = bam.header();
            for target_name in header.target_names() {
                let tid = header.tid(target_name).r()?;
                let target_name = String::from(std::str::from_utf8(target_name)?);
                let chr = annot.chrmap.get(&target_name).unwrap_or(&target_name);
                tidmap.insert(chr.clone(), tid);
            }
        }
        tidmaps.insert(bamfile.clone(), tidmap);
    }
    let tidmaps = Arc::new(tidmaps);

    let mut unmerged_exons = HashMap::<(String,String),Vec<Range<u64>>>::new();

    for (gene_row, gene) in annot.rows.iter().enumerate() {
        if options.gene_type.is_empty() || options.gene_type.contains(&gene.feature_type) {
            if let Some(feature_rows) = annot.row2children.get(&gene_row) {
                // get the transcript rows for this gene
                'transcript_row:
                for transcript_row in feature_rows {
                    let transcript = &annot.rows[*transcript_row];
                    // make sure this is a transcript type
                    // make sure transcript coordinates are consistent with the gene coordinates
                    if (options.transcript_type.is_empty() ||
                            options.transcript_type.contains(&transcript.feature_type)) &&
                        transcript.seqname == gene.seqname &&
                        transcript.strand == gene.strand
                    {
                        if let Some(exon_rows) = annot.row2children.get(transcript_row) {
                            // first get exon start/stop -> transcript associations
                            // and splice start/stop -> transcript associations
                            for exon_row in exon_rows {
                                let exon = &annot.rows[*exon_row];
                                if options.exon_type.is_empty() || options.exon_type.contains(&exon.feature_type) {
                                    let exon = &annot.rows[*exon_row];
                                    unmerged_exons.entry((exon.seqname.clone(),exon.strand.clone())).
                                        or_insert_with(Vec::new).push(exon.start-1..exon.end);
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    let mut merged_exons = HashMap::<(String,String),Vec<Range<u64>>>::new();
    let keys = unmerged_exons.keys().map(|k| k.clone()).collect::<Vec<_>>();
    for key in &keys {
        if let Some(ref mut unmerged) = unmerged_exons.get_mut(&key) {
            unmerged.sort_by_key(|u| u.start);
            let merged = merged_exons.entry((*key).clone()).or_insert_with(Vec::new);
            for exon in unmerged.iter() {
                let mut extend = false;
                if let Some(ref mut last) = merged.last_mut() {
                    if exon.start < last.end && last.start < exon.end {
                        last.end = exon.end;
                        extend = true;
                    }
                }
                if !extend {
                    merged.push(exon.clone());
                }
            }
        }
    }

    let num_cpus = num_cpus::get();
    let mut pair_futures = Vec::new();
    let pool = Arc::new(CpuPool::new(if options.cpu_threads==0 {num_cpus} else {options.cpu_threads}));
    for key in &keys {
        let exons = merged_exons[&key.clone()].clone();
        for exon in exons {
            let seqname = key.0.clone();
            let strand = key.1.clone();
            let chr = seqname.clone();
            let strand_is_plus = strand == "+";
            let bamfiles = bamfiles.clone();
            let bamstrand = bamstrand.clone();
            let tidmaps = tidmaps.clone();

            let pair_future = pool.spawn_fn(move ||->Result<Row> {
                let mut exon_cov = 0f64;
                let mut exon_reads = HashSet::<String>::new();
                for (i,bamfile) in bamfiles.iter().enumerate() {
                    let read1strand = bamstrand[i];
                    let tidmap = &tidmaps[bamfile];
                    if let Some(tid) = tidmap.get(&chr) {
                        let mut bam = IndexedReader::from_path(bamfile)?;
                        bam.fetch(*tid, exon.start as u32, exon.end as u32)?;
                        for read in bam.records() {
                            let read = read?;
                            // make sure the read's strand matches
                            if let Some(is_read1strand) = read1strand {
                                if !(((is_read1strand == read.is_first_in_template()) == !read.is_reverse()) == strand_is_plus) {
                                    continue;
                                }
                            }
                            let read_name = String::from(str::from_utf8(read.qname())?);
                            exon_reads.insert(read_name);
                            //eprintln!("Looking at read: {}", read_name);
                            let exons = cigar2exons(&read.cigar(), read.pos() as u64)?;
                            for e in exons {
                                if e.start < exon.end && exon.start < e.end {
                                    exon_cov +=
                                        (std::cmp::min(e.end, exon.end) -
                                        std::cmp::max(e.start, exon.start)) as f64;
                                }
                            }
                        }
                    }
                }
                let exon_length = exon.end-exon.start;
                let rpkm = (1e10f64 * exon_reads.len() as f64) / (total_reads as f64 * exon_length as f64);
                Ok(Row{
                    seqname: seqname,
                    strand: strand,
                    start: exon.start,
                    end: exon.end,
                    cov: OrderedFloat((exon_cov / exon_length as f64)),
                    rpkm: OrderedFloat(rpkm),
                })
            });
            pair_futures.push(pair_future);
        }
    }

    let mut rows = Vec::new();
    for future in pair_futures {
        match future.wait() {
            Ok(row) => {
                rows.push(row)
            }
            Err(ref e) => {
                eprintln!("Got Err in write_exon_cov: {:?}", e);
            }
        }
    }
    rows.sort();

    let mut output: BufWriter<Box<Write>> = BufWriter::new(
        if options.outfile == "-" { Box::new(stdout()) }
            else { Box::new(File::create(&options.outfile)?) });

    output.write_fmt(format_args!("{}\t{}\t{}\t{}\t{}\t{}\n",
                                  "seqname",
                                  "strand",
                                  "start",
                                  "end",
                                  "cov",
                                  "rpkm",
    ))?;
    for row in rows {
        output.write_fmt(format_args!("{}\t{}\t{}\t{}\t{}\t{}\n",
            row.seqname,
            row.strand,
            row.start,
            row.end,
            row.cov,
            row.rpkm,
        ))?;
    }
    Ok(())
}

fn run() -> Result<()> {
    let mut options = Options::from_args();
    // set defaults for feature types
    options.gene_type =
        (if options.gene_type.is_empty() { vec!["gene".to_string()] }
         else { options.gene_type.clone() }).into_iter().collect();
    //options.transcript_type =
    //    (if options.transcript_type.is_empty() { vec!["transcript".to_string()] }
    //     else { options.transcript_type.clone() }).into_iter().collect();
    options.exon_type =
        (if options.exon_type.is_empty() { vec!["exon".to_string()] }
         else { options.exon_type.clone() }).into_iter().collect();
    // set debug options if --debug flag is set

    // organize bamfiles by read strand type
    let mut bamfiles = Vec::<String>::new();
    bamfiles.append(&mut options.bam1.clone());
    bamfiles.append(&mut options.bam2.clone());
    bamfiles.append(&mut options.bam.clone());
    let mut bamstrand = Vec::<Option<bool>>::new();
    bamstrand.append(&mut options.bam1.iter().map(|_| Some(true)).collect());
    bamstrand.append(&mut options.bam2.iter().map(|_| Some(false)).collect());
    bamstrand.append(&mut options.bam.iter().map(|_| None).collect());
    // get the chromosome names and sizes from the first bam file
    if bamfiles.is_empty() {
        Options::clap().print_help()?;
        eprintln!("\n\nNo bam files were passed in!");
        std::process::exit(1)
    }
    let transcript_type = String::from("transcript");
    let mut annot = if let Some(annotfile_gff) = options.annotfile_gff.clone() {
        eprintln!("Reading annotation file {:?}", &annotfile_gff);
        IndexedAnnotation::from_gff(
            &annotfile_gff,
            &options.chrmap_file,
            &options.vizchrmap_file)?
    } else if let Some(annotfile_gtf) = options.annotfile_gtf.clone() {
        eprintln!("Reading annotation file {:?}", &annotfile_gtf);
        IndexedAnnotation::from_gtf(&annotfile_gtf,
            options.gene_type.get(0).r()?,
            options.transcript_type.get(0).unwrap_or(&transcript_type),
            &options.chrmap_file,
            &options.vizchrmap_file)?
    } else {
        Options::clap().print_help()?;
        eprintln!("\n\nNo annotation file was given!");
        std::process::exit(1);
    };
    eprintln!("Getting refseq lengths from bam file {:?}", &bamfiles[0]);
    let refs = match options.sizes_file.clone() {
        Some(sizes_file) => read_sizes_file(&sizes_file, &annot.chrmap)?,
        None => get_bam_refs(&bamfiles[0], &annot.chrmap)?,
    };
    annot.refs = refs;

    // get the total bam reads
    eprintln!("Running samtools idxstats to get total bam read counts");
    let total_reads = get_bam_total_reads(&bamfiles)?;
    eprintln!("Found {} total reads", total_reads);
    write_exon_cov(&options, &Arc::new(annot), total_reads, &bamfiles, &bamstrand)?;
    Ok(())
}

fn main() {
    // enable stack traces
    std::env::set_var("RUST_BACKTRACE", "full");

    if let Err(ref e) = run() {
        eprintln!("error: {}", e);
        eprintln!("backtrace: {:?}", e.backtrace());
        ::std::process::exit(1);
    }
}
<|endoftext|>


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

impl Ship {
    pub fn new(pos: Vec<[u8; 2]>) -> Ship {
        Ship {
            state: ShipState::Active,
            position: pos,
        }
    }

    pub fn pos(&self) -> &[[u8; 2]] {
        &self.position
    }

    pub fn is_active(&self) -> bool {
        self.state == ShipState::Active
    }

    pub fn is_sunk(&self) -> bool {
        self.state == ShipState::Sunk
    }

    pub fn sink(&mut self) {
        self.state = ShipState::Sunk;
    }
}

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

<commit_msg>Add Ship::len()<commit_after>pub struct Ship {
    state: ShipState,
    position: Vec<[u8; 2]>,
}

impl Ship {
    pub fn new(pos: Vec<[u8; 2]>) -> Ship {
        Ship {
            state: ShipState::Active,
            position: pos,
        }
    }

    pub fn pos(&self) -> &[[u8; 2]] {
        &self.position
    }

    pub fn len(&self) -> usize {
        self.position.len()
    }

    pub fn is_active(&self) -> bool {
        self.state == ShipState::Active
    }

    pub fn is_sunk(&self) -> bool {
        self.state == ShipState::Sunk
    }

    pub fn sink(&mut self) {
        self.state = ShipState::Sunk;
    }
}

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

<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Definitions and utilities for working with sinks.
//!
//! Sinks ("well-known flows" in terms of GRR nomenclature) are ever-existing
//! sessions on the GRR server that listen for various kinds of data. They are
//! a way to break away from the usual request-response workflow.
//!
//! For example, sinks can be used to notify the server about agent startup
//! (which is clearly not a response to a particular request) or to transfer
//! file blobs to a specialized storage.

use crate::action;
use crate::session;

/// Data sent to a server not belonging to a particular action call.
///
/// Sometimes the agent should send data to the server that is not connected
/// with a particular flow request. An example can be the agent startup data
/// sent to the server regardless of whether there was a request for it or not.
/// Another example can be file contents that are delivered to the server not
/// as part of the flow results but go to a separate store that has logic for
/// deduplicating data that we already collected in the past.
pub trait Parcel {
    /// Low-level Protocol Buffers type representing the parcel data.
    type Proto: protobuf::Message;

    /// A name of the corresponding RDF class in GRR.
    const RDF_NAME: &'static str;

    /// Converts the parcel to its low-level representation.
    fn into_proto(self) -> Self::Proto;
}

impl Parcel for () {

    type Proto = protobuf::well_known_types::Empty;

    // This implementation is intended to be used only in tests and we do not
    // really care about the `RDF_NAME` field there. And since GRR does not have
    // any RDF wrapper for empty type (except maybe `EmptyFlowArgs`, but this is
    // semantically different), we just leave it blank.
    const RDF_NAME: &'static str = "";

    fn into_proto(self) -> protobuf::well_known_types::Empty {
        protobuf::well_known_types::Empty::new()
    }
}

/// A parcel addressed to a particular server-side sink.
pub struct AddressedParcel<P: Parcel> {
    /// Destination of the parcel.
    sink: Sink,
    /// Actual parcel.
    parcel: P,
}

impl<P: Parcel> AddressedParcel<P> {

    /// Sink to which the parcel should be delivered to.
    pub fn sink(&self) -> Sink {
        self.sink
    }

    /// Unpacks the underlying parcel.
    pub fn unpack(self) -> P {
        self.parcel
    }
}

impl<P: Parcel> std::convert::TryInto<rrg_proto::jobs::GrrMessage> for AddressedParcel<P> {

    type Error = protobuf::ProtobufError;

    fn try_into(self) -> Result<rrg_proto::jobs::GrrMessage, protobuf::ProtobufError> {
        let mut message = rrg_proto::jobs::GrrMessage::new();
        message.set_session_id(String::from(self.sink.id));
        message.set_field_type(rrg_proto::jobs::GrrMessage_Type::MESSAGE);
        message.set_args_rdf_name(String::from(P::RDF_NAME));

        let proto = self.parcel.into_proto();
        message.set_args(protobuf::Message::write_to_bytes(&proto)?);

        Ok(message)
    }
}

/// Handle to a specific sink.
#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
pub struct Sink {
    /// An underlying identifier of the sink.
    id: &'static str,
}

impl Sink {

    // TODO: Refactor sinks to use custom `Message` type rather than `Response`
    // and make response always have `request_id` and `response_id` fields.

    /// Adresses a given `parcel` to this sink.
    pub fn address<P: Parcel>(&self, parcel: P) -> AddressedParcel<P> {
        AddressedParcel {
            sink: *self,
            parcel,
        }
    }

    // TODO: Remove the function below once only parcels are used.
    /// Wraps an action response to a sink-specific session response.
    pub fn wrap<R>(&self, response: R) -> session::Response<R>
    where
        R: action::Response,
    {
        session::Response {
            session_id: String::from(self.id),
            request_id: None,
            response_id: None,
            data: response,
        }
    }
}

/// A handle to the sink expecting startup information.
pub const STARTUP: Sink = Sink { id: "/flows/F:Startup" };

/// A handle to the transfer store sink.
pub const TRANSFER_STORE: Sink = Sink { id: "/flows/F:TransferStore" };
<commit_msg>Remove no longer used sink methods.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Definitions and utilities for working with sinks.
//!
//! Sinks ("well-known flows" in terms of GRR nomenclature) are ever-existing
//! sessions on the GRR server that listen for various kinds of data. They are
//! a way to break away from the usual request-response workflow.
//!
//! For example, sinks can be used to notify the server about agent startup
//! (which is clearly not a response to a particular request) or to transfer
//! file blobs to a specialized storage.

/// Data sent to a server not belonging to a particular action call.
///
/// Sometimes the agent should send data to the server that is not connected
/// with a particular flow request. An example can be the agent startup data
/// sent to the server regardless of whether there was a request for it or not.
/// Another example can be file contents that are delivered to the server not
/// as part of the flow results but go to a separate store that has logic for
/// deduplicating data that we already collected in the past.
pub trait Parcel {
    /// Low-level Protocol Buffers type representing the parcel data.
    type Proto: protobuf::Message;

    /// A name of the corresponding RDF class in GRR.
    const RDF_NAME: &'static str;

    /// Converts the parcel to its low-level representation.
    fn into_proto(self) -> Self::Proto;
}

impl Parcel for () {

    type Proto = protobuf::well_known_types::Empty;

    // This implementation is intended to be used only in tests and we do not
    // really care about the `RDF_NAME` field there. And since GRR does not have
    // any RDF wrapper for empty type (except maybe `EmptyFlowArgs`, but this is
    // semantically different), we just leave it blank.
    const RDF_NAME: &'static str = "";

    fn into_proto(self) -> protobuf::well_known_types::Empty {
        protobuf::well_known_types::Empty::new()
    }
}

/// A parcel addressed to a particular server-side sink.
pub struct AddressedParcel<P: Parcel> {
    /// Destination of the parcel.
    sink: Sink,
    /// Actual parcel.
    parcel: P,
}

impl<P: Parcel> AddressedParcel<P> {

    /// Sink to which the parcel should be delivered to.
    pub fn sink(&self) -> Sink {
        self.sink
    }

    /// Unpacks the underlying parcel.
    pub fn unpack(self) -> P {
        self.parcel
    }
}

impl<P: Parcel> std::convert::TryInto<rrg_proto::jobs::GrrMessage> for AddressedParcel<P> {

    type Error = protobuf::ProtobufError;

    fn try_into(self) -> Result<rrg_proto::jobs::GrrMessage, protobuf::ProtobufError> {
        let mut message = rrg_proto::jobs::GrrMessage::new();
        message.set_session_id(String::from(self.sink.id));
        message.set_field_type(rrg_proto::jobs::GrrMessage_Type::MESSAGE);
        message.set_args_rdf_name(String::from(P::RDF_NAME));

        let proto = self.parcel.into_proto();
        message.set_args(protobuf::Message::write_to_bytes(&proto)?);

        Ok(message)
    }
}

/// Handle to a specific sink.
#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
pub struct Sink {
    /// An underlying identifier of the sink.
    id: &'static str,
}

impl Sink {
    /// Adresses a given `parcel` to this sink.
    pub fn address<P: Parcel>(&self, parcel: P) -> AddressedParcel<P> {
        AddressedParcel {
            sink: *self,
            parcel,
        }
    }
}

/// A handle to the sink expecting startup information.
pub const STARTUP: Sink = Sink { id: "/flows/F:Startup" };

/// A handle to the transfer store sink.
pub const TRANSFER_STORE: Sink = Sink { id: "/flows/F:TransferStore" };
<|endoftext|>