<commit_before>extern crate libc;
extern crate xlib;

use self::libc::{ c_char, c_int, c_void };
use self::libc::funcs::c95::stdlib::malloc;
use self::xlib::{
    Display,
    Window,
    XCrossingEvent,
    XSetWindowBorder,
    XSetWindowBorderWidth,
    XDefaultRootWindow,
    XDisplayWidth,
    XDisplayHeight,
    XEnterWindowEvent,
    XFetchName,
    XLeaveWindowEvent,
    XMapRequestEvent,
    XMapWindow,
    XMotionEvent,
    XMoveWindow,
    XNextEvent,
    XOpenDisplay,
    XPending,
    XResizeWindow,
    XSelectInput
};

use std::ptr::null_mut;
use std::mem::transmute;
use std::mem::uninitialized;
use std::str::raw::c_str_to_static_slice;

use window_system::{ WindowSystem, WindowSystemEvent };
use window_system::{
    Enter,
    Leave,
    WindowCreated,
    WindowDestroyed,
    UnknownEvent
};

const KeyPress               : uint = 2;
const KeyRelease             : uint = 3;
const ButtonPress            : uint = 4;
const ButtonRelease          : uint = 5;
const MotionNotify           : uint = 6;
const EnterNotify            : uint = 7;
const LeaveNotify            : uint = 8;
const FocusIn                : uint = 9;
const FocusOut               : uint = 10;
const KeymapNotify           : uint = 11;
const Expose                 : uint = 12;
const GraphicsExpose         : uint = 13;
const NoExpose               : uint = 14;
const VisibilityNotify       : uint = 15;
const CreateNotify           : uint = 16;
const DestroyNotify          : uint = 17;
const UnmapNotify            : uint = 18;
const MapNotify              : uint = 19;
const MapRequest             : uint = 20;
const ReparentNotify         : uint = 21;
const ConfigureNotify        : uint = 22;
const ConfigureRequest       : uint = 23;
const GravityNotify          : uint = 24;
const ResizeRequest          : uint = 25;
const CirculateNotify        : uint = 26;
const CirculateRequest       : uint = 27;
const PropertyNotify         : uint = 28;
const SelectionClear         : uint = 29;
const SelectionRequest       : uint = 30;
const SelectionNotify        : uint = 31;
const ColormapNotify         : uint = 32;
const ClientMessage          : uint = 33;
const MappingNotify          : uint = 34;

pub struct XlibWindowSystem {
    display: *mut Display,
    root:    Window,
    event:   *mut c_void
}

impl XlibWindowSystem {
    pub fn new() -> XlibWindowSystem {
        unsafe {
            let display = XOpenDisplay(null_mut());
            let root    = XDefaultRootWindow(display);

            XSelectInput(display, root, 0x1E0030i64);

            XlibWindowSystem {
                display: display,
                root:    root,
                event:   malloc(256)
            }
        }
    }

    fn get_event_as<T>(&self) -> &T {
        unsafe {
            let event_ptr : *const T = transmute(self.event);
            let ref event = *event_ptr;
            event
        }
    }
}

impl WindowSystem for XlibWindowSystem {
    fn get_display_width(&self, screen: uint) -> u32 {
        unsafe {
            XDisplayWidth(self.display, screen as i32) as u32
        }
    }

    fn get_display_height(&self, screen: uint) -> u32 {
        unsafe {
            XDisplayHeight(self.display, screen as i32) as u32
        }
    }

    fn get_window_name(&self, window: u64) -> String {
        unsafe {
            let mut name : *mut c_char = uninitialized();
            let mut name_ptr : *mut *mut c_char = &mut name;
            XFetchName(self.display, window, name_ptr);
            String::from_str(c_str_to_static_slice(transmute(*name_ptr)))
        }
    }

    fn set_window_border_width(&mut self, window: u64, border_width: uint) {
        unsafe {
            XSetWindowBorderWidth(self.display, window, border_width as u32);
        }
    }

    fn set_window_border_color(&mut self, window: u64, border_color: uint) {
        unsafe {
            XSetWindowBorder(self.display, window, border_color as u64);
        }
    }

    fn resize_window(&mut self, window: u64, width: u32, height: u32) {
        unsafe {
            XResizeWindow(self.display, window, width, height);
        }
    }

    fn move_window(&mut self, window: u64, x: u32, y: u32) {
        unsafe {
            XMoveWindow(self.display, window, x as i32, y as i32);
        }
    }

    fn show_window(&mut self, window: u64) {
        unsafe {
            XMapWindow(self.display, window);
        }
    }

    fn event_pending(&self) -> bool {
        unsafe {
            XPending(self.display) != 0
        }
    }

    fn get_event(&mut self) -> WindowSystemEvent {
        unsafe {
            XNextEvent(self.display, self.event);

            let event_type : c_int = *self.get_event_as();

            match event_type as uint {
                MapRequest => {
                    let event : &XMapRequestEvent = self.get_event_as();
                    WindowCreated(event.window)
                },
                EnterNotify => {
                    let event : &XEnterWindowEvent = self.get_event_as();
                    Enter(event.window)
                },
                LeaveNotify => {
                    let event : &XLeaveWindowEvent = self.get_event_as();
                    Leave(event.window)

                },
                _  => {
                    UnknownEvent
                }
            }
        }
    }
}
<commit_msg>minimized unsafe region in get_event()<commit_after>extern crate libc;
extern crate xlib;

use self::libc::{ c_char, c_int, c_void };
use self::libc::funcs::c95::stdlib::malloc;
use self::xlib::{
    Display,
    Window,
    XCrossingEvent,
    XSetWindowBorder,
    XSetWindowBorderWidth,
    XDefaultRootWindow,
    XDisplayWidth,
    XDisplayHeight,
    XEnterWindowEvent,
    XFetchName,
    XLeaveWindowEvent,
    XMapRequestEvent,
    XMapWindow,
    XMotionEvent,
    XMoveWindow,
    XNextEvent,
    XOpenDisplay,
    XPending,
    XResizeWindow,
    XSelectInput
};

use std::ptr::null_mut;
use std::mem::transmute;
use std::mem::uninitialized;
use std::str::raw::c_str_to_static_slice;

use window_system::{ WindowSystem, WindowSystemEvent };
use window_system::{
    Enter,
    Leave,
    WindowCreated,
    WindowDestroyed,
    UnknownEvent
};

const KeyPress               : uint = 2;
const KeyRelease             : uint = 3;
const ButtonPress            : uint = 4;
const ButtonRelease          : uint = 5;
const MotionNotify           : uint = 6;
const EnterNotify            : uint = 7;
const LeaveNotify            : uint = 8;
const FocusIn                : uint = 9;
const FocusOut               : uint = 10;
const KeymapNotify           : uint = 11;
const Expose                 : uint = 12;
const GraphicsExpose         : uint = 13;
const NoExpose               : uint = 14;
const VisibilityNotify       : uint = 15;
const CreateNotify           : uint = 16;
const DestroyNotify          : uint = 17;
const UnmapNotify            : uint = 18;
const MapNotify              : uint = 19;
const MapRequest             : uint = 20;
const ReparentNotify         : uint = 21;
const ConfigureNotify        : uint = 22;
const ConfigureRequest       : uint = 23;
const GravityNotify          : uint = 24;
const ResizeRequest          : uint = 25;
const CirculateNotify        : uint = 26;
const CirculateRequest       : uint = 27;
const PropertyNotify         : uint = 28;
const SelectionClear         : uint = 29;
const SelectionRequest       : uint = 30;
const SelectionNotify        : uint = 31;
const ColormapNotify         : uint = 32;
const ClientMessage          : uint = 33;
const MappingNotify          : uint = 34;

pub struct XlibWindowSystem {
    display: *mut Display,
    root:    Window,
    event:   *mut c_void
}

impl XlibWindowSystem {
    pub fn new() -> XlibWindowSystem {
        unsafe {
            let display = XOpenDisplay(null_mut());
            let root    = XDefaultRootWindow(display);

            XSelectInput(display, root, 0x1E0030i64);

            XlibWindowSystem {
                display: display,
                root:    root,
                event:   malloc(256)
            }
        }
    }

    fn get_event_as<T>(&self) -> &T {
        unsafe {
            let event_ptr : *const T = transmute(self.event);
            let ref event = *event_ptr;
            event
        }
    }
}

impl WindowSystem for XlibWindowSystem {
    fn get_display_width(&self, screen: uint) -> u32 {
        unsafe {
            XDisplayWidth(self.display, screen as i32) as u32
        }
    }

    fn get_display_height(&self, screen: uint) -> u32 {
        unsafe {
            XDisplayHeight(self.display, screen as i32) as u32
        }
    }

    fn get_window_name(&self, window: u64) -> String {
        unsafe {
            let mut name : *mut c_char = uninitialized();
            let mut name_ptr : *mut *mut c_char = &mut name;
            XFetchName(self.display, window, name_ptr);
            String::from_str(c_str_to_static_slice(transmute(*name_ptr)))
        }
    }

    fn set_window_border_width(&mut self, window: u64, border_width: uint) {
        unsafe {
            XSetWindowBorderWidth(self.display, window, border_width as u32);
        }
    }

    fn set_window_border_color(&mut self, window: u64, border_color: uint) {
        unsafe {
            XSetWindowBorder(self.display, window, border_color as u64);
        }
    }

    fn resize_window(&mut self, window: u64, width: u32, height: u32) {
        unsafe {
            XResizeWindow(self.display, window, width, height);
        }
    }

    fn move_window(&mut self, window: u64, x: u32, y: u32) {
        unsafe {
            XMoveWindow(self.display, window, x as i32, y as i32);
        }
    }

    fn show_window(&mut self, window: u64) {
        unsafe {
            XMapWindow(self.display, window);
        }
    }

    fn event_pending(&self) -> bool {
        unsafe {
            XPending(self.display) != 0
        }
    }

    fn get_event(&mut self) -> WindowSystemEvent {
        unsafe {
            XNextEvent(self.display, self.event);
        }

        let event_type : c_int = *self.get_event_as();

        match event_type as uint {
            MapRequest => {
                let event : &XMapRequestEvent = self.get_event_as();
                WindowCreated(event.window)
            },
            EnterNotify => {
                let event : &XEnterWindowEvent = self.get_event_as();
                Enter(event.window)
            },
            LeaveNotify => {
                let event : &XLeaveWindowEvent = self.get_event_as();
                Leave(event.window)
            },
            _  => {
                UnknownEvent
            }
        }
    }
}
<|endoftext|>


<commit_before>trait Foo {
    fn f();
}

trait Bar : Foo {
    fn g();
}

struct A {
    x: int
}

impl A : Bar {
    fn g() { io::println("in g"); }
    fn f() { io::println("in f"); }
}

fn h<T:Foo>(a: &T) {
    a.f();
}

fn main() {
    let a = A { x: 3 };
    h(&a);
}

<commit_msg>test: XFAIL trait-inheritance-simple because of suspected 32-bit brokenness<commit_after>// xfail-fast
// xfail-test

// Broken on 32 bit, I guess?

trait Foo {
    fn f();
}

trait Bar : Foo {
    fn g();
}

struct A {
    x: int
}

impl A : Bar {
    fn g() { io::println("in g"); }
    fn f() { io::println("in f"); }
}

fn h<T:Foo>(a: &T) {
    a.f();
}

fn main() {
    let a = A { x: 3 };
    h(&a);
}

<|endoftext|>


<commit_before><commit_msg>libboot: Add dummy do_hypervisor_callback<commit_after>#[no_mangle]
pub unsafe extern "C" fn do_hypervisor_callback() {
    panic!("Hypercall callback");
}
<|endoftext|>


<commit_before>//! Panic support for libcore
//!
//! The core library cannot define panicking, but it does *declare* panicking. This
//! means that the functions inside of libcore are allowed to panic, but to be
//! useful an upstream crate must define panicking for libcore to use. The current
//! interface for panicking is:
//!
//! ```
//! # use std::fmt;
//! fn panic_impl(fmt: fmt::Arguments, file_line_col: &(&'static str, u32, u32)) -> !
//! # { loop {} }
//! ```
//!
//! This definition allows for panicking with any general message, but it does not
//! allow for failing with a `Box<Any>` value. The reason for this is that libcore
//! is not allowed to allocate.
//!
//! This module contains a few other panicking functions, but these are just the
//! necessary lang items for the compiler. All panics are funneled through this
//! one function. Currently, the actual symbol is declared in the standard
//! library, but the location of this may change over time.

#![allow(dead_code, missing_docs)]
#![unstable(feature = "core_panic",
            reason = "internal details of the implementation of the `panic!` \
                      and related macros",
            issue = "0")]

use crate::fmt;
use crate::panic::{Location, PanicInfo};

#[cold]
// never inline unless panic_immediate_abort to avoid code
// bloat at the call sites as much as possible
#[cfg_attr(not(feature="panic_immediate_abort"),inline(never))]
#[lang = "panic"]
pub fn panic(expr_file_line_col: &(&'static str, &'static str, u32, u32)) -> ! {
    if cfg!(feature = "panic_immediate_abort") {
        unsafe { super::intrinsics::abort() }
    }

    // Use Arguments::new_v1 instead of format_args!("{}", expr) to potentially
    // reduce size overhead. The format_args! macro uses str's Display trait to
    // write expr, which calls Formatter::pad, which must accommodate string
    // truncation and padding (even though none is used here). Using
    // Arguments::new_v1 may allow the compiler to omit Formatter::pad from the
    // output binary, saving up to a few kilobytes.
    let (expr, file, line, col) = *expr_file_line_col;
    panic_fmt(fmt::Arguments::new_v1(&[expr], &[]), &(file, line, col))
}

#[cold]
#[cfg_attr(not(feature="panic_immediate_abort"),inline(never))]
#[lang = "panic_bounds_check"]
fn panic_bounds_check(file_line_col: &(&'static str, u32, u32),
                     index: usize, len: usize) -> ! {
    if cfg!(feature = "panic_immediate_abort") {
        unsafe { super::intrinsics::abort() }
    }

    panic_fmt(format_args!("index out of bounds: the len is {} but the index is {}",
                           len, index), file_line_col)
}

#[cold]
#[cfg_attr(not(feature="panic_immediate_abort"),inline(never))]
#[cfg_attr(    feature="panic_immediate_abort" ,inline)]
pub fn panic_fmt(fmt: fmt::Arguments<'_>, file_line_col: &(&'static str, u32, u32)) -> ! {
    if cfg!(feature = "panic_immediate_abort") {
        unsafe { super::intrinsics::abort() }
    }

    // NOTE This function never crosses the FFI boundary; it's a Rust-to-Rust call
    #[cfg_attr(boostrap_stdarch_ignore_this, allow(improper_ctypes))]
    extern "Rust" {
        #[lang = "panic_impl"]
        fn panic_impl(pi: &PanicInfo<'_>) -> !;
    }

    let (file, line, col) = *file_line_col;
    let pi = PanicInfo::internal_constructor(
        Some(&fmt),
        Location::internal_constructor(file, line, col),
    );
    unsafe { panic_impl(&pi) }
}
<commit_msg>Fix rebase conflicts<commit_after>//! Panic support for libcore
//!
//! The core library cannot define panicking, but it does *declare* panicking. This
//! means that the functions inside of libcore are allowed to panic, but to be
//! useful an upstream crate must define panicking for libcore to use. The current
//! interface for panicking is:
//!
//! ```
//! # use std::fmt;
//! fn panic_impl(fmt: fmt::Arguments, file_line_col: &(&'static str, u32, u32)) -> !
//! # { loop {} }
//! ```
//!
//! This definition allows for panicking with any general message, but it does not
//! allow for failing with a `Box<Any>` value. The reason for this is that libcore
//! is not allowed to allocate.
//!
//! This module contains a few other panicking functions, but these are just the
//! necessary lang items for the compiler. All panics are funneled through this
//! one function. Currently, the actual symbol is declared in the standard
//! library, but the location of this may change over time.

#![allow(dead_code, missing_docs)]
#![unstable(feature = "core_panic",
            reason = "internal details of the implementation of the `panic!` \
                      and related macros",
            issue = "0")]

use crate::fmt;
use crate::panic::{Location, PanicInfo};

#[cold]
// never inline unless panic_immediate_abort to avoid code
// bloat at the call sites as much as possible
#[cfg_attr(not(feature="panic_immediate_abort"),inline(never))]
#[lang = "panic"]
pub fn panic(expr_file_line_col: &(&'static str, &'static str, u32, u32)) -> ! {
    if cfg!(feature = "panic_immediate_abort") {
        unsafe { super::intrinsics::abort() }
    }

    // Use Arguments::new_v1 instead of format_args!("{}", expr) to potentially
    // reduce size overhead. The format_args! macro uses str's Display trait to
    // write expr, which calls Formatter::pad, which must accommodate string
    // truncation and padding (even though none is used here). Using
    // Arguments::new_v1 may allow the compiler to omit Formatter::pad from the
    // output binary, saving up to a few kilobytes.
    let (expr, file, line, col) = *expr_file_line_col;
    panic_fmt(fmt::Arguments::new_v1(&[expr], &[]), &(file, line, col))
}

#[cold]
#[cfg_attr(not(feature="panic_immediate_abort"),inline(never))]
#[lang = "panic_bounds_check"]
fn panic_bounds_check(file_line_col: &(&'static str, u32, u32),
                     index: usize, len: usize) -> ! {
    if cfg!(feature = "panic_immediate_abort") {
        unsafe { super::intrinsics::abort() }
    }

    panic_fmt(format_args!("index out of bounds: the len is {} but the index is {}",
                           len, index), file_line_col)
}

#[cold]
#[cfg_attr(not(feature="panic_immediate_abort"),inline(never))]
#[cfg_attr(    feature="panic_immediate_abort" ,inline)]
pub fn panic_fmt(fmt: fmt::Arguments<'_>, file_line_col: &(&'static str, u32, u32)) -> ! {
    if cfg!(feature = "panic_immediate_abort") {
        unsafe { super::intrinsics::abort() }
    }

    // NOTE This function never crosses the FFI boundary; it's a Rust-to-Rust call
    #[cfg_attr(bootstrap, allow(improper_ctypes))]
    extern "Rust" {
        #[lang = "panic_impl"]
        fn panic_impl(pi: &PanicInfo<'_>) -> !;
    }

    let (file, line, col) = *file_line_col;
    let pi = PanicInfo::internal_constructor(
        Some(&fmt),
        Location::internal_constructor(file, line, col),
    );
    unsafe { panic_impl(&pi) }
}
<|endoftext|>


<commit_before>extern crate time;
extern crate rand;
extern crate timely;
extern crate radix_sort;

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

use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::dataflow::channels::pact::Exchange;

fn main() {

    // command-line args: numbers of nodes and edges in the random graph.
    let nodes: usize = std::env::args().nth(1).unwrap().parse().unwrap();
    let edges: usize = std::env::args().nth(2).unwrap().parse().unwrap();

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

        let index = root.index();
        let peers = root.peers();

        let seed: &[_] = &[1, 2, 3, root.index()];
        let mut rng: StdRng = SeedableRng::from_seed(seed);
        let mut sorter = RadixSorter::new();

        let mut offsets = Vec::new();
        let mut targets = Vec::new();
        let mut recents = Vec::new();
        let mut parents = vec![u32::max_value(); nodes / peers];

        let start = time::precise_time_s();

        root.scoped(move |scope| {

            // generate part of a random graph.
            let graph = (0..edges / peers)
                .map(move |_| (rng.gen_range(0u32, nodes as u32), rng.gen_range(0u32, nodes as u32)))
                .to_stream(scope);

            // define a loop variable, for the (node, root) pairs.
            let (handle, stream) = scope.loop_variable(usize::max_value(), 1);

            // use the stream of edges
            graph.binary_notify(
                &stream,
                Exchange::new(|x: &(u32, u32)| x.0 as u64),
                Exchange::new(|x: &(u32, u32)| x.0 as u64),
                "BFS",
                vec![],
                move |input1, input2, output, notify| {

                    // receive edges, start to sort them
                    input1.for_each(|time, data| {
                        notify.notify_at(time);
                        for &edge in data.iter() { sorter.push(edge, &|x| x.0); }
                    });

                    // receive (node, root) pairs, note any new ones.
                    input2.for_each(|time, data| {
                        notify.notify_at(time);
                        for &(node, root) in data.iter() {
                            // println!("received: {:?}, parents[{}] = {}", (node, root), node, parents[node as usize]);
                            if parents[(node as usize) / peers] == u32::max_value() {
                                // println!("keeping: {:?}", node);
                                parents[(node as usize) / peers] = root;
                                recents.push(node);
                            }
                        }
                    });

                    notify.for_each(|time, _count| {

                        // maybe process the graph
                        if time.inner == 0 {

                            if index == 0 {
                                println!("{}:\tstarting sort", time::precise_time_s() - start);
                            }

                            // construct the graph
                            offsets.push(0);
                            let mut prev_node = 0;
                            let sorted = sorter.finish(&|x| x.0);
                            for buffer in &sorted {
                                for &(node, edge) in buffer {
                                    while prev_node < node {
                                        prev_node += 1;
                                        offsets.push(targets.len())
                                    }
                                    targets.push(edge);
                                }
                            }
                            offsets.push(targets.len())

                        }

                        if index == 0 {
                            println!("{}:\tworker: {}, todo.len(): {}", time::precise_time_s() - start, index, recents.len());
                        }
                        let mut session = output.session(&time);
                        while let Some(next) = recents.pop() {
                            let lower = offsets[(next as usize) / peers];
                            let upper = offsets[((next as usize) / peers) + 1];
                            // println!("degree[{}]: {}", next, upper-lower);
                            for &target in &targets[lower..upper] {
                                session.give((target, next));
                            }
                        }

                    });
                }
            )
            .concat(&(0..1).map(|x| (x,x)).to_stream(scope))
            .connect_loop(handle);

        });

    });

}
<commit_msg>bfs bug<commit_after>extern crate time;
extern crate rand;
extern crate timely;
extern crate radix_sort;

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

use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::dataflow::channels::pact::Exchange;

fn main() {

    // command-line args: numbers of nodes and edges in the random graph.
    let nodes: usize = std::env::args().nth(1).unwrap().parse().unwrap();
    let edges: usize = std::env::args().nth(2).unwrap().parse().unwrap();

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

        let index = root.index();
        let peers = root.peers();

        let seed: &[_] = &[1, 2, 3, root.index()];
        let mut rng: StdRng = SeedableRng::from_seed(seed);
        let mut sorter = RadixSorter::new();

        let mut offsets = Vec::new();
        let mut targets = Vec::new();
        let mut recents = Vec::new();
        let mut parents = vec![u32::max_value(); nodes / peers];

        let start = time::precise_time_s();

        root.scoped(move |scope| {

            // generate part of a random graph.
            let graph = (0..edges / peers)
                .map(move |_| (rng.gen_range(0u32, nodes as u32), rng.gen_range(0u32, nodes as u32)))
                .to_stream(scope);

            // define a loop variable, for the (node, root) pairs.
            let (handle, stream) = scope.loop_variable(usize::max_value(), 1);

            // use the stream of edges
            graph.binary_notify(
                &stream,
                Exchange::new(|x: &(u32, u32)| x.0 as u64),
                Exchange::new(|x: &(u32, u32)| x.0 as u64),
                "BFS",
                vec![],
                move |input1, input2, output, notify| {

                    // receive edges, start to sort them
                    input1.for_each(|time, data| {
                        notify.notify_at(time);
                        for &edge in data.iter() { sorter.push(edge, &|x| x.0); }
                    });

                    // receive (node, root) pairs, note any new ones.
                    input2.for_each(|time, data| {
                        notify.notify_at(time);
                        for &(node, root) in data.iter() {
                            // println!("received: {:?}, parents[{}] = {}", (node, root), node, parents[node as usize]);
                            if parents[(node as usize) / peers] == u32::max_value() {
                                // println!("keeping: {:?}", node);
                                parents[(node as usize) / peers] = root;
                                recents.push(node);
                            }
                        }
                    });

                    notify.for_each(|time, _count| {

                        // maybe process the graph
                        if time.inner == 0 {

                            if index == 0 {
                                println!("{}:\tstarting sort", time::precise_time_s() - start);
                            }

                            // construct the graph
                            offsets.push(0);
                            let mut prev_node = 0;
                            let sorted = sorter.finish(&|x| x.0);
                            for buffer in &sorted {
                                for &(node, edge) in buffer {
                                    let node = node / peers as u32;
                                    while prev_node < node {
                                        prev_node += 1;
                                        offsets.push(targets.len())
                                    }
                                    targets.push(edge);
                                }
                            }
                            offsets.push(targets.len())

                        }

                        if index == 0 {
                            println!("{}:\tworker: {}, todo.len(): {}", time::precise_time_s() - start, index, recents.len());
                        }
                        let mut session = output.session(&time);
                        while let Some(next) = recents.pop() {
                            let lower = offsets[(next as usize) / peers];
                            let upper = offsets[((next as usize) / peers) + 1];
                            // println!("degree[{}]: {}", next, upper-lower);
                            for &target in &targets[lower..upper] {
                                session.give((target, next));
                            }
                        }

                    });
                }
            )
            .concat(&(0..1).map(|x| (x,x)).to_stream(scope))
            .connect_loop(handle);

        });

    });

}
<|endoftext|>


<commit_before>// Translation of the following Go program into Rust:
//
// ```
// func main() {
//     people := []string{"Anna", "Bob", "Cody", "Dave", "Eva"}
//     match := make(chan string, 1) // Make room for one unmatched send.
//     wg := new(sync.WaitGroup)
//     for _, name := range people {
//         wg.Add(1)
//         go Seek(name, match, wg)
//     }
//     wg.Wait()
//     select {
//     case name := <-match:
//         fmt.Printf("No one received %s’s message.\n", name)
//     default:
//         // There was no pending send operation.
//     }
// }
//
// // Seek either sends or receives, whichever possible, a name on the match
// // channel and notifies the wait group when done.
// func Seek(name string, match chan string, wg *sync.WaitGroup) {
//     select {
//     case peer := <-match:
//         fmt.Printf("%s received a message from %s.\n", name, peer)
//     case match <- name:
//         // Wait for someone to receive my message.
//     }
//     wg.Done()
// }
// ```
//
// Source: https://www.nada.kth.se/~snilsson/concurrency/#Match

extern crate crossbeam;
extern crate channel;

use channel::{select, Receiver, Sender};

fn main() {
    let people = vec!["Anna", "Bob", "Cody", "Dave", "Eva"];
    let (tx, rx) = channel::bounded(1); // Make room for one unmatched send.
    let (tx, rx) = (&tx, &rx);

    crossbeam::scope(|s| for name in people {
        s.spawn(move || seek(name, tx, rx));
    });

    if let Ok(name) = rx.try_recv() {
        println!("No one received {}’s message.", name);
    }
}

// Either sends or receives, whichever possible, a name on the channel.
fn seek<'a>(name: &'a str, tx: &Sender<&'a str>, rx: &Receiver<&'a str>) {
    loop {
        if let Ok(peer) = rx.select() {
            println!("{} received a message from {}.", name, peer);
            break;
        }
        if let Ok(_) = tx.select(name) {
            // Wait for someone to receive my message.
            break;
        }
    }
}
<commit_msg>Cosmetic fixes<commit_after>// Translation of the following Go program into Rust:
//
// ```
// func main() {
//     people := []string{"Anna", "Bob", "Cody", "Dave", "Eva"}
//     match := make(chan string, 1) // Make room for one unmatched send.
//     wg := new(sync.WaitGroup)
//     for _, name := range people {
//         wg.Add(1)
//         go Seek(name, match, wg)
//     }
//     wg.Wait()
//     select {
//     case name := <-match:
//         fmt.Printf("No one received %s’s message.\n", name)
//     default:
//         // There was no pending send operation.
//     }
// }
//
// // Seek either sends or receives, whichever possible, a name on the match
// // channel and notifies the wait group when done.
// func Seek(name string, match chan string, wg *sync.WaitGroup) {
//     select {
//     case peer := <-match:
//         fmt.Printf("%s received a message from %s.\n", name, peer)
//     case match <- name:
//         // Wait for someone to receive my message.
//     }
//     wg.Done()
// }
// ```
//
// Source: https://www.nada.kth.se/~snilsson/concurrency/#Match

extern crate crossbeam;
extern crate channel;

use channel::{select, Receiver, Sender};

fn main() {
    let people = vec!["Anna", "Bob", "Cody", "Dave", "Eva"];
    let (tx, rx) = channel::bounded(1); // Make room for one unmatched send.
    let (tx, rx) = (&tx, &rx);

    crossbeam::scope(|s| for name in people {
        s.spawn(move || seek(name, tx, rx));
    });

    if let Ok(name) = rx.try_recv() {
        println!("No one received {}’s message.", name);
    }
}

// Either sends or receives, whichever possible, a name on the channel.
fn seek<'a>(name: &'a str, tx: &Sender<&'a str>, rx: &Receiver<&'a str>) {
    loop {
        if let Ok(peer) = rx.select() {
            println!("{} received a message from {}.", name, peer);
            break;
        }
        if let Ok(()) = tx.select(name) {
            // Wait for someone to receive my message.
            break;
        }
    }
}
<|endoftext|>


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

//! A private parser implementation of IPv4, IPv6, and socket addresses.
//!
//! This module is "publicly exported" through the `FromStr` implementations
//! below.

use prelude::v1::*;

use str::FromStr;
use net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};

struct Parser<'a> {
    // parsing as ASCII, so can use byte array
    s: &'a [u8],
    pos: usize,
}

enum IpAddr {
    V4(Ipv4Addr),
    V6(Ipv6Addr),
}

impl<'a> Parser<'a> {
    fn new(s: &'a str) -> Parser<'a> {
        Parser {
            s: s.as_bytes(),
            pos: 0,
        }
    }

    fn is_eof(&self) -> bool {
        self.pos == self.s.len()
    }

    // Commit only if parser returns Some
    fn read_atomically<T, F>(&mut self, cb: F) -> Option<T> where
        F: FnOnce(&mut Parser) -> Option<T>,
    {
        let pos = self.pos;
        let r = cb(self);
        if r.is_none() {
            self.pos = pos;
        }
        r
    }

    // Commit only if parser read till EOF
    fn read_till_eof<T, F>(&mut self, cb: F) -> Option<T> where
        F: FnOnce(&mut Parser) -> Option<T>,
    {
        self.read_atomically(move |p| {
            match cb(p) {
                Some(x) => if p.is_eof() {Some(x)} else {None},
                None => None,
            }
        })
    }

    // Return result of first successful parser
    fn read_or<T>(&mut self, parsers: &mut [Box<FnMut(&mut Parser) -> Option<T>>])
               -> Option<T> {
        for pf in parsers.iter_mut() {
            match self.read_atomically(|p: &mut Parser| pf(p)) {
                Some(r) => return Some(r),
                None => {}
            }
        }
        None
    }

    // Apply 3 parsers sequentially
    fn read_seq_3<A, B, C, PA, PB, PC>(&mut self,
                                       pa: PA,
                                       pb: PB,
                                       pc: PC)
                                       -> Option<(A, B, C)> where
        PA: FnOnce(&mut Parser) -> Option<A>,
        PB: FnOnce(&mut Parser) -> Option<B>,
        PC: FnOnce(&mut Parser) -> Option<C>,
    {
        self.read_atomically(move |p| {
            let a = pa(p);
            let b = if a.is_some() { pb(p) } else { None };
            let c = if b.is_some() { pc(p) } else { None };
            match (a, b, c) {
                (Some(a), Some(b), Some(c)) => Some((a, b, c)),
                _ => None
            }
        })
    }

    // Read next char
    fn read_char(&mut self) -> Option<char> {
        if self.is_eof() {
            None
        } else {
            let r = self.s[self.pos] as char;
            self.pos += 1;
            Some(r)
        }
    }

    // Return char and advance iff next char is equal to requested
    fn read_given_char(&mut self, c: char) -> Option<char> {
        self.read_atomically(|p| {
            match p.read_char() {
                Some(next) if next == c => Some(next),
                _ => None,
            }
        })
    }

    // Read digit
    fn read_digit(&mut self, radix: u8) -> Option<u8> {
        fn parse_digit(c: char, radix: u8) -> Option<u8> {
            let c = c as u8;
            // assuming radix is either 10 or 16
            if c >= b'0' && c <= b'9' {
                Some(c - b'0')
            } else if radix > 10 && c >= b'a' && c < b'a' + (radix - 10) {
                Some(c - b'a' + 10)
            } else if radix > 10 && c >= b'A' && c < b'A' + (radix - 10) {
                Some(c - b'A' + 10)
            } else {
                None
            }
        }

        self.read_atomically(|p| {
            p.read_char().and_then(|c| parse_digit(c, radix))
        })
    }

    fn read_number_impl(&mut self, radix: u8, max_digits: u32, upto: u32) -> Option<u32> {
        let mut r = 0;
        let mut digit_count = 0;
        loop {
            match self.read_digit(radix) {
                Some(d) => {
                    r = r * (radix as u32) + (d as u32);
                    digit_count += 1;
                    if digit_count > max_digits || r >= upto {
                        return None
                    }
                }
                None => {
                    if digit_count == 0 {
                        return None
                    } else {
                        return Some(r)
                    }
                }
            };
        }
    }

    // Read number, failing if max_digits of number value exceeded
    fn read_number(&mut self, radix: u8, max_digits: u32, upto: u32) -> Option<u32> {
        self.read_atomically(|p| p.read_number_impl(radix, max_digits, upto))
    }

    fn read_ipv4_addr_impl(&mut self) -> Option<Ipv4Addr> {
        let mut bs = [0; 4];
        let mut i = 0;
        while i < 4 {
            if i != 0 && self.read_given_char('.').is_none() {
                return None;
            }

            let octet = self.read_number(10, 3, 0x100).map(|n| n as u8);
            match octet {
                Some(d) => bs[i] = d,
                None => return None,
            };
            i += 1;
        }
        Some(Ipv4Addr::new(bs[0], bs[1], bs[2], bs[3]))
    }

    // Read IPv4 address
    fn read_ipv4_addr(&mut self) -> Option<Ipv4Addr> {
        self.read_atomically(|p| p.read_ipv4_addr_impl())
    }

    fn read_ipv6_addr_impl(&mut self) -> Option<Ipv6Addr> {
        fn ipv6_addr_from_head_tail(head: &[u16], tail: &[u16]) -> Ipv6Addr {
            assert!(head.len() + tail.len() <= 8);
            let mut gs = [0; 8];
            gs.clone_from_slice(head);
            gs[(8 - tail.len()) .. 8].clone_from_slice(tail);
            Ipv6Addr::new(gs[0], gs[1], gs[2], gs[3], gs[4], gs[5], gs[6], gs[7])
        }

        fn read_groups(p: &mut Parser, groups: &mut [u16; 8], limit: usize)
                       -> (usize, bool) {
            let mut i = 0;
            while i < limit {
                if i < limit - 1 {
                    let ipv4 = p.read_atomically(|p| {
                        if i == 0 || p.read_given_char(':').is_some() {
                            p.read_ipv4_addr()
                        } else {
                            None
                        }
                    });
                    if let Some(v4_addr) = ipv4 {
                        let octets = v4_addr.octets();
                        groups[i + 0] = ((octets[0] as u16) << 8) | (octets[1] as u16);
                        groups[i + 1] = ((octets[2] as u16) << 8) | (octets[3] as u16);
                        return (i + 2, true);
                    }
                }

                let group = p.read_atomically(|p| {
                    if i == 0 || p.read_given_char(':').is_some() {
                        p.read_number(16, 4, 0x10000).map(|n| n as u16)
                    } else {
                        None
                    }
                });
                match group {
                    Some(g) => groups[i] = g,
                    None => return (i, false)
                }
                i += 1;
            }
            (i, false)
        }

        let mut head = [0; 8];
        let (head_size, head_ipv4) = read_groups(self, &mut head, 8);

        if head_size == 8 {
            return Some(Ipv6Addr::new(
                head[0], head[1], head[2], head[3],
                head[4], head[5], head[6], head[7]))
        }

        // IPv4 part is not allowed before `::`
        if head_ipv4 {
            return None
        }

        // read `::` if previous code parsed less than 8 groups
        if !self.read_given_char(':').is_some() || !self.read_given_char(':').is_some() {
            return None;
        }

        let mut tail = [0; 8];
        let (tail_size, _) = read_groups(self, &mut tail, 8 - head_size);
        Some(ipv6_addr_from_head_tail(&head[..head_size], &tail[..tail_size]))
    }

    fn read_ipv6_addr(&mut self) -> Option<Ipv6Addr> {
        self.read_atomically(|p| p.read_ipv6_addr_impl())
    }

    fn read_ip_addr(&mut self) -> Option<IpAddr> {
        let ipv4_addr = |p: &mut Parser| p.read_ipv4_addr().map(|v4| IpAddr::V4(v4));
        let ipv6_addr = |p: &mut Parser| p.read_ipv6_addr().map(|v6| IpAddr::V6(v6));
        self.read_or(&mut [Box::new(ipv4_addr), Box::new(ipv6_addr)])
    }

    fn read_socket_addr(&mut self) -> Option<SocketAddr> {
        let ip_addr = |p: &mut Parser| {
            let ipv4_p = |p: &mut Parser| p.read_ip_addr();
            let ipv6_p = |p: &mut Parser| {
                let open_br = |p: &mut Parser| p.read_given_char('[');
                let ip_addr = |p: &mut Parser| p.read_ipv6_addr();
                let clos_br = |p: &mut Parser| p.read_given_char(']');
                p.read_seq_3::<char, Ipv6Addr, char, _, _, _>(open_br, ip_addr, clos_br)
                        .map(|t| match t { (_, ip, _) => IpAddr::V6(ip) })
            };
            p.read_or(&mut [Box::new(ipv4_p), Box::new(ipv6_p)])
        };
        let colon = |p: &mut Parser| p.read_given_char(':');
        let port  = |p: &mut Parser| p.read_number(10, 5, 0x10000).map(|n| n as u16);

        // host, colon, port
        self.read_seq_3(ip_addr, colon, port).map(|t| {
            let (ip, _, port): (IpAddr, char, u16) = t;
            match ip {
                IpAddr::V4(ip) => SocketAddr::V4(SocketAddrV4::new(ip, port)),
                IpAddr::V6(ip) => SocketAddr::V6(SocketAddrV6::new(ip, port, 0, 0)),
            }
        })
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl FromStr for Ipv4Addr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<Ipv4Addr, AddrParseError> {
        match Parser::new(s).read_till_eof(|p| p.read_ipv4_addr()) {
            Some(s) => Ok(s),
            None => Err(AddrParseError(()))
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl FromStr for Ipv6Addr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<Ipv6Addr, AddrParseError> {
        match Parser::new(s).read_till_eof(|p| p.read_ipv6_addr()) {
            Some(s) => Ok(s),
            None => Err(AddrParseError(()))
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl FromStr for SocketAddr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<SocketAddr, AddrParseError> {
        match Parser::new(s).read_till_eof(|p| p.read_socket_addr()) {
            Some(s) => Ok(s),
            None => Err(AddrParseError(())),
        }
    }
}

/// An error returned when parsing an IP address or a socket address.
#[stable(feature = "rust1", since = "1.0.0")]
#[derive(Debug, Clone, PartialEq)]
pub struct AddrParseError(());
<commit_msg>impl `FromStr` for `IpAddr`<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A private parser implementation of IPv4, IPv6, and socket addresses.
//!
//! This module is "publicly exported" through the `FromStr` implementations
//! below.

use prelude::v1::*;

use str::FromStr;
use net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};

struct Parser<'a> {
    // parsing as ASCII, so can use byte array
    s: &'a [u8],
    pos: usize,
}

impl<'a> Parser<'a> {
    fn new(s: &'a str) -> Parser<'a> {
        Parser {
            s: s.as_bytes(),
            pos: 0,
        }
    }

    fn is_eof(&self) -> bool {
        self.pos == self.s.len()
    }

    // Commit only if parser returns Some
    fn read_atomically<T, F>(&mut self, cb: F) -> Option<T> where
        F: FnOnce(&mut Parser) -> Option<T>,
    {
        let pos = self.pos;
        let r = cb(self);
        if r.is_none() {
            self.pos = pos;
        }
        r
    }

    // Commit only if parser read till EOF
    fn read_till_eof<T, F>(&mut self, cb: F) -> Option<T> where
        F: FnOnce(&mut Parser) -> Option<T>,
    {
        self.read_atomically(move |p| {
            match cb(p) {
                Some(x) => if p.is_eof() {Some(x)} else {None},
                None => None,
            }
        })
    }

    // Return result of first successful parser
    fn read_or<T>(&mut self, parsers: &mut [Box<FnMut(&mut Parser) -> Option<T>>])
               -> Option<T> {
        for pf in parsers.iter_mut() {
            match self.read_atomically(|p: &mut Parser| pf(p)) {
                Some(r) => return Some(r),
                None => {}
            }
        }
        None
    }

    // Apply 3 parsers sequentially
    fn read_seq_3<A, B, C, PA, PB, PC>(&mut self,
                                       pa: PA,
                                       pb: PB,
                                       pc: PC)
                                       -> Option<(A, B, C)> where
        PA: FnOnce(&mut Parser) -> Option<A>,
        PB: FnOnce(&mut Parser) -> Option<B>,
        PC: FnOnce(&mut Parser) -> Option<C>,
    {
        self.read_atomically(move |p| {
            let a = pa(p);
            let b = if a.is_some() { pb(p) } else { None };
            let c = if b.is_some() { pc(p) } else { None };
            match (a, b, c) {
                (Some(a), Some(b), Some(c)) => Some((a, b, c)),
                _ => None
            }
        })
    }

    // Read next char
    fn read_char(&mut self) -> Option<char> {
        if self.is_eof() {
            None
        } else {
            let r = self.s[self.pos] as char;
            self.pos += 1;
            Some(r)
        }
    }

    // Return char and advance iff next char is equal to requested
    fn read_given_char(&mut self, c: char) -> Option<char> {
        self.read_atomically(|p| {
            match p.read_char() {
                Some(next) if next == c => Some(next),
                _ => None,
            }
        })
    }

    // Read digit
    fn read_digit(&mut self, radix: u8) -> Option<u8> {
        fn parse_digit(c: char, radix: u8) -> Option<u8> {
            let c = c as u8;
            // assuming radix is either 10 or 16
            if c >= b'0' && c <= b'9' {
                Some(c - b'0')
            } else if radix > 10 && c >= b'a' && c < b'a' + (radix - 10) {
                Some(c - b'a' + 10)
            } else if radix > 10 && c >= b'A' && c < b'A' + (radix - 10) {
                Some(c - b'A' + 10)
            } else {
                None
            }
        }

        self.read_atomically(|p| {
            p.read_char().and_then(|c| parse_digit(c, radix))
        })
    }

    fn read_number_impl(&mut self, radix: u8, max_digits: u32, upto: u32) -> Option<u32> {
        let mut r = 0;
        let mut digit_count = 0;
        loop {
            match self.read_digit(radix) {
                Some(d) => {
                    r = r * (radix as u32) + (d as u32);
                    digit_count += 1;
                    if digit_count > max_digits || r >= upto {
                        return None
                    }
                }
                None => {
                    if digit_count == 0 {
                        return None
                    } else {
                        return Some(r)
                    }
                }
            };
        }
    }

    // Read number, failing if max_digits of number value exceeded
    fn read_number(&mut self, radix: u8, max_digits: u32, upto: u32) -> Option<u32> {
        self.read_atomically(|p| p.read_number_impl(radix, max_digits, upto))
    }

    fn read_ipv4_addr_impl(&mut self) -> Option<Ipv4Addr> {
        let mut bs = [0; 4];
        let mut i = 0;
        while i < 4 {
            if i != 0 && self.read_given_char('.').is_none() {
                return None;
            }

            let octet = self.read_number(10, 3, 0x100).map(|n| n as u8);
            match octet {
                Some(d) => bs[i] = d,
                None => return None,
            };
            i += 1;
        }
        Some(Ipv4Addr::new(bs[0], bs[1], bs[2], bs[3]))
    }

    // Read IPv4 address
    fn read_ipv4_addr(&mut self) -> Option<Ipv4Addr> {
        self.read_atomically(|p| p.read_ipv4_addr_impl())
    }

    fn read_ipv6_addr_impl(&mut self) -> Option<Ipv6Addr> {
        fn ipv6_addr_from_head_tail(head: &[u16], tail: &[u16]) -> Ipv6Addr {
            assert!(head.len() + tail.len() <= 8);
            let mut gs = [0; 8];
            gs.clone_from_slice(head);
            gs[(8 - tail.len()) .. 8].clone_from_slice(tail);
            Ipv6Addr::new(gs[0], gs[1], gs[2], gs[3], gs[4], gs[5], gs[6], gs[7])
        }

        fn read_groups(p: &mut Parser, groups: &mut [u16; 8], limit: usize)
                       -> (usize, bool) {
            let mut i = 0;
            while i < limit {
                if i < limit - 1 {
                    let ipv4 = p.read_atomically(|p| {
                        if i == 0 || p.read_given_char(':').is_some() {
                            p.read_ipv4_addr()
                        } else {
                            None
                        }
                    });
                    if let Some(v4_addr) = ipv4 {
                        let octets = v4_addr.octets();
                        groups[i + 0] = ((octets[0] as u16) << 8) | (octets[1] as u16);
                        groups[i + 1] = ((octets[2] as u16) << 8) | (octets[3] as u16);
                        return (i + 2, true);
                    }
                }

                let group = p.read_atomically(|p| {
                    if i == 0 || p.read_given_char(':').is_some() {
                        p.read_number(16, 4, 0x10000).map(|n| n as u16)
                    } else {
                        None
                    }
                });
                match group {
                    Some(g) => groups[i] = g,
                    None => return (i, false)
                }
                i += 1;
            }
            (i, false)
        }

        let mut head = [0; 8];
        let (head_size, head_ipv4) = read_groups(self, &mut head, 8);

        if head_size == 8 {
            return Some(Ipv6Addr::new(
                head[0], head[1], head[2], head[3],
                head[4], head[5], head[6], head[7]))
        }

        // IPv4 part is not allowed before `::`
        if head_ipv4 {
            return None
        }

        // read `::` if previous code parsed less than 8 groups
        if !self.read_given_char(':').is_some() || !self.read_given_char(':').is_some() {
            return None;
        }

        let mut tail = [0; 8];
        let (tail_size, _) = read_groups(self, &mut tail, 8 - head_size);
        Some(ipv6_addr_from_head_tail(&head[..head_size], &tail[..tail_size]))
    }

    fn read_ipv6_addr(&mut self) -> Option<Ipv6Addr> {
        self.read_atomically(|p| p.read_ipv6_addr_impl())
    }

    fn read_ip_addr(&mut self) -> Option<IpAddr> {
        let ipv4_addr = |p: &mut Parser| p.read_ipv4_addr().map(|v4| IpAddr::V4(v4));
        let ipv6_addr = |p: &mut Parser| p.read_ipv6_addr().map(|v6| IpAddr::V6(v6));
        self.read_or(&mut [Box::new(ipv4_addr), Box::new(ipv6_addr)])
    }

    fn read_socket_addr(&mut self) -> Option<SocketAddr> {
        let ip_addr = |p: &mut Parser| {
            let ipv4_p = |p: &mut Parser| p.read_ip_addr();
            let ipv6_p = |p: &mut Parser| {
                let open_br = |p: &mut Parser| p.read_given_char('[');
                let ip_addr = |p: &mut Parser| p.read_ipv6_addr();
                let clos_br = |p: &mut Parser| p.read_given_char(']');
                p.read_seq_3::<char, Ipv6Addr, char, _, _, _>(open_br, ip_addr, clos_br)
                        .map(|t| match t { (_, ip, _) => IpAddr::V6(ip) })
            };
            p.read_or(&mut [Box::new(ipv4_p), Box::new(ipv6_p)])
        };
        let colon = |p: &mut Parser| p.read_given_char(':');
        let port  = |p: &mut Parser| p.read_number(10, 5, 0x10000).map(|n| n as u16);

        // host, colon, port
        self.read_seq_3(ip_addr, colon, port).map(|t| {
            let (ip, _, port): (IpAddr, char, u16) = t;
            match ip {
                IpAddr::V4(ip) => SocketAddr::V4(SocketAddrV4::new(ip, port)),
                IpAddr::V6(ip) => SocketAddr::V6(SocketAddrV6::new(ip, port, 0, 0)),
            }
        })
    }
}

#[unstable(feature = "ip_addr", reason = "recent addition")]
impl FromStr for IpAddr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<IpAddr, AddrParseError> {
        match Parser::new(s).read_till_eof(|p| p.read_ip_addr()) {
            Some(s) => Ok(s),
            None => Err(AddrParseError(()))
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl FromStr for Ipv4Addr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<Ipv4Addr, AddrParseError> {
        match Parser::new(s).read_till_eof(|p| p.read_ipv4_addr()) {
            Some(s) => Ok(s),
            None => Err(AddrParseError(()))
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl FromStr for Ipv6Addr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<Ipv6Addr, AddrParseError> {
        match Parser::new(s).read_till_eof(|p| p.read_ipv6_addr()) {
            Some(s) => Ok(s),
            None => Err(AddrParseError(()))
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl FromStr for SocketAddr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<SocketAddr, AddrParseError> {
        match Parser::new(s).read_till_eof(|p| p.read_socket_addr()) {
            Some(s) => Ok(s),
            None => Err(AddrParseError(())),
        }
    }
}

/// An error returned when parsing an IP address or a socket address.
#[stable(feature = "rust1", since = "1.0.0")]
#[derive(Debug, Clone, PartialEq)]
pub struct AddrParseError(());
<|endoftext|>


<commit_before>use std::error::Error;
use std::fs::{File, OpenOptions};
use std::io::BufReader;
use std::io::{prelude::*, SeekFrom};
use std::path::{Path, PathBuf};
use std::process;

use libsignify::{
    consts::DEFAULT_KDF_ROUNDS, Codeable, NewKeyOpts, PrivateKey, PublicKey, Signature,
};

use clap::Parser;

#[derive(Parser)]
#[clap(
    name = "signify",
    override_usage = r#"signify -h
    signify -G [-n] [-c <comment>] -p <pubkey> -s <seckey>
    signify -S [-e] [-x <sigfile>] -s <seckey> -m <message>
    signify -V [-e] [-x <sigfile>] -p <pubkey> -m <message>"#
)]
struct Args {
    #[clap(short = 'G', help = "Generate a new keypair.")]
    generate: bool,
    #[clap(short = 'S', help = "Sign the specified message file.")]
    sign: bool,
    #[clap(short = 'V', help = "Verify a message.")]
    verify: bool,

    #[clap(
        short = 'p',
        help = "Public key produced by -G, and used by -V to check a signature."
    )]
    pubkey: Option<PathBuf>,

    #[clap(
        short = 's',
        help = "Secret (private) key produced by -G, and used by -S to sign a message."
    )]
    seckey: Option<PathBuf>,

    #[clap(
        short = 'n',
        help = "Do not ask for a passphrase during key generation. Otherwise, signify will prompt the user for a passphrase to protect the secret key."
    )]
    skip_key_encryption: bool,

    #[clap(short = 'm')]
    message_path: Option<String>,

    #[clap(
        short = 'e',
        help = "When signing, embed the message after the signature. When verifying, extract the message from the signature."
    )]
    embed_message: bool,

    #[clap(
        short = 'x',
        help = "The signature file to create or verify. The default is <message>.sig."
    )]
    signature_path: Option<String>,

    #[clap(
        short = 'c',
        help = "Specify the comment to be added during key generation"
    )]
    comment: Option<String>,
}

fn write_base64_file<C: Codeable>(
    file: &mut File,
    comment: &str,
    data: &C,
) -> Result<(), Box<dyn Error>> {
    let contents = data.to_file_encoding(comment);
    file.write_all(&contents)?;

    Ok(())
}

fn read_base64_file<C: Codeable, R: Read>(
    reader: &mut BufReader<R>,
) -> Result<(C, u64), Box<dyn Error>> {
    let mut contents = String::with_capacity(1024);
    // Optimization: Read the two lines that have the comment and well structured data
    // instead of the entire file in case the message was large and embedded.
    reader.read_line(&mut contents)?;
    reader.read_line(&mut contents)?;

    let read = C::from_base64(&contents)?;
    Ok(read)
}

fn verify(
    pubkey_path: &Path,
    msg_path: &str,
    signature_path: Option<String>,
    embed: bool,
) -> Result<(), Box<dyn Error>> {
    let mut pubkey_file = BufReader::new(File::open(pubkey_path)?);
    let public_key: PublicKey = read_base64_file(&mut pubkey_file)?.0;

    let signature_path = match signature_path {
        Some(path) => path,
        None => format!("{}.sig", msg_path),
    };

    let mut sig_data = BufReader::new(File::open(&signature_path)?);

    let (signature, msg_data_pos): (Signature, u64) = read_base64_file(&mut sig_data)?;

    let mut msg = vec![];

    if embed {
        sig_data.seek(SeekFrom::Start(msg_data_pos))?;
        sig_data.read_to_end(&mut msg)?;
    } else {
        let mut msg_file = File::open(msg_path)?;
        msg_file.read_to_end(&mut msg)?;
    }

    public_key
        .verify(&msg, &signature)
        .map(|_| {
            println!("Signature Verified");
        })
        .map_err(Into::into)
}

fn sign(
    private_key_path: &Path,
    msg_path: &str,
    signature_path: Option<String>,
    embed: bool,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut secret_key = BufReader::new(File::open(private_key_path)?);
    let mut secret_key: PrivateKey = read_base64_file(&mut secret_key)?.0;

    if secret_key.is_encrypted() {
        let passphrase = read_passphrase(false)?;
        secret_key.decrypt_with_password(&passphrase)?;
    }

    let mut msg_file = File::open(&msg_path)?;
    let mut msg = vec![];
    msg_file.read_to_end(&mut msg)?;

    let signature_path = match signature_path {
        Some(path) => path,
        None => format!("{}.sig", msg_path),
    };

    let sig = secret_key.sign(&msg);

    let sig_comment = "signature from signify secret key";

    let mut file = OpenOptions::new()
        .write(true)
        .create_new(true)
        .open(&signature_path)?;

    write_base64_file(&mut file, sig_comment, &sig)?;

    if embed {
        file.write_all(&msg)?;
    }

    Ok(())
}

fn read_passphrase(confirm: bool) -> Result<String, Box<dyn std::error::Error>> {
    let passphrase = rpassword::prompt_password_stdout("passphrase: ")?;

    if confirm {
        let confirm_passphrase = rpassword::prompt_password_stdout("confirm passphrase: ")?;

        if passphrase != confirm_passphrase {
            return Err("passwords don't match".into());
        }
    }

    Ok(passphrase)
}

fn generate(
    pubkey_path: &Path,
    privkey_path: &Path,
    comment: Option<&str>,
    kdfrounds: Option<u32>,
) -> Result<(), Box<dyn std::error::Error>> {
    let comment = comment.unwrap_or("signify");

    let derivation_info = match kdfrounds {
        Some(kdf_rounds) => {
            let passphrase = read_passphrase(true)?;
            NewKeyOpts::Encrypted {
                passphrase,
                kdf_rounds,
            }
        }
        None => NewKeyOpts::NoEncryption,
    };

    // Store the private key
    let mut rng = rand_core::OsRng {};
    let private_key = PrivateKey::generate(&mut rng, derivation_info)?;

    let priv_comment = format!("{} secret key", comment);
    let mut file = OpenOptions::new()
        .write(true)
        .create_new(true)
        .open(privkey_path)?;

    write_base64_file(&mut file, &priv_comment, &private_key)?;

    // Store public key
    let public_key = private_key.public();

    let pub_comment = format!("{} public key", comment);
    let mut file = OpenOptions::new()
        .write(true)
        .create_new(true)
        .open(pubkey_path)?;

    write_base64_file(&mut file, &pub_comment, &public_key)?;

    Ok(())
}

fn human<T>(res: Result<T, Box<dyn std::error::Error>>) -> T {
    match res {
        Ok(val) => val,
        Err(e) => {
            println!("error: {}", e);
            process::exit(1);
        }
    }
}

fn unwrap_path<T>(kind: &'static str, path: Option<T>) -> T {
    match path {
        Some(p) => p,
        None => {
            println!("missing path to {}", kind);
            process::exit(1)
        }
    }
}

fn main() {
    let args = Args::parse();

    if args.verify {
        let public_key = unwrap_path("pubkey", args.pubkey);
        let message = unwrap_path("message", args.message_path);

        human(verify(
            &public_key,
            &message,
            args.signature_path,
            args.embed_message,
        ));

        return;
    }

    if args.generate {
        let public_key = unwrap_path("pubkey", args.pubkey);
        let private_key = unwrap_path("seckey", args.seckey);
        let rounds = if args.skip_key_encryption {
            None
        } else {
            Some(DEFAULT_KDF_ROUNDS)
        };

        human(generate(
            &public_key,
            &private_key,
            args.comment.as_deref(),
            rounds,
        ));
        return;
    }

    if args.sign {
        let private_key = unwrap_path("seckey", args.seckey);
        let msg_path = unwrap_path("message", args.message_path);

        human(sign(
            &private_key,
            &msg_path,
            args.signature_path,
            args.embed_message,
        ));
    }
}
<commit_msg>Print command usage when no other command was specified<commit_after>use std::error::Error;
use std::fs::{File, OpenOptions};
use std::io::BufReader;
use std::io::{prelude::*, SeekFrom};
use std::path::{Path, PathBuf};
use std::process;

use libsignify::{
    consts::DEFAULT_KDF_ROUNDS, Codeable, NewKeyOpts, PrivateKey, PublicKey, Signature,
};

use clap::{IntoApp, Parser};

#[derive(Parser)]
#[clap(
    name = "signify",
    override_usage = r#"signify -h
    signify -G [-n] [-c <comment>] -p <pubkey> -s <seckey>
    signify -S [-e] [-x <sigfile>] -s <seckey> -m <message>
    signify -V [-e] [-x <sigfile>] -p <pubkey> -m <message>"#
)]
struct Args {
    #[clap(short = 'G', help = "Generate a new keypair.")]
    generate: bool,
    #[clap(short = 'S', help = "Sign the specified message file.")]
    sign: bool,
    #[clap(short = 'V', help = "Verify a message.")]
    verify: bool,

    #[clap(
        short = 'p',
        help = "Public key produced by -G, and used by -V to check a signature."
    )]
    pubkey: Option<PathBuf>,

    #[clap(
        short = 's',
        help = "Secret (private) key produced by -G, and used by -S to sign a message."
    )]
    seckey: Option<PathBuf>,

    #[clap(
        short = 'n',
        help = "Do not ask for a passphrase during key generation. Otherwise, signify will prompt the user for a passphrase to protect the secret key."
    )]
    skip_key_encryption: bool,

    #[clap(short = 'm')]
    message_path: Option<String>,

    #[clap(
        short = 'e',
        help = "When signing, embed the message after the signature. When verifying, extract the message from the signature."
    )]
    embed_message: bool,

    #[clap(
        short = 'x',
        help = "The signature file to create or verify. The default is <message>.sig."
    )]
    signature_path: Option<String>,

    #[clap(
        short = 'c',
        help = "Specify the comment to be added during key generation"
    )]
    comment: Option<String>,
}

fn write_base64_file<C: Codeable>(
    file: &mut File,
    comment: &str,
    data: &C,
) -> Result<(), Box<dyn Error>> {
    let contents = data.to_file_encoding(comment);
    file.write_all(&contents)?;

    Ok(())
}

fn read_base64_file<C: Codeable, R: Read>(
    reader: &mut BufReader<R>,
) -> Result<(C, u64), Box<dyn Error>> {
    let mut contents = String::with_capacity(1024);
    // Optimization: Read the two lines that have the comment and well structured data
    // instead of the entire file in case the message was large and embedded.
    reader.read_line(&mut contents)?;
    reader.read_line(&mut contents)?;

    let read = C::from_base64(&contents)?;
    Ok(read)
}

fn verify(
    pubkey_path: &Path,
    msg_path: &str,
    signature_path: Option<String>,
    embed: bool,
) -> Result<(), Box<dyn Error>> {
    let mut pubkey_file = BufReader::new(File::open(pubkey_path)?);
    let public_key: PublicKey = read_base64_file(&mut pubkey_file)?.0;

    let signature_path = match signature_path {
        Some(path) => path,
        None => format!("{}.sig", msg_path),
    };

    let mut sig_data = BufReader::new(File::open(&signature_path)?);

    let (signature, msg_data_pos): (Signature, u64) = read_base64_file(&mut sig_data)?;

    let mut msg = vec![];

    if embed {
        sig_data.seek(SeekFrom::Start(msg_data_pos))?;
        sig_data.read_to_end(&mut msg)?;
    } else {
        let mut msg_file = File::open(msg_path)?;
        msg_file.read_to_end(&mut msg)?;
    }

    public_key
        .verify(&msg, &signature)
        .map(|_| {
            println!("Signature Verified");
        })
        .map_err(Into::into)
}

fn sign(
    private_key_path: &Path,
    msg_path: &str,
    signature_path: Option<String>,
    embed: bool,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut secret_key = BufReader::new(File::open(private_key_path)?);
    let mut secret_key: PrivateKey = read_base64_file(&mut secret_key)?.0;

    if secret_key.is_encrypted() {
        let passphrase = read_passphrase(false)?;
        secret_key.decrypt_with_password(&passphrase)?;
    }

    let mut msg_file = File::open(&msg_path)?;
    let mut msg = vec![];
    msg_file.read_to_end(&mut msg)?;

    let signature_path = match signature_path {
        Some(path) => path,
        None => format!("{}.sig", msg_path),
    };

    let sig = secret_key.sign(&msg);

    let sig_comment = "signature from signify secret key";

    let mut file = OpenOptions::new()
        .write(true)
        .create_new(true)
        .open(&signature_path)?;

    write_base64_file(&mut file, sig_comment, &sig)?;

    if embed {
        file.write_all(&msg)?;
    }

    Ok(())
}

fn read_passphrase(confirm: bool) -> Result<String, Box<dyn std::error::Error>> {
    let passphrase = rpassword::prompt_password_stdout("passphrase: ")?;

    if confirm {
        let confirm_passphrase = rpassword::prompt_password_stdout("confirm passphrase: ")?;

        if passphrase != confirm_passphrase {
            return Err("passwords don't match".into());
        }
    }

    Ok(passphrase)
}

fn generate(
    pubkey_path: &Path,
    privkey_path: &Path,
    comment: Option<&str>,
    kdfrounds: Option<u32>,
) -> Result<(), Box<dyn std::error::Error>> {
    let comment = comment.unwrap_or("signify");

    let derivation_info = match kdfrounds {
        Some(kdf_rounds) => {
            let passphrase = read_passphrase(true)?;
            NewKeyOpts::Encrypted {
                passphrase,
                kdf_rounds,
            }
        }
        None => NewKeyOpts::NoEncryption,
    };

    // Store the private key
    let mut rng = rand_core::OsRng {};
    let private_key = PrivateKey::generate(&mut rng, derivation_info)?;

    let priv_comment = format!("{} secret key", comment);
    let mut file = OpenOptions::new()
        .write(true)
        .create_new(true)
        .open(privkey_path)?;

    write_base64_file(&mut file, &priv_comment, &private_key)?;

    // Store public key
    let public_key = private_key.public();

    let pub_comment = format!("{} public key", comment);
    let mut file = OpenOptions::new()
        .write(true)
        .create_new(true)
        .open(pubkey_path)?;

    write_base64_file(&mut file, &pub_comment, &public_key)?;

    Ok(())
}

fn human<T>(res: Result<T, Box<dyn std::error::Error>>) -> T {
    match res {
        Ok(val) => val,
        Err(e) => {
            println!("error: {}", e);
            process::exit(1);
        }
    }
}

fn unwrap_path<T>(kind: &'static str, path: Option<T>) -> T {
    match path {
        Some(p) => p,
        None => {
            println!("missing path to {}", kind);
            process::exit(1)
        }
    }
}

fn main() {
    let args = Args::parse();

    if args.verify {
        let public_key = unwrap_path("pubkey", args.pubkey);
        let message = unwrap_path("message", args.message_path);

        human(verify(
            &public_key,
            &message,
            args.signature_path,
            args.embed_message,
        ));

        return;
    }

    if args.generate {
        let public_key = unwrap_path("pubkey", args.pubkey);
        let private_key = unwrap_path("seckey", args.seckey);
        let rounds = if args.skip_key_encryption {
            None
        } else {
            Some(DEFAULT_KDF_ROUNDS)
        };

        human(generate(
            &public_key,
            &private_key,
            args.comment.as_deref(),
            rounds,
        ));

        return;
    }

    if args.sign {
        let private_key = unwrap_path("seckey", args.seckey);
        let msg_path = unwrap_path("message", args.message_path);

        human(sign(
            &private_key,
            &msg_path,
            args.signature_path,
            args.embed_message,
        ));

        return;
    }

    // No command specified.
    println!("{}", Args::into_app().render_usage());
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{node::LocalNode, ChainInfo, FullNode, HealthCheckError, Node, Swarm, Validator};
use anyhow::{anyhow, Result};
use diem_config::config::NodeConfig;
use diem_genesis_tool::validator_builder::ValidatorBuilder;
use diem_sdk::{
    crypto::ed25519::Ed25519PrivateKey,
    types::{chain_id::ChainId, AccountKey, LocalAccount, PeerId},
};
use std::{
    collections::HashMap,
    convert::TryFrom,
    fs, mem, ops,
    path::{Path, PathBuf},
};
use tempfile::TempDir;

#[derive(Debug)]
pub enum SwarmDirectory {
    Persistent(PathBuf),
    Temporary(TempDir),
}

impl SwarmDirectory {
    pub fn persist(&mut self) {
        match self {
            SwarmDirectory::Persistent(_) => {}
            SwarmDirectory::Temporary(_) => {
                let mut temp = SwarmDirectory::Persistent(PathBuf::new());
                mem::swap(self, &mut temp);
                let _ = mem::replace(self, temp.into_persistent());
            }
        }
    }

    pub fn into_persistent(self) -> Self {
        match self {
            SwarmDirectory::Temporary(tempdir) => SwarmDirectory::Persistent(tempdir.into_path()),
            SwarmDirectory::Persistent(dir) => SwarmDirectory::Persistent(dir),
        }
    }
}

impl ops::Deref for SwarmDirectory {
    type Target = Path;

    fn deref(&self) -> &Self::Target {
        match self {
            SwarmDirectory::Persistent(dir) => dir.deref(),
            SwarmDirectory::Temporary(dir) => dir.path(),
        }
    }
}

impl AsRef<Path> for SwarmDirectory {
    fn as_ref(&self) -> &Path {
        match self {
            SwarmDirectory::Persistent(dir) => dir.as_ref(),
            SwarmDirectory::Temporary(dir) => dir.as_ref(),
        }
    }
}

#[derive(Debug)]
struct ValidatorNetworkAddressEncryptionKey {
    key: diem_sdk::types::network_address::encrypted::Key,
    version: diem_sdk::types::network_address::encrypted::KeyVersion,
}

pub struct LocalSwarmBuilder {
    diem_node_bin_path: PathBuf,
    template: NodeConfig,
    number_of_validators: usize,
    dir: Option<PathBuf>,
}

impl LocalSwarmBuilder {
    pub fn new<T: AsRef<Path>>(diem_node_bin_path: T) -> Self {
        Self {
            diem_node_bin_path: diem_node_bin_path.as_ref().into(),
            template: NodeConfig::default_for_validator(),
            number_of_validators: 1,
            dir: None,
        }
    }

    pub fn template(mut self, template: NodeConfig) -> Self {
        self.template = template;
        self
    }

    pub fn number_of_validators(mut self, number_of_validators: usize) -> Self {
        self.number_of_validators = number_of_validators;
        self
    }

    pub fn dir<T: AsRef<Path>>(mut self, dir: T) -> Self {
        self.dir = Some(dir.as_ref().into());
        self
    }

    pub fn build(self) -> Result<LocalSwarm> {
        let dir = if let Some(dir) = self.dir {
            if dir.exists() {
                fs::remove_dir_all(&dir)?;
            }
            fs::create_dir_all(&dir)?;
            SwarmDirectory::Persistent(dir)
        } else {
            SwarmDirectory::Temporary(TempDir::new()?)
        };

        let (root_keys, validators) = ValidatorBuilder::new(&dir)
            .num_validators(self.number_of_validators)
            .template(self.template)
            .build()?;
        let diem_node_bin_path = self.diem_node_bin_path;

        let validators = validators
            .into_iter()
            .map(|v| {
                let node = LocalNode::new(&diem_node_bin_path, v.name, v.directory)?;
                Ok((node.peer_id(), node))
            })
            .collect::<Result<HashMap<_, _>>>()?;

        let validator_network_address_encryption_key = ValidatorNetworkAddressEncryptionKey {
            key: root_keys.validator_network_address_encryption_key,
            version: root_keys.validator_network_address_encryption_key_version,
        };

        let root_account = LocalAccount::new(
            diem_sdk::types::account_config::diem_root_address(),
            AccountKey::from_private_key(root_keys.root_key),
            0,
        );

        // Designated dealer account is built to have the same private key as the TC account
        let designated_dealer_account = LocalAccount::new(
            diem_sdk::types::account_config::testnet_dd_account_address(),
            AccountKey::from_private_key(
                Ed25519PrivateKey::try_from(root_keys.treasury_compliance_key.to_bytes().as_ref())
                    .unwrap(),
            ),
            0,
        );

        let treasury_compliance_account = LocalAccount::new(
            diem_sdk::types::account_config::treasury_compliance_account_address(),
            AccountKey::from_private_key(root_keys.treasury_compliance_key),
            0,
        );

        Ok(LocalSwarm {
            validators,
            dir,
            validator_network_address_encryption_key,
            root_account,
            treasury_compliance_account,
            designated_dealer_account,
            chain_id: ChainId::test(),
        })
    }
}

#[derive(Debug)]
pub struct LocalSwarm {
    validators: HashMap<PeerId, LocalNode>,
    dir: SwarmDirectory,
    validator_network_address_encryption_key: ValidatorNetworkAddressEncryptionKey,
    root_account: LocalAccount,
    treasury_compliance_account: LocalAccount,
    designated_dealer_account: LocalAccount,
    chain_id: ChainId,
}

impl LocalSwarm {
    pub fn builder<T: AsRef<Path>>(diem_node_bin_path: T) -> LocalSwarmBuilder {
        LocalSwarmBuilder::new(diem_node_bin_path)
    }

    pub fn launch(&mut self) -> Result<()> {
        // Start all the validators
        for validator in self.validators.values_mut() {
            validator.start()?;
        }

        // Wait for all of them to startup
        self.wait_for_startup()?;
        self.wait_for_connectivity()?;

        Ok(())
    }

    fn wait_for_startup(&mut self) -> Result<()> {
        let num_attempts = 10;
        let mut done = vec![false; self.validators.len()];
        for i in 0..num_attempts {
            println!("Wait for startup attempt: {} of {}", i, num_attempts);
            for (node, done) in self.validators.values_mut().zip(done.iter_mut()) {
                if *done {
                    continue;
                }
                match node.health_check() {
                    Ok(()) => *done = true,

                    Err(HealthCheckError::Unknown(e)) => {
                        return Err(anyhow!("Diem node '{}' is not running: {}", node.name(), e));
                    }
                    Err(HealthCheckError::NotRunning) => {
                        return Err(anyhow!("Diem node '{}' is not running", node.name()));
                    }
                    Err(HealthCheckError::RpcFailure(e)) => {
                        println!("rpc error: {}", e);
                        continue;
                    }
                }
            }

            // Check if all the nodes have been successfully launched
            if done.iter().all(|status| *status) {
                return Ok(());
            }

            ::std::thread::sleep(::std::time::Duration::from_millis(1000));
        }

        Err(anyhow!("Launching Swarm timed out"))
    }

    fn wait_for_connectivity(&self) -> Result<()> {
        let expected_peers = self.validators.len().saturating_sub(1);

        let num_attempts = 60;
        for i in 0..num_attempts {
            println!("Wait for connectivity attempt: {}", i);

            if self
                .validators
                .values()
                .map(|node| Validator::check_connectivity(node, expected_peers))
                .collect::<Result<Vec<_>>>()?
                .into_iter()
                .all(|b| b)
            {
                return Ok(());
            }

            ::std::thread::sleep(::std::time::Duration::from_millis(1000));
        }

        Err(anyhow!("Waiting for connectivity timed out"))
    }
}

impl Swarm for LocalSwarm {
    fn health_check(&mut self) -> Result<()> {
        todo!()
    }

    fn validators<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a dyn Validator> + 'a> {
        Box::new(self.validators.values().map(|v| v as &'a dyn Validator))
    }

    fn validators_mut<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a mut dyn Validator> + 'a> {
        Box::new(
            self.validators
                .values_mut()
                .map(|v| v as &'a mut dyn Validator),
        )
    }

    fn validator(&self, _id: PeerId) -> Option<&dyn Validator> {
        todo!()
    }

    fn validator_mut(&mut self, _id: PeerId) -> Option<&mut dyn Validator> {
        todo!()
    }

    fn full_nodes<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a dyn FullNode> + 'a> {
        todo!()
    }

    fn full_nodes_mut<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a mut dyn FullNode> + 'a> {
        todo!()
    }

    fn full_node(&self, _id: PeerId) -> Option<&dyn FullNode> {
        todo!()
    }

    fn full_node_mut(&mut self, _id: PeerId) -> Option<&mut dyn FullNode> {
        todo!()
    }

    fn add_validator(&mut self, _id: PeerId) -> Result<PeerId> {
        todo!()
    }

    fn remove_validator(&mut self, _id: PeerId) -> Result<()> {
        todo!()
    }

    fn add_full_node(&mut self, _id: PeerId) -> Result<()> {
        todo!()
    }

    fn remove_full_node(&mut self, _id: PeerId) -> Result<()> {
        todo!()
    }

    fn chain_info(&mut self) -> ChainInfo<'_> {
        ChainInfo::new(
            &mut self.root_account,
            &mut self.treasury_compliance_account,
            &mut self.designated_dealer_account,
            self.validators
                .values()
                .map(|v| v.json_rpc_endpoint().to_string())
                .next()
                .unwrap(),
            self.chain_id,
        )
    }
}
<commit_msg>forge: implement querying validators by PeerId for LocalSwarm<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{node::LocalNode, ChainInfo, FullNode, HealthCheckError, Node, Swarm, Validator};
use anyhow::{anyhow, Result};
use diem_config::config::NodeConfig;
use diem_genesis_tool::validator_builder::ValidatorBuilder;
use diem_sdk::{
    crypto::ed25519::Ed25519PrivateKey,
    types::{chain_id::ChainId, AccountKey, LocalAccount, PeerId},
};
use std::{
    collections::HashMap,
    convert::TryFrom,
    fs, mem, ops,
    path::{Path, PathBuf},
};
use tempfile::TempDir;

#[derive(Debug)]
pub enum SwarmDirectory {
    Persistent(PathBuf),
    Temporary(TempDir),
}

impl SwarmDirectory {
    pub fn persist(&mut self) {
        match self {
            SwarmDirectory::Persistent(_) => {}
            SwarmDirectory::Temporary(_) => {
                let mut temp = SwarmDirectory::Persistent(PathBuf::new());
                mem::swap(self, &mut temp);
                let _ = mem::replace(self, temp.into_persistent());
            }
        }
    }

    pub fn into_persistent(self) -> Self {
        match self {
            SwarmDirectory::Temporary(tempdir) => SwarmDirectory::Persistent(tempdir.into_path()),
            SwarmDirectory::Persistent(dir) => SwarmDirectory::Persistent(dir),
        }
    }
}

impl ops::Deref for SwarmDirectory {
    type Target = Path;

    fn deref(&self) -> &Self::Target {
        match self {
            SwarmDirectory::Persistent(dir) => dir.deref(),
            SwarmDirectory::Temporary(dir) => dir.path(),
        }
    }
}

impl AsRef<Path> for SwarmDirectory {
    fn as_ref(&self) -> &Path {
        match self {
            SwarmDirectory::Persistent(dir) => dir.as_ref(),
            SwarmDirectory::Temporary(dir) => dir.as_ref(),
        }
    }
}

#[derive(Debug)]
struct ValidatorNetworkAddressEncryptionKey {
    key: diem_sdk::types::network_address::encrypted::Key,
    version: diem_sdk::types::network_address::encrypted::KeyVersion,
}

pub struct LocalSwarmBuilder {
    diem_node_bin_path: PathBuf,
    template: NodeConfig,
    number_of_validators: usize,
    dir: Option<PathBuf>,
}

impl LocalSwarmBuilder {
    pub fn new<T: AsRef<Path>>(diem_node_bin_path: T) -> Self {
        Self {
            diem_node_bin_path: diem_node_bin_path.as_ref().into(),
            template: NodeConfig::default_for_validator(),
            number_of_validators: 1,
            dir: None,
        }
    }

    pub fn template(mut self, template: NodeConfig) -> Self {
        self.template = template;
        self
    }

    pub fn number_of_validators(mut self, number_of_validators: usize) -> Self {
        self.number_of_validators = number_of_validators;
        self
    }

    pub fn dir<T: AsRef<Path>>(mut self, dir: T) -> Self {
        self.dir = Some(dir.as_ref().into());
        self
    }

    pub fn build(self) -> Result<LocalSwarm> {
        let dir = if let Some(dir) = self.dir {
            if dir.exists() {
                fs::remove_dir_all(&dir)?;
            }
            fs::create_dir_all(&dir)?;
            SwarmDirectory::Persistent(dir)
        } else {
            SwarmDirectory::Temporary(TempDir::new()?)
        };

        let (root_keys, validators) = ValidatorBuilder::new(&dir)
            .num_validators(self.number_of_validators)
            .template(self.template)
            .build()?;
        let diem_node_bin_path = self.diem_node_bin_path;

        let validators = validators
            .into_iter()
            .map(|v| {
                let node = LocalNode::new(&diem_node_bin_path, v.name, v.directory)?;
                Ok((node.peer_id(), node))
            })
            .collect::<Result<HashMap<_, _>>>()?;

        let validator_network_address_encryption_key = ValidatorNetworkAddressEncryptionKey {
            key: root_keys.validator_network_address_encryption_key,
            version: root_keys.validator_network_address_encryption_key_version,
        };

        let root_account = LocalAccount::new(
            diem_sdk::types::account_config::diem_root_address(),
            AccountKey::from_private_key(root_keys.root_key),
            0,
        );

        // Designated dealer account is built to have the same private key as the TC account
        let designated_dealer_account = LocalAccount::new(
            diem_sdk::types::account_config::testnet_dd_account_address(),
            AccountKey::from_private_key(
                Ed25519PrivateKey::try_from(root_keys.treasury_compliance_key.to_bytes().as_ref())
                    .unwrap(),
            ),
            0,
        );

        let treasury_compliance_account = LocalAccount::new(
            diem_sdk::types::account_config::treasury_compliance_account_address(),
            AccountKey::from_private_key(root_keys.treasury_compliance_key),
            0,
        );

        Ok(LocalSwarm {
            validators,
            dir,
            validator_network_address_encryption_key,
            root_account,
            treasury_compliance_account,
            designated_dealer_account,
            chain_id: ChainId::test(),
        })
    }
}

#[derive(Debug)]
pub struct LocalSwarm {
    validators: HashMap<PeerId, LocalNode>,
    dir: SwarmDirectory,
    validator_network_address_encryption_key: ValidatorNetworkAddressEncryptionKey,
    root_account: LocalAccount,
    treasury_compliance_account: LocalAccount,
    designated_dealer_account: LocalAccount,
    chain_id: ChainId,
}

impl LocalSwarm {
    pub fn builder<T: AsRef<Path>>(diem_node_bin_path: T) -> LocalSwarmBuilder {
        LocalSwarmBuilder::new(diem_node_bin_path)
    }

    pub fn launch(&mut self) -> Result<()> {
        // Start all the validators
        for validator in self.validators.values_mut() {
            validator.start()?;
        }

        // Wait for all of them to startup
        self.wait_for_startup()?;
        self.wait_for_connectivity()?;

        Ok(())
    }

    fn wait_for_startup(&mut self) -> Result<()> {
        let num_attempts = 10;
        let mut done = vec![false; self.validators.len()];
        for i in 0..num_attempts {
            println!("Wait for startup attempt: {} of {}", i, num_attempts);
            for (node, done) in self.validators.values_mut().zip(done.iter_mut()) {
                if *done {
                    continue;
                }
                match node.health_check() {
                    Ok(()) => *done = true,

                    Err(HealthCheckError::Unknown(e)) => {
                        return Err(anyhow!("Diem node '{}' is not running: {}", node.name(), e));
                    }
                    Err(HealthCheckError::NotRunning) => {
                        return Err(anyhow!("Diem node '{}' is not running", node.name()));
                    }
                    Err(HealthCheckError::RpcFailure(e)) => {
                        println!("rpc error: {}", e);
                        continue;
                    }
                }
            }

            // Check if all the nodes have been successfully launched
            if done.iter().all(|status| *status) {
                return Ok(());
            }

            ::std::thread::sleep(::std::time::Duration::from_millis(1000));
        }

        Err(anyhow!("Launching Swarm timed out"))
    }

    fn wait_for_connectivity(&self) -> Result<()> {
        let expected_peers = self.validators.len().saturating_sub(1);

        let num_attempts = 60;
        for i in 0..num_attempts {
            println!("Wait for connectivity attempt: {}", i);

            if self
                .validators
                .values()
                .map(|node| Validator::check_connectivity(node, expected_peers))
                .collect::<Result<Vec<_>>>()?
                .into_iter()
                .all(|b| b)
            {
                return Ok(());
            }

            ::std::thread::sleep(::std::time::Duration::from_millis(1000));
        }

        Err(anyhow!("Waiting for connectivity timed out"))
    }
}

impl Swarm for LocalSwarm {
    fn health_check(&mut self) -> Result<()> {
        todo!()
    }

    fn validators<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a dyn Validator> + 'a> {
        Box::new(self.validators.values().map(|v| v as &'a dyn Validator))
    }

    fn validators_mut<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a mut dyn Validator> + 'a> {
        Box::new(
            self.validators
                .values_mut()
                .map(|v| v as &'a mut dyn Validator),
        )
    }

    fn validator(&self, id: PeerId) -> Option<&dyn Validator> {
        self.validators.get(&id).map(|v| v as &dyn Validator)
    }

    fn validator_mut(&mut self, id: PeerId) -> Option<&mut dyn Validator> {
        self.validators
            .get_mut(&id)
            .map(|v| v as &mut dyn Validator)
    }

    fn full_nodes<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a dyn FullNode> + 'a> {
        todo!()
    }

    fn full_nodes_mut<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a mut dyn FullNode> + 'a> {
        todo!()
    }

    fn full_node(&self, _id: PeerId) -> Option<&dyn FullNode> {
        todo!()
    }

    fn full_node_mut(&mut self, _id: PeerId) -> Option<&mut dyn FullNode> {
        todo!()
    }

    fn add_validator(&mut self, _id: PeerId) -> Result<PeerId> {
        todo!()
    }

    fn remove_validator(&mut self, _id: PeerId) -> Result<()> {
        todo!()
    }

    fn add_full_node(&mut self, _id: PeerId) -> Result<()> {
        todo!()
    }

    fn remove_full_node(&mut self, _id: PeerId) -> Result<()> {
        todo!()
    }

    fn chain_info(&mut self) -> ChainInfo<'_> {
        ChainInfo::new(
            &mut self.root_account,
            &mut self.treasury_compliance_account,
            &mut self.designated_dealer_account,
            self.validators
                .values()
                .map(|v| v.json_rpc_endpoint().to_string())
                .next()
                .unwrap(),
            self.chain_id,
        )
    }
}
<|endoftext|>