<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{
    emitter::{
        stats::{DynamicStatsTracking, StatsAccumulator},
        wait_for_accounts_sequence,
    },
    transaction_generator::TransactionGenerator,
    EmitModeParams,
};
use aptos_logger::sample::Sampling;
use aptos_logger::{sample, sample::SampleRate, warn};
use aptos_rest_client::Client as RestClient;
use aptos_sdk::types::{transaction::SignedTransaction, vm_status::StatusCode, LocalAccount};
use core::{
    cmp::{max, min},
    result::Result::{Err, Ok},
    sync::atomic::{AtomicBool, Ordering},
    time::Duration,
};
use futures::future::join_all;
use itertools::Itertools;
use rand::seq::IteratorRandom;
use rand::Rng;
use std::sync::atomic::AtomicU64;
use std::{sync::Arc, time::Instant};
use tokio::time::sleep;

pub struct SubmissionWorker {
    pub(crate) accounts: Vec<LocalAccount>,
    client: RestClient,
    stop: Arc<AtomicBool>,
    params: EmitModeParams,
    stats: Arc<DynamicStatsTracking>,
    txn_generator: Box<dyn TransactionGenerator>,
    worker_index: usize,
    skip_latency_stats: bool,
    rng: ::rand::rngs::StdRng,
}

impl SubmissionWorker {
    pub fn new(
        accounts: Vec<LocalAccount>,
        client: RestClient,
        stop: Arc<AtomicBool>,
        params: EmitModeParams,
        stats: Arc<DynamicStatsTracking>,
        txn_generator: Box<dyn TransactionGenerator>,
        worker_index: usize,
        skip_latency_stats: bool,
        rng: ::rand::rngs::StdRng,
    ) -> Self {
        Self {
            accounts,
            client,
            stop,
            params,
            stats,
            txn_generator,
            worker_index,
            skip_latency_stats,
            rng,
        }
    }

    #[allow(clippy::collapsible_if)]
    pub(crate) async fn run(mut self) -> Vec<LocalAccount> {
        // Introduce a random jitter between, so that:
        //  - we don't hammer the rest APIs all at once.
        //  - allow for even spread for fixed TPS setup
        let start_sleep_duration = self.start_sleep_time();
        let start_time = Instant::now() + start_sleep_duration;

        self.sleep_check_done(start_sleep_duration).await;

        let wait_duration = Duration::from_millis(self.params.wait_millis);
        let mut wait_until = start_time;

        while !self.stop.load(Ordering::Relaxed) {
            let stats_clone = self.stats.clone();
            let loop_stats = stats_clone.get_cur();

            let loop_start_time = Arc::new(Instant::now());
            if wait_duration.as_secs() > 0
                && loop_start_time.duration_since(wait_until) > wait_duration
            {
                sample!(
                    SampleRate::Duration(Duration::from_secs(120)),
                    warn!(
                        "[{:?}] txn_emitter worker drifted out of sync too much: {}s",
                        self.client.path_prefix_string(),
                        loop_start_time.duration_since(wait_until).as_secs()
                    )
                );
            }
            // always add expected cycle duration, to not drift from expected pace.
            wait_until += wait_duration;

            let requests = self.gen_requests();

            let txn_expiration_time = requests
                .iter()
                .map(|txn| txn.expiration_timestamp_secs())
                .max()
                .unwrap_or(0);

            let num_requests = requests.len();
            let txn_offset_time = Arc::new(AtomicU64::new(0));

            join_all(
                requests
                    .chunks(self.params.max_submit_batch_size)
                    .map(|reqs| {
                        submit_transactions(
                            &self.client,
                            reqs,
                            loop_start_time.clone(),
                            txn_offset_time.clone(),
                            loop_stats,
                        )
                    }),
            )
            .await;

            if self.skip_latency_stats {
                // we also don't want to be stuck waiting for txn_expiration_time_secs
                // after stop is called, so we sleep until time or stop is set.
                self.sleep_check_done(Duration::from_secs(
                    self.params.txn_expiration_time_secs + 20,
                ))
                .await
            }

            self.update_stats(
                *loop_start_time,
                txn_offset_time.load(Ordering::Relaxed),
                num_requests,
                // skip latency if asked to check seq_num only once
                // even if we check more often due to stop (to not affect sampling)
                self.skip_latency_stats,
                txn_expiration_time,
                // if we don't care about latency, we can recheck less often.
                // generally, we should never need to recheck, as we wait enough time
                // before calling here, but in case of shutdown/or client we are talking
                // to being stale (having stale transaction_version), we might need to wait.
                Duration::from_millis(
                    if self.skip_latency_stats { 10 } else { 1 }
                        * self.params.check_account_sequence_sleep_millis,
                ),
                loop_stats,
            )
            .await;

            let now = Instant::now();
            if wait_until > now {
                self.sleep_check_done(wait_until - now).await;
            }
        }

        self.accounts
    }

    // returns true if it returned early
    async fn sleep_check_done(&self, duration: Duration) {
        let start_time = Instant::now();
        loop {
            sleep(Duration::from_secs(1)).await;
            if self.stop.load(Ordering::Relaxed) {
                return;
            }
            if start_time.elapsed() >= duration {
                return;
            }
        }
    }

    fn start_sleep_time(&mut self) -> Duration {
        let random_jitter_millis = if self.params.start_jitter_millis > 0 {
            self.rng.gen_range(0, self.params.start_jitter_millis)
        } else {
            0
        };
        Duration::from_millis(
            (self.params.start_offset_multiplier_millis * self.worker_index as f64) as u64
                + random_jitter_millis,
        )
    }

    /// This function assumes that num_requests == num_accounts, which is
    /// precisely how gen_requests works. If this changes, this code will
    /// need to be fixed.
    ///
    /// Note, the latency values are not accurate if --check-stats-at-end
    /// is used. There is no easy way around this accurately. As such, we
    /// don't update latency at all if that flag is set.
    async fn update_stats(
        &mut self,
        start_time: Instant,
        txn_offset_time: u64,
        num_requests: usize,
        skip_latency_stats: bool,
        txn_expiration_ts_secs: u64,
        check_account_sleep_duration: Duration,
        loop_stats: &StatsAccumulator,
    ) {
        assert_eq!(
            num_requests,
            self.params.transactions_per_account * self.accounts.len()
        );
        let (num_expired, sum_of_completion_timestamps_millis) = wait_for_accounts_sequence(
            start_time,
            &self.client,
            &mut self.accounts,
            self.params.transactions_per_account,
            txn_expiration_ts_secs,
            check_account_sleep_duration,
        )
        .await;

        let num_committed = num_requests - num_expired;

        if num_expired > 0 {
            loop_stats
                .expired
                .fetch_add(num_expired as u64, Ordering::Relaxed);
            sample!(
                SampleRate::Duration(Duration::from_secs(120)),
                warn!(
                    "[{:?}] Transactions were not committed before expiration: {:?}, for {:?}",
                    self.client.path_prefix_string(),
                    num_expired,
                    self.accounts
                        .iter()
                        .map(|a| a.address())
                        .collect::<Vec<_>>(),
                )
            );
        }

        if num_committed > 0 {
            let sum_latency = sum_of_completion_timestamps_millis
                - (txn_offset_time as u128 * num_committed as u128) / num_requests as u128;
            let avg_latency = (sum_latency / num_committed as u128) as u64;
            loop_stats
                .committed
                .fetch_add(num_committed as u64, Ordering::Relaxed);

            if !skip_latency_stats {
                loop_stats
                    .latency
                    .fetch_add(sum_latency as u64, Ordering::Relaxed);
                loop_stats
                    .latency_samples
                    .fetch_add(num_committed as u64, Ordering::Relaxed);
                loop_stats
                    .latencies
                    .record_data_point(avg_latency, num_committed as u64);
            }
        }
    }

    fn gen_requests(&mut self) -> Vec<SignedTransaction> {
        let batch_size = max(
            1,
            min(
                self.params.max_submit_batch_size / self.params.transactions_per_account,
                self.accounts.len(),
            ),
        );
        let accounts = self
            .accounts
            .iter_mut()
            .choose_multiple(&mut self.rng, batch_size);
        self.txn_generator
            .generate_transactions(accounts, self.params.transactions_per_account)
    }
}

pub async fn submit_transactions(
    client: &RestClient,
    txns: &[SignedTransaction],
    loop_start_time: Arc<Instant>,
    txn_offset_time: Arc<AtomicU64>,
    stats: &StatsAccumulator,
) {
    let cur_time = Instant::now();
    let offset = cur_time - *loop_start_time;
    txn_offset_time.fetch_add(
        txns.len() as u64 * offset.as_millis() as u64,
        Ordering::Relaxed,
    );
    stats
        .submitted
        .fetch_add(txns.len() as u64, Ordering::Relaxed);

    match client.submit_batch_bcs(txns).await {
        Err(e) => {
            stats
                .failed_submission
                .fetch_add(txns.len() as u64, Ordering::Relaxed);
            sample!(
                SampleRate::Duration(Duration::from_secs(120)),
                warn!(
                    "[{:?}] Failed to submit batch request: {:?}",
                    client.path_prefix_string(),
                    e
                )
            );
        }
        Ok(v) => {
            let failures = v.into_inner().transaction_failures;

            stats
                .failed_submission
                .fetch_add(failures.len() as u64, Ordering::Relaxed);

            sample!(SampleRate::Duration(Duration::from_secs(60)), {
                let by_error = failures
                    .iter()
                    .map(|f| {
                        f.error
                            .vm_error_code
                            .and_then(|c| StatusCode::try_from(c).ok())
                    })
                    .counts();
                if let Some(failure) = failures.first() {
                    let sender = txns[failure.transaction_index].sender();

                    let last_transactions =
                        if let Ok(account) = client.get_account_bcs(sender).await {
                            client
                                .get_account_transactions_bcs(
                                    sender,
                                    Some(account.into_inner().sequence_number() - 1),
                                    Some(5),
                                )
                                .await
                                .ok()
                                .map(|r| r.into_inner())
                        } else {
                            None
                        };

                    warn!(
                        "[{:?}] Failed to submit {} txns in a batch, first failure due to {:?}, for account {}, first asked: {}, failed seq nums: {:?}, failed error codes: {:?}, last transaction for account: {:?}",
                        client.path_prefix_string(),
                        failures.len(),
                        failure,
                        sender,
                        txns[0].sequence_number(),
                        failures.iter().map(|f| txns[f.transaction_index].sequence_number()).collect::<Vec<_>>(),
                        by_error,
                        last_transactions,
                    );
                }
            });
        }
    };
}
<commit_msg>[txn-emitter] Fix underflow when first request fails (#4567)<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{
    emitter::{
        stats::{DynamicStatsTracking, StatsAccumulator},
        wait_for_accounts_sequence,
    },
    transaction_generator::TransactionGenerator,
    EmitModeParams,
};
use aptos_logger::sample::Sampling;
use aptos_logger::{sample, sample::SampleRate, warn};
use aptos_rest_client::Client as RestClient;
use aptos_sdk::types::{transaction::SignedTransaction, vm_status::StatusCode, LocalAccount};
use core::{
    cmp::{max, min},
    result::Result::{Err, Ok},
    sync::atomic::{AtomicBool, Ordering},
    time::Duration,
};
use futures::future::join_all;
use itertools::Itertools;
use rand::seq::IteratorRandom;
use rand::Rng;
use std::sync::atomic::AtomicU64;
use std::{sync::Arc, time::Instant};
use tokio::time::sleep;

pub struct SubmissionWorker {
    pub(crate) accounts: Vec<LocalAccount>,
    client: RestClient,
    stop: Arc<AtomicBool>,
    params: EmitModeParams,
    stats: Arc<DynamicStatsTracking>,
    txn_generator: Box<dyn TransactionGenerator>,
    worker_index: usize,
    skip_latency_stats: bool,
    rng: ::rand::rngs::StdRng,
}

impl SubmissionWorker {
    pub fn new(
        accounts: Vec<LocalAccount>,
        client: RestClient,
        stop: Arc<AtomicBool>,
        params: EmitModeParams,
        stats: Arc<DynamicStatsTracking>,
        txn_generator: Box<dyn TransactionGenerator>,
        worker_index: usize,
        skip_latency_stats: bool,
        rng: ::rand::rngs::StdRng,
    ) -> Self {
        Self {
            accounts,
            client,
            stop,
            params,
            stats,
            txn_generator,
            worker_index,
            skip_latency_stats,
            rng,
        }
    }

    #[allow(clippy::collapsible_if)]
    pub(crate) async fn run(mut self) -> Vec<LocalAccount> {
        // Introduce a random jitter between, so that:
        //  - we don't hammer the rest APIs all at once.
        //  - allow for even spread for fixed TPS setup
        let start_sleep_duration = self.start_sleep_time();
        let start_time = Instant::now() + start_sleep_duration;

        self.sleep_check_done(start_sleep_duration).await;

        let wait_duration = Duration::from_millis(self.params.wait_millis);
        let mut wait_until = start_time;

        while !self.stop.load(Ordering::Relaxed) {
            let stats_clone = self.stats.clone();
            let loop_stats = stats_clone.get_cur();

            let loop_start_time = Arc::new(Instant::now());
            if wait_duration.as_secs() > 0
                && loop_start_time.duration_since(wait_until) > wait_duration
            {
                sample!(
                    SampleRate::Duration(Duration::from_secs(120)),
                    warn!(
                        "[{:?}] txn_emitter worker drifted out of sync too much: {}s",
                        self.client.path_prefix_string(),
                        loop_start_time.duration_since(wait_until).as_secs()
                    )
                );
            }
            // always add expected cycle duration, to not drift from expected pace.
            wait_until += wait_duration;

            let requests = self.gen_requests();

            let txn_expiration_time = requests
                .iter()
                .map(|txn| txn.expiration_timestamp_secs())
                .max()
                .unwrap_or(0);

            let num_requests = requests.len();
            let txn_offset_time = Arc::new(AtomicU64::new(0));

            join_all(
                requests
                    .chunks(self.params.max_submit_batch_size)
                    .map(|reqs| {
                        submit_transactions(
                            &self.client,
                            reqs,
                            loop_start_time.clone(),
                            txn_offset_time.clone(),
                            loop_stats,
                        )
                    }),
            )
            .await;

            if self.skip_latency_stats {
                // we also don't want to be stuck waiting for txn_expiration_time_secs
                // after stop is called, so we sleep until time or stop is set.
                self.sleep_check_done(Duration::from_secs(
                    self.params.txn_expiration_time_secs + 20,
                ))
                .await
            }

            self.update_stats(
                *loop_start_time,
                txn_offset_time.load(Ordering::Relaxed),
                num_requests,
                // skip latency if asked to check seq_num only once
                // even if we check more often due to stop (to not affect sampling)
                self.skip_latency_stats,
                txn_expiration_time,
                // if we don't care about latency, we can recheck less often.
                // generally, we should never need to recheck, as we wait enough time
                // before calling here, but in case of shutdown/or client we are talking
                // to being stale (having stale transaction_version), we might need to wait.
                Duration::from_millis(
                    if self.skip_latency_stats { 10 } else { 1 }
                        * self.params.check_account_sequence_sleep_millis,
                ),
                loop_stats,
            )
            .await;

            let now = Instant::now();
            if wait_until > now {
                self.sleep_check_done(wait_until - now).await;
            }
        }

        self.accounts
    }

    // returns true if it returned early
    async fn sleep_check_done(&self, duration: Duration) {
        let start_time = Instant::now();
        loop {
            sleep(Duration::from_secs(1)).await;
            if self.stop.load(Ordering::Relaxed) {
                return;
            }
            if start_time.elapsed() >= duration {
                return;
            }
        }
    }

    fn start_sleep_time(&mut self) -> Duration {
        let random_jitter_millis = if self.params.start_jitter_millis > 0 {
            self.rng.gen_range(0, self.params.start_jitter_millis)
        } else {
            0
        };
        Duration::from_millis(
            (self.params.start_offset_multiplier_millis * self.worker_index as f64) as u64
                + random_jitter_millis,
        )
    }

    /// This function assumes that num_requests == num_accounts, which is
    /// precisely how gen_requests works. If this changes, this code will
    /// need to be fixed.
    ///
    /// Note, the latency values are not accurate if --check-stats-at-end
    /// is used. There is no easy way around this accurately. As such, we
    /// don't update latency at all if that flag is set.
    async fn update_stats(
        &mut self,
        start_time: Instant,
        txn_offset_time: u64,
        num_requests: usize,
        skip_latency_stats: bool,
        txn_expiration_ts_secs: u64,
        check_account_sleep_duration: Duration,
        loop_stats: &StatsAccumulator,
    ) {
        assert_eq!(
            num_requests,
            self.params.transactions_per_account * self.accounts.len()
        );
        let (num_expired, sum_of_completion_timestamps_millis) = wait_for_accounts_sequence(
            start_time,
            &self.client,
            &mut self.accounts,
            self.params.transactions_per_account,
            txn_expiration_ts_secs,
            check_account_sleep_duration,
        )
        .await;

        let num_committed = num_requests - num_expired;

        if num_expired > 0 {
            loop_stats
                .expired
                .fetch_add(num_expired as u64, Ordering::Relaxed);
            sample!(
                SampleRate::Duration(Duration::from_secs(120)),
                warn!(
                    "[{:?}] Transactions were not committed before expiration: {:?}, for {:?}",
                    self.client.path_prefix_string(),
                    num_expired,
                    self.accounts
                        .iter()
                        .map(|a| a.address())
                        .collect::<Vec<_>>(),
                )
            );
        }

        if num_committed > 0 {
            let sum_latency = sum_of_completion_timestamps_millis
                - (txn_offset_time as u128 * num_committed as u128) / num_requests as u128;
            let avg_latency = (sum_latency / num_committed as u128) as u64;
            loop_stats
                .committed
                .fetch_add(num_committed as u64, Ordering::Relaxed);

            if !skip_latency_stats {
                loop_stats
                    .latency
                    .fetch_add(sum_latency as u64, Ordering::Relaxed);
                loop_stats
                    .latency_samples
                    .fetch_add(num_committed as u64, Ordering::Relaxed);
                loop_stats
                    .latencies
                    .record_data_point(avg_latency, num_committed as u64);
            }
        }
    }

    fn gen_requests(&mut self) -> Vec<SignedTransaction> {
        let batch_size = max(
            1,
            min(
                self.params.max_submit_batch_size / self.params.transactions_per_account,
                self.accounts.len(),
            ),
        );
        let accounts = self
            .accounts
            .iter_mut()
            .choose_multiple(&mut self.rng, batch_size);
        self.txn_generator
            .generate_transactions(accounts, self.params.transactions_per_account)
    }
}

pub async fn submit_transactions(
    client: &RestClient,
    txns: &[SignedTransaction],
    loop_start_time: Arc<Instant>,
    txn_offset_time: Arc<AtomicU64>,
    stats: &StatsAccumulator,
) {
    let cur_time = Instant::now();
    let offset = cur_time - *loop_start_time;
    txn_offset_time.fetch_add(
        txns.len() as u64 * offset.as_millis() as u64,
        Ordering::Relaxed,
    );
    stats
        .submitted
        .fetch_add(txns.len() as u64, Ordering::Relaxed);

    match client.submit_batch_bcs(txns).await {
        Err(e) => {
            stats
                .failed_submission
                .fetch_add(txns.len() as u64, Ordering::Relaxed);
            sample!(
                SampleRate::Duration(Duration::from_secs(120)),
                warn!(
                    "[{:?}] Failed to submit batch request: {:?}",
                    client.path_prefix_string(),
                    e
                )
            );
        }
        Ok(v) => {
            let failures = v.into_inner().transaction_failures;

            stats
                .failed_submission
                .fetch_add(failures.len() as u64, Ordering::Relaxed);

            sample!(SampleRate::Duration(Duration::from_secs(60)), {
                let by_error = failures
                    .iter()
                    .map(|f| {
                        f.error
                            .vm_error_code
                            .and_then(|c| StatusCode::try_from(c).ok())
                    })
                    .counts();
                if let Some(failure) = failures.first() {
                    let sender = txns[failure.transaction_index].sender();

                    let last_transactions =
                        if let Ok(account) = client.get_account_bcs(sender).await {
                            client
                                .get_account_transactions_bcs(
                                    sender,
                                    Some(account.into_inner().sequence_number().saturating_sub(1)),
                                    Some(5),
                                )
                                .await
                                .ok()
                                .map(|r| r.into_inner())
                        } else {
                            None
                        };

                    warn!(
                        "[{:?}] Failed to submit {} txns in a batch, first failure due to {:?}, for account {}, first asked: {}, failed seq nums: {:?}, failed error codes: {:?}, last transaction for account: {:?}",
                        client.path_prefix_string(),
                        failures.len(),
                        failure,
                        sender,
                        txns[0].sequence_number(),
                        failures.iter().map(|f| txns[f.transaction_index].sequence_number()).collect::<Vec<_>>(),
                        by_error,
                        last_transactions,
                    );
                }
            });
        }
    };
}
<|endoftext|>


<commit_before>extern crate regex;

use self::regex::{Error as RegexError, Regex};

use constant::{ME, WORKING_DIR_NAME};
use error::{CliError, CronError, CronErrorKind};
use lib::io::*;
use std::collections::BTreeSet;
use std::env;
use std::fmt::Display;
use std::io::Write;
use std::path::Path;
use std::process::{self, Stdio};
use std::str;

#[derive(Debug)]
pub struct Cron {
    exe    : String,
    upper  : String,
    my_area: String,
    lower  : String,
}

impl Cron {
    fn base_mark<D: Display>(keyword: &str, exec: D) -> String {
        format!(
            r#"###################################
# "{}" uses the lines {}.
# Please don't touch them and me!
###################################
{}"#,
            ME,
            keyword,
            exec,
        )
    }
    fn start_mark<S: AsRef<str>>(exe: S) -> String {
        Self::base_mark("from this", Self::patrol(exe))
    }
    fn end_mark() -> String {
        Self::base_mark("up to here", "")
    }

    fn patrol<S: AsRef<str>>(exe: S) -> String {
        format!("0 0 * * *\t{} patrol\n", exe.as_ref())
    }

    fn escape_asterisk<S: AsRef<str>>(s: S) -> String {
        let re = match Regex::new(r"\*") {
            Ok(re) => re,
            Err(_) => unreachable!("Mistake regular expression!!"),
        };
        re.replace_all(s.as_ref(), "\\*")
    }

    pub fn read() -> Result<Self, CliError> {
        let exe = try!(env::current_exe());
        let exe = try!(exe.to_str().ok_or(CronError::new(CronErrorKind::BinExistsInInvalidDir)));

        print_with_tag(Tag::Info, "Read cron");

        let result = process::Command::new("crontab")
            .arg("-l")
            .output();
        let output = try!(result);

        let contents = if output.status.success() {
            try!(str::from_utf8(&output.stdout).map_err(|_| CronError::new(CronErrorKind::InvalidCharacterCode)))
        } else {
            ""
        };

        let areas = format!(
            "^(?P<upper>(.|\n)*){}(?P<my_area>(.|\n)*){}(?P<lower>(.|\n)*)$",
            Self::escape_asterisk(Self::start_mark(exe)),
            Self::end_mark(),
        );
        let re = match Regex::new(&areas) {
            Ok(re) => re,
            Err(_) => unreachable!("Mistake the regular expression!!"),
        };

        let (upper, my_area, lower) = match re.captures(contents) {
            Some(caps) => match (caps.name("upper"), caps.name("my_area"), caps.name("lower")) {
                (Some(u), Some(m), Some(l)) => (u, m, l),
                _                           => unreachable!("Mistake regular expression!!"),
            },
            None => (contents, "", ""),
        };

        Ok(Cron {
            exe    : exe    .to_string(),
            upper  : upper  .to_string(),
            my_area: my_area.to_string(),
            lower  : lower  .to_string(),
        })
    }

    pub fn update(mut self, pairs: &[(&str, &str); 2]) -> Result<Self, CliError> {
        let current_dir = try!(Self::current_dir_string());
        self.delete(&current_dir);

        let my_new_area = pairs
            .iter()
            .fold(String::new(), |area, &(time, command)| {
                let line = format!("{}\tcd {} && {} {}\n", time, current_dir, self.exe, command);
                area + &line
            });

        self.my_area = self.my_area + &my_new_area;
        Ok(self)
    }

    fn re_for_matching_missing_dir_line<S: AsRef<str>>(&self, core: S) -> String {
        format!(r".*cd\s+{}\s+&&\s+{}.*\n", core.as_ref(), self.exe)
    }

    pub fn delete<S: AsRef<str>>(&mut self, dir: S) {
        let re = self.re_for_matching_missing_dir_line(dir);
        let re = match Regex::new(re.as_ref()) {
            Ok(re) => re,
            Err(_) => unreachable!("Wrong to use delete()!!"),
        };
        self.my_area = re.replace_all(&self.my_area, "");
    }

    pub fn discard(mut self) -> Result<Self, RegexError> {
        let re = self.re_for_matching_missing_dir_line(r"(?P<path>[^\s]+)");
        let re = match Regex::new(re.as_ref()) {
            Ok(re) => re,
            Err(_) => unreachable!("Mistake regular expression!!"),
        };
        let target_paths = re
            .captures_iter(&self.my_area)
            .map(|caps| match caps.name("path") {
                Some(p) => p.to_string(),
                None    => unreachable!("Mistake keyword for capturing!!"),
            })
            .filter(|path| !Path::new(path).join(WORKING_DIR_NAME).is_dir())
            .collect::<BTreeSet<String>>();

        if target_paths.len() != 0 {
            let re_core = target_paths
                .iter()
                .fold(String::new(), |core, path| core + "|" + path); // We want to use true "Iterator.reduce()".
            let re_core = format!("({})", re_core.trim_left_matches('|'));
            let re = self.re_for_matching_missing_dir_line(re_core);
            let re = try!(Regex::new(re.as_ref()));
            self.my_area = re.replace_all(&self.my_area, "");
        }

        Ok(self)
    }

    pub fn set(self) -> Result<(), CliError> {
        print_with_tag(Tag::Info, "Set new cron");

        let result = process::Command::new("crontab")
            .stdin(Stdio::piped())
            .spawn();
        let mut child = try!(result);

        let contents = if self.my_area_is_empty() {
            self.upper + &self.lower
        } else {
            self.upper
                + &Self::start_mark(&self.exe)
                + &self.my_area
                + &Self::end_mark()
                + &self.lower
        };
        match &mut child.stdin {
            &mut Some(ref mut stdin) => try!(stdin.write_all(contents.as_bytes())),
            &mut None                => unreachable!("Please set Stdio::piped()!!"),
        };

        let output = try!(child.wait_with_output());
        if output.status.success() {
            Ok(())
        } else {
            Err(From::from(CronError::new(CronErrorKind::FailedToWrite)))
        }
    }

    pub fn current_dir_string() -> Result<String, CronError> {
        let current_dir = match env::current_dir() {
            Ok(d)  => d,
            Err(_) => unreachable!("There is a problem for current directory!!"),
        };
        current_dir
            .to_str()
            .ok_or(CronError::new(CronErrorKind::InvalidPath))
            .map(|s| s.to_string())
    }

    fn my_area_is_empty(&self) -> bool {
        let re = match Regex::new(r"^\s*$") {
            Ok(re) => re,
            Err(_) => unreachable!("Mistake the regular expression!!"),
        };
        re.is_match(self.my_area.as_ref())
    }
}

#[test]
fn my_area_is_empty_should_return_true() {
    let empties = [
        "",
        " ",
        "\t",
        "\r",
        "\n",
        " \t\r\n",
    ];

    for empty in &empties {
        let cron = Cron {
            exe    : ""   .to_string(),
            upper  : ""   .to_string(),
            my_area: empty.to_string(),
            lower  : ""   .to_string(),
        };
        assert!(cron.my_area_is_empty());
    }
}
#[test]
fn my_area_is_empty_should_return_false() {
    let non_empties = [
        "a",
        "ab",
        " ab",
        "a b",
        "ab ",
        "\tab",
        "a\tb",
        "ab\t",
        "\rab",
        "a\rb",
        "ab\r",
        "\nab",
        "a\nb",
        "ab\n",
        " \t\r\nab",
        "a \t\r\nb",
        "ab \t\r\n",
    ];

    for non_empty in &non_empties {
        let cron = Cron {
            exe    : ""       .to_string(),
            upper  : ""       .to_string(),
            my_area: non_empty.to_string(),
            lower  : ""       .to_string(),
        };
        assert!(!cron.my_area_is_empty());
    }
}

#[test]
fn delete_should_success() {
    fn wrap_by_extra_content<S: AsRef<str>>(target: S) -> String {
        let extra_content = "when cd path_to_extra && path_to_bin command\n";
        format!("{}{}{}", extra_content, target.as_ref(), extra_content)
    }

    let mut cron = Cron {
        exe    : "path_to_exe".to_string(),
        upper  : "upper"      .to_string(),
        my_area: "middle"     .to_string(),
        lower  : "lower"      .to_string(),
    };

    let dir = "path_to_dir";

    let areas = vec![
        ("".to_string(), format!("cd {} && {}\n"             , dir, &cron.exe)),
        ("".to_string(), format!("when cd {} && {}\n"        , dir, &cron.exe)),
        ("".to_string(), format!("cd {} && {} command\n"     , dir, &cron.exe)),
        ("".to_string(), format!("when cd {} && {} command\n", dir, &cron.exe)),

        (wrap_by_extra_content(""), wrap_by_extra_content(format!("cd {} && {}\n"             , dir, &cron.exe))),
        (wrap_by_extra_content(""), wrap_by_extra_content(format!("when cd {} && {}\n"        , dir, &cron.exe))),
        (wrap_by_extra_content(""), wrap_by_extra_content(format!("cd {} && {} command\n"     , dir, &cron.exe))),
        (wrap_by_extra_content(""), wrap_by_extra_content(format!("when cd {} && {} command\n", dir, &cron.exe))),
    ];

    for (correct, area) in areas.into_iter() {
        cron.my_area = area;
        cron.delete(&dir);
        assert_eq!(correct, cron.my_area);
    }
}
#[test]
#[should_panic(expect = "entered unreachable code")]
fn delete_should_panic() {
    let mut cron = Cron {
        exe    : "path_to_bin".to_string(),
        upper  : "upper"      .to_string(),
        my_area: "middle"     .to_string(),
        lower  : "lower"      .to_string(),
    };

    cron.delete("() <= mistaken Regular Expression!");
}

#[test]
fn discard_should_remove_lines_including_non_existing_dir() {
    let mut cron = Cron {
        exe    : "path_to_bin".to_string(),
        upper  : "upper"      .to_string(),
        my_area: "middle"     .to_string(),
        lower  : "lower"      .to_string(),
    };

    cron.my_area = format!("
        cd /bin && {} command
        cd /bin && {} command2
        cd /path/to/non/existing/dir && {} command
        cd /path/to/non/existing/dir && {} command2
        cd /path/to/really/non/existing/dir && {} command
        cd /path/to/really/non/existing/dir && {} command2
    ", &cron.exe, &cron.exe, &cron.exe, &cron.exe, &cron.exe, &cron.exe);

    let cron = cron.discard();
    assert!(cron.is_ok());
    assert!(cron.unwrap().my_area_is_empty());
}
#[test]
fn discard_should_return_regex_err() {
    let mut cron = Cron {
        exe    : "path_to_bin".to_string(),
        upper  : "upper"      .to_string(),
        my_area: "middle"     .to_string(),
        lower  : "lower"      .to_string(),
    };

    cron.my_area = format!("
        cd /()/<=/mischief/by/bad/user && {} command
    ", &cron.exe);

    let cron = cron.discard();
    assert!(cron.is_err());
}
<commit_msg>[Remedy] "patrol" command delete lines which old bin's created<commit_after>extern crate regex;

use self::regex::{Error as RegexError, Regex};

use constant::{ME, WORKING_DIR_NAME};
use error::{CliError, CronError, CronErrorKind};
use lib::io::*;
use std::collections::BTreeSet;
use std::env;
use std::fmt::Display;
use std::io::Write;
use std::path::Path;
use std::process::{self, Stdio};
use std::str;

#[derive(Debug)]
pub struct Cron {
    exe    : String,
    upper  : String,
    my_area: String,
    lower  : String,
}

impl Cron {
    fn base_mark<D: Display>(keyword: &str, exec: D) -> String {
        format!(
            r#"###################################
# "{}" uses the lines {}.
# Please don't touch them and me!
###################################
{}"#,
            ME,
            keyword,
            exec,
        )
    }
    fn start_mark<D: Display>(addition: D) -> String {
        Self::base_mark("from this", addition)
    }
    fn start_mark_with_patrol<S: AsRef<str>>(exe: S) -> String {
        Self::start_mark(format!("0 12 * * *\t{} patrol\n", exe.as_ref()))
    } fn end_mark() -> String {
        Self::base_mark("up to here", "")
    }

    fn escape_asterisk<S: AsRef<str>>(s: S) -> String {
        let re = match Regex::new(r"\*") {
            Ok(re) => re,
            Err(_) => unreachable!("Mistake regular expression!!"),
        };
        re.replace_all(s.as_ref(), "\\*")
    }

    pub fn read() -> Result<Self, CliError> {
        let exe = try!(env::current_exe());
        let exe = try!(exe.to_str().ok_or(CronError::new(CronErrorKind::BinExistsInInvalidDir)));

        print_with_tag(Tag::Info, "Read cron");

        let result = process::Command::new("crontab")
            .arg("-l")
            .output();
        let output = try!(result);

        let contents = if output.status.success() {
            try!(str::from_utf8(&output.stdout).map_err(|_| CronError::new(CronErrorKind::InvalidCharacterCode)))
        } else {
            ""
        };

        let areas = format!(
            "^(?P<upper>(.|\n)*){}(?P<my_area>(.|\n)*){}(?P<lower>(.|\n)*)$",
            Self::escape_asterisk(Self::start_mark_with_patrol(exe)),
            Self::end_mark(),
        );
        let re = match Regex::new(&areas) {
            Ok(re) => re,
            Err(_) => unreachable!("Mistake the regular expression!!"),
        };

        let (upper, my_area, lower) = match re.captures(contents) {
            Some(caps) => match (caps.name("upper"), caps.name("my_area"), caps.name("lower")) {
                (Some(u), Some(m), Some(l)) => (u, m, l),
                _                           => unreachable!("Mistake regular expression!!"),
            },
            None => (contents, "", ""),
        };

        Ok(Cron {
            exe    : exe    .to_string(),
            upper  : upper  .to_string(),
            my_area: my_area.to_string(),
            lower  : lower  .to_string(),
        })
    }

    pub fn update(mut self, pairs: &[(&str, &str); 2]) -> Result<Self, CliError> {
        let current_dir = try!(Self::current_dir_string());
        self.delete(&current_dir);

        let my_new_area = pairs
            .iter()
            .fold(String::new(), |area, &(time, command)| {
                let line = format!("{}\tcd {} && {} {}\n", time, current_dir, self.exe, command);
                area + &line
            });

        self.my_area = self.my_area + &my_new_area;
        Ok(self)
    }

    fn re_for_matching_missing_dir_line<S: AsRef<str>>(&self, core: S) -> String {
        format!(r".*cd\s+{}\s+&&\s+{}.*\n", core.as_ref(), self.exe)
    }

    pub fn delete<S: AsRef<str>>(&mut self, dir: S) {
        let re = self.re_for_matching_missing_dir_line(dir);
        let re = match Regex::new(re.as_ref()) {
            Ok(re) => re,
            Err(_) => unreachable!("Wrong to use delete()!!"),
        };
        self.my_area = re.replace_all(&self.my_area, "");
    }

    pub fn discard(mut self) -> Result<Self, RegexError> {
        let re = self.re_for_matching_missing_dir_line(r"(?P<path>[^\s]+)");
        let re = match Regex::new(re.as_ref()) {
            Ok(re) => re,
            Err(_) => unreachable!("Mistake regular expression!!"),
        };
        let target_paths = re
            .captures_iter(&self.my_area)
            .map(|caps| match caps.name("path") {
                Some(p) => p.to_string(),
                None    => unreachable!("Mistake keyword for capturing!!"),
            })
            .filter(|path| !Path::new(path).join(WORKING_DIR_NAME).is_dir())
            .collect::<BTreeSet<String>>();

        if target_paths.len() != 0 {
            let re_core = target_paths
                .iter()
                .fold(String::new(), |core, path| core + "|" + path); // We want to use true "Iterator.reduce()".
            let re_core = format!("({})", re_core.trim_left_matches('|'));
            let re = self.re_for_matching_missing_dir_line(re_core);
            let re = try!(Regex::new(re.as_ref()));
            self.my_area = re.replace_all(&self.my_area, "");
        }

        let re = format!("{}(.|\n)*{}", Self::start_mark(""), Self::end_mark());
        let re = match Regex::new(re.as_ref()) {
            Ok(re) => re,
            Err(_) => unreachable!("Mistake regular expression!!"),
        };
        self.upper = re.replace_all(&self.upper, "");
        self.lower = re.replace_all(&self.lower, "");

        Ok(self)
    }

    pub fn set(self) -> Result<(), CliError> {
        print_with_tag(Tag::Info, "Set new cron");

        let result = process::Command::new("crontab")
            .stdin(Stdio::piped())
            .spawn();
        let mut child = try!(result);

        let contents = if self.my_area_is_empty() {
            self.upper + &self.lower
        } else {
            self.upper
                + &Self::start_mark_with_patrol(&self.exe)
                + &self.my_area
                + &Self::end_mark()
                + &self.lower
        };
        match &mut child.stdin {
            &mut Some(ref mut stdin) => try!(stdin.write_all(contents.as_bytes())),
            &mut None                => unreachable!("Please set Stdio::piped()!!"),
        };

        let output = try!(child.wait_with_output());
        if output.status.success() {
            Ok(())
        } else {
            Err(From::from(CronError::new(CronErrorKind::FailedToWrite)))
        }
    }

    pub fn current_dir_string() -> Result<String, CronError> {
        let current_dir = match env::current_dir() {
            Ok(d)  => d,
            Err(_) => unreachable!("There is a problem for current directory!!"),
        };
        current_dir
            .to_str()
            .ok_or(CronError::new(CronErrorKind::InvalidPath))
            .map(|s| s.to_string())
    }

    fn my_area_is_empty(&self) -> bool {
        let re = match Regex::new(r"^\s*$") {
            Ok(re) => re,
            Err(_) => unreachable!("Mistake the regular expression!!"),
        };
        re.is_match(self.my_area.as_ref())
    }
}

#[test]
fn my_area_is_empty_should_return_true() {
    let empties = [
        "",
        " ",
        "\t",
        "\r",
        "\n",
        " \t\r\n",
    ];

    for empty in &empties {
        let cron = Cron {
            exe    : ""   .to_string(),
            upper  : ""   .to_string(),
            my_area: empty.to_string(),
            lower  : ""   .to_string(),
        };
        assert!(cron.my_area_is_empty());
    }
}
#[test]
fn my_area_is_empty_should_return_false() {
    let non_empties = [
        "a",
        "ab",
        " ab",
        "a b",
        "ab ",
        "\tab",
        "a\tb",
        "ab\t",
        "\rab",
        "a\rb",
        "ab\r",
        "\nab",
        "a\nb",
        "ab\n",
        " \t\r\nab",
        "a \t\r\nb",
        "ab \t\r\n",
    ];

    for non_empty in &non_empties {
        let cron = Cron {
            exe    : ""       .to_string(),
            upper  : ""       .to_string(),
            my_area: non_empty.to_string(),
            lower  : ""       .to_string(),
        };
        assert!(!cron.my_area_is_empty());
    }
}

#[test]
fn delete_should_success() {
    fn wrap_by_extra_content<S: AsRef<str>>(target: S) -> String {
        let extra_content = "when cd path_to_extra && path_to_bin command\n";
        format!("{}{}{}", extra_content, target.as_ref(), extra_content)
    }

    let mut cron = Cron {
        exe    : "path_to_exe".to_string(),
        upper  : "upper"      .to_string(),
        my_area: "middle"     .to_string(),
        lower  : "lower"      .to_string(),
    };

    let dir = "path_to_dir";

    let areas = vec![
        ("".to_string(), format!("cd {} && {}\n"             , dir, &cron.exe)),
        ("".to_string(), format!("when cd {} && {}\n"        , dir, &cron.exe)),
        ("".to_string(), format!("cd {} && {} command\n"     , dir, &cron.exe)),
        ("".to_string(), format!("when cd {} && {} command\n", dir, &cron.exe)),

        (wrap_by_extra_content(""), wrap_by_extra_content(format!("cd {} && {}\n"             , dir, &cron.exe))),
        (wrap_by_extra_content(""), wrap_by_extra_content(format!("when cd {} && {}\n"        , dir, &cron.exe))),
        (wrap_by_extra_content(""), wrap_by_extra_content(format!("cd {} && {} command\n"     , dir, &cron.exe))),
        (wrap_by_extra_content(""), wrap_by_extra_content(format!("when cd {} && {} command\n", dir, &cron.exe))),
    ];

    for (correct, area) in areas.into_iter() {
        cron.my_area = area;
        cron.delete(&dir);
        assert_eq!(correct, cron.my_area);
    }
}
#[test]
#[should_panic(expect = "entered unreachable code")]
fn delete_should_panic() {
    let mut cron = Cron {
        exe    : "path_to_bin".to_string(),
        upper  : "upper"      .to_string(),
        my_area: "middle"     .to_string(),
        lower  : "lower"      .to_string(),
    };

    cron.delete("() <= mistaken Regular Expression!");
}

#[test]
fn discard_should_remove_lines_including_non_existing_dir() {
    let mut cron = Cron {
        exe    : "path_to_bin".to_string(),
        upper  : "upper"      .to_string(),
        my_area: "middle"     .to_string(),
        lower  : "lower"      .to_string(),
    };

    cron.my_area = format!("
        cd /bin && {} command
        cd /bin && {} command2
        cd /path/to/non/existing/dir && {} command
        cd /path/to/non/existing/dir && {} command2
        cd /path/to/really/non/existing/dir && {} command
        cd /path/to/really/non/existing/dir && {} command2
    ", &cron.exe, &cron.exe, &cron.exe, &cron.exe, &cron.exe, &cron.exe);

    let cron = cron.discard();
    assert!(cron.is_ok());
    let mut cron = cron.unwrap();
    assert!(cron.my_area_is_empty());

    cron.upper = Cron::start_mark("when path_to_bin command") + &Cron::end_mark();
    cron.lower = Cron::start_mark("when path_to_bin command") + &Cron::end_mark();

    let cron = cron.discard();
    assert!(cron.is_ok());
    let cron = cron.unwrap();
    assert_eq!("", cron.upper);
    assert_eq!("", cron.lower);
}
#[test]
fn discard_should_return_regex_err() {
    let mut cron = Cron {
        exe    : "path_to_bin".to_string(),
        upper  : "upper"      .to_string(),
        my_area: "middle"     .to_string(),
        lower  : "lower"      .to_string(),
    };

    cron.my_area = format!("
        cd /()/<=/mischief/by/bad/user && {} command
    ", &cron.exe);

    let cron = cron.discard();
    assert!(cron.is_err());
}
<|endoftext|>


<commit_before>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/*!

The `ToBytes` and `IterBytes` traits

*/

use io;
use io::Writer;
use option::{None, Option, Some};
use str;

pub type Cb<'self> = &'self fn(buf: &const [u8]) -> bool;

/**
 * A trait to implement in order to make a type hashable;
 * This works in combination with the trait `Hash::Hash`, and
 * may in the future be merged with that trait or otherwise
 * modified when default methods and trait inheritence are
 * completed.
 */
pub trait IterBytes {
    /**
     * Call the provided callback `f` one or more times with
     * byte-slices that should be used when computing a hash
     * value or otherwise "flattening" the structure into
     * a sequence of bytes. The `lsb0` parameter conveys
     * whether the caller is asking for little-endian bytes
     * (`true`) or big-endian (`false`); this should only be
     * relevant in implementations that represent a single
     * multi-byte datum such as a 32 bit integer or 64 bit
     * floating-point value. It can be safely ignored for
     * larger structured types as they are usually processed
     * left-to-right in declaration order, regardless of
     * underlying memory endianness.
     */
    fn iter_bytes(&self, lsb0: bool, f: Cb);
}

impl IterBytes for bool {
    #[inline(always)]
    fn iter_bytes(&self, _lsb0: bool, f: Cb) {
        f([
            *self as u8
        ]);
    }
}

impl IterBytes for u8 {
    #[inline(always)]
    fn iter_bytes(&self, _lsb0: bool, f: Cb) {
        f([
            *self
        ]);
    }
}

impl IterBytes for u16 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        if lsb0 {
            f([
                *self as u8,
                (*self >> 8) as u8
            ]);
        } else {
            f([
                (*self >> 8) as u8,
                *self as u8
            ]);
        }
    }
}

impl IterBytes for u32 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        if lsb0 {
            f([
                *self as u8,
                (*self >> 8) as u8,
                (*self >> 16) as u8,
                (*self >> 24) as u8,
            ]);
        } else {
            f([
                (*self >> 24) as u8,
                (*self >> 16) as u8,
                (*self >> 8) as u8,
                *self as u8
            ]);
        }
    }
}

impl IterBytes for u64 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        if lsb0 {
            f([
                *self as u8,
                (*self >> 8) as u8,
                (*self >> 16) as u8,
                (*self >> 24) as u8,
                (*self >> 32) as u8,
                (*self >> 40) as u8,
                (*self >> 48) as u8,
                (*self >> 56) as u8
            ]);
        } else {
            f([
                (*self >> 56) as u8,
                (*self >> 48) as u8,
                (*self >> 40) as u8,
                (*self >> 32) as u8,
                (*self >> 24) as u8,
                (*self >> 16) as u8,
                (*self >> 8) as u8,
                *self as u8
            ]);
        }
    }
}

impl IterBytes for i8 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as u8).iter_bytes(lsb0, f)
    }
}

impl IterBytes for i16 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as u16).iter_bytes(lsb0, f)
    }
}

impl IterBytes for i32 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as u32).iter_bytes(lsb0, f)
    }
}

impl IterBytes for i64 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as u64).iter_bytes(lsb0, f)
    }
}

impl IterBytes for char {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as u32).iter_bytes(lsb0, f)
    }
}

#[cfg(target_word_size = "32")]
pub mod x32 {
    use to_bytes::{Cb, IterBytes};

    impl IterBytes for uint {
        #[inline(always)]
        fn iter_bytes(&self, lsb0: bool, f: Cb) {
            (*self as u32).iter_bytes(lsb0, f)
        }
    }
}

#[cfg(target_word_size = "64")]
pub mod x64 {
    use to_bytes::{Cb, IterBytes};

    impl IterBytes for uint {
        #[inline(always)]
        fn iter_bytes(&self, lsb0: bool, f: Cb) {
            (*self as u64).iter_bytes(lsb0, f)
        }
    }
}

impl IterBytes for int {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as uint).iter_bytes(lsb0, f)
    }
}

impl<'self,A:IterBytes> IterBytes for &'self [A] {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        for (*self).each |elt| {
            do elt.iter_bytes(lsb0) |bytes| {
                f(bytes)
            }
        }
    }
}

impl<A:IterBytes,B:IterBytes> IterBytes for (A,B) {
  #[inline(always)]
  fn iter_bytes(&self, lsb0: bool, f: Cb) {
    match *self {
      (ref a, ref b) => {
        iter_bytes_2(a, b, lsb0, f);
      }
    }
  }
}

impl<A:IterBytes,B:IterBytes,C:IterBytes> IterBytes for (A,B,C) {
  #[inline(always)]
  fn iter_bytes(&self, lsb0: bool, f: Cb) {
    match *self {
      (ref a, ref b, ref c) => {
        iter_bytes_3(a, b, c, lsb0, f);
      }
    }
  }
}

// Move this to vec, probably.
fn borrow<'x,A>(a: &'x [A]) -> &'x [A] {
    a
}

impl<A:IterBytes> IterBytes for ~[A] {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        borrow(*self).iter_bytes(lsb0, f)
    }
}

impl<A:IterBytes> IterBytes for @[A] {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        borrow(*self).iter_bytes(lsb0, f)
    }
}

pub fn iter_bytes_2<A:IterBytes,B:IterBytes>(a: &A, b: &B,
                                            lsb0: bool, z: Cb) {
    let mut flag = true;
    a.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    b.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
}

pub fn iter_bytes_3<A: IterBytes,
                B: IterBytes,
                C: IterBytes>(a: &A, b: &B, c: &C,
                              lsb0: bool, z: Cb) {
    let mut flag = true;
    a.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    b.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    c.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
}

pub fn iter_bytes_4<A: IterBytes,
                B: IterBytes,
                C: IterBytes,
                D: IterBytes>(a: &A, b: &B, c: &C,
                              d: &D,
                              lsb0: bool, z: Cb) {
    let mut flag = true;
    a.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    b.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    c.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    d.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
}

pub fn iter_bytes_5<A: IterBytes,
                B: IterBytes,
                C: IterBytes,
                D: IterBytes,
                E: IterBytes>(a: &A, b: &B, c: &C,
                              d: &D, e: &E,
                              lsb0: bool, z: Cb) {
    let mut flag = true;
    a.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    b.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    c.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    d.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    e.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
}

pub fn iter_bytes_6<A: IterBytes,
                B: IterBytes,
                C: IterBytes,
                D: IterBytes,
                E: IterBytes,
                F: IterBytes>(a: &A, b: &B, c: &C,
                              d: &D, e: &E, f: &F,
                              lsb0: bool, z: Cb) {
    let mut flag = true;
    a.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    b.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    c.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    d.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    e.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    f.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
}

pub fn iter_bytes_7<A: IterBytes,
                B: IterBytes,
                C: IterBytes,
                D: IterBytes,
                E: IterBytes,
                F: IterBytes,
                G: IterBytes>(a: &A, b: &B, c: &C,
                              d: &D, e: &E, f: &F,
                              g: &G,
                              lsb0: bool, z: Cb) {
    let mut flag = true;
    a.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    b.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    c.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    d.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    e.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    f.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    g.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
}

impl<'self> IterBytes for &'self str {
    #[inline(always)]
    fn iter_bytes(&self, _lsb0: bool, f: Cb) {
        do str::byte_slice(*self) |bytes| {
            f(bytes);
        }
    }
}

impl IterBytes for ~str {
    #[inline(always)]
    fn iter_bytes(&self, _lsb0: bool, f: Cb) {
        do str::byte_slice(*self) |bytes| {
            f(bytes);
        }
    }
}

impl IterBytes for @str {
    #[inline(always)]
    fn iter_bytes(&self, _lsb0: bool, f: Cb) {
        do str::byte_slice(*self) |bytes| {
            f(bytes);
        }
    }
}

impl<A:IterBytes> IterBytes for Option<A> {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        match *self {
          Some(ref a) => iter_bytes_2(&0u8, a, lsb0, f),
          None => 1u8.iter_bytes(lsb0, f)
        }
    }
}

impl<'self,A:IterBytes> IterBytes for &'self A {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (**self).iter_bytes(lsb0, f);
    }
}

impl<A:IterBytes> IterBytes for @A {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (**self).iter_bytes(lsb0, f);
    }
}

impl<A:IterBytes> IterBytes for ~A {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (**self).iter_bytes(lsb0, f);
    }
}

// NB: raw-pointer IterBytes does _not_ dereference
// to the target; it just gives you the pointer-bytes.
impl<A> IterBytes for *const A {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as uint).iter_bytes(lsb0, f);
    }
}

trait ToBytes {
    fn to_bytes(&self, lsb0: bool) -> ~[u8];
}

impl<A:IterBytes> ToBytes for A {
    fn to_bytes(&self, lsb0: bool) -> ~[u8] {
        do io::with_bytes_writer |wr| {
            for self.iter_bytes(lsb0) |bytes| {
                wr.write(bytes)
            }
        }
    }
}
<commit_msg>FIX: Export to_bytes::ToBytes<commit_after>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/*!

The `ToBytes` and `IterBytes` traits

*/

use io;
use io::Writer;
use option::{None, Option, Some};
use str;

pub type Cb<'self> = &'self fn(buf: &const [u8]) -> bool;

/**
 * A trait to implement in order to make a type hashable;
 * This works in combination with the trait `Hash::Hash`, and
 * may in the future be merged with that trait or otherwise
 * modified when default methods and trait inheritence are
 * completed.
 */
pub trait IterBytes {
    /**
     * Call the provided callback `f` one or more times with
     * byte-slices that should be used when computing a hash
     * value or otherwise "flattening" the structure into
     * a sequence of bytes. The `lsb0` parameter conveys
     * whether the caller is asking for little-endian bytes
     * (`true`) or big-endian (`false`); this should only be
     * relevant in implementations that represent a single
     * multi-byte datum such as a 32 bit integer or 64 bit
     * floating-point value. It can be safely ignored for
     * larger structured types as they are usually processed
     * left-to-right in declaration order, regardless of
     * underlying memory endianness.
     */
    fn iter_bytes(&self, lsb0: bool, f: Cb);
}

impl IterBytes for bool {
    #[inline(always)]
    fn iter_bytes(&self, _lsb0: bool, f: Cb) {
        f([
            *self as u8
        ]);
    }
}

impl IterBytes for u8 {
    #[inline(always)]
    fn iter_bytes(&self, _lsb0: bool, f: Cb) {
        f([
            *self
        ]);
    }
}

impl IterBytes for u16 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        if lsb0 {
            f([
                *self as u8,
                (*self >> 8) as u8
            ]);
        } else {
            f([
                (*self >> 8) as u8,
                *self as u8
            ]);
        }
    }
}

impl IterBytes for u32 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        if lsb0 {
            f([
                *self as u8,
                (*self >> 8) as u8,
                (*self >> 16) as u8,
                (*self >> 24) as u8,
            ]);
        } else {
            f([
                (*self >> 24) as u8,
                (*self >> 16) as u8,
                (*self >> 8) as u8,
                *self as u8
            ]);
        }
    }
}

impl IterBytes for u64 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        if lsb0 {
            f([
                *self as u8,
                (*self >> 8) as u8,
                (*self >> 16) as u8,
                (*self >> 24) as u8,
                (*self >> 32) as u8,
                (*self >> 40) as u8,
                (*self >> 48) as u8,
                (*self >> 56) as u8
            ]);
        } else {
            f([
                (*self >> 56) as u8,
                (*self >> 48) as u8,
                (*self >> 40) as u8,
                (*self >> 32) as u8,
                (*self >> 24) as u8,
                (*self >> 16) as u8,
                (*self >> 8) as u8,
                *self as u8
            ]);
        }
    }
}

impl IterBytes for i8 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as u8).iter_bytes(lsb0, f)
    }
}

impl IterBytes for i16 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as u16).iter_bytes(lsb0, f)
    }
}

impl IterBytes for i32 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as u32).iter_bytes(lsb0, f)
    }
}

impl IterBytes for i64 {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as u64).iter_bytes(lsb0, f)
    }
}

impl IterBytes for char {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as u32).iter_bytes(lsb0, f)
    }
}

#[cfg(target_word_size = "32")]
pub mod x32 {
    use to_bytes::{Cb, IterBytes};

    impl IterBytes for uint {
        #[inline(always)]
        fn iter_bytes(&self, lsb0: bool, f: Cb) {
            (*self as u32).iter_bytes(lsb0, f)
        }
    }
}

#[cfg(target_word_size = "64")]
pub mod x64 {
    use to_bytes::{Cb, IterBytes};

    impl IterBytes for uint {
        #[inline(always)]
        fn iter_bytes(&self, lsb0: bool, f: Cb) {
            (*self as u64).iter_bytes(lsb0, f)
        }
    }
}

impl IterBytes for int {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as uint).iter_bytes(lsb0, f)
    }
}

impl<'self,A:IterBytes> IterBytes for &'self [A] {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        for (*self).each |elt| {
            do elt.iter_bytes(lsb0) |bytes| {
                f(bytes)
            }
        }
    }
}

impl<A:IterBytes,B:IterBytes> IterBytes for (A,B) {
  #[inline(always)]
  fn iter_bytes(&self, lsb0: bool, f: Cb) {
    match *self {
      (ref a, ref b) => {
        iter_bytes_2(a, b, lsb0, f);
      }
    }
  }
}

impl<A:IterBytes,B:IterBytes,C:IterBytes> IterBytes for (A,B,C) {
  #[inline(always)]
  fn iter_bytes(&self, lsb0: bool, f: Cb) {
    match *self {
      (ref a, ref b, ref c) => {
        iter_bytes_3(a, b, c, lsb0, f);
      }
    }
  }
}

// Move this to vec, probably.
fn borrow<'x,A>(a: &'x [A]) -> &'x [A] {
    a
}

impl<A:IterBytes> IterBytes for ~[A] {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        borrow(*self).iter_bytes(lsb0, f)
    }
}

impl<A:IterBytes> IterBytes for @[A] {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        borrow(*self).iter_bytes(lsb0, f)
    }
}

pub fn iter_bytes_2<A:IterBytes,B:IterBytes>(a: &A, b: &B,
                                            lsb0: bool, z: Cb) {
    let mut flag = true;
    a.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    b.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
}

pub fn iter_bytes_3<A: IterBytes,
                B: IterBytes,
                C: IterBytes>(a: &A, b: &B, c: &C,
                              lsb0: bool, z: Cb) {
    let mut flag = true;
    a.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    b.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    c.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
}

pub fn iter_bytes_4<A: IterBytes,
                B: IterBytes,
                C: IterBytes,
                D: IterBytes>(a: &A, b: &B, c: &C,
                              d: &D,
                              lsb0: bool, z: Cb) {
    let mut flag = true;
    a.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    b.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    c.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    d.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
}

pub fn iter_bytes_5<A: IterBytes,
                B: IterBytes,
                C: IterBytes,
                D: IterBytes,
                E: IterBytes>(a: &A, b: &B, c: &C,
                              d: &D, e: &E,
                              lsb0: bool, z: Cb) {
    let mut flag = true;
    a.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    b.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    c.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    d.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    e.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
}

pub fn iter_bytes_6<A: IterBytes,
                B: IterBytes,
                C: IterBytes,
                D: IterBytes,
                E: IterBytes,
                F: IterBytes>(a: &A, b: &B, c: &C,
                              d: &D, e: &E, f: &F,
                              lsb0: bool, z: Cb) {
    let mut flag = true;
    a.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    b.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    c.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    d.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    e.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    f.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
}

pub fn iter_bytes_7<A: IterBytes,
                B: IterBytes,
                C: IterBytes,
                D: IterBytes,
                E: IterBytes,
                F: IterBytes,
                G: IterBytes>(a: &A, b: &B, c: &C,
                              d: &D, e: &E, f: &F,
                              g: &G,
                              lsb0: bool, z: Cb) {
    let mut flag = true;
    a.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    b.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    c.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    d.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    e.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    f.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
    if !flag { return; }
    g.iter_bytes(lsb0, |bytes| {flag = z(bytes); flag});
}

impl<'self> IterBytes for &'self str {
    #[inline(always)]
    fn iter_bytes(&self, _lsb0: bool, f: Cb) {
        do str::byte_slice(*self) |bytes| {
            f(bytes);
        }
    }
}

impl IterBytes for ~str {
    #[inline(always)]
    fn iter_bytes(&self, _lsb0: bool, f: Cb) {
        do str::byte_slice(*self) |bytes| {
            f(bytes);
        }
    }
}

impl IterBytes for @str {
    #[inline(always)]
    fn iter_bytes(&self, _lsb0: bool, f: Cb) {
        do str::byte_slice(*self) |bytes| {
            f(bytes);
        }
    }
}

impl<A:IterBytes> IterBytes for Option<A> {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        match *self {
          Some(ref a) => iter_bytes_2(&0u8, a, lsb0, f),
          None => 1u8.iter_bytes(lsb0, f)
        }
    }
}

impl<'self,A:IterBytes> IterBytes for &'self A {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (**self).iter_bytes(lsb0, f);
    }
}

impl<A:IterBytes> IterBytes for @A {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (**self).iter_bytes(lsb0, f);
    }
}

impl<A:IterBytes> IterBytes for ~A {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (**self).iter_bytes(lsb0, f);
    }
}

// NB: raw-pointer IterBytes does _not_ dereference
// to the target; it just gives you the pointer-bytes.
impl<A> IterBytes for *const A {
    #[inline(always)]
    fn iter_bytes(&self, lsb0: bool, f: Cb) {
        (*self as uint).iter_bytes(lsb0, f);
    }
}

pub trait ToBytes {
    fn to_bytes(&self, lsb0: bool) -> ~[u8];
}

impl<A:IterBytes> ToBytes for A {
    fn to_bytes(&self, lsb0: bool) -> ~[u8] {
        do io::with_bytes_writer |wr| {
            for self.iter_bytes(lsb0) |bytes| {
                wr.write(bytes)
            }
        }
    }
}
<|endoftext|>


<commit_before>use core::marker::PhantomData;
use core::ops::{Index, IndexMut};

use mem::FrameAllocator;
use mem::paging::ENTRY_COUNT;
use mem::paging::entry::*;

pub const P4: *mut Table<Level4> = 0x000f_ffff_ffff_f000 as *mut _;

pub trait TableLevel {}
pub trait HierarchicalLevel: TableLevel {
    type NextLevel: TableLevel;
}

pub enum Level4 {}
pub enum Level3 {}
pub enum Level2 {}
pub enum Level1 {}

impl TableLevel for Level4 {}
impl TableLevel for Level3 {}
impl TableLevel for Level2 {}
impl TableLevel for Level1 {}
impl HierarchicalLevel for Level4 { type NextLevel = Level3; }
impl HierarchicalLevel for Level3 { type NextLevel = Level2; }
impl HierarchicalLevel for Level2 { type NextLevel = Level1; }

pub struct Table<L: TableLevel> {
    entries: [Entry; ENTRY_COUNT],
    level: PhantomData<L>
}

impl<L> Index<usize> for Table<L>
    where L: TableLevel
{
    type Output = Entry;

    fn index(&self, index: usize) -> &Entry {
        &self.entries[index]
    }
}

impl<L> IndexMut<usize> for Table<L>
    where L: TableLevel
{
    fn index_mut(&mut self, index: usize) -> &mut Entry {
        &mut self.entries[index]
    }
}

impl<L> Table<L>
    where L: TableLevel
{
    pub fn zero(&mut self) {
        for entry in self.entries.iter_mut() {
            entry.set_unused();
        }
    }
}

impl<L> Table<L>
    where L: HierarchicalLevel
{
    pub fn next_table(&self, index: usize) -> Option<&Table<L::NextLevel>> {
        self.next_table_address(index)
            .map(|address| unsafe { &*(address as *const _) })
    }

    pub fn next_table_mut(&mut self, index: usize) ->Option<&mut Table<L::NextLevel>> {
        self.next_table_address(index)
            .map(|address| unsafe { &mut *(address as *mut _) })
    }

    pub fn next_table_create<A>(&mut self, index: usize, allocator: &mut A) -> &mut Table<L::NextLevel>
        where A: FrameAllocator
    {
        if self.next_table(index).is_none() {
            assert!(!self.entries[index].flags().contains(HUGE_PAGE));
            let frame = allocator.alloc().expect("no frames available");
            self.entries[index].set(frame, PRESENT | WRITABLE);
            self.next_table_mut(index).unwrap().zero();
        }
        self.next_table_mut(index).unwrap()
    }

    // TODO: use physical/virtualaddress type alias?
    fn next_table_address(&self, index: usize) -> Option<usize> {
        let flags = self[index].flags();
        if flags.contains(PRESENT) && !flags.contains(HUGE_PAGE) {
            let table_address = self as *const _ as usize;
            Some((table_address << 9) | (index << 12))
        } else {
            None
        }
    }
}
<commit_msg>helps to use the right address<commit_after>use core::marker::PhantomData;
use core::ops::{Index, IndexMut};

use mem::FrameAllocator;
use mem::paging::ENTRY_COUNT;
use mem::paging::entry::*;

pub const P4: *mut Table<Level4> = 0xffff_ffff_ffff_f000 as *mut _;

pub trait TableLevel {}
pub trait HierarchicalLevel: TableLevel {
    type NextLevel: TableLevel;
}

pub enum Level4 {}
pub enum Level3 {}
pub enum Level2 {}
pub enum Level1 {}

impl TableLevel for Level4 {}
impl TableLevel for Level3 {}
impl TableLevel for Level2 {}
impl TableLevel for Level1 {}
impl HierarchicalLevel for Level4 { type NextLevel = Level3; }
impl HierarchicalLevel for Level3 { type NextLevel = Level2; }
impl HierarchicalLevel for Level2 { type NextLevel = Level1; }

pub struct Table<L: TableLevel> {
    entries: [Entry; ENTRY_COUNT],
    level: PhantomData<L>
}

impl<L> Index<usize> for Table<L>
    where L: TableLevel
{
    type Output = Entry;

    fn index(&self, index: usize) -> &Entry {
        &self.entries[index]
    }
}

impl<L> IndexMut<usize> for Table<L>
    where L: TableLevel
{
    fn index_mut(&mut self, index: usize) -> &mut Entry {
        &mut self.entries[index]
    }
}

impl<L> Table<L>
    where L: TableLevel
{
    pub fn zero(&mut self) {
        for entry in self.entries.iter_mut() {
            entry.set_unused();
        }
    }
}

impl<L> Table<L>
    where L: HierarchicalLevel
{
    pub fn next_table(&self, index: usize) -> Option<&Table<L::NextLevel>> {
        self.next_table_address(index)
            .map(|address| unsafe { &*(address as *const _) })
    }

    pub fn next_table_mut(&mut self, index: usize) ->Option<&mut Table<L::NextLevel>> {
        self.next_table_address(index)
            .map(|address| unsafe { &mut *(address as *mut _) })
    }

    pub fn next_table_create<A>(&mut self, index: usize, allocator: &mut A) -> &mut Table<L::NextLevel>
        where A: FrameAllocator
    {
        if self.next_table(index).is_none() {
            assert!(!self.entries[index].flags().contains(HUGE_PAGE));
            let frame = allocator.alloc().expect("no frames available");
            self.entries[index].set(frame, PRESENT | WRITABLE);
            self.next_table_mut(index).unwrap().zero();
        }
        self.next_table_mut(index).unwrap()
    }

    // TODO: use physical/virtualaddress type alias?
    fn next_table_address(&self, index: usize) -> Option<usize> {
        let flags = self[index].flags();
        if flags.contains(PRESENT) && !flags.contains(HUGE_PAGE) {
            let table_address = self as *const _ as usize;
            Some((table_address << 9) | (index << 12))
        } else {
            None
        }
    }
}
<|endoftext|>


<commit_before>//! High-level receiver for receiving P25 voice, data, and trunking messages.

use error::P25Error;
use message::data_unit::{DataUnitReceiver, ReceiverEvent};
use message::nid::NetworkId;
use message::status::StreamSymbol;
use trunking::tsbk::{TsbkFields, TsbkReceiver};
use voice::control::LinkControlFields;
use voice::crypto::CryptoControlFields;
use voice::frame::VoiceFrame;
use voice::header::{VoiceHeaderReceiver, VoiceHeaderFields};
use voice::term::VoiceLCTerminatorReceiver;
use stats::{Stats, HasStats};

use voice::frame_group::{
    FrameGroupEvent,
    VoiceCCFrameGroupReceiver,
    VoiceLCFrameGroupReceiver,
};

/// Events that can occur when receiving P25 messages.
pub enum MessageEvent {
    /// A runtime error occured.
    Error(P25Error),
    /// An NID at the start of a packet was decoded.
    PacketNID(NetworkId),
    /// A voice header was received.
    VoiceHeader(VoiceHeaderFields),
    /// A voice frame was received.
    VoiceFrame(VoiceFrame),
    /// A link control word was decoded.
    LinkControl(LinkControlFields),
    /// A crypto control word was decoded.
    CryptoControl(CryptoControlFields),
    /// A voice low-speed data fragment was decoded.
    LowSpeedDataFragment(u32),
    /// A trunking signalling packet was received.
    TrunkingControl(TsbkFields),
    /// A voice terminator link control was received.
    VoiceTerm(LinkControlFields),
}

/// Internal state of the state machine.
enum State {
    /// Waiting for an event from the lower-level state machine.
    Idle,
    /// Decoding a voice header packet.
    DecodeHeader(VoiceHeaderReceiver),
    /// Decoding a link control frame group packet.
    DecodeLCFrameGroup(VoiceLCFrameGroupReceiver),
    /// Decoding a crypto control frame group packet.
    DecodeCCFrameGroup(VoiceCCFrameGroupReceiver),
    /// Decoding a link control voice terminator.
    DecodeLCTerminator(VoiceLCTerminatorReceiver),
    /// Decoding a trunking signalling packet.
    DecodeTSBK(TsbkReceiver),
}

/// Action the state machine should take.
enum StateChange {
    /// Propagate an event.
    Event(MessageEvent),
    /// Propagate an event and change state.
    EventChange(MessageEvent, State),
    /// Do nothing.
    NoChange,
}

/// State machine for high-level message reception.
pub struct MessageReceiver {
    /// Lower-level stream receiver.
    pub recv: DataUnitReceiver,
    /// Current state.
    state: State,
    stats: Stats,
}

impl MessageReceiver {
    /// Create a new `MessageReceiver` in the initial state.
    pub fn new() -> MessageReceiver {
        MessageReceiver {
            recv: DataUnitReceiver::new(),
            state: State::Idle,
            stats: Stats::default(),
        }
    }

    /// Feed in a baseband sample, possibly producing a new event or message to be handled
    /// by the given handler.
    pub fn feed(&mut self, s: f32) -> Option<MessageEvent> {
        match self.handle(s) {
            StateChange::Event(e) => Some(e),
            StateChange::EventChange(e, s) => {
                self.state = s;
                Some(e)
            },
            StateChange::NoChange => None,
        }
    }

    /// Process the given sample and determine how to update state.
    fn handle(&mut self, s: f32) -> StateChange {
        use self::State::*;
        use self::StateChange::*;
        use message::nid::DataUnit::*;

        let event = match self.recv.feed(s) {
            Some(Ok(event)) => event,
            Some(Err(err)) => {
                self.recv.resync();
                return Event(MessageEvent::Error(err));
            },
            None => return NoChange,
        };

        self.stats.merge(&mut self.recv);

        let dibit = match event {
            ReceiverEvent::NetworkId(nid) => {
                let next = match nid.data_unit {
                    VoiceHeader =>
                        DecodeHeader(VoiceHeaderReceiver::new()),
                    VoiceSimpleTerminator => {
                        self.recv.flush_pads();
                        Idle
                    },
                    VoiceLCTerminator =>
                        DecodeLCTerminator(VoiceLCTerminatorReceiver::new()),
                    VoiceLCFrameGroup =>
                        DecodeLCFrameGroup(VoiceLCFrameGroupReceiver::new()),
                    VoiceCCFrameGroup =>
                        DecodeCCFrameGroup(VoiceCCFrameGroupReceiver::new()),
                    TrunkingSignaling =>
                        DecodeTSBK(TsbkReceiver::new()),
                    DataPacket => {
                        self.recv.resync();
                        Idle
                    },
                };

                return EventChange(MessageEvent::PacketNID(nid), next);
            },
            ReceiverEvent::Symbol(StreamSymbol::Status(_)) => return NoChange,
            ReceiverEvent::Symbol(StreamSymbol::Data(dibit)) => dibit,
        };

        let next = match self.state {
            DecodeHeader(ref mut head) => match head.feed(dibit) {
                Some(Ok(h)) => {
                    self.recv.flush_pads();
                    EventChange(MessageEvent::VoiceHeader(h), Idle)
                },
                Some(Err(err)) => {
                    self.recv.resync();
                    EventChange(MessageEvent::Error(err), Idle)
                },
                None => NoChange,
            },
            DecodeLCFrameGroup(ref mut fg) => match fg.feed(dibit) {
                Some(Ok(event)) => {
                    if fg.done() {
                        self.recv.flush_pads();
                    }

                    match event {
                        FrameGroupEvent::VoiceFrame(vf) =>
                            Event(MessageEvent::VoiceFrame(vf)),
                        FrameGroupEvent::Extra(lc) =>
                            Event(MessageEvent::LinkControl(lc)),
                        FrameGroupEvent::DataFragment(frag) =>
                            Event(MessageEvent::LowSpeedDataFragment(frag)),
                    }
                },
                Some(Err(err)) => {
                    self.recv.resync();
                    EventChange(MessageEvent::Error(err), Idle)
                },
                None => NoChange,
            },
            DecodeCCFrameGroup(ref mut fg) => match fg.feed(dibit) {
                Some(Ok(event)) => match event {
                    FrameGroupEvent::VoiceFrame(vf) => {
                        if fg.done() {
                            self.recv.flush_pads();
                        }

                        Event(MessageEvent::VoiceFrame(vf))
                    },
                    FrameGroupEvent::Extra(cc) =>
                        Event(MessageEvent::CryptoControl(cc)),
                    FrameGroupEvent::DataFragment(frag) =>
                        Event(MessageEvent::LowSpeedDataFragment(frag))
                },
                Some(Err(err)) => {
                    self.recv.resync();
                    EventChange(MessageEvent::Error(err), Idle)
                },
                None => NoChange,
            },
            DecodeLCTerminator(ref mut term) => match term.feed(dibit) {
                Some(Ok(lc)) => {
                    self.recv.flush_pads();
                    EventChange(MessageEvent::VoiceTerm(lc), Idle)
                },
                Some(Err(err)) => {
                    self.recv.resync();
                    EventChange(MessageEvent::Error(err), Idle)
                },
                None => NoChange,
            },
            DecodeTSBK(ref mut dec) => match dec.feed(dibit) {
                Some(Ok(tsbk)) => {
                    if tsbk.is_tail() {
                        self.recv.flush_pads();
                    }

                    Event(MessageEvent::TrunkingControl(tsbk))
                },
                Some(Err(err)) => {
                    self.recv.resync();
                    EventChange(MessageEvent::Error(err), Idle)
                },
                None => NoChange,
            },
            Idle => NoChange,
        };

        match self.state {
            DecodeHeader(ref mut head) => self.stats.merge(head),
            DecodeLCFrameGroup(ref mut fg) => self.stats.merge(fg),
            DecodeCCFrameGroup(ref mut fg) => self.stats.merge(fg),
            DecodeLCTerminator(ref mut term) => self.stats.merge(term),
            DecodeTSBK(_) | Idle => {},
        }

        next
    }
}

impl HasStats for MessageReceiver {
    fn stats(&mut self) -> &mut Stats { &mut self.stats }
}
<commit_msg>Make DataUnitReceiver private, add a resync method<commit_after>//! High-level receiver for receiving P25 voice, data, and trunking messages.

use error::P25Error;
use message::data_unit::{DataUnitReceiver, ReceiverEvent};
use message::nid::NetworkId;
use message::status::StreamSymbol;
use trunking::tsbk::{TsbkFields, TsbkReceiver};
use voice::control::LinkControlFields;
use voice::crypto::CryptoControlFields;
use voice::frame::VoiceFrame;
use voice::header::{VoiceHeaderReceiver, VoiceHeaderFields};
use voice::term::VoiceLCTerminatorReceiver;
use stats::{Stats, HasStats};

use voice::frame_group::{
    FrameGroupEvent,
    VoiceCCFrameGroupReceiver,
    VoiceLCFrameGroupReceiver,
};

/// Events that can occur when receiving P25 messages.
pub enum MessageEvent {
    /// A runtime error occured.
    Error(P25Error),
    /// An NID at the start of a packet was decoded.
    PacketNID(NetworkId),
    /// A voice header was received.
    VoiceHeader(VoiceHeaderFields),
    /// A voice frame was received.
    VoiceFrame(VoiceFrame),
    /// A link control word was decoded.
    LinkControl(LinkControlFields),
    /// A crypto control word was decoded.
    CryptoControl(CryptoControlFields),
    /// A voice low-speed data fragment was decoded.
    LowSpeedDataFragment(u32),
    /// A trunking signalling packet was received.
    TrunkingControl(TsbkFields),
    /// A voice terminator link control was received.
    VoiceTerm(LinkControlFields),
}

/// Internal state of the state machine.
enum State {
    /// Waiting for an event from the lower-level state machine.
    Idle,
    /// Decoding a voice header packet.
    DecodeHeader(VoiceHeaderReceiver),
    /// Decoding a link control frame group packet.
    DecodeLCFrameGroup(VoiceLCFrameGroupReceiver),
    /// Decoding a crypto control frame group packet.
    DecodeCCFrameGroup(VoiceCCFrameGroupReceiver),
    /// Decoding a link control voice terminator.
    DecodeLCTerminator(VoiceLCTerminatorReceiver),
    /// Decoding a trunking signalling packet.
    DecodeTSBK(TsbkReceiver),
}

/// Action the state machine should take.
enum StateChange {
    /// Propagate an event.
    Event(MessageEvent),
    /// Propagate an event and change state.
    EventChange(MessageEvent, State),
    /// Do nothing.
    NoChange,
}

/// State machine for high-level message reception.
pub struct MessageReceiver {
    /// Lower-level stream receiver.
    recv: DataUnitReceiver,
    /// Current state.
    state: State,
    stats: Stats,
}

impl MessageReceiver {
    /// Create a new `MessageReceiver` in the initial state.
    pub fn new() -> MessageReceiver {
        MessageReceiver {
            recv: DataUnitReceiver::new(),
            state: State::Idle,
            stats: Stats::default(),
        }
    }

    /// Force the receiver into frame synchronization.
    pub fn resync(&mut self) { self.recv.resync(); }

    /// Feed in a baseband sample, possibly producing a new event or message to be handled
    /// by the given handler.
    pub fn feed(&mut self, s: f32) -> Option<MessageEvent> {
        match self.handle(s) {
            StateChange::Event(e) => Some(e),
            StateChange::EventChange(e, s) => {
                self.state = s;
                Some(e)
            },
            StateChange::NoChange => None,
        }
    }

    /// Process the given sample and determine how to update state.
    fn handle(&mut self, s: f32) -> StateChange {
        use self::State::*;
        use self::StateChange::*;
        use message::nid::DataUnit::*;

        let event = match self.recv.feed(s) {
            Some(Ok(event)) => event,
            Some(Err(err)) => {
                self.recv.resync();
                return Event(MessageEvent::Error(err));
            },
            None => return NoChange,
        };

        self.stats.merge(&mut self.recv);

        let dibit = match event {
            ReceiverEvent::NetworkId(nid) => {
                let next = match nid.data_unit {
                    VoiceHeader =>
                        DecodeHeader(VoiceHeaderReceiver::new()),
                    VoiceSimpleTerminator => {
                        self.recv.flush_pads();
                        Idle
                    },
                    VoiceLCTerminator =>
                        DecodeLCTerminator(VoiceLCTerminatorReceiver::new()),
                    VoiceLCFrameGroup =>
                        DecodeLCFrameGroup(VoiceLCFrameGroupReceiver::new()),
                    VoiceCCFrameGroup =>
                        DecodeCCFrameGroup(VoiceCCFrameGroupReceiver::new()),
                    TrunkingSignaling =>
                        DecodeTSBK(TsbkReceiver::new()),
                    DataPacket => {
                        self.recv.resync();
                        Idle
                    },
                };

                return EventChange(MessageEvent::PacketNID(nid), next);
            },
            ReceiverEvent::Symbol(StreamSymbol::Status(_)) => return NoChange,
            ReceiverEvent::Symbol(StreamSymbol::Data(dibit)) => dibit,
        };

        let next = match self.state {
            DecodeHeader(ref mut head) => match head.feed(dibit) {
                Some(Ok(h)) => {
                    self.recv.flush_pads();
                    EventChange(MessageEvent::VoiceHeader(h), Idle)
                },
                Some(Err(err)) => {
                    self.recv.resync();
                    EventChange(MessageEvent::Error(err), Idle)
                },
                None => NoChange,
            },
            DecodeLCFrameGroup(ref mut fg) => match fg.feed(dibit) {
                Some(Ok(event)) => {
                    if fg.done() {
                        self.recv.flush_pads();
                    }

                    match event {
                        FrameGroupEvent::VoiceFrame(vf) =>
                            Event(MessageEvent::VoiceFrame(vf)),
                        FrameGroupEvent::Extra(lc) =>
                            Event(MessageEvent::LinkControl(lc)),
                        FrameGroupEvent::DataFragment(frag) =>
                            Event(MessageEvent::LowSpeedDataFragment(frag)),
                    }
                },
                Some(Err(err)) => {
                    self.recv.resync();
                    EventChange(MessageEvent::Error(err), Idle)
                },
                None => NoChange,
            },
            DecodeCCFrameGroup(ref mut fg) => match fg.feed(dibit) {
                Some(Ok(event)) => match event {
                    FrameGroupEvent::VoiceFrame(vf) => {
                        if fg.done() {
                            self.recv.flush_pads();
                        }

                        Event(MessageEvent::VoiceFrame(vf))
                    },
                    FrameGroupEvent::Extra(cc) =>
                        Event(MessageEvent::CryptoControl(cc)),
                    FrameGroupEvent::DataFragment(frag) =>
                        Event(MessageEvent::LowSpeedDataFragment(frag))
                },
                Some(Err(err)) => {
                    self.recv.resync();
                    EventChange(MessageEvent::Error(err), Idle)
                },
                None => NoChange,
            },
            DecodeLCTerminator(ref mut term) => match term.feed(dibit) {
                Some(Ok(lc)) => {
                    self.recv.flush_pads();
                    EventChange(MessageEvent::VoiceTerm(lc), Idle)
                },
                Some(Err(err)) => {
                    self.recv.resync();
                    EventChange(MessageEvent::Error(err), Idle)
                },
                None => NoChange,
            },
            DecodeTSBK(ref mut dec) => match dec.feed(dibit) {
                Some(Ok(tsbk)) => {
                    if tsbk.is_tail() {
                        self.recv.flush_pads();
                    }

                    Event(MessageEvent::TrunkingControl(tsbk))
                },
                Some(Err(err)) => {
                    self.recv.resync();
                    EventChange(MessageEvent::Error(err), Idle)
                },
                None => NoChange,
            },
            Idle => NoChange,
        };

        match self.state {
            DecodeHeader(ref mut head) => self.stats.merge(head),
            DecodeLCFrameGroup(ref mut fg) => self.stats.merge(fg),
            DecodeCCFrameGroup(ref mut fg) => self.stats.merge(fg),
            DecodeLCTerminator(ref mut term) => self.stats.merge(term),
            DecodeTSBK(_) | Idle => {},
        }

        next
    }
}

impl HasStats for MessageReceiver {
    fn stats(&mut self) -> &mut Stats { &mut self.stats }
}
<|endoftext|>