<commit_before>use rustc_data_structures::fx::FxHashSet;
use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use pathdiff::diff_paths;

use rustc_hir::def_id::CrateNum;
use rustc_middle::middle::cstore::LibSource;

pub struct RPathConfig<'a> {
    pub used_crates: &'a [(CrateNum, LibSource)],
    pub out_filename: PathBuf,
    pub is_like_osx: bool,
    pub has_rpath: bool,
    pub linker_is_gnu: bool,
    pub get_install_prefix_lib_path: &'a mut dyn FnMut() -> PathBuf,
}

pub fn get_rpath_flags(config: &mut RPathConfig<'_>) -> Vec<String> {
    // No rpath on windows
    if !config.has_rpath {
        return Vec::new();
    }

    debug!("preparing the RPATH!");

    let libs = config.used_crates.clone();
    let libs = libs.iter().filter_map(|&(_, ref l)| l.option()).collect::<Vec<_>>();
    let rpaths = get_rpaths(config, &libs);
    let mut flags = rpaths_to_flags(&rpaths);

    // Use DT_RUNPATH instead of DT_RPATH if available
    if config.linker_is_gnu {
        flags.push("-Wl,--enable-new-dtags".to_owned());
    }

    flags
}

fn rpaths_to_flags(rpaths: &[String]) -> Vec<String> {
    let mut ret = Vec::with_capacity(rpaths.len()); // the minimum needed capacity

    for rpath in rpaths {
        if rpath.contains(',') {
            ret.push("-Wl,-rpath".into());
            ret.push("-Xlinker".into());
            ret.push(rpath.clone());
        } else {
            ret.push(format!("-Wl,-rpath,{}", &(*rpath)));
        }
    }

    ret
}

fn get_rpaths(config: &mut RPathConfig<'_>, libs: &[PathBuf]) -> Vec<String> {
    debug!("output: {:?}", config.out_filename.display());
    debug!("libs:");
    for libpath in libs {
        debug!("    {:?}", libpath.display());
    }

    // Use relative paths to the libraries. Binaries can be moved
    // as long as they maintain the relative relationship to the
    // crates they depend on.
    let rel_rpaths = get_rpaths_relative_to_output(config, libs);

    // And a final backup rpath to the global library location.
    let fallback_rpaths = vec![get_install_prefix_rpath(config)];

    fn log_rpaths(desc: &str, rpaths: &[String]) {
        debug!("{} rpaths:", desc);
        for rpath in rpaths {
            debug!("    {}", *rpath);
        }
    }

    log_rpaths("relative", &rel_rpaths);
    log_rpaths("fallback", &fallback_rpaths);

    let mut rpaths = rel_rpaths;
    rpaths.extend_from_slice(&fallback_rpaths);

    // Remove duplicates
    minimize_rpaths(&rpaths)
}

fn get_rpaths_relative_to_output(config: &mut RPathConfig<'_>, libs: &[PathBuf]) -> Vec<String> {
    libs.iter().map(|a| get_rpath_relative_to_output(config, a)).collect()
}

fn get_rpath_relative_to_output(config: &mut RPathConfig<'_>, lib: &Path) -> String {
    // Mac doesn't appear to support $ORIGIN
    let prefix = if config.is_like_osx { "@loader_path" } else { "$ORIGIN" };

    let cwd = env::current_dir().unwrap();
    let mut lib = fs::canonicalize(&cwd.join(lib)).unwrap_or_else(|_| cwd.join(lib));
    lib.pop(); // strip filename
    let mut output = cwd.join(&config.out_filename);
    output.pop(); // strip filename
    let output = fs::canonicalize(&output).unwrap_or(output);
    let relative = path_relative_from(&lib, &output)
        .unwrap_or_else(|| panic!("couldn't create relative path from {:?} to {:?}", output, lib));
    // FIXME (#9639): This needs to handle non-utf8 paths
    format!("{}/{}", prefix, relative.to_str().expect("non-utf8 component in path"))
}

// This routine is adapted from the *old* Path's `path_relative_from`
// function, which works differently from the new `relative_from` function.
// In particular, this handles the case on unix where both paths are
// absolute but with only the root as the common directory.
fn path_relative_from(path: &Path, base: &Path) -> Option<PathBuf> {
    diff_paths(path, base)
}

fn get_install_prefix_rpath(config: &mut RPathConfig<'_>) -> String {
    let path = (config.get_install_prefix_lib_path)();
    let path = env::current_dir().unwrap().join(&path);
    // FIXME (#9639): This needs to handle non-utf8 paths
    path.to_str().expect("non-utf8 component in rpath").to_owned()
}

fn minimize_rpaths(rpaths: &[String]) -> Vec<String> {
    let mut set = FxHashSet::default();
    let mut minimized = Vec::new();
    for rpath in rpaths {
        if set.insert(rpath) {
            minimized.push(rpath.clone());
        }
    }
    minimized
}

#[cfg(all(unix, test))]
mod tests;
<commit_msg>Run fmt<commit_after>use pathdiff::diff_paths;
use rustc_data_structures::fx::FxHashSet;
use std::env;
use std::fs;
use std::path::{Path, PathBuf};

use rustc_hir::def_id::CrateNum;
use rustc_middle::middle::cstore::LibSource;

pub struct RPathConfig<'a> {
    pub used_crates: &'a [(CrateNum, LibSource)],
    pub out_filename: PathBuf,
    pub is_like_osx: bool,
    pub has_rpath: bool,
    pub linker_is_gnu: bool,
    pub get_install_prefix_lib_path: &'a mut dyn FnMut() -> PathBuf,
}

pub fn get_rpath_flags(config: &mut RPathConfig<'_>) -> Vec<String> {
    // No rpath on windows
    if !config.has_rpath {
        return Vec::new();
    }

    debug!("preparing the RPATH!");

    let libs = config.used_crates.clone();
    let libs = libs.iter().filter_map(|&(_, ref l)| l.option()).collect::<Vec<_>>();
    let rpaths = get_rpaths(config, &libs);
    let mut flags = rpaths_to_flags(&rpaths);

    // Use DT_RUNPATH instead of DT_RPATH if available
    if config.linker_is_gnu {
        flags.push("-Wl,--enable-new-dtags".to_owned());
    }

    flags
}

fn rpaths_to_flags(rpaths: &[String]) -> Vec<String> {
    let mut ret = Vec::with_capacity(rpaths.len()); // the minimum needed capacity

    for rpath in rpaths {
        if rpath.contains(',') {
            ret.push("-Wl,-rpath".into());
            ret.push("-Xlinker".into());
            ret.push(rpath.clone());
        } else {
            ret.push(format!("-Wl,-rpath,{}", &(*rpath)));
        }
    }

    ret
}

fn get_rpaths(config: &mut RPathConfig<'_>, libs: &[PathBuf]) -> Vec<String> {
    debug!("output: {:?}", config.out_filename.display());
    debug!("libs:");
    for libpath in libs {
        debug!("    {:?}", libpath.display());
    }

    // Use relative paths to the libraries. Binaries can be moved
    // as long as they maintain the relative relationship to the
    // crates they depend on.
    let rel_rpaths = get_rpaths_relative_to_output(config, libs);

    // And a final backup rpath to the global library location.
    let fallback_rpaths = vec![get_install_prefix_rpath(config)];

    fn log_rpaths(desc: &str, rpaths: &[String]) {
        debug!("{} rpaths:", desc);
        for rpath in rpaths {
            debug!("    {}", *rpath);
        }
    }

    log_rpaths("relative", &rel_rpaths);
    log_rpaths("fallback", &fallback_rpaths);

    let mut rpaths = rel_rpaths;
    rpaths.extend_from_slice(&fallback_rpaths);

    // Remove duplicates
    minimize_rpaths(&rpaths)
}

fn get_rpaths_relative_to_output(config: &mut RPathConfig<'_>, libs: &[PathBuf]) -> Vec<String> {
    libs.iter().map(|a| get_rpath_relative_to_output(config, a)).collect()
}

fn get_rpath_relative_to_output(config: &mut RPathConfig<'_>, lib: &Path) -> String {
    // Mac doesn't appear to support $ORIGIN
    let prefix = if config.is_like_osx { "@loader_path" } else { "$ORIGIN" };

    let cwd = env::current_dir().unwrap();
    let mut lib = fs::canonicalize(&cwd.join(lib)).unwrap_or_else(|_| cwd.join(lib));
    lib.pop(); // strip filename
    let mut output = cwd.join(&config.out_filename);
    output.pop(); // strip filename
    let output = fs::canonicalize(&output).unwrap_or(output);
    let relative = path_relative_from(&lib, &output)
        .unwrap_or_else(|| panic!("couldn't create relative path from {:?} to {:?}", output, lib));
    // FIXME (#9639): This needs to handle non-utf8 paths
    format!("{}/{}", prefix, relative.to_str().expect("non-utf8 component in path"))
}

// This routine is adapted from the *old* Path's `path_relative_from`
// function, which works differently from the new `relative_from` function.
// In particular, this handles the case on unix where both paths are
// absolute but with only the root as the common directory.
fn path_relative_from(path: &Path, base: &Path) -> Option<PathBuf> {
    diff_paths(path, base)
}

fn get_install_prefix_rpath(config: &mut RPathConfig<'_>) -> String {
    let path = (config.get_install_prefix_lib_path)();
    let path = env::current_dir().unwrap().join(&path);
    // FIXME (#9639): This needs to handle non-utf8 paths
    path.to_str().expect("non-utf8 component in rpath").to_owned()
}

fn minimize_rpaths(rpaths: &[String]) -> Vec<String> {
    let mut set = FxHashSet::default();
    let mut minimized = Vec::new();
    for rpath in rpaths {
        if set.insert(rpath) {
            minimized.push(rpath.clone());
        }
    }
    minimized
}

#[cfg(all(unix, test))]
mod tests;
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use syntax::{ast, attr};
use llvm::LLVMRustHasFeature;
use rustc::session::Session;
use rustc_trans::back::write::create_target_machine;
use syntax::parse::token::InternedString;
use syntax::parse::token::intern_and_get_ident as intern;
use libc::c_char;

/// Add `target_feature = "..."` cfgs for a variety of platform
/// specific features (SSE, NEON etc.).
///
/// This is performed by checking whether a whitelisted set of
/// features is available on the target machine, by querying LLVM.
pub fn add_configuration(cfg: &mut ast::CrateConfig, sess: &Session) {
    let target_machine = create_target_machine(sess);

    let arm_whitelist = [
        "neon\0",
        "vfp\0",
    ];

    let x86_whitelist = [
        "avx\0",
        "avx2\0",
        "sse\0",
        "sse2\0",
        "sse3\0",
        "sse4.1\0",
        "sse4.2\0",
        "ssse3\0",
    ];

    let whitelist = match &*sess.target.target.arch {
        "arm" => &arm_whitelist[..],
        "x86" | "x86_64" => &x86_whitelist[..],
        _ => &[][..],
    };

    let tf = InternedString::new("target_feature");
    for feat in whitelist {
        if unsafe { LLVMRustHasFeature(target_machine, feat.as_ptr() as *const c_char) } {
            cfg.push(attr::mk_name_value_item_str(tf.clone(), intern(feat)))
        }
    }
}
<commit_msg>Do not intern NUL terminators<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use syntax::{ast, attr};
use llvm::LLVMRustHasFeature;
use rustc::session::Session;
use rustc_trans::back::write::create_target_machine;
use syntax::parse::token::InternedString;
use syntax::parse::token::intern_and_get_ident as intern;
use libc::c_char;

/// Add `target_feature = "..."` cfgs for a variety of platform
/// specific features (SSE, NEON etc.).
///
/// This is performed by checking whether a whitelisted set of
/// features is available on the target machine, by querying LLVM.
pub fn add_configuration(cfg: &mut ast::CrateConfig, sess: &Session) {
    let target_machine = create_target_machine(sess);

    let arm_whitelist = [
        "neon\0",
        "vfp\0",
    ];

    let x86_whitelist = [
        "avx\0",
        "avx2\0",
        "sse\0",
        "sse2\0",
        "sse3\0",
        "sse4.1\0",
        "sse4.2\0",
        "ssse3\0",
    ];

    let whitelist = match &*sess.target.target.arch {
        "arm" => &arm_whitelist[..],
        "x86" | "x86_64" => &x86_whitelist[..],
        _ => &[][..],
    };

    let tf = InternedString::new("target_feature");
    for feat in whitelist {
        if unsafe { LLVMRustHasFeature(target_machine, feat.as_ptr() as *const c_char) } {
            cfg.push(attr::mk_name_value_item_str(tf.clone(), intern(&feat[..feat.len()-1])))
        }
    }
}
<|endoftext|>


<commit_before>
use pair::Pair;
use error::{NcclError, ErrorKind};
use token::{Token, TokenKind};

use std::io::Write;

pub struct Parser {
    current: usize,
    path: Vec<String>,
    tokens: Vec<Token>,
    pair: Pair,
    line: u64,
}

// nccl = value+
// value = name schema? nl key*
// key = indent (name nl | indent value)
// name = not colon
// nl = newline
// schema = ":" name

impl Parser {
    pub fn new(tokens: Vec<Token>) -> Self {
        Parser {
            current: 0,
            path: vec![],
            tokens: tokens,
            pair: Pair::new("__top_level__"),
            line: 1,
        }
    }

    pub fn parse(&mut self) -> Result<Pair, NcclError> {
        while !self.is_at_end() {
            self.value()?;
            self.advance();
        }

        Ok(self.pair.clone())
    }

    fn add(&mut self, value: String) {
        //self.pair.add_vec(self.path, value);
    }

    fn value(&mut self) -> Result<(), NcclError> {
        if !self.is_at_end() {
            self.name()?;

            if self.peek().kind == TokenKind::Colon {
                self.schema()?;
            }

            self.newline()?;

            while self.peek().kind != TokenKind::Indent {
                self.key()?;
            }

            Ok(())
        } else {
            Err(NcclError::new(ErrorKind::ParseError, "Expected value, found EOF", self.line))
        }
    }

    fn key(&mut self) -> Result<(), NcclError> {
        if !self.is_at_end() {
            self.indent()?;

            if self.peek().kind == TokenKind::Value {
                self.name()?;
                self.newline()?;
            } else if self.peek().kind == TokenKind::Indent {
                self.indent()?;
                self.value()?;
            }

            Ok(())
        } else {
            Err(NcclError::new(ErrorKind::ParseError, "Expected key, found EOF", self.line))
        }
    }

    fn schema(&mut self) -> Result<(), NcclError> {
        if !self.is_at_end() {
            self.colon()?;
            self.name()?;
            Ok(())
        } else {
            Err(NcclError::new(ErrorKind::ParseError, "Expected schema, found EOF", self.line))
        }
    }

    fn name(&mut self) -> Result<(), NcclError> {
        if !self.is_at_end() {
            self.advance();
            Ok(())
        } else {
            Err(NcclError::new(ErrorKind::ParseError, "Expected name, found EOF", self.line))
        }
    }

    fn indent(&mut self) -> Result<(), NcclError> {
        if !self.is_at_end() {
            self.advance();
            Ok(())
        } else {
            Err(NcclError::new(ErrorKind::ParseError, "Expected indent, found EOF", self.line))
        }
    }

    fn newline(&mut self) -> Result<(), NcclError> {
        if !self.is_at_end() {
            self.line += 1;
            self.advance();
            Ok(())
        } else {
            Err(NcclError::new(ErrorKind::ParseError, "Expected newline, found EOF", self.line))
        }
    }

    fn colon(&mut self) -> Result<(), NcclError> {
        if !self.is_at_end() {
            self.advance();
            Ok(())
        } else {
            Err(NcclError::new(ErrorKind::ParseError, "Expected colon, found EOF", self.line))
        }
    }

    fn matches(&mut self, kind: TokenKind) -> bool {
        if self.check(kind) {
            self.advance();
            true
        } else {
            false
        }
    }

    fn check(&mut self, kind: TokenKind) -> bool {
        if self.is_at_end() {
            false
        } else {
            self.peek().kind == kind
        }
    }

    fn advance(&mut self) -> &Token {
        if !self.is_at_end() {
            self.current += 1;
        }

        self.previous()
    }

    fn is_at_end(&mut self) -> bool {
        self.peek().kind == TokenKind::EOF
    }

    fn peek(&mut self) -> &Token {
        &self.tokens[self.current]
    }

    fn previous(&mut self) -> &Token {
        &self.tokens[self.current - 1]
    }

    fn consume(&mut self, kind: TokenKind, message: &str) -> Result<&Token, NcclError> {
        if self.check(kind) {
            Ok(self.advance())
        } else {
            Err(NcclError::new(ErrorKind::ParseError, message, self.line))
        }
    }
}

<commit_msg>finish up parse<commit_after>
use pair::Pair;
use error::NcclError;
use token::{Token, TokenKind};

#[derive(Debug)]
pub struct Parser {
    current: usize,
    path: Vec<String>,
    indent: usize,
    tokens: Vec<Token>,
    pair: Pair,
    line: u64,
}

impl Parser {
    pub fn new(tokens: Vec<Token>) -> Self {
        Parser {
            current: 0,
            path: vec![],
            indent: 0,
            tokens: tokens,
            pair: Pair::new("__top_level__"),
            line: 1,
        }
    }

    // faked you out with that Scanner, didn't I?
    // you thought this was going to be recursive descent. YOU WERE WRONG!
    pub fn parse(mut self) -> Result<Pair, Vec<NcclError>> {
        let mut i = 0;

        while i < self.tokens.len() {
            match self.tokens[i].kind {
                TokenKind::Value => { // add to path respective of self.index
                    if self.indent <= self.path.len() {
                        let mut new = self.path[0..self.indent].to_owned();
                        new.push(self.tokens[i].lexeme.clone());
                        self.path = new;
                    } else {
                        self.path.push(self.tokens[i].lexeme.clone());
                    }

                    self.pair.add_slice(&self.path);

                    if i + 2 <= self.tokens.len() && self.tokens[i + 2].kind == TokenKind::Value {
                        self.path.clear();
                        self.indent = 0;
                    }
                },

                TokenKind::Colon => { // TODO
                    unimplemented!();
                },

                TokenKind::Indent => { // set new self.index
                    let mut indent = 0;

                    while self.tokens[i].kind == TokenKind::Indent {
                        indent += 1;
                        i += 1;
                    }

                    i -= 1;

                    self.indent = indent;
                },

                TokenKind::Newline => { // reset self.index
                    self.indent = 0;
                },

                TokenKind::EOF => break,
            }
            i += 1;
        }

        Ok(self.pair)
    }
}

<|endoftext|>


<commit_before>use ast::*;
use scanner::{Lexeme, Position, Token, TokenWithContext};
use std::iter::Peekable;

/// This behave exactly as try! but wraps the returned result in a Some.
/// It's useful to remove some boilerplate in the code introduced by
/// the use of Option<Result<T, E>>
macro_rules! try_wrap_err {
    ($e:expr) => (match $e {Ok(e) => e, Err(e) => return Some(Err(e))})
}

#[derive(Debug)]
pub enum ParseError {
    UnexpectedEndOfFile,
    UnexpectedToken(Lexeme, Position),
    MissingSubexpression(Lexeme, Position),
    MissingClosingParen(Lexeme, Position),
    MissingSemicolon(Lexeme, Position),
}

pub fn parse(tokens: &[TokenWithContext]) -> Result<Vec<Statement>, Vec<ParseError>> {
    let mut iter = tokens.iter().peekable();
    let mut statements = Vec::new();
    let mut errors = Vec::new();
    while let Some(result) = parse_declaration(&mut iter) {
        match result {
            Ok(statement) => {
                statements.push(statement);
            }
            Err(error) => {
                errors.push(error);
                synchronise(&mut iter);
            }
        }
    }
    if errors.is_empty() {
        Ok(statements)
    } else {
        Err(errors)
    }
}

fn synchronise<'a, I>(tokens: &mut Peekable<I>)
    where I: Iterator<Item = &'a TokenWithContext>
{
    enum PositionInConstruct {
        Start,
        Body,
        End,
    };
    fn classify(token: &Token) -> PositionInConstruct {
        match *token {
            Token::Semicolon => PositionInConstruct::End,
            Token::Class | Token::Fun | Token::Var | Token::For | Token::If | Token::While |
            Token::Print | Token::Return => PositionInConstruct::Start,
            _ => PositionInConstruct::Body,
        }
    }
    while let Some(token_kind) = tokens.peek().map(|t| classify(&t.token)) {
        match token_kind {
            PositionInConstruct::Start => {
                break;
            }
            PositionInConstruct::Body => {
                let _ = tokens.next();
            }
            PositionInConstruct::End => {
                let _ = tokens.next();
                break;
            }
        }
    }
}

fn parse_declaration<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Statement, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    match tokens.peek().map(|t| &t.token) {
        Some(&Token::Var) => {
            let _ = tokens.next();
            parse_var_declaration(tokens)
        }
        Some(_) => parse_statement(tokens),
        None => None,
    }
}

fn parse_var_declaration<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Statement, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    unimplemented!()
}

fn parse_statement<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Statement, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    let result = match tokens.peek().map(|t| &t.token) {
        Some(&Token::Print) => {
            let _ = tokens.next();
            parse_print_statement(tokens)
        }
        Some(_) => parse_expression_statement(tokens),
        None => None,
    };
    if let Some(Ok(statement)) = result {
        match tokens.peek() {
            Some(&&TokenWithContext { token: Token::Semicolon, lexeme: _, position: _ }) => {
                let _ = tokens.next();
                Some(Ok(statement))
            }
            Some(&&TokenWithContext { token: _, ref lexeme, ref position }) => {
                Some(Err(ParseError::MissingSemicolon(lexeme.clone(), position.clone())))
            }
            None => Some(Err(ParseError::UnexpectedEndOfFile)),
        }
    } else {
        result
    }
}

fn parse_print_statement<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Statement, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    parse_expression(tokens).map(|r| r.map(|e| Statement::Print(e)))
}

fn parse_expression_statement<'a, I>(tokens: &mut Peekable<I>)
                                     -> Option<Result<Statement, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    parse_expression(tokens).map(|r| r.map(|e| Statement::Expression(e)))
}

fn parse_expression<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    parse_equality(tokens)
}

fn parse_binary<'a, I>(tokens: &mut Peekable<I>,
                       map_operator: &Fn(&Token) -> Option<BinaryOperator>,
                       parse_subexpression: &Fn(&mut Peekable<I>)
                                                -> Option<Result<Expr, ParseError>>)
                       -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    let mut expr;
    {
        let result = parse_subexpression(tokens);
        if let Some(Ok(subexpression)) = result {
            expr = subexpression;
        } else {
            return result;
        }
    };
    while let Some(Some(mapped_operator)) = tokens.peek().map(|pt| map_operator(&pt.token)) {
        let operator; // We know it's right, we read it just for error reporting
        {
            operator = tokens.next().unwrap();
        }
        let right = if let Some(result) = parse_subexpression(tokens) {
            try_wrap_err!(result)
        } else {
            return Some(Err(ParseError::MissingSubexpression(operator.lexeme.clone(),
                                                             operator.position)));
        };
        let binary_expression = BinaryExpr {
            left: expr,
            operator: mapped_operator,
            right: right,
        };
        expr = Expr::Binary(Box::new(binary_expression));
    }
    Some(Ok(expr))
}

fn parse_equality<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    fn map_operator(token: &Token) -> Option<BinaryOperator> {
        match *token {
            Token::BangEqual => Some(BinaryOperator::NotEqual),
            Token::EqualEqual => Some(BinaryOperator::Equal),
            _ => None,
        }
    }
    parse_binary(tokens, &map_operator, &parse_comparison)
}

fn parse_comparison<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    fn map_operator(token: &Token) -> Option<BinaryOperator> {
        match *token {
            Token::Greater => Some(BinaryOperator::Greater),
            Token::GreaterEqual => Some(BinaryOperator::GreaterEqual),
            Token::Less => Some(BinaryOperator::Less),
            Token::LessEqual => Some(BinaryOperator::LessEqual),
            _ => None,
        }
    }
    parse_binary(tokens, &map_operator, &parse_term)
}

fn parse_term<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    fn map_operator(token: &Token) -> Option<BinaryOperator> {
        match *token {
            Token::Minus => Some(BinaryOperator::Minus),
            Token::Plus => Some(BinaryOperator::Plus),
            _ => None,
        }
    }
    parse_binary(tokens, &map_operator, &parse_factor)
}

fn parse_factor<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    fn map_operator(token: &Token) -> Option<BinaryOperator> {
        match *token {
            Token::Slash => Some(BinaryOperator::Slash),
            Token::Star => Some(BinaryOperator::Star),
            _ => None,
        }
    }
    parse_binary(tokens, &map_operator, &parse_unary)
}

fn parse_unary<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    fn map_operator(token: &Token) -> Option<UnaryOperator> {
        match *token {
            Token::Minus => Some(UnaryOperator::Minus),
            Token::Bang => Some(UnaryOperator::Bang),
            _ => None,
        }
    }
    if let Some(Some(mapped_operator)) = tokens.peek().map(|pt| map_operator(&pt.token)) {
        let operator; // For error reporting
        {
            operator = tokens.next().unwrap();
        }
        let right = if let Some(result) = parse_unary(tokens) {
            try_wrap_err!(result)
        } else {
            return Some(Err(ParseError::MissingSubexpression(operator.lexeme.clone(),
                                                             operator.position)));
        };
        let unary_expression = UnaryExpr {
            operator: mapped_operator,
            right: right,
        };
        return Some(Ok(Expr::Unary(Box::new(unary_expression))));
    } else {
        parse_primary(tokens)
    }
}

fn parse_primary<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    let primary_token;
    {
        primary_token = tokens.next();
    };
    if let Some(primary_token) = primary_token {
        let parsed_expression = match primary_token.token {
            Token::False => Expr::Literal(Literal::BoolLiteral(false)),
            Token::True => Expr::Literal(Literal::BoolLiteral(true)),
            Token::Nil => Expr::Literal(Literal::NilLiteral),
            Token::NumberLiteral(n) => Expr::Literal(Literal::NumberLiteral(n)),
            Token::StringLiteral(ref s) => Expr::Literal(Literal::StringLiteral(s.clone())),
            Token::LeftParen => {
                let expr = if let Some(result) = parse_expression(tokens) {
                    try_wrap_err!(result)
                } else {
                    return Some(Err(ParseError::UnexpectedEndOfFile));
                };
                {
                    if let Some(token) = tokens.next() {
                        if token.token == Token::RightParen {
                            let grouping_expression = Grouping { expr: expr };
                            return Some(Ok(Expr::Grouping(Box::new(grouping_expression))));
                        } else {
                            return Some(Err(ParseError::MissingClosingParen(token.lexeme.clone(),
                                                                            token.position)));
                        }

                    }
                    return Some(Err(ParseError::UnexpectedEndOfFile));
                }
            }
            _ => {
                return Some(Err(ParseError::UnexpectedToken(primary_token.lexeme.clone(),
                                                            primary_token.position)));
            }
        };
        Some(Ok(parsed_expression))
    } else {
        None
    }
}

#[cfg(test)]
mod tests {
    use scanner::*;
    use parser::*;
    use pretty_printer::PrettyPrint;

    #[test]
    fn literal() {
        let string = String::from("123");
        let (tokens, _) = scan(&string);
        let expr = parse_expression(&mut tokens.iter().peekable()).unwrap().unwrap();
        assert_eq!(&string, &expr.pretty_print());
    }

    #[test]
    fn binary() {
        let (tokens, _) = scan(&"123+456");
        let expr = parse_expression(&mut tokens.iter().peekable()).unwrap().unwrap();
        assert_eq!("(+ 123 456)", &expr.pretty_print());
    }

    #[test]
    fn precedence_add_mul() {
        let (tokens, _) = scan(&"123+456*789");
        let expr = parse_expression(&mut tokens.iter().peekable()).unwrap().unwrap();
        assert_eq!("(+ 123 (* 456 789))", &expr.pretty_print());
    }

    #[test]
    fn precedence_mul_add() {
        let (tokens, _) = scan(&"123*456+789");
        let expr = parse_expression(&mut tokens.iter().peekable()).unwrap().unwrap();
        assert_eq!("(+ (* 123 456) 789)", &expr.pretty_print());
    }

    #[test]
    fn precedence_mul_add_unary() {
        let (tokens, _) = scan(&"-123*456+789");
        let expr = parse_expression(&mut tokens.iter().peekable()).unwrap().unwrap();
        assert_eq!("(+ (* (- 123) 456) 789)", &expr.pretty_print());
    }

    #[test]
    fn unclosed_group() {
        let (tokens, _) = scan(&"(2");
        assert!(parse(&tokens).is_err());
    }

    #[test]
    fn unopened_group() {
        let (tokens, _) = scan(&"2)");
        assert!(parse(&tokens).is_err());
    }
}
<commit_msg>Partial implementation of the function to parse variable declarations<commit_after>use ast::*;
use scanner::{Lexeme, Position, Token, TokenWithContext};
use std::iter::Peekable;

/// This behave exactly as try! but wraps the returned result in a Some.
/// It's useful to remove some boilerplate in the code introduced by
/// the use of Option<Result<T, E>>
macro_rules! try_wrap_err {
    ($e:expr) => (match $e {Ok(e) => e, Err(e) => return Some(Err(e))})
}

#[derive(Debug)]
pub enum ParseError {
    UnexpectedEndOfFile,
    UnexpectedToken(Lexeme, Position),
    MissingSubexpression(Lexeme, Position),
    MissingClosingParen(Lexeme, Position),
    MissingSemicolon(Lexeme, Position),
    MissingIdentifier(Lexeme, Position),
}

pub fn parse(tokens: &[TokenWithContext]) -> Result<Vec<Statement>, Vec<ParseError>> {
    let mut iter = tokens.iter().peekable();
    let mut statements = Vec::new();
    let mut errors = Vec::new();
    while let Some(result) = parse_declaration(&mut iter) {
        match result {
            Ok(statement) => {
                statements.push(statement);
            }
            Err(error) => {
                errors.push(error);
                synchronise(&mut iter);
            }
        }
    }
    if errors.is_empty() {
        Ok(statements)
    } else {
        Err(errors)
    }
}

fn synchronise<'a, I>(tokens: &mut Peekable<I>)
    where I: Iterator<Item = &'a TokenWithContext>
{
    enum PositionInConstruct {
        Start,
        Body,
        End,
    };
    fn classify(token: &Token) -> PositionInConstruct {
        match *token {
            Token::Semicolon => PositionInConstruct::End,
            Token::Class | Token::Fun | Token::Var | Token::For | Token::If | Token::While |
            Token::Print | Token::Return => PositionInConstruct::Start,
            _ => PositionInConstruct::Body,
        }
    }
    while let Some(token_kind) = tokens.peek().map(|t| classify(&t.token)) {
        match token_kind {
            PositionInConstruct::Start => {
                break;
            }
            PositionInConstruct::Body => {
                let _ = tokens.next();
            }
            PositionInConstruct::End => {
                let _ = tokens.next();
                break;
            }
        }
    }
}

fn parse_semicolon_terminated_statement<'a, I>(tokens: &mut Peekable<I>,
                                               parse_statement: &Fn(&mut Peekable<I>)
                                                                    -> Option<Result<Statement,
                                                                                ParseError>>)
                                               -> Option<Result<Statement, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    match parse_statement(tokens) {
        Some(Ok(statement)) => {
            match tokens.peek() {
                Some(&&TokenWithContext { token: Token::Semicolon, lexeme: _, position: _ }) => {
                    let _ = tokens.next();
                    Some(Ok(statement))
                }
                Some(&&TokenWithContext { token: _, ref lexeme, ref position }) => {
                    Some(Err(ParseError::MissingSemicolon(lexeme.clone(), position.clone())))
                }
                None => Some(Err(ParseError::UnexpectedEndOfFile)),
            }
        }
        e => e,
    }
}

fn parse_declaration<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Statement, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    match tokens.peek().map(|t| &t.token) {
        Some(&Token::Var) => {
            let _ = tokens.next();
            parse_semicolon_terminated_statement(tokens, &parse_var_declaration)
        }
        Some(_) => parse_semicolon_terminated_statement(tokens, &parse_statement),
        None => None,
    }
}

fn parse_var_declaration<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Statement, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    let identifier = match tokens.peek() {
        Some(&&TokenWithContext { token: Token::Identifier(ref identifier),
                                  lexeme: _,
                                  position: _ }) => {
            let _ = tokens.next();
            identifier.clone()
        }
        Some(&&TokenWithContext { token: _, ref lexeme, ref position }) => {
            return Some(Err(ParseError::MissingIdentifier(lexeme.clone(), position.clone())))
        }
        None => return Some(Err(ParseError::UnexpectedEndOfFile)),
    };
    let statement = if let Some(&&TokenWithContext { token: Token::Equal,
                                                     lexeme: _,
                                                     position: _ }) = tokens.peek() {
        unimplemented!();
    } else {
        unimplemented!();
    };
}

fn parse_statement<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Statement, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    match tokens.peek().map(|t| &t.token) {
        Some(&Token::Print) => {
            let _ = tokens.next();
            parse_print_statement(tokens)
        }
        Some(_) => parse_expression_statement(tokens),
        None => None,
    }
}

fn parse_print_statement<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Statement, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    parse_expression(tokens).map(|r| r.map(|e| Statement::Print(e)))
}

fn parse_expression_statement<'a, I>(tokens: &mut Peekable<I>)
                                     -> Option<Result<Statement, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    parse_expression(tokens).map(|r| r.map(|e| Statement::Expression(e)))
}

fn parse_expression<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    parse_equality(tokens)
}

fn parse_binary<'a, I>(tokens: &mut Peekable<I>,
                       map_operator: &Fn(&Token) -> Option<BinaryOperator>,
                       parse_subexpression: &Fn(&mut Peekable<I>)
                                                -> Option<Result<Expr, ParseError>>)
                       -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    let mut expr;
    {
        let result = parse_subexpression(tokens);
        if let Some(Ok(subexpression)) = result {
            expr = subexpression;
        } else {
            return result;
        }
    };
    while let Some(Some(mapped_operator)) = tokens.peek().map(|pt| map_operator(&pt.token)) {
        let operator; // We know it's right, we read it just for error reporting
        {
            operator = tokens.next().unwrap();
        }
        let right = if let Some(result) = parse_subexpression(tokens) {
            try_wrap_err!(result)
        } else {
            return Some(Err(ParseError::MissingSubexpression(operator.lexeme.clone(),
                                                             operator.position)));
        };
        let binary_expression = BinaryExpr {
            left: expr,
            operator: mapped_operator,
            right: right,
        };
        expr = Expr::Binary(Box::new(binary_expression));
    }
    Some(Ok(expr))
}

fn parse_equality<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    fn map_operator(token: &Token) -> Option<BinaryOperator> {
        match *token {
            Token::BangEqual => Some(BinaryOperator::NotEqual),
            Token::EqualEqual => Some(BinaryOperator::Equal),
            _ => None,
        }
    }
    parse_binary(tokens, &map_operator, &parse_comparison)
}

fn parse_comparison<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    fn map_operator(token: &Token) -> Option<BinaryOperator> {
        match *token {
            Token::Greater => Some(BinaryOperator::Greater),
            Token::GreaterEqual => Some(BinaryOperator::GreaterEqual),
            Token::Less => Some(BinaryOperator::Less),
            Token::LessEqual => Some(BinaryOperator::LessEqual),
            _ => None,
        }
    }
    parse_binary(tokens, &map_operator, &parse_term)
}

fn parse_term<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    fn map_operator(token: &Token) -> Option<BinaryOperator> {
        match *token {
            Token::Minus => Some(BinaryOperator::Minus),
            Token::Plus => Some(BinaryOperator::Plus),
            _ => None,
        }
    }
    parse_binary(tokens, &map_operator, &parse_factor)
}

fn parse_factor<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    fn map_operator(token: &Token) -> Option<BinaryOperator> {
        match *token {
            Token::Slash => Some(BinaryOperator::Slash),
            Token::Star => Some(BinaryOperator::Star),
            _ => None,
        }
    }
    parse_binary(tokens, &map_operator, &parse_unary)
}

fn parse_unary<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    fn map_operator(token: &Token) -> Option<UnaryOperator> {
        match *token {
            Token::Minus => Some(UnaryOperator::Minus),
            Token::Bang => Some(UnaryOperator::Bang),
            _ => None,
        }
    }
    if let Some(Some(mapped_operator)) = tokens.peek().map(|pt| map_operator(&pt.token)) {
        let operator; // For error reporting
        {
            operator = tokens.next().unwrap();
        }
        let right = if let Some(result) = parse_unary(tokens) {
            try_wrap_err!(result)
        } else {
            return Some(Err(ParseError::MissingSubexpression(operator.lexeme.clone(),
                                                             operator.position)));
        };
        let unary_expression = UnaryExpr {
            operator: mapped_operator,
            right: right,
        };
        return Some(Ok(Expr::Unary(Box::new(unary_expression))));
    } else {
        parse_primary(tokens)
    }
}

fn parse_primary<'a, I>(tokens: &mut Peekable<I>) -> Option<Result<Expr, ParseError>>
    where I: Iterator<Item = &'a TokenWithContext>
{
    let primary_token;
    {
        primary_token = tokens.next();
    };
    if let Some(primary_token) = primary_token {
        let parsed_expression = match primary_token.token {
            Token::False => Expr::Literal(Literal::BoolLiteral(false)),
            Token::True => Expr::Literal(Literal::BoolLiteral(true)),
            Token::Nil => Expr::Literal(Literal::NilLiteral),
            Token::NumberLiteral(n) => Expr::Literal(Literal::NumberLiteral(n)),
            Token::StringLiteral(ref s) => Expr::Literal(Literal::StringLiteral(s.clone())),
            Token::LeftParen => {
                let expr = if let Some(result) = parse_expression(tokens) {
                    try_wrap_err!(result)
                } else {
                    return Some(Err(ParseError::UnexpectedEndOfFile));
                };
                {
                    if let Some(token) = tokens.next() {
                        if token.token == Token::RightParen {
                            let grouping_expression = Grouping { expr: expr };
                            return Some(Ok(Expr::Grouping(Box::new(grouping_expression))));
                        } else {
                            return Some(Err(ParseError::MissingClosingParen(token.lexeme.clone(),
                                                                            token.position)));
                        }

                    }
                    return Some(Err(ParseError::UnexpectedEndOfFile));
                }
            }
            _ => {
                return Some(Err(ParseError::UnexpectedToken(primary_token.lexeme.clone(),
                                                            primary_token.position)));
            }
        };
        Some(Ok(parsed_expression))
    } else {
        None
    }
}

#[cfg(test)]
mod tests {
    use scanner::*;
    use parser::*;
    use pretty_printer::PrettyPrint;

    #[test]
    fn literal() {
        let string = String::from("123");
        let (tokens, _) = scan(&string);
        let expr = parse_expression(&mut tokens.iter().peekable()).unwrap().unwrap();
        assert_eq!(&string, &expr.pretty_print());
    }

    #[test]
    fn binary() {
        let (tokens, _) = scan(&"123+456");
        let expr = parse_expression(&mut tokens.iter().peekable()).unwrap().unwrap();
        assert_eq!("(+ 123 456)", &expr.pretty_print());
    }

    #[test]
    fn precedence_add_mul() {
        let (tokens, _) = scan(&"123+456*789");
        let expr = parse_expression(&mut tokens.iter().peekable()).unwrap().unwrap();
        assert_eq!("(+ 123 (* 456 789))", &expr.pretty_print());
    }

    #[test]
    fn precedence_mul_add() {
        let (tokens, _) = scan(&"123*456+789");
        let expr = parse_expression(&mut tokens.iter().peekable()).unwrap().unwrap();
        assert_eq!("(+ (* 123 456) 789)", &expr.pretty_print());
    }

    #[test]
    fn precedence_mul_add_unary() {
        let (tokens, _) = scan(&"-123*456+789");
        let expr = parse_expression(&mut tokens.iter().peekable()).unwrap().unwrap();
        assert_eq!("(+ (* (- 123) 456) 789)", &expr.pretty_print());
    }

    #[test]
    fn unclosed_group() {
        let (tokens, _) = scan(&"(2");
        assert!(parse(&tokens).is_err());
    }

    #[test]
    fn unopened_group() {
        let (tokens, _) = scan(&"2)");
        assert!(parse(&tokens).is_err());
    }
}
<|endoftext|>


<commit_before>use DocId;
use postings::{Postings, DocSet};
use std::cmp::Ordering;
use std::collections::BinaryHeap;
use query::MultiTermScorer;
use postings::ScoredDocSet;
use query::Scorer;
use fastfield::U32FastFieldReader;

#[derive(Eq, PartialEq)]
struct HeapItem(DocId, usize, u32);

impl HeapItem {
    pub fn ord(&self,) -> usize{
        self.1
    }
}

impl PartialOrd for HeapItem {
    fn partial_cmp(&self, other:&Self) -> Option<Ordering> {
        Some(self.cmp(&other))
    }
}

impl Ord for HeapItem {
    fn cmp(&self, other:&Self) -> Ordering {
         (self.0, self.1).cmp(&(other.0, other.1)).reverse()
    }
}

pub struct UnionPostings<TPostings: Postings> {
    fieldnorms_readers: Vec<U32FastFieldReader>,
    postings: Vec<TPostings>,
    queue: BinaryHeap<HeapItem>,
    doc: DocId,
    scorer: MultiTermScorer
}

impl<TPostings: Postings> UnionPostings<TPostings> {

    pub fn new(fieldnorms_reader: Vec<U32FastFieldReader>, mut postings: Vec<TPostings>, multi_term_scorer: MultiTermScorer) -> UnionPostings<TPostings> {
        let num_postings = postings.len();
        assert_eq!(fieldnorms_reader.len(), num_postings);

        for posting in &mut postings {
            assert!(posting.next());
        }
        let heap_items: Vec<HeapItem> = postings.iter()
            .map(|posting| {
                (posting.doc(), posting.term_freq())
            })
            .enumerate()
            .map(|(ord, (doc, tf))| {
                HeapItem(doc, ord, tf)
            })
            .collect();

        UnionPostings {
            fieldnorms_readers: fieldnorms_reader,
            postings: postings,
            queue: BinaryHeap::from(heap_items),
            doc: 0,
            scorer: multi_term_scorer
        }
    }

    fn advance_head(&mut self,) {
        let ord = self.queue.peek().unwrap().ord();
        let cur_postings = &mut self.postings[ord];
        if cur_postings.next() {
            let doc = cur_postings.doc();
            let tf = cur_postings.term_freq();
            self.queue.replace(HeapItem(doc, ord, tf));
        }
        else {
            self.queue.pop();
        }

    }

    fn get_field_norm(&self, ord:usize, doc:DocId) -> u32 {
        self.fieldnorms_readers[ord].get(doc)
    }

}

impl<TPostings: Postings> DocSet for UnionPostings<TPostings> {

    fn next(&mut self,) -> bool {
        self.scorer.clear();
        let cur_doc: DocId;
        match self.queue.peek() {
            Some(&HeapItem(doc, ord, tf)) => {
                cur_doc = doc;
                let fieldnorm = self.get_field_norm(ord, doc);
                self.scorer.update(ord, tf, fieldnorm);

            }
            None => {
                return false;
            }
        }
        self.advance_head();
        loop {
            match self.queue.peek() {
                Some(&HeapItem(peek_doc, peek_ord, peek_tf))  => {
                    if peek_doc != cur_doc {
                        break;
                    }
                    else {
                        let fieldnorm = self.get_field_norm(peek_ord, peek_doc);
                        self.scorer.update(peek_ord, peek_tf, fieldnorm);
                    }
                }
                None => { break; }
            }
            self.advance_head();
        }
        self.doc = cur_doc;
        return true;
    }

    // TODO implement a faster skip_next

    fn doc(&self,) -> DocId {
        self.doc
    }

    fn doc_freq(&self,) -> usize {
        panic!("Doc freq");
    }
}

impl<TPostings: Postings> ScoredDocSet for UnionPostings<TPostings> {
    fn score(&self,) -> f32 {
        self.scorer.score()
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use postings::{DocSet, VecPostings, ScoredDocSet};
    use query::MultiTermScorer;
    use directory::ReadOnlySource;
    use directory::SharedVec;
    use schema::Field;
    use fastfield::{U32FastFieldReader, U32FastFieldWriter, FastFieldSerializer};

    pub fn create_u32_fastfieldreader(field: Field, vals: Vec<u32>) -> U32FastFieldReader {
        let mut u32_field_writer = U32FastFieldWriter::new(field);
        for val in vals {
            u32_field_writer.add_val(val);
        }
        let data = SharedVec::new();
        let write: Box<SharedVec> = Box::new(data.clone());
        let mut serializer = FastFieldSerializer::new(write).unwrap();
        u32_field_writer.serialize(&mut serializer).unwrap();
        serializer.close().unwrap();
        U32FastFieldReader::open(ReadOnlySource::Anonymous(data.copy_vec())).unwrap()
    }

    fn abs_diff(left: f32, right: f32) -> f32 {
        (right - left).abs()
    }

    #[test]
    pub fn test_union_postings() {
        let left_fieldnorms = create_u32_fastfieldreader(Field(1), vec!(100,200,300));
        let right_fieldnorms = create_u32_fastfieldreader(Field(2), vec!(15,25,35));
        let left = VecPostings::new(vec!(1, 2, 3));
        let right = VecPostings::new(vec!(1, 3, 8));
        let multi_term_scorer = MultiTermScorer::new(vec!(1f32, 2f32), vec!(1f32, 4f32));
        let mut union = UnionPostings::new(
            vec!(left_fieldnorms, right_fieldnorms),
            vec!(left, right),
            multi_term_scorer
        );
        assert!(union.next());
        assert_eq!(union.doc(), 1);
        assert!(abs_diff(union.score(), 2.182179f32) < 0.001);
        assert!(union.next());
        assert_eq!(union.doc(), 2);
        assert!(abs_diff(union.score(), 0.2236068) < 0.001f32);
        assert!(union.next());
        assert_eq!(union.doc(), 3);
        assert!(union.next());
        assert!(abs_diff(union.score(), 0.8944272f32) < 0.001f32);
        assert_eq!(union.doc(), 8);
        assert!(!union.next());
    }

}

<commit_msg>Lighter HeapItem by putting term freq aside, and keeping ord as u32<commit_after>use DocId;
use postings::{Postings, DocSet};
use std::cmp::Ordering;
use std::collections::BinaryHeap;
use query::MultiTermScorer;
use postings::ScoredDocSet;
use query::Scorer;
use fastfield::U32FastFieldReader;
use std::iter;

#[derive(Eq, PartialEq)]
struct HeapItem(DocId, u32);

impl HeapItem {
    pub fn ord(&self,) -> usize{
        self.1 as usize
    }
}

impl PartialOrd for HeapItem {
    fn partial_cmp(&self, other:&Self) -> Option<Ordering> {
        Some(self.cmp(&other))
    }
}

impl Ord for HeapItem {
    fn cmp(&self, other:&Self) -> Ordering {
         (other.0).cmp(&self.0)
    }
}

pub struct UnionPostings<TPostings: Postings> {
    fieldnorms_readers: Vec<U32FastFieldReader>,
    postings: Vec<TPostings>,
    term_frequencies: Vec<u32>,
    queue: BinaryHeap<HeapItem>,
    doc: DocId,
    scorer: MultiTermScorer
}

impl<TPostings: Postings> UnionPostings<TPostings> {

    pub fn new(fieldnorms_reader: Vec<U32FastFieldReader>, mut postings: Vec<TPostings>, multi_term_scorer: MultiTermScorer) -> UnionPostings<TPostings> {
        let num_postings = postings.len();
        assert_eq!(fieldnorms_reader.len(), num_postings);

        for posting in &mut postings {
            assert!(posting.next());
        }
        let mut term_frequencies: Vec<u32> = iter::repeat(0u32).take(num_postings).collect();
        let heap_items: Vec<HeapItem> = postings
            .iter()
            .map(|posting| {
                (posting.doc(), posting.term_freq())
            })
            .enumerate()
            .map(|(ord, (doc, tf))| {
                term_frequencies[ord] = tf;
                HeapItem(doc, ord as u32)
            })
            .collect();

        UnionPostings {
            fieldnorms_readers: fieldnorms_reader,
            postings: postings,
            term_frequencies: term_frequencies,
            queue: BinaryHeap::from(heap_items),
            doc: 0,
            scorer: multi_term_scorer
        }
    }

    fn advance_head(&mut self,) {
        let ord = self.queue.peek().unwrap().ord();
        let cur_postings = &mut self.postings[ord];
        if cur_postings.next() {
            let doc = cur_postings.doc();
            self.term_frequencies[ord] = cur_postings.term_freq();
            self.queue.replace(HeapItem(doc, ord as u32));
        }
        else {
            self.queue.pop();
        }
    }

    fn get_field_norm(&self, ord:usize, doc:DocId) -> u32 {
        self.fieldnorms_readers[ord].get(doc)
    }

}

impl<TPostings: Postings> DocSet for UnionPostings<TPostings> {

    fn next(&mut self,) -> bool {
        self.scorer.clear();
        let cur_doc: DocId;
        match self.queue.peek() {
            Some(&HeapItem(doc, ord)) => {
                cur_doc = doc;
                let ord: usize = ord as usize;
                let fieldnorm = self.get_field_norm(ord, doc);
                let tf = self.term_frequencies[ord];
                self.scorer.update(ord, tf, fieldnorm);

            }
            None => {
                return false;
            }
        }
        self.advance_head();
        loop {
            match self.queue.peek() {
                Some(&HeapItem(peek_doc, peek_ord))  => {
                    if peek_doc != cur_doc {
                        break;
                    }
                    else {
                        let peek_ord: usize = peek_ord as usize;
                        let peek_tf = self.term_frequencies[peek_ord];
                        let fieldnorm = self.get_field_norm(peek_ord, peek_doc);
                        self.scorer.update(peek_ord, peek_tf, fieldnorm);
                    }
                }
                None => { break; }
            }
            self.advance_head();
        }
        self.doc = cur_doc;
        return true;
    }

    // TODO implement a faster skip_next

    fn doc(&self,) -> DocId {
        self.doc
    }

    fn doc_freq(&self,) -> usize {
        panic!("Doc freq");
    }
}

impl<TPostings: Postings> ScoredDocSet for UnionPostings<TPostings> {
    fn score(&self,) -> f32 {
        self.scorer.score()
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use postings::{DocSet, VecPostings, ScoredDocSet};
    use query::MultiTermScorer;
    use directory::ReadOnlySource;
    use directory::SharedVec;
    use schema::Field;
    use fastfield::{U32FastFieldReader, U32FastFieldWriter, FastFieldSerializer};

    pub fn create_u32_fastfieldreader(field: Field, vals: Vec<u32>) -> U32FastFieldReader {
        let mut u32_field_writer = U32FastFieldWriter::new(field);
        for val in vals {
            u32_field_writer.add_val(val);
        }
        let data = SharedVec::new();
        let write: Box<SharedVec> = Box::new(data.clone());
        let mut serializer = FastFieldSerializer::new(write).unwrap();
        u32_field_writer.serialize(&mut serializer).unwrap();
        serializer.close().unwrap();
        U32FastFieldReader::open(ReadOnlySource::Anonymous(data.copy_vec())).unwrap()
    }

    fn abs_diff(left: f32, right: f32) -> f32 {
        (right - left).abs()
    }

    #[test]
    pub fn test_union_postings() {
        let left_fieldnorms = create_u32_fastfieldreader(Field(1), vec!(100,200,300));
        let right_fieldnorms = create_u32_fastfieldreader(Field(2), vec!(15,25,35));
        let left = VecPostings::new(vec!(1, 2, 3));
        let right = VecPostings::new(vec!(1, 3, 8));
        let multi_term_scorer = MultiTermScorer::new(vec!(1f32, 2f32), vec!(1f32, 4f32));
        let mut union = UnionPostings::new(
            vec!(left_fieldnorms, right_fieldnorms),
            vec!(left, right),
            multi_term_scorer
        );
        assert!(union.next());
        assert_eq!(union.doc(), 1);
        assert!(abs_diff(union.score(), 2.182179f32) < 0.001);
        assert!(union.next());
        assert_eq!(union.doc(), 2);
        assert!(abs_diff(union.score(), 0.2236068) < 0.001f32);
        assert!(union.next());
        assert_eq!(union.doc(), 3);
        assert!(union.next());
        assert!(abs_diff(union.score(), 0.8944272f32) < 0.001f32);
        assert_eq!(union.doc(), 8);
        assert!(!union.next());
    }

}

<|endoftext|>


<commit_before>// Copyright 2017 Google Inc.
//
// 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
//
//     https://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.

#![cfg_attr(feature = "clippy", feature(plugin))]
#![cfg_attr(feature = "clippy", plugin(clippy))]

extern crate clap;
extern crate json;
extern crate rerast;

use std::io::{self, Write};
use rerast::{ArgBuilder, Config, RerastCompilerDriver};
use std::fs::OpenOptions;
use std::path::Path;

// Environment variables that we use to pass data from the outer invocation of cargo-rerast through
// to the inner invocation (where it acts as the rust compiler). Eventually we'll hopefully be able
// to query cargo as to what rustc commands are needed to build, then we can use those
// directly. Possibly we could just run "cargo check -v" and collect its output. Either way, we'd
// like to get rid of these as they add unnecessary complexity.
mod var_names {
    pub const RULES_FILE: &'static str = "RERAST_RULES_FILE";
    pub const RULES: &'static str = "RERAST_RULES";
    pub const DIFF_CMD: &'static str = "RERAST_DIFF_CMD";
    pub const DEBUG_SNIPPET: &'static str = "RERAST_DEBUG_SNIPPET";
}

// Queries cargo to find the file for each target and touches is, so that when we run cargo check
// with ourselves as the compiler, we can be sure that we'll get run. Hopefully eventually there'll
// be a nicer way to integrate with cargo such that we won't need to do this.
fn touch_source_files() -> io::Result<()> {
    let output = std::process::Command::new("cargo")
        .args(vec!["metadata", "--no-deps", "--format-version=1"])
        .stdout(std::process::Stdio::piped())
        .output()?;
    let metadata_str = std::str::from_utf8(output.stdout.as_slice()).map_err(|e| {
        io::Error::new(io::ErrorKind::InvalidData, e.to_string())
    })?;
    let parsed = json::parse(metadata_str).map_err(|e| {
        io::Error::new(io::ErrorKind::InvalidData, e.to_string())
    })?;
    for package in parsed["packages"].members() {
        for target in package["targets"].members() {
            if let Some(src_path) = target["src_path"].as_str() {
                println!("touching {}", src_path);
                OpenOptions::new()
                    .create(false)
                    .append(true)
                    .write(true)
                    .open(Path::new(src_path))?;
            }
        }
    }
    Ok(())
}

fn cargo_rerast() {
    let matches = clap::App::new("Rerast")
        .version("0.1")
        .author("David Lattimore <dml@google.com>")
        .subcommand(
            clap::SubCommand::with_name("rerast")
                .about("Replace Rust code based on typed, syntactic patterns")
                .args_from_usage(
                    "--rules_file=[FILE] 'Path to a rule file'
                     --use_statements=[USE_STATEMENTS] 'Use statements required by rule'
                     --placeholders=[PLACEHOLDERS] 'e.g. <T>(o: option<T>)'
                     --replace=[PATTERN] 'Pattern to search for'
                     --with=[PATTERN] 'Replacement code'
                     --debug_snippet=[CODE_SNIPPET] 'A snippet of code that you think should \
                                                     match or list_all to list all checked \
                                                     snippets.'
                                      \
                     --diff 'Diff changes'
                     --colordiff 'Diff changes with colordiff'",
                ),
        )
        .get_matches();
    if let Some(matches) = matches.subcommand_matches("rerast") {
        // User ran "cargo rerast", invoke cargo build with us as "rustc"
        let rules = if let (Some(search), Some(replacement)) =
            (matches.value_of("replace"), matches.value_of("with"))
        {
            let mut placeholders = matches.value_of("placeholders").unwrap_or("").to_owned();
            if !placeholders.contains('(') {
                placeholders = "(".to_owned() + &placeholders + ")";
            }
            matches.value_of("use_statements").unwrap_or("").to_owned() + "pub fn rule"
                + &placeholders + "{replace!(" + search + " => " + replacement + ");}"
        } else {
            "".to_owned()
        };
        let rules_file = matches.value_of("rules_file").unwrap_or("");
        if rules_file.is_empty() == rules.is_empty() {
            eprintln!("Must specify either --rules_file or both of --search and --replacement");
            std::process::exit(1);
        }
        if matches.is_present("diff") == matches.is_present("colordiff") {
            eprintln!("Currently either --diff or --colordiff must be specified");
            std::process::exit(1);
        }
        // We run cargo with RUSTC_WRAPPER set, so that it then uses us as the compiler. Really what
        // we'd like is to run cargo, asking it to build all the crate dependencies, then tell us
        // the commandline(s) to build all the targets in the current crate. This is pretty
        // ugly. Clippy does something similar, so there probably isn't a better way at the moment.
        // Unfortunately this hack has several consequences. (1) We're invoked separately for each
        // target, which precludes, or would at least make it hard to produce a summary of all
        // matches. (2) Cargo won't run us if the crate has already been built, requiring that we
        // touch the main source file in order to force it to actually run us. (3) If our
        // dependencies haven't already been built, we'll try to look for matches in them - not
        // really what we want. (4) Cargo sometimes prints messages that are either not helpful or
        // outright confusing - e.g. asking the user to add --verbose.
        touch_source_files().unwrap();
        let cargo_exit_status = std::process::Command::new("cargo")
            .env(
                "RUSTC_WRAPPER",
                std::env::current_exe().expect("env::current_exe() failed"),
            )
            .env(
                var_names::DEBUG_SNIPPET,
                matches.value_of("debug_snippet").unwrap_or(""),
            )
            .env(
                var_names::DIFF_CMD,
                if matches.is_present("diff") {
                    "diff"
                } else if matches.is_present("colordiff") {
                    "colordiff"
                } else {
                    ""
                },
            )
            .env(var_names::RULES, rules)
            .env(var_names::RULES_FILE, rules_file)
            .args(vec!["check"])
            .status()
            .expect("Failed to invoke cargo");
        // Return whatever exit code we got back from cargo or -1 if it was terminated by a signal.
        std::process::exit(cargo_exit_status.code().unwrap_or(-1));
    } else {
        panic!("Unknown mode");
    }
}

pub fn main() {
    let driver = RerastCompilerDriver::new(ArgBuilder::from_args(std::env::args().skip(1)));
    // If cargo is calling us to get compiler configuration or is compiling a dependent crate, then
    // just run the compiler normally.
    if driver.args().has_arg("--print=cfg") || driver.is_compiling_dependency() {
        let args: Vec<_> = std::env::args().skip(2).collect();
        std::process::Command::new("rustc")
            .args(args)
            .status()
            .expect("Failed to run rustc");
    } else if let (Ok(rules_filename), Ok(rules), Ok(debug_snippet)) = (
        std::env::var(var_names::RULES_FILE),
        std::env::var(var_names::RULES),
        std::env::var(var_names::DEBUG_SNIPPET),
    ) {
        let config = Config {
            debug_snippet: debug_snippet,
        };
        match driver.apply_rules_from_string_or_file(rules, &rules_filename, config) {
            Ok(output) => {
                if output.updated_files.is_empty() {
                    println!("No matches found");
                }
                for (filename, new_contents) in output.updated_files {
                    // rustfmt has a native diff built-in. If that were extracted into a separate
                    // crate, we could reuse that instead of calling out to an external diff.
                    if let Ok(diff_cmd) = std::env::var(var_names::DIFF_CMD) {
                        let mut diff_process = std::process::Command::new(diff_cmd)
                            .args(ArgBuilder::new().arg("-u").arg(filename).arg("-").build())
                            .stdin(std::process::Stdio::piped())
                            .spawn()
                            .expect("Failed to run diff command");
                        diff_process
                            .stdin
                            .as_mut()
                            .unwrap()
                            .write_all(new_contents.as_bytes())
                            .expect("Failed to pipe file contents to diff command");
                        diff_process
                            .wait()
                            .expect("Failed waiting for diff command");
                    } else {
                        println!("--- Updated {} ---", filename);
                        println!("{}", new_contents);
                    }
                }
            }
            Err(errors) => {
                eprintln!("{}", errors);
                std::process::exit(1);
            }
        }
    } else if let Some(arg1) = std::env::args().nth(1) {
        if arg1 == "rerast" {
            cargo_rerast();
        }
    }
}
<commit_msg>Fix matching in lib targets if the crate has already been built by running cargo clean<commit_after>// Copyright 2017 Google Inc.
//
// 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
//
//     https://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.

#![cfg_attr(feature = "clippy", feature(plugin))]
#![cfg_attr(feature = "clippy", plugin(clippy))]

extern crate clap;
extern crate json;
extern crate rerast;

use std::io::{self, Write};
use rerast::{ArgBuilder, Config, RerastCompilerDriver};

// Environment variables that we use to pass data from the outer invocation of cargo-rerast through
// to the inner invocation (where it acts as the rust compiler). Eventually we'll hopefully be able
// to query cargo as to what rustc commands are needed to build, then we can use those
// directly. Possibly we could just run "cargo check -v" and collect its output. Either way, we'd
// like to get rid of these as they add unnecessary complexity.
mod var_names {
    pub const RULES_FILE: &'static str = "RERAST_RULES_FILE";
    pub const RULES: &'static str = "RERAST_RULES";
    pub const DIFF_CMD: &'static str = "RERAST_DIFF_CMD";
    pub const DEBUG_SNIPPET: &'static str = "RERAST_DEBUG_SNIPPET";
}

// Queries cargo to find the name of the current crate, then runs cargo clean to
// clean up artifacts for that package (but not dependencies). This is necessary
// in order to ensure that all the files in the current crate actually get built
// when we run cargo check. Hopefully eventually there'll be a nicer way to
// integrate with cargo such that we won't need to do this.
fn clean_local_targets() -> io::Result<()> {
    let output = std::process::Command::new("cargo")
        .args(vec!["metadata", "--no-deps", "--format-version=1"])
        .stdout(std::process::Stdio::piped())
        .output()?;
    let metadata_str = std::str::from_utf8(output.stdout.as_slice()).map_err(|e| {
        io::Error::new(io::ErrorKind::InvalidData, e.to_string())
    })?;
    let parsed = json::parse(metadata_str).map_err(|e| {
        io::Error::new(io::ErrorKind::InvalidData, e.to_string())
    })?;
    for package in parsed["packages"].members() {
        if let Some(name) = package["name"].as_str() {
            std::process::Command::new("cargo")
                .args(vec!["clean", "--package", name])
                .status()?;
        }
    }
    Ok(())
}

fn cargo_rerast() {
    let matches = clap::App::new("Rerast")
        .version("0.1")
        .author("David Lattimore <dml@google.com>")
        .subcommand(
            clap::SubCommand::with_name("rerast")
                .about("Replace Rust code based on typed, syntactic patterns")
                .args_from_usage(
                    "--rules_file=[FILE] 'Path to a rule file'
                     --use_statements=[USE_STATEMENTS] 'Use statements required by rule'
                     --placeholders=[PLACEHOLDERS] 'e.g. <T>(o: option<T>)'
                     --replace=[PATTERN] 'Pattern to search for'
                     --with=[PATTERN] 'Replacement code'
                     --debug_snippet=[CODE_SNIPPET] 'A snippet of code that you think should \
                                                     match or list_all to list all checked \
                                                     snippets.'
                                      \
                     --diff 'Diff changes'
                     --colordiff 'Diff changes with colordiff'",
                ),
        )
        .get_matches();
    if let Some(matches) = matches.subcommand_matches("rerast") {
        // User ran "cargo rerast", invoke cargo build with us as "rustc"
        let rules = if let (Some(search), Some(replacement)) =
            (matches.value_of("replace"), matches.value_of("with"))
        {
            let mut placeholders = matches.value_of("placeholders").unwrap_or("").to_owned();
            if !placeholders.contains('(') {
                placeholders = "(".to_owned() + &placeholders + ")";
            }
            matches.value_of("use_statements").unwrap_or("").to_owned() + "pub fn rule"
                + &placeholders + "{replace!(" + search + " => " + replacement + ");}"
        } else {
            "".to_owned()
        };
        let rules_file = matches.value_of("rules_file").unwrap_or("");
        if rules_file.is_empty() == rules.is_empty() {
            eprintln!("Must specify either --rules_file or both of --search and --replacement");
            std::process::exit(1);
        }
        if matches.is_present("diff") == matches.is_present("colordiff") {
            eprintln!("Currently either --diff or --colordiff must be specified");
            std::process::exit(1);
        }
        // We run cargo with RUSTC_WRAPPER set, so that it then uses us as the compiler. Really what
        // we'd like is to run cargo, asking it to build all the crate dependencies, then tell us
        // the commandline(s) to build all the targets in the current crate. This is pretty
        // ugly. Clippy does something similar, so there probably isn't a better way at the moment.
        // Unfortunately this hack has several consequences. (1) We're invoked separately for each
        // target, which precludes, or would at least make it hard to produce a summary of all
        // matches. (2) Cargo won't run us if the crate has already been built, requiring that we
        // touch the main source file in order to force it to actually run us. (3) If our
        // dependencies haven't already been built, we'll try to look for matches in them - not
        // really what we want. (4) Cargo sometimes prints messages that are either not helpful or
        // outright confusing - e.g. asking the user to add --verbose.
        clean_local_targets().unwrap();
        let cargo_exit_status = std::process::Command::new("cargo")
            .env(
                "RUSTC_WRAPPER",
                std::env::current_exe().expect("env::current_exe() failed"),
            )
            .env(
                var_names::DEBUG_SNIPPET,
                matches.value_of("debug_snippet").unwrap_or(""),
            )
            .env(
                var_names::DIFF_CMD,
                if matches.is_present("diff") {
                    "diff"
                } else if matches.is_present("colordiff") {
                    "colordiff"
                } else {
                    ""
                },
            )
            .env(var_names::RULES, rules)
            .env(var_names::RULES_FILE, rules_file)
            .args(vec!["check"])
            .status()
            .expect("Failed to invoke cargo");
        // Return whatever exit code we got back from cargo or -1 if it was terminated by a signal.
        std::process::exit(cargo_exit_status.code().unwrap_or(-1));
    } else {
        panic!("Unknown mode");
    }
}

pub fn main() {
    let driver = RerastCompilerDriver::new(ArgBuilder::from_args(std::env::args().skip(1)));
    // If cargo is calling us to get compiler configuration or is compiling a dependent crate, then
    // just run the compiler normally.
    if driver.args().has_arg("--print=cfg") || driver.is_compiling_dependency() {
        let args: Vec<_> = std::env::args().skip(2).collect();
        std::process::Command::new("rustc")
            .args(args)
            .status()
            .expect("Failed to run rustc");
    } else if let (Ok(rules_filename), Ok(rules), Ok(debug_snippet)) = (
        std::env::var(var_names::RULES_FILE),
        std::env::var(var_names::RULES),
        std::env::var(var_names::DEBUG_SNIPPET),
    ) {
        let config = Config {
            debug_snippet: debug_snippet,
        };
        match driver.apply_rules_from_string_or_file(rules, &rules_filename, config) {
            Ok(output) => {
                if output.updated_files.is_empty() {
                    println!("No matches found");
                }
                for (filename, new_contents) in output.updated_files {
                    // rustfmt has a native diff built-in. If that were extracted into a separate
                    // crate, we could reuse that instead of calling out to an external diff.
                    if let Ok(diff_cmd) = std::env::var(var_names::DIFF_CMD) {
                        let mut diff_process = std::process::Command::new(diff_cmd)
                            .args(ArgBuilder::new().arg("-u").arg(filename).arg("-").build())
                            .stdin(std::process::Stdio::piped())
                            .spawn()
                            .expect("Failed to run diff command");
                        diff_process
                            .stdin
                            .as_mut()
                            .unwrap()
                            .write_all(new_contents.as_bytes())
                            .expect("Failed to pipe file contents to diff command");
                        diff_process
                            .wait()
                            .expect("Failed waiting for diff command");
                    } else {
                        println!("--- Updated {} ---", filename);
                        println!("{}", new_contents);
                    }
                }
            }
            Err(errors) => {
                eprintln!("{}", errors);
                std::process::exit(1);
            }
        }
    } else if let Some(arg1) = std::env::args().nth(1) {
        if arg1 == "rerast" {
            cargo_rerast();
        }
    }
}
<|endoftext|>


<commit_before>//! A small GPU compute shader demonstration.
//!
//! Here we use a compute shader to calculate the amplitude of `OSCILLATOR_COUNT` number of
//! oscillators. The oscillator amplitudes are then laid out across the screen using rectangles
//! with a gray value equal to the amplitude. Real-time interaction is demonstrated by providing
//! access to time, frequency (mouse `x`) and the number of oscillators via uniform data.

use nannou::prelude::*;
use std::sync::{Arc, Mutex};

struct Model {
    compute: Compute,
    oscillators: Arc<Mutex<Vec<f32>>>,
}

struct Compute {
    oscillator_buffer: wgpu::Buffer,
    oscillator_buffer_size: wgpu::BufferAddress,
    uniform_buffer: wgpu::Buffer,
    bind_group: wgpu::BindGroup,
    pipeline: wgpu::ComputePipeline,
}

#[repr(C)]
#[derive(Copy, Clone)]
pub struct Uniforms {
    time: f32,
    freq: f32,
    oscillator_count: u32,
}

const OSCILLATOR_COUNT: u32 = 128;

fn main() {
    nannou::app(model).update(update).run();
}

fn model(app: &App) -> Model {
    let w_id = app.new_window().size(1440, 512).view(view).build().unwrap();
    let window = app.window(w_id).unwrap();
    let device = window.swap_chain_device();

    // Create the compute shader module.
    let cs = include_bytes!("shaders/comp.spv");
    let cs_spirv = wgpu::read_spirv(std::io::Cursor::new(&cs[..])).unwrap();
    let cs_mod = device.create_shader_module(&cs_spirv);

    // Create the buffer that will store the result of our compute operation.
    let oscillator_buffer_size =
        (OSCILLATOR_COUNT as usize * std::mem::size_of::<f32>()) as wgpu::BufferAddress;
    let oscillator_buffer = device.create_buffer(&wgpu::BufferDescriptor {
        size: oscillator_buffer_size,
        usage: wgpu::BufferUsage::STORAGE
            | wgpu::BufferUsage::COPY_DST
            | wgpu::BufferUsage::COPY_SRC,
    });

    // Create the buffer that will store time.
    let uniforms = create_uniforms(app.time, app.mouse.x, window.rect());
    let uniform_buffer = device
        .create_buffer_mapped(1, wgpu::BufferUsage::UNIFORM | wgpu::BufferUsage::COPY_DST)
        .fill_from_slice(&[uniforms]);

    // Create the bind group and pipeline.
    let bind_group_layout = create_bind_group_layout(device);
    let bind_group = create_bind_group(
        device,
        &bind_group_layout,
        &oscillator_buffer,
        oscillator_buffer_size,
        &uniform_buffer,
    );
    let pipeline_layout = create_pipeline_layout(device, &bind_group_layout);
    let pipeline = create_compute_pipeline(device, &pipeline_layout, &cs_mod);

    let compute = Compute {
        oscillator_buffer,
        oscillator_buffer_size,
        uniform_buffer,
        bind_group,
        pipeline,
    };

    // The vector that we will write oscillator values to.
    let oscillators = Arc::new(Mutex::new(vec![0.0; OSCILLATOR_COUNT as usize]));

    Model {
        compute,
        oscillators,
    }
}

fn update(app: &App, model: &mut Model, _update: Update) {
    let window = app.main_window();
    let device = window.swap_chain_device();
    let win_rect = window.rect();
    let compute = &mut model.compute;

    // The buffer into which we'll read some data.
    let read_buffer = device.create_buffer(&wgpu::BufferDescriptor {
        size: compute.oscillator_buffer_size,
        usage: wgpu::BufferUsage::MAP_READ
            | wgpu::BufferUsage::COPY_DST
            | wgpu::BufferUsage::COPY_SRC,
    });

    // An update for the uniform buffer with the current time.
    let uniforms = create_uniforms(app.time, app.mouse.x, win_rect);
    let uniforms_size = std::mem::size_of::<Uniforms>() as wgpu::BufferAddress;
    let new_uniform_buffer = device
        .create_buffer_mapped(1, wgpu::BufferUsage::COPY_SRC)
        .fill_from_slice(&[uniforms]);

    // The encoder we'll use to encode the compute pass.
    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
    encoder.copy_buffer_to_buffer(
        &new_uniform_buffer,
        0,
        &compute.uniform_buffer,
        0,
        uniforms_size,
    );
    {
        let mut cpass = encoder.begin_compute_pass();
        cpass.set_pipeline(&compute.pipeline);
        cpass.set_bind_group(0, &compute.bind_group, &[]);
        cpass.dispatch(OSCILLATOR_COUNT as u32, 1, 1);
    }
    encoder.copy_buffer_to_buffer(
        &compute.oscillator_buffer,
        0,
        &read_buffer,
        0,
        compute.oscillator_buffer_size,
    );

    // Submit the compute pass to the device's queue.
    window
        .swap_chain_queue()
        .lock()
        .unwrap()
        .submit(&[encoder.finish()]);

    // Register a callback for reading the result of the compute pass.
    let oscillators = model.oscillators.clone();
    read_buffer.map_read_async(
        0,
        compute.oscillator_buffer_size,
        move |result: wgpu::BufferMapAsyncResult<&[f32]>| {
            if let Ok(mapping) = result {
                if let Ok(mut oscillators) = oscillators.lock() {
                    oscillators.copy_from_slice(&mapping.data);
                }
            }
        },
    );

    // Check for resource cleanups and mapping callbacks.
    device.poll(false);
}

fn view(app: &App, model: &Model, frame: &Frame) {
    frame.clear(BLACK);
    let draw = app.draw();
    let window = app.window(frame.window_id()).unwrap();
    let rect = window.rect();

    if let Ok(oscillators) = model.oscillators.lock() {
        let w = rect.w() / OSCILLATOR_COUNT as f32;
        let h = rect.h();
        let half_w = w * 0.5;
        for (i, &osc) in oscillators.iter().enumerate() {
            let x = half_w + map_range(i as u32, 0, OSCILLATOR_COUNT, rect.left(), rect.right());
            draw.rect().w_h(w, h).x(x).color(gray(osc));
        }
    }

    draw.to_frame(app, frame).unwrap();
}

fn create_uniforms(time: f32, mouse_x: f32, win_rect: geom::Rect) -> Uniforms {
    let freq = map_range(
        mouse_x,
        win_rect.left(),
        win_rect.right(),
        0.0,
        win_rect.w(),
    );
    let oscillator_count = OSCILLATOR_COUNT;
    Uniforms {
        time,
        freq,
        oscillator_count,
    }
}

fn create_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
    let oscillator_buffer_binding = wgpu::BindGroupLayoutBinding {
        binding: 0,
        visibility: wgpu::ShaderStage::COMPUTE,
        ty: wgpu::BindingType::StorageBuffer {
            dynamic: false,
            readonly: false,
        },
    };
    let uniforms_binding = wgpu::BindGroupLayoutBinding {
        binding: 1,
        visibility: wgpu::ShaderStage::COMPUTE,
        ty: wgpu::BindingType::UniformBuffer { dynamic: false },
    };
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        bindings: &[oscillator_buffer_binding, uniforms_binding],
    })
}

fn create_bind_group(
    device: &wgpu::Device,
    layout: &wgpu::BindGroupLayout,
    oscillator_buffer: &wgpu::Buffer,
    oscillator_buffer_size: wgpu::BufferAddress,
    uniform_buffer: &wgpu::Buffer,
) -> wgpu::BindGroup {
    let oscillator_buffer_binding = wgpu::Binding {
        binding: 0,
        resource: wgpu::BindingResource::Buffer {
            buffer: &oscillator_buffer,
            range: 0..oscillator_buffer_size,
        },
    };
    let uniforms_binding = wgpu::Binding {
        binding: 1,
        resource: wgpu::BindingResource::Buffer {
            buffer: &uniform_buffer,
            range: 0..std::mem::size_of::<Uniforms>() as wgpu::BufferAddress,
        },
    };
    let bindings = &[oscillator_buffer_binding, uniforms_binding];
    device.create_bind_group(&wgpu::BindGroupDescriptor { layout, bindings })
}

fn create_pipeline_layout(
    device: &wgpu::Device,
    bind_group_layout: &wgpu::BindGroupLayout,
) -> wgpu::PipelineLayout {
    device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        bind_group_layouts: &[&bind_group_layout],
    })
}

fn create_compute_pipeline(
    device: &wgpu::Device,
    layout: &wgpu::PipelineLayout,
    cs_mod: &wgpu::ShaderModule,
) -> wgpu::ComputePipeline {
    let compute_stage = wgpu::ProgrammableStageDescriptor {
        module: &cs_mod,
        entry_point: "main",
    };
    let desc = wgpu::ComputePipelineDescriptor {
        layout,
        compute_stage,
    };
    device.create_compute_pipeline(&desc)
}
<commit_msg>Add more details about poll to compute shader example<commit_after>//! A small GPU compute shader demonstration.
//!
//! Here we use a compute shader to calculate the amplitude of `OSCILLATOR_COUNT` number of
//! oscillators. The oscillator amplitudes are then laid out across the screen using rectangles
//! with a gray value equal to the amplitude. Real-time interaction is demonstrated by providing
//! access to time, frequency (mouse `x`) and the number of oscillators via uniform data.

use nannou::prelude::*;
use std::sync::{Arc, Mutex};

struct Model {
    compute: Compute,
    oscillators: Arc<Mutex<Vec<f32>>>,
}

struct Compute {
    oscillator_buffer: wgpu::Buffer,
    oscillator_buffer_size: wgpu::BufferAddress,
    uniform_buffer: wgpu::Buffer,
    bind_group: wgpu::BindGroup,
    pipeline: wgpu::ComputePipeline,
}

#[repr(C)]
#[derive(Copy, Clone)]
pub struct Uniforms {
    time: f32,
    freq: f32,
    oscillator_count: u32,
}

const OSCILLATOR_COUNT: u32 = 128;

fn main() {
    nannou::app(model).update(update).run();
}

fn model(app: &App) -> Model {
    let w_id = app.new_window().size(1440, 512).view(view).build().unwrap();
    let window = app.window(w_id).unwrap();
    let device = window.swap_chain_device();

    // Create the compute shader module.
    let cs = include_bytes!("shaders/comp.spv");
    let cs_spirv = wgpu::read_spirv(std::io::Cursor::new(&cs[..])).unwrap();
    let cs_mod = device.create_shader_module(&cs_spirv);

    // Create the buffer that will store the result of our compute operation.
    let oscillator_buffer_size =
        (OSCILLATOR_COUNT as usize * std::mem::size_of::<f32>()) as wgpu::BufferAddress;
    let oscillator_buffer = device.create_buffer(&wgpu::BufferDescriptor {
        size: oscillator_buffer_size,
        usage: wgpu::BufferUsage::STORAGE
            | wgpu::BufferUsage::COPY_DST
            | wgpu::BufferUsage::COPY_SRC,
    });

    // Create the buffer that will store time.
    let uniforms = create_uniforms(app.time, app.mouse.x, window.rect());
    let uniform_buffer = device
        .create_buffer_mapped(1, wgpu::BufferUsage::UNIFORM | wgpu::BufferUsage::COPY_DST)
        .fill_from_slice(&[uniforms]);

    // Create the bind group and pipeline.
    let bind_group_layout = create_bind_group_layout(device);
    let bind_group = create_bind_group(
        device,
        &bind_group_layout,
        &oscillator_buffer,
        oscillator_buffer_size,
        &uniform_buffer,
    );
    let pipeline_layout = create_pipeline_layout(device, &bind_group_layout);
    let pipeline = create_compute_pipeline(device, &pipeline_layout, &cs_mod);

    let compute = Compute {
        oscillator_buffer,
        oscillator_buffer_size,
        uniform_buffer,
        bind_group,
        pipeline,
    };

    // The vector that we will write oscillator values to.
    let oscillators = Arc::new(Mutex::new(vec![0.0; OSCILLATOR_COUNT as usize]));

    Model {
        compute,
        oscillators,
    }
}

fn update(app: &App, model: &mut Model, _update: Update) {
    let window = app.main_window();
    let device = window.swap_chain_device();
    let win_rect = window.rect();
    let compute = &mut model.compute;

    // The buffer into which we'll read some data.
    let read_buffer = device.create_buffer(&wgpu::BufferDescriptor {
        size: compute.oscillator_buffer_size,
        usage: wgpu::BufferUsage::MAP_READ
            | wgpu::BufferUsage::COPY_DST
            | wgpu::BufferUsage::COPY_SRC,
    });

    // An update for the uniform buffer with the current time.
    let uniforms = create_uniforms(app.time, app.mouse.x, win_rect);
    let uniforms_size = std::mem::size_of::<Uniforms>() as wgpu::BufferAddress;
    let new_uniform_buffer = device
        .create_buffer_mapped(1, wgpu::BufferUsage::COPY_SRC)
        .fill_from_slice(&[uniforms]);

    // The encoder we'll use to encode the compute pass.
    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
    encoder.copy_buffer_to_buffer(
        &new_uniform_buffer,
        0,
        &compute.uniform_buffer,
        0,
        uniforms_size,
    );
    {
        let mut cpass = encoder.begin_compute_pass();
        cpass.set_pipeline(&compute.pipeline);
        cpass.set_bind_group(0, &compute.bind_group, &[]);
        cpass.dispatch(OSCILLATOR_COUNT as u32, 1, 1);
    }
    encoder.copy_buffer_to_buffer(
        &compute.oscillator_buffer,
        0,
        &read_buffer,
        0,
        compute.oscillator_buffer_size,
    );

    // Submit the compute pass to the device's queue.
    window
        .swap_chain_queue()
        .lock()
        .unwrap()
        .submit(&[encoder.finish()]);

    // Register a callback for reading the result of the compute pass.
    let oscillators = model.oscillators.clone();
    read_buffer.map_read_async(
        0,
        compute.oscillator_buffer_size,
        move |result: wgpu::BufferMapAsyncResult<&[f32]>| {
            if let Ok(mapping) = result {
                if let Ok(mut oscillators) = oscillators.lock() {
                    oscillators.copy_from_slice(&mapping.data);
                }
            }
        },
    );

    // Check for resource cleanups and mapping callbacks.
    //
    // Note that this is not strictly necessary in our case, as the device we are using already
    // gets polled when nannou submits the command buffer for drawing and presentation after `view`
    // completes. If we were to use a standalone device to create our buffer and perform our
    // compute (rather than the device requested during window creation), calling `poll` regularly
    // would be a must.
    device.poll(false);
}

fn view(app: &App, model: &Model, frame: &Frame) {
    frame.clear(BLACK);
    let draw = app.draw();
    let window = app.window(frame.window_id()).unwrap();
    let rect = window.rect();

    if let Ok(oscillators) = model.oscillators.lock() {
        let w = rect.w() / OSCILLATOR_COUNT as f32;
        let h = rect.h();
        let half_w = w * 0.5;
        for (i, &osc) in oscillators.iter().enumerate() {
            let x = half_w + map_range(i as u32, 0, OSCILLATOR_COUNT, rect.left(), rect.right());
            draw.rect().w_h(w, h).x(x).color(gray(osc));
        }
    }

    draw.to_frame(app, frame).unwrap();
}

fn create_uniforms(time: f32, mouse_x: f32, win_rect: geom::Rect) -> Uniforms {
    let freq = map_range(
        mouse_x,
        win_rect.left(),
        win_rect.right(),
        0.0,
        win_rect.w(),
    );
    let oscillator_count = OSCILLATOR_COUNT;
    Uniforms {
        time,
        freq,
        oscillator_count,
    }
}

fn create_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
    let oscillator_buffer_binding = wgpu::BindGroupLayoutBinding {
        binding: 0,
        visibility: wgpu::ShaderStage::COMPUTE,
        ty: wgpu::BindingType::StorageBuffer {
            dynamic: false,
            readonly: false,
        },
    };
    let uniforms_binding = wgpu::BindGroupLayoutBinding {
        binding: 1,
        visibility: wgpu::ShaderStage::COMPUTE,
        ty: wgpu::BindingType::UniformBuffer { dynamic: false },
    };
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        bindings: &[oscillator_buffer_binding, uniforms_binding],
    })
}

fn create_bind_group(
    device: &wgpu::Device,
    layout: &wgpu::BindGroupLayout,
    oscillator_buffer: &wgpu::Buffer,
    oscillator_buffer_size: wgpu::BufferAddress,
    uniform_buffer: &wgpu::Buffer,
) -> wgpu::BindGroup {
    let oscillator_buffer_binding = wgpu::Binding {
        binding: 0,
        resource: wgpu::BindingResource::Buffer {
            buffer: &oscillator_buffer,
            range: 0..oscillator_buffer_size,
        },
    };
    let uniforms_binding = wgpu::Binding {
        binding: 1,
        resource: wgpu::BindingResource::Buffer {
            buffer: &uniform_buffer,
            range: 0..std::mem::size_of::<Uniforms>() as wgpu::BufferAddress,
        },
    };
    let bindings = &[oscillator_buffer_binding, uniforms_binding];
    device.create_bind_group(&wgpu::BindGroupDescriptor { layout, bindings })
}

fn create_pipeline_layout(
    device: &wgpu::Device,
    bind_group_layout: &wgpu::BindGroupLayout,
) -> wgpu::PipelineLayout {
    device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        bind_group_layouts: &[&bind_group_layout],
    })
}

fn create_compute_pipeline(
    device: &wgpu::Device,
    layout: &wgpu::PipelineLayout,
    cs_mod: &wgpu::ShaderModule,
) -> wgpu::ComputePipeline {
    let compute_stage = wgpu::ProgrammableStageDescriptor {
        module: &cs_mod,
        entry_point: "main",
    };
    let desc = wgpu::ComputePipelineDescriptor {
        layout,
        compute_stage,
    };
    device.create_compute_pipeline(&desc)
}
<|endoftext|>