<commit_before>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// 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.

//! Utilities for working with packets, eg. checksumming.

use ip::IpNextHeaderProtocol;
use pnet_macros_support::types::u16be;

use std::net::{Ipv4Addr, Ipv6Addr};
use std::slice;
use std::u8;

/// Convert a value to a byte array.
pub trait Octets {
    /// Output type - bytes array.
    type Output;

    /// Return a value as bytes (big-endian order).
    fn octets(&self) -> Self::Output;
}

impl Octets for u64 {
    type Output = [u8; 8];

    fn octets(&self) -> Self::Output {
        [(*self >> 56) as u8,
         (*self >> 48) as u8,
         (*self >> 40) as u8,
         (*self >> 32) as u8,
         (*self >> 24) as u8,
         (*self >> 16) as u8,
         (*self >> 8) as u8,
         *self as u8]
    }
}

impl Octets for u32 {
    type Output = [u8; 4];

    fn octets(&self) -> Self::Output {
        [(*self >> 24) as u8, (*self >> 16) as u8, (*self >> 8) as u8, *self as u8]
    }
}

impl Octets for u16 {
    type Output = [u8; 2];

    fn octets(&self) -> Self::Output {
        [(*self >> 8) as u8, *self as u8]
    }
}

impl Octets for u8 {
    type Output = [u8; 1];

    fn octets(&self) -> Self::Output {
        [*self]
    }
}

/// Calculates a checksum. Used by ipv4 and icmp. The two bytes starting at `skipword * 2` will be
/// ignored. Supposed to be the checksum field, which is regarded as zero during calculation. Will
/// panic if data is not u16-aligned.
pub fn checksum(data: &[u8], skipword: usize) -> u16be {
    if data.len() == 0 { return 0 }
    assert_eq!(0, data.as_ptr() as usize % 2, "Cannot sum mis-aligned words at {:p}", data.as_ptr());
    let sum = sum_be_words(data, skipword);
    finalize_checksum(sum)
}

fn finalize_checksum(mut sum: u32) -> u16be {
    while sum >> 16 != 0 {
        sum = (sum >> 16) + (sum & 0xFFFF);
    }
    !sum as u16
}

/// Calculate the checksum for a packet built on IPv4. Used by UDP and TCP.
pub fn ipv4_checksum(data: &[u8],
                     skipword: usize,
                     extra_data: &[u8],
                     source: &Ipv4Addr,
                     destination: &Ipv4Addr,
                     next_level_protocol: IpNextHeaderProtocol)
    -> u16be {
    let mut sum = 0u32;

    // Checksum pseudo-header
    sum += ipv4_word_sum(source);
    sum += ipv4_word_sum(destination);

    let IpNextHeaderProtocol(next_level_protocol) = next_level_protocol;
    sum += next_level_protocol as u32;

    let len = data.len() + extra_data.len();
    sum += len as u32;

    // Checksum packet header and data
    sum += sum_be_words(data, skipword);
    sum += sum_be_words(extra_data, extra_data.len() / 2);

    finalize_checksum(sum)
}

fn ipv4_word_sum(ip: &Ipv4Addr) -> u32 {
    let octets = ip.octets();
    ((octets[0] as u32) << 8 | octets[1] as u32) + ((octets[2] as u32) << 8 | octets[3] as u32)
}

/// Calculate the checksum for a packet built on IPv6.
pub fn ipv6_checksum(data: &[u8],
                     skipword: usize,
                     extra_data: &[u8],
                     source: &Ipv6Addr,
                     destination: &Ipv6Addr,
                     next_level_protocol: IpNextHeaderProtocol)
    -> u16be {
    let mut sum = 0u32;

    // Checksum pseudo-header
    sum += ipv6_word_sum(source);
    sum += ipv6_word_sum(destination);

    let IpNextHeaderProtocol(next_level_protocol) = next_level_protocol;
    sum += next_level_protocol as u32;

    let len = data.len() + extra_data.len();
    sum += len as u32;

    // Checksum packet header and data
    sum += sum_be_words(data, skipword);
    sum += sum_be_words(extra_data, extra_data.len() / 2);

    finalize_checksum(sum)
}

fn ipv6_word_sum(ip: &Ipv6Addr) -> u32 {
    ip.segments().iter().map(|x| *x as u32).sum()
}

/// Sum all words (16 bit chunks) in the given data. The word at word offset
/// `skipword` will be skipped. Each word is treated as big endian. Must be
/// called with u16-aligned data.
fn sum_be_words(data: &[u8], mut skipword: usize) -> u32 {
    if data.len() == 0 { return 0 }
    debug_assert_eq!(0, data.as_ptr() as usize % 2, "Cannot sum mis-aligned words at {:p}", data.as_ptr());
    let len = data.len();
    let wdata: &[u16] = unsafe { slice::from_raw_parts(data.as_ptr() as *const u16, len / 2) };
    skipword = ::std::cmp::min(skipword, wdata.len());

    let mut sum = 0u32;
    let mut i = 0;
    while i < skipword {
        sum += u16::from_be(unsafe { *wdata.get_unchecked(i) }) as u32;
        i += 1;
    }
    i += 1;
    while i < wdata.len() {
        sum += u16::from_be(unsafe { *wdata.get_unchecked(i) }) as u32;
        i += 1;
    }
    // If the length is odd, make sure to checksum the final byte
    if len & 1 != 0 {
        sum += (unsafe { *data.get_unchecked(len - 1) } as u32) << 8;
    }

    sum
}

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

    #[test]
    fn sum_be_words_different_skipwords() {
        let data = (0..11).collect::<Vec<u8>>();
        assert_eq!(7190, sum_be_words(&data, 1));
        assert_eq!(6676, sum_be_words(&data, 2));
        // Assert having the skipword outside the range gives correct and equal
        // results
        assert_eq!(7705, sum_be_words(&data, 99));
        assert_eq!(7705, sum_be_words(&data, 101));
    }
}

#[cfg(all(test, feature = "benchmark"))]
mod checksum_benchmarks {
    use super::checksum;
    use test::{Bencher, black_box};

    #[bench]
    fn bench_checksum_small(b: &mut Bencher) {
        let data = vec![99u8; 20];
        b.iter(|| checksum(black_box(&data), 5));
    }

    #[bench]
    fn bench_checksum_large(b: &mut Bencher) {
        let data = vec![123u8; 1024];
        b.iter(|| checksum(black_box(&data), 5));
    }
}

<commit_msg>Fix issue 399 by using u16::from_be_bytes so we aren't using any raw pointers.<commit_after>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// 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.

//! Utilities for working with packets, eg. checksumming.

use ip::IpNextHeaderProtocol;
use pnet_macros_support::types::u16be;

use std::net::{Ipv4Addr, Ipv6Addr};
use std::slice;
use std::u8;
use std::u16;
use std::convert::TryInto;

/// Convert a value to a byte array.
pub trait Octets {
    /// Output type - bytes array.
    type Output;

    /// Return a value as bytes (big-endian order).
    fn octets(&self) -> Self::Output;
}

impl Octets for u64 {
    type Output = [u8; 8];

    fn octets(&self) -> Self::Output {
        [(*self >> 56) as u8,
         (*self >> 48) as u8,
         (*self >> 40) as u8,
         (*self >> 32) as u8,
         (*self >> 24) as u8,
         (*self >> 16) as u8,
         (*self >> 8) as u8,
         *self as u8]
    }
}

impl Octets for u32 {
    type Output = [u8; 4];

    fn octets(&self) -> Self::Output {
        [(*self >> 24) as u8, (*self >> 16) as u8, (*self >> 8) as u8, *self as u8]
    }
}

impl Octets for u16 {
    type Output = [u8; 2];

    fn octets(&self) -> Self::Output {
        [(*self >> 8) as u8, *self as u8]
    }
}

impl Octets for u8 {
    type Output = [u8; 1];

    fn octets(&self) -> Self::Output {
        [*self]
    }
}

/// Calculates a checksum. Used by ipv4 and icmp. The two bytes starting at `skipword * 2` will be
/// ignored. Supposed to be the checksum field, which is regarded as zero during calculation.
pub fn checksum(data: &[u8], skipword: usize) -> u16be {
    if data.len() == 0 { return 0 }
    let sum = sum_be_words(data, skipword);
    finalize_checksum(sum)
}

fn finalize_checksum(mut sum: u32) -> u16be {
    while sum >> 16 != 0 {
        sum = (sum >> 16) + (sum & 0xFFFF);
    }
    !sum as u16
}

/// Calculate the checksum for a packet built on IPv4. Used by UDP and TCP.
pub fn ipv4_checksum(data: &[u8],
                     skipword: usize,
                     extra_data: &[u8],
                     source: &Ipv4Addr,
                     destination: &Ipv4Addr,
                     next_level_protocol: IpNextHeaderProtocol)
    -> u16be {
    let mut sum = 0u32;

    // Checksum pseudo-header
    sum += ipv4_word_sum(source);
    sum += ipv4_word_sum(destination);

    let IpNextHeaderProtocol(next_level_protocol) = next_level_protocol;
    sum += next_level_protocol as u32;

    let len = data.len() + extra_data.len();
    sum += len as u32;

    // Checksum packet header and data
    sum += sum_be_words(data, skipword);
    sum += sum_be_words(extra_data, extra_data.len() / 2);

    finalize_checksum(sum)
}

fn ipv4_word_sum(ip: &Ipv4Addr) -> u32 {
    let octets = ip.octets();
    ((octets[0] as u32) << 8 | octets[1] as u32) + ((octets[2] as u32) << 8 | octets[3] as u32)
}

/// Calculate the checksum for a packet built on IPv6.
pub fn ipv6_checksum(data: &[u8],
                     skipword: usize,
                     extra_data: &[u8],
                     source: &Ipv6Addr,
                     destination: &Ipv6Addr,
                     next_level_protocol: IpNextHeaderProtocol)
    -> u16be {
    let mut sum = 0u32;

    // Checksum pseudo-header
    sum += ipv6_word_sum(source);
    sum += ipv6_word_sum(destination);

    let IpNextHeaderProtocol(next_level_protocol) = next_level_protocol;
    sum += next_level_protocol as u32;

    let len = data.len() + extra_data.len();
    sum += len as u32;

    // Checksum packet header and data
    sum += sum_be_words(data, skipword);
    sum += sum_be_words(extra_data, extra_data.len() / 2);

    finalize_checksum(sum)
}

fn ipv6_word_sum(ip: &Ipv6Addr) -> u32 {
    ip.segments().iter().map(|x| *x as u32).sum()
}

/// Sum all words (16 bit chunks) in the given data. The word at word offset
/// `skipword` will be skipped. Each word is treated as big endian.
fn sum_be_words(data: &[u8], mut skipword: usize) -> u32 {
    if data.len() == 0 { return 0 }
    let len = data.len();
    let mut cur_data = &data[..];
    let mut sum = 0u32;
    let mut i = 0;
    while cur_data.len() >= 2 {
        if i != skipword {
            // It's safe to unwrap because we verified there are at least 2 bytes
            sum += u16::from_be_bytes(cur_data[0..2].try_into().unwrap()) as u32;
        }
        cur_data = &cur_data[2..];
        i += 1;
    }

    // If the length is odd, make sure to checksum the final byte
    if len & 1 != 0 {
        sum += (data[len - 1] as u32) << 8;
    }

    sum
}

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

    #[test]
    fn sum_be_words_different_skipwords() {
        let data = (0..11).collect::<Vec<u8>>();
        assert_eq!(7190, sum_be_words(&data, 1));
        assert_eq!(6676, sum_be_words(&data, 2));
        // Assert having the skipword outside the range gives correct and equal
        // results
        assert_eq!(7705, sum_be_words(&data, 99));
        assert_eq!(7705, sum_be_words(&data, 101));
    }
}

#[cfg(all(test, feature = "benchmark"))]
mod checksum_benchmarks {
    use super::checksum;
    use test::{Bencher, black_box};

    #[bench]
    fn bench_checksum_small(b: &mut Bencher) {
        let data = vec![99u8; 20];
        b.iter(|| checksum(black_box(&data), 5));
    }

    #[bench]
    fn bench_checksum_large(b: &mut Bencher) {
        let data = vec![123u8; 1024];
        b.iter(|| checksum(black_box(&data), 5));
    }
}

<|endoftext|>


<commit_before>use bytes::{BufMut, BytesMut};
use postgres_protocol::message::backend;
use postgres_protocol::message::frontend;
use socket2::{Domain, SockAddr, Socket, Type};
use std::io::{self, BufWriter, Read, Write};
use std::net::{SocketAddr, ToSocketAddrs};
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, RawFd};
#[cfg(windows)]
use std::os::windows::io::{AsRawSocket, RawSocket};
use std::result;
use std::time::Duration;

use error;
use params::{ConnectParams, Host};
use tls::TlsStream;
use {Result, TlsMode};

const INITIAL_CAPACITY: usize = 8 * 1024;

pub struct MessageStream {
    stream: BufWriter<Box<TlsStream>>,
    in_buf: BytesMut,
    out_buf: Vec<u8>,
}

impl MessageStream {
    pub fn new(stream: Box<TlsStream>) -> MessageStream {
        MessageStream {
            stream: BufWriter::new(stream),
            in_buf: BytesMut::with_capacity(INITIAL_CAPACITY),
            out_buf: vec![],
        }
    }

    pub fn get_ref(&self) -> &Box<TlsStream> {
        self.stream.get_ref()
    }

    pub fn write_message<F, E>(&mut self, f: F) -> result::Result<(), E>
    where
        F: FnOnce(&mut Vec<u8>) -> result::Result<(), E>,
        E: From<io::Error>,
    {
        self.out_buf.clear();
        f(&mut self.out_buf)?;
        self.stream.write_all(&self.out_buf).map_err(From::from)
    }

    pub fn read_message(&mut self) -> io::Result<backend::Message> {
        loop {
            match backend::Message::parse(&mut self.in_buf) {
                Ok(Some(message)) => return Ok(message),
                Ok(None) => self.read_in()?,
                Err(e) => return Err(e),
            }
        }
    }

    fn read_in(&mut self) -> io::Result<()> {
        self.in_buf.reserve(1);
        match self.stream
            .get_mut()
            .read(unsafe { self.in_buf.bytes_mut() })
        {
            Ok(0) => Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "unexpected EOF",
            )),
            Ok(n) => {
                unsafe { self.in_buf.advance_mut(n) };
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    pub fn read_message_timeout(
        &mut self,
        timeout: Duration,
    ) -> io::Result<Option<backend::Message>> {
        if self.in_buf.is_empty() {
            self.set_read_timeout(Some(timeout))?;
            let r = self.read_in();
            self.set_read_timeout(None)?;

            match r {
                Ok(()) => {}
                Err(ref e)
                    if e.kind() == io::ErrorKind::WouldBlock
                        || e.kind() == io::ErrorKind::TimedOut =>
                {
                    return Ok(None)
                }
                Err(e) => return Err(e),
            }
        }

        self.read_message().map(Some)
    }

    pub fn read_message_nonblocking(&mut self) -> io::Result<Option<backend::Message>> {
        if self.in_buf.is_empty() {
            self.set_nonblocking(true)?;
            let r = self.read_in();
            self.set_nonblocking(false)?;

            match r {
                Ok(()) => {}
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(None),
                Err(e) => return Err(e),
            }
        }

        self.read_message().map(Some)
    }

    pub fn flush(&mut self) -> io::Result<()> {
        self.stream.flush()
    }

    fn set_read_timeout(&self, timeout: Option<Duration>) -> io::Result<()> {
        self.stream.get_ref().get_ref().0.set_read_timeout(timeout)
    }

    fn set_nonblocking(&self, nonblock: bool) -> io::Result<()> {
        self.stream.get_ref().get_ref().0.set_nonblocking(nonblock)
    }
}

/// A connection to the Postgres server.
///
/// It implements `Read`, `Write` and `TlsStream`, as well as `AsRawFd` on
/// Unix platforms and `AsRawSocket` on Windows platforms.
#[derive(Debug)]
pub struct Stream(Socket);

impl Read for Stream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl Write for Stream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

impl TlsStream for Stream {
    fn get_ref(&self) -> &Stream {
        self
    }

    fn get_mut(&mut self) -> &mut Stream {
        self
    }
}

#[cfg(unix)]
impl AsRawFd for Stream {
    fn as_raw_fd(&self) -> RawFd {
        self.0.as_raw_fd()
    }
}

#[cfg(windows)]
impl AsRawSocket for Stream {
    fn as_raw_socket(&self) -> RawSocket {
        self.0.as_raw_socket()
    }
}

fn open_socket(params: &ConnectParams) -> Result<Socket> {
    let port = params.port();
    match *params.host() {
        Host::Tcp(ref host) => {
            let mut error = None;
            for addr in (&**host, port).to_socket_addrs()? {
                let domain = match addr {
                    SocketAddr::V4(_) => Domain::ipv4(),
                    SocketAddr::V6(_) => Domain::ipv6(),
                };
                let socket = Socket::new(domain, Type::stream(), None)?;
                if let Some(keepalive) = params.keepalive() {
                    socket.set_keepalive(Some(keepalive))?;
                }
                let addr = SockAddr::from(addr);
                let r = match params.connect_timeout() {
                    Some(timeout) => socket.connect_timeout(&addr, timeout),
                    None => socket.connect(&addr),
                };
                match r {
                    Ok(()) => return Ok(socket),
                    Err(e) => error = Some(e),
                }
            }

            Err(error
                .unwrap_or_else(|| {
                    io::Error::new(
                        io::ErrorKind::InvalidInput,
                        "could not resolve any addresses",
                    )
                })
                .into())
        }
        #[cfg(unix)]
        Host::Unix(ref path) => {
            let path = path.join(&format!(".s.PGSQL.{}", port));
            let socket = Socket::new(Domain::unix(), Type::stream(), None)?;
            let addr = SockAddr::unix(path)?;
            socket.connect(&addr)?;
            Ok(socket)
        }
        #[cfg(not(unix))]
        Host::Unix(..) => Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "unix sockets are not supported on this system",
        ).into()),
    }
}

pub fn initialize_stream(params: &ConnectParams, tls: TlsMode) -> Result<Box<TlsStream>> {
    let mut socket = Stream(open_socket(params)?);

    let (tls_required, handshaker) = match tls {
        TlsMode::None => return Ok(Box::new(socket)),
        TlsMode::Prefer(handshaker) => (false, handshaker),
        TlsMode::Require(handshaker) => (true, handshaker),
    };

    let mut buf = vec![];
    frontend::ssl_request(&mut buf);
    socket.write_all(&buf)?;
    socket.flush()?;

    let mut b = [0; 1];
    socket.read_exact(&mut b)?;
    if b[0] == b'N' {
        if tls_required {
            return Err(error::tls("the server does not support TLS".into()));
        } else {
            return Ok(Box::new(socket));
        }
    }

    let host = match *params.host() {
        Host::Tcp(ref host) => host,
        // Postgres doesn't support TLS over unix sockets
        Host::Unix(_) => return Err(::bad_response().into()),
    };

    handshaker.tls_handshake(host, socket).map_err(error::tls)
}
<commit_msg>Little cleanup<commit_after>use bytes::{BufMut, BytesMut};
use postgres_protocol::message::backend;
use postgres_protocol::message::frontend;
use socket2::{Domain, SockAddr, Socket, Type};
use std::io::{self, BufWriter, Read, Write};
use std::net::{SocketAddr, ToSocketAddrs};
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, RawFd};
#[cfg(windows)]
use std::os::windows::io::{AsRawSocket, RawSocket};
use std::result;
use std::time::Duration;

use error;
use params::{ConnectParams, Host};
use tls::TlsStream;
use {Result, TlsMode};

const INITIAL_CAPACITY: usize = 8 * 1024;

pub struct MessageStream {
    stream: BufWriter<Box<TlsStream>>,
    in_buf: BytesMut,
    out_buf: Vec<u8>,
}

impl MessageStream {
    pub fn new(stream: Box<TlsStream>) -> MessageStream {
        MessageStream {
            stream: BufWriter::new(stream),
            in_buf: BytesMut::with_capacity(INITIAL_CAPACITY),
            out_buf: vec![],
        }
    }

    pub fn get_ref(&self) -> &TlsStream {
        &**self.stream.get_ref()
    }

    pub fn write_message<F, E>(&mut self, f: F) -> result::Result<(), E>
    where
        F: FnOnce(&mut Vec<u8>) -> result::Result<(), E>,
        E: From<io::Error>,
    {
        self.out_buf.clear();
        f(&mut self.out_buf)?;
        self.stream.write_all(&self.out_buf).map_err(From::from)
    }

    pub fn read_message(&mut self) -> io::Result<backend::Message> {
        loop {
            match backend::Message::parse(&mut self.in_buf) {
                Ok(Some(message)) => return Ok(message),
                Ok(None) => self.read_in()?,
                Err(e) => return Err(e),
            }
        }
    }

    fn read_in(&mut self) -> io::Result<()> {
        self.in_buf.reserve(1);
        match self.stream
            .get_mut()
            .read(unsafe { self.in_buf.bytes_mut() })
        {
            Ok(0) => Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "unexpected EOF",
            )),
            Ok(n) => {
                unsafe { self.in_buf.advance_mut(n) };
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    pub fn read_message_timeout(
        &mut self,
        timeout: Duration,
    ) -> io::Result<Option<backend::Message>> {
        if self.in_buf.is_empty() {
            self.set_read_timeout(Some(timeout))?;
            let r = self.read_in();
            self.set_read_timeout(None)?;

            match r {
                Ok(()) => {}
                Err(ref e)
                    if e.kind() == io::ErrorKind::WouldBlock
                        || e.kind() == io::ErrorKind::TimedOut =>
                {
                    return Ok(None)
                }
                Err(e) => return Err(e),
            }
        }

        self.read_message().map(Some)
    }

    pub fn read_message_nonblocking(&mut self) -> io::Result<Option<backend::Message>> {
        if self.in_buf.is_empty() {
            self.set_nonblocking(true)?;
            let r = self.read_in();
            self.set_nonblocking(false)?;

            match r {
                Ok(()) => {}
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(None),
                Err(e) => return Err(e),
            }
        }

        self.read_message().map(Some)
    }

    pub fn flush(&mut self) -> io::Result<()> {
        self.stream.flush()
    }

    fn set_read_timeout(&self, timeout: Option<Duration>) -> io::Result<()> {
        self.stream.get_ref().get_ref().0.set_read_timeout(timeout)
    }

    fn set_nonblocking(&self, nonblock: bool) -> io::Result<()> {
        self.stream.get_ref().get_ref().0.set_nonblocking(nonblock)
    }
}

/// A connection to the Postgres server.
///
/// It implements `Read`, `Write` and `TlsStream`, as well as `AsRawFd` on
/// Unix platforms and `AsRawSocket` on Windows platforms.
#[derive(Debug)]
pub struct Stream(Socket);

impl Read for Stream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl Write for Stream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

impl TlsStream for Stream {
    fn get_ref(&self) -> &Stream {
        self
    }

    fn get_mut(&mut self) -> &mut Stream {
        self
    }
}

#[cfg(unix)]
impl AsRawFd for Stream {
    fn as_raw_fd(&self) -> RawFd {
        self.0.as_raw_fd()
    }
}

#[cfg(windows)]
impl AsRawSocket for Stream {
    fn as_raw_socket(&self) -> RawSocket {
        self.0.as_raw_socket()
    }
}

fn open_socket(params: &ConnectParams) -> Result<Socket> {
    let port = params.port();
    match *params.host() {
        Host::Tcp(ref host) => {
            let mut error = None;
            for addr in (&**host, port).to_socket_addrs()? {
                let domain = match addr {
                    SocketAddr::V4(_) => Domain::ipv4(),
                    SocketAddr::V6(_) => Domain::ipv6(),
                };
                let socket = Socket::new(domain, Type::stream(), None)?;
                if let Some(keepalive) = params.keepalive() {
                    socket.set_keepalive(Some(keepalive))?;
                }
                let addr = SockAddr::from(addr);
                let r = match params.connect_timeout() {
                    Some(timeout) => socket.connect_timeout(&addr, timeout),
                    None => socket.connect(&addr),
                };
                match r {
                    Ok(()) => return Ok(socket),
                    Err(e) => error = Some(e),
                }
            }

            Err(error
                .unwrap_or_else(|| {
                    io::Error::new(
                        io::ErrorKind::InvalidInput,
                        "could not resolve any addresses",
                    )
                })
                .into())
        }
        #[cfg(unix)]
        Host::Unix(ref path) => {
            let path = path.join(&format!(".s.PGSQL.{}", port));
            let socket = Socket::new(Domain::unix(), Type::stream(), None)?;
            let addr = SockAddr::unix(path)?;
            socket.connect(&addr)?;
            Ok(socket)
        }
        #[cfg(not(unix))]
        Host::Unix(..) => Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "unix sockets are not supported on this system",
        ).into()),
    }
}

pub fn initialize_stream(params: &ConnectParams, tls: TlsMode) -> Result<Box<TlsStream>> {
    let mut socket = Stream(open_socket(params)?);

    let (tls_required, handshaker) = match tls {
        TlsMode::None => return Ok(Box::new(socket)),
        TlsMode::Prefer(handshaker) => (false, handshaker),
        TlsMode::Require(handshaker) => (true, handshaker),
    };

    let mut buf = vec![];
    frontend::ssl_request(&mut buf);
    socket.write_all(&buf)?;
    socket.flush()?;

    let mut b = [0; 1];
    socket.read_exact(&mut b)?;
    if b[0] == b'N' {
        if tls_required {
            return Err(error::tls("the server does not support TLS".into()));
        } else {
            return Ok(Box::new(socket));
        }
    }

    let host = match *params.host() {
        Host::Tcp(ref host) => host,
        // Postgres doesn't support TLS over unix sockets
        Host::Unix(_) => return Err(::bad_response().into()),
    };

    handshaker.tls_handshake(host, socket).map_err(error::tls)
}
<|endoftext|>


<commit_before>extern crate rand;

use rand::Rng;

static BINARY: [&str; 2] = ["0", "1"];

pub trait Cromosome {
    fn new(size: usize, ) -> Cromosome;
    fn random(&mut self);
    fn fitness(&self) -> u32;
    fn crossover(&mut self, parent2: &mut Cromosome, crosspoint: usize);
    fn mutate(&mut self);
}

#[derive(Debug)]
pub struct Cromosome {
    size: usize,
    genotype: String,
}

impl Cromosome {
    pub fn new(size: usize) -> Cromosome {
        Cromosome {
            size: size,
            genotype: String::new(),
        }
    }

    pub fn random(&mut self) {
        let mut rng = rand::thread_rng();
        for _ in 0..self.size {
            self.genotype += BINARY[rng.gen::<usize>() % BINARY.len()];
        }
    }

    pub fn fitness(&self) -> u32 {
        let mut fitness = 0;
        let mut exp = self.genotype.len() as u32;

        for i in 0..self.genotype.len() {
            fitness += if self.genotype.chars().nth(i).unwrap() == '1' {
                2u32.pow(exp - 1)
            } else {
                0
            };
            exp -= 1;
        }

        fitness
    }

    pub fn crossover(&mut self, parent2: &mut Cromosome, crosspoint: usize) {
        if crosspoint > self.genotype.len() {
            println!("Not a valid crosspoint");
            return;
        }

        let remnant1 = self.genotype.split_off(crosspoint - 1);
        let remnant2 = parent2.genotype.split_off(crosspoint - 1);

        self.genotype += remnant2.as_ref();
        parent2.genotype += remnant1.as_ref();
    }

    pub fn mutate(&mut self) {
        println!("{}", self.genotype);
        let mut rng = rand::thread_rng();
        let pos = rng.gen::<usize>() % self.size;
        self.genotype.remove(pos);
        self.genotype
            .insert_str(pos, BINARY[rng.gen::<usize>() % BINARY.len()]);
        println!("{}", self.genotype);
    }
}

#[derive(Debug)]
pub struct Population<Cr> {
    pob_size: usize,
    crom_size: usize,
    population: Vec<Cr>,
    ruleta_vec: Vec<f32>,
}

impl Population {
    pub fn new(pob_size: usize, crom_size: usize) -> Population {
        Population {
            pob_size: pob_size,
            crom_size: crom_size,
            population: Vec::new(),
            ruleta_vec: Vec::new(),
        }
    }

    pub fn generate(&mut self) {
        println!("\nGenerando Población inicial");

        for _ in 0..self.pob_size {
            let mut crom = Cromosome::new(self.crom_size);
            crom.random();
            self.population.push(crom);
        }

        for i in 0..self.pob_size {
            println!("{}) {}", i + 1, self.population[i].genotype);
        }
    }

    pub fn eval(&self) {
        println!("Evaluando Individuos");
        for i in 0..self.pob_size {
            println!("{}) {} - {}",
                     i + 1,
                     self.population[i].genotype,
                     self.population[i].fitness());
        }
    }

    pub fn ruleta(&mut self) {
        println!("Selección de Individuos - Método de la Ruleta");
        let mut sum_ruleta = 0;
        self.ruleta_vec = Vec::new();

        for i in 0..self.pob_size {
            sum_ruleta += self.population[i].fitness();
        }

        for i in 0..self.pob_size {
            self.ruleta_vec
                .push((self.population[i].fitness() as f32 * 100.0) / sum_ruleta as f32);
        }

        for i in 0..self.pob_size {
            println!("{}) {} - {} - {}",
                     i + 1,
                     self.population[i].genotype,
                     self.population[i].fitness(),
                     self.ruleta_vec[i]);
        }
    }

    pub fn random_parent_index(&self) -> usize {
        let mut rng = rand::thread_rng();
        let num = rng.gen::<f32>() * 100.0;
        let mut s = 0.0;

        for i in 0..self.ruleta_vec.len() {
            s += self.ruleta_vec[i];
            if num < s {
                return i;
            }
        }

        self.ruleta_vec.len()
    }

    pub fn crossover(&mut self, crossover_prob: f32, crossover_point: usize) {
        let parent1 = self.random_parent_index();
        let mut parent2 = self.random_parent_index();

        while parent2 == parent1 {
            let mut rng = rand::thread_rng();
            parent2 = rng.gen::<usize>() % self.pob_size;
        }

        println!("Padre: {}", parent1 + 1);
        println!("Madre: {}", parent2 + 1);

        let mut child1 = Cromosome {
            size: self.population[parent1].size,
            genotype: self.population[parent1].genotype.clone(),
        };

        let mut child2 = Cromosome {
            size: self.population[parent2].size,
            genotype: self.population[parent2].genotype.clone(),
        };

        let mut rng = rand::thread_rng();
        if rng.gen::<f32>() < crossover_prob {
            println!("Cruzamiento");
            child1.crossover(&mut child2, crossover_point + 1);
        } else {
            println!("Sin cruzamiento");
        }

        self.population.push(child1);
        self.population.push(child2);
    }

    pub fn mutate(&mut self, chance: f32) {
        for i in 0..self.population.len() {
            let mut rng = rand::thread_rng();
            if rng.gen::<f32>() > chance {
                return;
            }
            println!("Mutó {}", i + 1);
            self.population[i].mutate();
        }
    }

    pub fn selection(&mut self) {
        println!("Selección de Siguiente Población");
        let mut new_population: Vec<(u32, String)> = Vec::new();
        for i in 0..self.population.len() {
            let fitness = self.population[i].fitness();
            let gen = self.population[i].genotype.clone();
            println!("{}) {} - {}", i + 1, gen, fitness);
            new_population.push((fitness, gen));
        }

        new_population.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());

        self.population.clear();
        for i in 0..self.pob_size {
            self.population
                .push(Cromosome {
                          size: new_population[i].1.len(),
                          genotype: new_population[i].1.clone(),
                      });
        }

        // println!("{:?}", new_population);
    }
}

#[derive(Debug)]
pub struct Solver {
    population: Population,
    iterations: usize,
    crossover_prob: f32,
    crossover_point: usize,
    mutation_prob: f32,
}

impl Solver {
    pub fn new(pob_size: usize,
               crom_size: usize,
               iterations: usize,
               crossover_prob: f32,
               crossover_point: usize,
               mutation_prob: f32)
               -> Solver {
        println!("run:");
        println!("Tamaño de la Población: {}", pob_size);
        println!("Tamaño de los Cromosomas: {}", crom_size);
        println!("Cantidad de Iteraciones: {}", iterations);
        println!("Probabilidad de Cruzamiento: {}", crossover_prob);
        println!("Cruzamiento de un Punto - Punto {}", crossover_point);
        println!("Probabilidad de Mutación: {}", mutation_prob);
        println!("Mutación Simple");

        let mut population = Population::new(pob_size, crom_size);
        population.generate();

        Solver {
            population: population,
            iterations: iterations,
            crossover_prob: crossover_prob,
            crossover_point: crossover_point,
            mutation_prob: mutation_prob,
        }
    }

    pub fn evolve(&mut self) {
        for i in 0..self.iterations {
            println!("\nIteración: {}", i);
            self.population.eval();
            self.population.ruleta();
            self.population
                .crossover(self.crossover_prob, self.crossover_point);
            self.population.mutate(self.mutation_prob);
            self.population
                .crossover(self.crossover_prob, self.crossover_point);
            self.population.mutate(self.mutation_prob);
            self.population.selection();
        }

        println!("\nIteración: {}", self.iterations);
        self.population.eval();

        println!("\nFin del Proceso");
    }
}

fn main() {
    let mut solver = Solver::new(6, 8, 100, 0.9, 3, 0.05);
    solver.evolve();
    println!("Hello, world!");
}
<commit_msg>Initial implementation of Practica 15<commit_after>extern crate rand;

use rand::{thread_rng, Rng};

fn get_hsp(w1: &String, w2: &String, match_scr: i32, mism_scr: i32, gap_scr: i32) -> i32 {
    let mut hsp = 0;

    let mut chars2 = w2.chars();
    for c1 in w1.chars() {
        let c2 = chars2.next().unwrap();
        hsp += if c1 == '_' || c2 == '_' {
            gap_scr
        } else if c1 == c2 {
            match_scr
        } else {
            mism_scr
        }
    }

    hsp
}

// fn remove_char(s: &mut String, ch: char) {

// }

fn ver_crossover(s1: &mut String, s2: &mut String, crosspoint: usize) {
    let mut post_s1 = s1.split_off(crosspoint);

    let chars = s1.replace("_", "");

    let mut ch_indx = 0;
    let mut indx = 0;
    while(ch_indx < chars.len()) {
        if s2.chars().nth(indx).unwrap() == chars.chars().nth(ch_indx).unwrap() {
            ch_indx += 1;
        }
        indx += 1;
    }

    let mut post_s2 = s2.split_off(indx);
    *s1 += post_s2.as_ref();
    *s2 += post_s1.as_ref();
}

fn adjust_to_size(len: usize, s: &mut String) {
    let mut s_len = s.len();
    let mut rng = thread_rng();
    while s_len != len {
        let indx = rng.gen_range(0, s_len);
        if s_len > len {
            if s.chars().nth(indx).unwrap() == '_' {
                s.remove(indx);
            }
        } else {
            s.insert_str(indx, "_");
        }
        s_len = s.len();
    }
}

// pub trait Cromosome {
//     fn new(size: usize, input: Vec<String>) -> Cromosome;
//     fn random(&mut self);
//     fn fitness(&self) -> u32;
//     fn crossover(&mut self, parent2: &mut Cromosome, crosspoint: usize);
//     fn mutate(&mut self);
// }

#[derive(Debug)]
pub struct AlignCromosome {
    size: usize,
    genotype: Vec<String>,
}

impl AlignCromosome {
    pub fn new(size: usize, input: Vec<&str>) -> AlignCromosome {
        let mut st_gen = Vec::new();
        for align in input.iter() {
            st_gen.push(String::from(*align));
        }

        let mut rng = thread_rng();
        for align in st_gen.iter_mut() {
            let mut current_len = align.len();
            while current_len < size {
                align.insert_str(rng.gen_range(0, current_len), "_");
                current_len = align.len();
            }
        }

        AlignCromosome {
            size: size,
            genotype: st_gen,
        }
    }

    pub fn random(&mut self) {
        // let mut rng = rand::thread_rng();
        // for _ in 0..self.size {
        //     self.genotype += BINARY[rng.gen::<usize>() % BINARY.len()];
        // }
        unimplemented!()
    }

    pub fn fitness(&self) -> i32 {
        let mut fitness = 0;
        let size = self.genotype.len();
        for i in 0..size {
            for j in i+1..size {
                fitness += get_hsp(&self.genotype[i], &self.genotype[j], 3, 1, 0);
            }
        }

        fitness
    }

    pub fn crossover(&mut self, parent2: &mut AlignCromosome, crosspoint: usize) {
        let len = self.genotype.len();
        if len != parent2.genotype.len() {
            panic!("Can't cross this pair of Cromosomes");
        }

        for i in 0..len {
            ver_crossover(&mut self.genotype[i], &mut parent2.genotype[i], crosspoint);
            adjust_to_size(self.size, &mut self.genotype[i]);
            adjust_to_size(self.size, &mut parent2.genotype[i]);
        }
    }

    pub fn mutate(&mut self) {
        println!("{}", self.genotype);
        let mut rng = rand::thread_rng();
        let pos = rng.gen::<usize>() % self.size;
        self.genotype.remove(pos);
        self.genotype
            .insert_str(pos, BINARY[rng.gen::<usize>() % BINARY.len()]);
        println!("{}", self.genotype);
    }
}

// #[derive(Debug)]
// pub struct Population<Cr> {
//     pob_size: usize,
//     crom_size: usize,
//     population: Vec<AlignCromosome>,
//     ruleta_vec: Vec<f32>,
// }

// impl Population {
//     pub fn new(pob_size: usize, crom_size: usize) -> Population {
//         Population {
//             pob_size: pob_size,
//             crom_size: crom_size,
//             population: Vec::new(),
//             ruleta_vec: Vec::new(),
//         }
//     }

//     pub fn generate(&mut self) {
//         println!("\nGenerando Población inicial");

//         for _ in 0..self.pob_size {
//             let mut crom = Cromosome::new(self.crom_size);
//             crom.random();
//             self.population.push(crom);
//         }

//         for i in 0..self.pob_size {
//             println!("{}) {}", i + 1, self.population[i].genotype);
//         }
//     }

//     pub fn eval(&self) {
//         println!("Evaluando Individuos");
//         for i in 0..self.pob_size {
//             println!("{}) {} - {}",
//                      i + 1,
//                      self.population[i].genotype,
//                      self.population[i].fitness());
//         }
//     }

//     pub fn ruleta(&mut self) {
//         println!("Selección de Individuos - Método de la Ruleta");
//         let mut sum_ruleta = 0;
//         self.ruleta_vec = Vec::new();

//         for i in 0..self.pob_size {
//             sum_ruleta += self.population[i].fitness();
//         }

//         for i in 0..self.pob_size {
//             self.ruleta_vec
//                 .push((self.population[i].fitness() as f32 * 100.0) / sum_ruleta as f32);
//         }

//         for i in 0..self.pob_size {
//             println!("{}) {} - {} - {}",
//                      i + 1,
//                      self.population[i].genotype,
//                      self.population[i].fitness(),
//                      self.ruleta_vec[i]);
//         }
//     }

//     pub fn random_parent_index(&self) -> usize {
//         let mut rng = rand::thread_rng();
//         let num = rng.gen::<f32>() * 100.0;
//         let mut s = 0.0;

//         for i in 0..self.ruleta_vec.len() {
//             s += self.ruleta_vec[i];
//             if num < s {
//                 return i;
//             }
//         }

//         self.ruleta_vec.len()
//     }

//     pub fn crossover(&mut self, crossover_prob: f32, crossover_point: usize) {
//         let parent1 = self.random_parent_index();
//         let mut parent2 = self.random_parent_index();

//         while parent2 == parent1 {
//             let mut rng = rand::thread_rng();
//             parent2 = rng.gen::<usize>() % self.pob_size;
//         }

//         println!("Padre: {}", parent1 + 1);
//         println!("Madre: {}", parent2 + 1);

//         let mut child1 = Cromosome {
//             size: self.population[parent1].size,
//             genotype: self.population[parent1].genotype.clone(),
//         };

//         let mut child2 = Cromosome {
//             size: self.population[parent2].size,
//             genotype: self.population[parent2].genotype.clone(),
//         };

//         let mut rng = rand::thread_rng();
//         if rng.gen::<f32>() < crossover_prob {
//             println!("Cruzamiento");
//             child1.crossover(&mut child2, crossover_point + 1);
//         } else {
//             println!("Sin cruzamiento");
//         }

//         self.population.push(child1);
//         self.population.push(child2);
//     }

//     pub fn mutate(&mut self, chance: f32) {
//         for i in 0..self.population.len() {
//             let mut rng = rand::thread_rng();
//             if rng.gen::<f32>() > chance {
//                 return;
//             }
//             println!("Mutó {}", i + 1);
//             self.population[i].mutate();
//         }
//     }

//     pub fn selection(&mut self) {
//         println!("Selección de Siguiente Población");
//         let mut new_population: Vec<(u32, String)> = Vec::new();
//         for i in 0..self.population.len() {
//             let fitness = self.population[i].fitness();
//             let gen = self.population[i].genotype.clone();
//             println!("{}) {} - {}", i + 1, gen, fitness);
//             new_population.push((fitness, gen));
//         }

//         new_population.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());

//         self.population.clear();
//         for i in 0..self.pob_size {
//             self.population
//                 .push(Cromosome {
//                           size: new_population[i].1.len(),
//                           genotype: new_population[i].1.clone(),
//                       });
//         }

//         // println!("{:?}", new_population);
//     }
// }

// #[derive(Debug)]
// pub struct Solver {
//     population: Population,
//     iterations: usize,
//     crossover_prob: f32,
//     crossover_point: usize,
//     mutation_prob: f32,
// }

// impl Solver {
//     pub fn new(pob_size: usize,
//                crom_size: usize,
//                iterations: usize,
//                crossover_prob: f32,
//                crossover_point: usize,
//                mutation_prob: f32)
//                -> Solver {
//         println!("run:");
//         println!("Tamaño de la Población: {}", pob_size);
//         println!("Tamaño de los Cromosomas: {}", crom_size);
//         println!("Cantidad de Iteraciones: {}", iterations);
//         println!("Probabilidad de Cruzamiento: {}", crossover_prob);
//         println!("Cruzamiento de un Punto - Punto {}", crossover_point);
//         println!("Probabilidad de Mutación: {}", mutation_prob);
//         println!("Mutación Simple");

//         let mut population = Population::new(pob_size, crom_size);
//         population.generate();

//         Solver {
//             population: population,
//             iterations: iterations,
//             crossover_prob: crossover_prob,
//             crossover_point: crossover_point,
//             mutation_prob: mutation_prob,
//         }
//     }

//     pub fn evolve(&mut self) {
//         for i in 0..self.iterations {
//             println!("\nIteración: {}", i);
//             self.population.eval();
//             self.population.ruleta();
//             self.population
//                 .crossover(self.crossover_prob, self.crossover_point);
//             self.population.mutate(self.mutation_prob);
//             self.population
//                 .crossover(self.crossover_prob, self.crossover_point);
//             self.population.mutate(self.mutation_prob);
//             self.population.selection();
//         }

//         println!("\nIteración: {}", self.iterations);
//         self.population.eval();

//         println!("\nFin del Proceso");
//     }
// }

fn main() {
    // let mut solver = Solver::new(6, 8, 100, 0.9, 3, 0.05);
    // solver.evolve();
    println!("Hello, world!");
    let mut cro1 = AlignCromosome::new(9, vec!["ADASDAAC", "AADTTC"]);
    let mut cro2 = AlignCromosome::new(9, vec!["ADASDAAC", "AADTTC"]);
    println!("{:#?}\n{}", cro1, cro1.fitness());
    println!("{:#?}\n{}", cro2, cro1.fitness());
    cro1.crossover(&mut cro2, 4);
    println!("{:#?}\n{}", cro1, cro1.fitness());
    println!("{:#?}\n{}", cro2, cro1.fitness());
    // println!("{:#?}", AlignCromosome::new(22, vec!["AGTCATTAATTGCGGTTAG", "CCAATTGTAGTT", "AGCATTCGTAGTT", "GTTCAAGGGAGTT"]));
}
<|endoftext|>


<commit_before>//- /libcore.rs crate:core
//! Signatures of traits, types and functions from the core lib for use in tests.
pub mod cmp {

    pub trait Ord {
        fn cmp(&self, other: &Self) -> Ordering;
        fn max(self, other: Self) -> Self;
        fn min(self, other: Self) -> Self;
        fn clamp(self, min: Self, max: Self) -> Self;
    }
}

pub mod convert {
    pub trait From<T> {
        fn from(t: T) -> Self;
    }

    pub trait Into<T> {
        pub fn into(self) -> T;
    }
}

pub mod option {
    pub enum Option<T> {
        None,
        Some(T),
    }
}

pub mod prelude {
    pub mod rust_2018 {
        pub use crate::{
            cmp::Ord,
            convert::{From, Into},
            default::Default,
            iter::{IntoIterator, Iterator},
            ops::{Fn, FnMut, FnOnce},
            option::Option::{self, *},
        };
    }
}
#[prelude_import]
pub use prelude::rust_2018::*;
//- /libstd.rs crate:std deps:core
//! Signatures of traits, types and functions from the std lib for use in tests.

/// Docs for return_keyword
mod return_keyword {}

/// Docs for prim_str
mod prim_str {}

pub use core::ops;
<commit_msg>minor: dead code<commit_after>//- /libcore.rs crate:core
//! Signatures of traits, types and functions from the core lib for use in tests.
pub mod cmp {

    pub trait Ord {
        fn cmp(&self, other: &Self) -> Ordering;
        fn max(self, other: Self) -> Self;
        fn min(self, other: Self) -> Self;
        fn clamp(self, min: Self, max: Self) -> Self;
    }
}

pub mod convert {
    pub trait From<T> {
        fn from(t: T) -> Self;
    }

    pub trait Into<T> {
        pub fn into(self) -> T;
    }
}

pub mod prelude {
    pub mod rust_2018 {
        pub use crate::{
            cmp::Ord,
            convert::{From, Into},
            default::Default,
            iter::{IntoIterator, Iterator},
            ops::{Fn, FnMut, FnOnce},
            option::Option::{self, *},
        };
    }
}
#[prelude_import]
pub use prelude::rust_2018::*;
//- /libstd.rs crate:std deps:core
//! Signatures of traits, types and functions from the std lib for use in tests.

/// Docs for return_keyword
mod return_keyword {}

/// Docs for prim_str
mod prim_str {}

pub use core::ops;
<|endoftext|>


<commit_before>//! A Rust interface for Objective-C blocks.
//!
//! For more information on the specifics of the block implementation, see
//! Clang's documentation: http://clang.llvm.org/docs/Block-ABI-Apple.html

use std::mem;
use std::ptr;
use libc::{c_int, c_ulong};

use runtime::Class;
use {Id, Message};

#[allow(improper_ctypes)]
#[link(name = "Foundation", kind = "framework")]
extern {
    static _NSConcreteStackBlock: Class;
}

/// An invoke function for a `Block`; this is the raw C function called by the
/// Objective-C runtime.
type BlockInvoke<A, R> = unsafe extern fn(*mut Block<A, R>, ...) -> R;

/// An invoke function for a `ConcreteBlock`; this is the raw C function called
/// by the Objective-C runtime.
pub type ConcreteBlockInvoke<A, R, C> =
    unsafe extern fn(*mut ConcreteBlock<A, R, C>, ...) -> R;

/// Types that may be used as the arguments to an Objective-C block.
pub trait BlockArguments {
    /// Calls the given `Block` with self as the arguments.
    fn call_block<R>(self, block: &Block<Self, R>) -> R;

    /// Returns an invoke function for a `ConcreteBlock` that takes this type
    /// of arguments.
    fn invoke_for_concrete_block<R, C>() -> ConcreteBlockInvoke<Self, R, C>;
}

macro_rules! block_args_impl(
    ($f:ident $(, $a:ident : $t:ident)*) => (
        impl<$($t),*> BlockArguments for ($($t,)*) {
            fn call_block<R>(self, block: &Block<($($t,)*), R>) -> R {
                let invoke: unsafe extern fn(*mut Block<($($t,)*), R> $(, $t)*) -> R = unsafe {
                    mem::transmute(block.invoke)
                };
                let ($($a,)*) = self;
                let block_ptr = block as *const _ as *mut _;
                unsafe {
                    invoke(block_ptr $(, $a)*)
                }
            }

            fn invoke_for_concrete_block<R, X>() ->
                    ConcreteBlockInvoke<($($t,)*), R, X> {
                unsafe extern fn $f<R, X $(, $t)*>(
                        block_ptr: *mut ConcreteBlock<($($t,)*), R, X>
                        $(, $a: $t)*) -> R {
                    let args = ($($a,)*);
                    let block = &*block_ptr;
                    (block.rust_invoke)(&block.context, args)
                }

                unsafe {
                    mem::transmute($f::<R, X $(, $t)*>)
                }
            }
        }
    );
);

block_args_impl!(concrete_block_invoke_args0);
block_args_impl!(concrete_block_invoke_args1, a: A);
block_args_impl!(concrete_block_invoke_args2, a: A, b: B);
block_args_impl!(concrete_block_invoke_args3, a: A, b: B, c: C);
block_args_impl!(concrete_block_invoke_args4, a: A, b: B, c: C, d: D);
block_args_impl!(concrete_block_invoke_args5, a: A, b: B, c: C, d: D, e: E);
block_args_impl!(concrete_block_invoke_args6, a: A, b: B, c: C, d: D, e: E, f: F);
block_args_impl!(concrete_block_invoke_args7, a: A, b: B, c: C, d: D, e: E, f: F, g: G);
block_args_impl!(concrete_block_invoke_args8, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H);
block_args_impl!(concrete_block_invoke_args9, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I);
block_args_impl!(concrete_block_invoke_args10, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J);
block_args_impl!(concrete_block_invoke_args11, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K);
block_args_impl!(concrete_block_invoke_args12, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K, l: L);

/// An Objective-C block that takes arguments of `A` when called and
/// returns a value of `R`.
#[repr(C)]
pub struct Block<A: BlockArguments, R> {
    isa: *const Class,
    flags: c_int,
    _reserved: c_int,
    invoke: BlockInvoke<A, R>,
}

impl<A: BlockArguments, R> Block<A, R> {
    /// Call self with the given arguments.
    pub fn call(&self, args: A) -> R {
        args.call_block(self)
    }
}

impl<A: BlockArguments, R> Message for Block<A, R> { }

/// An Objective-C block whose size is known at compile time and may be
/// constructed on the stack.
#[repr(C)]
pub struct ConcreteBlock<A: BlockArguments, R, C> {
    base: Block<A, R>,
    descriptor: Box<BlockDescriptor<ConcreteBlock<A, R, C>>>,
    rust_invoke: fn (&C, A) -> R,
    context: C,
}

impl<A: BlockArguments, R, C> ConcreteBlock<A, R, C> {
    /// Constructs a `ConcreteBlock` with the given invoke function and context.
    /// When the block is called, it will return the value that results from
    /// calling the invoke function with a reference to its context.
    pub fn new(invoke: fn (&C, A) -> R, context: C) -> ConcreteBlock<A, R, C> {
        let extern_invoke: ConcreteBlockInvoke<A, R, C> =
            BlockArguments::invoke_for_concrete_block();
        ConcreteBlock {
            base: Block {
                isa: &_NSConcreteStackBlock,
                // 1 << 25 = BLOCK_HAS_COPY_DISPOSE
                flags: 1 << 25,
                _reserved: 0,
                invoke: unsafe { mem::transmute(extern_invoke) },
            },
            descriptor: box BlockDescriptor::<ConcreteBlock<A, R, C>>::new(),
            rust_invoke: invoke,
            context: context,
        }
    }

    /// Copy self onto the heap.
    pub fn copy(self) -> Id<Block<A, R>> {
        unsafe {
            let block = msg_send![&self.base copy] as *mut Block<A, R>;
            // At this point, our copy helper has been run so the block will
            // be moved to the heap and we can forget the original block
            // because the heap block will drop in our dispose helper.
            mem::forget(self);
            Id::from_retained_ptr(block)
        }
    }
}

impl<A: BlockArguments, R, C: Clone> Clone for ConcreteBlock<A, R, C> {
    fn clone(&self) -> ConcreteBlock<A, R, C> {
        ConcreteBlock::new(self.rust_invoke, self.context.clone())
    }
}

impl<A: BlockArguments, R, C> Deref<Block<A, R>> for ConcreteBlock<A, R, C> {
    fn deref(&self) -> &Block<A, R> {
        &self.base
    }
}

unsafe extern fn block_context_dispose<B>(block: &mut B) {
    // Read the block onto the stack and let it drop
    ptr::read(block);
}

unsafe extern fn block_context_copy<B>(_dst: &mut B, _src: &B) {
    // The runtime memmoves the src block into the dst block, nothing to do
}

#[repr(C)]
struct BlockDescriptor<B> {
    _reserved: c_ulong,
    block_size: c_ulong,
    copy_helper: unsafe extern fn(&mut B, &B),
    dispose_helper: unsafe extern fn(&mut B),
}

impl<B> BlockDescriptor<B> {
    fn new() -> BlockDescriptor<B> {
        BlockDescriptor {
            _reserved: 0,
            block_size: mem::size_of::<B>() as c_ulong,
            copy_helper: block_context_copy::<B>,
            dispose_helper: block_context_dispose::<B>,
        }
    }
}

#[cfg(test)]
mod tests {
    use Id;
    use objc_test_utils;
    use super::{Block, ConcreteBlock};

    fn get_int_block() -> &'static Block<(), int> {
        unsafe {
            &*(objc_test_utils::get_int_block() as *const _)
        }
    }

    fn get_int_block_with(i: int) -> Id<Block<(), int>> {
        unsafe {
            let ptr = objc_test_utils::get_int_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn get_add_block() -> &'static Block<(int,), int> {
        unsafe {
            &*(objc_test_utils::get_add_block() as *const _)
        }
    }

    fn get_add_block_with(i: int) -> Id<Block<(int,), int>> {
        unsafe {
            let ptr = objc_test_utils::get_add_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn invoke_int_block(block: &Block<(), int>) -> int {
        let ptr = block as *const _;
        unsafe {
            objc_test_utils::invoke_int_block(ptr as *const _)
        }
    }

    fn invoke_add_block(block: &Block<(int,), int>, a: int) -> int {
        let ptr = block as *const _;
        unsafe {
            objc_test_utils::invoke_add_block(ptr as *const _, a)
        }
    }

    #[test]
    fn test_call_block() {
        let block = get_int_block();
        assert!(block.call(()) == 7);

        let block = get_int_block_with(13);
        assert!(block.call(()) == 13);
    }

    #[test]
    fn test_call_block_args() {
        let block = get_add_block();
        assert!(block.call((2,)) == 9);

        let block = get_add_block_with(13);
        assert!(block.call((2,)) == 15);
    }

    #[test]
    fn test_create_block() {
        fn block_get_int(context: &int, _args: ()) -> int {
            *context
        }

        let block = ConcreteBlock::new(block_get_int, 13);
        let result = invoke_int_block(&*block);
        assert!(result == 13);
    }

    #[test]
    fn test_create_block_args() {
        fn block_add_int(context: &int, (a,): (int,)) -> int {
            a + *context
        }

        let block = ConcreteBlock::new(block_add_int, 5);
        let result = invoke_add_block(&*block, 6);
        assert!(result == 11);
    }

    #[test]
    fn test_concrete_block_clone() {
        fn block_get_string_len(context: &String, _args: ()) -> uint {
            context.len()
        }

        let s = String::from_str("Hello!");
        let expected_len = s.len();
        let block = ConcreteBlock::new(block_get_string_len, s);
        assert!(block.call(()) == expected_len);

        let cloned = block.clone();
        assert!(cloned.call(()) == expected_len);

        drop(block);
        assert!(cloned.call(()) == expected_len);
    }

    #[test]
    fn test_concrete_block_copy() {
        fn block_get_string_len(context: &String, _args: ()) -> uint {
            context.len()
        }

        let s = String::from_str("Hello!");
        let expected_len = s.len();
        let block = ConcreteBlock::new(block_get_string_len, s);
        assert!(block.call(()) == expected_len);

        let copied = block.copy();
        assert!(copied.call(()) == expected_len);
    }
}
<commit_msg>Dried up some block code using Self.<commit_after>//! A Rust interface for Objective-C blocks.
//!
//! For more information on the specifics of the block implementation, see
//! Clang's documentation: http://clang.llvm.org/docs/Block-ABI-Apple.html

use std::mem;
use std::ptr;
use libc::{c_int, c_ulong};

use runtime::Class;
use {Id, Message};

#[allow(improper_ctypes)]
#[link(name = "Foundation", kind = "framework")]
extern {
    static _NSConcreteStackBlock: Class;
}

/// An invoke function for a `Block`; this is the raw C function called by the
/// Objective-C runtime.
type BlockInvoke<A, R> = unsafe extern fn(*mut Block<A, R>, ...) -> R;

/// An invoke function for a `ConcreteBlock`; this is the raw C function called
/// by the Objective-C runtime.
pub type ConcreteBlockInvoke<A, R, C> =
    unsafe extern fn(*mut ConcreteBlock<A, R, C>, ...) -> R;

/// Types that may be used as the arguments to an Objective-C block.
pub trait BlockArguments {
    /// Calls the given `Block` with self as the arguments.
    fn call_block<R>(self, block: &Block<Self, R>) -> R;

    /// Returns an invoke function for a `ConcreteBlock` that takes this type
    /// of arguments.
    fn invoke_for_concrete_block<R, C>() -> ConcreteBlockInvoke<Self, R, C>;
}

macro_rules! block_args_impl(
    ($f:ident $(, $a:ident : $t:ident)*) => (
        impl<$($t),*> BlockArguments for ($($t,)*) {
            fn call_block<R>(self, block: &Block<Self, R>) -> R {
                let invoke: unsafe extern fn(*mut Block<Self, R> $(, $t)*) -> R = unsafe {
                    mem::transmute(block.invoke)
                };
                let ($($a,)*) = self;
                let block_ptr = block as *const _ as *mut _;
                unsafe {
                    invoke(block_ptr $(, $a)*)
                }
            }

            fn invoke_for_concrete_block<R, X>() ->
                    ConcreteBlockInvoke<Self, R, X> {
                unsafe extern fn $f<R, X $(, $t)*>(
                        block_ptr: *mut ConcreteBlock<Self, R, X>
                        $(, $a: $t)*) -> R {
                    let args = ($($a,)*);
                    let block = &*block_ptr;
                    (block.rust_invoke)(&block.context, args)
                }

                unsafe {
                    mem::transmute($f::<R, X $(, $t)*>)
                }
            }
        }
    );
);

block_args_impl!(concrete_block_invoke_args0);
block_args_impl!(concrete_block_invoke_args1, a: A);
block_args_impl!(concrete_block_invoke_args2, a: A, b: B);
block_args_impl!(concrete_block_invoke_args3, a: A, b: B, c: C);
block_args_impl!(concrete_block_invoke_args4, a: A, b: B, c: C, d: D);
block_args_impl!(concrete_block_invoke_args5, a: A, b: B, c: C, d: D, e: E);
block_args_impl!(concrete_block_invoke_args6, a: A, b: B, c: C, d: D, e: E, f: F);
block_args_impl!(concrete_block_invoke_args7, a: A, b: B, c: C, d: D, e: E, f: F, g: G);
block_args_impl!(concrete_block_invoke_args8, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H);
block_args_impl!(concrete_block_invoke_args9, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I);
block_args_impl!(concrete_block_invoke_args10, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J);
block_args_impl!(concrete_block_invoke_args11, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K);
block_args_impl!(concrete_block_invoke_args12, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K, l: L);

/// An Objective-C block that takes arguments of `A` when called and
/// returns a value of `R`.
#[repr(C)]
pub struct Block<A: BlockArguments, R> {
    isa: *const Class,
    flags: c_int,
    _reserved: c_int,
    invoke: BlockInvoke<A, R>,
}

impl<A: BlockArguments, R> Block<A, R> {
    /// Call self with the given arguments.
    pub fn call(&self, args: A) -> R {
        args.call_block(self)
    }
}

impl<A: BlockArguments, R> Message for Block<A, R> { }

/// An Objective-C block whose size is known at compile time and may be
/// constructed on the stack.
#[repr(C)]
pub struct ConcreteBlock<A: BlockArguments, R, C> {
    base: Block<A, R>,
    descriptor: Box<BlockDescriptor<ConcreteBlock<A, R, C>>>,
    rust_invoke: fn (&C, A) -> R,
    context: C,
}

impl<A: BlockArguments, R, C> ConcreteBlock<A, R, C> {
    /// Constructs a `ConcreteBlock` with the given invoke function and context.
    /// When the block is called, it will return the value that results from
    /// calling the invoke function with a reference to its context.
    pub fn new(invoke: fn (&C, A) -> R, context: C) -> Self {
        let extern_invoke: ConcreteBlockInvoke<A, R, C> =
            BlockArguments::invoke_for_concrete_block();
        ConcreteBlock {
            base: Block {
                isa: &_NSConcreteStackBlock,
                // 1 << 25 = BLOCK_HAS_COPY_DISPOSE
                flags: 1 << 25,
                _reserved: 0,
                invoke: unsafe { mem::transmute(extern_invoke) },
            },
            descriptor: box BlockDescriptor::<Self>::new(),
            rust_invoke: invoke,
            context: context,
        }
    }

    /// Copy self onto the heap.
    pub fn copy(self) -> Id<Block<A, R>> {
        unsafe {
            let block = msg_send![&self.base copy] as *mut Block<A, R>;
            // At this point, our copy helper has been run so the block will
            // be moved to the heap and we can forget the original block
            // because the heap block will drop in our dispose helper.
            mem::forget(self);
            Id::from_retained_ptr(block)
        }
    }
}

impl<A: BlockArguments, R, C: Clone> Clone for ConcreteBlock<A, R, C> {
    fn clone(&self) -> Self {
        ConcreteBlock::new(self.rust_invoke, self.context.clone())
    }
}

impl<A: BlockArguments, R, C> Deref<Block<A, R>> for ConcreteBlock<A, R, C> {
    fn deref(&self) -> &Block<A, R> {
        &self.base
    }
}

unsafe extern fn block_context_dispose<B>(block: &mut B) {
    // Read the block onto the stack and let it drop
    ptr::read(block);
}

unsafe extern fn block_context_copy<B>(_dst: &mut B, _src: &B) {
    // The runtime memmoves the src block into the dst block, nothing to do
}

#[repr(C)]
struct BlockDescriptor<B> {
    _reserved: c_ulong,
    block_size: c_ulong,
    copy_helper: unsafe extern fn(&mut B, &B),
    dispose_helper: unsafe extern fn(&mut B),
}

impl<B> BlockDescriptor<B> {
    fn new() -> BlockDescriptor<B> {
        BlockDescriptor {
            _reserved: 0,
            block_size: mem::size_of::<B>() as c_ulong,
            copy_helper: block_context_copy::<B>,
            dispose_helper: block_context_dispose::<B>,
        }
    }
}

#[cfg(test)]
mod tests {
    use Id;
    use objc_test_utils;
    use super::{Block, ConcreteBlock};

    fn get_int_block() -> &'static Block<(), int> {
        unsafe {
            &*(objc_test_utils::get_int_block() as *const _)
        }
    }

    fn get_int_block_with(i: int) -> Id<Block<(), int>> {
        unsafe {
            let ptr = objc_test_utils::get_int_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn get_add_block() -> &'static Block<(int,), int> {
        unsafe {
            &*(objc_test_utils::get_add_block() as *const _)
        }
    }

    fn get_add_block_with(i: int) -> Id<Block<(int,), int>> {
        unsafe {
            let ptr = objc_test_utils::get_add_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn invoke_int_block(block: &Block<(), int>) -> int {
        let ptr = block as *const _;
        unsafe {
            objc_test_utils::invoke_int_block(ptr as *const _)
        }
    }

    fn invoke_add_block(block: &Block<(int,), int>, a: int) -> int {
        let ptr = block as *const _;
        unsafe {
            objc_test_utils::invoke_add_block(ptr as *const _, a)
        }
    }

    #[test]
    fn test_call_block() {
        let block = get_int_block();
        assert!(block.call(()) == 7);

        let block = get_int_block_with(13);
        assert!(block.call(()) == 13);
    }

    #[test]
    fn test_call_block_args() {
        let block = get_add_block();
        assert!(block.call((2,)) == 9);

        let block = get_add_block_with(13);
        assert!(block.call((2,)) == 15);
    }

    #[test]
    fn test_create_block() {
        fn block_get_int(context: &int, _args: ()) -> int {
            *context
        }

        let block = ConcreteBlock::new(block_get_int, 13);
        let result = invoke_int_block(&*block);
        assert!(result == 13);
    }

    #[test]
    fn test_create_block_args() {
        fn block_add_int(context: &int, (a,): (int,)) -> int {
            a + *context
        }

        let block = ConcreteBlock::new(block_add_int, 5);
        let result = invoke_add_block(&*block, 6);
        assert!(result == 11);
    }

    #[test]
    fn test_concrete_block_clone() {
        fn block_get_string_len(context: &String, _args: ()) -> uint {
            context.len()
        }

        let s = String::from_str("Hello!");
        let expected_len = s.len();
        let block = ConcreteBlock::new(block_get_string_len, s);
        assert!(block.call(()) == expected_len);

        let cloned = block.clone();
        assert!(cloned.call(()) == expected_len);

        drop(block);
        assert!(cloned.call(()) == expected_len);
    }

    #[test]
    fn test_concrete_block_copy() {
        fn block_get_string_len(context: &String, _args: ()) -> uint {
            context.len()
        }

        let s = String::from_str("Hello!");
        let expected_len = s.len();
        let block = ConcreteBlock::new(block_get_string_len, s);
        assert!(block.call(()) == expected_len);

        let copied = block.copy();
        assert!(copied.call(()) == expected_len);
    }
}
<|endoftext|>


<commit_before>use termcolor::{Color, ColorSpec};

use crate::diagnostic::{LabelStyle, Severity};

/// Configures how a diagnostic is rendered.
#[derive(Clone, Debug)]
pub struct Config {
    /// The display style to use when rendering diagnostics.
    /// Defaults to: [`DisplayStyle::Rich`].
    ///
    /// [`DisplayStyle::Rich`]: DisplayStyle::Rich
    pub display_style: DisplayStyle,
    /// Column width of tabs.
    /// Defaults to: `4`.
    pub tab_width: usize,
    /// Styles to use when rendering the diagnostic.
    pub styles: Styles,
    /// Characters to use when rendering the diagnostic.
    pub chars: Chars,
    /// The minimum number of lines to be shown after the line on which a multiline [`Label`] begins.
    ///
    /// Defaults to: `3`.
    pub start_context_lines: usize,
    /// The minimum number of lines to be shown before the line on which a multiline [`Label`] ends.
    ///
    /// Defaults to: `1`.
    pub end_context_lines: usize,
}

impl Default for Config {
    fn default() -> Config {
        Config {
            display_style: DisplayStyle::Rich,
            tab_width: 4,
            styles: Styles::default(),
            chars: Chars::default(),
            start_context_lines: 3,
            end_context_lines: 1,
        }
    }
}

/// The display style to use when rendering diagnostics.
#[derive(Clone, Debug)]
pub enum DisplayStyle {
    /// Output a richly formatted diagnostic, with source code previews.
    ///
    /// ```text
    /// error[E0001]: unexpected type in `+` application
    ///   ┌─ test:2:9
    ///   │
    /// 2 │ (+ test "")
    ///   │         ^^ expected `Int` but found `String`
    ///   │
    ///   = expected type `Int`
    ///        found type `String`
    ///
    /// error[E0002]: Bad config found
    ///
    /// ```
    Rich,
    /// Output a condensed diagnostic, with a line number, severity, message and notes (if any).
    ///
    /// ```text
    /// test:2:9: error[E0001]: unexpected type in `+` application
    /// = expected type `Int`
    ///      found type `String`
    ///
    /// error[E0002]: Bad config found
    /// ```
    Medium,
    /// Output a short diagnostic, with a line number, severity, and message.
    ///
    /// ```text
    /// test:2:9: error[E0001]: unexpected type in `+` application
    /// error[E0002]: Bad config found
    /// ```
    Short,
}

/// Styles to use when rendering the diagnostic.
#[derive(Clone, Debug)]
pub struct Styles {
    /// The style to use when rendering bug headers.
    /// Defaults to `fg:red bold intense`.
    pub header_bug: ColorSpec,
    /// The style to use when rendering error headers.
    /// Defaults to `fg:red bold intense`.
    pub header_error: ColorSpec,
    /// The style to use when rendering warning headers.
    /// Defaults to `fg:yellow bold intense`.
    pub header_warning: ColorSpec,
    /// The style to use when rendering note headers.
    /// Defaults to `fg:green bold intense`.
    pub header_note: ColorSpec,
    /// The style to use when rendering help headers.
    /// Defaults to `fg:cyan bold intense`.
    pub header_help: ColorSpec,
    /// The style to use when the main diagnostic message.
    /// Defaults to `bold intense`.
    pub header_message: ColorSpec,

    /// The style to use when rendering bug labels.
    /// Defaults to `fg:red`.
    pub primary_label_bug: ColorSpec,
    /// The style to use when rendering error labels.
    /// Defaults to `fg:red`.
    pub primary_label_error: ColorSpec,
    /// The style to use when rendering warning labels.
    /// Defaults to `fg:yellow`.
    pub primary_label_warning: ColorSpec,
    /// The style to use when rendering note labels.
    /// Defaults to `fg:green`.
    pub primary_label_note: ColorSpec,
    /// The style to use when rendering help labels.
    /// Defaults to `fg:cyan`.
    pub primary_label_help: ColorSpec,
    /// The style to use when rendering secondary labels.
    /// Defaults `fg:blue` (or `fg:cyan` on windows).
    pub secondary_label: ColorSpec,

    /// The style to use when rendering the line numbers.
    /// Defaults `fg:blue` (or `fg:cyan` on windows).
    pub line_number: ColorSpec,
    /// The style to use when rendering the source code borders.
    /// Defaults `fg:blue` (or `fg:cyan` on windows).
    pub source_border: ColorSpec,
    /// The style to use when rendering the note bullets.
    /// Defaults `fg:blue` (or `fg:cyan` on windows).
    pub note_bullet: ColorSpec,
}

impl Styles {
    /// The style used to mark a header at a given severity.
    pub fn header(&self, severity: Severity) -> &ColorSpec {
        match severity {
            Severity::Bug => &self.header_bug,
            Severity::Error => &self.header_error,
            Severity::Warning => &self.header_warning,
            Severity::Note => &self.header_note,
            Severity::Help => &self.header_help,
        }
    }

    /// The style used to mark a primary or secondary label at a given severity.
    pub fn label(&self, severity: Severity, label_style: LabelStyle) -> &ColorSpec {
        match (label_style, severity) {
            (LabelStyle::Primary, Severity::Bug) => &self.primary_label_bug,
            (LabelStyle::Primary, Severity::Error) => &self.primary_label_error,
            (LabelStyle::Primary, Severity::Warning) => &self.primary_label_warning,
            (LabelStyle::Primary, Severity::Note) => &self.primary_label_note,
            (LabelStyle::Primary, Severity::Help) => &self.primary_label_help,
            (LabelStyle::Secondary, _) => &self.secondary_label,
        }
    }

    #[doc(hidden)]
    pub fn with_blue(blue: Color) -> Styles {
        let header = ColorSpec::new().set_bold(true).set_intense(true).clone();

        Styles {
            header_bug: header.clone().set_fg(Some(Color::Red)).clone(),
            header_error: header.clone().set_fg(Some(Color::Red)).clone(),
            header_warning: header.clone().set_fg(Some(Color::Yellow)).clone(),
            header_note: header.clone().set_fg(Some(Color::Green)).clone(),
            header_help: header.clone().set_fg(Some(Color::Cyan)).clone(),
            header_message: header,

            primary_label_bug: ColorSpec::new().set_fg(Some(Color::Red)).clone(),
            primary_label_error: ColorSpec::new().set_fg(Some(Color::Red)).clone(),
            primary_label_warning: ColorSpec::new().set_fg(Some(Color::Yellow)).clone(),
            primary_label_note: ColorSpec::new().set_fg(Some(Color::Green)).clone(),
            primary_label_help: ColorSpec::new().set_fg(Some(Color::Cyan)).clone(),
            secondary_label: ColorSpec::new().set_fg(Some(blue)).clone(),

            line_number: ColorSpec::new().set_fg(Some(blue)).clone(),
            source_border: ColorSpec::new().set_fg(Some(blue)).clone(),
            note_bullet: ColorSpec::new().set_fg(Some(blue)).clone(),
        }
    }
}

impl Default for Styles {
    fn default() -> Styles {
        // Blue is really difficult to see on the standard windows command line
        #[cfg(windows)]
        const BLUE: Color = Color::Cyan;
        #[cfg(not(windows))]
        const BLUE: Color = Color::Blue;

        Self::with_blue(BLUE)
    }
}

/// Characters to use when rendering the diagnostic.
///
/// By using [`Chars::ascii()`] you can switch to an ASCII-only format suitable
/// for rendering on terminals that do not support box drawing characters.
#[derive(Clone, Debug)]
pub struct Chars {
    /// The characters to use for the top-left border of the snippet.
    /// Defaults to: `"┌─"` or `"-->"` with [`Chars::ascii()`].
    pub snippet_start: String,
    /// The character to use for the left border of the source.
    /// Defaults to: `'│'` or `'|'` with [`Chars::ascii()`].
    pub source_border_left: char,
    /// The character to use for the left border break of the source.
    /// Defaults to: `'·'` or `'.'` with [`Chars::ascii()`].
    pub source_border_left_break: char,

    /// The character to use for the note bullet.
    /// Defaults to: `'='`.
    pub note_bullet: char,

    /// The character to use for marking a single-line primary label.
    /// Defaults to: `'^'`.
    pub single_primary_caret: char,
    /// The character to use for marking a single-line secondary label.
    /// Defaults to: `'-'`.
    pub single_secondary_caret: char,

    /// The character to use for marking the start of a multi-line primary label.
    /// Defaults to: `'^'`.
    pub multi_primary_caret_start: char,
    /// The character to use for marking the end of a multi-line primary label.
    /// Defaults to: `'^'`.
    pub multi_primary_caret_end: char,
    /// The character to use for marking the start of a multi-line secondary label.
    /// Defaults to: `'\''`.
    pub multi_secondary_caret_start: char,
    /// The character to use for marking the end of a multi-line secondary label.
    /// Defaults to: `'\''`.
    pub multi_secondary_caret_end: char,
    /// The character to use for the top-left corner of a multi-line label.
    /// Defaults to: `'╭'` or `'/'` with [`Chars::ascii()`].
    pub multi_top_left: char,
    /// The character to use for the top of a multi-line label.
    /// Defaults to: `'─'` or `'-'` with [`Chars::ascii()`].
    pub multi_top: char,
    /// The character to use for the bottom-left corner of a multi-line label.
    /// Defaults to: `'╰'` or `'\\'` with [`Chars::ascii()`].
    pub multi_bottom_left: char,
    /// The character to use when marking the bottom of a multi-line label.
    /// Defaults to: `'─'` or `'-'` with [`Chars::ascii()`].
    pub multi_bottom: char,
    /// The character to use for the left of a multi-line label.
    /// Defaults to: `'│'` or `'|'` with [`Chars::ascii()`].
    pub multi_left: char,

    /// The character to use for the left of a pointer underneath a caret.
    /// Defaults to: `'│'` or `'|'` with [`Chars::ascii()`].
    pub pointer_left: char,
}

impl Default for Chars {
    fn default() -> Chars {
        Chars::box_drawing()
    }
}

impl Chars {
    /// A character set that uses Unicode box drawing characters.
    pub fn box_drawing() -> Chars {
        Chars {
            snippet_start: "┌─".into(),
            source_border_left: '│',
            source_border_left_break: '·',

            note_bullet: '=',

            single_primary_caret: '^',
            single_secondary_caret: '-',

            multi_primary_caret_start: '^',
            multi_primary_caret_end: '^',
            multi_secondary_caret_start: '\'',
            multi_secondary_caret_end: '\'',
            multi_top_left: '╭',
            multi_top: '─',
            multi_bottom_left: '╰',
            multi_bottom: '─',
            multi_left: '│',

            pointer_left: '│',
        }
    }

    /// A character set that only uses ASCII characters.
    ///
    /// This is useful if your terminal's font does not support box drawing
    /// characters well and results in output that looks similar to rustc's
    /// diagnostic output.
    pub fn ascii() -> Chars {
        Chars {
            snippet_start: "-->".into(),
            source_border_left: '|',
            source_border_left_break: '.',

            note_bullet: '=',

            single_primary_caret: '^',
            single_secondary_caret: '-',

            multi_primary_caret_start: '^',
            multi_primary_caret_end: '^',
            multi_secondary_caret_start: '\'',
            multi_secondary_caret_end: '\'',
            multi_top_left: '/',
            multi_top: '-',
            multi_bottom_left: '\\',
            multi_bottom: '-',
            multi_left: '|',

            pointer_left: '|',
        }
    }
}
<commit_msg>fix documentation: link, backslash not escaped<commit_after>use termcolor::{Color, ColorSpec};

use crate::diagnostic::{LabelStyle, Severity};

/// Configures how a diagnostic is rendered.
#[derive(Clone, Debug)]
pub struct Config {
    /// The display style to use when rendering diagnostics.
    /// Defaults to: [`DisplayStyle::Rich`].
    ///
    /// [`DisplayStyle::Rich`]: DisplayStyle::Rich
    pub display_style: DisplayStyle,
    /// Column width of tabs.
    /// Defaults to: `4`.
    pub tab_width: usize,
    /// Styles to use when rendering the diagnostic.
    pub styles: Styles,
    /// Characters to use when rendering the diagnostic.
    pub chars: Chars,
    /// The minimum number of lines to be shown after the line on which a multiline [`Label`] begins.
    ///
    /// Defaults to: `3`.
    ///
    /// [`Label`]: crate::diagnostic::Label
    pub start_context_lines: usize,
    /// The minimum number of lines to be shown before the line on which a multiline [`Label`] ends.
    ///
    /// Defaults to: `1`.
    ///
    /// [`Label`]: crate::diagnostic::Label
    pub end_context_lines: usize,
}

impl Default for Config {
    fn default() -> Config {
        Config {
            display_style: DisplayStyle::Rich,
            tab_width: 4,
            styles: Styles::default(),
            chars: Chars::default(),
            start_context_lines: 3,
            end_context_lines: 1,
        }
    }
}

/// The display style to use when rendering diagnostics.
#[derive(Clone, Debug)]
pub enum DisplayStyle {
    /// Output a richly formatted diagnostic, with source code previews.
    ///
    /// ```text
    /// error[E0001]: unexpected type in `+` application
    ///   ┌─ test:2:9
    ///   │
    /// 2 │ (+ test "")
    ///   │         ^^ expected `Int` but found `String`
    ///   │
    ///   = expected type `Int`
    ///        found type `String`
    ///
    /// error[E0002]: Bad config found
    ///
    /// ```
    Rich,
    /// Output a condensed diagnostic, with a line number, severity, message and notes (if any).
    ///
    /// ```text
    /// test:2:9: error[E0001]: unexpected type in `+` application
    /// = expected type `Int`
    ///      found type `String`
    ///
    /// error[E0002]: Bad config found
    /// ```
    Medium,
    /// Output a short diagnostic, with a line number, severity, and message.
    ///
    /// ```text
    /// test:2:9: error[E0001]: unexpected type in `+` application
    /// error[E0002]: Bad config found
    /// ```
    Short,
}

/// Styles to use when rendering the diagnostic.
#[derive(Clone, Debug)]
pub struct Styles {
    /// The style to use when rendering bug headers.
    /// Defaults to `fg:red bold intense`.
    pub header_bug: ColorSpec,
    /// The style to use when rendering error headers.
    /// Defaults to `fg:red bold intense`.
    pub header_error: ColorSpec,
    /// The style to use when rendering warning headers.
    /// Defaults to `fg:yellow bold intense`.
    pub header_warning: ColorSpec,
    /// The style to use when rendering note headers.
    /// Defaults to `fg:green bold intense`.
    pub header_note: ColorSpec,
    /// The style to use when rendering help headers.
    /// Defaults to `fg:cyan bold intense`.
    pub header_help: ColorSpec,
    /// The style to use when the main diagnostic message.
    /// Defaults to `bold intense`.
    pub header_message: ColorSpec,

    /// The style to use when rendering bug labels.
    /// Defaults to `fg:red`.
    pub primary_label_bug: ColorSpec,
    /// The style to use when rendering error labels.
    /// Defaults to `fg:red`.
    pub primary_label_error: ColorSpec,
    /// The style to use when rendering warning labels.
    /// Defaults to `fg:yellow`.
    pub primary_label_warning: ColorSpec,
    /// The style to use when rendering note labels.
    /// Defaults to `fg:green`.
    pub primary_label_note: ColorSpec,
    /// The style to use when rendering help labels.
    /// Defaults to `fg:cyan`.
    pub primary_label_help: ColorSpec,
    /// The style to use when rendering secondary labels.
    /// Defaults `fg:blue` (or `fg:cyan` on windows).
    pub secondary_label: ColorSpec,

    /// The style to use when rendering the line numbers.
    /// Defaults `fg:blue` (or `fg:cyan` on windows).
    pub line_number: ColorSpec,
    /// The style to use when rendering the source code borders.
    /// Defaults `fg:blue` (or `fg:cyan` on windows).
    pub source_border: ColorSpec,
    /// The style to use when rendering the note bullets.
    /// Defaults `fg:blue` (or `fg:cyan` on windows).
    pub note_bullet: ColorSpec,
}

impl Styles {
    /// The style used to mark a header at a given severity.
    pub fn header(&self, severity: Severity) -> &ColorSpec {
        match severity {
            Severity::Bug => &self.header_bug,
            Severity::Error => &self.header_error,
            Severity::Warning => &self.header_warning,
            Severity::Note => &self.header_note,
            Severity::Help => &self.header_help,
        }
    }

    /// The style used to mark a primary or secondary label at a given severity.
    pub fn label(&self, severity: Severity, label_style: LabelStyle) -> &ColorSpec {
        match (label_style, severity) {
            (LabelStyle::Primary, Severity::Bug) => &self.primary_label_bug,
            (LabelStyle::Primary, Severity::Error) => &self.primary_label_error,
            (LabelStyle::Primary, Severity::Warning) => &self.primary_label_warning,
            (LabelStyle::Primary, Severity::Note) => &self.primary_label_note,
            (LabelStyle::Primary, Severity::Help) => &self.primary_label_help,
            (LabelStyle::Secondary, _) => &self.secondary_label,
        }
    }

    #[doc(hidden)]
    pub fn with_blue(blue: Color) -> Styles {
        let header = ColorSpec::new().set_bold(true).set_intense(true).clone();

        Styles {
            header_bug: header.clone().set_fg(Some(Color::Red)).clone(),
            header_error: header.clone().set_fg(Some(Color::Red)).clone(),
            header_warning: header.clone().set_fg(Some(Color::Yellow)).clone(),
            header_note: header.clone().set_fg(Some(Color::Green)).clone(),
            header_help: header.clone().set_fg(Some(Color::Cyan)).clone(),
            header_message: header,

            primary_label_bug: ColorSpec::new().set_fg(Some(Color::Red)).clone(),
            primary_label_error: ColorSpec::new().set_fg(Some(Color::Red)).clone(),
            primary_label_warning: ColorSpec::new().set_fg(Some(Color::Yellow)).clone(),
            primary_label_note: ColorSpec::new().set_fg(Some(Color::Green)).clone(),
            primary_label_help: ColorSpec::new().set_fg(Some(Color::Cyan)).clone(),
            secondary_label: ColorSpec::new().set_fg(Some(blue)).clone(),

            line_number: ColorSpec::new().set_fg(Some(blue)).clone(),
            source_border: ColorSpec::new().set_fg(Some(blue)).clone(),
            note_bullet: ColorSpec::new().set_fg(Some(blue)).clone(),
        }
    }
}

impl Default for Styles {
    fn default() -> Styles {
        // Blue is really difficult to see on the standard windows command line
        #[cfg(windows)]
        const BLUE: Color = Color::Cyan;
        #[cfg(not(windows))]
        const BLUE: Color = Color::Blue;

        Self::with_blue(BLUE)
    }
}

/// Characters to use when rendering the diagnostic.
///
/// By using [`Chars::ascii()`] you can switch to an ASCII-only format suitable
/// for rendering on terminals that do not support box drawing characters.
#[derive(Clone, Debug)]
pub struct Chars {
    /// The characters to use for the top-left border of the snippet.
    /// Defaults to: `"┌─"` or `"-->"` with [`Chars::ascii()`].
    pub snippet_start: String,
    /// The character to use for the left border of the source.
    /// Defaults to: `'│'` or `'|'` with [`Chars::ascii()`].
    pub source_border_left: char,
    /// The character to use for the left border break of the source.
    /// Defaults to: `'·'` or `'.'` with [`Chars::ascii()`].
    pub source_border_left_break: char,

    /// The character to use for the note bullet.
    /// Defaults to: `'='`.
    pub note_bullet: char,

    /// The character to use for marking a single-line primary label.
    /// Defaults to: `'^'`.
    pub single_primary_caret: char,
    /// The character to use for marking a single-line secondary label.
    /// Defaults to: `'-'`.
    pub single_secondary_caret: char,

    /// The character to use for marking the start of a multi-line primary label.
    /// Defaults to: `'^'`.
    pub multi_primary_caret_start: char,
    /// The character to use for marking the end of a multi-line primary label.
    /// Defaults to: `'^'`.
    pub multi_primary_caret_end: char,
    /// The character to use for marking the start of a multi-line secondary label.
    /// Defaults to: `'\''`.
    pub multi_secondary_caret_start: char,
    /// The character to use for marking the end of a multi-line secondary label.
    /// Defaults to: `'\''`.
    pub multi_secondary_caret_end: char,
    /// The character to use for the top-left corner of a multi-line label.
    /// Defaults to: `'╭'` or `'/'` with [`Chars::ascii()`].
    pub multi_top_left: char,
    /// The character to use for the top of a multi-line label.
    /// Defaults to: `'─'` or `'-'` with [`Chars::ascii()`].
    pub multi_top: char,
    /// The character to use for the bottom-left corner of a multi-line label.
    /// Defaults to: `'╰'` or `'\'` with [`Chars::ascii()`].
    pub multi_bottom_left: char,
    /// The character to use when marking the bottom of a multi-line label.
    /// Defaults to: `'─'` or `'-'` with [`Chars::ascii()`].
    pub multi_bottom: char,
    /// The character to use for the left of a multi-line label.
    /// Defaults to: `'│'` or `'|'` with [`Chars::ascii()`].
    pub multi_left: char,

    /// The character to use for the left of a pointer underneath a caret.
    /// Defaults to: `'│'` or `'|'` with [`Chars::ascii()`].
    pub pointer_left: char,
}

impl Default for Chars {
    fn default() -> Chars {
        Chars::box_drawing()
    }
}

impl Chars {
    /// A character set that uses Unicode box drawing characters.
    pub fn box_drawing() -> Chars {
        Chars {
            snippet_start: "┌─".into(),
            source_border_left: '│',
            source_border_left_break: '·',

            note_bullet: '=',

            single_primary_caret: '^',
            single_secondary_caret: '-',

            multi_primary_caret_start: '^',
            multi_primary_caret_end: '^',
            multi_secondary_caret_start: '\'',
            multi_secondary_caret_end: '\'',
            multi_top_left: '╭',
            multi_top: '─',
            multi_bottom_left: '╰',
            multi_bottom: '─',
            multi_left: '│',

            pointer_left: '│',
        }
    }

    /// A character set that only uses ASCII characters.
    ///
    /// This is useful if your terminal's font does not support box drawing
    /// characters well and results in output that looks similar to rustc's
    /// diagnostic output.
    pub fn ascii() -> Chars {
        Chars {
            snippet_start: "-->".into(),
            source_border_left: '|',
            source_border_left_break: '.',

            note_bullet: '=',

            single_primary_caret: '^',
            single_secondary_caret: '-',

            multi_primary_caret_start: '^',
            multi_primary_caret_end: '^',
            multi_secondary_caret_start: '\'',
            multi_secondary_caret_end: '\'',
            multi_top_left: '/',
            multi_top: '-',
            multi_bottom_left: '\\',
            multi_bottom: '-',
            multi_left: '|',

            pointer_left: '|',
        }
    }
}
<|endoftext|>