<commit_before>
use std::fmt;

use parser::ast;
use parser::ast::{AstKind, AstNode};
use parser::input_stream::{InputStream, StreamPosition};
use parser::precedence::Precedence;
use parser::token_kind::TokenKind;
use parser::tokenizer::{Token, TokenError, TokenLocation, Tokenizer, TokenizerMode,
                        TokenizerPosition};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FullToken {
    kind: TokenKind,
    location: TokenLocation
}
impl FullToken {
    pub fn new(kind: TokenKind, location: TokenLocation) -> FullToken {
        FullToken {
            kind: kind,
            location: location
        }
    }

    pub fn location(&self) -> &TokenLocation {
        &self.location
    }
}
impl Token for FullToken {
    fn make(kind: TokenKind, location: TokenLocation) -> FullToken {
        FullToken::new(kind, location)
    }
    fn kind(&self) -> TokenKind {
        self.kind
    }
    fn start_offset(&self) -> StreamPosition {
        self.location.start_offset()
    }
    fn write_token(&self, w: &mut fmt::Write) -> Result<(), fmt::Error> {
        write!(w, "Token({}@{})", self.kind.name(), self.location.range_string())
    }
}

pub struct FullTokenizerMode {
}
impl TokenizerMode for FullTokenizerMode {
    type Tok = FullToken;

    fn cares_about_newline(&self) -> bool {
        false
    }
    fn note_newline(&mut self) {}
}

#[derive(Debug, Clone)]
pub enum ParseError {
    Unspecified,
    ErrorToken(FullToken),
    TokenizerError(TokenError),
    UnexpectedToken{expected:TokenKind, got:TokenKind},
    ExpectedVariableName,
    ExpectedCommaOrSemicolon,
    ExpectedExpression
}
pub type ParseResult<T> = Result<T, ParseError>;
pub type MaybeParseResult<T> = ParseResult<Option<T>>;

pub struct AstBuilder<STREAM: InputStream> {
    tokenizer: Tokenizer<STREAM, FullTokenizerMode>
}
impl<STREAM: InputStream> AstBuilder<STREAM> {
    pub fn new(stream: STREAM) -> AstBuilder<STREAM> {
        AstBuilder {
            tokenizer: Tokenizer::new(stream, FullTokenizerMode{})
        }
    }

    pub fn read_and_print_tokens(&mut self) {
        // Just read tokens and print them out until we're done, then return Error.
        loop {
            let token = self.next_token().unwrap();
            println!("Token: {}", token.kind().name());
            if token.kind().is_error() {
                panic!("Got token error: {:?}", self.tokenizer.get_error());
            }
            if token.kind().is_end() {
                break;
            }
        }
    }

    pub fn parse_program(&mut self) -> ParseResult<Box<ast::ProgramNode>> {
        let mut program_node = ast::ProgramNode::new();
        loop {
            match self.try_parse_statement()? {
                Some(boxed_source_element) => {
                    program_node.add_source_element(boxed_source_element);
                }
                None => {
                    break;
                }
            }
        }

        // Must have reached end of stream.
        self.must_expect_token(TokenKind::end())?;

        Ok(Box::new(program_node))
    }

    fn try_parse_statement(&mut self) -> MaybeParseResult<Box<AstNode>> {
        let begin_position = self.mark_position();
        let tok = self.next_token()?;
        if tok.kind().is_open_brace() {
            // FIXME: parse either a block or an object literal.
            // The spec says a '{' at the statement level can only be a block,
            // but practical implementations seem to allow bare object literals.
            return Ok(Some(Box::new(self.parse_block_statement()?)));
        }
        if tok.kind().is_var_keyword() {
            return Ok(Some(Box::new(self.parse_var_statement()?)));
        }
        if tok.kind().is_semicolon() {
            return Ok(Some(Box::new(ast::EmptyStmtNode::new())));
        }
        if tok.kind().is_if_keyword() {
            return Ok(Some(self.parse_if_statement()?));
        }

        if let Some(boxed_expr) = self.try_parse_expression(tok, Precedence::lowest())? {
            return Ok(Some(Box::new(ast::ExpressionStmtNode::new(boxed_expr))));
        }

        self.rewind_position(begin_position);
        Ok(None)
    }

    fn parse_block_statement(&mut self) -> ParseResult<ast::BlockStmtNode> {
        // FIXME: Parse list of statements.
        self.must_expect_token(TokenKind::close_brace())?;
        Ok(ast::BlockStmtNode::new())
    }

    fn parse_var_statement(&mut self) -> ParseResult<ast::VarStmtNode> {
        let mut var_statement = ast::VarStmtNode::new();
        loop {
            // FIXME: Support initializer expressions.
            // For now, we match only a VarName ("," VarName)* ";"
            let name_token = match self.expect_get_token(TokenKind::identifier())? {
                Some(token) => token,
                None => { return Err(ParseError::ExpectedVariableName); }
            };
            var_statement.add_variable(name_token);

            let next_tok = self.next_token()?;
            if next_tok.kind().is_semicolon() {
                break;
            }
            if next_tok.kind().is_comma() {
                continue;
            }
            return Err(ParseError::ExpectedCommaOrSemicolon);
        }
        Ok(var_statement)
    }

    fn parse_if_statement(&mut self) -> ParseResult<Box<AstNode>> {
        // FIXME: implement this after expression parsing is supported.
        panic!("parse_if_statement is not implemented.")
    }

    fn parse_expression(&mut self, precedence: Precedence) -> ParseResult<Box<AstNode>> {
        let token = self.next_token()?;
        if let Some(boxed_expr) = self.try_parse_expression(token, precedence)? {
            return Ok(boxed_expr);
        }

        Err(ParseError::ExpectedExpression)
    }

    fn try_parse_expression(&mut self, tok: FullToken, precedence: Precedence)
        -> MaybeParseResult<Box<AstNode>>
    {
        if tok.kind().is_identifier() {
            let name_expr = Box::new(ast::NameExpressionNode::new(tok));
            return Ok(Some(self.parse_rest_of_expression(name_expr, precedence)?));
        }
        Ok(None)
    }

    fn parse_rest_of_expression(&mut self, left_expr: Box<AstNode>, precedence: Precedence)
        -> ParseResult<Box<AstNode>>
    {
        assert!(left_expr.is_expression());

        let mut cur_expr = left_expr;
        loop {
            let position = self.mark_position();
            let tok = self.next_token()?;
            if tok.kind().is_comma() {
                if precedence >= Precedence::comma() {
                    self.rewind_position(position);
                    return Ok(cur_expr);
                }

                let right_expr = self.parse_expression(Precedence::comma())?;
                cur_expr = Box::new(ast::CommaExpressionNode::new(cur_expr, right_expr));
                continue;
            }

            if tok.kind().is_assignment_op() {
                // FIXME: Check that cur_expr is a valid lvalue expression.
                // Return syntax error if not.

                // Assignment associates right-to-left, so we use '>' for precedence
                // instead of '>='.
                if precedence > Precedence::assignment() {
                    self.rewind_position(position);
                    return Ok(cur_expr);
                }

                let right_expr = self.parse_expression(Precedence::assignment())?;
                cur_expr = Box::new(ast::AssignmentExpressionNode::new(tok, cur_expr, right_expr));
                continue;
            }

            if tok.kind().is_question() {
                if precedence > Precedence::conditional() {
                    self.rewind_position(position);
                    return Ok(cur_expr);
                }

                let if_expr = self.parse_expression(Precedence::assignment())?;
                self.must_expect_token(TokenKind::colon())?;
                let else_expr = self.parse_expression(Precedence::assignment())?;
                cur_expr = Box::new(ast::ConditionalExpressionNode::new(cur_expr, if_expr,
                                                                        else_expr));
                continue;
            }

            // Unknown token terminates expression.
            self.rewind_position(position);
            break;
        }

        Ok(cur_expr)
    }

    fn must_expect_token(&mut self, kind: TokenKind) -> ParseResult<()> {
        // Mark the position so we can backtrack.
        let position = self.mark_position();
        let token = self.next_token()?;
        if token.kind() == kind {
            Ok(())
        } else {
            self.rewind_position(position);
            Err(ParseError::UnexpectedToken {expected: kind, got: token.kind()})
        }
    }

    fn expect_token(&mut self, kind: TokenKind) -> ParseResult<bool> {
        // Mark the position so we can backtrack.
        let position = self.mark_position();
        let token = self.next_token()?;
        if token.kind() == kind {
            Ok(true)
        } else {
            self.rewind_position(position);
            Ok(false)
        }
    }
    fn expect_get_token(&mut self, kind: TokenKind) -> ParseResult<Option<FullToken>> {
        // Mark the position so we can backtrack.
        let position = self.mark_position();
        let token = self.next_token()?;
        if token.kind() == kind {
            Ok(Some(token))
        } else {
            self.rewind_position(position);
            Ok(None)
        }
    }

    fn mark_position(&mut self) -> TokenizerPosition {
        self.tokenizer.mark_position()
    }
    fn rewind_position(&mut self, position: TokenizerPosition) {
        self.tokenizer.rewind_position(position)
    }

    fn next_token(&mut self) -> ParseResult<FullToken> {
        self.next_token_impl(/* check_kw = */ true, /* want_newlines = */ false)
    }
    fn next_token_no_keywords(&mut self) -> ParseResult<FullToken> {
        self.next_token_impl(/* check_kw = */ true, /* want_newlines = */ false)
    }

    fn next_token_impl(&mut self, check_kw: bool, want_newlines: bool) -> ParseResult<FullToken> {
        loop {
            let token = self.tokenizer.next_token(/* check_kw = */ true);
            let kind = token.kind();
            // Ignore whitespace and comment and newline tokens.
            if kind.is_whitespace() || kind.is_comment() {
                continue;
            }
            if !want_newlines && kind.is_newline() {
                continue;
            }
            if kind.is_error() {
                return Err(ParseError::ErrorToken(token));
            }
            let kw_str = if check_kw { "kw" } else { "no-kw" };
            let nl_str = if want_newlines { "nl" } else { "no-nl" };
            println!("next_token({}, {}): {}", kw_str, nl_str, token.token_string());
            return Ok(token);
        }
    }
}
<commit_msg>Changed ast builder methods to always return boxed nodes, even if the concrete type is used.<commit_after>
use std::fmt;

use parser::ast;
use parser::ast::{AstKind, AstNode};
use parser::input_stream::{InputStream, StreamPosition};
use parser::precedence::Precedence;
use parser::token_kind::TokenKind;
use parser::tokenizer::{Token, TokenError, TokenLocation, Tokenizer, TokenizerMode,
                        TokenizerPosition};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FullToken {
    kind: TokenKind,
    location: TokenLocation
}
impl FullToken {
    pub fn new(kind: TokenKind, location: TokenLocation) -> FullToken {
        FullToken {
            kind: kind,
            location: location
        }
    }

    pub fn location(&self) -> &TokenLocation {
        &self.location
    }
}
impl Token for FullToken {
    fn make(kind: TokenKind, location: TokenLocation) -> FullToken {
        FullToken::new(kind, location)
    }
    fn kind(&self) -> TokenKind {
        self.kind
    }
    fn start_offset(&self) -> StreamPosition {
        self.location.start_offset()
    }
    fn write_token(&self, w: &mut fmt::Write) -> Result<(), fmt::Error> {
        write!(w, "Token({}@{})", self.kind.name(), self.location.range_string())
    }
}

pub struct FullTokenizerMode {
}
impl TokenizerMode for FullTokenizerMode {
    type Tok = FullToken;

    fn cares_about_newline(&self) -> bool {
        false
    }
    fn note_newline(&mut self) {}
}

#[derive(Debug, Clone)]
pub enum ParseError {
    Unspecified,
    ErrorToken(FullToken),
    TokenizerError(TokenError),
    UnexpectedToken{expected:TokenKind, got:TokenKind},
    ExpectedVariableName,
    ExpectedCommaOrSemicolon,
    ExpectedExpression
}
pub type ParseResult<T> = Result<T, ParseError>;
pub type MaybeParseResult<T> = ParseResult<Option<T>>;

pub struct AstBuilder<STREAM: InputStream> {
    tokenizer: Tokenizer<STREAM, FullTokenizerMode>
}
impl<STREAM: InputStream> AstBuilder<STREAM> {
    pub fn new(stream: STREAM) -> AstBuilder<STREAM> {
        AstBuilder {
            tokenizer: Tokenizer::new(stream, FullTokenizerMode{})
        }
    }

    pub fn read_and_print_tokens(&mut self) {
        // Just read tokens and print them out until we're done, then return Error.
        loop {
            let token = self.next_token().unwrap();
            println!("Token: {}", token.kind().name());
            if token.kind().is_error() {
                panic!("Got token error: {:?}", self.tokenizer.get_error());
            }
            if token.kind().is_end() {
                break;
            }
        }
    }

    pub fn parse_program(&mut self) -> ParseResult<Box<ast::ProgramNode>> {
        let mut program_node = ast::ProgramNode::new();
        loop {
            match self.try_parse_statement()? {
                Some(boxed_source_element) => {
                    program_node.add_source_element(boxed_source_element);
                }
                None => {
                    break;
                }
            }
        }

        // Must have reached end of stream.
        self.must_expect_token(TokenKind::end())?;

        Ok(Box::new(program_node))
    }

    fn try_parse_statement(&mut self) -> MaybeParseResult<Box<AstNode>> {
        let begin_position = self.mark_position();
        let tok = self.next_token()?;
        if tok.kind().is_open_brace() {
            // FIXME: parse either a block or an object literal.
            // The spec says a '{' at the statement level can only be a block,
            // but practical implementations seem to allow bare object literals.
            return Ok(Some(self.parse_block_statement()?));
        }
        if tok.kind().is_var_keyword() {
            return Ok(Some(self.parse_var_statement()?));
        }
        if tok.kind().is_semicolon() {
            return Ok(Some(Box::new(ast::EmptyStmtNode::new())));
        }
        if tok.kind().is_if_keyword() {
            return Ok(Some(self.parse_if_statement()?));
        }

        if let Some(boxed_expr) = self.try_parse_expression(tok, Precedence::lowest())? {
            return Ok(Some(Box::new(ast::ExpressionStmtNode::new(boxed_expr))));
        }

        self.rewind_position(begin_position);
        Ok(None)
    }

    fn parse_block_statement(&mut self) -> ParseResult<Box<ast::BlockStmtNode>> {
        // FIXME: Parse list of statements.
        self.must_expect_token(TokenKind::close_brace())?;
        Ok(Box::new(ast::BlockStmtNode::new()))
    }

    fn parse_var_statement(&mut self) -> ParseResult<Box<ast::VarStmtNode>> {
        let mut var_statement = Box::new(ast::VarStmtNode::new());
        loop {
            // FIXME: Support initializer expressions.
            // For now, we match only a VarName ("," VarName)* ";"
            let name_token = match self.expect_get_token(TokenKind::identifier())? {
                Some(token) => token,
                None => { return Err(ParseError::ExpectedVariableName); }
            };
            var_statement.add_variable(name_token);

            let next_tok = self.next_token()?;
            if next_tok.kind().is_semicolon() {
                break;
            }
            if next_tok.kind().is_comma() {
                continue;
            }
            return Err(ParseError::ExpectedCommaOrSemicolon);
        }
        Ok(var_statement)
    }

    fn parse_if_statement(&mut self) -> ParseResult<Box<AstNode>> {
        // FIXME: implement this after expression parsing is supported.
        panic!("parse_if_statement is not implemented.")
    }

    fn parse_expression(&mut self, precedence: Precedence) -> ParseResult<Box<AstNode>> {
        let token = self.next_token()?;
        if let Some(boxed_expr) = self.try_parse_expression(token, precedence)? {
            return Ok(boxed_expr);
        }

        Err(ParseError::ExpectedExpression)
    }

    fn try_parse_expression(&mut self, tok: FullToken, precedence: Precedence)
        -> MaybeParseResult<Box<AstNode>>
    {
        if tok.kind().is_identifier() {
            let name_expr = Box::new(ast::NameExpressionNode::new(tok));
            return Ok(Some(self.parse_rest_of_expression(name_expr, precedence)?));
        }
        Ok(None)
    }

    fn parse_rest_of_expression(&mut self, left_expr: Box<AstNode>, precedence: Precedence)
        -> ParseResult<Box<AstNode>>
    {
        assert!(left_expr.is_expression());

        let mut cur_expr = left_expr;
        loop {
            let position = self.mark_position();
            let tok = self.next_token()?;
            if tok.kind().is_comma() {
                if precedence >= Precedence::comma() {
                    self.rewind_position(position);
                    return Ok(cur_expr);
                }

                let right_expr = self.parse_expression(Precedence::comma())?;
                cur_expr = Box::new(ast::CommaExpressionNode::new(cur_expr, right_expr));
                continue;
            }

            if tok.kind().is_assignment_op() {
                // FIXME: Check that cur_expr is a valid lvalue expression.
                // Return syntax error if not.

                // Assignment associates right-to-left, so we use '>' for precedence
                // instead of '>='.
                if precedence > Precedence::assignment() {
                    self.rewind_position(position);
                    return Ok(cur_expr);
                }

                let right_expr = self.parse_expression(Precedence::assignment())?;
                cur_expr = Box::new(ast::AssignmentExpressionNode::new(tok, cur_expr, right_expr));
                continue;
            }

            if tok.kind().is_question() {
                if precedence > Precedence::conditional() {
                    self.rewind_position(position);
                    return Ok(cur_expr);
                }

                let if_expr = self.parse_expression(Precedence::assignment())?;
                self.must_expect_token(TokenKind::colon())?;
                let else_expr = self.parse_expression(Precedence::assignment())?;
                cur_expr = Box::new(ast::ConditionalExpressionNode::new(cur_expr, if_expr,
                                                                        else_expr));
                continue;
            }

            // Unknown token terminates expression.
            self.rewind_position(position);
            break;
        }

        Ok(cur_expr)
    }

    fn must_expect_token(&mut self, kind: TokenKind) -> ParseResult<()> {
        // Mark the position so we can backtrack.
        let position = self.mark_position();
        let token = self.next_token()?;
        if token.kind() == kind {
            Ok(())
        } else {
            self.rewind_position(position);
            Err(ParseError::UnexpectedToken {expected: kind, got: token.kind()})
        }
    }

    fn expect_token(&mut self, kind: TokenKind) -> ParseResult<bool> {
        // Mark the position so we can backtrack.
        let position = self.mark_position();
        let token = self.next_token()?;
        if token.kind() == kind {
            Ok(true)
        } else {
            self.rewind_position(position);
            Ok(false)
        }
    }
    fn expect_get_token(&mut self, kind: TokenKind) -> ParseResult<Option<FullToken>> {
        // Mark the position so we can backtrack.
        let position = self.mark_position();
        let token = self.next_token()?;
        if token.kind() == kind {
            Ok(Some(token))
        } else {
            self.rewind_position(position);
            Ok(None)
        }
    }

    fn mark_position(&mut self) -> TokenizerPosition {
        self.tokenizer.mark_position()
    }
    fn rewind_position(&mut self, position: TokenizerPosition) {
        self.tokenizer.rewind_position(position)
    }

    fn next_token(&mut self) -> ParseResult<FullToken> {
        self.next_token_impl(/* check_kw = */ true, /* want_newlines = */ false)
    }
    fn next_token_no_keywords(&mut self) -> ParseResult<FullToken> {
        self.next_token_impl(/* check_kw = */ true, /* want_newlines = */ false)
    }

    fn next_token_impl(&mut self, check_kw: bool, want_newlines: bool) -> ParseResult<FullToken> {
        loop {
            let token = self.tokenizer.next_token(/* check_kw = */ true);
            let kind = token.kind();
            // Ignore whitespace and comment and newline tokens.
            if kind.is_whitespace() || kind.is_comment() {
                continue;
            }
            if !want_newlines && kind.is_newline() {
                continue;
            }
            if kind.is_error() {
                return Err(ParseError::ErrorToken(token));
            }
            let kw_str = if check_kw { "kw" } else { "no-kw" };
            let nl_str = if want_newlines { "nl" } else { "no-nl" };
            println!("next_token({}, {}): {}", kw_str, nl_str, token.token_string());
            return Ok(token);
        }
    }
}
<|endoftext|>


<commit_before>extern crate libc;
use std::ptr;
use std::mem;
use std::slice::from_raw_parts;
use std::string::String;
use std::collections::HashMap;
use std::boxed::Box;
use x11::xlib;
use x11::keysym;

use super::config::Config;
use super::workspace::{ Workspace, Workspaces };
use super::handler;

fn get_text_property(display: *mut xlib::Display, window: xlib::Window, atom: xlib::Atom) -> Option<String>{
    unsafe{
        let mut prop: xlib::XTextProperty = mem::zeroed();
        let r = xlib::XGetTextProperty(display, window, &mut prop, atom);
        if r == 0 {
            None
        }else{
            let s = String::from_raw_parts(prop.value, prop.nitems as usize, prop.nitems as usize + 1 ).clone();
            let text = Some(s);
            xlib::XFree(prop.value as *mut libc::c_void);
            text
        }
    }
}

fn get_window_property(display: *mut xlib::Display, window: xlib::Window, atom: xlib::Atom) {
    unsafe {
        let mut actual_type_return: u64 = 0;
        let mut actual_format_return: i32 = 0;
        let mut nitem_return: libc::c_ulong = 0;
        let mut bytes_after_return: libc::c_ulong = 0;
        let mut prop_return: *mut libc::c_uchar = mem::zeroed();

        let r = xlib::XGetWindowProperty(display, window, xlib::XA_WM_NAME,
                                         0, 0xFFFFFFFF, 0, 0,
                                         &mut actual_type_return,
                                         &mut actual_format_return,
                                         &mut nitem_return,
                                         &mut bytes_after_return,
                                         &mut prop_return);

        debug!("result get wp {}", r);
        if r == 0 {

        }else{
            debug!("actual format return {}", actual_format_return);
            if actual_format_return == 0 {

            }
            else {
                // let s = from_raw_parts(prop_return as *const libc::c_ulong, nitems_return as usize).iter()
                //     .map(|&c| c as u64)
                //     .collect();
                let s = String::from_raw_parts(prop_return, nitem_return as usize, nitem_return as usize + 1 );
                println!("{}", s);

            }
        }
    }
}

pub struct WindowManager {
    pub display: *mut xlib::Display,
    pub screen_num: libc::c_int,
    pub root:    xlib::Window,

    pub workspaces: Workspaces,
    config: Config
}

impl WindowManager {
    pub fn new() -> WindowManager {
	unsafe {
	    let display = xlib::XOpenDisplay(ptr::null());
            if display == ptr::null_mut() {
                panic!("can't open display")
            }
	    let screen_num  = xlib::XDefaultScreen(display);
	    let root    = xlib::XRootWindow(display, screen_num);

	    let mut wm = WindowManager {
	        display: display,
                screen_num: screen_num,
	        root:    root,

                config: Config::load(),
                workspaces: Workspaces::new()
            };
            wm.workspaces.create('1');
            wm
	}
    }

    pub fn run(&mut self) {
        let mask = xlib::SubstructureRedirectMask | xlib::SubstructureNotifyMask;
        let keymask = xlib::KeyPressMask | xlib::KeyReleaseMask;
        unsafe{
            xlib::XSelectInput(self.display, self.root,
                               mask | keymask);
            xlib::XSync(self.display, 0);
        }

        loop {
            //handle events here
            let mut e = xlib::XEvent{
                pad: [0; 24],
            };

            unsafe {
                xlib::XNextEvent(self.display, &mut e);
                let t = e.get_type();
                match t {
                    xlib::CreateNotify => {
                        let event: xlib::XCreateWindowEvent = From::from(e);

                        if event.override_redirect != 0 {
                            continue;
                        }

                        // change attributes before display
                        let mut attrs: xlib::XSetWindowAttributes = mem::zeroed();
                        attrs.event_mask = xlib::KeyReleaseMask;
                        let valuemask = xlib::CWEventMask;
                        xlib::XChangeWindowAttributes(self.display, event.window, valuemask, &mut attrs);

                        // add to workspace
                        // let mut w = self.workspaces.get_mut(&'1');
                        // match w {
                        //     Some(workspace) => {
                        //         workspace.add(event.window);
                        //         workspace.config(self);
                        //     }
                        //     None => {}
                        // }

                        // get window property
                        let n = get_text_property(self.display, event.window, xlib::XA_WM_NAME);
                        match n {
                            Some(s) => {debug!("create window {} for {}", event.window, s);}
                            None => {}
                        }
                    }
                    xlib::DestroyNotify => {
                        debug!("destroy notify");
                        let event: xlib::XDestroyWindowEvent = From::from(e);
                    }
                    xlib::MapNotify => {
                        debug!("map notify");
                        let event: xlib::XMapEvent = From::from(e);
                        println!("event {}", event.event);
                        println!("w {}", event.window);
                    }
                    xlib::UnmapNotify => {
                        debug!("unmap notify");
                        let event: xlib::XUnmapEvent = From::from(e);
                        let workspace = self.workspaces.current_workspace();
                        match workspace.contain(event.window) {
                            Some(index) => {
                                workspace.remove(event.window);
                                workspace.config(self.display, self.screen_num);
                            }
                            None => {}
                        }
                    }
                    xlib::MapRequest => {
                        debug!("map request");
                        let event: xlib::XMapRequestEvent = From::from(e);
                        println!("w {}", event.window);
                        xlib::XMapWindow(self.display, event.window);

                        // add app top-level window to workspace
                        let mut attrs: xlib::XWindowAttributes = mem::zeroed();
                        xlib::XGetWindowAttributes(self.display, event.window, &mut attrs);

                        let mut window_return: xlib::Window = 0;
                        let transientfor_hint = xlib::XGetTransientForHint(self.display, event.window, &mut window_return);
                        println!("trans hint {}", transientfor_hint);
                        if attrs.override_redirect == 0 && transientfor_hint == 0 {
                            debug!("top level window");
                            let mut workspace = self.workspaces.current();
                            if workspace.contain(event.window).is_none() {
                                workspace.add(event.window);
                                workspace.config(self.display, self.screen_num);
                            }
                        }

                    }
                    xlib::ClientMessage => {
                        debug!("client message");
                        // let event = cast_event::<xlib::XClientMessageEvent>(&e);
                        let event: xlib::XClientMessageEvent = From::from(e);
                    }
                    xlib::KeyRelease => {
                        debug!("key release");
                        let mut event: xlib::XKeyEvent = From::from(e);

                        if event.state > 0 {
                            let mut sym: xlib::KeySym = 0;
                            let status: *mut xlib::XComposeStatus = ptr::null_mut();
                            xlib::XLookupString(&mut event, ptr::null_mut(), 0, &mut sym, status);

                            let b = handler::KeyBind {
                                key: sym,
                                mask: event.state,
                            };
                            debug!("key {} {}", event.state, sym);
                            // let mut h: Box<handler::Handler>;
                            match self.config.bindsyms.get_mut(&b) {
                                Some(handler) => {
                                    handler.handle(&mut self.workspaces, self.display, self.screen_num);
                                }
                                None => {
                                    println!("no bind");
                                    continue;
                                }
                            }
                        }
                    }
                    xlib::ConfigureNotify => {
                        debug!("configure notify");
                    }
                    xlib::ConfigureRequest =>{
                        debug!("configure request");
                        // let mut change = xlib::XWindowChanges {
                        //     x: event.x,
                        //     y: event.y,
                        //     width: event.width,
                        //     height: event.height,
                        //     border_width: event.border_width,
                        //     sibling: event.above,
                        //     stack_mode: event.detail
                        // };
                        // debug!("config x: {}, y: {}, width: {}, height: {}",
                        //        change.x, change.y, change.width, change.height);
                        let mut event: xlib::XConfigureRequestEvent = From::from(e);
                        // xlib::XConfigureWindow(event.display, event.window, event.value_mask as u32, &mut change);

                    }
                    _ => {
                        debug!("unhandled event {}", t);
                    }
                }
            }
        }
    }
}
<commit_msg>empty error handler<commit_after>extern crate libc;
use std::ptr;
use std::mem;
use std::slice::from_raw_parts;
use std::string::String;
use std::collections::HashMap;
use std::boxed::Box;
use x11::xlib;
use x11::keysym;

use super::config::Config;
use super::workspace::{ Workspace, Workspaces };
use super::handler;

unsafe extern fn error_handler(display: *mut xlib::Display, event: *mut xlib::XErrorEvent) -> libc::c_int {
    // match event.error_code {
    //     xlib::BadAtom => {
    //         println!("bad atom");
    //     }
    //     _ => {}
    // }
    1
}

fn get_text_property(display: *mut xlib::Display, window: xlib::Window, atom: xlib::Atom) -> Option<String>{
    unsafe{
        let mut prop: xlib::XTextProperty = mem::zeroed();
        let r = xlib::XGetTextProperty(display, window, &mut prop, atom);
        if r == 0 {
            None
        }else{
            let s = String::from_raw_parts(prop.value, prop.nitems as usize, prop.nitems as usize + 1 ).clone();
            let text = Some(s);
            xlib::XFree(prop.value as *mut libc::c_void);
            text
        }
    }
}

fn get_window_property(display: *mut xlib::Display, window: xlib::Window, atom: xlib::Atom) {
    unsafe {
        let mut actual_type_return: u64 = 0;
        let mut actual_format_return: i32 = 0;
        let mut nitem_return: libc::c_ulong = 0;
        let mut bytes_after_return: libc::c_ulong = 0;
        let mut prop_return: *mut libc::c_uchar = mem::zeroed();

        let r = xlib::XGetWindowProperty(display, window, xlib::XA_WM_NAME,
                                         0, 0xFFFFFFFF, 0, 0,
                                         &mut actual_type_return,
                                         &mut actual_format_return,
                                         &mut nitem_return,
                                         &mut bytes_after_return,
                                         &mut prop_return);

        debug!("result get wp {}", r);
        if r == 0 {

        }else{
            debug!("actual format return {}", actual_format_return);
            if actual_format_return == 0 {

            }
            else {
                // let s = from_raw_parts(prop_return as *const libc::c_ulong, nitems_return as usize).iter()
                //     .map(|&c| c as u64)
                //     .collect();
                let s = String::from_raw_parts(prop_return, nitem_return as usize, nitem_return as usize + 1 );
                println!("{}", s);

            }
        }
    }
}

pub struct WindowManager {
    pub display: *mut xlib::Display,
    pub screen_num: libc::c_int,
    pub root:    xlib::Window,

    pub workspaces: Workspaces,
    config: Config
}

impl WindowManager {
    pub fn new() -> WindowManager {
	unsafe {
	    let display = xlib::XOpenDisplay(ptr::null());
            if display == ptr::null_mut() {
                panic!("can't open display")
            }
	    let screen_num  = xlib::XDefaultScreen(display);
	    let root    = xlib::XRootWindow(display, screen_num);

	    let mut wm = WindowManager {
	        display: display,
                screen_num: screen_num,
	        root:    root,

                config: Config::load(),
                workspaces: Workspaces::new()
            };
            wm.workspaces.create('1');
            wm
	}
    }

    pub fn run(&mut self) {
        let mask = xlib::SubstructureRedirectMask | xlib::SubstructureNotifyMask;
        let keymask = xlib::KeyPressMask | xlib::KeyReleaseMask;
        unsafe{
            xlib::XSetErrorHandler(Some(error_handler));
            xlib::XSelectInput(self.display, self.root,
                               mask | keymask);
            xlib::XSync(self.display, 0);
        }

        loop {
            //handle events here
            let mut e = xlib::XEvent{
                pad: [0; 24],
            };

            unsafe {
                xlib::XNextEvent(self.display, &mut e);
                let t = e.get_type();
                match t {
                    xlib::CreateNotify => {
                        let event: xlib::XCreateWindowEvent = From::from(e);

                        if event.override_redirect != 0 {
                            continue;
                        }

                        // change attributes before display
                        let mut attrs: xlib::XSetWindowAttributes = mem::zeroed();
                        attrs.event_mask = xlib::KeyReleaseMask;
                        let valuemask = xlib::CWEventMask;
                        xlib::XChangeWindowAttributes(self.display, event.window, valuemask, &mut attrs);

                        // add to workspace
                        // let mut w = self.workspaces.get_mut(&'1');
                        // match w {
                        //     Some(workspace) => {
                        //         workspace.add(event.window);
                        //         workspace.config(self);
                        //     }
                        //     None => {}
                        // }

                        // get window property
                        let n = get_text_property(self.display, event.window, xlib::XA_WM_NAME);
                        match n {
                            Some(s) => {debug!("create window {} for {}", event.window, s);}
                            None => {}
                        }
                    }
                    xlib::DestroyNotify => {
                        debug!("destroy notify");
                        let event: xlib::XDestroyWindowEvent = From::from(e);
                    }
                    xlib::MapNotify => {
                        debug!("map notify");
                        let event: xlib::XMapEvent = From::from(e);
                        println!("event {}", event.event);
                        println!("w {}", event.window);
                    }
                    xlib::UnmapNotify => {
                        debug!("unmap notify");
                        let event: xlib::XUnmapEvent = From::from(e);
                        let workspace = self.workspaces.current_workspace();
                        match workspace.contain(event.window) {
                            Some(index) => {
                                workspace.remove(event.window);
                                workspace.config(self.display, self.screen_num);
                            }
                            None => {}
                        }
                    }
                    xlib::MapRequest => {
                        debug!("map request");
                        let event: xlib::XMapRequestEvent = From::from(e);
                        println!("w {}", event.window);
                        xlib::XMapWindow(self.display, event.window);

                        // add app top-level window to workspace
                        let mut attrs: xlib::XWindowAttributes = mem::zeroed();
                        xlib::XGetWindowAttributes(self.display, event.window, &mut attrs);

                        let mut window_return: xlib::Window = 0;
                        let transientfor_hint = xlib::XGetTransientForHint(self.display, event.window, &mut window_return);
                        println!("trans hint {}", transientfor_hint);
                        if attrs.override_redirect == 0 && transientfor_hint == 0 {
                            debug!("top level window");
                            let mut workspace = self.workspaces.current();
                            if workspace.contain(event.window).is_none() {
                                workspace.add(event.window);
                                workspace.config(self.display, self.screen_num);
                            }
                        }

                    }
                    xlib::ClientMessage => {
                        debug!("client message");
                        // let event = cast_event::<xlib::XClientMessageEvent>(&e);
                        let event: xlib::XClientMessageEvent = From::from(e);
                    }
                    xlib::KeyRelease => {
                        debug!("key release");
                        let mut event: xlib::XKeyEvent = From::from(e);

                        if event.state > 0 {
                            let mut sym: xlib::KeySym = 0;
                            let status: *mut xlib::XComposeStatus = ptr::null_mut();
                            xlib::XLookupString(&mut event, ptr::null_mut(), 0, &mut sym, status);

                            let b = handler::KeyBind {
                                key: sym,
                                mask: event.state,
                            };
                            debug!("key {} {}", event.state, sym);
                            // let mut h: Box<handler::Handler>;
                            match self.config.bindsyms.get_mut(&b) {
                                Some(handler) => {
                                    handler.handle(&mut self.workspaces, self.display, self.screen_num);
                                }
                                None => {
                                    println!("no bind");
                                    continue;
                                }
                            }
                        }
                    }
                    xlib::ConfigureNotify => {
                        debug!("configure notify");
                    }
                    xlib::ConfigureRequest =>{
                        debug!("configure request");
                        // let mut change = xlib::XWindowChanges {
                        //     x: event.x,
                        //     y: event.y,
                        //     width: event.width,
                        //     height: event.height,
                        //     border_width: event.border_width,
                        //     sibling: event.above,
                        //     stack_mode: event.detail
                        // };
                        // debug!("config x: {}, y: {}, width: {}, height: {}",
                        //        change.x, change.y, change.width, change.height);
                        let mut event: xlib::XConfigureRequestEvent = From::from(e);
                        // xlib::XConfigureWindow(event.display, event.window, event.value_mask as u32, &mut change);

                    }
                    _ => {
                        debug!("unhandled event {}", t);
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use any_pointer;
use MessageSize;
use capability::{Params, Request, RemotePromise, Results};
use std::cell::RefCell;
use std::rc::Rc;

pub trait ResponseHook {
    fn get<'a>(&'a self) -> ::Result<any_pointer::Reader<'a>>;
}

pub trait RequestHook {
    fn init<'a>(&'a mut self) -> any_pointer::Builder<'a>;
    fn send<'a>(self : Box<Self>) -> RemotePromise<any_pointer::Owned>;
}

pub trait ClientHook {
    fn new_call(&self,
                interface_id: u64,
                method_id: u16,
                size_hint: Option<MessageSize>)
                -> Request<any_pointer::Owned, any_pointer::Owned>;

    fn call(&self, interface_id: u64, method_id: u16,
            params: Box<ParamsHook>, results: Box<ResultsHook>);

    // HACK
    fn get_descriptor(&self) -> Box<::std::any::Any>;
}

#[cfg(feature = "rpc")]
pub trait ServerHook: 'static {
    fn new_client(server: Box<::capability::Server>) -> ::capability::Client;
}

pub trait ResultsHook {
    fn get<'a>(&'a mut self) -> any_pointer::Builder<'a>;
    fn fail(self: Box<Self>, message: String);
    fn unimplemented(self: Box<Self>);
    fn disconnected(self: Box<Self>);
    fn overloaded(self: Box<Self>);
}

pub trait ParamsHook {
    fn get<'a>(&'a self) -> any_pointer::Reader<'a>;
}

// Where should this live?
pub fn internal_get_typed_params<T>(typeless: Params<any_pointer::Owned>) -> Params<T> {
    Params { hook: typeless.hook, marker: ::std::marker::PhantomData }
}

pub fn internal_get_typed_results<T>(typeless: Results<any_pointer::Owned>) -> Results<T> {
    Results { hook: typeless.hook, marker: ::std::marker::PhantomData }
}

pub trait PipelineHook {
    fn get_pipelined_cap(&mut self, ops: &[PipelineOp]) -> Rc<RefCell<Box<ClientHook>>>;

    /// Version of get_pipelined_cap() passing the array by move. May avoid a copy in some cases.
    /// Default implementation just calls the other version.
    fn get_pipelined_cap_move(&mut self, ops: Vec<PipelineOp>) -> Rc<RefCell<Box<ClientHook>>> {
        self.get_pipelined_cap(&ops)
    }
}

#[derive(Clone, Copy)]
pub enum PipelineOp {
    Noop,
    GetPointerField(u16),
}

<commit_msg>get_brand()<commit_after>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use any_pointer;
use MessageSize;
use capability::{Params, Request, RemotePromise, Results};
use std::cell::RefCell;
use std::rc::Rc;

pub trait ResponseHook {
    fn get<'a>(&'a self) -> ::Result<any_pointer::Reader<'a>>;
}

pub trait RequestHook {
    fn init<'a>(&'a mut self) -> any_pointer::Builder<'a>;
    fn send<'a>(self : Box<Self>) -> RemotePromise<any_pointer::Owned>;
}

pub trait ClientHook {
    fn new_call(&self,
                interface_id: u64,
                method_id: u16,
                size_hint: Option<MessageSize>)
                -> Request<any_pointer::Owned, any_pointer::Owned>;

    fn call(&self, interface_id: u64, method_id: u16,
            params: Box<ParamsHook>, results: Box<ResultsHook>);

    fn get_brand(&self) -> usize;

    // HACK
    fn get_descriptor(&self) -> Box<::std::any::Any>;
}

#[cfg(feature = "rpc")]
pub trait ServerHook: 'static {
    fn new_client(server: Box<::capability::Server>) -> ::capability::Client;
}

pub trait ResultsHook {
    fn get<'a>(&'a mut self) -> any_pointer::Builder<'a>;
    fn fail(self: Box<Self>, message: String);
    fn unimplemented(self: Box<Self>);
    fn disconnected(self: Box<Self>);
    fn overloaded(self: Box<Self>);
}

pub trait ParamsHook {
    fn get<'a>(&'a self) -> any_pointer::Reader<'a>;
}

// Where should this live?
pub fn internal_get_typed_params<T>(typeless: Params<any_pointer::Owned>) -> Params<T> {
    Params { hook: typeless.hook, marker: ::std::marker::PhantomData }
}

pub fn internal_get_typed_results<T>(typeless: Results<any_pointer::Owned>) -> Results<T> {
    Results { hook: typeless.hook, marker: ::std::marker::PhantomData }
}

pub trait PipelineHook {
    fn get_pipelined_cap(&mut self, ops: &[PipelineOp]) -> Rc<RefCell<Box<ClientHook>>>;

    /// Version of get_pipelined_cap() passing the array by move. May avoid a copy in some cases.
    /// Default implementation just calls the other version.
    fn get_pipelined_cap_move(&mut self, ops: Vec<PipelineOp>) -> Rc<RefCell<Box<ClientHook>>> {
        self.get_pipelined_cap(&ops)
    }
}

#[derive(Clone, Copy)]
pub enum PipelineOp {
    Noop,
    GetPointerField(u16),
}

<|endoftext|>


<commit_before>use chrono::{NaiveDateTime, Utc};
use diesel::data_types::PgInterval;
use diesel::prelude::*;
use std::time::Duration;

use crate::schema::{publish_limit_buckets, publish_rate_overrides};
use crate::util::errors::{AppResult, TooManyRequests};

#[derive(Debug, Clone, Copy)]
pub struct PublishRateLimit {
    pub rate: Duration,
    pub burst: i32,
}

impl Default for PublishRateLimit {
    fn default() -> Self {
        Self {
            rate: Duration::from_secs(60) * 10,
            burst: 30,
        }
    }
}

#[derive(Queryable, Insertable, Debug, PartialEq, Clone, Copy)]
#[table_name = "publish_limit_buckets"]
#[allow(dead_code)] // Most fields only read in tests
struct Bucket {
    user_id: i32,
    tokens: i32,
    last_refill: NaiveDateTime,
}

impl PublishRateLimit {
    pub fn check_rate_limit(&self, uploader: i32, conn: &PgConnection) -> AppResult<()> {
        let bucket = self.take_token(uploader, Utc::now().naive_utc(), conn)?;
        if bucket.tokens >= 1 {
            Ok(())
        } else {
            Err(Box::new(TooManyRequests {
                retry_after: bucket.last_refill + chrono::Duration::from_std(self.rate).unwrap(),
            }))
        }
    }

    /// Refill a user's bucket as needed, take a token from it,
    /// and returns the result.
    ///
    /// The number of tokens remaining will always be between 0 and self.burst.
    /// If the number is 0, the request should be rejected, as the user doesn't
    /// have a token to take. Technically a "full" bucket would have
    /// `self.burst + 1` tokens in it, but that value would never be returned
    /// since we only refill buckets when trying to take a token from it.
    fn take_token(
        &self,
        uploader: i32,
        now: NaiveDateTime,
        conn: &PgConnection,
    ) -> QueryResult<Bucket> {
        use self::publish_limit_buckets::dsl::*;
        use diesel::sql_types::{Double, Interval, Text, Timestamp};

        sql_function!(fn date_part(x: Text, y: Timestamp) -> Double);
        sql_function! {
            #[sql_name = "date_part"]
            fn interval_part(x: Text, y: Interval) -> Double;
        }
        sql_function!(fn floor(x: Double) -> Integer);
        sql_function!(fn greatest<T>(x: T, y: T) -> T);
        sql_function!(fn least<T>(x: T, y: T) -> T);

        let burst: i32 = publish_rate_overrides::table
            .find(uploader)
            .select(publish_rate_overrides::burst)
            .first(conn)
            .optional()?
            .unwrap_or(self.burst);

        // Interval division is poorly defined in general (what is 1 month / 30 days?)
        // However, for the intervals we're dealing with, it is always well
        // defined, so we convert to an f64 of seconds to represent this.
        let tokens_to_add = floor(
            (date_part("epoch", now) - date_part("epoch", last_refill))
                / interval_part("epoch", self.refill_rate()),
        );

        diesel::insert_into(publish_limit_buckets)
            .values((user_id.eq(uploader), tokens.eq(burst), last_refill.eq(now)))
            .on_conflict(user_id)
            .do_update()
            .set((
                tokens.eq(least(burst, greatest(0, tokens - 1) + tokens_to_add)),
                last_refill
                    .eq(last_refill + self.refill_rate().into_sql::<Interval>() * tokens_to_add),
            ))
            .get_result(conn)
    }

    fn refill_rate(&self) -> PgInterval {
        use diesel::dsl::*;
        (self.rate.as_millis() as i64).milliseconds()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::email::Emails;
    use crate::test_util::*;

    #[test]
    fn take_token_with_no_bucket_creates_new_one() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let bucket = rate.take_token(new_user(&conn, "user1")?, now, &conn)?;
        let expected = Bucket {
            user_id: bucket.user_id,
            tokens: 10,
            last_refill: now,
        };
        assert_eq!(expected, bucket);

        let rate = PublishRateLimit {
            rate: Duration::from_millis(50),
            burst: 20,
        };
        let bucket = rate.take_token(new_user(&conn, "user2")?, now, &conn)?;
        let expected = Bucket {
            user_id: bucket.user_id,
            tokens: 20,
            last_refill: now,
        };
        assert_eq!(expected, bucket);
        Ok(())
    }

    #[test]
    fn take_token_with_existing_bucket_modifies_existing_bucket() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 5, now)?.user_id;
        let bucket = rate.take_token(user_id, now, &conn)?;
        let expected = Bucket {
            user_id,
            tokens: 4,
            last_refill: now,
        };
        assert_eq!(expected, bucket);
        Ok(())
    }

    #[test]
    fn take_token_after_delay_refills() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 5, now)?.user_id;
        let refill_time = now + chrono::Duration::seconds(2);
        let bucket = rate.take_token(user_id, refill_time, &conn)?;
        let expected = Bucket {
            user_id,
            tokens: 6,
            last_refill: refill_time,
        };
        assert_eq!(expected, bucket);
        Ok(())
    }

    #[test]
    fn refill_subsecond_rate() -> QueryResult<()> {
        let conn = pg_connection();
        // Subsecond rates have floating point rounding issues, so use a known
        // timestamp that rounds fine
        let now =
            NaiveDateTime::parse_from_str("2019-03-19T21:11:24.620401", "%Y-%m-%dT%H:%M:%S%.f")
                .unwrap();

        let rate = PublishRateLimit {
            rate: Duration::from_millis(100),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 5, now)?.user_id;
        let refill_time = now + chrono::Duration::milliseconds(300);
        let bucket = rate.take_token(user_id, refill_time, &conn)?;
        let expected = Bucket {
            user_id,
            tokens: 7,
            last_refill: refill_time,
        };
        assert_eq!(expected, bucket);
        Ok(())
    }

    #[test]
    fn last_refill_always_advanced_by_multiple_of_rate() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_millis(100),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 5, now)?.user_id;
        let bucket = rate.take_token(user_id, now + chrono::Duration::milliseconds(250), &conn)?;
        let expected_refill_time = now + chrono::Duration::milliseconds(200);
        let expected = Bucket {
            user_id,
            tokens: 6,
            last_refill: expected_refill_time,
        };
        assert_eq!(expected, bucket);
        Ok(())
    }

    #[test]
    fn zero_tokens_returned_when_user_has_no_tokens_left() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 1, now)?.user_id;
        let bucket = rate.take_token(user_id, now, &conn)?;
        let expected = Bucket {
            user_id,
            tokens: 0,
            last_refill: now,
        };
        assert_eq!(expected, bucket);

        let bucket = rate.take_token(user_id, now, &conn)?;
        assert_eq!(expected, bucket);
        Ok(())
    }

    #[test]
    fn a_user_with_no_tokens_gets_a_token_after_exactly_rate() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 0, now)?.user_id;
        let refill_time = now + chrono::Duration::seconds(1);
        let bucket = rate.take_token(user_id, refill_time, &conn)?;
        let expected = Bucket {
            user_id,
            tokens: 1,
            last_refill: refill_time,
        };
        assert_eq!(expected, bucket);

        Ok(())
    }

    #[test]
    fn tokens_never_refill_past_burst() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 8, now)?.user_id;
        let refill_time = now + chrono::Duration::seconds(4);
        let bucket = rate.take_token(user_id, refill_time, &conn)?;
        let expected = Bucket {
            user_id,
            tokens: 10,
            last_refill: refill_time,
        };
        assert_eq!(expected, bucket);

        Ok(())
    }

    #[test]
    fn override_is_used_instead_of_global_burst_if_present() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let user_id = new_user(&conn, "user1")?;
        let other_user_id = new_user(&conn, "user2")?;

        diesel::insert_into(publish_rate_overrides::table)
            .values((
                publish_rate_overrides::user_id.eq(user_id),
                publish_rate_overrides::burst.eq(20),
            ))
            .execute(&conn)?;

        let bucket = rate.take_token(user_id, now, &conn)?;
        let other_bucket = rate.take_token(other_user_id, now, &conn)?;

        assert_eq!(20, bucket.tokens);
        assert_eq!(10, other_bucket.tokens);
        Ok(())
    }

    fn new_user(conn: &PgConnection, gh_login: &str) -> QueryResult<i32> {
        use crate::models::NewUser;

        let user = NewUser {
            gh_login,
            ..NewUser::default()
        }
        .create_or_update(None, &Emails::new_in_memory(), conn)?;
        Ok(user.id)
    }

    fn new_user_bucket(
        conn: &PgConnection,
        tokens: i32,
        now: NaiveDateTime,
    ) -> QueryResult<Bucket> {
        diesel::insert_into(publish_limit_buckets::table)
            .values(Bucket {
                user_id: new_user(conn, "new_user")?,
                tokens,
                last_refill: now,
            })
            .get_result(conn)
    }

    /// Strips ns precision from `Utc::now`. PostgreSQL only has microsecond
    /// precision, but some platforms (notably Linux) provide nanosecond
    /// precision, meaning that round tripping through the database would
    /// change the value.
    fn now() -> NaiveDateTime {
        let now = Utc::now().naive_utc();
        let nanos = now.timestamp_subsec_nanos();
        now - chrono::Duration::nanoseconds(nanos.into())
    }
}
<commit_msg>Auto merge of #3544 - jtgeibel:rate-limiting, r=pietroalbini<commit_after>use chrono::{NaiveDateTime, Utc};
use diesel::data_types::PgInterval;
use diesel::prelude::*;
use std::time::Duration;

use crate::schema::{publish_limit_buckets, publish_rate_overrides};
use crate::util::errors::{AppResult, TooManyRequests};

#[derive(Debug, Clone, Copy)]
pub struct PublishRateLimit {
    pub rate: Duration,
    pub burst: i32,
}

impl Default for PublishRateLimit {
    fn default() -> Self {
        let minutes = dotenv::var("WEB_NEW_PKG_RATE_LIMIT_RATE_MINUTES")
            .unwrap_or_default()
            .parse()
            .ok()
            .unwrap_or(10);
        let burst = dotenv::var("WEB_NEW_PKG_RATE_LIMIT_BURST")
            .unwrap_or_default()
            .parse()
            .ok()
            .unwrap_or(5);
        Self {
            rate: Duration::from_secs(60) * minutes,
            burst,
        }
    }
}

#[derive(Queryable, Insertable, Debug, PartialEq, Clone, Copy)]
#[table_name = "publish_limit_buckets"]
#[allow(dead_code)] // Most fields only read in tests
struct Bucket {
    user_id: i32,
    tokens: i32,
    last_refill: NaiveDateTime,
}

impl PublishRateLimit {
    pub fn check_rate_limit(&self, uploader: i32, conn: &PgConnection) -> AppResult<()> {
        let bucket = self.take_token(uploader, Utc::now().naive_utc(), conn)?;
        if bucket.tokens >= 1 {
            Ok(())
        } else {
            Err(Box::new(TooManyRequests {
                retry_after: bucket.last_refill + chrono::Duration::from_std(self.rate).unwrap(),
            }))
        }
    }

    /// Refill a user's bucket as needed, take a token from it,
    /// and returns the result.
    ///
    /// The number of tokens remaining will always be between 0 and self.burst.
    /// If the number is 0, the request should be rejected, as the user doesn't
    /// have a token to take. Technically a "full" bucket would have
    /// `self.burst + 1` tokens in it, but that value would never be returned
    /// since we only refill buckets when trying to take a token from it.
    fn take_token(
        &self,
        uploader: i32,
        now: NaiveDateTime,
        conn: &PgConnection,
    ) -> QueryResult<Bucket> {
        use self::publish_limit_buckets::dsl::*;
        use diesel::sql_types::{Double, Interval, Text, Timestamp};

        sql_function!(fn date_part(x: Text, y: Timestamp) -> Double);
        sql_function! {
            #[sql_name = "date_part"]
            fn interval_part(x: Text, y: Interval) -> Double;
        }
        sql_function!(fn floor(x: Double) -> Integer);
        sql_function!(fn greatest<T>(x: T, y: T) -> T);
        sql_function!(fn least<T>(x: T, y: T) -> T);

        let burst: i32 = publish_rate_overrides::table
            .find(uploader)
            .select(publish_rate_overrides::burst)
            .first(conn)
            .optional()?
            .unwrap_or(self.burst);

        // Interval division is poorly defined in general (what is 1 month / 30 days?)
        // However, for the intervals we're dealing with, it is always well
        // defined, so we convert to an f64 of seconds to represent this.
        let tokens_to_add = floor(
            (date_part("epoch", now) - date_part("epoch", last_refill))
                / interval_part("epoch", self.refill_rate()),
        );

        diesel::insert_into(publish_limit_buckets)
            .values((user_id.eq(uploader), tokens.eq(burst), last_refill.eq(now)))
            .on_conflict(user_id)
            .do_update()
            .set((
                tokens.eq(least(burst, greatest(0, tokens - 1) + tokens_to_add)),
                last_refill
                    .eq(last_refill + self.refill_rate().into_sql::<Interval>() * tokens_to_add),
            ))
            .get_result(conn)
    }

    fn refill_rate(&self) -> PgInterval {
        use diesel::dsl::*;
        (self.rate.as_millis() as i64).milliseconds()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::email::Emails;
    use crate::test_util::*;

    #[test]
    fn take_token_with_no_bucket_creates_new_one() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let bucket = rate.take_token(new_user(&conn, "user1")?, now, &conn)?;
        let expected = Bucket {
            user_id: bucket.user_id,
            tokens: 10,
            last_refill: now,
        };
        assert_eq!(expected, bucket);

        let rate = PublishRateLimit {
            rate: Duration::from_millis(50),
            burst: 20,
        };
        let bucket = rate.take_token(new_user(&conn, "user2")?, now, &conn)?;
        let expected = Bucket {
            user_id: bucket.user_id,
            tokens: 20,
            last_refill: now,
        };
        assert_eq!(expected, bucket);
        Ok(())
    }

    #[test]
    fn take_token_with_existing_bucket_modifies_existing_bucket() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 5, now)?.user_id;
        let bucket = rate.take_token(user_id, now, &conn)?;
        let expected = Bucket {
            user_id,
            tokens: 4,
            last_refill: now,
        };
        assert_eq!(expected, bucket);
        Ok(())
    }

    #[test]
    fn take_token_after_delay_refills() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 5, now)?.user_id;
        let refill_time = now + chrono::Duration::seconds(2);
        let bucket = rate.take_token(user_id, refill_time, &conn)?;
        let expected = Bucket {
            user_id,
            tokens: 6,
            last_refill: refill_time,
        };
        assert_eq!(expected, bucket);
        Ok(())
    }

    #[test]
    fn refill_subsecond_rate() -> QueryResult<()> {
        let conn = pg_connection();
        // Subsecond rates have floating point rounding issues, so use a known
        // timestamp that rounds fine
        let now =
            NaiveDateTime::parse_from_str("2019-03-19T21:11:24.620401", "%Y-%m-%dT%H:%M:%S%.f")
                .unwrap();

        let rate = PublishRateLimit {
            rate: Duration::from_millis(100),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 5, now)?.user_id;
        let refill_time = now + chrono::Duration::milliseconds(300);
        let bucket = rate.take_token(user_id, refill_time, &conn)?;
        let expected = Bucket {
            user_id,
            tokens: 7,
            last_refill: refill_time,
        };
        assert_eq!(expected, bucket);
        Ok(())
    }

    #[test]
    fn last_refill_always_advanced_by_multiple_of_rate() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_millis(100),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 5, now)?.user_id;
        let bucket = rate.take_token(user_id, now + chrono::Duration::milliseconds(250), &conn)?;
        let expected_refill_time = now + chrono::Duration::milliseconds(200);
        let expected = Bucket {
            user_id,
            tokens: 6,
            last_refill: expected_refill_time,
        };
        assert_eq!(expected, bucket);
        Ok(())
    }

    #[test]
    fn zero_tokens_returned_when_user_has_no_tokens_left() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 1, now)?.user_id;
        let bucket = rate.take_token(user_id, now, &conn)?;
        let expected = Bucket {
            user_id,
            tokens: 0,
            last_refill: now,
        };
        assert_eq!(expected, bucket);

        let bucket = rate.take_token(user_id, now, &conn)?;
        assert_eq!(expected, bucket);
        Ok(())
    }

    #[test]
    fn a_user_with_no_tokens_gets_a_token_after_exactly_rate() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 0, now)?.user_id;
        let refill_time = now + chrono::Duration::seconds(1);
        let bucket = rate.take_token(user_id, refill_time, &conn)?;
        let expected = Bucket {
            user_id,
            tokens: 1,
            last_refill: refill_time,
        };
        assert_eq!(expected, bucket);

        Ok(())
    }

    #[test]
    fn tokens_never_refill_past_burst() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let user_id = new_user_bucket(&conn, 8, now)?.user_id;
        let refill_time = now + chrono::Duration::seconds(4);
        let bucket = rate.take_token(user_id, refill_time, &conn)?;
        let expected = Bucket {
            user_id,
            tokens: 10,
            last_refill: refill_time,
        };
        assert_eq!(expected, bucket);

        Ok(())
    }

    #[test]
    fn override_is_used_instead_of_global_burst_if_present() -> QueryResult<()> {
        let conn = pg_connection();
        let now = now();

        let rate = PublishRateLimit {
            rate: Duration::from_secs(1),
            burst: 10,
        };
        let user_id = new_user(&conn, "user1")?;
        let other_user_id = new_user(&conn, "user2")?;

        diesel::insert_into(publish_rate_overrides::table)
            .values((
                publish_rate_overrides::user_id.eq(user_id),
                publish_rate_overrides::burst.eq(20),
            ))
            .execute(&conn)?;

        let bucket = rate.take_token(user_id, now, &conn)?;
        let other_bucket = rate.take_token(other_user_id, now, &conn)?;

        assert_eq!(20, bucket.tokens);
        assert_eq!(10, other_bucket.tokens);
        Ok(())
    }

    fn new_user(conn: &PgConnection, gh_login: &str) -> QueryResult<i32> {
        use crate::models::NewUser;

        let user = NewUser {
            gh_login,
            ..NewUser::default()
        }
        .create_or_update(None, &Emails::new_in_memory(), conn)?;
        Ok(user.id)
    }

    fn new_user_bucket(
        conn: &PgConnection,
        tokens: i32,
        now: NaiveDateTime,
    ) -> QueryResult<Bucket> {
        diesel::insert_into(publish_limit_buckets::table)
            .values(Bucket {
                user_id: new_user(conn, "new_user")?,
                tokens,
                last_refill: now,
            })
            .get_result(conn)
    }

    /// Strips ns precision from `Utc::now`. PostgreSQL only has microsecond
    /// precision, but some platforms (notably Linux) provide nanosecond
    /// precision, meaning that round tripping through the database would
    /// change the value.
    fn now() -> NaiveDateTime {
        let now = Utc::now().naive_utc();
        let nanos = now.timestamp_subsec_nanos();
        now - chrono::Duration::nanoseconds(nanos.into())
    }
}
<|endoftext|>


<commit_before>//! This module contains the `EvalContext` methods for executing a single step of the interpreter.
//!
//! The main entry point is the `step` method.

use super::{
    CachedMir,
    ConstantId,
    EvalContext,
    Lvalue,
    ConstantKind,
    StackPopCleanup,
    Constant,
};
use error::EvalResult;
use rustc::mir::repr as mir;
use rustc::ty::{subst, self};
use rustc::hir::def_id::DefId;
use rustc::hir;
use rustc::mir::visit::{Visitor, LvalueContext};
use syntax::codemap::Span;
use std::rc::Rc;

impl<'a, 'tcx> EvalContext<'a, 'tcx> {
    /// Returns true as long as there are more things to do.
    pub fn step(&mut self) -> EvalResult<'tcx, bool> {
        if self.stack.is_empty() {
            return Ok(false);
        }

        let block = self.frame().block;
        let stmt_id = self.frame().stmt;
        let mir = self.mir();
        let basic_block = &mir.basic_blocks()[block];

        if let Some(stmt) = basic_block.statements.get(stmt_id) {
            let mut new = Ok(0);
            ConstantExtractor {
                span: stmt.source_info.span,
                substs: self.substs(),
                def_id: self.frame().def_id,
                ecx: self,
                mir: &mir,
                new_constants: &mut new,
            }.visit_statement(block, stmt, mir::Location {
                block: block,
                statement_index: stmt_id,
            });
            if new? == 0 {
                self.statement(stmt)?;
            }
            // if ConstantExtractor added new frames, we don't execute anything here
            // but await the next call to step
            return Ok(true);
        }

        let terminator = basic_block.terminator();
        let mut new = Ok(0);
        ConstantExtractor {
            span: terminator.source_info.span,
            substs: self.substs(),
            def_id: self.frame().def_id,
            ecx: self,
            mir: &mir,
            new_constants: &mut new,
        }.visit_terminator(block, terminator, mir::Location {
            block: block,
            statement_index: stmt_id,
        });
        if new? == 0 {
            self.terminator(terminator)?;
        }
        // if ConstantExtractor added new frames, we don't execute anything here
        // but await the next call to step
        Ok(true)
    }

    fn statement(&mut self, stmt: &mir::Statement<'tcx>) -> EvalResult<'tcx, ()> {
        trace!("{:?}", stmt);

        use rustc::mir::repr::StatementKind::*;
        match stmt.kind {
            Assign(ref lvalue, ref rvalue) => self.eval_rvalue_into_lvalue(rvalue, lvalue)?,
            SetDiscriminant { .. } => unimplemented!(),

            // Miri can safely ignore these. Only translation needs them.
            StorageLive(_) | StorageDead(_) => {}

            // Defined to do nothing. These are added by optimization passes, to avoid changing the
            // size of MIR constantly.
            Nop => {}
        }

        self.frame_mut().stmt += 1;
        Ok(())
    }

    fn terminator(&mut self, terminator: &mir::Terminator<'tcx>) -> EvalResult<'tcx, ()> {
        trace!("{:?}", terminator.kind);
        self.eval_terminator(terminator)?;
        if !self.stack.is_empty() {
            trace!("// {:?}", self.frame().block);
        }
        Ok(())
    }
}

// WARNING: make sure that any methods implemented on this type don't ever access ecx.stack
// this includes any method that might access the stack
// basically don't call anything other than `load_mir`, `alloc_ptr`, `push_stack_frame`
// The reason for this is, that `push_stack_frame` modifies the stack out of obvious reasons
struct ConstantExtractor<'a, 'b: 'a, 'tcx: 'b> {
    span: Span,
    ecx: &'a mut EvalContext<'b, 'tcx>,
    mir: &'a mir::Mir<'tcx>,
    def_id: DefId,
    substs: &'tcx subst::Substs<'tcx>,
    new_constants: &'a mut EvalResult<'tcx, u64>,
}

impl<'a, 'b, 'tcx> ConstantExtractor<'a, 'b, 'tcx> {
    fn global_item(&mut self, def_id: DefId, substs: &'tcx subst::Substs<'tcx>, span: Span, immutable: bool) {
        let cid = ConstantId {
            def_id: def_id,
            substs: substs,
            kind: ConstantKind::Global,
        };
        if self.ecx.statics.contains_key(&cid) {
            return;
        }
        self.try(|this| {
            let mir = this.ecx.load_mir(def_id)?;
            this.ecx.statics.insert(cid.clone(), Constant::uninitialized(mir.return_ty));
            let cleanup = if immutable && !mir.return_ty.type_contents(this.ecx.tcx).interior_unsafe() {
                StackPopCleanup::Freeze
            } else {
                StackPopCleanup::None
            };
            this.ecx.push_stack_frame(def_id, span, mir, substs, Lvalue::Static(cid.clone()), cleanup)
        });
    }
    fn try<F: FnOnce(&mut Self) -> EvalResult<'tcx, ()>>(&mut self, f: F) {
        if let Ok(ref mut n) = *self.new_constants {
            *n += 1;
        } else {
            return;
        }
        if let Err(e) = f(self) {
            *self.new_constants = Err(e);
        }
    }
}

impl<'a, 'b, 'tcx> Visitor<'tcx> for ConstantExtractor<'a, 'b, 'tcx> {
    fn visit_constant(&mut self, constant: &mir::Constant<'tcx>, location: mir::Location) {
        self.super_constant(constant, location);
        match constant.literal {
            // already computed by rustc
            mir::Literal::Value { .. } => {}
            mir::Literal::Item { def_id, substs } => {
                if let ty::TyFnDef(..) = constant.ty.sty {
                    // No need to do anything here,
                    // because the type is the actual function, not the signature of the function.
                    // Thus we can simply create a zero sized allocation in `evaluate_operand`
                } else {
                    self.global_item(def_id, substs, constant.span, true);
                }
            },
            mir::Literal::Promoted { index } => {
                let mir = self.mir.promoted[index].clone();
                let cid = ConstantId {
                    def_id: self.def_id,
                    substs: self.substs,
                    kind: ConstantKind::Promoted(index),
                };
                if self.ecx.statics.contains_key(&cid) {
                    return;
                }
                self.try(|this| {
                    let mir = CachedMir::Owned(Rc::new(mir));
                    let ty = this.ecx.monomorphize(mir.return_ty, this.substs);
                    this.ecx.statics.insert(cid.clone(), Constant::uninitialized(ty));
                    this.ecx.push_stack_frame(this.def_id,
                                              constant.span,
                                              mir,
                                              this.substs,
                                              Lvalue::Static(cid.clone()),
                                              StackPopCleanup::Freeze)
                });
            }
        }
    }

    fn visit_lvalue(
        &mut self,
        lvalue: &mir::Lvalue<'tcx>,
        context: LvalueContext<'tcx>,
        location: mir::Location
    ) {
        self.super_lvalue(lvalue, context, location);
        if let mir::Lvalue::Static(def_id) = *lvalue {
            let substs = subst::Substs::empty(self.ecx.tcx);
            let span = self.span;
            if let Some(node_item) = self.ecx.tcx.map.get_if_local(def_id) {
                if let hir::map::Node::NodeItem(&hir::Item { ref node, .. }) = node_item {
                    if let hir::ItemStatic(_, m, _) = *node {
                        self.global_item(def_id, substs, span, m == hir::MutImmutable);
                        return;
                    } else {
                        bug!("static def id doesn't point to static");
                    }
                } else {
                    bug!("static def id doesn't point to item");
                }
            } else {
                let def = self.ecx.tcx.sess.cstore.describe_def(def_id).expect("static not found");
                if let hir::def::Def::Static(_, mutable) = def {
                    self.global_item(def_id, substs, span, !mutable);
                } else {
                    bug!("static found but isn't a static: {:?}", def);
                }
            }
        }
    }
}
<commit_msg>constant ids are `Copy` now<commit_after>//! This module contains the `EvalContext` methods for executing a single step of the interpreter.
//!
//! The main entry point is the `step` method.

use super::{
    CachedMir,
    ConstantId,
    EvalContext,
    Lvalue,
    ConstantKind,
    StackPopCleanup,
    Constant,
};
use error::EvalResult;
use rustc::mir::repr as mir;
use rustc::ty::{subst, self};
use rustc::hir::def_id::DefId;
use rustc::hir;
use rustc::mir::visit::{Visitor, LvalueContext};
use syntax::codemap::Span;
use std::rc::Rc;

impl<'a, 'tcx> EvalContext<'a, 'tcx> {
    /// Returns true as long as there are more things to do.
    pub fn step(&mut self) -> EvalResult<'tcx, bool> {
        if self.stack.is_empty() {
            return Ok(false);
        }

        let block = self.frame().block;
        let stmt_id = self.frame().stmt;
        let mir = self.mir();
        let basic_block = &mir.basic_blocks()[block];

        if let Some(stmt) = basic_block.statements.get(stmt_id) {
            let mut new = Ok(0);
            ConstantExtractor {
                span: stmt.source_info.span,
                substs: self.substs(),
                def_id: self.frame().def_id,
                ecx: self,
                mir: &mir,
                new_constants: &mut new,
            }.visit_statement(block, stmt, mir::Location {
                block: block,
                statement_index: stmt_id,
            });
            if new? == 0 {
                self.statement(stmt)?;
            }
            // if ConstantExtractor added new frames, we don't execute anything here
            // but await the next call to step
            return Ok(true);
        }

        let terminator = basic_block.terminator();
        let mut new = Ok(0);
        ConstantExtractor {
            span: terminator.source_info.span,
            substs: self.substs(),
            def_id: self.frame().def_id,
            ecx: self,
            mir: &mir,
            new_constants: &mut new,
        }.visit_terminator(block, terminator, mir::Location {
            block: block,
            statement_index: stmt_id,
        });
        if new? == 0 {
            self.terminator(terminator)?;
        }
        // if ConstantExtractor added new frames, we don't execute anything here
        // but await the next call to step
        Ok(true)
    }

    fn statement(&mut self, stmt: &mir::Statement<'tcx>) -> EvalResult<'tcx, ()> {
        trace!("{:?}", stmt);

        use rustc::mir::repr::StatementKind::*;
        match stmt.kind {
            Assign(ref lvalue, ref rvalue) => self.eval_rvalue_into_lvalue(rvalue, lvalue)?,
            SetDiscriminant { .. } => unimplemented!(),

            // Miri can safely ignore these. Only translation needs them.
            StorageLive(_) | StorageDead(_) => {}

            // Defined to do nothing. These are added by optimization passes, to avoid changing the
            // size of MIR constantly.
            Nop => {}
        }

        self.frame_mut().stmt += 1;
        Ok(())
    }

    fn terminator(&mut self, terminator: &mir::Terminator<'tcx>) -> EvalResult<'tcx, ()> {
        trace!("{:?}", terminator.kind);
        self.eval_terminator(terminator)?;
        if !self.stack.is_empty() {
            trace!("// {:?}", self.frame().block);
        }
        Ok(())
    }
}

// WARNING: make sure that any methods implemented on this type don't ever access ecx.stack
// this includes any method that might access the stack
// basically don't call anything other than `load_mir`, `alloc_ptr`, `push_stack_frame`
// The reason for this is, that `push_stack_frame` modifies the stack out of obvious reasons
struct ConstantExtractor<'a, 'b: 'a, 'tcx: 'b> {
    span: Span,
    ecx: &'a mut EvalContext<'b, 'tcx>,
    mir: &'a mir::Mir<'tcx>,
    def_id: DefId,
    substs: &'tcx subst::Substs<'tcx>,
    new_constants: &'a mut EvalResult<'tcx, u64>,
}

impl<'a, 'b, 'tcx> ConstantExtractor<'a, 'b, 'tcx> {
    fn global_item(&mut self, def_id: DefId, substs: &'tcx subst::Substs<'tcx>, span: Span, immutable: bool) {
        let cid = ConstantId {
            def_id: def_id,
            substs: substs,
            kind: ConstantKind::Global,
        };
        if self.ecx.statics.contains_key(&cid) {
            return;
        }
        self.try(|this| {
            let mir = this.ecx.load_mir(def_id)?;
            this.ecx.statics.insert(cid, Constant::uninitialized(mir.return_ty));
            let cleanup = if immutable && !mir.return_ty.type_contents(this.ecx.tcx).interior_unsafe() {
                StackPopCleanup::Freeze
            } else {
                StackPopCleanup::None
            };
            this.ecx.push_stack_frame(def_id, span, mir, substs, Lvalue::Static(cid), cleanup)
        });
    }
    fn try<F: FnOnce(&mut Self) -> EvalResult<'tcx, ()>>(&mut self, f: F) {
        if let Ok(ref mut n) = *self.new_constants {
            *n += 1;
        } else {
            return;
        }
        if let Err(e) = f(self) {
            *self.new_constants = Err(e);
        }
    }
}

impl<'a, 'b, 'tcx> Visitor<'tcx> for ConstantExtractor<'a, 'b, 'tcx> {
    fn visit_constant(&mut self, constant: &mir::Constant<'tcx>, location: mir::Location) {
        self.super_constant(constant, location);
        match constant.literal {
            // already computed by rustc
            mir::Literal::Value { .. } => {}
            mir::Literal::Item { def_id, substs } => {
                if let ty::TyFnDef(..) = constant.ty.sty {
                    // No need to do anything here,
                    // because the type is the actual function, not the signature of the function.
                    // Thus we can simply create a zero sized allocation in `evaluate_operand`
                } else {
                    self.global_item(def_id, substs, constant.span, true);
                }
            },
            mir::Literal::Promoted { index } => {
                let mir = self.mir.promoted[index].clone();
                let cid = ConstantId {
                    def_id: self.def_id,
                    substs: self.substs,
                    kind: ConstantKind::Promoted(index),
                };
                if self.ecx.statics.contains_key(&cid) {
                    return;
                }
                self.try(|this| {
                    let mir = CachedMir::Owned(Rc::new(mir));
                    let ty = this.ecx.monomorphize(mir.return_ty, this.substs);
                    this.ecx.statics.insert(cid, Constant::uninitialized(ty));
                    this.ecx.push_stack_frame(this.def_id,
                                              constant.span,
                                              mir,
                                              this.substs,
                                              Lvalue::Static(cid),
                                              StackPopCleanup::Freeze)
                });
            }
        }
    }

    fn visit_lvalue(
        &mut self,
        lvalue: &mir::Lvalue<'tcx>,
        context: LvalueContext<'tcx>,
        location: mir::Location
    ) {
        self.super_lvalue(lvalue, context, location);
        if let mir::Lvalue::Static(def_id) = *lvalue {
            let substs = subst::Substs::empty(self.ecx.tcx);
            let span = self.span;
            if let Some(node_item) = self.ecx.tcx.map.get_if_local(def_id) {
                if let hir::map::Node::NodeItem(&hir::Item { ref node, .. }) = node_item {
                    if let hir::ItemStatic(_, m, _) = *node {
                        self.global_item(def_id, substs, span, m == hir::MutImmutable);
                        return;
                    } else {
                        bug!("static def id doesn't point to static");
                    }
                } else {
                    bug!("static def id doesn't point to item");
                }
            } else {
                let def = self.ecx.tcx.sess.cstore.describe_def(def_id).expect("static not found");
                if let hir::def::Def::Static(_, mutable) = def {
                    self.global_item(def_id, substs, span, !mutable);
                } else {
                    bug!("static found but isn't a static: {:?}", def);
                }
            }
        }
    }
}
<|endoftext|>