<commit_before>mod memory;
mod stack;
mod pc;
mod storage;
mod params;
mod eval;
mod commit;
pub mod errors;

pub use self::memory::{Memory, SeqMemory};
pub use self::stack::Stack;
pub use self::pc::{PC, Instruction};
pub use self::storage::{Storage, HashMapStorage};
pub use self::params::{Context, BlockHeader, Log, Patch};
pub use self::eval::{State, Machine, MachineStatus};
pub use self::commit::{AccountCommitment, Account};

use utils::bigint::M256;
use self::errors::{RequireError, CommitError, MachineError};

pub struct VM<M, S>(Vec<Machine<M, S>>);

#[derive(Debug, Clone)]
pub enum VMStatus {
    Running,
    ExitedOk,
    ExitedErr(MachineError),
}

impl<M: Memory + Default, S: Storage + Default + Clone> VM<M, S> {
    pub fn commit_account(&mut self, commitment: AccountCommitment<S>) -> Result<(), CommitError> {
        for machine in &mut self.0 {
            machine.commit_account(commitment.clone())?;
        }
        Ok(())
    }

    pub fn commit_blockhash(&mut self, number: M256, hash: M256) -> Result<(), CommitError> {
        for machine in &mut self.0 {
            machine.commit_blockhash(number, hash)?;
        }
        Ok(())
    }

    pub fn status(&self) -> VMStatus {
        match self.0[0].status() {
            MachineStatus::Running | MachineStatus::InvokeCall(_, _) => VMStatus::Running,
            MachineStatus::ExitedOk => VMStatus::ExitedOk,
            MachineStatus::ExitedErr(err) => VMStatus::ExitedErr(err),
        }
    }

    pub fn step(&mut self) -> Result<(), RequireError> {
        match self.0.last().unwrap().status().clone() {
            MachineStatus::Running => {
                self.0.last_mut().unwrap().step()
            },
            MachineStatus::ExitedOk | MachineStatus::ExitedErr(_) => {
                if self.0.len() <= 1 {
                    Ok(())
                } else {
                    let finished = self.0.pop().unwrap();
                    self.0.last_mut().unwrap().apply_sub(finished);
                    Ok(())
                }
            },
            MachineStatus::InvokeCall(context, _) => {
                let sub = self.0.last().unwrap().derive(context);
                self.0.push(sub);
                Ok(())
            },
        }
    }
}
<commit_msg>Problem: fire is not implemented Solution: create the fire impl for VM<commit_after>mod memory;
mod stack;
mod pc;
mod storage;
mod params;
mod eval;
mod commit;
pub mod errors;

pub use self::memory::{Memory, SeqMemory};
pub use self::stack::Stack;
pub use self::pc::{PC, Instruction};
pub use self::storage::{Storage, HashMapStorage};
pub use self::params::{Context, BlockHeader, Log, Patch};
pub use self::eval::{State, Machine, MachineStatus};
pub use self::commit::{AccountCommitment, Account};

use utils::bigint::M256;
use self::errors::{RequireError, CommitError, MachineError};

pub struct VM<M, S>(Vec<Machine<M, S>>);

#[derive(Debug, Clone)]
pub enum VMStatus {
    Running,
    ExitedOk,
    ExitedErr(MachineError),
}

impl<M: Memory + Default, S: Storage + Default + Clone> VM<M, S> {
    pub fn commit_account(&mut self, commitment: AccountCommitment<S>) -> Result<(), CommitError> {
        for machine in &mut self.0 {
            machine.commit_account(commitment.clone())?;
        }
        Ok(())
    }

    pub fn commit_blockhash(&mut self, number: M256, hash: M256) -> Result<(), CommitError> {
        for machine in &mut self.0 {
            machine.commit_blockhash(number, hash)?;
        }
        Ok(())
    }

    pub fn status(&self) -> VMStatus {
        match self.0[0].status() {
            MachineStatus::Running | MachineStatus::InvokeCall(_, _) => VMStatus::Running,
            MachineStatus::ExitedOk => VMStatus::ExitedOk,
            MachineStatus::ExitedErr(err) => VMStatus::ExitedErr(err),
        }
    }

    pub fn step(&mut self) -> Result<(), RequireError> {
        match self.0.last().unwrap().status().clone() {
            MachineStatus::Running => {
                self.0.last_mut().unwrap().step()
            },
            MachineStatus::ExitedOk | MachineStatus::ExitedErr(_) => {
                if self.0.len() <= 1 {
                    Ok(())
                } else {
                    let finished = self.0.pop().unwrap();
                    self.0.last_mut().unwrap().apply_sub(finished);
                    Ok(())
                }
            },
            MachineStatus::InvokeCall(context, _) => {
                let sub = self.0.last().unwrap().derive(context);
                self.0.push(sub);
                Ok(())
            },
        }
    }

    pub fn fire(&mut self) -> Result<(), RequireError> {
        loop {
            match self.status() {
                VMStatus::Running => self.step()?,
                VMStatus::ExitedOk | VMStatus::ExitedErr(_) => return Ok(()),
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::{ptr, mem, time, io};
use std::rc::Rc;
use std::net::{SocketAddr, Shutdown};

use bytes::{Bytes, BytesMut, Buf, BufMut};
use futures::{Future, Poll, Async};
use tokio_io::{AsyncRead, AsyncWrite};

use super::{h1, h2, utils, HttpHandler, IoStream};
use super::settings::WorkerSettings;

const HTTP2_PREFACE: [u8; 14] = *b"PRI * HTTP/2.0";

enum HttpProtocol<T: IoStream, H: 'static> {
    H1(h1::Http1<T, H>),
    H2(h2::Http2<T, H>),
    Unknown(Rc<WorkerSettings<H>>, Option<SocketAddr>, T, BytesMut),
}

enum ProtocolKind {
    Http1,
    Http2,
}

#[doc(hidden)]
pub struct HttpChannel<T, H> where T: IoStream, H: HttpHandler + 'static {
    proto: Option<HttpProtocol<T, H>>,
    node: Option<Node<HttpChannel<T, H>>>,
}

impl<T, H> HttpChannel<T, H> where T: IoStream, H: HttpHandler + 'static
{
    pub(crate) fn new(settings: Rc<WorkerSettings<H>>,
                      io: T, peer: Option<SocketAddr>, http2: bool) -> HttpChannel<T, H>
    {
        settings.add_channel();

        if http2 {
            HttpChannel {
                node: None, proto: Some(HttpProtocol::H2(
                    h2::Http2::new(settings, io, peer, Bytes::new()))) }
        } else {
            HttpChannel {
                node: None, proto: Some(HttpProtocol::Unknown(
                    settings, peer, io, BytesMut::with_capacity(4096))) }
        }
    }

    fn shutdown(&mut self) {
        match self.proto {
            Some(HttpProtocol::H1(ref mut h1)) => {
                let io = h1.io();
                let _ = IoStream::set_linger(io, Some(time::Duration::new(0, 0)));
                let _ = IoStream::shutdown(io, Shutdown::Both);
            }
            Some(HttpProtocol::H2(ref mut h2)) => {
                h2.shutdown()
            }
            _ => (),
        }
    }
}

impl<T, H> Future for HttpChannel<T, H> where T: IoStream, H: HttpHandler + 'static
{
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        if !self.node.is_none() {
            let el = self as *mut _;
            self.node = Some(Node::new(el));
            let _ = match self.proto {
                Some(HttpProtocol::H1(ref mut h1)) =>
                    self.node.as_ref().map(|n| h1.settings().head().insert(n)),
                Some(HttpProtocol::H2(ref mut h2)) =>
                    self.node.as_ref().map(|n| h2.settings().head().insert(n)),
                Some(HttpProtocol::Unknown(ref mut settings, _, _, _)) =>
                    self.node.as_ref().map(|n| settings.head().insert(n)),
                None => unreachable!(),
            };
        }

        let kind = match self.proto {
            Some(HttpProtocol::H1(ref mut h1)) => {
                let result = h1.poll();
                match result {
                    Ok(Async::Ready(())) | Err(_) => {
                        h1.settings().remove_channel();
                        self.node.as_mut().map(|n| n.remove());
                    },
                    _ => (),
                }
                return result
            },
            Some(HttpProtocol::H2(ref mut h2)) => {
                let result = h2.poll();
                match result {
                    Ok(Async::Ready(())) | Err(_) => {
                        h2.settings().remove_channel();
                        self.node.as_mut().map(|n| n.remove());
                    },
                    _ => (),
                }
                return result
            },
            Some(HttpProtocol::Unknown(ref mut settings, _, ref mut io, ref mut buf)) => {
                match utils::read_from_io(io, buf) {
                    Ok(Async::Ready(0)) | Err(_) => {
                        debug!("Ignored premature client disconnection");
                        settings.remove_channel();
                        self.node.as_mut().map(|n| n.remove());
                        return Err(())
                    },
                    _ => (),
                }

                if buf.len() >= 14 {
                    if buf[..14] == HTTP2_PREFACE[..] {
                        ProtocolKind::Http2
                    } else {
                        ProtocolKind::Http1
                    }
                } else {
                    return Ok(Async::NotReady);
                }
            },
            None => unreachable!(),
        };

        // upgrade to specific http protocol
        if let Some(HttpProtocol::Unknown(settings, addr, io, buf)) = self.proto.take() {
            match kind {
                ProtocolKind::Http1 => {
                    self.proto = Some(
                        HttpProtocol::H1(h1::Http1::new(settings, io, addr, buf)));
                    return self.poll()
                },
                ProtocolKind::Http2 => {
                    self.proto = Some(
                        HttpProtocol::H2(h2::Http2::new(settings, io, addr, buf.freeze())));
                    return self.poll()
                },
            }
        }
        unreachable!()
    }
}

pub(crate) struct Node<T>
{
    next: Option<*mut Node<()>>,
    prev: Option<*mut Node<()>>,
    element: *mut T,
}

impl<T> Node<T>
{
    fn new(el: *mut T) -> Self {
        Node {
            next: None,
            prev: None,
            element: el,
        }
    }

    fn insert<I>(&self, next: &Node<I>) {
        #[allow(mutable_transmutes)]
        unsafe {
            if let Some(ref next2) = self.next {
                let n: &mut Node<()> = mem::transmute(next2.as_ref().unwrap());
                n.prev = Some(next as *const _ as *mut _);
            }
            let slf: &mut Node<T> = mem::transmute(self);
            slf.next = Some(next as *const _ as *mut _);

            let next: &mut Node<T> = mem::transmute(next);
            next.prev = Some(slf as *const _ as *mut _);
        }
    }

    fn remove(&mut self) {
        unsafe {
            self.element = ptr::null_mut();
            let next = self.next.take();
            let mut prev = self.prev.take();

            if let Some(ref mut prev) = prev {
                prev.as_mut().unwrap().next = next;
            }
        }
    }
}

impl Node<()> {

    pub(crate) fn head() -> Self {
        Node {
            next: None,
            prev: None,
            element: ptr::null_mut(),
        }
    }

    pub(crate) fn traverse<T, H>(&self) where T: IoStream, H: HttpHandler + 'static {
        let mut next = self.next.as_ref();
        loop {
            if let Some(n) = next {
                unsafe {
                    let n: &Node<()> = mem::transmute(n.as_ref().unwrap());
                    next = n.next.as_ref();

                    if !n.element.is_null() {
                        let ch: &mut HttpChannel<T, H> = mem::transmute(
                            &mut *(n.element as *mut _));
                        ch.shutdown();
                    }
                }
            } else {
                return
            }
        }
    }
}

/// Wrapper for `AsyncRead + AsyncWrite` types
pub(crate) struct WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
   io: T,
}

impl<T> WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
    pub fn new(io: T) -> Self {
        WrapperStream{ io }
    }
}

impl<T> IoStream for WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
    #[inline]
    fn shutdown(&mut self, _: Shutdown) -> io::Result<()> {
        Ok(())
    }
    #[inline]
    fn set_nodelay(&mut self, _: bool) -> io::Result<()> {
        Ok(())
    }
    #[inline]
    fn set_linger(&mut self, _: Option<time::Duration>) -> io::Result<()> {
        Ok(())
    }
}

impl<T> io::Read for WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.io.read(buf)
    }
}

impl<T> io::Write for WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.io.write(buf)
    }
    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        self.io.flush()
    }
}

impl<T> AsyncRead for WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
    #[inline]
    fn read_buf<B: BufMut>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
        self.io.read_buf(buf)
    }
}

impl<T> AsyncWrite for WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
    #[inline]
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        self.io.shutdown()
    }
    #[inline]
    fn write_buf<B: Buf>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
        self.io.write_buf(buf)
    }
}
<commit_msg>explicitly set tcp nodelay<commit_after>use std::{ptr, mem, time, io};
use std::rc::Rc;
use std::net::{SocketAddr, Shutdown};

use bytes::{Bytes, BytesMut, Buf, BufMut};
use futures::{Future, Poll, Async};
use tokio_io::{AsyncRead, AsyncWrite};

use super::{h1, h2, utils, HttpHandler, IoStream};
use super::settings::WorkerSettings;

const HTTP2_PREFACE: [u8; 14] = *b"PRI * HTTP/2.0";

enum HttpProtocol<T: IoStream, H: 'static> {
    H1(h1::Http1<T, H>),
    H2(h2::Http2<T, H>),
    Unknown(Rc<WorkerSettings<H>>, Option<SocketAddr>, T, BytesMut),
}

enum ProtocolKind {
    Http1,
    Http2,
}

#[doc(hidden)]
pub struct HttpChannel<T, H> where T: IoStream, H: HttpHandler + 'static {
    proto: Option<HttpProtocol<T, H>>,
    node: Option<Node<HttpChannel<T, H>>>,
}

impl<T, H> HttpChannel<T, H> where T: IoStream, H: HttpHandler + 'static
{
    pub(crate) fn new(settings: Rc<WorkerSettings<H>>,
                      mut io: T, peer: Option<SocketAddr>, http2: bool) -> HttpChannel<T, H>
    {
        settings.add_channel();
        let _ = io.set_nodelay(true);

        if http2 {
            HttpChannel {
                node: None, proto: Some(HttpProtocol::H2(
                    h2::Http2::new(settings, io, peer, Bytes::new()))) }
        } else {
            HttpChannel {
                node: None, proto: Some(HttpProtocol::Unknown(
                    settings, peer, io, BytesMut::with_capacity(4096))) }
        }
    }

    fn shutdown(&mut self) {
        match self.proto {
            Some(HttpProtocol::H1(ref mut h1)) => {
                let io = h1.io();
                let _ = IoStream::set_linger(io, Some(time::Duration::new(0, 0)));
                let _ = IoStream::shutdown(io, Shutdown::Both);
            }
            Some(HttpProtocol::H2(ref mut h2)) => {
                h2.shutdown()
            }
            _ => (),
        }
    }
}

impl<T, H> Future for HttpChannel<T, H> where T: IoStream, H: HttpHandler + 'static
{
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        if !self.node.is_none() {
            let el = self as *mut _;
            self.node = Some(Node::new(el));
            let _ = match self.proto {
                Some(HttpProtocol::H1(ref mut h1)) =>
                    self.node.as_ref().map(|n| h1.settings().head().insert(n)),
                Some(HttpProtocol::H2(ref mut h2)) =>
                    self.node.as_ref().map(|n| h2.settings().head().insert(n)),
                Some(HttpProtocol::Unknown(ref mut settings, _, _, _)) =>
                    self.node.as_ref().map(|n| settings.head().insert(n)),
                None => unreachable!(),
            };
        }

        let kind = match self.proto {
            Some(HttpProtocol::H1(ref mut h1)) => {
                let result = h1.poll();
                match result {
                    Ok(Async::Ready(())) | Err(_) => {
                        h1.settings().remove_channel();
                        self.node.as_mut().map(|n| n.remove());
                    },
                    _ => (),
                }
                return result
            },
            Some(HttpProtocol::H2(ref mut h2)) => {
                let result = h2.poll();
                match result {
                    Ok(Async::Ready(())) | Err(_) => {
                        h2.settings().remove_channel();
                        self.node.as_mut().map(|n| n.remove());
                    },
                    _ => (),
                }
                return result
            },
            Some(HttpProtocol::Unknown(ref mut settings, _, ref mut io, ref mut buf)) => {
                match utils::read_from_io(io, buf) {
                    Ok(Async::Ready(0)) | Err(_) => {
                        debug!("Ignored premature client disconnection");
                        settings.remove_channel();
                        self.node.as_mut().map(|n| n.remove());
                        return Err(())
                    },
                    _ => (),
                }

                if buf.len() >= 14 {
                    if buf[..14] == HTTP2_PREFACE[..] {
                        ProtocolKind::Http2
                    } else {
                        ProtocolKind::Http1
                    }
                } else {
                    return Ok(Async::NotReady);
                }
            },
            None => unreachable!(),
        };

        // upgrade to specific http protocol
        if let Some(HttpProtocol::Unknown(settings, addr, io, buf)) = self.proto.take() {
            match kind {
                ProtocolKind::Http1 => {
                    self.proto = Some(
                        HttpProtocol::H1(h1::Http1::new(settings, io, addr, buf)));
                    return self.poll()
                },
                ProtocolKind::Http2 => {
                    self.proto = Some(
                        HttpProtocol::H2(h2::Http2::new(settings, io, addr, buf.freeze())));
                    return self.poll()
                },
            }
        }
        unreachable!()
    }
}

pub(crate) struct Node<T>
{
    next: Option<*mut Node<()>>,
    prev: Option<*mut Node<()>>,
    element: *mut T,
}

impl<T> Node<T>
{
    fn new(el: *mut T) -> Self {
        Node {
            next: None,
            prev: None,
            element: el,
        }
    }

    fn insert<I>(&self, next: &Node<I>) {
        #[allow(mutable_transmutes)]
        unsafe {
            if let Some(ref next2) = self.next {
                let n: &mut Node<()> = mem::transmute(next2.as_ref().unwrap());
                n.prev = Some(next as *const _ as *mut _);
            }
            let slf: &mut Node<T> = mem::transmute(self);
            slf.next = Some(next as *const _ as *mut _);

            let next: &mut Node<T> = mem::transmute(next);
            next.prev = Some(slf as *const _ as *mut _);
        }
    }

    fn remove(&mut self) {
        unsafe {
            self.element = ptr::null_mut();
            let next = self.next.take();
            let mut prev = self.prev.take();

            if let Some(ref mut prev) = prev {
                prev.as_mut().unwrap().next = next;
            }
        }
    }
}

impl Node<()> {

    pub(crate) fn head() -> Self {
        Node {
            next: None,
            prev: None,
            element: ptr::null_mut(),
        }
    }

    pub(crate) fn traverse<T, H>(&self) where T: IoStream, H: HttpHandler + 'static {
        let mut next = self.next.as_ref();
        loop {
            if let Some(n) = next {
                unsafe {
                    let n: &Node<()> = mem::transmute(n.as_ref().unwrap());
                    next = n.next.as_ref();

                    if !n.element.is_null() {
                        let ch: &mut HttpChannel<T, H> = mem::transmute(
                            &mut *(n.element as *mut _));
                        ch.shutdown();
                    }
                }
            } else {
                return
            }
        }
    }
}

/// Wrapper for `AsyncRead + AsyncWrite` types
pub(crate) struct WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
   io: T,
}

impl<T> WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
    pub fn new(io: T) -> Self {
        WrapperStream{ io }
    }
}

impl<T> IoStream for WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
    #[inline]
    fn shutdown(&mut self, _: Shutdown) -> io::Result<()> {
        Ok(())
    }
    #[inline]
    fn set_nodelay(&mut self, _: bool) -> io::Result<()> {
        Ok(())
    }
    #[inline]
    fn set_linger(&mut self, _: Option<time::Duration>) -> io::Result<()> {
        Ok(())
    }
}

impl<T> io::Read for WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.io.read(buf)
    }
}

impl<T> io::Write for WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.io.write(buf)
    }
    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        self.io.flush()
    }
}

impl<T> AsyncRead for WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
    #[inline]
    fn read_buf<B: BufMut>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
        self.io.read_buf(buf)
    }
}

impl<T> AsyncWrite for WrapperStream<T> where T: AsyncRead + AsyncWrite + 'static {
    #[inline]
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        self.io.shutdown()
    }
    #[inline]
    fn write_buf<B: Buf>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
        self.io.write_buf(buf)
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Handling log-probabilities.

use std::mem;
use std::f64;

pub use stats::{Prob, LogProb};

/// A factor to convert log-probabilities to PHRED-scale (phred = p * LOG_TO_PHRED_FACTOR).
const LOG_TO_PHRED_FACTOR: f64 = -4.3429448190325175; // -10 * 1 / ln(10)

/// A factor to convert PHRED-scale to log-probabilities (p = phred * PHRED_TO_LOG_FACTOR).
const PHRED_TO_LOG_FACTOR: f64 = -0.23025850929940456; // 1 / (-10 * log10(e))

/// Calculate log(1 - p) with p given in log space without loss of precision as described in
/// http://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf.
pub fn ln_1m_exp(p: LogProb) -> LogProb {
    if p < -0.693 {
        (-p.exp()).ln_1p()
    }
    else {
        (-p.exp_m1()).ln()
    }
}

/// Convert log scale probability to PHRED scale.
pub fn log_to_phred(p: LogProb) -> f64 {
    p * LOG_TO_PHRED_FACTOR
}

/// Convert PHRED scale probability to log scale.
pub fn phred_to_log(p: f64) -> LogProb {
    p * PHRED_TO_LOG_FACTOR
}

/// Calculate the sum of the given probabilities in a numerically stable way (Durbin 1998).
pub fn log_prob_sum(probs: &[LogProb]) -> LogProb {
    if probs.is_empty() {
        f64::NEG_INFINITY
    }
    else {
        let mut pmax = probs[0];
        let mut imax = 0;
        for (i, &p) in probs.iter().enumerate().skip(1) {
            if p > pmax {
                pmax = p;
                imax = i;
            }
        }
        if pmax == f64::NEG_INFINITY {
            f64::NEG_INFINITY
        }
        else if pmax == f64::INFINITY {
            f64::INFINITY
        }
        else {
            // TODO use sum() once it has been stabilized: .sum::<usize>()
            pmax + (probs.iter().enumerate().filter_map(|(i, p)| if i != imax { Some((p - pmax).exp()) } else { None }).fold(0.0, |s, e| s + e)).ln_1p()
        }
    }
}

/// Calculate the sum of the given probabilities in a numerically stable way (Durbin 1998).
pub fn log_prob_add(mut p0: LogProb, mut p1: LogProb) -> LogProb {
    if p1 > p0 {
        mem::swap(&mut p0, &mut p1);
    }
    if p0 == f64::NEG_INFINITY {
        f64::NEG_INFINITY
    }
    else if p0 == f64::INFINITY {
        f64::INFINITY
    }
    else {
        p0 + (p1 - p0).exp().ln_1p()
    }
}

/// Calcualte the cumulative sum of the given probabilities in a numerically stable way (Durbin 1998).
pub fn log_prob_cumsum(probs: &[LogProb]) -> Vec<LogProb> {
    probs.iter().scan(f64::NEG_INFINITY, |s, p| {
        *s = log_prob_add(*s, *p);
        Some(*s)
    }).collect()
}

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

    #[test]
    fn test_log_prob_sum() {
        let probs = [f64::NEG_INFINITY, 0.0, f64::NEG_INFINITY];
        assert_eq!(log_prob_sum(&probs), 0.0);
    }

    #[test]
    fn test_empty_log_prob_sum() {
        assert_eq!(log_prob_sum(&[]), f64::NEG_INFINITY);
    }
}
<commit_msg>Allow to use an iterator for log_prob_cumsum.<commit_after>// Copyright 2014 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Handling log-probabilities.

use std::mem;
use std::f64;

pub use stats::{Prob, LogProb};

/// A factor to convert log-probabilities to PHRED-scale (phred = p * LOG_TO_PHRED_FACTOR).
const LOG_TO_PHRED_FACTOR: f64 = -4.3429448190325175; // -10 * 1 / ln(10)

/// A factor to convert PHRED-scale to log-probabilities (p = phred * PHRED_TO_LOG_FACTOR).
const PHRED_TO_LOG_FACTOR: f64 = -0.23025850929940456; // 1 / (-10 * log10(e))

/// Calculate log(1 - p) with p given in log space without loss of precision as described in
/// http://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf.
pub fn ln_1m_exp(p: LogProb) -> LogProb {
    if p < -0.693 {
        (-p.exp()).ln_1p()
    }
    else {
        (-p.exp_m1()).ln()
    }
}

/// Convert log scale probability to PHRED scale.
pub fn log_to_phred(p: LogProb) -> f64 {
    p * LOG_TO_PHRED_FACTOR
}

/// Convert PHRED scale probability to log scale.
pub fn phred_to_log(p: f64) -> LogProb {
    p * PHRED_TO_LOG_FACTOR
}

/// Calculate the sum of the given probabilities in a numerically stable way (Durbin 1998).
pub fn log_prob_sum(probs: &[LogProb]) -> LogProb {
    if probs.is_empty() {
        f64::NEG_INFINITY
    }
    else {
        let mut pmax = probs[0];
        let mut imax = 0;
        for (i, &p) in probs.iter().enumerate().skip(1) {
            if p > pmax {
                pmax = p;
                imax = i;
            }
        }
        if pmax == f64::NEG_INFINITY {
            f64::NEG_INFINITY
        }
        else if pmax == f64::INFINITY {
            f64::INFINITY
        }
        else {
            // TODO use sum() once it has been stabilized: .sum::<usize>()
            pmax + (probs.iter().enumerate().filter_map(|(i, p)| if i != imax { Some((p - pmax).exp()) } else { None }).fold(0.0, |s, e| s + e)).ln_1p()
        }
    }
}

/// Calculate the sum of the given probabilities in a numerically stable way (Durbin 1998).
pub fn log_prob_add(mut p0: LogProb, mut p1: LogProb) -> LogProb {
    if p1 > p0 {
        mem::swap(&mut p0, &mut p1);
    }
    if p0 == f64::NEG_INFINITY {
        f64::NEG_INFINITY
    }
    else if p0 == f64::INFINITY {
        f64::INFINITY
    }
    else {
        p0 + (p1 - p0).exp().ln_1p()
    }
}

/// Calculate the cumulative sum of the given probabilities in a numerically stable way (Durbin 1998).
pub fn log_prob_cumsum<'a, I: Iterator<Item=&'a LogProb>>(probs: I) -> Vec<LogProb> {
    probs.scan(f64::NEG_INFINITY, |s, p| {
        *s = log_prob_add(*s, *p);
        Some(*s)
    }).collect()
}

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

    #[test]
    fn test_log_prob_sum() {
        let probs = [f64::NEG_INFINITY, 0.0, f64::NEG_INFINITY];
        assert_eq!(log_prob_sum(&probs), 0.0);
    }

    #[test]
    fn test_empty_log_prob_sum() {
        assert_eq!(log_prob_sum(&[]), f64::NEG_INFINITY);
    }
}
<|endoftext|>


<commit_before>use super::Blob;
use failure::Error;
use futures::Future;
use log::{error, warn};
use rusoto_core::region::Region;
use rusoto_credential::DefaultCredentialsProvider;
use rusoto_s3::{GetObjectRequest, PutObjectRequest, S3Client, S3};
use std::convert::TryInto;
use std::io::Read;
use time::Timespec;
use tokio::runtime::Runtime;

pub(crate) struct S3Backend<'a> {
    client: S3Client,
    bucket: &'a str,
    runtime: Runtime,
}

impl<'a> S3Backend<'a> {
    pub(crate) fn new(client: S3Client, bucket: &'a str) -> Self {
        Self {
            client,
            bucket,
            runtime: Runtime::new().unwrap(),
        }
    }

    pub(super) fn get(&self, path: &str) -> Result<Blob, Error> {
        let res = self
            .client
            .get_object(GetObjectRequest {
                bucket: self.bucket.to_string(),
                key: path.into(),
                ..Default::default()
            })
            .sync()?;

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::with_capacity(
            res.content_length
                .and_then(|l| l.try_into().ok())
                .unwrap_or(0),
        );
        b.read_to_end(&mut content).unwrap();

        let date_updated = parse_timespec(&res.last_modified.unwrap())?;

        Ok(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated,
            content,
        })
    }

    pub(super) fn store_batch(&mut self, batch: &[Blob]) -> Result<(), Error> {
        use futures::stream::FuturesUnordered;
        use futures::stream::Stream;
        let mut attempts = 0;

        loop {
            let mut futures = FuturesUnordered::new();
            for blob in batch {
                futures.push(
                    self.client
                        .put_object(PutObjectRequest {
                            bucket: self.bucket.to_string(),
                            key: blob.path.clone(),
                            body: Some(blob.content.clone().into()),
                            content_type: Some(blob.mime.clone()),
                            ..Default::default()
                        })
                        .inspect(|_| {
                            crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(1);
                        }),
                );
            }
            attempts += 1;

            match self.runtime.block_on(futures.map(drop).collect()) {
                // this batch was successful, start another batch if there are still more files
                Ok(_) => break,
                Err(err) => {
                    error!("failed to upload to s3: {:?}", err);
                    // if a futures error occurs, retry the batch
                    if attempts > 2 {
                        panic!("failed to upload 3 times, exiting");
                    }
                }
            }
        }
        Ok(())
    }
}

// public for testing
pub(crate) const TIME_FMT: &str = "%a, %d %b %Y %H:%M:%S %Z";

fn parse_timespec(raw: &str) -> Result<Timespec, Error> {
    Ok(time::strptime(raw, TIME_FMT)?.to_timespec())
}

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

    use crate::storage::s3::S3Backend;
    use rusoto_s3::{
        CreateBucketRequest, DeleteBucketRequest, DeleteObjectRequest, ListObjectsRequest, S3,
    };

    use std::cell::RefCell;
    use std::slice;

    pub(crate) struct TestS3(RefCell<S3Backend<'static>>);

    impl TestS3 {
        pub(crate) fn new() -> Self {
            // A random bucket name is generated and used for the current connection.
            // This allows each test to create a fresh bucket to test with.
            let bucket = format!("docs-rs-test-bucket-{}", rand::random::<u64>());
            let client = s3_client().unwrap();
            client
                .create_bucket(CreateBucketRequest {
                    bucket: bucket.clone(),
                    ..Default::default()
                })
                .sync()
                .expect("failed to create test bucket");
            let bucket = Box::leak(bucket.into_boxed_str());
            TestS3(RefCell::new(S3Backend::new(client, bucket)))
        }
        pub(crate) fn upload(&self, blobs: &[Blob]) -> Result<(), Error> {
            self.0.borrow_mut().store_batch(blobs)
        }
        fn assert_404(&self, path: &'static str) {
            use rusoto_core::RusotoError;
            use rusoto_s3::GetObjectError;

            let err = self.0.borrow().get(path).unwrap_err();
            match err
                .downcast_ref::<RusotoError<GetObjectError>>()
                .expect("wanted GetObject")
            {
                RusotoError::Unknown(http) => assert_eq!(http.status, 404),
                RusotoError::Service(GetObjectError::NoSuchKey(_)) => {}
                x => panic!("wrong error: {:?}", x),
            };
        }
        fn assert_blob(&self, blob: &Blob, path: &str) {
            let actual = self.0.borrow().get(path).unwrap();
            assert_eq!(blob.path, actual.path);
            assert_eq!(blob.content, actual.content);
            assert_eq!(blob.mime, actual.mime);
            // NOTE: this does _not_ compare the upload time since min.io doesn't allow this to be configured
        }
    }

    impl Drop for TestS3 {
        fn drop(&mut self) {
            // delete the bucket when the test ends
            // this has to delete all the objects in the bucket first or min.io will give an error
            let inner = self.0.borrow();
            let list_req = ListObjectsRequest {
                bucket: inner.bucket.to_owned(),
                ..Default::default()
            };
            let objects = inner.client.list_objects(list_req).sync().unwrap();
            assert!(!objects.is_truncated.unwrap_or(false));
            for path in objects.contents.unwrap() {
                let delete_req = DeleteObjectRequest {
                    bucket: inner.bucket.to_owned(),
                    key: path.key.unwrap(),
                    ..Default::default()
                };
                inner.client.delete_object(delete_req).sync().unwrap();
            }
            let delete_req = DeleteBucketRequest {
                bucket: inner.bucket.to_owned(),
            };
            inner
                .client
                .delete_bucket(delete_req)
                .sync()
                .expect("failed to delete test bucket");
        }
    }

    #[test]
    fn test_parse_timespec() {
        // Test valid conversions
        assert_eq!(
            parse_timespec("Thu, 1 Jan 1970 00:00:00 GMT").unwrap(),
            Timespec::new(0, 0)
        );
        assert_eq!(
            parse_timespec("Mon, 16 Apr 2018 04:33:50 GMT").unwrap(),
            Timespec::new(1523853230, 0)
        );

        // Test invalid conversion
        assert!(parse_timespec("foo").is_err());
    }

    #[test]
    fn test_get() {
        wrapper(|env| {
            let blob = Blob {
                path: "dir/foo.txt".into(),
                mime: "text/plain".into(),
                date_updated: Timespec::new(42, 0),
                content: "Hello world!".into(),
            };

            // Add a test file to the database
            let s3 = env.s3();
            s3.upload(slice::from_ref(&blob)).unwrap();

            // Test that the proper file was returned
            s3.assert_blob(&blob, "dir/foo.txt");

            // Test that other files are not returned
            s3.assert_404("dir/bar.txt");
            s3.assert_404("foo.txt");

            Ok(())
        });
    }

    #[test]
    fn test_store() {
        wrapper(|env| {
            let s3 = env.s3();
            let names = [
                "a",
                "b",
                "a_very_long_file_name_that_has_an.extension",
                "parent/child",
                "h/i/g/h/l/y/_/n/e/s/t/e/d/_/d/i/r/e/c/t/o/r/i/e/s",
            ];
            let blobs: Vec<_> = names
                .iter()
                .map(|&path| Blob {
                    path: path.into(),
                    mime: "text/plain".into(),
                    date_updated: Timespec::new(42, 0),
                    content: "Hello world!".into(),
                })
                .collect();
            s3.upload(&blobs).unwrap();
            for blob in &blobs {
                s3.assert_blob(blob, &blob.path);
            }
            Ok(())
        })
    }
    // NOTE: trying to upload a file ending with `/` will behave differently in test and prod.
    // NOTE: On s3, it will succeed and create a file called `/`.
    // NOTE: On min.io, it will fail with 'Object name contains unsupported characters.'
}

pub(crate) static S3_BUCKET_NAME: &str = "rust-docs-rs";

pub(crate) fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() && std::env::var_os("FORCE_S3").is_none() {
        return None;
    }
    let creds = match DefaultCredentialsProvider::new() {
        Ok(creds) => creds,
        Err(err) => {
            warn!("failed to retrieve AWS credentials: {}", err);
            return None;
        }
    };
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        creds,
        std::env::var("S3_ENDPOINT")
            .ok()
            .map(|e| Region::Custom {
                name: std::env::var("S3_REGION").unwrap_or_else(|_| "us-west-1".to_owned()),
                endpoint: e,
            })
            .unwrap_or(Region::UsWest1),
    ))
}
<commit_msg>Move some code around<commit_after>use super::Blob;
use failure::Error;
use futures::Future;
use log::{error, warn};
use rusoto_core::region::Region;
use rusoto_credential::DefaultCredentialsProvider;
use rusoto_s3::{GetObjectRequest, PutObjectRequest, S3Client, S3};
use std::convert::TryInto;
use std::io::Read;
use time::Timespec;
use tokio::runtime::Runtime;

pub(crate) static S3_BUCKET_NAME: &str = "rust-docs-rs";

pub(crate) struct S3Backend<'a> {
    client: S3Client,
    bucket: &'a str,
    runtime: Runtime,
}

impl<'a> S3Backend<'a> {
    pub(crate) fn new(client: S3Client, bucket: &'a str) -> Self {
        Self {
            client,
            bucket,
            runtime: Runtime::new().unwrap(),
        }
    }

    pub(super) fn get(&self, path: &str) -> Result<Blob, Error> {
        let res = self
            .client
            .get_object(GetObjectRequest {
                bucket: self.bucket.to_string(),
                key: path.into(),
                ..Default::default()
            })
            .sync()?;

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::with_capacity(
            res.content_length
                .and_then(|l| l.try_into().ok())
                .unwrap_or(0),
        );
        b.read_to_end(&mut content).unwrap();

        let date_updated = parse_timespec(&res.last_modified.unwrap())?;

        Ok(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated,
            content,
        })
    }

    pub(super) fn store_batch(&mut self, batch: &[Blob]) -> Result<(), Error> {
        use futures::stream::FuturesUnordered;
        use futures::stream::Stream;
        let mut attempts = 0;

        loop {
            let mut futures = FuturesUnordered::new();
            for blob in batch {
                futures.push(
                    self.client
                        .put_object(PutObjectRequest {
                            bucket: self.bucket.to_string(),
                            key: blob.path.clone(),
                            body: Some(blob.content.clone().into()),
                            content_type: Some(blob.mime.clone()),
                            ..Default::default()
                        })
                        .inspect(|_| {
                            crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(1);
                        }),
                );
            }
            attempts += 1;

            match self.runtime.block_on(futures.map(drop).collect()) {
                // this batch was successful, start another batch if there are still more files
                Ok(_) => break,
                Err(err) => {
                    error!("failed to upload to s3: {:?}", err);
                    // if a futures error occurs, retry the batch
                    if attempts > 2 {
                        panic!("failed to upload 3 times, exiting");
                    }
                }
            }
        }
        Ok(())
    }
}

// public for testing
pub(crate) const TIME_FMT: &str = "%a, %d %b %Y %H:%M:%S %Z";

fn parse_timespec(raw: &str) -> Result<Timespec, Error> {
    Ok(time::strptime(raw, TIME_FMT)?.to_timespec())
}

pub(crate) fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() && std::env::var_os("FORCE_S3").is_none() {
        return None;
    }
    let creds = match DefaultCredentialsProvider::new() {
        Ok(creds) => creds,
        Err(err) => {
            warn!("failed to retrieve AWS credentials: {}", err);
            return None;
        }
    };
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        creds,
        std::env::var("S3_ENDPOINT")
            .ok()
            .map(|e| Region::Custom {
                name: std::env::var("S3_REGION").unwrap_or_else(|_| "us-west-1".to_owned()),
                endpoint: e,
            })
            .unwrap_or(Region::UsWest1),
    ))
}

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

    use crate::storage::s3::S3Backend;
    use rusoto_s3::{
        CreateBucketRequest, DeleteBucketRequest, DeleteObjectRequest, ListObjectsRequest, S3,
    };

    use std::cell::RefCell;
    use std::slice;

    pub(crate) struct TestS3(RefCell<S3Backend<'static>>);

    impl TestS3 {
        pub(crate) fn new() -> Self {
            // A random bucket name is generated and used for the current connection.
            // This allows each test to create a fresh bucket to test with.
            let bucket = format!("docs-rs-test-bucket-{}", rand::random::<u64>());
            let client = s3_client().unwrap();
            client
                .create_bucket(CreateBucketRequest {
                    bucket: bucket.clone(),
                    ..Default::default()
                })
                .sync()
                .expect("failed to create test bucket");
            let bucket = Box::leak(bucket.into_boxed_str());
            TestS3(RefCell::new(S3Backend::new(client, bucket)))
        }
        pub(crate) fn upload(&self, blobs: &[Blob]) -> Result<(), Error> {
            self.0.borrow_mut().store_batch(blobs)
        }
        fn assert_404(&self, path: &'static str) {
            use rusoto_core::RusotoError;
            use rusoto_s3::GetObjectError;

            let err = self.0.borrow().get(path).unwrap_err();
            match err
                .downcast_ref::<RusotoError<GetObjectError>>()
                .expect("wanted GetObject")
            {
                RusotoError::Unknown(http) => assert_eq!(http.status, 404),
                RusotoError::Service(GetObjectError::NoSuchKey(_)) => {}
                x => panic!("wrong error: {:?}", x),
            };
        }
        fn assert_blob(&self, blob: &Blob, path: &str) {
            let actual = self.0.borrow().get(path).unwrap();
            assert_eq!(blob.path, actual.path);
            assert_eq!(blob.content, actual.content);
            assert_eq!(blob.mime, actual.mime);
            // NOTE: this does _not_ compare the upload time since min.io doesn't allow this to be configured
        }
    }

    impl Drop for TestS3 {
        fn drop(&mut self) {
            // delete the bucket when the test ends
            // this has to delete all the objects in the bucket first or min.io will give an error
            let inner = self.0.borrow();
            let list_req = ListObjectsRequest {
                bucket: inner.bucket.to_owned(),
                ..Default::default()
            };
            let objects = inner.client.list_objects(list_req).sync().unwrap();
            assert!(!objects.is_truncated.unwrap_or(false));
            for path in objects.contents.unwrap() {
                let delete_req = DeleteObjectRequest {
                    bucket: inner.bucket.to_owned(),
                    key: path.key.unwrap(),
                    ..Default::default()
                };
                inner.client.delete_object(delete_req).sync().unwrap();
            }
            let delete_req = DeleteBucketRequest {
                bucket: inner.bucket.to_owned(),
            };
            inner
                .client
                .delete_bucket(delete_req)
                .sync()
                .expect("failed to delete test bucket");
        }
    }

    #[test]
    fn test_parse_timespec() {
        // Test valid conversions
        assert_eq!(
            parse_timespec("Thu, 1 Jan 1970 00:00:00 GMT").unwrap(),
            Timespec::new(0, 0)
        );
        assert_eq!(
            parse_timespec("Mon, 16 Apr 2018 04:33:50 GMT").unwrap(),
            Timespec::new(1523853230, 0)
        );

        // Test invalid conversion
        assert!(parse_timespec("foo").is_err());
    }

    #[test]
    fn test_get() {
        wrapper(|env| {
            let blob = Blob {
                path: "dir/foo.txt".into(),
                mime: "text/plain".into(),
                date_updated: Timespec::new(42, 0),
                content: "Hello world!".into(),
            };

            // Add a test file to the database
            let s3 = env.s3();
            s3.upload(slice::from_ref(&blob)).unwrap();

            // Test that the proper file was returned
            s3.assert_blob(&blob, "dir/foo.txt");

            // Test that other files are not returned
            s3.assert_404("dir/bar.txt");
            s3.assert_404("foo.txt");

            Ok(())
        });
    }

    #[test]
    fn test_store() {
        wrapper(|env| {
            let s3 = env.s3();
            let names = [
                "a",
                "b",
                "a_very_long_file_name_that_has_an.extension",
                "parent/child",
                "h/i/g/h/l/y/_/n/e/s/t/e/d/_/d/i/r/e/c/t/o/r/i/e/s",
            ];
            let blobs: Vec<_> = names
                .iter()
                .map(|&path| Blob {
                    path: path.into(),
                    mime: "text/plain".into(),
                    date_updated: Timespec::new(42, 0),
                    content: "Hello world!".into(),
                })
                .collect();
            s3.upload(&blobs).unwrap();
            for blob in &blobs {
                s3.assert_blob(blob, &blob.path);
            }
            Ok(())
        })
    }
    // NOTE: trying to upload a file ending with `/` will behave differently in test and prod.
    // NOTE: On s3, it will succeed and create a file called `/`.
    // NOTE: On min.io, it will fail with 'Object name contains unsupported characters.'
}
<|endoftext|>


<commit_before>// Copyright (c) 2015, Johan Sköld.
// License: http://opensource.org/licenses/ISC

extern crate bgfx;
extern crate glutin;
extern crate libc;

use std::env;
use std::fs::File;
use std::io::Read;
use std::sync::mpsc::{Receiver, Sender};
use std::sync::mpsc;

use glutin::{Api, GlContext, GlRequest, Window, WindowBuilder};

/// Events received by the main thread, sent by the render thread.
#[derive(PartialEq, Eq, Hash, Debug)]
pub enum Event {
    /// Window close event.
    Close,

    /// Window size event.
    Size(u16, u16),
}

/// Example application.
pub struct Example {
    /// Receiver for events from the render thread.
    event_rx: Receiver<Event>,
}

impl Example {

    /// Handles events received from the render thread. If there are no events to process, returns
    /// instantly.
    ///
    /// Returns `true` if the app should exit.
    pub fn handle_events(&self,
                         bgfx: &bgfx::MainContext,
                         width: &mut u16,
                         height: &mut u16,
                         reset: bgfx::ResetFlags)
                         -> bool {
        let mut close = false;

        loop {
            let result = self.event_rx.try_recv();

            if !result.is_ok() {
                break;
            }

            match result.ok().unwrap() {
                Event::Close => close = true,
                Event::Size(w, h) => {
                    *width = w;
                    *height = h;
                    bgfx.reset(w, h, reset);
                }
            }
        }

        close
    }

}

/// Example data used by the render thread.
struct ExampleData {
    should_close: bool,

    /// The glutin window object.
    window: Window,

    /// Sender of events to the main thread.
    event_tx: Sender<Event>,
}

impl ExampleData {

    /// Process glutin events, and potentially forward them to the main thread.
    ///
    /// Returns `true` if the example should exit.
    fn process_events(&mut self) -> bool {
        for event in self.window.poll_events() {
            match event {
                glutin::Event::Closed => {
                    self.should_close = true;
                    self.event_tx.send(Event::Close).unwrap();
                }
                glutin::Event::Resized(w, h) => {
                    self.event_tx.send(Event::Size(w as u16, h as u16)).unwrap();
                }
                _ => {}
            }
        }

        self.should_close
    }

}

/// Loads the contents of a file and returns it.
fn load_file(name: &str) -> Vec<u8> {
    let mut data = Vec::new();
    let mut file = File::open(name).unwrap();
    file.read_to_end(&mut data).unwrap();
    data
}

/// Loads the two given shaders from disk, and creates a program using the new
/// shaders.
///
/// # Arguments
///
/// * ``
#[allow(dead_code)]
pub fn load_program<'a, 'b>(bgfx: &'a bgfx::MainContext,
                            vsh_name: &'b str,
                            fsh_name: &'b str)
                            -> bgfx::Program<'a> {
    let renderer = bgfx.get_renderer_type();
    let exe_path = env::current_exe().unwrap();
    let exe_stem = exe_path.file_stem().unwrap();
    let assets_path = format!("examples/assets/{}", exe_stem.to_str().unwrap());
    let vsh_path = format!("{}/{:?}/{}.bin", assets_path, renderer, vsh_name);
    let fsh_path = format!("{}/{:?}/{}.bin", assets_path, renderer, fsh_name);
    let vsh_mem = bgfx::Memory::copy(&load_file(&vsh_path));
    let fsh_mem = bgfx::Memory::copy(&load_file(&fsh_path));
    let vsh = bgfx::Shader::new(bgfx, vsh_mem);
    let fsh = bgfx::Shader::new(bgfx, fsh_mem);

    bgfx::Program::new(vsh, fsh)
}

/// Returns a new `bgfx::Application`.
///
/// # Arguments
///
/// * `window` - Reference to the glutin window object.
#[cfg(target_os = "linux")]
fn create_bgfx(window: &Window) -> bgfx::Config {
    use glutin::os::unix::WindowExt;
    let mut bgfx = bgfx::create();
    bgfx.display(window.get_xlib_display().unwrap());
    bgfx.window(window.get_xlib_window().unwrap());
    bgfx
}

#[cfg(target_os = "windows")]
fn create_bgfx(window: &Window) -> bgfx::Config {
    use glutin::os::windows::WindowExt;
    let mut bgfx = bgfx::create();
    bgfx.window(window.get_hwnd());
    bgfx
}

/// Runs an example.
///
/// # Arguments
///
/// * `width` - Initial width of the window, in pixels.
/// * `height` - Initial height of the window, in pixels.
/// * `main` - Function to act as the entry point for the example.
pub fn run_example<M>(width: u32, height: u32, main: M)
    where M: Send + 'static + FnOnce(bgfx::MainContext, &Example)
{
    let window = WindowBuilder::new()
        .with_dimensions(width, height)
        .with_gl(GlRequest::Specific(Api::OpenGl, (2, 1)))
        .with_title(String::from("BGFX"))
        .build()
        .expect("Failed to create window");

    unsafe { window.make_current().unwrap(); }

    // Create the channel used for communication between the main and render threads.
    let (event_tx, event_rx) = mpsc::channel::<Event>();

    // Initialize the example.
    let mut data = ExampleData {
        should_close: false,
        window: window,
        event_tx: event_tx,
    };

    let bgfx = create_bgfx(&data.window);

    // Main thread implementation.
    let main_thread = move |bgfx: bgfx::MainContext| {
        let example = Example { event_rx: event_rx };
        main(bgfx, &example);
    };

    // Render thread implementation.
    let render_thread = |bgfx: bgfx::RenderContext| {
        while !data.process_events() {
            bgfx.render_frame();
        }
    };

    // Run the application
    bgfx.run(main_thread, render_thread).unwrap();
}
<commit_msg>OpenGL 2.1 -> 3.1<commit_after>// Copyright (c) 2015, Johan Sköld.
// License: http://opensource.org/licenses/ISC

extern crate bgfx;
extern crate glutin;
extern crate libc;

use std::env;
use std::fs::File;
use std::io::Read;
use std::sync::mpsc::{Receiver, Sender};
use std::sync::mpsc;

use glutin::{Api, GlContext, GlRequest, Window, WindowBuilder};

/// Events received by the main thread, sent by the render thread.
#[derive(PartialEq, Eq, Hash, Debug)]
pub enum Event {
    /// Window close event.
    Close,

    /// Window size event.
    Size(u16, u16),
}

/// Example application.
pub struct Example {
    /// Receiver for events from the render thread.
    event_rx: Receiver<Event>,
}

impl Example {

    /// Handles events received from the render thread. If there are no events to process, returns
    /// instantly.
    ///
    /// Returns `true` if the app should exit.
    pub fn handle_events(&self,
                         bgfx: &bgfx::MainContext,
                         width: &mut u16,
                         height: &mut u16,
                         reset: bgfx::ResetFlags)
                         -> bool {
        let mut close = false;

        loop {
            let result = self.event_rx.try_recv();

            if !result.is_ok() {
                break;
            }

            match result.ok().unwrap() {
                Event::Close => close = true,
                Event::Size(w, h) => {
                    *width = w;
                    *height = h;
                    bgfx.reset(w, h, reset);
                }
            }
        }

        close
    }

}

/// Example data used by the render thread.
struct ExampleData {
    should_close: bool,

    /// The glutin window object.
    window: Window,

    /// Sender of events to the main thread.
    event_tx: Sender<Event>,
}

impl ExampleData {

    /// Process glutin events, and potentially forward them to the main thread.
    ///
    /// Returns `true` if the example should exit.
    fn process_events(&mut self) -> bool {
        for event in self.window.poll_events() {
            match event {
                glutin::Event::Closed => {
                    self.should_close = true;
                    self.event_tx.send(Event::Close).unwrap();
                }
                glutin::Event::Resized(w, h) => {
                    self.event_tx.send(Event::Size(w as u16, h as u16)).unwrap();
                }
                _ => {}
            }
        }

        self.should_close
    }

}

/// Loads the contents of a file and returns it.
fn load_file(name: &str) -> Vec<u8> {
    let mut data = Vec::new();
    let mut file = File::open(name).unwrap();
    file.read_to_end(&mut data).unwrap();
    data
}

/// Loads the two given shaders from disk, and creates a program using the new
/// shaders.
///
/// # Arguments
///
/// * ``
#[allow(dead_code)]
pub fn load_program<'a, 'b>(bgfx: &'a bgfx::MainContext,
                            vsh_name: &'b str,
                            fsh_name: &'b str)
                            -> bgfx::Program<'a> {
    let renderer = bgfx.get_renderer_type();
    let exe_path = env::current_exe().unwrap();
    let exe_stem = exe_path.file_stem().unwrap();
    let assets_path = format!("examples/assets/{}", exe_stem.to_str().unwrap());
    let vsh_path = format!("{}/{:?}/{}.bin", assets_path, renderer, vsh_name);
    let fsh_path = format!("{}/{:?}/{}.bin", assets_path, renderer, fsh_name);
    let vsh_mem = bgfx::Memory::copy(&load_file(&vsh_path));
    let fsh_mem = bgfx::Memory::copy(&load_file(&fsh_path));
    let vsh = bgfx::Shader::new(bgfx, vsh_mem);
    let fsh = bgfx::Shader::new(bgfx, fsh_mem);

    bgfx::Program::new(vsh, fsh)
}

/// Returns a new `bgfx::Application`.
///
/// # Arguments
///
/// * `window` - Reference to the glutin window object.
#[cfg(target_os = "linux")]
fn create_bgfx(window: &Window) -> bgfx::Config {
    use glutin::os::unix::WindowExt;
    let mut bgfx = bgfx::create();
    bgfx.display(window.get_xlib_display().unwrap());
    bgfx.window(window.get_xlib_window().unwrap());
    bgfx
}

#[cfg(target_os = "windows")]
fn create_bgfx(window: &Window) -> bgfx::Config {
    use glutin::os::windows::WindowExt;
    let mut bgfx = bgfx::create();
    bgfx.window(window.get_hwnd());
    bgfx
}

/// Runs an example.
///
/// # Arguments
///
/// * `width` - Initial width of the window, in pixels.
/// * `height` - Initial height of the window, in pixels.
/// * `main` - Function to act as the entry point for the example.
pub fn run_example<M>(width: u32, height: u32, main: M)
    where M: Send + 'static + FnOnce(bgfx::MainContext, &Example)
{
    let window = WindowBuilder::new()
        .with_dimensions(width, height)
        .with_gl(GlRequest::Specific(Api::OpenGl, (3, 1)))
        .with_title(String::from("BGFX"))
        .build()
        .expect("Failed to create window");

    unsafe { window.make_current().unwrap(); }

    // Create the channel used for communication between the main and render threads.
    let (event_tx, event_rx) = mpsc::channel::<Event>();

    // Initialize the example.
    let mut data = ExampleData {
        should_close: false,
        window: window,
        event_tx: event_tx,
    };

    let bgfx = create_bgfx(&data.window);

    // Main thread implementation.
    let main_thread = move |bgfx: bgfx::MainContext| {
        let example = Example { event_rx: event_rx };
        main(bgfx, &example);
    };

    // Render thread implementation.
    let render_thread = |bgfx: bgfx::RenderContext| {
        while !data.process_events() {
            bgfx.render_frame();
        }
    };

    // Run the application
    bgfx.run(main_thread, render_thread).unwrap();
}
<|endoftext|>


<commit_before>#![deny(deprecated)]
#![deny(missing_docs)]
#![deny(non_snake_case)]
#![deny(non_upper_case_globals)]

#![crate_name = "tut04"]

#![doc = "http://www.opengl-tutorial.org/beginners-tutorials/tutorial-4-a-colored-cube/"]

// Include SDL2 library.
extern crate sdl2;
// Include OpenGL library.
extern crate gl;

extern crate libc;

use gl::types::{GLenum, GLuint, GLsizei, GLchar};
use libc::types::common::c95::c_void;

extern crate tutcommon;

#[doc = "Module for GL drawing stuff."]
pub mod glscene;

extern "system" fn on_debug_message(source: GLenum, gltype: GLenum, id: GLuint, severity: GLenum, length: GLsizei, message: *const GLchar, user_param: *mut c_void) {
    println!("[OpenGL] {:?}", message);
}

fn main() {
    // Initialize SDL2:
    let sdl_context = sdl2::init(sdl2::INIT_VIDEO).unwrap();

    // Init OpenGL parameters:
    sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLMultiSampleBuffers, 1);
    sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLMultiSampleSamples, 4); // 4x antialiasing
    sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLContextMajorVersion, 3); // OpenGL 3.3
    sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLContextMinorVersion, 3);
    sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLContextFlags, sdl2::video::GL_CONTEXT_DEBUG.bits());
    sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLContextProfileMask
        , sdl2::video::GLProfile::GLCoreProfile as i32); // Don't use old OpenGL

    let window = sdl2::video::Window::new(&sdl_context, "Tutorial 04", sdl2::video::WindowPos::PosCentered
        , sdl2::video::WindowPos::PosCentered, 1024, 768, sdl2::video::OPENGL).unwrap();

    let gl_context = window.gl_create_context().unwrap();

    gl::load_with(|s| unsafe {
        std::mem::transmute(sdl2::video::gl_get_proc_address(s))
    });

    unsafe {
        gl::Enable(gl::DEBUG_OUTPUT_SYNCHRONOUS);
        gl::DebugMessageControl(gl::DONT_CARE, gl::DONT_CARE, gl::DONT_CARE, 0, std::ptr::null(), gl::TRUE);
        gl::DebugMessageCallback(on_debug_message, std::ptr::null());
    }

    unsafe {
        gl::ClearColor(0.0, 0.0, 0.4, 0.0);
        // Enable depth test
        gl::Enable(gl::DEPTH_TEST);
        // Accept fragment if it closer to the camera than the former one
        gl::DepthFunc(gl::LESS);
    }

    // init scene.
    let scene = glscene::GLScene::new();

    let mut event_pump = sdl_context.event_pump();
    'evloop : loop {
        unsafe {
            gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);
        }

        scene.draw();

        // Swap buffers.
        window.gl_swap_window();

        for event in event_pump.poll_iter() {
            // check if ESC key pressed or windows closed.
            match event {
                sdl2::event::Event::Quit { .. } => {
                    return;
                }
                sdl2::event::Event::KeyDown {
                            timestamp: _,
                            window_id: _,
                            keycode: _,
                            scancode,
                            keymod: _,
                            repeat: _,
                        } => {
                    if scancode == sdl2::scancode::ScanCode::Escape {
                        return;
                    }
                }
                _ => ()
            }
        }
    }
}

<commit_msg>Fix debug output.<commit_after>#![deny(deprecated)]
#![deny(missing_docs)]
#![deny(non_snake_case)]
#![deny(non_upper_case_globals)]

#![crate_name = "tut04"]

#![doc = "http://www.opengl-tutorial.org/beginners-tutorials/tutorial-4-a-colored-cube/"]

// Include SDL2 library.
extern crate sdl2;
// Include OpenGL library.
extern crate gl;

extern crate libc;

use gl::types::{GLenum, GLuint, GLsizei, GLchar};
use libc::types::common::c95::c_void;
use std::ffi::CStr;

extern crate tutcommon;

#[doc = "Module for GL drawing stuff."]
pub mod glscene;

extern "system" fn on_debug_message(source: GLenum, gltype: GLenum, id: GLuint, severity: GLenum, length: GLsizei, message: *const GLchar, user_param: *mut c_void) {
    let msg = unsafe {
        String::from_utf8_lossy(CStr::from_ptr(message).to_bytes())
    };

    println!("[OpenGL] {}", msg);
}

fn main() {
    // Initialize SDL2:
    let sdl_context = sdl2::init(sdl2::INIT_VIDEO).unwrap();

    // Init OpenGL parameters:
    sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLMultiSampleBuffers, 1);
    sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLMultiSampleSamples, 4); // 4x antialiasing
    sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLContextMajorVersion, 3); // OpenGL 3.3
    sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLContextMinorVersion, 3);
    sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLContextFlags, sdl2::video::GL_CONTEXT_DEBUG.bits());
    sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLContextProfileMask
        , sdl2::video::GLProfile::GLCoreProfile as i32); // Don't use old OpenGL

    let window = sdl2::video::Window::new(&sdl_context, "Tutorial 04", sdl2::video::WindowPos::PosCentered
        , sdl2::video::WindowPos::PosCentered, 1024, 768, sdl2::video::OPENGL).unwrap();

    let gl_context = window.gl_create_context().unwrap();

    gl::load_with(|s| unsafe {
        std::mem::transmute(sdl2::video::gl_get_proc_address(s))
    });

    unsafe {
        gl::Enable(gl::DEBUG_OUTPUT_SYNCHRONOUS);
        gl::DebugMessageControl(gl::DONT_CARE, gl::DONT_CARE, gl::DONT_CARE, 0, std::ptr::null(), gl::TRUE);
        gl::DebugMessageCallback(on_debug_message, std::ptr::null());
    }

    unsafe {
        gl::ClearColor(0.0, 0.0, 0.4, 0.0);
        // Enable depth test
        gl::Enable(gl::DEPTH_TEST);
        // Accept fragment if it closer to the camera than the former one
        gl::DepthFunc(gl::LESS);
    }

    // init scene.
    let scene = glscene::GLScene::new();

    let mut event_pump = sdl_context.event_pump();
    'evloop : loop {
        unsafe {
            gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);
        }

        scene.draw();

        // Swap buffers.
        window.gl_swap_window();

        for event in event_pump.poll_iter() {
            // check if ESC key pressed or windows closed.
            match event {
                sdl2::event::Event::Quit { .. } => {
                    return;
                }
                sdl2::event::Event::KeyDown {
                            timestamp: _,
                            window_id: _,
                            keycode: _,
                            scancode,
                            keymod: _,
                            repeat: _,
                        } => {
                    if scancode == sdl2::scancode::ScanCode::Escape {
                        return;
                    }
                }
                _ => ()
            }
        }
    }
}

<|endoftext|>


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

/// Key-Value container that uses DNS `TXT` RDATA as representation
///
/// The binary representation can be used as RDATA for `DNS-SD TXT
/// Records` (see [RFC 6763, section 6]).
///
/// Each entry results in one string in the `TXT` represenation; `TXT`
/// RDATA contains many (but at least one) possibly empty strings, each
/// up to 255 bytes.
///
/// Key and value are separated by the first `=` in an entry, and the
/// key must consist of printable ASCII characters (0x20...0x7E) apart
/// from `=`.  Keys should be 9 characters or fewer.
///
/// Values can be any binary string (but the total length of an entry
/// cannot exceed 255 bytes).
///
/// An entry also can have no value at all (which is different from
/// having an empty value) if there is no `=` separator in the entry.
///
/// [RFC 6763, section 6]: https://tools.ietf.org/html/rfc6763#section-6
///     "RFC 6763, 6. Data Syntax for DNS-SD TXT Records"
#[derive(Clone)]
pub struct TxtRecord(Vec<u8>);

impl TxtRecord {
	/// Constructs a new, empty `TxtRecord`.
	pub fn new() -> Self {
		TxtRecord(Vec::new())
	}

	/// Constructs a new, empty `TxtRecord` with the specified capacity.
	///
	/// The inserting operations will still reallocate if necessary.
	pub fn with_capacity(capacity: usize) -> Self {
		TxtRecord(Vec::with_capacity(capacity))
	}

	/// Reserves capacity for at least `additional` more bytes to be
	/// used by inserting operations.
	///
	/// Each entry requires 1 byte for the total length, the length
	/// of the key for the key; if there is a value 1 byte for the
	/// separator `=` and the length of the value for the value.
	pub fn reserve(&mut self, additional: usize) {
		self.0.reserve(additional);
	}

	/// Returns `true` if the `TxtRecord` contains no elements (both in
	/// bytes and key-value entries).
	pub fn is_empty(&self) -> bool {
		self.0.is_empty()
	}

	/// Clears the `TxtRecord`, removing all entries.
	pub fn clear(&mut self) {
		self.0.clear();
	}

	/// if not empty this returns valid TXT RDATA, otherwise just an
	/// empty slice.
	pub fn data(&self) -> &[u8] {
		&self.0
	}

	/// always returns valid TXT RDATA; when the container is empty it
	/// will return a TXT record with a single empty string (i.e.
	/// `&[0x00]`).
	pub fn rdata(&self) -> &[u8] {
		if self.0.is_empty() {
			&[0x00] // empty RDATA not allowed, use single empty chunk instead
		} else {
			&self.0
		}
	}

	fn _position_keys(&self) -> PositionKeyIter {
		PositionKeyIter {
			pos: 0,
			data: &self.0,
		}
	}

	/// Iterate over all `(key, value)` pairs.
	pub fn iter(&self) -> TxtRecordIter {
		TxtRecordIter {
			pos: 0,
			data: &self.0,
		}
	}

	/// Get value for entry with given key
	///
	/// Returns `None` if there is no such entry, `Some(None)` if the
	/// entry exists but has no value, and `Some(Some(value))` if the
	/// entry exists and has a value.
	#[cfg_attr(feature = "cargo-clippy", allow(option_option))]
	pub fn get(&self, key: &[u8]) -> Option<Option<&[u8]>> {
		self.iter().find(|&(k, _)| key == k).map(|(_, value)| value)
	}

	/// Remove entry with given key (if it exists)
	pub fn remove(&mut self, key: &[u8]) {
		if let Some((loc, _)) = self._position_keys().find(|&(_, k)| key == k) {
			self.0.drain(loc);
		}
	}

	/// Insert or update the entry with `key` to have the given value or on value
	pub fn set(
		&mut self,
		key: &[u8],
		value: Option<&[u8]>,
	) -> Result<(), TxtRecordError> {
		for &k in key {
			if k == b'=' || k < 0x20 || k > 0x7e {
				return Err(TxtRecordError::InvalidKey);
			}
		}
		let entry_len = key.len() + value.map(|v| v.len() + 1).unwrap_or(0);
		if entry_len > 255 {
			return Err(TxtRecordError::EntryTooLong);
		}
		self.remove(key);

		self.0.push(entry_len as u8);
		self.0.extend_from_slice(key);
		if let Some(value) = value {
			self.0.push(b'=');
			self.0.extend_from_slice(value);
		}

		Ok(())
	}

	/// Insert or update the entry with `key` to have no value
	pub fn set_no_value(&mut self, key: &[u8]) -> Result<(), TxtRecordError> {
		self.set(key, None)
	}

	/// Insert or update the entry with `key` to have the given value
	pub fn set_value(
		&mut self,
		key: &[u8],
		value: &[u8],
	) -> Result<(), TxtRecordError> {
		self.set(key, Some(value))
	}
}

impl Default for TxtRecord {
	fn default() -> Self {
		TxtRecord::new()
	}
}

impl<'a> IntoIterator for &'a TxtRecord {
	type IntoIter = TxtRecordIter<'a>;
	type Item = (&'a [u8], Option<&'a [u8]>);

	fn into_iter(self) -> Self::IntoIter {
		self.iter()
	}
}

/// Error returned when inserting new entries failed
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub enum TxtRecordError {
	/// Key contained invalid characters
	InvalidKey,
	/// Total entry would be longer than 255 bytes
	EntryTooLong,
}

struct PositionKeyIter<'a> {
	pos: usize,
	data: &'a [u8],
}

impl<'a> Iterator for PositionKeyIter<'a> {
	// (start..end, key)
	type Item = (Range<usize>, &'a [u8]);

	fn next(&mut self) -> Option<Self::Item> {
		if self.data.is_empty() {
			return None;
		}
		let len = self.data[0] as usize;
		let entry_pos = self.pos;
		let entry = &self.data[1..][..len];
		self.data = &self.data[len + 1..];
		self.pos += len + 1;

		Some(match entry.iter().position(|&b| b == b'=') {
			Some(pos) => (entry_pos..self.pos, &entry[..pos]),
			None => (entry_pos..self.pos, entry),
		})
	}
}

/// Iterator for entries in `TxtRecord`
///
/// Items are `(key, value)` pairs.
pub struct TxtRecordIter<'a> {
	pos: usize,
	data: &'a [u8],
}

impl<'a> Iterator for TxtRecordIter<'a> {
	// key, value
	type Item = (&'a [u8], Option<&'a [u8]>);

	fn next(&mut self) -> Option<Self::Item> {
		if self.data.is_empty() {
			return None;
		}
		let len = self.data[0] as usize;
		let entry = &self.data[1..][..len];
		self.data = &self.data[len + 1..];
		self.pos += len + 1;

		Some(match entry.iter().position(|&b| b == b'=') {
			Some(pos) => (&entry[..pos], Some(&entry[pos + 1..])),
			None => (entry, None),
		})
	}
}

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

	#[test]
	fn modifications() {
		let mut r = TxtRecord::new();
		assert!(r.is_empty());
		assert_eq!(r.data(), b"");
		assert_eq!(r.rdata(), b"\x00");

		r.set(b"foo", Some(b"bar")).unwrap();
		assert!(!r.is_empty());
		assert_eq!(r.data(), b"\x07foo=bar");
		assert_eq!(r.rdata(), b"\x07foo=bar");

		r.set(b"u", Some(b"vw")).unwrap();
		assert!(!r.is_empty());
		assert_eq!(r.data(), b"\x07foo=bar\x04u=vw");
		assert_eq!(r.rdata(), b"\x07foo=bar\x04u=vw");
		assert_eq!(
			r.iter().collect::<Vec<_>>(),
			vec![
				(b"foo" as &[u8], Some(b"bar" as &[u8])),
				(b"u", Some(b"vw")),
			]
		);

		r.set(b"foo", None).unwrap();
		assert!(!r.is_empty());
		assert_eq!(r.data(), b"\x04u=vw\x03foo");
		assert_eq!(r.rdata(), b"\x04u=vw\x03foo");
		assert_eq!(
			r.iter().collect::<Vec<_>>(),
			vec![(b"u" as &[u8], Some(b"vw" as &[u8])), (b"foo", None),]
		);

		r.set(b"foo", Some(b"bar")).unwrap();
		assert!(!r.is_empty());
		assert_eq!(r.data(), b"\x04u=vw\x07foo=bar");
		assert_eq!(r.rdata(), b"\x04u=vw\x07foo=bar");

		r.remove(b"foo");
		assert!(!r.is_empty());
		assert_eq!(r.data(), b"\x04u=vw");
		assert_eq!(r.rdata(), b"\x04u=vw");
	}
}
<commit_msg>add TxtRecord::parse[_vec]<commit_after>use std::ops::Range;

/// Key-Value container that uses DNS `TXT` RDATA as representation
///
/// The binary representation can be used as RDATA for `DNS-SD TXT
/// Records` (see [RFC 6763, section 6]).
///
/// Each entry results in one string in the `TXT` represenation; `TXT`
/// RDATA contains many (but at least one) possibly empty strings, each
/// up to 255 bytes.
///
/// Key and value are separated by the first `=` in an entry, and the
/// key must consist of printable ASCII characters (0x20...0x7E) apart
/// from `=`.  Keys should be 9 characters or fewer.
///
/// Values can be any binary string (but the total length of an entry
/// cannot exceed 255 bytes).
///
/// An entry also can have no value at all (which is different from
/// having an empty value) if there is no `=` separator in the entry.
///
/// [RFC 6763, section 6]: https://tools.ietf.org/html/rfc6763#section-6
///     "RFC 6763, 6. Data Syntax for DNS-SD TXT Records"
#[derive(Clone)]
pub struct TxtRecord(Vec<u8>);

impl TxtRecord {
	/// Constructs a new, empty `TxtRecord`.
	pub fn new() -> Self {
		TxtRecord(Vec::new())
	}

	/// Parse binary blob as TXT RDATA
	///
	/// Same as [`parse`] but takes ownership of buffer.
	///
	/// [`parse`]: #method.parse
	pub fn parse_vec(data: Vec<u8>) -> Option<Self> {
		if data.len() == 1 && data[0] == 0 {
			let mut data = data;
			data.clear();
			return Some(TxtRecord(data));
		}
		let mut pos = 0;
		while pos < data.len() {
			let len = data[pos] as usize;
			let new_pos = pos + 1 + len;
			if new_pos > data.len() {
				return None;
			}
			pos = new_pos;
		}
		Some(TxtRecord(data))
	}

	/// Parse some binary blob as TXT RDATA
	///
	/// A single empty string (encoded as `0x00`) gets decoded as "empty" `TxtRecord` (i.e. the
	/// reverse th `rdata()`); an empty slice is treated the same, although it wouldn't be valid
	/// RDATA.
	///
	/// This only fails when the length of a chunk exceeds the remaining data.
	pub fn parse(data: &[u8]) -> Option<Self> {
		Self::parse_vec(data.into())
	}

	/// Constructs a new, empty `TxtRecord` with the specified capacity.
	///
	/// The inserting operations will still reallocate if necessary.
	pub fn with_capacity(capacity: usize) -> Self {
		TxtRecord(Vec::with_capacity(capacity))
	}

	/// Reserves capacity for at least `additional` more bytes to be
	/// used by inserting operations.
	///
	/// Each entry requires 1 byte for the total length, the length
	/// of the key for the key; if there is a value 1 byte for the
	/// separator `=` and the length of the value for the value.
	pub fn reserve(&mut self, additional: usize) {
		self.0.reserve(additional);
	}

	/// Returns `true` if the `TxtRecord` contains no elements (both in
	/// bytes and key-value entries).
	pub fn is_empty(&self) -> bool {
		self.0.is_empty()
	}

	/// Clears the `TxtRecord`, removing all entries.
	pub fn clear(&mut self) {
		self.0.clear();
	}

	/// if not empty this returns valid TXT RDATA, otherwise just an
	/// empty slice.
	pub fn data(&self) -> &[u8] {
		&self.0
	}

	/// always returns valid TXT RDATA; when the container is empty it
	/// will return a TXT record with a single empty string (i.e.
	/// `&[0x00]`).
	pub fn rdata(&self) -> &[u8] {
		if self.0.is_empty() {
			&[0x00] // empty RDATA not allowed, use single empty chunk instead
		} else {
			&self.0
		}
	}

	fn _position_keys(&self) -> PositionKeyIter {
		PositionKeyIter {
			pos: 0,
			data: &self.0,
		}
	}

	/// Iterate over all `(key, value)` pairs.
	pub fn iter(&self) -> TxtRecordIter {
		TxtRecordIter {
			pos: 0,
			data: &self.0,
		}
	}

	/// Get value for entry with given key
	///
	/// Returns `None` if there is no such entry, `Some(None)` if the
	/// entry exists but has no value, and `Some(Some(value))` if the
	/// entry exists and has a value.
	#[cfg_attr(feature = "cargo-clippy", allow(option_option))]
	pub fn get(&self, key: &[u8]) -> Option<Option<&[u8]>> {
		self.iter().find(|&(k, _)| key == k).map(|(_, value)| value)
	}

	/// Remove entry with given key (if it exists)
	pub fn remove(&mut self, key: &[u8]) {
		if let Some((loc, _)) = self._position_keys().find(|&(_, k)| key == k) {
			self.0.drain(loc);
		}
	}

	/// Insert or update the entry with `key` to have the given value or on value
	pub fn set(
		&mut self,
		key: &[u8],
		value: Option<&[u8]>,
	) -> Result<(), TxtRecordError> {
		for &k in key {
			if k == b'=' || k < 0x20 || k > 0x7e {
				return Err(TxtRecordError::InvalidKey);
			}
		}
		let entry_len = key.len() + value.map(|v| v.len() + 1).unwrap_or(0);
		if entry_len > 255 {
			return Err(TxtRecordError::EntryTooLong);
		}
		self.remove(key);

		self.0.push(entry_len as u8);
		self.0.extend_from_slice(key);
		if let Some(value) = value {
			self.0.push(b'=');
			self.0.extend_from_slice(value);
		}

		Ok(())
	}

	/// Insert or update the entry with `key` to have no value
	pub fn set_no_value(&mut self, key: &[u8]) -> Result<(), TxtRecordError> {
		self.set(key, None)
	}

	/// Insert or update the entry with `key` to have the given value
	pub fn set_value(
		&mut self,
		key: &[u8],
		value: &[u8],
	) -> Result<(), TxtRecordError> {
		self.set(key, Some(value))
	}
}

impl Default for TxtRecord {
	fn default() -> Self {
		TxtRecord::new()
	}
}

impl<'a> IntoIterator for &'a TxtRecord {
	type IntoIter = TxtRecordIter<'a>;
	type Item = (&'a [u8], Option<&'a [u8]>);

	fn into_iter(self) -> Self::IntoIter {
		self.iter()
	}
}

/// Error returned when inserting new entries failed
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub enum TxtRecordError {
	/// Key contained invalid characters
	InvalidKey,
	/// Total entry would be longer than 255 bytes
	EntryTooLong,
}

struct PositionKeyIter<'a> {
	pos: usize,
	data: &'a [u8],
}

impl<'a> Iterator for PositionKeyIter<'a> {
	// (start..end, key)
	type Item = (Range<usize>, &'a [u8]);

	fn next(&mut self) -> Option<Self::Item> {
		if self.data.is_empty() {
			return None;
		}
		let len = self.data[0] as usize;
		let entry_pos = self.pos;
		let entry = &self.data[1..][..len];
		self.data = &self.data[len + 1..];
		self.pos += len + 1;

		Some(match entry.iter().position(|&b| b == b'=') {
			Some(pos) => (entry_pos..self.pos, &entry[..pos]),
			None => (entry_pos..self.pos, entry),
		})
	}
}

/// Iterator for entries in `TxtRecord`
///
/// Items are `(key, value)` pairs.
pub struct TxtRecordIter<'a> {
	pos: usize,
	data: &'a [u8],
}

impl<'a> Iterator for TxtRecordIter<'a> {
	// key, value
	type Item = (&'a [u8], Option<&'a [u8]>);

	fn next(&mut self) -> Option<Self::Item> {
		if self.data.is_empty() {
			return None;
		}
		let len = self.data[0] as usize;
		let entry = &self.data[1..][..len];
		self.data = &self.data[len + 1..];
		self.pos += len + 1;

		Some(match entry.iter().position(|&b| b == b'=') {
			Some(pos) => (&entry[..pos], Some(&entry[pos + 1..])),
			None => (entry, None),
		})
	}
}

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

	#[test]
	fn modifications() {
		let mut r = TxtRecord::new();
		assert!(r.is_empty());
		assert_eq!(r.data(), b"");
		assert_eq!(r.rdata(), b"\x00");

		r.set(b"foo", Some(b"bar")).unwrap();
		assert!(!r.is_empty());
		assert_eq!(r.data(), b"\x07foo=bar");
		assert_eq!(r.rdata(), b"\x07foo=bar");

		r.set(b"u", Some(b"vw")).unwrap();
		assert!(!r.is_empty());
		assert_eq!(r.data(), b"\x07foo=bar\x04u=vw");
		assert_eq!(r.rdata(), b"\x07foo=bar\x04u=vw");
		assert_eq!(
			r.iter().collect::<Vec<_>>(),
			vec![
				(b"foo" as &[u8], Some(b"bar" as &[u8])),
				(b"u", Some(b"vw")),
			]
		);

		r.set(b"foo", None).unwrap();
		assert!(!r.is_empty());
		assert_eq!(r.data(), b"\x04u=vw\x03foo");
		assert_eq!(r.rdata(), b"\x04u=vw\x03foo");
		assert_eq!(
			r.iter().collect::<Vec<_>>(),
			vec![(b"u" as &[u8], Some(b"vw" as &[u8])), (b"foo", None),]
		);

		r.set(b"foo", Some(b"bar")).unwrap();
		assert!(!r.is_empty());
		assert_eq!(r.data(), b"\x04u=vw\x07foo=bar");
		assert_eq!(r.rdata(), b"\x04u=vw\x07foo=bar");

		r.remove(b"foo");
		assert!(!r.is_empty());
		assert_eq!(r.data(), b"\x04u=vw");
		assert_eq!(r.rdata(), b"\x04u=vw");
	}
}
<|endoftext|>