<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// * (c) kwantam <kwantam@gmail.com>
// *     * 20150507 ~ added big_ routines to prevent overflow when num > 2^63
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use std::mem::swap;

pub fn gcd(mut a: u64, mut b: u64) -> u64 {
    while b > 0 {
        a %= b;
        swap(&mut a, &mut b);
    }
    a
}

pub(crate) trait Arithmetic: Copy + Sized {
    type I: Copy + Sized + Eq;

    fn new(m: u64) -> Self;
    fn modulus(&self) -> u64;
    fn from_u64(&self, n: u64) -> Self::I;
    fn to_u64(&self, n: Self::I) -> u64;
    fn add(&self, a: Self::I, b: Self::I) -> Self::I;
    fn mul(&self, a: Self::I, b: Self::I) -> Self::I;

    fn pow(&self, mut a: Self::I, mut b: u64) -> Self::I {
        let (_a, _b) = (a, b);
        let mut result = self.one();
        while b > 0 {
            if b & 1 != 0 {
                result = self.mul(result, a);
            }
            a = self.mul(a, a);
            b >>= 1;
        }

        // Check that r (reduced back to the usual representation) equals
        //  a^b % n, unless the latter computation overflows
        // Temporarily commented-out, as there u64::checked_pow is not available
        //  on the minimum supported Rust version, nor is an appropriate method
        //  for compiling the check conditionally.
        //debug_assert!(self
        //    .to_u64(_a)
        //    .checked_pow(_b as u32)
        //    .map(|r| r % self.modulus() == self.to_u64(result))
        //    .unwrap_or(true));

        result
    }

    fn one(&self) -> Self::I {
        self.from_u64(1)
    }
    fn minus_one(&self) -> Self::I {
        self.from_u64(self.modulus() - 1)
    }
    fn zero(&self) -> Self::I {
        self.from_u64(0)
    }
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct Montgomery {
    a: u64,
    n: u64,
}

impl Montgomery {
    /// computes x/R mod n efficiently
    fn reduce(&self, x: u64) -> u64 {
        // TODO: optimiiiiiiise
        let Montgomery { a, n } = self;
        let t = x.wrapping_mul(*a);
        let nt = (*n as u128) * (t as u128);
        let y = ((x as u128 + nt) >> 64) as u64;
        if y >= *n {
            y - n
        } else {
            y
        }
    }
}

impl Arithmetic for Montgomery {
    // Montgomery transform, R=2⁶⁴
    // Provides fast arithmetic mod n (n odd, u64)
    type I = u64;

    fn new(n: u64) -> Self {
        let a = inv_mod_u64(n).wrapping_neg();
        debug_assert_eq!(n.wrapping_mul(a), 1_u64.wrapping_neg());
        Montgomery { a, n }
    }

    fn modulus(&self) -> u64 {
        self.n
    }

    fn from_u64(&self, x: u64) -> Self::I {
        // TODO: optimise!
        assert!(x < self.n);
        let r = (((x as u128) << 64) % self.n as u128) as u64;
        debug_assert_eq!(x, self.to_u64(r));
        r
    }

    fn to_u64(&self, n: Self::I) -> u64 {
        self.reduce(n)
    }

    fn add(&self, a: Self::I, b: Self::I) -> Self::I {
        let r = a + b;

        // Check that r (reduced back to the usual representation) equals
        // a+b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2 = (((a_r as u128) + (b_r as u128)) % (self.n as u128)) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} + {} = [{}] + [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }

    fn mul(&self, a: Self::I, b: Self::I) -> Self::I {
        let r = self.reduce(a.wrapping_mul(b));

        // Check that r (reduced back to the usual representation) equals
        // a*b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2 = (((a_r as u128) * (b_r as u128)) % (self.n as u128)) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} * {} = [{}] * [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }
}

// extended Euclid algorithm
// precondition: a is odd
pub(crate) fn inv_mod_u64(a: u64) -> u64 {
    assert!(a % 2 == 1);
    let mut t = 0u64;
    let mut newt = 1u64;
    let mut r = 0u64;
    let mut newr = a;

    while newr != 0 {
        let quot = if r == 0 {
            // special case when we're just starting out
            // This works because we know that
            // a does not divide 2^64, so floor(2^64 / a) == floor((2^64-1) / a);
            std::u64::MAX
        } else {
            r
        } / newr;

        let newtp = t.wrapping_sub(quot.wrapping_mul(newt));
        t = newt;
        newt = newtp;

        let newrp = r.wrapping_sub(quot.wrapping_mul(newr));
        r = newr;
        newr = newrp;
    }

    assert_eq!(r, 1);
    t
}

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

    #[test]
    fn test_inverter() {
        // All odd integers from 1 to 20 000
        let mut test_values = (0..10_000u64).map(|i| 2 * i + 1);

        assert!(test_values.all(|x| x.wrapping_mul(inv_mod_u64(x)) == 1));
    }

    #[test]
    fn test_montgomery_add() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    println!("{n:?}, {x:?}, {y:?}", n = n, x = x, y = y);
                    assert_eq!((x + y) % n, m.to_u64(m.add(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_montgomery_mult() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    assert_eq!((x * y) % n, m.to_u64(m.mul(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_montgomery_roundtrip() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::new(n);
            for x in 0..n {
                let x_ = m.from_u64(x);
                assert_eq!(x, m.to_u64(x_));
            }
        }
    }
}
<commit_msg>factor::numeric::Montgomery: Fix overflow bug<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// * (c) kwantam <kwantam@gmail.com>
// *     * 20150507 ~ added big_ routines to prevent overflow when num > 2^63
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use std::mem::swap;

pub fn gcd(mut a: u64, mut b: u64) -> u64 {
    while b > 0 {
        a %= b;
        swap(&mut a, &mut b);
    }
    a
}

pub(crate) trait Arithmetic: Copy + Sized {
    type I: Copy + Sized + Eq;

    fn new(m: u64) -> Self;
    fn modulus(&self) -> u64;
    fn from_u64(&self, n: u64) -> Self::I;
    fn to_u64(&self, n: Self::I) -> u64;
    fn add(&self, a: Self::I, b: Self::I) -> Self::I;
    fn mul(&self, a: Self::I, b: Self::I) -> Self::I;

    fn pow(&self, mut a: Self::I, mut b: u64) -> Self::I {
        let (_a, _b) = (a, b);
        let mut result = self.one();
        while b > 0 {
            if b & 1 != 0 {
                result = self.mul(result, a);
            }
            a = self.mul(a, a);
            b >>= 1;
        }

        // Check that r (reduced back to the usual representation) equals
        //  a^b % n, unless the latter computation overflows
        // Temporarily commented-out, as there u64::checked_pow is not available
        //  on the minimum supported Rust version, nor is an appropriate method
        //  for compiling the check conditionally.
        //debug_assert!(self
        //    .to_u64(_a)
        //    .checked_pow(_b as u32)
        //    .map(|r| r % self.modulus() == self.to_u64(result))
        //    .unwrap_or(true));

        result
    }

    fn one(&self) -> Self::I {
        self.from_u64(1)
    }
    fn minus_one(&self) -> Self::I {
        self.from_u64(self.modulus() - 1)
    }
    fn zero(&self) -> Self::I {
        self.from_u64(0)
    }
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct Montgomery {
    a: u64,
    n: u64,
}

impl Montgomery {
    /// computes x/R mod n efficiently
    fn reduce(&self, x: u128) -> u64 {
        debug_assert!(x < (self.n as u128) << 64);
        // TODO: optimiiiiiiise
        let Montgomery { a, n } = self;
        let m = (x as u64).wrapping_mul(*a);
        let nm = (*n as u128) * (m as u128);
        let (xnm, overflow) = (x as u128).overflowing_add(nm); // x + n*m
        debug_assert_eq!(xnm % (1 << 64), 0);

        if !overflow {
            let y = (xnm >> 64) as u64 + overflow as u64; // (x + n*m) / R
            if y >= *n {
                y - n
            } else {
                y
            }
        } else {
            // y = (2¹²⁸ + xnm)/2⁶⁴ - n = xnm/2⁶⁴ + (2⁶⁴ - n)
            (xnm >> 64) as u64 + n.wrapping_neg()
        }
    }
}

impl Arithmetic for Montgomery {
    // Montgomery transform, R=2⁶⁴
    // Provides fast arithmetic mod n (n odd, u64)
    type I = u64;

    fn new(n: u64) -> Self {
        let a = inv_mod_u64(n).wrapping_neg();
        debug_assert_eq!(n.wrapping_mul(a), 1_u64.wrapping_neg());
        Montgomery { a, n }
    }

    fn modulus(&self) -> u64 {
        self.n
    }

    fn from_u64(&self, x: u64) -> Self::I {
        // TODO: optimise!
        assert!(x < self.n);
        let r = (((x as u128) << 64) % self.n as u128) as u64;
        debug_assert_eq!(x, self.to_u64(r));
        r
    }

    fn to_u64(&self, n: Self::I) -> u64 {
        self.reduce(n as u128)
    }

    fn add(&self, a: Self::I, b: Self::I) -> Self::I {
        let r = a + b;

        // Check that r (reduced back to the usual representation) equals
        // a+b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2 = (((a_r as u128) + (b_r as u128)) % (self.n as u128)) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} + {} = [{}] + [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }

    fn mul(&self, a: Self::I, b: Self::I) -> Self::I {
        let r = self.reduce((a as u128) * (b as u128));

        // Check that r (reduced back to the usual representation) equals
        // a*b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2 = (((a_r as u128) * (b_r as u128)) % (self.n as u128)) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} * {} = [{}] * [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }
}

// extended Euclid algorithm
// precondition: a is odd
pub(crate) fn inv_mod_u64(a: u64) -> u64 {
    assert!(a % 2 == 1);
    let mut t = 0u64;
    let mut newt = 1u64;
    let mut r = 0u64;
    let mut newr = a;

    while newr != 0 {
        let quot = if r == 0 {
            // special case when we're just starting out
            // This works because we know that
            // a does not divide 2^64, so floor(2^64 / a) == floor((2^64-1) / a);
            std::u64::MAX
        } else {
            r
        } / newr;

        let newtp = t.wrapping_sub(quot.wrapping_mul(newt));
        t = newt;
        newt = newtp;

        let newrp = r.wrapping_sub(quot.wrapping_mul(newr));
        r = newr;
        newr = newrp;
    }

    assert_eq!(r, 1);
    t
}

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

    #[test]
    fn test_inverter() {
        // All odd integers from 1 to 20 000
        let mut test_values = (0..10_000u64).map(|i| 2 * i + 1);

        assert!(test_values.all(|x| x.wrapping_mul(inv_mod_u64(x)) == 1));
    }

    #[test]
    fn test_montgomery_add() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    println!("{n:?}, {x:?}, {y:?}", n = n, x = x, y = y);
                    assert_eq!((x + y) % n, m.to_u64(m.add(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_montgomery_mult() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    assert_eq!((x * y) % n, m.to_u64(m.mul(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_montgomery_roundtrip() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::new(n);
            for x in 0..n {
                let x_ = m.from_u64(x);
                assert_eq!(x, m.to_u64(x_));
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Simple daemon
//!
//! This daemon will start web server, track new packages and build them

use std::{env, thread};
use std::process::exit;
use std::fs::File;
use std::io::Write;
use std::time::Duration;
use std::path::PathBuf;
use libc::fork;
use time;
use DocBuilderOptions;
use DocBuilder;
use utils::{update_sources, update_release_activity};

const DAEMON_PID_FILE_PATH: &'static str = "/var/run/cratesfyi.pid";

pub fn start_daemon() {
    // first check required environment variables
    for v in ["CRATESFYI_PREFIX",
              "CRATESFYI_PREFIX",
              "CRATESFYI_GITHUB_USERNAME",
              "CRATESFYI_GITHUB_ACCESSTOKEN"]
        .iter() {
        env::var(v).expect("Environment variable not found");
    }

    // check paths once
    opts().check_paths().unwrap();

    // fork the process
    let pid = unsafe { fork() };
    if pid > 0 {
        let mut file = File::create(DAEMON_PID_FILE_PATH).expect("Failed to create pid file");
        writeln!(&mut file, "{}", pid).expect("Failed to write pid");

        info!("cratesfyi {} daemon started on: {}", ::BUILD_VERSION, pid);
        exit(0);
    }

    // check new crates every 5 minutes
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(300));
            debug!("Checking new crates");
            let doc_builder = DocBuilder::new(opts());
            if let Err(e) = doc_builder.get_new_crates() {
                error!("Failed to get new crates: {}", e);
            }
        }
    });

    // build new crates every 3 minutes
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(180));

            let mut opts = opts();
            opts.skip_if_exists = true;

            // check lock file
            if opts.prefix.join("cratesfyi.lock").exists() {
                warn!("Lock file exits, skipping building new crates");
                continue;
            }

            // update index
            if let Err(e) = update_sources() {
                error!("Failed to update sources: {}", e);
                continue;
            }

            let mut doc_builder = DocBuilder::new(opts);
            if let Err(e) = doc_builder.load_cache() {
                error!("Failed to load cache: {}", e);
                continue;
            }

            debug!("Building new crates");
            if let Err(e) = doc_builder.build_packages_queue() {
                error!("Failed build new crates: {}", e);
            }

            if let Err(e) = doc_builder.save_cache() {
                error!("Failed to save cache: {}", e);
            }
        }
    });

    // update release activity everyday at 02:00
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(60));
            let now = time::now();
            if now.tm_hour == 2 && now.tm_min == 0 {
                info!("Updating release activity");
                if let Err(e) = update_release_activity() {
                    error!("Failed to update release activity: {}", e);
                }
            }
        }
    });

    // TODO: update ssl certificate every 3 months

    // at least start web server
    info!("Starting web server");
    ::start_web_server(None);
}

fn opts() -> DocBuilderOptions {
    let prefix = PathBuf::from(env::var("CRATESFYI_PREFIX")
                               .expect("CRATESFYI_PREFIX environment variable not found"));
    DocBuilderOptions::from_prefix(prefix)
}
<commit_msg>Update search index every 3 hours<commit_after>//! Simple daemon
//!
//! This daemon will start web server, track new packages and build them

use std::{env, thread};
use std::process::exit;
use std::fs::File;
use std::io::Write;
use std::time::Duration;
use std::path::PathBuf;
use libc::fork;
use time;
use DocBuilderOptions;
use DocBuilder;
use utils::{update_sources, update_release_activity};
use db::{connect_db, update_search_index};

const DAEMON_PID_FILE_PATH: &'static str = "/var/run/cratesfyi.pid";

pub fn start_daemon() {
    // first check required environment variables
    for v in ["CRATESFYI_PREFIX",
              "CRATESFYI_PREFIX",
              "CRATESFYI_GITHUB_USERNAME",
              "CRATESFYI_GITHUB_ACCESSTOKEN"]
        .iter() {
        env::var(v).expect("Environment variable not found");
    }

    // check paths once
    opts().check_paths().unwrap();

    // fork the process
    let pid = unsafe { fork() };
    if pid > 0 {
        let mut file = File::create(DAEMON_PID_FILE_PATH).expect("Failed to create pid file");
        writeln!(&mut file, "{}", pid).expect("Failed to write pid");

        info!("cratesfyi {} daemon started on: {}", ::BUILD_VERSION, pid);
        exit(0);
    }

    // check new crates every 5 minutes
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(300));
            debug!("Checking new crates");
            let doc_builder = DocBuilder::new(opts());
            if let Err(e) = doc_builder.get_new_crates() {
                error!("Failed to get new crates: {}", e);
            }
        }
    });

    // build new crates every 3 minutes
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(180));

            let mut opts = opts();
            opts.skip_if_exists = true;

            // check lock file
            if opts.prefix.join("cratesfyi.lock").exists() {
                warn!("Lock file exits, skipping building new crates");
                continue;
            }

            // update index
            if let Err(e) = update_sources() {
                error!("Failed to update sources: {}", e);
                continue;
            }

            let mut doc_builder = DocBuilder::new(opts);
            if let Err(e) = doc_builder.load_cache() {
                error!("Failed to load cache: {}", e);
                continue;
            }

            debug!("Building new crates");
            if let Err(e) = doc_builder.build_packages_queue() {
                error!("Failed build new crates: {}", e);
            }

            if let Err(e) = doc_builder.save_cache() {
                error!("Failed to save cache: {}", e);
            }
        }
    });

    // update release activity everyday at 02:00
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(60));
            let now = time::now();
            if now.tm_hour == 2 && now.tm_min == 0 {
                info!("Updating release activity");
                if let Err(e) = update_release_activity() {
                    error!("Failed to update release activity: {}", e);
                }
            }
        }
    });

    // update search index every 3 hours
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(60*60*3));
            let conn = connect_db().expect("Failed to connect database");
            if let Err(e) = update_search_index(&conn) {
                error!("Failed to update search index: {}", e);
            }
        }
    });

    // TODO: update ssl certificate every 3 months

    // at least start web server
    info!("Starting web server");
    ::start_web_server(None);
}

fn opts() -> DocBuilderOptions {
    let prefix = PathBuf::from(env::var("CRATESFYI_PREFIX")
                               .expect("CRATESFYI_PREFIX environment variable not found"));
    DocBuilderOptions::from_prefix(prefix)
}
<|endoftext|>


<commit_before>use countminsketch::CountMinSketch;
use std::collections::{BTreeSet, HashMap};
use std::collections::hash_map::Entry;
use std::cmp::Ordering;
use std::fmt;
use std::hash::Hash;
use std::rc::Rc;

struct TreeEntry<T>
where
    T: Eq + Ord,
{
    obj: Rc<T>,
    n: usize,
}

impl<T> Clone for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn clone(&self) -> Self {
        TreeEntry {
            obj: Rc::clone(&self.obj),
            n: self.n,
        }
    }
}

impl<T> PartialEq for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn eq(&self, other: &TreeEntry<T>) -> bool {
        self.obj == other.obj
    }
}

impl<T> Eq for TreeEntry<T>
where
    T: Eq + Ord,
{
}

impl<T> PartialOrd for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn partial_cmp(&self, other: &TreeEntry<T>) -> Option<Ordering> {
        match self.n.cmp(&other.n) {
            Ordering::Greater => Some(Ordering::Greater),
            Ordering::Less => Some(Ordering::Less),
            Ordering::Equal => Some(self.obj.cmp(&other.obj)),
        }
    }
}

impl<T> Ord for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn cmp(&self, other: &TreeEntry<T>) -> Ordering {
        match self.n.cmp(&other.n) {
            Ordering::Greater => Ordering::Greater,
            Ordering::Less => Ordering::Less,
            Ordering::Equal => self.obj.cmp(&other.obj),
        }
    }
}

/// Top-K implementation.
///
/// This data structure keeps the `k` most frequent data points of a stream.
#[derive(Clone)]
pub struct TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    cms: CountMinSketch,
    obj2count: HashMap<Rc<T>, usize>,
    tree: BTreeSet<TreeEntry<T>>,
    k: usize,
}

impl<T> TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    /// Create new Top-K data structure.
    ///
    /// - `k`: number of elements to remember
    /// - `cms`: CountMinSketch that is used for guessing the frequency of elements not currently
    ///   hold. Refer to its documentation about parameter selection.
    ///
    /// Panics if `k == 0`.
    pub fn new(k: usize, cms: CountMinSketch) -> TopK<T> {
        assert!(k > 0);
        TopK {
            cms: cms,
            obj2count: HashMap::new(),
            tree: BTreeSet::new(),
            k: k,
        }
    }

    /// Number of data points to remember.
    pub fn k(&self) -> usize {
        return self.k;
    }

    /// Observe a data point.
    pub fn add(&mut self, obj: T) {
        // create Rc so we can use obj in multiple data structs
        let rc = Rc::new(obj);

        // always increase CountMinSketch counts
        self.cms.add(&rc);

        // check if entry is in top K
        let size = self.obj2count.len();
        let mut update_obj2count: Option<(usize, Rc<T>)> = None;
        match self.obj2count.entry(Rc::clone(&rc)) {
            Entry::Occupied(mut o) => {
                // it is => increase exact counter
                let mut n = o.get_mut();
                *n += 1;

                // update tree data
                let mut entry = TreeEntry {
                    obj: Rc::clone(&rc),
                    n: *n - 1,
                };
                self.tree.remove(&entry);
                entry.n += 1;
                self.tree.insert(entry);
            }
            Entry::Vacant(v) => {
                // it's not => check capicity
                if size < self.k {
                    // space left => add to top k
                    v.insert(1);
                    self.tree.insert(TreeEntry {
                        obj: Rc::clone(&rc),
                        n: 1,
                    });
                } else {
                    // not enough space => check if it would be a top k element
                    let count = self.cms.query_point(&rc);
                    let min: TreeEntry<T> = (*self.tree.iter().next().unwrap()).clone();
                    if count > min.n {
                        // => kick out minimal element of top k
                        self.tree.remove(&min);
                        self.tree.insert(TreeEntry {
                            obj: Rc::clone(&rc),
                            n: count,
                        });

                        // delay obj2count changes, because it is currently borrowed
                        update_obj2count = Some((count, min.obj));
                    }
                }
            }
        }

        // trick the borrow checker
        if let Some((count, obj)) = update_obj2count {
            self.obj2count.insert(rc, count);
            self.obj2count.remove(&obj);
        }
    }

    /// Returns top-k values.
    ///
    /// The result may contain less than `k` values if less than `k` unique data points where
    /// observed.
    pub fn values(&self) -> Vec<T> {
        self.tree.iter().map(|x| (*x.obj).clone()).collect()
    }

    /// Check whether the sampler is empty (i.e. no data points observer so far)
    pub fn is_empty(&self) -> bool {
        self.obj2count.is_empty()
    }

    /// Clear sampler, so that it behaves like it never observed any data points.
    pub fn clear(&mut self) {
        self.cms.clear();
        self.tree.clear();
        self.obj2count.clear();
    }
}

impl<T> fmt::Debug for TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "TopK {{ k: {} }}", self.k)
    }
}

impl<T> Extend<T> for TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    fn extend<S: IntoIterator<Item = T>>(&mut self, iter: S) {
        for elem in iter {
            self.add(elem);
        }
    }
}

#[cfg(test)]
mod tests {
    use countminsketch::CountMinSketch;
    use super::TopK;

    #[test]
    fn getter() {
        let cms = CountMinSketch::<usize>::with_params(10, 20);
        let tk: TopK<usize> = TopK::new(2, cms);

        assert_eq!(tk.k(), 2);
    }

    #[test]
    fn add_1() {
        let cms = CountMinSketch::<usize>::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        tk.add(1);
        assert_eq!(tk.values(), vec![1]);
    }

    #[test]
    fn add_2_same() {
        let cms = CountMinSketch::<usize>::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        tk.add(1);
        tk.add(1);
        assert_eq!(tk.values(), vec![1]);
    }

    #[test]
    fn add_2_different() {
        let cms = CountMinSketch::<usize>::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        tk.add(1);
        tk.add(2);
        assert_eq!(tk.values(), vec![1, 2]);
    }

    #[test]
    fn add_n() {
        let cms = CountMinSketch::<usize>::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        for i in 0..5 {
            tk.add(i);
        }
        for _ in 0..100 {
            tk.add(99);
        }
        for _ in 0..100 {
            tk.add(100);
        }
        for i in 0..5 {
            tk.add(i);
        }
        assert_eq!(tk.values(), vec![99, 100]);
    }

    #[test]
    fn is_empty() {
        let cms = CountMinSketch::<usize>::with_params(10, 20);
        let mut tk = TopK::new(2, cms);
        assert_eq!(tk.is_empty(), true);

        tk.add(0);
        assert_eq!(tk.is_empty(), false);
    }

    #[test]
    fn clear() {
        let cms = CountMinSketch::<usize>::with_params(10, 20);
        let mut tk = TopK::new(2, cms);
        tk.add(0);

        tk.clear();
        assert_eq!(tk.is_empty(), true);

        tk.add(1);
        assert_eq!(tk.values(), vec![1]);
    }

    #[test]
    fn clone() {
        let cms = CountMinSketch::<usize>::with_params(10, 20);
        let mut tk1 = TopK::new(2, cms);
        tk1.add(0);

        let mut tk2 = tk1.clone();
        tk2.add(1);

        assert_eq!(tk1.values(), vec![0]);
        assert_eq!(tk2.values(), vec![0, 1]);
    }

    #[test]
    fn debug() {
        let cms = CountMinSketch::with_params(10, 20);
        let tk: TopK<usize> = TopK::new(2, cms);
        assert_eq!(format!("{:?}", tk), "TopK { k: 2 }");
    }

    #[test]
    fn extend() {
        let cms = CountMinSketch::with_params(10, 10);
        let mut tk = TopK::new(2, cms);

        tk.extend(vec![0, 1]);
        assert_eq!(tk.values(), vec![0, 1]);
    }
}
<commit_msg>topk: simplify code<commit_after>use countminsketch::CountMinSketch;
use std::collections::{BTreeSet, HashMap};
use std::collections::hash_map::Entry;
use std::cmp::Ordering;
use std::fmt;
use std::hash::Hash;
use std::rc::Rc;

struct TreeEntry<T>
where
    T: Eq + Ord,
{
    obj: Rc<T>,
    n: usize,
}

impl<T> Clone for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn clone(&self) -> Self {
        TreeEntry {
            obj: Rc::clone(&self.obj),
            n: self.n,
        }
    }
}

impl<T> PartialEq for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn eq(&self, other: &TreeEntry<T>) -> bool {
        self.obj == other.obj
    }
}

impl<T> Eq for TreeEntry<T>
where
    T: Eq + Ord,
{
}

impl<T> PartialOrd for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn partial_cmp(&self, other: &TreeEntry<T>) -> Option<Ordering> {
        match self.n.cmp(&other.n) {
            Ordering::Greater => Some(Ordering::Greater),
            Ordering::Less => Some(Ordering::Less),
            Ordering::Equal => Some(self.obj.cmp(&other.obj)),
        }
    }
}

impl<T> Ord for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn cmp(&self, other: &TreeEntry<T>) -> Ordering {
        self.partial_cmp(other).unwrap()
    }
}

/// Top-K implementation.
///
/// This data structure keeps the `k` most frequent data points of a stream.
#[derive(Clone)]
pub struct TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    cms: CountMinSketch,
    obj2count: HashMap<Rc<T>, usize>,
    tree: BTreeSet<TreeEntry<T>>,
    k: usize,
}

impl<T> TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    /// Create new Top-K data structure.
    ///
    /// - `k`: number of elements to remember
    /// - `cms`: CountMinSketch that is used for guessing the frequency of elements not currently
    ///   hold. Refer to its documentation about parameter selection.
    ///
    /// Panics if `k == 0`.
    pub fn new(k: usize, cms: CountMinSketch) -> TopK<T> {
        assert!(k > 0);
        TopK {
            cms: cms,
            obj2count: HashMap::new(),
            tree: BTreeSet::new(),
            k: k,
        }
    }

    /// Number of data points to remember.
    pub fn k(&self) -> usize {
        return self.k;
    }

    /// Observe a data point.
    pub fn add(&mut self, obj: T) {
        // create Rc so we can use obj in multiple data structs
        let rc = Rc::new(obj);

        // always increase CountMinSketch counts
        self.cms.add(&rc);

        // check if entry is in top K
        let size = self.obj2count.len();
        let mut update_obj2count: Option<(usize, Rc<T>)> = None;
        match self.obj2count.entry(Rc::clone(&rc)) {
            Entry::Occupied(mut o) => {
                // it is => increase exact counter
                let mut n = o.get_mut();
                *n += 1;

                // update tree data
                let mut entry = TreeEntry {
                    obj: Rc::clone(&rc),
                    n: *n - 1,
                };
                self.tree.remove(&entry);
                entry.n += 1;
                self.tree.insert(entry);
            }
            Entry::Vacant(v) => {
                // it's not => check capicity
                if size < self.k {
                    // space left => add to top k
                    v.insert(1);
                    self.tree.insert(TreeEntry {
                        obj: Rc::clone(&rc),
                        n: 1,
                    });
                } else {
                    // not enough space => check if it would be a top k element
                    let count = self.cms.query_point(&rc);
                    let min: TreeEntry<T> = (*self.tree.iter().next().unwrap()).clone();
                    if count > min.n {
                        // => kick out minimal element of top k
                        self.tree.remove(&min);
                        self.tree.insert(TreeEntry {
                            obj: Rc::clone(&rc),
                            n: count,
                        });

                        // delay obj2count changes, because it is currently borrowed
                        update_obj2count = Some((count, min.obj));
                    }
                }
            }
        }

        // trick the borrow checker
        if let Some((count, obj)) = update_obj2count {
            self.obj2count.insert(rc, count);
            self.obj2count.remove(&obj);
        }
    }

    /// Returns top-k values.
    ///
    /// The result may contain less than `k` values if less than `k` unique data points where
    /// observed.
    pub fn values(&self) -> Vec<T> {
        self.tree.iter().map(|x| (*x.obj).clone()).collect()
    }

    /// Check whether the sampler is empty (i.e. no data points observer so far)
    pub fn is_empty(&self) -> bool {
        self.obj2count.is_empty()
    }

    /// Clear sampler, so that it behaves like it never observed any data points.
    pub fn clear(&mut self) {
        self.cms.clear();
        self.tree.clear();
        self.obj2count.clear();
    }
}

impl<T> fmt::Debug for TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "TopK {{ k: {} }}", self.k)
    }
}

impl<T> Extend<T> for TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    fn extend<S: IntoIterator<Item = T>>(&mut self, iter: S) {
        for elem in iter {
            self.add(elem);
        }
    }
}

#[cfg(test)]
mod tests {
    use countminsketch::CountMinSketch;
    use super::TopK;

    #[test]
    fn getter() {
        let cms = CountMinSketch::with_params(10, 20);
        let tk: TopK<usize> = TopK::new(2, cms);

        assert_eq!(tk.k(), 2);
    }

    #[test]
    fn add_1() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        tk.add(1);
        assert_eq!(tk.values(), vec![1]);
    }

    #[test]
    fn add_2_same() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        tk.add(1);
        tk.add(1);
        assert_eq!(tk.values(), vec![1]);
    }

    #[test]
    fn add_2_different() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        tk.add(1);
        tk.add(2);
        assert_eq!(tk.values(), vec![1, 2]);
    }

    #[test]
    fn add_n() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        for i in 0..5 {
            tk.add(i);
        }
        for _ in 0..100 {
            tk.add(99);
        }
        for _ in 0..100 {
            tk.add(100);
        }
        for i in 0..5 {
            tk.add(i);
        }
        assert_eq!(tk.values(), vec![99, 100]);
    }

    #[test]
    fn is_empty() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);
        assert_eq!(tk.is_empty(), true);

        tk.add(0);
        assert_eq!(tk.is_empty(), false);
    }

    #[test]
    fn clear() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);
        tk.add(0);

        tk.clear();
        assert_eq!(tk.is_empty(), true);

        tk.add(1);
        assert_eq!(tk.values(), vec![1]);
    }

    #[test]
    fn clone() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk1 = TopK::new(2, cms);
        tk1.add(0);

        let mut tk2 = tk1.clone();
        tk2.add(1);

        assert_eq!(tk1.values(), vec![0]);
        assert_eq!(tk2.values(), vec![0, 1]);
    }

    #[test]
    fn debug() {
        let cms = CountMinSketch::with_params(10, 20);
        let tk: TopK<usize> = TopK::new(2, cms);
        assert_eq!(format!("{:?}", tk), "TopK { k: 2 }");
    }

    #[test]
    fn extend() {
        let cms = CountMinSketch::with_params(10, 10);
        let mut tk = TopK::new(2, cms);

        tk.extend(vec![0, 1]);
        assert_eq!(tk.values(), vec![0, 1]);
    }
}
<|endoftext|>


<commit_before>//! # Output utilities

pub use std::io::Write;

/// Similar to `println!(...)`, but write to `stderr` instead of `stdout`.
#[macro_export]
macro_rules! errln {
    ($($x:expr),*) => {
        writeln!(::std::io::stderr(), $($x,)*).unwrap();
    };
}

<commit_msg>fancy_status::utils: use Write trait inside errln! macro<commit_after>//! # Output utilities

/// Similar to `println!(...)`, but write to `stderr` instead of `stdout`.
#[macro_export]
macro_rules! errln {
    ($($x:expr),*) => {
        use std::io::Write;
        writeln!(::std::io::stderr(), $($x,)*).unwrap();
    };
}

<|endoftext|>


<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2015 dinowernli
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

use agent::Agent;
use bitstring::Bitstring;
use context_tree::Predictor;
use random::Random;
use types::{Observation, Reward};

#[test]
fn age() {
  let mut fake_random = FakeRandom::new();
  let fake_predictor = Box::new(FakePredictor::new());
  let mut agent = Agent::new(
      10, &mut fake_random, fake_predictor);

  assert_eq!(0, agent.age());
  agent.update(Observation(3), Reward(4.0));
  assert_eq!(1, agent.age());
}

// Fake predictor.

struct FakePredictor;

impl FakePredictor {
  fn new() -> Self {
    FakePredictor
  }
}

impl Predictor for FakePredictor {
  fn history_size(&self) -> usize { 17 }

  fn revert_to_history_size(&mut self, target_size: usize) {
    // Do nothing.
  }

  fn update(&mut self, bitstring: &Bitstring) {
    // Do nothing.
  }
}

// Fake random number generator.

struct FakeRandom {
  value: u64,
}

impl FakeRandom {
  fn new() -> Self {
    return FakeRandom {
      value: 4,
    };
  }
}

impl Random for FakeRandom {
  #[allow(unused_variables)]  // Ignores the limit.
  fn next_modulo(&mut self, limit: u64) -> u64 {
    return self.value;
  }
}
<commit_msg>Allow unused variables in mocks for testing.<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2015 dinowernli
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

use agent::Agent;
use bitstring::Bitstring;
use context_tree::Predictor;
use random::Random;
use types::{Observation, Reward};

#[test]
fn age() {
  let mut fake_random = FakeRandom::new();
  let fake_predictor = Box::new(FakePredictor::new());
  let mut agent = Agent::new(
      10, &mut fake_random, fake_predictor);

  assert_eq!(0, agent.age());
  agent.update(Observation(3), Reward(4.0));
  assert_eq!(1, agent.age());
}

// Fake predictor.

struct FakePredictor;

impl FakePredictor {
  fn new() -> Self {
    FakePredictor
  }
}

impl Predictor for FakePredictor {
  fn history_size(&self) -> usize { 17 }

  fn revert_to_history_size(&mut self, target_size: usize) {
    #![allow(unused_variables)]
    // Do nothing.
  }

  fn update(&mut self, bitstring: &Bitstring) {
    #![allow(unused_variables)]
    // Do nothing.
  }
}

// Fake random number generator.

struct FakeRandom {
  value: u64,
}

impl FakeRandom {
  fn new() -> Self {
    return FakeRandom {
      value: 4,
    };
  }
}

impl Random for FakeRandom {
  #[allow(unused_variables)]  // Ignores the limit.
  fn next_modulo(&mut self, limit: u64) -> u64 {
    return self.value;
  }
}
<|endoftext|>


<commit_before>#![feature(proc_macro, conservative_impl_trait, generators)]
extern crate bytes;
extern crate futures_await as futures;
extern crate tokio_core;
extern crate tokio_io;

mod cmd;
mod ftp;
mod error;

use std::io;
use futures::prelude::*;
use tokio_core::reactor::Core;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use tokio_core::reactor::Handle;
use tokio_core::net::TcpListener;

use tokio_core::net::TcpStream;

use futures::{Sink, Stream};
use futures::stream::SplitSink;

use tokio_io::AsyncRead;
use tokio_io::codec::Framed;
use codec::FtpCodec;
use error::Result;
use ftp::{Answer, ResultCode};

use cmd::Command;
impl Client {
    fn new(writer: Writer) -> Client {
        Client {
            writer,
        }
    }
    #[async]
    fn handle_cmd(mut self, cmd: Command) -> Result<Self> {
        return Ok(self);
    }
}

#[async]
fn client(stream: TcpStream) -> Result<()> {
    let (writer, reader) = stream.framed(FtpCodec).split();
    let writer = await!(writer.send(Answer::new(ResultCode::ServiceReadyForNewUser, "Welcome to this FTP Server!")))?;
    let mut client = Client::new(writer);
    #[async]
    for cmd in reader {
        client = await!(client.handle_cmd(cmd))?;
    }
    print!("Client closed");
    return Ok(());
}

#[async]
fn handle_client(stream: TcpStream) -> result::Result<(), ()> {
    await!(client(stream)).map_err(|error| println("Error handling client: {}", error))
}

#[async]
fn server(handle: Handle) -> io::Result<()> {
    let port = 1234;
    let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127,0,0,1)), port);
    let listener = TcpListener::bind(&addr,&handle)?;
    println!("Waiting clients on port {}...",port);
    #[async]
    for (stream, addr) in listener.incoming() {
        let address = format!("[address : {}]",addr);
        println!("New client {}", address);
        handle.spawn(handle_client(stream));
        println!("Waiting another client...");
    }
    return Ok(());
}

fn main() {
    let mut core = Core::new().expect("Cannot create tokio Core");
    let handle = core.handle();
    if let Err(error) = core.run(server(handle)) {
        println!("Error running the server: {}", error);
    }

}
<commit_msg>Fixed various errors<commit_after>#![feature(proc_macro, conservative_impl_trait, generators)]
extern crate bytes;
extern crate futures_await as futures;
extern crate tokio_core;
extern crate tokio_io;

mod cmd;
mod ftp;
mod error;

use std::io;
use futures::prelude::*;
use tokio_core::reactor::Core;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use tokio_core::reactor::Handle;
use tokio_core::net::TcpListener;

use tokio_core::net::TcpStream;

use futures::{Sink, Stream};
use futures::stream::SplitSink;

use tokio_io::AsyncRead;
use tokio_io::codec::Framed;
use codec::FtpCodec;
use error::Result;
use ftp::{Answer, ResultCode};

use cmd::Command;

impl Client {
    fn new(writer: Writer) -> Client {
        Client {
            writer,
        }
    }
    #[async]
    fn handle_cmd(mut self, cmd: Command) -> Result<Self> {
        return Ok(self);
    }
}

#[async]
fn client(stream: TcpStream) -> Result<()> {
    let (writer, reader) = stream.framed(FtpCodec).split();
    let writer = await!(writer.send(Answer::new(ResultCode::ServiceReadyForNewUser, "Welcome to this FTP Server!")))?;
    let mut client = Client::new(writer);
    #[async]
    for cmd in reader {
        client = await!(client.handle_cmd(cmd))?;
    }
    println!("Client closed");
    return Ok(());;
}

#[async]
fn handle_client(stream: TcpStream) -> result::Result<(), ()> {
    await!(client(stream)).map_err(|error| println!("Error handling client: {}", error))
}

#[async]
fn server(handle: Handle) -> io::Result<()> {
    let port = 1234;
    let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127,0,0,1)), port);
    let listener = TcpListener::bind(&addr,&handle)?;
    println!("Waiting clients on port {}...",port);
    #[async]
    for (stream, addr) in listener.incoming() {
        let address = format!("[address : {}]",addr);
        println!("New client {}", address);
        handle.spawn(handle_client(stream));
        println!("Waiting another client...");
    }
    return Ok(());
}

fn main() {
    let mut core = Core::new().expect("Cannot create tokio Core");
    let handle = core.handle();
    if let Err(error) = core.run(server(handle)) {
        println!("Error running the server: {}", error);
    }

}
<|endoftext|>


<commit_before>pub fn find_longest_common_prefix<T: Clone + Eq>(among: &[Vec<T>]) -> Option<Vec<T>> {
    if among.len() == 0 {
        return None;
    } else if among.len() == 1 {
        return Some(among[0].clone());
    }

    for s in among {
        if s.len() == 0 {
            return None;
        }
    }

    let shortest_word = among.iter().min_by_key(|x| x.len()).unwrap();

    let mut end = shortest_word.len();
    while end > 0 {
        let prefix = &shortest_word[..end];

        let mut failed = false;
        for s in among {
            if !s.starts_with(prefix) {
                failed = true;
                break;
            }
        }

        if !failed {
            return Some(prefix.into());
        }

        end -= 1;
    }

    None
}

fn is_ascii_letter(c: char) -> bool {
    (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')
}

pub fn remove_csi_codes(input: &str) -> String {
    // XXX: better way to do this
    input.split('\x1B')
        .flat_map(|x| x.chars().skip_while(|&c| c != '?' && !is_ascii_letter(c)).skip(1))
        .collect()
}
<commit_msg>fix csi code removal bug when input contains no csi codes<commit_after>pub fn find_longest_common_prefix<T: Clone + Eq>(among: &[Vec<T>]) -> Option<Vec<T>> {
    if among.len() == 0 {
        return None;
    } else if among.len() == 1 {
        return Some(among[0].clone());
    }

    for s in among {
        if s.len() == 0 {
            return None;
        }
    }

    let shortest_word = among.iter().min_by_key(|x| x.len()).unwrap();

    let mut end = shortest_word.len();
    while end > 0 {
        let prefix = &shortest_word[..end];

        let mut failed = false;
        for s in among {
            if !s.starts_with(prefix) {
                failed = true;
                break;
            }
        }

        if !failed {
            return Some(prefix.into());
        }

        end -= 1;
    }

    None
}

fn is_ascii_letter(c: char) -> bool {
    (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')
}

pub fn remove_csi_codes(input: &str) -> String {
    // XXX: better way to do this
    if !input.contains('\x1B') {
        return input.to_owned();
    }

    input.split('\x1B')
        .flat_map(|x| x.chars().skip_while(|&c| c != '?' && !is_ascii_letter(c)).skip(1))
        .collect()
}
<|endoftext|>


<commit_before>#![feature(rustc_private, custom_attribute)]
#![allow(unused_attributes)]

extern crate getopts;
extern crate miri;
extern crate rustc;
extern crate rustc_driver;
extern crate env_logger;
extern crate log_settings;
extern crate syntax;
#[macro_use] extern crate log;

use miri::{
    EvalContext,
    CachedMir,
    step,
    EvalError,
    Frame,
};
use rustc::session::Session;
use rustc_driver::{driver, CompilerCalls};
use rustc::ty::{TyCtxt, subst};
use rustc::mir::mir_map::MirMap;
use rustc::hir::def_id::DefId;

struct MiriCompilerCalls;

impl<'a> CompilerCalls<'a> for MiriCompilerCalls {
    fn build_controller(
        &mut self,
        _: &Session,
        _: &getopts::Matches
    ) -> driver::CompileController<'a> {
        let mut control = driver::CompileController::basic();

        control.after_analysis.callback = Box::new(|state| {
            state.session.abort_if_errors();
            interpret_start_points(state.tcx.unwrap(), state.mir_map.unwrap());
        });

        control
    }
}

fn interpret_start_points<'a, 'tcx>(
    tcx: TyCtxt<'a, 'tcx, 'tcx>,
    mir_map: &MirMap<'tcx>,
) {
    let initial_indentation = ::log_settings::settings().indentation;
    for (&id, mir) in &mir_map.map {
        for attr in tcx.map.attrs(id) {
            use syntax::attr::AttrMetaMethods;
            if attr.check_name("miri_run") {
                let item = tcx.map.expect_item(id);

                ::log_settings::settings().indentation = initial_indentation;

                debug!("Interpreting: {}", item.name);

                let mut ecx = EvalContext::new(tcx, mir_map);
                let substs = tcx.mk_substs(subst::Substs::empty());
                let return_ptr = ecx.alloc_ret_ptr(mir.return_ty, substs);

                ecx.push_stack_frame(tcx.map.local_def_id(id), mir.span, CachedMir::Ref(mir), substs, return_ptr);

                loop {
                    match (step(&mut ecx), return_ptr) {
                        (Ok(true), _) => {},
                        (Ok(false), Some(ptr)) => if log_enabled!(::log::LogLevel::Debug) {
                            ecx.memory().dump(ptr.alloc_id);
                            break;
                        },
                        (Ok(false), None) => {
                            warn!("diverging function returned");
                            break;
                        },
                        // FIXME: diverging functions can end up here in some future miri
                        (Err(e), _) => {
                            report(tcx, &ecx, e);
                            break;
                        },
                    }
                }
            }
        }
    }
}

fn report(tcx: TyCtxt, ecx: &EvalContext, e: EvalError) {
    let frame = ecx.stack().last().expect("stackframe was empty");
    let block = &frame.mir.basic_blocks()[frame.next_block];
    let span = if frame.stmt < block.statements.len() {
        block.statements[frame.stmt].source_info.span
    } else {
        block.terminator().source_info.span
    };
    let mut err = tcx.sess.struct_span_err(span, &e.to_string());
    for &Frame { def_id, substs, span, .. } in ecx.stack().iter().rev() {
        // FIXME(solson): Find a way to do this without this Display impl hack.
        use rustc::util::ppaux;
        use std::fmt;
        struct Instance<'tcx>(DefId, &'tcx subst::Substs<'tcx>);
        impl<'tcx> fmt::Display for Instance<'tcx> {
            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                ppaux::parameterized(f, self.1, self.0, ppaux::Ns::Value, &[],
                    |tcx| Some(tcx.lookup_item_type(self.0).generics))
            }
        }
        err.span_note(span, &format!("inside call to {}", Instance(def_id, substs)));
    }
    err.emit();
}

#[miri_run]
fn main() {
    init_logger();
    let args: Vec<String> = std::env::args().collect();
    rustc_driver::run_compiler(&args, &mut MiriCompilerCalls);
}

#[miri_run]
fn init_logger() {
    const NSPACES: usize = 40;
    let format = |record: &log::LogRecord| {
        // prepend spaces to indent the final string
        let indentation = log_settings::settings().indentation;
        format!("{lvl}:{module}{depth:2}{indent:<indentation$} {text}",
            lvl = record.level(),
            module = record.location().module_path(),
            depth = indentation / NSPACES,
            indentation = indentation % NSPACES,
            indent = "",
            text = record.args())
    };

    let mut builder = env_logger::LogBuilder::new();
    builder.format(format).filter(None, log::LogLevelFilter::Info);

    if std::env::var("MIRI_LOG").is_ok() {
        builder.parse(&std::env::var("MIRI_LOG").unwrap());
    }

    builder.init().unwrap();
}
<commit_msg>Fix infinite loop when debug trace is disabled.<commit_after>#![feature(rustc_private, custom_attribute)]
#![allow(unused_attributes)]

extern crate getopts;
extern crate miri;
extern crate rustc;
extern crate rustc_driver;
extern crate env_logger;
extern crate log_settings;
extern crate syntax;
#[macro_use] extern crate log;

use miri::{
    EvalContext,
    CachedMir,
    step,
    EvalError,
    Frame,
};
use rustc::session::Session;
use rustc_driver::{driver, CompilerCalls};
use rustc::ty::{TyCtxt, subst};
use rustc::mir::mir_map::MirMap;
use rustc::hir::def_id::DefId;

struct MiriCompilerCalls;

impl<'a> CompilerCalls<'a> for MiriCompilerCalls {
    fn build_controller(
        &mut self,
        _: &Session,
        _: &getopts::Matches
    ) -> driver::CompileController<'a> {
        let mut control = driver::CompileController::basic();

        control.after_analysis.callback = Box::new(|state| {
            state.session.abort_if_errors();
            interpret_start_points(state.tcx.unwrap(), state.mir_map.unwrap());
        });

        control
    }
}

fn interpret_start_points<'a, 'tcx>(
    tcx: TyCtxt<'a, 'tcx, 'tcx>,
    mir_map: &MirMap<'tcx>,
) {
    let initial_indentation = ::log_settings::settings().indentation;
    for (&id, mir) in &mir_map.map {
        for attr in tcx.map.attrs(id) {
            use syntax::attr::AttrMetaMethods;
            if attr.check_name("miri_run") {
                let item = tcx.map.expect_item(id);

                ::log_settings::settings().indentation = initial_indentation;

                debug!("Interpreting: {}", item.name);

                let mut ecx = EvalContext::new(tcx, mir_map);
                let substs = tcx.mk_substs(subst::Substs::empty());
                let return_ptr = ecx.alloc_ret_ptr(mir.return_ty, substs);

                ecx.push_stack_frame(tcx.map.local_def_id(id), mir.span, CachedMir::Ref(mir), substs, return_ptr);

                loop {
                    match step(&mut ecx) {
                        Ok(true) => {}
                        Ok(false) => {
                            match return_ptr {
                                Some(ptr) => if log_enabled!(::log::LogLevel::Debug) {
                                    ecx.memory().dump(ptr.alloc_id);
                                },
                                None => warn!("diverging function returned"),
                            }
                            break;
                        }
                        // FIXME: diverging functions can end up here in some future miri
                        Err(e) => {
                            report(tcx, &ecx, e);
                            break;
                        }
                    }
                }
            }
        }
    }
}

fn report(tcx: TyCtxt, ecx: &EvalContext, e: EvalError) {
    let frame = ecx.stack().last().expect("stackframe was empty");
    let block = &frame.mir.basic_blocks()[frame.next_block];
    let span = if frame.stmt < block.statements.len() {
        block.statements[frame.stmt].source_info.span
    } else {
        block.terminator().source_info.span
    };
    let mut err = tcx.sess.struct_span_err(span, &e.to_string());
    for &Frame { def_id, substs, span, .. } in ecx.stack().iter().rev() {
        // FIXME(solson): Find a way to do this without this Display impl hack.
        use rustc::util::ppaux;
        use std::fmt;
        struct Instance<'tcx>(DefId, &'tcx subst::Substs<'tcx>);
        impl<'tcx> fmt::Display for Instance<'tcx> {
            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                ppaux::parameterized(f, self.1, self.0, ppaux::Ns::Value, &[],
                    |tcx| Some(tcx.lookup_item_type(self.0).generics))
            }
        }
        err.span_note(span, &format!("inside call to {}", Instance(def_id, substs)));
    }
    err.emit();
}

#[miri_run]
fn main() {
    init_logger();
    let args: Vec<String> = std::env::args().collect();
    rustc_driver::run_compiler(&args, &mut MiriCompilerCalls);
}

#[miri_run]
fn init_logger() {
    const NSPACES: usize = 40;
    let format = |record: &log::LogRecord| {
        // prepend spaces to indent the final string
        let indentation = log_settings::settings().indentation;
        format!("{lvl}:{module}{depth:2}{indent:<indentation$} {text}",
            lvl = record.level(),
            module = record.location().module_path(),
            depth = indentation / NSPACES,
            indentation = indentation % NSPACES,
            indent = "",
            text = record.args())
    };

    let mut builder = env_logger::LogBuilder::new();
    builder.format(format).filter(None, log::LogLevelFilter::Info);

    if std::env::var("MIRI_LOG").is_ok() {
        builder.parse(&std::env::var("MIRI_LOG").unwrap());
    }

    builder.init().unwrap();
}
<|endoftext|>


<commit_before>// Copyright 2017 The Gfx-rs Developers.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Command Buffer device interface

use std::ops::{Deref, DerefMut};
use {state, pso, target};
use {IndexType, InstanceCount, VertexCount, Resources};

/// A universal clear color supporting integet formats
/// as well as the standard floating-point.
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd)]
pub enum ClearColor {
    /// Standard floating-point vec4 color
    Float([f32; 4]),
    /// Integer vector to clear ivec4 targets.
    Int([i32; 4]),
    /// Unsigned int vector to clear uvec4 targets.
    Uint([u32; 4]),
}

pub struct BufferCopy {
    pub src: usize,
    pub dest: usize,
    pub size: usize,
}

/// Optional instance parameters: (instance count, buffer offset)
pub type InstanceParams = (InstanceCount, VertexCount);

pub struct Encoder<'a, C: CommandBuffer + 'a>(&'a mut C);

impl<'a, C: CommandBuffer> Encoder<'a, C> {
    #[doc(hidden)]
    pub unsafe fn new(buffer: &'a mut C) -> Self {
        Encoder(buffer)
    }

    // TODO: return submission object
    pub fn finish(self) -> Submit<C> {
        Submit(unsafe { self.0.end() })
    }
}

impl<'a, C> Deref for Encoder<'a, C>
    where C: CommandBuffer
{
    type Target = C;

    fn deref(&self) -> &C {
        self.0
    }
}

impl<'a, C> DerefMut for Encoder<'a, C>
    where C: CommandBuffer
{
    fn deref_mut(&mut self) -> &mut C {
        self.0
    }
}

// TODO: DON'T take a reference! Probably needs to be done for each backend /:
pub struct Submit<C: CommandBuffer>(C::SubmitInfo);
impl<C: CommandBuffer> Submit<C> {
    #[doc(hidden)]
    pub unsafe fn get_info(&self) -> &C::SubmitInfo {
        &self.0
    }
}

pub trait GraphicsCommandBuffer<R: Resources> : PrimaryCommandBuffer<R> {
    fn clear_depth_stencil(&mut self, &R::DepthStencilView, Option<target::Depth>, Option<target::Stencil>);

    fn resolve_image(&mut self);

    fn bind_index_buffer(&mut self, &R::Buffer, IndexType);
    fn bind_vertex_buffers(&mut self, pso::VertexBufferSet<R>);

    fn set_viewports(&mut self, &[target::Rect]);
    fn set_scissors(&mut self, &[target::Rect]);
    fn set_ref_values(&mut self, state::RefValues);
}

pub trait RenderPassEncoder<C: GraphicsCommandBuffer<R>, R: Resources> {
    fn begin(&mut C, &R::RenderPass) -> Self;
    fn next_subpass(&mut self);
}

pub trait SubpassCommandBuffer<R: Resources> : SecondaryCommandBuffer<R> {
    fn clear_attachment(&mut self);
    fn draw(&mut self, start: VertexCount, count: VertexCount, Option<InstanceParams>);
    fn draw_indexed(&mut self, start: VertexCount, count: VertexCount, base: VertexCount, Option<InstanceParams>);
    fn draw_indirect(&mut self);
    fn draw_indexed_indirect(&mut self);

    fn bind_index_buffer(&mut self, &R::Buffer, IndexType);
    fn bind_vertex_buffers(&mut self, pso::VertexBufferSet<R>);

    fn set_viewports(&mut self, &[target::Rect]);
    fn set_scissors(&mut self, &[target::Rect]);
    fn set_ref_values(&mut self, state::RefValues);

    fn bind_pipeline(&mut self, &R::PipelineStateObject);
    fn bind_descriptor_sets(&mut self);
    fn push_constants(&mut self);
}

pub trait ComputeCommandBuffer<R: Resources> : ProcessingCommandBuffer<R> {
    fn dispatch(&mut self, u32, u32, u32);
    fn dispatch_indirect(&mut self);
}

pub trait ProcessingCommandBuffer<R: Resources> : TransferCommandBuffer<R> {
    fn clear_color(&mut self, &R::RenderTargetView, ClearColor);
    fn clear_buffer(&mut self);

    // TODO: consider splitting compute and graphics pso
    fn bind_pipeline(&mut self, &R::PipelineStateObject);
    fn bind_descriptor_sets(&mut self);
    fn push_constants(&mut self);
}

pub trait TransferCommandBuffer<R: Resources> : PrimaryCommandBuffer<R> {
    fn update_buffer(&mut self, &R::Buffer, data: &[u8], offset: usize);
    fn copy_buffer(&mut self, src: &R::Buffer, dest: &R::Buffer, &[BufferCopy]);
    fn copy_image(&mut self, src: &R::Image, dest: &R::Image);
    fn copy_buffer_to_image(&mut self);
    fn copy_image_to_buffer(&mut self);
}

pub trait PrimaryCommandBuffer<R: Resources>: CommandBuffer {
    fn pipeline_barrier(&mut self);
    fn execute_commands(&mut self);
}

pub trait SecondaryCommandBuffer<R: Resources>: CommandBuffer {
    fn pipeline_barrier(&mut self);
}

pub trait CommandBuffer {
    type SubmitInfo;

    #[doc(hidden)]
    unsafe fn end(&mut self) -> Self::SubmitInfo;
}

// Ignore for the moment (:
/*
// vk: primary/seconday | outside
fn set_event(&mut self); // vk: Graphics/Compute // d3d12:! emulation needed
// vk: primary/seconday | outside
fn reset_event(&mut self); // vk: Graphics/Compute
// vk: primary/seconday | inside/outside
fn wait_event(&mut self); // vk: Graphics/Compute

// vk: primary/seconday | inside/outside // d3d12: primary
fn begin_query(&mut self); // vk: Graphics/Compute // d3d12: BeginQuery
// vk: primary/seconday | inside/outside // d3d12: primary
fn end_query(&mut self); // vk: Graphics/Compute // d3d12: EndQuery
// vk: primary/seconday | outside
fn reset_query_pool(&mut self); // vk: Graphics/Compute
// vk: primary/seconday | inside/outside
fn write_timestamp(&mut self); // vk: Graphics/Compute
// vk: primary/seconday | outside
fn copy_query_pool_results(&mut self); // vk: Graphics/Compute
*/
<commit_msg>[ll] Cleanup some outdated TODOs<commit_after>// Copyright 2017 The Gfx-rs Developers.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Command Buffer device interface

use std::ops::{Deref, DerefMut};
use {state, pso, target};
use {IndexType, InstanceCount, VertexCount, Resources};

/// A universal clear color supporting integet formats
/// as well as the standard floating-point.
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd)]
pub enum ClearColor {
    /// Standard floating-point vec4 color
    Float([f32; 4]),
    /// Integer vector to clear ivec4 targets.
    Int([i32; 4]),
    /// Unsigned int vector to clear uvec4 targets.
    Uint([u32; 4]),
}

pub struct BufferCopy {
    pub src: usize,
    pub dest: usize,
    pub size: usize,
}

/// Optional instance parameters: (instance count, buffer offset)
pub type InstanceParams = (InstanceCount, VertexCount);

pub struct Encoder<'a, C: CommandBuffer + 'a>(&'a mut C);

impl<'a, C: CommandBuffer> Encoder<'a, C> {
    #[doc(hidden)]
    pub unsafe fn new(buffer: &'a mut C) -> Self {
        Encoder(buffer)
    }

    pub fn finish(self) -> Submit<C> {
        Submit(unsafe { self.0.end() })
    }
}

impl<'a, C> Deref for Encoder<'a, C>
    where C: CommandBuffer
{
    type Target = C;

    fn deref(&self) -> &C {
        self.0
    }
}

impl<'a, C> DerefMut for Encoder<'a, C>
    where C: CommandBuffer
{
    fn deref_mut(&mut self) -> &mut C {
        self.0
    }
}

pub struct Submit<C: CommandBuffer>(C::SubmitInfo);
impl<C: CommandBuffer> Submit<C> {
    #[doc(hidden)]
    pub unsafe fn get_info(&self) -> &C::SubmitInfo {
        &self.0
    }
}

pub trait GraphicsCommandBuffer<R: Resources> : PrimaryCommandBuffer<R> {
    fn clear_depth_stencil(&mut self, &R::DepthStencilView, Option<target::Depth>, Option<target::Stencil>);

    fn resolve_image(&mut self);

    fn bind_index_buffer(&mut self, &R::Buffer, IndexType);
    fn bind_vertex_buffers(&mut self, pso::VertexBufferSet<R>);

    fn set_viewports(&mut self, &[target::Rect]);
    fn set_scissors(&mut self, &[target::Rect]);
    fn set_ref_values(&mut self, state::RefValues);
}

pub trait RenderPassEncoder<C: GraphicsCommandBuffer<R>, R: Resources> {
    fn begin(&mut C, &R::RenderPass) -> Self;
    fn next_subpass(&mut self);
}

pub trait SubpassCommandBuffer<R: Resources> : SecondaryCommandBuffer<R> {
    fn clear_attachment(&mut self);
    fn draw(&mut self, start: VertexCount, count: VertexCount, Option<InstanceParams>);
    fn draw_indexed(&mut self, start: VertexCount, count: VertexCount, base: VertexCount, Option<InstanceParams>);
    fn draw_indirect(&mut self);
    fn draw_indexed_indirect(&mut self);

    fn bind_index_buffer(&mut self, &R::Buffer, IndexType);
    fn bind_vertex_buffers(&mut self, pso::VertexBufferSet<R>);

    fn set_viewports(&mut self, &[target::Rect]);
    fn set_scissors(&mut self, &[target::Rect]);
    fn set_ref_values(&mut self, state::RefValues);

    fn bind_pipeline(&mut self, &R::PipelineStateObject);
    fn bind_descriptor_sets(&mut self);
    fn push_constants(&mut self);
}

pub trait ComputeCommandBuffer<R: Resources> : ProcessingCommandBuffer<R> {
    fn dispatch(&mut self, u32, u32, u32);
    fn dispatch_indirect(&mut self);
}

pub trait ProcessingCommandBuffer<R: Resources> : TransferCommandBuffer<R> {
    fn clear_color(&mut self, &R::RenderTargetView, ClearColor);
    fn clear_buffer(&mut self);

    // TODO: consider splitting compute and graphics pso
    fn bind_pipeline(&mut self, &R::PipelineStateObject);
    fn bind_descriptor_sets(&mut self);
    fn push_constants(&mut self);
}

pub trait TransferCommandBuffer<R: Resources> : PrimaryCommandBuffer<R> {
    fn update_buffer(&mut self, &R::Buffer, data: &[u8], offset: usize);
    fn copy_buffer(&mut self, src: &R::Buffer, dest: &R::Buffer, &[BufferCopy]);
    fn copy_image(&mut self, src: &R::Image, dest: &R::Image);
    fn copy_buffer_to_image(&mut self);
    fn copy_image_to_buffer(&mut self);
}

pub trait PrimaryCommandBuffer<R: Resources>: CommandBuffer {
    fn pipeline_barrier(&mut self);
    fn execute_commands(&mut self);
}

pub trait SecondaryCommandBuffer<R: Resources>: CommandBuffer {
    fn pipeline_barrier(&mut self);
}

pub trait CommandBuffer {
    type SubmitInfo;

    #[doc(hidden)]
    unsafe fn end(&mut self) -> Self::SubmitInfo;
}

// Ignore for the moment (:
/*
// vk: primary/seconday | outside
fn set_event(&mut self); // vk: Graphics/Compute // d3d12:! emulation needed
// vk: primary/seconday | outside
fn reset_event(&mut self); // vk: Graphics/Compute
// vk: primary/seconday | inside/outside
fn wait_event(&mut self); // vk: Graphics/Compute

// vk: primary/seconday | inside/outside // d3d12: primary
fn begin_query(&mut self); // vk: Graphics/Compute // d3d12: BeginQuery
// vk: primary/seconday | inside/outside // d3d12: primary
fn end_query(&mut self); // vk: Graphics/Compute // d3d12: EndQuery
// vk: primary/seconday | outside
fn reset_query_pool(&mut self); // vk: Graphics/Compute
// vk: primary/seconday | inside/outside
fn write_timestamp(&mut self); // vk: Graphics/Compute
// vk: primary/seconday | outside
fn copy_query_pool_results(&mut self); // vk: Graphics/Compute
*/
<|endoftext|>


<commit_before>//! Decode Cryptographic Control (CC) packets.

use consts::CRYPTO_CONTROL_BYTES;
use util::slice_u16;

/// Buffer of bytes that represent a crypto control packet.
pub type Buf = [u8; CRYPTO_CONTROL_BYTES];

/// Information necessary to decrypt an encrypted message.
pub struct CryptoControlFields(Buf);

impl CryptoControlFields {
    /// Create a new `CryptoControlFields` decoder from the given bytes.
    pub fn new(buf: Buf) -> Self { CryptoControlFields(buf) }

    /// Initialization vector used internally by associated crypto algorithm.
    pub fn init(&self) -> &[u8] { &self.0[..9] }
    /// Type of crypto algorithm in use, if any.
    pub fn alg(&self) -> CryptoAlgorithm { CryptoAlgorithm::from_bits(self.0[9]) }
    /// Encryption key to use.
    pub fn key(&self) -> u16 { slice_u16(&self.0[10..]) }
}

/// Type of cryptographic algorithm.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum CryptoAlgorithm {
    Accordion,
    BatonEven,
    Firefly,
    Mayfly,
    Saville,
    BatonOdd,
    Unencrypted,
    Des,
    TripleDES,
    AES,
    Other(u8),
}

impl CryptoAlgorithm {
    /// Parse the given 8 bits into a crypto algorithm.
    pub fn from_bits(bits: u8) -> CryptoAlgorithm {
        use self::CryptoAlgorithm::*;

        match bits {
            0x00 => Accordion,
            0x01 => BatonEven,
            0x02 => Firefly,
            0x03 => Mayfly,
            0x04 => Saville,
            0x41 => BatonOdd,
            0x80 => Unencrypted,
            0x81 => Des,
            0x83 => TripleDES,
            0x84 => AES,
            b => Other(b),
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use super::CryptoAlgorithm::*;

    #[test]
    fn test_cypto() {
        let c = CryptoControlFields::new([
            0, 0, 0, 1, 0, 0, 0, 2, 0,
            0b10000100,
            0xDE, 0xAD,
        ]);

        assert_eq!(c.init(), &[0,0,0,1,0,0,0,2,0]);
        assert_eq!(c.alg(), AES);
        assert_eq!(c.key(), 0xDEAD);
    }
}
<commit_msg>s/AES/Aes/<commit_after>//! Decode Cryptographic Control (CC) packets.

use consts::CRYPTO_CONTROL_BYTES;
use util::slice_u16;

/// Buffer of bytes that represent a crypto control packet.
pub type Buf = [u8; CRYPTO_CONTROL_BYTES];

/// Information necessary to decrypt an encrypted message.
pub struct CryptoControlFields(Buf);

impl CryptoControlFields {
    /// Create a new `CryptoControlFields` decoder from the given bytes.
    pub fn new(buf: Buf) -> Self { CryptoControlFields(buf) }

    /// Initialization vector used internally by associated crypto algorithm.
    pub fn init(&self) -> &[u8] { &self.0[..9] }
    /// Type of crypto algorithm in use, if any.
    pub fn alg(&self) -> CryptoAlgorithm { CryptoAlgorithm::from_bits(self.0[9]) }
    /// Encryption key to use.
    pub fn key(&self) -> u16 { slice_u16(&self.0[10..]) }
}

/// Type of cryptographic algorithm.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum CryptoAlgorithm {
    Accordion,
    BatonEven,
    Firefly,
    Mayfly,
    Saville,
    BatonOdd,
    Unencrypted,
    Des,
    TripleDES,
    Aes,
    Other(u8),
}

impl CryptoAlgorithm {
    /// Parse the given 8 bits into a crypto algorithm.
    pub fn from_bits(bits: u8) -> CryptoAlgorithm {
        use self::CryptoAlgorithm::*;

        match bits {
            0x00 => Accordion,
            0x01 => BatonEven,
            0x02 => Firefly,
            0x03 => Mayfly,
            0x04 => Saville,
            0x41 => BatonOdd,
            0x80 => Unencrypted,
            0x81 => Des,
            0x83 => TripleDES,
            0x84 => Aes,
            b => Other(b),
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use super::CryptoAlgorithm::*;

    #[test]
    fn test_cypto() {
        let c = CryptoControlFields::new([
            0, 0, 0, 1, 0, 0, 0, 2, 0,
            0b10000100,
            0xDE, 0xAD,
        ]);

        assert_eq!(c.init(), &[0,0,0,1,0,0,0,2,0]);
        assert_eq!(c.alg(), Aes);
        assert_eq!(c.key(), 0xDEAD);
    }
}
<|endoftext|>


<commit_before>#![feature(phase, if_let)]

extern crate quire;
extern crate argparse;
extern crate serialize;
extern crate regex;
#[phase(plugin)] extern crate regex_macros;

extern crate config;
#[phase(plugin, link)] extern crate container;

use std::io::stderr;
use std::io::ALL_PERMISSIONS;
use std::io::{TypeSymlink, TypeDirectory, PathDoesntExist, FileNotFound};
use std::os::{getcwd, set_exit_status, self_exe_path};
use std::io::fs::{mkdir, copy, readlink, symlink};
use std::io::fs::PathExtensions;

use config::find_config;
use container::signal;
use container::monitor::{Monitor, Executor};
use container::container::{Command};
use container::mount::{mount_tmpfs, bind_mount, unmount};
use container::mount::{mount_ro_recursive, mount_pseudo};
use container::root::change_root;
use settings::{read_settings, MergedSettings};
use argparse::{ArgumentParser, Store, List};

mod settings;

pub fn make_mountpoint(project_root: &Path) -> Result<(), String> {
    let vagga_dir = project_root.join(".vagga");
    match vagga_dir.lstat() {
        Ok(stat) if stat.kind == TypeSymlink => {
            return Err(concat!("The `.vagga` dir can't be a symlink. ",
                               "Please run `unlink .vagga`").to_string());
        }
        Ok(stat) if stat.kind == TypeDirectory => {
            // ok
        }
        Ok(_) => {
            return Err(concat!("The `.vagga` must be a directory. ",
                               "Please run `unlink .vagga`").to_string());
        }
        Err(ref e) if e.kind == PathDoesntExist => {
            try!(mkdir(&vagga_dir, ALL_PERMISSIONS)
                .map_err(|e| format!("Can't create {}: {}",
                                     vagga_dir.display(), e)));
        }
        Err(ref e) => {
            return Err(format!("Can't stat `.vagga`: {}", e));
        }
    }
    let mnt_dir = vagga_dir.join(".mnt");
    match mnt_dir.lstat() {
        Ok(stat) if stat.kind == TypeSymlink => {
            return Err(concat!("The `.vagga/.mnt` dir can't be a symlink. ",
                               "Please run `unlink .vagga/.mnt`").to_string());
        }
        Ok(stat) if stat.kind == TypeDirectory => {
            // ok
        }
        Ok(_) => {
            return Err(concat!("The `.vagga/.mnt` must be a directory. ",
                               "Please run `unlink .vagga/.mnt`").to_string());
        }
        Err(ref e) if e.kind == PathDoesntExist => {
            try!(mkdir(&mnt_dir, ALL_PERMISSIONS)
                .map_err(|e| format!("Can't create {}: {}",
                                     mnt_dir.display(), e)));
        }
        Err(ref e) => {
            return Err(format!("Can't stat `.vagga/.mnt`: {}", e));
        }
    }
    return Ok(());
}

pub fn create_storage_dir(storage_dir: &Path, project_root: &Path)
    -> Result<Path, String>
{
    let name = match project_root.filename_str() {
        Some(name) => name,
        None => return Err(format!(
            "Project dir `{}` is either root or has bad characters",
            project_root.display())),
    };
    let path = storage_dir.join(name);
    if !path.exists() {
        return Ok(path);
    }
    for i in range(1i, 101i) {
        let result = format!("{}-{}", name, i);
        let path = storage_dir.join(result);
        if !path.exists() {
            return Ok(path);
        }
    }
    return Err(format!("Too many similar paths named {} in {}",
        name, storage_dir.display()));
}

pub fn setup_filesystem(project_root: &Path, settings: &MergedSettings)
    -> Result<(), String>
{
    let mnt_dir = project_root.join(".vagga/.mnt");
    try!(make_mountpoint(project_root));
    try!(mount_tmpfs(&mnt_dir, "size=10m"));

    let proc_dir = mnt_dir.join("proc");
    try_str!(mkdir(&proc_dir, ALL_PERMISSIONS));
    try!(mount_pseudo(&proc_dir, "proc", "", false));

    let dev_dir = mnt_dir.join("dev");
    try_str!(mkdir(&dev_dir, ALL_PERMISSIONS));
    try!(bind_mount(&Path::new("/dev"), &dev_dir));

    let vagga_dir = mnt_dir.join("vagga");
    try_str!(mkdir(&vagga_dir, ALL_PERMISSIONS));

    let bin_dir = vagga_dir.join("bin");
    try_str!(mkdir(&bin_dir, ALL_PERMISSIONS));
    try!(bind_mount(&self_exe_path().unwrap(), &bin_dir));
    try!(mount_ro_recursive(&bin_dir));

    let etc_dir = mnt_dir.join("etc");
    try_str!(mkdir(&etc_dir, ALL_PERMISSIONS));
    try!(copy(&Path::new("/etc/hosts"), &etc_dir.join("hosts"))
        .map_err(|e| format!("Error copying /etc/hosts: {}", e)));
    try!(copy(&Path::new("/etc/resolv.conf"), &etc_dir.join("resolv.conf"))
        .map_err(|e| format!("Error copying /etc/resolv.conf: {}", e)));

    let roots_dir = vagga_dir.join("roots");
    try_str!(mkdir(&roots_dir, ALL_PERMISSIONS));
    if let Some(ref dir) = settings.storage_dir {
        let lnkdir = project_root.join(".vagga/.lnk");
        let local_roots = match readlink(&lnkdir) {
            Ok(lnk) => {
                if let Some(name) = lnk.filename() {
                    let target = dir.join(name);
                    if Path::new(lnk.dirname()) != *dir {
                        return Err(concat!("You have set storage_dir to {}, ",
                            "but .vagga/.lnk points to {}. You probably need ",
                            "to run `ln -sfn {} .vagga/.lnk`").to_string());
                    }
                    if !lnkdir.exists() {
                        return Err(concat!("Your .vagga/.lnk points to a ",
                            "non-existent directory. Presumably you deleted ",
                            "dir {}. Just remove .vagga/.lnk now."
                            ).to_string());
                    }
                    target
                } else {
                    return Err(format!(concat!("Bad link .vagga/.lnk: {}.",
                        " You are pobably need to remove it now"),
                        lnk.display()));
                }
            }
            Err(ref e) if e.kind == FileNotFound => {
                let target = try!(create_storage_dir(dir, project_root));
                try_str!(mkdir(&target, ALL_PERMISSIONS));
                try_str!(symlink(&target, &lnkdir));
                target
            }
            Err(ref e) => {
                return Err(format!("Can't read link .vagga/.lnk: {}", e));
            }
        };
        try!(bind_mount(&local_roots, &roots_dir));
    } else {
        let local_roots = project_root.join(".vagga/.roots");
        if !local_roots.exists() {
            try_str!(mkdir(&local_roots, ALL_PERMISSIONS));
        }
        try!(bind_mount(&local_roots, &roots_dir));
    }

    let work_dir = mnt_dir.join("work");
    try_str!(mkdir(&work_dir, ALL_PERMISSIONS));
    try!(bind_mount(project_root, &work_dir));

    let old_root = vagga_dir.join("old_root");
    try_str!(mkdir(&old_root, ALL_PERMISSIONS));
    try!(change_root(&mnt_dir, &old_root));
    try!(unmount(&Path::new("/vagga/old_root")));

    return Ok(());
}

pub fn run() -> int {
    let mut err = stderr();
    let mut cmd: String = "".to_string();
    let mut args: Vec<String> = Vec::new();
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Runs a command in container, optionally builds container if that
            does not exists or outdated.

            Run `vagga` without arguments to see the list of commands.
            ");
        ap.refer(&mut cmd)
          .add_argument("command", box Store::<String>,
                "A vagga command to run")
          .required();
        ap.refer(&mut args)
          .add_argument("args", box List::<String>,
                "Arguments for the command");
        ap.stop_on_first_argument(true);
        match ap.parse_args() {
            Ok(()) => {}
            Err(0) => return 0,
            Err(_) => return 122,
        }
    }

    let workdir = getcwd();

    let (config, project_root) = match find_config(&workdir) {
        Ok(tup) => tup,
        Err(e) => {
            err.write_line(e.as_slice()).ok();
            return 126;
        }
    };
    let (ext_settings, int_settings) = match read_settings(&project_root)
    {
        Ok(tup) => tup,
        Err(e) => {
            err.write_line(e.as_slice()).ok();
            return 126;
        }
    };

    match setup_filesystem(&project_root, &ext_settings) {
        Ok(()) => {
            return 0;
        }
        Err(text) =>  {
            err.write_line(text.as_slice()).ok();
            return 121;
        }
    }
}

fn main() {
    signal::block_all();
    let val = run();
    set_exit_status(val);
}
<commit_msg>Factored out make_local_roots function<commit_after>#![feature(phase, if_let)]

extern crate quire;
extern crate argparse;
extern crate serialize;
extern crate regex;
#[phase(plugin)] extern crate regex_macros;

extern crate config;
#[phase(plugin, link)] extern crate container;

use std::io::stderr;
use std::io::ALL_PERMISSIONS;
use std::io::{TypeSymlink, TypeDirectory, PathDoesntExist, FileNotFound};
use std::os::{getcwd, set_exit_status, self_exe_path};
use std::io::fs::{mkdir, copy, readlink, symlink};
use std::io::fs::PathExtensions;

use config::find_config;
use container::signal;
use container::monitor::{Monitor, Executor};
use container::container::{Command};
use container::mount::{mount_tmpfs, bind_mount, unmount};
use container::mount::{mount_ro_recursive, mount_pseudo};
use container::root::change_root;
use settings::{read_settings, MergedSettings};
use argparse::{ArgumentParser, Store, List};

mod settings;

fn make_mountpoint(project_root: &Path) -> Result<(), String> {
    let vagga_dir = project_root.join(".vagga");
    match vagga_dir.lstat() {
        Ok(stat) if stat.kind == TypeSymlink => {
            return Err(concat!("The `.vagga` dir can't be a symlink. ",
                               "Please run `unlink .vagga`").to_string());
        }
        Ok(stat) if stat.kind == TypeDirectory => {
            // ok
        }
        Ok(_) => {
            return Err(concat!("The `.vagga` must be a directory. ",
                               "Please run `unlink .vagga`").to_string());
        }
        Err(ref e) if e.kind == PathDoesntExist => {
            try!(mkdir(&vagga_dir, ALL_PERMISSIONS)
                .map_err(|e| format!("Can't create {}: {}",
                                     vagga_dir.display(), e)));
        }
        Err(ref e) => {
            return Err(format!("Can't stat `.vagga`: {}", e));
        }
    }
    let mnt_dir = vagga_dir.join(".mnt");
    match mnt_dir.lstat() {
        Ok(stat) if stat.kind == TypeSymlink => {
            return Err(concat!("The `.vagga/.mnt` dir can't be a symlink. ",
                               "Please run `unlink .vagga/.mnt`").to_string());
        }
        Ok(stat) if stat.kind == TypeDirectory => {
            // ok
        }
        Ok(_) => {
            return Err(concat!("The `.vagga/.mnt` must be a directory. ",
                               "Please run `unlink .vagga/.mnt`").to_string());
        }
        Err(ref e) if e.kind == PathDoesntExist => {
            try!(mkdir(&mnt_dir, ALL_PERMISSIONS)
                .map_err(|e| format!("Can't create {}: {}",
                                     mnt_dir.display(), e)));
        }
        Err(ref e) => {
            return Err(format!("Can't stat `.vagga/.mnt`: {}", e));
        }
    }
    return Ok(());
}

fn create_storage_dir(storage_dir: &Path, project_root: &Path)
    -> Result<Path, String>
{
    let name = match project_root.filename_str() {
        Some(name) => name,
        None => return Err(format!(
            "Project dir `{}` is either root or has bad characters",
            project_root.display())),
    };
    let path = storage_dir.join(name);
    if !path.exists() {
        return Ok(path);
    }
    for i in range(1i, 101i) {
        let result = format!("{}-{}", name, i);
        let path = storage_dir.join(result);
        if !path.exists() {
            return Ok(path);
        }
    }
    return Err(format!("Too many similar paths named {} in {}",
        name, storage_dir.display()));
}

fn make_local_roots(project_root: &Path, settings: &MergedSettings)
    -> Result<Path, String>
{
    if let Some(ref dir) = settings.storage_dir {
        let lnkdir = project_root.join(".vagga/.lnk");
        match readlink(&lnkdir) {
            Ok(lnk) => {
                if let Some(name) = lnk.filename() {
                    let target = dir.join(name);
                    if Path::new(lnk.dirname()) != *dir {
                        return Err(concat!("You have set storage_dir to {}, ",
                            "but .vagga/.lnk points to {}. You probably need ",
                            "to run `ln -sfn {} .vagga/.lnk`").to_string());
                    }
                    if !lnkdir.exists() {
                        return Err(concat!("Your .vagga/.lnk points to a ",
                            "non-existent directory. Presumably you deleted ",
                            "dir {}. Just remove .vagga/.lnk now."
                            ).to_string());
                    }
                    return Ok(target);
                } else {
                    return Err(format!(concat!("Bad link .vagga/.lnk: {}.",
                        " You are pobably need to remove it now"),
                        lnk.display()));
                }
            }
            Err(ref e) if e.kind == FileNotFound => {
                let target = try!(create_storage_dir(dir, project_root));
                try_str!(mkdir(&target, ALL_PERMISSIONS));
                try_str!(symlink(&target, &lnkdir));
                return Ok(target)
            }
            Err(ref e) => {
                return Err(format!("Can't read link .vagga/.lnk: {}", e));
            }
        };
    } else {
        let local_roots = project_root.join(".vagga/.roots");
        if !local_roots.exists() {
            try_str!(mkdir(&local_roots, ALL_PERMISSIONS));
        }
        return Ok(local_roots);
    }
}

fn setup_filesystem(project_root: &Path, settings: &MergedSettings)
    -> Result<(), String>
{
    let mnt_dir = project_root.join(".vagga/.mnt");
    try!(make_mountpoint(project_root));
    try!(mount_tmpfs(&mnt_dir, "size=10m"));

    let proc_dir = mnt_dir.join("proc");
    try_str!(mkdir(&proc_dir, ALL_PERMISSIONS));
    try!(mount_pseudo(&proc_dir, "proc", "", false));

    let dev_dir = mnt_dir.join("dev");
    try_str!(mkdir(&dev_dir, ALL_PERMISSIONS));
    try!(bind_mount(&Path::new("/dev"), &dev_dir));

    let vagga_dir = mnt_dir.join("vagga");
    try_str!(mkdir(&vagga_dir, ALL_PERMISSIONS));

    let bin_dir = vagga_dir.join("bin");
    try_str!(mkdir(&bin_dir, ALL_PERMISSIONS));
    try!(bind_mount(&self_exe_path().unwrap(), &bin_dir));
    try!(mount_ro_recursive(&bin_dir));

    let etc_dir = mnt_dir.join("etc");
    try_str!(mkdir(&etc_dir, ALL_PERMISSIONS));
    try!(copy(&Path::new("/etc/hosts"), &etc_dir.join("hosts"))
        .map_err(|e| format!("Error copying /etc/hosts: {}", e)));
    try!(copy(&Path::new("/etc/resolv.conf"), &etc_dir.join("resolv.conf"))
        .map_err(|e| format!("Error copying /etc/resolv.conf: {}", e)));

    let roots_dir = vagga_dir.join("roots");
    try_str!(mkdir(&roots_dir, ALL_PERMISSIONS));
    {
        let local_roots = try!(make_local_roots(project_root, settings));
        try!(bind_mount(&local_roots, &roots_dir));
    }

    let work_dir = mnt_dir.join("work");
    try_str!(mkdir(&work_dir, ALL_PERMISSIONS));
    try!(bind_mount(project_root, &work_dir));

    let old_root = vagga_dir.join("old_root");
    try_str!(mkdir(&old_root, ALL_PERMISSIONS));
    try!(change_root(&mnt_dir, &old_root));
    try!(unmount(&Path::new("/vagga/old_root")));

    return Ok(());
}

pub fn run() -> int {
    let mut err = stderr();
    let mut cmd: String = "".to_string();
    let mut args: Vec<String> = Vec::new();
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Runs a command in container, optionally builds container if that
            does not exists or outdated.

            Run `vagga` without arguments to see the list of commands.
            ");
        ap.refer(&mut cmd)
          .add_argument("command", box Store::<String>,
                "A vagga command to run")
          .required();
        ap.refer(&mut args)
          .add_argument("args", box List::<String>,
                "Arguments for the command");
        ap.stop_on_first_argument(true);
        match ap.parse_args() {
            Ok(()) => {}
            Err(0) => return 0,
            Err(_) => return 122,
        }
    }

    let workdir = getcwd();

    let (config, project_root) = match find_config(&workdir) {
        Ok(tup) => tup,
        Err(e) => {
            err.write_line(e.as_slice()).ok();
            return 126;
        }
    };
    let (ext_settings, int_settings) = match read_settings(&project_root)
    {
        Ok(tup) => tup,
        Err(e) => {
            err.write_line(e.as_slice()).ok();
            return 126;
        }
    };

    match setup_filesystem(&project_root, &ext_settings) {
        Ok(()) => {
            return 0;
        }
        Err(text) =>  {
            err.write_line(text.as_slice()).ok();
            return 121;
        }
    }
}

fn main() {
    signal::block_all();
    let val = run();
    set_exit_status(val);
}
<|endoftext|>