<commit_before>use crate::interface::{Compiler, Result};
use crate::passes::{self, BoxedResolver, BoxedGlobalCtxt};

use rustc_incremental::DepGraphFuture;
use rustc_data_structures::sync::Lrc;
use rustc_codegen_utils::codegen_backend::CodegenBackend;
use rustc::session::config::{OutputFilenames, OutputType};
use rustc::util::common::{time, ErrorReported};
use rustc::hir;
use rustc::lint;
use rustc::session::Session;
use rustc::lint::LintStore;
use rustc::hir::def_id::LOCAL_CRATE;
use rustc::ty::steal::Steal;
use rustc::ty::ResolverOutputs;
use rustc::dep_graph::DepGraph;
use std::cell::{Ref, RefMut, RefCell};
use std::rc::Rc;
use std::any::Any;
use std::mem;
use syntax::{self, ast};

/// Represent the result of a query.
/// This result can be stolen with the `take` method and returned with the `give` method.
pub struct Query<T> {
    result: RefCell<Option<Result<T>>>,
}

impl<T> Query<T> {
    fn compute<F: FnOnce() -> Result<T>>(&self, f: F) -> Result<&Query<T>> {
        let mut result = self.result.borrow_mut();
        if result.is_none() {
            *result = Some(f());
        }
        result.as_ref().unwrap().as_ref().map(|_| self).map_err(|err| *err)
    }

    /// Takes ownership of the query result. Further attempts to take or peek the query
    /// result will panic unless it is returned by calling the `give` method.
    pub fn take(&self) -> T {
        self.result
            .borrow_mut()
            .take()
            .expect("missing query result")
            .unwrap()
    }

    /// Borrows the query result using the RefCell. Panics if the result is stolen.
    pub fn peek(&self) -> Ref<'_, T> {
        Ref::map(self.result.borrow(), |r| {
            r.as_ref().unwrap().as_ref().expect("missing query result")
        })
    }

    /// Mutably borrows the query result using the RefCell. Panics if the result is stolen.
    pub fn peek_mut(&self) -> RefMut<'_, T> {
        RefMut::map(self.result.borrow_mut(), |r| {
            r.as_mut().unwrap().as_mut().expect("missing query result")
        })
    }
}

impl<T> Default for Query<T> {
    fn default() -> Self {
        Query {
            result: RefCell::new(None),
        }
    }
}

pub struct Queries<'comp> {
    compiler: &'comp Compiler,

    dep_graph_future: Query<Option<DepGraphFuture>>,
    parse: Query<ast::Crate>,
    crate_name: Query<String>,
    register_plugins: Query<(ast::Crate, Lrc<LintStore>)>,
    expansion: Query<(ast::Crate, Steal<Rc<RefCell<BoxedResolver>>>, Lrc<LintStore>)>,
    dep_graph: Query<DepGraph>,
    lower_to_hir: Query<(Steal<hir::map::Forest>, Steal<ResolverOutputs>)>,
    prepare_outputs: Query<OutputFilenames>,
    global_ctxt: Query<BoxedGlobalCtxt>,
    ongoing_codegen: Query<Box<dyn Any>>,
}

impl<'comp> Queries<'comp> {
    pub fn new(compiler: &'comp Compiler) -> Queries<'comp> {
        Queries {
            compiler,
            dep_graph_future: Default::default(),
            parse: Default::default(),
            crate_name: Default::default(),
            register_plugins: Default::default(),
            expansion: Default::default(),
            dep_graph: Default::default(),
            lower_to_hir: Default::default(),
            prepare_outputs: Default::default(),
            global_ctxt: Default::default(),
            ongoing_codegen: Default::default(),
        }
    }

    fn session(&self) -> &Lrc<Session> {
        &self.compiler.sess
    }
    fn codegen_backend(&self) -> &Lrc<Box<dyn CodegenBackend>> {
        &self.compiler.codegen_backend()
    }

    pub fn dep_graph_future(&self) -> Result<&Query<Option<DepGraphFuture>>> {
        self.dep_graph_future.compute(|| {
            Ok(if self.session().opts.build_dep_graph() {
                Some(rustc_incremental::load_dep_graph(self.session()))
            } else {
                None
            })
        })
    }

    pub fn parse(&self) -> Result<&Query<ast::Crate>> {
        self.parse.compute(|| {
            passes::parse(self.session(), &self.compiler.input).map_err(
                |mut parse_error| {
                    parse_error.emit();
                    ErrorReported
                },
            )
        })
    }

    pub fn register_plugins(&self) -> Result<&Query<(ast::Crate, Lrc<LintStore>)>> {
        self.register_plugins.compute(|| {
            let crate_name = self.crate_name()?.peek().clone();
            let krate = self.parse()?.take();

            let empty: &(dyn Fn(&Session, &mut lint::LintStore) + Sync + Send) = &|_, _| {};
            let result = passes::register_plugins(
                self.session(),
                &*self.codegen_backend().metadata_loader(),
                self.compiler.register_lints
                    .as_ref()
                    .map(|p| &**p)
                    .unwrap_or_else(|| empty),
                krate,
                &crate_name,
            );

            // Compute the dependency graph (in the background). We want to do
            // this as early as possible, to give the DepGraph maximum time to
            // load before dep_graph() is called, but it also can't happen
            // until after rustc_incremental::prepare_session_directory() is
            // called, which happens within passes::register_plugins().
            self.dep_graph_future().ok();

            result
        })
    }

    pub fn crate_name(&self) -> Result<&Query<String>> {
        self.crate_name.compute(|| {
            Ok(match self.compiler.crate_name {
                Some(ref crate_name) => crate_name.clone(),
                None => {
                    let parse_result = self.parse()?;
                    let krate = parse_result.peek();
                    rustc_codegen_utils::link::find_crate_name(
                        Some(self.session()),
                        &krate.attrs,
                        &self.compiler.input
                    )
                }
            })
        })
    }

    pub fn expansion(
        &self
    ) -> Result<&Query<(ast::Crate, Steal<Rc<RefCell<BoxedResolver>>>, Lrc<LintStore>)>> {
        self.expansion.compute(|| {
            let crate_name = self.crate_name()?.peek().clone();
            let (krate, lint_store) = self.register_plugins()?.take();
            passes::configure_and_expand(
                self.session().clone(),
                lint_store.clone(),
                self.codegen_backend().metadata_loader(),
                krate,
                &crate_name,
            ).map(|(krate, resolver)| {
                (krate, Steal::new(Rc::new(RefCell::new(resolver))), lint_store)
            })
        })
    }

    pub fn dep_graph(&self) -> Result<&Query<DepGraph>> {
        self.dep_graph.compute(|| {
            Ok(match self.dep_graph_future()?.take() {
                None => DepGraph::new_disabled(),
                Some(future) => {
                    let (prev_graph, prev_work_products) =
                        time(self.session(), "blocked while dep-graph loading finishes", || {
                            future.open().unwrap_or_else(|e| rustc_incremental::LoadResult::Error {
                                message: format!("could not decode incremental cache: {:?}", e),
                            }).open(self.session())
                        });
                    DepGraph::new(prev_graph, prev_work_products)
                }
            })
        })
    }

    pub fn lower_to_hir(
        &self,
    ) -> Result<&Query<(Steal<hir::map::Forest>, Steal<ResolverOutputs>)>> {
        self.lower_to_hir.compute(|| {
            let expansion_result = self.expansion()?;
            let peeked = expansion_result.peek();
            let krate = &peeked.0;
            let resolver = peeked.1.steal();
            let lint_store = &peeked.2;
            let hir = Steal::new(resolver.borrow_mut().access(|resolver| {
                passes::lower_to_hir(
                    self.session(),
                    lint_store,
                    resolver,
                    &*self.dep_graph()?.peek(),
                    &krate
                )
            })?);
            Ok((hir, Steal::new(BoxedResolver::to_resolver_outputs(resolver))))
        })
    }

    pub fn prepare_outputs(&self) -> Result<&Query<OutputFilenames>> {
        self.prepare_outputs.compute(|| {
            let expansion_result = self.expansion()?;
            let (krate, boxed_resolver, _) = &*expansion_result.peek();
            let crate_name = self.crate_name()?;
            let crate_name = crate_name.peek();
            passes::prepare_outputs(
                self.session(), self.compiler, &krate, &boxed_resolver, &crate_name
            )
        })
    }

    pub fn global_ctxt(&self) -> Result<&Query<BoxedGlobalCtxt>> {
        self.global_ctxt.compute(|| {
            let crate_name = self.crate_name()?.peek().clone();
            let outputs = self.prepare_outputs()?.peek().clone();
            let lint_store = self.expansion()?.peek().2.clone();
            let hir = self.lower_to_hir()?;
            let hir = hir.peek();
            let (hir_forest, resolver_outputs) = &*hir;
            Ok(passes::create_global_ctxt(
                self.compiler,
                lint_store,
                hir_forest.steal(),
                resolver_outputs.steal(),
                outputs,
                &crate_name))
        })
    }

    pub fn ongoing_codegen(&self) -> Result<&Query<Box<dyn Any>>> {
        self.ongoing_codegen.compute(|| {
            let outputs = self.prepare_outputs()?;
            self.global_ctxt()?.peek_mut().enter(|tcx| {
                tcx.analysis(LOCAL_CRATE).ok();

                // Don't do code generation if there were any errors
                self.session().compile_status()?;

                Ok(passes::start_codegen(
                    &***self.codegen_backend(),
                    tcx,
                    &*outputs.peek()
                ))
            })
        })
    }

    pub fn linker(self) -> Result<Linker> {
        let dep_graph = self.dep_graph()?;
        let prepare_outputs = self.prepare_outputs()?;
        let ongoing_codegen = self.ongoing_codegen()?;

        let sess = self.session().clone();
        let codegen_backend = self.codegen_backend().clone();

        Ok(Linker {
            sess,
            dep_graph: dep_graph.take(),
            prepare_outputs: prepare_outputs.take(),
            ongoing_codegen: ongoing_codegen.take(),
            codegen_backend,
        })
    }
}

pub struct Linker {
    sess: Lrc<Session>,
    dep_graph: DepGraph,
    prepare_outputs: OutputFilenames,
    ongoing_codegen: Box<dyn Any>,
    codegen_backend: Lrc<Box<dyn CodegenBackend>>,
}

impl Linker {
    pub fn link(self) -> Result<()> {
        self.codegen_backend.join_codegen_and_link(
            self.ongoing_codegen,
            &self.sess,
            &self.dep_graph,
            &self.prepare_outputs,
        ).map_err(|_| ErrorReported)
    }
}

impl Compiler {
    // This method is different to all the other methods in `Compiler` because
    // it lacks a `Queries` entry. It's also not currently used. It does serve
    // as an example of how `Compiler` can be used, with additional steps added
    // between some passes. And see `rustc_driver::run_compiler` for a more
    // complex example.
    pub fn enter<'c, F, T>(&'c self, f: F) -> Result<T>
        where F: FnOnce(Queries<'c>) -> Result<T>
    {
        let queries = Queries::new(&self);
        f(queries)
    }

    // This method is different to all the other methods in `Compiler` because
    // it lacks a `Queries` entry. It's also not currently used. It does serve
    // as an example of how `Compiler` can be used, with additional steps added
    // between some passes. And see `rustc_driver::run_compiler` for a more
    // complex example.
    pub fn compile(&self) -> Result<()> {
        let linker = self.enter(|queries| {
            queries.prepare_outputs()?;

            if self.session().opts.output_types.contains_key(&OutputType::DepInfo)
                && self.session().opts.output_types.len() == 1
            {
                return Ok(None)
            }

            queries.global_ctxt()?;

            // Drop AST after creating GlobalCtxt to free memory.
            mem::drop(queries.expansion()?.take());

            queries.ongoing_codegen()?;

            // Drop GlobalCtxt after starting codegen to free memory.
            mem::drop(queries.global_ctxt()?.take());

            let linker = queries.linker()?;
            Ok(Some(linker))
        })?;

        if let Some(linker) = linker {
            linker.link()?
        }

        Ok(())
    }
}
<commit_msg>Remove wrong comment.<commit_after>use crate::interface::{Compiler, Result};
use crate::passes::{self, BoxedResolver, BoxedGlobalCtxt};

use rustc_incremental::DepGraphFuture;
use rustc_data_structures::sync::Lrc;
use rustc_codegen_utils::codegen_backend::CodegenBackend;
use rustc::session::config::{OutputFilenames, OutputType};
use rustc::util::common::{time, ErrorReported};
use rustc::hir;
use rustc::lint;
use rustc::session::Session;
use rustc::lint::LintStore;
use rustc::hir::def_id::LOCAL_CRATE;
use rustc::ty::steal::Steal;
use rustc::ty::ResolverOutputs;
use rustc::dep_graph::DepGraph;
use std::cell::{Ref, RefMut, RefCell};
use std::rc::Rc;
use std::any::Any;
use std::mem;
use syntax::{self, ast};

/// Represent the result of a query.
/// This result can be stolen with the `take` method and returned with the `give` method.
pub struct Query<T> {
    result: RefCell<Option<Result<T>>>,
}

impl<T> Query<T> {
    fn compute<F: FnOnce() -> Result<T>>(&self, f: F) -> Result<&Query<T>> {
        let mut result = self.result.borrow_mut();
        if result.is_none() {
            *result = Some(f());
        }
        result.as_ref().unwrap().as_ref().map(|_| self).map_err(|err| *err)
    }

    /// Takes ownership of the query result. Further attempts to take or peek the query
    /// result will panic unless it is returned by calling the `give` method.
    pub fn take(&self) -> T {
        self.result
            .borrow_mut()
            .take()
            .expect("missing query result")
            .unwrap()
    }

    /// Borrows the query result using the RefCell. Panics if the result is stolen.
    pub fn peek(&self) -> Ref<'_, T> {
        Ref::map(self.result.borrow(), |r| {
            r.as_ref().unwrap().as_ref().expect("missing query result")
        })
    }

    /// Mutably borrows the query result using the RefCell. Panics if the result is stolen.
    pub fn peek_mut(&self) -> RefMut<'_, T> {
        RefMut::map(self.result.borrow_mut(), |r| {
            r.as_mut().unwrap().as_mut().expect("missing query result")
        })
    }
}

impl<T> Default for Query<T> {
    fn default() -> Self {
        Query {
            result: RefCell::new(None),
        }
    }
}

pub struct Queries<'comp> {
    compiler: &'comp Compiler,

    dep_graph_future: Query<Option<DepGraphFuture>>,
    parse: Query<ast::Crate>,
    crate_name: Query<String>,
    register_plugins: Query<(ast::Crate, Lrc<LintStore>)>,
    expansion: Query<(ast::Crate, Steal<Rc<RefCell<BoxedResolver>>>, Lrc<LintStore>)>,
    dep_graph: Query<DepGraph>,
    lower_to_hir: Query<(Steal<hir::map::Forest>, Steal<ResolverOutputs>)>,
    prepare_outputs: Query<OutputFilenames>,
    global_ctxt: Query<BoxedGlobalCtxt>,
    ongoing_codegen: Query<Box<dyn Any>>,
}

impl<'comp> Queries<'comp> {
    pub fn new(compiler: &'comp Compiler) -> Queries<'comp> {
        Queries {
            compiler,
            dep_graph_future: Default::default(),
            parse: Default::default(),
            crate_name: Default::default(),
            register_plugins: Default::default(),
            expansion: Default::default(),
            dep_graph: Default::default(),
            lower_to_hir: Default::default(),
            prepare_outputs: Default::default(),
            global_ctxt: Default::default(),
            ongoing_codegen: Default::default(),
        }
    }

    fn session(&self) -> &Lrc<Session> {
        &self.compiler.sess
    }
    fn codegen_backend(&self) -> &Lrc<Box<dyn CodegenBackend>> {
        &self.compiler.codegen_backend()
    }

    pub fn dep_graph_future(&self) -> Result<&Query<Option<DepGraphFuture>>> {
        self.dep_graph_future.compute(|| {
            Ok(if self.session().opts.build_dep_graph() {
                Some(rustc_incremental::load_dep_graph(self.session()))
            } else {
                None
            })
        })
    }

    pub fn parse(&self) -> Result<&Query<ast::Crate>> {
        self.parse.compute(|| {
            passes::parse(self.session(), &self.compiler.input).map_err(
                |mut parse_error| {
                    parse_error.emit();
                    ErrorReported
                },
            )
        })
    }

    pub fn register_plugins(&self) -> Result<&Query<(ast::Crate, Lrc<LintStore>)>> {
        self.register_plugins.compute(|| {
            let crate_name = self.crate_name()?.peek().clone();
            let krate = self.parse()?.take();

            let empty: &(dyn Fn(&Session, &mut lint::LintStore) + Sync + Send) = &|_, _| {};
            let result = passes::register_plugins(
                self.session(),
                &*self.codegen_backend().metadata_loader(),
                self.compiler.register_lints
                    .as_ref()
                    .map(|p| &**p)
                    .unwrap_or_else(|| empty),
                krate,
                &crate_name,
            );

            // Compute the dependency graph (in the background). We want to do
            // this as early as possible, to give the DepGraph maximum time to
            // load before dep_graph() is called, but it also can't happen
            // until after rustc_incremental::prepare_session_directory() is
            // called, which happens within passes::register_plugins().
            self.dep_graph_future().ok();

            result
        })
    }

    pub fn crate_name(&self) -> Result<&Query<String>> {
        self.crate_name.compute(|| {
            Ok(match self.compiler.crate_name {
                Some(ref crate_name) => crate_name.clone(),
                None => {
                    let parse_result = self.parse()?;
                    let krate = parse_result.peek();
                    rustc_codegen_utils::link::find_crate_name(
                        Some(self.session()),
                        &krate.attrs,
                        &self.compiler.input
                    )
                }
            })
        })
    }

    pub fn expansion(
        &self
    ) -> Result<&Query<(ast::Crate, Steal<Rc<RefCell<BoxedResolver>>>, Lrc<LintStore>)>> {
        self.expansion.compute(|| {
            let crate_name = self.crate_name()?.peek().clone();
            let (krate, lint_store) = self.register_plugins()?.take();
            passes::configure_and_expand(
                self.session().clone(),
                lint_store.clone(),
                self.codegen_backend().metadata_loader(),
                krate,
                &crate_name,
            ).map(|(krate, resolver)| {
                (krate, Steal::new(Rc::new(RefCell::new(resolver))), lint_store)
            })
        })
    }

    pub fn dep_graph(&self) -> Result<&Query<DepGraph>> {
        self.dep_graph.compute(|| {
            Ok(match self.dep_graph_future()?.take() {
                None => DepGraph::new_disabled(),
                Some(future) => {
                    let (prev_graph, prev_work_products) =
                        time(self.session(), "blocked while dep-graph loading finishes", || {
                            future.open().unwrap_or_else(|e| rustc_incremental::LoadResult::Error {
                                message: format!("could not decode incremental cache: {:?}", e),
                            }).open(self.session())
                        });
                    DepGraph::new(prev_graph, prev_work_products)
                }
            })
        })
    }

    pub fn lower_to_hir(
        &self,
    ) -> Result<&Query<(Steal<hir::map::Forest>, Steal<ResolverOutputs>)>> {
        self.lower_to_hir.compute(|| {
            let expansion_result = self.expansion()?;
            let peeked = expansion_result.peek();
            let krate = &peeked.0;
            let resolver = peeked.1.steal();
            let lint_store = &peeked.2;
            let hir = Steal::new(resolver.borrow_mut().access(|resolver| {
                passes::lower_to_hir(
                    self.session(),
                    lint_store,
                    resolver,
                    &*self.dep_graph()?.peek(),
                    &krate
                )
            })?);
            Ok((hir, Steal::new(BoxedResolver::to_resolver_outputs(resolver))))
        })
    }

    pub fn prepare_outputs(&self) -> Result<&Query<OutputFilenames>> {
        self.prepare_outputs.compute(|| {
            let expansion_result = self.expansion()?;
            let (krate, boxed_resolver, _) = &*expansion_result.peek();
            let crate_name = self.crate_name()?;
            let crate_name = crate_name.peek();
            passes::prepare_outputs(
                self.session(), self.compiler, &krate, &boxed_resolver, &crate_name
            )
        })
    }

    pub fn global_ctxt(&self) -> Result<&Query<BoxedGlobalCtxt>> {
        self.global_ctxt.compute(|| {
            let crate_name = self.crate_name()?.peek().clone();
            let outputs = self.prepare_outputs()?.peek().clone();
            let lint_store = self.expansion()?.peek().2.clone();
            let hir = self.lower_to_hir()?;
            let hir = hir.peek();
            let (hir_forest, resolver_outputs) = &*hir;
            Ok(passes::create_global_ctxt(
                self.compiler,
                lint_store,
                hir_forest.steal(),
                resolver_outputs.steal(),
                outputs,
                &crate_name))
        })
    }

    pub fn ongoing_codegen(&self) -> Result<&Query<Box<dyn Any>>> {
        self.ongoing_codegen.compute(|| {
            let outputs = self.prepare_outputs()?;
            self.global_ctxt()?.peek_mut().enter(|tcx| {
                tcx.analysis(LOCAL_CRATE).ok();

                // Don't do code generation if there were any errors
                self.session().compile_status()?;

                Ok(passes::start_codegen(
                    &***self.codegen_backend(),
                    tcx,
                    &*outputs.peek()
                ))
            })
        })
    }

    pub fn linker(self) -> Result<Linker> {
        let dep_graph = self.dep_graph()?;
        let prepare_outputs = self.prepare_outputs()?;
        let ongoing_codegen = self.ongoing_codegen()?;

        let sess = self.session().clone();
        let codegen_backend = self.codegen_backend().clone();

        Ok(Linker {
            sess,
            dep_graph: dep_graph.take(),
            prepare_outputs: prepare_outputs.take(),
            ongoing_codegen: ongoing_codegen.take(),
            codegen_backend,
        })
    }
}

pub struct Linker {
    sess: Lrc<Session>,
    dep_graph: DepGraph,
    prepare_outputs: OutputFilenames,
    ongoing_codegen: Box<dyn Any>,
    codegen_backend: Lrc<Box<dyn CodegenBackend>>,
}

impl Linker {
    pub fn link(self) -> Result<()> {
        self.codegen_backend.join_codegen_and_link(
            self.ongoing_codegen,
            &self.sess,
            &self.dep_graph,
            &self.prepare_outputs,
        ).map_err(|_| ErrorReported)
    }
}

impl Compiler {
    pub fn enter<'c, F, T>(&'c self, f: F) -> Result<T>
        where F: FnOnce(Queries<'c>) -> Result<T>
    {
        let queries = Queries::new(&self);
        f(queries)
    }

    // This method is different to all the other methods in `Compiler` because
    // it lacks a `Queries` entry. It's also not currently used. It does serve
    // as an example of how `Compiler` can be used, with additional steps added
    // between some passes. And see `rustc_driver::run_compiler` for a more
    // complex example.
    pub fn compile(&self) -> Result<()> {
        let linker = self.enter(|queries| {
            queries.prepare_outputs()?;

            if self.session().opts.output_types.contains_key(&OutputType::DepInfo)
                && self.session().opts.output_types.len() == 1
            {
                return Ok(None)
            }

            queries.global_ctxt()?;

            // Drop AST after creating GlobalCtxt to free memory.
            mem::drop(queries.expansion()?.take());

            queries.ongoing_codegen()?;

            // Drop GlobalCtxt after starting codegen to free memory.
            mem::drop(queries.global_ctxt()?.take());

            let linker = queries.linker()?;
            Ok(Some(linker))
        })?;

        if let Some(linker) = linker {
            linker.link()?
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>//! This module analyzes crates to find call sites that can serve as examples in the documentation.

use crate::clean;
use crate::config;
use crate::formats;
use crate::formats::renderer::FormatRenderer;
use crate::html::render::Context;

use rustc_data_structures::fx::FxHashMap;
use rustc_hir::{
    self as hir,
    intravisit::{self, Visitor},
};
use rustc_interface::interface;
use rustc_macros::{Decodable, Encodable};
use rustc_middle::hir::map::Map;
use rustc_middle::ty::{self, TyCtxt};
use rustc_serialize::{
    opaque::{Decoder, FileEncoder},
    Decodable, Encodable,
};
use rustc_session::getopts;
use rustc_span::{
    def_id::{CrateNum, DefPathHash, LOCAL_CRATE},
    edition::Edition,
    BytePos, FileName, SourceFile,
};

use std::fs;
use std::path::PathBuf;

#[derive(Debug, Clone)]
crate struct ScrapeExamplesOptions {
    output_path: PathBuf,
    target_crates: Vec<String>,
}

impl ScrapeExamplesOptions {
    crate fn new(
        matches: &getopts::Matches,
        diag: &rustc_errors::Handler,
    ) -> Result<Option<Self>, i32> {
        let output_path = matches.opt_str("scrape-examples-output-path");
        let target_crates = matches.opt_strs("scrape-examples-target-crate");
        match (output_path, !target_crates.is_empty()) {
            (Some(output_path), true) => Ok(Some(ScrapeExamplesOptions {
                output_path: PathBuf::from(output_path),
                target_crates,
            })),
            (Some(_), false) | (None, true) => {
                diag.err("must use --scrape-examples-output-path and --scrape-examples-target-crate together");
                Err(1)
            }
            (None, false) => Ok(None),
        }
    }
}

#[derive(Encodable, Decodable, Debug, Clone)]
crate struct SyntaxRange {
    crate byte_span: (u32, u32),
    crate line_span: (usize, usize),
}

impl SyntaxRange {
    fn new(span: rustc_span::Span, file: &SourceFile) -> Self {
        let get_pos = |bytepos: BytePos| file.original_relative_byte_pos(bytepos).0;
        let get_line = |bytepos: BytePos| file.lookup_line(bytepos).unwrap();

        SyntaxRange {
            byte_span: (get_pos(span.lo()), get_pos(span.hi())),
            line_span: (get_line(span.lo()), get_line(span.hi())),
        }
    }
}

#[derive(Encodable, Decodable, Debug, Clone)]
crate struct CallLocation {
    crate call_expr: SyntaxRange,
    crate enclosing_item: SyntaxRange,
}

impl CallLocation {
    fn new(
        expr_span: rustc_span::Span,
        enclosing_item_span: rustc_span::Span,
        source_file: &SourceFile,
    ) -> Self {
        CallLocation {
            call_expr: SyntaxRange::new(expr_span, source_file),
            enclosing_item: SyntaxRange::new(enclosing_item_span, source_file),
        }
    }
}

#[derive(Encodable, Decodable, Debug, Clone)]
crate struct CallData {
    crate locations: Vec<CallLocation>,
    crate url: String,
    crate display_name: String,
    crate edition: Edition,
}

crate type FnCallLocations = FxHashMap<PathBuf, CallData>;
crate type AllCallLocations = FxHashMap<DefPathHash, FnCallLocations>;

/// Visitor for traversing a crate and finding instances of function calls.
struct FindCalls<'a, 'tcx> {
    tcx: TyCtxt<'tcx>,
    map: Map<'tcx>,
    cx: Context<'tcx>,
    target_crates: Vec<CrateNum>,
    calls: &'a mut AllCallLocations,
}

impl<'a, 'tcx> Visitor<'tcx> for FindCalls<'a, 'tcx>
where
    'tcx: 'a,
{
    type Map = Map<'tcx>;

    fn nested_visit_map(&mut self) -> intravisit::NestedVisitorMap<Self::Map> {
        intravisit::NestedVisitorMap::OnlyBodies(self.map)
    }

    fn visit_expr(&mut self, ex: &'tcx hir::Expr<'tcx>) {
        intravisit::walk_expr(self, ex);

        let tcx = self.tcx;

        // If we visit an item that contains an expression outside a function body,
        // then we need to exit before calling typeck (which will panic). See
        // test/run-make/rustdoc-scrape-examples-invalid-expr for an example.
        let hir = tcx.hir();
        let owner = hir.local_def_id_to_hir_id(ex.hir_id.owner);
        if hir.maybe_body_owned_by(owner).is_none() {
            return;
        }

        // Get type of function if expression is a function call
        let (ty, span) = match ex.kind {
            hir::ExprKind::Call(f, _) => {
                let types = tcx.typeck(ex.hir_id.owner);

                if let Some(ty) = types.node_type_opt(f.hir_id) {
                    (ty, ex.span)
                } else {
                    return;
                }
            }
            hir::ExprKind::MethodCall(_, _, _, span) => {
                let types = tcx.typeck(ex.hir_id.owner);
                let def_id = if let Some(def_id) = types.type_dependent_def_id(ex.hir_id) {
                    def_id
                } else {
                    return;
                };
                (tcx.type_of(def_id), span)
            }
            _ => {
                return;
            }
        };

        // If this span comes from a macro expansion, then the source code may not actually show
        // a use of the given item, so it would be a poor example. Hence, we skip all uses in macros.
        if span.from_expansion() {
            trace!("Rejecting expr from macro: {:?}", span);
            return;
        }

        // If the enclosing item has a span coming from a proc macro, then we also don't want to include
        // the example.
        let enclosing_item_span = tcx.hir().span_with_body(tcx.hir().get_parent_item(ex.hir_id));
        if enclosing_item_span.from_expansion() {
            trace!("Rejecting expr ({:?}) from macro item: {:?}", span, enclosing_item_span);
            return;
        }

        assert!(
            enclosing_item_span.contains(span),
            "Attempted to scrape call at [{:?}] whose enclosing item [{:?}] doesn't contain the span of the call.",
            span,
            enclosing_item_span
        );

        // Save call site if the function resolves to a concrete definition
        if let ty::FnDef(def_id, _) = ty.kind() {
            if self.target_crates.iter().all(|krate| *krate != def_id.krate) {
                trace!("Rejecting expr from crate not being documented: {:?}", span);
                return;
            }

            let file = tcx.sess.source_map().lookup_char_pos(span.lo()).file;
            let file_path = match file.name.clone() {
                FileName::Real(real_filename) => real_filename.into_local_path(),
                _ => None,
            };

            if let Some(file_path) = file_path {
                let abs_path = fs::canonicalize(file_path.clone()).unwrap();
                let cx = &self.cx;
                let mk_call_data = || {
                    let clean_span = crate::clean::types::Span::new(span);
                    let url = cx.href_from_span(clean_span, false).unwrap();
                    let display_name = file_path.display().to_string();
                    let edition = span.edition();
                    CallData { locations: Vec::new(), url, display_name, edition }
                };

                let fn_key = tcx.def_path_hash(*def_id);
                let fn_entries = self.calls.entry(fn_key).or_default();

                trace!("Including expr: {:?}", span);
                let location = CallLocation::new(span, enclosing_item_span, &file);
                fn_entries.entry(abs_path).or_insert_with(mk_call_data).locations.push(location);
            }
        }
    }
}

crate fn run(
    krate: clean::Crate,
    renderopts: config::RenderOptions,
    cache: formats::cache::Cache,
    tcx: TyCtxt<'_>,
    options: ScrapeExamplesOptions,
) -> interface::Result<()> {
    let inner = move || -> Result<(), String> {
        // Generates source files for examples
        let (cx, _) = Context::init(krate, renderopts, cache, tcx).map_err(|e| e.to_string())?;

        // Collect CrateIds corresponding to provided target crates
        // If two different versions of the crate in the dependency tree, then examples will be collcted from both.
        let all_crates = tcx
            .crates(())
            .iter()
            .chain([&LOCAL_CRATE])
            .map(|crate_num| (crate_num, tcx.crate_name(*crate_num)))
            .collect::<Vec<_>>();
        let target_crates = options
            .target_crates
            .into_iter()
            .map(|target| all_crates.iter().filter(move |(_, name)| name.as_str() == target))
            .flatten()
            .map(|(crate_num, _)| **crate_num)
            .collect::<Vec<_>>();

        debug!("All crates in TyCtxt: {:?}", all_crates);
        debug!("Scrape examples target_crates: {:?}", target_crates);

        // Run call-finder on all items
        let mut calls = FxHashMap::default();
        let mut finder = FindCalls { calls: &mut calls, tcx, map: tcx.hir(), cx, target_crates };
        tcx.hir().visit_all_item_likes(&mut finder.as_deep_visitor());

        // Sort call locations within a given file in document order
        for fn_calls in calls.values_mut() {
            for file_calls in fn_calls.values_mut() {
                file_calls.locations.sort_by_key(|loc| loc.call_expr.byte_span.0);
            }
        }

        // Save output to provided path
        let mut encoder = FileEncoder::new(options.output_path).map_err(|e| e.to_string())?;
        calls.encode(&mut encoder).map_err(|e| e.to_string())?;
        encoder.flush().map_err(|e| e.to_string())?;

        Ok(())
    };

    if let Err(e) = inner() {
        tcx.sess.fatal(&e);
    }

    Ok(())
}

// Note: the Handler must be passed in explicitly because sess isn't available while parsing options
crate fn load_call_locations(
    with_examples: Vec<String>,
    diag: &rustc_errors::Handler,
) -> Result<AllCallLocations, i32> {
    let inner = || {
        let mut all_calls: AllCallLocations = FxHashMap::default();
        for path in with_examples {
            let bytes = fs::read(&path).map_err(|e| format!("{} (for path {})", e, path))?;
            let mut decoder = Decoder::new(&bytes, 0);
            let calls = AllCallLocations::decode(&mut decoder)?;

            for (function, fn_calls) in calls.into_iter() {
                all_calls.entry(function).or_default().extend(fn_calls.into_iter());
            }
        }

        Ok(all_calls)
    };

    inner().map_err(|e: String| {
        diag.err(&format!("failed to load examples: {}", e));
        1
    })
}
<commit_msg>Add trace statements<commit_after>//! This module analyzes crates to find call sites that can serve as examples in the documentation.

use crate::clean;
use crate::config;
use crate::formats;
use crate::formats::renderer::FormatRenderer;
use crate::html::render::Context;

use rustc_data_structures::fx::FxHashMap;
use rustc_hir::{
    self as hir,
    intravisit::{self, Visitor},
};
use rustc_interface::interface;
use rustc_macros::{Decodable, Encodable};
use rustc_middle::hir::map::Map;
use rustc_middle::ty::{self, TyCtxt};
use rustc_serialize::{
    opaque::{Decoder, FileEncoder},
    Decodable, Encodable,
};
use rustc_session::getopts;
use rustc_span::{
    def_id::{CrateNum, DefPathHash, LOCAL_CRATE},
    edition::Edition,
    BytePos, FileName, SourceFile,
};

use std::fs;
use std::path::PathBuf;

#[derive(Debug, Clone)]
crate struct ScrapeExamplesOptions {
    output_path: PathBuf,
    target_crates: Vec<String>,
}

impl ScrapeExamplesOptions {
    crate fn new(
        matches: &getopts::Matches,
        diag: &rustc_errors::Handler,
    ) -> Result<Option<Self>, i32> {
        let output_path = matches.opt_str("scrape-examples-output-path");
        let target_crates = matches.opt_strs("scrape-examples-target-crate");
        match (output_path, !target_crates.is_empty()) {
            (Some(output_path), true) => Ok(Some(ScrapeExamplesOptions {
                output_path: PathBuf::from(output_path),
                target_crates,
            })),
            (Some(_), false) | (None, true) => {
                diag.err("must use --scrape-examples-output-path and --scrape-examples-target-crate together");
                Err(1)
            }
            (None, false) => Ok(None),
        }
    }
}

#[derive(Encodable, Decodable, Debug, Clone)]
crate struct SyntaxRange {
    crate byte_span: (u32, u32),
    crate line_span: (usize, usize),
}

impl SyntaxRange {
    fn new(span: rustc_span::Span, file: &SourceFile) -> Self {
        let get_pos = |bytepos: BytePos| file.original_relative_byte_pos(bytepos).0;
        let get_line = |bytepos: BytePos| file.lookup_line(bytepos).unwrap();

        SyntaxRange {
            byte_span: (get_pos(span.lo()), get_pos(span.hi())),
            line_span: (get_line(span.lo()), get_line(span.hi())),
        }
    }
}

#[derive(Encodable, Decodable, Debug, Clone)]
crate struct CallLocation {
    crate call_expr: SyntaxRange,
    crate enclosing_item: SyntaxRange,
}

impl CallLocation {
    fn new(
        expr_span: rustc_span::Span,
        enclosing_item_span: rustc_span::Span,
        source_file: &SourceFile,
    ) -> Self {
        CallLocation {
            call_expr: SyntaxRange::new(expr_span, source_file),
            enclosing_item: SyntaxRange::new(enclosing_item_span, source_file),
        }
    }
}

#[derive(Encodable, Decodable, Debug, Clone)]
crate struct CallData {
    crate locations: Vec<CallLocation>,
    crate url: String,
    crate display_name: String,
    crate edition: Edition,
}

crate type FnCallLocations = FxHashMap<PathBuf, CallData>;
crate type AllCallLocations = FxHashMap<DefPathHash, FnCallLocations>;

/// Visitor for traversing a crate and finding instances of function calls.
struct FindCalls<'a, 'tcx> {
    tcx: TyCtxt<'tcx>,
    map: Map<'tcx>,
    cx: Context<'tcx>,
    target_crates: Vec<CrateNum>,
    calls: &'a mut AllCallLocations,
}

impl<'a, 'tcx> Visitor<'tcx> for FindCalls<'a, 'tcx>
where
    'tcx: 'a,
{
    type Map = Map<'tcx>;

    fn nested_visit_map(&mut self) -> intravisit::NestedVisitorMap<Self::Map> {
        intravisit::NestedVisitorMap::OnlyBodies(self.map)
    }

    fn visit_expr(&mut self, ex: &'tcx hir::Expr<'tcx>) {
        intravisit::walk_expr(self, ex);

        let tcx = self.tcx;

        // If we visit an item that contains an expression outside a function body,
        // then we need to exit before calling typeck (which will panic). See
        // test/run-make/rustdoc-scrape-examples-invalid-expr for an example.
        let hir = tcx.hir();
        let owner = hir.local_def_id_to_hir_id(ex.hir_id.owner);
        if hir.maybe_body_owned_by(owner).is_none() {
            return;
        }

        // Get type of function if expression is a function call
        let (ty, span) = match ex.kind {
            hir::ExprKind::Call(f, _) => {
                let types = tcx.typeck(ex.hir_id.owner);

                if let Some(ty) = types.node_type_opt(f.hir_id) {
                    (ty, ex.span)
                } else {
                    trace!("node_type_opt({}) = None", f.hir_id);
                    return;
                }
            }
            hir::ExprKind::MethodCall(_, _, _, span) => {
                let types = tcx.typeck(ex.hir_id.owner);
                let def_id = if let Some(def_id) = types.type_dependent_def_id(ex.hir_id) {
                    def_id
                } else {
                    trace!("type_dependent_def_id({}) = None", ex.hir_id);
                    return;
                };
                (tcx.type_of(def_id), span)
            }
            _ => {
                return;
            }
        };

        // If this span comes from a macro expansion, then the source code may not actually show
        // a use of the given item, so it would be a poor example. Hence, we skip all uses in macros.
        if span.from_expansion() {
            trace!("Rejecting expr from macro: {:?}", span);
            return;
        }

        // If the enclosing item has a span coming from a proc macro, then we also don't want to include
        // the example.
        let enclosing_item_span = tcx.hir().span_with_body(tcx.hir().get_parent_item(ex.hir_id));
        if enclosing_item_span.from_expansion() {
            trace!("Rejecting expr ({:?}) from macro item: {:?}", span, enclosing_item_span);
            return;
        }

        assert!(
            enclosing_item_span.contains(span),
            "Attempted to scrape call at [{:?}] whose enclosing item [{:?}] doesn't contain the span of the call.",
            span,
            enclosing_item_span
        );

        // Save call site if the function resolves to a concrete definition
        if let ty::FnDef(def_id, _) = ty.kind() {
            if self.target_crates.iter().all(|krate| *krate != def_id.krate) {
                trace!("Rejecting expr from crate not being documented: {:?}", span);
                return;
            }

            let file = tcx.sess.source_map().lookup_char_pos(span.lo()).file;
            let file_path = match file.name.clone() {
                FileName::Real(real_filename) => real_filename.into_local_path(),
                _ => None,
            };

            if let Some(file_path) = file_path {
                let abs_path = fs::canonicalize(file_path.clone()).unwrap();
                let cx = &self.cx;
                let mk_call_data = || {
                    let clean_span = crate::clean::types::Span::new(span);
                    let url = cx.href_from_span(clean_span, false).unwrap();
                    let display_name = file_path.display().to_string();
                    let edition = span.edition();
                    CallData { locations: Vec::new(), url, display_name, edition }
                };

                let fn_key = tcx.def_path_hash(*def_id);
                let fn_entries = self.calls.entry(fn_key).or_default();

                trace!("Including expr: {:?}", span);
                let location = CallLocation::new(span, enclosing_item_span, &file);
                fn_entries.entry(abs_path).or_insert_with(mk_call_data).locations.push(location);
            }
        }
    }
}

crate fn run(
    krate: clean::Crate,
    renderopts: config::RenderOptions,
    cache: formats::cache::Cache,
    tcx: TyCtxt<'_>,
    options: ScrapeExamplesOptions,
) -> interface::Result<()> {
    let inner = move || -> Result<(), String> {
        // Generates source files for examples
        let (cx, _) = Context::init(krate, renderopts, cache, tcx).map_err(|e| e.to_string())?;

        // Collect CrateIds corresponding to provided target crates
        // If two different versions of the crate in the dependency tree, then examples will be collcted from both.
        let all_crates = tcx
            .crates(())
            .iter()
            .chain([&LOCAL_CRATE])
            .map(|crate_num| (crate_num, tcx.crate_name(*crate_num)))
            .collect::<Vec<_>>();
        let target_crates = options
            .target_crates
            .into_iter()
            .map(|target| all_crates.iter().filter(move |(_, name)| name.as_str() == target))
            .flatten()
            .map(|(crate_num, _)| **crate_num)
            .collect::<Vec<_>>();

        debug!("All crates in TyCtxt: {:?}", all_crates);
        debug!("Scrape examples target_crates: {:?}", target_crates);

        // Run call-finder on all items
        let mut calls = FxHashMap::default();
        let mut finder = FindCalls { calls: &mut calls, tcx, map: tcx.hir(), cx, target_crates };
        tcx.hir().visit_all_item_likes(&mut finder.as_deep_visitor());

        // Sort call locations within a given file in document order
        for fn_calls in calls.values_mut() {
            for file_calls in fn_calls.values_mut() {
                file_calls.locations.sort_by_key(|loc| loc.call_expr.byte_span.0);
            }
        }

        // Save output to provided path
        let mut encoder = FileEncoder::new(options.output_path).map_err(|e| e.to_string())?;
        calls.encode(&mut encoder).map_err(|e| e.to_string())?;
        encoder.flush().map_err(|e| e.to_string())?;

        Ok(())
    };

    if let Err(e) = inner() {
        tcx.sess.fatal(&e);
    }

    Ok(())
}

// Note: the Handler must be passed in explicitly because sess isn't available while parsing options
crate fn load_call_locations(
    with_examples: Vec<String>,
    diag: &rustc_errors::Handler,
) -> Result<AllCallLocations, i32> {
    let inner = || {
        let mut all_calls: AllCallLocations = FxHashMap::default();
        for path in with_examples {
            let bytes = fs::read(&path).map_err(|e| format!("{} (for path {})", e, path))?;
            let mut decoder = Decoder::new(&bytes, 0);
            let calls = AllCallLocations::decode(&mut decoder)?;

            for (function, fn_calls) in calls.into_iter() {
                all_calls.entry(function).or_default().extend(fn_calls.into_iter());
            }
        }

        Ok(all_calls)
    };

    inner().map_err(|e: String| {
        diag.err(&format!("failed to load examples: {}", e));
        1
    })
}
<|endoftext|>


<commit_before>use rocket_contrib::Json;
use rocket::request::FromParam;
use rocket::http::RawStr;
use ldap_hook::user::LdapUser;
use ldap_hook::config::LdapConfig;
use ldap_hook::ldap_search;

#[derive(PartialEq, Eq)]
pub enum LdapQueryType {
	Nick,
	UID,
}

impl<'a> FromParam<'a> for LdapQueryType {
	type Error = &'a RawStr;
	fn from_param(param: &'a RawStr) -> Result<Self, Self::Error> {
		if param == "nick" {
			Ok(LdapQueryType::Nick)
		} else if param == "uid" {
			Ok(LdapQueryType::UID)
		} else {
			Err(param)
		}
	}
}

// Dummy endpoints so that the front-end won't complain
#[get("/search/<query_type>")]
pub fn empty_search(query_type: LdapQueryType) -> Json<Vec<String>> { Json(vec![]) }

#[get("/search/<query_type>/<query>")]
pub fn search(query_type: LdapQueryType, query: String) -> Json<Vec<LdapUser>> {
	let id_list: Vec<String> = query.split(",").map(|s| s.to_string()).collect();

	let mut result_list: Vec<LdapUser> = vec![];

	// TODO: Load the config somewhere more sensible
	let ldap_config = match LdapConfig::load("config.toml") {
		Ok(ldap_config) => ldap_config,
		Err(msg) => {
			panic!("Could not load config.toml: {}", msg);
		}
	};

	for id in &id_list{
		let filter = if query_type == LdapQueryType::UID {
			format!("(uid={})", id)
		} else {
			format!("(nickname=*{}*)", id)
		};
		let res = ldap_search(&ldap_config, filter.as_str());

		match res.ok() {
			None => { break; }
			Some(list) => {
				for entry in list {
					if !result_list.contains(&entry) {
						result_list.push(entry)
					}
				}
			}
		}
	}

	Json(result_list)
}
<commit_msg>Prettify LDAP search filter<commit_after>use rocket_contrib::Json;
use rocket::request::FromParam;
use rocket::http::RawStr;
use ldap_hook::user::LdapUser;
use ldap_hook::config::LdapConfig;
use ldap_hook::ldap_search;

#[derive(PartialEq, Eq)]
pub enum LdapQueryType {
	Nick,
	UID,
}

impl<'a> FromParam<'a> for LdapQueryType {
	type Error = &'a RawStr;
	fn from_param(param: &'a RawStr) -> Result<Self, Self::Error> {
		if param == "nick" {
			Ok(LdapQueryType::Nick)
		} else if param == "uid" {
			Ok(LdapQueryType::UID)
		} else {
			Err(param)
		}
	}
}

// Dummy endpoints so that the front-end won't complain
#[get("/search/<query_type>")]
pub fn empty_search(query_type: LdapQueryType) -> Json<Vec<String>> { Json(vec![]) }

#[get("/search/<query_type>/<query>")]
pub fn search(query_type: LdapQueryType, query: String) -> Json<Vec<LdapUser>> {
	let id_list: Vec<String> = query.split(",").map(|s| s.to_string()).collect();

	let mut result_list: Vec<LdapUser> = vec![];

	// TODO: Load the config somewhere more sensible
	let ldap_config = match LdapConfig::load("config.toml") {
		Ok(ldap_config) => ldap_config,
		Err(msg) => {
			panic!("Could not load config.toml: {}", msg);
		}
	};

	for id in &id_list{
		let filter = match query_type {
			LdapQueryType::UID  => format!("(uid={})", id),
			LdapQueryType::Nick => format!("(nickname=*{}*)", id), // Basic fuzzy search when querying by nicknames
		};
		let res = ldap_search(&ldap_config, filter.as_str());

		match res.ok() {
			None => { break; }
			Some(list) => {
				for entry in list {
					if !result_list.contains(&entry) {
						result_list.push(entry)
					}
				}
			}
		}
	}

	Json(result_list)
}
<|endoftext|>


<commit_before>use indexmap::IndexMap;
use serde::ser::SerializeMap;
use serde::{de, ser};

use std::fmt;

use ast::InputValue;
use executor::ExecutionError;
use parser::{ParseError, SourcePosition, Spanning};
use validation::RuleError;
use {GraphQLError, Value};

impl ser::Serialize for ExecutionError {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        let mut map = serializer.serialize_map(Some(4))?;

        map.serialize_key("message")?;
        map.serialize_value(self.error().message())?;

        let locations = vec![self.location()];
        map.serialize_key("locations")?;
        map.serialize_value(&locations)?;

        map.serialize_key("path")?;
        map.serialize_value(self.path())?;

        if !self.error().data().is_null() {
            map.serialize_key("data")?;
            map.serialize_value(self.error().data())?;
        }

        map.end()
    }
}

impl<'a> ser::Serialize for GraphQLError<'a> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        match *self {
            GraphQLError::ParseError(ref err) => vec![err].serialize(serializer),
            GraphQLError::ValidationError(ref errs) => errs.serialize(serializer),
            GraphQLError::NoOperationProvided => {
                serializer.serialize_str("Must provide an operation")
            }
            GraphQLError::MultipleOperationsProvided => serializer
                .serialize_str("Must provide operation name if query contains multiple operations"),
            GraphQLError::UnknownOperationName => serializer.serialize_str("Unknown operation"),
        }
    }
}

impl<'de> de::Deserialize<'de> for InputValue {
    fn deserialize<D>(deserializer: D) -> Result<InputValue, D::Error>
    where
        D: de::Deserializer<'de>,
    {
        struct InputValueVisitor;

        impl<'de> de::Visitor<'de> for InputValueVisitor {
            type Value = InputValue;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("a valid input value")
            }

            fn visit_bool<E>(self, value: bool) -> Result<InputValue, E> {
                Ok(InputValue::boolean(value))
            }

            fn visit_i64<E>(self, value: i64) -> Result<InputValue, E>
            where
                E: de::Error,
            {
                if value >= i64::from(i32::min_value()) && value <= i64::from(i32::max_value()) {
                    Ok(InputValue::int(value as i32))
                } else {
                    // Browser's JSON.stringify serialize all numbers having no
                    // fractional part as integers (no decimal point), so we
                    // must parse large integers as floating point otherwise
                    // we would error on transferring large floating point
                    // numbers.
                    Ok(InputValue::float(value as f64))
                }
            }

            fn visit_u64<E>(self, value: u64) -> Result<InputValue, E>
            where
                E: de::Error,
            {
                if value <= i32::max_value() as u64 {
                    self.visit_i64(value as i64)
                } else {
                    // Browser's JSON.stringify serialize all numbers having no
                    // fractional part as integers (no decimal point), so we
                    // must parse large integers as floating point otherwise
                    // we would error on transferring large floating point
                    // numbers.
                    Ok(InputValue::float(value as f64))
                }
            }

            fn visit_f64<E>(self, value: f64) -> Result<InputValue, E> {
                Ok(InputValue::float(value))
            }

            fn visit_str<E>(self, value: &str) -> Result<InputValue, E>
            where
                E: de::Error,
            {
                self.visit_string(value.into())
            }

            fn visit_string<E>(self, value: String) -> Result<InputValue, E> {
                Ok(InputValue::string(value))
            }

            fn visit_none<E>(self) -> Result<InputValue, E> {
                Ok(InputValue::null())
            }

            fn visit_unit<E>(self) -> Result<InputValue, E> {
                Ok(InputValue::null())
            }

            fn visit_seq<V>(self, mut visitor: V) -> Result<InputValue, V::Error>
            where
                V: de::SeqAccess<'de>,
            {
                let mut values = Vec::new();

                while let Some(el) = visitor.next_element()? {
                    values.push(el);
                }

                Ok(InputValue::list(values))
            }

            fn visit_map<V>(self, mut visitor: V) -> Result<InputValue, V::Error>
            where
                V: de::MapAccess<'de>,
            {
                let mut values: IndexMap<String, InputValue> = IndexMap::new();

                while let Some((key, value)) = visitor.next_entry()? {
                    values.insert(key, value);
                }

                Ok(InputValue::object(values))
            }
        }

        deserializer.deserialize_any(InputValueVisitor)
    }
}

impl ser::Serialize for InputValue {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        match *self {
            InputValue::Null | InputValue::Variable(_) => serializer.serialize_unit(),
            InputValue::Int(v) => serializer.serialize_i64(i64::from(v)),
            InputValue::Float(v) => serializer.serialize_f64(v),
            InputValue::String(ref v) | InputValue::Enum(ref v) => serializer.serialize_str(v),
            InputValue::Boolean(v) => serializer.serialize_bool(v),
            InputValue::List(ref v) => v.iter()
                .map(|x| x.item.clone())
                .collect::<Vec<_>>()
                .serialize(serializer),
            InputValue::Object(ref v) => v.iter()
                .map(|&(ref k, ref v)| (k.item.clone(), v.item.clone()))
                .collect::<IndexMap<_, _>>()
                .serialize(serializer),
        }
    }
}

impl ser::Serialize for RuleError {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        let mut map = serializer.serialize_map(Some(2))?;

        map.serialize_key("message")?;
        map.serialize_value(self.message())?;

        map.serialize_key("locations")?;
        map.serialize_value(self.locations())?;

        map.end()
    }
}

impl ser::Serialize for SourcePosition {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        let mut map = serializer.serialize_map(Some(2))?;

        let line = self.line() + 1;
        map.serialize_key("line")?;
        map.serialize_value(&line)?;

        let column = self.column() + 1;
        map.serialize_key("column")?;
        map.serialize_value(&column)?;

        map.end()
    }
}

impl<'a> ser::Serialize for Spanning<ParseError<'a>> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        let mut map = serializer.serialize_map(Some(2))?;

        let message = format!("{}", self.item);
        map.serialize_key("message")?;
        map.serialize_value(&message)?;

        let mut location = IndexMap::new();
        location.insert("line".to_owned(), self.start.line() + 1);
        location.insert("column".to_owned(), self.start.column() + 1);

        let locations = vec![location];

        map.serialize_key("locations")?;
        map.serialize_value(&locations)?;

        map.end()
    }
}

impl ser::Serialize for Value {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        match *self {
            Value::Null => serializer.serialize_unit(),
            Value::Int(v) => serializer.serialize_i64(i64::from(v)),
            Value::Float(v) => serializer.serialize_f64(v),
            Value::String(ref v) => serializer.serialize_str(v),
            Value::Boolean(v) => serializer.serialize_bool(v),
            Value::List(ref v) => v.serialize(serializer),
            Value::Object(ref v) => v.serialize(serializer),
        }
    }
}

#[cfg(test)]
mod tests {
    use serde_json::from_str;
    use ast::InputValue;

    #[test]
    fn int() {
        assert_eq!(from_str::<InputValue>("1235").unwrap(),
                   InputValue::int(1235));
    }

    #[test]
    fn float() {
        assert_eq!(from_str::<InputValue>("2.0").unwrap(),
                   InputValue::float(2.0));
        // large value without a decimal part is also float
        assert_eq!(from_str::<InputValue>("123567890123").unwrap(),
                   InputValue::float(123567890123.0));
    }
}
<commit_msg>Consistent error serializing for GraphQLError (#207)<commit_after>use indexmap::IndexMap;
use serde::ser::SerializeMap;
use serde::{de, ser};

use std::fmt;

use ast::InputValue;
use executor::ExecutionError;
use parser::{ParseError, SourcePosition, Spanning};
use validation::RuleError;
use {GraphQLError, Value};

#[derive(Serialize)]
struct SerializeHelper {
    message: &'static str,
}

impl ser::Serialize for ExecutionError {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        let mut map = serializer.serialize_map(Some(4))?;

        map.serialize_key("message")?;
        map.serialize_value(self.error().message())?;

        let locations = vec![self.location()];
        map.serialize_key("locations")?;
        map.serialize_value(&locations)?;

        map.serialize_key("path")?;
        map.serialize_value(self.path())?;

        if !self.error().data().is_null() {
            map.serialize_key("data")?;
            map.serialize_value(self.error().data())?;
        }

        map.end()
    }
}

impl<'a> ser::Serialize for GraphQLError<'a> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        match *self {
            GraphQLError::ParseError(ref err) => vec![err].serialize(serializer),
            GraphQLError::ValidationError(ref errs) => errs.serialize(serializer),
            GraphQLError::NoOperationProvided => {
                [
                    SerializeHelper { message: "Must provide an operation" }
                ].serialize(serializer)
            }
            GraphQLError::MultipleOperationsProvided => {
                [SerializeHelper {
                    message: "Must provide operation name \
                    if query contains multiple operations",
                }].serialize(serializer)
            }
            GraphQLError::UnknownOperationName => {
                [
                    SerializeHelper { message: "Unknown operation" }
                ].serialize(serializer)
            }
        }
    }
}

impl<'de> de::Deserialize<'de> for InputValue {
    fn deserialize<D>(deserializer: D) -> Result<InputValue, D::Error>
    where
        D: de::Deserializer<'de>,
    {
        struct InputValueVisitor;

        impl<'de> de::Visitor<'de> for InputValueVisitor {
            type Value = InputValue;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("a valid input value")
            }

            fn visit_bool<E>(self, value: bool) -> Result<InputValue, E> {
                Ok(InputValue::boolean(value))
            }

            fn visit_i64<E>(self, value: i64) -> Result<InputValue, E>
            where
                E: de::Error,
            {
                if value >= i64::from(i32::min_value()) && value <= i64::from(i32::max_value()) {
                    Ok(InputValue::int(value as i32))
                } else {
                    // Browser's JSON.stringify serialize all numbers having no
                    // fractional part as integers (no decimal point), so we
                    // must parse large integers as floating point otherwise
                    // we would error on transferring large floating point
                    // numbers.
                    Ok(InputValue::float(value as f64))
                }
            }

            fn visit_u64<E>(self, value: u64) -> Result<InputValue, E>
            where
                E: de::Error,
            {
                if value <= i32::max_value() as u64 {
                    self.visit_i64(value as i64)
                } else {
                    // Browser's JSON.stringify serialize all numbers having no
                    // fractional part as integers (no decimal point), so we
                    // must parse large integers as floating point otherwise
                    // we would error on transferring large floating point
                    // numbers.
                    Ok(InputValue::float(value as f64))
                }
            }

            fn visit_f64<E>(self, value: f64) -> Result<InputValue, E> {
                Ok(InputValue::float(value))
            }

            fn visit_str<E>(self, value: &str) -> Result<InputValue, E>
            where
                E: de::Error,
            {
                self.visit_string(value.into())
            }

            fn visit_string<E>(self, value: String) -> Result<InputValue, E> {
                Ok(InputValue::string(value))
            }

            fn visit_none<E>(self) -> Result<InputValue, E> {
                Ok(InputValue::null())
            }

            fn visit_unit<E>(self) -> Result<InputValue, E> {
                Ok(InputValue::null())
            }

            fn visit_seq<V>(self, mut visitor: V) -> Result<InputValue, V::Error>
            where
                V: de::SeqAccess<'de>,
            {
                let mut values = Vec::new();

                while let Some(el) = visitor.next_element()? {
                    values.push(el);
                }

                Ok(InputValue::list(values))
            }

            fn visit_map<V>(self, mut visitor: V) -> Result<InputValue, V::Error>
            where
                V: de::MapAccess<'de>,
            {
                let mut values: IndexMap<String, InputValue> = IndexMap::new();

                while let Some((key, value)) = visitor.next_entry()? {
                    values.insert(key, value);
                }

                Ok(InputValue::object(values))
            }
        }

        deserializer.deserialize_any(InputValueVisitor)
    }
}

impl ser::Serialize for InputValue {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        match *self {
            InputValue::Null | InputValue::Variable(_) => serializer.serialize_unit(),
            InputValue::Int(v) => serializer.serialize_i64(i64::from(v)),
            InputValue::Float(v) => serializer.serialize_f64(v),
            InputValue::String(ref v) | InputValue::Enum(ref v) => serializer.serialize_str(v),
            InputValue::Boolean(v) => serializer.serialize_bool(v),
            InputValue::List(ref v) => v.iter()
                .map(|x| x.item.clone())
                .collect::<Vec<_>>()
                .serialize(serializer),
            InputValue::Object(ref v) => v.iter()
                .map(|&(ref k, ref v)| (k.item.clone(), v.item.clone()))
                .collect::<IndexMap<_, _>>()
                .serialize(serializer),
        }
    }
}

impl ser::Serialize for RuleError {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        let mut map = serializer.serialize_map(Some(2))?;

        map.serialize_key("message")?;
        map.serialize_value(self.message())?;

        map.serialize_key("locations")?;
        map.serialize_value(self.locations())?;

        map.end()
    }
}

impl ser::Serialize for SourcePosition {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        let mut map = serializer.serialize_map(Some(2))?;

        let line = self.line() + 1;
        map.serialize_key("line")?;
        map.serialize_value(&line)?;

        let column = self.column() + 1;
        map.serialize_key("column")?;
        map.serialize_value(&column)?;

        map.end()
    }
}

impl<'a> ser::Serialize for Spanning<ParseError<'a>> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        let mut map = serializer.serialize_map(Some(2))?;

        let message = format!("{}", self.item);
        map.serialize_key("message")?;
        map.serialize_value(&message)?;

        let mut location = IndexMap::new();
        location.insert("line".to_owned(), self.start.line() + 1);
        location.insert("column".to_owned(), self.start.column() + 1);

        let locations = vec![location];

        map.serialize_key("locations")?;
        map.serialize_value(&locations)?;

        map.end()
    }
}

impl ser::Serialize for Value {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        match *self {
            Value::Null => serializer.serialize_unit(),
            Value::Int(v) => serializer.serialize_i64(i64::from(v)),
            Value::Float(v) => serializer.serialize_f64(v),
            Value::String(ref v) => serializer.serialize_str(v),
            Value::Boolean(v) => serializer.serialize_bool(v),
            Value::List(ref v) => v.serialize(serializer),
            Value::Object(ref v) => v.serialize(serializer),
        }
    }
}

#[cfg(test)]
mod tests {
    use serde_json::from_str;
    use serde_json::to_string;
    use ast::InputValue;
    use super::GraphQLError;

    #[test]
    fn int() {
        assert_eq!(from_str::<InputValue>("1235").unwrap(),
                   InputValue::int(1235));
    }

    #[test]
    fn float() {
        assert_eq!(from_str::<InputValue>("2.0").unwrap(),
                   InputValue::float(2.0));
        // large value without a decimal part is also float
        assert_eq!(from_str::<InputValue>("123567890123").unwrap(),
                   InputValue::float(123567890123.0));
    }

    #[test]
    fn errors() {
        assert_eq!(to_string(&GraphQLError::UnknownOperationName).unwrap(),
                   r#"[{"message":"Unknown operation"}]"#);
    }
}
<|endoftext|>


<commit_before>use std::str;
use std::collections::{HashMap, BTreeSet};

use rocksdb::{Writable, IteratorMode, Direction, WriteBatch};
use byteorder::{ByteOrder, BigEndian, WriteBytesExt};

use RocksDBIndexStore;
use errors::{RocksDBReadError, RocksDBWriteError};
use document_index::DocRef;
use key_builder::KeyBuilder;
use search::doc_id_set::DocIdSet;

#[derive(Debug)]
pub enum SegmentMergeError {
    TooManyDocs,
    RocksDBReadError(RocksDBReadError),
    RocksDBWriteError(RocksDBWriteError),
}

impl From<RocksDBReadError> for SegmentMergeError {
    fn from(e: RocksDBReadError) -> SegmentMergeError {
        SegmentMergeError::RocksDBReadError(e)
    }
}

impl From<RocksDBWriteError> for SegmentMergeError {
    fn from(e: RocksDBWriteError) -> SegmentMergeError {
        SegmentMergeError::RocksDBWriteError(e)
    }
}

impl From<SegmentMergeError> for String {
    fn from(e: SegmentMergeError) -> String {
        format!("{:?}", e).to_string()
    }
}

impl RocksDBIndexStore {
    fn merge_segment_data(&self, source_segments: &Vec<u32>, dest_segment: u32, doc_ref_mapping: &HashMap<DocRef, u16>) -> Result<(), SegmentMergeError> {
        // Put source_segments in a BTreeSet as this is much faster for performing contains queries against
        let source_segments_btree = source_segments.iter().collect::<BTreeSet<_>>();

        // Merge the term directories
        // The term directory keys are ordered to be most convenient for retrieving all the segments
        // of for a term/field combination in one go (field/term/segment). So we don't end up pulling
        // in a lot of unwanted data, we firstly iterate the keys, it they one of the source segments
        // looking for then we load them and append them to our new segment.

        /// Converts term directory key strings "d1/2/3" into tuples of 3 i32s (1, 2, 3)
        fn parse_term_directory_key(key: &[u8]) -> (u32, u32, u32) {
            let mut nums_iter = key[1..].split(|b| *b == b'/').map(|s| str::from_utf8(s).unwrap().parse::<u32>().unwrap());
            (nums_iter.next().unwrap(), nums_iter.next().unwrap(), nums_iter.next().unwrap())
        }

        let mut current_td_key: Option<(u32, u32)> = None;
        let mut current_td = Vec::new();

        for k in self.db.keys_iterator(IteratorMode::From(b"d", Direction::Forward)) {
            if k[0] != b'd' {
                // No more term directories to merge
                break;
            }

            let (field, term, segment) = parse_term_directory_key(&k);

            if !source_segments_btree.contains(&segment) {
                continue;
            }

            if current_td_key != Some((field, term)) {
                // Finished current term directory. Write it to the DB and start the next one
                if let Some((field, term)) = current_td_key {
                    let kb = KeyBuilder::segment_dir_list(dest_segment, field, term);
                    if let Err(e) = self.db.put(&kb.key(), &current_td) {
                        return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
                    }
                    current_td.clear();
                }

                current_td_key = Some((field, term));
            }

            // Merge term directory into the new one (and remap the doc ids)
            match self.db.get(&k) {
                Ok(Some(docid_set)) => {
                    for doc_id in DocIdSet::FromRDB(docid_set).iter() {
                        let doc_ref = DocRef::from_segment_ord(segment, doc_id);
                        let new_doc_id = doc_ref_mapping.get(&doc_ref).unwrap();
                        current_td.write_u16::<BigEndian>(*new_doc_id).unwrap();
                    }
                }
                Ok(None) => {},  // FIXME
                Err(e) => return Err(RocksDBReadError::new(k.to_vec(), e).into())
            }
        }

        // All done, write the last term directory
        if let Some((field, term)) = current_td_key {
            let kb = KeyBuilder::segment_dir_list(dest_segment, field, term);
            if let Err(e) = self.db.put(&kb.key(), &current_td) {
                return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
            }
            current_td.clear();
        }

        // Merge the stored values
        // All stored value keys start with the segment id. So we need to:
        // - Iterate all stored value keys that are prefixed by one of the stored segment ids
        // - Remap their doc ids to the one in the new segment
        // - Write the value back with the new segment/doc ids in the key

        /// Converts stored value key strings "v1/2/3/v" into tuples of 3 i32s and a Vec<u8> (1, 2, 3, vec![b'v', b'a', b'l'])
        fn parse_stored_value_key(key: &[u8]) -> (u32, u32, u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let doc_id = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let field_ord = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let value_type = parts_iter.next().unwrap().to_vec();

            (segment, doc_id, field_ord, value_type)
        }

        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stored_values_prefix(*source_segment);
            for (k, v) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b'v' {
                    // No more stored values to move
                    break;
                }

                let (segment, doc_id, field, value_type) = parse_stored_value_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                // Remap doc id
                let doc_ref = DocRef::from_segment_ord(segment, doc_id as u16);
                let new_doc_id = doc_ref_mapping.get(&doc_ref).unwrap();

                // Write value into new segment
                let kb = KeyBuilder::stored_field_value(dest_segment, *new_doc_id, field, &value_type);
                if let Err(e) = self.db.put(&kb.key(), &v) {
                    return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
                }
            }
        }

        // Merge the statistics
        // Like stored values, these start with segment ids. But instead of just rewriting the
        // key, we need to sum up all the statistics across the segments being merged.

        let mut statistics = HashMap::new();

        /// Converts statistic key strings "s1/total_docs" into tuples of 1 i32 and a Vec<u8> (1, ['t', 'o', 't', ...])
        fn parse_statistic_key(key: &[u8]) -> (u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let statistic_name = parts_iter.next().unwrap().to_vec();

            (segment, statistic_name)
        }

        // Fetch and merge statistics
        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stat_prefix(*source_segment);
            for (k, v) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b's' {
                    // No more statistics to merge
                    break;
                }

                let (segment, statistic_name) = parse_statistic_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                let mut stat = statistics.entry(statistic_name).or_insert(0);
                *stat += BigEndian::read_i64(&v);
            }
        }

        // Write merged statistics to new segment
        for (stat_name, stat_value) in statistics {
            let kb = KeyBuilder::segment_stat(dest_segment, &stat_name);
            let mut val_bytes = [0; 8];
            BigEndian::write_i64(&mut val_bytes, stat_value);
            if let Err(e) = self.db.put(&kb.key(), &val_bytes) {
                return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
            }
        }

        // Note: Don't merge the deletion lists
        // Deletion lists can change at any time so we must lock the "document index"
        // before merging them so they can't be altered during merge. we cannot lock
        // this until the commit phase though.

        Ok(())
    }

    fn commit_segment_merge(&self, source_segments: &Vec<u32>, dest_segment: u32, doc_ref_mapping: &HashMap<DocRef, u16>) -> Result<(), SegmentMergeError> {
        let write_batch = WriteBatch::default();

        // Activate new segment
        let kb = KeyBuilder::segment_active(dest_segment);
        if let Err(e) = write_batch.put(&kb.key(), b"") {
            return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
        }

        // Deactivate old segments
        for source_segment in source_segments.iter() {
            // Activate new segment
            let kb = KeyBuilder::segment_active(*source_segment);
            if let Err(e) = write_batch.delete(&kb.key()) {
                return Err(RocksDBWriteError::new_delete(kb.key().to_vec(), e).into());
            }
        }

        // Update document index and commit
        // This will write the write batch
        try!(self.document_index.commit_segment_merge(&self.db, write_batch, source_segments, dest_segment, doc_ref_mapping));

        Ok(())
    }

    pub fn merge_segments(&self, source_segments: Vec<u32>) -> Result<u32, SegmentMergeError> {
        let dest_segment = try!(self.segments.new_segment(&self.db));

        // Generate a mapping between the ids of the documents in the old segments to the new one
        // This packs the id spaces of the old segments together:
        // For example, say we have to merge 3 segments with 100 documents each:
        //  - The first segment's ids will be the same as before
        //  - The second segment's ids will be remapped to 100 - 199
        //  - The third segment's ids will be remapped to 200 - 299

        let mut doc_ref_mapping: HashMap<DocRef, u16> = HashMap::new();
        let mut current_ord: u32 = 0;

        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stat(*source_segment, b"total_docs");
            let total_docs = match self.db.get(&kb.key()) {
                Ok(Some(total_docs_bytes)) => {
                    BigEndian::read_i64(&total_docs_bytes)
                }
                Ok(None) => continue,
                Err(e) => return Err(RocksDBReadError::new(kb.key().to_vec(), e).into())
            };

            for source_ord in 0..total_docs {
                if current_ord >= 65536 {
                    return Err(SegmentMergeError::TooManyDocs);
                }

                let from = DocRef::from_segment_ord(*source_segment, source_ord as u16);
                doc_ref_mapping.insert(from, current_ord as u16);
                current_ord += 1;
            }
        }

        // Merge segment data
        // Most of the heavy lifting happens here. This merges all the immutable parts of
        // the segment (which is everything but the deletion list). It does not activate the
        // segment.
        // This means that nothing bad will happen if it crashes half way through -- the
        // worst that could happen is we're left with a partially-written segment that we
        // have to clean up.
        try!(self.merge_segment_data(&source_segments, dest_segment, &doc_ref_mapping));

        // Commit the merge
        // This activates the new segment and updates the document index. Effectively committing
        // the merge.
        // Throughout this stage we need an exclusive lock to the document index. This is to
        // prevent documents in the source segments being deleted/updated so we don't accidentally
        // undelete them (this will block until the merge is complete so they delete/update from
        // the new segment).
        try!(self.commit_segment_merge(&source_segments, dest_segment, &doc_ref_mapping));

        Ok(dest_segment)
    }
}
<commit_msg>Pull in all term dictionaries instead of cherry-picking what we need<commit_after>use std::str;
use std::collections::{HashMap, BTreeSet};

use rocksdb::{Writable, IteratorMode, Direction, WriteBatch};
use byteorder::{ByteOrder, BigEndian, WriteBytesExt};

use RocksDBIndexStore;
use errors::{RocksDBReadError, RocksDBWriteError};
use document_index::DocRef;
use key_builder::KeyBuilder;
use search::doc_id_set::DocIdSet;

#[derive(Debug)]
pub enum SegmentMergeError {
    TooManyDocs,
    RocksDBReadError(RocksDBReadError),
    RocksDBWriteError(RocksDBWriteError),
}

impl From<RocksDBReadError> for SegmentMergeError {
    fn from(e: RocksDBReadError) -> SegmentMergeError {
        SegmentMergeError::RocksDBReadError(e)
    }
}

impl From<RocksDBWriteError> for SegmentMergeError {
    fn from(e: RocksDBWriteError) -> SegmentMergeError {
        SegmentMergeError::RocksDBWriteError(e)
    }
}

impl From<SegmentMergeError> for String {
    fn from(e: SegmentMergeError) -> String {
        format!("{:?}", e).to_string()
    }
}

impl RocksDBIndexStore {
    fn merge_segment_data(&self, source_segments: &Vec<u32>, dest_segment: u32, doc_ref_mapping: &HashMap<DocRef, u16>) -> Result<(), SegmentMergeError> {
        // Put source_segments in a BTreeSet as this is much faster for performing contains queries against
        let source_segments_btree = source_segments.iter().collect::<BTreeSet<_>>();

        // Merge the term directories
        // The term directory keys are ordered to be most convenient for retrieving all the segments
        // of for a term/field combination in one go (field/term/segment). So we don't end up pulling
        // in a lot of unwanted data, we firstly iterate the keys, it they one of the source segments
        // looking for then we load them and append them to our new segment.

        /// Converts term directory key strings "d1/2/3" into tuples of 3 i32s (1, 2, 3)
        fn parse_term_directory_key(key: &[u8]) -> (u32, u32, u32) {
            let mut nums_iter = key[1..].split(|b| *b == b'/').map(|s| str::from_utf8(s).unwrap().parse::<u32>().unwrap());
            (nums_iter.next().unwrap(), nums_iter.next().unwrap(), nums_iter.next().unwrap())
        }

        let mut current_td_key: Option<(u32, u32)> = None;
        let mut current_td = Vec::new();

        for (k, v) in self.db.iterator(IteratorMode::From(b"d", Direction::Forward)) {
            if k[0] != b'd' {
                // No more term directories to merge
                break;
            }

            let (field, term, segment) = parse_term_directory_key(&k);

            if !source_segments_btree.contains(&segment) {
                continue;
            }

            if current_td_key != Some((field, term)) {
                // Finished current term directory. Write it to the DB and start the next one
                if let Some((field, term)) = current_td_key {
                    let kb = KeyBuilder::segment_dir_list(dest_segment, field, term);
                    if let Err(e) = self.db.put(&kb.key(), &current_td) {
                        return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
                    }
                    current_td.clear();
                }

                current_td_key = Some((field, term));
            }

            // Merge term directory into the new one (and remap the doc ids)
            for doc_id in DocIdSet::Owned(v.to_vec()).iter() {
                let doc_ref = DocRef::from_segment_ord(segment, doc_id);
                let new_doc_id = doc_ref_mapping.get(&doc_ref).unwrap();
                current_td.write_u16::<BigEndian>(*new_doc_id).unwrap();
            }
        }

        // All done, write the last term directory
        if let Some((field, term)) = current_td_key {
            let kb = KeyBuilder::segment_dir_list(dest_segment, field, term);
            if let Err(e) = self.db.put(&kb.key(), &current_td) {
                return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
            }
            current_td.clear();
        }

        // Merge the stored values
        // All stored value keys start with the segment id. So we need to:
        // - Iterate all stored value keys that are prefixed by one of the stored segment ids
        // - Remap their doc ids to the one in the new segment
        // - Write the value back with the new segment/doc ids in the key

        /// Converts stored value key strings "v1/2/3/v" into tuples of 3 i32s and a Vec<u8> (1, 2, 3, vec![b'v', b'a', b'l'])
        fn parse_stored_value_key(key: &[u8]) -> (u32, u32, u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let doc_id = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let field_ord = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let value_type = parts_iter.next().unwrap().to_vec();

            (segment, doc_id, field_ord, value_type)
        }

        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stored_values_prefix(*source_segment);
            for (k, v) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b'v' {
                    // No more stored values to move
                    break;
                }

                let (segment, doc_id, field, value_type) = parse_stored_value_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                // Remap doc id
                let doc_ref = DocRef::from_segment_ord(segment, doc_id as u16);
                let new_doc_id = doc_ref_mapping.get(&doc_ref).unwrap();

                // Write value into new segment
                let kb = KeyBuilder::stored_field_value(dest_segment, *new_doc_id, field, &value_type);
                if let Err(e) = self.db.put(&kb.key(), &v) {
                    return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
                }
            }
        }

        // Merge the statistics
        // Like stored values, these start with segment ids. But instead of just rewriting the
        // key, we need to sum up all the statistics across the segments being merged.

        let mut statistics = HashMap::new();

        /// Converts statistic key strings "s1/total_docs" into tuples of 1 i32 and a Vec<u8> (1, ['t', 'o', 't', ...])
        fn parse_statistic_key(key: &[u8]) -> (u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let statistic_name = parts_iter.next().unwrap().to_vec();

            (segment, statistic_name)
        }

        // Fetch and merge statistics
        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stat_prefix(*source_segment);
            for (k, v) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b's' {
                    // No more statistics to merge
                    break;
                }

                let (segment, statistic_name) = parse_statistic_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                let mut stat = statistics.entry(statistic_name).or_insert(0);
                *stat += BigEndian::read_i64(&v);
            }
        }

        // Write merged statistics to new segment
        for (stat_name, stat_value) in statistics {
            let kb = KeyBuilder::segment_stat(dest_segment, &stat_name);
            let mut val_bytes = [0; 8];
            BigEndian::write_i64(&mut val_bytes, stat_value);
            if let Err(e) = self.db.put(&kb.key(), &val_bytes) {
                return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
            }
        }

        // Note: Don't merge the deletion lists
        // Deletion lists can change at any time so we must lock the "document index"
        // before merging them so they can't be altered during merge. we cannot lock
        // this until the commit phase though.

        Ok(())
    }

    fn commit_segment_merge(&self, source_segments: &Vec<u32>, dest_segment: u32, doc_ref_mapping: &HashMap<DocRef, u16>) -> Result<(), SegmentMergeError> {
        let write_batch = WriteBatch::default();

        // Activate new segment
        let kb = KeyBuilder::segment_active(dest_segment);
        if let Err(e) = write_batch.put(&kb.key(), b"") {
            return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
        }

        // Deactivate old segments
        for source_segment in source_segments.iter() {
            // Activate new segment
            let kb = KeyBuilder::segment_active(*source_segment);
            if let Err(e) = write_batch.delete(&kb.key()) {
                return Err(RocksDBWriteError::new_delete(kb.key().to_vec(), e).into());
            }
        }

        // Update document index and commit
        // This will write the write batch
        try!(self.document_index.commit_segment_merge(&self.db, write_batch, source_segments, dest_segment, doc_ref_mapping));

        Ok(())
    }

    pub fn merge_segments(&self, source_segments: Vec<u32>) -> Result<u32, SegmentMergeError> {
        let dest_segment = try!(self.segments.new_segment(&self.db));

        // Generate a mapping between the ids of the documents in the old segments to the new one
        // This packs the id spaces of the old segments together:
        // For example, say we have to merge 3 segments with 100 documents each:
        //  - The first segment's ids will be the same as before
        //  - The second segment's ids will be remapped to 100 - 199
        //  - The third segment's ids will be remapped to 200 - 299

        let mut doc_ref_mapping: HashMap<DocRef, u16> = HashMap::new();
        let mut current_ord: u32 = 0;

        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stat(*source_segment, b"total_docs");
            let total_docs = match self.db.get(&kb.key()) {
                Ok(Some(total_docs_bytes)) => {
                    BigEndian::read_i64(&total_docs_bytes)
                }
                Ok(None) => continue,
                Err(e) => return Err(RocksDBReadError::new(kb.key().to_vec(), e).into())
            };

            for source_ord in 0..total_docs {
                if current_ord >= 65536 {
                    return Err(SegmentMergeError::TooManyDocs);
                }

                let from = DocRef::from_segment_ord(*source_segment, source_ord as u16);
                doc_ref_mapping.insert(from, current_ord as u16);
                current_ord += 1;
            }
        }

        // Merge segment data
        // Most of the heavy lifting happens here. This merges all the immutable parts of
        // the segment (which is everything but the deletion list). It does not activate the
        // segment.
        // This means that nothing bad will happen if it crashes half way through -- the
        // worst that could happen is we're left with a partially-written segment that we
        // have to clean up.
        try!(self.merge_segment_data(&source_segments, dest_segment, &doc_ref_mapping));

        // Commit the merge
        // This activates the new segment and updates the document index. Effectively committing
        // the merge.
        // Throughout this stage we need an exclusive lock to the document index. This is to
        // prevent documents in the source segments being deleted/updated so we don't accidentally
        // undelete them (this will block until the merge is complete so they delete/update from
        // the new segment).
        try!(self.commit_segment_merge(&source_segments, dest_segment, &doc_ref_mapping));

        Ok(dest_segment)
    }
}
<|endoftext|>


<commit_before>use std::io;

use ::{BerSerialize, BerDeserialize, Asn1Info};
use tag;
use err;

// impl Asn1Info for Vec<u8> {
//   fn asn1_tag() -> tag::Tag {
//     tag::Tag {
//       class: tag::Class::Universal,
//       tagnum: 4u8.into(),
//       constructed: false,
//     }
//   }
//
//   fn asn1_type() -> tag::Type {
//     tag::Type::from("OCTET STRING")
//   }
// }
//
// impl BerSerialize for Vec<u8> {
//   fn serialize_value<E: ::BerEncRules, W: io::Write>
//       (&self, _: E, writer: &mut W) -> Result<(), err::EncodeError> {
//     try!(writer.write_all(self));
//     Ok(())
//   }
// }
//
// impl BerDeserialize for Vec<u8> {
//   fn deserialize_value<E: ::BerEncRules, I: Iterator<Item=io::Result<u8>>>
//       (_: E, reader: &mut I, len: tag::Len) -> Result<Self, err::DecodeError> {
//     let len_num = try!(len.as_num().ok_or(err::DecodeError::PrimIndef));
//     let bytes: Result<Vec<u8>, _> = reader.take(len_num as usize).collect();
//     Ok(try!(bytes))
//   }
// }
//
<commit_msg>Comment out prim.rs definition for now<commit_after>// use std::io;

// use ::{BerSerialize, BerDeserialize, Asn1Info};
// use tag;
// use err;

// impl Asn1Info for Vec<u8> {
//   fn asn1_tag() -> tag::Tag {
//     tag::Tag {
//       class: tag::Class::Universal,
//       tagnum: 4u8.into(),
//       constructed: false,
//     }
//   }
//
//   fn asn1_type() -> tag::Type {
//     tag::Type::from("OCTET STRING")
//   }
// }
//
// impl BerSerialize for Vec<u8> {
//   fn serialize_value<E: ::BerEncRules, W: io::Write>
//       (&self, _: E, writer: &mut W) -> Result<(), err::EncodeError> {
//     try!(writer.write_all(self));
//     Ok(())
//   }
// }
//
// impl BerDeserialize for Vec<u8> {
//   fn deserialize_value<E: ::BerEncRules, I: Iterator<Item=io::Result<u8>>>
//       (_: E, reader: &mut I, len: tag::Len) -> Result<Self, err::DecodeError> {
//     let len_num = try!(len.as_num().ok_or(err::DecodeError::PrimIndef));
//     let bytes: Result<Vec<u8>, _> = reader.take(len_num as usize).collect();
//     Ok(try!(bytes))
//   }
// }
//
<|endoftext|>