<commit_before>extern crate itertools;
use super::merge_policy::{MergePolicy, MergeCandidate};
use core::SegmentMeta;

pub struct LogMergePolicy;
use std::f64;

const LEVEL_LOG_SIZE: f64 = 0.75;

impl MergePolicy for LogMergePolicy {
    fn compute_merge_candidates(&self, segments: &[SegmentMeta]) -> Vec<MergeCandidate> {
        if segments.is_empty() {
            return Vec::new();
        }
        let mut size_sorted_tuples = segments.iter()
            .map(|x| x.num_docs)
            .enumerate()
            .collect::<Vec<(usize, u32)>>();

        size_sorted_tuples.sort_by(|x,y| y.cmp(x));

        let size_sorted_log_tuples: Vec<_> = size_sorted_tuples.iter()
            .map(|x| (x.0, (x.1 as f64).log2()))
            .collect();

        let (first_ind, first_score) = size_sorted_log_tuples[0];
        let mut current_max_log_size = first_score;
        let mut levels = Vec::new();
        levels.push(Vec::new());
        levels.last_mut().unwrap().push(first_ind);
        for &(ind, score) in (&size_sorted_log_tuples).iter().skip(1) {
            if score < (current_max_log_size - LEVEL_LOG_SIZE) {
                current_max_log_size = score;
                levels.push(Vec::new());
            }
            levels.last_mut().unwrap().push(ind);
        }

        let result = levels.iter()
            .map(|ind_vec| {
                MergeCandidate(ind_vec.iter()
                    .map(|&ind| segments[ind].segment_id)
                    .collect())
            })
            .collect();

        result
    }
}

impl Default for LogMergePolicy {
    fn default() -> LogMergePolicy {
        LogMergePolicy
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use indexer::merge_policy::MergePolicy;
    use core::{SegmentMeta, SegmentId};

    #[test]
    fn test_log_merge_policy_empty() {
        let y = Vec::new();
        let result_list = LogMergePolicy::default().compute_merge_candidates(&y);
        assert!(result_list.len() == 0);
    }

    #[test]
    fn test_log_merge_policy_pair() {
        let test_input = vec![SegmentMeta::new(SegmentId::generate_random(), 10),
                              SegmentMeta::new(SegmentId::generate_random(), 10)];
        let result_list = LogMergePolicy::default().compute_merge_candidates(&test_input);
        assert!(result_list.len() == 1);
    }

    #[test]
    fn test_log_merge_policy_levels() {
        // multiple levels all get merged correctly
        let test_input = vec![SegmentMeta::new(SegmentId::generate_random(), 10),
                              SegmentMeta::new(SegmentId::generate_random(), 10),
                              SegmentMeta::new(SegmentId::generate_random(), 1000),
                              SegmentMeta::new(SegmentId::generate_random(), 1000)];
        let result_list = LogMergePolicy::default().compute_merge_candidates(&test_input);
        assert!(result_list.len() == 2);
    }

    #[test]
    fn test_log_merge_policy_small() {
        // deal correctly with tiny segments?
    }

    // what else... deletions? threading/concurrency?
}
<commit_msg>Added a few more tests<commit_after>extern crate itertools;
use super::merge_policy::{MergePolicy, MergeCandidate};
use core::SegmentMeta;

pub struct LogMergePolicy;
use std::f64;

const LEVEL_LOG_SIZE: f64 = 0.75;

impl MergePolicy for LogMergePolicy {
    fn compute_merge_candidates(&self, segments: &[SegmentMeta]) -> Vec<MergeCandidate> {
        if segments.is_empty() {
            return Vec::new();
        }
        let mut size_sorted_tuples = segments.iter()
            .map(|x| x.num_docs)
            .enumerate()
            .collect::<Vec<(usize, u32)>>();

        size_sorted_tuples.sort_by(|x,y| y.cmp(x));

        let size_sorted_log_tuples: Vec<_> = size_sorted_tuples.iter()
            .map(|x| (x.0, (x.1 as f64).log2()))
            .collect();

        let (first_ind, first_score) = size_sorted_log_tuples[0];
        let mut current_max_log_size = first_score;
        let mut levels = Vec::new();
        levels.push(Vec::new());
        levels.last_mut().unwrap().push(first_ind);
        for &(ind, score) in (&size_sorted_log_tuples).iter().skip(1) {
            if score < (current_max_log_size - LEVEL_LOG_SIZE) {
                current_max_log_size = score;
                levels.push(Vec::new());
            }
            levels.last_mut().unwrap().push(ind);
        }

        let result = levels.iter()
            .map(|ind_vec| {
                MergeCandidate(ind_vec.iter()
                    .map(|&ind| segments[ind].segment_id)
                    .collect())
            })
            .collect();

        result
    }
}

impl Default for LogMergePolicy {
    fn default() -> LogMergePolicy {
        LogMergePolicy
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use indexer::merge_policy::MergePolicy;
    use core::{SegmentMeta, SegmentId};

    #[test]
    fn test_log_merge_policy_empty() {
        let y = Vec::new();
        let result_list = LogMergePolicy::default().compute_merge_candidates(&y);
        assert!(result_list.len() == 0);
    }

    #[test]
    fn test_log_merge_policy_pair() {
        let test_input = vec![SegmentMeta::new(SegmentId::generate_random(), 10),
                              SegmentMeta::new(SegmentId::generate_random(), 10)];
        let result_list = LogMergePolicy::default().compute_merge_candidates(&test_input);
        assert!(result_list.len() == 1);
    }

    #[test]
    fn test_log_merge_policy_levels() {
        // multiple levels all get merged correctly
        let test_input = vec![SegmentMeta::new(SegmentId::generate_random(), 10),
                              SegmentMeta::new(SegmentId::generate_random(), 10),
                              SegmentMeta::new(SegmentId::generate_random(), 1000),
                              SegmentMeta::new(SegmentId::generate_random(), 1000)];
        let result_list = LogMergePolicy::default().compute_merge_candidates(&test_input);
        assert!(result_list.len() == 2);
    }

    #[test]
    fn test_log_merge_policy_within_levels() {
        // multiple levels all get merged correctly
        let test_input = vec![SegmentMeta::new(SegmentId::generate_random(), 10),
                              SegmentMeta::new(SegmentId::generate_random(), 11),
                              SegmentMeta::new(SegmentId::generate_random(), 12),
                              SegmentMeta::new(SegmentId::generate_random(), 1000),
                              SegmentMeta::new(SegmentId::generate_random(), 1000)];
        let result_list = LogMergePolicy::default().compute_merge_candidates(&test_input);
        assert!(result_list.len() == 2);
    }
}
<|endoftext|>


<commit_before>use crate::{error, NetworkConfig};
/// This manages the discovery and management of peers.
///
/// Currently using discv5 for peer discovery.
///
use futures::prelude::*;
use libp2p::core::swarm::{
    ConnectedPoint, NetworkBehaviour, NetworkBehaviourAction, PollParameters,
};
use libp2p::core::{identity::Keypair, Multiaddr, PeerId, ProtocolsHandler};
use libp2p::discv5::{Discv5, Discv5Event};
use libp2p::enr::{Enr, EnrBuilder, NodeId};
use libp2p::multiaddr::Protocol;
use slog::{debug, info, o, warn};
use std::collections::HashSet;
use std::fs::File;
use std::io::prelude::*;
use std::str::FromStr;
use std::time::{Duration, Instant};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_timer::Delay;

/// Maximum seconds before searching for extra peers.
const MAX_TIME_BETWEEN_PEER_SEARCHES: u64 = 60;
/// Initial delay between peer searches.
const INITIAL_SEARCH_DELAY: u64 = 5;
/// Local ENR storage filename.
const ENR_FILENAME: &str = "enr.dat";

/// Lighthouse discovery behaviour. This provides peer management and discovery using the Discv5
/// libp2p protocol.
pub struct Discovery<TSubstream> {
    /// The peers currently connected to libp2p streams.
    connected_peers: HashSet<PeerId>,

    /// The target number of connected peers on the libp2p interface.
    max_peers: usize,

    /// The delay between peer discovery searches.
    peer_discovery_delay: Delay,

    /// Tracks the last discovery delay. The delay is doubled each round until the max
    /// time is reached.
    past_discovery_delay: u64,

    /// The TCP port for libp2p. Used to convert an updated IP address to a multiaddr. Note: This
    /// assumes that the external TCP port is the same as the internal TCP port if behind a NAT.
    //TODO: Improve NAT handling limit the above restriction
    tcp_port: u16,

    /// The discovery behaviour used to discover new peers.
    discovery: Discv5<TSubstream>,

    /// Logger for the discovery behaviour.
    log: slog::Logger,
}

impl<TSubstream> Discovery<TSubstream> {
    pub fn new(
        local_key: &Keypair,
        config: &NetworkConfig,
        log: &slog::Logger,
    ) -> error::Result<Self> {
        let log = log.new(o!("Service" => "Libp2p-Discovery"));

        // checks if current ENR matches that found on disk
        let local_enr = load_enr(local_key, config, &log)?;

        info!(log, "Local ENR: {}", local_enr.to_base64());
        debug!(log, "Local Node Id: {}", local_enr.node_id());

        let mut discovery = Discv5::new(local_enr, local_key.clone(), config.listen_address)
            .map_err(|e| format!("Discv5 service failed: {:?}", e))?;

        // Add bootnodes to routing table
        for bootnode_enr in config.boot_nodes.clone() {
            debug!(
                log,
                "Adding node to routing table: {}",
                bootnode_enr.node_id()
            );
            discovery.add_enr(bootnode_enr);
        }

        Ok(Self {
            connected_peers: HashSet::new(),
            max_peers: config.max_peers,
            peer_discovery_delay: Delay::new(Instant::now()),
            past_discovery_delay: INITIAL_SEARCH_DELAY,
            tcp_port: config.libp2p_port,
            discovery,
            log,
        })
    }

    /// Manually search for peers. This restarts the discovery round, sparking multiple rapid
    /// queries.
    pub fn discover_peers(&mut self) {
        self.past_discovery_delay = INITIAL_SEARCH_DELAY;
        self.find_peers();
    }

    /// Add an Enr to the routing table of the discovery mechanism.
    pub fn add_enr(&mut self, enr: Enr) {
        self.discovery.add_enr(enr);
    }

    /// Search for new peers using the underlying discovery mechanism.
    fn find_peers(&mut self) {
        // pick a random NodeId
        let random_node = NodeId::random();
        debug!(self.log, "Searching for peers...");
        self.discovery.find_node(random_node);

        // update the time until next discovery
        let delay = {
            if self.past_discovery_delay < MAX_TIME_BETWEEN_PEER_SEARCHES {
                self.past_discovery_delay *= 2;
                self.past_discovery_delay
            } else {
                MAX_TIME_BETWEEN_PEER_SEARCHES
            }
        };
        self.peer_discovery_delay
            .reset(Instant::now() + Duration::from_secs(delay));
    }
}

// Redirect all behaviour events to underlying discovery behaviour.
impl<TSubstream> NetworkBehaviour for Discovery<TSubstream>
where
    TSubstream: AsyncRead + AsyncWrite,
{
    type ProtocolsHandler = <Discv5<TSubstream> as NetworkBehaviour>::ProtocolsHandler;
    type OutEvent = <Discv5<TSubstream> as NetworkBehaviour>::OutEvent;

    fn new_handler(&mut self) -> Self::ProtocolsHandler {
        NetworkBehaviour::new_handler(&mut self.discovery)
    }

    fn addresses_of_peer(&mut self, peer_id: &PeerId) -> Vec<Multiaddr> {
        // Let discovery track possible known peers.
        self.discovery.addresses_of_peer(peer_id)
    }

    fn inject_connected(&mut self, peer_id: PeerId, _endpoint: ConnectedPoint) {
        self.connected_peers.insert(peer_id);
    }

    fn inject_disconnected(&mut self, peer_id: &PeerId, _endpoint: ConnectedPoint) {
        self.connected_peers.remove(peer_id);
    }

    fn inject_replaced(
        &mut self,
        _peer_id: PeerId,
        _closed: ConnectedPoint,
        _opened: ConnectedPoint,
    ) {
        // discv5 doesn't implement
    }

    fn inject_node_event(
        &mut self,
        _peer_id: PeerId,
        _event: <Self::ProtocolsHandler as ProtocolsHandler>::OutEvent,
    ) {
        // discv5 doesn't implement
    }

    fn poll(
        &mut self,
        params: &mut impl PollParameters,
    ) -> Async<
        NetworkBehaviourAction<
            <Self::ProtocolsHandler as ProtocolsHandler>::InEvent,
            Self::OutEvent,
        >,
    > {
        // search of peers if it is time
        loop {
            match self.peer_discovery_delay.poll() {
                Ok(Async::Ready(_)) => {
                    if self.connected_peers.len() < self.max_peers {
                        self.find_peers();
                    }
                }
                Ok(Async::NotReady) => break,
                Err(e) => {
                    warn!(self.log, "Discovery peer search failed: {:?}", e);
                }
            }
        }

        // Poll discovery
        loop {
            match self.discovery.poll(params) {
                Async::Ready(NetworkBehaviourAction::GenerateEvent(event)) => {
                    match event {
                        Discv5Event::Discovered(_enr) => {
                            // not concerned about FINDNODE results, rather the result of an entire
                            // query.
                        }
                        Discv5Event::SocketUpdated(socket) => {
                            info!(self.log, "Address updated"; "IP" => format!("{}",socket.ip()));
                            let mut address = Multiaddr::from(socket.ip());
                            address.push(Protocol::Tcp(self.tcp_port));
                            return Async::Ready(NetworkBehaviourAction::ReportObservedAddr {
                                address,
                            });
                        }
                        Discv5Event::FindNodeResult { closer_peers, .. } => {
                            debug!(self.log, "Discv5 query found {} peers", closer_peers.len());
                            if closer_peers.is_empty() {
                                debug!(self.log, "Discv5 random query yielded empty results");
                            }
                            for peer_id in closer_peers {
                                // if we need more peers, attempt a connection
                                if self.connected_peers.len() < self.max_peers
                                    && self.connected_peers.get(&peer_id).is_none()
                                {
                                    debug!(self.log, "Discv5: Peer discovered"; "Peer"=> format!("{:?}", peer_id));
                                    return Async::Ready(NetworkBehaviourAction::DialPeer {
                                        peer_id,
                                    });
                                }
                            }
                        }
                        _ => {}
                    }
                }
                // discv5 does not output any other NetworkBehaviourAction
                Async::Ready(_) => {}
                Async::NotReady => break,
            }
        }
        Async::NotReady
    }
}

/// Loads an ENR from file if it exists and matches the current NodeId and sequence number. If none
/// exists, generates a new one.
///
/// If an ENR exists, with the same NodeId and IP addresses, we use the disk-generated one as it's
/// ENR sequence will be equal or higher than a newly generated one.
fn load_enr(
    local_key: &Keypair,
    config: &NetworkConfig,
    log: &slog::Logger,
) -> Result<Enr, String> {
    // Build the local ENR.
    // Note: Discovery should update the ENR record's IP to the external IP as seen by the
    // majority of our peers.
    let mut local_enr = EnrBuilder::new()
        .ip(config.discovery_address.into())
        .tcp(config.libp2p_port)
        .udp(config.discovery_port)
        .build(&local_key)
        .map_err(|e| format!("Could not build Local ENR: {:?}", e))?;

    let enr_f = config.network_dir.join(ENR_FILENAME);
    if let Ok(mut enr_file) = File::open(enr_f.clone()) {
        let mut enr_string = String::new();
        match enr_file.read_to_string(&mut enr_string) {
            Err(_) => debug!(log, "Could not read ENR from file"),
            Ok(_) => {
                match Enr::from_str(&enr_string) {
                    Ok(enr) => {
                        debug!(log, "ENR found in file: {:?}", enr_f);

                        if enr.node_id() == local_enr.node_id() {
                            if enr.ip() == config.discovery_address.into()
                                && enr.tcp() == Some(config.libp2p_port)
                                && enr.udp() == Some(config.discovery_port)
                            {
                                debug!(log, "ENR loaded from file");
                                // the stored ENR has the same configuration, use it
                                return Ok(enr);
                            }

                            // same node id, different configuration - update the sequence number
                            let new_seq_no = enr.seq().checked_add(1).ok_or_else(|| "ENR sequence number on file is too large. Remove it to generate a new NodeId")?;
                            local_enr.set_seq(new_seq_no, local_key).map_err(|e| {
                                format!("Could not update ENR sequence number: {:?}", e)
                            })?;
                            debug!(log, "ENR sequence number increased to: {}", new_seq_no);
                        }
                    }
                    Err(e) => {
                        warn!(log, "ENR from file could not be decoded: {:?}", e);
                    }
                }
            }
        }
    }

    // write ENR to disk
    let _ = std::fs::create_dir_all(&config.network_dir);
    match File::create(enr_f.clone())
        .and_then(|mut f| f.write_all(&local_enr.to_base64().as_bytes()))
    {
        Ok(_) => {
            debug!(log, "ENR written to disk");
        }
        Err(e) => {
            warn!(
                log,
                "Could not write ENR to file: {:?}. Error: {}", enr_f, e
            );
        }
    }
    Ok(local_enr)
}
<commit_msg>Typo fixes<commit_after>use crate::{error, NetworkConfig};
/// This manages the discovery and management of peers.
///
/// Currently using discv5 for peer discovery.
///
use futures::prelude::*;
use libp2p::core::swarm::{
    ConnectedPoint, NetworkBehaviour, NetworkBehaviourAction, PollParameters,
};
use libp2p::core::{identity::Keypair, Multiaddr, PeerId, ProtocolsHandler};
use libp2p::discv5::{Discv5, Discv5Event};
use libp2p::enr::{Enr, EnrBuilder, NodeId};
use libp2p::multiaddr::Protocol;
use slog::{debug, info, o, warn};
use std::collections::HashSet;
use std::fs::File;
use std::io::prelude::*;
use std::str::FromStr;
use std::time::{Duration, Instant};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_timer::Delay;

/// Maximum seconds before searching for extra peers.
const MAX_TIME_BETWEEN_PEER_SEARCHES: u64 = 60;
/// Initial delay between peer searches.
const INITIAL_SEARCH_DELAY: u64 = 5;
/// Local ENR storage filename.
const ENR_FILENAME: &str = "enr.dat";

/// Lighthouse discovery behaviour. This provides peer management and discovery using the Discv5
/// libp2p protocol.
pub struct Discovery<TSubstream> {
    /// The peers currently connected to libp2p streams.
    connected_peers: HashSet<PeerId>,

    /// The target number of connected peers on the libp2p interface.
    max_peers: usize,

    /// The delay between peer discovery searches.
    peer_discovery_delay: Delay,

    /// Tracks the last discovery delay. The delay is doubled each round until the max
    /// time is reached.
    past_discovery_delay: u64,

    /// The TCP port for libp2p. Used to convert an updated IP address to a multiaddr. Note: This
    /// assumes that the external TCP port is the same as the internal TCP port if behind a NAT.
    //TODO: Improve NAT handling limit the above restriction
    tcp_port: u16,

    /// The discovery behaviour used to discover new peers.
    discovery: Discv5<TSubstream>,

    /// Logger for the discovery behaviour.
    log: slog::Logger,
}

impl<TSubstream> Discovery<TSubstream> {
    pub fn new(
        local_key: &Keypair,
        config: &NetworkConfig,
        log: &slog::Logger,
    ) -> error::Result<Self> {
        let log = log.new(o!("Service" => "Libp2p-Discovery"));

        // checks if current ENR matches that found on disk
        let local_enr = load_enr(local_key, config, &log)?;

        info!(log, "Local ENR: {}", local_enr.to_base64());
        debug!(log, "Local Node Id: {}", local_enr.node_id());

        let mut discovery = Discv5::new(local_enr, local_key.clone(), config.listen_address)
            .map_err(|e| format!("Discv5 service failed: {:?}", e))?;

        // Add bootnodes to routing table
        for bootnode_enr in config.boot_nodes.clone() {
            debug!(
                log,
                "Adding node to routing table: {}",
                bootnode_enr.node_id()
            );
            discovery.add_enr(bootnode_enr);
        }

        Ok(Self {
            connected_peers: HashSet::new(),
            max_peers: config.max_peers,
            peer_discovery_delay: Delay::new(Instant::now()),
            past_discovery_delay: INITIAL_SEARCH_DELAY,
            tcp_port: config.libp2p_port,
            discovery,
            log,
        })
    }

    /// Manually search for peers. This restarts the discovery round, sparking multiple rapid
    /// queries.
    pub fn discover_peers(&mut self) {
        self.past_discovery_delay = INITIAL_SEARCH_DELAY;
        self.find_peers();
    }

    /// Add an Enr to the routing table of the discovery mechanism.
    pub fn add_enr(&mut self, enr: Enr) {
        self.discovery.add_enr(enr);
    }

    /// Search for new peers using the underlying discovery mechanism.
    fn find_peers(&mut self) {
        // pick a random NodeId
        let random_node = NodeId::random();
        debug!(self.log, "Searching for peers...");
        self.discovery.find_node(random_node);

        // update the time until next discovery
        let delay = {
            if self.past_discovery_delay < MAX_TIME_BETWEEN_PEER_SEARCHES {
                self.past_discovery_delay *= 2;
                self.past_discovery_delay
            } else {
                MAX_TIME_BETWEEN_PEER_SEARCHES
            }
        };
        self.peer_discovery_delay
            .reset(Instant::now() + Duration::from_secs(delay));
    }
}

// Redirect all behaviour events to underlying discovery behaviour.
impl<TSubstream> NetworkBehaviour for Discovery<TSubstream>
where
    TSubstream: AsyncRead + AsyncWrite,
{
    type ProtocolsHandler = <Discv5<TSubstream> as NetworkBehaviour>::ProtocolsHandler;
    type OutEvent = <Discv5<TSubstream> as NetworkBehaviour>::OutEvent;

    fn new_handler(&mut self) -> Self::ProtocolsHandler {
        NetworkBehaviour::new_handler(&mut self.discovery)
    }

    fn addresses_of_peer(&mut self, peer_id: &PeerId) -> Vec<Multiaddr> {
        // Let discovery track possible known peers.
        self.discovery.addresses_of_peer(peer_id)
    }

    fn inject_connected(&mut self, peer_id: PeerId, _endpoint: ConnectedPoint) {
        self.connected_peers.insert(peer_id);
    }

    fn inject_disconnected(&mut self, peer_id: &PeerId, _endpoint: ConnectedPoint) {
        self.connected_peers.remove(peer_id);
    }

    fn inject_replaced(
        &mut self,
        _peer_id: PeerId,
        _closed: ConnectedPoint,
        _opened: ConnectedPoint,
    ) {
        // discv5 doesn't implement
    }

    fn inject_node_event(
        &mut self,
        _peer_id: PeerId,
        _event: <Self::ProtocolsHandler as ProtocolsHandler>::OutEvent,
    ) {
        // discv5 doesn't implement
    }

    fn poll(
        &mut self,
        params: &mut impl PollParameters,
    ) -> Async<
        NetworkBehaviourAction<
            <Self::ProtocolsHandler as ProtocolsHandler>::InEvent,
            Self::OutEvent,
        >,
    > {
        // search for peers if it is time
        loop {
            match self.peer_discovery_delay.poll() {
                Ok(Async::Ready(_)) => {
                    if self.connected_peers.len() < self.max_peers {
                        self.find_peers();
                    }
                }
                Ok(Async::NotReady) => break,
                Err(e) => {
                    warn!(self.log, "Discovery peer search failed: {:?}", e);
                }
            }
        }

        // Poll discovery
        loop {
            match self.discovery.poll(params) {
                Async::Ready(NetworkBehaviourAction::GenerateEvent(event)) => {
                    match event {
                        Discv5Event::Discovered(_enr) => {
                            // not concerned about FINDNODE results, rather the result of an entire
                            // query.
                        }
                        Discv5Event::SocketUpdated(socket) => {
                            info!(self.log, "Address updated"; "IP" => format!("{}",socket.ip()));
                            let mut address = Multiaddr::from(socket.ip());
                            address.push(Protocol::Tcp(self.tcp_port));
                            return Async::Ready(NetworkBehaviourAction::ReportObservedAddr {
                                address,
                            });
                        }
                        Discv5Event::FindNodeResult { closer_peers, .. } => {
                            debug!(self.log, "Discv5 query found {} peers", closer_peers.len());
                            if closer_peers.is_empty() {
                                debug!(self.log, "Discv5 random query yielded empty results");
                            }
                            for peer_id in closer_peers {
                                // if we need more peers, attempt a connection
                                if self.connected_peers.len() < self.max_peers
                                    && self.connected_peers.get(&peer_id).is_none()
                                {
                                    debug!(self.log, "Discv5: Peer discovered"; "Peer"=> format!("{:?}", peer_id));
                                    return Async::Ready(NetworkBehaviourAction::DialPeer {
                                        peer_id,
                                    });
                                }
                            }
                        }
                        _ => {}
                    }
                }
                // discv5 does not output any other NetworkBehaviourAction
                Async::Ready(_) => {}
                Async::NotReady => break,
            }
        }
        Async::NotReady
    }
}

/// Loads an ENR from file if it exists and matches the current NodeId and sequence number. If none
/// exists, generates a new one.
///
/// If an ENR exists, with the same NodeId and IP address, we use the disk-generated one as its
/// ENR sequence will be equal or higher than a newly generated one.
fn load_enr(
    local_key: &Keypair,
    config: &NetworkConfig,
    log: &slog::Logger,
) -> Result<Enr, String> {
    // Build the local ENR.
    // Note: Discovery should update the ENR record's IP to the external IP as seen by the
    // majority of our peers.
    let mut local_enr = EnrBuilder::new()
        .ip(config.discovery_address.into())
        .tcp(config.libp2p_port)
        .udp(config.discovery_port)
        .build(&local_key)
        .map_err(|e| format!("Could not build Local ENR: {:?}", e))?;

    let enr_f = config.network_dir.join(ENR_FILENAME);
    if let Ok(mut enr_file) = File::open(enr_f.clone()) {
        let mut enr_string = String::new();
        match enr_file.read_to_string(&mut enr_string) {
            Err(_) => debug!(log, "Could not read ENR from file"),
            Ok(_) => {
                match Enr::from_str(&enr_string) {
                    Ok(enr) => {
                        debug!(log, "ENR found in file: {:?}", enr_f);

                        if enr.node_id() == local_enr.node_id() {
                            if enr.ip() == config.discovery_address.into()
                                && enr.tcp() == Some(config.libp2p_port)
                                && enr.udp() == Some(config.discovery_port)
                            {
                                debug!(log, "ENR loaded from file");
                                // the stored ENR has the same configuration, use it
                                return Ok(enr);
                            }

                            // same node id, different configuration - update the sequence number
                            let new_seq_no = enr.seq().checked_add(1).ok_or_else(|| "ENR sequence number on file is too large. Remove it to generate a new NodeId")?;
                            local_enr.set_seq(new_seq_no, local_key).map_err(|e| {
                                format!("Could not update ENR sequence number: {:?}", e)
                            })?;
                            debug!(log, "ENR sequence number increased to: {}", new_seq_no);
                        }
                    }
                    Err(e) => {
                        warn!(log, "ENR from file could not be decoded: {:?}", e);
                    }
                }
            }
        }
    }

    // write ENR to disk
    let _ = std::fs::create_dir_all(&config.network_dir);
    match File::create(enr_f.clone())
        .and_then(|mut f| f.write_all(&local_enr.to_base64().as_bytes()))
    {
        Ok(_) => {
            debug!(log, "ENR written to disk");
        }
        Err(e) => {
            warn!(
                log,
                "Could not write ENR to file: {:?}. Error: {}", enr_f, e
            );
        }
    }
    Ok(local_enr)
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

#[feature(globs)];
#[feature(macro_rules)];

#[link(name = "benchmark", vers = "alpha", author = "dwrensha")];

#[crate_type = "bin"];

extern mod capnprust;

pub mod common;

pub mod carsales_capnp;
pub mod carsales;

pub mod catrank_capnp;
pub mod catrank;

pub mod eval_capnp;
pub mod eval;

mod Uncompressed {
    use capnprust;
    use std;

    pub fn write<T : std::rt::io::Writer>(writer: &mut T,
                                          message: &capnprust::message::MessageBuilder) {
        capnprust::serialize::writeMessage(writer, message);
    }

    pub fn newReader<U : std::rt::io::Reader, T>(
        inputStream : &mut U,
        options : capnprust::message::ReaderOptions,
        cont : &fn(v : &mut capnprust::message::MessageReader) -> T) -> T {
        capnprust::serialize::InputStreamMessageReader::new(inputStream, options, cont)
    }
}

mod Packed {
    use capnprust;
    use std;

    pub fn write<T : std::rt::io::Writer>(writer: &mut T,
                                          message: &capnprust::message::MessageBuilder) {
        capnprust::serialize_packed::writePackedMessage(writer, message);
    }

    pub fn newReader<U : std::rt::io::Reader, T>(
        inputStream : &mut U,
        options : capnprust::message::ReaderOptions,
        cont : &fn(v : &mut capnprust::message::MessageReader) -> T) -> T {
        capnprust::serialize::InputStreamMessageReader::new(
            &mut capnprust::serialize_packed::PackedInputStream{inner : inputStream},
            options, cont)
    }
}

macro_rules! passByObject(
    ( $testcase:ident, $iters:expr ) => ({
            let mut rng = common::FastRand::new();
            for _ in range(0, $iters) {
                let mut messageReq = capnprust::message::MessageBuilder::new_default();
                let mut messageRes = capnprust::message::MessageBuilder::new_default();

                let request = messageReq.initRoot::<$testcase::RequestBuilder>();
                let response = messageRes.initRoot::<$testcase::ResponseBuilder>();
                let expected = $testcase::setupRequest(&mut rng, request);

                do request.asReader |requestReader| {
                    $testcase::handleRequest(requestReader, response);
                }

                do response.asReader |responseReader| {
                    if (! $testcase::checkResponse(responseReader, expected)) {
                        println("Incorrect response.");
                    }
                }
            }
        });
    )

macro_rules! passByBytes(
    ( $testcase:ident, $compression:ident, $iters:expr ) => ({
            let mut rng = common::FastRand::new();
            for _ in range(0, $iters) {
                let mut messageReq = capnprust::message::MessageBuilder::new_default();
                let mut messageRes = capnprust::message::MessageBuilder::new_default();

                let request = messageReq.initRoot::<$testcase::RequestBuilder>();
                let response = messageRes.initRoot::<$testcase::ResponseBuilder>();
                let expected = $testcase::setupRequest(&mut rng, request);

                let requestBytes = do std::rt::io::mem::with_mem_writer |writer| {
                    $compression::write(writer, messageReq)
                };

                do $compression::newReader(
                      &mut std::rt::io::mem::BufReader::new(requestBytes),
                      capnprust::message::DEFAULT_READER_OPTIONS) |requestReader| {
                    let requestReader = $testcase::newRequestReader(requestReader.getRoot());
                    $testcase::handleRequest(requestReader, response);
                }

                let responseBytes = do std::rt::io::mem::with_mem_writer |writer| {
                    $compression::write(writer, messageRes);
                };

                do $compression::newReader(
                    &mut std::rt::io::mem::BufReader::new(responseBytes),
                    capnprust::message::DEFAULT_READER_OPTIONS) |responseReader| {
                    let responseReader = $testcase::newResponseReader(responseReader.getRoot());
                    if (! $testcase::checkResponse(responseReader, expected)) {
                        println("Incorrect response.");
                    }
                }
            }
        });
    )

macro_rules! server(
    ( $testcase:ident, $compression:ident, $iters:expr, $input:expr, $output:expr) => ({

        });
    )

macro_rules! syncClient(
    ( $testcase:ident, $iters:expr) => ({
            let mut rng = common::FastRand::new();
            for _ in range(0, $iters) {
                let mut messageReq = capnprust::message::MessageBuilder::new_default();
                let mut messageRes = capnprust::message::MessageBuilder::new_default();

                let request = messageReq.initRoot::<$testcase::RequestBuilder>();
                let _response = messageRes.initRoot::<$testcase::ResponseBuilder>();
                let _expected = $testcase::setupRequest(&mut rng, request);
                fail!("unimplemented");
            }
        });
    )

macro_rules! passByPipe(
    ( $testcase:ident, $compression:ident, $iters:expr) => ({
            use std::rt::io::process;

            // get a rustc crash if we put this in line below
            let io = ~[process::CreatePipe(true, false), // stdin
                       process::CreatePipe(false, true), // stdout
                       process::Ignored];

            let mut args = std::os::args();
            args[2] = ~"client";

            let config = process::ProcessConfig {
                program: "./benchmark/benchmark",
                args: args.slice(1, args.len()),
                env : None,
                cwd: None,
                io : io
            };
            match process::Process::new(config) {
                Some(ref mut p) => {
                    server!($testcase, $compression, $iters, p.io[1], p.io[0]);
                    println!("{}", p.wait());
                }
                None => {
                    println("bummer");
                }
            }
        });
    )

macro_rules! doTestcase(
    ( $testcase:ident, $mode:expr, $reuse:expr, $compression:ident, $iters:expr ) => ({
            match $mode {
                ~"object" => passByObject!($testcase, $iters),
                ~"bytes" => passByBytes!($testcase, $compression, $iters),
                ~"client" => syncClient!($testcase, $iters),
                ~"pipe" => passByPipe!($testcase, $compression, $iters),
                s => fail!("unrecognized mode: {}", s)
            }
        });
    )

macro_rules! doTestcase1(
    ( $testcase:expr, $mode:expr, $reuse:expr, $compression:ident, $iters:expr) => ({
            match $testcase {
                ~"carsales" => doTestcase!(carsales, $mode, $reuse, $compression, $iters),
                ~"catrank" => doTestcase!(catrank, $mode, $reuse, $compression, $iters),
                ~"eval" => doTestcase!(eval, $mode, $reuse, $compression, $iters),
                s => fail!("unrecognized test case: {}", s)
            }
        });
    )

pub fn main () {

    let args = std::os::args();

    if (args.len() != 6) {
        println!("USAGE: {} CASE MODE REUSE COMPRESSION ITERATION_COUNT", args[0]);
        std::os::set_exit_status(1);
        return;
    }

    let iters = match from_str::<u64>(args[5]) {
        Some (n) => n,
        None => {
            println!("Could not parse a u64 from: {}", args[5]);
            std::os::set_exit_status(1);
            return;
        }
    };

    match args[4] {
        ~"none" => doTestcase1!(args[1], args[2],  args[3], Uncompressed, iters),
        ~"packed" => doTestcase1!(args[1], args[2], args[3], Packed, iters),
        s => fail!("unrecognized compression: {}", s)
    }

}
<commit_msg>passByPipe appears to work. it's quite slow for packed messages<commit_after>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

#[feature(globs)];
#[feature(macro_rules)];

#[link(name = "benchmark", vers = "alpha", author = "dwrensha")];

#[crate_type = "bin"];

extern mod capnprust;

pub mod common;

pub mod carsales_capnp;
pub mod carsales;

pub mod catrank_capnp;
pub mod catrank;

pub mod eval_capnp;
pub mod eval;

mod Uncompressed {
    use capnprust;
    use std;

    pub fn write<T : std::rt::io::Writer>(writer: &mut T,
                                          message: &capnprust::message::MessageBuilder) {
        capnprust::serialize::writeMessage(writer, message);
    }

    pub fn newReader<U : std::rt::io::Reader, T>(
        inputStream : &mut U,
        options : capnprust::message::ReaderOptions,
        cont : &fn(v : &mut capnprust::message::MessageReader) -> T) -> T {
        capnprust::serialize::InputStreamMessageReader::new(inputStream, options, cont)
    }
}

mod Packed {
    use capnprust;
    use std;

    pub fn write<T : std::rt::io::Writer>(writer: &mut T,
                                          message: &capnprust::message::MessageBuilder) {
        capnprust::serialize_packed::writePackedMessage(writer, message);
    }

    pub fn newReader<U : std::rt::io::Reader, T>(
        inputStream : &mut U,
        options : capnprust::message::ReaderOptions,
        cont : &fn(v : &mut capnprust::message::MessageReader) -> T) -> T {
        capnprust::serialize::InputStreamMessageReader::new(
            &mut capnprust::serialize_packed::PackedInputStream{inner : inputStream},
            options, cont)
    }
}

macro_rules! passByObject(
    ( $testcase:ident, $iters:expr ) => ({
            let mut rng = common::FastRand::new();
            for _ in range(0, $iters) {
                let mut messageReq = capnprust::message::MessageBuilder::new_default();
                let mut messageRes = capnprust::message::MessageBuilder::new_default();

                let request = messageReq.initRoot::<$testcase::RequestBuilder>();
                let response = messageRes.initRoot::<$testcase::ResponseBuilder>();
                let expected = $testcase::setupRequest(&mut rng, request);

                do request.asReader |requestReader| {
                    $testcase::handleRequest(requestReader, response);
                }

                do response.asReader |responseReader| {
                    if (! $testcase::checkResponse(responseReader, expected)) {
                        fail!("Incorrect response.");
                    }
                }
            }
        });
    )

macro_rules! passByBytes(
    ( $testcase:ident, $compression:ident, $iters:expr ) => ({
            let mut rng = common::FastRand::new();
            for _ in range(0, $iters) {
                let mut messageReq = capnprust::message::MessageBuilder::new_default();
                let mut messageRes = capnprust::message::MessageBuilder::new_default();

                let request = messageReq.initRoot::<$testcase::RequestBuilder>();
                let response = messageRes.initRoot::<$testcase::ResponseBuilder>();
                let expected = $testcase::setupRequest(&mut rng, request);

                let requestBytes = do std::rt::io::mem::with_mem_writer |writer| {
                    $compression::write(writer, messageReq)
                };

                do $compression::newReader(
                      &mut std::rt::io::mem::BufReader::new(requestBytes),
                      capnprust::message::DEFAULT_READER_OPTIONS) |requestReader| {
                    let requestReader = $testcase::newRequestReader(requestReader.getRoot());
                    $testcase::handleRequest(requestReader, response);
                }

                let responseBytes = do std::rt::io::mem::with_mem_writer |writer| {
                    $compression::write(writer, messageRes);
                };

                do $compression::newReader(
                    &mut std::rt::io::mem::BufReader::new(responseBytes),
                    capnprust::message::DEFAULT_READER_OPTIONS) |responseReader| {
                    let responseReader = $testcase::newResponseReader(responseReader.getRoot());
                    if (! $testcase::checkResponse(responseReader, expected)) {
                        fail!("Incorrect response.");
                    }
                }
            }
        });
    )

macro_rules! server(
    ( $testcase:ident, $compression:ident, $iters:expr, $input:expr, $output:expr) => ({
            for _ in range(0, $iters) {
                let mut messageRes = capnprust::message::MessageBuilder::new_default();
                let response = messageRes.initRoot::<$testcase::ResponseBuilder>();
                do $compression::newReader(
                    &mut $input,
                    capnprust::message::DEFAULT_READER_OPTIONS) |requestReader| {
                    let requestReader = $testcase::newRequestReader(requestReader.getRoot());
                    $testcase::handleRequest(requestReader, response);
                }
                $compression::write(&mut $output, messageRes);
            }
        });
    )

macro_rules! syncClient(
    ( $testcase:ident, $compression:ident, $iters:expr) => ({
            let mut outStream = std::rt::io::stdout();
            let mut inStream = std::rt::io::stdin();
            let mut rng = common::FastRand::new();
            for _ in range(0, $iters) {
                let mut messageReq = capnprust::message::MessageBuilder::new_default();
                let request = messageReq.initRoot::<$testcase::RequestBuilder>();

                let expected = $testcase::setupRequest(&mut rng, request);
                $compression::write(&mut outStream, messageReq);

                do $compression::newReader(
                    &mut inStream,
                    capnprust::message::DEFAULT_READER_OPTIONS) |responseReader| {
                    let responseReader = $testcase::newResponseReader(responseReader.getRoot());
                    assert!($testcase::checkResponse(responseReader, expected));
                }

            }
        });
    )

macro_rules! passByPipe(
    ( $testcase:ident, $compression:ident, $iters:expr) => ({
            use std::rt::io::process;

            // get a rustc crash if we put this in line below
            let io = ~[process::CreatePipe(true, false), // stdin
                       process::CreatePipe(false, true), // stdout
                       process::Ignored];

            let mut args = std::os::args();
            args[2] = ~"client";

            let config = process::ProcessConfig {
                program: "./benchmark/benchmark",
                args: args.slice(1, args.len()),
                env : None,
                cwd: None,
                io : io
            };
            match process::Process::new(config) {
                Some(ref mut p) => {
                    server!($testcase, $compression, $iters, p.io[1], p.io[0]);
                    println!("{}", p.wait());
                }
                None => {
                    println("bummer");
                }
            }
        });
    )

macro_rules! doTestcase(
    ( $testcase:ident, $mode:expr, $reuse:expr, $compression:ident, $iters:expr ) => ({
            match $mode {
                ~"object" => passByObject!($testcase, $iters),
                ~"bytes" => passByBytes!($testcase, $compression, $iters),
                ~"client" => syncClient!($testcase, $compression, $iters),
                ~"pipe" => passByPipe!($testcase, $compression, $iters),
                s => fail!("unrecognized mode: {}", s)
            }
        });
    )

macro_rules! doTestcase1(
    ( $testcase:expr, $mode:expr, $reuse:expr, $compression:ident, $iters:expr) => ({
            match $testcase {
                ~"carsales" => doTestcase!(carsales, $mode, $reuse, $compression, $iters),
                ~"catrank" => doTestcase!(catrank, $mode, $reuse, $compression, $iters),
                ~"eval" => doTestcase!(eval, $mode, $reuse, $compression, $iters),
                s => fail!("unrecognized test case: {}", s)
            }
        });
    )

pub fn main () {

    let args = std::os::args();

    if (args.len() != 6) {
        println!("USAGE: {} CASE MODE REUSE COMPRESSION ITERATION_COUNT", args[0]);
        std::os::set_exit_status(1);
        return;
    }

    let iters = match from_str::<u64>(args[5]) {
        Some (n) => n,
        None => {
            println!("Could not parse a u64 from: {}", args[5]);
            std::os::set_exit_status(1);
            return;
        }
    };

    match args[4] {
        ~"none" => doTestcase1!(args[1], args[2],  args[3], Uncompressed, iters),
        ~"packed" => doTestcase1!(args[1], args[2], args[3], Packed, iters),
        s => fail!("unrecognized compression: {}", s)
    }

}
<|endoftext|>


<commit_before><commit_msg>Add Path test methods and tests<commit_after><|endoftext|>


<commit_before><commit_msg>Bounds checking for the development rasterizer<commit_after><|endoftext|>


<commit_before><commit_msg>mach.constants: update with some new osx and ios versions<commit_after><|endoftext|>


<commit_before><commit_msg>Will message can be any binary data<commit_after><|endoftext|>


<commit_before>import ast::{crate, expr_, mac_invoc,
                     mac_aq, mac_var};
import fold::*;
import visit::*;
import ext::base::*;
import ext::build::*;
import parse::parser;
import parse::parser::parse_from_source_str;

import print::*;
import io::*;

import codemap::span;

type aq_ctxt = @{lo: uint,
                 mut gather: [{lo: uint, hi: uint,
                                   e: @ast::expr,
                                   constr: str}]};
enum fragment {
    from_expr(@ast::expr),
    from_ty(@ast::ty)
}

iface qq_helper {
    fn span() -> span;
    fn visit(aq_ctxt, vt<aq_ctxt>);
    fn extract_mac() -> option<ast::mac_>;
    fn mk_parse_fn(ext_ctxt,span) -> @ast::expr;
    fn get_fold_fn() -> str;
}

impl of qq_helper for @ast::crate {
    fn span() -> span {self.span}
    fn visit(cx: aq_ctxt, v: vt<aq_ctxt>) {visit_crate(*self, cx, v);}
    fn extract_mac() -> option<ast::mac_> {fail}
    fn mk_parse_fn(cx: ext_ctxt, sp: span) -> @ast::expr {
        mk_path(cx, sp, ["syntax", "ext", "qquote", "parse_crate"])
    }
    fn get_fold_fn() -> str {"fold_crate"}
}
impl of qq_helper for @ast::expr {
    fn span() -> span {self.span}
    fn visit(cx: aq_ctxt, v: vt<aq_ctxt>) {visit_expr(self, cx, v);}
    fn extract_mac() -> option<ast::mac_> {
        alt (self.node) {
          ast::expr_mac({node: mac, _}) {some(mac)}
          _ {none}
        }
    }
    fn mk_parse_fn(cx: ext_ctxt, sp: span) -> @ast::expr {
        mk_path(cx, sp, ["syntax", "parse", "parser", "parse_expr"])
    }
    fn get_fold_fn() -> str {"fold_expr"}
}
impl of qq_helper for @ast::ty {
    fn span() -> span {self.span}
    fn visit(cx: aq_ctxt, v: vt<aq_ctxt>) {visit_ty(self, cx, v);}
    fn extract_mac() -> option<ast::mac_> {
        alt (self.node) {
          ast::ty_mac({node: mac, _}) {some(mac)}
          _ {none}
        }
    }
    fn mk_parse_fn(cx: ext_ctxt, sp: span) -> @ast::expr {
        mk_path(cx, sp, ["syntax", "ext", "qquote", "parse_ty"])
    }
    fn get_fold_fn() -> str {"fold_ty"}
}
impl of qq_helper for @ast::item {
    fn span() -> span {self.span}
    fn visit(cx: aq_ctxt, v: vt<aq_ctxt>) {visit_item(self, cx, v);}
    fn extract_mac() -> option<ast::mac_> {fail}
    fn mk_parse_fn(cx: ext_ctxt, sp: span) -> @ast::expr {
        mk_path(cx, sp, ["syntax", "ext", "qquote", "parse_item"])
    }
    fn get_fold_fn() -> str {"fold_item"}
}
impl of qq_helper for @ast::stmt {
    fn span() -> span {self.span}
    fn visit(cx: aq_ctxt, v: vt<aq_ctxt>) {visit_stmt(self, cx, v);}
    fn extract_mac() -> option<ast::mac_> {fail}
    fn mk_parse_fn(cx: ext_ctxt, sp: span) -> @ast::expr {
        mk_path(cx, sp, ["syntax", "ext", "qquote", "parse_stmt"])
    }
    fn get_fold_fn() -> str {"fold_stmt"}
}
impl of qq_helper for @ast::pat {
    fn span() -> span {self.span}
    fn visit(cx: aq_ctxt, v: vt<aq_ctxt>) {visit_pat(self, cx, v);}
    fn extract_mac() -> option<ast::mac_> {fail}
    fn mk_parse_fn(cx: ext_ctxt, sp: span) -> @ast::expr {
        mk_path(cx, sp, ["syntax", "parse", "parser", "parse_pat"])
    }
    fn get_fold_fn() -> str {"fold_pat"}
}

fn gather_anti_quotes<N: qq_helper>(lo: uint, node: N) -> aq_ctxt
{
    let v = @{visit_expr: {|node, &&cx, v|
                  visit_aq(node, "from_expr", cx, v)},
              visit_ty: {|node, &&cx, v|
                  visit_aq(node, "from_ty", cx, v)}
              with *default_visitor()};
    let cx = @{lo:lo, mut gather: []};
    node.visit(cx, mk_vt(v));
    // FIXME: Maybe this is an overkill (merge_sort), it might be better
    //   to just keep the gather array in sorted order ... (Issue #2250)
    cx.gather = std::sort::merge_sort({|a,b| a.lo < b.lo}, copy cx.gather);
    ret cx;
}

fn visit_aq<T:qq_helper>(node: T, constr: str, &&cx: aq_ctxt, v: vt<aq_ctxt>)
{
    alt (node.extract_mac()) {
      some(mac_aq(sp, e)) {
        cx.gather += [{lo: sp.lo - cx.lo, hi: sp.hi - cx.lo,
                       e: e, constr: constr}];
      }
      _ {node.visit(cx, v);}
    }
}

fn is_space(c: char) -> bool {
    parse::lexer::is_whitespace(c)
}

fn expand_ast(ecx: ext_ctxt, _sp: span,
              arg: ast::mac_arg, body: ast::mac_body)
    -> @ast::expr
{
    let mut what = "expr";
    option::iter(arg) {|arg|
        let args: [@ast::expr] =
            alt arg.node {
              ast::expr_vec(elts, _) { elts }
              _ {
                ecx.span_fatal
                    (_sp, "#ast requires arguments of the form `[...]`.")
              }
            };
        if vec::len::<@ast::expr>(args) != 1u {
            ecx.span_fatal(_sp, "#ast requires exactly one arg");
        }
        alt (args[0].node) {
          ast::expr_path(@{idents: id, _}) if vec::len(id) == 1u
              {what = id[0]}
          _ {ecx.span_fatal(args[0].span, "expected an identifier");}
        }
    }
    let body = get_mac_body(ecx,_sp,body);

    ret alt what {
      "crate" {finish(ecx, body, parse_crate)}
      "expr" {finish(ecx, body, parser::parse_expr)}
      "ty" {finish(ecx, body, parse_ty)}
      "item" {finish(ecx, body, parse_item)}
      "stmt" {finish(ecx, body, parse_stmt)}
      "pat" {finish(ecx, body, parser::parse_pat)}
      _ {ecx.span_fatal(_sp, "unsupported ast type")}
    };
}

fn parse_crate(p: parser) -> @ast::crate {
    parser::parse_crate_mod(p, [])
}

fn parse_ty(p: parser) -> @ast::ty {
    parser::parse_ty(p, false)
}

fn parse_stmt(p: parser) -> @ast::stmt {
    parser::parse_stmt(p, [])
}

fn parse_item(p: parser) -> @ast::item {
    alt parser::parse_item(p, [], ast::public) {
      some(item) { item }
      none       { fail "parse_item: parsing an item failed"; }
    }
}

fn finish<T: qq_helper>
    (ecx: ext_ctxt, body: ast::mac_body_, f: fn (p: parser) -> T)
    -> @ast::expr
{
    let cm = ecx.codemap();
    let str = @codemap::span_to_snippet(body.span, cm);
    #debug["qquote--str==%?", str];
    let fname = codemap::mk_substr_filename(cm, body.span);
    let node = parse_from_source_str
        (f, fname, codemap::fss_internal(body.span), str,
         ecx.cfg(), ecx.parse_sess());
    let loc = codemap::lookup_char_pos(cm, body.span.lo);

    let sp = node.span();
    let qcx = gather_anti_quotes(sp.lo, node);
    let cx = qcx;

    uint::range(1u, vec::len(cx.gather)) {|i|
        assert cx.gather[i-1u].lo < cx.gather[i].lo;
        // ^^ check that the vector is sorted
        assert cx.gather[i-1u].hi <= cx.gather[i].lo;
        // ^^ check that the spans are non-overlapping
    }

    let mut str2 = "";
    enum state {active, skip(uint), blank};
    let mut state = active;
    let mut i = 0u, j = 0u;
    let g_len = vec::len(cx.gather);
    str::chars_iter(*str) {|ch|
        if (j < g_len && i == cx.gather[j].lo) {
            assert ch == '$';
            let repl = #fmt("$%u ", j);
            state = skip(str::char_len(repl));
            str2 += repl;
        }
        alt state {
          active {str::push_char(str2, ch);}
          skip(1u) {state = blank;}
          skip(sk) {state = skip (sk-1u);}
          blank if is_space(ch) {str::push_char(str2, ch);}
          blank {str::push_char(str2, ' ');}
        }
        i += 1u;
        if (j < g_len && i == cx.gather[j].hi) {
            assert ch == ')';
            state = active;
            j += 1u;
        }
    }

    let cx = ecx;

    let cfg_call = {||
        mk_call_(cx, sp, mk_access(cx, sp, ["ext_cx"], "cfg"), [])
    };

    let parse_sess_call = {||
        mk_call_(cx, sp, mk_access(cx, sp, ["ext_cx"], "parse_sess"), [])
    };

    let pcall = mk_call(cx,sp,
                       ["syntax", "parse", "parser",
                        "parse_from_source_str"],
                       [node.mk_parse_fn(cx,sp),
                        mk_str(cx,sp, fname),
                        mk_call(cx,sp,
                                ["syntax","ext","qquote", "mk_file_substr"],
                                [mk_str(cx,sp, loc.file.name),
                                 mk_uint(cx,sp, loc.line),
                                 mk_uint(cx,sp, loc.col)]),
                        mk_unary(cx,sp, ast::box(ast::m_imm),
                                 mk_str(cx,sp, str2)),
                        cfg_call(),
                        parse_sess_call()]
                      );
    let mut rcall = pcall;
    if (g_len > 0u) {
        let gather = copy qcx.gather;
        rcall = mk_call(cx,sp,
                        ["syntax", "ext", "qquote", "replace"],
                        [pcall,
                         mk_vec_e(cx,sp, vec::map(gather) {|g|
                             mk_call(cx,sp,
                                     ["syntax", "ext", "qquote", g.constr],
                                     [g.e])}),
                         mk_path(cx,sp,
                                 ["syntax", "ext", "qquote",
                                  node.get_fold_fn()])]);
    }
    ret rcall;
}

fn replace<T>(node: T, repls: [fragment], ff: fn (ast_fold, T) -> T)
    -> T
{
    let aft = default_ast_fold();
    let f_pre = {fold_expr: bind replace_expr(repls, _, _, _,
                                              aft.fold_expr),
                 fold_ty: bind replace_ty(repls, _, _, _,
                                          aft.fold_ty)
                 with *aft};
    ret ff(make_fold(f_pre), node);
}
fn fold_crate(f: ast_fold, &&n: @ast::crate) -> @ast::crate {
    @f.fold_crate(*n)
}
fn fold_expr(f: ast_fold, &&n: @ast::expr) -> @ast::expr {f.fold_expr(n)}
fn fold_ty(f: ast_fold, &&n: @ast::ty) -> @ast::ty {f.fold_ty(n)}
fn fold_item(f: ast_fold, &&n: @ast::item) -> @ast::item {f.fold_item(n)}
fn fold_stmt(f: ast_fold, &&n: @ast::stmt) -> @ast::stmt {f.fold_stmt(n)}
fn fold_pat(f: ast_fold, &&n: @ast::pat) -> @ast::pat {f.fold_pat(n)}

fn replace_expr(repls: [fragment],
                e: ast::expr_, s: span, fld: ast_fold,
                orig: fn@(ast::expr_, span, ast_fold)->(ast::expr_, span))
    -> (ast::expr_, span)
{
    alt e {
      ast::expr_mac({node: mac_var(i), _}) {
        alt (repls[i]) {
          from_expr(r) {(r.node, r.span)}
          _ {fail /* fixme error message */}}}
      _ {orig(e,s,fld)}
    }
}

fn replace_ty(repls: [fragment],
                e: ast::ty_, s: span, fld: ast_fold,
                orig: fn@(ast::ty_, span, ast_fold)->(ast::ty_, span))
    -> (ast::ty_, span)
{
    alt e {
      ast::ty_mac({node: mac_var(i), _}) {
        alt (repls[i]) {
          from_ty(r) {(r.node, r.span)}
          _ {fail /* fixme error message */}}}
      _ {orig(e,s,fld)}
    }
}

fn print_expr(expr: @ast::expr) {
    let stdout = io::stdout();
    let pp = pprust::rust_printer(stdout);
    pprust::print_expr(pp, expr);
    pp::eof(pp.s);
    stdout.write_str("\n");
}

fn mk_file_substr(fname: str, line: uint, col: uint) -> codemap::file_substr {
    codemap::fss_external({filename: fname, line: line, col: col})
}

// Local Variables:
// mode: rust
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// End:
<commit_msg>port qquote to use dvec<commit_after>import ast::{crate, expr_, mac_invoc,
                     mac_aq, mac_var};
import fold::*;
import visit::*;
import ext::base::*;
import ext::build::*;
import parse::parser;
import parse::parser::parse_from_source_str;
import dvec::{dvec, extensions};

import print::*;
import io::*;

import codemap::span;

type aq_ctxt = @{lo: uint,
                 gather: dvec<{lo: uint, hi: uint,
                               e: @ast::expr,
                               constr: str}>};
enum fragment {
    from_expr(@ast::expr),
    from_ty(@ast::ty)
}

iface qq_helper {
    fn span() -> span;
    fn visit(aq_ctxt, vt<aq_ctxt>);
    fn extract_mac() -> option<ast::mac_>;
    fn mk_parse_fn(ext_ctxt,span) -> @ast::expr;
    fn get_fold_fn() -> str;
}

impl of qq_helper for @ast::crate {
    fn span() -> span {self.span}
    fn visit(cx: aq_ctxt, v: vt<aq_ctxt>) {visit_crate(*self, cx, v);}
    fn extract_mac() -> option<ast::mac_> {fail}
    fn mk_parse_fn(cx: ext_ctxt, sp: span) -> @ast::expr {
        mk_path(cx, sp, ["syntax", "ext", "qquote", "parse_crate"])
    }
    fn get_fold_fn() -> str {"fold_crate"}
}
impl of qq_helper for @ast::expr {
    fn span() -> span {self.span}
    fn visit(cx: aq_ctxt, v: vt<aq_ctxt>) {visit_expr(self, cx, v);}
    fn extract_mac() -> option<ast::mac_> {
        alt (self.node) {
          ast::expr_mac({node: mac, _}) {some(mac)}
          _ {none}
        }
    }
    fn mk_parse_fn(cx: ext_ctxt, sp: span) -> @ast::expr {
        mk_path(cx, sp, ["syntax", "parse", "parser", "parse_expr"])
    }
    fn get_fold_fn() -> str {"fold_expr"}
}
impl of qq_helper for @ast::ty {
    fn span() -> span {self.span}
    fn visit(cx: aq_ctxt, v: vt<aq_ctxt>) {visit_ty(self, cx, v);}
    fn extract_mac() -> option<ast::mac_> {
        alt (self.node) {
          ast::ty_mac({node: mac, _}) {some(mac)}
          _ {none}
        }
    }
    fn mk_parse_fn(cx: ext_ctxt, sp: span) -> @ast::expr {
        mk_path(cx, sp, ["syntax", "ext", "qquote", "parse_ty"])
    }
    fn get_fold_fn() -> str {"fold_ty"}
}
impl of qq_helper for @ast::item {
    fn span() -> span {self.span}
    fn visit(cx: aq_ctxt, v: vt<aq_ctxt>) {visit_item(self, cx, v);}
    fn extract_mac() -> option<ast::mac_> {fail}
    fn mk_parse_fn(cx: ext_ctxt, sp: span) -> @ast::expr {
        mk_path(cx, sp, ["syntax", "ext", "qquote", "parse_item"])
    }
    fn get_fold_fn() -> str {"fold_item"}
}
impl of qq_helper for @ast::stmt {
    fn span() -> span {self.span}
    fn visit(cx: aq_ctxt, v: vt<aq_ctxt>) {visit_stmt(self, cx, v);}
    fn extract_mac() -> option<ast::mac_> {fail}
    fn mk_parse_fn(cx: ext_ctxt, sp: span) -> @ast::expr {
        mk_path(cx, sp, ["syntax", "ext", "qquote", "parse_stmt"])
    }
    fn get_fold_fn() -> str {"fold_stmt"}
}
impl of qq_helper for @ast::pat {
    fn span() -> span {self.span}
    fn visit(cx: aq_ctxt, v: vt<aq_ctxt>) {visit_pat(self, cx, v);}
    fn extract_mac() -> option<ast::mac_> {fail}
    fn mk_parse_fn(cx: ext_ctxt, sp: span) -> @ast::expr {
        mk_path(cx, sp, ["syntax", "parse", "parser", "parse_pat"])
    }
    fn get_fold_fn() -> str {"fold_pat"}
}

fn gather_anti_quotes<N: qq_helper>(lo: uint, node: N) -> aq_ctxt
{
    let v = @{visit_expr: {|node, &&cx, v|
                  visit_aq(node, "from_expr", cx, v)},
              visit_ty: {|node, &&cx, v|
                  visit_aq(node, "from_ty", cx, v)}
              with *default_visitor()};
    let cx = @{lo:lo, gather: dvec()};
    node.visit(cx, mk_vt(v));
    // FIXME: Maybe this is an overkill (merge_sort), it might be better
    //   to just keep the gather array in sorted order ... (Issue #2250)
    cx.gather.swap { |v|
        vec::to_mut(std::sort::merge_sort({|a,b| a.lo < b.lo}, v))
    };
    ret cx;
}

fn visit_aq<T:qq_helper>(node: T, constr: str, &&cx: aq_ctxt, v: vt<aq_ctxt>)
{
    alt (node.extract_mac()) {
      some(mac_aq(sp, e)) {
        cx.gather.push({lo: sp.lo - cx.lo, hi: sp.hi - cx.lo,
                        e: e, constr: constr});
      }
      _ {node.visit(cx, v);}
    }
}

fn is_space(c: char) -> bool {
    parse::lexer::is_whitespace(c)
}

fn expand_ast(ecx: ext_ctxt, _sp: span,
              arg: ast::mac_arg, body: ast::mac_body)
    -> @ast::expr
{
    let mut what = "expr";
    option::iter(arg) {|arg|
        let args: [@ast::expr] =
            alt arg.node {
              ast::expr_vec(elts, _) { elts }
              _ {
                ecx.span_fatal
                    (_sp, "#ast requires arguments of the form `[...]`.")
              }
            };
        if vec::len::<@ast::expr>(args) != 1u {
            ecx.span_fatal(_sp, "#ast requires exactly one arg");
        }
        alt (args[0].node) {
          ast::expr_path(@{idents: id, _}) if vec::len(id) == 1u
              {what = id[0]}
          _ {ecx.span_fatal(args[0].span, "expected an identifier");}
        }
    }
    let body = get_mac_body(ecx,_sp,body);

    ret alt what {
      "crate" {finish(ecx, body, parse_crate)}
      "expr" {finish(ecx, body, parser::parse_expr)}
      "ty" {finish(ecx, body, parse_ty)}
      "item" {finish(ecx, body, parse_item)}
      "stmt" {finish(ecx, body, parse_stmt)}
      "pat" {finish(ecx, body, parser::parse_pat)}
      _ {ecx.span_fatal(_sp, "unsupported ast type")}
    };
}

fn parse_crate(p: parser) -> @ast::crate {
    parser::parse_crate_mod(p, [])
}

fn parse_ty(p: parser) -> @ast::ty {
    parser::parse_ty(p, false)
}

fn parse_stmt(p: parser) -> @ast::stmt {
    parser::parse_stmt(p, [])
}

fn parse_item(p: parser) -> @ast::item {
    alt parser::parse_item(p, [], ast::public) {
      some(item) { item }
      none       { fail "parse_item: parsing an item failed"; }
    }
}

fn finish<T: qq_helper>
    (ecx: ext_ctxt, body: ast::mac_body_, f: fn (p: parser) -> T)
    -> @ast::expr
{
    let cm = ecx.codemap();
    let str = @codemap::span_to_snippet(body.span, cm);
    #debug["qquote--str==%?", str];
    let fname = codemap::mk_substr_filename(cm, body.span);
    let node = parse_from_source_str
        (f, fname, codemap::fss_internal(body.span), str,
         ecx.cfg(), ecx.parse_sess());
    let loc = codemap::lookup_char_pos(cm, body.span.lo);

    let sp = node.span();
    let qcx = gather_anti_quotes(sp.lo, node);
    let cx = qcx;

    uint::range(1u, cx.gather.len()) {|i|
        assert cx.gather[i-1u].lo < cx.gather[i].lo;
        // ^^ check that the vector is sorted
        assert cx.gather[i-1u].hi <= cx.gather[i].lo;
        // ^^ check that the spans are non-overlapping
    }

    let mut str2 = "";
    enum state {active, skip(uint), blank};
    let mut state = active;
    let mut i = 0u, j = 0u;
    let g_len = cx.gather.len();
    str::chars_iter(*str) {|ch|
        if (j < g_len && i == cx.gather[j].lo) {
            assert ch == '$';
            let repl = #fmt("$%u ", j);
            state = skip(str::char_len(repl));
            str2 += repl;
        }
        alt state {
          active {str::push_char(str2, ch);}
          skip(1u) {state = blank;}
          skip(sk) {state = skip (sk-1u);}
          blank if is_space(ch) {str::push_char(str2, ch);}
          blank {str::push_char(str2, ' ');}
        }
        i += 1u;
        if (j < g_len && i == cx.gather[j].hi) {
            assert ch == ')';
            state = active;
            j += 1u;
        }
    }

    let cx = ecx;

    let cfg_call = {||
        mk_call_(cx, sp, mk_access(cx, sp, ["ext_cx"], "cfg"), [])
    };

    let parse_sess_call = {||
        mk_call_(cx, sp, mk_access(cx, sp, ["ext_cx"], "parse_sess"), [])
    };

    let pcall = mk_call(cx,sp,
                       ["syntax", "parse", "parser",
                        "parse_from_source_str"],
                       [node.mk_parse_fn(cx,sp),
                        mk_str(cx,sp, fname),
                        mk_call(cx,sp,
                                ["syntax","ext","qquote", "mk_file_substr"],
                                [mk_str(cx,sp, loc.file.name),
                                 mk_uint(cx,sp, loc.line),
                                 mk_uint(cx,sp, loc.col)]),
                        mk_unary(cx,sp, ast::box(ast::m_imm),
                                 mk_str(cx,sp, str2)),
                        cfg_call(),
                        parse_sess_call()]
                      );
    let mut rcall = pcall;
    if (g_len > 0u) {
        let gather = copy qcx.gather;
        rcall = mk_call(cx,sp,
                        ["syntax", "ext", "qquote", "replace"],
                        [pcall,
                         mk_vec_e(cx,sp, qcx.gather.map_to_vec {|g|
                             mk_call(cx,sp,
                                     ["syntax", "ext", "qquote", g.constr],
                                     [g.e])}),
                         mk_path(cx,sp,
                                 ["syntax", "ext", "qquote",
                                  node.get_fold_fn()])]);
    }
    ret rcall;
}

fn replace<T>(node: T, repls: [fragment], ff: fn (ast_fold, T) -> T)
    -> T
{
    let aft = default_ast_fold();
    let f_pre = {fold_expr: bind replace_expr(repls, _, _, _,
                                              aft.fold_expr),
                 fold_ty: bind replace_ty(repls, _, _, _,
                                          aft.fold_ty)
                 with *aft};
    ret ff(make_fold(f_pre), node);
}
fn fold_crate(f: ast_fold, &&n: @ast::crate) -> @ast::crate {
    @f.fold_crate(*n)
}
fn fold_expr(f: ast_fold, &&n: @ast::expr) -> @ast::expr {f.fold_expr(n)}
fn fold_ty(f: ast_fold, &&n: @ast::ty) -> @ast::ty {f.fold_ty(n)}
fn fold_item(f: ast_fold, &&n: @ast::item) -> @ast::item {f.fold_item(n)}
fn fold_stmt(f: ast_fold, &&n: @ast::stmt) -> @ast::stmt {f.fold_stmt(n)}
fn fold_pat(f: ast_fold, &&n: @ast::pat) -> @ast::pat {f.fold_pat(n)}

fn replace_expr(repls: [fragment],
                e: ast::expr_, s: span, fld: ast_fold,
                orig: fn@(ast::expr_, span, ast_fold)->(ast::expr_, span))
    -> (ast::expr_, span)
{
    alt e {
      ast::expr_mac({node: mac_var(i), _}) {
        alt (repls[i]) {
          from_expr(r) {(r.node, r.span)}
          _ {fail /* fixme error message */}}}
      _ {orig(e,s,fld)}
    }
}

fn replace_ty(repls: [fragment],
                e: ast::ty_, s: span, fld: ast_fold,
                orig: fn@(ast::ty_, span, ast_fold)->(ast::ty_, span))
    -> (ast::ty_, span)
{
    alt e {
      ast::ty_mac({node: mac_var(i), _}) {
        alt (repls[i]) {
          from_ty(r) {(r.node, r.span)}
          _ {fail /* fixme error message */}}}
      _ {orig(e,s,fld)}
    }
}

fn print_expr(expr: @ast::expr) {
    let stdout = io::stdout();
    let pp = pprust::rust_printer(stdout);
    pprust::print_expr(pp, expr);
    pp::eof(pp.s);
    stdout.write_str("\n");
}

fn mk_file_substr(fname: str, line: uint, col: uint) -> codemap::file_substr {
    codemap::fss_external({filename: fname, line: line, col: col})
}

// Local Variables:
// mode: rust
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// End:
<|endoftext|>