<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://www.rust-ci.org/sfackler/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,
                        ExtCtxt,
                        MacResult,
                        MRExpr,
                        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(~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]) -> MacResult {
    let entries = match parse_map(cx, tts) {
        Some(entries) => entries,
        None => return MacResult::dummy_expr(sp)
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return MacResult::dummy_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]) -> MacResult {
    let entries = match parse_set(cx, tts) {
        Some(entries) => entries,
        None => return MacResult::dummy_expr(sp)
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return MacResult::dummy_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])
                          -> MacResult {
    let entries = match parse_map(cx, tts) {
        Some(entries) => entries,
        None => return MacResult::dummy_expr(sp),
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return MacResult::dummy_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])
                          -> MacResult {
    let entries = match parse_set(cx, tts) {
        Some(entries) => entries,
        None => return MacResult::dummy_expr(sp)
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return MacResult::dummy_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;
        }
    }

    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)
              -> 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;
    MRExpr(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)
              -> MacResult {
    let map = match create_map(cx, sp, entries, state) {
        MRExpr(expr) => expr,
        _ => unreachable!(),
    };

    MRExpr(quote_expr!(cx, PhfSet { map: $map }))
}

fn create_ordered_map(cx: &mut ExtCtxt, sp: Span, entries: Vec<Entry>,
                      state: HashState) -> 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;
    MRExpr(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) -> MacResult {
    let map = match create_ordered_map(cx, sp, entries, state) {
        MRExpr(expr) => expr,
        _ => unreachable!(),
    };

    MRExpr(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>Actually do error checking in parse_set<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://www.rust-ci.org/sfackler/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,
                        ExtCtxt,
                        MacResult,
                        MRExpr,
                        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(~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]) -> MacResult {
    let entries = match parse_map(cx, tts) {
        Some(entries) => entries,
        None => return MacResult::dummy_expr(sp)
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return MacResult::dummy_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]) -> MacResult {
    let entries = match parse_set(cx, tts) {
        Some(entries) => entries,
        None => return MacResult::dummy_expr(sp)
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return MacResult::dummy_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])
                          -> MacResult {
    let entries = match parse_map(cx, tts) {
        Some(entries) => entries,
        None => return MacResult::dummy_expr(sp),
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return MacResult::dummy_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])
                          -> MacResult {
    let entries = match parse_set(cx, tts) {
        Some(entries) => entries,
        None => return MacResult::dummy_expr(sp)
    };

    if has_duplicates(cx, sp, entries.as_slice()) {
        return MacResult::dummy_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)
              -> 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;
    MRExpr(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)
              -> MacResult {
    let map = match create_map(cx, sp, entries, state) {
        MRExpr(expr) => expr,
        _ => unreachable!(),
    };

    MRExpr(quote_expr!(cx, PhfSet { map: $map }))
}

fn create_ordered_map(cx: &mut ExtCtxt, sp: Span, entries: Vec<Entry>,
                      state: HashState) -> 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;
    MRExpr(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) -> MacResult {
    let map = match create_ordered_map(cx, sp, entries, state) {
        MRExpr(expr) => expr,
        _ => unreachable!(),
    };

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

//! An implementation of `VatNetwork` for the common case of a client-server connection.

use capnp::message::ReaderOptions;
use capnp::capability::Promise;
use futures::{AsyncRead, AsyncWrite, FutureExt, TryFutureExt};
use futures::channel::oneshot;

use std::cell::RefCell;
use std::rc::{Rc, Weak};

pub type VatId = crate::rpc_twoparty_capnp::Side;

struct IncomingMessage {
    message: ::capnp::message::Reader<capnp::serialize::OwnedSegments>,
}

impl IncomingMessage {
    pub fn new(message: ::capnp::message::Reader<capnp::serialize::OwnedSegments>) -> IncomingMessage {
        IncomingMessage { message: message }
    }
}

impl crate::IncomingMessage for IncomingMessage {
    fn get_body<'a>(&'a self) -> ::capnp::Result<::capnp::any_pointer::Reader<'a>> {
        self.message.get_root()
    }
}

struct OutgoingMessage {
    message: ::capnp::message::Builder<::capnp::message::HeapAllocator>,
    sender: ::capnp_futures::Sender<Rc<::capnp::message::Builder<::capnp::message::HeapAllocator>>>,
}

impl crate::OutgoingMessage for OutgoingMessage {
    fn get_body<'a>(&'a mut self) -> ::capnp::Result<::capnp::any_pointer::Builder<'a>> {
        self.message.get_root()
    }

    fn get_body_as_reader<'a>(&'a self) -> ::capnp::Result<::capnp::any_pointer::Reader<'a>> {
        self.message.get_root_as_reader()
    }

    fn send(self: Box<Self>)
            ->
        (Promise<Rc<::capnp::message::Builder<::capnp::message::HeapAllocator>>, ::capnp::Error>,
         Rc<::capnp::message::Builder<::capnp::message::HeapAllocator>>)
    {
        let tmp = *self;
        let OutgoingMessage {message, mut sender} = tmp;
        let m = Rc::new(message);
        (Promise::from_future(sender.send(m.clone()).map_err(|e| e.into())), m)
    }

    fn take(self: Box<Self>)
            -> ::capnp::message::Builder<::capnp::message::HeapAllocator>
    {
        self.message
    }
}

struct ConnectionInner<T> where T: AsyncRead + 'static {
    input_stream: Rc<RefCell<Option<T>>>,
    sender: ::capnp_futures::Sender<Rc<::capnp::message::Builder<::capnp::message::HeapAllocator>>>,
    side: crate::rpc_twoparty_capnp::Side,
    receive_options: ReaderOptions,
    on_disconnect_fulfiller: Option<oneshot::Sender<()>>,
}

struct Connection<T> where T: AsyncRead + 'static {
    inner: Rc<RefCell<ConnectionInner<T>>>,
}

impl <T> Drop for ConnectionInner<T> where T: AsyncRead {
    fn drop(&mut self) {
        let maybe_fulfiller = ::std::mem::replace(&mut self.on_disconnect_fulfiller, None);
        match maybe_fulfiller {
            Some(fulfiller) => {
                let _ = fulfiller.send(());
            }
            None => unreachable!(),
        }
    }
}

impl <T> Connection<T> where T: AsyncRead {
    fn new(input_stream: T,
           sender: ::capnp_futures::Sender<Rc<::capnp::message::Builder<::capnp::message::HeapAllocator>>>,
           side: crate::rpc_twoparty_capnp::Side,
           receive_options: ReaderOptions,
           on_disconnect_fulfiller: oneshot::Sender<()>,
           ) -> Connection<T>
    {

        Connection {
            inner: Rc::new(RefCell::new(
                ConnectionInner {
                    input_stream: Rc::new(RefCell::new(Some(input_stream))),
                    sender: sender,
                    side: side,
                    receive_options: receive_options,
                    on_disconnect_fulfiller: Some(on_disconnect_fulfiller),
                })),
        }
    }
}

impl <T> crate::Connection<crate::rpc_twoparty_capnp::Side> for Connection<T>
    where T: AsyncRead + Unpin
{
    fn get_peer_vat_id(&self) -> crate::rpc_twoparty_capnp::Side {
        self.inner.borrow().side
    }

    fn new_outgoing_message(&mut self, _first_segment_word_size: u32) -> Box<dyn crate::OutgoingMessage> {
        Box::new(OutgoingMessage {
            message: ::capnp::message::Builder::new_default(),
            sender: self.inner.borrow().sender.clone(),
        })
    }

    fn receive_incoming_message(&mut self) -> Promise<Option<Box<dyn crate::IncomingMessage + 'static>>, ::capnp::Error> {
        let mut inner = self.inner.borrow_mut();
        let maybe_input_stream = ::std::mem::replace(&mut *inner.input_stream.borrow_mut(), None);
        let return_it_here = inner.input_stream.clone();
        match maybe_input_stream {
            Some(mut s) => {
                let receive_options = inner.receive_options;
                Promise::from_future(async move {
                    let maybe_message = ::capnp_futures::serialize::read_message(&mut s, receive_options).await?;
                    *return_it_here.borrow_mut() = Some(s);
                    Ok(maybe_message.map(|message|
                                         Box::new(IncomingMessage::new(message)) as Box<dyn crate::IncomingMessage>))
                })
            }
            None => {
                Promise::err(::capnp::Error::failed("this should not be possible".to_string()))
             //   unreachable!(),
            }
        }
    }

    fn shutdown(&mut self, result: ::capnp::Result<()>) -> Promise<(), ::capnp::Error> {
        Promise::from_future(self.inner.borrow_mut().sender.terminate(result).map_err(|e| e.into()))
    }
}

/// A vat networks with two parties, the client and the server.
pub struct VatNetwork<T> where T: AsyncRead + 'static + Unpin {
    connection: Option<Connection<T>>,

    // HACK
    weak_connection_inner: Weak<RefCell<ConnectionInner<T>>>,

    execution_driver: futures::future::Shared<Promise<(), ::capnp::Error>>,
    side: crate::rpc_twoparty_capnp::Side,
}

impl <T> VatNetwork<T> where T: AsyncRead + Unpin {
    /// Creates a new two-party vat network that will receive data on `input_stream` and send data on
    /// `output_stream`.
    ///
    /// `side` indicates whether this is the client or the server side of the connection. This has no
    /// effect on the data sent over the connection; it merely exists so that `RpcNetwork::bootstrap` knows
    /// whether to return the local or the remote bootstrap capability. `VatId` parameters like this one
    /// will make more sense once we have vat networks with more than two parties.
    ///
    /// The options in `receive_options` will be used when reading the messages that come in on `input_stream`.
    pub fn new<U>(input_stream: T,
               output_stream: U,
               side: crate::rpc_twoparty_capnp::Side,
               receive_options: ReaderOptions) -> VatNetwork<T>
        where U: AsyncWrite + 'static + Unpin,
    {

        let (fulfiller, disconnect_promise) = oneshot::channel();
        let disconnect_promise = disconnect_promise
            .map_err(|_| ::capnp::Error::disconnected("disconnected".into()));

        let (execution_driver, sender) = {
            let (tx, write_queue) = ::capnp_futures::write_queue(output_stream);

            // Don't use `.join()` here because we need to make sure to wait for `disconnect_promise` to
            // resolve even if `write_queue` resolves to an error.
            (Promise::from_future(
                write_queue
                    .then(move |r| disconnect_promise.then(move |_| futures::future::ready(r)).map_ok(|_| ()))).shared(),
             tx)
        };

        let connection = Connection::new(input_stream, sender, side, receive_options, fulfiller);
        let weak_inner = Rc::downgrade(&connection.inner);
        VatNetwork {
            connection: Some(connection),
            weak_connection_inner: weak_inner,
            execution_driver: execution_driver,
            side: side,
        }
    }
}

impl <T> crate::VatNetwork<VatId> for VatNetwork<T>
    where T: AsyncRead + Unpin
{
    fn connect(&mut self, host_id: VatId) -> Option<Box<dyn crate::Connection<VatId>>> {
        if host_id == self.side {
            None
        } else {
            let connection = ::std::mem::replace(&mut self.connection, None);
            match connection {
                Some(c) => {
                    Some(Box::new(c))
                } None => {
                    match self.weak_connection_inner.upgrade() {
                        Some(connection_inner) => {
                            Some(Box::new(Connection { inner: connection_inner }))
                        }
                        None => {
                            panic!("tried to reconnect a disconnected twoparty vat network.")
                        }
                    }
                }
            }
        }
    }

    fn accept(&mut self) -> Promise<Box<dyn crate::Connection<VatId>>, ::capnp::Error> {
        let connection = ::std::mem::replace(&mut self.connection, None);
        match connection {
            Some(c) => Promise::ok(Box::new(c) as Box<dyn crate::Connection<VatId>>),
            None => Promise::from_future(::futures::future::pending()),
        }
    }

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

//! An implementation of `VatNetwork` for the common case of a client-server connection.

use capnp::message::ReaderOptions;
use capnp::capability::Promise;
use futures::{AsyncRead, AsyncWrite, FutureExt, TryFutureExt};
use futures::channel::oneshot;

use std::cell::RefCell;
use std::rc::{Rc, Weak};

pub type VatId = crate::rpc_twoparty_capnp::Side;

struct IncomingMessage {
    message: ::capnp::message::Reader<capnp::serialize::OwnedSegments>,
}

impl IncomingMessage {
    pub fn new(message: ::capnp::message::Reader<capnp::serialize::OwnedSegments>) -> IncomingMessage {
        IncomingMessage { message: message }
    }
}

impl crate::IncomingMessage for IncomingMessage {
    fn get_body<'a>(&'a self) -> ::capnp::Result<::capnp::any_pointer::Reader<'a>> {
        self.message.get_root()
    }
}

struct OutgoingMessage {
    message: ::capnp::message::Builder<::capnp::message::HeapAllocator>,
    sender: ::capnp_futures::Sender<Rc<::capnp::message::Builder<::capnp::message::HeapAllocator>>>,
}

impl crate::OutgoingMessage for OutgoingMessage {
    fn get_body<'a>(&'a mut self) -> ::capnp::Result<::capnp::any_pointer::Builder<'a>> {
        self.message.get_root()
    }

    fn get_body_as_reader<'a>(&'a self) -> ::capnp::Result<::capnp::any_pointer::Reader<'a>> {
        self.message.get_root_as_reader()
    }

    fn send(self: Box<Self>)
            ->
        (Promise<Rc<::capnp::message::Builder<::capnp::message::HeapAllocator>>, ::capnp::Error>,
         Rc<::capnp::message::Builder<::capnp::message::HeapAllocator>>)
    {
        let tmp = *self;
        let OutgoingMessage {message, mut sender} = tmp;
        let m = Rc::new(message);
        (Promise::from_future(sender.send(m.clone()).map_err(|e| e.into())), m)
    }

    fn take(self: Box<Self>)
            -> ::capnp::message::Builder<::capnp::message::HeapAllocator>
    {
        self.message
    }
}

struct ConnectionInner<T> where T: AsyncRead + 'static {
    input_stream: Rc<RefCell<Option<T>>>,
    sender: ::capnp_futures::Sender<Rc<::capnp::message::Builder<::capnp::message::HeapAllocator>>>,
    side: crate::rpc_twoparty_capnp::Side,
    receive_options: ReaderOptions,
    on_disconnect_fulfiller: Option<oneshot::Sender<()>>,
}

struct Connection<T> where T: AsyncRead + 'static {
    inner: Rc<RefCell<ConnectionInner<T>>>,
}

impl <T> Drop for ConnectionInner<T> where T: AsyncRead {
    fn drop(&mut self) {
        let maybe_fulfiller = ::std::mem::replace(&mut self.on_disconnect_fulfiller, None);
        match maybe_fulfiller {
            Some(fulfiller) => {
                let _ = fulfiller.send(());
            }
            None => unreachable!(),
        }
    }
}

impl <T> Connection<T> where T: AsyncRead {
    fn new(input_stream: T,
           sender: ::capnp_futures::Sender<Rc<::capnp::message::Builder<::capnp::message::HeapAllocator>>>,
           side: crate::rpc_twoparty_capnp::Side,
           receive_options: ReaderOptions,
           on_disconnect_fulfiller: oneshot::Sender<()>,
           ) -> Connection<T>
    {

        Connection {
            inner: Rc::new(RefCell::new(
                ConnectionInner {
                    input_stream: Rc::new(RefCell::new(Some(input_stream))),
                    sender: sender,
                    side: side,
                    receive_options: receive_options,
                    on_disconnect_fulfiller: Some(on_disconnect_fulfiller),
                })),
        }
    }
}

impl <T> crate::Connection<crate::rpc_twoparty_capnp::Side> for Connection<T>
    where T: AsyncRead + Unpin
{
    fn get_peer_vat_id(&self) -> crate::rpc_twoparty_capnp::Side {
        self.inner.borrow().side
    }

    fn new_outgoing_message(&mut self, _first_segment_word_size: u32) -> Box<dyn crate::OutgoingMessage> {
        Box::new(OutgoingMessage {
            message: ::capnp::message::Builder::new_default(),
            sender: self.inner.borrow().sender.clone(),
        })
    }

    fn receive_incoming_message(&mut self) -> Promise<Option<Box<dyn crate::IncomingMessage + 'static>>, ::capnp::Error> {
        let mut inner = self.inner.borrow_mut();
        let maybe_input_stream = ::std::mem::replace(&mut *inner.input_stream.borrow_mut(), None);
        let return_it_here = inner.input_stream.clone();
        match maybe_input_stream {
            Some(mut s) => {
                let receive_options = inner.receive_options;
                Promise::from_future(async move {
                    let maybe_message = ::capnp_futures::serialize::read_message(&mut s, receive_options).await?;
                    *return_it_here.borrow_mut() = Some(s);
                    Ok(maybe_message.map(|message|
                                         Box::new(IncomingMessage::new(message)) as Box<dyn crate::IncomingMessage>))
                })
            }
            None => {
                Promise::err(::capnp::Error::failed("this should not be possible".to_string()))
             //   unreachable!(),
            }
        }
    }

    fn shutdown(&mut self, result: ::capnp::Result<()>) -> Promise<(), ::capnp::Error> {
        Promise::from_future(self.inner.borrow_mut().sender.terminate(result).map_err(|e| e.into()))
    }
}

/// A vat network with two parties, the client and the server.
pub struct VatNetwork<T> where T: AsyncRead + 'static + Unpin {
    connection: Option<Connection<T>>,

    // HACK
    weak_connection_inner: Weak<RefCell<ConnectionInner<T>>>,

    execution_driver: futures::future::Shared<Promise<(), ::capnp::Error>>,
    side: crate::rpc_twoparty_capnp::Side,
}

impl <T> VatNetwork<T> where T: AsyncRead + Unpin {
    /// Creates a new two-party vat network that will receive data on `input_stream` and send data on
    /// `output_stream`.
    ///
    /// `side` indicates whether this is the client or the server side of the connection. This has no
    /// effect on the data sent over the connection; it merely exists so that `RpcNetwork::bootstrap` knows
    /// whether to return the local or the remote bootstrap capability. `VatId` parameters like this one
    /// will make more sense once we have vat networks with more than two parties.
    ///
    /// The options in `receive_options` will be used when reading the messages that come in on `input_stream`.
    pub fn new<U>(input_stream: T,
               output_stream: U,
               side: crate::rpc_twoparty_capnp::Side,
               receive_options: ReaderOptions) -> VatNetwork<T>
        where U: AsyncWrite + 'static + Unpin,
    {

        let (fulfiller, disconnect_promise) = oneshot::channel();
        let disconnect_promise = disconnect_promise
            .map_err(|_| ::capnp::Error::disconnected("disconnected".into()));

        let (execution_driver, sender) = {
            let (tx, write_queue) = ::capnp_futures::write_queue(output_stream);

            // Don't use `.join()` here because we need to make sure to wait for `disconnect_promise` to
            // resolve even if `write_queue` resolves to an error.
            (Promise::from_future(
                write_queue
                    .then(move |r| disconnect_promise.then(move |_| futures::future::ready(r)).map_ok(|_| ()))).shared(),
             tx)
        };

        let connection = Connection::new(input_stream, sender, side, receive_options, fulfiller);
        let weak_inner = Rc::downgrade(&connection.inner);
        VatNetwork {
            connection: Some(connection),
            weak_connection_inner: weak_inner,
            execution_driver: execution_driver,
            side: side,
        }
    }
}

impl <T> crate::VatNetwork<VatId> for VatNetwork<T>
    where T: AsyncRead + Unpin
{
    fn connect(&mut self, host_id: VatId) -> Option<Box<dyn crate::Connection<VatId>>> {
        if host_id == self.side {
            None
        } else {
            let connection = ::std::mem::replace(&mut self.connection, None);
            match connection {
                Some(c) => {
                    Some(Box::new(c))
                } None => {
                    match self.weak_connection_inner.upgrade() {
                        Some(connection_inner) => {
                            Some(Box::new(Connection { inner: connection_inner }))
                        }
                        None => {
                            panic!("tried to reconnect a disconnected twoparty vat network.")
                        }
                    }
                }
            }
        }
    }

    fn accept(&mut self) -> Promise<Box<dyn crate::Connection<VatId>>, ::capnp::Error> {
        let connection = ::std::mem::replace(&mut self.connection, None);
        match connection {
            Some(c) => Promise::ok(Box::new(c) as Box<dyn crate::Connection<VatId>>),
            None => Promise::from_future(::futures::future::pending()),
        }
    }

    fn drive_until_shutdown(&mut self) -> Promise<(), ::capnp::Error> {
        Promise::from_future(self.execution_driver.clone())
    }
}
<|endoftext|>


<commit_before>use std::{fs, cmp};
use std::path::Path;
use std::io::{self, Read};

use reqwest::{Client as HClient, header};
use serde_json;
use prettytable::Table;
use url::Url;

use rpc::message::{CMessage, SMessage};
use rpc::criterion::{Criterion, Value, Operation};
use rpc::resource::{Resource, ResourceKind, SResourceUpdate};

use client::Client;
use error::{Result, ResultExt, ErrorKind};

pub fn add(
    mut c: Client,
    url: &str,
    files: Vec<&str>,
    dir: Option<&str>,
    start: bool,
) -> Result<()> {
    for file in files {
        add_file(&mut c, url, file, dir, start)?;
    }
    Ok(())
}

fn add_file(c: &mut Client, url: &str, file: &str, dir: Option<&str>, start: bool) -> Result<()> {
    let mut torrent = Vec::new();
    let mut f = fs::File::open(file).chain_err(|| ErrorKind::FileIO)?;
    f.read_to_end(&mut torrent).chain_err(|| ErrorKind::FileIO)?;

    let msg = CMessage::UploadTorrent {
        serial: c.next_serial(),
        size: torrent.len() as u64,
        path: dir.as_ref().map(|d| format!("{}", d)),
        start,
    };
    let token = if let SMessage::TransferOffer { token, .. } = c.rr(msg)? {
        token
    } else {
        bail!("Failed to receieve transfer offer from synapse!");
    };
    let client = HClient::new().chain_err(|| ErrorKind::HTTP)?;
    client
        .post(url)
        .chain_err(|| ErrorKind::HTTP)?
        .header(header::Authorization(header::Bearer { token }))
        .body(torrent)
        .send()
        .chain_err(|| ErrorKind::HTTP)?;

    if let SMessage::OResourcesExtant { .. } = c.recv()? {
    } else {
        bail!("Failed to receieve upload acknowledgement from synapse!");
    };

    Ok(())
}

pub fn del(mut c: Client, torrents: Vec<&str>) -> Result<()> {
    for torrent in torrents {
        del_torrent(&mut c, torrent)?;
    }
    Ok(())
}

fn del_torrent(c: &mut Client, torrent: &str) -> Result<()> {
    let resources = search_torrent_name(c, torrent)?;
    if resources.len() == 1 {
        let msg = CMessage::RemoveResource {
            serial: c.next_serial(),
            id: resources[0].id().to_owned(),
        };
        c.send(msg)?;
    } else if resources.is_empty() {
        eprintln!("Could not find any matching torrents for {}", torrent);
    } else {
        eprintln!(
            "Ambiguous results searching for {}. Potential alternatives include: ",
            torrent
        );
        for res in resources.into_iter().take(3) {
            if let Resource::Torrent(t) = res {
                eprintln!("{}", t.name);
            }
        }
    }
    Ok(())
}

pub fn dl(mut c: Client, url: &str, name: &str) -> Result<()> {
    let resources = search_torrent_name(&mut c, name)?;
    let files = if resources.len() == 1 {
        let msg = CMessage::FilterSubscribe {
            serial: c.next_serial(),
            kind: ResourceKind::File,
            criteria: vec![
                Criterion {
                    field: "torrent_id".to_owned(),
                    op: Operation::Eq,
                    value: Value::S(resources[0].id().to_owned()),
                },
            ],
        };
        if let SMessage::OResourcesExtant { ids, .. } = c.rr(msg)? {
            get_resources(&mut c, ids)?
        } else {
            bail!("Could not get files for torrent!");
        }
    } else if resources.is_empty() {
        eprintln!("Could not find any matching torrents for {}", name);
        return Ok(());
    } else {
        eprintln!(
            "Ambiguous results searching for {}. Potential alternatives include: ",
            name
        );
        for res in resources.into_iter().take(3) {
            if let Resource::Torrent(t) = res {
                eprintln!("{}", t.name);
            }
        }
        return Ok(());
    };

    for file in files {
        let mut dl_url = Url::parse(url).unwrap();
        dl_url.path_segments_mut().unwrap().push("dl").push(
            file.id(),
        );

        let client = HClient::new().chain_err(|| ErrorKind::HTTP)?;
        let mut resp = client
            .get(dl_url.as_str())
            .chain_err(|| ErrorKind::HTTP)?
            .send()
            .chain_err(|| ErrorKind::HTTP)?;
        if let Resource::File(f) = file {
            let p = Path::new(&f.path);
            if let Some(par) = p.parent() {
                fs::create_dir_all(par).chain_err(|| ErrorKind::FileIO)?;
            }
            let mut f = fs::File::create(p).chain_err(|| ErrorKind::FileIO)?;
            io::copy(&mut resp, &mut f).chain_err(|| ErrorKind::FileIO)?;
        } else {
            bail!("Expected a file resource");
        }
    }
    Ok(())
}

pub fn list(mut c: Client, kind: &str, crit: Vec<Criterion>, output: &str) -> Result<()> {
    let k = match kind {
        "torrent" => ResourceKind::Torrent,
        "tracker" => ResourceKind::Tracker,
        "peer" => ResourceKind::Peer,
        "piece" => ResourceKind::Piece,
        "file" => ResourceKind::File,
        "server" => ResourceKind::Server,
        _ => bail!("Unexpected resource kind {}", kind),
    };
    let results = search(&mut c, k, crit)?;
    if output == "text" {
        let mut table = Table::new();
        match k {
            ResourceKind::Torrent => {
                table.add_row(row!["Name", "Done", "DL", "UL", "DL RT", "UL RT", "Peers"]);
            }
            ResourceKind::Tracker => {
                table.add_row(row!["URL", "Torrent", "Error"]);
            }
            ResourceKind::Peer => {
                table.add_row(row!["IP", "Torrent", "DL RT", "UL RT"]);
            }
            ResourceKind::Piece => {
                table.add_row(row!["Torrent", "DLd", "Avail"]);
            }
            ResourceKind::File => {
                table.add_row(row!["Path", "Torrent", "Done", "Prio", "Avail"]);
            }
            ResourceKind::Server => {
                table.add_row(row!["DL RT", "UL RT"]);
            }
        }

        #[cfg_attr(rustfmt, rustfmt_skip)]
        for res in results {
            match k {
                ResourceKind::Torrent => {
                    let t = res.as_torrent();
                    table.add_row(row![
                        t.name,
                        format!("{:.2}%", t.progress * 100.),
                        fmt_bytes(t.transferred_down as f64),
                        fmt_bytes(t.transferred_up as f64),
                        fmt_bytes(t.rate_down as f64) + "/s",
                        fmt_bytes(t.rate_up as f64) + "/s",
                        t.peers
                    ]);
                }
                ResourceKind::Tracker => {
                    let t = res.as_tracker();
                    table.add_row(row![
                        t.url,
                        t.torrent_id,
                        t.error.as_ref().map(|s| s.as_str()).unwrap_or("")
                    ]);
                }
                ResourceKind::Peer => {
                    let p = res.as_peer();
                    let rd = fmt_bytes(p.rate_down as f64) + "/s";
                    let ru = fmt_bytes(p.rate_up as f64) + "/s";
                    table.add_row(row![p.ip, p.torrent_id, rd, ru]);
                }
                ResourceKind::Piece => {
                    let p = res.as_piece();
                    table.add_row(row![p.torrent_id, p.downloaded, p.available]);
                }
                ResourceKind::File => {
                    let f = res.as_file();
                    table.add_row(row![
                        f.path,
                        f.torrent_id,
                        format!("{:.2}%", f.progress as f64 * 100.),
                        f.priority,
                        format!("{:.2}%", f.availability as f64 * 100.)
                    ]);
                }
                ResourceKind::Server => {
                    let s = res.as_server();
                    let rd = fmt_bytes(s.rate_down as f64) + "/s";
                    let ru = fmt_bytes(s.rate_up as f64) + "/s";
                    table.add_row(row![rd, ru]);
                }
            }
        }
        table.printstd();
    } else {
        println!(
            "{}",
            serde_json::to_string_pretty(&results).chain_err(|| {
                ErrorKind::Serialization
            })?
        );
    }
    Ok(())
}

pub fn pause(mut c: Client, torrents: Vec<&str>) -> Result<()> {
    for torrent in torrents {
        pause_torrent(&mut c, torrent)?;
    }
    Ok(())
}

fn pause_torrent(c: &mut Client, torrent: &str) -> Result<()> {
    let resources = search_torrent_name(c, torrent)?;
    if resources.len() == 1 {
        let msg = CMessage::PauseTorrent {
            serial: c.next_serial(),
            id: resources[0].id().to_owned(),
        };
        c.send(msg)?;
    } else if resources.is_empty() {
        eprintln!("Could not find any matching torrents for {}", torrent);
    } else {
        eprintln!(
            "Ambiguous results searching for {}. Potential alternatives include: ",
            torrent
        );
        for res in resources.into_iter().take(3) {
            if let Resource::Torrent(t) = res {
                eprintln!("{}", t.name);
            }
        }
    }
    Ok(())
}

pub fn resume(mut c: Client, torrents: Vec<&str>) -> Result<()> {
    for torrent in torrents {
        resume_torrent(&mut c, torrent)?;
    }
    Ok(())
}

fn resume_torrent(c: &mut Client, torrent: &str) -> Result<()> {
    let resources = search_torrent_name(c, torrent)?;
    if resources.len() == 1 {
        let msg = CMessage::ResumeTorrent {
            serial: c.next_serial(),
            id: resources[0].id().to_owned(),
        };
        c.send(msg)?;
    } else if resources.is_empty() {
        eprintln!("Could not find any matching torrents for {}", torrent);
    } else {
        eprintln!(
            "Ambiguous results searching for {}. Potential alternatives include: ",
            torrent
        );
        for res in resources.into_iter().take(3) {
            if let Resource::Torrent(t) = res {
                eprintln!("{}", t.name);
            }
        }
    }
    Ok(())
}

fn search_torrent_name(c: &mut Client, name: &str) -> Result<Vec<Resource>> {
    search(
        c,
        ResourceKind::Torrent,
        vec![
            Criterion {
                field: "name".to_owned(),
                op: Operation::ILike,
                value: Value::S(format!("%{}%", name)),
            },
        ],
    )
}

fn search(c: &mut Client, kind: ResourceKind, criteria: Vec<Criterion>) -> Result<Vec<Resource>> {
    let s = c.next_serial();
    let msg = CMessage::FilterSubscribe {
        serial: s,
        kind,
        criteria,
    };
    if let SMessage::OResourcesExtant { ids, .. } = c.rr(msg)? {
        get_resources(c, ids)
    } else {
        bail!("Failed to receive extant resource list!");
    }
}

fn get_resources(c: &mut Client, ids: Vec<String>) -> Result<Vec<Resource>> {
    let msg = CMessage::Subscribe {
        serial: c.next_serial(),
        ids,
    };
    let resources = if let SMessage::UpdateResources { resources } = c.rr(msg)? {
        resources
    } else {
        bail!("Failed to received torrent resource list!");
    };

    let mut results = Vec::new();
    for r in resources {
        if let SResourceUpdate::OResource(res) = r {
            results.push(res);
        } else {
            bail!("Failed to received full resource!");
        }
    }
    Ok(results)
}

fn fmt_bytes(num: f64) -> String {
    let num = num.abs();
    let units = ["B", "kiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB"];
    if num < 1_f64 {
        return format!("{} {}", num, "B");
    }
    let delimiter = 1024_f64;
    let exponent = cmp::min(
        (num.ln() / delimiter.ln()).floor() as i32,
        (units.len() - 1) as i32,
    );
    let pretty_bytes = format!("{:.2}", num / delimiter.powi(exponent))
        .parse::<f64>()
        .unwrap() * 1_f64;
    let unit = units[exponent as usize];
    format!("{} {}", pretty_bytes, unit)
}
<commit_msg>Update sycli for new RPC spec<commit_after>use std::{fs, cmp};
use std::path::Path;
use std::io::{self, Read};

use reqwest::{Client as HClient, header};
use serde_json;
use prettytable::Table;
use url::Url;

use rpc::message::{CMessage, SMessage};
use rpc::criterion::{Criterion, Value, Operation};
use rpc::resource::{Resource, ResourceKind, SResourceUpdate};

use client::Client;
use error::{Result, ResultExt, ErrorKind};

pub fn add(
    mut c: Client,
    url: &str,
    files: Vec<&str>,
    dir: Option<&str>,
    start: bool,
) -> Result<()> {
    for file in files {
        add_file(&mut c, url, file, dir, start)?;
    }
    Ok(())
}

fn add_file(c: &mut Client, url: &str, file: &str, dir: Option<&str>, start: bool) -> Result<()> {
    let mut torrent = Vec::new();
    let mut f = fs::File::open(file).chain_err(|| ErrorKind::FileIO)?;
    f.read_to_end(&mut torrent).chain_err(|| ErrorKind::FileIO)?;

    let msg = CMessage::UploadTorrent {
        serial: c.next_serial(),
        size: torrent.len() as u64,
        path: dir.as_ref().map(|d| format!("{}", d)),
        start,
    };
    let token = if let SMessage::TransferOffer { token, .. } = c.rr(msg)? {
        token
    } else {
        bail!("Failed to receieve transfer offer from synapse!");
    };
    let client = HClient::new().chain_err(|| ErrorKind::HTTP)?;
    client
        .post(url)
        .chain_err(|| ErrorKind::HTTP)?
        .header(header::Authorization(header::Bearer { token }))
        .body(torrent)
        .send()
        .chain_err(|| ErrorKind::HTTP)?;

    if let SMessage::OResourcesExtant { .. } = c.recv()? {
    } else {
        bail!("Failed to receieve upload acknowledgement from synapse!");
    };

    Ok(())
}

pub fn del(mut c: Client, torrents: Vec<&str>) -> Result<()> {
    for torrent in torrents {
        del_torrent(&mut c, torrent)?;
    }
    Ok(())
}

fn del_torrent(c: &mut Client, torrent: &str) -> Result<()> {
    let resources = search_torrent_name(c, torrent)?;
    if resources.len() == 1 {
        let msg = CMessage::RemoveResource {
            serial: c.next_serial(),
            id: resources[0].id().to_owned(),
        };
        c.send(msg)?;
    } else if resources.is_empty() {
        eprintln!("Could not find any matching torrents for {}", torrent);
    } else {
        eprintln!(
            "Ambiguous results searching for {}. Potential alternatives include: ",
            torrent
        );
        for res in resources.into_iter().take(3) {
            if let Resource::Torrent(t) = res {
                eprintln!("{}", t.name.unwrap_or("[Unknown Magnet]".to_owned()));
            }
        }
    }
    Ok(())
}

pub fn dl(mut c: Client, url: &str, name: &str) -> Result<()> {
    let resources = search_torrent_name(&mut c, name)?;
    let files = if resources.len() == 1 {
        let msg = CMessage::FilterSubscribe {
            serial: c.next_serial(),
            kind: ResourceKind::File,
            criteria: vec![
                Criterion {
                    field: "torrent_id".to_owned(),
                    op: Operation::Eq,
                    value: Value::S(resources[0].id().to_owned()),
                },
            ],
        };
        if let SMessage::OResourcesExtant { ids, .. } = c.rr(msg)? {
            get_resources(&mut c, ids)?
        } else {
            bail!("Could not get files for torrent!");
        }
    } else if resources.is_empty() {
        eprintln!("Could not find any matching torrents for {}", name);
        return Ok(());
    } else {
        eprintln!(
            "Ambiguous results searching for {}. Potential alternatives include: ",
            name
        );
        for res in resources.into_iter().take(3) {
            if let Resource::Torrent(t) = res {
                eprintln!("{}", t.name.unwrap_or("[Unknown Magnet]".to_owned()));
            }
        }
        return Ok(());
    };

    for file in files {
        let mut dl_url = Url::parse(url).unwrap();
        dl_url.path_segments_mut().unwrap().push("dl").push(
            file.id(),
        );

        let client = HClient::new().chain_err(|| ErrorKind::HTTP)?;
        let mut resp = client
            .get(dl_url.as_str())
            .chain_err(|| ErrorKind::HTTP)?
            .send()
            .chain_err(|| ErrorKind::HTTP)?;
        if let Resource::File(f) = file {
            let p = Path::new(&f.path);
            if let Some(par) = p.parent() {
                fs::create_dir_all(par).chain_err(|| ErrorKind::FileIO)?;
            }
            let mut f = fs::File::create(p).chain_err(|| ErrorKind::FileIO)?;
            io::copy(&mut resp, &mut f).chain_err(|| ErrorKind::FileIO)?;
        } else {
            bail!("Expected a file resource");
        }
    }
    Ok(())
}

pub fn list(mut c: Client, kind: &str, crit: Vec<Criterion>, output: &str) -> Result<()> {
    let k = match kind {
        "torrent" => ResourceKind::Torrent,
        "tracker" => ResourceKind::Tracker,
        "peer" => ResourceKind::Peer,
        "piece" => ResourceKind::Piece,
        "file" => ResourceKind::File,
        "server" => ResourceKind::Server,
        _ => bail!("Unexpected resource kind {}", kind),
    };
    let results = search(&mut c, k, crit)?;
    if output == "text" {
        let mut table = Table::new();
        match k {
            ResourceKind::Torrent => {
                table.add_row(row!["Name", "Done", "DL", "UL", "DL RT", "UL RT", "Peers"]);
            }
            ResourceKind::Tracker => {
                table.add_row(row!["URL", "Torrent", "Error"]);
            }
            ResourceKind::Peer => {
                table.add_row(row!["IP", "Torrent", "DL RT", "UL RT"]);
            }
            ResourceKind::Piece => {
                table.add_row(row!["Torrent", "DLd", "Avail"]);
            }
            ResourceKind::File => {
                table.add_row(row!["Path", "Torrent", "Done", "Prio", "Avail"]);
            }
            ResourceKind::Server => {
                table.add_row(row!["DL RT", "UL RT"]);
            }
        }

        #[cfg_attr(rustfmt, rustfmt_skip)]
        for res in results {
            match k {
                ResourceKind::Torrent => {
                    let t = res.as_torrent();
                    table.add_row(row![
                        t.name.as_ref().map(|s| s.as_str()).unwrap_or("[Unknown Magnet]"),
                        format!("{:.2}%", t.progress * 100.),
                        fmt_bytes(t.transferred_down as f64),
                        fmt_bytes(t.transferred_up as f64),
                        fmt_bytes(t.rate_down as f64) + "/s",
                        fmt_bytes(t.rate_up as f64) + "/s",
                        t.peers
                    ]);
                }
                ResourceKind::Tracker => {
                    let t = res.as_tracker();
                    table.add_row(row![
                        t.url,
                        t.torrent_id,
                        t.error.as_ref().map(|s| s.as_str()).unwrap_or("")
                    ]);
                }
                ResourceKind::Peer => {
                    let p = res.as_peer();
                    let rd = fmt_bytes(p.rate_down as f64) + "/s";
                    let ru = fmt_bytes(p.rate_up as f64) + "/s";
                    table.add_row(row![p.ip, p.torrent_id, rd, ru]);
                }
                ResourceKind::Piece => {
                    let p = res.as_piece();
                    table.add_row(row![p.torrent_id, p.downloaded, p.available]);
                }
                ResourceKind::File => {
                    let f = res.as_file();
                    table.add_row(row![
                        f.path,
                        f.torrent_id,
                        format!("{:.2}%", f.progress as f64 * 100.),
                        f.priority,
                        format!("{:.2}%", f.availability as f64 * 100.)
                    ]);
                }
                ResourceKind::Server => {
                    let s = res.as_server();
                    let rd = fmt_bytes(s.rate_down as f64) + "/s";
                    let ru = fmt_bytes(s.rate_up as f64) + "/s";
                    table.add_row(row![rd, ru]);
                }
            }
        }
        table.printstd();
    } else {
        println!(
            "{}",
            serde_json::to_string_pretty(&results).chain_err(|| {
                ErrorKind::Serialization
            })?
        );
    }
    Ok(())
}

pub fn pause(mut c: Client, torrents: Vec<&str>) -> Result<()> {
    for torrent in torrents {
        pause_torrent(&mut c, torrent)?;
    }
    Ok(())
}

fn pause_torrent(c: &mut Client, torrent: &str) -> Result<()> {
    let resources = search_torrent_name(c, torrent)?;
    if resources.len() == 1 {
        let msg = CMessage::PauseTorrent {
            serial: c.next_serial(),
            id: resources[0].id().to_owned(),
        };
        c.send(msg)?;
    } else if resources.is_empty() {
        eprintln!("Could not find any matching torrents for {}", torrent);
    } else {
        eprintln!(
            "Ambiguous results searching for {}. Potential alternatives include: ",
            torrent
        );
        for res in resources.into_iter().take(3) {
            if let Resource::Torrent(t) = res {
                eprintln!("{}", t.name.unwrap_or("[Unknown Magnet]".to_owned()));
            }
        }
    }
    Ok(())
}

pub fn resume(mut c: Client, torrents: Vec<&str>) -> Result<()> {
    for torrent in torrents {
        resume_torrent(&mut c, torrent)?;
    }
    Ok(())
}

fn resume_torrent(c: &mut Client, torrent: &str) -> Result<()> {
    let resources = search_torrent_name(c, torrent)?;
    if resources.len() == 1 {
        let msg = CMessage::ResumeTorrent {
            serial: c.next_serial(),
            id: resources[0].id().to_owned(),
        };
        c.send(msg)?;
    } else if resources.is_empty() {
        eprintln!("Could not find any matching torrents for {}", torrent);
    } else {
        eprintln!(
            "Ambiguous results searching for {}. Potential alternatives include: ",
            torrent
        );
        for res in resources.into_iter().take(3) {
            if let Resource::Torrent(t) = res {
                eprintln!("{}", t.name.unwrap_or("[Unknown Magnet]".to_owned()));
            }
        }
    }
    Ok(())
}

fn search_torrent_name(c: &mut Client, name: &str) -> Result<Vec<Resource>> {
    search(
        c,
        ResourceKind::Torrent,
        vec![
            Criterion {
                field: "name".to_owned(),
                op: Operation::ILike,
                value: Value::S(format!("%{}%", name)),
            },
        ],
    )
}

fn search(c: &mut Client, kind: ResourceKind, criteria: Vec<Criterion>) -> Result<Vec<Resource>> {
    let s = c.next_serial();
    let msg = CMessage::FilterSubscribe {
        serial: s,
        kind,
        criteria,
    };
    if let SMessage::OResourcesExtant { ids, .. } = c.rr(msg)? {
        get_resources(c, ids)
    } else {
        bail!("Failed to receive extant resource list!");
    }
}

fn get_resources(c: &mut Client, ids: Vec<String>) -> Result<Vec<Resource>> {
    let msg = CMessage::Subscribe {
        serial: c.next_serial(),
        ids,
    };
    let resources = if let SMessage::UpdateResources { resources } = c.rr(msg)? {
        resources
    } else {
        bail!("Failed to received torrent resource list!");
    };

    let mut results = Vec::new();
    for r in resources {
        if let SResourceUpdate::OResource(res) = r {
            results.push(res);
        } else {
            bail!("Failed to received full resource!");
        }
    }
    Ok(results)
}

fn fmt_bytes(num: f64) -> String {
    let num = num.abs();
    let units = ["B", "kiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB"];
    if num < 1_f64 {
        return format!("{} {}", num, "B");
    }
    let delimiter = 1024_f64;
    let exponent = cmp::min(
        (num.ln() / delimiter.ln()).floor() as i32,
        (units.len() - 1) as i32,
    );
    let pretty_bytes = format!("{:.2}", num / delimiter.powi(exponent))
        .parse::<f64>()
        .unwrap() * 1_f64;
    let unit = units[exponent as usize];
    format!("{} {}", pretty_bytes, unit)
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::collections::HashMap;
use std::fmt;

use super::Format;
use super::english_cardinal_classifier;
use {Args, Message};

/// The set of [grammatical numbers] that we support.
///
/// [grammatical numbers]: https://en.wikipedia.org/wiki/Grammatical_number
#[derive(Clone,Copy,Debug,PartialEq)]
pub enum PluralCategory {
    /// Value is `0`.
    Zero,
    /// Value is `1`. In English, this corresponds to the "singular" form.
    One,
    /// Value is `2`.
    Two,
    /// Value is a few, more than `2`, but less than `many`. The exact
    /// range depends upon the locale.
    Few,
    /// Value is many, more than `few`. The exact range depends
    /// upon the locale.
    Many,
    /// Not one of the others. In English, this is used for the "plural"
    /// form.
    Other,
}

/// Format a value taking pluralization rules into account.
pub struct PluralFormat {
    /// The name of the variable whose value should be formatted.
    #[allow(dead_code)]
    variable_name: String,
    classifier: fn(i64) -> PluralCategory,
    literals: HashMap<i64, Message>,
    offset: i64,
    zero: Option<Message>,
    one: Option<Message>,
    two: Option<Message>,
    few: Option<Message>,
    many: Option<Message>,
    other: Message,
}

impl PluralFormat {
    /// Construct a `PluralFormat`.
    pub fn new(variable_name: &str, other: Message) -> Self {
        PluralFormat {
            variable_name: variable_name.to_string(),
            classifier: english_cardinal_classifier,
            literals: HashMap::new(),
            offset: 0,
            zero: None,
            one: None,
            two: None,
            few: None,
            many: None,
            other: other,
        }
    }

    /// Set the `message` to be used for a literal value.
    pub fn literal(mut self, literal: i64, message: Message) -> Self {
        self.literals.insert(literal, message);
        self
    }

    /// Apply an `offset`.
    pub fn offset(mut self, offset: i64) -> Self {
        self.offset = offset;
        self
    }

    /// Set the `message` for `PluralCategory::Zero`.
    pub fn zero(mut self, message: Message) -> Self {
        self.zero = Some(message);
        self
    }

    /// Set the `message` for `PluralCategory::One`.
    pub fn one(mut self, message: Message) -> Self {
        self.one = Some(message);
        self
    }

    /// Set the `message` for `PluralCategory::Two`.
    pub fn two(mut self, message: Message) -> Self {
        self.two = Some(message);
        self
    }

    /// Set the `message` for `PluralCategory::Few`.
    pub fn few(mut self, message: Message) -> Self {
        self.few = Some(message);
        self
    }

    /// Set the `message` for `PluralCategory::Many`.
    pub fn many(mut self, message: Message) -> Self {
        self.many = Some(message);
        self
    }

    /// Given a value adjusted by the `offset`, determine which `Message` to use.
    fn format_from_classifier(&self, offset_value: i64) -> &Message {
        let category = (self.classifier)(offset_value);
        match category {
            PluralCategory::Zero => self.zero.as_ref().unwrap_or(&self.other),
            PluralCategory::One => self.one.as_ref().unwrap_or(&self.other),
            PluralCategory::Two => self.two.as_ref().unwrap_or(&self.other),
            PluralCategory::Few => self.few.as_ref().unwrap_or(&self.other),
            PluralCategory::Many => self.many.as_ref().unwrap_or(&self.other),
            PluralCategory::Other => &self.other,
        }
    }

    /// Handle specialized behavior for `PlaceholderFormat` when formatting
    /// a `PluralFormat`.
    ///
    /// XXX: Can we fold this away and do something more general for
    /// `PlaceholderFormat`?
    fn format_plural_message(&self,
                             stream: &mut fmt::Write,
                             message: &Message,
                             _offset_value: i64,
                             args: &Args)
                             -> fmt::Result {
        for part in &message.parts {
            // XXX: Need to deal with PlaceholderFormat here and give it
            // the `offset_value`.
            try!(part.apply_format(stream, args));
        }
        Ok(())
    }
}

impl Format for PluralFormat {
    fn apply_format(&self, stream: &mut fmt::Write, args: &Args) -> fmt::Result {
        let value = 0;
        let offset_value = value - self.offset;
        let message = if !self.literals.is_empty() {
            if let Some(literal) = self.literals.get(&offset_value) {
                literal
            } else {
                self.format_from_classifier(offset_value)
            }
        } else {
            self.format_from_classifier(offset_value)
        };
        try!(self.format_plural_message(stream, message, offset_value, args));
        Ok(())
    }
}
<commit_msg>PluralFormat: Move all lookup into a single function.<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::collections::HashMap;
use std::fmt;

use super::Format;
use super::english_cardinal_classifier;
use {Args, Message};

/// The set of [grammatical numbers] that we support.
///
/// [grammatical numbers]: https://en.wikipedia.org/wiki/Grammatical_number
#[derive(Clone,Copy,Debug,PartialEq)]
pub enum PluralCategory {
    /// Value is `0`.
    Zero,
    /// Value is `1`. In English, this corresponds to the "singular" form.
    One,
    /// Value is `2`.
    Two,
    /// Value is a few, more than `2`, but less than `many`. The exact
    /// range depends upon the locale.
    Few,
    /// Value is many, more than `few`. The exact range depends
    /// upon the locale.
    Many,
    /// Not one of the others. In English, this is used for the "plural"
    /// form.
    Other,
}

/// Format a value taking pluralization rules into account.
pub struct PluralFormat {
    /// The name of the variable whose value should be formatted.
    #[allow(dead_code)]
    variable_name: String,
    classifier: fn(i64) -> PluralCategory,
    literals: HashMap<i64, Message>,
    offset: i64,
    zero: Option<Message>,
    one: Option<Message>,
    two: Option<Message>,
    few: Option<Message>,
    many: Option<Message>,
    other: Message,
}

impl PluralFormat {
    /// Construct a `PluralFormat`.
    pub fn new(variable_name: &str, other: Message) -> Self {
        PluralFormat {
            variable_name: variable_name.to_string(),
            classifier: english_cardinal_classifier,
            literals: HashMap::new(),
            offset: 0,
            zero: None,
            one: None,
            two: None,
            few: None,
            many: None,
            other: other,
        }
    }

    /// Set the `message` to be used for a literal value.
    pub fn literal(mut self, literal: i64, message: Message) -> Self {
        self.literals.insert(literal, message);
        self
    }

    /// Apply an `offset`.
    pub fn offset(mut self, offset: i64) -> Self {
        self.offset = offset;
        self
    }

    /// Set the `message` for `PluralCategory::Zero`.
    pub fn zero(mut self, message: Message) -> Self {
        self.zero = Some(message);
        self
    }

    /// Set the `message` for `PluralCategory::One`.
    pub fn one(mut self, message: Message) -> Self {
        self.one = Some(message);
        self
    }

    /// Set the `message` for `PluralCategory::Two`.
    pub fn two(mut self, message: Message) -> Self {
        self.two = Some(message);
        self
    }

    /// Set the `message` for `PluralCategory::Few`.
    pub fn few(mut self, message: Message) -> Self {
        self.few = Some(message);
        self
    }

    /// Set the `message` for `PluralCategory::Many`.
    pub fn many(mut self, message: Message) -> Self {
        self.many = Some(message);
        self
    }

    /// Given a value adjusted by the `offset`, determine which `Message` to use.
    fn lookup_message(&self, offset_value: i64) -> &Message {
        if let Some(literal) = self.literals.get(&offset_value) {
            literal
        } else {
            let category = (self.classifier)(offset_value);
            match category {
                PluralCategory::Zero => self.zero.as_ref().unwrap_or(&self.other),
                PluralCategory::One => self.one.as_ref().unwrap_or(&self.other),
                PluralCategory::Two => self.two.as_ref().unwrap_or(&self.other),
                PluralCategory::Few => self.few.as_ref().unwrap_or(&self.other),
                PluralCategory::Many => self.many.as_ref().unwrap_or(&self.other),
                PluralCategory::Other => &self.other,
            }
        }
    }

    /// Handle specialized behavior for `PlaceholderFormat` when formatting
    /// a `PluralFormat`.
    ///
    /// XXX: Can we fold this away and do something more general for
    /// `PlaceholderFormat`?
    fn format_plural_message(&self,
                             stream: &mut fmt::Write,
                             message: &Message,
                             _offset_value: i64,
                             args: &Args)
                             -> fmt::Result {
        for part in &message.parts {
            // XXX: Need to deal with PlaceholderFormat here and give it
            // the `offset_value`.
            try!(part.apply_format(stream, args));
        }
        Ok(())
    }
}

impl Format for PluralFormat {
    fn apply_format(&self, stream: &mut fmt::Write, args: &Args) -> fmt::Result {
        let value = 0;
        let offset_value = value - self.offset;
        let message = self.lookup_message(offset_value);
        try!(self.format_plural_message(stream, message, offset_value, args));
        Ok(())
    }
}
<|endoftext|>


<commit_before>pub mod segment;
pub mod agent;
pub mod gen;

use backend::obj;
use backend::obj::*;
use std::collections::HashMap;
use std::collections::HashSet;
use std::f32::consts;
use core::color;
use rand;
use core::color::ToRgb;
use core::geometry::*;
use backend::world::segment::*;
use backend::world::agent::Agent;
use backend::world::gen::*;
use num::FromPrimitive;

enum_from_primitive! {
	#[derive(Debug, Eq, PartialEq, Hash, Copy, Clone)]
	pub enum AgentType {
		Minion ,
		Spore,
		Player,
		FriendlyBullet,
		Enemy,
		EnemyBullet,
		Resource,
		Prop,
	}
}

// TODO: is there a better way to derive this?
const AGENT_TYPES: &'static [AgentType] = &[AgentType::Minion,
                                            AgentType::Spore,
                                            AgentType::Player,
                                            AgentType::FriendlyBullet,
                                            AgentType::Enemy,
                                            AgentType::EnemyBullet,
                                            AgentType::Resource,
                                            AgentType::Prop];

impl AgentType {
	pub fn all() -> &'static [AgentType] {
		AGENT_TYPES
	}
}

pub struct Swarm {
	seq: Id,
	agent_type: AgentType,
	rng: rand::ThreadRng,
	gen: Genome,
	agents: AgentMap,
}

impl Swarm {
	pub fn new(agent_type: AgentType) -> Swarm {
		Swarm {
			seq: 0,
			agent_type: agent_type,
			rng: rand::thread_rng(),
			gen: Genome::new(b"Rust-Oids are cool!"),
			agents: HashMap::new(),
		}
	}

	#[allow(dead_code)]
	pub fn type_of(&self) -> AgentType {
		self.agent_type
	}

	pub fn get(&self, id: Id) -> Option<&Agent> {
		self.agents.get(&id)
	}

	pub fn get_mut(&mut self, id: Id) -> Option<&mut agent::Agent> {
		self.agents.get_mut(&id)
	}

	pub fn mutate(&mut self) {
		self.gen = self.gen.mutate(&mut self.rng);
	}

	pub fn next_id(&mut self) -> Id {
		self.seq = self.seq + 1;
		self.seq << 8 | (self.agent_type as usize)
	}

	pub fn new_resource(&mut self, initial_pos: Position, charge: f32) -> Id {
		let albedo = color::YPbPr::new(0.5,
		                               self.gen.next_float(-0.5, 0.5),
		                               self.gen.next_float(-0.5, 0.5));
		let ball = self.gen.ball();
		let mut builder = agent::AgentBuilder::new(self.next_id(),
		                                           Material { density: 1.0, ..Default::default() },
		                                           Livery { albedo: albedo.to_rgba(), ..Default::default() },
		                                           self.gen.dna(),
		                                           segment::State::with_charge(charge, 0., charge));
		self.insert(builder.start(initial_pos, 0., &ball).build())
	}

	pub fn free_resources(&mut self, freed: &mut Vec<Agent>) {
		let mut dead = HashSet::new();

		for id in self.agents
			.iter()
			.filter(|&(_, agent)| !agent.state.is_alive())
			.map(|(&id, _)| id) {
			dead.insert(id);
		}
		for id in &dead {
			if let Some(agent) = self.agents.remove(&id) {
				freed.push(agent);
			}
		}
	}

	pub fn new_minion(&mut self, initial_pos: Position, charge: f32) -> Id {
		let albedo = color::Hsl::new(self.gen.next_float(0., 1.), 0.5, 0.5);
		let mut builder = agent::AgentBuilder::new(self.next_id(),
		                                           Material { density: 0.2, ..Default::default() },
		                                           Livery { albedo: albedo.to_rgba(), ..Default::default() },
		                                           self.gen.dna(),
		                                           segment::State::with_charge(0., charge, charge));
		let arm_shape = self.gen.star();
		let leg_shape = self.gen.star();
		let torso_shape = self.gen.npoly(5, true);
		let head_shape = self.gen.iso_triangle();
		let tail_shape = self.gen.vbar();
		let initial_angle = consts::PI / 2. + f32::atan2(initial_pos.y, initial_pos.x);

		let torso = builder.start(initial_pos, initial_angle, &torso_shape)
			.index();
		builder.addr(torso, 2, &arm_shape, ARM | JOINT | ACTUATOR | RUDDER)
			.addl(torso, -2, &arm_shape, ARM | JOINT | ACTUATOR | RUDDER);

		let head = builder.add(torso, 0, &head_shape, HEAD | SENSOR).index();
		builder.addr(head, 1, &head_shape, HEAD | ACTUATOR | RUDDER)
			.addl(head, 2, &head_shape, HEAD | ACTUATOR | RUDDER);

		let mut belly = torso;
		let mut belly_mid = torso_shape.mid();
		for _ in 0..self.gen.next_integer(0, 4) {
			let belly_shape = self.gen.poly(true);

			belly = builder.add(belly, belly_mid, &belly_shape, BELLY | JOINT).index();
			belly_mid = belly_shape.mid();
			if self.gen.next_integer(0, 4) == 0 {
				builder.addr(belly, 2, &arm_shape, ARM | ACTUATOR | RUDDER)
					.addl(belly, -2, &arm_shape, ARM | ACTUATOR | RUDDER);
			}
		}

		builder.addr(belly, belly_mid - 1, &leg_shape, LEG | ACTUATOR | THRUSTER)
			.addl(belly,
			      -(belly_mid - 1),
			      &leg_shape,
			      LEG | ACTUATOR | THRUSTER)
			.add(belly, belly_mid, &tail_shape, TAIL | ACTUATOR | BRAKE);
		self.mutate();
		self.insert(builder.build())
	}

	fn insert(&mut self, agent: Agent) -> Id {
		let id = agent.id();
		self.agents.insert(id, agent);
		id
	}

	pub fn kill(&mut self, id: &Id) {
		if let Some(ref mut agent) = self.agents.get_mut(id) {
			agent.state.die();
		}
	}

	pub fn agents(&self) -> &HashMap<Id, Agent> {
		&self.agents
	}

	pub fn agents_mut(&mut self) -> &mut HashMap<Id, Agent> {
		&mut self.agents
	}
}

pub type AgentMap = HashMap<Id, agent::Agent>;
pub type SwarmMap = HashMap<AgentType, Swarm>;

#[repr(packed)]
#[derive(Eq, Hash, PartialEq, Clone, Copy, Debug)]
pub struct AgentRefs {
	pub agent_id: obj::Id,
	pub segment_index: obj::SegmentIndex,
	pub bone_index: obj::BoneIndex,
}

impl Default for AgentRefs {
	fn default() -> AgentRefs {
		AgentRefs {
			agent_id: 0xdeadbeef,
			segment_index: 0,
			bone_index: 0,
		}
	}
}

impl AgentRefs {
	pub fn with_id(id: obj::Id) -> AgentRefs {
		AgentRefs { agent_id: id, ..Default::default() }
	}

	pub fn with_segment(id: obj::Id, segment_index: obj::SegmentIndex) -> AgentRefs {
		AgentRefs {
			agent_id: id,
			segment_index: segment_index,
			..Default::default()
		}
	}

	pub fn with_bone(id: obj::Id, segment_index: obj::SegmentIndex, bone_index: obj::BoneIndex) -> AgentRefs {
		AgentRefs {
			agent_id: id,
			segment_index: segment_index,
			bone_index: bone_index,
		}
	}

	pub fn no_bone(&self) -> AgentRefs {
		AgentRefs { bone_index: 0, ..*self }
	}
}

pub trait TypedAgent {
	fn type_of(&self) -> AgentType;
}

impl TypedAgent for Id {
	fn type_of(&self) -> AgentType {
		AgentType::from_usize(*self & 0xff).unwrap_or(AgentType::Prop)
	}
}

pub struct World {
	pub extent: Rect,
	pub fence: obj::Mesh,
	agents: HashMap<AgentType, Swarm>,
}

pub trait WorldState {
	fn agent(&self, id: obj::Id) -> Option<&Agent>;
}

impl WorldState for World {
	fn agent(&self, id: obj::Id) -> Option<&Agent> {
		self.agents.get(&id.type_of()).and_then(|m| m.get(id))
	}
}

pub struct Cleanup {
	pub freed: Box<[Agent]>,
}

impl World {
	pub fn new() -> Self {
		let mut agents = HashMap::new();
		let types = AgentType::all();
		for t in types {
			agents.insert(*t, Swarm::new(*t));
		}
		World {
			extent: Rect::new(-50., -50., 50., 50.),
			fence: Mesh::from_shape(Shape::new_ball(50.), Winding::CW),
			agents: agents,
		}
	}

	pub fn new_resource(&mut self, pos: Position) -> obj::Id {
		self.agents.get_mut(&AgentType::Resource).unwrap().new_resource(pos, 0.8)
	}

	pub fn new_minion(&mut self, pos: Position) -> obj::Id {
		self.agents.get_mut(&AgentType::Minion).unwrap().new_minion(pos, 0.3)
	}

	pub fn agent_mut(&mut self, id: obj::Id) -> Option<&mut Agent> {
		self.agents.get_mut(&id.type_of()).and_then(|m| m.get_mut(id))
	}

	pub fn agents_mut(&mut self, agent_type: AgentType) -> &mut AgentMap {
		&mut self.agents.get_mut(&agent_type).unwrap().agents
	}

	pub fn swarms(&self) -> &SwarmMap {
		&self.agents
	}

	pub fn swarms_mut(&mut self) -> &mut SwarmMap {
		&mut self.agents
	}

	pub fn cleanup(&mut self) -> Cleanup {
		let mut v = Vec::new();
		for (_, agents) in self.agents.iter_mut() {
			agents.free_resources(&mut v);
		}
		Cleanup { freed: v.into_boxed_slice() }
	}
}
<commit_msg>Moved Rng out of the Swarm<commit_after>pub mod segment;
pub mod agent;
pub mod gen;

use backend::obj;
use backend::obj::*;
use std::collections::HashMap;
use std::collections::HashSet;
use std::f32::consts;
use core::color;
use rand;
use core::color::ToRgb;
use core::geometry::*;
use backend::world::segment::*;
use backend::world::agent::Agent;
use backend::world::gen::*;
use num::FromPrimitive;

enum_from_primitive! {
	#[derive(Debug, Eq, PartialEq, Hash, Copy, Clone)]
	pub enum AgentType {
		Minion ,
		Spore,
		Player,
		FriendlyBullet,
		Enemy,
		EnemyBullet,
		Resource,
		Prop,
	}
}

// TODO: is there a better way to derive this?
const AGENT_TYPES: &'static [AgentType] = &[AgentType::Minion,
                                            AgentType::Spore,
                                            AgentType::Player,
                                            AgentType::FriendlyBullet,
                                            AgentType::Enemy,
                                            AgentType::EnemyBullet,
                                            AgentType::Resource,
                                            AgentType::Prop];

impl AgentType {
	pub fn all() -> &'static [AgentType] {
		AGENT_TYPES
	}
}

pub struct Swarm {
	seq: Id,
	agent_type: AgentType,
	gen: Genome,
	agents: AgentMap,
}

impl Swarm {
	pub fn new(agent_type: AgentType) -> Swarm {
		Swarm {
			seq: 0,
			agent_type: agent_type,
			gen: Genome::new(b"Rust-Oids are cool!"),
			agents: HashMap::new(),
		}
	}

	#[allow(dead_code)]
	pub fn type_of(&self) -> AgentType {
		self.agent_type
	}

	pub fn get(&self, id: Id) -> Option<&Agent> {
		self.agents.get(&id)
	}

	pub fn get_mut(&mut self, id: Id) -> Option<&mut agent::Agent> {
		self.agents.get_mut(&id)
	}

	pub fn mutate<R: rand::Rng>(&mut self, rng: &mut R) {
		self.gen = self.gen.mutate(rng);
	}

	pub fn next_id(&mut self) -> Id {
		self.seq = self.seq + 1;
		self.seq << 8 | (self.agent_type as usize)
	}

	pub fn new_resource(&mut self, initial_pos: Position, charge: f32) -> Id {
		let albedo = color::YPbPr::new(0.5,
		                               self.gen.next_float(-0.5, 0.5),
		                               self.gen.next_float(-0.5, 0.5));
		let ball = self.gen.ball();
		let mut builder = agent::AgentBuilder::new(self.next_id(),
		                                           Material { density: 1.0, ..Default::default() },
		                                           Livery { albedo: albedo.to_rgba(), ..Default::default() },
		                                           self.gen.dna(),
		                                           segment::State::with_charge(charge, 0., charge));
		self.insert(builder.start(initial_pos, 0., &ball).build())
	}

	pub fn free_resources(&mut self, freed: &mut Vec<Agent>) {
		let mut dead = HashSet::new();

		for id in self.agents
		              .iter()
		              .filter(|&(_, agent)| !agent.state.is_alive())
		              .map(|(&id, _)| id) {
			dead.insert(id);
		}
		for id in &dead {
			if let Some(agent) = self.agents.remove(&id) {
				freed.push(agent);
			}
		}
	}

	pub fn new_minion(&mut self, initial_pos: Position, charge: f32) -> Id {
		let albedo = color::Hsl::new(self.gen.next_float(0., 1.), 0.5, 0.5);
		let mut builder = agent::AgentBuilder::new(self.next_id(),
		                                           Material { density: 0.2, ..Default::default() },
		                                           Livery { albedo: albedo.to_rgba(), ..Default::default() },
		                                           self.gen.dna(),
		                                           segment::State::with_charge(0., charge, charge));
		let arm_shape = self.gen.star();
		let leg_shape = self.gen.star();
		let torso_shape = self.gen.npoly(5, true);
		let head_shape = self.gen.iso_triangle();
		let tail_shape = self.gen.vbar();
		let initial_angle = consts::PI / 2. + f32::atan2(initial_pos.y, initial_pos.x);

		let torso = builder.start(initial_pos, initial_angle, &torso_shape)
		                   .index();
		builder.addr(torso, 2, &arm_shape, ARM | JOINT | ACTUATOR | RUDDER)
		       .addl(torso, -2, &arm_shape, ARM | JOINT | ACTUATOR | RUDDER);

		let head = builder.add(torso, 0, &head_shape, HEAD | SENSOR).index();
		builder.addr(head, 1, &head_shape, HEAD | ACTUATOR | RUDDER)
		       .addl(head, 2, &head_shape, HEAD | ACTUATOR | RUDDER);

		let mut belly = torso;
		let mut belly_mid = torso_shape.mid();
		for _ in 0..self.gen.next_integer(0, 4) {
			let belly_shape = self.gen.poly(true);

			belly = builder.add(belly, belly_mid, &belly_shape, BELLY | JOINT).index();
			belly_mid = belly_shape.mid();
			if self.gen.next_integer(0, 4) == 0 {
				builder.addr(belly, 2, &arm_shape, ARM | ACTUATOR | RUDDER)
				       .addl(belly, -2, &arm_shape, ARM | ACTUATOR | RUDDER);
			}
		}

		builder.addr(belly, belly_mid - 1, &leg_shape, LEG | ACTUATOR | THRUSTER)
		       .addl(belly,
		             -(belly_mid - 1),
		             &leg_shape,
		             LEG | ACTUATOR | THRUSTER)
		       .add(belly, belly_mid, &tail_shape, TAIL | ACTUATOR | BRAKE);
		self.mutate(&mut rand::thread_rng());
		self.insert(builder.build())
	}

	fn insert(&mut self, agent: Agent) -> Id {
		let id = agent.id();
		self.agents.insert(id, agent);
		id
	}

	pub fn kill(&mut self, id: &Id) {
		if let Some(ref mut agent) = self.agents.get_mut(id) {
			agent.state.die();
		}
	}

	pub fn agents(&self) -> &HashMap<Id, Agent> {
		&self.agents
	}

	pub fn agents_mut(&mut self) -> &mut HashMap<Id, Agent> {
		&mut self.agents
	}
}

pub type AgentMap = HashMap<Id, agent::Agent>;
pub type SwarmMap = HashMap<AgentType, Swarm>;

#[repr(packed)]
#[derive(Eq, Hash, PartialEq, Clone, Copy, Debug)]
pub struct AgentRefs {
	pub agent_id: obj::Id,
	pub segment_index: obj::SegmentIndex,
	pub bone_index: obj::BoneIndex,
}

impl Default for AgentRefs {
	fn default() -> AgentRefs {
		AgentRefs {
			agent_id: 0xdeadbeef,
			segment_index: 0,
			bone_index: 0,
		}
	}
}

impl AgentRefs {
	pub fn with_id(id: obj::Id) -> AgentRefs {
		AgentRefs { agent_id: id, ..Default::default() }
	}

	pub fn with_segment(id: obj::Id, segment_index: obj::SegmentIndex) -> AgentRefs {
		AgentRefs {
			agent_id: id,
			segment_index: segment_index,
			..Default::default()
		}
	}

	pub fn with_bone(id: obj::Id, segment_index: obj::SegmentIndex, bone_index: obj::BoneIndex) -> AgentRefs {
		AgentRefs {
			agent_id: id,
			segment_index: segment_index,
			bone_index: bone_index,
		}
	}

	pub fn no_bone(&self) -> AgentRefs {
		AgentRefs { bone_index: 0, ..*self }
	}
}

pub trait TypedAgent {
	fn type_of(&self) -> AgentType;
}

impl TypedAgent for Id {
	fn type_of(&self) -> AgentType {
		AgentType::from_usize(*self & 0xff).unwrap_or(AgentType::Prop)
	}
}

pub struct World {
	pub extent: Rect,
	pub fence: obj::Mesh,
	agents: HashMap<AgentType, Swarm>,
}

pub trait WorldState {
	fn agent(&self, id: obj::Id) -> Option<&Agent>;
}

impl WorldState for World {
	fn agent(&self, id: obj::Id) -> Option<&Agent> {
		self.agents.get(&id.type_of()).and_then(|m| m.get(id))
	}
}

pub struct Cleanup {
	pub freed: Box<[Agent]>,
}

impl World {
	pub fn new() -> Self {
		let mut agents = HashMap::new();
		let types = AgentType::all();
		for t in types {
			agents.insert(*t, Swarm::new(*t));
		}
		World {
			extent: Rect::new(-50., -50., 50., 50.),
			fence: Mesh::from_shape(Shape::new_ball(50.), Winding::CW),
			agents: agents,
		}
	}

	pub fn new_resource(&mut self, pos: Position) -> obj::Id {
		self.agents.get_mut(&AgentType::Resource).unwrap().new_resource(pos, 0.8)
	}

	pub fn new_minion(&mut self, pos: Position) -> obj::Id {
		self.agents.get_mut(&AgentType::Minion).unwrap().new_minion(pos, 0.3)
	}

	pub fn agent_mut(&mut self, id: obj::Id) -> Option<&mut Agent> {
		self.agents.get_mut(&id.type_of()).and_then(|m| m.get_mut(id))
	}

	pub fn agents_mut(&mut self, agent_type: AgentType) -> &mut AgentMap {
		&mut self.agents.get_mut(&agent_type).unwrap().agents
	}

	pub fn swarms(&self) -> &SwarmMap {
		&self.agents
	}

	pub fn swarms_mut(&mut self) -> &mut SwarmMap {
		&mut self.agents
	}

	pub fn cleanup(&mut self) -> Cleanup {
		let mut v = Vec::new();
		for (_, agents) in self.agents.iter_mut() {
			agents.free_resources(&mut v);
		}
		Cleanup { freed: v.into_boxed_slice() }
	}
}
<|endoftext|>