<commit_before>#![feature(plugin)]
#![plugin(clippy)]

extern crate ansi_term;
extern crate clap;
extern crate env_logger;
#[macro_use]
extern crate log;
extern crate nix;
extern crate protobuf;
extern crate record_query;

use std::env;
use std::io;
use std::path;

use record_query as rq;

const ABOUT: &'static str = r#"
A tool for manipulating data records.  Similar in spirit to awk(1),
but on steroids.  Inspired by jq(1), but supports more record formats
and operations.

Records are read from stdin, processed, and written to stdout.  The
tool accepts a query in the custom rq query language as its main
command-line arguments.

Flags can be used to control the input and output formats of the tool.
By default, the input and output format is JSON.

See 'man rq' for in-depth documentation."#;

const AUTHOR: &'static str = "David Flemström <david.flemstrom@gmail.com>";

const INPUT_FORMAT_GROUP: &'static str = "input-format";
const OUTPUT_FORMAT_GROUP: &'static str = "output-format";

const PROTOBUF_CMD: &'static str = "protobuf";
const PROTOBUF_ADD_CMD: &'static str = "add";
const PROTOBUF_ADD_INPUT_ARG: &'static str = "protobuf-add-input";

const INPUT_JSON_ARG: &'static str = "input-json";
const OUTPUT_JSON_ARG: &'static str = "output-json";
const INPUT_PROTOBUF_ARG: &'static str = "input-protobuf";
const OUTPUT_PROTOBUF_ARG: &'static str = "output-protobuf";

const QUERY_ARG: &'static str = "query";
const VERBOSE_ARG: &'static str = "verbose";
const QUIET_ARG: &'static str = "quiet";

fn main() {
    use std::io::Read;

    let matches = match_args();

    setup_log(matches.occurrences_of(VERBOSE_ARG),
              matches.is_present(QUIET_ARG));

    let paths = rq::config::Paths::new().unwrap();

    if let Some(matches) = matches.subcommand_matches(PROTOBUF_CMD) {
        if let Some(matches) = matches.subcommand_matches(PROTOBUF_ADD_CMD) {
            let input = matches.value_of(PROTOBUF_ADD_INPUT_ARG).unwrap();
            rq::proto_index::add_file(&paths, path::Path::new(input)).unwrap();
        }
    } else {
        let query = rq::query::Query::parse(matches.value_of(QUERY_ARG).unwrap());

        let stdin = io::stdin();
        let mut input = stdin.lock();

        if matches.is_present(INPUT_PROTOBUF_ARG) {
            let name = matches.value_of(INPUT_PROTOBUF_ARG).unwrap();
            debug!("Input is protobuf with argument {}", name);

            let descriptors_proto = rq::proto_index::compile_descriptor_set(&paths).unwrap();
            let descriptors =
                rq::value::protobuf::descriptor::Descriptors::from_protobuf(&descriptors_proto);
            let stream = protobuf::CodedInputStream::new(&mut input);
            let values = rq::value::protobuf::ProtobufValues::new(descriptors,
                                                                  name.to_owned(),
                                                                  stream);
            run(values, query).unwrap_or_else(|e| error!("{:?}", e));
        } else {
            run(rq::value::json::JsonValues::new(input.bytes()), query)
                .unwrap_or_else(|e| error!("{:?}", e));
        }
    }
}

fn setup_log(verbose_level: u64, quiet: bool) {
    use ansi_term::ANSIStrings;
    use ansi_term::Colour;
    use ansi_term::Style;
    use log::LogLevel;
    use log::LogLevelFilter;
    use nix::unistd;
    use nix::sys::ioctl;

    let normal = Style::new();

    let format = move |record: &log::LogRecord| {
        if unistd::isatty(ioctl::libc::STDERR_FILENO).unwrap_or(false) {
            let (front, back) = match record.level() {
                LogLevel::Error => (Colour::Red.normal(), Colour::Red.dimmed()),
                LogLevel::Warn => (Colour::Yellow.normal(), Colour::Yellow.dimmed()),
                LogLevel::Info => (Colour::Blue.normal(), Colour::Blue.dimmed()),
                LogLevel::Debug => (Colour::Purple.normal(), Colour::Purple.dimmed()),
                LogLevel::Trace => (Colour::White.dimmed(), Colour::Black.normal()),
            };

            let strings = &[back.paint("["),
                            front.paint(format!("{}", record.level())),
                            back.paint("]"),
                            normal.paint(" "),
                            back.paint("["),
                            front.paint(record.location().module_path()),
                            back.paint("]"),
                            normal.paint(" "),
                            back.paint(format!("{}", record.args()))];

            format!("{}", ANSIStrings(strings))
        } else {
            format!("[{}] [{}] {}",
                    record.level(),
                    record.location().module_path(),
                    record.args())
        }
    };

    let mut builder = env_logger::LogBuilder::new();

    let filter = if quiet {
        LogLevelFilter::Off
    } else {
        match verbose_level {
            0 => LogLevelFilter::Info,
            1 => LogLevelFilter::Debug,
            _ => LogLevelFilter::Trace,
        }
    };

    builder.format(format).filter(None, filter);

    if let Ok(spec) = env::var("RUST_LOG") {
        builder.parse(&spec);
    }

    builder.init().unwrap();
}

#[cfg_attr(rustfmt, rustfmt_skip)]
fn match_args<'a>() -> clap::ArgMatches<'a> {
    use clap::AppSettings;
    use clap::Arg;
    use clap::ArgGroup;
    use clap::SubCommand;

    clap::App::new("rq - Record query")
        .version(env!("CARGO_PKG_VERSION"))
        .author(AUTHOR)
        .about(ABOUT)
        .setting(AppSettings::SubcommandsNegateReqs)
        .setting(AppSettings::ArgRequiredElseHelp)
        .setting(AppSettings::GlobalVersion)
        .setting(AppSettings::VersionlessSubcommands)
        .setting(AppSettings::UnifiedHelpMessage)
        .setting(AppSettings::ColoredHelp)
        .group(ArgGroup::with_name(INPUT_FORMAT_GROUP))
        .group(ArgGroup::with_name(OUTPUT_FORMAT_GROUP))
        .arg(Arg::with_name(INPUT_JSON_ARG)
             .group(INPUT_FORMAT_GROUP)
             .short("j")
             .long("input-json")
             .help("Input is white-space separated JSON values."))
        .arg(Arg::with_name(OUTPUT_JSON_ARG)
             .group(OUTPUT_FORMAT_GROUP)
             .short("J")
             .long("output-json")
             .help("Output should be formatted as JSON values."))
        .arg(Arg::with_name(INPUT_PROTOBUF_ARG)
             .group(INPUT_FORMAT_GROUP)
             .short("p")
             .long("input-protobuf")
             .takes_value(true)
             .value_name(".package.MessageType")
             .next_line_help(true)
             .help("Input is a single protocol buffer object.  The argument refers to the fully \
                    qualified name of the message type (including the leading '.')"))
        .arg(Arg::with_name(OUTPUT_PROTOBUF_ARG)
             .group(OUTPUT_FORMAT_GROUP)
             .short("P")
             .long("output-protobuf")
             .takes_value(true)
             .value_name(".package.MessageType")
             .next_line_help(true)
             .help("Output should be formatted as protocol buffer objects.  The argument refers \
                    to the fully qualified name of the message type (including the leading '.'), \
                    but if it is omitted and -p was used, the input schema is used instead."))
        .arg(Arg::with_name(QUERY_ARG)
             .required(true)
             .help("The query to apply."))
        .arg(Arg::with_name(VERBOSE_ARG)
             .short("v")
             .long("verbose")
             .multiple(true)
             .help("Increase the logging verbosity"))
        .arg(Arg::with_name(QUIET_ARG)
             .short("q")
             .long("quiet")
             .help("Disable default logging (the RUST_LOG env var is still respected)"))
        .subcommand(SubCommand::with_name(PROTOBUF_CMD)
                    .about("Control protobuf configuration and data")
                    .author(AUTHOR)
                    .subcommand(SubCommand::with_name(PROTOBUF_ADD_CMD)
                                .about("Add a schema file to the rq registry")
                                .author(AUTHOR)
                                .arg(Arg::with_name(PROTOBUF_ADD_INPUT_ARG)
                                     .required(true)
                                     .value_name("schema")
                                     .help("The path to a .proto file to add to the rq registry"))))
        .get_matches()
}

fn run<Iter>(input: Iter, query: rq::query::Query) -> rq::error::Result<()>
    where Iter: Iterator<Item = rq::error::Result<rq::value::Value>>
{
    use std::io::Write;
    let mut stdout = io::stdout();
    let context = rq::query::Context::new();

    debug!("Starting input consumption");

    for value in input {
        let value = try!(value);
        debug!("Consuming an input value");
        try!(query.evaluate(&context, value).to_json(&mut stdout));
        try!(stdout.write(&[10]));
        try!(stdout.flush());
    }
    Ok(())
}
<commit_msg>Source files should not be executable<commit_after>#![feature(plugin)]
#![plugin(clippy)]

extern crate ansi_term;
extern crate clap;
extern crate env_logger;
#[macro_use]
extern crate log;
extern crate nix;
extern crate protobuf;
extern crate record_query;

use std::env;
use std::io;
use std::path;

use record_query as rq;

const ABOUT: &'static str = r#"
A tool for manipulating data records.  Similar in spirit to awk(1),
but on steroids.  Inspired by jq(1), but supports more record formats
and operations.

Records are read from stdin, processed, and written to stdout.  The
tool accepts a query in the custom rq query language as its main
command-line arguments.

Flags can be used to control the input and output formats of the tool.
By default, the input and output format is JSON.

See 'man rq' for in-depth documentation."#;

const AUTHOR: &'static str = "David Flemström <david.flemstrom@gmail.com>";

const INPUT_FORMAT_GROUP: &'static str = "input-format";
const OUTPUT_FORMAT_GROUP: &'static str = "output-format";

const PROTOBUF_CMD: &'static str = "protobuf";
const PROTOBUF_ADD_CMD: &'static str = "add";
const PROTOBUF_ADD_INPUT_ARG: &'static str = "protobuf-add-input";

const INPUT_JSON_ARG: &'static str = "input-json";
const OUTPUT_JSON_ARG: &'static str = "output-json";
const INPUT_PROTOBUF_ARG: &'static str = "input-protobuf";
const OUTPUT_PROTOBUF_ARG: &'static str = "output-protobuf";

const QUERY_ARG: &'static str = "query";
const VERBOSE_ARG: &'static str = "verbose";
const QUIET_ARG: &'static str = "quiet";

fn main() {
    use std::io::Read;

    let matches = match_args();

    setup_log(matches.occurrences_of(VERBOSE_ARG),
              matches.is_present(QUIET_ARG));

    let paths = rq::config::Paths::new().unwrap();

    if let Some(matches) = matches.subcommand_matches(PROTOBUF_CMD) {
        if let Some(matches) = matches.subcommand_matches(PROTOBUF_ADD_CMD) {
            let input = matches.value_of(PROTOBUF_ADD_INPUT_ARG).unwrap();
            rq::proto_index::add_file(&paths, path::Path::new(input)).unwrap();
        }
    } else {
        let query = rq::query::Query::parse(matches.value_of(QUERY_ARG).unwrap());

        let stdin = io::stdin();
        let mut input = stdin.lock();

        if matches.is_present(INPUT_PROTOBUF_ARG) {
            let name = matches.value_of(INPUT_PROTOBUF_ARG).unwrap();
            debug!("Input is protobuf with argument {}", name);

            let descriptors_proto = rq::proto_index::compile_descriptor_set(&paths).unwrap();
            let descriptors =
                rq::value::protobuf::descriptor::Descriptors::from_protobuf(&descriptors_proto);
            let stream = protobuf::CodedInputStream::new(&mut input);
            let values = rq::value::protobuf::ProtobufValues::new(descriptors,
                                                                  name.to_owned(),
                                                                  stream);
            run(values, query).unwrap_or_else(|e| error!("{:?}", e));
        } else {
            run(rq::value::json::JsonValues::new(input.bytes()), query)
                .unwrap_or_else(|e| error!("{:?}", e));
        }
    }
}

fn setup_log(verbose_level: u64, quiet: bool) {
    use ansi_term::ANSIStrings;
    use ansi_term::Colour;
    use ansi_term::Style;
    use log::LogLevel;
    use log::LogLevelFilter;
    use nix::unistd;
    use nix::sys::ioctl;

    let normal = Style::new();

    let format = move |record: &log::LogRecord| {
        if unistd::isatty(ioctl::libc::STDERR_FILENO).unwrap_or(false) {
            let (front, back) = match record.level() {
                LogLevel::Error => (Colour::Red.normal(), Colour::Red.dimmed()),
                LogLevel::Warn => (Colour::Yellow.normal(), Colour::Yellow.dimmed()),
                LogLevel::Info => (Colour::Blue.normal(), Colour::Blue.dimmed()),
                LogLevel::Debug => (Colour::Purple.normal(), Colour::Purple.dimmed()),
                LogLevel::Trace => (Colour::White.dimmed(), Colour::Black.normal()),
            };

            let strings = &[back.paint("["),
                            front.paint(format!("{}", record.level())),
                            back.paint("]"),
                            normal.paint(" "),
                            back.paint("["),
                            front.paint(record.location().module_path()),
                            back.paint("]"),
                            normal.paint(" "),
                            back.paint(format!("{}", record.args()))];

            format!("{}", ANSIStrings(strings))
        } else {
            format!("[{}] [{}] {}",
                    record.level(),
                    record.location().module_path(),
                    record.args())
        }
    };

    let mut builder = env_logger::LogBuilder::new();

    let filter = if quiet {
        LogLevelFilter::Off
    } else {
        match verbose_level {
            0 => LogLevelFilter::Info,
            1 => LogLevelFilter::Debug,
            _ => LogLevelFilter::Trace,
        }
    };

    builder.format(format).filter(None, filter);

    if let Ok(spec) = env::var("RUST_LOG") {
        builder.parse(&spec);
    }

    builder.init().unwrap();
}

#[cfg_attr(rustfmt, rustfmt_skip)]
fn match_args<'a>() -> clap::ArgMatches<'a> {
    use clap::AppSettings;
    use clap::Arg;
    use clap::ArgGroup;
    use clap::SubCommand;

    clap::App::new("rq - Record query")
        .version(env!("CARGO_PKG_VERSION"))
        .author(AUTHOR)
        .about(ABOUT)
        .setting(AppSettings::SubcommandsNegateReqs)
        .setting(AppSettings::ArgRequiredElseHelp)
        .setting(AppSettings::GlobalVersion)
        .setting(AppSettings::VersionlessSubcommands)
        .setting(AppSettings::UnifiedHelpMessage)
        .setting(AppSettings::ColoredHelp)
        .group(ArgGroup::with_name(INPUT_FORMAT_GROUP))
        .group(ArgGroup::with_name(OUTPUT_FORMAT_GROUP))
        .arg(Arg::with_name(INPUT_JSON_ARG)
             .group(INPUT_FORMAT_GROUP)
             .short("j")
             .long("input-json")
             .help("Input is white-space separated JSON values."))
        .arg(Arg::with_name(OUTPUT_JSON_ARG)
             .group(OUTPUT_FORMAT_GROUP)
             .short("J")
             .long("output-json")
             .help("Output should be formatted as JSON values."))
        .arg(Arg::with_name(INPUT_PROTOBUF_ARG)
             .group(INPUT_FORMAT_GROUP)
             .short("p")
             .long("input-protobuf")
             .takes_value(true)
             .value_name(".package.MessageType")
             .next_line_help(true)
             .help("Input is a single protocol buffer object.  The argument refers to the fully \
                    qualified name of the message type (including the leading '.')"))
        .arg(Arg::with_name(OUTPUT_PROTOBUF_ARG)
             .group(OUTPUT_FORMAT_GROUP)
             .short("P")
             .long("output-protobuf")
             .takes_value(true)
             .value_name(".package.MessageType")
             .next_line_help(true)
             .help("Output should be formatted as protocol buffer objects.  The argument refers \
                    to the fully qualified name of the message type (including the leading '.'), \
                    but if it is omitted and -p was used, the input schema is used instead."))
        .arg(Arg::with_name(QUERY_ARG)
             .required(true)
             .help("The query to apply."))
        .arg(Arg::with_name(VERBOSE_ARG)
             .short("v")
             .long("verbose")
             .multiple(true)
             .help("Increase the logging verbosity"))
        .arg(Arg::with_name(QUIET_ARG)
             .short("q")
             .long("quiet")
             .help("Disable default logging (the RUST_LOG env var is still respected)"))
        .subcommand(SubCommand::with_name(PROTOBUF_CMD)
                    .about("Control protobuf configuration and data")
                    .author(AUTHOR)
                    .subcommand(SubCommand::with_name(PROTOBUF_ADD_CMD)
                                .about("Add a schema file to the rq registry")
                                .author(AUTHOR)
                                .arg(Arg::with_name(PROTOBUF_ADD_INPUT_ARG)
                                     .required(true)
                                     .value_name("schema")
                                     .help("The path to a .proto file to add to the rq registry"))))
        .get_matches()
}

fn run<Iter>(input: Iter, query: rq::query::Query) -> rq::error::Result<()>
    where Iter: Iterator<Item = rq::error::Result<rq::value::Value>>
{
    use std::io::Write;
    let mut stdout = io::stdout();
    let context = rq::query::Context::new();

    debug!("Starting input consumption");

    for value in input {
        let value = try!(value);
        debug!("Consuming an input value");
        try!(query.evaluate(&context, value).to_json(&mut stdout));
        try!(stdout.write(&[10]));
        try!(stdout.flush());
    }
    Ok(())
}
<|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.

/// Formatting of chained expressions, i.e. expressions which are chained by
/// dots: struct and enum field access, method calls, and try shorthand (?).
///
/// Instead of walking these subexpressions one-by-one, as is our usual strategy
/// for expression formatting, we collect maximal sequences of these expressions
/// and handle them simultaneously.
///
/// Whenever possible, the entire chain is put on a single line. If that fails,
/// we put each subexpression on a separate, much like the (default) function
/// argument function argument strategy.
///
/// Depends on config options: `chain_indent` is the indent to use for
/// blocks in the parent/root/base of the chain (and the rest of the chain's
/// alignment).
/// E.g., `let foo = { aaaa; bbb; ccc }.bar.baz();`, we would layout for the
/// following values of `chain_indent`:
/// Visual:
/// ```
/// let foo = {
///               aaaa;
///               bbb;
///               ccc
///           }
///           .bar
///           .baz();
/// ```
/// Inherit:
/// ```
/// let foo = {
///     aaaa;
///     bbb;
///     ccc
/// }
/// .bar
/// .baz();
/// ```
/// Tabbed:
/// ```
/// let foo = {
///         aaaa;
///         bbb;
///         ccc
///     }
///     .bar
///     .baz();
/// ```
///
/// If the first item in the chain is a block expression, we align the dots with
/// the braces.
/// Visual:
/// ```
/// let a = foo.bar
///            .baz()
///            .qux
/// ```
/// Inherit:
/// ```
/// let a = foo.bar
/// .baz()
/// .qux
/// ```
/// Tabbed:
/// ```
/// let a = foo.bar
///     .baz()
///     .qux
/// ```

use Shape;
use rewrite::{Rewrite, RewriteContext};
use utils::{wrap_str, first_line_width, last_line_width};
use expr::rewrite_call;
use config::IndentStyle;
use macros::convert_try_mac;

use std::iter;
use syntax::{ast, ptr};
use syntax::codemap::{mk_sp, Span};

pub fn rewrite_chain(expr: &ast::Expr, context: &RewriteContext, shape: Shape) -> Option<String> {
    debug!("rewrite_chain {:?}", shape);
    let total_span = expr.span;
    let (parent, subexpr_list) = make_subexpr_list(expr, context);

    // Bail out if the chain is just try sugar, i.e., an expression followed by
    // any number of `?`s.
    if chain_only_try(&subexpr_list) {
        return rewrite_try(&parent, subexpr_list.len(), context, shape);
    }

    // Parent is the first item in the chain, e.g., `foo` in `foo.bar.baz()`.
    let mut parent_shape = shape;
    if is_block_expr(&parent, "\n") {
        parent_shape = chain_indent(context, shape);
    }
    let parent_rewrite = try_opt!(parent.rewrite(context, parent_shape));

    // Decide how to layout the rest of the chain. `extend` is true if we can
    // put the first non-parent item on the same line as the parent.
    let (nested_shape, extend) = if !parent_rewrite.contains('\n') && is_continuable(&parent) {
        let nested_shape = if let ast::ExprKind::Try(..) = subexpr_list.last().unwrap().node {
            parent_shape.block_indent(context.config.tab_spaces)
        } else {
            chain_indent(context, shape.add_offset(parent_rewrite.len()))
        };
        (nested_shape, true)
    } else if is_block_expr(&parent, &parent_rewrite) {
        // The parent is a block, so align the rest of the chain with the closing
        // brace.
        (parent_shape, false)
    } else if parent_rewrite.contains('\n') {
        (chain_indent(context,
                      parent_shape.block_indent(context.config.tab_spaces)),
         false)
    } else {
        (shape.block_indent(context.config.tab_spaces), false)
    };

    let max_width = try_opt!((shape.width + shape.indent.width() + shape.offset)
                                 .checked_sub(nested_shape.indent.width() +
                                              nested_shape.offset));

    let other_child_shape = Shape {
        width: max_width,
        ..nested_shape
    };
    let first_child_shape = if extend {
        let mut shape = try_opt!(parent_shape.shrink_left(last_line_width(&parent_rewrite)));
        match context.config.chain_indent {
            IndentStyle::Visual => other_child_shape,
            IndentStyle::Block => {
                shape.offset = shape.offset.checked_sub(context.config.tab_spaces).unwrap_or(0);
                shape.indent.block_indent += context.config.tab_spaces;
                shape
            }
        }
    } else {
        other_child_shape
    };
    debug!("child_shapes {:?} {:?}",
           first_child_shape,
           other_child_shape);

    let child_shape_iter =
        Some(first_child_shape).into_iter().chain(::std::iter::repeat(other_child_shape)
                                                      .take(subexpr_list.len() - 1));
    let iter = subexpr_list.iter().rev().zip(child_shape_iter);
    let mut rewrites =
        try_opt!(iter.map(|(e, shape)| rewrite_chain_subexpr(e, total_span, context, shape))
                     .collect::<Option<Vec<_>>>());

    // Total of all items excluding the last.
    let almost_total = rewrites[..rewrites.len() - 1].iter().fold(0, |a, b| {
        a + first_line_width(b)
    }) + parent_rewrite.len();

    let veto_single_line = if subexpr_list.len() > context.config.chain_one_line_max - 1 {
        // -1 above because subexpr_list does not include the parent.
        true
    } else if context.config.take_source_hints && subexpr_list.len() > 1 {
        // Look at the source code. Unless all chain elements start on the same
        // line, we won't consider putting them on a single line either.
        let last_span = context.snippet(mk_sp(subexpr_list[1].span.hi, total_span.hi));
        let first_span = context.snippet(subexpr_list[1].span);
        let last_iter = last_span.chars().take_while(|c| c.is_whitespace());

        first_span.chars().chain(last_iter).any(|c| c == '\n')
    } else {
        false
    };

    let mut fits_single_line = !veto_single_line && almost_total <= shape.width;
    if fits_single_line {
        let len = rewrites.len();
        let (init, last) = rewrites.split_at_mut(len - 1);
        fits_single_line = init.iter().all(|s| !s.contains('\n'));

        if fits_single_line {
            fits_single_line = match expr.node {
                ref e @ ast::ExprKind::MethodCall(..) => {
                    rewrite_method_call_with_overflow(e,
                                                      &mut last[0],
                                                      almost_total,
                                                      total_span,
                                                      context,
                                                      shape)
                }
                _ => !last[0].contains('\n'),
            }
        }
    }

    let connector = if fits_single_line && !parent_rewrite.contains('\n') {
        // Yay, we can put everything on one line.
        String::new()
    } else {
        // Use new lines.
        format!("\n{}", nested_shape.indent.to_string(context.config))
    };

    let first_connector = if extend || subexpr_list.is_empty() {
        ""
    } else if let ast::ExprKind::Try(_) = subexpr_list[0].node {
        ""
    } else {
        &*connector
    };

    wrap_str(format!("{}{}{}",
                     parent_rewrite,
                     first_connector,
                     join_rewrites(&rewrites, &subexpr_list, &connector)),
             context.config.max_width,
             shape)
}

// True if the chain is only `?`s.
fn chain_only_try(exprs: &[ast::Expr]) -> bool {
    exprs.iter().all(|e| if let ast::ExprKind::Try(_) = e.node {
                         true
                     } else {
                         false
                     })
}

pub fn rewrite_try(expr: &ast::Expr,
                   try_count: usize,
                   context: &RewriteContext,
                   shape: Shape)
                   -> Option<String> {
    let sub_expr = try_opt!(expr.rewrite(context, try_opt!(shape.sub_width(try_count))));
    Some(format!("{}{}",
                 sub_expr,
                 iter::repeat("?").take(try_count).collect::<String>()))
}

fn join_rewrites(rewrites: &[String], subexps: &[ast::Expr], connector: &str) -> String {
    let mut rewrite_iter = rewrites.iter();
    let mut result = rewrite_iter.next().unwrap().clone();
    let mut subexpr_iter = subexps.iter().rev();
    subexpr_iter.next();

    for (rewrite, expr) in rewrite_iter.zip(subexpr_iter) {
        match expr.node {
            ast::ExprKind::Try(_) => (),
            _ => result.push_str(connector),
        };
        result.push_str(&rewrite[..]);
    }

    result
}

// States whether an expression's last line exclusively consists of closing
// parens, braces, and brackets in its idiomatic formatting.
fn is_block_expr(expr: &ast::Expr, repr: &str) -> bool {
    match expr.node {
        ast::ExprKind::Struct(..) |
        ast::ExprKind::While(..) |
        ast::ExprKind::WhileLet(..) |
        ast::ExprKind::If(..) |
        ast::ExprKind::IfLet(..) |
        ast::ExprKind::Block(..) |
        ast::ExprKind::Loop(..) |
        ast::ExprKind::ForLoop(..) |
        ast::ExprKind::Match(..) => repr.contains('\n'),
        ast::ExprKind::Paren(ref expr) |
        ast::ExprKind::Binary(_, _, ref expr) |
        ast::ExprKind::Index(_, ref expr) |
        ast::ExprKind::Unary(_, ref expr) => is_block_expr(expr, repr),
        _ => false,
    }
}

// Returns the root of the chain and a Vec of the prefixes of the rest of the chain.
// E.g., for input `a.b.c` we return (`a`, [`a.b.c`, `a.b`])
fn make_subexpr_list(expr: &ast::Expr, context: &RewriteContext) -> (ast::Expr, Vec<ast::Expr>) {
    let mut subexpr_list = vec![expr.clone()];

    while let Some(subexpr) = pop_expr_chain(subexpr_list.last().unwrap(), context) {
        subexpr_list.push(subexpr.clone());
    }

    let parent = subexpr_list.pop().unwrap();
    (parent, subexpr_list)
}

fn chain_indent(context: &RewriteContext, shape: Shape) -> Shape {
    match context.config.chain_indent {
        IndentStyle::Visual => shape.visual_indent(0),
        IndentStyle::Block => shape.block_indent(context.config.tab_spaces),
    }
}

fn rewrite_method_call_with_overflow(expr_kind: &ast::ExprKind,
                                     last: &mut String,
                                     almost_total: usize,
                                     total_span: Span,
                                     context: &RewriteContext,
                                     shape: Shape)
                                     -> bool {
    if let &ast::ExprKind::MethodCall(ref method_name, ref types, ref expressions) = expr_kind {
        let shape = match shape.shrink_left(almost_total) {
            Some(b) => b,
            None => return false,
        };
        let mut last_rewrite = rewrite_method_call(method_name.node,
                                                   types,
                                                   expressions,
                                                   total_span,
                                                   context,
                                                   shape);

        if let Some(ref mut s) = last_rewrite {
            ::std::mem::swap(s, last);
            true
        } else {
            false
        }
    } else {
        unreachable!();
    }
}

// Returns the expression's subexpression, if it exists. When the subexpr
// is a try! macro, we'll convert it to shorthand when the option is set.
fn pop_expr_chain(expr: &ast::Expr, context: &RewriteContext) -> Option<ast::Expr> {
    match expr.node {
        ast::ExprKind::MethodCall(_, _, ref expressions) => {
            Some(convert_try(&expressions[0], context))
        }
        ast::ExprKind::TupField(ref subexpr, _) |
        ast::ExprKind::Field(ref subexpr, _) |
        ast::ExprKind::Try(ref subexpr) => Some(convert_try(subexpr, context)),
        _ => None,
    }
}

fn convert_try(expr: &ast::Expr, context: &RewriteContext) -> ast::Expr {
    match expr.node {
        ast::ExprKind::Mac(ref mac) if context.config.use_try_shorthand => {
            if let Some(subexpr) = convert_try_mac(mac, context) {
                subexpr
            } else {
                expr.clone()
            }
        }
        _ => expr.clone(),
    }
}

// Rewrite the last element in the chain `expr`. E.g., given `a.b.c` we rewrite
// `.c`.
fn rewrite_chain_subexpr(expr: &ast::Expr,
                         span: Span,
                         context: &RewriteContext,
                         shape: Shape)
                         -> Option<String> {
    match expr.node {
        ast::ExprKind::MethodCall(ref method_name, ref types, ref expressions) => {
            rewrite_method_call(method_name.node, types, expressions, span, context, shape)
        }
        ast::ExprKind::Field(_, ref field) => {
            let s = format!(".{}", field.node);
            if s.len() <= shape.width {
                Some(s)
            } else {
                None
            }
        }
        ast::ExprKind::TupField(_, ref field) => {
            let s = format!(".{}", field.node);
            if s.len() <= shape.width {
                Some(s)
            } else {
                None
            }
        }
        ast::ExprKind::Try(_) => {
            if shape.width >= 1 {
                Some("?".into())
            } else {
                None
            }
        }
        _ => unreachable!(),
    }
}

// Determines if we can continue formatting a given expression on the same line.
fn is_continuable(expr: &ast::Expr) -> bool {
    match expr.node {
        ast::ExprKind::Path(..) => true,
        _ => false,
    }
}

fn rewrite_method_call(method_name: ast::Ident,
                       types: &[ptr::P<ast::Ty>],
                       args: &[ptr::P<ast::Expr>],
                       span: Span,
                       context: &RewriteContext,
                       shape: Shape)
                       -> Option<String> {
    let (lo, type_str) = if types.is_empty() {
        (args[0].span.hi, String::new())
    } else {
        let type_list: Vec<_> =
            try_opt!(types.iter().map(|ty| ty.rewrite(context, shape)).collect());

        let type_str = if context.config.spaces_within_angle_brackets && type_list.len() > 0 {
            format!("::< {} >", type_list.join(", "))
        } else {
            format!("::<{}>", type_list.join(", "))
        };

        (types.last()
             .unwrap()
             .span
             .hi,
         type_str)
    };

    let callee_str = format!(".{}{}", method_name, type_str);
    let span = mk_sp(lo, span.hi);

    rewrite_call(context, &callee_str, &args[1..], span, shape)
}
<commit_msg>Put the second element of a chain on its own line when block formatting.<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.

/// Formatting of chained expressions, i.e. expressions which are chained by
/// dots: struct and enum field access, method calls, and try shorthand (?).
///
/// Instead of walking these subexpressions one-by-one, as is our usual strategy
/// for expression formatting, we collect maximal sequences of these expressions
/// and handle them simultaneously.
///
/// Whenever possible, the entire chain is put on a single line. If that fails,
/// we put each subexpression on a separate, much like the (default) function
/// argument function argument strategy.
///
/// Depends on config options: `chain_indent` is the indent to use for
/// blocks in the parent/root/base of the chain (and the rest of the chain's
/// alignment).
/// E.g., `let foo = { aaaa; bbb; ccc }.bar.baz();`, we would layout for the
/// following values of `chain_indent`:
/// Visual:
/// ```
/// let foo = {
///               aaaa;
///               bbb;
///               ccc
///           }
///           .bar
///           .baz();
/// ```
/// Inherit:
/// ```
/// let foo = {
///     aaaa;
///     bbb;
///     ccc
/// }
/// .bar
/// .baz();
/// ```
/// Tabbed:
/// ```
/// let foo = {
///         aaaa;
///         bbb;
///         ccc
///     }
///     .bar
///     .baz();
/// ```
///
/// If the first item in the chain is a block expression, we align the dots with
/// the braces.
/// Visual:
/// ```
/// let a = foo.bar
///            .baz()
///            .qux
/// ```
/// Inherit:
/// ```
/// let a = foo.bar
/// .baz()
/// .qux
/// ```
/// Tabbed:
/// ```
/// let a = foo.bar
///     .baz()
///     .qux
/// ```

use Shape;
use rewrite::{Rewrite, RewriteContext};
use utils::{wrap_str, first_line_width, last_line_width};
use expr::rewrite_call;
use config::IndentStyle;
use macros::convert_try_mac;

use std::iter;
use syntax::{ast, ptr};
use syntax::codemap::{mk_sp, Span};

pub fn rewrite_chain(expr: &ast::Expr, context: &RewriteContext, shape: Shape) -> Option<String> {
    debug!("rewrite_chain {:?}", shape);
    let total_span = expr.span;
    let (parent, subexpr_list) = make_subexpr_list(expr, context);

    // Bail out if the chain is just try sugar, i.e., an expression followed by
    // any number of `?`s.
    if chain_only_try(&subexpr_list) {
        return rewrite_try(&parent, subexpr_list.len(), context, shape);
    }

    // Parent is the first item in the chain, e.g., `foo` in `foo.bar.baz()`.
    let mut parent_shape = shape;
    if is_block_expr(&parent, "\n") {
        parent_shape = chain_indent(context, shape);
    }
    let parent_rewrite = try_opt!(parent.rewrite(context, parent_shape));

    // Decide how to layout the rest of the chain. `extend` is true if we can
    // put the first non-parent item on the same line as the parent.
    let (nested_shape, extend) = if !parent_rewrite.contains('\n') && is_continuable(&parent) {
        let nested_shape = if let ast::ExprKind::Try(..) = subexpr_list.last().unwrap().node {
            parent_shape.block_indent(context.config.tab_spaces)
        } else {
            chain_indent(context, shape.add_offset(parent_rewrite.len()))
        };
        (nested_shape, context.config.chain_indent == IndentStyle::Visual ||
                       parent_rewrite.len() <= context.config.tab_spaces)
    } else if is_block_expr(&parent, &parent_rewrite) {
        // The parent is a block, so align the rest of the chain with the closing
        // brace.
        (parent_shape, false)
    } else if parent_rewrite.contains('\n') {
        (chain_indent(context,
                      parent_shape.block_indent(context.config.tab_spaces)),
         false)
    } else {
        (shape.block_indent(context.config.tab_spaces), false)
    };

    let max_width = try_opt!((shape.width + shape.indent.width() + shape.offset)
                                 .checked_sub(nested_shape.indent.width() +
                                              nested_shape.offset));

    let other_child_shape = Shape {
        width: max_width,
        ..nested_shape
    };
    let first_child_shape = if extend {
        let mut shape = try_opt!(parent_shape.shrink_left(last_line_width(&parent_rewrite)));
        match context.config.chain_indent {
            IndentStyle::Visual => other_child_shape,
            IndentStyle::Block => {
                shape.offset = shape.offset.checked_sub(context.config.tab_spaces).unwrap_or(0);
                shape.indent.block_indent += context.config.tab_spaces;
                shape
            }
        }
    } else {
        other_child_shape
    };
    debug!("child_shapes {:?} {:?}",
           first_child_shape,
           other_child_shape);

    let child_shape_iter =
        Some(first_child_shape).into_iter().chain(::std::iter::repeat(other_child_shape)
                                                      .take(subexpr_list.len() - 1));
    let iter = subexpr_list.iter().rev().zip(child_shape_iter);
    let mut rewrites =
        try_opt!(iter.map(|(e, shape)| rewrite_chain_subexpr(e, total_span, context, shape))
                     .collect::<Option<Vec<_>>>());

    // Total of all items excluding the last.
    let almost_total = rewrites[..rewrites.len() - 1].iter().fold(0, |a, b| {
        a + first_line_width(b)
    }) + parent_rewrite.len();

    let veto_single_line = if subexpr_list.len() > context.config.chain_one_line_max - 1 {
        // -1 above because subexpr_list does not include the parent.
        true
    } else if context.config.take_source_hints && subexpr_list.len() > 1 {
        // Look at the source code. Unless all chain elements start on the same
        // line, we won't consider putting them on a single line either.
        let last_span = context.snippet(mk_sp(subexpr_list[1].span.hi, total_span.hi));
        let first_span = context.snippet(subexpr_list[1].span);
        let last_iter = last_span.chars().take_while(|c| c.is_whitespace());

        first_span.chars().chain(last_iter).any(|c| c == '\n')
    } else {
        false
    };

    let mut fits_single_line = !veto_single_line && almost_total <= shape.width;
    if fits_single_line {
        let len = rewrites.len();
        let (init, last) = rewrites.split_at_mut(len - 1);
        fits_single_line = init.iter().all(|s| !s.contains('\n'));

        if fits_single_line {
            fits_single_line = match expr.node {
                ref e @ ast::ExprKind::MethodCall(..) => {
                    rewrite_method_call_with_overflow(e,
                                                      &mut last[0],
                                                      almost_total,
                                                      total_span,
                                                      context,
                                                      shape)
                }
                _ => !last[0].contains('\n'),
            }
        }
    }

    let connector = if fits_single_line && !parent_rewrite.contains('\n') {
        // Yay, we can put everything on one line.
        String::new()
    } else {
        // Use new lines.
        format!("\n{}", nested_shape.indent.to_string(context.config))
    };

    let first_connector = if extend || subexpr_list.is_empty() {
        ""
    } else if let ast::ExprKind::Try(_) = subexpr_list[0].node {
        ""
    } else {
        &*connector
    };

    wrap_str(format!("{}{}{}",
                     parent_rewrite,
                     first_connector,
                     join_rewrites(&rewrites, &subexpr_list, &connector)),
             context.config.max_width,
             shape)
}

// True if the chain is only `?`s.
fn chain_only_try(exprs: &[ast::Expr]) -> bool {
    exprs.iter().all(|e| if let ast::ExprKind::Try(_) = e.node {
                         true
                     } else {
                         false
                     })
}

pub fn rewrite_try(expr: &ast::Expr,
                   try_count: usize,
                   context: &RewriteContext,
                   shape: Shape)
                   -> Option<String> {
    let sub_expr = try_opt!(expr.rewrite(context, try_opt!(shape.sub_width(try_count))));
    Some(format!("{}{}",
                 sub_expr,
                 iter::repeat("?").take(try_count).collect::<String>()))
}

fn join_rewrites(rewrites: &[String], subexps: &[ast::Expr], connector: &str) -> String {
    let mut rewrite_iter = rewrites.iter();
    let mut result = rewrite_iter.next().unwrap().clone();
    let mut subexpr_iter = subexps.iter().rev();
    subexpr_iter.next();

    for (rewrite, expr) in rewrite_iter.zip(subexpr_iter) {
        match expr.node {
            ast::ExprKind::Try(_) => (),
            _ => result.push_str(connector),
        };
        result.push_str(&rewrite[..]);
    }

    result
}

// States whether an expression's last line exclusively consists of closing
// parens, braces, and brackets in its idiomatic formatting.
fn is_block_expr(expr: &ast::Expr, repr: &str) -> bool {
    match expr.node {
        ast::ExprKind::Struct(..) |
        ast::ExprKind::While(..) |
        ast::ExprKind::WhileLet(..) |
        ast::ExprKind::If(..) |
        ast::ExprKind::IfLet(..) |
        ast::ExprKind::Block(..) |
        ast::ExprKind::Loop(..) |
        ast::ExprKind::ForLoop(..) |
        ast::ExprKind::Match(..) => repr.contains('\n'),
        ast::ExprKind::Paren(ref expr) |
        ast::ExprKind::Binary(_, _, ref expr) |
        ast::ExprKind::Index(_, ref expr) |
        ast::ExprKind::Unary(_, ref expr) => is_block_expr(expr, repr),
        _ => false,
    }
}

// Returns the root of the chain and a Vec of the prefixes of the rest of the chain.
// E.g., for input `a.b.c` we return (`a`, [`a.b.c`, `a.b`])
fn make_subexpr_list(expr: &ast::Expr, context: &RewriteContext) -> (ast::Expr, Vec<ast::Expr>) {
    let mut subexpr_list = vec![expr.clone()];

    while let Some(subexpr) = pop_expr_chain(subexpr_list.last().unwrap(), context) {
        subexpr_list.push(subexpr.clone());
    }

    let parent = subexpr_list.pop().unwrap();
    (parent, subexpr_list)
}

fn chain_indent(context: &RewriteContext, shape: Shape) -> Shape {
    match context.config.chain_indent {
        IndentStyle::Visual => shape.visual_indent(0),
        IndentStyle::Block => shape.block_indent(context.config.tab_spaces),
    }
}

fn rewrite_method_call_with_overflow(expr_kind: &ast::ExprKind,
                                     last: &mut String,
                                     almost_total: usize,
                                     total_span: Span,
                                     context: &RewriteContext,
                                     shape: Shape)
                                     -> bool {
    if let &ast::ExprKind::MethodCall(ref method_name, ref types, ref expressions) = expr_kind {
        let shape = match shape.shrink_left(almost_total) {
            Some(b) => b,
            None => return false,
        };
        let mut last_rewrite = rewrite_method_call(method_name.node,
                                                   types,
                                                   expressions,
                                                   total_span,
                                                   context,
                                                   shape);

        if let Some(ref mut s) = last_rewrite {
            ::std::mem::swap(s, last);
            true
        } else {
            false
        }
    } else {
        unreachable!();
    }
}

// Returns the expression's subexpression, if it exists. When the subexpr
// is a try! macro, we'll convert it to shorthand when the option is set.
fn pop_expr_chain(expr: &ast::Expr, context: &RewriteContext) -> Option<ast::Expr> {
    match expr.node {
        ast::ExprKind::MethodCall(_, _, ref expressions) => {
            Some(convert_try(&expressions[0], context))
        }
        ast::ExprKind::TupField(ref subexpr, _) |
        ast::ExprKind::Field(ref subexpr, _) |
        ast::ExprKind::Try(ref subexpr) => Some(convert_try(subexpr, context)),
        _ => None,
    }
}

fn convert_try(expr: &ast::Expr, context: &RewriteContext) -> ast::Expr {
    match expr.node {
        ast::ExprKind::Mac(ref mac) if context.config.use_try_shorthand => {
            if let Some(subexpr) = convert_try_mac(mac, context) {
                subexpr
            } else {
                expr.clone()
            }
        }
        _ => expr.clone(),
    }
}

// Rewrite the last element in the chain `expr`. E.g., given `a.b.c` we rewrite
// `.c`.
fn rewrite_chain_subexpr(expr: &ast::Expr,
                         span: Span,
                         context: &RewriteContext,
                         shape: Shape)
                         -> Option<String> {
    match expr.node {
        ast::ExprKind::MethodCall(ref method_name, ref types, ref expressions) => {
            rewrite_method_call(method_name.node, types, expressions, span, context, shape)
        }
        ast::ExprKind::Field(_, ref field) => {
            let s = format!(".{}", field.node);
            if s.len() <= shape.width {
                Some(s)
            } else {
                None
            }
        }
        ast::ExprKind::TupField(_, ref field) => {
            let s = format!(".{}", field.node);
            if s.len() <= shape.width {
                Some(s)
            } else {
                None
            }
        }
        ast::ExprKind::Try(_) => {
            if shape.width >= 1 {
                Some("?".into())
            } else {
                None
            }
        }
        _ => unreachable!(),
    }
}

// Determines if we can continue formatting a given expression on the same line.
fn is_continuable(expr: &ast::Expr) -> bool {
    match expr.node {
        ast::ExprKind::Path(..) => true,
        _ => false,
    }
}

fn rewrite_method_call(method_name: ast::Ident,
                       types: &[ptr::P<ast::Ty>],
                       args: &[ptr::P<ast::Expr>],
                       span: Span,
                       context: &RewriteContext,
                       shape: Shape)
                       -> Option<String> {
    let (lo, type_str) = if types.is_empty() {
        (args[0].span.hi, String::new())
    } else {
        let type_list: Vec<_> =
            try_opt!(types.iter().map(|ty| ty.rewrite(context, shape)).collect());

        let type_str = if context.config.spaces_within_angle_brackets && type_list.len() > 0 {
            format!("::< {} >", type_list.join(", "))
        } else {
            format!("::<{}>", type_list.join(", "))
        };

        (types.last()
             .unwrap()
             .span
             .hi,
         type_str)
    };

    let callee_str = format!(".{}{}", method_name, type_str);
    let span = mk_sp(lo, span.hi);

    rewrite_call(context, &callee_str, &args[1..], span, shape)
}
<|endoftext|>


<commit_before>use crate::errors::BeaconChainError;
use ssz::{Decode, DecodeError, Encode};
use std::collections::HashMap;
use std::convert::TryInto;
use std::fs::{File, OpenOptions};
use std::io::{self, Read, Write};
use std::path::Path;
use types::{BeaconState, EthSpec, PublicKey, PublicKeyBytes, Validator};

/// Provides a mapping of `validator_index -> validator_publickey`.
///
/// This cache exists for two reasons:
///
/// 1. To avoid reading a `BeaconState` from disk each time we need a public key.
/// 2. To reduce the amount of public key _decompression_ required. A `BeaconState` stores public
///    keys in compressed form and they are needed in decompressed form for signature verification.
///    Decompression is expensive when many keys are involved.
///
/// The cache has a `persistence_file` that it uses to maintain a persistent, on-disk
/// copy of itself. This allows it to be restored between process invocations.
pub struct ValidatorPubkeyCache {
    pubkeys: Vec<PublicKey>,
    indices: HashMap<PublicKeyBytes, usize>,
    persitence_file: ValidatorPubkeyCacheFile,
}

impl ValidatorPubkeyCache {
    pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<Self, BeaconChainError> {
        ValidatorPubkeyCacheFile::open(&path)
            .and_then(ValidatorPubkeyCacheFile::into_cache)
            .map_err(Into::into)
    }

    /// Create a new public key cache using the keys in `state.validators`.
    ///
    /// Also creates a new persistence file, returning an error if there is already a file at
    /// `persistence_path`.
    pub fn new<T: EthSpec, P: AsRef<Path>>(
        state: &BeaconState<T>,
        persistence_path: P,
    ) -> Result<Self, BeaconChainError> {
        if persistence_path.as_ref().exists() {
            return Err(BeaconChainError::ValidatorPubkeyCacheFileError(format!(
                "Persistence file already exists: {:?}",
                persistence_path.as_ref()
            )));
        }

        let mut cache = Self {
            persitence_file: ValidatorPubkeyCacheFile::create(persistence_path)?,
            pubkeys: vec![],
            indices: HashMap::new(),
        };

        cache.import_new_pubkeys(state)?;

        Ok(cache)
    }

    /// Scan the given `state` and add any new validator public keys.
    ///
    /// Does not delete any keys from `self` if they don't appear in `state`.
    pub fn import_new_pubkeys<T: EthSpec>(
        &mut self,
        state: &BeaconState<T>,
    ) -> Result<(), BeaconChainError> {
        if state.validators.len() > self.pubkeys.len() {
            self.import(&state.validators[self.pubkeys.len()..])
        } else {
            Ok(())
        }
    }

    /// Adds zero or more validators to `self`.
    fn import(&mut self, validators: &[Validator]) -> Result<(), BeaconChainError> {
        self.pubkeys.reserve(validators.len());
        self.indices.reserve(validators.len());

        for v in validators.iter() {
            let i = self.pubkeys.len();

            if self.indices.contains_key(&v.pubkey) {
                return Err(BeaconChainError::DuplicateValidatorPublicKey);
            }

            // The item is written to disk (the persistence file) _before_ it is written into
            // the local struct.
            //
            // This means that a pubkey cache read from disk will always be equivalent to or
            // _later than_ the cache that was running in the previous instance of Lighthouse.
            //
            // The motivation behind this ordering is that we do not want to have states that
            // reference a pubkey that is not in our cache. However, it's fine to have pubkeys
            // that are never referenced in a state.
            self.persitence_file.append(i, &v.pubkey)?;

            self.pubkeys.push(
                (&v.pubkey)
                    .try_into()
                    .map_err(BeaconChainError::InvalidValidatorPubkeyBytes)?,
            );

            self.indices.insert(v.pubkey.clone(), i);
        }

        Ok(())
    }

    /// Get the public key for a validator with index `i`.
    pub fn get(&self, i: usize) -> Option<&PublicKey> {
        self.pubkeys.get(i)
    }

    /// Get the index of a validator with `pubkey`.
    pub fn get_index(&self, pubkey: &PublicKeyBytes) -> Option<usize> {
        self.indices.get(pubkey).copied()
    }

    /// Returns the number of validators in the cache.
    pub fn len(&self) -> usize {
        self.indices.len()
    }
}

/// Allows for maintaining an on-disk copy of the `ValidatorPubkeyCache`. The file is raw SSZ bytes
/// (not ASCII encoded).
///
/// ## Writes
///
/// Each entry is simply appended to the file.
///
/// ## Reads
///
/// The whole file is parsed as an SSZ "variable list" of objects.
///
/// This parsing method is possible because the items in the list are fixed-length SSZ objects.
struct ValidatorPubkeyCacheFile(File);

#[derive(Debug)]
enum Error {
    Io(io::Error),
    Ssz(DecodeError),
    PubkeyDecode(bls::Error),
    /// The file read from disk does not have a contiguous list of validator public keys. The file
    /// has become corrupted.
    InconsistentIndex {
        expected: Option<usize>,
        found: usize,
    },
}

impl From<Error> for BeaconChainError {
    fn from(e: Error) -> BeaconChainError {
        BeaconChainError::ValidatorPubkeyCacheFileError(format!("{:?}", e))
    }
}

impl ValidatorPubkeyCacheFile {
    /// Creates a file for reading and writing.
    pub fn create<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
        OpenOptions::new()
            .create_new(true)
            .write(true)
            .open(path)
            .map(Self)
            .map_err(Error::Io)
    }

    /// Opens an existing file for reading and writing.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
        OpenOptions::new()
            .read(true)
            .write(true)
            .create(false)
            .append(true)
            .open(path)
            .map(Self)
            .map_err(Error::Io)
    }

    /// Append a public key to file.
    ///
    /// The provided `index` should each be one greater than the previous and start at 0.
    /// Otherwise, the file will become corrupted and unable to be converted into a cache .
    pub fn append(&mut self, index: usize, pubkey: &PublicKeyBytes) -> Result<(), Error> {
        append_to_file(&mut self.0, index, pubkey)
    }

    /// Creates a `ValidatorPubkeyCache` by reading and parsing the underlying file.
    pub fn into_cache(mut self) -> Result<ValidatorPubkeyCache, Error> {
        let mut bytes = vec![];
        self.0.read_to_end(&mut bytes).map_err(Error::Io)?;

        let list: Vec<(usize, PublicKeyBytes)> = Vec::from_ssz_bytes(&bytes).map_err(Error::Ssz)?;

        let mut last = None;
        let mut pubkeys = Vec::with_capacity(list.len());
        let mut indices = HashMap::new();

        for (index, pubkey) in list {
            let expected = last.map(|n| n + 1);
            if expected.map_or(true, |expected| index == expected) {
                last = Some(index);
                pubkeys.push((&pubkey).try_into().map_err(Error::PubkeyDecode)?);
                indices.insert(pubkey, index);
            } else {
                return Err(Error::InconsistentIndex {
                    expected,
                    found: index,
                });
            }
        }

        Ok(ValidatorPubkeyCache {
            pubkeys,
            indices,
            persitence_file: self,
        })
    }
}

fn append_to_file(file: &mut File, index: usize, pubkey: &PublicKeyBytes) -> Result<(), Error> {
    let mut line = Vec::with_capacity(index.ssz_bytes_len() + pubkey.ssz_bytes_len());

    index.ssz_append(&mut line);
    pubkey.ssz_append(&mut line);

    file.write_all(&line).map_err(Error::Io)
}

#[cfg(test)]
mod test {
    use super::*;
    use tempfile::tempdir;
    use types::{
        test_utils::{generate_deterministic_keypair, TestingBeaconStateBuilder},
        BeaconState, EthSpec, Keypair, MainnetEthSpec,
    };

    fn get_state(validator_count: usize) -> (BeaconState<MainnetEthSpec>, Vec<Keypair>) {
        let spec = MainnetEthSpec::default_spec();
        let builder =
            TestingBeaconStateBuilder::from_deterministic_keypairs(validator_count, &spec);
        builder.build()
    }

    fn check_cache_get(cache: &ValidatorPubkeyCache, keypairs: &[Keypair]) {
        let validator_count = keypairs.len();

        for i in 0..validator_count + 1 {
            if i < validator_count {
                let pubkey = cache.get(i).expect("pubkey should be present");
                assert_eq!(pubkey, &keypairs[i].pk, "pubkey should match cache");

                let pubkey_bytes: PublicKeyBytes = pubkey.clone().into();

                assert_eq!(
                    i,
                    cache
                        .get_index(&pubkey_bytes)
                        .expect("should resolve index"),
                    "index should match cache"
                );
            } else {
                assert_eq!(
                    cache.get(i),
                    None,
                    "should not get pubkey for out of bounds index",
                );
            }
        }
    }

    #[test]
    fn basic_operation() {
        let (state, keypairs) = get_state(8);

        let dir = tempdir().expect("should create tempdir");
        let path = dir.path().join("cache.ssz");

        let mut cache = ValidatorPubkeyCache::new(&state, path).expect("should create cache");

        check_cache_get(&cache, &keypairs[..]);

        // Try adding a state with the same number of keypairs.
        let (state, keypairs) = get_state(8);
        cache
            .import_new_pubkeys(&state)
            .expect("should import pubkeys");
        check_cache_get(&cache, &keypairs[..]);

        // Try adding a state with less keypairs.
        let (state, _) = get_state(1);
        cache
            .import_new_pubkeys(&state)
            .expect("should import pubkeys");
        check_cache_get(&cache, &keypairs[..]);

        // Try adding a state with more keypairs.
        let (state, keypairs) = get_state(12);
        cache
            .import_new_pubkeys(&state)
            .expect("should import pubkeys");
        check_cache_get(&cache, &keypairs[..]);
    }

    #[test]
    fn persistence() {
        let (state, keypairs) = get_state(8);

        let dir = tempdir().expect("should create tempdir");
        let path = dir.path().join("cache.ssz");

        // Create a new cache.
        let cache = ValidatorPubkeyCache::new(&state, &path).expect("should create cache");
        check_cache_get(&cache, &keypairs[..]);
        drop(cache);

        // Re-init the cache from the file.
        let mut cache = ValidatorPubkeyCache::load_from_file(&path).expect("should open cache");
        check_cache_get(&cache, &keypairs[..]);

        // Add some more keypairs.
        let (state, keypairs) = get_state(12);
        cache
            .import_new_pubkeys(&state)
            .expect("should import pubkeys");
        check_cache_get(&cache, &keypairs[..]);
        drop(cache);

        // Re-init the cache from the file.
        let cache = ValidatorPubkeyCache::load_from_file(&path).expect("should open cache");
        check_cache_get(&cache, &keypairs[..]);
    }

    #[test]
    fn invalid_persisted_file() {
        let dir = tempdir().expect("should create tempdir");
        let path = dir.path().join("cache.ssz");
        let pubkey = generate_deterministic_keypair(0).pk.into();

        let mut file = File::create(&path).expect("should create file");
        append_to_file(&mut file, 0, &pubkey).expect("should write to file");
        drop(file);

        let cache = ValidatorPubkeyCache::load_from_file(&path).expect("should open cache");
        drop(cache);

        let mut file = OpenOptions::new()
            .write(true)
            .append(true)
            .open(&path)
            .expect("should open file");

        append_to_file(&mut file, 42, &pubkey).expect("should write bad data to file");
        drop(file);

        assert!(
            ValidatorPubkeyCache::load_from_file(&path).is_err(),
            "should not parse invalid file"
        );
    }
}
<commit_msg>Allow truncation of pubkey cache on creation (#1686)<commit_after>use crate::errors::BeaconChainError;
use ssz::{Decode, DecodeError, Encode};
use std::collections::HashMap;
use std::convert::TryInto;
use std::fs::{File, OpenOptions};
use std::io::{self, Read, Write};
use std::path::Path;
use types::{BeaconState, EthSpec, PublicKey, PublicKeyBytes, Validator};

/// Provides a mapping of `validator_index -> validator_publickey`.
///
/// This cache exists for two reasons:
///
/// 1. To avoid reading a `BeaconState` from disk each time we need a public key.
/// 2. To reduce the amount of public key _decompression_ required. A `BeaconState` stores public
///    keys in compressed form and they are needed in decompressed form for signature verification.
///    Decompression is expensive when many keys are involved.
///
/// The cache has a `persistence_file` that it uses to maintain a persistent, on-disk
/// copy of itself. This allows it to be restored between process invocations.
pub struct ValidatorPubkeyCache {
    pubkeys: Vec<PublicKey>,
    indices: HashMap<PublicKeyBytes, usize>,
    persitence_file: ValidatorPubkeyCacheFile,
}

impl ValidatorPubkeyCache {
    pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<Self, BeaconChainError> {
        ValidatorPubkeyCacheFile::open(&path)
            .and_then(ValidatorPubkeyCacheFile::into_cache)
            .map_err(Into::into)
    }

    /// Create a new public key cache using the keys in `state.validators`.
    ///
    /// Also creates a new persistence file, returning an error if there is already a file at
    /// `persistence_path`.
    pub fn new<T: EthSpec, P: AsRef<Path>>(
        state: &BeaconState<T>,
        persistence_path: P,
    ) -> Result<Self, BeaconChainError> {
        let mut cache = Self {
            persitence_file: ValidatorPubkeyCacheFile::create(persistence_path)?,
            pubkeys: vec![],
            indices: HashMap::new(),
        };

        cache.import_new_pubkeys(state)?;

        Ok(cache)
    }

    /// Scan the given `state` and add any new validator public keys.
    ///
    /// Does not delete any keys from `self` if they don't appear in `state`.
    pub fn import_new_pubkeys<T: EthSpec>(
        &mut self,
        state: &BeaconState<T>,
    ) -> Result<(), BeaconChainError> {
        if state.validators.len() > self.pubkeys.len() {
            self.import(&state.validators[self.pubkeys.len()..])
        } else {
            Ok(())
        }
    }

    /// Adds zero or more validators to `self`.
    fn import(&mut self, validators: &[Validator]) -> Result<(), BeaconChainError> {
        self.pubkeys.reserve(validators.len());
        self.indices.reserve(validators.len());

        for v in validators.iter() {
            let i = self.pubkeys.len();

            if self.indices.contains_key(&v.pubkey) {
                return Err(BeaconChainError::DuplicateValidatorPublicKey);
            }

            // The item is written to disk (the persistence file) _before_ it is written into
            // the local struct.
            //
            // This means that a pubkey cache read from disk will always be equivalent to or
            // _later than_ the cache that was running in the previous instance of Lighthouse.
            //
            // The motivation behind this ordering is that we do not want to have states that
            // reference a pubkey that is not in our cache. However, it's fine to have pubkeys
            // that are never referenced in a state.
            self.persitence_file.append(i, &v.pubkey)?;

            self.pubkeys.push(
                (&v.pubkey)
                    .try_into()
                    .map_err(BeaconChainError::InvalidValidatorPubkeyBytes)?,
            );

            self.indices.insert(v.pubkey.clone(), i);
        }

        Ok(())
    }

    /// Get the public key for a validator with index `i`.
    pub fn get(&self, i: usize) -> Option<&PublicKey> {
        self.pubkeys.get(i)
    }

    /// Get the index of a validator with `pubkey`.
    pub fn get_index(&self, pubkey: &PublicKeyBytes) -> Option<usize> {
        self.indices.get(pubkey).copied()
    }

    /// Returns the number of validators in the cache.
    pub fn len(&self) -> usize {
        self.indices.len()
    }
}

/// Allows for maintaining an on-disk copy of the `ValidatorPubkeyCache`. The file is raw SSZ bytes
/// (not ASCII encoded).
///
/// ## Writes
///
/// Each entry is simply appended to the file.
///
/// ## Reads
///
/// The whole file is parsed as an SSZ "variable list" of objects.
///
/// This parsing method is possible because the items in the list are fixed-length SSZ objects.
struct ValidatorPubkeyCacheFile(File);

#[derive(Debug)]
enum Error {
    Io(io::Error),
    Ssz(DecodeError),
    PubkeyDecode(bls::Error),
    /// The file read from disk does not have a contiguous list of validator public keys. The file
    /// has become corrupted.
    InconsistentIndex {
        expected: Option<usize>,
        found: usize,
    },
}

impl From<Error> for BeaconChainError {
    fn from(e: Error) -> BeaconChainError {
        BeaconChainError::ValidatorPubkeyCacheFileError(format!("{:?}", e))
    }
}

impl ValidatorPubkeyCacheFile {
    /// Creates a file for reading and writing.
    pub fn create<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
        OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(true)
            .open(path)
            .map(Self)
            .map_err(Error::Io)
    }

    /// Opens an existing file for reading and writing.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
        OpenOptions::new()
            .read(true)
            .write(true)
            .create(false)
            .append(true)
            .open(path)
            .map(Self)
            .map_err(Error::Io)
    }

    /// Append a public key to file.
    ///
    /// The provided `index` should each be one greater than the previous and start at 0.
    /// Otherwise, the file will become corrupted and unable to be converted into a cache .
    pub fn append(&mut self, index: usize, pubkey: &PublicKeyBytes) -> Result<(), Error> {
        append_to_file(&mut self.0, index, pubkey)
    }

    /// Creates a `ValidatorPubkeyCache` by reading and parsing the underlying file.
    pub fn into_cache(mut self) -> Result<ValidatorPubkeyCache, Error> {
        let mut bytes = vec![];
        self.0.read_to_end(&mut bytes).map_err(Error::Io)?;

        let list: Vec<(usize, PublicKeyBytes)> = Vec::from_ssz_bytes(&bytes).map_err(Error::Ssz)?;

        let mut last = None;
        let mut pubkeys = Vec::with_capacity(list.len());
        let mut indices = HashMap::new();

        for (index, pubkey) in list {
            let expected = last.map(|n| n + 1);
            if expected.map_or(true, |expected| index == expected) {
                last = Some(index);
                pubkeys.push((&pubkey).try_into().map_err(Error::PubkeyDecode)?);
                indices.insert(pubkey, index);
            } else {
                return Err(Error::InconsistentIndex {
                    expected,
                    found: index,
                });
            }
        }

        Ok(ValidatorPubkeyCache {
            pubkeys,
            indices,
            persitence_file: self,
        })
    }
}

fn append_to_file(file: &mut File, index: usize, pubkey: &PublicKeyBytes) -> Result<(), Error> {
    let mut line = Vec::with_capacity(index.ssz_bytes_len() + pubkey.ssz_bytes_len());

    index.ssz_append(&mut line);
    pubkey.ssz_append(&mut line);

    file.write_all(&line).map_err(Error::Io)
}

#[cfg(test)]
mod test {
    use super::*;
    use tempfile::tempdir;
    use types::{
        test_utils::{generate_deterministic_keypair, TestingBeaconStateBuilder},
        BeaconState, EthSpec, Keypair, MainnetEthSpec,
    };

    fn get_state(validator_count: usize) -> (BeaconState<MainnetEthSpec>, Vec<Keypair>) {
        let spec = MainnetEthSpec::default_spec();
        let builder =
            TestingBeaconStateBuilder::from_deterministic_keypairs(validator_count, &spec);
        builder.build()
    }

    fn check_cache_get(cache: &ValidatorPubkeyCache, keypairs: &[Keypair]) {
        let validator_count = keypairs.len();

        for i in 0..validator_count + 1 {
            if i < validator_count {
                let pubkey = cache.get(i).expect("pubkey should be present");
                assert_eq!(pubkey, &keypairs[i].pk, "pubkey should match cache");

                let pubkey_bytes: PublicKeyBytes = pubkey.clone().into();

                assert_eq!(
                    i,
                    cache
                        .get_index(&pubkey_bytes)
                        .expect("should resolve index"),
                    "index should match cache"
                );
            } else {
                assert_eq!(
                    cache.get(i),
                    None,
                    "should not get pubkey for out of bounds index",
                );
            }
        }
    }

    #[test]
    fn basic_operation() {
        let (state, keypairs) = get_state(8);

        let dir = tempdir().expect("should create tempdir");
        let path = dir.path().join("cache.ssz");

        let mut cache = ValidatorPubkeyCache::new(&state, path).expect("should create cache");

        check_cache_get(&cache, &keypairs[..]);

        // Try adding a state with the same number of keypairs.
        let (state, keypairs) = get_state(8);
        cache
            .import_new_pubkeys(&state)
            .expect("should import pubkeys");
        check_cache_get(&cache, &keypairs[..]);

        // Try adding a state with less keypairs.
        let (state, _) = get_state(1);
        cache
            .import_new_pubkeys(&state)
            .expect("should import pubkeys");
        check_cache_get(&cache, &keypairs[..]);

        // Try adding a state with more keypairs.
        let (state, keypairs) = get_state(12);
        cache
            .import_new_pubkeys(&state)
            .expect("should import pubkeys");
        check_cache_get(&cache, &keypairs[..]);
    }

    #[test]
    fn persistence() {
        let (state, keypairs) = get_state(8);

        let dir = tempdir().expect("should create tempdir");
        let path = dir.path().join("cache.ssz");

        // Create a new cache.
        let cache = ValidatorPubkeyCache::new(&state, &path).expect("should create cache");
        check_cache_get(&cache, &keypairs[..]);
        drop(cache);

        // Re-init the cache from the file.
        let mut cache = ValidatorPubkeyCache::load_from_file(&path).expect("should open cache");
        check_cache_get(&cache, &keypairs[..]);

        // Add some more keypairs.
        let (state, keypairs) = get_state(12);
        cache
            .import_new_pubkeys(&state)
            .expect("should import pubkeys");
        check_cache_get(&cache, &keypairs[..]);
        drop(cache);

        // Re-init the cache from the file.
        let cache = ValidatorPubkeyCache::load_from_file(&path).expect("should open cache");
        check_cache_get(&cache, &keypairs[..]);
    }

    #[test]
    fn invalid_persisted_file() {
        let dir = tempdir().expect("should create tempdir");
        let path = dir.path().join("cache.ssz");
        let pubkey = generate_deterministic_keypair(0).pk.into();

        let mut file = File::create(&path).expect("should create file");
        append_to_file(&mut file, 0, &pubkey).expect("should write to file");
        drop(file);

        let cache = ValidatorPubkeyCache::load_from_file(&path).expect("should open cache");
        drop(cache);

        let mut file = OpenOptions::new()
            .write(true)
            .append(true)
            .open(&path)
            .expect("should open file");

        append_to_file(&mut file, 42, &pubkey).expect("should write bad data to file");
        drop(file);

        assert!(
            ValidatorPubkeyCache::load_from_file(&path).is_err(),
            "should not parse invalid file"
        );
    }
}
<|endoftext|>