<commit_before>use std;
use rustc;

import std::fs;
import std::getopts;
import std::getopts::optopt;
import std::getopts::opt_present;
import std::getopts::opt_str;
import std::io;
import std::io::stdout;
import std::vec;
import std::ivec;
import std::str;
import std::uint;
import std::option;

import rustc::syntax::ast;
import rustc::syntax::fold;
import rustc::syntax::walk;
import rustc::syntax::codemap;
import rustc::syntax::parse::parser;
import rustc::syntax::print::pprust;

/*
// Imports for "the rest of driver::compile_input"
import driver = rustc::driver::rustc; // see https://github.com/graydon/rust/issues/624
import rustc::back::link;
import rustc::driver::rustc::time;
import rustc::driver::session;

import rustc::metadata::creader;
import rustc::metadata::cstore;
import rustc::syntax::parse::parser;
import rustc::syntax::parse::token;
import rustc::front;
import rustc::front::attr;
import rustc::middle;
import rustc::middle::trans;
import rustc::middle::resolve;
import rustc::middle::ty;
import rustc::middle::typeck;
import rustc::middle::tstate::ck;
import rustc::syntax::print::pp;
import rustc::util::ppaux;
import rustc::lib::llvm;
*/

fn read_whole_file(&str filename) -> str {
    str::unsafe_from_bytes(io::file_reader(filename).read_whole_stream())
}

fn write_file(&str filename, &str content) {
    io::file_writer(filename, [io::create]).write_str(content);
}

fn file_contains(&str filename, &str needle) -> bool {
    auto contents = read_whole_file(filename);
    ret str::find(contents, needle) != -1;
}

fn contains(&str haystack, &str needle) -> bool { str::find(haystack, needle) != -1 }

fn find_rust_files(&mutable str[] files, str path) {
    if (str::ends_with(path, ".rs")) {
        if (file_contains(path, "xfail-stage1")) {
            //log_err "Skipping " + path + " because it is marked as xfail-stage1";
        } else {
            files += ~[path];
        }
    } else if (fs::file_is_dir(path) && str::find(path, "compile-fail") == -1) {
        for (str p in fs::list_dir(path)) {
            find_rust_files(files, p);
        }
    }
}

fn safe_to_steal(ast::expr_ e) -> bool {
    alt (e) {
        // pretty-printer precedence issues -- https://github.com/graydon/rust/issues/670
        case (ast::expr_unary(_, _)) { false }
        case (ast::expr_lit(?lit)) {
            alt(lit.node) {
                case(ast::lit_str(_, _)) { true }
                case(ast::lit_char(_)) { true }
                case(ast::lit_int(_)) { false }
                case(ast::lit_uint(_)) { false }
                case(ast::lit_mach_int(_, _)) { false }
                case(ast::lit_float(_)) { false }
                case(ast::lit_mach_float(_, _)) { false }
                case(ast::lit_nil) { true }
                case(ast::lit_bool(_)) { true }
            }
        }
        case (ast::expr_cast(_, _)) { false }
        case (ast::expr_send(_, _)) { false }
        case (ast::expr_recv(_, _)) { false }
        case (ast::expr_assert(_)) { false }
        case (ast::expr_binary(_, _, _)) { false }
        case (ast::expr_assign(_, _)) { false }
        case (ast::expr_assign_op(_, _, _)) { false }

        // https://github.com/graydon/rust/issues/676
        case (ast::expr_ret(option::none)) { false }
        case (ast::expr_put(option::none)) { false }

        case (_) { true }
    }
}

fn steal_exprs(&ast::crate crate) -> ast::expr[] {
    let @mutable ast::expr[] exprs = @mutable ~[];
    // "Stash" is not type-parameterized because of the need for safe_to_steal
    fn stash_expr(@mutable ast::expr[] es, &@ast::expr e) { if (safe_to_steal(e.node)) { *es += ~[*e]; } else { /* now my indices are wrong :( */ } }
    auto v = rec(visit_expr_pre = bind stash_expr(exprs, _) with walk::default_visitor());
    walk::walk_crate(v, crate);
    *exprs
}

// https://github.com/graydon/rust/issues/652
fn safe_to_replace(ast::expr_ e) -> bool {
    alt (e) {
        case (ast::expr_if(_, _, _)) { false }
        case (ast::expr_block(_)) { false }
        case (_) { true }
    }
}

// Replace the |i|th expr (in fold order) of |crate| with |newexpr|.
fn replace_expr_in_crate(&ast::crate crate, uint i, ast::expr_ newexpr) -> ast::crate {
    let @mutable uint j = @mutable 0u;
    fn fold_expr_rep(@mutable uint j_, uint i_, &ast::expr_ newexpr_, &ast::expr_ original, fold::ast_fold fld) -> ast::expr_ {
      *j_ += 1u;
      if (i_ + 1u == *j_ && safe_to_replace(original)) {
        newexpr_
      } else {
        fold::noop_fold_expr(original, fld)
      }
    }
    auto afp = rec(fold_expr = bind fold_expr_rep(j, i, newexpr, _, _) with *fold::default_ast_fold());
    auto af = fold::make_fold(afp);
    let @ast::crate crate2 = @af.fold_crate(crate);
    fold::dummy_out(af); // work around a leak (https://github.com/graydon/rust/issues/651)
    *crate2
}

iter under(uint n) -> uint { let uint i = 0u; while (i < n) { put i; i += 1u; } }

fn devnull() -> io::writer { std::io::string_writer().get_writer() }

fn as_str(fn (io::writer) f) -> str { auto w = std::io::string_writer(); f(w.get_writer()); w.get_str() }

fn pp_variants(&ast::crate crate, &codemap::codemap cmap, &str filename) {
    auto exprs = steal_exprs(crate);
    auto exprsL = ivec::len(exprs);
    if (exprsL < 100u) {
        for each (uint i in under(uint::min(exprsL, 20u))) {
            log_err "Replacing... " + pprust::expr_to_str(@exprs.(i));
            for each (uint j in under(uint::min(exprsL, 5u))) {
                log_err "With... " + pprust::expr_to_str(@exprs.(j));
                auto crate2 = @replace_expr_in_crate(crate, i, exprs.(j).node);
                check_roundtrip(crate2, cmap, filename + ".4.rs");
            }
        }
    }
}

fn check_roundtrip(@ast::crate crate2, &codemap::codemap cmap, &str filename) {
    auto str3 = as_str(bind pprust::print_crate(cmap, crate2, filename, _, pprust::no_ann()));
    auto cm4 = codemap::new_codemap();
    if (true
      && !contains(str3, "][]") // https://github.com/graydon/rust/issues/669
      && !contains(str3, "][mutable]") // https://github.com/graydon/rust/issues/669
      && !contains(str3, "][mutable ]") // https://github.com/graydon/rust/issues/669
      && !contains(str3, "self") // crazy rules enforced by parser rather than typechecker?
      && !contains(str3, "spawn") // more precedence issues
      && !contains(str3, "bind") // more precedence issues?
       ) {
        auto crate4 = parser::parse_crate_from_source_str(filename, str3, ~[], cm4);
        // should compare crates at this point, but it's easier to compare strings
        auto str5 = as_str(bind pprust::print_crate(cmap, crate4, filename, _, pprust::no_ann()));
        if (!str::is_ascii(str3)) {
          log_err "Non-ASCII in " + filename; // why does non-ASCII work correctly with "rustc --pretty normal" but not here???
        } else if (str3 != str5) {
            write_file("round-trip-a.rs", str3);
            write_file("round-trip-b.rs", str5);
            std::run::run_program("kdiff3", ["round-trip-a.rs", "round-trip-b.rs"]);
            fail "Mismatch";
        }
   }
}

fn main(vec[str] args) {
    if (vec::len(args) != 2u) {
        log_err #fmt("usage: %s <testdir>", args.(0));
        ret;
    }
    auto files = ~[];
    auto root = args.(1);
    find_rust_files(files, root); // not using time here because that currently screws with passing-a-mutable-array
    log_err #fmt("%u files", ivec::len(files));

    for (str file in files) {
        log_err "=== " + file + " ===";
        auto cm = codemap::new_codemap();
        auto src = read_whole_file(file);
        auto crate = parser::parse_crate_from_source_str(file, src, ~[], cm);
        if (!contains(src, "#macro") // https://github.com/graydon/rust/issues/671
         && !str::ends_with(file, "block-expr-precedence.rs") // https://github.com/graydon/rust/issues/674
         && !str::ends_with(file, "syntax-extension-fmt.rs") // an issue where -2147483648 gains an extra negative sign each time through, which i can't reproduce using "rustc --pretty normal"???
) {
            check_roundtrip(crate, cm, file);
        }
        //pprust::print_crate(cm, crate, file, devnull(), pprust::no_ann());
        // Currently hits https://github.com/graydon/rust/issues/675
        //pp_variants(*crate, cm, file);
    }
}

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C $RBUILD 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
<commit_msg>Fix fuzzer to compile with ivec-ized pp module.<commit_after>use std;
use rustc;

import std::fs;
import std::getopts;
import std::getopts::optopt;
import std::getopts::opt_present;
import std::getopts::opt_str;
import std::ioivec;
import std::ioivec::stdout;
import std::vec;
import std::ivec;
import std::str;
import std::uint;
import std::option;

import rustc::syntax::ast;
import rustc::syntax::fold;
import rustc::syntax::walk;
import rustc::syntax::codemap;
import rustc::syntax::parse::parser;
import rustc::syntax::print::pprust;

/*
// Imports for "the rest of driver::compile_input"
import driver = rustc::driver::rustc; // see https://github.com/graydon/rust/issues/624
import rustc::back::link;
import rustc::driver::rustc::time;
import rustc::driver::session;

import rustc::metadata::creader;
import rustc::metadata::cstore;
import rustc::syntax::parse::parser;
import rustc::syntax::parse::token;
import rustc::front;
import rustc::front::attr;
import rustc::middle;
import rustc::middle::trans;
import rustc::middle::resolve;
import rustc::middle::ty;
import rustc::middle::typeck;
import rustc::middle::tstate::ck;
import rustc::syntax::print::pp;
import rustc::util::ppaux;
import rustc::lib::llvm;
*/

fn read_whole_file(&str filename) -> str {
    str::unsafe_from_bytes_ivec(ioivec::file_reader(filename).read_whole_stream())
}

fn write_file(&str filename, &str content) {
    ioivec::file_writer(filename, ~[ioivec::create]).write_str(content);
}

fn file_contains(&str filename, &str needle) -> bool {
    auto contents = read_whole_file(filename);
    ret str::find(contents, needle) != -1;
}

fn contains(&str haystack, &str needle) -> bool { str::find(haystack, needle) != -1 }

fn find_rust_files(&mutable str[] files, str path) {
    if (str::ends_with(path, ".rs")) {
        if (file_contains(path, "xfail-stage1")) {
            //log_err "Skipping " + path + " because it is marked as xfail-stage1";
        } else {
            files += ~[path];
        }
    } else if (fs::file_is_dir(path) && str::find(path, "compile-fail") == -1) {
        for (str p in fs::list_dir(path)) {
            find_rust_files(files, p);
        }
    }
}

fn safe_to_steal(ast::expr_ e) -> bool {
    alt (e) {
        // pretty-printer precedence issues -- https://github.com/graydon/rust/issues/670
        case (ast::expr_unary(_, _)) { false }
        case (ast::expr_lit(?lit)) {
            alt(lit.node) {
                case(ast::lit_str(_, _)) { true }
                case(ast::lit_char(_)) { true }
                case(ast::lit_int(_)) { false }
                case(ast::lit_uint(_)) { false }
                case(ast::lit_mach_int(_, _)) { false }
                case(ast::lit_float(_)) { false }
                case(ast::lit_mach_float(_, _)) { false }
                case(ast::lit_nil) { true }
                case(ast::lit_bool(_)) { true }
            }
        }
        case (ast::expr_cast(_, _)) { false }
        case (ast::expr_send(_, _)) { false }
        case (ast::expr_recv(_, _)) { false }
        case (ast::expr_assert(_)) { false }
        case (ast::expr_binary(_, _, _)) { false }
        case (ast::expr_assign(_, _)) { false }
        case (ast::expr_assign_op(_, _, _)) { false }

        // https://github.com/graydon/rust/issues/676
        case (ast::expr_ret(option::none)) { false }
        case (ast::expr_put(option::none)) { false }

        case (_) { true }
    }
}

fn steal_exprs(&ast::crate crate) -> ast::expr[] {
    let @mutable ast::expr[] exprs = @mutable ~[];
    // "Stash" is not type-parameterized because of the need for safe_to_steal
    fn stash_expr(@mutable ast::expr[] es, &@ast::expr e) { if (safe_to_steal(e.node)) { *es += ~[*e]; } else { /* now my indices are wrong :( */ } }
    auto v = rec(visit_expr_pre = bind stash_expr(exprs, _) with walk::default_visitor());
    walk::walk_crate(v, crate);
    *exprs
}

// https://github.com/graydon/rust/issues/652
fn safe_to_replace(ast::expr_ e) -> bool {
    alt (e) {
        case (ast::expr_if(_, _, _)) { false }
        case (ast::expr_block(_)) { false }
        case (_) { true }
    }
}

// Replace the |i|th expr (in fold order) of |crate| with |newexpr|.
fn replace_expr_in_crate(&ast::crate crate, uint i, ast::expr_ newexpr) -> ast::crate {
    let @mutable uint j = @mutable 0u;
    fn fold_expr_rep(@mutable uint j_, uint i_, &ast::expr_ newexpr_, &ast::expr_ original, fold::ast_fold fld) -> ast::expr_ {
      *j_ += 1u;
      if (i_ + 1u == *j_ && safe_to_replace(original)) {
        newexpr_
      } else {
        fold::noop_fold_expr(original, fld)
      }
    }
    auto afp = rec(fold_expr = bind fold_expr_rep(j, i, newexpr, _, _) with *fold::default_ast_fold());
    auto af = fold::make_fold(afp);
    let @ast::crate crate2 = @af.fold_crate(crate);
    fold::dummy_out(af); // work around a leak (https://github.com/graydon/rust/issues/651)
    *crate2
}

iter under(uint n) -> uint { let uint i = 0u; while (i < n) { put i; i += 1u; } }

fn devnull() -> ioivec::writer { std::ioivec::string_writer().get_writer() }

fn as_str(fn (ioivec::writer) f) -> str { auto w = std::ioivec::string_writer(); f(w.get_writer()); w.get_str() }

fn pp_variants(&ast::crate crate, &codemap::codemap cmap, &str filename) {
    auto exprs = steal_exprs(crate);
    auto exprsL = ivec::len(exprs);
    if (exprsL < 100u) {
        for each (uint i in under(uint::min(exprsL, 20u))) {
            log_err "Replacing... " + pprust::expr_to_str(@exprs.(i));
            for each (uint j in under(uint::min(exprsL, 5u))) {
                log_err "With... " + pprust::expr_to_str(@exprs.(j));
                auto crate2 = @replace_expr_in_crate(crate, i, exprs.(j).node);
                check_roundtrip(crate2, cmap, filename + ".4.rs");
            }
        }
    }
}

fn check_roundtrip(@ast::crate crate2, &codemap::codemap cmap, &str filename) {
    auto str3 = as_str(bind pprust::print_crate(cmap, crate2, filename, _, pprust::no_ann()));
    if (true
      && !contains(str3, "][]") // https://github.com/graydon/rust/issues/669
      && !contains(str3, "][mutable]") // https://github.com/graydon/rust/issues/669
      && !contains(str3, "][mutable ]") // https://github.com/graydon/rust/issues/669
      && !contains(str3, "self") // crazy rules enforced by parser rather than typechecker?
      && !contains(str3, "spawn") // more precedence issues
      && !contains(str3, "bind") // more precedence issues?
       ) {
        auto cm4 = codemap::new_codemap();
        auto crate4 = parser::parse_crate_from_source_str(filename, str3, ~[], cm4);
        // should compare crates at this point, but it's easier to compare strings
        auto str5 = as_str(bind pprust::print_crate(cm4, crate4, filename, _, pprust::no_ann()));
        if (!str::is_ascii(str3)) {
          log_err "Non-ASCII in " + filename; // why does non-ASCII work correctly with "rustc --pretty normal" but not here???
        } else if (str3 != str5) {
            write_file("round-trip-a.rs", str3);
            write_file("round-trip-b.rs", str5);
            std::run::run_program("kdiff3", ["round-trip-a.rs", "round-trip-b.rs"]);
            fail "Mismatch";
        }
   }
}

fn main(vec[str] args) {
    if (vec::len(args) != 2u) {
        log_err #fmt("usage: %s <testdir>", args.(0));
        ret;
    }
    auto files = ~[];
    auto root = args.(1);
    find_rust_files(files, root); // not using time here because that currently screws with passing-a-mutable-array
    log_err #fmt("%u files", ivec::len(files));

    for (str file in files) {
        log_err "=== " + file + " ===";
        auto cm = codemap::new_codemap();
        auto src = read_whole_file(file);
        auto crate = parser::parse_crate_from_source_str(file, src, ~[], cm);
        if (!contains(src, "#macro") // https://github.com/graydon/rust/issues/671
         && !str::ends_with(file, "block-expr-precedence.rs") // https://github.com/graydon/rust/issues/674
         && !str::ends_with(file, "syntax-extension-fmt.rs") // an issue where -2147483648 gains an extra negative sign each time through, which i can't reproduce using "rustc --pretty normal"???
) {
            check_roundtrip(crate, cm, file);
        }
        //pprust::print_crate(cm, crate, file, devnull(), pprust::no_ann());
        // Currently hits https://github.com/graydon/rust/issues/675
        //pp_variants(*crate, cm, file);
    }
}

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C $RBUILD 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
<|endoftext|>


<commit_before>use std::{
    collections::HashMap,
    path::{Path, PathBuf},
};

use hir::Crate;
use log::trace;
use petgraph::graph::{EdgeIndex, NodeIndex};
use ra_ap_hir::{self as hir, ModuleSource};
use ra_ap_ide_db::RootDatabase;
use ra_ap_vfs::Vfs;

use crate::graph::{orphans::add_orphan_nodes_to, Edge, Graph, Node};

#[derive(Clone, PartialEq, Debug)]
pub struct Options {
    pub focus_on: Option<String>,
    pub max_depth: Option<usize>,
    pub with_types: bool,
    pub with_orphans: bool,
    pub with_uses: bool,
    pub with_externs: bool,
}

#[derive(Debug)]
pub struct Builder<'a> {
    options: Options,
    db: &'a RootDatabase,
    vfs: &'a Vfs,
    graph: Graph,
    nodes: HashMap<String, NodeIndex>,
}

impl<'a> Builder<'a> {
    pub fn new(options: Options, db: &'a RootDatabase, vfs: &'a Vfs) -> Self {
        let graph = Graph::default();
        let nodes = HashMap::default();
        Self {
            options,
            db,
            vfs,
            graph,
            nodes,
        }
    }

    pub fn build(mut self, krate: Crate) -> anyhow::Result<Graph> {
        let root_module = krate.root_module(self.db);
        let root_module_def = hir::ModuleDef::Module(root_module);

        trace!("Scanning project ...");

        self.add_module_def(root_module_def, true, krate)?;

        Ok(self.graph)
    }

    fn add_module_def(
        &mut self,
        module_def: hir::ModuleDef,
        recursive: bool,
        krate: Crate,
    ) -> anyhow::Result<Option<NodeIndex>> {
        let is_module = matches!(module_def, hir::ModuleDef::Module(_));

        if !self.options.with_types && !is_module {
            return Ok(None);
        }

        let module_def_krate = {
            match module_def {
                hir::ModuleDef::Module(module) => Some(module),
                module_def @ _ => module_def.module(self.db),
            }
            .map(|module| module.krate())
        };

        let is_external = module_def_krate != Some(krate);

        let module_def = if is_external {
            if !self.options.with_externs {
                return Ok(None);
            }

            if let Some(module_def_krate) = module_def_krate {
                hir::ModuleDef::Module(module_def_krate.root_module(self.db))
            } else {
                module_def
            }
        } else {
            module_def
        };

        let path = self.module_path(module_def);

        // Check if we already have any node registered for this path,
        // otherwise add a node:

        let node_idx = match self.nodes.get(&path) {
            Some(node_idx) => *node_idx,
            None => self.add_module_node(module_def, module_def_krate, &path),
        };

        if !recursive {
            return Ok(Some(node_idx));
        }

        if let hir::ModuleDef::Module(module) = module_def {
            trace!("Scanning module {:?}", path);

            for (_name, scope_def) in module.scope(self.db, None) {
                let scope_module_def = if let hir::ScopeDef::ModuleDef(scope_module_def) = scope_def
                {
                    scope_module_def
                } else {
                    continue;
                };

                let is_module = matches!(scope_module_def, hir::ModuleDef::Module(_));

                if !self.options.with_types && !is_module {
                    continue;
                }

                let is_local = Some(module) == scope_module_def.module(self.db);

                if !self.options.with_uses && !is_local {
                    continue;
                }

                let scope_node_idx = self.add_module_def(scope_module_def, is_local, krate)?;

                if let Some(scope_node_idx) = scope_node_idx {
                    if self.graph.contains_edge(node_idx, scope_node_idx) {
                        // Avoid adding redundant edges:
                        continue;
                    }

                    let edge = if is_local { Edge::HasA } else { Edge::UsesA };
                    self.add_edge(node_idx, scope_node_idx, edge);
                }
            }

            if self.options.with_orphans {
                add_orphan_nodes_to(&mut self.graph, node_idx);
            }
        }

        Ok(Some(node_idx))
    }

    fn module_file(&self, module_source: hir::InFile<hir::ModuleSource>) -> Option<PathBuf> {
        let is_file_module: bool = match &module_source.value {
            ModuleSource::SourceFile(_) => true,
            ModuleSource::Module(_) => false,
        };

        if !is_file_module {
            return None;
        }

        let file_id = module_source.file_id.original_file(self.db);
        let vfs_path = self.vfs.file_path(file_id);
        let abs_path = vfs_path.as_path().expect("Could not convert to path");

        let path: &Path = &*abs_path;

        let file_extension = path.extension().and_then(|ext| ext.to_str());

        if file_extension != Some("rs") {
            return None;
        }

        Some(path.to_owned())
    }

    fn crate_name(&self, krate: hir::Crate) -> String {
        // Obtain the crate's declaration name:
        let display_name = &krate.display_name(self.db).unwrap();

        // Since a crate's name may contain `-` we canonicalize it by replacing with `_`:
        display_name.replace("-", "_")
    }

    fn module(&self, module_def: hir::ModuleDef) -> hir::Module {
        match module_def {
            hir::ModuleDef::Module(module) => module,
            _ => module_def.module(self.db).unwrap(),
        }
    }

    fn module_path(&self, module_def: hir::ModuleDef) -> String {
        let module = self.module(module_def);

        // Obtain the module's canonicalized name:
        let crate_name = self.crate_name(module.krate());

        let relative_canonical_path = module_def.canonical_path(self.db);

        match relative_canonical_path {
            Some(canonical_path) => format!("{}::{}", crate_name, canonical_path),
            None => crate_name,
        }
    }

    fn add_module_node(
        &mut self,
        module_def: hir::ModuleDef,
        krate: Option<hir::Crate>,
        module_path: &str,
    ) -> NodeIndex {
        trace!("Adding module node: {:?}", module_path);

        let module_path = self.module_path(module_def);

        let node = self.make_node(module_def, krate, &module_path);

        let node_idx = self.graph.add_node(node);
        self.nodes.insert(module_path, node_idx);

        node_idx
    }

    fn add_edge(&mut self, source_idx: NodeIndex, target_idx: NodeIndex, edge: Edge) -> EdgeIndex {
        trace!("Adding edge: {:?} -> {:?}", source_idx, target_idx);

        self.graph.add_edge(source_idx, target_idx, edge)
    }

    fn make_node(&self, hir: hir::ModuleDef, krate: Option<hir::Crate>, module_path: &str) -> Node {
        let path = module_path.to_owned();
        let file_path = {
            match hir {
                hir::ModuleDef::Module(module) => Some(module),
                _ => None,
            }
            .and_then(|module| {
                self.module_file(module.definition_source(self.db))
                    .map(Into::into)
            })
        };
        let hir = Some(hir);

        Node {
            path,
            file_path,
            hir,
            krate,
        }
    }
}
<commit_msg>Removed unused method argument from `GraphBuilder.add_module_node`<commit_after>use std::{
    collections::HashMap,
    path::{Path, PathBuf},
};

use hir::Crate;
use log::trace;
use petgraph::graph::{EdgeIndex, NodeIndex};
use ra_ap_hir::{self as hir, ModuleSource};
use ra_ap_ide_db::RootDatabase;
use ra_ap_vfs::Vfs;

use crate::graph::{orphans::add_orphan_nodes_to, Edge, Graph, Node};

#[derive(Clone, PartialEq, Debug)]
pub struct Options {
    pub focus_on: Option<String>,
    pub max_depth: Option<usize>,
    pub with_types: bool,
    pub with_orphans: bool,
    pub with_uses: bool,
    pub with_externs: bool,
}

#[derive(Debug)]
pub struct Builder<'a> {
    options: Options,
    db: &'a RootDatabase,
    vfs: &'a Vfs,
    graph: Graph,
    nodes: HashMap<String, NodeIndex>,
}

impl<'a> Builder<'a> {
    pub fn new(options: Options, db: &'a RootDatabase, vfs: &'a Vfs) -> Self {
        let graph = Graph::default();
        let nodes = HashMap::default();
        Self {
            options,
            db,
            vfs,
            graph,
            nodes,
        }
    }

    pub fn build(mut self, krate: Crate) -> anyhow::Result<Graph> {
        let root_module = krate.root_module(self.db);
        let root_module_def = hir::ModuleDef::Module(root_module);

        trace!("Scanning project ...");

        self.add_module_def(root_module_def, true, krate)?;

        Ok(self.graph)
    }

    fn add_module_def(
        &mut self,
        module_def: hir::ModuleDef,
        recursive: bool,
        krate: Crate,
    ) -> anyhow::Result<Option<NodeIndex>> {
        let is_module = matches!(module_def, hir::ModuleDef::Module(_));

        if !self.options.with_types && !is_module {
            return Ok(None);
        }

        let module_def_krate = {
            match module_def {
                hir::ModuleDef::Module(module) => Some(module),
                module_def @ _ => module_def.module(self.db),
            }
            .map(|module| module.krate())
        };

        let is_external = module_def_krate != Some(krate);

        let module_def = if is_external {
            if !self.options.with_externs {
                return Ok(None);
            }

            if let Some(module_def_krate) = module_def_krate {
                hir::ModuleDef::Module(module_def_krate.root_module(self.db))
            } else {
                module_def
            }
        } else {
            module_def
        };

        let path = self.module_path(module_def);

        // Check if we already have any node registered for this path,
        // otherwise add a node:

        let node_idx = match self.nodes.get(&path) {
            Some(node_idx) => *node_idx,
            None => self.add_module_node(module_def, module_def_krate, &path),
        };

        if !recursive {
            return Ok(Some(node_idx));
        }

        if let hir::ModuleDef::Module(module) = module_def {
            trace!("Scanning module {:?}", path);

            for (_name, scope_def) in module.scope(self.db, None) {
                let scope_module_def = if let hir::ScopeDef::ModuleDef(scope_module_def) = scope_def
                {
                    scope_module_def
                } else {
                    continue;
                };

                let is_module = matches!(scope_module_def, hir::ModuleDef::Module(_));

                if !self.options.with_types && !is_module {
                    continue;
                }

                let is_local = Some(module) == scope_module_def.module(self.db);

                if !self.options.with_uses && !is_local {
                    continue;
                }

                let scope_node_idx = self.add_module_def(scope_module_def, is_local, krate)?;

                if let Some(scope_node_idx) = scope_node_idx {
                    if self.graph.contains_edge(node_idx, scope_node_idx) {
                        // Avoid adding redundant edges:
                        continue;
                    }

                    let edge = if is_local { Edge::HasA } else { Edge::UsesA };
                    self.add_edge(node_idx, scope_node_idx, edge);
                }
            }

            if self.options.with_orphans {
                add_orphan_nodes_to(&mut self.graph, node_idx);
            }
        }

        Ok(Some(node_idx))
    }

    fn module_file(&self, module_source: hir::InFile<hir::ModuleSource>) -> Option<PathBuf> {
        let is_file_module: bool = match &module_source.value {
            ModuleSource::SourceFile(_) => true,
            ModuleSource::Module(_) => false,
        };

        if !is_file_module {
            return None;
        }

        let file_id = module_source.file_id.original_file(self.db);
        let vfs_path = self.vfs.file_path(file_id);
        let abs_path = vfs_path.as_path().expect("Could not convert to path");

        let path: &Path = &*abs_path;

        let file_extension = path.extension().and_then(|ext| ext.to_str());

        if file_extension != Some("rs") {
            return None;
        }

        Some(path.to_owned())
    }

    fn crate_name(&self, krate: hir::Crate) -> String {
        // Obtain the crate's declaration name:
        let display_name = &krate.display_name(self.db).unwrap();

        // Since a crate's name may contain `-` we canonicalize it by replacing with `_`:
        display_name.replace("-", "_")
    }

    fn module(&self, module_def: hir::ModuleDef) -> hir::Module {
        match module_def {
            hir::ModuleDef::Module(module) => module,
            _ => module_def.module(self.db).unwrap(),
        }
    }

    fn module_path(&self, module_def: hir::ModuleDef) -> String {
        let module = self.module(module_def);

        // Obtain the module's canonicalized name:
        let crate_name = self.crate_name(module.krate());

        let relative_canonical_path = module_def.canonical_path(self.db);

        match relative_canonical_path {
            Some(canonical_path) => format!("{}::{}", crate_name, canonical_path),
            None => crate_name,
        }
    }

    fn add_module_node(
        &mut self,
        module_def: hir::ModuleDef,
        krate: Option<hir::Crate>,
    ) -> NodeIndex {
        let module_path = self.module_path(module_def);

        trace!("Adding module node: {:?}", module_path);

        let node = self.make_node(module_def, krate, &module_path);

        let node_idx = self.graph.add_node(node);
        self.nodes.insert(module_path, node_idx);

        node_idx
    }

    fn add_edge(&mut self, source_idx: NodeIndex, target_idx: NodeIndex, edge: Edge) -> EdgeIndex {
        trace!("Adding edge: {:?} -> {:?}", source_idx, target_idx);

        self.graph.add_edge(source_idx, target_idx, edge)
    }

    fn make_node(&self, hir: hir::ModuleDef, krate: Option<hir::Crate>, module_path: &str) -> Node {
        let path = module_path.to_owned();
        let file_path = {
            match hir {
                hir::ModuleDef::Module(module) => Some(module),
                _ => None,
            }
            .and_then(|module| {
                self.module_file(module.definition_source(self.db))
                    .map(Into::into)
            })
        };
        let hir = Some(hir);

        Node {
            path,
            file_path,
            hir,
            krate,
        }
    }
}
<|endoftext|>


<commit_before>use super::btree::{BPlusTree, RTCursor as BPlusTreeCursor};
use super::sstable::*;
use super::*;
use client::AsyncClient;
use itertools::Itertools;
use parking_lot::RwLock;
use ram::segs::MAX_SEGMENT_SIZE;
use std::collections::BTreeSet;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use std::{mem, ptr};

const LEVEL_M_MAX_ELEMENTS_COUNT: usize = LEVEL_M * LEVEL_M * LEVEL_M;
const LEVEL_ELEMENTS_MULTIPLIER: usize = 100;
const LEVEL_DIFF_MULTIPLIER: usize = 10;
const LEVEL_M: usize = super::btree::NUM_KEYS;
const LEVEL_1: usize = LEVEL_M * LEVEL_DIFF_MULTIPLIER;
const LEVEL_2: usize = LEVEL_1 * LEVEL_DIFF_MULTIPLIER;
const LEVEL_3: usize = LEVEL_2 * LEVEL_DIFF_MULTIPLIER;
const LEVEL_4: usize = LEVEL_3 * LEVEL_DIFF_MULTIPLIER;

type LevelTrees = Vec<Box<SSLevelTree>>;

macro_rules! with_levels {
    ($($sym:ident, $level_size:ident;)*) => {
        $(
            type $sym = [EntryKey; $level_size];
            impl_sspage_slice!($sym, EntryKey, $level_size);
        )*

        fn init_lsm_level_trees(neb_client: &Arc<AsyncClient>) -> (LevelTrees, Vec<usize>, Vec<usize>) {
            let mut trees = LevelTrees::new();
            let mut max_elements_list = Vec::new();
            let mut max_elements = LEVEL_M_MAX_ELEMENTS_COUNT;
            let mut page_sizes = Vec::new();
            $(
                max_elements *= LEVEL_ELEMENTS_MULTIPLIER;
                trees.push(box LevelTree::<$sym>::new(neb_client));
                max_elements_list.push(max_elements);
                page_sizes.push($level_size);
                const_assert!($level_size * KEY_SIZE <= MAX_SEGMENT_SIZE);
            )*
            return (trees, max_elements_list, page_sizes);
        }
    };
}

with_levels!{
    L1, LEVEL_1;
    L2, LEVEL_2;
    L3, LEVEL_3;
    L4, LEVEL_4;
}

pub struct LSMTree {
    level_m: BPlusTree,
    trees: LevelTrees,
    // use Vec here for convenience
    max_sizes: Vec<usize>,
    page_sizes: Vec<usize>,
}

unsafe impl Send for LSMTree {}
unsafe impl Sync for LSMTree {}

impl LSMTree {
    pub fn new(neb_client: &Arc<AsyncClient>) -> Arc<Self> {
        let (trees, max_sizes, page_sizes) = init_lsm_level_trees(neb_client);
        let lsm_tree = Arc::new(LSMTree {
            level_m: BPlusTree::new(neb_client),
            trees,
            max_sizes,
            page_sizes,
        });
        let tree_clone = lsm_tree.clone();
        thread::spawn(move || {
            tree_clone.sentinel();
        });
        lsm_tree
    }

    pub fn insert(&self, mut key: EntryKey, id: &Id) -> Result<(), ()> {
        key_with_id(&mut key, id);
        self.level_m.insert(&key).map_err(|e| ())
    }

    pub fn remove(&self, mut key: EntryKey, id: &Id) -> Result<bool, ()> {
        key_with_id(&mut key, id);
        let m_deleted = self.level_m.remove(&key).map_err(|e| ())?;
        let levels_deleted = self
            .trees
            .iter()
            .map(|tree| tree.mark_deleted(&key))
            .collect_vec() // collect here to prevent short circuit
            .into_iter()
            .any(|d| d);
        Ok(m_deleted || levels_deleted)
    }

    pub fn seek(&self, mut key: EntryKey, ordering: Ordering) -> LSMTreeCursor {
        match ordering {
            Ordering::Forward => key_with_id(&mut key, &Id::unit_id()),
            Ordering::Backward => key_with_id(&mut key, &Id::new(::std::u64::MAX, ::std::u64::MAX)),
        };
        let mut cursors = vec![self.level_m.seek(&key, ordering).unwrap()];
        for tree in &self.trees {
            cursors.push(tree.seek(&key, ordering));
        }
        return LSMTreeCursor::new(
            cursors,
            self.trees.iter().map(|t| t.tombstones()).collect_vec(),
        );
    }

    fn sentinel(&self) {
        self.check_and_merge(
            LEVEL_M_MAX_ELEMENTS_COUNT,
            self.page_sizes[0],
            &mut BTreeSet::new(),
            &self.level_m,
            &*self.trees[0],
        );
        for i in 0..self.trees.len() - 1 {
            let upper = &*self.trees[i];
            let lower = &*self.trees[i + 1];
            let upper_tombstons_ref = upper.tombstones();
            let mut upper_tombstones = upper_tombstons_ref.write();
            self.check_and_merge(
                self.max_sizes[i],
                self.page_sizes[i + 1],
                &mut *upper_tombstones,
                upper,
                lower,
            );
        }
        thread::sleep(Duration::from_millis(100));
    }

    fn check_and_merge<U, L>(
        &self,
        max_elements: usize,
        upper_page_size: usize,
        upper_tombstones: &mut TombstonesInner,
        upper_level: &U,
        lower_level: &L,
    ) where
        U: MergeableTree + ?Sized,
        L: SSLevelTree + ?Sized,
    {
        let elements = upper_level.elements();
        if elements > max_elements {
            debug!(
                "upper level have elements {} exceeds {}, start level merging",
                elements, max_elements
            );
            let entries_to_merge = upper_page_size;
            let guard = upper_level.prepare_level_merge();
            let mut pages = vec![guard.last_page()];
            let mut collected_entries = pages.last().unwrap().keys().len();
            while collected_entries < entries_to_merge {
                let next_page = pages.last().unwrap().next();
                collected_entries += next_page.keys().len();
                pages.push(next_page);
            }
            let mut entries = vec![];
            for page in &pages {
                entries.append(&mut Vec::from(page.keys()));
            }
            debug!("going to merge {} entries", entries.len());
            lower_level.merge(entries.as_mut_slice(), upper_tombstones);
            guard.remove_pages(pages.iter().map(|p| p.keys()).collect_vec().as_slice())
        }
    }

    pub fn len(&self) -> usize {
        let mem_len = self.level_m.len();
        let levels_len_sum = self.trees.iter().map(|tree| tree.len()).sum::<usize>();
        return mem_len + levels_len_sum;
    }
}

pub struct LSMTreeCursor {
    level_cursors: Vec<Box<Cursor>>,
    tombstones: Vec<Arc<Tombstones>>,
}

impl LSMTreeCursor {
    fn new(cursors: Vec<Box<Cursor>>, tombstones: Vec<Arc<Tombstones>>) -> Self {
        LSMTreeCursor {
            tombstones,
            level_cursors: cursors,
        }
    }

    fn next_candidate(&mut self) -> (bool, usize) {
        let min_tree = self
            .level_cursors
            .iter()
            .enumerate()
            .map(|(i, cursor)| (i, cursor.current()))
            .filter_map(|(i, current)| current.map(|current_val| (i, current_val)))
            .min_by_key(|(i, val)| *val)
            .map(|(id, _)| id);
        if let Some(id) = min_tree {
            let min_has_next = self.level_cursors[id].next();
            if !min_has_next {
                let next = self
                    .level_cursors
                    .iter()
                    .enumerate()
                    .filter_map(|(level, cursor)| cursor.current().map(|key| (id, key)))
                    .min_by_key(|(level, key)| *key);
                if let Some((id, _)) = next {
                    return (true, id);
                }
            } else {
                return (true, id);
            }
        }
        return (false, 0);
    }
}

impl Cursor for LSMTreeCursor {
    fn next(&mut self) -> bool {
        loop {
            let (has_next, candidate_level) = self.next_candidate();
            if !has_next {
                break;
            }
            if candidate_level == 0 {
                return has_next;
            }
            let candidate_key = self.level_cursors[candidate_level].current().unwrap();
            if self.tombstones[candidate_level + 1]
                .read()
                .contains(candidate_key)
            {
                // marked in tombstone, skip
                continue;
            }
            return true;
        }
        return false;
    }

    fn current(&self) -> Option<&EntryKey> {
        self.level_cursors.iter().filter_map(|c| c.current()).min()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use byteorder::BigEndian;
    use byteorder::WriteBytesExt;
    use client;
    use futures::prelude::*;
    use rand::distributions::Uniform;
    use rand::thread_rng;
    use rand::Rng;
    use rayon::prelude::*;
    use server::NebServer;
    use server::ServerOptions;
    use std::env;
    use std::io::Cursor as StdCursor;

    fn u64_to_slice(n: u64) -> [u8; 8] {
        let mut key_slice = [0u8; 8];
        {
            let mut cursor = StdCursor::new(&mut key_slice[..]);
            cursor.write_u64::<BigEndian>(n);
        };
        key_slice
    }

    #[test]
    pub fn insertions() {
        env_logger::init();
        let server_group = "sstable_index_init";
        let server_addr = String::from("127.0.0.1:5300");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 1,
                memory_size: 8 * 1024 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
            },
            &server_addr,
            &server_group,
        );
        let client = Arc::new(
            client::AsyncClient::new(&server.rpc, &vec![server_addr], server_group).unwrap(),
        );
        client.new_schema_with_id(super::get_schema()).wait();
        let tree = LSMTree::new(&client);
        let num = env::var("LSM_TREE_TEST_ITEMS")
            // this value cannot do anything useful to the test
            // must arrange a long-term test to cover every levels
            .unwrap_or("1000".to_string())
            .parse::<u64>()
            .unwrap();

        (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let key = SmallVec::from_slice(&key_slice);
            tree.insert(key, &id).unwrap();
        });

        if tree.len() > LEVEL_M_MAX_ELEMENTS_COUNT {
            debug!("Sleep 5 minute for level merge");
            thread::sleep(Duration::from_secs(5 * 60));
        }

        debug!("Start validations");
        (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
            let mut rng = rand::rngs::OsRng::new().unwrap();
            let die_range = Uniform::new_inclusive(1, 6);
            let mut roll_die = rng.sample_iter(&die_range);
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let mut key = SmallVec::from_slice(&key_slice);
            debug!("checking: {}", i);
            let mut cursor = tree.seek(key.clone(), Ordering::Forward);
            key_with_id(&mut key, &id);
            assert_eq!(cursor.current(), Some(&key), "{}", i);
            if roll_die.next().unwrap() == 6 {
                for j in i..num {
                    let id = Id::new(0, j);
                    let key_slice = u64_to_slice(j);
                    let mut key = SmallVec::from_slice(&key_slice);
                    key_with_id(&mut key, &id);
                    assert_eq!(cursor.current(), Some(&key), "{}/{}", i, j);
                    assert_eq!(cursor.next(), j != num - 1, "{}/{}", i, j);
                }
            }
        });
    }
}
<commit_msg>reduce memory requirement for lsm tree test<commit_after>use super::btree::{BPlusTree, RTCursor as BPlusTreeCursor};
use super::sstable::*;
use super::*;
use client::AsyncClient;
use itertools::Itertools;
use parking_lot::RwLock;
use ram::segs::MAX_SEGMENT_SIZE;
use std::collections::BTreeSet;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use std::{mem, ptr};

const LEVEL_M_MAX_ELEMENTS_COUNT: usize = LEVEL_M * LEVEL_M * LEVEL_M;
const LEVEL_ELEMENTS_MULTIPLIER: usize = 10;
const LEVEL_DIFF_MULTIPLIER: usize = 10;
const LEVEL_M: usize = super::btree::NUM_KEYS;
const LEVEL_1: usize = LEVEL_M * LEVEL_DIFF_MULTIPLIER;
const LEVEL_2: usize = LEVEL_1 * LEVEL_DIFF_MULTIPLIER;
const LEVEL_3: usize = LEVEL_2 * LEVEL_DIFF_MULTIPLIER;
const LEVEL_4: usize = LEVEL_3 * LEVEL_DIFF_MULTIPLIER;

type LevelTrees = Vec<Box<SSLevelTree>>;

macro_rules! with_levels {
    ($($sym:ident, $level_size:ident;)*) => {
        $(
            type $sym = [EntryKey; $level_size];
            impl_sspage_slice!($sym, EntryKey, $level_size);
        )*

        fn init_lsm_level_trees(neb_client: &Arc<AsyncClient>) -> (LevelTrees, Vec<usize>, Vec<usize>) {
            let mut trees = LevelTrees::new();
            let mut max_elements_list = Vec::new();
            let mut max_elements = LEVEL_M_MAX_ELEMENTS_COUNT;
            let mut page_sizes = Vec::new();
            $(
                max_elements *= LEVEL_ELEMENTS_MULTIPLIER;
                trees.push(box LevelTree::<$sym>::new(neb_client));
                max_elements_list.push(max_elements);
                page_sizes.push($level_size);
                const_assert!($level_size * KEY_SIZE <= MAX_SEGMENT_SIZE);
            )*
            return (trees, max_elements_list, page_sizes);
        }
    };
}

with_levels!{
    L1, LEVEL_1;
    L2, LEVEL_2;
    L3, LEVEL_3;
    L4, LEVEL_4;
}

pub struct LSMTree {
    level_m: BPlusTree,
    trees: LevelTrees,
    // use Vec here for convenience
    max_sizes: Vec<usize>,
    page_sizes: Vec<usize>,
}

unsafe impl Send for LSMTree {}
unsafe impl Sync for LSMTree {}

impl LSMTree {
    pub fn new(neb_client: &Arc<AsyncClient>) -> Arc<Self> {
        let (trees, max_sizes, page_sizes) = init_lsm_level_trees(neb_client);
        let lsm_tree = Arc::new(LSMTree {
            level_m: BPlusTree::new(neb_client),
            trees,
            max_sizes,
            page_sizes,
        });
        let tree_clone = lsm_tree.clone();
        thread::spawn(move || {
            tree_clone.sentinel();
        });
        lsm_tree
    }

    pub fn insert(&self, mut key: EntryKey, id: &Id) -> Result<(), ()> {
        key_with_id(&mut key, id);
        self.level_m.insert(&key).map_err(|e| ())
    }

    pub fn remove(&self, mut key: EntryKey, id: &Id) -> Result<bool, ()> {
        key_with_id(&mut key, id);
        let m_deleted = self.level_m.remove(&key).map_err(|e| ())?;
        let levels_deleted = self
            .trees
            .iter()
            .map(|tree| tree.mark_deleted(&key))
            .collect_vec() // collect here to prevent short circuit
            .into_iter()
            .any(|d| d);
        Ok(m_deleted || levels_deleted)
    }

    pub fn seek(&self, mut key: EntryKey, ordering: Ordering) -> LSMTreeCursor {
        match ordering {
            Ordering::Forward => key_with_id(&mut key, &Id::unit_id()),
            Ordering::Backward => key_with_id(&mut key, &Id::new(::std::u64::MAX, ::std::u64::MAX)),
        };
        let mut cursors = vec![self.level_m.seek(&key, ordering).unwrap()];
        for tree in &self.trees {
            cursors.push(tree.seek(&key, ordering));
        }
        return LSMTreeCursor::new(
            cursors,
            self.trees.iter().map(|t| t.tombstones()).collect_vec(),
        );
    }

    fn sentinel(&self) {
        self.check_and_merge(
            LEVEL_M_MAX_ELEMENTS_COUNT,
            self.page_sizes[0],
            &mut BTreeSet::new(),
            &self.level_m,
            &*self.trees[0],
        );
        for i in 0..self.trees.len() - 1 {
            let upper = &*self.trees[i];
            let lower = &*self.trees[i + 1];
            let upper_tombstons_ref = upper.tombstones();
            let mut upper_tombstones = upper_tombstons_ref.write();
            self.check_and_merge(
                self.max_sizes[i],
                self.page_sizes[i + 1],
                &mut *upper_tombstones,
                upper,
                lower,
            );
        }
        thread::sleep(Duration::from_millis(100));
    }

    fn check_and_merge<U, L>(
        &self,
        max_elements: usize,
        upper_page_size: usize,
        upper_tombstones: &mut TombstonesInner,
        upper_level: &U,
        lower_level: &L,
    ) where
        U: MergeableTree + ?Sized,
        L: SSLevelTree + ?Sized,
    {
        let elements = upper_level.elements();
        if elements > max_elements {
            debug!(
                "upper level have elements {} exceeds {}, start level merging",
                elements, max_elements
            );
            let entries_to_merge = upper_page_size;
            let guard = upper_level.prepare_level_merge();
            let mut pages = vec![guard.last_page()];
            let mut collected_entries = pages.last().unwrap().keys().len();
            while collected_entries < entries_to_merge {
                let next_page = pages.last().unwrap().next();
                collected_entries += next_page.keys().len();
                pages.push(next_page);
            }
            let mut entries = vec![];
            for page in &pages {
                entries.append(&mut Vec::from(page.keys()));
            }
            debug!("going to merge {} entries", entries.len());
            lower_level.merge(entries.as_mut_slice(), upper_tombstones);
            guard.remove_pages(pages.iter().map(|p| p.keys()).collect_vec().as_slice())
        }
    }

    pub fn len(&self) -> usize {
        let mem_len = self.level_m.len();
        let levels_len_sum = self.trees.iter().map(|tree| tree.len()).sum::<usize>();
        return mem_len + levels_len_sum;
    }
}

pub struct LSMTreeCursor {
    level_cursors: Vec<Box<Cursor>>,
    tombstones: Vec<Arc<Tombstones>>,
}

impl LSMTreeCursor {
    fn new(cursors: Vec<Box<Cursor>>, tombstones: Vec<Arc<Tombstones>>) -> Self {
        LSMTreeCursor {
            tombstones,
            level_cursors: cursors,
        }
    }

    fn next_candidate(&mut self) -> (bool, usize) {
        let min_tree = self
            .level_cursors
            .iter()
            .enumerate()
            .map(|(i, cursor)| (i, cursor.current()))
            .filter_map(|(i, current)| current.map(|current_val| (i, current_val)))
            .min_by_key(|(i, val)| *val)
            .map(|(id, _)| id);
        if let Some(id) = min_tree {
            let min_has_next = self.level_cursors[id].next();
            if !min_has_next {
                let next = self
                    .level_cursors
                    .iter()
                    .enumerate()
                    .filter_map(|(level, cursor)| cursor.current().map(|key| (id, key)))
                    .min_by_key(|(level, key)| *key);
                if let Some((id, _)) = next {
                    return (true, id);
                }
            } else {
                return (true, id);
            }
        }
        return (false, 0);
    }
}

impl Cursor for LSMTreeCursor {
    fn next(&mut self) -> bool {
        loop {
            let (has_next, candidate_level) = self.next_candidate();
            if !has_next {
                break;
            }
            if candidate_level == 0 {
                return has_next;
            }
            let candidate_key = self.level_cursors[candidate_level].current().unwrap();
            if self.tombstones[candidate_level + 1]
                .read()
                .contains(candidate_key)
            {
                // marked in tombstone, skip
                continue;
            }
            return true;
        }
        return false;
    }

    fn current(&self) -> Option<&EntryKey> {
        self.level_cursors.iter().filter_map(|c| c.current()).min()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use byteorder::BigEndian;
    use byteorder::WriteBytesExt;
    use client;
    use futures::prelude::*;
    use rand::distributions::Uniform;
    use rand::thread_rng;
    use rand::Rng;
    use rayon::prelude::*;
    use server::NebServer;
    use server::ServerOptions;
    use std::env;
    use std::io::Cursor as StdCursor;

    fn u64_to_slice(n: u64) -> [u8; 8] {
        let mut key_slice = [0u8; 8];
        {
            let mut cursor = StdCursor::new(&mut key_slice[..]);
            cursor.write_u64::<BigEndian>(n);
        };
        key_slice
    }

    #[test]
    pub fn insertions() {
        env_logger::init();
        let server_group = "sstable_index_init";
        let server_addr = String::from("127.0.0.1:5300");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 1,
                memory_size: 4 * 1024 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
            },
            &server_addr,
            &server_group,
        );
        let client = Arc::new(
            client::AsyncClient::new(&server.rpc, &vec![server_addr], server_group).unwrap(),
        );
        client.new_schema_with_id(super::get_schema()).wait();
        let tree = LSMTree::new(&client);
        let num = env::var("LSM_TREE_TEST_ITEMS")
            // this value cannot do anything useful to the test
            // must arrange a long-term test to cover every levels
            .unwrap_or("1000".to_string())
            .parse::<u64>()
            .unwrap();

        (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let key = SmallVec::from_slice(&key_slice);
            tree.insert(key, &id).unwrap();
        });

        if tree.len() > LEVEL_M_MAX_ELEMENTS_COUNT {
            debug!("Sleep 5 minute for level merge");
            thread::sleep(Duration::from_secs(5 * 60));
        }

        debug!("Start validations");
        (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
            let mut rng = rand::rngs::OsRng::new().unwrap();
            let die_range = Uniform::new_inclusive(1, 6);
            let mut roll_die = rng.sample_iter(&die_range);
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let mut key = SmallVec::from_slice(&key_slice);
            debug!("checking: {}", i);
            let mut cursor = tree.seek(key.clone(), Ordering::Forward);
            key_with_id(&mut key, &id);
            assert_eq!(cursor.current(), Some(&key), "{}", i);
            if roll_die.next().unwrap() == 6 {
                for j in i..num {
                    let id = Id::new(0, j);
                    let key_slice = u64_to_slice(j);
                    let mut key = SmallVec::from_slice(&key_slice);
                    key_with_id(&mut key, &id);
                    assert_eq!(cursor.current(), Some(&key), "{}/{}", i, j);
                    assert_eq!(cursor.next(), j != num - 1, "{}/{}", i, j);
                }
            }
        });
    }
}
<|endoftext|>


<commit_before>use bifrost::utils::fut_exec::wait;
use client::AsyncClient;
use dovahkiin::types::custom_types::bytes::SmallBytes;
use dovahkiin::types::custom_types::id::Id;
use dovahkiin::types::custom_types::map::Map;
use dovahkiin::types::type_id_of;
use dovahkiin::types::value::ToValue;
use futures::prelude::*;
use index::btree::Ordering;
use index::EntryKey;
use index::Slice;
use index::{id_from_key, key_with_id};
use itertools::Itertools;
use parking_lot::Mutex;
use parking_lot::RwLock;
use ram::cell::Cell;
use ram::schema::Field;
use ram::schema::Schema;
use ram::types::*;
use std::cell::RefCell;
use std::collections::btree_map::Range;
use std::collections::hash_map::DefaultHasher;
use std::collections::{BTreeMap, BTreeSet, LinkedList};
use std::fmt::Debug;
use std::marker::PhantomData;
use std::mem;
use std::sync::Arc;
use utils::lru_cache::LRUCache;

// LevelTree items cannot been added or removed individually
// Items must been merged from higher level in bulk
// Deletion will be performed when merging to or from higher level tree
// Because tree update will not to be performed in parallel. Unlike memtable, a single r/w lock
// should be sufficient. Thus concurrency control will be simple and efficient.

type Tombstones = BTreeSet<EntryKey>;
type PageCache<S> = Arc<Mutex<LRUCache<Id, Arc<SSPage<S>>>>>;

const PAGE_SCHEMA: &'static str = "NEB_SSTABLE_PAGE";

lazy_static! {
    static ref PAGE_SCHEMA_ID: u32 = key_hash(PAGE_SCHEMA) as u32;
}

pub struct LevelTree<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    neb_client: Arc<AsyncClient>,
    index: BTreeMap<EntryKey, Id>,
    tombstones: Tombstones,
    pages: PageCache<S>,
    marker: PhantomData<S>,
}
struct SSPage<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    id: Id,
    slice: S,
    len: usize,
}

impl<S> SSPage<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    pub fn from_cell(cell: Cell) -> Self {
        let mut slice = S::init();
        let slice_len = slice.as_slice().len();
        let cell_id = cell.id();
        let keys = cell.data;
        let keys_len = keys.len().unwrap();
        let keys_array = if let Value::PrimArray(PrimitiveArray::SmallBytes(ref array)) = keys {
            array
        } else {
            panic!()
        };
        for (i, val) in keys_array.iter().enumerate() {
            slice.as_slice()[i] = EntryKey::from(val.as_slice());
        }
        Self {
            id: cell_id,
            slice,
            len: keys_array.len(),
        }
    }

    pub fn to_cell(&self) -> Cell {
        let value = self
            .slice
            .as_slice_immute()
            .iter()
            .map(|key| SmallBytes::from_vec(key.as_slice().to_vec()))
            .take(self.len)
            .collect_vec()
            .value();
        Cell::new_with_id(*PAGE_SCHEMA_ID, &self.id, value)
    }

    pub fn new() -> Self {
        SSPage {
            slice: S::init(),
            id: Id::unit_id(),
            len: 0,
        }
    }
}

impl<S> LevelTree<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    pub fn new(neb_client: &Arc<AsyncClient>) -> Self {
        let client = neb_client.clone();
        let mut index = BTreeMap::new();
        Self {
            neb_client: neb_client.clone(),
            index,
            tombstones: Tombstones::new(),
            // TODO: carefully set the capacity
            pages: Arc::new(Mutex::new(LRUCache::new(
                100,
                move |id| {
                    client
                        .read_cell(*id)
                        .wait()
                        .unwrap()
                        .map(|cell| Arc::new(SSPage::from_cell(cell)))
                        .ok()
                },
                |_, _| {
                    // will do nothing on evict for pages will be write to back storage once updated
                },
            ))),
            marker: PhantomData {},
        }
    }

    pub fn seek(&self, key: &EntryKey, ordering: Ordering) -> impl Cursor {
        debug!("Seeking for {:?}, in index len {}", key, self.index.len());
        let mut range = match ordering {
            Ordering::Forward => self.index.range::<EntryKey, _>(key..),
            Ordering::Backward => self.index.range::<EntryKey, _>(..=key),
        };
        let first = match ordering {
            Ordering::Forward => range.next(),
            Ordering::Backward => range.next_back(),
        }
        .map(|(_, page_id)| {
            debug!("First page id is {:?}", page_id);
            self.pages.lock().get_or_fetch(page_id).unwrap().clone()
        });

        let (pos, page_len) = if let Some(ref page) = first {
            debug!("First page has {} elements", page.len);
            (
                page.slice.as_slice_immute()[..page.len]
                    .binary_search(key)
                    .unwrap_or_else(|i| i),
                page.len,
            )
        } else {
            (0, 0)
        };

        debug!("First page search pos is {}", pos);

        let mut cursor = RTCursor {
            ordering,
            range,
            page_len,
            current: first,
            index: pos,
            cache: self.pages.clone(),
        };

        if cursor.current.is_some() {
            match ordering {
                Ordering::Forward => {
                    if pos >= page_len {
                        cursor.next();
                    }
                }
                Ordering::Backward => {}
            }
        }

        cursor
    }

    pub fn merge(&mut self, mut source: Vec<EntryKey>, tombstones: &mut Tombstones) {
        let mut pages_cache = self.pages.lock();
        let (keys_to_removed, mut merged, mut ids_to_reuse) = {
            let (target_keys, target_pages): (Vec<_>, Vec<_>) = {
                let first = source.first().unwrap();
                let last = source.last().unwrap();
                self.index
                    .range::<EntryKey, _>(first..=last)
                    .map(|(k, id)| (k, pages_cache.get_or_fetch(id).unwrap().clone()))
                    .unzip()
            };
            let mut merged: Vec<SSPage<S>> = vec![SSPage::new()];
            let page_size = merged[0].slice.len();
            let target_page_ids = target_pages.iter().map(|p| p.id).collect::<LinkedList<_>>();
            let mut target = target_pages
                .iter()
                .map(|page| &page.slice.as_slice_immute()[..page.len])
                .flat_map(|x| x)
                .map(|x| (*x).clone())
                .collect_vec();

            let src_len = source.len();
            let tgt_len = target.len();

            debug!("Src len {}, target len {}", src_len, tgt_len);

            let mut srci = 0;
            let mut tgti = 0;

            while srci < src_len || tgti < tgt_len {
                let mut from_source = false;
                let mut lowest: &mut EntryKey = if srci >= src_len {
                    // source overflowed
                    tgti += 1;
                    &mut target[tgti - 1]
                } else if tgti >= tgt_len {
                    // target overflowed
                    srci += 1;
                    from_source = true;
                    &mut source[srci - 1]
                } else {
                    // normal case, compare
                    let t = &mut target[tgti];
                    let s = &mut source[srci];
                    if t <= s {
                        tgti += 1;
                        t
                    } else {
                        srci += 1;
                        from_source = true;
                        s
                    }
                };

                if from_source {
                    if tombstones.remove(lowest) {
                        continue;
                    }
                } else {
                    if self.tombstones.remove(lowest) {
                        continue;
                    }
                }

                let merging_page_len = {
                    let merging_page = merged.last_mut().unwrap();
                    let page_len = merging_page.len;
                    debug!("i is {}, picking {:?}", page_len, lowest);
                    mem::swap(&mut merging_page.slice.as_slice()[page_len], lowest);
                    merging_page.len += 1;
                    merging_page.len
                };

                if merging_page_len >= page_size {
                    merged.push(SSPage::new());
                }
            }
            if merged.last().unwrap().len == 0 {
                let merged_len = merged.len();
                merged.remove(merged_len - 1);
            }
            (
                target_keys.into_iter().map(|x| x.clone()).collect_vec(),
                merged,
                target_page_ids,
            )
        };

        debug!("Merged {} pages", merged.len());

        for k in keys_to_removed {
            debug!("Remove stale page index {:?}", k);
            self.index.remove(&k).unwrap();
        }

        for id in &ids_to_reuse {
            debug!("Remove stale page cache {:?}", id);
            pages_cache.remove(id);
        }

        for mut page in merged {
            let reuse_id = ids_to_reuse.pop_back();
            let id = reuse_id.unwrap_or_else(|| Id::rand());
            let index_key = page.slice.as_slice_immute()[0].clone();
            page.id = id;
            let cell = page.to_cell();
            if reuse_id.is_some() {
                debug!("Reuse page id by updating {:?}", cell.id());
                self.neb_client.update_cell(cell).wait().unwrap();
            } else {
                debug!("Inserting page with id {:?}", cell.id());
                self.neb_client.write_cell(cell).wait().unwrap();
            }

            debug!("Insert page index key {:?}, id {:?}", index_key, id);
            self.index.insert(index_key, id);
        }
    }
}

pub trait SortableEntrySlice: Sized + Slice<Item = EntryKey> {}

pub trait Cursor<'a> {
    fn next(&mut self) -> bool;
    fn current(&self) -> Option<&EntryKey>;
}

pub struct RTCursor<'a, S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    index: usize,
    ordering: Ordering,
    page_len: usize,
    range: Range<'a, EntryKey, Id>,
    current: Option<Arc<SSPage<S>>>,
    cache: PageCache<S>,
}

impl<'a, S> Cursor<'a> for RTCursor<'a, S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    fn next(&mut self) -> bool {
        if self.current.is_some() {
            match self.ordering {
                Ordering::Forward => {
                    if self.index + 1 >= self.page_len {
                        if let Some((_, page_id)) = self.range.next() {
                            let next_page = self.get_page(page_id);
                            self.current = Some(next_page);
                            self.index = 0;
                        } else {
                            self.current = None;
                            return false;
                        }
                    } else {
                        self.index += 1;
                    }
                }
                Ordering::Backward => {
                    if self.index == 0 {
                        if let Some((_, page_id)) = self.range.next_back() {
                            let next_page = self.get_page(page_id);
                            self.current = Some(next_page);
                            self.index = self.page_len - 1;
                        } else {
                            self.current = None;
                            return false;
                        }
                    } else {
                        self.index -= 1;
                    }
                }
            }
            true
        } else {
            false
        }
    }

    fn current(&self) -> Option<&EntryKey> {
        if self.index < 0 || self.index >= self.page_len {
            return None;
        }
        if let Some(ref page) = self.current {
            Some(&page.slice.as_slice_immute()[self.index])
        } else {
            None
        }
    }
}

impl<'a, S> RTCursor<'a, S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    fn get_page(&self, id: &Id) -> Arc<SSPage<S>> {
        self.cache.lock().get_or_fetch(id).unwrap().clone()
    }
}

pub fn get_schema() -> Schema {
    Schema {
        id: *PAGE_SCHEMA_ID,
        name: String::from(PAGE_SCHEMA),
        key_field: None,
        str_key_field: None,
        is_dynamic: false,
        fields: Field::new("*", type_id_of(Type::SmallBytes), false, true, None),
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use client;
    use futures::prelude::*;
    use index::sstable::LevelTree;
    use server::NebServer;
    use server::ServerOptions;
    use std::ptr;
    use std::sync::Arc;

    type SmallPage = [EntryKey; 5];

    impl_sspage_slice!(SmallPage, EntryKey, 5);

    #[test]
    pub fn init() {
        env_logger::init();
        let server_group = "sstable_index_init";
        let server_addr = String::from("127.0.0.1:5200");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 1,
                memory_size: 16 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
            },
            &server_addr,
            &server_group,
        );
        let client = Arc::new(
            client::AsyncClient::new(&server.rpc, &vec![server_addr], server_group).unwrap(),
        );
        client.new_schema_with_id(super::get_schema()).wait();

        let mut tree: LevelTree<SmallPage> = LevelTree::new(&client);
        let id = Id::unit_id();
        let key = smallvec![1, 2, 3, 4, 5, 6];
        let mut tombstones = BTreeSet::new();
        info!("test insertion");

        let mut key_id = key.clone();
        key_with_id(&mut key_id, &id);

        tree.merge(vec![key_id], &mut tombstones);
        let mut cursor = tree.seek(&key, Ordering::Forward);
        assert_eq!(id_from_key(cursor.current().unwrap()), id);
    }
}
<commit_msg>refactor for parallel and trait friendly<commit_after>use bifrost::utils::fut_exec::wait;
use client::AsyncClient;
use dovahkiin::types::custom_types::bytes::SmallBytes;
use dovahkiin::types::custom_types::id::Id;
use dovahkiin::types::custom_types::map::Map;
use dovahkiin::types::type_id_of;
use dovahkiin::types::value::ToValue;
use futures::prelude::*;
use index::btree::Ordering;
use index::EntryKey;
use index::Slice;
use index::{id_from_key, key_with_id};
use itertools::Itertools;
use parking_lot::Mutex;
use parking_lot::RwLock;
use ram::cell::Cell;
use ram::schema::Field;
use ram::schema::Schema;
use ram::types::*;
use std::cell::RefCell;
use std::collections::btree_map::Range;
use std::collections::hash_map::DefaultHasher;
use std::collections::{BTreeMap, BTreeSet, LinkedList};
use std::fmt::Debug;
use std::marker::PhantomData;
use std::mem;
use std::ops::Bound::*;
use std::ops::RangeBounds;
use std::sync::Arc;
use utils::lru_cache::LRUCache;

// LevelTree items cannot been added or removed individually
// Items must been merged from higher level in bulk
// Deletion will be performed when merging to or from higher level tree
// Because tree update will not to be performed in parallel. Unlike memtable, a single r/w lock
// should be sufficient. Thus concurrency control will be simple and efficient.

type Tombstones = BTreeSet<EntryKey>;
type PageCache<S> = Arc<Mutex<LRUCache<Id, Arc<SSPage<S>>>>>;
type PageIndex = Arc<RwLock<BTreeMap<EntryKey, Id>>>;

const PAGE_SCHEMA: &'static str = "NEB_SSTABLE_PAGE";

lazy_static! {
    static ref PAGE_SCHEMA_ID: u32 = key_hash(PAGE_SCHEMA) as u32;
}

pub struct LevelTree<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    neb_client: Arc<AsyncClient>,
    index: PageIndex,
    tombstones: Tombstones,
    pages: PageCache<S>,
    marker: PhantomData<S>,
}
struct SSPage<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    id: Id,
    slice: S,
    len: usize,
}

impl<S> SSPage<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    pub fn from_cell(cell: Cell) -> Self {
        let mut slice = S::init();
        let slice_len = slice.as_slice().len();
        let cell_id = cell.id();
        let keys = cell.data;
        let keys_len = keys.len().unwrap();
        let keys_array = if let Value::PrimArray(PrimitiveArray::SmallBytes(ref array)) = keys {
            array
        } else {
            panic!()
        };
        for (i, val) in keys_array.iter().enumerate() {
            slice.as_slice()[i] = EntryKey::from(val.as_slice());
        }
        Self {
            id: cell_id,
            slice,
            len: keys_array.len(),
        }
    }

    pub fn to_cell(&self) -> Cell {
        let value = self
            .slice
            .as_slice_immute()
            .iter()
            .map(|key| SmallBytes::from_vec(key.as_slice().to_vec()))
            .take(self.len)
            .collect_vec()
            .value();
        Cell::new_with_id(*PAGE_SCHEMA_ID, &self.id, value)
    }

    pub fn new() -> Self {
        SSPage {
            slice: S::init(),
            id: Id::unit_id(),
            len: 0,
        }
    }
}

pub trait Tree {
    fn seek(&self, key: &EntryKey, ordering: Ordering) -> Box<Cursor>;
    fn merge(&mut self, mut source: Vec<EntryKey>, tombstones: &mut Tombstones);
}

impl<S> LevelTree<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice + 'static,
{
    pub fn new(neb_client: &Arc<AsyncClient>) -> Self {
        let client = neb_client.clone();
        let mut index = Arc::new(RwLock::new(BTreeMap::new()));
        Self {
            neb_client: neb_client.clone(),
            index,
            tombstones: Tombstones::new(),
            // TODO: carefully set the capacity
            pages: Arc::new(Mutex::new(LRUCache::new(
                100,
                move |id| {
                    client
                        .read_cell(*id)
                        .wait()
                        .unwrap()
                        .map(|cell| Arc::new(SSPage::from_cell(cell)))
                        .ok()
                },
                |_, _| {
                    // will do nothing on evict for pages will be write to back storage once updated
                },
            ))),
            marker: PhantomData {},
        }
    }

    pub fn seek(&self, key: &EntryKey, ordering: Ordering) -> Box<Cursor> {
        let index = self.index.read();
        debug!("Seeking for {:?}, in index len {}", key, index.len());
        let mut range = match ordering {
            Ordering::Forward => index.range::<EntryKey, _>(key..),
            Ordering::Backward => index.range::<EntryKey, _>(..=key),
        };
        let first = match ordering {
            Ordering::Forward => range.next(),
            Ordering::Backward => range.next_back(),
        }
        .map(|(_, page_id)| {
            debug!("First page id is {:?}", page_id);
            self.pages.lock().get_or_fetch(page_id).unwrap().clone()
        });

        let (pos, page_len) = if let Some(ref page) = first {
            debug!("First page has {} elements", page.len);
            (
                page.slice.as_slice_immute()[..page.len]
                    .binary_search(key)
                    .unwrap_or_else(|i| i),
                page.len,
            )
        } else {
            (0, 0)
        };

        debug!("First page search pos is {}", pos);

        let mut cursor = RTCursor {
            ordering,
            page_len,
            current: first,
            pos,
            cache: self.pages.clone(),
            index: self.index.clone(),
        };

        if cursor.current.is_some() {
            match ordering {
                Ordering::Forward => {
                    if pos >= page_len {
                        cursor.next();
                    }
                }
                Ordering::Backward => {}
            }
        }

        box cursor
    }

    pub fn merge(&mut self, mut source: Vec<EntryKey>, tombstones: &mut Tombstones) {
        let mut pages_cache = self.pages.lock();
        let mut index = self.index.write();
        let (keys_to_removed, mut merged, mut ids_to_reuse) = {
            let (target_keys, target_pages): (Vec<_>, Vec<_>) = {
                let first = source.first().unwrap();
                let last = source.last().unwrap();
                index
                    .range::<EntryKey, _>(first..=last)
                    .map(|(k, id)| (k, pages_cache.get_or_fetch(id).unwrap().clone()))
                    .unzip()
            };
            let mut merged: Vec<SSPage<S>> = vec![SSPage::new()];
            let page_size = merged[0].slice.len();
            let target_page_ids = target_pages.iter().map(|p| p.id).collect::<LinkedList<_>>();
            let mut target = target_pages
                .iter()
                .map(|page| &page.slice.as_slice_immute()[..page.len])
                .flat_map(|x| x)
                .map(|x| (*x).clone())
                .collect_vec();

            let src_len = source.len();
            let tgt_len = target.len();

            debug!("Src len {}, target len {}", src_len, tgt_len);

            let mut srci = 0;
            let mut tgti = 0;

            while srci < src_len || tgti < tgt_len {
                let mut from_source = false;
                let mut lowest: &mut EntryKey = if srci >= src_len {
                    // source overflowed
                    tgti += 1;
                    &mut target[tgti - 1]
                } else if tgti >= tgt_len {
                    // target overflowed
                    srci += 1;
                    from_source = true;
                    &mut source[srci - 1]
                } else {
                    // normal case, compare
                    let t = &mut target[tgti];
                    let s = &mut source[srci];
                    if t <= s {
                        tgti += 1;
                        t
                    } else {
                        srci += 1;
                        from_source = true;
                        s
                    }
                };

                if from_source {
                    if tombstones.remove(lowest) {
                        continue;
                    }
                } else {
                    if self.tombstones.remove(lowest) {
                        continue;
                    }
                }

                let merging_page_len = {
                    let merging_page = merged.last_mut().unwrap();
                    let page_len = merging_page.len;
                    debug!("i is {}, picking {:?}", page_len, lowest);
                    mem::swap(&mut merging_page.slice.as_slice()[page_len], lowest);
                    merging_page.len += 1;
                    merging_page.len
                };

                if merging_page_len >= page_size {
                    merged.push(SSPage::new());
                }
            }
            if merged.last().unwrap().len == 0 {
                let merged_len = merged.len();
                merged.remove(merged_len - 1);
            }
            (
                target_keys.into_iter().map(|x| x.clone()).collect_vec(),
                merged,
                target_page_ids,
            )
        };

        debug!("Merged {} pages", merged.len());

        for k in keys_to_removed {
            debug!("Remove stale page index {:?}", k);
            index.remove(&k).unwrap();
        }

        for id in &ids_to_reuse {
            debug!("Remove stale page cache {:?}", id);
            pages_cache.remove(id);
        }

        for mut page in merged {
            let reuse_id = ids_to_reuse.pop_back();
            let id = reuse_id.unwrap_or_else(|| Id::rand());
            let index_key = page.slice.as_slice_immute()[0].clone();
            page.id = id;
            let cell = page.to_cell();
            if reuse_id.is_some() {
                debug!("Reuse page id by updating {:?}", cell.id());
                self.neb_client.update_cell(cell).wait().unwrap();
            } else {
                debug!("Inserting page with id {:?}", cell.id());
                self.neb_client.write_cell(cell).wait().unwrap();
            }

            debug!("Insert page index key {:?}, id {:?}", index_key, id);
            index.insert(index_key, id);
        }
    }
}

pub trait SortableEntrySlice: Sized + Slice<Item = EntryKey> {}

pub trait Cursor {
    fn next(&mut self) -> bool;
    fn current(&self) -> Option<&EntryKey>;
}

pub struct RTCursor<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    pos: usize,
    ordering: Ordering,
    page_len: usize,
    current: Option<Arc<SSPage<S>>>,
    cache: PageCache<S>,
    index: PageIndex,
}

impl<S> Cursor for RTCursor<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    fn next(&mut self) -> bool {
        if self.current.is_some() {
            match self.ordering {
                Ordering::Forward => {
                    if self.pos + 1 >= self.page_len {
                        if let Some((_, page_id)) = self
                            .index
                            .read()
                            .range::<EntryKey, _>((Excluded(self.current().unwrap()), Unbounded))
                            .next()
                        {
                            let next_page = self.get_page(page_id);
                            self.current = Some(next_page);
                            self.pos = 0;
                        } else {
                            self.current = None;
                            return false;
                        }
                    } else {
                        self.pos += 1;
                    }
                }
                Ordering::Backward => {
                    if self.pos == 0 {
                        if let Some((_, page_id)) = self
                            .index
                            .read()
                            .range::<EntryKey, _>((Unbounded, Excluded(self.current().unwrap())))
                            .next_back()
                        {
                            let next_page = self.get_page(page_id);
                            self.current = Some(next_page);
                            self.pos = self.page_len - 1;
                        } else {
                            self.current = None;
                            return false;
                        }
                    } else {
                        self.pos -= 1;
                    }
                }
            }
            true
        } else {
            false
        }
    }

    fn current(&self) -> Option<&EntryKey> {
        if self.pos < 0 || self.pos >= self.page_len {
            return None;
        }
        if let Some(ref page) = self.current {
            Some(&page.slice.as_slice_immute()[self.pos])
        } else {
            None
        }
    }
}

impl<S> RTCursor<S>
where
    S: Slice<Item = EntryKey> + SortableEntrySlice,
{
    fn get_page(&self, id: &Id) -> Arc<SSPage<S>> {
        self.cache.lock().get_or_fetch(id).unwrap().clone()
    }
}

pub fn get_schema() -> Schema {
    Schema {
        id: *PAGE_SCHEMA_ID,
        name: String::from(PAGE_SCHEMA),
        key_field: None,
        str_key_field: None,
        is_dynamic: false,
        fields: Field::new("*", type_id_of(Type::SmallBytes), false, true, None),
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use client;
    use futures::prelude::*;
    use index::sstable::LevelTree;
    use server::NebServer;
    use server::ServerOptions;
    use std::ptr;
    use std::sync::Arc;

    type SmallPage = [EntryKey; 5];

    impl_sspage_slice!(SmallPage, EntryKey, 5);

    #[test]
    pub fn init() {
        env_logger::init();
        let server_group = "sstable_index_init";
        let server_addr = String::from("127.0.0.1:5200");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 1,
                memory_size: 16 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
            },
            &server_addr,
            &server_group,
        );
        let client = Arc::new(
            client::AsyncClient::new(&server.rpc, &vec![server_addr], server_group).unwrap(),
        );
        client.new_schema_with_id(super::get_schema()).wait();

        let mut tree: LevelTree<SmallPage> = LevelTree::new(&client);
        let id = Id::unit_id();
        let key = smallvec![1, 2, 3, 4, 5, 6];
        let mut tombstones = BTreeSet::new();
        info!("test insertion");

        let mut key_id = key.clone();
        key_with_id(&mut key_id, &id);

        tree.merge(vec![key_id], &mut tombstones);
        let mut cursor = tree.seek(&key, Ordering::Forward);
        assert_eq!(id_from_key(cursor.current().unwrap()), id);
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 Pierre Talbot (IRCAM)

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::iter::repeat;
use propagator::Propagator;
use variable::Variable;
use dependencies::VarEventDependencies;
use agenda::Agenda;
use std::rc::Rc;
use std::cell::RefCell;

pub struct Solver<V, D, A> {
  propagators: Vec<Box<Propagator + 'static>>,
  variables: Vec<Rc<RefCell<V>>>,
  deps: D,
  agenda: A
}

impl<'d, V, D, A> Solver<V, D, A> where
 V: Variable,
 D: VarEventDependencies<'d>,
 A: Agenda
{
  pub fn new() -> Solver<V, D, A> {
    Solver {
      propagators: vec![],
      variables: vec![],
      deps: VarEventDependencies::new(0,0),
      agenda: Agenda::new(0)
    }
  }

  pub fn newvar(&mut self, dom: <V as Variable>::Domain) -> Rc<RefCell<V>> {
    let var_idx = self.variables.len();
    self.variables.push(Rc::new(RefCell::new(Variable::new(var_idx as u32, dom))));
    self.variables[var_idx].clone()
  }
}

#[cfg(test)]
mod test {
  use super::*;
  use propagator::*;
  use fd::var::*;
  use agenda::RelaxedFifoAgenda;
  use dependencies::VarEventDepsVector;

  type FDSolver = Solver<FDVar, VarEventDepsVector, RelaxedFifoAgenda>;

  #[test]
  fn newvar_test() {
    let mut solver: FDSolver = Solver::new();
    let var1 = solver.newvar(Interval::new(1,4));
    let var2 = solver.newvar(Interval::new(1,4));
  }
}
<commit_msg>Prepare the deps and agenda for the main solving algorithm.<commit_after>// Copyright 2015 Pierre Talbot (IRCAM)

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::iter::repeat;
use propagator::Propagator;
use variable::Variable;
use dependencies::VarEventDependencies;
use agenda::Agenda;
use event::VarEvent;
use std::rc::Rc;
use std::cell::RefCell;

pub enum Status {
  Satisfiable,
  Unsatisfiable,
  Unknown
}

pub struct Solver<V: Variable, D, A> {
  propagators: Vec<Box<Propagator<Event=V::Event> + 'static>>,
  variables: Vec<Rc<RefCell<V>>>,
  deps: D,
  agenda: A
}

impl<'d, V, D, A> Solver<V, D, A> where
 V: Variable,
 D: VarEventDependencies<'d>,
 A: Agenda
{
  pub fn new() -> Solver<V, D, A> {
    Solver {
      propagators: vec![],
      variables: vec![],
      deps: VarEventDependencies::new(0,0),
      agenda: Agenda::new(0)
    }
  }

  pub fn newvar(&mut self, dom: <V as Variable>::Domain) -> Rc<RefCell<V>> {
    let var_idx = self.variables.len();
    self.variables.push(Rc::new(RefCell::new(Variable::new(var_idx as u32, dom))));
    self.variables[var_idx].clone()
  }

  pub fn add(&mut self, p: Box<Propagator<Event=<V as Variable>::Event> + 'static>) {
    self.propagators.push(p);
  }

  pub fn solve(&mut self) -> Status {
    self.prepare();
    Status::Satisfiable
  }

  fn prepare(&mut self) {
    self.init_deps();
    self.init_agenda();
  }

  fn init_deps(&mut self) {
    let dummy: Option<<V as Variable>::Event> = None;
    self.deps = VarEventDependencies::new(self.variables.len(), VarEvent::size(dummy));
    for (p_idx, p) in self.propagators.iter().enumerate() {
      for (v, ev) in p.dependencies().into_iter() {
        self.deps.subscribe(v as usize, ev, p_idx);
      }
    }
  }

  fn init_agenda(&mut self) {
    self.agenda = Agenda::new(self.propagators.len());
  }
}

#[cfg(test)]
mod test {
  use super::*;
  use fd::var::*;
  use fd::propagator::*;
  use propagator::Propagator;
  use agenda::RelaxedFifoAgenda;
  use dependencies::VarEventDepsVector;

  type FDSolver = Solver<FDVar, VarEventDepsVector, RelaxedFifoAgenda>;

  #[test]
  fn newvar_test() {
    let mut solver: FDSolver = Solver::new();
    let var1 = solver.newvar(Interval::new(1,4));
    let var2 = solver.newvar(Interval::new(1,4));

    let prop: Box<Propagator<Event=FDEvent>> = Box::new(XLessThanY::new(var1, var2));
    solver.add(prop);
  }
}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate test;
extern crate voronoi;

#[cfg(test)]
mod tests {
    use test::{Bencher, black_box};
    use voronoi::in_circle::*;
    use voronoi::geometry::*;

    #[bench]
    fn bench_triangle_in_circle_test(b: &mut Bencher) {
        let t = Triangle::new(
            Point2D::new(0.0, 0.0),
            Point2D::new(0.0, 1.0),
            Point2D::new(1.0, 1.0),
        );

        let test_point = Point2D::new(0.0, 0.0);

        b.iter(|| {
            let t_box = black_box(&t);
            let test_box = black_box(&test_point);

            t_box.in_circle_test(test_box)
        });
    }

    #[bench]
    fn bench_triangle_point_location_test(b: &mut Bencher) {
        let t = Triangle::new(
            Point2D::new(0.0, 0.0),
            Point2D::new(0.0, 10.0),
            Point2D::new(10.0, 1.0),
        );

        let test_point = Point2D::new(2.0, 2.0);

        b.iter(|| {
            let t_box = black_box(&t);
            let test_box = black_box(&test_point);

            t_box.locate(test_box)
        });
    }
}
<commit_msg>Add benchmark for in-circle test of tetrahedron.<commit_after>#![feature(test)]

extern crate test;
extern crate voronoi;

#[cfg(test)]
mod tests {
    use test::{Bencher, black_box};
    use voronoi::in_circle::*;
    use voronoi::geometry::*;

    #[bench]
    fn bench_triangle_in_circle_test(b: &mut Bencher) {
        let t = Triangle::new(
            Point2D::new(0.0, 0.0),
            Point2D::new(0.0, 1.0),
            Point2D::new(1.0, 1.0),
        );

        let test_point = Point2D::new(0.0, 0.0);

        b.iter(|| {
            let t_box = black_box(&t);
            let test_box = black_box(&test_point);

            t_box.in_circle_test(test_box)
        });
    }

    #[bench]
    fn bench_tetrahedron_in_circle_test(b: &mut Bencher) {
        let t = Tetrahedron::new(
            Point3D::new(1.0, 0.0, 0.0),
            Point3D::new(0.0, 1.0, 0.0),
            Point3D::new(0.0, 0.0, 1.0),
            Point3D::new(0.0, 0.0, 0.0),
        );

        let test_point = Point3D::new(0.0, 0.0, 0.0);

        b.iter(|| {
            let t_box = black_box(&t);
            let test_box = black_box(&test_point);

            t_box.in_circle_test(test_box)
        });
    }

    #[bench]
    fn bench_triangle_point_location_test(b: &mut Bencher) {
        let t = Triangle::new(
            Point2D::new(0.0, 0.0),
            Point2D::new(0.0, 10.0),
            Point2D::new(10.0, 1.0),
        );

        let test_point = Point2D::new(2.0, 2.0);

        b.iter(|| {
            let t_box = black_box(&t);
            let test_box = black_box(&test_point);

            t_box.locate(test_box)
        });
    }
}
<|endoftext|>