<commit_before>// Copyright 2015-2017 Intecture Developers.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

//! Endpoint for managing packages.
//!
//! A package is represented by the `Package` struct, which is idempotent. This
//! means you can execute it repeatedly and it'll only run as needed.

pub mod providers;

use command::{ExitStatus, parse_body_stream};
use errors::*;
use futures::{future, Future};
use futures::stream::Stream;
use host::Host;
use remote::{Request, Response};
use self::providers::PackageProvider;

/// Represents a system package to be managed for a host.
///
///## Examples
///
pub struct Package<H: Host> {
    host: H,
    provider: Option<Box<PackageProvider>>,
    name: String,
}

impl<H: Host + 'static> Package<H> {
    /// Create a new `Package` with the default `PackageProvider`.
    pub fn new(host: &H, name: &str) -> Package<H> {
        Package {
            host: host.clone(),
            provider: None,
            name: name.into(),
        }
    }

    /// Create a new `Package` with the specified `PackageProvider`.
    ///
    ///## Example
    ///```
    ///extern crate futures;
    ///extern crate intecture_api;
    ///extern crate tokio_core;
    ///
    ///use futures::Future;
    ///use intecture_api::package::providers::Yum;
    ///use intecture_api::prelude::*;
    ///use tokio_core::reactor::Core;
    ///
    ///# fn main() {
    ///let mut core = Core::new().unwrap();
    ///let handle = core.handle();
    ///
    ///let host = Local::new(&handle).wait().unwrap();
    ///
    ///Package::with_provider(&host, Yum, "nginx");
    ///# }
    pub fn with_provider<P>(host: &H, provider: P, name: &str) -> Package<H>
        where P: PackageProvider + 'static
    {
        Package {
            host: host.clone(),
            provider: Some(Box::new(provider)),
            name: name.into(),
        }
    }

    /// Check if the package is installed.
    pub fn installed(&self) -> Box<Future<Item = bool, Error = Error>> {
        let request = Request::PackageInstalled(self.provider.as_ref().map(|p| p.name()), self.name.clone());
        Box::new(self.host.request(request)
            .chain_err(|| ErrorKind::Request { endpoint: "Package", func: "installed" })
            .map(|msg| {
                match msg.into_inner() {
                    Response::Bool(b) => b,
                    _ => unreachable!(),
                }
            }))
    }

    /// Install the package.
    ///
    /// This function is idempotent, which is represented by the type
    /// `Future<Item = Option<..>, ...>`. Thus if it returns `Option::None`
    /// then the package is already installed, and if it returns `Option::Some`
    /// then Intecture is attempting to install the package.
    ///
    /// If this fn returns `Option::Some<..>`, the nested tuple will hold
    /// handles to the live output and the result of the installation. Under
    /// the hood this reuses the `Command` endpoint, so see
    /// [`Command` docs](../command/struct.Command.html) for detailed
    /// usage.
    ///
    ///# Example
    ///
    /// Install a package and print the result.
    ///
    ///```no_run
    ///extern crate futures;
    ///extern crate intecture_api;
    ///extern crate tokio_core;
    ///
    ///use futures::{future, Future, Stream};
    ///use intecture_api::errors::Error;
    ///use intecture_api::prelude::*;
    ///use tokio_core::reactor::Core;
    ///
    ///# fn main() {
    ///let mut core = Core::new().unwrap();
    ///let handle = core.handle();
    ///
    ///let host = Local::new(&handle).wait().unwrap();
    ///
    ///let nginx = Package::new(&host, "nginx");
    ///let result = nginx.install().and_then(|status| {
    ///    match status {
    ///        // We're performing the install
    ///        Some((stream, status)) => Box::new(stream.fold(String::new(), |mut acc, line| {
    ///                acc.push_str(&line);
    ///                future::ok::<_, Error>(acc)
    ///            })
    ///            .join(status)
    ///            .map(|(output, status)| {
    ///                if status.success {
    ///                    println!("Installed");
    ///                } else {
    ///                    println!("Failed with output: {}", output);
    ///                }
    ///            })) as Box<Future<Item = _, Error = Error>>,
    ///        None => {
    ///            println!("Already installed");
    ///            Box::new(future::ok(()))
    ///        },
    ///    }
    ///});
    ///
    ///core.run(result).unwrap();
    ///# }
    ///```
    pub fn install(&self) -> Box<Future<Item = Option<(
            Box<Stream<Item = String, Error = Error>>,
            Box<Future<Item = ExitStatus, Error = Error>>
        )>, Error = Error>>
    {
        let host = self.host.clone();
        let provider = self.provider.as_ref().map(|p| p.name());
        let name = self.name.clone();

        Box::new(self.installed()
            .and_then(move |installed| {
                if installed {
                    Box::new(future::ok(None)) as Box<Future<Item = _, Error = Error>>
                } else {
                    Box::new(host.request(Request::PackageInstall(provider, name))
                        .chain_err(|| ErrorKind::Request { endpoint: "Package", func: "install" })
                        .map(|msg| {
                            Some(parse_body_stream(msg))
                        }))
                }
            }))
    }

    /// Uninstall the package.
    ///
    /// This function is idempotent, which is represented by the type
    /// `Future<Item = Option<..>, ...>`. Thus if it returns `Option::None`
    /// then the package is already uninstalled, and if it returns
    /// `Option::Some` then Intecture is attempting to uninstall the package.
    ///
    /// If this fn returns `Option::Some<..>`, the nested tuple will hold
    /// handles to the live output and the result of the deinstallation. Under
    /// the hood this reuses the `Command` endpoint, so see
    /// [`Command` docs](../command/struct.Command.html) for detailed
    /// usage.
    pub fn uninstall(&self) -> Box<Future<Item = Option<(
            Box<Stream<Item = String, Error = Error>>,
            Box<Future<Item = ExitStatus, Error = Error>>
        )>, Error = Error>>
    {
        let host = self.host.clone();
        let provider = self.provider.as_ref().map(|p| p.name());
        let name = self.name.clone();

        Box::new(self.installed()
            .and_then(move |installed| {
                if installed {
                    Box::new(host.request(Request::PackageUninstall(provider, name))
                        .chain_err(|| ErrorKind::Request { endpoint: "Package", func: "uninstall" })
                        .map(|msg| {
                            Some(parse_body_stream(msg))
                        }))
                } else {
                    Box::new(future::ok(None)) as Box<Future<Item = _, Error = Error>>
                }
            }))
    }
}
<commit_msg>Moved Package::install() example to Package struct docs<commit_after>// Copyright 2015-2017 Intecture Developers.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

//! Endpoint for managing packages.
//!
//! A package is represented by the `Package` struct, which is idempotent. This
//! means you can execute it repeatedly and it'll only run as needed.

pub mod providers;

use command::{ExitStatus, parse_body_stream};
use errors::*;
use futures::{future, Future};
use futures::stream::Stream;
use host::Host;
use remote::{Request, Response};
use self::providers::PackageProvider;

/// Represents a system package to be managed for a host.
///
///# Example
///
/// Install a package and print the result.
///
///```no_run
///extern crate futures;
///extern crate intecture_api;
///extern crate tokio_core;
///
///use futures::{future, Future, Stream};
///use intecture_api::errors::Error;
///use intecture_api::prelude::*;
///use tokio_core::reactor::Core;
///
///# fn main() {
///let mut core = Core::new().unwrap();
///let handle = core.handle();
///
///let host = Local::new(&handle).wait().unwrap();
///
///let nginx = Package::new(&host, "nginx");
///let result = nginx.install().and_then(|status| {
///    match status {
///        // We're performing the install
///        Some((stream, status)) => Box::new(stream.fold(String::new(), |mut acc, line| {
///                acc.push_str(&line);
///                future::ok::<_, Error>(acc)
///            })
///            .join(status)
///            .map(|(output, status)| {
///                if status.success {
///                    println!("Installed");
///                } else {
///                    println!("Failed with output: {}", output);
///                }
///            })) as Box<Future<Item = _, Error = Error>>,
///
///        // This package is already installed
///        None => {
///            println!("Already installed");
///            Box::new(future::ok(()))
///        },
///    }
///});
///
///core.run(result).unwrap();
///# }
///```
pub struct Package<H: Host> {
    host: H,
    provider: Option<Box<PackageProvider>>,
    name: String,
}

impl<H: Host + 'static> Package<H> {
    /// Create a new `Package` with the default `PackageProvider`.
    pub fn new(host: &H, name: &str) -> Package<H> {
        Package {
            host: host.clone(),
            provider: None,
            name: name.into(),
        }
    }

    /// Create a new `Package` with the specified `PackageProvider`.
    ///
    ///## Example
    ///```
    ///extern crate futures;
    ///extern crate intecture_api;
    ///extern crate tokio_core;
    ///
    ///use futures::Future;
    ///use intecture_api::package::providers::Yum;
    ///use intecture_api::prelude::*;
    ///use tokio_core::reactor::Core;
    ///
    ///# fn main() {
    ///let mut core = Core::new().unwrap();
    ///let handle = core.handle();
    ///
    ///let host = Local::new(&handle).wait().unwrap();
    ///
    ///Package::with_provider(&host, Yum, "nginx");
    ///# }
    pub fn with_provider<P>(host: &H, provider: P, name: &str) -> Package<H>
        where P: PackageProvider + 'static
    {
        Package {
            host: host.clone(),
            provider: Some(Box::new(provider)),
            name: name.into(),
        }
    }

    /// Check if the package is installed.
    pub fn installed(&self) -> Box<Future<Item = bool, Error = Error>> {
        let request = Request::PackageInstalled(self.provider.as_ref().map(|p| p.name()), self.name.clone());
        Box::new(self.host.request(request)
            .chain_err(|| ErrorKind::Request { endpoint: "Package", func: "installed" })
            .map(|msg| {
                match msg.into_inner() {
                    Response::Bool(b) => b,
                    _ => unreachable!(),
                }
            }))
    }

    /// Install the package.
    ///
    /// This function is idempotent, which is represented by the type
    /// `Future<Item = Option<..>, ...>`. Thus if it returns `Option::None`
    /// then the package is already installed, and if it returns `Option::Some`
    /// then Intecture is attempting to install the package.
    ///
    /// If this fn returns `Option::Some<..>`, the nested tuple will hold
    /// handles to the live output and the result of the installation. Under
    /// the hood this reuses the `Command` endpoint, so see
    /// [`Command` docs](../command/struct.Command.html) for detailed
    /// usage.
    pub fn install(&self) -> Box<Future<Item = Option<(
            Box<Stream<Item = String, Error = Error>>,
            Box<Future<Item = ExitStatus, Error = Error>>
        )>, Error = Error>>
    {
        let host = self.host.clone();
        let provider = self.provider.as_ref().map(|p| p.name());
        let name = self.name.clone();

        Box::new(self.installed()
            .and_then(move |installed| {
                if installed {
                    Box::new(future::ok(None)) as Box<Future<Item = _, Error = Error>>
                } else {
                    Box::new(host.request(Request::PackageInstall(provider, name))
                        .chain_err(|| ErrorKind::Request { endpoint: "Package", func: "install" })
                        .map(|msg| {
                            Some(parse_body_stream(msg))
                        }))
                }
            }))
    }

    /// Uninstall the package.
    ///
    /// This function is idempotent, which is represented by the type
    /// `Future<Item = Option<..>, ...>`. Thus if it returns `Option::None`
    /// then the package is already uninstalled, and if it returns
    /// `Option::Some` then Intecture is attempting to uninstall the package.
    ///
    /// If this fn returns `Option::Some<..>`, the nested tuple will hold
    /// handles to the live output and the result of the deinstallation. Under
    /// the hood this reuses the `Command` endpoint, so see
    /// [`Command` docs](../command/struct.Command.html) for detailed
    /// usage.
    pub fn uninstall(&self) -> Box<Future<Item = Option<(
            Box<Stream<Item = String, Error = Error>>,
            Box<Future<Item = ExitStatus, Error = Error>>
        )>, Error = Error>>
    {
        let host = self.host.clone();
        let provider = self.provider.as_ref().map(|p| p.name());
        let name = self.name.clone();

        Box::new(self.installed()
            .and_then(move |installed| {
                if installed {
                    Box::new(host.request(Request::PackageUninstall(provider, name))
                        .chain_err(|| ErrorKind::Request { endpoint: "Package", func: "uninstall" })
                        .map(|msg| {
                            Some(parse_body_stream(msg))
                        }))
                } else {
                    Box::new(future::ok(None)) as Box<Future<Item = _, Error = Error>>
                }
            }))
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 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.

// This test case tests the incremental compilation hash (ICH) implementation
// for `while` loops.

// The general pattern followed here is: Change one thing between rev1 and rev2
// and make sure that the hash has changed, then change nothing between rev2 and
// rev3 and make sure that the hash has not changed.

// must-compile-successfully
// revisions: cfail1 cfail2 cfail3
// compile-flags: -Z query-dep-graph

#![allow(warnings)]
#![feature(rustc_attrs)]
#![crate_type="rlib"]

// Change loop body ------------------------------------------------------------
#[cfg(cfail1)]
fn change_loop_body() {
    let mut _x = 0;
    while true {
        _x = 1;
        break;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(label="Hir", cfg="cfail2")]
#[rustc_clean(label="Hir", cfg="cfail3")]
#[rustc_dirty(label="HirBody", cfg="cfail2")]
#[rustc_clean(label="HirBody", cfg="cfail3")]
fn change_loop_body() {
    let mut _x = 0;
    while true {
        _x = 2;
        break;
    }
}

// Change loop body ------------------------------------------------------------
#[cfg(cfail1)]
fn change_loop_condition() {
    let mut _x = 0;
    while true {
        _x = 1;
        break;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(label="Hir", cfg="cfail2")]
#[rustc_clean(label="Hir", cfg="cfail3")]
#[rustc_dirty(label="HirBody", cfg="cfail2")]
#[rustc_clean(label="HirBody", cfg="cfail3")]
fn change_loop_condition() {
    let mut _x = 0;
    while false {
        _x = 1;
        break;
    }
}

// Add break -------------------------------------------------------------------
#[cfg(cfail1)]
fn add_break() {
    let mut _x = 0;
    while true {
        _x = 1;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(label="Hir", cfg="cfail2")]
#[rustc_clean(label="Hir", cfg="cfail3")]
#[rustc_dirty(label="HirBody", cfg="cfail2")]
#[rustc_clean(label="HirBody", cfg="cfail3")]
fn add_break() {
    let mut _x = 0;
    while true {
        _x = 1;
        break;
    }
}

// Add loop label --------------------------------------------------------------
#[cfg(cfail1)]
fn add_loop_label() {
    let mut _x = 0;
    while true {
        _x = 1;
        break;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(label="Hir", cfg="cfail2")]
#[rustc_clean(label="Hir", cfg="cfail3")]
#[rustc_dirty(label="HirBody", cfg="cfail2")]
#[rustc_clean(label="HirBody", cfg="cfail3")]
fn add_loop_label() {
    let mut _x = 0;
    'label: while true {
        _x = 1;
        break;
    }
}

// Add loop label to break -----------------------------------------------------
#[cfg(cfail1)]
fn add_loop_label_to_break() {
    let mut _x = 0;
    'label: while true {
        _x = 1;
        break;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(label="Hir", cfg="cfail2")]
#[rustc_clean(label="Hir", cfg="cfail3")]
#[rustc_dirty(label="HirBody", cfg="cfail2")]
#[rustc_clean(label="HirBody", cfg="cfail3")]
fn add_loop_label_to_break() {
    let mut _x = 0;
    'label: while true {
        _x = 1;
        break 'label;
    }
}

// Change break label ----------------------------------------------------------
#[cfg(cfail1)]
fn change_break_label() {
    let mut _x = 0;
    'outer: while true {
        'inner: while true {
            _x = 1;
            break 'inner;
        }
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(label="Hir", cfg="cfail2")]
#[rustc_clean(label="Hir", cfg="cfail3")]
#[rustc_dirty(label="HirBody", cfg="cfail2")]
#[rustc_clean(label="HirBody", cfg="cfail3")]
fn change_break_label() {
    let mut _x = 0;
    'outer: while true {
        'inner: while true {
            _x = 1;
            break 'outer;
        }
    }
}

// Add loop label to continue --------------------------------------------------
#[cfg(cfail1)]
fn add_loop_label_to_continue() {
    let mut _x = 0;
    'label: while true {
        _x = 1;
        continue;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(label="Hir", cfg="cfail2")]
#[rustc_clean(label="Hir", cfg="cfail3")]
#[rustc_dirty(label="HirBody", cfg="cfail2")]
#[rustc_clean(label="HirBody", cfg="cfail3")]
fn add_loop_label_to_continue() {
    let mut _x = 0;
    'label: while true {
        _x = 1;
        continue 'label;
    }
}

// Change continue label ----------------------------------------------------------
#[cfg(cfail1)]
fn change_continue_label() {
    let mut _x = 0;
    'outer: while true {
        'inner: while true {
            _x = 1;
            continue 'inner;
        }
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(label="Hir", cfg="cfail2")]
#[rustc_clean(label="Hir", cfg="cfail3")]
#[rustc_dirty(label="HirBody", cfg="cfail2")]
#[rustc_clean(label="HirBody", cfg="cfail3")]
fn change_continue_label() {
    let mut _x = 0;
    'outer: while true {
        'inner: while true {
            _x = 1;
            continue 'outer;
        }
    }
}

// Change continue to break ----------------------------------------------------
#[cfg(cfail1)]
fn change_continue_to_break() {
    let mut _x = 0;
    while true {
        _x = 1;
        continue;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(label="Hir", cfg="cfail2")]
#[rustc_clean(label="Hir", cfg="cfail3")]
#[rustc_dirty(label="HirBody", cfg="cfail2")]
#[rustc_clean(label="HirBody", cfg="cfail3")]
fn change_continue_to_break() {
    let mut _x = 0;
    while true {
        _x = 1;
        break;
    }
}
<commit_msg>Update 'while loop' fingerprint hash tests<commit_after>// Copyright 2016 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.

// This test case tests the incremental compilation hash (ICH) implementation
// for `while` loops.

// The general pattern followed here is: Change one thing between rev1 and rev2
// and make sure that the hash has changed, then change nothing between rev2 and
// rev3 and make sure that the hash has not changed.

// must-compile-successfully
// revisions: cfail1 cfail2 cfail3
// compile-flags: -Z query-dep-graph

#![allow(warnings)]
#![feature(rustc_attrs)]
#![crate_type="rlib"]

// Change loop body ------------------------------------------------------------
#[cfg(cfail1)]
pub fn change_loop_body() {
    let mut _x = 0;
    while true {
        _x = 1;
        break;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(cfg="cfail2", except="HirBody, MirValidated, MirOptimized")]
#[rustc_clean(cfg="cfail3")]
pub fn change_loop_body() {
    let mut _x = 0;
    while true {
        _x = 2;
        break;
    }
}

// Change loop body ------------------------------------------------------------
#[cfg(cfail1)]
pub fn change_loop_condition() {
    let mut _x = 0;
    while true {
        _x = 1;
        break;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(cfg="cfail2", except="HirBody, MirValidated, MirOptimized")]
#[rustc_clean(cfg="cfail3")]
pub fn change_loop_condition() {
    let mut _x = 0;
    while false {
        _x = 1;
        break;
    }
}

// Add break -------------------------------------------------------------------
#[cfg(cfail1)]
pub fn add_break() {
    let mut _x = 0;
    while true {
        _x = 1;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(cfg="cfail2", except="HirBody, MirValidated, MirOptimized, TypeckTables")]
#[rustc_clean(cfg="cfail3")]
pub fn add_break() {
    let mut _x = 0;
    while true {
        _x = 1;
        break;
    }
}

// Add loop label --------------------------------------------------------------
#[cfg(cfail1)]
pub fn add_loop_label() {
    let mut _x = 0;
    while true {
        _x = 1;
        break;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(cfg="cfail2", except="HirBody")]
#[rustc_clean(cfg="cfail3")]
pub fn add_loop_label() {
    let mut _x = 0;
    'label: while true {
        _x = 1;
        break;
    }
}

// Add loop label to break -----------------------------------------------------
#[cfg(cfail1)]
pub fn add_loop_label_to_break() {
    let mut _x = 0;
    'label: while true {
        _x = 1;
        break;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(cfg="cfail2", except="HirBody")]
#[rustc_clean(cfg="cfail3")]
pub fn add_loop_label_to_break() {
    let mut _x = 0;
    'label: while true {
        _x = 1;
        break 'label;
    }
}

// Change break label ----------------------------------------------------------
#[cfg(cfail1)]
pub fn change_break_label() {
    let mut _x = 0;
    'outer: while true {
        'inner: while true {
            _x = 1;
            break 'inner;
        }
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(cfg="cfail2", except="HirBody, MirValidated, MirOptimized")]
#[rustc_clean(cfg="cfail3")]
pub fn change_break_label() {
    let mut _x = 0;
    'outer: while true {
        'inner: while true {
            _x = 1;
            break 'outer;
        }
    }
}

// Add loop label to continue --------------------------------------------------
#[cfg(cfail1)]
pub fn add_loop_label_to_continue() {
    let mut _x = 0;
    'label: while true {
        _x = 1;
        continue;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(cfg="cfail2", except="HirBody")]
#[rustc_clean(cfg="cfail3")]
pub fn add_loop_label_to_continue() {
    let mut _x = 0;
    'label: while true {
        _x = 1;
        continue 'label;
    }
}

// Change continue label ----------------------------------------------------------
#[cfg(cfail1)]
pub fn change_continue_label() {
    let mut _x = 0;
    'outer: while true {
        'inner: while true {
            _x = 1;
            continue 'inner;
        }
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(cfg="cfail2", except="HirBody, MirValidated")]
#[rustc_clean(cfg="cfail3")]
pub fn change_continue_label() {
    let mut _x = 0;
    'outer: while true {
        'inner: while true {
            _x = 1;
            continue 'outer;
        }
    }
}

// Change continue to break ----------------------------------------------------
#[cfg(cfail1)]
pub fn change_continue_to_break() {
    let mut _x = 0;
    while true {
        _x = 1;
        continue;
    }
}

#[cfg(not(cfail1))]
#[rustc_clean(cfg="cfail2", except="HirBody, MirValidated, MirOptimized")]
#[rustc_clean(cfg="cfail3")]
pub fn change_continue_to_break() {
    let mut _x = 0;
    while true {
        _x = 1;
        break;
    }
}
<|endoftext|>


<commit_before>extern crate ktensor as k;
use k::{Arc, Vec2, Tensor, Context, Graph, State, Variable};
use std::path::{Path};
use std::fs::{File};
use std::io::{BufReader, Read};
use std::error::{Error};
use std::cmp;

fn read_u8(reader: &mut Read) -> u8 {
    use std::mem;

    let mut buf: [u8; 1] = [0];
    match reader.read_exact(&mut buf).map(|_| {
        let data: u8 = unsafe { mem::transmute(buf) };
        data
    }) {
        Err(reason) => panic!("failed to read u8: {}", Error::description(&reason)),
        Ok(byte)    => byte,
    }
}

fn read_u32(reader: &mut Read) -> u32 {
    use std::mem;

    let mut buf: [u8; 4] = [0, 0, 0, 0];
    match reader.read_exact(&mut buf).map(|_| {
        let data: u32 = unsafe { mem::transmute(buf) };
        data
    }) {
        Err(reason) => panic!("failed to read u32: {}", Error::description(&reason)),
        Ok(byte)    => byte,
    }
}

fn read_mnist(labels_path: &Path, labels_checknum: u32, data_path: &Path, data_checknum: u32, batch_size: usize, samples: Option<usize>) {
    let labels_file = match File::open(&labels_path) {
        Err(reason) => panic!("failed to open {}: {}", labels_path.display(), Error::description(&reason)),
        Ok(file)    => file,
    };
    let data_file = match File::open(&data_path) {
        Err(reason) => panic!("failed to open {}: {}", data_path.display(), Error::description(&reason)),
        Ok(file)    => file,
    };
    let ref mut labels_reader = BufReader::new(labels_file);
    let ref mut data_reader = BufReader::new(data_file);

    let labels_magic = u32::from_be(read_u32(labels_reader));
    let data_magic = u32::from_be(read_u32(data_reader));

    assert_eq!(labels_magic, labels_checknum);
    assert_eq!(data_magic, data_checknum);

    let labels_count = u32::from_be(read_u32(labels_reader)) as usize;
    let data_count = u32::from_be(read_u32(data_reader)) as usize;
    let sample_count = cmp::min(labels_count, data_count);
    let sample_count = cmp::min(sample_count, samples.unwrap_or(sample_count));

    let rows = u32::from_be(read_u32(data_reader)) as usize;
    let columns = u32::from_be(read_u32(data_reader)) as usize;

    let mut sample_vec: Vec<(Vec<u8>, Vec<u8>)> = Vec::with_capacity(sample_count/batch_size);
    for _ in 0..sample_count/batch_size {
        let mut sample_data = vec![0u8; batch_size*rows*columns];
        match data_reader.read_exact(sample_data.as_mut()) {
            Err(reason) => panic!("failed to read data byte array: {}", Error::description(&reason)),
            Ok(_)       => (),
        }
        let mut sample_labels = vec![0u8; batch_size];
        match labels_reader.read_exact(sample_labels.as_mut()) {
            Err(reason) => panic!("failed to read labels byte array: {}", Error::description(&reason)),
            Ok(_)       => (),
        }
        sample_vec.push((sample_data, sample_labels));
    }
}

#[test]
#[ignore]
fn mnist(){
    let train_labels_path = Path::new("data/train-labels-idx1-ubyte");
    let train_data_path = Path::new("data/train-images-idx3-ubyte");
    read_mnist(&train_labels_path, 2049, &train_data_path, 2051, 16, Some(4096));
}
<commit_msg>some more mnist<commit_after>extern crate ktensor as k;
use k::{Arc, Vec2, Tensor, Context, Graph, State, Variable};
use std::path::{Path};
use std::fs::{File};
use std::io::{BufReader, Read};
use std::error::{Error};
use std::cmp;

fn read_u8(reader: &mut Read) -> u8 {
    use std::mem;

    let mut buf: [u8; 1] = [0];
    match reader.read_exact(&mut buf).map(|_| {
        let data: u8 = unsafe { mem::transmute(buf) };
        data
    }) {
        Err(reason) => panic!("failed to read u8: {}", Error::description(&reason)),
        Ok(byte)    => byte,
    }
}

fn read_u32(reader: &mut Read) -> u32 {
    use std::mem;

    let mut buf: [u8; 4] = [0, 0, 0, 0];
    match reader.read_exact(&mut buf).map(|_| {
        let data: u32 = unsafe { mem::transmute(buf) };
        data
    }) {
        Err(reason) => panic!("failed to read u32: {}", Error::description(&reason)),
        Ok(byte)    => byte,
    }
}

fn read_mnist(labels_path: &Path, labels_checknum: u32, data_path: &Path, data_checknum: u32, batch_size: usize, samples: Option<usize>) {
    let labels_file = match File::open(&labels_path) {
        Err(reason) => panic!("failed to open {}: {}", labels_path.display(), Error::description(&reason)),
        Ok(file)    => file,
    };
    let data_file = match File::open(&data_path) {
        Err(reason) => panic!("failed to open {}: {}", data_path.display(), Error::description(&reason)),
        Ok(file)    => file,
    };
    let ref mut labels_reader = BufReader::new(labels_file);
    let ref mut data_reader = BufReader::new(data_file);

    let labels_magic = u32::from_be(read_u32(labels_reader));
    let data_magic = u32::from_be(read_u32(data_reader));

    assert_eq!(labels_magic, labels_checknum);
    assert_eq!(data_magic, data_checknum);

    let labels_count = u32::from_be(read_u32(labels_reader)) as usize;
    let data_count = u32::from_be(read_u32(data_reader)) as usize;
    let sample_count = cmp::min(labels_count, data_count);
    let sample_count = cmp::min(sample_count, samples.unwrap_or(sample_count));

    let rows = u32::from_be(read_u32(data_reader)) as usize;
    let columns = u32::from_be(read_u32(data_reader)) as usize;

    let mut sample_vec: Vec<(Vec<u8>, Vec<u8>)> = Vec::with_capacity(sample_count/batch_size);
    for _ in 0..sample_count/batch_size {
        let mut sample_data = vec![0u8; batch_size*rows*columns];
        match data_reader.read_exact(sample_data.as_mut()) {
            Err(reason) => panic!("failed to read data byte array: {}", Error::description(&reason)),
            Ok(_)       => (),
        }
        let mut sample_labels = vec![0u8; batch_size];
        match labels_reader.read_exact(sample_labels.as_mut()) {
            Err(reason) => panic!("failed to read labels byte array: {}", Error::description(&reason)),
            Ok(_)       => (),
        }
        sample_vec.push((sample_data, sample_labels));
    }
}

#[test]
#[ignore]
fn mnist(){
    let train_labels_path = Path::new("data/train-labels-idx1-ubyte");
    let train_data_path = Path::new("data/train-images-idx3-ubyte");
    read_mnist(&train_labels_path, 2049, &train_data_path, 2051, 16, Some(4096));

    ///////////////
    // Variables //
    ///////////////

    let input_x = Arc::new(Variable::new("input_x".to_string(), Vec2(0, 28 * 28)));
    let target_y = Arc::new(Variable::new("target_y".to_string(), Vec2(0, 10)));

    ///////////
    // Graph //
    ///////////

    let layers: usize = 2;
    let mut states = Vec::<Arc<State>>::with_capacity(2 * layers);
    let mut graph_head: Arc<Graph<f64>> = input_x.clone();

    {
        let w = Arc::new(State::new(format!("weight_w_{}", 1), Vec2(28 * 28, 16)));
        let b = Arc::new(State::new(format!("weight_b_{}", 1), Vec2(1, 16)));

        let dot = Arc::new(k::op::dot::<f64>(format!("layer_{}_dot", 1), graph_head.clone(), w.clone()));
        let add = Arc::new(k::op::add::<f64>(format!("layer_{}_add", 1), dot.clone(), b.clone()));

        let relu = Arc::new(k::op::relu_f64(format!("layer_{}_relu", 1), add.clone()));

        graph_head = relu.clone();

        states.push(w.clone());
        states.push(b.clone());
    }

    {
        let w = Arc::new(State::new(format!("weight_w_{}", 1), Vec2(2, 4)));
        let b = Arc::new(State::new(format!("weight_b_{}", 1), Vec2(1, 4)));

        let dot = Arc::new(k::op::dot::<f64>(format!("layer_{}_dot", 1), graph_head.clone(), w.clone()));
        let add = Arc::new(k::op::add::<f64>(format!("layer_{}_add", 1), dot.clone(), b.clone()));

        let relu = Arc::new(k::op::relu_f64(format!("layer_{}_relu", 1), add.clone()));

        graph_head = relu.clone();

        states.push(w.clone());
        states.push(b.clone());
    }
}
<|endoftext|>


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

use crate::{
    ChainInfo, FullNode, HealthCheckError, LocalNode, LocalVersion, Node, NodeExt, Swarm, SwarmExt,
    Validator, Version,
};
use anyhow::{anyhow, bail, Result};
use diem_config::config::NodeConfig;
use diem_genesis_tool::{fullnode_builder::FullnodeConfig, validator_builder::ValidatorBuilder};
use diem_sdk::{
    crypto::ed25519::Ed25519PrivateKey,
    types::{
        chain_id::ChainId, transaction::Transaction, waypoint::Waypoint, AccountKey, LocalAccount,
        PeerId,
    },
};
use std::{
    collections::HashMap,
    convert::TryFrom,
    fs, mem,
    num::NonZeroUsize,
    ops,
    path::{Path, PathBuf},
    sync::Arc,
    time::{Duration, Instant},
};
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 {
    versions: Arc<HashMap<Version, LocalVersion>>,
    initial_version: Option<Version>,
    template: NodeConfig,
    number_of_validators: NonZeroUsize,
    dir: Option<PathBuf>,
}

impl LocalSwarmBuilder {
    pub fn new(versions: Arc<HashMap<Version, LocalVersion>>) -> Self {
        Self {
            versions,
            initial_version: None,
            template: NodeConfig::default_for_validator(),
            number_of_validators: NonZeroUsize::new(1).unwrap(),
            dir: None,
        }
    }

    pub fn initial_version(mut self, initial_version: Version) -> Self {
        self.initial_version = Some(initial_version);
        self
    }

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

    pub fn number_of_validators(mut self, number_of_validators: NonZeroUsize) -> 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<R>(self, rng: R) -> Result<LocalSwarm>
    where
        R: ::rand::RngCore + ::rand::CryptoRng,
    {
        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, genesis, genesis_waypoint, validators) = ValidatorBuilder::new(
            &dir,
            diem_framework_releases::current_module_blobs().to_vec(),
        )
        .num_validators(self.number_of_validators)
        .template(self.template)
        .build(rng)?;

        // Get the initial version to start the nodes with, either the one provided or fallback to
        // using the the latest version
        let versions = self.versions;
        let initial_version = self.initial_version.unwrap_or_else(|| {
            versions
                .iter()
                .max_by(|v1, v2| v1.0.cmp(v2.0))
                .unwrap()
                .0
                .clone()
        });
        let version = versions.get(&initial_version).unwrap();

        let validators = validators
            .into_iter()
            .map(|v| {
                let node = LocalNode::new(version.to_owned(), 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),
            1,
        );

        // 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 {
            node_name_counter: validators.len() as u64,
            genesis,
            genesis_waypoint,
            versions,
            validators,
            fullnodes: HashMap::new(),
            dir,
            validator_network_address_encryption_key,
            root_account,
            treasury_compliance_account,
            designated_dealer_account,
            chain_id: ChainId::test(),
        })
    }
}

#[derive(Debug)]
pub struct LocalSwarm {
    node_name_counter: u64,
    genesis: Transaction,
    genesis_waypoint: Waypoint,
    versions: Arc<HashMap<Version, LocalVersion>>,
    validators: HashMap<PeerId, LocalNode>,
    fullnodes: 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(versions: Arc<HashMap<Version, LocalVersion>>) -> LocalSwarmBuilder {
        LocalSwarmBuilder::new(versions)
    }

    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
        let deadline = Instant::now() + Duration::from_secs(60);
        self.wait_for_startup()?;
        self.wait_for_connectivity(deadline)?;
        self.liveness_check(deadline)?;

        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"))
    }

    pub fn add_validator_fullnode(
        &mut self,
        version: &Version,
        template: NodeConfig,
        validator_peer_id: PeerId,
    ) -> Result<PeerId> {
        let validator = self
            .validators
            .get_mut(&validator_peer_id)
            .ok_or_else(|| anyhow!("no validator with peer_id: {}", validator_peer_id))?;

        if self.fullnodes.contains_key(&validator_peer_id) {
            bail!("VFN for validator {} already configured", validator_peer_id);
        }

        let mut validator_config = validator.config().clone();
        let name = self.node_name_counter.to_string();
        self.node_name_counter += 1;
        let fullnode_config = FullnodeConfig::validator_fullnode(
            name,
            self.dir.as_ref(),
            template,
            &mut validator_config,
            &self.genesis_waypoint,
            &self.genesis,
        )?;

        // Since the validator's config has changed we need to save it
        validator_config.save(validator.config_path())?;
        *validator.config_mut() = validator_config;
        validator.restart()?;

        let version = self.versions.get(version).unwrap();
        let mut fullnode = LocalNode::new(
            version.to_owned(),
            fullnode_config.name,
            fullnode_config.directory,
        )?;

        let peer_id = fullnode.peer_id();
        assert_eq!(peer_id, validator_peer_id);
        fullnode.start()?;

        self.fullnodes.insert(peer_id, fullnode);

        Ok(peer_id)
    }

    fn add_fullnode(&mut self, version: &Version, template: NodeConfig) -> Result<PeerId> {
        let name = self.node_name_counter.to_string();
        self.node_name_counter += 1;
        let fullnode_config = FullnodeConfig::public_fullnode(
            name,
            self.dir.as_ref(),
            template,
            &self.genesis_waypoint,
            &self.genesis,
        )?;

        let version = self.versions.get(version).unwrap();
        let mut fullnode = LocalNode::new(
            version.to_owned(),
            fullnode_config.name,
            fullnode_config.directory,
        )?;

        let peer_id = fullnode.peer_id();
        fullnode.start()?;

        self.fullnodes.insert(peer_id, fullnode);

        Ok(peer_id)
    }

    pub fn chain_id(&self) -> ChainId {
        self.chain_id
    }

    pub fn validator(&self, peer_id: PeerId) -> Option<&LocalNode> {
        self.validators.get(&peer_id)
    }

    pub fn validator_mut(&mut self, peer_id: PeerId) -> Option<&mut LocalNode> {
        self.validators.get_mut(&peer_id)
    }

    pub fn validators(&self) -> impl Iterator<Item = &LocalNode> {
        self.validators.values()
    }

    pub fn validators_mut(&mut self) -> impl Iterator<Item = &mut LocalNode> {
        self.validators.values_mut()
    }

    pub fn dir(&self) -> &Path {
        self.dir.as_ref()
    }
}

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

    fn validators<'a>(&'a 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 upgrade_validator(&mut self, id: PeerId, version: &Version) -> Result<()> {
        let version = self
            .versions
            .get(version)
            .cloned()
            .ok_or_else(|| anyhow!("Invalid version: {:?}", version))?;
        let validator = self
            .validators
            .get_mut(&id)
            .ok_or_else(|| anyhow!("Invalid id: {}", id))?;
        validator.upgrade(version)
    }

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

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

    fn full_node(&self, id: PeerId) -> Option<&dyn FullNode> {
        self.fullnodes.get(&id).map(|v| v as &dyn FullNode)
    }

    fn full_node_mut(&mut self, id: PeerId) -> Option<&mut dyn FullNode> {
        self.fullnodes.get_mut(&id).map(|v| v as &mut dyn FullNode)
    }

    fn add_validator(&mut self, _version: &Version, _template: NodeConfig) -> Result<PeerId> {
        todo!()
    }

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

    fn add_full_node(&mut self, version: &Version, template: NodeConfig) -> Result<PeerId> {
        self.add_fullnode(version, template)
    }

    fn remove_full_node(&mut self, id: PeerId) -> Result<()> {
        if let Some(mut fullnode) = self.fullnodes.remove(&id) {
            fullnode.stop();
        }

        Ok(())
    }

    fn versions<'a>(&'a self) -> Box<dyn Iterator<Item = Version> + 'a> {
        Box::new(self.versions.keys().cloned())
    }

    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,
        )
    }

    fn logs_location(&mut self) -> String {
        self.dir.persist();
        self.dir.display().to_string()
    }
}
<commit_msg>forge: persist logs on panic<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    ChainInfo, FullNode, HealthCheckError, LocalNode, LocalVersion, Node, NodeExt, Swarm, SwarmExt,
    Validator, Version,
};
use anyhow::{anyhow, bail, Result};
use diem_config::config::NodeConfig;
use diem_genesis_tool::{fullnode_builder::FullnodeConfig, validator_builder::ValidatorBuilder};
use diem_sdk::{
    crypto::ed25519::Ed25519PrivateKey,
    types::{
        chain_id::ChainId, transaction::Transaction, waypoint::Waypoint, AccountKey, LocalAccount,
        PeerId,
    },
};
use std::{
    collections::HashMap,
    convert::TryFrom,
    fs, mem,
    num::NonZeroUsize,
    ops,
    path::{Path, PathBuf},
    sync::Arc,
    time::{Duration, Instant},
};
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 {
    versions: Arc<HashMap<Version, LocalVersion>>,
    initial_version: Option<Version>,
    template: NodeConfig,
    number_of_validators: NonZeroUsize,
    dir: Option<PathBuf>,
}

impl LocalSwarmBuilder {
    pub fn new(versions: Arc<HashMap<Version, LocalVersion>>) -> Self {
        Self {
            versions,
            initial_version: None,
            template: NodeConfig::default_for_validator(),
            number_of_validators: NonZeroUsize::new(1).unwrap(),
            dir: None,
        }
    }

    pub fn initial_version(mut self, initial_version: Version) -> Self {
        self.initial_version = Some(initial_version);
        self
    }

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

    pub fn number_of_validators(mut self, number_of_validators: NonZeroUsize) -> 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<R>(self, rng: R) -> Result<LocalSwarm>
    where
        R: ::rand::RngCore + ::rand::CryptoRng,
    {
        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, genesis, genesis_waypoint, validators) = ValidatorBuilder::new(
            &dir,
            diem_framework_releases::current_module_blobs().to_vec(),
        )
        .num_validators(self.number_of_validators)
        .template(self.template)
        .build(rng)?;

        // Get the initial version to start the nodes with, either the one provided or fallback to
        // using the the latest version
        let versions = self.versions;
        let initial_version = self.initial_version.unwrap_or_else(|| {
            versions
                .iter()
                .max_by(|v1, v2| v1.0.cmp(v2.0))
                .unwrap()
                .0
                .clone()
        });
        let version = versions.get(&initial_version).unwrap();

        let validators = validators
            .into_iter()
            .map(|v| {
                let node = LocalNode::new(version.to_owned(), 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),
            1,
        );

        // 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 {
            node_name_counter: validators.len() as u64,
            genesis,
            genesis_waypoint,
            versions,
            validators,
            fullnodes: HashMap::new(),
            dir,
            validator_network_address_encryption_key,
            root_account,
            treasury_compliance_account,
            designated_dealer_account,
            chain_id: ChainId::test(),
        })
    }
}

#[derive(Debug)]
pub struct LocalSwarm {
    node_name_counter: u64,
    genesis: Transaction,
    genesis_waypoint: Waypoint,
    versions: Arc<HashMap<Version, LocalVersion>>,
    validators: HashMap<PeerId, LocalNode>,
    fullnodes: 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(versions: Arc<HashMap<Version, LocalVersion>>) -> LocalSwarmBuilder {
        LocalSwarmBuilder::new(versions)
    }

    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
        let deadline = Instant::now() + Duration::from_secs(60);
        self.wait_for_startup()?;
        self.wait_for_connectivity(deadline)?;
        self.liveness_check(deadline)?;

        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"))
    }

    pub fn add_validator_fullnode(
        &mut self,
        version: &Version,
        template: NodeConfig,
        validator_peer_id: PeerId,
    ) -> Result<PeerId> {
        let validator = self
            .validators
            .get_mut(&validator_peer_id)
            .ok_or_else(|| anyhow!("no validator with peer_id: {}", validator_peer_id))?;

        if self.fullnodes.contains_key(&validator_peer_id) {
            bail!("VFN for validator {} already configured", validator_peer_id);
        }

        let mut validator_config = validator.config().clone();
        let name = self.node_name_counter.to_string();
        self.node_name_counter += 1;
        let fullnode_config = FullnodeConfig::validator_fullnode(
            name,
            self.dir.as_ref(),
            template,
            &mut validator_config,
            &self.genesis_waypoint,
            &self.genesis,
        )?;

        // Since the validator's config has changed we need to save it
        validator_config.save(validator.config_path())?;
        *validator.config_mut() = validator_config;
        validator.restart()?;

        let version = self.versions.get(version).unwrap();
        let mut fullnode = LocalNode::new(
            version.to_owned(),
            fullnode_config.name,
            fullnode_config.directory,
        )?;

        let peer_id = fullnode.peer_id();
        assert_eq!(peer_id, validator_peer_id);
        fullnode.start()?;

        self.fullnodes.insert(peer_id, fullnode);

        Ok(peer_id)
    }

    fn add_fullnode(&mut self, version: &Version, template: NodeConfig) -> Result<PeerId> {
        let name = self.node_name_counter.to_string();
        self.node_name_counter += 1;
        let fullnode_config = FullnodeConfig::public_fullnode(
            name,
            self.dir.as_ref(),
            template,
            &self.genesis_waypoint,
            &self.genesis,
        )?;

        let version = self.versions.get(version).unwrap();
        let mut fullnode = LocalNode::new(
            version.to_owned(),
            fullnode_config.name,
            fullnode_config.directory,
        )?;

        let peer_id = fullnode.peer_id();
        fullnode.start()?;

        self.fullnodes.insert(peer_id, fullnode);

        Ok(peer_id)
    }

    pub fn chain_id(&self) -> ChainId {
        self.chain_id
    }

    pub fn validator(&self, peer_id: PeerId) -> Option<&LocalNode> {
        self.validators.get(&peer_id)
    }

    pub fn validator_mut(&mut self, peer_id: PeerId) -> Option<&mut LocalNode> {
        self.validators.get_mut(&peer_id)
    }

    pub fn validators(&self) -> impl Iterator<Item = &LocalNode> {
        self.validators.values()
    }

    pub fn validators_mut(&mut self) -> impl Iterator<Item = &mut LocalNode> {
        self.validators.values_mut()
    }

    pub fn dir(&self) -> &Path {
        self.dir.as_ref()
    }
}

impl Drop for LocalSwarm {
    fn drop(&mut self) {
        // If panicking, persist logs
        if std::thread::panicking() {
            eprintln!("Logs located at {}", self.logs_location());
        }
    }
}

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

    fn validators<'a>(&'a 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 upgrade_validator(&mut self, id: PeerId, version: &Version) -> Result<()> {
        let version = self
            .versions
            .get(version)
            .cloned()
            .ok_or_else(|| anyhow!("Invalid version: {:?}", version))?;
        let validator = self
            .validators
            .get_mut(&id)
            .ok_or_else(|| anyhow!("Invalid id: {}", id))?;
        validator.upgrade(version)
    }

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

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

    fn full_node(&self, id: PeerId) -> Option<&dyn FullNode> {
        self.fullnodes.get(&id).map(|v| v as &dyn FullNode)
    }

    fn full_node_mut(&mut self, id: PeerId) -> Option<&mut dyn FullNode> {
        self.fullnodes.get_mut(&id).map(|v| v as &mut dyn FullNode)
    }

    fn add_validator(&mut self, _version: &Version, _template: NodeConfig) -> Result<PeerId> {
        todo!()
    }

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

    fn add_full_node(&mut self, version: &Version, template: NodeConfig) -> Result<PeerId> {
        self.add_fullnode(version, template)
    }

    fn remove_full_node(&mut self, id: PeerId) -> Result<()> {
        if let Some(mut fullnode) = self.fullnodes.remove(&id) {
            fullnode.stop();
        }

        Ok(())
    }

    fn versions<'a>(&'a self) -> Box<dyn Iterator<Item = Version> + 'a> {
        Box::new(self.versions.keys().cloned())
    }

    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,
        )
    }

    fn logs_location(&mut self) -> String {
        self.dir.persist();
        self.dir.display().to_string()
    }
}
<|endoftext|>


<commit_before>use super::arc_wake::ArcWake;
use core::mem;
use core::task::{Waker, RawWaker, RawWakerVTable};
use alloc::sync::Arc;

pub(super) fn waker_vtable<W: ArcWake>() -> &'static RawWakerVTable {
    &RawWakerVTable::new(
        clone_arc_raw::<W>,
        wake_arc_raw::<W>,
        wake_by_ref_arc_raw::<W>,
        drop_arc_raw::<W>,
    )
}

/// Creates a [`Waker`] from an `Arc<impl ArcWake>`.
///
/// The returned [`Waker`] will call
/// [`ArcWake.wake()`](ArcWake::wake) if awoken.
pub fn waker<W>(wake: Arc<W>) -> Waker
where
    W: ArcWake,
{
    let ptr = Arc::into_raw(wake) as *const ();

    unsafe {
        Waker::from_raw(RawWaker::new(ptr, waker_vtable::<W>()))
    }
}

// FIXME: panics on Arc::clone / refcount changes could wreak havoc on the
// code here. We should guard against this by aborting.

#[allow(clippy::redundant_clone)] // The clone here isn't actually redundant.
unsafe fn increase_refcount<T: ArcWake>(data: *const ()) {
    // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
    let arc = mem::ManuallyDrop::new(Arc::<T>::from_raw(data as *const T));
    // Now increase refcount, but don't drop new refcount either
    let _arc_clone: mem::ManuallyDrop<_> = arc.clone();
}

// used by `waker_ref`
unsafe fn clone_arc_raw<T: ArcWake>(data: *const ()) -> RawWaker {
    increase_refcount::<T>(data);
    RawWaker::new(data, waker_vtable::<T>())
}

unsafe fn wake_arc_raw<T: ArcWake>(data: *const ()) {
    let arc: Arc<T> = Arc::from_raw(data as *const T);
    ArcWake::wake(arc);
}

// used by `waker_ref`
unsafe fn wake_by_ref_arc_raw<T: ArcWake>(data: *const ()) {
    // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
    let arc = mem::ManuallyDrop::new(Arc::<T>::from_raw(data as *const T));
    ArcWake::wake_by_ref(&arc);
}

unsafe fn drop_arc_raw<T: ArcWake>(data: *const ()) {
    drop(Arc::<T>::from_raw(data as *const T))
}
<commit_msg>Fix UB due to missing 'static on task::waker (#2206)<commit_after>use super::arc_wake::ArcWake;
use core::mem;
use core::task::{Waker, RawWaker, RawWakerVTable};
use alloc::sync::Arc;

pub(super) fn waker_vtable<W: ArcWake>() -> &'static RawWakerVTable {
    &RawWakerVTable::new(
        clone_arc_raw::<W>,
        wake_arc_raw::<W>,
        wake_by_ref_arc_raw::<W>,
        drop_arc_raw::<W>,
    )
}

/// Creates a [`Waker`] from an `Arc<impl ArcWake>`.
///
/// The returned [`Waker`] will call
/// [`ArcWake.wake()`](ArcWake::wake) if awoken.
pub fn waker<W>(wake: Arc<W>) -> Waker
where
    W: ArcWake + 'static,
{
    let ptr = Arc::into_raw(wake) as *const ();

    unsafe {
        Waker::from_raw(RawWaker::new(ptr, waker_vtable::<W>()))
    }
}

// FIXME: panics on Arc::clone / refcount changes could wreak havoc on the
// code here. We should guard against this by aborting.

#[allow(clippy::redundant_clone)] // The clone here isn't actually redundant.
unsafe fn increase_refcount<T: ArcWake>(data: *const ()) {
    // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
    let arc = mem::ManuallyDrop::new(Arc::<T>::from_raw(data as *const T));
    // Now increase refcount, but don't drop new refcount either
    let _arc_clone: mem::ManuallyDrop<_> = arc.clone();
}

// used by `waker_ref`
unsafe fn clone_arc_raw<T: ArcWake>(data: *const ()) -> RawWaker {
    increase_refcount::<T>(data);
    RawWaker::new(data, waker_vtable::<T>())
}

unsafe fn wake_arc_raw<T: ArcWake>(data: *const ()) {
    let arc: Arc<T> = Arc::from_raw(data as *const T);
    ArcWake::wake(arc);
}

// used by `waker_ref`
unsafe fn wake_by_ref_arc_raw<T: ArcWake>(data: *const ()) {
    // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
    let arc = mem::ManuallyDrop::new(Arc::<T>::from_raw(data as *const T));
    ArcWake::wake_by_ref(&arc);
}

unsafe fn drop_arc_raw<T: ArcWake>(data: *const ()) {
    drop(Arc::<T>::from_raw(data as *const T))
}
<|endoftext|>


<commit_before>use std::error::Error;
use std::collections::{VecDeque, vec_deque};
use std::str;
use std::cmp;
use std::mem;

use itertools::Itertools;
use rust_htslib::bcf;

pub fn match_variants(matchbcf: &str, max_dist: u32, max_len_diff: u32) -> Result<(), Box<Error>> {
    let inbcf = try!(bcf::Reader::new(&"-"));
    let mut header = bcf::Header::with_template(&inbcf.header);

    header.push_record(
        format!("##INFO=<ID=MATCHING,Number=A,Type=Integer,\
        Description=\"For each alternative allele, -1 if it does not match a variant in another VCF/BCF. \
        If it matches a variant, an id i>=0 is points to the i-th variant in the VCF/BCF (counting each \
        alternative allele separately). For indels, matching is fuzzy: distance of centres <= {}, difference of \
        lengths <= {}\">", max_dist, max_len_diff).as_bytes()
    );
    header.push_record(
        format!("##rust-bio-tools={}", env!("CARGO_PKG_VERSION")).as_bytes()
    );
    header.push_record(
        b"##rust-bio-tools-subcommand=vcf-match"
    );
    let mut outbcf = try!(bcf::Writer::new(&"-", &header, false, false));
    let mut buffer = RecordBuffer::new(try!(bcf::Reader::new(&matchbcf)), max_dist);

    let mut rec = bcf::Record::new();
    let mut i = 0;
    loop {
        if let Err(e) = inbcf.read(&mut rec) {
            if e.is_eof() {
                break;
            }
            return Err(Box::new(e));
        }
        outbcf.translate(&mut rec);

        if let Some(rid) = rec.rid() {
            let chrom = inbcf.header.rid2name(rid);
            let pos = rec.pos();
            // move buffer to pos
            try!(buffer.fill(chrom, pos));

            let var = try!(Variant::new(&mut rec, &mut i));
            let matching = var.alleles.iter().map(|a| {
                for v in buffer.iter() {
                    if let Some(id) = var.matches(v, a, max_dist, max_len_diff) {
                        return id as i32;
                    }
                }
                -1
            }).collect_vec();

            try!(rec.push_info_integer(b"MATCHING", &matching));
            try!(outbcf.write(&rec));
        } else {
            try!(outbcf.write(&rec));
        }

        if (i) % 1000 == 0 {
            info!("{} variants written.", i);
        }
    }
    info!("{} variants written.", i);

    Ok(())
}

#[derive(Debug)]
pub struct Variant {
    id: u32,
    rid: u32,
    pos: u32,
    alleles: Vec<VariantType>
}

impl Variant {

    pub fn new(rec: &mut bcf::Record, id: &mut u32) -> Result<Self, Box<Error>> {
        let pos = rec.pos();

        let svlen = if let Ok(Some(svlen)) = rec.info(b"SVLEN").integer() {
            Some(svlen.to_owned())
        } else { None };
        let svtype = if let Ok(Some(svtype)) = rec.info(b"SVTYPE").string() {
            Some(svtype.into_iter().map(|s| s.to_owned()).collect_vec())
        } else { None };

        let mut _alleles = Vec::new();
        let alleles = rec.alleles();
        let refallele = alleles[0];
        for (i, &a) in alleles[1..].iter().enumerate() {
            let vartype = match (&svlen, &svtype) {

                (&Some(ref svlen), &Some(ref svtype)) => {
                    let svlen = svlen.get(i).unwrap_or(&svlen[0]).abs();
                    let svtype = svtype.get(i).unwrap_or(&svtype[0]);
                    match &svtype[..] {
                        b"INS" => VariantType::Insertion(svlen as u32),
                        b"DEL" => VariantType::Deletion(svlen as u32),
                        t => return Err(Box::new(MatchError::UnsupportedVariant(try!(str::from_utf8(t)).to_owned())))
                    }
                },

                (&Some(ref svlen), _) if a[0] == b'<' => {
                    let svlen = svlen.get(i).unwrap_or(&svlen[0]).abs();
                    match a {
                        b"<DEL>" => VariantType::Deletion(svlen as u32),
                        b"<INS>" => VariantType::Insertion(svlen as u32),
                        a => return Err(Box::new(MatchError::UnsupportedVariant(try!(str::from_utf8(a)).to_owned())))
                    }
                },

                _ => {
                    if a.len() < refallele.len() {
                        VariantType::Deletion((refallele.len() - a.len()) as u32)
                    } else if a.len() > refallele.len() {
                        VariantType::Insertion((a.len() - refallele.len()) as u32)
                    } else if a.len() == 1 {
                        VariantType::SNV(a[0])
                    } else {
                        return Err(Box::new(MatchError::UnsupportedVariant("complex".to_owned())))
                    }
                }
            };
            _alleles.push(vartype);
        }
        let var = Variant {
            id: *id,
            rid: rec.rid().unwrap(),
            pos: pos,
            alleles: _alleles
        };
        *id += alleles.len() as u32 - 1;
        Ok(var)
    }

    pub fn centerpoint(&self, allele: &VariantType) -> u32 {
        match allele {
            &VariantType::SNV(_) => self.pos,
            &VariantType::Insertion(_) => self.pos,
            &VariantType::Deletion(len) => (self.pos as f64 + len as f64 / 2.0) as u32
        }
    }

    pub fn matches(&self, other: &Variant, allele: &VariantType, max_dist: u32, max_len_diff: u32) -> Option<u32> {
        for (j, b) in other.alleles.iter().enumerate() {
            let dist = (self.centerpoint(allele) as i32 - other.centerpoint(b) as i32).abs() as u32;
            match (allele, b) {
                (&VariantType::SNV(a), &VariantType::SNV(b)) => {
                    if a == b && dist == 0 {
                        return Some(other.id(j));
                    }
                },
                (&VariantType::Insertion(l1), &VariantType::Insertion(l2)) |
                (&VariantType::Deletion(l1), &VariantType::Deletion(l2)) => {
                    if (l1 as i32 - l2 as i32).abs() as u32 <= max_len_diff && dist <= max_dist {
                        return Some(other.id(j));
                    }
                },
                _ => continue
            }
        }
        None
    }

    pub fn id(&self, allele: usize) -> u32 {
        self.id + allele as u32
    }
}

#[derive(Debug)]
pub enum VariantType {
    SNV(u8),
    Insertion(u32),
    Deletion(u32)
}

quick_error! {
    #[derive(Debug)]
    pub enum MatchError {
        UnsupportedVariant(vartype: String) {
            description("unsupported variant")
            display("variant type {} is not supported", vartype)
        }
    }
}

pub struct RecordBuffer {
    reader: bcf::Reader,
    ringbuffer: VecDeque<Variant>,
    ringbuffer2: VecDeque<Variant>,
    window: u32,
    i: u32
}

impl RecordBuffer {
    pub fn new(reader: bcf::Reader, window: u32) -> Self {
        RecordBuffer {
            reader: reader,
            ringbuffer: VecDeque::new(),
            ringbuffer2: VecDeque::new(),
            window: window,
            i: 0
        }
    }

    pub fn last_rid(&self) -> Option<u32> {
        self.ringbuffer.back().map(|var| var.rid)
    }

    pub fn next_rid(&self) -> Option<u32> {
        self.ringbuffer2.back().map(|var| var.rid)
    }

    fn swap_buffers(&mut self) {
        // swap with buffer for next rid
        mem::swap(&mut self.ringbuffer2, &mut self.ringbuffer);
        // clear second buffer
        self.ringbuffer2.clear();
    }

    fn drain_left(&mut self, rid: u32, window_start: u32) {
        // remove records too far left or from wrong rid
        let to_remove = self.ringbuffer.iter().take_while(|var| var.pos < window_start || var.rid != rid).count();
        self.ringbuffer.drain(..to_remove);
    }

    pub fn fill(&mut self, chrom: &[u8], pos: u32) -> Result<(), Box<Error>> {
        let window_start = cmp::max(pos as i32 - self.window as i32 - 1, 0) as u32;
        let window_end = pos + self.window;
        let rid = try!(self.reader.header.name2rid(chrom));

        match (self.last_rid(), self.next_rid()) {
            (Some(last_rid), _) => {
                if last_rid != rid {
                    self.swap_buffers();
                } else {
                    self.drain_left(rid, window_start);
                }
            },
            (_, Some(_)) => {
                self.swap_buffers();
                self.drain_left(rid, window_start);
            },
            _ => ()
        }

        if !self.ringbuffer2.is_empty() {
            // We have already read beyond the current rid. Hence we can't extend to the right for
            // this rid.
            return Ok(())
        }

        // extend to the right
        let mut rec = bcf::Record::new();
        loop {
            if let Err(e) = self.reader.read(&mut rec) {
                if e.is_eof() {
                    break;
                }
                return Err(Box::new(e));
            }
            let pos = rec.pos();
            if let Some(rec_rid) = rec.rid() {
                if rec_rid == rid {
                    if pos > window_end {
                        // Record is beyond our window. Store it anyways but stop.
                        self.ringbuffer.push_back(try!(Variant::new(&mut rec, &mut self.i)));
                        break;
                    } else if pos >= window_start {
                        // Record is within out window.
                        self.ringbuffer.push_back(try!(Variant::new(&mut rec, &mut self.i)));
                    }
                } else if rec_rid > rid {
                    // record comes from next rid. Store it in second buffer but stop filling.
                    self.ringbuffer2.push_back(try!(Variant::new(&mut rec, &mut self.i)));
                    break;
                } else {
                    // Record comes from previous rid. Ignore it.
                    continue;
                }
            } else {
                // skip records without proper rid
                continue;
            }
        }

        Ok(())
    }

    pub fn iter(&self) -> vec_deque::Iter<Variant> {
        self.ringbuffer.iter()
    }
}
<commit_msg>Be tolerant in case of complex variants.<commit_after>use std::error::Error;
use std::collections::{VecDeque, vec_deque};
use std::str;
use std::cmp;
use std::mem;

use itertools::Itertools;
use rust_htslib::bcf;

pub fn match_variants(matchbcf: &str, max_dist: u32, max_len_diff: u32) -> Result<(), Box<Error>> {
    let inbcf = try!(bcf::Reader::new(&"-"));
    let mut header = bcf::Header::with_template(&inbcf.header);

    header.push_record(
        format!("##INFO=<ID=MATCHING,Number=A,Type=Integer,\
        Description=\"For each alternative allele, -1 if it does not match a variant in another VCF/BCF. \
        If it matches a variant, an id i>=0 is points to the i-th variant in the VCF/BCF (counting each \
        alternative allele separately). For indels, matching is fuzzy: distance of centres <= {}, difference of \
        lengths <= {}\">", max_dist, max_len_diff).as_bytes()
    );
    header.push_record(
        format!("##rust-bio-tools={}", env!("CARGO_PKG_VERSION")).as_bytes()
    );
    header.push_record(
        b"##rust-bio-tools-subcommand=vcf-match"
    );
    let mut outbcf = try!(bcf::Writer::new(&"-", &header, false, false));
    let mut buffer = RecordBuffer::new(try!(bcf::Reader::new(&matchbcf)), max_dist);

    let mut rec = bcf::Record::new();
    let mut i = 0;
    loop {
        if let Err(e) = inbcf.read(&mut rec) {
            if e.is_eof() {
                break;
            }
            return Err(Box::new(e));
        }
        outbcf.translate(&mut rec);

        if let Some(rid) = rec.rid() {
            let chrom = inbcf.header.rid2name(rid);
            let pos = rec.pos();
            // move buffer to pos
            try!(buffer.fill(chrom, pos));

            let var = try!(Variant::new(&mut rec, &mut i));
            let matching = var.alleles.iter().map(|a| {
                for v in buffer.iter() {
                    if let Some(id) = var.matches(v, a, max_dist, max_len_diff) {
                        return id as i32;
                    }
                }
                -1
            }).collect_vec();

            try!(rec.push_info_integer(b"MATCHING", &matching));
            try!(outbcf.write(&rec));
        } else {
            try!(outbcf.write(&rec));
        }

        if (i) % 1000 == 0 {
            info!("{} variants written.", i);
        }
    }
    info!("{} variants written.", i);

    Ok(())
}

#[derive(Debug)]
pub struct Variant {
    id: u32,
    rid: u32,
    pos: u32,
    alleles: Vec<VariantType>
}

impl Variant {

    pub fn new(rec: &mut bcf::Record, id: &mut u32) -> Result<Self, Box<Error>> {
        let pos = rec.pos();

        let svlen = if let Ok(Some(svlen)) = rec.info(b"SVLEN").integer() {
            Some(svlen.to_owned())
        } else { None };
        let svtype = if let Ok(Some(svtype)) = rec.info(b"SVTYPE").string() {
            Some(svtype.into_iter().map(|s| s.to_owned()).collect_vec())
        } else { None };

        let mut _alleles = Vec::new();
        let alleles = rec.alleles();
        let refallele = alleles[0];
        for (i, &a) in alleles[1..].iter().enumerate() {
            let vartype = match (&svlen, &svtype) {

                (&Some(ref svlen), &Some(ref svtype)) => {
                    let svlen = svlen.get(i).unwrap_or(&svlen[0]).abs();
                    let svtype = svtype.get(i).unwrap_or(&svtype[0]);
                    match &svtype[..] {
                        b"INS" => VariantType::Insertion(svlen as u32),
                        b"DEL" => VariantType::Deletion(svlen as u32),
                        t => {
                            info!("Unsupported variant {}", try!(str::from_utf8(t)));
                            VariantType::Complex(svlen as u32)
                        }
                    }
                },

                (&Some(ref svlen), _) if a[0] == b'<' => {
                    let svlen = svlen.get(i).unwrap_or(&svlen[0]).abs();
                    match a {
                        b"<DEL>" => VariantType::Deletion(svlen as u32),
                        b"<INS>" => VariantType::Insertion(svlen as u32),
                        a => {
                            info!("Unsupported variant {}", try!(str::from_utf8(a)));
                            VariantType::Complex(svlen as u32)
                        }
                    }
                },

                _ => {
                    if a.len() < refallele.len() {
                        VariantType::Deletion((refallele.len() - a.len()) as u32)
                    } else if a.len() > refallele.len() {
                        VariantType::Insertion((a.len() - refallele.len()) as u32)
                    } else if a.len() == 1 {
                        VariantType::SNV(a[0])
                    } else {
                        info!("Unsupported variant {} -> {}", try!(str::from_utf8(refallele)), try!(str::from_utf8(a)));
                        VariantType::Complex(a.len() as u32)
                    }
                }
            };
            _alleles.push(vartype);
        }
        let var = Variant {
            id: *id,
            rid: rec.rid().unwrap(),
            pos: pos,
            alleles: _alleles
        };
        *id += alleles.len() as u32 - 1;
        Ok(var)
    }

    pub fn centerpoint(&self, allele: &VariantType) -> u32 {
        match allele {
            &VariantType::SNV(_) => self.pos,
            &VariantType::Insertion(_) => self.pos,
            &VariantType::Deletion(len) => (self.pos as f64 + len as f64 / 2.0) as u32,
            &VariantType::Complex(_) => panic!("Unsupported variant.")
        }
    }

    pub fn matches(&self, other: &Variant, allele: &VariantType, max_dist: u32, max_len_diff: u32) -> Option<u32> {
        for (j, b) in other.alleles.iter().enumerate() {
            let dist = (self.centerpoint(allele) as i32 - other.centerpoint(b) as i32).abs() as u32;
            match (allele, b) {
                (&VariantType::SNV(a), &VariantType::SNV(b)) => {
                    if a == b && dist == 0 {
                        return Some(other.id(j));
                    }
                },
                (&VariantType::Insertion(l1), &VariantType::Insertion(l2)) |
                (&VariantType::Deletion(l1), &VariantType::Deletion(l2)) => {
                    if (l1 as i32 - l2 as i32).abs() as u32 <= max_len_diff && dist <= max_dist {
                        return Some(other.id(j));
                    }
                },
                // TODO: for now, ignore complex variants
                _ => continue
            }
        }
        None
    }

    pub fn id(&self, allele: usize) -> u32 {
        self.id + allele as u32
    }
}

#[derive(Debug)]
pub enum VariantType {
    SNV(u8),
    Insertion(u32),
    Deletion(u32),
    Complex(u32)
}

quick_error! {
    #[derive(Debug)]
    pub enum MatchError {
        UnsupportedVariant(vartype: String) {
            description("unsupported variant")
            display("variant type {} is not supported", vartype)
        }
    }
}

pub struct RecordBuffer {
    reader: bcf::Reader,
    ringbuffer: VecDeque<Variant>,
    ringbuffer2: VecDeque<Variant>,
    window: u32,
    i: u32
}

impl RecordBuffer {
    pub fn new(reader: bcf::Reader, window: u32) -> Self {
        RecordBuffer {
            reader: reader,
            ringbuffer: VecDeque::new(),
            ringbuffer2: VecDeque::new(),
            window: window,
            i: 0
        }
    }

    pub fn last_rid(&self) -> Option<u32> {
        self.ringbuffer.back().map(|var| var.rid)
    }

    pub fn next_rid(&self) -> Option<u32> {
        self.ringbuffer2.back().map(|var| var.rid)
    }

    fn swap_buffers(&mut self) {
        // swap with buffer for next rid
        mem::swap(&mut self.ringbuffer2, &mut self.ringbuffer);
        // clear second buffer
        self.ringbuffer2.clear();
    }

    fn drain_left(&mut self, rid: u32, window_start: u32) {
        // remove records too far left or from wrong rid
        let to_remove = self.ringbuffer.iter().take_while(|var| var.pos < window_start || var.rid != rid).count();
        self.ringbuffer.drain(..to_remove);
    }

    pub fn fill(&mut self, chrom: &[u8], pos: u32) -> Result<(), Box<Error>> {
        let window_start = cmp::max(pos as i32 - self.window as i32 - 1, 0) as u32;
        let window_end = pos + self.window;
        let rid = try!(self.reader.header.name2rid(chrom));

        match (self.last_rid(), self.next_rid()) {
            (Some(last_rid), _) => {
                if last_rid != rid {
                    self.swap_buffers();
                } else {
                    self.drain_left(rid, window_start);
                }
            },
            (_, Some(_)) => {
                self.swap_buffers();
                self.drain_left(rid, window_start);
            },
            _ => ()
        }

        if !self.ringbuffer2.is_empty() {
            // We have already read beyond the current rid. Hence we can't extend to the right for
            // this rid.
            return Ok(())
        }

        // extend to the right
        let mut rec = bcf::Record::new();
        loop {
            if let Err(e) = self.reader.read(&mut rec) {
                if e.is_eof() {
                    break;
                }
                return Err(Box::new(e));
            }
            let pos = rec.pos();
            if let Some(rec_rid) = rec.rid() {
                if rec_rid == rid {
                    if pos > window_end {
                        // Record is beyond our window. Store it anyways but stop.
                        self.ringbuffer.push_back(try!(Variant::new(&mut rec, &mut self.i)));
                        break;
                    } else if pos >= window_start {
                        // Record is within out window.
                        self.ringbuffer.push_back(try!(Variant::new(&mut rec, &mut self.i)));
                    }
                } else if rec_rid > rid {
                    // record comes from next rid. Store it in second buffer but stop filling.
                    self.ringbuffer2.push_back(try!(Variant::new(&mut rec, &mut self.i)));
                    break;
                } else {
                    // Record comes from previous rid. Ignore it.
                    continue;
                }
            } else {
                // skip records without proper rid
                continue;
            }
        }

        Ok(())
    }

    pub fn iter(&self) -> vec_deque::Iter<Variant> {
        self.ringbuffer.iter()
    }
}
<|endoftext|>


<commit_before>use hir_def::{
    attr::Attrs,
    db::DefDatabase,
    docs::Documentation,
    resolver::{HasResolver, Resolver},
    AdtId, FunctionId, GenericDefId, ModuleId, StaticId, TraitId, VariantId,
};
use hir_ty::db::HirDatabase;
use stdx::impl_from;

use crate::{
    doc_links::Resolvable, Adt, Const, Enum, EnumVariant, Field, Function, GenericDef, ImplDef,
    Local, MacroDef, Module, ModuleDef, Static, Struct, Trait, TypeAlias, TypeParam, Union,
    VariantDef,
};

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum AttrDef {
    Module(Module),
    Field(Field),
    Adt(Adt),
    Function(Function),
    EnumVariant(EnumVariant),
    Static(Static),
    Const(Const),
    Trait(Trait),
    TypeAlias(TypeAlias),
    MacroDef(MacroDef),
}

impl_from!(
    Module,
    Field,
    Adt(Struct, Enum, Union),
    EnumVariant,
    Static,
    Const,
    Function,
    Trait,
    TypeAlias,
    MacroDef
    for AttrDef
);

pub trait HasAttrs {
    fn attrs(self, db: &dyn HirDatabase) -> Attrs;
    fn docs(self, db: &dyn HirDatabase) -> Option<Documentation>;
}

impl<T: Into<AttrDef>> HasAttrs for T {
    fn attrs(self, db: &dyn HirDatabase) -> Attrs {
        let def: AttrDef = self.into();
        db.attrs(def.into())
    }
    fn docs(self, db: &dyn HirDatabase) -> Option<Documentation> {
        let def: AttrDef = self.into();
        db.documentation(def.into())
    }
}

impl Resolvable for ModuleDef {
    fn resolver<D: DefDatabase + HirDatabase>(&self, db: &D) -> Option<Resolver> {
        Some(match self {
            ModuleDef::Module(m) => ModuleId::from(m.clone()).resolver(db),
            ModuleDef::Function(f) => FunctionId::from(f.clone()).resolver(db),
            ModuleDef::Adt(adt) => AdtId::from(adt.clone()).resolver(db),
            ModuleDef::EnumVariant(ev) => {
                GenericDefId::from(GenericDef::from(ev.clone())).resolver(db)
            }
            ModuleDef::Const(c) => GenericDefId::from(GenericDef::from(c.clone())).resolver(db),
            ModuleDef::Static(s) => StaticId::from(s.clone()).resolver(db),
            ModuleDef::Trait(t) => TraitId::from(t.clone()).resolver(db),
            ModuleDef::TypeAlias(t) => ModuleId::from(t.module(db)).resolver(db),
            // FIXME: This should be a resolver relative to `std/core`
            ModuleDef::BuiltinType(_t) => None?,
        })
    }

    fn try_into_module_def(self) -> Option<ModuleDef> {
        Some(self)
    }
}

impl Resolvable for TypeParam {
    fn resolver<D: DefDatabase + HirDatabase>(&self, db: &D) -> Option<Resolver> {
        Some(Into::<ModuleId>::into(self.module(db)).resolver(db))
    }

    fn try_into_module_def(self) -> Option<ModuleDef> {
        None
    }
}

impl Resolvable for MacroDef {
    fn resolver<D: DefDatabase + HirDatabase>(&self, db: &D) -> Option<Resolver> {
        Some(Into::<ModuleId>::into(self.module(db)?).resolver(db))
    }

    fn try_into_module_def(self) -> Option<ModuleDef> {
        None
    }
}

impl Resolvable for Field {
    fn resolver<D: DefDatabase + HirDatabase>(&self, db: &D) -> Option<Resolver> {
        Some(Into::<VariantId>::into(Into::<VariantDef>::into(self.parent_def(db))).resolver(db))
    }

    fn try_into_module_def(self) -> Option<ModuleDef> {
        None
    }
}

impl Resolvable for ImplDef {
    fn resolver<D: DefDatabase + HirDatabase>(&self, db: &D) -> Option<Resolver> {
        Some(Into::<ModuleId>::into(self.module(db)).resolver(db))
    }

    fn try_into_module_def(self) -> Option<ModuleDef> {
        None
    }
}

impl Resolvable for Local {
    fn resolver<D: DefDatabase + HirDatabase>(&self, db: &D) -> Option<Resolver> {
        Some(Into::<ModuleId>::into(self.module(db)).resolver(db))
    }

    fn try_into_module_def(self) -> Option<ModuleDef> {
        None
    }
}
<commit_msg>Simplify<commit_after>//! Attributes & documentation for hir types.
use hir_def::{
    attr::Attrs,
    db::DefDatabase,
    docs::Documentation,
    resolver::{HasResolver, Resolver},
    AdtId, FunctionId, GenericDefId, ModuleId, StaticId, TraitId, VariantId,
};
use hir_ty::db::HirDatabase;
use stdx::impl_from;

use crate::{
    doc_links::Resolvable, Adt, Const, Enum, EnumVariant, Field, Function, GenericDef, ImplDef,
    Local, MacroDef, Module, ModuleDef, Static, Struct, Trait, TypeAlias, TypeParam, Union,
};

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum AttrDef {
    Module(Module),
    Field(Field),
    Adt(Adt),
    Function(Function),
    EnumVariant(EnumVariant),
    Static(Static),
    Const(Const),
    Trait(Trait),
    TypeAlias(TypeAlias),
    MacroDef(MacroDef),
}

impl_from!(
    Module,
    Field,
    Adt(Struct, Enum, Union),
    EnumVariant,
    Static,
    Const,
    Function,
    Trait,
    TypeAlias,
    MacroDef
    for AttrDef
);

pub trait HasAttrs {
    fn attrs(self, db: &dyn HirDatabase) -> Attrs;
    fn docs(self, db: &dyn HirDatabase) -> Option<Documentation>;
}

impl<T: Into<AttrDef>> HasAttrs for T {
    fn attrs(self, db: &dyn HirDatabase) -> Attrs {
        let def: AttrDef = self.into();
        db.attrs(def.into())
    }
    fn docs(self, db: &dyn HirDatabase) -> Option<Documentation> {
        let def: AttrDef = self.into();
        db.documentation(def.into())
    }
}

impl Resolvable for ModuleDef {
    fn resolver<D: DefDatabase + HirDatabase>(&self, db: &D) -> Option<Resolver> {
        Some(match self {
            ModuleDef::Module(m) => ModuleId::from(m.clone()).resolver(db),
            ModuleDef::Function(f) => FunctionId::from(f.clone()).resolver(db),
            ModuleDef::Adt(adt) => AdtId::from(adt.clone()).resolver(db),
            ModuleDef::EnumVariant(ev) => {
                GenericDefId::from(GenericDef::from(ev.clone())).resolver(db)
            }
            ModuleDef::Const(c) => GenericDefId::from(GenericDef::from(c.clone())).resolver(db),
            ModuleDef::Static(s) => StaticId::from(s.clone()).resolver(db),
            ModuleDef::Trait(t) => TraitId::from(t.clone()).resolver(db),
            ModuleDef::TypeAlias(t) => ModuleId::from(t.module(db)).resolver(db),
            // FIXME: This should be a resolver relative to `std/core`
            ModuleDef::BuiltinType(_t) => None?,
        })
    }

    fn try_into_module_def(self) -> Option<ModuleDef> {
        Some(self)
    }
}

impl Resolvable for TypeParam {
    fn resolver<D: DefDatabase + HirDatabase>(&self, db: &D) -> Option<Resolver> {
        Some(ModuleId::from(self.module(db)).resolver(db))
    }

    fn try_into_module_def(self) -> Option<ModuleDef> {
        None
    }
}

impl Resolvable for MacroDef {
    fn resolver<D: DefDatabase + HirDatabase>(&self, db: &D) -> Option<Resolver> {
        Some(ModuleId::from(self.module(db)?).resolver(db))
    }

    fn try_into_module_def(self) -> Option<ModuleDef> {
        None
    }
}

impl Resolvable for Field {
    fn resolver<D: DefDatabase + HirDatabase>(&self, db: &D) -> Option<Resolver> {
        Some(VariantId::from(self.parent_def(db)).resolver(db))
    }

    fn try_into_module_def(self) -> Option<ModuleDef> {
        None
    }
}

impl Resolvable for ImplDef {
    fn resolver<D: DefDatabase + HirDatabase>(&self, db: &D) -> Option<Resolver> {
        Some(ModuleId::from(self.module(db)).resolver(db))
    }

    fn try_into_module_def(self) -> Option<ModuleDef> {
        None
    }
}

impl Resolvable for Local {
    fn resolver<D: DefDatabase + HirDatabase>(&self, db: &D) -> Option<Resolver> {
        Some(ModuleId::from(self.module(db)).resolver(db))
    }

    fn try_into_module_def(self) -> Option<ModuleDef> {
        None
    }
}
<|endoftext|>