<commit_before>use std::fmt::{Show,Formatter};
use std::fmt;
use std::cast::transmute;
use std::str::CharRange;

pub struct TokenValue {
    token: Token,
    value: ~str
}

impl TokenValue {
    pub fn new(tok: Token, val: ~str) -> TokenValue {
        TokenValue {
            token: tok,
            value: val
        }
    }
}

struct LexerIterator<'a> {
    string: &'a str,
    pos: uint,
    line: uint,
    next: uint,
    current: Option<char>,
    peek: Option<char>,
    previous: Option<char>
}

impl<'a> LexerIterator<'a> {
    pub fn new(string: &'a str) -> LexerIterator<'a> {
        LexerIterator {
            string: string,
            pos: 0,
            next: 0,
            line: 1,
            current: None,
            peek: None,
            previous: None
        }
    }

    pub fn next(&mut self) -> Option<char> {
        if self.string.len() != 0 && self.pos < self.string.len() {
            let CharRange {ch, next} = self.string.char_range_at(self.pos);
            self.pos = next;

            if self.current.is_some() {
                self.previous = self.current;
            }

            self.current = Some(ch);
            Some(ch)
        } else {
            None
        }
    }

    pub fn peek(&mut self) -> Option<char> {
        if self.string.len() != 0 && self.pos < self.string.len() {
            let CharRange {ch, next} = self.string.char_range_at(self.pos);
            self.peek = Some(ch);
            Some(ch)
        } else {
            None
        }
    }

}

/// Lexer tokens.
pub enum Token {
    PLUS,
    MINUS,
    INTEGER, // Any type of integer
    IDEN, // Identifier
    LPAREN,
    RPAREN,
    COLON
}

/// Lexer
pub struct Lexer<'a> {
    iter: *mut LexerIterator<'a>,
    done: bool
}

impl<'a> Lexer<'a> {
    pub fn new(input: &'a str) -> Lexer<'a> {
        Lexer {
            iter: unsafe { transmute(~LexerIterator::new(input)) },
            done: false
        }
    }

    pub fn next_token(&mut self) -> Option<TokenValue> {
        unsafe {
            for c in *self.iter {
                println!("{}", c);
                match c {
                    ' ' => continue,
                    '\r' => continue,
                    '\n' => continue,
                    '(' => return Some(TokenValue::new(LPAREN, ~"(")),
                    ')' => return Some(TokenValue::new(RPAREN, ~")")),
                    '+' => return Some(TokenValue::new(PLUS, ~"+")),
                    '-' => return Some(TokenValue::new(MINUS, ~"-")),
                    ':' => return Some(TokenValue::new(COLON, ~":")),
                    '9' => return Some(TokenValue::new(INTEGER, ~"9")),
                    '8' => return Some(TokenValue::new(INTEGER, ~"8")),
                    '7' => return Some(TokenValue::new(INTEGER, ~"7")),
                    '6' => return Some(TokenValue::new(INTEGER, ~"6")),
                    '5' => return Some(TokenValue::new(INTEGER, ~"5")),
                    '4' => return Some(TokenValue::new(INTEGER, ~"4")),
                    '3' => return Some(TokenValue::new(INTEGER, ~"3")),
                    '2' => return Some(TokenValue::new(INTEGER, ~"2")),
                    '1' => return Some(TokenValue::new(INTEGER, ~"1")),
                    '0' => return Some(TokenValue::new(INTEGER, ~"0")),
                    _ => {

                        let mut current = ~"";
                        let mut previous = ~"";

                        // Push the current character
                        current.push_char(c);

                        loop {
                            // Try the next character. If we get a valid identifier
                            // than we can move the position of the iterator.
                            match (*self.iter).peek() {
                                Some(character) => {
                                    println!("{}", character);
                                    // Lets save the current buffer as the previous.
                                    // This ensures that if this character makes the invalidation,
                                    // we can return the previously valid string.
                                    previous = current.clone();

                                    // Push the new chracter onto the current buffer.
                                    current.push_char(character);

                                    if is_iden(current) {
                                        (*self.iter).next();
                                    }

                                    if !is_iden(current) && is_iden(previous) {
                                        return Some(TokenValue::new(IDEN, previous));
                                    }
                                },
                                None => { break }
                            }
                        }

                        if !is_iden(current) && is_iden(previous) {
                            return Some(TokenValue::new(IDEN, previous));
                        }
                    }
                }
            }
        }
        None
    }
}

// An iden is identified by:
// /^([A-Za-z-+*^][A-Za-z0-9-_]*)$/
pub fn is_iden<'b>(ch: &'b str) -> bool {
    // Create a new iterator
    let mut iter = ch.chars().peekable();

    // Verify the first character.
    if !verify_iden_start(iter.next().unwrap()) {
        return false;
    }

    for x in iter {
        match x {
            '\n' => return false,
            '\r' => return false,
            ' ' => return false,
            _  =>  if !verify_iden_start(x) { return false }
        }
    }

    true
}

pub fn verify_iden_start(ch: char) -> bool {
    match ch {
        'A' => true,
        'a' => true,
        'B' => true,
        'b' => true,
        'C' => true,
        'c' => true,
        'D' => true,
        'd' => true,
        'E' => true,
        'e' => true,
        'F' => true,
        'f' => true,
        'G' => true,
        'g' => true,
        'H' => true,
        'h' => true,
        'I' => true,
        'i' => true,
        'J' => true,
        'j' => true,
        'K' => true,
        'k' => true,
        'L' => true,
        'l' => true,
        'M' => true,
        'm' => true,
        'N' => true,
        'n' => true,
        'O' => true,
        'o' => true,
        'P' => true,
        'p' => true,
        'Q' => true,
        'q' => true,
        'R' => true,
        'r' => true,
        'S' => true,
        's' => true,
        'T' => true,
        't' => true,
        'U' => true,
        'u' => true,
        'V' => true,
        'v' => true,
        'W' => true,
        'w' => true,
        'Y' => true,
        'y' => true,
        'Z' => true,
        'z' => true,
        '+' => true,
        '-' => true,
        '*' => true,
        _ => false
    }
}

impl Eq for Token {
    fn eq(&self, other: &Token) -> bool {
        match *self {
            PLUS => match *other { PLUS => true, _ => false },
            MINUS => match *other { MINUS => true, _ => false },
            INTEGER => match *other { INTEGER => true, _ => false },
            IDEN => match *other { IDEN => true, _ => false },
            LPAREN => match *other { LPAREN => true, _ => false },
            RPAREN => match *other { RPAREN => true, _ => false },
            COLON => match *other { COLON => true, _ => false }
        }
    }
}

impl Show for Token {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let c = match *self {
            PLUS => &"+",
            MINUS => &"-",
            INTEGER => &"integer",
            IDEN => &"iden",
            LPAREN => &"(",
            RPAREN => &")",
            COLON => ":"
        };

        write!(f.buf, "({})", c)
    }
}

#[cfg(test)]
mod test {
    use super::Lexer;
    use super::LPAREN;
    use super::TokenValue;

    #[test]
    fn next_token() {
        // Should spit out two tokens (LPAREN,RPAREN)
        let mut lex = Lexer::new(&"()");
        assert!(lex.next_token().unwrap().value == ~"(");
        assert!(lex.next_token().unwrap().value == ~")");
    }

    #[test]
    fn return_complex_tokens() {
        let mut lex = Lexer::new(&"(+(():830");

        assert!(lex.next_token().unwrap().value == ~"(");
        assert!(lex.next_token().unwrap().value == ~"+");
        assert!(lex.next_token().unwrap().value == ~"(");
        assert!(lex.next_token().unwrap().value == ~"(");
        assert!(lex.next_token().unwrap().value == ~")");
        assert!(lex.next_token().unwrap().value == ~":");
        assert!(lex.next_token().unwrap().value == ~"8");
        assert!(lex.next_token().unwrap().value == ~"3");
        assert!(lex.next_token().unwrap().value == ~"0");
    }

    #[test]
    fn should_tokenize_identifier() {
        let mut lex = Lexer::new(&"(let)");
        println!("")
        assert_eq!(lex.next_token().unwrap().value, ~"(");
        assert_eq!(lex.next_token().unwrap().value, ~"let");
        assert_eq!(lex.next_token().unwrap().value, ~")");
    }
}<commit_msg>tests pass<commit_after>use std::fmt::{Show,Formatter};
use std::fmt;
use std::cast::transmute;
use std::str::CharRange;

pub struct TokenValue {
    token: Token,
    value: ~str
}

impl TokenValue {
    pub fn new(tok: Token, val: ~str) -> TokenValue {
        TokenValue {
            token: tok,
            value: val
        }
    }
}

struct LexerIterator<'a> {
    string: &'a str,
    pos: uint,
    line: uint,
    next: uint,
    current: Option<char>,
    peek: Option<char>,
    previous: Option<char>
}

impl<'a> LexerIterator<'a> {
    pub fn new(string: &'a str) -> LexerIterator<'a> {
        LexerIterator {
            string: string,
            pos: 0,
            next: 0,
            line: 1,
            current: None,
            peek: None,
            previous: None
        }
    }

    pub fn next(&mut self) -> Option<char> {
        if self.string.len() != 0 && self.pos < self.string.len() {
            let CharRange {ch, next} = self.string.char_range_at(self.pos);
            self.pos = next;

            if self.current.is_some() {
                self.previous = self.current;
            }

            self.current = Some(ch);
            Some(ch)
        } else {
            None
        }
    }

    pub fn peek(&mut self) -> Option<char> {
        if self.string.len() != 0 && self.pos < self.string.len() {
            let CharRange {ch, next} = self.string.char_range_at(self.pos);
            self.peek = Some(ch);
            Some(ch)
        } else {
            None
        }
    }

}

/// Lexer tokens.
pub enum Token {
    PLUS,
    MINUS,
    INTEGER, // Any type of integer
    IDEN, // Identifier
    LPAREN,
    RPAREN,
    COLON
}

/// Lexer
pub struct Lexer<'a> {
    iter: *mut LexerIterator<'a>,
    done: bool
}

impl<'a> Lexer<'a> {
    pub fn new(input: &'a str) -> Lexer<'a> {
        Lexer {
            iter: unsafe { transmute(~LexerIterator::new(input)) },
            done: false
        }
    }

    pub fn next_token(&mut self) -> Option<TokenValue> {
        unsafe {
            for c in *self.iter {
                match c {
                    ' ' => continue,
                    '\r' => continue,
                    '\n' => continue,
                    '(' => return Some(TokenValue::new(LPAREN, ~"(")),
                    ')' => return Some(TokenValue::new(RPAREN, ~")")),
                    '+' => return Some(TokenValue::new(PLUS, ~"+")),
                    '-' => return Some(TokenValue::new(MINUS, ~"-")),
                    ':' => return Some(TokenValue::new(COLON, ~":")),
                    '9' => return Some(TokenValue::new(INTEGER, ~"9")),
                    '8' => return Some(TokenValue::new(INTEGER, ~"8")),
                    '7' => return Some(TokenValue::new(INTEGER, ~"7")),
                    '6' => return Some(TokenValue::new(INTEGER, ~"6")),
                    '5' => return Some(TokenValue::new(INTEGER, ~"5")),
                    '4' => return Some(TokenValue::new(INTEGER, ~"4")),
                    '3' => return Some(TokenValue::new(INTEGER, ~"3")),
                    '2' => return Some(TokenValue::new(INTEGER, ~"2")),
                    '1' => return Some(TokenValue::new(INTEGER, ~"1")),
                    '0' => return Some(TokenValue::new(INTEGER, ~"0")),
                    _ => {

                        let mut current = ~"";
                        let mut previous = ~"";

                        // Push the current character
                        current.push_char(c);

                        loop {
                            // Try the next character. If we get a valid identifier
                            // than we can move the position of the iterator.
                            match (*self.iter).peek() {
                                Some(character) => {
                                    // Lets save the current buffer as the previous.
                                    // This ensures that if this character makes the invalidation,
                                    // we can return the previously valid string.
                                    previous = current.clone();

                                    // Push the new chracter onto the current buffer.
                                    current.push_char(character);

                                    if is_iden(current) {
                                        (*self.iter).next();
                                    }

                                    if !is_iden(current) && is_iden(previous) {
                                        return Some(TokenValue::new(IDEN, previous));
                                    }
                                },
                                None => { break }
                            }
                        }

                        if !is_iden(current) && is_iden(previous) {
                            return Some(TokenValue::new(IDEN, previous));
                        }
                    }
                }
            }
        }
        None
    }
}

// An iden is identified by:
// /^([A-Za-z-+*^][A-Za-z0-9-_]*)$/
pub fn is_iden<'b>(ch: &'b str) -> bool {
    // Create a new iterator
    let mut iter = ch.chars().peekable();

    // Verify the first character.
    if !verify_iden_start(iter.next().unwrap()) {
        return false;
    }

    for x in iter {
        match x {
            '\n' => return false,
            '\r' => return false,
            ' ' => return false,
            _  =>  if !verify_iden_start(x) { return false }
        }
    }

    true
}

pub fn verify_iden_start(ch: char) -> bool {
    match ch {
        'A' => true,
        'a' => true,
        'B' => true,
        'b' => true,
        'C' => true,
        'c' => true,
        'D' => true,
        'd' => true,
        'E' => true,
        'e' => true,
        'F' => true,
        'f' => true,
        'G' => true,
        'g' => true,
        'H' => true,
        'h' => true,
        'I' => true,
        'i' => true,
        'J' => true,
        'j' => true,
        'K' => true,
        'k' => true,
        'L' => true,
        'l' => true,
        'M' => true,
        'm' => true,
        'N' => true,
        'n' => true,
        'O' => true,
        'o' => true,
        'P' => true,
        'p' => true,
        'Q' => true,
        'q' => true,
        'R' => true,
        'r' => true,
        'S' => true,
        's' => true,
        'T' => true,
        't' => true,
        'U' => true,
        'u' => true,
        'V' => true,
        'v' => true,
        'W' => true,
        'w' => true,
        'Y' => true,
        'y' => true,
        'Z' => true,
        'z' => true,
        '+' => true,
        '-' => true,
        '*' => true,
        _ => false
    }
}

impl Eq for Token {
    fn eq(&self, other: &Token) -> bool {
        match *self {
            PLUS => match *other { PLUS => true, _ => false },
            MINUS => match *other { MINUS => true, _ => false },
            INTEGER => match *other { INTEGER => true, _ => false },
            IDEN => match *other { IDEN => true, _ => false },
            LPAREN => match *other { LPAREN => true, _ => false },
            RPAREN => match *other { RPAREN => true, _ => false },
            COLON => match *other { COLON => true, _ => false }
        }
    }
}

impl Show for Token {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let c = match *self {
            PLUS => &"+",
            MINUS => &"-",
            INTEGER => &"integer",
            IDEN => &"iden",
            LPAREN => &"(",
            RPAREN => &")",
            COLON => ":"
        };

        write!(f.buf, "({})", c)
    }
}

#[cfg(test)]
mod test {
    use super::Lexer;
    use super::LPAREN;
    use super::TokenValue;

    #[test]
    fn next_token() {
        // Should spit out two tokens (LPAREN,RPAREN)
        let mut lex = Lexer::new(&"()");
        assert!(lex.next_token().unwrap().value == ~"(");
        assert!(lex.next_token().unwrap().value == ~")");
    }

    #[test]
    fn return_complex_tokens() {
        let mut lex = Lexer::new(&"(+(():830");

        assert!(lex.next_token().unwrap().value == ~"(");
        assert!(lex.next_token().unwrap().value == ~"+");
        assert!(lex.next_token().unwrap().value == ~"(");
        assert!(lex.next_token().unwrap().value == ~"(");
        assert!(lex.next_token().unwrap().value == ~")");
        assert!(lex.next_token().unwrap().value == ~":");
        assert!(lex.next_token().unwrap().value == ~"8");
        assert!(lex.next_token().unwrap().value == ~"3");
        assert!(lex.next_token().unwrap().value == ~"0");
    }

    #[test]
    fn should_tokenize_identifier() {
        let mut lex = Lexer::new(&"(let)");
        assert_eq!(lex.next_token().unwrap().value, ~"(");
        assert_eq!(lex.next_token().unwrap().value, ~"let");
        assert_eq!(lex.next_token().unwrap().value, ~")");
    }

    #[test]
    fn should_forget_whitespace() {
        let mut lex = Lexer::new(&"( )");
        assert_eq!(lex.next_token().unwrap().value, ~"(");
        assert_eq!(lex.next_token().unwrap().value, ~")");
    }
}<|endoftext|>


<commit_before>#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum XmlSpace {
    Default,
    Preserve
}

/// Implements `xml:space` handling per the SVG spec
///
/// Normalizes a string as it comes out of the XML parser's handler
/// for character data according to the SVG rules in
/// https://www.w3.org/TR/SVG/text.html#WhiteSpace
pub fn xml_space_normalize(mode: XmlSpace, s: &str) -> String {
    match mode {
        XmlSpace::Default => normalize_default(s),
        XmlSpace::Preserve => normalize_preserve(s)
    }
}

fn normalize_default(s: &str) -> String {
    #[derive(PartialEq)]
    enum State {
        SpacesAtStart,
        NonSpace,
        SpacesAtMiddle
    }

    let mut result = String::new();
    let mut state = State::SpacesAtStart;

    for ch in s.chars() {
        match ch {
            '\n' => continue,
            '\t' | ' ' => {
                match state {
                    State::SpacesAtStart  => continue,
                    State::NonSpace       => { state = State::SpacesAtMiddle; },
                    State::SpacesAtMiddle => continue,
                }
            },

            _ => {
                if state == State::SpacesAtMiddle {
                    result.push(' ');
                }

                result.push(ch);
                state = State::NonSpace;
            }
        }
    }

    result
}

fn normalize_preserve(s: &str) -> String {
    let mut result = String::new();

    for ch in s.chars() {
        match ch {
            '\n' | '\t' => result.push(' '),

            _ => result.push(ch),
        }
    }

    result
}

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

    #[test]
    fn xml_space_default() {
        assert_eq!(xml_space_normalize(XmlSpace::Default, "\n    WS example\n    indented lines\n  "),
                   "WS example indented lines");
        assert_eq!(xml_space_normalize(XmlSpace::Default, "\n  \t  \tWS \t\t\texample\n  \t  indented lines\t\t  \n  "),
                   "WS example indented lines");

        assert_eq!(xml_space_normalize(XmlSpace::Default, "\nWS example\nnon-indented lines\n  "),
                   "WS examplenon-indented lines");
    }

    #[test]
    fn xml_space_preserve() {
        assert_eq!(xml_space_normalize(XmlSpace::Preserve, "\n    WS example\n    indented lines\n  "),
                   "     WS example     indented lines   ");
        assert_eq!(xml_space_normalize(XmlSpace::Preserve, "\n  \t  \tWS \t\t\texample\n  \t  indented lines\t\t  \n  "),
                   "       WS    example      indented lines       ");

    }
}
<commit_msg>space::normalize_preserve() - Use iterator functions instead of making fire by rubbing sticks<commit_after>#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum XmlSpace {
    Default,
    Preserve
}

/// Implements `xml:space` handling per the SVG spec
///
/// Normalizes a string as it comes out of the XML parser's handler
/// for character data according to the SVG rules in
/// https://www.w3.org/TR/SVG/text.html#WhiteSpace
pub fn xml_space_normalize(mode: XmlSpace, s: &str) -> String {
    match mode {
        XmlSpace::Default => normalize_default(s),
        XmlSpace::Preserve => normalize_preserve(s)
    }
}

fn normalize_default(s: &str) -> String {
    #[derive(PartialEq)]
    enum State {
        SpacesAtStart,
        NonSpace,
        SpacesAtMiddle
    }

    let mut result = String::new();
    let mut state = State::SpacesAtStart;

    for ch in s.chars() {
        match ch {
            '\n' => continue,
            '\t' | ' ' => {
                match state {
                    State::SpacesAtStart  => continue,
                    State::NonSpace       => { state = State::SpacesAtMiddle; },
                    State::SpacesAtMiddle => continue,
                }
            },

            _ => {
                if state == State::SpacesAtMiddle {
                    result.push(' ');
                }

                result.push(ch);
                state = State::NonSpace;
            }
        }
    }

    result
}

fn normalize_preserve(s: &str) -> String {
    s.chars()
        .map(|ch| {
            match ch {
                '\n' | '\t' => ' ',

                c => c
            }
        })
        .collect()
}

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

    #[test]
    fn xml_space_default() {
        assert_eq!(xml_space_normalize(XmlSpace::Default, "\n    WS example\n    indented lines\n  "),
                   "WS example indented lines");
        assert_eq!(xml_space_normalize(XmlSpace::Default, "\n  \t  \tWS \t\t\texample\n  \t  indented lines\t\t  \n  "),
                   "WS example indented lines");

        assert_eq!(xml_space_normalize(XmlSpace::Default, "\nWS example\nnon-indented lines\n  "),
                   "WS examplenon-indented lines");
    }

    #[test]
    fn xml_space_preserve() {
        assert_eq!(xml_space_normalize(XmlSpace::Preserve, "\n    WS example\n    indented lines\n  "),
                   "     WS example     indented lines   ");
        assert_eq!(xml_space_normalize(XmlSpace::Preserve, "\n  \t  \tWS \t\t\texample\n  \t  indented lines\t\t  \n  "),
                   "       WS    example      indented lines       ");

    }
}
<|endoftext|>


<commit_before>use std::any::Any;
use std::any::TypeId;
use std::fmt;
use std::io::Read;
use std::io::Write;

#[cfg(feature = "bytes")]
use bytes::Bytes;

use crate::clear::Clear;
use crate::error::ProtobufError;
use crate::error::ProtobufResult;
use crate::reflect::MessageDescriptor;
use crate::reflect::ProtobufValue;
use crate::stream::CodedInputStream;
use crate::stream::CodedOutputStream;
use crate::stream::WithCodedInputStream;
use crate::stream::WithCodedOutputStream;
use crate::unknown::UnknownFields;

/// Trait implemented for all generated structs for protobuf messages.
///
/// Also, generated messages implement `Clone + Default + PartialEq`
pub trait Message: fmt::Debug + Clear + Send + Sync + ProtobufValue {
    /// Message descriptor for this message, used for reflection.
    fn descriptor(&self) -> &'static MessageDescriptor;

    /// True iff all required fields are initialized.
    /// Always returns `true` for protobuf 3.
    fn is_initialized(&self) -> bool;

    /// Update this message object with fields read from given stream.
    fn merge_from(&mut self, is: &mut CodedInputStream) -> ProtobufResult<()>;

    /// Write message to the stream.
    ///
    /// Sizes of this messages and nested messages must be cached
    /// by calling `compute_size` prior to this call.
    fn write_to_with_cached_sizes(&self, os: &mut CodedOutputStream) -> ProtobufResult<()>;

    /// Compute and cache size of this message and all nested messages
    fn compute_size(&self) -> u32;

    /// Get size previously computed by `compute_size`.
    fn get_cached_size(&self) -> u32;

    /// Write the message to the stream.
    ///
    /// Results in error if message is not fully initialized.
    fn write_to(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        self.check_initialized()?;

        // cache sizes
        self.compute_size();
        // TODO: reserve additional
        self.write_to_with_cached_sizes(os)?;

        Ok(())
    }

    /// Write the message to the stream prepending the message with message length
    /// encoded as varint.
    fn write_length_delimited_to(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        let size = self.compute_size();
        os.write_raw_varint32(size)?;
        self.write_to_with_cached_sizes(os)?;

        // TODO: assert we've written same number of bytes as computed

        Ok(())
    }

    /// Write the message to the vec, prepend the message with message length
    /// encoded as varint.
    fn write_length_delimited_to_vec(&self, vec: &mut Vec<u8>) -> ProtobufResult<()> {
        let mut os = CodedOutputStream::vec(vec);
        self.write_length_delimited_to(&mut os)?;
        os.flush()?;
        Ok(())
    }

    /// Update this message object with fields read from given stream.
    fn merge_from_bytes(&mut self, bytes: &[u8]) -> ProtobufResult<()> {
        let mut is = CodedInputStream::from_bytes(bytes);
        self.merge_from(&mut is)
    }

    /// Check if all required fields of this object are initialized.
    fn check_initialized(&self) -> ProtobufResult<()> {
        if !self.is_initialized() {
            Err(ProtobufError::MessageNotInitialized(
                self.descriptor().name(),
            ))
        } else {
            Ok(())
        }
    }

    /// Write the message to the writer.
    fn write_to_writer(&self, w: &mut dyn Write) -> ProtobufResult<()> {
        w.with_coded_output_stream(|os| self.write_to(os))
    }

    /// Write the message to bytes vec.
    fn write_to_vec(&self, v: &mut Vec<u8>) -> ProtobufResult<()> {
        v.with_coded_output_stream(|os| self.write_to(os))
    }

    /// Write the message to bytes vec.
    ///
    /// > **Note**: You can use [`parse_from_bytes`](crate::parse_from_bytes)
    /// to do the reverse.
    fn write_to_bytes(&self) -> ProtobufResult<Vec<u8>> {
        self.check_initialized()?;

        let size = self.compute_size() as usize;
        let mut v = Vec::with_capacity(size);
        // skip zerofill
        unsafe {
            v.set_len(size);
        }
        {
            let mut os = CodedOutputStream::bytes(&mut v);
            self.write_to_with_cached_sizes(&mut os)?;
            os.check_eof();
        }
        Ok(v)
    }

    /// Write the message to the writer, prepend the message with message length
    /// encoded as varint.
    fn write_length_delimited_to_writer(&self, w: &mut dyn Write) -> ProtobufResult<()> {
        w.with_coded_output_stream(|os| self.write_length_delimited_to(os))
    }

    /// Write the message to the bytes vec, prepend the message with message length
    /// encoded as varint.
    fn write_length_delimited_to_bytes(&self) -> ProtobufResult<Vec<u8>> {
        let mut v = Vec::new();
        v.with_coded_output_stream(|os| self.write_length_delimited_to(os))?;
        Ok(v)
    }

    /// Get a reference to unknown fields.
    fn get_unknown_fields(&self) -> &UnknownFields;
    /// Get a mutable reference to unknown fields.
    fn mut_unknown_fields(&mut self) -> &mut UnknownFields;

    /// Create an empty message object.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>() {
    /// let m = MyMessage::new();
    /// # }
    /// ```
    fn new() -> Self
    where
        Self: Sized;

    /// Get message descriptor for message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>() {
    /// let descriptor = MyMessage::descriptor_static();
    /// assert_eq!("MyMessage", descriptor.name());
    /// # }
    /// ```
    fn descriptor_static() -> &'static MessageDescriptor
    where
        Self: Sized,
    {
        panic!(
            "descriptor_static is not implemented for message, \
             LITE_RUNTIME must be used"
        );
    }

    /// Return a pointer to default immutable message with static lifetime.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>() {
    /// let m: &MyMessage = MyMessage::default_instance();
    /// # }
    /// ```
    fn default_instance() -> &'static Self
    where
        Self: Sized;
}

impl dyn Message {
    /// Downcast `Box<dyn Message>` to specific message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>(message: Box<Message>) {
    /// let m: Box<Message> = message;
    /// let m: Box<MyMessage> = Message::downcast_box(m).unwrap();
    /// # }
    /// ```
    pub fn downcast_box<T: Any>(self: Box<Self>) -> Result<Box<T>, Box<dyn Message>> {
        if Any::type_id(&*self) == TypeId::of::<T>() {
            unsafe {
                let raw: *mut dyn Message = Box::into_raw(self);
                Ok(Box::from_raw(raw as *mut T))
            }
        } else {
            Err(self)
        }
    }

    /// Downcast `&dyn Message` to specific message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>(message: &dyn Message) {
    /// let m: &dyn Message = message;
    /// let m: &MyMessage = Message::downcast_ref(m).unwrap();
    /// # }
    /// ```
    pub fn downcast_ref<'a, M: Message + 'a>(&'a self) -> Option<&'a M> {
        if Any::type_id(&*self) == TypeId::of::<M>() {
            unsafe { Some(&*(self as *const dyn Message as *const M)) }
        } else {
            None
        }
    }

    /// Downcast `&mut dyn Message` to specific message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>(message: &mut Message) {
    /// let m: &mut Message = message;
    /// let m: &mut MyMessage = Message::downcast_mut(m).unwrap();
    /// # }
    /// ```
    pub fn downcast_mut<'a, M: Message + 'a>(&'a mut self) -> Option<&'a mut M> {
        if Any::type_id(&*self) == TypeId::of::<M>() {
            unsafe { Some(&mut *(self as *mut dyn Message as *mut M)) }
        } else {
            None
        }
    }
}

impl Clone for Box<dyn Message> {
    fn clone(&self) -> Self {
        use std::ops::Deref;
        self.descriptor().clone(self.deref())
    }
}

#[cfg(off)] // don't need it
impl PartialEq for Box<Message> {
    fn eq(&self, other: &Box<Message>) -> bool {
        use std::ops::Deref;
        self.descriptor() == other.descriptor() && self.descriptor().eq(self.deref(), other.deref())
    }
}

/// Parse message from stream.
pub fn parse_from<M: Message>(is: &mut CodedInputStream) -> ProtobufResult<M> {
    let mut r: M = Message::new();
    r.merge_from(is)?;
    r.check_initialized()?;
    Ok(r)
}

/// Parse message from reader.
/// Parse stops on EOF or when error encountered.
pub fn parse_from_reader<M: Message>(reader: &mut dyn Read) -> ProtobufResult<M> {
    reader.with_coded_input_stream(|is| parse_from::<M>(is))
}

/// Parse message from byte array.
pub fn parse_from_bytes<M: Message>(bytes: &[u8]) -> ProtobufResult<M> {
    bytes.with_coded_input_stream(|is| parse_from::<M>(is))
}

/// Parse message from `Bytes` object.
/// Resulting message may share references to the passed bytes object.
#[cfg(feature = "bytes")]
pub fn parse_from_carllerche_bytes<M: Message>(bytes: &Bytes) -> ProtobufResult<M> {
    // Call trait explicitly to avoid accidental construction from `&[u8]`
    WithCodedInputStream::with_coded_input_stream(bytes, |is| parse_from::<M>(is))
}
<commit_msg>Use dyn before trait in rustdoc<commit_after>use std::any::Any;
use std::any::TypeId;
use std::fmt;
use std::io::Read;
use std::io::Write;

#[cfg(feature = "bytes")]
use bytes::Bytes;

use crate::clear::Clear;
use crate::error::ProtobufError;
use crate::error::ProtobufResult;
use crate::reflect::MessageDescriptor;
use crate::reflect::ProtobufValue;
use crate::stream::CodedInputStream;
use crate::stream::CodedOutputStream;
use crate::stream::WithCodedInputStream;
use crate::stream::WithCodedOutputStream;
use crate::unknown::UnknownFields;

/// Trait implemented for all generated structs for protobuf messages.
///
/// Also, generated messages implement `Clone + Default + PartialEq`
pub trait Message: fmt::Debug + Clear + Send + Sync + ProtobufValue {
    /// Message descriptor for this message, used for reflection.
    fn descriptor(&self) -> &'static MessageDescriptor;

    /// True iff all required fields are initialized.
    /// Always returns `true` for protobuf 3.
    fn is_initialized(&self) -> bool;

    /// Update this message object with fields read from given stream.
    fn merge_from(&mut self, is: &mut CodedInputStream) -> ProtobufResult<()>;

    /// Write message to the stream.
    ///
    /// Sizes of this messages and nested messages must be cached
    /// by calling `compute_size` prior to this call.
    fn write_to_with_cached_sizes(&self, os: &mut CodedOutputStream) -> ProtobufResult<()>;

    /// Compute and cache size of this message and all nested messages
    fn compute_size(&self) -> u32;

    /// Get size previously computed by `compute_size`.
    fn get_cached_size(&self) -> u32;

    /// Write the message to the stream.
    ///
    /// Results in error if message is not fully initialized.
    fn write_to(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        self.check_initialized()?;

        // cache sizes
        self.compute_size();
        // TODO: reserve additional
        self.write_to_with_cached_sizes(os)?;

        Ok(())
    }

    /// Write the message to the stream prepending the message with message length
    /// encoded as varint.
    fn write_length_delimited_to(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        let size = self.compute_size();
        os.write_raw_varint32(size)?;
        self.write_to_with_cached_sizes(os)?;

        // TODO: assert we've written same number of bytes as computed

        Ok(())
    }

    /// Write the message to the vec, prepend the message with message length
    /// encoded as varint.
    fn write_length_delimited_to_vec(&self, vec: &mut Vec<u8>) -> ProtobufResult<()> {
        let mut os = CodedOutputStream::vec(vec);
        self.write_length_delimited_to(&mut os)?;
        os.flush()?;
        Ok(())
    }

    /// Update this message object with fields read from given stream.
    fn merge_from_bytes(&mut self, bytes: &[u8]) -> ProtobufResult<()> {
        let mut is = CodedInputStream::from_bytes(bytes);
        self.merge_from(&mut is)
    }

    /// Check if all required fields of this object are initialized.
    fn check_initialized(&self) -> ProtobufResult<()> {
        if !self.is_initialized() {
            Err(ProtobufError::MessageNotInitialized(
                self.descriptor().name(),
            ))
        } else {
            Ok(())
        }
    }

    /// Write the message to the writer.
    fn write_to_writer(&self, w: &mut dyn Write) -> ProtobufResult<()> {
        w.with_coded_output_stream(|os| self.write_to(os))
    }

    /// Write the message to bytes vec.
    fn write_to_vec(&self, v: &mut Vec<u8>) -> ProtobufResult<()> {
        v.with_coded_output_stream(|os| self.write_to(os))
    }

    /// Write the message to bytes vec.
    ///
    /// > **Note**: You can use [`parse_from_bytes`](crate::parse_from_bytes)
    /// to do the reverse.
    fn write_to_bytes(&self) -> ProtobufResult<Vec<u8>> {
        self.check_initialized()?;

        let size = self.compute_size() as usize;
        let mut v = Vec::with_capacity(size);
        // skip zerofill
        unsafe {
            v.set_len(size);
        }
        {
            let mut os = CodedOutputStream::bytes(&mut v);
            self.write_to_with_cached_sizes(&mut os)?;
            os.check_eof();
        }
        Ok(v)
    }

    /// Write the message to the writer, prepend the message with message length
    /// encoded as varint.
    fn write_length_delimited_to_writer(&self, w: &mut dyn Write) -> ProtobufResult<()> {
        w.with_coded_output_stream(|os| self.write_length_delimited_to(os))
    }

    /// Write the message to the bytes vec, prepend the message with message length
    /// encoded as varint.
    fn write_length_delimited_to_bytes(&self) -> ProtobufResult<Vec<u8>> {
        let mut v = Vec::new();
        v.with_coded_output_stream(|os| self.write_length_delimited_to(os))?;
        Ok(v)
    }

    /// Get a reference to unknown fields.
    fn get_unknown_fields(&self) -> &UnknownFields;
    /// Get a mutable reference to unknown fields.
    fn mut_unknown_fields(&mut self) -> &mut UnknownFields;

    /// Create an empty message object.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>() {
    /// let m = MyMessage::new();
    /// # }
    /// ```
    fn new() -> Self
    where
        Self: Sized;

    /// Get message descriptor for message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>() {
    /// let descriptor = MyMessage::descriptor_static();
    /// assert_eq!("MyMessage", descriptor.name());
    /// # }
    /// ```
    fn descriptor_static() -> &'static MessageDescriptor
    where
        Self: Sized,
    {
        panic!(
            "descriptor_static is not implemented for message, \
             LITE_RUNTIME must be used"
        );
    }

    /// Return a pointer to default immutable message with static lifetime.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>() {
    /// let m: &MyMessage = MyMessage::default_instance();
    /// # }
    /// ```
    fn default_instance() -> &'static Self
    where
        Self: Sized;
}

impl dyn Message {
    /// Downcast `Box<dyn Message>` to specific message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>(message: Box<dyn Message>) {
    /// let m: Box<dyn Message> = message;
    /// let m: Box<MyMessage> = Message::downcast_box(m).unwrap();
    /// # }
    /// ```
    pub fn downcast_box<T: Any>(self: Box<Self>) -> Result<Box<T>, Box<dyn Message>> {
        if Any::type_id(&*self) == TypeId::of::<T>() {
            unsafe {
                let raw: *mut dyn Message = Box::into_raw(self);
                Ok(Box::from_raw(raw as *mut T))
            }
        } else {
            Err(self)
        }
    }

    /// Downcast `&dyn Message` to specific message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>(message: &dyn Message) {
    /// let m: &dyn Message = message;
    /// let m: &MyMessage = Message::downcast_ref(m).unwrap();
    /// # }
    /// ```
    pub fn downcast_ref<'a, M: Message + 'a>(&'a self) -> Option<&'a M> {
        if Any::type_id(&*self) == TypeId::of::<M>() {
            unsafe { Some(&*(self as *const dyn Message as *const M)) }
        } else {
            None
        }
    }

    /// Downcast `&mut dyn Message` to specific message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>(message: &mut dyn Message) {
    /// let m: &mut dyn Message = message;
    /// let m: &mut MyMessage = Message::downcast_mut(m).unwrap();
    /// # }
    /// ```
    pub fn downcast_mut<'a, M: Message + 'a>(&'a mut self) -> Option<&'a mut M> {
        if Any::type_id(&*self) == TypeId::of::<M>() {
            unsafe { Some(&mut *(self as *mut dyn Message as *mut M)) }
        } else {
            None
        }
    }
}

impl Clone for Box<dyn Message> {
    fn clone(&self) -> Self {
        use std::ops::Deref;
        self.descriptor().clone(self.deref())
    }
}

#[cfg(off)] // don't need it
impl PartialEq for Box<Message> {
    fn eq(&self, other: &Box<Message>) -> bool {
        use std::ops::Deref;
        self.descriptor() == other.descriptor() && self.descriptor().eq(self.deref(), other.deref())
    }
}

/// Parse message from stream.
pub fn parse_from<M: Message>(is: &mut CodedInputStream) -> ProtobufResult<M> {
    let mut r: M = Message::new();
    r.merge_from(is)?;
    r.check_initialized()?;
    Ok(r)
}

/// Parse message from reader.
/// Parse stops on EOF or when error encountered.
pub fn parse_from_reader<M: Message>(reader: &mut dyn Read) -> ProtobufResult<M> {
    reader.with_coded_input_stream(|is| parse_from::<M>(is))
}

/// Parse message from byte array.
pub fn parse_from_bytes<M: Message>(bytes: &[u8]) -> ProtobufResult<M> {
    bytes.with_coded_input_stream(|is| parse_from::<M>(is))
}

/// Parse message from `Bytes` object.
/// Resulting message may share references to the passed bytes object.
#[cfg(feature = "bytes")]
pub fn parse_from_carllerche_bytes<M: Message>(bytes: &Bytes) -> ProtobufResult<M> {
    // Call trait explicitly to avoid accidental construction from `&[u8]`
    WithCodedInputStream::with_coded_input_stream(bytes, |is| parse_from::<M>(is))
}
<|endoftext|>


<commit_before>// This example drives c-ares using an mio::EventLoop.
extern crate c_ares;
extern crate mio;

use std::collections::HashSet;
use std::mem;
use std::os::unix::io;
use std::sync::mpsc;
use std::thread;

// Messages for the event loop.
enum CAresHandlerMessage {
    // 'Notify me when this file descriptor becomes readable, or writable'
    // The first bool is for 'readable' and the second is for 'writable'.  It's
    // allowed to set both of these - or neither, meaning 'I am no longer
    // interested in this file descriptor'.
    RegisterInterest(io::RawFd, bool, bool),

    // 'Shut down'.
    ShutDown,
}

struct CAresEventHandler {
    // Since the event handler owns the Channel, it won't be possible to submit
    // further queries once the event loop is running.  If you want to do that,
    // use an Arc<Mutex<c_ares::Channel>>.
    ares_channel: c_ares::Channel,
    tracked_fds: HashSet<io::RawFd>,
}

impl mio::Handler for CAresEventHandler {
    type Timeout = ();
    type Message = CAresHandlerMessage;

    // mio notifies us that a file descriptor is readable, so we tell the
    // Channel the same.
    fn readable(
        &mut self,
        _event_loop: &mut mio::EventLoop<CAresEventHandler>,
        token: mio::Token,
        _read_hint: mio::ReadHint) {
        let fd = token.as_usize() as io::RawFd;
        self.ares_channel.process_fd(fd, c_ares::INVALID_FD);
    }

    // mio notifies us that a file descriptor is writable, so we tell the
    // Channel the same.
    fn writable(
        &mut self,
        _event_loop: &mut mio::EventLoop<CAresEventHandler>,
        token: mio::Token) {
        let fd = token.as_usize() as io::RawFd;
        self.ares_channel.process_fd(c_ares::INVALID_FD, fd);
    }

    // Process received messages.  Either:
    // - we're asked to register interest (or non-interest) in a file
    // descriptor
    // - we're asked to shut down the event loop.
    fn notify(
        &mut self,
        event_loop:&mut mio::EventLoop<CAresEventHandler>,
        msg: Self::Message) {
        match msg {
            CAresHandlerMessage::RegisterInterest(fd, readable, writable) => {
                let io = mio::Io::new(fd);
                if !readable && !writable {
                    self.tracked_fds.remove(&fd);
                    event_loop
                        .deregister(&io)
                        .ok()
                        .expect("failed to deregister interest");
                } else {
                    let mut interest = mio::Interest::none();
                    if readable {
                        interest = interest | mio::Interest::readable();
                    }
                    if writable {
                        interest = interest | mio::Interest::writable();
                    }
                    let token = mio::Token(fd as usize);
                    let register_result = if !self.tracked_fds.insert(fd) {
                        event_loop
                            .reregister(
                                &io,
                                token,
                                interest,
                                mio::PollOpt::level())
                    } else {
                        event_loop
                            .register_opt(
                                &io,
                                token,
                                interest,
                                mio::PollOpt::level())
                    };
                    register_result.ok().expect("failed to register interest");
                }

                // Don't accidentally close the file descriptor by dropping io!
                mem::forget(io);
            },

            CAresHandlerMessage::ShutDown => event_loop.shutdown(),
        }
    }

    // We run a recurring timer so that we can spot non-responsive servers.
    //
    // In that case we won't get a callback saying that anything is happening
    // on any file descriptor, but nevertheless need to give the Channel an
    // opportunity to notice that it has timed-out requests pending.
    fn timeout(
        &mut self,
        event_loop: &mut mio::EventLoop<CAresEventHandler>,
        _timeout: Self::Timeout) {
        self.ares_channel.process_fd(c_ares::INVALID_FD, c_ares::INVALID_FD);
        event_loop.timeout_ms((), 500).unwrap();
    }
}

impl CAresEventHandler {
    fn new(ares_channel: c_ares::Channel) -> CAresEventHandler {
        CAresEventHandler {
            ares_channel: ares_channel,
            tracked_fds: HashSet::new(),
        }
    }
}

fn print_a_result(result: Result<c_ares::AResult, c_ares::AresError>) {
    match result {
        Err(e) => {
            let err_string = c_ares::str_error(e);
            println!("Lookup failed with error '{:}'", err_string);
        }
        Ok(result) => {
            println!("Successful lookup...");
            for addr in &result.ip_addrs {
                println!("{:}", addr);
            }
        }
    }
}

fn print_aaaa_result(result: Result<c_ares::AAAAResult, c_ares::AresError>) {
    match result {
        Err(e) => {
            let err_string = c_ares::str_error(e);
            println!("Lookup failed with error '{:}'", err_string);
        }
        Ok(result) => {
            println!("Successful lookup...");
            for addr in &result.ip_addrs {
                println!("{:}", addr);
            }
        }
    }
}

fn main() {
    // Create an event loop, and a c_ares::Channel.
    let mut event_loop = mio::EventLoop::new()
        .ok()
        .expect("failed to create event loop");
    let event_loop_channel = event_loop.channel();
    let event_loop_channel_clone = event_loop_channel.clone();
    let sock_callback = move |fd: io::RawFd, readable: bool, writable: bool| {
        event_loop_channel_clone
            .send(CAresHandlerMessage::RegisterInterest(fd, readable, writable))
            .ok()
            .expect("Failed to send RegisterInterest");
    };
    let mut ares_channel = c_ares::Channel::new(sock_callback)
        .ok()
        .expect("Failed to create channel");

    // Set up a couple of queries.
    let (results_tx, results_rx) = mpsc::channel();
    let tx = results_tx.clone();
    ares_channel.query_a("apple.com", move |result| {
        print_a_result(result);
        tx.send(()).unwrap()
    });

    let tx = results_tx.clone();
    ares_channel.query_aaaa("google.com", move |result| {
        print_aaaa_result(result);
        tx.send(()).unwrap()
    });

    // Set the first instance of the recurring timer on the event loop.
    event_loop.timeout_ms((), 500).unwrap();

    // Kick off the event loop.
    let mut event_handler = CAresEventHandler::new(ares_channel);
    let handle = thread::spawn(move || {
        event_loop
            .run(&mut event_handler)
            .ok()
            .expect("failed to run event loop")
    });

    // Wait for results to roll in.
    for _ in 0..2 {
        results_rx.recv().unwrap();
    }

    // Shut down the event loop and wait for it to finish.
    event_loop_channel
        .send(CAresHandlerMessage::ShutDown)
        .ok()
        .expect("failed to shut down event loop");
    handle.join().unwrap();
}
<commit_msg>Upgrade to mio 0.4.0<commit_after>// This example drives c-ares using an mio::EventLoop.
extern crate c_ares;
extern crate mio;

use std::collections::HashSet;
use std::mem;
use std::os::unix::io;
use std::sync::mpsc;
use std::thread;

// Messages for the event loop.
enum CAresHandlerMessage {
    // 'Notify me when this file descriptor becomes readable, or writable'
    // The first bool is for 'readable' and the second is for 'writable'.  It's
    // allowed to set both of these - or neither, meaning 'I am no longer
    // interested in this file descriptor'.
    RegisterInterest(io::RawFd, bool, bool),

    // 'Shut down'.
    ShutDown,
}

struct CAresEventHandler {
    // Since the event handler owns the Channel, it won't be possible to submit
    // further queries once the event loop is running.  If you want to do that,
    // use an Arc<Mutex<c_ares::Channel>>.
    ares_channel: c_ares::Channel,
    tracked_fds: HashSet<io::RawFd>,
}

impl mio::Handler for CAresEventHandler {
    type Timeout = ();
    type Message = CAresHandlerMessage;

    // mio notifies us that a file descriptor is readable or writable, so we
    // tell the Channel the same.
    fn ready(
        &mut self,
        _event_loop: &mut mio::EventLoop<CAresEventHandler>,
        token: mio::Token,
        events: mio::EventSet) {
        let fd = token.as_usize() as io::RawFd;
        let read_fd = if events.is_readable() { fd } else { c_ares::INVALID_FD };
        let write_fd = if events.is_writable() { fd } else { c_ares::INVALID_FD };
        self.ares_channel.process_fd(read_fd, write_fd);
    }

    // Process received messages.  Either:
    // - we're asked to register interest (or non-interest) in a file
    // descriptor
    // - we're asked to shut down the event loop.
    fn notify(
        &mut self,
        event_loop:&mut mio::EventLoop<CAresEventHandler>,
        msg: Self::Message) {
        match msg {
            CAresHandlerMessage::RegisterInterest(fd, readable, writable) => {
                let io = mio::Io::from(fd);
                if !readable && !writable {
                    self.tracked_fds.remove(&fd);
                    event_loop
                        .deregister(&io)
                        .ok()
                        .expect("failed to deregister interest");
                } else {
                    let mut interest = mio::EventSet::none();
                    if readable {
                        interest = interest | mio::EventSet::readable();
                    }
                    if writable {
                        interest = interest | mio::EventSet::writable();
                    }
                    let token = mio::Token(fd as usize);
                    let register_result = if !self.tracked_fds.insert(fd) {
                        event_loop
                            .reregister(
                                &io,
                                token,
                                interest,
                                mio::PollOpt::level())
                    } else {
                        event_loop
                            .register_opt(
                                &io,
                                token,
                                interest,
                                mio::PollOpt::level())
                    };
                    register_result.ok().expect("failed to register interest");
                }

                // Don't accidentally close the file descriptor by dropping io!
                mem::forget(io);
            },

            CAresHandlerMessage::ShutDown => event_loop.shutdown(),
        }
    }

    // We run a recurring timer so that we can spot non-responsive servers.
    //
    // In that case we won't get a callback saying that anything is happening
    // on any file descriptor, but nevertheless need to give the Channel an
    // opportunity to notice that it has timed-out requests pending.
    fn timeout(
        &mut self,
        event_loop: &mut mio::EventLoop<CAresEventHandler>,
        _timeout: Self::Timeout) {
        self.ares_channel.process_fd(c_ares::INVALID_FD, c_ares::INVALID_FD);
        event_loop.timeout_ms((), 500).unwrap();
    }
}

impl CAresEventHandler {
    fn new(ares_channel: c_ares::Channel) -> CAresEventHandler {
        CAresEventHandler {
            ares_channel: ares_channel,
            tracked_fds: HashSet::new(),
        }
    }
}

fn print_a_result(result: Result<c_ares::AResult, c_ares::AresError>) {
    match result {
        Err(e) => {
            let err_string = c_ares::str_error(e);
            println!("Lookup failed with error '{:}'", err_string);
        }
        Ok(result) => {
            println!("Successful lookup...");
            for addr in &result.ip_addrs {
                println!("{:}", addr);
            }
        }
    }
}

fn print_aaaa_result(result: Result<c_ares::AAAAResult, c_ares::AresError>) {
    match result {
        Err(e) => {
            let err_string = c_ares::str_error(e);
            println!("Lookup failed with error '{:}'", err_string);
        }
        Ok(result) => {
            println!("Successful lookup...");
            for addr in &result.ip_addrs {
                println!("{:}", addr);
            }
        }
    }
}

fn main() {
    // Create an event loop, and a c_ares::Channel.
    let mut event_loop = mio::EventLoop::new()
        .ok()
        .expect("failed to create event loop");
    let event_loop_channel = event_loop.channel();
    let event_loop_channel_clone = event_loop_channel.clone();
    let sock_callback = move |fd: io::RawFd, readable: bool, writable: bool| {
        event_loop_channel_clone
            .send(CAresHandlerMessage::RegisterInterest(fd, readable, writable))
            .ok()
            .expect("Failed to send RegisterInterest");
    };
    let mut ares_channel = c_ares::Channel::new(sock_callback)
        .ok()
        .expect("Failed to create channel");

    // Set up a couple of queries.
    let (results_tx, results_rx) = mpsc::channel();
    let tx = results_tx.clone();
    ares_channel.query_a("apple.com", move |result| {
        print_a_result(result);
        tx.send(()).unwrap()
    });

    let tx = results_tx.clone();
    ares_channel.query_aaaa("google.com", move |result| {
        print_aaaa_result(result);
        tx.send(()).unwrap()
    });

    // Set the first instance of the recurring timer on the event loop.
    event_loop.timeout_ms((), 500).unwrap();

    // Kick off the event loop.
    let mut event_handler = CAresEventHandler::new(ares_channel);
    let handle = thread::spawn(move || {
        event_loop
            .run(&mut event_handler)
            .ok()
            .expect("failed to run event loop")
    });

    // Wait for results to roll in.
    for _ in 0..2 {
        results_rx.recv().unwrap();
    }

    // Shut down the event loop and wait for it to finish.
    event_loop_channel
        .send(CAresHandlerMessage::ShutDown)
        .ok()
        .expect("failed to shut down event loop");
    handle.join().unwrap();
}
<|endoftext|>


<commit_before>trait Magic {
    fn magic_num(&self) -> usize;
}

struct Foobar {
    name: String,
}

impl Magic for Foobar {
    fn magic_num(&self) -> usize {
        return if self.name.len() == 0 { 2 } else { 33 };
    }
}

impl Magic for bool {
    fn magic_num(&self) -> usize {
        return if *self { 3 } else { 54 };
    }
}

fn report_magic<M: Magic>(v: &M) {
    println!("magic num: {}", v.magic_num());
}

fn main() {
    let fb = Foobar {
        name: "Joe".to_string(),
    };

    report_magic(&fb);

    let tv = false;
    report_magic(&tv);

    println!("{}", true.magic_num());
}
<commit_msg>Example of implementing extension trait on Vec<commit_after>trait Magic {
    fn magic_num(&self) -> usize;
}

struct Foobar {
    name: String,
}

impl Magic for Foobar {
    fn magic_num(&self) -> usize {
        return if self.name.len() == 0 { 2 } else { 33 };
    }
}

impl Magic for bool {
    fn magic_num(&self) -> usize {
        return if *self { 3 } else { 54 };
    }
}

impl<T> Magic for Vec<T> {
    fn magic_num(&self) -> usize {
        return if self.len() == 0 { 10 } else { 5 };
    }
}

fn report_magic<M: Magic>(v: &M) {
    println!("magic num: {}", v.magic_num());
}

fn main() {
    let fb = Foobar {
        name: "Joe".to_string(),
    };

    report_magic(&fb);

    let tv = false;
    report_magic(&tv);

    println!("{}", true.magic_num());

    let vv = vec![1, 2, 3];
    report_magic(&vv);
}
<|endoftext|>


<commit_before><commit_msg>Add description to accessor stride assert<commit_after><|endoftext|>


<commit_before>use gl;
use gl::types::*;
use texture::{Texture2dBinder, TextureCubeMapBinder};
use types::GLError;

unsafe fn _active_texture(idx: u32) {
    gl::ActiveTexture(gl::TEXTURE0 + (idx as GLenum));
    dbg_gl_error! {
        GLError::InvalidEnum => "`texture` is out of bounds (expected to be GL_TEXTUREi, 0 <= i < GL_MAX_TEXTURE_IMAGE_UNITS)",
        _ => "Unknown error"
    }
}

pub trait TextureUnit {
    fn idx() -> u32;

    fn active(&mut self) -> TextureUnitBinding {
        unsafe { _active_texture(Self::idx()); }
        TextureUnitBinding {
            texture_2d: Texture2dBinder,
            texture_cube_map: TextureCubeMapBinder
        }
    }
}

// TODO: Use a macro, or const generic parameters:
// https://github.com/rust-lang/rfcs/issues/273
// https://github.com/rust-lang/rfcs/issues/1038
pub struct TextureUnit0;
pub struct TextureUnit1;
pub struct TextureUnit2;
pub struct TextureUnit3;
pub struct TextureUnit4;
pub struct TextureUnit5;
pub struct TextureUnit6;
pub struct TextureUnit7;

impl TextureUnit for TextureUnit0 { fn idx() -> u32 { 0 } }
impl TextureUnit for TextureUnit1 { fn idx() -> u32 { 1 } }
impl TextureUnit for TextureUnit2 { fn idx() -> u32 { 2 } }
impl TextureUnit for TextureUnit3 { fn idx() -> u32 { 3 } }
impl TextureUnit for TextureUnit4 { fn idx() -> u32 { 4 } }
impl TextureUnit for TextureUnit5 { fn idx() -> u32 { 5 } }
impl TextureUnit for TextureUnit6 { fn idx() -> u32 { 6 } }
impl TextureUnit for TextureUnit7 { fn idx() -> u32 { 7 } }

// NOTE: Ensure the number of each texture unit matches its index in the tuple
pub struct TextureUnits(pub TextureUnit0, pub TextureUnit1, pub TextureUnit2,
                        pub TextureUnit3, pub TextureUnit4, pub TextureUnit5,
                        pub TextureUnit6, pub TextureUnit7);

impl TextureUnits {
    pub unsafe fn current() -> TextureUnits {
        TextureUnits(TextureUnit0, TextureUnit1, TextureUnit2, TextureUnit3,
                     TextureUnit4, TextureUnit5, TextureUnit6, TextureUnit7)
    }

    pub unsafe fn active_nth(&self, idx: u32) -> TextureUnitBinding {
        _active_texture(idx);
        TextureUnitBinding {
            texture_2d: Texture2dBinder,
            texture_cube_map: TextureCubeMapBinder
        }
    }
}

pub struct TextureUnitBinding {
    pub texture_2d: Texture2dBinder,
    pub texture_cube_map: TextureCubeMapBinder
}

#[macro_export]
macro_rules! bind_texture_cube_map {
    ($gl_tex_unit:expr, $texture:expr) => {
        $gl_tex_unit.texture_cube_map.bind($texture)
    }
}

#[macro_export]
macro_rules! bind_texture_2d {
    ($gl_tex_unit:expr, $texture:expr) => {
        $gl_tex_unit.texture_2d.bind($texture)
    }
}
<commit_msg>Add `TextureUnitBinding::current_at_idx` private constructor<commit_after>use gl;
use gl::types::*;
use texture::{Texture2dBinder, TextureCubeMapBinder};
use types::GLError;

unsafe fn _active_texture(idx: u32) {
    gl::ActiveTexture(gl::TEXTURE0 + (idx as GLenum));
    dbg_gl_error! {
        GLError::InvalidEnum => "`texture` is out of bounds (expected to be GL_TEXTUREi, 0 <= i < GL_MAX_TEXTURE_IMAGE_UNITS)",
        _ => "Unknown error"
    }
}

pub trait TextureUnit {
    fn idx() -> u32;

    fn active(&mut self) -> TextureUnitBinding {
        unsafe {
            _active_texture(Self::idx());
            TextureUnitBinding::current_at_idx(Self::idx())
        }
    }
}

// TODO: Use a macro, or const generic parameters:
// https://github.com/rust-lang/rfcs/issues/273
// https://github.com/rust-lang/rfcs/issues/1038
pub struct TextureUnit0;
pub struct TextureUnit1;
pub struct TextureUnit2;
pub struct TextureUnit3;
pub struct TextureUnit4;
pub struct TextureUnit5;
pub struct TextureUnit6;
pub struct TextureUnit7;

impl TextureUnit for TextureUnit0 { fn idx() -> u32 { 0 } }
impl TextureUnit for TextureUnit1 { fn idx() -> u32 { 1 } }
impl TextureUnit for TextureUnit2 { fn idx() -> u32 { 2 } }
impl TextureUnit for TextureUnit3 { fn idx() -> u32 { 3 } }
impl TextureUnit for TextureUnit4 { fn idx() -> u32 { 4 } }
impl TextureUnit for TextureUnit5 { fn idx() -> u32 { 5 } }
impl TextureUnit for TextureUnit6 { fn idx() -> u32 { 6 } }
impl TextureUnit for TextureUnit7 { fn idx() -> u32 { 7 } }

// NOTE: Ensure the number of each texture unit matches its index in the tuple
pub struct TextureUnits(pub TextureUnit0, pub TextureUnit1, pub TextureUnit2,
                        pub TextureUnit3, pub TextureUnit4, pub TextureUnit5,
                        pub TextureUnit6, pub TextureUnit7);

impl TextureUnits {
    pub unsafe fn current() -> TextureUnits {
        TextureUnits(TextureUnit0, TextureUnit1, TextureUnit2, TextureUnit3,
                     TextureUnit4, TextureUnit5, TextureUnit6, TextureUnit7)
    }

    pub unsafe fn active_nth(&self, idx: u32) -> TextureUnitBinding {
        _active_texture(idx);
        TextureUnitBinding::current_at_idx(idx)
    }
}

pub struct TextureUnitBinding {
    pub texture_2d: Texture2dBinder,
    pub texture_cube_map: TextureCubeMapBinder
}

impl TextureUnitBinding {
    unsafe fn current_at_idx(_idx: u32) -> Self {
        TextureUnitBinding {
            texture_2d: Texture2dBinder,
            texture_cube_map: TextureCubeMapBinder
        }
    }
}

#[macro_export]
macro_rules! bind_texture_cube_map {
    ($gl_tex_unit:expr, $texture:expr) => {
        $gl_tex_unit.texture_cube_map.bind($texture)
    }
}

#[macro_export]
macro_rules! bind_texture_2d {
    ($gl_tex_unit:expr, $texture:expr) => {
        $gl_tex_unit.texture_2d.bind($texture)
    }
}
<|endoftext|>


<commit_before>extern crate mles_utils;

use std::thread;
use std::sync::mpsc::{Sender, Receiver};
use std::sync::mpsc::channel;
use std::net::TcpStream;
use std::net::TcpListener;
use std::collections::HashMap;
use std::io::Write;
use std::time::Duration;
use std::option::Option;
use std::str;
use mles_utils::*;

fn process_channel(tx: Sender<Sender<TcpStream>>, removedtx: Sender<String>, this_channel: String ) {
    let mut cnt = 0;
    let (thr_tx, thr_rx): (Sender<TcpStream>, Receiver<TcpStream>) = channel();
    println!("Spawned: New channel created!");
    let mut users = HashMap::new();
    let mut messages: Vec<Vec<u8>> = Vec::new();
    tx.send(thr_tx.clone()).unwrap();
    loop {
        let mut removals = Vec::new();
        let mut newuser = true;
        while newuser {
            match thr_rx.try_recv() {
                Ok(val) => {
                    let mut thr: TcpStream = val;
                    cnt += 1;
                    println!("Adding user {}", cnt);
                    users.insert(cnt, thr.try_clone().unwrap());

                    /* If a new user, all push messages to her */
                    for buf in &messages {
                        thr.write(buf.as_slice()).unwrap();
                    }
                    newuser = true;
                },
                    Err(_) => { newuser = false; }
            }
        }
        for (user, thr_socket) in &users {
            let stream = thr_socket.try_clone().unwrap();
            let tuple = read_n(&stream, 4);
            let status = tuple.0;
            match status {
                Ok(0) => {
                    removals.push(user.clone());
                },
                    _ => {}
            }
            let mut buf = tuple.1;
            if 0 == buf.len() {
                continue;
            }
            if read_hdr_type(buf.as_slice()) != 'M' as u32 {
                continue;
            }
            let hdr_len = read_hdr_len(buf.as_slice()) as u64;
            if 0 == hdr_len {
                continue;
            }
            let tuple = read_n(&stream, hdr_len);
            let status = tuple.0;
            match status {
                Ok(0) => {
                    removals.push(user.clone());
                },
                    _ => {}
            }
            let payload = tuple.1;
            if payload.len() != (hdr_len as usize) {
                continue;
            }
            buf.extend(payload);
            for (another_user, mut thr_sock) in &users {
                if user != another_user {
                    thr_sock.write(buf.as_slice()).unwrap();
                }
            }
            /* Add to local db */
            messages.push(buf);
        }
        for removal in &removals {
            println!("Removing user {}", removal);
            users.remove(removal);
        }
        if cnt > 0 && users.is_empty() {
            removedtx.send(this_channel).unwrap();
            break;
        }
    }
}

fn main() {
    let address = "0.0.0.0:8081";
    let listener = TcpListener::bind(&address).unwrap();
    let mut spawned = HashMap::new();
    let option: Option<Duration> = Some(Duration::from_millis(50));

    let addr = listener.local_addr().unwrap();
    println!("Listening for connections on {}", addr);
    let (tx, rx) = channel();
    let (removedtx, removedrx) = channel();

    for socket in listener.incoming() {
        /* Check first has anybody removed channels */
        let mut remrx = true;
        while remrx {
            match removedrx.try_recv() {
                Ok(val) => {
                    let removed_channel: String = val;
                    println!("Removing unused channel {}", removed_channel.as_str());
                    spawned.remove(&removed_channel);
                    remrx = true;
                },
                Err(_) => { remrx = false; }
            }
        }
        /* 1. Read incoming msg channel
         * 2. If it does not exist, spawn new thread
         * 3. Send socket to thread
         */
        let socket = socket.unwrap();
        let stream = socket.try_clone().unwrap();
        let tuple = read_n(&stream, 4);
        let status = tuple.0;
        match status {
            Ok(0) => {
               continue;
            },
            Ok(_) => {},
            _ => {
               continue;
            },
        }
        let buf = tuple.1;
        if 0 == buf.len() {
            continue;
        }
        if read_hdr_type(buf.as_slice()) != 'M' as u32 {
            println!("Incorrect payload type");
            continue;
        }
        let tuple = read_n(&stream, read_hdr_len(buf.as_slice()) as u64);
        let status = tuple.0;
        match status {
            Ok(0) => {
               continue;
            },
            Ok(_) => {},
            _ => {
               continue;
            },
        }
        let buf = tuple.1;
        let decoded_msg = message_decode(buf.as_slice());
        if 0 == decoded_msg.channel.len() {
            continue;
        }
        socket.set_nodelay(true);
        socket.set_read_timeout(option);
        if !spawned.contains_key(decoded_msg.channel.as_str()) {
            let tx = tx.clone();
            let removedtx = removedtx.clone();
            let this_channel = decoded_msg.channel.clone();
            thread::spawn(move|| process_channel(tx, removedtx, this_channel));

            let thr_feed = rx.recv().unwrap();
            println!("Channel {}", decoded_msg.channel.as_str());
            spawned.insert(decoded_msg.channel.clone(), thr_feed);
        }
        let thr_socket = spawned.get_mut(&decoded_msg.channel).unwrap();
        thr_socket.send(socket).unwrap();
    }
}
<commit_msg>Drop old main.<commit_after><|endoftext|>


<commit_before>#[cfg(feature = "trace_log")]
extern crate env_logger;

extern crate radeco_lib;
extern crate r2pipe;
extern crate r2api;
extern crate base64;
extern crate rustyline;

use r2pipe::{R2, R2Pipe};
use radeco_lib::analysis::cse::cse::CSE;
use radeco_lib::analysis::interproc::fixcall::CallFixer;
use radeco_lib::analysis::sccp;
use radeco_lib::backend::lang_c::c_cfg::CCFGVerifier;
use radeco_lib::backend::lang_c::c_cfg::ctrl_flow_struct;
use radeco_lib::backend::lang_c::c_cfg_builder;
use radeco_lib::frontend::radeco_containers::*;
use radeco_lib::middle::{dce, dot};
use radeco_lib::middle::ir_writer;
use radeco_lib::middle::regfile::SubRegisterFile;
use radeco_lib::middle::ssa::ssastorage::SSAStorage;
use radeco_lib::middle::ssa::verifier;
use rustyline::{CompletionType, Config, EditMode, Editor};
use rustyline::completion::Completer;
use rustyline::completion::FilenameCompleter;
use rustyline::error::ReadlineError;
use std::cell::RefCell;
use std::collections::HashMap;
use std::env;

use std::fs;
use std::process;
use std::rc::Rc;
use std::str;
use std::sync::Arc;

mod scheme {
    pub const HTTP: &'static str = "http://";
    pub const TCP: &'static str = "tcp://";
}

// On unix platforms you can use ANSI escape sequences
#[cfg(unix)]
const PROMPT: &'static str = "\x1b[1;32m>>\x1b[0m ";

// Windows consoles typically don't support ANSI escape sequences out
// of the box
#[cfg(windows)]
const PROMPT: &'static str = ">> ";

#[derive(Default)]
struct Completes {
    file_completer: FilenameCompleter,
}

impl Completer for Completes {
    fn complete(&self, line: &str, _pos: usize) -> rustyline::Result<(usize, Vec<String>)> {
        // TODO Completion for function names
        let cmds = vec![
            command::HELP,
            command::LOAD,
            command::CONNECT,
            command::FNLIST,
            command::ANALYZE,
            command::DOT,
            command::IR,
            command::DECOMPILE,
        ];
        let mut ret: Vec<String> = cmds.into_iter()
            .filter(|s| s.starts_with(line))
            .map(|s| s.to_string())
            .collect();
        match self.file_completer.complete(line, _pos) {
            Ok((n, mut ss)) => {
                let mut completed_lines = ss.into_iter()
                    .map(|s| format!("{}{}", &line[..n], s))
                    .collect();
                ret.append(&mut completed_lines);
            }
            Err(_) => {}
        }
        Ok((0, ret))
    }
}

fn main() {
    let config = Config::builder()
        .history_ignore_space(true)
        .completion_type(CompletionType::List)
        .edit_mode(EditMode::Emacs)
        .build();
    let mut rl = Editor::with_config(config);
    rl.set_helper(Some((Completes::default(), ())));
    let mut proj = env::args().nth(1).and_then(|ref s| {
        if is_http(s) {
            load_proj_http(&s[scheme::HTTP.len()..])
        } else if is_tcp(s) {
            load_proj_tcp(&s[scheme::TCP.len()..])
        } else if is_file(s) {
            Ok(load_proj_by_path(s))
        } else {
            Err("Invalid argument")
        }.ok()
    });
    loop {
        let readline = rl.readline(PROMPT);
        match readline {
            Ok(line) => {
                let mut terms = line.split_whitespace();
                let o1 = terms.next();
                let o2 = terms.next();
                cmd(o1, o2, &mut proj);
            }
            Err(ReadlineError::Interrupted) |
            Err(ReadlineError::Eof) => break,
            Err(err) => {
                println!("Error: {:?}", err);
                break;
            }
        }
    }
}

mod command {
    pub const HELP: &'static str = "help";
    pub const LOAD: &'static str = "load";
    pub const CONNECT: &'static str = "connect";
    pub const FNLIST: &'static str = "fn_list";
    pub const ANALYZE: &'static str = "analyze";
    pub const DOT: &'static str = "dot";
    pub const IR: &'static str = "ir";
    pub const DECOMPILE: &'static str = "decompile";
}

fn help() {
    let width = 30;
    println!("{:width$}    Show this help", command::HELP, width = width);
    println!(
        "{:width$}    Load binary",
        format!("{} path", command::LOAD),
        width = width
    );
    println!(
        "{:width$}    Connect to radare2 server",
        format!("{} (http|tcp)://<url>", command::CONNECT),
        width = width
    );
    println!(
        "{:width$}    Show function list",
        command::FNLIST,
        width = width
    );
    println!(
        "{:width$}    Analyze <func>",
        format!("{} <func>", command::ANALYZE),
        width = width
    );
    println!(
        "{:width$}    Emit IR of <func>",
        format!("{} <func>", command::IR),
        width = width
    );
    println!(
        "{:width$}    Emit graph of the IR in Graphviz dot",
        format!("{} <func>", command::DOT),
        width = width
    );
    println!(
        "{:width$}    Decompile <func>",
        format!("{} <func>", command::DECOMPILE),
        width = width
    );
}

fn cmd(op1: Option<&str>, op2: Option<&str>, proj_opt: &mut Option<RadecoProject>) {
    match (op1, op2) {
        (Some(command::HELP), _) => {
            help();
            return;
        }
        (Some(command::LOAD), Some(path)) => {
            if is_file(path) {
                *proj_opt = Some(load_proj_by_path(path));
                return;
            } else {
                println!("{} is not found.", path);
                return;
            }
        }
        (Some(command::CONNECT), Some(url)) => {
            let p_opt = if is_http(&url) {
                load_proj_http(&url[scheme::HTTP.len()..])
            } else if is_tcp(&url) {
                load_proj_tcp(&url[scheme::TCP.len()..])
            } else {
                Err("Invalid url")
            };
            match p_opt {
                Ok(proj) => *proj_opt = Some(proj),
                Err(msg) => println!("{}", msg),
            }
            return;
        }
        _ => {}
    };
    let mut proj = match proj_opt {
        Some(proj) => proj,
        None => {
            println!("Load a project first");
            return;
        }
    };
    match (op1, op2) {
        (Some(command::ANALYZE), Some("*")) => {
            let rfns = proj.iter_mut().map(|i| i.module).flat_map(|rmod| {
                rmod.functions.values_mut()
            });
            for rfn in rfns {
                analyze(rfn);
            }
        }
        (Some(command::FNLIST), _) => {
            let funcs = fn_list(&proj);
            println!("{}", funcs.join("\n"));
        }
        // TODO Show list of dependency information of analyses
        // TODO Add command for individual analyses
        (Some(command::ANALYZE), Some(f)) => {
            if let Some(rfn) = get_function_mut(f, &mut proj) {
                analyze(rfn);
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::DOT), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                println!("{}", emit_dot(rfn.ssa()));
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::IR), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                println!("{}", emit_ir(rfn));
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::DECOMPILE), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                let rmod = proj.iter().map(|i| i.module).next().unwrap();
                let func_name_map = func_names(&rmod);
                let strings = strings(&rmod);
                println!("{}", decompile(rfn, &func_name_map, &strings));
            } else {
                println!("{} is not found", f);
            }
        }
        _ => {
            println!(
                "Invalid command {} {}",
                op1.unwrap_or(""),
                op2.unwrap_or("")
            );
        }
    }
}

fn load_proj_by_path(path: &str) -> RadecoProject {
    let mut p = ProjectLoader::new().path(path).load();
    let regfile = p.regfile().clone();
    for mut xy in p.iter_mut() {
        analyze_mod(regfile.clone(), xy.module);
    }
    p
}

fn load_proj_tcp(url: &str) -> Result<RadecoProject, &'static str> {
    let r2p = R2Pipe::tcp(url)?;
    Ok(load_project_by_r2pipe(r2p))
}

fn load_proj_http(url: &str) -> Result<RadecoProject, &'static str> {
    let r2p = R2Pipe::http(url)?;
    Ok(load_project_by_r2pipe(r2p))
}

fn load_project_by_r2pipe(r2p: R2Pipe) -> RadecoProject {
    let r2 = R2::from(r2p);
    let r2w = Rc::new(RefCell::new(r2));
    let mut p = ProjectLoader::new().source(Rc::new(r2w)).load();
    let regfile = p.regfile().clone();
    for mut xy in p.iter_mut() {
        analyze_mod(regfile.clone(), xy.module);
    }
    p
}

fn fn_list(proj: &RadecoProject) -> Vec<String> {
    proj.iter()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values())
        .map(|s| s.name.to_string())
        .collect()
}

fn get_function<'a>(name: &str, proj: &'a RadecoProject) -> Option<&'a RadecoFunction> {
    proj.iter()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values())
        .filter(|rfn| rfn.name == name)
        .next()
}

fn get_function_mut<'a>(name: &str, proj: &'a mut RadecoProject) -> Option<&'a mut RadecoFunction> {
    proj.iter_mut()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values_mut())
        .filter(|rfn| rfn.name == name)
        .next()
}

fn analyze_mod(regfile: Arc<SubRegisterFile>, rmod: &mut RadecoModule) {
    // Analyze preserved for all functions.
    {
        println!("[*] Fixing Callee Information");
        let bp_name = regfile.get_name_by_alias(&"BP".to_string());
        let bp_name = bp_name.map(|s| s.to_owned());
        let sp_name = regfile.get_name_by_alias(&"SP".to_string());
        let sp_name = sp_name.map(|s| s.to_owned());
        let mut callfixer = CallFixer::new(rmod, bp_name, sp_name);
        callfixer.rounded_analysis();
    }

    // Fix call sites
    radeco_lib::analysis::functions::fix_ssa_opcalls::go(rmod);

    // Infer calling conventions
    radeco_lib::analysis::functions::infer_regusage::run(rmod, &*regfile);
}

fn analyze(rfn: &mut RadecoFunction) {
    println!("[+] Analyzing: {} @ {:#x}", rfn.name, rfn.offset);
    {
        println!("  [*] Eliminating Dead Code");
        dce::collect(rfn.ssa_mut());
    }
    let mut ssa = {
        // Constant Propagation (sccp)
        println!("  [*] Propagating Constants");
        let mut analyzer = sccp::Analyzer::new(rfn.ssa_mut());
        analyzer.analyze();
        analyzer.emit_ssa()
    };
    {
        println!("  [*] Eliminating More DeadCode");
        dce::collect(&mut ssa);
    }
    *rfn.ssa_mut() = ssa;
    {
        // Common SubExpression Elimination (cse)
        println!("  [*] Eliminating Common SubExpressions");
        let mut cse = CSE::new(rfn.ssa_mut());
        cse.run();
    }
    {
        // Verify SSA
        println!("  [*] Verifying SSA's Validity");
        match verifier::verify(rfn.ssa()) {
            Err(e) => {
                println!("  [*] Found Error: {}", e);
                process::exit(255);
            }
            Ok(_) => {}
        }
    }
}

fn emit_ir(rfn: &RadecoFunction) -> String {
    println!("  [*] Writing out IR");
    let mut res = String::new();
    ir_writer::emit_il(&mut res, Some(rfn.name.to_string()), rfn.ssa()).unwrap();
    res
}

fn emit_dot(ssa: &SSAStorage) -> String {
    dot::emit_dot(ssa)
}

fn decompile(
    rfn: &RadecoFunction,
    func_name_map: &HashMap<u64, String>,
    strings: &HashMap<u64, String>,
) -> String {
    let c_cfg = c_cfg_builder::recover_c_cfg(rfn, func_name_map, strings);
    if let Err(err) = CCFGVerifier::verify(&c_cfg) {
        println!("CCFG verification failed");
        println!("{}", err);
    }
    ctrl_flow_struct::structure_and_convert(c_cfg)
        .unwrap()
        .print()
}

fn is_file(path: &str) -> bool {
    fs::metadata(path).map(|f| f.is_file()).unwrap_or(false)
}

fn is_http(url: &str) -> bool {
    url.starts_with(scheme::HTTP)
}

fn is_tcp(url: &str) -> bool {
    url.starts_with(scheme::TCP)
}

fn func_names(rmod: &RadecoModule) -> HashMap<u64, String> {
    rmod.functions
        .iter()
        .map(|(&addr, f)| (addr, f.name.to_string()))
        .collect()
}

fn strings(rmod: &RadecoModule) -> HashMap<u64, String> {
    rmod.strings()
        .iter()
        .filter(|ref s| s.paddr.is_some() && s.string.is_some())
        .cloned()
        .map(|s| {
            let (addr, _s) = (s.vaddr.unwrap(), s.string.unwrap());
            let bytes = base64::decode(&_s).unwrap_or(Vec::new());
            let ret_string = match str::from_utf8(bytes.as_slice()) {
                Ok(v) => v.to_string(),
                Err(_e) => _s,
            };
            (addr, ret_string)
        })
        .collect()
}
<commit_msg>Enable history<commit_after>#[cfg(feature = "trace_log")]
extern crate env_logger;

extern crate radeco_lib;
extern crate r2pipe;
extern crate r2api;
extern crate base64;
extern crate rustyline;

use r2pipe::{R2, R2Pipe};
use radeco_lib::analysis::cse::cse::CSE;
use radeco_lib::analysis::interproc::fixcall::CallFixer;
use radeco_lib::analysis::sccp;
use radeco_lib::backend::lang_c::c_cfg::CCFGVerifier;
use radeco_lib::backend::lang_c::c_cfg::ctrl_flow_struct;
use radeco_lib::backend::lang_c::c_cfg_builder;
use radeco_lib::frontend::radeco_containers::*;
use radeco_lib::middle::{dce, dot};
use radeco_lib::middle::ir_writer;
use radeco_lib::middle::regfile::SubRegisterFile;
use radeco_lib::middle::ssa::ssastorage::SSAStorage;
use radeco_lib::middle::ssa::verifier;
use rustyline::{CompletionType, Config, EditMode, Editor};
use rustyline::completion::Completer;
use rustyline::completion::FilenameCompleter;
use rustyline::error::ReadlineError;
use std::cell::RefCell;
use std::collections::HashMap;
use std::env;

use std::fs;
use std::process;
use std::rc::Rc;
use std::str;
use std::sync::Arc;

mod scheme {
    pub const HTTP: &'static str = "http://";
    pub const TCP: &'static str = "tcp://";
}

// On unix platforms you can use ANSI escape sequences
#[cfg(unix)]
const PROMPT: &'static str = "\x1b[1;32m>>\x1b[0m ";

// Windows consoles typically don't support ANSI escape sequences out
// of the box
#[cfg(windows)]
const PROMPT: &'static str = ">> ";

#[derive(Default)]
struct Completes {
    file_completer: FilenameCompleter,
}

impl Completer for Completes {
    fn complete(&self, line: &str, _pos: usize) -> rustyline::Result<(usize, Vec<String>)> {
        // TODO Completion for function names
        let cmds = vec![
            command::HELP,
            command::LOAD,
            command::CONNECT,
            command::FNLIST,
            command::ANALYZE,
            command::DOT,
            command::IR,
            command::DECOMPILE,
        ];
        let mut ret: Vec<String> = cmds.into_iter()
            .filter(|s| s.starts_with(line))
            .map(|s| s.to_string())
            .collect();
        match self.file_completer.complete(line, _pos) {
            Ok((n, mut ss)) => {
                let mut completed_lines = ss.into_iter()
                    .map(|s| format!("{}{}", &line[..n], s))
                    .collect();
                ret.append(&mut completed_lines);
            }
            Err(_) => {}
        }
        Ok((0, ret))
    }
}

fn main() {
    let config = Config::builder()
        .auto_add_history(true)
        .history_ignore_space(true)
        .completion_type(CompletionType::List)
        .edit_mode(EditMode::Emacs)
        .build();
    let mut rl = Editor::with_config(config);
    rl.set_helper(Some((Completes::default(), ())));
    let mut proj = env::args().nth(1).and_then(|ref s| {
        if is_http(s) {
            load_proj_http(&s[scheme::HTTP.len()..])
        } else if is_tcp(s) {
            load_proj_tcp(&s[scheme::TCP.len()..])
        } else if is_file(s) {
            Ok(load_proj_by_path(s))
        } else {
            Err("Invalid argument")
        }.ok()
    });
    loop {
        let readline = rl.readline(PROMPT);
        match readline {
            Ok(line) => {
                let mut terms = line.split_whitespace();
                let o1 = terms.next();
                let o2 = terms.next();
                cmd(o1, o2, &mut proj);
            }
            Err(ReadlineError::Interrupted) |
            Err(ReadlineError::Eof) => break,
            Err(err) => {
                println!("Error: {:?}", err);
                break;
            }
        }
    }
}

mod command {
    pub const HELP: &'static str = "help";
    pub const LOAD: &'static str = "load";
    pub const CONNECT: &'static str = "connect";
    pub const FNLIST: &'static str = "fn_list";
    pub const ANALYZE: &'static str = "analyze";
    pub const DOT: &'static str = "dot";
    pub const IR: &'static str = "ir";
    pub const DECOMPILE: &'static str = "decompile";
}

fn help() {
    let width = 30;
    println!("{:width$}    Show this help", command::HELP, width = width);
    println!(
        "{:width$}    Load binary",
        format!("{} path", command::LOAD),
        width = width
    );
    println!(
        "{:width$}    Connect to radare2 server",
        format!("{} (http|tcp)://<url>", command::CONNECT),
        width = width
    );
    println!(
        "{:width$}    Show function list",
        command::FNLIST,
        width = width
    );
    println!(
        "{:width$}    Analyze <func>",
        format!("{} <func>", command::ANALYZE),
        width = width
    );
    println!(
        "{:width$}    Emit IR of <func>",
        format!("{} <func>", command::IR),
        width = width
    );
    println!(
        "{:width$}    Emit graph of the IR in Graphviz dot",
        format!("{} <func>", command::DOT),
        width = width
    );
    println!(
        "{:width$}    Decompile <func>",
        format!("{} <func>", command::DECOMPILE),
        width = width
    );
}

fn cmd(op1: Option<&str>, op2: Option<&str>, proj_opt: &mut Option<RadecoProject>) {
    match (op1, op2) {
        (Some(command::HELP), _) => {
            help();
            return;
        }
        (Some(command::LOAD), Some(path)) => {
            if is_file(path) {
                *proj_opt = Some(load_proj_by_path(path));
                return;
            } else {
                println!("{} is not found.", path);
                return;
            }
        }
        (Some(command::CONNECT), Some(url)) => {
            let p_opt = if is_http(&url) {
                load_proj_http(&url[scheme::HTTP.len()..])
            } else if is_tcp(&url) {
                load_proj_tcp(&url[scheme::TCP.len()..])
            } else {
                Err("Invalid url")
            };
            match p_opt {
                Ok(proj) => *proj_opt = Some(proj),
                Err(msg) => println!("{}", msg),
            }
            return;
        }
        _ => {}
    };
    let mut proj = match proj_opt {
        Some(proj) => proj,
        None => {
            println!("Load a project first");
            return;
        }
    };
    match (op1, op2) {
        (Some(command::ANALYZE), Some("*")) => {
            let rfns = proj.iter_mut().map(|i| i.module).flat_map(|rmod| {
                rmod.functions.values_mut()
            });
            for rfn in rfns {
                analyze(rfn);
            }
        }
        (Some(command::FNLIST), _) => {
            let funcs = fn_list(&proj);
            println!("{}", funcs.join("\n"));
        }
        // TODO Show list of dependency information of analyses
        // TODO Add command for individual analyses
        (Some(command::ANALYZE), Some(f)) => {
            if let Some(rfn) = get_function_mut(f, &mut proj) {
                analyze(rfn);
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::DOT), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                println!("{}", emit_dot(rfn.ssa()));
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::IR), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                println!("{}", emit_ir(rfn));
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::DECOMPILE), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                let rmod = proj.iter().map(|i| i.module).next().unwrap();
                let func_name_map = func_names(&rmod);
                let strings = strings(&rmod);
                println!("{}", decompile(rfn, &func_name_map, &strings));
            } else {
                println!("{} is not found", f);
            }
        }
        _ => {
            println!(
                "Invalid command {} {}",
                op1.unwrap_or(""),
                op2.unwrap_or("")
            );
        }
    }
}

fn load_proj_by_path(path: &str) -> RadecoProject {
    let mut p = ProjectLoader::new().path(path).load();
    let regfile = p.regfile().clone();
    for mut xy in p.iter_mut() {
        analyze_mod(regfile.clone(), xy.module);
    }
    p
}

fn load_proj_tcp(url: &str) -> Result<RadecoProject, &'static str> {
    let r2p = R2Pipe::tcp(url)?;
    Ok(load_project_by_r2pipe(r2p))
}

fn load_proj_http(url: &str) -> Result<RadecoProject, &'static str> {
    let r2p = R2Pipe::http(url)?;
    Ok(load_project_by_r2pipe(r2p))
}

fn load_project_by_r2pipe(r2p: R2Pipe) -> RadecoProject {
    let r2 = R2::from(r2p);
    let r2w = Rc::new(RefCell::new(r2));
    let mut p = ProjectLoader::new().source(Rc::new(r2w)).load();
    let regfile = p.regfile().clone();
    for mut xy in p.iter_mut() {
        analyze_mod(regfile.clone(), xy.module);
    }
    p
}

fn fn_list(proj: &RadecoProject) -> Vec<String> {
    proj.iter()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values())
        .map(|s| s.name.to_string())
        .collect()
}

fn get_function<'a>(name: &str, proj: &'a RadecoProject) -> Option<&'a RadecoFunction> {
    proj.iter()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values())
        .filter(|rfn| rfn.name == name)
        .next()
}

fn get_function_mut<'a>(name: &str, proj: &'a mut RadecoProject) -> Option<&'a mut RadecoFunction> {
    proj.iter_mut()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values_mut())
        .filter(|rfn| rfn.name == name)
        .next()
}

fn analyze_mod(regfile: Arc<SubRegisterFile>, rmod: &mut RadecoModule) {
    // Analyze preserved for all functions.
    {
        println!("[*] Fixing Callee Information");
        let bp_name = regfile.get_name_by_alias(&"BP".to_string());
        let bp_name = bp_name.map(|s| s.to_owned());
        let sp_name = regfile.get_name_by_alias(&"SP".to_string());
        let sp_name = sp_name.map(|s| s.to_owned());
        let mut callfixer = CallFixer::new(rmod, bp_name, sp_name);
        callfixer.rounded_analysis();
    }

    // Fix call sites
    radeco_lib::analysis::functions::fix_ssa_opcalls::go(rmod);

    // Infer calling conventions
    radeco_lib::analysis::functions::infer_regusage::run(rmod, &*regfile);
}

fn analyze(rfn: &mut RadecoFunction) {
    println!("[+] Analyzing: {} @ {:#x}", rfn.name, rfn.offset);
    {
        println!("  [*] Eliminating Dead Code");
        dce::collect(rfn.ssa_mut());
    }
    let mut ssa = {
        // Constant Propagation (sccp)
        println!("  [*] Propagating Constants");
        let mut analyzer = sccp::Analyzer::new(rfn.ssa_mut());
        analyzer.analyze();
        analyzer.emit_ssa()
    };
    {
        println!("  [*] Eliminating More DeadCode");
        dce::collect(&mut ssa);
    }
    *rfn.ssa_mut() = ssa;
    {
        // Common SubExpression Elimination (cse)
        println!("  [*] Eliminating Common SubExpressions");
        let mut cse = CSE::new(rfn.ssa_mut());
        cse.run();
    }
    {
        // Verify SSA
        println!("  [*] Verifying SSA's Validity");
        match verifier::verify(rfn.ssa()) {
            Err(e) => {
                println!("  [*] Found Error: {}", e);
                process::exit(255);
            }
            Ok(_) => {}
        }
    }
}

fn emit_ir(rfn: &RadecoFunction) -> String {
    println!("  [*] Writing out IR");
    let mut res = String::new();
    ir_writer::emit_il(&mut res, Some(rfn.name.to_string()), rfn.ssa()).unwrap();
    res
}

fn emit_dot(ssa: &SSAStorage) -> String {
    dot::emit_dot(ssa)
}

fn decompile(
    rfn: &RadecoFunction,
    func_name_map: &HashMap<u64, String>,
    strings: &HashMap<u64, String>,
) -> String {
    let c_cfg = c_cfg_builder::recover_c_cfg(rfn, func_name_map, strings);
    if let Err(err) = CCFGVerifier::verify(&c_cfg) {
        println!("CCFG verification failed");
        println!("{}", err);
    }
    ctrl_flow_struct::structure_and_convert(c_cfg)
        .unwrap()
        .print()
}

fn is_file(path: &str) -> bool {
    fs::metadata(path).map(|f| f.is_file()).unwrap_or(false)
}

fn is_http(url: &str) -> bool {
    url.starts_with(scheme::HTTP)
}

fn is_tcp(url: &str) -> bool {
    url.starts_with(scheme::TCP)
}

fn func_names(rmod: &RadecoModule) -> HashMap<u64, String> {
    rmod.functions
        .iter()
        .map(|(&addr, f)| (addr, f.name.to_string()))
        .collect()
}

fn strings(rmod: &RadecoModule) -> HashMap<u64, String> {
    rmod.strings()
        .iter()
        .filter(|ref s| s.paddr.is_some() && s.string.is_some())
        .cloned()
        .map(|s| {
            let (addr, _s) = (s.vaddr.unwrap(), s.string.unwrap());
            let bytes = base64::decode(&_s).unwrap_or(Vec::new());
            let ret_string = match str::from_utf8(bytes.as_slice()) {
                Ok(v) => v.to_string(),
                Err(_e) => _s,
            };
            (addr, ret_string)
        })
        .collect()
}
<|endoftext|>