<commit_before>use crate::cell::UnsafeCell;
use crate::mem;
use crate::mem::MaybeUninit;
use crate::sync::atomic::{AtomicU32, Ordering};
use crate::sys::cloudabi::abi;

extern "C" {
    #[thread_local]
    static __pthread_thread_id: abi::tid;
}

#[thread_local]
static mut RDLOCKS_ACQUIRED: u32 = 0;

pub struct RWLock {
    lock: UnsafeCell<AtomicU32>,
}

pub unsafe fn raw(r: &RWLock) -> *mut AtomicU32 {
    r.lock.get()
}

unsafe impl Send for RWLock {}
unsafe impl Sync for RWLock {}

const NEW: RWLock = RWLock { lock: UnsafeCell::new(AtomicU32::new(abi::LOCK_UNLOCKED.0)) };

impl RWLock {
    pub const fn new() -> RWLock {
        NEW
    }

    pub unsafe fn try_read(&self) -> bool {
        let lock = self.lock.get();
        let mut old = abi::LOCK_UNLOCKED.0;
        while let Err(cur) =
            (*lock).compare_exchange_weak(old, old + 1, Ordering::Acquire, Ordering::Relaxed)
        {
            if (cur & abi::LOCK_WRLOCKED.0) != 0 {
                // Another thread already has a write lock.
                assert_ne!(
                    old & !abi::LOCK_KERNEL_MANAGED.0,
                    __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
                    "Attempted to acquire a read lock while holding a write lock"
                );
                return false;
            } else if (old & abi::LOCK_KERNEL_MANAGED.0) != 0 && RDLOCKS_ACQUIRED == 0 {
                // Lock has threads waiting for the lock. Only acquire
                // the lock if we have already acquired read locks. In
                // that case, it is justified to acquire this lock to
                // prevent a deadlock.
                return false;
            }
            old = cur;
        }

        RDLOCKS_ACQUIRED += 1;
        true
    }

    pub unsafe fn read(&self) {
        if !self.try_read() {
            // Call into the kernel to acquire a read lock.
            let lock = self.lock.get();
            let subscription = abi::subscription {
                type_: abi::eventtype::LOCK_RDLOCK,
                union: abi::subscription_union {
                    lock: abi::subscription_lock {
                        lock: lock as *mut abi::lock,
                        lock_scope: abi::scope::PRIVATE,
                    },
                },
                ..mem::zeroed()
            };
            let mut event = MaybeUninit::<abi::event>::uninit();
            let mut nevents = MaybeUninit::<usize>::uninit();
            let ret = abi::poll(&subscription, event.as_mut_ptr(), 1, nevents.as_mut_ptr());
            assert_eq!(ret, abi::errno::SUCCESS, "Failed to acquire read lock");
            let event = event.assume_init();
            assert_eq!(event.error, abi::errno::SUCCESS, "Failed to acquire read lock");

            RDLOCKS_ACQUIRED += 1;
        }
    }

    pub unsafe fn read_unlock(&self) {
        // Perform a read unlock. We can do this in userspace, except when
        // other threads are blocked and we are performing the last unlock.
        // In that case, call into the kernel.
        //
        // Other threads may attempt to increment the read lock count,
        // meaning that the call into the kernel could be spurious. To
        // prevent this from happening, upgrade to a write lock first. This
        // allows us to call into the kernel, having the guarantee that the
        // lock value will not change in the meantime.
        assert!(RDLOCKS_ACQUIRED > 0, "Bad lock count");
        let mut old = 1;
        loop {
            let lock = self.lock.get();
            if old == 1 | abi::LOCK_KERNEL_MANAGED.0 {
                // Last read lock while threads are waiting. Attempt to upgrade
                // to a write lock before calling into the kernel to unlock.
                if let Err(cur) = (*lock).compare_exchange_weak(
                    old,
                    __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0 | abi::LOCK_KERNEL_MANAGED.0,
                    Ordering::Acquire,
                    Ordering::Relaxed,
                ) {
                    old = cur;
                } else {
                    // Call into the kernel to unlock.
                    let ret = abi::lock_unlock(lock as *mut abi::lock, abi::scope::PRIVATE);
                    assert_eq!(ret, abi::errno::SUCCESS, "Failed to write unlock a rwlock");
                    break;
                }
            } else {
                // No threads waiting or not the last read lock. Just decrement
                // the read lock count.
                assert_ne!(old & !abi::LOCK_KERNEL_MANAGED.0, 0, "This rwlock is not locked");
                assert_eq!(
                    old & abi::LOCK_WRLOCKED.0,
                    0,
                    "Attempted to read-unlock a write-locked rwlock"
                );
                if let Err(cur) = (*lock).compare_exchange_weak(
                    old,
                    old - 1,
                    Ordering::Acquire,
                    Ordering::Relaxed,
                ) {
                    old = cur;
                } else {
                    break;
                }
            }
        }

        RDLOCKS_ACQUIRED -= 1;
    }

    pub unsafe fn try_write(&self) -> bool {
        // Attempt to acquire the lock.
        let lock = self.lock.get();
        if let Err(old) = (*lock).compare_exchange(
            abi::LOCK_UNLOCKED.0,
            __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
            Ordering::Acquire,
            Ordering::Relaxed,
        ) {
            // Failure. Crash upon recursive acquisition.
            assert_ne!(
                old & !abi::LOCK_KERNEL_MANAGED.0,
                __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
                "Attempted to recursive write-lock a rwlock",
            );
            false
        } else {
            // Success.
            true
        }
    }

    pub unsafe fn write(&self) {
        if !self.try_write() {
            // Call into the kernel to acquire a write lock.
            let lock = self.lock.get();
            let subscription = abi::subscription {
                type_: abi::eventtype::LOCK_WRLOCK,
                union: abi::subscription_union {
                    lock: abi::subscription_lock {
                        lock: lock as *mut abi::lock,
                        lock_scope: abi::scope::PRIVATE,
                    },
                },
                ..mem::zeroed()
            };
            let mut event = MaybeUninit::<abi::event>::uninit();
            let mut nevents = MaybeUninit::<usize>::uninit();
            let ret = abi::poll(&subscription, event.as_mut_ptr(), 1, nevents.as_mut_ptr());
            assert_eq!(ret, abi::errno::SUCCESS, "Failed to acquire write lock");
            let event = event.assume_init();
            assert_eq!(event.error, abi::errno::SUCCESS, "Failed to acquire write lock");
        }
    }

    pub unsafe fn write_unlock(&self) {
        let lock = self.lock.get();
        assert_eq!(
            (*lock).load(Ordering::Relaxed) & !abi::LOCK_KERNEL_MANAGED.0,
            __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
            "This rwlock is not write-locked by this thread"
        );

        if !(*lock)
            .compare_exchange(
                __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
                abi::LOCK_UNLOCKED.0,
                Ordering::Release,
                Ordering::Relaxed,
            )
            .is_ok()
        {
            // Lock is managed by kernelspace. Call into the kernel
            // to unblock waiting threads.
            let ret = abi::lock_unlock(lock as *mut abi::lock, abi::scope::PRIVATE);
            assert_eq!(ret, abi::errno::SUCCESS, "Failed to write unlock a rwlock");
        }
    }

    pub unsafe fn destroy(&self) {
        let lock = self.lock.get();
        assert_eq!(
            (*lock).load(Ordering::Relaxed),
            abi::LOCK_UNLOCKED.0,
            "Attempted to destroy locked rwlock"
        );
    }
}
<commit_msg>Get rid of UnsafeCell in cloudabi rwlock.<commit_after>use crate::mem;
use crate::mem::MaybeUninit;
use crate::sync::atomic::{AtomicU32, Ordering};
use crate::sys::cloudabi::abi;

extern "C" {
    #[thread_local]
    static __pthread_thread_id: abi::tid;
}

#[thread_local]
static mut RDLOCKS_ACQUIRED: u32 = 0;

pub struct RWLock {
    lock: AtomicU32,
}

pub unsafe fn raw(r: &RWLock) -> *mut AtomicU32 {
    &r.lock as *const AtomicU32 as *mut AtomicU32
}

unsafe impl Send for RWLock {}
unsafe impl Sync for RWLock {}

impl RWLock {
    pub const fn new() -> RWLock {
        RWLock { lock: AtomicU32::new(abi::LOCK_UNLOCKED.0) }
    }

    pub unsafe fn try_read(&self) -> bool {
        let mut old = abi::LOCK_UNLOCKED.0;
        while let Err(cur) =
            self.lock.compare_exchange_weak(old, old + 1, Ordering::Acquire, Ordering::Relaxed)
        {
            if (cur & abi::LOCK_WRLOCKED.0) != 0 {
                // Another thread already has a write lock.
                assert_ne!(
                    old & !abi::LOCK_KERNEL_MANAGED.0,
                    __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
                    "Attempted to acquire a read lock while holding a write lock"
                );
                return false;
            } else if (old & abi::LOCK_KERNEL_MANAGED.0) != 0 && RDLOCKS_ACQUIRED == 0 {
                // Lock has threads waiting for the lock. Only acquire
                // the lock if we have already acquired read locks. In
                // that case, it is justified to acquire this lock to
                // prevent a deadlock.
                return false;
            }
            old = cur;
        }

        RDLOCKS_ACQUIRED += 1;
        true
    }

    pub unsafe fn read(&self) {
        if !self.try_read() {
            // Call into the kernel to acquire a read lock.
            let subscription = abi::subscription {
                type_: abi::eventtype::LOCK_RDLOCK,
                union: abi::subscription_union {
                    lock: abi::subscription_lock {
                        lock: &self.lock as *const AtomicU32 as *mut abi::lock,
                        lock_scope: abi::scope::PRIVATE,
                    },
                },
                ..mem::zeroed()
            };
            let mut event = MaybeUninit::<abi::event>::uninit();
            let mut nevents = MaybeUninit::<usize>::uninit();
            let ret = abi::poll(&subscription, event.as_mut_ptr(), 1, nevents.as_mut_ptr());
            assert_eq!(ret, abi::errno::SUCCESS, "Failed to acquire read lock");
            let event = event.assume_init();
            assert_eq!(event.error, abi::errno::SUCCESS, "Failed to acquire read lock");

            RDLOCKS_ACQUIRED += 1;
        }
    }

    pub unsafe fn read_unlock(&self) {
        // Perform a read unlock. We can do this in userspace, except when
        // other threads are blocked and we are performing the last unlock.
        // In that case, call into the kernel.
        //
        // Other threads may attempt to increment the read lock count,
        // meaning that the call into the kernel could be spurious. To
        // prevent this from happening, upgrade to a write lock first. This
        // allows us to call into the kernel, having the guarantee that the
        // lock value will not change in the meantime.
        assert!(RDLOCKS_ACQUIRED > 0, "Bad lock count");
        let mut old = 1;
        loop {
            if old == 1 | abi::LOCK_KERNEL_MANAGED.0 {
                // Last read lock while threads are waiting. Attempt to upgrade
                // to a write lock before calling into the kernel to unlock.
                if let Err(cur) = self.lock.compare_exchange_weak(
                    old,
                    __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0 | abi::LOCK_KERNEL_MANAGED.0,
                    Ordering::Acquire,
                    Ordering::Relaxed,
                ) {
                    old = cur;
                } else {
                    // Call into the kernel to unlock.
                    let ret = abi::lock_unlock(
                        &self.lock as *const AtomicU32 as *mut abi::lock,
                        abi::scope::PRIVATE,
                    );
                    assert_eq!(ret, abi::errno::SUCCESS, "Failed to write unlock a rwlock");
                    break;
                }
            } else {
                // No threads waiting or not the last read lock. Just decrement
                // the read lock count.
                assert_ne!(old & !abi::LOCK_KERNEL_MANAGED.0, 0, "This rwlock is not locked");
                assert_eq!(
                    old & abi::LOCK_WRLOCKED.0,
                    0,
                    "Attempted to read-unlock a write-locked rwlock"
                );
                if let Err(cur) = self.lock.compare_exchange_weak(
                    old,
                    old - 1,
                    Ordering::Acquire,
                    Ordering::Relaxed,
                ) {
                    old = cur;
                } else {
                    break;
                }
            }
        }

        RDLOCKS_ACQUIRED -= 1;
    }

    pub unsafe fn try_write(&self) -> bool {
        // Attempt to acquire the lock.
        if let Err(old) = self.lock.compare_exchange(
            abi::LOCK_UNLOCKED.0,
            __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
            Ordering::Acquire,
            Ordering::Relaxed,
        ) {
            // Failure. Crash upon recursive acquisition.
            assert_ne!(
                old & !abi::LOCK_KERNEL_MANAGED.0,
                __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
                "Attempted to recursive write-lock a rwlock",
            );
            false
        } else {
            // Success.
            true
        }
    }

    pub unsafe fn write(&self) {
        if !self.try_write() {
            // Call into the kernel to acquire a write lock.
            let subscription = abi::subscription {
                type_: abi::eventtype::LOCK_WRLOCK,
                union: abi::subscription_union {
                    lock: abi::subscription_lock {
                        lock: &self.lock as *const AtomicU32 as *mut abi::lock,
                        lock_scope: abi::scope::PRIVATE,
                    },
                },
                ..mem::zeroed()
            };
            let mut event = MaybeUninit::<abi::event>::uninit();
            let mut nevents = MaybeUninit::<usize>::uninit();
            let ret = abi::poll(&subscription, event.as_mut_ptr(), 1, nevents.as_mut_ptr());
            assert_eq!(ret, abi::errno::SUCCESS, "Failed to acquire write lock");
            let event = event.assume_init();
            assert_eq!(event.error, abi::errno::SUCCESS, "Failed to acquire write lock");
        }
    }

    pub unsafe fn write_unlock(&self) {
        assert_eq!(
            self.lock.load(Ordering::Relaxed) & !abi::LOCK_KERNEL_MANAGED.0,
            __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
            "This rwlock is not write-locked by this thread"
        );

        if !self
            .lock
            .compare_exchange(
                __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
                abi::LOCK_UNLOCKED.0,
                Ordering::Release,
                Ordering::Relaxed,
            )
            .is_ok()
        {
            // Lock is managed by kernelspace. Call into the kernel
            // to unblock waiting threads.
            let ret = abi::lock_unlock(
                &self.lock as *const AtomicU32 as *mut abi::lock,
                abi::scope::PRIVATE,
            );
            assert_eq!(ret, abi::errno::SUCCESS, "Failed to write unlock a rwlock");
        }
    }

    pub unsafe fn destroy(&self) {
        assert_eq!(
            self.lock.load(Ordering::Relaxed),
            abi::LOCK_UNLOCKED.0,
            "Attempted to destroy locked rwlock"
        );
    }
}
<|endoftext|>


<commit_before>use {Distribution, Generator};

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    /// The size of the probabilty vector.
    pub k: usize,
    /// The probability vector.
    pub p: Vec<f64>,
}

fn is_prob_vec(p: &[f64]) -> bool {
    let mut sum = 0.;
    for &p_i in p.iter() {
        if p_i < 0.|| p_i > 1. { return false; }
        sum += p_i;
    }
    (sum - 1.).abs() <= 1e-12
}

impl Categorical {
    /// Create a Categorical distribution with success probability `p`.
    ///
    /// # Panics
    ///
    /// Panics if `p < 0` or `p > 1`.
    #[inline]
    pub fn new(p: &[f64]) -> Categorical {
        debug_assert!(is_prob_vec(p), "Categorical::new() is called with p not a probabilty vector");
        Categorical { k: p.len(),  p: p.to_vec()}
    }
}

impl Distribution for Categorical {
    type Value = usize;

    #[inline]
    fn mean(&self) -> f64 {
        // sum_{i=0}^k i p_i
        self.p.iter().enumerate().fold(0., |acc, b| acc + b.0 as f64 * b.1)
    }

    #[inline]
    fn var(&self) -> f64 {
        let mu = self.mean();
        self.p.iter().enumerate().fold(
            0., |acc, b| acc + (b.0 as f64 - mu).powi(2) * b.1
        )
    }

    #[inline]
    fn skewness(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let skew = self.p.iter().enumerate().fold(
            0., |acc, b| acc + (b.0 as f64 - mu).powi(3) * b.1
        );
        skew / (sigma2 * sigma2.sqrt())
    }

    fn kurtosis(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let kurt = self.p.iter().enumerate().fold(
            0., |acc, b| acc + (b.0 as f64 - mu).powi(4) * b.1
        );
        kurt / sigma2.powi(2) - 3.
    }

    fn median(&self) -> f64 {
        if self.p[0] > 0.5 { return 0.0; }
        else if self.p[0] == 0.5 { return 0.5; }
        let mut sum = 0.;
        for i in 0..self.k {
            sum += self.p[i];
            if sum == 0.5 { return (2 * i - 1) as f64 / 2.; }
            else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }

    fn modes(&self) -> Vec<Self::Value> {
        let mut m = Vec::new();
        let mut max = 0.;
        for (i, &p_i) in self.p.iter().enumerate() {
            if p_i == max { m.push(i); }
            if p_i > max {
                max = p_i;
                m = vec![i];
            }
        }
        m
    }

    #[inline]
    fn entropy(&self) -> f64 {
        - self.p.iter().fold(0., |acc, p_i| acc + p_i * p_i.ln())
    }

    #[inline]
    fn cdf(&self, x: Self::Value) -> f64 {
        if x >= self.k - 1 { 1. }
        else {
            self.p.iter().take(x + 1).fold(0., |a, b| a + b)
        }
    }

    fn inv_cdf(&self, p: f64) -> Self::Value {
        debug_assert!(0.0 <= p && p <= 1.0, "inv_cdf is called with p outside of [0, 1]");
        if p == 0. {
            // return the first non-zero index
            return self.p.iter().enumerate().find(|&v| *v.1 > 0.).unwrap().0;
        }
        let mut sum = 0.;
        for i in 0..self.k {
            sum += self.p[i];
            if sum >= p || sum == 1. { return i; }
        }
        self.k - 1
    }

    #[inline]
    fn pdf(&self, x: Self::Value) -> f64 {
        self.p[x]
    }

    #[inline]
    fn sample<G: Generator>(&self, generator: &mut G) -> Self::Value {
        self.inv_cdf(generator.gen::<f64>())
    }
}

#[cfg(test)]
mod tests {
    use assert;

    use {Distribution, Sampler};
    use distributions::Categorical;

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1./$k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    #[should_panic]
    fn invalid_prob_vec_1() {
        new!([0.6, 0.6, 0.6]);
    }

    #[test]
    #[should_panic]
    fn invalid_prob_vec_2() {
        new!([0.6, 0.6, 0.6, -0.2]);
    }

    #[test]
    #[should_panic]
    fn invalid_prob_vec_3() {
        new!([1.2, -0.2]);
    }

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).mean(), 1.5);
    }

    #[test]
    fn var() {
        assert_eq!(new!(equal 3).var(), 2./3.);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).var(), 11./12.);
    }

    #[test]
    fn sd() {
        assert_eq!(new!(equal 2).sd(), 0.5);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).sd(), 0.9574271077563381);
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).skewness(), 0.);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).modes(), vec![1, 2]);
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn pdf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d1 = new!(p);
        assert::equal(&(0..4).map(|x| d1.pdf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d2 = new!(equal 3);
        assert::equal(&(0..3).map(|x| d2.pdf(x)).collect::<Vec<_>>(), &vec![1./3.; 3])
    }

    #[test]
    fn cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        assert::equal(&(0..4).map(|x| d1.cdf(x)).collect::<Vec<_>>(),
                      &vec![0.0, 0.75, 1.0, 1.0]);

        let d2 = new!(equal 3);
        assert::equal(&(0..3).map(|x| d2.cdf(x)).collect::<Vec<_>>(),
                      &vec![1./3., 2./3., 1.]);
    }

    #[test]
    fn inv_cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        let p1 = vec![0.0, 0.75, 0.7500001, 1.0];
        assert::equal(&p1.iter().map(|&p| d1.inv_cdf(p)).collect::<Vec<_>>(),
                      &vec![1, 1, 2, 2]);

        let d2 = new!(equal 3);
        let p2 = vec![0.0, 0.5, 0.75, 1.0];
        assert::equal(&p2.iter().map(|&p| d2.inv_cdf(p)).collect::<Vec<_>>(),
                      &vec![0, 1, 2, 2]);

    }

    #[test]
    fn sample() {
        let mut generator = ::generator();

        // Discrete Uniform(1, 2)
        let sum = Sampler(&new!([0.0, 0.5, 0.5]), &mut generator)
            .take(100)
            .fold(0, |a, b| a + b);

        assert!(100 <= sum && sum <= 200);
    }
}
<commit_msg>Moved the probability vector check inside the debug scope<commit_after>use {Distribution, Generator};

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    /// The size of the probabilty vector.
    pub k: usize,
    /// The probability vector.
    pub p: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// # Panics
    ///
    /// Panics if `p[i] < 0` or `p[i] > 1` for any `i = 0..k`, or if
    /// `sum(p) != 1`.
    #[inline]
    pub fn new(p: &[f64]) -> Categorical {
        debug_assert!({
            fn is_prob_vec(p: &[f64]) -> bool {
                // Check if p is a probability vector.
                let mut sum = 0.;
                for &p_i in p.iter() {
                    if p_i < 0.|| p_i > 1. { return false; }
                    sum += p_i;
                }
                (sum - 1.).abs() <= 1e-12
            }
            is_prob_vec(&p)
        }, "Categorical::new() is called with p not a probabilty vector");

        Categorical { k: p.len(),  p: p.to_vec()}
    }
}

impl Distribution for Categorical {
    type Value = usize;

    #[inline]
    fn mean(&self) -> f64 {
        // sum_{i=0}^k i p_i
        self.p.iter().enumerate().fold(0., |acc, b| acc + b.0 as f64 * b.1)
    }

    #[inline]
    fn var(&self) -> f64 {
        let mu = self.mean();
        self.p.iter().enumerate().fold(
            0., |acc, b| acc + (b.0 as f64 - mu).powi(2) * b.1
        )
    }

    #[inline]
    fn skewness(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let skew = self.p.iter().enumerate().fold(
            0., |acc, b| acc + (b.0 as f64 - mu).powi(3) * b.1
        );
        skew / (sigma2 * sigma2.sqrt())
    }

    fn kurtosis(&self) -> f64 {
        let mu = self.mean();
        let sigma2 = self.var();
        let kurt = self.p.iter().enumerate().fold(
            0., |acc, b| acc + (b.0 as f64 - mu).powi(4) * b.1
        );
        kurt / sigma2.powi(2) - 3.
    }

    fn median(&self) -> f64 {
        if self.p[0] > 0.5 { return 0.0; }
        else if self.p[0] == 0.5 { return 0.5; }
        let mut sum = 0.;
        for i in 0..self.k {
            sum += self.p[i];
            if sum == 0.5 { return (2 * i - 1) as f64 / 2.; }
            else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }

    fn modes(&self) -> Vec<Self::Value> {
        let mut m = Vec::new();
        let mut max = 0.;
        for (i, &p_i) in self.p.iter().enumerate() {
            if p_i == max { m.push(i); }
            if p_i > max {
                max = p_i;
                m = vec![i];
            }
        }
        m
    }

    #[inline]
    fn entropy(&self) -> f64 {
        - self.p.iter().fold(0., |acc, p_i| acc + p_i * p_i.ln())
    }

    #[inline]
    fn cdf(&self, x: Self::Value) -> f64 {
        if x >= self.k - 1 { 1. }
        else {
            self.p.iter().take(x + 1).fold(0., |a, b| a + b)
        }
    }

    fn inv_cdf(&self, p: f64) -> Self::Value {
        debug_assert!(0.0 <= p && p <= 1.0, "inv_cdf is called with p outside of [0, 1]");
        if p == 0. {
            // return the first non-zero index
            return self.p.iter().enumerate().find(|&v| *v.1 > 0.).unwrap().0;
        }
        let mut sum = 0.;
        for i in 0..self.k {
            sum += self.p[i];
            if sum >= p || sum == 1. { return i; }
        }
        self.k - 1
    }

    #[inline]
    fn pdf(&self, x: Self::Value) -> f64 {
        self.p[x]
    }

    #[inline]
    fn sample<G: Generator>(&self, generator: &mut G) -> Self::Value {
        self.inv_cdf(generator.gen::<f64>())
    }
}

#[cfg(test)]
mod tests {
    use assert;

    use {Distribution, Sampler};
    use distributions::Categorical;

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1./$k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    #[should_panic]
    fn invalid_prob_vec_1() {
        new!([0.6, 0.6, 0.6]);
    }

    #[test]
    #[should_panic]
    fn invalid_prob_vec_2() {
        new!([0.6, 0.6, 0.6, -0.2]);
    }

    #[test]
    #[should_panic]
    fn invalid_prob_vec_3() {
        new!([1.2, -0.2]);
    }

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).mean(), 1.5);
    }

    #[test]
    fn var() {
        assert_eq!(new!(equal 3).var(), 2./3.);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).var(), 11./12.);
    }

    #[test]
    fn sd() {
        assert_eq!(new!(equal 2).sd(), 0.5);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).sd(), 0.9574271077563381);
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).skewness(), 0.);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1./6., 1./3., 1./3., 1./6.]).modes(), vec![1, 2]);
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn pdf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d1 = new!(p);
        assert::equal(&(0..4).map(|x| d1.pdf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d2 = new!(equal 3);
        assert::equal(&(0..3).map(|x| d2.pdf(x)).collect::<Vec<_>>(), &vec![1./3.; 3])
    }

    #[test]
    fn cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        assert::equal(&(0..4).map(|x| d1.cdf(x)).collect::<Vec<_>>(),
                      &vec![0.0, 0.75, 1.0, 1.0]);

        let d2 = new!(equal 3);
        assert::equal(&(0..3).map(|x| d2.cdf(x)).collect::<Vec<_>>(),
                      &vec![1./3., 2./3., 1.]);
    }

    #[test]
    fn inv_cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        let p1 = vec![0.0, 0.75, 0.7500001, 1.0];
        assert::equal(&p1.iter().map(|&p| d1.inv_cdf(p)).collect::<Vec<_>>(),
                      &vec![1, 1, 2, 2]);

        let d2 = new!(equal 3);
        let p2 = vec![0.0, 0.5, 0.75, 1.0];
        assert::equal(&p2.iter().map(|&p| d2.inv_cdf(p)).collect::<Vec<_>>(),
                      &vec![0, 1, 2, 2]);

    }

    #[test]
    fn sample() {
        let mut generator = ::generator();

        // Discrete Uniform(1, 2)
        let sum = Sampler(&new!([0.0, 0.5, 0.5]), &mut generator)
            .take(100)
            .fold(0, |a, b| a + b);

        assert!(100 <= sum && sum <= 200);
    }
}
<|endoftext|>


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

#![cfg(not(test))]

extern crate env_logger;
extern crate getopts;
extern crate rustfmt_nightly as rustfmt;

use std::{env, error};
use std::fs::File;
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use std::str::FromStr;

use getopts::{Matches, Options};

use rustfmt::{run, Input, Summary};
use rustfmt::file_lines::FileLines;
use rustfmt::config::{get_toml_path, Config, WriteMode};

type FmtError = Box<error::Error + Send + Sync>;
type FmtResult<T> = std::result::Result<T, FmtError>;

/// Rustfmt operations.
enum Operation {
    /// Format files and their child modules.
    Format {
        files: Vec<PathBuf>,
        config_path: Option<PathBuf>,
        minimal_config_path: Option<String>,
    },
    /// Print the help message.
    Help,
    // Print version information
    Version,
    /// Print detailed configuration help.
    ConfigHelp,
    /// Output default config to a file
    ConfigOutputDefault { path: String },
    /// No file specified, read from stdin
    Stdin {
        input: String,
        config_path: Option<PathBuf>,
    },
}

/// Parsed command line options.
#[derive(Clone, Debug, Default)]
struct CliOptions {
    skip_children: bool,
    verbose: bool,
    write_mode: Option<WriteMode>,
    file_lines: FileLines, // Default is all lines in all files.
}

impl CliOptions {
    fn from_matches(matches: &Matches) -> FmtResult<CliOptions> {
        let mut options = CliOptions::default();
        options.skip_children = matches.opt_present("skip-children");
        options.verbose = matches.opt_present("verbose");

        if let Some(ref write_mode) = matches.opt_str("write-mode") {
            if let Ok(write_mode) = WriteMode::from_str(write_mode) {
                options.write_mode = Some(write_mode);
            } else {
                return Err(FmtError::from(
                    format!("Invalid write-mode: {}", write_mode),
                ));
            }
        }

        if let Some(ref file_lines) = matches.opt_str("file-lines") {
            options.file_lines = file_lines.parse()?;
        }

        Ok(options)
    }

    fn apply_to(self, config: &mut Config) {
        config.set().skip_children(self.skip_children);
        config.set().verbose(self.verbose);
        config.set().file_lines(self.file_lines);
        if let Some(write_mode) = self.write_mode {
            config.set().write_mode(write_mode);
        }
    }
}

/// read the given config file path recursively if present else read the project file path
fn match_cli_path_or_file(
    config_path: Option<PathBuf>,
    input_file: &Path,
) -> FmtResult<(Config, Option<PathBuf>)> {
    if let Some(config_file) = config_path {
        let toml = Config::from_toml_path(config_file.as_ref())?;
        return Ok((toml, Some(config_file)));
    }
    Config::from_resolved_toml_path(input_file).map_err(FmtError::from)
}

fn make_opts() -> Options {
    let mut opts = Options::new();
    opts.optflag("h", "help", "show this message");
    opts.optflag("V", "version", "show version information");
    opts.optflag("v", "verbose", "print verbose output");
    opts.optopt(
        "",
        "write-mode",
        "how to write output (not usable when piping from stdin)",
        "[replace|overwrite|display|plain|diff|coverage|checkstyle]",
    );
    opts.optflag("", "skip-children", "don't reformat child modules");

    opts.optflag(
        "",
        "config-help",
        "show details of rustfmt configuration options",
    );
    opts.optopt(
        "",
        "dump-default-config",
        "Dumps the default configuration to a file and exits.",
        "PATH",
    );
    opts.optopt(
        "",
        "dump-minimal-config",
        "Dumps configuration options that were checked during formatting to a file.",
        "PATH",
    );
    opts.optopt(
        "",
        "config-path",
        "Recursively searches the given path for the rustfmt.toml config file. If not \
         found reverts to the input file path",
        "[Path for the configuration file]",
    );
    opts.optopt(
        "",
        "file-lines",
        "Format specified line ranges. See README for more detail on the JSON format.",
        "JSON",
    );

    opts
}

fn execute(opts: &Options) -> FmtResult<Summary> {
    let matches = opts.parse(env::args().skip(1))?;

    match determine_operation(&matches)? {
        Operation::Help => {
            print_usage(opts, "");
            Summary::print_exit_codes();
            Ok(Summary::default())
        }
        Operation::Version => {
            print_version();
            Ok(Summary::default())
        }
        Operation::ConfigHelp => {
            Config::print_docs();
            Ok(Summary::default())
        }
        Operation::ConfigOutputDefault { path } => {
            let mut file = File::create(path)?;
            let toml = Config::default().all_options().to_toml()?;
            file.write_all(toml.as_bytes())?;
            Ok(Summary::default())
        }
        Operation::Stdin { input, config_path } => {
            // try to read config from local directory
            let (mut config, _) =
                match_cli_path_or_file(config_path, &env::current_dir().unwrap())?;

            // write_mode is always Plain for Stdin.
            config.set().write_mode(WriteMode::Plain);

            // parse file_lines
            if let Some(ref file_lines) = matches.opt_str("file-lines") {
                config.set().file_lines(file_lines.parse()?);
                for f in config.file_lines().files() {
                    if f != "stdin" {
                        println!("Warning: Extra file listed in file_lines option '{}'", f);
                    }
                }
            }

            Ok(run(Input::Text(input), &config))
        }
        Operation::Format {
            files,
            config_path,
            minimal_config_path,
        } => {
            let options = CliOptions::from_matches(&matches)?;

            for f in options.file_lines.files() {
                if !files.contains(&PathBuf::from(f)) {
                    println!("Warning: Extra file listed in file_lines option '{}'", f);
                }
            }

            let mut config = Config::default();
            // Load the config path file if provided
            if let Some(config_file) = config_path.as_ref() {
                config = Config::from_toml_path(config_file.as_ref())?;
            };

            if options.verbose {
                if let Some(path) = config_path.as_ref() {
                    println!("Using rustfmt config file {}", path.display());
                }
            }

            let mut error_summary = Summary::default();
            for file in files {
                if !file.exists() {
                    println!("Error: file `{}` does not exist", file.to_str().unwrap());
                    error_summary.add_operational_error();
                } else if file.is_dir() {
                    println!("Error: `{}` is a directory", file.to_str().unwrap());
                    error_summary.add_operational_error();
                } else {
                    // Check the file directory if the config-path could not be read or not provided
                    if config_path.is_none() {
                        let (config_tmp, path_tmp) =
                            Config::from_resolved_toml_path(file.parent().unwrap())?;
                        if options.verbose {
                            if let Some(path) = path_tmp.as_ref() {
                                println!(
                                    "Using rustfmt config file {} for {}",
                                    path.display(),
                                    file.display()
                                );
                            }
                        }
                        config = config_tmp;
                    }

                    options.clone().apply_to(&mut config);
                    error_summary.add(run(Input::File(file), &config));
                }
            }

            // If we were given a path via dump-minimal-config, output any options
            // that were used during formatting as TOML.
            if let Some(path) = minimal_config_path {
                let mut file = File::create(path)?;
                let toml = config.used_options().to_toml()?;
                file.write_all(toml.as_bytes())?;
            }

            Ok(error_summary)
        }
    }
}

fn main() {
    let _ = env_logger::init();

    let opts = make_opts();

    let exit_code = match execute(&opts) {
        Ok(summary) => {
            if summary.has_operational_errors() {
                1
            } else if summary.has_parsing_errors() {
                2
            } else if summary.has_formatting_errors() {
                3
            } else if summary.has_diff {
                // should only happen in diff mode
                4
            } else {
                assert!(summary.has_no_errors());
                0
            }
        }
        Err(e) => {
            print_usage(&opts, &e.to_string());
            1
        }
    };
    // Make sure standard output is flushed before we exit.
    std::io::stdout().flush().unwrap();

    // Exit with given exit code.
    //
    // NOTE: This immediately terminates the process without doing any cleanup,
    // so make sure to finish all necessary cleanup before this is called.
    std::process::exit(exit_code);
}

fn print_usage(opts: &Options, reason: &str) {
    let reason = format!(
        "{}\n\nusage: {} [options] <file>...",
        reason,
        env::args_os().next().unwrap().to_string_lossy()
    );
    println!("{}", opts.usage(&reason));
}

fn print_version() {
    println!(
        "{}-nightly{}",
        env!("CARGO_PKG_VERSION"),
        include_str!(concat!(env!("OUT_DIR"), "/commit-info.txt"))
    )
}

fn determine_operation(matches: &Matches) -> FmtResult<Operation> {
    if matches.opt_present("h") {
        return Ok(Operation::Help);
    }

    if matches.opt_present("config-help") {
        return Ok(Operation::ConfigHelp);
    }

    if let Some(path) = matches.opt_str("dump-default-config") {
        return Ok(Operation::ConfigOutputDefault { path });
    }

    if matches.opt_present("version") {
        return Ok(Operation::Version);
    }

    let config_path_not_found = |path: &str| -> FmtResult<Operation> {
        Err(FmtError::from(format!(
            "Error: unable to find a config file for the given path: `{}`",
            path
        )))
    };

    // Read the config_path and convert to parent dir if a file is provided.
    // If a config file cannot be found from the given path, return error.
    let config_path: Option<PathBuf> = match matches.opt_str("config-path").map(PathBuf::from) {
        Some(ref path) if !path.exists() => return config_path_not_found(path.to_str().unwrap()),
        Some(ref path) if path.is_dir() => {
            let config_file_path = get_toml_path(path)?;
            if config_file_path.is_some() {
                config_file_path
            } else {
                return config_path_not_found(path.to_str().unwrap());
            }
        }
        path => path,
    };

    // If no path is given, we won't output a minimal config.
    let minimal_config_path = matches.opt_str("dump-minimal-config");

    // if no file argument is supplied, read from stdin
    if matches.free.is_empty() {
        let mut buffer = String::new();
        io::stdin().read_to_string(&mut buffer)?;

        return Ok(Operation::Stdin {
            input: buffer,
            config_path: config_path,
        });
    }

    let files: Vec<_> = matches
        .free
        .iter()
        .map(|s| {
            let p = PathBuf::from(s);
            // we will do comparison later, so here tries to canonicalize first
            // to get the expected behavior.
            p.canonicalize().unwrap_or(p)
        })
        .collect();

    Ok(Operation::Format {
        files: files,
        config_path: config_path,
        minimal_config_path: minimal_config_path,
    })
}
<commit_msg>output --dump-default-config to stdout if no path is given<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![cfg(not(test))]

extern crate env_logger;
extern crate getopts;
extern crate rustfmt_nightly as rustfmt;

use std::{env, error};
use std::fs::File;
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use std::str::FromStr;

use getopts::{HasArg, Matches, Occur, Options};

use rustfmt::{run, Input, Summary};
use rustfmt::file_lines::FileLines;
use rustfmt::config::{get_toml_path, Config, WriteMode};

type FmtError = Box<error::Error + Send + Sync>;
type FmtResult<T> = std::result::Result<T, FmtError>;

/// Rustfmt operations.
enum Operation {
    /// Format files and their child modules.
    Format {
        files: Vec<PathBuf>,
        config_path: Option<PathBuf>,
        minimal_config_path: Option<String>,
    },
    /// Print the help message.
    Help,
    // Print version information
    Version,
    /// Print detailed configuration help.
    ConfigHelp,
    /// Output default config to a file, or stdout if None
    ConfigOutputDefault { path: Option<String> },
    /// No file specified, read from stdin
    Stdin {
        input: String,
        config_path: Option<PathBuf>,
    },
}

/// Parsed command line options.
#[derive(Clone, Debug, Default)]
struct CliOptions {
    skip_children: bool,
    verbose: bool,
    write_mode: Option<WriteMode>,
    file_lines: FileLines, // Default is all lines in all files.
}

impl CliOptions {
    fn from_matches(matches: &Matches) -> FmtResult<CliOptions> {
        let mut options = CliOptions::default();
        options.skip_children = matches.opt_present("skip-children");
        options.verbose = matches.opt_present("verbose");

        if let Some(ref write_mode) = matches.opt_str("write-mode") {
            if let Ok(write_mode) = WriteMode::from_str(write_mode) {
                options.write_mode = Some(write_mode);
            } else {
                return Err(FmtError::from(
                    format!("Invalid write-mode: {}", write_mode),
                ));
            }
        }

        if let Some(ref file_lines) = matches.opt_str("file-lines") {
            options.file_lines = file_lines.parse()?;
        }

        Ok(options)
    }

    fn apply_to(self, config: &mut Config) {
        config.set().skip_children(self.skip_children);
        config.set().verbose(self.verbose);
        config.set().file_lines(self.file_lines);
        if let Some(write_mode) = self.write_mode {
            config.set().write_mode(write_mode);
        }
    }
}

/// read the given config file path recursively if present else read the project file path
fn match_cli_path_or_file(
    config_path: Option<PathBuf>,
    input_file: &Path,
) -> FmtResult<(Config, Option<PathBuf>)> {
    if let Some(config_file) = config_path {
        let toml = Config::from_toml_path(config_file.as_ref())?;
        return Ok((toml, Some(config_file)));
    }
    Config::from_resolved_toml_path(input_file).map_err(FmtError::from)
}

fn make_opts() -> Options {
    let mut opts = Options::new();
    opts.optflag("h", "help", "show this message");
    opts.optflag("V", "version", "show version information");
    opts.optflag("v", "verbose", "print verbose output");
    opts.optopt(
        "",
        "write-mode",
        "how to write output (not usable when piping from stdin)",
        "[replace|overwrite|display|plain|diff|coverage|checkstyle]",
    );
    opts.optflag("", "skip-children", "don't reformat child modules");

    opts.optflag(
        "",
        "config-help",
        "show details of rustfmt configuration options",
    );
    opts.opt(
        "",
        "dump-default-config",
        "Dumps the default configuration to a file and exits. PATH defaults to rustfmt.toml if \
         omitted.",
        "PATH",
        HasArg::Maybe,
        Occur::Optional,
    );
    opts.optopt(
        "",
        "dump-minimal-config",
        "Dumps configuration options that were checked during formatting to a file.",
        "PATH",
    );
    opts.optopt(
        "",
        "config-path",
        "Recursively searches the given path for the rustfmt.toml config file. If not \
         found reverts to the input file path",
        "[Path for the configuration file]",
    );
    opts.optopt(
        "",
        "file-lines",
        "Format specified line ranges. See README for more detail on the JSON format.",
        "JSON",
    );

    opts
}

fn execute(opts: &Options) -> FmtResult<Summary> {
    let matches = opts.parse(env::args().skip(1))?;

    match determine_operation(&matches)? {
        Operation::Help => {
            print_usage(opts, "");
            Summary::print_exit_codes();
            Ok(Summary::default())
        }
        Operation::Version => {
            print_version();
            Ok(Summary::default())
        }
        Operation::ConfigHelp => {
            Config::print_docs();
            Ok(Summary::default())
        }
        Operation::ConfigOutputDefault { path } => {
            let toml = Config::default().all_options().to_toml()?;
            if let Some(path) = path {
                let mut file = File::create(path)?;
                file.write_all(toml.as_bytes())?;
            } else {
                io::stdout().write_all(toml.as_bytes())?;
            }
            Ok(Summary::default())
        }
        Operation::Stdin { input, config_path } => {
            // try to read config from local directory
            let (mut config, _) =
                match_cli_path_or_file(config_path, &env::current_dir().unwrap())?;

            // write_mode is always Plain for Stdin.
            config.set().write_mode(WriteMode::Plain);

            // parse file_lines
            if let Some(ref file_lines) = matches.opt_str("file-lines") {
                config.set().file_lines(file_lines.parse()?);
                for f in config.file_lines().files() {
                    if f != "stdin" {
                        println!("Warning: Extra file listed in file_lines option '{}'", f);
                    }
                }
            }

            Ok(run(Input::Text(input), &config))
        }
        Operation::Format {
            files,
            config_path,
            minimal_config_path,
        } => {
            let options = CliOptions::from_matches(&matches)?;

            for f in options.file_lines.files() {
                if !files.contains(&PathBuf::from(f)) {
                    println!("Warning: Extra file listed in file_lines option '{}'", f);
                }
            }

            let mut config = Config::default();
            // Load the config path file if provided
            if let Some(config_file) = config_path.as_ref() {
                config = Config::from_toml_path(config_file.as_ref())?;
            };

            if options.verbose {
                if let Some(path) = config_path.as_ref() {
                    println!("Using rustfmt config file {}", path.display());
                }
            }

            let mut error_summary = Summary::default();
            for file in files {
                if !file.exists() {
                    println!("Error: file `{}` does not exist", file.to_str().unwrap());
                    error_summary.add_operational_error();
                } else if file.is_dir() {
                    println!("Error: `{}` is a directory", file.to_str().unwrap());
                    error_summary.add_operational_error();
                } else {
                    // Check the file directory if the config-path could not be read or not provided
                    if config_path.is_none() {
                        let (config_tmp, path_tmp) =
                            Config::from_resolved_toml_path(file.parent().unwrap())?;
                        if options.verbose {
                            if let Some(path) = path_tmp.as_ref() {
                                println!(
                                    "Using rustfmt config file {} for {}",
                                    path.display(),
                                    file.display()
                                );
                            }
                        }
                        config = config_tmp;
                    }

                    options.clone().apply_to(&mut config);
                    error_summary.add(run(Input::File(file), &config));
                }
            }

            // If we were given a path via dump-minimal-config, output any options
            // that were used during formatting as TOML.
            if let Some(path) = minimal_config_path {
                let mut file = File::create(path)?;
                let toml = config.used_options().to_toml()?;
                file.write_all(toml.as_bytes())?;
            }

            Ok(error_summary)
        }
    }
}

fn main() {
    let _ = env_logger::init();

    let opts = make_opts();

    let exit_code = match execute(&opts) {
        Ok(summary) => {
            if summary.has_operational_errors() {
                1
            } else if summary.has_parsing_errors() {
                2
            } else if summary.has_formatting_errors() {
                3
            } else if summary.has_diff {
                // should only happen in diff mode
                4
            } else {
                assert!(summary.has_no_errors());
                0
            }
        }
        Err(e) => {
            print_usage(&opts, &e.to_string());
            1
        }
    };
    // Make sure standard output is flushed before we exit.
    std::io::stdout().flush().unwrap();

    // Exit with given exit code.
    //
    // NOTE: This immediately terminates the process without doing any cleanup,
    // so make sure to finish all necessary cleanup before this is called.
    std::process::exit(exit_code);
}

fn print_usage(opts: &Options, reason: &str) {
    let reason = format!(
        "{}\n\nusage: {} [options] <file>...",
        reason,
        env::args_os().next().unwrap().to_string_lossy()
    );
    println!("{}", opts.usage(&reason));
}

fn print_version() {
    println!(
        "{}-nightly{}",
        env!("CARGO_PKG_VERSION"),
        include_str!(concat!(env!("OUT_DIR"), "/commit-info.txt"))
    )
}

fn determine_operation(matches: &Matches) -> FmtResult<Operation> {
    if matches.opt_present("h") {
        return Ok(Operation::Help);
    }

    if matches.opt_present("config-help") {
        return Ok(Operation::ConfigHelp);
    }

    if matches.opt_present("dump-default-config") {
        // NOTE for some reason when configured with HasArg::Maybe + Occur::Optional opt_default
        // doesn't recognize `--foo bar` as a long flag with an argument but as a long flag with no
        // argument *plus* a free argument. Thus we check for that case in this branch -- this is
        // required for backward compatibility.
        if let Some(path) = matches.free.get(0) {
            return Ok(Operation::ConfigOutputDefault {
                path: Some(path.clone()),
            });
        } else {
            return Ok(Operation::ConfigOutputDefault {
                path: matches.opt_str("dump-default-config"),
            });
        }
    }

    if matches.opt_present("version") {
        return Ok(Operation::Version);
    }

    let config_path_not_found = |path: &str| -> FmtResult<Operation> {
        Err(FmtError::from(format!(
            "Error: unable to find a config file for the given path: `{}`",
            path
        )))
    };

    // Read the config_path and convert to parent dir if a file is provided.
    // If a config file cannot be found from the given path, return error.
    let config_path: Option<PathBuf> = match matches.opt_str("config-path").map(PathBuf::from) {
        Some(ref path) if !path.exists() => return config_path_not_found(path.to_str().unwrap()),
        Some(ref path) if path.is_dir() => {
            let config_file_path = get_toml_path(path)?;
            if config_file_path.is_some() {
                config_file_path
            } else {
                return config_path_not_found(path.to_str().unwrap());
            }
        }
        path => path,
    };

    // If no path is given, we won't output a minimal config.
    let minimal_config_path = matches.opt_str("dump-minimal-config");

    // if no file argument is supplied, read from stdin
    if matches.free.is_empty() {
        let mut buffer = String::new();
        io::stdin().read_to_string(&mut buffer)?;

        return Ok(Operation::Stdin {
            input: buffer,
            config_path: config_path,
        });
    }

    let files: Vec<_> = matches
        .free
        .iter()
        .map(|s| {
            let p = PathBuf::from(s);
            // we will do comparison later, so here tries to canonicalize first
            // to get the expected behavior.
            p.canonicalize().unwrap_or(p)
        })
        .collect();

    Ok(Operation::Format {
        files: files,
        config_path: config_path,
        minimal_config_path: minimal_config_path,
    })
}
<|endoftext|>


<commit_before>extern crate futures;
extern crate text_ui;
extern crate tokio;
extern crate tokio_serde_msgpack;
extern crate tokio_uds;
extern crate weaver;

use text_ui::app::App;
use text_ui::backend::{color, Backend, Color};
use text_ui::pane::Pane;
//use text_ui::widget::DbgDump;
use text_ui::widget::Widget;
use text_ui::widget::{shared, Line, Linear, Readline, Shared, Text};
use text_ui::{text_to_lines, Event, Input, Key, Position, Size};

use std::sync::{Arc, RwLock};
use std::thread;

use tokio::prelude::Future;
use weaver::{WeaverClient, WeaverCommand, WeaverNotification, WeaverState};

struct WeaverStateWidget {
    state: Shared<WeaverState>,
    selected: Option<u32>,
}

impl WeaverStateWidget {
    pub fn new(state: Shared<WeaverState>) -> Self {
        let selected = None;
        WeaverStateWidget { state, selected }
    }

    pub fn find_cmd_by_index(&self, i: usize) -> Option<String> {
        let state = self.state.read().unwrap();
        let rv = state
            .command_history
            .iter()
            .rev()
            .skip(i)
            .next()
            .map(|(_, cmd)| cmd.cmd.clone());
        rv
    }
}

fn render_command_summary(
    cmd: &WeaverCommand,
    width: usize,
    maxlines: usize,
    selected: bool,
) -> Pane {
    let mut pane = Pane::new_width(width);
    let (icon, style) = match cmd.status {
        None => ('…', "command.running"),
        Some(0) => ('✔', "command.success"),
        _ => ('X', "command.failed"),
    };
    let status_pane = Pane::new_styled(
        Position::new(0, 0),
        Size::new(1, 1),
        vec![icon.to_string()],
        style,
    );
    pane.push_child(status_pane);
    let subwidth = width - 1;
    let command_line = text_to_lines(cmd.cmd.clone(), subwidth);
    let pos = Position::new(1, 0);
    let textlen = command_line.len();
    let mut offset = textlen;
    let size = Size::new(subwidth as u16, textlen as u16);
    let prefix = match selected {
        true => "selected.",
        false => "",
    };
    pane.push_child(Pane::new_styled(
        pos,
        size,
        command_line,
        &format!("{}command", prefix),
    ));
    if cmd.stdout.len() > 0 {
        let mut stdout = text_to_lines(cmd.stdout.clone(), subwidth);
        let pos = Position::new(1, offset as u16);
        let mut textlen = stdout.len();
        if textlen > maxlines {
            stdout = stdout.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        offset += textlen;
        let size = Size::new(subwidth as u16, textlen as u16);
        pane.push_child(Pane::new_styled(
            pos,
            size,
            stdout,
            &format!("{}stdout", prefix),
        ));
    }
    if cmd.stderr.len() > 0 {
        let mut stderr = text_to_lines(cmd.stderr.clone(), subwidth);
        let pos = Position::new(1, offset as u16);
        let mut textlen = stderr.len();
        if textlen > maxlines {
            stderr = stderr.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        let size = Size::new(subwidth as u16, textlen as u16);
        pane.push_child(Pane::new_styled(
            pos,
            size,
            stderr,
            &format!("{}stderr", prefix),
        ));
    }
    pane
}

fn render_command_detail(cmd: &WeaverCommand, size: Size) -> Pane {
    let mut pane = Pane::new_width(size.width as usize);
    let (icon, style) = match cmd.status {
        None => ('…', "command.running"),
        Some(0) => ('✔', "command.success"),
        _ => ('X', "command.failed"),
    };
    let status_pane = Pane::new_styled(
        Position::new(0, 0),
        Size::new(1, 1),
        vec![icon.to_string()],
        style,
    );
    pane.push_child(status_pane);
    let subwidth: usize = size.width as usize - 1;
    let command_line = text_to_lines(cmd.cmd.clone(), subwidth);
    let pos = Position::new(1, 0);
    let textlen = command_line.len();
    let mut offset = textlen;
    let pane_size = Size::new(subwidth as u16, textlen as u16);
    let prefix = "selected.";
    let mut maxlines = size.height as usize - textlen;
    pane.push_child(Pane::new_styled(
        pos,
        pane_size,
        command_line,
        &format!("{}command", prefix),
    ));
    if cmd.stdout.len() > 0 {
        let mut stdout = text_to_lines(cmd.stdout.clone(), subwidth);
        let pos = Position::new(1, offset as u16);
        let mut textlen = stdout.len();
        if textlen > maxlines {
            stdout = stdout.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        maxlines -= textlen;
        offset += textlen;
        let pane_size = Size::new(subwidth as u16, textlen as u16);
        pane.push_child(Pane::new_styled(
            pos,
            pane_size,
            stdout,
            &format!("{}stdout", prefix),
        ));
    }
    if cmd.stderr.len() > 0 {
        let mut stderr = text_to_lines(cmd.stderr.clone(), subwidth);
        let pos = Position::new(1, offset as u16);
        let mut textlen = stderr.len();
        if textlen > maxlines {
            stderr = stderr.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        let pane_size = Size::new(subwidth as u16, textlen as u16);
        pane.push_child(Pane::new_styled(
            pos,
            pane_size,
            stderr,
            &format!("{}stderr", prefix),
        ));
    }
    pane
}

impl Widget for WeaverStateWidget {
    fn render_children(&self, size: Size) -> Option<Vec<Pane>> {
        let height = size.height as usize;
        let mut ctr: usize = 0;
        let state = self.state.read().unwrap();
        let mut children: Vec<Pane> = vec![];
        let mut i = 0;
        let child_width: usize = match self.selected {
            None => size.width as usize,
            Some(_) => size.width as usize / 2,
        };
        for (_i, cmd) in state.command_history.iter().rev() {
            let selected: bool = match self.selected {
                None => false,
                Some(idx) => idx == i,
            };
            let mut child = render_command_summary(cmd, child_width, 10, selected);
            let offset = child.size.height as usize;

            ctr += offset;
            if ctr >= height {
                let spillover = (ctr - height) as u16;
                let clip_pos = Position::new(0, spillover);
                let clip_size = Size::new(child_width as u16, child.size.height - spillover);
                child = child.clip(clip_pos, clip_size).unwrap();
                child.position = Position::new(0, 0);
                ctr = height;
            }
            let child = child.offset(Position::new(0, (height - ctr) as u16));
            children.push(child);
            if selected {
                let child_pos = Position::new(child_width as u16, 0);
                let child_size = Size::new(size.width - (child_width as u16), size.height);
                let child = render_command_detail(cmd, child_size).offset(child_pos);
                children.push(child);
            }
            if ctr == height {
                break;
            }
            i += 1;
        }
        Some(children)
    }
}

struct WeaverTui {
    log: Shared<Text>,
    input: Shared<Readline>,
    vbox: Shared<Linear>,
    content: Shared<Linear>,
    state: Shared<WeaverState>,
    statew: Shared<WeaverStateWidget>,
    show_debug: bool,
}

impl WeaverTui {
    fn new(state: Arc<RwLock<WeaverState>>) -> WeaverTui {
        let log = shared(Text::new(vec![]));
        let input = shared(Readline::new());
        let state: Shared<WeaverState> = state.into();
        let statew = shared(WeaverStateWidget::new(state.clone()));
        //let dbgdump = shared(DbgDump::new(&state));
        let show_debug = false;
        let mut contentbox = Linear::hbox();
        //contentbox.push(&dbgdump);
        contentbox.push(&statew);
        let content = shared(contentbox);
        let mut mainbox = Linear::vbox();
        mainbox.push(&content);
        mainbox.push(&input);
        let vbox = shared(mainbox);
        WeaverTui {
            log,
            input,
            vbox,
            content,
            state,
            statew,
            show_debug,
        }
    }

    fn toggle_debug(&mut self) {
        let mut content = self.content.write().unwrap();
        match self.show_debug {
            true => {
                self.show_debug = false;
                content.contents.truncate(0);
                content.push(&self.statew);
            }
            false => {
                self.show_debug = true;
                content.contents.truncate(0);
                content.push(&self.statew);
                content.push(&shared(Line::vertical()));
                content.push(&self.log);
            }
        }
    }

    fn submit_input(&mut self) {
        let text = self.input.write().unwrap().finalize();
        if text.len() > 0 {
            self.state
                .write()
                .unwrap()
                .run_command(text.clone())
                .unwrap();
        }
        self.statew.write().unwrap().selected = None;
    }

    fn log_msg(&mut self, msg: &str) {
        let lines: Vec<String> = msg.lines().map(|l| l.to_owned()).collect();
        self.log.write().unwrap().lines.extend(lines);
    }

    fn input(&mut self, key: Key) {
        match key {
            Key::Char('\n') => self.submit_input(),
            Key::Alt('\r') => self.input.write().unwrap().process_key(Key::Char('\n')),
            Key::Up => {
                let mut statew = self.statew.write().unwrap();
                statew.selected = match statew.selected.take() {
                    None => Some(0),
                    Some(i) => Some(i + 1),
                };
                if let Some(cmd) = statew.find_cmd_by_index(statew.selected.unwrap() as usize) {
                    self.input.write().unwrap().set_line(&cmd);
                };
            }
            Key::Down => {
                let mut statew = self.statew.write().unwrap();
                statew.selected = match statew.selected.take() {
                    None => None,
                    Some(0) => None,
                    Some(i) => Some(i - 1),
                };
                match statew.selected {
                    Some(i) => if let Some(cmd) = statew.find_cmd_by_index(i as usize) {
                        self.input.write().unwrap().set_line(&cmd);
                    },
                    None => self.input.write().unwrap().set_line(""),
                }
            }
            k => self.input.write().unwrap().process_key(k),
        }
    }
}

impl App for WeaverTui {
    type UI = Shared<Linear>;
    type MyEvent = WeaverNotification;
    fn widget(&self) -> Self::UI {
        self.vbox.clone()
    }
    fn handle_event(&mut self, event: Event<Self::MyEvent>) -> Result<(), Option<String>> {
        match event {
            Event::InputEvent(i) => match i {
                Input::Key(Key::Esc) => Err(None),
                Input::Key(Key::Alt('d')) => {
                    self.toggle_debug();
                    Ok(())
                }
                Input::Key(k) => {
                    self.input(k);
                    Ok(())
                }
                _ => Ok(()),
            },
            Event::AppEvent(_) => {
                self.log_msg(&format!("{:?}", event));
                Ok(())
            }
        }
    }
    fn style(&self, name: &str) -> (Option<Box<Color>>, Option<Box<Color>>) {
        match name {
            "command" => (None, Some(Box::new(color::Rgb(16, 16, 32)))),
            "stderr" => (None, Some(Box::new(color::Rgb(32, 16, 16)))),
            "stdout" => (None, Some(Box::new(color::Rgb(16, 32, 16)))),
            "input" => (None, Some(Box::new(color::Rgb(32, 32, 32)))),
            "command.running" => (
                Some(Box::new(color::LightYellow)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "command.success" => (
                Some(Box::new(color::LightGreen)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "command.failed" => (Some(Box::new(color::LightRed)), None),
            "selected.command" => (
                Some(Box::new(color::LightWhite)),
                Some(Box::new(color::Rgb(32, 32, 128))),
            ),
            "selected.stderr" => (
                Some(Box::new(color::LightWhite)),
                Some(Box::new(color::Rgb(64, 16, 16))),
            ),
            "selected.stdout" => (
                Some(Box::new(color::LightWhite)),
                Some(Box::new(color::Rgb(16, 64, 16))),
            ),
            "selected.command.running" => (
                Some(Box::new(color::LightYellow)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "selected.command.success" => (
                Some(Box::new(color::LightGreen)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "selected.command.failed" => (Some(Box::new(color::LightRed)), None),
            _ => (None, None),
        }
    }
}

fn main() {
    let be = Backend::new();
    let sender = be.sender.clone();
    let weaver = WeaverClient::new(sender);
    let mut app = WeaverTui::new(weaver.state.clone());
    thread::spawn(move || {
        tokio::run(weaver.map_err(|e| panic!("Client Error: {:#?}", e)));
    });
    app.log_msg("Esc to exit");
    be.run_app(&mut app);
}
<commit_msg>u16 -> usize<commit_after>extern crate futures;
extern crate text_ui;
extern crate tokio;
extern crate tokio_serde_msgpack;
extern crate tokio_uds;
extern crate weaver;

use text_ui::app::App;
use text_ui::backend::{color, Backend, Color};
use text_ui::pane::Pane;
//use text_ui::widget::DbgDump;
use text_ui::widget::Widget;
use text_ui::widget::{shared, Line, Linear, Readline, Shared, Text};
use text_ui::{text_to_lines, Event, Input, Key, Position, Size};

use std::sync::{Arc, RwLock};
use std::thread;

use tokio::prelude::Future;
use weaver::{WeaverClient, WeaverCommand, WeaverNotification, WeaverState};

struct WeaverStateWidget {
    state: Shared<WeaverState>,
    selected: Option<usize>,
}

impl WeaverStateWidget {
    pub fn new(state: Shared<WeaverState>) -> Self {
        let selected = None;
        WeaverStateWidget { state, selected }
    }

    pub fn find_cmd_by_index(&self, i: usize) -> Option<String> {
        let state = self.state.read().unwrap();
        let rv = state
            .command_history
            .iter()
            .rev()
            .skip(i)
            .next()
            .map(|(_, cmd)| cmd.cmd.clone());
        rv
    }
}

fn render_command_summary(
    cmd: &WeaverCommand,
    width: usize,
    maxlines: usize,
    selected: bool,
) -> Pane {
    let mut pane = Pane::new_width(width);
    let (icon, style) = match cmd.status {
        None => ('…', "command.running"),
        Some(0) => ('✔', "command.success"),
        _ => ('X', "command.failed"),
    };
    let status_pane = Pane::new_styled(
        Position::new(0, 0),
        Size::new(1, 1),
        vec![icon.to_string()],
        style,
    );
    pane.push_child(status_pane);
    let subwidth = width - 1;
    let command_line = text_to_lines(cmd.cmd.clone(), subwidth);
    let pos = Position::new(1, 0);
    let textlen = command_line.len();
    let mut offset = textlen;
    let size = Size::new(subwidth, textlen);
    let prefix = match selected {
        true => "selected.",
        false => "",
    };
    pane.push_child(Pane::new_styled(
        pos,
        size,
        command_line,
        &format!("{}command", prefix),
    ));
    if cmd.stdout.len() > 0 {
        let mut stdout = text_to_lines(cmd.stdout.clone(), subwidth);
        let pos = Position::new(1, offset);
        let mut textlen = stdout.len();
        if textlen > maxlines {
            stdout = stdout.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        offset += textlen;
        let size = Size::new(subwidth, textlen);
        pane.push_child(Pane::new_styled(
            pos,
            size,
            stdout,
            &format!("{}stdout", prefix),
        ));
    }
    if cmd.stderr.len() > 0 {
        let mut stderr = text_to_lines(cmd.stderr.clone(), subwidth);
        let pos = Position::new(1, offset);
        let mut textlen = stderr.len();
        if textlen > maxlines {
            stderr = stderr.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        let size = Size::new(subwidth, textlen);
        pane.push_child(Pane::new_styled(
            pos,
            size,
            stderr,
            &format!("{}stderr", prefix),
        ));
    }
    pane
}

fn render_command_detail(cmd: &WeaverCommand, size: Size) -> Pane {
    let mut pane = Pane::new_width(size.width);
    let (icon, style) = match cmd.status {
        None => ('…', "command.running"),
        Some(0) => ('✔', "command.success"),
        _ => ('X', "command.failed"),
    };
    let status_pane = Pane::new_styled(
        Position::new(0, 0),
        Size::new(1, 1),
        vec![icon.to_string()],
        style,
    );
    pane.push_child(status_pane);
    let subwidth: usize = size.width - 1;
    let command_line = text_to_lines(cmd.cmd.clone(), subwidth);
    let pos = Position::new(1, 0);
    let textlen = command_line.len();
    let mut offset = textlen;
    let pane_size = Size::new(subwidth, textlen);
    let prefix = "selected.";
    let mut maxlines = size.height - textlen;
    pane.push_child(Pane::new_styled(
        pos,
        pane_size,
        command_line,
        &format!("{}command", prefix),
    ));
    if cmd.stdout.len() > 0 {
        let mut stdout = text_to_lines(cmd.stdout.clone(), subwidth);
        let pos = Position::new(1, offset);
        let mut textlen = stdout.len();
        if textlen > maxlines {
            stdout = stdout.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        maxlines -= textlen;
        offset += textlen;
        let pane_size = Size::new(subwidth, textlen);
        pane.push_child(Pane::new_styled(
            pos,
            pane_size,
            stdout,
            &format!("{}stdout", prefix),
        ));
    }
    if cmd.stderr.len() > 0 {
        let mut stderr = text_to_lines(cmd.stderr.clone(), subwidth);
        let pos = Position::new(1, offset);
        let mut textlen = stderr.len();
        if textlen > maxlines {
            stderr = stderr.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        let pane_size = Size::new(subwidth, textlen);
        pane.push_child(Pane::new_styled(
            pos,
            pane_size,
            stderr,
            &format!("{}stderr", prefix),
        ));
    }
    pane
}

impl Widget for WeaverStateWidget {
    fn render_children(&self, size: Size) -> Option<Vec<Pane>> {
        let height = size.height;
        let mut ctr: usize = 0;
        let state = self.state.read().unwrap();
        let mut children: Vec<Pane> = vec![];
        let mut i = 0;
        let child_width: usize = match self.selected {
            None => size.width,
            Some(_) => size.width / 2,
        };
        for (_i, cmd) in state.command_history.iter().rev() {
            let selected: bool = match self.selected {
                None => false,
                Some(idx) => idx == i,
            };
            let mut child = render_command_summary(cmd, child_width, 10, selected);
            let offset = child.size.height;

            ctr += offset;
            if ctr >= height {
                let spillover = ctr - height;
                let clip_pos = Position::new(0, spillover);
                let clip_size = Size::new(child_width, child.size.height - spillover);
                child = child.clip(clip_pos, clip_size).unwrap();
                child.position = Position::new(0, 0);
                ctr = height;
            }
            let child = child.offset(Position::new(0, height - ctr));
            children.push(child);
            if selected {
                let child_pos = Position::new(child_width, 0);
                let child_size = Size::new(size.width - child_width, size.height);
                let child = render_command_detail(cmd, child_size).offset(child_pos);
                children.push(child);
            }
            if ctr == height {
                break;
            }
            i += 1;
        }
        Some(children)
    }
}

struct WeaverTui {
    log: Shared<Text>,
    input: Shared<Readline>,
    vbox: Shared<Linear>,
    content: Shared<Linear>,
    state: Shared<WeaverState>,
    statew: Shared<WeaverStateWidget>,
    show_debug: bool,
}

impl WeaverTui {
    fn new(state: Arc<RwLock<WeaverState>>) -> WeaverTui {
        let log = shared(Text::new(vec![]));
        let input = shared(Readline::new());
        let state: Shared<WeaverState> = state.into();
        let statew = shared(WeaverStateWidget::new(state.clone()));
        //let dbgdump = shared(DbgDump::new(&state));
        let show_debug = false;
        let mut contentbox = Linear::hbox();
        //contentbox.push(&dbgdump);
        contentbox.push(&statew);
        let content = shared(contentbox);
        let mut mainbox = Linear::vbox();
        mainbox.push(&content);
        mainbox.push(&input);
        let vbox = shared(mainbox);
        WeaverTui {
            log,
            input,
            vbox,
            content,
            state,
            statew,
            show_debug,
        }
    }

    fn toggle_debug(&mut self) {
        let mut content = self.content.write().unwrap();
        match self.show_debug {
            true => {
                self.show_debug = false;
                content.contents.truncate(0);
                content.push(&self.statew);
            }
            false => {
                self.show_debug = true;
                content.contents.truncate(0);
                content.push(&self.statew);
                content.push(&shared(Line::vertical()));
                content.push(&self.log);
            }
        }
    }

    fn submit_input(&mut self) {
        let text = self.input.write().unwrap().finalize();
        if text.len() > 0 {
            self.state
                .write()
                .unwrap()
                .run_command(text.clone())
                .unwrap();
        }
        self.statew.write().unwrap().selected = None;
    }

    fn log_msg(&mut self, msg: &str) {
        let lines: Vec<String> = msg.lines().map(|l| l.to_owned()).collect();
        self.log.write().unwrap().lines.extend(lines);
    }

    fn input(&mut self, key: Key) {
        match key {
            Key::Char('\n') => self.submit_input(),
            Key::Alt('\r') => self.input.write().unwrap().process_key(Key::Char('\n')),
            Key::Up => {
                let mut statew = self.statew.write().unwrap();
                statew.selected = match statew.selected.take() {
                    None => Some(0),
                    Some(i) => Some(i + 1),
                };
                if let Some(cmd) = statew.find_cmd_by_index(statew.selected.unwrap()) {
                    self.input.write().unwrap().set_line(&cmd);
                };
            }
            Key::Down => {
                let mut statew = self.statew.write().unwrap();
                statew.selected = match statew.selected.take() {
                    None => None,
                    Some(0) => None,
                    Some(i) => Some(i - 1),
                };
                match statew.selected {
                    Some(i) => if let Some(cmd) = statew.find_cmd_by_index(i) {
                        self.input.write().unwrap().set_line(&cmd);
                    },
                    None => self.input.write().unwrap().set_line(""),
                }
            }
            k => self.input.write().unwrap().process_key(k),
        }
    }
}

impl App for WeaverTui {
    type UI = Shared<Linear>;
    type MyEvent = WeaverNotification;
    fn widget(&self) -> Self::UI {
        self.vbox.clone()
    }
    fn handle_event(&mut self, event: Event<Self::MyEvent>) -> Result<(), Option<String>> {
        match event {
            Event::InputEvent(i) => match i {
                Input::Key(Key::Esc) => Err(None),
                Input::Key(Key::Alt('d')) => {
                    self.toggle_debug();
                    Ok(())
                }
                Input::Key(k) => {
                    self.input(k);
                    Ok(())
                }
                _ => Ok(()),
            },
            Event::AppEvent(_) => {
                self.log_msg(&format!("{:?}", event));
                Ok(())
            }
        }
    }
    fn style(&self, name: &str) -> (Option<Box<Color>>, Option<Box<Color>>) {
        match name {
            "command" => (None, Some(Box::new(color::Rgb(16, 16, 32)))),
            "stderr" => (None, Some(Box::new(color::Rgb(32, 16, 16)))),
            "stdout" => (None, Some(Box::new(color::Rgb(16, 32, 16)))),
            "input" => (None, Some(Box::new(color::Rgb(32, 32, 32)))),
            "command.running" => (
                Some(Box::new(color::LightYellow)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "command.success" => (
                Some(Box::new(color::LightGreen)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "command.failed" => (Some(Box::new(color::LightRed)), None),
            "selected.command" => (
                Some(Box::new(color::LightWhite)),
                Some(Box::new(color::Rgb(32, 32, 128))),
            ),
            "selected.stderr" => (
                Some(Box::new(color::LightWhite)),
                Some(Box::new(color::Rgb(64, 16, 16))),
            ),
            "selected.stdout" => (
                Some(Box::new(color::LightWhite)),
                Some(Box::new(color::Rgb(16, 64, 16))),
            ),
            "selected.command.running" => (
                Some(Box::new(color::LightYellow)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "selected.command.success" => (
                Some(Box::new(color::LightGreen)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "selected.command.failed" => (Some(Box::new(color::LightRed)), None),
            _ => (None, None),
        }
    }
}

fn main() {
    let be = Backend::new();
    let sender = be.sender.clone();
    let weaver = WeaverClient::new(sender);
    let mut app = WeaverTui::new(weaver.state.clone());
    thread::spawn(move || {
        tokio::run(weaver.map_err(|e| panic!("Client Error: {:#?}", e)));
    });
    app.log_msg("Esc to exit");
    be.run_app(&mut app);
}
<|endoftext|>