<commit_before>#[doc = "Finds docs for reexported items and duplicates them"];

import std::map;
import std::map::hashmap;
import rustc::syntax::ast;
import rustc::syntax::ast_util;
import rustc::util::common;
import rustc::middle::ast_map;

export mk_pass;

fn mk_pass() -> pass {
    {
        name: "reexport",
        f: run
    }
}

type def_set = map::set<ast::def_id>;
type def_map = map::hashmap<ast::def_id, doc::itemtag>;
type path_map = map::hashmap<str, [(str, doc::itemtag)]>;

fn run(srv: astsrv::srv, doc: doc::doc) -> doc::doc {

    // First gather the set of defs that are used as reexports
    let def_set = build_reexport_def_set(srv);

    // Now find the docs that go with those defs
    let def_map = build_reexport_def_map(srv, doc, def_set);

    // Now create a map that tells us where to insert the duplicated
    // docs into the existing doc tree
    let path_map = build_reexport_path_map(srv, def_map);

    // Finally update the doc tree
    merge_reexports(doc, path_map)
}

// Hash maps are not sendable so converting them back and forth
// to association lists. Yuck.
fn to_assoc_list<K:copy, V:copy>(
    map: map::hashmap<K, V>
) -> [(K, V)] {

    let vec = [];
    map.items {|k, v|
        vec += [(k, v)];
    }
    ret vec;
}

fn from_assoc_list<K:copy, V:copy>(
    list: [(K, V)],
    new_hash: fn() -> map::hashmap<K, V>
) -> map::hashmap<K, V> {

    let map = new_hash();
    vec::iter(list) {|elt|
        let (k, v) = elt;
        map.insert(k, v);
    }
    ret map;
}

fn from_def_assoc_list<V:copy>(
    list: [(ast::def_id, V)]
) -> map::hashmap<ast::def_id, V> {
    from_assoc_list(list, bind common::new_def_hash())
}

fn from_str_assoc_list<V:copy>(
    list: [(str, V)]
) -> map::hashmap<str, V> {
    from_assoc_list(list, bind map::str_hash())
}

fn build_reexport_def_set(srv: astsrv::srv) -> def_set {
    let assoc_list = astsrv::exec(srv) {|ctxt|
        let def_set = common::new_def_hash();
        ctxt.exp_map.items {|_id, defs|
            for def in defs {
                if def.reexp {
                    def_set.insert(def.id, ());
                }
            }
        }
        to_assoc_list(def_set)
    };

    from_def_assoc_list(assoc_list)
}

fn build_reexport_def_map(
    srv: astsrv::srv,
    doc: doc::doc,
    def_set: def_set
) -> def_map {

    type ctxt = {
        srv: astsrv::srv,
        def_set: def_set,
        def_map: def_map
    };

    let ctxt = {
        srv: srv,
        def_set: def_set,
        def_map: common::new_def_hash()
    };

    // FIXME: Do a parallel fold
    let fold = fold::fold({
        fold_mod: fold_mod,
        fold_nmod: fold_nmod
        with *fold::default_seq_fold(ctxt)
    });

    fold.fold_doc(fold, doc);

    ret ctxt.def_map;

    fn fold_mod(fold: fold::fold<ctxt>, doc: doc::moddoc) -> doc::moddoc {
        let doc = fold::default_seq_fold_mod(fold, doc);

        for item in doc.items {
            let def_id = ast_util::local_def(item.id());
            if fold.ctxt.def_set.contains_key(def_id) {
                fold.ctxt.def_map.insert(def_id, item);
            }
        }

        ret doc;
    }

    fn fold_nmod(fold: fold::fold<ctxt>, doc: doc::nmoddoc) -> doc::nmoddoc {
        let doc = fold::default_seq_fold_nmod(fold, doc);

        for fndoc in doc.fns {
            let def_id = ast_util::local_def(fndoc.id());
            if fold.ctxt.def_set.contains_key(def_id) {
                fold.ctxt.def_map.insert(def_id, doc::fntag(fndoc));
            }
        }

        ret doc;
    }
}

fn build_reexport_path_map(srv: astsrv::srv, -def_map: def_map) -> path_map {

    // This is real unfortunate. Lots of copying going on here
    let def_assoc_list = to_assoc_list(def_map);
    #debug("def_map: %?", def_assoc_list);

    let assoc_list = astsrv::exec(srv) {|ctxt|

        let def_map = from_def_assoc_list(def_assoc_list);
        let path_map = map::str_hash();

        ctxt.exp_map.items {|exp_id, defs|
            let path = alt check ctxt.ast_map.get(exp_id) {
              ast_map::node_export(_, path) { path }
            };
          // should be a constraint on the node_export constructor
          // that guarantees path is non-empty
            let name = alt check vec::last(*path) {
              ast_map::path_name(nm) { nm }
            };
            let modpath = ast_map::path_to_str(vec::init(*path));

            let reexportdocs = [];
            for def in defs {
                if !def.reexp { cont; }
                alt def_map.find(def.id) {
                  some(itemtag) {
                    reexportdocs += [(name, itemtag)];
                  }
                  _ {}
                }
            }

            if reexportdocs.len() > 0u {
                option::may(path_map.find(modpath)) {|docs|
                    reexportdocs = docs + vec::filter(reexportdocs, {|x|
                        !vec::contains(docs, x)
                    });
                }
                path_map.insert(modpath, reexportdocs);
                #debug("path_map entry: %? - %?",
                       modpath, (name, reexportdocs));
            }
        }

        to_assoc_list(path_map)
    };

    from_str_assoc_list(assoc_list)
}

fn merge_reexports(
    doc: doc::doc,
    path_map: path_map
) -> doc::doc {

    let fold = fold::fold({
        fold_mod: fold_mod
        with *fold::default_seq_fold(path_map)
    });

    ret fold.fold_doc(fold, doc);

    fn fold_mod(fold: fold::fold<path_map>, doc: doc::moddoc) -> doc::moddoc {
        let doc = fold::default_seq_fold_mod(fold, doc);

        let is_topmod = doc.id() == rustc::syntax::ast::crate_node_id;

        // In the case of the top mod, it really doesn't have a name;
        // the name we have here is actually the crate name
        let path = if is_topmod {
            doc.path()
        } else {
            doc.path() + [doc.name()]
        };

        let new_items = get_new_items(path, fold.ctxt);
        #debug("merging into %?: %?", path, new_items);

        {
            items: (doc.items + new_items)
            with doc
        }
    }

    fn get_new_items(path: [str], path_map: path_map) -> [doc::itemtag] {
        #debug("looking for reexports in path %?", path);
        alt path_map.find(str::connect(path, "::")) {
          some(name_docs) {
            vec::foldl([], name_docs) {|v, name_doc|
                let (name, doc) = name_doc;
                v + [reexport_doc(doc, name)]
            }
          }
          none { [] }
        }
    }

    fn reexport_doc(doc: doc::itemtag, name: str) -> doc::itemtag {
        alt doc {
          doc::modtag(doc @ {item, _}) {
            doc::modtag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::nmodtag(_) { fail }
          doc::consttag(doc @ {item, _}) {
            doc::consttag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::fntag(doc @ {item, _}) {
            doc::fntag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::enumtag(doc @ {item, _}) {
            doc::enumtag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::restag(doc @ {item, _}) {
            doc::restag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::ifacetag(doc @ {item, _}) {
            doc::ifacetag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::impltag(doc @ {item, _}) {
            doc::impltag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::tytag(doc @ {item, _}) {
            doc::tytag({
                item: reexport(item, name)
                with doc
            })
          }
        }
    }

    fn reexport(doc: doc::itemdoc, name: str) -> doc::itemdoc {
        {
            name: name,
            reexport: true
            with doc
        }
    }
}

#[test]
fn should_duplicate_reexported_items() {
    let source = "mod a { export b; fn b() { } } \
                  mod c { import a::b; export b; }";
    let doc = test::mk_doc(source);
    assert doc.cratemod().mods()[1].fns()[0].name() == "b";
}

#[test]
fn should_mark_reepxorts_as_such() {
    let source = "mod a { export b; fn b() { } } \
                  mod c { import a::b; export b; }";
    let doc = test::mk_doc(source);
    assert doc.cratemod().mods()[1].fns()[0].item.reexport == true;
}

#[test]
fn should_duplicate_reexported_native_fns() {
    let source = "native mod a { fn b(); } \
                  mod c { import a::b; export b; }";
    let doc = test::mk_doc(source);
    assert doc.cratemod().mods()[0].fns()[0].name() == "b";
}

#[test]
fn should_duplicate_multiple_reexported_items() {
    let source = "mod a { \
                  export b; export c; \
                  fn b() { } fn c() { } \
                  } \
                  mod d { \
                  import a::b; import a::c; \
                  export b; export c; \
                  }";
    astsrv::from_str(source) {|srv|
        let doc = extract::from_srv(srv, "");
        let doc = path_pass::mk_pass().f(srv, doc);
        let doc = run(srv, doc);
        // Reexports may not be in any specific order
        let doc = sort_item_name_pass::mk_pass().f(srv, doc);
        assert doc.cratemod().mods()[1].fns()[0].name() == "b";
        assert doc.cratemod().mods()[1].fns()[1].name() == "c";
    }
}

#[test]
fn should_rename_items_reexported_with_different_names() {
    let source = "mod a { export b; fn b() { } } \
                  mod c { import x = a::b; export x; }";
    let doc = test::mk_doc(source);
    assert doc.cratemod().mods()[1].fns()[0].name() == "x";
}

#[test]
fn should_reexport_in_topmod() {
    fn mk_doc(source: str) -> doc::doc {
        astsrv::from_str(source) {|srv|
            let doc = extract::from_srv(srv, "core");
            let doc = path_pass::mk_pass().f(srv, doc);
            run(srv, doc)
        }
    }
    let source = "import option::{some, none}; \
                  import option = option::t; \
                  export option, some, none; \
                  mod option { \
                  enum t { some, none } \
                  }";
    let doc = mk_doc(source);
    assert doc.cratemod().enums()[0].name() == "option";
}

#[test]
fn should_not_reexport_multiple_times() {
    let source = "import option = option::t; \
                  export option; \
                  export option; \
                  mod option { \
                  enum t { none, some } \
                  }";
    let doc = test::mk_doc(source);
    assert vec::len(doc.cratemod().enums()) == 1u;
}

#[cfg(test)]
mod test {
    fn mk_doc(source: str) -> doc::doc {
        astsrv::from_str(source) {|srv|
            let doc = extract::from_srv(srv, "");
            let doc = path_pass::mk_pass().f(srv, doc);
            run(srv, doc)
        }
    }
}
<commit_msg>rustdoc: Document impl reexports<commit_after>#[doc = "Finds docs for reexported items and duplicates them"];

import std::map;
import std::map::hashmap;
import std::list;
import rustc::syntax::ast;
import rustc::syntax::ast_util;
import rustc::util::common;
import rustc::middle::ast_map;
import rustc::syntax::visit;
import rustc::syntax::codemap;

export mk_pass;

fn mk_pass() -> pass {
    {
        name: "reexport",
        f: run
    }
}

type def_set = map::set<ast::def_id>;
type def_map = map::hashmap<ast::def_id, doc::itemtag>;
type path_map = map::hashmap<str, [(str, doc::itemtag)]>;

fn run(srv: astsrv::srv, doc: doc::doc) -> doc::doc {

    // First gather the set of defs that are used as reexports
    let def_set = build_reexport_def_set(srv);

    // Now find the docs that go with those defs
    let def_map = build_reexport_def_map(srv, doc, def_set);

    // Now create a map that tells us where to insert the duplicated
    // docs into the existing doc tree
    let path_map = build_reexport_path_map(srv, def_map);

    // Finally update the doc tree
    merge_reexports(doc, path_map)
}

// Hash maps are not sendable so converting them back and forth
// to association lists. Yuck.
fn to_assoc_list<K:copy, V:copy>(
    map: map::hashmap<K, V>
) -> [(K, V)] {

    let vec = [];
    map.items {|k, v|
        vec += [(k, v)];
    }
    ret vec;
}

fn from_assoc_list<K:copy, V:copy>(
    list: [(K, V)],
    new_hash: fn() -> map::hashmap<K, V>
) -> map::hashmap<K, V> {

    let map = new_hash();
    vec::iter(list) {|elt|
        let (k, v) = elt;
        map.insert(k, v);
    }
    ret map;
}

fn from_def_assoc_list<V:copy>(
    list: [(ast::def_id, V)]
) -> map::hashmap<ast::def_id, V> {
    from_assoc_list(list, bind common::new_def_hash())
}

fn from_str_assoc_list<V:copy>(
    list: [(str, V)]
) -> map::hashmap<str, V> {
    from_assoc_list(list, bind map::str_hash())
}

fn build_reexport_def_set(srv: astsrv::srv) -> def_set {
    let assoc_list = astsrv::exec(srv) {|ctxt|
        let def_set = common::new_def_hash();
        ctxt.exp_map.items {|_id, defs|
            for def in defs {
                if def.reexp {
                    def_set.insert(def.id, ());
                }
            }
        }
        for def in find_reexport_impls(ctxt) {
            def_set.insert(def, ());
        }
        to_assoc_list(def_set)
    };

    from_def_assoc_list(assoc_list)
}

fn find_reexport_impls(ctxt: astsrv::ctxt) -> [ast::def_id] {
    let defs = @mut [];
    let visitor = @{
        visit_mod: bind visit_mod(ctxt, defs, _, _, _)
        with *visit::default_simple_visitor()
    };
    let visitor = visit::mk_simple_visitor(visitor);
    visit::visit_crate(*ctxt.ast, (), visitor);
    ret *defs;

    fn visit_mod(
        ctxt: astsrv::ctxt,
        defs: @mut [ast::def_id],
        m: ast::_mod,
        _sp: codemap::span,
        mod_id: ast::node_id
    ) {
        let all_impls = all_impls(m);
        alt check ctxt.impl_map.get(mod_id) {
          list::cons(impls, @list::nil) {
            for i in *impls {
                // This impl is not an item in the current mod
                if !all_impls.contains_key(i.did) {
                    // Ignore external impls because I don't
                    // know what to do with them yet
                    if i.did.crate == ast::local_crate {
                        *defs += [i.did]
                    }
                }
            }
          }
        }
    }
}

fn all_impls(m: ast::_mod) -> map::set<ast::def_id> {
    let all_impls = common::new_def_hash();
    for item in m.items {
        alt item.node {
          ast::item_impl(_, _, _, _) {
            all_impls.insert(ast_util::local_def(item.id), ());
          }
          _ { }
        }
    }
    ret all_impls;
}

fn build_reexport_def_map(
    srv: astsrv::srv,
    doc: doc::doc,
    def_set: def_set
) -> def_map {

    type ctxt = {
        srv: astsrv::srv,
        def_set: def_set,
        def_map: def_map
    };

    let ctxt = {
        srv: srv,
        def_set: def_set,
        def_map: common::new_def_hash()
    };

    // FIXME: Do a parallel fold
    let fold = fold::fold({
        fold_mod: fold_mod,
        fold_nmod: fold_nmod
        with *fold::default_seq_fold(ctxt)
    });

    fold.fold_doc(fold, doc);

    ret ctxt.def_map;

    fn fold_mod(fold: fold::fold<ctxt>, doc: doc::moddoc) -> doc::moddoc {
        let doc = fold::default_seq_fold_mod(fold, doc);

        for item in doc.items {
            let def_id = ast_util::local_def(item.id());
            if fold.ctxt.def_set.contains_key(def_id) {
                fold.ctxt.def_map.insert(def_id, item);
            }
        }

        ret doc;
    }

    fn fold_nmod(fold: fold::fold<ctxt>, doc: doc::nmoddoc) -> doc::nmoddoc {
        let doc = fold::default_seq_fold_nmod(fold, doc);

        for fndoc in doc.fns {
            let def_id = ast_util::local_def(fndoc.id());
            if fold.ctxt.def_set.contains_key(def_id) {
                fold.ctxt.def_map.insert(def_id, doc::fntag(fndoc));
            }
        }

        ret doc;
    }
}

fn build_reexport_path_map(srv: astsrv::srv, -def_map: def_map) -> path_map {

    // This is real unfortunate. Lots of copying going on here
    let def_assoc_list = to_assoc_list(def_map);
    #debug("def_map: %?", def_assoc_list);

    let assoc_list = astsrv::exec(srv) {|ctxt|

        let def_map = from_def_assoc_list(def_assoc_list);
        let path_map = map::str_hash();

        ctxt.exp_map.items {|exp_id, defs|
            let path = alt check ctxt.ast_map.get(exp_id) {
              ast_map::node_export(_, path) { path }
            };
          // should be a constraint on the node_export constructor
          // that guarantees path is non-empty
            let name = alt check vec::last(*path) {
              ast_map::path_name(nm) { nm }
            };
            let modpath = ast_map::path_to_str(vec::init(*path));

            let reexportdocs = [];
            for def in defs {
                if !def.reexp { cont; }
                alt def_map.find(def.id) {
                  some(itemtag) {
                    reexportdocs += [(name, itemtag)];
                  }
                  _ {}
                }
            }

            if reexportdocs.len() > 0u {
                option::may(path_map.find(modpath)) {|docs|
                    reexportdocs = docs + vec::filter(reexportdocs, {|x|
                        !vec::contains(docs, x)
                    });
                }
                path_map.insert(modpath, reexportdocs);
                #debug("path_map entry: %? - %?",
                       modpath, (name, reexportdocs));
            }
        }

        for (path, doc) in find_reexport_impl_docs(ctxt, def_map) {
            let docs = alt path_map.find(path) {
              some(docs) { docs + [(doc)] }
              none { [doc] }
            };
            path_map.insert(path, docs);
        }

        to_assoc_list(path_map)
    };

    from_str_assoc_list(assoc_list)
}

fn find_reexport_impl_docs(
    ctxt: astsrv::ctxt,
    def_map: def_map
) -> [(str, (str, doc::itemtag))] {
    let docs = @mut [];
    let visitor = @{
        visit_mod: bind visit_mod(ctxt, def_map, docs, _, _, _)
        with *visit::default_simple_visitor()
    };
    let visitor = visit::mk_simple_visitor(visitor);
    visit::visit_crate(*ctxt.ast, (), visitor);
    ret *docs;

    fn visit_mod(
        ctxt: astsrv::ctxt,
        def_map: def_map,
        docs: @mut [(str, (str, doc::itemtag))],
        m: ast::_mod,
        _sp: codemap::span,
        mod_id: ast::node_id
    ) {
        let all_impls = all_impls(m);
        alt check ctxt.impl_map.get(mod_id) {
          list::cons(impls, @list::nil) {
            for i in *impls {
                // This impl is not an item in the current mod
                if !all_impls.contains_key(i.did) {
                    // Ignore external impls because I don't
                    // know what to do with them yet
                    if i.did.crate == ast::local_crate {
                        let path = alt ctxt.ast_map.find(mod_id) {
                          some(ast_map::node_item(item, path)) {
                            let path = ast_map::path_to_str(*path);
                            if str::is_empty(path) {
                                item.ident
                            } else {
                                path + "::" + item.ident
                            }
                          }
                          _ {
                            assert mod_id == ast::crate_node_id;
                            ""
                          }
                        };
                        let ident = i.ident;
                        let doc = alt check def_map.find(i.did) {
                          some(doc) { doc }
                        };
                        *docs += [(path, (ident, doc))];
                    }
                }
            }
          }
        }
    }
}

fn merge_reexports(
    doc: doc::doc,
    path_map: path_map
) -> doc::doc {

    let fold = fold::fold({
        fold_mod: fold_mod
        with *fold::default_seq_fold(path_map)
    });

    ret fold.fold_doc(fold, doc);

    fn fold_mod(fold: fold::fold<path_map>, doc: doc::moddoc) -> doc::moddoc {
        let doc = fold::default_seq_fold_mod(fold, doc);

        let is_topmod = doc.id() == rustc::syntax::ast::crate_node_id;

        // In the case of the top mod, it really doesn't have a name;
        // the name we have here is actually the crate name
        let path = if is_topmod {
            doc.path()
        } else {
            doc.path() + [doc.name()]
        };

        let new_items = get_new_items(path, fold.ctxt);
        #debug("merging into %?: %?", path, new_items);

        {
            items: (doc.items + new_items)
            with doc
        }
    }

    fn get_new_items(path: [str], path_map: path_map) -> [doc::itemtag] {
        #debug("looking for reexports in path %?", path);
        alt path_map.find(str::connect(path, "::")) {
          some(name_docs) {
            vec::foldl([], name_docs) {|v, name_doc|
                let (name, doc) = name_doc;
                v + [reexport_doc(doc, name)]
            }
          }
          none { [] }
        }
    }

    fn reexport_doc(doc: doc::itemtag, name: str) -> doc::itemtag {
        alt doc {
          doc::modtag(doc @ {item, _}) {
            doc::modtag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::nmodtag(_) { fail }
          doc::consttag(doc @ {item, _}) {
            doc::consttag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::fntag(doc @ {item, _}) {
            doc::fntag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::enumtag(doc @ {item, _}) {
            doc::enumtag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::restag(doc @ {item, _}) {
            doc::restag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::ifacetag(doc @ {item, _}) {
            doc::ifacetag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::impltag(doc @ {item, _}) {
            doc::impltag({
                item: reexport(item, name)
                with doc
            })
          }
          doc::tytag(doc @ {item, _}) {
            doc::tytag({
                item: reexport(item, name)
                with doc
            })
          }
        }
    }

    fn reexport(doc: doc::itemdoc, name: str) -> doc::itemdoc {
        {
            name: name,
            reexport: true
            with doc
        }
    }
}

#[test]
fn should_duplicate_reexported_items() {
    let source = "mod a { export b; fn b() { } } \
                  mod c { import a::b; export b; }";
    let doc = test::mk_doc(source);
    assert doc.cratemod().mods()[1].fns()[0].name() == "b";
}

#[test]
fn should_mark_reepxorts_as_such() {
    let source = "mod a { export b; fn b() { } } \
                  mod c { import a::b; export b; }";
    let doc = test::mk_doc(source);
    assert doc.cratemod().mods()[1].fns()[0].item.reexport == true;
}

#[test]
fn should_duplicate_reexported_impls() {
    let source = "mod a { impl b for int { fn c() { } } } \
                  mod d { import a::b; export b; }";
    let doc = test::mk_doc(source);
    assert doc.cratemod().mods()[1].impls()[0].name() == "b";
}

#[test]
fn should_duplicate_reexported_impls_deep() {
    let source = "mod a { impl b for int { fn c() { } } } \
                  mod d { mod e { import a::b; export b; } }";
    let doc = test::mk_doc(source);
    assert doc.cratemod().mods()[1].mods()[0].impls()[0].name() == "b";
}

#[test]
fn should_duplicate_reexported_impls_crate() {
    let source = "import a::b; export b; \
                  mod a { impl b for int { fn c() { } } }";
    let doc = test::mk_doc(source);
    assert doc.cratemod().impls()[0].name() == "b";
}

#[test]
fn should_duplicate_reexported_native_fns() {
    let source = "native mod a { fn b(); } \
                  mod c { import a::b; export b; }";
    let doc = test::mk_doc(source);
    assert doc.cratemod().mods()[0].fns()[0].name() == "b";
}

#[test]
fn should_duplicate_multiple_reexported_items() {
    let source = "mod a { \
                  export b; export c; \
                  fn b() { } fn c() { } \
                  } \
                  mod d { \
                  import a::b; import a::c; \
                  export b; export c; \
                  }";
    astsrv::from_str(source) {|srv|
        let doc = extract::from_srv(srv, "");
        let doc = path_pass::mk_pass().f(srv, doc);
        let doc = run(srv, doc);
        // Reexports may not be in any specific order
        let doc = sort_item_name_pass::mk_pass().f(srv, doc);
        assert doc.cratemod().mods()[1].fns()[0].name() == "b";
        assert doc.cratemod().mods()[1].fns()[1].name() == "c";
    }
}

#[test]
fn should_rename_items_reexported_with_different_names() {
    let source = "mod a { export b; fn b() { } } \
                  mod c { import x = a::b; export x; }";
    let doc = test::mk_doc(source);
    assert doc.cratemod().mods()[1].fns()[0].name() == "x";
}

#[test]
fn should_reexport_in_topmod() {
    fn mk_doc(source: str) -> doc::doc {
        astsrv::from_str(source) {|srv|
            let doc = extract::from_srv(srv, "core");
            let doc = path_pass::mk_pass().f(srv, doc);
            run(srv, doc)
        }
    }
    let source = "import option::{some, none}; \
                  import option = option::t; \
                  export option, some, none; \
                  mod option { \
                  enum t { some, none } \
                  }";
    let doc = mk_doc(source);
    assert doc.cratemod().enums()[0].name() == "option";
}

#[test]
fn should_not_reexport_multiple_times() {
    let source = "import option = option::t; \
                  export option; \
                  export option; \
                  mod option { \
                  enum t { none, some } \
                  }";
    let doc = test::mk_doc(source);
    assert vec::len(doc.cratemod().enums()) == 1u;
}

#[cfg(test)]
mod test {
    fn mk_doc(source: str) -> doc::doc {
        astsrv::from_str(source) {|srv|
            let doc = extract::from_srv(srv, "");
            let doc = path_pass::mk_pass().f(srv, doc);
            run(srv, doc)
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     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,
// See the License for the specific language governing permissions and
// limitations under the License.

#![allow(dead_code)]

use server::Node;
use server::transport::{ServerRaftStoreRouter, RaftStoreRouter};
use raftstore::store::{Transport, Peekable};
use raftstore::errors::Error as RaftServerError;
use raftstore::coprocessor::RegionSnapshot;
use util::HandyRwLock;
use kvproto::raft_cmdpb::{RaftCmdRequest, RaftCmdResponse, RaftRequestHeader, Request, Response,
                          GetRequest, CmdType, SeekRequest, DeleteRequest, PutRequest};
use kvproto::errorpb;
use kvproto::kvrpcpb::Context;

use pd::PdClient;
use uuid::Uuid;
use std::sync::{Arc, RwLock, Mutex, Condvar};
use std::fmt::{self, Formatter, Debug};
use std::io::Error as IoError;
use std::time::Duration;
use std::result;
use rocksdb::DB;
use protobuf::RepeatedField;

use storage::engine;
use super::{Engine, Modify, Snapshot};
use storage::{Key, Value, KvPair};

const DEFAULT_TIMEOUT_SECS: u64 = 5;

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        RequestFailed(e: errorpb::Error) {
            from()
            description(e.get_message())
        }
        Io(e: IoError) {
            from()
            cause(e)
            description(e.description())
        }
        RocksDb(reason: String) {
            description(reason)
        }
        Server(e: RaftServerError) {
            from()
            cause(e)
            description(e.description())
        }
        InvalidResponse(reason: String) {
            description(reason)
        }
        InvalidRequest(reason: String) {
            description(reason)
        }
        Timeout(d: Duration) {
            description("request timeout")
            display("timeout after {:?}", d)
        }
    }
}

pub type Result<T> = result::Result<T, Error>;

impl From<Error> for engine::Error {
    fn from(e: Error) -> engine::Error {
        match e {
            Error::RequestFailed(e) => engine::Error::Request(e),
            e => box_err!(e),
        }
    }
}

/// RaftKv is a storage engine base on RaftKvServer.
pub struct RaftKv<T: PdClient + 'static, Trans: Transport + 'static> {
    node: Node<T, Trans>,
    db: Arc<DB>,
    router: Arc<RwLock<ServerRaftStoreRouter>>,
}

impl<T: PdClient, Trans: Transport> RaftKv<T, Trans> {
    /// Create a RaftKv using specified configuration.
    pub fn new(node: Node<T, Trans>, db: Arc<DB>) -> RaftKv<T, Trans> {
        let router = node.raft_store_router();
        RaftKv {
            node: node,
            db: db,
            router: router,
        }
    }

    // TODO: reuse msg code.
    pub fn call_command(&self, request: RaftCmdRequest) -> Result<RaftCmdResponse> {
        let resp: Option<RaftCmdResponse> = None;
        let pair = Arc::new((Mutex::new(resp), Condvar::new()));
        let pair2 = pair.clone();
        let timeout = Duration::from_secs(DEFAULT_TIMEOUT_SECS);

        try!(self.router.rl().send_command(request,
                                           box move |resp| {
                                               let &(ref lock, ref cvar) = &*pair2;
                                               let mut v = lock.lock().unwrap();
                                               *v = Some(resp);
                                               cvar.notify_one();
                                               Ok(())
                                           }));

        let &(ref lock, ref cvar) = &*pair;
        let mut v = lock.lock().unwrap();
        while v.is_none() {
            let (resp, timeout_res) = cvar.wait_timeout(v, timeout).unwrap();
            if timeout_res.timed_out() {
                return Err(Error::Timeout(timeout));
            }

            v = resp
        }

        Ok(v.take().unwrap())
    }

    fn exec_cmd_request(&self, req: RaftCmdRequest) -> Result<RaftCmdResponse> {
        let uuid = req.get_header().get_uuid().to_vec();
        let l = req.get_requests().len();

        // Only when tx is closed will recv return Err, which should never happen.
        let mut resp = try!(self.call_command(req));
        if resp.get_header().get_uuid() != &*uuid {
            return Err(Error::InvalidResponse("response is not correct!!!".to_owned()));
        }
        if resp.get_header().has_error() {
            return Err(Error::RequestFailed(resp.mut_header().take_error()));
        }
        if l != resp.get_responses().len() {
            return Err(Error::InvalidResponse("response count is not equal to requests, \
                                               something must go wrong."
                                                  .to_owned()));
        }
        Ok(resp)
    }

    fn new_request_header(&self, ctx: &Context) -> RaftRequestHeader {
        let mut header = RaftRequestHeader::new();
        header.set_region_id(ctx.get_region_id());
        header.set_peer(ctx.get_peer().clone());
        header.set_region_epoch(ctx.get_region_epoch().clone());
        header.set_uuid(Uuid::new_v4().as_bytes().to_vec());
        header
    }

    fn exec_requests(&self, ctx: &Context, reqs: Vec<Request>) -> Result<Vec<Response>> {
        let header = self.new_request_header(ctx);
        let mut cmd = RaftCmdRequest::new();
        cmd.set_header(header);
        cmd.set_requests(RepeatedField::from_vec(reqs));
        let mut resp = try!(self.exec_cmd_request(cmd));
        Ok(resp.take_responses().to_vec())
    }

    fn exec_request(&self, ctx: &Context, req: Request) -> Result<Response> {
        let resps = self.exec_requests(ctx, vec![req]);
        resps.map(|mut v| v.remove(0))
    }
}

fn invalid_resp_type(exp: CmdType, act: CmdType) -> engine::Error {
    Error::InvalidResponse(format!("cmd type not match, want {:?}, got {:?}!", exp, act)).into()
}

impl<T: PdClient, Trans: Transport> Debug for RaftKv<T, Trans> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "RaftKv")
    }
}

impl<T: PdClient, Trans: Transport> Engine for RaftKv<T, Trans> {
    fn get(&self, ctx: &Context, key: &Key) -> engine::Result<Option<Value>> {
        let mut get = GetRequest::new();
        get.set_key(key.raw().clone());
        let mut req = Request::new();
        req.set_cmd_type(CmdType::Get);
        req.set_get(get);
        let mut resp = try!(self.exec_request(ctx, req));
        if resp.get_cmd_type() != CmdType::Get {
            return Err(invalid_resp_type(CmdType::Get, resp.get_cmd_type()));
        }
        let get_resp = resp.mut_get();
        if get_resp.has_value() {
            Ok(Some(get_resp.take_value()))
        } else {
            Ok(None)
        }
    }

    fn seek(&self, ctx: &Context, key: &Key) -> engine::Result<Option<KvPair>> {
        let mut seek = SeekRequest::new();
        seek.set_key(key.raw().clone());
        let mut req = Request::new();
        req.set_cmd_type(CmdType::Seek);
        req.set_seek(seek);
        let mut resp = try!(self.exec_request(ctx, req));
        if resp.get_cmd_type() != CmdType::Seek {
            return Err(invalid_resp_type(CmdType::Seek, resp.get_cmd_type()));
        }
        let seek_resp = resp.mut_seek();
        if seek_resp.has_key() {
            Ok(Some((seek_resp.take_key(), seek_resp.take_value())))
        } else {
            Ok(None)
        }
    }

    fn reverse_seek(&self, _: &Context, _: &Key) -> engine::Result<Option<KvPair>> {
        unimplemented!();
    }

    fn write(&self, ctx: &Context, mut modifies: Vec<Modify>) -> engine::Result<()> {
        if modifies.len() == 0 {
            return Ok(());
        }
        let mut reqs = Vec::with_capacity(modifies.len());
        while !modifies.is_empty() {
            let m = modifies.pop().unwrap();
            let mut req = Request::new();
            match m {
                Modify::Delete(k) => {
                    let mut delete = DeleteRequest::new();
                    delete.set_key(k.raw().clone());
                    req.set_cmd_type(CmdType::Delete);
                    req.set_delete(delete);
                }
                Modify::Put((k, v)) => {
                    let mut put = PutRequest::new();
                    put.set_key(k.raw().clone());
                    put.set_value(v);
                    req.set_cmd_type(CmdType::Put);
                    req.set_put(put);
                }
            }
            reqs.push(req);
        }
        try!(self.exec_requests(ctx, reqs));
        Ok(())
    }

    fn snapshot<'a>(&'a self, ctx: &Context) -> engine::Result<Box<Snapshot + 'a>> {
        let mut req = Request::new();
        req.set_cmd_type(CmdType::Snap);
        let mut resp = try!(self.exec_request(ctx, req));
        if resp.get_cmd_type() != CmdType::Snap {
            return Err(invalid_resp_type(CmdType::Snap, resp.get_cmd_type()));
        }
        let region = resp.take_snap().take_region();
        // TODO figure out a way to create snapshot when apply
        let snap = RegionSnapshot::from_raw(self.db.as_ref(), region);
        Ok(box snap)
    }
}

impl<'a> Snapshot for RegionSnapshot<'a> {
    fn get(&self, key: &Key) -> engine::Result<Option<Value>> {
        let v = box_try!(self.get_value(key.raw()));
        Ok(v.map(|v| v.to_vec()))
    }

    fn seek(&self, key: &Key) -> engine::Result<Option<KvPair>> {
        let pair = box_try!(self.seek(key.raw()));
        Ok(pair)
    }

    fn reverse_seek(&self, _: &Key) -> engine::Result<Option<KvPair>> {
        unimplemented!();
    }
}
<commit_msg>raftkv: tiny clean up. (#481)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     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,
// See the License for the specific language governing permissions and
// limitations under the License.

#![allow(dead_code)]

use server::Node;
use server::transport::{ServerRaftStoreRouter, RaftStoreRouter};
use raftstore::store::{Transport, Peekable};
use raftstore::errors::Error as RaftServerError;
use raftstore::coprocessor::RegionSnapshot;
use util::HandyRwLock;
use kvproto::raft_cmdpb::{RaftCmdRequest, RaftCmdResponse, RaftRequestHeader, Request, Response,
                          GetRequest, CmdType, SeekRequest, DeleteRequest, PutRequest};
use kvproto::errorpb;
use kvproto::kvrpcpb::Context;

use pd::PdClient;
use uuid::Uuid;
use std::sync::{Arc, RwLock, Mutex, Condvar};
use std::fmt::{self, Formatter, Debug};
use std::io::Error as IoError;
use std::time::Duration;
use std::result;
use rocksdb::DB;
use protobuf::RepeatedField;

use storage::engine;
use super::{Engine, Modify, Snapshot};
use storage::{Key, Value, KvPair};

const DEFAULT_TIMEOUT_SECS: u64 = 5;

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        RequestFailed(e: errorpb::Error) {
            from()
            description(e.get_message())
        }
        Io(e: IoError) {
            from()
            cause(e)
            description(e.description())
        }
        RocksDb(reason: String) {
            description(reason)
        }
        Server(e: RaftServerError) {
            from()
            cause(e)
            description(e.description())
        }
        InvalidResponse(reason: String) {
            description(reason)
        }
        InvalidRequest(reason: String) {
            description(reason)
        }
        Timeout(d: Duration) {
            description("request timeout")
            display("timeout after {:?}", d)
        }
    }
}

pub type Result<T> = result::Result<T, Error>;

impl From<Error> for engine::Error {
    fn from(e: Error) -> engine::Error {
        match e {
            Error::RequestFailed(e) => engine::Error::Request(e),
            e => box_err!(e),
        }
    }
}

/// RaftKv is a storage engine base on RaftKvServer.
pub struct RaftKv<T: PdClient + 'static, Trans: Transport + 'static> {
    node: Node<T, Trans>,
    db: Arc<DB>,
    router: Arc<RwLock<ServerRaftStoreRouter>>,
}

impl<T: PdClient, Trans: Transport> RaftKv<T, Trans> {
    /// Create a RaftKv using specified configuration.
    pub fn new(node: Node<T, Trans>, db: Arc<DB>) -> RaftKv<T, Trans> {
        let router = node.raft_store_router();
        RaftKv {
            node: node,
            db: db,
            router: router,
        }
    }

    // TODO: reuse msg code.
    pub fn call_command(&self, request: RaftCmdRequest) -> Result<RaftCmdResponse> {
        let resp = None;
        let pair = Arc::new((Mutex::new(resp), Condvar::new()));
        let pair2 = pair.clone();
        let timeout = Duration::from_secs(DEFAULT_TIMEOUT_SECS);

        try!(self.router.rl().send_command(request,
                                           box move |resp| {
                                               let &(ref lock, ref cvar) = &*pair2;
                                               let mut v = lock.lock().unwrap();
                                               *v = Some(resp);
                                               cvar.notify_one();
                                               Ok(())
                                           }));

        let &(ref lock, ref cvar) = &*pair;
        let mut v = lock.lock().unwrap();
        while v.is_none() {
            let (resp, timeout_res) = cvar.wait_timeout(v, timeout).unwrap();
            if timeout_res.timed_out() {
                return Err(Error::Timeout(timeout));
            }

            v = resp
        }

        Ok(v.take().unwrap())
    }

    fn exec_cmd_request(&self, req: RaftCmdRequest) -> Result<RaftCmdResponse> {
        let uuid = req.get_header().get_uuid().to_vec();
        let l = req.get_requests().len();

        // Only when tx is closed will recv return Err, which should never happen.
        let mut resp = try!(self.call_command(req));
        if resp.get_header().get_uuid() != &*uuid {
            return Err(Error::InvalidResponse("response is not correct!!!".to_owned()));
        }
        if resp.get_header().has_error() {
            return Err(Error::RequestFailed(resp.mut_header().take_error()));
        }
        if l != resp.get_responses().len() {
            return Err(Error::InvalidResponse("response count is not equal to requests, \
                                               something must go wrong."
                                                  .to_owned()));
        }
        Ok(resp)
    }

    fn new_request_header(&self, ctx: &Context) -> RaftRequestHeader {
        let mut header = RaftRequestHeader::new();
        header.set_region_id(ctx.get_region_id());
        header.set_peer(ctx.get_peer().clone());
        header.set_region_epoch(ctx.get_region_epoch().clone());
        header.set_uuid(Uuid::new_v4().as_bytes().to_vec());
        header
    }

    fn exec_requests(&self, ctx: &Context, reqs: Vec<Request>) -> Result<Vec<Response>> {
        let header = self.new_request_header(ctx);
        let mut cmd = RaftCmdRequest::new();
        cmd.set_header(header);
        cmd.set_requests(RepeatedField::from_vec(reqs));
        let mut resp = try!(self.exec_cmd_request(cmd));
        Ok(resp.take_responses().to_vec())
    }

    fn exec_request(&self, ctx: &Context, req: Request) -> Result<Response> {
        let resps = self.exec_requests(ctx, vec![req]);
        resps.map(|mut v| v.remove(0))
    }
}

fn invalid_resp_type(exp: CmdType, act: CmdType) -> engine::Error {
    Error::InvalidResponse(format!("cmd type not match, want {:?}, got {:?}!", exp, act)).into()
}

impl<T: PdClient, Trans: Transport> Debug for RaftKv<T, Trans> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "RaftKv")
    }
}

impl<T: PdClient, Trans: Transport> Engine for RaftKv<T, Trans> {
    fn get(&self, ctx: &Context, key: &Key) -> engine::Result<Option<Value>> {
        let mut get = GetRequest::new();
        get.set_key(key.raw().clone());
        let mut req = Request::new();
        req.set_cmd_type(CmdType::Get);
        req.set_get(get);
        let mut resp = try!(self.exec_request(ctx, req));
        if resp.get_cmd_type() != CmdType::Get {
            return Err(invalid_resp_type(CmdType::Get, resp.get_cmd_type()));
        }
        let get_resp = resp.mut_get();
        if get_resp.has_value() {
            Ok(Some(get_resp.take_value()))
        } else {
            Ok(None)
        }
    }

    fn seek(&self, ctx: &Context, key: &Key) -> engine::Result<Option<KvPair>> {
        let mut seek = SeekRequest::new();
        seek.set_key(key.raw().clone());
        let mut req = Request::new();
        req.set_cmd_type(CmdType::Seek);
        req.set_seek(seek);
        let mut resp = try!(self.exec_request(ctx, req));
        if resp.get_cmd_type() != CmdType::Seek {
            return Err(invalid_resp_type(CmdType::Seek, resp.get_cmd_type()));
        }
        let seek_resp = resp.mut_seek();
        if seek_resp.has_key() {
            Ok(Some((seek_resp.take_key(), seek_resp.take_value())))
        } else {
            Ok(None)
        }
    }

    fn reverse_seek(&self, _: &Context, _: &Key) -> engine::Result<Option<KvPair>> {
        unimplemented!();
    }

    fn write(&self, ctx: &Context, mut modifies: Vec<Modify>) -> engine::Result<()> {
        if modifies.len() == 0 {
            return Ok(());
        }
        let mut reqs = Vec::with_capacity(modifies.len());
        while !modifies.is_empty() {
            let m = modifies.pop().unwrap();
            let mut req = Request::new();
            match m {
                Modify::Delete(k) => {
                    let mut delete = DeleteRequest::new();
                    delete.set_key(k.raw().clone());
                    req.set_cmd_type(CmdType::Delete);
                    req.set_delete(delete);
                }
                Modify::Put((k, v)) => {
                    let mut put = PutRequest::new();
                    put.set_key(k.raw().clone());
                    put.set_value(v);
                    req.set_cmd_type(CmdType::Put);
                    req.set_put(put);
                }
            }
            reqs.push(req);
        }
        try!(self.exec_requests(ctx, reqs));
        Ok(())
    }

    fn snapshot<'a>(&'a self, ctx: &Context) -> engine::Result<Box<Snapshot + 'a>> {
        let mut req = Request::new();
        req.set_cmd_type(CmdType::Snap);
        let mut resp = try!(self.exec_request(ctx, req));
        if resp.get_cmd_type() != CmdType::Snap {
            return Err(invalid_resp_type(CmdType::Snap, resp.get_cmd_type()));
        }
        let region = resp.take_snap().take_region();
        // TODO figure out a way to create snapshot when apply
        let snap = RegionSnapshot::from_raw(self.db.as_ref(), region);
        Ok(box snap)
    }
}

impl<'a> Snapshot for RegionSnapshot<'a> {
    fn get(&self, key: &Key) -> engine::Result<Option<Value>> {
        let v = box_try!(self.get_value(key.raw()));
        Ok(v.map(|v| v.to_vec()))
    }

    fn seek(&self, key: &Key) -> engine::Result<Option<KvPair>> {
        let pair = box_try!(self.seek(key.raw()));
        Ok(pair)
    }

    fn reverse_seek(&self, _: &Key) -> engine::Result<Option<KvPair>> {
        unimplemented!();
    }
}
<|endoftext|>


<commit_before>use futures_core::task::{Context, Poll};
#[cfg(feature = "read-initializer")]
use futures_io::Initializer;
use futures_io::{AsyncBufRead, AsyncRead, IoSliceMut};
use pin_utils::{unsafe_pinned, unsafe_unpinned};
use std::fmt;
use std::io;
use std::pin::Pin;

/// Reader for the [`chain`](super::AsyncReadExt::chain) method.
#[must_use = "readers do nothing unless polled"]
pub struct Chain<T, U> {
    first: T,
    second: U,
    done_first: bool,
}

impl<T, U> Unpin for Chain<T, U>
where
    T: Unpin,
    U: Unpin,
{
}

impl<T, U> Chain<T, U>
where
    T: AsyncRead,
    U: AsyncRead,
{
    unsafe_pinned!(first: T);
    unsafe_pinned!(second: U);
    unsafe_unpinned!(done_first: bool);

    pub(super) fn new(first: T, second: U) -> Self {
        Self {
            first,
            second,
            done_first: false,
        }
    }

    /// Gets references to the underlying readers in this `Chain`.
    pub fn get_ref(&self) -> (&T, &U) {
        (&self.first, &self.second)
    }

    /// Gets mutable references to the underlying readers in this `Chain`.
    ///
    /// Care should be taken to avoid modifying the internal I/O state of the
    /// underlying readers as doing so may corrupt the internal state of this
    /// `Chain`.
    pub fn get_mut(&mut self) -> (&mut T, &mut U) {
        (&mut self.first, &mut self.second)
    }

    /// Gets pinned mutable references to the underlying readers in this `Chain`.
    ///
    /// Care should be taken to avoid modifying the internal I/O state of the
    /// underlying readers as doing so may corrupt the internal state of this
    /// `Chain`.
    pub fn get_pin_mut(self: Pin<&mut Self>) -> (Pin<&mut T>, Pin<&mut U>) {
        unsafe {
            let Self { first, second, .. } = self.get_unchecked_mut();
            (Pin::new_unchecked(first), Pin::new_unchecked(second))
        }
    }

    /// Consumes the `Chain`, returning the wrapped readers.
    pub fn into_inner(self) -> (T, U) {
        (self.first, self.second)
    }
}

impl<T, U> fmt::Debug for Chain<T, U>
where
    T: fmt::Debug,
    U: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Chain")
            .field("t", &self.first)
            .field("u", &self.second)
            .finish()
    }
}

impl<T, U> AsyncRead for Chain<T, U>
where
    T: AsyncRead,
    U: AsyncRead,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        if !self.done_first {
            match ready!(self.as_mut().first().poll_read(cx, buf)?) {
                0 if !buf.is_empty() => *self.as_mut().done_first() = true,
                n => return Poll::Ready(Ok(n)),
            }
        }
        self.second().poll_read(cx, buf)
    }

    fn poll_read_vectored(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &mut [IoSliceMut<'_>],
    ) -> Poll<io::Result<usize>> {
        if !self.done_first {
            let n = ready!(self.as_mut().first().poll_read_vectored(cx, bufs)?);
            if n == 0 && bufs.iter().any(|b| !b.is_empty()) {
                *self.as_mut().done_first() = true
            } else {
                return Poll::Ready(Ok(n));
            }
        }
        self.second().poll_read_vectored(cx, bufs)
    }

    #[cfg(feature = "read-initializer")]
    unsafe fn initializer(&self) -> Initializer {
        let initializer = self.first.initializer();
        if initializer.should_initialize() {
            initializer
        } else {
            self.second.initializer()
        }
    }
}

impl<T, U> AsyncBufRead for Chain<T, U>
where
    T: AsyncBufRead,
    U: AsyncBufRead,
{
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        let Self {
            first,
            second,
            done_first,
        } = unsafe { self.get_unchecked_mut() };
        let first = unsafe { Pin::new_unchecked(first) };
        let second = unsafe { Pin::new_unchecked(second) };

        if !*done_first {
            match ready!(first.poll_fill_buf(cx)?) {
                buf if buf.is_empty() => {
                    *done_first = true;
                }
                buf => return Poll::Ready(Ok(buf)),
            }
        }
        second.poll_fill_buf(cx)
    }

    fn consume(self: Pin<&mut Self>, amt: usize) {
        if !self.done_first {
            self.first().consume(amt)
        } else {
            self.second().consume(amt)
        }
    }
}
<commit_msg>Fixed `impl Debug` for `Chain`<commit_after>use futures_core::task::{Context, Poll};
#[cfg(feature = "read-initializer")]
use futures_io::Initializer;
use futures_io::{AsyncBufRead, AsyncRead, IoSliceMut};
use pin_utils::{unsafe_pinned, unsafe_unpinned};
use std::fmt;
use std::io;
use std::pin::Pin;

/// Reader for the [`chain`](super::AsyncReadExt::chain) method.
#[must_use = "readers do nothing unless polled"]
pub struct Chain<T, U> {
    first: T,
    second: U,
    done_first: bool,
}

impl<T, U> Unpin for Chain<T, U>
where
    T: Unpin,
    U: Unpin,
{
}

impl<T, U> Chain<T, U>
where
    T: AsyncRead,
    U: AsyncRead,
{
    unsafe_pinned!(first: T);
    unsafe_pinned!(second: U);
    unsafe_unpinned!(done_first: bool);

    pub(super) fn new(first: T, second: U) -> Self {
        Self {
            first,
            second,
            done_first: false,
        }
    }

    /// Gets references to the underlying readers in this `Chain`.
    pub fn get_ref(&self) -> (&T, &U) {
        (&self.first, &self.second)
    }

    /// Gets mutable references to the underlying readers in this `Chain`.
    ///
    /// Care should be taken to avoid modifying the internal I/O state of the
    /// underlying readers as doing so may corrupt the internal state of this
    /// `Chain`.
    pub fn get_mut(&mut self) -> (&mut T, &mut U) {
        (&mut self.first, &mut self.second)
    }

    /// Gets pinned mutable references to the underlying readers in this `Chain`.
    ///
    /// Care should be taken to avoid modifying the internal I/O state of the
    /// underlying readers as doing so may corrupt the internal state of this
    /// `Chain`.
    pub fn get_pin_mut(self: Pin<&mut Self>) -> (Pin<&mut T>, Pin<&mut U>) {
        unsafe {
            let Self { first, second, .. } = self.get_unchecked_mut();
            (Pin::new_unchecked(first), Pin::new_unchecked(second))
        }
    }

    /// Consumes the `Chain`, returning the wrapped readers.
    pub fn into_inner(self) -> (T, U) {
        (self.first, self.second)
    }
}

impl<T, U> fmt::Debug for Chain<T, U>
where
    T: fmt::Debug,
    U: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Chain")
            .field("t", &self.first)
            .field("u", &self.second)
            .field("done_first", &self.done_first)
            .finish()
    }
}

impl<T, U> AsyncRead for Chain<T, U>
where
    T: AsyncRead,
    U: AsyncRead,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        if !self.done_first {
            match ready!(self.as_mut().first().poll_read(cx, buf)?) {
                0 if !buf.is_empty() => *self.as_mut().done_first() = true,
                n => return Poll::Ready(Ok(n)),
            }
        }
        self.second().poll_read(cx, buf)
    }

    fn poll_read_vectored(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &mut [IoSliceMut<'_>],
    ) -> Poll<io::Result<usize>> {
        if !self.done_first {
            let n = ready!(self.as_mut().first().poll_read_vectored(cx, bufs)?);
            if n == 0 && bufs.iter().any(|b| !b.is_empty()) {
                *self.as_mut().done_first() = true
            } else {
                return Poll::Ready(Ok(n));
            }
        }
        self.second().poll_read_vectored(cx, bufs)
    }

    #[cfg(feature = "read-initializer")]
    unsafe fn initializer(&self) -> Initializer {
        let initializer = self.first.initializer();
        if initializer.should_initialize() {
            initializer
        } else {
            self.second.initializer()
        }
    }
}

impl<T, U> AsyncBufRead for Chain<T, U>
where
    T: AsyncBufRead,
    U: AsyncBufRead,
{
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        let Self {
            first,
            second,
            done_first,
        } = unsafe { self.get_unchecked_mut() };
        let first = unsafe { Pin::new_unchecked(first) };
        let second = unsafe { Pin::new_unchecked(second) };

        if !*done_first {
            match ready!(first.poll_fill_buf(cx)?) {
                buf if buf.is_empty() => {
                    *done_first = true;
                }
                buf => return Poll::Ready(Ok(buf)),
            }
        }
        second.poll_fill_buf(cx)
    }

    fn consume(self: Pin<&mut Self>, amt: usize) {
        if !self.done_first {
            self.first().consume(amt)
        } else {
            self.second().consume(amt)
        }
    }
}
<|endoftext|>


<commit_before>use super::utils::types::{ Sha256Digest, Address };
use super::utils::bls::PublicKey;

pub struct ValidatorRecord {
    pub pubkey: PublicKey,
    pub withdrawal_shard: u16,
    pub withdrawal_address: Address,
    pub randao_commitment: Sha256Digest,
    pub balance: u64,
    pub switch_dynasty: u64
}

impl ValidatorRecord {
    pub fn new(pubkey: PublicKey,
               withdrawal_shard: u16,
               withdrawal_address: Address,
               randao_commitment: Sha256Digest,
               balance: u64,
               switch_dynasty: u64) -> ValidatorRecord {
        ValidatorRecord {
            pubkey: pubkey,
            withdrawal_shard: withdrawal_shard,
            withdrawal_address: withdrawal_address,
            randao_commitment: randao_commitment,
            balance: balance,
            switch_dynasty: switch_dynasty
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate rand;

    use super::*;
    use super::super::utils::bls::Keypair;
    use self::rand::{ SeedableRng, XorShiftRng };

    #[test]
    fn test_new() {
        let mut rng = XorShiftRng::from_seed([0xbc4f6d44, 0xd62f276c, 0xb963afd0, 0x5455863d]);
        let keypair = Keypair::generate(&mut rng);
        let withdrawal_shard = 1;
        let withdrawal_address = Address::random();
        let randao_commitment = Sha256Digest::random();
        let balance = 100;
        let switch_dynasty = 10;

        let v = ValidatorRecord::new(
            keypair.public,
            withdrawal_shard,
            withdrawal_address,
            randao_commitment,
            balance,
            switch_dynasty);
        // TODO: figure out how to compare keys
        // assert_eq!(v.pubkey, keypair.public);
        assert_eq!(v.withdrawal_shard, withdrawal_shard);
        assert_eq!(v.withdrawal_address, withdrawal_address);
        assert_eq!(v.randao_commitment, randao_commitment);
        assert_eq!(v.balance, balance);
        assert_eq!(v.switch_dynasty, switch_dynasty);
    }
}
<commit_msg>Add RLP serialization to validator_record<commit_after>use super::utils::types::{ Sha256Digest, Address };
use super::utils::bls::PublicKey;
use super::rlp::{ RlpStream, Encodable };

pub struct ValidatorRecord {
    pub pubkey: PublicKey,
    pub withdrawal_shard: u16,
    pub withdrawal_address: Address,
    pub randao_commitment: Sha256Digest,
    pub balance: u64,
    pub switch_dynasty: u64
}

impl ValidatorRecord {
    pub fn new(pubkey: PublicKey,
               withdrawal_shard: u16,
               withdrawal_address: Address,
               randao_commitment: Sha256Digest,
               balance: u64,
               switch_dynasty: u64) -> ValidatorRecord {
        ValidatorRecord {
            pubkey: pubkey,
            withdrawal_shard: withdrawal_shard,
            withdrawal_address: withdrawal_address,
            randao_commitment: randao_commitment,
            balance: balance,
            switch_dynasty: switch_dynasty
        }
    }
}

/*
 * RLP Encoding
 */
impl Encodable for ValidatorRecord {
    fn rlp_append(&self, s: &mut RlpStream) {
        // s.append(&self.pubkey);      // TODO: serialize this
        s.append(&self.withdrawal_shard);
        s.append(&self.withdrawal_address);
        s.append(&self.randao_commitment);
        s.append(&self.balance);
        s.append(&self.switch_dynasty);
    }
}

#[cfg(test)]
mod tests {
    use super::super::rlp;
    extern crate rand;

    use super::*;
    use super::super::utils::bls::Keypair;
    use self::rand::{ SeedableRng, XorShiftRng };

    fn get_keypair() -> Keypair {
        let mut rng = XorShiftRng::from_seed([0xbc4f6d44,
                                             0xd62f276c,
                                             0xb963afd0,
                                             0x5455863d]);
        Keypair::generate(&mut rng)
    }

    #[test]
    fn test_new() {
        let keypair = get_keypair();;
        let withdrawal_shard = 1;
        let withdrawal_address = Address::random();
        let randao_commitment = Sha256Digest::random();
        let balance = 100;
        let switch_dynasty = 10;

        let v = ValidatorRecord::new(
            keypair.public,
            withdrawal_shard,
            withdrawal_address,
            randao_commitment,
            balance,
            switch_dynasty);
        // TODO: figure out how to compare keys
        // assert_eq!(v.pubkey, keypair.public);
        assert_eq!(v.withdrawal_shard, withdrawal_shard);
        assert_eq!(v.withdrawal_address, withdrawal_address);
        assert_eq!(v.randao_commitment, randao_commitment);
        assert_eq!(v.balance, balance);
        assert_eq!(v.switch_dynasty, switch_dynasty);
    }

    #[test]
    fn test_serialization() {
        let keypair = get_keypair();
        let v = ValidatorRecord {
            pubkey: keypair.public,
            withdrawal_shard: 100,
            withdrawal_address: Address::zero(),
            randao_commitment: Sha256Digest::zero(),
            balance: 120,
            switch_dynasty: 30
        };
        let e = rlp::encode(&v);
        assert_eq!(e.len(), 57);    // TODO: fix when pubkey is serialized
        // TODO: test for serialized pubkey
        assert_eq!(e[0], 100);
        assert_eq!(e[1], 148);
        assert_eq!(e[2..22], [0; 20]);
        assert_eq!(e[22], 160);
        assert_eq!(e[23..55], [0; 32]);
        assert_eq!(e[55], 120);
        assert_eq!(e[56], 30);
    }
}
<|endoftext|>


<commit_before>extern crate rand;
extern crate timely;
extern crate differential_dataflow;

use rand::{Rng, SeedableRng, StdRng};

use timely::dataflow::*;
use timely::dataflow::operators::probe::Handle;

use differential_dataflow::input::Input;
use differential_dataflow::Collection;
use differential_dataflow::operators::*;
use differential_dataflow::lattice::Lattice;
use differential_dataflow::operators::iterate::Variable;
use differential_dataflow::operators::arrange::ArrangeByKey;
use differential_dataflow::operators::arrange::ArrangeBySelf;

type Node = usize;
type Edge = (Node, Node);

fn main() {

    let nodes: usize = std::env::args().nth(1).unwrap().parse().unwrap();
    let edges: usize = std::env::args().nth(2).unwrap().parse().unwrap();
    let batch: usize = std::env::args().nth(3).unwrap().parse().unwrap();
    let inspect: bool = std::env::args().any(|x| x == "inspect");

    // Our setting involves four read query types, and two updatable base relations.
    //
    //  Q1: Point lookup: reads "state" associated with a node.
    //  Q2: One-hop lookup: reads "state" associated with neighbors of a node.
    //  Q3: Two-hop lookup: reads "state" associated with n-of-n's of a node.
    //  Q4: Shortest path: reports hop count between two query nodes.
    //
    //  R1: "State": a pair of (node, T) for some type T that I don't currently know.
    //  R2: "Graph": pairs (node, node) indicating linkage between the two nodes.

    timely::execute_from_args(std::env::args().skip(3), move |worker| {

        let index = worker.index();
        let peers = worker.peers();
        let timer = ::std::time::Instant::now();

        // define BFS dataflow; return handles to roots and edges inputs
        let mut probe = Handle::new();

        let (mut q1, mut q2, mut q3, mut q4, mut state, mut graph) = worker.dataflow(|scope| {

            let (q1_input, q1) = scope.new_collection();
            let (q2_input, q2) = scope.new_collection();
            let (q3_input, q3) = scope.new_collection();
            let (q4_input, q4) = scope.new_collection();

            let (state_input, state) = scope.new_collection();
            let (graph_input, graph) = scope.new_collection();

            let state_indexed = state.arrange_by_key();
            let graph_indexed = graph.map(|(src, dst)| (dst, src))
                                     .concat(&graph)
                                     .arrange_by_key();

            // Q1: Point lookups on `state`:
            q1  .arrange_by_self()
                .join_core(&state_indexed, |&query, &(), &state| Some((query, state)))
                .filter(move |_| inspect)
                .inspect(|x| println!("Q1: {:?}", x))
                .probe_with(&mut probe);

            // Q2: One-hop lookups on `state`:
            q2  .arrange_by_self()
                .join_core(&graph_indexed, |&query, &(), &friend| Some((friend, query)))
                .join_core(&state_indexed, |_friend, &query, &state| Some((query, state)))
                .filter(move |_| inspect)
                .inspect(|x| println!("Q2: {:?}", x))
                .probe_with(&mut probe);

            // Q3: Two-hop lookups on `state`:
            q3  .arrange_by_self()
                .join_core(&graph_indexed, |&query, &(), &friend| Some((friend, query)))
                .join_core(&graph_indexed, |_friend, &query, &friend2| Some((friend2, query)))
                .join_core(&state_indexed, |_friend2, &query, &state| Some((query, state)))
                .filter(move |_| inspect)
                .consolidate()
                .inspect(|x| println!("Q3: {:?}", x))
                .probe_with(&mut probe);

            // Q4: Shortest path queries:
            bidijkstra(&graph_indexed, &graph_indexed, &q4)
                .filter(move |_| inspect)
                .inspect(|x| println!("Q4: {:?}", x))
                .probe_with(&mut probe);

            (q1_input, q2_input, q3_input, q4_input, state_input, graph_input)
        });

        let seed: &[_] = &[1, 2, 3, index];
        let mut rng1: StdRng = SeedableRng::from_seed(seed);    // rng for edge additions
        // let mut rng2: StdRng = SeedableRng::from_seed(seed);    // rng for edge deletions
        let mut rng3: StdRng = SeedableRng::from_seed(seed);    // rng for queries

        if index == 0 { println!("performing workload on random graph with {} nodes, {} edges:", nodes, edges); }

        let worker_edges = edges/peers + if index < (edges % peers) { 1 } else { 0 };
        for _ in 0 .. worker_edges {
            graph.insert((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)));
        }
        for node in 0 .. nodes {
            if node % peers == index {
                state.insert((node, node));
            }
        }

        q1.advance_to(1);       q1.flush();     // q1 queries start now.
        q2.advance_to(1001);    q2.flush();     // q2 queries start here.
        q3.advance_to(2001);    q3.flush();     // q3 queries start here.
        q4.advance_to(3001);    q4.flush();     // q4 queries start here.
        state.close();                          // no changes to state.
        graph.close();                          // no changes to graph.

        // finish graph loading work.
        while probe.less_than(q1.time()) { worker.step(); }

        if index == 0 { println!("{:?}\tgraph loaded", timer.elapsed()); }

        let worker_batch = batch/peers + if index < batch % peers { 1 } else { 0 };

        // Q1 testing:
        let timer_q1 = ::std::time::Instant::now();
        for round in 1 .. 1001 {
            for _ in 0 .. worker_batch {
                q1.insert(rng3.gen_range(0, nodes));
            }
            q1.advance_to(round);
            q1.flush();
            while probe.less_than(q1.time()) { worker.step(); }
        }
        if index == 0 { println!("{:?}\tq1 eval complete; avg: {:?}", timer.elapsed(), timer_q1.elapsed()/1000); }
        q1.close();

        // Q2 testing:
        let timer_q2 = ::std::time::Instant::now();
        for round in 1001 .. 2001 {
            for _ in 0 .. worker_batch {
                q2.insert(rng3.gen_range(0, nodes));
            }
            q2.advance_to(round);
            q2.flush();
            while probe.less_than(q2.time()) { worker.step(); }
        }
        if index == 0 { println!("{:?}\tq2 eval complete; avg: {:?}", timer.elapsed(), timer_q2.elapsed()/1000); }
        q2.close();

        // Q3 testing:
        let timer_q3 = ::std::time::Instant::now();
        for round in 2001 .. 3001 {
            for _ in 0 .. worker_batch {
                q3.insert(rng3.gen_range(0, nodes));
            }
            q3.advance_to(round);
            q3.flush();
            while probe.less_than(q3.time()) { worker.step(); }
        }
        if index == 0 { println!("{:?}\tq3 eval complete; avg: {:?}", timer.elapsed(), timer_q3.elapsed()/1000); }
        q3.close();

        // Q4 testing:
        let timer_q4 = ::std::time::Instant::now();
        for round in 3001 .. 4001 {
            for _ in 0 .. worker_batch {
               q4.insert((rng3.gen_range(0, nodes), rng3.gen_range(0, nodes)));
            }
            q4.advance_to(round);
            q4.flush();
            while probe.less_than(q4.time()) { worker.step(); }
        }
        if index == 0 { println!("{:?}\tq4 eval complete; avg: {:?}", timer.elapsed(), timer_q4.elapsed()/1000); }
        q4.close();

    }).unwrap();
}

use differential_dataflow::trace::implementations::ord::OrdValSpine as DefaultValTrace;
use differential_dataflow::operators::arrange::TraceAgent;
use differential_dataflow::operators::arrange::Arranged;

type Arrange<G: Scope, K, V, R> = Arranged<G, K, V, R, TraceAgent<K, V, G::Timestamp, R, DefaultValTrace<K, V, G::Timestamp, R>>>;

// returns pairs (n, s) indicating node n can be reached from a root in s steps.
fn bidijkstra<G: Scope>(
    forward_graph: &Arrange<G, Node, Node, isize>,
    reverse_graph: &Arrange<G, Node, Node, isize>,
    goals: &Collection<G, (Node, Node)>) -> Collection<G, ((Node, Node), u32)>
where G::Timestamp: Lattice+Ord {

    goals.scope().scoped(|inner| {

        // Our plan is to start evolving distances from both sources and destinations.
        // The evolution from a source or destination should continue as long as there
        // is a corresponding destination or source that has not yet been reached.

        // forward and reverse (node, (root, dist))
        let forward = Variable::from(goals.map(|(x,_)| (x,(x,0))).enter(inner));
        let reverse = Variable::from(goals.map(|(_,y)| (y,(y,0))).enter(inner));

        let goals = goals.enter(inner);
        let forward_graph = forward_graph.enter(inner);
        let reverse_graph = reverse_graph.enter(inner);

        // Let's determine which (src, dst) pairs are ready to return.
        //
        //   done(src, dst) := forward(src, med), reverse(dst, med), goal(src, dst).
        //
        // This is a cyclic join, which should scare us a bunch.
        let reached =
        forward
            .join_map(&reverse, |_, &(src,d1), &(dst,d2)| ((src, dst), d1 + d2))
            .group(|_key, s, t| t.push((*s[0].0, 1)))
            .semijoin(&goals);

        let active =
        reached
            .negate()
            .map(|(srcdst,_)| srcdst)
            .concat(&goals)
            .consolidate();

        // Let's expand out forward queries that are active.
        let forward_active = active.map(|(x,_y)| x).distinct();
        let forward_next =
        forward
            .map(|(med, (src, dist))| (src, (med, dist)))
            .semijoin(&forward_active)
            .map(|(src, (med, dist))| (med, (src, dist)))
            .join_core(&forward_graph, |_med, &(src, dist), &next| Some((next, (src, dist+1))))
            // .join_map(&edges, |_med, &(src, dist), &next| (next, (src, dist+1)))
            .concat(&forward)
            .map(|(next, (src, dist))| ((next, src), dist))
            .group(|_key, s, t| t.push((*s[0].0, 1)))
            .map(|((next, src), dist)| (next, (src, dist)));

        forward.set(&forward_next);

        // Let's expand out reverse queries that are active.
        let reverse_active = active.map(|(_x,y)| y).distinct();
        let reverse_next =
        reverse
            .map(|(med, (rev, dist))| (rev, (med, dist)))
            .semijoin(&reverse_active)
            .map(|(rev, (med, dist))| (med, (rev, dist)))
            .join_core(&reverse_graph, |_med, &(rev, dist), &next| Some((next, (rev, dist+1))))
            // .join_map(&edges.map(|(x,y)| (y,x)), |_med, &(rev, dist), &next| (next, (rev, dist+1)))
            .concat(&reverse)
            .map(|(next, (rev, dist))| ((next, rev), dist))
            .group(|_key, s, t| t.push((*s[0].0, 1)))
            .map(|((next,rev), dist)| (next, (rev, dist)));

        reverse.set(&reverse_next);

        reached.leave()
    })
}<commit_msg>graph example update<commit_after>extern crate rand;
extern crate timely;
extern crate differential_dataflow;

use rand::{Rng, SeedableRng, StdRng};

use timely::dataflow::*;
use timely::dataflow::operators::probe::Handle;

use differential_dataflow::input::Input;
use differential_dataflow::Collection;
use differential_dataflow::operators::*;
use differential_dataflow::lattice::Lattice;
use differential_dataflow::operators::iterate::Variable;
use differential_dataflow::operators::arrange::ArrangeByKey;
use differential_dataflow::operators::arrange::ArrangeBySelf;

type Node = usize;

fn main() {

    let nodes: usize = std::env::args().nth(1).unwrap().parse().unwrap();
    let edges: usize = std::env::args().nth(2).unwrap().parse().unwrap();
    let batch: usize = std::env::args().nth(3).unwrap().parse().unwrap();
    let inspect: bool = std::env::args().any(|x| x == "inspect");

    // Our setting involves four read query types, and two updatable base relations.
    //
    //  Q1: Point lookup: reads "state" associated with a node.
    //  Q2: One-hop lookup: reads "state" associated with neighbors of a node.
    //  Q3: Two-hop lookup: reads "state" associated with n-of-n's of a node.
    //  Q4: Shortest path: reports hop count between two query nodes.
    //
    //  R1: "State": a pair of (node, T) for some type T that I don't currently know.
    //  R2: "Graph": pairs (node, node) indicating linkage between the two nodes.

    timely::execute_from_args(std::env::args().skip(3), move |worker| {

        let index = worker.index();
        let peers = worker.peers();
        let timer = ::std::time::Instant::now();

        // define BFS dataflow; return handles to roots and edges inputs
        let mut probe = Handle::new();

        let (mut q1, mut q2, mut q3, mut q4, mut state, mut graph) = worker.dataflow(|scope| {

            let (q1_input, q1) = scope.new_collection();
            let (q2_input, q2) = scope.new_collection();
            let (q3_input, q3) = scope.new_collection();
            let (q4_input, q4) = scope.new_collection();

            let (state_input, state) = scope.new_collection();
            let (graph_input, graph) = scope.new_collection();

            let state_indexed = state.arrange_by_key();
            let graph_indexed = graph.map(|(src, dst)| (dst, src))
                                     .concat(&graph)
                                     .arrange_by_key();

            // Q1: Point lookups on `state`:
            q1  .arrange_by_self()
                .join_core(&state_indexed, |&query, &(), &state| Some((query, state)))
                .filter(move |_| inspect)
                .inspect(|x| println!("Q1: {:?}", x))
                .probe_with(&mut probe);

            // Q2: One-hop lookups on `state`:
            q2  .arrange_by_self()
                .join_core(&graph_indexed, |&query, &(), &friend| Some((friend, query)))
                .join_core(&state_indexed, |_friend, &query, &state| Some((query, state)))
                .filter(move |_| inspect)
                .inspect(|x| println!("Q2: {:?}", x))
                .probe_with(&mut probe);

            // Q3: Two-hop lookups on `state`:
            q3  .arrange_by_self()
                .join_core(&graph_indexed, |&query, &(), &friend| Some((friend, query)))
                .join_core(&graph_indexed, |_friend, &query, &friend2| Some((friend2, query)))
                .join_core(&state_indexed, |_friend2, &query, &state| Some((query, state)))
                .filter(move |_| inspect)
                .consolidate()
                .inspect(|x| println!("Q3: {:?}", x))
                .probe_with(&mut probe);

            // Q4: Shortest path queries:
            bidijkstra(&graph_indexed, &graph_indexed, &q4)
                .filter(move |_| inspect)
                .inspect(|x| println!("Q4: {:?}", x))
                .probe_with(&mut probe);

            (q1_input, q2_input, q3_input, q4_input, state_input, graph_input)
        });

        let seed: &[_] = &[1, 2, 3, index];
        let mut rng1: StdRng = SeedableRng::from_seed(seed);    // rng for edge additions
        // let mut rng2: StdRng = SeedableRng::from_seed(seed);    // rng for edge deletions
        let mut rng3: StdRng = SeedableRng::from_seed(seed);    // rng for queries

        if index == 0 { println!("performing workload on random graph with {} nodes, {} edges:", nodes, edges); }

        let worker_edges = edges/peers + if index < (edges % peers) { 1 } else { 0 };
        for _ in 0 .. worker_edges {
            graph.insert((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)));
        }
        for node in 0 .. nodes {
            if node % peers == index {
                state.insert((node, node));
            }
        }

        q1.advance_to(1);       q1.flush();     // q1 queries start now.
        q2.advance_to(1001);    q2.flush();     // q2 queries start here.
        q3.advance_to(2001);    q3.flush();     // q3 queries start here.
        q4.advance_to(3001);    q4.flush();     // q4 queries start here.
        state.close();                          // no changes to state.
        graph.close();                          // no changes to graph.

        // finish graph loading work.
        while probe.less_than(q1.time()) { worker.step(); }

        if index == 0 { println!("{:?}\tgraph loaded", timer.elapsed()); }

        let worker_batch = batch/peers + if index < batch % peers { 1 } else { 0 };

        // Q1 testing:
        let timer_q1 = ::std::time::Instant::now();
        for round in 1 .. 1001 {
            for _ in 0 .. worker_batch {
                q1.insert(rng3.gen_range(0, nodes));
            }
            q1.advance_to(round);
            q1.flush();
            while probe.less_than(q1.time()) { worker.step(); }
        }
        if index == 0 { println!("{:?}\tq1 eval complete; avg: {:?}", timer.elapsed(), timer_q1.elapsed()/1000); }
        q1.close();

        // Q2 testing:
        let timer_q2 = ::std::time::Instant::now();
        for round in 1001 .. 2001 {
            for _ in 0 .. worker_batch {
                q2.insert(rng3.gen_range(0, nodes));
            }
            q2.advance_to(round);
            q2.flush();
            while probe.less_than(q2.time()) { worker.step(); }
        }
        if index == 0 { println!("{:?}\tq2 eval complete; avg: {:?}", timer.elapsed(), timer_q2.elapsed()/1000); }
        q2.close();

        // Q3 testing:
        let timer_q3 = ::std::time::Instant::now();
        for round in 2001 .. 3001 {
            for _ in 0 .. worker_batch {
                q3.insert(rng3.gen_range(0, nodes));
            }
            q3.advance_to(round);
            q3.flush();
            while probe.less_than(q3.time()) { worker.step(); }
        }
        if index == 0 { println!("{:?}\tq3 eval complete; avg: {:?}", timer.elapsed(), timer_q3.elapsed()/1000); }
        q3.close();

        // Q4 testing:
        let timer_q4 = ::std::time::Instant::now();
        for round in 3001 .. 4001 {
            for _ in 0 .. worker_batch {
               q4.insert((rng3.gen_range(0, nodes), rng3.gen_range(0, nodes)));
            }
            q4.advance_to(round);
            q4.flush();
            while probe.less_than(q4.time()) { worker.step(); }
        }
        if index == 0 { println!("{:?}\tq4 eval complete; avg: {:?}", timer.elapsed(), timer_q4.elapsed()/1000); }
        q4.close();

    }).unwrap();
}

use differential_dataflow::trace::implementations::ord::OrdValSpine as DefaultValTrace;
use differential_dataflow::operators::arrange::TraceAgent;
use differential_dataflow::operators::arrange::Arranged;

type Arrange<G: Scope, K, V, R> = Arranged<G, K, V, R, TraceAgent<K, V, G::Timestamp, R, DefaultValTrace<K, V, G::Timestamp, R>>>;

// returns pairs (n, s) indicating node n can be reached from a root in s steps.
fn bidijkstra<G: Scope>(
    forward_graph: &Arrange<G, Node, Node, isize>,
    reverse_graph: &Arrange<G, Node, Node, isize>,
    goals: &Collection<G, (Node, Node)>) -> Collection<G, ((Node, Node), u32)>
where G::Timestamp: Lattice+Ord {

    goals.scope().scoped(|inner| {

        // Our plan is to start evolving distances from both sources and destinations.
        // The evolution from a source or destination should continue as long as there
        // is a corresponding destination or source that has not yet been reached.

        // forward and reverse (node, (root, dist))
        let forward = Variable::from(goals.map(|(x,_)| (x,(x,0))).enter(inner));
        let reverse = Variable::from(goals.map(|(_,y)| (y,(y,0))).enter(inner));

        let goals = goals.enter(inner);
        let forward_graph = forward_graph.enter(inner);
        let reverse_graph = reverse_graph.enter(inner);

        // Let's determine which (src, dst) pairs are ready to return.
        //
        //   done(src, dst) := forward(src, med), reverse(dst, med), goal(src, dst).
        //
        // This is a cyclic join, which should scare us a bunch.
        let reached =
        forward
            .join_map(&reverse, |_, &(src,d1), &(dst,d2)| ((src, dst), d1 + d2))
            .group(|_key, s, t| t.push((*s[0].0, 1)))
            .semijoin(&goals);

        let active =
        reached
            .negate()
            .map(|(srcdst,_)| srcdst)
            .concat(&goals)
            .consolidate();

        // Let's expand out forward queries that are active.
        let forward_active = active.map(|(x,_y)| x).distinct();
        let forward_next =
        forward
            .map(|(med, (src, dist))| (src, (med, dist)))
            .semijoin(&forward_active)
            .map(|(src, (med, dist))| (med, (src, dist)))
            .join_core(&forward_graph, |_med, &(src, dist), &next| Some((next, (src, dist+1))))
            .concat(&forward)
            .map(|(next, (src, dist))| ((next, src), dist))
            .group(|_key, s, t| t.push((*s[0].0, 1)))
            .map(|((next, src), dist)| (next, (src, dist)));

        forward.set(&forward_next);

        // Let's expand out reverse queries that are active.
        let reverse_active = active.map(|(_x,y)| y).distinct();
        let reverse_next =
        reverse
            .map(|(med, (rev, dist))| (rev, (med, dist)))
            .semijoin(&reverse_active)
            .map(|(rev, (med, dist))| (med, (rev, dist)))
            .join_core(&reverse_graph, |_med, &(rev, dist), &next| Some((next, (rev, dist+1))))
            .concat(&reverse)
            .map(|(next, (rev, dist))| ((next, rev), dist))
            .group(|_key, s, t| t.push((*s[0].0, 1)))
            .map(|((next,rev), dist)| (next, (rev, dist)));

        reverse.set(&reverse_next);

        reached.leave()
    })
}<|endoftext|>