<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

//! Defines the merge plan for executing partitions in parallel and then merging the results
//! into a single partition

use std::any::Any;
use std::sync::Arc;

use futures::channel::mpsc;
use futures::sink::SinkExt;
use futures::stream::StreamExt;
use futures::Stream;

use async_trait::async_trait;

use arrow::record_batch::RecordBatch;
use arrow::{
    datatypes::SchemaRef,
    error::{ArrowError, Result as ArrowResult},
};

use super::RecordBatchStream;
use crate::error::{DataFusionError, Result};
use crate::physical_plan::ExecutionPlan;
use crate::physical_plan::Partitioning;

use super::SendableRecordBatchStream;
use pin_project_lite::pin_project;

/// Merge execution plan executes partitions in parallel and combines them into a single
/// partition. No guarantees are made about the order of the resulting partition.
#[derive(Debug)]
pub struct MergeExec {
    /// Input execution plan
    input: Arc<dyn ExecutionPlan>,
}

impl MergeExec {
    /// Create a new MergeExec
    pub fn new(input: Arc<dyn ExecutionPlan>) -> Self {
        MergeExec { input }
    }

    /// Input execution plan
    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
        &self.input
    }
}

#[async_trait]
impl ExecutionPlan for MergeExec {
    /// Return a reference to Any that can be used for downcasting
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn schema(&self) -> SchemaRef {
        self.input.schema()
    }

    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
        vec![self.input.clone()]
    }

    /// Get the output partitioning of this plan
    fn output_partitioning(&self) -> Partitioning {
        Partitioning::UnknownPartitioning(1)
    }

    fn with_new_children(
        &self,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        match children.len() {
            1 => Ok(Arc::new(MergeExec::new(children[0].clone()))),
            _ => Err(DataFusionError::Internal(
                "MergeExec wrong number of children".to_string(),
            )),
        }
    }

    async fn execute(&self, partition: usize) -> Result<SendableRecordBatchStream> {
        // MergeExec produces a single partition
        if 0 != partition {
            return Err(DataFusionError::Internal(format!(
                "MergeExec invalid partition {}",
                partition
            )));
        }

        let input_partitions = self.input.output_partitioning().partition_count();
        match input_partitions {
            0 => Err(DataFusionError::Internal(
                "MergeExec requires at least one input partition".to_owned(),
            )),
            1 => {
                // bypass any threading if there is a single partition
                self.input.execute(0).await
            }
            _ => {
                // use a stream that allows each sender to put in at
                // least one result in an attempt to maximize
                // parallelism.
                let (sender, receiver) =
                    mpsc::channel::<ArrowResult<RecordBatch>>(input_partitions);

                // spawn independent tasks whose resulting streams (of batches)
                // are sent to the channel for consumption.
                for part_i in (0..input_partitions) {
                    let input = self.input.clone();
                    let mut sender = sender.clone();
                    tokio::spawn(async move {
                        let mut stream = match input.execute(part_i).await {
                            Err(e) => {
                                // If send fails, plan being torn
                                // down, no place to send the error
                                let arrow_error = ArrowError::ExternalError(Box::new(e));
                                sender.send(Err(arrow_error)).await.ok();
                                return;
                            }
                            Ok(stream) => stream,
                        };

                        while let Some(item) = stream.next().await {
                            // If send fails, plan being torn down,
                            // there is no place to send the error
                            sender.send(item).await.ok();
                        }
                    });
                }

                Ok(Box::pin(MergeStream {
                    input: receiver,
                    schema: self.schema(),
                }))
            }
        }
    }
}

pin_project! {
    struct MergeStream {
        schema: SchemaRef,
        #[pin]
        input: mpsc::Receiver<ArrowResult<RecordBatch>>,
    }
}

impl Stream for MergeStream {
    type Item = ArrowResult<RecordBatch>;

    fn poll_next(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        let this = self.project();
        this.input.poll_next(cx)
    }
}

impl RecordBatchStream for MergeStream {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use crate::physical_plan::common;
    use crate::physical_plan::csv::{CsvExec, CsvReadOptions};
    use crate::test;

    #[tokio::test]
    async fn merge() -> Result<()> {
        let schema = test::aggr_test_schema();

        let num_partitions = 4;
        let path =
            test::create_partitioned_csv("aggregate_test_100.csv", num_partitions)?;

        let csv =
            CsvExec::try_new(&path, CsvReadOptions::new().schema(&schema), None, 1024)?;

        // input should have 4 partitions
        assert_eq!(csv.output_partitioning().partition_count(), num_partitions);

        let merge = MergeExec::new(Arc::new(csv));

        // output of MergeExec should have a single partition
        assert_eq!(merge.output_partitioning().partition_count(), 1);

        // the result should contain 4 batches (one per input partition)
        let iter = merge.execute(0).await?;
        let batches = common::collect(iter).await?;
        assert_eq!(batches.len(), num_partitions);

        // there should be a total of 100 rows
        let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum();
        assert_eq!(row_count, 100);

        Ok(())
    }
}
<commit_msg>ARROW-11881: [Rust][DataFusion] Fix clippy lint<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

//! Defines the merge plan for executing partitions in parallel and then merging the results
//! into a single partition

use std::any::Any;
use std::sync::Arc;

use futures::channel::mpsc;
use futures::sink::SinkExt;
use futures::stream::StreamExt;
use futures::Stream;

use async_trait::async_trait;

use arrow::record_batch::RecordBatch;
use arrow::{
    datatypes::SchemaRef,
    error::{ArrowError, Result as ArrowResult},
};

use super::RecordBatchStream;
use crate::error::{DataFusionError, Result};
use crate::physical_plan::ExecutionPlan;
use crate::physical_plan::Partitioning;

use super::SendableRecordBatchStream;
use pin_project_lite::pin_project;

/// Merge execution plan executes partitions in parallel and combines them into a single
/// partition. No guarantees are made about the order of the resulting partition.
#[derive(Debug)]
pub struct MergeExec {
    /// Input execution plan
    input: Arc<dyn ExecutionPlan>,
}

impl MergeExec {
    /// Create a new MergeExec
    pub fn new(input: Arc<dyn ExecutionPlan>) -> Self {
        MergeExec { input }
    }

    /// Input execution plan
    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
        &self.input
    }
}

#[async_trait]
impl ExecutionPlan for MergeExec {
    /// Return a reference to Any that can be used for downcasting
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn schema(&self) -> SchemaRef {
        self.input.schema()
    }

    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
        vec![self.input.clone()]
    }

    /// Get the output partitioning of this plan
    fn output_partitioning(&self) -> Partitioning {
        Partitioning::UnknownPartitioning(1)
    }

    fn with_new_children(
        &self,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        match children.len() {
            1 => Ok(Arc::new(MergeExec::new(children[0].clone()))),
            _ => Err(DataFusionError::Internal(
                "MergeExec wrong number of children".to_string(),
            )),
        }
    }

    async fn execute(&self, partition: usize) -> Result<SendableRecordBatchStream> {
        // MergeExec produces a single partition
        if 0 != partition {
            return Err(DataFusionError::Internal(format!(
                "MergeExec invalid partition {}",
                partition
            )));
        }

        let input_partitions = self.input.output_partitioning().partition_count();
        match input_partitions {
            0 => Err(DataFusionError::Internal(
                "MergeExec requires at least one input partition".to_owned(),
            )),
            1 => {
                // bypass any threading if there is a single partition
                self.input.execute(0).await
            }
            _ => {
                // use a stream that allows each sender to put in at
                // least one result in an attempt to maximize
                // parallelism.
                let (sender, receiver) =
                    mpsc::channel::<ArrowResult<RecordBatch>>(input_partitions);

                // spawn independent tasks whose resulting streams (of batches)
                // are sent to the channel for consumption.
                for part_i in 0..input_partitions {
                    let input = self.input.clone();
                    let mut sender = sender.clone();
                    tokio::spawn(async move {
                        let mut stream = match input.execute(part_i).await {
                            Err(e) => {
                                // If send fails, plan being torn
                                // down, no place to send the error
                                let arrow_error = ArrowError::ExternalError(Box::new(e));
                                sender.send(Err(arrow_error)).await.ok();
                                return;
                            }
                            Ok(stream) => stream,
                        };

                        while let Some(item) = stream.next().await {
                            // If send fails, plan being torn down,
                            // there is no place to send the error
                            sender.send(item).await.ok();
                        }
                    });
                }

                Ok(Box::pin(MergeStream {
                    input: receiver,
                    schema: self.schema(),
                }))
            }
        }
    }
}

pin_project! {
    struct MergeStream {
        schema: SchemaRef,
        #[pin]
        input: mpsc::Receiver<ArrowResult<RecordBatch>>,
    }
}

impl Stream for MergeStream {
    type Item = ArrowResult<RecordBatch>;

    fn poll_next(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        let this = self.project();
        this.input.poll_next(cx)
    }
}

impl RecordBatchStream for MergeStream {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use crate::physical_plan::common;
    use crate::physical_plan::csv::{CsvExec, CsvReadOptions};
    use crate::test;

    #[tokio::test]
    async fn merge() -> Result<()> {
        let schema = test::aggr_test_schema();

        let num_partitions = 4;
        let path =
            test::create_partitioned_csv("aggregate_test_100.csv", num_partitions)?;

        let csv =
            CsvExec::try_new(&path, CsvReadOptions::new().schema(&schema), None, 1024)?;

        // input should have 4 partitions
        assert_eq!(csv.output_partitioning().partition_count(), num_partitions);

        let merge = MergeExec::new(Arc::new(csv));

        // output of MergeExec should have a single partition
        assert_eq!(merge.output_partitioning().partition_count(), 1);

        // the result should contain 4 batches (one per input partition)
        let iter = merge.execute(0).await?;
        let batches = common::collect(iter).await?;
        assert_eq!(batches.len(), num_partitions);

        // there should be a total of 100 rows
        let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum();
        assert_eq!(row_count, 100);

        Ok(())
    }
}
<|endoftext|>


<commit_before>use core::cmp::{Eq, Ord};
use iter::BaseIter;

pub trait PersistentMap<K: Copy Eq Ord, V: Copy> {
  pure fn get(k: K) -> @Option<V>;
  pure fn put(k: K, v: V) -> self;
  pure fn delete(k: K) -> self;
}

enum RBColor {
  Red,
  Black
}

pub enum RBMap<K: Copy Eq Ord, V: Copy> {
  //   color    left         k  value       right
  Tree(RBColor, @RBMap<K,V>, K, @Option<V>, @RBMap<K,V>),
  Leaf
}

pure fn RBMap<K: Copy Eq Ord, V: Copy>(key: K, value: V) -> @RBMap<K, V> {
  @Tree(Red, @Leaf, key, @Some(value), @Leaf)
}

impl<K: Copy Eq Ord, V: Copy> @RBMap<K, V> {
  pure fn modify(k: K, new_value: @Option<V>) -> @RBMap<K, V> {
    match self {
      @Leaf => @Tree(Red, self, k, new_value, self),
      @Tree(color, left, key, original_value, right) => {
        if (k.lt(&key)) {
          self.balanceLeft(color, left.modify(k, new_value), key, original_value, right)
        } else if (k == key) {
          @Tree(color, left, k, new_value, right)
        } else {
          self.balanceRight(color, left, key, original_value, right.modify(k, new_value))
        }
      }
    }
  }

  pure fn balanceLeft(c: RBColor, l: @RBMap<K,V>, k: K, v: @Option<V>, r: @RBMap<K,V>) -> @RBMap<K,V> {
    match (c, l, k, v, r) {
      (Black,@Tree(Red,@Tree(Red,a,xK,xV,b),yK,yV,c),zK,zV,d) =>
        @Tree(Red,@Tree(Black,a,xK,xV,b),yK,yV,@Tree(Black,c,zK,zV,d)),
      (c,a,xK,xV,b) =>
        @Tree(c,a,xK,xV,b)
    }
  }

  pure fn balanceRight(c: RBColor, l: @RBMap<K,V>, k: K, v: @Option<V>, r: @RBMap<K,V>) -> @RBMap<K,V> {
    match (c, l, k, v, r) {
      (Black, @Tree(Red, t1, k1, v1, t2), k0, v0, @Tree(Red, t3, k2, v2, t4)) =>
        @Tree(Red, @Tree(Black, t1, k1, v1, t2), k0, v0, @Tree(Black, t3, k2, v2, t4)),
      (c, l, k, v, @Tree(Red, t1, k1, v1, t2)) =>
        @Tree(c, @Tree(Red, l, k, v, t1), k1, v1, t2),
      (c,a,xK,xV,b) =>
        @Tree(c,a,xK,xV,b)
    }
  }
}

/**
 * A purely functional Left-Leaning Red-Black Tree.
 */
impl<K: Copy Eq Ord, V: Copy> @RBMap<K, V> : PersistentMap<K, V> {
  pure fn get(k: K) -> @Option<V> {
    match self {
      @Leaf => @None,
      @Tree(_, left, key, maybe_value, right) => {
        if (k.lt(&key)) {
          left.get(k)
        } else if (k.gt(&key)) {
          right.get(k)
        } else {
          maybe_value
        }
      }
    }
  }

  pure fn put(k: K, new_value: V) -> @RBMap<K, V> {
    self.modify(k, @Some(new_value))
  }

  pure fn delete(k: K) -> @RBMap<K, V> {
    self.modify(k, @None)
  }
}

impl<K: Copy Eq Ord, V: Copy> RBMap<K, V>: iter::BaseIter<(&K, &V)> {
  pure fn size_hint() -> Option<uint> {
    None
  }

  pure fn each(f: fn(&(&K, &V)) -> bool) {
    match self {
      Leaf => (),
      Tree(_, left, key, maybe_value, right) => {
        left.each(f);
        match maybe_value {
          @Some(value) => f(&(&key, &value)),
          @None => false,
        };
        right.each(f);
      }
    }
  }
}

#[test]
fn test_rb_tree() {
  let v1 = RBMap("stevej", 150);
  let v2 = v1.put("thatstacy", 187);

  assert(v2.get("stevej") == @Some(150));
  assert(v2.get("thatstacy") == @Some(187));

  let v3 = v2.delete("stevej");
  assert(v3.get("stevej") == @None);
  assert(v3.get("thatstacy") == @Some(187));

  let v4 = v2.put("jeremy", 16);

  let v5 = v4.put("ev", 20);
  let v6 = v5.put("zhana", 51573);

  assert(v6.get("jeremy") == @Some(16));
  assert(v6.get("stevej") == @Some(150));
  assert(v6.get("thatstacy") == @Some(187));
  assert(v6.get("ev") == @Some(20));
  assert(v6.get("zhana") == @Some(51573));
}

#[test]
fn test_base_iter_each() {
  let v1 = RBMap(1, 0);
  let v2 = v1.put(4, 0);
  let v3 = v2.put(3, 0);
  let v4 = v3.put(5, 0);
  let v5 = v4.put(2, 0);

  let n = @mut 1;

  fn t(n: @mut int, kv: &(&int, &int)) -> bool{
    match kv {
      &(k, _) => {
        assert (*n == *k);
        *n += 1;
      }
    }
    false
  }

  v5.each(|z| t(n, z));
}
<commit_msg>make the anonymous trait private<commit_after>use core::cmp::{Eq, Ord};
use iter::BaseIter;

pub trait PersistentMap<K: Copy Eq Ord, V: Copy> {
  pure fn get(k: K) -> @Option<V>;
  pure fn put(k: K, v: V) -> self;
  pure fn delete(k: K) -> self;
}

enum RBColor {
  Red,
  Black
}

pub enum RBMap<K: Copy Eq Ord, V: Copy> {
  //   color    left         k  value       right
  Tree(RBColor, @RBMap<K,V>, K, @Option<V>, @RBMap<K,V>),
  Leaf
}

pure fn RBMap<K: Copy Eq Ord, V: Copy>(key: K, value: V) -> @RBMap<K, V> {
  @Tree(Red, @Leaf, key, @Some(value), @Leaf)
}

priv impl<K: Copy Eq Ord, V: Copy> @RBMap<K, V> {
  pure fn modify(k: K, new_value: @Option<V>) -> @RBMap<K, V> {
    match self {
      @Leaf => @Tree(Red, self, k, new_value, self),
      @Tree(color, left, key, original_value, right) => {
        if (k.lt(&key)) {
          self.balanceLeft(color, left.modify(k, new_value), key, original_value, right)
        } else if (k == key) {
          @Tree(color, left, k, new_value, right)
        } else {
          self.balanceRight(color, left, key, original_value, right.modify(k, new_value))
        }
      }
    }
  }

  pure fn balanceLeft(c: RBColor, l: @RBMap<K,V>, k: K, v: @Option<V>, r: @RBMap<K,V>) -> @RBMap<K,V> {
    match (c, l, k, v, r) {
      (Black,@Tree(Red,@Tree(Red,a,xK,xV,b),yK,yV,c),zK,zV,d) =>
        @Tree(Red,@Tree(Black,a,xK,xV,b),yK,yV,@Tree(Black,c,zK,zV,d)),
      (c,a,xK,xV,b) =>
        @Tree(c,a,xK,xV,b)
    }
  }

  pure fn balanceRight(c: RBColor, l: @RBMap<K,V>, k: K, v: @Option<V>, r: @RBMap<K,V>) -> @RBMap<K,V> {
    match (c, l, k, v, r) {
      (Black, @Tree(Red, t1, k1, v1, t2), k0, v0, @Tree(Red, t3, k2, v2, t4)) =>
        @Tree(Red, @Tree(Black, t1, k1, v1, t2), k0, v0, @Tree(Black, t3, k2, v2, t4)),
      (c, l, k, v, @Tree(Red, t1, k1, v1, t2)) =>
        @Tree(c, @Tree(Red, l, k, v, t1), k1, v1, t2),
      (c,a,xK,xV,b) =>
        @Tree(c,a,xK,xV,b)
    }
  }
}

/**
 * A purely functional Left-Leaning Red-Black Tree.
 */
impl<K: Copy Eq Ord, V: Copy> @RBMap<K, V> : PersistentMap<K, V> {
  pure fn get(k: K) -> @Option<V> {
    match self {
      @Leaf => @None,
      @Tree(_, left, key, maybe_value, right) => {
        if (k.lt(&key)) {
          left.get(k)
        } else if (k.gt(&key)) {
          right.get(k)
        } else {
          maybe_value
        }
      }
    }
  }

  pure fn put(k: K, new_value: V) -> @RBMap<K, V> {
    self.modify(k, @Some(new_value))
  }

  pure fn delete(k: K) -> @RBMap<K, V> {
    self.modify(k, @None)
  }
}

impl<K: Copy Eq Ord, V: Copy> RBMap<K, V>: iter::BaseIter<(&K, &V)> {
  pure fn size_hint() -> Option<uint> {
    None
  }

  pure fn each(f: fn(&(&K, &V)) -> bool) {
    match self {
      Leaf => (),
      Tree(_, left, key, maybe_value, right) => {
        left.each(f);
        match maybe_value {
          @Some(value) => f(&(&key, &value)),
          @None => false,
        };
        right.each(f);
      }
    }
  }
}

#[test]
fn test_rb_tree() {
  let v1 = RBMap("stevej", 150);
  let v2 = v1.put("thatstacy", 187);

  assert(v2.get("stevej") == @Some(150));
  assert(v2.get("thatstacy") == @Some(187));

  let v3 = v2.delete("stevej");
  assert(v3.get("stevej") == @None);
  assert(v3.get("thatstacy") == @Some(187));

  let v4 = v2.put("jeremy", 16);

  let v5 = v4.put("ev", 20);
  let v6 = v5.put("zhana", 51573);

  assert(v6.get("jeremy") == @Some(16));
  assert(v6.get("stevej") == @Some(150));
  assert(v6.get("thatstacy") == @Some(187));
  assert(v6.get("ev") == @Some(20));
  assert(v6.get("zhana") == @Some(51573));
}

#[test]
fn test_base_iter_each() {
  let v1 = RBMap(1, 0);
  let v2 = v1.put(4, 0);
  let v3 = v2.put(3, 0);
  let v4 = v3.put(5, 0);
  let v5 = v4.put(2, 0);

  let n = @mut 1;

  fn t(n: @mut int, kv: &(&int, &int)) -> bool{
    match kv {
      &(k, _) => {
        assert (*n == *k);
        *n += 1;
      }
    }
    false
  }

  v5.each(|z| t(n, z));
}
<|endoftext|>


<commit_before>extern crate rand;

use rand::Rng;

use std::net::SocketAddr;

use messages::{Any, RawMessage, Connect, Status, Message, RequestPeers};
use events::Channel;
use super::{NodeHandler, RequestData, ExternalMessage, NodeTimeout, Height};

impl<S> NodeHandler<S>
    where S: Channel<ApplicationEvent = ExternalMessage, Timeout = NodeTimeout>
{
    pub fn handle_message(&mut self, raw: RawMessage) {
        // TODO: check message headers (network id, protocol version)
        // FIXME: call message.verify method
        //     if !raw.verify() {
        //         return;
        //     }

        //FIXME: add whitelist verify public_key

        let msg = Any::from_raw(raw).unwrap();
        match msg {
            Any::Connect(msg) => self.handle_connect(msg),
            Any::Status(msg) => self.handle_status(msg),
            Any::Consensus(msg) => self.handle_consensus(msg),
            Any::Request(msg) => self.handle_request(msg),
            Any::Block(msg) => self.handle_block(msg),
            Any::Transaction(msg) => self.handle_tx(msg),
        }
    }

    pub fn handle_connected(&mut self, addr: &SocketAddr) {
        info!("Connected to: {}", addr);
        let message = self.state.our_connect_message().clone();
        self.send_to_addr(addr, message.raw());
    }

    pub fn handle_disconnected(&mut self, addr: &SocketAddr) {
        info!("Disconnected from: {}", addr);
        let need_reconnect = self.state.remove_peer_with_addr(addr);
        if need_reconnect {
            self.connect(addr);
        }
    }

    pub fn handle_connect(&mut self, message: Connect) {
        // TODO add spam protection
        let address = message.addr();
        if address == self.state.our_connect_message().addr() {
            return;
        }
        // Check if we have another connect message from peer with the given public_key
        let public_key = *message.pub_key();
        let mut need_connect = true;
        if let Some(saved_message) = self.state.peers().get(&public_key) {
            if saved_message.time() > message.time() {
                error!("Received outdated Connect message from {}", address);
                return;
            } else if saved_message.time() < message.time() {
                need_connect = saved_message.addr() != message.addr();
            } else if saved_message.addr() != message.addr() {
                error!("Received weird Connect message from {}", address);
                return;
            }
        }
        info!("Received Connect message from {}, {}",
              address,
              need_connect);
        self.state.add_peer(public_key, message);
        if need_connect {
            // TODO: reduce double sending of connect message
            info!("Send Connect message to {}", address);
            self.connect(&address);
        }
    }

    pub fn handle_status(&mut self, msg: Status) {
        let height = self.state.height();
        trace!("HANDLE STATUS: current height = {}, msg height = {}", height, msg.height());

        // Handle message from future height
        if msg.height() > height {
            let peer = msg.from();

            //verify message
            if !msg.verify_signature(peer) {
                return;
            }

            // Check validator height info
            if msg.height() > self.state.node_height(peer) {
                // Update validator height
                self.state.set_node_height(*peer, msg.height());
            }

            // Request block
            self.request(RequestData::Block(height), *peer);
        }
    }

    pub fn handle_request_peers(&mut self, msg: RequestPeers) {
        let peers: Vec<Connect> = self.state.peers().iter().map(|(_, b)| b.clone()).collect();
        trace!("HANDLE REQUEST PEERS: Sending {:?} peers to {:?}", peers, msg.from());

        for peer in peers {
            self.send_to_peer(*msg.from(), peer.raw());
        }
    }

    pub fn handle_status_timeout(&mut self, height: Height) {
        if self.state.height() == height {
            self.broadcast_status();
            self.add_status_timeout();
        }
    }

    pub fn handle_peer_exchange_timeout(&mut self) {
        if !self.state.peers().is_empty() {
            let to = self.state.peers().len();
            let gen_peer_id = || -> usize {
                let mut rng = rand::thread_rng();
                rng.gen_range(0, to)
            };

            let peer = self.state
                .peers()
                .iter()
                .map(|x| x.1.clone())
                .nth(gen_peer_id())
                .unwrap();
            let peer = peer.clone();
            let msg = RequestPeers::new(self.state.public_key(),
                                        peer.pub_key(),
                                        self.state.secret_key());
            trace!("Request peers from peer with addr {:?}", peer.addr());
            self.send_to_peer(*peer.pub_key(), msg.raw());
        }
        self.add_peer_exchange_timeout();
    }

    pub fn broadcast_status(&mut self) {
        let hash = self.blockchain.last_hash().unwrap();
        let status = Status::new(self.state.public_key(),
                                 self.state.height(),
                                 &hash,
                                 self.state.secret_key());
        trace!("Broadcast status: {:?}", status);
        self.broadcast(status.raw());
    }
}
<commit_msg>Remove unwrap in the 'handle_message' function<commit_after>extern crate rand;

use rand::Rng;

use std::net::SocketAddr;

use messages::{Any, RawMessage, Connect, Status, Message, RequestPeers};
use events::Channel;
use super::{NodeHandler, RequestData, ExternalMessage, NodeTimeout, Height};

impl<S> NodeHandler<S>
    where S: Channel<ApplicationEvent = ExternalMessage, Timeout = NodeTimeout>
{
    pub fn handle_message(&mut self, raw: RawMessage) {
        // TODO: check message headers (network id, protocol version)
        // FIXME: call message.verify method
        //     if !raw.verify() {
        //         return;
        //     }

        //FIXME: add whitelist verify public_key

        match Any::from_raw(raw) {
            Ok(Any::Connect(msg)) => self.handle_connect(msg),
            Ok(Any::Status(msg)) => self.handle_status(msg),
            Ok(Any::Consensus(msg)) => self.handle_consensus(msg),
            Ok(Any::Request(msg)) => self.handle_request(msg),
            Ok(Any::Block(msg)) => self.handle_block(msg),
            Ok(Any::Transaction(msg)) => self.handle_tx(msg),
            Err(err) => {
                // TODO: Replace by `err.description()` after #103 is merged.
                error!("Invalid message received: {:?}", err);
            }
        }
    }

    pub fn handle_connected(&mut self, addr: &SocketAddr) {
        info!("Connected to: {}", addr);
        let message = self.state.our_connect_message().clone();
        self.send_to_addr(addr, message.raw());
    }

    pub fn handle_disconnected(&mut self, addr: &SocketAddr) {
        info!("Disconnected from: {}", addr);
        let need_reconnect = self.state.remove_peer_with_addr(addr);
        if need_reconnect {
            self.connect(addr);
        }
    }

    pub fn handle_connect(&mut self, message: Connect) {
        // TODO add spam protection
        let address = message.addr();
        if address == self.state.our_connect_message().addr() {
            return;
        }
        // Check if we have another connect message from peer with the given public_key
        let public_key = *message.pub_key();
        let mut need_connect = true;
        if let Some(saved_message) = self.state.peers().get(&public_key) {
            if saved_message.time() > message.time() {
                error!("Received outdated Connect message from {}", address);
                return;
            } else if saved_message.time() < message.time() {
                need_connect = saved_message.addr() != message.addr();
            } else if saved_message.addr() != message.addr() {
                error!("Received weird Connect message from {}", address);
                return;
            }
        }
        info!("Received Connect message from {}, {}",
              address,
              need_connect);
        self.state.add_peer(public_key, message);
        if need_connect {
            // TODO: reduce double sending of connect message
            info!("Send Connect message to {}", address);
            self.connect(&address);
        }
    }

    pub fn handle_status(&mut self, msg: Status) {
        let height = self.state.height();
        trace!("HANDLE STATUS: current height = {}, msg height = {}", height, msg.height());

        // Handle message from future height
        if msg.height() > height {
            let peer = msg.from();

            //verify message
            if !msg.verify_signature(peer) {
                return;
            }

            // Check validator height info
            if msg.height() > self.state.node_height(peer) {
                // Update validator height
                self.state.set_node_height(*peer, msg.height());
            }

            // Request block
            self.request(RequestData::Block(height), *peer);
        }
    }

    pub fn handle_request_peers(&mut self, msg: RequestPeers) {
        let peers: Vec<Connect> = self.state.peers().iter().map(|(_, b)| b.clone()).collect();
        trace!("HANDLE REQUEST PEERS: Sending {:?} peers to {:?}", peers, msg.from());

        for peer in peers {
            self.send_to_peer(*msg.from(), peer.raw());
        }
    }

    pub fn handle_status_timeout(&mut self, height: Height) {
        if self.state.height() == height {
            self.broadcast_status();
            self.add_status_timeout();
        }
    }

    pub fn handle_peer_exchange_timeout(&mut self) {
        if !self.state.peers().is_empty() {
            let to = self.state.peers().len();
            let gen_peer_id = || -> usize {
                let mut rng = rand::thread_rng();
                rng.gen_range(0, to)
            };

            let peer = self.state
                .peers()
                .iter()
                .map(|x| x.1.clone())
                .nth(gen_peer_id())
                .unwrap();
            let peer = peer.clone();
            let msg = RequestPeers::new(self.state.public_key(),
                                        peer.pub_key(),
                                        self.state.secret_key());
            trace!("Request peers from peer with addr {:?}", peer.addr());
            self.send_to_peer(*peer.pub_key(), msg.raw());
        }
        self.add_peer_exchange_timeout();
    }

    pub fn broadcast_status(&mut self) {
        let hash = self.blockchain.last_hash().unwrap();
        let status = Status::new(self.state.public_key(),
                                 self.state.height(),
                                 &hash,
                                 self.state.secret_key());
        trace!("Broadcast status: {:?}", status);
        self.broadcast(status.raw());
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::{Read, Write, Error};
use std::path::{Path, PathBuf};
use std::fs;
use std::fs::Metadata;
use std::ffi::CString;
use std::os::unix::fs::{PermissionsExt, MetadataExt, symlink};
use std::os::unix::io::AsRawFd;
use std::os::unix::ffi::OsStrExt;

use nix;
use libc::{uid_t, gid_t, c_int, utime, utimbuf};
use nix::fcntl::{flock, FlockArg};

use path_util::ToCString;

extern "C" {
    fn lchown(path: *const i8, owner: uid_t, group: gid_t) -> c_int;
}

pub struct Lock {
    file: fs::File,
}

pub fn read_visible_entries(dir: &Path) -> Result<Vec<PathBuf>, Error> {
    let mut res = vec!();
    for entry_ in fs::read_dir(dir)? {
        let entry = entry_?;
        if !entry.file_name()[..].to_str().map(|x| x.starts_with("."))
            .unwrap_or(false)
        {
            res.push(entry.path().to_path_buf());
        }
    }
    Ok(res)
}

pub struct Dir<'a> {
    path: &'a Path,
    recursive: bool,
    mode: Option<u32>,
    uid: Option<uid_t>,
    gid: Option<gid_t>,
}

impl<'a> Dir<'a> {
    pub fn new<P: ?Sized>(path: &'a P) -> Dir<'a>
        where P: AsRef<Path>
    {
        Dir {
            path: path.as_ref(),
            recursive: false,
            mode: None,
            uid: None,
            gid: None,
        }
    }

    pub fn recursive(&mut self, recursive: bool) -> &mut Dir<'a> {
        self.recursive = recursive;
        self
    }

    pub fn mode(&mut self, mode: u32) -> &mut Dir<'a> {
        self.mode = Some(mode);
        self
    }

    pub fn uid(&mut self, uid: uid_t) -> &mut Dir<'a> {
        self.uid = Some(uid);
        self
    }

    pub fn gid(&mut self, gid: gid_t) -> &mut Dir<'a> {
        self.gid = Some(gid);
        self
    }

    pub fn create(&self) -> Result<(), Error> {
        self._create(self.path, true)
    }

    fn _create(&self, path: &Path, is_last: bool)
        -> Result<(), Error>
    {
        if path.is_dir() {
            return Ok(())
        }
        if self.recursive {
            match path.parent() {
                Some(p) if p != path => {
                    self._create(p, false)?;
                }
                _ => {}
            }
        }
        fs::create_dir(path)?;
        let mode = if is_last { self.mode } else { None };
        fs::set_permissions(path,
            fs::Permissions::from_mode(mode.unwrap_or(0o755)))?;
        if is_last {
            if self.uid.is_some() || self.gid.is_some() {
                let uid = if let Some(uid) = self.uid {
                    uid
                } else {
                    path.symlink_metadata()?.uid()
                };
                let gid = if let Some(gid) = self.gid {
                    gid
                } else {
                    path.symlink_metadata()?.gid()
                };
                set_owner_group(path, uid, gid)?;
            }
        }
        Ok(())
    }
}

pub fn safe_ensure_dir(dir: &Path) -> Result<(), String> {
    match fs::symlink_metadata(dir) {
        Ok(ref stat) if stat.file_type().is_symlink() => {
            return Err(format!(concat!("The `{0:?}` dir can't be a symlink. ",
                               "Please run `unlink {0:?}`"), dir));
        }
        Ok(ref stat) if stat.file_type().is_dir() => {
            // ok
        }
        Ok(_) => {
            return Err(format!(concat!("The {0:?} must be a directory. ",
                               "Please run `unlink {0:?}`"), dir));
        }
        Err(ref e) if e.kind() == io::ErrorKind::NotFound => {
            try_msg!(Dir::new(dir).create(),
                "Can't create {dir:?}: {err}", dir=dir);
        }
        Err(ref e) => {
            return Err(format!("Can't stat {:?}: {}", dir, e));
        }
    }
    return Ok(());
}

pub fn ensure_symlink(target: &Path, linkpath: &Path) -> Result<(), io::Error>
{
    match symlink(target, linkpath) {
        Ok(()) => Ok(()),
        Err(e) => {
            if e.kind() == io::ErrorKind::AlreadyExists {
                match fs::read_link(linkpath) {
                    Ok(ref path) if Path::new(path) == target => Ok(()),
                    Ok(_) => Err(e),
                    Err(e) => Err(e),
                }
            } else  {
                Err(e)
            }
        }
    }
}

pub fn copy<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q) -> io::Result<()>
{
    _copy(from.as_ref(), to.as_ref(), None)
}

pub fn copy_with_mode<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q, mode: u32)
    -> io::Result<()>
{
    _copy(from.as_ref(), to.as_ref(), Some(mode))
}

fn _copy(from: &Path, to: &Path, mode: Option<u32>) -> io::Result<()> {
    if !from.is_file() {
        return Err(io::Error::new(io::ErrorKind::InvalidInput,
                              "the source path is not an existing regular file"))
    }

    let mut reader = fs::File::open(from)?;
    let mut writer = fs::File::create(to)?;

    copy_stream(&mut reader, &mut writer)?;

    let perm = match mode {
        Some(mode) => {
            fs::Permissions::from_mode(mode)
        },
        None => {
            reader.metadata()?.permissions()
        }
    };
    fs::set_permissions(to, perm)?;
    Ok(())
}

pub fn copy_stream(reader: &mut Read, writer: &mut Write)
    -> io::Result<()>
{
    // Smaller buffer on the stack
    // Because rust musl has very small stack (80k)
    // Allocating buffer on heap for each copy is too slow
    let mut buf = [0; 32768];
    loop {
        let len = match reader.read(&mut buf) {
            Ok(0) => break,
            Ok(len) => len,
            Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
            Err(e) => return Err(e),
        };
        writer.write_all(&buf[..len])?;
    }
    Ok(())
}

impl Lock {
    pub fn exclusive<P: AsRef<Path>>(p: P) -> Result<Lock, Error> {
        let f = fs::File::create(p)?;
        flock(f.as_raw_fd(), FlockArg::LockExclusiveNonblock)
            .map_err(|e| match e {
                nix::Error::Sys(code) => Error::from_raw_os_error(code as i32),
                nix::Error::InvalidPath => unreachable!(),
            })?;
        Ok(Lock {
            file: f,
        })
    }
    pub fn exclusive_wait<P: AsRef<Path>>(path: P, message: &str)
        -> Result<Lock, Error>
    {
        let f = fs::File::create(path)?;
        match flock(f.as_raw_fd(), FlockArg::LockExclusiveNonblock) {
            Ok(()) => {}
            Err(nix::Error::Sys(nix::Errno::EAGAIN)) => {
                warn!("{}", message);
                flock(f.as_raw_fd(), FlockArg::LockExclusive)
                    .map_err(|e| match e {
                        nix::Error::Sys(code) => {
                            Error::from_raw_os_error(code as i32)
                        },
                        nix::Error::InvalidPath => unreachable!(),
                    })?;
            }
            Err(nix::Error::Sys(code)) => {
                return Err(Error::from_raw_os_error(code as i32))
            }
            Err(nix::Error::InvalidPath) => unreachable!(),
        }
        Ok(Lock {
            file: f,
        })
    }
}

impl Drop for Lock {
    fn drop(&mut self) {
        flock(self.file.as_raw_fd(), FlockArg::Unlock)
            .map_err(|e| error!("Couldn't unlock file: {:?}", e)).ok();
    }
}

pub fn check_stream_hashsum(mut reader: &mut Read, sha256: &String)
    -> Result<(), String>
{
    use sha2::Sha256;
    use digest_writer::Writer;
    use digest::hex;

    let mut hash = Writer::new(Sha256::new());
    try_msg!(io::copy(&mut reader, &mut hash),
        "Error when calculating hashsum: {err}");
    let hash_str = format!("{:x}", hex(&hash.into_inner()));
    if !hash_str.starts_with(sha256) {
        return Err(format!("Hashsum mismatch: expected {} but was {}",
            sha256, hash_str));
    }
    Ok(())
}

pub fn force_symlink(target: &Path, linkpath: &Path)
    -> Result<(), io::Error>
{
    let tmpname = linkpath.with_extension(".vagga.tmp.link~");
    symlink(target, &tmpname)?;
    fs::rename(&tmpname, linkpath)?;
    Ok(())
}

pub fn set_owner_group(target: &Path, uid: uid_t, gid: gid_t)
    -> Result<(), io::Error>
{
    let c_target = target.to_cstring();
    let rc = unsafe { lchown( c_target.as_ptr(), uid, gid) };
    if rc != 0 {
        warn!("Can't chown {:?}: {}", target, io::Error::last_os_error());
        Ok(())
        // Err(io::Error::last_os_error())
    } else {
        Ok(())
    }
}

/// Shallow copy of file/dir/symlink
///
/// The owner_uid/owner_gid parameters optionally override the owner
///
/// Returns false if path is a directory
pub fn shallow_copy(src: &Path, src_stat: &Metadata, dest: &Path,
    owner_uid: Option<uid_t>, owner_gid: Option<gid_t>,
    mode: Option<u32>)
    -> Result<bool, io::Error>
{
    let src_type = src_stat.file_type();
    let mut is_dir = false;
    if src_type.is_dir() {
        is_dir = true;
        let nstat = dest.symlink_metadata();
        // We don't change permissions and owner of already created directories
        //
        // There are couple of reasons:
        // 1. We consider such overlaps a "merge" of a directory, so it's
        //    unclear which one should override the permissions
        // 2. Changing permissions of the root directory of the copied system
        //    is usually counter-intuitive (so we contradict rsync)
        // 3. Some directories (/proc, /sys, /dev) are mount points and we
        //    can't change the permissions
        if nstat.is_err() {
            Dir::new(dest)
                .mode(mode.unwrap_or(src_stat.mode()))
                .uid(owner_uid.unwrap_or(src_stat.uid()))
                .gid(owner_gid.unwrap_or(src_stat.gid()))
                .create()?;
        }
    } else if src_type.is_symlink() {
        let value = fs::read_link(&src)?;
        force_symlink(&value, dest)?;
        set_owner_group(dest,
            owner_uid.unwrap_or(src_stat.uid()),
            owner_gid.unwrap_or(src_stat.gid()))?;
    } else {
        match mode {
            Some(mode) => copy_with_mode(src, dest, mode)?,
            None => copy(src, dest)?,
        }
        set_owner_group(dest,
            owner_uid.unwrap_or(src_stat.uid()),
            owner_gid.unwrap_or(src_stat.gid()))?;
    }
    Ok(!is_dir)
}

pub fn copy_utime<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q)
    -> io::Result<()>
{
    let metadata = fs::metadata(from.as_ref())?;
    let filename = CString::new(to.as_ref().as_os_str().as_bytes()).unwrap();
    let utimes = utimbuf {
        actime: metadata.atime(),
        modtime: metadata.mtime(),
    };
    let rc = unsafe { utime(filename.as_ptr(), &utimes) };
    if rc != 0 {
        return Err(io::Error::last_os_error());
    }
    Ok(())
}
<commit_msg>Add log of elapsed time to lock waiting<commit_after>use std::io;
use std::io::{Read, Write, Error};
use std::path::{Path, PathBuf};
use std::fs;
use std::fs::Metadata;
use std::ffi::CString;
use std::time::{Instant, Duration};
use std::os::unix::fs::{PermissionsExt, MetadataExt, symlink};
use std::os::unix::io::AsRawFd;
use std::os::unix::ffi::OsStrExt;

use nix;
use libc::{uid_t, gid_t, c_int, utime, utimbuf};
use nix::fcntl::{flock, FlockArg};

use path_util::ToCString;

extern "C" {
    fn lchown(path: *const i8, owner: uid_t, group: gid_t) -> c_int;
}

pub struct Lock {
    file: fs::File,
}

pub fn read_visible_entries(dir: &Path) -> Result<Vec<PathBuf>, Error> {
    let mut res = vec!();
    for entry_ in fs::read_dir(dir)? {
        let entry = entry_?;
        if !entry.file_name()[..].to_str().map(|x| x.starts_with("."))
            .unwrap_or(false)
        {
            res.push(entry.path().to_path_buf());
        }
    }
    Ok(res)
}

pub struct Dir<'a> {
    path: &'a Path,
    recursive: bool,
    mode: Option<u32>,
    uid: Option<uid_t>,
    gid: Option<gid_t>,
}

impl<'a> Dir<'a> {
    pub fn new<P: ?Sized>(path: &'a P) -> Dir<'a>
        where P: AsRef<Path>
    {
        Dir {
            path: path.as_ref(),
            recursive: false,
            mode: None,
            uid: None,
            gid: None,
        }
    }

    pub fn recursive(&mut self, recursive: bool) -> &mut Dir<'a> {
        self.recursive = recursive;
        self
    }

    pub fn mode(&mut self, mode: u32) -> &mut Dir<'a> {
        self.mode = Some(mode);
        self
    }

    pub fn uid(&mut self, uid: uid_t) -> &mut Dir<'a> {
        self.uid = Some(uid);
        self
    }

    pub fn gid(&mut self, gid: gid_t) -> &mut Dir<'a> {
        self.gid = Some(gid);
        self
    }

    pub fn create(&self) -> Result<(), Error> {
        self._create(self.path, true)
    }

    fn _create(&self, path: &Path, is_last: bool)
        -> Result<(), Error>
    {
        if path.is_dir() {
            return Ok(())
        }
        if self.recursive {
            match path.parent() {
                Some(p) if p != path => {
                    self._create(p, false)?;
                }
                _ => {}
            }
        }
        fs::create_dir(path)?;
        let mode = if is_last { self.mode } else { None };
        fs::set_permissions(path,
            fs::Permissions::from_mode(mode.unwrap_or(0o755)))?;
        if is_last {
            if self.uid.is_some() || self.gid.is_some() {
                let uid = if let Some(uid) = self.uid {
                    uid
                } else {
                    path.symlink_metadata()?.uid()
                };
                let gid = if let Some(gid) = self.gid {
                    gid
                } else {
                    path.symlink_metadata()?.gid()
                };
                set_owner_group(path, uid, gid)?;
            }
        }
        Ok(())
    }
}

pub fn safe_ensure_dir(dir: &Path) -> Result<(), String> {
    match fs::symlink_metadata(dir) {
        Ok(ref stat) if stat.file_type().is_symlink() => {
            return Err(format!(concat!("The `{0:?}` dir can't be a symlink. ",
                               "Please run `unlink {0:?}`"), dir));
        }
        Ok(ref stat) if stat.file_type().is_dir() => {
            // ok
        }
        Ok(_) => {
            return Err(format!(concat!("The {0:?} must be a directory. ",
                               "Please run `unlink {0:?}`"), dir));
        }
        Err(ref e) if e.kind() == io::ErrorKind::NotFound => {
            try_msg!(Dir::new(dir).create(),
                "Can't create {dir:?}: {err}", dir=dir);
        }
        Err(ref e) => {
            return Err(format!("Can't stat {:?}: {}", dir, e));
        }
    }
    return Ok(());
}

pub fn ensure_symlink(target: &Path, linkpath: &Path) -> Result<(), io::Error>
{
    match symlink(target, linkpath) {
        Ok(()) => Ok(()),
        Err(e) => {
            if e.kind() == io::ErrorKind::AlreadyExists {
                match fs::read_link(linkpath) {
                    Ok(ref path) if Path::new(path) == target => Ok(()),
                    Ok(_) => Err(e),
                    Err(e) => Err(e),
                }
            } else  {
                Err(e)
            }
        }
    }
}

pub fn copy<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q) -> io::Result<()>
{
    _copy(from.as_ref(), to.as_ref(), None)
}

pub fn copy_with_mode<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q, mode: u32)
    -> io::Result<()>
{
    _copy(from.as_ref(), to.as_ref(), Some(mode))
}

fn _copy(from: &Path, to: &Path, mode: Option<u32>) -> io::Result<()> {
    if !from.is_file() {
        return Err(io::Error::new(io::ErrorKind::InvalidInput,
                              "the source path is not an existing regular file"))
    }

    let mut reader = fs::File::open(from)?;
    let mut writer = fs::File::create(to)?;

    copy_stream(&mut reader, &mut writer)?;

    let perm = match mode {
        Some(mode) => {
            fs::Permissions::from_mode(mode)
        },
        None => {
            reader.metadata()?.permissions()
        }
    };
    fs::set_permissions(to, perm)?;
    Ok(())
}

pub fn copy_stream(reader: &mut Read, writer: &mut Write)
    -> io::Result<()>
{
    // Smaller buffer on the stack
    // Because rust musl has very small stack (80k)
    // Allocating buffer on heap for each copy is too slow
    let mut buf = [0; 32768];
    loop {
        let len = match reader.read(&mut buf) {
            Ok(0) => break,
            Ok(len) => len,
            Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
            Err(e) => return Err(e),
        };
        writer.write_all(&buf[..len])?;
    }
    Ok(())
}

impl Lock {
    pub fn exclusive<P: AsRef<Path>>(p: P) -> Result<Lock, Error> {
        let f = fs::File::create(p)?;
        flock(f.as_raw_fd(), FlockArg::LockExclusiveNonblock)
            .map_err(|e| match e {
                nix::Error::Sys(code) => Error::from_raw_os_error(code as i32),
                nix::Error::InvalidPath => unreachable!(),
            })?;
        Ok(Lock {
            file: f,
        })
    }
    pub fn exclusive_wait<P: AsRef<Path>>(path: P, message: &str)
        -> Result<Lock, Error>
    {
        let f = fs::File::create(path)?;
        match flock(f.as_raw_fd(), FlockArg::LockExclusiveNonblock) {
            Ok(()) => {}
            Err(nix::Error::Sys(nix::Errno::EAGAIN)) => {
                warn!("{}", message);
                let lock_start = Instant::now();
                flock(f.as_raw_fd(), FlockArg::LockExclusive)
                    .map_err(|e| match e {
                        nix::Error::Sys(code) => {
                            Error::from_raw_os_error(code as i32)
                        },
                        nix::Error::InvalidPath => unreachable!(),
                    })?;
                let elapsed = lock_start.elapsed();
                if elapsed > Duration::new(5, 0) {
                    warn!("Lock was held {} seconds. Proceeding...",
                        elapsed.as_secs());
                }
            }
            Err(nix::Error::Sys(code)) => {
                return Err(Error::from_raw_os_error(code as i32))
            }
            Err(nix::Error::InvalidPath) => unreachable!(),
        }
        Ok(Lock {
            file: f,
        })
    }
}

impl Drop for Lock {
    fn drop(&mut self) {
        flock(self.file.as_raw_fd(), FlockArg::Unlock)
            .map_err(|e| error!("Couldn't unlock file: {:?}", e)).ok();
    }
}

pub fn check_stream_hashsum(mut reader: &mut Read, sha256: &String)
    -> Result<(), String>
{
    use sha2::Sha256;
    use digest_writer::Writer;
    use digest::hex;

    let mut hash = Writer::new(Sha256::new());
    try_msg!(io::copy(&mut reader, &mut hash),
        "Error when calculating hashsum: {err}");
    let hash_str = format!("{:x}", hex(&hash.into_inner()));
    if !hash_str.starts_with(sha256) {
        return Err(format!("Hashsum mismatch: expected {} but was {}",
            sha256, hash_str));
    }
    Ok(())
}

pub fn force_symlink(target: &Path, linkpath: &Path)
    -> Result<(), io::Error>
{
    let tmpname = linkpath.with_extension(".vagga.tmp.link~");
    symlink(target, &tmpname)?;
    fs::rename(&tmpname, linkpath)?;
    Ok(())
}

pub fn set_owner_group(target: &Path, uid: uid_t, gid: gid_t)
    -> Result<(), io::Error>
{
    let c_target = target.to_cstring();
    let rc = unsafe { lchown( c_target.as_ptr(), uid, gid) };
    if rc != 0 {
        warn!("Can't chown {:?}: {}", target, io::Error::last_os_error());
        Ok(())
        // Err(io::Error::last_os_error())
    } else {
        Ok(())
    }
}

/// Shallow copy of file/dir/symlink
///
/// The owner_uid/owner_gid parameters optionally override the owner
///
/// Returns false if path is a directory
pub fn shallow_copy(src: &Path, src_stat: &Metadata, dest: &Path,
    owner_uid: Option<uid_t>, owner_gid: Option<gid_t>,
    mode: Option<u32>)
    -> Result<bool, io::Error>
{
    let src_type = src_stat.file_type();
    let mut is_dir = false;
    if src_type.is_dir() {
        is_dir = true;
        let nstat = dest.symlink_metadata();
        // We don't change permissions and owner of already created directories
        //
        // There are couple of reasons:
        // 1. We consider such overlaps a "merge" of a directory, so it's
        //    unclear which one should override the permissions
        // 2. Changing permissions of the root directory of the copied system
        //    is usually counter-intuitive (so we contradict rsync)
        // 3. Some directories (/proc, /sys, /dev) are mount points and we
        //    can't change the permissions
        if nstat.is_err() {
            Dir::new(dest)
                .mode(mode.unwrap_or(src_stat.mode()))
                .uid(owner_uid.unwrap_or(src_stat.uid()))
                .gid(owner_gid.unwrap_or(src_stat.gid()))
                .create()?;
        }
    } else if src_type.is_symlink() {
        let value = fs::read_link(&src)?;
        force_symlink(&value, dest)?;
        set_owner_group(dest,
            owner_uid.unwrap_or(src_stat.uid()),
            owner_gid.unwrap_or(src_stat.gid()))?;
    } else {
        match mode {
            Some(mode) => copy_with_mode(src, dest, mode)?,
            None => copy(src, dest)?,
        }
        set_owner_group(dest,
            owner_uid.unwrap_or(src_stat.uid()),
            owner_gid.unwrap_or(src_stat.gid()))?;
    }
    Ok(!is_dir)
}

pub fn copy_utime<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q)
    -> io::Result<()>
{
    let metadata = fs::metadata(from.as_ref())?;
    let filename = CString::new(to.as_ref().as_os_str().as_bytes()).unwrap();
    let utimes = utimbuf {
        actime: metadata.atime(),
        modtime: metadata.mtime(),
    };
    let rc = unsafe { utime(filename.as_ptr(), &utimes) };
    if rc != 0 {
        return Err(io::Error::last_os_error());
    }
    Ok(())
}
<|endoftext|>


<commit_before>// (c) 2016 Productize SPRL <joost@productize.be>

use std::io;

use Sexp;
use symbolic_expressions::Result;
use symbolic_expressions::Formatter;

// custom symbolic_expressions formatter that aims to be
// kicad compatible

#[derive(PartialEq)]
enum Indent {
    Not,
    Before,
    BeforeAfter,
    BeforeDouble,
    BeforeDoubleAfter,
    BeforeDoubleAfterDouble,
}

pub struct KicadFormatter {
    indent:i64,
    stack:Vec<Option<(String,Indent)>>,
    ind:Vec<u8>,
    pts_xy_count:i64,
}

impl KicadFormatter {

    pub fn new(initial_indent_level:i64) -> KicadFormatter {
        KicadFormatter {
            indent:initial_indent_level,
            stack:vec![],
            ind:vec![b' ',b' '], // two spaces
            pts_xy_count:0,
        }
    }

    fn is(&self, what:&'static str) -> bool {
        self.stack.iter().any(
            |x:&Option<(String,Indent)>| {
                if let Some((ref y,_)) = *x {
                    y == what
                } else {
                    false
                }
            })
    }

    fn parent_is(&self, what:&'static str) -> bool {
        if let Some(s) = self.stack.last() {
            if let Some((ref t,_)) = *s {
                return t == what
            }
        }
        false
    }

    fn indent<W:io::Write>(&self, writer:&mut W, nls:i64) -> Result<()> {
        for _ in 0..nls {
            try!(writer.write_all(b"\n"));
        }
        for _ in 0..self.indent {
            try!(writer.write_all(&self.ind));
        }
        Ok(())
    }

    fn indent_plus<W:io::Write>(&mut self, writer:&mut W, nls:i64) -> Result<()> {
        self.indent+=1;
        let res = self.indent(writer, nls);
        self.indent-=1;
        res
    }

    fn want_indent_module(&self, ele:&str) -> Indent {
        if !self.is("module") {
            return Indent::Not
        }
        if self.parent_is("module") {
            match ele {
                "at" | "descr" | "fp_line" | "fp_poly" |
                "pad" => return Indent::Before,
                "model" | "fp_text" => return Indent::BeforeAfter,
                _ => (),
            }
        }
        if self.parent_is("fp_text") | self.parent_is("gr_text") {
            if let "effects" = ele {
                return Indent::Before
            }
        }
        if self.parent_is("pts") {
            if let "xy" = ele {
                if self.pts_xy_count > 0 && self.pts_xy_count % 4 == 0 {
                    return Indent::Before
                }
            }
        }
        if self.parent_is("model") {
            match ele {
                "at" | "scale" | "rotate" => {
                    return Indent::Before
                },
                _ => (),
            }
        }
        Indent::Not
    }

    fn want_indent_layout(&self, ele:&str, e:&Sexp) -> Indent {
        if !self.is("kicad_pcb") {
            return Indent::Not
        }
        if self.parent_is("kicad_pcb") {
            match ele {
                "page" |
                "module" |
                "gr_arc"  | "gr_circle"
                    => return Indent::BeforeDouble,
                "net" | "gr_line" | "segment" | "via"
                    => return Indent::Before,
                "layers" | "gr_text" | "dimension" | "zone"
                    => return Indent::BeforeAfter,
                "setup"
                    => return Indent::BeforeDoubleAfterDouble,
                "general" | "net_class"
                    => return Indent::BeforeDoubleAfter,
                _ => (),
            }
        }
        if self.parent_is("general") {
            return Indent::Before
        }
        if self.parent_is("layers") {
            return Indent::Before
        }
        if self.parent_is("setup") {
            return Indent::Before
        }
        if self.parent_is("pcbplotparams") {
            return Indent::Before
        }
        if self.parent_is("net_class") {
            return Indent::Before
        }
        if self.parent_is("dimension") {
            match ele {
                "gr_text" | "feature1" |
                "feature2" | "crossbar" |
                "arrow1a" | "arrow1b" |
                "arrow2a" | "arrow2b" => {
                    return Indent::Before
                },
                _ => (),
            }
        }
        if self.parent_is("zone") {
            match ele {
                "connect_pads" | "min_thickness" | "fill" |
                "polygon" | "filled_polygon"
                    => return Indent::Before,
                _ => (),
            }
        }
        if self.parent_is("polygon") | self.parent_is("filled_polygon") {
            return Indent::Before
        }
        Indent::Not
    }

    fn want_indent(&self, value:&Sexp) -> Indent {
        let first = match *value {
            Sexp::List(ref l) => {
                if l.is_empty() {
                    return Indent::Not
                }
                (&l[0]).clone()
            },
            Sexp::Empty => return Indent::Not,
            Sexp::String(ref l) => Sexp::String(l.clone()),
        };
        if let Sexp::String(ref ele) = first {
            let i = self.want_indent_module(ele);
            if i != Indent::Not {
                return i
            }
            let i = self.want_indent_layout(ele, value);
            if i != Indent::Not {
                return i
            }
        }
        Indent::Not
    }
}

impl Formatter for KicadFormatter {

    fn open<W>(&mut self, writer: &mut W, value:Option<&Sexp>) -> Result<()>
        where W: io::Write
    {
        let mut ele = String::new();
        // if first element is string
        if let Some(ref sexp) = value {
            if let Sexp::String(ref s) = **sexp {
                ele.push_str(s);
            }
        }
        let exp = Sexp::String(ele.clone());
        let want_indent = self.want_indent(&exp);
        if want_indent != Indent::Not {
            self.indent += 1;
            if want_indent == Indent::BeforeDouble
            || want_indent == Indent::BeforeDoubleAfterDouble
            || want_indent == Indent::BeforeDoubleAfter {
                try!(self.indent(writer, 2));
            } else {
                try!(self.indent(writer, 1));
            }
        }

        // special handling for breaking after 4 elements of xy
        if let "pts" = &ele[..] {
            self.pts_xy_count = 0;
        }
        if self.parent_is("pts") {
            if let "xy" = &ele[..] {
                self.pts_xy_count += 1;
                if self.pts_xy_count == 5 {
                    self.pts_xy_count = 1;
                }
            }
        }

        if !ele.is_empty() {
            self.stack.push(Some((ele, want_indent)))
        } else {
            self.stack.push(None)
        }
        writer.write_all(b"(").map_err(From::from)
    }

    fn element<W>(&mut self, writer: &mut W, value:&Sexp) -> Result<()>
        where W: io::Write
    {
        // get rid of the space if we will be putting a newline next
        if self.want_indent(value) == Indent::Not {
            try!(writer.write_all(b" "));
        } else if let Sexp::String(_) = *value {
            try!(writer.write_all(b" "));
        }
        Ok(())

    }

    fn close<W>(&mut self, writer: &mut W) -> Result<()>
        where W: io::Write
    {
        if let Some(Some((s, want_indent))) = self.stack.pop() {
            if want_indent != Indent::Not {
                self.indent -= 1
            }
            match want_indent {
                Indent::Not |
                Indent::Before |
                Indent::BeforeDouble => (),
                Indent::BeforeDoubleAfter |
                Indent::BeforeAfter |
                Indent::BeforeDoubleAfterDouble
                    => {
                    try!(self.indent_plus(writer, 1));
                }
            }
            // TODO: remove
            if let "module" = &s[..] {
                try!(self.indent(writer, 1));
            }
            try!(writer.write_all(b")"));
            if want_indent == Indent::BeforeDoubleAfterDouble {
                try!(writer.write_all(b"\n"));
            }
        } else {
            try!(writer.write_all(b")"));
        }
        Ok(())
    }
}

<commit_msg>rework indent<commit_after>// (c) 2016 Productize SPRL <joost@productize.be>

use std::io;

use Sexp;
use symbolic_expressions::Result;
use symbolic_expressions::Formatter;

// custom symbolic_expressions formatter that aims to be
// kicad compatible

struct Indent2 {
    newline_before:i64,
    closing_on_newline:bool,
    newline_after:i64,
}

impl Indent2 {
    fn before() -> Indent2 {
        Indent2 {
            newline_before:1,
            closing_on_newline:false,
            newline_after:0,
        }
    }
    fn before_after() -> Indent2 {
        Indent2 {
            newline_before:1,
            closing_on_newline:true,
            newline_after:0,
        }
    }
    fn before_double() -> Indent2 {
        Indent2 {
            newline_before:2,
            closing_on_newline:false,
            newline_after:0,
        }
    }
    fn before_double_after() -> Indent2 {
        Indent2 {
            newline_before:2,
            closing_on_newline:true,
            newline_after:0,
        }
    }
    fn before_double_after_double() -> Indent2 {
        Indent2 {
            newline_before:2,
            closing_on_newline:true,
            newline_after:1,
        }
    }
}

pub struct KicadFormatter {
    indent:i64,
    stack:Vec<Option<(String,Option<Indent2>)>>,
    ind:Vec<u8>,
    pts_xy_count:i64,
}

impl KicadFormatter {

    pub fn new(initial_indent_level:i64) -> KicadFormatter {
        KicadFormatter {
            indent:initial_indent_level,
            stack:vec![],
            ind:vec![b' ',b' '], // two spaces
            pts_xy_count:0,
        }
    }

    fn is(&self, what:&'static str) -> bool {
        for x in &self.stack {
            if let Some((ref name,_)) = *x {
                if name == what {
                    return true
                }
            }
        }
        false
    }

    fn parent_is(&self, what:&'static str) -> bool {
        if let Some(s) = self.stack.last() {
            if let Some((ref t,_)) = *s {
                return t == what
            }
        }
        false
    }

    fn indent<W:io::Write>(&self, writer:&mut W, nls:i64) -> Result<()> {
        for _ in 0..nls {
            try!(writer.write_all(b"\n"));
        }
        for _ in 0..self.indent {
            try!(writer.write_all(&self.ind));
        }
        Ok(())
    }

    fn indent_plus<W:io::Write>(&mut self, writer:&mut W, nls:i64) -> Result<()> {
        self.indent+=1;
        let res = self.indent(writer, nls);
        self.indent-=1;
        res
    }

    fn want_indent_module(&self, ele:&str) -> Option<Indent2> {
        //if !self.is("module") {
        //    return None
        //}
        if self.parent_is("module") {
            match ele {
                "at" | "descr" | "fp_line" | "fp_poly" |
                "pad" => return Some(Indent2::before()),
                "model" | "fp_text" => return Some(Indent2::before_after()),
                _ => (),
            }
        }
        if self.parent_is("fp_text") | self.parent_is("gr_text") {
            if let "effects" = ele {
                return Some(Indent2::before())
            }
        }
        if self.parent_is("pts") {
            if let "xy" = ele {
                if self.pts_xy_count > 0 && self.pts_xy_count % 4 == 0 {
                    return Some(Indent2::before())
                }
            }
        }
        if self.parent_is("model") {
            match ele {
                "at" | "scale" | "rotate" => {
                    return Some(Indent2::before())
                },
                _ => (),
            }
        }
        None
    }

    fn want_indent_layout(&self, ele:&str, e:&Sexp) -> Option<Indent2> {
        if !self.is("kicad_pcb") {
            return None
        }
        if self.parent_is("kicad_pcb") {
            match ele {
                "page" |
                "module" |
                "gr_arc"  | "gr_circle"
                    => return Some(Indent2::before_double()),
                "net" | "gr_line" | "segment" | "via"
                    => return Some(Indent2::before()),
                "layers" | "gr_text" | "dimension" | "zone"
                    => return Some(Indent2::before_after()),
                "setup"
                    => return Some(Indent2::before_double_after_double()),
                "general" | "net_class"
                    => return Some(Indent2::before_double_after()),
                _ => (),
            }
        }
        if self.parent_is("general") {
            return Some(Indent2::before())
        }
        if self.parent_is("layers") {
            return Some(Indent2::before())
        }
        if self.parent_is("setup") {
            return Some(Indent2::before())
        }
        if self.parent_is("pcbplotparams") {
            return Some(Indent2::before())
        }
        if self.parent_is("net_class") {
            return Some(Indent2::before())
        }
        if self.parent_is("dimension") {
            match ele {
                "gr_text" | "feature1" |
                "feature2" | "crossbar" |
                "arrow1a" | "arrow1b" |
                "arrow2a" | "arrow2b" => {
                    return Some(Indent2::before())
                },
                _ => (),
            }
        }
        if self.parent_is("zone") {
            match ele {
                "connect_pads" | "min_thickness" | "fill" |
                "polygon" | "filled_polygon"
                    => return Some(Indent2::before()),
                _ => (),
            }
        }
        if self.parent_is("polygon") | self.parent_is("filled_polygon") {
            return Some(Indent2::before())
        }
        None
    }

    fn want_indent(&self, value:&Sexp) -> Option<Indent2> {
        let first = match *value {
            Sexp::List(ref l) => {
                if l.is_empty() {
                    return None
                }
                (&l[0]).clone()
            },
            Sexp::Empty => return None,
            Sexp::String(ref l) => Sexp::String(l.clone()),
        };
        if let Sexp::String(ref ele) = first {
            let i = self.want_indent_module(ele);
            if i.is_some() {
                return i
            }
            let i = self.want_indent_layout(ele, value);
            if i.is_some() {
                return i
            }
        }
        None
    }
}

impl Formatter for KicadFormatter {

    fn open<W>(&mut self, writer: &mut W, value:Option<&Sexp>) -> Result<()>
        where W: io::Write
    {
        let mut ele = String::new();
        // if first element is string
        if let Some(ref sexp) = value {
            if let Sexp::String(ref s) = **sexp {
                ele.push_str(s);
            }
        }
        let exp = Sexp::String(ele.clone());
        let want_indent = self.want_indent(&exp);
        if let Some(ref want_indent) = want_indent {
            self.indent += 1;
            if want_indent.newline_before > 0 {
                try!(self.indent(writer, want_indent.newline_before));
            }
        }

        // special handling for breaking after 4 elements of xy
        if let "pts" = &ele[..] {
            self.pts_xy_count = 0;
        }
        if self.parent_is("pts") {
            if let "xy" = &ele[..] {
                self.pts_xy_count += 1;
                if self.pts_xy_count == 5 {
                    self.pts_xy_count = 1;
                }
            }
        }

        if !ele.is_empty() {
            self.stack.push(Some((ele, want_indent)))
        } else {
            self.stack.push(None)
        }
        writer.write_all(b"(").map_err(From::from)
    }

    fn element<W>(&mut self, writer: &mut W, value:&Sexp) -> Result<()>
        where W: io::Write
    {
        // get rid of the space if we will be putting a newline next
        if self.want_indent(value).is_none() {
            try!(writer.write_all(b" "));
        } else if let Sexp::String(_) = *value {
            try!(writer.write_all(b" "));
        }
        Ok(())

    }

    fn close<W>(&mut self, writer: &mut W) -> Result<()>
        where W: io::Write
    {
        if let Some(Some((s, want_indent))) = self.stack.pop() {
            if let Some(indent) = want_indent {
                self.indent -= 1;
                if indent.closing_on_newline {
                    try!(self.indent_plus(writer, 1));
                }
                // special handling of toplevel module...
                // which doesn't work, because it is not indented
                if &s == "module" && self.stack.is_empty() {
                    try!(writer.write_all(b"\n"));
                }
                try!(writer.write_all(b")"));
                for _ in 0..indent.newline_after {
                    try!(writer.write_all(b"\n"));
                }
                return Ok(())
            } else {
                if self.stack.is_empty() && &s == "module" {
                    try!(writer.write_all(b"\n"));
                }
            }
        }
        try!(writer.write_all(b")"));
        Ok(())
    }
}

<|endoftext|>


<commit_before>//! Create or redefine SQL functions
use std::ffi::CStr;
use std::mem;
use std::str;
use libc::{c_int, c_double, c_char, c_void};

use ffi;
pub use ffi::sqlite3_context;
pub use ffi::sqlite3_value;
pub use ffi::sqlite3_value_type;
pub use ffi::sqlite3_value_numeric_type;

use types::Null;

use {SqliteResult, SqliteError, SqliteConnection, str_to_cstring, InnerSqliteConnection};

/// A trait for types that can be converted into the result of an SQL function.
pub trait ToResult {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context);
}

macro_rules! raw_to_impl(
    ($t:ty, $f:ident) => (
        impl ToResult for $t {
            unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
                ffi::$f(ctx, *self)
            }
        }
    )
);

raw_to_impl!(c_int, sqlite3_result_int);
raw_to_impl!(i64, sqlite3_result_int64);
raw_to_impl!(c_double, sqlite3_result_double);

impl<'a> ToResult for bool {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        match *self {
            true => ffi::sqlite3_result_int(ctx, 1),
            _ => ffi::sqlite3_result_int(ctx, 0),
        }
    }
}

impl<'a> ToResult for &'a str {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        let length = self.len();
        if length > ::std::i32::MAX as usize {
            ffi::sqlite3_result_error_toobig(ctx);
            return;
        }
        match str_to_cstring(self) {
            Ok(c_str) => {
                ffi::sqlite3_result_text(ctx,
                                         c_str.as_ptr(),
                                         length as c_int,
                                         ffi::SQLITE_TRANSIENT())
            }
            Err(_) => ffi::sqlite3_result_error_code(ctx, ffi::SQLITE_MISUSE), // TODO sqlite3_result_error
        }
    }
}

impl ToResult for String {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        (&self[..]).set_result(ctx)
    }
}

impl<'a> ToResult for &'a [u8] {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        if self.len() > ::std::i32::MAX as usize {
            ffi::sqlite3_result_error_toobig(ctx);
            return;
        }
        ffi::sqlite3_result_blob(ctx,
                                 mem::transmute(self.as_ptr()),
                                 self.len() as c_int,
                                 ffi::SQLITE_TRANSIENT())
    }
}

impl ToResult for Vec<u8> {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        (&self[..]).set_result(ctx)
    }
}

impl<T: ToResult> ToResult for Option<T> {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        match *self {
            None => ffi::sqlite3_result_null(ctx),
            Some(ref t) => t.set_result(ctx),
        }
    }
}

impl ToResult for Null {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        ffi::sqlite3_result_null(ctx)
    }
}

// sqlite3_result_error_code, c_int
// sqlite3_result_error_nomem
// sqlite3_result_error_toobig
// sqlite3_result_error, *const c_char, c_int
// sqlite3_result_zeroblob
// sqlite3_result_value

/// A trait for types that can be created from a SQLite function parameter value.
pub trait FromValue: Sized {
    unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<Self>;

    /// FromValue types can implement this method and use sqlite3_value_type to check that
    /// the type reported by SQLite matches a type suitable for Self. This method is used
    /// by `???` to confirm that the parameter contains a valid type before
    /// attempting to retrieve the value.
    unsafe fn parameter_has_valid_sqlite_type(_: *mut sqlite3_value) -> bool {
        true
    }
}

macro_rules! raw_from_impl(
    ($t:ty, $f:ident, $c:expr) => (
        impl FromValue for $t {
            unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<$t> {
                Ok(ffi::$f(v))
            }

            unsafe fn parameter_has_valid_sqlite_type(v: *mut sqlite3_value) -> bool {
                sqlite3_value_numeric_type(v) == $c
            }
        }
    )
);

raw_from_impl!(c_int, sqlite3_value_int, ffi::SQLITE_INTEGER);
raw_from_impl!(i64, sqlite3_value_int64, ffi::SQLITE_INTEGER);

impl FromValue for bool {
    unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<bool> {
        match ffi::sqlite3_value_int(v) {
            0 => Ok(false),
            _ => Ok(true),
        }
    }

    unsafe fn parameter_has_valid_sqlite_type(v: *mut sqlite3_value) -> bool {
        sqlite3_value_numeric_type(v) == ffi::SQLITE_INTEGER
    }
}

impl FromValue for c_double {
    unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<c_double> {
        Ok(ffi::sqlite3_value_double(v))
    }

    unsafe fn parameter_has_valid_sqlite_type(v: *mut sqlite3_value) -> bool {
        sqlite3_value_numeric_type(v) == ffi::SQLITE_FLOAT ||
        sqlite3_value_numeric_type(v) == ffi::SQLITE_INTEGER
    }
}

impl FromValue for String {
    unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<String> {
        let c_text = ffi::sqlite3_value_text(v);
        if c_text.is_null() {
            Ok("".to_string())
        } else {
            let c_slice = CStr::from_ptr(c_text as *const c_char).to_bytes();
            let utf8_str = str::from_utf8(c_slice);
            utf8_str.map(|s| s.to_string())
                    .map_err(|e| {
                        SqliteError {
                            code: 0,
                            message: e.to_string(),
                        }
                    })
        }
    }

    unsafe fn parameter_has_valid_sqlite_type(v: *mut sqlite3_value) -> bool {
        sqlite3_value_type(v) == ffi::SQLITE_TEXT
    }
}

impl FromValue for Vec<u8> {
    unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<Vec<u8>> {
        use std::slice::from_raw_parts;
        let c_blob = ffi::sqlite3_value_blob(v);
        let len = ffi::sqlite3_value_bytes(v);

        assert!(len >= 0,
                "unexpected negative return from sqlite3_value_bytes");
        let len = len as usize;

        Ok(from_raw_parts(mem::transmute(c_blob), len).to_vec())
    }

    unsafe fn parameter_has_valid_sqlite_type(v: *mut sqlite3_value) -> bool {
        sqlite3_value_type(v) == ffi::SQLITE_BLOB
    }
}

impl<T: FromValue> FromValue for Option<T> {
    unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<Option<T>> {
        if sqlite3_value_type(v) == ffi::SQLITE_NULL {
            Ok(None)
        } else {
            FromValue::parameter_value(v).map(|t| Some(t))
        }
    }

    unsafe fn parameter_has_valid_sqlite_type(v: *mut sqlite3_value) -> bool {
        sqlite3_value_type(v) == ffi::SQLITE_NULL || T::parameter_has_valid_sqlite_type(v)
    }
}

// sqlite3_user_data
// sqlite3_get_auxdata
// sqlite3_set_auxdata

pub trait ScalarFunction: FnMut(*mut sqlite3_context, c_int, *mut *mut sqlite3_value) {}
impl<F: FnMut(*mut sqlite3_context, c_int, *mut *mut sqlite3_value)> ScalarFunction for F {}

impl SqliteConnection {
    pub fn create_scalar_function<F>(&self,
                                     fn_name: &str,
                                     n_arg: c_int,
                                     deterministic: bool,
                                     x_func: F)
                                     -> SqliteResult<()>
        where F: ScalarFunction
    {
        self.db.borrow_mut().create_scalar_function(fn_name, n_arg, deterministic, x_func)
    }
}

impl InnerSqliteConnection {
    pub fn create_scalar_function<F>(&mut self,
                                     fn_name: &str,
                                     n_arg: c_int,
                                     deterministic: bool,
                                     x_func: F)
                                     -> SqliteResult<()>
        where F: ScalarFunction
    {
        extern "C" fn free_boxed_closure<F>(p: *mut c_void)
            where F: ScalarFunction
        {
            let _: Box<F> = unsafe { Box::from_raw(mem::transmute(p)) };
        }

        extern "C" fn call_boxed_closure<F>(ctx: *mut sqlite3_context,
                                            argc: c_int,
                                            argv: *mut *mut sqlite3_value)
            where F: ScalarFunction
        {
            unsafe {
                let boxed_f: *mut F = mem::transmute(ffi::sqlite3_user_data(ctx));
                assert!(!boxed_f.is_null(), "Internal error - null function pointer");
                (*boxed_f)(ctx, argc, argv);
            }
        }

        let boxed_f: *mut F = Box::into_raw(Box::new(x_func));
        let c_name = try!(str_to_cstring(fn_name));
        let mut flags = ffi::SQLITE_UTF8;
        if deterministic {
            flags |= ffi::SQLITE_DETERMINISTIC;
        }
        let r = unsafe {
            ffi::sqlite3_create_function_v2(self.db(),
                                            c_name.as_ptr(),
                                            n_arg,
                                            flags,
                                            mem::transmute(boxed_f),
                                            Some(call_boxed_closure::<F>),
                                            None,
                                            None,
                                            Some(free_boxed_closure::<F>))
        };
        self.decode_result(r)
    }
}

#[cfg(test)]
mod test {
    extern crate regex;

    use std::boxed::Box;
    use std::ffi::CString;
    use std::mem;
    use libc::{c_int, c_double, c_void};
    use self::regex::Regex;

    use SqliteConnection;
    use ffi;
    use ffi::sqlite3_context;
    use ffi::sqlite3_value;
    use functions::{FromValue, ToResult};

    fn half(ctx: *mut sqlite3_context, _: c_int, argv: *mut *mut sqlite3_value) {
        unsafe {
            let arg = *argv.offset(0);
            if c_double::parameter_has_valid_sqlite_type(arg) {
                let value = c_double::parameter_value(arg).unwrap() / 2f64;
                value.set_result(ctx);
            } else {
                ffi::sqlite3_result_error_code(ctx, ffi::SQLITE_MISMATCH);
            }
        }
    }

    #[test]
    fn test_function_half() {
        let db = SqliteConnection::open_in_memory().unwrap();
        db.create_scalar_function("half", 1, true, half).unwrap();
        let result = db.query_row("SELECT half(6)", &[], |r| r.get::<f64>(0));

        assert_eq!(3f64, result.unwrap());
    }

    extern "C" fn regexp_free(raw: *mut c_void) {
        let _: Box<Regex> = unsafe { Box::from_raw(mem::transmute(raw)) };
    }

    fn regexp(ctx: *mut sqlite3_context, _: c_int, argv: *mut *mut sqlite3_value) {
        unsafe {
            let mut re_ptr = ffi::sqlite3_get_auxdata(ctx, 0) as *const Regex;
            let need_re = re_ptr.is_null();
            if need_re {
                let raw = String::parameter_value(*argv.offset(0));
                if raw.is_err() {
                    let msg = CString::new(format!("{}", raw.unwrap_err())).unwrap();
                    ffi::sqlite3_result_error(ctx, msg.as_ptr(), -1);
                    return;
                }
                let comp = Regex::new(raw.unwrap().as_ref());
                if comp.is_err() {
                    let msg = CString::new(format!("{}", comp.unwrap_err())).unwrap();
                    ffi::sqlite3_result_error(ctx, msg.as_ptr(), -1);
                    return;
                }
                let re = Box::new(comp.unwrap());
                re_ptr = Box::into_raw(re);
            }

            let text = String::parameter_value(*argv.offset(1));
            if text.is_ok() {
                let text = text.unwrap();
                (*re_ptr).is_match(text.as_ref()).set_result(ctx);
            }

            if need_re {
                ffi::sqlite3_set_auxdata(ctx, 0, mem::transmute(re_ptr), Some(regexp_free));
            }
        }
    }

    #[test]
    fn test_function_regexp() {
        let db = SqliteConnection::open_in_memory().unwrap();
        db.create_scalar_function("regexp", 2, true, regexp).unwrap();
        let result = db.query_row("SELECT regexp('l.s[aeiouy]', 'lisa')",
                                  &[],
                                  |r| r.get::<bool>(0));

        assert_eq!(true, result.unwrap());
    }
}
<commit_msg>Introduce Context wrapper for user-defined functions.<commit_after>//! Create or redefine SQL functions
use std::ffi::CStr;
use std::mem;
use std::str;
use libc::{c_int, c_double, c_char, c_void};

use ffi;
pub use ffi::sqlite3_context;
pub use ffi::sqlite3_value;
pub use ffi::sqlite3_value_type;
pub use ffi::sqlite3_value_numeric_type;

use types::Null;

use {SqliteResult, SqliteError, SqliteConnection, str_to_cstring, InnerSqliteConnection};

/// A trait for types that can be converted into the result of an SQL function.
pub trait ToResult {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context);
}

macro_rules! raw_to_impl(
    ($t:ty, $f:ident) => (
        impl ToResult for $t {
            unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
                ffi::$f(ctx, *self)
            }
        }
    )
);

raw_to_impl!(c_int, sqlite3_result_int);
raw_to_impl!(i64, sqlite3_result_int64);
raw_to_impl!(c_double, sqlite3_result_double);

impl<'a> ToResult for bool {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        match *self {
            true => ffi::sqlite3_result_int(ctx, 1),
            _ => ffi::sqlite3_result_int(ctx, 0),
        }
    }
}

impl<'a> ToResult for &'a str {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        let length = self.len();
        if length > ::std::i32::MAX as usize {
            ffi::sqlite3_result_error_toobig(ctx);
            return;
        }
        match str_to_cstring(self) {
            Ok(c_str) => {
                ffi::sqlite3_result_text(ctx,
                                         c_str.as_ptr(),
                                         length as c_int,
                                         ffi::SQLITE_TRANSIENT())
            }
            Err(_) => ffi::sqlite3_result_error_code(ctx, ffi::SQLITE_MISUSE), // TODO sqlite3_result_error
        }
    }
}

impl ToResult for String {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        (&self[..]).set_result(ctx)
    }
}

impl<'a> ToResult for &'a [u8] {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        if self.len() > ::std::i32::MAX as usize {
            ffi::sqlite3_result_error_toobig(ctx);
            return;
        }
        ffi::sqlite3_result_blob(ctx,
                                 mem::transmute(self.as_ptr()),
                                 self.len() as c_int,
                                 ffi::SQLITE_TRANSIENT())
    }
}

impl ToResult for Vec<u8> {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        (&self[..]).set_result(ctx)
    }
}

impl<T: ToResult> ToResult for Option<T> {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        match *self {
            None => ffi::sqlite3_result_null(ctx),
            Some(ref t) => t.set_result(ctx),
        }
    }
}

impl ToResult for Null {
    unsafe fn set_result(&self, ctx: *mut sqlite3_context) {
        ffi::sqlite3_result_null(ctx)
    }
}

// sqlite3_result_error_code, c_int
// sqlite3_result_error_nomem
// sqlite3_result_error_toobig
// sqlite3_result_error, *const c_char, c_int
// sqlite3_result_zeroblob
// sqlite3_result_value

/// A trait for types that can be created from a SQLite function parameter value.
pub trait FromValue: Sized {
    unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<Self>;

    /// FromValue types can implement this method and use sqlite3_value_type to check that
    /// the type reported by SQLite matches a type suitable for Self. This method is used
    /// by `???` to confirm that the parameter contains a valid type before
    /// attempting to retrieve the value.
    unsafe fn parameter_has_valid_sqlite_type(_: *mut sqlite3_value) -> bool {
        true
    }
}

macro_rules! raw_from_impl(
    ($t:ty, $f:ident, $c:expr) => (
        impl FromValue for $t {
            unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<$t> {
                Ok(ffi::$f(v))
            }

            unsafe fn parameter_has_valid_sqlite_type(v: *mut sqlite3_value) -> bool {
                sqlite3_value_numeric_type(v) == $c
            }
        }
    )
);

raw_from_impl!(c_int, sqlite3_value_int, ffi::SQLITE_INTEGER);
raw_from_impl!(i64, sqlite3_value_int64, ffi::SQLITE_INTEGER);

impl FromValue for bool {
    unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<bool> {
        match ffi::sqlite3_value_int(v) {
            0 => Ok(false),
            _ => Ok(true),
        }
    }

    unsafe fn parameter_has_valid_sqlite_type(v: *mut sqlite3_value) -> bool {
        sqlite3_value_numeric_type(v) == ffi::SQLITE_INTEGER
    }
}

impl FromValue for c_double {
    unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<c_double> {
        Ok(ffi::sqlite3_value_double(v))
    }

    unsafe fn parameter_has_valid_sqlite_type(v: *mut sqlite3_value) -> bool {
        sqlite3_value_numeric_type(v) == ffi::SQLITE_FLOAT ||
        sqlite3_value_numeric_type(v) == ffi::SQLITE_INTEGER
    }
}

impl FromValue for String {
    unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<String> {
        let c_text = ffi::sqlite3_value_text(v);
        if c_text.is_null() {
            Ok("".to_string())
        } else {
            let c_slice = CStr::from_ptr(c_text as *const c_char).to_bytes();
            let utf8_str = str::from_utf8(c_slice);
            utf8_str.map(|s| s.to_string())
                    .map_err(|e| {
                        SqliteError {
                            code: 0,
                            message: e.to_string(),
                        }
                    })
        }
    }

    unsafe fn parameter_has_valid_sqlite_type(v: *mut sqlite3_value) -> bool {
        sqlite3_value_type(v) == ffi::SQLITE_TEXT
    }
}

impl FromValue for Vec<u8> {
    unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<Vec<u8>> {
        use std::slice::from_raw_parts;
        let c_blob = ffi::sqlite3_value_blob(v);
        let len = ffi::sqlite3_value_bytes(v);

        assert!(len >= 0,
                "unexpected negative return from sqlite3_value_bytes");
        let len = len as usize;

        Ok(from_raw_parts(mem::transmute(c_blob), len).to_vec())
    }

    unsafe fn parameter_has_valid_sqlite_type(v: *mut sqlite3_value) -> bool {
        sqlite3_value_type(v) == ffi::SQLITE_BLOB
    }
}

impl<T: FromValue> FromValue for Option<T> {
    unsafe fn parameter_value(v: *mut sqlite3_value) -> SqliteResult<Option<T>> {
        if sqlite3_value_type(v) == ffi::SQLITE_NULL {
            Ok(None)
        } else {
            FromValue::parameter_value(v).map(|t| Some(t))
        }
    }

    unsafe fn parameter_has_valid_sqlite_type(v: *mut sqlite3_value) -> bool {
        sqlite3_value_type(v) == ffi::SQLITE_NULL || T::parameter_has_valid_sqlite_type(v)
    }
}

// sqlite3_user_data
// sqlite3_get_auxdata
// sqlite3_set_auxdata

unsafe extern "C" fn free_boxed_value<T>(p: *mut c_void) {
    let _: Box<T> = Box::from_raw(mem::transmute(p));
}

pub struct Context {
    pub ctx: *mut sqlite3_context,
}

impl Context {
    pub fn set_aux<T>(&self, arg: c_int, value: T) {
        let boxed = Box::into_raw(Box::new(value));
        unsafe {
            ffi::sqlite3_set_auxdata(self.ctx,
                                     arg,
                                     mem::transmute(boxed),
                                     Some(mem::transmute(free_boxed_value::<T>)))
        };
    }

    pub unsafe fn get_aux<T>(&self, arg: c_int) -> Option<&T> {
        let p = ffi::sqlite3_get_auxdata(self.ctx, arg) as *mut T;
        if p.is_null() {
            None
        } else {
            Some(&*p)
        }
    }
}

pub trait ScalarFunction: FnMut(&Context, c_int, *mut *mut sqlite3_value) {}
impl<F: FnMut(&Context, c_int, *mut *mut sqlite3_value)> ScalarFunction for F {}

impl SqliteConnection {
    pub fn create_scalar_function<F>(&self,
                                     fn_name: &str,
                                     n_arg: c_int,
                                     deterministic: bool,
                                     x_func: F)
                                     -> SqliteResult<()>
        where F: ScalarFunction
    {
        self.db.borrow_mut().create_scalar_function(fn_name, n_arg, deterministic, x_func)
    }
}

impl InnerSqliteConnection {
    pub fn create_scalar_function<F>(&mut self,
                                     fn_name: &str,
                                     n_arg: c_int,
                                     deterministic: bool,
                                     x_func: F)
                                     -> SqliteResult<()>
        where F: ScalarFunction
    {
        extern "C" fn call_boxed_closure<F>(ctx: *mut sqlite3_context,
                                            argc: c_int,
                                            argv: *mut *mut sqlite3_value)
            where F: ScalarFunction
        {
            let ctx = Context { ctx: ctx };
            unsafe {
                let boxed_f: *mut F = mem::transmute(ffi::sqlite3_user_data(ctx.ctx));
                assert!(!boxed_f.is_null(), "Internal error - null function pointer");
                (*boxed_f)(&ctx, argc, argv);
            }
        }

        let boxed_f: *mut F = Box::into_raw(Box::new(x_func));
        let c_name = try!(str_to_cstring(fn_name));
        let mut flags = ffi::SQLITE_UTF8;
        if deterministic {
            flags |= ffi::SQLITE_DETERMINISTIC;
        }
        let r = unsafe {
            ffi::sqlite3_create_function_v2(self.db(),
                                            c_name.as_ptr(),
                                            n_arg,
                                            flags,
                                            mem::transmute(boxed_f),
                                            Some(call_boxed_closure::<F>),
                                            None,
                                            None,
                                            Some(mem::transmute(free_boxed_value::<F>)))
        };
        self.decode_result(r)
    }
}

#[cfg(test)]
mod test {
    extern crate regex;

    use std::ffi::CString;
    use libc::{c_int, c_double};
    use self::regex::Regex;

    use SqliteConnection;
    use ffi;
    use ffi::sqlite3_value;
    use functions::{Context, FromValue, ToResult};

    fn half(ctx: &Context, _: c_int, argv: *mut *mut sqlite3_value) {
        unsafe {
            let arg = *argv.offset(0);
            if c_double::parameter_has_valid_sqlite_type(arg) {
                let value = c_double::parameter_value(arg).unwrap() / 2f64;
                value.set_result(ctx.ctx);
            } else {
                ffi::sqlite3_result_error_code(ctx.ctx, ffi::SQLITE_MISMATCH);
            }
        }
    }

    #[test]
    fn test_function_half() {
        let db = SqliteConnection::open_in_memory().unwrap();
        db.create_scalar_function("half", 1, true, half).unwrap();
        let result = db.query_row("SELECT half(6)", &[], |r| r.get::<f64>(0));

        assert_eq!(3f64, result.unwrap());
    }

    fn regexp(ctx: &Context, _: c_int, argv: *mut *mut sqlite3_value) {
        let saved_re: Option<&Regex> = unsafe { ctx.get_aux(0) };
        let new_re = match saved_re {
            None => unsafe {
                let raw = String::parameter_value(*argv.offset(0));
                if raw.is_err() {
                    let msg = CString::new(format!("{}", raw.unwrap_err())).unwrap();
                    ffi::sqlite3_result_error(ctx.ctx, msg.as_ptr(), -1);
                    return;
                }
                let comp = Regex::new(raw.unwrap().as_ref());
                if comp.is_err() {
                    let msg = CString::new(format!("{}", comp.unwrap_err())).unwrap();
                    ffi::sqlite3_result_error(ctx.ctx, msg.as_ptr(), -1);
                    return;
                }
                Some(comp.unwrap())
            },
            Some(_) => None,
        };

        {
            let re = saved_re.unwrap_or_else(|| new_re.as_ref().unwrap());

            let text = unsafe { String::parameter_value(*argv.offset(1)) };
            if text.is_ok() {
                let text = text.unwrap();
                unsafe {
                    re.is_match(text.as_ref()).set_result(ctx.ctx);
                }
            } else {
                let msg = CString::new(format!("{}", text.unwrap_err())).unwrap();
                unsafe {
                    ffi::sqlite3_result_error(ctx.ctx, msg.as_ptr(), -1);
                }
                return;
            }
        }

        if let Some(re) = new_re {
            ctx.set_aux(0, re);
        }
    }

    #[test]
    fn test_function_regexp() {
        let db = SqliteConnection::open_in_memory().unwrap();
        db.create_scalar_function("regexp", 2, true, regexp).unwrap();
        let result = db.query_row("SELECT regexp('l.s[aeiouy]', 'lisa')",
                                  &[],
                                  |r| r.get::<bool>(0));

        assert_eq!(true, result.unwrap());
    }
}
<|endoftext|>