<commit_before>
extern crate ndarray;

use ndarray::prelude::*;

#[test]
fn broadcast_1()
{
    let a_dim = (2, 4, 2, 2);
    let b_dim = (2, 1, 2, 1);
    let a = RcArray::linspace(0., 1., a_dim.size()).reshape(a_dim);
    let b = RcArray::linspace(0., 1., b_dim.size()).reshape(b_dim);
    assert!(b.broadcast(a.dim()).is_some());

    let c_dim = (2, 1);
    let c = RcArray::linspace(0., 1., c_dim.size()).reshape(c_dim);
    assert!(c.broadcast(1).is_none());
    assert!(c.broadcast(()).is_none());
    assert!(c.broadcast((2, 1)).is_some());
    assert!(c.broadcast((2, 2)).is_some());
    assert!(c.broadcast((32, 2, 1)).is_some());
    assert!(c.broadcast((32, 1, 2)).is_none());

    /* () can be broadcast to anything */
    let z = RcArray::<f32,_>::zeros(());
    assert!(z.broadcast(()).is_some());
    assert!(z.broadcast(1).is_some());
    assert!(z.broadcast(3).is_some());
    assert!(z.broadcast((7,2,9)).is_some());
}

#[test]
fn test_add()
{
    let a_dim = (2, 4, 2, 2);
    let b_dim = (2, 1, 2, 1);
    let mut a = RcArray::linspace(0.0, 1., a_dim.size()).reshape(a_dim);
    let b = RcArray::linspace(0.0, 1., b_dim.size()).reshape(b_dim);
    a += &b;
    let t = RcArray::from_elem((), 1.0f32);
    a += &t;
}

#[test] #[should_panic]
fn test_add_incompat()
{
    let a_dim = (2, 4, 2, 2);
    let mut a = RcArray::linspace(0.0, 1., a_dim.size()).reshape(a_dim);
    let incompat = RcArray::from_elem(3, 1.0f32);
    a += &incompat;
}

#[test]
fn test_broadcast() {
    let (_, n, k) = (16, 16, 16);
    let x1 = 1.;
    // b0 broadcast 1 -> n, k
    let x = Array::from_vec(vec![x1]);
    let b0 = x.broadcast((n, k)).unwrap();
    // b1 broadcast n -> n, k
    let b1 = Array::from_elem(n, x1);
    let b1 = b1.broadcast((n, k)).unwrap();
    // b2 is n, k
    let b2 = Array::from_elem((n, k), x1);

    println!("b0=\n{:?}", b0);
    println!("b1=\n{:?}", b1);
    println!("b2=\n{:?}", b2);
    assert_eq!(b0, b1);
    assert_eq!(b0, b2);
}

#[test]
fn test_broadcast_1d() {
    let n = 16;
    let x1 = 1.;
    // b0 broadcast 1 -> n
    let x = Array::from_vec(vec![x1]);
    let b0 = x.broadcast(n).unwrap();
    let b2 = Array::from_elem(n, x1);

    println!("b0=\n{:?}", b0);
    println!("b2=\n{:?}", b2);
    assert_eq!(b0, b2);
}
<commit_msg>Update tests/broadcast.rs<commit_after>
extern crate ndarray;

use ndarray::prelude::*;

#[test]
fn broadcast_1()
{
    let a_dim = (2, 4, 2, 2);
    let b_dim = (2, 1, 2, 1);
    let a = RcArray::linspace(0., 1., a_dim.size()).reshape(a_dim);
    let b = RcArray::linspace(0., 1., b_dim.size()).reshape(b_dim);
    assert!(b.broadcast(a.dim()).is_some());

    let c_dim = [2, 1];
    let c = RcArray::linspace(0., 1., c_dim.size()).reshape(c_dim);
    assert!(c.broadcast(1).is_none());
    assert!(c.broadcast(()).is_none());
    assert!(c.broadcast((2, 1)).is_some());
    assert!(c.broadcast((2, 2)).is_some());
    assert!(c.broadcast((32, 2, 1)).is_some());
    assert!(c.broadcast((32, 1, 2)).is_none());

    /* () can be broadcast to anything */
    let z = RcArray::<f32,_>::zeros(());
    assert!(z.broadcast(()).is_some());
    assert!(z.broadcast(1).is_some());
    assert!(z.broadcast(3).is_some());
    assert!(z.broadcast((7,2,9)).is_some());
}

#[test]
fn test_add()
{
    let a_dim = (2, 4, 2, 2);
    let b_dim = (2, 1, 2, 1);
    let mut a = RcArray::linspace(0.0, 1., a_dim.size()).reshape(a_dim);
    let b = RcArray::linspace(0.0, 1., b_dim.size()).reshape(b_dim);
    a += &b;
    let t = RcArray::from_elem((), 1.0f32);
    a += &t;
}

#[test] #[should_panic]
fn test_add_incompat()
{
    let a_dim = (2, 4, 2, 2);
    let mut a = RcArray::linspace(0.0, 1., a_dim.size()).reshape(a_dim);
    let incompat = RcArray::from_elem(3, 1.0f32);
    a += &incompat;
}

#[test]
fn test_broadcast() {
    let (_, n, k) = (16, 16, 16);
    let x1 = 1.;
    // b0 broadcast 1 -> n, k
    let x = Array::from_vec(vec![x1]);
    let b0 = x.broadcast((n, k)).unwrap();
    // b1 broadcast n -> n, k
    let b1 = Array::from_elem(n, x1);
    let b1 = b1.broadcast((n, k)).unwrap();
    // b2 is n, k
    let b2 = Array::from_elem((n, k), x1);

    println!("b0=\n{:?}", b0);
    println!("b1=\n{:?}", b1);
    println!("b2=\n{:?}", b2);
    assert_eq!(b0, b1);
    assert_eq!(b0, b2);
}

#[test]
fn test_broadcast_1d() {
    let n = 16;
    let x1 = 1.;
    // b0 broadcast 1 -> n
    let x = Array::from_vec(vec![x1]);
    let b0 = x.broadcast(n).unwrap();
    let b2 = Array::from_elem(n, x1);

    println!("b0=\n{:?}", b0);
    println!("b2=\n{:?}", b2);
    assert_eq!(b0, b2);
}
<|endoftext|>


<commit_before>use crate::mem::{size_of, transmute_copy};
use crate::ptr::write_bytes;

pub(super) trait SpecFill<T> {
    fn spec_fill(&mut self, value: T);
}

impl<T: Clone> SpecFill<T> for [T] {
    default fn spec_fill(&mut self, value: T) {
        if let Some((last, elems)) = self.split_last_mut() {
            for el in elems {
                el.clone_from(&value);
            }

            *last = value
        }
    }
}

impl<T: Copy> SpecFill<T> for [T] {
    fn spec_fill(&mut self, value: T) {
        if size_of::<T>() == 1 {
            // SAFETY: The size_of check above ensures that values are 1 byte wide, as required
            // for the transmute and write_bytes
            unsafe {
                let value: u8 = transmute_copy(&value);
                write_bytes(self.as_mut_ptr(), value, self.len());
            }
        } else {
            for item in self.iter_mut() {
                *item = value;
            }
        }
    }
}
<commit_msg>Auto merge of #87892 - rust-lang:spec-fill-size-one-bye, r=the8472<commit_after>pub(super) trait SpecFill<T> {
    fn spec_fill(&mut self, value: T);
}

impl<T: Clone> SpecFill<T> for [T] {
    default fn spec_fill(&mut self, value: T) {
        if let Some((last, elems)) = self.split_last_mut() {
            for el in elems {
                el.clone_from(&value);
            }

            *last = value
        }
    }
}

impl<T: Copy> SpecFill<T> for [T] {
    fn spec_fill(&mut self, value: T) {
        for item in self.iter_mut() {
            *item = value;
        }
    }
}
<|endoftext|>


<commit_before>//! Handling column descriptions
//!
//! Columns are represented as
//! [`ConcreteColumnDescription`](struct.ConcreteColumnDescription.html). This is constructed
//! through the builder pattern, by creating a [`ColumnDescription`](struct.ColumnDescription.html)
//! and calling [`create`](struct.ColumnDescription.html#method.create)
use errors::Result;
use std::str::FromStr;

/// Description for new columns
#[derive(Debug, Clone)]
pub struct ColumnDescription {
    /// Name of the column
    pub name: String,

    /// Type of the data, see the cfitsio documentation
    pub data_type: Option<ColumnDataDescription>,
}

/// Concrete representation of the description of a column
#[derive(Debug, Clone)]
pub struct ConcreteColumnDescription {
    /// Name of the column
    pub name: String,

    /// Type of the data, see the cfitsio documentation
    pub data_type: ColumnDataDescription,
}

impl ColumnDescription {
    /// Create a new [`ColumnDescription`](struct.ColumnDescription.html) from a name
    pub fn new<T: Into<String>>(name: T) -> Self {
        ColumnDescription {
            name: name.into(),
            data_type: None,
        }
    }

    /// Add a data type to the column description
    pub fn with_type(&mut self, typ: ColumnDataType) -> &mut ColumnDescription {
        self.data_type = Some(ColumnDataDescription::scalar(typ));
        self
    }

    /// Make the column repeat
    pub fn that_repeats(&mut self, repeat: usize) -> &mut ColumnDescription {
        if let Some(ref mut desc) = self.data_type {
            desc.repeat = repeat;
        }
        self
    }

    /// Define the column width
    pub fn with_width(&mut self, width: usize) -> &mut ColumnDescription {
        if let Some(ref mut desc) = self.data_type {
            desc.width = width;
        }
        self
    }

    /// Render the [`ColumnDescription`](struct.ColumnDescription.html) into a
    /// [`ConcreteColumnDescription`](struct.ConcreteColumnDescription.html)
    pub fn create(&self) -> Result<ConcreteColumnDescription> {
        match self.data_type {
            Some(ref d) => {
                Ok(ConcreteColumnDescription {
                    name: self.name.clone(),
                    data_type: d.clone(),
                })
            }
            None => {
                Err(
                    "No data type given. Ensure the `with_type` method has been called.".into(),
                )
            }
        }
    }
}

/// Description of the column data
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ColumnDataDescription {
    /// Does the column contain multiple values?
    pub repeat: usize,

    /// How wide is the column?
    pub width: usize,

    /// What data type does the column store?
    pub typ: ColumnDataType,
}

impl ColumnDataDescription {
    /// Create a new column data description
    pub fn new(typ: ColumnDataType, repeat: usize, width: usize) -> Self {
        ColumnDataDescription {
            repeat: repeat,
            width: width,
            typ: typ,
        }
    }

    /// Shortcut for creating a scalar column
    pub fn scalar(typ: ColumnDataType) -> Self {
        ColumnDataDescription::new(typ, 1, 1)
    }

    /// Shortcut for creating a vector column
    pub fn vector(typ: ColumnDataType, repeat: usize) -> Self {
        ColumnDataDescription::new(typ, repeat, 1)
    }

    /// Set the repeat count
    /* XXX These two methods force a call to clone which is wasteful of memory. I do not know if
     * this means that memory is leaked, or that destructors are needlessly called (I suspect the
     * latter) but it is fairly wasteful. On the other hand, it's unlikely this sort of thing will
     * be called in performance-critical code, and is more likely a one-time definition. I will
     * leave it for now - SRW 2017-03-07
     * */
    pub fn repeats(&mut self, repeat: usize) -> Result<Self> {
        if repeat == 0 {
            return Err("repeat parameter must be > 0".into());
        } else {
            self.repeat = repeat;
            Ok(self.clone())
        }
    }

    /// Set the width of the column
    pub fn width(&mut self, width: usize) -> Result<Self> {
        if width == 0 {
            return Err("width parameter must be > 0".into())
        } else {
            self.width = width;
            Ok(self.clone())
        }
    }
}

impl From<ColumnDataDescription> for String {
    fn from(orig: ColumnDataDescription) -> String {
        match orig.typ {
            ColumnDataType::Text => {
                if orig.width > 1 {
                    format!(
                        "{repeat}{data_type}{width}",
                        data_type = String::from(orig.typ),
                        repeat = orig.repeat,
                        width = orig.width
                    )
                } else {
                    format!(
                        "{repeat}{data_type}",
                        data_type = String::from(orig.typ),
                        repeat = orig.repeat
                    )
                }
            }
            _ => {
                format!(
                    "{repeat}{data_type}",
                    data_type = String::from(orig.typ),
                    repeat = orig.repeat
                )
            }
        }
    }
}

/// Types a column can represent
#[allow(missing_docs)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ColumnDataType {
    Int,
    Float,
    Text,
    Double,
    Short,
    Long,
    String,
}

impl From<ColumnDataType> for String {
    fn from(orig: ColumnDataType) -> String {
        use self::ColumnDataType::*;

        match orig {
            Int => "J",
            Float => "E",
            Text | String => "A",
            Double => "D",
            Short => "I",
            Long => "K",
        }.to_string()
    }
}

impl FromStr for ColumnDataDescription {
    type Err = Box<::std::error::Error>;

    fn from_str(s: &str) -> ::std::result::Result<Self, Self::Err> {
        let chars: Vec<_> = s.chars().collect();

        let mut repeat_str = Vec::new();
        let mut last_position = 0;
        for c in &chars {
            if c.is_digit(10) {
                repeat_str.push(c);
                last_position += 1;
            } else {
                break;
            }
        }

        let repeat = if repeat_str.is_empty() {
            1
        } else {
            /* TODO: in nightly the following line works
            let repeat_str: String = repeat_str.into_iter().collect(); */
            let repeat_str: String = repeat_str.into_iter().cloned().collect();
            repeat_str.parse::<usize>()?
        };

        let data_type_char = chars[last_position];
        last_position += 1;

        let mut width_str = Vec::new();
        for c in chars.iter().skip(last_position) {
            if c.is_digit(10) {
                width_str.push(c);
            } else {
                break;
            }
        }

        /* TODO: validate that the whole string has been used up */

        let width = if width_str.is_empty() {
            1
        } else {
            /* TODO: in nightly the following line works
            let width_str: String = width_str.into_iter().collect(); */
            let width_str: String = width_str.into_iter().cloned().collect();
            width_str.parse::<usize>()?
        };

        let data_type = match data_type_char {
            'E' => ColumnDataType::Float,
            'J' => ColumnDataType::Int,
            'D' => ColumnDataType::Double,
            'I' => ColumnDataType::Short,
            'K' => ColumnDataType::Long,
            'A' => ColumnDataType::String,
            _ => {
                panic!(
                    "Have not implemented str -> ColumnDataType for {}",
                    data_type_char
                )
            }
        };

        Ok(ColumnDataDescription {
            repeat: repeat,
            typ: data_type,
            width: width,
        })
    }
}

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

    #[test]
    fn test_column_data_descriptions_builder_pattern() {
        let desc = ColumnDataDescription::scalar(ColumnDataType::Int)
            .width(100)
            .and_then(|mut d| d.repeats(5))
            .unwrap();
        assert_eq!(desc.repeat, 5);
        assert_eq!(desc.width, 100);
    }

    #[test]
    fn width_after_repeates() {
        let desc = ColumnDataDescription::scalar(ColumnDataType::Int)
            .repeats(5)
            .and_then(|mut d| d.width(100))
            .unwrap();
        assert_eq!(desc.repeat, 5);
        assert_eq!(desc.width, 100);
    }

    #[test]
    fn from_impls() {
        {
            let desc = ColumnDataDescription::scalar(ColumnDataType::Int).repeats(5).unwrap();
            assert_eq!(String::from(desc), "5J");
        }

        {
            let desc = ColumnDataDescription::scalar(ColumnDataType::Float);
            assert_eq!(String::from(desc), "1E");
        }

        {
            let desc = ColumnDataDescription::scalar(ColumnDataType::Text).width(100).unwrap();
            assert_eq!(String::from(desc), "1A100");
        }
    }

    #[test]
    fn parsing() {
        let s = "1E";
        assert_eq!(
            s.parse::<ColumnDataDescription>().unwrap(),
            ColumnDataDescription {
                repeat: 1,
                width: 1,
                typ: ColumnDataType::Float,
            }
        );
    }

    #[test]
    fn parse_many_repeats() {
        let s = "100E";
        assert_eq!(
            s.parse::<ColumnDataDescription>().unwrap(),
            ColumnDataDescription {
                repeat: 100,
                width: 1,
                typ: ColumnDataType::Float,
            }
        );
    }

    #[test]
    fn parse_with_width() {
        let s = "1E26";
        assert_eq!(
            s.parse::<ColumnDataDescription>().unwrap(),
            ColumnDataDescription {
                repeat: 1,
                width: 26,
                typ: ColumnDataType::Float,
            }
        );
    }

    #[test]
    fn creating_data_description() {
        let concrete_desc = ColumnDescription::new("FOO")
            .with_type(ColumnDataType::Int)
            .that_repeats(10)
            .create()
            .unwrap();
        assert_eq!(concrete_desc.name, "FOO".to_string());
        assert_eq!(concrete_desc.data_type.repeat, 10);
        assert_eq!(concrete_desc.data_type.width, 1);

        /* Do not call `with_type` */
        let bad_desc = ColumnDescription::new("FOO").create();
        assert!(bad_desc.is_err());
    }
}
<commit_msg>Remove nightly-only features<commit_after>//! Handling column descriptions
//!
//! Columns are represented as
//! [`ConcreteColumnDescription`](struct.ConcreteColumnDescription.html). This is constructed
//! through the builder pattern, by creating a [`ColumnDescription`](struct.ColumnDescription.html)
//! and calling [`create`](struct.ColumnDescription.html#method.create)
use errors::Result;
use std::str::FromStr;

/// Description for new columns
#[derive(Debug, Clone)]
pub struct ColumnDescription {
    /// Name of the column
    pub name: String,

    /// Type of the data, see the cfitsio documentation
    pub data_type: Option<ColumnDataDescription>,
}

/// Concrete representation of the description of a column
#[derive(Debug, Clone)]
pub struct ConcreteColumnDescription {
    /// Name of the column
    pub name: String,

    /// Type of the data, see the cfitsio documentation
    pub data_type: ColumnDataDescription,
}

impl ColumnDescription {
    /// Create a new [`ColumnDescription`](struct.ColumnDescription.html) from a name
    pub fn new<T: Into<String>>(name: T) -> Self {
        ColumnDescription {
            name: name.into(),
            data_type: None,
        }
    }

    /// Add a data type to the column description
    pub fn with_type(&mut self, typ: ColumnDataType) -> &mut ColumnDescription {
        self.data_type = Some(ColumnDataDescription::scalar(typ));
        self
    }

    /// Make the column repeat
    pub fn that_repeats(&mut self, repeat: usize) -> &mut ColumnDescription {
        if let Some(ref mut desc) = self.data_type {
            desc.repeat = repeat;
        }
        self
    }

    /// Define the column width
    pub fn with_width(&mut self, width: usize) -> &mut ColumnDescription {
        if let Some(ref mut desc) = self.data_type {
            desc.width = width;
        }
        self
    }

    /// Render the [`ColumnDescription`](struct.ColumnDescription.html) into a
    /// [`ConcreteColumnDescription`](struct.ConcreteColumnDescription.html)
    pub fn create(&self) -> Result<ConcreteColumnDescription> {
        match self.data_type {
            Some(ref d) => {
                Ok(ConcreteColumnDescription {
                    name: self.name.clone(),
                    data_type: d.clone(),
                })
            }
            None => {
                Err(
                    "No data type given. Ensure the `with_type` method has been called.".into(),
                )
            }
        }
    }
}

/// Description of the column data
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ColumnDataDescription {
    /// Does the column contain multiple values?
    pub repeat: usize,

    /// How wide is the column?
    pub width: usize,

    /// What data type does the column store?
    pub typ: ColumnDataType,
}

impl ColumnDataDescription {
    /// Create a new column data description
    pub fn new(typ: ColumnDataType, repeat: usize, width: usize) -> Self {
        ColumnDataDescription {
            repeat: repeat,
            width: width,
            typ: typ,
        }
    }

    /// Shortcut for creating a scalar column
    pub fn scalar(typ: ColumnDataType) -> Self {
        ColumnDataDescription::new(typ, 1, 1)
    }

    /// Shortcut for creating a vector column
    pub fn vector(typ: ColumnDataType, repeat: usize) -> Self {
        ColumnDataDescription::new(typ, repeat, 1)
    }

    /// Set the repeat count
    /* XXX These two methods force a call to clone which is wasteful of memory. I do not know if
     * this means that memory is leaked, or that destructors are needlessly called (I suspect the
     * latter) but it is fairly wasteful. On the other hand, it's unlikely this sort of thing will
     * be called in performance-critical code, and is more likely a one-time definition. I will
     * leave it for now - SRW 2017-03-07
     * */
    pub fn repeats(&mut self, repeat: usize) -> Result<Self> {
        if repeat == 0 {
            return Err("repeat parameter must be > 0".into());
        } else {
            self.repeat = repeat;
            Ok(self.clone())
        }
    }

    /// Set the width of the column
    pub fn width(&mut self, width: usize) -> Result<Self> {
        if width == 0 {
            return Err("width parameter must be > 0".into())
        } else {
            self.width = width;
            Ok(self.clone())
        }
    }
}

impl From<ColumnDataDescription> for String {
    fn from(orig: ColumnDataDescription) -> String {
        match orig.typ {
            ColumnDataType::Text => {
                if orig.width > 1 {
                    format!(
                        "{repeat}{data_type}{width}",
                        data_type = String::from(orig.typ),
                        repeat = orig.repeat,
                        width = orig.width
                    )
                } else {
                    format!(
                        "{repeat}{data_type}",
                        data_type = String::from(orig.typ),
                        repeat = orig.repeat
                    )
                }
            }
            _ => {
                format!(
                    "{repeat}{data_type}",
                    data_type = String::from(orig.typ),
                    repeat = orig.repeat
                )
            }
        }
    }
}

/// Types a column can represent
#[allow(missing_docs)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ColumnDataType {
    Int,
    Float,
    Text,
    Double,
    Short,
    Long,
    String,
}

impl From<ColumnDataType> for String {
    fn from(orig: ColumnDataType) -> String {
        use self::ColumnDataType::*;

        match orig {
            Int => "J",
            Float => "E",
            Text | String => "A",
            Double => "D",
            Short => "I",
            Long => "K",
        }.to_string()
    }
}

impl FromStr for ColumnDataDescription {
    type Err = Box<::std::error::Error>;

    fn from_str(s: &str) -> ::std::result::Result<Self, Self::Err> {
        let chars: Vec<_> = s.chars().collect();

        let mut repeat_str = Vec::new();
        let mut last_position = 0;
        for c in &chars {
            if c.is_digit(10) {
                repeat_str.push(c);
                last_position += 1;
            } else {
                break;
            }
        }

        let repeat = if repeat_str.is_empty() {
            1
        } else {
            let repeat_str: String = repeat_str.into_iter().collect();
            repeat_str.parse::<usize>()?
        };

        let data_type_char = chars[last_position];
        last_position += 1;

        let mut width_str = Vec::new();
        for c in chars.iter().skip(last_position) {
            if c.is_digit(10) {
                width_str.push(c);
            } else {
                break;
            }
        }

        /* TODO: validate that the whole string has been used up */

        let width = if width_str.is_empty() {
            1
        } else {
            let width_str: String = width_str.into_iter().collect();
            width_str.parse::<usize>()?
        };

        let data_type = match data_type_char {
            'E' => ColumnDataType::Float,
            'J' => ColumnDataType::Int,
            'D' => ColumnDataType::Double,
            'I' => ColumnDataType::Short,
            'K' => ColumnDataType::Long,
            'A' => ColumnDataType::String,
            _ => {
                panic!(
                    "Have not implemented str -> ColumnDataType for {}",
                    data_type_char
                )
            }
        };

        Ok(ColumnDataDescription {
            repeat: repeat,
            typ: data_type,
            width: width,
        })
    }
}

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

    #[test]
    fn test_column_data_descriptions_builder_pattern() {
        let desc = ColumnDataDescription::scalar(ColumnDataType::Int)
            .width(100)
            .and_then(|mut d| d.repeats(5))
            .unwrap();
        assert_eq!(desc.repeat, 5);
        assert_eq!(desc.width, 100);
    }

    #[test]
    fn width_after_repeates() {
        let desc = ColumnDataDescription::scalar(ColumnDataType::Int)
            .repeats(5)
            .and_then(|mut d| d.width(100))
            .unwrap();
        assert_eq!(desc.repeat, 5);
        assert_eq!(desc.width, 100);
    }

    #[test]
    fn from_impls() {
        {
            let desc = ColumnDataDescription::scalar(ColumnDataType::Int).repeats(5).unwrap();
            assert_eq!(String::from(desc), "5J");
        }

        {
            let desc = ColumnDataDescription::scalar(ColumnDataType::Float);
            assert_eq!(String::from(desc), "1E");
        }

        {
            let desc = ColumnDataDescription::scalar(ColumnDataType::Text).width(100).unwrap();
            assert_eq!(String::from(desc), "1A100");
        }
    }

    #[test]
    fn parsing() {
        let s = "1E";
        assert_eq!(
            s.parse::<ColumnDataDescription>().unwrap(),
            ColumnDataDescription {
                repeat: 1,
                width: 1,
                typ: ColumnDataType::Float,
            }
        );
    }

    #[test]
    fn parse_many_repeats() {
        let s = "100E";
        assert_eq!(
            s.parse::<ColumnDataDescription>().unwrap(),
            ColumnDataDescription {
                repeat: 100,
                width: 1,
                typ: ColumnDataType::Float,
            }
        );
    }

    #[test]
    fn parse_with_width() {
        let s = "1E26";
        assert_eq!(
            s.parse::<ColumnDataDescription>().unwrap(),
            ColumnDataDescription {
                repeat: 1,
                width: 26,
                typ: ColumnDataType::Float,
            }
        );
    }

    #[test]
    fn creating_data_description() {
        let concrete_desc = ColumnDescription::new("FOO")
            .with_type(ColumnDataType::Int)
            .that_repeats(10)
            .create()
            .unwrap();
        assert_eq!(concrete_desc.name, "FOO".to_string());
        assert_eq!(concrete_desc.data_type.repeat, 10);
        assert_eq!(concrete_desc.data_type.width, 1);

        /* Do not call `with_type` */
        let bad_desc = ColumnDescription::new("FOO").create();
        assert!(bad_desc.is_err());
    }
}
<|endoftext|>


<commit_before><commit_msg>refactor: ignore unused error<commit_after><|endoftext|>


<commit_before>//! Define a system components to create.
//!
//! This interface could use a lot of improvement.

use database::{AvailableComponents, DataBase};
use error::{GrafenCliError, Result, UIErrorKind};
use ui::utils;
use ui::utils::{CommandList, CommandParser};

use grafen::system::Coord;
use std::error::Error;

#[derive(Clone, Copy, Debug)]
/// User commands for defining the system.
enum Command {
    DefineSystem,
    RemoveSystem,
    SwapSystems,
    QuitAndSave,
    QuitWithoutSaving,
}

/// Edit the list of system definitions to construct from.
pub fn user_menu(database: &DataBase, mut system_defs: &mut Vec<AvailableComponents>)
        -> Result<()> {
    let command_list: CommandList<Command> = vec![
        ("d", Command::DefineSystem, "Define a system to create"),
        ("r", Command::RemoveSystem, "Remove a system from the list"),
        ("s", Command::SwapSystems, "Swap the order of two systems"),
        ("f", Command::QuitAndSave, "Finalize editing and return"),
        ("a", Command::QuitWithoutSaving, "Abort and discard changes to list")
    ];
    let commands = CommandParser::from_list(command_list);

    let backup = system_defs.clone();

    loop {
        describe_system_definitions(&system_defs);
        commands.print_menu();
        let input = utils::get_input_string("Selection")?;
        println!("");

        if let Some((cmd, tail)) = commands.get_selection_and_tail(&input) {
            match cmd {
                Command::DefineSystem => {
                    match create_definition(&database) {
                        Ok(def) => system_defs.push(def),
                        Err(err) => println!("Could not create definition: {}", err.description()),
                    }
                },
                Command::RemoveSystem => {
                    match utils::remove_item(&mut system_defs, &tail) {
                        Ok(i) => println!("Removed system at index {}.", i),
                        Err(err) => println!("Could not remove system: {}", err.description()),
                    }
                },
                Command::SwapSystems => {
                    match utils::swap_items(&mut system_defs, &tail) {
                        Ok((i, j)) => println!("Swapped system at index {} with system at {}.",
                                               i, j),
                        Err(err) => println!("Could not swap systems: {}", err.description()),
                    }
                },
                Command::QuitAndSave => {
                    return Ok(());
                },
                Command::QuitWithoutSaving => {
                    system_defs.clear();
                    system_defs.extend_from_slice(&backup);

                    return Ok(());
                },
            }
        } else {
            println!("Not a valid selection.");
        }

        println!("");
    }
}

/// Print the current system definitions to stdout.
pub fn describe_system_definitions(system_defs: &[AvailableComponents]) {
    if system_defs.is_empty() {
        println!("(No systems have been defined)");
    } else {
        println!("System definitions:");
        for (i, def) in system_defs.iter().enumerate() {
            println!("{}. {}", i, def.describe_long());
        }
    }

    println!("");
}

fn create_definition(database: &DataBase) -> Result<AvailableComponents> {
    //let mut definition = select_substrate(&database).map(|def| def.clone())?;
    let mut definition = select_component(&database).map(|def| def.clone())?;

    match &mut definition {
        &mut AvailableComponents::Sheet(ref mut conf) => {
            let position = select_position()?;
            let size = select_size()?;

            conf.position = Some(position);
            conf.size = Some(size);
        },
        &mut AvailableComponents::Cylinder(ref mut conf) => {
            let position = select_position()?;
            let radius = utils::get_and_parse_string_single("Set radius")?;
            let height = utils::get_and_parse_string_single("Set height")?;

            conf.position = Some(position);
            conf.radius = Some(radius);
            conf.height = Some(height);
        },
    }

    Ok(definition)
}

fn select_component(database: &DataBase) -> Result<&AvailableComponents> {
    println!("Available components:");
    for (i, sub) in database.component_defs.iter().enumerate() {
        println!("{}. {}", i, &sub.describe());
    }
    println!("");

    let selection = utils::get_input_string("Select component")?;
    let index = utils::parse_string_for_index(&selection, &database.component_defs)?;

    database.component_defs
        .get(index)
        .ok_or(GrafenCliError::UIError(format!("'{}' is not a valid index", &selection)))

        /*
    selection
        .parse::<usize>()
        .map_err(|_| UIErrorKind::BadValue(format!("'{}' is not a valid index", &selection)))
        .and_then(|n| {
            database.component_defs
                .get(n)
                .map(|def| def.clone())
                .ok_or(UIErrorKind::BadValue(format!("No component with index {} exists", n)))
        })
        .map_err(|err| GrafenCliError::from(err))
        */
}

fn select_position() -> Result<Coord> {
    let selection = utils::get_input_string("Change position (default: (0.0, 0.0, 0.0))")?;
    if selection.is_empty() {
        return Ok(Coord::new(0.0, 0.0, 0.0));
    }

    let coords = utils::parse_string(&selection)?;
    let &x = coords.get(0).ok_or(UIErrorKind::BadValue("3 positions are required".to_string()))?;
    let &y = coords.get(1).ok_or(UIErrorKind::BadValue("3 positions are required".to_string()))?;
    let &z = coords.get(2).ok_or(UIErrorKind::BadValue("3 positions are required".to_string()))?;

    Ok(Coord::new(x, y, z))
}

fn select_size() -> Result<(f64, f64)> {
    let size = utils::get_and_parse_string("Set size")?;
    let &dx = size.get(0).ok_or(UIErrorKind::BadValue("2 values are required".to_string()))?;
    let &dy = size.get(1).ok_or(UIErrorKind::BadValue("2 values are required".to_string()))?;

    Ok((dx, dy))
}
<commit_msg>Rework `create_definition` to use better matching<commit_after>//! Define a system components to create.
//!
//! This interface could use a lot of improvement.

use database::{AvailableComponents, DataBase};
use error::{GrafenCliError, Result, UIErrorKind};
use ui::utils;
use ui::utils::CommandParser;

use grafen::system::Coord;
use std::error::Error;

#[derive(Clone, Copy, Debug)]
/// User commands for defining the system.
enum Command {
    DefineSystem,
    RemoveSystem,
    SwapSystems,
    QuitAndSave,
    QuitWithoutSaving,
}

/// Edit the list of system definitions to construct from.
pub fn user_menu(database: &DataBase, mut system_defs: &mut Vec<AvailableComponents>)
        -> Result<()> {
    let commands = command_parser!(
        ("d", Command::DefineSystem, "Define a system to create"),
        ("r", Command::RemoveSystem, "Remove a system from the list"),
        ("s", Command::SwapSystems, "Swap the order of two systems"),
        ("f", Command::QuitAndSave, "Finalize editing and return"),
        ("a", Command::QuitWithoutSaving, "Abort and discard changes to list")
    );

    let backup = system_defs.clone();

    loop {
        describe_system_definitions(&system_defs);
        commands.print_menu();
        let input = utils::get_input_string("Selection")?;
        println!("");

        if let Some((cmd, tail)) = commands.get_selection_and_tail(&input) {
            match cmd {
                Command::DefineSystem => {
                    match create_definition(&database) {
                        Ok(def) => system_defs.push(def),
                        Err(err) => println!("Could not create definition: {}", err.description()),
                    }
                },
                Command::RemoveSystem => {
                    match utils::remove_item(&mut system_defs, &tail) {
                        Ok(i) => println!("Removed system at index {}.", i),
                        Err(err) => println!("Could not remove system: {}", err.description()),
                    }
                },
                Command::SwapSystems => {
                    match utils::swap_items(&mut system_defs, &tail) {
                        Ok((i, j)) => println!("Swapped system at index {} with system at {}.",
                                               i, j),
                        Err(err) => println!("Could not swap systems: {}", err.description()),
                    }
                },
                Command::QuitAndSave => {
                    break Ok(())
                },
                Command::QuitWithoutSaving => {
                    system_defs.clear();
                    system_defs.extend_from_slice(&backup);

                    break Ok(())
                },
            }
        } else {
            println!("Not a valid selection.");
        }

        println!("");
    }
}

/// Print the current system definitions to stdout.
pub fn describe_system_definitions(system_defs: &[AvailableComponents]) {
    if system_defs.is_empty() {
        println!("(No systems have been defined)");
    } else {
        println!("System definitions:");
        for (i, def) in system_defs.iter().enumerate() {
            println!("{}. {}", i, def.describe_long());
        }
    }
    println!("");
}

/// Prompt the user to fill in the missing information for a definition.
fn create_definition(database: &DataBase) -> Result<AvailableComponents> {
    use database::AvailableComponents::*;

    match select_component(&database) {
        Ok(Sheet(mut def)) => {
            let position = select_position()?;
            let size = select_size()?;

            def.position = Some(position);
            def.size = Some(size);

            Ok(Sheet(def))
        },
        Ok(Cylinder(mut def)) => {
            let position = select_position()?;
            let radius = utils::get_and_parse_string_single("Set radius (nm)")?;
            let height = utils::get_and_parse_string_single("Set height (nm)")?;

            def.position = Some(position);
            def.radius = Some(radius);
            def.height = Some(height);

            Ok(Cylinder(def))
        },
        err @ Err(_) => err,
    }
}

/// Prompt the user for a component from the list in the `DataBase`.
fn select_component(database: &DataBase) -> Result<AvailableComponents> {
    println!("Available components:");
    for (i, sub) in database.component_defs.iter().enumerate() {
        println!("{}. {}", i, &sub.describe());
    }
    println!("");

    let selection = utils::get_input_string("Select component")?;
    let index = utils::parse_string_for_index(&selection, &database.component_defs)?;

    database.component_defs
        .get(index)
        .ok_or(GrafenCliError::UIError(format!("'{}' is not a valid index", &selection)))
        .map(|comp| comp.clone())
}

/// Get a `Coord` either from the user or by default at (0, 0, 0). Ugly?
fn select_position() -> Result<Coord> {
    let selection = utils::get_input_string("Change position (default: (0.0, 0.0, 0.0))")?;
    if selection.is_empty() {
        return Ok(Coord::new(0.0, 0.0, 0.0));
    }

    let coords = utils::parse_string(&selection)?;
    let &x = coords.get(0).ok_or(UIErrorKind::BadValue("3 positions are required".to_string()))?;
    let &y = coords.get(1).ok_or(UIErrorKind::BadValue("3 positions are required".to_string()))?;
    let &z = coords.get(2).ok_or(UIErrorKind::BadValue("3 positions are required".to_string()))?;

    Ok(Coord::new(x, y, z))
}

/// Get a 2D size. Ugly.
fn select_size() -> Result<(f64, f64)> {
    let size = utils::get_and_parse_string("Set size (nm^2)")?;
    let &dx = size.get(0).ok_or(UIErrorKind::BadValue("2 values are required".to_string()))?;
    let &dy = size.get(1).ok_or(UIErrorKind::BadValue("2 values are required".to_string()))?;

    Ok((dx, dy))
}
<|endoftext|>


<commit_before>extern crate lcm;
extern crate example;
use example::exlcm;

fn main() {
    let mut lcm = lcm::Lcm::new().unwrap();

    lcm.subscribe("EXAMPLE",
                  Box::new(|msg: exlcm::Example| println!("Received message: {:?}", msg)));

    loop {
        lcm.handle().unwrap_or_else(|e| {
            println!("Error handling message: {}", e);
        })
    }
}
<commit_msg>Remove box from callback.<commit_after>extern crate lcm;
extern crate example;
use example::exlcm;

fn main() {
    let mut lcm = lcm::Lcm::new().unwrap();

    lcm.subscribe("EXAMPLE",
                  |msg: exlcm::Example| println!("Received message: {:?}", msg));

    loop {
        lcm.handle().unwrap_or_else(|e| {
            println!("Error handling message: {}", e);
        })
    }
}
<|endoftext|>


<commit_before>use std::{
    collections::HashSet,
    fmt::Display,
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use log::{debug, error, info, trace, warn};
use tokio::{net::TcpListener, select, time};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, OpCode, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{sys::create_udp_socket, utils::try_timeout},
};

use self::upstream::{ProxyUpstream, Upstream};

pub(crate) mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || match iter.next() {
        Some(label) => std::str::from_utf8(label).unwrap_or("*"),
        None => "0", // zero fill the missing labels
    };
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0],
                segments[1],
                segments[2],
                segments[3],
                segments[4],
                segments[5],
                segments[6],
                segments[7],
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

fn check_name_in_proxy_list(acl: &AccessControl, name: &Name) -> Option<bool> {
    if name.is_fqdn() {
        // remove the last dot from FQDN
        let mut name = name.to_ascii();
        name.pop();
        acl.check_host_in_proxy_list(&name)
    } else {
        // unconditionally use default for PQDNs
        Some(acl.is_default_in_proxy_list())
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: Option<&AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = check_name_in_proxy_list(acl, query.name());
            if result == None && acl.is_ip_empty() && acl.is_host_empty() {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: Option<&AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        if let Ok(ref local_response) = local_response {
            let mut names = HashSet::new();
            names.insert(query.name());
            macro_rules! examine_name {
                ($name:expr, $is_answer:expr) => {{
                    names.insert($name);
                    if $is_answer {
                        if let Some(value) = check_name_in_proxy_list(acl, $name) {
                            value
                        } else {
                            acl.is_default_in_proxy_list()
                        }
                    } else {
                        acl.is_default_in_proxy_list()
                    }
                }};
            }
            macro_rules! examine_record {
                ($rec:ident, $is_answer:expr) => {
                    if let RData::CNAME(ref name) = $rec.rdata() {
                        if $is_answer {
                            if let Some(value) = check_name_in_proxy_list(acl, name) {
                                return value;
                            }
                        }
                        names.insert(name);
                        continue;
                    }
                    if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                        warn!(
                            "local DNS response has inconsistent answer type {} for query {}",
                            $rec.record_type(),
                            query
                        );
                        return true;
                    }
                    let forward = match $rec.rdata() {
                        RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V4(*ip)),
                        RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V6(*ip)),
                        // MX records cause type A additional section processing for the host specified by EXCHANGE.
                        RData::MX(ref mx) => examine_name!(mx.exchange(), $is_answer),
                        // NS records cause both the usual additional section processing to locate a type A record...
                        RData::NS(ref name) => examine_name!(name, $is_answer),
                        RData::PTR(_) => unreachable!(),
                        _ => acl.is_default_in_proxy_list(),
                    };
                    if !forward {
                        return false;
                    }
                };
            }
            for rec in local_response.answers() {
                if !names.contains(rec.name()) {
                    warn!(
                        "local DNS response contains unexpected name {} for query {}",
                        rec.name(),
                        query
                    );
                    return true;
                }
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                if names.contains(rec.name()) {
                    examine_record!(rec, false);
                }
            }
        }
        true
    } else {
        unreachable!()
    }
}

struct DnsRelay<Remote>
where
    Remote: Upstream,
{
    context: SharedContext,
    remote_upstream: Remote,
}

impl<Remote> DnsRelay<Remote>
where
    Remote: Upstream,
{
    fn new(context: SharedContext, remote_upstream: Remote) -> DnsRelay<Remote> {
        DnsRelay {
            context,
            remote_upstream,
        }
    }

    async fn acl_lookup(&self, query: &Query) -> (io::Result<Message>, bool) {
        let acl = self.context.acl();
        let local = self.context.local_dns();
        let remote = &self.remote_upstream;

        // Start querying name servers
        debug!("DNS lookup {:?} {}", query.query_type(), query.name());

        let remote_response_fut = try_timeout(remote.lookup(&self.context, query), Some(Duration::from_secs(5)));
        let local_response_fut = try_timeout(local.lookup(&self.context, query), Some(Duration::from_secs(5)));

        match should_forward_by_query(acl, query) {
            Some(true) => {
                let remote_response = remote_response_fut.await;
                trace!("pick remote response (query): {:?}", remote_response);
                return (remote_response, true);
            }
            Some(false) => {
                let local_response = local_response_fut.await;
                trace!("pick local response (query): {:?}", local_response);
                return (local_response, false);
            }
            None => (),
        }

        let decider = async {
            let local_response = local_response_fut.await;
            if should_forward_by_response(acl, &local_response, query) {
                None
            } else {
                Some(local_response)
            }
        };
        tokio::pin!(remote_response_fut, decider);
        let mut use_remote = false;
        let mut remote_response = None;
        loop {
            tokio::select! {
                response = &mut remote_response_fut, if remote_response.is_none() => {
                    if use_remote {
                        trace!("pick remote response (response): {:?}", response);
                        return (response, true);
                    } else {
                        remote_response = Some(response);
                    }
                }
                decision = &mut decider, if !use_remote => {
                    if let Some(local_response) = decision {
                        trace!("pick local response (response): {:?}", local_response);
                        return (local_response, false);
                    } else if let Some(remote_response) = remote_response {
                        trace!("pick remote response (response): {:?}", remote_response);
                        return (remote_response, true);
                    } else {
                        use_remote = true;
                    }
                }
                else => unreachable!(),
            }
        }
    }

    async fn resolve(&self, request: Message) -> Message {
        let mut message = Message::new();
        message.set_id(request.id());
        message.set_recursion_desired(true);
        message.set_recursion_available(true);
        message.set_message_type(MessageType::Response);
        if !request.recursion_desired() {
            message.set_recursion_desired(false);
            message.set_response_code(ResponseCode::NotImp);
        } else if request.op_code() != OpCode::Query || request.message_type() != MessageType::Query {
            message.set_response_code(ResponseCode::NotImp);
        } else if request.query_count() > 0 {
            let (r, forward) = self.acl_lookup(&request.queries()[0]).await;
            if let Ok(result) = r {
                for rec in result.answers() {
                    trace!("dns answer: {:?}", rec);
                    match rec.rdata() {
                        RData::A(ref ip) => {
                            self.context
                                .add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward)
                                .await
                        }
                        RData::AAAA(ref ip) => {
                            self.context
                                .add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward)
                                .await
                        }
                        _ => (),
                    }
                }
                message = result;
                message.set_id(request.id());
            } else {
                message.set_response_code(ResponseCode::ServFail);
            }
        }
        message
    }
}

async fn run_tcp<Remote>(relay: Arc<DnsRelay<Remote>>, bind_addr: SocketAddr) -> io::Result<()>
where
    Remote: Upstream + Display + Send + Sync + 'static,
{
    let listener = TcpListener::bind(&bind_addr).await?;

    let actual_local_addr = listener.local_addr()?;
    info!(
        "shadowsocks DNS relay (TCP) listening on {}, local {}, remote {}",
        actual_local_addr,
        relay.context.local_dns(),
        relay.remote_upstream
    );

    loop {
        let (mut stream, src) = match listener.accept().await {
            Ok(s) => s,
            Err(err) => {
                error!("accept failed with error: {}", err);
                time::sleep(Duration::from_secs(1)).await;
                continue;
            }
        };

        let relay = relay.clone();
        tokio::spawn(async move {
            match upstream::read_message(&mut stream).await {
                Ok(request) => {
                    trace!("received src: {}, query: {:?}", src, request);
                    let message = relay.resolve(request).await;
                    trace!("DNS src: {}, final response: {:?}", src, message);
                    if let Err(err) = upstream::write_message(&mut stream, &message).await {
                        error!("failed to write DNS response, error: {}", err);
                    }
                }
                Err(e) => error!("failed to parse TCP query message from {}, error: {:?}", src, e),
            }
        });
    }
}

async fn run_udp<Remote>(relay: Arc<DnsRelay<Remote>>, bind_addr: SocketAddr) -> io::Result<()>
where
    Remote: Upstream + Display + Send + Sync + 'static,
{
    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!(
        "shadowsocks DNS relay (UDP) listening on {}, local {}, remote {}",
        actual_local_addr,
        relay.context.local_dns(),
        relay.remote_upstream
    );

    let rx = Arc::new(socket);
    let tx = rx.clone();

    loop {
        let mut req_buffer = [0u8; 512];
        let (n, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        trace!("DNS recv {} bytes from {}", n, src);

        let relay = relay.clone();
        let tx = tx.clone();

        tokio::spawn(async move {
            let req_buffer = &req_buffer[..n];

            let request = match Message::from_vec(&req_buffer) {
                Ok(x) => x,
                Err(e) => {
                    error!("failed to parse UDP query message, error: {:?}", e);
                    return;
                }
            };

            trace!("DNS query from {}, {:?}", src, request);
            let message = relay.resolve(request).await;
            trace!("DNS response to {}, {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(err) = tx.send_to(&res_buffer, &src).await {
                        error!("DNS send back UDP error: {}", err);
                    }
                }
            }
        });
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let bind_addr = match context.config().dns_bind_addr {
        None => {
            if matches!(context.config().config_type, ConfigType::DnsLocal) {
                match context.config().local_addr {
                    None => panic!("DNS relay requires dns_bind_addr or local_addr"),
                    Some(ref addr) => addr.bind_addr(&context).await?,
                }
            } else {
                panic!("Local dns relay requires dns_bind_addr");
            }
        }
        Some(ref bind_addr) => bind_addr.bind_addr(&context).await?,
    };

    let remote_dns_addr = context
        .config()
        .remote_dns_addr
        .clone()
        .expect("remote query DNS address");

    let proxy_upstream = ProxyUpstream::new(context.clone(), remote_dns_addr).await;
    let proxy_relay = Arc::new(DnsRelay::new(context.clone(), proxy_upstream));

    let udp_dns = run_udp(proxy_relay.clone(), bind_addr);
    let tcp_dns = run_tcp(proxy_relay, bind_addr);

    select! {
        res = udp_dns => res,
        res = tcp_dns => res,
    }
}
<commit_msg>DNS UDP listener sleep 1s if error occurs<commit_after>use std::{
    collections::HashSet,
    fmt::Display,
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use log::{debug, error, info, trace, warn};
use tokio::{net::TcpListener, select, time};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, OpCode, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{sys::create_udp_socket, utils::try_timeout},
};

use self::upstream::{ProxyUpstream, Upstream};

pub(crate) mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || match iter.next() {
        Some(label) => std::str::from_utf8(label).unwrap_or("*"),
        None => "0", // zero fill the missing labels
    };
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0],
                segments[1],
                segments[2],
                segments[3],
                segments[4],
                segments[5],
                segments[6],
                segments[7],
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

fn check_name_in_proxy_list(acl: &AccessControl, name: &Name) -> Option<bool> {
    if name.is_fqdn() {
        // remove the last dot from FQDN
        let mut name = name.to_ascii();
        name.pop();
        acl.check_host_in_proxy_list(&name)
    } else {
        // unconditionally use default for PQDNs
        Some(acl.is_default_in_proxy_list())
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: Option<&AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = check_name_in_proxy_list(acl, query.name());
            if result == None && acl.is_ip_empty() && acl.is_host_empty() {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: Option<&AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        if let Ok(ref local_response) = local_response {
            let mut names = HashSet::new();
            names.insert(query.name());
            macro_rules! examine_name {
                ($name:expr, $is_answer:expr) => {{
                    names.insert($name);
                    if $is_answer {
                        if let Some(value) = check_name_in_proxy_list(acl, $name) {
                            value
                        } else {
                            acl.is_default_in_proxy_list()
                        }
                    } else {
                        acl.is_default_in_proxy_list()
                    }
                }};
            }
            macro_rules! examine_record {
                ($rec:ident, $is_answer:expr) => {
                    if let RData::CNAME(ref name) = $rec.rdata() {
                        if $is_answer {
                            if let Some(value) = check_name_in_proxy_list(acl, name) {
                                return value;
                            }
                        }
                        names.insert(name);
                        continue;
                    }
                    if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                        warn!(
                            "local DNS response has inconsistent answer type {} for query {}",
                            $rec.record_type(),
                            query
                        );
                        return true;
                    }
                    let forward = match $rec.rdata() {
                        RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V4(*ip)),
                        RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V6(*ip)),
                        // MX records cause type A additional section processing for the host specified by EXCHANGE.
                        RData::MX(ref mx) => examine_name!(mx.exchange(), $is_answer),
                        // NS records cause both the usual additional section processing to locate a type A record...
                        RData::NS(ref name) => examine_name!(name, $is_answer),
                        RData::PTR(_) => unreachable!(),
                        _ => acl.is_default_in_proxy_list(),
                    };
                    if !forward {
                        return false;
                    }
                };
            }
            for rec in local_response.answers() {
                if !names.contains(rec.name()) {
                    warn!(
                        "local DNS response contains unexpected name {} for query {}",
                        rec.name(),
                        query
                    );
                    return true;
                }
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                if names.contains(rec.name()) {
                    examine_record!(rec, false);
                }
            }
        }
        true
    } else {
        unreachable!()
    }
}

struct DnsRelay<Remote>
where
    Remote: Upstream,
{
    context: SharedContext,
    remote_upstream: Remote,
}

impl<Remote> DnsRelay<Remote>
where
    Remote: Upstream,
{
    fn new(context: SharedContext, remote_upstream: Remote) -> DnsRelay<Remote> {
        DnsRelay {
            context,
            remote_upstream,
        }
    }

    async fn acl_lookup(&self, query: &Query) -> (io::Result<Message>, bool) {
        let acl = self.context.acl();
        let local = self.context.local_dns();
        let remote = &self.remote_upstream;

        // Start querying name servers
        debug!("DNS lookup {:?} {}", query.query_type(), query.name());

        let remote_response_fut = try_timeout(remote.lookup(&self.context, query), Some(Duration::from_secs(5)));
        let local_response_fut = try_timeout(local.lookup(&self.context, query), Some(Duration::from_secs(5)));

        match should_forward_by_query(acl, query) {
            Some(true) => {
                let remote_response = remote_response_fut.await;
                trace!("pick remote response (query): {:?}", remote_response);
                return (remote_response, true);
            }
            Some(false) => {
                let local_response = local_response_fut.await;
                trace!("pick local response (query): {:?}", local_response);
                return (local_response, false);
            }
            None => (),
        }

        let decider = async {
            let local_response = local_response_fut.await;
            if should_forward_by_response(acl, &local_response, query) {
                None
            } else {
                Some(local_response)
            }
        };
        tokio::pin!(remote_response_fut, decider);
        let mut use_remote = false;
        let mut remote_response = None;
        loop {
            tokio::select! {
                response = &mut remote_response_fut, if remote_response.is_none() => {
                    if use_remote {
                        trace!("pick remote response (response): {:?}", response);
                        return (response, true);
                    } else {
                        remote_response = Some(response);
                    }
                }
                decision = &mut decider, if !use_remote => {
                    if let Some(local_response) = decision {
                        trace!("pick local response (response): {:?}", local_response);
                        return (local_response, false);
                    } else if let Some(remote_response) = remote_response {
                        trace!("pick remote response (response): {:?}", remote_response);
                        return (remote_response, true);
                    } else {
                        use_remote = true;
                    }
                }
                else => unreachable!(),
            }
        }
    }

    async fn resolve(&self, request: Message) -> Message {
        let mut message = Message::new();
        message.set_id(request.id());
        message.set_recursion_desired(true);
        message.set_recursion_available(true);
        message.set_message_type(MessageType::Response);
        if !request.recursion_desired() {
            message.set_recursion_desired(false);
            message.set_response_code(ResponseCode::NotImp);
        } else if request.op_code() != OpCode::Query || request.message_type() != MessageType::Query {
            message.set_response_code(ResponseCode::NotImp);
        } else if request.query_count() > 0 {
            let (r, forward) = self.acl_lookup(&request.queries()[0]).await;
            if let Ok(result) = r {
                for rec in result.answers() {
                    trace!("dns answer: {:?}", rec);
                    match rec.rdata() {
                        RData::A(ref ip) => {
                            self.context
                                .add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward)
                                .await
                        }
                        RData::AAAA(ref ip) => {
                            self.context
                                .add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward)
                                .await
                        }
                        _ => (),
                    }
                }
                message = result;
                message.set_id(request.id());
            } else {
                message.set_response_code(ResponseCode::ServFail);
            }
        }
        message
    }
}

async fn run_tcp<Remote>(relay: Arc<DnsRelay<Remote>>, bind_addr: SocketAddr) -> io::Result<()>
where
    Remote: Upstream + Display + Send + Sync + 'static,
{
    let listener = TcpListener::bind(&bind_addr).await?;

    let actual_local_addr = listener.local_addr()?;
    info!(
        "shadowsocks DNS relay (TCP) listening on {}, local {}, remote {}",
        actual_local_addr,
        relay.context.local_dns(),
        relay.remote_upstream
    );

    loop {
        let (mut stream, src) = match listener.accept().await {
            Ok(s) => s,
            Err(err) => {
                error!("accept failed with error: {}", err);
                time::sleep(Duration::from_secs(1)).await;
                continue;
            }
        };

        let relay = relay.clone();
        tokio::spawn(async move {
            match upstream::read_message(&mut stream).await {
                Ok(request) => {
                    trace!("received src: {}, query: {:?}", src, request);
                    let message = relay.resolve(request).await;
                    trace!("DNS src: {}, final response: {:?}", src, message);
                    if let Err(err) = upstream::write_message(&mut stream, &message).await {
                        error!("failed to write DNS response, error: {}", err);
                    }
                }
                Err(e) => error!("failed to parse TCP query message from {}, error: {:?}", src, e),
            }
        });
    }
}

async fn run_udp<Remote>(relay: Arc<DnsRelay<Remote>>, bind_addr: SocketAddr) -> io::Result<()>
where
    Remote: Upstream + Display + Send + Sync + 'static,
{
    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!(
        "shadowsocks DNS relay (UDP) listening on {}, local {}, remote {}",
        actual_local_addr,
        relay.context.local_dns(),
        relay.remote_upstream
    );

    let rx = Arc::new(socket);
    let tx = rx.clone();

    loop {
        let mut req_buffer = [0u8; 512];
        let (n, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                time::sleep(Duration::from_secs(1)).await;
                continue;
            }
        };

        trace!("DNS recv {} bytes from {}", n, src);

        let relay = relay.clone();
        let tx = tx.clone();

        tokio::spawn(async move {
            let req_buffer = &req_buffer[..n];

            let request = match Message::from_vec(&req_buffer) {
                Ok(x) => x,
                Err(e) => {
                    error!("failed to parse UDP query message, error: {:?}", e);
                    return;
                }
            };

            trace!("DNS query from {}, {:?}", src, request);
            let message = relay.resolve(request).await;
            trace!("DNS response to {}, {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(err) = tx.send_to(&res_buffer, &src).await {
                        error!("DNS send back UDP error: {}", err);
                    }
                }
            }
        });
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let bind_addr = match context.config().dns_bind_addr {
        None => {
            if matches!(context.config().config_type, ConfigType::DnsLocal) {
                match context.config().local_addr {
                    None => panic!("DNS relay requires dns_bind_addr or local_addr"),
                    Some(ref addr) => addr.bind_addr(&context).await?,
                }
            } else {
                panic!("Local dns relay requires dns_bind_addr");
            }
        }
        Some(ref bind_addr) => bind_addr.bind_addr(&context).await?,
    };

    let remote_dns_addr = context
        .config()
        .remote_dns_addr
        .clone()
        .expect("remote query DNS address");

    let proxy_upstream = ProxyUpstream::new(context.clone(), remote_dns_addr).await;
    let proxy_relay = Arc::new(DnsRelay::new(context.clone(), proxy_upstream));

    let udp_dns = run_udp(proxy_relay.clone(), bind_addr);
    let tcp_dns = run_tcp(proxy_relay, bind_addr);

    select! {
        res = udp_dns => res,
        res = tcp_dns => res,
    }
}
<|endoftext|>


<commit_before>// The Computer Language Benchmarks Game
// http://benchmarksgame.alioth.debian.org/
//
// contributed by the Rust Project Developers
// contributed by TeXitoi

#![feature(libc, io, core, std_misc)]

extern crate libc;

use std::old_io::stdio::{stdin_raw, stdout_raw};
use std::old_io::{IoResult, EndOfFile};
use std::ptr::copy_memory;
use std::thread::Thread;

struct Tables {
    table8: [u8;1 << 8],
    table16: [u16;1 << 16]
}

impl Tables {
    fn new() -> Tables {
        let mut table8 = [0;1 << 8];
        for (i, v) in table8.iter_mut().enumerate() {
            *v = Tables::computed_cpl8(i as u8);
        }
        let mut table16 = [0;1 << 16];
        for (i, v) in table16.iter_mut().enumerate() {
            *v = (table8[i & 255] as u16) << 8 |
                 table8[i >> 8]  as u16;
        }
        Tables { table8: table8, table16: table16 }
    }

    fn computed_cpl8(c: u8) -> u8 {
        match c {
            b'A' | b'a' => b'T',
            b'C' | b'c' => b'G',
            b'G' | b'g' => b'C',
            b'T' | b't' => b'A',
            b'U' | b'u' => b'A',
            b'M' | b'm' => b'K',
            b'R' | b'r' => b'Y',
            b'W' | b'w' => b'W',
            b'S' | b's' => b'S',
            b'Y' | b'y' => b'R',
            b'K' | b'k' => b'M',
            b'V' | b'v' => b'B',
            b'H' | b'h' => b'D',
            b'D' | b'd' => b'H',
            b'B' | b'b' => b'V',
            b'N' | b'n' => b'N',
            i => i,
        }
    }

    /// Retreives the complement for `i`.
    fn cpl8(&self, i: u8) -> u8 {
        self.table8[i as usize]
    }

    /// Retreives the complement for `i`.
    fn cpl16(&self, i: u16) -> u16 {
        self.table16[i as usize]
    }
}

/// Reads all remaining bytes from the stream.
fn read_to_end<R: Reader>(r: &mut R) -> IoResult<Vec<u8>> {
    // As reading the input stream in memory is a bottleneck, we tune
    // Reader::read_to_end() with a fast growing policy to limit
    // recopies.  If MREMAP_RETAIN is implemented in the linux kernel
    // and jemalloc use it, this trick will become useless.
    const CHUNK: usize = 64 * 1024;

    let mut vec = Vec::with_capacity(CHUNK);
    loop {
        // workaround: very fast growing
        let len = vec.len();
        if vec.capacity() - len < CHUNK {
            let cap = vec.capacity();
            let mult = if cap < 256 * 1024 * 1024 {
                16
            } else {
                2
            };
            vec.reserve_exact(mult * cap - len);
        }
        match r.push_at_least(1, CHUNK, &mut vec) {
            Ok(_) => {}
            Err(ref e) if e.kind == EndOfFile => break,
            Err(e) => return Err(e)
        }
    }
    Ok(vec)
}

/// Finds the first position at which `b` occurs in `s`.
fn memchr(h: &[u8], n: u8) -> Option<usize> {
    use libc::{c_void, c_int, size_t};
    let res = unsafe {
        libc::memchr(h.as_ptr() as *const c_void, n as c_int, h.len() as size_t)
    };
    if res.is_null() {
        None
    } else {
        Some(res as usize - h.as_ptr() as usize)
    }
}

/// A mutable iterator over DNA sequences
struct MutDnaSeqs<'a> { s: &'a mut [u8] }
fn mut_dna_seqs<'a>(s: &'a mut [u8]) -> MutDnaSeqs<'a> {
    MutDnaSeqs { s: s }
}
impl<'a> Iterator for MutDnaSeqs<'a> {
    type Item = &'a mut [u8];

    fn next(&mut self) -> Option<&'a mut [u8]> {
        let tmp = std::mem::replace(&mut self.s, &mut []);
        let tmp = match memchr(tmp, b'\n') {
            Some(i) => &mut tmp[i + 1 ..],
            None => return None,
        };
        let (seq, tmp) = match memchr(tmp, b'>') {
            Some(i) => tmp.split_at_mut(i),
            None => {
                let len = tmp.len();
                tmp.split_at_mut(len)
            }
        };
        self.s = tmp;
        Some(seq)
    }
}

/// Length of a normal line without the terminating \n.
const LINE_LEN: usize = 60;

/// Compute the reverse complement.
fn reverse_complement(seq: &mut [u8], tables: &Tables) {
    let len = seq.len() - 1;
    let seq = &mut seq[..len];// Drop the last newline
    let off = LINE_LEN - len % (LINE_LEN + 1);
    let mut i = LINE_LEN;
    while i < len {
        unsafe {
            copy_memory(seq.as_mut_ptr().offset((i - off + 1) as isize),
                        seq.as_ptr().offset((i - off) as isize), off);
            *seq.get_unchecked_mut(i - off) = b'\n';
        }
        i += LINE_LEN + 1;
    }

    let div = len / 4;
    let rem = len % 4;
    unsafe {
        let mut left = seq.as_mut_ptr() as *mut u16;
        // This is slow if len % 2 != 0 but still faster than bytewise operations.
        let mut right = seq.as_mut_ptr().offset(len as isize - 2) as *mut u16;
        let end = left.offset(div as isize);
        while left != end {
            let tmp = tables.cpl16(*left);
            *left = tables.cpl16(*right);
            *right = tmp;
            left = left.offset(1);
            right = right.offset(-1);
        }

        let end = end as *mut u8;
        match rem {
            1 => *end = tables.cpl8(*end),
            2 => {
                let tmp = tables.cpl8(*end);
                *end = tables.cpl8(*end.offset(1));
                *end.offset(1) = tmp;
            },
            3 => {
                *end.offset(1) = tables.cpl8(*end.offset(1));
                let tmp = tables.cpl8(*end);
                *end = tables.cpl8(*end.offset(2));
                *end.offset(2) = tmp;
            },
            _ => { },
        }
    }
}

struct Racy<T>(T);

unsafe impl<T: 'static> Send for Racy<T> {}

/// Executes a closure in parallel over the given iterator over mutable slice.
/// The closure `f` is run in parallel with an element of `iter`.
fn parallel<'a, I, T, F>(iter: I, f: F)
        where T: 'a+Send + Sync,
              I: Iterator<Item=&'a mut [T]>,
              F: Fn(&mut [T]) + Sync {
    use std::mem;
    use std::raw::Repr;

    iter.map(|chunk| {
        // Need to convert `f` and `chunk` to something that can cross the task
        // boundary.
        let f = Racy(&f as *const F as *const usize);
        let raw = Racy(chunk.repr());
        Thread::scoped(move|| {
            let f = f.0 as *const F;
            unsafe { (*f)(mem::transmute(raw.0)) }
        })
    }).collect::<Vec<_>>();
}

fn main() {
    let mut data = read_to_end(&mut stdin_raw()).unwrap();
    let tables = &Tables::new();
    parallel(mut_dna_seqs(&mut*data), |&: seq| reverse_complement(seq, tables));
    stdout_raw().write_all(&mut*data).unwrap();
}
<commit_msg>update reverse-complement from rust repo<commit_after>// The Computer Language Benchmarks Game
// http://benchmarksgame.alioth.debian.org/
//
// contributed by the Rust Project Developers
// contributed by TeXitoi

#![feature(libc, core, old_io)]

extern crate libc;

use std::old_io::stdio::{stdin_raw, stdout_raw};
use std::old_io::{IoResult, EndOfFile};
use std::ptr::copy_memory;
use std::thread;

struct Tables {
    table8: [u8;1 << 8],
    table16: [u16;1 << 16]
}

impl Tables {
    fn new() -> Tables {
        let mut table8 = [0;1 << 8];
        for (i, v) in table8.iter_mut().enumerate() {
            *v = Tables::computed_cpl8(i as u8);
        }
        let mut table16 = [0;1 << 16];
        for (i, v) in table16.iter_mut().enumerate() {
            *v = (table8[i & 255] as u16) << 8 |
                 table8[i >> 8]  as u16;
        }
        Tables { table8: table8, table16: table16 }
    }

    fn computed_cpl8(c: u8) -> u8 {
        match c {
            b'A' | b'a' => b'T',
            b'C' | b'c' => b'G',
            b'G' | b'g' => b'C',
            b'T' | b't' => b'A',
            b'U' | b'u' => b'A',
            b'M' | b'm' => b'K',
            b'R' | b'r' => b'Y',
            b'W' | b'w' => b'W',
            b'S' | b's' => b'S',
            b'Y' | b'y' => b'R',
            b'K' | b'k' => b'M',
            b'V' | b'v' => b'B',
            b'H' | b'h' => b'D',
            b'D' | b'd' => b'H',
            b'B' | b'b' => b'V',
            b'N' | b'n' => b'N',
            i => i,
        }
    }

    /// Retrieves the complement for `i`.
    fn cpl8(&self, i: u8) -> u8 {
        self.table8[i as usize]
    }

    /// Retrieves the complement for `i`.
    fn cpl16(&self, i: u16) -> u16 {
        self.table16[i as usize]
    }
}

/// Reads all remaining bytes from the stream.
fn read_to_end<R: Reader>(r: &mut R) -> IoResult<Vec<u8>> {
    // As reading the input stream in memory is a bottleneck, we tune
    // Reader::read_to_end() with a fast growing policy to limit
    // recopies.  If MREMAP_RETAIN is implemented in the linux kernel
    // and jemalloc use it, this trick will become useless.
    const CHUNK: usize = 64 * 1024;

    let mut vec = Vec::with_capacity(CHUNK);
    loop {
        // workaround: very fast growing
        let len = vec.len();
        if vec.capacity() - len < CHUNK {
            let cap = vec.capacity();
            let mult = if cap < 256 * 1024 * 1024 {
                16
            } else {
                2
            };
            vec.reserve_exact(mult * cap - len);
        }
        match r.push_at_least(1, CHUNK, &mut vec) {
            Ok(_) => {}
            Err(ref e) if e.kind == EndOfFile => break,
            Err(e) => return Err(e)
        }
    }
    Ok(vec)
}

/// Finds the first position at which `b` occurs in `s`.
fn memchr(h: &[u8], n: u8) -> Option<usize> {
    use libc::{c_void, c_int, size_t};
    let res = unsafe {
        libc::memchr(h.as_ptr() as *const c_void, n as c_int, h.len() as size_t)
    };
    if res.is_null() {
        None
    } else {
        Some(res as usize - h.as_ptr() as usize)
    }
}

/// A mutable iterator over DNA sequences
struct MutDnaSeqs<'a> { s: &'a mut [u8] }
fn mut_dna_seqs<'a>(s: &'a mut [u8]) -> MutDnaSeqs<'a> {
    MutDnaSeqs { s: s }
}
impl<'a> Iterator for MutDnaSeqs<'a> {
    type Item = &'a mut [u8];

    fn next(&mut self) -> Option<&'a mut [u8]> {
        let tmp = std::mem::replace(&mut self.s, &mut []);
        let tmp = match memchr(tmp, b'\n') {
            Some(i) => &mut tmp[i + 1 ..],
            None => return None,
        };
        let (seq, tmp) = match memchr(tmp, b'>') {
            Some(i) => tmp.split_at_mut(i),
            None => {
                let len = tmp.len();
                tmp.split_at_mut(len)
            }
        };
        self.s = tmp;
        Some(seq)
    }
}

/// Length of a normal line without the terminating \n.
const LINE_LEN: usize = 60;

/// Compute the reverse complement.
fn reverse_complement(seq: &mut [u8], tables: &Tables) {
    let len = seq.len() - 1;
    let seq = &mut seq[..len];// Drop the last newline
    let off = LINE_LEN - len % (LINE_LEN + 1);
    let mut i = LINE_LEN;
    while i < len {
        unsafe {
            copy_memory(seq.as_mut_ptr().offset((i - off + 1) as isize),
                        seq.as_ptr().offset((i - off) as isize), off);
            *seq.get_unchecked_mut(i - off) = b'\n';
        }
        i += LINE_LEN + 1;
    }

    let div = len / 4;
    let rem = len % 4;
    unsafe {
        let mut left = seq.as_mut_ptr() as *mut u16;
        // This is slow if len % 2 != 0 but still faster than bytewise operations.
        let mut right = seq.as_mut_ptr().offset(len as isize - 2) as *mut u16;
        let end = left.offset(div as isize);
        while left != end {
            let tmp = tables.cpl16(*left);
            *left = tables.cpl16(*right);
            *right = tmp;
            left = left.offset(1);
            right = right.offset(-1);
        }

        let end = end as *mut u8;
        match rem {
            1 => *end = tables.cpl8(*end),
            2 => {
                let tmp = tables.cpl8(*end);
                *end = tables.cpl8(*end.offset(1));
                *end.offset(1) = tmp;
            },
            3 => {
                *end.offset(1) = tables.cpl8(*end.offset(1));
                let tmp = tables.cpl8(*end);
                *end = tables.cpl8(*end.offset(2));
                *end.offset(2) = tmp;
            },
            _ => { },
        }
    }
}

/// Executes a closure in parallel over the given iterator over mutable slice.
/// The closure `f` is run in parallel with an element of `iter`.
fn parallel<'a, I: Iterator, F>(iter: I, ref f: F)
        where I::Item: Send + 'a,
              F: Fn(I::Item) + Sync + 'a {
    iter.map(|x| {
        thread::scoped(move|| {
            f(x)
        })
    }).collect::<Vec<_>>();
}

fn main() {
    let mut data = read_to_end(&mut stdin_raw()).unwrap();
    let tables = &Tables::new();
    parallel(mut_dna_seqs(&mut data), |seq| reverse_complement(seq, tables));
    stdout_raw().write_all(&data).unwrap();
}
<|endoftext|>