<commit_before>use std::slice;

use super::{StreamEncrypt, StreamDecrypt};

use byteorder::{ByteOrder, LittleEndian};

const ROUNDS: usize = 20;
const STATE_WORDS: usize = 16;
const STATE_BYTES: usize = STATE_WORDS * 4;

#[derive(Copy, Clone)]
struct State([u32; STATE_WORDS]);

macro_rules! quarter_round {
    ($a:expr, $b:expr, $c:expr, $d:expr) => {{
        $a = $a.wrapping_add($b); $d = $d ^ $a; $d = $d.rotate_left(16);
        $c = $c.wrapping_add($d); $b = $b ^ $c; $b = $b.rotate_left(12);
        $a = $a.wrapping_add($b); $d = $d ^ $a; $d = $d.rotate_left( 8);
        $c = $c.wrapping_add($d); $b = $b ^ $c; $b = $b.rotate_left( 7);
    }}
}

macro_rules! double_round {
    ($x:expr) => {{
        // Column round
        quarter_round!($x[ 0], $x[ 4], $x[ 8], $x[12]);
        quarter_round!($x[ 1], $x[ 5], $x[ 9], $x[13]);
        quarter_round!($x[ 2], $x[ 6], $x[10], $x[14]);
        quarter_round!($x[ 3], $x[ 7], $x[11], $x[15]);
        // Diagonal round
        quarter_round!($x[ 0], $x[ 5], $x[10], $x[15]);
        quarter_round!($x[ 1], $x[ 6], $x[11], $x[12]);
        quarter_round!($x[ 2], $x[ 7], $x[ 8], $x[13]);
        quarter_round!($x[ 3], $x[ 4], $x[ 9], $x[14]);
    }}
}

impl State {
    fn expand(key: &[u8], nonce: &[u8], position: u32) -> Self {
        let mut state = [0u32; STATE_WORDS];

        state[0] = 0x61707865;
        state[1] = 0x3320646e;
        state[2] = 0x79622d32;
        state[3] = 0x6b206574;

        for (state, chunk) in state[4..12].iter_mut().zip(key.chunks(4)) {
            *state = LittleEndian::read_u32(chunk);
        }

        state[12] = position;

        for (state, chunk) in state[13..16].iter_mut().zip(nonce.chunks(4)) {
            *state = LittleEndian::read_u32(chunk);
        }

        State(state)
    }

    fn update(&mut self, output: &mut [u32]) {
        let mut state = self.0;

        for _ in 0..ROUNDS / 2 {
            double_round!(state);
        }

        for i in 0..STATE_WORDS {
            output[i] = self.0[i].wrapping_add(state[i]);
        }

        self.0[12] = self.0[12] + 1;
    }
}

pub struct ChaCha20 {
    state: State,
    buffer: [u32; STATE_WORDS],
    index: usize,
}

impl ChaCha20 {
    pub fn init(key: &[u8], nonce: &[u8], position: u32) -> Self {
        ChaCha20 {
            state: State::expand(key.as_ref(), nonce.as_ref(), position),
            buffer: [0; STATE_WORDS],
            index: STATE_BYTES,
        }
    }

    pub fn new<Key, Nonce>(key: Key, nonce: Nonce) -> Self
        where Key: AsRef<[u8]>,
              Nonce: AsRef<[u8]>
    {
        Self::init(key.as_ref(), nonce.as_ref(), 1)
    }

    fn update(&mut self) {
        self.state.update(&mut self.buffer[..]);

        self.index = 0;
    }

    fn crypt(&mut self, input: &[u8], output: &mut [u8]) {
        assert_eq!(input.len(), output.len());

        if self.index == STATE_BYTES {
            self.update()
        }

        let buffer = unsafe {
            slice::from_raw_parts(self.buffer.as_ptr() as *const u8, STATE_BYTES)
        };

        for i in self.index..input.len() {
            output[i] = input[i] ^ buffer[i];
        }

        self.index = input.len();
    }
}

impl StreamEncrypt for ChaCha20 {
    fn encrypt_stream<I, O>(&mut self, input: I, mut output: O)
        where I: AsRef<[u8]>,
              O: AsMut<[u8]>
    {
        assert_eq!(input.as_ref().len(), output.as_mut().len());

        let input = input.as_ref().chunks(STATE_BYTES);
        let output = output.as_mut().chunks_mut(STATE_BYTES);
        for (i, o) in input.zip(output) {
            self.crypt(i, o)
        }
    }
}

impl StreamDecrypt for ChaCha20 {
    fn decrypt_stream<I, O>(&mut self, input: I, mut output: O)
        where I: AsRef<[u8]>,
              O: AsMut<[u8]>
    {
        assert_eq!(input.as_ref().len(), output.as_mut().len());

        let input = input.as_ref().chunks(STATE_BYTES);
        let output = output.as_mut().chunks_mut(STATE_BYTES);
        for (i, o) in input.zip(output) {
            self.crypt(i, o)
        }
    }
}
<commit_msg>fix(crypto::stream::chacha20): Fix buffer utilisation<commit_after>use std::slice;

use super::{StreamEncrypt, StreamDecrypt};

use byteorder::{ByteOrder, LittleEndian};

const ROUNDS: usize = 20;
const STATE_WORDS: usize = 16;
const STATE_BYTES: usize = STATE_WORDS * 4;

#[derive(Copy, Clone)]
struct State([u32; STATE_WORDS]);

macro_rules! quarter_round {
    ($a:expr, $b:expr, $c:expr, $d:expr) => {{
        $a = $a.wrapping_add($b); $d = $d ^ $a; $d = $d.rotate_left(16);
        $c = $c.wrapping_add($d); $b = $b ^ $c; $b = $b.rotate_left(12);
        $a = $a.wrapping_add($b); $d = $d ^ $a; $d = $d.rotate_left( 8);
        $c = $c.wrapping_add($d); $b = $b ^ $c; $b = $b.rotate_left( 7);
    }}
}

macro_rules! double_round {
    ($x:expr) => {{
        // Column round
        quarter_round!($x[ 0], $x[ 4], $x[ 8], $x[12]);
        quarter_round!($x[ 1], $x[ 5], $x[ 9], $x[13]);
        quarter_round!($x[ 2], $x[ 6], $x[10], $x[14]);
        quarter_round!($x[ 3], $x[ 7], $x[11], $x[15]);
        // Diagonal round
        quarter_round!($x[ 0], $x[ 5], $x[10], $x[15]);
        quarter_round!($x[ 1], $x[ 6], $x[11], $x[12]);
        quarter_round!($x[ 2], $x[ 7], $x[ 8], $x[13]);
        quarter_round!($x[ 3], $x[ 4], $x[ 9], $x[14]);
    }}
}

impl State {
    fn expand(key: &[u8], nonce: &[u8], position: u32) -> Self {
        let mut state = [0u32; STATE_WORDS];

        state[0] = 0x61707865;
        state[1] = 0x3320646e;
        state[2] = 0x79622d32;
        state[3] = 0x6b206574;

        for (state, chunk) in state[4..12].iter_mut().zip(key.chunks(4)) {
            *state = LittleEndian::read_u32(chunk);
        }

        state[12] = position;

        for (state, chunk) in state[13..16].iter_mut().zip(nonce.chunks(4)) {
            *state = LittleEndian::read_u32(chunk);
        }

        State(state)
    }

    fn update(&mut self, output: &mut [u32]) {
        let mut state = self.0;

        for _ in 0..ROUNDS / 2 {
            double_round!(state);
        }

        for i in 0..STATE_WORDS {
            output[i] = self.0[i].wrapping_add(state[i]);
        }

        self.0[12] = self.0[12] + 1;
    }
}

pub struct ChaCha20 {
    state: State,
    buffer: [u32; STATE_WORDS],
    index: usize,
}

impl ChaCha20 {
    pub fn init(key: &[u8], nonce: &[u8], position: u32) -> Self {
        ChaCha20 {
            state: State::expand(key.as_ref(), nonce.as_ref(), position),
            buffer: [0; STATE_WORDS],
            index: STATE_BYTES,
        }
    }

    pub fn new<Key, Nonce>(key: Key, nonce: Nonce) -> Self
        where Key: AsRef<[u8]>,
              Nonce: AsRef<[u8]>
    {
        Self::init(key.as_ref(), nonce.as_ref(), 1)
    }

    fn update(&mut self) {
        self.state.update(&mut self.buffer[..]);

        self.index = 0;
    }

    fn crypt(&mut self, input: &[u8], output: &mut [u8]) {
        if self.index == STATE_BYTES {
            self.update()
        }

        let buffer = unsafe {
            slice::from_raw_parts(self.buffer.as_ptr() as *const u8, STATE_BYTES)
        };

        for i in self.index..input.len() {
            output[i] = input[i] ^ buffer[i];
        }

        self.index = input.len();
    }
}

impl StreamEncrypt for ChaCha20 {
    fn encrypt_stream<I, O>(&mut self, input: I, mut output: O)
        where I: AsRef<[u8]>,
              O: AsMut<[u8]>
    {
        assert_eq!(input.as_ref().len(), output.as_mut().len());
        let input = input.as_ref();
        let output = output.as_mut();

        let from = STATE_BYTES - self.index;

        if from > 0 {
            self.crypt(&input[..from], &mut output[..from]);
        }

        for (i, o) in input[from..]
                          .chunks(STATE_BYTES)
                          .zip(output[from..].chunks_mut(STATE_BYTES)) {
            self.crypt(i, o)
        }
    }
}

impl StreamDecrypt for ChaCha20 {
    fn decrypt_stream<I, O>(&mut self, input: I, mut output: O)
        where I: AsRef<[u8]>,
              O: AsMut<[u8]>
    {
        assert_eq!(input.as_ref().len(), output.as_mut().len());

        let input = input.as_ref().chunks(STATE_BYTES);
        let output = output.as_mut().chunks_mut(STATE_BYTES);
        for (i, o) in input.zip(output) {
            self.crypt(i, o)
        }
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::fs::File;
use std::path::PathBuf;

use simplelog::{CombinedLogger, LevelFilter, TerminalMode, TermLogger};

use crate::core::set_env;
use crate::pathenv::PathVec;

pub fn setup_logging(debug: bool) {
    let log_level = if debug { LevelFilter::Debug } else { LevelFilter::Info };
    let config = simplelog::ConfigBuilder::new()
        .set_thread_level(LevelFilter::Trace)
        .set_target_level(LevelFilter::Trace)
        .build();
    CombinedLogger::init(
        vec![
            TermLogger::new(log_level,
                            config,
                            TerminalMode::Stderr).unwrap()
        ]
    ).unwrap();
}

pub fn setup_environment(nmk_dir: &PathBuf) {
    let init_vim = nmk_dir.join("vim").join("init.vim");
    let zdotdir = nmk_dir.join("zsh");
    set_env("NMK_DIR", nmk_dir);

    if let Some(path) = init_vim
        .to_str()
        .map(|s|
            s.to_string().replace(" ", r"\ ")
        ) {
        set_env("VIMINIT", format!("source {}", path));
    }
    set_env("ZDOTDIR", zdotdir);

    env::remove_var("VIRTUAL_ENV");

    set_env("NMK_BIN", env::current_exe().unwrap());
}

pub fn setup_preferred_editor() {
    const EDITOR: &str = "EDITOR";
    if env::var_os(EDITOR).is_none() {
        let mut editors = ["nvim", "vim"].iter();
        if let Some(editor) = editors.find(|bin| which::which(bin).is_ok()) {
            set_env(EDITOR, editor);
            debug!("using {} as prefer editor", editor);
        }
    }
}

pub fn setup_path(nmk_dir: &PathBuf) {
    const PATH: &str = "PATH";
    let mut p = PathVec::parse(env::var_os(PATH).expect("$PATH not found"));
    p.push_front(nmk_dir.join("local").join("bin"));
    p.push_front(nmk_dir.join("bin"));
    p = p.unique().no_version_managers();
    if log_enabled!(log::Level::Debug) {
        for (i, path) in p.iter().enumerate() {
            debug!("{}[{}]={:?}", PATH, i + 1, path);
        }
    }
    set_env(PATH, p.make());
}

pub fn setup_ld_library_path(nmk_dir: &PathBuf) {
    const LD: &str = "LD_LIBRARY_PATH";

    let local_lib_dir = nmk_dir.join("local").join("lib");
    if local_lib_dir.exists() {
        let mut ps = match env::var_os(LD) {
            Some(path) => {
                debug!("{}: {:?}", LD, path);
                PathVec::parse(path)
            }
            None => PathVec::new(),
        };
        ps.push_front(local_lib_dir);
        let next_ld = ps.make();
        set_env(LD, next_ld);
    }
}

pub fn dir() -> PathBuf {
    const NMK_DIR: &str = "NMK_DIR";
    let path = match env::var_os(NMK_DIR) {
        Some(s) => PathBuf::from(s),
        None => dirs::home_dir()
            .expect("Can't find home directory")
            .join(".nmk"),
    };
    if !path.exists() {
        panic!(format!("{:?} doesn't exist", path));
    }
    path
}

pub fn print_message_of_the_day() {
    let mut stdout = std::io::stdout();
    ["/var/run/motd.dynamic", "/etc/motd"]
        .iter()
        .map(PathBuf::from)
        .filter(|p| p.exists())
        .flat_map(File::open)
        .for_each(|mut f| {
            std::io::copy(&mut f, &mut stdout).expect("fail to print motd");
        });
}<commit_msg>nmk.rs: refactoring<commit_after>use std::env;
use std::fs::File;
use std::path::PathBuf;

use simplelog::{LevelFilter, TerminalMode, TermLogger};

use crate::core::set_env;
use crate::pathenv::PathVec;

pub fn setup_logging(debug: bool) {
    let log_level = if debug { LevelFilter::Debug } else { LevelFilter::Info };
    let config = simplelog::ConfigBuilder::new()
        .set_thread_level(LevelFilter::Trace)
        .set_target_level(LevelFilter::Trace)
        .build();
    TermLogger::init(log_level,
                     config,
                     TerminalMode::Stderr).unwrap();
}

pub fn setup_environment(nmk_dir: &PathBuf) {
    let init_vim = nmk_dir.join("vim").join("init.vim");
    let zdotdir = nmk_dir.join("zsh");
    set_env("NMK_DIR", nmk_dir);

    let quote_first_space = |s: &str| s.to_string().replace(" ", r"\ ");
    if let Some(path) = init_vim.to_str().map(quote_first_space) {
        set_env("VIMINIT", format!("source {}", path));
    }
    set_env("ZDOTDIR", zdotdir);

    env::remove_var("VIRTUAL_ENV");

    set_env("NMK_BIN", env::current_exe().unwrap());
}

pub fn setup_preferred_editor() {
    const EDITOR: &str = "EDITOR";
    match env::var_os(EDITOR) {
        Some(editor) => {
            debug!("using {:?} as preferred editor", editor);
        }
        None => {
            let mut editors = ["nvim", "vim"].iter();
            if let Some(editor) = editors.find(|bin| which::which(bin).is_ok()) {
                set_env(EDITOR, editor);
                debug!("using {} as preferred editor", editor);
            }
        }
    }
}

pub fn setup_path(nmk_dir: &PathBuf) {
    const PATH: &str = "PATH";
    let mut p = PathVec::parse(env::var_os(PATH).expect("$PATH not found"));
    p.push_front(nmk_dir.join("local").join("bin"));
    p.push_front(nmk_dir.join("bin"));
    p = p.unique().no_version_managers();
    if log_enabled!(log::Level::Debug) {
        for (i, path) in p.iter().enumerate() {
            debug!("{}[{}]={:?}", PATH, i + 1, path);
        }
    }
    set_env(PATH, p.make());
}

pub fn setup_ld_library_path(nmk_dir: &PathBuf) {
    const LD: &str = "LD_LIBRARY_PATH";

    let local_lib_dir = nmk_dir.join("local").join("lib");
    if local_lib_dir.exists() {
        let mut ps = match env::var_os(LD) {
            Some(path) => {
                debug!("{}: {:?}", LD, path);
                PathVec::parse(path)
            }
            None => PathVec::new(),
        };
        ps.push_front(local_lib_dir);
        let next_ld = ps.make();
        set_env(LD, next_ld);
    }
}

pub fn dir() -> PathBuf {
    const NMK_DIR: &str = "NMK_DIR";
    let path = match env::var_os(NMK_DIR) {
        Some(s) => PathBuf::from(s),
        None => dirs::home_dir()
            .expect("Can't find home directory")
            .join(".nmk"),
    };
    if !path.exists() {
        panic!(format!("{:?} doesn't exist", path));
    }
    path
}

pub fn print_message_of_the_day() {
    let mut stdout = std::io::stdout();
    ["/var/run/motd.dynamic", "/etc/motd"]
        .iter()
        .map(PathBuf::from)
        .filter(|p| p.exists())
        .flat_map(File::open)
        .for_each(|mut f| {
            std::io::copy(&mut f, &mut stdout).expect("fail to print motd");
        });
}<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
// Copyright 2016-2017 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The binomial distribution.

use Rng;
use distributions::Distribution;

/// The binomial distribution `Binomial(n, p)`.
///
/// This distribution has density function:
/// `f(k) = n!/(k! (n-k)!) p^k (1-p)^(n-k)` for `k >= 0`.
///
/// # Example
///
/// ```
/// use rand::distributions::{Binomial, Distribution};
///
/// let bin = Binomial::new(20, 0.3);
/// let v = bin.sample(&mut rand::thread_rng());
/// println!("{} is from a binomial distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Binomial {
    /// Number of trials.
    n: u64,
    /// Probability of success.
    p: f64,
}

impl Binomial {
    /// Construct a new `Binomial` with the given shape parameters `n` (number
    /// of trials) and `p` (probability of success).
    ///
    /// Panics if `p < 0` or `p > 1`.
    pub fn new(n: u64, p: f64) -> Binomial {
        assert!(p >= 0.0, "Binomial::new called with p < 0");
        assert!(p <= 1.0, "Binomial::new called with p > 1");
        Binomial { n, p }
    }
}

/// Convert a `f64` to an `i64`, panicing on overflow.
// In the future (Rust 1.34), this might be replaced with `TryFrom`.
fn f64_to_i64(x: f64) -> i64 {
    assert!(x < (::std::i64::MAX as f64));
    x as i64
}

impl Distribution<u64> for Binomial {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
        // Handle these values directly.
        if self.p == 0.0 {
            return 0;
        } else if self.p == 1.0 {
            return self.n;
        }

        // The binomial distribution is symmetrical with respect to p -> 1-p,
        // k -> n-k switch p so that it is less than 0.5 - this allows for lower
        // expected values we will just invert the result at the end
        let p = if self.p <= 0.5 {
            self.p
        } else {
            1.0 - self.p
        };

        let result;
        let q = 1. - p;

        // For small n * min(p, 1 - p), the BINV algorithm based on the inverse
        // transformation of the binomial distribution is efficient. Otherwise,
        // the BTPE algorithm is used.
        //
        // Voratas Kachitvichyanukul and Bruce W. Schmeiser. 1988. Binomial
        // random variate generation. Commun. ACM 31, 2 (February 1988),
        // 216-222. http://dx.doi.org/10.1145/42372.42381

        // Threshold for prefering the BINV algorithm. The paper suggests 10,
        // Ranlib uses 30, and GSL uses 14.
        const BINV_THRESHOLD: f64 = 10.;

        if (self.n as f64) * p < BINV_THRESHOLD &&
           self.n <= (::std::i32::MAX as u64) {
            // Use the BINV algorithm.
            let s = p / q;
            let a = ((self.n + 1) as f64) * s;
            let mut r = q.powi(self.n as i32);
            let mut u: f64 = rng.gen();
            let mut x = 0;
            while u > r as f64 {
                u -= r;
                x += 1;
                r *= a / (x as f64) - s;
            }
            result = x;
        } else {
            // Use the BTPE algorithm.

            // Threshold for using the squeeze algorithm. This can be freely
            // chosen based on performance. Ranlib and GSL use 20.
            const SQUEEZE_THRESHOLD: i64 = 20;

            // Step 0: Calculate constants as functions of `n` and `p`.
            let n = self.n as f64;
            let np = n * p;
            let npq = np * q;
            let f_m = np + p;
            let m = f64_to_i64(f_m);
            // radius of triangle region, since height=1 also area of region
            let p1 = (2.195 * npq.sqrt() - 4.6 * q).floor() + 0.5;
            // tip of triangle
            let x_m = (m as f64) + 0.5;
            // left edge of triangle
            let x_l = x_m - p1;
            // right edge of triangle
            let x_r = x_m + p1;
            let c = 0.134 + 20.5 / (15.3 + (m as f64));
            // p1 + area of parallelogram region
            let p2 = p1 * (1. + 2. * c);

            fn lambda(a: f64) -> f64 {
                a * (1. + 0.5 * a)
            }

            let lambda_l = lambda((f_m - x_l) / (f_m - x_l * p));
            let lambda_r = lambda((x_r - f_m) / (x_r * q));
            // p1 + area of left tail
            let p3 = p2 + c / lambda_l;
            // p1 + area of right tail
            let p4 = p3 + c / lambda_r;

            // return value
            let mut y: i64;

            loop {
                // Step 1: Generate `u` for selecting the region. If region 1 is
                // selected, generate a triangularly distributed variate.
                let u = rng.gen_range(0., p4);
                let mut v = rng.gen_range(0., 1.);
                if !(u > p1) {
                    y = f64_to_i64(x_m - p1 * v + u);
                    break;
                }

                if !(u > p2) {
                    // Step 2: Region 2, parallelograms. Check if region 2 is
                    // used. If so, generate `y`.
                    let x = x_l + (u - p1) / c;
                    v = v * c + 1.0 - (x - x_m).abs() / p1;
                    if v > 1. {
                        continue;
                    } else {
                        y = f64_to_i64(x);
                    }
                } else if !(u > p3) {
                    // Step 3: Region 3, left exponential tail.
                    y = f64_to_i64(x_l + v.ln() / lambda_l);
                    if y < 0 {
                        continue;
                    } else {
                        v *= (u - p2) * lambda_l;
                    }
                } else {
                    // Step 4: Region 4, right exponential tail.
                    y = f64_to_i64(x_r - v.ln() / lambda_r);
                    if y > 0 && (y as u64) > self.n {
                        continue;
                    } else {
                        v *= (u - p3) * lambda_r;
                    }
                }

                // Step 5: Acceptance/rejection comparison.

                // Step 5.0: Test for appropriate method of evaluating f(y).
                let k = (y - m).abs();
                if !(k > SQUEEZE_THRESHOLD && (k as f64) < 0.5 * npq - 1.) {
                    // Step 5.1: Evaluate f(y) via the recursive relationship. Start the
                    // search from the mode.
                    let s = p / q;
                    let a = s * (n + 1.);
                    let mut f = 1.0;
                    if m < y {
                        let mut i = m;
                        loop {
                            i += 1;
                            f *= a / (i as f64) - s;
                            if i == y {
                                break;
                            }
                        }
                    } else if m > y {
                        let mut i = y;
                        loop {
                            i += 1;
                            f /= a / (i as f64) - s;
                            if i == m {
                                break;
                            }
                        }
                    }
                    if v > f {
                        continue;
                    } else {
                        break;
                    }
                }

                // Step 5.2: Squeezing. Check the value of ln(v) againts upper and
                // lower bound of ln(f(y)).
                let k = k as f64;
                let rho = (k / npq) * ((k * (k / 3. + 0.625) + 1./6.) / npq + 0.5);
                let t = -0.5 * k*k / npq;
                let alpha = v.ln();
                if alpha < t - rho {
                    break;
                }
                if alpha > t + rho {
                    continue;
                }

                // Step 5.3: Final acceptance/rejection test.
                let x1 = (y + 1) as f64;
                let f1 = (m + 1) as f64;
                let z = (f64_to_i64(n) + 1 - m) as f64;
                let w = (f64_to_i64(n) - y + 1) as f64;

                fn stirling(a: f64) -> f64 {
                    let a2 = a * a;
                    (13860. - (462. - (132. - (99. - 140. / a2) / a2) / a2) / a2) / a / 166320.
                }

                if alpha > x_m * (f1 / x1).ln()
                    + (n - (m as f64) + 0.5) * (z / w).ln()
                    + ((y - m) as f64) * (w * p / (x1 * q)).ln()
                    // We use the signs from the GSL implementation, which are
                    // different than the ones in the reference. According to
                    // the GSL authors, the new signs were verified to be
                    // correct by one of the original designers of the
                    // algorithm.
                    + stirling(f1) + stirling(z) - stirling(x1) - stirling(w)
                {
                    continue;
                }

                break;
            }
            assert!(y >= 0);
            result = y as u64;
        }

        // Invert the result for p < 0.5.
        if p != self.p {
            self.n - result
        } else {
            result
        }
    }
}

#[cfg(test)]
mod test {
    use Rng;
    use distributions::Distribution;
    use super::Binomial;

    fn test_binomial_mean_and_variance<R: Rng>(n: u64, p: f64, rng: &mut R) {
        let binomial = Binomial::new(n, p);

        let expected_mean = n as f64 * p;
        let expected_variance = n as f64 * p * (1.0 - p);

        let mut results = [0.0; 1000];
        for i in results.iter_mut() { *i = binomial.sample(rng) as f64; }

        let mean = results.iter().sum::<f64>() / results.len() as f64;
        assert!((mean as f64 - expected_mean).abs() < expected_mean / 50.0,
                "mean: {}, expected_mean: {}", mean, expected_mean);

        let variance =
            results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>()
            / results.len() as f64;
        assert!((variance - expected_variance).abs() < expected_variance / 10.0,
                "variance: {}, expected_variance: {}", variance, expected_variance);
    }

    #[test]
    fn test_binomial() {
        let mut rng = ::test::rng(351);
        test_binomial_mean_and_variance(150, 0.1, &mut rng);
        test_binomial_mean_and_variance(70, 0.6, &mut rng);
        test_binomial_mean_and_variance(40, 0.5, &mut rng);
        test_binomial_mean_and_variance(20, 0.7, &mut rng);
        test_binomial_mean_and_variance(20, 0.5, &mut rng);
    }

    #[test]
    fn test_binomial_end_points() {
        let mut rng = ::test::rng(352);
        assert_eq!(rng.sample(Binomial::new(20, 0.0)), 0);
        assert_eq!(rng.sample(Binomial::new(20, 1.0)), 20);
    }

    #[test]
    #[should_panic]
    fn test_binomial_invalid_lambda_neg() {
        Binomial::new(20, -10.0);
    }
}
<commit_msg>Binomial: Move distribution initialization out of loop<commit_after>// Copyright 2018 Developers of the Rand project.
// Copyright 2016-2017 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The binomial distribution.

use Rng;
use distributions::{Distribution, Uniform};

/// The binomial distribution `Binomial(n, p)`.
///
/// This distribution has density function:
/// `f(k) = n!/(k! (n-k)!) p^k (1-p)^(n-k)` for `k >= 0`.
///
/// # Example
///
/// ```
/// use rand::distributions::{Binomial, Distribution};
///
/// let bin = Binomial::new(20, 0.3);
/// let v = bin.sample(&mut rand::thread_rng());
/// println!("{} is from a binomial distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Binomial {
    /// Number of trials.
    n: u64,
    /// Probability of success.
    p: f64,
}

impl Binomial {
    /// Construct a new `Binomial` with the given shape parameters `n` (number
    /// of trials) and `p` (probability of success).
    ///
    /// Panics if `p < 0` or `p > 1`.
    pub fn new(n: u64, p: f64) -> Binomial {
        assert!(p >= 0.0, "Binomial::new called with p < 0");
        assert!(p <= 1.0, "Binomial::new called with p > 1");
        Binomial { n, p }
    }
}

/// Convert a `f64` to an `i64`, panicing on overflow.
// In the future (Rust 1.34), this might be replaced with `TryFrom`.
fn f64_to_i64(x: f64) -> i64 {
    assert!(x < (::std::i64::MAX as f64));
    x as i64
}

impl Distribution<u64> for Binomial {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
        // Handle these values directly.
        if self.p == 0.0 {
            return 0;
        } else if self.p == 1.0 {
            return self.n;
        }

        // The binomial distribution is symmetrical with respect to p -> 1-p,
        // k -> n-k switch p so that it is less than 0.5 - this allows for lower
        // expected values we will just invert the result at the end
        let p = if self.p <= 0.5 {
            self.p
        } else {
            1.0 - self.p
        };

        let result;
        let q = 1. - p;

        // For small n * min(p, 1 - p), the BINV algorithm based on the inverse
        // transformation of the binomial distribution is efficient. Otherwise,
        // the BTPE algorithm is used.
        //
        // Voratas Kachitvichyanukul and Bruce W. Schmeiser. 1988. Binomial
        // random variate generation. Commun. ACM 31, 2 (February 1988),
        // 216-222. http://dx.doi.org/10.1145/42372.42381

        // Threshold for prefering the BINV algorithm. The paper suggests 10,
        // Ranlib uses 30, and GSL uses 14.
        const BINV_THRESHOLD: f64 = 10.;

        if (self.n as f64) * p < BINV_THRESHOLD &&
           self.n <= (::std::i32::MAX as u64) {
            // Use the BINV algorithm.
            let s = p / q;
            let a = ((self.n + 1) as f64) * s;
            let mut r = q.powi(self.n as i32);
            let mut u: f64 = rng.gen();
            let mut x = 0;
            while u > r as f64 {
                u -= r;
                x += 1;
                r *= a / (x as f64) - s;
            }
            result = x;
        } else {
            // Use the BTPE algorithm.

            // Threshold for using the squeeze algorithm. This can be freely
            // chosen based on performance. Ranlib and GSL use 20.
            const SQUEEZE_THRESHOLD: i64 = 20;

            // Step 0: Calculate constants as functions of `n` and `p`.
            let n = self.n as f64;
            let np = n * p;
            let npq = np * q;
            let f_m = np + p;
            let m = f64_to_i64(f_m);
            // radius of triangle region, since height=1 also area of region
            let p1 = (2.195 * npq.sqrt() - 4.6 * q).floor() + 0.5;
            // tip of triangle
            let x_m = (m as f64) + 0.5;
            // left edge of triangle
            let x_l = x_m - p1;
            // right edge of triangle
            let x_r = x_m + p1;
            let c = 0.134 + 20.5 / (15.3 + (m as f64));
            // p1 + area of parallelogram region
            let p2 = p1 * (1. + 2. * c);

            fn lambda(a: f64) -> f64 {
                a * (1. + 0.5 * a)
            }

            let lambda_l = lambda((f_m - x_l) / (f_m - x_l * p));
            let lambda_r = lambda((x_r - f_m) / (x_r * q));
            // p1 + area of left tail
            let p3 = p2 + c / lambda_l;
            // p1 + area of right tail
            let p4 = p3 + c / lambda_r;

            // return value
            let mut y: i64;

            let gen_u = Uniform::new(0., p4);
            let gen_v = Uniform::new(0., 1.);

            loop {
                // Step 1: Generate `u` for selecting the region. If region 1 is
                // selected, generate a triangularly distributed variate.
                let u = gen_u.sample(rng);
                let mut v = gen_v.sample(rng);
                if !(u > p1) {
                    y = f64_to_i64(x_m - p1 * v + u);
                    break;
                }

                if !(u > p2) {
                    // Step 2: Region 2, parallelograms. Check if region 2 is
                    // used. If so, generate `y`.
                    let x = x_l + (u - p1) / c;
                    v = v * c + 1.0 - (x - x_m).abs() / p1;
                    if v > 1. {
                        continue;
                    } else {
                        y = f64_to_i64(x);
                    }
                } else if !(u > p3) {
                    // Step 3: Region 3, left exponential tail.
                    y = f64_to_i64(x_l + v.ln() / lambda_l);
                    if y < 0 {
                        continue;
                    } else {
                        v *= (u - p2) * lambda_l;
                    }
                } else {
                    // Step 4: Region 4, right exponential tail.
                    y = f64_to_i64(x_r - v.ln() / lambda_r);
                    if y > 0 && (y as u64) > self.n {
                        continue;
                    } else {
                        v *= (u - p3) * lambda_r;
                    }
                }

                // Step 5: Acceptance/rejection comparison.

                // Step 5.0: Test for appropriate method of evaluating f(y).
                let k = (y - m).abs();
                if !(k > SQUEEZE_THRESHOLD && (k as f64) < 0.5 * npq - 1.) {
                    // Step 5.1: Evaluate f(y) via the recursive relationship. Start the
                    // search from the mode.
                    let s = p / q;
                    let a = s * (n + 1.);
                    let mut f = 1.0;
                    if m < y {
                        let mut i = m;
                        loop {
                            i += 1;
                            f *= a / (i as f64) - s;
                            if i == y {
                                break;
                            }
                        }
                    } else if m > y {
                        let mut i = y;
                        loop {
                            i += 1;
                            f /= a / (i as f64) - s;
                            if i == m {
                                break;
                            }
                        }
                    }
                    if v > f {
                        continue;
                    } else {
                        break;
                    }
                }

                // Step 5.2: Squeezing. Check the value of ln(v) againts upper and
                // lower bound of ln(f(y)).
                let k = k as f64;
                let rho = (k / npq) * ((k * (k / 3. + 0.625) + 1./6.) / npq + 0.5);
                let t = -0.5 * k*k / npq;
                let alpha = v.ln();
                if alpha < t - rho {
                    break;
                }
                if alpha > t + rho {
                    continue;
                }

                // Step 5.3: Final acceptance/rejection test.
                let x1 = (y + 1) as f64;
                let f1 = (m + 1) as f64;
                let z = (f64_to_i64(n) + 1 - m) as f64;
                let w = (f64_to_i64(n) - y + 1) as f64;

                fn stirling(a: f64) -> f64 {
                    let a2 = a * a;
                    (13860. - (462. - (132. - (99. - 140. / a2) / a2) / a2) / a2) / a / 166320.
                }

                if alpha > x_m * (f1 / x1).ln()
                    + (n - (m as f64) + 0.5) * (z / w).ln()
                    + ((y - m) as f64) * (w * p / (x1 * q)).ln()
                    // We use the signs from the GSL implementation, which are
                    // different than the ones in the reference. According to
                    // the GSL authors, the new signs were verified to be
                    // correct by one of the original designers of the
                    // algorithm.
                    + stirling(f1) + stirling(z) - stirling(x1) - stirling(w)
                {
                    continue;
                }

                break;
            }
            assert!(y >= 0);
            result = y as u64;
        }

        // Invert the result for p < 0.5.
        if p != self.p {
            self.n - result
        } else {
            result
        }
    }
}

#[cfg(test)]
mod test {
    use Rng;
    use distributions::Distribution;
    use super::Binomial;

    fn test_binomial_mean_and_variance<R: Rng>(n: u64, p: f64, rng: &mut R) {
        let binomial = Binomial::new(n, p);

        let expected_mean = n as f64 * p;
        let expected_variance = n as f64 * p * (1.0 - p);

        let mut results = [0.0; 1000];
        for i in results.iter_mut() { *i = binomial.sample(rng) as f64; }

        let mean = results.iter().sum::<f64>() / results.len() as f64;
        assert!((mean as f64 - expected_mean).abs() < expected_mean / 50.0,
                "mean: {}, expected_mean: {}", mean, expected_mean);

        let variance =
            results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>()
            / results.len() as f64;
        assert!((variance - expected_variance).abs() < expected_variance / 10.0,
                "variance: {}, expected_variance: {}", variance, expected_variance);
    }

    #[test]
    fn test_binomial() {
        let mut rng = ::test::rng(351);
        test_binomial_mean_and_variance(150, 0.1, &mut rng);
        test_binomial_mean_and_variance(70, 0.6, &mut rng);
        test_binomial_mean_and_variance(40, 0.5, &mut rng);
        test_binomial_mean_and_variance(20, 0.7, &mut rng);
        test_binomial_mean_and_variance(20, 0.5, &mut rng);
    }

    #[test]
    fn test_binomial_end_points() {
        let mut rng = ::test::rng(352);
        assert_eq!(rng.sample(Binomial::new(20, 0.0)), 0);
        assert_eq!(rng.sample(Binomial::new(20, 1.0)), 20);
    }

    #[test]
    #[should_panic]
    fn test_binomial_invalid_lambda_neg() {
        Binomial::new(20, -10.0);
    }
}
<|endoftext|>


<commit_before>use std::time::{SystemTime, UNIX_EPOCH};
use std::sync::{Mutex, Arc};
use std::cell::RefCell;
use std::str::FromStr;
use std::path::*;

pub mod block;
pub mod index;
pub mod store;

pub use self::store::BlockStore;
pub use self::index::Indexer;
pub use self::block::*;

use crypto::*;
use asset::*;
use tx::*;

pub struct Blockchain {
    pub indexer: Arc<Indexer>,
    pub store: Mutex<RefCell<BlockStore>>
}

impl Blockchain {

    ///
    /// Creates a new `Blockchain` with an associated indexer and backing
    /// storage is automatically created based on the given `path`.
    ///
    pub fn new(path: &Path) -> Blockchain {
        let indexer = Arc::new(Indexer::new(&Path::join(path, "index")));
        let store = BlockStore::new(&Path::join(path, "blklog"), Arc::clone(&indexer));
        Blockchain {
            indexer,
            store: Mutex::new(RefCell::new(store))
        }
    }

    pub fn get_chain_head(&self) -> Arc<SignedBlock> {
        let lock = self.store.lock().unwrap();
        let store = lock.borrow();
        let height = store.get_chain_height();
        store.get(height).expect("Failed to get blockchain head")
    }

    pub fn get_block(&self, height: u64) -> Option<Arc<SignedBlock>> {
        let lock = self.store.lock().unwrap();
        let store = lock.borrow();
        store.get(height)
    }

    pub fn insert_block(&self, block: SignedBlock) {
        let lock = self.store.lock().unwrap();
        let store = &mut lock.borrow_mut();
        store.insert(block);
    }

    ///
    /// Create and store a generated genesis block. By default the provided
    /// `minter_key` will be the genesis minter and will stake with a small
    /// amount of GOLD tokens.
    ///
    pub fn create_genesis_block(&mut self, minter_key: &KeyPair) {
        use sodiumoxide::crypto::hash::sha256::Digest;

        println!("=> Generating new block chain");
        let staker_key = KeyPair::gen_keypair();
        println!("=> Staker private key: {}", staker_key.1.to_wif());
        println!("=> Staker public key: {}", staker_key.0.to_wif());

        let timestamp = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() as u32;

        let bond_tx = BondTx {
            base: Tx {
                tx_type: TxType::BOND,
                fee: Asset::from_str("0 GOLD").unwrap(),
                timestamp,
                signature_pairs: Vec::new()
            },
            minter: staker_key.0.clone(),
            staker: staker_key.0.clone(),
            bond_fee: EMPTY_GOLD,
            stake_amt: Asset::from_str("1 GOLD").unwrap()
        };

        let transactions = {
            let mut vec = Vec::new();
            vec.push(TxVariant::RewardTx(RewardTx {
                base: Tx {
                    tx_type: TxType::REWARD,
                    fee: Asset::from_str("0 GOLD").unwrap(),
                    timestamp,
                    signature_pairs: Vec::new()
                },
                to: staker_key.0.clone(),
                rewards: vec![Asset::from_str("1 GOLD").unwrap()]
            }));
            vec.push(TxVariant::BondTx(bond_tx.clone()));
            vec
        };

        let block = (Block {
            height: 0,
            previous_hash: Digest::from_slice(&[0u8; 32]).unwrap(),
            tx_merkle_root: Digest::from_slice(&[0u8; 32]).unwrap(),
            timestamp: timestamp as u32,
            transactions
        }).sign(&minter_key);

        let lock = self.store.lock().unwrap();
        let store = &mut lock.borrow_mut();
        store.insert_genesis(block);
        self.indexer.set_bond(&bond_tx);
    }
}
<commit_msg>Fix genesis block using staker key instead of minter key<commit_after>use std::time::{SystemTime, UNIX_EPOCH};
use std::sync::{Mutex, Arc};
use std::cell::RefCell;
use std::str::FromStr;
use std::path::*;

pub mod block;
pub mod index;
pub mod store;

pub use self::store::BlockStore;
pub use self::index::Indexer;
pub use self::block::*;

use crypto::*;
use asset::*;
use tx::*;

pub struct Blockchain {
    pub indexer: Arc<Indexer>,
    pub store: Mutex<RefCell<BlockStore>>
}

impl Blockchain {

    ///
    /// Creates a new `Blockchain` with an associated indexer and backing
    /// storage is automatically created based on the given `path`.
    ///
    pub fn new(path: &Path) -> Blockchain {
        let indexer = Arc::new(Indexer::new(&Path::join(path, "index")));
        let store = BlockStore::new(&Path::join(path, "blklog"), Arc::clone(&indexer));
        Blockchain {
            indexer,
            store: Mutex::new(RefCell::new(store))
        }
    }

    pub fn get_chain_head(&self) -> Arc<SignedBlock> {
        let lock = self.store.lock().unwrap();
        let store = lock.borrow();
        let height = store.get_chain_height();
        store.get(height).expect("Failed to get blockchain head")
    }

    pub fn get_block(&self, height: u64) -> Option<Arc<SignedBlock>> {
        let lock = self.store.lock().unwrap();
        let store = lock.borrow();
        store.get(height)
    }

    pub fn insert_block(&self, block: SignedBlock) {
        let lock = self.store.lock().unwrap();
        let store = &mut lock.borrow_mut();
        store.insert(block);
    }

    ///
    /// Create and store a generated genesis block. By default the provided
    /// `minter_key` will be the genesis minter and will stake with a small
    /// amount of GOLD tokens.
    ///
    pub fn create_genesis_block(&mut self, minter_key: &KeyPair) {
        use sodiumoxide::crypto::hash::sha256::Digest;

        println!("=> Generating new block chain");
        let staker_key = KeyPair::gen_keypair();
        println!("=> Staker private key: {}", staker_key.1.to_wif());
        println!("=> Staker public key: {}", staker_key.0.to_wif());

        let timestamp = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() as u32;

        let bond_tx = BondTx {
            base: Tx {
                tx_type: TxType::BOND,
                fee: Asset::from_str("0 GOLD").unwrap(),
                timestamp,
                signature_pairs: Vec::new()
            },
            minter: minter_key.0.clone(),
            staker: staker_key.0.clone(),
            bond_fee: EMPTY_GOLD,
            stake_amt: Asset::from_str("1 GOLD").unwrap()
        };

        let transactions = {
            let mut vec = Vec::new();
            vec.push(TxVariant::RewardTx(RewardTx {
                base: Tx {
                    tx_type: TxType::REWARD,
                    fee: Asset::from_str("0 GOLD").unwrap(),
                    timestamp,
                    signature_pairs: Vec::new()
                },
                to: staker_key.0.clone(),
                rewards: vec![Asset::from_str("1 GOLD").unwrap()]
            }));
            vec.push(TxVariant::BondTx(bond_tx.clone()));
            vec
        };

        let block = (Block {
            height: 0,
            previous_hash: Digest::from_slice(&[0u8; 32]).unwrap(),
            tx_merkle_root: Digest::from_slice(&[0u8; 32]).unwrap(),
            timestamp: timestamp as u32,
            transactions
        }).sign(&minter_key);

        let lock = self.store.lock().unwrap();
        let store = &mut lock.borrow_mut();
        store.insert_genesis(block);
        self.indexer.set_bond(&bond_tx);
    }
}
<|endoftext|>


<commit_before>//! Dynamic library facilities.
//!
//! A simple wrapper over the platform's dynamic library facilities

use std::ffi::CString;
use std::path::Path;

pub struct DynamicLibrary {
    handle: *mut u8,
}

impl Drop for DynamicLibrary {
    fn drop(&mut self) {
        unsafe { dl::close(self.handle) }
    }
}

impl DynamicLibrary {
    /// Lazily open a dynamic library.
    pub fn open(filename: &Path) -> Result<DynamicLibrary, String> {
        let maybe_library = dl::open(filename.as_os_str());

        // The dynamic library must not be constructed if there is
        // an error opening the library so the destructor does not
        // run.
        match maybe_library {
            Err(err) => Err(err),
            Ok(handle) => Ok(DynamicLibrary { handle }),
        }
    }

    /// Accesses the value at the symbol of the dynamic library.
    pub unsafe fn symbol<T>(&self, symbol: &str) -> Result<*mut T, String> {
        // This function should have a lifetime constraint of 'a on
        // T but that feature is still unimplemented

        let raw_string = CString::new(symbol).unwrap();
        let maybe_symbol_value = dl::symbol(self.handle, raw_string.as_ptr());

        // The value must not be constructed if there is an error so
        // the destructor does not run.
        match maybe_symbol_value {
            Err(err) => Err(err),
            Ok(symbol_value) => Ok(symbol_value as *mut T),
        }
    }
}

#[cfg(test)]
mod tests;

#[cfg(unix)]
mod dl {
    use std::ffi::{CString, OsStr};
    use std::os::unix::prelude::*;

    // As of the 2017 revision of the POSIX standard (IEEE 1003.1-2017), it is
    // implementation-defined whether `dlerror` is thread-safe (in which case it returns the most
    // recent error in the calling thread) or not thread-safe (in which case it returns the most
    // recent error in *any* thread).
    //
    // There's no easy way to tell what strategy is used by a given POSIX implementation, so we
    // lock around all calls that can modify `dlerror` in this module lest we accidentally read an
    // error from a different thread. This is bulletproof when we are the *only* code using the
    // dynamic library APIs at a given point in time. However, it's still possible for us to race
    // with other code (see #74469) on platforms where `dlerror` is not thread-safe.
    mod error {
        use std::ffi::CStr;
        use std::lazy::SyncLazy;
        use std::sync::{Mutex, MutexGuard};

        pub fn lock() -> MutexGuard<'static, Guard> {
            static LOCK: SyncLazy<Mutex<Guard>> = SyncLazy::new(|| Mutex::new(Guard { _priv: () }));
            LOCK.lock().unwrap()
        }

        pub struct Guard {
            _priv: (),
        }

        impl Guard {
            pub fn get(&mut self) -> Result<(), String> {
                let msg = unsafe { libc::dlerror() };
                if msg.is_null() {
                    Ok(())
                } else {
                    let msg = unsafe { CStr::from_ptr(msg as *const _) };
                    Err(msg.to_string_lossy().into_owned())
                }
            }

            pub fn clear(&mut self) {
                let _ = unsafe { libc::dlerror() };
            }
        }
    }

    pub(super) fn open(filename: &OsStr) -> Result<*mut u8, String> {
        let s = CString::new(filename.as_bytes()).unwrap();

        let mut dlerror = error::lock();
        let ret = unsafe { libc::dlopen(s.as_ptr(), libc::RTLD_LAZY) } as *mut u8;

        if !ret.is_null() {
            return Ok(ret);
        }

        // A NULL return from `dlopen` indicates that an error has definitely occurred, so if
        // nothing is in `dlerror`, we are racing with another thread that has stolen our error
        // message. See the explanation on the `dl::error` module for more information.
        dlerror.get().and_then(|()| Err("Unknown error".to_string()))
    }

    pub(super) unsafe fn symbol(
        handle: *mut u8,
        symbol: *const libc::c_char,
    ) -> Result<*mut u8, String> {
        let mut dlerror = error::lock();

        // Unlike `dlopen`, it's possible for `dlsym` to return NULL without overwriting `dlerror`.
        // Because of this, we clear `dlerror` before calling `dlsym` to avoid picking up a stale
        // error message by accident.
        dlerror.clear();

        let ret = libc::dlsym(handle as *mut libc::c_void, symbol) as *mut u8;

        if !ret.is_null() {
            return Ok(ret);
        }

        // If `dlsym` returns NULL but there is nothing in `dlerror` it means one of two things:
        // - We tried to load a symbol mapped to address 0. This is not technically an error but is
        //   unlikely to occur in practice and equally unlikely to be handled correctly by calling
        //   code. Therefore we treat it as an error anyway.
        // - An error has occurred, but we are racing with another thread that has stolen our error
        //   message. See the explanation on the `dl::error` module for more information.
        dlerror.get().and_then(|()| Err("Tried to load symbol mapped to address 0".to_string()))
    }

    pub(super) unsafe fn close(handle: *mut u8) {
        libc::dlclose(handle as *mut libc::c_void);
    }
}

#[cfg(windows)]
mod dl {
    use std::ffi::OsStr;
    use std::io;
    use std::os::windows::prelude::*;
    use std::ptr;

    use winapi::shared::minwindef::HMODULE;
    use winapi::um::errhandlingapi::SetThreadErrorMode;
    use winapi::um::libloaderapi::{FreeLibrary, GetProcAddress, LoadLibraryW};
    use winapi::um::winbase::SEM_FAILCRITICALERRORS;

    pub(super) fn open(filename: &OsStr) -> Result<*mut u8, String> {
        // disable "dll load failed" error dialog.
        let prev_error_mode = unsafe {
            let new_error_mode = SEM_FAILCRITICALERRORS;
            let mut prev_error_mode = 0;
            let result = SetThreadErrorMode(new_error_mode, &mut prev_error_mode);
            if result == 0 {
                return Err(io::Error::last_os_error().to_string());
            }
            prev_error_mode
        };

        let filename_str: Vec<_> = filename.encode_wide().chain(Some(0)).collect();
        let result = unsafe { LoadLibraryW(filename_str.as_ptr()) } as *mut u8;
        let result = ptr_result(result);

        unsafe {
            SetThreadErrorMode(prev_error_mode, ptr::null_mut());
        }

        result
    }

    pub(super) unsafe fn symbol(
        handle: *mut u8,
        symbol: *const libc::c_char,
    ) -> Result<*mut u8, String> {
        let ptr = GetProcAddress(handle as HMODULE, symbol) as *mut u8;
        ptr_result(ptr)
    }

    pub(super) unsafe fn close(handle: *mut u8) {
        FreeLibrary(handle as HMODULE);
    }

    fn ptr_result<T>(ptr: *mut T) -> Result<*mut T, String> {
        if ptr.is_null() { Err(io::Error::last_os_error().to_string()) } else { Ok(ptr) }
    }
}
<commit_msg>Explicitly pass `RTLD_LOCAL` to `dlopen`<commit_after>//! Dynamic library facilities.
//!
//! A simple wrapper over the platform's dynamic library facilities

use std::ffi::CString;
use std::path::Path;

pub struct DynamicLibrary {
    handle: *mut u8,
}

impl Drop for DynamicLibrary {
    fn drop(&mut self) {
        unsafe { dl::close(self.handle) }
    }
}

impl DynamicLibrary {
    /// Lazily open a dynamic library.
    pub fn open(filename: &Path) -> Result<DynamicLibrary, String> {
        let maybe_library = dl::open(filename.as_os_str());

        // The dynamic library must not be constructed if there is
        // an error opening the library so the destructor does not
        // run.
        match maybe_library {
            Err(err) => Err(err),
            Ok(handle) => Ok(DynamicLibrary { handle }),
        }
    }

    /// Accesses the value at the symbol of the dynamic library.
    pub unsafe fn symbol<T>(&self, symbol: &str) -> Result<*mut T, String> {
        // This function should have a lifetime constraint of 'a on
        // T but that feature is still unimplemented

        let raw_string = CString::new(symbol).unwrap();
        let maybe_symbol_value = dl::symbol(self.handle, raw_string.as_ptr());

        // The value must not be constructed if there is an error so
        // the destructor does not run.
        match maybe_symbol_value {
            Err(err) => Err(err),
            Ok(symbol_value) => Ok(symbol_value as *mut T),
        }
    }
}

#[cfg(test)]
mod tests;

#[cfg(unix)]
mod dl {
    use std::ffi::{CString, OsStr};
    use std::os::unix::prelude::*;

    // As of the 2017 revision of the POSIX standard (IEEE 1003.1-2017), it is
    // implementation-defined whether `dlerror` is thread-safe (in which case it returns the most
    // recent error in the calling thread) or not thread-safe (in which case it returns the most
    // recent error in *any* thread).
    //
    // There's no easy way to tell what strategy is used by a given POSIX implementation, so we
    // lock around all calls that can modify `dlerror` in this module lest we accidentally read an
    // error from a different thread. This is bulletproof when we are the *only* code using the
    // dynamic library APIs at a given point in time. However, it's still possible for us to race
    // with other code (see #74469) on platforms where `dlerror` is not thread-safe.
    mod error {
        use std::ffi::CStr;
        use std::lazy::SyncLazy;
        use std::sync::{Mutex, MutexGuard};

        pub fn lock() -> MutexGuard<'static, Guard> {
            static LOCK: SyncLazy<Mutex<Guard>> = SyncLazy::new(|| Mutex::new(Guard { _priv: () }));
            LOCK.lock().unwrap()
        }

        pub struct Guard {
            _priv: (),
        }

        impl Guard {
            pub fn get(&mut self) -> Result<(), String> {
                let msg = unsafe { libc::dlerror() };
                if msg.is_null() {
                    Ok(())
                } else {
                    let msg = unsafe { CStr::from_ptr(msg as *const _) };
                    Err(msg.to_string_lossy().into_owned())
                }
            }

            pub fn clear(&mut self) {
                let _ = unsafe { libc::dlerror() };
            }
        }
    }

    pub(super) fn open(filename: &OsStr) -> Result<*mut u8, String> {
        let s = CString::new(filename.as_bytes()).unwrap();

        let mut dlerror = error::lock();
        let ret = unsafe { libc::dlopen(s.as_ptr(), libc::RTLD_LAZY | libc::RTLD_LOCAL) };

        if !ret.is_null() {
            return Ok(ret.cast());
        }

        // A NULL return from `dlopen` indicates that an error has definitely occurred, so if
        // nothing is in `dlerror`, we are racing with another thread that has stolen our error
        // message. See the explanation on the `dl::error` module for more information.
        dlerror.get().and_then(|()| Err("Unknown error".to_string()))
    }

    pub(super) unsafe fn symbol(
        handle: *mut u8,
        symbol: *const libc::c_char,
    ) -> Result<*mut u8, String> {
        let mut dlerror = error::lock();

        // Unlike `dlopen`, it's possible for `dlsym` to return NULL without overwriting `dlerror`.
        // Because of this, we clear `dlerror` before calling `dlsym` to avoid picking up a stale
        // error message by accident.
        dlerror.clear();

        let ret = libc::dlsym(handle as *mut libc::c_void, symbol);

        if !ret.is_null() {
            return Ok(ret.cast());
        }

        // If `dlsym` returns NULL but there is nothing in `dlerror` it means one of two things:
        // - We tried to load a symbol mapped to address 0. This is not technically an error but is
        //   unlikely to occur in practice and equally unlikely to be handled correctly by calling
        //   code. Therefore we treat it as an error anyway.
        // - An error has occurred, but we are racing with another thread that has stolen our error
        //   message. See the explanation on the `dl::error` module for more information.
        dlerror.get().and_then(|()| Err("Tried to load symbol mapped to address 0".to_string()))
    }

    pub(super) unsafe fn close(handle: *mut u8) {
        libc::dlclose(handle as *mut libc::c_void);
    }
}

#[cfg(windows)]
mod dl {
    use std::ffi::OsStr;
    use std::io;
    use std::os::windows::prelude::*;
    use std::ptr;

    use winapi::shared::minwindef::HMODULE;
    use winapi::um::errhandlingapi::SetThreadErrorMode;
    use winapi::um::libloaderapi::{FreeLibrary, GetProcAddress, LoadLibraryW};
    use winapi::um::winbase::SEM_FAILCRITICALERRORS;

    pub(super) fn open(filename: &OsStr) -> Result<*mut u8, String> {
        // disable "dll load failed" error dialog.
        let prev_error_mode = unsafe {
            let new_error_mode = SEM_FAILCRITICALERRORS;
            let mut prev_error_mode = 0;
            let result = SetThreadErrorMode(new_error_mode, &mut prev_error_mode);
            if result == 0 {
                return Err(io::Error::last_os_error().to_string());
            }
            prev_error_mode
        };

        let filename_str: Vec<_> = filename.encode_wide().chain(Some(0)).collect();
        let result = unsafe { LoadLibraryW(filename_str.as_ptr()) } as *mut u8;
        let result = ptr_result(result);

        unsafe {
            SetThreadErrorMode(prev_error_mode, ptr::null_mut());
        }

        result
    }

    pub(super) unsafe fn symbol(
        handle: *mut u8,
        symbol: *const libc::c_char,
    ) -> Result<*mut u8, String> {
        let ptr = GetProcAddress(handle as HMODULE, symbol) as *mut u8;
        ptr_result(ptr)
    }

    pub(super) unsafe fn close(handle: *mut u8) {
        FreeLibrary(handle as HMODULE);
    }

    fn ptr_result<T>(ptr: *mut T) -> Result<*mut T, String> {
        if ptr.is_null() { Err(io::Error::last_os_error().to_string()) } else { Ok(ptr) }
    }
}
<|endoftext|>


<commit_before>extern crate proc_macro;
extern crate proc_macro2;
#[macro_use]
extern crate syn;
#[macro_use]
extern crate quote;

use proc_macro::TokenStream;
use syn::{ItemFn, Stmt, Item, Block, Expr,
          spanned::Spanned,
          fold::Fold};

/// Parses a list of variable names separated by commas
///
/// This is how the compiler passes in arguments to our attribute -- it is
/// everything inside the delimiters after the attribute name.
///```rust,ignore
///     #[wlroots_dehandle(a, b, c)]
///```
struct Args;

impl Fold for Args {
    fn fold_block(&mut self, block: Block) -> Block {
        build_block(block.stmts.iter(), self)
    }
}

/// Attribute to automatically call the `run` method on handles with the
/// remaining block of code.
///
/// The name of the variable you want to use as the upgraded handle should be
/// provided as an argument to the attribute. It does not need to be the same
/// as the handle variable.
///
/// The syntax in the code should be `use $handle as $upgraded_handle`.
/// E.g the variable in the code that stores the handle should go on the
/// **left** and the variable you used in the attribute declaration should
/// go on the **right**.
///
/// # Panics
/// If the handle is invalid (e.g. default constructed, or is a dangling
/// handle) then your code will `panic!`.
///
/// If this is undesirable, please use the non-proc macro `with_handles!`.
///
/// # Example
///
/// ```rust,ignore
/// impl InputManagerHandler for InputManager {
///     #[wlroots_dehandle]
///     fn keyboard_added(&mut self,
///                       compositor_handle: CompositorHandle,
///                       keyboard: KeyboardHandle)
///                       -> Option<Box<Keyboard Handler>> {
///         {
///             #[dehandle] let compositor = compositor_handle;
///             #[dehandle] let keyboard = keyboard;
///             let server: &mut ::Server = compositor.into();
///             server.keyboards.push(keyboard.weak_reference());
///             // Now that we have at least one keyboard, update the seat capabilities.
///             #[dehandle] let seat = &server.seat.seat;
///             let mut capabilities = seat.capabilities();
///             capabilities.insert(Capability::Keyboard);
///             seat.set_capabilities(capabilities);
///             seat.set_keyboard(keyboard.input_device());
///         }
///         // Due to some weird closure inference rules, this has to be outside
///         // of the above block.
///         Some(Box::new(::Keyboard))
///     }
/// }
/// ```
#[proc_macro_attribute]
pub fn wlroots_dehandle(_args: TokenStream, input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ItemFn);
    let output = Args.fold_item_fn(input);
    TokenStream::from(quote!(#output))
}

fn build_block(mut input: std::slice::Iter<Stmt>, args: &mut Args) -> Block {
    let mut output = vec![];
    let mut inner = None;
    let mut is_try = false;
    while let Some(stmt) = input.next().cloned() {
        use syn::{Pat, punctuated::Pair};
        match stmt.clone() {
            // Recurse into function body
            Stmt::Item(Item::Fn(mut function)) => {
                let inner_block = function.block.clone();
                *function.block = build_block(inner_block.stmts.iter(), args);
                output.push(Stmt::Item(Item::Fn(function)))
            },
            Stmt::Local(mut local) => {
                // Ensure attribute is prefaced here
                let mut dehandle = false;
                for attribute in &local.attrs {
                    let meta = attribute.parse_meta();
                    match meta {
                        Ok(syn::Meta::Word(name)) => {
                            if name.to_string() == "dehandle" {
                                dehandle = true;
                                break;
                            }
                        },
                        _ => {}
                    }
                };
                let left_side = local.pats.first().map(Pair::into_value).cloned();
                let right_side = local.init.clone();
                match (dehandle, left_side, right_side) {
                    (true,
                     Some(Pat::Ident(dehandle_name)),
                     Some((_, body))) => {
                        let mut body = *body;
                        is_try = match body.clone() {
                            syn::Expr::Try(syn::ExprTry { expr, .. }) => {
                                body = *expr.clone();
                                true
                            },
                            _ => false
                        };
                        inner = Some((body, dehandle_name));
                        break;
                    },
                    // Recurse into let call
                    (false, _, Some((_, body))) => {
                        let body = build_block_expr(*body.clone(), args);
                        let stream = quote_spanned!(stmt.span()=> #body);
                        let body: Expr = syn::parse_quote::parse(stream.into());
                        local.init.as_mut().unwrap().1 = Box::new(body);
                        output.push(Stmt::Local(local))
                    },
                    _ => output.push(Stmt::Local(local))
                }
            },
            Stmt::Expr(expr) => {
                let body = build_block_expr(expr, args);
                output.push(syn::parse_quote::parse(quote_spanned!(stmt.span()=> {#body}).into()))
            }
            Stmt::Semi(expr, _) => {
                let body = build_block_expr(expr, args);
                output.push(syn::parse_quote::parse(quote_spanned!(stmt.span()=> {#body;}).into()))
            }
            _ => output.push(stmt)
        }
    }
    if let Some((handle, dehandle)) = inner {
        let inner_block = build_block(input, args);
        let handle_call = if !is_try {
            syn::parse_quote::parse(
                quote_spanned!(handle.span()=>
                               {(#handle).run(|#dehandle|{
                                   #inner_block
                               }).expect(concat!("Could not upgrade handle ",
                                                 stringify!(#handle), " to ",
                                                 stringify!(#dehandle)))}).into())
        } else {
            syn::parse_quote::parse(
                quote_spanned!(handle.span()=>
                               {(#handle).run(|#dehandle|{
                                   #inner_block
                               })?}).into())
        };
        output.push(handle_call);
    }
    parse_quote!({#(#output)*})
}

/// Tries to build a block from the expression.
fn build_block_expr(expr: Expr, args: &mut Args) -> Expr {
    match expr {
        Expr::Block(block) => {
            let block = build_block(block.block.stmts.iter(), args);
            syn::parse_quote::parse(quote_spanned!(block.span()=> #block))
        }
        Expr::Let(mut let_expr) => {
            *let_expr.expr = build_block_expr(*let_expr.expr.clone(), args);
            Expr::Let(let_expr)
        },
        Expr::If(mut if_expr) => {
            let then_branch = if_expr.then_branch.clone();
            let then_parsed = syn::parse_quote::parse(
                quote_spanned!(then_branch.span()=> #then_branch));
            let then_branch = build_block_expr(then_parsed, args);
            if_expr.then_branch = syn::parse_quote::parse(
                quote_spanned!(then_branch.span()=> #then_branch));
            if_expr.else_branch = match if_expr.else_branch.clone() {
                None => if_expr.else_branch,
                Some((token, else_branch)) => {
                    Some((token, Box::new( build_block_expr(*else_branch, args))))
                }
            };
            Expr::If(if_expr)
        },
        Expr::While(mut while_expr) => {
            let body = while_expr.body.clone();
            let body = build_block_expr(syn::parse_quote::parse(
                quote_spanned!(body.span()=> #body)), args);
            while_expr.body = parse_quote!(#body);
            Expr::While(while_expr)
        },
        Expr::ForLoop(mut for_expr) => {
            let body = for_expr.body.clone();
            let body = build_block_expr(parse_quote!(#body), args);
            for_expr.body = parse_quote!(#body);
            Expr::ForLoop(for_expr)
        },
        Expr::Loop(mut loop_expr) => {
            let body = loop_expr.body.clone();
            let body = build_block_expr(parse_quote!(#body), args);
            loop_expr.body = parse_quote!(#body);
            Expr::Loop(loop_expr)
        },
        Expr::Match(mut match_expr) => {
            for arm in &mut match_expr.arms {
                *arm.body = build_block_expr(*arm.body.clone(), args)
            }
            Expr::Match(match_expr)
        },
        Expr::Struct(mut struct_expr) => {
            for field in &mut struct_expr.fields {
                field.expr = build_block_expr(field.expr.clone(), args);
            }
            Expr::Struct(struct_expr)
        },
        Expr::Call(mut call_expr) => {
            for arg in &mut call_expr.args {
                *arg = build_block_expr(arg.clone(), args);
            }
            Expr::Call(call_expr)
        },
        Expr::MethodCall(mut call_expr) => {
            for arg in &mut call_expr.args {
                *arg = build_block_expr(arg.clone(), args);
            }
            Expr::MethodCall(call_expr)
        },
        Expr::Closure(mut closure_expr) => {
            *closure_expr.body = build_block_expr(*closure_expr.body.clone(),
                                                  args);
            Expr::Closure(closure_expr)
        },
        Expr::Unsafe(mut unsafe_expr) => {
            unsafe_expr.block = build_block(unsafe_expr.block.stmts.iter(),
                                            args);
            Expr::Unsafe(unsafe_expr)
        },
        Expr::Assign(mut assign_expr) => {
            *assign_expr.right = build_block_expr(*assign_expr.right.clone(),
                                                  args);
            Expr::Assign(assign_expr)
        },
        Expr::AssignOp(mut assign_expr) => {
            *assign_expr.right = build_block_expr(*assign_expr.right.clone(),
                                                  args);
            Expr::AssignOp(assign_expr)
        },
        Expr::Break(mut break_expr) => {
            match break_expr.expr {
                None => {},
                Some(ref mut expr) => {
                    **expr = build_block_expr(*expr.clone(), args);
                }
            }
            Expr::Break(break_expr)
        },
        Expr::Return(mut return_expr) => {
            match return_expr.expr {
                None => {},
                Some(ref mut expr) => {
                    **expr = build_block_expr(*expr.clone(), args);
                }
            }
            Expr::Return(return_expr)
        },
        Expr::Reference(mut reference_expr) => {
            *reference_expr.expr = build_block_expr(*reference_expr.expr.clone(),
                                                    args);
            Expr::Reference(reference_expr)
        },
        Expr::Paren(mut paren_expr) => {
            *paren_expr.expr = build_block_expr(*paren_expr.expr.clone(), args);
            Expr::Paren(paren_expr)
        },
        Expr::Unary(mut unary_expr) => {
            *unary_expr.expr = build_block_expr(*unary_expr.expr.clone(), args);
            Expr::Unary(unary_expr)
        },
        Expr::Binary(mut binary_expr) => {
            *binary_expr.left = build_block_expr(*binary_expr.left.clone(), args);
            *binary_expr.right = build_block_expr(*binary_expr.right.clone(), args);
            Expr::Binary(binary_expr)
        },
        v => {
            v
        }
    }
}
<commit_msg>Fixed up the documentation for wlroots-dehandle<commit_after>extern crate proc_macro;
extern crate proc_macro2;
#[macro_use]
extern crate syn;
#[macro_use]
extern crate quote;

use proc_macro::TokenStream;
use syn::{ItemFn, Stmt, Item, Block, Expr,
          spanned::Spanned,
          fold::Fold};

/// Parses a list of variable names separated by commas
///
/// This is how the compiler passes in arguments to our attribute -- it is
/// everything inside the delimiters after the attribute name.
///```rust,ignore
///     #[wlroots_dehandle(a, b, c)]
///```
struct Args;

impl Fold for Args {
    fn fold_block(&mut self, block: Block) -> Block {
        build_block(block.stmts.iter(), self)
    }
}

/// Attribute to automatically call the `run` method on handles with the
/// remaining block of code.
///
/// The syntax in the code to denote a handle being accessed is
/// `#[dehandle] let $upgraded_handle = $handle`.
///
///
/// # Panics
/// If the handle is invalid (e.g. default constructed, borrowed multiple times,
/// or it is a dangling handle) then your code will `panic!`.
///
/// If this is undesirable, please append `?` to the end like so:
/// `#[dehandle] let $upgraded_handle = $handle?`. This will make it behave
/// as you expect (i.e. the `Err` is returned early and the return type of the
/// function should be `HandleResult<T>`)
///
/// # Example
///
/// ```rust,ignore
/// impl InputManagerHandler for InputManager {
///     #[wlroots_dehandle]
///     fn keyboard_added(&mut self,
///                       compositor_handle: CompositorHandle,
///                       keyboard: KeyboardHandle)
///                       -> Option<Box<Keyboard Handler>> {
///         {
///             #[dehandle] let compositor = compositor_handle;
///             #[dehandle] let keyboard = keyboard;
///             let server: &mut ::Server = compositor.into();
///             server.keyboards.push(keyboard.weak_reference());
///             // Now that we have at least one keyboard, update the seat capabilities.
///             #[dehandle] let seat = &server.seat.seat;
///             let mut capabilities = seat.capabilities();
///             capabilities.insert(Capability::Keyboard);
///             seat.set_capabilities(capabilities);
///             seat.set_keyboard(keyboard.input_device());
///         }
///         // Due to some weird closure inference rules, this has to be outside
///         // of the above block.
///         Some(Box::new(::Keyboard))
///     }
/// }
/// ```
#[proc_macro_attribute]
pub fn wlroots_dehandle(_args: TokenStream, input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ItemFn);
    let output = Args.fold_item_fn(input);
    TokenStream::from(quote!(#output))
}

fn build_block(mut input: std::slice::Iter<Stmt>, args: &mut Args) -> Block {
    let mut output = vec![];
    let mut inner = None;
    let mut is_try = false;
    while let Some(stmt) = input.next().cloned() {
        use syn::{Pat, punctuated::Pair};
        match stmt.clone() {
            // Recurse into function body
            Stmt::Item(Item::Fn(mut function)) => {
                let inner_block = function.block.clone();
                *function.block = build_block(inner_block.stmts.iter(), args);
                output.push(Stmt::Item(Item::Fn(function)))
            },
            Stmt::Local(mut local) => {
                // Ensure attribute is prefaced here
                let mut dehandle = false;
                for attribute in &local.attrs {
                    let meta = attribute.parse_meta();
                    match meta {
                        Ok(syn::Meta::Word(name)) => {
                            if name.to_string() == "dehandle" {
                                dehandle = true;
                                break;
                            }
                        },
                        _ => {}
                    }
                };
                let left_side = local.pats.first().map(Pair::into_value).cloned();
                let right_side = local.init.clone();
                match (dehandle, left_side, right_side) {
                    (true,
                     Some(Pat::Ident(dehandle_name)),
                     Some((_, body))) => {
                        let mut body = *body;
                        is_try = match body.clone() {
                            syn::Expr::Try(syn::ExprTry { expr, .. }) => {
                                body = *expr.clone();
                                true
                            },
                            _ => false
                        };
                        inner = Some((body, dehandle_name));
                        break;
                    },
                    // Recurse into let call
                    (false, _, Some((_, body))) => {
                        let body = build_block_expr(*body.clone(), args);
                        let stream = quote_spanned!(stmt.span()=> #body);
                        let body: Expr = syn::parse_quote::parse(stream.into());
                        local.init.as_mut().unwrap().1 = Box::new(body);
                        output.push(Stmt::Local(local))
                    },
                    _ => output.push(Stmt::Local(local))
                }
            },
            Stmt::Expr(expr) => {
                let body = build_block_expr(expr, args);
                output.push(syn::parse_quote::parse(quote_spanned!(stmt.span()=> {#body}).into()))
            }
            Stmt::Semi(expr, _) => {
                let body = build_block_expr(expr, args);
                output.push(syn::parse_quote::parse(quote_spanned!(stmt.span()=> {#body;}).into()))
            }
            _ => output.push(stmt)
        }
    }
    if let Some((handle, dehandle)) = inner {
        let inner_block = build_block(input, args);
        let handle_call = if !is_try {
            syn::parse_quote::parse(
                quote_spanned!(handle.span()=>
                               {(#handle).run(|#dehandle|{
                                   #inner_block
                               }).expect(concat!("Could not upgrade handle ",
                                                 stringify!(#handle), " to ",
                                                 stringify!(#dehandle)))}).into())
        } else {
            syn::parse_quote::parse(
                quote_spanned!(handle.span()=>
                               {(#handle).run(|#dehandle|{
                                   #inner_block
                               })?}).into())
        };
        output.push(handle_call);
    }
    parse_quote!({#(#output)*})
}

/// Tries to build a block from the expression.
fn build_block_expr(expr: Expr, args: &mut Args) -> Expr {
    match expr {
        Expr::Block(block) => {
            let block = build_block(block.block.stmts.iter(), args);
            syn::parse_quote::parse(quote_spanned!(block.span()=> #block))
        }
        Expr::Let(mut let_expr) => {
            *let_expr.expr = build_block_expr(*let_expr.expr.clone(), args);
            Expr::Let(let_expr)
        },
        Expr::If(mut if_expr) => {
            let then_branch = if_expr.then_branch.clone();
            let then_parsed = syn::parse_quote::parse(
                quote_spanned!(then_branch.span()=> #then_branch));
            let then_branch = build_block_expr(then_parsed, args);
            if_expr.then_branch = syn::parse_quote::parse(
                quote_spanned!(then_branch.span()=> #then_branch));
            if_expr.else_branch = match if_expr.else_branch.clone() {
                None => if_expr.else_branch,
                Some((token, else_branch)) => {
                    Some((token, Box::new( build_block_expr(*else_branch, args))))
                }
            };
            Expr::If(if_expr)
        },
        Expr::While(mut while_expr) => {
            let body = while_expr.body.clone();
            let body = build_block_expr(syn::parse_quote::parse(
                quote_spanned!(body.span()=> #body)), args);
            while_expr.body = parse_quote!(#body);
            Expr::While(while_expr)
        },
        Expr::ForLoop(mut for_expr) => {
            let body = for_expr.body.clone();
            let body = build_block_expr(parse_quote!(#body), args);
            for_expr.body = parse_quote!(#body);
            Expr::ForLoop(for_expr)
        },
        Expr::Loop(mut loop_expr) => {
            let body = loop_expr.body.clone();
            let body = build_block_expr(parse_quote!(#body), args);
            loop_expr.body = parse_quote!(#body);
            Expr::Loop(loop_expr)
        },
        Expr::Match(mut match_expr) => {
            for arm in &mut match_expr.arms {
                *arm.body = build_block_expr(*arm.body.clone(), args)
            }
            Expr::Match(match_expr)
        },
        Expr::Struct(mut struct_expr) => {
            for field in &mut struct_expr.fields {
                field.expr = build_block_expr(field.expr.clone(), args);
            }
            Expr::Struct(struct_expr)
        },
        Expr::Call(mut call_expr) => {
            for arg in &mut call_expr.args {
                *arg = build_block_expr(arg.clone(), args);
            }
            Expr::Call(call_expr)
        },
        Expr::MethodCall(mut call_expr) => {
            for arg in &mut call_expr.args {
                *arg = build_block_expr(arg.clone(), args);
            }
            Expr::MethodCall(call_expr)
        },
        Expr::Closure(mut closure_expr) => {
            *closure_expr.body = build_block_expr(*closure_expr.body.clone(),
                                                  args);
            Expr::Closure(closure_expr)
        },
        Expr::Unsafe(mut unsafe_expr) => {
            unsafe_expr.block = build_block(unsafe_expr.block.stmts.iter(),
                                            args);
            Expr::Unsafe(unsafe_expr)
        },
        Expr::Assign(mut assign_expr) => {
            *assign_expr.right = build_block_expr(*assign_expr.right.clone(),
                                                  args);
            Expr::Assign(assign_expr)
        },
        Expr::AssignOp(mut assign_expr) => {
            *assign_expr.right = build_block_expr(*assign_expr.right.clone(),
                                                  args);
            Expr::AssignOp(assign_expr)
        },
        Expr::Break(mut break_expr) => {
            match break_expr.expr {
                None => {},
                Some(ref mut expr) => {
                    **expr = build_block_expr(*expr.clone(), args);
                }
            }
            Expr::Break(break_expr)
        },
        Expr::Return(mut return_expr) => {
            match return_expr.expr {
                None => {},
                Some(ref mut expr) => {
                    **expr = build_block_expr(*expr.clone(), args);
                }
            }
            Expr::Return(return_expr)
        },
        Expr::Reference(mut reference_expr) => {
            *reference_expr.expr = build_block_expr(*reference_expr.expr.clone(),
                                                    args);
            Expr::Reference(reference_expr)
        },
        Expr::Paren(mut paren_expr) => {
            *paren_expr.expr = build_block_expr(*paren_expr.expr.clone(), args);
            Expr::Paren(paren_expr)
        },
        Expr::Unary(mut unary_expr) => {
            *unary_expr.expr = build_block_expr(*unary_expr.expr.clone(), args);
            Expr::Unary(unary_expr)
        },
        Expr::Binary(mut binary_expr) => {
            *binary_expr.left = build_block_expr(*binary_expr.left.clone(), args);
            *binary_expr.right = build_block_expr(*binary_expr.right.clone(), args);
            Expr::Binary(binary_expr)
        },
        v => {
            v
        }
    }
}
<|endoftext|>