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

use crate::core::global;
use crate::core::libtx::proof;
use crate::keychain::{ExtKeychain, Identifier, Keychain};
use crate::libwallet::internal::{keys, updater};
use crate::libwallet::types::*;
use crate::libwallet::Error;
use crate::util::secp::{key::SecretKey, pedersen};
use std::collections::HashMap;

/// Utility struct for return values from below
#[derive(Clone)]
struct OutputResult {
	///
	pub commit: pedersen::Commitment,
	///
	pub key_id: Identifier,
	///
	pub n_child: u32,
	///
	pub value: u64,
	///
	pub height: u64,
	///
	pub lock_height: u64,
	///
	pub is_coinbase: bool,
	///
	pub blinding: SecretKey,
}

fn identify_utxo_outputs<T, C, K>(
	wallet: &mut T,
	outputs: Vec<(pedersen::Commitment, pedersen::RangeProof, bool, u64)>,
) -> Result<Vec<OutputResult>, Error>
where
	T: WalletBackend<C, K>,
	C: NodeClient,
	K: Keychain,
{
	let mut wallet_outputs: Vec<OutputResult> = Vec::new();

	warn!(
		"Scanning {} outputs in the current Grin utxo set",
		outputs.len(),
	);

	for output in outputs.iter() {
		let (commit, proof, is_coinbase, height) = output;
		// attempt to unwind message from the RP and get a value
		// will fail if it's not ours
		let info = proof::rewind(wallet.keychain(), *commit, None, *proof)?;

		if !info.success {
			continue;
		}

		let lock_height = if *is_coinbase {
			*height + global::coinbase_maturity()
		} else {
			*height
		};

		// TODO: Output paths are always going to be length 3 for now, but easy enough to grind
		// through to find the right path if required later
		let key_id = Identifier::from_serialized_path(3u8, &info.message.as_bytes());

		info!(
			"Output found: {:?}, amount: {:?}, parent_key_id: {:?}",
			commit, info.value, key_id
		);

		wallet_outputs.push(OutputResult {
			commit: *commit,
			key_id: key_id.clone(),
			n_child: key_id.to_path().last_path_index(),
			value: info.value,
			height: *height,
			lock_height: lock_height,
			is_coinbase: *is_coinbase,
			blinding: info.blinding,
		});
	}
	Ok(wallet_outputs)
}

fn collect_chain_outputs<T, C, K>(wallet: &mut T) -> Result<Vec<OutputResult>, Error>
where
	T: WalletBackend<C, K>,
	C: NodeClient,
	K: Keychain,
{
	let batch_size = 1000;
	let mut start_index = 1;
	let mut result_vec: Vec<OutputResult> = vec![];
	loop {
		let (highest_index, last_retrieved_index, outputs) = wallet
			.w2n_client()
			.get_outputs_by_pmmr_index(start_index, batch_size)?;
		warn!(
			"Checking {} outputs, up to index {}. (Highest index: {})",
			outputs.len(),
			highest_index,
			last_retrieved_index,
		);

		result_vec.append(&mut identify_utxo_outputs(wallet, outputs.clone())?);

		if highest_index == last_retrieved_index {
			break;
		}
		start_index = last_retrieved_index + 1;
	}
	Ok(result_vec)
}

///
fn restore_missing_output<T, C, K>(
	wallet: &mut T,
	output: OutputResult,
	found_parents: &mut HashMap<Identifier, u32>,
) -> Result<(), Error>
where
	T: WalletBackend<C, K>,
	C: NodeClient,
	K: Keychain,
{
	let mut batch = wallet.batch()?;

	let parent_key_id = output.key_id.parent_path();
	if !found_parents.contains_key(&parent_key_id) {
		found_parents.insert(parent_key_id.clone(), 0);
	}

	let log_id = batch.next_tx_log_id(&parent_key_id)?;
	let entry_type = match output.is_coinbase {
		true => TxLogEntryType::ConfirmedCoinbase,
		false => TxLogEntryType::TxReceived,
	};

	let mut t = TxLogEntry::new(parent_key_id.clone(), entry_type, log_id);
	t.confirmed = true;
	t.amount_credited = output.value;
	t.num_outputs = 1;
	t.update_confirmation_ts();
	batch.save_tx_log_entry(t, &parent_key_id)?;

	let _ = batch.save(OutputData {
		root_key_id: parent_key_id.clone(),
		key_id: output.key_id,
		n_child: output.n_child,
		value: output.value,
		status: OutputStatus::Unspent,
		height: output.height,
		lock_height: output.lock_height,
		is_coinbase: output.is_coinbase,
		tx_log_entry: Some(log_id),
	});

	let max_child_index = found_parents.get(&parent_key_id).unwrap().clone();
	if output.n_child >= max_child_index {
		found_parents.insert(parent_key_id.clone(), output.n_child);
	}

	batch.commit()?;
	Ok(())
}

///
fn cancel_tx_log_entry<T, C, K>(wallet: &mut T, output: &OutputData) -> Result<(), Error>
where
	T: WalletBackend<C, K>,
	C: NodeClient,
	K: Keychain,
{
	let parent_key_id = output.key_id.parent_path();
	let updated_tx_entry = if output.tx_log_entry.is_some() {
		let entries = updater::retrieve_txs(
			wallet,
			output.tx_log_entry.clone(),
			None,
			Some(&parent_key_id),
		)?;
		if entries.len() > 0 {
			let mut entry = entries[0].clone();
			match entry.tx_type {
				TxLogEntryType::TxSent => entry.tx_type = TxLogEntryType::TxSentCancelled,
				TxLogEntryType::TxReceived => entry.tx_type = TxLogEntryType::TxReceivedCancelled,
				_ => {}
			}
			Some(entry)
		} else {
			None
		}
	} else {
		None
	};
	let mut batch = wallet.batch()?;
	if let Some(t) = updated_tx_entry {
		batch.save_tx_log_entry(t, &parent_key_id)?;
	}
	batch.commit()?;
	Ok(())
}

/// Check / repair wallet contents
/// assume wallet contents have been freshly updated with contents
/// of latest block
pub fn check_repair<T, C, K>(wallet: &mut T) -> Result<(), Error>
where
	T: WalletBackend<C, K>,
	C: NodeClient,
	K: Keychain,
{
	// First, get a definitive list of outputs we own from the chain
	warn!("Starting wallet check.");
	let chain_outs = collect_chain_outputs(wallet)?;
	warn!(
		"Identified {} wallet_outputs as belonging to this wallet",
		chain_outs.len(),
	);

	// Now, get all outputs owned by this wallet (regardless of account)
	let wallet_outputs = {
		let res = updater::retrieve_outputs(&mut *wallet, true, None, None)?;
		res
	};

	// check all definitive outputs exist in the wallet outputs
	let mut missing_outs = vec![];
	let mut accidental_spend_outs = vec![];
	let mut locked_outs = vec![];
	for deffo in chain_outs.into_iter() {
		let matched_out = wallet_outputs.iter().find(|wo| wo.0.key_id == deffo.key_id);
		match matched_out {
			Some(s) => {
				if s.0.status == OutputStatus::Spent {
					accidental_spend_outs.push((s.0.clone(), deffo.clone()));
				}
				if s.0.status == OutputStatus::Locked {
					locked_outs.push((s.0.clone(), deffo));
				}
			}
			None => missing_outs.push(deffo),
		}
	}

	// mark problem spent outputs as unspent (confirmed against a short-lived fork, for example)
	for m in accidental_spend_outs.into_iter() {
		let mut o = m.0;
		warn!(
			"Output for {} with ID {} ({:?}) marked as spent but exists in UTXO set. \
			 Marking unspent and cancelling any associated transaction log entries.",
			o.value, o.key_id, m.1.commit,
		);
		o.status = OutputStatus::Unspent;
		// any transactions associated with this should be cancelled
		cancel_tx_log_entry(wallet, &o)?;
		let mut batch = wallet.batch()?;
		batch.save(o)?;
		batch.commit()?;
	}

	let mut found_parents: HashMap<Identifier, u32> = HashMap::new();

	// Restore missing outputs, adding transaction for it back to the log
	for m in missing_outs.into_iter() {
		warn!(
			"Confirmed output for {} with ID {} ({:?}) exists in UTXO set but not in wallet. \
			 Restoring.",
			m.value, m.key_id, m.commit,
		);
		restore_missing_output(wallet, m, &mut found_parents)?;
	}

	for m in locked_outs.into_iter() {
		let mut o = m.0;
		warn!(
			"Confirmed output for {} with ID {} ({:?}) exists in UTXO set and is locked. \
			 Unlocking and cancelling associated transaction log entries.",
			o.value, o.key_id, m.1.commit,
		);
		o.status = OutputStatus::Unspent;
		cancel_tx_log_entry(wallet, &o)?;
		let mut batch = wallet.batch()?;
		batch.save(o)?;
		batch.commit()?;
	}

	Ok(())
}

/// Restore a wallet
pub fn restore<T, C, K>(wallet: &mut T) -> Result<(), Error>
where
	T: WalletBackend<C, K>,
	C: NodeClient,
	K: Keychain,
{
	// Don't proceed if wallet_data has anything in it
	let is_empty = wallet.iter().next().is_none();
	if !is_empty {
		error!("Not restoring. Please back up and remove existing db directory first.");
		return Ok(());
	}

	warn!("Starting restore.");

	let result_vec = collect_chain_outputs(wallet)?;

	warn!(
		"Identified {} wallet_outputs as belonging to this wallet",
		result_vec.len(),
	);

	let mut found_parents: HashMap<Identifier, u32> = HashMap::new();
	// Now save what we have

	for output in result_vec {
		restore_missing_output(wallet, output, &mut found_parents)?;
	}

	// restore labels, account paths and child derivation indices
	let label_base = "account";
	let mut index = 1;
	for (path, max_child_index) in found_parents.iter() {
		if *path == ExtKeychain::derive_key_id(2, 0, 0, 0, 0) {
			//default path already exists
			continue;
		}
		let label = format!("{}_{}", label_base, index);
		keys::set_acct_path(wallet, &label, path)?;
		index = index + 1;
		{
			let mut batch = wallet.batch()?;
			batch.save_child_index(path, max_child_index + 1)?;
		}
	}
	Ok(())
}
<commit_msg>Delete Unconfirmed outputs as part of wallet check_repair (#2261)<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//! Functions to restore a wallet's outputs from just the master seed

use crate::core::global;
use crate::core::libtx::proof;
use crate::keychain::{ExtKeychain, Identifier, Keychain};
use crate::libwallet::internal::{keys, updater};
use crate::libwallet::types::*;
use crate::libwallet::Error;
use crate::util::secp::{key::SecretKey, pedersen};
use std::collections::HashMap;

/// Utility struct for return values from below
#[derive(Clone)]
struct OutputResult {
	///
	pub commit: pedersen::Commitment,
	///
	pub key_id: Identifier,
	///
	pub n_child: u32,
	///
	pub value: u64,
	///
	pub height: u64,
	///
	pub lock_height: u64,
	///
	pub is_coinbase: bool,
	///
	pub blinding: SecretKey,
}

fn identify_utxo_outputs<T, C, K>(
	wallet: &mut T,
	outputs: Vec<(pedersen::Commitment, pedersen::RangeProof, bool, u64)>,
) -> Result<Vec<OutputResult>, Error>
where
	T: WalletBackend<C, K>,
	C: NodeClient,
	K: Keychain,
{
	let mut wallet_outputs: Vec<OutputResult> = Vec::new();

	warn!(
		"Scanning {} outputs in the current Grin utxo set",
		outputs.len(),
	);

	for output in outputs.iter() {
		let (commit, proof, is_coinbase, height) = output;
		// attempt to unwind message from the RP and get a value
		// will fail if it's not ours
		let info = proof::rewind(wallet.keychain(), *commit, None, *proof)?;

		if !info.success {
			continue;
		}

		let lock_height = if *is_coinbase {
			*height + global::coinbase_maturity()
		} else {
			*height
		};

		// TODO: Output paths are always going to be length 3 for now, but easy enough to grind
		// through to find the right path if required later
		let key_id = Identifier::from_serialized_path(3u8, &info.message.as_bytes());

		info!(
			"Output found: {:?}, amount: {:?}, parent_key_id: {:?}",
			commit, info.value, key_id
		);

		wallet_outputs.push(OutputResult {
			commit: *commit,
			key_id: key_id.clone(),
			n_child: key_id.to_path().last_path_index(),
			value: info.value,
			height: *height,
			lock_height: lock_height,
			is_coinbase: *is_coinbase,
			blinding: info.blinding,
		});
	}
	Ok(wallet_outputs)
}

fn collect_chain_outputs<T, C, K>(wallet: &mut T) -> Result<Vec<OutputResult>, Error>
where
	T: WalletBackend<C, K>,
	C: NodeClient,
	K: Keychain,
{
	let batch_size = 1000;
	let mut start_index = 1;
	let mut result_vec: Vec<OutputResult> = vec![];
	loop {
		let (highest_index, last_retrieved_index, outputs) = wallet
			.w2n_client()
			.get_outputs_by_pmmr_index(start_index, batch_size)?;
		warn!(
			"Checking {} outputs, up to index {}. (Highest index: {})",
			outputs.len(),
			highest_index,
			last_retrieved_index,
		);

		result_vec.append(&mut identify_utxo_outputs(wallet, outputs.clone())?);

		if highest_index == last_retrieved_index {
			break;
		}
		start_index = last_retrieved_index + 1;
	}
	Ok(result_vec)
}

///
fn restore_missing_output<T, C, K>(
	wallet: &mut T,
	output: OutputResult,
	found_parents: &mut HashMap<Identifier, u32>,
) -> Result<(), Error>
where
	T: WalletBackend<C, K>,
	C: NodeClient,
	K: Keychain,
{
	let mut batch = wallet.batch()?;

	let parent_key_id = output.key_id.parent_path();
	if !found_parents.contains_key(&parent_key_id) {
		found_parents.insert(parent_key_id.clone(), 0);
	}

	let log_id = batch.next_tx_log_id(&parent_key_id)?;
	let entry_type = match output.is_coinbase {
		true => TxLogEntryType::ConfirmedCoinbase,
		false => TxLogEntryType::TxReceived,
	};

	let mut t = TxLogEntry::new(parent_key_id.clone(), entry_type, log_id);
	t.confirmed = true;
	t.amount_credited = output.value;
	t.num_outputs = 1;
	t.update_confirmation_ts();
	batch.save_tx_log_entry(t, &parent_key_id)?;

	let _ = batch.save(OutputData {
		root_key_id: parent_key_id.clone(),
		key_id: output.key_id,
		n_child: output.n_child,
		value: output.value,
		status: OutputStatus::Unspent,
		height: output.height,
		lock_height: output.lock_height,
		is_coinbase: output.is_coinbase,
		tx_log_entry: Some(log_id),
	});

	let max_child_index = found_parents.get(&parent_key_id).unwrap().clone();
	if output.n_child >= max_child_index {
		found_parents.insert(parent_key_id.clone(), output.n_child);
	}

	batch.commit()?;
	Ok(())
}

///
fn cancel_tx_log_entry<T, C, K>(wallet: &mut T, output: &OutputData) -> Result<(), Error>
where
	T: WalletBackend<C, K>,
	C: NodeClient,
	K: Keychain,
{
	let parent_key_id = output.key_id.parent_path();
	let updated_tx_entry = if output.tx_log_entry.is_some() {
		let entries = updater::retrieve_txs(
			wallet,
			output.tx_log_entry.clone(),
			None,
			Some(&parent_key_id),
		)?;
		if entries.len() > 0 {
			let mut entry = entries[0].clone();
			match entry.tx_type {
				TxLogEntryType::TxSent => entry.tx_type = TxLogEntryType::TxSentCancelled,
				TxLogEntryType::TxReceived => entry.tx_type = TxLogEntryType::TxReceivedCancelled,
				_ => {}
			}
			Some(entry)
		} else {
			None
		}
	} else {
		None
	};
	let mut batch = wallet.batch()?;
	if let Some(t) = updated_tx_entry {
		batch.save_tx_log_entry(t, &parent_key_id)?;
	}
	batch.commit()?;
	Ok(())
}

/// Check / repair wallet contents
/// assume wallet contents have been freshly updated with contents
/// of latest block
pub fn check_repair<T, C, K>(wallet: &mut T) -> Result<(), Error>
where
	T: WalletBackend<C, K>,
	C: NodeClient,
	K: Keychain,
{
	// First, get a definitive list of outputs we own from the chain
	warn!("Starting wallet check.");
	let chain_outs = collect_chain_outputs(wallet)?;
	warn!(
		"Identified {} wallet_outputs as belonging to this wallet",
		chain_outs.len(),
	);

	// Now, get all outputs owned by this wallet (regardless of account)
	let wallet_outputs = {
		let res = updater::retrieve_outputs(&mut *wallet, true, None, None)?;
		res
	};

	// check all definitive outputs exist in the wallet outputs
	let mut missing_outs = vec![];
	let mut accidental_spend_outs = vec![];
	let mut locked_outs = vec![];
	for deffo in chain_outs.into_iter() {
		let matched_out = wallet_outputs.iter().find(|wo| wo.0.key_id == deffo.key_id);
		match matched_out {
			Some(s) => {
				if s.0.status == OutputStatus::Spent {
					accidental_spend_outs.push((s.0.clone(), deffo.clone()));
				}
				if s.0.status == OutputStatus::Locked {
					locked_outs.push((s.0.clone(), deffo.clone()));
				}
			}
			None => missing_outs.push(deffo),
		}
	}

	// mark problem spent outputs as unspent (confirmed against a short-lived fork, for example)
	for m in accidental_spend_outs.into_iter() {
		let mut o = m.0;
		warn!(
			"Output for {} with ID {} ({:?}) marked as spent but exists in UTXO set. \
			 Marking unspent and cancelling any associated transaction log entries.",
			o.value, o.key_id, m.1.commit,
		);
		o.status = OutputStatus::Unspent;
		// any transactions associated with this should be cancelled
		cancel_tx_log_entry(wallet, &o)?;
		let mut batch = wallet.batch()?;
		batch.save(o)?;
		batch.commit()?;
	}

	let mut found_parents: HashMap<Identifier, u32> = HashMap::new();

	// Restore missing outputs, adding transaction for it back to the log
	for m in missing_outs.into_iter() {
		warn!(
			"Confirmed output for {} with ID {} ({:?}) exists in UTXO set but not in wallet. \
			 Restoring.",
			m.value, m.key_id, m.commit,
		);
		restore_missing_output(wallet, m, &mut found_parents)?;
	}

	// Unlock locked outputs
	for m in locked_outs.into_iter() {
		let mut o = m.0;
		warn!(
			"Confirmed output for {} with ID {} ({:?}) exists in UTXO set and is locked. \
			 Unlocking and cancelling associated transaction log entries.",
			o.value, o.key_id, m.1.commit,
		);
		o.status = OutputStatus::Unspent;
		cancel_tx_log_entry(wallet, &o)?;
		let mut batch = wallet.batch()?;
		batch.save(o)?;
		batch.commit()?;
	}

	let unconfirmed_outs: Vec<&(OutputData, pedersen::Commitment)> = wallet_outputs
		.iter()
		.filter(|o| o.0.status == OutputStatus::Unconfirmed)
		.collect();
	// Delete unconfirmed outputs
	for m in unconfirmed_outs.into_iter() {
		let o = m.0.clone();
		warn!(
			"Unconfirmed output for {} with ID {} ({:?}) not in UTXO set. \
			 Deleting and cancelling associated transaction log entries.",
			o.value, o.key_id, m.1,
		);
		cancel_tx_log_entry(wallet, &o)?;
		let mut batch = wallet.batch()?;
		batch.delete(&o.key_id)?;
		batch.commit()?;
	}

	Ok(())
}

/// Restore a wallet
pub fn restore<T, C, K>(wallet: &mut T) -> Result<(), Error>
where
	T: WalletBackend<C, K>,
	C: NodeClient,
	K: Keychain,
{
	// Don't proceed if wallet_data has anything in it
	let is_empty = wallet.iter().next().is_none();
	if !is_empty {
		error!("Not restoring. Please back up and remove existing db directory first.");
		return Ok(());
	}

	warn!("Starting restore.");

	let result_vec = collect_chain_outputs(wallet)?;

	warn!(
		"Identified {} wallet_outputs as belonging to this wallet",
		result_vec.len(),
	);

	let mut found_parents: HashMap<Identifier, u32> = HashMap::new();
	// Now save what we have

	for output in result_vec {
		restore_missing_output(wallet, output, &mut found_parents)?;
	}

	// restore labels, account paths and child derivation indices
	let label_base = "account";
	let mut index = 1;
	for (path, max_child_index) in found_parents.iter() {
		if *path == ExtKeychain::derive_key_id(2, 0, 0, 0, 0) {
			//default path already exists
			continue;
		}
		let label = format!("{}_{}", label_base, index);
		keys::set_acct_path(wallet, &label, path)?;
		index = index + 1;
		{
			let mut batch = wallet.batch()?;
			batch.save_child_index(path, max_child_index + 1)?;
		}
	}
	Ok(())
}
<|endoftext|>


<commit_before>#![feature(generators)]

// run-pass

use std::mem::size_of_val;

fn main() {
    // Generator taking a `Copy`able resume arg.
    let gen_copy = |mut x: usize| {
        loop {
            drop(x);
            x = yield;
        }
    };

    // Generator taking a non-`Copy` resume arg.
    let gen_move = |mut x: Box<usize>| {
        loop {
            drop(x);
            x = yield;
        }
    };

    // Neither of these generators have the resume arg live across the `yield`, so they should be
    // 4 Bytes in size (only storing the discriminant)
    assert_eq!(size_of_val(&gen_copy), 4);
    assert_eq!(size_of_val(&gen_move), 4);
}
<commit_msg>Fix rebase fallout<commit_after>#![feature(generators)]

// run-pass

use std::mem::size_of_val;

fn main() {
    // Generator taking a `Copy`able resume arg.
    let gen_copy = |mut x: usize| {
        loop {
            drop(x);
            x = yield;
        }
    };

    // Generator taking a non-`Copy` resume arg.
    let gen_move = |mut x: Box<usize>| {
        loop {
            drop(x);
            x = yield;
        }
    };

    // Neither of these generators have the resume arg live across the `yield`, so they should be
    // 4 Bytes in size (only storing the discriminant)
    assert_eq!(size_of_val(&gen_copy), 1);
    assert_eq!(size_of_val(&gen_move), 1);
}
<|endoftext|>


<commit_before>use walkdir::WalkDir;
use std::path::PathBuf;
use std::{env, fs};

fn main() {
    // The src directory (we are in src/tools/error_index_generator)
    // Note that we could skip one of the .. but this ensures we at least loosely find the right
    // directory.
    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
    let dest = out_dir.join("error_codes.rs");
    let mut idx = 0;
    for entry in WalkDir::new("../../../src") {
        let entry = entry.unwrap();
        if entry.file_name() == "error_codes.rs" {
            println!("cargo:rerun-if-changed={}", entry.path().to_str().unwrap());
            let file = fs::read_to_string(entry.path()).unwrap()
                .replace("syntax::register_diagnostics!", "register_diagnostics!");
            let contents = format!("(|| {{\n{}\n}})()", file);

            fs::write(&out_dir.join(&format!("error_{}.rs", idx)), &contents).unwrap();

            idx += 1;
        }
    }

    let mut all = String::new();
    all.push_str(r###"
fn register_all() -> Vec<(&'static str, Option<&'static str>)> {
    let mut long_codes: Vec<(&'static str, Option<&'static str>)> = Vec::new();
    macro_rules! register_diagnostics {
        ($($ecode:ident: $message:expr,)*) => (
            register_diagnostics!{$($ecode:$message,)* ;}
        );

        ($($ecode:ident: $message:expr,)* ; $($code:ident,)*) => (
            $(
                {long_codes.extend([
                    (stringify!($ecode), Some($message)),
                ].iter());}
            )*
            $(
                {long_codes.extend([
                    stringify!($code),
                ].iter().cloned().map(|s| (s, None)).collect::<Vec<_>>());}
            )*
        )
    }
"###);
    for idx in 0..idx {
        all.push_str(&format!(r#"include!(concat!(env!("OUT_DIR"), "/error_{}.rs"));"#, idx));
        all.push_str("\n");
    }
    all.push_str("\nlong_codes\n");
    all.push_str("}\n");

    fs::write(&dest, all).unwrap();
}
<commit_msg>Fix error codes index generation<commit_after>use walkdir::WalkDir;
use std::path::PathBuf;
use std::{env, fs};

fn main() {
    // The src directory (we are in src/tools/error_index_generator)
    // Note that we could skip one of the .. but this ensures we at least loosely find the right
    // directory.
    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
    let dest = out_dir.join("error_codes.rs");
    let mut idx = 0;
    for entry in WalkDir::new("../../../src") {
        let entry = entry.unwrap();
        if entry.file_name() == "error_codes.rs" {
            println!("cargo:rerun-if-changed={}", entry.path().to_str().unwrap());
            let file = fs::read_to_string(entry.path()).unwrap()
                .replace("crate::register_diagnostics!", "register_diagnostics!")
                .replace(": include_str!(\"./",
                         ": include_str!(\"../../../../../../../../src/librustc_error_codes/");
            let contents = format!("(|| {{\n{}\n}})()", file);

            fs::write(&out_dir.join(&format!("error_{}.rs", idx)), &contents).unwrap();

            idx += 1;
        }
    }

    let mut all = String::new();
    all.push_str(r###"
fn register_all() -> Vec<(&'static str, Option<&'static str>)> {
    let mut long_codes: Vec<(&'static str, Option<&'static str>)> = Vec::new();
    macro_rules! register_diagnostics {
        ($($ecode:ident: $message:expr,)*) => (
            register_diagnostics!{$($ecode:$message,)* ;}
        );

        ($($ecode:ident: $message:expr,)* ; $($code:ident,)*) => (
            $(
                {long_codes.extend([
                    (stringify!($ecode), Some($message)),
                ].iter());}
            )*
            $(
                {long_codes.extend([
                    stringify!($code),
                ].iter().cloned().map(|s| (s, None)).collect::<Vec<_>>());}
            )*
        )
    }
"###);
    for idx in 0..idx {
        all.push_str(&format!(r#"include!(concat!(env!("OUT_DIR"), "/error_{}.rs"));"#, idx));
        all.push_str("\n");
    }
    all.push_str("\nlong_codes\n");
    all.push_str("}\n");

    fs::write(&dest, all).unwrap();
}
<|endoftext|>


<commit_before>/// This is a small server which is intended to run inside of an emulator or
/// on a remote test device. This server pairs with the `remote-test-client`
/// program in this repository. The `remote-test-client` connects to this
/// server over a TCP socket and performs work such as:
///
/// 1. Pushing shared libraries to the server
/// 2. Running tests through the server
///
/// The server supports running tests concurrently and also supports tests
/// themselves having support libraries. All data over the TCP sockets is in a
/// basically custom format suiting our needs.

use std::cmp;
use std::env;
use std::fs::{self, File, Permissions};
use std::io::prelude::*;
use std::io::{self, BufReader};
use std::net::{TcpListener, TcpStream};
use std::os::unix::prelude::*;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::str;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;

macro_rules! t {
    ($e:expr) => (match $e {
        Ok(e) => e,
        Err(e) => panic!("{} failed with {}", stringify!($e), e),
    })
}

static TEST: AtomicUsize = AtomicUsize::new(0);

struct Config {
    pub remote: bool,
    pub verbose: bool,
}

impl Config {
    pub fn default() -> Config {
        Config {
            remote: false,
            verbose: false,
        }
    }

    pub fn parse_args() -> Config {
        let mut config = Config::default();

        let args = env::args().skip(1);
        for argument in args {
            match &argument[..] {
                "remote" => {
                    config.remote = true;
                },
                "verbose" | "-v" => {
                    config.verbose = true;
                }
                arg => panic!("unknown argument: {}", arg),
            }
        }

        config
    }
}

fn main() {
    println!("starting test server");

    let config = Config::parse_args();

    let bind_addr = if cfg!(target_os = "android") || config.remote {
        "0.0.0.0:12345"
    } else {
        "10.0.2.15:12345"
    };

    let (listener, work) = if cfg!(target_os = "android") {
        (t!(TcpListener::bind(bind_addr)), "/data/tmp/work")
    } else {
        (t!(TcpListener::bind(bind_addr)), "/tmp/work")
    };
    println!("listening!");

    let work = Path::new(work);
    t!(fs::create_dir_all(work));

    let lock = Arc::new(Mutex::new(()));

    for socket in listener.incoming() {
        let mut socket = t!(socket);
        let mut buf = [0; 4];
        if socket.read_exact(&mut buf).is_err() {
            continue
        }
        if &buf[..] == b"ping" {
            t!(socket.write_all(b"pong"));
        } else if &buf[..] == b"push" {
            handle_push(socket, work);
        } else if &buf[..] == b"run " {
            let lock = lock.clone();
            thread::spawn(move || handle_run(socket, work, &lock));
        } else {
            panic!("unknown command {:?}", buf);
        }
    }
}

fn handle_push(socket: TcpStream, work: &Path) {
    let mut reader = BufReader::new(socket);
    recv(&work, &mut reader);

    let mut socket = reader.into_inner();
    t!(socket.write_all(b"ack "));
}

struct RemoveOnDrop<'a> {
    inner: &'a Path,
}

impl<'a> Drop for RemoveOnDrop<'a> {
    fn drop(&mut self) {
        t!(fs::remove_dir_all(self.inner));
    }
}

fn handle_run(socket: TcpStream, work: &Path, lock: &Mutex<()>) {
    let mut arg = Vec::new();
    let mut reader = BufReader::new(socket);

    // Allocate ourselves a directory that we'll delete when we're done to save
    // space.
    let n = TEST.fetch_add(1, Ordering::SeqCst);
    let path = work.join(format!("test{}", n));
    t!(fs::create_dir(&path));
    let _a = RemoveOnDrop { inner: &path };

    // First up we'll get a list of arguments delimited with 0 bytes. An empty
    // argument means that we're done.
    let mut args = Vec::new();
    while t!(reader.read_until(0, &mut arg)) > 1 {
        args.push(t!(str::from_utf8(&arg[..arg.len() - 1])).to_string());
        arg.truncate(0);
    }

    // Next we'll get a bunch of env vars in pairs delimited by 0s as well
    let mut env = Vec::new();
    arg.truncate(0);
    while t!(reader.read_until(0, &mut arg)) > 1 {
        let key_len = arg.len() - 1;
        let val_len = t!(reader.read_until(0, &mut arg)) - 1;
        {
            let key = &arg[..key_len];
            let val = &arg[key_len + 1..][..val_len];
            let key = t!(str::from_utf8(key)).to_string();
            let val = t!(str::from_utf8(val)).to_string();
            env.push((key, val));
        }
        arg.truncate(0);
    }

    // The section of code from here down to where we drop the lock is going to
    // be a critical section for us. On Linux you can't execute a file which is
    // open somewhere for writing, as you'll receive the error "text file busy".
    // Now here we never have the text file open for writing when we spawn it,
    // so why do we still need a critical section?
    //
    // Process spawning first involves a `fork` on Unix, which clones all file
    // descriptors into the child process. This means that it's possible for us
    // to open the file for writing (as we're downloading it), then some other
    // thread forks, then we close the file and try to exec. At that point the
    // other thread created a child process with the file open for writing, and
    // we attempt to execute it, so we get an error.
    //
    // This race is resolve by ensuring that only one thread can writ ethe file
    // and spawn a child process at once. Kinda an unfortunate solution, but we
    // don't have many other choices with this sort of setup!
    //
    // In any case the lock is acquired here, before we start writing any files.
    // It's then dropped just after we spawn the child. That way we don't lock
    // the execution of the child, just the creation of its files.
    let lock = lock.lock();

    // Next there's a list of dynamic libraries preceded by their filenames.
    while t!(reader.fill_buf())[0] != 0 {
        recv(&path, &mut reader);
    }
    assert_eq!(t!(reader.read(&mut [0])), 1);

    // Finally we'll get the binary. The other end will tell us how big the
    // binary is and then we'll download it all to the exe path we calculated
    // earlier.
    let exe = recv(&path, &mut reader);

    let mut cmd = Command::new(&exe);
    for arg in args {
        cmd.arg(arg);
    }
    for (k, v) in env {
        cmd.env(k, v);
    }

    // Support libraries were uploaded to `work` earlier, so make sure that's
    // in `LD_LIBRARY_PATH`. Also include our own current dir which may have
    // had some libs uploaded.
    cmd.env("LD_LIBRARY_PATH",
            format!("{}:{}", work.display(), path.display()));

    // Spawn the child and ferry over stdout/stderr to the socket in a framed
    // fashion (poor man's style)
    let mut child = t!(cmd.stdin(Stdio::null())
                          .stdout(Stdio::piped())
                          .stderr(Stdio::piped())
                          .spawn());
    drop(lock);
    let mut stdout = child.stdout.take().unwrap();
    let mut stderr = child.stderr.take().unwrap();
    let socket = Arc::new(Mutex::new(reader.into_inner()));
    let socket2 = socket.clone();
    let thread = thread::spawn(move || my_copy(&mut stdout, 0, &*socket2));
    my_copy(&mut stderr, 1, &*socket);
    thread.join().unwrap();

    // Finally send over the exit status.
    let status = t!(child.wait());
    let (which, code) = match status.code() {
        Some(n) => (0, n),
        None => (1, status.signal().unwrap()),
    };
    t!(socket.lock().unwrap().write_all(&[
        which,
        (code >> 24) as u8,
        (code >> 16) as u8,
        (code >>  8) as u8,
        (code >>  0) as u8,
    ]));
}

fn recv<B: BufRead>(dir: &Path, io: &mut B) -> PathBuf {
    let mut filename = Vec::new();
    t!(io.read_until(0, &mut filename));

    // We've got some tests with *really* long names. We try to name the test
    // executable the same on the target as it is on the host to aid with
    // debugging, but the targets we're emulating are often more restrictive
    // than the hosts as well.
    //
    // To ensure we can run a maximum number of tests without modifications we
    // just arbitrarily truncate the filename to 50 bytes. That should
    // hopefully allow us to still identify what's running while staying under
    // the filesystem limits.
    let len = cmp::min(filename.len() - 1, 50);
    let dst = dir.join(t!(str::from_utf8(&filename[..len])));
    let amt = read_u32(io) as u64;
    t!(io::copy(&mut io.take(amt),
                &mut t!(File::create(&dst))));
    t!(fs::set_permissions(&dst, Permissions::from_mode(0o755)));
    dst
}

fn my_copy(src: &mut dyn Read, which: u8, dst: &Mutex<dyn Write>) {
    let mut b = [0; 1024];
    loop {
        let n = t!(src.read(&mut b));
        let mut dst = dst.lock().unwrap();
        t!(dst.write_all(&[
            which,
            (n >> 24) as u8,
            (n >> 16) as u8,
            (n >>  8) as u8,
            (n >>  0) as u8,
        ]));
        if n > 0 {
            t!(dst.write_all(&b[..n]));
        } else {
            break
        }
    }
}

fn read_u32(r: &mut dyn Read) -> u32 {
    let mut len = [0; 4];
    t!(r.read_exact(&mut len));
    ((len[0] as u32) << 24) |
    ((len[1] as u32) << 16) |
    ((len[2] as u32) <<  8) |
    ((len[3] as u32) <<  0)
}
<commit_msg>Fix typo<commit_after>/// This is a small server which is intended to run inside of an emulator or
/// on a remote test device. This server pairs with the `remote-test-client`
/// program in this repository. The `remote-test-client` connects to this
/// server over a TCP socket and performs work such as:
///
/// 1. Pushing shared libraries to the server
/// 2. Running tests through the server
///
/// The server supports running tests concurrently and also supports tests
/// themselves having support libraries. All data over the TCP sockets is in a
/// basically custom format suiting our needs.

use std::cmp;
use std::env;
use std::fs::{self, File, Permissions};
use std::io::prelude::*;
use std::io::{self, BufReader};
use std::net::{TcpListener, TcpStream};
use std::os::unix::prelude::*;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::str;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;

macro_rules! t {
    ($e:expr) => (match $e {
        Ok(e) => e,
        Err(e) => panic!("{} failed with {}", stringify!($e), e),
    })
}

static TEST: AtomicUsize = AtomicUsize::new(0);

struct Config {
    pub remote: bool,
    pub verbose: bool,
}

impl Config {
    pub fn default() -> Config {
        Config {
            remote: false,
            verbose: false,
        }
    }

    pub fn parse_args() -> Config {
        let mut config = Config::default();

        let args = env::args().skip(1);
        for argument in args {
            match &argument[..] {
                "remote" => {
                    config.remote = true;
                },
                "verbose" | "-v" => {
                    config.verbose = true;
                }
                arg => panic!("unknown argument: {}", arg),
            }
        }

        config
    }
}

fn main() {
    println!("starting test server");

    let config = Config::parse_args();

    let bind_addr = if cfg!(target_os = "android") || config.remote {
        "0.0.0.0:12345"
    } else {
        "10.0.2.15:12345"
    };

    let (listener, work) = if cfg!(target_os = "android") {
        (t!(TcpListener::bind(bind_addr)), "/data/tmp/work")
    } else {
        (t!(TcpListener::bind(bind_addr)), "/tmp/work")
    };
    println!("listening!");

    let work = Path::new(work);
    t!(fs::create_dir_all(work));

    let lock = Arc::new(Mutex::new(()));

    for socket in listener.incoming() {
        let mut socket = t!(socket);
        let mut buf = [0; 4];
        if socket.read_exact(&mut buf).is_err() {
            continue
        }
        if &buf[..] == b"ping" {
            t!(socket.write_all(b"pong"));
        } else if &buf[..] == b"push" {
            handle_push(socket, work);
        } else if &buf[..] == b"run " {
            let lock = lock.clone();
            thread::spawn(move || handle_run(socket, work, &lock));
        } else {
            panic!("unknown command {:?}", buf);
        }
    }
}

fn handle_push(socket: TcpStream, work: &Path) {
    let mut reader = BufReader::new(socket);
    recv(&work, &mut reader);

    let mut socket = reader.into_inner();
    t!(socket.write_all(b"ack "));
}

struct RemoveOnDrop<'a> {
    inner: &'a Path,
}

impl<'a> Drop for RemoveOnDrop<'a> {
    fn drop(&mut self) {
        t!(fs::remove_dir_all(self.inner));
    }
}

fn handle_run(socket: TcpStream, work: &Path, lock: &Mutex<()>) {
    let mut arg = Vec::new();
    let mut reader = BufReader::new(socket);

    // Allocate ourselves a directory that we'll delete when we're done to save
    // space.
    let n = TEST.fetch_add(1, Ordering::SeqCst);
    let path = work.join(format!("test{}", n));
    t!(fs::create_dir(&path));
    let _a = RemoveOnDrop { inner: &path };

    // First up we'll get a list of arguments delimited with 0 bytes. An empty
    // argument means that we're done.
    let mut args = Vec::new();
    while t!(reader.read_until(0, &mut arg)) > 1 {
        args.push(t!(str::from_utf8(&arg[..arg.len() - 1])).to_string());
        arg.truncate(0);
    }

    // Next we'll get a bunch of env vars in pairs delimited by 0s as well
    let mut env = Vec::new();
    arg.truncate(0);
    while t!(reader.read_until(0, &mut arg)) > 1 {
        let key_len = arg.len() - 1;
        let val_len = t!(reader.read_until(0, &mut arg)) - 1;
        {
            let key = &arg[..key_len];
            let val = &arg[key_len + 1..][..val_len];
            let key = t!(str::from_utf8(key)).to_string();
            let val = t!(str::from_utf8(val)).to_string();
            env.push((key, val));
        }
        arg.truncate(0);
    }

    // The section of code from here down to where we drop the lock is going to
    // be a critical section for us. On Linux you can't execute a file which is
    // open somewhere for writing, as you'll receive the error "text file busy".
    // Now here we never have the text file open for writing when we spawn it,
    // so why do we still need a critical section?
    //
    // Process spawning first involves a `fork` on Unix, which clones all file
    // descriptors into the child process. This means that it's possible for us
    // to open the file for writing (as we're downloading it), then some other
    // thread forks, then we close the file and try to exec. At that point the
    // other thread created a child process with the file open for writing, and
    // we attempt to execute it, so we get an error.
    //
    // This race is resolve by ensuring that only one thread can write the file
    // and spawn a child process at once. Kinda an unfortunate solution, but we
    // don't have many other choices with this sort of setup!
    //
    // In any case the lock is acquired here, before we start writing any files.
    // It's then dropped just after we spawn the child. That way we don't lock
    // the execution of the child, just the creation of its files.
    let lock = lock.lock();

    // Next there's a list of dynamic libraries preceded by their filenames.
    while t!(reader.fill_buf())[0] != 0 {
        recv(&path, &mut reader);
    }
    assert_eq!(t!(reader.read(&mut [0])), 1);

    // Finally we'll get the binary. The other end will tell us how big the
    // binary is and then we'll download it all to the exe path we calculated
    // earlier.
    let exe = recv(&path, &mut reader);

    let mut cmd = Command::new(&exe);
    for arg in args {
        cmd.arg(arg);
    }
    for (k, v) in env {
        cmd.env(k, v);
    }

    // Support libraries were uploaded to `work` earlier, so make sure that's
    // in `LD_LIBRARY_PATH`. Also include our own current dir which may have
    // had some libs uploaded.
    cmd.env("LD_LIBRARY_PATH",
            format!("{}:{}", work.display(), path.display()));

    // Spawn the child and ferry over stdout/stderr to the socket in a framed
    // fashion (poor man's style)
    let mut child = t!(cmd.stdin(Stdio::null())
                          .stdout(Stdio::piped())
                          .stderr(Stdio::piped())
                          .spawn());
    drop(lock);
    let mut stdout = child.stdout.take().unwrap();
    let mut stderr = child.stderr.take().unwrap();
    let socket = Arc::new(Mutex::new(reader.into_inner()));
    let socket2 = socket.clone();
    let thread = thread::spawn(move || my_copy(&mut stdout, 0, &*socket2));
    my_copy(&mut stderr, 1, &*socket);
    thread.join().unwrap();

    // Finally send over the exit status.
    let status = t!(child.wait());
    let (which, code) = match status.code() {
        Some(n) => (0, n),
        None => (1, status.signal().unwrap()),
    };
    t!(socket.lock().unwrap().write_all(&[
        which,
        (code >> 24) as u8,
        (code >> 16) as u8,
        (code >>  8) as u8,
        (code >>  0) as u8,
    ]));
}

fn recv<B: BufRead>(dir: &Path, io: &mut B) -> PathBuf {
    let mut filename = Vec::new();
    t!(io.read_until(0, &mut filename));

    // We've got some tests with *really* long names. We try to name the test
    // executable the same on the target as it is on the host to aid with
    // debugging, but the targets we're emulating are often more restrictive
    // than the hosts as well.
    //
    // To ensure we can run a maximum number of tests without modifications we
    // just arbitrarily truncate the filename to 50 bytes. That should
    // hopefully allow us to still identify what's running while staying under
    // the filesystem limits.
    let len = cmp::min(filename.len() - 1, 50);
    let dst = dir.join(t!(str::from_utf8(&filename[..len])));
    let amt = read_u32(io) as u64;
    t!(io::copy(&mut io.take(amt),
                &mut t!(File::create(&dst))));
    t!(fs::set_permissions(&dst, Permissions::from_mode(0o755)));
    dst
}

fn my_copy(src: &mut dyn Read, which: u8, dst: &Mutex<dyn Write>) {
    let mut b = [0; 1024];
    loop {
        let n = t!(src.read(&mut b));
        let mut dst = dst.lock().unwrap();
        t!(dst.write_all(&[
            which,
            (n >> 24) as u8,
            (n >> 16) as u8,
            (n >>  8) as u8,
            (n >>  0) as u8,
        ]));
        if n > 0 {
            t!(dst.write_all(&b[..n]));
        } else {
            break
        }
    }
}

fn read_u32(r: &mut dyn Read) -> u32 {
    let mut len = [0; 4];
    t!(r.read_exact(&mut len));
    ((len[0] as u32) << 24) |
    ((len[1] as u32) << 16) |
    ((len[2] as u32) <<  8) |
    ((len[3] as u32) <<  0)
}
<|endoftext|>


<commit_before>use na::{self, Isometry2, Vector2};

use ncollide2d::query;
use ncollide2d::shape::{Compound, Cuboid, ShapeHandle};

#[test]
fn normal() {
  let c1 = Cuboid::new(Vector2::new(1.0, 1.0));
  let c2 = Cuboid::new(Vector2::new(1.0, 1.0));

  let contact = query::contact(
    &Isometry2::new(Vector2::new(10.5, 10.5), na::zero()),
    &c2,
    &Isometry2::new(Vector2::new(10.0, 10.0), na::zero()),
    &c1,
    1.0,
  );

  assert!(contact.is_some()); // ok
}

#[test]
fn compound() {
  let c1 = Cuboid::new(Vector2::new(1.0, 1.0));
  let c2 = Cuboid::new(Vector2::new(1.0, 1.0));

  let shapes = vec![(
    Isometry2::new(Vector2::new(10.0, 10.0), na::zero()),
    ShapeHandle::new(c1),
  )];
  let compound = Compound::new(shapes);

  let contact = query::contact(
    &Isometry2::new(Vector2::new(10.5, 10.5), na::zero()),
    &c2,
    &Isometry2::identity(),
    &compound,
    1.0,
  );

  assert!(contact.is_some()); // FAILED
}
<commit_msg>Add issue number for tests.<commit_after>use na::{self, Isometry2, Vector2};

use ncollide2d::query;
use ncollide2d::shape::{Compound, Cuboid, ShapeHandle};

// Issue #171.
#[test]
fn normal() {
  let c1 = Cuboid::new(Vector2::new(1.0, 1.0));
  let c2 = Cuboid::new(Vector2::new(1.0, 1.0));

  let contact = query::contact(
    &Isometry2::new(Vector2::new(10.5, 10.5), na::zero()),
    &c2,
    &Isometry2::new(Vector2::new(10.0, 10.0), na::zero()),
    &c1,
    1.0,
  );

  assert!(contact.is_some()); // ok
}

// Issue #171.
#[test]
fn compound() {
  let c1 = Cuboid::new(Vector2::new(1.0, 1.0));
  let c2 = Cuboid::new(Vector2::new(1.0, 1.0));

  let shapes = vec![(
    Isometry2::new(Vector2::new(10.0, 10.0), na::zero()),
    ShapeHandle::new(c1),
  )];
  let compound = Compound::new(shapes);

  let contact = query::contact(
    &Isometry2::new(Vector2::new(10.5, 10.5), na::zero()),
    &c2,
    &Isometry2::identity(),
    &compound,
    1.0,
  );

  assert!(contact.is_some()); // FAILED
}
<|endoftext|>


<commit_before>extern crate matroska;
extern crate av_format;
extern crate av_data;
extern crate circular;
extern crate nom;

use nom::{Err, Offset};
use circular::Buffer;
use std::env;
use std::io::Read;
use std::fs::File;

use matroska::ebml::ebml_header;
use matroska::elements::{segment, segment_element, SegmentElement};

fn main() {
    let mut args = env::args();
    let _ = args.next().expect("first arg is program path");
    let filename = args.next().expect("expected file path");

    //println!("filename: {}", filename);
    run(&filename).expect("should parse file correctly");
}

fn run(filename: &str) -> std::io::Result<()> {
    let mut file = File::open(filename)?;

    let capacity = 5242880;
    let mut b = Buffer::with_capacity(capacity);

    // we write into the `&mut[u8]` returned by `space()`
    let sz = file.read(b.space()).expect("should write");
    b.fill(sz);
    //eprintln!("write {:#?}", sz);

    let length = {
        let res = ebml_header(b.data());
        if let Ok((remaining, header)) = res {
            //eprintln!("parsed header: {:#?}", h);
            println!("+ EBML head");
            println!("|+ EBML version: {}", header.version);
            println!("|+ EBML read version: {}", header.read_version);
            println!("|+ EBML maximum ID length: {}", header.max_id_length);
            println!("|+ EBML maximum size length: {}", header.max_size_length);
            println!("|+ Doc type: {}", header.doc_type);
            println!("|+ Doc type read version: {}", header.doc_type_read_version);

            b.data().offset(remaining)
        } else {
            panic!("couldn't parse header");
        }
    };

    let mut _consumed = length;
    //println!("consumed {} bytes", length);
    b.consume(length);

    let length = {
        let res = segment(b.data());
        if let Ok((remaining, segment)) = res {
            //eprintln!("parsed segment: {:#?}", h);
            println!("+ Segment, size {}", segment.1.unwrap_or(0));

            b.data().offset(remaining)
        } else {
            panic!("couldn't parse header");
        }
    };

    //eprintln!("consumed {} bytes", length);
    b.consume(length);

    // handle first elements
    let mut seek_head = None;
    let mut info = None;
    let mut tracks = None;

    loop {
        if seek_head.is_some() && info.is_some() && tracks.is_some() {
            break;
        }

        // refill the buffer
        let sz = file.read(b.space()).expect("should read");
        b.fill(sz);
        /*
        eprintln!(
            "refill: {} more bytes, available data: {} bytes, consumed: {} bytes",
            sz,
            b.available_data(),
            consumed
        );
        */

        // if there's no more available data in the buffer after a write, that means we reached
        // the end of the file
        if b.available_data() == 0 {
            panic!("no more data to read or parse, stopping the reading loop");
        }

        let offset = {
            let (i, element) = match segment_element(b.data()) {
                Ok((i, o)) => (i, o),
                Err(Err::Error(e)) |
                Err(Err::Failure(e)) => panic!("failed parsing: {:?}", e),
                Err(Err::Incomplete(_)) => continue,
            };

            match element {
                SegmentElement::SeekHead(s) => {
                    //eprintln!("got seek head: {:#?}", s);
                    println!("|+ Seek head");
                    if seek_head.is_some() {
                        panic!("already got a SeekHead element");
                    } else {
                        seek_head = Some(s);
                    }
                },
                SegmentElement::Info(i) => {
                    println!("|+ Segment Information");
                    println!("| + timestamp scale: {}", i.timecode_scale);
                    println!("| + multiplexing application: {}", i.muxing_app);
                    println!("| + writing application: {}", i.writing_app);
                    println!("| + segment UID: {:?}", i.segment_uid.as_ref().unwrap_or(&Vec::new()));
                    println!("| + duration: {}s", i.duration.unwrap_or(0f64) / 1000f64);
                    if info.is_some() {
                        panic!("already got an Info element");
                    } else {
                        info = Some(i);
                    }
                },
                SegmentElement::Tracks(t) => {
                    //eprintln!("got tracks: {:#?}", t);
                    println!("|+ Segment tracks");
                    for tr in t.tracks.iter() {
                      println!("| + A track");
                      println!("|  + Track number: {}", tr.track_number);
                      println!("|  + Track UID: {}", tr.track_uid);
                      println!("|  + Track type: {}", tr.track_type);
                      println!("|  + Lacing flag: {}", tr.flag_lacing.unwrap_or(0));
                      println!("|  + Default flag: {}", tr.flag_default.unwrap_or(0));
                      println!("|  + Language: {}", tr.language.as_ref().unwrap_or(&"".to_string()));
                      println!("|  + Codec ID: {}", tr.codec_id);
                      println!("|  + Codec private: length {}", tr.codec_private.as_ref().map(|v| v.len()).unwrap_or(0));

                      if let Some(ref v) = tr.video {
                        println!("|  + Video track");
                        println!("|    + Pixel width: {}", v.pixel_width);
                        println!("|    + Pixel height: {}", v.pixel_height);
                        if let Some(inter) = v.flag_interlaced {
                          println!("|    + Interlaced: {}", inter);
                        }
                        if let Some(width) = v.display_width {
                          println!("|    + Display width: {}", width);
                        }
                        if let Some(height) = v.display_height {
                          println!("|    + Display height: {}", height);
                        }
                        if let Some(unit) = v.display_unit {
                          println!("|    + Display unit: {}", unit);
                        }
                      }

                      if let Some(ref a) = tr.audio {
                        println!("|  + Audio track");
                        println!("|    + Sampling frequency: {}", a.sampling_frequency);
                        if let Some(frequency) = a.output_sampling_frequency {
                          println!("|    + Output sampling freqeuncy: {}", frequency);
                        }
                        println!("|    + Channels: {}", a.channels);
                        if let Some(ref channel_positions) = a.channel_positions {
                          println!("|    + Channel position: {:?}", channel_positions);
                        }
                        if let Some(bit_depth) = a.bit_depth {
                          println!("|    + Bit depth: {}", bit_depth);
                        }
                      }
                    }
                    if tracks.is_some() {
                        panic!("already got a Tracks element");
                    } else {
                        tracks = Some(t);
                    }
                },
                /*SegmentElement::Cluster(c) => {
                    println!("|+ Cluster");
                    eprintln!("got a cluster: {:#?}", c);
                },*/
                SegmentElement::Void => {
                    println!("|+ EbmlVoid");
                },
                el => {
                    panic!("got unexpected element: {:#?}", el);
                }
            }

            b.data().offset(i)
        };
        _consumed += offset;
        //eprintln!("consumed {} bytes", length);
        b.consume(offset);
    }

    loop {
        if b.available_space() == 0 {
          b.shift();
          if b.available_space() == 0 {
            println!("buffer is already full,  cannot refill");
            break;
          }
        }

        // refill the buffer
        let sz = file.read(b.space()).expect("should read");
        b.fill(sz);

        /*
        eprintln!(
            "refill: {} more bytes, available data: {} bytes, consumed: {} bytes",
            sz,
            b.available_data(),
            consumed
        );
        */

        // if there's no more available data in the buffer after a write, that means we reached
        // the end of the file
        if b.available_data() == 0 {
            //panic!("no more data to read or parse, stopping the reading loop");
            break;
        }

        let offset = {
            let (i, element) = match segment_element(b.data()) {
                Ok((i, o)) => (i, o),
                Err(Err::Error(e)) |
                Err(Err::Failure(e)) => panic!("failed parsing: {:?}", e),
                Err(Err::Incomplete(_)) => continue,
            };

            match element {
                SegmentElement::SeekHead(_) | SegmentElement::Info(_) | SegmentElement::Tracks(_) => {
                  panic!("unexpected seek head, info or tracks element: {:?}", element);
                },
                SegmentElement::Cluster(c) => {
                  println!("|+ Cluster");
                  //eprintln!("got a cluster: {:#?}", c);
                },
                SegmentElement::Void => {
                  println!("|+ EbmlVoid");
                },
                SegmentElement::Tags(_) => {
                  println!("|+ Tags");
                }
                SegmentElement::Cues(_) => {
                  println!("|+ Cues");
                }
                el => {
                    panic!("got unexpected element: {:#?}", el);
                }
            }

            b.data().offset(i)
        };
        _consumed += offset;
        //eprintln!("consumed {} bytes", length);
        b.consume(offset);
    }

    Ok(())
}
<commit_msg>more info in matroska_info<commit_after>extern crate matroska;
extern crate av_format;
extern crate av_data;
extern crate circular;
extern crate nom;

use nom::{Err, Offset};
use circular::Buffer;
use std::env;
use std::io::Read;
use std::fs::File;

use matroska::ebml::ebml_header;
use matroska::elements::{segment, segment_element, SegmentElement};

fn main() {
    let mut args = env::args();
    let _ = args.next().expect("first arg is program path");
    let filename = args.next().expect("expected file path");

    //println!("filename: {}", filename);
    run(&filename).expect("should parse file correctly");
}

fn run(filename: &str) -> std::io::Result<()> {
    let mut file = File::open(filename)?;

    let capacity = 5242880;
    let mut b = Buffer::with_capacity(capacity);

    // we write into the `&mut[u8]` returned by `space()`
    let sz = file.read(b.space()).expect("should write");
    b.fill(sz);
    //eprintln!("write {:#?}", sz);

    let length = {
        let res = ebml_header(b.data());
        if let Ok((remaining, header)) = res {
            //eprintln!("parsed header: {:#?}", h);
            println!("+ EBML head");
            println!("|+ EBML version: {}", header.version);
            println!("|+ EBML read version: {}", header.read_version);
            println!("|+ EBML maximum ID length: {}", header.max_id_length);
            println!("|+ EBML maximum size length: {}", header.max_size_length);
            println!("|+ Doc type: {}", header.doc_type);
            println!("|+ Doc type read version: {}", header.doc_type_read_version);

            b.data().offset(remaining)
        } else {
            panic!("couldn't parse header");
        }
    };

    let mut _consumed = length;
    //println!("consumed {} bytes", length);
    b.consume(length);

    let length = {
        let res = segment(b.data());
        if let Ok((remaining, segment)) = res {
            //eprintln!("parsed segment: {:#?}", h);
            println!("+ Segment, size {}", segment.1.unwrap_or(0));

            b.data().offset(remaining)
        } else {
            panic!("couldn't parse header");
        }
    };

    //eprintln!("consumed {} bytes", length);
    b.consume(length);

    // handle first elements
    let mut seek_head = None;
    let mut info = None;
    let mut tracks = None;

    loop {
        if seek_head.is_some() && info.is_some() && tracks.is_some() {
            break;
        }

        println!("available space: {}", b.available_space());
        if b.available_space() == 0 {
          b.shift();
          if b.available_space() == 0 {
            println!("buffer is already full,  cannot refill");
            break;
          }
        }

        // refill the buffer
        let sz = file.read(b.space()).expect("should read");
        b.fill(sz);
        /*
        eprintln!(
            "refill: {} more bytes, available data: {} bytes, consumed: {} bytes",
            sz,
            b.available_data(),
            consumed
        );
        */

        // if there's no more available data in the buffer after a write, that means we reached
        // the end of the file
        if b.available_data() == 0 {
            panic!("no more data to read or parse, stopping the reading loop");
        }

        let offset = {
            let (i, element) = match segment_element(b.data()) {
                Ok((i, o)) => (i, o),
                Err(Err::Error(e)) |
                Err(Err::Failure(e)) => panic!("failed parsing: {:?}", e),
                Err(Err::Incomplete(_)) => continue,
            };

            match element {
                SegmentElement::SeekHead(s) => {
                    //eprintln!("got seek head: {:#?}", s);
                    println!("|+ Seek head");
                    if seek_head.is_some() {
                        panic!("already got a SeekHead element");
                    } else {
                        seek_head = Some(s);
                    }
                },
                SegmentElement::Info(i) => {
                    println!("|+ Segment Information");
                    println!("| + timestamp scale: {}", i.timecode_scale);
                    println!("| + multiplexing application: {}", i.muxing_app);
                    println!("| + writing application: {}", i.writing_app);
                    println!("| + segment UID: {:?}", i.segment_uid.as_ref().unwrap_or(&Vec::new()));
                    println!("| + duration: {}s", i.duration.unwrap_or(0f64) / 1000f64);
                    if info.is_some() {
                        panic!("already got an Info element");
                    } else {
                        info = Some(i);
                    }
                },
                SegmentElement::Tracks(t) => {
                    //eprintln!("got tracks: {:#?}", t);
                    println!("|+ Segment tracks");
                    for tr in t.tracks.iter() {
                      println!("| + A track");
                      println!("|  + Track number: {}", tr.track_number);
                      println!("|  + Track UID: {}", tr.track_uid);
                      println!("|  + Track type: {}", tr.track_type);
                      println!("|  + Lacing flag: {}", tr.flag_lacing.unwrap_or(0));
                      println!("|  + Default flag: {}", tr.flag_default.unwrap_or(0));
                      println!("|  + Language: {}", tr.language.as_ref().unwrap_or(&"".to_string()));
                      println!("|  + Codec ID: {}", tr.codec_id);
                      println!("|  + Codec private: length {}", tr.codec_private.as_ref().map(|v| v.len()).unwrap_or(0));

                      if let Some(ref v) = tr.video {
                        println!("|  + Video track");
                        println!("|    + Pixel width: {}", v.pixel_width);
                        println!("|    + Pixel height: {}", v.pixel_height);
                        if let Some(inter) = v.flag_interlaced {
                          println!("|    + Interlaced: {}", inter);
                        }
                        if let Some(width) = v.display_width {
                          println!("|    + Display width: {}", width);
                        }
                        if let Some(height) = v.display_height {
                          println!("|    + Display height: {}", height);
                        }
                        if let Some(unit) = v.display_unit {
                          println!("|    + Display unit: {}", unit);
                        }
                      }

                      if let Some(ref a) = tr.audio {
                        println!("|  + Audio track");
                        println!("|    + Sampling frequency: {}", a.sampling_frequency);
                        if let Some(frequency) = a.output_sampling_frequency {
                          println!("|    + Output sampling freqeuncy: {}", frequency);
                        }
                        println!("|    + Channels: {}", a.channels);
                        if let Some(ref channel_positions) = a.channel_positions {
                          println!("|    + Channel position: {:?}", channel_positions);
                        }
                        if let Some(bit_depth) = a.bit_depth {
                          println!("|    + Bit depth: {}", bit_depth);
                        }
                      }
                    }
                    if tracks.is_some() {
                        panic!("already got a Tracks element");
                    } else {
                        tracks = Some(t);
                    }
                },
                /*SegmentElement::Cluster(c) => {
                    println!("|+ Cluster");
                    eprintln!("got a cluster: {:#?}", c);
                },*/
                SegmentElement::Void => {
                    println!("|+ EbmlVoid");
                },
                el => {
                    panic!("got unexpected element: {:#?}", el);
                }
            }

            b.data().offset(i)
        };
        _consumed += offset;
        //eprintln!("consumed {} bytes", length);
        b.consume(offset);
    }

    loop {
        if b.available_space() == 0 {
          b.shift();
          if b.available_space() == 0 {
            println!("buffer is already full,  cannot refill");
            break;
          }
        }

        // refill the buffer
        let sz = file.read(b.space()).expect("should read");
        b.fill(sz);

        /*
        eprintln!(
            "refill: {} more bytes, available data: {} bytes, consumed: {} bytes",
            sz,
            b.available_data(),
            consumed
        );
        */

        // if there's no more available data in the buffer after a write, that means we reached
        // the end of the file
        if b.available_data() == 0 {
            //panic!("no more data to read or parse, stopping the reading loop");
            break;
        }

        let offset = {
            let (i, element) = match segment_element(b.data()) {
                Ok((i, o)) => (i, o),
                Err(Err::Error(e)) |
                Err(Err::Failure(e)) => panic!("failed parsing: {:?}", e),
                Err(Err::Incomplete(_)) => continue,
            };

            match element {
                SegmentElement::SeekHead(_) | SegmentElement::Info(_) | SegmentElement::Tracks(_) => {
                  panic!("unexpected seek head, info or tracks element: {:?}", element);
                },
                SegmentElement::Cluster(c) => {
                  println!("|+ Cluster");
                  println!("|+   timecode: {}", c.timecode);
                  println!("|+   silent_tracks: {:?}", c.silent_tracks);
                  println!("|+   position: {:?}", c.position);
                  println!("|+   prev_size: {:?}", c.prev_size);
                  println!("|+   simple_block: {} elements", c.simple_block.len());
                  println!("|+   block_group: {} elements", c.block_group.len());
                  println!("|+   encrypted_block: {:?} bytes", c.encrypted_block.as_ref().map(|s| s.len()));
                  //eprintln!("got a cluster: {:#?}", c);
                },
                SegmentElement::Void => {
                  println!("|+ EbmlVoid");
                },
                SegmentElement::Tags(_) => {
                  println!("|+ Tags");
                }
                SegmentElement::Cues(_) => {
                  println!("|+ Cues");
                }
                el => {
                    panic!("got unexpected element: {:#?}", el);
                }
            }

            b.data().offset(i)
        };
        _consumed += offset;
        //eprintln!("consumed {} bytes", length);
        b.consume(offset);
    }

    Ok(())
}
<|endoftext|>


<commit_before>use buffer::{Position, Range, range};

/// A more concise expression for ranges spanning complete lines.
pub struct LineRange {
    start: usize,
    end:   usize,
}

impl LineRange {
    pub fn start(&self) -> usize {
        self.start
    }

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

    /// Converts the line range to a regular, zero-offset range.
    ///
    /// # Examples
    ///
    /// ```
    /// use scribe::buffer::{LineRange, Position, Range, range};
    ///
    /// // Builder a line range.
    /// let line_range = LineRange{ start: 10, end: 14 };
    ///
    /// // Ensure that the resulting range is a zero-based equivalent.
    /// assert_eq!(line_range.to_range(), range::new(
    ///     Position{ line: 10, offset: 0 },
    ///     Position{ line: 14, offset:0 }
    /// ));
    /// ```
    pub fn to_range(&self) -> Range {
        range::new(
            Position{ line: self.start, offset: 0 },
            Position{ line: self.end, offset:0 }
        )
    }

pub fn new(start: usize, end: usize) -> LineRange {
    if start < end {
        LineRange{ start: start, end: end }
    } else {
        LineRange{ start: end, end: start }
    }
}

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

    #[test]
    fn new_does_not_swap_values_if_end_does_not_precede_start() {
        let range = new(0, 1);

        assert_eq!(range.start(), 0);
        assert_eq!(range.end(), 1);
    }

    #[test]
    fn new_swaps_start_and_end_when_end_precedes_start() {
        let range = new(1, 0);

        assert_eq!(range.start(), 0);
        assert_eq!(range.end(), 1);
    }
}
<commit_msg>Implement includes method for LineRange type.<commit_after>use buffer::{Position, Range, range};

/// A more concise expression for ranges spanning complete lines.
pub struct LineRange {
    start: usize,
    end:   usize,
}

impl LineRange {
    pub fn start(&self) -> usize {
        self.start
    }

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

    /// Converts the line range to a regular, zero-offset range.
    ///
    /// # Examples
    ///
    /// ```
    /// use scribe::buffer::{line_range, Position, range};
    ///
    /// // Builder a line range.
    /// let line_range = line_range::new(10, 14);
    ///
    /// // Ensure that the resulting range is a zero-based equivalent.
    /// assert_eq!(line_range.to_range(), range::new(
    ///     Position{ line: 10, offset: 0 },
    ///     Position{ line: 14, offset:0 }
    /// ));
    /// ```
    pub fn to_range(&self) -> Range {
        range::new(
            Position{ line: self.start, offset: 0 },
            Position{ line: self.end, offset:0 }
        )
    }

    /// Whether or not the line range includes the specified line.
    /// The range is exclusive, such that its ending line is not included.
    ///
    /// # Examples
    ///
    /// ```
    /// use scribe::buffer::line_range;
    ///
    /// // Builder a line range.
    /// let line_range = line_range::new(10, 14);
    ///
    /// assert!(line_range.includes(11));
    /// assert!(!line_range.includes(14));
    /// ```
    pub fn includes(&self, line: usize) -> bool {
        line >= self.start() && line < self.end()
    }
}

pub fn new(start: usize, end: usize) -> LineRange {
    if start < end {
        LineRange{ start: start, end: end }
    } else {
        LineRange{ start: end, end: start }
    }
}

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

    #[test]
    fn new_does_not_swap_values_if_end_does_not_precede_start() {
        let range = new(0, 1);

        assert_eq!(range.start(), 0);
        assert_eq!(range.end(), 1);
    }

    #[test]
    fn new_swaps_start_and_end_when_end_precedes_start() {
        let range = new(1, 0);

        assert_eq!(range.start(), 0);
        assert_eq!(range.end(), 1);
    }
}
<|endoftext|>


<commit_before>extern crate rand;
extern crate timely;
extern crate timely_sort;
extern crate differential_dataflow;
extern crate vec_map;
extern crate graph_map;

use rand::{Rng, SeedableRng, StdRng};

use timely::dataflow::*;
use timely::dataflow::operators::*;

use differential_dataflow::{Collection, AsCollection, Data};
use differential_dataflow::operators::*;
use differential_dataflow::collection::LeastUpperBound;

use graph_map::GraphMMap;

type Node = u32;
type Edge = (Node, Node);

fn main() {

    // define a new computational scope, in which to run reach
    timely::execute_from_args(std::env::args(), move |computation| {

        // contruct iterative transaction dataflow, attach reachability query to graph.
        let (mut raw, mut trans, mut query, mut roots, query_probe, roots_probe) = computation.scoped(|scope| {

            // transactional updates to edge collection.
            let (raw_input, raw) = scope.new_input();
            let raw = raw.as_collection();

            // transactional updates to edge collection.
            let (trans_input, trans) = scope.new_input();
            let trans = trans.as_collection();

            // graph contains committed edges, will persist across epochs.
            let (handle, graph) = scope.loop_variable(u64::max_value(), 1);
            let graph = graph.as_collection();

            let aborts = transactional(&trans, &graph);

            // restrict writes to those that commit.
            let writes = trans.filter(|x| !x.1)
                              .map(|x| (x.0, x.2))
                              .antijoin(&aborts)
                              .map(|(_,edge)| edge);

            // define the graph as itself plus any committed changes.
            let graph = raw.concat(&writes);

            graph.inner.connect_loop(handle);

            // read queries against edge collection
            let (query_input, query) = scope.new_input();
            let query = query.as_collection();

            // query the graph, return a probe.
            let query_probe = graph.semijoin_u(&query)
                                   .probe().0;

            // issue reachability queries from various roots
            let (roots_input, roots) = scope.new_input();
            let roots = roots.as_collection();

            // do reach on the graph from the roots, report counts at each distance.
            let result = reach(&graph, &roots);
            let roots_probe = result//.map(|(_,l)| l)
                                    //.consolidate_by(|&x| x)
                                    .probe().0;

            (raw_input, trans_input, query_input, roots_input, query_probe, roots_probe)
        });

        let filename = std::env::args().nth(1).unwrap();
        let graph = GraphMMap::new(&filename);

        if computation.index() == 0 {
            for node in 0 .. graph.nodes() {
                for &edge in graph.edges(node) {
                    raw.send(((node as u32, edge as u32), 1));
                }
            }
        }
        raw.close();

        let next = trans.epoch() + 2;
        trans.advance_to(next);
        query.advance_to(next);
        roots.advance_to(next);

        // do computation for a bit, until we see outputs from each.
        while query_probe.lt(&trans.time()) || roots_probe.lt(&trans.time()) {
            computation.step();
        }

        let seed: &[_] = &[1, 2, 3, 4];
        let mut rng: StdRng = SeedableRng::from_seed(seed);    // rng for edge additions

        let block: u32 = std::env::args().nth(2).unwrap().parse().unwrap();

        let mut reads = Vec::with_capacity(1000);
        while reads.len() < reads.capacity() {
            let mut buffer = vec![];
            for i in 0 .. block as usize {
                let key = rng.gen_range(0, graph.nodes() as u32);
                if i % computation.peers() == computation.index() {
                    buffer.push(key);
                }
            }
            reads.push(buffer);
        }

        let timer = ::std::time::Instant::now();

        for buffer in &reads {
            for &entry in buffer {
                query.send((entry, 1));
            }

            let next = trans.epoch() + 1;
            trans.advance_to(next);
            query.advance_to(next);
            roots.advance_to(next);

            // do computation for a bit, until we see outputs from each.
            while query_probe.lt(&trans.time()) || roots_probe.lt(&trans.time()) {
                computation.step();
            }
        }

        query.close();
        println!("query elapsed: {:?} for 1,000 x {}", timer.elapsed(), block);

        let mut travs = Vec::with_capacity(100);
        while travs.len() < travs.capacity() {
            let mut temp = vec![];
            for _ in 0 .. block {
                temp.push(rng.gen_range(0, graph.nodes() as u32));
            }
            travs.push(temp);
        }

        let mut latencies = vec![];

        let timer = ::std::time::Instant::now();

        for buffer in &travs {
            let inner_timer = ::std::time::Instant::now();
            for &src in buffer {
                roots.send((src, 1));
            }

            let next = trans.epoch() + 1;
            trans.advance_to(next);
            // query.advance_to(next);
            roots.advance_to(next);

            // do computation for a bit, until we see outputs from each.
            while query_probe.lt(&trans.time()) || roots_probe.lt(&trans.time()) {
                computation.step();
            }

            let elapsed = inner_timer.elapsed();
            latencies.push(elapsed.as_secs() * 1000000000 + elapsed.subsec_nanos() as u64);
        }

        roots.close();
        println!("travs elapsed: {:?} for 1,000 x {}", timer.elapsed(), block);

        if computation.index() == 0 {
            latencies[50 ..].sort();
            for &x in latencies[50..].iter() {
                println!("{}", (x as f64) / 1000000000.0f64);
            }
        }

        let mut writes = Vec::with_capacity(1000);
        let index = computation.index();
        let peers = computation.peers();
        let mut t_id = index;

        while writes.len() < writes.capacity() {
            let mut edge_set = vec![];
            for _ in 0 .. block {
                edge_set.push((rng.gen_range(0, graph.nodes() as u32), rng.gen_range(0, graph.nodes() as u32)));
            }
            let mut buffer = vec![];
            for i in 0 .. block as usize {
                let edge1 = edge_set[rng.gen_range(0, edge_set.len())];
                let edge2 = edge_set[rng.gen_range(0, edge_set.len())];
                if i % computation.peers() == computation.index() {
                    buffer.push((t_id, true, edge1));
                    buffer.push((t_id, false, edge2));
                }
                t_id += peers;
            }
            writes.push(buffer);
        }

        let timer = ::std::time::Instant::now();

        for buffer in &writes {
            for &entry in buffer {
                trans.send((entry, 1));
            }

            let next = trans.epoch() + 1;
            trans.advance_to(next);
            // query.advance_to(next);
            // roots.advance_to(next);

            // do computation for a bit, until we see outputs from each.
            while query_probe.lt(&trans.time()) || roots_probe.lt(&trans.time()) {
                computation.step();
            }
        }

        println!("trans elapsed: {:?} for 1,000 x {}", timer.elapsed(), block);

    }).unwrap();
}

// returns pairs (n, s) indicating node n can be reached from node s.
fn reach<G>(edges: &Collection<G, Edge>, roots: &Collection<G, Node>) -> Collection<G, (Node, Node)>
where G: Scope, G::Timestamp: LeastUpperBound {

    // initialize roots as reaching themselves
    let nodes = roots.map(|x| (x, x));

    // repeatedly update reachable nodes
    nodes.iterate(|inner| {

        let edges = edges.enter(&inner.scope());
        let nodes = nodes.enter(&inner.scope());

        inner.join_map(&edges, |_k,&s,&d| (d, s))
             .concat(&nodes)
             .distinct()
     })
}

fn transactional<G, D>(trans: &Collection<G, (usize, bool, D)>, state: &Collection<G, D>) -> Collection<G, usize>
where G: Scope, D: Data+Default, G::Timestamp: LeastUpperBound {

    // we develop the set of transaction ids that must abort,
    // starting from the optimistic assumption that none must.

    trans.filter(|_| false)
         .map(|_| 0)
         .iterate(|abort| {

            // bring in transactions and initial state.
            let trans = trans.enter(&abort.scope());
            let state = state.enter(&abort.scope())
                             .map(|x| (x, 0));

            // reads -> (location, trans_id)
            let reads = trans.filter(|x| x.1 == true)
                             .map(|x| (x.2, x.0));

            // writes -> (location, trans_id), filtered by aborts
            let writes = trans.filter(|x| x.1 == false)
                              .map(|x| (x.0, x.2))
                              .antijoin_u(&abort)
                              .map(|x| (x.1, x.0));

            // cogroup may be the right way to do this,
            // but we can do a join for now.
            let lookup = writes.concat(&state)
                               .join_map(&reads, |key, &tid1, &tid2| (key.clone(), tid1, tid2))
                               .filter(|&(_, tid1, tid2)| tid1 < tid2)
                               .map(|(k, _, tid2)| (k, tid2))
                               .distinct();

            reads.map(|x| (x,()))
                 .antijoin(&lookup)
                 .map(|((_,t),())| t)
                 .distinct_u()
        })

}
<commit_msg>tidies example code<commit_after>extern crate rand;
extern crate timely;
extern crate timely_sort;
extern crate differential_dataflow;
extern crate vec_map;
extern crate graph_map;

use rand::{Rng, SeedableRng, StdRng};

use timely::dataflow::*;
use timely::dataflow::operators::*;

use differential_dataflow::{Collection, AsCollection, Data};
use differential_dataflow::operators::*;
use differential_dataflow::collection::LeastUpperBound;

use graph_map::GraphMMap;

type Node = u32;
type Edge = (Node, Node);

fn main() {

    // define a new computational scope, in which to run reach
    timely::execute_from_args(std::env::args(), move |computation| {

        // contruct iterative transaction dataflow, attach reachability query to graph.
        let (mut raw, mut trans, mut query, mut roots, query_probe, roots_probe) = computation.scoped(|scope| {

            // transactional updates to edge collection.
            let (raw_input, raw) = scope.new_input();
            let raw = raw.as_collection();

            // transactional updates to edge collection.
            let (trans_input, trans) = scope.new_input();
            let trans = trans.as_collection();

            // graph contains committed edges, will persist across epochs.
            let (handle, graph) = scope.loop_variable(u64::max_value(), 1);
            let graph = graph.as_collection();

            let aborts = transactional(&trans, &graph);

            // restrict writes to those that commit.
            let writes = trans.filter(|x| !x.1)
                              .map(|x| (x.0, x.2))
                              .antijoin(&aborts)
                              .map(|(_,edge)| edge);

            // define the graph as itself plus any committed changes.
            let graph = raw.concat(&writes);

            graph.inner.connect_loop(handle);

            // read queries against edge collection
            let (query_input, query) = scope.new_input();
            let query = query.as_collection();

            // query the graph, return a probe.
            let query_probe = graph.semijoin_u(&query)
                                   .probe().0;

            // issue reachability queries from various roots
            let (roots_input, roots) = scope.new_input();
            let roots = roots.as_collection();

            // do reach on the graph from the roots, report counts at each distance.
            let result = reach(&graph, &roots);
            let roots_probe = result.probe().0;

            (raw_input, trans_input, query_input, roots_input, query_probe, roots_probe)
        });

        let filename = std::env::args().nth(1).unwrap();
        let graph = GraphMMap::new(&filename);

        if computation.index() == 0 {
            for node in 0 .. graph.nodes() {
                for &edge in graph.edges(node) {
                    raw.send(((node as u32, edge as u32), 1));
                }
            }
        }
        raw.close();

        let next = trans.epoch() + 2;
        trans.advance_to(next);
        query.advance_to(next);
        roots.advance_to(next);

        // do computation for a bit, until we see outputs from each.
        while query_probe.lt(&trans.time()) || roots_probe.lt(&trans.time()) {
            computation.step();
        }

        let seed: &[_] = &[1, 2, 3, 4];
        let mut rng: StdRng = SeedableRng::from_seed(seed);    // rng for edge additions

        let block: u32 = std::env::args().nth(2).unwrap().parse().unwrap();

        let mut reads = Vec::with_capacity(1000);
        while reads.len() < reads.capacity() {
            let mut buffer = vec![];
            for i in 0 .. block as usize {
                let key = rng.gen_range(0, graph.nodes() as u32);
                if i % computation.peers() == computation.index() {
                    buffer.push(key);
                }
            }
            reads.push(buffer);
        }

        let timer = ::std::time::Instant::now();

        for buffer in &reads {
            for &entry in buffer {
                query.send((entry, 1));
            }

            let next = trans.epoch() + 1;
            trans.advance_to(next);
            query.advance_to(next);
            roots.advance_to(next);

            // do computation for a bit, until we see outputs from each.
            while query_probe.lt(&trans.time()) || roots_probe.lt(&trans.time()) {
                computation.step();
            }
        }

        query.close();
        println!("query elapsed: {:?} for 1,000 x {}", timer.elapsed(), block);

        let mut travs = Vec::with_capacity(1000);
        while travs.len() < travs.capacity() {
            let mut temp = vec![];
            for _ in 0 .. block {
                temp.push(rng.gen_range(0, graph.nodes() as u32));
            }
            travs.push(temp);
        }

        let mut latencies = vec![];

        let timer = ::std::time::Instant::now();

        for buffer in &travs {
            let inner_timer = ::std::time::Instant::now();
            for &src in buffer {
                roots.send((src, 1));
            }

            let next = trans.epoch() + 1;
            trans.advance_to(next);
            // query.advance_to(next);
            roots.advance_to(next);

            // do computation for a bit, until we see outputs from each.
            while query_probe.lt(&trans.time()) || roots_probe.lt(&trans.time()) {
                computation.step();
            }

            let elapsed = inner_timer.elapsed();
            latencies.push(elapsed.as_secs() * 1000000000 + elapsed.subsec_nanos() as u64);
        }

        roots.close();
        println!("travs elapsed: {:?} for 1,000 x {}", timer.elapsed(), block);

        if computation.index() == 0 {
            latencies[500 ..].sort();
            for &x in latencies[500..].iter() {
                println!("{}", (x as f64) / 1000000000.0f64);
            }
        }

        let mut writes = Vec::with_capacity(1000);
        let index = computation.index();
        let peers = computation.peers();
        let mut t_id = index;

        while writes.len() < writes.capacity() {
            let mut edge_set = vec![];
            for _ in 0 .. block {
                edge_set.push((rng.gen_range(0, graph.nodes() as u32), rng.gen_range(0, graph.nodes() as u32)));
            }
            let mut buffer = vec![];
            for i in 0 .. block as usize {
                let edge1 = edge_set[rng.gen_range(0, edge_set.len())];
                let edge2 = edge_set[rng.gen_range(0, edge_set.len())];
                if i % computation.peers() == computation.index() {
                    buffer.push((t_id, true, edge1));
                    buffer.push((t_id, false, edge2));
                }
                t_id += peers;
            }
            writes.push(buffer);
        }

        let timer = ::std::time::Instant::now();

        for buffer in &writes {
            for &entry in buffer {
                trans.send((entry, 1));
            }

            let next = trans.epoch() + 1;
            trans.advance_to(next);
            // query.advance_to(next);
            // roots.advance_to(next);

            // do computation for a bit, until we see outputs from each.
            while query_probe.lt(&trans.time()) || roots_probe.lt(&trans.time()) {
                computation.step();
            }
        }

        println!("trans elapsed: {:?} for 1,000 x {}", timer.elapsed(), block);

    }).unwrap();
}

// returns pairs (n, s) indicating node n can be reached from node s.
fn reach<G>(edges: &Collection<G, Edge>, roots: &Collection<G, Node>) -> Collection<G, (Node, Node)>
where G: Scope, G::Timestamp: LeastUpperBound {

    // initialize roots as reaching themselves
    let nodes = roots.map(|x| (x, x));

    // repeatedly update reachable nodes
    nodes.iterate(|inner| {

        let edges = edges.enter(&inner.scope());
        let nodes = nodes.enter(&inner.scope());

        inner.join_map(&edges, |_k,&s,&d| (d, s))
             .concat(&nodes)
             .distinct()
     })
}

fn transactional<G, D>(trans: &Collection<G, (usize, bool, D)>, state: &Collection<G, D>) -> Collection<G, usize>
where G: Scope, D: Data+Default, G::Timestamp: LeastUpperBound {

    // we develop the set of transaction ids that must abort,
    // starting from the optimistic assumption that none must.

    trans.filter(|_| false)
         .map(|_| 0)
         .iterate(|abort| {

            // bring in transactions and initial state.
            let trans = trans.enter(&abort.scope());
            let state = state.enter(&abort.scope())
                             .map(|x| (x, 0));

            // reads -> (location, trans_id)
            let reads = trans.filter(|x| x.1 == true)
                             .map(|x| (x.2, x.0));

            // writes -> (location, trans_id), filtered by aborts
            let writes = trans.filter(|x| x.1 == false)
                              .map(|x| (x.0, x.2))
                              .antijoin_u(&abort)
                              .map(|x| (x.1, x.0));

            // cogroup may be the right way to do this,
            // but we can do a join for now.
            let lookup = writes.concat(&state)
                               .join_map(&reads, |key, &tid1, &tid2| (key.clone(), tid1, tid2))
                               .filter(|&(_, tid1, tid2)| tid1 < tid2)
                               .map(|(k, _, tid2)| (k, tid2))
                               .distinct();

            reads.map(|x| (x,()))
                 .antijoin(&lookup)
                 .map(|((_,t),())| t)
                 .distinct_u()
        })

}
<|endoftext|>