<commit_before>use std::time;
use std::ops;

use chrono;
use chrono::{Timelike};

#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
pub struct FloatDuration {
    secs: f64,
}

impl FloatDuration {
    pub fn new(secs: f64) -> FloatDuration {
        FloatDuration {secs}
    }

    pub fn days(days: f64) -> FloatDuration {
        FloatDuration {secs: days * (3600.0 * 24.0)}
    }
    pub fn hours(hours: f64) -> FloatDuration {
        FloatDuration {secs: hours * 3600.0}
    }
    pub fn minutes(mins: f64) -> FloatDuration {
        FloatDuration {secs: mins * 60.0}
    }
    pub fn seconds(secs: f64) -> FloatDuration {
        FloatDuration {secs}
    }
    pub fn milliseconds(millis: f64) -> FloatDuration {
        FloatDuration {secs: millis / 1000.0}
    }
    pub fn microseconds(micros: f64) -> FloatDuration {
        FloatDuration {secs: micros / 1.0e6}
    }
    pub fn nanoseconds(nanos: f64) -> FloatDuration {
        FloatDuration {secs: nanos / 1.0e9}
    }

    pub fn as_days(&self) -> f64 {
        self.secs / (3600.0*24.0)
    }
    pub fn as_hours(&self) -> f64 {
        self.secs / 3600.0
    }
    pub fn as_minutes(&self) -> f64 {
        self.secs / 60.0
    }
    pub fn as_seconds(&self) -> f64 {
        self.secs
    }
    pub fn as_milliseconds(&self) -> f64 {
        self.secs * 1.0e3
    }
    pub fn as_microseconds(&self) -> f64 {
        self.secs * 1.0e6
    }
    pub fn as_nanoseconds(&self) -> f64 {
        self.secs * 1.0e9
    }

    pub fn as_std(&self) -> Option<time::Duration> {
        if self.secs.is_sign_negative() {
            None
        } else {
            let seconds = self.secs.trunc();
            let nanos = self.secs.fract() * 1e9;

            Some(time::Duration::new(seconds as u64, nanos as u32))
        }
    }

    pub fn from_std(duration: time::Duration) -> FloatDuration {
        return FloatDuration::seconds(
            (duration.as_secs() as f64) * (duration.subsec_nanos() as f64)/1.0e-9);
    }
}

impl ops::Neg for FloatDuration {
    type Output = FloatDuration;

    fn neg(self) -> FloatDuration {
        FloatDuration {secs: -self.secs}
    }
}

impl ops::Add<FloatDuration> for FloatDuration {
    type Output = FloatDuration;

    fn add(self, rhs: FloatDuration) -> FloatDuration {
        FloatDuration {secs: self.secs + rhs.secs}
    }
}
impl ops::Sub<FloatDuration> for FloatDuration {
    type Output = FloatDuration;

    fn sub(self, rhs: FloatDuration) -> FloatDuration {
        FloatDuration {secs: self.secs - rhs.secs}
    }
}

impl ops::Mul<f64> for FloatDuration {
    type Output = FloatDuration;

    fn mul(self, rhs: f64) -> FloatDuration {
        FloatDuration {secs: self.secs * rhs}
    }
}
impl ops::Div<f64> for FloatDuration {
    type Output = FloatDuration;

    fn div(self, rhs: f64) -> FloatDuration {
        FloatDuration {secs: self.secs / rhs}
    }
}

impl ops::AddAssign<FloatDuration> for FloatDuration {
    fn add_assign(&mut self, rhs: FloatDuration) {
        self.secs += rhs.secs;
    }
}
impl ops::SubAssign<FloatDuration> for FloatDuration {
    fn sub_assign(&mut self, rhs: FloatDuration) {
        self.secs -= rhs.secs;
    }
}

impl ops::MulAssign<f64> for FloatDuration {
    fn mul_assign(&mut self, rhs: f64) {
        self.secs *= rhs;
    }
}
impl ops::DivAssign<f64> for FloatDuration {
    fn div_assign(&mut self, rhs: f64) {
        self.secs /= rhs;
    }
}
<commit_msg>Further fleshed out duration classes. Added tests.<commit_after>use std::time;
use std::fmt;
use std::ops;

pub const SECS_PER_DAY: f64 = 60.0*60.0*24.0;
pub const SECS_PER_HOUR: f64 = 60.0*60.0;
pub const SECS_PER_MINUTE: f64 = 60.0;

#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
pub struct FloatDuration {
    secs: f64,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct DecomposedTime {
    pub days: u32,
    pub hours: u32,
    pub minutes: u32,
    pub seconds: u32,
    pub fractional_seconds: f64,
    pub sign: i8,
}

impl FloatDuration {
    pub fn new(secs: f64) -> FloatDuration {
        FloatDuration {secs}
    }

    pub fn days(days: f64) -> FloatDuration {
        FloatDuration {secs: days * (3600.0 * 24.0)}
    }
    pub fn hours(hours: f64) -> FloatDuration {
        FloatDuration {secs: hours * 3600.0}
    }
    pub fn minutes(mins: f64) -> FloatDuration {
        FloatDuration {secs: mins * 60.0}
    }
    pub fn seconds(secs: f64) -> FloatDuration {
        FloatDuration {secs}
    }
    pub fn milliseconds(millis: f64) -> FloatDuration {
        FloatDuration {secs: millis / 1000.0}
    }
    pub fn microseconds(micros: f64) -> FloatDuration {
        FloatDuration {secs: micros / 1.0e6}
    }
    pub fn nanoseconds(nanos: f64) -> FloatDuration {
        FloatDuration {secs: nanos / 1.0e9}
    }

    pub fn decompose(&self) -> DecomposedTime {
        let mut secs_left = self.secs;

        let days = (secs_left / SECS_PER_DAY).trunc();
        secs_left -= days*SECS_PER_DAY;
        let hours = (secs_left / SECS_PER_HOUR).trunc();
        secs_left -= hours*SECS_PER_HOUR;
        let minutes = (secs_left / SECS_PER_MINUTE).trunc();
        secs_left -= hours*SECS_PER_MINUTE;
        let seconds = secs_left.trunc();
        let fract = secs_left.fract();

        DecomposedTime {
            days: days as u32,
            hours: hours as u32,
            minutes: minutes as u32,
            seconds: seconds as u32,
            fractional_seconds: fract,
            sign: self.secs.signum() as i8,
        }
    }

    pub fn as_days(&self) -> f64 {
        self.secs / SECS_PER_DAY
    }
    pub fn as_hours(&self) -> f64 {
        self.secs / SECS_PER_HOUR
    }
    pub fn as_minutes(&self) -> f64 {
        self.secs / SECS_PER_MINUTE
    }
    pub fn as_seconds(&self) -> f64 {
        self.secs
    }
    pub fn as_milliseconds(&self) -> f64 {
        self.secs * 1.0e3
    }
    pub fn as_microseconds(&self) -> f64 {
        self.secs * 1.0e6
    }
    pub fn as_nanoseconds(&self) -> f64 {
        self.secs * 1.0e9
    }

    pub fn zero() -> FloatDuration {
        FloatDuration {secs: 0.0}
    }
    pub fn is_zero(&self) -> bool {
        self.secs == 0.0
    }
    pub fn is_positive(&self) -> bool {
        self.secs.is_sign_positive()
    }
    pub fn is_negative(&self) -> bool {
        self.secs.is_sign_negative()
    }

    pub fn as_std(&self) -> Option<time::Duration> {
        if self.secs.is_sign_negative() {
            None
        } else {
            let seconds = self.secs.trunc();
            let nanos = self.secs.fract() * 1e9;

            Some(time::Duration::new(seconds as u64, nanos as u32))
        }
    }

    pub fn from_std(duration: time::Duration) -> FloatDuration {
        return FloatDuration::seconds(
            (duration.as_secs() as f64) + (duration.subsec_nanos() as f64)*1.0e-9);
    }

    pub fn from_decomposed(decomposed: &DecomposedTime) -> FloatDuration {
        let seconds = (decomposed.days as f64)*SECS_PER_DAY
            + (decomposed.hours as f64)*SECS_PER_HOUR
            + (decomposed.minutes as f64)*SECS_PER_MINUTE
            + (decomposed.seconds as f64)
            + decomposed.fractional_seconds;

        FloatDuration {
            secs: seconds
        }
    }
}

impl DecomposedTime {
    pub fn from_components(days: u32, hours: u32, minutes: u32,
               seconds: u32, fractional_seconds: f64) -> DecomposedTime
    {
        DecomposedTime {
            days, hours, minutes, seconds, fractional_seconds, sign: 1}
    }

    pub fn negate(mut self) -> DecomposedTime {
        self.sign *= -1;
        self
    }
}

impl fmt::Display for FloatDuration {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let decompose_time = self.decompose();

        if decompose_time.days > 0 {
            write!(fmt, "{}d ", decompose_time.days as u64)?;
        }

        if decompose_time.sign.is_negative() {
            write!(fmt, "-")?;
        }
        write!(fmt, "{:02}:{:02}:{:02}.{}",
            decompose_time.hours,
            decompose_time.minutes,
            decompose_time.seconds,
            decompose_time.fractional_seconds)
    }
}

impl ops::Neg for FloatDuration {
    type Output = FloatDuration;

    fn neg(self) -> FloatDuration {
        FloatDuration {secs: -self.secs}
    }
}

impl ops::Add<FloatDuration> for FloatDuration {
    type Output = FloatDuration;

    fn add(self, rhs: FloatDuration) -> FloatDuration {
        FloatDuration {secs: self.secs + rhs.secs}
    }
}
impl ops::Sub<FloatDuration> for FloatDuration {
    type Output = FloatDuration;

    fn sub(self, rhs: FloatDuration) -> FloatDuration {
        FloatDuration {secs: self.secs - rhs.secs}
    }
}

impl ops::Mul<f64> for FloatDuration {
    type Output = FloatDuration;

    fn mul(self, rhs: f64) -> FloatDuration {
        FloatDuration {secs: self.secs * rhs}
    }
}
impl ops::Div<f64> for FloatDuration {
    type Output = FloatDuration;

    fn div(self, rhs: f64) -> FloatDuration {
        FloatDuration {secs: self.secs / rhs}
    }
}

impl ops::AddAssign<FloatDuration> for FloatDuration {
    fn add_assign(&mut self, rhs: FloatDuration) {
        self.secs += rhs.secs;
    }
}
impl ops::SubAssign<FloatDuration> for FloatDuration {
    fn sub_assign(&mut self, rhs: FloatDuration) {
        self.secs -= rhs.secs;
    }
}

impl ops::MulAssign<f64> for FloatDuration {
    fn mul_assign(&mut self, rhs: f64) {
        self.secs *= rhs;
    }
}
impl ops::DivAssign<f64> for FloatDuration {
    fn div_assign(&mut self, rhs: f64) {
        self.secs /= rhs;
    }
}

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

    #[test]
    fn test_construct() {
        let duration1 = FloatDuration::hours(3.0);
        assert_eq!(duration1.as_hours(), 3.0);
        assert_eq!(duration1.as_minutes(), 180.0);
        assert_eq!(duration1.as_seconds(), 180.0*60.0);
        assert_eq!(duration1.as_days(), 3.0/24.0);
        assert_eq!(duration1.as_milliseconds(), 180.0*60.0*1000.0);

        let duration2 = FloatDuration::milliseconds(55.0);
        assert_eq!(duration2.as_seconds(), 0.055);
        assert_eq!(duration2.as_milliseconds(), 55.0);
        assert_eq!(duration2.as_microseconds(), 55000.0);
        assert_eq!(duration2.as_nanoseconds(), 55000000.0);
        assert!(!duration2.is_zero());

        let duration3 = FloatDuration::zero();
        assert!(duration3.is_zero());

        assert_eq!(FloatDuration::days(1.5), FloatDuration::hours(36.0));
        assert_eq!(FloatDuration::minutes(30.0), FloatDuration::hours(0.5));
        assert_eq!(FloatDuration::seconds(180.0), FloatDuration::minutes(3.0));
        assert_eq!(FloatDuration::seconds(3.5), FloatDuration::milliseconds(3500.0));
        assert_eq!(FloatDuration::microseconds(300.0), FloatDuration::milliseconds(0.30));
        assert_eq!(FloatDuration::nanoseconds(1000.0), FloatDuration::microseconds(1.0));
    }

    #[test]
    fn test_arithmetic() {
        let duration1 = FloatDuration::minutes(5.0) + FloatDuration::seconds(30.0);
        assert_eq!(duration1, FloatDuration::seconds(330.0));

        let duration2 = FloatDuration::hours(3.0) * 2.5;
        assert_eq!(duration2, FloatDuration::hours(7.5));

        let duration3 = FloatDuration::days(3.0) / 3.0 - FloatDuration::hours(2.0);
        assert_eq!(duration3, FloatDuration::hours(22.0));

        let duration4 = FloatDuration::zero() + FloatDuration::milliseconds(500.0)
            + FloatDuration::microseconds(500.0);
        assert_eq!(duration4, FloatDuration::microseconds(500500.0));

    }

    #[test]
    fn test_std_conversion() {
        let duration1 = FloatDuration::minutes(5.0);
        let std_duration1 = duration1.as_std().unwrap();
        assert!(duration1.is_positive());
        assert_eq!(std_duration1, time::Duration::new(300, 0));
        assert_eq!(FloatDuration::from_std(std_duration1), duration1);

        let duration2 = FloatDuration::hours(-2.0);
        assert!(duration2.is_negative());
        assert_eq!(duration2.as_std(), None);
        let std_duration2 = (-duration2).as_std().unwrap();
        assert_eq!(std_duration2, time::Duration::new(3600*2, 0));
        assert_eq!(FloatDuration::from_std(std_duration2), -duration2);
    }
}
<|endoftext|>


<commit_before>//! Stream encoder
use std::future::Future;
use std::io::{self, Write};
use std::pin::Pin;
use std::task::{Context, Poll};

use actix_threadpool::{run, CpuFuture};
use brotli2::write::BrotliEncoder;
use bytes::Bytes;
use flate2::write::{GzEncoder, ZlibEncoder};
use futures_core::ready;
use pin_project::{pin_project, project};

use crate::body::{Body, BodySize, MessageBody, ResponseBody};
use crate::http::header::{ContentEncoding, CONTENT_ENCODING};
use crate::http::{HeaderValue, StatusCode};
use crate::{Error, ResponseHead};

use super::Writer;

const INPLACE: usize = 1024;

#[pin_project]
pub struct Encoder<B> {
    eof: bool,
    #[pin]
    body: EncoderBody<B>,
    encoder: Option<ContentEncoder>,
    fut: Option<CpuFuture<ContentEncoder, io::Error>>,
}

impl<B: MessageBody> Encoder<B> {
    pub fn response(
        encoding: ContentEncoding,
        head: &mut ResponseHead,
        body: ResponseBody<B>,
    ) -> ResponseBody<Encoder<B>> {
        let can_encode = !(head.headers().contains_key(&CONTENT_ENCODING)
            || head.status == StatusCode::SWITCHING_PROTOCOLS
            || head.status == StatusCode::NO_CONTENT
            || encoding == ContentEncoding::Identity
            || encoding == ContentEncoding::Auto);

        let body = match body {
            ResponseBody::Other(b) => match b {
                Body::None => return ResponseBody::Other(Body::None),
                Body::Empty => return ResponseBody::Other(Body::Empty),
                Body::Bytes(buf) => {
                    if can_encode {
                        EncoderBody::Bytes(buf)
                    } else {
                        return ResponseBody::Other(Body::Bytes(buf));
                    }
                }
                Body::Message(stream) => EncoderBody::BoxedStream(stream),
            },
            ResponseBody::Body(stream) => EncoderBody::Stream(stream),
        };

        if can_encode {
            // Modify response body only if encoder is not None
            if let Some(enc) = ContentEncoder::encoder(encoding) {
                update_head(encoding, head);
                head.no_chunking(false);
                return ResponseBody::Body(Encoder {
                    body,
                    eof: false,
                    fut: None,
                    encoder: Some(enc),
                });
            }
        }
        ResponseBody::Body(Encoder {
            body,
            eof: false,
            fut: None,
            encoder: None,
        })
    }
}

#[pin_project]
enum EncoderBody<B> {
    Bytes(Bytes),
    Stream(#[pin] B),
    BoxedStream(#[pin] Box<dyn MessageBody>),
}

impl<B: MessageBody> MessageBody for EncoderBody<B> {
    fn size(&self) -> BodySize {
        match self {
            EncoderBody::Bytes(ref b) => b.size(),
            EncoderBody::Stream(ref b) => b.size(),
            EncoderBody::BoxedStream(ref b) => b.size(),
        }
    }

    #[project]
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes, Error>>> {
        #[project]
        match self.project() {
            EncoderBody::Bytes(b) => {
                if b.is_empty() {
                    Poll::Ready(None)
                } else {
                    Poll::Ready(Some(Ok(std::mem::replace(b, Bytes::new()))))
                }
            }
            EncoderBody::Stream(b) => b.poll_next(cx),
            EncoderBody::BoxedStream(b) => b.poll_next(cx),
        }
    }
}

impl<B: MessageBody> MessageBody for Encoder<B> {
    fn size(&self) -> BodySize {
        if self.encoder.is_none() {
            self.body.size()
        } else {
            BodySize::Stream
        }
    }

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes, Error>>> {
        let mut this = self.project();
        loop {
            if *this.eof {
                return Poll::Ready(None);
            }

            if let Some(ref mut fut) = this.fut {
                let mut encoder = match ready!(Pin::new(fut).poll(cx)) {
                    Ok(item) => item,
                    Err(e) => return Poll::Ready(Some(Err(e.into()))),
                };
                let chunk = encoder.take();
                *this.encoder = Some(encoder);
                this.fut.take();
                if !chunk.is_empty() {
                    return Poll::Ready(Some(Ok(chunk)));
                }
            }

            let result = this.body.as_mut().poll_next(cx);

            match result {
                Poll::Ready(Some(Ok(chunk))) => {
                    if let Some(mut encoder) = this.encoder.take() {
                        if chunk.len() < INPLACE {
                            encoder.write(&chunk)?;
                            let chunk = encoder.take();
                            *this.encoder = Some(encoder);
                            if !chunk.is_empty() {
                                return Poll::Ready(Some(Ok(chunk)));
                            }
                        } else {
                            *this.fut = Some(run(move || {
                                encoder.write(&chunk)?;
                                Ok(encoder)
                            }));
                        }
                    } else {
                        return Poll::Ready(Some(Ok(chunk)));
                    }
                }
                Poll::Ready(None) => {
                    if let Some(encoder) = this.encoder.take() {
                        let chunk = encoder.finish()?;
                        if chunk.is_empty() {
                            return Poll::Ready(None);
                        } else {
                            *this.eof = true;
                            return Poll::Ready(Some(Ok(chunk)));
                        }
                    } else {
                        return Poll::Ready(None);
                    }
                }
                val => return val,
            }
        }
    }
}

fn update_head(encoding: ContentEncoding, head: &mut ResponseHead) {
    head.headers_mut().insert(
        CONTENT_ENCODING,
        HeaderValue::from_static(encoding.as_str()),
    );
}

enum ContentEncoder {
    Deflate(ZlibEncoder<Writer>),
    Gzip(GzEncoder<Writer>),
    Br(BrotliEncoder<Writer>),
}

impl ContentEncoder {
    fn encoder(encoding: ContentEncoding) -> Option<Self> {
        match encoding {
            ContentEncoding::Deflate => Some(ContentEncoder::Deflate(ZlibEncoder::new(
                Writer::new(),
                flate2::Compression::fast(),
            ))),
            ContentEncoding::Gzip => Some(ContentEncoder::Gzip(GzEncoder::new(
                Writer::new(),
                flate2::Compression::fast(),
            ))),
            ContentEncoding::Br => {
                Some(ContentEncoder::Br(BrotliEncoder::new(Writer::new(), 3)))
            }
            _ => None,
        }
    }

    #[inline]
    pub(crate) fn take(&mut self) -> Bytes {
        match *self {
            ContentEncoder::Br(ref mut encoder) => encoder.get_mut().take(),
            ContentEncoder::Deflate(ref mut encoder) => encoder.get_mut().take(),
            ContentEncoder::Gzip(ref mut encoder) => encoder.get_mut().take(),
        }
    }

    fn finish(self) -> Result<Bytes, io::Error> {
        match self {
            ContentEncoder::Br(encoder) => match encoder.finish() {
                Ok(writer) => Ok(writer.buf.freeze()),
                Err(err) => Err(err),
            },
            ContentEncoder::Gzip(encoder) => match encoder.finish() {
                Ok(writer) => Ok(writer.buf.freeze()),
                Err(err) => Err(err),
            },
            ContentEncoder::Deflate(encoder) => match encoder.finish() {
                Ok(writer) => Ok(writer.buf.freeze()),
                Err(err) => Err(err),
            },
        }
    }

    fn write(&mut self, data: &[u8]) -> Result<(), io::Error> {
        match *self {
            ContentEncoder::Br(ref mut encoder) => match encoder.write_all(data) {
                Ok(_) => Ok(()),
                Err(err) => {
                    trace!("Error decoding br encoding: {}", err);
                    Err(err)
                }
            },
            ContentEncoder::Gzip(ref mut encoder) => match encoder.write_all(data) {
                Ok(_) => Ok(()),
                Err(err) => {
                    trace!("Error decoding gzip encoding: {}", err);
                    Err(err)
                }
            },
            ContentEncoder::Deflate(ref mut encoder) => match encoder.write_all(data) {
                Ok(_) => Ok(()),
                Err(err) => {
                    trace!("Error decoding deflate encoding: {}", err);
                    Err(err)
                }
            },
        }
    }
}
<commit_msg>Fix EncoderBody enum to align with Body::Message<commit_after>//! Stream encoder
use std::future::Future;
use std::io::{self, Write};
use std::pin::Pin;
use std::task::{Context, Poll};

use actix_threadpool::{run, CpuFuture};
use brotli2::write::BrotliEncoder;
use bytes::Bytes;
use flate2::write::{GzEncoder, ZlibEncoder};
use futures_core::ready;
use pin_project::{pin_project, project};

use crate::body::{Body, BodySize, MessageBody, ResponseBody};
use crate::http::header::{ContentEncoding, CONTENT_ENCODING};
use crate::http::{HeaderValue, StatusCode};
use crate::{Error, ResponseHead};

use super::Writer;

const INPLACE: usize = 1024;

#[pin_project]
pub struct Encoder<B> {
    eof: bool,
    #[pin]
    body: EncoderBody<B>,
    encoder: Option<ContentEncoder>,
    fut: Option<CpuFuture<ContentEncoder, io::Error>>,
}

impl<B: MessageBody> Encoder<B> {
    pub fn response(
        encoding: ContentEncoding,
        head: &mut ResponseHead,
        body: ResponseBody<B>,
    ) -> ResponseBody<Encoder<B>> {
        let can_encode = !(head.headers().contains_key(&CONTENT_ENCODING)
            || head.status == StatusCode::SWITCHING_PROTOCOLS
            || head.status == StatusCode::NO_CONTENT
            || encoding == ContentEncoding::Identity
            || encoding == ContentEncoding::Auto);

        let body = match body {
            ResponseBody::Other(b) => match b {
                Body::None => return ResponseBody::Other(Body::None),
                Body::Empty => return ResponseBody::Other(Body::Empty),
                Body::Bytes(buf) => {
                    if can_encode {
                        EncoderBody::Bytes(buf)
                    } else {
                        return ResponseBody::Other(Body::Bytes(buf));
                    }
                }
                Body::Message(stream) => EncoderBody::BoxedStream(stream),
            },
            ResponseBody::Body(stream) => EncoderBody::Stream(stream),
        };

        if can_encode {
            // Modify response body only if encoder is not None
            if let Some(enc) = ContentEncoder::encoder(encoding) {
                update_head(encoding, head);
                head.no_chunking(false);
                return ResponseBody::Body(Encoder {
                    body,
                    eof: false,
                    fut: None,
                    encoder: Some(enc),
                });
            }
        }
        ResponseBody::Body(Encoder {
            body,
            eof: false,
            fut: None,
            encoder: None,
        })
    }
}

#[pin_project]
enum EncoderBody<B> {
    Bytes(Bytes),
    Stream(#[pin] B),
    BoxedStream(Box<dyn MessageBody + Unpin>),
}

impl<B: MessageBody> MessageBody for EncoderBody<B> {
    fn size(&self) -> BodySize {
        match self {
            EncoderBody::Bytes(ref b) => b.size(),
            EncoderBody::Stream(ref b) => b.size(),
            EncoderBody::BoxedStream(ref b) => b.size(),
        }
    }

    #[project]
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes, Error>>> {
        #[project]
        match self.project() {
            EncoderBody::Bytes(b) => {
                if b.is_empty() {
                    Poll::Ready(None)
                } else {
                    Poll::Ready(Some(Ok(std::mem::replace(b, Bytes::new()))))
                }
            }
            EncoderBody::Stream(b) => b.poll_next(cx),
            EncoderBody::BoxedStream(ref mut b) => Pin::new(b.as_mut()).poll_next(cx),
        }
    }
}

impl<B: MessageBody> MessageBody for Encoder<B> {
    fn size(&self) -> BodySize {
        if self.encoder.is_none() {
            self.body.size()
        } else {
            BodySize::Stream
        }
    }

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes, Error>>> {
        let mut this = self.project();
        loop {
            if *this.eof {
                return Poll::Ready(None);
            }

            if let Some(ref mut fut) = this.fut {
                let mut encoder = match ready!(Pin::new(fut).poll(cx)) {
                    Ok(item) => item,
                    Err(e) => return Poll::Ready(Some(Err(e.into()))),
                };
                let chunk = encoder.take();
                *this.encoder = Some(encoder);
                this.fut.take();
                if !chunk.is_empty() {
                    return Poll::Ready(Some(Ok(chunk)));
                }
            }

            let result = this.body.as_mut().poll_next(cx);

            match result {
                Poll::Ready(Some(Ok(chunk))) => {
                    if let Some(mut encoder) = this.encoder.take() {
                        if chunk.len() < INPLACE {
                            encoder.write(&chunk)?;
                            let chunk = encoder.take();
                            *this.encoder = Some(encoder);
                            if !chunk.is_empty() {
                                return Poll::Ready(Some(Ok(chunk)));
                            }
                        } else {
                            *this.fut = Some(run(move || {
                                encoder.write(&chunk)?;
                                Ok(encoder)
                            }));
                        }
                    } else {
                        return Poll::Ready(Some(Ok(chunk)));
                    }
                }
                Poll::Ready(None) => {
                    if let Some(encoder) = this.encoder.take() {
                        let chunk = encoder.finish()?;
                        if chunk.is_empty() {
                            return Poll::Ready(None);
                        } else {
                            *this.eof = true;
                            return Poll::Ready(Some(Ok(chunk)));
                        }
                    } else {
                        return Poll::Ready(None);
                    }
                }
                val => return val,
            }
        }
    }
}

fn update_head(encoding: ContentEncoding, head: &mut ResponseHead) {
    head.headers_mut().insert(
        CONTENT_ENCODING,
        HeaderValue::from_static(encoding.as_str()),
    );
}

enum ContentEncoder {
    Deflate(ZlibEncoder<Writer>),
    Gzip(GzEncoder<Writer>),
    Br(BrotliEncoder<Writer>),
}

impl ContentEncoder {
    fn encoder(encoding: ContentEncoding) -> Option<Self> {
        match encoding {
            ContentEncoding::Deflate => Some(ContentEncoder::Deflate(ZlibEncoder::new(
                Writer::new(),
                flate2::Compression::fast(),
            ))),
            ContentEncoding::Gzip => Some(ContentEncoder::Gzip(GzEncoder::new(
                Writer::new(),
                flate2::Compression::fast(),
            ))),
            ContentEncoding::Br => {
                Some(ContentEncoder::Br(BrotliEncoder::new(Writer::new(), 3)))
            }
            _ => None,
        }
    }

    #[inline]
    pub(crate) fn take(&mut self) -> Bytes {
        match *self {
            ContentEncoder::Br(ref mut encoder) => encoder.get_mut().take(),
            ContentEncoder::Deflate(ref mut encoder) => encoder.get_mut().take(),
            ContentEncoder::Gzip(ref mut encoder) => encoder.get_mut().take(),
        }
    }

    fn finish(self) -> Result<Bytes, io::Error> {
        match self {
            ContentEncoder::Br(encoder) => match encoder.finish() {
                Ok(writer) => Ok(writer.buf.freeze()),
                Err(err) => Err(err),
            },
            ContentEncoder::Gzip(encoder) => match encoder.finish() {
                Ok(writer) => Ok(writer.buf.freeze()),
                Err(err) => Err(err),
            },
            ContentEncoder::Deflate(encoder) => match encoder.finish() {
                Ok(writer) => Ok(writer.buf.freeze()),
                Err(err) => Err(err),
            },
        }
    }

    fn write(&mut self, data: &[u8]) -> Result<(), io::Error> {
        match *self {
            ContentEncoder::Br(ref mut encoder) => match encoder.write_all(data) {
                Ok(_) => Ok(()),
                Err(err) => {
                    trace!("Error decoding br encoding: {}", err);
                    Err(err)
                }
            },
            ContentEncoder::Gzip(ref mut encoder) => match encoder.write_all(data) {
                Ok(_) => Ok(()),
                Err(err) => {
                    trace!("Error decoding gzip encoding: {}", err);
                    Err(err)
                }
            },
            ContentEncoder::Deflate(ref mut encoder) => match encoder.write_all(data) {
                Ok(_) => Ok(()),
                Err(err) => {
                    trace!("Error decoding deflate encoding: {}", err);
                    Err(err)
                }
            },
        }
    }
}
<|endoftext|>


<commit_before>//! Construct planar sheets.

use surface::distribution::PoissonDistribution;
use surface::lattice::Lattice;
use surface::LatticeType;
use surface::LatticeType::*;

use coord::{Coord, Direction, Periodic, Translate,
    rotate_planar_coords_to_alignment};
use describe::{unwrap_name, Describe};
use error::{GrafenError, Result};
use iterator::{AtomIterator, AtomIterItem};
use system::*;
use volume::pbc_multiply_volume;

impl_component![Sheet];
impl_translate![Circle, Sheet];

#[derive(Clone, Debug, Deserialize, Serialize)]
/// A rectangular sheet.
pub struct Sheet {
    /// Name of component.
    pub name: Option<String>,
    /// Optional residue placed at each coordinate. If not set the sheet describes
    ///  a general collection of coordinates.
    pub residue: Option<Residue>,
    /// Lattice type used to construct the surface structure.
    pub lattice: LatticeType,
    /// Standard deviation along z of coordinates. Added to the coordinates when `construct`
    /// is called.
    pub std_z: Option<f64>,
    #[serde(skip)]
    /// Origin of the sheet. Located in the lower-left position of it.
    pub origin: Coord,
    /// Normal vector of the sheet.
    pub normal: Direction,
    #[serde(skip)]
    /// Length of the sheet along the x axis.
    pub length: f64,
    #[serde(skip)]
    /// Length of the sheet along the y axis.
    pub width: f64,
    #[serde(skip)]
    /// List of coordinates belonging to the sheet. Relative to the `origin.
    pub coords: Vec<Coord>,
}

impl Sheet {
    /// Construct the sheet coordinates and return the object.
    ///
    /// # Errors
    /// Returns an error if either the length or width is non-positive.
    pub fn construct(self) -> Result<Sheet> {
        if self.length <= 0.0 || self.width <= 0.0 {
            return Err(
                GrafenError::RunError("cannot create a substrate of negative size".to_string())
            );
        }

        let mut coords_lattice = match self.lattice {
            Hexagonal { a } => {
                Lattice::hexagonal(a)
                    .with_size(self.length, self.width)
                    .finalize()
            },
            Triclinic { a, b, gamma } => {
                Lattice::triclinic(a, b, gamma.to_radians())
                    .with_size(self.length, self.width)
                    .finalize()
            },
            PoissonDisc { density } => {
                // The factor 1/sqrt(pi) comes from the area and the factor sqrt(2)
                // is a magic number which roughly gives the correct density. It works!
                use std::f64::consts::PI;
                let rmin = (2.0 / (PI * density)).sqrt();

                PoissonDistribution::new(rmin, self.length, self.width)
            },
        };

        if let Some(std) = self.std_z {
            coords_lattice = coords_lattice.uniform_distribution(std);
        };

        let (length, width, _) = coords_lattice.box_size.to_tuple();

        let coords = match self.normal {
            Direction::Z => coords_lattice.coords,
            direction => {
                rotate_planar_coords_to_alignment(&coords_lattice.coords, Direction::Z, direction)
            },
        };

        Ok(Sheet {
            length,
            width,
            coords,
            .. self
        })
    }

    /// Calculate the box size. The height of a pure sheet is set to 0.1 (nm)
    /// as a lower limit of the system.
    fn calc_box_size(&self) -> Coord {
        Coord::new(self.length, self.width, 0.1)
    }

    /// Cut a circle out of coordinates in the sheet.
    pub fn to_circle(&self, radius: f64) -> Circle {
        let center = Coord::new(radius, radius, 0.0);
        let box_size = Coord::new(self.length, self.width, 0.0);

        // Assert that we have a large enough sheet to cut a circle from
        let pbc_multiples = (
            (2.0 * radius / self.length).ceil() as usize,
            (2.0 * radius / self.width).ceil() as usize
        );

        let coords = match pbc_multiples {
            (1, 1) => {
                cut_circle(&self.coords, center, box_size, radius)
            },
            (nx, ny) => {
                let box_size = box_size.pbc_multiply(nx, ny, 1);
                let coords = self.pbc_multiply(nx, ny, 1).coords;

                cut_circle(&coords, center, box_size, radius)
            },
        };

        Circle {
            residue: self.residue.clone(),
            origin: self.origin,
            radius,
            coords,
        }
    }
}

impl Describe for Sheet {
    fn describe(&self) -> String {
        format!("{} (Rectangular sheet of size ({:.2}, {:.2}) at {})",
            unwrap_name(&self.name), self.length, self.width, self.origin)
    }

    fn describe_short(&self) -> String {
        format!("{} (Sheet)", unwrap_name(&self.name))
    }
}

impl Periodic for Sheet {
    /// Clone sheet coordinates into PBC multiples.
    fn pbc_multiply(&self, nx: usize, ny: usize, _: usize) -> Sheet {
        let current_size = Coord::new(self.length, self.width, 0.0);
        let coords = pbc_multiply_volume(&self.coords, current_size, nx, ny, 1);

        Sheet {
            length: nx as f64 * self.length,
            width: ny as f64 * self.width,
            coords,
            .. self.clone()
        }

    }
}

#[derive(Clone, Debug)]
/// A 2D circular sheet. For now used only to cut caps for the `Cylinder` construction.
pub struct Circle {
    residue: Option<Residue>,
    origin: Coord,
    radius: f64,
    pub coords: Vec<Coord>,
}

/// Cut a set of coordinates into a circle with input radius in the x-y plane.
fn cut_circle(coords: &[Coord], center: Coord, box_size: Coord, radius: f64) -> Vec<Coord> {
    coords.iter()
        .map(|&coord| coord.with_pbc(box_size) - center)
        .filter(|coord| {
            let (dr, _) = coord.distance_cylindrical(Coord::ORIGO, Direction::Z);
            dr <= radius
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::f64::consts::PI;

    fn setup_sheet(length: f64, width: f64, lattice: &LatticeType) -> Sheet {
        Sheet {
            name: None,
            residue: None,
            lattice: lattice.clone(),
            std_z: None,
            origin: Coord::default(),
            normal: Direction::Z,
            length,
            width,
            coords: vec![],
        }
    }

    #[test]
    fn create_sheet_with_density() {
        let length = 10.0;
        let width = 5.0;
        let density = 10.0;

        let sheet = setup_sheet(length, width, &PoissonDisc { density }).construct().unwrap();

        // There should be almost no risk of the density being off by more than 10%
        // since 500 points were created.
        let output_density = sheet.coords.len() as f64 / (length * width);
        assert!((output_density - density).abs() / density < 0.1);
    }

    #[test]
    fn create_regular_sheets_updates_sheet_size_to_match_pbc() {
        let a = 1.1;
        let b = 1.4;
        let lattice = Triclinic { a, b, gamma: 90.0 };

        // The length and width do not perfectly match the lattice constants
        let length = 2.0;
        let width = 1.0;

        let sheet = setup_sheet(length, width, &lattice).construct().unwrap();

        // The final size has been rounded to the closest multiple of the spacing
        assert_eq!(2.0 * a, sheet.length);
        assert_eq!(1.0 * b, sheet.width);
    }

    #[test]
    fn create_sheets_with_negative_input_size_returns_error() {
        let lattice = PoissonDisc { density: 10.0 };

        assert!(setup_sheet(-1.0, 1.0, &lattice).construct().is_err());
        assert!(setup_sheet(1.0, -1.0, &lattice).construct().is_err());
        assert!(setup_sheet(1.0, 1.0, &lattice).construct().is_ok());
    }

    #[test]
    fn variance_is_added_if_requested() {
        let length = 10.0;
        let width = 5.0;
        let lattice = PoissonDisc { density: 10.0 };
        let sheet = Sheet {
            std_z: Some(1.0),
            .. setup_sheet(length, width, &lattice)
        }.construct().unwrap();

        let num_coords = sheet.coords.len();
        let var = sheet.coords
            .iter()
            .map(|&Coord { x: _, y: _, z }| z.abs())
            .sum::<f64>() / (num_coords as f64);

        assert!(var > 0.0);
    }

    #[test]
    fn calc_box_size_of_sheet() {
        let length = 5.0;
        let width = 4.0;
        let lattice = Hexagonal { a: 0.1 };

        let sheet = setup_sheet(length, width, &lattice).construct().unwrap();

        assert_eq!(Coord::new(sheet.length, sheet.width, 0.1), sheet.calc_box_size());
    }

    #[test]
    fn cut_a_sheet_into_a_circle() {
        let radius = 4.0;

        let density = 10.0;
        let lattice = PoissonDisc { density };

        let circle = setup_sheet(10.0, 10.0, &lattice)
            .construct()
            .unwrap()
            .to_circle(radius);

        // Our density should match within 10%.
        let expected = PI * radius.powi(2) * density;
        assert!((circle.coords.len() as f64 - expected).abs() / expected < 0.1);

        // The coordinates are centered around origo.
        for coord in circle.coords {
            let (x, y, _) = coord.to_tuple();

            assert!(x.powi(2) + y.powi(2) <= radius.powi(2));
        }
    }

    #[test]
    fn small_sheets_extend_periodically_if_the_circle_is_too_large() {
        let radius = 4.0;

        let density = 10.0;
        let lattice = PoissonDisc { density };

        // This sheet will not contain the entire circle
        let circle = setup_sheet(6.0, 6.0, &lattice)
            .construct()
            .unwrap()
            .to_circle(radius);

        // Our density should match within 10%.
        let expected = PI * radius.powi(2) * density;
        assert!((circle.coords.len() as f64 - expected).abs() / expected < 0.1);
    }

    #[test]
    fn sheet_can_be_constructed_along_x_and_y() {
        let length = 10.0;
        let density = 10.0;
        let min_expected_coords = (length * length * density) as usize / 2;

        let sheet_x = Sheet {
            normal: Direction::X,
            .. setup_sheet(length, length, &PoissonDisc { density })
        }.construct().unwrap();

        assert!(sheet_x.coords.len() > min_expected_coords);
        for coord in sheet_x.coords {
            assert_eq!(coord.x, 0.0);
        }

        let sheet_y = Sheet {
            normal: Direction::Y,
            .. setup_sheet(length, length, &PoissonDisc { density })
        }.construct().unwrap();

        assert!(sheet_y.coords.len() > min_expected_coords);
        for coord in sheet_y.coords {
            assert_eq!(coord.y, 0.0);
        }

        let sheet_z = Sheet {
            normal: Direction::Z,
            .. setup_sheet(length, length, &PoissonDisc { density })
        }.construct().unwrap();

        assert!(sheet_z.coords.len() > min_expected_coords);
        for coord in sheet_z.coords {
            assert_eq!(coord.z, 0.0);
        }
    }
}
<commit_msg>Calculating the box size of `Sheet`s accounts for the normal<commit_after>//! Construct planar sheets.

use surface::distribution::PoissonDistribution;
use surface::lattice::Lattice;
use surface::LatticeType;
use surface::LatticeType::*;

use coord::{Coord, Direction, Periodic, Translate,
    rotate_planar_coords_to_alignment};
use describe::{unwrap_name, Describe};
use error::{GrafenError, Result};
use iterator::{AtomIterator, AtomIterItem};
use system::*;
use volume::pbc_multiply_volume;

impl_component![Sheet];
impl_translate![Circle, Sheet];

#[derive(Clone, Debug, Deserialize, Serialize)]
/// A rectangular sheet.
pub struct Sheet {
    /// Name of component.
    pub name: Option<String>,
    /// Optional residue placed at each coordinate. If not set the sheet describes
    ///  a general collection of coordinates.
    pub residue: Option<Residue>,
    /// Lattice type used to construct the surface structure.
    pub lattice: LatticeType,
    /// Standard deviation along z of coordinates. Added to the coordinates when `construct`
    /// is called.
    pub std_z: Option<f64>,
    #[serde(skip)]
    /// Origin of the sheet. Located in the lower-left position of it.
    pub origin: Coord,
    /// Normal vector of the sheet.
    pub normal: Direction,
    #[serde(skip)]
    /// Length of the sheet along the x axis.
    pub length: f64,
    #[serde(skip)]
    /// Length of the sheet along the y axis.
    pub width: f64,
    #[serde(skip)]
    /// List of coordinates belonging to the sheet. Relative to the `origin.
    pub coords: Vec<Coord>,
}

impl Sheet {
    /// Construct the sheet coordinates and return the object.
    ///
    /// # Errors
    /// Returns an error if either the length or width is non-positive.
    pub fn construct(self) -> Result<Sheet> {
        if self.length <= 0.0 || self.width <= 0.0 {
            return Err(
                GrafenError::RunError("cannot create a substrate of negative size".to_string())
            );
        }

        let mut coords_lattice = match self.lattice {
            Hexagonal { a } => {
                Lattice::hexagonal(a)
                    .with_size(self.length, self.width)
                    .finalize()
            },
            Triclinic { a, b, gamma } => {
                Lattice::triclinic(a, b, gamma.to_radians())
                    .with_size(self.length, self.width)
                    .finalize()
            },
            PoissonDisc { density } => {
                // The factor 1/sqrt(pi) comes from the area and the factor sqrt(2)
                // is a magic number which roughly gives the correct density. It works!
                use std::f64::consts::PI;
                let rmin = (2.0 / (PI * density)).sqrt();

                PoissonDistribution::new(rmin, self.length, self.width)
            },
        };

        if let Some(std) = self.std_z {
            coords_lattice = coords_lattice.uniform_distribution(std);
        };

        let (length, width, _) = coords_lattice.box_size.to_tuple();

        let coords = match self.normal {
            Direction::Z => coords_lattice.coords,
            direction => {
                rotate_planar_coords_to_alignment(&coords_lattice.coords, Direction::Z, direction)
            },
        };

        Ok(Sheet {
            length,
            width,
            coords,
            .. self
        })
    }

    /// Calculate the box size. The height of a pure sheet is set to 0.1 (nm)
    /// as a lower limit of the system.
    fn calc_box_size(&self) -> Coord {
        let margin = 0.1;

        match self.normal {
            Direction::X => Coord::new(margin, self.width, self.length),
            Direction::Y => Coord::new(self.length, margin, self.width),
            Direction::Z => Coord::new(self.length, self.width, margin),
        }
    }

    /// Cut a circle out of coordinates in the sheet.
    pub fn to_circle(&self, radius: f64) -> Circle {
        let center = Coord::new(radius, radius, 0.0);
        let box_size = Coord::new(self.length, self.width, 0.0);

        // Assert that we have a large enough sheet to cut a circle from
        let pbc_multiples = (
            (2.0 * radius / self.length).ceil() as usize,
            (2.0 * radius / self.width).ceil() as usize
        );

        let coords = match pbc_multiples {
            (1, 1) => {
                cut_circle(&self.coords, center, box_size, radius)
            },
            (nx, ny) => {
                let box_size = box_size.pbc_multiply(nx, ny, 1);
                let coords = self.pbc_multiply(nx, ny, 1).coords;

                cut_circle(&coords, center, box_size, radius)
            },
        };

        Circle {
            residue: self.residue.clone(),
            origin: self.origin,
            radius,
            coords,
        }
    }
}

impl Describe for Sheet {
    fn describe(&self) -> String {
        format!("{} (Rectangular sheet of size ({:.2}, {:.2}) at {})",
            unwrap_name(&self.name), self.length, self.width, self.origin)
    }

    fn describe_short(&self) -> String {
        format!("{} (Sheet)", unwrap_name(&self.name))
    }
}

impl Periodic for Sheet {
    /// Clone sheet coordinates into PBC multiples.
    fn pbc_multiply(&self, nx: usize, ny: usize, _: usize) -> Sheet {
        let current_size = Coord::new(self.length, self.width, 0.0);
        let coords = pbc_multiply_volume(&self.coords, current_size, nx, ny, 1);

        Sheet {
            length: nx as f64 * self.length,
            width: ny as f64 * self.width,
            coords,
            .. self.clone()
        }

    }
}

#[derive(Clone, Debug)]
/// A 2D circular sheet. For now used only to cut caps for the `Cylinder` construction.
pub struct Circle {
    residue: Option<Residue>,
    origin: Coord,
    radius: f64,
    pub coords: Vec<Coord>,
}

/// Cut a set of coordinates into a circle with input radius in the x-y plane.
fn cut_circle(coords: &[Coord], center: Coord, box_size: Coord, radius: f64) -> Vec<Coord> {
    coords.iter()
        .map(|&coord| coord.with_pbc(box_size) - center)
        .filter(|coord| {
            let (dr, _) = coord.distance_cylindrical(Coord::ORIGO, Direction::Z);
            dr <= radius
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::f64::consts::PI;

    fn setup_sheet(length: f64, width: f64, lattice: &LatticeType) -> Sheet {
        Sheet {
            name: None,
            residue: None,
            lattice: lattice.clone(),
            std_z: None,
            origin: Coord::default(),
            normal: Direction::Z,
            length,
            width,
            coords: vec![],
        }
    }

    #[test]
    fn create_sheet_with_density() {
        let length = 10.0;
        let width = 5.0;
        let density = 10.0;

        let sheet = setup_sheet(length, width, &PoissonDisc { density }).construct().unwrap();

        // There should be almost no risk of the density being off by more than 10%
        // since 500 points were created.
        let output_density = sheet.coords.len() as f64 / (length * width);
        assert!((output_density - density).abs() / density < 0.1);
    }

    #[test]
    fn create_regular_sheets_updates_sheet_size_to_match_pbc() {
        let a = 1.1;
        let b = 1.4;
        let lattice = Triclinic { a, b, gamma: 90.0 };

        // The length and width do not perfectly match the lattice constants
        let length = 2.0;
        let width = 1.0;

        let sheet = setup_sheet(length, width, &lattice).construct().unwrap();

        // The final size has been rounded to the closest multiple of the spacing
        assert_eq!(2.0 * a, sheet.length);
        assert_eq!(1.0 * b, sheet.width);
    }

    #[test]
    fn create_sheets_with_negative_input_size_returns_error() {
        let lattice = PoissonDisc { density: 10.0 };

        assert!(setup_sheet(-1.0, 1.0, &lattice).construct().is_err());
        assert!(setup_sheet(1.0, -1.0, &lattice).construct().is_err());
        assert!(setup_sheet(1.0, 1.0, &lattice).construct().is_ok());
    }

    #[test]
    fn variance_is_added_if_requested() {
        let length = 10.0;
        let width = 5.0;
        let lattice = PoissonDisc { density: 10.0 };
        let sheet = Sheet {
            std_z: Some(1.0),
            .. setup_sheet(length, width, &lattice)
        }.construct().unwrap();

        let num_coords = sheet.coords.len();
        let var = sheet.coords
            .iter()
            .map(|&Coord { x: _, y: _, z }| z.abs())
            .sum::<f64>() / (num_coords as f64);

        assert!(var > 0.0);
    }

    #[test]
    fn calc_box_size_of_sheet_accounts_for_sheet_normal() {
        let length = 5.0;
        let width = 4.0;
        let lattice = Hexagonal { a: 0.1 };

        let margin = 0.1;

        let sheet = Sheet {
            normal: Direction::X,
            .. setup_sheet(length, width, &lattice)
        }.construct().unwrap();

        assert_eq!(Coord::new(margin, sheet.width, sheet.length), sheet.calc_box_size());

        let sheet = Sheet {
            normal: Direction::Y,
            .. setup_sheet(length, width, &lattice)
        }.construct().unwrap();

        assert_eq!(Coord::new(sheet.length, margin, sheet.width), sheet.calc_box_size());

        let sheet = Sheet {
            normal: Direction::Z,
            .. setup_sheet(length, width, &lattice)
        }.construct().unwrap();

        assert_eq!(Coord::new(sheet.length, sheet.width, 0.1), sheet.calc_box_size());
    }

    #[test]
    fn cut_a_sheet_into_a_circle() {
        let radius = 4.0;

        let density = 10.0;
        let lattice = PoissonDisc { density };

        let circle = setup_sheet(10.0, 10.0, &lattice)
            .construct()
            .unwrap()
            .to_circle(radius);

        // Our density should match within 10%.
        let expected = PI * radius.powi(2) * density;
        assert!((circle.coords.len() as f64 - expected).abs() / expected < 0.1);

        // The coordinates are centered around origo.
        for coord in circle.coords {
            let (x, y, _) = coord.to_tuple();

            assert!(x.powi(2) + y.powi(2) <= radius.powi(2));
        }
    }

    #[test]
    fn small_sheets_extend_periodically_if_the_circle_is_too_large() {
        let radius = 4.0;

        let density = 10.0;
        let lattice = PoissonDisc { density };

        // This sheet will not contain the entire circle
        let circle = setup_sheet(6.0, 6.0, &lattice)
            .construct()
            .unwrap()
            .to_circle(radius);

        // Our density should match within 10%.
        let expected = PI * radius.powi(2) * density;
        assert!((circle.coords.len() as f64 - expected).abs() / expected < 0.1);
    }

    #[test]
    fn sheet_can_be_constructed_along_x_and_y() {
        let length = 10.0;
        let density = 10.0;
        let min_expected_coords = (length * length * density) as usize / 2;

        let sheet_x = Sheet {
            normal: Direction::X,
            .. setup_sheet(length, length, &PoissonDisc { density })
        }.construct().unwrap();

        assert!(sheet_x.coords.len() > min_expected_coords);
        for coord in sheet_x.coords {
            assert_eq!(coord.x, 0.0);
        }

        let sheet_y = Sheet {
            normal: Direction::Y,
            .. setup_sheet(length, length, &PoissonDisc { density })
        }.construct().unwrap();

        assert!(sheet_y.coords.len() > min_expected_coords);
        for coord in sheet_y.coords {
            assert_eq!(coord.y, 0.0);
        }

        let sheet_z = Sheet {
            normal: Direction::Z,
            .. setup_sheet(length, length, &PoissonDisc { density })
        }.construct().unwrap();

        assert!(sheet_z.coords.len() > min_expected_coords);
        for coord in sheet_z.coords {
            assert_eq!(coord.z, 0.0);
        }
    }
}
<|endoftext|>


<commit_before>use std::{mem, ptr};
use libc;

use c;
use low;

type FfiType_      = *mut low::ffi_type;
type FfiTypeArray_ = *mut FfiType_;

#[derive(Debug)]
pub struct FfiType(FfiType_);

#[derive(Debug)]
pub struct FfiTypeArray {
    ptr: FfiTypeArray_,
    len: usize,
}

/// Creates a null-terminated array of FfiType_. Takes ownership of
/// the elements.
unsafe fn ffi_type_array_create(elements: Vec<FfiType>) -> FfiTypeArray_ {
    let size = elements.len();
    let array = libc::malloc((size+1) * mem::size_of::<FfiType_>())
                    as FfiTypeArray_;

    for i in 0 .. size {
        *array.offset(i as isize) = elements[i].0;
    }
    *array.offset(size as isize) = ptr::null::<low::ffi_type>() as FfiType_;

    for t in elements {
        mem::forget(t);
    }

    println!("ffi_type_array_create(...) = {:?}", array);

    array
}

/// Creates a struct ffi_type with the given elements. Takes ownership
/// of the elements.
unsafe fn ffi_type_struct_create(elements: Vec<FfiType>) -> FfiType_ {
    println!("ffi_type_array_create({:?})", elements);
    let array    = ffi_type_array_create(elements);
    let ffi_type = libc::malloc(mem::size_of::<low::ffi_type>())
                       as FfiType_;

    (*ffi_type).size      = 0;
    (*ffi_type).alignment = 0;
    (*ffi_type).type_     = c::ffi_type_enum::STRUCT as u16;
    (*ffi_type).elements  = array;

    println!("ffi_type_array_create(...) = {:?}", ffi_type);

    ffi_type
}

unsafe fn ffi_type_array_clone(ffi_types: FfiTypeArray_) -> FfiTypeArray_ {
    let mut current = ffi_types;
    let mut count   = 0;
    while !(*current).is_null() {
        current = current.offset(1);
        count += 1;
    }

    let new = libc::malloc((count+1) * mem::size_of::<FfiType_>())
                    as FfiTypeArray_;

    for i in 0 .. count {
        *new.offset(i as isize) = ffi_type_clone(*ffi_types.offset(i as isize));
    }
    *new.offset(count as isize) = ptr::null::<low::ffi_type>() as FfiType_;

    new
}

unsafe fn ffi_type_clone(ffi_type: FfiType_) -> FfiType_ {
    unimplemented!();
}

/// Destroys an array of FfiType_ and all of its elements.
unsafe fn ffi_type_array_destroy(ffi_types: FfiTypeArray_) {
    println!("ffi_type_array_destroy({:?})", ffi_types);
    let mut current = ffi_types;
    while !(*current).is_null() {
        ffi_type_destroy(*current);
        current = current.offset(1);
    }

    libc::free(ffi_types as *mut libc::c_void);
}

/// Destroys an FfiType_ if it was dynamically allocated.
unsafe fn ffi_type_destroy(ffi_type: FfiType_) {
    println!("ffi_type_destroy({:?})", ffi_type);
    if (*ffi_type).type_ == c::ffi_type_enum::STRUCT as u16 {
        ffi_type_array_destroy((*ffi_type).elements);
        libc::free(ffi_type as *mut libc::c_void);
    }
}

impl Drop for FfiType {
    fn drop(&mut self) {
        unsafe { ffi_type_destroy(self.0) }
    }
}

impl Drop for FfiTypeArray {
    fn drop(&mut self) {
        unsafe { ffi_type_array_destroy(self.ptr) }
    }
}

impl FfiType {
    pub fn void() -> Self {
        FfiType(unsafe { &mut low::ffi_type_void })
    }

    pub fn uint8() -> Self {
        FfiType(unsafe { &mut low::ffi_type_uint8 })
    }

    pub fn sint8() -> Self {
        FfiType(unsafe { &mut low::ffi_type_sint8 })
    }

    pub fn uint16() -> Self {
        FfiType(unsafe { &mut low::ffi_type_uint16 })
    }

    pub fn sint16() -> Self {
        FfiType(unsafe { &mut low::ffi_type_sint16 })
    }

    pub fn uint32() -> Self {
        FfiType(unsafe { &mut low::ffi_type_uint32 })
    }

    pub fn sint32() -> Self {
        FfiType(unsafe { &mut low::ffi_type_sint32 })
    }

    pub fn uint64() -> Self {
        FfiType(unsafe { &mut low::ffi_type_uint64 })
    }

    pub fn sint64() -> Self {
        FfiType(unsafe { &mut low::ffi_type_sint64 })
    }

    pub fn float() -> Self {
        FfiType(unsafe { &mut low::ffi_type_float })
    }

    pub fn double() -> Self {
        FfiType(unsafe { &mut low::ffi_type_double })
    }

    pub fn pointer() -> Self {
        FfiType(unsafe { &mut low::ffi_type_pointer })
    }

    pub fn longdouble() -> Self {
        FfiType(unsafe { &mut low::ffi_type_longdouble })
    }

    pub fn complex_float() -> Self {
        FfiType(unsafe { &mut low::ffi_type_complex_float })
    }

    pub fn complex_double() -> Self {
        FfiType(unsafe { &mut low::ffi_type_complex_double })
    }

    pub fn complex_longdouble() -> Self {
        FfiType(unsafe { &mut low::ffi_type_complex_longdouble })
    }

    pub fn structure(fields: Vec<FfiType>) -> Self {
        println!("FfiType::structure({:?})", fields);
        unsafe {
            FfiType(ffi_type_struct_create(fields))
        }
    }

    pub fn as_raw_ptr(&self) -> *mut low::ffi_type {
        self.0
    }
}

impl FfiTypeArray {
    pub fn new(types: Vec<FfiType>) -> Self {
        let len = types.len();
        unsafe {
            FfiTypeArray {
                ptr: ffi_type_array_create(types),
                len: len,
            }
        }
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn as_raw_ptr(&self) -> *mut *mut low::ffi_type {
        self.ptr
    }
}

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

    #[test]
    fn create_uint64() {
        FfiType::uint64();
    }

    #[test]
    fn create_struct() {
        FfiType::structure(vec![FfiType::sint64(),
                                FfiType::sint64(),
                                FfiType::uint64()]);
    }

}
<commit_msg>implemented ffi_type_clone<commit_after>use std::{mem, ptr};
use libc;

use c;
use low;

type FfiType_      = *mut low::ffi_type;
type FfiTypeArray_ = *mut FfiType_;

#[derive(Debug)]
pub struct FfiType(FfiType_);

#[derive(Debug)]
pub struct FfiTypeArray {
    ptr: FfiTypeArray_,
    len: usize,
}

/// Creates a null-terminated array of FfiType_. Takes ownership of
/// the elements.
unsafe fn ffi_type_array_create(elements: Vec<FfiType>) -> FfiTypeArray_ {
    let size = elements.len();
    let array = libc::malloc((size+1) * mem::size_of::<FfiType_>())
                    as FfiTypeArray_;

    for i in 0 .. size {
        *array.offset(i as isize) = elements[i].0;
    }
    *array.offset(size as isize) = ptr::null::<low::ffi_type>() as FfiType_;

    for t in elements {
        mem::forget(t);
    }

    println!("ffi_type_array_create(...) = {:?}", array);

    array
}

/// Creates a struct ffi_type with the given elements. Takes ownership
/// of the elements.
unsafe fn ffi_type_struct_create(elements: Vec<FfiType>) -> FfiType_ {
    println!("ffi_type_array_create({:?})", elements);
    let array    = ffi_type_array_create(elements);
    let ffi_type = libc::malloc(mem::size_of::<low::ffi_type>())
                       as FfiType_;

    (*ffi_type).size      = 0;
    (*ffi_type).alignment = 0;
    (*ffi_type).type_     = c::ffi_type_enum::STRUCT as u16;
    (*ffi_type).elements  = array;

    println!("ffi_type_array_create(...) = {:?}", ffi_type);

    ffi_type
}

unsafe fn ffi_type_array_clone(ffi_types: FfiTypeArray_) -> FfiTypeArray_ {
    let mut current = ffi_types;
    let mut count   = 0;
    while !(*current).is_null() {
        current = current.offset(1);
        count += 1;
    }

    let new = libc::malloc((count+1) * mem::size_of::<FfiType_>())
                    as FfiTypeArray_;

    for i in 0 .. count {
        *new.offset(i as isize) = ffi_type_clone(*ffi_types.offset(i as isize));
    }
    *new.offset(count as isize) = ptr::null::<low::ffi_type>() as FfiType_;

    new
}

unsafe fn ffi_type_clone(ffi_type: FfiType_) -> FfiType_ {
    if (*ffi_type).type_ == c::ffi_type_enum::STRUCT as u16 {
        let elements = ffi_type_array_clone((*ffi_type).elements);
        let new = libc::malloc(mem::size_of::<low::ffi_type>()) as FfiType_;

        (*new).size      = 0;
        (*new).alignment = 0;
        (*new).type_     = c::ffi_type_enum::STRUCT as u16;
        (*new).elements  = elements;

        new
    } else {
        ffi_type
    }
}

/// Destroys an array of FfiType_ and all of its elements.
unsafe fn ffi_type_array_destroy(ffi_types: FfiTypeArray_) {
    println!("ffi_type_array_destroy({:?})", ffi_types);
    let mut current = ffi_types;
    while !(*current).is_null() {
        ffi_type_destroy(*current);
        current = current.offset(1);
    }

    libc::free(ffi_types as *mut libc::c_void);
}

/// Destroys an FfiType_ if it was dynamically allocated.
unsafe fn ffi_type_destroy(ffi_type: FfiType_) {
    println!("ffi_type_destroy({:?})", ffi_type);
    if (*ffi_type).type_ == c::ffi_type_enum::STRUCT as u16 {
        ffi_type_array_destroy((*ffi_type).elements);
        libc::free(ffi_type as *mut libc::c_void);
    }
}

impl Drop for FfiType {
    fn drop(&mut self) {
        unsafe { ffi_type_destroy(self.0) }
    }
}

impl Drop for FfiTypeArray {
    fn drop(&mut self) {
        unsafe { ffi_type_array_destroy(self.ptr) }
    }
}

impl FfiType {
    pub fn void() -> Self {
        FfiType(unsafe { &mut low::ffi_type_void })
    }

    pub fn uint8() -> Self {
        FfiType(unsafe { &mut low::ffi_type_uint8 })
    }

    pub fn sint8() -> Self {
        FfiType(unsafe { &mut low::ffi_type_sint8 })
    }

    pub fn uint16() -> Self {
        FfiType(unsafe { &mut low::ffi_type_uint16 })
    }

    pub fn sint16() -> Self {
        FfiType(unsafe { &mut low::ffi_type_sint16 })
    }

    pub fn uint32() -> Self {
        FfiType(unsafe { &mut low::ffi_type_uint32 })
    }

    pub fn sint32() -> Self {
        FfiType(unsafe { &mut low::ffi_type_sint32 })
    }

    pub fn uint64() -> Self {
        FfiType(unsafe { &mut low::ffi_type_uint64 })
    }

    pub fn sint64() -> Self {
        FfiType(unsafe { &mut low::ffi_type_sint64 })
    }

    pub fn float() -> Self {
        FfiType(unsafe { &mut low::ffi_type_float })
    }

    pub fn double() -> Self {
        FfiType(unsafe { &mut low::ffi_type_double })
    }

    pub fn pointer() -> Self {
        FfiType(unsafe { &mut low::ffi_type_pointer })
    }

    pub fn longdouble() -> Self {
        FfiType(unsafe { &mut low::ffi_type_longdouble })
    }

    pub fn complex_float() -> Self {
        FfiType(unsafe { &mut low::ffi_type_complex_float })
    }

    pub fn complex_double() -> Self {
        FfiType(unsafe { &mut low::ffi_type_complex_double })
    }

    pub fn complex_longdouble() -> Self {
        FfiType(unsafe { &mut low::ffi_type_complex_longdouble })
    }

    pub fn structure(fields: Vec<FfiType>) -> Self {
        println!("FfiType::structure({:?})", fields);
        unsafe {
            FfiType(ffi_type_struct_create(fields))
        }
    }

    pub fn as_raw_ptr(&self) -> *mut low::ffi_type {
        self.0
    }
}

impl FfiTypeArray {
    pub fn new(types: Vec<FfiType>) -> Self {
        let len = types.len();
        unsafe {
            FfiTypeArray {
                ptr: ffi_type_array_create(types),
                len: len,
            }
        }
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn as_raw_ptr(&self) -> *mut *mut low::ffi_type {
        self.ptr
    }
}

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

    #[test]
    fn create_uint64() {
        FfiType::uint64();
    }

    #[test]
    fn create_struct() {
        FfiType::structure(vec![FfiType::sint64(),
                                FfiType::sint64(),
                                FfiType::uint64()]);
    }

}
<|endoftext|>


<commit_before>// -*- rust-indent-offset: 2 -*-
// More compact, just for slides.

pub type Pair<'a> = (int, &'a str);

#[deriving(Show, Clone)]
pub struct Config<'a>(pub Vec<Pair<'a>>);

static DIVISOR_MIN: int = 2;
static DIVISOR_MAX: int = 100;

fn validate_pair(&(d, _): &Pair) {
  assert!(d >= DIVISOR_MIN,
          "divisor {} must be >= {}", d, DIVISOR_MIN);
  assert!(d <= DIVISOR_MAX,
          "divisor {} must be <= {}", d, DIVISOR_MAX);
}

impl<'a> Config<'a> {
  pub fn new(pairs: Vec<Pair>) -> Config {
    for pair in pairs.iter() {
      validate_pair(pair);
    }
    Config(pairs)
  }
}

// Make the decision to allocate a String here.
fn rule(&(n, word): &Pair, i: int) -> Option<String> {
  if i % n == 0 {
    Some(word.to_string())
  } else {
    None
  }
}

fn add_option(a1: Option<String>, a2: Option<String>)
              -> Option<String> {
  match (a1, a2) {
    (Some(s1), None)     => Some(s1),
    (None,     Some(s2)) => Some(s2),
    (Some(s1), Some(s2)) => Some(s1 + s2),
    (None,     None)     => None,
  }
}

// TODO Use closures to "compile", when Rust supports that.
pub fn evaluate(Config(pairs): Config, i: int)
                -> String {
  let combined: Option<String> = pairs.iter()
            .map(|pair| rule(pair, i))
            .fold(None, add_option);
  combined.unwrap_or_else(|| i.to_string())
}

#[cfg(test)]
mod test {
  use super::{DIVISOR_MIN, DIVISOR_MAX, Config, evaluate};
  use quickcheck::TestResult;

/*
  // TODO Create own Arbitrary instances for Config.
  // Too complex in Rust for presentation, unfortunately.
  use quickcheck::{Arbitrary, Gen, Shrinker};

  // Trivial in Scala and Haskell.
  impl Arbitrary for Config {
    fn arbitrary<G: Gen>(g: &mut G) -> Config {
      // TODO
      fail!()
    }
    fn shrink(&self) -> Box<Shrinker<Config>> {
      // TODO
      fail!()
    }
  }
*/

  // TODO nest the checks, staging the check of i.
  // Trivial in Scala and Haskell.
  #[quickcheck]
  fn d1_but_not_d2((d1, w1): (int, String),
                   (d2, w2): (int, String),
                   i: int) -> TestResult {
    if (d1 >= DIVISOR_MIN && d1 <= DIVISOR_MAX)
      && (d2 >= DIVISOR_MIN && d2 <= DIVISOR_MAX) {
      let config = Config::new(vec![(d1, w1.as_slice()),
                                    (d2, w2.as_slice())]);

      if i % d1 == 0 && i % d2 != 0 {
        TestResult::from_bool(evaluate(config, i) == w1)
      } else {
        TestResult::discard()
      }
    } else {
      TestResult::discard()
    }
  }

  // TODO the other three cases are similar.
}
<commit_msg>Change s1 + s2 to s1.append(s2.as_slice()) for efficiency.<commit_after>// -*- rust-indent-offset: 2 -*-
// More compact, just for slides.

pub type Pair<'a> = (int, &'a str);

#[deriving(Show, Clone)]
pub struct Config<'a>(pub Vec<Pair<'a>>);

static DIVISOR_MIN: int = 2;
static DIVISOR_MAX: int = 100;

fn validate_pair(&(d, _): &Pair) {
  assert!(d >= DIVISOR_MIN,
          "divisor {} must be >= {}", d, DIVISOR_MIN);
  assert!(d <= DIVISOR_MAX,
          "divisor {} must be <= {}", d, DIVISOR_MAX);
}

impl<'a> Config<'a> {
  pub fn new(pairs: Vec<Pair>) -> Config {
    for pair in pairs.iter() {
      validate_pair(pair);
    }
    Config(pairs)
  }
}

// Make the decision to allocate a String here.
fn rule(&(n, word): &Pair, i: int) -> Option<String> {
  if i % n == 0 {
    Some(word.to_string())
  } else {
    None
  }
}

fn add_option(a1: Option<String>, a2: Option<String>)
              -> Option<String> {
  match (a1, a2) {
    (Some(s1), None)     => Some(s1),
    (None,     Some(s2)) => Some(s2),
    (Some(s1), Some(s2)) => Some(s1.append(s2.as_slice())),
    (None,     None)     => None,
  }
}

// TODO Use closures to "compile", when Rust supports that.
pub fn evaluate(Config(pairs): Config, i: int)
                -> String {
  let combined: Option<String> = pairs.iter()
            .map(|pair| rule(pair, i))
            .fold(None, add_option);
  combined.unwrap_or_else(|| i.to_string())
}

#[cfg(test)]
mod test {
  use super::{DIVISOR_MIN, DIVISOR_MAX, Config, evaluate};
  use quickcheck::TestResult;

/*
  // TODO Create own Arbitrary instances for Config.
  // Too complex in Rust for presentation, unfortunately.
  use quickcheck::{Arbitrary, Gen, Shrinker};

  // Trivial in Scala and Haskell.
  impl Arbitrary for Config {
    fn arbitrary<G: Gen>(g: &mut G) -> Config {
      // TODO
      fail!()
    }
    fn shrink(&self) -> Box<Shrinker<Config>> {
      // TODO
      fail!()
    }
  }
*/

  // TODO nest the checks, staging the check of i.
  // Trivial in Scala and Haskell.
  #[quickcheck]
  fn d1_but_not_d2((d1, w1): (int, String),
                   (d2, w2): (int, String),
                   i: int) -> TestResult {
    if (d1 >= DIVISOR_MIN && d1 <= DIVISOR_MAX)
      && (d2 >= DIVISOR_MIN && d2 <= DIVISOR_MAX) {
      let config = Config::new(vec![(d1, w1.as_slice()),
                                    (d2, w2.as_slice())]);

      if i % d1 == 0 && i % d2 != 0 {
        TestResult::from_bool(evaluate(config, i) == w1)
      } else {
        TestResult::discard()
      }
    } else {
      TestResult::discard()
    }
  }

  // TODO the other three cases are similar.
}
<|endoftext|>


<commit_before>use std::io::Write;
use std::fmt::Display;
use std::fs::{PathExt, File, create_dir_all, set_permissions};
use std::os::unix::ffi::OsStrExt;
use std::os::unix::fs::{symlink, MetadataExt, PermissionsExt};
use std::path::Path;
use std::collections::{BTreeMap, BTreeSet};

use cronparse::Limited;
use cronparse::crontab::{CrontabEntry, SystemCrontabEntry, UserCrontabEntry};
use cronparse::schedule::{Schedule, Period, Calendar};
use cronparse::interval::Interval;

use pgs_files::passwd;

macro_rules! try_ {
    ($exp:expr) => {
        match $exp {
            Ok(value) => value,
            Err(err) => { error!("{}", err); return; }
        }
    }
}

pub fn generate_systemd_units(entry: CrontabEntry, env: &BTreeMap<String, String>, path: &Path, dstdir: &Path) {
    use cronparse::crontab::CrontabEntry::*;

    info!("generating units for {}: {:?}, {:?}", path.display(), entry, env);

    let owner = try_!(path.metadata()).uid();

    let mut persistent = env.get("PERSISTENT").and_then(|v| match &**v {
        "yes" | "true" | "1" => Some(true),
        "auto" | "" => None,
        _ => Some(false)
    }).unwrap_or_else(|| match entry {
        Anacron(_) | User(UserCrontabEntry { sched: Schedule::Period(_), .. }) | System(SystemCrontabEntry { sched: Schedule::Period(_), .. }) => true,
        _ => false
    });

    let batch = env.get("BATCH").map(|v| match &**v {
        "yes" | "true" | "1" => true,
        _ => false
    }).unwrap_or(false);

    let random_delay = env.get("RANDOM_DELAY").and_then(|v| v.parse::<u64>().ok()).unwrap_or(1);
    let mut delay = env.get("DELAY").and_then(|v| v.parse::<u64>().ok()).unwrap_or(0);
    let hour = env.get("START_HOURS_RANGE").and_then(|v| v.splitn(1, '-').next().and_then(|v| v.parse::<u64>().ok())).unwrap_or(0);
    let shell = env.get("SHELL").map(|v| &**v).unwrap_or("/bin/sh");

    let schedule = entry.period().and_then(|period| match *period {
        Period::Reboot => {
            persistent = false;
            if delay == 0 {
                delay = 1;
            }
            None
        },
        Period::Minutely => {
            persistent = false;
            Some("@minutely".to_owned())
        },
        Period::Hourly => {
            if delay == 0 {
                Some("@hourly".to_owned())
            } else {
                Some(format!("*-*-* *:{}:0", delay))
            }
        },
        Period::Midnight => {
            if delay == 0 {
                Some("@daily".to_owned())
            } else {
                Some(format!("*-*-* 0:{}:0", delay))
            }
        },
        Period::Daily => {
            if delay == 0 && hour == 0 {
                Some("@daily".to_owned())
            } else {
                Some(format!("*-*-* {}:{}:0", hour, delay))
            }
        },
        Period::Weekly => {
            if delay == 0 && hour == 0 {
                Some("@weekly".to_owned())
            } else {
                Some(format!("Mon *-*-* {}:{}:0", hour, delay))
            }
        },
        Period::Monthly => {
            if delay == 0 && hour == 0 {
                Some("@monthly".to_owned())
            } else {
                Some(format!("*-*-1 {}:{}:0", hour, delay))
            }
        },
        Period::Quaterly => {
            if delay == 0 && hour == 0 {
                Some("@quaterly".to_owned())
            } else {
                Some(format!("*-1,4,7,10-1 {}:{}:0", hour, delay))
            }
        },
        Period::Biannually => {
            if delay == 0 && hour == 0 {
                Some("@semi-annually".to_owned())
            } else {
                Some(format!("*-1,7-1 {}:{}:0", hour, delay))
            }
        },
        Period::Yearly => {
            if delay == 0 && hour == 0 {
                Some("@yearly".to_owned())
            } else {
                Some(format!("*-1-1 {}:{}:0", hour, delay))
            }
        },
        Period::Days(days) => {
            // workaround for anacrontab
            if days > 31 {
                Some(format!("*-1/{}-1 {}:{}:0", days / 30, hour, delay))
            } else {
                Some(format!("*-*-1/{} {}:{}:0", days, hour, delay))
            }
        },
    }).or_else(|| entry.calendar().and_then(|cal| {
        let Calendar {
            ref dows,
            ref days,
            ref mons,
            ref hrs,
            ref mins
        } = *cal;

        Some(format!("{} *-{}-{} {}:{}:00",
                     linearize(&**dows),
                     linearize(&**mons),
                     linearize(&**days),
                     linearize(&**hrs),
                     linearize(&**mins)))
    }));

    if let Some(cmd) = entry.command() {
        let mut md5ctx = ::md5::Context::new();
        md5ctx.consume(path.as_os_str().as_bytes());
        if let Some(ref schedule) = schedule {
            md5ctx.consume(schedule.as_bytes());
        }
        md5ctx.consume(cmd.as_bytes());
        let md5hex = tohex(&md5ctx.compute());

        let service_unit_name = format!("cronjob-{}.service", md5hex);
        let timer_unit_name = format!("cronjob-{}.timer", md5hex);

        let service_unit_path = dstdir.join(&service_unit_name);
        let timer_unit_path = dstdir.join(&timer_unit_name);
        let cron_target_wants_path = dstdir.join("cron.target.wants");
        try_!(create_dir_all(&cron_target_wants_path));

        let user = entry.user().and_then(passwd::get_entry_by_name).or_else(|| passwd::get_entry_by_uid(owner)).unwrap();

        let command = if Path::new(cmd).is_file() {
            cmd.to_owned()
        } else {
            let script_command_path = dstdir.join(format!("cronjob-{}.sh", md5hex));

            debug!("generating script {:?} from {:?}", script_command_path, path);
            {
                let mut script_command_file = try_!(File::create(&script_command_path));
                try_!(writeln!(script_command_file, "#!{}", shell));
                try_!(writeln!(script_command_file, "{}", cmd));
            }

            let mut perms = try_!(script_command_path.metadata()).permissions();
            perms.set_mode(0o755);
            try_!(set_permissions(&script_command_path, perms));
            script_command_path.to_str().unwrap().to_owned()
        };

        debug!("generating timer {:?} from {:?}", timer_unit_path, path);
        {
            let mut timer_unit_file = try_!(File::create(&timer_unit_path));

            try_!(writeln!(timer_unit_file, r###"[Unit]
Description=[Timer] "{entry}"
Documentation=man:systemd-crontab-generator(8)
PartOf=cron.target
RefuseManualStart=true
RefuseManualStop=true
SourcePath={source_crontab_path}

[Timer]
Unit={service_unit_name}
Persistent={persistent}"###,
                entry = entry,
                source_crontab_path = path.display(),
                service_unit_name = service_unit_name,
                persistent = persistent,
                ));

            if let Some(schedule) = schedule {
                try_!(writeln!(timer_unit_file, "OnCalendar={}", schedule));
            } else {
                try_!(writeln!(timer_unit_file, "OnBootSec={}m", delay));
            }

            if random_delay != 1 {
                try_!(writeln!(timer_unit_file, "AccuracySec={}m", random_delay));
            }
        }
        try_!(symlink(timer_unit_path, cron_target_wants_path.join(timer_unit_name)));

        debug!("generating service {:?} from {:?}", service_unit_path, path);
        {
            let mut service_unit_file = try_!(File::create(service_unit_path));

            try_!(writeln!(service_unit_file, r###"[Unit]
Description=[Cron] "{entry}"
Documentation=man:systemd-crontab-generator(8)
RefuseManualStart=true
RefuseManualStop=true
SourcePath={source_crontab_path}

[Service]
Type=oneshot
IgnoreSIGPIPE=false
ExecStart={command}"###,
                entry = entry,
                source_crontab_path = path.display(),
                command = command,
                ));

            if user.uid != 0 {
                try_!(writeln!(service_unit_file, "User={}", user.name));
            }

            if let Some(group) = entry.group() {
                try_!(writeln!(service_unit_file, "Group={}", group));
            }
            if batch {
                try_!(writeln!(service_unit_file, "CPUSchedulingPolicy=idle"));
                try_!(writeln!(service_unit_file, "IOSchedulingClass=idle"));
            }

            if !env.is_empty() {
                try_!(write!(service_unit_file, "Environment="));
                for (name, value) in env.iter() {
                    try_!(write!(service_unit_file, r#""{}={}""#, name, value));
                }
                try_!(write!(service_unit_file, "\n"));
            }
        }
    }
}

fn linearize<T: Limited + Display>(input: &[Interval<T>]) -> String {
    if input.len() == 1 && input[0] == Interval::Full(1) {
        "*".to_owned()
    } else {
        let mut output = String::new();
        for part in input.iter().flat_map(|v| v.iter()).collect::<BTreeSet<_>>().iter() {
            output.push_str(&*part.to_string());
            output.push(',');
        }
        output.pop();
        output
    }
}

fn tohex(input: &[u8]) -> String {
    #[inline]
    fn hex(d: u8) -> char {
        match d {
            0...9 => (d + 0x30) as char,
            10...15 => (d + 0x57) as char,
            _ => unreachable!("unexpected value: {}", d)
        }
    }

    let mut buf = String::with_capacity(32);
    for b in input.into_iter() {
        buf.push(hex(b >> 4));
        buf.push(hex(b & 0xf));
    }
    buf
}

<commit_msg>log entry in a nicer form<commit_after>use std::io::Write;
use std::fmt::Display;
use std::fs::{PathExt, File, create_dir_all, set_permissions};
use std::os::unix::ffi::OsStrExt;
use std::os::unix::fs::{symlink, MetadataExt, PermissionsExt};
use std::path::Path;
use std::collections::{BTreeMap, BTreeSet};

use cronparse::Limited;
use cronparse::crontab::{CrontabEntry, SystemCrontabEntry, UserCrontabEntry};
use cronparse::schedule::{Schedule, Period, Calendar};
use cronparse::interval::Interval;

use pgs_files::passwd;

macro_rules! try_ {
    ($exp:expr) => {
        match $exp {
            Ok(value) => value,
            Err(err) => { error!("{}", err); return; }
        }
    }
}

pub fn generate_systemd_units(entry: CrontabEntry, env: &BTreeMap<String, String>, path: &Path, dstdir: &Path) {
    use cronparse::crontab::CrontabEntry::*;

    info!("generating units for {}: \"{}\", {:?}", path.display(), entry, env);

    let owner = try_!(path.metadata()).uid();

    let mut persistent = env.get("PERSISTENT").and_then(|v| match &**v {
        "yes" | "true" | "1" => Some(true),
        "auto" | "" => None,
        _ => Some(false)
    }).unwrap_or_else(|| match entry {
        Anacron(_) | User(UserCrontabEntry { sched: Schedule::Period(_), .. }) | System(SystemCrontabEntry { sched: Schedule::Period(_), .. }) => true,
        _ => false
    });

    let batch = env.get("BATCH").map(|v| match &**v {
        "yes" | "true" | "1" => true,
        _ => false
    }).unwrap_or(false);

    let random_delay = env.get("RANDOM_DELAY").and_then(|v| v.parse::<u64>().ok()).unwrap_or(1);
    let mut delay = env.get("DELAY").and_then(|v| v.parse::<u64>().ok()).unwrap_or(0);
    let hour = env.get("START_HOURS_RANGE").and_then(|v| v.splitn(1, '-').next().and_then(|v| v.parse::<u64>().ok())).unwrap_or(0);
    let shell = env.get("SHELL").map(|v| &**v).unwrap_or("/bin/sh");

    let schedule = entry.period().and_then(|period| match *period {
        Period::Reboot => {
            persistent = false;
            if delay == 0 {
                delay = 1;
            }
            None
        },
        Period::Minutely => {
            persistent = false;
            Some("@minutely".to_owned())
        },
        Period::Hourly => {
            if delay == 0 {
                Some("@hourly".to_owned())
            } else {
                Some(format!("*-*-* *:{}:0", delay))
            }
        },
        Period::Midnight => {
            if delay == 0 {
                Some("@daily".to_owned())
            } else {
                Some(format!("*-*-* 0:{}:0", delay))
            }
        },
        Period::Daily => {
            if delay == 0 && hour == 0 {
                Some("@daily".to_owned())
            } else {
                Some(format!("*-*-* {}:{}:0", hour, delay))
            }
        },
        Period::Weekly => {
            if delay == 0 && hour == 0 {
                Some("@weekly".to_owned())
            } else {
                Some(format!("Mon *-*-* {}:{}:0", hour, delay))
            }
        },
        Period::Monthly => {
            if delay == 0 && hour == 0 {
                Some("@monthly".to_owned())
            } else {
                Some(format!("*-*-1 {}:{}:0", hour, delay))
            }
        },
        Period::Quaterly => {
            if delay == 0 && hour == 0 {
                Some("@quaterly".to_owned())
            } else {
                Some(format!("*-1,4,7,10-1 {}:{}:0", hour, delay))
            }
        },
        Period::Biannually => {
            if delay == 0 && hour == 0 {
                Some("@semi-annually".to_owned())
            } else {
                Some(format!("*-1,7-1 {}:{}:0", hour, delay))
            }
        },
        Period::Yearly => {
            if delay == 0 && hour == 0 {
                Some("@yearly".to_owned())
            } else {
                Some(format!("*-1-1 {}:{}:0", hour, delay))
            }
        },
        Period::Days(days) => {
            // workaround for anacrontab
            if days > 31 {
                Some(format!("*-1/{}-1 {}:{}:0", days / 30, hour, delay))
            } else {
                Some(format!("*-*-1/{} {}:{}:0", days, hour, delay))
            }
        },
    }).or_else(|| entry.calendar().and_then(|cal| {
        let Calendar {
            ref dows,
            ref days,
            ref mons,
            ref hrs,
            ref mins
        } = *cal;

        Some(format!("{} *-{}-{} {}:{}:00",
                     linearize(&**dows),
                     linearize(&**mons),
                     linearize(&**days),
                     linearize(&**hrs),
                     linearize(&**mins)))
    }));

    if let Some(cmd) = entry.command() {
        let mut md5ctx = ::md5::Context::new();
        md5ctx.consume(path.as_os_str().as_bytes());
        if let Some(ref schedule) = schedule {
            md5ctx.consume(schedule.as_bytes());
        }
        md5ctx.consume(cmd.as_bytes());
        let md5hex = tohex(&md5ctx.compute());

        let service_unit_name = format!("cronjob-{}.service", md5hex);
        let timer_unit_name = format!("cronjob-{}.timer", md5hex);

        let service_unit_path = dstdir.join(&service_unit_name);
        let timer_unit_path = dstdir.join(&timer_unit_name);
        let cron_target_wants_path = dstdir.join("cron.target.wants");
        try_!(create_dir_all(&cron_target_wants_path));

        let user = entry.user().and_then(passwd::get_entry_by_name).or_else(|| passwd::get_entry_by_uid(owner)).unwrap();

        let command = if Path::new(cmd).is_file() {
            cmd.to_owned()
        } else {
            let script_command_path = dstdir.join(format!("cronjob-{}.sh", md5hex));

            debug!("generating script {:?} from {:?}", script_command_path, path);
            {
                let mut script_command_file = try_!(File::create(&script_command_path));
                try_!(writeln!(script_command_file, "#!{}", shell));
                try_!(writeln!(script_command_file, "{}", cmd));
            }

            let mut perms = try_!(script_command_path.metadata()).permissions();
            perms.set_mode(0o755);
            try_!(set_permissions(&script_command_path, perms));
            script_command_path.to_str().unwrap().to_owned()
        };

        debug!("generating timer {:?} from {:?}", timer_unit_path, path);
        {
            let mut timer_unit_file = try_!(File::create(&timer_unit_path));

            try_!(writeln!(timer_unit_file, r###"[Unit]
Description=[Timer] "{entry}"
Documentation=man:systemd-crontab-generator(8)
PartOf=cron.target
RefuseManualStart=true
RefuseManualStop=true
SourcePath={source_crontab_path}

[Timer]
Unit={service_unit_name}
Persistent={persistent}"###,
                entry = entry,
                source_crontab_path = path.display(),
                service_unit_name = service_unit_name,
                persistent = persistent,
                ));

            if let Some(schedule) = schedule {
                try_!(writeln!(timer_unit_file, "OnCalendar={}", schedule));
            } else {
                try_!(writeln!(timer_unit_file, "OnBootSec={}m", delay));
            }

            if random_delay != 1 {
                try_!(writeln!(timer_unit_file, "AccuracySec={}m", random_delay));
            }
        }
        try_!(symlink(timer_unit_path, cron_target_wants_path.join(timer_unit_name)));

        debug!("generating service {:?} from {:?}", service_unit_path, path);
        {
            let mut service_unit_file = try_!(File::create(service_unit_path));

            try_!(writeln!(service_unit_file, r###"[Unit]
Description=[Cron] "{entry}"
Documentation=man:systemd-crontab-generator(8)
RefuseManualStart=true
RefuseManualStop=true
SourcePath={source_crontab_path}

[Service]
Type=oneshot
IgnoreSIGPIPE=false
ExecStart={command}"###,
                entry = entry,
                source_crontab_path = path.display(),
                command = command,
                ));

            if user.uid != 0 {
                try_!(writeln!(service_unit_file, "User={}", user.name));
            }

            if let Some(group) = entry.group() {
                try_!(writeln!(service_unit_file, "Group={}", group));
            }
            if batch {
                try_!(writeln!(service_unit_file, "CPUSchedulingPolicy=idle"));
                try_!(writeln!(service_unit_file, "IOSchedulingClass=idle"));
            }

            if !env.is_empty() {
                try_!(write!(service_unit_file, "Environment="));
                for (name, value) in env.iter() {
                    try_!(write!(service_unit_file, r#""{}={}""#, name, value));
                }
                try_!(write!(service_unit_file, "\n"));
            }
        }
    }
}

fn linearize<T: Limited + Display>(input: &[Interval<T>]) -> String {
    if input.len() == 1 && input[0] == Interval::Full(1) {
        "*".to_owned()
    } else {
        let mut output = String::new();
        for part in input.iter().flat_map(|v| v.iter()).collect::<BTreeSet<_>>().iter() {
            output.push_str(&*part.to_string());
            output.push(',');
        }
        output.pop();
        output
    }
}

fn tohex(input: &[u8]) -> String {
    #[inline]
    fn hex(d: u8) -> char {
        match d {
            0...9 => (d + 0x30) as char,
            10...15 => (d + 0x57) as char,
            _ => unreachable!("unexpected value: {}", d)
        }
    }

    let mut buf = String::with_capacity(32);
    for b in input.into_iter() {
        buf.push(hex(b >> 4));
        buf.push(hex(b & 0xf));
    }
    buf
}

<|endoftext|>