<commit_before>//! Compiler plugin for Rust-PHF
//!
//! See the documentation for the `phf` crate for more details.
#![crate_id="github.com/sfackler/rust-phf/phf_mac"]
#![crate_type="dylib"]
#![doc(html_root_url="http://sfackler.github.io/rust-phf/doc")]
#![feature(managed_boxes, macro_registrar, quote)]

extern crate collections;
extern crate rand;
extern crate syntax;
extern crate time;
extern crate phf;

use collections::HashMap;
use std::os;
use syntax::ast;
use syntax::ast::{Name, TokenTree, LitStr, Expr, ExprVec, ExprLit};
use syntax::codemap::Span;
use syntax::ext::base::{SyntaxExtension,
                        DummyResult,
                        ExtCtxt,
                        MacResult,
                        MacExpr,
                        NormalTT,
                        BasicMacroExpander};
use syntax::parse;
use syntax::parse::token;
use syntax::parse::token::{InternedString, COMMA, EOF, FAT_ARROW};
use rand::{Rng, SeedableRng, XorShiftRng};

static DEFAULT_LAMBDA: uint = 5;

static FIXED_SEED: [u32, ..4] = [3141592653, 589793238, 462643383, 2795028841];

#[macro_registrar]
#[doc(hidden)]
pub fn macro_registrar(register: |Name, SyntaxExtension|) {
    let reg = |name, fn_| {
        register(token::intern(name),
                 NormalTT(box BasicMacroExpander {
                     expander: fn_,
                     span: None
                 },
                 None));
    };
    reg("phf_map", expand_phf_map);
    reg("phf_set", expand_phf_set);
    reg("phf_ordered_map", expand_phf_ordered_map);
    reg("phf_ordered_set", expand_phf_ordered_set);
}

struct Entry {
    key_str: InternedString,
    key: @Expr,
    value: @Expr
}

struct HashState {
    k1: u64,
    k2: u64,
    disps: Vec<(uint, uint)>,
    map: Vec<uint>,
}

fn expand_phf_map(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree])
                  -> Box<MacResult> {
    let entries = match parse_map(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp)
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, entries.as_slice());

    create_map(cx, sp, entries, state)
}

fn expand_phf_set(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree])
                  -> Box<MacResult> {
    let entries = match parse_set(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp)
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, entries.as_slice());

    create_set(cx, sp, entries, state)
}

fn expand_phf_ordered_map(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree])
                          -> Box<MacResult> {
    let entries = match parse_map(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp),
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, entries.as_slice());

    create_ordered_map(cx, sp, entries, state)
}

fn expand_phf_ordered_set(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree])
                          -> Box<MacResult> {
    let entries = match parse_set(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp)
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, entries.as_slice());

    create_ordered_set(cx, sp, entries, state)
}

fn parse_map(cx: &mut ExtCtxt, tts: &[TokenTree]) -> Option<Vec<Entry>> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(),
                                                Vec::from_slice(tts));
    let mut entries = Vec::new();

    let mut bad = false;
    while parser.token != EOF {
        let key = cx.expand_expr(parser.parse_expr());
        let key_str = parse_str(cx, key).unwrap_or_else(|| {
            bad = true;
            InternedString::new("")
        });

        if !parser.eat(&FAT_ARROW) {
            cx.span_err(parser.span, "expected `=>`");
            return None;
        }

        let value = parser.parse_expr();

        entries.push(Entry {
            key_str: key_str,
            key: key,
            value: value
        });

        if !parser.eat(&COMMA) && parser.token != EOF {
            cx.span_err(parser.span, "expected `,`");
            return None;
        }
    }

    if entries.len() > phf::MAX_SIZE {
        cx.span_err(parser.span,
                    format!("maps with more than {} entries are not supported",
                            phf::MAX_SIZE));
        return None;
    }

    if bad {
        return None;
    }

    Some(entries)
}

fn parse_set(cx: &mut ExtCtxt, tts: &[TokenTree]) -> Option<Vec<Entry>> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(),
                                                Vec::from_slice(tts));
    let mut entries = Vec::new();
    let value = quote_expr!(&*cx, ());

    let mut bad = false;
    while parser.token != EOF {
        let key = cx.expand_expr(parser.parse_expr());
        let key_str = parse_str(cx, key).unwrap_or_else(|| {
            bad = true;
            InternedString::new("")
        });

        entries.push(Entry {
            key_str: key_str,
            key: key,
            value: value,
        });

        if !parser.eat(&COMMA) && parser.token != EOF {
            cx.span_err(parser.span, "expected `,`");
            return None;
        }
    }

    if entries.len() > phf::MAX_SIZE {
        cx.span_err(parser.span,
                    format!("maps with more than {} entries are not supported",
                            phf::MAX_SIZE));
        return None;
    }

    if bad {
        return None;
    }

    Some(entries)
}

fn parse_str(cx: &mut ExtCtxt, e: &Expr) -> Option<InternedString> {
    match e.node {
        ExprLit(lit) => {
            match lit.node {
                LitStr(ref s, _) => Some(s.clone()),
                _ => {
                    cx.span_err(e.span, "expected string literal");
                    None
                }
            }
        }
        _ => {
            cx.span_err(e.span, "expected string literal");
            None
        }
    }
}

fn has_duplicates(cx: &mut ExtCtxt, sp: Span, entries: &[Entry]) -> bool {
    let mut dups = false;
    let mut strings = HashMap::new();
    for entry in entries.iter() {
        strings.insert_or_update_with(entry.key_str.clone(), (entry, true),
                                      |_, &(orig, ref mut first)| {
            if *first {
                cx.span_err(sp, format!("duplicate key \"{}\"",
                                        entry.key_str));
                cx.span_note(orig.key.span, "one occurrence here");
                *first = false;
            }
            cx.span_note(entry.key.span, "one occurrence here");
            dups = true;
        });
    }

    dups
}

fn generate_hash(cx: &mut ExtCtxt, sp: Span, entries: &[Entry]) -> HashState {
    let mut rng: XorShiftRng = SeedableRng::from_seed(FIXED_SEED);
    let start = time::precise_time_s();
    let state;
    loop {
        match try_generate_hash(entries.as_slice(), &mut rng) {
            Some(s) => {
                state = s;
                break;
            }
            None => {}
        }
    }
    let time = time::precise_time_s() - start;
    if os::getenv("PHF_STATS").is_some() {
        cx.span_note(sp, format!("PHF generation took {} seconds", time));
    }

    state
}

fn try_generate_hash(entries: &[Entry], rng: &mut XorShiftRng)
                     -> Option<HashState> {
    struct Bucket {
        idx: uint,
        keys: Vec<uint>,
    }

    struct Hashes {
        g: uint,
        f1: uint,
        f2: uint,
    }

    let k1 = rng.gen();
    let k2 = rng.gen();

    let hashes: Vec<Hashes> = entries.iter().map(|entry| {
        let (g, f1, f2) = phf::hash(entry.key_str.get(), k1, k2);
        Hashes {
            g: g,
            f1: f1,
            f2: f2
        }
    }).collect();

    let buckets_len = (entries.len() + DEFAULT_LAMBDA - 1) / DEFAULT_LAMBDA;
    let mut buckets = Vec::from_fn(buckets_len,
                                   |i| Bucket { idx: i, keys: Vec::new() });

    for (i, hash) in hashes.iter().enumerate() {
        buckets.get_mut(hash.g % buckets_len).keys.push(i);
    }

    // Sort descending
    buckets.sort_by(|a, b| b.keys.len().cmp(&a.keys.len()));

    let table_len = entries.len();
    let mut map = Vec::from_elem(table_len, None);
    let mut disps = Vec::from_elem(buckets_len, (0u, 0u));
    let mut try_map = HashMap::new();
    'buckets: for bucket in buckets.iter() {
        for d1 in range(0, table_len) {
            'disps_l: for d2 in range(0, table_len) {
                try_map.clear();
                for &key in bucket.keys.iter() {
                    let idx = phf::displace(hashes.get(key).f1,
                                            hashes.get(key).f2,
                                            d1,
                                            d2) % table_len;
                    if map.get(idx).is_some() || try_map.find(&idx).is_some() {
                        continue 'disps_l;
                    }
                    try_map.insert(idx, key);
                }

                // We've picked a good set of disps
                *disps.get_mut(bucket.idx) = (d1, d2);
                for (&idx, &key) in try_map.iter() {
                    *map.get_mut(idx) = Some(key);
                }
                continue 'buckets;
            }
        }

        // Unable to find displacements for a bucket
        return None;
    }

    Some(HashState {
        k1: k1,
        k2: k2,
        disps: disps,
        map: map.move_iter().map(|i| i.unwrap()).collect(),
    })
}

fn create_map(cx: &mut ExtCtxt, sp: Span, entries: Vec<Entry>, state: HashState)
              -> Box<MacResult> {
    let disps = state.disps.iter().map(|&(d1, d2)| {
        quote_expr!(&*cx, ($d1, $d2))
    }).collect();
    let disps = create_slice_expr(disps, sp);

    let entries = state.map.iter().map(|&idx| {
        let &Entry { key, value, .. } = entries.get(idx);
        quote_expr!(&*cx, ($key, $value))
    }).collect();
    let entries = create_slice_expr(entries, sp);

    let k1 = state.k1;
    let k2 = state.k2;
    MacExpr::new(quote_expr!(cx, PhfMap {
        k1: $k1,
        k2: $k2,
        disps: &'static $disps,
        entries: &'static $entries,
    }))
}

fn create_set(cx: &mut ExtCtxt, sp: Span, entries: Vec<Entry>, state: HashState)
              -> Box<MacResult> {
    let map = create_map(cx, sp, entries, state).make_expr().unwrap();
    MacExpr::new(quote_expr!(cx, PhfSet { map: $map }))
}

fn create_ordered_map(cx: &mut ExtCtxt, sp: Span, entries: Vec<Entry>,
                      state: HashState) -> Box<MacResult> {
    let disps = state.disps.iter().map(|&(d1, d2)| {
        quote_expr!(&*cx, ($d1, $d2))
    }).collect();
    let disps = create_slice_expr(disps, sp);

    let idxs = state.map.iter().map(|&idx| quote_expr!(&*cx, $idx)).collect();
    let idxs = create_slice_expr(idxs, sp);

    let entries = entries.iter().map(|&Entry { key, value, .. }| {
        quote_expr!(&*cx, ($key, $value))
    }).collect();
    let entries = create_slice_expr(entries, sp);

    let k1 = state.k1;
    let k2 = state.k2;
    MacExpr::new(quote_expr!(cx, PhfOrderedMap {
        k1: $k1,
        k2: $k2,
        disps: &'static $disps,
        idxs: &'static $idxs,
        entries: &'static $entries,
    }))
}

fn create_ordered_set(cx: &mut ExtCtxt, sp: Span, entries: Vec<Entry>,
                      state: HashState) -> Box<MacResult> {
    let map = create_ordered_map(cx, sp, entries, state).make_expr().unwrap();
    MacExpr::new(quote_expr!(cx, PhfOrderedSet { map: $map }))
}

fn create_slice_expr(vec: Vec<@Expr>, sp: Span) -> @Expr {
    @Expr {
        id: ast::DUMMY_NODE_ID,
        node: ExprVec(vec),
        span: sp
    }
}
<commit_msg>Update for format! changes<commit_after>//! Compiler plugin for Rust-PHF
//!
//! See the documentation for the `phf` crate for more details.
#![crate_id="github.com/sfackler/rust-phf/phf_mac"]
#![crate_type="dylib"]
#![doc(html_root_url="http://sfackler.github.io/rust-phf/doc")]
#![feature(managed_boxes, macro_registrar, quote)]

extern crate collections;
extern crate rand;
extern crate syntax;
extern crate time;
extern crate phf;

use collections::HashMap;
use std::os;
use syntax::ast;
use syntax::ast::{Name, TokenTree, LitStr, Expr, ExprVec, ExprLit};
use syntax::codemap::Span;
use syntax::ext::base::{SyntaxExtension,
                        DummyResult,
                        ExtCtxt,
                        MacResult,
                        MacExpr,
                        NormalTT,
                        BasicMacroExpander};
use syntax::parse;
use syntax::parse::token;
use syntax::parse::token::{InternedString, COMMA, EOF, FAT_ARROW};
use rand::{Rng, SeedableRng, XorShiftRng};

static DEFAULT_LAMBDA: uint = 5;

static FIXED_SEED: [u32, ..4] = [3141592653, 589793238, 462643383, 2795028841];

#[macro_registrar]
#[doc(hidden)]
pub fn macro_registrar(register: |Name, SyntaxExtension|) {
    let reg = |name, fn_| {
        register(token::intern(name),
                 NormalTT(box BasicMacroExpander {
                     expander: fn_,
                     span: None
                 },
                 None));
    };
    reg("phf_map", expand_phf_map);
    reg("phf_set", expand_phf_set);
    reg("phf_ordered_map", expand_phf_ordered_map);
    reg("phf_ordered_set", expand_phf_ordered_set);
}

struct Entry {
    key_str: InternedString,
    key: @Expr,
    value: @Expr
}

struct HashState {
    k1: u64,
    k2: u64,
    disps: Vec<(uint, uint)>,
    map: Vec<uint>,
}

fn expand_phf_map(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree])
                  -> Box<MacResult> {
    let entries = match parse_map(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp)
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, entries.as_slice());

    create_map(cx, sp, entries, state)
}

fn expand_phf_set(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree])
                  -> Box<MacResult> {
    let entries = match parse_set(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp)
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, entries.as_slice());

    create_set(cx, sp, entries, state)
}

fn expand_phf_ordered_map(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree])
                          -> Box<MacResult> {
    let entries = match parse_map(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp),
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, entries.as_slice());

    create_ordered_map(cx, sp, entries, state)
}

fn expand_phf_ordered_set(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree])
                          -> Box<MacResult> {
    let entries = match parse_set(cx, tts) {
        Some(entries) => entries,
        None => return DummyResult::expr(sp)
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return DummyResult::expr(sp);
    }

    let state = generate_hash(cx, sp, entries.as_slice());

    create_ordered_set(cx, sp, entries, state)
}

fn parse_map(cx: &mut ExtCtxt, tts: &[TokenTree]) -> Option<Vec<Entry>> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(),
                                                Vec::from_slice(tts));
    let mut entries = Vec::new();

    let mut bad = false;
    while parser.token != EOF {
        let key = cx.expand_expr(parser.parse_expr());
        let key_str = parse_str(cx, key).unwrap_or_else(|| {
            bad = true;
            InternedString::new("")
        });

        if !parser.eat(&FAT_ARROW) {
            cx.span_err(parser.span, "expected `=>`");
            return None;
        }

        let value = parser.parse_expr();

        entries.push(Entry {
            key_str: key_str,
            key: key,
            value: value
        });

        if !parser.eat(&COMMA) && parser.token != EOF {
            cx.span_err(parser.span, "expected `,`");
            return None;
        }
    }

    if entries.len() > phf::MAX_SIZE {
        cx.span_err(parser.span,
                    format!("maps with more than {} entries are not supported",
                            phf::MAX_SIZE).as_slice());
        return None;
    }

    if bad {
        return None;
    }

    Some(entries)
}

fn parse_set(cx: &mut ExtCtxt, tts: &[TokenTree]) -> Option<Vec<Entry>> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(),
                                                Vec::from_slice(tts));
    let mut entries = Vec::new();
    let value = quote_expr!(&*cx, ());

    let mut bad = false;
    while parser.token != EOF {
        let key = cx.expand_expr(parser.parse_expr());
        let key_str = parse_str(cx, key).unwrap_or_else(|| {
            bad = true;
            InternedString::new("")
        });

        entries.push(Entry {
            key_str: key_str,
            key: key,
            value: value,
        });

        if !parser.eat(&COMMA) && parser.token != EOF {
            cx.span_err(parser.span, "expected `,`");
            return None;
        }
    }

    if entries.len() > phf::MAX_SIZE {
        cx.span_err(parser.span,
                    format!("maps with more than {} entries are not supported",
                            phf::MAX_SIZE).as_slice());
        return None;
    }

    if bad {
        return None;
    }

    Some(entries)
}

fn parse_str(cx: &mut ExtCtxt, e: &Expr) -> Option<InternedString> {
    match e.node {
        ExprLit(lit) => {
            match lit.node {
                LitStr(ref s, _) => Some(s.clone()),
                _ => {
                    cx.span_err(e.span, "expected string literal");
                    None
                }
            }
        }
        _ => {
            cx.span_err(e.span, "expected string literal");
            None
        }
    }
}

fn has_duplicates(cx: &mut ExtCtxt, sp: Span, entries: &[Entry]) -> bool {
    let mut dups = false;
    let mut strings = HashMap::new();
    for entry in entries.iter() {
        strings.insert_or_update_with(entry.key_str.clone(), (entry, true),
                                      |_, &(orig, ref mut first)| {
            if *first {
                cx.span_err(sp, format!("duplicate key \"{}\"",
                                        entry.key_str).as_slice());
                cx.span_note(orig.key.span, "one occurrence here");
                *first = false;
            }
            cx.span_note(entry.key.span, "one occurrence here");
            dups = true;
        });
    }

    dups
}

fn generate_hash(cx: &mut ExtCtxt, sp: Span, entries: &[Entry]) -> HashState {
    let mut rng: XorShiftRng = SeedableRng::from_seed(FIXED_SEED);
    let start = time::precise_time_s();
    let state;
    loop {
        match try_generate_hash(entries.as_slice(), &mut rng) {
            Some(s) => {
                state = s;
                break;
            }
            None => {}
        }
    }
    let time = time::precise_time_s() - start;
    if os::getenv("PHF_STATS").is_some() {
        cx.span_note(sp, format!("PHF generation took {} seconds", time)
                            .as_slice());
    }

    state
}

fn try_generate_hash(entries: &[Entry], rng: &mut XorShiftRng)
                     -> Option<HashState> {
    struct Bucket {
        idx: uint,
        keys: Vec<uint>,
    }

    struct Hashes {
        g: uint,
        f1: uint,
        f2: uint,
    }

    let k1 = rng.gen();
    let k2 = rng.gen();

    let hashes: Vec<Hashes> = entries.iter().map(|entry| {
        let (g, f1, f2) = phf::hash(entry.key_str.get(), k1, k2);
        Hashes {
            g: g,
            f1: f1,
            f2: f2
        }
    }).collect();

    let buckets_len = (entries.len() + DEFAULT_LAMBDA - 1) / DEFAULT_LAMBDA;
    let mut buckets = Vec::from_fn(buckets_len,
                                   |i| Bucket { idx: i, keys: Vec::new() });

    for (i, hash) in hashes.iter().enumerate() {
        buckets.get_mut(hash.g % buckets_len).keys.push(i);
    }

    // Sort descending
    buckets.sort_by(|a, b| b.keys.len().cmp(&a.keys.len()));

    let table_len = entries.len();
    let mut map = Vec::from_elem(table_len, None);
    let mut disps = Vec::from_elem(buckets_len, (0u, 0u));
    let mut try_map = HashMap::new();
    'buckets: for bucket in buckets.iter() {
        for d1 in range(0, table_len) {
            'disps_l: for d2 in range(0, table_len) {
                try_map.clear();
                for &key in bucket.keys.iter() {
                    let idx = phf::displace(hashes.get(key).f1,
                                            hashes.get(key).f2,
                                            d1,
                                            d2) % table_len;
                    if map.get(idx).is_some() || try_map.find(&idx).is_some() {
                        continue 'disps_l;
                    }
                    try_map.insert(idx, key);
                }

                // We've picked a good set of disps
                *disps.get_mut(bucket.idx) = (d1, d2);
                for (&idx, &key) in try_map.iter() {
                    *map.get_mut(idx) = Some(key);
                }
                continue 'buckets;
            }
        }

        // Unable to find displacements for a bucket
        return None;
    }

    Some(HashState {
        k1: k1,
        k2: k2,
        disps: disps,
        map: map.move_iter().map(|i| i.unwrap()).collect(),
    })
}

fn create_map(cx: &mut ExtCtxt, sp: Span, entries: Vec<Entry>, state: HashState)
              -> Box<MacResult> {
    let disps = state.disps.iter().map(|&(d1, d2)| {
        quote_expr!(&*cx, ($d1, $d2))
    }).collect();
    let disps = create_slice_expr(disps, sp);

    let entries = state.map.iter().map(|&idx| {
        let &Entry { key, value, .. } = entries.get(idx);
        quote_expr!(&*cx, ($key, $value))
    }).collect();
    let entries = create_slice_expr(entries, sp);

    let k1 = state.k1;
    let k2 = state.k2;
    MacExpr::new(quote_expr!(cx, PhfMap {
        k1: $k1,
        k2: $k2,
        disps: &'static $disps,
        entries: &'static $entries,
    }))
}

fn create_set(cx: &mut ExtCtxt, sp: Span, entries: Vec<Entry>, state: HashState)
              -> Box<MacResult> {
    let map = create_map(cx, sp, entries, state).make_expr().unwrap();
    MacExpr::new(quote_expr!(cx, PhfSet { map: $map }))
}

fn create_ordered_map(cx: &mut ExtCtxt, sp: Span, entries: Vec<Entry>,
                      state: HashState) -> Box<MacResult> {
    let disps = state.disps.iter().map(|&(d1, d2)| {
        quote_expr!(&*cx, ($d1, $d2))
    }).collect();
    let disps = create_slice_expr(disps, sp);

    let idxs = state.map.iter().map(|&idx| quote_expr!(&*cx, $idx)).collect();
    let idxs = create_slice_expr(idxs, sp);

    let entries = entries.iter().map(|&Entry { key, value, .. }| {
        quote_expr!(&*cx, ($key, $value))
    }).collect();
    let entries = create_slice_expr(entries, sp);

    let k1 = state.k1;
    let k2 = state.k2;
    MacExpr::new(quote_expr!(cx, PhfOrderedMap {
        k1: $k1,
        k2: $k2,
        disps: &'static $disps,
        idxs: &'static $idxs,
        entries: &'static $entries,
    }))
}

fn create_ordered_set(cx: &mut ExtCtxt, sp: Span, entries: Vec<Entry>,
                      state: HashState) -> Box<MacResult> {
    let map = create_ordered_map(cx, sp, entries, state).make_expr().unwrap();
    MacExpr::new(quote_expr!(cx, PhfOrderedSet { map: $map }))
}

fn create_slice_expr(vec: Vec<@Expr>, sp: Span) -> @Expr {
    @Expr {
        id: ast::DUMMY_NODE_ID,
        node: ExprVec(vec),
        span: sp
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code, unused)]

use lift::{Higher, SemiGroup, Monoid, Functor, Monad};
use functor::*;
use monoid::*;
use semigroup::*;
<commit_msg>added monad<commit_after>#![allow(dead_code, unused)]

use lift::{Higher, SemiGroup, Monoid, Functor, Monad};
use functor::*;
use monoid::*;
use semigroup::*;
use monad::*;
<|endoftext|>


<commit_before>use command::Command;
use vm::VM;

pub struct Program {
    instructions: Vec<Command>,
    instruction_pointer: Option<usize>,
    current_depth: u64,
    pub status: ProgramStatus,
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum ProgramStatus {
    Normal,
    Seeking(u64),
}

#[derive(Clone, Copy, Debug)]
pub struct ProgramPatch {
    pub instruction_pointer: usize,
    pub current_depth: u64,
    pub status: ProgramStatus,
}

fn find_by<T, F>(list: &[T], start: Option<usize>, mut predicate: F) -> Option<usize>
where F: FnMut(&T) -> bool {
    let start = start.unwrap_or(0);
    for (offset, element) in list.iter().enumerate() {
        if predicate(element) { return Some(start + offset); }
    }

    return None;
}

impl Program {
    pub fn new() -> Program {
        Program {
            instructions: Vec::new(),
            instruction_pointer: None,
            current_depth: 0,
            status: ProgramStatus::Normal,
        }
    }

    pub fn append(&mut self, instructions: &[Command]) {
        self.instructions.extend(instructions.iter().cloned());
        if self.instruction_pointer.is_none() {
            self.instruction_pointer = Some(0);
        }
    }

    pub fn execute(&mut self, vm: &mut VM) {
        match (self.instruction_pointer, self.status) {
            (Some(index), ProgramStatus::Seeking(goal)) if index < self.instructions.len() => {

            }

            (Some(start), ProgramStatus::Normal) if start < self.instructions.len() => {
                for index in start..self.instructions.len() {
                }
            }

            _ => {}
        }
    }

    fn update (&mut self, patch: ProgramPatch) -> &mut Self {
        self.instruction_pointer = Some(patch.instruction_pointer);
        self.current_depth = patch.current_depth;
        self.status = patch.status;
        return self;
    }

    fn attempt_forward_jump(commands: &Vec<Command>, index: usize, depth: u64) -> ProgramPatch {
        let goal_depth = depth;
        let mut current_depth = depth;

        let index = find_by(commands, Some(index), |command| {
            match command {
                &Command::JumpBackward => current_depth -= 1,
                &Command::JumpForward => current_depth += 1,
                _ => {},
            };
            return current_depth == goal_depth;
        });

        return match index {
            Some(index) => ProgramPatch {
                instruction_pointer: index,
                status: ProgramStatus::Normal,
                current_depth
            },

            None => ProgramPatch {
                instruction_pointer: commands.len(),
                status: ProgramStatus::Seeking(goal_depth),
                current_depth,
            }
        };
    }
}
<commit_msg>Implement JumpForward handling in Program.execute<commit_after>use command::Command;
use vm::VM;

pub struct Program {
    instructions: Vec<Command>,
    instruction_pointer: Option<usize>,
    current_depth: u64,
    pub status: ProgramStatus,
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum ProgramStatus {
    Normal,
    Seeking(u64),
}

#[derive(Clone, Copy, Debug)]
pub struct ProgramPatch {
    pub instruction_pointer: usize,
    pub current_depth: u64,
    pub status: ProgramStatus,
}

fn find_by<T, F>(list: &[T], start: Option<usize>, mut predicate: F) -> Option<usize>
where F: FnMut(&T) -> bool {
    let start = start.unwrap_or(0);
    for (offset, element) in list.iter().enumerate() {
        if predicate(element) { return Some(start + offset); }
    }

    return None;
}

impl Program {
    pub fn new() -> Program {
        Program {
            instructions: Vec::new(),
            instruction_pointer: None,
            current_depth: 0,
            status: ProgramStatus::Normal,
        }
    }

    pub fn append(&mut self, instructions: &[Command]) {
        self.instructions.extend(instructions.iter().cloned());
        if self.instruction_pointer.is_none() {
            self.instruction_pointer = Some(0);
        }
    }

    pub fn execute(&mut self, vm: &mut VM) {
        match (self.instruction_pointer, self.status) {
            (Some(index), ProgramStatus::Seeking(goal)) if index < self.instructions.len() => {

            }

            (Some(start), ProgramStatus::Normal) if start < self.instructions.len() => {
                for index in start..self.instructions.len() {
                    match self.instructions[index] {
                        Command::JumpBackward => {},

                        Command::JumpForward => {
                            let patch = Program::attempt_forward_jump(&self.instructions, index, self.current_depth);
                            return self.update(patch).execute(vm);
                        },

                        command => vm.apply(command),
                    }
                }
            }

            _ => {}
        }
    }

    fn update (&mut self, patch: ProgramPatch) -> &mut Self {
        self.instruction_pointer = Some(patch.instruction_pointer);
        self.current_depth = patch.current_depth;
        self.status = patch.status;
        return self;
    }

    fn attempt_forward_jump(commands: &Vec<Command>, index: usize, depth: u64) -> ProgramPatch {
        let goal_depth = depth;
        let mut current_depth = depth;

        let index = find_by(commands, Some(index), |command| {
            match command {
                &Command::JumpBackward => current_depth -= 1,
                &Command::JumpForward => current_depth += 1,
                _ => {},
            };
            return current_depth == goal_depth;
        });

        return match index {
            Some(index) => ProgramPatch {
                instruction_pointer: index,
                status: ProgramStatus::Normal,
                current_depth
            },

            None => ProgramPatch {
                instruction_pointer: commands.len(),
                status: ProgramStatus::Seeking(goal_depth),
                current_depth,
            }
        };
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Amanieu d'Antras
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

use crate::raw_mutex::RawMutex;
use core::num::NonZeroUsize;
use lock_api::{self, GetThreadId};

/// Implementation of the `GetThreadId` trait for `lock_api::ReentrantMutex`.
pub struct RawThreadId;

unsafe impl GetThreadId for RawThreadId {
    const INIT: RawThreadId = RawThreadId;

    fn nonzero_thread_id(&self) -> NonZeroUsize {
        // The address of a thread-local variable is guaranteed to be unique to the
        // current thread, and is also guaranteed to be non-zero. The variable has to have a
        // non-zero size to guarantee it has a unique address for each thread.
        thread_local!(static KEY: u8 = 0);
        KEY.with(|x| {
            NonZeroUsize::new(x as *const _ as usize)
                .expect("thread-local variable address is null")
        })
    }
}

/// A mutex which can be recursively locked by a single thread.
///
/// This type is identical to `Mutex` except for the following points:
///
/// - Locking multiple times from the same thread will work correctly instead of
///   deadlocking.
/// - `ReentrantMutexGuard` does not give mutable references to the locked data.
///   Use a `RefCell` if you need this.
///
/// See [`Mutex`](struct.Mutex.html) for more details about the underlying mutex
/// primitive.
pub type ReentrantMutex<T> = lock_api::ReentrantMutex<RawMutex, RawThreadId, T>;

/// An RAII implementation of a "scoped lock" of a reentrant mutex. When this structure
/// is dropped (falls out of scope), the lock will be unlocked.
///
/// The data protected by the mutex can be accessed through this guard via its
/// `Deref` implementation.
pub type ReentrantMutexGuard<'a, T> = lock_api::ReentrantMutexGuard<'a, RawMutex, RawThreadId, T>;

/// An RAII mutex guard returned by `ReentrantMutexGuard::map`, which can point to a
/// subfield of the protected data.
///
/// The main difference between `MappedReentrantMutexGuard` and `ReentrantMutexGuard` is that the
/// former doesn't support temporarily unlocking and re-locking, since that
/// could introduce soundness issues if the locked object is modified by another
/// thread.
pub type MappedReentrantMutexGuard<'a, T> =
    lock_api::MappedReentrantMutexGuard<'a, RawMutex, RawThreadId, T>;

#[cfg(test)]
mod tests {
    use crate::ReentrantMutex;
    use std::cell::RefCell;
    use std::sync::Arc;
    use std::thread;

    #[cfg(feature = "serde")]
    use bincode::{deserialize, serialize};

    #[test]
    fn smoke() {
        let m = ReentrantMutex::new(());
        {
            let a = m.lock();
            {
                let b = m.lock();
                {
                    let c = m.lock();
                    assert_eq!(*c, ());
                }
                assert_eq!(*b, ());
            }
            assert_eq!(*a, ());
        }
    }

    #[test]
    fn is_mutex() {
        let m = Arc::new(ReentrantMutex::new(RefCell::new(0)));
        let m2 = m.clone();
        let lock = m.lock();
        let child = thread::spawn(move || {
            let lock = m2.lock();
            assert_eq!(*lock.borrow(), 4950);
        });
        for i in 0..100 {
            let lock = m.lock();
            *lock.borrow_mut() += i;
        }
        drop(lock);
        child.join().unwrap();
    }

    #[test]
    fn trylock_works() {
        let m = Arc::new(ReentrantMutex::new(()));
        let m2 = m.clone();
        let _lock = m.try_lock();
        let _lock2 = m.try_lock();
        thread::spawn(move || {
            let lock = m2.try_lock();
            assert!(lock.is_none());
        })
        .join()
        .unwrap();
        let _lock3 = m.try_lock();
    }

    #[test]
    fn test_reentrant_mutex_debug() {
        let mutex = ReentrantMutex::new(vec![0u8, 10]);

        assert_eq!(format!("{:?}", mutex), "ReentrantMutex { data: [0, 10] }");
    }

    #[cfg(feature = "serde")]
    #[test]
    fn test_serde() {
        let contents: Vec<u8> = vec![0, 1, 2];
        let mutex = ReentrantMutex::new(contents.clone());

        let serialized = serialize(&mutex).unwrap();
        let deserialized: ReentrantMutex<Vec<u8>> = deserialize(&serialized).unwrap();

        assert_eq!(*(mutex.lock()), *(deserialized.lock()));
        assert_eq!(contents, *(deserialized.lock()));
    }
}
<commit_msg>Test with integer instead of unit type<commit_after>// Copyright 2016 Amanieu d'Antras
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

use crate::raw_mutex::RawMutex;
use core::num::NonZeroUsize;
use lock_api::{self, GetThreadId};

/// Implementation of the `GetThreadId` trait for `lock_api::ReentrantMutex`.
pub struct RawThreadId;

unsafe impl GetThreadId for RawThreadId {
    const INIT: RawThreadId = RawThreadId;

    fn nonzero_thread_id(&self) -> NonZeroUsize {
        // The address of a thread-local variable is guaranteed to be unique to the
        // current thread, and is also guaranteed to be non-zero. The variable has to have a
        // non-zero size to guarantee it has a unique address for each thread.
        thread_local!(static KEY: u8 = 0);
        KEY.with(|x| {
            NonZeroUsize::new(x as *const _ as usize)
                .expect("thread-local variable address is null")
        })
    }
}

/// A mutex which can be recursively locked by a single thread.
///
/// This type is identical to `Mutex` except for the following points:
///
/// - Locking multiple times from the same thread will work correctly instead of
///   deadlocking.
/// - `ReentrantMutexGuard` does not give mutable references to the locked data.
///   Use a `RefCell` if you need this.
///
/// See [`Mutex`](struct.Mutex.html) for more details about the underlying mutex
/// primitive.
pub type ReentrantMutex<T> = lock_api::ReentrantMutex<RawMutex, RawThreadId, T>;

/// An RAII implementation of a "scoped lock" of a reentrant mutex. When this structure
/// is dropped (falls out of scope), the lock will be unlocked.
///
/// The data protected by the mutex can be accessed through this guard via its
/// `Deref` implementation.
pub type ReentrantMutexGuard<'a, T> = lock_api::ReentrantMutexGuard<'a, RawMutex, RawThreadId, T>;

/// An RAII mutex guard returned by `ReentrantMutexGuard::map`, which can point to a
/// subfield of the protected data.
///
/// The main difference between `MappedReentrantMutexGuard` and `ReentrantMutexGuard` is that the
/// former doesn't support temporarily unlocking and re-locking, since that
/// could introduce soundness issues if the locked object is modified by another
/// thread.
pub type MappedReentrantMutexGuard<'a, T> =
    lock_api::MappedReentrantMutexGuard<'a, RawMutex, RawThreadId, T>;

#[cfg(test)]
mod tests {
    use crate::ReentrantMutex;
    use std::cell::RefCell;
    use std::sync::Arc;
    use std::thread;

    #[cfg(feature = "serde")]
    use bincode::{deserialize, serialize};

    #[test]
    fn smoke() {
        let m = ReentrantMutex::new(2);
        {
            let a = m.lock();
            {
                let b = m.lock();
                {
                    let c = m.lock();
                    assert_eq!(*c, 2);
                }
                assert_eq!(*b, 2);
            }
            assert_eq!(*a, 2);
        }
    }

    #[test]
    fn is_mutex() {
        let m = Arc::new(ReentrantMutex::new(RefCell::new(0)));
        let m2 = m.clone();
        let lock = m.lock();
        let child = thread::spawn(move || {
            let lock = m2.lock();
            assert_eq!(*lock.borrow(), 4950);
        });
        for i in 0..100 {
            let lock = m.lock();
            *lock.borrow_mut() += i;
        }
        drop(lock);
        child.join().unwrap();
    }

    #[test]
    fn trylock_works() {
        let m = Arc::new(ReentrantMutex::new(()));
        let m2 = m.clone();
        let _lock = m.try_lock();
        let _lock2 = m.try_lock();
        thread::spawn(move || {
            let lock = m2.try_lock();
            assert!(lock.is_none());
        })
        .join()
        .unwrap();
        let _lock3 = m.try_lock();
    }

    #[test]
    fn test_reentrant_mutex_debug() {
        let mutex = ReentrantMutex::new(vec![0u8, 10]);

        assert_eq!(format!("{:?}", mutex), "ReentrantMutex { data: [0, 10] }");
    }

    #[cfg(feature = "serde")]
    #[test]
    fn test_serde() {
        let contents: Vec<u8> = vec![0, 1, 2];
        let mutex = ReentrantMutex::new(contents.clone());

        let serialized = serialize(&mutex).unwrap();
        let deserialized: ReentrantMutex<Vec<u8>> = deserialize(&serialized).unwrap();

        assert_eq!(*(mutex.lock()), *(deserialized.lock()));
        assert_eq!(contents, *(deserialized.lock()));
    }
}
<|endoftext|>


<commit_before>use ::syn::Ident;
use ::yaml_rust::Yaml;

use infer::TokensResult;

fn fn_ident(method: &str) -> Ident {
    Ident::new(method)
}

pub(super) fn infer_v3(path_st: &Ident, method: &str, schema: &Yaml) -> TokensResult {
    let method_fn = fn_ident(method);

    let tokens = quote! {
        impl Schema for #path_st {
            fn #method_fn(&self) {
                panic!("Not implemented!")
            }
        }
    };

    Ok(tokens)
}
<commit_msg>Validate HTTP method string<commit_after>use ::syn::Ident;
use ::yaml_rust::Yaml;

use infer::TokensResult;

const METHODS: &'static [&'static str] = &[
    "DELETE",
    "GET",
    "HEAD",
    "PATCH",
    "POST",
    "PUT",
];

pub(super) fn valid(method: &str) -> bool {
    METHODS.contains(&method.to_uppercase().as_str())
}

fn fn_ident(method: &str) -> Ident {
    assert!(valid(method), "Invalid method: {}", method);
    Ident::new(method.to_lowercase())
}

pub(super) fn infer_v3(path_st: &Ident, method: &str, schema: &Yaml) -> TokensResult {
    let method_fn = fn_ident(method);

    let tokens = quote! {
        impl Schema for #path_st {
            fn #method_fn(&self) {
                panic!("Not implemented!")
            }
        }
    };

    Ok(tokens)
}
<|endoftext|>


<commit_before>use std::env::home_dir;
use std::fs::{create_dir_all, OpenOptions, read_dir, remove_dir_all, remove_file, rename};
use std::io::Write;
use std::ffi::{OsStr, OsString};
use std::path::{Path, PathBuf};

use error::{ErrorKind, Result};
use process::{exec_command, run_command};

const DEFAULT_HOME_DIR: &str = ".monger";
const DEFAULT_BIN_DIR: &str = "mongodb-versions";
const DEFAULT_DB_DIR: &str = "db";

#[derive(Debug)]
pub struct Fs {
    home_dir: PathBuf,
    bin_dir: PathBuf,
    db_dir: PathBuf,
}

impl Fs {
    pub fn new() -> FsBuilder {
        Default::default()
    }

    #[inline]
    fn get_file<P: AsRef<Path>>(&self, filename: P) -> PathBuf {
        self.home_dir.join(filename)
    }

    #[inline]
    fn get_bin_file_rel<P: AsRef<Path>>(&self, filename: P) -> PathBuf {
        self.bin_dir.join(filename)
    }

    #[inline]
    fn get_bin_file_abs<P: AsRef<Path>>(&self, filename: P) -> PathBuf {
        self.get_file(self.get_bin_file_rel(filename))
    }

    #[inline]
    fn get_bin_dir(&self) -> PathBuf {
        self.get_file(self.bin_dir.as_path())
    }

    #[inline]
    fn get_db_file_rel<P: AsRef<Path>>(&self, filename: P) -> PathBuf {
        self.db_dir.join(filename)
    }

    #[inline]
    fn get_version_dir(&self, version: &str) -> PathBuf {
        self.get_bin_file_abs(version)
    }

    #[inline]
    fn get_version_bin_dir<P: AsRef<Path>>(&self, version: P) -> PathBuf {
        self.get_bin_file_abs(version.as_ref().join("bin"))
    }

    fn create(&self) -> Result<()> {
        create_dir_all(self.get_file(self.bin_dir.as_path()).as_path())?;
        create_dir_all(self.get_file(self.db_dir.as_path()).as_path())?;

        Ok(())
    }

    fn decompress_download<P: AsRef<Path>>(&self, filename: P, dirname: P, version: P) -> Result<()> {
        run_command(
            "tar",
            vec!["xf".as_ref(), filename.as_ref().as_os_str()],
            self.get_bin_dir(),
        )?;

        let old_name = self.get_bin_file_abs(dirname);
        let new_name = self.get_bin_file_abs(version);
        rename(old_name, new_name)?;

        Ok(())
    }

    fn delete_file<P: AsRef<Path>>(&self, filename: P) -> Result<()> {
        let path = self.home_dir.join(filename);

        if !path.exists() {
            return Ok(());
        }

        remove_file(path)?;

        Ok(())
    }

    fn delete_directory<P: AsRef<Path>>(&self, dirname: P) -> Result<bool> {
        let path = self.home_dir.join(dirname);

        if !path.exists() {
            return Ok(false);
        }

        remove_dir_all(path)?;

        Ok(true)
    }

    fn write_file<P: AsRef<Path>>(&self, filename: P, bytes: &[u8]) -> Result<()> {
        self.create()?;

        let filepath = self.home_dir.join(filename);
        let mut file = OpenOptions::new().write(true).create(true).open(filepath)?;
        file.write_all(bytes)?;

        Ok(())
    }

    #[inline]
    pub fn version_exists(&self, version: &str) -> bool {
        self.get_version_dir(version).is_dir()
    }

    pub fn create_or_get_db_dir(&self, version: &str) -> Result<PathBuf> {
        let db_dir = self.get_file(self.get_db_file_rel(version));
        create_dir_all(db_dir.as_path())?;
        Ok(db_dir)
    }

    pub fn delete_mongodb_version(&self, version: &str) -> Result<bool> {
        self.delete_directory(self.get_version_dir(version))
    }

    pub fn write_mongodb_download(
        &self,
        filename: &str,
        dirname: &str,
        bytes: &[u8],
        version: &str,
    ) -> Result<()> {
        let bin_file = self.get_bin_file_rel(filename);

        println!("writing {}...", bin_file.display());
        self.write_file(&bin_file, bytes)?;

        println!("decompressing...");
        self.decompress_download(filename, dirname, version)?;

        println!("cleaning up...");
        self.delete_file(&bin_file)?;

        Ok(())
    }

    pub fn list_versions(&self) -> Result<Vec<OsString>> {
        self.create()?;

        let mut versions = Vec::new();

        for e in read_dir(self.get_bin_dir())? {
            let entry = e?;

            if entry.file_type()?.is_dir() {
                versions.push(entry.file_name());
            }
        }

        Ok(versions)
    }

    pub fn exec<I, S>(&self, binary_name: &str, args: I, version: &str) -> Result<()>
    where
        I: IntoIterator<Item = S>,
        S: AsRef<OsStr>,
    {
        exec_command(
            &format!("./{}", binary_name),
            args,
            self.get_version_bin_dir(version),
        )
    }
}

#[derive(Debug, Default)]
pub struct FsBuilder {
    home_dir: Option<String>,
    bin_dir: Option<String>,
    db_dir: Option<String>,
}

impl FsBuilder {
    #[allow(dead_code)]
    pub fn with_home_dir(&mut self, home_dir: &str) -> &mut Self {
        self.home_dir = Some(home_dir.into());
        self
    }

    #[allow(dead_code)]
    pub fn with_bin_dir(&mut self, bin_dir: &str) -> &mut Self {
        self.bin_dir = Some(bin_dir.into());
        self
    }

    #[allow(dead_code)]
    pub fn with_db_dir(&mut self, db_dir: &str) -> &mut Self {
        self.db_dir = Some(db_dir.into());
        self
    }

    pub fn build(self) -> Result<Fs> {
        match home_dir() {
            Some(mut home_dir) => {
                home_dir.push(self.home_dir.unwrap_or(DEFAULT_HOME_DIR.into()));

                let bin_dir = Path::new(&self.bin_dir.unwrap_or(DEFAULT_BIN_DIR.into()))
                    .to_path_buf();

                let db_dir = Path::new(&self.db_dir.unwrap_or(DEFAULT_DB_DIR.into())).to_path_buf();

                Ok(Fs {
                    home_dir,
                    bin_dir,
                    db_dir,
                })
            }
            None => bail!(ErrorKind::UnknownHomeDirectory),
        }
    }
}
<commit_msg>format fix<commit_after>use std::env::home_dir;
use std::fs::{create_dir_all, OpenOptions, read_dir, remove_dir_all, remove_file, rename};
use std::io::Write;
use std::ffi::{OsStr, OsString};
use std::path::{Path, PathBuf};

use error::{ErrorKind, Result};
use process::{exec_command, run_command};

const DEFAULT_HOME_DIR: &str = ".monger";
const DEFAULT_BIN_DIR: &str = "mongodb-versions";
const DEFAULT_DB_DIR: &str = "db";

#[derive(Debug)]
pub struct Fs {
    home_dir: PathBuf,
    bin_dir: PathBuf,
    db_dir: PathBuf,
}

impl Fs {
    pub fn new() -> FsBuilder {
        Default::default()
    }

    #[inline]
    fn get_file<P: AsRef<Path>>(&self, filename: P) -> PathBuf {
        self.home_dir.join(filename)
    }

    #[inline]
    fn get_bin_file_rel<P: AsRef<Path>>(&self, filename: P) -> PathBuf {
        self.bin_dir.join(filename)
    }

    #[inline]
    fn get_bin_file_abs<P: AsRef<Path>>(&self, filename: P) -> PathBuf {
        self.get_file(self.get_bin_file_rel(filename))
    }

    #[inline]
    fn get_bin_dir(&self) -> PathBuf {
        self.get_file(self.bin_dir.as_path())
    }

    #[inline]
    fn get_db_file_rel<P: AsRef<Path>>(&self, filename: P) -> PathBuf {
        self.db_dir.join(filename)
    }

    #[inline]
    fn get_version_dir(&self, version: &str) -> PathBuf {
        self.get_bin_file_abs(version)
    }

    #[inline]
    fn get_version_bin_dir<P: AsRef<Path>>(&self, version: P) -> PathBuf {
        self.get_bin_file_abs(version.as_ref().join("bin"))
    }

    fn create(&self) -> Result<()> {
        create_dir_all(self.get_file(self.bin_dir.as_path()).as_path())?;
        create_dir_all(self.get_file(self.db_dir.as_path()).as_path())?;

        Ok(())
    }

    fn decompress_download<P: AsRef<Path>>(
        &self,
        filename: P,
        dirname: P,
        version: P,
    ) -> Result<()> {
        run_command(
            "tar",
            vec!["xf".as_ref(), filename.as_ref().as_os_str()],
            self.get_bin_dir(),
        )?;

        let old_name = self.get_bin_file_abs(dirname);
        let new_name = self.get_bin_file_abs(version);
        rename(old_name, new_name)?;

        Ok(())
    }

    fn delete_file<P: AsRef<Path>>(&self, filename: P) -> Result<()> {
        let path = self.home_dir.join(filename);

        if !path.exists() {
            return Ok(());
        }

        remove_file(path)?;

        Ok(())
    }

    fn delete_directory<P: AsRef<Path>>(&self, dirname: P) -> Result<bool> {
        let path = self.home_dir.join(dirname);

        if !path.exists() {
            return Ok(false);
        }

        remove_dir_all(path)?;

        Ok(true)
    }

    fn write_file<P: AsRef<Path>>(&self, filename: P, bytes: &[u8]) -> Result<()> {
        self.create()?;

        let filepath = self.home_dir.join(filename);
        let mut file = OpenOptions::new().write(true).create(true).open(filepath)?;
        file.write_all(bytes)?;

        Ok(())
    }

    #[inline]
    pub fn version_exists(&self, version: &str) -> bool {
        self.get_version_dir(version).is_dir()
    }

    pub fn create_or_get_db_dir(&self, version: &str) -> Result<PathBuf> {
        let db_dir = self.get_file(self.get_db_file_rel(version));
        create_dir_all(db_dir.as_path())?;
        Ok(db_dir)
    }

    pub fn delete_mongodb_version(&self, version: &str) -> Result<bool> {
        self.delete_directory(self.get_version_dir(version))
    }

    pub fn write_mongodb_download(
        &self,
        filename: &str,
        dirname: &str,
        bytes: &[u8],
        version: &str,
    ) -> Result<()> {
        let bin_file = self.get_bin_file_rel(filename);

        println!("writing {}...", bin_file.display());
        self.write_file(&bin_file, bytes)?;

        println!("decompressing...");
        self.decompress_download(filename, dirname, version)?;

        println!("cleaning up...");
        self.delete_file(&bin_file)?;

        Ok(())
    }

    pub fn list_versions(&self) -> Result<Vec<OsString>> {
        self.create()?;

        let mut versions = Vec::new();

        for e in read_dir(self.get_bin_dir())? {
            let entry = e?;

            if entry.file_type()?.is_dir() {
                versions.push(entry.file_name());
            }
        }

        Ok(versions)
    }

    pub fn exec<I, S>(&self, binary_name: &str, args: I, version: &str) -> Result<()>
    where
        I: IntoIterator<Item = S>,
        S: AsRef<OsStr>,
    {
        exec_command(
            &format!("./{}", binary_name),
            args,
            self.get_version_bin_dir(version),
        )
    }
}

#[derive(Debug, Default)]
pub struct FsBuilder {
    home_dir: Option<String>,
    bin_dir: Option<String>,
    db_dir: Option<String>,
}

impl FsBuilder {
    #[allow(dead_code)]
    pub fn with_home_dir(&mut self, home_dir: &str) -> &mut Self {
        self.home_dir = Some(home_dir.into());
        self
    }

    #[allow(dead_code)]
    pub fn with_bin_dir(&mut self, bin_dir: &str) -> &mut Self {
        self.bin_dir = Some(bin_dir.into());
        self
    }

    #[allow(dead_code)]
    pub fn with_db_dir(&mut self, db_dir: &str) -> &mut Self {
        self.db_dir = Some(db_dir.into());
        self
    }

    pub fn build(self) -> Result<Fs> {
        match home_dir() {
            Some(mut home_dir) => {
                home_dir.push(self.home_dir.unwrap_or(DEFAULT_HOME_DIR.into()));

                let bin_dir = Path::new(&self.bin_dir.unwrap_or(DEFAULT_BIN_DIR.into()))
                    .to_path_buf();

                let db_dir = Path::new(&self.db_dir.unwrap_or(DEFAULT_DB_DIR.into())).to_path_buf();

                Ok(Fs {
                    home_dir,
                    bin_dir,
                    db_dir,
                })
            }
            None => bail!(ErrorKind::UnknownHomeDirectory),
        }
    }
}
<|endoftext|>


<commit_before>// CITA
// Copyright 2016-2017 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation,
// either version 3 of the License, or (at your option) any
// later version.

// This program is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE. See the GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use bytes::BytesMut;
use citaprotocol::pubsub_message_to_network_message;
use config;
use config::NetConfig;
use libproto::{Message, OperateType};
use native_tls::{self, TlsConnector};
use notify::DebouncedEvent;
use std::collections::HashMap;
use std::convert::TryInto;
use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::net::SocketAddr;
use std::net::{Shutdown, TcpStream};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use tokio;
use tokio::prelude::*;
use tokio::util::FutureExt;
use tokio_tls;

const TIMEOUT: u64 = 15;
const ROOT_CERT_FILE: &str = "rootCA.crt";

pub enum RealStream {
    CryptStream(tokio_tls::TlsStream<tokio::net::TcpStream>),
    NormalStream(TcpStream),
}

impl RealStream {
    pub fn shutdown(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.shutdown().map(|_| ()),
            RealStream::NormalStream(ref mut tcp) => tcp.shutdown(Shutdown::Both),
        }
    }
}

impl io::Write for RealStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.write(buf),
            RealStream::NormalStream(ref mut tcp) => tcp.write(buf),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.flush(),
            RealStream::NormalStream(ref mut tcp) => tcp.flush(),
        }
    }
}

pub enum Task {
    Broadcast((String, Message)),
    Update(NetConfig),
    NewTCP((u32, SocketAddr, RealStream, String)),
}

/// Manager unconnected address
struct Manager {
    need_connect: Vec<(u32, SocketAddr, String)>,
    connect_receiver: Receiver<(u32, SocketAddr, String)>,
    task_sender: Sender<Task>,
    enable_tls: bool,
}

fn generate_tls_connector(path: &str) -> Option<TlsConnector> {
    let mut file = File::open(path).expect("Not has cert file");
    let mut pem = vec![];
    file.read_to_end(&mut pem).expect("Cert File read error");
    let root_ca =
        native_tls::Certificate::from_pem(&pem).expect("Cert file content not right, pem ?");

    native_tls::TlsConnector::builder()
        .min_protocol_version(Some(native_tls::Protocol::Tlsv11))
        .max_protocol_version(Some(native_tls::Protocol::Tlsv12))
        .add_root_certificate(root_ca)
        .build()
        .ok()
}

impl Manager {
    fn new(
        task_sender: Sender<Task>,
        connect_receiver: Receiver<(u32, SocketAddr, String)>,
        enable_tls: bool,
    ) -> Self {
        Manager {
            need_connect: Vec::new(),
            connect_receiver,
            task_sender,
            enable_tls,
        }
    }

    fn run(&mut self) {
        let (tls_connector, mut rt) = if self.enable_tls {
            let tls_connector = generate_tls_connector(ROOT_CERT_FILE);
            let rt = tokio::runtime::current_thread::Runtime::new().ok();
            if tls_connector.is_none() || rt.is_none() {
                panic!("TLS connector not generated");
            }
            (tls_connector, rt)
        } else {
            (None, None)
        };

        loop {
            while let Ok(message) = self.connect_receiver.try_recv() {
                self.need_connect.push(message);
            }
            let mut new_need_connect = Vec::new();
            while let Some((id, addr, common_name)) = self.need_connect.pop() {
                match tls_connector.clone() {
                    Some(tls_connect) => {
                        let common_name_clone = common_name.clone();
                        let task = tokio::net::TcpStream::connect(&addr)
                            .and_then(move |socket| {
                                tokio_tls::TlsConnector::from(tls_connect)
                                    .connect(&common_name_clone, socket)
                                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                            })
                            .timeout(Duration::from_secs(TIMEOUT));
                        match rt.as_mut().unwrap().block_on(task) {
                            Ok(tls) => {
                                self.task_sender
                                    .send(Task::NewTCP((
                                        id,
                                        addr,
                                        RealStream::CryptStream(tls),
                                        common_name.clone(),
                                    )))
                                    .unwrap();
                            }
                            Err(e) => {
                                warn!("TLS connect {} failed, error: {}", addr, e);
                                new_need_connect.push((id, addr, common_name));
                            }
                        };
                    }
                    None => match TcpStream::connect_timeout(&addr, Duration::from_secs(TIMEOUT)) {
                        Ok(tcp) => {
                            self.task_sender
                                .send(Task::NewTCP((
                                    id,
                                    addr,
                                    RealStream::NormalStream(tcp),
                                    common_name.clone(),
                                )))
                                .unwrap();
                        }
                        Err(e) => {
                            warn!(
                                "Node{}, {} unable to establish connection, error: {}",
                                id, addr, e
                            );
                            new_need_connect.push((id, addr, common_name));
                        }
                    },
                }
            }
            self.need_connect = new_need_connect;

            if !self.need_connect.is_empty() {
                debug!(
                    "Complete a round of attempts to connect, \
                     left {} address for the next round of processing",
                    self.need_connect.len()
                );
            }

            thread::sleep(Duration::from_secs(TIMEOUT));
        }
    }
}

/// Manage p2p networks
pub struct Connections {
    id_card: u32,
    /// list of peer: id, addr, tcp_connect
    peers: HashMap<(u32, SocketAddr, String), RealStream>,
    pub is_pause: Arc<AtomicBool>,
    pub connect_number: Arc<AtomicUsize>,
    task_receiver: Receiver<Task>,
    connect_sender: Sender<(u32, SocketAddr, String)>,
}

impl Connections {
    pub fn new(config: &config::NetConfig) -> (Self, Sender<Task>) {
        let id_card = config.id_card.unwrap();
        let (task_sender, task_receiver) = channel();
        let (connect_sender, connect_receiver) = channel();

        let connect_task_sender = task_sender.clone();
        let enable_tls = config.enable_tls.unwrap_or(false);
        thread::spawn(move || {
            Manager::new(connect_task_sender, connect_receiver, enable_tls).run()
        });

        if let Some(peers) = config.peers.as_ref() {
            for peer in peers.iter() {
                let id_card: u32 = peer.id_card.unwrap();
                let addr = format!("{}:{}", peer.ip.clone().unwrap(), peer.port.unwrap())
                    .parse()
                    .unwrap();
                connect_sender
                    .send((id_card, addr, peer.common_name.clone().unwrap_or_default()))
                    .unwrap();
            }
        }

        (
            Connections {
                id_card,
                peers: HashMap::new(),
                is_pause: Arc::new(AtomicBool::new(false)),
                connect_number: Arc::new(AtomicUsize::new(0)),
                task_receiver,
                connect_sender,
            },
            task_sender,
        )
    }

    fn is_send(id_card: u32, origin: u32, operate: OperateType) -> bool {
        match operate {
            OperateType::Broadcast => true,
            OperateType::Single => id_card == origin,
            OperateType::Subtract => id_card != origin,
        }
    }

    pub fn run(&mut self) {
        loop {
            match self
                .task_receiver
                .recv_timeout(Duration::from_secs(TIMEOUT))
            {
                Ok(task) => match task {
                    Task::Broadcast((key, message)) => {
                        if !self.is_pause.load(Ordering::SeqCst) {
                            self.broadcast(key, message)
                        }
                    }
                    Task::NewTCP(tcp) => {
                        self.peers.insert((tcp.0, tcp.1, tcp.3), tcp.2);
                        self.connect_number.fetch_add(1, Ordering::Relaxed);
                    }
                    Task::Update(config) => self.update(config),
                },
                Err(_) => {
                    self.heart_beat();
                }
            }
        }
    }

    fn update(&mut self, config: config::NetConfig) {
        // Update configuration
        match config.peers {
            Some(peers) => {
                let config_peers = peers
                    .into_iter()
                    .map(|peer| {
                        let id_card: u32 = peer.id_card.unwrap();
                        let addr = format!("{}:{}", peer.ip.unwrap(), peer.port.unwrap())
                            .parse()
                            .unwrap();
                        (id_card, addr, peer.common_name.unwrap_or_default())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                let remove_peers = self
                    .peers
                    .keys()
                    .filter(|peer| !config_peers.contains(&peer))
                    .map(|ref peer| {
                        info!("Remove peer {}, {},{}", peer.0, peer.1, peer.2);
                        (peer.0, peer.1, peer.2.clone())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                config_peers
                    .into_iter()
                    .filter(|peer| {
                        self.peers
                            .keys()
                            .find(|&current_peer| current_peer == peer)
                            .is_none()
                    })
                    .for_each(|peer| {
                        info!("Add peer {}, {}", peer.0, peer.1);
                        let _ = self.connect_sender.send(peer);
                    });

                self.close(Some(remove_peers), false);
            }
            None => {
                info!("clear all peers after update!");
                self.close::<Vec<(u32, SocketAddr, String)>>(None, false);
            }
        }
    }

    fn broadcast(&mut self, key: String, mut msg: Message) {
        let origin = msg.get_origin();
        let operate = msg.get_operate();
        msg.set_origin(self.id_card);

        trace!("Broadcast msg {:?} from key {}", msg, key);
        let msg_bytes: Vec<u8> = msg.try_into().unwrap();

        let mut buf = BytesMut::with_capacity(4 + 4 + 1 + key.len() + msg_bytes.len());
        pubsub_message_to_network_message(&mut buf, Some((key, msg_bytes)));

        let mut peers = Vec::new();
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            if Connections::is_send(peer.0, origin, operate) {
                match stream.write_all(&buf) {
                    Ok(_) => {
                        peers.push(peer.0);
                    }
                    Err(e) => {
                        warn!("Node{} {} is shutdown, err: {}", peer.0, peer.1, e);
                        remove_peers.push(peer.clone());
                    }
                }
            }
        }
        self.close(Some(remove_peers), true);
        trace!(
            "{:?} broadcast msg to nodes {:?} {:?}",
            self.id_card,
            operate,
            peers
        );
    }

    fn close<T: ::std::iter::IntoIterator<Item = (u32, SocketAddr, String)>>(
        &mut self,
        peers: Option<T>,
        need_reconnect: bool,
    ) {
        match peers {
            Some(peers) => peers.into_iter().for_each(|peer| {
                if let Some(mut stream) = self.peers.remove(&peer) {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                    if need_reconnect {
                        self.connect_sender.send(peer).unwrap();
                    }
                    self.connect_number.fetch_sub(1, Ordering::Relaxed);
                }
            }),
            None => {
                self.peers.iter_mut().for_each(|(peer, stream)| {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                });
                if need_reconnect {
                    self.peers.iter().for_each(|(ref peer, _)| {
                        self.connect_sender
                            .send((peer.0, peer.1, peer.2.clone()))
                            .unwrap();
                    })
                }
                self.peers.clear();
                self.connect_number.store(0, Ordering::Relaxed);
            }
        }
    }

    fn heart_beat(&mut self) {
        let mut buf = BytesMut::with_capacity(4 + 4);
        pubsub_message_to_network_message(&mut buf, None);
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            match stream.write(&buf) {
                Ok(_) => {
                    let _ = stream.flush();
                }
                Err(e) => {
                    warn!(
                        "Node{} {} {} is shutdown, err: {}",
                        peer.0, peer.1, peer.2, e
                    );
                    remove_peers.push((peer.0, peer.1, peer.2.clone()));
                }
            };
        }
        self.close(Some(remove_peers), true);
    }
}

pub fn manage_connect(config_path: &str, rx: Receiver<DebouncedEvent>, task_send: Sender<Task>) {
    let config = String::from(config_path);

    thread::spawn(move || loop {
        match rx.recv() {
            Ok(event) => match event {
                DebouncedEvent::Create(path_buf) | DebouncedEvent::Write(path_buf) => {
                    if path_buf.is_file() {
                        let file_name = path_buf.file_name().unwrap().to_str().unwrap();
                        if file_name == config.as_str() {
                            info!("file {} changed, will auto reload!", file_name);
                            let config = NetConfig::new(config.as_str());
                            let _ = task_send.send(Task::Update(config));
                        }
                    }
                }
                _ => trace!("file notify event: {:?}", event),
            },
            Err(e) => warn!("watch error: {:?}", e),
        }
    });
}

#[cfg(test)]
mod test {
    use super::Connections;
    use libproto::OperateType;
    #[test]
    fn is_send_msg() {
        assert!(Connections::is_send(0, 0, OperateType::Broadcast));
        assert!(Connections::is_send(0, 1, OperateType::Broadcast));

        assert!(Connections::is_send(0, 0, OperateType::Single));
        assert!(!Connections::is_send(0, 1, OperateType::Single));

        assert!(!Connections::is_send(0, 0, OperateType::Subtract));
        assert!(Connections::is_send(0, 1, OperateType::Subtract));
    }
}
<commit_msg>support hostname<commit_after>// CITA
// Copyright 2016-2017 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation,
// either version 3 of the License, or (at your option) any
// later version.

// This program is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE. See the GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use bytes::BytesMut;
use citaprotocol::pubsub_message_to_network_message;
use config;
use config::NetConfig;
use libproto::{Message, OperateType};
use native_tls::{self, TlsConnector};
use notify::DebouncedEvent;
use std::collections::HashMap;
use std::convert::TryInto;
use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::net::{Shutdown, TcpStream};
use std::net::{SocketAddr, ToSocketAddrs};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use tokio;
use tokio::prelude::*;
use tokio::util::FutureExt;
use tokio_tls;

const TIMEOUT: u64 = 15;
const ROOT_CERT_FILE: &str = "rootCA.crt";

pub enum RealStream {
    CryptStream(tokio_tls::TlsStream<tokio::net::TcpStream>),
    NormalStream(TcpStream),
}

impl RealStream {
    pub fn shutdown(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.shutdown().map(|_| ()),
            RealStream::NormalStream(ref mut tcp) => tcp.shutdown(Shutdown::Both),
        }
    }
}

impl io::Write for RealStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.write(buf),
            RealStream::NormalStream(ref mut tcp) => tcp.write(buf),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.flush(),
            RealStream::NormalStream(ref mut tcp) => tcp.flush(),
        }
    }
}

pub enum Task {
    Broadcast((String, Message)),
    Update(NetConfig),
    NewTCP((u32, SocketAddr, RealStream, String)),
}

/// Manager unconnected address
struct Manager {
    need_connect: Vec<(u32, SocketAddr, String)>,
    connect_receiver: Receiver<(u32, SocketAddr, String)>,
    task_sender: Sender<Task>,
    enable_tls: bool,
}

fn generate_tls_connector(path: &str) -> Option<TlsConnector> {
    let mut file = File::open(path).expect("Not has cert file");
    let mut pem = vec![];
    file.read_to_end(&mut pem).expect("Cert File read error");
    let root_ca =
        native_tls::Certificate::from_pem(&pem).expect("Cert file content not right, pem ?");

    native_tls::TlsConnector::builder()
        .min_protocol_version(Some(native_tls::Protocol::Tlsv11))
        .max_protocol_version(Some(native_tls::Protocol::Tlsv12))
        .add_root_certificate(root_ca)
        .build()
        .ok()
}

impl Manager {
    fn new(
        task_sender: Sender<Task>,
        connect_receiver: Receiver<(u32, SocketAddr, String)>,
        enable_tls: bool,
    ) -> Self {
        Manager {
            need_connect: Vec::new(),
            connect_receiver,
            task_sender,
            enable_tls,
        }
    }

    fn run(&mut self) {
        let (tls_connector, mut rt) = if self.enable_tls {
            let tls_connector = generate_tls_connector(ROOT_CERT_FILE);
            let rt = tokio::runtime::current_thread::Runtime::new().ok();
            if tls_connector.is_none() || rt.is_none() {
                panic!("TLS connector not generated");
            }
            (tls_connector, rt)
        } else {
            (None, None)
        };

        loop {
            while let Ok(message) = self.connect_receiver.try_recv() {
                self.need_connect.push(message);
            }
            let mut new_need_connect = Vec::new();
            while let Some((id, addr, common_name)) = self.need_connect.pop() {
                match tls_connector.clone() {
                    Some(tls_connect) => {
                        let common_name_clone = common_name.clone();
                        let task = tokio::net::TcpStream::connect(&addr)
                            .and_then(move |socket| {
                                tokio_tls::TlsConnector::from(tls_connect)
                                    .connect(&common_name_clone, socket)
                                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                            })
                            .timeout(Duration::from_secs(TIMEOUT));
                        match rt.as_mut().unwrap().block_on(task) {
                            Ok(tls) => {
                                self.task_sender
                                    .send(Task::NewTCP((
                                        id,
                                        addr,
                                        RealStream::CryptStream(tls),
                                        common_name.clone(),
                                    )))
                                    .unwrap();
                            }
                            Err(e) => {
                                warn!("TLS connect {} failed, error: {}", addr, e);
                                new_need_connect.push((id, addr, common_name));
                            }
                        };
                    }
                    None => match TcpStream::connect_timeout(&addr, Duration::from_secs(TIMEOUT)) {
                        Ok(tcp) => {
                            self.task_sender
                                .send(Task::NewTCP((
                                    id,
                                    addr,
                                    RealStream::NormalStream(tcp),
                                    common_name.clone(),
                                )))
                                .unwrap();
                        }
                        Err(e) => {
                            warn!(
                                "Node{}, {} unable to establish connection, error: {}",
                                id, addr, e
                            );
                            new_need_connect.push((id, addr, common_name));
                        }
                    },
                }
            }
            self.need_connect = new_need_connect;

            if !self.need_connect.is_empty() {
                debug!(
                    "Complete a round of attempts to connect, \
                     left {} address for the next round of processing",
                    self.need_connect.len()
                );
            }

            thread::sleep(Duration::from_secs(TIMEOUT));
        }
    }
}

/// Manage p2p networks
pub struct Connections {
    id_card: u32,
    /// list of peer: id, addr, tcp_connect
    peers: HashMap<(u32, SocketAddr, String), RealStream>,
    pub is_pause: Arc<AtomicBool>,
    pub connect_number: Arc<AtomicUsize>,
    task_receiver: Receiver<Task>,
    connect_sender: Sender<(u32, SocketAddr, String)>,
}

impl Connections {
    pub fn new(config: &config::NetConfig) -> (Self, Sender<Task>) {
        let id_card = config.id_card.unwrap();
        let (task_sender, task_receiver) = channel();
        let (connect_sender, connect_receiver) = channel();

        let connect_task_sender = task_sender.clone();
        let enable_tls = config.enable_tls.unwrap_or(false);
        thread::spawn(move || {
            Manager::new(connect_task_sender, connect_receiver, enable_tls).run()
        });

        if let Some(peers) = config.peers.as_ref() {
            for peer in peers.iter() {
                let id_card: u32 = peer.id_card.unwrap();
                let addr = format!("{}:{}", peer.ip.clone().unwrap(), peer.port.unwrap())
                    .to_socket_addrs()
                    .unwrap()
                    .next()
                    .unwrap();
                connect_sender
                    .send((id_card, addr, peer.common_name.clone().unwrap_or_default()))
                    .unwrap();
            }
        }

        (
            Connections {
                id_card,
                peers: HashMap::new(),
                is_pause: Arc::new(AtomicBool::new(false)),
                connect_number: Arc::new(AtomicUsize::new(0)),
                task_receiver,
                connect_sender,
            },
            task_sender,
        )
    }

    fn is_send(id_card: u32, origin: u32, operate: OperateType) -> bool {
        match operate {
            OperateType::Broadcast => true,
            OperateType::Single => id_card == origin,
            OperateType::Subtract => id_card != origin,
        }
    }

    pub fn run(&mut self) {
        loop {
            match self
                .task_receiver
                .recv_timeout(Duration::from_secs(TIMEOUT))
            {
                Ok(task) => match task {
                    Task::Broadcast((key, message)) => {
                        if !self.is_pause.load(Ordering::SeqCst) {
                            self.broadcast(key, message)
                        }
                    }
                    Task::NewTCP(tcp) => {
                        self.peers.insert((tcp.0, tcp.1, tcp.3), tcp.2);
                        self.connect_number.fetch_add(1, Ordering::Relaxed);
                    }
                    Task::Update(config) => self.update(config),
                },
                Err(_) => {
                    self.heart_beat();
                }
            }
        }
    }

    fn update(&mut self, config: config::NetConfig) {
        // Update configuration
        match config.peers {
            Some(peers) => {
                let config_peers = peers
                    .into_iter()
                    .map(|peer| {
                        let id_card: u32 = peer.id_card.unwrap();
                        let addr = format!("{}:{}", peer.ip.unwrap(), peer.port.unwrap())
                            .to_socket_addrs()
                            .unwrap()
                            .next()
                            .unwrap();
                        (id_card, addr, peer.common_name.unwrap_or_default())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                let remove_peers = self
                    .peers
                    .keys()
                    .filter(|peer| !config_peers.contains(&peer))
                    .map(|ref peer| {
                        info!("Remove peer {}, {},{}", peer.0, peer.1, peer.2);
                        (peer.0, peer.1, peer.2.clone())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                config_peers
                    .into_iter()
                    .filter(|peer| {
                        self.peers
                            .keys()
                            .find(|&current_peer| current_peer == peer)
                            .is_none()
                    })
                    .for_each(|peer| {
                        info!("Add peer {}, {}", peer.0, peer.1);
                        let _ = self.connect_sender.send(peer);
                    });

                self.close(Some(remove_peers), false);
            }
            None => {
                info!("clear all peers after update!");
                self.close::<Vec<(u32, SocketAddr, String)>>(None, false);
            }
        }
    }

    fn broadcast(&mut self, key: String, mut msg: Message) {
        let origin = msg.get_origin();
        let operate = msg.get_operate();
        msg.set_origin(self.id_card);

        trace!("Broadcast msg {:?} from key {}", msg, key);
        let msg_bytes: Vec<u8> = msg.try_into().unwrap();

        let mut buf = BytesMut::with_capacity(4 + 4 + 1 + key.len() + msg_bytes.len());
        pubsub_message_to_network_message(&mut buf, Some((key, msg_bytes)));

        let mut peers = Vec::new();
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            if Connections::is_send(peer.0, origin, operate) {
                match stream.write_all(&buf) {
                    Ok(_) => {
                        peers.push(peer.0);
                    }
                    Err(e) => {
                        warn!("Node{} {} is shutdown, err: {}", peer.0, peer.1, e);
                        remove_peers.push(peer.clone());
                    }
                }
            }
        }
        self.close(Some(remove_peers), true);
        trace!(
            "{:?} broadcast msg to nodes {:?} {:?}",
            self.id_card,
            operate,
            peers
        );
    }

    fn close<T: ::std::iter::IntoIterator<Item = (u32, SocketAddr, String)>>(
        &mut self,
        peers: Option<T>,
        need_reconnect: bool,
    ) {
        match peers {
            Some(peers) => peers.into_iter().for_each(|peer| {
                if let Some(mut stream) = self.peers.remove(&peer) {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                    if need_reconnect {
                        self.connect_sender.send(peer).unwrap();
                    }
                    self.connect_number.fetch_sub(1, Ordering::Relaxed);
                }
            }),
            None => {
                self.peers.iter_mut().for_each(|(peer, stream)| {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                });
                if need_reconnect {
                    self.peers.iter().for_each(|(ref peer, _)| {
                        self.connect_sender
                            .send((peer.0, peer.1, peer.2.clone()))
                            .unwrap();
                    })
                }
                self.peers.clear();
                self.connect_number.store(0, Ordering::Relaxed);
            }
        }
    }

    fn heart_beat(&mut self) {
        let mut buf = BytesMut::with_capacity(4 + 4);
        pubsub_message_to_network_message(&mut buf, None);
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            match stream.write(&buf) {
                Ok(_) => {
                    let _ = stream.flush();
                }
                Err(e) => {
                    warn!(
                        "Node{} {} {} is shutdown, err: {}",
                        peer.0, peer.1, peer.2, e
                    );
                    remove_peers.push((peer.0, peer.1, peer.2.clone()));
                }
            };
        }
        self.close(Some(remove_peers), true);
    }
}

pub fn manage_connect(config_path: &str, rx: Receiver<DebouncedEvent>, task_send: Sender<Task>) {
    let config = String::from(config_path);

    thread::spawn(move || loop {
        match rx.recv() {
            Ok(event) => match event {
                DebouncedEvent::Create(path_buf) | DebouncedEvent::Write(path_buf) => {
                    if path_buf.is_file() {
                        let file_name = path_buf.file_name().unwrap().to_str().unwrap();
                        if file_name == config.as_str() {
                            info!("file {} changed, will auto reload!", file_name);
                            let config = NetConfig::new(config.as_str());
                            let _ = task_send.send(Task::Update(config));
                        }
                    }
                }
                _ => trace!("file notify event: {:?}", event),
            },
            Err(e) => warn!("watch error: {:?}", e),
        }
    });
}

#[cfg(test)]
mod test {
    use super::Connections;
    use libproto::OperateType;
    #[test]
    fn is_send_msg() {
        assert!(Connections::is_send(0, 0, OperateType::Broadcast));
        assert!(Connections::is_send(0, 1, OperateType::Broadcast));

        assert!(Connections::is_send(0, 0, OperateType::Single));
        assert!(!Connections::is_send(0, 1, OperateType::Single));

        assert!(!Connections::is_send(0, 0, OperateType::Subtract));
        assert!(Connections::is_send(0, 1, OperateType::Subtract));
    }
}
<|endoftext|>


<commit_before>use lalafell::error::*;
use lalafell::commands::prelude::*;
use lalafell::commands::MentionOrId;

use rand::{Rng, thread_rng};

use serenity::model::channel::{ChannelType, PermissionOverwrite, PermissionOverwriteType};
use serenity::model::permissions::Permissions;

use unicase::UniCase;

use std::sync::Arc;

#[derive(BotCommand)]
pub struct ReportCommand;

#[derive(Debug, StructOpt)]
#[structopt(about = "Create a private channel with a member")]
pub struct Params {
  #[structopt(help = "The member to assign the role to")]
  who: MentionOrId,
}

impl HasParams for ReportCommand {
  type Params = Params;
}

impl<'a> PublicChannelCommand<'a> for ReportCommand {
  fn run(&self, _: &Context, msg: &Message, guild_id: GuildId, _: Arc<RwLock<GuildChannel>>, params: &[&str]) -> CommandResult<'a> {
    let member = guild_id.member(&msg.author).chain_err(|| "could not get member")?;
    if !member.permissions().chain_err(|| "could not get permissions")?.manage_roles() {
      return Err(ExternalCommandFailure::default()
        .message(|e: CreateEmbed| e
          .title("Not enough permissions.")
          .description("You don't have enough permissions to use this command."))
        .wrap());
    }

    let params = self.params("report", params)?;

    let guild = guild_id.find().chain_err(|| "could not find guild in cache")?;

    let reports_name = UniCase::new("Reports");
    let category = guild.read().channels
      .iter()
      .find(|(_, x)| {
        let channel = x.read();
        channel.kind == ChannelType::Category && UniCase::new(channel.name.as_str()) == reports_name
      })
      .map(|(&x, _)| x);

    let everyone = match guild.read().roles.values().find(|r| r.name == "@everyone") {
      Some(r) => r.id,
      None => return Err("No `@everyone` role?".into()),
    };
    let moderator_name = UniCase::new("moderator");
    let moderator = match guild.read().roles.values().find(|r| UniCase::new(r.name.as_str()) == moderator_name) {
        Some(r) => r.id,
        None => return Err("No `moderator` role.".into()),
      };

    let deny_everyone = PermissionOverwrite {
      allow: Permissions::empty(),
      deny: Permissions::READ_MESSAGES | Permissions::READ_MESSAGE_HISTORY,
      kind: PermissionOverwriteType::Role(everyone),
    };

    let allow_moderators = PermissionOverwrite {
      allow: Permissions::READ_MESSAGES,
      deny: Permissions::empty(),
      kind: PermissionOverwriteType::Role(moderator),
    };

    let allow_reporter = PermissionOverwrite {
      allow: Permissions::READ_MESSAGES,
      deny: Permissions::empty(),
      kind: PermissionOverwriteType::Member(*params.who),
    };

    let chars: String = thread_rng().gen_ascii_chars().take(7).collect();
    let channel_name = format!("report_{}", chars);
    let channel = guild_id
      .create_channel(&channel_name, ChannelType::Text, category)
      .chain_err(|| "could not create channel")?;

    channel.create_permission(&deny_everyone).chain_err(|| "could not deny @everyone")?;
    channel.create_permission(&allow_moderators).chain_err(|| "could not allow moderators")?;
    channel.create_permission(&allow_reporter).chain_err(|| "could not allow reporter")?;

    Ok(CommandSuccess::default())
  }
}
<commit_msg>fix: don't deny message history in report channels<commit_after>use lalafell::error::*;
use lalafell::commands::prelude::*;
use lalafell::commands::MentionOrId;

use rand::{Rng, thread_rng};

use serenity::model::channel::{ChannelType, PermissionOverwrite, PermissionOverwriteType};
use serenity::model::permissions::Permissions;

use unicase::UniCase;

use std::sync::Arc;

#[derive(BotCommand)]
pub struct ReportCommand;

#[derive(Debug, StructOpt)]
#[structopt(about = "Create a private channel with a member")]
pub struct Params {
  #[structopt(help = "The member to assign the role to")]
  who: MentionOrId,
}

impl HasParams for ReportCommand {
  type Params = Params;
}

impl<'a> PublicChannelCommand<'a> for ReportCommand {
  fn run(&self, _: &Context, msg: &Message, guild_id: GuildId, _: Arc<RwLock<GuildChannel>>, params: &[&str]) -> CommandResult<'a> {
    let member = guild_id.member(&msg.author).chain_err(|| "could not get member")?;
    if !member.permissions().chain_err(|| "could not get permissions")?.manage_roles() {
      return Err(ExternalCommandFailure::default()
        .message(|e: CreateEmbed| e
          .title("Not enough permissions.")
          .description("You don't have enough permissions to use this command."))
        .wrap());
    }

    let params = self.params("report", params)?;

    let guild = guild_id.find().chain_err(|| "could not find guild in cache")?;

    let reports_name = UniCase::new("Reports");
    let category = guild.read().channels
      .iter()
      .find(|(_, x)| {
        let channel = x.read();
        channel.kind == ChannelType::Category && UniCase::new(channel.name.as_str()) == reports_name
      })
      .map(|(&x, _)| x);

    let everyone = match guild.read().roles.values().find(|r| r.name == "@everyone") {
      Some(r) => r.id,
      None => return Err("No `@everyone` role?".into()),
    };
    let moderator_name = UniCase::new("moderator");
    let moderator = match guild.read().roles.values().find(|r| UniCase::new(r.name.as_str()) == moderator_name) {
        Some(r) => r.id,
        None => return Err("No `moderator` role.".into()),
      };

    let deny_everyone = PermissionOverwrite {
      allow: Permissions::empty(),
      deny: Permissions::READ_MESSAGES,
      kind: PermissionOverwriteType::Role(everyone),
    };

    let allow_moderators = PermissionOverwrite {
      allow: Permissions::READ_MESSAGES,
      deny: Permissions::empty(),
      kind: PermissionOverwriteType::Role(moderator),
    };

    let allow_reporter = PermissionOverwrite {
      allow: Permissions::READ_MESSAGES,
      deny: Permissions::empty(),
      kind: PermissionOverwriteType::Member(*params.who),
    };

    let chars: String = thread_rng().gen_ascii_chars().take(7).collect();
    let channel_name = format!("report_{}", chars);
    let channel = guild_id
      .create_channel(&channel_name, ChannelType::Text, category)
      .chain_err(|| "could not create channel")?;

    channel.create_permission(&deny_everyone).chain_err(|| "could not deny @everyone")?;
    channel.create_permission(&allow_moderators).chain_err(|| "could not allow moderators")?;
    channel.create_permission(&allow_reporter).chain_err(|| "could not allow reporter")?;

    Ok(CommandSuccess::default())
  }
}
<|endoftext|>


<commit_before>use std::f32::consts::PI;

/// Collision

/// The maximum number of contact points between two convex shapes. Do
/// not change this value.
pub const MAX_MANIFOLD_POINTS: i32 = 2;

/// The maximum number of vertices on a convex polygon.
pub const MAX_POLYGON_VERTICES: i32 = 8;

/// The is used to flatten AABBs in the dynamic tree. This allows proxies
/// to move by a small amount without triggering a tree adjustment.
/// This is in meters.
const _AABB_EXTENSION: f32 = 0.1;

/// This is used to fatten AABBs in the dynamic tree. This is used to predict
/// the future position based on the current displacement.
/// This is a dimenionless multiplier.
const _AABB_MULTIPLIER: f32 = 2.0;

/// A small length used as a collision and constraint tolerance. Usually it is
/// chosen to be numerically significant, but visually insignificant.
const _LINEAR_SLOP: f32 = 0.005;

/// A small angle used as a collision and constraint tolerance. Usually it is
/// chosen to be numerically significant, but visually insignificant.
const _ANGULAR_SLOP: f32 = (2.0 / 180.0 * PI);

/// The radius of the polygon/edge shape skin. This should not be modified. Making
/// this smaller means polygons will have an insufficient buffer for continuous collision.
/// Making it larger may create artifacts for vertex collision.
pub const POLYGON_RADIUS: f32 = (2.0 * _LINEAR_SLOP);

/// Maximum number of sub-steps per contact in continuous physics simulation.
const _MAX_SUB_STEPS: i32 = 8;

/// Dynamics

/// Maximum number of contacts to be handled to solve a TOI impact.
const _MAX_TOI_CONTACTS: i32 = 32;

/// A velocity threshold for elastic collisions. Any collisions with a relative linear
/// velocity below this threshold will be treated as inelastic.
const _VELOCITY_THRESHOLD: f32 = 1.0;

/// The maximum linear position correction used when solving constraints. This helps to
/// prevent overshoot.
const _MAX_LINEAR_CORRECTION: f32 = 0.2;

/// The maximum linear velocity of a body. This limit is very large and is used
/// to prevent numerical problems. You shouldn't need to adjust this.
const _MAX_TRANSLATION: f32 = 2.0;
const _MAX_TRANSLATION_SQUARED: f32 = _MAX_TRANSLATION * _MAX_TRANSLATION;

/// The maximum linear velocity of a body. This limit is very large and is used
/// to prevent numerical problems. You shouldn't need to adjust this.
const _MAX_ROTATION: f32 = 0.5 * PI;
const _MAX_ROTATION_SQUARED: f32 = _MAX_ROTATION * _MAX_ROTATION;

/// This scale factor controls how fast overlap is resolved. Ideally this would be 1 so
/// that overlap is removed in one time step. However using values close to 1 often lead
/// to overshoot.
const _BAUMGARTE: f32 = 0.2;
const _TOI_BAUGARTE: f32 = 0.75;

/// Sleep

/// The time that a body must be still before it will go to sleep.
const _TIME_TO_SLEEP: f32 = 0.5;

/// A body cannot sleep if its linear velocity is above this tolerance.
const _LINEAR_SLEEP_TOLERANCE: f32 = 0.01;

/// A body cannot sleep if its angular velocity is above this tolerance.
const _ANGULAR_SLEEP_TOLERANCE: f32 = (2.0 / 180.0 * PI);
<commit_msg>Adding Max Angular Correction Constant<commit_after>use std::f32::consts::PI;

/// Collision

/// The maximum number of contact points between two convex shapes. Do
/// not change this value.
pub const MAX_MANIFOLD_POINTS: i32 = 2;

/// The maximum number of vertices on a convex polygon.
pub const MAX_POLYGON_VERTICES: i32 = 8;

/// The is used to flatten AABBs in the dynamic tree. This allows proxies
/// to move by a small amount without triggering a tree adjustment.
/// This is in meters.
const _AABB_EXTENSION: f32 = 0.1;

/// This is used to fatten AABBs in the dynamic tree. This is used to predict
/// the future position based on the current displacement.
/// This is a dimenionless multiplier.
const _AABB_MULTIPLIER: f32 = 2.0;

/// A small length used as a collision and constraint tolerance. Usually it is
/// chosen to be numerically significant, but visually insignificant.
const _LINEAR_SLOP: f32 = 0.005;

/// A small angle used as a collision and constraint tolerance. Usually it is
/// chosen to be numerically significant, but visually insignificant.
const _ANGULAR_SLOP: f32 = (2.0 / 180.0 * PI);

/// The radius of the polygon/edge shape skin. This should not be modified. Making
/// this smaller means polygons will have an insufficient buffer for continuous collision.
/// Making it larger may create artifacts for vertex collision.
pub const POLYGON_RADIUS: f32 = (2.0 * _LINEAR_SLOP);

/// Maximum number of sub-steps per contact in continuous physics simulation.
const _MAX_SUB_STEPS: i32 = 8;

/// Dynamics

/// Maximum number of contacts to be handled to solve a TOI impact.
const _MAX_TOI_CONTACTS: i32 = 32;

/// A velocity threshold for elastic collisions. Any collisions with a relative linear
/// velocity below this threshold will be treated as inelastic.
const _VELOCITY_THRESHOLD: f32 = 1.0;

/// The maximum linear position correction used when solving constraints. This helps to
/// prevent overshoot.
const _MAX_LINEAR_CORRECTION: f32 = 0.2;

/// The maximum angular position correction used when solving constraints. This helps to
/// prevent overshoot.
const _MAX_ANGULAR_CORRECTION: f32 = (8.0 / 180.0 * PI);

/// The maximum linear velocity of a body. This limit is very large and is used
/// to prevent numerical problems. You shouldn't need to adjust this.
const _MAX_TRANSLATION: f32 = 2.0;
const _MAX_TRANSLATION_SQUARED: f32 = _MAX_TRANSLATION * _MAX_TRANSLATION;

/// The maximum linear velocity of a body. This limit is very large and is used
/// to prevent numerical problems. You shouldn't need to adjust this.
const _MAX_ROTATION: f32 = 0.5 * PI;
const _MAX_ROTATION_SQUARED: f32 = _MAX_ROTATION * _MAX_ROTATION;

/// This scale factor controls how fast overlap is resolved. Ideally this would be 1 so
/// that overlap is removed in one time step. However using values close to 1 often lead
/// to overshoot.
const _BAUMGARTE: f32 = 0.2;
const _TOI_BAUGARTE: f32 = 0.75;

/// Sleep

/// The time that a body must be still before it will go to sleep.
const _TIME_TO_SLEEP: f32 = 0.5;

/// A body cannot sleep if its linear velocity is above this tolerance.
const _LINEAR_SLEEP_TOLERANCE: f32 = 0.01;

/// A body cannot sleep if its angular velocity is above this tolerance.
const _ANGULAR_SLEEP_TOLERANCE: f32 = (2.0 / 180.0 * PI);
<|endoftext|>