<commit_before>//! Group records by a key, and apply a reduction function.
//!
//! The `group` operators act on data that can be viewed as pairs `(key, val)`. They group records
//! with the same key, and apply user supplied functions to the key and a list of values, which are
//! expected to populate a list of output values.
//!
//! Several variants of `group` exist which allow more precise control over how grouping is done.
//! For example, the `_by` suffixed variants take arbitrary data, but require a key-value selector
//! to be applied to each record. The `_u` suffixed variants use unsigned integers as keys, and
//! will use a dense array rather than a `HashMap` to store their keys.
//!
//! The list of values are presented as an iterator which internally merges sorted lists of values.
//! This ordering can be exploited in several cases to avoid computation when only the first few
//! elements are required.
//!
//! #Examples
//!
//! This example groups a stream of `(key,val)` pairs by `key`, and yields only the most frequently
//! occurring value for each key.
//!
//! ```ignore
//! stream.group(|key, vals, output| {
//!     let (mut max_val, mut max_wgt) = vals.peek().unwrap();
//!     for (val, wgt) in vals {
//!         if wgt > max_wgt {
//!             max_wgt = wgt;
//!             max_val = val;
//!         }
//!     }
//!     output.push((max_val.clone(), max_wgt));
//! })
//! ```

use std::rc::Rc;
use std::default::Default;
use std::hash::{Hash, Hasher};
use std::collections::HashMap;
use std::fmt::Debug;
use std::ops::DerefMut;

use itertools::Itertools;

use ::Data;
use timely::dataflow::*;
use timely::dataflow::operators::{Map, Binary};
use timely::dataflow::channels::pact::Exchange;
use timely::drain::DrainExt;

use collection::{LeastUpperBound, Lookup, Trace, Offset};
use collection::trace::CollectionIterator;

use iterators::coalesce::Coalesce;
use radix_sort::{RadixSorter, Unsigned};
use collection::compact::Compact;

/// Extension trait for the `group_by` and `group_by_u` differential dataflow methods.
pub trait CoGroupBy<G: Scope, K: Data, V1: Data> : Binary<G, ((K,V1), i32)>+Map<G, ((K,V1), i32)>
where G::Timestamp: LeastUpperBound {

    /// The lowest level `cogroup` implementation, which is parameterized by the type of storage to
    /// use for mapping keys `K` to `Offset`, an internal `CollectionTrace` type. This method should
    /// probably rarely be used directly.
    fn cogroup_by_inner<
        D:     Data,
        V2:    Data+Default,
        V3:    Data+Default,
        U:     Unsigned+Default,
        KH:    Fn(&K)->U+'static,
        Look:  Lookup<K, Offset>+'static,
        LookG: Fn(u64)->Look,
        Logic: Fn(&K, &mut CollectionIterator<V1>, &mut CollectionIterator<V2>, &mut Vec<(V3, i32)>)+'static,
        Reduc: Fn(&K, &V3)->D+'static,
    >
    (&self, other: &Stream<G, ((K,V2),i32)>, key_h: KH, reduc: Reduc, look: LookG, logic: Logic) -> Stream<G, (D, i32)> {

        let mut source1 = Trace::new(look(0));
        let mut source2 = Trace::new(look(0));
        let mut result = Trace::new(look(0));

        // A map from times to received (key, val, wgt) triples.
        let mut inputs1 = Vec::new();
        let mut inputs2 = Vec::new();

        // A map from times to a list of keys that need processing at that time.
        let mut to_do = Vec::new();

        // temporary storage for operator implementations to populate
        let mut buffer = vec![];
        let mut heap1 = vec![];
        let mut heap2 = vec![];
        let mut heap3 = vec![];

        let key_h = Rc::new(key_h);
        let key_1 = key_h.clone();
        let key_2 = key_h.clone();

        // create an exchange channel based on the supplied Fn(&D1)->u64.
        let exch1 = Exchange::new(move |&(ref x,_)| key_1(&x.0).as_u64());
        let exch2 = Exchange::new(move |&(ref x,_)| key_2(&x.0).as_u64());

        let mut sorter1 = RadixSorter::new();
        let mut sorter2 = RadixSorter::new();

        // fabricate a data-parallel operator using the `unary_notify` pattern.
        self.binary_notify(other, exch1, exch2, "CoGroupBy", vec![], move |input1, input2, output, notificator| {

            // 1. read each input, and stash it in our staging area
            while let Some((time, data)) = input1.next() {
                notificator.notify_at(&time);
                inputs1.entry_or_insert(time.clone(), || Vec::new())
                       .push(::std::mem::replace(data.deref_mut(), Vec::new()));
            }

            // 1. read each input, and stash it in our staging area
            while let Some((time, data)) = input2.next() {
                notificator.notify_at(&time);
                inputs2.entry_or_insert(time.clone(), || Vec::new())
                       .push(::std::mem::replace(data.deref_mut(), Vec::new()));
            }

            // 2. go through each time of interest that has reached completion
            // times are interesting either because we received data, or because we conclude
            // in the processing of a time that a future time will be interesting.
            while let Some((index, _count)) = notificator.next() {

                // 2a. fetch any data associated with this time.
                if let Some(mut queue) = inputs1.remove_key(&index) {

                    // sort things; radix if many, .sort_by if few.
                    let compact = if queue.len() > 1 {
                        for element in queue.into_iter() {
                            sorter1.extend(element.into_iter(), &|x| key_h(&(x.0).0));
                        }
                        let mut sorted = sorter1.finish(&|x| key_h(&(x.0).0));
                        let result = Compact::from_radix(&mut sorted, &|k| key_h(k));
                        sorted.truncate(256);
                        sorter1.recycle(sorted);
                        result
                    }
                    else {
                        let mut vec = queue.pop().unwrap();
                        let mut vec = vec.drain_temp().collect::<Vec<_>>();
                        vec.sort_by(|x,y| key_h(&(x.0).0).cmp(&key_h((&(y.0).0))));
                        Compact::from_radix(&mut vec![vec], &|k| key_h(k))
                    };

                    if let Some(compact) = compact {

                        for key in &compact.keys {
                            for time in source1.interesting_times(key, index.clone()).iter() {
                                let mut queue = to_do.entry_or_insert((*time).clone(), || { notificator.notify_at(time); Vec::new() });
                                queue.push((*key).clone());
                            }
                        }

                        source1.set_difference(index.clone(), compact);
                    }
                }

                // 2a. fetch any data associated with this time.
                if let Some(mut queue) = inputs2.remove_key(&index) {

                    // sort things; radix if many, .sort_by if few.
                    let compact = if queue.len() > 1 {
                        for element in queue.into_iter() {
                            sorter2.extend(element.into_iter(), &|x| key_h(&(x.0).0));
                        }
                        let mut sorted = sorter2.finish(&|x| key_h(&(x.0).0));
                        let result = Compact::from_radix(&mut sorted, &|k| key_h(k));
                        sorted.truncate(256);
                        sorter2.recycle(sorted);
                        result
                    }
                    else {
                        let mut vec = queue.pop().unwrap();
                        let mut vec = vec.drain_temp().collect::<Vec<_>>();
                        vec.sort_by(|x,y| key_h(&(x.0).0).cmp(&key_h((&(y.0).0))));
                        Compact::from_radix(&mut vec![vec], &|k| key_h(k))
                    };

                    if let Some(compact) = compact {

                        for key in &compact.keys {
                            for time in source2.interesting_times(key, index.clone()).iter() {
                                let mut queue = to_do.entry_or_insert((*time).clone(), || { notificator.notify_at(time); Vec::new() });
                                queue.push((*key).clone());
                            }
                        }

                        source2.set_difference(index.clone(), compact);
                    }
                }

                // we may need to produce output at index
                let mut session = output.session(&index);

                    // 2b. We must now determine for each interesting key at this time, how does the
                    // currently reported output match up with what we need as output. Should we send
                    // more output differences, and what are they?

                // Much of this logic used to hide in `OperatorTrace` and `CollectionTrace`.
                // They are now gone and simpler, respectively.
                if let Some(mut keys) = to_do.remove_key(&index) {

                    // we would like these keys in a particular order.
                    // TODO : use a radix sort since we have `key_h`.
                    keys.sort_by(|x,y| (key_h(&x), x).cmp(&(key_h(&y), y)));
                    keys.dedup();

                    // accumulations for installation into result
                    let mut accumulation = Compact::new(0,0);

                    for key in keys {

                        // acquire an iterator over the collection at `time`.
                        let mut input1 = unsafe { source1.get_collection_using(&key, &index, &mut heap1) };
                        let mut input2 = unsafe { source2.get_collection_using(&key, &index, &mut heap2) };

                        // if we have some data, invoke logic to populate self.dst
                        if input1.peek().is_some() || input2.peek().is_some() { logic(&key, &mut input1, &mut input2, &mut buffer); }

                        buffer.sort_by(|x,y| x.0.cmp(&y.0));

                        // push differences in to Compact.
                        let mut compact = accumulation.session();
                        for (val, wgt) in Coalesce::coalesce(unsafe { result.get_collection_using(&key, &index, &mut heap3) }
                                                                   .map(|(v, w)| (v,-w))
                                                                   .merge_by(buffer.iter().map(|&(ref v, w)| (v, w)), |x,y| {
                                                                        x.0 <= y.0
                                                                   }))
                        {
                            session.give((reduc(&key, val), wgt));
                            compact.push(val.clone(), wgt);
                        }
                        compact.done(key);
                        buffer.clear();
                    }

                    if accumulation.vals.len() > 0 {
                        // println!("group2");
                        result.set_difference(index.clone(), accumulation);
                    }
                }
            }
        })
    }
}
<commit_msg>cogroup implementation; untested<commit_after>//! Group records by a key, and apply a reduction function.
//!
//! The `group` operators act on data that can be viewed as pairs `(key, val)`. They group records
//! with the same key, and apply user supplied functions to the key and a list of values, which are
//! expected to populate a list of output values.
//!
//! Several variants of `group` exist which allow more precise control over how grouping is done.
//! For example, the `_by` suffixed variants take arbitrary data, but require a key-value selector
//! to be applied to each record. The `_u` suffixed variants use unsigned integers as keys, and
//! will use a dense array rather than a `HashMap` to store their keys.
//!
//! The list of values are presented as an iterator which internally merges sorted lists of values.
//! This ordering can be exploited in several cases to avoid computation when only the first few
//! elements are required.
//!
//! #Examples
//!
//! This example groups a stream of `(key,val)` pairs by `key`, and yields only the most frequently
//! occurring value for each key.
//!
//! ```ignore
//! stream.group(|key, vals, output| {
//!     let (mut max_val, mut max_wgt) = vals.peek().unwrap();
//!     for (val, wgt) in vals {
//!         if wgt > max_wgt {
//!             max_wgt = wgt;
//!             max_val = val;
//!         }
//!     }
//!     output.push((max_val.clone(), max_wgt));
//! })
//! ```

use std::rc::Rc;
use std::default::Default;
use std::hash::Hasher;
use std::ops::DerefMut;

use itertools::Itertools;

use ::Data;
use timely::dataflow::*;
use timely::dataflow::operators::{Map, Binary};
use timely::dataflow::channels::pact::Exchange;
use timely::drain::DrainExt;

use collection::{LeastUpperBound, Lookup, Trace, Offset};
use collection::trace::CollectionIterator;

use iterators::coalesce::Coalesce;
use radix_sort::{RadixSorter, Unsigned};
use collection::compact::Compact;

/// Extension trait for the `group_by` and `group_by_u` differential dataflow methods.
pub trait CoGroupBy<G: Scope, K: Data, V1: Data> : Binary<G, ((K,V1), i32)>+Map<G, ((K,V1), i32)>
where G::Timestamp: LeastUpperBound {

    /// A primitive binary version of `group_by`, which acts on a `Stream<((K,V1),i32)` and a `Stream<((K,V2),i32)`.
    ///
    /// The two streams must already be key-value pairs, which is too bad. Also, in addition to the
    /// normal arguments (another stream, a hash for the key, a reduction function, and per-key logic),
    /// the user must specify a function implmenting `Fn(u64) -> Look`, where `Look: Lookup<K, Offset>` is something you shouldn't have to know about yet.
    /// The right thing to use here, for the moment, is `|_| HashMap::new()`.
    ///
    /// There are better options if you know your key is an unsigned integer, namely `|x| (Vec::new(), x)`.
    fn cogroup_by_inner<
        D:     Data,
        V2:    Data+Default,
        V3:    Data+Default,
        U:     Unsigned+Default,
        KH:    Fn(&K)->U+'static,
        Look:  Lookup<K, Offset>+'static,
        LookG: Fn(u64)->Look,
        Logic: Fn(&K, &mut CollectionIterator<V1>, &mut CollectionIterator<V2>, &mut Vec<(V3, i32)>)+'static,
        Reduc: Fn(&K, &V3)->D+'static,
    >
    (&self, other: &Stream<G, ((K,V2),i32)>, key_h: KH, reduc: Reduc, look: LookG, logic: Logic) -> Stream<G, (D, i32)> {

        let mut source1 = Trace::new(look(0));
        let mut source2 = Trace::new(look(0));
        let mut result = Trace::new(look(0));

        // A map from times to received (key, val, wgt) triples.
        let mut inputs1 = Vec::new();
        let mut inputs2 = Vec::new();

        // A map from times to a list of keys that need processing at that time.
        let mut to_do = Vec::new();

        // temporary storage for operator implementations to populate
        let mut buffer = vec![];
        let mut heap1 = vec![];
        let mut heap2 = vec![];
        let mut heap3 = vec![];

        let key_h = Rc::new(key_h);
        let key_1 = key_h.clone();
        let key_2 = key_h.clone();

        // create an exchange channel based on the supplied Fn(&D1)->u64.
        let exch1 = Exchange::new(move |&((ref k, _),_)| key_1(k).as_u64());
        let exch2 = Exchange::new(move |&((ref k, _),_)| key_2(k).as_u64());

        let mut sorter1 = RadixSorter::new();
        let mut sorter2 = RadixSorter::new();

        // fabricate a data-parallel operator using the `unary_notify` pattern.
        self.binary_notify(other, exch1, exch2, "CoGroupBy", vec![], move |input1, input2, output, notificator| {

            // 1. read each input, and stash it in our staging area
            while let Some((time, data)) = input1.next() {
                notificator.notify_at(&time);
                inputs1.entry_or_insert(time.clone(), || Vec::new())
                       .push(::std::mem::replace(data.deref_mut(), Vec::new()));
            }

            // 1. read each input, and stash it in our staging area
            while let Some((time, data)) = input2.next() {
                notificator.notify_at(&time);
                inputs2.entry_or_insert(time.clone(), || Vec::new())
                       .push(::std::mem::replace(data.deref_mut(), Vec::new()));
            }

            // 2. go through each time of interest that has reached completion
            // times are interesting either because we received data, or because we conclude
            // in the processing of a time that a future time will be interesting.
            while let Some((index, _count)) = notificator.next() {

                // 2a. fetch any data associated with this time.
                if let Some(mut queue) = inputs1.remove_key(&index) {

                    // sort things; radix if many, .sort_by if few.
                    let compact = if queue.len() > 1 {
                        for element in queue.into_iter() {
                            sorter1.extend(element.into_iter(), &|x| key_h(&(x.0).0));
                        }
                        let mut sorted = sorter1.finish(&|x| key_h(&(x.0).0));
                        let result = Compact::from_radix(&mut sorted, &|k| key_h(k));
                        sorted.truncate(256);
                        sorter1.recycle(sorted);
                        result
                    }
                    else {
                        let mut vec = queue.pop().unwrap();
                        let mut vec = vec.drain_temp().collect::<Vec<_>>();
                        vec.sort_by(|x,y| key_h(&(x.0).0).cmp(&key_h((&(y.0).0))));
                        Compact::from_radix(&mut vec![vec], &|k| key_h(k))
                    };

                    if let Some(compact) = compact {

                        for key in &compact.keys {
                            for time in source1.interesting_times(key, index.clone()).iter() {
                                let mut queue = to_do.entry_or_insert((*time).clone(), || { notificator.notify_at(time); Vec::new() });
                                queue.push((*key).clone());
                            }
                        }

                        source1.set_difference(index.clone(), compact);
                    }
                }

                // 2a. fetch any data associated with this time.
                if let Some(mut queue) = inputs2.remove_key(&index) {

                    // sort things; radix if many, .sort_by if few.
                    let compact = if queue.len() > 1 {
                        for element in queue.into_iter() {
                            sorter2.extend(element.into_iter(), &|x| key_h(&(x.0).0));
                        }
                        let mut sorted = sorter2.finish(&|x| key_h(&(x.0).0));
                        let result = Compact::from_radix(&mut sorted, &|k| key_h(k));
                        sorted.truncate(256);
                        sorter2.recycle(sorted);
                        result
                    }
                    else {
                        let mut vec = queue.pop().unwrap();
                        let mut vec = vec.drain_temp().collect::<Vec<_>>();
                        vec.sort_by(|x,y| key_h(&(x.0).0).cmp(&key_h((&(y.0).0))));
                        Compact::from_radix(&mut vec![vec], &|k| key_h(k))
                    };

                    if let Some(compact) = compact {

                        for key in &compact.keys {
                            for time in source2.interesting_times(key, index.clone()).iter() {
                                let mut queue = to_do.entry_or_insert((*time).clone(), || { notificator.notify_at(time); Vec::new() });
                                queue.push((*key).clone());
                            }
                        }

                        source2.set_difference(index.clone(), compact);
                    }
                }

                // we may need to produce output at index
                let mut session = output.session(&index);

                    // 2b. We must now determine for each interesting key at this time, how does the
                    // currently reported output match up with what we need as output. Should we send
                    // more output differences, and what are they?

                // Much of this logic used to hide in `OperatorTrace` and `CollectionTrace`.
                // They are now gone and simpler, respectively.
                if let Some(mut keys) = to_do.remove_key(&index) {

                    // we would like these keys in a particular order.
                    // TODO : use a radix sort since we have `key_h`.
                    keys.sort_by(|x,y| (key_h(&x), x).cmp(&(key_h(&y), y)));
                    keys.dedup();

                    // accumulations for installation into result
                    let mut accumulation = Compact::new(0,0);

                    for key in keys {

                        // acquire an iterator over the collection at `time`.
                        let mut input1 = unsafe { source1.get_collection_using(&key, &index, &mut heap1) };
                        let mut input2 = unsafe { source2.get_collection_using(&key, &index, &mut heap2) };

                        // if we have some data, invoke logic to populate self.dst
                        if input1.peek().is_some() || input2.peek().is_some() { logic(&key, &mut input1, &mut input2, &mut buffer); }

                        buffer.sort_by(|x,y| x.0.cmp(&y.0));

                        // push differences in to Compact.
                        let mut compact = accumulation.session();
                        for (val, wgt) in Coalesce::coalesce(unsafe { result.get_collection_using(&key, &index, &mut heap3) }
                                                                   .map(|(v, w)| (v,-w))
                                                                   .merge_by(buffer.iter().map(|&(ref v, w)| (v, w)), |x,y| {
                                                                        x.0 <= y.0
                                                                   }))
                        {
                            session.give((reduc(&key, val), wgt));
                            compact.push(val.clone(), wgt);
                        }
                        compact.done(key);
                        buffer.clear();
                    }

                    if accumulation.vals.len() > 0 {
                        // println!("group2");
                        result.set_difference(index.clone(), accumulation);
                    }
                }
            }
        })
    }
}
<|endoftext|>


<commit_before>use beacon_chain::{
    fork_choice::BitwiseLMDGhost,
    slot_clock::SystemTimeSlotClock,
    store::{DiskStore, MemoryStore, Store},
    BeaconChain, BeaconChainTypes,
};
use fork_choice::ForkChoice;
use slog::{info, Logger};
use slot_clock::SlotClock;
use std::sync::Arc;
use tree_hash::TreeHash;
use types::{
    test_utils::TestingBeaconStateBuilder, BeaconBlock, EthSpec, FewValidatorsEthSpec, Hash256,
};

/// Provides a new, initialized `BeaconChain`
pub trait InitialiseBeaconChain<T: BeaconChainTypes> {
    fn initialise_beacon_chain(store: Arc<T::Store>, log: Logger) -> BeaconChain<T>;
}

/// A testnet-suitable BeaconChainType, using `MemoryStore`.
#[derive(Clone)]
pub struct TestnetMemoryBeaconChainTypes;

impl BeaconChainTypes for TestnetMemoryBeaconChainTypes {
    type Store = MemoryStore;
    type SlotClock = SystemTimeSlotClock;
    type ForkChoice = BitwiseLMDGhost<Self::Store, Self::EthSpec>;
    type EthSpec = FewValidatorsEthSpec;
}

impl<T: BeaconChainTypes> InitialiseBeaconChain<T> for TestnetMemoryBeaconChainTypes {
    fn initialise_beacon_chain(store: Arc<T::Store>, log: Logger) -> BeaconChain<T> {
        maybe_load_from_store_for_testnet::<_, T::Store, T::EthSpec>(store, log)
    }
}

/// A testnet-suitable BeaconChainType, using `DiskStore`.
#[derive(Clone)]
pub struct TestnetDiskBeaconChainTypes;

impl BeaconChainTypes for TestnetDiskBeaconChainTypes {
    type Store = DiskStore;
    type SlotClock = SystemTimeSlotClock;
    type ForkChoice = BitwiseLMDGhost<Self::Store, Self::EthSpec>;
    type EthSpec = FewValidatorsEthSpec;
}

impl<T: BeaconChainTypes> InitialiseBeaconChain<T> for TestnetDiskBeaconChainTypes {
    fn initialise_beacon_chain(store: Arc<T::Store>, log: Logger) -> BeaconChain<T> {
        maybe_load_from_store_for_testnet::<_, T::Store, T::EthSpec>(store, log)
    }
}

/// Loads a `BeaconChain` from `store`, if it exists. Otherwise, create a new chain from genesis.
fn maybe_load_from_store_for_testnet<T, U: Store, V: EthSpec>(
    store: Arc<U>,
    log: Logger,
) -> BeaconChain<T>
where
    T: BeaconChainTypes<Store = U>,
    T::ForkChoice: ForkChoice<U>,
{
    if let Ok(Some(beacon_chain)) = BeaconChain::from_store(store.clone()) {
        info!(
            log,
            "Loaded BeaconChain from store";
            "slot" => beacon_chain.state.read().slot,
            "best_slot" => beacon_chain.best_slot(),
        );

        beacon_chain
    } else {
        info!(log, "Initializing new BeaconChain from genesis");
        let spec = T::EthSpec::spec();

        let state_builder =
            TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(8, &spec);
        let (genesis_state, _keypairs) = state_builder.build();

        let mut genesis_block = BeaconBlock::empty(&spec);
        genesis_block.state_root = Hash256::from_slice(&genesis_state.tree_hash_root());

        // Slot clock
        let slot_clock = T::SlotClock::new(
            spec.genesis_slot,
            genesis_state.genesis_time,
            spec.seconds_per_slot,
        );
        // Choose the fork choice
        let fork_choice = T::ForkChoice::new(store.clone());

        // Genesis chain
        //TODO: Handle error correctly
        BeaconChain::from_genesis(
            store,
            slot_clock,
            genesis_state,
            genesis_block,
            spec.clone(),
            fork_choice,
        )
        .expect("Terminate if beacon chain generation fails")
    }
}
<commit_msg>Use `LighthouseTestnet` params, not `FewValidators`<commit_after>use beacon_chain::{
    fork_choice::BitwiseLMDGhost,
    slot_clock::SystemTimeSlotClock,
    store::{DiskStore, MemoryStore, Store},
    BeaconChain, BeaconChainTypes,
};
use fork_choice::ForkChoice;
use slog::{info, Logger};
use slot_clock::SlotClock;
use std::sync::Arc;
use tree_hash::TreeHash;
use types::{
    test_utils::TestingBeaconStateBuilder, BeaconBlock, EthSpec, Hash256, LighthouseTestnetEthSpec,
};

/// Provides a new, initialized `BeaconChain`
pub trait InitialiseBeaconChain<T: BeaconChainTypes> {
    fn initialise_beacon_chain(store: Arc<T::Store>, log: Logger) -> BeaconChain<T>;
}

/// A testnet-suitable BeaconChainType, using `MemoryStore`.
#[derive(Clone)]
pub struct TestnetMemoryBeaconChainTypes;

impl BeaconChainTypes for TestnetMemoryBeaconChainTypes {
    type Store = MemoryStore;
    type SlotClock = SystemTimeSlotClock;
    type ForkChoice = BitwiseLMDGhost<Self::Store, Self::EthSpec>;
    type EthSpec = LighthouseTestnetEthSpec;
}

impl<T: BeaconChainTypes> InitialiseBeaconChain<T> for TestnetMemoryBeaconChainTypes {
    fn initialise_beacon_chain(store: Arc<T::Store>, log: Logger) -> BeaconChain<T> {
        maybe_load_from_store_for_testnet::<_, T::Store, T::EthSpec>(store, log)
    }
}

/// A testnet-suitable BeaconChainType, using `DiskStore`.
#[derive(Clone)]
pub struct TestnetDiskBeaconChainTypes;

impl BeaconChainTypes for TestnetDiskBeaconChainTypes {
    type Store = DiskStore;
    type SlotClock = SystemTimeSlotClock;
    type ForkChoice = BitwiseLMDGhost<Self::Store, Self::EthSpec>;
    type EthSpec = LighthouseTestnetEthSpec;
}

impl<T: BeaconChainTypes> InitialiseBeaconChain<T> for TestnetDiskBeaconChainTypes {
    fn initialise_beacon_chain(store: Arc<T::Store>, log: Logger) -> BeaconChain<T> {
        maybe_load_from_store_for_testnet::<_, T::Store, T::EthSpec>(store, log)
    }
}

/// Loads a `BeaconChain` from `store`, if it exists. Otherwise, create a new chain from genesis.
fn maybe_load_from_store_for_testnet<T, U: Store, V: EthSpec>(
    store: Arc<U>,
    log: Logger,
) -> BeaconChain<T>
where
    T: BeaconChainTypes<Store = U>,
    T::ForkChoice: ForkChoice<U>,
{
    if let Ok(Some(beacon_chain)) = BeaconChain::from_store(store.clone()) {
        info!(
            log,
            "Loaded BeaconChain from store";
            "slot" => beacon_chain.state.read().slot,
            "best_slot" => beacon_chain.best_slot(),
        );

        beacon_chain
    } else {
        info!(log, "Initializing new BeaconChain from genesis");
        let spec = T::EthSpec::spec();

        let state_builder =
            TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(8, &spec);
        let (genesis_state, _keypairs) = state_builder.build();

        let mut genesis_block = BeaconBlock::empty(&spec);
        genesis_block.state_root = Hash256::from_slice(&genesis_state.tree_hash_root());

        // Slot clock
        let slot_clock = T::SlotClock::new(
            spec.genesis_slot,
            genesis_state.genesis_time,
            spec.seconds_per_slot,
        );
        // Choose the fork choice
        let fork_choice = T::ForkChoice::new(store.clone());

        // Genesis chain
        //TODO: Handle error correctly
        BeaconChain::from_genesis(
            store,
            slot_clock,
            genesis_state,
            genesis_block,
            spec.clone(),
            fork_choice,
        )
        .expect("Terminate if beacon chain generation fails")
    }
}
<|endoftext|>


<commit_before>extern crate made_up;
extern crate log;
use std::env;
fn main() {
    log::set_logger(|max_log_level| {
                        max_log_level.set(::log::LogLevelFilter::Debug);
                        Box::new(SimpleLogger)
                    })
            .unwrap();
    // NOTE Not a very robust way to do it
    let dir = env::args().nth(1).expect("Usage: made-up <root_directory>");
    made_up::generate_site(dir).unwrap();
}

use log::{LogLevel, LogRecord, LogMetadata};

struct SimpleLogger;

impl ::log::Log for SimpleLogger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= LogLevel::Info
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            println!("{} - {}", record.level(), record.args());
        }
    }
}
<commit_msg>Allow for no argument execution<commit_after>extern crate made_up;
extern crate log;
use std::env;
fn main() {
    log::set_logger(|max_log_level| {
                        max_log_level.set(::log::LogLevelFilter::Debug);
                        Box::new(SimpleLogger)
                    })
            .unwrap();
    // NOTE Not a very robust way to do it
    // TODO Add in help display on -h or --help
    // .expect("Usage: made-up <root_directory>");
    let dir = env::args().nth(1).unwrap_or(".");
    made_up::generate_site(dir).unwrap_or_else(|x| println!("{:?}", x));
}

use log::{LogLevel, LogRecord, LogMetadata};

struct SimpleLogger;

impl ::log::Log for SimpleLogger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= LogLevel::Info
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            println!("{} - {}", record.level(), record.args());
        }
    }
}
<|endoftext|>


<commit_before>#![cfg(feature = "getopts")]
extern crate getopts;
extern crate markov;

use std::env::args;
use std::path::Path;
use getopts::Options;
use markov::Chain;

#[cfg(not(test))]
fn main() {
    markov_gen(args().collect()).iter().map(|s| println!("{}", s)).count();
}

/// Generates a number of strings using a markov chain on specified inputs. This is designed
/// primarily for command-line usage. The arguments are expected to be paths to files to be fed
/// into the chain. Additionally, the argument `-n #` is supported to specify the number of phrases
/// to be generated. This number must be a positive, non-zero integer. `-o #` is also supported to
/// specify the order of Markov chain to be used. Note `-o` must be specified before any file
/// names and must also be a positive, non-zero integer.
///
/// Some valid usages of this function:
/// `markov_gen(vec!["test".to_owned()])`
/// `markov_gen(vec!["test".to_owned(), "-n".to_owned(), "3".to_owned()])`
/// `markov_gen(vec!["-n".to_owned(), "3".to_owned(), "test".to_owned()])`
/// `markov_gen(vec!["-o".to_owned(), "2".to_owned(), "test".to_owned()])`
///
/// Some invalid usages of this function:
/// `markov_gen(vec!["-n".to_owned(), "3".to_owned()])`
/// `markov_gen(vec!["test".to_owned(), "-n".to_owned(), "0".to_owned()])`
/// `markov_gen(vec!["test".to_owned(), "-n".to_owned(), "test".to_owned()])`
/// `markov_gen(vec!["test".to_owned(), "-o".to_owned(), "3".to_owned()])`
/// `markov_gen(vec!["-o".to_owned(), "0".to_owned()], "test".to_owned())`
fn markov_gen(args: Vec<String>) -> Vec<String> {
    let mut opts = Options::new();
    opts.optopt("n", "count", "set the number of phrases to generate", "3");
    opts.optopt("o", "order", "set the order of the Markov chain", "2");
    opts.optflag("h", "help", "print this help menu");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(e) => panic!(e.to_string()),
    };
    if matches.opt_present("h") {
        let brief = format!("Usage: {} FILE1 [FILE2 FILE3 ...] [options]", args[0]);
        print!("{}", opts.usage(&brief));
        Vec::new()
    } else {
        let mut chain: Chain<String> = Chain::new();
        let count = match matches.opt_str("n") {
            Some(arg) => match arg.parse() {
                Ok(n) if n > 0 => n,
                _ => panic!("Expected positive integer argument to -n, found {}.", &arg),
            },
            None => 1
        };
        if let Some(arg) = matches.opt_str("o") {
            match arg.parse() {
                Ok(n) if n > 0 => { chain.order(n); },
                _ => panic!("Expected positive integer argument to -n, found {}.", &arg),
            }
        }
        for path in matches.free.iter() {
            chain.feed_file(Path::new(&path));
        }
        if chain.is_empty() { panic!("No files were fed into the chain.") }
        chain.str_iter_for(count).collect()
    }
}

#[cfg(test)]
mod test {
    use super::markov_gen;
    use std::borrow::ToOwned;

    #[test]
    fn gen_default() {
        assert_eq!(markov_gen(vec!["test".to_owned()]).len(), 1)
    }

    #[test]
    fn gen_number_after() {
        assert_eq!(markov_gen(
            vec!["test".to_owned(), "-n".to_owned(), "3".to_owned()]
        ).len(), 3)
    }

    #[test]
    fn gen_before_number() {
        assert_eq!(markov_gen(
            vec!["-n".to_owned(), "3".to_owned(), "test".to_owned()]
        ).len(), 3)
    }

    #[test]
    #[should_panic(message = "No files were fed into the chain.")]
    fn gen_invalid_no_files() {
        markov_gen(vec!["-n".to_owned(), "3".to_owned()]);
    }

    #[test]
    #[should_panic(message = "Expected positive integer argument to -n, found 0.")]
    fn gen_invalid_n_arg_zero() {
        markov_gen(vec!["test".to_owned(), "-n".to_owned(), "0".to_owned()]);
    }

    #[test]
    #[should_panic(message = "Expected positive integer argument to -n, found test.")]
    fn gen_invalid_n_arg_string() {
        markov_gen(vec!["test".to_owned(), "-n".to_owned(), "test".to_owned()]);
    }
}
<commit_msg>Updated help menu text.<commit_after>#![cfg(feature = "getopts")]
extern crate getopts;
extern crate markov;

use std::env::args;
use std::path::Path;
use getopts::Options;
use markov::Chain;

#[cfg(not(test))]
fn main() {
    markov_gen(args().collect()).iter().map(|s| println!("{}", s)).count();
}

/// Generates a number of strings using a markov chain on specified inputs. This is designed
/// primarily for command-line usage. The arguments are expected to be paths to files to be fed
/// into the chain. Additionally, the argument `-n #` is supported to specify the number of phrases
/// to be generated. This number must be a positive, non-zero integer. `-o #` is also supported to
/// specify the order of Markov chain to be used. Note `-o` must be specified before any file
/// names and must also be a positive, non-zero integer.
///
/// Some valid usages of this function:
/// `markov_gen(vec!["test".to_owned()])`
/// `markov_gen(vec!["test".to_owned(), "-n".to_owned(), "3".to_owned()])`
/// `markov_gen(vec!["-n".to_owned(), "3".to_owned(), "test".to_owned()])`
/// `markov_gen(vec!["-o".to_owned(), "2".to_owned(), "test".to_owned()])`
///
/// Some invalid usages of this function:
/// `markov_gen(vec!["-n".to_owned(), "3".to_owned()])`
/// `markov_gen(vec!["test".to_owned(), "-n".to_owned(), "0".to_owned()])`
/// `markov_gen(vec!["test".to_owned(), "-n".to_owned(), "test".to_owned()])`
/// `markov_gen(vec!["test".to_owned(), "-o".to_owned(), "3".to_owned()])`
/// `markov_gen(vec!["-o".to_owned(), "0".to_owned()], "test".to_owned())`
fn markov_gen(args: Vec<String>) -> Vec<String> {
    let mut opts = Options::new();
    opts.optopt("n", "count", "set the number of phrases to generate", "COUNT");
    opts.optopt("o", "order", "set the order of the Markov chain", "ORDER");
    opts.optflag("h", "help", "print this help menu");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(e) => panic!(e.to_string()),
    };
    if matches.opt_present("h") {
        let brief = format!("Usage: {} FILE1 [FILE2 FILE3 ...] [options]", args[0]);
        print!("{}", opts.usage(&brief));
        Vec::new()
    } else {
        let mut chain: Chain<String> = Chain::new();
        let count = match matches.opt_str("n") {
            Some(arg) => match arg.parse() {
                Ok(n) if n > 0 => n,
                _ => panic!("Expected positive integer argument to -n, found {}.", &arg),
            },
            None => 1
        };
        if let Some(arg) = matches.opt_str("o") {
            match arg.parse() {
                Ok(n) if n > 0 => { chain.order(n); },
                _ => panic!("Expected positive integer argument to -n, found {}.", &arg),
            }
        }
        for path in matches.free.iter() {
            chain.feed_file(Path::new(&path));
        }
        if chain.is_empty() { panic!("No files were fed into the chain.") }
        chain.str_iter_for(count).collect()
    }
}

#[cfg(test)]
mod test {
    use super::markov_gen;
    use std::borrow::ToOwned;

    #[test]
    fn gen_default() {
        assert_eq!(markov_gen(vec!["test".to_owned()]).len(), 1)
    }

    #[test]
    fn gen_number_after() {
        assert_eq!(markov_gen(
            vec!["test".to_owned(), "-n".to_owned(), "3".to_owned()]
        ).len(), 3)
    }

    #[test]
    fn gen_before_number() {
        assert_eq!(markov_gen(
            vec!["-n".to_owned(), "3".to_owned(), "test".to_owned()]
        ).len(), 3)
    }

    #[test]
    #[should_panic(message = "No files were fed into the chain.")]
    fn gen_invalid_no_files() {
        markov_gen(vec!["-n".to_owned(), "3".to_owned()]);
    }

    #[test]
    #[should_panic(message = "Expected positive integer argument to -n, found 0.")]
    fn gen_invalid_n_arg_zero() {
        markov_gen(vec!["test".to_owned(), "-n".to_owned(), "0".to_owned()]);
    }

    #[test]
    #[should_panic(message = "Expected positive integer argument to -n, found test.")]
    fn gen_invalid_n_arg_string() {
        markov_gen(vec!["test".to_owned(), "-n".to_owned(), "test".to_owned()]);
    }
}
<|endoftext|>


<commit_before>use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

use bit_vec::BitVec;

struct HashIter<'a, T>
where
    T: 'a,
{
    m: usize,
    k: usize,
    i: usize,
    obj: &'a T,
}

impl<'a, T> HashIter<'a, T>
where
    T: 'a,
{
    fn new(m: usize, k: usize, obj: &'a T) -> HashIter<'a, T> {
        HashIter {
            m: m,
            k: k,
            i: 0,
            obj: obj,
        }
    }
}

impl<'a, T> Iterator for HashIter<'a, T>
where
    T: 'a + Hash,
{
    type Item = usize;

    fn next(&mut self) -> Option<usize> {
        if self.i < self.k {
            let mut hasher = DefaultHasher::new();
            hasher.write_usize(self.i);
            self.obj.hash(&mut hasher);
            let x = (hasher.finish() as usize) % self.m;

            self.i += 1;

            Some(x)
        } else {
            None
        }
    }
}

pub struct BloomFilter {
    bv: BitVec,
    k: usize,
}

impl BloomFilter {
    pub fn with_params(m: usize, k: usize) -> BloomFilter {
        BloomFilter {
            bv: BitVec::from_elem(m, false),
            k: k,
        }
    }

    pub fn with_properties(n: usize, p: f64) -> BloomFilter {
        assert!(n > 0);
        assert!(p > 0.);
        assert!(p < 1.);

        let k = (-p.log2()) as usize;
        let ln2 = (2f64).ln();
        let m = (-((n as f64) * p.ln()) / (ln2 * ln2)) as usize;

        BloomFilter::with_params(m, k)
    }

    pub fn get_k(&self) -> usize {
        self.k
    }

    pub fn get_m(&self) -> usize {
        self.bv.len()
    }

    pub fn add<T>(&mut self, obj: &T)
    where
        T: Hash,
    {
        for pos in HashIter::new(self.bv.len(), self.k, obj) {
            self.bv.set(pos, true);
        }
    }

    pub fn query<T>(&self, obj: &T) -> bool
    where
        T: Hash,
    {
        for pos in HashIter::new(self.bv.len(), self.k, obj) {
            if !self.bv.get(pos).unwrap() {
                return false;
            }
        }
        true
    }

    pub fn clear(&mut self) {
        self.bv.clear()
    }

    pub fn union(&mut self, other: &BloomFilter) {
        assert!(self.k == other.k);
        assert!(self.bv.len() == other.bv.len());
        self.bv.union(&other.bv);
    }

    pub fn guess_n(&self) -> usize {
        let m = self.bv.len() as f64;
        let k = self.k as f64;
        let x = self.bv.iter().filter(|x| *x).count() as f64;

        (-m / k * (1. - x / m).ln()) as usize
    }
}

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

    #[test]
    fn getter() {
        let bf = BloomFilter::with_params(100, 2);
        assert!(bf.get_k() == 2);
        assert!(bf.get_m() == 100);
    }

    #[test]
    fn empty() {
        let bf = BloomFilter::with_params(100, 2);
        assert!(!bf.query(&1));
    }

    #[test]
    fn add() {
        let mut bf = BloomFilter::with_params(100, 2);

        bf.add(&1);
        assert!(bf.query(&1));
        assert!(!bf.query(&2));
    }

    #[test]
    fn clear() {
        let mut bf = BloomFilter::with_params(100, 2);

        bf.add(&1);
        bf.clear();
        assert!(!bf.query(&1));
    }

    #[test]
    fn union() {
        let mut bf1 = BloomFilter::with_params(100, 2);
        bf1.add(&1);
        assert!(bf1.query(&1));
        assert!(!bf1.query(&2));
        assert!(!bf1.query(&3));

        let mut bf2 = BloomFilter::with_params(100, 2);
        bf2.add(&2);
        assert!(!bf2.query(&1));
        assert!(bf2.query(&2));
        assert!(!bf2.query(&3));

        bf1.union(&bf2);
        assert!(bf1.query(&1));
        assert!(bf1.query(&2));
        assert!(!bf1.query(&3));
    }

    #[test]
    fn with_properties() {
        let bf = BloomFilter::with_properties(1000, 0.1);
        assert!(bf.get_k() == 3);
        assert!(bf.get_m() == 4792);
    }

    #[test]
    fn guess_n() {
        let mut bf = BloomFilter::with_params(100, 2);
        assert!(bf.guess_n() == 0);

        bf.add(&1);
        assert!(bf.guess_n() == 1);

        bf.add(&1);
        assert!(bf.guess_n() == 1);

        bf.add(&2);
        assert!(bf.guess_n() == 2);
    }
}
<commit_msg>bloomfilter: docs<commit_after>use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

use bit_vec::BitVec;

struct HashIter<'a, T>
where
    T: 'a,
{
    m: usize,
    k: usize,
    i: usize,
    obj: &'a T,
}

impl<'a, T> HashIter<'a, T>
where
    T: 'a,
{
    fn new(m: usize, k: usize, obj: &'a T) -> HashIter<'a, T> {
        HashIter {
            m: m,
            k: k,
            i: 0,
            obj: obj,
        }
    }
}

impl<'a, T> Iterator for HashIter<'a, T>
where
    T: 'a + Hash,
{
    type Item = usize;

    fn next(&mut self) -> Option<usize> {
        if self.i < self.k {
            let mut hasher = DefaultHasher::new();
            hasher.write_usize(self.i);
            self.obj.hash(&mut hasher);
            let x = (hasher.finish() as usize) % self.m;

            self.i += 1;

            Some(x)
        } else {
            None
        }
    }
}

/// Simple implementation of a [BloomFilter](https://en.wikipedia.org/wiki/Bloom_filter)
pub struct BloomFilter {
    bv: BitVec,
    k: usize,
}

impl BloomFilter {
    /// Create new, empty BloomFilter with internal parameters.
    ///
    /// - `k` is the number of hash functions
    /// - `m` is the number of bits used to store state
    pub fn with_params(m: usize, k: usize) -> BloomFilter {
        BloomFilter {
            bv: BitVec::from_elem(m, false),
            k: k,
        }
    }

    /// Create new, empty BloomFilter with given properties.
    ///
    /// - `n` number of unique elements the BloomFilter is expected to hold, must be `> 0`
    /// - `p` false positive property when querying the BloomFilter after adding `n` unique
    ///   elements, must be `> 0` and `< 1`
    ///
    /// Panics if the parameters are not in range.
    pub fn with_properties(n: usize, p: f64) -> BloomFilter {
        assert!(n > 0);
        assert!(p > 0.);
        assert!(p < 1.);

        let k = (-p.log2()) as usize;
        let ln2 = (2f64).ln();
        let m = (-((n as f64) * p.ln()) / (ln2 * ln2)) as usize;

        BloomFilter::with_params(m, k)
    }

    /// Get `k` (number of hash functions).
    pub fn get_k(&self) -> usize {
        self.k
    }

    /// Get `m` (number of stored bits).
    pub fn get_m(&self) -> usize {
        self.bv.len()
    }

    /// Add new element to the BloomFilter.
    ///
    /// If the same element is added multiple times or if an element results in the same hash
    /// signature, this method does not have any effect.
    pub fn add<T>(&mut self, obj: &T)
    where
        T: Hash,
    {
        for pos in HashIter::new(self.bv.len(), self.k, obj) {
            self.bv.set(pos, true);
        }
    }

    /// Guess if the given element was added to the BloomFilter.
    pub fn query<T>(&self, obj: &T) -> bool
    where
        T: Hash,
    {
        for pos in HashIter::new(self.bv.len(), self.k, obj) {
            if !self.bv.get(pos).unwrap() {
                return false;
            }
        }
        true
    }

    /// Clear state of the BloomFilter, so that it behaves like a fresh one.
    pub fn clear(&mut self) {
        self.bv.clear()
    }

    /// Add the entire content of another bloomfilter to this BloomFilter.
    ///
    /// The result is the same as adding all elements added to `other` to `self` in the first
    /// place.
    ///
    /// Panics if `k` and `m` of the two BloomFilters are not identical.
    pub fn union(&mut self, other: &BloomFilter) {
        assert!(self.k == other.k);
        assert!(self.bv.len() == other.bv.len());
        self.bv.union(&other.bv);
    }

    /// Guess the number of unique elements added to the BloomFilter.
    pub fn guess_n(&self) -> usize {
        let m = self.bv.len() as f64;
        let k = self.k as f64;
        let x = self.bv.iter().filter(|x| *x).count() as f64;

        (-m / k * (1. - x / m).ln()) as usize
    }
}

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

    #[test]
    fn getter() {
        let bf = BloomFilter::with_params(100, 2);
        assert!(bf.get_k() == 2);
        assert!(bf.get_m() == 100);
    }

    #[test]
    fn empty() {
        let bf = BloomFilter::with_params(100, 2);
        assert!(!bf.query(&1));
    }

    #[test]
    fn add() {
        let mut bf = BloomFilter::with_params(100, 2);

        bf.add(&1);
        assert!(bf.query(&1));
        assert!(!bf.query(&2));
    }

    #[test]
    fn clear() {
        let mut bf = BloomFilter::with_params(100, 2);

        bf.add(&1);
        bf.clear();
        assert!(!bf.query(&1));
    }

    #[test]
    fn union() {
        let mut bf1 = BloomFilter::with_params(100, 2);
        bf1.add(&1);
        assert!(bf1.query(&1));
        assert!(!bf1.query(&2));
        assert!(!bf1.query(&3));

        let mut bf2 = BloomFilter::with_params(100, 2);
        bf2.add(&2);
        assert!(!bf2.query(&1));
        assert!(bf2.query(&2));
        assert!(!bf2.query(&3));

        bf1.union(&bf2);
        assert!(bf1.query(&1));
        assert!(bf1.query(&2));
        assert!(!bf1.query(&3));
    }

    #[test]
    fn with_properties() {
        let bf = BloomFilter::with_properties(1000, 0.1);
        assert!(bf.get_k() == 3);
        assert!(bf.get_m() == 4792);
    }

    #[test]
    fn guess_n() {
        let mut bf = BloomFilter::with_params(100, 2);
        assert!(bf.guess_n() == 0);

        bf.add(&1);
        assert!(bf.guess_n() == 1);

        bf.add(&1);
        assert!(bf.guess_n() == 1);

        bf.add(&2);
        assert!(bf.guess_n() == 2);
    }
}
<|endoftext|>


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

#[cfg(feature="discovery")]
use std::thread;

use std::thread::JoinHandle;

use context::ApplicationContext;

#[cfg(feature = "discovery")]
use scanner::discovery;

#[cfg(feature="discovery")]
use utils;

use utils::logger::{BoxLogger, Logger};

#[cfg(feature="discovery")]
use utils::logger::Severity;

use futures::{BoxFuture, Future, Poll, Stream};

use futures::sync::mpsc;
use futures::sync::mpsc::{UnboundedReceiver, UnboundedSender};

use time;

/// Network scan period.
const NETWORK_SCAN_PERIOD: f64 = 300.0;

/// Different command types that the command handler can receive.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum Command {
    ResetServiceTable,
    ScanNetwork,
    PeriodicNetworkScan,
}

/// Command handler event.
#[derive(Debug, Copy, Clone)]
enum Event {
    Command(Command),
    ScanCompleted,
}

type CommandReceiver = UnboundedReceiver<Event>;

type CommandSender = UnboundedSender<Event>;

/// A channel for sending commands across threads. The channel is cloneable and every copy of it
/// will send commands to the same handler.
#[derive(Clone)]
pub struct CommandChannel {
    tx: CommandSender,
}

impl CommandChannel {
    /// Create a new command channel.
    fn new(tx: CommandSender) -> CommandChannel {
        CommandChannel {
            tx: tx,
        }
    }

    /// Send a given command.
    pub fn send(&self, cmd: Command) {
        self.tx.send(Event::Command(cmd))
            .expect("broken command channel");
    }
}

/// Command handler context.
struct CommandHandlerContext {
    app_context:  ApplicationContext,
    logger:       BoxLogger,
    cmd_sender:   CommandSender,
    last_nw_scan: f64,
    scanner:      Option<JoinHandle<()>>,
}

impl CommandHandlerContext {
    /// Create a new command handler context.
    fn new(app_context: ApplicationContext, cmd_sender: CommandSender) -> CommandHandlerContext {
        let logger = app_context.get_logger();

        let t = time::precise_time_s();

        CommandHandlerContext {
            app_context:  app_context,
            logger:       logger,
            cmd_sender:   cmd_sender,
            last_nw_scan: t - 300.0,
            scanner:      None,
        }
    }

    /// Process a given command handler event.
    fn proces_event(&mut self, event: Event) {
        match event {
            Event::Command(cmd)  => self.process_command(cmd),
            Event::ScanCompleted => self.scan_completed(),
        }
    }

    /// Process a given command.
    fn process_command(&mut self, cmd: Command) {
        match cmd {
            Command::ResetServiceTable   => self.reset_service_table(),
            Command::ScanNetwork         => self.scan_network(),
            Command::PeriodicNetworkScan => self.periodic_network_scan(),
        }
    }

    /// Reset service table.
    fn reset_service_table(&mut self) {
        self.app_context.reset_service_table()
    }

    /// Trigger periodic netork scan.
    fn periodic_network_scan(&mut self) {
        let t = time::precise_time_s();

        if (self.last_nw_scan + NETWORK_SCAN_PERIOD) < t {
            self.scan_network();
        }
    }

    #[cfg(feature="discovery")]
    /// Trigger a network scan.
    fn scan_network(&mut self) {
        if !self.app_context.get_discovery() ||
            self.scanner.is_some() {
            return;
        }

        let app_context = self.app_context.clone();
        let cmd_sender  = self.cmd_sender.clone();

        let handle = thread::spawn(move || {
            network_scanner_thread(app_context);

            cmd_sender.send(Event::ScanCompleted)
                .unwrap();
        });

        self.scanner = Some(handle);
    }

    #[cfg(not(feature="discovery"))]
    /// Dummy network scan.
    fn scan_network(&mut self) {
    }

    /// Cleanup the scanner context.
    fn scan_completed(&mut self) {
        if let Some(handle) = self.scanner.take() {
            if let Err(_) = handle.join() {
                log_warn!(self.logger, "network scanner thread panicked");
            }
        }

        self.last_nw_scan = time::precise_time_s();
    }
}

/// Command handler. It implements the future trait and it's designed to be used in combination
/// with tokio event loop.
pub struct CommandHandler {
    handler: BoxFuture<(), ()>,
}

impl CommandHandler {
    /// Create a new command handler.
    fn new(
        app_context: ApplicationContext,
        rx: CommandReceiver,
        tx: CommandSender) -> CommandHandler {
        let mut context = CommandHandlerContext::new(app_context, tx);

        let handler = rx.for_each(move |event| {
            context.proces_event(event);
            Ok(())
        });

        CommandHandler {
            handler: handler.boxed(),
        }
    }
}

impl Future for CommandHandler {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.handler.poll()
    }
}

/// Create a new channel-handler pair.
pub fn new(app_context: ApplicationContext) -> (CommandChannel, CommandHandler) {
    let (tx, rx) = mpsc::unbounded();

    let rx = CommandHandler::new(app_context, rx, tx.clone());
    let tx = CommandChannel::new(tx);

    (tx, rx)
}

#[cfg(feature = "discovery")]
/// Run device discovery and update the service table.
fn network_scanner_thread(mut app_context: ApplicationContext) {
    let mut logger = app_context.get_logger();

    let rtsp_paths_file = app_context.get_rtsp_paths_file();
    let mjpeg_paths_file = app_context.get_mjpeg_paths_file();

    app_context.set_scanning(true);

    log_info!(logger, "looking for local services...");

    let result = utils::result_or_log(&mut logger, Severity::WARN,
        "network scanner error",
        discovery::scan_network(
            &rtsp_paths_file,
            &mjpeg_paths_file));

    if let Some(result) = result {
        let services = result.services()
            .map(|svc| svc.clone())
            .collect::<Vec<_>>();

        let count = services.len();

        app_context.update_service_table(services);
        app_context.set_scan_result(result);

        log_info!(logger, "{} services found, current service table: {}",
            count, app_context.get_service_table());
    }

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

#[cfg(feature="discovery")]
use std::thread;

use std::thread::JoinHandle;

use context::ApplicationContext;

#[cfg(feature = "discovery")]
use scanner::discovery;

#[cfg(feature="discovery")]
use utils;

use utils::logger::{BoxLogger, Logger};

#[cfg(feature="discovery")]
use utils::logger::Severity;

use futures::{Future, Poll, Stream};

use futures::sync::mpsc;
use futures::sync::mpsc::{UnboundedReceiver, UnboundedSender};

use time;

/// Network scan period.
const NETWORK_SCAN_PERIOD: f64 = 300.0;

/// Different command types that the command handler can receive.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum Command {
    ResetServiceTable,
    ScanNetwork,
    PeriodicNetworkScan,
}

/// Command handler event.
#[derive(Debug, Copy, Clone)]
enum Event {
    Command(Command),
    ScanCompleted,
}

type CommandReceiver = UnboundedReceiver<Event>;

type CommandSender = UnboundedSender<Event>;

/// A channel for sending commands across threads. The channel is cloneable and every copy of it
/// will send commands to the same handler.
#[derive(Clone)]
pub struct CommandChannel {
    tx: CommandSender,
}

impl CommandChannel {
    /// Create a new command channel.
    fn new(tx: CommandSender) -> CommandChannel {
        CommandChannel {
            tx: tx,
        }
    }

    /// Send a given command.
    pub fn send(&self, cmd: Command) {
        self.tx.unbounded_send(Event::Command(cmd))
            .expect("broken command channel");
    }
}

/// Command handler context.
struct CommandHandlerContext {
    app_context:  ApplicationContext,
    logger:       BoxLogger,
    cmd_sender:   CommandSender,
    last_nw_scan: f64,
    scanner:      Option<JoinHandle<()>>,
}

impl CommandHandlerContext {
    /// Create a new command handler context.
    fn new(app_context: ApplicationContext, cmd_sender: CommandSender) -> CommandHandlerContext {
        let logger = app_context.get_logger();

        let t = time::precise_time_s();

        CommandHandlerContext {
            app_context:  app_context,
            logger:       logger,
            cmd_sender:   cmd_sender,
            last_nw_scan: t - 300.0,
            scanner:      None,
        }
    }

    /// Process a given command handler event.
    fn proces_event(&mut self, event: Event) {
        match event {
            Event::Command(cmd)  => self.process_command(cmd),
            Event::ScanCompleted => self.scan_completed(),
        }
    }

    /// Process a given command.
    fn process_command(&mut self, cmd: Command) {
        match cmd {
            Command::ResetServiceTable   => self.reset_service_table(),
            Command::ScanNetwork         => self.scan_network(),
            Command::PeriodicNetworkScan => self.periodic_network_scan(),
        }
    }

    /// Reset service table.
    fn reset_service_table(&mut self) {
        self.app_context.reset_service_table()
    }

    /// Trigger periodic netork scan.
    fn periodic_network_scan(&mut self) {
        let t = time::precise_time_s();

        if (self.last_nw_scan + NETWORK_SCAN_PERIOD) < t {
            self.scan_network();
        }
    }

    #[cfg(feature="discovery")]
    /// Trigger a network scan.
    fn scan_network(&mut self) {
        if !self.app_context.get_discovery() ||
            self.scanner.is_some() {
            return;
        }

        let app_context = self.app_context.clone();
        let cmd_sender  = self.cmd_sender.clone();

        let handle = thread::spawn(move || {
            network_scanner_thread(app_context);

            cmd_sender.unbounded_send(Event::ScanCompleted)
                .unwrap();
        });

        self.scanner = Some(handle);
    }

    #[cfg(not(feature="discovery"))]
    /// Dummy network scan.
    fn scan_network(&mut self) {
    }

    /// Cleanup the scanner context.
    fn scan_completed(&mut self) {
        if let Some(handle) = self.scanner.take() {
            if let Err(_) = handle.join() {
                log_warn!(self.logger, "network scanner thread panicked");
            }
        }

        self.last_nw_scan = time::precise_time_s();
    }
}

/// Command handler. It implements the future trait and it's designed to be used in combination
/// with tokio event loop.
pub struct CommandHandler {
    handler: Box<Future<Item=(), Error=()>>,
}

impl CommandHandler {
    /// Create a new command handler.
    fn new(
        app_context: ApplicationContext,
        rx: CommandReceiver,
        tx: CommandSender) -> CommandHandler {
        let mut context = CommandHandlerContext::new(app_context, tx);

        let handler = rx.for_each(move |event| {
            context.proces_event(event);
            Ok(())
        });

        CommandHandler {
            handler: Box::new(handler),
        }
    }
}

impl Future for CommandHandler {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.handler.poll()
    }
}

/// Create a new channel-handler pair.
pub fn new(app_context: ApplicationContext) -> (CommandChannel, CommandHandler) {
    let (tx, rx) = mpsc::unbounded();

    let rx = CommandHandler::new(app_context, rx, tx.clone());
    let tx = CommandChannel::new(tx);

    (tx, rx)
}

#[cfg(feature = "discovery")]
/// Run device discovery and update the service table.
fn network_scanner_thread(mut app_context: ApplicationContext) {
    let mut logger = app_context.get_logger();

    let rtsp_paths_file = app_context.get_rtsp_paths_file();
    let mjpeg_paths_file = app_context.get_mjpeg_paths_file();

    app_context.set_scanning(true);

    log_info!(logger, "looking for local services...");

    let result = utils::result_or_log(&mut logger, Severity::WARN,
        "network scanner error",
        discovery::scan_network(
            &rtsp_paths_file,
            &mjpeg_paths_file));

    if let Some(result) = result {
        let services = result.services()
            .map(|svc| svc.clone())
            .collect::<Vec<_>>();

        let count = services.len();

        app_context.update_service_table(services);
        app_context.set_scan_result(result);

        log_info!(logger, "{} services found, current service table: {}",
            count, app_context.get_service_table());
    }

    app_context.set_scanning(false);
}
<|endoftext|>


<commit_before>extern crate brotli2;
extern crate flate2;
extern crate zopfli;

use structs::*;
use errors::*;
use std::fs::File;
use std::io::{Read, Write};

use std::path::Path;

impl CompressionFormat for CompressionAlgorithm {
    fn extension(&self) -> &'static str {
        match self {
            &CompressionAlgorithm::Brotli => "br",
            &CompressionAlgorithm::GZip => "gz",
            &CompressionAlgorithm::WebP => "webp",
            &CompressionAlgorithm::Zopfli => "gz",
        }
    }
}

impl FileCompressor for CompressionAlgorithm {
    fn compress(&self, src: &Path, dst: &Path, quality: Option<u8>) -> Result<()> {
        match self {
            &CompressionAlgorithm::GZip => gzip_compress(src, dst, quality),
            &CompressionAlgorithm::Brotli => brotli_compress(src, dst),
            &CompressionAlgorithm::WebP => webp_compress(src, dst, quality),
            &CompressionAlgorithm::Zopfli => zopfli_compress(src, dst),
            // _ => bail!("Compression algorithm not implemented!"),
        }
    }
}

fn gzip_compress(src_path: &Path, dst_path: &Path, quality: Option<u8>) -> Result<()> {
    let mut src = File::open(src_path)?;
    let dst = File::create(dst_path)?;

    let level = match quality {
        Some(0) => flate2::Compression::None,
        Some(1) => flate2::Compression::Fast,
        Some(9) => flate2::Compression::Best,
        _ => flate2::Compression::Default,
    };

    let mut encoder = flate2::write::GzEncoder::new(dst, level);
    let mut buf = [0u8; 1024];
    loop {
        let bytes_read = src.read(&mut buf).chain_err(|| "Error reading from source file!")?;
        match bytes_read {
            0 => break, //end of file
            l => encoder.write_all(&buf[0..l]).chain_err(|| "Fatal gzip encoder error!")?,
        };
    }

    Ok(())
}

fn brotli_compress(src_path: &Path, dst_path: &Path) -> Result<()> {
    let mut src = File::open(src_path)?;
    let dst = File::create(dst_path)?;

    let mut buf = [0u8; 1024];
    let mut encoder = brotli2::write::BrotliEncoder::new(dst, 6);
    loop {
        let bytes_read = src.read(&mut buf).chain_err(|| "Error reading from source file!")?;
        match bytes_read {
            0 => break, //end of file
            l => encoder.write_all(&buf[0..l]).chain_err(|| "Fatal gzip encoder error!")?,
        };
    }

    Ok(())
}

fn zopfli_compress(src_path: &Path, dst_path: &Path) -> Result<()> {
    let mut src = File::open(src_path)?;
    let mut dst = File::create(dst_path)?;

    let mut src_data = Vec::<u8>::new();
    src.read_to_end(&mut src_data)?;

    zopfli::compress(&zopfli::Options::default(), &zopfli::Format::Gzip, &mut src_data, &mut dst)?;

    Ok(())
}

fn webp_compress(src_path: &Path, dst_path: &Path, quality: Option<u8>) -> Result<()> {
    use std::process::Command;

    let output = Command::new("cwebp")
        .arg("-q")
        .arg(quality.unwrap_or(90).to_string())
        .arg(src_path.as_os_str())
        .arg("-o")
        .arg(dst_path.as_os_str())
        .output()
        .chain_err(|| "Error executing cwebp!")?;

    if !output.status.success() {
        bail!("Error compressing via webp: {:?}", output);
    }

    Ok(())
}
<commit_msg>--quality range validation for gzip compression<commit_after>extern crate brotli2;
extern crate flate2;
extern crate zopfli;

use structs::*;
use errors::*;
use std::fs::File;
use std::io::{Read, Write};

use std::path::Path;

impl CompressionFormat for CompressionAlgorithm {
    fn extension(&self) -> &'static str {
        match self {
            &CompressionAlgorithm::Brotli => "br",
            &CompressionAlgorithm::GZip => "gz",
            &CompressionAlgorithm::WebP => "webp",
            &CompressionAlgorithm::Zopfli => "gz",
        }
    }
}

impl FileCompressor for CompressionAlgorithm {
    fn compress(&self, src: &Path, dst: &Path, quality: Option<u8>) -> Result<()> {
        match self {
            &CompressionAlgorithm::GZip => gzip_compress(src, dst, quality),
            &CompressionAlgorithm::Brotli => brotli_compress(src, dst),
            &CompressionAlgorithm::WebP => webp_compress(src, dst, quality),
            &CompressionAlgorithm::Zopfli => zopfli_compress(src, dst),
            // _ => bail!("Compression algorithm not implemented!"),
        }
    }
}

fn gzip_compress(src_path: &Path, dst_path: &Path, quality: Option<u8>) -> Result<()> {
    let mut src = File::open(src_path)?;
    let dst = File::create(dst_path)?;

    let level = match quality {
        None => flate2::Compression::Default,
        Some(0) => flate2::Compression::None,
        Some(1) => flate2::Compression::Fast,
        Some(2...6) => flate2::Compression::Default,
        Some(3...9) => flate2::Compression::Best,
        _ => bail!("Invalid --quality parameter specified!"),
    };

    let mut encoder = flate2::write::GzEncoder::new(dst, level);
    let mut buf = [0u8; 1024];
    loop {
        let bytes_read = src.read(&mut buf).chain_err(|| "Error reading from source file!")?;
        match bytes_read {
            0 => break, //end of file
            l => encoder.write_all(&buf[0..l]).chain_err(|| "Fatal gzip encoder error!")?,
        };
    }

    Ok(())
}

fn brotli_compress(src_path: &Path, dst_path: &Path) -> Result<()> {
    let mut src = File::open(src_path)?;
    let dst = File::create(dst_path)?;

    let mut buf = [0u8; 1024];
    let mut encoder = brotli2::write::BrotliEncoder::new(dst, 6);
    loop {
        let bytes_read = src.read(&mut buf).chain_err(|| "Error reading from source file!")?;
        match bytes_read {
            0 => break, //end of file
            l => encoder.write_all(&buf[0..l]).chain_err(|| "Fatal gzip encoder error!")?,
        };
    }

    Ok(())
}

fn zopfli_compress(src_path: &Path, dst_path: &Path) -> Result<()> {
    let mut src = File::open(src_path)?;
    let mut dst = File::create(dst_path)?;

    let mut src_data = Vec::<u8>::new();
    src.read_to_end(&mut src_data)?;

    zopfli::compress(&zopfli::Options::default(), &zopfli::Format::Gzip, &mut src_data, &mut dst)?;

    Ok(())
}

fn webp_compress(src_path: &Path, dst_path: &Path, quality: Option<u8>) -> Result<()> {
    use std::process::Command;

    let output = Command::new("cwebp")
        .arg("-q")
        .arg(quality.unwrap_or(90).to_string())
        .arg(src_path.as_os_str())
        .arg("-o")
        .arg(dst_path.as_os_str())
        .output()
        .chain_err(|| "Error executing cwebp!")?;

    if !output.status.success() {
        bail!("Error compressing via webp: {:?}", output);
    }

    Ok(())
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use super::stream::Stream;

pub struct Database {
    tables: HashMap<String, Stream>
}

impl Database {
    pub fn new() -> Database {
        Database{
            tables: HashMap::new()
        }
    }

    fn create_table(&mut self, name: &str)  {
        let tmp = Stream::new();
        self.tables.insert(name.to_string(), tmp);
    }
}

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

    #[test]
    fn create_table() {
        let mut db = Database::new();
        db.create_table("Jon");
    }
}
<commit_msg>get mutable stream<commit_after>use std::collections::HashMap;
use super::stream::Stream;

pub struct Database {
    tables: HashMap<String, Stream>
}

impl Database {
    pub fn new() -> Database {
        Database{
            tables: HashMap::new()
        }
    }

    fn create_table(&mut self, name: &str)  {
        let tmp = Stream::new();
        self.tables.insert(name.to_string(), tmp);
    }

    fn get_table_mut(&mut self, name: &str) -> Option<&mut Stream> {
        self.tables.get_mut(name)
    }
}

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

    #[test]
    fn create_table() {
        let mut db = Database::new();
        db.create_table("Jon");
    }

}
<|endoftext|>


<commit_before>extern crate websocket;

use std;
use std::sync::{Arc, Mutex};
use std::sync::mpsc;
use std::thread;
use std::time::Duration;

use websocket::{Message, Sender, Receiver};
use websocket::message::Type;
use websocket::client::request::Url;

use hyper::header::{Cookie};

use rustc_serialize::json;
use rustc_serialize::json::Json;

use lila;

// making a move
// out {"t":"move","d":{"from":"e2","to":"e4","b":1}}
// in {t: "ack"}
// in {"v":1,"t":"move","d":{"uci":"e2e4","san":"e4","fen":"rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR","ply":1,"clock":{"white":172800,"black":172800},"dests":{"b8":"a6c6","g8":"f6h6","h7":"h6h5","d7":"d6d5","g7":"g6g5","c7":"c6c5","f7":"f6f5","b7":"b6b5","e7":"e6e5","a7":"a6a5"},"crazyhouse":{"pockets":[{},{}]}}}
// other player moved
// in {"v":2,"t":"move","d":{"uci":"g7g5","san":"g5","fen":"rnbqkbnr/pppppp1p/8/6p1/4P3/8/PPPP1PPP/RNBQKBNR","ply":2,"clock":{"white":172800,"black":172800},"dests":{"a2":"a3a4","g1":"f3h3e2","d1":"e2f3g4h5","e1":"e2","d2":"d3d4","b1":"a3c3","e4":"e5","f1":"e2d3c4b5a6","h2":"h3h4","b2":"b3b4","f2":"f3f4","c2":"c3c4","g2":"g3g4"},"crazyhouse":{"pockets":[{},{}]}}}

#[derive(RustcEncodable)]
struct PingPacket {
    t: String,
    v: i64
}

impl PingPacket {
    pub fn new(version: i64) -> PingPacket {
        PingPacket {
            t: "p".to_string(),
            v: version
        }
    }

    pub fn to_message(&self) -> Message<'static> {
        Message::text(json::encode(self).unwrap())
    }
}

#[allow(dead_code)]
#[derive(RustcEncodable)]
pub struct MovePacket {
    t: String,
    d: Dest,
}

#[allow(dead_code)]
#[derive(RustcEncodable)]
pub struct Dest {
    pub from: String,
    pub to: String,
    // pub: i8, // bool, but 0 or 1
}

pub fn connect(session: &lila::Session, base_url: String, sri: String, pov: Arc<Mutex<super::Pov>>) {

    let url;
    {
        let pov = pov.lock().unwrap();
        let version = match pov.player.version {
            Some(v) => v,
            None => 0
        };
        url = Url::parse(&format!("ws://{}{}?sri={}&version={}", base_url, pov.url.socket.clone(), sri, version)).unwrap();
    }

    // TODO: this unwrap fails when url is wrong, port for example
    let mut request = websocket::Client::connect(url).unwrap();
    request.headers.set(Cookie::from_cookie_jar(&session.cjar));

    let response = request.send().unwrap();

    match response.validate() {
        Ok(()) => (),
        Err(e) => {
            error!("reponse validation error: {:?}", e);
            debug!("reponse status: {:?}, headers: {:?}", response.status, response.headers);
            return;
        }
    };

    //println!("Successfully connected");

    let (mut sender, mut receiver) = response.begin().split();

    let (tx, rx) = mpsc::channel();

    let tx_1 = tx.clone();

    let _send_loop = thread::spawn(move || {
        loop {
            let message: Message = match rx.recv() {
                Ok(m) => m,
                Err(e) => {
                    error!("Send loop, could not read next message: {:?}", e);
                    return;
                }
            };
            trace!("Transmit raw: {:?}", message);
            debug!("Transmit msg: {}",
                   &std::str::from_utf8(&message.payload).unwrap());
            match message.opcode {
                Type::Close => {
                    let _ = sender.send_message(&message);
                    return;
                }
                _ => (),
            }
            match sender.send_message(&message) {
                Ok(()) => (),
                Err(e) => {
                    error!("Send loop: could not send message: {:?}", e);
                    let _ = sender.send_message(&Message::close());
                    return;
                }
            }
        }
    });

    let pov_1 = pov.clone();
    let _receive_loop = thread::spawn(move || {

        let handle = |obj: json::Object| {
            let mut pov = pov_1.lock().unwrap();
            match obj.get("v") {
                Some(&Json::U64(v)) => {
                    pov.player.version = Some(v as i64);
                },
                _ => ()
            }
            match obj.get("t").and_then(|t| t.as_string()) {
                Some("move") => {
                    let d = obj.get("d").unwrap().as_object().unwrap();
                    let fen = d.get("fen").unwrap().as_string().unwrap();
                    pov.game.fen = fen.to_string();
                    //cli::render_fen(fen);
                },
                Some("clock") => {
                    let d = obj.get("d").unwrap().as_object().unwrap();
                    match pov.clock.as_mut() {
                        Some(clock) => {
                            clock.white = d.get("white").unwrap().as_f64().unwrap();
                            clock.black = d.get("black").unwrap().as_f64().unwrap();
                        },
                        None => ()
                    };
                },
                Some("end") => {
                    let _ = tx_1.send(Message::close());
                    // exit
                },
                Some(ref t) => warn!("unhandled: {}, {:?}", t, obj),
                _ => ()
            }
        };

        for message in receiver.incoming_messages() {
            let message: Message = match message {
                Ok(m) => m,
                Err(e) => {
                    warn!("Receive loop error: {:?}", e);
                    let _ = tx_1.send(Message::close());
                    return;
                }
            };
            match message.opcode {
                Type::Close => {
                    let _ = tx_1.send(Message::close());
                    return;
                }
                Type::Ping => match tx_1.send(Message::pong(message.payload)) {
                    Ok(()) => (),
                    Err(e) => {
                        warn!("Could not send pong: {:?}", e);
                        return;
                    }
                },
                Type::Text => {
                    let str_payload = std::str::from_utf8(&message.payload).unwrap();
                    trace!("Received raw: {:?}", &message.payload);
                    trace!("Received str: {}", str_payload);
                    let json = Json::from_str(str_payload).unwrap();
                    debug!("Received obj: {:?}", json);
                    if json.is_object() {
                        let obj = json.as_object().unwrap();
                        match obj.get("t").and_then(|t| t.as_string()) {
                            Some("n") => { // pong
                                // track time between pings and these pongs for latency
                            },
                            Some("b") => { // batch
                                for item in obj.get("d").unwrap().as_array().unwrap().iter() {
                                   let obj = item.as_object().unwrap();
                                   handle(obj.clone());
                                }
                            }
                            _ => {
                                handle(obj.clone());
                            }
                        }
                    }
                }
                _ => debug!("Unhandled message: {:?}", message),
            }
        }
    });

    let tx_2 = tx.clone();

    thread::spawn(move || loop {
        thread::sleep(Duration::from_millis(2000));

        pov.lock().ok()
            .and_then(|p| p.player.version )
            .map(|v| PingPacket::new(v).to_message())
            .map(|p| tx_2.send(p).unwrap());
    });

    //let _ = send_loop.join();
    //let _ = receive_loop.join();
}

<commit_msg>sprinkle some types into socket handler<commit_after>extern crate websocket;

use std;
use std::sync::{Arc, Mutex};
use std::sync::mpsc;
use std::thread;
use std::time::Duration;
use std::collections::BTreeMap;

use websocket::{Message, Sender, Receiver};
use websocket::message::Type;
use websocket::client::request::Url;

use hyper::header::{Cookie};

use rustc_serialize::json;
use rustc_serialize::json::Json;
use rustc_serialize::Decodable;

use lila;
use game::Clock;

// making a move
// out {"t":"move","d":{"from":"e2","to":"e4","b":1}}
// in {t: "ack"}
// in {"v":1,"t":"move","d":{"uci":"e2e4","san":"e4","fen":"rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR","ply":1,"clock":{"white":172800,"black":172800},"dests":{"b8":"a6c6","g8":"f6h6","h7":"h6h5","d7":"d6d5","g7":"g6g5","c7":"c6c5","f7":"f6f5","b7":"b6b5","e7":"e6e5","a7":"a6a5"},"crazyhouse":{"pockets":[{},{}]}}}
// other player moved
// in {"v":2,"t":"move","d":{"uci":"g7g5","san":"g5","fen":"rnbqkbnr/pppppp1p/8/6p1/4P3/8/PPPP1PPP/RNBQKBNR","ply":2,"clock":{"white":172800,"black":172800},"dests":{"a2":"a3a4","g1":"f3h3e2","d1":"e2f3g4h5","e1":"e2","d2":"d3d4","b1":"a3c3","e4":"e5","f1":"e2d3c4b5a6","h2":"h3h4","b2":"b3b4","f2":"f3f4","c2":"c3c4","g2":"g3g4"},"crazyhouse":{"pockets":[{},{}]}}}

enum LilaMessage {
    Pong,
    Move(Move),
    Clock(Clock),
}

impl LilaMessage {
    fn decode(obj: &BTreeMap<String, json::Json>) -> Vec<LilaMessage> {
        fn decode<T: Decodable>(data: json::Json) -> T {
            let mut decoder = json::Decoder::new(data);
            Decodable::decode(&mut decoder).unwrap()
        }
        let mut list = Vec::new();
        let data = obj.get("d");
        match (obj.get("t").and_then(|t| t.as_string()), data) {
            // TODO: gone, crowd, end, tvSelect, challenges, drop,
            // following_enters, following_leaves, following_onlines,
            // following_playing, following_stopped_plaing,
            // message, analysisProgress, reload, and more
            (Some("n"), _) => {
                list.push(LilaMessage::Pong);
            },
            (Some("move"), Some(data)) => {
                list.push(LilaMessage::Move(decode(data.to_owned())));
            },
            (Some("clock"), Some(data)) => {
                list.push(LilaMessage::Clock(decode(data.to_owned())));
            },
            (Some("b"), Some(data)) => { // batch
                for item in data.as_array().unwrap().iter() {
                   let obj = item.as_object().unwrap();
                   list.extend(LilaMessage::decode(obj));
                }
            },
            (Some(ref t), d) => {
                warn!("unhandled: {}, {:?}", t, d);
            },
            _ => {
                warn!("unhandled: Missing type");
            }
        }
        list
    }

}

#[derive(RustcDecodable)]
struct Move {
    clock: Option<Clock>,
    fen: String,
}

#[derive(RustcEncodable)]
struct PingPacket {
    t: String,
    v: i64
}

impl PingPacket {
    pub fn new(version: i64) -> PingPacket {
        PingPacket {
            t: "p".to_string(),
            v: version
        }
    }

    pub fn to_message(&self) -> Message<'static> {
        Message::text(json::encode(self).unwrap())
    }
}

#[allow(dead_code)]
#[derive(RustcEncodable)]
pub struct MovePacket {
    t: String,
    d: Dest,
}

#[allow(dead_code)]
#[derive(RustcEncodable)]
pub struct Dest {
    pub from: String,
    pub to: String,
    // pub: i8, // bool, but 0 or 1
}

pub fn connect(session: &lila::Session, base_url: String, sri: String, pov: Arc<Mutex<super::Pov>>) {

    let url;
    {
        let pov = pov.lock().unwrap();
        let version = match pov.player.version {
            Some(v) => v,
            None => 0
        };
        url = Url::parse(&format!("ws://{}{}?sri={}&version={}", base_url, pov.url.socket.clone(), sri, version)).unwrap();
    }

    // TODO: this unwrap fails when url is wrong, port for example
    let mut request = websocket::Client::connect(url).unwrap();
    request.headers.set(Cookie::from_cookie_jar(&session.cjar));

    let response = request.send().unwrap();

    match response.validate() {
        Ok(()) => (),
        Err(e) => {
            error!("reponse validation error: {:?}", e);
            debug!("reponse status: {:?}, headers: {:?}", response.status, response.headers);
            return;
        }
    };

    //println!("Successfully connected");

    let (mut sender, mut receiver) = response.begin().split();

    let (tx, rx) = mpsc::channel();

    let tx_1 = tx.clone();

    let _send_loop = thread::spawn(move || {
        loop {
            let message: Message = match rx.recv() {
                Ok(m) => m,
                Err(e) => {
                    error!("Send loop, could not read next message: {:?}", e);
                    return;
                }
            };
            trace!("Transmit raw: {:?}", message);
            debug!("Transmit msg: {}",
                   &std::str::from_utf8(&message.payload).unwrap());
            match message.opcode {
                Type::Close => {
                    let _ = sender.send_message(&message);
                    return;
                }
                _ => (),
            }
            match sender.send_message(&message) {
                Ok(()) => (),
                Err(e) => {
                    error!("Send loop: could not send message: {:?}", e);
                    let _ = sender.send_message(&Message::close());
                    return;
                }
            }
        }
    });

    let pov_1 = pov.clone();
    let _receive_loop = thread::spawn(move || {

        let handle = |msg: LilaMessage, version: Option<u64>| {
            let mut pov = pov_1.lock().unwrap();
            if let Some(v) = version {
                pov.player.version = Some(v as i64);
            };
            match msg {
                LilaMessage::Pong => {
                    // TODO
                },
                LilaMessage::Move(m) => {
                    pov.game.fen = m.fen;
                    if let Some(c) = m.clock {
                        pov.clock = Some(c);
                    };
                },
                LilaMessage::Clock(c) => {
                    pov.clock = Some(c);
                },
                //LilaMessage::End => tx_1.send(Message::close()).unwrap(),
            };
        };

        for message in receiver.incoming_messages() {
            let message: Message = match message {
                Ok(m) => m,
                Err(e) => {
                    warn!("Receive loop error: {:?}", e);
                    let _ = tx_1.send(Message::close());
                    return;
                }
            };
            match message.opcode {
                Type::Close => {
                    let _ = tx_1.send(Message::close());
                    return;
                }
                Type::Ping => match tx_1.send(Message::pong(message.payload)) {
                    Ok(()) => (),
                    Err(e) => {
                        warn!("Could not send pong: {:?}", e);
                        return;
                    }
                },
                Type::Text => {
                    let str_payload = std::str::from_utf8(&message.payload).unwrap();
                    trace!("Received raw: {:?}", &message.payload);
                    trace!("Received str: {}", str_payload);
                    let json = Json::from_str(str_payload).unwrap();
                    debug!("Received obj: {:?}", json);
                    if json.is_object() {
                        let obj = json.as_object().unwrap();
                        let version = obj.get("v").and_then(|v| v.as_u64());
                        for m in LilaMessage::decode(&obj) {
                            handle(m, version);
                        };
                    }
                }
                _ => debug!("Unhandled message: {:?}", message),
            }
        }
    });

    let tx_2 = tx.clone();

    thread::spawn(move || loop {
        thread::sleep(Duration::from_millis(2000));

        pov.lock().ok()
            .and_then(|p| p.player.version )
            .map(|v| PingPacket::new(v).to_message())
            .map(|p| tx_2.send(p).unwrap());
    });

    //let _ = send_loop.join();
    //let _ = receive_loop.join();
}

<|endoftext|>