<commit_before>use std::str::Chars;
use errors::Reporter;
use position::{Position, Span, Spanned};
use tokens::{Token, Keyword, Number};
use ast::{Signedness, Size};

pub fn lex(source: &str, reporter: Reporter) -> Vec<Spanned<Token>> {
    let mut lexer = Lexer::new(source, reporter);
    let mut tokens = Vec::new();
    while let Some(tok) = lexer.next_token() {
        tokens.push(tok);
    }
    tokens
}

struct Lexer<'a> {
    chars: Chars<'a>,
    next_char: Option<char>,
    current_pos: Position,
    reporter: Reporter,
    previous_error: bool,
}

impl<'a> Lexer<'a> {
    fn new(source: &'a str, reporter: Reporter) -> Self {
        let mut chars = source.chars();
        let next_char = chars.next();
        Lexer {
            chars,
            next_char,
            current_pos: Position::new(1, 1),
            reporter,
            previous_error: false,
        }
    }

    fn peek(&self) -> Option<char> {
        self.next_char
    }

    fn advance(&mut self) {
        match self.next_char {
            Some('\n') => {
                self.current_pos.line += 1;
                self.current_pos.column = 1;
            }
            Some(_) => {
                self.current_pos.column += 1;
            }
            None => {}
        }
        self.next_char = self.chars.next();
    }

    fn check(&mut self, ch: char) -> bool {
        if self.peek() == Some(ch) {
            self.advance();
            true
        } else {
            false
        }
    }

    fn consume(&mut self) -> Option<char> {
        let result = self.peek();
        self.advance();
        result
    }

    fn consume_hex_digit(&mut self) -> Option<u8> {
        match self.peek() {
            Some(ch) if ch.is_digit(16) => {
                self.consume().map(|c| c.to_digit(16).unwrap() as u8)
            }
            _ => {
                None
            }
        }
    }

    fn single_char(&mut self, tok: Token) -> Spanned<Token> {
        let start = self.current_pos;
        self.advance();
        let span = start.span_to(self.current_pos);
        Spanned::new(tok, span)
    }

    fn two_char(&mut self, ch: char, tok: Token) -> Spanned<Token> {
        let start = self.current_pos;
        let first = self.consume().expect("char disappeared");
        let tok = if self.check(ch) {
            tok
        } else {
            let span = start.span_to(self.current_pos.forward(1));
            self.reporter
                .error(format!("unknown token: `{}`", first), span)
                .span_note(span, format!("maybe you wanted `{}{}`?", first, ch))
                .build();
            Token::Error
        };
        let span = start.span_to(self.current_pos);
        Spanned::new(tok, span)
    }

    fn test_second(&mut self, ch: char, if_one: Token, if_two: Token) -> Spanned<Token> {
        let start = self.current_pos;
        self.advance();
        let tok = if self.check(ch) { if_two } else { if_one };
        let span = start.span_to(self.current_pos);
        Spanned::new(tok, span)
    }

    fn skip_line_comment(&mut self) {
        loop {
            match self.consume() {
                Some('\n') | None => break,
                _ => {}
            }
        }
    }

    fn skip_block_comment(&mut self) {
        // make span of starting `/*`
        let opener_span = Span::new(
            self.current_pos.backwards(2),
            self.current_pos
        );
        let mut depth = 1;
        while depth > 0 {
            match self.consume() {
                Some('*') => {
                    if self.consume() == Some('/') {
                        depth -= 1;
                    }
                }
                Some('/') => {
                    if self.consume() == Some('*') {
                        depth += 1;
                    }
                }
                None => {
                    self.reporter
                        .error("unterminated block comment", opener_span)
                        .span(opener_span)
                        .build();
                    break;
                }
                _ => {}
            }
        }
    }

    fn slash_or_comment(&mut self) -> Option<Spanned<Token>> {
        let start = self.current_pos;
        self.advance();
        if self.check('/') {
            self.skip_line_comment();
            None
        } else if self.check('*') {
            self.skip_block_comment();
            None
        } else {
            let span = start.span_to(self.current_pos);
            Some(Spanned::new(Token::Slash, span))
        }
    }

    fn lex_number(&mut self) -> Spanned<Token> {
        let mut string = String::new();
        let start = self.current_pos;
        loop {
            match self.peek() {
                Some(ch) if ch.is_alphanumeric() => {
                    string.push(ch);
                    self.advance();
                }
                _ => break,
            }
        }
        let span = start.span_to(self.current_pos);
        let tok = match parse_number(&string) {
            Ok(num) => Token::Number(num),
            Err(err) => {
                let msg = err.to_string();
                self.reporter
                    .error(msg, span)
                    .span(span)
                    .build();
                Token::Error
            }
        };
        Spanned::new(tok, span)
    }

    fn lex_name(&mut self) -> Spanned<Token> {
        let mut string = String::new();
        let start = self.current_pos;
        loop {
            match self.peek() {
                Some(ch) if is_ident_char(ch) => {
                    string.push(ch);
                    self.advance();
                }
                _ => break,
            }
        }
        let span = start.span_to(self.current_pos);
        let tok = if let Some(tok) = keyword(&string) {
            tok
        } else {
            Token::Ident(string)
        };
        Spanned::new(tok, span)
    }

    fn lex_string_raw(&mut self, closing: char) -> Spanned<Option<Vec<u8>>> {
        let mut string = Vec::new();
        let start_pos = self.current_pos;
        self.advance();
        loop {
            match self.peek() {
                Some(ch) if ch == closing => {
                    self.advance();
                    let span = start_pos.span_to(self.current_pos);
                    return Spanned::new(Some(string), span);
                }
                Some('\\') => {
                    let escape_start = self.current_pos;
                    self.advance();
                    match self.consume() {
                        Some('\\') => string.push(b'\\'),
                        Some('\'') => string.push(b'\''),
                        Some('"') => string.push(b'"'),
                        Some('n') => string.push(b'\n'),
                        Some('x') => {
                            let byte = self
                                .consume_hex_digit()
                                .and_then(|high| {
                                    self.consume_hex_digit()
                                        .map(|low| (high << 4) | low)
                                });
                            match byte {
                                Some(byte) => string.push(byte),
                                None => {
                                    let span = self.current_pos.span_to(self.current_pos.forward(1));
                                    self.reporter
                                        .error("`x` should be followed by two hex digits", span)
                                        .span(span)
                                        .build();
                                }
                            }
                        }
                        _ => {
                            let span = escape_start.span_to(self.current_pos);
                            self.reporter
                                .error("invalid escape sequence", span)
                                .span(span)
                                .build();
                        }
                    }
                }
                Some('\n') | Some('\r') | None => {
                    let span = start_pos.span_to(self.current_pos);
                    self.reporter
                        .error("unterminated string", span)
                        .span(span)
                        .build();
                    return Spanned::new(None, span);
                }
                Some(ch) => {
                    let ch = ch as u32;
                    if 32 <= ch && ch < 127 {
                        string.push(ch as u8);
                        self.advance();
                    } else {
                        let start = self.current_pos;
                        self.advance();
                        let span = start.span_to(self.current_pos);
                        self.reporter
                            .error(format!("unknown char in string (codepoint: {})", ch), span)
                            .span(span)
                            .build();
                    }
                }
            }
        }
    }

    fn lex_string(&mut self) -> Spanned<Token> {
        Spanned::map(self.lex_string_raw('"'), |t| {
            t.map(Token::Str).unwrap_or(Token::Error)
        })
    }

    fn lex_char(&mut self) -> Spanned<Token> {
        let tok = self.lex_string_raw('\'');
        let span = Spanned::span(&tok);
        if let Some(value) = Spanned::into_value(tok) {
            if value.len() != 1 {
                self.reporter
                    .error(format!("char literal must have 1 char, but it has {}", value.len()), span)
                    .span(span)
                    .build();
                Spanned::new(Token::Error, span)
            } else {
                Spanned::new(Token::Char(value[0]), span)
            }
        } else {
            Spanned::new(Token::Error, span)
        }
    }

    fn next_raw_token(&mut self) -> Option<Spanned<Token>> {
        loop {
            return Some(match self.peek() {
                Some(' ') | Some('\t') | Some('\r') | Some('\n') => {
                    self.advance();
                    continue;
                }
                Some(ch) if ch.is_digit(10) => self.lex_number(),
                Some(ch) if is_ident_char(ch) => self.lex_name(),
                Some('.') => self.single_char(Token::Dot),
                Some(',') => self.single_char(Token::Comma),
                Some(';') => self.single_char(Token::Semicolon),
                Some('*') => self.single_char(Token::Star),
                Some('+') => self.single_char(Token::Plus),
                Some('%') => self.single_char(Token::Percent),
                Some('(') => self.single_char(Token::LeftParen),
                Some(')') => self.single_char(Token::RightParen),
                Some('{') => self.single_char(Token::LeftBrace),
                Some('}') => self.single_char(Token::RightBrace),
                Some('|') => self.two_char('|', Token::Or),
                Some('!') => self.test_second('=', Token::Not, Token::NotEqual),
                Some('&') => self.test_second('&', Token::Ampersand, Token::And),
                Some(':') => self.test_second(':', Token::Colon, Token::DoubleColon),
                Some('<') => self.test_second('=', Token::Less, Token::LessEqual),
                Some('>') => self.test_second('=', Token::Greater, Token::GreaterEqual),
                Some('-') => self.test_second('>', Token::Minus, Token::Arrow),
                Some('=') => self.test_second('=', Token::Assign, Token::Equal),
                Some('/') => {
                    if let Some(tok) = self.slash_or_comment() {
                        tok
                    } else {
                        continue;
                    }
                }
                Some('"') => self.lex_string(),
                Some('\'') => self.lex_char(),
                Some(ch) => {
                    let tok = self.single_char(Token::Error);
                    if !self.previous_error {
                        let codepoint = ch as u32;
                        let msg = if 32 <= codepoint && codepoint < 127 {
                            format!("unknown char: `{}`", ch)
                        } else {
                            format!("unknown char (codepoint: {})", codepoint)
                        };
                        let span = Spanned::span(&tok);
                        self.reporter
                            .error(msg, span)
                            .span(span)
                            .build();
                    }
                    tok
                }
                None => return None,
            });
        }
    }

    fn next_token(&mut self) -> Option<Spanned<Token>> {
        while let Some(tok) = self.next_raw_token() {
            if Spanned::value(&tok) == &Token::Error {
                // don't report two lexing errors in a row
                if !self.previous_error {
                    return Some(tok)
                }
            } else {
                self.previous_error = false;
                return Some(tok);
            }
        }
        None
    }
}

fn is_ident_char(ch: char) -> bool {
    // TODO: grammar allows only /[a-zA-Z0-9_]/,
    // but here we allow funky unicode stuff
    ch.is_alphanumeric() || ch == '_'
}

fn keyword(s: &str) -> Option<Token> {
    match s {
        "extern" => Some(Token::Keyword(Keyword::Extern)),
        "fn" => Some(Token::Keyword(Keyword::Fn)),
        "struct" => Some(Token::Keyword(Keyword::Struct)),
        "if" => Some(Token::Keyword(Keyword::If)),
        "else" => Some(Token::Keyword(Keyword::Else)),
        "loop" => Some(Token::Keyword(Keyword::Loop)),
        "while" => Some(Token::Keyword(Keyword::While)),
        "continue" => Some(Token::Keyword(Keyword::Continue)),
        "break" => Some(Token::Keyword(Keyword::Break)),
        "let" => Some(Token::Keyword(Keyword::Let)),
        "return" => Some(Token::Keyword(Keyword::Return)),
        "true" => Some(Token::Bool(true)),
        "false" => Some(Token::Bool(false)),
        "_" => Some(Token::Underscore),
        "i8" => Some(Token::Keyword(Keyword::I8)),
        "u8" => Some(Token::Keyword(Keyword::U8)),
        "i16" => Some(Token::Keyword(Keyword::I16)),
        "u16" => Some(Token::Keyword(Keyword::U16)),
        "i32" => Some(Token::Keyword(Keyword::I32)),
        "u32" => Some(Token::Keyword(Keyword::U32)),
        "bool" => Some(Token::Keyword(Keyword::Bool)),
        _ => None,
    }
}

enum ParseNumberError {
    BadInt,
    TooLarge,
    BadBitCount,
}

impl ::std::fmt::Display for ParseNumberError {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let msg = match *self {
            ParseNumberError::BadInt => "invalid int literal",
            ParseNumberError::TooLarge => "int literal is too big",
            ParseNumberError::BadBitCount => "invalid int size",
        };
        write!(f, "{}", msg)
    }
}

fn parse_number(s: &str) -> Result<Number, ParseNumberError> {
    match s.find(|c| c == 'i' || c == 'u') {
        Some(index) => {
            // we have a suffix
            // parse prefix as number, parse rest as bit count
            let prefix = s.get(..index).unwrap();
            let suffix = s.get((index + 1)..).unwrap();
            let value = parse_simple_number(prefix)?;
            let signedness = if s.get(index..(index + 1)).unwrap() == "i" {
                Signedness::Signed
            } else {
                Signedness::Unsigned
            };
            let size = match suffix {
                "8" => Size::Bit8,
                "16" => Size::Bit16,
                "32" => Size::Bit32,
                _ => return Err(ParseNumberError::BadBitCount),
            };
            Ok(Number {
                value,
                signedness: Some(signedness),
                size: Some(size),
            })
        }
        None => {
            // no suffix, parse simple number
            let value = parse_simple_number(s)?;
            Ok(Number {
                value,
                signedness: None,
                size: None,
            })
        }
    }
}

fn parse_simple_number(s: &str) -> Result<u64, ParseNumberError> {
    if s.len() == 0 {
        return Err(ParseNumberError::BadInt);
    }
    if s.len() > 1 && s.chars().next() == Some('0') {
        // we have a leading zero, and number is longer than one digit
        return Err(ParseNumberError::BadInt);
    }
    let mut result = 0u64;
    for ch in s.chars() {
        match ch.to_digit(10) {
            Some(digit) => {
                result = result
                    .checked_mul(10)
                    .and_then(|n| n.checked_add(digit as u64))
                    .ok_or(ParseNumberError::TooLarge)?;
            }
            None => {
                return Err(ParseNumberError::BadInt);
            }
        }
    }
    Ok(result)
}
<commit_msg>fix nested block comment lexing<commit_after>use std::str::Chars;
use errors::Reporter;
use position::{Position, Span, Spanned};
use tokens::{Token, Keyword, Number};
use ast::{Signedness, Size};

pub fn lex(source: &str, reporter: Reporter) -> Vec<Spanned<Token>> {
    let mut lexer = Lexer::new(source, reporter);
    let mut tokens = Vec::new();
    while let Some(tok) = lexer.next_token() {
        tokens.push(tok);
    }
    tokens
}

struct Lexer<'a> {
    chars: Chars<'a>,
    next_char: Option<char>,
    current_pos: Position,
    reporter: Reporter,
    previous_error: bool,
}

impl<'a> Lexer<'a> {
    fn new(source: &'a str, reporter: Reporter) -> Self {
        let mut chars = source.chars();
        let next_char = chars.next();
        Lexer {
            chars,
            next_char,
            current_pos: Position::new(1, 1),
            reporter,
            previous_error: false,
        }
    }

    fn peek(&self) -> Option<char> {
        self.next_char
    }

    fn advance(&mut self) {
        match self.next_char {
            Some('\n') => {
                self.current_pos.line += 1;
                self.current_pos.column = 1;
            }
            Some(_) => {
                self.current_pos.column += 1;
            }
            None => {}
        }
        self.next_char = self.chars.next();
    }

    fn check(&mut self, ch: char) -> bool {
        if self.peek() == Some(ch) {
            self.advance();
            true
        } else {
            false
        }
    }

    fn consume(&mut self) -> Option<char> {
        let result = self.peek();
        self.advance();
        result
    }

    fn consume_hex_digit(&mut self) -> Option<u8> {
        match self.peek() {
            Some(ch) if ch.is_digit(16) => {
                self.consume().map(|c| c.to_digit(16).unwrap() as u8)
            }
            _ => {
                None
            }
        }
    }

    fn single_char(&mut self, tok: Token) -> Spanned<Token> {
        let start = self.current_pos;
        self.advance();
        let span = start.span_to(self.current_pos);
        Spanned::new(tok, span)
    }

    fn two_char(&mut self, ch: char, tok: Token) -> Spanned<Token> {
        let start = self.current_pos;
        let first = self.consume().expect("char disappeared");
        let tok = if self.check(ch) {
            tok
        } else {
            let span = start.span_to(self.current_pos.forward(1));
            self.reporter
                .error(format!("unknown token: `{}`", first), span)
                .span_note(span, format!("maybe you wanted `{}{}`?", first, ch))
                .build();
            Token::Error
        };
        let span = start.span_to(self.current_pos);
        Spanned::new(tok, span)
    }

    fn test_second(&mut self, ch: char, if_one: Token, if_two: Token) -> Spanned<Token> {
        let start = self.current_pos;
        self.advance();
        let tok = if self.check(ch) { if_two } else { if_one };
        let span = start.span_to(self.current_pos);
        Spanned::new(tok, span)
    }

    fn skip_line_comment(&mut self) {
        loop {
            match self.consume() {
                Some('\n') | None => break,
                _ => {}
            }
        }
    }

    fn skip_block_comment(&mut self) {
        // make span of starting `/*`
        let opener_span = Span::new(
            self.current_pos.backwards(2),
            self.current_pos
        );
        let mut depth = 1;
        while depth > 0 {
            match self.consume() {
                Some('*') => {
                    loop {
                        match self.consume() {
                            Some('/') => {
                                depth -= 1;
                                break;
                            }
                            Some('*') => {}
                            _ => break,
                        }
                    }
                }
                Some('/') => {
                    if self.consume() == Some('*') {
                        depth += 1;
                    }
                }
                None => {
                    self.reporter
                        .error("unterminated block comment", opener_span)
                        .span(opener_span)
                        .build();
                    break;
                }
                _ => {}
            }
        }
    }

    fn slash_or_comment(&mut self) -> Option<Spanned<Token>> {
        let start = self.current_pos;
        self.advance();
        if self.check('/') {
            self.skip_line_comment();
            None
        } else if self.check('*') {
            self.skip_block_comment();
            None
        } else {
            let span = start.span_to(self.current_pos);
            Some(Spanned::new(Token::Slash, span))
        }
    }

    fn lex_number(&mut self) -> Spanned<Token> {
        let mut string = String::new();
        let start = self.current_pos;
        loop {
            match self.peek() {
                Some(ch) if ch.is_alphanumeric() => {
                    string.push(ch);
                    self.advance();
                }
                _ => break,
            }
        }
        let span = start.span_to(self.current_pos);
        let tok = match parse_number(&string) {
            Ok(num) => Token::Number(num),
            Err(err) => {
                let msg = err.to_string();
                self.reporter
                    .error(msg, span)
                    .span(span)
                    .build();
                Token::Error
            }
        };
        Spanned::new(tok, span)
    }

    fn lex_name(&mut self) -> Spanned<Token> {
        let mut string = String::new();
        let start = self.current_pos;
        loop {
            match self.peek() {
                Some(ch) if is_ident_char(ch) => {
                    string.push(ch);
                    self.advance();
                }
                _ => break,
            }
        }
        let span = start.span_to(self.current_pos);
        let tok = if let Some(tok) = keyword(&string) {
            tok
        } else {
            Token::Ident(string)
        };
        Spanned::new(tok, span)
    }

    fn lex_string_raw(&mut self, closing: char) -> Spanned<Option<Vec<u8>>> {
        let mut string = Vec::new();
        let start_pos = self.current_pos;
        self.advance();
        loop {
            match self.peek() {
                Some(ch) if ch == closing => {
                    self.advance();
                    let span = start_pos.span_to(self.current_pos);
                    return Spanned::new(Some(string), span);
                }
                Some('\\') => {
                    let escape_start = self.current_pos;
                    self.advance();
                    match self.consume() {
                        Some('\\') => string.push(b'\\'),
                        Some('\'') => string.push(b'\''),
                        Some('"') => string.push(b'"'),
                        Some('n') => string.push(b'\n'),
                        Some('x') => {
                            let byte = self
                                .consume_hex_digit()
                                .and_then(|high| {
                                    self.consume_hex_digit()
                                        .map(|low| (high << 4) | low)
                                });
                            match byte {
                                Some(byte) => string.push(byte),
                                None => {
                                    let span = self.current_pos.span_to(self.current_pos.forward(1));
                                    self.reporter
                                        .error("`x` should be followed by two hex digits", span)
                                        .span(span)
                                        .build();
                                }
                            }
                        }
                        _ => {
                            let span = escape_start.span_to(self.current_pos);
                            self.reporter
                                .error("invalid escape sequence", span)
                                .span(span)
                                .build();
                        }
                    }
                }
                Some('\n') | Some('\r') | None => {
                    let span = start_pos.span_to(self.current_pos);
                    self.reporter
                        .error("unterminated string", span)
                        .span(span)
                        .build();
                    return Spanned::new(None, span);
                }
                Some(ch) => {
                    let ch = ch as u32;
                    if 32 <= ch && ch < 127 {
                        string.push(ch as u8);
                        self.advance();
                    } else {
                        let start = self.current_pos;
                        self.advance();
                        let span = start.span_to(self.current_pos);
                        self.reporter
                            .error(format!("unknown char in string (codepoint: {})", ch), span)
                            .span(span)
                            .build();
                    }
                }
            }
        }
    }

    fn lex_string(&mut self) -> Spanned<Token> {
        Spanned::map(self.lex_string_raw('"'), |t| {
            t.map(Token::Str).unwrap_or(Token::Error)
        })
    }

    fn lex_char(&mut self) -> Spanned<Token> {
        let tok = self.lex_string_raw('\'');
        let span = Spanned::span(&tok);
        if let Some(value) = Spanned::into_value(tok) {
            if value.len() != 1 {
                self.reporter
                    .error(format!("char literal must have 1 char, but it has {}", value.len()), span)
                    .span(span)
                    .build();
                Spanned::new(Token::Error, span)
            } else {
                Spanned::new(Token::Char(value[0]), span)
            }
        } else {
            Spanned::new(Token::Error, span)
        }
    }

    fn next_raw_token(&mut self) -> Option<Spanned<Token>> {
        loop {
            return Some(match self.peek() {
                Some(' ') | Some('\t') | Some('\r') | Some('\n') => {
                    self.advance();
                    continue;
                }
                Some(ch) if ch.is_digit(10) => self.lex_number(),
                Some(ch) if is_ident_char(ch) => self.lex_name(),
                Some('.') => self.single_char(Token::Dot),
                Some(',') => self.single_char(Token::Comma),
                Some(';') => self.single_char(Token::Semicolon),
                Some('*') => self.single_char(Token::Star),
                Some('+') => self.single_char(Token::Plus),
                Some('%') => self.single_char(Token::Percent),
                Some('(') => self.single_char(Token::LeftParen),
                Some(')') => self.single_char(Token::RightParen),
                Some('{') => self.single_char(Token::LeftBrace),
                Some('}') => self.single_char(Token::RightBrace),
                Some('|') => self.two_char('|', Token::Or),
                Some('!') => self.test_second('=', Token::Not, Token::NotEqual),
                Some('&') => self.test_second('&', Token::Ampersand, Token::And),
                Some(':') => self.test_second(':', Token::Colon, Token::DoubleColon),
                Some('<') => self.test_second('=', Token::Less, Token::LessEqual),
                Some('>') => self.test_second('=', Token::Greater, Token::GreaterEqual),
                Some('-') => self.test_second('>', Token::Minus, Token::Arrow),
                Some('=') => self.test_second('=', Token::Assign, Token::Equal),
                Some('/') => {
                    if let Some(tok) = self.slash_or_comment() {
                        tok
                    } else {
                        continue;
                    }
                }
                Some('"') => self.lex_string(),
                Some('\'') => self.lex_char(),
                Some(ch) => {
                    let tok = self.single_char(Token::Error);
                    if !self.previous_error {
                        let codepoint = ch as u32;
                        let msg = if 32 <= codepoint && codepoint < 127 {
                            format!("unknown char: `{}`", ch)
                        } else {
                            format!("unknown char (codepoint: {})", codepoint)
                        };
                        let span = Spanned::span(&tok);
                        self.reporter
                            .error(msg, span)
                            .span(span)
                            .build();
                    }
                    tok
                }
                None => return None,
            });
        }
    }

    fn next_token(&mut self) -> Option<Spanned<Token>> {
        while let Some(tok) = self.next_raw_token() {
            if Spanned::value(&tok) == &Token::Error {
                // don't report two lexing errors in a row
                if !self.previous_error {
                    return Some(tok)
                }
            } else {
                self.previous_error = false;
                return Some(tok);
            }
        }
        None
    }
}

fn is_ident_char(ch: char) -> bool {
    // TODO: grammar allows only /[a-zA-Z0-9_]/,
    // but here we allow funky unicode stuff
    ch.is_alphanumeric() || ch == '_'
}

fn keyword(s: &str) -> Option<Token> {
    match s {
        "extern" => Some(Token::Keyword(Keyword::Extern)),
        "fn" => Some(Token::Keyword(Keyword::Fn)),
        "struct" => Some(Token::Keyword(Keyword::Struct)),
        "if" => Some(Token::Keyword(Keyword::If)),
        "else" => Some(Token::Keyword(Keyword::Else)),
        "loop" => Some(Token::Keyword(Keyword::Loop)),
        "while" => Some(Token::Keyword(Keyword::While)),
        "continue" => Some(Token::Keyword(Keyword::Continue)),
        "break" => Some(Token::Keyword(Keyword::Break)),
        "let" => Some(Token::Keyword(Keyword::Let)),
        "return" => Some(Token::Keyword(Keyword::Return)),
        "true" => Some(Token::Bool(true)),
        "false" => Some(Token::Bool(false)),
        "_" => Some(Token::Underscore),
        "i8" => Some(Token::Keyword(Keyword::I8)),
        "u8" => Some(Token::Keyword(Keyword::U8)),
        "i16" => Some(Token::Keyword(Keyword::I16)),
        "u16" => Some(Token::Keyword(Keyword::U16)),
        "i32" => Some(Token::Keyword(Keyword::I32)),
        "u32" => Some(Token::Keyword(Keyword::U32)),
        "bool" => Some(Token::Keyword(Keyword::Bool)),
        _ => None,
    }
}

enum ParseNumberError {
    BadInt,
    TooLarge,
    BadBitCount,
}

impl ::std::fmt::Display for ParseNumberError {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let msg = match *self {
            ParseNumberError::BadInt => "invalid int literal",
            ParseNumberError::TooLarge => "int literal is too big",
            ParseNumberError::BadBitCount => "invalid int size",
        };
        write!(f, "{}", msg)
    }
}

fn parse_number(s: &str) -> Result<Number, ParseNumberError> {
    match s.find(|c| c == 'i' || c == 'u') {
        Some(index) => {
            // we have a suffix
            // parse prefix as number, parse rest as bit count
            let prefix = s.get(..index).unwrap();
            let suffix = s.get((index + 1)..).unwrap();
            let value = parse_simple_number(prefix)?;
            let signedness = if s.get(index..(index + 1)).unwrap() == "i" {
                Signedness::Signed
            } else {
                Signedness::Unsigned
            };
            let size = match suffix {
                "8" => Size::Bit8,
                "16" => Size::Bit16,
                "32" => Size::Bit32,
                _ => return Err(ParseNumberError::BadBitCount),
            };
            Ok(Number {
                value,
                signedness: Some(signedness),
                size: Some(size),
            })
        }
        None => {
            // no suffix, parse simple number
            let value = parse_simple_number(s)?;
            Ok(Number {
                value,
                signedness: None,
                size: None,
            })
        }
    }
}

fn parse_simple_number(s: &str) -> Result<u64, ParseNumberError> {
    if s.len() == 0 {
        return Err(ParseNumberError::BadInt);
    }
    if s.len() > 1 && s.chars().next() == Some('0') {
        // we have a leading zero, and number is longer than one digit
        return Err(ParseNumberError::BadInt);
    }
    let mut result = 0u64;
    for ch in s.chars() {
        match ch.to_digit(10) {
            Some(digit) => {
                result = result
                    .checked_mul(10)
                    .and_then(|n| n.checked_add(digit as u64))
                    .ok_or(ParseNumberError::TooLarge)?;
            }
            None => {
                return Err(ParseNumberError::BadInt);
            }
        }
    }
    Ok(result)
}
<|endoftext|>


<commit_before>use prolog::and_stack::*;
use prolog::instructions::*;

use std::ops::IndexMut;

type Trail = Vec<(Ref, HeapCellValue)>;

pub(crate) struct RedirectInfo {
    trail: Trail
}

pub(crate) trait CopierTarget: IndexMut<usize, Output=HeapCellValue>
{
    fn source(&self) -> usize;
    fn threshold(&self) -> usize;
    fn push(&mut self, HeapCellValue);
    fn store(&self, Addr) -> Addr;
    fn deref(&self, Addr) -> Addr;
    fn stack(&mut self) -> &mut AndStack;

    fn unwind_trail(&mut self, redirect: RedirectInfo)
    {
        for (r, hcv) in redirect.trail {
            match r {
                Ref::AttrVar(hc) | Ref::HeapCell(hc) => self[hc] = hcv.clone(),
                Ref::StackCell(fr, sc) => self.stack()[fr][sc] = hcv.as_addr(0)
            }
        }
    }

    fn reinstantiate_var(&mut self, ra: Addr, scan: usize, trail: &mut Trail)
    {
        match ra {
            Addr::HeapCell(hc) => {
                self[scan] = HeapCellValue::Addr(Addr::HeapCell(scan));
                self[hc] = HeapCellValue::Addr(Addr::HeapCell(scan));
                trail.push((Ref::HeapCell(hc),
                            HeapCellValue::Addr(Addr::HeapCell(hc))));
            },
            Addr::StackCell(fr, sc) => {
                self[scan] = HeapCellValue::Addr(Addr::HeapCell(scan));
                self.stack()[fr][sc] = Addr::HeapCell(scan);
                trail.push((Ref::StackCell(fr, sc),
                            HeapCellValue::Addr(Addr::StackCell(fr, sc))));
            },
            Addr::AttrVar(hc) => {
                self[scan] = HeapCellValue::Addr(Addr::AttrVar(scan));
                self[hc] = HeapCellValue::Addr(Addr::AttrVar(scan));
                trail.push((Ref::AttrVar(hc),
                            HeapCellValue::Addr(Addr::AttrVar(hc))));
            },
            _ => {}
        }
    }

    // duplicate_term_impl(L1, L2) uses Cheney's algorithm to copy the term
    // at L1 to L2. trail is kept to restore the innards of L1 after
    // it's been copied to L2.
    fn duplicate_term_impl(&mut self, addr: Addr) -> RedirectInfo
    {
        let mut trail = Trail::new();
        let mut scan = self.source();
        let old_h = self.threshold();

        self.push(HeapCellValue::Addr(addr));

        while scan < self.threshold() {
            match self[scan].clone() {
                HeapCellValue::NamedStr(..) =>
                    scan += 1,
                HeapCellValue::Addr(a) =>
                    match a.clone() {
                        Addr::Lis(a) => {
                            if let HeapCellValue::Addr(Addr::Lis(b)) = self[a].clone() {
                                if b >= old_h {
                                    self[scan] = HeapCellValue::Addr(Addr::Lis(b));
                                    scan += 1;
                                    continue;
                                }
                            }

                            let threshold = self.threshold();
                            self[scan] = HeapCellValue::Addr(Addr::Lis(threshold));

                            let hcv = self[a].clone();
                            self.push(hcv.clone());

                            let ra = hcv.as_addr(threshold);
                            let rd = self.store(self.deref(ra));

                            match rd.clone() {
                                Addr::HeapCell(hc) if hc >= old_h =>
                                    self[threshold] = HeapCellValue::Addr(rd),
                                ra @ Addr::AttrVar(_) | ra @ Addr::HeapCell(..) | ra @ Addr::StackCell(..) =>
                                    if ra == rd {
                                        self.reinstantiate_var(ra, threshold, &mut trail);
                                    } else {
                                        self[threshold] = HeapCellValue::Addr(ra);
                                    },
                                _ => {
                                    trail.push((Ref::HeapCell(a), self[a].clone()));
                                    self[a] = HeapCellValue::Addr(Addr::Lis(threshold))
                                }
                            };

                            let hcv = self[a+1].clone();
                            self.push(hcv);

                            scan += 1;
                        },
                        Addr::AttrVar(_) | Addr::HeapCell(_) | Addr::StackCell(_, _) => {
                            let ra = a;
                            let rd = self.store(self.deref(ra.clone()));

                            match rd.clone() {
                                Addr::AttrVar(h) | Addr::HeapCell(h) if h >= old_h => {
                                    self[scan] = HeapCellValue::Addr(rd);
                                    scan += 1;
                                },
                                Addr::AttrVar(h) if ra == rd => {
                                    let threshold = self.threshold();

                                    self.push(HeapCellValue::Addr(Addr::AttrVar(threshold)));

                                    let list_val = self[h+1].clone();
                                    self.push(list_val);

                                    self[scan] = HeapCellValue::Addr(Addr::AttrVar(threshold));
                                    self[h] = HeapCellValue::Addr(Addr::AttrVar(threshold));

                                    trail.push((Ref::AttrVar(h),
                                                HeapCellValue::Addr(Addr::AttrVar(h))));
                                },
                                _ if ra == rd => {
                                    self.reinstantiate_var(ra, scan, &mut trail);
                                    scan += 1;
                                },
                                _ => self[scan] = HeapCellValue::Addr(rd)
                            };
                        },
                        Addr::Str(s) => {
                            match self[s].clone() {
                                HeapCellValue::NamedStr(arity, name, fixity) => {
                                    let threshold = self.threshold();

                                    self[scan] = HeapCellValue::Addr(Addr::Str(threshold));
                                    self[s] = HeapCellValue::Addr(Addr::Str(threshold));

                                    trail.push((Ref::HeapCell(s),
                                                HeapCellValue::NamedStr(arity, name.clone(), fixity)));

                                    self.push(HeapCellValue::NamedStr(arity, name, fixity));

                                    for i in 0 .. arity {
                                        let hcv = self[s + 1 + i].clone();
                                        self.push(hcv);
                                    }
                                },
                                HeapCellValue::Addr(Addr::Str(o)) =>
                                    self[scan] = HeapCellValue::Addr(Addr::Str(o)),
                                _ => {}
                            };

                            scan += 1;
                        },
                        Addr::Con(_) => scan += 1
                    }
            }
        }

        RedirectInfo { trail }
    }

    fn duplicate_term(&mut self, addr: Addr)
    {
        let redirect = self.duplicate_term_impl(addr);
        self.unwind_trail(redirect);
    }
}
<commit_msg>add support for copying attributed variables in copy_term<commit_after>use prolog::and_stack::*;
use prolog::instructions::*;

use std::ops::IndexMut;

type Trail = Vec<(Ref, HeapCellValue)>;

pub(crate) struct RedirectInfo {
    trail: Trail
}

pub(crate) trait CopierTarget: IndexMut<usize, Output=HeapCellValue>
{
    fn source(&self) -> usize;
    fn threshold(&self) -> usize;
    fn push(&mut self, HeapCellValue);
    fn store(&self, Addr) -> Addr;
    fn deref(&self, Addr) -> Addr;
    fn stack(&mut self) -> &mut AndStack;

    fn unwind_trail(&mut self, redirect: RedirectInfo)
    {
        for (r, hcv) in redirect.trail {
            match r {
                Ref::AttrVar(hc) | Ref::HeapCell(hc) => self[hc] = hcv.clone(),
                Ref::StackCell(fr, sc) => self.stack()[fr][sc] = hcv.as_addr(0)
            }
        }
    }

    fn reinstantiate_var(&mut self, ra: Addr, scan: usize, trail: &mut Trail)
    {
        match ra {
            Addr::HeapCell(hc) => {
                self[scan] = HeapCellValue::Addr(Addr::HeapCell(scan));
                self[hc] = HeapCellValue::Addr(Addr::HeapCell(scan));
                trail.push((Ref::HeapCell(hc),
                            HeapCellValue::Addr(Addr::HeapCell(hc))));
            },
            Addr::StackCell(fr, sc) => {
                self[scan] = HeapCellValue::Addr(Addr::HeapCell(scan));
                self.stack()[fr][sc] = Addr::HeapCell(scan);
                trail.push((Ref::StackCell(fr, sc),
                            HeapCellValue::Addr(Addr::StackCell(fr, sc))));
            },
            Addr::AttrVar(hc) => {
                self[scan] = HeapCellValue::Addr(Addr::AttrVar(scan));
                self[hc] = HeapCellValue::Addr(Addr::AttrVar(scan));
                trail.push((Ref::AttrVar(hc),
                            HeapCellValue::Addr(Addr::AttrVar(hc))));
            },
            _ => {}
        }
    }

    // duplicate_term_impl(L1, L2) uses Cheney's algorithm to copy the term
    // at L1 to L2. trail is kept to restore the innards of L1 after
    // it's been copied to L2.
    fn duplicate_term_impl(&mut self, addr: Addr) -> RedirectInfo
    {
        let mut trail = Trail::new();
        let mut scan = self.source();
        let old_h = self.threshold();

        self.push(HeapCellValue::Addr(addr));

        while scan < self.threshold() {
            match self[scan].clone() {
                HeapCellValue::NamedStr(..) =>
                    scan += 1,
                HeapCellValue::Addr(a) =>
                    match a.clone() {
                        Addr::Lis(a) => {
                            if let HeapCellValue::Addr(Addr::Lis(b)) = self[a].clone() {
                                if b >= old_h {
                                    self[scan] = HeapCellValue::Addr(Addr::Lis(b));
                                    scan += 1;
                                    continue;
                                }
                            }

                            let threshold = self.threshold();
                            self[scan] = HeapCellValue::Addr(Addr::Lis(threshold));

                            let hcv = self[a].clone();
                            self.push(hcv.clone());

                            let ra = hcv.as_addr(threshold);
                            let rd = self.store(self.deref(ra));

                            match rd.clone() {
                                Addr::HeapCell(hc) if hc >= old_h =>
                                    self[threshold] = HeapCellValue::Addr(rd),
                                ra @ Addr::AttrVar(_) | ra @ Addr::HeapCell(..) | ra @ Addr::StackCell(..) =>
                                    if ra == rd {
                                        self.reinstantiate_var(ra, threshold, &mut trail);
                                    } else {
                                        self[threshold] = HeapCellValue::Addr(ra);
                                    },
                                _ => {
                                    trail.push((Ref::HeapCell(a), self[a].clone()));
                                    self[a] = HeapCellValue::Addr(Addr::Lis(threshold))
                                }
                            };

                            let hcv = self[a+1].clone();
                            self.push(hcv);

                            scan += 1;
                        },
                        Addr::AttrVar(_) | Addr::HeapCell(_) | Addr::StackCell(_, _) => {
                            let ra = a;
                            let rd = self.store(self.deref(ra.clone()));

                            match rd.clone() {
                                Addr::AttrVar(h) | Addr::HeapCell(h) if h >= old_h => {
                                    self[scan] = HeapCellValue::Addr(rd);
                                    scan += 1;
                                },
                                Addr::AttrVar(h) if ra == rd => {
                                    let threshold = self.threshold();
                                    self.push(HeapCellValue::Addr(Addr::AttrVar(threshold)));

                                    let list_val = self[h+1].clone();
                                    self.push(list_val);

                                    self.reinstantiate_var(ra, threshold, &mut trail);
                                    self[scan] = HeapCellValue::Addr(Addr::AttrVar(threshold));
                                },
                                _ if ra == rd => {
                                    self.reinstantiate_var(ra, scan, &mut trail);
                                    scan += 1;
                                },
                                _ => self[scan] = HeapCellValue::Addr(rd)
                            };
                        },
                        Addr::Str(s) => {
                            match self[s].clone() {
                                HeapCellValue::NamedStr(arity, name, fixity) => {
                                    let threshold = self.threshold();

                                    self[scan] = HeapCellValue::Addr(Addr::Str(threshold));
                                    self[s] = HeapCellValue::Addr(Addr::Str(threshold));

                                    trail.push((Ref::HeapCell(s),
                                                HeapCellValue::NamedStr(arity, name.clone(), fixity)));

                                    self.push(HeapCellValue::NamedStr(arity, name, fixity));

                                    for i in 0 .. arity {
                                        let hcv = self[s + 1 + i].clone();
                                        self.push(hcv);
                                    }
                                },
                                HeapCellValue::Addr(Addr::Str(o)) =>
                                    self[scan] = HeapCellValue::Addr(Addr::Str(o)),
                                _ => {}
                            };

                            scan += 1;
                        },
                        Addr::Con(_) => scan += 1
                    }
            }
        }

        RedirectInfo { trail }
    }

    fn duplicate_term(&mut self, addr: Addr)
    {
        let redirect = self.duplicate_term_impl(addr);
        self.unwind_trail(redirect);
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]
use raft::errors::{Result, Error};
use raft::Storage;
use protobuf::{self, RepeatedField};
use proto::raftpb::{HardState, Entry, EntryType, Message, Snapshot, MessageType, ConfChange,
                    ConfChangeType, ConfState};
use raft::raft::{Config, Raft, SoftState, INVALID_ID};
use raft::Status;

#[derive(Debug, Default)]
struct Peer {
    id: u64,
    context: Vec<u8>,
}

#[derive(Debug, PartialEq, Copy, Clone)]
enum SnapshotStatus {
    SnapshotFinish,
    SnapshotFailure,
}

fn is_local_msg(m: &Message) -> bool {
    match m.get_msg_type() {
        MessageType::MsgHup |
        MessageType::MsgBeat |
        MessageType::MsgUnreachable |
        MessageType::MsgSnapStatus |
        MessageType::MsgCheckQuorum => true,
        _ => false,
    }
}

fn is_response_msg(m: &Message) -> bool {
    match m.get_msg_type() {
        MessageType::MsgAppendResponse |
        MessageType::MsgRequestVoteResponse |
        MessageType::MsgHeartbeatResponse |
        MessageType::MsgUnreachable => true,
        _ => false,
    }
}

fn is_empty_snap(s: &Snapshot) -> bool {
    s.get_metadata().get_index() == 0
}

// Ready encapsulates the entries and messages that are ready to read,
// be saved to stable storage, committed or sent to other peers.
// All fields in Ready are read-only.
#[derive(Default)]
pub struct Ready {
    // The current volatile state of a Node.
    // SoftState will be nil if there is no update.
    // It is not required to consume or store SoftState.
    ss: Option<SoftState>,

    // The current state of a Node to be saved to stable storage BEFORE
    // Messages are sent.
    // HardState will be equal to empty state if there is no update.
    hs: Option<HardState>,

    // Entries specifies entries to be saved to stable storage BEFORE
    // Messages are sent.
    entries: Vec<Entry>,

    // Snapshot specifies the snapshot to be saved to stable storage.
    snapshot: Snapshot,

    // CommittedEntries specifies entries to be committed to a
    // store/state-machine. These have previously been committed to stable
    // store.
    committed_entries: Vec<Entry>,

    // Messages specifies outbound messages to be sent AFTER Entries are
    // committed to stable storage.
    // If it contains a MsgSnap message, the application MUST report back to raft
    // when the snapshot has been received or has failed by calling ReportSnapshot.
    messages: Vec<Message>,
}

impl Ready {
    fn new<T: Storage + Default>(raft: &mut Raft<T>,
                                 prev_ss: &SoftState,
                                 prev_hs: &HardState)
                                 -> Ready {
        let mut rd = Ready {
            entries: raft.raft_log.unstable_entries().unwrap_or(&vec![]).to_vec(),
            committed_entries: raft.raft_log.next_entries().unwrap_or(vec![]),
            messages: raft.msgs.drain(..).collect(),
            ..Default::default()
        };
        let ss = raft.soft_state();
        if &ss != prev_ss {
            rd.ss = Some(ss);
        }
        let hs = raft.hard_state();
        if &hs != prev_hs {
            rd.hs = Some(hs);
        }
        if raft.raft_log.get_unstable().snapshot.is_some() {
            rd.snapshot = raft.raft_log.get_unstable().snapshot.clone().unwrap();
        }
        rd
    }
}

// RawNode is a thread-unsafe Node.
// The methods of this struct correspond to the methods of Node and are described
// more fully there.
#[derive(Default)]
pub struct RawNode<T: Storage + Default> {
    raft: Raft<T>,
    prev_ss: SoftState,
    prev_hs: HardState,
}

impl<T: Storage + Default> RawNode<T> {
    // NewRawNode returns a new RawNode given configuration and a list of raft peers.
    fn new(config: &Config<T>, peers: &[Peer]) -> Result<RawNode<T>> {
        assert!(config.id != 0, "config.id must not be zero");
        let r = Raft::new(config);
        let mut rn = RawNode { raft: r, ..Default::default() };
        let last_index = rn.raft.get_store().last_index().expect("");
        if last_index == 0 {
            rn.raft.become_follower(1, INVALID_ID);
            let mut ents = Vec::with_capacity(peers.len());
            for (i, peer) in peers.iter().enumerate() {
                let mut cc = ConfChange::new();
                cc.set_change_type(ConfChangeType::ConfChangeAddNode);
                cc.set_node_id(peer.id);
                cc.set_context(peer.context.clone());
                let data = protobuf::Message::write_to_bytes(&cc)
                               .expect("unexpected marshal error");
                let mut e = Entry::new();
                e.set_entry_type(EntryType::EntryConfChange);
                e.set_term(1);
                e.set_index(i as u64 + 1);
                e.set_data(data);
                ents.push(e);
            }
            rn.raft.raft_log.append(&ents);
            rn.raft.raft_log.committed = ents.len() as u64;
            for peer in peers {
                rn.raft.add_node(peer.id);
            }
        }
        rn.prev_ss = rn.raft.soft_state();
        rn.prev_hs = rn.raft.hard_state();
        Ok(rn)
    }

    fn commit_ready(&mut self, rd: Ready) {
        if rd.ss.is_some() {
            self.prev_ss = rd.ss.unwrap();
        }
        if let Some(e) = rd.hs {
            if e != HardState::new() {
                self.prev_hs = e;
            }
        }
        if self.prev_hs.get_commit() != 0 {
            // In most cases, prevHardSt and rd.HardState will be the same
            // because when there are new entries to apply we just sent a
            // HardState with an updated Commit value. However, on initial
            // startup the two are different because we don't send a HardState
            // until something changes, but we do send any un-applied but
            // committed entries (and previously-committed entries may be
            // incorporated into the snapshot, even if rd.CommittedEntries is
            // empty). Therefore we mark all committed entries as applied
            // whether they were included in rd.HardState or not.
            self.raft.raft_log.applied_to(self.prev_hs.get_commit());
        }
        if rd.entries.len() > 0 {
            let e = rd.entries.last().unwrap();
            self.raft.raft_log.stable_to(e.get_index(), e.get_term());
        }
        if rd.snapshot != Snapshot::new() {
            self.raft.raft_log.stable_snap_to(rd.snapshot.get_metadata().get_index());
        }
    }

    // Tick advances the internal logical clock by a single tick.
    fn tick(&mut self) {
        self.raft.tick();
    }

    // Campaign causes this RawNode to transition to candidate state.
    fn campaign(&mut self) -> Result<()> {
        let mut m = Message::new();
        m.set_msg_type(MessageType::MsgHup);
        self.raft.step(m)
    }

    // Propose proposes data be appended to the raft log.
    fn propose(&mut self, data: Vec<u8>) -> Result<()> {
        let mut m = Message::new();
        m.set_msg_type(MessageType::MsgPropose);
        m.set_from(self.raft.id);
        let mut e = Entry::new();
        e.set_data(data);
        m.set_entries(RepeatedField::from_vec(vec![e]));
        self.raft.step(m)
    }

    // ProposeConfChange proposes a config change.
    fn propose_conf_change(&mut self, cc: ConfChange) -> Result<()> {
        let data = try!(protobuf::Message::write_to_bytes(&cc));
        let mut m = Message::new();
        m.set_msg_type(MessageType::MsgPropose);
        let mut e = Entry::new();
        e.set_entry_type(EntryType::EntryConfChange);
        e.set_data(data);
        m.set_entries(RepeatedField::from_vec(vec![e]));
        self.raft.step(m)
    }

    fn apply_conf_change(&mut self, cc: ConfChange) -> ConfState {
        if cc.get_node_id() == INVALID_ID {
            self.raft.reset_pending_conf();
            let mut cs = ConfState::new();
            cs.set_nodes(self.raft.nodes());
            return cs;
        }
        let nid = cc.get_node_id();
        assert!(cc.has_change_type(), "unexpected conf type");
        match cc.get_change_type() {
            ConfChangeType::ConfChangeAddNode => self.raft.add_node(nid),
            ConfChangeType::ConfChangeRemoveNode => self.raft.remove_node(nid),
            ConfChangeType::ConfChangeUpdateNode => self.raft.reset_pending_conf(),
        }
        let mut cs = ConfState::new();
        cs.set_nodes(self.raft.nodes());
        cs
    }

    // Step advances the state machine using the given message.
    fn step(&mut self, m: Message) -> Result<()> {
        // ignore unexpected local messages receiving over network
        if is_local_msg(&m) {
            return Err(Error::StepLocalMsg);
        }
        if self.raft.prs.contains_key(&m.get_from()) || is_response_msg(&m) {
            return self.raft.step(m);
        }
        Err(Error::StepPeerNotFound)
    }

    // Ready returns the current point-in-time state of this RawNode.
    fn ready(&mut self) -> Ready {
        Ready::new(&mut self.raft, &self.prev_ss, &self.prev_hs)
    }

    // HasReady called when RawNode user need to check if any Ready pending.
    // Checking logic in this method should be consistent with Ready.containsUpdates().
    fn has_ready(&self) -> bool {
        let raft = &self.raft;
        if raft.soft_state() != self.prev_ss {
            return true;
        }
        if raft.hs != HardState::new() && raft.hs != self.prev_hs {
            return true;
        }
        if Some(true) == raft.raft_log.get_unstable().snapshot.as_ref().map(|s| !is_empty_snap(s)) {
            return true;
        }
        if raft.msgs.len() > 0 || raft.raft_log.unstable_entries().is_some() ||
           raft.raft_log.has_next_entries() {
            return true;
        }
        false
    }

    // Advance notifies the RawNode that the application has applied and saved progress in the
    // last Ready results.
    fn advance(&mut self, rd: Ready) {
        self.commit_ready(rd);
    }

    // Status returns the current status of the given group.
    fn status(&self) -> Status {
        Status::new(&self.raft)
    }

    // ReportUnreachable reports the given node is not reachable for the last send.
    fn report_unreachable(&mut self, id: u64) {
        let mut m = Message::new();
        m.set_msg_type(MessageType::MsgUnreachable);
        m.set_from(id);
        // we don't care if it is ok actually
        self.raft.step(m).is_ok();
    }

    // ReportSnapshot reports the status of the sent snapshot.
    fn report_snapshot(&mut self, id: u64, status: SnapshotStatus) {
        let rej = status == SnapshotStatus::SnapshotFailure;
        let mut m = Message::new();
        m.set_msg_type(MessageType::MsgSnapStatus);
        m.set_from(id);
        m.set_reject(rej);
        // we don't care if it is ok actually
        self.raft.step(m).is_ok();
    }
}
<commit_msg>raft: public for outer use.<commit_after>#![allow(dead_code)]
use raft::errors::{Result, Error};
use raft::Storage;
use protobuf::{self, RepeatedField};
use proto::raftpb::{HardState, Entry, EntryType, Message, Snapshot, MessageType, ConfChange,
                    ConfChangeType, ConfState};
use raft::raft::{Config, Raft, SoftState, INVALID_ID};
use raft::Status;

#[derive(Debug, Default)]
pub struct Peer {
    pub id: u64,
    pub context: Vec<u8>,
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum SnapshotStatus {
    SnapshotFinish,
    SnapshotFailure,
}

fn is_local_msg(m: &Message) -> bool {
    match m.get_msg_type() {
        MessageType::MsgHup |
        MessageType::MsgBeat |
        MessageType::MsgUnreachable |
        MessageType::MsgSnapStatus |
        MessageType::MsgCheckQuorum => true,
        _ => false,
    }
}

fn is_response_msg(m: &Message) -> bool {
    match m.get_msg_type() {
        MessageType::MsgAppendResponse |
        MessageType::MsgRequestVoteResponse |
        MessageType::MsgHeartbeatResponse |
        MessageType::MsgUnreachable => true,
        _ => false,
    }
}

fn is_empty_snap(s: &Snapshot) -> bool {
    s.get_metadata().get_index() == 0
}

// Ready encapsulates the entries and messages that are ready to read,
// be saved to stable storage, committed or sent to other peers.
// All fields in Ready are read-only.
#[derive(Default)]
pub struct Ready {
    // The current volatile state of a Node.
    // SoftState will be nil if there is no update.
    // It is not required to consume or store SoftState.
    pub ss: Option<SoftState>,

    // The current state of a Node to be saved to stable storage BEFORE
    // Messages are sent.
    // HardState will be equal to empty state if there is no update.
    pub hs: Option<HardState>,

    // Entries specifies entries to be saved to stable storage BEFORE
    // Messages are sent.
    pub entries: Vec<Entry>,

    // Snapshot specifies the snapshot to be saved to stable storage.
    pub snapshot: Snapshot,

    // CommittedEntries specifies entries to be committed to a
    // store/state-machine. These have previously been committed to stable
    // store.
    pub committed_entries: Vec<Entry>,

    // Messages specifies outbound messages to be sent AFTER Entries are
    // committed to stable storage.
    // If it contains a MsgSnap message, the application MUST report back to raft
    // when the snapshot has been received or has failed by calling ReportSnapshot.
    pub messages: Vec<Message>,
}

impl Ready {
    fn new<T: Storage + Default>(raft: &mut Raft<T>,
                                 prev_ss: &SoftState,
                                 prev_hs: &HardState)
                                 -> Ready {
        let mut rd = Ready {
            entries: raft.raft_log.unstable_entries().unwrap_or(&vec![]).to_vec(),
            committed_entries: raft.raft_log.next_entries().unwrap_or(vec![]),
            messages: raft.msgs.drain(..).collect(),
            ..Default::default()
        };
        let ss = raft.soft_state();
        if &ss != prev_ss {
            rd.ss = Some(ss);
        }
        let hs = raft.hard_state();
        if &hs != prev_hs {
            rd.hs = Some(hs);
        }
        if raft.raft_log.get_unstable().snapshot.is_some() {
            rd.snapshot = raft.raft_log.get_unstable().snapshot.clone().unwrap();
        }
        rd
    }
}

// RawNode is a thread-unsafe Node.
// The methods of this struct correspond to the methods of Node and are described
// more fully there.
#[derive(Default)]
pub struct RawNode<T: Storage + Default> {
    raft: Raft<T>,
    prev_ss: SoftState,
    prev_hs: HardState,
}

impl<T: Storage + Default> RawNode<T> {
    // NewRawNode returns a new RawNode given configuration and a list of raft peers.
    pub fn new(config: &Config<T>, peers: &[Peer]) -> Result<RawNode<T>> {
        assert!(config.id != 0, "config.id must not be zero");
        let r = Raft::new(config);
        let mut rn = RawNode { raft: r, ..Default::default() };
        let last_index = rn.raft.get_store().last_index().expect("");
        if last_index == 0 {
            rn.raft.become_follower(1, INVALID_ID);
            let mut ents = Vec::with_capacity(peers.len());
            for (i, peer) in peers.iter().enumerate() {
                let mut cc = ConfChange::new();
                cc.set_change_type(ConfChangeType::ConfChangeAddNode);
                cc.set_node_id(peer.id);
                cc.set_context(peer.context.clone());
                let data = protobuf::Message::write_to_bytes(&cc)
                               .expect("unexpected marshal error");
                let mut e = Entry::new();
                e.set_entry_type(EntryType::EntryConfChange);
                e.set_term(1);
                e.set_index(i as u64 + 1);
                e.set_data(data);
                ents.push(e);
            }
            rn.raft.raft_log.append(&ents);
            rn.raft.raft_log.committed = ents.len() as u64;
            for peer in peers {
                rn.raft.add_node(peer.id);
            }
        }
        rn.prev_ss = rn.raft.soft_state();
        rn.prev_hs = rn.raft.hard_state();
        Ok(rn)
    }

    fn commit_ready(&mut self, rd: Ready) {
        if rd.ss.is_some() {
            self.prev_ss = rd.ss.unwrap();
        }
        if let Some(e) = rd.hs {
            if e != HardState::new() {
                self.prev_hs = e;
            }
        }
        if self.prev_hs.get_commit() != 0 {
            // In most cases, prevHardSt and rd.HardState will be the same
            // because when there are new entries to apply we just sent a
            // HardState with an updated Commit value. However, on initial
            // startup the two are different because we don't send a HardState
            // until something changes, but we do send any un-applied but
            // committed entries (and previously-committed entries may be
            // incorporated into the snapshot, even if rd.CommittedEntries is
            // empty). Therefore we mark all committed entries as applied
            // whether they were included in rd.HardState or not.
            self.raft.raft_log.applied_to(self.prev_hs.get_commit());
        }
        if rd.entries.len() > 0 {
            let e = rd.entries.last().unwrap();
            self.raft.raft_log.stable_to(e.get_index(), e.get_term());
        }
        if rd.snapshot != Snapshot::new() {
            self.raft.raft_log.stable_snap_to(rd.snapshot.get_metadata().get_index());
        }
    }

    // Tick advances the internal logical clock by a single tick.
    pub fn tick(&mut self) {
        self.raft.tick();
    }

    // Campaign causes this RawNode to transition to candidate state.
    pub fn campaign(&mut self) -> Result<()> {
        let mut m = Message::new();
        m.set_msg_type(MessageType::MsgHup);
        self.raft.step(m)
    }

    // Propose proposes data be appended to the raft log.
    pub fn propose(&mut self, data: Vec<u8>) -> Result<()> {
        let mut m = Message::new();
        m.set_msg_type(MessageType::MsgPropose);
        m.set_from(self.raft.id);
        let mut e = Entry::new();
        e.set_data(data);
        m.set_entries(RepeatedField::from_vec(vec![e]));
        self.raft.step(m)
    }

    // ProposeConfChange proposes a config change.
    pub fn propose_conf_change(&mut self, cc: ConfChange) -> Result<()> {
        let data = try!(protobuf::Message::write_to_bytes(&cc));
        let mut m = Message::new();
        m.set_msg_type(MessageType::MsgPropose);
        let mut e = Entry::new();
        e.set_entry_type(EntryType::EntryConfChange);
        e.set_data(data);
        m.set_entries(RepeatedField::from_vec(vec![e]));
        self.raft.step(m)
    }

    pub fn apply_conf_change(&mut self, cc: ConfChange) -> ConfState {
        if cc.get_node_id() == INVALID_ID {
            self.raft.reset_pending_conf();
            let mut cs = ConfState::new();
            cs.set_nodes(self.raft.nodes());
            return cs;
        }
        let nid = cc.get_node_id();
        assert!(cc.has_change_type(), "unexpected conf type");
        match cc.get_change_type() {
            ConfChangeType::ConfChangeAddNode => self.raft.add_node(nid),
            ConfChangeType::ConfChangeRemoveNode => self.raft.remove_node(nid),
            ConfChangeType::ConfChangeUpdateNode => self.raft.reset_pending_conf(),
        }
        let mut cs = ConfState::new();
        cs.set_nodes(self.raft.nodes());
        cs
    }

    // Step advances the state machine using the given message.
    pub fn step(&mut self, m: Message) -> Result<()> {
        // ignore unexpected local messages receiving over network
        if is_local_msg(&m) {
            return Err(Error::StepLocalMsg);
        }
        if self.raft.prs.contains_key(&m.get_from()) || is_response_msg(&m) {
            return self.raft.step(m);
        }
        Err(Error::StepPeerNotFound)
    }

    // Ready returns the current point-in-time state of this RawNode.
    pub fn ready(&mut self) -> Ready {
        Ready::new(&mut self.raft, &self.prev_ss, &self.prev_hs)
    }

    // HasReady called when RawNode user need to check if any Ready pending.
    // Checking logic in this method should be consistent with Ready.containsUpdates().
    pub fn has_ready(&self) -> bool {
        let raft = &self.raft;
        if raft.soft_state() != self.prev_ss {
            return true;
        }
        if raft.hs != HardState::new() && raft.hs != self.prev_hs {
            return true;
        }
        if Some(true) == raft.raft_log.get_unstable().snapshot.as_ref().map(|s| !is_empty_snap(s)) {
            return true;
        }
        if raft.msgs.len() > 0 || raft.raft_log.unstable_entries().is_some() ||
           raft.raft_log.has_next_entries() {
            return true;
        }
        false
    }

    // Advance notifies the RawNode that the application has applied and saved progress in the
    // last Ready results.
    pub fn advance(&mut self, rd: Ready) {
        self.commit_ready(rd);
    }

    // Status returns the current status of the given group.
    pub fn status(&self) -> Status {
        Status::new(&self.raft)
    }

    // ReportUnreachable reports the given node is not reachable for the last send.
    pub fn report_unreachable(&mut self, id: u64) {
        let mut m = Message::new();
        m.set_msg_type(MessageType::MsgUnreachable);
        m.set_from(id);
        // we don't care if it is ok actually
        self.raft.step(m).is_ok();
    }

    // ReportSnapshot reports the status of the sent snapshot.
    pub fn report_snapshot(&mut self, id: u64, status: SnapshotStatus) {
        let rej = status == SnapshotStatus::SnapshotFailure;
        let mut m = Message::new();
        m.set_msg_type(MessageType::MsgSnapStatus);
        m.set_from(id);
        m.set_reject(rej);
        // we don't care if it is ok actually
        self.raft.step(m).is_ok();
    }
}
<|endoftext|>


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

use tokio::io;
use tokio::io::AsyncWriteExt;
use tokio::net::{TcpListener, TcpStream};
use tracing::{error, info, warn};

use crate::config::Config;
use crate::socket;

use super::Error;

pub struct InboundPassthrough {
    cfg: Config,
}

impl InboundPassthrough {
    pub(crate) fn new(cfg: Config) -> InboundPassthrough {
        InboundPassthrough { cfg }
    }
    pub(super) async fn run(self) {
        let tcp_listener: TcpListener = TcpListener::bind(self.cfg.inbound_plaintext_addr)
            .await
            .expect("failed to bind");
        match socket::set_transparent(&tcp_listener) {
            Err(_e) => info!("running without transparent mode"),
            _ => info!("running with transparent mode"),
        };

        info!(
            "inbound plaintext listener established {}",
            tcp_listener.local_addr().unwrap()
        );

        loop {
            // Asynchronously wait for an inbound socket.
            let socket = tcp_listener.accept().await;
            match socket {
                Ok((mut stream, remote)) => {
                    info!("accepted inbound plaintext connection from {}", remote);
                    tokio::spawn(async move {
                        if let Err(e) = Self::proxy_inbound_plaintext(&mut stream).await {
                            warn!("plaintext proxying failed {}", e)
                        }
                    });
                }
                Err(e) => error!("Failed TCP handshake {}", e),
            }
        }
    }

    async fn proxy_inbound_plaintext(inbound: &mut TcpStream) -> Result<(), Error> {
        let orig = socket::orig_dst_addr_or_default(inbound);
        let mut outbound = TcpStream::connect(orig).await?;

        let (mut ri, mut wi) = inbound.split();
        let (mut ro, mut wo) = outbound.split();

        let client_to_server = async {
            io::copy(&mut ri, &mut wo).await?;
            wo.shutdown().await
        };

        let server_to_client = async {
            io::copy(&mut ro, &mut wi).await?;
            wi.shutdown().await
        };

        tokio::try_join!(client_to_server, server_to_client)?;

        info!("proxy inbound plaintext complete");
        Ok(())
    }
}
<commit_msg>Use zero-copy in passthrough mode. (#179)<commit_after>// Copyright Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use tokio::net::{TcpListener, TcpStream};
use tracing::{error, info, warn};

use crate::config::Config;
use crate::socket;
use crate::socket::relay;

use super::Error;

pub struct InboundPassthrough {
    cfg: Config,
}

impl InboundPassthrough {
    pub(crate) fn new(cfg: Config) -> InboundPassthrough {
        InboundPassthrough { cfg }
    }
    pub(super) async fn run(self) {
        let tcp_listener: TcpListener = TcpListener::bind(self.cfg.inbound_plaintext_addr)
            .await
            .expect("failed to bind");
        match socket::set_transparent(&tcp_listener) {
            Err(_e) => info!("running without transparent mode"),
            _ => info!("running with transparent mode"),
        };

        info!(
            "inbound plaintext listener established {}",
            tcp_listener.local_addr().unwrap()
        );

        loop {
            // Asynchronously wait for an inbound socket.
            let socket = tcp_listener.accept().await;
            match socket {
                Ok((mut stream, remote)) => {
                    info!("accepted inbound plaintext connection from {}", remote);
                    tokio::spawn(async move {
                        if let Err(e) = Self::proxy_inbound_plaintext(&mut stream).await {
                            warn!("plaintext proxying failed {}", e)
                        }
                    });
                }
                Err(e) => error!("Failed TCP handshake {}", e),
            }
        }
    }

    async fn proxy_inbound_plaintext(inbound: &mut TcpStream) -> Result<(), Error> {
        let orig = socket::orig_dst_addr_or_default(inbound);
        let mut outbound = TcpStream::connect(orig).await?;
        relay(inbound, &mut outbound).await?;
        info!("proxy inbound plaintext complete");
        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::f32::consts::PI;
use std::sync::{Arc, Mutex, mpsc::{Sender, Receiver, TryRecvError}};

use cgmath::{Point3, Vector3, Vector4, prelude::*};

use glium::Rect;

use rand::{self, prelude::*};

use crate::camera::Camera;
use crate::material::Material;
use crate::pt_renderer::{Intersect, Ray, RenderCoordinator};
use crate::scene::Scene;
use crate::triangle::{Hit, RTTriangle};

const EPSILON: f32 = 1e-5;

pub struct RenderWorker {
    scene: Arc<Scene>,
    camera: Camera,
    coordinator: Arc<Mutex<RenderCoordinator>>,
    message_rx: Receiver<()>,
    result_tx: Sender<(Rect, Vec<f32>)>,
}

impl RenderWorker {
    pub fn new(scene: Arc<Scene>, camera: Camera, coordinator: Arc<Mutex<RenderCoordinator>>,
               message_rx: Receiver<()>, result_tx: Sender<(Rect, Vec<f32>)>) -> RenderWorker {
        RenderWorker {
            scene, camera, coordinator,
            message_rx, result_tx,
        }
    }

    pub fn run(&self) {
        let (width, height) = {
            let coordinator = self.coordinator.lock().unwrap();
            (coordinator.width, coordinator.height)
        };
        let clip_to_world = self.camera.get_world_to_clip().invert().unwrap();
        loop {
            match self.message_rx.try_recv() {
                Err(TryRecvError::Empty) => (),
                Ok(_) => return,
                Err(TryRecvError::Disconnected) => {
                    println!("Threads were not properly stopped before disconnecting channel!");
                    return;
                }
            }
            if let Some(rect) = self.coordinator.lock().unwrap().next_block() {
                let mut block = vec![0.0f32; (3 * rect.width * rect.height) as usize];
                for h in 0..rect.height {
                    for w in 0..rect.width {
                        // TODO: implement multiple samples per pixel
                        let dx: f32 = rand::random();
                        let dy: f32 = rand::random();
                        let clip_x = 2.0 * ((rect.left + w) as f32 + dx) / width as f32 - 1.0;
                        let clip_y = 2.0 * ((rect.bottom + h) as f32 + dy) / height as f32 - 1.0;
                        let clip_p = Vector4::new(clip_x, clip_y, 1.0, 1.0);
                        let world_p = clip_to_world * clip_p;
                        let dir = ((world_p / world_p.w).truncate() - self.camera.pos.to_vec())
                            .normalize();
                        let ray = Ray::new(self.camera.pos, dir, 100.0);
                        let c = self.trace_ray(&ray, 0);
                        let pixel_i = 3 * (h * rect.width + w) as usize;
                        block[pixel_i]     = c.x;
                        block[pixel_i + 1] = c.y;
                        block[pixel_i + 2] = c.z;
                    }
                }
                self.result_tx.send((rect, block)).expect("Receiver closed!");
            } else {
                return;
            }
        }
    }

    fn trace_ray(&self, ray: &Ray, bounce: u32) -> Vector3<f32> {
        let mut c = Vector3::zero();
        if let Some(hit) = self.find_hit(&ray) {
            let material = &self.scene.materials[hit.tri.material_i];
            let mut normal = hit.normal();
            // Flip the normal if its pointing to the opposite side from the hit
            if normal.dot(ray.dir) > 0.0 {
                normal *= -1.0;
            }
            if bounce == 0 {
                if let Some(emissive) = material.emissive {
                    c += normal.dot(-ray.dir).max(0.0) * emissive;
                }
            }
            let (light, light_pos, light_normal, light_pdf) = self.sample_light();
            let bump_pos = hit.pos() + EPSILON * normal;
            let hit_to_light = light_pos - bump_pos;
            let light_dir = hit_to_light.normalize();
            let shadow_ray = Ray::new(bump_pos, light_dir, hit_to_light.magnitude() - EPSILON);
            if shadow_ray.length > 0.1 && self.find_hit(&shadow_ray).is_none() {
                let cos_l = light_normal.dot(-light_dir).max(0.0);
                let cos_t = normal.dot(light_dir).max(0.0);
                let light_material = &self.scene.materials[light.material_i];
                let emissive = light_material.emissive.expect("Light wasn't emissive");
                c += emissive.mul_element_wise(self.brdf(&ray, &shadow_ray, material))
                    * cos_l * cos_t / (hit_to_light.magnitude2() * light_pdf);
            }
            if bounce < 10 {
                let (new_dir, pdf) = self.sample_dir(normal);
                let new_ray = Ray::new(bump_pos, new_dir, 100.0);
                c += normal.dot(new_dir) * self.brdf(&ray, &new_ray, material)
                    .mul_element_wise(self.trace_ray(&new_ray, bounce + 1)) / pdf;
            }
        }
        c
    }

    fn brdf(&self, _in_ray: &Ray, _out_ray: &Ray, material: &Material) -> Vector3<f32> {
        material.diffuse / PI
    }

    fn sample_dir(&self, normal: Vector3<f32>) -> (Vector3<f32>, f32) {
        let dir = 2.0 * PI * rand::random::<f32>();
        let length: f32 = rand::random();
        let x = length * dir.cos();
        let y = length * dir.sin();
        let z = (1.0 - length.powi(2)).sqrt();
        let nx = if normal.x.abs() > normal.y.abs() {
            Vector3::new(normal.z, 0.0, -normal.x).normalize()
        } else {
            Vector3::new(0.0, -normal.z, normal.y).normalize()
        };
        let ny = normal.cross(nx);
        (x * nx + y * ny + z * normal, z / PI)
    }

    pub fn sample_light(&self) -> (&RTTriangle, Point3<f32>, Vector3<f32>, f32) {
        if self.scene.lights.is_empty() {
            panic!("Rendered scene has no lights!");
        } else {
            let i = rand::thread_rng().gen_range(0, self.scene.lights.len());
            let light = &self.scene.lights[i];
            let pdf = 1.0 / (self.scene.lights.len() as f32 * light.area());
            let (point, normal) = light.random_point();
            (light, point, normal, pdf)
        }
    }

    fn find_hit(&self, ray: &Ray) -> Option<Hit> {
        let bvh = &self.scene.bvh;
        let mut node_stack = Vec::new();
        node_stack.push((bvh.root(), 0.0f32));
        let mut closest_hit: Option<Hit> = None;
        while let Some((node, t)) = node_stack.pop() {
            // We've already found closer hit
            if closest_hit.as_ref().map_or(false, |hit| hit.t <= t) { continue }
            if node.is_leaf() {
                for tri in &self.scene.triangles[node.start_i..node.end_i] {
                    if let Some(hit) = tri.intersect(&ray) {
                        if let Some(closest) = closest_hit.take() {
                            if hit.t < closest.t {
                                closest_hit = Some(hit);
                            } else {
                                closest_hit = Some(closest);
                            }
                        } else {
                            closest_hit = Some(hit);
                        }
                    }
                }
            } else {
                // TODO: put closest hit on top of the stack
                let (left, right) = bvh.get_children(node).expect("Non leaf node had no child nodes!");
                if let Some(t_left) = left.intersect(&ray) {
                    node_stack.push((left, t_left));
                }
                if let Some(t_right) = right.intersect(&ray) {
                    node_stack.push((right, t_right));
                }
            }
        }
        closest_hit
    }
}
<commit_msg>Don't reallocate node_stack for each ray<commit_after>use std::f32::consts::PI;
use std::sync::{Arc, Mutex, mpsc::{Sender, Receiver, TryRecvError}};

use cgmath::{Point3, Vector3, Vector4, prelude::*};

use glium::Rect;

use rand::{self, prelude::*};

use crate::bvh::BVHNode;
use crate::camera::Camera;
use crate::material::Material;
use crate::pt_renderer::{Intersect, Ray, RenderCoordinator};
use crate::scene::Scene;
use crate::triangle::{Hit, RTTriangle};

const EPSILON: f32 = 1e-5;

pub struct RenderWorker {
    scene: Arc<Scene>,
    camera: Camera,
    coordinator: Arc<Mutex<RenderCoordinator>>,
    message_rx: Receiver<()>,
    result_tx: Sender<(Rect, Vec<f32>)>,
}

impl RenderWorker {
    pub fn new(scene: Arc<Scene>, camera: Camera, coordinator: Arc<Mutex<RenderCoordinator>>,
               message_rx: Receiver<()>, result_tx: Sender<(Rect, Vec<f32>)>) -> RenderWorker {
        RenderWorker {
            scene, camera, coordinator,
            message_rx, result_tx,
        }
    }

    pub fn run(&self) {
        let (width, height) = {
            let coordinator = self.coordinator.lock().unwrap();
            (coordinator.width, coordinator.height)
        };
        let clip_to_world = self.camera.get_world_to_clip().invert().unwrap();
        let mut node_stack: Vec<(&BVHNode, f32)> = Vec::new();
        loop {
            match self.message_rx.try_recv() {
                Err(TryRecvError::Empty) => (),
                Ok(_) => return,
                Err(TryRecvError::Disconnected) => {
                    println!("Threads were not properly stopped before disconnecting channel!");
                    return;
                }
            }
            if let Some(rect) = self.coordinator.lock().unwrap().next_block() {
                let mut block = vec![0.0f32; (3 * rect.width * rect.height) as usize];
                for h in 0..rect.height {
                    for w in 0..rect.width {
                        // TODO: implement multiple samples per pixel
                        let dx: f32 = rand::random();
                        let dy: f32 = rand::random();
                        let clip_x = 2.0 * ((rect.left + w) as f32 + dx) / width as f32 - 1.0;
                        let clip_y = 2.0 * ((rect.bottom + h) as f32 + dy) / height as f32 - 1.0;
                        let clip_p = Vector4::new(clip_x, clip_y, 1.0, 1.0);
                        let world_p = clip_to_world * clip_p;
                        let dir = ((world_p / world_p.w).truncate() - self.camera.pos.to_vec())
                            .normalize();
                        let ray = Ray::new(self.camera.pos, dir, 100.0);
                        let c = self.trace_ray(&ray, &mut node_stack, 0);
                        let pixel_i = 3 * (h * rect.width + w) as usize;
                        block[pixel_i]     = c.x;
                        block[pixel_i + 1] = c.y;
                        block[pixel_i + 2] = c.z;
                    }
                }
                self.result_tx.send((rect, block)).expect("Receiver closed!");
            } else {
                return;
            }
        }
    }

    fn trace_ray(&'a self, ray: &Ray, node_stack: &mut Vec<(&'a BVHNode, f32)>, bounce: u32) -> Vector3<f32> {
        let mut c = Vector3::zero();
        if let Some(hit) = self.find_hit(&ray, node_stack) {
            let material = &self.scene.materials[hit.tri.material_i];
            let mut normal = hit.normal();
            // Flip the normal if its pointing to the opposite side from the hit
            if normal.dot(ray.dir) > 0.0 {
                normal *= -1.0;
            }
            if bounce == 0 {
                if let Some(emissive) = material.emissive {
                    c += normal.dot(-ray.dir).max(0.0) * emissive;
                }
            }
            let (light, light_pos, light_normal, light_pdf) = self.sample_light();
            let bump_pos = hit.pos() + EPSILON * normal;
            let hit_to_light = light_pos - bump_pos;
            let light_dir = hit_to_light.normalize();
            let shadow_ray = Ray::new(bump_pos, light_dir, hit_to_light.magnitude() - EPSILON);
            if shadow_ray.length > 0.1 && self.find_hit(&shadow_ray, node_stack).is_none() {
                let cos_l = light_normal.dot(-light_dir).max(0.0);
                let cos_t = normal.dot(light_dir).max(0.0);
                let light_material = &self.scene.materials[light.material_i];
                let emissive = light_material.emissive.expect("Light wasn't emissive");
                c += emissive.mul_element_wise(self.brdf(&ray, &shadow_ray, material))
                    * cos_l * cos_t / (hit_to_light.magnitude2() * light_pdf);
            }
            if bounce < 10 {
                let (new_dir, pdf) = self.sample_dir(normal);
                let new_ray = Ray::new(bump_pos, new_dir, 100.0);
                c += normal.dot(new_dir) * self.brdf(&ray, &new_ray, material)
                    .mul_element_wise(self.trace_ray(&new_ray, node_stack, bounce + 1)) / pdf;
            }
        }
        c
    }

    fn brdf(&self, _in_ray: &Ray, _out_ray: &Ray, material: &Material) -> Vector3<f32> {
        material.diffuse / PI
    }

    fn sample_dir(&self, normal: Vector3<f32>) -> (Vector3<f32>, f32) {
        let dir = 2.0 * PI * rand::random::<f32>();
        let length: f32 = rand::random();
        let x = length * dir.cos();
        let y = length * dir.sin();
        let z = (1.0 - length.powi(2)).sqrt();
        let nx = if normal.x.abs() > normal.y.abs() {
            Vector3::new(normal.z, 0.0, -normal.x).normalize()
        } else {
            Vector3::new(0.0, -normal.z, normal.y).normalize()
        };
        let ny = normal.cross(nx);
        (x * nx + y * ny + z * normal, z / PI)
    }

    pub fn sample_light(&self) -> (&RTTriangle, Point3<f32>, Vector3<f32>, f32) {
        if self.scene.lights.is_empty() {
            panic!("Rendered scene has no lights!");
        } else {
            let i = rand::thread_rng().gen_range(0, self.scene.lights.len());
            let light = &self.scene.lights[i];
            let pdf = 1.0 / (self.scene.lights.len() as f32 * light.area());
            let (point, normal) = light.random_point();
            (light, point, normal, pdf)
        }
    }

    fn find_hit(&'a self, ray: &Ray, node_stack: &mut Vec<(&'a BVHNode, f32)>) -> Option<Hit> {
        let bvh = &self.scene.bvh;
        node_stack.push((bvh.root(), 0.0f32));
        let mut closest_hit: Option<Hit> = None;
        while let Some((node, t)) = node_stack.pop() {
            // We've already found closer hit
            if closest_hit.as_ref().map_or(false, |hit| hit.t <= t) { continue }
            if node.is_leaf() {
                for tri in &self.scene.triangles[node.start_i..node.end_i] {
                    if let Some(hit) = tri.intersect(&ray) {
                        if let Some(closest) = closest_hit.take() {
                            if hit.t < closest.t {
                                closest_hit = Some(hit);
                            } else {
                                closest_hit = Some(closest);
                            }
                        } else {
                            closest_hit = Some(hit);
                        }
                    }
                }
            } else {
                // TODO: put closest hit on top of the stack
                let (left, right) = bvh.get_children(node).expect("Non leaf node had no child nodes!");
                if let Some(t_left) = left.intersect(&ray) {
                    node_stack.push((left, t_left));
                }
                if let Some(t_right) = right.intersect(&ray) {
                    node_stack.push((right, t_right));
                }
            }
        }
        closest_hit
    }
}
<|endoftext|>


<commit_before>use tiled::Map;

use crate::data::{get_map_tile, load_map_file};
use crate::game::constants::{MAP_FILE_PATH, TILE_MAP_BUF_LENGTH, TILES_PCS_H, TILES_PCS_W};
use crate::shaders::TileMapData;

fn calc_index(x_pos: usize, y_pos: usize) -> usize {
  (y_pos * TILES_PCS_W) + x_pos
}

fn populate_tile_map<'a>(tiles: &'a mut Vec<TileMapData>, map: &Map) -> &'a mut Vec<TileMapData> {
  for y_pos in 0..TILES_PCS_H {
    for x_pos in 0..TILES_PCS_W {
      let map_val = get_map_tile(map, 0, x_pos, y_pos) - 1;
      let idx = calc_index(x_pos, y_pos);

      if idx < TILE_MAP_BUF_LENGTH {
        tiles[idx] =
          TileMapData::new([map_val as f32, 0.0, 0.0, 0.0]);
      } else if idx < TILE_MAP_BUF_LENGTH * 2 {
        tiles[idx - TILE_MAP_BUF_LENGTH] =
          TileMapData::new([tiles[idx - TILE_MAP_BUF_LENGTH].data[0], map_val as f32, 0.0, 0.0]);
      } else if idx < TILE_MAP_BUF_LENGTH * 3 {
        tiles[idx - TILE_MAP_BUF_LENGTH * 2] =
          TileMapData::new([tiles[idx - TILE_MAP_BUF_LENGTH * 2].data[0], tiles[idx - TILE_MAP_BUF_LENGTH * 2].data[1], map_val as f32, 0.0]);
      } else {
        tiles[idx - TILE_MAP_BUF_LENGTH * 3] =
          TileMapData::new([tiles[idx - TILE_MAP_BUF_LENGTH * 3].data[0], tiles[idx - TILE_MAP_BUF_LENGTH * 3].data[1], tiles[idx - TILE_MAP_BUF_LENGTH * 3].data[2], map_val as f32]);
      }
    }
  }
  tiles
}

pub struct Terrain {
  pub tiles: Vec<TileMapData>,
  pub tile_sets: [Map; 1],
  pub curr_tile_set_idx: usize,
}

impl Terrain {
  pub fn new() -> Terrain {
    let mut map_data = Vec::with_capacity(TILE_MAP_BUF_LENGTH);

    for _ in 0..TILE_MAP_BUF_LENGTH {
      map_data.push(TileMapData::new_empty());
    }

    let map_a = load_map_file(MAP_FILE_PATH);

    Terrain {
      tiles: populate_tile_map(&mut map_data, &map_a).to_vec(),
      tile_sets: [map_a],
      curr_tile_set_idx: 0,
    }
  }
}
<commit_msg>Fix tile_map initial allocation size<commit_after>use tiled::Map;

use crate::data::{get_map_tile, load_map_file};
use crate::game::constants::{MAP_FILE_PATH, TILE_MAP_BUF_LENGTH, TILES_PCS_H, TILES_PCS_W};
use crate::shaders::TileMapData;

fn calc_index(x_pos: usize, y_pos: usize) -> usize {
  (y_pos * TILES_PCS_W) + x_pos
}

fn populate_tile_map<'a>(tiles: &'a mut Vec<TileMapData>, map: &Map) -> &'a mut Vec<TileMapData> {
  for y_pos in 0..TILES_PCS_H {
    for x_pos in 0..TILES_PCS_W {
      let map_val = get_map_tile(map, 0, x_pos, y_pos) - 1;
      let idx = calc_index(x_pos, y_pos);

      if idx < TILE_MAP_BUF_LENGTH {
        tiles[idx] =
          TileMapData::new([map_val as f32, 0.0, 0.0, 0.0]);
      } else if idx < TILE_MAP_BUF_LENGTH * 2 {
        tiles[idx - TILE_MAP_BUF_LENGTH] =
          TileMapData::new([tiles[idx - TILE_MAP_BUF_LENGTH].data[0], map_val as f32, 0.0, 0.0]);
      } else if idx < TILE_MAP_BUF_LENGTH * 3 {
        tiles[idx - TILE_MAP_BUF_LENGTH * 2] =
          TileMapData::new([tiles[idx - TILE_MAP_BUF_LENGTH * 2].data[0], tiles[idx - TILE_MAP_BUF_LENGTH * 2].data[1], map_val as f32, 0.0]);
      } else {
        tiles[idx - TILE_MAP_BUF_LENGTH * 3] =
          TileMapData::new([tiles[idx - TILE_MAP_BUF_LENGTH * 3].data[0], tiles[idx - TILE_MAP_BUF_LENGTH * 3].data[1], tiles[idx - TILE_MAP_BUF_LENGTH * 3].data[2], map_val as f32]);
      }
    }
  }
  tiles
}

pub struct Terrain {
  pub tiles: Vec<TileMapData>,
  pub tile_sets: [Map; 1],
  pub curr_tile_set_idx: usize,
}

impl Terrain {
  pub fn new() -> Terrain {
    let mut map_data = Vec::with_capacity(TILES_PCS_H * TILES_PCS_W);

    for _ in 0..(TILES_PCS_H * TILES_PCS_W) {
      map_data.push(TileMapData::new_empty());
    }

    let map_a = load_map_file(MAP_FILE_PATH);

    Terrain {
      tiles: populate_tile_map(&mut map_data, &map_a).to_vec(),
      tile_sets: [map_a],
      curr_tile_set_idx: 0,
    }
  }
}
<|endoftext|>


<commit_before>use piston_window;

pub enum InputResult {
    PushScreen(Box<Screen>),
    PopScreen,
}

pub type InputResults = Vec<InputResult>;

pub struct Context {
    pub assets_path: ::std::path::PathBuf,
}

pub trait Screen {
    fn on_input(&mut self,
                input: &piston_window::Input,
                window: &piston_window::PistonWindow,
                context: &mut Context)
                -> InputResults;
    fn on_draw(&mut self, args: &piston_window::RenderArgs, window: &piston_window::PistonWindow);
    fn on_update(&mut self, args: &piston_window::UpdateArgs);
}
<commit_msg>ExitGame in InputResult<commit_after>use piston_window;

pub enum InputResult {
    PushScreen(Box<Screen>),
    PopScreen,
    ExitGame,
}

pub type InputResults = Vec<InputResult>;

pub struct Context {
    pub assets_path: ::std::path::PathBuf,
}

pub trait Screen {
    fn on_input(&mut self,
                input: &piston_window::Input,
                window: &piston_window::PistonWindow,
                context: &mut Context)
                -> InputResults;
    fn on_draw(&mut self, args: &piston_window::RenderArgs, window: &piston_window::PistonWindow);
    fn on_update(&mut self, args: &piston_window::UpdateArgs);
}
<|endoftext|>