<commit_before>//! # Socket code around `libc`
//!
//! ## Notes
//!
//! This module provides a low level one-to-one mapping between `libc` system call wrappers
//! with defaults specific to netlink sockets as well as a higher level API for simplification
//! of netlink code.

use std::io;
use std::os::unix::io::{AsRawFd,IntoRawFd,RawFd};
use std::marker::PhantomData;
use std::mem::{zeroed,size_of};
#[cfg(feature = "stream")]
use std::time::Duration;

use libc::{self,c_int,c_void};
#[cfg(feature = "evented")]
use mio::{self,Evented};
#[cfg(feature = "stream")]
use tokio::prelude::{Async,Stream};

#[cfg(feature = "stream")]
use {Nl,MemRead,MemWrite,MAX_NL_LENGTH};
#[cfg(not(feature = "stream"))]
use {Nl,MemRead,MemWrite};
use ffi::NlFamily;
#[cfg(feature = "stream")]
use nlhdr::NlHdr;

/// Handle for the socket file descriptor
#[cfg(feature = "stream")]
pub struct NlSocket<I, P> {
    fd: c_int,
    poll: mio::Poll,
    data_type: PhantomData<I>,
    data_payload: PhantomData<P>,
}

/// Handle for the socket file descriptor
#[cfg(not(feature = "stream"))]
pub struct NlSocket<I, P> {
    fd: c_int,
    data_type: PhantomData<I>,
    data_payload: PhantomData<P>,
}

impl<I, P> NlSocket<I, P> where I: Nl, P: Nl {
    /// Wrapper around `socket()` syscall filling in the netlink-specific information
    #[cfg(feature = "stream")]
    pub fn new(proto: NlFamily) -> Result<Self, io::Error> {
        let fd = match unsafe {
            libc::socket(libc::AF_NETLINK, libc::SOCK_RAW, proto.into())
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }?;
        let poll = mio::Poll::new()?;
        let socket = NlSocket { fd, poll, data_type: PhantomData, data_payload: PhantomData };
        socket.register(&socket.poll, mio::Token(0), mio::Ready::readable(), mio::PollOpt::edge())?;
        Ok(socket)
    }

    /// Wrapper around `socket()` syscall filling in the netlink-specific information
    #[cfg(not(feature = "stream"))]
    pub fn new(proto: NlFamily) -> Result<Self, io::Error> {
        let fd = match unsafe {
            libc::socket(libc::AF_NETLINK, libc::SOCK_RAW, proto.into())
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }?;
        Ok(NlSocket { fd, data_type: PhantomData, data_payload: PhantomData })
    }

    /// Use this function to bind to a netlink ID and subscribe to groups. See netlink(7)
    /// man pages for more information on netlink IDs and groups.
    pub fn bind(&mut self, pid: Option<u32>, groups: Option<u32>) -> Result<(), io::Error> {
        let mut nladdr = unsafe { zeroed::<libc::sockaddr_nl>() };
        nladdr.nl_family = libc::AF_NETLINK as u16;
        nladdr.nl_pid = pid.unwrap_or(0);
        nladdr.nl_groups = groups.unwrap_or(0);
        match unsafe {
            libc::bind(self.fd, &nladdr as *const _ as *const libc::sockaddr,
                       size_of::<libc::sockaddr_nl>() as u32)
        } {
            i if i >= 0 => Ok(()),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Send message encoded as byte slice to the netlink ID specified in the netlink header
    /// (`nl::nlhdr::NlHdr`).
    pub fn send(&mut self, buf: MemRead, flags: i32) -> Result<isize, io::Error> {
        match unsafe {
            libc::send(self.fd, buf.as_slice() as *const _ as *const c_void, buf.len(), flags)
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Receive message encoded as byte slice from the netlink socket.
    pub fn recv<'a>(&mut self, mut buf: MemWrite<'a>, flags: i32)
            -> Result<MemRead<'a>, io::Error> {
        match unsafe {
            libc::recv(self.fd, buf.as_mut_slice() as *mut _ as *mut c_void, buf.len(), flags)
        } {
            i if i >= 0 => Ok(buf.shrink(i as usize).into()),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Equivalent of `socket` and `bind` calls.
    pub fn connect(proto: NlFamily, pid: Option<u32>, groups: Option<u32>)
                   -> Result<Self, io::Error> {
        let mut s = try!(NlSocket::new(proto));
        try!(s.bind(pid, groups));
        Ok(s)
    }
}

impl<I, P> AsRawFd for NlSocket<I, P> {
    fn as_raw_fd(&self) -> RawFd {
        self.fd
    }
}

impl<I, P> IntoRawFd for NlSocket<I, P> {
    fn into_raw_fd(self) -> RawFd {
        self.fd
    }
}

#[cfg(feature = "stream")]
impl<I, P> Stream for NlSocket<I, P> where I: Nl, P: Nl {
    type Item = NlHdr<I, P>;
    type Error = ();

    fn poll(&mut self) -> Result<Async<Option<Self::Item>>, Self::Error> {
        let mem = MemWrite::new_vec(Some(MAX_NL_LENGTH));
        let mut events = mio::Events::with_capacity(1);
        match self.poll.poll(&mut events, Some(Duration::from_secs(0))) {
            Ok(_) => (),
            Err(_) => return Err(()),
        };
        if let Some(event) = events.iter().nth(0) {
            if !event.readiness().is_readable() {
                return Ok(Async::NotReady);
            }
        }
        let mut mem_read = match self.recv(mem, 0) {
            Ok(mr) => mr,
            Err(_) => return Err(()),
        };
        if mem_read.as_slice().len() == 0 {
            return Ok(Async::Ready(None));
        }
        let hdr = match NlHdr::<I, P>::deserialize(&mut mem_read) {
            Ok(h) => h,
            Err(_) => return Err(()),
        };
        Ok(Async::Ready(Some(hdr)))
    }
}

#[cfg(feature = "evented")]
impl<I, P> Evented for NlSocket<I, P> {
    fn register(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready,
                opts: mio::PollOpt) -> io::Result<()> {
        poll.register(&mio::unix::EventedFd(&self.as_raw_fd()), token, interest, opts)
    }

    fn reregister(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready,
                  opts: mio::PollOpt) -> io::Result<()> {
        poll.reregister(&mio::unix::EventedFd(&self.as_raw_fd()), token, interest, opts)
    }

    fn deregister(&self, poll: &mio::Poll) -> io::Result<()> {
        poll.deregister(&mio::unix::EventedFd(&self.as_raw_fd()))
    }
}

impl<I, P> Drop for NlSocket<I, P> {
    /// Closes underlying file descriptor to avoid file descriptor leaks.
    fn drop(&mut self) {
        unsafe { libc::close(self.fd); }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use ffi::Nlmsg;
    use genlhdr::GenlHdr;

    #[test]
    fn test_socket_creation() {
        match NlSocket::<Nlmsg, GenlHdr>::connect(NlFamily::Generic, None, None) {
            Err(_) => panic!(),
            _ => (),
        }
    }
}
<commit_msg>Fix docs<commit_after>//! # Socket code around `libc`
//!
//! ## Notes
//!
//! This module provides a low level one-to-one mapping between `libc` system call wrappers
//! with defaults specific to netlink sockets as well as a higher level API for simplification
//! of netlink code.

use std::io;
use std::os::unix::io::{AsRawFd,IntoRawFd,RawFd};
use std::marker::PhantomData;
use std::mem::{zeroed,size_of};
#[cfg(feature = "stream")]
use std::time::Duration;

use libc::{self,c_int,c_void};
#[cfg(feature = "evented")]
use mio::{self,Evented};
#[cfg(feature = "stream")]
use tokio::prelude::{Async,Stream};

#[cfg(feature = "stream")]
use {Nl,MemRead,MemWrite,MAX_NL_LENGTH};
#[cfg(not(feature = "stream"))]
use {Nl,MemRead,MemWrite};
use ffi::NlFamily;
#[cfg(feature = "stream")]
use nlhdr::NlHdr;

/// Handle for the socket file descriptor
#[cfg(feature = "stream")]
pub struct NlSocket<I, P> {
    fd: c_int,
    poll: mio::Poll,
    data_type: PhantomData<I>,
    data_payload: PhantomData<P>,
}

/// Handle for the socket file descriptor
#[cfg(not(feature = "stream"))]
pub struct NlSocket<I, P> {
    fd: c_int,
    data_type: PhantomData<I>,
    data_payload: PhantomData<P>,
}

impl<I, P> NlSocket<I, P> where I: Nl, P: Nl {
    /// Wrapper around `socket()` syscall filling in the netlink-specific information
    #[cfg(feature = "stream")]
    pub fn new(proto: NlFamily) -> Result<Self, io::Error> {
        let fd = match unsafe {
            libc::socket(libc::AF_NETLINK, libc::SOCK_RAW, proto.into())
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }?;
        let poll = mio::Poll::new()?;
        let socket = NlSocket { fd, poll, data_type: PhantomData, data_payload: PhantomData };
        socket.register(&socket.poll, mio::Token(0), mio::Ready::readable(), mio::PollOpt::edge())?;
        Ok(socket)
    }

    /// Wrapper around `socket()` syscall filling in the netlink-specific information
    #[cfg(not(feature = "stream"))]
    pub fn new(proto: NlFamily) -> Result<Self, io::Error> {
        let fd = match unsafe {
            libc::socket(libc::AF_NETLINK, libc::SOCK_RAW, proto.into())
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }?;
        Ok(NlSocket { fd, data_type: PhantomData, data_payload: PhantomData })
    }

    /// Use this function to bind to a netlink ID and subscribe to groups. See netlink(7)
    /// man pages for more information on netlink IDs and groups.
    pub fn bind(&mut self, pid: Option<u32>, groups: Option<u32>) -> Result<(), io::Error> {
        let mut nladdr = unsafe { zeroed::<libc::sockaddr_nl>() };
        nladdr.nl_family = libc::AF_NETLINK as u16;
        nladdr.nl_pid = pid.unwrap_or(0);
        nladdr.nl_groups = groups.unwrap_or(0);
        match unsafe {
            libc::bind(self.fd, &nladdr as *const _ as *const libc::sockaddr,
                       size_of::<libc::sockaddr_nl>() as u32)
        } {
            i if i >= 0 => Ok(()),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Send message encoded as byte slice to the netlink ID specified in the netlink header
    /// (`neli::nlhdr::NlHdr`).
    pub fn send(&mut self, buf: MemRead, flags: i32) -> Result<isize, io::Error> {
        match unsafe {
            libc::send(self.fd, buf.as_slice() as *const _ as *const c_void, buf.len(), flags)
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Receive message encoded as byte slice from the netlink socket.
    pub fn recv<'a>(&mut self, mut buf: MemWrite<'a>, flags: i32)
            -> Result<MemRead<'a>, io::Error> {
        match unsafe {
            libc::recv(self.fd, buf.as_mut_slice() as *mut _ as *mut c_void, buf.len(), flags)
        } {
            i if i >= 0 => Ok(buf.shrink(i as usize).into()),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Equivalent of `socket` and `bind` calls.
    pub fn connect(proto: NlFamily, pid: Option<u32>, groups: Option<u32>)
                   -> Result<Self, io::Error> {
        let mut s = try!(NlSocket::new(proto));
        try!(s.bind(pid, groups));
        Ok(s)
    }
}

impl<I, P> AsRawFd for NlSocket<I, P> {
    fn as_raw_fd(&self) -> RawFd {
        self.fd
    }
}

impl<I, P> IntoRawFd for NlSocket<I, P> {
    fn into_raw_fd(self) -> RawFd {
        self.fd
    }
}

#[cfg(feature = "stream")]
impl<I, P> Stream for NlSocket<I, P> where I: Nl, P: Nl {
    type Item = NlHdr<I, P>;
    type Error = ();

    fn poll(&mut self) -> Result<Async<Option<Self::Item>>, Self::Error> {
        let mem = MemWrite::new_vec(Some(MAX_NL_LENGTH));
        let mut events = mio::Events::with_capacity(1);
        match self.poll.poll(&mut events, Some(Duration::from_secs(0))) {
            Ok(_) => (),
            Err(_) => return Err(()),
        };
        if let Some(event) = events.iter().nth(0) {
            if !event.readiness().is_readable() {
                return Ok(Async::NotReady);
            }
        }
        let mut mem_read = match self.recv(mem, 0) {
            Ok(mr) => mr,
            Err(_) => return Err(()),
        };
        if mem_read.as_slice().len() == 0 {
            return Ok(Async::Ready(None));
        }
        let hdr = match NlHdr::<I, P>::deserialize(&mut mem_read) {
            Ok(h) => h,
            Err(_) => return Err(()),
        };
        Ok(Async::Ready(Some(hdr)))
    }
}

#[cfg(feature = "evented")]
impl<I, P> Evented for NlSocket<I, P> {
    fn register(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready,
                opts: mio::PollOpt) -> io::Result<()> {
        poll.register(&mio::unix::EventedFd(&self.as_raw_fd()), token, interest, opts)
    }

    fn reregister(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready,
                  opts: mio::PollOpt) -> io::Result<()> {
        poll.reregister(&mio::unix::EventedFd(&self.as_raw_fd()), token, interest, opts)
    }

    fn deregister(&self, poll: &mio::Poll) -> io::Result<()> {
        poll.deregister(&mio::unix::EventedFd(&self.as_raw_fd()))
    }
}

impl<I, P> Drop for NlSocket<I, P> {
    /// Closes underlying file descriptor to avoid file descriptor leaks.
    fn drop(&mut self) {
        unsafe { libc::close(self.fd); }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use ffi::Nlmsg;
    use genlhdr::GenlHdr;

    #[test]
    fn test_socket_creation() {
        match NlSocket::<Nlmsg, GenlHdr>::connect(NlFamily::Generic, None, None) {
            Err(_) => panic!(),
            _ => (),
        }
    }
}
<|endoftext|>


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

use std::f32::consts::{PI, PI_2};
use std::num::Float;
use super::{Quad, Triangle, Polygon};
use super::Polygon::{PolyTri, PolyQuad};
use super::generators::{SharedVertex, IndexedPolygon};

/// Represents a sphere with radius of 1, centered at (0, 0, 0)
#[derive(Copy)]
pub struct SphereUV {
    u: usize,
    v: usize,
    sub_u: usize,
    sub_v: usize
}

impl SphereUV {
    /// Create a new sphere.
    /// `u` is the number of points across the equator of the sphere.
    /// `v` is the number of points from pole to pole.
    pub fn new(u: usize, v: usize) -> SphereUV {
        SphereUV {
            u: 0,
            v: 0,
            sub_u: u,
            sub_v: v
        }
    }

    fn vert(&self, u: usize, v: usize) -> (f32, f32, f32) {
        let u = (u as f32 / self.sub_u as f32) * PI_2;
        let v = (v as f32 / self.sub_v as f32) * PI;

        (u.cos() * v.sin(),
         u.sin() * v.sin(),
         v.cos())
    }
}

impl Iterator for SphereUV {
    type Item = Polygon<(f32, f32, f32)>;

    fn next(&mut self) -> Option<Polygon<(f32, f32, f32)>> {
        if self.u == self.sub_u {
            self.u = 0;
            self.v += 1;
            if self.v == self.sub_v {
                return None;
            }
        }

        let x = self.vert(self.u,   self.v);
        let y = self.vert(self.u,   self.v+1);
        let z = self.vert(self.u+1, self.v+1);
        let w = self.vert(self.u+1, self.v);
        let v = self.v;
        self.u += 1;

        if v == 0 {
            Some(PolyTri(Triangle::new(x, y, z)))
        } else if v == self.sub_v - 1 {
            Some(PolyTri(Triangle::new(z, w, x)))
        } else {
            Some(PolyQuad(Quad::new(x, y, z, w)))
        }
    }
}

impl SharedVertex<(f32, f32, f32)> for SphereUV {
    fn shared_vertex(&self, idx: usize) -> (f32, f32, f32) {
        if idx == 0 {
            self.vert(0, 0)
        } else if idx == self.shared_vertex_count() - 1 {
            self.vert(0, self.sub_v)
        } else {
            // since the bottom verts all map to the same
            // we jump over them in index space
            let idx = idx - 1;
            let u = idx % (self.sub_u);
            let v = idx / (self.sub_u);
            self.vert(u, v+1)
        }
    }

    fn shared_vertex_count(&self) -> usize {
        (self.sub_v - 1) * (self.sub_u) + 2
    }
}

impl IndexedPolygon<Polygon<usize>> for SphereUV {
    fn indexed_polygon(&self, idx: usize) -> Polygon<usize> {
        let u = idx % self.sub_u;
        let v = idx / self.sub_u;

        let f = |&: u: usize, v: usize| {
            if v == 0 {
                0
            } else if self.sub_v == v {
                (self.sub_v-1) * (self.sub_u) + 1
            } else {
                (v-1) * self.sub_u + (u % self.sub_u) + 1
            }
        };

        if v == 0 {
            PolyTri(Triangle::new(f(u,   v),
                                  f(u,   v+1),
                                  f(u+1, v+1)))
        } else if self.sub_v - 1 == v {
            PolyTri(Triangle::new(f(u+1, v+1),
                                  f(u+1, v),
                                  f(u,   v)))
        } else {
            PolyQuad(Quad::new(f(u,   v),
                               f(u,   v+1),
                               f(u+1, v+1),
                               f(u+1, v)))
        }
    }

    fn indexed_polygon_count(&self) -> usize {
        self.sub_v * self.sub_u
    }
}

<commit_msg>fix closure inference for the current nightly<commit_after>//   Copyright Colin Sherratt 2014
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.

use std::f32::consts::{PI, PI_2};
use std::num::Float;
use super::{Quad, Triangle, Polygon};
use super::Polygon::{PolyTri, PolyQuad};
use super::generators::{SharedVertex, IndexedPolygon};

/// Represents a sphere with radius of 1, centered at (0, 0, 0)
#[derive(Copy)]
pub struct SphereUV {
    u: usize,
    v: usize,
    sub_u: usize,
    sub_v: usize
}

impl SphereUV {
    /// Create a new sphere.
    /// `u` is the number of points across the equator of the sphere.
    /// `v` is the number of points from pole to pole.
    pub fn new(u: usize, v: usize) -> SphereUV {
        SphereUV {
            u: 0,
            v: 0,
            sub_u: u,
            sub_v: v
        }
    }

    fn vert(&self, u: usize, v: usize) -> (f32, f32, f32) {
        let u = (u as f32 / self.sub_u as f32) * PI_2;
        let v = (v as f32 / self.sub_v as f32) * PI;

        (u.cos() * v.sin(),
         u.sin() * v.sin(),
         v.cos())
    }
}

impl Iterator for SphereUV {
    type Item = Polygon<(f32, f32, f32)>;

    fn next(&mut self) -> Option<Polygon<(f32, f32, f32)>> {
        if self.u == self.sub_u {
            self.u = 0;
            self.v += 1;
            if self.v == self.sub_v {
                return None;
            }
        }

        let x = self.vert(self.u,   self.v);
        let y = self.vert(self.u,   self.v+1);
        let z = self.vert(self.u+1, self.v+1);
        let w = self.vert(self.u+1, self.v);
        let v = self.v;
        self.u += 1;

        if v == 0 {
            Some(PolyTri(Triangle::new(x, y, z)))
        } else if v == self.sub_v - 1 {
            Some(PolyTri(Triangle::new(z, w, x)))
        } else {
            Some(PolyQuad(Quad::new(x, y, z, w)))
        }
    }
}

impl SharedVertex<(f32, f32, f32)> for SphereUV {
    fn shared_vertex(&self, idx: usize) -> (f32, f32, f32) {
        if idx == 0 {
            self.vert(0, 0)
        } else if idx == self.shared_vertex_count() - 1 {
            self.vert(0, self.sub_v)
        } else {
            // since the bottom verts all map to the same
            // we jump over them in index space
            let idx = idx - 1;
            let u = idx % (self.sub_u);
            let v = idx / (self.sub_u);
            self.vert(u, v+1)
        }
    }

    fn shared_vertex_count(&self) -> usize {
        (self.sub_v - 1) * (self.sub_u) + 2
    }
}

impl IndexedPolygon<Polygon<usize>> for SphereUV {
    fn indexed_polygon(&self, idx: usize) -> Polygon<usize> {
        let u = idx % self.sub_u;
        let v = idx / self.sub_u;

        let f = |u: usize, v: usize| {
            if v == 0 {
                0
            } else if self.sub_v == v {
                (self.sub_v-1) * (self.sub_u) + 1
            } else {
                (v-1) * self.sub_u + (u % self.sub_u) + 1
            }
        };

        if v == 0 {
            PolyTri(Triangle::new(f(u,   v),
                                  f(u,   v+1),
                                  f(u+1, v+1)))
        } else if self.sub_v - 1 == v {
            PolyTri(Triangle::new(f(u+1, v+1),
                                  f(u+1, v),
                                  f(u,   v)))
        } else {
            PolyQuad(Quad::new(f(u,   v),
                               f(u,   v+1),
                               f(u+1, v+1),
                               f(u+1, v)))
        }
    }

    fn indexed_polygon_count(&self) -> usize {
        self.sub_v * self.sub_u
    }
}

<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// Format string literals.

use regex::Regex;
use unicode_segmentation::UnicodeSegmentation;

use config::Config;
use shape::Shape;
use utils::wrap_str;

const MIN_STRING: usize = 10;

/// Describes the layout of a piece of text.
pub struct StringFormat<'a> {
    /// The opening sequence of characters for the piece of text
    pub opener: &'a str,
    /// The closing sequence of characters for the piece of text
    pub closer: &'a str,
    /// The opening sequence of characters for a line
    pub line_start: &'a str,
    /// The closing sequence of characters for a line
    pub line_end: &'a str,
    /// The allocated box to fit the text into
    pub shape: Shape,
    /// Trim trailing whitespaces
    pub trim_end: bool,
    pub config: &'a Config,
}

impl<'a> StringFormat<'a> {
    pub fn new(shape: Shape, config: &'a Config) -> StringFormat<'a> {
        StringFormat {
            opener: "\"",
            closer: "\"",
            line_start: " ",
            line_end: "\\",
            shape,
            trim_end: false,
            config,
        }
    }

    /// Returns the maximum number of graphemes that is possible on a line while taking the
    /// indentation into account.
    ///
    /// If we cannot put at least a single character per line, the rewrite won't succeed.
    fn max_chars_with_indent(&self) -> Option<usize> {
        Some(
            self.shape
                .width
                .checked_sub(self.opener.len() + self.line_end.len() + 1)?
                + 1,
        )
    }

    /// Like max_chars_with_indent but the indentation is not substracted.
    /// This allows to fit more graphemes from the string on a line when
    /// SnippetState::Overflow.
    fn max_chars_without_indent(&self) -> Option<usize> {
        Some(self.config.max_width().checked_sub(self.line_end.len())?)
    }
}

pub fn rewrite_string<'a>(orig: &str, fmt: &StringFormat<'a>) -> Option<String> {
    let max_chars_with_indent = fmt.max_chars_with_indent()?;
    let max_chars_without_indent = fmt.max_chars_without_indent()?;
    let indent = fmt.shape.indent.to_string_with_newline(fmt.config);

    // Strip line breaks.
    // With this regex applied, all remaining whitespaces are significant
    let strip_line_breaks_re = Regex::new(r"([^\\](\\\\)*)\\[\n\r][[:space:]]*").unwrap();
    let stripped_str = strip_line_breaks_re.replace_all(orig, "$1");

    let graphemes = UnicodeSegmentation::graphemes(&*stripped_str, false).collect::<Vec<&str>>();

    // `cur_start` is the position in `orig` of the start of the current line.
    let mut cur_start = 0;
    let mut result = String::with_capacity(
        stripped_str
            .len()
            .checked_next_power_of_two()
            .unwrap_or(usize::max_value()),
    );
    result.push_str(fmt.opener);

    // Snip a line at a time from `stripped_str` until it is used up. Push the snippet
    // onto result.
    let mut cur_max_chars = max_chars_with_indent;
    loop {
        // All the input starting at cur_start fits on the current line
        if graphemes.len() - cur_start <= cur_max_chars {
            result.push_str(&graphemes[cur_start..].join(""));
            break;
        }

        // The input starting at cur_start needs to be broken
        match break_string(cur_max_chars, fmt.trim_end, &graphemes[cur_start..]) {
            SnippetState::LineEnd(line, len) => {
                result.push_str(&line);
                result.push_str(fmt.line_end);
                result.push_str(&indent);
                result.push_str(fmt.line_start);
                cur_max_chars = max_chars_with_indent;
                cur_start += len;
            }
            SnippetState::Overflow(line, len) => {
                result.push_str(&line);
                cur_max_chars = max_chars_without_indent;
                cur_start += len;
            }
            SnippetState::EndOfInput(line) => {
                result.push_str(&line);
                break;
            }
        }
    }

    result.push_str(fmt.closer);
    wrap_str(result, fmt.config.max_width(), fmt.shape)
}

/// Result of breaking a string so it fits in a line and the state it ended in.
/// The state informs about what to do with the snippet and how to continue the breaking process.
#[derive(Debug, PartialEq)]
enum SnippetState {
    /// The input could not be broken and so rewriting the string is finished.
    EndOfInput(String),
    /// The input could be broken and the returned snippet should be ended with a
    /// `[StringFormat::line_end]`. The next snippet needs to be indented.
    /// The returned string is the line to print out and the number is the length that got read in
    /// the text being rewritten. That length may be greater than the returned string if trailing
    /// whitespaces got trimmed.
    LineEnd(String, usize),
    /// The input could be broken but the returned snippet should not be ended with a
    /// `[StringFormat::line_end]` because the whitespace is significant. Therefore, the next
    /// snippet should not be indented.
    Overflow(String, usize),
}

/// Break the input string at a boundary character around the offset `max_chars`. A boundary
/// character is either a punctuation or a whitespace.
fn break_string(max_chars: usize, trim_end: bool, input: &[&str]) -> SnippetState {
    let break_at = |index /* grapheme at index is included */| {
        // Take in any whitespaces to the left/right of `input[index]` and
        // check if there is a line feed, in which case whitespaces needs to be kept.
        let mut index_minus_ws = index;
        for (i, grapheme) in input[0..=index].iter().enumerate().rev() {
            if !is_whitespace(grapheme) {
                index_minus_ws = i;
                break;
            }
        }
        // Take into account newlines occuring in input[0..=index], i.e., the possible next new
        // line. If there is one, then text after it could be rewritten in a way that the available
        // space is fully used.
        for (i, grapheme) in input[0..=index].iter().enumerate() {
            if is_line_feed(grapheme) {
                if i < index_minus_ws || !trim_end {
                    return SnippetState::Overflow(input[0..=i].join("").to_string(), i + 1);
                }
                break;
            }
        }

        let mut index_plus_ws = index;
        for (i, grapheme) in input[index + 1..].iter().enumerate() {
            if !trim_end && is_line_feed(grapheme) {
                return SnippetState::Overflow(
                    input[0..=index + 1 + i].join("").to_string(),
                    index + 2 + i,
                );
            } else if !is_whitespace(grapheme) {
                index_plus_ws = index + i;
                break;
            }
        }

        if trim_end {
            SnippetState::LineEnd(
                input[0..=index_minus_ws].join("").to_string(),
                index_plus_ws + 1,
            )
        } else {
            SnippetState::LineEnd(
                input[0..=index_plus_ws].join("").to_string(),
                index_plus_ws + 1,
            )
        }
    };

    // Find the position in input for breaking the string
    match input[0..max_chars]
        .iter()
        .rposition(|grapheme| is_whitespace(grapheme))
    {
        // Found a whitespace and what is on its left side is big enough.
        Some(index) if index >= MIN_STRING => break_at(index),
        // No whitespace found, try looking for a punctuation instead
        _ => match input[0..max_chars]
            .iter()
            .rposition(|grapheme| is_punctuation(grapheme))
        {
            // Found a punctuation and what is on its left side is big enough.
            Some(index) if index >= MIN_STRING => break_at(index),
            // Either no boundary character was found to the left of `input[max_chars]`, or the line
            // got too small. We try searching for a boundary character to the right.
            _ => match input[max_chars..]
                .iter()
                .position(|grapheme| is_whitespace(grapheme) || is_punctuation(grapheme))
            {
                // A boundary was found after the line limit
                Some(index) => break_at(max_chars + index),
                // No boundary to the right, the input cannot be broken
                None => SnippetState::EndOfInput(input.join("").to_string()),
            },
        },
    }
}

fn is_line_feed(grapheme: &str) -> bool {
    grapheme.as_bytes()[0] == b'\n'
}

fn is_whitespace(grapheme: &str) -> bool {
    grapheme.chars().all(|c| c.is_whitespace())
}

fn is_punctuation(grapheme: &str) -> bool {
    match grapheme.as_bytes()[0] {
        b':' | b',' | b';' | b'.' => true,
        _ => false,
    }
}

#[cfg(test)]
mod test {
    use super::{break_string, rewrite_string, SnippetState, StringFormat};
    use config::Config;
    use shape::{Indent, Shape};
    use unicode_segmentation::UnicodeSegmentation;

    #[test]
    fn issue343() {
        let config = Default::default();
        let fmt = StringFormat::new(Shape::legacy(2, Indent::empty()), &config);
        rewrite_string("eq_", &fmt);
    }

    #[test]
    fn should_break_on_whitespace() {
        let string = "Placerat felis. Mauris porta ante sagittis purus.";
        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(20, false, &graphemes[..]),
            SnippetState::LineEnd("Placerat felis. ".to_string(), 16)
        );
        assert_eq!(
            break_string(20, true, &graphemes[..]),
            SnippetState::LineEnd("Placerat felis.".to_string(), 16)
        );
    }

    #[test]
    fn should_break_on_punctuation() {
        let string = "Placerat_felis._Mauris_porta_ante_sagittis_purus.";
        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(20, false, &graphemes[..]),
            SnippetState::LineEnd("Placerat_felis.".to_string(), 15)
        );
    }

    #[test]
    fn should_break_forward() {
        let string = "Venenatis_tellus_vel_tellus. Aliquam aliquam dolor at justo.";
        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(20, false, &graphemes[..]),
            SnippetState::LineEnd("Venenatis_tellus_vel_tellus. ".to_string(), 29)
        );
        assert_eq!(
            break_string(20, true, &graphemes[..]),
            SnippetState::LineEnd("Venenatis_tellus_vel_tellus.".to_string(), 29)
        );
    }

    #[test]
    fn nothing_to_break() {
        let string = "Venenatis_tellus_vel_tellus";
        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(20, false, &graphemes[..]),
            SnippetState::EndOfInput("Venenatis_tellus_vel_tellus".to_string())
        );
    }

    #[test]
    fn significant_whitespaces() {
        let string = "Neque in sem.      \n      Pellentesque tellus augue.";
        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(15, false, &graphemes[..]),
            SnippetState::Overflow("Neque in sem.      \n".to_string(), 20)
        );
        assert_eq!(
            break_string(25, false, &graphemes[..]),
            SnippetState::Overflow("Neque in sem.      \n".to_string(), 20)
        );
        // if `StringFormat::line_end` is true, then the line feed does not matter anymore
        assert_eq!(
            break_string(15, true, &graphemes[..]),
            SnippetState::LineEnd("Neque in sem.".to_string(), 26)
        );
        assert_eq!(
            break_string(25, true, &graphemes[..]),
            SnippetState::LineEnd("Neque in sem.".to_string(), 26)
        );
    }

    #[test]
    fn big_whitespace() {
        let string = "Neque in sem.            Pellentesque tellus augue.";
        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(20, false, &graphemes[..]),
            SnippetState::LineEnd("Neque in sem.            ".to_string(), 25)
        );
        assert_eq!(
            break_string(20, true, &graphemes[..]),
            SnippetState::LineEnd("Neque in sem.".to_string(), 25)
        );
    }

    #[test]
    fn newline_in_candidate_line() {
        let string = "Nulla\nconsequat erat at massa. Vivamus id mi.";

        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(25, false, &graphemes[..]),
            SnippetState::Overflow("Nulla\n".to_string(), 6)
        );
        assert_eq!(
            break_string(25, true, &graphemes[..]),
            SnippetState::Overflow("Nulla\n".to_string(), 6)
        );

        let mut config: Config = Default::default();
        config.set().max_width(27);
        let fmt = StringFormat::new(Shape::legacy(25, Indent::empty()), &config);
        let rewritten_string = rewrite_string(string, &fmt);
        assert_eq!(
            rewritten_string,
            Some("\"Nulla\nconsequat erat at massa. \\\n Vivamus id mi.\"".to_string())
        );
    }
}
<commit_msg>rewrite_string: trim trailing whitespace in case the last_line fits in the given shape<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// Format string literals.

use regex::Regex;
use unicode_segmentation::UnicodeSegmentation;

use config::Config;
use shape::Shape;
use utils::wrap_str;

const MIN_STRING: usize = 10;

/// Describes the layout of a piece of text.
pub struct StringFormat<'a> {
    /// The opening sequence of characters for the piece of text
    pub opener: &'a str,
    /// The closing sequence of characters for the piece of text
    pub closer: &'a str,
    /// The opening sequence of characters for a line
    pub line_start: &'a str,
    /// The closing sequence of characters for a line
    pub line_end: &'a str,
    /// The allocated box to fit the text into
    pub shape: Shape,
    /// Trim trailing whitespaces
    pub trim_end: bool,
    pub config: &'a Config,
}

impl<'a> StringFormat<'a> {
    pub fn new(shape: Shape, config: &'a Config) -> StringFormat<'a> {
        StringFormat {
            opener: "\"",
            closer: "\"",
            line_start: " ",
            line_end: "\\",
            shape,
            trim_end: false,
            config,
        }
    }

    /// Returns the maximum number of graphemes that is possible on a line while taking the
    /// indentation into account.
    ///
    /// If we cannot put at least a single character per line, the rewrite won't succeed.
    fn max_chars_with_indent(&self) -> Option<usize> {
        Some(
            self.shape
                .width
                .checked_sub(self.opener.len() + self.line_end.len() + 1)?
                + 1,
        )
    }

    /// Like max_chars_with_indent but the indentation is not substracted.
    /// This allows to fit more graphemes from the string on a line when
    /// SnippetState::Overflow.
    fn max_chars_without_indent(&self) -> Option<usize> {
        Some(self.config.max_width().checked_sub(self.line_end.len())?)
    }
}

pub fn rewrite_string<'a>(orig: &str, fmt: &StringFormat<'a>) -> Option<String> {
    let max_chars_with_indent = fmt.max_chars_with_indent()?;
    let max_chars_without_indent = fmt.max_chars_without_indent()?;
    let indent = fmt.shape.indent.to_string_with_newline(fmt.config);

    // Strip line breaks.
    // With this regex applied, all remaining whitespaces are significant
    let strip_line_breaks_re = Regex::new(r"([^\\](\\\\)*)\\[\n\r][[:space:]]*").unwrap();
    let stripped_str = strip_line_breaks_re.replace_all(orig, "$1");

    let graphemes = UnicodeSegmentation::graphemes(&*stripped_str, false).collect::<Vec<&str>>();

    // `cur_start` is the position in `orig` of the start of the current line.
    let mut cur_start = 0;
    let mut result = String::with_capacity(
        stripped_str
            .len()
            .checked_next_power_of_two()
            .unwrap_or(usize::max_value()),
    );
    result.push_str(fmt.opener);

    // Snip a line at a time from `stripped_str` until it is used up. Push the snippet
    // onto result.
    let mut cur_max_chars = max_chars_with_indent;
    loop {
        // All the input starting at cur_start fits on the current line
        if graphemes.len() - cur_start <= cur_max_chars {
            let last_line = graphemes[cur_start..].join("");
            if fmt.trim_end {
                result.push_str(&last_line.trim_right());
            } else {
                result.push_str(&last_line);
            }
            break;
        }

        // The input starting at cur_start needs to be broken
        match break_string(cur_max_chars, fmt.trim_end, &graphemes[cur_start..]) {
            SnippetState::LineEnd(line, len) => {
                result.push_str(&line);
                result.push_str(fmt.line_end);
                result.push_str(&indent);
                result.push_str(fmt.line_start);
                cur_max_chars = max_chars_with_indent;
                cur_start += len;
            }
            SnippetState::Overflow(line, len) => {
                result.push_str(&line);
                cur_max_chars = max_chars_without_indent;
                cur_start += len;
            }
            SnippetState::EndOfInput(line) => {
                result.push_str(&line);
                break;
            }
        }
    }

    result.push_str(fmt.closer);
    wrap_str(result, fmt.config.max_width(), fmt.shape)
}

/// Result of breaking a string so it fits in a line and the state it ended in.
/// The state informs about what to do with the snippet and how to continue the breaking process.
#[derive(Debug, PartialEq)]
enum SnippetState {
    /// The input could not be broken and so rewriting the string is finished.
    EndOfInput(String),
    /// The input could be broken and the returned snippet should be ended with a
    /// `[StringFormat::line_end]`. The next snippet needs to be indented.
    /// The returned string is the line to print out and the number is the length that got read in
    /// the text being rewritten. That length may be greater than the returned string if trailing
    /// whitespaces got trimmed.
    LineEnd(String, usize),
    /// The input could be broken but the returned snippet should not be ended with a
    /// `[StringFormat::line_end]` because the whitespace is significant. Therefore, the next
    /// snippet should not be indented.
    Overflow(String, usize),
}

/// Break the input string at a boundary character around the offset `max_chars`. A boundary
/// character is either a punctuation or a whitespace.
fn break_string(max_chars: usize, trim_end: bool, input: &[&str]) -> SnippetState {
    let break_at = |index /* grapheme at index is included */| {
        // Take in any whitespaces to the left/right of `input[index]` and
        // check if there is a line feed, in which case whitespaces needs to be kept.
        let mut index_minus_ws = index;
        for (i, grapheme) in input[0..=index].iter().enumerate().rev() {
            if !is_whitespace(grapheme) {
                index_minus_ws = i;
                break;
            }
        }
        // Take into account newlines occuring in input[0..=index], i.e., the possible next new
        // line. If there is one, then text after it could be rewritten in a way that the available
        // space is fully used.
        for (i, grapheme) in input[0..=index].iter().enumerate() {
            if is_line_feed(grapheme) {
                if i < index_minus_ws || !trim_end {
                    return SnippetState::Overflow(input[0..=i].join("").to_string(), i + 1);
                }
                break;
            }
        }

        let mut index_plus_ws = index;
        for (i, grapheme) in input[index + 1..].iter().enumerate() {
            if !trim_end && is_line_feed(grapheme) {
                return SnippetState::Overflow(
                    input[0..=index + 1 + i].join("").to_string(),
                    index + 2 + i,
                );
            } else if !is_whitespace(grapheme) {
                index_plus_ws = index + i;
                break;
            }
        }

        if trim_end {
            SnippetState::LineEnd(
                input[0..=index_minus_ws].join("").to_string(),
                index_plus_ws + 1,
            )
        } else {
            SnippetState::LineEnd(
                input[0..=index_plus_ws].join("").to_string(),
                index_plus_ws + 1,
            )
        }
    };

    // Find the position in input for breaking the string
    match input[0..max_chars]
        .iter()
        .rposition(|grapheme| is_whitespace(grapheme))
    {
        // Found a whitespace and what is on its left side is big enough.
        Some(index) if index >= MIN_STRING => break_at(index),
        // No whitespace found, try looking for a punctuation instead
        _ => match input[0..max_chars]
            .iter()
            .rposition(|grapheme| is_punctuation(grapheme))
        {
            // Found a punctuation and what is on its left side is big enough.
            Some(index) if index >= MIN_STRING => break_at(index),
            // Either no boundary character was found to the left of `input[max_chars]`, or the line
            // got too small. We try searching for a boundary character to the right.
            _ => match input[max_chars..]
                .iter()
                .position(|grapheme| is_whitespace(grapheme) || is_punctuation(grapheme))
            {
                // A boundary was found after the line limit
                Some(index) => break_at(max_chars + index),
                // No boundary to the right, the input cannot be broken
                None => SnippetState::EndOfInput(input.join("").to_string()),
            },
        },
    }
}

fn is_line_feed(grapheme: &str) -> bool {
    grapheme.as_bytes()[0] == b'\n'
}

fn is_whitespace(grapheme: &str) -> bool {
    grapheme.chars().all(|c| c.is_whitespace())
}

fn is_punctuation(grapheme: &str) -> bool {
    match grapheme.as_bytes()[0] {
        b':' | b',' | b';' | b'.' => true,
        _ => false,
    }
}

#[cfg(test)]
mod test {
    use super::{break_string, rewrite_string, SnippetState, StringFormat};
    use config::Config;
    use shape::{Indent, Shape};
    use unicode_segmentation::UnicodeSegmentation;

    #[test]
    fn issue343() {
        let config = Default::default();
        let fmt = StringFormat::new(Shape::legacy(2, Indent::empty()), &config);
        rewrite_string("eq_", &fmt);
    }

    #[test]
    fn should_break_on_whitespace() {
        let string = "Placerat felis. Mauris porta ante sagittis purus.";
        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(20, false, &graphemes[..]),
            SnippetState::LineEnd("Placerat felis. ".to_string(), 16)
        );
        assert_eq!(
            break_string(20, true, &graphemes[..]),
            SnippetState::LineEnd("Placerat felis.".to_string(), 16)
        );
    }

    #[test]
    fn should_break_on_punctuation() {
        let string = "Placerat_felis._Mauris_porta_ante_sagittis_purus.";
        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(20, false, &graphemes[..]),
            SnippetState::LineEnd("Placerat_felis.".to_string(), 15)
        );
    }

    #[test]
    fn should_break_forward() {
        let string = "Venenatis_tellus_vel_tellus. Aliquam aliquam dolor at justo.";
        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(20, false, &graphemes[..]),
            SnippetState::LineEnd("Venenatis_tellus_vel_tellus. ".to_string(), 29)
        );
        assert_eq!(
            break_string(20, true, &graphemes[..]),
            SnippetState::LineEnd("Venenatis_tellus_vel_tellus.".to_string(), 29)
        );
    }

    #[test]
    fn nothing_to_break() {
        let string = "Venenatis_tellus_vel_tellus";
        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(20, false, &graphemes[..]),
            SnippetState::EndOfInput("Venenatis_tellus_vel_tellus".to_string())
        );
    }

    #[test]
    fn significant_whitespaces() {
        let string = "Neque in sem.      \n      Pellentesque tellus augue.";
        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(15, false, &graphemes[..]),
            SnippetState::Overflow("Neque in sem.      \n".to_string(), 20)
        );
        assert_eq!(
            break_string(25, false, &graphemes[..]),
            SnippetState::Overflow("Neque in sem.      \n".to_string(), 20)
        );
        // if `StringFormat::line_end` is true, then the line feed does not matter anymore
        assert_eq!(
            break_string(15, true, &graphemes[..]),
            SnippetState::LineEnd("Neque in sem.".to_string(), 26)
        );
        assert_eq!(
            break_string(25, true, &graphemes[..]),
            SnippetState::LineEnd("Neque in sem.".to_string(), 26)
        );
    }

    #[test]
    fn big_whitespace() {
        let string = "Neque in sem.            Pellentesque tellus augue.";
        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(20, false, &graphemes[..]),
            SnippetState::LineEnd("Neque in sem.            ".to_string(), 25)
        );
        assert_eq!(
            break_string(20, true, &graphemes[..]),
            SnippetState::LineEnd("Neque in sem.".to_string(), 25)
        );
    }

    #[test]
    fn newline_in_candidate_line() {
        let string = "Nulla\nconsequat erat at massa. Vivamus id mi.";

        let graphemes = UnicodeSegmentation::graphemes(&*string, false).collect::<Vec<&str>>();
        assert_eq!(
            break_string(25, false, &graphemes[..]),
            SnippetState::Overflow("Nulla\n".to_string(), 6)
        );
        assert_eq!(
            break_string(25, true, &graphemes[..]),
            SnippetState::Overflow("Nulla\n".to_string(), 6)
        );

        let mut config: Config = Default::default();
        config.set().max_width(27);
        let fmt = StringFormat::new(Shape::legacy(25, Indent::empty()), &config);
        let rewritten_string = rewrite_string(string, &fmt);
        assert_eq!(
            rewritten_string,
            Some("\"Nulla\nconsequat erat at massa. \\\n Vivamus id mi.\"".to_string())
        );
    }

    #[test]
    fn last_line_fit_with_trailing_whitespaces() {
        let string = "Vivamus id mi.  ";
        let config: Config = Default::default();
        let mut fmt = StringFormat::new(Shape::legacy(25, Indent::empty()), &config);

        fmt.trim_end = true;
        let rewritten_string = rewrite_string(string, &fmt);
        assert_eq!(rewritten_string, Some("\"Vivamus id mi.\"".to_string()));

        fmt.trim_end = false; // default value of trim_end
        let rewritten_string = rewrite_string(string, &fmt);
        assert_eq!(rewritten_string, Some("\"Vivamus id mi.  \"".to_string()));
    }
}
<|endoftext|>


<commit_before>use serde_json;

use op::Op;

pub struct Update {
    rev: u64,
    ops: Vec<Op>,
}

impl Update {
    pub fn from_value(value: &serde_json::Value) -> Update {
        let object = value.as_object().unwrap();
        let ops = object.get("ops").unwrap().as_array().unwrap().iter().map(
            |op| Op::from_value(op)
        ).collect();
        Update {
            rev: object.get("rev").unwrap().as_u64().unwrap(),
            ops: ops,
        }
    }
}
<commit_msg>Makes all fields of Update struct public<commit_after>use serde_json;

use op::Op;

pub struct Update {
    pub rev: u64,
    pub ops: Vec<Op>,
}

impl Update {
    pub fn from_value(value: &serde_json::Value) -> Update {
        let object = value.as_object().unwrap();
        let ops = object.get("ops").unwrap().as_array().unwrap().iter().map(
            |op| Op::from_value(op)
        ).collect();
        Update {
            rev: object.get("rev").unwrap().as_u64().unwrap(),
            ops: ops,
        }
    }
}
<|endoftext|>


<commit_before>//! Utility functions to write nifti images.

use std::fs::File;
use std::io::BufWriter;
use std::ops::{Div, Sub};
use std::path::Path;

use byteorder::{LittleEndian, WriteBytesExt};
use flate2::write::GzEncoder;
use flate2::Compression;
use ndarray::{ArrayBase, ArrayView, Axis, Data, Dimension, RemoveAxis, ScalarOperand};
use num_traits::FromPrimitive;
use safe_transmute::{guarded_transmute_to_bytes_pod_many, PodTransmutable};

use {
    header::MAGIC_CODE_NIP1, util::is_gz_file, volume::element::DataElement, NiftiHeader,
    NiftiType, Result,
};

// TODO make this configurable. The Nifti standard does not specify a specific field for endianness,
// but it is encoded in `dim[0]`. "if dim[0] is outside range 1..7, then swap".
type B = LittleEndian;

/// Write a nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields. The voxels
/// intensity will be subtracted by `scl_slope` and divided by `scl_inter`. If `reference` is not
/// given, a default `NiftiHeader` will be built and written.
///
/// In all cases, the `dim`, `datatype` and `bitpix` fields will depend only on `data`, not on the
/// header. In other words, the `datatype` defined in `reference` will be ignored.
pub fn write_nifti<P, A, S, D>(
    path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = A>,
    A: Copy,
    A: DataElement,
    A: Div<Output = A>,
    A: FromPrimitive,
    A: PodTransmutable,
    A: ScalarOperand,
    A: Sub<Output = A>,
    D: Dimension + RemoveAxis,
{
    let header = build_header(data, reference, A::DATA_TYPE);

    // Need the transpose for fortran ordering used in nifti file format.
    let data = data.t();

    let f = File::create(&path)?;
    let mut writer = BufWriter::new(f);
    if is_gz_file(&path) {
        let mut e = GzEncoder::new(writer, Compression::fast());
        write_header(&mut e, &header)?;
        write_data(&mut e, &header, data)?;
        let _ = e.finish()?; // Must use result
    } else {
        write_header(&mut writer, &header)?;
        write_data(&mut writer, &header, data)?;
    }
    Ok(())
}

/// Write a RGB nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields, except those
/// necessary to be recognized as a RGB image. `scl_slope` will be set to 1.0 and `scl_inter` to
/// 0.0.  If `reference` is not given, a default `NiftiHeader` will be built and written.
pub fn write_rgb_nifti<P, S, D>(
    path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = [u8; 3]>,
    D: Dimension + RemoveAxis,
{
    // The `scl_slope` and `scl_inter` fields are ignored on the Rgb24 type.
    let mut header = build_header(data, reference, NiftiType::Rgb24);
    header.scl_slope = 1.0;
    header.scl_inter = 0.0;

    // Need the transpose for fortran used in nifti file format.
    let data = data.t();

    let f = File::create(&path)?;
    let mut writer = BufWriter::new(f);
    if is_gz_file(&path) {
        let mut e = GzEncoder::new(writer, Compression::fast());
        write_header(&mut e, &header)?;
        write_slices(&mut e, data)?;
        let _ = e.finish()?; // Must use result
    } else {
        write_header(&mut writer, &header)?;
        write_slices(&mut writer, data)?;
    }
    Ok(())
}

fn build_header<T, D>(
    data: &ArrayBase<T, D>,
    reference: Option<&NiftiHeader>,
    datatype: NiftiType,
) -> NiftiHeader
where
    T: Data,
    D: Dimension,
{
    let mut dim = [1; 8];
    dim[0] = data.ndim() as u16;
    for (i, s) in data.shape().iter().enumerate() {
        dim[i + 1] = *s as u16;
    }

    // If no reference header is given, use the default.
    let reference = match reference {
        Some(r) => r.clone(),
        None => {
            let mut header = NiftiHeader::default();
            header.pixdim = [1.0; 8];
            header.sform_code = 2;
            header.srow_x = [1.0, 0.0, 0.0, 0.0];
            header.srow_y = [0.0, 1.0, 0.0, 0.0];
            header.srow_z = [0.0, 0.0, 1.0, 0.0];
            header
        }
    };

    NiftiHeader {
        dim,
        sizeof_hdr: 348,
        datatype: datatype as i16,
        bitpix: (datatype.size_of() * 8) as i16,
        vox_offset: 352.0,
        magic: *MAGIC_CODE_NIP1,
        // All other fields are copied from reference header
        ..reference
    }
}

fn write_header<W>(writer: &mut W, header: &NiftiHeader) -> Result<()>
where
    W: WriteBytesExt,
{
    writer.write_i32::<B>(header.sizeof_hdr)?;
    writer.write_all(&header.data_type)?;
    writer.write_all(&header.db_name)?;
    writer.write_i32::<B>(header.extents)?;
    writer.write_i16::<B>(header.session_error)?;
    writer.write_u8(header.regular)?;
    writer.write_u8(header.dim_info)?;
    for s in &header.dim {
        writer.write_u16::<B>(*s)?;
    }
    writer.write_f32::<B>(header.intent_p1)?;
    writer.write_f32::<B>(header.intent_p2)?;
    writer.write_f32::<B>(header.intent_p3)?;
    writer.write_i16::<B>(header.intent_code)?;
    writer.write_i16::<B>(header.datatype)?;
    writer.write_i16::<B>(header.bitpix)?;
    writer.write_i16::<B>(header.slice_start)?;
    for f in &header.pixdim {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_f32::<B>(header.vox_offset)?;
    writer.write_f32::<B>(header.scl_slope)?;
    writer.write_f32::<B>(header.scl_inter)?;
    writer.write_i16::<B>(header.slice_end)?;
    writer.write_u8(header.slice_code)?;
    writer.write_u8(header.xyzt_units)?;
    writer.write_f32::<B>(header.cal_max)?;
    writer.write_f32::<B>(header.cal_min)?;
    writer.write_f32::<B>(header.slice_duration)?;
    writer.write_f32::<B>(header.toffset)?;
    writer.write_i32::<B>(header.glmax)?;
    writer.write_i32::<B>(header.glmin)?;

    writer.write_all(&header.descrip)?;
    writer.write_all(&header.aux_file)?;
    writer.write_i16::<B>(header.qform_code)?;
    writer.write_i16::<B>(header.sform_code)?;
    for f in &[
        header.quatern_b,
        header.quatern_c,
        header.quatern_d,
        header.quatern_x,
        header.quatern_y,
        header.quatern_z,
    ] {
        writer.write_f32::<B>(*f)?;
    }
    for f in header
        .srow_x
        .iter()
        .chain(&header.srow_y)
        .chain(&header.srow_z)
    {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_all(&header.intent_name)?;
    writer.write_all(&header.magic)?;

    // Empty 4 bytes after the header
    // TODO Support writing extension data.
    writer.write_u32::<B>(0)?;

    Ok(())
}

/// Write the data in 'f' order.
///
/// Like NiBabel, we iterate by "slice" to improve speed and use less memory.
fn write_data<T, D, W>(writer: &mut W, header: &NiftiHeader, data: ArrayView<T, D>) -> Result<()>
where
    T: Clone + PodTransmutable,
    T: Div<Output = T>,
    T: FromPrimitive,
    T: PodTransmutable,
    T: ScalarOperand,
    T: Sub<Output = T>,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    // `1.0x + 0.0` would give the same results, but we avoid a lot of divisions
    let slope = if header.scl_slope == 0.0 {
        1.0
    } else {
        header.scl_slope
    };
    if slope != 1.0 || header.scl_inter != 0.0 {
        // TODO Use linear transformation like when reading. An scl_slope of 0.5 would turn all
        // voxel values to 0 if we pass an ndarray of integers.
        let slope = T::from_f32(slope).unwrap();
        let inter = T::from_f32(header.scl_inter).unwrap();
        for arr_data in data.axis_iter(Axis(0)) {
            write_slice(writer, arr_data.sub(inter).div(slope))?;
        }
    } else {
        write_slices(writer, data)?;
    }
    Ok(())
}

fn write_slices<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    let mut iter = data.axis_iter(Axis(0));
    if let Some(arr_data) = iter.next() {
        // Keep slice voxels in a separate array to ensure `C` ordering even after `into_shape`.
        let mut slice = arr_data.to_owned();
        write_slice(writer, slice.view())?;
        for arr_data in iter {
            slice.assign(&arr_data);
            write_slice(writer, slice.view())?;
        }
    }
    Ok(())
}

fn write_slice<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension,
    W: WriteBytesExt,
{
    let len = data.len();
    let arr_data = data.into_shape(len).unwrap();
    let slice = arr_data.as_slice().unwrap();
    writer.write_all(guarded_transmute_to_bytes_pod_many(slice))?;
    Ok(())
}
<commit_msg>Support writing hdr files<commit_after>//! Utility functions to write nifti images.

use std::fs::File;
use std::io::BufWriter;
use std::ops::{Div, Sub};
use std::path::{Path, PathBuf};

use byteorder::{LittleEndian, WriteBytesExt};
use flate2::write::GzEncoder;
use flate2::Compression;
use ndarray::{ArrayBase, ArrayView, Axis, Data, Dimension, RemoveAxis, ScalarOperand};
use num_traits::FromPrimitive;
use safe_transmute::{guarded_transmute_to_bytes_pod_many, PodTransmutable};

use {
    header::{MAGIC_CODE_NIP1, MAGIC_CODE_NI1},
    util::{is_gz_file, is_hdr_file},
    volume::element::DataElement,
    NiftiHeader, NiftiType, Result,
};

// TODO make this configurable. The Nifti standard does not specify a specific field for endianness,
// but it is encoded in `dim[0]`. "if dim[0] is outside range 1..7, then swap".
type B = LittleEndian;

/// Write a nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields. The voxels
/// intensity will be subtracted by `scl_slope` and divided by `scl_inter`. If `reference` is not
/// given, a default `NiftiHeader` will be built and written.
///
/// In all cases, the `dim`, `datatype` and `bitpix` fields will depend only on `data`, not on the
/// header. In other words, the `datatype` defined in `reference` will be ignored.
pub fn write_nifti<P, A, S, D>(
    path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = A>,
    A: Copy,
    A: DataElement,
    A: Div<Output = A>,
    A: FromPrimitive,
    A: PodTransmutable,
    A: ScalarOperand,
    A: Sub<Output = A>,
    D: Dimension + RemoveAxis,
{
    let compression_level = Compression::fast();
    let is_gz = is_gz_file(&path);
    let (header, header_path, data_path) =
        prepare_header_and_paths(path, data, reference, A::DATA_TYPE);

    // Need the transpose for fortran ordering used in nifti file format.
    let data = data.t();

    let header_file = File::create(&header_path)?;
    let mut header_writer = BufWriter::new(header_file);
    if header.vox_offset > 0.0 {
        if is_gz {
            let mut e = GzEncoder::new(header_writer, compression_level);
            write_header(&mut e, &header)?;
            write_data(&mut e, &header, data)?;
            let _ = e.finish()?;
        } else {
            write_header(&mut header_writer, &header)?;
            write_data(&mut header_writer, &header, data)?;
        }
    } else {
        let data_file = File::create(&data_path)?;
        let mut data_writer = BufWriter::new(data_file);
        if is_gz {
            let mut e = GzEncoder::new(header_writer, compression_level);
            write_header(&mut e, &header)?;
            let _ = e.finish()?;

            let mut e = GzEncoder::new(data_writer, compression_level);
            write_data(&mut e, &header, data)?;
            let _ = e.finish()?;
        } else {
            write_header(&mut header_writer, &header)?;
            write_data(&mut data_writer, &header, data)?;
        }
    }

    Ok(())
}

/// Write a RGB nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields, except those
/// necessary to be recognized as a RGB image. `scl_slope` will be set to 1.0 and `scl_inter` to
/// 0.0.  If `reference` is not given, a default `NiftiHeader` will be built and written.
pub fn write_rgb_nifti<P, S, D>(
    path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = [u8; 3]>,
    D: Dimension + RemoveAxis,
{
    let is_gz = is_gz_file(&path);
    let (mut header, header_path, _) =
        prepare_header_and_paths(path, data, reference, NiftiType::Rgb24);

    // The `scl_slope` and `scl_inter` fields are ignored on the Rgb24 type.
    header.scl_slope = 1.0;
    header.scl_inter = 0.0;

    // Need the transpose for fortran used in nifti file format.
    let data = data.t();

    let f = File::create(&header_path)?;
    let mut writer = BufWriter::new(f);
    if is_gz {
        let mut e = GzEncoder::new(writer, Compression::fast());
        write_header(&mut e, &header)?;
        write_slices(&mut e, data)?;
        let _ = e.finish()?; // Must use result
    } else {
        write_header(&mut writer, &header)?;
        write_slices(&mut writer, data)?;
    }
    Ok(())
}

fn prepare_header_and_paths<P, T, D>(
    path: P,
    data: &ArrayBase<T, D>,
    reference: Option<&NiftiHeader>,
    datatype: NiftiType,
) -> (NiftiHeader, P, PathBuf)
where
    P: AsRef<Path>,
    T: Data,
    D: Dimension,
{
    let mut dim = [1; 8];
    dim[0] = data.ndim() as u16;
    for (i, s) in data.shape().iter().enumerate() {
        dim[i + 1] = *s as u16;
    }

    // If no reference header is given, use the default.
    let reference = match reference {
        Some(r) => r.clone(),
        None => {
            let mut header = NiftiHeader::default();
            header.pixdim = [1.0; 8];
            header.sform_code = 2;
            header.srow_x = [1.0, 0.0, 0.0, 0.0];
            header.srow_y = [0.0, 1.0, 0.0, 0.0];
            header.srow_z = [0.0, 0.0, 1.0, 0.0];
            header
        }
    };

    let mut header = NiftiHeader {
        dim,
        sizeof_hdr: 348,
        datatype: datatype as i16,
        bitpix: (datatype.size_of() * 8) as i16,
        vox_offset: 352.0,
        magic: *MAGIC_CODE_NIP1,
        // All other fields are copied from reference header
        ..reference
    };

    let mut path_buf = PathBuf::from(path.as_ref());
    let (header_path, data_path) = if is_hdr_file(&path) {
        header.vox_offset = 0.0;
        header.magic = *MAGIC_CODE_NI1;
        let data_path = if is_gz_file(&path) {
            let _ = path_buf.set_extension("");
            path_buf.with_extension("img.gz")
        } else {
            path_buf.with_extension("img")
        };
        (path, data_path)
    } else {
        (path, path_buf)
    };

    (header, header_path, data_path)
}

fn write_header<W>(writer: &mut W, header: &NiftiHeader) -> Result<()>
where
    W: WriteBytesExt,
{
    writer.write_i32::<B>(header.sizeof_hdr)?;
    writer.write_all(&header.data_type)?;
    writer.write_all(&header.db_name)?;
    writer.write_i32::<B>(header.extents)?;
    writer.write_i16::<B>(header.session_error)?;
    writer.write_u8(header.regular)?;
    writer.write_u8(header.dim_info)?;
    for s in &header.dim {
        writer.write_u16::<B>(*s)?;
    }
    writer.write_f32::<B>(header.intent_p1)?;
    writer.write_f32::<B>(header.intent_p2)?;
    writer.write_f32::<B>(header.intent_p3)?;
    writer.write_i16::<B>(header.intent_code)?;
    writer.write_i16::<B>(header.datatype)?;
    writer.write_i16::<B>(header.bitpix)?;
    writer.write_i16::<B>(header.slice_start)?;
    for f in &header.pixdim {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_f32::<B>(header.vox_offset)?;
    writer.write_f32::<B>(header.scl_slope)?;
    writer.write_f32::<B>(header.scl_inter)?;
    writer.write_i16::<B>(header.slice_end)?;
    writer.write_u8(header.slice_code)?;
    writer.write_u8(header.xyzt_units)?;
    writer.write_f32::<B>(header.cal_max)?;
    writer.write_f32::<B>(header.cal_min)?;
    writer.write_f32::<B>(header.slice_duration)?;
    writer.write_f32::<B>(header.toffset)?;
    writer.write_i32::<B>(header.glmax)?;
    writer.write_i32::<B>(header.glmin)?;

    writer.write_all(&header.descrip)?;
    writer.write_all(&header.aux_file)?;
    writer.write_i16::<B>(header.qform_code)?;
    writer.write_i16::<B>(header.sform_code)?;
    for f in &[
        header.quatern_b,
        header.quatern_c,
        header.quatern_d,
        header.quatern_x,
        header.quatern_y,
        header.quatern_z,
    ] {
        writer.write_f32::<B>(*f)?;
    }
    for f in header
        .srow_x
        .iter()
        .chain(&header.srow_y)
        .chain(&header.srow_z)
    {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_all(&header.intent_name)?;
    writer.write_all(&header.magic)?;

    // Empty 4 bytes after the header
    // TODO Support writing extension data.
    writer.write_u32::<B>(0)?;

    Ok(())
}

/// Write the data in 'f' order.
///
/// Like NiBabel, we iterate by "slice" to improve speed and use less memory.
fn write_data<T, D, W>(writer: &mut W, header: &NiftiHeader, data: ArrayView<T, D>) -> Result<()>
where
    T: Clone + PodTransmutable,
    T: Div<Output = T>,
    T: FromPrimitive,
    T: PodTransmutable,
    T: ScalarOperand,
    T: Sub<Output = T>,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    // `1.0x + 0.0` would give the same results, but we avoid a lot of divisions
    let slope = if header.scl_slope == 0.0 {
        1.0
    } else {
        header.scl_slope
    };
    if slope != 1.0 || header.scl_inter != 0.0 {
        // TODO Use linear transformation like when reading. An scl_slope of 0.5 would turn all
        // voxel values to 0 if we pass an ndarray of integers.
        let slope = T::from_f32(slope).unwrap();
        let inter = T::from_f32(header.scl_inter).unwrap();
        for arr_data in data.axis_iter(Axis(0)) {
            write_slice(writer, arr_data.sub(inter).div(slope))?;
        }
    } else {
        write_slices(writer, data)?;
    }
    Ok(())
}

fn write_slices<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    let mut iter = data.axis_iter(Axis(0));
    if let Some(arr_data) = iter.next() {
        // Keep slice voxels in a separate array to ensure `C` ordering even after `into_shape`.
        let mut slice = arr_data.to_owned();
        write_slice(writer, slice.view())?;
        for arr_data in iter {
            slice.assign(&arr_data);
            write_slice(writer, slice.view())?;
        }
    }
    Ok(())
}

fn write_slice<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension,
    W: WriteBytesExt,
{
    let len = data.len();
    let arr_data = data.into_shape(len).unwrap();
    let slice = arr_data.as_slice().unwrap();
    writer.write_all(guarded_transmute_to_bytes_pod_many(slice))?;
    Ok(())
}
<|endoftext|>


<commit_before>#![feature(core)]
#![feature(str_words)]

extern crate mmap;
extern crate time;
extern crate core;
extern crate timely;
extern crate columnar;
extern crate dataflow_join;

extern crate docopt;
use docopt::Docopt;

use std::io::{BufRead, stdin};
use std::rc::Rc;
use std::cell::RefCell;
use std::thread;

use core::cmp::Ordering::*;

use dataflow_join::*;
use dataflow_join::graph::{GraphTrait, GraphMMap, GraphExtenderExt, gallop};

use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::progress::nested::product::Product;
use timely::example_static::*;

use timely::communication::*;
use timely::communication::pact::Pipeline;

static USAGE: &'static str = "
Usage: triangles dataflow <source> <workers> [--inspect] [--interactive]
       triangles autorun <source> <workers>
       triangles compute <source>
       triangles help
";

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    if args.get_bool("dataflow") {
        let inspect = args.get_bool("--inspect");
        let interactive = args.get_bool("--interactive");
        let workers = if let Ok(threads) = args.get_str("<workers>").parse() { threads }
                      else { panic!("invalid setting for workers: {}", args.get_str("-t")) };;
        println!("starting triangles dataflow with {:?} worker{}; inspection: {:?}, interactive: {:?}",
                    workers, if workers == 1 { "" } else { "s" }, inspect, interactive);
        let source = args.get_str("<source>");

        triangles_multi(ProcessCommunicator::new_vector(workers), |_, _| GraphMMap::new(&source), inspect, interactive);
    }
    if args.get_bool("autorun") {
        let workers = if let Ok(threads) = args.get_str("<workers>").parse() { threads }
                      else { panic!("invalid setting for workers: {}", args.get_str("-t")) };;
        println!("starting triangles dataflow with {:?} worker{}; autorun",
                    workers, if workers == 1 { "" } else { "s" });
        let source = args.get_str("<source>");
        triangles_auto_multi(ProcessCommunicator::new_vector(workers), |_, _| GraphMMap::new(&source));
    }
    if args.get_bool("compute") {
        let source = args.get_str("<source>");
        let graph = GraphMMap::new(&source);
        println!("triangles: {:?}", raw_triangles(&graph));
    }

    if args.get_bool("help") {
        println!("the code presently assumes you have access to the livejournal graph, from:");
        println!("   https://snap.stanford.edu/data/soc-LiveJournal1.html");
        println!("");
        println!("before you can use it you will need to \"digest\" it using example digest.rs");
        println!("digest will overwrite <target>.targets and <target>.offsets, so be careful!");
        println!("at least, it will once you edit the code to uncomment the line.");
        println!("Once you grok the binary format, you can totally use other graphs too!");
    }
}

fn raw_triangles<G: GraphTrait<Target=u32>>(graph: &G) -> u64 {
    let mut count = 0;
    for a in (0..graph.nodes()) {
        let aaa = graph.edges(a);
        for &b in aaa {
            let bbb = graph.edges(b as usize);
            count += if aaa.len() < bbb.len() { intersect(aaa.clone(), bbb) }
                     else                     { intersect(bbb, aaa.clone()) };
        }
    }
    count
}

fn intersect<E: Ord>(mut aaa: &[E], mut bbb: &[E]) -> u64 {
    let mut count = 0;
    // magic gallop overhead # is 4
    if aaa.len() < bbb.len() / 4 {
        for a in aaa {
            bbb = gallop(bbb, a);
            if bbb.len() > 0 && &bbb[0] == a { count += 1; }
        }
    }
    else {
        while aaa.len() > 0 && bbb.len() > 0 {
            match aaa[0].cmp(&bbb[0]) {
                Greater => { bbb = &bbb[1..]; },
                Less    => { aaa = &aaa[1..]; },
                Equal   => { aaa = &aaa[1..];
                             bbb = &bbb[1..];
                             count += 1;
                           },
            }
        }
    }
    count
}

fn triangles_multi<C, G, F>(communicators: Vec<C>, loader: F, inspect: bool, interactive: bool)
where C: Communicator+Send, G: GraphTrait<Target=u32>, F: Fn(u64, u64)->G+Send+Sync {
    let mut guards = Vec::new();
    let loader = &loader;
    for communicator in communicators.into_iter() {
        guards.push(thread::Builder::new().name(format!("worker thread {}", communicator.index()))
                                          .scoped(move || triangles(communicator, loader, inspect, interactive))
                                          .unwrap());
    }
}

fn triangles<C, G, F>(communicator: C, loader: &F, inspect: bool, interactive: bool)
where C: Communicator, G: GraphTrait<Target=u32>, F: Fn(u64, u64)->G {
    let comm_index = communicator.index();
    let comm_peers = communicator.peers();

    let graph = Rc::new(RefCell::new(loader(comm_index, comm_peers)));

    let mut root = GraphRoot::new(communicator);
    let mut input = { // new scope to avoid long borrow on root

        let mut builder = root.new_subgraph();
        let (input, stream) = builder.new_input::<u32>();

        // extend u32s to pairs, then pairs to triples.
        let triangles = builder.enable(&stream)
                               .extend(vec![&graph.extend_using(|| { |&a| a as u64 } )])
                               .flat_map(|(p, es)| es.into_iter().map(move |e| (p, e)))
                               .extend(vec![&graph.extend_using(|| { |&(a,_)| a as u64 }),
                                            &graph.extend_using(|| { |&(_,b)| b as u64 })]);

        if inspect { triangles.inspect(|x| println!("triangles: {:?}", x)); }

        input
    };

    if interactive {
        let mut stdinput = stdin();
        let mut line = String::new();
        for index in (0..) {
            stdinput.read_line(&mut line).unwrap();
            let start = time::precise_time_ns();
            if let Some(word) = line.words().next() {
                let read = word.parse();
                if let Ok(number) = read {
                    input.send_messages(&Product::new((), index as u64), vec![number]);
                }
            }
            line.clear();

            input.advance(&Product::new((), index as u64), &Product::new((), index as u64 + 1));
            for _ in (0..5) { root.step(); }

            println!("elapsed: {:?}us", (time::precise_time_ns() - start)/1000);
        }
    }
    else {
        let step = 1000;
        let nodes = graph.borrow().nodes() - 1;
        let limit = (nodes / step) + 1;
        for index in (0..limit) {
            let index = index as usize;
            let data = ((index * step) .. ((index + 1) * step))
                           .filter(|&x| x as u64 % comm_peers == comm_index)
                           .filter(|&x| x < nodes)
                           .map(|x| x as u32)
                           .collect();
            input.send_messages(&Product::new((), index as u64), data);
            input.advance(&Product::new((), index as u64), &Product::new((), index as u64 + 1));
            root.step();
        }

        input.close_at(&Product::new((), limit as u64));
        while root.step() { }
    }
}

fn triangles_auto_multi<C, G, F>(communicators: Vec<C>, loader: F)
where C: Communicator+Send, G: GraphTrait<Target=u32>, F: Fn(u64, u64)->G+Send+Sync {
    let mut guards = Vec::new();
    let loader = &loader;
    for communicator in communicators.into_iter() {
        guards.push(thread::Builder::new().name(format!("worker thread {}", communicator.index()))
                                          .scoped(move || triangles_auto(communicator, loader))
                                          .unwrap());
    }
}

fn triangles_auto<C, G, F>(communicator: C, loader: &F)
where C: Communicator, G: GraphTrait<Target=u32>, F: Fn(u64, u64)->G {
    let index = communicator.index();
    let peers = communicator.peers();

    let batch = 100;
    let graph = Rc::new(RefCell::new(loader(index, peers)));
    let nodes = graph.borrow().nodes() as u64;

    let mut root = GraphRoot::new(communicator);

    {   // new scope to avoid long borrow on root

        let mut builder = root.new_subgraph::<u64>();

        let (feedback, stream) = builder.loop_variable(Product::new((), nodes + 1), Local(1));

        stream.enable(&mut builder)
              .unary_notify(Pipeline, format!("Input"), vec![Product::new((), 0)], move |handle| {
                  if let Some((time, _count)) = handle.notificator.next() {
                      if time.inner < (nodes / batch) {
                          handle.notificator.notify_at(&Product::new((), time.inner + 1));
                      }
                      let mut session = handle.output.session(&time);
                      for next in (0..(batch/peers)) {
                          session.give(time.inner * batch + next * peers + index);
                      }
                  }})
              .extend(vec![&graph.extend_using(|| { |&a| a as u64 } )])
              .flat_map(|(p, es)| es.into_iter().map(move |e| (p, e)))
              .extend(vec![&graph.extend_using(|| { |&(a,_)| a as u64 }),
                           &graph.extend_using(|| { |&(_,b)| b as u64 })])
              .filter(|_| false)
              .connect_loop(feedback);
    }

    while root.step() { }
}

// // fn _quads<'a, G, G2>(stream: &mut Stream<'a, G, ((u32, u32), u32)>, graph: &Rc<RefCell<G2>>) ->
// //                                                             Stream<'a, G, (((u32, u32), u32), u32)>
// // where G: GraphBuilder+'a, G2: GraphTrait<Target=u32> {
// //     //
// //     stream.extend(vec![&graph.extend_using(|| { |&((a,_),_)| a as u64 }),
// //                        &graph.extend_using(|| { |&((_,b),_)| b as u64 }),
// //                        &graph.extend_using(|| { |&((_,_),c)| c as u64 })])
// //           .flat_map(|(p,es)| es.into_iter().map(move |e| (p.clone(), e)))
// //     //
// //     //  let mut fives = quads.extend(vec![Box::new(fragment.extend_using(|| { |&(((a,_),_),_)| a as u64 })),
// //     //                                    Box::new(fragment.extend_using(|| { |&(((_,b),_),_)| b as u64 })),
// //     //                                    Box::new(fragment.extend_using(|| { |&(((_,_),c),_)| c as u64 })),
// //     //                                    Box::new(fragment.extend_using(|| { |&(((_,_),_),d)| d as u64 }))]).flatten();
// //     //
// //     // let mut sixes = fives.extend(vec![Box::new(fragment.extend_using(|| { |&((((a,_),_),_),_)| a as u64 })),
// //     //                                   Box::new(fragment.extend_using(|| { |&((((_,b),_),_),_)| b as u64 })),
// //     //                                   Box::new(fragment.extend_using(|| { |&((((_,_),c),_),_)| c as u64 })),
// //     //                                   Box::new(fragment.extend_using(|| { |&((((_,_),_),d),_)| d as u64 })),
// //     //                                   Box::new(fragment.extend_using(|| { |&((((_,_),_),_),e)| e as u64 }))]).flatten();
// //     //
// //     // sixes.observe();
// // }
<commit_msg>made step size a command line parameter<commit_after>#![feature(core)]
#![feature(str_words)]

extern crate mmap;
extern crate time;
extern crate core;
extern crate timely;
extern crate columnar;
extern crate dataflow_join;

extern crate docopt;
use docopt::Docopt;

use std::io::{BufRead, stdin};
use std::rc::Rc;
use std::cell::RefCell;
use std::thread;

use core::cmp::Ordering::*;

use dataflow_join::*;
use dataflow_join::graph::{GraphTrait, GraphMMap, GraphExtenderExt, gallop};

use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::progress::nested::product::Product;
use timely::example_static::*;

use timely::communication::*;
use timely::communication::pact::Pipeline;

static USAGE: &'static str = "
Usage: triangles dataflow <source> <workers> <stepsize> [--inspect] [--interactive]
       triangles autorun <source> <workers> <stepsize>
       triangles compute <source>
       triangles help
";

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    if args.get_bool("dataflow") {
        let inspect = args.get_bool("--inspect");
        let interactive = args.get_bool("--interactive");
        let workers = if let Ok(threads) = args.get_str("<workers>").parse() { threads }
                      else { panic!("invalid setting for workers: {}", args.get_str("<workers>")) };
        let stepsize = if let Ok(size) = args.get_str("<stepsize>").parse() { size }
                      else { panic!("invalid setting for stepsize: {}", args.get_str("<stepsize>")) };
        println!("starting triangles dataflow with {:?} worker{}; inspection: {:?}, interactive: {:?}",
                    workers, if workers == 1 { "" } else { "s" }, inspect, interactive);
        let source = args.get_str("<source>");

        triangles_multi(ProcessCommunicator::new_vector(workers), |_, _| GraphMMap::new(&source), inspect, interactive, stepsize);
    }
    if args.get_bool("autorun") {
        let workers = if let Ok(threads) = args.get_str("<workers>").parse() { threads }
                      else { panic!("invalid setting for workers: {}", args.get_str("-t")) };
      let stepsize = if let Ok(size) = args.get_str("<stepsize>").parse() { size }
                    else { panic!("invalid setting for stepsize: {}", args.get_str("<stepsize>")) };
        println!("starting triangles dataflow with {:?} worker{}; autorun",
                    workers, if workers == 1 { "" } else { "s" });
        let source = args.get_str("<source>");
        triangles_auto_multi(ProcessCommunicator::new_vector(workers), |_, _| GraphMMap::new(&source), stepsize);
    }
    if args.get_bool("compute") {
        let source = args.get_str("<source>");
        let graph = GraphMMap::new(&source);
        println!("triangles: {:?}", raw_triangles(&graph));
    }

    if args.get_bool("help") {
        println!("the code presently assumes you have access to the livejournal graph, from:");
        println!("   https://snap.stanford.edu/data/soc-LiveJournal1.html");
        println!("");
        println!("before you can use it you will need to \"digest\" it using example digest.rs");
        println!("digest will overwrite <target>.targets and <target>.offsets, so be careful!");
        println!("at least, it will once you edit the code to uncomment the line.");
        println!("Once you grok the binary format, you can totally use other graphs too!");
    }
}

fn raw_triangles<G: GraphTrait<Target=u32>>(graph: &G) -> u64 {
    let mut count = 0;
    for a in (0..graph.nodes()) {
        let aaa = graph.edges(a);
        for &b in aaa {
            let bbb = graph.edges(b as usize);
            count += if aaa.len() < bbb.len() { intersect(aaa.clone(), bbb) }
                     else                     { intersect(bbb, aaa.clone()) };
        }
    }
    count
}

fn intersect<E: Ord>(mut aaa: &[E], mut bbb: &[E]) -> u64 {
    let mut count = 0;
    // magic gallop overhead # is 4
    if aaa.len() < bbb.len() / 4 {
        for a in aaa {
            bbb = gallop(bbb, a);
            if bbb.len() > 0 && &bbb[0] == a { count += 1; }
        }
    }
    else {
        while aaa.len() > 0 && bbb.len() > 0 {
            match aaa[0].cmp(&bbb[0]) {
                Greater => { bbb = &bbb[1..]; },
                Less    => { aaa = &aaa[1..]; },
                Equal   => { aaa = &aaa[1..];
                             bbb = &bbb[1..];
                             count += 1;
                           },
            }
        }
    }
    count
}

fn triangles_multi<C, G, F>(communicators: Vec<C>, loader: F, inspect: bool, interactive: bool, step_size: u64)
where C: Communicator+Send, G: GraphTrait<Target=u32>, F: Fn(u64, u64)->G+Send+Sync {
    let mut guards = Vec::new();
    let loader = &loader;
    for communicator in communicators.into_iter() {
        guards.push(thread::Builder::new().name(format!("worker thread {}", communicator.index()))
                                          .scoped(move || triangles(communicator, loader, inspect, interactive, step_size))
                                          .unwrap());
    }
}

fn triangles<C, G, F>(communicator: C, loader: &F, inspect: bool, interactive: bool, step_size: u64)
where C: Communicator, G: GraphTrait<Target=u32>, F: Fn(u64, u64)->G {
    let comm_index = communicator.index();
    let comm_peers = communicator.peers();

    let graph = Rc::new(RefCell::new(loader(comm_index, comm_peers)));

    let mut root = GraphRoot::new(communicator);
    let mut input = { // new scope to avoid long borrow on root

        let mut builder = root.new_subgraph();
        let (input, stream) = builder.new_input::<u32>();

        // extend u32s to pairs, then pairs to triples.
        let triangles = builder.enable(&stream)
                               .extend(vec![&graph.extend_using(|| { |&a| a as u64 } )])
                               .flat_map(|(p, es)| es.into_iter().map(move |e| (p, e)))
                               .extend(vec![&graph.extend_using(|| { |&(a,_)| a as u64 }),
                                            &graph.extend_using(|| { |&(_,b)| b as u64 })]);

        if inspect { triangles.inspect(|x| println!("triangles: {:?}", x)); }

        input
    };

    if interactive {
        let mut stdinput = stdin();
        let mut line = String::new();
        for index in (0..) {
            stdinput.read_line(&mut line).unwrap();
            let start = time::precise_time_ns();
            if let Some(word) = line.words().next() {
                let read = word.parse();
                if let Ok(number) = read {
                    input.send_messages(&Product::new((), index as u64), vec![number]);
                }
            }
            line.clear();

            input.advance(&Product::new((), index as u64), &Product::new((), index as u64 + 1));
            for _ in (0..5) { root.step(); }

            println!("elapsed: {:?}us", (time::precise_time_ns() - start)/1000);
        }
    }
    else {
        let step = step_size as usize;
        let nodes = graph.borrow().nodes() - 1;
        let limit = (nodes / step) + 1;
        for index in (0..limit) {
            let index = index as usize;
            let data = ((index * step) .. ((index + 1) * step))
                           .filter(|&x| x as u64 % comm_peers == comm_index)
                           .filter(|&x| x < nodes)
                           .map(|x| x as u32)
                           .collect();
            input.send_messages(&Product::new((), index as u64), data);
            input.advance(&Product::new((), index as u64), &Product::new((), index as u64 + 1));
            root.step();
        }

        input.close_at(&Product::new((), limit as u64));
        while root.step() { }
    }
}

fn triangles_auto_multi<C, G, F>(communicators: Vec<C>, loader: F, step_size: u64)
where C: Communicator+Send, G: GraphTrait<Target=u32>, F: Fn(u64, u64)->G+Send+Sync {
    let mut guards = Vec::new();
    let loader = &loader;
    for communicator in communicators.into_iter() {
        guards.push(thread::Builder::new().name(format!("worker thread {}", communicator.index()))
                                          .scoped(move || triangles_auto(communicator, loader, step_size))
                                          .unwrap());
    }
}

fn triangles_auto<C, G, F>(communicator: C, loader: &F, step_size: u64)
where C: Communicator, G: GraphTrait<Target=u32>, F: Fn(u64, u64)->G {
    let index = communicator.index();
    let peers = communicator.peers();

    let batch = step_size;
    let graph = Rc::new(RefCell::new(loader(index, peers)));
    let nodes = graph.borrow().nodes() as u64;

    let mut root = GraphRoot::new(communicator);

    {   // new scope to avoid long borrow on root

        let mut builder = root.new_subgraph::<u64>();

        let (feedback, stream) = builder.loop_variable(Product::new((), nodes + 1), Local(1));

        stream.enable(&mut builder)
              .unary_notify(Pipeline, format!("Input"), vec![Product::new((), 0)], move |handle| {
                  if let Some((time, _count)) = handle.notificator.next() {
                      if time.inner < (nodes / batch) {
                          handle.notificator.notify_at(&Product::new((), time.inner + 1));
                      }
                      let mut session = handle.output.session(&time);
                      for next in (0..(batch/peers)) {
                          session.give(time.inner * batch + next * peers + index);
                      }
                  }})
              .extend(vec![&graph.extend_using(|| { |&a| a as u64 } )])
              .flat_map(|(p, es)| es.into_iter().map(move |e| (p, e)))
              .extend(vec![&graph.extend_using(|| { |&(a,_)| a as u64 }),
                           &graph.extend_using(|| { |&(_,b)| b as u64 })])
              .filter(|_| false)
              .connect_loop(feedback);
    }

    while root.step() { }
}

// // fn _quads<'a, G, G2>(stream: &mut Stream<'a, G, ((u32, u32), u32)>, graph: &Rc<RefCell<G2>>) ->
// //                                                             Stream<'a, G, (((u32, u32), u32), u32)>
// // where G: GraphBuilder+'a, G2: GraphTrait<Target=u32> {
// //     //
// //     stream.extend(vec![&graph.extend_using(|| { |&((a,_),_)| a as u64 }),
// //                        &graph.extend_using(|| { |&((_,b),_)| b as u64 }),
// //                        &graph.extend_using(|| { |&((_,_),c)| c as u64 })])
// //           .flat_map(|(p,es)| es.into_iter().map(move |e| (p.clone(), e)))
// //     //
// //     //  let mut fives = quads.extend(vec![Box::new(fragment.extend_using(|| { |&(((a,_),_),_)| a as u64 })),
// //     //                                    Box::new(fragment.extend_using(|| { |&(((_,b),_),_)| b as u64 })),
// //     //                                    Box::new(fragment.extend_using(|| { |&(((_,_),c),_)| c as u64 })),
// //     //                                    Box::new(fragment.extend_using(|| { |&(((_,_),_),d)| d as u64 }))]).flatten();
// //     //
// //     // let mut sixes = fives.extend(vec![Box::new(fragment.extend_using(|| { |&((((a,_),_),_),_)| a as u64 })),
// //     //                                   Box::new(fragment.extend_using(|| { |&((((_,b),_),_),_)| b as u64 })),
// //     //                                   Box::new(fragment.extend_using(|| { |&((((_,_),c),_),_)| c as u64 })),
// //     //                                   Box::new(fragment.extend_using(|| { |&((((_,_),_),d),_)| d as u64 })),
// //     //                                   Box::new(fragment.extend_using(|| { |&((((_,_),_),_),e)| e as u64 }))]).flatten();
// //     //
// //     // sixes.observe();
// // }
<|endoftext|>


<commit_before>
use std::iter::Filter;
use std::ops::RangeFrom;

use common::divisible;

#[allow(dead_code)]
pub fn is_prime(n: u64) -> bool {
    let half = n / 2 + 1;
    let mut range = 2..half; // TODO: find out why this needs to be 'mut'
    n != 1 && !range.any(|r| divisible(n, r))
}

#[allow(dead_code)]
pub fn factorize(n: u64) -> Vec<u64> {
    let mut factors: Vec<u64> = vec![];
    let mut current = n;
    let mut i = 2;
    loop {
        if divisible(current, i) {
            factors.push(i);
            current = current / i;
            if current == 1 {
                return factors;
            }
        } else {
            i = i + 1;
        }
    }
}

#[allow(dead_code)]
pub fn largest_prime_factor(n: u64) -> u64 {
    let factors = factorize(n);
    match factors.last() {
        Some(&f) => f,
        _ => 1
    }
}

// TODO: finish functional way

#[allow(dead_code)]
#[allow(unused_variables)]
pub fn factorize_fun(n: u64) -> Vec<u64> {
    let primes = (1..).filter(|&n| is_prime(n));
    let five = primes.take(5).collect::<Vec<u64>>();
    //print!("!!!!!!!! primes: {:?}", five);
    //let a = prime_factors(n, primes);
    vec![]
}

#[allow(dead_code)]
#[allow(unused_variables)]
fn prime_factors(n: u64, series: Filter<RangeFrom<u64>, u64>) -> Vec<u64> {
    match (n, series) {
        (1, _)  => vec![],
        (m, ps) => {
            let (q, r) = (m/2, m/2);
            vec![]
        }
    }
}

// \TODO

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    //#[should_panic(expected="assertion failed")]
    fn one_is_not_prime() {
        assert!(!is_prime(1));
    }

    #[test]
    fn there_are_168_primes_under_1000() {
        let primes = (1..1000).filter(|n| is_prime(*n)).count();
        assert_eq!(primes, 168);
    }

    #[test]
    fn example_factors() {
        assert_eq!(factorize(13195), vec![5, 7, 13, 29]);
    }

    #[test]
    fn factors_repeating() { // TODO: this is not what I wanted
        factorize_fun(12);
        assert_eq!(factorize(12), vec![2, 2, 3]);
    }

    #[test]
    fn example() {
        assert_eq!(largest_prime_factor(13195), 29);
    }

    #[test]
    fn solution() {
        assert_eq!(largest_prime_factor(600851475143), 6857);
    }
}
<commit_msg>factorization with iterator, unfinished<commit_after>
use common::divisible;
use common::quot_rem;

#[allow(dead_code)]
pub fn is_prime(n: u64) -> bool {
    let half = n / 2 + 1;
    let mut range = 2..half; // TODO: find out why this needs to be 'mut'
    n != 1 && !range.any(|r| divisible(n, r))
}

#[allow(dead_code)]
pub fn factorize(n: u64) -> Vec<u64> {
    let mut factors: Vec<u64> = vec![];
    let mut current = n;
    let mut i = 2;
    loop {
        if divisible(current, i) {
            factors.push(i);
            current = current / i;
            if current == 1 {
                return factors;
            }
        } else {
            i = i + 1;
        }
    }
}

#[allow(dead_code)]
pub fn largest_prime_factor(n: u64) -> u64 {
    let factors = factorize(n);
    match factors.last() {
        Some(&f) => f,
        _ => 1,
    }
}

// TODO: finish functional way

#[allow(dead_code)]
#[allow(unused_variables)]
pub fn factorize_fun(n: u64) -> Vec<u64> {
    let primes = (1..).filter(|&n| is_prime(n));
    let five = primes.take(5).collect::<Vec<u64>>();
    // print!("!!!!!!!! primes: {:?}", five);
    // let a = prime_factors(n, primes);
    vec![]
}

#[allow(dead_code)]
#[allow(unused_variables)]
#[allow(unused_mut)]
fn prime_factors(n: u64, factors: Vec<u64>) -> Vec<u64> {
    match (n, factors) {
        (1, _) => vec![],
        (m, fs) => {
            let mut f = fs[0];
            let (q, r) = quot_rem(m, f);
            if m < f * f {
                return vec![m];
            } else if r == 0 {
                let mut f_vec = vec![f];
                let mut rest = prime_factors(q, fs);
                //let f_vec: Vec<u64> = f_vec.append(&mut rest);
                //return f_vec;

            } else {
                let f2: Vec<u64> = prime_factors(m, fs);
                return f2;
            }
            vec![]
        }
    }
}

// \TODO

#[allow(dead_code)]
pub struct Primes {
    primes: Vec<u64>,
    current: u64,
}

#[allow(dead_code)]
pub fn primes() -> Primes {
    Primes {
        primes: vec![],
        current: 2,
    }
}

impl Iterator for Primes {
    type Item = u64;
    fn next(&mut self) -> Option<u64> {
        for i in self.current..u64::max_value() {
            if self.primes.iter().all(|x| i % x != 0) {
                self.primes.push(i);
                self.current = i + 1;
                return Some(i);
            }
        }
        panic!("Integer overflow");
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    // #[should_panic(expected="assertion failed")]
    fn one_is_not_prime() {
        assert!(!is_prime(1));
    }

    #[test]
    fn there_are_168_primes_under_1000() {
        let primes = (1..1000).filter(|n| is_prime(*n)).count();
        assert_eq!(primes, 168);
    }

    #[test]
    fn example_factors() {
        assert_eq!(factorize(13195), vec![5, 7, 13, 29]);
    }

    #[test]
    fn factors_repeating() {
        // TODO: this is not what I wanted
        assert_eq!(factorize(12), vec![2, 2, 3]);
    }

    #[test]
    fn example() {
        assert_eq!(largest_prime_factor(13195), 29);
    }

    #[test]
    fn solution() {
        assert_eq!(largest_prime_factor(600851475143), 6857);
    }
}
<|endoftext|>