<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

//! This provides a simple networking substrate between a client and server. It is assumed that all
//! operations are blocking and return only complete blocks of data. The intended use case has the
//! server blocking on read.  Upon receiving a payload during a read, the server should process the
//! payload, write a response, and then block on read again. The client should block on read after
//! performing a write. Upon errors or remote disconnections, the call (read, write) will return an
//! error to let the caller know of the event. A follow up call will result in the service
//! attempting to either reconnect in the case of a client or accept a new client in the case of a
//! server.
//!
//! Internally both the client and server leverage a NetworkStream that communications in blocks
//! where a block is a length prefixed array of bytes.

use libra_logger::{debug, trace};
use std::{
    io::{Read, Write},
    net::{Shutdown, SocketAddr, TcpListener, TcpStream},
    thread, time,
};
use thiserror::Error;

#[derive(Debug, Error)]
pub enum Error {
    #[error("Already called shutdown")]
    AlreadyShutdown,
    #[error("Found data that is too large to decode: {0}")]
    DataTooLarge(usize),
    #[error("Internal network error:")]
    NetworkError(#[from] std::io::Error),
    #[error("No active stream")]
    NoActiveStream,
    #[error("Remote stream cleanly closed")]
    RemoteStreamClosed,
}

pub struct NetworkClient {
    server: SocketAddr,
    stream: Option<NetworkStream>,
}

impl NetworkClient {
    pub fn new(server: SocketAddr) -> Self {
        Self {
            server,
            stream: None,
        }
    }

    /// Blocking read until able to successfully read an entire message
    pub fn read(&mut self) -> Result<Vec<u8>, Error> {
        let stream = self.server()?;
        let result = stream.read();
        if result.is_err() {
            debug!("On read, upstream peer disconnected, setting stream to None");
            self.stream = None;
        }
        result
    }

    /// Shutdown the internal network stream
    pub fn shutdown(&mut self) -> Result<(), Error> {
        debug!("Shutdown called");
        let stream = self.stream.take().ok_or_else(|| Error::NoActiveStream)?;
        stream.shutdown()?;
        Ok(())
    }

    /// Blocking write until able to successfully send an entire message
    pub fn write(&mut self, data: &[u8]) -> Result<(), Error> {
        let stream = self.server()?;
        let result = stream.write(data);
        if result.is_err() {
            debug!("On write, upstream peer disconnected, setting stream to None");
            self.stream = None;
        }
        result
    }

    fn server(&mut self) -> Result<&mut NetworkStream, Error> {
        if self.stream.is_none() {
            debug!("Attempting to connect to upstream {}", self.server);
            let mut stream = TcpStream::connect(self.server);

            let sleeptime = time::Duration::from_millis(100);
            while let Err(e) = stream {
                debug!("Failed to connect to upstream {} {:?}", self.server, e);
                thread::sleep(sleeptime);
                stream = TcpStream::connect(self.server);
            }

            let stream = stream?;
            stream.set_nodelay(true)?;
            self.stream = Some(NetworkStream::new(stream));
            debug!("Connection established to upstream {}", self.server);
        }

        self.stream.as_mut().ok_or_else(|| Error::NoActiveStream)
    }
}

pub struct NetworkServer {
    listener: Option<TcpListener>,
    stream: Option<NetworkStream>,
}

impl NetworkServer {
    pub fn new(listen: SocketAddr) -> Self {
        let listener = TcpListener::bind(listen);
        Self {
            listener: Some(listener.unwrap()),
            stream: None,
        }
    }

    /// If there isn't already a downstream client, it accepts. Otherwise it
    /// blocks until able to successfully read an entire message
    pub fn read(&mut self) -> Result<Vec<u8>, Error> {
        let stream = self.client()?;
        let result = stream.read();
        if result.is_err() {
            debug!("On read, downstream peer disconnected, setting stream to None");
            self.stream = None;
        }
        result
    }

    /// Shutdown the internal network stream
    pub fn shutdown(&mut self) -> Result<(), Error> {
        debug!("Shutdown called");
        self.listener.take().ok_or_else(|| Error::AlreadyShutdown)?;
        let stream = self.stream.take().ok_or_else(|| Error::NoActiveStream)?;
        stream.shutdown()?;
        Ok(())
    }

    /// If there isn't already a downstream client, it accepts. Otherwise it
    /// blocks until it is able to successfully send an entire message.
    pub fn write(&mut self, data: &[u8]) -> Result<(), Error> {
        let stream = self.client()?;
        let result = stream.write(data);
        if result.is_err() {
            debug!("On write, downstream peer disconnected, setting stream to None");
            self.stream = None;
        }
        result
    }

    fn client(&mut self) -> Result<&mut NetworkStream, Error> {
        if self.stream.is_none() {
            debug!("Waiting for downstream to connect");
            let listener = self
                .listener
                .as_mut()
                .ok_or_else(|| Error::AlreadyShutdown)?;
            let (stream, stream_addr) = listener.accept()?;
            debug!("Connection established with downstream {}", stream_addr);
            stream.set_nodelay(true)?;
            self.stream = Some(NetworkStream::new(stream));
        }

        self.stream.as_mut().ok_or_else(|| Error::NoActiveStream)
    }
}

struct NetworkStream {
    stream: TcpStream,
    buffer: Vec<u8>,
    temp_buffer: [u8; 1024],
}

impl NetworkStream {
    pub fn new(stream: TcpStream) -> Self {
        Self {
            stream,
            buffer: Vec::new(),
            temp_buffer: [0; 1024],
        }
    }

    /// Blocking read until able to successfully read an entire message
    pub fn read(&mut self) -> Result<Vec<u8>, Error> {
        let result = self.read_buffer();
        if !result.is_empty() {
            return Ok(result);
        }

        loop {
            trace!("Attempting to read from stream");
            let read = self.stream.read(&mut self.temp_buffer)?;
            trace!("Read {} bytes from stream", read);
            if read == 0 {
                return Err(Error::RemoteStreamClosed);
            }
            self.buffer.extend(self.temp_buffer[..read].to_vec());
            let result = self.read_buffer();
            if !result.is_empty() {
                trace!("Found a message in the stream");
                return Ok(result);
            }
            trace!("Did not find a message yet, reading again");
        }
    }

    /// Terminate the socket
    pub fn shutdown(&self) -> Result<(), Error> {
        Ok(self.stream.shutdown(Shutdown::Both)?)
    }

    /// Blocking write until able to successfully send an entire message
    pub fn write(&mut self, data: &[u8]) -> Result<(), Error> {
        let u32_max = u32::max_value() as usize;
        if u32_max <= data.len() {
            return Err(Error::DataTooLarge(data.len()));
        }
        let data_len = data.len() as u32;
        trace!("Attempting to write length, {},  to the stream", data_len);
        self.write_all(&data_len.to_le_bytes())?;
        trace!("Attempting to write data, {},  to the stream", data_len);
        self.write_all(data)?;
        trace!(
            "Successfully wrote length, {}, and data to the stream",
            data_len
        );
        Ok(())
    }

    /// Data sent on a TCP socket may not necessarily be delivered at the exact time. So a read may
    /// only include a subset of what was sent. This wraps around the TCP read buffer to ensure
    /// that only full messages are received.
    fn read_buffer(&mut self) -> Vec<u8> {
        if self.buffer.len() < 4 {
            return Vec::new();
        }

        let mut u32_bytes = [0; 4];
        u32_bytes.copy_from_slice(&self.buffer[..4]);
        let data_size = u32::from_le_bytes(u32_bytes) as usize;

        let remaining_data = &self.buffer[4..];
        if remaining_data.len() < data_size {
            return Vec::new();
        }

        let returnable_data = remaining_data[..data_size].to_vec();
        self.buffer = remaining_data[data_size..].to_vec();
        returnable_data
    }

    /// Writing to a TCP socket will take in as much data as the underlying buffer has space for.
    /// This wraps around that buffer and blocks until all the data has been pushed.
    fn write_all(&mut self, data: &[u8]) -> Result<(), Error> {
        let mut unwritten = &data[..];
        let mut total_written = 0;

        while !unwritten.is_empty() {
            let written = self.stream.write(unwritten)?;
            total_written += written;
            unwritten = &data[total_written..];
        }
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use libra_config::utils;
    use std::net::{IpAddr, Ipv4Addr, SocketAddr};

    #[test]
    fn test_ping() {
        let server_port = utils::get_available_port();
        let server_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), server_port);
        let mut server = NetworkServer::new(server_addr);
        let mut client = NetworkClient::new(server_addr);

        let data = vec![0, 1, 2, 3];
        client.write(&data).unwrap();
        let result = server.read().unwrap();
        assert_eq!(data, result);

        let data = vec![4, 5, 6, 7];
        server.write(&data).unwrap();
        let result = client.read().unwrap();
        assert_eq!(data, result);
    }

    #[test]
    fn test_client_shutdown() {
        let server_port = utils::get_available_port();
        let server_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), server_port);
        let mut server = NetworkServer::new(server_addr);
        let mut client = NetworkClient::new(server_addr);

        let data = vec![0, 1, 2, 3];
        client.write(&data).unwrap();
        let result = server.read().unwrap();
        assert_eq!(data, result);

        client.shutdown().unwrap();
        let mut client = NetworkClient::new(server_addr);
        assert!(server.read().is_err());

        let data = vec![4, 5, 6, 7];
        client.write(&data).unwrap();
        let result = server.read().unwrap();
        assert_eq!(data, result);
    }

    #[test]
    fn test_server_shutdown() {
        let server_port = utils::get_available_port();
        let server_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), server_port);
        let mut server = NetworkServer::new(server_addr);
        let mut client = NetworkClient::new(server_addr);

        let data = vec![0, 1, 2, 3];
        client.write(&data).unwrap();
        let result = server.read().unwrap();
        assert_eq!(data, result);

        server.shutdown().unwrap();
        let mut server = NetworkServer::new(server_addr);

        let data = vec![4, 5, 6, 7];
        // We aren't notified immediately that a server has shutdown, but it happens eventually
        while client.write(&data).is_ok() {}

        let data = vec![8, 9, 10, 11];
        client.write(&data).unwrap();
        let result = server.read().unwrap();
        assert_eq!(data, result);
    }

    #[test]
    fn test_write_two_messages_buffered() {
        let server_port = utils::get_available_port();
        let server_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), server_port);
        let mut server = NetworkServer::new(server_addr);
        let mut client = NetworkClient::new(server_addr);
        let data1 = vec![0, 1, 2, 3];
        let data2 = vec![4, 5, 6, 7];
        client.write(&data1).unwrap();
        client.write(&data2).unwrap();
        let result1 = server.read().unwrap();
        let result2 = server.read().unwrap();
        assert_eq!(data1, result1);
        assert_eq!(data2, result2);
    }
}
<commit_msg>[secure-net] add timeouts<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

//! This provides a simple networking substrate between a client and server. It is assumed that all
//! operations are blocking and return only complete blocks of data. The intended use case has the
//! server blocking on read.  Upon receiving a payload during a read, the server should process the
//! payload, write a response, and then block on read again. The client should block on read after
//! performing a write. Upon errors or remote disconnections, the call (read, write) will return an
//! error to let the caller know of the event. A follow up call will result in the service
//! attempting to either reconnect in the case of a client or accept a new client in the case of a
//! server.
//!
//! Internally both the client and server leverage a NetworkStream that communications in blocks
//! where a block is a length prefixed array of bytes.

use libra_logger::{debug, trace};
use std::{
    io::{Read, Write},
    net::{Shutdown, SocketAddr, TcpListener, TcpStream},
    thread, time,
};
use thiserror::Error;

/// Read, Write, Connect timeout in seconds.
#[cfg(not(test))]
const TIMEOUT: u64 = 30;
#[cfg(test)]
const TIMEOUT: u64 = 5;

#[derive(Debug, Error)]
pub enum Error {
    #[error("Already called shutdown")]
    AlreadyShutdown,
    #[error("Found data that is too large to decode: {0}")]
    DataTooLarge(usize),
    #[error("Internal network error:")]
    NetworkError(#[from] std::io::Error),
    #[error("No active stream")]
    NoActiveStream,
    #[error("Remote stream cleanly closed")]
    RemoteStreamClosed,
}

pub struct NetworkClient {
    server: SocketAddr,
    stream: Option<NetworkStream>,
}

impl NetworkClient {
    pub fn new(server: SocketAddr) -> Self {
        Self {
            server,
            stream: None,
        }
    }

    /// Blocking read until able to successfully read an entire message
    pub fn read(&mut self) -> Result<Vec<u8>, Error> {
        let stream = self.server()?;
        let result = stream.read();
        if result.is_err() {
            debug!("On read, upstream peer disconnected, setting stream to None");
            self.stream = None;
        }
        result
    }

    /// Shutdown the internal network stream
    pub fn shutdown(&mut self) -> Result<(), Error> {
        debug!("Shutdown called");
        let stream = self.stream.take().ok_or_else(|| Error::NoActiveStream)?;
        stream.shutdown()?;
        Ok(())
    }

    /// Blocking write until able to successfully send an entire message
    pub fn write(&mut self, data: &[u8]) -> Result<(), Error> {
        let stream = self.server()?;
        let result = stream.write(data);
        if result.is_err() {
            debug!("On write, upstream peer disconnected, setting stream to None");
            self.stream = None;
        }
        result
    }

    fn server(&mut self) -> Result<&mut NetworkStream, Error> {
        if self.stream.is_none() {
            let timeout = std::time::Duration::from_secs(TIMEOUT);
            debug!("Attempting to connect to upstream {}", self.server);
            let mut stream = TcpStream::connect_timeout(&self.server, timeout);

            let sleeptime = time::Duration::from_millis(100);
            while let Err(e) = stream {
                debug!("Failed to connect to upstream {} {:?}", self.server, e);
                thread::sleep(sleeptime);
                stream = TcpStream::connect_timeout(&self.server, timeout);
            }

            let stream = stream?;
            stream.set_nodelay(true)?;
            self.stream = Some(NetworkStream::new(stream));
            debug!("Connection established to upstream {}", self.server);
        }

        self.stream.as_mut().ok_or_else(|| Error::NoActiveStream)
    }
}

pub struct NetworkServer {
    listener: Option<TcpListener>,
    stream: Option<NetworkStream>,
}

impl NetworkServer {
    pub fn new(listen: SocketAddr) -> Self {
        let listener = TcpListener::bind(listen);
        Self {
            listener: Some(listener.unwrap()),
            stream: None,
        }
    }

    /// If there isn't already a downstream client, it accepts. Otherwise it
    /// blocks until able to successfully read an entire message
    pub fn read(&mut self) -> Result<Vec<u8>, Error> {
        let stream = self.client()?;
        let result = stream.read();
        if result.is_err() {
            debug!("On read, downstream peer disconnected, setting stream to None");
            self.stream = None;
        }
        result
    }

    /// Shutdown the internal network stream
    pub fn shutdown(&mut self) -> Result<(), Error> {
        debug!("Shutdown called");
        self.listener.take().ok_or_else(|| Error::AlreadyShutdown)?;
        let stream = self.stream.take().ok_or_else(|| Error::NoActiveStream)?;
        stream.shutdown()?;
        Ok(())
    }

    /// If there isn't already a downstream client, it accepts. Otherwise it
    /// blocks until it is able to successfully send an entire message.
    pub fn write(&mut self, data: &[u8]) -> Result<(), Error> {
        let stream = self.client()?;
        let result = stream.write(data);
        if result.is_err() {
            debug!("On write, downstream peer disconnected, setting stream to None");
            self.stream = None;
        }
        result
    }

    fn client(&mut self) -> Result<&mut NetworkStream, Error> {
        if self.stream.is_none() {
            debug!("Waiting for downstream to connect");
            let listener = self
                .listener
                .as_mut()
                .ok_or_else(|| Error::AlreadyShutdown)?;
            let (stream, stream_addr) = listener.accept()?;
            debug!("Connection established with downstream {}", stream_addr);
            stream.set_nodelay(true)?;
            self.stream = Some(NetworkStream::new(stream));
        }

        self.stream.as_mut().ok_or_else(|| Error::NoActiveStream)
    }
}

struct NetworkStream {
    stream: TcpStream,
    buffer: Vec<u8>,
    temp_buffer: [u8; 1024],
}

impl NetworkStream {
    pub fn new(stream: TcpStream) -> Self {
        let timeout = Some(std::time::Duration::from_secs(TIMEOUT));
        // These only fail if a duration of 0 is passed in.
        stream.set_read_timeout(timeout).unwrap();
        stream.set_write_timeout(timeout).unwrap();

        Self {
            stream,
            buffer: Vec::new(),
            temp_buffer: [0; 1024],
        }
    }

    /// Blocking read until able to successfully read an entire message
    pub fn read(&mut self) -> Result<Vec<u8>, Error> {
        let result = self.read_buffer();
        if !result.is_empty() {
            return Ok(result);
        }

        loop {
            trace!("Attempting to read from stream");
            let read = self.stream.read(&mut self.temp_buffer)?;
            trace!("Read {} bytes from stream", read);
            if read == 0 {
                return Err(Error::RemoteStreamClosed);
            }
            self.buffer.extend(self.temp_buffer[..read].to_vec());
            let result = self.read_buffer();
            if !result.is_empty() {
                trace!("Found a message in the stream");
                return Ok(result);
            }
            trace!("Did not find a message yet, reading again");
        }
    }

    /// Terminate the socket
    pub fn shutdown(&self) -> Result<(), Error> {
        Ok(self.stream.shutdown(Shutdown::Both)?)
    }

    /// Blocking write until able to successfully send an entire message
    pub fn write(&mut self, data: &[u8]) -> Result<(), Error> {
        let u32_max = u32::max_value() as usize;
        if u32_max <= data.len() {
            return Err(Error::DataTooLarge(data.len()));
        }
        let data_len = data.len() as u32;
        trace!("Attempting to write length, {},  to the stream", data_len);
        self.write_all(&data_len.to_le_bytes())?;
        trace!("Attempting to write data, {},  to the stream", data_len);
        self.write_all(data)?;
        trace!(
            "Successfully wrote length, {}, and data to the stream",
            data_len
        );
        Ok(())
    }

    /// Data sent on a TCP socket may not necessarily be delivered at the exact time. So a read may
    /// only include a subset of what was sent. This wraps around the TCP read buffer to ensure
    /// that only full messages are received.
    fn read_buffer(&mut self) -> Vec<u8> {
        if self.buffer.len() < 4 {
            return Vec::new();
        }

        let mut u32_bytes = [0; 4];
        u32_bytes.copy_from_slice(&self.buffer[..4]);
        let data_size = u32::from_le_bytes(u32_bytes) as usize;

        let remaining_data = &self.buffer[4..];
        if remaining_data.len() < data_size {
            return Vec::new();
        }

        let returnable_data = remaining_data[..data_size].to_vec();
        self.buffer = remaining_data[data_size..].to_vec();
        returnable_data
    }

    /// Writing to a TCP socket will take in as much data as the underlying buffer has space for.
    /// This wraps around that buffer and blocks until all the data has been pushed.
    fn write_all(&mut self, data: &[u8]) -> Result<(), Error> {
        let mut unwritten = &data[..];
        let mut total_written = 0;

        while !unwritten.is_empty() {
            let written = self.stream.write(unwritten)?;
            total_written += written;
            unwritten = &data[total_written..];
        }
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use libra_config::utils;
    use std::net::{IpAddr, Ipv4Addr, SocketAddr};

    #[test]
    fn test_ping() {
        let server_port = utils::get_available_port();
        let server_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), server_port);
        let mut server = NetworkServer::new(server_addr);
        let mut client = NetworkClient::new(server_addr);

        let data = vec![0, 1, 2, 3];
        client.write(&data).unwrap();
        let result = server.read().unwrap();
        assert_eq!(data, result);

        let data = vec![4, 5, 6, 7];
        server.write(&data).unwrap();
        let result = client.read().unwrap();
        assert_eq!(data, result);
    }

    #[test]
    fn test_client_shutdown() {
        let server_port = utils::get_available_port();
        let server_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), server_port);
        let mut server = NetworkServer::new(server_addr);
        let mut client = NetworkClient::new(server_addr);

        let data = vec![0, 1, 2, 3];
        client.write(&data).unwrap();
        let result = server.read().unwrap();
        assert_eq!(data, result);

        client.shutdown().unwrap();
        let mut client = NetworkClient::new(server_addr);
        assert!(server.read().is_err());

        let data = vec![4, 5, 6, 7];
        client.write(&data).unwrap();
        let result = server.read().unwrap();
        assert_eq!(data, result);
    }

    #[test]
    fn test_server_shutdown() {
        let server_port = utils::get_available_port();
        let server_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), server_port);
        let mut server = NetworkServer::new(server_addr);
        let mut client = NetworkClient::new(server_addr);

        let data = vec![0, 1, 2, 3];
        client.write(&data).unwrap();
        let result = server.read().unwrap();
        assert_eq!(data, result);

        server.shutdown().unwrap();
        let mut server = NetworkServer::new(server_addr);

        let data = vec![4, 5, 6, 7];
        // We aren't notified immediately that a server has shutdown, but it happens eventually
        while client.write(&data).is_ok() {}

        let data = vec![8, 9, 10, 11];
        client.write(&data).unwrap();
        let result = server.read().unwrap();
        assert_eq!(data, result);
    }

    #[test]
    fn test_write_two_messages_buffered() {
        let server_port = utils::get_available_port();
        let server_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), server_port);
        let mut server = NetworkServer::new(server_addr);
        let mut client = NetworkClient::new(server_addr);
        let data1 = vec![0, 1, 2, 3];
        let data2 = vec![4, 5, 6, 7];
        client.write(&data1).unwrap();
        client.write(&data2).unwrap();
        let result1 = server.read().unwrap();
        let result2 = server.read().unwrap();
        assert_eq!(data1, result1);
        assert_eq!(data2, result2);
    }

    #[test]
    fn test_server_timeout() {
        let server_port = utils::get_available_port();
        let server_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), server_port);
        let mut server = NetworkServer::new(server_addr);
        let mut client = NetworkClient::new(server_addr);
        let data1 = vec![0, 1, 2, 3];
        let data2 = vec![4, 5, 6, 7];

        // First client, success
        client.write(&data1).unwrap();
        let result1 = server.read().unwrap();
        assert_eq!(data1, result1);

        // Timedout
        server.read().unwrap_err();

        // New client, success, note the previous client connection is still active, the server is
        // actively letting it go due to lack of activity.
        let mut client2 = NetworkClient::new(server_addr);
        client2.write(&data2).unwrap();
        let result2 = server.read().unwrap();
        assert_eq!(data2, result2);
    }

    #[test]
    fn test_client_timeout() {
        let server_port = utils::get_available_port();
        let server_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), server_port);
        let mut server = NetworkServer::new(server_addr);
        let mut client = NetworkClient::new(server_addr);
        let data1 = vec![0, 1, 2, 3];
        let data2 = vec![4, 5, 6, 7];

        // First server success
        client.write(&data1).unwrap();
        let result1 = server.read().unwrap();
        assert_eq!(data1, result1);

        // Timedout, it is hard to simulate a client receiving a write timeout
        client.read().unwrap_err();

        // Clean up old Server listener but keep the stream online. Start a new server, which will
        // be the one the client now connects to.
        server.listener = None;
        let mut server2 = NetworkServer::new(server_addr);

        // Client starts a new stream, success
        client.write(&data2).unwrap();
        let result2 = server2.read().unwrap();
        assert_eq!(data2, result2);
    }
}
<|endoftext|>


<commit_before>use std;
use std::cmp::min;

use arrayvec::ArrayVec;

use coefs::Coefficients;
use consts::MAX_HARMONICS;
use params::BaseParams;
use prev::PrevFrame;

#[derive(Clone)]
pub struct Spectrals(ArrayVec<[f32; MAX_HARMONICS]>);

impl Spectrals {
    pub fn new(coefs: &Coefficients, params: &BaseParams, prev: &PrevFrame) -> Spectrals {
        let mut spectral = ArrayVec::new();

        let indexes = |x: u32| {
            let k = prev.params.harmonics as f32 / params.harmonics as f32 * x as f32;
            (k.trunc() as usize, k.fract())
        };

        let pred = (0.03 * params.harmonics as f32 - 0.05).max(0.4).min(0.7);
        let pred_scaled = pred / params.harmonics as f32;

        let sum = (1...params.harmonics).map(|lambda| {
            let (k, dec) = indexes(lambda);
            (1.0 - dec) * prev.spectrals.get(k).log2() +
                dec * prev.spectrals.get(k + 1).log2()
        }).fold(0.0, |s, x| s + x);

        for l in 1...params.harmonics {
            let (k, dec) = indexes(l);

            spectral.push((
                coefs.get(l as usize) +
                    pred * (1.0 - dec) * prev.spectrals.get(k).log2() +
                    pred * dec * prev.spectrals.get(k + 1).log2() -
                    pred_scaled * sum
            ).exp2());
        }

        Spectrals(spectral)
    }

    pub fn get(&self, l: usize) -> f32 {
        if l == 0 {
            1.0
        } else {
            self.0[min(l, self.0.len()) - 1]
        }
    }
}

impl std::ops::Deref for Spectrals {
    type Target = ArrayVec<[f32; MAX_HARMONICS]>;
    fn deref(&self) -> &Self::Target { &self.0 }
}

impl Default for Spectrals {
    fn default() -> Spectrals {
        Spectrals((0..MAX_HARMONICS).map(|_| 1.0).collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use descramble::{Bootstrap, descramble};
    use gain::Gains;
    use coefs::Coefficients;
    use params::BaseParams;
    use prev::{PrevFrame};

    #[test]
    fn test_spectrals() {
        let chunks = [
            0b001000010010,
            0b110011001100,
            0b111000111000,
            0b111111111111,
            0b10100110101,
            0b00101111010,
            0b01110111011,
            0b00001000,
        ];

        let b = Bootstrap::new(&chunks);
        let p = BaseParams::new(b.unwrap_period());
        let (amps, _, gain_idx) = descramble(&chunks, &p);
        let g = Gains::new(gain_idx, &amps, &p);
        let c = Coefficients::new(&g, &amps, &p);
        let mut prev = PrevFrame::default();
        let s = Spectrals::new(&c, &p, &prev);

        assert_eq!(s.0.len(), 16);

        assert!((s.get(0) - 1.0).abs() < 0.000001);
        assert!((s.get(1) - 0.5306769781475001).abs() < 0.000001);
        assert!((s.get(2) - 0.3535007196676076).abs() < 0.000001);
        assert!((s.get(3) - 0.9173875577243951).abs() < 0.000001);
        assert!((s.get(4) - 0.13169278935782622).abs() < 0.000001);
        assert!((s.get(5) - 4.438599873836148).abs() < 0.00001);
        assert!((s.get(6) - 0.6796441620283217).abs() < 0.000001);
        assert!((s.get(7) - 0.9439604687782126).abs() < 0.000001);
        assert!((s.get(8) - 10.646341109175768).abs() < 0.000001);
        assert!((s.get(9) - 1.3058035207362984).abs() < 0.000001);
        assert!((s.get(10) - 1.566168983101457).abs() < 0.000001);
        assert!((s.get(11) - 8.325591838823936).abs() < 0.00001);
        assert!((s.get(12) - 1.5943520720633444).abs() < 0.000001);
        assert!((s.get(13) - 1.556318849691248).abs() < 0.000001);
        assert!((s.get(14) - 15.882295806333524).abs() < 0.00001);
        assert!((s.get(15) - 11.386537444591385).abs() < 0.00001);
        assert!((s.get(16) - 12.498792573794846).abs() < 0.00001);
        assert!((s.get(17) - 12.498792573794846).abs() < 0.00001);
        assert!((s.get(30) - 12.498792573794846).abs() < 0.00001);
        assert!((s.get(56) - 12.498792573794846).abs() < 0.00001);

        prev.spectrals = s;
        let s = Spectrals::new(&c, &p, &prev);

        assert!((s.get(1) - 0.18026009141598204).abs() < 0.000001);
        assert!((s.get(2) - 0.09466759399487303).abs() < 0.000001);
        assert!((s.get(3) - 0.546543586730317).abs() < 0.000001);
        assert!((s.get(4) - 0.11240949805284323).abs() < 0.000001);
        assert!((s.get(5) - 2.664303862945382).abs() < 0.000001);
        assert!((s.get(6) - 0.7356498197494585).abs() < 0.000001);
        assert!((s.get(7) - 0.6722918529743475).abs() < 0.000001);
        assert!((s.get(8) - 15.748045510026174).abs() < 0.000001);
        assert!((s.get(9) - 2.010522620838216).abs() < 0.000001);
        assert!((s.get(10) - 2.411402725277654).abs() < 0.000001);
        assert!((s.get(11) - 12.818766727159018).abs() < 0.00001);
        assert!((s.get(12) - 2.4547957296486485).abs() < 0.000001);
        assert!((s.get(13) - 2.396236648815911).abs() < 0.000001);
        assert!((s.get(14) - 24.45369037714974).abs() < 0.0001);
        assert!((s.get(15) - 17.53165062111628).abs() < 0.0001);
        assert!((s.get(16) - 19.244170201509185).abs() < 0.0001);
    }
}
<commit_msg>Use collect rather than repeated push<commit_after>use std;
use std::cmp::min;

use arrayvec::ArrayVec;

use coefs::Coefficients;
use consts::MAX_HARMONICS;
use params::BaseParams;
use prev::PrevFrame;

#[derive(Clone)]
pub struct Spectrals(ArrayVec<[f32; MAX_HARMONICS]>);

impl Spectrals {
    pub fn new(coefs: &Coefficients, params: &BaseParams, prev: &PrevFrame) -> Spectrals {
        let mut spectral = ArrayVec::new();

        let indexes = |x: u32| {
            let k = prev.params.harmonics as f32 / params.harmonics as f32 * x as f32;
            (k.trunc() as usize, k.fract())
        };

        let pred = (0.03 * params.harmonics as f32 - 0.05).max(0.4).min(0.7);
        let pred_scaled = pred / params.harmonics as f32;

        let sum = (1...params.harmonics).map(|lambda| {
            let (k, dec) = indexes(lambda);
            (1.0 - dec) * prev.spectrals.get(k).log2() +
                dec * prev.spectrals.get(k + 1).log2()
        }).fold(0.0, |s, x| s + x);

        Spectrals((1...params.harmonics).map(|l| {
            let (k, dec) = indexes(l);

            (
                coefs.get(l as usize) +
                    pred * (1.0 - dec) * prev.spectrals.get(k).log2() +
                    pred * dec * prev.spectrals.get(k + 1).log2() -
                    pred_scaled * sum
            ).exp2()
        }).collect())
    }

    pub fn get(&self, l: usize) -> f32 {
        if l == 0 {
            1.0
        } else {
            self.0[min(l, self.0.len()) - 1]
        }
    }
}

impl std::ops::Deref for Spectrals {
    type Target = ArrayVec<[f32; MAX_HARMONICS]>;
    fn deref(&self) -> &Self::Target { &self.0 }
}

impl Default for Spectrals {
    fn default() -> Spectrals {
        Spectrals((0..MAX_HARMONICS).map(|_| 1.0).collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use descramble::{Bootstrap, descramble};
    use gain::Gains;
    use coefs::Coefficients;
    use params::BaseParams;
    use prev::{PrevFrame};

    #[test]
    fn test_spectrals() {
        let chunks = [
            0b001000010010,
            0b110011001100,
            0b111000111000,
            0b111111111111,
            0b10100110101,
            0b00101111010,
            0b01110111011,
            0b00001000,
        ];

        let b = Bootstrap::new(&chunks);
        let p = BaseParams::new(b.unwrap_period());
        let (amps, _, gain_idx) = descramble(&chunks, &p);
        let g = Gains::new(gain_idx, &amps, &p);
        let c = Coefficients::new(&g, &amps, &p);
        let mut prev = PrevFrame::default();
        let s = Spectrals::new(&c, &p, &prev);

        assert_eq!(s.0.len(), 16);

        assert!((s.get(0) - 1.0).abs() < 0.000001);
        assert!((s.get(1) - 0.5306769781475001).abs() < 0.000001);
        assert!((s.get(2) - 0.3535007196676076).abs() < 0.000001);
        assert!((s.get(3) - 0.9173875577243951).abs() < 0.000001);
        assert!((s.get(4) - 0.13169278935782622).abs() < 0.000001);
        assert!((s.get(5) - 4.438599873836148).abs() < 0.00001);
        assert!((s.get(6) - 0.6796441620283217).abs() < 0.000001);
        assert!((s.get(7) - 0.9439604687782126).abs() < 0.000001);
        assert!((s.get(8) - 10.646341109175768).abs() < 0.000001);
        assert!((s.get(9) - 1.3058035207362984).abs() < 0.000001);
        assert!((s.get(10) - 1.566168983101457).abs() < 0.000001);
        assert!((s.get(11) - 8.325591838823936).abs() < 0.00001);
        assert!((s.get(12) - 1.5943520720633444).abs() < 0.000001);
        assert!((s.get(13) - 1.556318849691248).abs() < 0.000001);
        assert!((s.get(14) - 15.882295806333524).abs() < 0.00001);
        assert!((s.get(15) - 11.386537444591385).abs() < 0.00001);
        assert!((s.get(16) - 12.498792573794846).abs() < 0.00001);
        assert!((s.get(17) - 12.498792573794846).abs() < 0.00001);
        assert!((s.get(30) - 12.498792573794846).abs() < 0.00001);
        assert!((s.get(56) - 12.498792573794846).abs() < 0.00001);

        prev.spectrals = s;
        let s = Spectrals::new(&c, &p, &prev);

        assert!((s.get(1) - 0.18026009141598204).abs() < 0.000001);
        assert!((s.get(2) - 0.09466759399487303).abs() < 0.000001);
        assert!((s.get(3) - 0.546543586730317).abs() < 0.000001);
        assert!((s.get(4) - 0.11240949805284323).abs() < 0.000001);
        assert!((s.get(5) - 2.664303862945382).abs() < 0.000001);
        assert!((s.get(6) - 0.7356498197494585).abs() < 0.000001);
        assert!((s.get(7) - 0.6722918529743475).abs() < 0.000001);
        assert!((s.get(8) - 15.748045510026174).abs() < 0.000001);
        assert!((s.get(9) - 2.010522620838216).abs() < 0.000001);
        assert!((s.get(10) - 2.411402725277654).abs() < 0.000001);
        assert!((s.get(11) - 12.818766727159018).abs() < 0.00001);
        assert!((s.get(12) - 2.4547957296486485).abs() < 0.000001);
        assert!((s.get(13) - 2.396236648815911).abs() < 0.000001);
        assert!((s.get(14) - 24.45369037714974).abs() < 0.0001);
        assert!((s.get(15) - 17.53165062111628).abs() < 0.0001);
        assert!((s.get(16) - 19.244170201509185).abs() < 0.0001);
    }
}
<|endoftext|>


<commit_before>use std;
use std::f32::consts::PI;

use arrayvec::ArrayVec;
use num::complex::Complex32;
use collect_slice::CollectSlice;

use consts::SAMPLES;
use descramble::VoiceDecisions;
use enhance::EnhancedSpectrals;
use noise::Noise;
use params::BaseParams;
use window;

// Result of equation 121
const SCALING_COEF: f32 = 146.6432708443356;

pub struct UnvoicedDFT {
    dfts: ArrayVec<[Complex32; 256]>,
    scale: f32,
}

impl UnvoicedDFT {
    pub fn new(lower: isize, upper: isize, noise: &Noise, window: &window::Window)
        -> UnvoicedDFT
    {
        let dfts = (lower..upper).map(|m| {
            let mut nclone = noise.clone();

            (-104..105).map(|n| {
                nclone.next() * window.get(n) *
                    Complex32::new(0.0, -2.0 / 256.0 * PI * m as f32 * n as f32).exp()
            }).fold(Complex32::new(0.0, 0.0), |s, x| s + x)
        }).collect::<ArrayVec<[Complex32; 256]>>();

        let sum = dfts.iter().map(|dft| dft.norm_sqr()).fold(0.0, |s, x| s + x);

        UnvoicedDFT {
            dfts: dfts,
            scale: SCALING_COEF / (sum / (upper - lower) as f32).sqrt(),
        }
    }

    pub fn scale(&self, spectral: f32) -> f32 { self.scale * spectral }
}

impl std::ops::Deref for UnvoicedDFT {
    type Target = ArrayVec<[Complex32; 256]>;
    fn deref(&self) -> &Self::Target { &self.dfts }
}

fn edges(l: usize, params: &BaseParams) -> (isize, isize) {
    let edge = |inner: f32| {
        256.0 / (2.0 * PI) * inner * params.fundamental
    };

    (
        edge(l as f32 - 0.5).ceil() as isize,
        edge(l as f32 + 0.5).ceil() as isize,
    )
}

pub struct UnvoicedParts([Complex32; 256]);

impl UnvoicedParts {
    pub fn new(params: &BaseParams, voice: &VoiceDecisions, enhanced: &EnhancedSpectrals)
        -> UnvoicedParts
    {
        let mut parts = [Complex32::new(0.0, 0.0); 256];

        let noise = Noise::new();
        let window = window::synthesis_trunc();

        for (l, &m) in enhanced.iter().enumerate() {
            let l = l + 1;

            if voice.is_voiced(l) {
                continue;
            }

            let (lower, upper) = edges(l, params);

            let pos = UnvoicedDFT::new(lower, upper, &noise, &window);
            let pscale = pos.scale(m);

            let neg = UnvoicedDFT::new(-upper+1, -lower+1, &noise, &window);
            let nscale = neg.scale(m);

            pos.iter().map(|&dft| {
                pscale * dft
            }).collect_slice_checked(&mut parts[128 + lower as usize..128 + upper as usize]);

            neg.iter().map(|&dft| {
                nscale * dft
            }).collect_slice_checked(&mut parts[128 - upper as usize + 1..128 - lower as usize + 1]);
        }

        UnvoicedParts(parts)
    }

    // -128 <= m < 128
    fn get(&self, m: isize) -> Complex32 { self.0[(m + 128) as usize] }

    pub fn idft(&self, n: isize) -> f32 {
        if n < -128 || n > 127 {
            return 0.0;
        }

        (-128..128).map(|m| {
            self.get(m) *
                Complex32::new(0.0, 2.0 / 256.0 * PI * m as f32 * n as f32).exp()
        }).fold(Complex32::new(0.0, 0.0), |s, x| s + x).re / 256.0
    }
}

impl Default for UnvoicedParts {
    fn default() -> UnvoicedParts {
        UnvoicedParts([Complex32::new(0.0, 0.0); 256])
    }
}

pub struct Unvoiced<'a, 'b> {
    cur: &'a UnvoicedParts,
    prev: &'b UnvoicedParts,
    window: window::Window,
}

impl<'a, 'b> Unvoiced<'a, 'b> {
    pub fn new(cur: &'a UnvoicedParts, prev: &'b UnvoicedParts) -> Unvoiced<'a, 'b> {
        Unvoiced {
            cur: cur,
            prev: prev,
            window: window::synthesis_full(),
        }
    }

    pub fn get(&self, n: usize) -> f32 {
        let n = n as isize;

        let numer = self.window.get(n) * self.prev.idft(n) +
            self.window.get(n - SAMPLES as isize) * self.cur.idft(n - SAMPLES as isize);
        let denom = self.window.get(n).powi(2) +
            self.window.get(n - SAMPLES as isize).powi(2);

        numer / denom
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use super::edges;
    use descramble::Bootstrap;
    use params::BaseParams;

    #[test]
    fn test_edges() {
        let chunks = [
            0b001000010010,
            0b110011001100,
            0b111000111000,
            0b111111111111,
            0b10100110101,
            0b00101111010,
            0b01110111011,
            0b00001000,
        ];

        let b = Bootstrap::new(&chunks);
        let p = BaseParams::new(b.unwrap_period());

        assert!((p.fundamental - 0.17575344).abs() < 0.000001);

        let (lower, upper) = edges(1, &p);
        assert_eq!(lower, 4);
        assert_eq!(upper, 11);

        let (lower, upper) = edges(2, &p);
        assert_eq!(lower, 11);
        assert_eq!(upper, 18);

        let (lower, upper) = edges(5, &p);
        assert_eq!(lower, 33);
        assert_eq!(upper, 40);

        let (lower, upper) = edges(8, &p);
        assert_eq!(lower, 54);
        assert_eq!(upper, 61);

        let (lower, upper) = edges(16, &p);
        assert_eq!(lower, 111);
        assert_eq!(upper, 119);
    }
}
<commit_msg>Use an average for the unvoiced scaling coefficient<commit_after>use std;
use std::f32::consts::PI;

use arrayvec::ArrayVec;
use num::complex::Complex32;
use collect_slice::CollectSlice;

use consts::SAMPLES;
use descramble::VoiceDecisions;
use enhance::EnhancedSpectrals;
use noise::Noise;
use params::BaseParams;
use window;

// Result of equation 121
const SCALING_NUMER: f32 = 146.6432708443356;
const SCALING_DENOM: f32 = 3.33198263;
const SCALING_COEF: f32 = SCALING_NUMER / SCALING_DENOM;

pub struct UnvoicedDFT(ArrayVec<[Complex32; 256]>);

impl UnvoicedDFT {
    pub fn new(lower: isize, upper: isize, noise: &Noise, window: &window::Window)
        -> UnvoicedDFT
    {
        UnvoicedDFT((lower..upper).map(|m| {
            let mut nclone = noise.clone();

            (-104..105).map(|n| {
                nclone.next() * window.get(n) *
                    Complex32::new(0.0, -2.0 / 256.0 * PI * m as f32 * n as f32).exp()
            }).fold(Complex32::new(0.0, 0.0), |s, x| s + x)
        }).collect())
    }
}

impl std::ops::Deref for UnvoicedDFT {
    type Target = ArrayVec<[Complex32; 256]>;
    fn deref(&self) -> &Self::Target { &self.0 }
}

fn edges(l: usize, params: &BaseParams) -> (isize, isize) {
    let edge = |inner: f32| {
        256.0 / (2.0 * PI) * inner * params.fundamental
    };

    (
        edge(l as f32 - 0.5).ceil() as isize,
        edge(l as f32 + 0.5).ceil() as isize,
    )
}

pub struct UnvoicedParts([Complex32; 256]);

impl UnvoicedParts {
    pub fn new(params: &BaseParams, voice: &VoiceDecisions, enhanced: &EnhancedSpectrals)
        -> UnvoicedParts
    {
        let mut parts = [Complex32::new(0.0, 0.0); 256];

        let noise = Noise::new();
        let window = window::synthesis_trunc();

        for l in 1..params.harmonics as usize+1 {
            if voice.is_voiced(l) {
                continue;
            }

            let (lower, upper) = edges(l, params);
            let scale = SCALING_COEF * enhanced.get(l);

            let pos = UnvoicedDFT::new(lower, upper, &noise, &window);
            let neg = UnvoicedDFT::new(-upper+1, -lower+1, &noise, &window);

            pos.iter()
               .map(|&dft| scale * dft)
               .collect_slice(&mut parts[128 + lower as usize..]);

            neg.iter()
               .map(|&dft| scale * dft)
               .collect_slice(&mut parts[129 - upper as usize..]);
        }

        UnvoicedParts(parts)
    }

    // -128 <= m < 128
    fn get(&self, m: isize) -> Complex32 { self.0[(m + 128) as usize] }

    pub fn idft(&self, n: isize) -> f32 {
        if n < -128 || n > 127 {
            return 0.0;
        }

        (-128..128).map(|m| {
            self.get(m) *
                Complex32::new(0.0, 2.0 / 256.0 * PI * m as f32 * n as f32).exp()
        }).fold(Complex32::new(0.0, 0.0), |s, x| s + x).re / 256.0
    }
}

impl Default for UnvoicedParts {
    fn default() -> UnvoicedParts {
        UnvoicedParts([Complex32::new(0.0, 0.0); 256])
    }
}

pub struct Unvoiced<'a, 'b> {
    cur: &'a UnvoicedParts,
    prev: &'b UnvoicedParts,
    window: window::Window,
}

impl<'a, 'b> Unvoiced<'a, 'b> {
    pub fn new(cur: &'a UnvoicedParts, prev: &'b UnvoicedParts) -> Unvoiced<'a, 'b> {
        Unvoiced {
            cur: cur,
            prev: prev,
            window: window::synthesis_full(),
        }
    }

    pub fn get(&self, n: usize) -> f32 {
        let n = n as isize;

        let numer = self.window.get(n) * self.prev.idft(n) +
            self.window.get(n - SAMPLES as isize) * self.cur.idft(n - SAMPLES as isize);
        let denom = self.window.get(n).powi(2) +
            self.window.get(n - SAMPLES as isize).powi(2);

        numer / denom
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use super::edges;
    use descramble::Bootstrap;
    use params::BaseParams;

    #[test]
    fn test_edges() {
        let chunks = [
            0b001000010010,
            0b110011001100,
            0b111000111000,
            0b111111111111,
            0b10100110101,
            0b00101111010,
            0b01110111011,
            0b00001000,
        ];

        let b = Bootstrap::new(&chunks);
        let p = BaseParams::new(b.unwrap_period());

        assert!((p.fundamental - 0.17575344).abs() < 0.000001);

        let (lower, upper) = edges(1, &p);
        assert_eq!(lower, 4);
        assert_eq!(upper, 11);

        let (lower, upper) = edges(2, &p);
        assert_eq!(lower, 11);
        assert_eq!(upper, 18);

        let (lower, upper) = edges(5, &p);
        assert_eq!(lower, 33);
        assert_eq!(upper, 40);

        let (lower, upper) = edges(8, &p);
        assert_eq!(lower, 54);
        assert_eq!(upper, 61);

        let (lower, upper) = edges(16, &p);
        assert_eq!(lower, 111);
        assert_eq!(upper, 119);
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2017  Sebastian Wiesner <swiesner@lunaryorn.com>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

//! Provide wordlists.
//!
//! This module contains the built-in list of words, and provides helpers to collect statistics
//! about a wordlist.

// Include the wordlist in the binary
static EFF_WORDLIST: &'static str = include_str!("eff_large_diceware_wordlist.txt");

/// Get the list of built-in words.
///
/// This function returns the [large diceware wordlist from the EFF][1] ([download][2]).  This list
/// contains 7776 common English words between 3 and 9 characters long, at an average of 7
/// characters per word.
///
/// The wordlist is freely available under the CC BY 3.0 US license; see
/// <https://www.eff.org/copyright> for more information about the copyright of the EFF.
///
/// [1]: https://www.eff.org/deeplinks/2016/07/new-wordlists-random-passphrases
/// [2]: https://www.eff.org/files/2016/07/18/eff_large_wordlist.txt
pub fn builtin_words() -> Vec<&'static str> {
    EFF_WORDLIST.lines().collect()
}

/// Provide statistics about a wordlist.
pub struct WordlistStatistics {
    /// The total number of words in the list.
    pub number_of_words: usize,
    /// The length of the shortest word.
    pub min_word_length: usize,
    /// The length of the longest word.
    pub max_word_length: usize,
    /// The average length of a word.
    pub avg_word_length: f64,
    /// The median of the length of all words.
    pub med_word_length: usize,
}

fn median(x: &[usize]) -> usize {
    let n = x.len();
    if n % 2 == 1 {
        x[(n - 1) / 2]
    } else {
        (x[n / 2] + x[(n / 2) - 1]) / 2
    }
}

impl WordlistStatistics {
    /// Collect statistics for the given wordlist.
    pub fn from_words<'a, W, T>(words: W) -> WordlistStatistics
        where W: IntoIterator<Item = &'a T>,
              T: AsRef<str> + 'a
    {
        let mut lengths: Vec<usize> = words
            .into_iter()
            .map(|w| w.as_ref().chars().count())
            .collect();
        if lengths.is_empty() {
            WordlistStatistics {
                number_of_words: 0,
                min_word_length: 0,
                max_word_length: 0,
                avg_word_length: 0.0,
                med_word_length: 0,
            }
        } else {
            lengths.sort();
            let sum_of_lengths: usize = lengths.iter().sum();
            WordlistStatistics {
                number_of_words: lengths.len(),
                min_word_length: lengths[0],
                max_word_length: lengths[lengths.len() - 1],
                avg_word_length: (sum_of_lengths as f64) / lengths.len() as f64,
                med_word_length: median(&lengths),
            }
        }
    }
}

#[cfg(test)]
mod test {
    use super::builtin_words;
    use std::collections::HashSet;

    #[test]
    fn has_7776_words() {
        assert_eq!(builtin_words().len(), 7776);
    }

    #[test]
    fn contains_no_duplicate_words() {
        let mut seen_words = HashSet::with_capacity(8000);
        let mut duplicate_words = HashSet::new();
        for word in builtin_words() {
            if !seen_words.insert(word) {
                duplicate_words.insert(word);
            }
        }

        assert!(duplicate_words.is_empty(),
                "Duplicate words found: {}",
                duplicate_words
                    .into_iter()
                    .collect::<Vec<_>>()
                    .join(" "));
    }

    #[test]
    fn contains_no_empty_words() {
        for word in builtin_words() {
            assert!(!word.is_empty(), "Got empty word");
        }
    }

    #[test]
    fn contains_no_words_with_whitespace() {
        for word in builtin_words() {
            assert!(!word.contains(|c: char| c.is_whitespace()),
                    "Word {} contained whitespace!",
                    word);
        }
    }
}
<commit_msg>Elide redundant 'static lifetime<commit_after>/*
 * Copyright (C) 2017  Sebastian Wiesner <swiesner@lunaryorn.com>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

//! Provide wordlists.
//!
//! This module contains the built-in list of words, and provides helpers to collect statistics
//! about a wordlist.

// Include the wordlist in the binary
static EFF_WORDLIST: &str = include_str!("eff_large_diceware_wordlist.txt");

/// Get the list of built-in words.
///
/// This function returns the [large diceware wordlist from the EFF][1] ([download][2]).  This list
/// contains 7776 common English words between 3 and 9 characters long, at an average of 7
/// characters per word.
///
/// The wordlist is freely available under the CC BY 3.0 US license; see
/// <https://www.eff.org/copyright> for more information about the copyright of the EFF.
///
/// [1]: https://www.eff.org/deeplinks/2016/07/new-wordlists-random-passphrases
/// [2]: https://www.eff.org/files/2016/07/18/eff_large_wordlist.txt
pub fn builtin_words() -> Vec<&'static str> {
    EFF_WORDLIST.lines().collect()
}

/// Provide statistics about a wordlist.
pub struct WordlistStatistics {
    /// The total number of words in the list.
    pub number_of_words: usize,
    /// The length of the shortest word.
    pub min_word_length: usize,
    /// The length of the longest word.
    pub max_word_length: usize,
    /// The average length of a word.
    pub avg_word_length: f64,
    /// The median of the length of all words.
    pub med_word_length: usize,
}

fn median(x: &[usize]) -> usize {
    let n = x.len();
    if n % 2 == 1 {
        x[(n - 1) / 2]
    } else {
        (x[n / 2] + x[(n / 2) - 1]) / 2
    }
}

impl WordlistStatistics {
    /// Collect statistics for the given wordlist.
    pub fn from_words<'a, W, T>(words: W) -> WordlistStatistics
        where W: IntoIterator<Item = &'a T>,
              T: AsRef<str> + 'a
    {
        let mut lengths: Vec<usize> = words
            .into_iter()
            .map(|w| w.as_ref().chars().count())
            .collect();
        if lengths.is_empty() {
            WordlistStatistics {
                number_of_words: 0,
                min_word_length: 0,
                max_word_length: 0,
                avg_word_length: 0.0,
                med_word_length: 0,
            }
        } else {
            lengths.sort();
            let sum_of_lengths: usize = lengths.iter().sum();
            WordlistStatistics {
                number_of_words: lengths.len(),
                min_word_length: lengths[0],
                max_word_length: lengths[lengths.len() - 1],
                avg_word_length: (sum_of_lengths as f64) / lengths.len() as f64,
                med_word_length: median(&lengths),
            }
        }
    }
}

#[cfg(test)]
mod test {
    use super::builtin_words;
    use std::collections::HashSet;

    #[test]
    fn has_7776_words() {
        assert_eq!(builtin_words().len(), 7776);
    }

    #[test]
    fn contains_no_duplicate_words() {
        let mut seen_words = HashSet::with_capacity(8000);
        let mut duplicate_words = HashSet::new();
        for word in builtin_words() {
            if !seen_words.insert(word) {
                duplicate_words.insert(word);
            }
        }

        assert!(duplicate_words.is_empty(),
                "Duplicate words found: {}",
                duplicate_words
                    .into_iter()
                    .collect::<Vec<_>>()
                    .join(" "));
    }

    #[test]
    fn contains_no_empty_words() {
        for word in builtin_words() {
            assert!(!word.is_empty(), "Got empty word");
        }
    }

    #[test]
    fn contains_no_words_with_whitespace() {
        for word in builtin_words() {
            assert!(!word.contains(|c: char| c.is_whitespace()),
                    "Word {} contained whitespace!",
                    word);
        }
    }
}
<|endoftext|>


<commit_before>use x86::__m128i;
use simd_llvm::{simd_shuffle2, simd_shuffle4};
use v128::*;

#[cfg(test)]
use stdsimd_test::assert_instr;

/// Alternatively add and subtract packed single-precision (32-bit)
/// floating-point elements in `a` to/from packed elements in `b`.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(addsubps))]
pub unsafe fn _mm_addsub_ps(a: f32x4, b: f32x4) -> f32x4 {
    addsubps(a, b)
}

/// Alternatively add and subtract packed double-precision (64-bit)
/// floating-point elements in `a` to/from packed elements in `b`.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(addsubpd))]
pub unsafe fn _mm_addsub_pd(a: f64x2, b: f64x2) -> f64x2 {
    addsubpd(a, b)
}

/// Horizontally add adjacent pairs of double-precision (64-bit)
/// floating-point elements in `a` and `b`, and pack the results.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(haddpd))]
pub unsafe fn _mm_hadd_pd(a: f64x2, b: f64x2) -> f64x2 {
    haddpd(a, b)
}

/// Horizontally add adjacent pairs of single-precision (32-bit)
/// floating-point elements in `a` and `b`, and pack the results.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(haddps))]
pub unsafe fn _mm_hadd_ps(a: f32x4, b: f32x4) -> f32x4 {
    haddps(a, b)
}

/// Horizontally subtract adjacent pairs of double-precision (64-bit)
/// floating-point elements in `a` and `b`, and pack the results.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(hsubpd))]
pub unsafe fn _mm_hsub_pd(a: f64x2, b: f64x2) -> f64x2 {
    hsubpd(a, b)
}

/// Horizontally add adjacent pairs of single-precision (32-bit)
/// floating-point elements in `a` and `b`, and pack the results.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(hsubps))]
pub unsafe fn _mm_hsub_ps(a: f32x4, b: f32x4) -> f32x4 {
    hsubps(a, b)
}

/// Load 128-bits of integer data from unaligned memory.
/// This intrinsic may perform better than `_mm_loadu_si128`
/// when the data crosses a cache line boundary.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(lddqu))]
pub unsafe fn _mm_lddqu_si128(mem_addr: *const __m128i) -> __m128i {
    lddqu(mem_addr as *const _)
}

/// Duplicate the low double-precision (64-bit) floating-point element
/// from `a`.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(movddup))]
pub unsafe fn _mm_movedup_pd(a: f64x2) -> f64x2 {
    simd_shuffle2(a, a, [0, 0])
}

/// Duplicate odd-indexed single-precision (32-bit) floating-point elements
/// from `a`.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(movshdup))]
pub unsafe fn _mm_movehdup_ps(a: f32x4) -> f32x4 {
    simd_shuffle4(a, a, [1, 1, 3, 3])
}

#[allow(improper_ctypes)]
extern {
    #[link_name = "llvm.x86.sse3.addsub.ps"]
    fn addsubps(a: f32x4, b: f32x4) -> f32x4;
    #[link_name = "llvm.x86.sse3.addsub.pd"]
    fn addsubpd(a: f64x2, b: f64x2) -> f64x2;
    #[link_name = "llvm.x86.sse3.hadd.pd"]
    fn haddpd(a: f64x2, b: f64x2) -> f64x2;
    #[link_name = "llvm.x86.sse3.hadd.ps"]
    fn haddps(a: f32x4, b: f32x4) -> f32x4;
    #[link_name = "llvm.x86.sse3.hsub.pd"]
    fn hsubpd(a: f64x2, b: f64x2) -> f64x2;
    #[link_name = "llvm.x86.sse3.hsub.ps"]
    fn hsubps(a: f32x4, b: f32x4) -> f32x4;
    #[link_name = "llvm.x86.sse3.ldu.dq"]
    fn lddqu(mem_addr: *const i8) -> __m128i;
}

#[cfg(test)]
mod tests {
    use stdsimd_test::simd_test;

    use v128::*;
    use x86::sse3 as sse3;

    #[simd_test = "sse3"]
    unsafe fn _mm_addsub_ps() {
        let a = f32x4::new(-1.0, 5.0, 0.0, -10.0);
        let b = f32x4::new(-100.0, 20.0, 0.0, -5.0);
        let r = sse3::_mm_addsub_ps(a, b);
        assert_eq!(r, f32x4::new(99.0, 25.0, 0.0, -15.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_addsub_pd() {
        let a = f64x2::new(-1.0, 5.0);
        let b = f64x2::new(-100.0, 20.0);
        let r = sse3::_mm_addsub_pd(a, b);
        assert_eq!(r, f64x2::new(99.0, 25.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_hadd_pd() {
        let a = f64x2::new(-1.0, 5.0);
        let b = f64x2::new(-100.0, 20.0);
        let r = sse3::_mm_hadd_pd(a, b);
        assert_eq!(r, f64x2::new(4.0, -80.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_hadd_ps() {
        let a = f32x4::new(-1.0, 5.0, 0.0, -10.0);
        let b = f32x4::new(-100.0, 20.0, 0.0, -5.0);
        let r = sse3::_mm_hadd_ps(a, b);
        assert_eq!(r, f32x4::new(4.0, -10.0, -80.0, -5.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_hsub_pd() {
        let a = f64x2::new(-1.0, 5.0);
        let b = f64x2::new(-100.0, 20.0);
        let r = sse3::_mm_hsub_pd(a, b);
        assert_eq!(r, f64x2::new(-6.0, -120.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_hsub_ps() {
        let a = f32x4::new(-1.0, 5.0, 0.0, -10.0);
        let b = f32x4::new(-100.0, 20.0, 0.0, -5.0);
        let r = sse3::_mm_hsub_ps(a, b);
        assert_eq!(r, f32x4::new(-6.0, 10.0, -120.0, 5.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_lddqu_si128() {
        let a = i8x16::new(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16);
        let r = sse3::_mm_lddqu_si128(&a);
        assert_eq!(a, r);
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_movedup_pd() {
        let a = f64x2::new(-1.0, 5.0);
        let r = sse3::_mm_movedup_pd(a);
        assert_eq!(r, f64x2::new(-1.0, -1.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_movehdup_ps() {
        let a = f32x4::new(-1.0, 5.0, 0.0, -10.0);
        let r = sse3::_mm_movehdup_ps(a);
        assert_eq!(r, f32x4::new(5.0, 5.0, -10.0, -10.0));
    }
}<commit_msg>sse3: _mm_moveldup_ps<commit_after>use x86::__m128i;
use simd_llvm::{simd_shuffle2, simd_shuffle4};
use v128::*;

#[cfg(test)]
use stdsimd_test::assert_instr;

/// Alternatively add and subtract packed single-precision (32-bit)
/// floating-point elements in `a` to/from packed elements in `b`.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(addsubps))]
pub unsafe fn _mm_addsub_ps(a: f32x4, b: f32x4) -> f32x4 {
    addsubps(a, b)
}

/// Alternatively add and subtract packed double-precision (64-bit)
/// floating-point elements in `a` to/from packed elements in `b`.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(addsubpd))]
pub unsafe fn _mm_addsub_pd(a: f64x2, b: f64x2) -> f64x2 {
    addsubpd(a, b)
}

/// Horizontally add adjacent pairs of double-precision (64-bit)
/// floating-point elements in `a` and `b`, and pack the results.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(haddpd))]
pub unsafe fn _mm_hadd_pd(a: f64x2, b: f64x2) -> f64x2 {
    haddpd(a, b)
}

/// Horizontally add adjacent pairs of single-precision (32-bit)
/// floating-point elements in `a` and `b`, and pack the results.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(haddps))]
pub unsafe fn _mm_hadd_ps(a: f32x4, b: f32x4) -> f32x4 {
    haddps(a, b)
}

/// Horizontally subtract adjacent pairs of double-precision (64-bit)
/// floating-point elements in `a` and `b`, and pack the results.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(hsubpd))]
pub unsafe fn _mm_hsub_pd(a: f64x2, b: f64x2) -> f64x2 {
    hsubpd(a, b)
}

/// Horizontally add adjacent pairs of single-precision (32-bit)
/// floating-point elements in `a` and `b`, and pack the results.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(hsubps))]
pub unsafe fn _mm_hsub_ps(a: f32x4, b: f32x4) -> f32x4 {
    hsubps(a, b)
}

/// Load 128-bits of integer data from unaligned memory.
/// This intrinsic may perform better than `_mm_loadu_si128`
/// when the data crosses a cache line boundary.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(lddqu))]
pub unsafe fn _mm_lddqu_si128(mem_addr: *const __m128i) -> __m128i {
    lddqu(mem_addr as *const _)
}

/// Duplicate the low double-precision (64-bit) floating-point element
/// from `a`.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(movddup))]
pub unsafe fn _mm_movedup_pd(a: f64x2) -> f64x2 {
    simd_shuffle2(a, a, [0, 0])
}

/// Duplicate odd-indexed single-precision (32-bit) floating-point elements
/// from `a`.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(movshdup))]
pub unsafe fn _mm_movehdup_ps(a: f32x4) -> f32x4 {
    simd_shuffle4(a, a, [1, 1, 3, 3])
}

/// Duplicate even-indexed single-precision (32-bit) floating-point elements
/// from `a`.
#[inline(always)]
#[target_feature = "+sse3"]
#[cfg_attr(test, assert_instr(movsldup))]
pub unsafe fn _mm_moveldup_ps(a: f32x4) -> f32x4 {
    simd_shuffle4(a, a, [0, 0, 2, 2])
}

#[allow(improper_ctypes)]
extern {
    #[link_name = "llvm.x86.sse3.addsub.ps"]
    fn addsubps(a: f32x4, b: f32x4) -> f32x4;
    #[link_name = "llvm.x86.sse3.addsub.pd"]
    fn addsubpd(a: f64x2, b: f64x2) -> f64x2;
    #[link_name = "llvm.x86.sse3.hadd.pd"]
    fn haddpd(a: f64x2, b: f64x2) -> f64x2;
    #[link_name = "llvm.x86.sse3.hadd.ps"]
    fn haddps(a: f32x4, b: f32x4) -> f32x4;
    #[link_name = "llvm.x86.sse3.hsub.pd"]
    fn hsubpd(a: f64x2, b: f64x2) -> f64x2;
    #[link_name = "llvm.x86.sse3.hsub.ps"]
    fn hsubps(a: f32x4, b: f32x4) -> f32x4;
    #[link_name = "llvm.x86.sse3.ldu.dq"]
    fn lddqu(mem_addr: *const i8) -> __m128i;
}

#[cfg(test)]
mod tests {
    use stdsimd_test::simd_test;

    use v128::*;
    use x86::sse3 as sse3;

    #[simd_test = "sse3"]
    unsafe fn _mm_addsub_ps() {
        let a = f32x4::new(-1.0, 5.0, 0.0, -10.0);
        let b = f32x4::new(-100.0, 20.0, 0.0, -5.0);
        let r = sse3::_mm_addsub_ps(a, b);
        assert_eq!(r, f32x4::new(99.0, 25.0, 0.0, -15.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_addsub_pd() {
        let a = f64x2::new(-1.0, 5.0);
        let b = f64x2::new(-100.0, 20.0);
        let r = sse3::_mm_addsub_pd(a, b);
        assert_eq!(r, f64x2::new(99.0, 25.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_hadd_pd() {
        let a = f64x2::new(-1.0, 5.0);
        let b = f64x2::new(-100.0, 20.0);
        let r = sse3::_mm_hadd_pd(a, b);
        assert_eq!(r, f64x2::new(4.0, -80.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_hadd_ps() {
        let a = f32x4::new(-1.0, 5.0, 0.0, -10.0);
        let b = f32x4::new(-100.0, 20.0, 0.0, -5.0);
        let r = sse3::_mm_hadd_ps(a, b);
        assert_eq!(r, f32x4::new(4.0, -10.0, -80.0, -5.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_hsub_pd() {
        let a = f64x2::new(-1.0, 5.0);
        let b = f64x2::new(-100.0, 20.0);
        let r = sse3::_mm_hsub_pd(a, b);
        assert_eq!(r, f64x2::new(-6.0, -120.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_hsub_ps() {
        let a = f32x4::new(-1.0, 5.0, 0.0, -10.0);
        let b = f32x4::new(-100.0, 20.0, 0.0, -5.0);
        let r = sse3::_mm_hsub_ps(a, b);
        assert_eq!(r, f32x4::new(-6.0, 10.0, -120.0, 5.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_lddqu_si128() {
        let a = i8x16::new(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16);
        let r = sse3::_mm_lddqu_si128(&a);
        assert_eq!(a, r);
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_movedup_pd() {
        let a = f64x2::new(-1.0, 5.0);
        let r = sse3::_mm_movedup_pd(a);
        assert_eq!(r, f64x2::new(-1.0, -1.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_movehdup_ps() {
        let a = f32x4::new(-1.0, 5.0, 0.0, -10.0);
        let r = sse3::_mm_movehdup_ps(a);
        assert_eq!(r, f32x4::new(5.0, 5.0, -10.0, -10.0));
    }

    #[simd_test = "sse3"]
    unsafe fn _mm_moveldup_ps() {
        let a = f32x4::new(-1.0, 5.0, 0.0, -10.0);
        let r = sse3::_mm_moveldup_ps(a);
        assert_eq!(r, f32x4::new(-1.0, -1.0, 0.0, 0.0));
    }
}<|endoftext|>