<commit_before>// Copyright (C) 2017, 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0
use argh::FromArgs;
use slog_scope::info;
use std::path::PathBuf;

#[derive(FromArgs)]
/// Top-level command.
struct TopLevel {
    #[argh(subcommand)]
    entry_point: EntryPoints,
}

#[derive(FromArgs)]
#[argh(subcommand)]
enum EntryPoints {
    Client(ClientOptions),
    Server(ServerOptions),
}

#[derive(FromArgs)]
/// Client subcommand
#[argh(subcommand, name = "client")]
struct ClientOptions {
    #[argh(subcommand)]
    commands: ClientCommands,
}

#[derive(FromArgs)]
#[argh(subcommand)]
enum ClientCommands {
    Info(Info),
    Log(Log),
    Probe(Probe),
    AbortDownload(AbortDownload),
    LocalInstall(LocalInstall),
    RemoteInstall(RemoteInstall),
}

#[derive(FromArgs)]
/// Fetches information about the current state of the agent
#[argh(subcommand, name = "info")]
struct Info {}

#[derive(FromArgs)]
/// Fetches the available log entries for the last update cycle
#[argh(subcommand, name = "log")]
struct Log {}

#[derive(FromArgs)]
/// Checks if the server has a new update for this device.
///
/// A custom server for the update cycle can be specified via the ´--server´
#[argh(subcommand, name = "probe")]
struct Probe {
    /// custom address to try probe
    #[argh(option)]
    server: Option<String>,
}

#[derive(FromArgs)]
/// Abort current running download
#[argh(subcommand, name = "abort-download")]
struct AbortDownload {}

#[derive(FromArgs)]
/// Request agent to install a local update package
#[argh(subcommand, name = "local-install")]
struct LocalInstall {
    /// path to the update package
    #[argh(positional)]
    file: PathBuf,
}

#[derive(FromArgs)]
/// Request agent to download and install a package from a direct URL
#[argh(subcommand, name = "remote-install")]
struct RemoteInstall {
    /// the URL to get the update package
    #[argh(positional)]
    url: String,
}

#[derive(FromArgs)]
/// Server subcommand
#[argh(subcommand, name = "server")]
struct ServerOptions {
    /// increase the verboseness level
    #[argh(option, short = 'v', from_str_fn(verbosity_level), default = "slog::Level::Info")]
    verbosity: slog::Level,

    /// configuration file to use (defaults to "/etc/updatehub.conf")
    #[argh(option, short = 'c', default = "PathBuf::from(\"/etc/updatehub.conf\")")]
    config: PathBuf,
}

fn verbosity_level(value: &str) -> Result<slog::Level, String> {
    use std::str::FromStr;
    slog::Level::from_str(value).map_err(|_| format!("failed to parse verbosity level: {}", value))
}

async fn server_main(cmd: ServerOptions) -> updatehub::Result<()> {
    let _guard = updatehub::logger::init(cmd.verbosity);
    info!("starting UpdateHub Agent {}", updatehub::version());

    updatehub::run(&cmd.config).await?;

    Ok(())
}

async fn client_main(cmd: ClientCommands) -> updatehub::Result<()> {
    let client = sdk::Client::new("localhost:8080");

    match cmd {
        ClientCommands::Info(_) => println!("{:#?}", client.info().await),
        ClientCommands::Log(_) => {
            for entry in client.log().await?.into_iter() {
                let level = match entry.level {
                    sdk::api::log::Level::Critical => "CRIT",
                    sdk::api::log::Level::Error => "ERRO",
                    sdk::api::log::Level::Warning => "WARN",
                    sdk::api::log::Level::Info => "INFO",
                    sdk::api::log::Level::Debug => "DEBG",
                    sdk::api::log::Level::Trace => "TRCE",
                };

                println!(
                    "{timestamp} {level} {msg}",
                    timestamp = entry.time,
                    level = level,
                    msg = entry.message
                );
            }
        }
        ClientCommands::Probe(Probe { server }) => {
            let response = client.probe(server).await?;

            match (response.update_available, response.try_again_in) {
                (true, None) => println!("Update available. The update is running in background."),
                (false, None) => println!("There are no updates available."),
                (false, Some(t)) => {
                    println!("Server replied asking us to try again in {} seconds", t);
                }
                (true, Some(_)) => {
                    unreachable!("We cannot have an update and receive a try again interval");
                }
            }
        }
        ClientCommands::AbortDownload(_) => println!("{:#?}", client.abort_download().await),
        ClientCommands::LocalInstall(LocalInstall { file }) => {
            let file =
                if file.is_absolute() { file } else { std::env::current_dir().unwrap().join(file) };
            println!("{:#?}", client.local_install(&file).await)
        }
        ClientCommands::RemoteInstall(RemoteInstall { url }) => {
            println!("{:#?}", client.remote_install(&url).await)
        }
    }

    Ok(())
}

#[async_std::main]
async fn main() {
    let cmd: TopLevel = argh::from_env();

    let res = match cmd.entry_point {
        EntryPoints::Client(client) => client_main(client.commands).await,
        EntryPoints::Server(cmd) => server_main(cmd).await,
    };

    if let Err(e) = res {
        eprintln!("{}", e);
        std::process::exit(1);
    }
}
<commit_msg>updatehub: Change to print the debug response instead result<commit_after>// Copyright (C) 2017, 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0
use argh::FromArgs;
use slog_scope::info;
use std::path::PathBuf;

#[derive(FromArgs)]
/// Top-level command.
struct TopLevel {
    #[argh(subcommand)]
    entry_point: EntryPoints,
}

#[derive(FromArgs)]
#[argh(subcommand)]
enum EntryPoints {
    Client(ClientOptions),
    Server(ServerOptions),
}

#[derive(FromArgs)]
/// Client subcommand
#[argh(subcommand, name = "client")]
struct ClientOptions {
    #[argh(subcommand)]
    commands: ClientCommands,
}

#[derive(FromArgs)]
#[argh(subcommand)]
enum ClientCommands {
    Info(Info),
    Log(Log),
    Probe(Probe),
    AbortDownload(AbortDownload),
    LocalInstall(LocalInstall),
    RemoteInstall(RemoteInstall),
}

#[derive(FromArgs)]
/// Fetches information about the current state of the agent
#[argh(subcommand, name = "info")]
struct Info {}

#[derive(FromArgs)]
/// Fetches the available log entries for the last update cycle
#[argh(subcommand, name = "log")]
struct Log {}

#[derive(FromArgs)]
/// Checks if the server has a new update for this device.
///
/// A custom server for the update cycle can be specified via the ´--server´
#[argh(subcommand, name = "probe")]
struct Probe {
    /// custom address to try probe
    #[argh(option)]
    server: Option<String>,
}

#[derive(FromArgs)]
/// Abort current running download
#[argh(subcommand, name = "abort-download")]
struct AbortDownload {}

#[derive(FromArgs)]
/// Request agent to install a local update package
#[argh(subcommand, name = "local-install")]
struct LocalInstall {
    /// path to the update package
    #[argh(positional)]
    file: PathBuf,
}

#[derive(FromArgs)]
/// Request agent to download and install a package from a direct URL
#[argh(subcommand, name = "remote-install")]
struct RemoteInstall {
    /// the URL to get the update package
    #[argh(positional)]
    url: String,
}

#[derive(FromArgs)]
/// Server subcommand
#[argh(subcommand, name = "server")]
struct ServerOptions {
    /// increase the verboseness level
    #[argh(option, short = 'v', from_str_fn(verbosity_level), default = "slog::Level::Info")]
    verbosity: slog::Level,

    /// configuration file to use (defaults to "/etc/updatehub.conf")
    #[argh(option, short = 'c', default = "PathBuf::from(\"/etc/updatehub.conf\")")]
    config: PathBuf,
}

fn verbosity_level(value: &str) -> Result<slog::Level, String> {
    use std::str::FromStr;
    slog::Level::from_str(value).map_err(|_| format!("failed to parse verbosity level: {}", value))
}

async fn server_main(cmd: ServerOptions) -> updatehub::Result<()> {
    let _guard = updatehub::logger::init(cmd.verbosity);
    info!("starting UpdateHub Agent {}", updatehub::version());

    updatehub::run(&cmd.config).await?;

    Ok(())
}

async fn client_main(cmd: ClientCommands) -> updatehub::Result<()> {
    let client = sdk::Client::new("localhost:8080");

    match cmd {
        ClientCommands::Info(_) => println!("{:#?}", client.info().await?),
        ClientCommands::Log(_) => {
            for entry in client.log().await?.into_iter() {
                let level = match entry.level {
                    sdk::api::log::Level::Critical => "CRIT",
                    sdk::api::log::Level::Error => "ERRO",
                    sdk::api::log::Level::Warning => "WARN",
                    sdk::api::log::Level::Info => "INFO",
                    sdk::api::log::Level::Debug => "DEBG",
                    sdk::api::log::Level::Trace => "TRCE",
                };

                println!(
                    "{timestamp} {level} {msg}",
                    timestamp = entry.time,
                    level = level,
                    msg = entry.message
                );
            }
        }
        ClientCommands::Probe(Probe { server }) => {
            let response = client.probe(server).await?;

            match (response.update_available, response.try_again_in) {
                (true, None) => println!("Update available. The update is running in background."),
                (false, None) => println!("There are no updates available."),
                (false, Some(t)) => {
                    println!("Server replied asking us to try again in {} seconds", t);
                }
                (true, Some(_)) => {
                    unreachable!("We cannot have an update and receive a try again interval");
                }
            }
        }
        ClientCommands::AbortDownload(_) => println!("{:#?}", client.abort_download().await?),
        ClientCommands::LocalInstall(LocalInstall { file }) => {
            let file =
                if file.is_absolute() { file } else { std::env::current_dir().unwrap().join(file) };
            println!("{:#?}", client.local_install(&file).await?)
        }
        ClientCommands::RemoteInstall(RemoteInstall { url }) => {
            println!("{:#?}", client.remote_install(&url).await?)
        }
    }

    Ok(())
}

#[async_std::main]
async fn main() {
    let cmd: TopLevel = argh::from_env();

    let res = match cmd.entry_point {
        EntryPoints::Client(client) => client_main(client.commands).await,
        EntryPoints::Server(cmd) => server_main(cmd).await,
    };

    if let Err(e) = res {
        eprintln!("{}", e);
        std::process::exit(1);
    }
}
<|endoftext|>


<commit_before>//! Encoding of PPM Images

use std::io;
use std::io::Write;

use color;
use color::ColorType:: {
    Gray,
    Palette,
    GrayA,
    RGB,
    RGBA,
};

/// A representation of a PPM encoder.
pub struct PPMEncoder<'a, W: 'a> {
    w: &'a mut W
}

impl<'a, W: Write> PPMEncoder<'a, W> {
    /// Create a new PPMEncoder from the Writer ```w```.
    /// This function takes ownership of the Writer.
    pub fn new(w: &mut W) -> PPMEncoder<W> {
        PPMEncoder { w: w }
    }

    /// Encode the buffer ```im``` as a PPM image.
    /// ```width``` and ```height``` are the dimensions of the buffer.
    /// ```color``` is the buffers ColorType.
    pub fn encode(&mut self, im: &[u8], width: u32, height: u32, color: color::ColorType) -> io::Result<()> {
        try!(self.write_magic_number());
        try!(self.write_metadata(width, height, color));

        self.write_image(im, color, width, height)
    }

    fn write_magic_number(&mut self) -> io::Result<()> {
        write!(self.w, "P6\n")
    }

    fn write_metadata(&mut self, width: u32, height: u32, pixel_type: color::ColorType) -> io::Result<()> {
        let m = max_pixel_value(pixel_type);
        write!(self.w, "{0} {1}\n{2}\n", width, height, m)
    }

    fn write_image(
        &mut self,
        buf: &[u8],
        pixel_type: color::ColorType,
        width: u32,
        height: u32) -> io::Result<()> {

        assert!(!buf.is_empty());
        match pixel_type {
            Gray(8) => {
                for i in 0..(width * height) as usize {
                    try!(self.w.write_all(&[buf[i]]));
                    try!(self.w.write_all(&[buf[i]]));
                    try!(self.w.write_all(&[buf[i]]));
                }
            }
            RGB(8) | RGB(16) => try!(self.w.write_all(buf)),
            RGBA(8) => {
                for x in buf.chunks(4) {
                    try!(self.w.write_all(&[x[0]]));
                    try!(self.w.write_all(&[x[1]]));
                    try!(self.w.write_all(&[x[2]]));
                }
            }
            a => panic!(format!("not implemented: {:?}", a))
        }

        Ok(())
    }
}

fn max_pixel_value(pixel_type: color::ColorType) -> u16 {
    let max = match pixel_type {
        Gray(n) | RGB(n) | Palette(n) | GrayA(n) | RGBA(n) => 2u32.pow(n as u32) - 1,
    };

    if max > 65535 {
      panic!("PPM: Trying to encode image with more than 16bit per pixel");
    } else {
      max as u16
    }
}
<commit_msg>Fix PPM encoder ownership comment<commit_after>//! Encoding of PPM Images

use std::io;
use std::io::Write;

use color;
use color::ColorType:: {
    Gray,
    Palette,
    GrayA,
    RGB,
    RGBA,
};

/// A representation of a PPM encoder.
pub struct PPMEncoder<'a, W: 'a> {
    w: &'a mut W
}

impl<'a, W: Write> PPMEncoder<'a, W> {
    /// Create a new PPMEncoder from the Writer ```w```.
    pub fn new(w: &mut W) -> PPMEncoder<W> {
        PPMEncoder { w: w }
    }

    /// Encode the buffer ```im``` as a PPM image.
    /// ```width``` and ```height``` are the dimensions of the buffer.
    /// ```color``` is the buffers ColorType.
    pub fn encode(&mut self, im: &[u8], width: u32, height: u32, color: color::ColorType) -> io::Result<()> {
        try!(self.write_magic_number());
        try!(self.write_metadata(width, height, color));

        self.write_image(im, color, width, height)
    }

    fn write_magic_number(&mut self) -> io::Result<()> {
        write!(self.w, "P6\n")
    }

    fn write_metadata(&mut self, width: u32, height: u32, pixel_type: color::ColorType) -> io::Result<()> {
        let m = max_pixel_value(pixel_type);
        write!(self.w, "{0} {1}\n{2}\n", width, height, m)
    }

    fn write_image(
        &mut self,
        buf: &[u8],
        pixel_type: color::ColorType,
        width: u32,
        height: u32) -> io::Result<()> {

        assert!(!buf.is_empty());
        match pixel_type {
            Gray(8) => {
                for i in 0..(width * height) as usize {
                    try!(self.w.write_all(&[buf[i]]));
                    try!(self.w.write_all(&[buf[i]]));
                    try!(self.w.write_all(&[buf[i]]));
                }
            }
            RGB(8) | RGB(16) => try!(self.w.write_all(buf)),
            RGBA(8) => {
                for x in buf.chunks(4) {
                    try!(self.w.write_all(&[x[0]]));
                    try!(self.w.write_all(&[x[1]]));
                    try!(self.w.write_all(&[x[2]]));
                }
            }
            a => panic!(format!("not implemented: {:?}", a))
        }

        Ok(())
    }
}

fn max_pixel_value(pixel_type: color::ColorType) -> u16 {
    let max = match pixel_type {
        Gray(n) | RGB(n) | Palette(n) | GrayA(n) | RGBA(n) => 2u32.pow(n as u32) - 1,
    };

    if max > 65535 {
      panic!("PPM: Trying to encode image with more than 16bit per pixel");
    } else {
      max as u16
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt::{self, Display, Formatter};
use std::boxed::FnBox;
use std::time::Instant;
use std::sync::{Arc, RwLock};

use mio::Token;
use futures::{Async, Future, Poll, Stream};
use futures::stream::{self, Once};
use grpc::{ChannelBuilder, Environment, WriteFlags};
use kvproto::raft_serverpb::SnapshotChunk;
use kvproto::raft_serverpb::RaftMessage;
use kvproto::tikvpb_grpc::TikvClient;

use raftstore::store::{SnapEntry, SnapKey, SnapManager, Snapshot};
use util::threadpool::{DefaultContext, ThreadPool, ThreadPoolBuilder};
use util::worker::Runnable;
use util::buf::PipeBuffer;
use util::security::SecurityManager;
use util::collections::{HashMap, HashMapEntry as Entry};
use util::HandyRwLock;

use super::metrics::*;
use super::{Error, Result};
use super::transport::RaftStoreRouter;

pub type Callback = Box<FnBox(Result<()>) + Send>;

const DEFAULT_SENDER_POOL_SIZE: usize = 3;

/// `Task` that `Runner` can handle.
///
/// `Register` register a pending snapshot file with token;
/// `Write` write data to snapshot file;
/// `Close` save the snapshot file;
/// `Discard` discard all the unsaved changes made to snapshot file;
/// `SendTo` send the snapshot file to specified address.
pub enum Task {
    Register(Token, RaftMessage),
    Write(Token, PipeBuffer),
    Close(Token),
    Discard(Token),
    SendTo {
        addr: String,
        msg: RaftMessage,
        cb: Callback,
    },
}

impl Display for Task {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        match *self {
            Task::Register(token, ref meta) => write!(f, "Register {:?} token: {:?}", meta, token),
            Task::Write(token, _) => write!(f, "Write snap for {:?}", token),
            Task::Close(token) => write!(f, "Close file {:?}", token),
            Task::Discard(token) => write!(f, "Discard file {:?}", token),
            Task::SendTo {
                ref addr, ref msg, ..
            } => write!(f, "SendTo Snap[to: {}, snap: {:?}]", addr, msg),
        }
    }
}

struct SnapChunk {
    snap: Arc<RwLock<Box<Snapshot>>>,
    remain_bytes: usize,
}

const SNAP_CHUNK_LEN: usize = 1024 * 1024;

impl Stream for SnapChunk {
    type Item = (SnapshotChunk, WriteFlags);
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Error> {
        let mut buf = match self.remain_bytes {
            0 => return Ok(Async::Ready(None)),
            n if n > SNAP_CHUNK_LEN => vec![0; SNAP_CHUNK_LEN],
            n => vec![0; n],
        };
        let result = self.snap.wl().read_exact(buf.as_mut_slice());
        match result {
            Ok(_) => {
                self.remain_bytes -= buf.len();
                let mut chunk = SnapshotChunk::new();
                chunk.set_data(buf);
                Ok(Async::Ready(Some((
                    chunk,
                    WriteFlags::default().buffer_hint(true),
                ))))
            }
            Err(e) => Err(box_err!("failed to read snapshot chunk: {}", e)),
        }
    }
}

/// Send the snapshot to specified address.
///
/// It will first send the normal raft snapshot message and then send the snapshot file.
fn send_snap(
    env: Arc<Environment>,
    mgr: SnapManager,
    security_mgr: Arc<SecurityManager>,
    addr: &str,
    msg: RaftMessage,
) -> Result<()> {
    assert!(msg.get_message().has_snapshot());
    let timer = Instant::now();

    let send_timer = SEND_SNAP_HISTOGRAM.start_coarse_timer();

    let key = {
        let snap = msg.get_message().get_snapshot();
        SnapKey::from_snap(snap)?
    };
    mgr.register(key.clone(), SnapEntry::Sending);
    defer!({
        mgr.deregister(&key, &SnapEntry::Sending);
    });
    let s = box_try!(mgr.get_snapshot_for_sending(&key));
    if !s.exists() {
        return Err(box_err!("missing snap file: {:?}", s.path()));
    }
    let total_size = s.total_size()?;

    // snapshot file has been validated when created, so no need to validate again.
    let s = Arc::new(RwLock::new(s));

    let chunks = {
        let snap_chunk = SnapChunk {
            snap: Arc::clone(&s),
            remain_bytes: total_size as usize,
        };
        let first: Once<(SnapshotChunk, _), Error> = stream::once({
            let mut chunk = SnapshotChunk::new();
            chunk.set_message(msg);
            Ok((chunk, WriteFlags::default().buffer_hint(true)))
        });
        first.chain(snap_chunk)
    };

    let cb = ChannelBuilder::new(env);
    let channel = security_mgr.connect(cb, addr);
    let client = TikvClient::new(channel);
    let (sink, receiver) = client.snapshot()?;
    let send = chunks.forward(sink);
    let res = send.and_then(|_| receiver.map_err(Error::from))
        .and_then(|_| {
            info!(
                "[region {}] sent snapshot {} [size: {}, dur: {:?}]",
                key.region_id,
                key,
                total_size,
                timer.elapsed()
            );
            s.wl().delete();
            Ok(())
        })
        .wait()
        .map_err(Error::from);

    send_timer.observe_duration();
    res
}

pub struct Runner<R: RaftStoreRouter + 'static> {
    env: Arc<Environment>,
    snap_mgr: SnapManager,
    files: HashMap<Token, (Box<Snapshot>, RaftMessage)>,
    pool: ThreadPool<DefaultContext>,
    raft_router: R,
    security_mgr: Arc<SecurityManager>,
}

impl<R: RaftStoreRouter + 'static> Runner<R> {
    pub fn new(
        env: Arc<Environment>,
        snap_mgr: SnapManager,
        r: R,
        security_mgr: Arc<SecurityManager>,
    ) -> Runner<R> {
        Runner {
            env: env,
            snap_mgr: snap_mgr,
            files: map![],
            pool: ThreadPoolBuilder::with_default_factory(thd_name!("snap sender"))
                .thread_count(DEFAULT_SENDER_POOL_SIZE)
                .build(),
            raft_router: r,
            security_mgr: security_mgr,
        }
    }
}

impl<R: RaftStoreRouter + 'static> Runnable<Task> for Runner<R> {
    fn run(&mut self, task: Task) {
        match task {
            Task::Register(token, meta) => {
                SNAP_TASK_COUNTER.with_label_values(&["register"]).inc();
                let mgr = self.snap_mgr.clone();
                let key = match SnapKey::from_snap(meta.get_message().get_snapshot()) {
                    Ok(k) => k,
                    Err(e) => {
                        error!("failed to create snap key for token {:?}: {:?}", token, e);
                        return;
                    }
                };
                match mgr.get_snapshot_for_receiving(
                    &key,
                    meta.get_message().get_snapshot().get_data(),
                ) {
                    Ok(snap) => {
                        if snap.exists() {
                            info!(
                                "snapshot file {} already exists, skip receiving.",
                                snap.path()
                            );
                            if let Err(e) = self.raft_router.send_raft_msg(meta) {
                                error!("send snapshot for key {} token {:?}: {:?}", key, token, e);
                            }
                            return;
                        }
                        debug!("begin to receive snap {:?}", meta);
                        mgr.register(key, SnapEntry::Receiving);
                        self.files.insert(token, (snap, meta));
                    }
                    Err(e) => {
                        error!(
                            "failed to create snapshot file for token {:?}: {:?}",
                            token, e
                        );
                        return;
                    }
                }
            }
            Task::Write(token, mut data) => {
                SNAP_TASK_COUNTER.with_label_values(&["write"]).inc();
                match self.files.entry(token) {
                    Entry::Occupied(mut e) => {
                        if let Err(err) = data.write_all_to(&mut e.get_mut().0) {
                            error!(
                                "failed to write data to snapshot file {} for token {:?}: {:?}",
                                e.get_mut().0.path(),
                                token,
                                err
                            );
                            let (_, msg) = e.remove();
                            let key = SnapKey::from_snap(msg.get_message().get_snapshot()).unwrap();
                            self.snap_mgr.deregister(&key, &SnapEntry::Receiving);
                        }
                    }
                    Entry::Vacant(_) => error!("invalid snap token {:?}", token),
                }
            }
            Task::Close(token) => {
                SNAP_TASK_COUNTER.with_label_values(&["close"]).inc();
                match self.files.remove(&token) {
                    Some((mut snap, msg)) => {
                        let key = SnapKey::from_snap(msg.get_message().get_snapshot()).unwrap();
                        info!("saving snapshot to {}", snap.path());
                        defer!({
                            self.snap_mgr.deregister(&key, &SnapEntry::Receiving);
                        });
                        if let Err(e) = snap.save() {
                            error!(
                                "failed to save snapshot file {} for token {:?}: {:?}",
                                snap.path(),
                                token,
                                e
                            );
                            return;
                        }
                        if let Err(e) = self.raft_router.send_raft_msg(msg) {
                            error!("send snapshot for token {:?} err {:?}", token, e);
                        }
                    }
                    None => error!("invalid snap token {:?}", token),
                }
            }
            Task::Discard(token) => {
                SNAP_TASK_COUNTER.with_label_values(&["discard"]).inc();
                if let Some((_, msg)) = self.files.remove(&token) {
                    debug!("discard snapshot: {:?}", msg);
                    // because token is inserted, following can't panic.
                    let key = SnapKey::from_snap(msg.get_message().get_snapshot()).unwrap();
                    self.snap_mgr.deregister(&key, &SnapEntry::Receiving);
                }
            }
            Task::SendTo { addr, msg, cb } => {
                SNAP_TASK_COUNTER.with_label_values(&["send"]).inc();
                let env = Arc::clone(&self.env);
                let mgr = self.snap_mgr.clone();
                let security_mgr = Arc::clone(&self.security_mgr);
                self.pool.execute(move |_| {
                    let res = send_snap(env, mgr, security_mgr, &addr, msg);
                    if res.is_err() {
                        error!("failed to send snap to {}: {:?}", addr, res);
                    }
                    cb(res)
                });
            }
        }
    }
}
<commit_msg>make send snapshots async (#2850)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt::{self, Display, Formatter};
use std::boxed::FnBox;
use std::time::{Duration, Instant};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

use mio::Token;
use futures::{future, Async, Future, Poll, Stream};
use futures_cpupool::{Builder as CpuPoolBuilder, CpuPool};
use grpc::{ChannelBuilder, Environment, WriteFlags};
use kvproto::raft_serverpb::SnapshotChunk;
use kvproto::raft_serverpb::RaftMessage;
use kvproto::tikvpb_grpc::TikvClient;

use raftstore::store::{SnapEntry, SnapKey, SnapManager, Snapshot};
use util::DeferContext;
use util::worker::Runnable;
use util::buf::PipeBuffer;
use util::security::SecurityManager;
use util::collections::{HashMap, HashMapEntry as Entry};

use super::metrics::*;
use super::{Error, Result};
use super::transport::RaftStoreRouter;

pub type Callback = Box<FnBox(Result<()>) + Send>;

const DEFAULT_SENDER_POOL_SIZE: usize = 3;
// How many snapshots can be sent concurrently.
const MAX_SENDER_CONCURRENT: usize = 64;

/// `Task` that `Runner` can handle.
///
/// `Register` register a pending snapshot file with token;
/// `Write` write data to snapshot file;
/// `Close` save the snapshot file;
/// `Discard` discard all the unsaved changes made to snapshot file;
/// `SendTo` send the snapshot file to specified address.
pub enum Task {
    Register(Token, RaftMessage),
    Write(Token, PipeBuffer),
    Close(Token),
    Discard(Token),
    SendTo {
        addr: String,
        msg: RaftMessage,
        cb: Callback,
    },
}

impl Display for Task {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        match *self {
            Task::Register(token, ref meta) => write!(f, "Register {:?} token: {:?}", meta, token),
            Task::Write(token, _) => write!(f, "Write snap for {:?}", token),
            Task::Close(token) => write!(f, "Close file {:?}", token),
            Task::Discard(token) => write!(f, "Discard file {:?}", token),
            Task::SendTo {
                ref addr, ref msg, ..
            } => write!(f, "SendTo Snap[to: {}, snap: {:?}]", addr, msg),
        }
    }
}

struct SnapChunk {
    first: Option<SnapshotChunk>,
    snap: Box<Snapshot>,
    remain_bytes: usize,
}

const SNAP_CHUNK_LEN: usize = 1024 * 1024;

impl Stream for SnapChunk {
    type Item = (SnapshotChunk, WriteFlags);
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Error> {
        if let Some(t) = self.first.take() {
            let write_flags = WriteFlags::default().buffer_hint(true);
            return Ok(Async::Ready(Some((t, write_flags))));
        }

        let mut buf = match self.remain_bytes {
            0 => return Ok(Async::Ready(None)),
            n if n > SNAP_CHUNK_LEN => vec![0; SNAP_CHUNK_LEN],
            n => vec![0; n],
        };
        let result = self.snap.read_exact(buf.as_mut_slice());
        match result {
            Ok(_) => {
                self.remain_bytes -= buf.len();
                let mut chunk = SnapshotChunk::new();
                chunk.set_data(buf);
                Ok(Async::Ready(Some((
                    chunk,
                    WriteFlags::default().buffer_hint(true),
                ))))
            }
            Err(e) => Err(box_err!("failed to read snapshot chunk: {}", e)),
        }
    }
}

struct SendStat {
    key: SnapKey,
    total_size: u64,
    elapsed: Duration,
}

/// Send the snapshot to specified address.
///
/// It will first send the normal raft snapshot message and then send the snapshot file.
fn send_snap(
    env: Arc<Environment>,
    mgr: SnapManager,
    security_mgr: Arc<SecurityManager>,
    addr: &str,
    msg: RaftMessage,
) -> Result<impl Future<Item = SendStat, Error = Error>> {
    assert!(msg.get_message().has_snapshot());
    let timer = Instant::now();

    let send_timer = SEND_SNAP_HISTOGRAM.start_coarse_timer();

    let key = {
        let snap = msg.get_message().get_snapshot();
        SnapKey::from_snap(snap)?
    };

    mgr.register(key.clone(), SnapEntry::Sending);
    let deregister = {
        let (mgr, key) = (mgr.clone(), key.clone());
        DeferContext::new(move || mgr.deregister(&key, &SnapEntry::Sending))
    };

    let s = box_try!(mgr.get_snapshot_for_sending(&key));
    if !s.exists() {
        return Err(box_err!("missing snap file: {:?}", s.path()));
    }
    let total_size = s.total_size()?;

    let chunks = {
        let mut first_chunk = SnapshotChunk::new();
        first_chunk.set_message(msg);

        SnapChunk {
            first: Some(first_chunk),
            snap: s,
            remain_bytes: total_size as usize,
        }
    };

    let cb = ChannelBuilder::new(env);
    let channel = security_mgr.connect(cb, addr);
    let client = TikvClient::new(channel);
    let (sink, receiver) = client.snapshot()?;

    let send = chunks.forward(sink).map_err(Error::from);
    let send = send.and_then(|(s, _)| receiver.map_err(Error::from).map(|_| s))
        .then(move |result| {
            send_timer.observe_duration();
            drop(deregister);
            drop(client);
            result.map(|s| {
                s.snap.delete();
                // TODO: improve it after rustc resolves the bug.
                // Call `info` in the closure directly will cause rustc
                // panic with `Cannot create local mono-item for DefId`.
                SendStat {
                    key: key,
                    total_size: total_size,
                    elapsed: timer.elapsed(),
                }
            })
        });
    Ok(send)
}

pub struct Runner<R: RaftStoreRouter + 'static> {
    env: Arc<Environment>,
    snap_mgr: SnapManager,
    files: HashMap<Token, (Box<Snapshot>, RaftMessage)>,
    pool: CpuPool,
    raft_router: R,
    security_mgr: Arc<SecurityManager>,
    sending_count: Arc<AtomicUsize>,
}

impl<R: RaftStoreRouter + 'static> Runner<R> {
    pub fn new(
        env: Arc<Environment>,
        snap_mgr: SnapManager,
        r: R,
        security_mgr: Arc<SecurityManager>,
    ) -> Runner<R> {
        Runner {
            env: env,
            snap_mgr: snap_mgr,
            files: map![],
            pool: CpuPoolBuilder::new()
                .name_prefix(thd_name!("snap sender"))
                .pool_size(DEFAULT_SENDER_POOL_SIZE)
                .create(),
            raft_router: r,
            security_mgr: security_mgr,
            sending_count: Arc::new(AtomicUsize::new(0)),
        }
    }
}

impl<R: RaftStoreRouter + 'static> Runnable<Task> for Runner<R> {
    fn run(&mut self, task: Task) {
        match task {
            Task::Register(token, meta) => {
                SNAP_TASK_COUNTER.with_label_values(&["register"]).inc();
                let mgr = self.snap_mgr.clone();
                let key = match SnapKey::from_snap(meta.get_message().get_snapshot()) {
                    Ok(k) => k,
                    Err(e) => {
                        error!("failed to create snap key for token {:?}: {:?}", token, e);
                        return;
                    }
                };
                match mgr.get_snapshot_for_receiving(
                    &key,
                    meta.get_message().get_snapshot().get_data(),
                ) {
                    Ok(snap) => {
                        if snap.exists() {
                            info!(
                                "snapshot file {} already exists, skip receiving.",
                                snap.path()
                            );
                            if let Err(e) = self.raft_router.send_raft_msg(meta) {
                                error!("send snapshot for key {} token {:?}: {:?}", key, token, e);
                            }
                            return;
                        }
                        debug!("begin to receive snap {:?}", meta);
                        mgr.register(key, SnapEntry::Receiving);
                        self.files.insert(token, (snap, meta));
                    }
                    Err(e) => {
                        error!(
                            "failed to create snapshot file for token {:?}: {:?}",
                            token, e
                        );
                        return;
                    }
                }
            }
            Task::Write(token, mut data) => {
                SNAP_TASK_COUNTER.with_label_values(&["write"]).inc();
                match self.files.entry(token) {
                    Entry::Occupied(mut e) => {
                        if let Err(err) = data.write_all_to(&mut e.get_mut().0) {
                            error!(
                                "failed to write data to snapshot file {} for token {:?}: {:?}",
                                e.get_mut().0.path(),
                                token,
                                err
                            );
                            let (_, msg) = e.remove();
                            let key = SnapKey::from_snap(msg.get_message().get_snapshot()).unwrap();
                            self.snap_mgr.deregister(&key, &SnapEntry::Receiving);
                        }
                    }
                    Entry::Vacant(_) => error!("invalid snap token {:?}", token),
                }
            }
            Task::Close(token) => {
                SNAP_TASK_COUNTER.with_label_values(&["close"]).inc();
                match self.files.remove(&token) {
                    Some((mut snap, msg)) => {
                        let key = SnapKey::from_snap(msg.get_message().get_snapshot()).unwrap();
                        info!("saving snapshot to {}", snap.path());
                        defer!({
                            self.snap_mgr.deregister(&key, &SnapEntry::Receiving);
                        });
                        if let Err(e) = snap.save() {
                            error!(
                                "failed to save snapshot file {} for token {:?}: {:?}",
                                snap.path(),
                                token,
                                e
                            );
                            return;
                        }
                        if let Err(e) = self.raft_router.send_raft_msg(msg) {
                            error!("send snapshot for token {:?} err {:?}", token, e);
                        }
                    }
                    None => error!("invalid snap token {:?}", token),
                }
            }
            Task::Discard(token) => {
                SNAP_TASK_COUNTER.with_label_values(&["discard"]).inc();
                if let Some((_, msg)) = self.files.remove(&token) {
                    debug!("discard snapshot: {:?}", msg);
                    // because token is inserted, following can't panic.
                    let key = SnapKey::from_snap(msg.get_message().get_snapshot()).unwrap();
                    self.snap_mgr.deregister(&key, &SnapEntry::Receiving);
                }
            }
            Task::SendTo { addr, msg, cb } => {
                if self.sending_count.load(Ordering::SeqCst) >= MAX_SENDER_CONCURRENT {
                    warn!(
                        "too many sending snapshot tasks, drop SendTo Snap[to: {}, snap: {:?}]",
                        addr, msg
                    );
                    cb(Err(Error::Other("Too many sending snapshot tasks".into())));
                    return;
                }
                SNAP_TASK_COUNTER.with_label_values(&["send"]).inc();

                let env = Arc::clone(&self.env);
                let mgr = self.snap_mgr.clone();
                let security_mgr = Arc::clone(&self.security_mgr);
                let sending_count = Arc::clone(&self.sending_count);
                sending_count.fetch_add(1, Ordering::SeqCst);

                let f = future::result(send_snap(env, mgr, security_mgr, &addr, msg))
                    .flatten()
                    .then(move |res| {
                        match res {
                            Ok(stat) => {
                                info!(
                                    "[region {}] sent snapshot {} [size: {}, dur: {:?}]",
                                    stat.key.region_id, stat.key, stat.total_size, stat.elapsed,
                                );
                                cb(Ok(()));
                            }
                            Err(e) => {
                                error!("failed to send snap to {}: {:?}", addr, e);
                                cb(Err(e));
                            }
                        };
                        sending_count.fetch_sub(1, Ordering::SeqCst);
                        future::ok::<_, ()>(())
                    });

                self.pool.spawn(f).forget();
            }
        }
    }
}
<|endoftext|>


<commit_before>#[cfg(feature = "trace_log")]
extern crate env_logger;

extern crate radeco_lib;
extern crate r2pipe;
extern crate r2api;
extern crate base64;
extern crate rustyline;

use r2pipe::{R2, R2Pipe};
use radeco_lib::analysis::cse::cse::CSE;
use radeco_lib::analysis::interproc::fixcall::CallFixer;
use radeco_lib::analysis::sccp;
use radeco_lib::backend::lang_c::c_cfg::CCFGVerifier;
use radeco_lib::backend::lang_c::c_cfg_builder;
use radeco_lib::frontend::radeco_containers::*;
use radeco_lib::middle::{dce, dot};
use radeco_lib::middle::ir_writer;
use radeco_lib::middle::regfile::SubRegisterFile;
use radeco_lib::middle::ssa::ssastorage::SSAStorage;
use radeco_lib::middle::ssa::verifier;
use rustyline::Editor;

use rustyline::error::ReadlineError;
use std::cell::RefCell;

use std::collections::HashMap;
use std::process;
use std::rc::Rc;
use std::str;
use std::sync::Arc;

// On unix platforms you can use ANSI escape sequences
#[cfg(unix)]
const PROMPT: &'static str = "\x1b[1;32m>>\x1b[0m ";

// Windows consoles typically don't support ANSI escape sequences out
// of the box
#[cfg(windows)]
const PROMPT: &'static str = ">> ";

fn main() {
    let mut rl = Editor::<()>::new();
    let mut proj = None;
    loop {
        let readline = rl.readline(PROMPT);
        match readline {
            Ok(line) => {
                let mut terms = line.split_whitespace();
                let o1 = terms.next();
                let o2 = terms.next();
                cmd(o1, o2, &mut proj);
            }
            Err(ReadlineError::Interrupted) |
            Err(ReadlineError::Eof) => break,
            Err(err) => {
                println!("Error: {:?}", err);
                break;
            }
        }
    }
}

mod command {
    pub const LOAD: &'static str = "load";
    pub const CONNECT: &'static str = "connect";
    pub const FNLIST: &'static str = "fn_list";
    pub const ANALYZE: &'static str = "analyze";
    pub const DOT: &'static str = "dot";
    pub const IR: &'static str = "ir";
    pub const DECOMPILE: &'static str = "decompile";
}

fn cmd(op1: Option<&str>, op2: Option<&str>, proj_opt: &mut Option<RadecoProject>) {
    match (op1, op2) {
        (Some(command::LOAD), Some(path)) => {
            *proj_opt = Some(load_proj_by_path(path));
            return;
        }
        (Some(command::CONNECT), Some(port)) => {
            match load_proj_by_tcp(port) {
                Ok(proj) => *proj_opt = Some(proj),
                Err(msg) => println!("{}", msg),
            }
            return;
        }
        _ => {}
    };
    let mut proj = match proj_opt {
        Some(proj) => proj,
        None => {
            println!("Load a project first");
            return;
        }
    };
    match (op1, op2) {
        (Some(command::ANALYZE), Some("*")) => {
            let rfns = proj.iter_mut().map(|i| i.module).flat_map(|rmod| {
                rmod.functions.values_mut()
            });
            for rfn in rfns {
                analyze(rfn);
            }
        }
        (Some(command::FNLIST), _) => {
            let funcs = fn_list(&proj);
            println!("{}", funcs.join("\n"));
        }
        // TODO Show list of dependency information of analyses
        // TODO Add command for individual analyses
        (Some(command::ANALYZE), Some(f)) => {
            if let Some(rfn) = get_function_mut(f, &mut proj) {
                analyze(rfn);
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::DOT), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                println!("{}", emit_dot(rfn.ssa()));
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::IR), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                println!("{}", emit_ir(rfn));
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::DECOMPILE), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                let rmod = proj.iter().map(|i| i.module).next().unwrap();
                let func_name_map = rmod.functions
                    .iter()
                    .map(|(&addr, f)| (addr, f.name.to_string()))
                    .collect();
                let strings = rmod.strings()
                    .iter()
                    .filter(|ref s| s.vaddr.is_some() && s.string.is_some())
                    .cloned()
                    .map(|s| {
                        let (addr, _s) = (s.vaddr.unwrap(), s.string.unwrap());
                        let bytes = base64::decode(&_s).unwrap_or(Vec::new());
                        let ret_string = match str::from_utf8(bytes.as_slice()) {
                            Ok(v) => v.to_string(),
                            Err(_e) => _s,
                        };
                        (addr, ret_string)
                    })
                    .collect();
                println!("{}", decompile(rfn, &func_name_map, &strings));
            } else {
                println!("{} is not found", f);
            }
        }
        _ => {}
    }
}

fn load_proj_by_path(path: &str) -> RadecoProject {
    let mut p = ProjectLoader::new().path(path).load();
    let regfile = p.regfile().clone();
    for mut xy in p.iter_mut() {
        analyze_mod(regfile.clone(), xy.module);
    }
    p
}

fn load_proj_by_tcp(port: &str) -> Result<RadecoProject, &'static str> {
    let r2p = R2Pipe::tcp(&format!("localhost:{}", port))?;
    let r2 = R2::from(r2p);
    let r2w = Rc::new(RefCell::new(r2));
    let mut p = ProjectLoader::new().source(Rc::new(r2w)).load();
    let regfile = p.regfile().clone();
    for mut xy in p.iter_mut() {
        analyze_mod(regfile.clone(), xy.module);
    }
    Ok(p)
}

fn fn_list(proj: &RadecoProject) -> Vec<String> {
    proj.iter()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values())
        .map(|s| s.name.to_string())
        .collect()
}

fn get_function<'a>(name: &str, proj: &'a RadecoProject) -> Option<&'a RadecoFunction> {
    proj.iter()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values())
        .filter(|rfn| rfn.name == name)
        .next()
}

fn get_function_mut<'a>(name: &str, proj: &'a mut RadecoProject) -> Option<&'a mut RadecoFunction> {
    proj.iter_mut()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values_mut())
        .filter(|rfn| rfn.name == name)
        .next()
}

fn analyze_mod(regfile: Arc<SubRegisterFile>, rmod: &mut RadecoModule) {
    // Analyze preserved for all functions.
    {
        println!("[*] Fixing Callee Information");
        let bp_name = regfile.get_name_by_alias(&"BP".to_string());
        let bp_name = bp_name.map(|s| s.to_owned());
        let sp_name = regfile.get_name_by_alias(&"SP".to_string());
        let sp_name = sp_name.map(|s| s.to_owned());
        let mut callfixer = CallFixer::new(rmod, bp_name, sp_name);
        callfixer.rounded_analysis();
    }

    // Fix call sites
    radeco_lib::analysis::functions::fix_ssa_opcalls::go(rmod);

    // Infer calling conventions
    radeco_lib::analysis::functions::infer_regusage::run(rmod, &*regfile);
}

fn analyze(rfn: &mut RadecoFunction) {
    println!("[+] Analyzing: {} @ {:#x}", rfn.name, rfn.offset);
    {
        println!("  [*] Eliminating Dead Code");
        dce::collect(rfn.ssa_mut());
    }
    let mut ssa = {
        // Constant Propagation (sccp)
        println!("  [*] Propagating Constants");
        let mut analyzer = sccp::Analyzer::new(rfn.ssa_mut());
        analyzer.analyze();
        analyzer.emit_ssa()
    };
    {
        println!("  [*] Eliminating More DeadCode");
        dce::collect(&mut ssa);
    }
    *rfn.ssa_mut() = ssa;
    {
        // Common SubExpression Elimination (cse)
        println!("  [*] Eliminating Common SubExpressions");
        let mut cse = CSE::new(rfn.ssa_mut());
        cse.run();
    }
    {
        // Verify SSA
        println!("  [*] Verifying SSA's Validity");
        match verifier::verify(rfn.ssa()) {
            Err(e) => {
                println!("  [*] Found Error: {}", e);
                process::exit(255);
            }
            Ok(_) => {}
        }
    }
}

fn emit_ir(rfn: &RadecoFunction) -> String {
    println!("  [*] Writing out IR");
    let mut res = String::new();
    ir_writer::emit_il(&mut res, Some(rfn.name.to_string()), rfn.ssa()).unwrap();
    res
}

fn emit_dot(ssa: &SSAStorage) -> String {
    dot::emit_dot(ssa)
}

fn decompile(
    rfn: &RadecoFunction,
    func_name_map: &HashMap<u64, String>,
    strings: &HashMap<u64, String>,
) -> String {
    let c_cfg = c_cfg_builder::recover_c_cfg(rfn, func_name_map, strings);
    if let Err(err) = CCFGVerifier::verify(&c_cfg) {
        println!("CCFG verification failed");
        println!("{}", err);
    }
    c_cfg.to_c_ast().print()
}
<commit_msg>Apply ctrl_flow_struct<commit_after>#[cfg(feature = "trace_log")]
extern crate env_logger;

extern crate radeco_lib;
extern crate r2pipe;
extern crate r2api;
extern crate base64;
extern crate rustyline;

use r2pipe::{R2, R2Pipe};
use radeco_lib::analysis::cse::cse::CSE;
use radeco_lib::analysis::interproc::fixcall::CallFixer;
use radeco_lib::analysis::sccp;
use radeco_lib::backend::lang_c::c_cfg::CCFGVerifier;
use radeco_lib::backend::lang_c::c_cfg::ctrl_flow_struct;
use radeco_lib::backend::lang_c::c_cfg_builder;
use radeco_lib::frontend::radeco_containers::*;
use radeco_lib::middle::{dce, dot};
use radeco_lib::middle::ir_writer;
use radeco_lib::middle::regfile::SubRegisterFile;
use radeco_lib::middle::ssa::ssastorage::SSAStorage;
use radeco_lib::middle::ssa::verifier;
use rustyline::Editor;

use rustyline::error::ReadlineError;
use std::cell::RefCell;

use std::collections::HashMap;
use std::process;
use std::rc::Rc;
use std::str;
use std::sync::Arc;

// On unix platforms you can use ANSI escape sequences
#[cfg(unix)]
const PROMPT: &'static str = "\x1b[1;32m>>\x1b[0m ";

// Windows consoles typically don't support ANSI escape sequences out
// of the box
#[cfg(windows)]
const PROMPT: &'static str = ">> ";

fn main() {
    let mut rl = Editor::<()>::new();
    let mut proj = None;
    loop {
        let readline = rl.readline(PROMPT);
        match readline {
            Ok(line) => {
                let mut terms = line.split_whitespace();
                let o1 = terms.next();
                let o2 = terms.next();
                cmd(o1, o2, &mut proj);
            }
            Err(ReadlineError::Interrupted) |
            Err(ReadlineError::Eof) => break,
            Err(err) => {
                println!("Error: {:?}", err);
                break;
            }
        }
    }
}

mod command {
    pub const LOAD: &'static str = "load";
    pub const CONNECT: &'static str = "connect";
    pub const FNLIST: &'static str = "fn_list";
    pub const ANALYZE: &'static str = "analyze";
    pub const DOT: &'static str = "dot";
    pub const IR: &'static str = "ir";
    pub const DECOMPILE: &'static str = "decompile";
}

fn cmd(op1: Option<&str>, op2: Option<&str>, proj_opt: &mut Option<RadecoProject>) {
    match (op1, op2) {
        (Some(command::LOAD), Some(path)) => {
            *proj_opt = Some(load_proj_by_path(path));
            return;
        }
        (Some(command::CONNECT), Some(port)) => {
            match load_proj_by_tcp(port) {
                Ok(proj) => *proj_opt = Some(proj),
                Err(msg) => println!("{}", msg),
            }
            return;
        }
        _ => {}
    };
    let mut proj = match proj_opt {
        Some(proj) => proj,
        None => {
            println!("Load a project first");
            return;
        }
    };
    match (op1, op2) {
        (Some(command::ANALYZE), Some("*")) => {
            let rfns = proj.iter_mut().map(|i| i.module).flat_map(|rmod| {
                rmod.functions.values_mut()
            });
            for rfn in rfns {
                analyze(rfn);
            }
        }
        (Some(command::FNLIST), _) => {
            let funcs = fn_list(&proj);
            println!("{}", funcs.join("\n"));
        }
        // TODO Show list of dependency information of analyses
        // TODO Add command for individual analyses
        (Some(command::ANALYZE), Some(f)) => {
            if let Some(rfn) = get_function_mut(f, &mut proj) {
                analyze(rfn);
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::DOT), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                println!("{}", emit_dot(rfn.ssa()));
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::IR), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                println!("{}", emit_ir(rfn));
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::DECOMPILE), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                let rmod = proj.iter().map(|i| i.module).next().unwrap();
                let func_name_map = rmod.functions
                    .iter()
                    .map(|(&addr, f)| (addr, f.name.to_string()))
                    .collect();
                let strings = rmod.strings()
                    .iter()
                    .filter(|ref s| s.vaddr.is_some() && s.string.is_some())
                    .cloned()
                    .map(|s| {
                        let (addr, _s) = (s.vaddr.unwrap(), s.string.unwrap());
                        let bytes = base64::decode(&_s).unwrap_or(Vec::new());
                        let ret_string = match str::from_utf8(bytes.as_slice()) {
                            Ok(v) => v.to_string(),
                            Err(_e) => _s,
                        };
                        (addr, ret_string)
                    })
                    .collect();
                println!("{}", decompile(rfn, &func_name_map, &strings));
            } else {
                println!("{} is not found", f);
            }
        }
        _ => {}
    }
}

fn load_proj_by_path(path: &str) -> RadecoProject {
    let mut p = ProjectLoader::new().path(path).load();
    let regfile = p.regfile().clone();
    for mut xy in p.iter_mut() {
        analyze_mod(regfile.clone(), xy.module);
    }
    p
}

fn load_proj_by_tcp(port: &str) -> Result<RadecoProject, &'static str> {
    let r2p = R2Pipe::tcp(&format!("localhost:{}", port))?;
    let r2 = R2::from(r2p);
    let r2w = Rc::new(RefCell::new(r2));
    let mut p = ProjectLoader::new().source(Rc::new(r2w)).load();
    let regfile = p.regfile().clone();
    for mut xy in p.iter_mut() {
        analyze_mod(regfile.clone(), xy.module);
    }
    Ok(p)
}

fn fn_list(proj: &RadecoProject) -> Vec<String> {
    proj.iter()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values())
        .map(|s| s.name.to_string())
        .collect()
}

fn get_function<'a>(name: &str, proj: &'a RadecoProject) -> Option<&'a RadecoFunction> {
    proj.iter()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values())
        .filter(|rfn| rfn.name == name)
        .next()
}

fn get_function_mut<'a>(name: &str, proj: &'a mut RadecoProject) -> Option<&'a mut RadecoFunction> {
    proj.iter_mut()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values_mut())
        .filter(|rfn| rfn.name == name)
        .next()
}

fn analyze_mod(regfile: Arc<SubRegisterFile>, rmod: &mut RadecoModule) {
    // Analyze preserved for all functions.
    {
        println!("[*] Fixing Callee Information");
        let bp_name = regfile.get_name_by_alias(&"BP".to_string());
        let bp_name = bp_name.map(|s| s.to_owned());
        let sp_name = regfile.get_name_by_alias(&"SP".to_string());
        let sp_name = sp_name.map(|s| s.to_owned());
        let mut callfixer = CallFixer::new(rmod, bp_name, sp_name);
        callfixer.rounded_analysis();
    }

    // Fix call sites
    radeco_lib::analysis::functions::fix_ssa_opcalls::go(rmod);

    // Infer calling conventions
    radeco_lib::analysis::functions::infer_regusage::run(rmod, &*regfile);
}

fn analyze(rfn: &mut RadecoFunction) {
    println!("[+] Analyzing: {} @ {:#x}", rfn.name, rfn.offset);
    {
        println!("  [*] Eliminating Dead Code");
        dce::collect(rfn.ssa_mut());
    }
    let mut ssa = {
        // Constant Propagation (sccp)
        println!("  [*] Propagating Constants");
        let mut analyzer = sccp::Analyzer::new(rfn.ssa_mut());
        analyzer.analyze();
        analyzer.emit_ssa()
    };
    {
        println!("  [*] Eliminating More DeadCode");
        dce::collect(&mut ssa);
    }
    *rfn.ssa_mut() = ssa;
    {
        // Common SubExpression Elimination (cse)
        println!("  [*] Eliminating Common SubExpressions");
        let mut cse = CSE::new(rfn.ssa_mut());
        cse.run();
    }
    {
        // Verify SSA
        println!("  [*] Verifying SSA's Validity");
        match verifier::verify(rfn.ssa()) {
            Err(e) => {
                println!("  [*] Found Error: {}", e);
                process::exit(255);
            }
            Ok(_) => {}
        }
    }
}

fn emit_ir(rfn: &RadecoFunction) -> String {
    println!("  [*] Writing out IR");
    let mut res = String::new();
    ir_writer::emit_il(&mut res, Some(rfn.name.to_string()), rfn.ssa()).unwrap();
    res
}

fn emit_dot(ssa: &SSAStorage) -> String {
    dot::emit_dot(ssa)
}

fn decompile(
    rfn: &RadecoFunction,
    func_name_map: &HashMap<u64, String>,
    strings: &HashMap<u64, String>,
) -> String {
    let c_cfg = c_cfg_builder::recover_c_cfg(rfn, func_name_map, strings);
    if let Err(err) = CCFGVerifier::verify(&c_cfg) {
        println!("CCFG verification failed");
        println!("{}", err);
    }
    ctrl_flow_struct::structure_and_convert(c_cfg)
        .unwrap()
        .print()
}
<|endoftext|>


<commit_before>// Copyright 2014-2016 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Handling log-probabilities.

use std::mem;
use std::f64;
use std::iter;
use std::ops::{Add, Sub, Mul, Div};
use num::traits::cast;
use num::NumCast;

use itertools::linspace;
use itertools::Itertools;
use itertools::misc::ToFloat;

/// A factor to convert log-probabilities to PHRED-scale (phred = p * `LOG_TO_PHRED_FACTOR`).
const LOG_TO_PHRED_FACTOR: f64 = -4.3429448190325175; // -10 * 1 / ln(10)

/// A factor to convert PHRED-scale to log-probabilities (p = phred * `PHRED_TO_LOG_FACTOR`).
const PHRED_TO_LOG_FACTOR: f64 = -0.23025850929940456; // 1 / (-10 * log10(e))

/// Calculate log(1 - p) with p given in log space without loss of precision as described in
/// http://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf.
fn ln_1m_exp(p: f64) -> f64 {
    assert!(p <= 0.0);
    if p < -0.693 {
        (-p.exp()).ln_1p()
    } else {
        (-p.exp_m1()).ln()
    }
}

/// A newtype for probabilities.
custom_derive! {
    #[derive(
        NewtypeFrom,
        NewtypeDeref,
        NewtypeAdd(*),
        NewtypeSub(*),
        NewtypeMul(*),
        NewtypeDiv(*),
        PartialEq,
        PartialOrd,
        Copy,
        Clone,
        Debug
    )]
    pub struct Prob(f64);
}

/// A newtype for log-scale probabilities.
custom_derive! {
    #[derive(
        NewtypeFrom,
        NewtypeDeref,
        NewtypeAdd(*),
        NewtypeSub(*),
        PartialEq,
        PartialOrd,
        Copy,
        Clone,
        Debug
    )]
    pub struct LogProb(f64);
}

/// A newtype for PHRED-scale probabilities.
custom_derive! {
    #[derive(
        NewtypeFrom,
        NewtypeDeref,
        NewtypeAdd(*),
        NewtypeSub(*),
        PartialEq,
        PartialOrd,
        Copy,
        Clone,
        Debug
    )]
    pub struct PHREDProb(f64);
}

/// Iterator returned by scans over logprobs.
pub type ScanIter<I> = iter::Scan<<I as IntoIterator>::IntoIter, LogProb, fn(&mut LogProb, LogProb) -> Option<LogProb>>;

impl LogProb {
    pub fn ln_zero() -> LogProb {
        LogProb(f64::NEG_INFINITY)
    }

    pub fn ln_one() -> LogProb {
        LogProb(0.0)
    }

    pub fn ln_one_minus_exp(&self) -> LogProb {
        LogProb(ln_1m_exp(**self))
    }

    pub fn ln_sum_exp(probs: &[LogProb]) -> LogProb {
        if probs.is_empty() {
            Self::ln_zero()
        } else {
            let mut pmax = probs[0];
            let mut imax = 0;
            for (i, &p) in probs.iter().enumerate().skip(1) {
                if p > pmax {
                    pmax = p;
                    imax = i;
                }
            }
            if pmax == Self::ln_zero() {
                Self::ln_zero()
            } else if *pmax == f64::INFINITY {
                LogProb(f64::INFINITY)
            } else {
                // TODO use sum() once it has been stabilized: .sum::<usize>()
                pmax + LogProb(
                    (probs.iter()
                          .enumerate()
                          .filter_map(|(i, p)| {
                              if i == imax {
                                  None
                              } else {
                                  Some((p - pmax).exp())
                              }
                          })
                          .fold(0.0, |s, e| s + e)
                     ).ln_1p()
                 )
            }
        }
    }

    pub fn ln_add_exp(self, other: LogProb) -> LogProb {
        let (mut p0, mut p1) = (self, other);
        if p1 > p0 {
            mem::swap(&mut p0, &mut p1);
        }
        if p0 == Self::ln_zero() {
            Self::ln_zero()
        } else if *p0 == f64::INFINITY {
            LogProb(f64::INFINITY)
        } else {
            p0 + LogProb((p1 - p0).exp().ln_1p())
        }
    }

    pub fn ln_sub_exp(self, other: LogProb) -> LogProb {
        let (p0, p1) = (self, other);
        assert!(p0 >= p1,
                "Subtraction would lead to negative probability, which is undefined in log space.");
        if relative_eq!(*p0, *p1) || p0 == Self::ln_zero() {
            // the first case leads to zero,
            // in the second case p0 and p1 are -inf, which is fine
            Self::ln_zero()
        } else if *p0 == f64::INFINITY {
            LogProb(f64::INFINITY)
        } else {
            p0 + (p1 - p0).ln_one_minus_exp()
        }
    }

    /// Calculate the cumulative sum of the given probabilities in a numerically stable way (Durbin 1998).
    pub fn ln_cumsum_exp<I: IntoIterator<Item = LogProb>>(probs: I) -> ScanIter<I> {
        probs.into_iter().scan(Self::ln_zero(), Self::scan_ln_add_exp)
    }

    /// Integrate numerically stable over given log-space density in the interval [a, b]. Uses the trapezoidal rule with n grid points.
    pub fn ln_integrate_exp<T, D>(density: &D, a: T, b: T, n: usize) -> LogProb
        where T: NumCast + Copy + Add<Output=T> + Sub<Output=T> + Div<Output=T> + Mul<Output=T>, D: Fn(T) -> LogProb, usize: ToFloat<T>
    {
        let mut probs = linspace(a, b, n).dropping(1).dropping_back(1).map(|v| LogProb(*density(v) + 2.0f64.ln())).collect_vec();
        probs.push(density(a));
        probs.push(density(b));
        let width: f64 = cast(b - a).unwrap();

        LogProb(*Self::ln_sum_exp(&probs) + width.ln() - (2.0 * (n - 1) as f64).ln())
    }

    fn scan_ln_add_exp(s: &mut LogProb, p: LogProb) -> Option<LogProb> {
        *s = s.ln_add_exp(p);
        Some(*s)
    }
}

impl From<LogProb> for Prob {
    fn from(p: LogProb) -> Prob {
        Prob(p.exp())
    }
}

impl From<PHREDProb> for Prob {
    fn from(p: PHREDProb) -> Prob {
        Prob(10.0f64.powf(-*p / 10.0))
    }
}

impl From<Prob> for LogProb {
    fn from(p: Prob) -> LogProb {
        LogProb(p.ln())
    }
}

impl From<PHREDProb> for LogProb {
    fn from(p: PHREDProb) -> LogProb {
        LogProb(*p * PHRED_TO_LOG_FACTOR)
    }
}

impl From<LogProb> for PHREDProb {
    fn from(p: LogProb) -> PHREDProb {
        PHREDProb(*p * LOG_TO_PHRED_FACTOR)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::f64;
    use itertools::Itertools;

    #[test]
    fn test_sum() {
        let probs = [LogProb::ln_zero(), LogProb::ln_one(), LogProb::ln_zero()];
        assert_eq!(LogProb::ln_sum_exp(&probs), LogProb::ln_one());
    }

    #[test]
    fn test_empty_sum() {
        assert_eq!(LogProb::ln_sum_exp(&[]), LogProb::ln_zero());
    }

    #[test]
    fn test_cumsum() {
        let probs = vec![LogProb::ln_zero(), LogProb(0.01f64.ln()), LogProb(0.001f64.ln())];
        assert_eq!(LogProb::ln_cumsum_exp(probs).collect_vec(),
                   [LogProb::ln_zero(), LogProb(0.01f64.ln()), LogProb(0.011f64.ln())]);
    }

    #[test]
    fn test_sub() {
        assert_eq!(LogProb::ln_one().ln_sub_exp(LogProb::ln_one()), LogProb::ln_zero());
        assert_relative_eq!(*LogProb::ln_one().ln_sub_exp(LogProb(0.5f64.ln())), *LogProb(0.5f64.ln()));
    }

    #[test]
    fn test_one_minus() {
        assert_eq!(LogProb::ln_zero().ln_one_minus_exp(), LogProb::ln_one());
        assert_eq!(LogProb::ln_one().ln_one_minus_exp(), LogProb::ln_zero());
    }

    #[test]
    fn test_integrate() {
        let density = |_| LogProb(0.1f64.ln());
        let prob = LogProb::ln_integrate_exp(&density, 0.0, 10.0, 5);
        assert_relative_eq!(*prob, *LogProb::ln_one(), epsilon=0.0000001);
    }
}
<commit_msg>Add docs and examples.<commit_after>// Copyright 2014-2016 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Handling log-probabilities.

use std::mem;
use std::f64;
use std::iter;
use std::ops::{Add, Sub, Mul, Div};
use num::traits::cast;
use num::NumCast;

use itertools::linspace;
use itertools::Itertools;
use itertools::misc::ToFloat;

/// A factor to convert log-probabilities to PHRED-scale (phred = p * `LOG_TO_PHRED_FACTOR`).
const LOG_TO_PHRED_FACTOR: f64 = -4.3429448190325175; // -10 * 1 / ln(10)

/// A factor to convert PHRED-scale to log-probabilities (p = phred * `PHRED_TO_LOG_FACTOR`).
const PHRED_TO_LOG_FACTOR: f64 = -0.23025850929940456; // 1 / (-10 * log10(e))

/// Calculate log(1 - p) with p given in log space without loss of precision as described in
/// http://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf.
fn ln_1m_exp(p: f64) -> f64 {
    assert!(p <= 0.0);
    if p < -0.693 {
        (-p.exp()).ln_1p()
    } else {
        (-p.exp_m1()).ln()
    }
}

custom_derive! {
    #[derive(
        NewtypeFrom,
        NewtypeDeref,
        NewtypeAdd(*),
        NewtypeSub(*),
        NewtypeMul(*),
        NewtypeDiv(*),
        PartialEq,
        PartialOrd,
        Copy,
        Clone,
        Debug,
        Default
    )]
    /// A newtype for probabilities.
    ///
    /// # Example
    ///
    /// ```
    /// #[macro_use]
    /// extern crate approx;
    /// # extern crate bio;
    /// # fn main() {
    /// use bio::stats::Prob;
    ///
    /// let p = Prob(0.5);
    /// let q = Prob(0.2);
    ///
    /// assert_relative_eq!(*(p + q), *Prob(0.7));
    /// # }
    /// ```
    pub struct Prob(pub f64);
}

custom_derive! {
    #[derive(
        NewtypeFrom,
        NewtypeDeref,
        NewtypeAdd(*),
        NewtypeSub(*),
        PartialEq,
        PartialOrd,
        Copy,
        Clone,
        Debug
    )]
    /// A newtype for log-scale probabilities.
    ///
    /// # Example
    ///
    /// ```
    /// #[macro_use]
    /// extern crate approx;
    /// # extern crate bio;
    /// # fn main() {
    /// use bio::stats::{LogProb, Prob};
    ///
    /// // convert from probability
    /// let p = LogProb::from(Prob(0.5));
    /// // convert manually
    /// let q = LogProb(0.2f64.ln());
    /// // obtain zero probability in log-space
    /// let o = LogProb::ln_one();
    ///
    /// assert_relative_eq!(*Prob::from(p.ln_add_exp(q) + o), *Prob(0.7));
    /// # }
    /// ```
    pub struct LogProb(pub f64);
}

custom_derive! {
    #[derive(
        NewtypeFrom,
        NewtypeDeref,
        NewtypeAdd(*),
        NewtypeSub(*),
        PartialEq,
        PartialOrd,
        Copy,
        Clone,
        Debug
    )]
    /// A newtype for PHRED-scale probabilities.
    ///
    /// # Example
    ///
    /// ```
    /// #[macro_use]
    /// extern crate approx;
    /// # extern crate bio;
    /// # fn main() {
    /// use bio::stats::{PHREDProb, Prob};
    ///
    /// let p = PHREDProb::from(Prob(0.5));
    ///
    /// assert_relative_eq!(*Prob::from(p), *Prob(0.5));
    /// # }
    /// ```
    pub struct PHREDProb(pub f64);
}

/// Iterator returned by scans over logprobs.
pub type ScanIter<I> = iter::Scan<<I as IntoIterator>::IntoIter, LogProb, fn(&mut LogProb, LogProb) -> Option<LogProb>>;

impl LogProb {
    /// Log-space representation of Pr=0
    pub fn ln_zero() -> LogProb {
        LogProb(f64::NEG_INFINITY)
    }

    /// Log-space representation of Pr=1
    pub fn ln_one() -> LogProb {
        LogProb(0.0)
    }

    /// Numerically stable calculation of 1 - p in log-space.
    pub fn ln_one_minus_exp(&self) -> LogProb {
        LogProb(ln_1m_exp(**self))
    }

    /// Numerically stable sum of probabilities in log-space.
    pub fn ln_sum_exp(probs: &[LogProb]) -> LogProb {
        if probs.is_empty() {
            Self::ln_zero()
        } else {
            let mut pmax = probs[0];
            let mut imax = 0;
            for (i, &p) in probs.iter().enumerate().skip(1) {
                if p > pmax {
                    pmax = p;
                    imax = i;
                }
            }
            if pmax == Self::ln_zero() {
                Self::ln_zero()
            } else if *pmax == f64::INFINITY {
                LogProb(f64::INFINITY)
            } else {
                // TODO use sum() once it has been stabilized: .sum::<usize>()
                pmax + LogProb(
                    (probs.iter()
                          .enumerate()
                          .filter_map(|(i, p)| {
                              if i == imax {
                                  None
                              } else {
                                  Some((p - pmax).exp())
                              }
                          })
                          .fold(0.0, |s, e| s + e)
                     ).ln_1p()
                 )
            }
        }
    }

    /// Numerically stable addition probabilities in log-space.
    pub fn ln_add_exp(self, other: LogProb) -> LogProb {
        let (mut p0, mut p1) = (self, other);
        if p1 > p0 {
            mem::swap(&mut p0, &mut p1);
        }
        if p0 == Self::ln_zero() {
            Self::ln_zero()
        } else if *p0 == f64::INFINITY {
            LogProb(f64::INFINITY)
        } else {
            p0 + LogProb((p1 - p0).exp().ln_1p())
        }
    }

    /// Numerically stable subtraction of probabilities in log-space.
    pub fn ln_sub_exp(self, other: LogProb) -> LogProb {
        let (p0, p1) = (self, other);
        assert!(p0 >= p1,
                "Subtraction would lead to negative probability, which is undefined in log space.");
        if relative_eq!(*p0, *p1) || p0 == Self::ln_zero() {
            // the first case leads to zero,
            // in the second case p0 and p1 are -inf, which is fine
            Self::ln_zero()
        } else if *p0 == f64::INFINITY {
            LogProb(f64::INFINITY)
        } else {
            p0 + (p1 - p0).ln_one_minus_exp()
        }
    }

    /// Calculate the cumulative sum of the given probabilities in a numerically stable way (Durbin 1998).
    pub fn ln_cumsum_exp<I: IntoIterator<Item = LogProb>>(probs: I) -> ScanIter<I> {
        probs.into_iter().scan(Self::ln_zero(), Self::scan_ln_add_exp)
    }

    /// Integrate numerically stable over given log-space density in the interval [a, b]. Uses the trapezoidal rule with n grid points.
    pub fn ln_integrate_exp<T, D>(density: &D, a: T, b: T, n: usize) -> LogProb
        where T: NumCast + Copy + Add<Output=T> + Sub<Output=T> + Div<Output=T> + Mul<Output=T>, D: Fn(T) -> LogProb, usize: ToFloat<T>
    {
        let mut probs = linspace(a, b, n).dropping(1).dropping_back(1).map(|v| LogProb(*density(v) + 2.0f64.ln())).collect_vec();
        probs.push(density(a));
        probs.push(density(b));
        let width: f64 = cast(b - a).unwrap();

        LogProb(*Self::ln_sum_exp(&probs) + width.ln() - (2.0 * (n - 1) as f64).ln())
    }

    fn scan_ln_add_exp(s: &mut LogProb, p: LogProb) -> Option<LogProb> {
        *s = s.ln_add_exp(p);
        Some(*s)
    }
}

impl From<LogProb> for Prob {
    fn from(p: LogProb) -> Prob {
        Prob(p.exp())
    }
}

impl From<PHREDProb> for Prob {
    fn from(p: PHREDProb) -> Prob {
        Prob(10.0f64.powf(-*p / 10.0))
    }
}

impl From<Prob> for LogProb {
    fn from(p: Prob) -> LogProb {
        LogProb(p.ln())
    }
}

impl From<PHREDProb> for LogProb {
    fn from(p: PHREDProb) -> LogProb {
        LogProb(*p * PHRED_TO_LOG_FACTOR)
    }
}

impl From<Prob> for PHREDProb {
    fn from(p: Prob) -> PHREDProb {
        PHREDProb(-10.0 * p.log10())
    }
}

impl From<LogProb> for PHREDProb {
    fn from(p: LogProb) -> PHREDProb {
        PHREDProb(*p * LOG_TO_PHRED_FACTOR)
    }
}

impl Default for LogProb {
    fn default() -> LogProb {
        LogProb::ln_zero()
    }
}

impl Default for PHREDProb {
    fn default() -> PHREDProb {
        PHREDProb::from(Prob(0.0))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use itertools::Itertools;

    #[test]
    fn test_sum() {
        let probs = [LogProb::ln_zero(), LogProb::ln_one(), LogProb::ln_zero()];
        assert_eq!(LogProb::ln_sum_exp(&probs), LogProb::ln_one());
    }

    #[test]
    fn test_empty_sum() {
        assert_eq!(LogProb::ln_sum_exp(&[]), LogProb::ln_zero());
    }

    #[test]
    fn test_cumsum() {
        let probs = vec![LogProb::ln_zero(), LogProb(0.01f64.ln()), LogProb(0.001f64.ln())];
        assert_eq!(LogProb::ln_cumsum_exp(probs).collect_vec(),
                   [LogProb::ln_zero(), LogProb(0.01f64.ln()), LogProb(0.011f64.ln())]);
    }

    #[test]
    fn test_sub() {
        assert_eq!(LogProb::ln_one().ln_sub_exp(LogProb::ln_one()), LogProb::ln_zero());
        assert_relative_eq!(*LogProb::ln_one().ln_sub_exp(LogProb(0.5f64.ln())), *LogProb(0.5f64.ln()));
    }

    #[test]
    fn test_one_minus() {
        assert_eq!(LogProb::ln_zero().ln_one_minus_exp(), LogProb::ln_one());
        assert_eq!(LogProb::ln_one().ln_one_minus_exp(), LogProb::ln_zero());
    }

    #[test]
    fn test_integrate() {
        let density = |_| LogProb(0.1f64.ln());
        let prob = LogProb::ln_integrate_exp(&density, 0.0, 10.0, 5);
        assert_relative_eq!(*prob, *LogProb::ln_one(), epsilon=0.0000001);
    }
}
<|endoftext|>


<commit_before>use grammar::*;
use grammar::Statement::*;
use grammar::Expression::*;

/// The `Transformable` trait provides an interace for instances of grammar
/// to alter the AST, either by mutating self, or by returning a new node.
///
/// NOTE: Returning `None` means no changes are necessary!
trait Transformable {
    fn transform(&mut self) {}
}

impl Transformable for Parameter {}

impl Transformable for Expression {
    fn transform(&mut self) {
        *self = match *self {
            ArrowFunctionExpression {
                ref mut params,
                ref mut body,
            } => {
                params.transform();
                body.transform();

                // return on feature switch
                // return;

                bind_this(FunctionExpression {
                    name: None,
                    params: params.clone(),
                    body: match *body.clone() {
                        BlockStatement { body }   => body,
                        ExpressionStatement(expr) => vec![
                            ReturnStatement(expr)
                        ],
                        statement => {
                            panic!("Invalid arrow function body {:#?}", statement);
                        }
                    },
                })
            },

            CallExpression {
                ref mut callee,
                ref mut arguments,
            } => {
                callee.transform();
                arguments.transform();
                return;
            },

            _ => return,
        }
    }
}

impl<T: Transformable> Transformable for Vec<T> {
    fn transform(&mut self) {
        for item in self.iter_mut() {
            item.transform();
        }
    }
}

impl Transformable for Statement {
    fn transform(&mut self) {
        match *self {
            VariableDeclarationStatement {
                ref mut kind,
                ..
            } => {
                *kind = VariableDeclarationKind::Var;
            },

            ExpressionStatement(ref mut expression) => {
                expression.transform();
            },

            IfStatement {
                ref mut test,
                ref mut consequent,
                ref mut alternate,
                ..
            } => {
                test.transform();
                consequent.transform();
                if let Some(ref mut alternate) = *alternate {
                    alternate.transform();
                }
            },

            BlockStatement {
                ref mut body,
            } => {
                body.transform();
            },

            _ => {},
        }
    }
}

fn bind_this(function: Expression) -> Expression {
    CallExpression {
        callee: Box::new(MemberExpression {
            object: Box::new(function),
            property: Box::new(MemberKey::Literal("bind".to_string())),
        }),
        arguments: vec![ThisExpression]
    }
}

pub fn transform(program: &mut Program) {
    program.body.transform();
}
<commit_msg>Only bind `this` when necessary<commit_after>use grammar::*;
use grammar::Statement::*;
use grammar::Expression::*;
use grammar::ClassMember::*;

/// The `Transformable` trait provides an interace for instances of grammar
/// to alter the AST, either by mutating self, or by returning a new node.
///
/// NOTE: Returning `None` means no changes are necessary!
trait Transformable {
    fn transform(&mut self) {}

    fn contains_this(&self) -> bool {
        false
    }
}

impl Transformable for Parameter {}

impl Transformable for Expression {
    fn transform(&mut self) {
        *self = match *self {
            ArrowFunctionExpression {
                ref mut params,
                ref mut body,
            } => {
                params.transform();
                body.transform();

                // return on feature switch
                // return;

                let body = match *body.clone() {
                    BlockStatement { body }   => body,
                    ExpressionStatement(expr) => vec![
                        ReturnStatement(expr)
                    ],
                    statement => {
                        panic!("Invalid arrow function body {:#?}", statement);
                    }
                };

                let bind = body.contains_this();

                let function = FunctionExpression {
                    name: None,
                    params: params.clone(),
                    body: body,
                };

                if bind {
                    bind_this(function)
                } else {
                    function
                }
            },

            CallExpression {
                ref mut callee,
                ref mut arguments,
            } => {
                callee.transform();
                arguments.transform();
                return;
            },

            _ => return,
        }
    }

    fn contains_this(&self) -> bool {
        match *self {
            ThisExpression => true,

            MemberExpression {
                ref object,
                ..
            } => object.contains_this(),

            CallExpression {
                ref callee,
                ref arguments,
            } => callee.contains_this() || arguments.contains_this(),

            _ => false,
        }
    }
}

impl Transformable for ClassMember {
    fn transform(&mut self) {
        match *self {
            ClassConstructor {
                ref mut params,
                ref mut body,
            } => {
                params.transform();
                body.transform();
            },

            ClassMethod {
                ref mut params,
                ref mut body,
                ..
            } => {
                params.transform();
                body.transform();
            },

            ClassProperty {
                ref mut value,
                ..
            } => {
                value.transform();
            }
        }
    }
}

impl Transformable for Statement {
    fn transform(&mut self) {
        *self = match *self {
            VariableDeclarationStatement {
                ref mut kind,
                ..
            } => {
                *kind = VariableDeclarationKind::Var;
                return;
            },

            ExpressionStatement(ref mut expression) => {
                expression.transform();
                return;
            },

            IfStatement {
                ref mut test,
                ref mut consequent,
                ref mut alternate,
                ..
            } => {
                test.transform();
                consequent.transform();
                if let Some(ref mut alternate) = *alternate {
                    alternate.transform();
                }
                return;
            },

            BlockStatement {
                ref mut body,
            } => {
                body.transform();
                return;
            },

            ClassStatement {
                ref mut body,
                ..
            } => {
                body.transform();
                return;
            }

            _ => return,
        }
    }

    fn contains_this(&self) -> bool {
        match *self {
            ExpressionStatement(ref expression) => expression.contains_this(),

            ReturnStatement(ref expression) => expression.contains_this(),

            _ => false,
        }
    }
}

impl<T: Transformable> Transformable for Vec<T> {
    fn transform(&mut self) {
        for item in self.iter_mut() {
            item.transform();
        }
    }

    fn contains_this(&self) -> bool {
        for item in self {
            if item.contains_this() {
                return true;
            }
        }
        return false;
    }
}

fn bind_this(function: Expression) -> Expression {
    CallExpression {
        callee: Box::new(MemberExpression {
            object: Box::new(function),
            property: Box::new(MemberKey::Literal("bind".to_string())),
        }),
        arguments: vec![ThisExpression]
    }
}

pub fn transform(program: &mut Program) {
    program.body.transform();
}
<|endoftext|>


<commit_before>#![deny(rust_2018_idioms)]

use clap::{crate_version};

use std::env;
use std::path::{Path, PathBuf};

use clap::{App, ArgMatches, SubCommand, AppSettings};

use mdbook_1::{MDBook as MDBook1};
use mdbook_1::errors::{Result as Result1};

use mdbook::MDBook;
use mdbook::errors::{Result as Result3};

#[cfg(all(target_arch = "x86_64", target_os = "linux"))]
use mdbook::renderer::RenderContext;

#[cfg(all(target_arch = "x86_64", target_os = "linux"))]
use mdbook_linkcheck::{self, errors::BrokenLinks};
use failure::Error;

#[cfg(not(all(target_arch = "x86_64", target_os = "linux")))]
use failure::bail;

fn main() {
    let d_message = "-d, --dest-dir=[dest-dir]
'The output directory for your book{n}(Defaults to ./book when omitted)'";
    let dir_message = "[dir]
'A directory for your book{n}(Defaults to Current Directory when omitted)'";
    let vers_message = "-m, --mdbook-vers=[md-version]
'The version of mdbook to use for your book{n}(Defaults to 1 when omitted)'";

    let matches = App::new("rustbook")
                    .about("Build a book with mdBook")
                    .author("Steve Klabnik <steve@steveklabnik.com>")
                    .version(&*format!("v{}", crate_version!()))
                    .setting(AppSettings::SubcommandRequired)
                    .subcommand(SubCommand::with_name("build")
                        .about("Build the book from the markdown files")
                        .arg_from_usage(d_message)
                        .arg_from_usage(dir_message)
                        .arg_from_usage(vers_message))
                    .subcommand(SubCommand::with_name("linkcheck")
                        .about("Run linkcheck with mdBook 3")
                        .arg_from_usage(dir_message))
                    .get_matches();

    // Check which subcomamnd the user ran...
    match matches.subcommand() {
        ("build", Some(sub_matches)) => {
            match sub_matches.value_of("mdbook-vers") {
                None | Some("1") => {
                    if let Err(e) = build_1(sub_matches) {
                        eprintln!("Error: {}", e);

                        for cause in e.iter().skip(1) {
                            eprintln!("\tCaused By: {}", cause);
                        }

                        ::std::process::exit(101);
                    }
                }
                Some("2") | Some("3") => {
                    if let Err(e) = build(sub_matches) {
                        eprintln!("Error: {}", e);

                        for cause in e.iter().skip(1) {
                            eprintln!("\tCaused By: {}", cause);
                        }

                        ::std::process::exit(101);
                    }
                }
                _ => {
                    panic!("Invalid mdBook version! Select '1' or '2' or '3'");
                }
            };
        },
        ("linkcheck", Some(sub_matches)) => {
            if let Err(err) = linkcheck(sub_matches) {
                eprintln!("Error: {}", err);

                #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
                {
                    if let Ok(broken_links) = err.downcast::<BrokenLinks>() {
                        for cause in broken_links.links().iter() {
                            eprintln!("\tCaused By: {}", cause);
                        }
                    }
                }

                ::std::process::exit(101);
            }
        },
        (_, _) => unreachable!(),
    };
}

#[cfg(all(target_arch = "x86_64", target_os = "linux"))]
pub fn linkcheck(args: &ArgMatches<'_>) -> Result<(), Error> {
    let book_dir = get_book_dir(args);
    let book = MDBook::load(&book_dir).unwrap();
    let cfg = book.config;
    let render_ctx = RenderContext::new(&book_dir, book.book, cfg, &book_dir);

    mdbook_linkcheck::check_links(&render_ctx)
}

#[cfg(not(all(target_arch = "x86_64", target_os = "linux")))]
pub fn linkcheck(args: &ArgMatches<'_>) -> Result<(), Error> {
    bail!("mdbook-linkcheck only works on x86_64 linux targets.");
}

// Build command implementation
pub fn build_1(args: &ArgMatches<'_>) -> Result1<()> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook1::load(&book_dir)?;

    // Set this to allow us to catch bugs in advance.
    book.config.build.create_missing = false;

    if let Some(dest_dir) = args.value_of("dest-dir") {
        book.config.build.build_dir = PathBuf::from(dest_dir);
    }

    book.build()?;

    Ok(())
}

// Build command implementation
pub fn build(args: &ArgMatches<'_>) -> Result3<()> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::load(&book_dir)?;

    // Set this to allow us to catch bugs in advance.
    book.config.build.create_missing = false;

    if let Some(dest_dir) = args.value_of("dest-dir") {
        book.config.build.build_dir = PathBuf::from(dest_dir);
    }

    book.build()?;

    Ok(())
}

fn get_book_dir(args: &ArgMatches<'_>) -> PathBuf {
    if let Some(dir) = args.value_of("dir") {
        // Check if path is relative from current dir, or absolute...
        let p = Path::new(dir);
        if p.is_relative() {
            env::current_dir().unwrap().join(dir)
        } else {
            p.to_path_buf()
        }
    } else {
        env::current_dir().unwrap()
    }
}
<commit_msg>Ignore unused variable for non-linux builds<commit_after>#![deny(rust_2018_idioms)]

use clap::{crate_version};

use std::env;
use std::path::{Path, PathBuf};

use clap::{App, ArgMatches, SubCommand, AppSettings};

use mdbook_1::{MDBook as MDBook1};
use mdbook_1::errors::{Result as Result1};

use mdbook::MDBook;
use mdbook::errors::{Result as Result3};

#[cfg(all(target_arch = "x86_64", target_os = "linux"))]
use mdbook::renderer::RenderContext;

#[cfg(all(target_arch = "x86_64", target_os = "linux"))]
use mdbook_linkcheck::{self, errors::BrokenLinks};
use failure::Error;

#[cfg(not(all(target_arch = "x86_64", target_os = "linux")))]
use failure::bail;

fn main() {
    let d_message = "-d, --dest-dir=[dest-dir]
'The output directory for your book{n}(Defaults to ./book when omitted)'";
    let dir_message = "[dir]
'A directory for your book{n}(Defaults to Current Directory when omitted)'";
    let vers_message = "-m, --mdbook-vers=[md-version]
'The version of mdbook to use for your book{n}(Defaults to 1 when omitted)'";

    let matches = App::new("rustbook")
                    .about("Build a book with mdBook")
                    .author("Steve Klabnik <steve@steveklabnik.com>")
                    .version(&*format!("v{}", crate_version!()))
                    .setting(AppSettings::SubcommandRequired)
                    .subcommand(SubCommand::with_name("build")
                        .about("Build the book from the markdown files")
                        .arg_from_usage(d_message)
                        .arg_from_usage(dir_message)
                        .arg_from_usage(vers_message))
                    .subcommand(SubCommand::with_name("linkcheck")
                        .about("Run linkcheck with mdBook 3")
                        .arg_from_usage(dir_message))
                    .get_matches();

    // Check which subcomamnd the user ran...
    match matches.subcommand() {
        ("build", Some(sub_matches)) => {
            match sub_matches.value_of("mdbook-vers") {
                None | Some("1") => {
                    if let Err(e) = build_1(sub_matches) {
                        eprintln!("Error: {}", e);

                        for cause in e.iter().skip(1) {
                            eprintln!("\tCaused By: {}", cause);
                        }

                        ::std::process::exit(101);
                    }
                }
                Some("2") | Some("3") => {
                    if let Err(e) = build(sub_matches) {
                        eprintln!("Error: {}", e);

                        for cause in e.iter().skip(1) {
                            eprintln!("\tCaused By: {}", cause);
                        }

                        ::std::process::exit(101);
                    }
                }
                _ => {
                    panic!("Invalid mdBook version! Select '1' or '2' or '3'");
                }
            };
        },
        ("linkcheck", Some(sub_matches)) => {
            if let Err(err) = linkcheck(sub_matches) {
                eprintln!("Error: {}", err);

                #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
                {
                    if let Ok(broken_links) = err.downcast::<BrokenLinks>() {
                        for cause in broken_links.links().iter() {
                            eprintln!("\tCaused By: {}", cause);
                        }
                    }
                }

                ::std::process::exit(101);
            }
        },
        (_, _) => unreachable!(),
    };
}

#[cfg(all(target_arch = "x86_64", target_os = "linux"))]
pub fn linkcheck(args: &ArgMatches<'_>) -> Result<(), Error> {
    let book_dir = get_book_dir(args);
    let book = MDBook::load(&book_dir).unwrap();
    let cfg = book.config;
    let render_ctx = RenderContext::new(&book_dir, book.book, cfg, &book_dir);

    mdbook_linkcheck::check_links(&render_ctx)
}

#[cfg(not(all(target_arch = "x86_64", target_os = "linux")))]
pub fn linkcheck(_args: &ArgMatches<'_>) -> Result<(), Error> {
    bail!("mdbook-linkcheck only works on x86_64 linux targets.");
}

// Build command implementation
pub fn build_1(args: &ArgMatches<'_>) -> Result1<()> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook1::load(&book_dir)?;

    // Set this to allow us to catch bugs in advance.
    book.config.build.create_missing = false;

    if let Some(dest_dir) = args.value_of("dest-dir") {
        book.config.build.build_dir = PathBuf::from(dest_dir);
    }

    book.build()?;

    Ok(())
}

// Build command implementation
pub fn build(args: &ArgMatches<'_>) -> Result3<()> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::load(&book_dir)?;

    // Set this to allow us to catch bugs in advance.
    book.config.build.create_missing = false;

    if let Some(dest_dir) = args.value_of("dest-dir") {
        book.config.build.build_dir = PathBuf::from(dest_dir);
    }

    book.build()?;

    Ok(())
}

fn get_book_dir(args: &ArgMatches<'_>) -> PathBuf {
    if let Some(dir) = args.value_of("dir") {
        // Check if path is relative from current dir, or absolute...
        let p = Path::new(dir);
        if p.is_relative() {
            env::current_dir().unwrap().join(dir)
        } else {
            p.to_path_buf()
        }
    } else {
        env::current_dir().unwrap()
    }
}
<|endoftext|>