<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    constants,
    error::Error,
    secure_backend::StorageLocation::{LocalStorage, RemoteStorage},
    SecureBackends,
};
use libra_config::config::HANDSHAKE_VERSION;
use libra_crypto::{ed25519::Ed25519PublicKey, x25519, ValidCryptoMaterial};
use libra_global_constants::{
    CONSENSUS_KEY, FULLNODE_NETWORK_KEY, OPERATOR_KEY, VALIDATOR_NETWORK_KEY,
};
use libra_network_address::{
    encrypted::{RawEncNetworkAddress, TEST_ROOT_KEY, TEST_ROOT_KEY_VERSION},
    NetworkAddress, RawNetworkAddress,
};
use libra_secure_storage::{CryptoStorage, KVStorage, Storage, Value};
use libra_secure_time::{RealTimeService, TimeService};
use libra_types::{
    account_address::{self, AccountAddress},
    transaction::{RawTransaction, SignedTransaction, Transaction},
};
use std::{convert::TryFrom, time::Duration};
use structopt::StructOpt;

// TODO(davidiw) add operator_address, since that will eventually be the identity producing this.
#[derive(Debug, StructOpt)]
pub struct ValidatorConfig {
    #[structopt(long)]
    owner_address: AccountAddress,
    #[structopt(long)]
    validator_address: NetworkAddress,
    #[structopt(long)]
    fullnode_address: NetworkAddress,
    #[structopt(flatten)]
    backends: SecureBackends,
}

impl ValidatorConfig {
    pub fn execute(self) -> Result<Transaction, Error> {
        let mut local_storage = self.backends.local.create_storage(LocalStorage)?;

        // Step 1) Retrieve keys from local storage
        let consensus_key = ed25519_from_storage(CONSENSUS_KEY, &local_storage)?;
        let fullnode_network_key = x25519_from_storage(FULLNODE_NETWORK_KEY, &local_storage)?;
        let validator_network_key = x25519_from_storage(VALIDATOR_NETWORK_KEY, &local_storage)?;
        let operator_key = ed25519_from_storage(OPERATOR_KEY, &local_storage)?;

        // TODO(davidiw): In genesis this is irrelevant -- afterward we need to obtain the
        // current sequence number by querying the blockchain.
        let sequence_number = 0;

        // TODO(davidiw): This is currently not supported
        // let sender = self.owner_address;
        let sender = account_address::from_public_key(&operator_key);

        // TODO(philiphayes): actually get the real keys from storage
        let root_key = TEST_ROOT_KEY;
        let key_version = TEST_ROOT_KEY_VERSION;
        // only supports one address for now
        let addr_idx = 0;

        // append ln-noise-ik and ln-handshake protocols to base network addresses

        let validator_address = self
            .validator_address
            .clone()
            .append_prod_protos(validator_network_key, HANDSHAKE_VERSION);
        let raw_validator_address = RawNetworkAddress::try_from(&validator_address)
            .map_err(|e| Error::UnexpectedError(format!("(raw_validator_address) {}", e)))?;
        let enc_validator_address = raw_validator_address
            .encrypt(&root_key, key_version, &sender, sequence_number, addr_idx)
            .map_err(|e| Error::UnexpectedError(format!("(enc_validator_address) {}", e)))?;
        let raw_enc_validator_address = RawEncNetworkAddress::try_from(&enc_validator_address)
            .map_err(|e| Error::UnexpectedError(format!("(raw_enc_validator_address) {}", e)))?;

        let fullnode_address = self
            .fullnode_address
            .clone()
            .append_prod_protos(fullnode_network_key, HANDSHAKE_VERSION);
        let raw_fullnode_address = RawNetworkAddress::try_from(&fullnode_address)
            .map_err(|e| Error::UnexpectedError(format!("(raw_fullnode_address) {}", e)))?;

        // Step 2) Generate transaction

        // TODO(philiphayes): remove network identity pubkey field from struct when
        // transition complete
        let script = transaction_builder::encode_set_validator_config_script(
            sender,
            consensus_key.to_bytes().to_vec(),
            validator_network_key.to_bytes(),
            raw_enc_validator_address.into(),
            fullnode_network_key.to_bytes(),
            raw_fullnode_address.into(),
        );
        let expiration_time = RealTimeService::new().now() + constants::TXN_EXPIRATION_SECS;
        let raw_transaction = RawTransaction::new_script(
            sender,
            sequence_number,
            script,
            constants::MAX_GAS_AMOUNT,
            constants::GAS_UNIT_PRICE,
            constants::GAS_CURRENCY_CODE.to_owned(),
            Duration::from_secs(expiration_time),
        );
        let signature = local_storage
            .sign(OPERATOR_KEY, &raw_transaction)
            .map_err(|e| {
                Error::LocalStorageSigningError("validator-config", OPERATOR_KEY, e.to_string())
            })?;
        let signed_txn = SignedTransaction::new(raw_transaction, operator_key, signature);
        let txn = Transaction::UserTransaction(signed_txn);

        // Step 3) Submit to remote storage

        if let Some(remote_config) = self.backends.remote {
            let mut remote_storage = remote_config.create_storage(RemoteStorage)?;
            let txn = Value::Transaction(txn.clone());
            remote_storage
                .set(constants::VALIDATOR_CONFIG, txn)
                .map_err(|e| {
                    Error::RemoteStorageWriteError(constants::VALIDATOR_CONFIG, e.to_string())
                })?;
        }

        Ok(txn)
    }
}

fn ed25519_from_storage(
    key_name: &'static str,
    storage: &Storage,
) -> Result<Ed25519PublicKey, Error> {
    Ok(storage
        .get_public_key(key_name)
        .map_err(|e| Error::LocalStorageReadError(key_name, e.to_string()))?
        .public_key)
}

fn x25519_from_storage(
    key_name: &'static str,
    storage: &Storage,
) -> Result<x25519::PublicKey, Error> {
    let edkey = ed25519_from_storage(key_name, storage)?;
    x25519::PublicKey::from_ed25519_public_bytes(&edkey.to_bytes())
        .map_err(|e| Error::UnexpectedError(e.to_string()))
}
<commit_msg>[config][management] clarify network address related errors<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    constants,
    error::Error,
    secure_backend::StorageLocation::{LocalStorage, RemoteStorage},
    SecureBackends,
};
use libra_config::config::HANDSHAKE_VERSION;
use libra_crypto::{ed25519::Ed25519PublicKey, x25519, ValidCryptoMaterial};
use libra_global_constants::{
    CONSENSUS_KEY, FULLNODE_NETWORK_KEY, OPERATOR_KEY, VALIDATOR_NETWORK_KEY,
};
use libra_network_address::{
    encrypted::{RawEncNetworkAddress, TEST_ROOT_KEY, TEST_ROOT_KEY_VERSION},
    NetworkAddress, RawNetworkAddress,
};
use libra_secure_storage::{CryptoStorage, KVStorage, Storage, Value};
use libra_secure_time::{RealTimeService, TimeService};
use libra_types::{
    account_address::{self, AccountAddress},
    transaction::{RawTransaction, SignedTransaction, Transaction},
};
use std::{convert::TryFrom, time::Duration};
use structopt::StructOpt;

// TODO(davidiw) add operator_address, since that will eventually be the identity producing this.
#[derive(Debug, StructOpt)]
pub struct ValidatorConfig {
    #[structopt(long)]
    owner_address: AccountAddress,
    #[structopt(long)]
    validator_address: NetworkAddress,
    #[structopt(long)]
    fullnode_address: NetworkAddress,
    #[structopt(flatten)]
    backends: SecureBackends,
}

impl ValidatorConfig {
    pub fn execute(self) -> Result<Transaction, Error> {
        let mut local_storage = self.backends.local.create_storage(LocalStorage)?;

        // Step 1) Retrieve keys from local storage
        let consensus_key = ed25519_from_storage(CONSENSUS_KEY, &local_storage)?;
        let fullnode_network_key = x25519_from_storage(FULLNODE_NETWORK_KEY, &local_storage)?;
        let validator_network_key = x25519_from_storage(VALIDATOR_NETWORK_KEY, &local_storage)?;
        let operator_key = ed25519_from_storage(OPERATOR_KEY, &local_storage)?;

        // TODO(davidiw): In genesis this is irrelevant -- afterward we need to obtain the
        // current sequence number by querying the blockchain.
        let sequence_number = 0;

        // TODO(davidiw): This is currently not supported
        // let sender = self.owner_address;
        let sender = account_address::from_public_key(&operator_key);

        // TODO(philiphayes): actually get the real keys from storage
        let root_key = TEST_ROOT_KEY;
        let key_version = TEST_ROOT_KEY_VERSION;
        // only supports one address for now
        let addr_idx = 0;

        // append ln-noise-ik and ln-handshake protocols to base network addresses
        // and encrypt the validator address.

        let validator_address = self
            .validator_address
            .clone()
            .append_prod_protos(validator_network_key, HANDSHAKE_VERSION);
        let raw_validator_address =
            RawNetworkAddress::try_from(&validator_address).map_err(|e| {
                Error::UnexpectedError(format!(
                    "error serializing validator address: \"{}\", error: {}",
                    validator_address, e
                ))
            })?;
        let enc_validator_address = raw_validator_address
            .encrypt(&root_key, key_version, &sender, sequence_number, addr_idx)
            .map_err(|e| {
                Error::UnexpectedError(format!(
                    "error encrypting validator address: \"{}\", error: {}",
                    validator_address, e
                ))
            })?;
        let raw_enc_validator_address = RawEncNetworkAddress::try_from(&enc_validator_address)
            .map_err(|e| {
                Error::UnexpectedError(format!(
                    "error serializing encrypted validator address: {:?}, error: {}",
                    enc_validator_address, e
                ))
            })?;

        let fullnode_address = self
            .fullnode_address
            .clone()
            .append_prod_protos(fullnode_network_key, HANDSHAKE_VERSION);
        let raw_fullnode_address = RawNetworkAddress::try_from(&fullnode_address).map_err(|e| {
            Error::UnexpectedError(format!(
                "error serializing fullnode address: \"{}\", error: {}",
                fullnode_address, e
            ))
        })?;

        // Step 2) Generate transaction

        // TODO(philiphayes): remove network identity pubkey field from struct when
        // transition complete
        let script = transaction_builder::encode_set_validator_config_script(
            sender,
            consensus_key.to_bytes().to_vec(),
            validator_network_key.to_bytes(),
            raw_enc_validator_address.into(),
            fullnode_network_key.to_bytes(),
            raw_fullnode_address.into(),
        );
        let expiration_time = RealTimeService::new().now() + constants::TXN_EXPIRATION_SECS;
        let raw_transaction = RawTransaction::new_script(
            sender,
            sequence_number,
            script,
            constants::MAX_GAS_AMOUNT,
            constants::GAS_UNIT_PRICE,
            constants::GAS_CURRENCY_CODE.to_owned(),
            Duration::from_secs(expiration_time),
        );
        let signature = local_storage
            .sign(OPERATOR_KEY, &raw_transaction)
            .map_err(|e| {
                Error::LocalStorageSigningError("validator-config", OPERATOR_KEY, e.to_string())
            })?;
        let signed_txn = SignedTransaction::new(raw_transaction, operator_key, signature);
        let txn = Transaction::UserTransaction(signed_txn);

        // Step 3) Submit to remote storage

        if let Some(remote_config) = self.backends.remote {
            let mut remote_storage = remote_config.create_storage(RemoteStorage)?;
            let txn = Value::Transaction(txn.clone());
            remote_storage
                .set(constants::VALIDATOR_CONFIG, txn)
                .map_err(|e| {
                    Error::RemoteStorageWriteError(constants::VALIDATOR_CONFIG, e.to_string())
                })?;
        }

        Ok(txn)
    }
}

fn ed25519_from_storage(
    key_name: &'static str,
    storage: &Storage,
) -> Result<Ed25519PublicKey, Error> {
    Ok(storage
        .get_public_key(key_name)
        .map_err(|e| Error::LocalStorageReadError(key_name, e.to_string()))?
        .public_key)
}

fn x25519_from_storage(
    key_name: &'static str,
    storage: &Storage,
) -> Result<x25519::PublicKey, Error> {
    let edkey = ed25519_from_storage(key_name, storage)?;
    x25519::PublicKey::from_ed25519_public_bytes(&edkey.to_bytes())
        .map_err(|e| Error::UnexpectedError(e.to_string()))
}
<|endoftext|>


<commit_before>use super::*;
use std::collections::{BTreeMap, VecDeque};

#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub struct TermId {
    address: Vec<usize>,
    compos_var_pos: usize,
}

impl From<(Vec<usize>, usize)> for TermId {
    fn from((address, compos_var_pos): (Vec<usize>, usize)) -> TermId {
        TermId { address, compos_var_pos }
    }
}

pub fn identify_terminals<A>(tree_map: &GornTree<Composition<A>>)
        -> (GornTree<Composition<TermId>>, BTreeMap<TermId, A>)
    where A: Clone,
{
    let mut identified_tree_map = GornTree::new();
    let mut terminal_map = BTreeMap::new();

    for (address, composition) in tree_map {
        let vec_compos = &composition.composition;
        let mut identified_compos = Vec::new();
        let mut compos_var_pos = 0;

        for component in vec_compos {
            let mut identified_compon = Vec::new();

            for variable in component {
                match variable {
                    &VarT::Var(x, y) => {
                        identified_compon.push(VarT::Var(x, y));
                    },
                    &VarT::T(ref terminal) => {
                        let terminal_id = TermId { address: address.clone(), compos_var_pos };
                        identified_compon.push(VarT::T(terminal_id.clone()));
                        terminal_map.insert(terminal_id, terminal.clone());
                    },
                };

                compos_var_pos += 1;
            }

            identified_compos.push(identified_compon);
        }

        identified_tree_map.insert(address.clone(), Composition::from(identified_compos));
    }

    (identified_tree_map, terminal_map)
}

pub fn to_negra<H, T, W>(tree_map: &GornTree<PMCFGRule<H, T, W>>, sentence_id: usize)
        -> String
    where H: Clone + ToString,
          T: Clone + ToString,
{
    let negra_vector = to_negra_vector(&tree_map);
    let mut output = format!("#BOS {}\n", sentence_id);

    for (word, tag, parent) in negra_vector {
        output.push_str(&format!("{}\t{}\t--\t--\t{}\n", word, tag, parent));
    }

    output.push_str(&format!("#EOS {}", sentence_id));
    output
}

pub fn to_negra_vector<H, T, W>(tree_map: &GornTree<PMCFGRule<H, T, W>>)
        -> Vec<(String, String, usize)>
    where H: Clone + ToString,
          T: Clone + ToString,
{
    let (term_map, nonterminal_map) = to_term(&tree_map);
    // TODO: Enforce negra grammar restrictions
    let (identified_tree_map, terminal_map) = identify_terminals(&term_map);
    let evaluated_compos = evaluate(&identified_tree_map);

    let mut negra_vector = Vec::new();
    let mut rule_queue = VecDeque::new();
    let mut rule_number_map = GornTree::new();
    let mut rule_counter = 1;

    for component in evaluated_compos.composition {
        for variable in component {
            match variable {
                VarT::Var(_, _) => {
                    panic!("Nonterminals must not appear in a fully evaluated configuration!");
                },
                VarT::T(terminal_id) => {
                    let terminal_symbol = terminal_map.get(&terminal_id).unwrap();
                    let address = terminal_id.address;
                    let rule_label = nonterminal_map.get(&address).unwrap();

                    let mut parent_address = address;
                    let parent_number = if let None = parent_address.pop() {
                        panic!("Terminals must have a nonterminal-only rule as their parent!");
                    } else {
                        get_rule_number(&parent_address, &mut rule_queue, &mut rule_number_map,
                                        &mut rule_counter)
                    };
                    negra_vector.push((terminal_symbol.to_string(), rule_label.to_string(), parent_number));
                }
            }
        }
    }

    while let Some((address, rule_number)) = rule_queue.pop_front() {
        let mut parent_address = address.clone();

        let parent_number = if let None = parent_address.pop() {
            0
        } else {
            get_rule_number(&parent_address, &mut rule_queue, &mut rule_number_map, &mut rule_counter)
        };

        let rule_label = nonterminal_map.get(&address).unwrap();
        negra_vector.push((format!("#{}", rule_number), rule_label.to_string(), parent_number));
    }

    negra_vector
}

fn get_rule_number(address: &Vec<usize>, rule_queue: &mut VecDeque<(Vec<usize>, usize)>, rule_number_map: &mut GornTree<usize>, rule_counter: &mut usize)
        -> usize
{
    if let Some(rule_number) = rule_number_map.get(address) {
        return *rule_number
    }

    let rule_number = *rule_counter;
    *rule_counter = *rule_counter + 1;

    rule_number_map.insert(address.clone(), rule_number.clone());
    rule_queue.push_back((address.clone(), rule_number));
    rule_number
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::tests::*;
    use self::VarT::{Var, T};

    #[test]
    fn test_identify_terminals() {
        let mut tree_map = GornTree::new();
        tree_map.insert(vec![], Composition::from(vec![
            vec![Var(0, 0), T("a"), Var(0, 1), T("b")]
        ]));
        tree_map.insert(vec![0], Composition::from(vec![
            vec![Var(1, 0)],
            vec![T("c")]
        ]));
        tree_map.insert(vec![0, 1], Composition::from(vec![
            vec![T("d")]
        ]));

        let mut identified_tree_map = GornTree::new();
        identified_tree_map.insert(vec![], Composition::from(vec![
            vec![Var(0, 0), T(TermId::from((vec![], 1))), Var(0, 1), T(TermId::from((vec![], 3)))]
        ]));
        identified_tree_map.insert(vec![0], Composition::from(vec![
            vec![Var(1, 0)],
            vec![T(TermId::from((vec![0], 1)))]
        ]));
        identified_tree_map.insert(vec![0, 1], Composition::from(vec![
            vec![T(TermId::from((vec![0, 1], 0)))]
        ]));

        let mut terminal_map = BTreeMap::new();
        terminal_map.insert(TermId::from((vec![], 1)), "a");
        terminal_map.insert(TermId::from((vec![], 3)), "b");
        terminal_map.insert(TermId::from((vec![0], 1)), "c");
        terminal_map.insert(TermId::from((vec![0, 1], 0)), "d");

        assert_eq!((identified_tree_map, terminal_map), identify_terminals(&tree_map));
    }

    #[test]
    fn test_to_negra_vector() {
        let tree_map = example_tree_map();
        let negra_vector = vec![
            (String::from("a"), String::from("a"), 1),
            (String::from("a"), String::from("a"), 2),
            (String::from("b"), String::from("b"), 3),
            (String::from("c"), String::from("c"), 1),
            (String::from("c"), String::from("c"), 2),
            (String::from("d"), String::from("d"), 3),
            (String::from("#1"), String::from("A"), 4),
            (String::from("#2"), String::from("A"), 1),
            (String::from("#3"), String::from("B"), 4),
            (String::from("#4"), String::from("S"), 0)
        ];

        assert_eq!(negra_vector, to_negra_vector(&tree_map));
    }
}
<commit_msg>src/pmcfg/negra.rs: enforced negra grammar restrictions<commit_after>use super::*;
use std::collections::{BTreeMap, VecDeque};

#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub struct TermId {
    address: Vec<usize>,
    compos_var_pos: usize,
}

impl From<(Vec<usize>, usize)> for TermId {
    fn from((address, compos_var_pos): (Vec<usize>, usize)) -> TermId {
        TermId { address, compos_var_pos }
    }
}

pub fn identify_terminals<A>(tree_map: &GornTree<Composition<A>>)
        -> (GornTree<Composition<TermId>>, BTreeMap<TermId, A>)
    where A: Clone,
{
    let mut identified_tree_map = GornTree::new();
    let mut terminal_map = BTreeMap::new();

    for (address, composition) in tree_map {
        let vec_compos = &composition.composition;
        let mut identified_compos = Vec::new();
        let mut compos_var_pos = 0;

        for component in vec_compos {
            let mut identified_compon = Vec::new();

            for variable in component {
                match variable {
                    &VarT::Var(x, y) => {
                        identified_compon.push(VarT::Var(x, y));
                    },
                    &VarT::T(ref terminal) => {
                        let terminal_id = TermId { address: address.clone(), compos_var_pos };
                        identified_compon.push(VarT::T(terminal_id.clone()));
                        terminal_map.insert(terminal_id, terminal.clone());
                    },
                };

                compos_var_pos += 1;
            }

            identified_compos.push(identified_compon);
        }

        identified_tree_map.insert(address.clone(), Composition::from(identified_compos));
    }

    (identified_tree_map, terminal_map)
}

pub fn to_negra<H, T, W>(tree_map: &GornTree<PMCFGRule<H, T, W>>, sentence_id: usize)
        -> String
    where H: Clone + ToString,
          T: Clone + ToString,
{
    if !meets_negra_criteria(&tree_map) {
        panic!("The given tree does not meet the negra criteria! All rules must either consist \
                only of nonterminals or of exactly one terminal symbol.");
    }

    let negra_vector = to_negra_vector(&tree_map);
    let mut output = format!("#BOS {}\n", sentence_id);

    for (word, tag, parent) in negra_vector {
        output.push_str(&format!("{}\t{}\t--\t--\t{}\n", word, tag, parent));
    }

    output.push_str(&format!("#EOS {}", sentence_id));
    output
}

pub fn meets_negra_criteria<H, T, W>(tree_map: &GornTree<PMCFGRule<H, T, W>>)
        -> bool
{
    for (_address, rule) in tree_map {
        let &PMCFGRule { head: _, tail: _, ref composition, weight: _ } = rule;
        let mut contains_nonterminal = false;
        let mut contains_terminal = false;

        for component in &composition.composition {
            for variable in component {
                match variable {
                    &VarT::Var(_, _) => {
                        if contains_terminal {
                            return false;
                        }

                        contains_nonterminal = true;
                    },
                    &VarT::T(_) => {
                        if contains_nonterminal || contains_terminal {
                            return false;
                        }

                        contains_terminal = true;
                    },
                }
            }
        }
    }

    true
}

pub fn to_negra_vector<H, T, W>(tree_map: &GornTree<PMCFGRule<H, T, W>>)
        -> Vec<(String, String, usize)>
    where H: Clone + ToString,
          T: Clone + ToString,
{
    let (term_map, nonterminal_map) = to_term(&tree_map);
    let (identified_tree_map, terminal_map) = identify_terminals(&term_map);
    let evaluated_compos = evaluate(&identified_tree_map);

    let mut negra_vector = Vec::new();
    let mut rule_queue = VecDeque::new();
    let mut rule_number_map = GornTree::new();
    let mut rule_counter = 1;

    for component in evaluated_compos.composition {
        for variable in component {
            match variable {
                VarT::Var(_, _) => {
                    panic!("Nonterminals must not appear in a fully evaluated configuration!");
                },
                VarT::T(terminal_id) => {
                    let terminal_symbol = terminal_map.get(&terminal_id).unwrap();
                    let address = terminal_id.address;
                    let rule_label = nonterminal_map.get(&address).unwrap();

                    let mut parent_address = address;
                    let parent_number = if let None = parent_address.pop() {
                        panic!("Terminals must have a nonterminal-only rule as their parent!");
                    } else {
                        get_rule_number(&parent_address, &mut rule_queue, &mut rule_number_map,
                                        &mut rule_counter)
                    };
                    negra_vector.push((terminal_symbol.to_string(), rule_label.to_string(), parent_number));
                }
            }
        }
    }

    while let Some((address, rule_number)) = rule_queue.pop_front() {
        let mut parent_address = address.clone();

        let parent_number = if let None = parent_address.pop() {
            0
        } else {
            get_rule_number(&parent_address, &mut rule_queue, &mut rule_number_map, &mut rule_counter)
        };

        let rule_label = nonterminal_map.get(&address).unwrap();
        negra_vector.push((format!("#{}", rule_number), rule_label.to_string(), parent_number));
    }

    negra_vector
}

fn get_rule_number(address: &Vec<usize>, rule_queue: &mut VecDeque<(Vec<usize>, usize)>, rule_number_map: &mut GornTree<usize>, rule_counter: &mut usize)
        -> usize
{
    if let Some(rule_number) = rule_number_map.get(address) {
        return *rule_number
    }

    let rule_number = *rule_counter;
    *rule_counter = *rule_counter + 1;

    rule_number_map.insert(address.clone(), rule_number.clone());
    rule_queue.push_back((address.clone(), rule_number));
    rule_number
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::tests::*;
    use self::VarT::{Var, T};

    #[test]
    fn test_identify_terminals() {
        let mut tree_map = GornTree::new();
        tree_map.insert(vec![], Composition::from(vec![
            vec![Var(0, 0), T("a"), Var(0, 1), T("b")]
        ]));
        tree_map.insert(vec![0], Composition::from(vec![
            vec![Var(1, 0)],
            vec![T("c")]
        ]));
        tree_map.insert(vec![0, 1], Composition::from(vec![
            vec![T("d")]
        ]));

        let mut identified_tree_map = GornTree::new();
        identified_tree_map.insert(vec![], Composition::from(vec![
            vec![Var(0, 0), T(TermId::from((vec![], 1))), Var(0, 1), T(TermId::from((vec![], 3)))]
        ]));
        identified_tree_map.insert(vec![0], Composition::from(vec![
            vec![Var(1, 0)],
            vec![T(TermId::from((vec![0], 1)))]
        ]));
        identified_tree_map.insert(vec![0, 1], Composition::from(vec![
            vec![T(TermId::from((vec![0, 1], 0)))]
        ]));

        let mut terminal_map = BTreeMap::new();
        terminal_map.insert(TermId::from((vec![], 1)), "a");
        terminal_map.insert(TermId::from((vec![], 3)), "b");
        terminal_map.insert(TermId::from((vec![0], 1)), "c");
        terminal_map.insert(TermId::from((vec![0, 1], 0)), "d");

        assert_eq!((identified_tree_map, terminal_map), identify_terminals(&tree_map));
    }

    #[test]
    fn test_to_negra_vector() {
        let tree_map = example_tree_map();
        let negra_vector = vec![
            (String::from("a"), String::from("a"), 1),
            (String::from("a"), String::from("a"), 2),
            (String::from("b"), String::from("b"), 3),
            (String::from("c"), String::from("c"), 1),
            (String::from("c"), String::from("c"), 2),
            (String::from("d"), String::from("d"), 3),
            (String::from("#1"), String::from("A"), 4),
            (String::from("#2"), String::from("A"), 1),
            (String::from("#3"), String::from("B"), 4),
            (String::from("#4"), String::from("S"), 0)
        ];

        assert_eq!(negra_vector, to_negra_vector(&tree_map));
    }
}
<|endoftext|>


<commit_before>use prolog_parser::ast::*;
use prolog_parser::parser::*;
use prolog_parser::tabled_rc::TabledData;

use prolog::forms::*;
use prolog::iterators::*;
use prolog::machine::machine_indices::*;
use prolog::machine::machine_state::MachineState;

use std::collections::VecDeque;
use std::io::Read;

type SubtermDeque = VecDeque<(usize, usize)>;

impl<'a> TermRef<'a> {
    fn as_addr(&self, h: usize) -> Addr {
        match self {
            &TermRef::AnonVar(_) | &TermRef::Var(..) => Addr::HeapCell(h),
            &TermRef::Cons(..) => Addr::HeapCell(h),
            &TermRef::Constant(_, _, c) => Addr::Con(c.clone()),
            &TermRef::Clause(..) => Addr::Str(h),
        }
    }
}

pub type PrologStream = ParsingStream<Box<Read>>;

#[cfg(feature = "readline_rs_compat")]
pub mod readline
{
    use prolog_parser::ast::*;
    use readline_rs_compat::readline::*;
    use std::io::{Error, Read};

    #[derive(Clone, Copy)]
    pub enum LineMode {
        Single,
        Multi
    }

    pub struct ReadlineStream {
        pending_input: String
    }

    impl ReadlineStream {
        #[inline]
        fn new(pending_input: String) -> Self {
            ReadlineStream { pending_input }
        }

        fn call_readline(&mut self, prompt: &str, buf: &mut [u8]) -> std::io::Result<usize> {
            match readline_rl(prompt) {
                Some(text) => {
                    self.pending_input += &text;
                    Ok(self.write_to_buf(buf))
                },
                None => Err(Error::last_os_error())
            }
        }

        fn split_pending(&mut self, buf: &mut [u8], split_idx: usize) -> usize {
            let (outgoing, _) = self.pending_input.split_at(split_idx);

            for (idx, b) in outgoing.bytes().enumerate() {
                buf[idx] = b;
            }

            outgoing.len()
        }

        fn write_to_buf(&mut self, buf: &mut [u8]) -> usize {
            let split_idx = std::cmp::min(self.pending_input.len(), buf.len());
            let output_len = self.split_pending(buf, split_idx);

            if split_idx < self.pending_input.len() {
                self.pending_input = self.pending_input[split_idx ..].to_string();
            } else {
                self.pending_input.clear();
            }

            output_len
        }
    }

    impl Read for ReadlineStream {
        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
            if self.pending_input.is_empty() {
                self.call_readline("", buf)
            } else {
                Ok(self.write_to_buf(buf))
            }
        }
    }

    static mut LINE_MODE: LineMode = LineMode::Single;
    static mut END_OF_LINE: bool = false;

    pub fn set_line_mode(mode: LineMode) {
        unsafe {
            LINE_MODE = mode;
            END_OF_LINE = false;
            rl_done = 0;
        }
    }

    unsafe extern "C" fn bind_end_chord(_: i32, _: i32) -> i32 {
        if let LineMode::Multi = LINE_MODE {
            rl_done = 1;
        }

        0
    }

    unsafe extern "C" fn bind_cr(_: i32, _: i32) -> i32 {
        if let LineMode::Single = LINE_MODE {
            insert_text_rl("\n");
            println!("");
            rl_done = 1;
        } else {
            insert_text_rl("\n");
        }

        0
    }

    pub fn readline_initialize() {
        let rc = initialize_rl(); // initialize editline.

        if rc != 0 {
            panic!("initialize_rl() failed with return code {}", rc);
        }

        bind_key_rl('\t' as i32, rl_insert); // just insert tabs when typed.
        bind_key_rl('\n' as i32, bind_cr);
        bind_key_rl('\r' as i32, bind_cr);
        bind_keyseq_rl("\\C-d", bind_end_chord);
    }

    #[inline]
    pub fn input_stream() -> ::PrologStream {
        let reader: Box<Read> = Box::new(ReadlineStream::new(String::from("")));
        parsing_stream(reader)
    }
}

#[cfg(not(feature = "readline_rs_compat"))]
pub mod readline
{
    use prolog_parser::ast::*;
    use std::io::{BufReader, Read, Stdin, stdin};

    struct StdinWrapper {
        buf: BufReader<Stdin>
    }

    impl Read for StdinWrapper {
        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
            self.buf.read(buf)
        }
    }

    pub fn read_batch(_: &str) -> Result<Vec<u8>, ::SessionError> {
        let mut buf = vec![];

        let stdin = stdin();
        let mut stdin = stdin.lock();

        match stdin.read_to_end(&mut buf) {
            Ok(_) => Ok(buf),
            _ => Err(::SessionError::UserPrompt)
        }
    }

    #[inline]
    pub fn input_stream() -> ::PrologStream {
        let reader: Box<Read> = Box::new(StdinWrapper { buf: BufReader::new(stdin()) });
        parsing_stream(reader)
    }
}

impl MachineState {
    pub fn read(&mut self, inner: &mut PrologStream, atom_tbl: TabledData<Atom>, op_dir: &OpDir)
                -> Result<TermWriteResult, ParserError>
    {
        let mut parser = Parser::new(inner, atom_tbl, self.flags);
        let term = parser.read_term(composite_op!(op_dir))?;

        Ok(write_term_to_heap(&term, self))
    }
}

fn push_stub_addr(machine_st: &mut MachineState) {
    let h = machine_st.heap.h;
    machine_st.heap.push(HeapCellValue::Addr(Addr::HeapCell(h)));
}

fn modify_head_of_queue(machine_st: &mut MachineState, queue: &mut SubtermDeque, term: TermRef, h: usize)
{
    if let Some((arity, site_h)) = queue.pop_front() {
        machine_st.heap[site_h] = HeapCellValue::Addr(term.as_addr(h));

        if arity > 1 {
            queue.push_front((arity - 1, site_h + 1));
        }
    }
}

pub struct TermWriteResult {
    pub(crate) heap_loc: usize,
    pub(crate) var_dict: HeapVarDict,
}

pub(crate)
fn write_term_to_heap(term: &Term, machine_st: &mut MachineState) -> TermWriteResult
{
    let heap_loc = machine_st.heap.h;

    let mut queue = SubtermDeque::new();
    let mut var_dict = HeapVarDict::new();

    for term in breadth_first_iter(term, true) {
        let h = machine_st.heap.h;

        match &term {
            &TermRef::Cons(lvl, ..) => {
                queue.push_back((2, h+1));
                machine_st.heap.push(HeapCellValue::Addr(Addr::Lis(h+1)));

                push_stub_addr(machine_st);
                push_stub_addr(machine_st);

                if let Level::Root = lvl {
                    continue;
                }
            },
            &TermRef::Clause(lvl, _, ref ct, subterms) => {
                queue.push_back((subterms.len(), h+1));
                let named = HeapCellValue::NamedStr(subterms.len(), ct.name(), ct.spec());

                machine_st.heap.push(named);

                for _ in 0 .. subterms.len() {
                    push_stub_addr(machine_st);
                }

                if let Level::Root = lvl {
                    continue;
                }
            },
            &TermRef::AnonVar(Level::Root) | &TermRef::Constant(Level::Root, ..) =>
                machine_st.heap.push(HeapCellValue::Addr(term.as_addr(h))),
            &TermRef::Var(Level::Root, ..) =>
                machine_st.heap.push(HeapCellValue::Addr(term.as_addr(h))),
            &TermRef::AnonVar(_) => {
                if let Some((arity, site_h)) = queue.pop_front() {
                    if arity > 1 {
                        queue.push_front((arity - 1, site_h + 1));
                    }
                }

                continue;
            },
            &TermRef::Var(_, _, ref var) => {
                if let Some((arity, site_h)) = queue.pop_front() {
                    if let Some(addr) = var_dict.get(var).cloned() {
                        machine_st.heap[site_h] = HeapCellValue::Addr(addr);
                    } else {
                        var_dict.insert(var.clone(), Addr::HeapCell(site_h));
                    }

                    if arity > 1 {
                        queue.push_front((arity - 1, site_h + 1));
                    }
                }

                continue;
            },
            _ => {}
        };

        modify_head_of_queue(machine_st, &mut queue, term, h);
    }

    TermWriteResult { heap_loc, var_dict }
}
<commit_msg>fix read module if readline is not present<commit_after>use prolog_parser::ast::*;
use prolog_parser::parser::*;
use prolog_parser::tabled_rc::TabledData;

use prolog::forms::*;
use prolog::iterators::*;
use prolog::machine::machine_indices::*;
use prolog::machine::machine_state::MachineState;

use std::collections::VecDeque;
use std::io::Read;

type SubtermDeque = VecDeque<(usize, usize)>;

impl<'a> TermRef<'a> {
    fn as_addr(&self, h: usize) -> Addr {
        match self {
            &TermRef::AnonVar(_) | &TermRef::Var(..) => Addr::HeapCell(h),
            &TermRef::Cons(..) => Addr::HeapCell(h),
            &TermRef::Constant(_, _, c) => Addr::Con(c.clone()),
            &TermRef::Clause(..) => Addr::Str(h),
        }
    }
}

pub type PrologStream = ParsingStream<Box<Read>>;

#[cfg(feature = "readline_rs_compat")]
pub mod readline
{
    use prolog_parser::ast::*;
    use readline_rs_compat::readline::*;
    use std::io::{Error, Read};

    #[derive(Clone, Copy)]
    pub enum LineMode {
        Single,
        Multi
    }

    pub struct ReadlineStream {
        pending_input: String
    }

    impl ReadlineStream {
        #[inline]
        fn new(pending_input: String) -> Self {
            ReadlineStream { pending_input }
        }

        fn call_readline(&mut self, prompt: &str, buf: &mut [u8]) -> std::io::Result<usize> {
            match readline_rl(prompt) {
                Some(text) => {
                    self.pending_input += &text;
                    Ok(self.write_to_buf(buf))
                },
                None => Err(Error::last_os_error())
            }
        }

        fn split_pending(&mut self, buf: &mut [u8], split_idx: usize) -> usize {
            let (outgoing, _) = self.pending_input.split_at(split_idx);

            for (idx, b) in outgoing.bytes().enumerate() {
                buf[idx] = b;
            }

            outgoing.len()
        }

        fn write_to_buf(&mut self, buf: &mut [u8]) -> usize {
            let split_idx = std::cmp::min(self.pending_input.len(), buf.len());
            let output_len = self.split_pending(buf, split_idx);

            if split_idx < self.pending_input.len() {
                self.pending_input = self.pending_input[split_idx ..].to_string();
            } else {
                self.pending_input.clear();
            }

            output_len
        }
    }

    impl Read for ReadlineStream {
        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
            if self.pending_input.is_empty() {
                self.call_readline("", buf)
            } else {
                Ok(self.write_to_buf(buf))
            }
        }
    }

    static mut LINE_MODE: LineMode = LineMode::Single;
    static mut END_OF_LINE: bool = false;

    pub fn set_line_mode(mode: LineMode) {
        unsafe {
            LINE_MODE = mode;
            END_OF_LINE = false;
            rl_done = 0;
        }
    }

    unsafe extern "C" fn bind_end_chord(_: i32, _: i32) -> i32 {
        if let LineMode::Multi = LINE_MODE {
            rl_done = 1;
        }

        0
    }

    unsafe extern "C" fn bind_cr(_: i32, _: i32) -> i32 {
        if let LineMode::Single = LINE_MODE {
            insert_text_rl("\n");
            println!("");
            rl_done = 1;
        } else {
            insert_text_rl("\n");
        }

        0
    }

    pub fn readline_initialize() {
        let rc = initialize_rl(); // initialize editline.

        if rc != 0 {
            panic!("initialize_rl() failed with return code {}", rc);
        }

        bind_key_rl('\t' as i32, rl_insert); // just insert tabs when typed.
        bind_key_rl('\n' as i32, bind_cr);
        bind_key_rl('\r' as i32, bind_cr);
        bind_keyseq_rl("\\C-d", bind_end_chord);
    }

    #[inline]
    pub fn input_stream() -> ::PrologStream {
        let reader: Box<Read> = Box::new(ReadlineStream::new(String::from("")));
        parsing_stream(reader)
    }
}

#[cfg(not(feature = "readline_rs_compat"))]
pub mod readline
{
    use prolog_parser::ast::*;
    use std::io::{BufReader, Read, Stdin, stdin};

    struct StdinWrapper {
        buf: BufReader<Stdin>
    }

    impl Read for StdinWrapper {
        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
            self.buf.read(buf)
        }
    }

    #[inline]
    pub fn input_stream() -> ::PrologStream {
        let reader: Box<Read> = Box::new(StdinWrapper { buf: BufReader::new(stdin()) });
        parsing_stream(reader)
    }
}

impl MachineState {
    pub fn read(&mut self, inner: &mut PrologStream, atom_tbl: TabledData<Atom>, op_dir: &OpDir)
                -> Result<TermWriteResult, ParserError>
    {
        let mut parser = Parser::new(inner, atom_tbl, self.flags);
        let term = parser.read_term(composite_op!(op_dir))?;

        Ok(write_term_to_heap(&term, self))
    }
}

fn push_stub_addr(machine_st: &mut MachineState) {
    let h = machine_st.heap.h;
    machine_st.heap.push(HeapCellValue::Addr(Addr::HeapCell(h)));
}

fn modify_head_of_queue(machine_st: &mut MachineState, queue: &mut SubtermDeque, term: TermRef, h: usize)
{
    if let Some((arity, site_h)) = queue.pop_front() {
        machine_st.heap[site_h] = HeapCellValue::Addr(term.as_addr(h));

        if arity > 1 {
            queue.push_front((arity - 1, site_h + 1));
        }
    }
}

pub struct TermWriteResult {
    pub(crate) heap_loc: usize,
    pub(crate) var_dict: HeapVarDict,
}

pub(crate)
fn write_term_to_heap(term: &Term, machine_st: &mut MachineState) -> TermWriteResult
{
    let heap_loc = machine_st.heap.h;

    let mut queue = SubtermDeque::new();
    let mut var_dict = HeapVarDict::new();

    for term in breadth_first_iter(term, true) {
        let h = machine_st.heap.h;

        match &term {
            &TermRef::Cons(lvl, ..) => {
                queue.push_back((2, h+1));
                machine_st.heap.push(HeapCellValue::Addr(Addr::Lis(h+1)));

                push_stub_addr(machine_st);
                push_stub_addr(machine_st);

                if let Level::Root = lvl {
                    continue;
                }
            },
            &TermRef::Clause(lvl, _, ref ct, subterms) => {
                queue.push_back((subterms.len(), h+1));
                let named = HeapCellValue::NamedStr(subterms.len(), ct.name(), ct.spec());

                machine_st.heap.push(named);

                for _ in 0 .. subterms.len() {
                    push_stub_addr(machine_st);
                }

                if let Level::Root = lvl {
                    continue;
                }
            },
            &TermRef::AnonVar(Level::Root) | &TermRef::Constant(Level::Root, ..) =>
                machine_st.heap.push(HeapCellValue::Addr(term.as_addr(h))),
            &TermRef::Var(Level::Root, ..) =>
                machine_st.heap.push(HeapCellValue::Addr(term.as_addr(h))),
            &TermRef::AnonVar(_) => {
                if let Some((arity, site_h)) = queue.pop_front() {
                    if arity > 1 {
                        queue.push_front((arity - 1, site_h + 1));
                    }
                }

                continue;
            },
            &TermRef::Var(_, _, ref var) => {
                if let Some((arity, site_h)) = queue.pop_front() {
                    if let Some(addr) = var_dict.get(var).cloned() {
                        machine_st.heap[site_h] = HeapCellValue::Addr(addr);
                    } else {
                        var_dict.insert(var.clone(), Addr::HeapCell(site_h));
                    }

                    if arity > 1 {
                        queue.push_front((arity - 1, site_h + 1));
                    }
                }

                continue;
            },
            _ => {}
        };

        modify_head_of_queue(machine_st, &mut queue, term, h);
    }

    TermWriteResult { heap_loc, var_dict }
}
<|endoftext|>


<commit_before>// Copyright (C) 2020 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

mod address;

use super::{
    DirectDownload, EntryPoint, Metadata, PrepareLocalInstall, Result, RuntimeSettings, Settings,
    State, StateChangeImpl, Validation,
};
use async_std::{prelude::FutureExt, sync};
use slog_scope::trace;

pub(crate) use address::{
    AbortDownloadResponse, Addr, Message, ProbeResponse, Response, StateResponse,
};

pub(super) struct StateMachine {
    state: State,
    context: Context,
}

pub struct Context {
    pub(super) communication: Channel<(Message, sync::Sender<Result<Response>>)>,
    pub(super) waker: Channel<()>,
    pub settings: Settings,
    pub runtime_settings: RuntimeSettings,
    pub firmware: Metadata,
}

pub(super) struct Channel<T> {
    sender: sync::Sender<T>,
    receiver: sync::Receiver<T>,
}

impl<T> Channel<T> {
    fn new(cap: usize) -> Self {
        let (sender, receiver) = sync::channel(cap);
        Channel { sender, receiver }
    }
}

impl Context {
    pub(crate) fn new(
        settings: Settings,
        runtime_settings: RuntimeSettings,
        firmware: Metadata,
    ) -> Self {
        Context {
            communication: Channel::new(10),
            waker: Channel::new(1),
            settings,
            runtime_settings,
            firmware,
        }
    }

    pub(super) fn server_address(&self) -> &str {
        self.runtime_settings
            .custom_server_address()
            .unwrap_or(&self.settings.network.server_address)
    }
}

#[derive(Debug)]
pub(super) enum StepTransition {
    Delayed(std::time::Duration),
    Immediate,
    Never,
}

impl StateMachine {
    pub(super) fn new(
        state: State,
        settings: Settings,
        runtime_settings: RuntimeSettings,
        firmware: Metadata,
    ) -> Self {
        StateMachine { state, context: Context::new(settings, runtime_settings, firmware) }
    }

    pub(super) fn address(&self) -> Addr {
        Addr {
            message: self.context.communication.sender.clone(),
            waker: self.context.waker.sender.clone(),
        }
    }

    pub(super) async fn start(mut self) {
        loop {
            // Since the loop is already currently running, we can
            // discharges any wake message received.
            let _ = self.context.waker.receiver.try_recv();

            self.consume_pending_communication().await;

            let (state, transition) = self
                .state
                .move_to_next_state(&mut self.context)
                .await
                .unwrap_or_else(|e| (State::from(e), StepTransition::Immediate));
            self.state = state;

            match transition {
                StepTransition::Immediate => {}
                StepTransition::Delayed(t) => {
                    trace!("delaying transition for: {} seconds", t.as_secs());
                    let waker = self.context.waker.receiver.clone();
                    async_std::task::sleep(t)
                        .race(async {
                            let _ = waker.recv().await;
                        })
                        .race(self.await_communication())
                        .await;
                }
                StepTransition::Never => {
                    trace!("stopping transition until awoken");
                    let _ = self
                        .context
                        .waker
                        .receiver
                        .clone()
                        .recv()
                        .race(async {
                            self.await_communication().await;
                            Ok(())
                        })
                        .await;
                }
            }
        }
    }

    async fn consume_pending_communication(&mut self) {
        while let Ok((msg, responder)) = self.context.communication.receiver.try_recv() {
            self.handle_communication(msg, responder).await;
        }
    }

    async fn await_communication(&mut self) {
        while let Ok((msg, responder)) = self.context.communication.receiver.recv().await {
            self.handle_communication(msg, responder).await;
        }
    }

    async fn handle_communication(
        &mut self,
        msg: address::Message,
        responder: sync::Sender<Result<address::Response>>,
    ) {
        trace!("Received external request: {:?}", msg);

        let response = match msg {
            address::Message::Info => {
                let state = self.state.name().to_owned();
                Ok(address::Response::Info(sdk::api::info::Response {
                    state,
                    version: crate::version().to_string(),
                    config: self.context.settings.0.clone(),
                    firmware: self.context.firmware.0.clone(),
                    runtime_settings: self.context.runtime_settings.0.clone(),
                }))
            }
            address::Message::Probe(custom_server) => {
                self.handle_probe_request(custom_server).await.map(|r| address::Response::Probe(r))
            }
            address::Message::AbortDownload => {
                if self.state.is_handling_download() {
                    self.state = State::EntryPoint(EntryPoint {});
                    Ok(address::Response::AbortDownload(
                        address::AbortDownloadResponse::RequestAccepted,
                    ))
                } else {
                    Ok(address::Response::AbortDownload(
                        address::AbortDownloadResponse::InvalidState,
                    ))
                }
            }
            address::Message::LocalInstall(update_file) => {
                let state = self.state.name().to_owned();

                if self.state.is_preemptive_state() {
                    crate::logger::start_memory_logging();
                    self.context.waker.sender.send(()).await;

                    self.state = State::PrepareLocalInstall(PrepareLocalInstall { update_file });

                    Ok(address::Response::LocalInstall(address::StateResponse::RequestAccepted(
                        state,
                    )))
                } else {
                    Ok(address::Response::LocalInstall(address::StateResponse::InvalidState(state)))
                }
            }
            address::Message::RemoteInstall(url) => {
                let state = self.state.name().to_owned();

                if self.state.is_preemptive_state() {
                    crate::logger::start_memory_logging();
                    self.context.waker.sender.send(()).await;

                    self.state = State::DirectDownload(DirectDownload { url });

                    Ok(address::Response::RemoteInstall(address::StateResponse::RequestAccepted(
                        state,
                    )))
                } else {
                    Ok(address::Response::RemoteInstall(address::StateResponse::InvalidState(
                        state,
                    )))
                }
            }
        };

        responder.send(response).await;
    }

    async fn handle_probe_request(
        &mut self,
        custom_server: Option<String>,
    ) -> Result<address::ProbeResponse> {
        use chrono::Utc;
        use cloud::api::ProbeResponse;

        if !self.state.is_preemptive_state() {
            let state = self.state.name().to_owned();
            return Ok(address::ProbeResponse::Busy(state));
        }

        if let Some(server_address) = custom_server {
            self.context.runtime_settings.set_custom_server_address(&server_address);
        }

        match crate::CloudClient::new(&self.context.server_address())
            .probe(
                self.context.runtime_settings.retries() as u64,
                self.context.firmware.as_cloud_metadata(),
            )
            .await?
        {
            ProbeResponse::ExtraPoll(s) => Ok(address::ProbeResponse::Delayed(s)),

            ProbeResponse::NoUpdate => {
                self.context.waker.sender.send(()).await;

                // Store timestamp of last polling
                self.context.runtime_settings.set_last_polling(Utc::now())?;
                self.state = State::EntryPoint(EntryPoint {});
                Ok(address::ProbeResponse::Unavailable)
            }

            ProbeResponse::Update(package, sign) => {
                self.context.waker.sender.send(()).await;

                // Store timestamp of last polling
                self.context.runtime_settings.set_last_polling(Utc::now())?;
                self.state = State::Validation(Validation { package, sign });
                Ok(address::ProbeResponse::Available)
            }
        }
    }
}
<commit_msg>updatehub: CommunicationState trait<commit_after>// Copyright (C) 2020 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

mod address;

use super::{
    DirectDownload, EntryPoint, Metadata, PrepareLocalInstall, Result, RuntimeSettings, Settings,
    State, StateChangeImpl, Validation,
};
use async_std::{prelude::FutureExt, sync};
use slog_scope::trace;
use std::path::PathBuf;

pub(crate) use address::{
    AbortDownloadResponse, Addr, Message, ProbeResponse, Response, StateResponse,
};

pub(super) struct StateMachine {
    state: State,
    context: Context,
}

pub struct Context {
    pub(super) communication: Channel<(Message, sync::Sender<Result<Response>>)>,
    pub(super) waker: Channel<()>,
    pub settings: Settings,
    pub runtime_settings: RuntimeSettings,
    pub firmware: Metadata,
}

pub(super) struct Channel<T> {
    pub(super) sender: sync::Sender<T>,
    pub(super) receiver: sync::Receiver<T>,
}

impl<T> Channel<T> {
    fn new(cap: usize) -> Self {
        let (sender, receiver) = sync::channel(cap);
        Channel { sender, receiver }
    }
}

impl CommunicationState for State {}

#[async_trait::async_trait]
pub(super) trait CommunicationState: StateChangeImpl {
    async fn handle_communication(
        &self,
        msg: address::Message,
        responder: sync::Sender<Result<address::Response>>,
        context: &mut Context,
    ) -> Option<State> {
        trace!("Received external request: {:?}", msg);

        let res = match msg {
            address::Message::Info => {
                let state = self.name().to_owned();
                Ok((
                    address::Response::Info(sdk::api::info::Response {
                        state,
                        version: crate::version().to_string(),
                        config: context.settings.0.clone(),
                        firmware: context.firmware.0.clone(),
                        runtime_settings: context.runtime_settings.0.clone(),
                    }),
                    None,
                ))
            }
            address::Message::Probe(custom_server) => self
                .handle_probe(context, custom_server)
                .await
                .map(|(res, st)| (address::Response::Probe(res), st)),
            address::Message::AbortDownload => self
                .handle_abort_download(context)
                .await
                .map(|(res, st)| (address::Response::AbortDownload(res), st)),
            address::Message::LocalInstall(update_file) => self
                .handle_local_install(context, update_file)
                .await
                .map(|(res, st)| (address::Response::LocalInstall(res), st)),
            address::Message::RemoteInstall(url) => self
                .handle_remote_install(context, url)
                .await
                .map(|(res, st)| (address::Response::RemoteInstall(res), st)),
        };

        match res {
            Ok((response, state)) => {
                responder.send(Ok(response)).await;
                state
            }
            Err(e) => {
                responder.send(Err(e)).await;
                None
            }
        }
    }

    async fn handle_probe(
        &self,
        context: &mut Context,
        custom_server: Option<String>,
    ) -> Result<(address::ProbeResponse, Option<State>)> {
        use chrono::Utc;
        use cloud::api::ProbeResponse;

        if !self.is_preemptive_state() {
            return Ok((address::ProbeResponse::Busy(self.name().to_owned()), None));
        }

        if let Some(server_address) = custom_server {
            context.runtime_settings.set_custom_server_address(&server_address);
        }

        match crate::CloudClient::new(&context.server_address())
            .probe(context.runtime_settings.retries() as u64, context.firmware.as_cloud_metadata())
            .await?
        {
            ProbeResponse::ExtraPoll(s) => Ok((address::ProbeResponse::Delayed(s), None)),

            ProbeResponse::NoUpdate => {
                context.waker.sender.send(()).await;

                // Store timestamp of last polling
                context.runtime_settings.set_last_polling(Utc::now())?;
                Ok((address::ProbeResponse::Unavailable, Some(State::EntryPoint(EntryPoint {}))))
            }

            ProbeResponse::Update(package, sign) => {
                context.waker.sender.send(()).await;

                // Store timestamp of last polling
                context.runtime_settings.set_last_polling(Utc::now())?;
                Ok((
                    address::ProbeResponse::Available,
                    Some(State::Validation(Validation { package, sign })),
                ))
            }
        }
    }

    async fn handle_abort_download(
        &self,
        _: &Context,
    ) -> Result<(address::AbortDownloadResponse, Option<State>)> {
        if self.is_handling_download() {
            Ok((
                address::AbortDownloadResponse::RequestAccepted,
                Some(State::EntryPoint(EntryPoint {})),
            ))
        } else {
            Ok((address::AbortDownloadResponse::InvalidState, None))
        }
    }

    async fn handle_local_install(
        &self,
        context: &Context,
        update_file: PathBuf,
    ) -> Result<(address::StateResponse, Option<State>)> {
        let name = self.name().to_owned();
        if self.is_preemptive_state() {
            crate::logger::start_memory_logging();
            context.waker.sender.send(()).await;

            Ok((
                address::StateResponse::RequestAccepted(name),
                Some(State::PrepareLocalInstall(PrepareLocalInstall { update_file })),
            ))
        } else {
            Ok((address::StateResponse::InvalidState(name), None))
        }
    }

    async fn handle_remote_install(
        &self,
        context: &Context,
        url: String,
    ) -> Result<(address::StateResponse, Option<State>)> {
        let name = self.name().to_owned();

        if self.is_preemptive_state() {
            crate::logger::start_memory_logging();
            context.waker.sender.send(()).await;

            Ok((
                address::StateResponse::RequestAccepted(name),
                Some(State::DirectDownload(DirectDownload { url })),
            ))
        } else {
            Ok((address::StateResponse::InvalidState(name), None))
        }
    }
}

impl Context {
    pub(crate) fn new(
        settings: Settings,
        runtime_settings: RuntimeSettings,
        firmware: Metadata,
    ) -> Self {
        Context {
            communication: Channel::new(10),
            waker: Channel::new(1),
            settings,
            runtime_settings,
            firmware,
        }
    }

    pub(super) fn server_address(&self) -> &str {
        self.runtime_settings
            .custom_server_address()
            .unwrap_or(&self.settings.network.server_address)
    }
}

#[derive(Debug)]
pub(super) enum StepTransition {
    Delayed(std::time::Duration),
    Immediate,
    Never,
}

impl StateMachine {
    pub(super) fn new(
        state: State,
        settings: Settings,
        runtime_settings: RuntimeSettings,
        firmware: Metadata,
    ) -> Self {
        StateMachine { state, context: Context::new(settings, runtime_settings, firmware) }
    }

    pub(super) fn address(&self) -> Addr {
        Addr {
            message: self.context.communication.sender.clone(),
            waker: self.context.waker.sender.clone(),
        }
    }

    pub(super) async fn start(mut self) {
        loop {
            // Since the loop is already currently running, we can
            // discharges any wake message received.
            let _ = self.context.waker.receiver.try_recv();

            self.consume_pending_communication().await;

            let (state, transition) = self
                .state
                .move_to_next_state(&mut self.context)
                .await
                .unwrap_or_else(|e| (State::from(e), StepTransition::Immediate));
            self.state = state;

            match transition {
                StepTransition::Immediate => {}
                StepTransition::Delayed(t) => {
                    trace!("delaying transition for: {} seconds", t.as_secs());
                    let waker = self.context.waker.receiver.clone();
                    async_std::task::sleep(t)
                        .race(async {
                            let _ = waker.recv().await;
                        })
                        .race(self.await_communication())
                        .await;
                }
                StepTransition::Never => {
                    trace!("stopping transition until awoken");
                    let _ = self
                        .context
                        .waker
                        .receiver
                        .clone()
                        .recv()
                        .race(async {
                            self.await_communication().await;
                            Ok(())
                        })
                        .await;
                }
            }
        }
    }

    async fn consume_pending_communication(&mut self) {
        while let Ok((msg, responder)) = self.context.communication.receiver.try_recv() {
            if let Some(new_state) =
                self.state.handle_communication(msg, responder, &mut self.context).await
            {
                self.state = new_state;
            }
        }
    }

    async fn await_communication(&mut self) {
        while let Ok((msg, responder)) = self.context.communication.receiver.recv().await {
            if let Some(new_state) =
                self.state.handle_communication(msg, responder, &mut self.context).await
            {
                self.state = new_state;
            }
        }
    }
}
<|endoftext|>


<commit_before>// TODO: Make this non public
pub mod file_naming;

use std::collections::HashMap;
use std::fmt::{Display, Error, Formatter};
use self::file_naming::{FileName, FileNameInfo, FileType, FileNameParser};

pub struct BackupSet {
    pub file_type : FileType,
    pub time : String,
    pub start_time : String,
    pub end_time : String,
    pub compressed : bool,
    pub encrypted : bool,
    pub partial : bool,
    pub manifest_path : String,
    volumes_paths : HashMap<i32, String>,
    info_set : bool,    // true if informations are set
}

impl BackupSet {
    pub fn new() -> Self {
        BackupSet{
            file_type : FileType::Full,
            time : String::new(),
            start_time : String::new(),
            end_time : String::new(),
            compressed : false,
            encrypted : false,
            partial : false,
            manifest_path : String::new(),
            volumes_paths : HashMap::new(),
            info_set : false
        }
    }

    /// Add a filename to given set. Return true if it fits.
    ///
    /// The filename will match the given set if it has the right
    /// times and is of the right type. The information will be set
    /// from the first filename given.
    pub fn add_filename(&mut self, file_info : &FileNameInfo) -> bool {
        let pr = &file_info.info;
        let fname = file_info.file_name;

        if !self.info_set {
            self.set_info(&pr);
        }
        else {
            // check if the file is from the same backup set
            if self.file_type != pr.file_type ||
                self.time != pr.time ||
                self.start_time != pr.start_time ||
                self.end_time != pr.end_time {
                    return false;
            }
            else {
                // fix encrypted flag
                if self.encrypted != pr.encrypted &&
                    self.partial && pr.encrypted {
                        self.encrypted = pr.encrypted;
                    }
            }
        }
        // set manifest or volume number
        if pr.manifest {
            self.manifest_path = fname.to_owned();
        }
        else {
            self.volumes_paths.insert(pr.volume_number, fname.to_owned());
        }
        true
    }

    pub fn is_complete(&self) -> bool {
        !self.manifest_path.is_empty()
    }

    pub fn get_time(&self) -> &str {
        if self.time.is_empty() {
            self.end_time.as_ref()
        }
        else {
            self.time.as_ref()
        }
    }

    fn set_info(&mut self, fname : &FileName) {
        self.file_type = fname.file_type;
        self.time = fname.time.clone();
        self.start_time = fname.start_time.clone();
        self.end_time = fname.end_time.clone();
        self.compressed = fname.compressed;
        self.encrypted = fname.encrypted;
        self.partial = fname.partial;
        self.info_set = true;
    }
}

impl Display for BackupSet {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        match self.file_type {
            FileType::Full => {
                try!(write!(f, "Full, time: {}", self.time));
            },
            FileType::Inc => {
                try!(write!(f, "Incremental, start time: {}, end time: {}", self.start_time, self.end_time));
            },
            _ => { }
        }
        if self.compressed {
            try!(write!(f, ", compressed"));
        }
        if self.encrypted {
            try!(write!(f, ", encrypted"));
        }
        if self.partial {
            try!(write!(f, ", partial"));
        }
        if !self.manifest_path.is_empty() {
            try!(write!(f, "\n manifest: {}", self.manifest_path));
        }
        if !self.volumes_paths.is_empty() {
            try!(write!(f, "\n volumes:\n"));
            for (_, vol) in self.volumes_paths.iter() {
                try!(write!(f, " {}", vol));
            }
        }
        Ok(())
    }
}

pub struct BackupChain {
    pub fullset : BackupSet,
    pub incset_list : Vec<BackupSet>,
    pub start_time : String,
    pub end_time : String
}

impl BackupChain {
    /// Create a new BackupChain starting from a full backup set.
    pub fn new(fullset : BackupSet) -> Self {
        assert_eq!(fullset.file_type, FileType::Full);
        let time = fullset.time.clone();

        BackupChain{
            fullset : fullset,
            incset_list : Vec::new(),
            start_time : time.clone(),
            end_time : time
        }
    }

    /// Adds the given incremental backup element to the backup chain if possible,
    /// returns it back otherwise.
    pub fn add_inc(&mut self, incset : BackupSet) -> Option<BackupSet> {
        if self.end_time == incset.start_time {
            self.end_time = incset.end_time.clone();
            self.incset_list.push(incset);
            None
        }
        else {
            // replace the last element if the end time comes before
            let replace_last = self.incset_list.last().map_or(false,
                |last| incset.start_time == last.start_time && incset.end_time > last.end_time);
            if replace_last {
                self.end_time = incset.end_time.clone();
                self.incset_list.pop();
                self.incset_list.push(incset);
                None
            }
            else {
                // ignore the given incremental backup set
                Some(incset)
            }
        }
    }
}

impl Display for BackupChain {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        try!(write!(f, "start time: {}, end time: {}\n{}",
                    &self.start_time, &self.end_time, &self.fullset));
        for inc in self.incset_list.iter() {
            try!(write!(f, "\n{}", inc));
        }
        Ok(())
    }
}

pub struct SignatureChain {
    pub fullsig : String,
    pub inclist : Vec<String>,
    pub start_time : String,
    pub end_time : String
}

impl SignatureChain {
    /// Create a new SignatureChain starting from a full signature.
    pub fn new(fname : &str, pr : &FileName) -> Self {
        SignatureChain {
            fullsig : fname.to_owned(),
            inclist : Vec::new(),
            start_time : pr.time.clone(),
            end_time : pr.time.clone()
        }
    }

    pub fn from_filename_info(fname_info : &FileNameInfo) -> Self {
        Self::new(fname_info.file_name, &fname_info.info)
    }

    /// Adds the given incremental signature to the signature chain if possible,
    /// returns false otherwise.
    pub fn add_new_sig(&mut self, fname : &FileNameInfo) -> bool {
        if fname.info.file_type != FileType::NewSig {
            false
        }
        else {
            self.inclist.push(fname.file_name.to_owned());
            self.end_time = fname.info.end_time.clone();
            true
        }
    }
}

pub type FileNameList<'a> = Vec<&'a str>;
pub type BackupChains = Vec<BackupChain>;
pub type SignatureChains = Vec<SignatureChain>;
type BackupSetList = Vec<BackupSet>;
type FileNameInfos<'a> = Vec<FileNameInfo<'a>>;

pub struct CollectionsStatus {
    pub backup_chains : BackupChains,
    pub sig_chains : SignatureChains
}

impl CollectionsStatus {
    pub fn new() -> Self {
        CollectionsStatus {
            backup_chains : Vec::new(),
            sig_chains : Vec::new()
        }
    }

    pub fn from_filename_list(filename_list : &FileNameList) -> Self {
        let mut result = Self::new();
        let filename_info_list = Self::compute_filename_infos(&filename_list);
        result.compute_backup_chains(&filename_info_list);
        result.compute_signature_chains(&filename_info_list);
        result
    }

    fn compute_filename_infos<'a>(filename_list : &'a FileNameList) -> FileNameInfos<'a> {
        let mut result = Vec::new();
        let parser = FileNameParser::new();
        for name in filename_list {
            if let Some(info) = parser.parse(&name) {
                result.push(FileNameInfo::new(&name, info));
            }
        }
        result
    }

    fn compute_backup_chains(&mut self, filename_list : &FileNameInfos) {
        let sets = Self::compute_backup_sets(filename_list);
        self.add_to_backup_chains(sets);
    }

    fn compute_backup_sets(filename_list : &FileNameInfos) -> BackupSetList {
        let mut sets = BackupSetList::new();
        for fileinfo in filename_list.iter() {
            let mut inserted = false;
            for set in sets.iter_mut() {
                if set.add_filename(&fileinfo) {
                    inserted = true;
                    break;
                }
            }
            if !inserted {
                let mut new_set = BackupSet::new();
                new_set.add_filename(&fileinfo);
                sets.push(new_set);
            }
        }
        // sort by time
        sets.sort_by(|a, b| a.get_time().cmp(b.get_time()));
        sets
    }

    fn add_to_backup_chains(&mut self, set_list : BackupSetList) {
        for set in set_list.into_iter() {
            match set.file_type {
                FileType::Full => {
                    let new_chain = BackupChain::new(set);
                    self.backup_chains.push(new_chain);
                }
                FileType::Inc => {
                    let mut rejected_set = Some(set);
                    for chain in self.backup_chains.iter_mut() {
                        rejected_set = chain.add_inc(rejected_set.unwrap());
                        if rejected_set.is_none() {
                            break;
                        }
                    }
                    if let Some(_) = rejected_set {
                        // TODO: add to orphaned sets
                    }
                }
                _ => { continue; }
            }
        }
        // sort by end time
        self.backup_chains.sort_by(|a, b| a.end_time.cmp(&b.end_time));
    }

    fn compute_signature_chains(&mut self, filename_list : &FileNameInfos) {
        // create a new signature chain for each fill signature
        self.sig_chains = filename_list.iter()
            .filter(|f| f.info.file_type == FileType::FullSig)
            .map(|f| SignatureChain::from_filename_info(f))
            .collect();
        // and collect all the new signatures, sorted by start time
        let mut new_sig = filename_list.iter()
            .filter(|f| f.info.file_type == FileType::NewSig)
            .collect::<Vec<_>>();
        new_sig.sort_by(|a, b| a.info.start_time.cmp(&b.info.start_time));

        // add the new signatures to signature chains
        for sig in new_sig.into_iter() {
            let mut added = false;
            for chain in self.sig_chains.iter_mut() {
                if chain.add_new_sig(&sig) {
                    added = true;
                    break;
                }
            }
            if !added {
                // TODO: add to orphaned filenames
            }
        }
    }
}

impl Display for CollectionsStatus {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        for backup_chain in &self.backup_chains {
            try!(backup_chain.fmt(f));
        }
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::{BackupSet, CollectionsStatus};
    use super::file_naming::{FileType, FileNameParser};

    #[test]
    fn parse_and_add() {
        let full1_name = "duplicity-full.20150617T182545Z.vol1.difftar.gz";
        let manifest1_name = "duplicity-full.20150617T182545Z.manifest";
        let inc1_name = "duplicity-inc.20150617T182629Z.to.20150617T182650Z.vol1.difftar.gz";

        let parser = FileNameParser::new();
        let full1 = parser.parse(full1_name).unwrap();
        let manifest1 = parser.parse(manifest1_name).unwrap();
        let inc1 = parser.parse(inc1_name).unwrap();

        let mut set = BackupSet::new();
        // add to set
        assert!(set.add_filename(full1_name, &full1));
        assert!(set.add_filename(manifest1_name, &manifest1));
        assert!(!set.add_filename(inc1_name, &inc1));
        // test results
        assert_eq!(set.file_type, FileType::Full);
        assert_eq!(set.time, "20150617t182545z");
        assert_eq!(set.start_time, "");
        assert_eq!(set.end_time, "");
        assert!(set.compressed);
        assert!(!set.encrypted);
        assert!(!set.partial);
        assert_eq!(set.manifest_path, manifest1_name);
    }

    #[test]
    fn collection_status() {
        let filename_list = vec![
            "duplicity-full.20150617T182545Z.manifest",
            "duplicity-full.20150617T182545Z.vol1.difftar.gz",
            "duplicity-full-signatures.20150617T182545Z.sigtar.gz",
            "duplicity-inc.20150617T182545Z.to.20150617T182629Z.manifest",
            "duplicity-inc.20150617T182545Z.to.20150617T182629Z.vol1.difftar.gz",
            "duplicity-inc.20150617T182629Z.to.20150617T182650Z.manifest",
            "duplicity-inc.20150617T182629Z.to.20150617T182650Z.vol1.difftar.gz",
            "duplicity-new-signatures.20150617T182545Z.to.20150617T182629Z.sigtar.gz",
            "duplicity-new-signatures.20150617T182629Z.to.20150617T182650Z.sigtar.gz"
        ];
        let collection_status = CollectionsStatus::from_filename_list(&filename_list);
        println!("collection: {}", collection_status);
    }
}
<commit_msg>Fixed tests<commit_after>// TODO: Make this non public
pub mod file_naming;

use std::collections::HashMap;
use std::fmt::{Display, Error, Formatter};
use self::file_naming::{FileName, FileNameInfo, FileType, FileNameParser};

pub struct BackupSet {
    pub file_type : FileType,
    pub time : String,
    pub start_time : String,
    pub end_time : String,
    pub compressed : bool,
    pub encrypted : bool,
    pub partial : bool,
    pub manifest_path : String,
    volumes_paths : HashMap<i32, String>,
    info_set : bool,    // true if informations are set
}

impl BackupSet {
    pub fn new() -> Self {
        BackupSet{
            file_type : FileType::Full,
            time : String::new(),
            start_time : String::new(),
            end_time : String::new(),
            compressed : false,
            encrypted : false,
            partial : false,
            manifest_path : String::new(),
            volumes_paths : HashMap::new(),
            info_set : false
        }
    }

    /// Add a filename to given set. Return true if it fits.
    ///
    /// The filename will match the given set if it has the right
    /// times and is of the right type. The information will be set
    /// from the first filename given.
    pub fn add_filename(&mut self, file_info : &FileNameInfo) -> bool {
        let pr = &file_info.info;
        let fname = file_info.file_name;

        if !self.info_set {
            self.set_info(&pr);
        }
        else {
            // check if the file is from the same backup set
            if self.file_type != pr.file_type ||
                self.time != pr.time ||
                self.start_time != pr.start_time ||
                self.end_time != pr.end_time {
                    return false;
            }
            else {
                // fix encrypted flag
                if self.encrypted != pr.encrypted &&
                    self.partial && pr.encrypted {
                        self.encrypted = pr.encrypted;
                    }
            }
        }
        // set manifest or volume number
        if pr.manifest {
            self.manifest_path = fname.to_owned();
        }
        else {
            self.volumes_paths.insert(pr.volume_number, fname.to_owned());
        }
        true
    }

    pub fn is_complete(&self) -> bool {
        !self.manifest_path.is_empty()
    }

    pub fn get_time(&self) -> &str {
        if self.time.is_empty() {
            self.end_time.as_ref()
        }
        else {
            self.time.as_ref()
        }
    }

    fn set_info(&mut self, fname : &FileName) {
        self.file_type = fname.file_type;
        self.time = fname.time.clone();
        self.start_time = fname.start_time.clone();
        self.end_time = fname.end_time.clone();
        self.compressed = fname.compressed;
        self.encrypted = fname.encrypted;
        self.partial = fname.partial;
        self.info_set = true;
    }
}

impl Display for BackupSet {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        match self.file_type {
            FileType::Full => {
                try!(write!(f, "Full, time: {}", self.time));
            },
            FileType::Inc => {
                try!(write!(f, "Incremental, start time: {}, end time: {}", self.start_time, self.end_time));
            },
            _ => { }
        }
        if self.compressed {
            try!(write!(f, ", compressed"));
        }
        if self.encrypted {
            try!(write!(f, ", encrypted"));
        }
        if self.partial {
            try!(write!(f, ", partial"));
        }
        if !self.manifest_path.is_empty() {
            try!(write!(f, "\n manifest: {}", self.manifest_path));
        }
        if !self.volumes_paths.is_empty() {
            try!(write!(f, "\n volumes:\n"));
            for (_, vol) in self.volumes_paths.iter() {
                try!(write!(f, " {}", vol));
            }
        }
        Ok(())
    }
}

pub struct BackupChain {
    pub fullset : BackupSet,
    pub incset_list : Vec<BackupSet>,
    pub start_time : String,
    pub end_time : String
}

impl BackupChain {
    /// Create a new BackupChain starting from a full backup set.
    pub fn new(fullset : BackupSet) -> Self {
        assert_eq!(fullset.file_type, FileType::Full);
        let time = fullset.time.clone();

        BackupChain{
            fullset : fullset,
            incset_list : Vec::new(),
            start_time : time.clone(),
            end_time : time
        }
    }

    /// Adds the given incremental backup element to the backup chain if possible,
    /// returns it back otherwise.
    pub fn add_inc(&mut self, incset : BackupSet) -> Option<BackupSet> {
        if self.end_time == incset.start_time {
            self.end_time = incset.end_time.clone();
            self.incset_list.push(incset);
            None
        }
        else {
            // replace the last element if the end time comes before
            let replace_last = self.incset_list.last().map_or(false,
                |last| incset.start_time == last.start_time && incset.end_time > last.end_time);
            if replace_last {
                self.end_time = incset.end_time.clone();
                self.incset_list.pop();
                self.incset_list.push(incset);
                None
            }
            else {
                // ignore the given incremental backup set
                Some(incset)
            }
        }
    }
}

impl Display for BackupChain {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        try!(write!(f, "start time: {}, end time: {}\n{}",
                    &self.start_time, &self.end_time, &self.fullset));
        for inc in self.incset_list.iter() {
            try!(write!(f, "\n{}", inc));
        }
        Ok(())
    }
}

pub struct SignatureChain {
    pub fullsig : String,
    pub inclist : Vec<String>,
    pub start_time : String,
    pub end_time : String
}

impl SignatureChain {
    /// Create a new SignatureChain starting from a full signature.
    pub fn new(fname : &str, pr : &FileName) -> Self {
        SignatureChain {
            fullsig : fname.to_owned(),
            inclist : Vec::new(),
            start_time : pr.time.clone(),
            end_time : pr.time.clone()
        }
    }

    pub fn from_filename_info(fname_info : &FileNameInfo) -> Self {
        Self::new(fname_info.file_name, &fname_info.info)
    }

    /// Adds the given incremental signature to the signature chain if possible,
    /// returns false otherwise.
    pub fn add_new_sig(&mut self, fname : &FileNameInfo) -> bool {
        if fname.info.file_type != FileType::NewSig {
            false
        }
        else {
            self.inclist.push(fname.file_name.to_owned());
            self.end_time = fname.info.end_time.clone();
            true
        }
    }
}

pub type FileNameList<'a> = Vec<&'a str>;
pub type BackupChains = Vec<BackupChain>;
pub type SignatureChains = Vec<SignatureChain>;
type BackupSetList = Vec<BackupSet>;
type FileNameInfos<'a> = Vec<FileNameInfo<'a>>;

pub struct CollectionsStatus {
    pub backup_chains : BackupChains,
    pub sig_chains : SignatureChains
}

impl CollectionsStatus {
    pub fn new() -> Self {
        CollectionsStatus {
            backup_chains : Vec::new(),
            sig_chains : Vec::new()
        }
    }

    pub fn from_filename_list(filename_list : &FileNameList) -> Self {
        let mut result = Self::new();
        let filename_info_list = Self::compute_filename_infos(&filename_list);
        result.compute_backup_chains(&filename_info_list);
        result.compute_signature_chains(&filename_info_list);
        result
    }

    fn compute_filename_infos<'a>(filename_list : &'a FileNameList) -> FileNameInfos<'a> {
        let mut result = Vec::new();
        let parser = FileNameParser::new();
        for name in filename_list {
            if let Some(info) = parser.parse(&name) {
                result.push(FileNameInfo::new(&name, info));
            }
        }
        result
    }

    fn compute_backup_chains(&mut self, filename_list : &FileNameInfos) {
        let sets = Self::compute_backup_sets(filename_list);
        self.add_to_backup_chains(sets);
    }

    fn compute_backup_sets(filename_list : &FileNameInfos) -> BackupSetList {
        let mut sets = BackupSetList::new();
        for fileinfo in filename_list.iter() {
            let mut inserted = false;
            for set in sets.iter_mut() {
                if set.add_filename(&fileinfo) {
                    inserted = true;
                    break;
                }
            }
            if !inserted {
                let mut new_set = BackupSet::new();
                new_set.add_filename(&fileinfo);
                sets.push(new_set);
            }
        }
        // sort by time
        sets.sort_by(|a, b| a.get_time().cmp(b.get_time()));
        sets
    }

    fn add_to_backup_chains(&mut self, set_list : BackupSetList) {
        for set in set_list.into_iter() {
            match set.file_type {
                FileType::Full => {
                    let new_chain = BackupChain::new(set);
                    self.backup_chains.push(new_chain);
                }
                FileType::Inc => {
                    let mut rejected_set = Some(set);
                    for chain in self.backup_chains.iter_mut() {
                        rejected_set = chain.add_inc(rejected_set.unwrap());
                        if rejected_set.is_none() {
                            break;
                        }
                    }
                    if let Some(_) = rejected_set {
                        // TODO: add to orphaned sets
                    }
                }
                _ => { continue; }
            }
        }
        // sort by end time
        self.backup_chains.sort_by(|a, b| a.end_time.cmp(&b.end_time));
    }

    fn compute_signature_chains(&mut self, filename_list : &FileNameInfos) {
        // create a new signature chain for each fill signature
        self.sig_chains = filename_list.iter()
            .filter(|f| f.info.file_type == FileType::FullSig)
            .map(|f| SignatureChain::from_filename_info(f))
            .collect();
        // and collect all the new signatures, sorted by start time
        let mut new_sig = filename_list.iter()
            .filter(|f| f.info.file_type == FileType::NewSig)
            .collect::<Vec<_>>();
        new_sig.sort_by(|a, b| a.info.start_time.cmp(&b.info.start_time));

        // add the new signatures to signature chains
        for sig in new_sig.into_iter() {
            let mut added = false;
            for chain in self.sig_chains.iter_mut() {
                if chain.add_new_sig(&sig) {
                    added = true;
                    break;
                }
            }
            if !added {
                // TODO: add to orphaned filenames
            }
        }
    }
}

impl Display for CollectionsStatus {
    fn fmt(&self, f : &mut Formatter) -> Result<(), Error> {
        for backup_chain in &self.backup_chains {
            try!(backup_chain.fmt(f));
        }
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::{BackupSet, CollectionsStatus};
    use super::file_naming::{FileType, FileNameInfo, FileNameParser};

    #[test]
    fn parse_and_add() {
        let full1_name = "duplicity-full.20150617T182545Z.vol1.difftar.gz";
        let manifest1_name = "duplicity-full.20150617T182545Z.manifest";
        let inc1_name = "duplicity-inc.20150617T182629Z.to.20150617T182650Z.vol1.difftar.gz";

        let parser = FileNameParser::new();
        let full1 = FileNameInfo::new(full1_name, parser.parse(full1_name).unwrap());
        let manifest1 = FileNameInfo::new(manifest1_name, parser.parse(manifest1_name).unwrap());
        let inc1 = FileNameInfo::new(inc1_name, parser.parse(inc1_name).unwrap());

        let mut set = BackupSet::new();
        // add to set
        assert!(set.add_filename(&full1));
        assert!(set.add_filename(&manifest1));
        assert!(!set.add_filename(&inc1));
        // test results
        assert_eq!(set.file_type, FileType::Full);
        assert_eq!(set.time, "20150617t182545z");
        assert_eq!(set.start_time, "");
        assert_eq!(set.end_time, "");
        assert!(set.compressed);
        assert!(!set.encrypted);
        assert!(!set.partial);
        assert_eq!(set.manifest_path, manifest1_name);
    }

    #[test]
    fn collection_status() {
        let filename_list = vec![
            "duplicity-full.20150617T182545Z.manifest",
            "duplicity-full.20150617T182545Z.vol1.difftar.gz",
            "duplicity-full-signatures.20150617T182545Z.sigtar.gz",
            "duplicity-inc.20150617T182545Z.to.20150617T182629Z.manifest",
            "duplicity-inc.20150617T182545Z.to.20150617T182629Z.vol1.difftar.gz",
            "duplicity-inc.20150617T182629Z.to.20150617T182650Z.manifest",
            "duplicity-inc.20150617T182629Z.to.20150617T182650Z.vol1.difftar.gz",
            "duplicity-new-signatures.20150617T182545Z.to.20150617T182629Z.sigtar.gz",
            "duplicity-new-signatures.20150617T182629Z.to.20150617T182650Z.sigtar.gz"
        ];
        let collection_status = CollectionsStatus::from_filename_list(&filename_list);
        println!("collection: {}", collection_status);
    }
}
<|endoftext|>


<commit_before>fn f(x: f32) -> f32 {
    return x * x - x - 1.0;
}

fn midpoint_rule(a: f32, b: f32, n: i32) -> f32 {
    let h = (b - a) / n as f32;
    let mut r = 0.0;
    let mut x = a + h / 2.0;

    for _ in 0..n {
        r += f(x);
        x += h;
    }
    return r * h;
}

fn trapezoid_rule(a: f32, b: f32, n: i32) -> f32 {
    let h = (b - a) / n as f32;
    let mut r = 0.0;
    let mut x = a + h;
    let mut y1 = f(a);

    for _ in 0..n {
        let y2 = f(x);
        r += y1 + y2;
        x += h;
        y1 = y2;
    }
    return r * h / 2.0;
}

fn main() {
    println!("Midpoint Rule: {}", midpoint_rule(-2.0, 4.0, 1000));
    println!("Trapezoid Rule: {}", trapezoid_rule(-2.0, 4.0, 1000));
}
<commit_msg>Apply Tuple<commit_after>struct Interval(f32, f32);

fn f(x: f32) -> f32 {
    return x * x - x - 1.0;
}

fn midpoint_rule(i: Interval, n: i32) -> f32 {
    let h = (i.1 - i.0) / n as f32;
    let mut r = 0.0;
    let mut x = i.0 + h / 2.0;

    for _ in 0..n {
        r += f(x);
        x += h;
    }
    return r * h;
}

fn trapezoid_rule(i: Interval, n: i32) -> f32 {
    let h = (i.1 - i.0) / n as f32;
    let mut r = 0.0;
    let mut x = i.0 + h;
    let mut y1 = f(i.0);

    for _ in 0..n {
        let y2 = f(x);
        r += y1 + y2;
        x += h;
        y1 = y2;
    }
    return r * h / 2.0;
}

fn main() {
    println!("Midpoint Rule: {}", midpoint_rule(Interval(-2.0, 4.0), 1000));
    println!("Trapezoid Rule: {}", trapezoid_rule(Interval(-2.0, 4.0), 1000));
}
<|endoftext|>