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

#![forbid(unsafe_code)]
use std::{fmt, time::Duration};

use async_trait::async_trait;
use diem_types::{account_address::AccountAddress, transaction::Transaction};
use futures::channel::{mpsc, oneshot};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tokio::time::timeout;

const MEMPOOL_NOTIFICATION_CHANNEL_SIZE: usize = 1;

#[derive(Clone, Debug, Deserialize, Error, PartialEq, Serialize)]
pub enum Error {
    #[error("Commit notification failed: {0}")]
    CommitNotificationError(String),
    #[error("Hit the timeout waiting for mempool to respond to the notification!")]
    TimeoutWaitingForMempool,
    #[error("Unexpected error encountered: {0}")]
    UnexpectedErrorEncountered(String),
}

/// The interface between state sync and mempool, allowing state sync to notify
/// mempool of events (e.g., newly committed transactions).
#[async_trait]
pub trait MempoolNotificationSender: Send {
    /// Notify mempool of the newly committed transactions at the specified block timestamp.
    async fn notify_new_commit(
        &mut self,
        committed_transactions: Vec<Transaction>,
        block_timestamp_usecs: u64,
        notification_timeout_ms: u64,
    ) -> Result<(), Error>;
}

/// The state sync component responsible for notifying mempool.
///
/// Note: When a MempoolNotifier instance is created, mempool must take and
/// listen to the receiver in the corresponding MempoolNotificationListener.
#[derive(Debug)]
pub struct MempoolNotifier {
    notification_sender: mpsc::Sender<MempoolCommitNotification>,
}

impl MempoolNotifier {
    /// Returns a new MempoolNotifier and MempoolNotificationListener (to be
    /// used in conjuction with one another).
    pub fn new() -> (Self, MempoolNotificationListener) {
        let (notification_sender, notification_receiver) =
            mpsc::channel(MEMPOOL_NOTIFICATION_CHANNEL_SIZE);

        let mempool_notifier = MempoolNotifier {
            notification_sender,
        };
        let mempool_listener = MempoolNotificationListener::new(notification_receiver);
        (mempool_notifier, mempool_listener)
    }
}

#[async_trait]
impl MempoolNotificationSender for MempoolNotifier {
    async fn notify_new_commit(
        &mut self,
        transactions: Vec<Transaction>,
        block_timestamp_usecs: u64,
        notification_timeout_ms: u64,
    ) -> Result<(), Error> {
        // Get only user transactions from committed transactions
        let user_transactions: Vec<CommittedTransaction> = transactions
            .iter()
            .filter_map(|transaction| match transaction {
                Transaction::UserTransaction(signed_txn) => Some(CommittedTransaction {
                    sender: signed_txn.sender(),
                    sequence_number: signed_txn.sequence_number(),
                }),
                _ => None,
            })
            .collect();

        // Only send a notification if user transactions have been committed
        if user_transactions.is_empty() {
            return Ok(());
        }

        // Construct a oneshot channel to receive a mempool response
        let (callback, callback_receiver) = oneshot::channel();
        let commit_notification = MempoolCommitNotification {
            transactions: user_transactions,
            block_timestamp_usecs,
            callback,
        };

        // Send the notification to mempool
        if let Err(error) = self.notification_sender.try_send(commit_notification) {
            return Err(Error::CommitNotificationError(format!(
                "Failed to notify mempool of committed transactions! Error: {:?}",
                error
            )));
        }

        // Handle any responses or a timeout
        if let Ok(response) = timeout(
            Duration::from_millis(notification_timeout_ms),
            callback_receiver,
        )
        .await
        {
            match response {
                Ok(Ok(MempoolNotificationResponse::Success)) => Ok(()),
                Ok(Err(error)) => Err(Error::UnexpectedErrorEncountered(format!(
                    "Unexpected response from mempool! Error: {:?}",
                    error
                ))),
                Err(error) => Err(Error::UnexpectedErrorEncountered(format!("{:?}", error))),
            }
        } else {
            Err(Error::TimeoutWaitingForMempool)
        }
    }
}

/// A notification for newly committed transactions sent by state sync to mempool.
#[derive(Debug)]
pub struct MempoolCommitNotification {
    pub transactions: Vec<CommittedTransaction>,
    pub block_timestamp_usecs: u64, // The timestamp of the committed block.
    pub(crate) callback: oneshot::Sender<Result<MempoolNotificationResponse, Error>>,
}

impl fmt::Display for MempoolCommitNotification {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "MempoolCommitNotification [block_timestamp_usecs: {}, txns: {:?}]",
            self.block_timestamp_usecs, self.transactions
        )
    }
}

/// A successfully executed and committed user transaction.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CommittedTransaction {
    pub sender: AccountAddress,
    pub sequence_number: u64,
}

impl fmt::Display for CommittedTransaction {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}:{}", self.sender, self.sequence_number,)
    }
}

/// The mempool component responsible for responding to state sync notifications.
#[derive(Debug)]
pub struct MempoolNotificationListener {
    pub notification_receiver: mpsc::Receiver<MempoolCommitNotification>,
}

impl MempoolNotificationListener {
    pub fn new(notification_receiver: mpsc::Receiver<MempoolCommitNotification>) -> Self {
        MempoolNotificationListener {
            notification_receiver,
        }
    }

    /// Respond (succesfully) to the commit notification previously sent by state sync.
    pub async fn ack_commit_notification(
        &mut self,
        mempool_commit_notification: MempoolCommitNotification,
    ) -> Result<(), Error> {
        mempool_commit_notification
            .callback
            .send(Ok(MempoolNotificationResponse::Success))
            .map_err(|error| Error::UnexpectedErrorEncountered(format!("{:?}", error)))
    }
}

/// A response from mempool for a notification.
///
/// Note: failure responses are not currently used.
#[derive(Debug)]
enum MempoolNotificationResponse {
    Success,
}

#[cfg(test)]
mod tests {
    use crate::{CommittedTransaction, Error, MempoolNotificationSender, MempoolNotifier};
    use diem_crypto::{ed25519::Ed25519PrivateKey, HashValue, PrivateKey, SigningKey, Uniform};
    use diem_types::{
        account_address::AccountAddress,
        block_metadata::BlockMetadata,
        chain_id::ChainId,
        transaction::{
            ChangeSet, RawTransaction, Script, SignedTransaction, Transaction, TransactionPayload,
            WriteSetPayload,
        },
        write_set::WriteSetMut,
    };
    use futures::executor::block_on;
    use tokio::runtime::{Builder, Runtime};

    #[test]
    fn test_mempool_not_listening() {
        // Create runtime and mempool notifier
        let runtime = create_runtime();
        let _enter = runtime.enter();
        let (mut mempool_notifier, mut mempool_listener) = MempoolNotifier::new();

        // Send a notification and expect a timeout (no listener)
        let notify_result =
            block_on(mempool_notifier.notify_new_commit(vec![create_user_transaction()], 0, 1000));
        assert!(matches!(
            notify_result,
            Err(Error::TimeoutWaitingForMempool)
        ));

        // Drop the receiver and try again
        mempool_listener.notification_receiver.close();
        let notify_result =
            block_on(mempool_notifier.notify_new_commit(vec![create_user_transaction()], 0, 1000));
        assert!(matches!(
            notify_result,
            Err(Error::CommitNotificationError(_))
        ));
    }

    #[test]
    fn test_zero_timeout() {
        // Create runtime and mempool notifier
        let runtime = create_runtime();
        let _enter = runtime.enter();
        let (mut mempool_notifier, mut _mempool_listener) = MempoolNotifier::new();

        // Send a notification and expect a timeout (zero timeout)
        let notify_result =
            block_on(mempool_notifier.notify_new_commit(vec![create_user_transaction()], 0, 0));
        assert!(matches!(
            notify_result,
            Err(Error::TimeoutWaitingForMempool)
        ));
    }

    #[test]
    fn test_no_transactions() {
        // Create runtime and mempool notifier
        let runtime = create_runtime();
        let _enter = runtime.enter();
        let (mut mempool_notifier, mut _mempool_listener) = MempoolNotifier::new();

        // Send a notification and verify no timeout because no notification was sent!
        let notify_result = block_on(mempool_notifier.notify_new_commit(vec![], 0, 1000));
        assert!(notify_result.is_ok());
    }

    #[test]
    fn test_transaction_filtering() {
        // Create runtime and mempool notifier
        let runtime = create_runtime();
        let _enter = runtime.enter();
        let (mut mempool_notifier, mut _mempool_listener) = MempoolNotifier::new();

        // Create several transactions that should be filtered out
        let mut transactions = vec![];
        for _ in 0..5 {
            transactions.push(create_block_metadata_transaction())
        }
        for _ in 0..5 {
            transactions.push(create_genesis_transaction())
        }

        // Send a notification and verify no timeout because no notification was sent!
        let notify_result =
            block_on(mempool_notifier.notify_new_commit(transactions.clone(), 0, 1000));
        assert!(notify_result.is_ok());

        // Send another notification with a single user transaction now included.
        transactions.push(create_user_transaction());
        let notify_result = block_on(mempool_notifier.notify_new_commit(transactions, 0, 1000));
        assert!(matches!(
            notify_result,
            Err(Error::TimeoutWaitingForMempool)
        ));
    }

    #[test]
    fn test_commit_notification_arrives() {
        // Create runtime and mempool notifier
        let runtime = create_runtime();
        let _enter = runtime.enter();
        let (mut mempool_notifier, mut mempool_listener) = MempoolNotifier::new();

        // Send a notification
        let user_transaction = create_user_transaction();
        let transactions = vec![user_transaction.clone()];
        let block_timestamp_usecs = 101;
        let _ =
            block_on(mempool_notifier.notify_new_commit(transactions, block_timestamp_usecs, 1000));

        // Verify the notification arrives at the receiver
        match mempool_listener.notification_receiver.try_next() {
            Ok(Some(mempool_commit_notification)) => match user_transaction {
                Transaction::UserTransaction(signed_transaction) => {
                    assert_eq!(
                        mempool_commit_notification.transactions,
                        vec![CommittedTransaction {
                            sender: signed_transaction.sender(),
                            sequence_number: signed_transaction.sequence_number(),
                        }]
                    );
                    assert_eq!(
                        mempool_commit_notification.block_timestamp_usecs,
                        block_timestamp_usecs
                    );
                }
                result => panic!("Expected user transaction but got: {:?}", result),
            },
            result => panic!("Expected mempool commit notification but got: {:?}", result),
        };
    }

    #[test]
    fn test_mempool_success_response() {
        // Create runtime and mempool notifier
        let runtime = create_runtime();
        let _enter = runtime.enter();
        let (mut mempool_notifier, mut mempool_listener) = MempoolNotifier::new();

        // Spawn a new thread to handle any messages on the receiver
        let _handler = std::thread::spawn(move || loop {
            if let Ok(Some(mempool_commit_notification)) =
                mempool_listener.notification_receiver.try_next()
            {
                let _result =
                    block_on(mempool_listener.ack_commit_notification(mempool_commit_notification));
            }
        });

        // Send a notification and verify a successful response
        let notify_result = block_on(mempool_notifier.notify_new_commit(
            vec![create_user_transaction()],
            101,
            1000,
        ));
        assert!(notify_result.is_ok());
    }

    fn create_user_transaction() -> Transaction {
        let private_key = Ed25519PrivateKey::generate_for_testing();
        let public_key = private_key.public_key();

        let transaction_payload = TransactionPayload::Script(Script::new(vec![], vec![], vec![]));
        let raw_transaction = RawTransaction::new(
            AccountAddress::random(),
            0,
            transaction_payload,
            0,
            0,
            "".into(),
            0,
            ChainId::new(10),
        );
        let signed_transaction = SignedTransaction::new(
            raw_transaction.clone(),
            public_key,
            private_key.sign(&raw_transaction),
        );

        Transaction::UserTransaction(signed_transaction)
    }

    fn create_block_metadata_transaction() -> Transaction {
        Transaction::BlockMetadata(BlockMetadata::new(
            HashValue::new([0; HashValue::LENGTH]),
            1,
            300000001,
            vec![],
            AccountAddress::random(),
        ))
    }

    fn create_genesis_transaction() -> Transaction {
        Transaction::GenesisTransaction(WriteSetPayload::Direct(ChangeSet::new(
            WriteSetMut::new(vec![])
                .freeze()
                .expect("freeze cannot fail"),
            vec![],
        )))
    }

    fn create_runtime() -> Runtime {
        Builder::new_multi_thread().enable_all().build().unwrap()
    }
}
<commit_msg>[State Sync] Remove uneccessary result from mempool notifier.<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]
use std::{fmt, time::Duration};

use async_trait::async_trait;
use diem_types::{account_address::AccountAddress, transaction::Transaction};
use futures::channel::{mpsc, oneshot};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tokio::time::timeout;

const MEMPOOL_NOTIFICATION_CHANNEL_SIZE: usize = 1;

#[derive(Clone, Debug, Deserialize, Error, PartialEq, Serialize)]
pub enum Error {
    #[error("Commit notification failed: {0}")]
    CommitNotificationError(String),
    #[error("Hit the timeout waiting for mempool to respond to the notification!")]
    TimeoutWaitingForMempool,
    #[error("Unexpected error encountered: {0}")]
    UnexpectedErrorEncountered(String),
}

/// The interface between state sync and mempool, allowing state sync to notify
/// mempool of events (e.g., newly committed transactions).
#[async_trait]
pub trait MempoolNotificationSender: Send {
    /// Notify mempool of the newly committed transactions at the specified block timestamp.
    async fn notify_new_commit(
        &mut self,
        committed_transactions: Vec<Transaction>,
        block_timestamp_usecs: u64,
        notification_timeout_ms: u64,
    ) -> Result<(), Error>;
}

/// The state sync component responsible for notifying mempool.
///
/// Note: When a MempoolNotifier instance is created, mempool must take and
/// listen to the receiver in the corresponding MempoolNotificationListener.
#[derive(Debug)]
pub struct MempoolNotifier {
    notification_sender: mpsc::Sender<MempoolCommitNotification>,
}

impl MempoolNotifier {
    /// Returns a new MempoolNotifier and MempoolNotificationListener (to be
    /// used in conjuction with one another).
    pub fn new() -> (Self, MempoolNotificationListener) {
        let (notification_sender, notification_receiver) =
            mpsc::channel(MEMPOOL_NOTIFICATION_CHANNEL_SIZE);

        let mempool_notifier = MempoolNotifier {
            notification_sender,
        };
        let mempool_listener = MempoolNotificationListener::new(notification_receiver);
        (mempool_notifier, mempool_listener)
    }
}

#[async_trait]
impl MempoolNotificationSender for MempoolNotifier {
    async fn notify_new_commit(
        &mut self,
        transactions: Vec<Transaction>,
        block_timestamp_usecs: u64,
        notification_timeout_ms: u64,
    ) -> Result<(), Error> {
        // Get only user transactions from committed transactions
        let user_transactions: Vec<CommittedTransaction> = transactions
            .iter()
            .filter_map(|transaction| match transaction {
                Transaction::UserTransaction(signed_txn) => Some(CommittedTransaction {
                    sender: signed_txn.sender(),
                    sequence_number: signed_txn.sequence_number(),
                }),
                _ => None,
            })
            .collect();

        // Only send a notification if user transactions have been committed
        if user_transactions.is_empty() {
            return Ok(());
        }

        // Construct a oneshot channel to receive a mempool response
        let (callback, callback_receiver) = oneshot::channel();
        let commit_notification = MempoolCommitNotification {
            transactions: user_transactions,
            block_timestamp_usecs,
            callback,
        };

        // Send the notification to mempool
        if let Err(error) = self.notification_sender.try_send(commit_notification) {
            return Err(Error::CommitNotificationError(format!(
                "Failed to notify mempool of committed transactions! Error: {:?}",
                error
            )));
        }

        // Handle any responses or a timeout
        if let Ok(response) = timeout(
            Duration::from_millis(notification_timeout_ms),
            callback_receiver,
        )
        .await
        {
            match response {
                Ok(MempoolNotificationResponse::Success) => Ok(()),
                Err(error) => Err(Error::UnexpectedErrorEncountered(format!("{:?}", error))),
            }
        } else {
            Err(Error::TimeoutWaitingForMempool)
        }
    }
}

/// A notification for newly committed transactions sent by state sync to mempool.
#[derive(Debug)]
pub struct MempoolCommitNotification {
    pub transactions: Vec<CommittedTransaction>,
    pub block_timestamp_usecs: u64, // The timestamp of the committed block.
    pub(crate) callback: oneshot::Sender<MempoolNotificationResponse>,
}

impl fmt::Display for MempoolCommitNotification {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "MempoolCommitNotification [block_timestamp_usecs: {}, txns: {:?}]",
            self.block_timestamp_usecs, self.transactions
        )
    }
}

/// A successfully executed and committed user transaction.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CommittedTransaction {
    pub sender: AccountAddress,
    pub sequence_number: u64,
}

impl fmt::Display for CommittedTransaction {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}:{}", self.sender, self.sequence_number,)
    }
}

/// The mempool component responsible for responding to state sync notifications.
#[derive(Debug)]
pub struct MempoolNotificationListener {
    pub notification_receiver: mpsc::Receiver<MempoolCommitNotification>,
}

impl MempoolNotificationListener {
    pub fn new(notification_receiver: mpsc::Receiver<MempoolCommitNotification>) -> Self {
        MempoolNotificationListener {
            notification_receiver,
        }
    }

    /// Respond (succesfully) to the commit notification previously sent by state sync.
    pub async fn ack_commit_notification(
        &mut self,
        mempool_commit_notification: MempoolCommitNotification,
    ) -> Result<(), Error> {
        mempool_commit_notification
            .callback
            .send(MempoolNotificationResponse::Success)
            .map_err(|error| Error::UnexpectedErrorEncountered(format!("{:?}", error)))
    }
}

/// A response from mempool for a notification.
///
/// Note: failure responses are not currently used.
#[derive(Debug)]
enum MempoolNotificationResponse {
    Success,
}

#[cfg(test)]
mod tests {
    use crate::{CommittedTransaction, Error, MempoolNotificationSender, MempoolNotifier};
    use diem_crypto::{ed25519::Ed25519PrivateKey, HashValue, PrivateKey, SigningKey, Uniform};
    use diem_types::{
        account_address::AccountAddress,
        block_metadata::BlockMetadata,
        chain_id::ChainId,
        transaction::{
            ChangeSet, RawTransaction, Script, SignedTransaction, Transaction, TransactionPayload,
            WriteSetPayload,
        },
        write_set::WriteSetMut,
    };
    use futures::executor::block_on;
    use tokio::runtime::{Builder, Runtime};

    #[test]
    fn test_mempool_not_listening() {
        // Create runtime and mempool notifier
        let runtime = create_runtime();
        let _enter = runtime.enter();
        let (mut mempool_notifier, mut mempool_listener) = MempoolNotifier::new();

        // Send a notification and expect a timeout (no listener)
        let notify_result =
            block_on(mempool_notifier.notify_new_commit(vec![create_user_transaction()], 0, 1000));
        assert!(matches!(
            notify_result,
            Err(Error::TimeoutWaitingForMempool)
        ));

        // Drop the receiver and try again
        mempool_listener.notification_receiver.close();
        let notify_result =
            block_on(mempool_notifier.notify_new_commit(vec![create_user_transaction()], 0, 1000));
        assert!(matches!(
            notify_result,
            Err(Error::CommitNotificationError(_))
        ));
    }

    #[test]
    fn test_zero_timeout() {
        // Create runtime and mempool notifier
        let runtime = create_runtime();
        let _enter = runtime.enter();
        let (mut mempool_notifier, mut _mempool_listener) = MempoolNotifier::new();

        // Send a notification and expect a timeout (zero timeout)
        let notify_result =
            block_on(mempool_notifier.notify_new_commit(vec![create_user_transaction()], 0, 0));
        assert!(matches!(
            notify_result,
            Err(Error::TimeoutWaitingForMempool)
        ));
    }

    #[test]
    fn test_no_transactions() {
        // Create runtime and mempool notifier
        let runtime = create_runtime();
        let _enter = runtime.enter();
        let (mut mempool_notifier, mut _mempool_listener) = MempoolNotifier::new();

        // Send a notification and verify no timeout because no notification was sent!
        let notify_result = block_on(mempool_notifier.notify_new_commit(vec![], 0, 1000));
        assert!(notify_result.is_ok());
    }

    #[test]
    fn test_transaction_filtering() {
        // Create runtime and mempool notifier
        let runtime = create_runtime();
        let _enter = runtime.enter();
        let (mut mempool_notifier, mut _mempool_listener) = MempoolNotifier::new();

        // Create several transactions that should be filtered out
        let mut transactions = vec![];
        for _ in 0..5 {
            transactions.push(create_block_metadata_transaction())
        }
        for _ in 0..5 {
            transactions.push(create_genesis_transaction())
        }

        // Send a notification and verify no timeout because no notification was sent!
        let notify_result =
            block_on(mempool_notifier.notify_new_commit(transactions.clone(), 0, 1000));
        assert!(notify_result.is_ok());

        // Send another notification with a single user transaction now included.
        transactions.push(create_user_transaction());
        let notify_result = block_on(mempool_notifier.notify_new_commit(transactions, 0, 1000));
        assert!(matches!(
            notify_result,
            Err(Error::TimeoutWaitingForMempool)
        ));
    }

    #[test]
    fn test_commit_notification_arrives() {
        // Create runtime and mempool notifier
        let runtime = create_runtime();
        let _enter = runtime.enter();
        let (mut mempool_notifier, mut mempool_listener) = MempoolNotifier::new();

        // Send a notification
        let user_transaction = create_user_transaction();
        let transactions = vec![user_transaction.clone()];
        let block_timestamp_usecs = 101;
        let _ =
            block_on(mempool_notifier.notify_new_commit(transactions, block_timestamp_usecs, 1000));

        // Verify the notification arrives at the receiver
        match mempool_listener.notification_receiver.try_next() {
            Ok(Some(mempool_commit_notification)) => match user_transaction {
                Transaction::UserTransaction(signed_transaction) => {
                    assert_eq!(
                        mempool_commit_notification.transactions,
                        vec![CommittedTransaction {
                            sender: signed_transaction.sender(),
                            sequence_number: signed_transaction.sequence_number(),
                        }]
                    );
                    assert_eq!(
                        mempool_commit_notification.block_timestamp_usecs,
                        block_timestamp_usecs
                    );
                }
                result => panic!("Expected user transaction but got: {:?}", result),
            },
            result => panic!("Expected mempool commit notification but got: {:?}", result),
        };
    }

    #[test]
    fn test_mempool_success_response() {
        // Create runtime and mempool notifier
        let runtime = create_runtime();
        let _enter = runtime.enter();
        let (mut mempool_notifier, mut mempool_listener) = MempoolNotifier::new();

        // Spawn a new thread to handle any messages on the receiver
        let _handler = std::thread::spawn(move || loop {
            if let Ok(Some(mempool_commit_notification)) =
                mempool_listener.notification_receiver.try_next()
            {
                let _result =
                    block_on(mempool_listener.ack_commit_notification(mempool_commit_notification));
            }
        });

        // Send a notification and verify a successful response
        let notify_result = block_on(mempool_notifier.notify_new_commit(
            vec![create_user_transaction()],
            101,
            1000,
        ));
        assert!(notify_result.is_ok());
    }

    fn create_user_transaction() -> Transaction {
        let private_key = Ed25519PrivateKey::generate_for_testing();
        let public_key = private_key.public_key();

        let transaction_payload = TransactionPayload::Script(Script::new(vec![], vec![], vec![]));
        let raw_transaction = RawTransaction::new(
            AccountAddress::random(),
            0,
            transaction_payload,
            0,
            0,
            "".into(),
            0,
            ChainId::new(10),
        );
        let signed_transaction = SignedTransaction::new(
            raw_transaction.clone(),
            public_key,
            private_key.sign(&raw_transaction),
        );

        Transaction::UserTransaction(signed_transaction)
    }

    fn create_block_metadata_transaction() -> Transaction {
        Transaction::BlockMetadata(BlockMetadata::new(
            HashValue::new([0; HashValue::LENGTH]),
            1,
            300000001,
            vec![],
            AccountAddress::random(),
        ))
    }

    fn create_genesis_transaction() -> Transaction {
        Transaction::GenesisTransaction(WriteSetPayload::Direct(ChangeSet::new(
            WriteSetMut::new(vec![])
                .freeze()
                .expect("freeze cannot fail"),
            vec![],
        )))
    }

    fn create_runtime() -> Runtime {
        Builder::new_multi_thread().enable_all().build().unwrap()
    }
}
<|endoftext|>


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

use aptos_crypto::HashValue;
use aptos_logger::info;
use aptos_types::{
    state_store::{state_key::StateKey, state_value::StateValue},
    transaction::Version,
};
use executor_types::StateSnapshotDelta;
use scratchpad::SparseMerkleTree;
use std::{
    collections::VecDeque,
    sync::{mpsc, Arc},
};
use storage_interface::DbWriter;

const NUM_COMMITTED_SMTS_TO_CACHE: usize = 5;
const NUM_COMMITS_TO_BATCH: usize = 20;

pub struct StateCommitter {
    commit_receiver: mpsc::Receiver<StateSnapshotDelta>,
    db: Arc<dyn DbWriter>,

    // keep some recently committed SMTs in mem as a naive cache
    cache_queue: VecDeque<SparseMerkleTree<StateValue>>,
    version: Version,
    smt: SparseMerkleTree<StateValue>,
    committed_smt: SparseMerkleTree<StateValue>,
    updates: Vec<(HashValue, (HashValue, StateKey))>,
    num_pending_commits: usize,
}

impl StateCommitter {
    pub fn new(
        commit_receiver: mpsc::Receiver<StateSnapshotDelta>,
        db: Arc<dyn DbWriter>,
        committed_smt: SparseMerkleTree<StateValue>,
    ) -> Self {
        let mut cache_queue = VecDeque::new();
        cache_queue.push_back(committed_smt.clone());

        Self {
            commit_receiver,
            db,

            cache_queue,
            version: 0,
            smt: committed_smt.clone(),
            committed_smt,
            updates: Vec::new(),
            num_pending_commits: 0,
        }
    }

    pub fn run(mut self) {
        while let Ok(StateSnapshotDelta {
            version,
            smt,
            jmt_updates,
        }) = self.commit_receiver.recv()
        {
            self.version = version;
            self.smt = smt;
            self.updates.extend(jmt_updates.into_iter());
            self.num_pending_commits += 1;

            if self.num_pending_commits >= NUM_COMMITS_TO_BATCH {
                self.commit();
            }
        }
        println!("Final state commit...");
        self.commit();
    }

    fn commit(&mut self) {
        // commit
        let mut to_commit = Vec::new();
        std::mem::swap(&mut to_commit, &mut self.updates);
        let node_hashes = self
            .smt
            .clone()
            .freeze()
            .new_node_hashes_since(&self.committed_smt.clone().freeze());
        self.db
            .save_state_snapshot(to_commit, Some(&node_hashes), self.version)
            .unwrap();
        info!(version = self.version, "State snapshot saved.");

        // reset pending updates
        self.num_pending_commits = 0;
        self.committed_smt = self.smt.clone();

        // cache maintenance
        if self.cache_queue.len() >= NUM_COMMITTED_SMTS_TO_CACHE {
            self.cache_queue.pop_front();
        }
        self.cache_queue.push_back(self.smt.clone());
    }
}
<commit_msg>[executor-benchmark] make state commit size adaptive<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use aptos_crypto::HashValue;
use aptos_logger::info;
use aptos_types::{
    state_store::{state_key::StateKey, state_value::StateValue},
    transaction::Version,
};
use executor_types::StateSnapshotDelta;
use scratchpad::SparseMerkleTree;
use std::{
    collections::VecDeque,
    sync::{mpsc, mpsc::TryRecvError, Arc},
    time,
};
use storage_interface::DbWriter;

const NUM_COMMITTED_SMTS_TO_CACHE: usize = 5;
const NUM_MIN_COMMITS_TO_BATCH: usize = 20;

pub struct StateCommitter {
    commit_receiver: mpsc::Receiver<StateSnapshotDelta>,
    db: Arc<dyn DbWriter>,

    // keep some recently committed SMTs in mem as a naive cache
    cache_queue: VecDeque<SparseMerkleTree<StateValue>>,
    version: Version,
    smt: SparseMerkleTree<StateValue>,
    committed_smt: SparseMerkleTree<StateValue>,
    updates: Vec<(HashValue, (HashValue, StateKey))>,
    num_pending_commits: usize,
}

impl StateCommitter {
    pub fn new(
        commit_receiver: mpsc::Receiver<StateSnapshotDelta>,
        db: Arc<dyn DbWriter>,
        committed_smt: SparseMerkleTree<StateValue>,
    ) -> Self {
        let mut cache_queue = VecDeque::new();
        cache_queue.push_back(committed_smt.clone());

        Self {
            commit_receiver,
            db,

            cache_queue,
            version: 0,
            smt: committed_smt.clone(),
            committed_smt,
            updates: Vec::new(),
            num_pending_commits: 0,
        }
    }

    pub fn run(mut self) {
        loop {
            match self.commit_receiver.try_recv() {
                Ok(StateSnapshotDelta {
                    version,
                    smt,
                    jmt_updates,
                }) => {
                    self.version = version;
                    self.smt = smt;
                    self.updates.extend(jmt_updates.into_iter());
                    self.num_pending_commits += 1;
                }
                Err(TryRecvError::Empty) => {
                    if self.num_pending_commits < NUM_MIN_COMMITS_TO_BATCH {
                        std::thread::sleep(time::Duration::from_secs(1));
                    } else {
                        self.commit();
                    }
                }
                Err(TryRecvError::Disconnected) => {
                    println!("Final state commit...");
                    self.commit();
                    return;
                }
            }
        }
    }

    fn commit(&mut self) {
        // commit
        info!(
            num_pending_commits = self.num_pending_commits,
            version = self.version,
            "Committing state.",
        );
        let mut to_commit = Vec::new();
        std::mem::swap(&mut to_commit, &mut self.updates);
        let node_hashes = self
            .smt
            .clone()
            .freeze()
            .new_node_hashes_since(&self.committed_smt.clone().freeze());
        self.db
            .save_state_snapshot(to_commit, Some(&node_hashes), self.version)
            .unwrap();
        info!("Committing state. Saved.");

        // reset pending updates
        self.num_pending_commits = 0;
        self.committed_smt = self.smt.clone();

        // cache maintenance
        if self.cache_queue.len() >= NUM_COMMITTED_SMTS_TO_CACHE {
            self.cache_queue.pop_front();
        }
        self.cache_queue.push_back(self.smt.clone());
    }
}
<|endoftext|>


<commit_before>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! System configuration loading for windows

use std::str::FromStr;
use std::net::SocketAddr;

use ipconfig::get_adapters;
use ipconfig::computer::{get_search_list, get_domain, is_round_robin_enabled};

use trust_dns::rr::Name;

use config::{ResolverConfig, ResolverOpts, NameServerConfig, Protocol};
use error::*;

macro_rules! map_ipconfig_error {
    ($result:expr) => ($result.map_err(|e| ResolveErrorKind::Msg("failed to read from registry: {}", e).into()));
}

/// Returns the name servers of the computer (of all adapters)
fn get_name_servers() -> ResolveResult<Vec<NameServerConfig>> {
    let adapters = map_ipconfig_error!(get_adapters())?;
    let mut name_servers = vec![];

    for dns_server in adapters.iter().flat_map(|adapter| adapter.dns_servers().iter()) {
        let socket_addr = SocketAddr::new(*dns_server, 53);
        name_servers.push(NameServerConfig {
            socket_addr,
            protocol: Protocol::Udp,
        });
        name_servers.push(NameServerConfig {
            socket_addr,
            protocol: Protocol::Tcp,
        });
    };
    Ok(name_servers)
}

pub(crate) fn read_system_conf() -> ResolveResult<(ResolverConfig, ResolverOpts)> {
    let name_servers = get_name_servers()?;

    let search_list: Vec<Name> = map_ipconfig_error!(get_search_list())?
        .iter()
        .map(|x| Name::from_str(x))
        .collect::<Result<Vec<_>, _>>()?;

    let domain = match map_ipconfig_error!(get_domain())? {
        Some(domain) => Name::from_str(&domain)?,
        None => Name::root(),
    };

    let config = ResolverConfig::from_parts(domain, search_list, name_servers);

    let rotate = map_ipconfig_error!(is_round_robin_enabled())?;

    let opts = ResolverOpts {
        rotate,
        ..Default::default()
    };
    Ok((config, opts))
}
<commit_msg>fix error type on windows.<commit_after>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! System configuration loading for windows

use std::str::FromStr;
use std::net::SocketAddr;

use ipconfig::get_adapters;
use ipconfig::computer::{get_search_list, get_domain, is_round_robin_enabled};

use trust_dns::rr::Name;

use config::{ResolverConfig, ResolverOpts, NameServerConfig, Protocol};
use error::*;

macro_rules! map_ipconfig_error {
    ($result:expr) => (
            $result.map_err(|e| ResolveErrorKind::Msg(format!("failed to read from registry: {}", e)).into())
        );
}

/// Returns the name servers of the computer (of all adapters)
fn get_name_servers() -> ResolveResult<Vec<NameServerConfig>> {
    let adapters = map_ipconfig_error!(get_adapters())?;
    let mut name_servers = vec![];

    for dns_server in adapters.iter().flat_map(|adapter| adapter.dns_servers().iter()) {
        let socket_addr = SocketAddr::new(*dns_server, 53);
        name_servers.push(NameServerConfig {
            socket_addr,
            protocol: Protocol::Udp,
        });
        name_servers.push(NameServerConfig {
            socket_addr,
            protocol: Protocol::Tcp,
        });
    };
    Ok(name_servers)
}

pub(crate) fn read_system_conf() -> ResolveResult<(ResolverConfig, ResolverOpts)> {
    let name_servers = get_name_servers()?;

    let search_list: Vec<Name> = map_ipconfig_error!(get_search_list())?
        .iter()
        .map(|x| Name::from_str(x))
        .collect::<Result<Vec<_>, _>>()?;

    let domain = match map_ipconfig_error!(get_domain())? {
        Some(domain) => Name::from_str(&domain)?,
        None => Name::root(),
    };

    let config = ResolverConfig::from_parts(domain, search_list, name_servers);

    let rotate = map_ipconfig_error!(is_round_robin_enabled())?;

    let opts = ResolverOpts {
        rotate,
        ..Default::default()
    };
    Ok((config, opts))
}
<|endoftext|>


<commit_before>use super::{BeaconChain, ClientDB, SlotClock};
use types::{AttestationData, Hash256};

#[derive(Debug, PartialEq)]
pub enum Error {
    SlotTooOld,
    PresentSlotUnknown,
    StateError,
}

impl<T, U> BeaconChain<T, U>
where
    T: ClientDB,
    U: SlotClock,
{
    pub fn produce_attestation_data(&self, shard: u64) -> Result<AttestationData, Error> {
        let justified_slot = self.justified_slot();
        let justified_block_root = self
            .state
            .read()
            .get_block_root(justified_slot, &self.spec)
            .ok_or_else(|| Error::SlotTooOld)?
            .clone();

        let epoch_boundary_root = self
            .state
            .read()
            .get_block_root(
                self.state.read().current_epoch_start_slot(&self.spec),
                &self.spec,
            )
            .ok_or_else(|| Error::SlotTooOld)?
            .clone();

        Ok(AttestationData {
            slot: self.state.read().slot,
            shard,
            beacon_block_root: self.head().beacon_block_root.clone(),
            epoch_boundary_root,
            shard_block_root: Hash256::zero(),
            latest_crosslink_root: Hash256::zero(),
            justified_slot,
            justified_block_root,
        })
    }
}
<commit_msg>Tidy attestation_production and create docstrings<commit_after>use super::{BeaconChain, ClientDB, SlotClock};
use types::{AttestationData, Hash256};

#[derive(Debug, PartialEq)]
pub enum Error {
    /// The `justified_block_root` is unknown. This is an internal error.
    UnknownJustifiedRoot,
    /// The `epoch_boundary_root` is unknown. This is an internal error.
    UnknownBoundaryRoot,
}

impl<T, U> BeaconChain<T, U>
where
    T: ClientDB,
    U: SlotClock,
{
    /// Produce an `AttestationData` that is valid for the present `slot` and given `shard`.
    pub fn produce_attestation_data(&self, shard: u64) -> Result<AttestationData, Error> {
        let justified_slot = self.justified_slot();
        let justified_block_root = self
            .state
            .read()
            .get_block_root(justified_slot, &self.spec)
            .ok_or_else(|| Error::UnknownJustifiedRoot)?
            .clone();

        let epoch_boundary_root = self
            .state
            .read()
            .get_block_root(
                self.state.read().current_epoch_start_slot(&self.spec),
                &self.spec,
            )
            .ok_or_else(|| Error::UnknownBoundaryRoot)?
            .clone();

        Ok(AttestationData {
            slot: self.state.read().slot,
            shard,
            beacon_block_root: self.head().beacon_block_root.clone(),
            epoch_boundary_root,
            shard_block_root: Hash256::zero(),
            latest_crosslink_root: Hash256::zero(),
            justified_slot,
            justified_block_root,
        })
    }
}
<|endoftext|>


<commit_before>use rustc_infer::infer::TyCtxtInferExt;
use rustc_middle::traits::query::NoSolution;
use rustc_middle::ty::query::Providers;
use rustc_middle::ty::subst::GenericArg;
use rustc_middle::ty::{self, ParamEnvAnd, TyCtxt};
use rustc_trait_selection::traits::query::normalize::AtExt;
use rustc_trait_selection::traits::{Normalized, ObligationCause};
use std::sync::atomic::Ordering;

crate fn provide(p: &mut Providers) {
    *p = Providers { normalize_generic_arg_after_erasing_regions, ..*p };
}

fn normalize_generic_arg_after_erasing_regions<'tcx>(
    tcx: TyCtxt<'tcx>,
    goal: ParamEnvAnd<'tcx, GenericArg<'tcx>>,
) -> GenericArg<'tcx> {
    debug!("normalize_generic_arg_after_erasing_regions(goal={:#?})", goal);

    let ParamEnvAnd { param_env, value } = goal;
    tcx.sess.perf_stats.normalize_generic_arg_after_erasing_regions.fetch_add(1, Ordering::Relaxed);
    tcx.infer_ctxt().enter(|infcx| {
        let cause = ObligationCause::dummy();
        match infcx.at(&cause, param_env).normalize(value) {
            Ok(Normalized { value: normalized_value, obligations: normalized_obligations }) => {
                // We don't care about the `obligations`; they are
                // always only region relations, and we are about to
                // erase those anyway:
                debug_assert_eq!(
                    normalized_obligations.iter().find(|p| not_outlives_predicate(&p.predicate)),
                    None,
                );

                let normalized_value = infcx.resolve_vars_if_possible(normalized_value);
                infcx.tcx.erase_regions(normalized_value)
            }
            Err(NoSolution) => bug!("could not fully normalize `{:?}`", value),
        }
    })
}

fn not_outlives_predicate(p: &ty::Predicate<'tcx>) -> bool {
    match p.skip_binders() {
        ty::PredicateAtom::RegionOutlives(..) | ty::PredicateAtom::TypeOutlives(..) => false,
        ty::PredicateAtom::Trait(..)
        | ty::PredicateAtom::Projection(..)
        | ty::PredicateAtom::WellFormed(..)
        | ty::PredicateAtom::ObjectSafe(..)
        | ty::PredicateAtom::ClosureKind(..)
        | ty::PredicateAtom::Subtype(..)
        | ty::PredicateAtom::ConstEvaluatable(..)
        | ty::PredicateAtom::ConstEquate(..)
        | ty::PredicateAtom::TypeWellFormedFromEnv(..) => true,
    }
}
<commit_msg>Auto merge of #77616 - jyn514:no-normalize, r=lcnr<commit_after>use rustc_infer::infer::TyCtxtInferExt;
use rustc_middle::traits::query::NoSolution;
use rustc_middle::ty::query::Providers;
use rustc_middle::ty::subst::GenericArg;
use rustc_middle::ty::{self, ParamEnvAnd, TyCtxt, TypeFoldable};
use rustc_trait_selection::traits::query::normalize::AtExt;
use rustc_trait_selection::traits::{Normalized, ObligationCause};
use std::sync::atomic::Ordering;

crate fn provide(p: &mut Providers) {
    *p = Providers { normalize_generic_arg_after_erasing_regions, ..*p };
}

fn normalize_generic_arg_after_erasing_regions<'tcx>(
    tcx: TyCtxt<'tcx>,
    goal: ParamEnvAnd<'tcx, GenericArg<'tcx>>,
) -> GenericArg<'tcx> {
    debug!("normalize_generic_arg_after_erasing_regions(goal={:#?})", goal);

    let ParamEnvAnd { param_env, value } = goal;
    tcx.sess.perf_stats.normalize_generic_arg_after_erasing_regions.fetch_add(1, Ordering::Relaxed);
    tcx.infer_ctxt().enter(|infcx| {
        let cause = ObligationCause::dummy();
        match infcx.at(&cause, param_env).normalize(value) {
            Ok(Normalized { value: normalized_value, obligations: normalized_obligations }) => {
                // We don't care about the `obligations`; they are
                // always only region relations, and we are about to
                // erase those anyway:
                debug_assert_eq!(
                    normalized_obligations.iter().find(|p| not_outlives_predicate(&p.predicate)),
                    None,
                );

                let resolved_value = infcx.resolve_vars_if_possible(normalized_value);
                // It's unclear when `resolve_vars` would have an effect in a
                // fresh `InferCtxt`. If this assert does trigger, it will give
                // us a test case.
                debug_assert_eq!(normalized_value, resolved_value);
                let erased = infcx.tcx.erase_regions(resolved_value);
                debug_assert!(!erased.needs_infer(), "{:?}", erased);
                erased
            }
            Err(NoSolution) => bug!("could not fully normalize `{:?}`", value),
        }
    })
}

fn not_outlives_predicate(p: &ty::Predicate<'tcx>) -> bool {
    match p.skip_binders() {
        ty::PredicateAtom::RegionOutlives(..) | ty::PredicateAtom::TypeOutlives(..) => false,
        ty::PredicateAtom::Trait(..)
        | ty::PredicateAtom::Projection(..)
        | ty::PredicateAtom::WellFormed(..)
        | ty::PredicateAtom::ObjectSafe(..)
        | ty::PredicateAtom::ClosureKind(..)
        | ty::PredicateAtom::Subtype(..)
        | ty::PredicateAtom::ConstEvaluatable(..)
        | ty::PredicateAtom::ConstEquate(..)
        | ty::PredicateAtom::TypeWellFormedFromEnv(..) => true,
    }
}
<|endoftext|>


<commit_before>use decoders::RawImage;
use decoders::cfa::CFA;
use imageops::*;

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct OpDemosaic {
  cfa: String,
}

impl OpDemosaic {
  pub fn new(img: &RawImage) -> OpDemosaic {
    OpDemosaic{
      cfa: img.cropped_cfa().to_string(),
    }
  }
}

impl<'a> ImageOp<'a> for OpDemosaic {
  fn name(&self) -> &str {"demosaic"}
  fn run(&self, pipeline: &mut PipelineGlobals, inid: BufHash, outid: BufHash) {
    let buf = pipeline.cache.get(&inid).unwrap();

    let (scale, nwidth, nheight) = if pipeline.settings.maxwidth == 0 || pipeline.settings.maxheight == 0 {
      (1.0, buf.width, buf.height)
    } else {
      // Do the calculations manually to avoid off-by-one errors from floating point rounding
      let xscale = (buf.width as f32) / (pipeline.settings.maxwidth as f32);
      let yscale = (buf.height as f32) / (pipeline.settings.maxheight as f32);
      if yscale > xscale {
        (yscale, ((buf.width as f32)/yscale) as usize, pipeline.settings.maxheight)
      } else {
        (xscale, pipeline.settings.maxwidth, ((buf.height as f32)/xscale) as usize)
      }
    };

    let cfa = CFA::new(&self.cfa);
    let minscale = match cfa.width {
      2  => 2.0,  // RGGB/RGBE bayer
      6  => 3.0,  // x-trans is 6 wide but has all colors in every 3x3 block
      8  => 2.0,  // Canon pro 70 has a 8x2 patern that has all four colors every 2x2 block
      12 => 12.0, // some crazy sensor I haven't actually encountered, use full block
      _  => 2.0,  // default
    };

    if scale <= 1.0 && buf.colors == 4 {
      // We want full size and the image is already 4 color, pass it through
      pipeline.cache.alias(inid, outid);
    } else if buf.colors == 4 {
      // Scale down a 4 colour image
      let buf = scale_down(&buf, nwidth, nheight);
      pipeline.cache.put(outid, buf, 1);
    } else if scale >= minscale {
      // We're scaling down enough that each pixel has all four colors under it so do the
      // demosaic and scale down in one go
      let buf = scaled(cfa, &buf, nwidth, nheight);
      pipeline.cache.put(outid, buf, 1);
    } else {
      // We're in a close to full scale output that needs full demosaic and possibly
      // minimal scale down
      let fullsize = full(cfa, &buf);
      let buf = if scale > 1.0 {
        scale_down(&fullsize, nwidth, nheight)
      } else {
        fullsize
      };
      pipeline.cache.put(outid, buf, 1);
    }
  }
}

pub fn full(cfa: CFA, buf: &OpBuffer) -> OpBuffer {
  let mut out = OpBuffer::new(buf.width, buf.height, 4);

  let offsets33: [(isize,isize);9] = [
    (-1,-1), (-1, 0), (-1, 1),
    ( 0,-1), ( 0, 0), ( 0, 1),
    ( 1,-1), ( 1, 0), ( 1, 1),
  ];

  // Initialize a lookup table for the colors of each pixel in a 3x3 grid
  let mut lookups = [[[0;9];48];48];
  for (row, line) in lookups.iter_mut().enumerate() {
    for (col, colors) in line.iter_mut().enumerate() {
      let pixcolor = cfa.color_at(row, col);

      for (i, o) in offsets33.iter().enumerate() {
        let (dy, dx) = *o;
        let row = (48+dy) as usize + row;
        let col = (48+dx) as usize + col;
        let ocolor = cfa.color_at(row, col);
        colors[i] = if ocolor != pixcolor || (dx == 0 && dy == 0) { ocolor } else { 4 };
      }
    }
  }

  // Now calculate RGBE for each pixel based on the lookup table
  out.mutate_lines(&(|line: &mut [f32], row| {
    for (col, pix) in line.chunks_mut(4).enumerate() {
      let ref colors = lookups[row%48][col%48];
      let mut sums = [0f32;5];
      let mut counts = [0f32;5];

      for (i, o) in offsets33.iter().enumerate() {
        let (dy, dx) = *o;
        let row = row as isize + dy;
        let col = col as isize + dx;
        if row >= 0 && row < (buf.height as isize) &&
           col >= 0 && col < (buf.width as isize) {
          sums[colors[i]] += buf.data[(row as usize)*buf.width+(col as usize)];
          counts[colors[i]] += 1.0;
        }
      }

      for c in 0..4 {
        if counts[c] > 0.0 {
          pix[c] = sums[c] / counts[c];
        }
      }
    }
  }));

  out
}

fn calc_skips(idx: usize, idxmax: usize, skip: f32) -> (usize, usize, f32, f32) {
  let from = (idx as f32)*skip;
  let fromback = from.floor();
  let fromfactor = 1.0 - (from-fromback).fract();

  let to = ((idx+1) as f32)*skip;
  let toforward = (idxmax as f32).min(to.ceil());
  let tofactor = (toforward-to).fract();

  (fromback as usize, toforward as usize, fromfactor, tofactor)
}

pub fn scaled(cfa: CFA, buf: &OpBuffer, nwidth: usize, nheight: usize) -> OpBuffer {
  let mut out = OpBuffer::new(nwidth, nheight, 4);

  let rowskip = (buf.width as f32) / (nwidth as f32);
  let colskip = (buf.height as f32) / (nheight as f32);

  // Go around the image averaging blocks of pixels
  out.mutate_lines(&(|line: &mut [f32], row| {
    for col in 0..nwidth {
      let mut sums: [f32; 4] = [0.0;4];
      let mut counts: [f32; 4] = [0.0;4];
      let (fromrow, torow, topfactor, bottomfactor) = calc_skips(row, buf.height, rowskip);
      for y in fromrow..torow {
        let (fromcol, tocol, leftfactor, rightfactor) = calc_skips(col, buf.width, colskip);
        for x in fromcol..tocol {
          let factor = {
            (if y == fromrow {topfactor} else if y == torow {bottomfactor} else {1.0}) *
            (if x == fromcol {leftfactor} else if x == tocol {rightfactor} else {1.0})
          };

          let c = cfa.color_at(y, x);
          sums[c] += (buf.data[y*buf.width+x] as f32) * factor;
          counts[c] += factor;
        }
      }

      for c in 0..4 {
        if counts[c] > 0.0 {
          line[col*4+c] = sums[c] / counts[c];
        }
      }
    }
  }));

  out
}

pub fn scale_down(buf: &OpBuffer, nwidth: usize, nheight: usize) -> OpBuffer {
  assert_eq!(buf.colors, 4); // When we're scaling down we're always at 4 cpp

  let mut out = OpBuffer::new(nwidth, nheight, 4);
  let rowskip = (buf.width as f32) / (nwidth as f32);
  let colskip = (buf.height as f32) / (nheight as f32);

  // Go around the image averaging blocks of pixels
  out.mutate_lines(&(|line: &mut [f32], row| {
    for col in 0..nwidth {
      let mut sums: [f32; 4] = [0.0;4];
      let mut counts: [f32; 4] = [0.0;4];
      let (fromrow, torow, topfactor, bottomfactor) = calc_skips(row, buf.height, rowskip);
      for y in fromrow..torow {
        let (fromcol, tocol, leftfactor, rightfactor) = calc_skips(col, buf.width, colskip);
        for x in fromcol..tocol {
          let factor = {
            (if y == fromrow {topfactor} else if y == torow {bottomfactor} else {1.0}) *
            (if x == fromcol {leftfactor} else if x == tocol {rightfactor} else {1.0})
          };

          for c in 0..4 {
            sums[c] += buf.data[(y*buf.width+x)*4 + c] * factor;
            counts[c] += factor;
          }
        }
      }

      for c in 0..4 {
        if counts[c] > 0.0 {
          line[col*4+c] = sums[c] / counts[c];
        }
      }
    }
  }));

  out
}
<commit_msg>demosaic: slight naming cleanup<commit_after>use decoders::RawImage;
use decoders::cfa::CFA;
use imageops::*;

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct OpDemosaic {
  cfa: String,
}

impl OpDemosaic {
  pub fn new(img: &RawImage) -> OpDemosaic {
    OpDemosaic{
      cfa: img.cropped_cfa().to_string(),
    }
  }
}

impl<'a> ImageOp<'a> for OpDemosaic {
  fn name(&self) -> &str {"demosaic"}
  fn run(&self, pipeline: &mut PipelineGlobals, inid: BufHash, outid: BufHash) {
    let buf = pipeline.cache.get(&inid).unwrap();

    let (scale, nwidth, nheight) = if pipeline.settings.maxwidth == 0 || pipeline.settings.maxheight == 0 {
      (1.0, buf.width, buf.height)
    } else {
      // Do the calculations manually to avoid off-by-one errors from floating point rounding
      let xscale = (buf.width as f32) / (pipeline.settings.maxwidth as f32);
      let yscale = (buf.height as f32) / (pipeline.settings.maxheight as f32);
      if yscale > xscale {
        (yscale, ((buf.width as f32)/yscale) as usize, pipeline.settings.maxheight)
      } else {
        (xscale, pipeline.settings.maxwidth, ((buf.height as f32)/xscale) as usize)
      }
    };

    let cfa = CFA::new(&self.cfa);
    let minscale = match cfa.width {
      2  => 2.0,  // RGGB/RGBE bayer
      6  => 3.0,  // x-trans is 6 wide but has all colors in every 3x3 block
      8  => 2.0,  // Canon pro 70 has a 8x2 patern that has all four colors every 2x2 block
      12 => 12.0, // some crazy sensor I haven't actually encountered, use full block
      _  => 2.0,  // default
    };

    if scale <= 1.0 && buf.colors == 4 {
      // We want full size and the image is already 4 color, pass it through
      pipeline.cache.alias(inid, outid);
    } else if buf.colors == 4 {
      // Scale down a 4 colour image
      let buf = scale_down(&buf, nwidth, nheight);
      pipeline.cache.put(outid, buf, 1);
    } else if scale >= minscale {
      // We're scaling down enough that each pixel has all four colors under it so do the
      // demosaic and scale down in one go
      let buf = scaled(cfa, &buf, nwidth, nheight);
      pipeline.cache.put(outid, buf, 1);
    } else {
      // We're in a close to full scale output that needs full demosaic and possibly
      // minimal scale down
      let fullsize = full(cfa, &buf);
      let buf = if scale > 1.0 {
        scale_down(&fullsize, nwidth, nheight)
      } else {
        fullsize
      };
      pipeline.cache.put(outid, buf, 1);
    }
  }
}

pub fn full(cfa: CFA, buf: &OpBuffer) -> OpBuffer {
  let mut out = OpBuffer::new(buf.width, buf.height, 4);

  let offsets3x3: [(isize,isize);9] = [
    (-1,-1), (-1, 0), (-1, 1),
    ( 0,-1), ( 0, 0), ( 0, 1),
    ( 1,-1), ( 1, 0), ( 1, 1),
  ];

  // Initialize a lookup table for the colors of each pixel in a 3x3 grid
  let mut lookups = [[[0;9];48];48];
  for (row, line) in lookups.iter_mut().enumerate() {
    for (col, colors) in line.iter_mut().enumerate() {
      let pixcolor = cfa.color_at(row, col);

      for (i, o) in offsets3x3.iter().enumerate() {
        let (dy, dx) = *o;
        let row = (48+dy) as usize + row;
        let col = (48+dx) as usize + col;
        let ocolor = cfa.color_at(row, col);
        colors[i] = if ocolor != pixcolor || (dx == 0 && dy == 0) { ocolor } else { 4 };
      }
    }
  }

  // Now calculate RGBE for each pixel based on the lookup table
  out.mutate_lines(&(|line: &mut [f32], row| {
    for (col, pix) in line.chunks_mut(4).enumerate() {
      let ref colors = lookups[row%48][col%48];
      let mut sums = [0f32;5];
      let mut counts = [0f32;5];

      for (i, o) in offsets3x3.iter().enumerate() {
        let (dy, dx) = *o;
        let row = row as isize + dy;
        let col = col as isize + dx;
        if row >= 0 && row < (buf.height as isize) &&
           col >= 0 && col < (buf.width as isize) {
          sums[colors[i]] += buf.data[(row as usize)*buf.width+(col as usize)];
          counts[colors[i]] += 1.0;
        }
      }

      for c in 0..4 {
        if counts[c] > 0.0 {
          pix[c] = sums[c] / counts[c];
        }
      }
    }
  }));

  out
}

fn calc_skips(idx: usize, idxmax: usize, skip: f32) -> (usize, usize, f32, f32) {
  let from = (idx as f32)*skip;
  let fromback = from.floor();
  let fromfactor = 1.0 - (from-fromback).fract();

  let to = ((idx+1) as f32)*skip;
  let toforward = (idxmax as f32).min(to.ceil());
  let tofactor = (toforward-to).fract();

  (fromback as usize, toforward as usize, fromfactor, tofactor)
}

pub fn scaled(cfa: CFA, buf: &OpBuffer, nwidth: usize, nheight: usize) -> OpBuffer {
  let mut out = OpBuffer::new(nwidth, nheight, 4);

  let rowskip = (buf.width as f32) / (nwidth as f32);
  let colskip = (buf.height as f32) / (nheight as f32);

  // Go around the image averaging blocks of pixels
  out.mutate_lines(&(|line: &mut [f32], row| {
    for col in 0..nwidth {
      let mut sums: [f32; 4] = [0.0;4];
      let mut counts: [f32; 4] = [0.0;4];
      let (fromrow, torow, topfactor, bottomfactor) = calc_skips(row, buf.height, rowskip);
      for y in fromrow..torow {
        let (fromcol, tocol, leftfactor, rightfactor) = calc_skips(col, buf.width, colskip);
        for x in fromcol..tocol {
          let factor = {
            (if y == fromrow {topfactor} else if y == torow {bottomfactor} else {1.0}) *
            (if x == fromcol {leftfactor} else if x == tocol {rightfactor} else {1.0})
          };

          let c = cfa.color_at(y, x);
          sums[c] += (buf.data[y*buf.width+x] as f32) * factor;
          counts[c] += factor;
        }
      }

      for c in 0..4 {
        if counts[c] > 0.0 {
          line[col*4+c] = sums[c] / counts[c];
        }
      }
    }
  }));

  out
}

pub fn scale_down(buf: &OpBuffer, nwidth: usize, nheight: usize) -> OpBuffer {
  assert_eq!(buf.colors, 4); // When we're scaling down we're always at 4 cpp

  let mut out = OpBuffer::new(nwidth, nheight, 4);
  let rowskip = (buf.width as f32) / (nwidth as f32);
  let colskip = (buf.height as f32) / (nheight as f32);

  // Go around the image averaging blocks of pixels
  out.mutate_lines(&(|line: &mut [f32], row| {
    for col in 0..nwidth {
      let mut sums: [f32; 4] = [0.0;4];
      let mut counts: [f32; 4] = [0.0;4];
      let (fromrow, torow, topfactor, bottomfactor) = calc_skips(row, buf.height, rowskip);
      for y in fromrow..torow {
        let (fromcol, tocol, leftfactor, rightfactor) = calc_skips(col, buf.width, colskip);
        for x in fromcol..tocol {
          let factor = {
            (if y == fromrow {topfactor} else if y == torow {bottomfactor} else {1.0}) *
            (if x == fromcol {leftfactor} else if x == tocol {rightfactor} else {1.0})
          };

          for c in 0..4 {
            sums[c] += buf.data[(y*buf.width+x)*4 + c] * factor;
            counts[c] += factor;
          }
        }
      }

      for c in 0..4 {
        if counts[c] > 0.0 {
          line[col*4+c] = sums[c] / counts[c];
        }
      }
    }
  }));

  out
}
<|endoftext|>


<commit_before>//! Digital signatures and collections of signatures.

use std::collections::HashMap;

use base64::{encode_config, STANDARD_NO_PAD};

use crate::{split_id, Algorithm, Error, SplitError};

/// A digital signature.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct Signature {
    /// The cryptographic algorithm that generated this signature.
    pub algorithm: Algorithm,

    /// The signature data.
    pub signature: Vec<u8>,

    /// The "version" of the key identifier for the public key used to generate this signature.
    pub version: String,
}

impl Signature {
    /// Creates a signature from raw bytes.
    ///
    /// While a signature can be created directly using struct literal syntax, this constructor can
    /// be used to automatically determine the algorithm and version from a key identifier in the
    /// form *algorithm:version*, e.g. "ed25519:1".
    ///
    /// This constructor will ensure that the version does not contain characters that violate the
    /// guidelines in the specification. Because it may be necessary to represent signatures with
    /// versions that don't adhere to these guidelines, it's possible to simply use the struct
    /// literal syntax to construct a `Signature` with an arbitrary key.
    ///
    /// # Parameters
    ///
    /// * id: A key identifier, e.g. "ed25519:1".
    /// * bytes: The digital signature, as a series of bytes.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    ///
    /// * The key ID specifies an unknown algorithm.
    /// * The key ID is malformed.
    /// * The key ID contains a version with invalid characters.
    pub fn new(id: &str, bytes: &[u8]) -> Result<Self, Error> {
        let (algorithm, version) = split_id(id).map_err(|split_error| match split_error {
            SplitError::InvalidLength(length) => Error::new(format!(
                "malformed signature ID: expected exactly \
                 2 segment separated by a colon, found {}",
                length
            )),
            SplitError::InvalidVersion(version) => Error::new(format!(
                "malformed signature ID: expected version to contain only \
                 characters in the character set `[a-zA-Z0-9_]`, found `{}`",
                version
            )),
            SplitError::UnknownAlgorithm(algorithm) => {
                Error::new(format!("unknown algorithm: {}", algorithm))
            }
        })?;

        Ok(Self { algorithm, signature: bytes.to_vec(), version })
    }

    /// The algorithm used to generate the signature.
    pub fn algorithm(&self) -> Algorithm {
        self.algorithm
    }

    /// The raw bytes of the signature.
    pub fn as_bytes(&self) -> &[u8] {
        self.signature.as_slice()
    }

    /// A Base64 encoding of the signature.
    ///
    /// Uses the standard character set with no padding.
    pub fn base64(&self) -> String {
        encode_config(self.signature.as_slice(), STANDARD_NO_PAD)
    }

    /// The key identifier, a string containing the signature algorithm and the key "version"
    /// separated by a colon, e.g. "ed25519:1".
    pub fn id(&self) -> String {
        format!("{}:{}", self.algorithm, self.version)
    }

    /// The "version" of the key used for this signature.
    ///
    /// Versions are used as an identifier to distinguish signatures generated from different keys
    /// but using the same algorithm on the same homeserver.
    pub fn version(&self) -> &str {
        &self.version
    }
}

/// A map from entity names to sets of digital signatures created by that entity.
///
/// "Entity" is generally a homeserver, e.g. "example.com".
pub type SignatureMap = HashMap<String, SignatureSet>;

/// A set of digital signatures created by a single homeserver.
///
/// This is represented as a map from signing key ID to Base64-encoded signature.
pub type SignatureSet = HashMap<String, String>;

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

    #[test]
    fn valid_key_id() {
        assert!(Signature::new("ed25519:abcdef", &[]).is_ok());
    }

    #[test]
    fn invalid_valid_key_id_length() {
        assert!(Signature::new("ed25519:abcdef:123456", &[]).is_err());
    }

    #[test]
    fn invalid_key_id_version() {
        assert!(Signature::new("ed25519:abc!def", &[]).is_err());
    }

    #[test]
    fn invalid_key_id_algorithm() {
        assert!(Signature::new("foobar:abcdef", &[]).is_err());
    }
}
<commit_msg>signatures: Replace remaining uses of HashMap with BTreeMap<commit_after>//! Digital signatures and collections of signatures.

use std::collections::BTreeMap;

use base64::{encode_config, STANDARD_NO_PAD};

use crate::{split_id, Algorithm, Error, SplitError};

/// A digital signature.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct Signature {
    /// The cryptographic algorithm that generated this signature.
    pub algorithm: Algorithm,

    /// The signature data.
    pub signature: Vec<u8>,

    /// The "version" of the key identifier for the public key used to generate this signature.
    pub version: String,
}

impl Signature {
    /// Creates a signature from raw bytes.
    ///
    /// While a signature can be created directly using struct literal syntax, this constructor can
    /// be used to automatically determine the algorithm and version from a key identifier in the
    /// form *algorithm:version*, e.g. "ed25519:1".
    ///
    /// This constructor will ensure that the version does not contain characters that violate the
    /// guidelines in the specification. Because it may be necessary to represent signatures with
    /// versions that don't adhere to these guidelines, it's possible to simply use the struct
    /// literal syntax to construct a `Signature` with an arbitrary key.
    ///
    /// # Parameters
    ///
    /// * id: A key identifier, e.g. "ed25519:1".
    /// * bytes: The digital signature, as a series of bytes.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    ///
    /// * The key ID specifies an unknown algorithm.
    /// * The key ID is malformed.
    /// * The key ID contains a version with invalid characters.
    pub fn new(id: &str, bytes: &[u8]) -> Result<Self, Error> {
        let (algorithm, version) = split_id(id).map_err(|split_error| match split_error {
            SplitError::InvalidLength(length) => Error::new(format!(
                "malformed signature ID: expected exactly \
                 2 segment separated by a colon, found {}",
                length
            )),
            SplitError::InvalidVersion(version) => Error::new(format!(
                "malformed signature ID: expected version to contain only \
                 characters in the character set `[a-zA-Z0-9_]`, found `{}`",
                version
            )),
            SplitError::UnknownAlgorithm(algorithm) => {
                Error::new(format!("unknown algorithm: {}", algorithm))
            }
        })?;

        Ok(Self { algorithm, signature: bytes.to_vec(), version })
    }

    /// The algorithm used to generate the signature.
    pub fn algorithm(&self) -> Algorithm {
        self.algorithm
    }

    /// The raw bytes of the signature.
    pub fn as_bytes(&self) -> &[u8] {
        self.signature.as_slice()
    }

    /// A Base64 encoding of the signature.
    ///
    /// Uses the standard character set with no padding.
    pub fn base64(&self) -> String {
        encode_config(self.signature.as_slice(), STANDARD_NO_PAD)
    }

    /// The key identifier, a string containing the signature algorithm and the key "version"
    /// separated by a colon, e.g. "ed25519:1".
    pub fn id(&self) -> String {
        format!("{}:{}", self.algorithm, self.version)
    }

    /// The "version" of the key used for this signature.
    ///
    /// Versions are used as an identifier to distinguish signatures generated from different keys
    /// but using the same algorithm on the same homeserver.
    pub fn version(&self) -> &str {
        &self.version
    }
}

/// A map from entity names to sets of digital signatures created by that entity.
///
/// "Entity" is generally a homeserver, e.g. "example.com".
pub type SignatureMap = BTreeMap<String, SignatureSet>;

/// A set of digital signatures created by a single homeserver.
///
/// This is represented as a map from signing key ID to Base64-encoded signature.
pub type SignatureSet = BTreeMap<String, String>;

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

    #[test]
    fn valid_key_id() {
        assert!(Signature::new("ed25519:abcdef", &[]).is_ok());
    }

    #[test]
    fn invalid_valid_key_id_length() {
        assert!(Signature::new("ed25519:abcdef:123456", &[]).is_err());
    }

    #[test]
    fn invalid_key_id_version() {
        assert!(Signature::new("ed25519:abc!def", &[]).is_err());
    }

    #[test]
    fn invalid_key_id_algorithm() {
        assert!(Signature::new("foobar:abcdef", &[]).is_err());
    }
}
<|endoftext|>


<commit_before>// Edge Operation L-System Genome

pub mod edgeop;
mod expr_op;
mod cond_op;

use evo::prob::{Probability, ProbabilityValue};
use evo::crossover::linear_2point_crossover_random;
use evo::nsga2::{self, FitnessEval, Mate, MultiObjective3};
use rand::Rng;
use rand::distributions::{IndependentSample, Weighted};
use rand::distributions::range::Range;
use graph_annealing::owned_weighted_choice::OwnedWeightedChoice;
use graph_annealing::fitness_function::FitnessFunction;
use graph_annealing::goal::Goal;
use std::str::FromStr;
use triadic_census::OptDenseDigraph;
use lindenmayer_system::{Alphabet, Condition, Symbol, SymbolString, LSystem, Rule, apply_first_rule};
use lindenmayer_system::symbol::Sym2;
use lindenmayer_system::expr::Expr;
use self::edgeop::{EdgeOp, edgeops_to_graph};
use self::expr_op::{ConstExprOp, ExprOp, random_const_expr, random_expr};
use self::cond_op::{CondOp, random_cond};
use simple_parallel::Pool;
use crossbeam;
use std::collections::BTreeMap;

/// Element-wise Mutation operation.
defops!{MutOp;
    // No mutation (copy element)
    Copy,
    // Insert new operation
    Insert,
    // Remove an operation
    Remove,
    // Modify the operation
    ModifyOp,
    // Modify a parameter value
    ModifyParam,
    // Modify both operation and parameter
    Replace
}

/// Variation operators
defops!{VarOp;
    // No variation. Reproduce exactly
    Copy,

    // Mutate
    Mutate,

    // 2-point Linear crossover
    LinearCrossover2,

    // Uniform crossover
    UniformCrossover
}

type RuleId = u32;

// The alphabet of terminal and non-terminals we use.
// Non-terminals are our rules, while terminals are the EdgeOps.
#[derive(Debug, PartialEq, Eq, Clone)]
enum EdgeAlphabet {
    Terminal(EdgeOp),
    NonTerminal(RuleId),
}

impl Alphabet for EdgeAlphabet {}

// We use 2-ary symbols, i.e. symbols with two parameters.
type Sym = Sym2<EdgeAlphabet, f32>;

/// Rules can only be stored for NonTerminals
#[derive(Clone, Debug)]
struct System {
    rules: BTreeMap<RuleId, Vec<Rule<Sym>>>,
}

impl System {
    fn new() -> System {
        System {
            rules: BTreeMap::new(),
        }
    }
}

impl LSystem<Sym> for System {
    fn apply_first_rule(&self, sym: &Sym) -> Option<SymbolString<Sym>> {
        match sym.symbol() {
            // We don't store rules for terminal symbols.
            &EdgeAlphabet::Terminal(_) => {
                None
            }

            &EdgeAlphabet::NonTerminal(id) => {
                self.rules.get(&id).and_then(|rules| apply_first_rule(&rules[..], sym))
            }
        }
    }
}

pub struct SymbolGenerator {
    pub max_expr_depth: usize,
    pub expr_weighted_op: OwnedWeightedChoice<ExprOp>,
    pub expr_weighted_op_max_depth: OwnedWeightedChoice<ExprOp>,

    // terminal symbols
    pub edge_weighted_op: OwnedWeightedChoice<EdgeOp>,

    pub nonterminal_symbols: Range<u32>,

    /// The probability with which a terminal value is choosen.
    pub prob_terminal: Probability,

    pub const_expr_weighted_op: OwnedWeightedChoice<ConstExprOp>,
}

impl SymbolGenerator {
    /// Generates a random production (right hand-side of a rule).
    pub fn gen_symbolstring<R, S>(&self,
                                  rng: &mut R,
                                  len: usize,
                                  arity: usize,
                                  num_params: usize)
                                  -> SymbolString<S>
        where R: Rng,
              S: Symbol<A = EdgeAlphabet, T = f32>
    {
        SymbolString((0..len)
                         .into_iter()
                         .map(|_| self.gen_symbol(rng, arity, num_params))
                         .collect())
    }

    pub fn gen_symbol<R, S>(&self, rng: &mut R, arity: usize, num_params: usize) -> S
        where R: Rng,
              S: Symbol<A = EdgeAlphabet, T = f32>
    {
        S::from_iter(self.gen_symbol_value(rng),
                     (0..arity).into_iter().map(|_| self.gen_expr(rng, num_params)))

    }

    fn gen_expr<R: Rng>(&self, rng: &mut R, num_params: usize) -> Expr<f32> {
        random_expr(rng,
                    num_params,
                    self.max_expr_depth,
                    &self.expr_weighted_op,
                    &self.expr_weighted_op_max_depth)
    }

    fn gen_terminal<R: Rng>(&self, rng: &mut R) -> EdgeAlphabet {
        EdgeAlphabet::Terminal(self.edge_weighted_op.ind_sample(rng))
    }

    fn gen_nonterminal<R: Rng>(&self, rng: &mut R) -> EdgeAlphabet {
        EdgeAlphabet::NonTerminal(self.nonterminal_symbols.ind_sample(rng))
    }

    fn gen_symbol_value<R: Rng>(&self, rng: &mut R) -> EdgeAlphabet {
        if rng.gen::<ProbabilityValue>().is_probable_with(self.prob_terminal) {
            self.gen_terminal(rng)
        } else {
            self.gen_nonterminal(rng)
        }
    }

    /// Generate a simple condition like:
    ///     Arg(n) or 0.0 [>=] or [<=] constant expr
    fn gen_simple_rule_condition<R: Rng>(&self, rng: &mut R, arity: usize) -> Condition<f32> {
        let lhs = if arity > 0 {
            Expr::Arg(rng.gen_range(0, arity))
        } else {
            Expr::Const(0.0)
        };

        let rhs = random_const_expr(rng, &self.const_expr_weighted_op);

        if rng.gen::<bool>() {
            Condition::GreaterEqual(Box::new(lhs), Box::new(rhs))
        } else {
            Condition::LessEqual(Box::new(lhs), Box::new(rhs))
        }
    }

    // Generate a random rule with `symbol` and `arity` parameters.
    //
    // fn gen_rule<R: Rng>(&self, symbol: EdgeAlphabet, arity: usize) {
    // }
}

// When we need a value within (0, 1], we simply cut off the integer part.
// Max depth.

// # Genome interpretation:
//
//     * We use rule with number 0 as axiom.
//       The arguments for the axiom are fixed and passed in from the command line.
//       For simplicity, these arguments are duplicated in each Genome in `axiom_params`.
//       One could also apply genetic operations onto these.
//
//
#[derive(Clone, Debug)]
pub struct Genome {
    system: System,
}

pub struct Toolbox<N, E> {
    goal: Goal<N, E>,
    pool: Pool,
    fitness_functions: (FitnessFunction, FitnessFunction, FitnessFunction),

    weighted_op: OwnedWeightedChoice<EdgeOp>,
    weighted_var_op: OwnedWeightedChoice<VarOp>,
    weighted_mut_op: OwnedWeightedChoice<MutOp>,
    prob_mutate_elem: Probability,

    /// Arguments to the axiom rule.
    pub axiom_args: Vec<Expr<f32>>,

    /// Maximum number of iterations of the L-system  (XXX: Limit also based on generated length)
    pub iterations: usize,
}

impl<N: Clone + Default, E: Clone + Default> Toolbox<N, E> {
    pub fn new(goal: Goal<N, E>,
               pool: Pool,
               fitness_functions: (FitnessFunction, FitnessFunction, FitnessFunction),
               weighted_op: Vec<Weighted<EdgeOp>>,
               weighted_var_op: Vec<Weighted<VarOp>>,
               weighted_mut_op: Vec<Weighted<MutOp>>,
               prob_mutate_elem: Probability)
               -> Toolbox<N, E> {

        Toolbox {
            goal: goal,
            pool: pool,
            fitness_functions: fitness_functions,
            prob_mutate_elem: prob_mutate_elem,
            weighted_op: OwnedWeightedChoice::new(weighted_op),
            weighted_var_op: OwnedWeightedChoice::new(weighted_var_op),
            weighted_mut_op: OwnedWeightedChoice::new(weighted_mut_op),
            axiom_args: vec![],
            iterations: 10,
        }
    }

    pub fn mutate<R: Rng>(&self, rng: &mut R, ind: &Genome) -> Genome {
        // let mut mut_ind = Vec::with_capacity(ind.len() + 1);
        //
        // for edge_op in ind.edge_ops.iter() {
        // let new_op = if rng.gen::<ProbabilityValue>().is_probable_with(self.prob_mutate_elem) {
        // match self.weighted_mut_op.ind_sample(rng) {
        // MutOp::Copy => edge_op.clone(),
        // MutOp::Insert => {
        // Insert a new op before the current operation
        // mut_ind.push(self.generate_random_edge_operation(rng));
        // edge_op.clone()
        // }
        // MutOp::Remove => {
        // remove current operation
        // continue;
        // }
        // MutOp::ModifyOp => (self.weighted_op.ind_sample(rng), edge_op.1),
        // MutOp::ModifyParam => (edge_op.0, rng.gen::<f32>()),
        // MutOp::Replace => self.generate_random_edge_operation(rng),
        // }
        // } else {
        // edge_op.clone()
        // };
        // mut_ind.push(new_op);
        // }
        //

        Genome {
            system: System::new(),
        }
    }

    fn generate_random_edge_operation<R: Rng>(&self, rng: &mut R) -> (EdgeOp, f32) {
        generate_random_edge_operation(&self.weighted_op, rng)
    }

    // There are many parameters that influence the creation of a random
    // genome:
    //
    //     - Number of rules
    //     - Arity of symbols
    //     - Axiom and Length
    //     - Length of a production rule
    //     - Number of (condition, successor) pairs per rule.
    //     - Complexity of expression in Symbol
    //
    //     - Number of Iterations.
    //
    //     There are many
    //
    pub fn random_genome<R: Rng>(&self, rng: &mut R) -> Genome {
        Genome {
            system: System::new(),
        }

    }
}

impl<N:Clone+Sync+Default,E:Clone+Sync+Default> FitnessEval<Genome, MultiObjective3<f32>> for Toolbox<N,E> {
    /// Evaluates the fitness of a Genome population.
    fn fitness(&mut self, pop: &[Genome]) -> Vec<MultiObjective3<f32>> {
        let pool = &mut self.pool;
        let goal = &self.goal;
        let axiom_args = &self.axiom_args[..];
        let iterations = self.iterations;

        let fitness_functions = self.fitness_functions;

        crossbeam::scope(|scope| {
            pool.map(scope, pop, |ind| {
                let edge_ops = ind.to_edge_ops(axiom_args, iterations);
                let g = edgeops_to_graph(&edge_ops);

                //let g = ind.to_graph();
                MultiObjective3::from((goal.apply_fitness_function(fitness_functions.0, &g),
                                       goal.apply_fitness_function(fitness_functions.1, &g),
                                       goal.apply_fitness_function(fitness_functions.2, &g)))

            })
                .collect()
        })
    }
}

impl<N: Clone + Default, E: Clone + Default> Mate<Genome> for Toolbox<N, E> {
    // p1 is potentially "better" than p2
    fn mate<R: Rng>(&mut self, rng: &mut R, p1: &Genome, p2: &Genome) -> Genome {

        match self.weighted_var_op.ind_sample(rng) {
            VarOp::Copy => p1.clone(),
            VarOp::Mutate => self.mutate(rng, p1),
            VarOp::LinearCrossover2 => {
                Genome {
                    // edge_ops: linear_2point_crossover_random(rng,
                    // &p1.edge_ops[..],
                    // &p2.edge_ops[..]),
                    system: System::new(),
                }
            }
            VarOp::UniformCrossover => {
                panic!("TODO");
            }
        }
    }
}

impl Genome {
    /// Develops the L-system into a vector of edge operations
    pub fn to_edge_ops(&self, axiom_args: &[Expr<f32>], iterations: usize) -> Vec<(EdgeOp, f32)> {
        let mut edge_ops = vec![];

        //system

        return edge_ops;
    }
}

#[inline]
fn generate_random_edge_operation<R: Rng>(weighted_op: &OwnedWeightedChoice<EdgeOp>,
                                          rng: &mut R)
                                          -> (EdgeOp, f32) {
    (weighted_op.ind_sample(rng), rng.gen::<f32>())
}
<commit_msg>implement L-system production<commit_after>// Edge Operation L-System Genome

pub mod edgeop;
mod expr_op;
mod cond_op;

use evo::prob::{Probability, ProbabilityValue};
use evo::crossover::linear_2point_crossover_random;
use evo::nsga2::{self, FitnessEval, Mate, MultiObjective3};
use rand::Rng;
use rand::distributions::{IndependentSample, Weighted};
use rand::distributions::range::Range;
use graph_annealing::owned_weighted_choice::OwnedWeightedChoice;
use graph_annealing::fitness_function::FitnessFunction;
use graph_annealing::goal::Goal;
use std::str::FromStr;
use triadic_census::OptDenseDigraph;
use lindenmayer_system::{Alphabet, Condition, Symbol, SymbolString, LSystem, Rule, apply_first_rule};
use lindenmayer_system::symbol::Sym2;
use lindenmayer_system::expr::Expr;
use self::edgeop::{EdgeOp, edgeops_to_graph};
use self::expr_op::{ConstExprOp, ExprOp, random_const_expr, random_expr};
use self::cond_op::{CondOp, random_cond};
use simple_parallel::Pool;
use crossbeam;
use std::collections::BTreeMap;

/// Element-wise Mutation operation.
defops!{MutOp;
    // No mutation (copy element)
    Copy,
    // Insert new operation
    Insert,
    // Remove an operation
    Remove,
    // Modify the operation
    ModifyOp,
    // Modify a parameter value
    ModifyParam,
    // Modify both operation and parameter
    Replace
}

/// Variation operators
defops!{VarOp;
    // No variation. Reproduce exactly
    Copy,

    // Mutate
    Mutate,

    // 2-point Linear crossover
    LinearCrossover2,

    // Uniform crossover
    UniformCrossover
}

type RuleId = u32;

// The alphabet of terminal and non-terminals we use.
// Non-terminals are our rules, while terminals are the EdgeOps.
#[derive(Debug, PartialEq, Eq, Clone)]
enum EdgeAlphabet {
    Terminal(EdgeOp),
    NonTerminal(RuleId),
}

impl Alphabet for EdgeAlphabet {}

// We use 2-ary symbols, i.e. symbols with two parameters.
type Sym = Sym2<EdgeAlphabet, f32>;

/// Rules can only be stored for NonTerminals
#[derive(Clone, Debug)]
struct System {
    rules: BTreeMap<RuleId, Vec<Rule<Sym>>>,
}

impl System {
    fn new() -> System {
        System {
            rules: BTreeMap::new(),
        }
    }
}

impl LSystem<Sym> for System {
    fn apply_first_rule(&self, sym: &Sym) -> Option<SymbolString<Sym>> {
        match sym.symbol() {
            // We don't store rules for terminal symbols.
            &EdgeAlphabet::Terminal(_) => {
                None
            }

            &EdgeAlphabet::NonTerminal(id) => {
                self.rules.get(&id).and_then(|rules| apply_first_rule(&rules[..], sym))
            }
        }
    }
}

pub struct SymbolGenerator {
    pub max_expr_depth: usize,
    pub expr_weighted_op: OwnedWeightedChoice<ExprOp>,
    pub expr_weighted_op_max_depth: OwnedWeightedChoice<ExprOp>,

    // terminal symbols
    pub edge_weighted_op: OwnedWeightedChoice<EdgeOp>,

    pub nonterminal_symbols: Range<u32>,

    /// The probability with which a terminal value is choosen.
    pub prob_terminal: Probability,

    pub const_expr_weighted_op: OwnedWeightedChoice<ConstExprOp>,
}

impl SymbolGenerator {
    /// Generates a random production (right hand-side of a rule).
    pub fn gen_symbolstring<R, S>(&self,
                                  rng: &mut R,
                                  len: usize,
                                  arity: usize,
                                  num_params: usize)
                                  -> SymbolString<S>
        where R: Rng,
              S: Symbol<A = EdgeAlphabet, T = f32>
    {
        SymbolString((0..len)
                         .into_iter()
                         .map(|_| self.gen_symbol(rng, arity, num_params))
                         .collect())
    }

    pub fn gen_symbol<R, S>(&self, rng: &mut R, arity: usize, num_params: usize) -> S
        where R: Rng,
              S: Symbol<A = EdgeAlphabet, T = f32>
    {
        S::from_iter(self.gen_symbol_value(rng),
                     (0..arity).into_iter().map(|_| self.gen_expr(rng, num_params)))

    }

    fn gen_expr<R: Rng>(&self, rng: &mut R, num_params: usize) -> Expr<f32> {
        random_expr(rng,
                    num_params,
                    self.max_expr_depth,
                    &self.expr_weighted_op,
                    &self.expr_weighted_op_max_depth)
    }

    fn gen_terminal<R: Rng>(&self, rng: &mut R) -> EdgeAlphabet {
        EdgeAlphabet::Terminal(self.edge_weighted_op.ind_sample(rng))
    }

    fn gen_nonterminal<R: Rng>(&self, rng: &mut R) -> EdgeAlphabet {
        EdgeAlphabet::NonTerminal(self.nonterminal_symbols.ind_sample(rng))
    }

    fn gen_symbol_value<R: Rng>(&self, rng: &mut R) -> EdgeAlphabet {
        if rng.gen::<ProbabilityValue>().is_probable_with(self.prob_terminal) {
            self.gen_terminal(rng)
        } else {
            self.gen_nonterminal(rng)
        }
    }

    /// Generate a simple condition like:
    ///     Arg(n) or 0.0 [>=] or [<=] constant expr
    fn gen_simple_rule_condition<R: Rng>(&self, rng: &mut R, arity: usize) -> Condition<f32> {
        let lhs = if arity > 0 {
            Expr::Arg(rng.gen_range(0, arity))
        } else {
            Expr::Const(0.0)
        };

        let rhs = random_const_expr(rng, &self.const_expr_weighted_op);

        if rng.gen::<bool>() {
            Condition::GreaterEqual(Box::new(lhs), Box::new(rhs))
        } else {
            Condition::LessEqual(Box::new(lhs), Box::new(rhs))
        }
    }

    // Generate a random rule with `symbol` and `arity` parameters.
    //
    // fn gen_rule<R: Rng>(&self, symbol: EdgeAlphabet, arity: usize) {
    // }
}

// When we need a value within (0, 1], we simply cut off the integer part.
// Max depth.

// # Genome interpretation:
//
//     * We use rule with number 0 as axiom.
//       The arguments for the axiom are fixed and passed in from the command line.
//       For simplicity, these arguments are duplicated in each Genome in `axiom_params`.
//       One could also apply genetic operations onto these.
//
//
#[derive(Clone, Debug)]
pub struct Genome {
    system: System,
}

pub struct Toolbox<N, E> {
    goal: Goal<N, E>,
    pool: Pool,
    fitness_functions: (FitnessFunction, FitnessFunction, FitnessFunction),

    weighted_op: OwnedWeightedChoice<EdgeOp>,
    weighted_var_op: OwnedWeightedChoice<VarOp>,
    weighted_mut_op: OwnedWeightedChoice<MutOp>,
    prob_mutate_elem: Probability,

    /// Arguments to the axiom rule.
    pub axiom_args: Vec<Expr<f32>>,

    /// Maximum number of iterations of the L-system  (XXX: Limit also based on generated length)
    pub iterations: usize,
}

impl<N: Clone + Default, E: Clone + Default> Toolbox<N, E> {
    pub fn new(goal: Goal<N, E>,
               pool: Pool,
               fitness_functions: (FitnessFunction, FitnessFunction, FitnessFunction),
               weighted_op: Vec<Weighted<EdgeOp>>,
               weighted_var_op: Vec<Weighted<VarOp>>,
               weighted_mut_op: Vec<Weighted<MutOp>>,
               prob_mutate_elem: Probability)
               -> Toolbox<N, E> {

        Toolbox {
            goal: goal,
            pool: pool,
            fitness_functions: fitness_functions,
            prob_mutate_elem: prob_mutate_elem,
            weighted_op: OwnedWeightedChoice::new(weighted_op),
            weighted_var_op: OwnedWeightedChoice::new(weighted_var_op),
            weighted_mut_op: OwnedWeightedChoice::new(weighted_mut_op),
            axiom_args: vec![],
            iterations: 10,
        }
    }

    pub fn mutate<R: Rng>(&self, rng: &mut R, ind: &Genome) -> Genome {
        // let mut mut_ind = Vec::with_capacity(ind.len() + 1);
        //
        // for edge_op in ind.edge_ops.iter() {
        // let new_op = if rng.gen::<ProbabilityValue>().is_probable_with(self.prob_mutate_elem) {
        // match self.weighted_mut_op.ind_sample(rng) {
        // MutOp::Copy => edge_op.clone(),
        // MutOp::Insert => {
        // Insert a new op before the current operation
        // mut_ind.push(self.generate_random_edge_operation(rng));
        // edge_op.clone()
        // }
        // MutOp::Remove => {
        // remove current operation
        // continue;
        // }
        // MutOp::ModifyOp => (self.weighted_op.ind_sample(rng), edge_op.1),
        // MutOp::ModifyParam => (edge_op.0, rng.gen::<f32>()),
        // MutOp::Replace => self.generate_random_edge_operation(rng),
        // }
        // } else {
        // edge_op.clone()
        // };
        // mut_ind.push(new_op);
        // }
        //

        Genome {
            system: System::new(),
        }
    }

    /*
    fn generate_random_edge_operation<R: Rng>(&self, rng: &mut R) -> (EdgeOp, f32) {
        generate_random_edge_operation(&self.weighted_op, rng)
    }
    */

    // There are many parameters that influence the creation of a random
    // genome:
    //
    //     - Number of rules
    //     - Arity of symbols
    //     - Axiom and Length
    //     - Length of a production rule
    //     - Number of (condition, successor) pairs per rule.
    //     - Complexity of expression in Symbol
    //
    //     - Number of Iterations.
    //
    //     There are many
    //
    pub fn random_genome<R: Rng>(&self, rng: &mut R) -> Genome {
        Genome {
            system: System::new(),
        }

    }
}

impl<N:Clone+Sync+Default,E:Clone+Sync+Default> FitnessEval<Genome, MultiObjective3<f32>> for Toolbox<N,E> {
    /// Evaluates the fitness of a Genome population.
    fn fitness(&mut self, pop: &[Genome]) -> Vec<MultiObjective3<f32>> {
        let pool = &mut self.pool;
        let goal = &self.goal;
        let axiom_args = &self.axiom_args[..];
        let iterations = self.iterations;

        let fitness_functions = self.fitness_functions;

        crossbeam::scope(|scope| {
            pool.map(scope, pop, |ind| {
                let edge_ops = ind.to_edge_ops(axiom_args, iterations);
                let g = edgeops_to_graph(&edge_ops);

                //let g = ind.to_graph();
                MultiObjective3::from((goal.apply_fitness_function(fitness_functions.0, &g),
                                       goal.apply_fitness_function(fitness_functions.1, &g),
                                       goal.apply_fitness_function(fitness_functions.2, &g)))

            })
                .collect()
        })
    }
}

impl<N: Clone + Default, E: Clone + Default> Mate<Genome> for Toolbox<N, E> {
    // p1 is potentially "better" than p2
    fn mate<R: Rng>(&mut self, rng: &mut R, p1: &Genome, p2: &Genome) -> Genome {

        match self.weighted_var_op.ind_sample(rng) {
            VarOp::Copy => p1.clone(),
            VarOp::Mutate => self.mutate(rng, p1),
            VarOp::LinearCrossover2 => {
                Genome {
                    // edge_ops: linear_2point_crossover_random(rng,
                    // &p1.edge_ops[..],
                    // &p2.edge_ops[..]),
                    system: System::new(),
                }
            }
            VarOp::UniformCrossover => {
                panic!("TODO");
            }
        }
    }
}

impl Genome {
    /// Develops the L-system into a vector of edge operations
    pub fn to_edge_ops(&self, axiom_args: &[Expr<f32>], iterations: usize) -> Vec<(EdgeOp, f32)> {
        let axiom = SymbolString(vec![Sym2::new_parametric(EdgeAlphabet::NonTerminal(0), (axiom_args[0].clone(), axiom_args[1].clone()))]);
        println!("axiom: {:?}", axiom);

        // XXX: limit lenght
        let (s, iter) = self.system.develop(axiom, iterations);
        println!("produced string: {:?}", s);
        println!("after iterations: {:?}", iter);

        let edge_ops: Vec<_> = s.0.into_iter().filter_map(|op| {
            match op.symbol() {
                &EdgeAlphabet::Terminal(ref edge_op) => {
                    if let Some(&Expr::Const(param)) = op.args().get(0) {
                        Some((edge_op.clone(), param))
                    } else {
                        println!("Invalid parameter");
                        None
                    }
                }
                _ => None
            }
        }).collect();

        return edge_ops;
    }
}

/*
#[inline]
fn generate_random_edge_operation<R: Rng>(weighted_op: &OwnedWeightedChoice<EdgeOp>,
                                          rng: &mut R)
                                          -> (EdgeOp, f32) {
    (weighted_op.ind_sample(rng), rng.gen::<f32>())
}
*/
<|endoftext|>