<commit_before>/*
 * Copyright (C) 2015 Benjamin Fry <benjaminfry@me.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#![allow(missing_docs)]

//! All defined errors for Trust-DNS

mod dnssec_error;
//mod encode_error;
mod client_error;
mod lexer_error;
mod parse_error;

pub use self::dnssec_error::Error as DnsSecError;
//pub use self::encode_error::Error as EncodeError;
pub use self::client_error::Error as ClientError;
pub use self::lexer_error::Error as LexerError;
pub use self::parse_error::Error as ParseError;

pub use self::dnssec_error::ErrorKind as DnsSecErrorKind;
//pub use self::encode_error::ErrorKind as EncodeErrorKind;
pub use self::client_error::ErrorKind as ClientErrorKind;
pub use self::lexer_error::ErrorKind as LexerErrorKind;
pub use self::parse_error::ErrorKind as ParseErrorKind;

pub use self::dnssec_error::ChainErr as DnsSecChainErr;
//pub use self::encode_error::ChainErr as EncodeChainErr;
pub use self::client_error::ChainErr as ClientChainErr;
pub use self::lexer_error::ChainErr as LexerChainErr;
pub use self::parse_error::ChainErr as ParseChainErr;

pub use self::dnssec_error::Result as DnsSecResult;
//pub type EncodeResult = self::encode_error::Result<()>;
pub use self::client_error::Result as ClientResult;
pub use self::lexer_error::Result as LexerResult;
pub use self::parse_error::Result as ParseResult;
<commit_msg>remove commented out errors<commit_after>/*
 * Copyright (C) 2015 Benjamin Fry <benjaminfry@me.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#![allow(missing_docs)]

//! All defined errors for Trust-DNS

mod dnssec_error;
mod client_error;
mod lexer_error;
mod parse_error;

pub use self::dnssec_error::Error as DnsSecError;
pub use self::client_error::Error as ClientError;
pub use self::lexer_error::Error as LexerError;
pub use self::parse_error::Error as ParseError;

pub use self::dnssec_error::ErrorKind as DnsSecErrorKind;
pub use self::client_error::ErrorKind as ClientErrorKind;
pub use self::lexer_error::ErrorKind as LexerErrorKind;
pub use self::parse_error::ErrorKind as ParseErrorKind;

pub use self::dnssec_error::ChainErr as DnsSecChainErr;
pub use self::client_error::ChainErr as ClientChainErr;
pub use self::lexer_error::ChainErr as LexerChainErr;
pub use self::parse_error::ChainErr as ParseChainErr;

pub use self::dnssec_error::Result as DnsSecResult;
pub use self::client_error::Result as ClientResult;
pub use self::lexer_error::Result as LexerResult;
pub use self::parse_error::Result as ParseResult;
<|endoftext|>


<commit_before>use anyhow::{Context, Error};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{
        tcp::{OwnedReadHalf, OwnedWriteHalf},
        TcpListener, TcpStream,
    },
    runtime::Runtime,
    sync::broadcast::Sender,
};
use url::Url;

pub(crate) struct ChaosProxy {
    address: SocketAddr,
    backend_address: SocketAddr,

    runtime: Runtime,
    listener: TcpListener,

    break_networking_send: Sender<()>,
    restore_networking_send: Sender<()>,
}

impl ChaosProxy {
    pub(crate) fn new(backend_address: SocketAddr) -> Result<Arc<Self>, Error> {
        let runtime = Runtime::new().expect("failed to create Tokio runtime");
        let listener = runtime.block_on(TcpListener::bind("127.0.0.1:0"))?;

        let (break_networking_send, _) = tokio::sync::broadcast::channel(16);
        let (restore_networking_send, _) = tokio::sync::broadcast::channel(16);

        let instance = Arc::new(ChaosProxy {
            address: listener.local_addr()?,
            backend_address,

            listener,
            runtime,

            break_networking_send,
            restore_networking_send,
        });

        let instance_clone = instance.clone();
        instance.runtime.spawn(async move {
            if let Err(err) = instance_clone.server_loop().await {
                eprintln!("ChaosProxy server error: {err}");
            }
        });

        Ok(instance)
    }

    pub(crate) fn proxy_database_url(url: &str) -> Result<(Arc<Self>, String), Error> {
        let mut db_url = Url::parse(url).context("failed to parse database url")?;
        let backend_addr = db_url
            .socket_addrs(|| Some(5432))
            .context("could not resolve database url")?
            .get(0)
            .copied()
            .ok_or_else(|| anyhow::anyhow!("the database url does not point to any IP"))?;

        let instance = ChaosProxy::new(backend_addr).unwrap();
        db_url.set_ip_host(instance.address.ip()).unwrap();
        db_url.set_port(Some(instance.address.port())).unwrap();
        Ok((instance, db_url.into()))
    }

    pub(crate) fn break_networking(&self) {
        self.break_networking_send
            .send(())
            .expect("failed to send the break_networking message");
    }

    pub(crate) fn restore_networking(&self) {
        self.restore_networking_send
            .send(())
            .expect("failed to send the restore_networking message");
    }

    async fn server_loop(self: Arc<Self>) -> Result<(), Error> {
        let mut break_networking_recv = self.break_networking_send.subscribe();
        let mut restore_networking_recv = self.restore_networking_send.subscribe();

        loop {
            let (client_read, client_write) = tokio::select! {
                accepted = self.listener.accept() => accepted?.0.into_split(),

                // When networking is broken stop accepting connections until restore_networking()
                _ = break_networking_recv.recv() => {
                    let _ = restore_networking_recv.recv().await;
                    continue;
                },
            };
            let (backend_read, backend_write) = TcpStream::connect(&self.backend_address)
                .await?
                .into_split();

            let self_clone = self.clone();
            self.runtime.spawn(async move {
                if let Err(err) = self_clone.proxy_data(client_read, backend_write).await {
                    eprintln!("ChaosProxy connection error: {err}");
                }
            });

            let self_clone = self.clone();
            tokio::spawn(async move {
                if let Err(err) = self_clone.proxy_data(backend_read, client_write).await {
                    eprintln!("ChaosProxy connection error: {err}");
                }
            });
        }
    }

    async fn proxy_data(
        &self,
        mut from: OwnedReadHalf,
        mut to: OwnedWriteHalf,
    ) -> Result<(), Error> {
        let mut break_connections_recv = self.break_networking_send.subscribe();
        let mut buf = [0; 1024];

        loop {
            tokio::select! {
                len = from.read(&mut buf) => {
                    let len = len?;
                    if len == 0 {
                        // EOF, the socket was closed
                        return Ok(());
                    }
                    to.write_all(&buf[0..len]).await?;
                }
                _ = break_connections_recv.recv() => {
                    to.shutdown().await?;
                    return Ok(());
                }
            }
        }
    }
}
<commit_msg>actually close and reopen the port in chaosproxy<commit_after>use anyhow::{Context, Error};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{
        tcp::{OwnedReadHalf, OwnedWriteHalf},
        TcpListener, TcpStream,
    },
    runtime::Runtime,
    sync::broadcast::Sender,
};
use url::Url;

pub(crate) struct ChaosProxy {
    address: SocketAddr,
    backend_address: SocketAddr,

    runtime: Runtime,

    break_networking_send: Sender<()>,
    restore_networking_send: Sender<()>,
}

impl ChaosProxy {
    pub(crate) fn new(backend_address: SocketAddr) -> Result<Arc<Self>, Error> {
        let runtime = Runtime::new().expect("failed to create Tokio runtime");
        let listener = runtime.block_on(TcpListener::bind("127.0.0.1:0"))?;

        let (break_networking_send, _) = tokio::sync::broadcast::channel(16);
        let (restore_networking_send, _) = tokio::sync::broadcast::channel(16);

        let instance = Arc::new(ChaosProxy {
            address: listener.local_addr()?,
            backend_address,

            runtime,

            break_networking_send,
            restore_networking_send,
        });

        let instance_clone = instance.clone();
        instance.runtime.spawn(async move {
            if let Err(err) = instance_clone.server_loop(listener).await {
                eprintln!("ChaosProxy server error: {err}");
            }
        });

        Ok(instance)
    }

    pub(crate) fn proxy_database_url(url: &str) -> Result<(Arc<Self>, String), Error> {
        let mut db_url = Url::parse(url).context("failed to parse database url")?;
        let backend_addr = db_url
            .socket_addrs(|| Some(5432))
            .context("could not resolve database url")?
            .get(0)
            .copied()
            .ok_or_else(|| anyhow::anyhow!("the database url does not point to any IP"))?;

        let instance = ChaosProxy::new(backend_addr).unwrap();
        db_url.set_ip_host(instance.address.ip()).unwrap();
        db_url.set_port(Some(instance.address.port())).unwrap();
        Ok((instance, db_url.into()))
    }

    pub(crate) fn break_networking(&self) {
        self.break_networking_send
            .send(())
            .expect("failed to send the break_networking message");
    }

    pub(crate) fn restore_networking(&self) {
        self.restore_networking_send
            .send(())
            .expect("failed to send the restore_networking message");
    }

    async fn server_loop(self: Arc<Self>, initial_listener: TcpListener) -> Result<(), Error> {
        let mut listener = Some(initial_listener);

        let mut break_networking_recv = self.break_networking_send.subscribe();
        let mut restore_networking_recv = self.restore_networking_send.subscribe();

        loop {
            if let Some(l) = &listener {
                tokio::select! {
                    accepted = l.accept() => {
                        self.clone().accept_connection(accepted?.0).await?;
                    },

                    _ = break_networking_recv.recv() => {
                        // Setting the listener to `None` results in the listener being dropped,
                        // which closes the network port. A new listener will be established when
                        // networking is restored.
                        listener = None;
                    },
                };
            } else {
                let _ = restore_networking_recv.recv().await;
                listener = Some(TcpListener::bind(self.address).await?);
            }
        }
    }

    async fn accept_connection(self: Arc<Self>, accepted: TcpStream) -> Result<(), Error> {
        let (client_read, client_write) = accepted.into_split();
        let (backend_read, backend_write) = TcpStream::connect(&self.backend_address)
            .await?
            .into_split();

        let self_clone = self.clone();
        self.runtime.spawn(async move {
            if let Err(err) = self_clone.proxy_data(client_read, backend_write).await {
                eprintln!("ChaosProxy connection error: {err}");
            }
        });

        let self_clone = self.clone();
        tokio::spawn(async move {
            if let Err(err) = self_clone.proxy_data(backend_read, client_write).await {
                eprintln!("ChaosProxy connection error: {err}");
            }
        });

        Ok(())
    }

    async fn proxy_data(
        &self,
        mut from: OwnedReadHalf,
        mut to: OwnedWriteHalf,
    ) -> Result<(), Error> {
        let mut break_connections_recv = self.break_networking_send.subscribe();
        let mut buf = [0; 1024];

        loop {
            tokio::select! {
                len = from.read(&mut buf) => {
                    let len = len?;
                    if len == 0 {
                        // EOF, the socket was closed
                        return Ok(());
                    }
                    to.write_all(&buf[0..len]).await?;
                }
                _ = break_connections_recv.recv() => {
                    to.shutdown().await?;
                    return Ok(());
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>/// mush graph backend and abstraction to swap out with custom backend

use uuid::Uuid;
use std::collections::{HashMap,HashSet,VecDeque};

/// unidirectional edge, use two edges for bidirectional/undirected graph
/*
struct GraphEdge {
    to: Uuid,
    maybe_weight: Option<f64>,
}
impl GraphEdge {
    // this should be private, this should be called by a graph method
    fn new (uuid: Uuid, weight: Option<f64>) -> GraphEdge {
        GraphEdge { to: uuid,
                    maybe_weight:weight }
    }
    pub fn get_directed (&self) -> &Uuid {
        &self.to
    }
    pub fn get_weight (&self) -> Option<f64> {
        self.maybe_weight
    }
}*/

#[derive(Debug)]
struct GraphNode {
    uuid: Uuid,
    edges: HashMap<Uuid,Option<f64>>, //to,weight
}
impl GraphNode {
    fn new (uuid: Uuid) -> GraphNode {
        GraphNode { uuid: uuid,
                    edges: HashMap::new() }
    }
    /// direct the node towards another node
    // todo: rename me! sounds too similar to unidirectional
    fn direct (&mut self, to:&Uuid, weight: Option<f64>) -> bool {
        self.edges.insert(*to,weight).is_some()
    }
    fn undirect (&mut self, to:&Uuid) -> bool {
        self.edges.remove(to).is_some()
    }
}

//----

pub struct Graph {
    nodes: HashMap<Uuid,GraphNode>,

    // todo: as traits
    is_weighted: bool,
    is_directed: bool,
}
impl Graph {
    pub fn default() -> Graph {
        Graph { nodes: HashMap::new(),
                is_weighted: false,
                is_directed: true, }
    }

    /// manual accessors
    fn get_mut(&mut self, uuid: &Uuid) -> Option<&mut GraphNode> {
        self.nodes.get_mut(uuid)
    }
    fn get(&mut self, uuid: &Uuid) -> Option<&GraphNode> {
        self.nodes.get(uuid)
    }

    pub fn add (&mut self) -> Uuid { //todo: maybe_edges fn arg
        let uuid = Uuid::new_v4();
        let n = GraphNode::new(uuid);
        self.nodes.insert(uuid,n);
        uuid
    }
    pub fn remove(&mut self, node: &Uuid) { //-> Option<GraphNode> {
        self.nodes.remove(node);
    }

    pub fn direct(&mut self, from: &Uuid, to: &Uuid) -> bool {
        let mut r = false;
        if let Some(f) = self.nodes.get_mut(from) {
            r = f.direct(to,None);
        }

        if self.is_directed { return r }
        else if r {
            if let Some(t) = self.nodes.get_mut(to) {
                r = t.direct(from,None);
            }
        }

        return r
    }
    pub fn undirect(&mut self, from: &Uuid, to: &Uuid) -> bool {
        let mut r = false;
        if let Some(f) = self.nodes.get_mut(from) {
            r = f.undirect(to);
        }

        if self.is_directed { return r }
        else if r {
            if let Some(t) = self.nodes.get_mut(to) {
                r = t.undirect(from);
            }
        }

        return r
    }

    // search functions
    // todo: consider weights between nodes to direct search, cycle-detection
    pub fn get_path(&mut self, s: GraphSearch) -> Option<Vec<Uuid>> {
        let mut visited: HashSet<Uuid> = HashSet::new();
        let mut result = vec!();

        match s {
            GraphSearch::Depth(from,to) => {
                let mut stack = vec!();

                stack.push(from);
                visited.insert(from);
                result.push(from);

                let mut cursor = Some(from);

                while cursor.is_some() {
                    if let Some(ref node) = self.get(&cursor.unwrap()) {
                        //get first unvisited node
                        let not_visited = node.edges.iter().find(|&(&n,v)| !visited.contains(&n));

                        if let Some((&n,w)) = not_visited {
                            stack.push(n);
                            visited.insert(n);
                            result.push(n);

                            if let Some(to_node) = to {
                                if n == to_node { break; }
                            }

                            cursor = Some(n);
                        }
                        else { cursor = stack.pop(); }
                    }
                    else { cursor = stack.pop(); }
                }

                if let Some(to_node) = to {
                    if result.contains(&to_node) {
                        return Some(result)
                    }
                }

                return None
            },
            GraphSearch::Breadth(from,to) => { // breadth first search, uses a queue
                let mut queue = VecDeque::new();

                queue.push_back(from);
                visited.insert(from);
                result.push(from);

                let mut cursor = Some(from);

                while cursor.is_some() {
                    if let Some(ref node) = self.get(&cursor.unwrap()) {
                        //get unvisted nodes to queue up
                        let not_visited: Vec<Option<(Uuid,Option<f64>)>> = node.edges.iter().map(|(&n,&v)| {
                            if !visited.contains(&n) {
                                Some((n,v))
                            }
                            else { None }
                        }).collect();

                        for maybe_node in not_visited {
                            if let Some((n,w)) = maybe_node {
                                queue.push_back(n);
                                visited.insert(n);
                                result.push(n);

                                if let Some(to_node) = to {
                                    if n == to_node { break; }
                                }

                                cursor = Some(n);
                            }
                        }

                        cursor = queue.pop_front();
                    }
                    else { cursor = queue.pop_front(); }
                }

                if let Some(to_node) = to {
                    if result.contains(&to_node) {
                        return Some(result)
                    }
                }

                return None
            },
            _ => None, // todo: djk algo
        }

    }

    pub fn get_cycle(&self, s: GraphSearch) -> Option<Vec<&Uuid>> {
        match s {
            GraphSearch::Depth(from,to) => { None }, // todo: use dfs
            _ => None,
        }
    }

    /// get immediate next node from list of connected nodes for the current node
    pub fn get_next(&self, from: &Uuid) -> Option<&Uuid> {
        if let Some(n) = self.nodes.get(from) {
            if let Some(next_id) = n.edges.iter().next() {
                return Some(next_id.0) // grab uuid key
            }
        }

        return None

    }

    fn is_connected() -> bool { false }
    fn is_complete() -> bool { false }

    fn get_path_shortest(&self) -> bool {
        if !self.is_weighted { false } //must be weighted
        else { false } //todo: use bfs
    }
}

// todo: impl as trait instead?
pub enum GraphSearch {
    Depth(Uuid,Option<Uuid>), // used on part of graph for reachability, and all of graph for cycle-detection
    Breadth(Uuid,Option<Uuid>), // used on part of graph for reachability, and (unweighted) for shortest path
    Dijkstra(Uuid,Uuid), // used on part of graph (weighted) for shortest path
}

pub struct GraphBuilder(Graph);

impl GraphBuilder {
    pub fn new() -> GraphBuilder {
        GraphBuilder(Graph::default())
    }

    pub fn directed(mut self, d: bool) -> GraphBuilder {
        self.0.is_directed = d;
        self
    }
    pub fn weighted(mut self, w: bool) -> GraphBuilder {
        self.0.is_weighted = w;
        self
    }
    pub fn build(mut self) -> Graph {
        self.0
    }
}

// ----
// tests
// ----
#[cfg(test)]
mod tests {
    extern crate test;

    use ::{Graph,GraphSearch};

    #[test]
    fn test_basic() {
        let mut graph = Graph::default();
        let mut nodes = vec!();
        for i in 0..5 {
            nodes.push(graph.add());
        }

        graph.get_mut(&nodes[0]).unwrap().direct(&nodes[1],None); //note there is no verification that node[1] exists when doing this manually

        graph.direct(&nodes[3],&nodes[0]);
        graph.direct(&nodes[4],&nodes[3]);
        let n2 = graph.remove(&nodes[2]);
        assert!(!graph.direct(&nodes[2],&nodes[3]));

        let r = graph.get_path(GraphSearch::Depth(nodes[0],Some(nodes[4])));
        assert!(!r.is_some());

        let r = graph.get_path(GraphSearch::Depth(nodes[4],Some(nodes[0])));
        assert_eq!(r.unwrap().len(), 3);

        let r = graph.get_path(GraphSearch::Breadth(nodes[4],Some(nodes[0])));
        assert!(r.is_some());

    }

}
<commit_msg>cycle detection, needs a few tests to make sure this is correct<commit_after>/// mush graph backend and abstraction to swap out with custom backend

use uuid::Uuid;
use std::collections::{HashMap,HashSet,VecDeque};

/// unidirectional edge, use two edges for bidirectional/undirected graph
/*
struct GraphEdge {
    to: Uuid,
    maybe_weight: Option<f64>,
}
impl GraphEdge {
    // this should be private, this should be called by a graph method
    fn new (uuid: Uuid, weight: Option<f64>) -> GraphEdge {
        GraphEdge { to: uuid,
                    maybe_weight:weight }
    }
    pub fn get_directed (&self) -> &Uuid {
        &self.to
    }
    pub fn get_weight (&self) -> Option<f64> {
        self.maybe_weight
    }
}*/

#[derive(Debug)]
struct GraphNode {
    uuid: Uuid,
    edges: HashMap<Uuid,Option<f64>>, //to,weight
}
impl GraphNode {
    fn new (uuid: Uuid) -> GraphNode {
        GraphNode { uuid: uuid,
                    edges: HashMap::new() }
    }
    /// direct the node towards another node
    // todo: rename me! sounds too similar to unidirectional
    fn direct (&mut self, to:&Uuid, weight: Option<f64>) -> bool {
        self.edges.insert(*to,weight).is_some()
    }
    fn undirect (&mut self, to:&Uuid) -> bool {
        self.edges.remove(to).is_some()
    }
}

//----

pub struct Graph {
    nodes: HashMap<Uuid,GraphNode>,

    // todo: as traits
    is_weighted: bool,
    is_directed: bool,
}
impl Graph {
    pub fn default() -> Graph {
        Graph { nodes: HashMap::new(),
                is_weighted: false,
                is_directed: true, }
    }

    /// manual accessors
    fn get_mut(&mut self, uuid: &Uuid) -> Option<&mut GraphNode> {
        self.nodes.get_mut(uuid)
    }
    fn get(&mut self, uuid: &Uuid) -> Option<&GraphNode> {
        self.nodes.get(uuid)
    }

    pub fn add (&mut self) -> Uuid { //todo: maybe_edges fn arg
        let uuid = Uuid::new_v4();
        let n = GraphNode::new(uuid);
        self.nodes.insert(uuid,n);
        uuid
    }
    pub fn remove(&mut self, node: &Uuid) { //-> Option<GraphNode> {
        self.nodes.remove(node);
    }

    pub fn direct(&mut self, from: &Uuid, to: &Uuid) -> bool {
        let mut r = false;
        if let Some(f) = self.nodes.get_mut(from) {
            r = f.direct(to,None);
        }

        if self.is_directed { return r }
        else if r {
            if let Some(t) = self.nodes.get_mut(to) {
                r = t.direct(from,None);
            }
        }

        return r
    }
    pub fn undirect(&mut self, from: &Uuid, to: &Uuid) -> bool {
        let mut r = false;
        if let Some(f) = self.nodes.get_mut(from) {
            r = f.undirect(to);
        }

        if self.is_directed { return r }
        else if r {
            if let Some(t) = self.nodes.get_mut(to) {
                r = t.undirect(from);
            }
        }

        return r
    }

    // search functions
    // todo: consider weights between nodes to direct search, cycle-detection
    pub fn get_path(&mut self, s: GraphSearch) -> Option<Vec<Uuid>> {
        let mut visited: HashSet<Uuid> = HashSet::new();
        let mut result = vec!();

        match s {
            GraphSearch::Depth(from,to) => {
                let mut stack = vec!();

                stack.push(from);
                visited.insert(from);
                result.push(from);

                let mut cursor = Some(from);

                while cursor.is_some() {
                    if let Some(ref node) = self.get(&cursor.unwrap()) {
                        //get first unvisited node
                        let not_visited = node.edges.iter().find(|&(&n,v)| !visited.contains(&n));

                        if let Some((&n,w)) = not_visited {
                            stack.push(n);
                            visited.insert(n);
                            result.push(n);

                            if let Some(to_node) = to {
                                if n == to_node { break; }
                            }

                            cursor = Some(n);
                        }
                        else { cursor = stack.pop(); }
                    }
                    else { cursor = stack.pop(); }
                }

                if let Some(to_node) = to {
                    if result.contains(&to_node) {
                        return Some(result)
                    }
                }

                return None
            },
            GraphSearch::Breadth(from,to) => { // breadth first search, uses a queue
                let mut queue = VecDeque::new();

                queue.push_back(from);
                visited.insert(from);
                result.push(from);

                let mut cursor = Some(from);

                while cursor.is_some() {
                    if let Some(ref node) = self.get(&cursor.unwrap()) {
                        //get unvisted nodes to queue up
                        let not_visited: Vec<Option<(Uuid,Option<f64>)>> = node.edges.iter().map(|(&n,&v)| {
                            if !visited.contains(&n) {
                                Some((n,v))
                            }
                            else { None }
                        }).collect();

                        for maybe_node in not_visited {
                            if let Some((n,w)) = maybe_node {
                                queue.push_back(n);
                                visited.insert(n);
                                result.push(n);

                                if let Some(to_node) = to {
                                    if n == to_node { break; }
                                }

                                cursor = Some(n);
                            }
                        }

                        cursor = queue.pop_front();
                    }
                    else { cursor = queue.pop_front(); }
                }

                if let Some(to_node) = to {
                    if result.contains(&to_node) {
                        return Some(result)
                    }
                }

                return None
            },
            _ => None, // todo: djk algo
        }

    }

    pub fn get_cycle(&mut self, s: GraphSearch) -> Option<Vec<Uuid>> {
        match s {
            GraphSearch::Depth(from,to) => {
                //const NEW: u8 = 0;
                const SEARCH: u8 = 1; // currently searching in this node
                const VISIT: u8 = 2; // has visited this node

                let mut status: HashMap<Uuid,u8> = HashMap::new();
                let mut result = vec!();
                let mut stack = vec!();

                stack.push(from);
                //status.insert(from,SEARCH);
                result.push(from);

                let mut cursor = Some(from);

                while cursor.is_some() {
                    status.insert(cursor.unwrap(),SEARCH);
                    println!("{:?}",status);
                    if let Some(ref node) = self.get(&cursor.unwrap()) {
                        //get first unvisited node
                        let not_visited = node.edges.iter().find(|&(&n,v)| !status.contains_key(&n));

                        let has_searching = node.edges.iter().find(|&(&n,v)| {
                            if let Some(_v) = status.get(&n) {
                                if *_v == SEARCH { true } //cycles found
                                else { false }
                            }
                            else { false }
                        });

                        if let Some((&n,w)) = has_searching {
                            result.push(n);
                            return Some(result);
                        }

                        if let Some((&n,w)) = not_visited {
                            stack.push(n);
                            status.insert(n,VISIT);
                            result.push(n);

                            if let Some(to_node) = to {
                                if n == to_node { break; }
                            }

                            cursor = Some(n);
                        }
                        else { cursor = stack.pop(); }
                    }
                    else { cursor = stack.pop(); }
                }

                return None
            },
            _ => None,
        }
    }

    /// get immediate next node from list of connected nodes for the current node
    pub fn get_next(&self, from: &Uuid) -> Option<&Uuid> {
        if let Some(n) = self.nodes.get(from) {
            if let Some(next_id) = n.edges.iter().next() {
                return Some(next_id.0) // grab uuid key
            }
        }

        return None

    }

    fn is_connected() -> bool { false }
    fn is_complete() -> bool { false }

    fn get_path_shortest(&self) -> bool {
        if !self.is_weighted { false } //must be weighted
        else { false } //todo: use bfs
    }
}

// todo: impl as trait instead?
pub enum GraphSearch {
    Depth(Uuid,Option<Uuid>), // used on part of graph for reachability, and all of graph for cycle-detection
    Breadth(Uuid,Option<Uuid>), // used on part of graph for reachability, and (unweighted) for shortest path
    Dijkstra(Uuid,Uuid), // used on part of graph (weighted) for shortest path
}

pub struct GraphBuilder(Graph);

impl GraphBuilder {
    pub fn new() -> GraphBuilder {
        GraphBuilder(Graph::default())
    }

    pub fn directed(mut self, d: bool) -> GraphBuilder {
        self.0.is_directed = d;
        self
    }
    pub fn weighted(mut self, w: bool) -> GraphBuilder {
        self.0.is_weighted = w;
        self
    }
    pub fn build(mut self) -> Graph {
        self.0
    }
}

// ----
// tests
// ----
#[cfg(test)]
mod tests {
    extern crate test;

    use ::{Graph,GraphSearch};

    #[test]
    fn test_basic() {
        let mut graph = Graph::default();
        let mut nodes = vec!();
        for i in 0..5 {
            nodes.push(graph.add());
        }

        graph.get_mut(&nodes[0]).unwrap().direct(&nodes[1],None); //note there is no verification that node[1] exists when doing this manually

        graph.direct(&nodes[3],&nodes[0]);
        graph.direct(&nodes[4],&nodes[3]);
        let n2 = graph.remove(&nodes[2]);
        assert!(!graph.direct(&nodes[2],&nodes[3]));

        let r = graph.get_path(GraphSearch::Depth(nodes[0],Some(nodes[4])));
        assert!(!r.is_some());

        let r = graph.get_path(GraphSearch::Depth(nodes[4],Some(nodes[0])));
        assert_eq!(r.unwrap().len(), 3);

        let r = graph.get_path(GraphSearch::Breadth(nodes[4],Some(nodes[0])));
        assert!(r.is_some());

        let r = graph.get_cycle(GraphSearch::Depth(nodes[4],Some(nodes[0])));
        assert!(!r.is_some());

        //find cycles
        let r = graph.get_cycle(GraphSearch::Depth(nodes[4],Some(nodes[0])));
        assert!(!r.is_some());
        let n2 = graph.add();
        graph.direct(&n2,&nodes[3]);
        graph.direct(&nodes[0],&n2);
        let r = graph.get_cycle(GraphSearch::Depth(nodes[4],None));
        assert!(r.is_some());

    }

}
<|endoftext|>


<commit_before>// Copyright 2014 Christopher Schröder, Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

pub mod record;
pub mod header;

use std::ffi;
use std::path;
use std::ffi::AsOsStr;
use std::os::unix::prelude::OsStrExt;

use htslib;

/// A trait for a BAM reader with a read method.
pub trait BAMRead {
    /// Read next BAM record into given record.
    /// Use this method in combination with a single allocated record to avoid the reallocations
    /// occurring with the iterator.
    ///
    /// # Arguments
    ///
    /// * `record` - the record to be filled
    fn read(&self, record: &mut record::Record) -> Result<(), ReadError>;
}

/// A BAM reader.
pub struct BAMReader {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl BAMReader {
    /// Create a new BAMReader.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    pub fn new<P: path::AsPath>(path: P) -> Self {
        let f = unsafe {
            htslib::bgzf_open(
                path.as_path().as_os_str().to_cstring().unwrap().as_ptr(),
                ffi::CString::new(b"r").unwrap().as_ptr()
            )
        };
        let header = unsafe { htslib::bam_hdr_read(f) };
        BAMReader { f : f, header : header }
    }

    /// Iterator over the records of the BAM file.
    pub fn records(self) -> Records {
        Records { bam: self }
    }
}

impl BAMRead for BAMReader {
    fn read(&self, record: &mut record::Record) -> Result<(), ReadError> {
        match unsafe { htslib::bam_read1(self.f, &mut record.inner) } {
            -1 => Err(ReadError::EOF),
            -2 => Err(ReadError::Truncated),
            -4 => Err(ReadError::Invalid),
            _  => Ok(())
        }
    }
}

impl Drop for BAMReader {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

/// Wrapper for opening a BAM file.
fn bgzf_open<P: path::AsPath>(path: &P, mode: &[u8]) -> *mut htslib::Struct_BGZF {
    unsafe {
        htslib::bgzf_open(
            path.as_path().as_os_str().to_cstring().unwrap().as_ptr(),
            ffi::CString::new(mode).unwrap().as_ptr()
        )
    }
}

/// A BAM writer.
pub struct BAMWriter {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl BAMWriter {
    /// Create a new BAM file.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    /// * `header` - header definition to use
    pub fn new<P: path::AsPath>(path: &P, header: &header::Header) -> Self {
        let f = bgzf_open(path, b"w");
        let header_record = unsafe {
            let header_string = header.to_bytes();
            htslib::sam_hdr_parse(
                header_string.len() as i32,
                ffi::CString::new(header_string).unwrap().as_ptr()
            )
        };
        unsafe { htslib::bam_hdr_write(f, header_record); }
        BAMWriter { f: f, header: header_record }
    }

    /// Create a new BAM file from template.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    /// * `template` - the template BAM. Use "-" for stdin.
    pub fn with_template<P: path::AsPath, T: path::AsPath>(path: &P, template: &T) -> Self {
        let t = bgzf_open(template, b"r");
        let header = unsafe { htslib::bam_hdr_read(t) };
        unsafe { htslib::bgzf_close(t); }
        let f = bgzf_open(path, b"w");
        unsafe { htslib::bam_hdr_write(f, header); }
        BAMWriter { f: f, header: header }
    }

    /// Write record to BAM.
    ///
    /// # Arguments
    ///
    /// * `record` - the record to write
    pub fn write(&mut self, record: &record::Record) -> Result<(), ()> {
        if unsafe { htslib::bam_write1(self.f, &record.inner) } == -1 {
            Err(())
        }
        else {
            Ok(())
        }
    }
}

impl Drop for BAMWriter {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

/// Iterator over the records of a BAM.
pub struct Records {
    bam: BAMReader
}

impl Iterator for Records {
    type Item = Result<record::Record, ReadError>;

    fn next(&mut self) -> Option<Result<record::Record, ReadError>> {
        let mut record = record::Record::new();
        match self.bam.read(&mut record) {
            Err(ReadError::EOF) => None,
            Ok(())   => Some(Ok(record)),
            Err(err) => Some(Err(err))
        }
    }
}

pub enum ReadError {
    Truncated,
    Invalid,
    EOF,
}

#[cfg(test)]
mod tests {
    use super::record::*;
    use super::*;
    use std::str;

    #[test]
    fn test_read() {
        let names = [b"I", b"II.14978392", b"III", b"IV", b"V", b"VI"];
        let flags = [16u16, 16u16, 16u16, 16u16, 16u16, 2048u16];
        let seqs = [
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"ACTAAGCCTAAGCCTAAGCCTAAGCCAATTATCGATTTCTGAAAAAATTATCGAATTTTCTAGAAATTTTGCAAATTTTTTCATAAAATTATCGATTTTA",
        ];
        let cigars = [
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(100000), Cigar::Match(73)],
        ];

        let bam = BAMReader::new("test.bam");

        for (i, record) in bam.records().enumerate() {
            let rec = record.ok().expect("Expected valid record");
            println!("{}", str::from_utf8(rec.qname()).ok().unwrap());
            assert_eq!(rec.qname(), names[i]);
            assert_eq!(rec.flags(), flags[i]);
            assert_eq!(rec.seq().as_bytes(), seqs[i]);
            assert_eq!(rec.cigar(), cigars[i]);
        }
    }

    #[test]
    fn test_set_record() {
        let qname = b"I";
        let cigar = [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)];
        let seq =  b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGC";
        let qual = b"JJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJ";

        let mut rec = record::Record::new();
        rec.set_reverse();
        rec.set(qname, &cigar, seq, qual);
        rec.push_aux(b"NM", &Aux::Integer(15));
        assert_eq!(rec.qname(), qname);
        assert_eq!(rec.cigar(), cigar);
        assert_eq!(rec.seq().as_bytes(), seq);
        assert_eq!(rec.qual(), qual);
        assert!(rec.is_reverse());
        assert_eq!(rec.aux(b"NM").unwrap(), Aux::Integer(15));
    }
}
<commit_msg>Use wrapper.<commit_after>// Copyright 2014 Christopher Schröder, Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

pub mod record;
pub mod header;

use std::ffi;
use std::path;
use std::ffi::AsOsStr;
use std::os::unix::prelude::OsStrExt;

use htslib;

/// A trait for a BAM reader with a read method.
pub trait BAMRead {
    /// Read next BAM record into given record.
    /// Use this method in combination with a single allocated record to avoid the reallocations
    /// occurring with the iterator.
    ///
    /// # Arguments
    ///
    /// * `record` - the record to be filled
    fn read(&self, record: &mut record::Record) -> Result<(), ReadError>;
}

/// A BAM reader.
pub struct BAMReader {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl BAMReader {
    /// Create a new BAMReader.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    pub fn new<P: path::AsPath>(path: P) -> Self {
        let f = bgzf_open(path, b"r");
        let header = unsafe { htslib::bam_hdr_read(f) };
        BAMReader { f : f, header : header }
    }

    /// Iterator over the records of the BAM file.
    pub fn records(self) -> Records {
        Records { bam: self }
    }
}

impl BAMRead for BAMReader {
    fn read(&self, record: &mut record::Record) -> Result<(), ReadError> {
        match unsafe { htslib::bam_read1(self.f, &mut record.inner) } {
            -1 => Err(ReadError::EOF),
            -2 => Err(ReadError::Truncated),
            -4 => Err(ReadError::Invalid),
            _  => Ok(())
        }
    }
}

impl Drop for BAMReader {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

/// A BAM writer.
pub struct BAMWriter {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl BAMWriter {
    /// Create a new BAM file.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    /// * `header` - header definition to use
    pub fn new<P: path::AsPath>(path: &P, header: &header::Header) -> Self {
        let f = bgzf_open(path, b"w");
        let header_record = unsafe {
            let header_string = header.to_bytes();
            htslib::sam_hdr_parse(
                header_string.len() as i32,
                ffi::CString::new(header_string).unwrap().as_ptr()
            )
        };
        unsafe { htslib::bam_hdr_write(f, header_record); }
        BAMWriter { f: f, header: header_record }
    }

    /// Create a new BAM file from template.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    /// * `template` - the template BAM. Use "-" for stdin.
    pub fn with_template<P: path::AsPath, T: path::AsPath>(path: &P, template: &T) -> Self {
        let t = bgzf_open(template, b"r");
        let header = unsafe { htslib::bam_hdr_read(t) };
        unsafe { htslib::bgzf_close(t); }
        let f = bgzf_open(path, b"w");
        unsafe { htslib::bam_hdr_write(f, header); }
        BAMWriter { f: f, header: header }
    }

    /// Write record to BAM.
    ///
    /// # Arguments
    ///
    /// * `record` - the record to write
    pub fn write(&mut self, record: &record::Record) -> Result<(), ()> {
        if unsafe { htslib::bam_write1(self.f, &record.inner) } == -1 {
            Err(())
        }
        else {
            Ok(())
        }
    }
}

impl Drop for BAMWriter {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

/// Iterator over the records of a BAM.
pub struct Records {
    bam: BAMReader
}

impl Iterator for Records {
    type Item = Result<record::Record, ReadError>;

    fn next(&mut self) -> Option<Result<record::Record, ReadError>> {
        let mut record = record::Record::new();
        match self.bam.read(&mut record) {
            Err(ReadError::EOF) => None,
            Ok(())   => Some(Ok(record)),
            Err(err) => Some(Err(err))
        }
    }
}

pub enum ReadError {
    Truncated,
    Invalid,
    EOF,
}

/// Wrapper for opening a BAM file.
fn bgzf_open<P: path::AsPath>(path: &P, mode: &[u8]) -> *mut htslib::Struct_BGZF {
    unsafe {
        htslib::bgzf_open(
            path.as_path().as_os_str().to_cstring().unwrap().as_ptr(),
            ffi::CString::new(mode).unwrap().as_ptr()
        )
    }
}

#[cfg(test)]
mod tests {
    use super::record::*;
    use super::*;
    use std::str;

    #[test]
    fn test_read() {
        let names = [b"I", b"II.14978392", b"III", b"IV", b"V", b"VI"];
        let flags = [16u16, 16u16, 16u16, 16u16, 16u16, 2048u16];
        let seqs = [
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"ACTAAGCCTAAGCCTAAGCCTAAGCCAATTATCGATTTCTGAAAAAATTATCGAATTTTCTAGAAATTTTGCAAATTTTTTCATAAAATTATCGATTTTA",
        ];
        let cigars = [
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(100000), Cigar::Match(73)],
        ];

        let bam = BAMReader::new("test.bam");

        for (i, record) in bam.records().enumerate() {
            let rec = record.ok().expect("Expected valid record");
            println!("{}", str::from_utf8(rec.qname()).ok().unwrap());
            assert_eq!(rec.qname(), names[i]);
            assert_eq!(rec.flags(), flags[i]);
            assert_eq!(rec.seq().as_bytes(), seqs[i]);
            assert_eq!(rec.cigar(), cigars[i]);
        }
    }

    #[test]
    fn test_set_record() {
        let qname = b"I";
        let cigar = [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)];
        let seq =  b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGC";
        let qual = b"JJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJ";

        let mut rec = record::Record::new();
        rec.set_reverse();
        rec.set(qname, &cigar, seq, qual);
        rec.push_aux(b"NM", &Aux::Integer(15));
        assert_eq!(rec.qname(), qname);
        assert_eq!(rec.cigar(), cigar);
        assert_eq!(rec.seq().as_bytes(), seq);
        assert_eq!(rec.qual(), qual);
        assert!(rec.is_reverse());
        assert_eq!(rec.aux(b"NM").unwrap(), Aux::Integer(15));
    }
}
<|endoftext|>


<commit_before>//! The Brainfuck interpreter executable.
//!
//! ```
//! USAGE:
//!     bfi [FLAGS] [OPTIONS] [--] [FILE]...
//!
//! FLAGS:
//!         --ast          Interpret the unoptimized AST
//!         --byte         Compile AST to bytecode
//!     -h, --help         Prints help information
//!         --jit          JIT to native x64 (default)
//!         --llvm         JIT using LLVM
//!         --peep         Interpret the peephole-optimized AST
//!         --rle          Interpret the run-length encoded the AST
//!     -u, --unchecked    Omit memory bounds checks in JIT
//!     -V, --version      Prints version information
//!
//! OPTIONS:
//!     -e, --expr <CODE>...    BF code to execute
//!     -s, --size <SIZE>       Memory size in bytes (default 30,000)
//!
//! ARGS:
//!     <FILE>...    The source file(s) to interpret
//! ```
//!
//! See [the library crate documentation](../bf/index.html) for more.

extern crate bf;
extern crate clap;

use std::io::{self, Read, Write};
use std::fs::File;
use std::process::exit;

use clap::{Arg, App};

use bf::ast;
use bf::traits::*;

const VERSION: &str = "0.4.6";

#[derive(Debug, Clone)]
struct Options {
    program_text:  Vec<u8>,
    memory_size:   Option<usize>,
    compiler_pass: Pass,
    unchecked:     bool,
}

#[derive(Debug, Clone, Copy)]
enum Pass {
    Ast,
    Rle,
    Bytecode,
    Peephole,
    #[cfg(feature = "jit")]
    Jit,
    #[cfg(feature = "llvm")]
    Llvm,
}

fn main() {
    let options = get_options();

    let program = parse(&options);

    match options.compiler_pass {
        Pass::Ast => {
            interpret(&*program, &options);
        }

        Pass::Rle => {
            let program = program.rle_compile();
            interpret(&*program, &options);
        }

        Pass::Peephole => {
            let program = program.peephole_compile();
            interpret(&*program, &options);
        }

        Pass::Bytecode => {
            let program = program.bytecode_compile();
            interpret(&*program, &options);
        }

        #[cfg(feature = "jit")]
        Pass::Jit => {
            let program = program.jit_compile(!options.unchecked);
            interpret(&program, &options);
        }

        #[cfg(feature = "llvm")]
        Pass::Llvm => {
            program.llvm_run(options.memory_size)
                .unwrap_or_else(|e| error_exit(3, format!("runtime error: {}.", e)));
        }
    }
}

fn parse(options: &Options) -> Box<ast::Program> {
    ast::parse_program(&options.program_text)
        .unwrap_or_else(|e| error_exit(2, format!("syntax error: {}.", e)))
}

fn interpret<P: Interpretable + ?Sized>(program: &P, options: &Options) {
    program.interpret_stdin(options.memory_size)
        .unwrap_or_else(|e| error_exit(3, format!("runtime error: {}.", e)))
}

#[cfg(feature = "jit")]
const DEFAULT_PASS: Pass = Pass::Jit;

#[cfg(not(feature = "jit"))]
const DEFAULT_PASS: Pass = Pass::Peephole;

fn get_options() -> Options {
    let mut result = Options {
        program_text:  Vec::new(),
        memory_size:   None,
        compiler_pass: DEFAULT_PASS,
        unchecked:     false,
    };

    let matches = build_clap_app().get_matches();

    if let Some(size) = matches.value_of("size") {
        let size = size.parse()
            .unwrap_or_else(|e|
                error_exit(1, format!("error: could not parse memory size: {}.", e)));
        if size == 0 {
            error_exit(1, "error: memory size must be at least 1.".to_owned());
        }
        result.memory_size = Some(size);
    }

    if matches.is_present("jit") {
        #[cfg(feature = "jit")]
        let _ = result.compiler_pass = Pass::Jit;
    } else if matches.is_present("llvm") {
        #[cfg(feature = "llvm")]
        let _ = result.compiler_pass = Pass::Llvm;
    } else if matches.is_present("byte") {
        result.compiler_pass = Pass::Bytecode;
    } else if matches.is_present("peep") {
        result.compiler_pass = Pass::Peephole;
    } else if matches.is_present("rle") {
        result.compiler_pass = Pass::Rle;
    } else if matches.is_present("ast") {
        result.compiler_pass = Pass::Ast;
    }

    if matches.is_present("unchecked") {
        result.unchecked = true;
    }

    if let Some(exprs) = matches.values_of("expr") {
        for e in exprs {
            result.program_text.extend(e.as_bytes());
        }
    } else if let Some(files) = matches.values_of("FILE") {
        for f in files {
            let mut file = File::open(f)
                .unwrap_or_else(|e| error_exit(1, format!("{}: ‘{}’.", e, f)));
            file.read_to_end(&mut result.program_text)
                .unwrap_or_else(|e| error_exit(1, format!("{}: ‘{}’.", e, f)));
        }
    } else {
        error_exit(1, "error: no program given.".to_owned());
    }

    result
}

fn build_clap_app() -> App<'static, 'static> {
    let app = App::new("bfi")
        .version(VERSION)
        .author("Jesse A. Tov <jesse.tov@gmail.com>")
        .about("A Brainfuck interpreter")
        .arg(Arg::with_name("expr")
            .short("e")
            .long("expr")
            .value_name("CODE")
            .help("BF code to execute")
            .multiple(true)
            .takes_value(true)
            .conflicts_with("FILE"))
        .arg(Arg::with_name("FILE")
            .help("The source file(s) to interpret")
            .multiple(true)
            .conflicts_with("expr")
            .index(1))
        .arg(Arg::with_name("size")
            .short("s")
            .long("size")
            .value_name("SIZE")
            .help("Memory size in bytes (default 30,000)")
            .takes_value(true))
        .arg(Arg::with_name("ast")
            .long("ast")
            .help("Interpret the unoptimized AST")
            .conflicts_with_all(&["rle", "peep", "byte", "jit", "llvm"]))
        .arg(Arg::with_name("rle")
            .long("rle")
            .help("Interpret the run-length encoded the AST")
            .conflicts_with_all(&["ast", "peep", "byte", "jit", "llvm"]))
        .arg(Arg::with_name("peep")
            .long("peep")
            .help(
                if cfg!(feature = "jit") {
                    "Interpret the peephole-optimized AST"
                } else {
                    "Interpret the peephole-optimized AST (default)"
                })
            .conflicts_with_all(&["ast", "rle", "byte", "jit", "llvm"]))
        .arg(Arg::with_name("byte")
            .long("byte")
            .help("Compile AST to bytecode")
            .conflicts_with_all(&["ast", "rle", "peep", "jit", "llvm"]));

    #[cfg(feature = "llvm")]
    let app = app
        .arg(Arg::with_name("llvm")
            .long("llvm")
            .help("JIT using LLVM")
            .conflicts_with_all(&["ast", "rle", "peep", "byte", "jit"]));

    #[cfg(feature = "jit")]
    let app = app
        .arg(Arg::with_name("jit")
            .long("jit")
            .help("JIT to native x64 (default)")
            .conflicts_with_all(&["ast", "rle", "peep", "byte", "llvm"]));

    #[cfg(any(feature = "jit"))]
    let app = app
        .arg(Arg::with_name("unchecked")
            .short("u")
            .long("unchecked")
            .help("Omit memory bounds checks in JIT")
            .conflicts_with_all(&["ast", "rle", "peep", "byte", "llvm"]));

    app
}

fn error_exit(code: i32, msg: String) -> ! {
    eprintln!("bfi: {}", msg);
    exit(code)
}

<commit_msg>Clippy cleaner.<commit_after>//! The Brainfuck interpreter executable.
//!
//! ```
//! USAGE:
//!     bfi [FLAGS] [OPTIONS] [--] [FILE]...
//!
//! FLAGS:
//!         --ast          Interpret the unoptimized AST
//!         --byte         Compile AST to bytecode
//!     -h, --help         Prints help information
//!         --jit          JIT to native x64 (default)
//!         --llvm         JIT using LLVM
//!         --peep         Interpret the peephole-optimized AST
//!         --rle          Interpret the run-length encoded the AST
//!     -u, --unchecked    Omit memory bounds checks in JIT
//!     -V, --version      Prints version information
//!
//! OPTIONS:
//!     -e, --expr <CODE>...    BF code to execute
//!     -s, --size <SIZE>       Memory size in bytes (default 30,000)
//!
//! ARGS:
//!     <FILE>...    The source file(s) to interpret
//! ```
//!
//! See [the library crate documentation](../bf/index.html) for more.

extern crate bf;
extern crate clap;

use std::io::Read;
use std::fs::File;
use std::process::exit;

use clap::{Arg, App};

use bf::ast;
use bf::traits::*;

const VERSION: &str = "0.4.6";

#[derive(Debug, Clone)]
struct Options {
    program_text:  Vec<u8>,
    memory_size:   Option<usize>,
    compiler_pass: Pass,
    unchecked:     bool,
}

#[derive(Debug, Clone, Copy)]
enum Pass {
    Ast,
    Rle,
    Bytecode,
    Peephole,
    #[cfg(feature = "jit")]
    Jit,
    #[cfg(feature = "llvm")]
    Llvm,
}

fn main() {
    let options = get_options();

    let program = parse(&options);

    match options.compiler_pass {
        Pass::Ast => {
            interpret(&*program, &options);
        }

        Pass::Rle => {
            let program = program.rle_compile();
            interpret(&*program, &options);
        }

        Pass::Peephole => {
            let program = program.peephole_compile();
            interpret(&*program, &options);
        }

        Pass::Bytecode => {
            let program = program.bytecode_compile();
            interpret(&*program, &options);
        }

        #[cfg(feature = "jit")]
        Pass::Jit => {
            let program = program.jit_compile(!options.unchecked);
            interpret(&program, &options);
        }

        #[cfg(feature = "llvm")]
        Pass::Llvm => {
            program.llvm_run(options.memory_size)
                .unwrap_or_else(|e| error_exit(3, &format!("runtime error: {}.", e)));
        }
    }
}

fn parse(options: &Options) -> Box<ast::Program> {
    ast::parse_program(&options.program_text)
        .unwrap_or_else(|e| error_exit(2, &format!("syntax error: {}.", e)))
}

fn interpret<P: Interpretable + ?Sized>(program: &P, options: &Options) {
    program.interpret_stdin(options.memory_size)
        .unwrap_or_else(|e| error_exit(3, &format!("runtime error: {}.", e)))
}

#[cfg(feature = "jit")]
const DEFAULT_PASS: Pass = Pass::Jit;

#[cfg(not(feature = "jit"))]
const DEFAULT_PASS: Pass = Pass::Peephole;

fn get_options() -> Options {
    let mut result = Options {
        program_text:  Vec::new(),
        memory_size:   None,
        compiler_pass: DEFAULT_PASS,
        unchecked:     false,
    };

    let matches = build_clap_app().get_matches();

    if let Some(size) = matches.value_of("size") {
        let size = size.parse()
            .unwrap_or_else(|e|
                error_exit(1, &format!("error: could not parse memory size: {}.", e)));
        if size == 0 {
            error_exit(1, "error: memory size must be at least 1.");
        }
        result.memory_size = Some(size);
    }

    if matches.is_present("jit") {
        #[cfg(feature = "jit")]
        let _ = result.compiler_pass = Pass::Jit;
    } else if matches.is_present("llvm") {
        #[cfg(feature = "llvm")]
        let _ = result.compiler_pass = Pass::Llvm;
    } else if matches.is_present("byte") {
        result.compiler_pass = Pass::Bytecode;
    } else if matches.is_present("peep") {
        result.compiler_pass = Pass::Peephole;
    } else if matches.is_present("rle") {
        result.compiler_pass = Pass::Rle;
    } else if matches.is_present("ast") {
        result.compiler_pass = Pass::Ast;
    }

    if matches.is_present("unchecked") {
        result.unchecked = true;
    }

    if let Some(exprs) = matches.values_of("expr") {
        for e in exprs {
            result.program_text.extend(e.as_bytes());
        }
    } else if let Some(files) = matches.values_of("FILE") {
        for f in files {
            let mut file = File::open(f)
                .unwrap_or_else(|e| error_exit(1, &format!("{}: ‘{}’.", e, f)));
            file.read_to_end(&mut result.program_text)
                .unwrap_or_else(|e| error_exit(1, &format!("{}: ‘{}’.", e, f)));
        }
    } else {
        error_exit(1, "error: no program given.");
    }

    result
}

fn build_clap_app() -> App<'static, 'static> {
    let app = App::new("bfi")
        .version(VERSION)
        .author("Jesse A. Tov <jesse.tov@gmail.com>")
        .about("A Brainfuck interpreter")
        .arg(Arg::with_name("expr")
            .short("e")
            .long("expr")
            .value_name("CODE")
            .help("BF code to execute")
            .multiple(true)
            .takes_value(true)
            .conflicts_with("FILE"))
        .arg(Arg::with_name("FILE")
            .help("The source file(s) to interpret")
            .multiple(true)
            .conflicts_with("expr")
            .index(1))
        .arg(Arg::with_name("size")
            .short("s")
            .long("size")
            .value_name("SIZE")
            .help("Memory size in bytes (default 30,000)")
            .takes_value(true))
        .arg(Arg::with_name("ast")
            .long("ast")
            .help("Interpret the unoptimized AST")
            .conflicts_with_all(&["rle", "peep", "byte", "jit", "llvm"]))
        .arg(Arg::with_name("rle")
            .long("rle")
            .help("Interpret the run-length encoded the AST")
            .conflicts_with_all(&["ast", "peep", "byte", "jit", "llvm"]))
        .arg(Arg::with_name("peep")
            .long("peep")
            .help(
                if cfg!(feature = "jit") {
                    "Interpret the peephole-optimized AST"
                } else {
                    "Interpret the peephole-optimized AST (default)"
                })
            .conflicts_with_all(&["ast", "rle", "byte", "jit", "llvm"]))
        .arg(Arg::with_name("byte")
            .long("byte")
            .help("Compile AST to bytecode")
            .conflicts_with_all(&["ast", "rle", "peep", "jit", "llvm"]));

    #[cfg(feature = "llvm")]
    let app = app
        .arg(Arg::with_name("llvm")
            .long("llvm")
            .help("JIT using LLVM")
            .conflicts_with_all(&["ast", "rle", "peep", "byte", "jit"]));

    #[cfg(feature = "jit")]
    let app = app
        .arg(Arg::with_name("jit")
            .long("jit")
            .help("JIT to native x64 (default)")
            .conflicts_with_all(&["ast", "rle", "peep", "byte", "llvm"]));

    #[cfg(any(feature = "jit"))]
    let app = app
        .arg(Arg::with_name("unchecked")
            .short("u")
            .long("unchecked")
            .help("Omit memory bounds checks in JIT")
            .conflicts_with_all(&["ast", "rle", "peep", "byte", "llvm"]));

    app
}

fn error_exit(code: i32, msg: &str) -> ! {
    eprintln!("bfi: {}", msg);
    exit(code)
}

<|endoftext|>


<commit_before>
extern crate cdrs;

use cdrs::client::CDRS;
use cdrs::query::QueryBuilder;
use cdrs::authenticators::NoneAuthenticator;
use cdrs::compression::Compression;
use cdrs::consistency::Consistency;
use cdrs::transport::TransportPlain;
// default credentials
const _ADDR: &'static str = "127.0.0.1:9042";

#[test]
fn it_works() {
    const _ADDR: &'static str = "127.0.0.1:9042";
    let authenticator = NoneAuthenticator;
    let tcp_transport = TransportPlain::new(_ADDR);
    assert_eq!(tcp_transport.is_ok(), true);

    let client = CDRS::new(tcp_transport.unwrap(), authenticator);
    let mut session = client.start(Compression::None).unwrap();
    let drop_ks = "DROP KEYSPACE my_ks;";
    let create_ks_cql = "CREATE KEYSPACE my_ks WITH REPLICATION = { 'class' : 'SimpleStrategy', 'replication_factor' : 1 } ;";
    let create_table_cql = "CREATE TABLE my_ks.users (
        user_name varchar PRIMARY KEY,
        password varchar,
        gender varchar,
        session_token varchar,
        state varchar,
        birth_year bigint
    );";

    let with_tracing = false;
    let with_warnings = false;

    let drop_ks_query = QueryBuilder::new(drop_ks)
        .consistency(Consistency::One)
        .finalize();

    let drop_ks_query_result =  session.query(drop_ks_query, with_tracing, with_warnings) ;

    assert_eq!(drop_ks_query_result.is_ok(), true);

    match drop_ks_query_result {
        Ok(ref res) => println!("keyspace dropped: {:?}", res.get_body()),
        Err(ref err) => println!("Error occured: {:?}", err)
    }

    let create_ks_query = QueryBuilder::new(create_ks_cql)
        .consistency(Consistency::One)
        .finalize();

    let create_ks_query_result =  session.query(create_ks_query, with_tracing, with_warnings);

    assert_eq!(create_ks_query_result.is_ok(), true);

    match create_ks_query_result {
        Ok(ref res) => println!("keyspace created: {:?}", res.get_body()),
        Err(ref err) => println!("Error occured: {:?}", err)
    }

    let create_table_query = QueryBuilder::new(create_table_cql)
        .consistency(Consistency::One)
        .finalize();
    match session.query(create_table_query, with_tracing, with_warnings) {
        Ok(ref res) => println!("table created: {:?}", res.get_body()),
        Err(ref err) => println!("Error occured: {:?}", err)
    }
}
<commit_msg>drop added after create<commit_after>
extern crate cdrs;

use cdrs::client::CDRS;
use cdrs::query::QueryBuilder;
use cdrs::authenticators::NoneAuthenticator;
use cdrs::compression::Compression;
use cdrs::consistency::Consistency;
use cdrs::transport::TransportPlain;
// default credentials
const _ADDR: &'static str = "127.0.0.1:9042";

#[test]
fn it_works() {
    const _ADDR: &'static str = "127.0.0.1:9042";
    let authenticator = NoneAuthenticator;
    let tcp_transport = TransportPlain::new(_ADDR);
    assert_eq!(tcp_transport.is_ok(), true);

    let client = CDRS::new(tcp_transport.unwrap(), authenticator);
    let mut session = client.start(Compression::None).unwrap();
    let drop_ks = "DROP KEYSPACE my_ks;";
    let create_ks_cql = "CREATE KEYSPACE my_ks WITH REPLICATION = { 'class' : 'SimpleStrategy', 'replication_factor' : 1 } ;";
    let create_table_cql = "CREATE TABLE my_ks.users (
        user_name varchar PRIMARY KEY,
        password varchar,
        gender varchar,
        session_token varchar,
        state varchar,
        birth_year bigint
    );";

    let with_tracing = false;
    let with_warnings = false;

    let create_ks_query = QueryBuilder::new(create_ks_cql)
        .consistency(Consistency::One)
        .finalize();

    let create_ks_query_result =  session.query(create_ks_query, with_tracing, with_warnings);

    assert_eq!(create_ks_query_result.is_ok(), true);

    match create_ks_query_result {
        Ok(ref res) => println!("keyspace created: {:?}", res.get_body()),
        Err(ref err) => println!("Error occured: {:?}", err)
    }

    let create_table_query = QueryBuilder::new(create_table_cql)
        .consistency(Consistency::One)
        .finalize();
    match session.query(create_table_query, with_tracing, with_warnings) {
        Ok(ref res) => println!("table created: {:?}", res.get_body()),
        Err(ref err) => println!("Error occured: {:?}", err)
    }

    let drop_ks_query = QueryBuilder::new(drop_ks)
        .consistency(Consistency::One)
        .finalize();

    let drop_ks_query_result =  session.query(drop_ks_query, with_tracing, with_warnings) ;

    assert_eq!(drop_ks_query_result.is_ok(), true);

    match drop_ks_query_result {
        Ok(ref res) => println!("keyspace dropped: {:?}", res.get_body()),
        Err(ref err) => println!("Error occured: {:?}", err)
    }
}
<|endoftext|>


<commit_before>extern crate core;
#[macro_use]
extern crate error_chain;
extern crate image;
#[macro_use]
extern crate lazy_static;
extern crate pest;
#[macro_use]
extern crate pest_derive;
extern crate regex;
#[macro_use]
extern crate structopt;

mod commands;
mod basm;
mod beast;
mod mnemonic;

// use clap::{App, AppSettings, Arg, SubCommand};
use commands::{Lang, PackingType};
use core::error::*;
use std::path::PathBuf;
use structopt::StructOpt;

fn main() {
    if let Err(ref e) = run() {
        println!("error: {}", e);

        for e in e.iter().skip(1) {
            println!("caused by: {}", e);
        }

        // The backtrace is not always generated. Try to run this example
        // with `RUST_BACKTRACE=1`.
        if let Some(backtrace) = e.backtrace() {
            println!("backtrace: {:?}", backtrace);
        }

        ::std::process::exit(1);
    }
}

#[derive(StructOpt, Debug)]
#[structopt(name = "hudson", about = "The bakervm toolkit")]
enum Opt {
    #[structopt(name = "compile",
                about = "compiles a compatible source file into a bakerVM image", alias = "c")]
    Compile {
        #[structopt(short = "l")]
        lang: Option<Lang>,
        #[structopt(short = "o", parse(from_os_str))]
        output: Option<PathBuf>,
        #[structopt(parse(from_os_str))]
        input: PathBuf,
    },
    #[structopt(name = "pack", about = "Write texture functions from images", alias = "p")]
    Pack {
        #[structopt(long = "type", short = "t", value_name = "type")]
        packing_type: Option<PackingType>,
        #[structopt(short = "o", parse(from_os_str))]
        output: Option<PathBuf>,
        #[structopt(parse(from_os_str))]
        input: PathBuf,
    },
}

fn run() -> Result<()> {
    let opt = Opt::from_args();

    match opt {
        Opt::Compile {
            lang,
            input,
            output,
        } => commands::compile(lang, input, output)?,
        Opt::Pack {
            packing_type,
            input,
            output,
        } => commands::pack(packing_type, input, output)?,
    }

    Ok(())
}
<commit_msg>Adjusted case<commit_after>extern crate core;
#[macro_use]
extern crate error_chain;
extern crate image;
#[macro_use]
extern crate lazy_static;
extern crate pest;
#[macro_use]
extern crate pest_derive;
extern crate regex;
#[macro_use]
extern crate structopt;

mod commands;
mod basm;
mod beast;
mod mnemonic;

// use clap::{App, AppSettings, Arg, SubCommand};
use commands::{Lang, PackingType};
use core::error::*;
use std::path::PathBuf;
use structopt::StructOpt;

fn main() {
    if let Err(ref e) = run() {
        println!("error: {}", e);

        for e in e.iter().skip(1) {
            println!("caused by: {}", e);
        }

        // The backtrace is not always generated. Try to run this example
        // with `RUST_BACKTRACE=1`.
        if let Some(backtrace) = e.backtrace() {
            println!("backtrace: {:?}", backtrace);
        }

        ::std::process::exit(1);
    }
}

#[derive(StructOpt, Debug)]
#[structopt(name = "hudson", about = "The bakerVM toolkit")]
enum Opt {
    #[structopt(name = "compile",
                about = "compiles a compatible source file into a bakerVM image", alias = "c")]
    Compile {
        #[structopt(short = "l")]
        lang: Option<Lang>,
        #[structopt(short = "o", parse(from_os_str))]
        output: Option<PathBuf>,
        #[structopt(parse(from_os_str))]
        input: PathBuf,
    },
    #[structopt(name = "pack", about = "Write texture functions from images", alias = "p")]
    Pack {
        #[structopt(long = "type", short = "t", value_name = "type")]
        packing_type: Option<PackingType>,
        #[structopt(short = "o", parse(from_os_str))]
        output: Option<PathBuf>,
        #[structopt(parse(from_os_str))]
        input: PathBuf,
    },
}

fn run() -> Result<()> {
    let opt = Opt::from_args();

    match opt {
        Opt::Compile {
            lang,
            input,
            output,
        } => commands::compile(lang, input, output)?,
        Opt::Pack {
            packing_type,
            input,
            output,
        } => commands::pack(packing_type, input, output)?,
    }

    Ok(())
}
<|endoftext|>


<commit_before>use latch::{LatchProbe, SpinLatch};
#[allow(unused_imports)]
use log::Event::*;
use job::StackJob;
use registry::{self, WorkerThread};
use std::any::Any;
use unwind;

#[cfg(test)]
mod test;

/// The `join` function takes two closures and *potentially* runs them
/// in parallel. It returns a pair of the results from those closures.
///
/// Conceptually, calling `join()` is similar to spawning two threads,
/// one executing each of the two closures. However, the
/// implementation is quite different and incurs very low
/// overhead. The underlying technique is called "work stealing": the
/// Rayon runtime uses a fixed pool of worker threads and attempts to
/// only execute code in parallel when there are idle CPUs to handle
/// it.
///
/// ### Warning about blocking I/O
///
/// The assumption is that the closures given to `join()` are
/// CPU-bound tasks that to do not perform I/O or other blocking
/// operations. If you do perform I/O, and that I/O should block
/// (e.g., waiting for a network request), the overall performance may
/// be poor.  Moreover, if you cause one closure to be blocked waiting
/// on another (for example, using a channel), that could lead to a
/// deadlock.
///
/// ### Panics
///
/// No matter what happens, both closures will always be executed.  If
/// a single closure panics, whether it be the first or second
/// closure, that panic will be propagated and hence `join()` will
/// panic with the same panic value. If both closures panic, `join()`
/// will panic with the panic value from the first closure.
pub fn join<A, B, RA, RB>(oper_a: A, oper_b: B) -> (RA, RB)
    where A: FnOnce() -> RA + Send,
          B: FnOnce() -> RB + Send,
          RA: Send,
          RB: Send
{
    registry::in_worker(|worker_thread| unsafe {
        log!(Join { worker: worker_thread.index() });

        // Create virtual wrapper for task b; this all has to be
        // done here so that the stack frame can keep it all live
        // long enough.
        let job_b = StackJob::new(oper_b, SpinLatch::new());
        let job_b_ref = job_b.as_job_ref();
        worker_thread.push(job_b_ref);

        // Execute task a; hopefully b gets stolen in the meantime.
        let result_a = match unwind::halt_unwinding(oper_a) {
            Ok(v) => v,
            Err(err) => join_recover_from_panic(worker_thread, &job_b.latch, err),
        };

        // Now that task A has finished, try to pop job B from the
        // local stack.  It may already have been popped by job A; it
        // may also have been stolen. There may also be some tasks
        // pushed on top of it in the stack, and we will have to pop
        // those off to get to it.
        while !job_b.latch.probe() {
            if let Some(job) = worker_thread.take_local_job() {
                if job == job_b_ref {
                    // Found it! Let's run it.
                    //
                    // Note that this could panic, but it's ok if we unwind here.
                    log!(PoppedRhs { worker: worker_thread.index() });
                    let result_b = job_b.run_inline();
                    return (result_a, result_b);
                } else {
                    log!(PoppedJob { worker: worker_thread.index() });
                    worker_thread.execute(job);
                }
            } else {
                // Local deque is empty. Time to steal from other
                // threads.
                log!(LostJob { worker: worker_thread.index() });
                worker_thread.wait_until(&job_b.latch);
                debug_assert!(job_b.latch.probe());
                break;
            }
        }

        return (result_a, job_b.into_result());
    })
}

/// If job A panics, we still cannot return until we are sure that job
/// B is complete. This is because it may contain references into the
/// enclosing stack frame(s).
#[cold] // cold path
unsafe fn join_recover_from_panic(worker_thread: &WorkerThread,
                                  job_b_latch: &SpinLatch,
                                  err: Box<Any + Send>)
                                  -> !
{
    worker_thread.wait_until(job_b_latch);
    unwind::resume_unwinding(err)
}
<commit_msg>Fix a typo<commit_after>use latch::{LatchProbe, SpinLatch};
#[allow(unused_imports)]
use log::Event::*;
use job::StackJob;
use registry::{self, WorkerThread};
use std::any::Any;
use unwind;

#[cfg(test)]
mod test;

/// The `join` function takes two closures and *potentially* runs them
/// in parallel. It returns a pair of the results from those closures.
///
/// Conceptually, calling `join()` is similar to spawning two threads,
/// one executing each of the two closures. However, the
/// implementation is quite different and incurs very low
/// overhead. The underlying technique is called "work stealing": the
/// Rayon runtime uses a fixed pool of worker threads and attempts to
/// only execute code in parallel when there are idle CPUs to handle
/// it.
///
/// ### Warning about blocking I/O
///
/// The assumption is that the closures given to `join()` are
/// CPU-bound tasks that do not perform I/O or other blocking
/// operations. If you do perform I/O, and that I/O should block
/// (e.g., waiting for a network request), the overall performance may
/// be poor.  Moreover, if you cause one closure to be blocked waiting
/// on another (for example, using a channel), that could lead to a
/// deadlock.
///
/// ### Panics
///
/// No matter what happens, both closures will always be executed.  If
/// a single closure panics, whether it be the first or second
/// closure, that panic will be propagated and hence `join()` will
/// panic with the same panic value. If both closures panic, `join()`
/// will panic with the panic value from the first closure.
pub fn join<A, B, RA, RB>(oper_a: A, oper_b: B) -> (RA, RB)
    where A: FnOnce() -> RA + Send,
          B: FnOnce() -> RB + Send,
          RA: Send,
          RB: Send
{
    registry::in_worker(|worker_thread| unsafe {
        log!(Join { worker: worker_thread.index() });

        // Create virtual wrapper for task b; this all has to be
        // done here so that the stack frame can keep it all live
        // long enough.
        let job_b = StackJob::new(oper_b, SpinLatch::new());
        let job_b_ref = job_b.as_job_ref();
        worker_thread.push(job_b_ref);

        // Execute task a; hopefully b gets stolen in the meantime.
        let result_a = match unwind::halt_unwinding(oper_a) {
            Ok(v) => v,
            Err(err) => join_recover_from_panic(worker_thread, &job_b.latch, err),
        };

        // Now that task A has finished, try to pop job B from the
        // local stack.  It may already have been popped by job A; it
        // may also have been stolen. There may also be some tasks
        // pushed on top of it in the stack, and we will have to pop
        // those off to get to it.
        while !job_b.latch.probe() {
            if let Some(job) = worker_thread.take_local_job() {
                if job == job_b_ref {
                    // Found it! Let's run it.
                    //
                    // Note that this could panic, but it's ok if we unwind here.
                    log!(PoppedRhs { worker: worker_thread.index() });
                    let result_b = job_b.run_inline();
                    return (result_a, result_b);
                } else {
                    log!(PoppedJob { worker: worker_thread.index() });
                    worker_thread.execute(job);
                }
            } else {
                // Local deque is empty. Time to steal from other
                // threads.
                log!(LostJob { worker: worker_thread.index() });
                worker_thread.wait_until(&job_b.latch);
                debug_assert!(job_b.latch.probe());
                break;
            }
        }

        return (result_a, job_b.into_result());
    })
}

/// If job A panics, we still cannot return until we are sure that job
/// B is complete. This is because it may contain references into the
/// enclosing stack frame(s).
#[cold] // cold path
unsafe fn join_recover_from_panic(worker_thread: &WorkerThread,
                                  job_b_latch: &SpinLatch,
                                  err: Box<Any + Send>)
                                  -> !
{
    worker_thread.wait_until(job_b_latch);
    unwind::resume_unwinding(err)
}
<|endoftext|>


<commit_before>mod common;

use std::env::{self, VarError};
use std::mem::size_of;
use std::process::{exit, Child, Command, ExitStatus};
use std::thread;
use std::time::Duration;

use rand::Rng;

use sled::Config;

use common::cleanup;

const TEST_ENV_VAR: &str = "SLED_CRASH_TEST";
const N_TESTS: usize = 100;
const CYCLE: usize = 256;

// test names, also used as dir names
const RECOVERY_WITH_SNAPSHOT: &str = "crash_recovery_with_runtime_snapshot";
const RECOVERY_NO_SNAPSHOT: &str = "crash_recovery_no_runtime_snapshot";

fn main() {
    match env::var(TEST_ENV_VAR) {
        Err(VarError::NotPresent) => {
            test_crash_recovery_with_runtime_snapshot();
            test_crash_recovery_no_runtime_snapshot();
        }

        Ok(ref s) if s == RECOVERY_WITH_SNAPSHOT => run_with_snapshot(s),
        Ok(ref s) if s == RECOVERY_NO_SNAPSHOT => run_without_snapshot(s),

        Ok(_) | Err(_) => panic!("invalid crash test case"),
    }
}

/// Verifies that the keys in the tree are correctly recovered.
/// Panics if they are incorrect.
/// Returns the key that should be resumed at, and the current cycle value.
fn verify(tree: &sled::Tree) -> (u32, u32) {
    // key 0 should always be the highest value, as that's where we increment
    // at some point, it might go down by one
    // it should never return, or go down again after that
    let mut iter = tree.iter();
    let highest = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return (0, 0),
    };

    let highest_vec = u32_to_vec(highest);

    // find how far we got
    let mut contiguous: u32 = 0;
    let mut lowest = 0;
    for res in iter {
        let (_k, v) = res.unwrap();
        if &v[..4] == &highest_vec[..4] {
            contiguous += 1;
        } else {
            let expected = if highest == 0 {
                CYCLE as u32 - 1
            } else {
                (highest - 1) % CYCLE as u32
            };
            let actual = slice_to_u32(&*v);
            assert_eq!(expected, actual);
            lowest = actual;
            break;
        }
    }

    // ensure nothing changes after this point
    let low_beginning = u32_to_vec(contiguous + 1);

    for res in tree.range(&*low_beginning..) {
        let (k, v): (sled::IVec, _) = res.unwrap();
        assert_eq!(
            slice_to_u32(&*v),
            lowest,
            "expected key {} to have value {}, instead it had value {} in db: {:?}",
            slice_to_u32(&*k),
            lowest,
            slice_to_u32(&*v),
            tree
        );
    }

    (contiguous, highest)
}

fn u32_to_vec(u: u32) -> Vec<u8> {
    let buf: [u8; size_of::<u32>()] = u.to_be_bytes();
    buf.to_vec()
}

fn slice_to_u32(b: &[u8]) -> u32 {
    let mut buf = [0u8; size_of::<u32>()];
    buf.copy_from_slice(&b[..size_of::<u32>()]);

    u32::from_be_bytes(buf)
}

fn spawn_killah() {
    thread::spawn(|| {
        let runtime = rand::thread_rng().gen_range(0, 200);
        thread::sleep(Duration::from_millis(runtime));
        exit(9);
    });
}

fn run(config: Config) {
    common::setup_logger();

    let crash_during_initialization = rand::thread_rng().gen_bool(0.1);

    if crash_during_initialization {
        spawn_killah();
    }

    let tree = config.open().unwrap();

    if !crash_during_initialization {
        spawn_killah();
    }

    let (key, highest) = verify(&tree);

    let mut hu = ((highest as usize) * CYCLE) + key as usize;
    assert_eq!(hu % CYCLE, key as usize);
    assert_eq!(hu / CYCLE, highest as usize);

    loop {
        hu += 1;

        if hu / CYCLE >= CYCLE {
            hu = 0;
        }

        let key = u32_to_vec((hu % CYCLE) as u32);

        let mut value = u32_to_vec((hu / CYCLE) as u32);
        let additional_len = rand::thread_rng().gen_range(0, 1000);
        value.append(&mut vec![0u8; additional_len]);

        tree.insert(&key, value).unwrap();
    }
}

fn run_with_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(5000)
        .segment_size(1024);

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_without_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(1 << 56)
        .segment_size(1024);

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_child_process(test_name: &str) -> Child {
    let bin = env::current_exe().expect("could not get test binary path");

    env::set_var(TEST_ENV_VAR, test_name);

    Command::new(bin).env(TEST_ENV_VAR, test_name).spawn().expect(&format!(
        "could not spawn child process for {} test",
        test_name
    ))
}

fn handle_child_exit_status(dir: &str, status: ExitStatus) {
    let code = status.code();

    if code.is_none() || code.unwrap() != 9 {
        cleanup(dir);
        panic!("{} test child exited abnormally", dir);
    }
}

fn handle_child_wait_err(dir: &str, e: std::io::Error) {
    cleanup(dir);

    panic!("error waiting for {} test child: {}", dir, e);
}

fn test_crash_recovery_with_runtime_snapshot() {
    let dir = RECOVERY_WITH_SNAPSHOT;
    cleanup(dir);

    for _ in 0..N_TESTS {
        let mut child = run_child_process(RECOVERY_WITH_SNAPSHOT);

        child.wait().map(|status| handle_child_exit_status(dir, status));
    }

    cleanup(dir);
}

fn test_crash_recovery_no_runtime_snapshot() {
    let dir = RECOVERY_NO_SNAPSHOT;
    cleanup(dir);

    for _ in 0..N_TESTS {
        let mut child = run_child_process(RECOVERY_NO_SNAPSHOT);

        child
            .wait()
            .map(|status| handle_child_exit_status(dir, status))
            .map_err(|e| handle_child_wait_err(dir, e));
    }

    cleanup(dir);
}
<commit_msg>Handle with_snapshot test child wait error<commit_after>mod common;

use std::env::{self, VarError};
use std::mem::size_of;
use std::process::{exit, Child, Command, ExitStatus};
use std::thread;
use std::time::Duration;

use rand::Rng;

use sled::Config;

use common::cleanup;

const TEST_ENV_VAR: &str = "SLED_CRASH_TEST";
const N_TESTS: usize = 100;
const CYCLE: usize = 256;

// test names, also used as dir names
const RECOVERY_WITH_SNAPSHOT: &str = "crash_recovery_with_runtime_snapshot";
const RECOVERY_NO_SNAPSHOT: &str = "crash_recovery_no_runtime_snapshot";

fn main() {
    match env::var(TEST_ENV_VAR) {
        Err(VarError::NotPresent) => {
            test_crash_recovery_with_runtime_snapshot();
            test_crash_recovery_no_runtime_snapshot();
        }

        Ok(ref s) if s == RECOVERY_WITH_SNAPSHOT => run_with_snapshot(s),
        Ok(ref s) if s == RECOVERY_NO_SNAPSHOT => run_without_snapshot(s),

        Ok(_) | Err(_) => panic!("invalid crash test case"),
    }
}

/// Verifies that the keys in the tree are correctly recovered.
/// Panics if they are incorrect.
/// Returns the key that should be resumed at, and the current cycle value.
fn verify(tree: &sled::Tree) -> (u32, u32) {
    // key 0 should always be the highest value, as that's where we increment
    // at some point, it might go down by one
    // it should never return, or go down again after that
    let mut iter = tree.iter();
    let highest = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return (0, 0),
    };

    let highest_vec = u32_to_vec(highest);

    // find how far we got
    let mut contiguous: u32 = 0;
    let mut lowest = 0;
    for res in iter {
        let (_k, v) = res.unwrap();
        if &v[..4] == &highest_vec[..4] {
            contiguous += 1;
        } else {
            let expected = if highest == 0 {
                CYCLE as u32 - 1
            } else {
                (highest - 1) % CYCLE as u32
            };
            let actual = slice_to_u32(&*v);
            assert_eq!(expected, actual);
            lowest = actual;
            break;
        }
    }

    // ensure nothing changes after this point
    let low_beginning = u32_to_vec(contiguous + 1);

    for res in tree.range(&*low_beginning..) {
        let (k, v): (sled::IVec, _) = res.unwrap();
        assert_eq!(
            slice_to_u32(&*v),
            lowest,
            "expected key {} to have value {}, instead it had value {} in db: {:?}",
            slice_to_u32(&*k),
            lowest,
            slice_to_u32(&*v),
            tree
        );
    }

    (contiguous, highest)
}

fn u32_to_vec(u: u32) -> Vec<u8> {
    let buf: [u8; size_of::<u32>()] = u.to_be_bytes();
    buf.to_vec()
}

fn slice_to_u32(b: &[u8]) -> u32 {
    let mut buf = [0u8; size_of::<u32>()];
    buf.copy_from_slice(&b[..size_of::<u32>()]);

    u32::from_be_bytes(buf)
}

fn spawn_killah() {
    thread::spawn(|| {
        let runtime = rand::thread_rng().gen_range(0, 200);
        thread::sleep(Duration::from_millis(runtime));
        exit(9);
    });
}

fn run(config: Config) {
    common::setup_logger();

    let crash_during_initialization = rand::thread_rng().gen_bool(0.1);

    if crash_during_initialization {
        spawn_killah();
    }

    let tree = config.open().unwrap();

    if !crash_during_initialization {
        spawn_killah();
    }

    let (key, highest) = verify(&tree);

    let mut hu = ((highest as usize) * CYCLE) + key as usize;
    assert_eq!(hu % CYCLE, key as usize);
    assert_eq!(hu / CYCLE, highest as usize);

    loop {
        hu += 1;

        if hu / CYCLE >= CYCLE {
            hu = 0;
        }

        let key = u32_to_vec((hu % CYCLE) as u32);

        let mut value = u32_to_vec((hu / CYCLE) as u32);
        let additional_len = rand::thread_rng().gen_range(0, 1000);
        value.append(&mut vec![0u8; additional_len]);

        tree.insert(&key, value).unwrap();
    }
}

fn run_with_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(5000)
        .segment_size(1024);

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_without_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(1 << 56)
        .segment_size(1024);

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_child_process(test_name: &str) -> Child {
    let bin = env::current_exe().expect("could not get test binary path");

    env::set_var(TEST_ENV_VAR, test_name);

    Command::new(bin).env(TEST_ENV_VAR, test_name).spawn().expect(&format!(
        "could not spawn child process for {} test",
        test_name
    ))
}

fn handle_child_exit_status(dir: &str, status: ExitStatus) {
    let code = status.code();

    if code.is_none() || code.unwrap() != 9 {
        cleanup(dir);
        panic!("{} test child exited abnormally", dir);
    }
}

fn handle_child_wait_err(dir: &str, e: std::io::Error) {
    cleanup(dir);

    panic!("error waiting for {} test child: {}", dir, e);
}

fn test_crash_recovery_with_runtime_snapshot() {
    let dir = RECOVERY_WITH_SNAPSHOT;
    cleanup(dir);

    for _ in 0..N_TESTS {
        let mut child = run_child_process(RECOVERY_WITH_SNAPSHOT);

        child
            .wait()
            .map(|status| handle_child_exit_status(dir, status))
            .map_err(|e| handle_child_wait_err(dir, e))
            .unwrap();
    }

    cleanup(dir);
}

fn test_crash_recovery_no_runtime_snapshot() {
    let dir = RECOVERY_NO_SNAPSHOT;
    cleanup(dir);

    for _ in 0..N_TESTS {
        let mut child = run_child_process(RECOVERY_NO_SNAPSHOT);

        child
            .wait()
            .map(|status| handle_child_exit_status(dir, status))
            .map_err(|e| handle_child_wait_err(dir, e));
    }

    cleanup(dir);
}
<|endoftext|>