<commit_before>/// The validator service. Connects to a beacon node and signs blocks when required.
use crate::attester_service::{AttestationGrpcClient, AttesterService};
use crate::block_producer_service::{BeaconBlockGrpcClient, BlockProducerService};
use crate::config::Config as ValidatorConfig;
use crate::duties::{DutiesManager, DutiesManagerService, EpochDutiesMap};
use attester::test_utils::EpochMap;
use attester::{test_utils::LocalSigner as AttesterLocalSigner, Attester};
use block_proposer::{test_utils::LocalSigner as BlockProposerLocalSigner, BlockProducer};
use bls::Keypair;
use grpcio::{ChannelBuilder, EnvBuilder};
use protos::services::Empty;
use protos::services_grpc::{
    AttestationServiceClient, BeaconBlockServiceClient, BeaconNodeServiceClient,
    ValidatorServiceClient,
};
use slog::{debug, info, warn};
use slot_clock::{SlotClock, SystemTimeSlotClock};
use std::ops::Sub;
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};
use tokio::prelude::*;
use tokio::runtime::Builder;
use tokio::timer::Interval;
use tokio_timer::clock::Clock;
use types::{Epoch, Fork, Slot};

//TODO: This service should be simplified in the future. Can be made more steamlined.

/// The validator service. This is the main thread that executes and maintains validator
/// duties.
pub struct Service {
    /// The node we currently connected to.
    connected_node_version: String,
    /// The chain id we are processing on.
    chain_id: u16,
    /// The fork state we processing on.
    fork: Fork,
    /// The slot clock keeping track of time.
    slot_clock: Arc<SystemTimeSlotClock>,
    /// The current slot we are processing.
    current_slot: Slot,
    /// Seconds until the next slot. This is used for initializing the tokio timer interval.
    seconds_to_next_slot: Duration,
    // GRPC Clients
    /// The beacon block GRPC client.
    beacon_block_client: Arc<BeaconBlockServiceClient>,
    /// The validator GRPC client.
    validator_client: Arc<ValidatorServiceClient>,
    /// The attester GRPC client.
    attester_client: Arc<AttestationServiceClient>,
    /// The validator client logger.
    log: slog::Logger,
}

impl Service {
    ///  Initial connection to the beacon node to determine its properties.
    ///
    ///  This tries to connect to a beacon node. Once connected, it initialised the gRPC clients
    ///  and returns an instance of the service.
    fn initialize_service(config: &ValidatorConfig, log: slog::Logger) -> Self {
        // initialise the beacon node client to check for a connection

        let env = Arc::new(EnvBuilder::new().build());
        // Beacon node gRPC beacon node endpoints.
        let beacon_node_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(BeaconNodeServiceClient::new(ch))
        };

        // retrieve node information
        let node_info = loop {
            let info = match beacon_node_client.info(&Empty::new()) {
                Err(e) => {
                    warn!(log, "Could not connect to node. Error: {}", e);
                    info!(log, "Retrying in 5 seconds...");
                    std::thread::sleep(Duration::from_secs(5));
                    continue;
                }
                Ok(info) => break info,
            };
        };

        // build requisite objects to form Self
        let genesis_time = node_info.get_genesis_time();

        info!(log,"Beacon node connected"; "Node Version" => node_info.version.clone(), "Chain ID" => node_info.chain_id, "Genesis time" => genesis_time);

        let proto_fork = node_info.get_fork();
        let mut previous_version: [u8; 4] = [0; 4];
        let mut current_version: [u8; 4] = [0; 4];
        previous_version.copy_from_slice(&proto_fork.get_previous_version()[..4]);
        current_version.copy_from_slice(&proto_fork.get_current_version()[..4]);
        let fork = Fork {
            previous_version,
            current_version,
            epoch: Epoch::from(proto_fork.get_epoch()),
        };

        // build the validator slot clock
        let slot_clock = {
            let clock = SystemTimeSlotClock::new(genesis_time, config.spec.seconds_per_slot)
                .expect("Unable to instantiate SystemTimeSlotClock.");
            Arc::new(clock)
        };

        // initialize the RPC clients

        // Beacon node gRPC beacon block endpoints.
        let beacon_block_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(BeaconBlockServiceClient::new(ch))
        };

        // Beacon node gRPC validator endpoints.
        let validator_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(ValidatorServiceClient::new(ch))
        };

        //Beacon node gRPC attester endpoints.
        let attester_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(AttestationServiceClient::new(ch))
        };

        //TODO: Add error chain. Handle errors
        let current_slot = slot_clock.present_slot().unwrap().unwrap().sub(1);

        // calculate seconds to the next slot
        let seconds_to_next_slot = {
            let syslot_time = SystemTime::now();
            let duration_since_epoch = syslot_time.duration_since(SystemTime::UNIX_EPOCH).unwrap();
            let mut secs_to_slot = None;
            if let Some(duration_since_genesis) =
                duration_since_epoch.checked_sub(Duration::from_secs(genesis_time))
            {
                // seconds till next slot
                secs_to_slot = duration_since_genesis
                    .as_secs()
                    .checked_rem(config.spec.seconds_per_slot);
            }
            secs_to_slot.unwrap_or_else(|| 0)
        };

        Self {
            connected_node_version: node_info.version,
            chain_id: node_info.chain_id as u16,
            fork,
            slot_clock,
            current_slot,
            seconds_to_next_slot: Duration::from_secs(seconds_to_next_slot),
            beacon_block_client,
            validator_client,
            attester_client,
            log,
        }
    }

    /// Initialise the service then run the core thread.
    pub fn start(config: ValidatorConfig, log: slog::Logger) {
        // connect to the node and retrieve its properties and initialize the gRPC clients
        let service = Service::initialize_service(&config, log);

        // we have connected to a node and established its parameters. Spin up the core service

        // set up the validator service runtime
        let mut runtime = Builder::new()
            .clock(Clock::system())
            .name_prefix("validator-client-")
            .build()
            .unwrap();

        // set up the validator work interval - start at next slot and proceed every slot
        // TODO: Error chain handle errors.
        let interval = {
            // Set the interval to start at the next slot, and every slot after
            let slot_duration = Duration::from_secs(config.spec.seconds_per_slot);
            //TODO: Handle checked add correctly
            Interval::new(Instant::now() + service.seconds_to_next_slot, slot_duration)
        };

        // kick off core service

        // generate keypairs

        // TODO: keypairs are randomly generated; they should be loaded from a file or generated.
        // https://github.com/sigp/lighthouse/issues/160
        let keypairs = Arc::new(vec![Keypair::random()]);

        // build requisite objects to pass to core thread.
        let duties_map = Arc::new(EpochDutiesMap::new(config.spec.slots_per_epoch));
        let epoch_map_for_attester = Arc::new(EpochMap::new(config.spec.slots_per_epoch));
        let manager = DutiesManager {
            duties_map,
            pubkeys: keypairs.iter().map(|keypair| keypair.pk.clone()).collect(),
            spec: Arc::new(config.spec),
            slot_clock: service.slot_clock.clone(),
            beacon_node: service.validator_client.clone(),
        };

        runtime.block_on(interval.for_each(move |_| {
            // update duties
            debug!(service.log, "Processing new slot...");
            manager.poll();
            Ok(())
        }));
    }

    /*

    let duties_map = Arc::new(EpochDutiesMap::new(spec.slots_per_epoch));
    let epoch_map_for_attester = Arc::new(EpochMap::new(spec.slots_per_epoch));

    for keypair in keypairs {
        info!(self.log, "Starting validator services"; "validator" => keypair.pk.concatenated_hex_id());

        // Spawn a new thread to maintain the validator's `EpochDuties`.
        let duties_manager_thread = {
            let spec = spec.clone();
            let duties_map = duties_map.clone();
            let slot_clock = self.slot_clock.clone();
            let log = self.log.clone();
            let beacon_node = self.validator_client.clone();
            let pubkey = keypair.pk.clone();
            thread::spawn(move || {
                let manager = DutiesManager {
                    duties_map,
                    pubkey,
                    spec,
                    slot_clock,
                    beacon_node,
                };
                let mut duties_manager_service = DutiesManagerService {
                    manager,
                    poll_interval_millis,
                    log,
                };

                duties_manager_service.run();
            })
        };

        // Spawn a new thread to perform block production for the validator.
        let producer_thread = {
            let spec = spec.clone();
            let signer = Arc::new(BlockProposerLocalSigner::new(keypair.clone()));
            let duties_map = duties_map.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let client = Arc::new(BeaconBlockGrpcClient::new(beacon_block_grpc_client.clone()));
            thread::spawn(move || {
                let block_producer =
                    BlockProducer::new(spec, duties_map, slot_clock, client, signer);
                let mut block_producer_service = BlockProducerService {
                    block_producer,
                    poll_interval_millis,
                    log,
                };

                block_producer_service.run();
            })
        };

        // Spawn a new thread for attestation for the validator.
        let attester_thread = {
            let signer = Arc::new(AttesterLocalSigner::new(keypair.clone()));
            let epoch_map = epoch_map_for_attester.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let client = Arc::new(AttestationGrpcClient::new(attester_grpc_client.clone()));
            thread::spawn(move || {
                let attester = Attester::new(epoch_map, slot_clock, client, signer);
                let mut attester_service = AttesterService {
                    attester,
                    poll_interval_millis,
                    log,
                };

                attester_service.run();
            })
        };

        threads.push((duties_manager_thread, producer_thread, attester_thread));
    }

    // Naively wait for all the threads to complete.
    for tuple in threads {
        let (manager, producer, attester) = tuple;
        let _ = producer.join();
        let _ = manager.join();
        let _ = attester.join();
    }
    */
}
<commit_msg>Adds microsecond duration to validator client<commit_after>/// The validator service. Connects to a beacon node and signs blocks when required.
use crate::attester_service::{AttestationGrpcClient, AttesterService};
use crate::block_producer_service::{BeaconBlockGrpcClient, BlockProducerService};
use crate::config::Config as ValidatorConfig;
use crate::duties::{DutiesManager, DutiesManagerService, EpochDutiesMap};
use attester::test_utils::EpochMap;
use attester::{test_utils::LocalSigner as AttesterLocalSigner, Attester};
use block_proposer::{test_utils::LocalSigner as BlockProposerLocalSigner, BlockProducer};
use bls::Keypair;
use grpcio::{ChannelBuilder, EnvBuilder};
use protos::services::Empty;
use protos::services_grpc::{
    AttestationServiceClient, BeaconBlockServiceClient, BeaconNodeServiceClient,
    ValidatorServiceClient,
};
use slog::{debug, info, warn};
use slot_clock::{SlotClock, SystemTimeSlotClock};
use std::ops::Sub;
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};
use tokio::prelude::*;
use tokio::runtime::Builder;
use tokio::timer::Interval;
use tokio_timer::clock::Clock;
use types::{Epoch, Fork, Slot};

//TODO: This service should be simplified in the future. Can be made more steamlined.

/// The validator service. This is the main thread that executes and maintains validator
/// duties.
pub struct Service {
    /// The node we currently connected to.
    connected_node_version: String,
    /// The chain id we are processing on.
    chain_id: u16,
    /// The fork state we processing on.
    fork: Fork,
    /// The slot clock keeping track of time.
    slot_clock: Arc<SystemTimeSlotClock>,
    /// The current slot we are processing.
    current_slot: Slot,
    /// Micro seconds until the next slot. This is used for initializing the tokio timer interval.
    micros_to_next_slot: Duration,
    // GRPC Clients
    /// The beacon block GRPC client.
    beacon_block_client: Arc<BeaconBlockServiceClient>,
    /// The validator GRPC client.
    validator_client: Arc<ValidatorServiceClient>,
    /// The attester GRPC client.
    attester_client: Arc<AttestationServiceClient>,
    /// The validator client logger.
    log: slog::Logger,
}

impl Service {
    ///  Initial connection to the beacon node to determine its properties.
    ///
    ///  This tries to connect to a beacon node. Once connected, it initialised the gRPC clients
    ///  and returns an instance of the service.
    fn initialize_service(config: &ValidatorConfig, log: slog::Logger) -> Self {
        // initialise the beacon node client to check for a connection

        let env = Arc::new(EnvBuilder::new().build());
        // Beacon node gRPC beacon node endpoints.
        let beacon_node_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(BeaconNodeServiceClient::new(ch))
        };

        // retrieve node information
        let node_info = loop {
            let info = match beacon_node_client.info(&Empty::new()) {
                Err(e) => {
                    warn!(log, "Could not connect to node. Error: {}", e);
                    info!(log, "Retrying in 5 seconds...");
                    std::thread::sleep(Duration::from_secs(5));
                    continue;
                }
                Ok(info) => break info,
            };
        };

        // build requisite objects to form Self
        let genesis_time = node_info.get_genesis_time();
        let genesis_time = 1_549_935_547;

        info!(log,"Beacon node connected"; "Node Version" => node_info.version.clone(), "Chain ID" => node_info.chain_id, "Genesis time" => genesis_time);

        let proto_fork = node_info.get_fork();
        let mut previous_version: [u8; 4] = [0; 4];
        let mut current_version: [u8; 4] = [0; 4];
        previous_version.copy_from_slice(&proto_fork.get_previous_version()[..4]);
        current_version.copy_from_slice(&proto_fork.get_current_version()[..4]);
        let fork = Fork {
            previous_version,
            current_version,
            epoch: Epoch::from(proto_fork.get_epoch()),
        };

        // build the validator slot clock
        let slot_clock = {
            let clock = SystemTimeSlotClock::new(genesis_time, config.spec.seconds_per_slot)
                .expect("Unable to instantiate SystemTimeSlotClock.");
            Arc::new(clock)
        };

        // initialize the RPC clients

        // Beacon node gRPC beacon block endpoints.
        let beacon_block_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(BeaconBlockServiceClient::new(ch))
        };

        // Beacon node gRPC validator endpoints.
        let validator_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(ValidatorServiceClient::new(ch))
        };

        //Beacon node gRPC attester endpoints.
        let attester_client = {
            let ch = ChannelBuilder::new(env.clone()).connect(&config.server);
            Arc::new(AttestationServiceClient::new(ch))
        };

        //TODO: Add error chain. Handle errors
        let current_slot = slot_clock.present_slot().unwrap().unwrap().sub(1);

        // calculate seconds to the next slot
        let micros_to_next_slot = {
            let syslot_time = SystemTime::now();
            let duration_since_epoch = syslot_time.duration_since(SystemTime::UNIX_EPOCH).unwrap();
            debug!(log, "Duration since unix epoch {:?}", duration_since_epoch);
            let mut micros_to_slot = None;
            if let Some(duration_since_genesis) =
                duration_since_epoch.checked_sub(Duration::from_secs(genesis_time))
            {
                // seconds till next slot
                debug!(log, "Genesis Time {:?}", genesis_time);
                debug!(log, "Duration since genesis {:?}", duration_since_genesis);
                micros_to_slot = duration_since_genesis
                    .as_secs()
                    .checked_rem(config.spec.seconds_per_slot);
            }
            micros_to_slot.unwrap_or_else(|| 0)
            /*
            let duration_to_slot = duration_since_genesis
                .checked_sub(Duration::from(
                    duration_since_genesis
                        .checked_div(config.spec.seconds_per_slot as u64)
                        .unwrap()
                        .as_secs()
                        .checked_mul(config.spec.seconds_per_slot)
                        .unwrap(),
                ))
                .unwrap();
                */
        };

        info!(log, ""; "Micro Seconds to next slot"=>micros_to_next_slot);

        Self {
            connected_node_version: node_info.version,
            chain_id: node_info.chain_id as u16,
            fork,
            slot_clock,
            current_slot,
            micros_to_next_slot: Duration::from_micros(micros_to_next_slot),
            beacon_block_client,
            validator_client,
            attester_client,
            log,
        }
    }

    /// Initialise the service then run the core thread.
    pub fn start(config: ValidatorConfig, log: slog::Logger) {
        // connect to the node and retrieve its properties and initialize the gRPC clients
        let service = Service::initialize_service(&config, log);

        // we have connected to a node and established its parameters. Spin up the core service

        // set up the validator service runtime
        let mut runtime = Builder::new()
            .clock(Clock::system())
            .name_prefix("validator-client-")
            .build()
            .unwrap();

        // set up the validator work interval - start at next slot and proceed every slot
        // TODO: Error chain handle errors.
        let interval = {
            // Set the interval to start at the next slot, and every slot after
            let slot_duration = Duration::from_secs(config.spec.seconds_per_slot);
            //TODO: Handle checked add correctly
            Interval::new(Instant::now() + service.micros_to_next_slot, slot_duration)
        };

        // kick off core service

        // generate keypairs

        // TODO: keypairs are randomly generated; they should be loaded from a file or generated.
        // https://github.com/sigp/lighthouse/issues/160
        let keypairs = Arc::new(vec![Keypair::random()]);

        // build requisite objects to pass to core thread.
        let duties_map = Arc::new(EpochDutiesMap::new(config.spec.slots_per_epoch));
        let epoch_map_for_attester = Arc::new(EpochMap::new(config.spec.slots_per_epoch));
        let manager = DutiesManager {
            duties_map,
            pubkeys: keypairs.iter().map(|keypair| keypair.pk.clone()).collect(),
            spec: Arc::new(config.spec),
            slot_clock: service.slot_clock.clone(),
            beacon_node: service.validator_client.clone(),
        };

        runtime.block_on(interval.for_each(move |_| {
            // update duties
            debug!(
                service.log,
                "Processing slot: {}",
                service.slot_clock.present_slot().unwrap().unwrap().as_u64()
            );
            manager.poll();
            Ok(())
        }));
    }

    /*

    let duties_map = Arc::new(EpochDutiesMap::new(spec.slots_per_epoch));
    let epoch_map_for_attester = Arc::new(EpochMap::new(spec.slots_per_epoch));

    for keypair in keypairs {
        info!(self.log, "Starting validator services"; "validator" => keypair.pk.concatenated_hex_id());

        // Spawn a new thread to maintain the validator's `EpochDuties`.
        let duties_manager_thread = {
            let spec = spec.clone();
            let duties_map = duties_map.clone();
            let slot_clock = self.slot_clock.clone();
            let log = self.log.clone();
            let beacon_node = self.validator_client.clone();
            let pubkey = keypair.pk.clone();
            thread::spawn(move || {
                let manager = DutiesManager {
                    duties_map,
                    pubkey,
                    spec,
                    slot_clock,
                    beacon_node,
                };
                let mut duties_manager_service = DutiesManagerService {
                    manager,
                    poll_interval_millis,
                    log,
                };

                duties_manager_service.run();
            })
        };

        // Spawn a new thread to perform block production for the validator.
        let producer_thread = {
            let spec = spec.clone();
            let signer = Arc::new(BlockProposerLocalSigner::new(keypair.clone()));
            let duties_map = duties_map.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let client = Arc::new(BeaconBlockGrpcClient::new(beacon_block_grpc_client.clone()));
            thread::spawn(move || {
                let block_producer =
                    BlockProducer::new(spec, duties_map, slot_clock, client, signer);
                let mut block_producer_service = BlockProducerService {
                    block_producer,
                    poll_interval_millis,
                    log,
                };

                block_producer_service.run();
            })
        };

        // Spawn a new thread for attestation for the validator.
        let attester_thread = {
            let signer = Arc::new(AttesterLocalSigner::new(keypair.clone()));
            let epoch_map = epoch_map_for_attester.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let client = Arc::new(AttestationGrpcClient::new(attester_grpc_client.clone()));
            thread::spawn(move || {
                let attester = Attester::new(epoch_map, slot_clock, client, signer);
                let mut attester_service = AttesterService {
                    attester,
                    poll_interval_millis,
                    log,
                };

                attester_service.run();
            })
        };

        threads.push((duties_manager_thread, producer_thread, attester_thread));
    }

    // Naively wait for all the threads to complete.
    for tuple in threads {
        let (manager, producer, attester) = tuple;
        let _ = producer.join();
        let _ = manager.join();
        let _ = attester.join();
    }
    */
}
<|endoftext|>


<commit_before>use coords::Coord;

pub struct SystemBox<T> {
    pub dimensions: Coord,
    pub coords: Vec<T>
}

pub type Grid = SystemBox<Coord>;

/// Return a hexagonal grid of input size and base length.
pub fn hexagonal_grid(size_x: f64, size_y: f64, base_length: f64, z0: f64)
        -> Grid {
    // Calculate the box dimension and number of base vector replications
    let spacing = get_hexagonal_spacing(base_length);
    let (nx, ny) = get_num_replications(size_x, size_y, spacing);

    Grid {
        dimensions: get_system_dimensions(nx, ny, spacing, z0),
        coords: gen_hexagonal_grid(nx, ny, base_length, spacing)
    }
}

fn get_hexagonal_spacing(base_length: f64) -> Coord {
    let dx = base_length*f64::sqrt(3.0)/2.0;
    let dy = base_length*0.5;

    // The spacing to the next set of four points is two times the
    // move to the first point along x and adding another bond spacing
    // to the last point along y.
    Coord::new(2.0*dx, 2.0*dy + 2.0*base_length, 0.0)
}

fn get_num_replications(size_x: f64, size_y: f64, spacing: Coord) -> (u64, u64) {
    (f64::round(size_x/spacing.x) as u64, f64::round(size_y/spacing.y) as u64)
}

fn get_system_dimensions(nx: u64, ny: u64, spacing: Coord, z0: f64) -> Coord {
    Coord::new((nx as f64)*spacing.x, (ny as f64)*spacing.y, z0)
}

/// Generate a hexagonal grid of nx*ny base vectors. The distance
/// between each grid point is the input base length.
///
/// Each base vector consists of four grid points since this
/// is the amount required to create a periodically replicating
/// hexagonal grid.
fn gen_hexagonal_grid(nx: u64, ny: u64, base_length: f64, spacing: Coord) -> Vec<Coord> {
    // Starting at (0.0, 0.0) we construct these four base points
    // in the order of going counter-clockwise with an angle of 30 degrees,
    // then one full bond spacing above that, and finally counter-clockwise
    // with 150 degrees.
    let dx = base_length*f64::sqrt(3.0)/2.0;
    let dy = base_length*0.5;

    let base_coords = vec![
        Coord::new(0.0, 0.0,    0.0),
        Coord::new(dx,  dy,     0.0),
        Coord::new(dx,  dy + base_length, 0.0),
        Coord::new(0.0, 2.0*dy + base_length, 0.0)
    ];

    let mut grid_points = Vec::new();
    for i in 0..nx {
        for j in 0..ny {
            for coord in &base_coords {
                grid_points.push(coord.at_index(i, j, 0, &spacing));
            }
        }
    }

    grid_points
}

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

    #[test]
    fn gen_a_small_hexagonal_grid() {
        // A hexagonal grid must be replicated from four base points
        // which means that a 2-by-2 input replications should result
        // in 16 grid points. Assert that the final four are
        // correctly placed.
        let spacing = get_hexagonal_spacing(1.0);
        let grid = gen_hexagonal_grid(2, 2, 1.0, spacing);

        // We generate the base grid points used in the creation
        // of the grid. The angle from the first point to the second
        // is 30 degrees which gives us the difference along x and y
        // since cos(30) = sqrt(3)/2 and sin(30) = 0.5
        let dx = f64::sqrt(3.0)/2.0;
        let dy = 0.5;
        let base = vec![
            Coord::new(0.0, 0.0,          0.0),
            Coord::new(dx,  dy,           0.0),
            Coord::new(dx,  dy + 1.0,     0.0),
            Coord::new(0.0, 2.0*dy + 1.0, 0.0)
        ];
        assert_eq!(base, &grid[0..4]);

        let spacing = Coord::new(2.0*dx, 2.0+2.0*dy, 0.0);
        let expect_last: Vec<Coord> = base.iter().map(|c| c.at_index(1, 1, 0, &spacing)).collect();
        assert_eq!(expect_last, &grid[12..]);
    }
}
<commit_msg>Write a more general Lattice framework for grid construction<commit_after>use coords::Coord;

pub struct SystemBox<T> {
    pub dimensions: Coord,
    pub coords: Vec<T>
}

pub type Grid = SystemBox<Coord>;

enum CrystalSystem {
    Hexagonal { length: f64 },
    Triclinic { a: f64, b: f64, gamma: f64 }
}

struct Crystal {
    a: f64,
    b: f64,
    gamma: f64,
}

struct Spacing ( f64, f64, f64);

impl Crystal {
    fn from_system(input: CrystalSystem) -> Crystal {
        let pi = ::std::f64::consts::PI;

        match input {
            CrystalSystem::Hexagonal { length } => Crystal {
                a: length,
                b: length,
                gamma: 2.0*pi/3.0
            },
            CrystalSystem::Triclinic { a, b, gamma } => Crystal {
                a: a,
                b: b,
                gamma: gamma
            }
        }
    }

    fn spacing(&self) -> Spacing {
        let dx = self.a;
        let dy = self.b * self.gamma.sin();
        let dx_per_row = self.b * self.gamma.cos();

        Spacing(dx, dy, dx_per_row)
    }
}

struct Lattice {
    box_size: Coord,
    coords: Vec<Coord>,
}

impl Lattice {
    fn new(crystal: &Crystal, nx: u64, ny: u64) -> Lattice {
        let Spacing(dx, dy, dx_per_row) = crystal.spacing();

        let box_size = Coord { x: (nx as f64)*dx, y: (ny as f64)*dy, z: 0.0 };
        let coords = (0..ny)
            .flat_map(|row| {
                (0..nx).map(move |col| Coord {
                    x: (col as f64)*dx + (row as f64)*dx_per_row,
                    y: (row as f64)*dy,
                    z: 0.0,
                })
            })
            .collect();

        Lattice { box_size: box_size, coords: coords }
    }

    fn from_size(crystal: &Crystal, size_x: f64, size_y: f64) -> Lattice {
        let Spacing(dx, dy, _) = crystal.spacing();
        let (nx, ny) = ((size_x/dx).round() as u64, (size_y/dy).round() as u64);

        Lattice::new(&crystal, nx, ny)
    }
}

/// Return a hexagonal grid of input size and base length.
pub fn hexagonal_grid(size_x: f64, size_y: f64, base_length: f64, z0: f64)
        -> Grid {
    // Calculate the box dimension and number of base vector replications
    let spacing = get_hexagonal_spacing(base_length);
    let (nx, ny) = get_num_replications(size_x, size_y, spacing);

    Grid {
        dimensions: get_system_dimensions(nx, ny, spacing, z0),
        coords: gen_hexagonal_grid(nx, ny, base_length, spacing)
    }
}

fn get_hexagonal_spacing(base_length: f64) -> Coord {
    let dx = base_length*f64::sqrt(3.0)/2.0;
    let dy = base_length*0.5;

    // The spacing to the next set of four points is two times the
    // move to the first point along x and adding another bond spacing
    // to the last point along y.
    Coord::new(2.0*dx, 2.0*dy + 2.0*base_length, 0.0)
}

fn get_num_replications(size_x: f64, size_y: f64, spacing: Coord) -> (u64, u64) {
    (f64::round(size_x/spacing.x) as u64, f64::round(size_y/spacing.y) as u64)
}

fn get_system_dimensions(nx: u64, ny: u64, spacing: Coord, z0: f64) -> Coord {
    Coord::new((nx as f64)*spacing.x, (ny as f64)*spacing.y, z0)
}

/// Generate a hexagonal grid of nx*ny base vectors. The distance
/// between each grid point is the input base length.
///
/// Each base vector consists of four grid points since this
/// is the amount required to create a periodically replicating
/// hexagonal grid.
fn gen_hexagonal_grid(nx: u64, ny: u64, base_length: f64, spacing: Coord) -> Vec<Coord> {
    // Starting at (0.0, 0.0) we construct these four base points
    // in the order of going counter-clockwise with an angle of 30 degrees,
    // then one full bond spacing above that, and finally counter-clockwise
    // with 150 degrees.
    let dx = base_length*f64::sqrt(3.0)/2.0;
    let dy = base_length*0.5;

    let base_coords = vec![
        Coord::new(0.0, 0.0,    0.0),
        Coord::new(dx,  dy,     0.0),
        Coord::new(dx,  dy + base_length, 0.0),
        Coord::new(0.0, 2.0*dy + base_length, 0.0)
    ];

    let mut grid_points = Vec::new();
    for i in 0..nx {
        for j in 0..ny {
            for coord in &base_coords {
                grid_points.push(coord.at_index(i, j, 0, &spacing));
            }
        }
    }

    grid_points
}

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

    #[test]
    fn hexagonal_crystal() {
        let system = CrystalSystem::Hexagonal { length: 1.0 };
        let crystal = Crystal::from_system(system);
        assert_eq!(1.0, crystal.a);
        assert_eq!(1.0, crystal.b);
        assert_eq!(2.0*f64::consts::PI/3.0, crystal.gamma);
    }

    #[test]
    fn triclinic_crystal() {
        let system = CrystalSystem::Triclinic { a: 1.0, b: 2.0, gamma: 3.0 };
        let crystal = Crystal::from_system(system);
        assert_eq!(1.0, crystal.a);
        assert_eq!(2.0, crystal.b);
        assert_eq!(3.0, crystal.gamma);
    }

    #[test]
    fn triclinic_lattice() {
        let dx = 1.0;
        let angle = f64::consts::PI/3.0; // 60 degrees
        let crystal = Crystal { a: dx, b: dx, gamma: angle };
        let lattice = Lattice::new(&crystal, 3, 2);

        // Calculate shifts for x and y when shifting along y
        let dy = dx*f64::sin(angle);
        let dx_per_y = dx*f64::cos(angle);

        // Check the dimensions
        assert_eq!(Coord { x: 3.0*dx, y: 2.0*dy, z: 0.0 }, lattice.box_size);

        // ... and the coordinates
        let mut iter = lattice.coords.iter();
        assert_eq!(Some(&Coord { x: 0.0,               y: 0.0, z: 0.0 }), iter.next());
        assert_eq!(Some(&Coord { x: dx,                y: 0.0, z: 0.0 }), iter.next());
        assert_eq!(Some(&Coord { x: 2.0*dx,            y: 0.0, z: 0.0 }), iter.next());
        assert_eq!(Some(&Coord { x: dx_per_y,          y: dy,  z: 0.0 }), iter.next());
        assert_eq!(Some(&Coord { x: dx_per_y + dx,     y: dy,  z: 0.0 }), iter.next());
        assert_eq!(Some(&Coord { x: dx_per_y + 2.0*dx, y: dy,  z: 0.0 }), iter.next());
        assert_eq!(None, iter.next());
    }

    #[test]
    fn lattice_from_size() {
        // This should result in a 2-by-2 lattice
        let crystal = Crystal { a: 1.0, b: 0.5, gamma: f64::consts::PI/2.0 };
        let lattice = Lattice::from_size(&crystal, 2.1, 0.9);
        let expected = Lattice::new(&crystal, 2, 2);

        assert_eq!(expected.coords, lattice.coords);
        assert_eq!(expected.box_size, lattice.box_size);
    }

    #[test]
    fn crystal_spacing() {
        let crystal = Crystal { a: 1.0, b: 3.0, gamma: f64::consts::PI/3.0 };
        let Spacing(dx, dy, dx_per_row) = crystal.spacing();

        assert_eq!(1.0, dx);
        assert_eq!(3.0*f64::sqrt(3.0)/2.0, dy);
        assert!((1.5 - dx_per_row).abs() < 1e-6);
    }
}
<|endoftext|>


<commit_before>use rust_base58::ToBase58;
use self::HashFnType::*;

use openssl::crypto::hash as openssl_hash;

// https://github.com/jbenet/multihash
pub enum HashFnType {
    SHA1 = 0x11,
    SHA2_256 = 0x12,
    SHA2_512 = 0x13,
    SHA3 = 0x14,
    BLAKE2B = 0x40,
    BLAKE2S = 0x41,
}

impl HashFnType {
    pub fn from_u8(x: u8) -> Option<HashFnType> {
        match x {
            0x11 => Some(SHA1),
            0x12 => Some(SHA2_256),
            0x13 => Some(SHA2_512),
            0x14 => Some(SHA3),
            0x40 => Some(BLAKE2B),
            0x41 => Some(BLAKE2S),
            _ => None,
        }
    }
}

// Takes bytes and a HashFnType, returns hash of the bytes
pub fn multihash<'a>(data: &'a[u8], hash_type: HashFnType) -> VecMultiHash {
    let mut openssl_type: openssl_hash::Type = match hash_type {
        SHA1 => openssl_hash::Type::SHA1,
        SHA2_256 => openssl_hash::Type::SHA256,
        SHA2_512 => openssl_hash::Type::SHA512,
        _ => panic!("That hash function is not yet implemented. Sorry"),
    };

    let mut hashed = openssl_hash::hash(openssl_type, data);

    VecMultiHash::encode(&hashed[..], hash_type as u8).unwrap()
}

struct HashFnTypeData {
    name: &'static str,
    default_len: u8,
}

fn hashfn_data(hft: &HashFnType) -> HashFnTypeData {
    match *hft {
        SHA1 =>
            HashFnTypeData { name: "sha1", default_len: 20 },
        SHA2_256 =>
            HashFnTypeData { name: "sha2-256", default_len: 32 },
        SHA2_512 =>
            HashFnTypeData { name: "sha2-512", default_len: 64 },
        SHA3 =>
            HashFnTypeData { name: "sha3", default_len: 64 },
        BLAKE2B =>
            HashFnTypeData { name: "blake2b", default_len: 64 },
        BLAKE2S =>
            HashFnTypeData { name: "blake2s", default_len: 32 },
    }
}

pub struct VecMultiHash {
    vec: Vec<u8>,
}

impl VecMultiHash {
    pub fn new() -> VecMultiHash {
        VecMultiHash { vec: Vec::new() }
    }

    // For constructing a VecMultiHash directly from a vector.
    // Function assumes that first byte is hash function code,
    // second is digest length. It's up to you.
    pub fn from_vec(v: Vec<u8>) -> VecMultiHash {
        VecMultiHash { vec: v }
    }

    pub fn encode<'a>(data: &'a [u8], code: u8) -> Result<VecMultiHash, EncodeError> {
        match HashFnType::from_u8(code) {
            None => return Err(EncodeError::UnknownCode(code)),
            Some(_) => {},
        }

        let size = data.len();
        if size > 127 {
            return Err(EncodeError::NotSupported(size));
        }

        let mut v = Vec::with_capacity(size + 2);
        v.push(code);
        v.push(size as u8);
        v.push_all(data);
        Ok(VecMultiHash::from_vec(v))
    }

    pub fn decode<'a>(&'a self) -> Result<DecodedMultiHash<'a>, DecodeError> {
        let mhlen = self.vec.len();
        if mhlen < 3 {
            return Err(DecodeError::TooShort);
        } else if mhlen > 129 {
            return Err(DecodeError::TooLong);
        } else {
            let digest_len = (mhlen - 2) as u8;
            if digest_len != self.vec[1] {
                return Err(DecodeError::InvalidDigestLength(self.vec[1], digest_len));
            }
        }

        let fn_code = match HashFnType::from_u8(self.vec[0]) {
            None => return Err(DecodeError::UnknownCode(self.vec[0])),
            Some(code) => code,
        };

        let hash_fn = hashfn_data(&fn_code);

        let decoded = DecodedMultiHash {
            code: fn_code,
            name: hash_fn.name,
            length: self.vec[1],
            digest: &self.vec[2..],
        };
        Ok(decoded)

    }

    pub fn vec(self) -> Vec<u8> {
        self.vec
    }

    pub fn to_base58_string(&self) -> String {
        self.vec.to_base58()
    }
}

pub struct DecodedMultiHash<'a> {
    code: HashFnType,
    name: &'static str,
    length: u8,
    digest: &'a [u8],
}

#[derive(Debug)]
pub enum DecodeError {
    UnknownCode(u8),
    TooShort,
    TooLong,
    InvalidDigestLength(u8, u8), // (stated, actual)
}

#[derive(Debug)]
pub enum EncodeError {
    UnknownCode(u8),
    NotSupported(usize),
}

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

    #[test]
    fn test_hash() {
        multihash::hash(b"ABC", multihash::HashFnType::SHA1);
    }

}
<commit_msg>Improve description of `multihash`<commit_after>use rust_base58::ToBase58;
use self::HashFnType::*;

use openssl::crypto::hash as openssl_hash;

// https://github.com/jbenet/multihash
pub enum HashFnType {
    SHA1 = 0x11,
    SHA2_256 = 0x12,
    SHA2_512 = 0x13,
    SHA3 = 0x14,
    BLAKE2B = 0x40,
    BLAKE2S = 0x41,
}

impl HashFnType {
    pub fn from_u8(x: u8) -> Option<HashFnType> {
        match x {
            0x11 => Some(SHA1),
            0x12 => Some(SHA2_256),
            0x13 => Some(SHA2_512),
            0x14 => Some(SHA3),
            0x40 => Some(BLAKE2B),
            0x41 => Some(BLAKE2S),
            _ => None,
        }
    }
}

// Takes bytes and a HashFnType, returns multihash of the bytes
pub fn multihash<'a>(data: &'a[u8], hash_type: HashFnType) -> VecMultiHash {
    let mut openssl_type: openssl_hash::Type = match hash_type {
        SHA1 => openssl_hash::Type::SHA1,
        SHA2_256 => openssl_hash::Type::SHA256,
        SHA2_512 => openssl_hash::Type::SHA512,
        _ => panic!("That hash function is not yet implemented. Sorry"),
    };

    let mut hashed = openssl_hash::hash(openssl_type, data);

    VecMultiHash::encode(&hashed[..], hash_type as u8).unwrap()
}

struct HashFnTypeData {
    name: &'static str,
    default_len: u8,
}

fn hashfn_data(hft: &HashFnType) -> HashFnTypeData {
    match *hft {
        SHA1 =>
            HashFnTypeData { name: "sha1", default_len: 20 },
        SHA2_256 =>
            HashFnTypeData { name: "sha2-256", default_len: 32 },
        SHA2_512 =>
            HashFnTypeData { name: "sha2-512", default_len: 64 },
        SHA3 =>
            HashFnTypeData { name: "sha3", default_len: 64 },
        BLAKE2B =>
            HashFnTypeData { name: "blake2b", default_len: 64 },
        BLAKE2S =>
            HashFnTypeData { name: "blake2s", default_len: 32 },
    }
}

pub struct VecMultiHash {
    vec: Vec<u8>,
}

impl VecMultiHash {
    pub fn new() -> VecMultiHash {
        VecMultiHash { vec: Vec::new() }
    }

    // For constructing a VecMultiHash directly from a vector.
    // Function assumes that first byte is hash function code,
    // second is digest length. It's up to you.
    pub fn from_vec(v: Vec<u8>) -> VecMultiHash {
        VecMultiHash { vec: v }
    }

    pub fn encode<'a>(data: &'a [u8], code: u8) -> Result<VecMultiHash, EncodeError> {
        match HashFnType::from_u8(code) {
            None => return Err(EncodeError::UnknownCode(code)),
            Some(_) => {},
        }

        let size = data.len();
        if size > 127 {
            return Err(EncodeError::NotSupported(size));
        }

        let mut v = Vec::with_capacity(size + 2);
        v.push(code);
        v.push(size as u8);
        v.push_all(data);
        Ok(VecMultiHash::from_vec(v))
    }

    pub fn decode<'a>(&'a self) -> Result<DecodedMultiHash<'a>, DecodeError> {
        let mhlen = self.vec.len();
        if mhlen < 3 {
            return Err(DecodeError::TooShort);
        } else if mhlen > 129 {
            return Err(DecodeError::TooLong);
        } else {
            let digest_len = (mhlen - 2) as u8;
            if digest_len != self.vec[1] {
                return Err(DecodeError::InvalidDigestLength(self.vec[1], digest_len));
            }
        }

        let fn_code = match HashFnType::from_u8(self.vec[0]) {
            None => return Err(DecodeError::UnknownCode(self.vec[0])),
            Some(code) => code,
        };

        let hash_fn = hashfn_data(&fn_code);

        let decoded = DecodedMultiHash {
            code: fn_code,
            name: hash_fn.name,
            length: self.vec[1],
            digest: &self.vec[2..],
        };
        Ok(decoded)

    }

    pub fn vec(self) -> Vec<u8> {
        self.vec
    }

    pub fn to_base58_string(&self) -> String {
        self.vec.to_base58()
    }
}

pub struct DecodedMultiHash<'a> {
    code: HashFnType,
    name: &'static str,
    length: u8,
    digest: &'a [u8],
}

#[derive(Debug)]
pub enum DecodeError {
    UnknownCode(u8),
    TooShort,
    TooLong,
    InvalidDigestLength(u8, u8), // (stated, actual)
}

#[derive(Debug)]
pub enum EncodeError {
    UnknownCode(u8),
    NotSupported(usize),
}

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

    #[test]
    fn test_hash() {
        multihash::hash(b"ABC", multihash::HashFnType::SHA1);
    }

}
<|endoftext|>


<commit_before>use cdchunking::{Chunker, ChunkInput, ZPAQ};
use rusqlite;
use rusqlite::{Connection, Transaction};
use rusqlite::types::ToSql;
use sha1::Sha1;
use std::fs::File;
use std::path::{Path, PathBuf};

use crate::{Error, HashDigest};

const SCHEMA: &'static str = "
    CREATE TABLE version(
        name VARCHAR(8) NOT NULL,
        version VARCHAR(16) NOT NULL
    );
    INSERT INTO version(name, version) VALUES('rrsync', '0.1');

    CREATE TABLE files(
        file_id INTEGER NOT NULL PRIMARY KEY,
        name VARCHAR(512) NOT NULL,
        modified DATETIME NOT NULL
    );
    CREATE INDEX idx_files_name ON files(name);

    CREATE TABLE blocks(
        hash VARCHAR(40) NOT NULL,
        file_id INTEGER NOT NULL,
        offset INTEGER NOT NULL,
        PRIMARY KEY(file_id, offset)
    );
    CREATE INDEX idx_blocks_hash ON blocks(hash);
    CREATE INDEX idx_blocks_file ON blocks(file_id);
    CREATE INDEX idx_blocks_file_offset ON blocks(file_id, offset);
";

fn get_block(
    db: &Connection,
    hash: &HashDigest,
) -> Result<Option<(PathBuf, usize)>, Error>
{
    let mut stmt = db.prepare(
        "
        SELECT files.name, blocks.offset
        FROM blocks
        INNER JOIN files ON blocks.file_id = files.file_id
        WHERE blocks.hash = ?;
        ",
    )?;
    let mut rows = stmt.query(&[hash as &dyn ToSql])?;
    if let Some(row) = rows.next() {
        let row = row?;
        let path: String = row.get(0);
        let path: PathBuf = path.into();
        let offset: i64 = row.get(1);
        let offset = offset as usize;
        Ok(Some((path, offset)))
    } else {
        Ok(None)
    }
}

/// Index of files and blocks
pub struct Index {
    db: Connection,
}

impl Index {
    /// Open an index from a file
    pub fn open(filename: &Path) -> Result<Index, Error> {
        let exists = filename.exists();
        let db = Connection::open(filename)?;
        if !exists {
            warn!("Database doesn't exist, creating tables...");
            db.execute_batch(SCHEMA)?;
        }
        Ok(Index { db })
    }

    /// Open an in-memory index
    pub fn open_in_memory() -> Result<Index, Error> {
        let db = Connection::open_in_memory()?;
        db.execute_batch(SCHEMA)?;
        Ok(Index { db })
    }

    /// Try to find a block in the indexed files
    pub fn get_block(
        &self,
        hash: &HashDigest,
    ) -> Result<Option<(PathBuf, usize)>, Error>
    {
        get_block(&self.db, hash)
    }

    /// Start a transaction to update the index
    pub fn transaction<'a>(
        &'a mut self
    ) -> Result<IndexTransaction<'a>, rusqlite::Error>
    {
        let tx = self.db.transaction()?;
        Ok(IndexTransaction { tx })
    }
}

/// A transaction on the index, for safety and performance
pub struct IndexTransaction<'a> {
    tx: Transaction<'a>,
}

const ZPAQ_BITS: usize = 13; // 13 bits = 8 KiB block average
const MAX_BLOCK_SIZE: usize = 1 << 15; // 32 KiB

impl<'a> IndexTransaction<'a> {
    /// Add a file to the index
    ///
    /// This returns a tuple `(file_id, up_to_date)` where `file_id` can be
    /// used to insert blocks, and `up_to_date` indicates whether the file's
    /// modification date has changed and it should be re-indexed.
    pub fn add_file(
        &mut self,
        name: &Path,
        modified: chrono::DateTime<chrono::Utc>,
    ) -> Result<(u32, bool), Error>
    {
        let mut stmt = self.tx.prepare(
            "
            SELECT file_id, modified FROM files
            WHERE name = ?;
            ",
        )?;
        let mut rows = stmt.query(&[name.to_str().expect("encoding")])?;
        if let Some(row) = rows.next() {
            let row = row?;
            let file_id: u32 = row.get(0);
            let old_modified: chrono::DateTime<chrono::Utc> = row.get(1);
            if old_modified != modified {
                info!("Resetting file {:?}, modified", name);
                // Delete blocks
                self.tx.execute(
                    "
                    DELETE FROM blocks WHERE file_id = ?;
                    ",
                    &[&file_id],
                )?;
                // Update modification time
                self.tx.execute(
                    "
                    UPDATE files SET modified = ? WHERE file_id = ?;
                    ",
                    &[&modified as &dyn ToSql, &file_id],
                )?;
                Ok((file_id, false))
            } else {
                debug!("File {:?} up to date", name);
                Ok((file_id, true))
            }
        } else {
            info!("Inserting new file {:?}", name);
            self.tx.execute(
                "
                INSERT INTO files(name, modified)
                VALUES(?, ?);
                ",
                &[&name.to_str().expect("encoding") as &dyn ToSql, &modified],
            )?;
            let file_id = self.tx.last_insert_rowid();
            Ok((file_id as u32, false))
        }
    }

    /// Remove a file and all its blocks from the index
    pub fn remove_file(
        &mut self,
        file_id: u32,
    ) -> Result<(), Error>
    {
        self.tx.execute(
            "
            DELETE FROM blocks WHERE file_id = ?;
            ",
            &[&file_id],
        )?;
        self.tx.execute(
            "
            DELETE FROM files WHERE file_id = ?;
            ",
            &[&file_id],
        )?;
        Ok(())
    }

    /// Get a list of all the files in the index
    pub fn list_files(&self) -> Result<Vec<(u32, PathBuf)>, Error> {
        let mut stmt = self.tx.prepare(
            "
            SELECT file_id, name FROM files;
            ",
        )?;
        let mut rows = stmt.query(rusqlite::NO_PARAMS)?;
        let mut results = Vec::new();
        loop {
            match rows.next() {
                Some(Ok(row)) => {
                    let path: String = row.get(1);
                    results.push((row.get(0), path.into()))
                }
                Some(Err(e)) => return Err(e.into()),
                None => break,
            }
        }
        Ok(results)
    }

    /// Add a block to the index
    pub fn add_block(
        &mut self,
        hash: &HashDigest,
        file_id: u32,
        offset: usize,
    ) -> Result<(), Error>
    {
        self.tx.execute(
            "
            INSERT INTO blocks(hash, file_id, offset)
            VALUES(?, ?, ?);
            ",
            &[&hash as &dyn ToSql, &file_id, &(offset as i64)],
        )?;
        Ok(())
    }

    /// Try to find a block in the indexed files
    pub fn get_block(
        &self,
        hash: &HashDigest,
    ) -> Result<Option<(PathBuf, usize)>, Error>
    {
        get_block(&self.tx, hash)
    }

    /// Cut up a file into blocks and add them to the index
    pub fn index_file(
        &mut self,
        name: &Path,
    ) -> Result<(), Error>
    {
        let file = File::open(name)?;
        let (file_id, up_to_date) = self.add_file(
            name,
            file.metadata()?.modified()?.into(),
        )?;
        if !up_to_date {
            // Use ZPAQ to cut the stream into blocks
            let chunker = Chunker::new(
                ZPAQ::new(ZPAQ_BITS) // 13 bits = 8 KiB block average
            ).max_size(MAX_BLOCK_SIZE);
            let mut chunk_iterator = chunker.stream(file);
            let mut start_offset = 0;
            let mut offset = 0;
            let mut sha1 = Sha1::new();
            while let Some(chunk) = chunk_iterator.read() {
                match chunk? {
                    ChunkInput::Data(d) => {
                        sha1.update(d);
                        offset += d.len();
                    }
                    ChunkInput::End => {
                        let digest = HashDigest(sha1.digest().bytes());
                        debug!(
                            "Adding block, offset={}, size={}, sha1={}",
                            start_offset, offset - start_offset, digest,
                        );
                        self.add_block(&digest, file_id, start_offset)?;
                        start_offset = offset;
                        sha1.reset();
                    }
                }
            }
        }
        Ok(())
    }

    /// Commit the transaction
    pub fn commit(self) -> Result<(), rusqlite::Error> {
        self.tx.commit()
    }
}

#[cfg(test)]
mod tests {
    use std::io::Write;
    use tempfile::NamedTempFile;

    use crate::HashDigest;
    use super::{Index, MAX_BLOCK_SIZE};

    #[test]
    fn test() {
        let mut file = NamedTempFile::new().expect("tempfile");
        for i in 0..2000 {
            write!(file, "Line {}\n", i + 1).expect("tempfile");
        }
        for _ in 0..2000 {
            write!(file, "Test content\n").expect("tempfile");
        }
        file.flush().expect("tempfile");
        let mut index = Index::open_in_memory().expect("db");
        {
            let mut tx = index.transaction().expect("db");
            tx.index_file(file.path()).expect("index");
            tx.commit().expect("db");
        }
        assert!(
            index.get_block(&HashDigest(*b"12345678901234567890"))
                .expect("get")
                .is_none()
        );
        let block1 = index.get_block(&HashDigest(
            *b"\xfb\x5e\xf7\xeb\xad\xd8\x2c\x80\x85\xc5\
               \xff\x63\x82\x36\x22\xba\xe0\xe2\x63\xf6"
        )).expect("get");
        assert_eq!(
            block1,
            Some((file.path().into(), 0)),
        );
        let block2 = index.get_block(&HashDigest(
            *b"\x57\x0d\x8b\x30\xfc\xfd\x58\x5e\x41\x27\
               \xb5\x61\xf5\xec\xd3\x76\xff\x4d\x01\x01"
        )).expect("get");
        assert_eq!(
            block2,
            Some((file.path().into(), 11579)),
        );
        let block3 = index.get_block(&HashDigest(
            *b"\xb9\xa8\xc2\x64\x1a\xf2\xcf\x8f\xd8\xf3\
               \x6a\x24\x56\xa3\xea\xa9\x5c\x02\x91\x27"
        )).expect("get");
        assert_eq!(
            block3,
            Some((file.path().into(), 44347)),
        );
        assert_eq!(block3.unwrap().1 - block2.unwrap().1, MAX_BLOCK_SIZE);
    }
}
<commit_msg>Store size of blocks in index<commit_after>use cdchunking::{Chunker, ChunkInput, ZPAQ};
use rusqlite;
use rusqlite::{Connection, Transaction};
use rusqlite::types::ToSql;
use sha1::Sha1;
use std::fs::File;
use std::path::{Path, PathBuf};

use crate::{Error, HashDigest};

const SCHEMA: &'static str = "
    CREATE TABLE version(
        name VARCHAR(8) NOT NULL,
        version VARCHAR(16) NOT NULL
    );
    INSERT INTO version(name, version) VALUES('rrsync', '0.1');

    CREATE TABLE files(
        file_id INTEGER NOT NULL PRIMARY KEY,
        name VARCHAR(512) NOT NULL,
        modified DATETIME NOT NULL
    );
    CREATE INDEX idx_files_name ON files(name);

    CREATE TABLE blocks(
        hash VARCHAR(40) NOT NULL,
        file_id INTEGER NOT NULL,
        offset INTEGER NOT NULL,
        size INTEGER NOT NULL,
        PRIMARY KEY(file_id, offset)
    );
    CREATE INDEX idx_blocks_hash ON blocks(hash);
    CREATE INDEX idx_blocks_file ON blocks(file_id);
    CREATE INDEX idx_blocks_file_offset ON blocks(file_id, offset);
";

fn get_block(
    db: &Connection,
    hash: &HashDigest,
) -> Result<Option<(PathBuf, usize, usize)>, Error>
{
    let mut stmt = db.prepare(
        "
        SELECT files.name, blocks.offset, blocks.size
        FROM blocks
        INNER JOIN files ON blocks.file_id = files.file_id
        WHERE blocks.hash = ?;
        ",
    )?;
    let mut rows = stmt.query(&[hash as &dyn ToSql])?;
    if let Some(row) = rows.next() {
        let row = row?;
        let path: String = row.get(0);
        let path: PathBuf = path.into();
        let offset: i64 = row.get(1);
        let offset = offset as usize;
        let size: i64 = row.get(2);
        let size = size as usize;
        Ok(Some((path, offset, size)))
    } else {
        Ok(None)
    }
}

/// Index of files and blocks
pub struct Index {
    db: Connection,
}

impl Index {
    /// Open an index from a file
    pub fn open(filename: &Path) -> Result<Index, Error> {
        let exists = filename.exists();
        let db = Connection::open(filename)?;
        if !exists {
            warn!("Database doesn't exist, creating tables...");
            db.execute_batch(SCHEMA)?;
        }
        Ok(Index { db })
    }

    /// Open an in-memory index
    pub fn open_in_memory() -> Result<Index, Error> {
        let db = Connection::open_in_memory()?;
        db.execute_batch(SCHEMA)?;
        Ok(Index { db })
    }

    /// Try to find a block in the indexed files
    pub fn get_block(
        &self,
        hash: &HashDigest,
    ) -> Result<Option<(PathBuf, usize, usize)>, Error>
    {
        get_block(&self.db, hash)
    }

    /// Start a transaction to update the index
    pub fn transaction<'a>(
        &'a mut self
    ) -> Result<IndexTransaction<'a>, rusqlite::Error>
    {
        let tx = self.db.transaction()?;
        Ok(IndexTransaction { tx })
    }
}

/// A transaction on the index, for safety and performance
pub struct IndexTransaction<'a> {
    tx: Transaction<'a>,
}

const ZPAQ_BITS: usize = 13; // 13 bits = 8 KiB block average
const MAX_BLOCK_SIZE: usize = 1 << 15; // 32 KiB

impl<'a> IndexTransaction<'a> {
    /// Add a file to the index
    ///
    /// This returns a tuple `(file_id, up_to_date)` where `file_id` can be
    /// used to insert blocks, and `up_to_date` indicates whether the file's
    /// modification date has changed and it should be re-indexed.
    pub fn add_file(
        &mut self,
        name: &Path,
        modified: chrono::DateTime<chrono::Utc>,
    ) -> Result<(u32, bool), Error>
    {
        let mut stmt = self.tx.prepare(
            "
            SELECT file_id, modified FROM files
            WHERE name = ?;
            ",
        )?;
        let mut rows = stmt.query(&[name.to_str().expect("encoding")])?;
        if let Some(row) = rows.next() {
            let row = row?;
            let file_id: u32 = row.get(0);
            let old_modified: chrono::DateTime<chrono::Utc> = row.get(1);
            if old_modified != modified {
                info!("Resetting file {:?}, modified", name);
                // Delete blocks
                self.tx.execute(
                    "
                    DELETE FROM blocks WHERE file_id = ?;
                    ",
                    &[&file_id],
                )?;
                // Update modification time
                self.tx.execute(
                    "
                    UPDATE files SET modified = ? WHERE file_id = ?;
                    ",
                    &[&modified as &dyn ToSql, &file_id],
                )?;
                Ok((file_id, false))
            } else {
                debug!("File {:?} up to date", name);
                Ok((file_id, true))
            }
        } else {
            info!("Inserting new file {:?}", name);
            self.tx.execute(
                "
                INSERT INTO files(name, modified)
                VALUES(?, ?);
                ",
                &[&name.to_str().expect("encoding") as &dyn ToSql, &modified],
            )?;
            let file_id = self.tx.last_insert_rowid();
            Ok((file_id as u32, false))
        }
    }

    /// Remove a file and all its blocks from the index
    pub fn remove_file(
        &mut self,
        file_id: u32,
    ) -> Result<(), Error>
    {
        self.tx.execute(
            "
            DELETE FROM blocks WHERE file_id = ?;
            ",
            &[&file_id],
        )?;
        self.tx.execute(
            "
            DELETE FROM files WHERE file_id = ?;
            ",
            &[&file_id],
        )?;
        Ok(())
    }

    /// Get a list of all the files in the index
    pub fn list_files(&self) -> Result<Vec<(u32, PathBuf)>, Error> {
        let mut stmt = self.tx.prepare(
            "
            SELECT file_id, name FROM files;
            ",
        )?;
        let mut rows = stmt.query(rusqlite::NO_PARAMS)?;
        let mut results = Vec::new();
        loop {
            match rows.next() {
                Some(Ok(row)) => {
                    let path: String = row.get(1);
                    results.push((row.get(0), path.into()))
                }
                Some(Err(e)) => return Err(e.into()),
                None => break,
            }
        }
        Ok(results)
    }

    /// Add a block to the index
    pub fn add_block(
        &mut self,
        hash: &HashDigest,
        file_id: u32,
        offset: usize,
        size: usize,
    ) -> Result<(), Error>
    {
        self.tx.execute(
            "
            INSERT INTO blocks(hash, file_id, offset, size)
            VALUES(?, ?, ?, ?);
            ",
            &[&hash as &dyn ToSql, &file_id, &(offset as i64), &(size as i64)],
        )?;
        Ok(())
    }

    /// Get a list of all the blocks in a specific file
    pub fn list_file_blocks(&self, file_id: u32) -> Result<Vec<(HashDigest, usize, usize)>, Error> {
        let mut stmt = self.tx.prepare(
            "
            SELECT hash, offset, size FROM blocks
            WHERE file_id = ?;
            ",
        )?;
        let mut rows = stmt.query(&[file_id])?;
        let mut results = Vec::new();
        loop {
            match rows.next() {
                Some(Ok(row)) => {
                    let offset: i64 = row.get(1);
                    let size: i64 = row.get(2);
                    results.push((row.get(0), offset as usize, size as usize))
                }
                Some(Err(e)) => return Err(e.into()),
                None => break,
            }
        }
        Ok(results)
    }

    /// Try to find a block in the indexed files
    pub fn get_block(
        &self,
        hash: &HashDigest,
    ) -> Result<Option<(PathBuf, usize, usize)>, Error>
    {
        get_block(&self.tx, hash)
    }

    /// Cut up a file into blocks and add them to the index
    pub fn index_file(
        &mut self,
        name: &Path,
    ) -> Result<(), Error>
    {
        let file = File::open(name)?;
        let (file_id, up_to_date) = self.add_file(
            name,
            file.metadata()?.modified()?.into(),
        )?;
        if !up_to_date {
            // Use ZPAQ to cut the stream into blocks
            let chunker = Chunker::new(
                ZPAQ::new(ZPAQ_BITS) // 13 bits = 8 KiB block average
            ).max_size(MAX_BLOCK_SIZE);
            let mut chunk_iterator = chunker.stream(file);
            let mut start_offset = 0;
            let mut offset = 0;
            let mut sha1 = Sha1::new();
            while let Some(chunk) = chunk_iterator.read() {
                match chunk? {
                    ChunkInput::Data(d) => {
                        sha1.update(d);
                        offset += d.len();
                    }
                    ChunkInput::End => {
                        let digest = HashDigest(sha1.digest().bytes());
                        let size = offset - start_offset;
                        debug!(
                            "Adding block, offset={}, size={}, sha1={}",
                            start_offset, size, digest,
                        );
                        self.add_block(&digest, file_id, start_offset, size)?;
                        start_offset = offset;
                        sha1.reset();
                    }
                }
            }
        }
        Ok(())
    }

    /// Commit the transaction
    pub fn commit(self) -> Result<(), rusqlite::Error> {
        self.tx.commit()
    }
}

#[cfg(test)]
mod tests {
    use std::io::Write;
    use tempfile::NamedTempFile;

    use crate::HashDigest;
    use super::{Index, MAX_BLOCK_SIZE};

    #[test]
    fn test() {
        let mut file = NamedTempFile::new().expect("tempfile");
        for i in 0..2000 {
            write!(file, "Line {}\n", i + 1).expect("tempfile");
        }
        for _ in 0..2000 {
            write!(file, "Test content\n").expect("tempfile");
        }
        file.flush().expect("tempfile");
        let mut index = Index::open_in_memory().expect("db");
        {
            let mut tx = index.transaction().expect("db");
            tx.index_file(file.path()).expect("index");
            tx.commit().expect("db");
        }
        assert!(
            index.get_block(&HashDigest(*b"12345678901234567890"))
                .expect("get")
                .is_none()
        );
        let block1 = index.get_block(&HashDigest(
            *b"\xfb\x5e\xf7\xeb\xad\xd8\x2c\x80\x85\xc5\
               \xff\x63\x82\x36\x22\xba\xe0\xe2\x63\xf6"
        )).expect("get");
        assert_eq!(
            block1,
            Some((file.path().into(), 0, 11579)),
        );
        let block2 = index.get_block(&HashDigest(
            *b"\x57\x0d\x8b\x30\xfc\xfd\x58\x5e\x41\x27\
               \xb5\x61\xf5\xec\xd3\x76\xff\x4d\x01\x01"
        )).expect("get");
        assert_eq!(
            block2,
            Some((file.path().into(), 11579, 32768)),
        );
        let block3 = index.get_block(&HashDigest(
            *b"\xb9\xa8\xc2\x64\x1a\xf2\xcf\x8f\xd8\xf3\
               \x6a\x24\x56\xa3\xea\xa9\x5c\x02\x91\x27"
        )).expect("get");
        assert_eq!(
            block3,
            Some((file.path().into(), 44347, 546)),
        );
        assert_eq!(block3.unwrap().1 - block2.unwrap().1, MAX_BLOCK_SIZE);
    }
}
<|endoftext|>


<commit_before>//! An abstract input state object that gets fed user
//! events and updates itself based on a set of key
//! bindings.
//! The idea is threefold:
//!
//! * Have a layer of abstract key bindings rather than
//! looking at concrete event types
//! * Use this to be able to abstract away differences
//! between keyboards, joysticks and game controllers
//! (rather based on Unity3D),
//! * Do some tweening of input axes and stuff just for
//! fun maybe.
//!
//! Right now ggez doesn't handle joysticks or controllers
//! anyway, so.

use std::cmp::PartialOrd;
use std::hash::Hash;
use std::collections::BTreeMap;
use std::collections::HashMap;
use ggez::event::*;

// Okay, but how does it actually work?
// Basically we have to bind input events to buttons and axes.
// Input events can be keys, mouse buttons/motion, or eventually
// joystick/controller inputs.  Mouse delta can be mapped to axes too.
//
// https://docs.unity3d.com/Manual/ConventionalGameInput.html has useful
// descriptions of the exact behavior of axes.
//
// So to think about this more clearly, here are the default bindings:
//
// W, ↑: +Y axis
// A, ←: -X axis
// S, ↓: -Y axis
// D, →: +X axis
// Enter, z, LMB: Button 1
// Shift, x, MMB: Button 2
// Ctrl,  c, RMB: Button 3
//
// Easy way?  Hash map of event -> axis/button bindings.

#[derive(Debug, Hash, Eq, PartialEq, Copy, Clone)]
enum InputEvent {
    KeyEvent(Keycode),
    MouseButtonEvent,
}

#[derive(Debug, Eq, Ord, PartialEq, PartialOrd, Copy, Clone)]
enum InputEffect<Axes, Buttons>
    where Axes: Eq + Ord + Clone,
          Buttons: Eq + Ord + Clone
{
    Axis(Axes, bool),
    Button(Buttons),
}

#[derive(Debug)]
struct AxisStatus {
    position: f64,
    current_direction: f64,
}

#[derive(Debug)]
pub struct InputManager<Axes, Buttons>
    where Axes: Eq + Ord + Hash + Clone,
          Buttons: Eq + Ord + Hash + Clone
{
    // Once EnumSet is stable it should be used for these
    // instead of BTreeMap. ♥?
    // Binding of keys to input values.
    bindings: HashMap<InputEvent, InputEffect<Axes, Buttons>>,
    // Input state for axes
    axes: BTreeMap<Axes, AxisStatus>,
    // Input states for buttons
    buttons: BTreeMap<Buttons, bool>,
}

impl<Axes, Buttons> InputManager<Axes, Buttons>
    where Axes: Eq + Ord + Hash + Clone,
          Buttons: Eq + Ord + Hash + Clone
{
    pub fn new() -> Self {
        InputManager {
            bindings: HashMap::new(),
            axes: BTreeMap::new(),
            buttons: BTreeMap::new(),
        }
    }

    pub fn bind_key_to_axis(mut self, keycode: Keycode, axis: Axes, positive: bool) -> Self {

        self.bindings.insert(InputEvent::KeyEvent(keycode),
                             InputEffect::Axis(axis, positive));
        self
    }

    pub fn bind_key_to_button(mut self, keycode: Keycode, button: Buttons) -> Self {
        self.bindings.insert(InputEvent::KeyEvent(keycode), InputEffect::Button(button));
        self
    }

    pub fn update_keydown(&mut self, keycode: Option<Keycode>) {
        if let Some(keycode) = keycode {
            let effect = {
                if let Some(e) = self.bindings.get(&InputEvent::KeyEvent(keycode)) {
                    e.clone()
                } else {
                    return;
                }
            };
            self.update_effect(effect, true);
        }
    }

    pub fn update_keyup(&mut self, keycode: Option<Keycode>) {
        if let Some(keycode) = keycode {
            let effect = {
                if let Some(e) = self.bindings.get(&InputEvent::KeyEvent(keycode)) {
                    e.clone()
                } else {
                    return;
                }
            };
            self.update_effect(effect, false);
        }
    }

    fn update_effect(&mut self, effect: InputEffect<Axes, Buttons>, started: bool) {
        match effect {
            InputEffect::Axis(axis, direction) => {
                let direction_float = if direction {
                    1.0
                } else {
                    -1.0
                };
                let default_status = AxisStatus {
                    position: 0.0,
                    current_direction: direction_float,
                };
                let axis_status = self.axes.entry(axis).or_insert(default_status);
                axis_status.current_direction = direction_float;
            }
            InputEffect::Button(button) => {
                let button_pressed = self.buttons.entry(button).or_insert(started);
                *button_pressed = started;

            }
        }
    }

    pub fn mouse_position() {}

    pub fn mouse_scroll_delta() {}

    pub fn get_axis() {}

    pub fn get_axis_raw() {}

    pub fn get_button() {}

    pub fn get_button_down() {}

    pub fn get_button_up() {}

    pub fn get_mouse_button() {}

    pub fn get_mouse_button_down() {}

    pub fn get_mouse_button_up() {}

    pub fn reset_input_axes() {}
}

mod tests {

    #[ignore(unused_imports)]
    use ggez::event::*;
    #[ignore(unused_imports)]
    use super::*;

    #[ignore(dead_code)]
    #[derive(Hash, Copy, Clone, Eq, Ord, PartialEq, PartialOrd)]
    enum Buttons {
        Up,
        Left,
        Right,
        Down,
    }
    #[test]
    fn test_input_events() {
        let im = InputManager::<(), Buttons>::new().bind_key_to_button(Keycode::W, Buttons::Up);
    }
}
<commit_msg>Slowly adding more functionality to input.<commit_after>//! An abstract input state object that gets fed user
//! events and updates itself based on a set of key
//! bindings.
//! The idea is threefold:
//!
//! * Have a layer of abstract key bindings rather than
//! looking at concrete event types
//! * Use this to be able to abstract away differences
//! between keyboards, joysticks and game controllers
//! (rather based on Unity3D),
//! * Do some tweening of input axes and stuff just for
//! fun maybe.
//!
//! Right now ggez doesn't handle joysticks or controllers
//! anyway, so.

use std::cmp::PartialOrd;
use std::hash::Hash;
use std::collections::BTreeMap;
use std::collections::HashMap;
use ggez::event::*;

// Okay, but how does it actually work?
// Basically we have to bind input events to buttons and axes.
// Input events can be keys, mouse buttons/motion, or eventually
// joystick/controller inputs.  Mouse delta can be mapped to axes too.
//
// https://docs.unity3d.com/Manual/ConventionalGameInput.html has useful
// descriptions of the exact behavior of axes.
//
// So to think about this more clearly, here are the default bindings:
//
// W, ↑: +Y axis
// A, ←: -X axis
// S, ↓: -Y axis
// D, →: +X axis
// Enter, z, LMB: Button 1
// Shift, x, MMB: Button 2
// Ctrl,  c, RMB: Button 3
//
// Easy way?  Hash map of event -> axis/button bindings.

#[derive(Debug, Hash, Eq, PartialEq, Copy, Clone)]
enum InputEvent {
    KeyEvent(Keycode),
    MouseButtonEvent,
}

#[derive(Debug, Eq, Ord, PartialEq, PartialOrd, Copy, Clone)]
enum InputEffect<Axes, Buttons>
    where Axes: Eq + Ord + Clone,
          Buttons: Eq + Ord + Clone
{
    Axis(Axes, bool),
    Button(Buttons),
}

#[derive(Debug)]
struct AxisStatus {
    position: f64,
    current_direction: f64,
}

impl Default for AxisStatus {
    fn default() -> Self {
        AxisStatus {
            position: 0.0,
            current_direction: 0.0,
        }
    }
}

#[derive(Debug)]
pub struct InputManager<Axes, Buttons>
    where Axes: Eq + Ord + Hash + Clone,
          Buttons: Eq + Ord + Hash + Clone
{
    // Once EnumSet is stable it should be used for these
    // instead of BTreeMap. ♥?
    // Binding of keys to input values.
    bindings: HashMap<InputEvent, InputEffect<Axes, Buttons>>,
    // Input state for axes
    axes: BTreeMap<Axes, AxisStatus>,
    // Input states for buttons
    buttons: BTreeMap<Buttons, bool>,
}

impl<Axes, Buttons> InputManager<Axes, Buttons>
    where Axes: Eq + Ord + Hash + Clone,
          Buttons: Eq + Ord + Hash + Clone
{
    pub fn new() -> Self {
        InputManager {
            bindings: HashMap::new(),
            axes: BTreeMap::new(),
            buttons: BTreeMap::new(),
        }
    }

    pub fn bind_key_to_axis(mut self, keycode: Keycode, axis: Axes, positive: bool) -> Self {

        self.bindings.insert(InputEvent::KeyEvent(keycode),
                             InputEffect::Axis(axis, positive));
        self
    }

    pub fn bind_key_to_button(mut self, keycode: Keycode, button: Buttons) -> Self {
        self.bindings.insert(InputEvent::KeyEvent(keycode), InputEffect::Button(button));
        self
    }

    pub fn update_keydown(&mut self, keycode: Option<Keycode>) {
        if let Some(keycode) = keycode {
            let effect = {
                if let Some(e) = self.bindings.get(&InputEvent::KeyEvent(keycode)) {
                    e.clone()
                } else {
                    return;
                }
            };
            self.update_effect(effect, true);
        }
    }

    pub fn update_keyup(&mut self, keycode: Option<Keycode>) {
        if let Some(keycode) = keycode {
            let effect = {
                if let Some(e) = self.bindings.get(&InputEvent::KeyEvent(keycode)) {
                    e.clone()
                } else {
                    return;
                }
            };
            self.update_effect(effect, false);
        }
    }

    fn update_effect(&mut self, effect: InputEffect<Axes, Buttons>, started: bool) {
        match effect {
            InputEffect::Axis(axis, direction) => {
                let direction_float = if direction {
                    1.0
                } else {
                    -1.0
                };
                let default_status = AxisStatus {
                    position: 0.0,
                    current_direction: direction_float,
                };
                let axis_status = self.axes.entry(axis).or_insert(default_status);
                axis_status.current_direction = direction_float;
            }
            InputEffect::Button(button) => {
                let button_pressed = self.buttons.entry(button).or_insert(started);
                *button_pressed = started;

            }
        }
    }

    pub fn get_axis(&self, axis: &Axes) -> f64 {
        if let Some(ax) = self.axes.get(axis) {
            ax.position
        } else {
            0.0
        }
    }

    pub fn get_axis_raw() {}

    pub fn get_button(&self, axis: &Buttons) -> bool {
        if let Some(pressed) = self.buttons.get(axis) {
            *pressed
        } else {
            false
        }
    }

    pub fn mouse_position() {}

    pub fn mouse_scroll_delta() {}

    pub fn get_button_down() {}

    pub fn get_button_up() {}

    pub fn get_mouse_button() {}

    pub fn get_mouse_button_down() {}

    pub fn get_mouse_button_up() {}

    pub fn reset_input_axes() {}
}

mod tests {

    #[ignore(unused_imports)]
    use ggez::event::*;
    #[ignore(unused_imports)]
    use super::*;

    #[ignore(dead_code)]
    #[derive(Hash, Copy, Clone, Eq, Ord, PartialEq, PartialOrd)]
    enum Buttons {
        Up,
        Left,
        Right,
        Down,
    }
    #[test]
    fn test_input_events() {
        let mut im = InputManager::<(), Buttons>::new().bind_key_to_button(Keycode::W, Buttons::Up);

        im.update_keydown(Some(Keycode::W));
        assert!(im.get_button(&Buttons::Up));
        im.update_keyup(Some(Keycode::W));
        assert!(!im.get_button(&Buttons::Up));
    }
}
<|endoftext|>


<commit_before>//! Input sessions for simplified collection updates.
//!
//! Although users can directly manipulate timely dataflow streams as collection inputs,
//! the `InputSession` type can make this more efficient and less error-prone. Specifically,
//! the type batches up updates with their logical times and ships them with coarsened
//! timely dataflow capabilities, exposing more concurrency to the operator implementations
//! than are evident from the logical times, which appear to execute in sequence.

use timely::communication::Allocate;

use timely::progress::Timestamp;
use timely::progress::timestamp::RootTimestamp;
use timely::progress::nested::product::Product;
use timely::dataflow::operators::Input as TimelyInput;
use timely::dataflow::operators::input::Handle;
use timely::dataflow::scopes::{Child, Root};

use ::{Data, Diff};
use collection::{Collection, AsCollection};

/// Create a new collection and input handle to control the collection.
pub trait Input<'a, A: Allocate, T: Timestamp+Ord> {
    /// Create a new collection and input handle to subsequently control the collection.
    ///
    /// # Examples
    ///
    /// ```
    /// extern crate timely;
    /// extern crate differential_dataflow;
    ///
    /// use timely::Configuration;
    /// use differential_dataflow::input::Input;
    ///
    /// fn main() {
    ///     ::timely::execute(Configuration::Thread, |worker| {
    ///
    ///			let (mut handle, probe) = worker.dataflow::<(),_,_>(|scope| {
    ///				// create input handle and collection.
    ///				let (handle, data) = scope.new_collection();
    ///         	let probe = data.map(|x| x * 2)
    ///				            	.inspect(|x| println!("{:?}", x))
    ///				            	.probe();
    ///				(handle, probe)
    ///     	});
    ///
    ///			handle.insert(1);
    ///			handle.insert(5);
    ///
    ///		}).unwrap();
    /// }
    /// ```
    fn new_collection<D, R>(&mut self) -> (InputSession<T, D, R>, Collection<Child<'a, Root<A>, T>, D, R>)
    where D: Data, R: Diff;
    /// Create a new collection and input handle from initial data.
    ///
    /// # Examples
    ///
    /// ```
    /// extern crate timely;
    /// extern crate differential_dataflow;
    ///
    /// use timely::Configuration;
    /// use differential_dataflow::input::Input;
    ///
    /// fn main() {
    ///     ::timely::execute(Configuration::Thread, |worker| {
    ///
    ///			let (mut handle, probe) = worker.dataflow::<(),_,_>(|scope| {
    ///				// create input handle and collection.
    ///				let (handle, data) = scope.new_collection_from(0 .. 10);
    ///         	let probe = data.map(|x| x * 2)
    ///				            	.inspect(|x| println!("{:?}", x))
    ///				            	.probe();
    ///				(handle, probe)
    ///     	});
    ///
    ///			handle.insert(1);
    ///			handle.insert(5);
    ///
    ///		}).unwrap();
    /// }
    /// ```
    fn new_collection_from<I>(&mut self, data: I) -> (InputSession<T, I::Item, isize>, Collection<Child<'a, Root<A>, T>, I::Item, isize>)
    where I: IntoIterator+'static, I::Item: Data;
}

impl<'a, A: Allocate, T: Timestamp+Ord> Input<'a, A, T> for Child<'a, Root<A>, T> {
    fn new_collection<D, R>(&mut self) -> (InputSession<T, D, R>, Collection<Child<'a, Root<A>, T>, D, R>)
    where D: Data, R: Diff{
		let (handle, stream) = self.new_input();
		(InputSession::from(handle), stream.as_collection())
    }
    fn new_collection_from<I>(&mut self, data: I) -> (InputSession<T, I::Item, isize>, Collection<Child<'a, Root<A>, T>, I::Item, isize>)
    where I: IntoIterator+'static, I::Item: Data {

    	use timely::dataflow::operators::ToStream;

		let (handle, stream) = self.new_input();
		let source = data.into_iter().map(|d| (d, Default::default(), 1)).to_stream(self).as_collection();

		(InputSession::from(handle), stream.as_collection().concat(&source))
    }
}

/// An input session wrapping a single timely dataflow capability.
///
/// Each timely dataflow message has a corresponding capability, which is a logical time in the
/// timely dataflow system. Differential dataflow updates can happen at a much higher rate than
/// timely dataflow's progress tracking infrastructure supports, because the logical times are
/// promoted to data and updates are batched together. The `InputSession` type does this batching.
///
/// # Examples
///
/// ```
/// extern crate timely;
/// extern crate differential_dataflow;
///
/// use timely::Configuration;
/// use differential_dataflow::input::Input;
///
/// fn main() {
///     ::timely::execute(Configuration::Thread, |worker| {
///
///			let (mut handle, probe) = worker.dataflow(|scope| {
///				// create input handle and collection.
///				let (handle, data) = scope.new_collection_from(0 .. 10);
///         	let probe = data.map(|x| x * 2)
///				            	.inspect(|x| println!("{:?}", x))
///				            	.probe();
///				(handle, probe)
///     	});
///
///			handle.insert(3);
///			handle.advance_to(1);
///			handle.insert(5);
///			handle.advance_to(2);
///			handle.flush();
///
///			while probe.less_than(handle.time()) {
///				worker.step();
///			}
///
///			handle.remove(5);
///			handle.advance_to(3);
///			handle.flush();
///
///			while probe.less_than(handle.time()) {
///				worker.step();
///			}
///
///		}).unwrap();
/// }
/// ```
pub struct InputSession<T: Timestamp+Clone, D: Data, R: Diff> {
	time: Product<RootTimestamp, T>,
	buffer: Vec<(D, Product<RootTimestamp, T>, R)>,
	handle: Handle<T,(D,Product<RootTimestamp, T>,R)>,
}

impl<T: Timestamp+Clone, D: Data> InputSession<T, D, isize> {
	/// Adds an element to the collection.
	pub fn insert(&mut self, element: D) { self.update(element, 1); }
	/// Removes an element from the collection.
	pub fn remove(&mut self, element: D) { self.update(element,-1); }
}

// impl<T: Timestamp+Clone, D: Data> InputSession<T, D, i64> {
//     /// Adds an element to the collection.
//     pub fn insert(&mut self, element: D) { self.update(element, 1); }
//     /// Removes an element from the collection.
//     pub fn remove(&mut self, element: D) { self.update(element,-1); }
// }

// impl<T: Timestamp+Clone, D: Data> InputSession<T, D, i32> {
//     /// Adds an element to the collection.
//     pub fn insert(&mut self, element: D) { self.update(element, 1); }
//     /// Removes an element from the collection.
//     pub fn remove(&mut self, element: D) { self.update(element,-1); }
// }

impl<'a, T: Timestamp+Clone, D: Data, R: Diff> InputSession<T, D, R> {

	/// Creates a new session from a reference to an input handle.
	pub fn from(handle: Handle<T,(D,Product<RootTimestamp, T>,R)>) -> Self {
		InputSession {
			time: handle.time().clone(),
			buffer: Vec::new(),
			handle: handle,
		}
	}

	/// Adds to the weight of an element in the collection.
	pub fn update(&mut self, element: D, change: R) {
        if self.buffer.len() == self.buffer.capacity() {
            if self.buffer.len() > 0 {
                self.handle.send_batch(&mut self.buffer);
            }
            // TODO : This is a fairly arbitrary choice; should probably use `Context::default_size()` or such.
            self.buffer.reserve(1024);
        }
		self.buffer.push((element, self.time.clone(), change));
	}

	/// Forces buffered data into the timely dataflow input, and advances its time to match that of the session.
	///
	/// It is important to call `flush` before expecting timely dataflow to report progress. Until this method is
	/// called, all updates may still be in internal buffers and not exposed to timely dataflow. Once the method is
	/// called, all buffers are flushed and timely dataflow is advised that some logical times are no longer possible.
	pub fn flush(&mut self) {
		self.handle.send_batch(&mut self.buffer);
		if self.handle.epoch().less_than(&self.time.inner) {
			self.handle.advance_to(self.time.inner.clone());
		}
	}

	/// Advances the logical time for future records.
	///
	/// Importantly, this method does **not** immediately inform timely dataflow of the change. This happens only when
	/// the session is dropped or flushed. It is not correct to use this time as a basis for a computation's `step_while`
	/// method unless the session has just been flushed.
	pub fn advance_to(&mut self, time: T) {
		assert!(self.handle.epoch().less_equal(&time));
		assert!(&self.time.inner.less_equal(&time));
		self.time = Product::new(RootTimestamp, time);
	}

	/// Reveals the current time of the session.
	pub fn epoch(&self) -> &T { &self.time.inner }
	/// Reveals the current time of the session.
	pub fn time(&self) -> &Product<RootTimestamp, T> { &self.time }

	/// Closes the input, flushing and sealing the wrapped timely input.
	pub fn close(self) { }
}

impl<T: Timestamp+Clone, D: Data, R: Diff> Drop for InputSession<T, D, R> {
	fn drop(&mut self) {
		self.flush();
	}
}
<commit_msg>input ergonomics<commit_after>//! Input sessions for simplified collection updates.
//!
//! Although users can directly manipulate timely dataflow streams as collection inputs,
//! the `InputSession` type can make this more efficient and less error-prone. Specifically,
//! the type batches up updates with their logical times and ships them with coarsened
//! timely dataflow capabilities, exposing more concurrency to the operator implementations
//! than are evident from the logical times, which appear to execute in sequence.

use timely::communication::Allocate;

use timely::progress::Timestamp;
use timely::progress::timestamp::RootTimestamp;
use timely::progress::nested::product::Product;
use timely::dataflow::operators::Input as TimelyInput;
use timely::dataflow::operators::input::Handle;
use timely::dataflow::scopes::{Child, Root};

use ::{Data, Diff};
use collection::{Collection, AsCollection};

/// Create a new collection and input handle to control the collection.
pub trait Input<'a, A: Allocate, T: Timestamp+Ord> {
    /// Create a new collection and input handle to subsequently control the collection.
    ///
    /// # Examples
    ///
    /// ```
    /// extern crate timely;
    /// extern crate differential_dataflow;
    ///
    /// use timely::Configuration;
    /// use differential_dataflow::input::Input;
    ///
    /// fn main() {
    ///     ::timely::execute(Configuration::Thread, |worker| {
    ///
    ///			let (mut handle, probe) = worker.dataflow::<(),_,_>(|scope| {
    ///				// create input handle and collection.
    ///				let (handle, data) = scope.new_collection();
    ///         	let probe = data.map(|x| x * 2)
    ///				            	.inspect(|x| println!("{:?}", x))
    ///				            	.probe();
    ///				(handle, probe)
    ///     	});
    ///
    ///			handle.insert(1);
    ///			handle.insert(5);
    ///
    ///		}).unwrap();
    /// }
    /// ```
    fn new_collection<D, R>(&mut self) -> (InputSession<T, D, R>, Collection<Child<'a, Root<A>, T>, D, R>)
    where D: Data, R: Diff;
    /// Create a new collection and input handle from initial data.
    ///
    /// # Examples
    ///
    /// ```
    /// extern crate timely;
    /// extern crate differential_dataflow;
    ///
    /// use timely::Configuration;
    /// use differential_dataflow::input::Input;
    ///
    /// fn main() {
    ///     ::timely::execute(Configuration::Thread, |worker| {
    ///
    ///			let (mut handle, probe) = worker.dataflow::<(),_,_>(|scope| {
    ///				// create input handle and collection.
    ///				let (handle, data) = scope.new_collection_from(0 .. 10);
    ///         	let probe = data.map(|x| x * 2)
    ///				            	.inspect(|x| println!("{:?}", x))
    ///				            	.probe();
    ///				(handle, probe)
    ///     	});
    ///
    ///			handle.insert(1);
    ///			handle.insert(5);
    ///
    ///		}).unwrap();
    /// }
    /// ```
    fn new_collection_from<I>(&mut self, data: I) -> (InputSession<T, I::Item, isize>, Collection<Child<'a, Root<A>, T>, I::Item, isize>)
    where I: IntoIterator+'static, I::Item: Data;
}

impl<'a, A: Allocate, T: Timestamp+Ord> Input<'a, A, T> for Child<'a, Root<A>, T> {
    fn new_collection<D, R>(&mut self) -> (InputSession<T, D, R>, Collection<Child<'a, Root<A>, T>, D, R>)
    where D: Data, R: Diff{
		let (handle, stream) = self.new_input();
		(InputSession::from(handle), stream.as_collection())
    }
    fn new_collection_from<I>(&mut self, data: I) -> (InputSession<T, I::Item, isize>, Collection<Child<'a, Root<A>, T>, I::Item, isize>)
    where I: IntoIterator+'static, I::Item: Data {

    	use timely::dataflow::operators::ToStream;

		let (handle, stream) = self.new_input();
		let source = data.into_iter().map(|d| (d, Default::default(), 1)).to_stream(self).as_collection();

		(InputSession::from(handle), stream.as_collection().concat(&source))
    }
}

/// An input session wrapping a single timely dataflow capability.
///
/// Each timely dataflow message has a corresponding capability, which is a logical time in the
/// timely dataflow system. Differential dataflow updates can happen at a much higher rate than
/// timely dataflow's progress tracking infrastructure supports, because the logical times are
/// promoted to data and updates are batched together. The `InputSession` type does this batching.
///
/// # Examples
///
/// ```
/// extern crate timely;
/// extern crate differential_dataflow;
///
/// use timely::Configuration;
/// use differential_dataflow::input::Input;
///
/// fn main() {
///     ::timely::execute(Configuration::Thread, |worker| {
///
///			let (mut handle, probe) = worker.dataflow(|scope| {
///				// create input handle and collection.
///				let (handle, data) = scope.new_collection_from(0 .. 10);
///         	let probe = data.map(|x| x * 2)
///				            	.inspect(|x| println!("{:?}", x))
///				            	.probe();
///				(handle, probe)
///     	});
///
///			handle.insert(3);
///			handle.advance_to(1);
///			handle.insert(5);
///			handle.advance_to(2);
///			handle.flush();
///
///			while probe.less_than(handle.time()) {
///				worker.step();
///			}
///
///			handle.remove(5);
///			handle.advance_to(3);
///			handle.flush();
///
///			while probe.less_than(handle.time()) {
///				worker.step();
///			}
///
///		}).unwrap();
/// }
/// ```
pub struct InputSession<T: Timestamp+Clone, D: Data, R: Diff> {
	time: Product<RootTimestamp, T>,
	buffer: Vec<(D, Product<RootTimestamp, T>, R)>,
	handle: Handle<T,(D,Product<RootTimestamp, T>,R)>,
}

impl<T: Timestamp+Clone, D: Data> InputSession<T, D, isize> {
	/// Adds an element to the collection.
	pub fn insert(&mut self, element: D) { self.update(element, 1); }
	/// Removes an element from the collection.
	pub fn remove(&mut self, element: D) { self.update(element,-1); }
}

// impl<T: Timestamp+Clone, D: Data> InputSession<T, D, i64> {
//     /// Adds an element to the collection.
//     pub fn insert(&mut self, element: D) { self.update(element, 1); }
//     /// Removes an element from the collection.
//     pub fn remove(&mut self, element: D) { self.update(element,-1); }
// }

// impl<T: Timestamp+Clone, D: Data> InputSession<T, D, i32> {
//     /// Adds an element to the collection.
//     pub fn insert(&mut self, element: D) { self.update(element, 1); }
//     /// Removes an element from the collection.
//     pub fn remove(&mut self, element: D) { self.update(element,-1); }
// }

impl<T: Timestamp+Clone, D: Data, R: Diff> InputSession<T, D, R> {

    /// Introduces a handle as collection.
    pub fn to_collection<'a, A: Allocate>(&mut self, scope: &mut Child<'a, Root<A>, T>) -> Collection<Child<'a, Root<A>, T>, D, R> {
        scope
            .input_from(&mut self.handle)
            .as_collection()
    }

    /// Allocates a new input handle.
    pub fn new() -> Self {
        let handle = Handle::new();
        InputSession {
            time: handle.time().clone(),
            buffer: Vec::new(),
            handle,
        }
    }

	/// Creates a new session from a reference to an input handle.
	pub fn from(handle: Handle<T,(D,Product<RootTimestamp, T>,R)>) -> Self {
		InputSession {
			time: handle.time().clone(),
			buffer: Vec::new(),
			handle: handle,
		}
	}

	/// Adds to the weight of an element in the collection.
	pub fn update(&mut self, element: D, change: R) {
        if self.buffer.len() == self.buffer.capacity() {
            if self.buffer.len() > 0 {
                self.handle.send_batch(&mut self.buffer);
            }
            // TODO : This is a fairly arbitrary choice; should probably use `Context::default_size()` or such.
            self.buffer.reserve(1024);
        }
		self.buffer.push((element, self.time.clone(), change));
	}

    /// Adds to the weight of an element in the collection at a future time.
    pub fn update_at(&mut self, element: D, time: T, change: R) {
        assert!(self.time.inner.less_equal(&time));
        if self.buffer.len() == self.buffer.capacity() {
            if self.buffer.len() > 0 {
                self.handle.send_batch(&mut self.buffer);
            }
            // TODO : This is a fairly arbitrary choice; should probably use `Context::default_size()` or such.
            self.buffer.reserve(1024);
        }
        self.buffer.push((element, RootTimestamp::new(time), change));
    }

	/// Forces buffered data into the timely dataflow input, and advances its time to match that of the session.
	///
	/// It is important to call `flush` before expecting timely dataflow to report progress. Until this method is
	/// called, all updates may still be in internal buffers and not exposed to timely dataflow. Once the method is
	/// called, all buffers are flushed and timely dataflow is advised that some logical times are no longer possible.
	pub fn flush(&mut self) {
		self.handle.send_batch(&mut self.buffer);
		if self.handle.epoch().less_than(&self.time.inner) {
			self.handle.advance_to(self.time.inner.clone());
		}
	}

	/// Advances the logical time for future records.
	///
	/// Importantly, this method does **not** immediately inform timely dataflow of the change. This happens only when
	/// the session is dropped or flushed. It is not correct to use this time as a basis for a computation's `step_while`
	/// method unless the session has just been flushed.
	pub fn advance_to(&mut self, time: T) {
		assert!(self.handle.epoch().less_equal(&time));
		assert!(&self.time.inner.less_equal(&time));
		self.time = Product::new(RootTimestamp, time);
	}

	/// Reveals the current time of the session.
	pub fn epoch(&self) -> &T { &self.time.inner }
	/// Reveals the current time of the session.
	pub fn time(&self) -> &Product<RootTimestamp, T> { &self.time }

	/// Closes the input, flushing and sealing the wrapped timely input.
	pub fn close(self) { }
}

impl<T: Timestamp+Clone, D: Data, R: Diff> Drop for InputSession<T, D, R> {
	fn drop(&mut self) {
		self.flush();
	}
}
<|endoftext|>